From 8c062fa481aac9c7ef3fdc85fca1586a7cf686de Mon Sep 17 00:00:00 2001 From: Nicky Leach Date: Wed, 4 Nov 2026 21:02:22 -0800 Subject: [PATCH] fix(native): deliver plugin VM Runner usage to Core (#19162) --- doc/execution-semantics.md | 7 +++ .../native-session-executor.test.ts | 41 +++++++++++++- .../native-runtime/native-session-executor.ts | 56 +++++++++++-------- .../native-usage-delivery.integration.test.ts | 38 +++++++++++++ .../native-runtime/native-usage-delivery.ts | 18 +++--- 5 files changed, 129 insertions(+), 31 deletions(-) diff --git a/doc/execution-semantics.md b/doc/execution-semantics.md index e3f8e6a1e48..03845bd8fef 100644 --- a/doc/execution-semantics.md +++ b/doc/execution-semantics.md @@ -1727,6 +1727,13 @@ session has no durable path for this, so its late usage stays with the run that active when it arrives, or the next run. The total stays exact, but the attribution can be off by those tokens. +Every remote Runner transport, including one on a plugin-owned VM, delivers usage to +Core as events commit. Reports still in the Runner's outbox when a run finishes +settle on that run when the next run drains them, even if it has no receipts yet. +Codex reports the whole run's delta, so one such report can settle a finished run +whose terminal event was already acknowledged. Other providers report per turn and +still need their terminal event. + ### Session settlement after a run Attaching a later run to a session asks whether anything is live or waiting: an diff --git a/server/src/services/native-runtime/native-session-executor.test.ts b/server/src/services/native-runtime/native-session-executor.test.ts index e69c2b0876d..d0bb518a579 100644 --- a/server/src/services/native-runtime/native-session-executor.test.ts +++ b/server/src/services/native-runtime/native-session-executor.test.ts @@ -165,6 +165,7 @@ type BackendFactoryOptions = { type RunnerTransportOptions = { onAccountingEvent?: (event: PrpEvent) => Promise; + onRecoveryAccountingEvent?: unknown; environment?: NodeJS.ProcessEnv; acpxRuntimeDirectory?: string; adoptExistingRunner?: { pid: number; isAlive: () => Promise | boolean }; @@ -7905,13 +7906,51 @@ describe("native terminal-turn accounting", () => { await retainedHook(secondReport); expect(secondCapture).toHaveBeenCalledExactlyOnceWith(secondReport); expect(firstCapture).toHaveBeenCalledTimes(1); - await expect(retainedHook(firstReport)).rejects.toThrow("native_usage_run_binding_mismatch"); + await retainedHook(firstReport); + expect(firstCapture).toHaveBeenCalledTimes(2); + expect(secondCapture).toHaveBeenCalledTimes(1); + await expect(retainedHook({ ...firstReport, runId: "unrelated-run" })).rejects.toThrow("native_usage_run_binding_mismatch"); } finally { if (previous === undefined) delete process.env.PAPERCLIP_RUNNER_STATE_DIR; else process.env.PAPERCLIP_RUNNER_STATE_DIR = previous; await rm(directory, { recursive: true, force: true }); } }); + it("delivers plugin task usage to the producer that owns the run", async () => { + const directory = await mkdtemp(join(tmpdir(), "native-plugin-usage-")); + const previous = process.env.PAPERCLIP_RUNNER_STATE_DIR; + process.env.PAPERCLIP_RUNNER_STATE_DIR = directory; + const current = { ...execution, provider: { kind: "codex", model: "gpt-6-astra" }, + binding: { ...execution.binding, runId: "plugin-usage-run", executionWorkspaceId: "plugin-usage-workspace" }, + session: { ...execution.session, normalizedSessionId: "plugin-usage-session" } } as NativeExecutionInput; + const launcher = vi.fn(), registration = vi.fn(); + const driver = { prepareLaunch: vi.fn(), runnerProcessLauncher: launcher, controlPlaneRegistration: registration }; + const owner = vi.spyOn(pluginTaskRuntime, "createPluginRunnerTaskOwner").mockReturnValue( + { launcher, registration, driver, readWorkspaceFile: vi.fn(), isRetired: () => false, rebind: vi.fn() } as never); + const task = vi.fn(async (operation: unknown) => { + environmentTaskOperationSchema.parse(operation); + return { kind: "description", launchPreparationVersion: 1, projectAuthCheckVersion: 1, taskId: "task", + runner: { version: "fixture", digest: `sha256:${"a".repeat(64)}` } }; + }); + const capture = vi.fn(async (_event: PrpEvent) => {}); + state.createBackend.mockClear(); state.createTransport.mockClear(); + try { + await createRunnerdBackend({ db: leaseDb(current), execution: current, runnerInstanceId: "plugin-producer", runnerIngressAuthorized: true, + onAccountingEvent: capture, runnerExecutionTarget: { kind: "remote", transport: "sandbox", providerKey: "plugin", + leaseId: "plugin-usage-lease", remoteCwd: "/unused", runnerTask: task } }); + state.createBackend.mock.calls.at(-1)![1].codexTransportFactory!(); + const options = state.createTransport.mock.calls.at(-1)![0]; + const report = { eventType: "usage.reported", runId: current.binding.runId, sourceInstanceId: "plugin-producer" } as PrpEvent; + await options.onAccountingEvent!(report); + expect(capture).toHaveBeenCalledExactlyOnceWith(report); + expect(options.onRecoveryAccountingEvent).toEqual(expect.any(Function)); + } finally { + owner.mockRestore(); + if (previous === undefined) delete process.env.PAPERCLIP_RUNNER_STATE_DIR; else process.env.PAPERCLIP_RUNNER_STATE_DIR = previous; + await rm(directory, { recursive: true, force: true }); + } + }); + it("uses the canonical recorder for partial and completed usage when no onUsage callback is supplied", async () => { const updates: Array<{ table: unknown; values: Record }> = []; state.execute.mockReset().mockImplementationOnce(async () => { diff --git a/server/src/services/native-runtime/native-session-executor.ts b/server/src/services/native-runtime/native-session-executor.ts index 54ed773ea8b..d1e9d671973 100644 --- a/server/src/services/native-runtime/native-session-executor.ts +++ b/server/src/services/native-runtime/native-session-executor.ts @@ -1840,11 +1840,16 @@ class SessionAuthorityEpoch { readonly runId: string; #authority: PaperclipRunnerToolAuthority; #revoked = false; + readonly #drainingRun?: { runId: string; capture?: (event: PrpEvent) => Promise }; constructor(runId: string, authority: PaperclipRunnerToolAuthority, private readonly toolTrace?: NativeToolTrace, - private readonly onAccountingEvent?: (event: PrpEvent) => Promise) { + private readonly onAccountingEvent?: (event: PrpEvent) => Promise, previous?: SessionAuthorityEpoch) { this.runId = runId; this.#authority = authority; + if (previous) { + this.#drainingRun = previous.runId === runId ? previous.#drainingRun + : { runId: previous.runId, capture: previous.onAccountingEvent }; + } } revoke(): void { @@ -1861,8 +1866,8 @@ class SessionAuthorityEpoch { async acceptAccountingEvent(event: PrpEvent) { if (!isNativeUsageAccountingEvent(event)) return; this.#assertCurrent(); - if (event.runId !== this.runId) throw new Error("native_usage_run_binding_mismatch"); - await this.onAccountingEvent?.(event); + if (event.runId !== this.runId && event.runId !== this.#drainingRun?.runId) throw new Error("native_usage_run_binding_mismatch"); + await (event.runId === this.runId ? this.onAccountingEvent : this.#drainingRun?.capture)?.(event); this.#assertCurrent(); } @@ -13300,6 +13305,7 @@ async function createRunnerdBackendWithinSessionClaim( authority, input.toolTrace, input.onAccountingEvent, + sessionToolAuthorityEpochs.get(sessionScopeId), ); let dynamicTools: Awaited< ReturnType @@ -13342,6 +13348,28 @@ async function createRunnerdBackendWithinSessionClaim( nativeSessionConfigDigest(previous, target.kind) === nativeSessionConfigDigest(input.execution, target.kind), }) : undefined; + const accountingTransportOptions = { + onAccountingEvent: input.onAccountingEvent ? async (event: PrpEvent) => { + const current = sessionToolAuthorityEpochs.get(sessionScopeId); + if (!current) throw new Error("native_session_authority_unavailable"); + if (input.remoteWarmRecovery && event.runId === input.remoteWarmRecovery.identity.runId) { + if (!recoveryUsage) throw new Error("native_recovery_usage_source_unproven"); + const source = input.remoteWarmRecovery.identity; + await current.acceptRecoveryAccountingEvent(event, () => recoveryUsage(event, { + runnerInstanceId: source.runnerInstanceId, environmentLeaseId: source.environmentLeaseId, + runId: source.runId, normalizedSessionId: source.normalizedSessionId, + turnId: source.turnId, itemId: source.itemId, + })); + return; + } + await current.acceptAccountingEvent(event); + } : undefined, + onRecoveryAccountingEvent: input.onAccountingEvent ? async (event: PrpEvent, source: Parameters>[1]) => { + const current = sessionToolAuthorityEpochs.get(sessionScopeId); + if (!current || !recoveryUsage) throw new Error("native_recovery_usage_source_unproven"); + await current.acceptRecoveryAccountingEvent(event, () => recoveryUsage(event, source)); + } : undefined, + }; const effectiveRunnerInstanceId = durableBinding?.runnerInstanceId ?? input.runnerInstanceId; const effectiveEnvironmentLeaseId = @@ -13448,6 +13476,7 @@ async function createRunnerdBackendWithinSessionClaim( return nativeCompletionFeedback(input.db, current.runId, result); }, codexTransportFactory: recovery => createRunnerdCodexTransport({ provider: provider.kind, stateDirectory: root, + ...accountingTransportOptions, ...(input.recoveryPurpose === "stop_cleanup" ? { recoveryPurpose: "stop_cleanup" as const } : {}), measureOperation: process.env.OTEL_EXPORTER_OTLP_ENDPOINT?.trim() ? measureRunnerOperation : undefined, recordTimingReceipt: process.env.OTEL_EXPORTER_OTLP_ENDPOINT?.trim() ? recordRunnerTimingReceipt : undefined, @@ -14830,26 +14859,7 @@ async function createRunnerdBackendWithinSessionClaim( onSpawn: input.onSpawn, ...(input.execution.provider.kind === "codex" ? { disableCodexWebSearch: "webSearch" in input.execution.provider && input.execution.provider.webSearch === "disabled" } : {}), homeMode: preconfiguredHome ? "preconfigured" : "isolated", - onAccountingEvent: input.onAccountingEvent ? async event => { - const current = sessionToolAuthorityEpochs.get(sessionScopeId); - if (!current) throw new Error("native_session_authority_unavailable"); - if (input.remoteWarmRecovery && event.runId === input.remoteWarmRecovery.identity.runId) { - if (!recoveryUsage) throw new Error("native_recovery_usage_source_unproven"); - const source = input.remoteWarmRecovery.identity; - await current.acceptRecoveryAccountingEvent(event, () => recoveryUsage(event, { - runnerInstanceId: source.runnerInstanceId, environmentLeaseId: source.environmentLeaseId, - runId: source.runId, normalizedSessionId: source.normalizedSessionId, - turnId: source.turnId, itemId: source.itemId, - })); - return; - } - await current.acceptAccountingEvent(event); - } : undefined, - onRecoveryAccountingEvent: input.onAccountingEvent ? async (event, source) => { - const current = sessionToolAuthorityEpochs.get(sessionScopeId); - if (!current || !recoveryUsage) throw new Error("native_recovery_usage_source_unproven"); - await current.acceptRecoveryAccountingEvent(event, () => recoveryUsage(event, source)); - } : undefined, + ...accountingTransportOptions, onCredentialsUpdated: (event) => appendNativeCredentialUpdate(input.db, { companyId: input.execution.binding.companyId, agentId: input.execution.binding.agentId, diff --git a/server/src/services/native-runtime/native-usage-delivery.integration.test.ts b/server/src/services/native-runtime/native-usage-delivery.integration.test.ts index f49ff9b41ec..af9998ef7e3 100644 --- a/server/src/services/native-runtime/native-usage-delivery.integration.test.ts +++ b/server/src/services/native-runtime/native-usage-delivery.integration.test.ts @@ -59,6 +59,44 @@ describe("producer-owned native usage", () => { return { input, source, report, done }; }; + it.each([ + { name: "a report whose terminal was already acknowledged", types: ["usage.reported"] }, + { name: "a report and its terminal", types: ["usage.reported", "turn.completed"] }, + ])("settles a finished run's outbox with no receipts yet: $name", async ({ types }) => { + const input = await seed(); + await db.update(heartbeatRuns).set({ runnerInstanceId: input.runnerInstanceId }).where(eq(heartbeatRuns.id, input.runId)); + await createNativeUsageDelivery(db, input); + const { billing, ...source } = input; + await db.update(heartbeatRuns).set({ status: "succeeded", finishedAt: new Date(), + usageJson: sql`${heartbeatRuns.usageJson} || ${JSON.stringify(billing)}::jsonb` }).where(eq(heartbeatRuns.id, input.runId)); + const continuation = await createNativeUsageContinuationDelivery(db, source); + for (const [index, type] of types.entries()) await continuation.capture(event(input, index + 1, type)); + expect(await accountRunCost(db, input.runId)).toBe(true); + expect(await db.select().from(costEvents).where(eq(costEvents.heartbeatRunId, input.runId))) + .toEqual([expect.objectContaining({ biller: billing.biller, billingType: billing.billingType, outputTokens: 10 })]); + }); + + it("keeps a lone report open while the run has not finished", async () => { + const input = await seed(); + await db.update(heartbeatRuns).set({ runnerInstanceId: input.runnerInstanceId }).where(eq(heartbeatRuns.id, input.runId)); + await createNativeUsageDelivery(db, input); + const { billing, ...source } = input; + await db.update(heartbeatRuns).set({ usageJson: sql`${heartbeatRuns.usageJson} || ${JSON.stringify(billing)}::jsonb` }) + .where(eq(heartbeatRuns.id, input.runId)); + const continuation = await createNativeUsageContinuationDelivery(db, source); + await continuation.capture(event(input, 1, "usage.reported")); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, input.runId)); + expect(run!.usageJson).toMatchObject({ accountingReceiptReady: false }); + }); + + it("refuses to continue a run with no receipts and no recorded classification", async () => { + const input = await seed(); + await db.update(heartbeatRuns).set({ runnerInstanceId: input.runnerInstanceId }).where(eq(heartbeatRuns.id, input.runId)); + await createNativeUsageDelivery(db, input); + const { billing: _billing, ...source } = input; + await expect(createNativeUsageContinuationDelivery(db, source)).rejects.toThrow("native_usage_continuation_receipts_missing"); + }); + it("continues exact finalized source receipts without acquiring ownership or resetting accounting", async () => { const { input, source, report, done } = await finalizedSource(); const before = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, input.runId)); diff --git a/server/src/services/native-runtime/native-usage-delivery.ts b/server/src/services/native-runtime/native-usage-delivery.ts index 412cfa51c29..364b7319c18 100644 --- a/server/src/services/native-runtime/native-usage-delivery.ts +++ b/server/src/services/native-runtime/native-usage-delivery.ts @@ -46,20 +46,23 @@ export async function createNativeUsageDelivery(db: Db, input: NativeUsageDelive * and reconciliation rules; this grants no business/tool execution authority. */ export async function createNativeUsageContinuationDelivery(db: Db, input: Omit) { const sourceId = uuid(`${input.runId}:${input.runnerInstanceId}`); - const [source] = await db.select({ runnerInstanceId: heartbeatRuns.runnerInstanceId, usage: heartbeatRuns.usageJson }) + const [source] = await db.select({ runnerInstanceId: heartbeatRuns.runnerInstanceId, usage: heartbeatRuns.usageJson, finishedAt: heartbeatRuns.finishedAt }) .from(heartbeatRuns).where(and(eq(heartbeatRuns.companyId, input.companyId), eq(heartbeatRuns.id, input.runId))); if (source?.runnerInstanceId !== input.runnerInstanceId || source.usage?.accountingReceiptOwner !== "native_producer" || source.usage.accountingReceiptSourceId !== sourceId) throw new Error("native_usage_continuation_source_mismatch"); const saved = await loadNativeUsageReceipts(db, input, sourceId); - if (!saved.length) throw new Error("native_usage_continuation_receipts_missing"); - const classification = (receipt: AdapterUsageCheckpoint) => ({ + const classification = (receipt: Pick) => ({ provider: receipt.provider, biller: receipt.biller, billingType: receipt.billingType, }); const receipts = saved.map(row => (row.receipt as UsageReceiptEnvelope).receipt as AdapterUsageCheckpoint); - const billing = classification(receipts[0]); + const recorded = source.usage as Partial; + if (!receipts.length && !recorded.biller) throw new Error("native_usage_continuation_receipts_missing"); + const billing = classification(receipts[0] ?? recorded as AdapterUsageCheckpoint); if (receipts.some(receipt => receiptFingerprint(classification(receipt)) !== receiptFingerprint(billing) || receipt.model !== (input.provider.model ?? "unknown"))) throw new Error("native_usage_source_conflict"); - return createNativeUsageReceiver(db, { ...input, billing }, sourceId, saved, true); + const terminalAcknowledged = source.finishedAt !== null && input.provider.kind === "codex" + && !saved.some(row => terminal.has((row.receipt as UsageReceiptEnvelope).producerEvent?.eventType ?? "")); + return createNativeUsageReceiver(db, { ...input, billing }, sourceId, saved, true, terminalAcknowledged); } function loadNativeUsageReceipts(db: Db, input: Pick, sourceId: string) { @@ -109,7 +112,7 @@ async function captureLateUsageForOwner(db: Db, input: NativeUsageDeliveryInput, } function createNativeUsageReceiver(db: Db, input: NativeUsageDeliveryInput, sourceId: string, - saved: Awaited>, verifySaved = false) { + saved: Awaited>, verifySaved = false, terminalAcknowledged = false) { const accepted = new Map(); const turnAccounting = ["opencode", "acpx"].includes(input.provider.kind) ? createNativeTurnAccounting(input.provider) : null; let observedUsage: Record | null = null; @@ -133,7 +136,8 @@ function createNativeUsageReceiver(db: Db, input: NativeUsageDeliveryInput, sour observedUsage = interpretedUsage; observedTurn = normalizeNativeUsage(observedUsage, { inputIncludesCacheReads: input.provider.kind === "codex" }) ? accountingTurn : undefined; } - return { usage: observedUsage, complete: terminal.has(fact.eventType) && observedTurn !== undefined && observedTurn === accountingTurn }; + const closesRun = terminal.has(fact.eventType) || (terminalAcknowledged && interpretedUsage !== undefined); + return { usage: observedUsage, complete: closesRun && observedTurn !== undefined && observedTurn === accountingTurn }; }; for (const row of saved) { const envelope = row.receipt as UsageReceiptEnvelope;