Skip to content
2 changes: 2 additions & 0 deletions src/main.ts
Original file line number Diff line number Diff line change
Expand Up @@ -407,6 +407,8 @@ export {
*/
export * as visionApi from './gen/service/vision/v1/vision_pb';

export { WorldStateStoreClient } from './services/world-state-store';

export {
GenericClient as GenericServiceClient,
type Generic as GenericService,
Expand Down
2 changes: 2 additions & 0 deletions src/services/world-state-store.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
export { WorldStateStoreClient } from './world-state-store/client';
export type { WorldStateStore } from './world-state-store/world-state-store';
130 changes: 130 additions & 0 deletions src/services/world-state-store/client.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
// @vitest-environment happy-dom

import { createClient, createRouterTransport } from '@connectrpc/connect';
import { Struct } from '@bufbuild/protobuf';
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { WorldStateStoreService } from '../../gen/service/worldstatestore/v1/world_state_store_connect';
import {
GetTransformResponse,
ListUUIDsResponse,
StreamTransformChangesResponse,
} from '../../gen/service/worldstatestore/v1/world_state_store_pb';
import { RobotClient } from '../../robot';
import { WorldStateStoreClient } from './client';
import { TransformChangeType } from '../../gen/service/worldstatestore/v1/world_state_store_pb';
import { Transform, PoseInFrame, Pose } from '../../gen/common/v1/common_pb';

vi.mock('../../robot');

const worldStateStoreClientName = 'test-world-state-store';

let worldStateStore: WorldStateStoreClient;

const mockUuids = [new Uint8Array([1, 2, 3, 4]), new Uint8Array([5, 6, 7, 8])];
const mockTransform = new Transform({
referenceFrame: 'test-frame',
poseInObserverFrame: new PoseInFrame({
referenceFrame: 'observer-frame',
pose: new Pose({
x: 10,
y: 20,
z: 30,
oX: 0,
oY: 0,
oZ: 1,
theta: 90,
}),
}),
uuid: mockUuids[0],
});

describe('WorldStateStoreClient Tests', () => {
beforeEach(() => {
const mockTransport = createRouterTransport(({ service }) => {
service(WorldStateStoreService, {
listUUIDs: () => new ListUUIDsResponse({ uuids: mockUuids }),
getTransform: () =>
new GetTransformResponse({ transform: mockTransform }),
streamTransformChanges: async function* mockStream() {
// Add await to satisfy linter
await Promise.resolve();
yield new StreamTransformChangesResponse({
changeType: TransformChangeType.ADDED,
transform: mockTransform,
});
yield new StreamTransformChangesResponse({
changeType: TransformChangeType.UPDATED,
transform: mockTransform,
updatedFields: { paths: ['pose_in_observer_frame'] },
});
},
doCommand: () => ({ result: Struct.fromJson({ success: true }) }),
});
});

RobotClient.prototype.createServiceClient = vi
.fn()
.mockImplementation(() =>
createClient(WorldStateStoreService, mockTransport)
);
worldStateStore = new WorldStateStoreClient(
new RobotClient('host'),
worldStateStoreClientName
);
});

describe('listUUIDs', () => {
it('returns all transform UUIDs', async () => {
const expected = mockUuids;

await expect(worldStateStore.listUUIDs()).resolves.toStrictEqual(
expected
);
});
});

describe('getTransform', () => {
it('returns a transform by UUID', async () => {
const uuid = new Uint8Array([1, 2, 3, 4]);
const expected = mockTransform;

await expect(worldStateStore.getTransform(uuid)).resolves.toStrictEqual(
expected
);
});
});

describe('streamTransformChanges', () => {
it('streams transform changes', async () => {
const stream = worldStateStore.streamTransformChanges();
const results = [];

for await (const result of stream) {
results.push(result);
}

expect(results).toHaveLength(2);
expect(results[0]).toEqual({
changeType: TransformChangeType.ADDED,
transform: mockTransform,
updatedFields: undefined,
});
expect(results[1]).toEqual({
changeType: TransformChangeType.UPDATED,
transform: mockTransform,
updatedFields: { paths: ['pose_in_observer_frame'] },
});
});
});

describe('doCommand', () => {
it('executes arbitrary commands', async () => {
const command = Struct.fromJson({ test: 'value' });
const expected = { success: true };

await expect(worldStateStore.doCommand(command)).resolves.toStrictEqual(
expected
);
});
});
});
95 changes: 95 additions & 0 deletions src/services/world-state-store/client.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
import { Struct, type JsonValue } from '@bufbuild/protobuf';
import type { CallOptions, Client } from '@connectrpc/connect';
import { WorldStateStoreService } from '../../gen/service/worldstatestore/v1/world_state_store_connect';
import {
GetTransformRequest,
ListUUIDsRequest,
StreamTransformChangesRequest,
} from '../../gen/service/worldstatestore/v1/world_state_store_pb';
import type { RobotClient } from '../../robot';
import type { Options } from '../../types';
import { doCommandFromClient } from '../../utils';
import type { WorldStateStore } from './world-state-store';

/**
* A gRPC-web client for a WorldStateStore service.
*
* @group Clients
*/
export class WorldStateStoreClient implements WorldStateStore {
private client: Client<typeof WorldStateStoreService>;
public readonly name: string;
private readonly options: Options;
public callOptions: CallOptions = { headers: {} as Record<string, string> };

constructor(client: RobotClient, name: string, options: Options = {}) {
this.client = client.createServiceClient(WorldStateStoreService);
this.name = name;
this.options = options;
}

async listUUIDs(extra = {}, callOptions = this.callOptions) {
const request = new ListUUIDsRequest({
name: this.name,
extra: Struct.fromJson(extra),
});

this.options.requestLogger?.(request);

const response = await this.client.listUUIDs(request, callOptions);
return response.uuids;
}

async getTransform(
uuid: Uint8Array,
extra = {},
callOptions = this.callOptions
) {
const request = new GetTransformRequest({
name: this.name,
uuid,
extra: Struct.fromJson(extra),
});

this.options.requestLogger?.(request);

const response = await this.client.getTransform(request, callOptions);
if (!response.transform) {
throw new Error('No transform returned from server');
}

return response.transform;
}

async *streamTransformChanges(extra = {}, callOptions = this.callOptions) {
const request = new StreamTransformChangesRequest({
name: this.name,
extra: Struct.fromJson(extra),
});

this.options.requestLogger?.(request);

const stream = this.client.streamTransformChanges(request, callOptions);

for await (const response of stream) {
yield {
changeType: response.changeType,
transform: response.transform,
updatedFields: response.updatedFields,
};
}
}

async doCommand(
command: Struct,
callOptions = this.callOptions
): Promise<JsonValue> {
return doCommandFromClient(
this.client.doCommand,
this.name,
command,
this.options,
callOptions
);
}
}
2 changes: 2 additions & 0 deletions src/services/world-state-store/types.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
import type { PlainMessage } from '@bufbuild/protobuf';
import type { Geometry } from '../../gen/common/v1/common_pb';
78 changes: 78 additions & 0 deletions src/services/world-state-store/world-state-store.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
import type { Struct } from '@bufbuild/protobuf';
import type { Transform, Resource } from '../../types';
import type { TransformChangeType } from '../../gen/service/worldstatestore/v1/world_state_store_pb';

/**
* A service that manages world state transforms, allowing storage and retrieval
* of spatial relationships between reference frames.
*/
export interface WorldStateStore extends Resource {
/**
* ListUUIDs returns all world state transform UUIDs.
*
* @example
*
* ```ts
* const worldStateStore = new VIAM.WorldStateStoreClient(
* machine,
* 'builtin'
* );
*
* // Get all transform UUIDs
* const uuids = await worldStateStore.listUUIDs();
* ```
*
* @param extra - Additional arguments to the method
*/
listUUIDs: (extra?: Struct) => Promise<Uint8Array[]>;

/**
* GetTransform returns a world state transform by UUID.
*
* @example
*
* ```ts
* const worldStateStore = new VIAM.WorldStateStoreClient(
* machine,
* 'builtin'
* );
*
* // Get a specific transform by UUID
* const transform = await worldStateStore.getTransform(uuid);
* ```
*
* @param uuid - The UUID of the transform to retrieve
* @param extra - Additional arguments to the method
*/
getTransform: (uuid: Uint8Array, extra?: Struct) => Promise<Transform>;

/**
* StreamTransformChanges streams changes to world state transforms.
*
* @example
*
* ```ts
* const worldStateStore = new VIAM.WorldStateStoreClient(
* machine,
* 'builtin'
* );
*
* // Stream transform changes
* const stream = worldStateStore.streamTransformChanges();
* for await (const change of stream) {
* console.log(
* 'Transform change:',
* change.changeType,
* change.transform
* );
* }
* ```
*
* @param extra - Additional arguments to the method
*/
streamTransformChanges: (extra?: Struct) => AsyncIterable<{
changeType: TransformChangeType;
transform?: Transform;
updatedFields?: { paths: string[] } | undefined;
}>;
}
Loading