diff --git a/apps/api/src/llm-token.test.ts b/apps/api/src/llm-token.test.ts new file mode 100644 index 000000000..e61a109ab --- /dev/null +++ b/apps/api/src/llm-token.test.ts @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + acquireEndpoint: vi.fn(), + acquireToken: vi.fn(), + modelClient: vi.fn(() => ({ path: vi.fn() })), + azureKeyCredential: vi.fn(function (this: any, key: string) { + this.key = key; + }), +})); + +vi.mock("shared", () => ({ + TokenManagerClient: class { + acquireEndpoint = mocks.acquireEndpoint; + acquireToken = mocks.acquireToken; + }, +})); + +vi.mock("@azure-rest/ai-inference", () => ({ + default: mocks.modelClient, +})); + +vi.mock("@azure/core-auth", () => ({ + AzureKeyCredential: mocks.azureKeyCredential, +})); + +import { acquireInferenceClient } from "./llm-token.js"; + +describe("acquireInferenceClient", () => { + const originalEnv = process.env; + + beforeEach(() => { + vi.clearAllMocks(); + process.env = { ...originalEnv }; + delete process.env.AZURE_AI_INFERENCE_ENDPOINT; + delete process.env.AZURE_AI_INFERENCE_API_KEY; + delete process.env.GITHUB_MODELS_API_KEY; + delete process.env.GITHUB_TOKEN; + process.env.TOKEN_MANAGER_URL = "http://token-manager:3000"; + }); + + afterEach(() => { + process.env = originalEnv; + }); + + it("uses the Token Manager structured endpoint response", async () => { + mocks.acquireEndpoint.mockResolvedValue({ + endpoint: "https://foundry.example.com/models", + apiKey: "foundry-key", + deployment: "gpt-4.1-mini", + }); + + const result = await acquireInferenceClient(); + + expect(mocks.acquireEndpoint).toHaveBeenCalledWith("azure-ai-inference"); + expect(mocks.acquireToken).not.toHaveBeenCalledWith("azure-ai-inference"); + expect(mocks.azureKeyCredential).toHaveBeenCalledWith("foundry-key"); + expect(mocks.modelClient).toHaveBeenCalledWith( + "https://foundry.example.com/models", + expect.anything() + ); + expect(result).toMatchObject({ + endpoint: "https://foundry.example.com/models", + source: "azure-ai-foundry", + via: "azure-ai-foundry-token-manager", + model: "gpt-4.1-mini", + }); + }); +}); diff --git a/apps/api/src/llm-token.ts b/apps/api/src/llm-token.ts index b9d080357..99e7063e5 100644 --- a/apps/api/src/llm-token.ts +++ b/apps/api/src/llm-token.ts @@ -20,10 +20,7 @@ */ import ModelClient, { type ModelClient as ModelClientType } from "@azure-rest/ai-inference"; import { AzureKeyCredential } from "@azure/core-auth"; -import { - TokenManagerClient, - parseAzureAiFoundrySecret, -} from "shared"; +import { TokenManagerClient } from "shared"; const GITHUB_MODELS_ENDPOINT = "https://models.inference.ai.azure.com"; @@ -189,13 +186,7 @@ function logInferenceAcquired(handle: InferenceClientHandle): void { * * Returns null when no token-manager is configured, when no key is * registered for the capability, when the request fails, or when the - * registered secret is malformed. - * - * Note: when only AZURE_AI_INFERENCE_API_KEY is set (without the matching - * endpoint), TokenManagerClient's env-var shortcut returns the bare key, - * which parseAzureAiFoundrySecret rejects as malformed — so we fall - * through to the next backend. The both-vars-set case is already handled - * in acquireInferenceClient before this is called. + * registered endpoint is malformed. */ async function tryAcquireFoundryFromTokenManager(): Promise<{ endpoint: string; @@ -206,8 +197,12 @@ async function tryAcquireFoundryFromTokenManager(): Promise<{ if (!client) return null; try { - const raw = await client.acquireToken("azure-ai-inference"); - return parseAzureAiFoundrySecret(raw); + const endpoint = await client.acquireEndpoint("azure-ai-inference"); + return { + endpoint: endpoint.endpoint, + apiKey: endpoint.apiKey, + model: endpoint.deployment, + }; } catch { return null; } diff --git a/apps/token-manager/src/endpoint-routes.test.ts b/apps/token-manager/src/endpoint-routes.test.ts new file mode 100644 index 000000000..4f309b216 --- /dev/null +++ b/apps/token-manager/src/endpoint-routes.test.ts @@ -0,0 +1,210 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import express from "express"; +import type { Collection, FindCursor, WithId } from "mongodb"; +import type { KeyDocument } from "shared"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { SecretStore } from "./keyvault-store.js"; +import { createKeyRouter } from "./routes.js"; + +function makeEndpoint(overrides: Partial = {}): KeyDocument { + return { + _id: "foundry-1", + type: "azure-ai-foundry", + capabilities: ["azure-ai-inference"], + secretName: "token-azure-ai-foundry-foundry", + lastValidationStatus: "valid", + enabled: true, + acquireCount: 0, + createdAt: new Date("2026-01-01T00:00:00Z"), + ...overrides, + }; +} + +function makeCollection(docs: KeyDocument[]) { + return { + find: vi.fn((filter: Record) => { + const matches = docs.filter((doc) => { + if (filter.enabled !== undefined && doc.enabled !== filter.enabled) return false; + if (filter.lastValidationStatus && doc.lastValidationStatus !== filter.lastValidationStatus) { + return false; + } + if (filter.deletedAt?.$exists === false && doc.deletedAt) return false; + if (filter.capabilities?.$in && !filter.capabilities.$in.some( + (capability: string) => doc.capabilities.includes(capability as any) + )) return false; + if (filter.type?.$in && !filter.type.$in.includes(doc.type)) return false; + return true; + }); + return { + toArray: vi.fn(async () => matches), + } as unknown as FindCursor>; + }), + updateOne: vi.fn(async () => ({ matchedCount: 1, modifiedCount: 1 })), + }; +} + +function makeStore(secrets: Record) { + return { + getSecret: vi.fn(async (name: string) => { + const value = secrets[name]; + if (value === undefined) throw new Error(`Secret '${name}' not found`); + return value; + }), + }; +} + +function makeApp( + collection: ReturnType, + store: ReturnType +) { + const app = express(); + app.use(express.json()); + app.use(createKeyRouter( + collection as unknown as Collection, + store as unknown as SecretStore + )); + return app; +} + +async function request( + app: express.Express, + path: string, + body: unknown +): Promise<{ status: number; body: any }> { + return new Promise((resolve) => { + const req = { + method: "POST", + url: path, + headers: { "content-type": "application/json" }, + body, + params: {}, + query: {}, + get: (header: string) => header === "content-type" ? "application/json" : undefined, + } as unknown as express.Request; + let status = 200; + const res = { + status(code: number) { + status = code; + return this; + }, + json(data: unknown) { + resolve({ status, body: data }); + }, + setHeader() { + return this; + }, + getHeader() { + return undefined; + }, + end() { + resolve({ status, body: null }); + }, + } as unknown as express.Response; + + (app as any).handle(req, res, (error: Error) => { + resolve({ status: 500, body: { error: error.message } }); + }); + }); +} + +describe("structured endpoint acquisition", () => { + let endpoint: KeyDocument; + let collection: ReturnType; + let store: ReturnType; + let app: express.Express; + + beforeEach(() => { + endpoint = makeEndpoint(); + collection = makeCollection([endpoint]); + store = makeStore({ + [endpoint.secretName]: JSON.stringify({ + endpoint: "https://foundry.example.com/models/", + apiKey: "foundry-key", + model: "gpt-4.1-mini", + }), + }); + app = makeApp(collection, store); + }); + + it("returns a typed endpoint and records the acquisition", async () => { + const response = await request(app, "/api/v1/endpoints/acquire", { + capability: "azure-ai-inference", + }); + + expect(response).toEqual({ + status: 200, + body: { + endpoint: "https://foundry.example.com/models", + apiKey: "foundry-key", + deployment: "gpt-4.1-mini", + }, + }); + expect(store.getSecret).toHaveBeenCalledWith(endpoint.secretName); + expect(collection.updateOne).toHaveBeenCalledWith( + { _id: endpoint._id }, + expect.objectContaining({ $inc: { acquireCount: 1 } }) + ); + }); + + it("returns 404 when no eligible endpoint exists", async () => { + collection = makeCollection([makeEndpoint({ enabled: false })]); + app = makeApp(collection, store); + + const response = await request(app, "/api/v1/endpoints/acquire", { + capability: "azure-ai-inference", + }); + + expect(response.status).toBe(404); + expect(response.body.error).toContain("No valid endpoints available"); + expect(store.getSecret).not.toHaveBeenCalled(); + }); + + it("returns 422 and does not count malformed stored credentials", async () => { + store = makeStore({ [endpoint.secretName]: "not-json" }); + app = makeApp(collection, store); + + const response = await request(app, "/api/v1/endpoints/acquire", { + capability: "azure-ai-inference", + }); + + expect(response.status).toBe(422); + expect(response.body.error).toContain("malformed"); + expect(collection.updateOne).not.toHaveBeenCalled(); + }); + + it("rejects unsupported endpoint capabilities", async () => { + const response = await request(app, "/api/v1/endpoints/acquire", { + capability: "copilot-sdk", + }); + + expect(response.status).toBe(400); + expect(response.body.error).toContain("Invalid endpoint capability"); + }); + + it.each([undefined, null])("rejects a missing or null request body (%#)", async (body) => { + const response = await request(app, "/api/v1/endpoints/acquire", body); + + expect(response.status).toBe(400); + expect(response.body.error).toContain("Invalid endpoint capability"); + expect(collection.find).not.toHaveBeenCalled(); + expect(store.getSecret).not.toHaveBeenCalled(); + }); + + it("preserves raw key acquisition for existing callers", async () => { + const response = await request(app, "/api/v1/keys/acquire", { + capability: "azure-ai-inference", + }); + + expect(response.status).toBe(200); + expect(response.body.value).toBe( + JSON.stringify({ + endpoint: "https://foundry.example.com/models/", + apiKey: "foundry-key", + model: "gpt-4.1-mini", + }) + ); + expect(response.body.keyType).toBe("azure-ai-foundry"); + }); +}); diff --git a/apps/token-manager/src/routes.ts b/apps/token-manager/src/routes.ts index 2766f4d28..9cfb4a361 100644 --- a/apps/token-manager/src/routes.ts +++ b/apps/token-manager/src/routes.ts @@ -12,7 +12,12 @@ import { CreateKeyRequest, UpdateKeyRequest, AcquireKeyRequest, + AcquireEndpointRequest, + AcquireEndpointResponse, + EndpointCapability, + EndpointType, deriveSecretName, + parseAzureAiFoundrySecret, KEY_TYPES, isKeyType, } from "shared"; @@ -30,6 +35,8 @@ const VALID_CAPABILITIES: KeyCapability[] = [ "anthropic-api", "azure-ai-inference", ]; +const VALID_ENDPOINT_CAPABILITIES: EndpointCapability[] = ["azure-ai-inference"]; +const VALID_ENDPOINT_TYPES: EndpointType[] = ["azure-ai-foundry"]; export function createKeyRouter( collection: Collection, @@ -353,5 +360,70 @@ export function createKeyRouter( } }); + // ────────────────────────────────────────────── + // POST /api/v1/endpoints/acquire — Structured endpoint acquisition (internal) + // ────────────────────────────────────────────── + router.post("/api/v1/endpoints/acquire", async (req, res, next) => { + try { + const body = req.body as Partial | null | undefined; + + if ( + !body?.capability || + !VALID_ENDPOINT_CAPABILITIES.includes(body.capability) + ) { + res.status(400).json({ + error: `Invalid endpoint capability. Must be one of: ${VALID_ENDPOINT_CAPABILITIES.join(", ")}`, + }); + return; + } + + const endpoints = await collection + .find({ + capabilities: { $in: [body.capability] }, + type: { $in: VALID_ENDPOINT_TYPES }, + enabled: true, + lastValidationStatus: "valid", + deletedAt: { $exists: false }, + }) + .toArray(); + + if (endpoints.length === 0) { + res.status(404).json({ + error: `No valid endpoints available for capability '${body.capability}'`, + }); + return; + } + + const selected = roundRobin.next(`endpoint:${body.capability}`, endpoints); + const value = await store.getSecret(selected.secretName); + const parsed = parseAzureAiFoundrySecret(value); + + if (!parsed) { + res.status(422).json({ + error: `Stored endpoint credential '${selected._id}' is malformed`, + }); + return; + } + + collection + .updateOne( + { _id: selected._id }, + { $inc: { acquireCount: 1 }, $set: { lastAcquiredAt: new Date() } } + ) + .catch((err) => + console.error(`[routes] Failed to update acquireCount for ${selected._id}:`, err) + ); + + const response: AcquireEndpointResponse = { + endpoint: parsed.endpoint, + apiKey: parsed.apiKey, + ...(parsed.model ? { deployment: parsed.model } : {}), + }; + res.json(response); + } catch (err) { + next(err); + } + }); + return router; } diff --git a/docs/architecture/token-manager.md b/docs/architecture/token-manager.md index 544d09667..de434d924 100644 --- a/docs/architecture/token-manager.md +++ b/docs/architecture/token-manager.md @@ -103,10 +103,17 @@ may contain a `requestProfile` field; the parser accepts and ignores it. ### Acquisition Fallback Chain -`TokenManagerClient.acquireToken(capability)` resolves a token in three tiers, returning the first one that succeeds: +`TokenManagerClient.acquireToken(capability)` resolves a token in two tiers, returning the first one that succeeds: 1. **Capability-specific env var** — e.g. `GITHUB_MODELS_TOKEN` for `github-models`, `ANTHROPIC_API_KEY` for `claude-code-cli`. Mapping lives in `KEY_CAPABILITY_ENV_VARS`. -3. **Token Manager HTTP service** — `GET /tokens/acquire?capability=...` against `TOKEN_MANAGER_URL`. This is the only tier used in production K8s deployments, where no static token env vars are mounted. +2. **Token Manager HTTP service** — `POST /api/v1/keys/acquire` against `TOKEN_MANAGER_URL`. This is the only tier used in production K8s deployments, where no static token env vars are mounted. + +Structured credentials use `TokenManagerClient.acquireEndpoint(capability)`. +It calls `POST /api/v1/endpoints/acquire` and returns an object containing +`endpoint`, `apiKey`, and an optional `deployment`. Initially this is a typed +projection over existing valid `azure-ai-foundry` keys, so registrations and +legacy `acquireToken()` callers remain compatible while consumers no longer +parse Key Vault JSON themselves. The API's discovery endpoint uses this chain via the `github-api-token` helper to acquire a `github-public-api` token per call, so each request can be served by a different token in round-robin. @@ -166,7 +173,17 @@ stateDiagram-v2 | Method | Endpoint | Description | |--------|----------|-------------| -| `GET` | `/api/v1/keys/acquire?capability=X` | Acquire a key for given capability | +| `POST` | `/api/v1/keys/acquire` | Acquire a key for the capability in the request body | + +### Endpoint Acquisition + +| Method | Endpoint | Description | +|--------|----------|-------------| +| `POST` | `/api/v1/endpoints/acquire` | Acquire a structured endpoint for a supported capability | + +The endpoint response contains `{ endpoint, apiKey, deployment? }`. The +service performs the stored JSON validation and returns HTTP 422 for malformed +endpoint credentials rather than requiring each consumer to parse the secret. The acquire endpoint uses round-robin selection among valid, enabled keys that provide the requested capability. @@ -260,10 +277,10 @@ first source that succeeds: `AZURE_AI_INFERENCE_ENDPOINT` + `AZURE_AI_INFERENCE_API_KEY`. Endpoint URLs missing the `/models` suffix are auto-corrected with a warning. -2. **Azure AI Foundry via the Token Manager** — fetched directly via - `POST {TOKEN_MANAGER_URL}/api/v1/keys/acquire {capability: - "azure-ai-inference"}` so the helper receives the full JSON blob - (endpoint + key + model), not just the API key. The capability +2. **Azure AI Foundry via the Token Manager** — fetched via + `TokenManagerClient.acquireEndpoint("azure-ai-inference")`, which calls + `POST {TOKEN_MANAGER_URL}/api/v1/endpoints/acquire` and returns a typed + endpoint + key + deployment response. The capability `azure-ai-inference` is derived from any registered `azure-ai-foundry` key. 3. **GitHub Models** — `GITHUB_MODELS_API_KEY` env var, then diff --git a/packages/shared/src/token-manager/client.test.ts b/packages/shared/src/token-manager/client.test.ts index 41c519879..903013ed6 100644 --- a/packages/shared/src/token-manager/client.test.ts +++ b/packages/shared/src/token-manager/client.test.ts @@ -200,4 +200,70 @@ describe("TokenManagerClient", () => { ); }); }); + + describe("acquireEndpoint", () => { + it("returns a typed endpoint from the Token Manager", async () => { + const result = { + endpoint: "https://foundry.example.com/models", + apiKey: "foundry-key", + deployment: "gpt-4.1-mini", + }; + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue({ + ok: true, + json: async () => result, + } as Response); + + const client = new TokenManagerClient("http://token-manager:80"); + + await expect(client.acquireEndpoint("azure-ai-inference")).resolves.toEqual(result); + expect(fetchSpy).toHaveBeenCalledWith( + "http://token-manager:80/api/v1/endpoints/acquire", + expect.objectContaining({ + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ capability: "azure-ai-inference" }), + }) + ); + }); + + it("throws when no Token Manager URL is configured", async () => { + delete process.env.TOKEN_MANAGER_URL; + const client = new TokenManagerClient(); + + await expect(client.acquireEndpoint("azure-ai-inference")).rejects.toThrow( + /No endpoint available.*TOKEN_MANAGER_URL/ + ); + }); + + it("includes the server response when acquisition fails", async () => { + vi.spyOn(globalThis, "fetch").mockResolvedValue({ + ok: false, + status: 404, + text: async () => "No valid endpoints available", + } as Response); + const client = new TokenManagerClient("http://token-manager:80"); + + await expect(client.acquireEndpoint("azure-ai-inference")).rejects.toThrow( + /Endpoint acquisition failed.*azure-ai-inference.*404.*No valid endpoints/ + ); + }); + + it.each([ + null, + {}, + { endpoint: "", apiKey: "key" }, + { endpoint: "https://foundry.example.com/models", apiKey: "" }, + { endpoint: "https://foundry.example.com/models", apiKey: "key", deployment: 42 }, + ])("rejects a malformed endpoint response %#", async (result) => { + vi.spyOn(globalThis, "fetch").mockResolvedValue({ + ok: true, + json: async () => result, + } as Response); + const client = new TokenManagerClient("http://token-manager:80"); + + await expect(client.acquireEndpoint("azure-ai-inference")).rejects.toThrow( + /Invalid endpoint response.*azure-ai-inference/ + ); + }); + }); }); diff --git a/packages/shared/src/token-manager/client.ts b/packages/shared/src/token-manager/client.ts index ba4f7aef7..3af834e5e 100644 --- a/packages/shared/src/token-manager/client.ts +++ b/packages/shared/src/token-manager/client.ts @@ -3,8 +3,10 @@ import { AcquireAccountResponse, + AcquireEndpointResponse, AcquireKeyResponse, AccountType, + EndpointCapability, KEY_CAPABILITY_ENV_VARS, KeyCapability, KeyType, @@ -139,6 +141,60 @@ export class TokenManagerClient { return result; } + /** + * Acquire a structured endpoint for the given capability. + * + * Unlike acquireToken(), this method has no single-value environment + * fallback: endpoint credentials require both an endpoint and an API key. + * Callers that support environment configuration should resolve the pair + * before calling this method. + * + * @throws Error if no endpoint is available or the response is malformed. + */ + async acquireEndpoint(capability: EndpointCapability): Promise { + if (!this.baseUrl) { + throw new Error( + `No endpoint available for capability '${capability}': TOKEN_MANAGER_URL is not configured` + ); + } + + const url = `${this.baseUrl}/api/v1/endpoints/acquire`; + const response = await fetch(url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ capability }), + signal: AbortSignal.timeout(10_000), + }); + + if (!response.ok) { + const errorBody = await response.text().catch(() => "unknown error"); + throw new Error( + `Endpoint acquisition failed for capability '${capability}' (HTTP ${response.status}): ${errorBody}` + ); + } + + const result: unknown = await response.json(); + if ( + typeof result !== "object" || + result === null || + !("endpoint" in result) || + typeof result.endpoint !== "string" || + result.endpoint.trim() === "" || + !("apiKey" in result) || + typeof result.apiKey !== "string" || + result.apiKey.trim() === "" || + ("deployment" in result && + result.deployment !== undefined && + typeof result.deployment !== "string") + ) { + throw new Error( + `Invalid endpoint response for capability '${capability}': missing or invalid credentials` + ); + } + + return result as AcquireEndpointResponse; + } + /** * Acquire account credentials for the given account type. * diff --git a/packages/shared/src/token-manager/types.ts b/packages/shared/src/token-manager/types.ts index d49e9c0b8..f96abc065 100644 --- a/packages/shared/src/token-manager/types.ts +++ b/packages/shared/src/token-manager/types.ts @@ -37,6 +37,16 @@ export function isKeyType(value: unknown): value is KeyType { export type KeyCapability = "github-models" | "github-public-api" | "copilot-models" | "copilot-sdk" | "copilot-cli" | "claude-code-cli" | "anthropic-api" | "azure-ai-inference"; +/** + * A structured endpoint provider stored by the Token Manager. + */ +export type EndpointType = "azure-ai-foundry"; + +/** + * A capability that can be acquired as a structured endpoint. + */ +export type EndpointCapability = "azure-ai-inference"; + /** * Validation status of a key. */ @@ -86,6 +96,22 @@ export interface AcquireKeyResponse { expiresAt?: Date; } +/** + * Request body for POST /api/v1/endpoints/acquire. + */ +export interface AcquireEndpointRequest { + capability: EndpointCapability; +} + +/** + * Structured response from POST /api/v1/endpoints/acquire. + */ +export interface AcquireEndpointResponse { + endpoint: string; + apiKey: string; + deployment?: string; +} + /** * Result of validating a key against its provider's API. */