diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index fcc30286453..c72e1dff56b 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -1984,12 +1984,12 @@ console.log(JSON.stringify({ type: "result", subtype: "success", result: "hello" expect(execute).not.toHaveBeenCalled(); } finally { execute.mockRestore(); } }); - it("defers a new plugin task's capability check to native launch and still checks retained tasks", async () => { + it("defers new and retained plugin task capability checks to native launch", async () => { const runnerTask = vi.fn().mockResolvedValue({ kind: "description", taskId: "task", projectAuthCheckVersion: 1, runner: { version: "fixture", digest: `sha256:${"a".repeat(64)}` } }); const target = { kind: "remote", transport: "sandbox", providerKey: "plugin", remoteCwd: ".", runnerTask, sandboxLeaseAcquisition: { outcome: "created", providerLeaseId: "task" } } as const; - const options = { deferNewRunnerTaskCapabilityCheck: true }; + const options = { deferRunnerTaskCapabilityCheck: true }; await expect(assertManagedAiProjectAuth({}, "openai", target, options)).resolves.toBeUndefined(); await expect(assertManagedAiProjectAuth({ args: ["--api-key=override"] }, "openai", target, options)).rejects.toThrow("overrides"); @@ -1997,6 +1997,8 @@ console.log(JSON.stringify({ type: "result", subtype: "success", result: "hello" const retained = { ...target, sandboxLeaseAcquisition: { ...target.sandboxLeaseAcquisition, outcome: "resumed" as const } }; await expect(assertManagedAiProjectAuth({}, "openai", retained, options)).resolves.toBeUndefined(); + expect(runnerTask).not.toHaveBeenCalled(); + await expect(assertManagedAiProjectAuth({}, "openai", retained)).resolves.toBeUndefined(); expect(runnerTask).toHaveBeenCalledExactlyOnceWith({ kind: "describe" }); }); it("keeps personal defaults separate and does not replace the first default", async () => { diff --git a/server/src/__tests__/environment-task-authority.integration.test.ts b/server/src/__tests__/environment-task-authority.integration.test.ts index ce9ad8b6be6..0f1ad4c9ff7 100644 --- a/server/src/__tests__/environment-task-authority.integration.test.ts +++ b/server/src/__tests__/environment-task-authority.integration.test.ts @@ -67,7 +67,7 @@ databaseSuite("environment task lease authority", () => { try { await expect(environmentRuntimeService(db, { pluginWorkerManager: workers }).task({ companyId, leaseId, operation: { kind: "describe" }, executionContext: context })).resolves.toMatchObject({ kind: "description", taskId: "replacement-task" }); - expect(calls).toEqual([...(mode === "writeback" ? ["status"] : []), "status", "status", "complete", "environmentAcquireLease", "describe"]); + expect(calls).toEqual([...(mode === "writeback" ? ["status"] : []), "status", "complete", "environmentAcquireLease", "describe"]); expect(context.pending).toBe(false); const [lease] = await db.select().from(environmentLeases).where(eq(environmentLeases.id, leaseId)); expect(lease.providerLeaseId).toBe("replacement-task"); @@ -189,7 +189,7 @@ databaseSuite("environment task lease authority", () => { taskId: replaced ? "replacement-task" : "tombstoned-task" }); const [lease] = await db.select().from(environmentLeases).where(eq(environmentLeases.id, leaseId)); if (replaced) { - expect(calls).toEqual(["status:tombstoned-task", "status:tombstoned-task", "complete:tombstoned-task", "environmentAcquireLease", "describe:replacement-task"]); + expect(calls).toEqual(["status:tombstoned-task", "complete:tombstoned-task", "environmentAcquireLease", "describe:replacement-task"]); expect(lease.providerLeaseId).toBe("replacement-task"); expect(lease.metadata).toMatchObject({ runnerTaskExecutionStopped: false, runnerTaskDispatched: false, runnerTaskSessionRetained: false, runnerTaskSourceRecovery: { schema: "fixture" } }); diff --git a/server/src/services/ai-connection-runtime.ts b/server/src/services/ai-connection-runtime.ts index 130d58638c2..ade2d0e0f53 100644 --- a/server/src/services/ai-connection-runtime.ts +++ b/server/src/services/ai-connection-runtime.ts @@ -103,7 +103,7 @@ export async function assertManagedAiProjectAuth( config: Record, provider: AiConnectionBinding["provider"], target?: AdapterExecutionTarget | null, - options: { deferNewRunnerTaskCapabilityCheck?: boolean } = {}, + options: { deferRunnerTaskCapabilityCheck?: boolean } = {}, ) { const extraArgs = [ ...(Array.isArray(config.extraArgs) ? config.extraArgs : []), @@ -132,7 +132,7 @@ export async function assertManagedAiProjectAuth( const pattern = "apiKeyHelper|ANTHROPIC_API_KEY|ANTHROPIC_AUTH_TOKEN|CLAUDE_CODE_OAUTH_TOKEN|OPENAI_API_KEY|model_provider[[:space:]]*=|env_key[[:space:]]*=|experimental_bearer_token|cli_auth_credentials_store"; if (target?.kind === "remote" && target.transport === "sandbox" && target.runnerTask && files.length) { - if (options.deferNewRunnerTaskCapabilityCheck && target.sandboxLeaseAcquisition?.outcome === "created") return; + if (options.deferRunnerTaskCapabilityCheck) return; const description = environmentTaskResultSchema.parse(await target.runnerTask({ kind: "describe" })); if (description.kind !== "description" || description.projectAuthCheckVersion !== 1) { throw unprocessable("The plugin environment cannot validate managed project authentication", { code: "ai_connection_incompatible" }); diff --git a/server/src/services/environment-task-runtime.ts b/server/src/services/environment-task-runtime.ts index 7c8a787a953..12f7e54913d 100644 --- a/server/src/services/environment-task-runtime.ts +++ b/server/src/services/environment-task-runtime.ts @@ -279,6 +279,14 @@ export async function executeEnvironmentTask(db: Db, workers: PluginWorkerManage sql`${environmentLeases.metadata}->'runnerTaskIdentity' = ${JSON.stringify(savedIdentity)}::jsonb`)).returning({id:environmentLeases.id}); if (marked.length !== 1) throw new Error("Plugin warm task attachment lost authority"); } + async function recordTaskStop(result: Extract, {kind:"status" | "copilot_metadata"}>) { + await db.update(environmentLeases).set({ + metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || ${JSON.stringify({ runnerTaskExecutionStopped: true, + runnerTaskSessionRetained: result.kind === "status" && result.sessionStateAvailable === true, + ...(result.kind === "status" && result.preparationStopped === true ? { runnerTaskPreparationStopped: true } : {}) })}::jsonb`, + }).where(and(eq(environmentLeases.id, lease.id), eq(environmentLeases.companyId, input.companyId), + eq(environmentLeases.providerLeaseId, taskId))); + } try { const admissionRejection = metadata.runnerTaskAdmissionRejection as { taskId?: unknown } | undefined; const stoppedWithoutSession = !!run && metadata.runnerTaskExecutionStopped === true && @@ -291,8 +299,8 @@ export async function executeEnvironmentTask(db: Db, workers: PluginWorkerManage })); if (state.kind === "status" && state.executionStopped === true && (metadata.runnerTaskReused === true || state.sessionStateAvailable !== true)) { - await executeEnvironmentTask(db, workers, {...input, operation:{kind:"status"}}); - await executeEnvironmentTask(db, workers, {...input, operation:{kind:"complete"}}); + await recordTaskStop(state); + await executeEnvironmentTask(db, workers, {...input, expectedTaskId:taskId, operation:{kind:"complete"}}); // Preserve the predecessor's stop proof for checkpoint authorization. if (typeof metadata.runnerTaskPreviousLeaseId === "string") await db.update(environmentLeases).set({ metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || ${JSON.stringify({runnerTaskExecutionStopped:true, @@ -466,12 +474,7 @@ export async function executeEnvironmentTask(db: Db, workers: PluginWorkerManage } } if ((parsed.kind === "status" || parsed.kind === "copilot_metadata") && parsed.executionStopped === true) { - await db.update(environmentLeases).set({ - metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || ${JSON.stringify({ runnerTaskExecutionStopped: true, - runnerTaskSessionRetained: parsed.kind === "status" && parsed.sessionStateAvailable === true, - ...(parsed.kind === "status" && parsed.preparationStopped === true ? { runnerTaskPreparationStopped: true } : {}) })}::jsonb`, - }).where(and(eq(environmentLeases.id, lease.id), eq(environmentLeases.companyId, input.companyId), - eq(environmentLeases.providerLeaseId, taskId))); + await recordTaskStop(parsed); } if (operation.kind === "complete") { if (freshWorkspace && metadata.runnerTaskExecutionStopped === true) diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index eb94fc0cb15..6bfc03289d0 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -2947,7 +2947,7 @@ export function heartbeatService( if (managedAiRuntime && aiBinding) { try { await assertManagedAiProjectAuth({ ...resolvedConfig, cwd: executionWorkspace.cwd }, aiBinding.provider, executionTarget, { - deferNewRunnerTaskCapabilityCheck: agent.adapterType === "paperclip_runner" && !runOptions.nativeRestartRecovery, + deferRunnerTaskCapabilityCheck: agent.adapterType === "paperclip_runner" && !runOptions.nativeRestartRecovery, }); } catch (error) { throw classifyManagedProjectAuthFailure(error, agent.id); 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 b0450623921..fe2a000a75f 100644 --- a/server/src/services/native-runtime/native-session-executor.test.ts +++ b/server/src/services/native-runtime/native-session-executor.test.ts @@ -17722,7 +17722,8 @@ describe("runnerd provider runtime wiring", () => { } }); - it.each(["codex", "opencode", "acpx"] as const)("uses a task-only %s driver without exposing Core paths or transferring the agent home", async (kind) => { + it.each([["codex", undefined], ["opencode", undefined], ["acpx", undefined], ["codex", "task"], ["codex", "previous-task"]] as const)( + "uses a task-only %s driver with a prior description for %s without exposing Core paths or transferring the agent home", async (kind, describedTask) => { const taskExecution = { ...execution, provider: kind === "codex" ? execution.provider : kind === "opencode" ? { kind, model: "openai/test", permissionMode: "deny" } : { kind, agent: "claude", model: "claude-sonnet-4-6", permissionMode: "deny-all", claudeEffort: "high" } } as NativeExecutionInput; @@ -17735,12 +17736,16 @@ describe("runnerd provider runtime wiring", () => { const readWorkspaceFile = vi.fn(async () => Buffer.from("proof")); const release = vi.fn(async () => {}), resume = vi.fn(async () => "resumed" as const); const owner = vi.spyOn(pluginTaskRuntime, "createPluginRunnerTaskOwner").mockReturnValue({ launcher, registration, driver, readWorkspaceFile, release, resume, isRetired: () => false, retireSupersededTask: async () => false, retiredBeforeClaim: () => false, credentialsUpdated: vi.fn(), rebind:vi.fn() }); + const description = { kind: "description" as const, launchPreparationVersion: 1 as const, projectAuthCheckVersion: 1 as const, + releaseResumeVersion: 1 as const, ...(kind === "codex" ? { workspaceFileReadVersion: 1 as const } : {}), + taskId: "task", runner: { version: "fixture", digest: `sha256:${"a".repeat(64)}` } }; const task = vi.fn(async (operation: unknown) => { environmentTaskOperationSchema.parse(operation); - return { kind: "description", launchPreparationVersion: 1, projectAuthCheckVersion: 1, releaseResumeVersion: 1, ...(kind === "codex" ? { workspaceFileReadVersion: 1 } : {}), taskId: "task", runner: { version: "fixture", digest: `sha256:${"a".repeat(64)}` } }; + return description; }); const target = { kind: "remote", transport: "sandbox", providerKey: "plugin", leaseId: "plugin-task-lease", - remoteCwd: "/unused", runnerTask: task, runnerTaskProjectIds: ["10000000-0000-4000-8000-000000000001"] } as const; + remoteCwd: "/unused", runnerTask: task, runnerTaskProjectIds: ["10000000-0000-4000-8000-000000000001"], + sandboxLeaseAcquisition: { outcome: "created", providerLeaseId: "task" } } as const; const updates: Array<{ table: unknown; values: Record }> = []; state.createBackend.mockClear(); state.createTransport.mockClear(); try { @@ -17748,6 +17753,7 @@ describe("runnerd provider runtime wiring", () => { .rejects.toThrow("runner_ingress_unavailable"); await createRunnerdBackend({ db: leaseDb(execution, {}, {}, updates), execution: taskExecution, runnerInstanceId: "runner", runnerExecutionTarget: target, runnerIngressAuthorized: true, + ...(describedTask ? { pluginTaskDescription: { ...description, taskId: describedTask } } : {}), runnerEnvironment: { HOME: "/controller/private", PATH: "/controller/bin", OPENAI_API_KEY: "fixture-secret", PAPERCLIP_NATIVE_MCP_NAME: "assigned", PAPERCLIP_NATIVE_MCP_URL: "http://localhost/mcp/gateways/assigned-test", PAPERCLIP_NATIVE_MCP_TOKEN: "fixture-mcp-secret" } }); @@ -17770,7 +17776,7 @@ describe("runnerd provider runtime wiring", () => { for (const key of ["environment", "codexCommand", "codexArgs", "runnerRuntimePaths", "readRunnerState", "archiveExternalRunnerState"]) { expect(options).not.toHaveProperty(key); } - expect(task.mock.calls).toEqual([[{ kind: "describe" }]]); + expect(task.mock.calls).toEqual(describedTask === "task" ? [] : [[{ kind: "describe" }]]); expect(release).not.toHaveBeenCalled(); expect(resume).not.toHaveBeenCalled(); const binding = state.toolAuthorityDefinitions.mock.calls.at(-1)![0] as unknown as { diff --git a/server/src/services/native-runtime/native-session-executor.ts b/server/src/services/native-runtime/native-session-executor.ts index 89b14a11a2c..d474b8f4c40 100644 --- a/server/src/services/native-runtime/native-session-executor.ts +++ b/server/src/services/native-runtime/native-session-executor.ts @@ -10686,6 +10686,7 @@ async function executePaperclipNativeSessionAttempt( }); let retainedTransition: VerifiedWarmTransitionBinding | undefined; let remoteWarmRecovery: RemoteWarmSessionRecovery | undefined; + let pluginTaskDescription: Extract,{kind:"description"}> | undefined; const cachedReleasedOwner = warmNativeSessions.get(nativeSessionScopeKey(input.execution)); if (cachedReleasedOwner?.released) { const proof = releasedWarmRunner(readWarmNativeEnvelope(input.execution)); @@ -10728,7 +10729,9 @@ async function executePaperclipNativeSessionAttempt( if (taskTarget.transport === "sandbox" && taskTarget.runnerTask && taskTarget.leaseId && durableIdentityMatchesSession(identity,input.execution) && identity.runId !== input.execution.binding.runId) { // Inspection can replace a stopped retained task before deciding between adoption and checkpoint recovery. - await taskTarget.runnerTask({kind:"describe"}); + const description = environmentTaskResultSchema.parse(await taskTarget.runnerTask({kind:"describe"})); + if (description.kind !== "description") throw new Error("Plugin Runner description unavailable"); + pluginTaskDescription = description; if (taskTarget.sandboxLeaseAcquisition?.outcome !== "resumed") { await recoverPluginRunnerCheckpoint(input, root); } @@ -11709,7 +11712,6 @@ async function executePaperclipNativeSessionAttempt( ); let existingWarmSession: NativeSession | undefined; let reusingReleasedRunner = false; - let pluginTaskDescription: Extract,{kind:"description"}> | undefined; let managedCredentialSession: NativeSession | undefined; let githubAccess: NativeGitHubAccess | undefined; let releaseGitHubRun: (() => void) | undefined; @@ -11747,7 +11749,9 @@ async function executePaperclipNativeSessionAttempt( if (entry) entry.preparingRunId = undefined; const taskTarget=input.runnerExecutionTarget; if ((entry?.pluginTaskId || remoteWarmRecovery?.releasedSource) && taskTarget?.kind === "remote" && taskTarget.transport === "sandbox" && taskTarget.runnerTask) { - const description=environmentTaskResultSchema.parse(await (executionContext + const priorDescription = pluginTaskDescription?.taskId === taskTarget.sandboxLeaseAcquisition?.providerLeaseId + ? pluginTaskDescription : undefined; + const description=priorDescription ?? environmentTaskResultSchema.parse(await (executionContext ? executionContext.wait(() => taskTarget.runnerTask!({kind:"describe"}, { executionContext })) : taskTarget.runnerTask({kind:"describe"}))); if(description.kind !== "description") throw new Error("Plugin Runner description unavailable"); @@ -12058,7 +12062,7 @@ async function executePaperclipNativeSessionAttempt( const description = environmentTaskResultSchema.parse(await (executionContext ? executionContext.wait(inspect) : inspect())); if (description.kind !== "description" || description.taskId === retainedTaskId) return false; existingWarmSession = undefined; - pluginTaskDescription = undefined; + pluginTaskDescription = description; await finishStartupContextForFallback(); persistedWarmSession = loadWarmNativeCheckpoint(input.execution, warmConfigDigest, target.kind); runnerdBackend = await createExecutionBackend(); @@ -16013,7 +16017,9 @@ async function createRunnerdBackendWithinSessionClaim( if(stopped.kind!=="status" || stopped.executionStopped!==true) throw new Error("Plugin Runner termination is unconfirmed; reconcile the saved task"); await pluginTask({kind:"complete"}); } - const description = input.pluginTaskDescription ?? environmentTaskResultSchema.parse(await pluginTask({ kind: "describe" })); + const priorDescription = input.pluginTaskDescription?.taskId === remoteTarget.sandboxLeaseAcquisition?.providerLeaseId + ? input.pluginTaskDescription : undefined; + const description = priorDescription ?? environmentTaskResultSchema.parse(await pluginTask({ kind: "describe" })); if (description.kind !== "description" || description.launchPreparationVersion !== 1 || description.projectAuthCheckVersion !== 1) throw new Error("Plugin Runner launch preparation v1 unavailable"); const saved = durableBinding ? durableIdentity as RunnerdDurableIdentity : undefined; const prior = saved ? {