feat: add per-session spend limits

packages/control-plane/src/router.policy.test.ts

@@ -12,11 +12,11 @@ function routeFor(method: string, path: string) {
 
 describe("route policy table", () => {
   it("publishes the complete canonical route catalog", () => {
-    expect(routes).toHaveLength(171);
+    expect(routes).toHaveLength(172);
 
     const paths = routes.map((route) => route.path);
-    expect(new Set(paths).size).toBe(130);
-    expect(new Set(routes.map((route) => `${route.method}:${route.path}`)).size).toBe(171);
+    expect(new Set(paths).size).toBe(131);
+    expect(new Set(routes.map((route) => `${route.method}:${route.path}`)).size).toBe(172);
 
     for (const route of routes) {
       expect(route.pattern.source).toBe(parsePattern(route.path).source);

packages/control-plane/src/routes/session-runtime-proxy.test.ts

@@ -1,11 +1,12 @@
-import { describe, expect, it, vi } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import { SessionInternalPaths } from "../session/contracts";
 import type { PermissionId } from "@open-inspect/shared/rbac";
 import type { RequestContext } from "./shared";
 import type { SqlDatabase } from "../db/sql-database";
 import { sessionRuntimeProxyRoutes } from "./session-runtime-proxy";
 import type { Env } from "../types";
 import { TEST_BACKGROUND_TASK_CONTEXT } from "../router.test-support";
+import { SessionIndexStore } from "../db/session-index";
 
 function createCtx(
   db: SqlDatabase = {} as SqlDatabase,
@@ -53,7 +54,63 @@ function getHandler(method: string, path: string) {
   throw new Error(`No route found for ${method} ${path}`);
 }
 
+afterEach(() => vi.restoreAllMocks());
+
 describe("session runtime proxy routes", () => {
+  it("forwards budget updates with verified user identity", async () => {
+    vi.spyOn(SessionIndexStore.prototype, "get").mockResolvedValue({
+      id: "session-1",
+      userId: "user-1",
+    } as Awaited<ReturnType<SessionIndexStore["get"]>>);
+    const requests: Request[] = [];
+    const fetch = vi.fn(async (request: Request) => {
+      requests.push(request);
+      return Response.json({ maxSessionCostUsd: 20 });
+    });
+    const path = "/sessions/session-1/budget";
+    const { handler, match, route } = getHandler("PATCH", path);
+
+    const response = await handler(
+      new Request(`https://test.local${path}`, {
+        method: "PATCH",
+        headers: { "content-type": "application/json" },
+        body: JSON.stringify({ maxCostUsd: 20 }),
+      }),
+      createEnv(fetch),
+      match,
+      createCtx({} as SqlDatabase, ["sessions.lifecycle"])
+    );
+
+    expect(response.status).toBe(200);
+    expect(route.authentication.kind).toBe("user");
+    expect(new URL(requests[0].url).pathname).toBe(SessionInternalPaths.budget);
+    await expect(requests[0].json()).resolves.toEqual({ maxCostUsd: 20 });
+  });
+
+  it("rejects budget updates from a non-owner", async () => {
+    vi.spyOn(SessionIndexStore.prototype, "get").mockResolvedValue({
+      id: "session-1",
+      userId: "owner-1",
+    } as Awaited<ReturnType<SessionIndexStore["get"]>>);
+    const fetch = vi.fn(async () => Response.json({ maxSessionCostUsd: 20 }));
+    const path = "/sessions/session-1/budget";
+    const { handler, match } = getHandler("PATCH", path);
+
+    const response = await handler(
+      new Request(`https://test.local${path}`, {
+        method: "PATCH",
+        headers: { "content-type": "application/json" },
+        body: JSON.stringify({ maxCostUsd: 20 }),
+      }),
+      createEnv(fetch),
+      match,
+      createCtx({} as SqlDatabase, ["sessions.lifecycle"])
+    );
+
+    expect(response.status).toBe(403);
+    expect(fetch).not.toHaveBeenCalled();
+  });
+
   it("forwards sandbox access for users", async () => {
     const requests: Request[] = [];
     const fetch = vi.fn(async (request: Request) => {

packages/control-plane/src/routes/session-runtime-proxy.ts

@@ -3,6 +3,7 @@ import type {
   SessionParticipantProfilesResponse,
   SessionParticipantProfile,
 } from "@open-inspect/shared/types/sessions";
+import { sessionBudgetUpdateSchema } from "@open-inspect/shared/types/session-api";
 import {
   redactSessionSnapshotSandboxAccess,
   sessionSnapshotSchema,
@@ -315,6 +316,38 @@ function lifecycleProxyRoute(
   );
 }
 
+const budgetProxyRoute = defineRoute(
+  SCM_AGNOSTIC_HUMAN_USER_ROUTE,
+  sessionRoute({
+    method: "PATCH",
+    path: "/sessions/:id/budget",
+    authorization: requirePermission("sessions.lifecycle"),
+    handler: async (request, _env, match, ctx) => {
+      const sessionId = getSessionId(match);
+      if (sessionId instanceof Response) return sessionId;
+
+      const rawBody = await parseJsonBody<unknown>(request);
+      if (rawBody instanceof Response) return rawBody;
+      if (!isObjectBody(rawBody)) return error("Invalid budget request", 400);
+
+      const parsed = sessionBudgetUpdateSchema.safeParse(rawBody);
+      if (!parsed.success) return error("Invalid budget request", 400);
+
+      const session = await new SessionIndexStore(ctx.db).get(sessionId);
+      if (!session) return error("Session not found", 404);
+      if (!ctx.authorization?.userId || session.userId !== ctx.authorization.userId) {
+        return error("Only the session owner can change the cost limit", 403);
+      }
+
+      return ctx.sessionRuntime.fetch(sessionId, SessionInternalPaths.budget, {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify(parsed.data),
+      });
+    },
+  })
+);
+
 export const sessionRuntimeProxyRoutes: Route[] = [
   simpleProxyRoute({
     policy: SCM_AGNOSTIC_HUMAN_USER_ROUTE,
@@ -432,4 +465,5 @@ export const sessionRuntimeProxyRoutes: Route[] = [
   lifecycleProxyRoute("PATCH", "/sessions/:id/title", SessionInternalPaths.updateTitle),
   lifecycleProxyRoute("POST", "/sessions/:id/archive", SessionInternalPaths.archive),
   lifecycleProxyRoute("POST", "/sessions/:id/unarchive", SessionInternalPaths.unarchive),
+  budgetProxyRoute,
 ];

packages/control-plane/src/sandbox/lifecycle/manager.test.ts

@@ -87,6 +87,10 @@ function createMockSession(overrides: Partial<SessionRow> = {}): SessionRow {
     code_server_enabled: 0,
     vnc_enabled: 0,
     total_cost: 0,
+    max_cost_usd: null,
+    cost_warning_sent: 0,
+    budget_exhausted: 0,
+    cost_tracking_unavailable: 0,
     sandbox_settings: null,
     environment_id: null,
     created_at: Date.now() - 60000,

packages/control-plane/src/sandbox/lifecycle/manager.ts

@@ -307,6 +307,7 @@ export interface SandboxLifecycle {
 export type UnresponsiveSandboxTrigger =
   | "prompt_dispatch_send_failed"
   | "stop_send_failed"
+  | "stop_alarm_failed"
   | "stop_confirmation_timeout";
 
 export type SandboxAlarmResult = "no_action" | "sandbox_failed" | "sandbox_terminated";
@@ -1462,6 +1463,7 @@ export class SandboxLifecycleManager implements SandboxLifecycle {
     const closeReason = {
       prompt_dispatch_send_failed: "Prompt dispatch send failed",
       stop_send_failed: "Stop command send failed",
+      stop_alarm_failed: "Stop confirmation alarm failed",
       stop_confirmation_timeout: "Stop confirmation timed out",
     }[trigger];
     this.wsManager.detachSandboxWebSocket(1011, closeReason);

packages/control-plane/src/sandbox/settings.test.ts

@@ -126,6 +126,36 @@ describe("normalizeSandboxSettings", () => {
     ).toEqual({ terminalEnabled: true });
   });
 
+  it("accepts valid session cost settings", () => {
+    expect(
+      normalizeSandboxSettings({ maxSessionCostUsd: 12.5, costWarningThresholdPct: 99 })
+    ).toEqual({ maxSessionCostUsd: 12.5, costWarningThresholdPct: 99 });
+  });
+
+  it.each([
+    { maxSessionCostUsd: 0 },
+    { maxSessionCostUsd: -1 },
+    { maxSessionCostUsd: Number.NaN },
+    { costWarningThresholdPct: 0 },
+    { costWarningThresholdPct: 99.5 },
+    { costWarningThresholdPct: 100 },
+  ])("rejects invalid session cost settings %#", (settings) => {
+    expect(() => normalizeSandboxSettings(settings)).toThrow(SandboxSettingsValidationError);
+  });
+
+  it("omits invalid session cost settings while preserving valid siblings", () => {
+    expect(
+      normalizeSandboxSettings(
+        {
+          maxSessionCostUsd: -1,
+          costWarningThresholdPct: 100,
+          terminalEnabled: true,
+        },
+        { invalid: "omit" }
+      )
+    ).toEqual({ terminalEnabled: true });
+  });
+
   it("accepts valid service ports", () => {
     expect(
       normalizeSandboxSettings({ codeServerPort: 8081, vncPort: 6081, terminalPort: 7000 })

packages/control-plane/src/sandbox/settings.ts

@@ -150,6 +150,28 @@ export function normalizeSandboxSettings(
     result.buildTimeoutSeconds = buildTimeoutSeconds;
   }
 
+  const maxSessionCostUsd = normalizePositiveNumberSetting(
+    settings.maxSessionCostUsd,
+    "maxSessionCostUsd",
+    reject
+  );
+  if (maxSessionCostUsd !== undefined) {
+    result.maxSessionCostUsd = maxSessionCostUsd;
+  }
+
+  const costWarningThresholdPct = normalizePositiveIntegerSetting(
+    settings.costWarningThresholdPct,
+    "costWarningThresholdPct",
+    reject
+  );
+  if (costWarningThresholdPct !== undefined) {
+    if (costWarningThresholdPct > 99) {
+      reject("costWarningThresholdPct must be an integer between 1 and 99");
+    } else {
+      result.costWarningThresholdPct = costWarningThresholdPct;
+    }
+  }
+
   checkPortCollisions(result, reject);
 
   return result;
@@ -271,3 +293,16 @@ function normalizePositiveIntegerSetting(
   }
   return value;
 }
+
+function normalizePositiveNumberSetting(
+  value: unknown,
+  name: string,
+  reject: (message: string) => false
+): number | undefined {
+  if (value === undefined) return undefined;
+  if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) {
+    reject(`${name} must be a positive finite number`);
+    return undefined;
+  }
+  return value;
+}

packages/control-plane/src/session/alarm/handler.test.ts

@@ -11,6 +11,8 @@ function createHandler() {
   };
   const messageQueue = {
     failStuckProcessingMessage: vi.fn<() => Promise<void>>().mockResolvedValue(),
+  };
+  const executionStop = {
     recoverStopConfirmationTimeout: vi.fn<() => Promise<void>>().mockResolvedValue(),
     resumeAfterSandboxTermination: vi.fn<() => Promise<void>>().mockResolvedValue(),
   };
@@ -37,6 +39,7 @@ function createHandler() {
   const handler = createAlarmHandler({
     repository: repository as unknown as MessageRepository,
     messageQueue,
+    executionStop,
     lifecycleManager,
     terminalMessageProjection,
     alarmScheduler,
@@ -49,6 +52,7 @@ function createHandler() {
     handler,
     repository,
     messageQueue,
+    executionStop,
     lifecycleManager,
     terminalMessageProjection,
     alarmScheduler,
@@ -59,28 +63,35 @@ function createHandler() {
 
 describe("createAlarmHandler", () => {
   it("delegates to lifecycle manager when no processing message exists", async () => {
-    const { handler, repository, messageQueue, lifecycleManager, alarmScheduler, now } =
-      createHandler();
+    const {
+      handler,
+      repository,
+      messageQueue,
+      executionStop,
+      lifecycleManager,
+      alarmScheduler,
+      now,
+    } = createHandler();
     repository.getProcessingMessageWithStartedAt.mockReturnValue(null);
 
     await handler.handle();
 
     expect(now).not.toHaveBeenCalled();
     expect(alarmScheduler.schedule).not.toHaveBeenCalled();
     expect(messageQueue.failStuckProcessingMessage).not.toHaveBeenCalled();
-    expect(messageQueue.recoverStopConfirmationTimeout).toHaveBeenCalledOnce();
+    expect(executionStop.recoverStopConfirmationTimeout).toHaveBeenCalledOnce();
     expect(lifecycleManager.handleAlarm).toHaveBeenCalledTimes(1);
   });
 
   it("retries a deferred terminal message projection before anything else", async () => {
-    const { handler, repository, messageQueue, terminalMessageProjection } = createHandler();
+    const { handler, repository, executionStop, terminalMessageProjection } = createHandler();
     repository.getProcessingMessageWithStartedAt.mockReturnValue(null);
 
     await handler.handle();
 
     expect(terminalMessageProjection.flushPending).toHaveBeenCalledOnce();
     expect(terminalMessageProjection.flushPending.mock.invocationCallOrder[0]).toBeLessThan(
-      messageQueue.recoverStopConfirmationTimeout.mock.invocationCallOrder[0]
+      executionStop.recoverStopConfirmationTimeout.mock.invocationCallOrder[0]
     );
   });
 
@@ -116,6 +127,7 @@ describe("createAlarmHandler", () => {
       earliest: vi.fn(() => null),
       cancelled: vi.fn(() => false),
       setPending: vi.fn(),
+      setPendingEarliest: vi.fn(),
       activate: vi.fn(),
       clear: vi.fn(),
       beginDelivery: vi.fn(() => null),
@@ -135,13 +147,16 @@ describe("createAlarmHandler", () => {
     };
     const messageQueue = {
       failStuckProcessingMessage: vi.fn<() => Promise<void>>().mockResolvedValue(),
+    };
+    const executionStop = {
       recoverStopConfirmationTimeout: vi.fn<() => Promise<void>>().mockResolvedValue(),
       resumeAfterSandboxTermination: vi.fn<() => Promise<void>>().mockResolvedValue(),
     };
 
     const handler = createAlarmHandler({
       repository: repository as unknown as MessageRepository,
       messageQueue,
+      executionStop,
       lifecycleManager,
       terminalMessageProjection: { flushPending: vi.fn(async () => {}) },
       alarmScheduler,
@@ -179,24 +194,24 @@ describe("createAlarmHandler", () => {
   });
 
   it("fails stuck work without resuming after a connecting timeout", async () => {
-    const { handler, repository, messageQueue, lifecycleManager } = createHandler();
+    const { handler, repository, messageQueue, executionStop, lifecycleManager } = createHandler();
     repository.getProcessingMessageWithStartedAt.mockReturnValue(null);
     lifecycleManager.handleAlarm.mockResolvedValue("sandbox_failed");
 
     await handler.handle();
 
     expect(messageQueue.failStuckProcessingMessage).toHaveBeenCalledOnce();
-    expect(messageQueue.resumeAfterSandboxTermination).not.toHaveBeenCalled();
+    expect(executionStop.resumeAfterSandboxTermination).not.toHaveBeenCalled();
   });
 
   it("fails stuck work and resumes after lifecycle termination", async () => {
-    const { handler, repository, messageQueue, lifecycleManager } = createHandler();
+    const { handler, repository, messageQueue, executionStop, lifecycleManager } = createHandler();
     repository.getProcessingMessageWithStartedAt.mockReturnValue(null);
     lifecycleManager.handleAlarm.mockResolvedValue("sandbox_terminated");
 
     await handler.handle();
 
     expect(messageQueue.failStuckProcessingMessage).toHaveBeenCalledOnce();
-    expect(messageQueue.resumeAfterSandboxTermination).toHaveBeenCalledOnce();
+    expect(executionStop.resumeAfterSandboxTermination).toHaveBeenCalledOnce();
   });
 });