@aprovan/chat-backend 0.1.0-dev.4d82df8

package/src/types.ts

@@ -0,0 +1,29 @@
+import type { CognitoAccessTokenPayload } from "aws-jwt-verify/jwt-model";
+
+export interface WorkspaceLimits {
+  dailyChatCap: number;
+  maxModels: string[];
+  maxToolSteps: number;
+  maxTokensPerRequest: number;
+}
+
+export interface WorkspaceFeatures {
+  advancedTools: boolean;
+  customPrompts: boolean;
+}
+
+export interface WorkspaceItem {
+  workspaceId: string;
+  name: string;
+  plan: string;
+  limits: WorkspaceLimits;
+  features: WorkspaceFeatures;
+  createdAt: string;
+  updatedAt: string;
+}
+
+export type AppVariables = {
+  claims: CognitoAccessTokenPayload;
+  workspaceId: string;
+  workspace: WorkspaceItem;
+};

package/test/app.test.ts

@@ -0,0 +1,62 @@
+import { describe, it, expect, vi, beforeAll } from "vitest";
+import { createChatApp } from "../src/app";
+
+vi.mock("aws-jwt-verify", () => {
+  const verify = vi.fn();
+  return {
+    CognitoJwtVerifier: {
+      create: vi.fn(() => ({ verify })),
+    },
+  };
+});
+
+vi.mock("@aws-sdk/client-dynamodb", () => ({
+  DynamoDBClient: vi.fn(() => ({})),
+}));
+
+vi.mock("@aws-sdk/lib-dynamodb", () => ({
+  DynamoDBDocumentClient: {
+    from: vi.fn(() => ({ send: vi.fn() })),
+  },
+  QueryCommand: vi.fn((input) => input),
+  GetCommand: vi.fn((input) => input),
+  PutCommand: vi.fn((input) => input),
+  BatchGetCommand: vi.fn((input) => input),
+}));
+
+beforeAll(() => {
+  process.env["COGNITO_USER_POOL_ID"] = "us-east-1_test";
+  process.env["COGNITO_CLIENT_ID"] = "test-client-id";
+  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", () => {
+  const app = createChatApp();
+
+  it("GET /health returns ok without auth", async () => {
+    const res = await app.request("/health");
+    expect(res.status).toBe(200);
+    const body = await res.json();
+    expect(body).toEqual({ status: "ok" });
+  });
+
+  it("GET /api/* returns 401 without Authorization header", async () => {
+    const res = await app.request("/api/chat");
+    expect(res.status).toBe(401);
+  });
+
+  it("GET /api/* returns 401 with invalid Bearer token", async () => {
+    const { CognitoJwtVerifier } = await import("aws-jwt-verify");
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    (CognitoJwtVerifier.create({} as any) as any).verify.mockRejectedValueOnce(
+      new Error("invalid token"),
+    );
+    const res = await app.request("/api/chat", {
+      headers: { Authorization: "Bearer bad.token" },
+    });
+    expect(res.status).toBe(401);
+  });
+});

package/test/gateway-session.test.ts

@@ -0,0 +1,211 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import {
+  getGatewaySession,
+  evictGatewaySession,
+  resetGatewaySessionCache,
+  getCachedTools,
+  setCachedTools,
+  evictCachedTools,
+  GatewaySessionError,
+} from "../src/gateway-session";
+import type { CognitoAccessTokenPayload } from "aws-jwt-verify/jwt-model";
+
+const GATEWAY_URL = "https://gateway.test";
+const WORKSPACE_ID = "ws-abc";
+const COGNITO_TOKEN = "cognito-access-token";
+const USER_SUB = "user-sub-123";
+const TOKEN_EXP = Math.floor(Date.now() / 1000) + 3600;
+
+function makeClaims(overrides?: Partial<CognitoAccessTokenPayload>): CognitoAccessTokenPayload {
+  return {
+    sub: USER_SUB,
+    exp: TOKEN_EXP,
+    iss: "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_test",
+    client_id: "test-client",
+    username: "testuser",
+    token_use: "access",
+    iat: Math.floor(Date.now() / 1000),
+    jti: "jti-test",
+    auth_time: Math.floor(Date.now() / 1000),
+    scope: "",
+    version: 2,
+    origin_jti: "",
+    ...overrides,
+  } as unknown as CognitoAccessTokenPayload;
+}
+
+beforeEach(() => {
+  resetGatewaySessionCache();
+  process.env["GATEWAY_URL"] = GATEWAY_URL;
+});
+
+afterEach(() => {
+  resetGatewaySessionCache();
+  vi.restoreAllMocks();
+});
+
+describe("getGatewaySession", () => {
+  it("calls POST /auth/sessions and returns the cognito token", async () => {
+    const fetchSpy = vi.spyOn(globalThis, "fetch").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(session.expires_at).toBe(TOKEN_EXP);
+    expect(fetchSpy).toHaveBeenCalledWith(
+      `${GATEWAY_URL}/auth/sessions`,
+      expect.objectContaining({
+        method: "POST",
+        headers: expect.objectContaining({ Authorization: `Bearer ${COGNITO_TOKEN}` }),
+        body: JSON.stringify({ workspace_id: WORKSPACE_ID }),
+      }),
+    );
+  });
+
+  it("caches the session and does not re-call the gateway on second call", async () => {
+    const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(
+      new Response(JSON.stringify({ workspace_id: WORKSPACE_ID }), { status: 200 }),
+    );
+
+    const s1 = await getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN);
+    const s2 = await getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN);
+
+    expect(s1.token).toBe(s2.token);
+    expect(fetchSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it("evicts and re-exchanges after TTL expires", async () => {
+    const expiredClaims = makeClaims({ exp: Math.floor(Date.now() / 1000) - 1 });
+
+    const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(
+      new Response(JSON.stringify({ workspace_id: WORKSPACE_ID }), { status: 200 }),
+    );
+
+    await getGatewaySession(expiredClaims, WORKSPACE_ID, COGNITO_TOKEN);
+    await getGatewaySession(expiredClaims, WORKSPACE_ID, COGNITO_TOKEN);
+
+    expect(fetchSpy).toHaveBeenCalledTimes(2);
+  });
+
+  it("retries once on 401 and returns session if retry succeeds", async () => {
+    const fetchSpy = vi.spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response("Unauthorized", { status: 401 }))
+      .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 401", async () => {
+    vi.spyOn(globalThis, "fetch").mockResolvedValue(
+      new Response("Unauthorized", { status: 401 }),
+    );
+
+    await expect(
+      getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN),
+    ).rejects.toThrow(GatewaySessionError);
+  });
+
+  it("throws GatewaySessionError on 403", async () => {
+    vi.spyOn(globalThis, "fetch").mockResolvedValueOnce(
+      new Response("Forbidden", { status: 403 }),
+    );
+
+    await expect(
+      getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN),
+    ).rejects.toThrow(GatewaySessionError);
+  });
+
+  it("evictGatewaySession removes the cached entry", async () => {
+    const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(
+      new Response(JSON.stringify({ workspace_id: WORKSPACE_ID }), { status: 200 }),
+    );
+
+    await getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN);
+    evictGatewaySession(USER_SUB);
+    await getGatewaySession(makeClaims(), WORKSPACE_ID, COGNITO_TOKEN);
+
+    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/auth.test.ts

@@ -0,0 +1,87 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { Hono } from "hono";
+import { authMiddleware } from "../../src/middleware/auth";
+import type { AppVariables } from "../../src/types";
+import type { CognitoAccessTokenPayload } from "aws-jwt-verify/jwt-model";
+
+// Mock aws-jwt-verify at the module level
+vi.mock("aws-jwt-verify", () => {
+  const verify = vi.fn();
+  return {
+    CognitoJwtVerifier: {
+      create: vi.fn(() => ({ verify })),
+    },
+  };
+});
+
+// Access the mocked verify function
+const { CognitoJwtVerifier } = await import("aws-jwt-verify");
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const mockVerify = (CognitoJwtVerifier.create({} as any) as any).verify as ReturnType<typeof vi.fn>;
+
+function buildApp() {
+  const app = new Hono<{ Variables: AppVariables }>();
+  app.use("/protected", authMiddleware);
+  app.get("/protected", (c) => c.json({ sub: c.get("claims").sub }));
+  return app;
+}
+
+const fakeClaims = {
+  sub: "user-sub-123",
+  token_use: "access",
+  iss: "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_abc",
+  aud: "client-id",
+  exp: Math.floor(Date.now() / 1000) + 3600,
+  iat: Math.floor(Date.now() / 1000),
+} as unknown as CognitoAccessTokenPayload;
+
+describe("authMiddleware", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    process.env["COGNITO_USER_POOL_ID"] = "us-east-1_abc";
+    process.env["COGNITO_CLIENT_ID"] = "client-id";
+  });
+
+  it("passes with a valid Bearer token and sets claims", async () => {
+    mockVerify.mockResolvedValueOnce(fakeClaims);
+    const app = buildApp();
+    const res = await app.request("/protected", {
+      headers: { Authorization: "Bearer valid.token.here" },
+    });
+    expect(res.status).toBe(200);
+    const body = await res.json();
+    expect(body).toEqual({ sub: "user-sub-123" });
+  });
+
+  it("returns 401 when Authorization header is missing", async () => {
+    const app = buildApp();
+    const res = await app.request("/protected");
+    expect(res.status).toBe(401);
+  });
+
+  it("returns 401 when Authorization header is not Bearer", async () => {
+    const app = buildApp();
+    const res = await app.request("/protected", {
+      headers: { Authorization: "Basic dXNlcjpwYXNz" },
+    });
+    expect(res.status).toBe(401);
+  });
+
+  it("returns 401 for an expired token", async () => {
+    mockVerify.mockRejectedValueOnce(new Error("Token expired"));
+    const app = buildApp();
+    const res = await app.request("/protected", {
+      headers: { Authorization: "Bearer expired.token.here" },
+    });
+    expect(res.status).toBe(401);
+  });
+
+  it("returns 401 for wrong audience", async () => {
+    mockVerify.mockRejectedValueOnce(new Error("Token has expired or audience mismatch"));
+    const app = buildApp();
+    const res = await app.request("/protected", {
+      headers: { Authorization: "Bearer wrong.audience.token" },
+    });
+    expect(res.status).toBe(401);
+  });
+});

package/test/middleware/plan.test.ts

@@ -0,0 +1,99 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { Hono } from "hono";
+import type { AppVariables, WorkspaceItem } 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 })),
+  },
+  GetCommand: vi.fn((input) => input),
+}));
+
+const { planMiddleware, getWorkspace } = await import("../../src/middleware/plan");
+
+const freeWorkspace: WorkspaceItem = {
+  workspaceId: "ws-plan-test",
+  name: "Test Workspace",
+  plan: "free",
+  limits: {
+    dailyChatCap: 50,
+    maxModels: ["openrouter/auto"],
+    maxToolSteps: 5,
+    maxTokensPerRequest: 4096,
+  },
+  features: {
+    advancedTools: false,
+    customPrompts: false,
+  },
+  createdAt: "2026-07-01T00:00:00Z",
+  updatedAt: "2026-07-01T00:00:00Z",
+};
+
+function buildApp(workspaceId: string) {
+  const app = new Hono<{ Variables: AppVariables }>();
+  app.use("/protected", async (c, next) => {
+    c.set("claims", { sub: "user-sub" } as unknown as CognitoAccessTokenPayload);
+    c.set("workspaceId", workspaceId);
+    await next();
+  });
+  app.use("/protected", planMiddleware);
+  app.get("/protected", (c) =>
+    c.json({ plan: c.get("workspace").plan }),
+  );
+  return app;
+}
+
+describe("planMiddleware", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    process.env["WORKSPACE_TABLE_NAME"] = "gateway-prd-use1-workspaces";
+    process.env["AWS_REGION"] = "us-east-1";
+  });
+
+  it("passes and sets workspace for a valid free plan", async () => {
+    mockSend.mockResolvedValueOnce({ Item: { ...freeWorkspace, workspaceId: "ws-free" } });
+    const app = buildApp("ws-free");
+    const res = await app.request("/protected");
+    expect(res.status).toBe(200);
+    const body = await res.json();
+    expect(body).toEqual({ plan: "free" });
+  });
+
+  it("returns 404 when workspace is not found", async () => {
+    mockSend.mockResolvedValueOnce({ Item: undefined });
+    const app = buildApp("ws-missing");
+    const res = await app.request("/protected");
+    expect(res.status).toBe(404);
+    const body = await res.json();
+    expect(body).toMatchObject({ error: "Workspace not found" });
+  });
+
+  it("returns 402 when chat feature is explicitly disabled", async () => {
+    const restrictedWorkspace = {
+      ...freeWorkspace,
+      workspaceId: "ws-restricted",
+      features: { ...freeWorkspace.features, chat: false },
+    };
+    mockSend.mockResolvedValueOnce({ Item: restrictedWorkspace });
+    const app = buildApp("ws-restricted");
+    const res = await app.request("/protected");
+    expect(res.status).toBe(402);
+    const body = await res.json();
+    expect(body).toMatchObject({ error: expect.stringContaining("plan") });
+  });
+
+  it("uses cache on subsequent calls for the same workspaceId", async () => {
+    mockSend.mockResolvedValueOnce({
+      Item: { ...freeWorkspace, workspaceId: "ws-cache-test" },
+    });
+    await getWorkspace("ws-cache-test");
+    await getWorkspace("ws-cache-test");
+    expect(mockSend).toHaveBeenCalledTimes(1);
+  });
+});

package/test/middleware/workspace.test.ts

@@ -0,0 +1,122 @@
+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),
+}));
+
+const {
+  workspaceMiddleware,
+  listWorkspaceMemberships,
+  resetMembershipCache,
+} = await import("../../src/middleware/workspace");
+const { resetSessionCache } = await import("../../src/session");
+
+const MEMBERSHIPS_TABLE = "gateway-prd-use1-memberships";
+const SESSIONS_TABLE = "test-user-sessions";
+
+type MockCommand = { TableName?: string };
+
+function buildApp(userSub: string) {
+  const fakeClaims = { sub: userSub } as unknown as CognitoAccessTokenPayload;
+  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") }));
+  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["AWS_REGION"] = "us-east-1";
+  });
+
+  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" }] });
+    });
+
+    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 = 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");
+  });
+
+  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 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 = 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" }] });
+
+    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);
+  });
+});