fix(native): deliver plugin VM Runner usage to Core (#19162)

This commit is contained in:
Nicky Leach authored and GitHub committed 2026-11-04 21:02:22 -08:00
1 parent 53937be283
commit 8c062fa481
5 files changed
+129 -31

No files matched your search

+7
View File
@@ -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
@@ -165,6 +165,7 @@ type BackendFactoryOptions = {
type RunnerTransportOptions = {
onAccountingEvent?: (event: PrpEvent) => Promise<void>;
onRecoveryAccountingEvent?: unknown;
environment?: NodeJS.ProcessEnv;
acpxRuntimeDirectory?: string;
adoptExistingRunner?: { pid: number; isAlive: () => Promise<boolean> | 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<string, unknown> }> = [];
state.execute.mockReset().mockImplementationOnce(async () => {
@@ -1840,11 +1840,16 @@ class SessionAuthorityEpoch {
readonly runId: string;
#authority: PaperclipRunnerToolAuthority;
#revoked = false;
readonly #drainingRun?: { runId: string; capture?: (event: PrpEvent) => Promise<void> };
constructor(runId: string, authority: PaperclipRunnerToolAuthority, private readonly toolTrace?: NativeToolTrace,
private readonly onAccountingEvent?: (event: PrpEvent) => Promise<void>) {
private readonly onAccountingEvent?: (event: PrpEvent) => Promise<void>, 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<SessionAuthorityEpoch["definitions"]>
@@ -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<NonNullable<typeof recoveryUsage>>[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,
@@ -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));
@@ -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<NativeUsageDeliveryInput, "billing">) {
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<AdapterUsageCheckpoint, "provider" | "biller" | "billingType">) => ({
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<AdapterUsageCheckpoint>;
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<NativeUsageDeliveryInput, "companyId" | "runId">, sourceId: string) {
@@ -109,7 +112,7 @@ async function captureLateUsageForOwner(db: Db, input: NativeUsageDeliveryInput,
}
function createNativeUsageReceiver(db: Db, input: NativeUsageDeliveryInput, sourceId: string,
saved: Awaited<ReturnType<typeof loadNativeUsageReceipts>>, verifySaved = false) {
saved: Awaited<ReturnType<typeof loadNativeUsageReceipts>>, verifySaved = false, terminalAcknowledged = false) {
const accepted = new Map<string, UsageReceiptEnvelope>();
const turnAccounting = ["opencode", "acpx"].includes(input.provider.kind) ? createNativeTurnAccounting(input.provider) : null;
let observedUsage: Record<string, unknown> | 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;