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 ac377b99445..fea3271f66d 100644 --- a/server/src/services/native-runtime/native-session-executor.test.ts +++ b/server/src/services/native-runtime/native-session-executor.test.ts @@ -1,6 +1,7 @@ import { onRetainedNativeSessionCleanupRelease } from "../retained-native-session-cleanup-notifications.js"; import { readIdleLocalWork, markIdleStartupComplete, markIdleIngressTracked } from "../idle-local-work.js"; import * as runnerCheckpoint from "./native-runner-checkpoint.js"; +import * as stagedAttachmentRecovery from "./native-staged-attachment-recovery.js"; import * as controllerRecovery from "./native-controller-recovery.js"; import { DurablePrpControlPlane } from "../../vendor/paperclip-runner/index.js"; import * as environmentTaskRuntime from "../environment-task-runtime.js"; @@ -12317,6 +12318,55 @@ describe("runnerd provider runtime wiring", () => { ).toThrow("scope_closed"); }); + it.each(["attached", "changed-admission", "changed-template", "changed-template-digest", "changed-preparation"] as const)( + "verifies frozen staged admission after normal wake attachment preparation: %s", async mode => { + const current = scopedNativeExecution(`staged-attachment-${mode}`); + current.task.prompt = "Inspect attachment ref:original"; + const root = await initializeScopedRunnerRoot(current, "attachment-runner"); + const sourceIdentity = { runnerInstanceId: "attachment-runner", environmentLeaseId: "attachment-lease", + runId: "attachment-source", normalizedSessionId: current.session.normalizedSessionId!, turnId: "source-turn", itemId: "source-item" }; + const targetIdentity = { ...sourceIdentity, runId: current.binding.runId, turnId: "target-turn", itemId: "target-item" }; + const controller = new DurablePrpControlPlane({ stateDirectory: join(root, "control-plane"), identity: sourceIdentity, + expectedRunnerVersion: "fixture", expectedRunnerDigest: `sha256:${"a".repeat(64)}` }); + controller.stageRunAttachment(targetIdentity, { provider: { kind: "codex", instructions: "prepared attachment grant" } }); + const path = join(root, "control-plane/control-plane-state.json"); + const saved = JSON.parse(await readFile(path, "utf8")); + const canonicalTemplateDigest = saved.stagedRunAttachment.templateDigest; + if (mode === "changed-template" || mode === "changed-template-digest") { + saved.stagedRunAttachment.runAttachTemplate.provider.instructions = "substituted grant"; + if (mode === "changed-template-digest") saved.stagedRunAttachment.templateDigest = `sha256:${nativeSha256(saved.stagedRunAttachment.runAttachTemplate)}`; + await writeFile(path, JSON.stringify(saved)); + } + state.stageNativeRunnerWakeAttachments.mockImplementationOnce(async () => { + if (mode === "changed-preparation") current.task.prompt = "Inspect attachment ref:mutated-during-preparation"; + return { attachments: [{ id: "attachment-original" }], cleanup: async () => undefined } as never; + }); + state.renderNativeRunnerStagedAttachmentPrompt.mockReturnValueOnce("Materialized attachment: .paperclip-inbound/current/original"); + const admission = structuredClone(current); + if (mode === "changed-admission") admission.task.prompt = "Inspect attachment ref:changed"; + const reached = new Error("verified staged admission reached recheck"); + const restore = vi.spyOn(runnerCheckpoint, "restoreNativeRunnerCheckpoint").mockResolvedValue(null); + const header = vi.spyOn(runnerCheckpoint, "loadNativeRunnerCheckpointHeader").mockResolvedValue("staged header"); + const verify = vi.spyOn(stagedAttachmentRecovery, "resolveNativeStagedAttachmentRecovery") + .mockResolvedValueOnce({ currentExecution: admission, sourceIdentity, targetIdentity, + templateDigest: canonicalTemplateDigest } as never) + .mockRejectedValueOnce(reached); + try { + const resume = executePaperclipNativeSession({ db: leaseDb(current), execution: current, useRunnerd: true, + runnerInstanceId: "attachment-runner", restartRecovery: { kind: "reattach_remote_runner", runId: current.binding.runId, + leaseOwner: "controller", controllerGeneration: 2, providerAttempt: 1, restartKind: "hard", recoveryRequestId: null, + remote: { kind: "sandbox", leaseId: "attachment-lease", providerLeaseId: "sandbox", remoteCwd: "/workspace" } } as never, + runnerExecutionTarget: { kind: "remote", transport: "sandbox", leaseId: "attachment-lease", remoteCwd: "/workspace", + runner: { execute: vi.fn() }, sandboxLeaseAcquisition: { outcome: "resumed", providerLeaseId: "sandbox" } } as never }); + if (mode === "attached") { await expect(resume).rejects.toBe(reached); expect(verify).toHaveBeenCalledTimes(2); } + else await expect(resume).rejects.toThrow(mode === "changed-template" ? "staged_attachment_binding_mismatch" + : mode === "changed-preparation" ? "native_attachment_admission_changed" : "staged_attachment_recovery_unproven"); + expect(state.stageNativeRunnerWakeAttachments).toHaveBeenCalledWith(expect.objectContaining({ binding: expect.objectContaining({ + companyId: current.binding.companyId, issueId: current.binding.issueId, agentId: current.binding.agentId, runId: current.binding.runId, + }) })); + } finally { restore.mockRestore(); header.mockRestore(); verify.mockRestore(); await controller.stop(); } + }); + it("passes the run checkpoint active turn into restart recovery", async () => { state.createBackend.mockClear(); state.createTransport.mockClear(); diff --git a/server/src/services/native-runtime/native-session-executor.ts b/server/src/services/native-runtime/native-session-executor.ts index 8dd4cb27d68..604bff67972 100644 --- a/server/src/services/native-runtime/native-session-executor.ts +++ b/server/src/services/native-runtime/native-session-executor.ts @@ -5091,6 +5091,7 @@ async function recoverPluginRunnerCheckpoint(input: { async function migrateRunnerdStateRootForExecution(input: { db: Db; execution: NativeExecutionInput; + preparedExecutionAdmission?: NativePreparedExecutionAdmission; allowVerifiedBackup: boolean; retainedWarmOwnerRunId?: string; allowLocalRecovery: boolean; @@ -5678,8 +5679,11 @@ function readStagedRunnerdAttachment(root: string) { /** A staged attachment has not published a target session. Re-establish its * exact source and target from private durable evidence before any remote RPC. */ +type NativePreparedExecutionAdmission = Readonly<{ admissionSha256: string; preparedSha256: string }>; + async function verifyStagedAttachmentRestart(input: { db: Db; execution: NativeExecutionInput; restartRecovery?: NativeRestartRecoveryClaim; + preparedExecutionAdmission?: NativePreparedExecutionAdmission; runnerExecutionTarget?: AdapterExecutionTarget | null; }, expected?: NativeStagedAttachmentRecovery): Promise { const root = scopedRunnerdStateRoot(input.execution); @@ -5701,7 +5705,9 @@ async function verifyStagedAttachmentRestart(input: { matchesExecution: (previous, execution) => nativeSessionScopeKey(previous) === nativeSessionScopeKey(execution) && nativeSessionConfigDigest(previous, "remote") === nativeSessionConfigDigest(execution, "remote"), }); - if (!verified || nativeSha256(verified.currentExecution) !== nativeSha256(input.execution) || + const executionSha256 = nativeSha256(input.execution); + if (!verified || nativeSha256(verified.currentExecution) !== (input.preparedExecutionAdmission?.admissionSha256 ?? executionSha256) || + (input.preparedExecutionAdmission && executionSha256 !== input.preparedExecutionAdmission.preparedSha256) || verified.templateDigest !== staged.templateDigest || nativeSha256(verified.sourceIdentity) !== nativeSha256(staged.sourceIdentity) || (localStage && nativeSha256(readStagedRunnerdAttachment(root)) !== nativeSha256(localStage))) return deny(); @@ -7958,6 +7964,8 @@ export async function executePaperclipNativeSession(input: { refreshTools?: boolean; db: Db; execution: NativeExecutionInput; + /** Private server preparation proof; never persisted or accepted from a tool. */ + preparedExecutionAdmission?: NativePreparedExecutionAdmission; runnerInstanceId: string; /** Trusted task identity from the heartbeat orchestration. */ conversationMode?: boolean; @@ -8123,6 +8131,7 @@ export async function executePaperclipNativeSession(input: { const needsWarmRetirement = [...warmNativeSessions].some(([scope, entry]) => scope !== sessionScopeId && entry.ownerScope === ownerScope); const prepareAttachments = async () => { + const admissionSha256 = nativeSha256(input.execution); const targetKind = input.runnerExecutionTarget?.kind ?? "local"; const chatAttachmentReadScope = new NativeChatAttachmentReadScope({ db: input.db, @@ -8155,16 +8164,16 @@ export async function executePaperclipNativeSession(input: { const stagedPrompt = renderNativeRunnerStagedAttachmentPrompt( attachmentStage.attachments, ); + if (nativeSha256(input.execution) !== admissionSha256) throw new Error("native_attachment_admission_changed"); if (stagedPrompt) { + const preparedExecution = parseNativeExecutionInput({ + ...input.execution, + task: { ...input.execution.task, prompt: `${input.execution.task.prompt}\n\n${stagedPrompt}` }, + }); preparedInput = { ...preparedInput, - execution: parseNativeExecutionInput({ - ...input.execution, - task: { - ...input.execution.task, - prompt: `${input.execution.task.prompt}\n\n${stagedPrompt}`, - }, - }), + execution: preparedExecution, + preparedExecutionAdmission: { admissionSha256, preparedSha256: nativeSha256(preparedExecution) }, }; } return preparedInput; @@ -8384,6 +8393,7 @@ async function executePaperclipNativeSessionWithinScope( retainedTransition = await migrateRunnerdStateRootForExecution({ db: input.db, execution: input.execution, + preparedExecutionAdmission: input.preparedExecutionAdmission, allowVerifiedBackup: input.runnerExecutionTarget?.kind === "remote", retainedWarmOwnerRunId: idleWarmNativeSessionOwnerRunId(input) ?? remoteWarmRecovery?.identity.runId, allowLocalRecovery: input.runnerExecutionTarget?.kind !== "remote", @@ -12121,6 +12131,7 @@ export async function createRunnerdBackend(input: { controllerPersistence?: { bootId: string; runnerArtifact?: { version: string; digest: string } }; onAccountingEvent?: (event: PrpEvent) => Promise; execution: NativeExecutionInput; + preparedExecutionAdmission?: NativePreparedExecutionAdmission; runnerInstanceId: string; chatAttachmentReadScope?: NativeChatAttachmentReadScope; restartRecovery?: NativeRestartRecoveryClaim; @@ -12195,6 +12206,7 @@ export async function createRunnerdBackend(input: { retainedTransition = await migrateRunnerdStateRootForExecution({ db: input.db, execution: input.execution, + preparedExecutionAdmission: input.preparedExecutionAdmission, allowVerifiedBackup: input.runnerExecutionTarget?.kind === "remote", retainedWarmOwnerRunId: undefined, allowLocalRecovery: input.runnerExecutionTarget?.kind !== "remote",