@opengeni/api-router 3.3.0 → 3.4.0

package/src/access-grant-rls.ts

@@ -49,6 +49,55 @@ export function creditDebitAttributionForGrant(grant: AccessGrant): CreditDebitA
       : { kind: "unknown" };
 }
 
+/**
+ * Start `work` alongside an authorization `gate` instead of after it, so a
+ * request pays one database wait for both. The gate still decides: when it
+ * refuses, its error is the answer and `work`'s result or failure is dropped
+ * unread. Only use this for reads whose result is disclosed after the gate.
+ */
+export async function alongsideAuthorization<G, W>(
+  gate: Promise<G>,
+  work: () => Promise<W>,
+): Promise<[G, W]> {
+  const pending = work();
+  // Observed here so a refused gate never leaves an unhandled rejection.
+  pending.catch(() => undefined);
+  const decided = await gate;
+  return [decided, await pending];
+}
+
+type AgentLearningActor = {
+  kind: "agent";
+  sessionId: string;
+  turnId: string;
+  attemptId: string;
+  executionGeneration: number;
+};
+
+/** The attempt a well-formed agent-attempt grant claims to be; null otherwise. */
+function claimedLearningActor(grant: AccessGrant, sessionId: string): AgentLearningActor | null {
+  const turnId = grant.metadata?.turnId;
+  const attemptId = grant.metadata?.attemptId;
+  const executionGeneration = grant.metadata?.executionGeneration;
+  if (
+    typeof turnId !== "string" ||
+    typeof attemptId !== "string" ||
+    typeof executionGeneration !== "number"
+  ) {
+    return null;
+  }
+  return { kind: "agent", sessionId, turnId, attemptId, executionGeneration };
+}
+
+function sameLearningActor(left: AgentLearningActor, right: AgentLearningActor): boolean {
+  return (
+    left.sessionId === right.sessionId &&
+    left.turnId === right.turnId &&
+    left.attemptId === right.attemptId &&
+    left.executionGeneration === right.executionGeneration
+  );
+}
+
 export async function withAccessGrantSessionRlsContext<T>(
   deps: Pick<ApiRouteDeps, "db">,
   grant: AccessGrant,
@@ -63,19 +112,33 @@ export async function withAccessGrantSessionRlsContext<T>(
   if (typeof callerSessionId !== "string") {
     throw agentAttemptAuthorityError("malformed");
   }
-  try {
-    const actor = await requireLiveAgentAttemptAuthorization(deps.db, grant, callerSessionId);
-    const learning = await freezeAgentLearningPolicy(deps.db, {
+  const freezeLearning = (actor: AgentLearningActor) =>
+    freezeAgentLearningPolicy(deps.db, {
       accountId: grant.accountId,
       workspaceId: grant.workspaceId,
-      actor: {
-        kind: "agent",
-        sessionId: actor.callerSessionId,
-        turnId: actor.turnId,
-        attemptId: actor.attemptId,
-        executionGeneration: actor.executionGeneration,
-      },
+      actor,
     });
+  try {
+    // The learning policy is keyed by the exact attempt, which a live actor
+    // takes verbatim from the grant's claims. Read it for the claimed attempt
+    // while the liveness check runs; it is used only once that check passes
+    // and names the same attempt, and is otherwise read again for the actor.
+    const claimed = claimedLearningActor(grant, callerSessionId);
+    const [actor, speculative] = await alongsideAuthorization(
+      requireLiveAgentAttemptAuthorization(deps.db, grant, callerSessionId),
+      async () => (claimed ? await freezeLearning(claimed) : null),
+    );
+    const verified: AgentLearningActor = {
+      kind: "agent",
+      sessionId: actor.callerSessionId,
+      turnId: actor.turnId,
+      attemptId: actor.attemptId,
+      executionGeneration: actor.executionGeneration,
+    };
+    const learning =
+      claimed && speculative && sameLearningActor(claimed, verified)
+        ? speculative
+        : await freezeLearning(verified);
     return await withCreditDebitAttribution(
       {
         kind: "turn",
@@ -101,3 +164,28 @@ export async function withAccessGrantSessionRlsContext<T>(
     throw error;
   }
 }
+
+/**
+ * Prove that an agent-attempt grant is the exact live attempt of its bound
+ * session, with the same refusal as {@link withAccessGrantSessionRlsContext},
+ * without installing that request's RLS actor context. For requests that read
+ * and disclose nothing scoped (MCP protocol handshakes), so they need the
+ * liveness fence but none of the actor's data scope.
+ */
+export async function requireLiveAgentAttemptGrant(
+  deps: Pick<ApiRouteDeps, "db">,
+  grant: AccessGrant,
+): Promise<void> {
+  const callerSessionId = grant.metadata?.sessionId;
+  if (grant.principalKind !== "agent_attempt" || typeof callerSessionId !== "string") {
+    throw agentAttemptAuthorityError("malformed");
+  }
+  try {
+    await requireLiveAgentAttemptAuthorization(deps.db, grant, callerSessionId);
+  } catch (error) {
+    if (error instanceof SessionAuthorizationDeniedError) {
+      throw agentAttemptAuthorityError(error.reason, error);
+    }
+    throw error;
+  }
+}

package/src/app.ts

@@ -110,6 +110,7 @@ import {
   dbSql,
   getManagedAuthSessionSetSnapshot,
   getWorkspace,
+  getWorkspaceSettings,
   reapManagedAuthIsolatedSessions,
   reapExpiredManagedAuthSessionSets,
   resolveSessionMemoryAgentScope,
@@ -224,7 +225,11 @@ import {
   assertOrganizationUserSetupQueryTransportConfigured,
 } from "./auth/organization-user-setup";
 import { createApiSandboxClient, makeResumeBoxById } from "./sandbox/access";
-import { buildOpenGeniMcpServer } from "./mcp/server";
+import {
+  buildOpenGeniMcpProtocolServer,
+  buildOpenGeniMcpServer,
+  isOpenGeniMcpProtocolOnlyMessage,
+} from "./mcp/server";
 import {
   buildWorkspaceToolGatewayMcpServer,
   approveWorkspaceToolGatewayCall,
@@ -273,7 +278,11 @@ import {
   localBrowserRequestHost,
   markLocalInternalDispatch,
 } from "./http/local-browser-boundary";
-import { withAccessGrantSessionRlsContext } from "./access-grant-rls";
+import {
+  alongsideAuthorization,
+  requireLiveAgentAttemptGrant,
+  withAccessGrantSessionRlsContext,
+} from "./access-grant-rls";
 import { registerCapabilityRoutes } from "./routes/capabilities";
 import { registerCatalogAssetRoutes } from "./routes/catalog-assets";
 import { registerCodexRoutes } from "./routes/codex";
@@ -1795,6 +1804,28 @@ export function createAppComposition(deps: AppDependencies): {
       throw error;
     }
     const grant = authorization.grant;
+    // MCP handshakes (initialize, ping, notifications) read and disclose
+    // nothing scoped. A session-bound agent attempt still proves it is the
+    // live attempt, but skips the workspace, catalog, Memory and admin reads
+    // that building the tool surface costs on every tools/* POST.
+    if (
+      c.req.method === "POST" &&
+      grant.principalKind === "agent_attempt" &&
+      typeof grant.metadata?.sessionId === "string" &&
+      (await isProtocolOnlyMcpRequest(boundedRequest))
+    ) {
+      await requireLiveAgentAttemptGrant(routeDeps, grant);
+      const transport = new WebStandardStreamableHTTPServerTransport({
+        enableJsonResponse: true,
+      });
+      const mcp = buildOpenGeniMcpProtocolServer(grant);
+      try {
+        await mcp.connect(transport);
+        return await handleMcpRequestWithClientAbort(transport, boundedRequest, c.req.raw.signal);
+      } finally {
+        await mcp.close().catch(() => undefined);
+      }
+    }
     // Admin actions dispatch back into this API as a person. Capture the
     // context now, before this request takes on the agent's database and
     // billing scope, so none of it carries into those actions.
@@ -1804,7 +1835,8 @@ export function createAppComposition(deps: AppDependencies): {
     return await withSessionAuthorizationReadReuse((reads) =>
       withAccessGrantSessionRlsContext(routeDeps, grant, async () => {
         const boundSessionId = grant.metadata?.sessionId;
-        if (typeof boundSessionId === "string") {
+        const authorizeBoundSession = async () => {
+          if (typeof boundSessionId !== "string") return;
           try {
             await requireSessionAuthorization(routeDeps, grant, {
               sessionId: boundSessionId,
@@ -1822,16 +1854,14 @@ export function createAppComposition(deps: AppDependencies): {
             }
             throw error;
           }
-        }
-        if (c.req.method === "GET") return unsupportedGet();
-        reads.handOffToToolDispatch();
-        const workspace = await getWorkspace(routeDeps.db, workspaceId);
-        const workspaceMemoryEnabled = resolveWorkspaceMemoryEnabled(workspace?.settings);
-        const workspaceMemoryPromptMode = resolveWorkspaceMemoryPromptMode();
+        };
         const transport = new WebStandardStreamableHTTPServerTransport({
           enableJsonResponse: true,
         });
-        if (!grantUsesAttemptScopedMcp(grant)) {
+        if (c.req.method === "GET" || !grantUsesAttemptScopedMcp(grant)) {
+          await authorizeBoundSession();
+          if (c.req.method === "GET") return unsupportedGet();
+          reads.handOffToToolDispatch();
           const prepared = await prepareWorkspaceToolGateway(routeDeps, authorization);
           const mcp = buildWorkspaceToolGatewayMcpServer(prepared, grant, routeDeps.observability);
           try {
@@ -1845,39 +1875,58 @@ export function createAppComposition(deps: AppDependencies): {
             await Promise.allSettled([mcp.close(), prepared.close()]);
           }
         }
-        const mcpDeps = await resolveWorkspaceMcpRouteDeps(routeDeps, grant);
-        // The bound session's frozen Memory selector (migration 0427) decides
-        // which Memory tools the attempt receives and which typed layers they
-        // read and write. A missing row resolves to no Memory tools.
-        const sessionMemory =
-          typeof boundSessionId === "string"
-            ? ((await resolveSessionMemoryAgentScope(
-                routeDeps.db,
-                workspaceId,
-                boundSessionId,
-                grant.metadata,
-              )) ?? {
-                mode: "off" as const,
-                userSubjectId: null,
-                rootSessionId: null,
-              })
-            : null;
-        // Checked live: allowed by the organization, given to this session,
-        // and the person who gave it is still an owner or admin.
-        const sessionAdminAccess =
-          typeof boundSessionId === "string" &&
-          (await resolveSessionAdminAuthority(routeDeps.db, {
-            accountId: grant.accountId,
-            workspaceId,
-            sessionId: boundSessionId,
-          }))
-            ? {
-                origin: new URL(c.req.url).origin,
-                dispatch: (request: Request) =>
-                  outsideAgentContext(async () => await app.fetch(request, c.env)),
-                signal: c.req.raw.signal,
-              }
-            : null;
+        // The tool surface is four independent reads: workspace settings, the
+        // model catalog, the session's frozen Memory selector and its admin
+        // access. Read them together, alongside the bound-session check, which
+        // still decides before any of them reaches the response.
+        const [, [workspaceSettings, mcpDeps, sessionMemory, sessionAdminAuthority]] =
+          await alongsideAuthorization(authorizeBoundSession(), () =>
+            Promise.all([
+              getWorkspaceSettings(routeDeps.db, workspaceId),
+              resolveWorkspaceMcpRouteDeps(routeDeps, grant),
+              // The bound session's frozen Memory selector (migration 0427)
+              // decides which Memory tools the attempt receives and which typed
+              // layers they read and write. A missing row resolves to no Memory
+              // tools.
+              typeof boundSessionId === "string"
+                ? resolveSessionMemoryAgentScope(
+                    routeDeps.db,
+                    workspaceId,
+                    boundSessionId,
+                    grant.metadata,
+                  ).then(
+                    (scope) =>
+                      scope ?? {
+                        mode: "off" as const,
+                        userSubjectId: null,
+                        rootSessionId: null,
+                      },
+                  )
+                : null,
+              // Checked live: allowed by the organization, given to this
+              // session, and the person who gave it is still an owner or admin.
+              typeof boundSessionId === "string"
+                ? resolveSessionAdminAuthority(routeDeps.db, {
+                    accountId: grant.accountId,
+                    workspaceId,
+                    sessionId: boundSessionId,
+                  })
+                : false,
+            ]),
+          );
+        reads.handOffToToolDispatch();
+        const workspaceMemoryEnabled = resolveWorkspaceMemoryEnabled(
+          workspaceSettings ?? undefined,
+        );
+        const workspaceMemoryPromptMode = resolveWorkspaceMemoryPromptMode();
+        const sessionAdminAccess = sessionAdminAuthority
+          ? {
+              origin: new URL(c.req.url).origin,
+              dispatch: (request: Request) =>
+                outsideAgentContext(async () => await app.fetch(request, c.env)),
+              signal: c.req.raw.signal,
+            }
+          : null;
         const mcp = buildOpenGeniMcpServer(mcpDeps, grant, {
           requestOrigin: new URL(c.req.url).origin,
           workspaceMemoryEnabled,
@@ -2494,6 +2543,19 @@ async function validateCodexAccountTarget(request: Request): Promise<void> {
   }
 }
 
+/** Whether a POST carries only MCP protocol messages that touch no tool. */
+async function isProtocolOnlyMcpRequest(request: Request): Promise<boolean> {
+  let body: unknown;
+  try {
+    body = await request.clone().json();
+  } catch {
+    // The transport answers malformed bodies itself, on the full path.
+    return false;
+  }
+  const messages = Array.isArray(body) ? body : [body];
+  return messages.length > 0 && messages.every(isOpenGeniMcpProtocolOnlyMessage);
+}
+
 async function requireMcpAccessGrantAuthorization(
   c: Parameters<typeof requireAccessGrant>[0],
   deps: ApiRouteDeps,

package/src/billing-stripe-customer.ts

@@ -0,0 +1,40 @@
+import type Stripe from "stripe";
+
+export interface StripeCustomerTransaction {
+  read(): Promise<string | null>;
+  account(): Promise<{ name: string } | null>;
+  save(customer: Stripe.Customer): Promise<void>;
+}
+
+/** The adapter must hold the account advisory lock for this whole callback,
+ * including provider I/O and persistence. The stable provider key also recovers
+ * creation when Stripe succeeds but the database transaction rolls back. */
+export async function serializedStripeCustomer(input: {
+  stripe: Stripe;
+  accountId: string;
+  provider: string;
+  email?: string;
+  locked: (run: (tx: StripeCustomerTransaction) => Promise<string>) => Promise<string>;
+}): Promise<string> {
+  return await input.locked(async (tx) => {
+    const existing = await tx.read();
+    if (existing) return existing;
+    const account = await tx.account();
+    if (!account) throw new Error("account not found");
+    const created = await input.stripe.customers.create(
+      {
+        metadata: { opengeni_account_id: input.accountId },
+      },
+      { idempotencyKey: `billing-customer:${input.provider}:${input.accountId}` },
+    );
+    // Creation parameters must stay identical even if a different administrator
+    // retries after a rollback, or the organization is renamed in between.
+    // Profile updates are idempotent sets and do not create another customer.
+    const customer = await input.stripe.customers.update(created.id, {
+      name: account.name,
+      ...(input.email ? { email: input.email } : {}),
+    });
+    await tx.save(customer);
+    return customer.id;
+  });
+}

package/src/billing-subscription-store.ts

@@ -0,0 +1,121 @@
+import type { BillingPlanBalance, SubscriptionPlan } from "@opengeni/contracts";
+import {
+  assertBillingSubscriptionOperationRoot,
+  beginBillingSubscriptionOperation,
+  completeBillingSubscriptionOperation,
+  getBillingSubscriptionOperation,
+  nestedPostgresSqlState,
+  expirePlanCreditPeriod,
+  getBillingCustomer,
+  getManagedAccount,
+  getOrganizationPlanBalance,
+  readSubscriptionPlanCatalog,
+  replacePlanCreditPeriod,
+  upsertOrganizationPlan,
+  withRlsContext,
+  type Database,
+  type OrganizationPlanStateInput,
+  type PlanCreditPeriodInput,
+  type BillingSubscriptionOperation,
+  type BillingSubscriptionOperationKind,
+} from "@opengeni/db";
+import { HTTPException } from "hono/http-exception";
+import { organizationPlans, planCreditPeriods } from "@opengeni/db/schema";
+import { and, eq, isNull, sql } from "drizzle-orm";
+
+export type StoredSubscriptionPlan = {
+  plan_key: string;
+  stripe_subscription_id: string | null;
+};
+
+export interface SubscriptionAccountStore {
+  operation(key: string): Promise<BillingSubscriptionOperation | null>;
+  beginOperation(input: {
+    operationKey: string;
+    kind: BillingSubscriptionOperationKind;
+    request: Record<string, unknown>;
+    intent: Record<string, unknown>;
+  }): Promise<BillingSubscriptionOperation>;
+  completeOperation(key: string, result: Record<string, unknown>): Promise<void>;
+  exists(): Promise<boolean>;
+  state(): Promise<StoredSubscriptionPlan | null>;
+  customer(provider: string): Promise<string | null>;
+  balance(): Promise<BillingPlanBalance | null>;
+  save(input: Omit<OrganizationPlanStateInput, "accountId">): Promise<void>;
+  replace(input: Omit<PlanCreditPeriodInput, "accountId">): Promise<void>;
+  expire(subscriptionId: string): Promise<void>;
+}
+
+export interface SubscriptionStore {
+  catalog(): Promise<{ revision: number; plans: SubscriptionPlan[] }>;
+  locked<T>(accountId: string, run: (store: SubscriptionAccountStore) => Promise<T>): Promise<T>;
+}
+
+/** Reuse the SQL seams in one outer transaction, membership/org fence before
+ * the credit-account lock. Stripe reads also run under this fence: concurrent
+ * deliveries cannot commit an older provider snapshot after a newer one. */
+export function subscriptionStore(db: Database): SubscriptionStore {
+  assertBillingSubscriptionOperationRoot(db);
+  return {
+    catalog: () => readSubscriptionPlanCatalog(db),
+    locked: (accountId, run) =>
+      withRlsContext(db, { accountId, workspaceId: null }, async (tx) => {
+        await tx.execute(
+          sql`select pg_advisory_xact_lock(hashtextextended(${`organization-membership:${accountId}`}, 0))`,
+        );
+        return await run({
+          operation: (key) => getBillingSubscriptionOperation(tx, accountId, key),
+          beginOperation: async (input) => {
+            try {
+              return await beginBillingSubscriptionOperation(tx, { ...input, accountId });
+            } catch (error) {
+              if (nestedPostgresSqlState(error) === "P0S02")
+                throw new HTTPException(409, {
+                  message: "Retry the pending subscription operation before starting another",
+                });
+              throw error;
+            }
+          },
+          completeOperation: async (operationKey, result) => {
+            await completeBillingSubscriptionOperation(tx, { accountId, operationKey, result });
+          },
+          exists: async () => Boolean(await getManagedAccount(tx, accountId)),
+          state: async () => {
+            const [row] = await tx
+              .select({
+                plan_key: organizationPlans.planKey,
+                stripe_subscription_id: organizationPlans.stripeSubscriptionId,
+              })
+              .from(organizationPlans)
+              .where(eq(organizationPlans.accountId, accountId))
+              .limit(1);
+            return row ?? null;
+          },
+          customer: async (provider) =>
+            (await getBillingCustomer(tx, accountId, provider))?.providerCustomerId ?? null,
+          balance: () => getOrganizationPlanBalance(tx, accountId),
+          save: async (input) => {
+            await upsertOrganizationPlan(tx, { ...input, accountId });
+          },
+          replace: async (input) => {
+            await replacePlanCreditPeriod(tx, { ...input, accountId });
+          },
+          expire: async (subscriptionId) => {
+            const periods = await tx
+              .select({ id: planCreditPeriods.id })
+              .from(planCreditPeriods)
+              .where(
+                and(
+                  eq(planCreditPeriods.accountId, accountId),
+                  eq(planCreditPeriods.stripeSubscriptionId, subscriptionId),
+                  isNull(planCreditPeriods.expiredAt),
+                ),
+              );
+            for (const period of periods) {
+              await expirePlanCreditPeriod(tx, { accountId, periodId: period.id, force: true });
+            }
+          },
+        });
+      }),
+  };
+}