diff --git a/package-lock.json b/package-lock.json index 0f1c74796..a9b93a99f 100644 --- a/package-lock.json +++ b/package-lock.json @@ -37,6 +37,7 @@ "prettier-plugin-jsdoc": "^1.1.1", "typedoc": "^0.27.9", "typescript": "^5.3.3", + "uuid-tool": "^2.0.3", "vite": "^5.0.10", "vitest": "^1.1.0", "yaml": "^2.3.3" @@ -5652,6 +5653,13 @@ "dev": true, "license": "MIT" }, + "node_modules/uuid-tool": { + "version": "2.0.3", + "resolved": "https://registry.npmjs.org/uuid-tool/-/uuid-tool-2.0.3.tgz", + "integrity": "sha512-En9cTJm+bLgjjfoIVEW91f0Irr0HZ9mqzK0BwntzZzOg82bOUEIzIw8pNwvC+dhVnnJ8x+fQV8TQ+q6te1ximA==", + "dev": true, + "license": "MIT" + }, "node_modules/validate-npm-package-license": { "version": "3.0.4", "resolved": "https://registry.npmjs.org/validate-npm-package-license/-/validate-npm-package-license-3.0.4.tgz", diff --git a/package.json b/package.json index 836cfd1c5..29fdd4be9 100644 --- a/package.json +++ b/package.json @@ -80,6 +80,7 @@ "prettier-plugin-jsdoc": "^1.1.1", "typedoc": "^0.27.9", "typescript": "^5.3.3", + "uuid-tool": "^2.0.3", "vite": "^5.0.10", "vitest": "^1.1.0", "yaml": "^2.3.3" diff --git a/src/main.ts b/src/main.ts index 2e9c3b31e..2dc59e00d 100644 --- a/src/main.ts +++ b/src/main.ts @@ -407,6 +407,8 @@ export { */ export * as visionApi from './gen/service/vision/v1/vision_pb'; +export * from './services/world-state-store'; + export { GenericClient as GenericServiceClient, type Generic as GenericService, diff --git a/src/robot/client.ts b/src/robot/client.ts index 08c0b6e30..1967e7440 100644 --- a/src/robot/client.ts +++ b/src/robot/client.ts @@ -36,6 +36,7 @@ import type { Robot } from './robot'; import SessionManager from './session-manager'; import { MLModelService } from '../gen/service/mlmodel/v1/mlmodel_connect'; import type { AccessToken, Credential } from '../main'; +import { WorldStateStoreService } from '../gen/service/worldstatestore/v1/world_state_store_connect'; import { assertExists } from '../assert'; interface ICEServer { @@ -238,6 +239,10 @@ export class RobotClient extends EventDispatcher implements Robot { private slamServiceClient: Client | undefined; + private worldStateStoreServiceClient: + | Client + | undefined; + private currentRetryAttempt = 0; constructor( @@ -478,6 +483,13 @@ export class RobotClient extends EventDispatcher implements Robot { return this.slamServiceClient; } + get worldStateStoreService() { + this.worldStateStoreServiceClient ??= this.createServiceClient( + WorldStateStoreService + ); + return this.worldStateStoreServiceClient; + } + createServiceClient(svcType: T): Client { assertExists(this.clientTransport, RobotClient.notConnectedYetStr); return createClient(svcType, this.clientTransport); diff --git a/src/services/world-state-store/client.spec.ts b/src/services/world-state-store/client.spec.ts new file mode 100644 index 000000000..d769c75a6 --- /dev/null +++ b/src/services/world-state-store/client.spec.ts @@ -0,0 +1,132 @@ +// @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'; +import { transformWithUUID, uuidToString } from './world-state-store'; + +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.map((uuid) => uuidToString(uuid)); + + await expect(worldStateStore.listUUIDs()).resolves.toStrictEqual( + expected + ); + }); + }); + + describe('getTransform', () => { + it('returns a transform by UUID', async () => { + const uuid = '01020304'; + const expected = mockTransform; + + await expect(worldStateStore.getTransform(uuid)).resolves.toStrictEqual({ + ...expected, + uuidString: uuid, + }); + }); + }); + + 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: transformWithUUID(mockTransform), + updatedFields: undefined, + }); + expect(results[1]).toEqual({ + changeType: TransformChangeType.UPDATED, + transform: transformWithUUID(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 + ); + }); + }); +}); diff --git a/src/services/world-state-store/client.ts b/src/services/world-state-store/client.ts new file mode 100644 index 000000000..4cedb8eea --- /dev/null +++ b/src/services/world-state-store/client.ts @@ -0,0 +1,103 @@ +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'; +import { + transformWithUUID, + uuidFromString, + uuidToString, +} from './world-state-store'; +import type { TransformChangeEvent } from './types'; + +/** + * A gRPC-web client for a WorldStateStore service. + * + * @group Clients + */ +export class WorldStateStoreClient implements WorldStateStore { + private client: Client; + public readonly name: string; + private readonly options: Options; + public callOptions: CallOptions = { headers: {} as Record }; + + 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.map((uuid) => uuidToString(uuid)); + } + + async getTransform(uuid: string, extra = {}, callOptions = this.callOptions) { + const request = new GetTransformRequest({ + name: this.name, + uuid: uuidFromString(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 transformWithUUID(response.transform); + } + + async *streamTransformChanges( + extra = {}, + callOptions = this.callOptions + ): AsyncGenerator { + 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) { + if (!response.transform) { + continue; + } + + yield { + ...response, + transform: transformWithUUID(response.transform), + }; + } + } + + async doCommand( + command: Struct, + callOptions = this.callOptions + ): Promise { + return doCommandFromClient( + this.client.doCommand, + this.name, + command, + this.options, + callOptions + ); + } +} diff --git a/src/services/world-state-store/index.ts b/src/services/world-state-store/index.ts new file mode 100644 index 000000000..e1d42c583 --- /dev/null +++ b/src/services/world-state-store/index.ts @@ -0,0 +1,6 @@ +export { WorldStateStoreClient } from './client'; +export type { WorldStateStore } from './world-state-store'; +export { transformWithUUID } from './world-state-store'; +export type { TransformChangeEvent, TransformWithUUID } from './types'; + +export { TransformChangeType } from '../../gen/service/worldstatestore/v1/world_state_store_pb'; diff --git a/src/services/world-state-store/types.ts b/src/services/world-state-store/types.ts new file mode 100644 index 000000000..9e3e41b90 --- /dev/null +++ b/src/services/world-state-store/types.ts @@ -0,0 +1,14 @@ +import type { PlainMessage } from '@bufbuild/protobuf'; +import type { Transform } from '../../gen/common/v1/common_pb'; +import type { StreamTransformChangesResponse } from '../../gen/service/worldstatestore/v1/world_state_store_pb'; + +export interface TransformWithUUID extends PlainMessage { + uuidString: string; +} + +export type TransformChangeEvent = Omit< + PlainMessage, + 'transform' +> & { + transform: TransformWithUUID | undefined; +}; diff --git a/src/services/world-state-store/world-state-store.ts b/src/services/world-state-store/world-state-store.ts new file mode 100644 index 000000000..94a01a75d --- /dev/null +++ b/src/services/world-state-store/world-state-store.ts @@ -0,0 +1,90 @@ +import type { Struct } from '@bufbuild/protobuf'; +import type { Transform, Resource } from '../../types'; +import type { TransformChangeEvent, TransformWithUUID } from './types'; +import { UuidTool } from 'uuid-tool'; + +/** + * 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; + + /** + * 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: string, extra?: Struct) => Promise; + + /** + * StreamTransformChanges streams changes to world state transforms. + * + * The returned transform will only contain the fields that have changed. + * + * @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 + ) => AsyncGenerator; +} + +export const uuidToString = (uuid: Uint8Array) => UuidTool.toString([...uuid]); +export const uuidFromString = (uuid: string): Uint8Array => + new Uint8Array(UuidTool.toBytes(uuid)); + +export const transformWithUUID = (transform: Transform): TransformWithUUID => { + return { + ...transform, + uuidString: uuidToString(transform.uuid), + }; +};