feat(chat): Enforce a hard conversation spend cap

packages/junior/src/chat/agent/index.ts

@@ -13,6 +13,7 @@
 import { Agent, type AgentLoopTurnUpdate } from "@earendil-works/pi-agent-core";
 import type { FileUpload } from "chat";
 import { botConfig } from "@/chat/config";
+import { getConversationStore } from "@/chat/db";
 import {
   extractGenAiUsageAttributes,
   extractGenAiUsageSummary,
@@ -74,6 +75,11 @@ import {
 } from "@/chat/services/turn-result";
 import { isProviderRetryError } from "@/chat/services/provider-error";
 import { nextProviderRetry } from "@/chat/services/provider-retry";
+import {
+  enforceTurnSpendLimit,
+  isTurnSpendLimitError,
+  TURN_SPEND_LIMIT_RESPONSE,
+} from "@/chat/services/spend-limit";
 import { annotateTurnDeadlineToolResult } from "@/chat/tool-support/turn-deadline-result";
 import {
   configuredTurnRoute,
@@ -269,8 +275,11 @@ async function executeAgentRunInPrivacyContext(
   let lastKnownSandboxRef = state.sandboxRef;
   let mcpToolManager: McpToolManager | undefined;
   let connectedMcpProviders = new Set<string>();
+  let agent: Agent | undefined;
+  let spendLimitPiMessages: PiMessage[] | undefined;
   let turnUsage: AgentTurnUsage | undefined;
   let priorPhaseUsage: AgentTurnUsage | undefined;
+  let retryUsage: AgentTurnUsage | undefined;
   const configuredReasoningLevel =
     policy.reasoningLevel ?? botConfig.reasoningLevel;
   let turnRoute: TurnRoute | undefined = configuredReasoningLevel
@@ -379,6 +388,13 @@ async function executeAgentRunInPrivacyContext(
       conversationId,
       turnId,
     });
+    const conversationUsage = (
+      await getConversationStore().get({ conversationId })
+    )?.usage;
+    enforceTurnSpendLimit({
+      maxSpendUsd: botConfig.maxSpendUsd,
+      usage: conversationUsage,
+    });
     // Mirror the committed provenance prefix the turn session record owns. A
     // fresh run may already include batched parked input committed before the
     // agent starts, then adds the current actor's turn-start instruction.
@@ -566,7 +582,6 @@ async function executeAgentRunInPrivacyContext(
     const generatedFiles: FileUpload[] = [];
     const artifactStatePatch: Partial<ThreadArtifactsState> = {};
     const toolCalls: string[] = [];
-    let agent: Agent | undefined;
     // Handoff becomes live only after its replacement epoch commits. This
     // pending value then drives the one-way model/context swap at Pi's boundary.
     let pendingHandoff:
@@ -598,6 +613,26 @@ async function executeAgentRunInPrivacyContext(
       );
       return hasAgentTurnUsage(usage) ? usage : undefined;
     };
+    const enforceSpendLimit = (
+      messages: PiMessage[],
+      latestUsage?: AgentTurnUsage,
+    ): void => {
+      const currentBoundaryUsage = usageSinceCurrentBoundary(messages);
+      spendLimitPiMessages = [...messages];
+      turnUsage = addAgentTurnUsage(
+        priorPhaseUsage,
+        retryUsage,
+        currentBoundaryUsage,
+      );
+      enforceTurnSpendLimit({
+        maxSpendUsd: botConfig.maxSpendUsd,
+        usage: latestUsage,
+      });
+      enforceTurnSpendLimit({
+        maxSpendUsd: botConfig.maxSpendUsd,
+        usage: addAgentTurnUsage(conversationUsage, turnUsage),
+      });
+    };
     /** Commit the durable handoff epoch before staging its in-memory model swap. */
     const scheduleHandoff = async (args: {
       profile: ModelProfile;
@@ -850,6 +885,7 @@ async function executeAgentRunInPrivacyContext(
     let assistantMessageDeliveryError:
       | AssistantMessageDeliveryError
       | undefined;
+    let pendingTurnSpendLimitError: Error | undefined;
     /** Deliver one completed tool-free visible message before the agent advances. */
     const deliverAssistantMessage = async (
       message: Parameters<typeof extractAssistantText>[0],
@@ -943,6 +979,11 @@ async function executeAgentRunInPrivacyContext(
       streamFn: createTracedStreamFn({ conversationPrivacy }),
       steeringMode: "all",
       beforeToolCall: async ({ assistantMessage }) => {
+        const usage = extractGenAiUsageSummary(assistantMessage);
+        enforceSpendLimit(
+          currentAgentMessages(),
+          hasAgentTurnUsage(usage) ? usage : undefined,
+        );
         const toolCalls = assistantMessage.content.filter(
           (part) => part.type === "toolCall",
         );
@@ -964,6 +1005,7 @@ async function executeAgentRunInPrivacyContext(
           : undefined,
       prepareNextTurnWithContext: async (nextTurn, hookSignal) => {
         try {
+          enforceSpendLimit(nextTurn.context.messages as PiMessage[]);
           const handoffUpdate = applyPendingHandoff();
           const pendingMessages = await drainSteeringMessages();
           const capacityUpdate = await applyActiveContextCompaction(
@@ -1016,10 +1058,17 @@ async function executeAgentRunInPrivacyContext(
         ) {
           return;
         }
-        const containsHandoff = event.message.content.some(
-          (part) => part.type === "toolCall" && part.name === HANDOFF_TOOL_NAME,
-        );
-        if (containsHandoff) {
+        try {
+          const usage = extractGenAiUsageSummary(event.message);
+          enforceSpendLimit(
+            currentAgentMessages(),
+            hasAgentTurnUsage(usage) ? usage : undefined,
+          );
+        } catch (error) {
+          if (!isTurnSpendLimitError(error)) {
+            throw error;
+          }
+          pendingTurnSpendLimitError = error;
           return;
         }
         return deliverAssistantMessage(event.message);
@@ -1220,14 +1269,17 @@ async function executeAgentRunInPrivacyContext(
               freshPromptMessage,
             ]);
           }
+          enforceSpendLimit(currentAgentMessages());
           run =
             shouldPromptAgent && !capacityUpdate
               ? agent!.prompt(freshPromptMessage)
               : agent!.continue();
-          let retryUsage: AgentTurnUsage | undefined;
           try {
             for (let attempt = 0; ; attempt += 1) {
               await runAgentStep(run);
+              if (pendingTurnSpendLimitError) {
+                throw pendingTurnSpendLimitError;
+              }
               if (assistantMessageDeliveryError) {
                 throw assistantMessageDeliveryError;
               }
@@ -1247,6 +1299,9 @@ async function executeAgentRunInPrivacyContext(
               const currentUsage = hasAgentTurnUsage(usageSummary)
                 ? usageSummary
                 : undefined;
+              if (lastAssistant?.stopReason === "error") {
+                enforceSpendLimit(currentAgentMessages(), currentUsage);
+              }
               const currentPhaseUsage = addAgentTurnUsage(
                 retryUsage,
                 currentUsage,
@@ -1367,6 +1422,35 @@ async function executeAgentRunInPrivacyContext(
       result,
     };
   } catch (error) {
+    if (isTurnSpendLimitError(error)) {
+      if (delivery) {
+        await delivery.onAssistantMessage({ text: TURN_SPEND_LIMIT_RESPONSE });
+      }
+      return {
+        status: "completed",
+        result: {
+          text: TURN_SPEND_LIMIT_RESPONSE,
+          sandboxRef: lastKnownSandboxRef,
+          piMessages: spendLimitPiMessages,
+          diagnostics: {
+            outcome: "success",
+            modelId: activeModelId,
+            assistantMessageCount: 1,
+            ...(turnRoute
+              ? {
+                  reasoningLevel: turnRoute.reasoningLevel,
+                }
+              : {}),
+            toolCalls: [],
+            toolResultCount: 0,
+            toolErrorCount: 0,
+            usedPrimaryText: true,
+            durationMs: Date.now() - replyStartedAtMs,
+            usage: turnUsage,
+          },
+        },
+      };
+    }
     if (
       error instanceof AssistantMessageDeliveryError &&
       !(error.originalError instanceof RetryableDeliveryError)

packages/junior/src/chat/config.ts

@@ -50,6 +50,7 @@ export interface BotConfig {
   fastModelId: string;
   imageGenerationModelId: string;
   loadingMessages: string[];
+  maxSpendUsd?: number;
   profiles: Readonly<Record<string, ExecutionProfileConfig>>;
   reasoningLevel?: TurnReasoningLevel;
   visionModelId?: string;
@@ -148,6 +149,22 @@ function parseLoadingMessages(rawValue: string | undefined): string[] {
   });
 }
 
+function parseOptionalPositiveNumber(
+  envName: string,
+  rawValue: string | undefined,
+): number | undefined {
+  const trimmed = toOptionalTrimmed(rawValue);
+  if (trimmed === undefined) {
+    return undefined;
+  }
+
+  const value = Number(trimmed);
+  if (!Number.isFinite(value) || value <= 0) {
+    throw new Error(`${envName} must be a positive number`);
+  }
+  return value;
+}
+
 function parseOptionalPositiveInteger(
   envName: string,
   rawValue: string | undefined,
@@ -317,6 +334,7 @@ function readBotConfig(
       validateEmbeddingModelId(env.AI_EMBEDDING_MODEL) ??
       DEFAULT_EMBEDDING_MODEL_ID,
     loadingMessages: parseLoadingMessages(env.JUNIOR_LOADING_MESSAGES),
+    maxSpendUsd: parseOptionalPositiveNumber("MAX_SPEND", env.MAX_SPEND),
     visionModelId: validateGatewayModelId(env.AI_VISION_MODEL),
     maxSlicesPerTurn: MAX_SLICES_PER_TURN,
     turnTimeoutMs: parseAgentTurnTimeoutMs(

packages/junior/src/chat/conversations/sql/store.ts

@@ -335,6 +335,7 @@ function conversationFromRow(readRow: ConversationReadRow): Conversation {
     ...(row.channelName ? { channelName: row.channelName } : {}),
     ...(source ? { source } : {}),
     ...(row.title ? { title: row.title } : {}),
+    ...(row.usage ? { usage: row.usage } : {}),
     ...(msFromDate(row.transcriptPurgedAt) !== undefined
       ? { transcriptPurgedAtMs: msFromDate(row.transcriptPurgedAt) }
       : {}),

packages/junior/src/chat/conversations/store.ts

@@ -46,6 +46,8 @@ export interface Conversation {
   source?: ConversationSource;
   title?: string;
   updatedAtMs: number;
+  /** Cumulative model usage used by reporting and runtime controls. */
+  usage?: AgentTurnUsage;
   /**
    * When retention purged this conversation's content. Set means messages and
    * events were deleted wholesale; reporting presents the transcript as expired

packages/junior/src/chat/services/spend-limit.ts

@@ -0,0 +1,58 @@
+import type { AgentTurnUsage } from "@/chat/usage";
+
+export const TURN_SPEND_LIMIT_RESPONSE =
+  "I stopped this conversation because it reached its configured spend limit.";
+
+/** Terminal failure raised when one agent turn reaches its configured USD cap. */
+export class TurnSpendLimitExceededError extends Error {
+  constructor(readonly maxSpendUsd: number) {
+    super(`Agent turn reached spend limit ($${maxSpendUsd.toFixed(6)})`);
+    this.name = "TurnSpendLimitExceededError";
+  }
+}
+
+/** Terminal failure raised when a configured cap cannot verify provider spend. */
+export class TurnSpendCostUnavailableError extends Error {
+  constructor() {
+    super("Agent turn provider usage omitted cost data");
+    this.name = "TurnSpendCostUnavailableError";
+  }
+}
+
+/** Return whether an error should stop the turn with the static spend-limit response. */
+export function isTurnSpendLimitError(
+  error: unknown,
+): error is TurnSpendLimitExceededError | TurnSpendCostUnavailableError {
+  return (
+    error instanceof TurnSpendLimitExceededError ||
+    error instanceof TurnSpendCostUnavailableError
+  );
+}
+
+/** Return the provider-reported total USD cost, deriving it from components when needed. */
+export function agentTurnCostUsd(usage: AgentTurnUsage | undefined): number {
+  if (!usage?.cost) return 0;
+  if (usage.cost.total !== undefined) return usage.cost.total;
+  return (
+    (usage.cost.input ?? 0) +
+    (usage.cost.output ?? 0) +
+    (usage.cost.cacheRead ?? 0) +
+    (usage.cost.cacheWrite ?? 0)
+  );
+}
+
+/** Throw once reported turn cost reaches the cap or cannot be verified. */
+export function enforceTurnSpendLimit(args: {
+  maxSpendUsd: number | undefined;
+  usage: AgentTurnUsage | undefined;
+}): void {
+  if (args.maxSpendUsd === undefined || args.usage === undefined) {
+    return;
+  }
+  if (!args.usage.cost || Object.keys(args.usage.cost).length === 0) {
+    throw new TurnSpendCostUnavailableError();
+  }
+  if (agentTurnCostUsd(args.usage) >= args.maxSpendUsd) {
+    throw new TurnSpendLimitExceededError(args.maxSpendUsd);
+  }
+}

packages/junior/src/chat/usage.ts

@@ -91,10 +91,16 @@ export function addAgentTurnUsage(
       reasoningTokens = (reasoningTokens ?? 0) + reasoning;
     }
     if (usage.cost) {
-      for (const field of [...COST_COMPONENT_FIELDS, "total"] as const) {
+      let componentCostTotal: number | undefined;
+      for (const field of COST_COMPONENT_FIELDS) {
         const value = getFiniteCost(usage.cost[field]);
         if (value === undefined) continue;
         cost[field] = addCost(cost[field], value);
+        componentCostTotal = addCost(componentCostTotal, value);
+      }
+      const total = getFiniteCost(usage.cost.total) ?? componentCostTotal;
+      if (total !== undefined) {
+        cost.total = addCost(cost.total, total);
       }
     }
     const usageComponentTotal = getComponentTotal(usage);

packages/junior/tests/component/config/chat-config.test.ts

@@ -380,6 +380,25 @@ describe("chat config", () => {
     expect(botConfig.maxSlicesPerTurn).toBe(100);
   });
 
+  it("leaves the spend cap disabled when MAX_SPEND is unset", async () => {
+    delete process.env.MAX_SPEND;
+    const { botConfig } = await loadConfig();
+    expect(botConfig.maxSpendUsd).toBeUndefined();
+  });
+
+  it("parses MAX_SPEND as a positive USD amount", async () => {
+    process.env.MAX_SPEND = "1.25";
+    const { botConfig } = await loadConfig();
+    expect(botConfig.maxSpendUsd).toBe(1.25);
+  });
+
+  it("rejects an invalid MAX_SPEND", async () => {
+    process.env.MAX_SPEND = "0";
+    await expect(loadConfig()).rejects.toThrow(
+      "MAX_SPEND must be a positive number",
+    );
+  });
+
   it("uses default AGENT_TURN_TIMEOUT_MS when env var is unset", async () => {
     delete process.env.AGENT_TURN_TIMEOUT_MS;
     const { botConfig } = await loadConfig();