fix(billing): carry the Pi charge onto a session its agent failed

A backend that returns a result carrying `totals.errorMessage` has already
run: the sandbox was billed and the sink holds the charge. But that path
throws instead of reaching `buildOutput`, which is what publishes the cost,
so the charge was accumulated and then dropped — lost revenue rather than
an over-charge.

Both failure paths now carry it on the error they raise, the same way the
Function handler carries its tool cost, so `handleBlockError` can pick it
up. An agent that ran and then reported a failure consumed the same tokens
and sandbox seconds as one that succeeded, which is why the cost
computation is now shared between the two rather than duplicated.

Also corrects the sink's doc comment. Local mode does fill it — the agent
runs on the caller's own machine and costs Sim nothing, but a
`function_execute` among the Sim tools it calls bills its own remote
sandbox into the same total.

The new case was confirmed to fail without the attach.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: icecrasher321

apps/sim/executor/handlers/pi/core/backend.ts

@@ -176,11 +176,15 @@ export interface PiRunContext {
   /**
    * Where a backend reports the cost of Sim-provisioned compute it used.
    *
-   * Only the cloud modes have any: they run the agent in a Sim-paid sandbox,
-   * while local mode drives the caller's own machine over SSH and costs Sim
-   * nothing. The handler folds whatever lands here into the block's `toolCost`,
-   * which is what keeps a BYOK Pi run — model unbilled by definition — from
-   * reporting no cost at all for a session that ran for tens of minutes.
+   * Both modes can fill it, from different sources. Cloud modes run the agent in
+   * a Sim-paid sandbox and report that session. Local mode drives the caller's
+   * own machine over SSH, so the agent itself costs Sim nothing — but the Sim
+   * tools it calls still run here, and a `function_execute` among them bills its
+   * own remote sandbox into the same total.
+   *
+   * The handler folds whatever lands here into the block's `toolCost`, which is
+   * what keeps a BYOK Pi run — model unbilled by definition — from reporting no
+   * cost at all for compute Sim actually paid for.
    */
   sandboxCost?: SandboxCostSink
 }

apps/sim/executor/handlers/pi/pi-handler.test.ts

@@ -114,6 +114,7 @@ vi.mock('@/blocks/utils', () => ({
 import type { PiRunContext } from '@/executor/handlers/pi/core/backend'
 import { PiBlockHandler, parsePiReviewMentions } from '@/executor/handlers/pi/pi-handler'
 import type { ExecutionContext, StreamingExecution } from '@/executor/types'
+import { readTrustedExecutionCost } from '@/executor/utils/errors'
 import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
 import type { SerializedBlock } from '@/serializer/types'
 
@@ -313,6 +314,35 @@ describe('PiBlockHandler', () => {
     expect(output.cost).toEqual({ input: 0, output: 0, toolCost: 0.0842, total: 0.0842 })
   })
 
+  it('keeps the sandbox charge on a cloud session whose agent reported an error', async () => {
+    // The backend returned, so the sandbox was billed and the sink holds the
+    // charge — but this path throws instead of reaching buildOutput, which is
+    // what would otherwise have published it.
+    mockRunCloud.mockImplementation(async (_params: unknown, context: PiRunContext) => {
+      if (context.sandboxCost) context.sandboxCost.total += 0.0631
+      return {
+        totals: { finalText: '', inputTokens: 0, outputTokens: 0, errorMessage: 'agent gave up' },
+      }
+    })
+
+    const error = await handler
+      .execute(ctx(), block, {
+        mode: 'cloud',
+        task: 'do it',
+        model: 'claude',
+        owner: 'o',
+        repo: 'r',
+        githubToken: 'ghp',
+      })
+      .catch((thrown: unknown) => thrown)
+
+    expect(error).toBeInstanceOf(Error)
+    // `toolCost` is absent by design: the trusted envelope validates exactly the
+    // three numeric fields it will let cross the handler boundary. `total` is
+    // what the ledger bills on, and it carries the sandbox charge intact.
+    expect(readTrustedExecutionCost(error)).toEqual({ input: 0, output: 0, total: 0.0631 })
+  })
+
   it('leaves a cloud run that provisioned no sandbox uncharged', async () => {
     const output = (await handler.execute(ctx(), block, {
       mode: 'cloud',

apps/sim/executor/handlers/pi/pi-handler.ts

@@ -37,6 +37,7 @@ import {
   type PiMemoryConfig,
   resolvePiSkills,
 } from '@/executor/handlers/pi/core/context'
+import type { PiRunTotals } from '@/executor/handlers/pi/core/events'
 import { streamTextForEvent } from '@/executor/handlers/pi/core/events'
 import {
   computePiCost,
@@ -54,7 +55,9 @@ import type {
   NormalizedBlockOutput,
   StreamingExecution,
 } from '@/executor/types'
+import { attachTrustedExecutionCost } from '@/executor/utils/errors'
 import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal'
+import type { ModelCost } from '@/providers/cost-policy'
 import { isPiSupportedProvider, resolvePiModelId } from '@/providers/pi-providers'
 import { getProviderFromModel } from '@/providers/utils'
 import type { SerializedBlock } from '@/serializer/types'
@@ -155,6 +158,34 @@ export function parsePiReviewMentions(value: unknown): string[] {
   return mentions
 }
 
+/**
+ * What a Pi block charges: its model tokens plus the Sim-paid sandbox compute.
+ *
+ * Sandbox cost rides in `toolCost` so it survives a BYOK run — the model side is
+ * zero by definition there, and the ledger bills a model row on `total > 0`.
+ * Folding it in is what makes a BYOK Pi session bill for the provider time it
+ * actually consumed instead of nothing at all.
+ *
+ * Shared with the failure path deliberately: an agent that ran and then reported
+ * an error consumed exactly the same tokens and sandbox seconds as one that
+ * succeeded, so both have to arrive at the same number.
+ */
+function buildPiCost(
+  model: string,
+  isBYOK: boolean,
+  totals: PiRunTotals,
+  sandboxCost: number
+): ModelCost {
+  const modelCost = computePiCost(model, totals.inputTokens, totals.outputTokens, isBYOK)
+  if (sandboxCost <= 0) return modelCost
+
+  return {
+    ...modelCost,
+    toolCost: sandboxCost,
+    total: modelCost.total + sandboxCost,
+  }
+}
+
 export class PiBlockHandler implements BlockHandler {
   canHandle(block: SerializedBlock): boolean {
     return block.metadata?.id === BlockType.PI
@@ -480,21 +511,7 @@ export class PiBlockHandler implements BlockHandler {
   ): NormalizedBlockOutput {
     const { totals } = result
     const endTime = Date.now()
-    const modelCost = computePiCost(model, totals.inputTokens, totals.outputTokens, isBYOK)
-    /*
-     * Sandbox compute rides in `toolCost` so it survives a BYOK run: the model
-     * side is zero by definition there, and the ledger bills a model row on
-     * `total > 0`. Folding it in is what makes a BYOK Pi session bill for the
-     * E2B time it actually consumed instead of nothing at all.
-     */
-    const cost =
-      sandboxCost > 0
-        ? {
-            ...modelCost,
-            toolCost: sandboxCost,
-            total: modelCost.total + sandboxCost,
-          }
-        : modelCost
+    const cost = buildPiCost(model, isBYOK, totals, sandboxCost)
     return {
       content: totals.finalText,
       model,
@@ -574,7 +591,12 @@ export class PiBlockHandler implements BlockHandler {
               sandboxCost,
             })
             if (result.totals.errorMessage) {
-              controller.error(new Error(result.totals.errorMessage))
+              const error = new Error(result.totals.errorMessage)
+              attachTrustedExecutionCost(
+                error,
+                buildPiCost(params.model, params.isBYOK, result.totals, sandboxCost.total)
+              )
+              controller.error(error)
               return
             }
             if (params.mode === 'cloud_plan' && result.totals.finalText) {
@@ -626,7 +648,19 @@ export class PiBlockHandler implements BlockHandler {
       sandboxCost,
     })
     if (result.totals.errorMessage) {
-      throw new Error(result.totals.errorMessage)
+      /*
+       * The backend returned, so the sandbox was billed and the sink holds the
+       * charge — but this throw skips `buildOutput`, which is what would have
+       * published it. Carrying the cost on the error is what keeps a session
+       * whose agent reported a failure from being run for free, the same way the
+       * Function handler carries its tool cost onto the error it raises.
+       */
+      const error = new Error(result.totals.errorMessage)
+      attachTrustedExecutionCost(
+        error,
+        buildPiCost(params.model, params.isBYOK, result.totals, sandboxCost.total)
+      )
+      throw error
     }
     if (memoryConfig) {
       await appendPiMemory(