@aprovan/chat-backend 0.1.0-dev.4d82df8 → 0.1.0-dev.879ed82

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

@@ -307,55 +307,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(),

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

package/test/routes/workspaces.test.ts

@@ -1,164 +0,0 @@
-import { describe, it, expect, vi, beforeEach } from "vitest";
-import { Hono } from "hono";
-import type { AppVariables } from "../../src/types";
-import type { CognitoAccessTokenPayload } from "aws-jwt-verify/jwt-model";
-
-vi.mock("@aws-sdk/client-dynamodb", () => ({
-  DynamoDBClient: vi.fn(() => ({})),
-}));
-
-const mockSend = vi.fn();
-vi.mock("@aws-sdk/lib-dynamodb", () => ({
-  DynamoDBDocumentClient: {
-    from: vi.fn(() => ({ send: mockSend })),
-  },
-  QueryCommand: vi.fn((input) => input),
-  GetCommand: vi.fn((input) => input),
-  PutCommand: vi.fn((input) => input),
-  BatchGetCommand: vi.fn((input) => input),
-}));
-
-const { workspacesRoute } = await import("../../src/routes/workspaces");
-const { resetMembershipCache } = await import("../../src/middleware/workspace");
-const { resetSessionCache } = await import("../../src/session");
-
-const MEMBERSHIPS_TABLE = "test-memberships";
-const WORKSPACE_TABLE = "test-workspaces";
-const SESSIONS_TABLE = "test-user-sessions";
-
-type MockCommand = { TableName?: string; RequestItems?: Record<string, unknown> };
-
-const fakeClaims = { sub: "user-ws-test" } as unknown as CognitoAccessTokenPayload;
-
-function buildApp() {
-  const app = new Hono<{ Variables: AppVariables }>();
-  app.use("/*", async (c, next) => {
-    c.set("claims", fakeClaims);
-    await next();
-  });
-  app.route("/workspaces", workspacesRoute);
-  return app;
-}
-
-describe("GET /workspaces", () => {
-  beforeEach(() => {
-    vi.clearAllMocks();
-    resetMembershipCache();
-    resetSessionCache();
-    process.env["MEMBERSHIPS_TABLE_NAME"] = MEMBERSHIPS_TABLE;
-    process.env["WORKSPACE_TABLE_NAME"] = WORKSPACE_TABLE;
-    process.env["USER_SESSIONS_TABLE_NAME"] = SESSIONS_TABLE;
-    process.env["AWS_REGION"] = "us-east-1";
-  });
-
-  it("returns workspace list with active flag", async () => {
-    mockSend.mockImplementation((cmd: MockCommand) => {
-      if (cmd.TableName === SESSIONS_TABLE) {
-        return Promise.resolve({ Item: { activeWorkspaceId: "ws-b" } });
-      }
-      if (cmd.TableName === MEMBERSHIPS_TABLE) {
-        return Promise.resolve({ Items: [{ workspaceId: "ws-a" }, { workspaceId: "ws-b" }] });
-      }
-      if (cmd.RequestItems?.[WORKSPACE_TABLE]) {
-        return Promise.resolve({
-          Responses: {
-            [WORKSPACE_TABLE]: [
-              { workspaceId: "ws-a", name: "Alpha", plan: "free" },
-              { workspaceId: "ws-b", name: "Beta", plan: "pro" },
-            ],
-          },
-        });
-      }
-      return Promise.resolve({});
-    });
-
-    const app = buildApp();
-    const res = await app.request("/workspaces");
-    expect(res.status).toBe(200);
-    const body = await res.json();
-    expect(body.activeWorkspaceId).toBe("ws-b");
-    const beta = body.workspaces.find((w: { workspaceId: string }) => w.workspaceId === "ws-b");
-    expect(beta?.active).toBe(true);
-  });
-
-  it("falls back to first membership when no session preference", async () => {
-    mockSend.mockImplementation((cmd: MockCommand) => {
-      if (cmd.TableName === SESSIONS_TABLE) {
-        return Promise.resolve({ Item: undefined });
-      }
-      if (cmd.TableName === MEMBERSHIPS_TABLE) {
-        return Promise.resolve({ Items: [{ workspaceId: "ws-a" }] });
-      }
-      if (cmd.RequestItems?.[WORKSPACE_TABLE]) {
-        return Promise.resolve({
-          Responses: { [WORKSPACE_TABLE]: [{ workspaceId: "ws-a", name: "Alpha", plan: "free" }] },
-        });
-      }
-      return Promise.resolve({});
-    });
-
-    const app = buildApp();
-    const res = await app.request("/workspaces");
-    const body = await res.json();
-    expect(body.activeWorkspaceId).toBe("ws-a");
-    expect(body.workspaces[0].active).toBe(true);
-  });
-});
-
-describe("PUT /workspaces/active", () => {
-  beforeEach(() => {
-    vi.clearAllMocks();
-    resetMembershipCache();
-    resetSessionCache();
-    process.env["MEMBERSHIPS_TABLE_NAME"] = MEMBERSHIPS_TABLE;
-    process.env["WORKSPACE_TABLE_NAME"] = WORKSPACE_TABLE;
-    process.env["USER_SESSIONS_TABLE_NAME"] = SESSIONS_TABLE;
-    process.env["AWS_REGION"] = "us-east-1";
-  });
-
-  it("sets active workspace and returns it", async () => {
-    mockSend.mockImplementation((cmd: MockCommand) => {
-      if (cmd.TableName === MEMBERSHIPS_TABLE) {
-        return Promise.resolve({ Items: [{ workspaceId: "ws-a" }, { workspaceId: "ws-b" }] });
-      }
-      return Promise.resolve({});
-    });
-
-    const app = buildApp();
-    const res = await app.request("/workspaces/active", {
-      method: "PUT",
-      headers: { "Content-Type": "application/json" },
-      body: JSON.stringify({ workspaceId: "ws-b" }),
-    });
-    expect(res.status).toBe(200);
-    const body = await res.json();
-    expect(body.activeWorkspaceId).toBe("ws-b");
-  });
-
-  it("returns 403 when workspace is not in user memberships", async () => {
-    mockSend.mockImplementation((cmd: MockCommand) => {
-      if (cmd.TableName === MEMBERSHIPS_TABLE) {
-        return Promise.resolve({ Items: [{ workspaceId: "ws-a" }] });
-      }
-      return Promise.resolve({});
-    });
-
-    const app = buildApp();
-    const res = await app.request("/workspaces/active", {
-      method: "PUT",
-      headers: { "Content-Type": "application/json" },
-      body: JSON.stringify({ workspaceId: "ws-other" }),
-    });
-    expect(res.status).toBe(403);
-  });
-
-  it("returns 400 for invalid request body", async () => {
-    const app = buildApp();
-    const res = await app.request("/workspaces/active", {
-      method: "PUT",
-      headers: { "Content-Type": "application/json" },
-      body: JSON.stringify({ workspaceId: "" }),
-    });
-    expect(res.status).toBe(400);
-  });
-});

package/test/session.test.ts

@@ -1,56 +0,0 @@
-import { describe, it, expect, vi, beforeEach } from "vitest";
-
-vi.mock("@aws-sdk/client-dynamodb", () => ({
-  DynamoDBClient: vi.fn(() => ({})),
-}));
-
-const mockSend = vi.fn();
-vi.mock("@aws-sdk/lib-dynamodb", () => ({
-  DynamoDBDocumentClient: {
-    from: vi.fn(() => ({ send: mockSend })),
-  },
-  GetCommand: vi.fn((input) => input),
-  PutCommand: vi.fn((input) => input),
-}));
-
-const { getSessionWorkspaceId, setSessionWorkspaceId, resetSessionCache } =
-  await import("../src/session");
-
-describe("session", () => {
-  beforeEach(() => {
-    vi.clearAllMocks();
-    resetSessionCache();
-    process.env["USER_SESSIONS_TABLE_NAME"] = "test-user-sessions";
-    process.env["AWS_REGION"] = "us-east-1";
-  });
-
-  it("returns null when no session record exists", async () => {
-    mockSend.mockResolvedValueOnce({ Item: undefined });
-    const result = await getSessionWorkspaceId("user-abc");
-    expect(result).toBeNull();
-  });
-
-  it("returns activeWorkspaceId from session record", async () => {
-    mockSend.mockResolvedValueOnce({ Item: { userSub: "user-abc", activeWorkspaceId: "ws-xyz" } });
-    const result = await getSessionWorkspaceId("user-abc");
-    expect(result).toBe("ws-xyz");
-  });
-
-  it("caches the result and skips DDB on second call", async () => {
-    mockSend.mockResolvedValueOnce({ Item: { activeWorkspaceId: "ws-cached" } });
-    await getSessionWorkspaceId("user-cached");
-    await getSessionWorkspaceId("user-cached");
-    expect(mockSend).toHaveBeenCalledTimes(1);
-  });
-
-  it("setSessionWorkspaceId writes to DDB and updates cache", async () => {
-    mockSend.mockResolvedValueOnce({});
-    await setSessionWorkspaceId("user-set", "ws-new");
-    expect(mockSend).toHaveBeenCalledTimes(1);
-
-    // Second read should come from cache, not DDB
-    const result = await getSessionWorkspaceId("user-set");
-    expect(result).toBe("ws-new");
-    expect(mockSend).toHaveBeenCalledTimes(1);
-  });
-});