zopu-code/packages/agents/src/auth.ts

/**
 * Authenticated Flue route middleware for the organization-scoped global agent.
 *
 * Every request to the Zopu agent routes through here. The middleware:
 *
 * 1. Requires a Bearer access token (the Better Auth/Convex JWT).
 * 2. Creates a fresh, request-scoped {@link ConvexHttpClient} and sets the
 *    token on it — never on a shared/global client. Auth is per-request.
 * 3. Resolves (ensuring) the caller's current personal organization.
 * 4. Requires the agent instance id (`:id`) to equal that organization id, so
 *    one organization cannot use or observe another's agent instance.
 * 5. Authorizes GET/SSE observation routes as strictly as POST.
 * 6. For POST only: captures exact user-message evidence *before* Flue
 *    admission (`beginUserMessage`, status `admitting`), then patches it to
 *    `admitted` with the Flue receipt submission id, or `failed` on error.
 *    Observation requests (GET/HEAD) never create evidence.
 *
 * The organization id, conversation id, message id, role, and timestamp are
 * never trusted from the client or an agent tool — they are resolved
 * server-side. Only the raw message text travels from the user, and it is
 * stored verbatim (never trimmed or rewritten).
 */
import { parseAgentEnv } from "@code/env/agent";
import type { AgentRouteHandler } from "@flue/runtime";
import { ConvexHttpClient } from "convex/browser";
import { makeFunctionReference } from "convex/server";

// ---------------------------------------------------------------------------
// Typed function references (no generated backend bindings imported).
// ---------------------------------------------------------------------------

interface OrganizationView {
  readonly _id: string;
  readonly _creationTime: number;
  readonly name: string;
  readonly kind: "personal" | "team";
  readonly createdBy: string;
  readonly createdAt: number;
}

interface ConversationMessageView {
  readonly _id: string;
  readonly _creationTime: number;
  readonly organizationId: string;
  readonly conversationId: string;
  readonly messageId: string;
  readonly submissionId: string | null;
  readonly clientRequestId: string;
  readonly role: "user";
  readonly rawText: string;
  readonly status: "admitting" | "admitted" | "failed";
  readonly createdAt: number;
}

const ensurePersonalOrganization = makeFunctionReference<
  "mutation",
  Record<string, never>,
  OrganizationView
>("organizations:ensurePersonalOrganization");

const beginUserMessage = makeFunctionReference<
  "mutation",
  {
    readonly organizationId: string;
    readonly clientRequestId: string;
    readonly rawText: string;
  },
  ConversationMessageView
>("conversationMessages:beginUserMessage");

const markAdmitted = makeFunctionReference<
  "mutation",
  {
    readonly organizationId: string;
    readonly clientRequestId: string;
    readonly submissionId: string;
  },
  ConversationMessageView
>("conversationMessages:markAdmitted");

const markFailed = makeFunctionReference<
  "mutation",
  {
    readonly organizationId: string;
    readonly clientRequestId: string;
  },
  ConversationMessageView
>("conversationMessages:markFailed");

// ---------------------------------------------------------------------------
// Request-scoped authenticated Convex client.
// ---------------------------------------------------------------------------

const { CONVEX_URL } = parseAgentEnv(process.env);

/**
 * A short-lived authenticated Convex client bound to one HTTP request. The
 * caller's JWT is set on this instance only and discarded when the request
 * ends. We never mutate auth on a shared/global Convex client.
 */
export const createAuthenticatedClient = (
  accessToken: string
): ConvexHttpClient => {
  const client = new ConvexHttpClient(CONVEX_URL);
  client.setAuth(accessToken);
  return client;
};

/**
 * Minimal structural shape the header helpers read from an incoming request.
 * Using this avoids a structural clash between the global (Bun) `Request`/
 * `Headers` and the undici types that Hono's context exposes, while staying
 * type-safe.
 */
interface HeaderSource {
  readonly headers: {
    readonly get: (name: string) => string | null;
  };
}

/**
 * Extract and validate the Bearer access token from the request. Returns null
 * when absent or malformed so the caller can produce a clean 401.
 */
export const extractBearerToken = (request: HeaderSource): string | null => {
  const header = request.headers.get("authorization");
  if (!header) {
    return null;
  }
  const match = /^Bearer\s+(?<token>\S+)$/u.exec(header);
  return match?.groups?.token ?? null;
};

/**
 * Read the stable per-send request id supplied by the web client. This is
 * generated once when the user sends a message and reused across fetch/stream
 * retries, so it is a reliable idempotency key for evidence capture.
 */
const extractClientRequestId = (request: HeaderSource): string | null => {
  const header = request.headers.get("x-zopu-request-id");
  if (!header || header.length === 0) {
    return null;
  }
  return header;
};

/**
 * Safely parse the direct agent payload to extract the exact user message
 * text. Mirrors Flue's `DirectAgentPayload` shape `{ message, images? }`.
 * Returns null for non-JSON or missing `message` so the request is rejected
 * with 400 rather than crashing the middleware.
 */
const extractMessageText = async (
  request: HeaderSource & { readonly json: () => Promise<unknown> }
): Promise<string | null> => {
  const contentType = request.headers.get("content-type") ?? "";
  if (!contentType.includes("application/json")) {
    return null;
  }
  try {
    const body = (await request.json()) as unknown;
    if (
      typeof body !== "object" ||
      body === null ||
      !("message" in body) ||
      typeof (body as { message?: unknown }).message !== "string"
    ) {
      return null;
    }
    return (body as { message: string }).message;
  } catch {
    return null;
  }
};

/**
 * Read the Flue admission receipt submission id from the response Flue
 * produced after `next()`. The response body is `202 { streamUrl, offset,
 * submissionId }`. Returns null when the body has no submission id (e.g. an
 * error response).
 */
const extractSubmissionId = async (
  response: HeaderSource & {
    readonly clone: () => { readonly json: () => Promise<unknown> };
  }
): Promise<string | null> => {
  const contentType = response.headers.get("content-type") ?? "";
  if (!contentType.includes("application/json")) {
    return null;
  }
  try {
    const body = (await response.clone().json()) as unknown;
    if (
      typeof body === "object" &&
      body !== null &&
      "submissionId" in body &&
      typeof (body as { submissionId?: unknown }).submissionId === "string"
    ) {
      return (body as { submissionId: string }).submissionId;
    }
  } catch {
    // Non-JSON or unreadable body: no submission id to record.
  }
  return null;
};

/**
 * Admit a direct submission through Flue and record the admission outcome on
 * the normalized evidence row. On admission failure, the evidence is marked
 * failed and the original error is rethrown — never swallowed. On success,
 * the Flue receipt submission id (if present) patches the evidence to
 * admitted.
 */
const admitAndRecordEvidence = async (
  continueAdmission: () => Promise<void>,
  client: ConvexHttpClient,
  getResponse: () => Response | undefined,
  evidence: {
    readonly organizationId: string;
    readonly clientRequestId: string;
  }
): Promise<void> => {
  try {
    await continueAdmission();
  } catch (error) {
    try {
      await client.mutation(markFailed, evidence);
    } catch {
      // Best-effort: the original error is the primary concern.
    }
    throw error;
  }

  // Patch the evidence with the Flue receipt submission id. If the response
  // carries no submission id (unexpected), leave the evidence in its begun
  // state rather than guessing.
  const response = getResponse();
  if (!response) {
    return;
  }
  const submissionId = await extractSubmissionId(response);
  if (!submissionId) {
    return;
  }
  try {
    await client.mutation(markAdmitted, {
      ...evidence,
      submissionId,
    });
  } catch {
    // Best-effort: the admission already succeeded and the response is
    // finalized. A failure to patch metadata must not break the request.
  }
};

/**
 * The authenticated Flue route middleware. Authorizes every method, and for
 * POST captures normalized user-message evidence around Flue admission.
 */
export const authenticatedAgentRoute: AgentRouteHandler = async (c, next) => {
  const accessToken = extractBearerToken(c.req.raw);
  if (!accessToken) {
    return c.json({ error: "Unauthorized" }, 401);
  }

  const client = createAuthenticatedClient(accessToken);

  // Resolve (ensuring) the caller's current personal organization.
  let organization: OrganizationView;
  try {
    organization = await client.mutation(ensurePersonalOrganization, {});
  } catch {
    return c.json({ error: "Unauthorized" }, 401);
  }

  // The agent instance id must be the caller's organization id. This is the
  // hard tenancy boundary: one organization cannot use or observe another's
  // agent instance.
  const instanceId = c.req.param("id") ?? "";
  if (instanceId !== organization._id) {
    return c.json({ error: "Forbidden" }, 403);
  }

  // Observation routes (GET/HEAD) are authorized but never create evidence.
  if (c.req.method !== "POST") {
    return next();
  }

  // --- Evidence capture (POST only) -------------------------------------

  const clientRequestId = extractClientRequestId(c.req.raw);
  const messageText = await extractMessageText(c.req.raw.clone());

  if (!clientRequestId) {
    return c.json({ error: "Missing x-zopu-request-id" }, 400);
  }
  if (messageText === null) {
    return c.json({ error: "Invalid message payload" }, 400);
  }

  // Begin evidence before Flue admission so the message is durable even if
  // admission fails or the process is interrupted.
  try {
    await client.mutation(beginUserMessage, {
      clientRequestId,
      organizationId: organization._id,
      rawText: messageText,
    });
  } catch {
    return c.json({ error: "Evidence capture failed" }, 500);
  }

  // Admit through Flue. After next(), c.res holds the admission response.
  // Evidence is patched from the admission receipt, or marked failed without
  // swallowing the original error.
  await admitAndRecordEvidence(next, client, () => c.res, {
    clientRequestId,
    organizationId: organization._id,
  });
  return c.res;
};