feat(token-manager): add structured endpoint acquisition

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",
+    });
+  });
+});

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;
   }

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> = {}): 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<string, any>) => {
+      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<WithId<KeyDocument>>;
+    }),
+    updateOne: vi.fn(async () => ({ matchedCount: 1, modifiedCount: 1 })),
+  };
+}
+
+function makeStore(secrets: Record<string, string>) {
+  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<typeof makeCollection>,
+  store: ReturnType<typeof makeStore>
+) {
+  const app = express();
+  app.use(express.json());
+  app.use(createKeyRouter(
+    collection as unknown as Collection<KeyDocument>,
+    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<typeof makeCollection>;
+  let store: ReturnType<typeof makeStore>;
+  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");
+  });
+});