feat(harness): secure session model configuration

2026-09-07 15:55:30 +08:00
parent 85739f58ed
commit f9ea4ce429
9 changed files with 511 additions and 1 deletions
@@ -0,0 +1,106 @@
+import { describe, expect, it } from "vitest";
+import * as harness from "../src/index.js";
+
+type Input = {
+  providerId: string;
+  api: "openai-responses" | "openai-completions" | "anthropic-messages";
+  apiKey?: string;
+  baseUrl?: string;
+  modelId: string;
+  displayName?: string;
+  supportsImages: boolean;
+  contextWindow?: number;
+  maxTokens?: number;
+};
+type Resolved = { connection: Input; apiKey?: string; source: "session" | "environment" };
+type Store = {
+  configure(sessionId: string, input: Input): Promise<void>;
+  getRedacted(sessionId: string): Record<string, unknown>;
+  resolveConnection(sessionId: string): Resolved;
+  clearCredentials(sessionId: string): void;
+};
+
+function input(overrides: Partial<Input> = {}): Input {
+  return {
+    providerId: "training-openai",
+    api: "openai-responses",
+    apiKey: "sk-test-secret",
+    baseUrl: "https://api.example.com/v1",
+    modelId: "training-model",
+    displayName: "Training Model",
+    supportsImages: true,
+    contextWindow: 32_000,
+    maxTokens: 4_096,
+    ...overrides,
+  };
+}
+
+function createStore(environment?: Input): Store {
+  const Constructor = Reflect.get(harness, "InMemoryModelConnectionStore") as unknown as
+    | (new (options?: Record<string, unknown>) => Store)
+    | undefined;
+  expect(Constructor, "InMemoryModelConnectionStore must be exported").toBeDefined();
+  return new Constructor!({
+    ...(environment === undefined ? {} : { environment }),
+    validateUrl: (url: string) => Promise.resolve({ normalizedUrl: url, addresses: ["93.184.216.34"], redirectPolicy: "same-origin" }),
+  });
+}
+
+describe("InMemoryModelConnectionStore", () => {
+  it("never returns the raw key", async () => {
+    const store = createStore();
+    await store.configure("ses_1", input());
+
+    const json = JSON.stringify(store.getRedacted("ses_1"));
+    expect(json).not.toContain("sk-test-secret");
+    expect(store.getRedacted("ses_1").keyHint).toBe("••••cret");
+  });
+
+  it("uses session config before environment and falls back after an idempotent clear", async () => {
+    const store = createStore(input({ providerId: "env-provider", apiKey: "env-secret", modelId: "env-model" }));
+    await store.configure("ses_1", input({ apiKey: "session-secret" }));
+    expect(store.resolveConnection("ses_1")).toMatchObject({ source: "session", apiKey: "session-secret" });
+
+    store.clearCredentials("ses_1");
+    store.clearCredentials("ses_1");
+    expect(store.resolveConnection("ses_1")).toMatchObject({ source: "environment", apiKey: "env-secret" });
+  });
+
+  it("rotates keys without exposing the old credential", async () => {
+    const store = createStore();
+    await store.configure("ses_1", input({ apiKey: "first-secret" }));
+    await store.configure("ses_1", input({ apiKey: "second-secret" }));
+
+    expect(store.resolveConnection("ses_1").apiKey).toBe("second-secret");
+    expect(JSON.stringify(store.getRedacted("ses_1"))).not.toContain("first-secret");
+  });
+
+  it("fails explicitly when no connection is configured", () => {
+    const store = createStore();
+    expect(() => store.resolveConnection("ses_missing")).toThrowError(
+      expect.objectContaining({ code: "MODEL_NOT_CONFIGURED" }),
+    );
+  });
+
+  it.each([
+    input({ providerId: "anthropic", api: "openai-responses" }),
+    input({ providerId: "openai", api: "anthropic-messages" }),
+  ])("rejects a built-in provider/API mismatch", async (invalid) => {
+    await expect(createStore().configure("ses_1", invalid)).rejects.toMatchObject({
+      code: "MODEL_API_MISMATCH",
+    });
+  });
+
+  it("creates a Pi runtime with only the configured model and in-memory credentials", async () => {
+    const createRuntime = Reflect.get(harness, "createPiModelRuntime") as
+      | ((resolved: Resolved) => Promise<{ getModel(providerId: string, modelId: string): unknown; getAuth(providerId: string): Promise<unknown> }>)
+      | undefined;
+    expect(createRuntime, "createPiModelRuntime must be exported").toBeDefined();
+    const resolved = { connection: input(), apiKey: "runtime-secret", source: "session" as const };
+
+    const runtime = await createRuntime!(resolved);
+
+    expect(runtime.getModel("training-openai", "training-model")).toBeDefined();
+    expect(await runtime.getAuth("training-openai")).toMatchObject({ auth: { apiKey: "runtime-secret" } });
+  });
+});