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

package/src/gateway-session.ts

@@ -0,0 +1,148 @@
+/**
+ * Gateway session client.
+ *
+ * The registry gateway authenticates every request with the caller's Cognito
+ * access token and requires the active workspace to be persisted in DDB via
+ * `POST /auth/sessions` before other endpoints will accept requests.
+ *
+ * GatewaySessionClient handles that one-time setup per user and returns the
+ * Cognito token itself as the bearer to use on subsequent gateway calls. The
+ * result is cached in-memory per user sub for the remaining lifetime of the
+ * Cognito token.
+ */
+
+import type { CognitoAccessTokenPayload } from "aws-jwt-verify/jwt-model";
+
+interface SessionEntry {
+  /** The Cognito access token to use as bearer for gateway calls. */
+  token: string;
+  /** Unix seconds — mirrors the Cognito token's `exp` claim. */
+  expires_at: number;
+}
+
+const sessionCache = new Map<string, SessionEntry>();
+
+/**
+ * Per-session tool-list cache. Keyed by userSub, parallel to sessionCache.
+ * Cleared on session eviction so a fresh session always re-fetches tools.
+ */
+const toolsCache = new Map<string, unknown[]>();
+
+export function getCachedTools(sub: string): unknown[] | undefined {
+  return toolsCache.get(sub);
+}
+
+export function setCachedTools(sub: string, tools: unknown[]): void {
+  toolsCache.set(sub, tools);
+}
+
+export function evictCachedTools(sub: string): void {
+  toolsCache.delete(sub);
+}
+
+function gatewayUrl(): string {
+  const url = process.env["GATEWAY_URL"];
+  if (!url) throw new Error("GATEWAY_URL is not set");
+  return url.replace(/\/$/, "");
+}
+
+/**
+ * Establish (or reuse) a gateway session for the caller.
+ *
+ * 1. Returns the cached entry if still valid.
+ * 2. Calls `POST /auth/sessions` on the gateway to register the active
+ *    workspace for this user.
+ * 3. Returns the Cognito token as the bearer — the gateway uses it directly.
+ * 4. On 401, evicts the cache and retries once.
+ */
+export async function getGatewaySession(
+  claims: CognitoAccessTokenPayload,
+  workspaceId: string,
+  cognitoToken: string,
+): Promise<SessionEntry> {
+  const sub = claims.sub;
+  const now = Math.floor(Date.now() / 1000);
+
+  const cached = sessionCache.get(sub);
+  if (cached && cached.expires_at > now + 60) {
+    return cached;
+  }
+
+  const entry = await exchangeSession(cognitoToken, workspaceId, claims.exp as number);
+  sessionCache.set(sub, entry);
+  return entry;
+}
+
+async function exchangeSession(
+  cognitoToken: string,
+  workspaceId: string,
+  exp: number,
+): Promise<SessionEntry> {
+  let res: Response;
+  try {
+    res = await callAuthSessions(cognitoToken, workspaceId);
+  } catch {
+    // Network error — retry once after 200ms.
+    await new Promise<void>((r) => setTimeout(r, 200));
+    res = await callAuthSessions(cognitoToken, workspaceId);
+  }
+
+  if (res.status === 401) {
+    // Retry once — token may have just been refreshed by the caller.
+    const retryRes = await callAuthSessions(cognitoToken, workspaceId);
+    if (!retryRes.ok) {
+      throw new GatewaySessionError(retryRes.status, await retryRes.text());
+    }
+    return { token: cognitoToken, expires_at: exp };
+  }
+
+  if (res.status >= 500) {
+    // 5xx — retry once after 200ms.
+    await new Promise<void>((r) => setTimeout(r, 200));
+    const retryRes = await callAuthSessions(cognitoToken, workspaceId);
+    if (!retryRes.ok) {
+      throw new GatewaySessionError(retryRes.status, await retryRes.text());
+    }
+    return { token: cognitoToken, expires_at: exp };
+  }
+
+  if (!res.ok) {
+    throw new GatewaySessionError(res.status, await res.text());
+  }
+
+  return { token: cognitoToken, expires_at: exp };
+}
+
+async function callAuthSessions(
+  cognitoToken: string,
+  workspaceId: string,
+): Promise<Response> {
+  return fetch(`${gatewayUrl()}/auth/sessions`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${cognitoToken}`,
+    },
+    body: JSON.stringify({ workspace_id: workspaceId }),
+  });
+}
+
+export function evictGatewaySession(sub: string): void {
+  sessionCache.delete(sub);
+  toolsCache.delete(sub);
+}
+
+export function resetGatewaySessionCache(): void {
+  sessionCache.clear();
+  toolsCache.clear();
+}
+
+export class GatewaySessionError extends Error {
+  constructor(
+    public readonly status: number,
+    message: string,
+  ) {
+    super(`Gateway session exchange failed (${status}): ${message}`);
+    this.name = "GatewaySessionError";
+  }
+}

package/src/index.ts

@@ -0,0 +1,19 @@
+import { serve } from "@hono/node-server";
+import { createChatApp, initPostHog } from "./app.js";
+import { parseEnv } from "./env.js";
+
+export { createChatApp, initPostHog } from "./app.js";
+export type { ChatApp } from "./app.js";
+export { handler } from "./lambda.js";
+export type { LambdaEvent, LambdaContext } from "hono/aws-lambda";
+
+const env = parseEnv(process.env);
+
+initPostHog(env);
+
+if (env.NODE_ENV !== "test") {
+  const app = createChatApp();
+  serve({ fetch: app.fetch, port: env.PORT }, () => {
+    console.log(`Chat API listening on :${env.PORT}`);
+  });
+}

package/src/lambda.ts

@@ -0,0 +1,17 @@
+import { streamHandle } from "hono/aws-lambda";
+import { createChatApp, initPostHog } from "./app.js";
+import { parseEnv } from "./env.js";
+import type { LambdaEvent } from "hono/aws-lambda";
+
+const env = parseEnv(process.env);
+
+// Initialize module-scope PostHog singletons once per cold start
+initPostHog(env);
+
+const app = createChatApp();
+
+export const handler = streamHandle(app) as (
+  event: LambdaEvent,
+  context: unknown,
+  callback: unknown,
+) => Promise<unknown>;

package/src/middleware/auth.ts

@@ -0,0 +1,38 @@
+import { CognitoJwtVerifier } from "aws-jwt-verify";
+import type { CognitoAccessTokenPayload } from "aws-jwt-verify/jwt-model";
+import type { MiddlewareHandler } from "hono";
+import type { AppVariables } from "../types";
+
+interface JwtVerifier {
+  verify(token: string): Promise<CognitoAccessTokenPayload>;
+}
+
+let verifier: JwtVerifier | null = null;
+
+function getVerifier(): JwtVerifier {
+  if (!verifier) {
+    verifier = CognitoJwtVerifier.create({
+      userPoolId: process.env["COGNITO_USER_POOL_ID"]!,
+      clientId: process.env["COGNITO_CLIENT_ID"]!,
+      tokenUse: "access",
+    });
+  }
+  return verifier;
+}
+
+export const authMiddleware: MiddlewareHandler<{ Variables: AppVariables }> =
+  async (c, next) => {
+    const authHeader = c.req.header("Authorization");
+    if (!authHeader?.startsWith("Bearer ")) {
+      return c.json({ error: "Unauthorized" }, 401);
+    }
+
+    const token = authHeader.slice(7);
+    try {
+      const payload = await getVerifier().verify(token);
+      c.set("claims", payload);
+      return next();
+    } catch {
+      return c.json({ error: "Unauthorized" }, 401);
+    }
+  };

package/src/middleware/plan.ts

@@ -0,0 +1,67 @@
+import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
+import { DynamoDBDocumentClient, GetCommand } from "@aws-sdk/lib-dynamodb";
+import type { MiddlewareHandler } from "hono";
+import type { AppVariables, WorkspaceItem } from "../types";
+
+const WORKSPACE_CACHE_TTL_MS = 60_000;
+
+interface CacheEntry {
+  workspace: WorkspaceItem;
+  fetchedAt: number;
+}
+
+const workspaceCache = new Map<string, CacheEntry>();
+
+let ddbClient: DynamoDBDocumentClient | null = null;
+
+function getDdb() {
+  if (!ddbClient) {
+    ddbClient = DynamoDBDocumentClient.from(
+      new DynamoDBClient({ region: process.env["AWS_REGION"] }),
+    );
+  }
+  return ddbClient;
+}
+
+export async function getWorkspace(workspaceId: string): Promise<WorkspaceItem | null> {
+  const now = Date.now();
+  const cached = workspaceCache.get(workspaceId);
+  if (cached && now - cached.fetchedAt < WORKSPACE_CACHE_TTL_MS) {
+    return cached.workspace;
+  }
+
+  const result = await getDdb().send(
+    new GetCommand({
+      TableName: process.env["WORKSPACE_TABLE_NAME"]!,
+      Key: { workspaceId },
+    }),
+  );
+
+  if (!result.Item) return null;
+
+  const workspace = result.Item as WorkspaceItem;
+  workspaceCache.set(workspaceId, { workspace, fetchedAt: now });
+  return workspace;
+}
+
+export const planMiddleware: MiddlewareHandler<{ Variables: AppVariables }> =
+  async (c, next) => {
+    const workspaceId = c.get("workspaceId");
+    const workspace = await getWorkspace(workspaceId);
+    if (!workspace) {
+      return c.json({ error: "Workspace not found" }, 404);
+    }
+
+    // 402 Payment Required when chat is not included in the plan.
+    // Currently all plans include chat; this gate is a forward-looking hook
+    // for future restricted plans where chat is a paid add-on.
+    if ("chat" in workspace.features && !workspace.features.chat) {
+      return c.json(
+        { error: "Chat is not available on your current plan", plan: workspace.plan },
+        402,
+      );
+    }
+
+    c.set("workspace", workspace);
+    return next();
+  };

package/src/middleware/workspace.ts

@@ -0,0 +1,96 @@
+import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
+import { DynamoDBDocumentClient, QueryCommand } from "@aws-sdk/lib-dynamodb";
+import { getSessionWorkspaceId } from "../session.js";
+import type { AppVariables } from "../types.js";
+import type { MiddlewareHandler } from "hono";
+
+ 
+
+const MEMBERSHIP_CACHE_TTL_MS = 300_000;
+
+interface CacheEntry {
+  workspaceIds: string[];
+  fetchedAt: number;
+}
+
+const membershipCache = new Map<string, CacheEntry>();
+
+let ddbClient: DynamoDBDocumentClient | null = null;
+
+function getDdb() {
+  if (!ddbClient) {
+    ddbClient = DynamoDBDocumentClient.from(
+      new DynamoDBClient({ region: process.env["AWS_REGION"] }),
+    );
+  }
+  return ddbClient;
+}
+
+/** Returns all workspace IDs the user is a member of. Empty array = no memberships. */
+export async function listWorkspaceMemberships(
+  userSub: string,
+): Promise<string[]> {
+  const now = Date.now();
+  const cached = membershipCache.get(userSub);
+  if (cached && now - cached.fetchedAt < MEMBERSHIP_CACHE_TTL_MS) {
+    return cached.workspaceIds;
+  }
+
+  const result = await getDdb().send(
+    new QueryCommand({
+      TableName: process.env["MEMBERSHIPS_TABLE_NAME"]!,
+      IndexName: "ByUserSub",
+      KeyConditionExpression: "userSub = :sub",
+      ExpressionAttributeValues: { ":sub": userSub },
+    }),
+  );
+
+  const workspaceIds = (result.Items ?? []).map(
+    (item) => (item as { workspaceId: string }).workspaceId,
+  );
+  membershipCache.set(userSub, { workspaceIds, fetchedAt: now });
+  return workspaceIds;
+}
+
+export function clearMembershipCache(userSub: string): void {
+  membershipCache.delete(userSub);
+}
+
+export function resetMembershipCache(): void {
+  membershipCache.clear();
+}
+
+/**
+ * Resolves the active workspace ID for the user:
+ * 1. Check the session table for an explicit activeWorkspaceId.
+ * 2. Verify the user still has membership for that workspace.
+ * 3. Fall back to the first membership if no session preference or membership expired.
+ */
+export async function resolveWorkspaceId(
+  userSub: string,
+): Promise<string | null> {
+  const [sessionWsId, workspaceIds] = await Promise.all([
+    getSessionWorkspaceId(userSub),
+    listWorkspaceMemberships(userSub),
+  ]);
+
+  if (workspaceIds.length === 0) return null;
+
+  if (sessionWsId && workspaceIds.includes(sessionWsId)) {
+    return sessionWsId;
+  }
+
+  return workspaceIds[0] ?? null;
+}
+
+export const workspaceMiddleware: MiddlewareHandler<{
+  Variables: AppVariables;
+}> = async (c, next) => {
+  const claims = c.get("claims");
+  const workspaceId = await resolveWorkspaceId(claims.sub);
+  if (!workspaceId) {
+    return c.json({ error: "No workspace membership" }, 403);
+  }
+  c.set("workspaceId", workspaceId);
+  return next();
+};

package/src/posthog.ts

@@ -0,0 +1,48 @@
+import { Prompts, type PromptResult } from "@posthog/ai";
+import { PostHog } from "posthog-node";
+import { FALLBACK_PROMPTS } from "./fallback-prompts.js";
+import type { Env } from "./env.js";
+
+// Module-scope singletons — survive Lambda warm invocations
+let _client: PostHog | null = null;
+let _prompts: Prompts | null = null;
+
+export function initPostHog(env: Env): void {
+  if (!env.POSTHOG_PROJECT_API_KEY || !env.POSTHOG_PERSONAL_API_KEY) return;
+
+  _client = new PostHog(env.POSTHOG_PROJECT_API_KEY, {
+    host: env.POSTHOG_HOST,
+    personalApiKey: env.POSTHOG_PERSONAL_API_KEY,
+  });
+
+  _prompts = new Prompts({ posthog: _client });
+}
+
+export async function getPrompt(
+  name: string,
+  cacheTtlSeconds = 300,
+): Promise<PromptResult> {
+  const fallback = FALLBACK_PROMPTS[name] ?? "";
+
+  if (!_prompts) {
+    return { source: "code_fallback", prompt: fallback, name: undefined, version: undefined };
+  }
+
+  return _prompts.get(name, { cacheTtlSeconds, fallback });
+}
+
+export function compilePrompt(
+  template: string,
+  vars: Record<string, string>,
+): string {
+  if (_prompts) {
+    return _prompts.compile(template, vars);
+  }
+  return template.replace(/\{\{(\w+)\}\}/g, (_, key) => vars[key] ?? "");
+}
+
+export function getPostHogClient(): PostHog | null {
+  return _client;
+}
+
+export type { PostHog };

package/src/providers/openrouter.ts

@@ -0,0 +1,37 @@
+import {
+  SecretsManagerClient,
+  GetSecretValueCommand,
+} from "@aws-sdk/client-secrets-manager";
+import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
+
+let secretsClient: SecretsManagerClient | null = null;
+let cachedKey: string | null = null;
+
+function getSecretsClient(): SecretsManagerClient {
+  if (!secretsClient) {
+    secretsClient = new SecretsManagerClient({
+      region: process.env["AWS_REGION"],
+    });
+  }
+  return secretsClient;
+}
+
+export async function getOpenRouterKey(): Promise<string> {
+  if (cachedKey) return cachedKey;
+
+  const result = await getSecretsClient().send(
+    new GetSecretValueCommand({
+      SecretId: process.env["OPENROUTER_SECRET_ARN"]!,
+    }),
+  );
+  cachedKey = result.SecretString!;
+  return cachedKey!;
+}
+
+export function createOpenRouterProvider(apiKey: string) {
+  return createOpenAICompatible({
+    name: "openrouter",
+    apiKey,
+    baseURL: "https://openrouter.ai/api/v1",
+  });
+}