From f7bbf1d8424ab8ae02a1a957abfcb592a91e74d5 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 8 Oct 2026 16:05:15 +0200 Subject: [PATCH] feat(shared): migrate genie to modular @databricks/sdk-genie Route the facade's genie accessor through a lazily built modular GenieClient (sdk-genie 0.54.0, compatible with sdk-core/auth/options 0.51.0) and port the genie connector to its camelCase model. - Poll getConversationMessage ourselves instead of using the SDK waiter: the modular wait() has no onProgress, which the SSE status events need. Cadence, timeout and the "failed to reach COMPLETED state" errors match the legacy waiter. - Convert the query statement response back to the snake_case wire shape with bigint counts as numbers, so the SSE contract and JSON.stringify keep working. - Patch sdk-genie: its model types data_array as protobuf ListValue[], so the unmarshal rejected every real JSON_ARRAY query result. - Match RESOURCE_DOES_NOT_EXIST on ApiError .code / .errorCode too. The CLI's listSpaces/getSpace calls stay on the legacy client. Co-authored-by: Isaac Signed-off-by: MarioCadenas # Conflicts: # packages/shared/src/workspace-client/client.ts # packages/shared/src/workspace-client/modular.ts # packages/shared/src/workspace-client/tests/modular.test.ts # packages/shared/src/workspace-client/types.ts --- .../server/testing-kit.integration.test.ts | 2 +- .../api/appkit/Interface.WorkspaceClient.md | 4 +- docs/docs/plugins/testing.md | 4 +- knip.json | 1 + packages/appkit/package.json | 1 + .../appkit/src/connectors/genie/client.ts | 212 ++++++--- .../src/connectors/genie/tests/client.test.ts | 352 ++++++++------- .../src/plugins/genie/tests/genie.test.ts | 417 ++++++++++-------- .../tests/mock-workspace-client.test.ts | 2 +- packages/appkit/src/workspace-client/index.ts | 9 +- packages/shared/package.json | 1 + .../shared/src/workspace-client/client.ts | 11 +- packages/shared/src/workspace-client/index.ts | 2 - .../shared/src/workspace-client/legacy.ts | 8 +- .../shared/src/workspace-client/modular.ts | 17 + .../workspace-client/tests/modular.test.ts | 33 +- packages/shared/src/workspace-client/types.ts | 7 +- patches/@databricks__sdk-genie@0.54.0.patch | 26 ++ pnpm-lock.yaml | 20 + pnpm-workspace.yaml | 1 + 20 files changed, 686 insertions(+), 444 deletions(-) create mode 100644 patches/@databricks__sdk-genie@0.54.0.patch diff --git a/apps/dev-playground/server/testing-kit.integration.test.ts b/apps/dev-playground/server/testing-kit.integration.test.ts index 5d0443eca..fa644ba4f 100644 --- a/apps/dev-playground/server/testing-kit.integration.test.ts +++ b/apps/dev-playground/server/testing-kit.integration.test.ts @@ -68,7 +68,7 @@ describe("createMockWorkspaceClient — build the fake client yourself", () => { }); expect(getMock(client, "jobs.getRun")).toHaveBeenCalledWith({ run_id: 1 }); await expect( - client.genie.getMessage({ id: "m-1" }), + client.genie.genieGetConversationMessage({ messageId: "m-1" }), ).resolves.toBeUndefined(); }); }); diff --git a/docs/docs/api/appkit/Interface.WorkspaceClient.md b/docs/docs/api/appkit/Interface.WorkspaceClient.md index 26a680581..9d3fec157 100644 --- a/docs/docs/api/appkit/Interface.WorkspaceClient.md +++ b/docs/docs/api/appkit/Interface.WorkspaceClient.md @@ -56,10 +56,10 @@ UC Volumes / Files API. ### genie ```ts -readonly genie: GenieService; +readonly genie: GenieClient; ``` -Genie / dashboards. +Genie (modular SDK). *** diff --git a/docs/docs/plugins/testing.md b/docs/docs/plugins/testing.md index 6d06d355b..7a7797c01 100644 --- a/docs/docs/plugins/testing.md +++ b/docs/docs/plugins/testing.md @@ -343,7 +343,7 @@ await mock.attach(plugin); ``` `options` is: -- `responses` — seed the mock workspace client with responses keyed by dotted path (`"jobs.getRun"`, `"genie.getMessage"`). A value can be static or a function of call arguments and the abort signal. +- `responses` — seed the mock workspace client with responses keyed by dotted path (`"jobs.getRun"`, `"genie.genieGetConversationMessage"`). A value can be static or a function of call arguments and the abort signal. - `env` — set environment variables scoped to the test; they are restored on plugin detach. - `strict` — throw if a handler calls an undeclared workspace-client path (instead of silently resolving `undefined`). The built-in defaults still count as declared. @@ -513,7 +513,7 @@ const client = createMockWorkspaceClient({ }); await client.jobs.getRun({ run_id: 1 }); // → { state: "TERMINATED" } -await client.genie.getMessage({ id: "m-1" }); // → undefined, does not throw +await client.genie.genieGetConversationMessage({ messageId: "m-1" }); // → undefined, does not throw ``` `createTestApp` installs one of these for you, so reach for it directly only when you're driving a plugin through `createTestPluginContext` or `mockServiceContext`. diff --git a/knip.json b/knip.json index b46121129..c57c8be24 100644 --- a/knip.json +++ b/knip.json @@ -13,6 +13,7 @@ "@databricks/sdk-auth", "@databricks/sdk-core", "@databricks/sdk-experimental", + "@databricks/sdk-genie", "@databricks/sdk-options", "@databricks/sdk-scim", "@databricks/sdk-statementexecution", diff --git a/packages/appkit/package.json b/packages/appkit/package.json index fb1ce78ec..a72554fcd 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -74,6 +74,7 @@ "@databricks/sdk-auth": "0.51.0", "@databricks/sdk-core": "0.51.0", "@databricks/sdk-experimental": "0.17.0", + "@databricks/sdk-genie": "0.54.0", "@databricks/sdk-options": "0.51.0", "@databricks/sdk-scim": "0.51.0", "@databricks/sdk-statementexecution": "0.52.0", diff --git a/packages/appkit/src/connectors/genie/client.ts b/packages/appkit/src/connectors/genie/client.ts index ec8a987d1..ed77a5998 100644 --- a/packages/appkit/src/connectors/genie/client.ts +++ b/packages/appkit/src/connectors/genie/client.ts @@ -1,13 +1,7 @@ import { createLogger } from "../../logging"; -import { - type GenieMessage, - Time, - TimeUnits, - type Waiter, - type WorkspaceClient, -} from "../../workspace-client"; +import type { GenieMessage, WorkspaceClient } from "../../workspace-client"; import { genieConnectorDefaults } from "./defaults"; -import { pollWaiter } from "./poll-waiter"; +import { type Pollable, pollWaiter } from "./poll-waiter"; import type { GenieAttachmentResponse, GenieConversationHistoryResponse, @@ -26,7 +20,11 @@ const GenieErrors = { QUERY_RESULT_FAILED: "Failed to fetch query result", } as const; -type CreateMessageWaiter = Waiter; +type CreateMessageWaiter = Pollable; + +// Legacy SDK waiter defaults, kept so polling cadence is unchanged. +const DEFAULT_WAIT_TIMEOUT_MS = 10 * 60_000; +const MAX_POLL_INTERVAL_MS = 10_000; interface GenieConnectorConfig { timeout?: number; @@ -35,38 +33,72 @@ interface GenieConnectorConfig { function mapAttachments(message: GenieMessage): GenieAttachmentResponse[] { return ( - message.attachments?.map((att) => ({ - attachmentId: att.attachment_id, - query: att.query - ? { - title: att.query.title, - description: att.query.description, - query: att.query.query, - statementId: att.query.statement_id, - } - : undefined, - text: att.text ? { content: att.text.content } : undefined, - suggestedQuestions: att.suggested_questions?.questions, + message.attachments?.map(({ attachmentId, attachment }) => ({ + attachmentId, + query: + attachment?.$case === "query" + ? { + title: attachment.query.title, + description: attachment.query.description, + query: attachment.query.query, + statementId: attachment.query.statementId, + } + : undefined, + text: + attachment?.$case === "text" + ? { content: attachment.text.content } + : undefined, + suggestedQuestions: + attachment?.$case === "suggestedQuestions" + ? attachment.suggestedQuestions.questions + : undefined, })) ?? [] ); } function toMessageResponse(message: GenieMessage): GenieMessageResponse { return { - messageId: message.message_id, - conversationId: message.conversation_id, - spaceId: message.space_id, + messageId: message.messageId ?? "", + conversationId: message.conversationId ?? "", + spaceId: message.spaceId ?? "", status: message.status ?? "COMPLETED", - content: message.content, + content: message.content ?? "", attachments: mapAttachments(message), error: message.error?.error, }; } +/** + * The modular SDK returns the statement response camelCased with int64 fields as + * `bigint`. The SSE contract (`GenieStatementResponse`, read by appkit-ui) is the + * raw snake_case API shape, and `JSON.stringify` throws on `bigint`, so convert + * back here. Recurses into objects only; `data_array` rows pass through as-is. + */ +function toWireShape(value: unknown): unknown { + if (typeof value === "bigint") return Number(value); + if (Array.isArray(value)) return value.map(toWireShape); + if (value === null || typeof value !== "object") return value; + return Object.fromEntries( + Object.entries(value).map(([key, v]) => [ + key.replace(/[A-Z]/g, (c) => `_${c.toLowerCase()}`), + toWireShape(v), + ]), + ); +} + function classifyGenieError(error: unknown): string { const message = error instanceof Error ? error.message : String(error); + // Modular ApiError carries the code on `.code`, legacy on `.errorCode`. + const { code, errorCode } = (error ?? {}) as { + code?: unknown; + errorCode?: unknown; + }; - if (message.includes("RESOURCE_DOES_NOT_EXIST")) { + if ( + code === "RESOURCE_DOES_NOT_EXIST" || + errorCode === "RESOURCE_DOES_NOT_EXIST" || + message.includes("RESOURCE_DOES_NOT_EXIST") + ) { return GenieErrors.SPACE_ACCESS_DENIED; } @@ -100,26 +132,74 @@ export class GenieConnector { conversationId: string; messageId: string; }> { - if (conversationId) { - const waiter = await workspaceClient.genie.createMessage({ - space_id: spaceId, - conversation_id: conversationId, - content, - }); - return { - messageWaiter: waiter, - conversationId, - messageId: waiter.message_id ?? "", - }; - } - const start = await workspaceClient.genie.startConversation({ - space_id: spaceId, - content, - }); + const started = conversationId + ? await workspaceClient.genie.genieCreateConversationMessage({ + spaceId, + conversationId, + content, + }) + : await workspaceClient.genie.genieStartConversation({ + spaceId, + content, + }); return { - messageWaiter: start as unknown as CreateMessageWaiter, - conversationId: start.conversation_id, - messageId: start.message_id, + messageWaiter: this.messagePoller( + workspaceClient, + spaceId, + started.conversationId, + started.messageId, + ), + conversationId: started.conversationId, + messageId: started.messageId, + }; + } + + /** + * Polls `getConversationMessage` until COMPLETED. Replaces the SDK waiter: the + * modular `wait()` has no `onProgress`, which the SSE `status` events need. It + * mirrors the legacy waiter: progress on every poll, backoff of `attempt` + * seconds + 50-750ms jitter capped at 10s, a 10 minute default timeout, and the + * same `failed to reach COMPLETED state` errors `classifyGenieError` matches. + */ + private messagePoller( + workspaceClient: WorkspaceClient, + spaceId: string, + conversationId: string, + messageId: string, + ): CreateMessageWaiter { + return { + async wait(options) { + const timeout = + typeof options?.timeout === "number" + ? options.timeout + : DEFAULT_WAIT_TIMEOUT_MS; + const deadline = Date.now() + timeout; + let lastStatus: string | undefined; + for (let attempt = 1; Date.now() < deadline; attempt++) { + const message = + await workspaceClient.genie.genieGetConversationMessage({ + spaceId, + conversationId, + messageId, + }); + await options?.onProgress?.(message); + lastStatus = message.status; + if (lastStatus === "COMPLETED") return message; + if (lastStatus === "FAILED") { + throw new Error("failed to reach COMPLETED state, got FAILED"); + } + const jitter = 50 + Math.random() * 700; + await new Promise((resolve) => + setTimeout( + resolve, + Math.min(attempt * 1000 + jitter, MAX_POLL_INTERVAL_MS), + ), + ); + } + throw new Error( + `timed out: failed to reach COMPLETED state, got ${lastStatus}`, + ); + }, }; } @@ -128,9 +208,7 @@ export class GenieConnector { options?: { timeout?: number }, ): Promise { const timeout = options?.timeout ?? this.config.timeout; - const waitOptions = - timeout > 0 ? { timeout: new Time(timeout, TimeUnits.milliseconds) } : {}; - return messageWaiter.wait(waitOptions); + return messageWaiter.wait(timeout > 0 ? { timeout } : {}); } async listConversationMessages( @@ -145,18 +223,18 @@ export class GenieConnector { const pageSize = options?.pageSize ?? genieConnectorDefaults.initialPageSize; - const response = await workspaceClient.genie.listConversationMessages({ - space_id: spaceId, - conversation_id: conversationId, - page_size: pageSize, - ...(options?.pageToken ? { page_token: options.pageToken } : {}), + const response = await workspaceClient.genie.genieListConversationMessages({ + spaceId, + conversationId, + pageSize, + ...(options?.pageToken ? { pageToken: options.pageToken } : {}), }); const messages = (response.messages ?? []).reverse().map(toMessageResponse); return { messages, - nextPageToken: response.next_page_token ?? null, + nextPageToken: response.nextPageToken ?? null, }; } @@ -169,13 +247,13 @@ export class GenieConnector { _signal?: AbortSignal, ): Promise { const response = - await workspaceClient.genie.getMessageAttachmentQueryResult({ - space_id: spaceId, - conversation_id: conversationId, - message_id: messageId, - attachment_id: attachmentId, + await workspaceClient.genie.genieGetMessageAttachmentQueryResult({ + spaceId, + conversationId, + messageId, + attachmentId, }); - return response.statement_response as GenieStatementResponse; + return toWireShape(response.statementResponse) as GenieStatementResponse; } async *streamSendMessage( @@ -206,10 +284,7 @@ export class GenieConnector { const timeout = options?.timeout != null ? options.timeout : this.config.timeout; - const waitOptions = - timeout > 0 - ? { timeout: new Time(timeout, TimeUnits.milliseconds) } - : {}; + const waitOptions = timeout > 0 ? { timeout } : {}; let completedMessage!: GenieMessage; for await (const event of pollWaiter(messageWaiter, waitOptions)) { @@ -402,11 +477,10 @@ export class GenieConnector { while (true) { if (signal?.aborted) return; - const message = await workspaceClient.genie.getMessage({ - space_id: spaceId, - conversation_id: conversationId, - message_id: messageId, - }); + const message = await workspaceClient.genie.genieGetConversationMessage( + { spaceId, conversationId, messageId }, + { signal }, + ); if (message.status && message.status !== lastStatus) { lastStatus = message.status; diff --git a/packages/appkit/src/connectors/genie/tests/client.test.ts b/packages/appkit/src/connectors/genie/tests/client.test.ts index 50ae69f2b..287bef9d1 100644 --- a/packages/appkit/src/connectors/genie/tests/client.test.ts +++ b/packages/appkit/src/connectors/genie/tests/client.test.ts @@ -20,9 +20,9 @@ async function collect( function makeGenieMessage(overrides: Partial = {}): GenieMessage { return { - message_id: "msg-1", - conversation_id: "conv-1", - space_id: "space-1", + messageId: "msg-1", + conversationId: "conv-1", + spaceId: "space-1", status: "COMPLETED", content: "Hello from Genie", attachments: [], @@ -36,12 +36,15 @@ function makeGenieMessageWithQuery( return makeGenieMessage({ attachments: [ { - attachment_id: "att-1", - query: { - title: "Sales Query", - description: "Total sales", - query: "SELECT sum(amount) FROM sales", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Sales Query", + description: "Total sales", + query: "SELECT sum(amount) FROM sales", + statementId: "stmt-1", + }, }, }, ], @@ -53,37 +56,26 @@ function makeGenieMessageWithQuery( function createMockWorkspaceClient() { return { genie: { - startConversation: vi.fn(), - createMessage: vi.fn(), - getMessage: vi.fn(), - listConversationMessages: vi.fn(), - getMessageAttachmentQueryResult: vi.fn(), + genieStartConversation: vi.fn(), + genieCreateConversationMessage: vi.fn(), + genieGetConversationMessage: vi.fn(), + genieListConversationMessages: vi.fn(), + genieGetMessageAttachmentQueryResult: vi.fn(), }, } as any; } -/** - * Builds a mock waiter whose `.wait()` invokes `onProgress` for each - * progress value, then resolves with the final result. - */ -function createMockWaiter(opts: { - progressValues?: Partial[]; - result: GenieMessage; -}) { - return { - wait: vi.fn().mockImplementation(async (options: any = {}) => { - if (opts.progressValues) { - for (const value of opts.progressValues) { - if (options.onProgress) { - await options.onProgress(value); - } - } - } - return opts.result; - }), - message_id: opts.result.message_id, - conversation_id: opts.result.conversation_id, - }; +/** Stubs the start call; every poll then returns `result`. */ +function mockStart( + ws: ReturnType, + result: GenieMessage, + method = "genieStartConversation", +) { + ws.genie[method].mockResolvedValue({ + messageId: result.messageId, + conversationId: result.conversationId, + }); + ws.genie.genieGetConversationMessage.mockResolvedValue(result); } // --------------------------------------------------------------------------- @@ -105,17 +97,14 @@ describe("GenieConnector", () => { describe("streamSendMessage", () => { test("yields message_start, status updates, then message_result", async () => { + vi.useFakeTimers(); const completedMsg = makeGenieMessage(); - const waiter = createMockWaiter({ - progressValues: [ - { status: "EXECUTING_QUERY" }, - { status: "COMPLETED" }, - ], - result: completedMsg, - }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, completedMsg); + ws.genie.genieGetConversationMessage.mockResolvedValueOnce( + makeGenieMessage({ status: "EXECUTING_QUERY" }), + ); - const events = await collect( + const pending = collect( connector.streamSendMessage( ws, "space-1", @@ -123,6 +112,9 @@ describe("GenieConnector", () => { undefined, ), ); + await vi.runAllTimersAsync(); + const events = await pending; + vi.useRealTimers(); expect(events[0]).toEqual({ type: "message_start", @@ -144,51 +136,59 @@ describe("GenieConnector", () => { test("new conversation calls startConversation", async () => { const completedMsg = makeGenieMessage(); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, completedMsg); await collect( connector.streamSendMessage(ws, "space-1", "hello", undefined), ); - expect(ws.genie.startConversation).toHaveBeenCalledWith({ - space_id: "space-1", + expect(ws.genie.genieStartConversation).toHaveBeenCalledWith({ + spaceId: "space-1", content: "hello", }); - expect(ws.genie.createMessage).not.toHaveBeenCalled(); + expect(ws.genie.genieCreateConversationMessage).not.toHaveBeenCalled(); }); test("existing conversation calls createMessage", async () => { const completedMsg = makeGenieMessage(); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.createMessage.mockResolvedValue(waiter); + mockStart(ws, completedMsg, "genieCreateConversationMessage"); await collect( connector.streamSendMessage(ws, "space-1", "hello", "conv-existing"), ); - expect(ws.genie.createMessage).toHaveBeenCalledWith({ - space_id: "space-1", - conversation_id: "conv-existing", + expect(ws.genie.genieCreateConversationMessage).toHaveBeenCalledWith({ + spaceId: "space-1", + conversationId: "conv-existing", content: "hello", }); - expect(ws.genie.startConversation).not.toHaveBeenCalled(); + expect(ws.genie.genieStartConversation).not.toHaveBeenCalled(); }); test("emits query_result for attachments with statementIds", async () => { const completedMsg = makeGenieMessageWithQuery(); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, completedMsg); + // The modular SDK hands back camelCase with int64 counts as bigint; the + // SSE payload must be the snake_case wire shape and JSON-serializable. + ws.genie.genieGetMessageAttachmentQueryResult.mockResolvedValue({ + statementResponse: { + statementId: "stmt-1", + manifest: { + schema: { columns: [{ name: "total", typeName: "DOUBLE" }] }, + totalRowCount: 1n, + }, + result: { dataArray: [["1234.56", null]], rowCount: 1n }, + }, + }); const statementResponse = { + statement_id: "stmt-1", manifest: { schema: { columns: [{ name: "total", type_name: "DOUBLE" }] }, + total_row_count: 1, }, - result: { data_array: [["1234.56"]] }, + result: { data_array: [["1234.56", null]], row_count: 1 }, }; - ws.genie.getMessageAttachmentQueryResult.mockResolvedValue({ - statement_response: statementResponse, - }); const events = await collect( connector.streamSendMessage(ws, "space-1", "query", undefined), @@ -201,10 +201,25 @@ describe("GenieConnector", () => { statementId: "stmt-1", data: statementResponse, }); + expect(() => JSON.stringify(queryResult)).not.toThrow(); + }); + + test("maps a FAILED poll to the table permissions error", async () => { + mockStart(ws, makeGenieMessage({ status: "FAILED" })); + + const events = await collect( + connector.streamSendMessage(ws, "space-1", "q", undefined), + ); + + expect(events.at(-1)).toEqual({ + type: "error", + error: + "You may not have access to the data tables. Please verify your table permissions.", + }); }); test("yields error event on SDK failure", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("Network timeout"), ); @@ -216,7 +231,7 @@ describe("GenieConnector", () => { }); test("classifies RESOURCE_DOES_NOT_EXIST as access denied", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("RESOURCE_DOES_NOT_EXIST: space not found"), ); @@ -234,9 +249,8 @@ describe("GenieConnector", () => { test("emits error event when query result fetch fails", async () => { const completedMsg = makeGenieMessageWithQuery(); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.startConversation.mockResolvedValue(waiter); - ws.genie.getMessageAttachmentQueryResult.mockRejectedValue( + mockStart(ws, completedMsg); + ws.genie.genieGetMessageAttachmentQueryResult.mockRejectedValue( new Error("statement expired"), ); @@ -258,12 +272,12 @@ describe("GenieConnector", () => { describe("streamConversation", () => { test("yields message_result for each message, then history_info", async () => { - ws.genie.listConversationMessages.mockResolvedValue({ + ws.genie.genieListConversationMessages.mockResolvedValue({ messages: [ - makeGenieMessage({ message_id: "m1", content: "first" }), - makeGenieMessage({ message_id: "m2", content: "second" }), + makeGenieMessage({ messageId: "m1", content: "first" }), + makeGenieMessage({ messageId: "m2", content: "second" }), ], - next_page_token: null, + nextPageToken: null, }); const events = await collect( @@ -286,39 +300,45 @@ describe("GenieConnector", () => { }); test("fetches query results in parallel when includeQueryResults=true", async () => { - ws.genie.listConversationMessages.mockResolvedValue({ + ws.genie.genieListConversationMessages.mockResolvedValue({ messages: [ makeGenieMessageWithQuery({ - message_id: "m1", + messageId: "m1", attachments: [ { - attachment_id: "att-a", - query: { - title: "Q1", - query: "SELECT 1", - statement_id: "stmt-a", + attachmentId: "att-a", + attachment: { + $case: "query", + query: { + title: "Q1", + query: "SELECT 1", + statementId: "stmt-a", + }, }, }, { - attachment_id: "att-b", - query: { - title: "Q2", - query: "SELECT 2", - statement_id: "stmt-b", + attachmentId: "att-b", + attachment: { + $case: "query", + query: { + title: "Q2", + query: "SELECT 2", + statementId: "stmt-b", + }, }, }, ], }), ], - next_page_token: null, + nextPageToken: null, }); const stmtResponse = { manifest: { schema: { columns: [] } }, - result: { data_array: [] }, + result: { dataArray: [] }, }; - ws.genie.getMessageAttachmentQueryResult.mockResolvedValue({ - statement_response: stmtResponse, + ws.genie.genieGetMessageAttachmentQueryResult.mockResolvedValue({ + statementResponse: stmtResponse, }); const events = await collect( @@ -329,13 +349,15 @@ describe("GenieConnector", () => { const queryResults = events.filter((e) => e.type === "query_result"); expect(queryResults).toHaveLength(2); - expect(ws.genie.getMessageAttachmentQueryResult).toHaveBeenCalledTimes(2); + expect( + ws.genie.genieGetMessageAttachmentQueryResult, + ).toHaveBeenCalledTimes(2); }); test("skips query results when includeQueryResults=false", async () => { - ws.genie.listConversationMessages.mockResolvedValue({ + ws.genie.genieListConversationMessages.mockResolvedValue({ messages: [makeGenieMessageWithQuery()], - next_page_token: null, + nextPageToken: null, }); const events = await collect( @@ -345,44 +367,52 @@ describe("GenieConnector", () => { ); expect(events.filter((e) => e.type === "query_result")).toHaveLength(0); - expect(ws.genie.getMessageAttachmentQueryResult).not.toHaveBeenCalled(); + expect( + ws.genie.genieGetMessageAttachmentQueryResult, + ).not.toHaveBeenCalled(); }); test("handles partial query result failures via Promise.allSettled", async () => { - ws.genie.listConversationMessages.mockResolvedValue({ + ws.genie.genieListConversationMessages.mockResolvedValue({ messages: [ makeGenieMessage({ - message_id: "m1", + messageId: "m1", attachments: [ { - attachment_id: "att-ok", - query: { - title: "OK", - query: "SELECT 1", - statement_id: "stmt-ok", + attachmentId: "att-ok", + attachment: { + $case: "query", + query: { + title: "OK", + query: "SELECT 1", + statementId: "stmt-ok", + }, }, }, { - attachment_id: "att-fail", - query: { - title: "Fail", - query: "SELECT 2", - statement_id: "stmt-fail", + attachmentId: "att-fail", + attachment: { + $case: "query", + query: { + title: "Fail", + query: "SELECT 2", + statementId: "stmt-fail", + }, }, }, ], }), ], - next_page_token: null, + nextPageToken: null, }); const stmtResponse = { manifest: { schema: { columns: [] } }, - result: { data_array: [] }, + result: { dataArray: [] }, }; - ws.genie.getMessageAttachmentQueryResult - .mockResolvedValueOnce({ statement_response: stmtResponse }) + ws.genie.genieGetMessageAttachmentQueryResult + .mockResolvedValueOnce({ statementResponse: stmtResponse }) .mockRejectedValueOnce(new Error("statement expired")); const events = await collect( @@ -400,7 +430,7 @@ describe("GenieConnector", () => { }); test("yields error when listConversationMessages fails", async () => { - ws.genie.listConversationMessages.mockRejectedValue( + ws.genie.genieListConversationMessages.mockRejectedValue( new Error("RESOURCE_DOES_NOT_EXIST: conv not found"), ); @@ -423,7 +453,7 @@ describe("GenieConnector", () => { describe("streamGetMessage", () => { test("polls until COMPLETED, yields status + message_result", async () => { - ws.genie.getMessage + ws.genie.genieGetConversationMessage .mockResolvedValueOnce(makeGenieMessage({ status: "EXECUTING_QUERY" })) .mockResolvedValueOnce(makeGenieMessage({ status: "COMPLETED" })); @@ -439,11 +469,11 @@ describe("GenieConnector", () => { }); expect(events[1]).toEqual({ type: "status", status: "COMPLETED" }); expect(events[2]).toMatchObject({ type: "message_result" }); - expect(ws.genie.getMessage).toHaveBeenCalledTimes(2); + expect(ws.genie.genieGetConversationMessage).toHaveBeenCalledTimes(2); }); test("polls until FAILED, yields status + message_result", async () => { - ws.genie.getMessage + ws.genie.genieGetConversationMessage .mockResolvedValueOnce(makeGenieMessage({ status: "EXECUTING_QUERY" })) .mockResolvedValueOnce( makeGenieMessage({ @@ -472,7 +502,7 @@ describe("GenieConnector", () => { test("respects abort signal", async () => { const controller = new AbortController(); - ws.genie.getMessage.mockResolvedValue( + ws.genie.genieGetConversationMessage.mockResolvedValue( makeGenieMessage({ status: "EXECUTING_QUERY" }), ); @@ -500,7 +530,9 @@ describe("GenieConnector", () => { }); test("yields error when getMessage throws", async () => { - ws.genie.getMessage.mockRejectedValue(new Error("service unavailable")); + ws.genie.genieGetConversationMessage.mockRejectedValue( + new Error("service unavailable"), + ); const events = await collect( connector.streamGetMessage(ws, "space-1", "conv-1", "msg-1", { @@ -512,7 +544,7 @@ describe("GenieConnector", () => { }); test("does not duplicate status events for same status", async () => { - ws.genie.getMessage + ws.genie.genieGetConversationMessage .mockResolvedValueOnce(makeGenieMessage({ status: "EXECUTING_QUERY" })) .mockResolvedValueOnce(makeGenieMessage({ status: "EXECUTING_QUERY" })) .mockResolvedValueOnce(makeGenieMessage({ status: "COMPLETED" })); @@ -538,11 +570,10 @@ describe("GenieConnector", () => { describe("sendMessage", () => { test("returns completed message response", async () => { const completedMsg = makeGenieMessage({ - message_id: "msg-42", - conversation_id: "conv-new", + messageId: "msg-42", + conversationId: "conv-new", }); - const waiter = createMockWaiter({ result: completedMsg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, completedMsg); const result = await connector.sendMessage( ws, @@ -566,17 +597,17 @@ describe("GenieConnector", () => { // listConversationMessages reverses the SDK response, so mock data // is ordered newest-first (as the SDK returns) and results are // oldest-first after reversal. - ws.genie.listConversationMessages + ws.genie.genieListConversationMessages .mockResolvedValueOnce({ messages: [ - makeGenieMessage({ message_id: "m2" }), - makeGenieMessage({ message_id: "m1" }), + makeGenieMessage({ messageId: "m2" }), + makeGenieMessage({ messageId: "m1" }), ], - next_page_token: "page2", + nextPageToken: "page2", }) .mockResolvedValueOnce({ - messages: [makeGenieMessage({ message_id: "m3" })], - next_page_token: null, + messages: [makeGenieMessage({ messageId: "m3" })], + nextPageToken: null, }); const result = await connector.getConversation(ws, "space-1", "conv-1"); @@ -587,7 +618,7 @@ describe("GenieConnector", () => { "m2", "m3", ]); - expect(ws.genie.listConversationMessages).toHaveBeenCalledTimes(2); + expect(ws.genie.genieListConversationMessages).toHaveBeenCalledTimes(2); }); test("respects maxMessages limit", async () => { @@ -596,13 +627,13 @@ describe("GenieConnector", () => { maxMessages: 2, }); - ws.genie.listConversationMessages.mockResolvedValueOnce({ + ws.genie.genieListConversationMessages.mockResolvedValueOnce({ messages: [ - makeGenieMessage({ message_id: "m1" }), - makeGenieMessage({ message_id: "m2" }), - makeGenieMessage({ message_id: "m3" }), + makeGenieMessage({ messageId: "m1" }), + makeGenieMessage({ messageId: "m2" }), + makeGenieMessage({ messageId: "m3" }), ], - next_page_token: "page2", + nextPageToken: "page2", }); const result = await smallConnector.getConversation( @@ -614,7 +645,7 @@ describe("GenieConnector", () => { // Should be sliced to maxMessages expect(result.messages).toHaveLength(2); // Should NOT fetch a second page since length already >= maxMessages - expect(ws.genie.listConversationMessages).toHaveBeenCalledTimes(1); + expect(ws.genie.genieListConversationMessages).toHaveBeenCalledTimes(1); }); }); @@ -625,14 +656,13 @@ describe("GenieConnector", () => { describe("mapAttachments", () => { test("handles query attachments", async () => { const msg = makeGenieMessageWithQuery(); - const waiter = createMockWaiter({ result: msg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, msg); // We drive through streamSendMessage to exercise mapAttachments - ws.genie.getMessageAttachmentQueryResult.mockResolvedValue({ - statement_response: { + ws.genie.genieGetMessageAttachmentQueryResult.mockResolvedValue({ + statementResponse: { manifest: { schema: { columns: [] } }, - result: { data_array: [] }, + result: { dataArray: [] }, }, }); @@ -658,13 +688,15 @@ describe("GenieConnector", () => { const msg = makeGenieMessage({ attachments: [ { - attachment_id: "att-text", - text: { content: "Here is the explanation" }, + attachmentId: "att-text", + attachment: { + $case: "text", + text: { content: "Here is the explanation" }, + }, }, ], }); - const waiter = createMockWaiter({ result: msg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, msg); const events = await collect( connector.streamSendMessage(ws, "space-1", "q", undefined), @@ -683,15 +715,17 @@ describe("GenieConnector", () => { const msg = makeGenieMessage({ attachments: [ { - attachment_id: "att-sq", - suggested_questions: { - questions: ["What is X?", "Show me Y"], + attachmentId: "att-sq", + attachment: { + $case: "suggestedQuestions", + suggestedQuestions: { + questions: ["What is X?", "Show me Y"], + }, }, }, ], }); - const waiter = createMockWaiter({ result: msg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, msg); const events = await collect( connector.streamSendMessage(ws, "space-1", "q", undefined), @@ -708,8 +742,7 @@ describe("GenieConnector", () => { test("returns empty array when message has no attachments", async () => { const msg = makeGenieMessage({ attachments: undefined }); - const waiter = createMockWaiter({ result: msg }); - ws.genie.startConversation.mockResolvedValue(waiter); + mockStart(ws, msg); const events = await collect( connector.streamSendMessage(ws, "space-1", "q", undefined), @@ -726,7 +759,7 @@ describe("GenieConnector", () => { describe("classifyGenieError", () => { test("maps RESOURCE_DOES_NOT_EXIST to space access denied", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("RESOURCE_DOES_NOT_EXIST: space xyz"), ); @@ -740,8 +773,25 @@ describe("GenieConnector", () => { }); }); + test("maps an ApiError .code of RESOURCE_DOES_NOT_EXIST to space access denied", async () => { + ws.genie.genieStartConversation.mockRejectedValue( + Object.assign(new Error("Space xyz not found"), { + code: "RESOURCE_DOES_NOT_EXIST", + }), + ); + + const events = await collect( + connector.streamSendMessage(ws, "space-1", "hi", undefined), + ); + + expect(events[0]).toEqual({ + type: "error", + error: "You don't have access to this Genie Space.", + }); + }); + test("maps failed-to-reach-COMPLETED + FAILED to table permissions", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("failed to reach COMPLETED state, got FAILED"), ); @@ -757,7 +807,7 @@ describe("GenieConnector", () => { }); test("passes through unknown error messages", async () => { - ws.genie.startConversation.mockRejectedValue( + ws.genie.genieStartConversation.mockRejectedValue( new Error("something unexpected"), ); @@ -772,7 +822,7 @@ describe("GenieConnector", () => { }); test("handles non-Error throwable", async () => { - ws.genie.startConversation.mockRejectedValue("string error"); + ws.genie.genieStartConversation.mockRejectedValue("string error"); const events = await collect( connector.streamSendMessage(ws, "space-1", "hi", undefined), diff --git a/packages/appkit/src/plugins/genie/tests/genie.test.ts b/packages/appkit/src/plugins/genie/tests/genie.test.ts index 5cd03a348..09b26d4c1 100644 --- a/packages/appkit/src/plugins/genie/tests/genie.test.ts +++ b/packages/appkit/src/plugins/genie/tests/genie.test.ts @@ -19,68 +19,73 @@ import type { IGenieConfig } from "../types"; useTestCache(); function createMockGenieService() { - const getMessageAttachmentQueryResult = vi.fn(); - - const createWaiter = ( + const genieGetMessageAttachmentQueryResult = vi.fn(); + // Messages the start/create stubs "created"; the poll serves them back. + const messages = new Map(); + const register = ( conversationId: string, messageId: string, attachments: any[] = [], status = "COMPLETED", - ) => ({ - wait: vi.fn().mockImplementation(async ({ onProgress }: any) => { - if (onProgress) { - await onProgress({ status: "ASKING_AI" }); - await onProgress({ status: "EXECUTING_QUERY" }); - } - return { - message_id: messageId, - conversation_id: conversationId, - space_id: "test-space-id", - content: "Here are your results", - status, - attachments, - error: undefined, - }; - }), - }); + ) => { + messages.set(messageId, { + messageId, + conversationId, + spaceId: "test-space-id", + content: "Here are your results", + status, + attachments, + error: undefined, + }); + return { conversationId, messageId }; + }; + + const genieGetConversationMessage = vi + .fn() + .mockImplementation(async ({ messageId }: any) => messages.get(messageId)); - const startConversation = vi.fn().mockImplementation(async () => ({ - conversation_id: "new-conv-id", - message_id: "new-msg-id", - ...createWaiter("new-conv-id", "new-msg-id", [ + const genieStartConversation = vi.fn().mockImplementation(async () => + register("new-conv-id", "new-msg-id", [ { - attachment_id: "att-1", - query: { - title: "Top Customers", - description: "Query for top customers", - query: "SELECT * FROM customers", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Top Customers", + description: "Query for top customers", + query: "SELECT * FROM customers", + statementId: "stmt-1", + }, }, }, ]), - })); + ); - const createMessage = vi.fn().mockImplementation(async () => - createWaiter("existing-conv-id", "followup-msg-id", [ + const genieCreateConversationMessage = vi.fn().mockImplementation(async () => + register("existing-conv-id", "followup-msg-id", [ { - attachment_id: "att-2", - query: { - title: "Follow-up Query", - query: "SELECT * FROM orders", - statement_id: "stmt-2", + attachmentId: "att-2", + attachment: { + $case: "query", + query: { + title: "Follow-up Query", + query: "SELECT * FROM orders", + statementId: "stmt-2", + }, }, }, ]), ); - const listConversationMessages = vi.fn(); + const genieListConversationMessages = vi.fn(); return { - startConversation, - createMessage, - getMessageAttachmentQueryResult, - listConversationMessages, - createWaiter, + genieStartConversation, + genieCreateConversationMessage, + genieGetConversationMessage, + genieGetMessageAttachmentQueryResult, + genieListConversationMessages, + register, }; } @@ -102,11 +107,11 @@ describe("Genie Plugin", () => { mockGenieService = createMockGenieService(); - mockGenieService.getMessageAttachmentQueryResult.mockResolvedValue({ - statement_response: { + mockGenieService.genieGetMessageAttachmentQueryResult.mockResolvedValue({ + statementResponse: { status: { state: "SUCCEEDED" }, result: { - data_array: [ + dataArray: [ ["Acme Corp", "1000000"], ["Globex", "500000"], ], @@ -114,8 +119,8 @@ describe("Genie Plugin", () => { manifest: { schema: { columns: [ - { name: "customer", type_name: "STRING" }, - { name: "revenue", type_name: "DECIMAL" }, + { name: "customer", typeName: "STRING" }, + { name: "revenue", typeName: "DECIMAL" }, ], }, }, @@ -238,9 +243,9 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect(mockRes.status).not.toHaveBeenCalledWith(404); - expect(mockGenieService.startConversation).toHaveBeenCalledWith( + expect(mockGenieService.genieStartConversation).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", + spaceId: "test-space-id", content: "What are my top customers?", }), ); @@ -291,12 +296,16 @@ describe("Genie Plugin", () => { }, }); const mockRes = createMockResponse(); + // One non-terminal poll before COMPLETED (costs one ~1s backoff sleep). + mockGenieService.genieGetConversationMessage.mockResolvedValueOnce({ + status: "ASKING_AI", + }); await handler(mockReq, mockRes); - expect(mockGenieService.startConversation).toHaveBeenCalledWith( + expect(mockGenieService.genieStartConversation).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", + spaceId: "test-space-id", content: "What are my top customers?", }), ); @@ -357,15 +366,17 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); - expect(mockGenieService.createMessage).toHaveBeenCalledWith( + expect( + mockGenieService.genieCreateConversationMessage, + ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "existing-conv-id", + spaceId: "test-space-id", + conversationId: "existing-conv-id", content: "Show me more details", }), ); - expect(mockGenieService.startConversation).not.toHaveBeenCalled(); + expect(mockGenieService.genieStartConversation).not.toHaveBeenCalled(); const writeCalls = mockRes.write.mock.calls.map((call: any[]) => call[0]); const allWritten = writeCalls.join(""); @@ -379,51 +390,46 @@ describe("Genie Plugin", () => { describe("multiple attachments", () => { test("should yield query_result for each query attachment", async () => { // Override startConversation to return multiple query attachments - mockGenieService.startConversation.mockImplementation(async () => ({ - conversation_id: "multi-conv-id", - message_id: "multi-msg-id", - wait: vi.fn().mockImplementation(async ({ onProgress }: any) => { - if (onProgress) { - await onProgress({ status: "ASKING_AI" }); - } - return { - message_id: "multi-msg-id", - conversation_id: "multi-conv-id", - space_id: "test-space-id", - content: "Here are two queries", - status: "COMPLETED", - attachments: [ - { - attachment_id: "att-q1", - query: { - title: "Query 1", - query: "SELECT 1", - statement_id: "stmt-q1", - }, - }, - { - attachment_id: "att-q2", - query: { - title: "Query 2", - query: "SELECT 2", - statement_id: "stmt-q2", - }, + mockGenieService.genieStartConversation.mockImplementation(async () => + mockGenieService.register("multi-conv-id", "multi-msg-id", [ + { + attachmentId: "att-q1", + attachment: { + $case: "query", + query: { + title: "Query 1", + query: "SELECT 1", + statementId: "stmt-q1", }, - { - attachment_id: "att-text", - text: { content: "Some explanation" }, + }, + }, + { + attachmentId: "att-q2", + attachment: { + $case: "query", + query: { + title: "Query 2", + query: "SELECT 2", + statementId: "stmt-q2", }, - ], - }; - }), - })); + }, + }, + { + attachmentId: "att-text", + attachment: { + $case: "text", + text: { content: "Some explanation" }, + }, + }, + ]), + ); - mockGenieService.getMessageAttachmentQueryResult + mockGenieService.genieGetMessageAttachmentQueryResult .mockResolvedValueOnce({ - statement_response: { result: { data: [["row1"]] } }, + statementResponse: { result: { data: [["row1"]] } }, }) .mockResolvedValueOnce({ - statement_response: { result: { data: [["row2"]] } }, + statementResponse: { result: { data: [["row2"]] } }, }); const plugin = new GeniePlugin(config); @@ -446,18 +452,18 @@ describe("Genie Plugin", () => { // getMessageAttachmentQueryResult should be called twice (once per query attachment) expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledTimes(2); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( - expect.objectContaining({ attachment_id: "att-q1" }), + expect.objectContaining({ attachmentId: "att-q1" }), ); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( - expect.objectContaining({ attachment_id: "att-q2" }), + expect.objectContaining({ attachmentId: "att-q2" }), ); const writeCalls = mockRes.write.mock.calls.map((call: any[]) => call[0]); @@ -473,7 +479,7 @@ describe("Genie Plugin", () => { describe("error handling", () => { test("should yield error event on SDK failure", async () => { - mockGenieService.startConversation.mockRejectedValue( + mockGenieService.genieStartConversation.mockRejectedValue( new Error("Genie service unavailable"), ); @@ -519,9 +525,9 @@ describe("Genie Plugin", () => { } function mockMessages(messages: any[]) { - mockGenieService.listConversationMessages.mockResolvedValue({ + mockGenieService.genieListConversationMessages.mockResolvedValue({ messages, - next_page_token: undefined, + nextPageToken: undefined, }); } @@ -551,26 +557,29 @@ describe("Genie Plugin", () => { test("should stream message_result events for each message", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "What are the top customers?", status: "COMPLETED", attachments: [], }, { - message_id: "msg-2", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-2", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Here are the results", status: "COMPLETED", attachments: [ { - attachment_id: "att-1", - query: { - title: "Top Customers", - query: "SELECT * FROM customers", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Top Customers", + query: "SELECT * FROM customers", + statementId: "stmt-1", + }, }, }, ], @@ -591,11 +600,13 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); - expect(mockGenieService.listConversationMessages).toHaveBeenCalledWith( + expect( + mockGenieService.genieListConversationMessages, + ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "conv-123", - page_size: genieConnectorDefaults.initialPageSize, + spaceId: "test-space-id", + conversationId: "conv-123", + pageSize: genieConnectorDefaults.initialPageSize, }), ); @@ -621,18 +632,21 @@ describe("Genie Plugin", () => { test("should stream query_result events when includeQueryResults is true (default)", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Results", status: "COMPLETED", attachments: [ { - attachment_id: "att-1", - query: { - title: "Query 1", - query: "SELECT 1", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Query 1", + query: "SELECT 1", + statementId: "stmt-1", + }, }, }, ], @@ -654,13 +668,13 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "conv-123", - message_id: "msg-1", - attachment_id: "att-1", + spaceId: "test-space-id", + conversationId: "conv-123", + messageId: "msg-1", + attachmentId: "att-1", }), ); @@ -675,18 +689,21 @@ describe("Genie Plugin", () => { test("should NOT stream query_result events when includeQueryResults is false", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Results", status: "COMPLETED", attachments: [ { - attachment_id: "att-1", - query: { - title: "Query 1", - query: "SELECT 1", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Query 1", + query: "SELECT 1", + statementId: "stmt-1", + }, }, }, ], @@ -710,7 +727,7 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).not.toHaveBeenCalled(); const writeCalls = mockRes.write.mock.calls.map((call: any[]) => call[0]); @@ -722,18 +739,18 @@ describe("Genie Plugin", () => { }); test("should fetch only one page and emit history_info with nextPageToken", async () => { - mockGenieService.listConversationMessages.mockResolvedValueOnce({ + mockGenieService.genieListConversationMessages.mockResolvedValueOnce({ messages: [ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Most recent message", status: "COMPLETED", attachments: [], }, ], - next_page_token: "page-2-token", + nextPageToken: "page-2-token", }); const plugin = new GeniePlugin(config); @@ -753,15 +770,17 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); // Should only fetch one page (lazy loading) - expect(mockGenieService.listConversationMessages).toHaveBeenCalledTimes( - 1, - ); + expect( + mockGenieService.genieListConversationMessages, + ).toHaveBeenCalledTimes(1); - expect(mockGenieService.listConversationMessages).toHaveBeenCalledWith( + expect( + mockGenieService.genieListConversationMessages, + ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "conv-123", - page_size: genieConnectorDefaults.initialPageSize, + spaceId: "test-space-id", + conversationId: "conv-123", + pageSize: genieConnectorDefaults.initialPageSize, }), ); @@ -778,9 +797,9 @@ describe("Genie Plugin", () => { test("should emit history_info with null nextPageToken when no more pages", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Only message", status: "COMPLETED", attachments: [], @@ -838,7 +857,7 @@ describe("Genie Plugin", () => { }); test("should yield error event on SDK failure", async () => { - mockGenieService.listConversationMessages.mockRejectedValue( + mockGenieService.genieListConversationMessages.mockRejectedValue( new Error("Conversation not found"), ); @@ -867,35 +886,41 @@ describe("Genie Plugin", () => { test("should fetch query results in parallel for multiple attachments across messages", async () => { mockMessages([ { - message_id: "msg-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "First query", status: "COMPLETED", attachments: [ { - attachment_id: "att-1", - query: { - title: "Query 1", - query: "SELECT 1", - statement_id: "stmt-1", + attachmentId: "att-1", + attachment: { + $case: "query", + query: { + title: "Query 1", + query: "SELECT 1", + statementId: "stmt-1", + }, }, }, ], }, { - message_id: "msg-2", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-2", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Second query", status: "COMPLETED", attachments: [ { - attachment_id: "att-2", - query: { - title: "Query 2", - query: "SELECT 2", - statement_id: "stmt-2", + attachmentId: "att-2", + attachment: { + $case: "query", + query: { + title: "Query 2", + query: "SELECT 2", + statementId: "stmt-2", + }, }, }, ], @@ -917,23 +942,23 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledTimes(2); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( expect.objectContaining({ - message_id: "msg-1", - attachment_id: "att-1", + messageId: "msg-1", + attachmentId: "att-1", }), ); expect( - mockGenieService.getMessageAttachmentQueryResult, + mockGenieService.genieGetMessageAttachmentQueryResult, ).toHaveBeenCalledWith( expect.objectContaining({ - message_id: "msg-2", - attachment_id: "att-2", + messageId: "msg-2", + attachmentId: "att-2", }), ); @@ -950,7 +975,7 @@ describe("Genie Plugin", () => { describe("error classification", () => { test("should return user-friendly message for RESOURCE_DOES_NOT_EXIST error", async () => { - mockGenieService.startConversation.mockRejectedValue( + mockGenieService.genieStartConversation.mockRejectedValue( new Error( "RESOURCE_DOES_NOT_EXIST: No Genie space found with id test-space-id", ), @@ -984,7 +1009,7 @@ describe("Genie Plugin", () => { }); test("should return user-friendly message for FAILED state error (table access denied)", async () => { - mockGenieService.startConversation.mockRejectedValue( + mockGenieService.genieStartConversation.mockRejectedValue( new Error( "failed to reach COMPLETED state, got FAILED: [object Object]", ), @@ -1018,7 +1043,7 @@ describe("Genie Plugin", () => { }); test("should return user-friendly message for RESOURCE_DOES_NOT_EXIST on getConversation", async () => { - mockGenieService.listConversationMessages.mockRejectedValue( + mockGenieService.genieListConversationMessages.mockRejectedValue( new Error( "RESOURCE_DOES_NOT_EXIST: No Genie space found with id test-space-id", ), @@ -1078,9 +1103,9 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); expect(mockRes.status).not.toHaveBeenCalledWith(404); - expect(mockGenieService.startConversation).toHaveBeenCalledWith( + expect(mockGenieService.genieStartConversation).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "env-space-id", + spaceId: "env-space-id", content: "test question", }), ); @@ -1162,18 +1187,18 @@ describe("Genie Plugin", () => { describe("getConversation with pageToken", () => { test("should pass pageToken through to streamConversation", async () => { - mockGenieService.listConversationMessages.mockResolvedValueOnce({ + mockGenieService.genieListConversationMessages.mockResolvedValueOnce({ messages: [ { - message_id: "msg-old-1", - conversation_id: "conv-123", - space_id: "test-space-id", + messageId: "msg-old-1", + conversationId: "conv-123", + spaceId: "test-space-id", content: "Older message", status: "COMPLETED", attachments: [], }, ], - next_page_token: "next-token-abc", + nextPageToken: "next-token-abc", }); const plugin = new GeniePlugin(config); @@ -1197,11 +1222,13 @@ describe("Genie Plugin", () => { await handler(mockReq, mockRes); - expect(mockGenieService.listConversationMessages).toHaveBeenCalledWith( + expect( + mockGenieService.genieListConversationMessages, + ).toHaveBeenCalledWith( expect.objectContaining({ - space_id: "test-space-id", - conversation_id: "conv-123", - page_token: "some-page-token", + spaceId: "test-space-id", + conversationId: "conv-123", + pageToken: "some-page-token", }), ); @@ -1215,7 +1242,7 @@ describe("Genie Plugin", () => { }); test("should yield error event when paginated request fails", async () => { - mockGenieService.listConversationMessages.mockRejectedValue( + mockGenieService.genieListConversationMessages.mockRejectedValue( new Error("Page token expired"), ); diff --git a/packages/appkit/src/testing/tests/mock-workspace-client.test.ts b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts index 5ec90a0d9..1dada4f5c 100644 --- a/packages/appkit/src/testing/tests/mock-workspace-client.test.ts +++ b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts @@ -21,7 +21,7 @@ describe("createMockWorkspaceClient", () => { // Never-crash is the headline claim, so all nine are asserted, not sampled. test.each([ ["files", "listDirectory", undefined], - ["genie", "getMessage", undefined], + ["genie", "genieGetConversationMessage", undefined], ["jobs", "getRun", undefined], ["servingEndpoints", "get", undefined], ["warehouses", "getWarehouse", { state: "RUNNING" }], diff --git a/packages/appkit/src/workspace-client/index.ts b/packages/appkit/src/workspace-client/index.ts index a7d7e1c34..39e400c6a 100644 --- a/packages/appkit/src/workspace-client/index.ts +++ b/packages/appkit/src/workspace-client/index.ts @@ -5,14 +5,7 @@ * behind a stable facade. */ -export { - ApiError, - ConfigError, - Context, - createWorkspaceClient, - Time, - TimeUnits, -} from "shared"; +export { ApiError, ConfigError, Context, createWorkspaceClient } from "shared"; // Forwards every wrapper type — legacy service namespaces (files/jobs/serving), // the client option/waiter types, and the modular SDK client + model types // (warehouses, statementExecution). `sql` is gone: its statement + warehouse diff --git a/packages/shared/package.json b/packages/shared/package.json index 867a9f2b4..6002e4693 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -51,6 +51,7 @@ "@databricks/sdk-auth": "0.51.0", "@databricks/sdk-core": "0.51.0", "@databricks/sdk-experimental": "0.17.0", + "@databricks/sdk-genie": "0.54.0", "@databricks/sdk-options": "0.51.0", "@databricks/sdk-scim": "0.51.0", "@databricks/sdk-statementexecution": "0.52.0", diff --git a/packages/shared/src/workspace-client/client.ts b/packages/shared/src/workspace-client/client.ts index 29494bb3c..1552392e5 100644 --- a/packages/shared/src/workspace-client/client.ts +++ b/packages/shared/src/workspace-client/client.ts @@ -18,6 +18,8 @@ import { buildWarehousesClient, buildWorkspaceAuth, type ScimClient, + buildGenieClient, + type GenieClient, type StatementExecutionClient, type WarehousesClient, type WorkspaceAuth, @@ -32,6 +34,7 @@ export class AppKitWorkspaceClient implements WorkspaceClient { #statementExecution?: StatementExecutionClient; #auth?: WorkspaceAuth; #currentUser?: ScimClient; + #genie?: GenieClient; constructor(opts: WorkspaceClientOptions) { this.#opts = opts; @@ -49,8 +52,12 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#warehouses; } - get genie() { - return this.#getLegacy().genie; + // Migrated to the modular SDK — built lazily, independent of the legacy client. + get genie(): GenieClient { + if (!this.#genie) { + this.#genie = buildGenieClient(this.#opts); + } + return this.#genie; } get jobs() { diff --git a/packages/shared/src/workspace-client/index.ts b/packages/shared/src/workspace-client/index.ts index 2b981ebec..619c24d43 100644 --- a/packages/shared/src/workspace-client/index.ts +++ b/packages/shared/src/workspace-client/index.ts @@ -9,9 +9,7 @@ export { createWorkspaceClient } from "./factory"; export type { CancellationToken, ClientOptions, - GenieMessage, LegacyWorkspaceClient, - Waiter, WorkspaceClientOptions, } from "./legacy"; // SDK value + type re-exports so AppKit modules import them from the wrapper. diff --git a/packages/shared/src/workspace-client/legacy.ts b/packages/shared/src/workspace-client/legacy.ts index 8784fbe57..745cf0976 100644 --- a/packages/shared/src/workspace-client/legacy.ts +++ b/packages/shared/src/workspace-client/legacy.ts @@ -77,7 +77,7 @@ export type { // // AppKit modules import these from the wrapper instead of the SDK so the // boundary rule holds. `Context` bridges AbortSignal → CancellationToken -// (serving/jobs/sql-warehouse); `Time`/`TimeUnits` drive genie polling; +// (serving/jobs/sql-warehouse); `Time`/`TimeUnits` are re-exported by `shared`'s index; // `ConfigError` is matched in service-context's auth-failure handling. // // These are sourced off the namespace import rather than `export { ... } from` @@ -90,9 +90,3 @@ export type { export const { ConfigError, Context, TimeUnits, loadConfigFile } = SDK; export const Time = SDK.Time ?? (SDK as unknown as { default: typeof SDK }).default.Time; - -// Deep-import types used by the genie connector's waiter idiom. Not exposed -// on the SDK's top-level index, so re-exported here to keep the genie -// connector off a direct `@databricks/sdk-experimental/dist/**` import. -export type { GenieMessage } from "@databricks/sdk-experimental/dist/apis/dashboards"; -export type { Waiter } from "@databricks/sdk-experimental/dist/wait"; diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index 377836417..c4baade6b 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -16,6 +16,12 @@ * (`patches/@databricks__sdk-statementexecution@0.46.0.patch`) that restores the * undocumented Reyden `attachment` response field, which the SDK's generated * unmarshal transform would otherwise strip. + * + * NOTE: genie relies on a pinned pnpm patch + * (`patches/@databricks__sdk-genie@0.54.0.patch`): the generated model types the + * query result's `data_array` as protobuf `ListValue[]` (`{ values: [...] }`), but + * the API returns plain `JSON_ARRAY` rows (`[["a", null], ...]`), so the unmarshal + * schema rejected every real query result. The patch restores `(string | null)[][]`. */ import { type Credentials, @@ -35,6 +41,7 @@ import { newFetchHttpClient, } from "@databricks/sdk-core/http"; import { resolve } from "@databricks/sdk-core/profiles"; +import { GenieClient } from "@databricks/sdk-genie/v1"; import type { ClientOptions } from "@databricks/sdk-options/client"; import { ScimClient } from "@databricks/sdk-scim/v1"; import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; @@ -321,7 +328,13 @@ export function buildScimClient(opts: WorkspaceClientOptions): ScimClient { return new ScimClient(mapToClientOptions(opts)); } +/** Build a modular Genie client from wrapper options. */ +export function buildGenieClient(opts: WorkspaceClientOptions): GenieClient { + return new GenieClient(mapToClientOptions(opts)); +} + // ── Client type re-exports (for the facade accessor types) ─────────────── +export type { GenieClient } from "@databricks/sdk-genie/v1"; export type { ScimClient } from "@databricks/sdk-scim/v1"; export type { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; export type { WarehousesClient } from "@databricks/sdk-warehouses/v1"; @@ -346,6 +359,10 @@ export type { StatementStatus, StatementStatus_State, } from "@databricks/sdk-statementexecution/v1"; +export type { + GenieGetMessageQueryResultResponse, + GenieMessage, +} from "@databricks/sdk-genie/v1"; export type { EndpointHealth, EndpointInfo, diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index 868059ec1..8577c73ba 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -33,6 +33,12 @@ vi.mock("@databricks/sdk-warehouses/v1", () => ({ vi.mock("@databricks/sdk-statementexecution/v1", () => ({ StatementExecutionClient: vi.fn().mockImplementation((opts) => ({ opts })), })); +vi.mock("@databricks/sdk-genie/v1", () => ({ + GenieClient: vi.fn().mockImplementation((opts) => { + ctorOpts.push(opts); + return { opts }; + }), +})); vi.mock("@databricks/sdk-auth/credentials", () => ({ newPatCredentials: vi.fn((token: string) => { patTokens.push(token); @@ -69,7 +75,11 @@ vi.mock("@databricks/sdk-core/http", () => ({ })); import { ApiError } from "../errors"; -import { buildWarehousesClient, buildWorkspaceAuth } from "../modular"; +import { + buildGenieClient, + buildWarehousesClient, + buildWorkspaceAuth, +} from "../modular"; /** Drive the wrapped httpClient with one request and return the UA it set. */ async function sentUserAgent( @@ -250,6 +260,27 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { expect(ua).toBe("@databricks/appkit/0.64.0"); }); + test("genie (asUser) uses the OBO token as PAT and keeps the AppKit User-Agent", async () => { + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + buildGenieClient({ + token: "user-token", + host: "https://x", + clientOptions: { + product: "@databricks/appkit", + productVersion: "0.64.0", + }, + } as never); + expect(ctorOpts[0].credentials).toMatchObject({ + kind: "pat", + token: "user-token", + }); + expect(m2mOpts).toEqual([]); + expect(await sentUserAgent(ctorOpts[0].httpClient)).toBe( + "@databricks/appkit/0.64.0", + ); + }); + test("no product configured (build-time) → no httpClient override (SDK default UA)", () => { buildWarehousesClient({ host: "https://x" }); expect(ctorOpts[0].httpClient).toBeUndefined(); diff --git a/packages/shared/src/workspace-client/types.ts b/packages/shared/src/workspace-client/types.ts index bce944dc1..75b8b80c0 100644 --- a/packages/shared/src/workspace-client/types.ts +++ b/packages/shared/src/workspace-client/types.ts @@ -19,6 +19,7 @@ import type { WarehousesClient, WorkspaceAuth, ScimClient, + GenieClient, } from "./modular"; // Legacy SDK type namespaces for un-migrated services, re-exported so AppKit @@ -30,7 +31,7 @@ import type { // that request has no `skip_cannot_use` filter, so it could pick a warehouse // the caller can't use. Statement + warehouse service types come from `./modular`. export type { files, jobs, serving, sql } from "@databricks/sdk-experimental"; -// Modular SDK client + model types (warehouses, statementExecution). +// Modular SDK client + model types (warehouses, statementExecution, genie). export type * from "./modular"; /** @@ -48,8 +49,8 @@ export interface WorkspaceClient extends WorkspaceAuth { /** SQL Warehouses (modular SDK). */ readonly warehouses: WarehousesClient; - /** Genie / dashboards. */ - readonly genie: LegacyWorkspaceClient["genie"]; + /** Genie (modular SDK). */ + readonly genie: GenieClient; /** Jobs. */ readonly jobs: LegacyWorkspaceClient["jobs"]; diff --git a/patches/@databricks__sdk-genie@0.54.0.patch b/patches/@databricks__sdk-genie@0.54.0.patch new file mode 100644 index 000000000..09e3a03ba --- /dev/null +++ b/patches/@databricks__sdk-genie@0.54.0.patch @@ -0,0 +1,26 @@ +diff --git a/dist/v1/model.d.ts b/dist/v1/model.d.ts +index ceb8ef56c1d7e4912851dbdce891623f4e8fc93c..4ff6fa02396a17485630132cf1a5bae4a3ada512 100644 +--- a/dist/v1/model.d.ts ++++ b/dist/v1/model.d.ts +@@ -1581,7 +1581,7 @@ interface ResultData { + * The `JSON_ARRAY` format is an array of arrays of values, where each non-null value is + * formatted as a string. Null values are encoded as JSON `null`. + */ +- dataArray?: ListValue[] | undefined; ++ dataArray?: (string | null)[][] | undefined; + /** The position within the sequence of result set chunks. */ + chunkIndex?: number | undefined; + /** The starting row offset within the result set. */ +diff --git a/dist/v1/model.js b/dist/v1/model.js +index 8770051f3102f70625db5758166d7fc1636afcb7..d20c134274fde0ffbf7d67a8346765fa7cba7f76 100644 +--- a/dist/v1/model.js ++++ b/dist/v1/model.js +@@ -846,7 +846,7 @@ const unmarshalResultSchema = z.object({ + })); + const unmarshalResultDataSchema = z.object({ + external_links: z.array(z.lazy(() => unmarshalExternalLinkSchema)).optional(), +- data_array: z.array(z.lazy(() => unmarshalListValueSchema)).optional(), ++ data_array: z.array(z.array(z.string().nullable())).optional(), + chunk_index: z.number().optional(), + row_offset: z.union([ + z.number(), diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index ba0ae702c..fd69b375f 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -12,6 +12,7 @@ overrides: size-sensor: 1.0.3 patchedDependencies: + '@databricks/sdk-genie@0.54.0': 664f3e13b155eb1b5499486a4f6fc1ab9a6a502e2c6b200d06749bd25dd9ae54 '@databricks/sdk-statementexecution@0.52.0': 9377a46b883b548d47116662bf73b159a4da51ec28b9ff1896a7e4a7841c240c importers: @@ -288,6 +289,9 @@ importers: '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-genie': + specifier: 0.54.0 + version: 0.54.0(patch_hash=664f3e13b155eb1b5499486a4f6fc1ab9a6a502e2c6b200d06749bd25dd9ae54) '@databricks/sdk-options': specifier: 0.51.0 version: 0.51.0 @@ -621,6 +625,9 @@ importers: '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-genie': + specifier: 0.54.0 + version: 0.54.0(patch_hash=664f3e13b155eb1b5499486a4f6fc1ab9a6a502e2c6b200d06749bd25dd9ae54) '@databricks/sdk-options': specifier: 0.51.0 version: 0.51.0 @@ -2014,6 +2021,10 @@ packages: resolution: {integrity: sha512-dOJIt4F2nBk6HKObnv7Xbmy/qLYTy2835qhXSuW0Qw1QAXui9plmCet1KqG3yeQcMTyncWGbnhjGdQi8GEGQSA==} engines: {node: '>=22.0', npm: '>=10.0.0'} + '@databricks/sdk-genie@0.54.0': + resolution: {integrity: sha512-xT+dyuXocgwPJ5VA1Kr5/dSdKzzR9hRAkEOybeO2i+v891kuuFzVd/M9OM31PtO2/YZp95vdFH5Sb+jVGtm4nw==} + engines: {node: '>=22.0.0'} + '@databricks/sdk-options@0.51.0': resolution: {integrity: sha512-p5uBh64Y1onnvwlEwfRsmKmGRI23Y4Ly8fZiCgNUwVDe4Z9d0ctI0mo1v4gnDqlmBaBml3TZfqn5BJBROCXsEw==} engines: {node: '>=22.0.0'} @@ -14326,6 +14337,15 @@ snapshots: transitivePeerDependencies: - supports-color + '@databricks/sdk-genie@0.54.0(patch_hash=664f3e13b155eb1b5499486a4f6fc1ab9a6a502e2c6b200d06749bd25dd9ae54)': + dependencies: + '@databricks/sdk-auth': 0.51.0 + '@databricks/sdk-core': 0.51.0 + '@databricks/sdk-options': 0.51.0 + '@js-temporal/polyfill': 0.5.1 + json-bigint: 1.0.0 + zod: 4.3.6 + '@databricks/sdk-options@0.51.0': dependencies: '@databricks/sdk-auth': 0.51.0 diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index bc863eab6..3f48fc9af 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -19,4 +19,5 @@ allowBuilds: esbuild: true protobufjs: false patchedDependencies: + '@databricks/sdk-genie@0.54.0': patches/@databricks__sdk-genie@0.54.0.patch '@databricks/sdk-statementexecution@0.52.0': patches/@databricks__sdk-statementexecution@0.52.0.patch