diff --git a/packages/ai-config/providers.schema.json b/packages/ai-config/providers.schema.json index f62ebe7..e893939 100644 --- a/packages/ai-config/providers.schema.json +++ b/packages/ai-config/providers.schema.json @@ -1431,6 +1431,9 @@ "account": { "type": "string" }, + "connectionName": { + "type": "string" + }, "home": { "type": "string" }, @@ -5669,6 +5672,9 @@ "account": { "type": "string" }, + "connectionName": { + "type": "string" + }, "home": { "type": "string" }, diff --git a/packages/ai-config/src/schema.ts b/packages/ai-config/src/schema.ts index 081988a..fe62c8b 100644 --- a/packages/ai-config/src/schema.ts +++ b/packages/ai-config/src/schema.ts @@ -146,6 +146,13 @@ export const snowflakeConfigSchema = z * managed directory. Non-secret. */ home: z.string().optional(), + /** + * Name of the `connections.toml` connection to use (non-secret). Selects + * one entry from the file `home` points at. Consumed by + * `@assistant/node`'s Snowflake resolver on Node platforms; ignored in + * Positron, which defers Snowflake credentials to `vscode.authentication`. + */ + connectionName: z.string().optional(), }) .strict(); diff --git a/packages/ai-config/src/types.ts b/packages/ai-config/src/types.ts index f6cd144..b717b83 100644 --- a/packages/ai-config/src/types.ts +++ b/packages/ai-config/src/types.ts @@ -125,7 +125,7 @@ export interface ResolvedConnection { positaiLogin?: { host?: string; clientId?: string; scope?: string }; aws?: { region?: string; profile?: string }; googleCloud?: { project?: string; location?: string }; - snowflake?: { account?: string; host?: string; home?: string }; + snowflake?: { account?: string; host?: string; home?: string; connectionName?: string }; databricks?: { host?: string }; } diff --git a/packages/ai-credentials/src/types/credentials.ts b/packages/ai-credentials/src/types/credentials.ts index 4a8948f..2a4c164 100644 --- a/packages/ai-credentials/src/types/credentials.ts +++ b/packages/ai-credentials/src/types/credentials.ts @@ -39,6 +39,27 @@ export interface ApiKeyCredentials { apiKey: string; baseUrl?: string; customHeaders?: Record; + /** + * Snowflake Cortex only. Present iff `apiKey` holds a Snowflake **session + * token** (from external-browser SSO) that must be sent as + * `Authorization: Snowflake Token="..."` rather than a Bearer token — its + * presence is the session-auth discriminant; absence means Bearer. Ignored + * by every other provider. + * + * `sessionConnectionIdentity` is an **opaque, client-bound token** minted by + * the credential resolver. It encodes not just the connections.toml connection + * name but a snapshot of the connection this token was acquired from (its + * endpoint), so the reauthentication hook can refresh *that exact* connection + * on expiry (not whatever connection is currently selected) and reject a + * refresh whose re-resolved connection no longer matches the bound snapshot — + * e.g. the same-named connection was edited in place to a different account. + * Consumers treat it as opaque; only the resolver encodes/decodes it. Grouped + * so the flag and its required identity are co-present or absent together. + * + * Synthesized at credential-resolution time and never persisted, so it does + * not affect the on-disk credential format. + */ + snowflake?: { sessionConnectionIdentity: string }; } /** diff --git a/packages/ai-provider-bridge/src/model-clients/SnowflakeClient.ts b/packages/ai-provider-bridge/src/model-clients/SnowflakeClient.ts index b5176c1..982d49b 100644 --- a/packages/ai-provider-bridge/src/model-clients/SnowflakeClient.ts +++ b/packages/ai-provider-bridge/src/model-clients/SnowflakeClient.ts @@ -27,15 +27,164 @@ import { import type { ModelClient, ModelClientChatParams } from "./ModelClient"; import { createOpenAICompatibleFetch } from "./openai-compat-fetch"; +/** + * How the credential's token authenticates with Snowflake Cortex: + * - `"bearer"`: an API/PAT token sent as `Authorization: Bearer` (or `x-api-key`). + * - `"session"`: a session token (external-browser SSO) sent as + * `Authorization: Snowflake Token="..."`. + */ +export type SnowflakeAuthScheme = "bearer" | "session"; + +type FetchFn = (url: string | URL | globalThis.Request, init?: RequestInit) => Promise; + +/** + * Reauthenticate an expired Snowflake **session** for a specific client-bound + * connection identity, returning a fresh session token. Pre-built by the Node + * caller and threaded in via the provider factory; the bridge never constructs + * it. Only session auth uses it. + */ +export type SnowflakeSessionReauth = (sessionConnectionIdentity: string) => Promise; + +/** + * Client-bound session refresh: which connections.toml connection this client's + * token was acquired from, and how to reauthenticate *it* (not whatever + * connection is currently selected) when its session token expires mid-request. + */ +export interface SnowflakeSessionRefresh { + connectionIdentity: string; + reauthenticate: SnowflakeSessionReauth; +} + +/** + * Snowflake error code for an expired — but renewable — session token. Cortex + * returns it with HTTP 200 and does not perform the operation, so it must be + * detected on the response body, not the status. + */ +const SESSION_EXPIRED_CODE = "390112"; + +/** True if a Cortex response body signals an expired session token (390112). */ +function isSessionExpiredBody(bodyText: string): boolean { + // `"code"` may be a JSON string ("390112") or number (390112) — match both. + return new RegExp(`"code"\\s*:\\s*"?${SESSION_EXPIRED_CODE}"?`).test(bodyText); +} + +/** Copy headers, dropping framing headers invalidated by rebuilding the body. */ +function headersWithoutFraming(headers: Headers): Headers { + const copy = new Headers(headers); + copy.delete("content-length"); + copy.delete("content-encoding"); + return copy; +} + +/** + * Classify a Cortex response as a **pre-stream** `390112` (expired session) error + * without disturbing the streaming success path. + * + * Per the Cortex REST contract a response is either a `text/event-stream` (the + * model output) or a JSON error envelope — so branching on the content type is + * exact. A stream is returned **untouched**: the wrapper never reads it, so there + * is no first-token latency, and model output is never scanned for an error code + * that could legitimately appear inside it. Any non-stream response is a small + * envelope: it is buffered in full (so a body split across transport chunks is + * still whole), checked for `390112`, and rebuilt so the SDK — or the reauth + * failure path — can still read it. + * + * A `390112` that instead appeared mid-SSE after emitted output would not be + * caught here — by policy that case surfaces as a retryable failure rather than + * replaying partial output (see the plan's Phase 5 decision gate). + */ +async function peekSessionExpiry( + response: Response, +): Promise<{ expired: boolean; response: Response }> { + const contentType = response.headers.get("content-type")?.toLowerCase() ?? ""; + // A streaming success is the common case: hand it back without reading a byte. + if (contentType.includes("text/event-stream")) { + return { expired: false, response }; + } + // Non-stream response: a small error envelope. Buffer it whole and classify, + // rebuilding the body so it stays readable downstream. + const bodyText = await response.text(); + return { + expired: isSessionExpiredBody(bodyText), + response: new Response(bodyText, { + status: response.status, + statusText: response.statusText, + headers: headersWithoutFraming(response.headers), + }), + }; +} + +/** + * Wrap a fetch so every request authenticates with a Snowflake **session token** + * (`Authorization: Snowflake Token="..."`), replacing any Bearer/x-api-key + * header the SDK set. + * + * When `refresh` is supplied, the wrapper also handles session-token expiry: a + * pre-stream `390112` triggers a single transparent reauthenticate-and-retry of + * *this client's* connection. Without `refresh` it behaves as a plain auth + * rewrite (there is nothing to retry with). + */ +function createSnowflakeSessionFetch( + sessionToken: string, + delegate: FetchFn, + refresh?: SnowflakeSessionRefresh, +): FetchFn { + let currentToken = sessionToken; + const withSessionAuth = (init?: RequestInit): RequestInit => { + const headers = new Headers(init?.headers); + headers.delete("x-api-key"); + headers.set("Authorization", `Snowflake Token="${currentToken}"`); + return { ...init, headers }; + }; + + return async (url, init) => { + const response = await delegate(url, withSessionAuth(init)); + if (!refresh) { + return response; + } + + const peeked = await peekSessionExpiry(response); + if (!peeked.expired) { + return peeked.response; + } + + // Pre-stream 390112: reauthenticate this connection and retry exactly once. + // If reauth fails, surface the original error body to the caller. + let freshToken: string; + try { + freshToken = await refresh.reauthenticate(refresh.connectionIdentity); + } catch { + return peeked.response; + } + currentToken = freshToken; + return delegate(url, withSessionAuth(init)); + }; +} + export class SnowflakeClient implements ModelClient { - private readonly bearerToken: string; + private readonly token: string; private readonly baseUrl: string; + private readonly authScheme: SnowflakeAuthScheme; private readonly customHeaders?: Record; + private readonly sessionRefresh?: SnowflakeSessionRefresh; - constructor(bearerToken: string, baseUrl: string, customHeaders?: Record) { - this.bearerToken = bearerToken; + constructor( + token: string, + baseUrl: string, + authScheme: SnowflakeAuthScheme, + customHeaders?: Record, + sessionRefresh?: SnowflakeSessionRefresh, + ) { + this.token = token; this.baseUrl = baseUrl; + this.authScheme = authScheme; this.customHeaders = customHeaders; + this.sessionRefresh = sessionRefresh; + } + + /** True when the token is a session token needing the `Snowflake Token=` scheme. */ + private get isSessionAuth(): boolean { + return this.authScheme === "session"; } async chat(params: ModelClientChatParams): Promise> { @@ -71,11 +220,20 @@ export class SnowflakeClient implements ModelClient { baseUrl: string, ): Promise> { const headers = safeSdkCustomHeaders(this.customHeaders); - const provider = createAnthropic({ - authToken: this.bearerToken, - baseURL: baseUrl, - ...(headers && { headers }), - }); + const provider = this.isSessionAuth + ? createAnthropic({ + // Auth is applied by the session fetch wrapper; this placeholder key + // just satisfies the SDK (its x-api-key header is stripped there). + apiKey: "session-auth", + baseURL: baseUrl, + fetch: createSnowflakeSessionFetch(this.token, globalThis.fetch, this.sessionRefresh), + ...(headers && { headers }), + }) + : createAnthropic({ + authToken: this.token, + baseURL: baseUrl, + ...(headers && { headers }), + }); const model = provider(params.model); const { abortController, cleanup } = createAbortControllerFromToken(params.cancellationToken); @@ -127,10 +285,18 @@ export class SnowflakeClient implements ModelClient { params: ModelClientChatParams, baseUrl: string, ): Promise> { + // The compat fetch applies OpenAI-spec fix-ups. For session auth we keep a + // non-empty apiKey so it does NOT strip the Authorization header, and wrap + // it so the outer fetch installs the `Snowflake Token=` header last. + const compatFetch = this.isSessionAuth + ? createOpenAICompatibleFetch("Snowflake", "session-auth", this.customHeaders) + : createOpenAICompatibleFetch("Snowflake", this.token, this.customHeaders); const provider = createOpenAI({ - apiKey: this.bearerToken || "sk-placeholder", + apiKey: this.token || "sk-placeholder", baseURL: baseUrl, - fetch: createOpenAICompatibleFetch("Snowflake", this.bearerToken, this.customHeaders), + fetch: this.isSessionAuth + ? createSnowflakeSessionFetch(this.token, compatFetch, this.sessionRefresh) + : compatFetch, }); const model = provider.chat(params.model); diff --git a/packages/ai-provider-bridge/src/model-clients/__tests__/snowflake-session-auth.test.ts b/packages/ai-provider-bridge/src/model-clients/__tests__/snowflake-session-auth.test.ts new file mode 100644 index 0000000..f04fe63 --- /dev/null +++ b/packages/ai-provider-bridge/src/model-clients/__tests__/snowflake-session-auth.test.ts @@ -0,0 +1,393 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +// Capture the options each SDK factory receives so we can assert on the auth +// scheme and drive the installed fetch wrapper. `vi.hoisted` lets these mocks +// exist before the hoisted `vi.mock` factories run. +const { createAnthropic, createOpenAI } = vi.hoisted(() => ({ + createAnthropic: vi.fn(() => vi.fn(() => ({}))), + createOpenAI: vi.fn(() => ({ chat: vi.fn(() => ({})) })), +})); + +vi.mock("@ai-sdk/anthropic", () => ({ createAnthropic })); +vi.mock("@ai-sdk/openai", () => ({ createOpenAI })); +vi.mock("ai", () => ({ streamText: vi.fn(() => ({ fullStream: {} })) })); +// Bypass the stream-conversion + abort plumbing; we only care about the auth +// headers the client's fetch wrapper produces. +vi.mock("../ai-sdk-helpers", () => ({ + convertAiSdkStreamToPlatform: vi.fn(() => (async function* () {})()), + createAbortControllerFromToken: vi.fn(() => ({ + abortController: new AbortController(), + cleanup: vi.fn(), + })), + createStepLogger: vi.fn(() => undefined), +})); + +import { ProviderRegistry } from "../../providers/ProviderRegistry"; +import { + registerSnowflakeCortexProvider, + type SnowflakeProviderCallbacks, +} from "../../providers/snowflake-cortex-provider"; +import type { CancellationToken, Logger, ProviderCredentials } from "../../types"; +import type { ModelClient, ModelClientChatParams } from "../ModelClient"; +import { SnowflakeClient, type SnowflakeSessionRefresh } from "../SnowflakeClient"; + +const cancellationToken: CancellationToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose() {} }), +}; + +const mockLogger: Logger = { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + trace: vi.fn(), +}; + +const BASE_URL = "https://acct.snowflakecomputing.com/api/v2/cortex"; +const SESSION_TOKEN = "sess-tok-123"; + +// Claude models route through the Anthropic Messages API; others through OpenAI. +const CLAUDE_MODEL = "claude-opus-4-7"; +const OPENAI_MODEL = "openai-gpt-5.2"; + +type SdkFetch = (url: string | URL | Request, init?: RequestInit) => Promise; + +interface AnthropicOptions { + apiKey?: string; + authToken?: string; + baseURL?: string; + fetch?: SdkFetch; + headers?: Record; +} + +interface OpenAIOptions { + apiKey?: string; + baseURL?: string; + fetch?: SdkFetch; +} + +function anthropicOptions(): AnthropicOptions { + return createAnthropic.mock.calls[0]?.[0] as AnthropicOptions; +} + +function openaiOptions(): OpenAIOptions { + return createOpenAI.mock.calls[0]?.[0] as OpenAIOptions; +} + +const params = (model: string): ModelClientChatParams => ({ + model, + messages: [], + maxOutputTokens: 1024, + cancellationToken, +}); + +/** Headers seen by the terminal `globalThis.fetch` after the wrappers run. */ +function outgoingHeaders(spy: ReturnType): Headers { + return new Headers((spy.mock.calls[0]?.[1] as RequestInit | undefined)?.headers); +} + +describe("SnowflakeClient auth schemes", () => { + beforeEach(() => { + createAnthropic.mockClear(); + createOpenAI.mockClear(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("bearer scheme: Anthropic route uses authToken, no fetch override, forwards custom headers verbatim", async () => { + await new SnowflakeClient("bearer-xyz", BASE_URL, "bearer", { "x-gateway": "keep-me" }).chat( + params(CLAUDE_MODEL), + ); + + const opts = anthropicOptions(); + expect(opts.authToken).toBe("bearer-xyz"); + expect(opts.apiKey).toBeUndefined(); + expect(opts.fetch).toBeUndefined(); + // customHeaders are additive gateway headers only — passed through untouched. + expect(opts.headers).toEqual({ "x-gateway": "keep-me" }); + }); + + it("session scheme: Anthropic route sends `Snowflake Token=` auth and strips x-api-key", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(new Response(null)); + + await new SnowflakeClient(SESSION_TOKEN, BASE_URL, "session", { "x-gateway": "keep-me" }).chat( + params(CLAUDE_MODEL), + ); + + const opts = anthropicOptions(); + // The placeholder key satisfies the SDK; real auth is applied by the wrapper. + expect(opts.apiKey).toBe("session-auth"); + expect(opts.authToken).toBeUndefined(); + expect(opts.headers).toEqual({ "x-gateway": "keep-me" }); + + // Drive the installed fetch wrapper as the SDK would (SDK sets x-api-key). + await opts.fetch?.("https://x/v1/messages", { + headers: { "x-api-key": "session-auth" }, + }); + + const sent = outgoingHeaders(fetchSpy); + expect(sent.get("authorization")).toBe(`Snowflake Token="${SESSION_TOKEN}"`); + expect(sent.has("x-api-key")).toBe(false); + }); + + it("session scheme: OpenAI route sends `Snowflake Token=` auth through the compat fetch and keeps gateway headers", async () => { + const fetchSpy = vi + .spyOn(globalThis, "fetch") + .mockResolvedValue(new Response(null, { headers: { "content-type": "application/json" } })); + + await new SnowflakeClient(SESSION_TOKEN, BASE_URL, "session", { "x-gateway": "keep-me" }).chat( + params(OPENAI_MODEL), + ); + + const opts = openaiOptions(); + expect(opts.fetch).toBeDefined(); + + await opts.fetch?.("https://x/chat/completions", { + headers: { Authorization: "Bearer session-auth", "x-api-key": "session-auth" }, + body: JSON.stringify({ max_tokens: 10, messages: [] }), + }); + + const sent = outgoingHeaders(fetchSpy); + // Session wrapper installs `Snowflake Token=` last, over the SDK's Bearer. + expect(sent.get("authorization")).toBe(`Snowflake Token="${SESSION_TOKEN}"`); + expect(sent.has("x-api-key")).toBe(false); + // The additive gateway header still reaches the wire. + expect(sent.get("x-gateway")).toBe("keep-me"); + }); +}); + +// Cover the production seam in the registered client factory: the mapping from +// the `credentials.snowflake` group to the client's auth scheme. Constructing +// SnowflakeClient directly (above) bypasses this, so a factory regression could +// leave the client tests green while the product silently uses Bearer auth. +describe("Snowflake provider factory seam", () => { + function createClient(credentials: ProviderCredentials): ModelClient { + const registry = new ProviderRegistry(mockLogger); + registerSnowflakeCortexProvider(registry, mockLogger); + const client = registry.getClientForProvider("snowflake-cortex", credentials); + if (!client) throw new Error("expected a Snowflake client from the factory"); + return client; + } + + beforeEach(() => { + createAnthropic.mockClear(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("maps a present `snowflake` group to session auth (`Snowflake Token=`)", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(new Response(null)); + + const client = createClient({ + type: "apikey", + apiKey: SESSION_TOKEN, + baseUrl: BASE_URL, + snowflake: { sessionConnectionIdentity: "dev" }, + }); + await client.chat(params(CLAUDE_MODEL)); + + const opts = anthropicOptions(); + expect(opts.apiKey).toBe("session-auth"); + expect(opts.authToken).toBeUndefined(); + + await opts.fetch?.("https://x/v1/messages", { headers: { "x-api-key": "session-auth" } }); + expect(outgoingHeaders(fetchSpy).get("authorization")).toBe( + `Snowflake Token="${SESSION_TOKEN}"`, + ); + }); + + it("maps an unflagged credential to Bearer auth", async () => { + const client = createClient({ type: "apikey", apiKey: "bearer-xyz", baseUrl: BASE_URL }); + await client.chat(params(CLAUDE_MODEL)); + + const opts = anthropicOptions(); + expect(opts.authToken).toBe("bearer-xyz"); + expect(opts.fetch).toBeUndefined(); + }); +}); + +// Session-token expiry (390112) detect + retry in the fetch wrapper. Cortex +// returns HTTP 200 with a `390112` body when a session token has expired; the +// wrapper must reauthenticate the *client-bound* connection and retry once, +// transparently, on the pre-stream case (see the plan's Phase 5 decision gate). +describe("SnowflakeClient session-token expiry (390112)", () => { + const REFRESH_IDENTITY = "dev"; + const FRESH_TOKEN = "fresh-session-tok"; + const EXPIRED_BODY = JSON.stringify({ code: "390112", message: "session token has expired" }); + + /** + * A Cortex error response: HTTP 200 with a JSON envelope. Per the REST + * contract, errors (including 390112) are `application/json`, distinct from the + * `text/event-stream` of a real model stream — which is how the wrapper tells a + * pre-stream error from streamed output without reading the stream body. + */ + function jsonError(body: string): Response { + return new Response(body, { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + + /** A streaming (`text/event-stream`) success response. */ + function eventStream(body: string): Response { + return new Response(body, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + } + + beforeEach(() => { + createAnthropic.mockClear(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + function refresh( + reauthenticate: SnowflakeSessionRefresh["reauthenticate"], + ): SnowflakeSessionRefresh { + return { connectionIdentity: REFRESH_IDENTITY, reauthenticate }; + } + + /** Build the client's installed fetch wrapper for the Anthropic route. */ + async function sessionFetch(sessionRefresh: SnowflakeSessionRefresh): Promise { + await new SnowflakeClient(SESSION_TOKEN, BASE_URL, "session", undefined, sessionRefresh).chat( + params(CLAUDE_MODEL), + ); + const wrapper = anthropicOptions().fetch; + if (!wrapper) throw new Error("expected a session fetch wrapper"); + return wrapper; + } + + function authHeader(spy: ReturnType, callIndex: number): string | null { + return new Headers((spy.mock.calls[callIndex]?.[1] as RequestInit | undefined)?.headers).get( + "authorization", + ); + } + + it("pre-stream 390112: reauthenticates the client-bound connection and retries once with the fresh token", async () => { + const reauthenticate = vi.fn(async () => FRESH_TOKEN); + const fetchSpy = vi + .spyOn(globalThis, "fetch") + .mockResolvedValueOnce(jsonError(EXPIRED_BODY)) + .mockResolvedValueOnce(eventStream("data: ok\n\n")); + + const wrapper = await sessionFetch(refresh(reauthenticate)); + const result = await wrapper("https://x/v1/messages", { + headers: { "x-api-key": "session-auth" }, + }); + + // Refreshed *this* connection, not whatever might be selected. + expect(reauthenticate).toHaveBeenCalledWith(REFRESH_IDENTITY); + expect(fetchSpy).toHaveBeenCalledTimes(2); + // First attempt used the original token; the retry used the refreshed one. + expect(authHeader(fetchSpy, 0)).toBe(`Snowflake Token="${SESSION_TOKEN}"`); + expect(authHeader(fetchSpy, 1)).toBe(`Snowflake Token="${FRESH_TOKEN}"`); + // The caller receives the successful retry stream, not the 390112 body. + expect(await result.text()).toContain("ok"); + }); + + it("event-stream success: passes the response through without reading its body (no first-token latency)", async () => { + // The wrapper must not buffer a streaming success while deciding whether it + // is a 390112 error — content-type alone is enough. Prove it never touches + // the body: a low-volume stream that keeps its connection open would + // otherwise have its first tokens withheld. + const reauthenticate = vi.fn(async () => FRESH_TOKEN); + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode("data: hello\n\n")); + // Deliberately left open — no close(), no further chunks. + }, + }); + const response = new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(response); + + const wrapper = await sessionFetch(refresh(reauthenticate)); + const result = await wrapper("https://x/v1/messages", { + headers: { "x-api-key": "session-auth" }, + }); + + expect(reauthenticate).not.toHaveBeenCalled(); + expect(fetchSpy).toHaveBeenCalledTimes(1); + // Same response instance, body untouched — the SDK (not the wrapper) reads it. + expect(result).toBe(response); + expect(response.bodyUsed).toBe(false); + }); + + it("event-stream success: streamed bytes reach the caller intact", async () => { + const reauthenticate = vi.fn(async () => FRESH_TOKEN); + const body = "data: hello\n\ndata: world\n\n"; + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(eventStream(body)); + + const wrapper = await sessionFetch(refresh(reauthenticate)); + const result = await wrapper("https://x/v1/messages", { + headers: { "x-api-key": "session-auth" }, + }); + + expect(reauthenticate).not.toHaveBeenCalled(); + expect(fetchSpy).toHaveBeenCalledTimes(1); + expect(await result.text()).toBe(body); + }); + + it("reauth failure: surfaces the original 390112 error and does not retry", async () => { + const reauthenticate = vi.fn(async () => { + throw new Error("SSO window closed"); + }); + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(jsonError(EXPIRED_BODY)); + + const wrapper = await sessionFetch(refresh(reauthenticate)); + const result = await wrapper("https://x/v1/messages", { + headers: { "x-api-key": "session-auth" }, + }); + + expect(reauthenticate).toHaveBeenCalledTimes(1); + expect(fetchSpy).toHaveBeenCalledTimes(1); + // Original error body is preserved for the caller to surface. + expect(await result.text()).toContain("390112"); + }); + + it("factory seam: the client-bound identity comes from the credentials, not the current selection", async () => { + // A client built from connection "A" must refresh "A" even though the hook + // is a generic "refresh whatever identity you're given" — the identity is + // captured from the credential at construction, so a later selection of "B" + // cannot redirect this in-flight request's retry. + const reauthenticate = vi.fn(async () => FRESH_TOKEN); + const callbacks: SnowflakeProviderCallbacks = { reauthenticateSession: reauthenticate }; + const registry = new ProviderRegistry(mockLogger); + registerSnowflakeCortexProvider(registry, mockLogger, callbacks); + const client = registry.getClientForProvider("snowflake-cortex", { + type: "apikey", + apiKey: SESSION_TOKEN, + baseUrl: BASE_URL, + snowflake: { sessionConnectionIdentity: "A" }, + }); + if (!client) throw new Error("expected a Snowflake client from the factory"); + + const fetchSpy = vi + .spyOn(globalThis, "fetch") + .mockResolvedValueOnce(jsonError(EXPIRED_BODY)) + .mockResolvedValueOnce(eventStream("data: ok\n\n")); + + await client.chat(params(CLAUDE_MODEL)); + await anthropicOptions().fetch?.("https://x/v1/messages", { + headers: { "x-api-key": "session-auth" }, + }); + + expect(reauthenticate).toHaveBeenCalledWith("A"); + expect(fetchSpy).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/ai-provider-bridge/src/provider-registration.ts b/packages/ai-provider-bridge/src/provider-registration.ts index 7b049b8..b5449e8 100644 --- a/packages/ai-provider-bridge/src/provider-registration.ts +++ b/packages/ai-provider-bridge/src/provider-registration.ts @@ -10,6 +10,7 @@ import type { BedrockProviderCallbacks } from "./providers/bedrock-provider"; import type { GoogleVertexProviderCallbacks } from "./providers/google-vertex-provider"; import type { ProviderRegistry } from "./providers/ProviderRegistry"; +import type { SnowflakeProviderCallbacks } from "./providers/snowflake-cortex-provider"; import type { Logger, ProviderId } from "./types"; export interface ProviderRegistrationConfig { @@ -26,6 +27,7 @@ export interface ProviderRegistrationConfig { /** Pre-built by the caller. The bridge must NOT construct these. */ bedrockCallbacks?: BedrockProviderCallbacks; googleVertexCallbacks?: GoogleVertexProviderCallbacks; + snowflakeCallbacks?: SnowflakeProviderCallbacks; } /** diff --git a/packages/ai-provider-bridge/src/providers.ts b/packages/ai-provider-bridge/src/providers.ts index eec8c60..b3f654b 100644 --- a/packages/ai-provider-bridge/src/providers.ts +++ b/packages/ai-provider-bridge/src/providers.ts @@ -26,6 +26,7 @@ export { registerOpenAIProvider } from "./providers/openai-provider"; export { registerOpenRouterProvider } from "./providers/openrouter-provider"; export { registerPositAiProvider } from "./providers/positai-provider"; export { registerSnowflakeCortexProvider } from "./providers/snowflake-cortex-provider"; +export type { SnowflakeProviderCallbacks } from "./providers/snowflake-cortex-provider"; // Provider registration orchestrator export { registerAllProviders } from "./register-all-providers"; @@ -78,7 +79,7 @@ export { OllamaClient, ollamaThinkParam } from "./model-clients/OllamaClient"; export { OpenAIClient } from "./model-clients/OpenAIClient"; export { OpenRouterClient } from "./model-clients/OpenRouterClient"; export { PositAiClient } from "./model-clients/PositAiClient"; -export { SnowflakeClient } from "./model-clients/SnowflakeClient"; +export { SnowflakeClient, type SnowflakeAuthScheme } from "./model-clients/SnowflakeClient"; // AI SDK helpers export * from "./model-clients/ai-sdk-helpers"; diff --git a/packages/ai-provider-bridge/src/providers/snowflake-cortex-provider.ts b/packages/ai-provider-bridge/src/providers/snowflake-cortex-provider.ts index 5030b0e..9abbee6 100644 --- a/packages/ai-provider-bridge/src/providers/snowflake-cortex-provider.ts +++ b/packages/ai-provider-bridge/src/providers/snowflake-cortex-provider.ts @@ -5,11 +5,25 @@ import { getAnthropicModelCapabilities, getOpenAIModelCapabilities } from "ai-config"; import { SnowflakeClient } from "../model-clients/SnowflakeClient"; +import type { SnowflakeSessionReauth } from "../model-clients/SnowflakeClient"; import type { Logger, ModelInfo, ProviderCredentials } from "../types"; import type { ProviderRegistry } from "./ProviderRegistry"; const IMAGE_MEDIA_TYPES = ["image/png", "image/jpeg", "image/gif", "image/webp"] as const; +/** + * Platform-provided Snowflake hooks. Pre-built by the Node caller and threaded in + * through registration; the bridge must NOT construct these. + */ +export interface SnowflakeProviderCallbacks { + /** + * Reauthenticate an expired Snowflake **session** for a specific client-bound + * connection identity, returning a fresh session token. Only session auth + * (external-browser SSO) uses it; Bearer credentials never expire mid-request. + */ + reauthenticateSession: SnowflakeSessionReauth; +} + /** * Helper to build a Claude model entry (Anthropic Messages API protocol). * All Claude models share the same capabilities on Snowflake Cortex. @@ -90,7 +104,11 @@ const SNOWFLAKE_MODELS: ModelInfo[] = [ openaiModel("openai-gpt-5.1", "GPT-5.1", { supportsImages: true }), ]; -export function registerSnowflakeCortexProvider(registry: ProviderRegistry, logger: Logger): void { +export function registerSnowflakeCortexProvider( + registry: ProviderRegistry, + logger: Logger, + callbacks?: SnowflakeProviderCallbacks, +): void { // Static model fetcher — Snowflake Cortex serves a known set of models. const fetcher = async (credentials: ProviderCredentials): Promise => { if (credentials.type !== "apikey") { @@ -112,6 +130,22 @@ export function registerSnowflakeCortexProvider(registry: ProviderRegistry, logg if (credentials.type !== "apikey") { throw new Error(`Snowflake provider requires API key credentials, got: ${credentials.type}`); } - return new SnowflakeClient(credentials.apiKey, credentials.baseUrl!, credentials.customHeaders); + // Session auth wires a client-bound refresh so an expired token retries the + // connection *this* token came from, not whatever is currently selected. + const session = credentials.snowflake; + const sessionRefresh = + session && callbacks + ? { + connectionIdentity: session.sessionConnectionIdentity, + reauthenticate: callbacks.reauthenticateSession, + } + : undefined; + return new SnowflakeClient( + credentials.apiKey, + credentials.baseUrl!, + session ? "session" : "bearer", + credentials.customHeaders, + sessionRefresh, + ); }); } diff --git a/packages/ai-provider-bridge/src/register-all-providers.ts b/packages/ai-provider-bridge/src/register-all-providers.ts index fc6e3ee..03daed6 100644 --- a/packages/ai-provider-bridge/src/register-all-providers.ts +++ b/packages/ai-provider-bridge/src/register-all-providers.ts @@ -53,8 +53,8 @@ type ProviderRegistrar = ( * id set equals PROVIDER_IDS (the single source of truth): a mislabeled, duplicated, or missing * id here would silently corrupt `allowedProviders` filtering, which keys on these labels. * - * Only positai/bedrock/google-vertex need a wrapper to thread config into a non-uniform - * signature; the rest reference their `(registry, logger)` register fn directly. + * Only positai/bedrock/google-vertex/snowflake-cortex need a wrapper to thread config into a + * non-uniform signature; the rest reference their `(registry, logger)` register fn directly. */ export const PROVIDER_REGISTRARS: readonly [ProviderId, ProviderRegistrar][] = [ [ @@ -81,7 +81,11 @@ export const PROVIDER_REGISTRARS: readonly [ProviderId, ProviderRegistrar][] = [ ["gemini", registerGeminiProvider], ["openai-compatible", registerOpenAICompatibleProvider], ["ms-foundry", registerFoundryProvider], - ["snowflake-cortex", registerSnowflakeCortexProvider], + [ + "snowflake-cortex", + (registry, logger, config) => + registerSnowflakeCortexProvider(registry, logger, config.snowflakeCallbacks), + ], ["deepseek", registerDeepSeekProvider], ["databricks", registerDatabricksProvider], ];