diff --git a/packages/ai-config/providers.schema.json b/packages/ai-config/providers.schema.json index 850e9ad..7ea1093 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" + }, "host": { "type": "string" } @@ -5431,6 +5434,9 @@ "account": { "type": "string" }, + "connectionName": { + "type": "string" + }, "host": { "type": "string" } diff --git a/packages/ai-config/src/schema.ts b/packages/ai-config/src/schema.ts index def9143..9a4c958 100644 --- a/packages/ai-config/src/schema.ts +++ b/packages/ai-config/src/schema.ts @@ -139,6 +139,12 @@ export const snowflakeConfigSchema = z .object({ account: z.string().optional(), host: z.string().optional(), + /** + * Name of the `connections.toml` connection to use (non-secret). 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 ece907d..ad4cfc0 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 }; + snowflake?: { account?: string; host?: string; connectionName?: 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..3a5fae3 100644 --- a/packages/ai-provider-bridge/src/model-clients/SnowflakeClient.ts +++ b/packages/ai-provider-bridge/src/model-clients/SnowflakeClient.ts @@ -27,6 +27,33 @@ import { import type { ModelClient, ModelClientChatParams } from "./ModelClient"; import { createOpenAICompatibleFetch } from "./openai-compat-fetch"; +/** Header carrying the Snowflake auth scheme (real token types are forwarded). */ +const TOKEN_TYPE_HEADER = "X-Snowflake-Authorization-Token-Type"; +/** + * Internal sentinel token-type: the credential's `apiKey` is a Snowflake session + * token (from external-browser SSO) that must be sent as + * `Authorization: Snowflake Token="..."` rather than a Bearer token. This value + * is consumed by the client and never forwarded to Snowflake. + */ +const SESSION_TOKEN_TYPE = "SESSION"; + +type FetchFn = (url: string | URL | globalThis.Request, init?: RequestInit) => Promise; + +/** + * Wrap a fetch so every request authenticates with a Snowflake **session token** + * (`Authorization: Snowflake Token="..."`). Removes any Bearer/x-api-key header + * the SDK set and the internal token-type sentinel before the request goes out. + */ +function createSnowflakeSessionFetch(sessionToken: string, delegate: FetchFn): FetchFn { + return async (url, init) => { + const headers = new Headers(init?.headers); + headers.delete("x-api-key"); + headers.delete(TOKEN_TYPE_HEADER); + headers.set("Authorization", `Snowflake Token="${sessionToken}"`); + return delegate(url, { ...init, headers }); + }; +} + export class SnowflakeClient implements ModelClient { private readonly bearerToken: string; private readonly baseUrl: string; @@ -38,6 +65,19 @@ export class SnowflakeClient implements ModelClient { this.customHeaders = customHeaders; } + /** True when the token is a session token needing the `Snowflake Token=` scheme. */ + private get isSessionAuth(): boolean { + return this.customHeaders?.[TOKEN_TYPE_HEADER] === SESSION_TOKEN_TYPE; + } + + /** customHeaders to forward, with the internal session sentinel removed. */ + private forwardedHeaders(): Record | undefined { + if (!this.customHeaders || !this.isSessionAuth) return this.customHeaders; + const rest = { ...this.customHeaders }; + delete rest[TOKEN_TYPE_HEADER]; + return Object.keys(rest).length > 0 ? rest : undefined; + } + async chat(params: ModelClientChatParams): Promise> { const effectiveBaseUrl = params.baseUrl ?? this.baseUrl; @@ -70,12 +110,21 @@ export class SnowflakeClient implements ModelClient { params: ModelClientChatParams, baseUrl: string, ): Promise> { - const headers = safeSdkCustomHeaders(this.customHeaders); - const provider = createAnthropic({ - authToken: this.bearerToken, - baseURL: baseUrl, - ...(headers && { headers }), - }); + const headers = safeSdkCustomHeaders(this.forwardedHeaders()); + 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.bearerToken, globalThis.fetch), + ...(headers && { headers }), + }) + : createAnthropic({ + authToken: this.bearerToken, + baseURL: baseUrl, + ...(headers && { headers }), + }); const model = provider(params.model); const { abortController, cleanup } = createAbortControllerFromToken(params.cancellationToken); @@ -127,10 +176,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.forwardedHeaders()) + : createOpenAICompatibleFetch("Snowflake", this.bearerToken, this.customHeaders); const provider = createOpenAI({ apiKey: this.bearerToken || "sk-placeholder", baseURL: baseUrl, - fetch: createOpenAICompatibleFetch("Snowflake", this.bearerToken, this.customHeaders), + fetch: this.isSessionAuth + ? createSnowflakeSessionFetch(this.bearerToken, compatFetch) + : compatFetch, }); const model = provider.chat(params.model);