diff --git a/package-lock.json b/package-lock.json index 3c531cf..9f72bfe 100644 --- a/package-lock.json +++ b/package-lock.json @@ -6,6 +6,7 @@ "packages": { "": { "name": "ai-lib", + "license": "MIT", "workspaces": [ "packages/*" ], @@ -6258,6 +6259,7 @@ }, "packages/ai-config": { "version": "0.0.1", + "license": "MIT", "dependencies": { "jsonc-parser": "^3.3.1", "proper-lockfile": "^4.1.2", @@ -6309,22 +6311,9 @@ "@typescript/typescript-win32-x64": "7.0.2" } }, - "packages/ai-credential-store": { - "version": "0.0.1", - "extraneous": true, - "dependencies": { - "chokidar": "^5.0.0", - "proper-lockfile": "^4.1.2" - }, - "devDependencies": { - "@types/node": "22.x", - "@types/proper-lockfile": "^4.1.4", - "typescript": "^6.0.2", - "vitest": "^3.2.1" - } - }, "packages/ai-credentials": { "version": "0.0.1", + "license": "MIT", "dependencies": { "chokidar": "^5.0.0", "proper-lockfile": "^4.1.2", diff --git a/packages/ai-config/providers.schema.json b/packages/ai-config/providers.schema.json index 8817b2e..c73fe57 100644 --- a/packages/ai-config/providers.schema.json +++ b/packages/ai-config/providers.schema.json @@ -482,6 +482,236 @@ }, "type": "object" }, + "connect": { + "additionalProperties": false, + "properties": { + "baseUrl": { + "type": "string" + }, + "customHeaders": { + "additionalProperties": { + "type": "string" + }, + "propertyNames": { + "type": "string" + }, + "type": "object" + }, + "enabled": { + "type": "boolean" + }, + "endpoint": { + "type": "string" + }, + "endpoints": { + "additionalProperties": { + "type": "string" + }, + "propertyNames": { + "enum": [ + "anthropic-messages", + "openai-chat", + "openai-responses", + "mlflow-responses", + "bedrock-converse", + "google-generative" + ], + "type": "string" + }, + "type": "object" + }, + "models": { + "additionalProperties": false, + "properties": { + "allow": { + "items": { + "type": "string" + }, + "type": "array" + }, + "custom": { + "items": { + "additionalProperties": false, + "properties": { + "baseUrl": { + "type": "string" + }, + "family": { + "type": "string" + }, + "id": { + "minLength": 1, + "type": "string" + }, + "maxContextLength": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "maxInputTokens": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "maxOutputTokens": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "name": { + "minLength": 1, + "type": "string" + }, + "protocol": { + "enum": [ + "anthropic-messages", + "openai-chat", + "openai-responses", + "mlflow-responses", + "bedrock-converse", + "google-generative" + ], + "type": "string" + }, + "supportedInputMediaTypes": { + "items": { + "type": "string" + }, + "type": "array" + }, + "supportsImages": { + "type": "boolean" + }, + "supportsToolResultImages": { + "type": "boolean" + }, + "supportsTools": { + "type": "boolean" + }, + "supportsWebSearch": { + "type": "boolean" + }, + "thinkingEffortLevels": { + "items": { + "type": "string" + }, + "type": "array" + } + }, + "required": [ + "id", + "name", + "maxContextLength", + "supportsTools", + "supportsImages", + "supportsToolResultImages", + "supportsWebSearch" + ], + "type": "object" + }, + "type": "array" + }, + "deny": { + "items": { + "type": "string" + }, + "type": "array" + }, + "discovery": { + "enum": [ + "auto", + "off" + ], + "type": "string" + }, + "overrides": { + "additionalProperties": { + "additionalProperties": false, + "properties": { + "baseUrl": { + "type": "string" + }, + "family": { + "type": "string" + }, + "maxContextLength": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "maxInputTokens": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "maxOutputTokens": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "name": { + "type": "string" + }, + "protocol": { + "enum": [ + "anthropic-messages", + "openai-chat", + "openai-responses", + "mlflow-responses", + "bedrock-converse", + "google-generative" + ], + "type": "string" + }, + "supportedInputMediaTypes": { + "items": { + "type": "string" + }, + "type": "array" + }, + "supportsImages": { + "type": "boolean" + }, + "supportsToolResultImages": { + "type": "boolean" + }, + "supportsTools": { + "type": "boolean" + }, + "supportsWebSearch": { + "type": "boolean" + }, + "thinkingEffortLevels": { + "items": { + "type": "string" + }, + "type": "array" + } + }, + "type": "object" + }, + "propertyNames": { + "type": "string" + }, + "type": "object" + } + }, + "type": "object" + }, + "protocol": { + "enum": [ + "anthropic-messages", + "openai-chat", + "openai-responses", + "mlflow-responses", + "bedrock-converse", + "google-generative" + ], + "type": "string" + } + }, + "type": "object" + }, "copilot": { "additionalProperties": false, "properties": { diff --git a/packages/ai-config/src/build-catalog.ts b/packages/ai-config/src/build-catalog.ts index 4e06735..0856b0c 100644 --- a/packages/ai-config/src/build-catalog.ts +++ b/packages/ai-config/src/build-catalog.ts @@ -58,6 +58,7 @@ const BUILTIN_CLIENT_KIND = { databricks: "databricks", litellm: "litellm", portkey: "portkey", + connect: "connect", } as const satisfies Record; /** diff --git a/packages/ai-config/src/connection-env.ts b/packages/ai-config/src/connection-env.ts index 84f797e..04ba9ad 100644 --- a/packages/ai-config/src/connection-env.ts +++ b/packages/ai-config/src/connection-env.ts @@ -86,6 +86,9 @@ const CONNECTION_ENV_MAPPINGS: Partial { + it("borrows limits, family, and media types from the Anthropic-on-Bedrock table", () => { + const id = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"; + const claude = getAnthropicModelCapabilities(id)!; + + const caps = getConnectBedrockModelCapabilities(id); + + expect(caps.family).toBe(claude.family); + expect(caps.family).toBeDefined(); + expect(caps.maxContextLength).toBe(claude.maxContextLength); + expect(caps.maxOutputTokens).toBe(claude.maxOutputTokens); + expect(caps.maxInputTokens).toBe(caps.maxContextLength - caps.maxOutputTokens); + expect(caps.supportedInputMediaTypes).toEqual(claude.supportedInputMediaTypes); + expect(caps).toMatchObject({ + supportsTools: true, + supportsImages: true, + supportsToolResultImages: true, + supportsWebSearch: false, + }); + }); + + it("falls back to conservative limits and image flags for an id the Anthropic table cannot answer", () => { + const caps = getConnectBedrockModelCapabilities("mistral.mistral-large-2407-v1:0"); + + expect(caps.family).toBeUndefined(); + expect(caps.thinkingEffortLevels).toBeUndefined(); + expect(caps.supportedInputMediaTypes).toBeUndefined(); + expect(caps.maxContextLength).toBe(200_000); + expect(caps.maxOutputTokens).toBe(4_096); + expect(caps.maxInputTokens).toBe(caps.maxContextLength - caps.maxOutputTokens); + // Vision is not the norm across Bedrock chat models; over-reporting it + // produces hard API errors, so unknown ids must not claim image support. + expect(caps.supportsImages).toBe(false); + expect(caps.supportsToolResultImages).toBe(false); + expect(caps.supportsTools).toBe(true); + }); +}); diff --git a/packages/ai-config/src/model-capabilities/connect-helpers.ts b/packages/ai-config/src/model-capabilities/connect-helpers.ts new file mode 100644 index 0000000..0faa37f --- /dev/null +++ b/packages/ai-config/src/model-capabilities/connect-helpers.ts @@ -0,0 +1,83 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import { getAnthropicModelCapabilities } from "./anthropic-helpers.js"; + +/** + * Bedrock models offered through a Posit Connect gateway route. Declared + * rather than discovered: Connect's Bedrock gateway proxies Bedrock *Runtime* + * operations only and rejects ListFoundationModels/ListInferenceProfiles + * (connect/src/connect/gateway/bedrock/provider.go), so live discovery would + * have to bypass the gateway and hit AWS directly with Connect-minted + * credentials. + */ +export const CONNECT_BEDROCK_MODELS: readonly { id: string; name: string }[] = [ + { id: "us.anthropic.claude-sonnet-4-5-20250929-v1:0", name: "Claude Sonnet 4.5" }, + { id: "us.anthropic.claude-opus-4-1-20250805-v1:0", name: "Claude Opus 4.1" }, + { id: "us.anthropic.claude-3-5-haiku-20241022-v1:0", name: "Claude 3.5 Haiku" }, +]; + +export const CONNECT_BEDROCK_MODEL_IDS: readonly string[] = CONNECT_BEDROCK_MODELS.map( + (model) => model.id, +); + +/** + * Context window and output cap for an id outside the Anthropic table — only + * reachable via a user-configured override of the declared list. Deliberately + * conservative: an under-reported cap wastes budget, an over-reported one + * produces API errors. + */ +const FALLBACK_CONTEXT_LENGTH = 200_000; +const FALLBACK_MAX_OUTPUT_TOKENS = 4_096; + +/** + * Capabilities of a Bedrock model served through a Connect gateway route. + * Every field is resolved — `family`, `thinkingEffortLevels`, and + * `supportedInputMediaTypes` are `undefined` when the Anthropic table knows + * nothing about the id, not absent. + */ +export interface ConnectBedrockModelCapabilities { + family: string | undefined; + thinkingEffortLevels: string[] | undefined; + maxContextLength: number; + maxInputTokens: number; + maxOutputTokens: number; + supportsTools: boolean; + supportsImages: boolean; + supportsToolResultImages: boolean; + supportedInputMediaTypes: string[] | undefined; + supportsWebSearch: false; +} + +/** + * Resolve the capabilities a Connect Bedrock gateway serves a model with. + * + * Tool/image/web-search flags mirror the bridge's own live-discovery defaults + * for Anthropic-on-Bedrock models (bedrock-provider.ts); token limits, family, + * thinking levels, and media types come from the Anthropic-on-Bedrock table, + * which answers every Claude id. Accepts any model id, so callers never have + * to handle an unknown-model case (see the fallbacks above). For non-Claude + * ids the image flags go conservative — vision is not the norm across Bedrock + * chat models and over-reporting produces hard API errors — while tool use is, + * so `supportsTools` stays on. + */ +export function getConnectBedrockModelCapabilities( + modelId: string, +): ConnectBedrockModelCapabilities { + const claude = getAnthropicModelCapabilities(modelId); + const maxContextLength = claude?.maxContextLength ?? FALLBACK_CONTEXT_LENGTH; + const maxOutputTokens = claude?.maxOutputTokens ?? FALLBACK_MAX_OUTPUT_TOKENS; + return { + family: claude?.family, + thinkingEffortLevels: claude?.thinkingEffortLevels, + maxContextLength, + maxInputTokens: maxContextLength - maxOutputTokens, + maxOutputTokens, + supportsTools: true, + supportsImages: claude !== undefined, + supportsToolResultImages: claude !== undefined, + supportedInputMediaTypes: claude?.supportedInputMediaTypes, + supportsWebSearch: false, + }; +} diff --git a/packages/ai-config/src/schema.ts b/packages/ai-config/src/schema.ts index 6bb6fe2..83763bb 100644 --- a/packages/ai-config/src/schema.ts +++ b/packages/ai-config/src/schema.ts @@ -276,6 +276,7 @@ const BUILTIN_CONNECTION_SECTIONS = { databricks: ["databricks"], litellm: [], portkey: [], + connect: [], } as const satisfies Record; /** diff --git a/packages/ai-config/src/vocabulary.ts b/packages/ai-config/src/vocabulary.ts index 2890441..0082600 100644 --- a/packages/ai-config/src/vocabulary.ts +++ b/packages/ai-config/src/vocabulary.ts @@ -36,6 +36,7 @@ export const BUILTIN_PROVIDER_IDS = [ "databricks", "litellm", "portkey", + "connect", ] as const; export type BuiltinProviderId = (typeof BUILTIN_PROVIDER_IDS)[number]; @@ -93,6 +94,7 @@ export const CLIENT_KIND_VALUES = [ "databricks", "litellm", "portkey", + "connect", ] as const; export type ClientKind = (typeof CLIENT_KIND_VALUES)[number]; @@ -111,8 +113,8 @@ export type ClientKind = (typeof CLIENT_KIND_VALUES)[number]; * **equal**. * * This is a strict subset of {@link CLIENT_KIND_VALUES}. Product-bound kinds - * (`positai`, `copilot`, `databricks`) remain excluded because their auth flows - * are not generic custom-provider flows. + * (`positai`, `copilot`, `databricks`, `connect`) remain excluded because + * their auth flows are not generic custom-provider flows. */ export const SUPPORTED_CUSTOM_CLIENT_KIND_VALUES = [ "openai-compatible", diff --git a/packages/ai-credentials/src/store-backend/providerEnvMappings.ts b/packages/ai-credentials/src/store-backend/providerEnvMappings.ts index 9f5d578..475bda8 100644 --- a/packages/ai-credentials/src/store-backend/providerEnvMappings.ts +++ b/packages/ai-credentials/src/store-backend/providerEnvMappings.ts @@ -86,4 +86,9 @@ export const PROVIDER_ENV_MAPPINGS: Record = { portkey: { apiKey: "PORTKEY_API_KEY", }, + // The standard Posit Connect API-key variable (rsconnect/connectapi + // convention); pairs with ai-config's POSIT_CONNECT_URL connection var. + connect: { + apiKey: "CONNECT_API_KEY", + }, }; diff --git a/packages/ai-provider-bridge/src/model-clients/BedrockClient.ts b/packages/ai-provider-bridge/src/model-clients/BedrockClient.ts index a916e65..620a4e9 100644 --- a/packages/ai-provider-bridge/src/model-clients/BedrockClient.ts +++ b/packages/ai-provider-bridge/src/model-clients/BedrockClient.ts @@ -17,6 +17,7 @@ import type { LanguageModelV3 } from "@ai-sdk/provider"; import { streamText } from "ai"; import { createAwsCredentialProvider } from "../aws-credentials"; +import { safeSdkCustomHeaders } from "../custom-headers"; import { resolveBedrockTransport, type BedrockTransport } from "../providers/bedrock-transport"; import { sanitizeToolCallIdsForAnthropic } from "../tool-call-ids"; import { @@ -54,6 +55,18 @@ export interface BedrockClientConfig { accessKeyId?: string; secretAccessKey?: string; sessionToken?: string; + /** Extra request headers (e.g. for a header-gated proxy in front of a gateway). */ + customHeaders?: Record; + /** + * Allow an explicit `baseUrl` (see {@link ModelClientChatParams.baseUrl}) to + * override the resolved FIPS runtime endpoint. Defaults to false: an + * unrecognized override under FIPS is rejected rather than silently sent to + * an endpoint with no FIPS guarantee. Callers that route through a trusted, + * admin-configured gateway (e.g. the Posit Connect provider) opt in + * explicitly, since that redirect is deliberate rather than accidental + * misconfiguration. + */ + allowBaseUrlUnderFips?: boolean; } export class BedrockClient implements ModelClient { @@ -216,9 +229,17 @@ export class BedrockClient implements ModelClient { * - Anthropic models use `createBedrockAnthropic` (native Anthropic InvokeModel API * through Bedrock) for full feature parity including prompt caching via * `providerOptions.anthropic.cacheControl`. - * - OpenAI protocols use Bedrock Mantle. Only these routes honor `baseUrl`. + * - OpenAI protocols use Bedrock Mantle. * - All other models use `createAmazonBedrock` (Converse API). * + * The Anthropic and Converse routes honor an explicit `baseUrl` (e.g. a + * Connect gateway route), falling back to the resolved AWS runtime endpoint + * otherwise. When FIPS endpoints are mandated, an explicit `baseUrl` is + * rejected unless {@link BedrockClientConfig.allowBaseUrlUnderFips} opts in + * (then the override is logged). Mantle also honors an explicit `baseUrl` + * but has no runtime-endpoint fallback and is vetoed entirely under FIPS, + * with no override. + * * When an explicit `protocol` is provided, it takes precedence over the * model-ID heuristic. */ @@ -229,6 +250,7 @@ export class BedrockClient implements ModelClient { baseUrl?: string, ): LanguageModelV3 { const credentialProvider = createAwsCredentialProvider(this.config); + const headers = safeSdkCustomHeaders(this.config.customHeaders); if (protocol === "openai-chat" || protocol === "openai-responses") { if (!transport.mantleEnabled) { @@ -239,6 +261,7 @@ export class BedrockClient implements ModelClient { const mantle = createBedrockMantle({ region: this.config.region, baseURL: baseUrl, + headers, credentialProvider, // Enforce the AWS-credentials-only contract. Without this explicit // opt-out, a stale AWS_BEARER_TOKEN_BEDROCK overrides SigV4. @@ -247,6 +270,18 @@ export class BedrockClient implements ModelClient { return protocol === "openai-chat" ? mantle.chat(modelId) : mantle.responses(modelId); } + if (baseUrl && transport.useFipsEndpoint) { + if (!this.config.allowBaseUrlUnderFips) { + throw new Error( + `Bedrock base URL override to ${baseUrl} is not permitted while AWS FIPS endpoints ` + + `are enforced (resolved FIPS endpoint: ${transport.runtimeBaseUrl}).`, + ); + } + this.logger?.warn( + `[Bedrock] Explicit base URL ${baseUrl} overrides the FIPS runtime endpoint ${transport.runtimeBaseUrl}`, + ); + } + const useAnthropicApi = protocol ? protocol === "anthropic-messages" : isAnthropicModel(modelId); @@ -254,14 +289,16 @@ export class BedrockClient implements ModelClient { if (useAnthropicApi) { return createBedrockAnthropic({ region: this.config.region, - baseURL: transport.runtimeBaseUrl, + baseURL: baseUrl ?? transport.runtimeBaseUrl, + headers, credentialProvider, })(modelId); } return createAmazonBedrock({ region: this.config.region, - baseURL: transport.runtimeBaseUrl, + baseURL: baseUrl ?? transport.runtimeBaseUrl, + headers, credentialProvider, })(modelId); } diff --git a/packages/ai-provider-bridge/src/model-clients/__tests__/bedrock-mantle-routing.test.ts b/packages/ai-provider-bridge/src/model-clients/__tests__/bedrock-mantle-routing.test.ts index 7133de3..2ea3caf 100644 --- a/packages/ai-provider-bridge/src/model-clients/__tests__/bedrock-mantle-routing.test.ts +++ b/packages/ai-provider-bridge/src/model-clients/__tests__/bedrock-mantle-routing.test.ts @@ -141,6 +141,35 @@ describe("Bedrock Mantle protocol routing", () => { expect(createBedrockMantle).not.toHaveBeenCalled(); }); + it("honors an explicit baseUrl on the Converse and Anthropic routes", async () => { + await client.chat( + params({ + model: "amazon.nova-pro", + protocol: "bedrock-converse", + baseUrl: "https://connect.example.com/__gateway__/bedrock/guid", + }), + ); + expect(createAmazonBedrock).toHaveBeenCalledWith( + expect.objectContaining({ + baseURL: "https://connect.example.com/__gateway__/bedrock/guid", + }), + ); + + vi.clearAllMocks(); + await client.chat( + params({ + model: "anthropic.claude-sonnet-4-6", + protocol: "anthropic-messages", + baseUrl: "https://connect.example.com/__gateway__/bedrock/guid", + }), + ); + expect(createBedrockAnthropic).toHaveBeenCalledWith( + expect.objectContaining({ + baseURL: "https://connect.example.com/__gateway__/bedrock/guid", + }), + ); + }); + it("routes both runtime factories through the resolved FIPS host", async () => { resolveBedrockTransport.mockResolvedValue({ useFipsEndpoint: true, @@ -170,6 +199,113 @@ describe("Bedrock Mantle protocol routing", () => { ); }); + it("rejects an explicit baseUrl override under FIPS by default", async () => { + resolveBedrockTransport.mockResolvedValue({ + useFipsEndpoint: true, + runtimeBaseUrl: "https://bedrock-runtime-fips.us-gov-west-1.amazonaws.com", + mantleEnabled: false, + }); + const fipsClient = new BedrockClient({ + region: "us-gov-west-1", + accessKeyId: "key", + secretAccessKey: "secret", + }); + + await expect( + fipsClient.chat( + params({ + model: "amazon.nova-pro", + protocol: "bedrock-converse", + baseUrl: "https://connect.example.com/__gateway__/bedrock/guid", + }), + ), + ).rejects.toThrow(/not permitted while AWS FIPS endpoints are enforced/); + expect(createAmazonBedrock).not.toHaveBeenCalled(); + }); + + it("warns but proceeds when an explicit baseUrl overrides FIPS with allowBaseUrlUnderFips", async () => { + resolveBedrockTransport.mockResolvedValue({ + useFipsEndpoint: true, + runtimeBaseUrl: "https://bedrock-runtime-fips.us-gov-west-1.amazonaws.com", + mantleEnabled: false, + }); + const logger = { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + trace: vi.fn(), + }; + const fipsClient = new BedrockClient( + { + region: "us-gov-west-1", + accessKeyId: "key", + secretAccessKey: "secret", + allowBaseUrlUnderFips: true, + }, + logger, + ); + + await fipsClient.chat( + params({ + model: "amazon.nova-pro", + protocol: "bedrock-converse", + baseUrl: "https://connect.example.com/__gateway__/bedrock/guid", + }), + ); + + expect(createAmazonBedrock).toHaveBeenCalledWith( + expect.objectContaining({ + baseURL: "https://connect.example.com/__gateway__/bedrock/guid", + }), + ); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining("overrides the FIPS runtime endpoint"), + ); + }); + + it("forwards customHeaders to every factory route", async () => { + const customHeaders = { "x-proxy-token": "t" }; + const headeredClient = new BedrockClient({ + region: "us-east-2", + accessKeyId: "key", + secretAccessKey: "secret", + customHeaders, + }); + + await headeredClient.chat(params({ model: "amazon.nova-pro", protocol: "bedrock-converse" })); + expect(createAmazonBedrock).toHaveBeenCalledWith( + expect.objectContaining({ headers: customHeaders }), + ); + + await headeredClient.chat( + params({ model: "anthropic.claude-sonnet-4-6", protocol: "anthropic-messages" }), + ); + expect(createBedrockAnthropic).toHaveBeenCalledWith( + expect.objectContaining({ headers: customHeaders }), + ); + + await headeredClient.chat(params({ model: "openai.gpt-5.5", protocol: "openai-chat" })); + expect(createBedrockMantle).toHaveBeenCalledWith( + expect.objectContaining({ headers: customHeaders }), + ); + }); + + it("strips SDK-managed header names out of customHeaders before they reach the factory", async () => { + const headeredClient = new BedrockClient({ + region: "us-east-2", + accessKeyId: "key", + secretAccessKey: "secret", + customHeaders: { Authorization: "should-not-leak", "x-proxy-token": "t" }, + }); + + await headeredClient.chat(params({ model: "amazon.nova-pro", protocol: "bedrock-converse" })); + + expect(createAmazonBedrock).toHaveBeenCalledWith( + expect.objectContaining({ headers: { "x-proxy-token": "t" } }), + ); + }); + it("rejects Mantle protocols when FIPS endpoints are enabled", async () => { resolveBedrockTransport.mockResolvedValue({ useFipsEndpoint: true, diff --git a/packages/ai-provider-bridge/src/providers.ts b/packages/ai-provider-bridge/src/providers.ts index 221fed0..1ce5b08 100644 --- a/packages/ai-provider-bridge/src/providers.ts +++ b/packages/ai-provider-bridge/src/providers.ts @@ -13,6 +13,12 @@ export { registerAnthropicProvider, registerCustomAnthropicProvider, } from "./providers/anthropic-provider"; +export { registerConnectProvider, shapeConnectIntegrations } from "./providers/connect-provider"; +export type { + ConnectAwsCredentialResult, + ConnectIntegration, + ConnectProviderCallbacks, +} from "./providers/connect-provider"; export { registerCopilotProvider } from "./providers/copilot-provider"; export { registerDatabricksProvider } from "./providers/databricks-provider"; export { diff --git a/packages/ai-provider-bridge/src/providers/__tests__/cached-model-fetcher.test.ts b/packages/ai-provider-bridge/src/providers/__tests__/cached-model-fetcher.test.ts index 6daab1c..7569f55 100644 --- a/packages/ai-provider-bridge/src/providers/__tests__/cached-model-fetcher.test.ts +++ b/packages/ai-provider-bridge/src/providers/__tests__/cached-model-fetcher.test.ts @@ -76,6 +76,80 @@ describe("createCachedModelFetcher — fetchFresh variant", () => { }); }); +describe("createCachedModelFetcher — credential partitions", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + function keyedFetcher(overrides?: { + ttl?: number; + maxCacheEntries?: number; + fetchFresh?: (credentials: ApiKeyCredentials) => Promise; + }) { + const fetchFresh = + overrides?.fetchFresh ?? + vi.fn(async (creds: ApiKeyCredentials) => [model(`fresh-${creds.apiKey}`)]); + const fetcher = createCachedModelFetcher({ + providerId: "test-provider", + hasCredentials: (creds) => Boolean(creds.apiKey), + cacheKey: async (creds) => `fingerprint-${creds.apiKey}`, + fetchFresh, + fallbackModels: [model("fallback")], + ttl: overrides?.ttl ?? 60_000, + maxCacheEntries: overrides?.maxCacheEntries ?? 32, + logger, + }); + return { fetcher, fetchFresh }; + } + + it("partitions model lists by an asynchronous credential fingerprint", async () => { + const { fetcher, fetchFresh } = keyedFetcher(); + const first = { ...credentials, apiKey: "first" }; + const second = { ...credentials, apiKey: "second" }; + + expect((await fetcher(first)).map((entry) => entry.id)).toEqual(["fresh-first"]); + expect((await fetcher(second)).map((entry) => entry.id)).toEqual(["fresh-second"]); + expect((await fetcher(first)).map((entry) => entry.id)).toEqual(["fresh-first"]); + expect(fetchFresh).toHaveBeenCalledTimes(2); + }); + + it("keeps stale fallback inside its credential partition", async () => { + const failing = new Set(); + const { fetcher } = keyedFetcher({ + ttl: 0, + fetchFresh: async (creds) => { + if (failing.has(creds.apiKey)) throw new Error("unavailable"); + return [model(`fresh-${creds.apiKey}`)]; + }, + }); + const first = { ...credentials, apiKey: "first" }; + const second = { ...credentials, apiKey: "second" }; + + await fetcher(first); + await fetcher(second); + failing.add("first"); + + expect((await fetcher(first)).map((entry) => entry.id)).toEqual(["fresh-first"]); + }); + + it("clears every partition and evicts the oldest entry at the configured bound", async () => { + const { fetcher, fetchFresh } = keyedFetcher({ maxCacheEntries: 2 }); + const first = { ...credentials, apiKey: "first" }; + const second = { ...credentials, apiKey: "second" }; + const third = { ...credentials, apiKey: "third" }; + + await fetcher(first); + await fetcher(second); + await fetcher(third); + await fetcher(first); + expect(fetchFresh).toHaveBeenCalledTimes(4); + + fetcher.clearCache?.(); + await fetcher(second); + expect(fetchFresh).toHaveBeenCalledTimes(5); + }); +}); + describe("createCachedModelFetcher — discovery deadline", () => { afterEach(() => { vi.useRealTimers(); diff --git a/packages/ai-provider-bridge/src/providers/__tests__/connect-provider.test.ts b/packages/ai-provider-bridge/src/providers/__tests__/connect-provider.test.ts new file mode 100644 index 0000000..e16151b --- /dev/null +++ b/packages/ai-provider-bridge/src/providers/__tests__/connect-provider.test.ts @@ -0,0 +1,714 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import { CONNECT_BEDROCK_MODEL_IDS, CONNECT_BEDROCK_MODELS } from "ai-config"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +const chats = vi.hoisted(() => ({ + anthropic: vi.fn(async () => (async function* () {})()), + bedrock: vi.fn(async () => (async function* () {})()), +})); + +vi.mock("../../model-clients/AnthropicClient", () => ({ + AnthropicClient: vi.fn(function () { + return { chat: chats.anthropic }; + }), +})); +vi.mock("../../model-clients/BedrockClient", () => ({ + BedrockClient: vi.fn(function () { + return { chat: chats.bedrock }; + }), +})); + +import { AnthropicClient } from "../../model-clients/AnthropicClient"; +import { BedrockClient } from "../../model-clients/BedrockClient"; +import type { CancellationToken, Logger, ProviderCredentials } from "../../types"; +import { + registerConnectProvider, + shapeConnectIntegrations, + type ConnectProviderCallbacks, +} from "../connect-provider"; +import { ProviderRegistry } from "../ProviderRegistry"; + +const logger: Logger = { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + trace: vi.fn(), +}; + +const CONNECT_URL = "https://connect.example.com"; +const ANTHROPIC_GUID = "aaaa1111-1111-1111-1111-111111111111"; +const AWS_GUID = "bbbb2222-2222-2222-2222-222222222222"; +const ANTHROPIC_GATEWAY = `${CONNECT_URL}/__gateway__/anthropic/${ANTHROPIC_GUID}/v1`; +const BEDROCK_GATEWAY = `${CONNECT_URL}/__gateway__/bedrock/${AWS_GUID}`; +// Minted prefixes always embed the full guid; see mintIntegrationPrefix. +const ANTHROPIC_PREFIX = `connect-anthropic-prod-${ANTHROPIC_GUID}`; +const AWS_PREFIX = `connect-bedrock-team-${AWS_GUID}`; + +const INTEGRATION_RECORDS = [ + { + guid: ANTHROPIC_GUID, + template: "anthropic", + name: "Anthropic Prod", + description: "Team key", + }, + { + guid: AWS_GUID, + template: "aws", + name: "Bedrock Team", + description: null, + auth_type: "Viewer", + config: { sts_region: "us-west-2" }, + }, + // Service Account auth mints against content items, not the calling user. + { + guid: "cccc3333-3333-3333-3333-333333333333", + template: "aws", + name: "Service Bedrock", + auth_type: "Service Account", + config: {}, + }, + { + guid: "dddd4444-4444-4444-4444-444444444444", + template: "github", + name: "GitHub", + }, +]; + +const ANTHROPIC_MODELS_BODY = { + data: [{ id: "claude-sonnet-4-5-20250929", display_name: "Claude Sonnet 4.5" }], +}; + +const credentials = { type: "apikey", apiKey: "tok", baseUrl: CONNECT_URL } as const; + +const cancellationToken: CancellationToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose() {} }), +}; + +function json(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { status }); +} + +/** Route the two discovery endpoints; anything else is a test bug. */ +function stubDiscoveryFetch(routes: Record Response> = {}): ReturnType { + const fetchMock = vi.fn(async (url: string) => { + const route = { + [`${CONNECT_URL}/__api__/v1/oauth/integrations`]: () => json(INTEGRATION_RECORDS), + [`${ANTHROPIC_GATEWAY}/models`]: () => json(ANTHROPIC_MODELS_BODY), + ...routes, + }[url]; + if (!route) throw new Error(`Unexpected fetch: ${url}`); + return route(); + }); + vi.stubGlobal("fetch", fetchMock); + return fetchMock; +} + +function mintSuccess(region = "us-west-2") { + return vi.fn(async () => ({ + ok: true as const, + credentials: { + type: "aws-credentials" as const, + region, + accessKeyId: "AKIA", + secretAccessKey: "SECRET", + sessionToken: "SESSION", + }, + })); +} + +beforeEach(() => { + vi.clearAllMocks(); +}); + +afterEach(() => { + vi.unstubAllGlobals(); +}); + +describe("shapeConnectIntegrations", () => { + it("keeps supported templates even from a verbatim allowlist, requires Viewer auth for aws, and drops guid-less records", () => { + // "github" is in the allowlist but has no shaping rule, so it is still dropped. + const shaped = shapeConnectIntegrations( + [...INTEGRATION_RECORDS, { template: "anthropic", name: "No Guid" }], + CONNECT_URL, + ["anthropic", "aws", "github"], + ); + + expect(shaped.map((integration) => integration.idPrefix)).toEqual([ + ANTHROPIC_PREFIX, + AWS_PREFIX, + ]); + expect(shaped[0]).toMatchObject({ + guid: ANTHROPIC_GUID, + baseUrl: ANTHROPIC_GATEWAY, + loginUrl: `${CONNECT_URL}/__oauth__/integrations/${ANTHROPIC_GUID}/login`, + }); + expect(shaped[0].region).toBeUndefined(); + expect(shaped[1]).toMatchObject({ + guid: AWS_GUID, + baseUrl: BEDROCK_GATEWAY, + region: "us-west-2", + }); + }); + + it("mints prefixes from name, then description, then template, always suffixed with the guid", () => { + const shaped = shapeConnectIntegrations( + [ + { guid: "guid-one", template: "anthropic", name: "Anthropic" }, + { guid: "guid-two", template: "anthropic", name: "Anthropic" }, + { guid: "guid-three", template: "anthropic", name: "", description: "Fallback Desc" }, + { guid: "guid-four", template: "anthropic", name: "", description: "" }, + ], + CONNECT_URL, + ["anthropic"], + ); + + // Two integrations sharing a name still mint distinct prefixes — the + // guid, not the slug, is what makes a prefix unique. + expect(shaped.map((integration) => integration.idPrefix)).toEqual([ + "connect-anthropic-guid-one", + "connect-anthropic-guid-two", + "connect-fallback-desc-guid-three", + "connect-anthropic-guid-four", + ]); + }); +}); + +describe("connect model fetcher", () => { + function registryWithProvider(callbacks?: ConnectProviderCallbacks): ProviderRegistry { + const registry = new ProviderRegistry(logger); + registerConnectProvider(registry, logger, callbacks); + return registry; + } + + it("returns no models when the Connect server URL is missing", async () => { + const fetchMock = stubDiscoveryFetch(); + const models = await registryWithProvider().getModelsForProvider("connect", { + type: "apikey", + apiKey: "tok", + }); + + expect(models).toEqual([]); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("returns no models for the wrong credential type", async () => { + stubDiscoveryFetch(); + const models = await registryWithProvider().getModelsForProvider("connect", { + type: "oauth", + accessToken: "tok", + }); + + expect(models).toEqual([]); + }); + + it("discovers integrations and namespaces each gateway's models", async () => { + const fetchMock = stubDiscoveryFetch(); + // Trailing slash on the configured URL must not produce double-slash requests. + const models = await registryWithProvider({ + getAwsCredentials: vi.fn(), + }).getModelsForProvider("connect", { + ...credentials, + baseUrl: `${CONNECT_URL}/`, + }); + + expect(fetchMock).toHaveBeenCalledWith( + `${CONNECT_URL}/__api__/v1/oauth/integrations`, + expect.objectContaining({ headers: { Authorization: "Key tok" } }), + ); + expect(fetchMock).toHaveBeenCalledWith( + `${ANTHROPIC_GATEWAY}/models`, + expect.objectContaining({ + headers: { "x-api-key": "tok", "anthropic-version": "2023-06-01" }, + }), + ); + + const anthropicModel = models.find( + (model) => model.id === `${ANTHROPIC_PREFIX}/claude-sonnet-4-5-20250929`, + ); + expect(anthropicModel).toMatchObject({ + name: "Claude Sonnet 4.5 (Anthropic Prod)", + providerId: "connect", + vendor: "anthropic", + protocol: "anthropic-messages", + baseUrl: ANTHROPIC_GATEWAY, + supportsWebSearch: true, + }); + + // Every declared Connect Bedrock model is a recognized Claude id today, + // so it declares anthropic-messages — the route BedrockClient actually + // takes for it — even though it is served through the bedrock gateway. + const bedrockModels = models.filter((model) => model.id.startsWith(`${AWS_PREFIX}/`)); + expect(bedrockModels.map((model) => model.id)).toEqual( + CONNECT_BEDROCK_MODEL_IDS.map((id) => `${AWS_PREFIX}/${id}`), + ); + // Display names come from the declared table's human-readable name plus + // the integration label, never the raw wire id. + const [firstDeclared] = CONNECT_BEDROCK_MODELS; + expect(models.find((model) => model.id === `${AWS_PREFIX}/${firstDeclared.id}`)?.name).toBe( + `${firstDeclared.name} (Bedrock Team)`, + ); + for (const model of bedrockModels) { + expect(model).toMatchObject({ + providerId: "connect", + baseUrl: BEDROCK_GATEWAY, + supportsWebSearch: false, + protocol: "anthropic-messages", + }); + } + + expect(models).toHaveLength(1 + CONNECT_BEDROCK_MODEL_IDS.length); + }); + + it("skips AWS-backed integrations and warns when no credential callback is provided", async () => { + stubDiscoveryFetch(); + const models = await registryWithProvider().getModelsForProvider("connect", credentials); + + expect(models.map((model) => model.id)).toEqual([ + `${ANTHROPIC_PREFIX}/claude-sonnet-4-5-20250929`, + ]); + expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining("no AWS credential callback")); + }); + + it("isolates a failing integration's discovery from the others", async () => { + stubDiscoveryFetch({ + [`${ANTHROPIC_GATEWAY}/models`]: () => json({ error: "boom" }, 500), + }); + const models = await registryWithProvider({ + getAwsCredentials: vi.fn(), + }).getModelsForProvider("connect", credentials); + + expect(models.map((model) => model.id)).toEqual( + CONNECT_BEDROCK_MODEL_IDS.map((id) => `${AWS_PREFIX}/${id}`), + ); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining('discovery failed for integration "Anthropic Prod"'), + ); + }); + + it("intersects the host template allowlist with the supported set", async () => { + stubDiscoveryFetch(); + const callbacks: ConnectProviderCallbacks = { + getAwsCredentials: vi.fn(), + templates: () => ["anthropic", "github"], + }; + const models = await registryWithProvider(callbacks).getModelsForProvider( + "connect", + credentials, + ); + + expect(models.every((model) => model.protocol === "anthropic-messages")).toBe(true); + expect(models).toHaveLength(1); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining("unsupported integration template(s): github"), + ); + }); +}); + +describe("connect chat routing", () => { + function chatParams(model: string, protocol?: "anthropic-messages" | "bedrock-converse") { + return { + model, + messages: [], + cancellationToken, + protocol, + baseUrl: protocol === "bedrock-converse" ? BEDROCK_GATEWAY : ANTHROPIC_GATEWAY, + }; + } + + function clientWithoutDiscovery( + callbacks?: ConnectProviderCallbacks, + creds: ProviderCredentials = credentials, + ) { + const registry = new ProviderRegistry(logger); + registerConnectProvider(registry, logger, callbacks); + const client = registry.getClientForProvider("connect", creds); + expect(client).not.toBeNull(); + return client!; + } + + async function clientAfterDiscovery( + callbacks?: ConnectProviderCallbacks, + creds: ProviderCredentials = credentials, + ) { + stubDiscoveryFetch(); + const registry = new ProviderRegistry(logger); + registerConnectProvider(registry, logger, callbacks); + await registry.getModelsForProvider("connect", creds); + const client = registry.getClientForProvider("connect", creds); + expect(client).not.toBeNull(); + return client!; + } + + it("routes anthropic-messages to an AnthropicClient spending the token at the gateway", async () => { + const client = await clientAfterDiscovery(); + + await client.chat( + chatParams(`${ANTHROPIC_PREFIX}/claude-sonnet-4-5-20250929`, "anthropic-messages"), + ); + + expect(AnthropicClient).toHaveBeenCalledWith( + { apiKey: "tok" }, + ANTHROPIC_GATEWAY, + undefined, + logger, + ); + expect(chats.anthropic).toHaveBeenCalledWith( + expect.objectContaining({ + model: "claude-sonnet-4-5-20250929", + baseUrl: ANTHROPIC_GATEWAY, + }), + ); + }); + + it("routes bedrock-converse through per-request STS credentials without forcing a protocol", async () => { + const getAwsCredentials = mintSuccess(); + const client = await clientAfterDiscovery({ getAwsCredentials }); + const modelId = CONNECT_BEDROCK_MODEL_IDS[0]; + + await client.chat(chatParams(`${AWS_PREFIX}/${modelId}`, "bedrock-converse")); + + expect(getAwsCredentials).toHaveBeenCalledWith( + expect.objectContaining({ guid: AWS_GUID, region: "us-west-2" }), + expect.any(AbortSignal), + ); + expect(BedrockClient).toHaveBeenCalledWith( + expect.objectContaining({ + region: "us-west-2", + accessKeyId: "AKIA", + secretAccessKey: "SECRET", + sessionToken: "SESSION", + // Connect's gateway redirect is a deliberate, trusted override, so + // it must reach BedrockClient even when FIPS endpoints are enforced. + allowBaseUrlUnderFips: true, + }), + logger, + ); + // The model-id heuristic must keep us.anthropic.* ids on the native + // Anthropic route, so no explicit protocol may be forwarded. + expect(chats.bedrock).toHaveBeenCalledTimes(1); + const delegated = chats.bedrock.mock.calls[0][0] as { + model: string; + protocol?: string; + baseUrl?: string; + }; + expect(delegated.model).toBe(modelId); + expect(delegated.baseUrl).toBe(BEDROCK_GATEWAY); + expect(delegated.protocol).toBeUndefined(); + }); + + it("routes stamped models statelessly when the integration cache is empty", async () => { + const getAwsCredentials = mintSuccess(); + const client = clientWithoutDiscovery({ getAwsCredentials }); + const modelId = CONNECT_BEDROCK_MODEL_IDS[0]; + + await client.chat(chatParams(`${AWS_PREFIX}/${modelId}`, "bedrock-converse")); + + // No discovery ran, so the integration is synthesized from the stamp. + expect(getAwsCredentials).toHaveBeenCalledWith( + expect.objectContaining({ guid: AWS_GUID, template: "aws" }), + expect.any(AbortSignal), + ); + expect(BedrockClient).toHaveBeenCalledWith( + expect.objectContaining({ region: "us-west-2" }), + logger, + ); + const delegated = chats.bedrock.mock.calls[0][0] as { model: string; baseUrl?: string }; + expect(delegated.model).toBe(modelId); + expect(delegated.baseUrl).toBe(BEDROCK_GATEWAY); + + await client.chat( + chatParams(`${ANTHROPIC_PREFIX}/claude-sonnet-4-5-20250929`, "anthropic-messages"), + ); + expect(chats.anthropic).toHaveBeenCalledWith( + expect.objectContaining({ baseUrl: ANTHROPIC_GATEWAY }), + ); + }); + + it("rejects a stamped model discovered against a different Connect server", async () => { + const client = clientWithoutDiscovery(); + + await expect( + client.chat({ + model: `${ANTHROPIC_PREFIX}/claude-sonnet-4-5-20250929`, + messages: [], + cancellationToken, + baseUrl: `https://other.example.com/__gateway__/anthropic/${ANTHROPIC_GUID}/v1`, + }), + ).rejects.toThrow(/other\.example\.com.*Refresh the model list/s); + expect(chats.anthropic).not.toHaveBeenCalled(); + }); + + it("rejects a baseUrl that is not a Connect gateway route", async () => { + const client = clientWithoutDiscovery(); + + await expect( + client.chat({ + model: `${ANTHROPIC_PREFIX}/claude-sonnet-4-5-20250929`, + messages: [], + cancellationToken, + baseUrl: `${CONNECT_URL}/some/other/api`, + }), + ).rejects.toThrow(/not a Connect gateway URL/); + }); + + it("resolves an unstamped override model through the discovery cache", async () => { + const client = await clientAfterDiscovery(); + + await client.chat({ + model: `${ANTHROPIC_PREFIX}/claude-3-haiku-20240307`, + messages: [], + cancellationToken, + }); + + expect(chats.anthropic).toHaveBeenCalledWith( + expect.objectContaining({ + model: "claude-3-haiku-20240307", + baseUrl: ANTHROPIC_GATEWAY, + }), + ); + }); + + it("resolves an unstamped model through the cache when ai-config falls back to the bare Connect root", async () => { + // ai-config's resolver stamps a model with the provider's own baseUrl + // when the model carries none of its own — the bare Connect server root, + // not a real gateway URL. That fallback must be treated like "no + // baseUrl" and fall through to the cache, not rejected as an unknown + // gateway route. + const client = await clientAfterDiscovery(); + + await client.chat({ + model: `${ANTHROPIC_PREFIX}/claude-3-haiku-20240307`, + messages: [], + cancellationToken, + baseUrl: CONNECT_URL, + }); + + expect(chats.anthropic).toHaveBeenCalledWith( + expect.objectContaining({ + model: "claude-3-haiku-20240307", + baseUrl: ANTHROPIC_GATEWAY, + }), + ); + }); + + it("asks for a model-list refresh when an unstamped model's prefix is unknown", async () => { + const client = await clientAfterDiscovery(); + + await expect( + client.chat({ + model: "connect-nonexistent/claude-sonnet-4-5-20250929", + messages: [], + cancellationToken, + }), + ).rejects.toThrow(/Refresh the model list/); + }); + + it("never resolves an unstamped model against a cache populated by different credentials", async () => { + // Discovery under one session's token must not leak its integrations to + // a second session that never discovered against this provider — even + // against the same Connect server. + stubDiscoveryFetch(); + const registry = new ProviderRegistry(logger); + registerConnectProvider(registry, logger); + await registry.getModelsForProvider("connect", credentials); + + const otherSession = { ...credentials, apiKey: "tok-other-user" }; + const client = registry.getClientForProvider("connect", otherSession)!; + + await expect( + client.chat({ + model: `${ANTHROPIC_PREFIX}/claude-3-haiku-20240307`, + messages: [], + cancellationToken, + }), + ).rejects.toThrow(/Refresh the model list/); + expect(chats.anthropic).not.toHaveBeenCalled(); + }); + + it("keeps unstamped routing available for interleaved credential sessions", async () => { + stubDiscoveryFetch(); + const registry = new ProviderRegistry(logger); + registerConnectProvider(registry, logger); + const firstSession = credentials; + const secondSession = { ...credentials, apiKey: "tok-other-user" }; + + await registry.getModelsForProvider("connect", firstSession); + await registry.getModelsForProvider("connect", secondSession); + // This is a model-cache hit for the first session. Routing state must + // remain available without forcing another network discovery. + await registry.getModelsForProvider("connect", firstSession); + + const firstClient = registry.getClientForProvider("connect", firstSession)!; + await firstClient.chat({ + model: `${ANTHROPIC_PREFIX}/claude-3-haiku-20240307`, + messages: [], + cancellationToken, + }); + + expect(chats.anthropic).toHaveBeenCalledWith( + expect.objectContaining({ + model: "claude-3-haiku-20240307", + baseUrl: ANTHROPIC_GATEWAY, + }), + ); + }); + + it("drops integrations deleted on the server at the next discovery", async () => { + stubDiscoveryFetch(); + const registry = new ProviderRegistry(logger); + registerConnectProvider(registry, logger); + await registry.getModelsForProvider("connect", credentials); + const client = registry.getClientForProvider("connect", credentials)!; + + await client.chat({ + model: `${ANTHROPIC_PREFIX}/claude-3-haiku-20240307`, + messages: [], + cancellationToken, + }); + expect(chats.anthropic).toHaveBeenCalledTimes(1); + + stubDiscoveryFetch({ + [`${CONNECT_URL}/__api__/v1/oauth/integrations`]: () => json([]), + }); + registry.clearModelCache("connect"); + await registry.getModelsForProvider("connect", credentials); + + await expect( + client.chat({ + model: `${ANTHROPIC_PREFIX}/claude-3-haiku-20240307`, + messages: [], + cancellationToken, + }), + ).rejects.toThrow(/Refresh the model list/); + }); + + it("forwards an anthropic-messages override on an AWS-backed integration to Bedrock", async () => { + const getAwsCredentials = mintSuccess(); + const client = await clientAfterDiscovery({ getAwsCredentials }); + const modelId = CONNECT_BEDROCK_MODEL_IDS[0]; + + // The template selects the transport: the gateway still requires SigV4, + // so the override picks the wire format inside BedrockClient instead of + // re-routing to the token-spending Anthropic path. + await client.chat({ + model: `${AWS_PREFIX}/${modelId}`, + messages: [], + cancellationToken, + protocol: "anthropic-messages", + baseUrl: BEDROCK_GATEWAY, + }); + + expect(chats.anthropic).not.toHaveBeenCalled(); + expect(chats.bedrock).toHaveBeenCalledWith( + expect.objectContaining({ protocol: "anthropic-messages", baseUrl: BEDROCK_GATEWAY }), + ); + }); + + it("surfaces the failure code and login URL when AWS credentials cannot be minted", async () => { + const loginUrl = `${CONNECT_URL}/__oauth__/integrations/${AWS_GUID}/login`; + const client = await clientAfterDiscovery({ + getAwsCredentials: vi.fn(async () => ({ + ok: false as const, + code: "oauth_session_required", + detail: "No active session.", + loginUrl, + })), + }); + + await expect( + client.chat(chatParams(`${AWS_PREFIX}/${CONNECT_BEDROCK_MODEL_IDS[0]}`, "bedrock-converse")), + ).rejects.toThrow(new RegExp(`oauth_session_required.*${loginUrl}`)); + expect(chats.bedrock).not.toHaveBeenCalled(); + }); + + it("throws rather than fall back to ambient AWS credentials when the mint is incomplete", async () => { + const client = await clientAfterDiscovery({ + getAwsCredentials: vi.fn(async () => ({ + ok: true as const, + credentials: { type: "aws-credentials" as const, region: "us-west-2" }, + })), + }); + + await expect( + client.chat(chatParams(`${AWS_PREFIX}/${CONNECT_BEDROCK_MODEL_IDS[0]}`, "bedrock-converse")), + ).rejects.toThrow(/incomplete AWS credentials/); + expect(chats.bedrock).not.toHaveBeenCalled(); + }); + + it("signs with the integration's sts_region over the minted region and warns", async () => { + const client = await clientAfterDiscovery({ getAwsCredentials: mintSuccess("us-east-1") }); + + await client.chat( + chatParams(`${AWS_PREFIX}/${CONNECT_BEDROCK_MODEL_IDS[0]}`, "bedrock-converse"), + ); + + expect(BedrockClient).toHaveBeenCalledWith( + expect.objectContaining({ region: "us-west-2" }), + logger, + ); + expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining("sts_region")); + }); + + it("passes provider customHeaders through to BedrockClient", async () => { + const customHeaders = { "x-proxy-token": "t" }; + const client = await clientAfterDiscovery( + { getAwsCredentials: mintSuccess() }, + { + ...credentials, + customHeaders, + }, + ); + + await client.chat( + chatParams(`${AWS_PREFIX}/${CONNECT_BEDROCK_MODEL_IDS[0]}`, "bedrock-converse"), + ); + + expect(BedrockClient).toHaveBeenCalledWith(expect.objectContaining({ customHeaders }), logger); + }); + + it("leaves ids whose first segment is not a connect- prefix unsplit", async () => { + const arn = + "arn:aws:bedrock:us-west-2:123456789012:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0"; + const client = await clientAfterDiscovery({ getAwsCredentials: mintSuccess() }); + + await client.chat({ + model: arn, + messages: [], + cancellationToken, + protocol: "bedrock-converse", + baseUrl: BEDROCK_GATEWAY, + }); + + const delegated = chats.bedrock.mock.calls[0][0] as { model: string }; + expect(delegated.model).toBe(arn); + }); + + it("rejects protocols the Connect gateways cannot serve", async () => { + const client = await clientAfterDiscovery({ getAwsCredentials: mintSuccess() }); + + await expect( + client.chat({ + model: `${ANTHROPIC_PREFIX}/claude-sonnet-4-5-20250929`, + messages: [], + cancellationToken, + protocol: "openai-chat", + baseUrl: ANTHROPIC_GATEWAY, + }), + ).rejects.toThrow(/cannot route protocol "openai-chat"/); + + await expect( + client.chat({ + model: `${AWS_PREFIX}/${CONNECT_BEDROCK_MODEL_IDS[0]}`, + messages: [], + cancellationToken, + protocol: "openai-chat", + baseUrl: BEDROCK_GATEWAY, + }), + ).rejects.toThrow(/cannot route protocol "openai-chat"/); + }); +}); diff --git a/packages/ai-provider-bridge/src/providers/cached-model-fetcher.ts b/packages/ai-provider-bridge/src/providers/cached-model-fetcher.ts index 7ac0068..9250a6a 100644 --- a/packages/ai-provider-bridge/src/providers/cached-model-fetcher.ts +++ b/packages/ai-provider-bridge/src/providers/cached-model-fetcher.ts @@ -38,6 +38,24 @@ interface CachedModelFetcherCommonConfig boolean; + /** + * Optional: partition the cache by credential identity rather than sharing + * one entry across every call. Providers whose credentials can legitimately + * change between calls to the same registered fetcher (e.g. Connect, where + * `baseUrl` names a different server per credential set) should return a + * stable opaque fingerprint of the parts of `T` that determine the fetched + * model list. Never return a credential secret itself. Omitted means every + * call shares one cache entry, matching prior behavior. + */ + cacheKey?: (credentials: T) => string | Promise; + + /** + * Maximum credential-partitioned entries retained by this fetcher. The + * oldest entry is evicted when the bound is exceeded. Defaults to 32 when + * `cacheKey` is supplied and 1 otherwise. + */ + maxCacheEntries?: number; + /** * Optional: Enrich models with additional data after initial fetch * Useful for providers that need multiple API calls per model (e.g., Ollama /api/show) @@ -158,15 +176,32 @@ export type ClearableModelFetcher = ((credentials: ProviderCredentials) => Promi * `clearCache()` may still answer its own caller but must not repopulate a * newer cache generation. */ +const DEFAULT_CACHE_KEY = "__default__"; +const DEFAULT_MAX_KEYED_CACHE_ENTRIES = 32; + +interface CacheEntry { + models: ModelInfo[]; + fetchedAt: number; +} + export function createCachedModelFetcher( config: CachedModelFetcherConfig, ): ClearableModelFetcher { const TTL = config.ttl ?? DEFAULT_TTL; const discoveryDeadlineMs = config.discoveryDeadlineMs ?? DEFAULT_DISCOVERY_DEADLINE_MS; - let lastFetch = 0; - let cachedModels: ModelInfo[] | null = null; + const maxCacheEntries = + config.maxCacheEntries ?? (config.cacheKey ? DEFAULT_MAX_KEYED_CACHE_ENTRIES : 1); + if (!Number.isInteger(maxCacheEntries) || maxCacheEntries < 1) { + throw new Error("maxCacheEntries must be a positive integer"); + } + // Keyed by `config.cacheKey` (default: one shared entry, matching prior + // behavior) so credentials that resolve to different backends never share + // a cached model list. Insertion order supplies bounded FIFO eviction; + // refreshing an existing key moves it to the back. + const cache = new Map(); // Generation guard: a fetch that spans clearCache() may still answer its // own caller, but must not repopulate the cache over a newer generation. + // Shared across keys — clearCache() invalidates the whole fetcher. let generation = 0; const fetcher: ClearableModelFetcher = async ( @@ -182,11 +217,14 @@ export function createCachedModelFetcher maxCacheEntries) { + const oldestKey = cache.keys().next().value; + if (oldestKey === undefined) break; + cache.delete(oldestKey); + } config.logger.info(`${logPrefix} Fetched ${freshModels.length} models from API`); } else { config.logger.debug( @@ -272,9 +315,9 @@ export function createCachedModelFetcher { generation += 1; - cachedModels = null; - lastFetch = 0; + cache.clear(); }; return fetcher; diff --git a/packages/ai-provider-bridge/src/providers/connect-provider.ts b/packages/ai-provider-bridge/src/providers/connect-provider.ts new file mode 100644 index 0000000..8fd5859 --- /dev/null +++ b/packages/ai-provider-bridge/src/providers/connect-provider.ts @@ -0,0 +1,772 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +/** + * Posit Connect provider. + * + * One built-in provider fronts every OAuth integration a Connect server + * offers. The provider credential is the session's federated token (an + * `apikey`), and `baseUrl` is the Connect server URL. Discovery spends the + * token on `GET {baseUrl}/__api__/v1/oauth/integrations`, keeps the + * integrations whose `template` is allowlisted, and namespaces each + * integration's models as `connect-/` so one flat model list + * can route back to the right gateway. + * + * Chat routing is stateless for discovery-stamped models: the stamped gateway + * URL embeds the Connect server, the template (which selects the transport), + * and the integration guid, so routing never depends on discovery-time state + * that a re-registration, restart, or re-discovery could invalidate. The + * integration cache exists only to resolve user-configured models that carry + * no stamp. + * + * Two templates are shaped today: + * - `anthropic` — a pass-through reverse proxy on + * `{connect}/__gateway__/anthropic/{guid}/v1` that swaps the federated + * token for the integration's real key. It has no path validator, so + * `GET /v1/models` reaches api.anthropic.com and models are discovered + * live per integration. + * - `aws` — a SigV4-verifying proxy on `{connect}/__gateway__/bedrock/{guid}` + * that serves Bedrock *Runtime* operations only, so models are declared + * from ai-config's Connect Bedrock table rather than discovered. Requests + * are signed with per-integration STS credentials minted by the host + * (see {@link ConnectProviderCallbacks.getAwsCredentials}). + */ + +import type { ResolvedProviderId } from "ai-config"; +import { + CONNECT_BEDROCK_MODELS, + getAnthropicModelCapabilities, + getConnectBedrockModelCapabilities, +} from "ai-config"; + +import { additiveHeaderRecord } from "../custom-headers"; +import { createAbortControllerFromToken } from "../model-clients/ai-sdk-helpers"; +import { AnthropicClient } from "../model-clients/AnthropicClient"; +import { BedrockClient } from "../model-clients/BedrockClient"; +import type { ModelClient, ModelClientChatParams } from "../model-clients/ModelClient"; +import type { ApiKeyCredentials, AwsCredentials, Logger, LMStreamPart, ModelInfo } from "../types"; +import { normalizeProtocol } from "../types"; +import { createCachedModelFetcher } from "./cached-model-fetcher"; +import type { ClientFactory, ProviderRegistry } from "./ProviderRegistry"; + +const DEFAULT_TEMPLATES = ["anthropic", "aws"] as const; +/** + * Templates this module knows how to shape into gateway-backed models. A + * host-configured allowlist is intersected with this set rather than trusted + * outright — a template outside it has no shaping rule at all. + */ +const SUPPORTED_TEMPLATES: ReadonlySet = new Set(DEFAULT_TEMPLATES); +/** + * Connect's AWS template mints credentials for the *calling user* only under + * Viewer auth (`connect/src/connect/auth/oauth2/templates_store.go`); Service + * Account auth mints against a content item, which a user session is not. + * Shaping one would produce models that always fail. + */ +const VIEWER_AUTH_TYPE = "Viewer"; +/** Connect's own default when an AWS integration names no region. */ +const DEFAULT_STS_REGION = "us-east-1"; +const INTEGRATIONS_PATH = "/__api__/v1/oauth/integrations"; +const ANTHROPIC_VERSION_HEADER = "2023-06-01"; +const CONNECT_CACHE_MAX_ENTRIES = 32; + +/** One allowlisted integration, shaped from Connect's integrations endpoint. */ +export interface ConnectIntegration { + /** + * The model-id namespace for this integration (`connect-`); every + * model it serves is listed as `/`. See + * {@link mintIntegrationPrefix}. + */ + readonly idPrefix: string; + readonly guid: string; + readonly template: string; + readonly name: string; + readonly description: string; + /** The integration's gateway route; see {@link gatewayBaseUrl}. */ + readonly baseUrl: string; + /** + * AWS region for `aws` integrations, read from the integration's own + * `config.sts_region`. Connect verifies inbound SigV4 against this exact + * value, so it must come from the record, never from local config. + * `undefined` for non-AWS templates. + */ + readonly region?: string; + /** + * `{connect}/__oauth__/integrations/{guid}/login` — Connect's interactive + * login for this integration. + */ + readonly loginUrl: string; +} + +/** A successfully minted set of per-integration AWS credentials. */ +export interface ConnectAwsCredentialSuccess { + ok: true; + credentials: AwsCredentials; + expiresAt?: string; +} + +/** A failure to mint AWS credentials, translated from the host's wire form. */ +export interface ConnectAwsCredentialFailure { + ok: false; + /** Machine-readable failure code, e.g. `"oauth_session_required"`. */ + code: string; + detail?: string; + /** The integration's login URL when signing in would fix the failure. */ + loginUrl?: string; +} + +export type ConnectAwsCredentialResult = ConnectAwsCredentialSuccess | ConnectAwsCredentialFailure; + +/** + * Platform-provided Connect hooks. Pre-built by the Node caller and threaded + * in through registration; the bridge must NOT construct these. + */ +export interface ConnectProviderCallbacks { + /** + * Mint per-integration AWS credentials for an `aws`-template integration + * (rserver's `/connect_aws_credentials`, keyed by the integration's guid). + * Called per request — the credentials are short-lived STS material. + * `signal` aborts when the chat is cancelled; honor it so a cancelled chat + * is not held behind a hung credential exchange. + */ + getAwsCredentials( + integration: ConnectIntegration, + signal?: AbortSignal, + ): Promise; + /** + * The admin-configured template allowlist. The bridge intersects it with + * its supported set; absent means the default (`anthropic`, `aws`). + */ + templates?(): readonly string[]; +} + +// --------------------------------------------------------------------------- +// Integration shaping (ported from @assistant/rstudio's connectSession.ts, +// re-scoped from provider ids to model-id prefixes) +// --------------------------------------------------------------------------- + +function slugify(input: string): string { + return input + .toLowerCase() + .replace(/[\s_]+/g, "-") + .replace(/[^a-z0-9-]/g, "") + .replace(/-{2,}/g, "-") + .replace(/^-|-$/g, ""); +} + +/** + * Derive a human-readable model-id prefix from the integration's admin-facing + * name (e.g. `connect-anthropic-superuser-`), falling back to + * description, then template, when the name is blank. The guid is always + * embedded, so two integrations can never mint the same prefix regardless of + * record order — chat routing keys on the stamped gateway URL regardless, but + * this keeps the prefix itself collision-free too. + */ +function mintIntegrationPrefix(input: { + name: string; + description: string; + template: string; + guid: string; +}): string { + const slug = slugify(input.name) || slugify(input.description) || slugify(input.template); + return slug ? `connect-${slug}-${input.guid}` : `connect-${input.guid}`; +} + +/** + * The gateway route for an integration. Anthropic's route is versioned + * (`/v1`) because the Anthropic client appends bare paths; Bedrock's is not — + * the AWS SDK appends `/model/{id}/{action}` itself. + */ +function gatewayBaseUrl(connectUrl: string, template: string, guid: string): string { + return template === "aws" + ? `${connectUrl}/__gateway__/bedrock/${guid}` + : `${connectUrl}/__gateway__/anthropic/${guid}/v1`; +} + +/** Connect's interactive login for one integration; same shape for every template. */ +function integrationLoginUrl(connectUrl: string, guid: string): string { + return `${connectUrl}/__oauth__/integrations/${guid}/login`; +} + +/** `name`/`description`/`template` are NullString on the wire; tolerate `null` as `""`. */ +function nullString(value: unknown): string { + return typeof value === "string" ? value : ""; +} + +/** `config` is a string map on the wire, with secrets already stripped by Connect. */ +function configString(record: Record, key: string): string | undefined { + const config = record.config; + if (typeof config !== "object" || config === null) return undefined; + const value = (config as Record)[key]; + return typeof value === "string" && value.length > 0 ? value : undefined; +} + +/** + * Shape the integrations endpoint's raw body into allowlisted + * {@link ConnectIntegration}s. Exported so hosts that resolve prefixes back to + * integrations (e.g. a credential transport) mint identical slugs; templates + * outside {@link SUPPORTED_TEMPLATES} are dropped here regardless of the + * allowlist, so a host may pass the admin allowlist verbatim and still shape + * the exact records (and prefixes) the bridge shapes. + */ +export function shapeConnectIntegrations( + body: readonly unknown[], + connectUrl: string, + templates: readonly string[], +): ConnectIntegration[] { + const integrations: ConnectIntegration[] = []; + for (const entry of body) { + if (typeof entry !== "object" || entry === null) continue; + const record = entry as Record; + const guid = typeof record.guid === "string" ? record.guid : ""; + if (!guid) continue; + const template = nullString(record.template); + if (!templates.includes(template) || !SUPPORTED_TEMPLATES.has(template)) continue; + if (template === "aws" && nullString(record.auth_type) !== VIEWER_AUTH_TYPE) continue; + const name = nullString(record.name); + const description = nullString(record.description); + const idPrefix = mintIntegrationPrefix({ name, description, template, guid }); + integrations.push({ + idPrefix, + guid, + template, + name, + description, + baseUrl: gatewayBaseUrl(connectUrl, template, guid), + loginUrl: integrationLoginUrl(connectUrl, guid), + ...(template === "aws" + ? { region: configString(record, "sts_region") ?? DEFAULT_STS_REGION } + : {}), + }); + } + return integrations; +} + +// --------------------------------------------------------------------------- +// Model discovery +// --------------------------------------------------------------------------- + +/** + * Credential-scoped routing facts for unstamped configured models. Entries + * use the same bound and FIFO policy as the model cache, and clear with it. + */ +class ConnectIntegrationCache { + private readonly entries = new Map>(); + private generation = 0; + + currentGeneration(): number { + return this.generation; + } + + replace( + credentialKey: string, + integrations: readonly ConnectIntegration[], + startedInGeneration: number, + ): void { + if (this.generation !== startedInGeneration) return; + this.entries.delete(credentialKey); + this.entries.set( + credentialKey, + new Map(integrations.map((integration) => [integration.idPrefix, integration])), + ); + while (this.entries.size > CONNECT_CACHE_MAX_ENTRIES) { + const oldestKey = this.entries.keys().next().value; + if (oldestKey === undefined) break; + this.entries.delete(oldestKey); + } + } + + get(credentialKey: string, prefix: string): ConnectIntegration | undefined { + return this.entries.get(credentialKey)?.get(prefix); + } + + findByGuid(credentialKey: string, guid: string): ConnectIntegration | undefined { + for (const integration of this.entries.get(credentialKey)?.values() ?? []) { + if (integration.guid === guid) return integration; + } + return undefined; + } + + clear(): void { + this.generation += 1; + this.entries.clear(); + } +} + +/** Opaque fingerprint identifying the server and token that populated an entry. */ +async function connectCredentialKey( + connectUrl: string, + credentials: ApiKeyCredentials, +): Promise { + const input = new TextEncoder().encode(`${connectUrl}${credentials.apiKey}`); + const digest = new Uint8Array(await globalThis.crypto.subtle.digest("SHA-256", input)); + let fingerprint = ""; + for (const byte of digest) { + fingerprint += byte.toString(16).padStart(2, "0"); + } + return fingerprint; +} + +function resolveTemplates( + configured: readonly string[] | undefined, + logger: Logger, +): readonly string[] { + if (!configured) return DEFAULT_TEMPLATES; + const unsupported = configured.filter((template) => !SUPPORTED_TEMPLATES.has(template)); + if (unsupported.length > 0) { + logger.warn( + `[connect] Ignoring unsupported integration template(s): ${unsupported.join(", ")}; ` + + `only ${[...SUPPORTED_TEMPLATES].join(", ")} can be shaped into models today.`, + ); + } + return configured.filter((template) => SUPPORTED_TEMPLATES.has(template)); +} + +async function fetchIntegrationRecords( + connectUrl: string, + credentials: ApiKeyCredentials, + signal: AbortSignal, +): Promise { + const headers = additiveHeaderRecord( + { Authorization: `Key ${credentials.apiKey}` }, + credentials.customHeaders, + ); + const response = await fetch(`${connectUrl}${INTEGRATIONS_PATH}`, { headers, signal }); + if (!response.ok) { + throw new Error(`Connect integrations endpoint returned ${response.status}`); + } + const body: unknown = await response.json(); + if (!Array.isArray(body)) { + throw new Error("Connect integrations response was not a JSON array"); + } + return body; +} + +/** + * Display-name suffix tying a model to its integration — one flat list serves + * every integration, and two integrations can offer the same model. The + * unique idPrefix stands in when the admin left the name blank. + */ +function integrationLabel(integration: ConnectIntegration): string { + return integration.name || integration.idPrefix; +} + +/** + * The declared Bedrock models for one `aws`-template integration. Models + * recognized by the Anthropic-on-Bedrock table declare `anthropic-messages` + * — the route {@link BedrockClient}'s heuristic actually takes for them — so + * host-side behavior keyed on the declared protocol (e.g. explicit + * prompt-cache markers) matches the wire format in use; anything else falls + * back to `bedrock-converse`. + */ +function bedrockGatewayModels( + providerId: ResolvedProviderId, + integration: ConnectIntegration, +): ModelInfo[] { + return CONNECT_BEDROCK_MODELS.map((model) => ({ + id: `${integration.idPrefix}/${model.id}`, + name: `${model.name} (${integrationLabel(integration)})`, + providerId, + vendor: "anthropic", + ...getConnectBedrockModelCapabilities(model.id), + protocol: + getAnthropicModelCapabilities(model.id) !== undefined + ? ("anthropic-messages" as const) + : ("bedrock-converse" as const), + baseUrl: integration.baseUrl, + })); +} + +/** + * Live per-integration discovery through the Anthropic gateway, which proxies + * `GET /models` to api.anthropic.com with the integration's real key. + */ +async function discoverAnthropicGatewayModels( + providerId: ResolvedProviderId, + integration: ConnectIntegration, + credentials: ApiKeyCredentials, + signal: AbortSignal, +): Promise { + const headers = additiveHeaderRecord( + { "x-api-key": credentials.apiKey, "anthropic-version": ANTHROPIC_VERSION_HEADER }, + credentials.customHeaders, + ); + const response = await fetch(`${integration.baseUrl}/models`, { headers, signal }); + if (!response.ok) { + throw new Error(`Anthropic gateway returned ${response.status}`); + } + const body: unknown = await response.json(); + const data = + typeof body === "object" && body !== null ? (body as Record).data : undefined; + if (!Array.isArray(data)) { + throw new Error("Anthropic gateway model list had no data array"); + } + + const models: ModelInfo[] = []; + for (const entry of data) { + if (typeof entry !== "object" || entry === null) continue; + const record = entry as Record; + if (typeof record.id !== "string" || record.id.length === 0) continue; + const displayName = typeof record.display_name === "string" ? record.display_name : record.id; + models.push({ + id: `${integration.idPrefix}/${record.id}`, + name: `${displayName} (${integrationLabel(integration)})`, + providerId, + vendor: "anthropic", + family: undefined, + maxInputTokens: 200_000, + maxOutputTokens: 16_000, + supportsTools: true, + supportsImages: true, + supportsToolResultImages: true, + supportedInputMediaTypes: [ + "image/png", + "image/jpeg", + "image/gif", + "image/webp", + "application/pdf", + ], + maxContextLength: 200_000, + ...getAnthropicModelCapabilities(record.id), + // The gateway forwards to api.anthropic.com, so provider-native web + // search works exactly as it does against Anthropic directly. + supportsWebSearch: true, + protocol: "anthropic-messages" as const, + baseUrl: integration.baseUrl, + }); + } + return models; +} + +function createConnectModelFetcher( + providerId: ResolvedProviderId, + logger: Logger, + cache: ConnectIntegrationCache, + callbacks?: ConnectProviderCallbacks, +) { + const fetcher = createCachedModelFetcher({ + providerId, + // The token AND the Connect server URL must both be present; without a + // baseUrl there is nothing to discover against (Snowflake pattern). + hasCredentials: (credentials) => Boolean(credentials.apiKey && credentials.baseUrl), + // Different sessions against the same (or different) Connect server must + // never share a cached model list — the list reflects which + // integrations THIS token can see. + cacheKey: (credentials) => + connectCredentialKey(credentials.baseUrl!.replace(/\/+$/, ""), credentials), + maxCacheEntries: CONNECT_CACHE_MAX_ENTRIES, + fetchFresh: async (credentials, signal) => { + const cacheGeneration = cache.currentGeneration(); + const connectUrl = credentials.baseUrl!.replace(/\/+$/, ""); + const credentialKey = await connectCredentialKey(connectUrl, credentials); + const templates = resolveTemplates(callbacks?.templates?.(), logger); + const records = await fetchIntegrationRecords(connectUrl, credentials, signal); + const integrations = shapeConnectIntegrations(records, connectUrl, templates); + + if (!callbacks && integrations.some((integration) => integration.template === "aws")) { + logger.warn( + "[connect] Skipping AWS-backed integrations: no AWS credential callback was provided, " + + "so their gateway requests could never be signed.", + ); + } + const modelLists = await Promise.all( + integrations.map(async (integration): Promise => { + if (integration.template === "aws") { + return callbacks ? bedrockGatewayModels(providerId, integration) : []; + } + // One misconfigured integration must not hide the others, so + // per-integration discovery failures are isolated — but a deadline + // abort fails the whole discovery (the fetcher has already fallen + // back), so it must not be swallowed as a per-integration failure. + try { + return await discoverAnthropicGatewayModels( + providerId, + integration, + credentials, + signal, + ); + } catch (error) { + if (signal.aborted) throw error; + const message = error instanceof Error ? error.message : String(error); + logger.warn( + `[connect] Model discovery failed for integration "${integrationLabel(integration)}": ${message}`, + ); + return []; + } + }), + ); + const models = modelLists.flat(); + if (signal.aborted) { + throw signal.reason ?? new Error("Connect model discovery was aborted"); + } + cache.replace(credentialKey, integrations, cacheGeneration); + return models; + }, + fallbackModels: [], + logger, + }); + const clearModelCache = fetcher.clearCache; + fetcher.clearCache = () => { + cache.clear(); + clearModelCache?.(); + }; + return fetcher; +} + +// --------------------------------------------------------------------------- +// Client +// --------------------------------------------------------------------------- + +/** + * Split a namespaced `connect-/` id on the FIRST `/` only — + * the minted prefix never contains one, but Bedrock model ids (ARNs) may. + * A first segment that is not a minted `connect-` prefix (e.g. a raw ARN's + * `arn:aws:...`) leaves the whole id as the wire model. + */ +function splitConnectModelId(model: string): { prefix?: string; wireModel: string } { + const separator = model.indexOf("/"); + if (separator <= 0 || !model.slice(0, separator).startsWith("connect-")) { + return { wireModel: model }; + } + return { prefix: model.slice(0, separator), wireModel: model.slice(separator + 1) }; +} + +/** + * Recover the routing facts a discovery stamp encodes. The gateway URL shape + * ({@link gatewayBaseUrl}) embeds the Connect server, the proxy kind (which + * names the template), and the integration guid, which is what makes + * stamped-model routing stateless. Returns `undefined` for URLs that are not + * Connect gateway routes. + */ +function parseGatewayBaseUrl( + baseUrl: string, +): { connectUrl: string; template: string; guid: string } | undefined { + const match = /^(.+?)\/__gateway__\/(anthropic|bedrock)\/([^/]+)/.exec(baseUrl); + if (!match) return undefined; + return { + connectUrl: match[1], + template: match[2] === "bedrock" ? "aws" : "anthropic", + guid: match[3], + }; +} + +/** + * Routes each request to the integration's gateway: `anthropic-messages` + * spends the federated token as `x-api-key` (the gateway swaps it for the + * real key), `bedrock-converse` signs with per-request STS credentials minted + * through {@link ConnectProviderCallbacks.getAwsCredentials}. + */ +class ConnectClient implements ModelClient { + constructor( + private readonly credentials: ApiKeyCredentials, + private readonly cache: ConnectIntegrationCache, + private readonly logger: Logger, + private readonly callbacks?: ConnectProviderCallbacks, + ) {} + + async chat(params: ModelClientChatParams): Promise> { + const connectUrl = this.credentials.baseUrl?.replace(/\/+$/, ""); + if (!connectUrl) { + throw new Error("Connect provider credentials carry no Connect server URL."); + } + const { prefix, wireModel } = splitConnectModelId(params.model); + + // A discovery-stamped baseUrl is the source of truth: it names the + // server, template, and guid, so routing never depends on cache state. + // An unstamped model resolves to the bare Connect root as ai-config's + // fallback (the provider's own configured baseUrl), which is not a real + // stamp — treat it the same as no baseUrl at all and fall back to the + // cache. Only user-configured models without a real stamp reach that + // fallback. + const hasStampedBaseUrl = + Boolean(params.baseUrl) && params.baseUrl!.replace(/\/+$/, "") !== connectUrl; + + let integration: ConnectIntegration; + if (hasStampedBaseUrl) { + const parsed = parseGatewayBaseUrl(params.baseUrl!); + if (!parsed) { + throw new Error( + `Connect provider model "${params.model}" has base URL "${params.baseUrl}", ` + + `which is not a Connect gateway URL.`, + ); + } + if (parsed.connectUrl !== connectUrl) { + throw new Error( + `Connect provider model "${params.model}" was discovered against ${parsed.connectUrl}, ` + + `but the provider now points at ${connectUrl}. Refresh the model list and try again.`, + ); + } + integration = await this.resolveIntegration(parsed, prefix); + } else { + const credentialKey = await connectCredentialKey(connectUrl, this.credentials); + const cached = prefix ? this.cache.get(credentialKey, prefix) : undefined; + if (!cached) { + throw new Error( + `Connect provider has no gateway base URL for model "${params.model}". ` + + `Refresh the model list and try again.`, + ); + } + integration = cached; + } + + const protocol = normalizeProtocol(params.protocol); + // An unstamped model's baseUrl (when present at all) is ai-config's bare + // Connect-root fallback, not a real gateway route — always route through + // the resolved integration's baseUrl in that case. + const baseUrl = hasStampedBaseUrl ? params.baseUrl! : integration.baseUrl; + + // The template selects the transport — a protocol override can pick the + // wire format within it, but never re-routes off the integration's + // gateway (an aws gateway always requires SigV4, whatever the protocol). + if (integration.template === "aws") { + if (protocol && protocol !== "bedrock-converse" && protocol !== "anthropic-messages") { + throw new Error( + `Connect provider cannot route protocol "${params.protocol}" for model "${params.model}"; ` + + `an AWS-backed integration supports bedrock-converse and anthropic-messages.`, + ); + } + return this.bedrockChat( + params, + wireModel, + baseUrl, + integration, + protocol === "anthropic-messages" ? "anthropic-messages" : undefined, + ); + } + + if (protocol && protocol !== "anthropic-messages") { + throw new Error( + `Connect provider cannot route protocol "${params.protocol}" for model "${params.model}"; ` + + `an Anthropic-backed integration supports only anthropic-messages.`, + ); + } + const client = new AnthropicClient( + { apiKey: this.credentials.apiKey }, + baseUrl, + this.credentials.customHeaders, + this.logger, + ); + return client.chat({ ...params, model: wireModel, baseUrl }); + } + + /** + * Resolve the integration record for a stamped model: prefer the cached + * record (it carries the admin-facing name and `sts_region`) when it came + * from these same credentials, else synthesize one from the stamp so + * routing survives an empty, replaced, or foreign-credential cache. + */ + private async resolveIntegration( + parsed: { connectUrl: string; template: string; guid: string }, + prefix: string | undefined, + ): Promise { + const credentialKey = await connectCredentialKey(parsed.connectUrl, this.credentials); + const cached = this.cache.findByGuid(credentialKey, parsed.guid); + if (cached) return cached; + return { + idPrefix: prefix ?? `connect-${parsed.guid}`, + guid: parsed.guid, + template: parsed.template, + name: "", + description: "", + baseUrl: gatewayBaseUrl(parsed.connectUrl, parsed.template, parsed.guid), + loginUrl: integrationLoginUrl(parsed.connectUrl, parsed.guid), + }; + } + + private async bedrockChat( + params: ModelClientChatParams, + wireModel: string, + baseUrl: string, + integration: ConnectIntegration, + protocol: "anthropic-messages" | undefined, + ): Promise> { + if (!this.callbacks) { + throw new Error( + "Connect provider has no AWS credential callback; Bedrock-backed integrations are unavailable.", + ); + } + // The mint is a network exchange of its own; tie it to the chat's + // cancellation so an abandoned request is not held behind it. + const { abortController, cleanup } = createAbortControllerFromToken(params.cancellationToken); + let result: ConnectAwsCredentialResult; + try { + result = await this.callbacks.getAwsCredentials(integration, abortController.signal); + } finally { + cleanup(); + } + if (!result.ok) { + const detail = result.detail ? ` ${result.detail}` : ""; + const login = result.loginUrl ? ` Sign in at ${result.loginUrl} and try again.` : ""; + throw new Error( + `Posit Connect could not mint AWS credentials (${result.code}).${detail}${login}`, + ); + } + const aws = result.credentials; + if (!aws.accessKeyId || !aws.secretAccessKey) { + // An incomplete key set would send BedrockClient to the ambient AWS + // credential chain — the wrong identity for a gateway that verifies + // the minted one. + throw new Error( + `Posit Connect returned incomplete AWS credentials for integration ` + + `"${integrationLabel(integration)}".`, + ); + } + // Connect verifies inbound SigV4 against the integration record's + // sts_region, so prefer it; the minted material's region covers a + // synthesized record that has none. + const region = integration.region ?? aws.region; + if (integration.region && aws.region && integration.region !== aws.region) { + this.logger.warn( + `[connect] Integration "${integrationLabel(integration)}" declares sts_region ` + + `${integration.region} but the minted credentials name ${aws.region}; signing with ${region}.`, + ); + } + const client = new BedrockClient( + { + region, + accessKeyId: aws.accessKeyId, + secretAccessKey: aws.secretAccessKey, + sessionToken: aws.sessionToken, + customHeaders: this.credentials.customHeaders, + // Routing through Connect's gateway is a deliberate, admin-configured + // redirect (not an accidental override), so it overrides even a FIPS + // runtime endpoint. + allowBaseUrlUnderFips: true, + }, + this.logger, + ); + // Absent protocol: BedrockClient's model-id heuristic keeps + // `us.anthropic.*` ids on the native Anthropic (InvokeModel) route, which + // the gateway also allows and which supports thinking. + return client.chat({ ...params, model: wireModel, baseUrl, protocol }); + } +} + +function createConnectClientFactory( + logger: Logger, + cache: ConnectIntegrationCache, + callbacks?: ConnectProviderCallbacks, +): ClientFactory { + return (credentials) => { + if (credentials.type !== "apikey") { + throw new Error(`Connect provider requires API key credentials, got: ${credentials.type}`); + } + return new ConnectClient(credentials, cache, logger, callbacks); + }; +} + +export function registerConnectProvider( + registry: ProviderRegistry, + logger: Logger, + callbacks?: ConnectProviderCallbacks, +): void { + // Shared between the fetcher (writer) and the client (reader) so chat can + // resolve user-configured models that carry no discovery stamp; stamped + // models route from their gateway URL alone. + const cache = new ConnectIntegrationCache(); + registry.registerModelFetcher( + "connect", + createConnectModelFetcher("connect", logger, cache, callbacks), + ); + registry.registerClientFactory("connect", createConnectClientFactory(logger, cache, callbacks)); +} diff --git a/packages/ai-provider-bridge/src/register-all-providers.ts b/packages/ai-provider-bridge/src/register-all-providers.ts index d63f42e..c4c835b 100644 --- a/packages/ai-provider-bridge/src/register-all-providers.ts +++ b/packages/ai-provider-bridge/src/register-all-providers.ts @@ -15,6 +15,10 @@ import { registerBedrockProvider, type BedrockProviderCallbacks, } from "./providers/bedrock-provider"; +import { + registerConnectProvider, + type ConnectProviderCallbacks, +} from "./providers/connect-provider"; import { registerCopilotProvider } from "./providers/copilot-provider"; import { registerDatabricksProvider } from "./providers/databricks-provider"; import { registerDeepSeekProvider } from "./providers/deepseek-provider"; @@ -49,6 +53,7 @@ export interface ProviderRegistrationConfig { bedrockCallbacks?: BedrockProviderCallbacks; googleVertexCallbacks?: GoogleVertexProviderCallbacks; snowflakeCallbacks?: SnowflakeProviderCallbacks; + connectCallbacks?: ConnectProviderCallbacks; } /** @@ -90,6 +95,8 @@ const PROVIDER_REGISTRARS = { databricks: registerDatabricksProvider, litellm: registerLitellmProvider, portkey: registerPortkeyProvider, + connect: (registry, logger, config) => + registerConnectProvider(registry, logger, config.connectCallbacks), } satisfies Record; /** diff --git a/packages/ai-provider-bridge/src/types.ts b/packages/ai-provider-bridge/src/types.ts index 54543f0..1feafcf 100644 --- a/packages/ai-provider-bridge/src/types.ts +++ b/packages/ai-provider-bridge/src/types.ts @@ -57,6 +57,7 @@ export const PROVIDER_IDS = [ "databricks", "litellm", "portkey", + "connect", ] as const; /**