@aprovan/chat-backend 0.1.0-dev.4d82df8 → 0.1.0-dev.78c0b14

package/test/app.test.ts

@@ -20,8 +20,6 @@ vi.mock("@aws-sdk/lib-dynamodb", () => ({
   },
   QueryCommand: vi.fn((input) => input),
   GetCommand: vi.fn((input) => input),
-  PutCommand: vi.fn((input) => input),
-  BatchGetCommand: vi.fn((input) => input),
 }));
 
 beforeAll(() => {
@@ -30,7 +28,6 @@ beforeAll(() => {
   process.env["AWS_REGION"] = "us-east-1";
   process.env["WORKSPACE_TABLE_NAME"] = "test-workspaces";
   process.env["MEMBERSHIPS_TABLE_NAME"] = "test-memberships";
-  process.env["USER_SESSIONS_TABLE_NAME"] = "test-user-sessions";
 });
 
 describe("chat app", () => {

package/test/gateway-session.test.ts

@@ -3,9 +3,6 @@ import {
   getGatewaySession,
   evictGatewaySession,
   resetGatewaySessionCache,
-  getCachedTools,
-  setCachedTools,
-  evictCachedTools,
   GatewaySessionError,
 } from "../src/gateway-session";
 import type { CognitoAccessTokenPayload } from "aws-jwt-verify/jwt-model";
@@ -132,80 +129,4 @@ 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

@@ -3,6 +3,7 @@ import { Hono } from "hono";
 import type { AppVariables } from "../../src/types";
 import type { CognitoAccessTokenPayload } from "aws-jwt-verify/jwt-model";
 
+// Mock the DDB client before importing the middleware
 vi.mock("@aws-sdk/client-dynamodb", () => ({
   DynamoDBClient: vi.fn(() => ({})),
 }));
@@ -13,110 +14,88 @@ vi.mock("@aws-sdk/lib-dynamodb", () => ({
     from: vi.fn(() => ({ send: mockSend })),
   },
   QueryCommand: vi.fn((input) => input),
-  GetCommand: vi.fn((input) => input),
-  PutCommand: vi.fn((input) => input),
 }));
 
-const {
-  workspaceMiddleware,
-  listWorkspaceMemberships,
-  resetMembershipCache,
-} = await import("../../src/middleware/workspace");
-const { resetSessionCache } = await import("../../src/session");
+const { workspaceMiddleware, resolveWorkspaceId } = await import(
+  "../../src/middleware/workspace"
+);
 
-const MEMBERSHIPS_TABLE = "gateway-prd-use1-memberships";
-const SESSIONS_TABLE = "test-user-sessions";
+const fakeClaims = {
+  sub: "user-sub-123",
+} as unknown as CognitoAccessTokenPayload;
 
-type MockCommand = { TableName?: string };
-
-function buildApp(userSub: string) {
-  const fakeClaims = { sub: userSub } as unknown as CognitoAccessTokenPayload;
+function buildApp() {
   const app = new Hono<{ Variables: AppVariables }>();
   app.use("/protected", async (c, next) => {
     c.set("claims", fakeClaims);
     await next();
   });
   app.use("/protected", workspaceMiddleware);
-  app.get("/protected", (c) => c.json({ workspaceId: c.get("workspaceId") }));
+  app.get("/protected", (c) =>
+    c.json({ workspaceId: c.get("workspaceId") }),
+  );
   return app;
 }
 
 describe("workspaceMiddleware", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    resetMembershipCache();
-    resetSessionCache();
-    process.env["MEMBERSHIPS_TABLE_NAME"] = MEMBERSHIPS_TABLE;
-    process.env["USER_SESSIONS_TABLE_NAME"] = SESSIONS_TABLE;
+    process.env["MEMBERSHIPS_TABLE_NAME"] = "gateway-prd-use1-memberships";
     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.
   });
 
-  it("uses session activeWorkspaceId when membership is valid", async () => {
-    mockSend.mockImplementation((cmd: MockCommand) => {
-      if (cmd.TableName === SESSIONS_TABLE) {
-        return Promise.resolve({ Item: { activeWorkspaceId: "ws-b" } });
-      }
-      return Promise.resolve({ Items: [{ workspaceId: "ws-a" }, { workspaceId: "ws-b" }] });
+  it("resolves workspaceId from DDB and sets it on context", async () => {
+    mockSend.mockResolvedValueOnce({
+      Items: [{ workspaceId: "ws-abc", userSub: "user-sub-456" }],
     });
 
-    const app = buildApp("user-session");
-    const res = await app.request("/protected");
-    expect(res.status).toBe(200);
-    const body = await res.json();
-    expect(body.workspaceId).toBe("ws-b");
-  });
-
-  it("falls back to first membership when no session preference exists", async () => {
-    mockSend.mockImplementation((cmd: MockCommand) => {
-      if (cmd.TableName === SESSIONS_TABLE) {
-        return Promise.resolve({ Item: undefined });
-      }
-      return Promise.resolve({ Items: [{ workspaceId: "ws-first" }] });
+    const app = new Hono<{ Variables: AppVariables }>();
+    const localClaims = { sub: "user-sub-456" } 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 app = buildApp("user-no-session");
     const res = await app.request("/protected");
     expect(res.status).toBe(200);
     const body = await res.json();
-    expect(body.workspaceId).toBe("ws-first");
+    expect(body).toEqual({ workspaceId: "ws-abc" });
   });
 
-  it("falls back to first membership when session workspace is not in memberships", async () => {
-    mockSend.mockImplementation((cmd: MockCommand) => {
-      if (cmd.TableName === SESSIONS_TABLE) {
-        return Promise.resolve({ Item: { activeWorkspaceId: "ws-stale" } });
-      }
-      return Promise.resolve({ Items: [{ workspaceId: "ws-a" }] });
-    });
-
-    const app = buildApp("user-stale");
-    const res = await app.request("/protected");
-    expect(res.status).toBe(200);
-    const body = await res.json();
-    expect(body.workspaceId).toBe("ws-a");
-  });
+  it("returns 403 when no membership exists", async () => {
+    mockSend.mockResolvedValueOnce({ Items: [] });
 
-  it("returns 403 when user has no workspace memberships", async () => {
-    mockSend.mockImplementation((cmd: MockCommand) => {
-      if (cmd.TableName === SESSIONS_TABLE) {
-        return Promise.resolve({ Item: undefined });
-      }
-      return Promise.resolve({ Items: [] });
+    const app = new Hono<{ Variables: AppVariables }>();
+    const localClaims = {
+      sub: "user-no-ws",
+    } 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 app = buildApp("user-no-ws");
     const res = await app.request("/protected");
     expect(res.status).toBe(403);
     const body = await res.json();
     expect(body).toMatchObject({ error: "No workspace membership" });
   });
 
-  it("uses membership cache on subsequent calls for the same sub", async () => {
-    mockSend.mockResolvedValue({ Items: [{ workspaceId: "ws-cached" }] });
+  it("uses cache on subsequent calls for the same sub", async () => {
+    mockSend.mockResolvedValueOnce({
+      Items: [{ workspaceId: "ws-cached", userSub: "user-cached" }],
+    });
+
+    // First call — hits DDB
+    await resolveWorkspaceId("user-cached");
+    // Second call — should use cache
+    await resolveWorkspaceId("user-cached");
 
-    await listWorkspaceMemberships("user-cached-q");
-    await listWorkspaceMemberships("user-cached-q");
-    // Only one DDB send for memberships (cache hit on second call)
     expect(mockSend).toHaveBeenCalledTimes(1);
   });
 });

package/test/routes/chat.test.ts

@@ -12,8 +12,6 @@ const {
   mockProviderFactory,
   mockGetGatewaySession,
   mockEvictGatewaySession,
-  mockGetCachedTools,
-  mockSetCachedTools,
   mockGetPrompt,
   mockCompilePrompt,
   mockGetPostHogClient,
@@ -27,8 +25,6 @@ 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}}",
@@ -55,8 +51,6 @@ 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) {
@@ -188,8 +182,6 @@ 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({
@@ -307,55 +299,6 @@ 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(),
@@ -541,47 +484,5 @@ 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,
-      );
-    });
   });
 });

package/src/routes/workspaces.ts

@@ -1,89 +0,0 @@
-import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
-import { BatchGetCommand, DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb";
-import { zValidator } from "@hono/zod-validator";
-import { Hono } from "hono";
-import { z } from "zod";
-import { listWorkspaceMemberships } from "../middleware/workspace.js";
-import { getSessionWorkspaceId, setSessionWorkspaceId } from "../session.js";
-import type { AppVariables, WorkspaceItem } from "../types.js";
-
-let ddbClient: DynamoDBDocumentClient | null = null;
-
-function getDdb() {
-  if (!ddbClient) {
-    ddbClient = DynamoDBDocumentClient.from(
-      new DynamoDBClient({ region: process.env["AWS_REGION"] }),
-    );
-  }
-  return ddbClient;
-}
-
-async function batchGetWorkspaces(
-  workspaceIds: string[],
-): Promise<WorkspaceItem[]> {
-  if (workspaceIds.length === 0) return [];
-  const tableName = process.env["WORKSPACE_TABLE_NAME"]!;
-  const result = await getDdb().send(
-    new BatchGetCommand({
-      RequestItems: {
-        [tableName]: { Keys: workspaceIds.map((id) => ({ workspaceId: id })) },
-      },
-    }),
-  );
-  return (result.Responses?.[tableName] ?? []) as WorkspaceItem[];
-}
-
-export const workspacesRoute = new Hono<{ Variables: AppVariables }>();
-
-workspacesRoute.get("/", async (c) => {
-  const claims = c.get("claims");
-  const userSub = claims.sub;
-
-  const [workspaceIds, activeWorkspaceId] = await Promise.all([
-    listWorkspaceMemberships(userSub),
-    getSessionWorkspaceId(userSub),
-  ]);
-
-  const workspaces = await batchGetWorkspaces(workspaceIds);
-
-  const effectiveActiveId =
-    activeWorkspaceId && workspaceIds.includes(activeWorkspaceId)
-      ? activeWorkspaceId
-      : (workspaceIds[0] ?? null);
-
-  const sorted = workspaceIds
-    .map((id) => workspaces.find((w) => w.workspaceId === id))
-    .filter((w): w is WorkspaceItem => w !== undefined);
-
-  return c.json({
-    workspaces: sorted.map((w) => ({
-      workspaceId: w.workspaceId,
-      name: w.name,
-      plan: w.plan,
-      active: w.workspaceId === effectiveActiveId,
-    })),
-    activeWorkspaceId: effectiveActiveId,
-  });
-});
-
-const setActiveSchema = z.object({
-  workspaceId: z.string().min(1),
-});
-
-workspacesRoute.put(
-  "/active",
-  zValidator("json", setActiveSchema),
-  async (c) => {
-    const claims = c.get("claims");
-    const userSub = claims.sub;
-    const { workspaceId } = c.req.valid("json");
-
-    const workspaceIds = await listWorkspaceMemberships(userSub);
-    if (!workspaceIds.includes(workspaceId)) {
-      return c.json({ error: "Not a member of this workspace" }, 403);
-    }
-
-    await setSessionWorkspaceId(userSub, workspaceId);
-    return c.json({ activeWorkspaceId: workspaceId });
-  },
-);

package/src/session.ts

@@ -1,80 +0,0 @@
-import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
-import {
-  DynamoDBDocumentClient,
-  GetCommand,
-  PutCommand,
-} from "@aws-sdk/lib-dynamodb";
-
-const SESSION_CACHE_TTL_MS = 15_000;
-
-interface SessionCacheEntry {
-  activeWorkspaceId: string | null;
-  fetchedAt: number;
-}
-
-const sessionCache = new Map<string, SessionCacheEntry>();
-
-let ddbClient: DynamoDBDocumentClient | null = null;
-
-function getDdb() {
-  if (!ddbClient) {
-    ddbClient = DynamoDBDocumentClient.from(
-      new DynamoDBClient({ region: process.env["AWS_REGION"] }),
-    );
-  }
-  return ddbClient;
-}
-
-function tableName(): string {
-  return process.env["USER_SESSIONS_TABLE_NAME"]!;
-}
-
-export async function getSessionWorkspaceId(
-  userSub: string,
-): Promise<string | null> {
-  const now = Date.now();
-  const cached = sessionCache.get(userSub);
-  if (cached && now - cached.fetchedAt < SESSION_CACHE_TTL_MS) {
-    return cached.activeWorkspaceId;
-  }
-
-  const result = await getDdb().send(
-    new GetCommand({
-      TableName: tableName(),
-      Key: { userSub },
-    }),
-  );
-
-  const activeWorkspaceId =
-    (result.Item?.activeWorkspaceId as string | undefined) ?? null;
-  sessionCache.set(userSub, { activeWorkspaceId, fetchedAt: now });
-  return activeWorkspaceId;
-}
-
-export async function setSessionWorkspaceId(
-  userSub: string,
-  workspaceId: string,
-): Promise<void> {
-  await getDdb().send(
-    new PutCommand({
-      TableName: tableName(),
-      Item: {
-        userSub,
-        activeWorkspaceId: workspaceId,
-        updatedAt: new Date().toISOString(),
-      },
-    }),
-  );
-  sessionCache.set(userSub, {
-    activeWorkspaceId: workspaceId,
-    fetchedAt: Date.now(),
-  });
-}
-
-export function clearSessionCache(userSub: string): void {
-  sessionCache.delete(userSub);
-}
-
-export function resetSessionCache(): void {
-  sessionCache.clear();
-}