@aprovan/chat-backend 0.1.0-dev.93c7b6a → 0.1.0-dev.9c336a0

package/test/gateway-session.test.ts

@@ -3,6 +3,9 @@ import {
   getGatewaySession,
   evictGatewaySession,
   resetGatewaySessionCache,
+  getCachedTools,
+  setCachedTools,
+  evictCachedTools,
   GatewaySessionError,
 } from "../src/gateway-session";
 import type { CognitoAccessTokenPayload } from "aws-jwt-verify/jwt-model";
@@ -129,4 +132,80 @@ describe("getGatewaySession", () => {
 
     expect(fetchSpy).toHaveBeenCalledTimes(2);
   });
+
+  it("retries once after 200ms on 5xx and returns session if retry succeeds", async () => {
+    const fetchSpy = vi.spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response("Service Unavailable", { status: 503 }))
+      .mockResolvedValueOnce(
+        new Response(JSON.stringify({ workspace_id: WORKSPACE_ID }), { status: 200 }),
+      );
+
+    const session = await getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN);
+    expect(session.token).toBe(COGNITO_TOKEN);
+    expect(fetchSpy).toHaveBeenCalledTimes(2);
+  });
+
+  it("throws GatewaySessionError on persistent 5xx after retry", async () => {
+    vi.spyOn(globalThis, "fetch").mockResolvedValue(
+      new Response("Service Unavailable", { status: 503 }),
+    );
+
+    await expect(
+      getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN),
+    ).rejects.toThrow(GatewaySessionError);
+  });
+
+  it("retries once after 200ms on network error and returns session if retry succeeds", async () => {
+    const fetchSpy = vi.spyOn(globalThis, "fetch")
+      .mockRejectedValueOnce(new TypeError("fetch failed"))
+      .mockResolvedValueOnce(
+        new Response(JSON.stringify({ workspace_id: WORKSPACE_ID }), { status: 200 }),
+      );
+
+    const session = await getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN);
+    expect(session.token).toBe(COGNITO_TOKEN);
+    expect(fetchSpy).toHaveBeenCalledTimes(2);
+  });
+});
+
+describe("tool cache", () => {
+  beforeEach(() => {
+    resetGatewaySessionCache();
+  });
+
+  afterEach(() => {
+    resetGatewaySessionCache();
+  });
+
+  it("getCachedTools returns undefined when nothing cached", () => {
+    expect(getCachedTools("sub-123")).toBeUndefined();
+  });
+
+  it("setCachedTools stores tools and getCachedTools retrieves them", () => {
+    const tools = [{ name: "github_repos_list" }];
+    setCachedTools("sub-123", tools);
+    expect(getCachedTools("sub-123")).toEqual(tools);
+  });
+
+  it("evictCachedTools removes tools for a sub", () => {
+    setCachedTools("sub-123", [{ name: "github_repos_list" }]);
+    evictCachedTools("sub-123");
+    expect(getCachedTools("sub-123")).toBeUndefined();
+  });
+
+  it("evictGatewaySession also evicts the tool cache for that sub", async () => {
+    vi.spyOn(globalThis, "fetch").mockResolvedValue(
+      new Response(JSON.stringify({ workspace_id: WORKSPACE_ID }), { status: 200 }),
+    );
+    await getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN);
+    setCachedTools(USER_SUB, [{ name: "tool" }]);
+    evictGatewaySession(USER_SUB);
+    expect(getCachedTools(USER_SUB)).toBeUndefined();
+  });
+
+  it("resetGatewaySessionCache also clears the tool cache", () => {
+    setCachedTools("sub-a", [{ name: "tool" }]);
+    resetGatewaySessionCache();
+    expect(getCachedTools("sub-a")).toBeUndefined();
+  });
 });

package/test/middleware/workspace.test.ts

@@ -14,9 +14,10 @@ vi.mock("@aws-sdk/lib-dynamodb", () => ({
     from: vi.fn(() => ({ send: mockSend })),
   },
   QueryCommand: vi.fn((input) => input),
+  GetCommand: vi.fn((input) => input),
 }));
 
-const { workspaceMiddleware, resolveWorkspaceId } = await import(
+const { workspaceMiddleware, resolveWorkspaceId, evictWorkspaceCache } = await import(
   "../../src/middleware/workspace"
 );
 
@@ -41,15 +42,70 @@ describe("workspaceMiddleware", () => {
   beforeEach(() => {
     vi.clearAllMocks();
     process.env["MEMBERSHIPS_TABLE_NAME"] = "gateway-prd-use1-memberships";
+    process.env["USERS_TABLE_NAME"] = "gateway-prd-use1-users";
     process.env["AWS_REGION"] = "us-east-1";
-    // Clear the module-level cache between tests by resetting the module state.
-    // resolveWorkspaceId memoizes per sub; tests use unique subs or clear cache.
+    // Clear module-level cache between tests
+    evictWorkspaceCache("user-sub-123");
+    evictWorkspaceCache("user-sub-456");
+    evictWorkspaceCache("user-no-ws");
+    evictWorkspaceCache("user-cached");
+    evictWorkspaceCache("user-users-table");
+    evictWorkspaceCache("user-fallback");
   });
 
-  it("resolves workspaceId from DDB and sets it on context", async () => {
+  it("uses activeWorkspaceId from Users table when present", async () => {
+    // Users table returns activeWorkspaceId — no Memberships query needed
     mockSend.mockResolvedValueOnce({
-      Items: [{ workspaceId: "ws-abc", userSub: "user-sub-456" }],
+      Item: { sub: "user-users-table", activeWorkspaceId: "ws-from-users" },
+    });
+
+    const app = new Hono<{ Variables: AppVariables }>();
+    const localClaims = { sub: "user-users-table" } as unknown as CognitoAccessTokenPayload;
+    app.use("/protected", async (c, next) => {
+      c.set("claims", localClaims);
+      await next();
+    });
+    app.use("/protected", workspaceMiddleware);
+    app.get("/protected", (c) => c.json({ workspaceId: c.get("workspaceId") }));
+
+    const res = await app.request("/protected");
+    expect(res.status).toBe(200);
+    const body = await res.json();
+    expect(body).toEqual({ workspaceId: "ws-from-users" });
+    // Only one DDB call (Users Get), no Memberships query
+    expect(mockSend).toHaveBeenCalledTimes(1);
+  });
+
+  it("falls back to Memberships when Users table has no activeWorkspaceId", async () => {
+    // Users Get returns no item; Memberships query returns ws-from-memberships
+    mockSend
+      .mockResolvedValueOnce({}) // Users GetCommand → no item
+      .mockResolvedValueOnce({
+        Items: [{ workspaceId: "ws-from-memberships", userSub: "user-fallback" }],
+      }); // Memberships QueryCommand
+
+    const app = new Hono<{ Variables: AppVariables }>();
+    const localClaims = { sub: "user-fallback" } as unknown as CognitoAccessTokenPayload;
+    app.use("/protected", async (c, next) => {
+      c.set("claims", localClaims);
+      await next();
     });
+    app.use("/protected", workspaceMiddleware);
+    app.get("/protected", (c) => c.json({ workspaceId: c.get("workspaceId") }));
+
+    const res = await app.request("/protected");
+    expect(res.status).toBe(200);
+    const body = await res.json();
+    expect(body).toEqual({ workspaceId: "ws-from-memberships" });
+    expect(mockSend).toHaveBeenCalledTimes(2);
+  });
+
+  it("resolves workspaceId from DDB and sets it on context", async () => {
+    mockSend
+      .mockResolvedValueOnce({}) // Users GetCommand → no item
+      .mockResolvedValueOnce({
+        Items: [{ workspaceId: "ws-abc", userSub: "user-sub-456" }],
+      });
 
     const app = new Hono<{ Variables: AppVariables }>();
     const localClaims = { sub: "user-sub-456" } as unknown as CognitoAccessTokenPayload;
@@ -66,8 +122,10 @@ describe("workspaceMiddleware", () => {
     expect(body).toEqual({ workspaceId: "ws-abc" });
   });
 
-  it("returns 403 when no membership exists", async () => {
-    mockSend.mockResolvedValueOnce({ Items: [] });
+  it("returns 403 when no membership exists and Users table has no activeWorkspaceId", async () => {
+    mockSend
+      .mockResolvedValueOnce({}) // Users GetCommand → no item
+      .mockResolvedValueOnce({ Items: [] }); // Memberships QueryCommand → empty
 
     const app = new Hono<{ Variables: AppVariables }>();
     const localClaims = {
@@ -88,7 +146,7 @@ describe("workspaceMiddleware", () => {
 
   it("uses cache on subsequent calls for the same sub", async () => {
     mockSend.mockResolvedValueOnce({
-      Items: [{ workspaceId: "ws-cached", userSub: "user-cached" }],
+      Item: { activeWorkspaceId: "ws-cached" },
     });
 
     // First call — hits DDB
@@ -98,4 +156,17 @@ describe("workspaceMiddleware", () => {
 
     expect(mockSend).toHaveBeenCalledTimes(1);
   });
+
+  it("evictWorkspaceCache clears the cache so next call re-fetches", async () => {
+    mockSend
+      .mockResolvedValueOnce({ Item: { activeWorkspaceId: "ws-v1" } })
+      .mockResolvedValueOnce({ Item: { activeWorkspaceId: "ws-v2" } });
+
+    await resolveWorkspaceId("user-cached");
+    evictWorkspaceCache("user-cached");
+    const result = await resolveWorkspaceId("user-cached");
+
+    expect(mockSend).toHaveBeenCalledTimes(2);
+    expect(result).toBe("ws-v2");
+  });
 });

package/test/routes/chat.test.ts

@@ -12,6 +12,8 @@ const {
   mockProviderFactory,
   mockGetGatewaySession,
   mockEvictGatewaySession,
+  mockGetCachedTools,
+  mockSetCachedTools,
   mockGetPrompt,
   mockCompilePrompt,
   mockGetPostHogClient,
@@ -25,6 +27,8 @@ const {
     mockProviderFactory,
     mockGetGatewaySession: vi.fn(),
     mockEvictGatewaySession: vi.fn(),
+    mockGetCachedTools: vi.fn().mockReturnValue(undefined),
+    mockSetCachedTools: vi.fn(),
     mockGetPrompt: vi.fn().mockResolvedValue({
       source: "code_fallback",
       prompt: "System: {{compilers}} {{tool_docs}}",
@@ -51,6 +55,8 @@ vi.mock("../../src/providers/openrouter.js", () => ({
 vi.mock("../../src/gateway-session.js", () => ({
   getGatewaySession: mockGetGatewaySession,
   evictGatewaySession: mockEvictGatewaySession,
+  getCachedTools: mockGetCachedTools,
+  setCachedTools: mockSetCachedTools,
   GatewaySessionError: class GatewaySessionError extends Error {
     status: number;
     constructor(status: number, message: string) {
@@ -182,6 +188,8 @@ describe("POST /chat", () => {
     mockProviderFactory.mockReturnValue(mockModel);
     mockCreateOpenRouterProvider.mockReturnValue(mockProviderFactory);
     mockGetOpenRouterKey.mockResolvedValue("test-key");
+    mockGetCachedTools.mockReturnValue(undefined);
+    mockSetCachedTools.mockReset();
 
     // Reset prompt + tool-docs mocks to defaults
     mockGetPrompt.mockResolvedValue({
@@ -299,6 +307,55 @@ describe("POST /chat", () => {
     expect(mockDoStream).toHaveBeenCalledTimes(1);
   });
 
+  it("returns 403 when requested model is not in workspace.limits.maxModels", async () => {
+    const app = buildApp();
+    const res = await app.request("/chat", {
+      method: "POST",
+      headers: validHeaders,
+      body: JSON.stringify({
+        id: "chat-1",
+        messages: [{ role: "user", parts: [{ type: "text", text: "Hello" }], id: "msg-1" }],
+        trigger: "submit-message",
+        model: "anthropic/claude-opus-4",
+      }),
+    });
+    expect(res.status).toBe(403);
+    const body = await res.json();
+    expect(body).toMatchObject({ error: expect.stringContaining("Model not allowed") });
+    expect(mockDoStream).not.toHaveBeenCalled();
+  });
+
+  it("uses the requested model when it is in workspace.limits.maxModels", async () => {
+    mockDoStream.mockResolvedValueOnce({
+      stream: makeSuccessStream(),
+      rawResponse: { headers: {} },
+    });
+
+    const proWorkspace: WorkspaceItem = {
+      ...fakeWorkspace,
+      plan: "pro",
+      limits: {
+        ...fakeWorkspace.limits,
+        maxModels: ["openrouter/auto", "anthropic/claude-opus-4"],
+      },
+    };
+
+    const app = buildApp(proWorkspace);
+    const res = await app.request("/chat", {
+      method: "POST",
+      headers: validHeaders,
+      body: JSON.stringify({
+        id: "chat-1",
+        messages: [{ role: "user", parts: [{ type: "text", text: "Hello" }], id: "msg-1" }],
+        trigger: "submit-message",
+        model: "anthropic/claude-opus-4",
+      }),
+    });
+
+    expect(res.status).toBe(200);
+    expect(mockProviderFactory).toHaveBeenCalledWith("anthropic/claude-opus-4");
+  });
+
   it("selects model from workspace.limits.maxModels[0]", async () => {
     mockDoStream.mockResolvedValueOnce({
       stream: makeSuccessStream(),
@@ -484,5 +541,47 @@ describe("POST /chat", () => {
         expect.any(String),
       );
     });
+
+    it("uses cached tools when available — skips gateway fetch", async () => {
+      mockGetCachedTools.mockReturnValue(MOCK_TOOLS_RESPONSE.tools);
+      mockDoStream.mockResolvedValueOnce({
+        stream: makeSuccessStream(),
+        rawResponse: { headers: {} },
+      });
+
+      const app = buildApp();
+      const res = await app.request("/chat", {
+        method: "POST",
+        headers: validHeaders,
+        body: validBody,
+      });
+
+      expect(res.status).toBe(200);
+      // No gateway tools fetch should have occurred
+      expect(mockFetch).not.toHaveBeenCalled();
+    });
+
+    it("caches fetched tools via setCachedTools", async () => {
+      mockFetch.mockResolvedValueOnce(
+        new Response(JSON.stringify(MOCK_TOOLS_RESPONSE), { status: 200 }),
+      );
+      mockDoStream.mockResolvedValueOnce({
+        stream: makeSuccessStream(),
+        rawResponse: { headers: {} },
+      });
+
+      const app = buildApp();
+      await app.request("/chat", {
+        method: "POST",
+        headers: validHeaders,
+        body: validBody,
+      });
+
+      expect(mockSetCachedTools).toHaveBeenCalledOnce();
+      expect(mockSetCachedTools).toHaveBeenCalledWith(
+        fakeClaims.sub,
+        MOCK_TOOLS_RESPONSE.tools,
+      );
+    });
   });
 });