From 33cc567f75ae8a3fecf4442d20772e62944f1552 Mon Sep 17 00:00:00 2001 From: Nicky Leach Date: Tue, 20 Oct 2026 09:14:23 -0700 Subject: [PATCH] Add lease-scoped Runner task hooks to environment plugins (#15406) ## Thinking Path > - Paperclip coordinates AI agents and their runs. > - Environment plugins own remote capacity. > - A plugin needs a typed API to start and stop its installed Runner. > - Core must derive authority from the saved lease. > - A lost reply or controller restart must not lose cleanup authority. > - This PR adds the worker contract and the host task service. > - Native execution and Runner changes are separate PRs. ## Linked Issues or Issue Description **Subsystem affected** Plugin SDK and environment lease services. **Current behavior** Environment plugins have no typed Runner task hook. Core cannot dispatch these tasks or recover their credential handoff through a saved lease. **Proposed behavior** Add validated task operations, worker negotiation, company and project checks, retained-session authorization, and durable cleanup. Keep credentials out of lease metadata. **Additional context** This PR is the first part of a smaller series. Runner installed-home support is in #16977. Native execution and setup are in #16978. Both are stacked drafts. This PR does not enable native plugin execution, change default environments, add a route, or add a database migration. ## What Changed - Add `supportsTasks` and `onEnvironmentTask` with seven typed operations. - Validate worker requests and responses. Reject shell commands and file transfer operations. - Pin dispatch and cleanup to the saved company, lease, plugin, and driver. - Authorize projects and retained sessions. Validate the current run separately from the retained PRP identity. - Save dispatch before the provider call. Fence admission during cleanup. Require observed process termination. - Save account writeback authority on the lease. Recover Codex refreshes through the existing account and secret services. - Document the task API. Keep the Runner, heartbeat, setup, workspace, and provider changes in later PRs. ## Verification - Full `pnpm -r typecheck` and `pnpm build` passed on the separated API scope. - 201 focused tests passed: SDK worker dispatch, task validation, PostgreSQL task lineage, credential recovery, and existing environment cleanup. - The lineage test covers a third run, repeated same-run replacements, a stale submitting run, and exact termination receipts. - Initial commands hit sandbox socket/IPC restrictions. The permitted runs executed PostgreSQL and passed. A typecheck attempt lacked Cargo in PATH; the run with the installed toolchain passed. - The earlier combined implementation had a full local run with 22,286 passing tests, 120 skips, and eight failed assertions. Task fixture failures were fixed before this split. Remaining host Git and temporary-path assertions are recorded in the previous PR revision. This split does not claim a fresh full local suite pass. - The standard Greptile review scored this API head 5/5 with no unresolved threads. CI is green. The single Apex review also scored this head 5/5 with no actionable findings. - Fresh API CI lost a self-hosted runner during server shard 9. GitHub reported lost communication, with no failed test step. The affected jobs passed on retry without code changes. - CI on the prior combined head failed an unchanged workspace dev-server readiness case. No unrelated test fix is included here. Fresh CI is required for this smaller head. - This scope does not run a provider or a Cloud test. The user owns deployment and the intermediate Cloud test. ## Risks - The API is opt-in. Native plugin execution requires the later integrations. - Providers must make task submission idempotent and prove descendant termination. - A failed status read or credential export keeps the lease available for cleanup. - Codex refresh writeback requires the same active grant, secret reference, and session epoch. Revocation or account replacement discards the old refresh. ## Model Used OpenAI GPT-6 through Codex. Used source inspection, code execution, tests, and GitHub tools. The exact served model ID and context-window size are not exposed in this session. ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with `Fixes: #` / `Closes #` / `Refs #` OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip` URLs) - [x] My branch name describes the change (e.g. `docs/...`, `fix/...`) and contains no internal Paperclip ticket id or instance-derived details - [x] I have run local checks and recorded failures and their results above - [x] I have added or updated tests where applicable - [x] I have updated relevant documentation to reflect my changes - [x] I have considered and documented any risks above - [x] All Paperclip CI gates are green - [x] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge Co-authored-by: Paperclip --- doc/plugins/PLUGIN_AUTHORING_GUIDE.md | 7 + doc/plugins/PLUGIN_ENVIRONMENT_TASKS.md | 95 ++++++++ packages/plugins/sdk/src/define-plugin.ts | 4 + packages/plugins/sdk/src/environment-tasks.ts | 114 ++++++++++ packages/plugins/sdk/src/index.ts | 2 + packages/plugins/sdk/src/protocol.ts | 3 + packages/plugins/sdk/src/worker-rpc-host.ts | 9 + packages/shared/src/types/plugin.ts | 2 + packages/shared/src/validators/plugin.ts | 1 + ...ronment-task-authority.integration.test.ts | 108 ++++++++++ .../environment-task-runtime.test.ts | 199 +++++++++++++++++ .../plugin-environment-driver-seam.test.ts | 49 +++++ .../plugin-task-credential-writeback.test.ts | 117 ++++++++++ server/src/services/ai-connection-runtime.ts | 5 + server/src/services/environment-runtime.ts | 79 +++++-- .../src/services/environment-task-runtime.ts | 202 ++++++++++++++++++ .../plugin-task-credential-writeback.ts | 143 +++++++++++++ 17 files changed, 1124 insertions(+), 15 deletions(-) create mode 100644 doc/plugins/PLUGIN_ENVIRONMENT_TASKS.md create mode 100644 packages/plugins/sdk/src/environment-tasks.ts create mode 100644 server/src/__tests__/environment-task-authority.integration.test.ts create mode 100644 server/src/__tests__/environment-task-runtime.test.ts create mode 100644 server/src/__tests__/plugin-task-credential-writeback.test.ts create mode 100644 server/src/services/environment-task-runtime.ts create mode 100644 server/src/services/plugin-task-credential-writeback.ts diff --git a/doc/plugins/PLUGIN_AUTHORING_GUIDE.md b/doc/plugins/PLUGIN_AUTHORING_GUIDE.md index f4b4c06a8d3..e6913a203df 100644 --- a/doc/plugins/PLUGIN_AUTHORING_GUIDE.md +++ b/doc/plugins/PLUGIN_AUTHORING_GUIDE.md @@ -791,3 +791,10 @@ The host resets plugin state on account/company changes and keeps its built-in menu when no unique contribution exists, discovery fails, the module is missing, or rendering throws. The slot is a React-only contract; do not use a custom element export. This replaces only the menu, not company policy or authorization. + +## Runner task worker hooks + +A plugin can declare `supportsTasks` on an environment driver and implement +`onEnvironmentTask`. See the [task API](PLUGIN_ENVIRONMENT_TASKS.md) for the +operations, validation, and lease rules. This hook does not add a shell or file +transfer API. Native agent execution is a separate integration. diff --git a/doc/plugins/PLUGIN_ENVIRONMENT_TASKS.md b/doc/plugins/PLUGIN_ENVIRONMENT_TASKS.md new file mode 100644 index 00000000000..efbd7b3aa4a --- /dev/null +++ b/doc/plugins/PLUGIN_ENVIRONMENT_TASKS.md @@ -0,0 +1,95 @@ +# Environment task API + +An environment driver can provide a Runner task through a worker hook. +The task API has no shell command or file transfer operation. + +This API is a server and worker contract. It does not select an environment +for a native agent or connect a Runner to a run. Those integrations are separate. + +## Declare the hook + +Set `supportsTasks: true` on the environment driver in the plugin manifest. +Implement `onEnvironmentTask` with the SDK types: + +```ts +import type { + PluginEnvironmentTaskParams, + PluginEnvironmentTaskResult, +} from "@paperclipai/plugin-sdk"; + +type TaskHook = ( + params: PluginEnvironmentTaskParams, +) => Promise; +``` + +Core requires the `environment.drivers.register` capability, a ready plugin, +the manifest declaration, and the worker's `environmentTask` method. +The SDK validates each operation and its result before it returns the result. +See [the schemas](../../packages/plugins/sdk/src/environment-tasks.ts) for all fields. + +## Operations + +| Operation | Result | Required behavior | +| --- | --- | --- | +| `describe` | `description` | Return the installed Runner version and SHA-256 digest. | +| `submit` | `accepted` | Start the same opaque task at most once. Use the structured launch inputs. | +| `connection` | `connection` | Return an authenticated WSS endpoint and its generation. | +| `status` | `status` | Report the phase and observed process termination. | +| `credentials` | `credentials` | Return the selected account after process termination. | +| `stop` | `accepted` | Stop the task and its descendants. Fence a delayed submission. | +| `complete` | `accepted` | Release task resources after Core saves required results. | + +Each result must contain the task ID from the request. Core rejects a result +with a different ID or result kind. A WSS URL must have no user information, +query, or fragment. Return authentication in the secret headers. + +`submit` contains the PRP version range, Runner identity, bootstrap ticket, +authorized project IDs, and optional launch limits. It can also contain a saved +task ID for session reuse and a typed Codex or Claude credential. Core does not +supply an executable command or a guest path. + +## Lease authority + +Core calls `environmentRuntimeService.task` with a company ID, a saved lease ID, +and an operation. It derives the plugin and driver from the lease. A change to +the environment configuration cannot redirect cleanup to a new provider. + +Before submission, Core checks the active lease, run, issue, and every project +in the company. It checks the submitting run ID separately from a retained PRP +identity. Session reuse requires the same agent, environment, workspace, and +provider. It also requires a retained source task and proof of termination. + +An internal setup probe can use a private, time-limited lease without a run row. +Core checks the saved probe identity, rejects project mounts, and limits the +Runner lifetime to the remaining probe time. + +## Failure and cleanup + +A failed submission reply does not prove that the task failed to start. +Core saves the dispatch marker before it calls the worker. Retry against the +same task ID. Do not allocate another task to resolve an uncertain reply. + +An accepted stop does not prove termination. Set `executionStopped: true` +only after the Runner and all descendants have stopped. A terminal phase alone +does not permit Core to release execution authority. + +Core fences new submissions when cleanup starts. It stops the saved task, +checks termination, saves credential results, and completes the task before it +releases or destroys the environment lease. An uncertain result leaves a saved +lease for the cleanup sweep. Cleanup can use that lease after the environment +or run row is deleted. + +## Credential writeback + +Account selection and credential revision have separate identities. Keep the +selected account fixed for the task. A credential result of `null` means that +cancellation prevented task admission. It must not mean that an export failed. + +Core saves the account grant, secret reference, and session epoch on the lease. +It saves no credential bytes or temporary path in lease metadata. After process +termination, it can recover a Codex refresh even if Core restarted. It uses the +existing account comparison and secret service. A revoked grant, replaced +account, or changed session epoch discards the old refresh. + +Keep file synchronization and backup outside this API. The provider must retain +recoverable credential state until task completion succeeds. diff --git a/packages/plugins/sdk/src/define-plugin.ts b/packages/plugins/sdk/src/define-plugin.ts index 3d7d81a31b5..13096cbcf8d 100644 --- a/packages/plugins/sdk/src/define-plugin.ts +++ b/packages/plugins/sdk/src/define-plugin.ts @@ -1,3 +1,4 @@ +import type { PluginEnvironmentTaskParams, PluginEnvironmentTaskResult } from "./environment-tasks.js"; import type { AgentLifecycleRequest, AgentLifecycleResult } from "@paperclipai/shared"; import type { AiConnectionRouterRequest, AiConnectionRouterResult } from "@paperclipai/shared"; /** @@ -401,6 +402,9 @@ export interface PluginDefinition { params: PluginEnvironmentValidateConfigParams, ): Promise; + /** Provide Runner execution: prepare resources, start Runner, connect via PRP, and manage task lifecycle. */ + onEnvironmentTask?(params: PluginEnvironmentTaskParams): Promise; + /** Called to test reachability or readiness of a plugin-hosted environment. */ onEnvironmentProbe?( params: PluginEnvironmentProbeParams, diff --git a/packages/plugins/sdk/src/environment-tasks.ts b/packages/plugins/sdk/src/environment-tasks.ts new file mode 100644 index 00000000000..ae2cd9e7eed --- /dev/null +++ b/packages/plugins/sdk/src/environment-tasks.ts @@ -0,0 +1,114 @@ +import { z } from "zod"; +import type { PluginEnvironmentDriverBaseParams, PluginEnvironmentLease } from "./protocol.js"; + +// Provider IDs are opaque; providers validate their own addressing constraints. +const providerId = z.string().min(1); +// PRP identity.schema.json stableId. +const runnerIdentity = z.string().regex(/^[A-Za-z0-9][A-Za-z0-9._:-]{0,159}$/); +const identifier = z.string().regex(/^[a-zA-Z0-9][a-zA-Z0-9_-]{0,95}$/); +const secureUrl = z.string().url().refine(value => { + const url = new URL(value); + return url.protocol === "wss:" && !url.username && !url.password && !url.search && !url.hash; +}, "Expected a credential-free WSS URL"); + +const credentialValue = z.string().min(1).max(65_536).refine(value => !value.includes("\0")); +/** Only provider authentication is transferred. Paths and executable settings remain provider-owned. */ +export const environmentTaskCredentialSchema = z.discriminatedUnion("type", [ + z.object({ type: z.literal("codex-auth-json"), authJson: credentialValue }).strict(), + z.object({ type: z.literal("claude-env"), environment: z.object({ + CLAUDE_CODE_OAUTH_TOKEN: credentialValue.optional(), ANTHROPIC_API_KEY: credentialValue.optional(), + }).strict().refine(env => Object.values(env).filter(Boolean).length === 1, "Select one Claude credential") }).strict(), +]); +export type PluginEnvironmentTaskCredential = z.infer; + +/** One plugin-provided Runner execution attempt with a PRP connection. Secrets are transient RPC input. */ +export const environmentTaskOperationSchema = z.discriminatedUnion("kind", [ + z.object({ kind: z.literal("describe") }).strict(), + z.object({ + kind: z.literal("submit"), + projectIds: z.array(z.string().uuid()).max(64).refine(ids => new Set(ids).size === ids.length, "Duplicate project IDs"), + runner: z.object({ + /** Inclusive PRP version range supported by the submitting client. */ + protocolMin: z.number().int().min(1).max(Number.MAX_SAFE_INTEGER), + protocolMax: z.number().int().min(1).max(Number.MAX_SAFE_INTEGER), + harness: identifier, + runnerId: runnerIdentity, leaseId: runnerIdentity, runId: runnerIdentity, + sessionId: runnerIdentity, turnId: runnerIdentity, itemId: runnerIdentity, + }).strict().refine(value => value.protocolMin <= value.protocolMax, "Invalid PRP version range"), + bootstrapTicket: z.string().min(1).max(65_536), + /** Opaque prior provider task; Core authorizes the saved session before supplying it. */ + resumeFrom: z.string().min(1).optional(), + /** Require the externally installed account to match this Core-selected credential revision. */ + credentialIdentity: z.string().min(1).max(1024).optional(), + /** Stable account selection; distinct from its current credential revision. */ + accountIdentity: z.string().min(1).max(1024).optional(), + credential: environmentTaskCredentialSchema.optional(), + /** Core supplies these protected launch inputs for native agent runs. */ + launch: z.object({ + runnerVersion: z.string().min(1), runnerDigest: z.string().regex(/^sha256:[a-f0-9]{64}$/), + maxOutboxBytes: z.number().int().positive(), p0ReserveBytes: z.number().int().positive(), + maxRuntimeMs: z.number().int().nonnegative().optional(), maxLifetimeMs: z.number().int().nonnegative().optional(), + reconnectGraceMs: z.number().int().nonnegative().optional(), + lifecyclePolicy: z.object({ mode: z.literal("per_turn"), idleTimeoutMs: z.null() }).strict(), + }).strict().optional(), + }).strict(), + z.object({ kind: z.literal("credentials"), accountIdentity: z.string().min(1).max(1024) }).strict(), + z.object({ kind: z.literal("status") }).strict(), + z.object({ kind: z.literal("connection") }).strict(), + z.object({ kind: z.literal("complete") }).strict(), + z.object({ kind: z.literal("stop") }).strict(), +]); + +export const environmentTaskResultSchema = z.discriminatedUnion("kind", [ + z.object({ kind: z.literal("accepted"), taskId: providerId }).strict(), + z.object({ kind: z.literal("credentials"), taskId: providerId, /** Null proves that cancellation fenced admission before the task could start. It must not mean an unavailable export. */ + credential: environmentTaskCredentialSchema.nullable() }).strict(), + z.object({ kind: z.literal("description"), taskId: providerId, + runner: z.object({ version: z.string().min(1), digest: z.string().regex(/^sha256:[a-f0-9]{64}$/) }).strict(), + }).strict(), + z.object({ + kind: z.literal("status"), taskId: providerId, + phase: z.enum(["preparing", "running", "completed", "failed", "cancelled", "interrupted"]), + exitCode: z.number().int().optional(), + /** Provider observed the task and all descendants stopped; terminal phase alone is insufficient. */ + executionStopped: z.boolean().optional(), + /** The provider retains this task's suspended session for a later launch. */ + sessionStateAvailable: z.boolean().optional(), + }).strict(), + z.object({ + kind: z.literal("connection"), taskId: providerId, + endpoint: z.object({ + kind: z.literal("authenticated_websocket"), websocketUrl: secureUrl, + generation: providerId, + secretHeaders: z.array(z.object({ + name: z.string().regex(/^[!#$%&'*+.^_`|~0-9A-Za-z-]+$/), + value: z.string().min(1).max(65_536).regex(/^[^\r\n]+$/), + }).strict()).max(16), + }).strict(), + }).strict(), +]); + +export type PluginEnvironmentTaskOperation = z.infer; +export type PluginEnvironmentTaskResult = z.infer; + +export interface PluginEnvironmentTaskParams extends Omit { + /** Null during cleanup after the environment is deleted. */ + environmentId: string | null; + lease: PluginEnvironmentLease; + /** Provider-issued task identifier persisted in the lease, stable across ambiguous submission retries. */ + taskId: string; + /** Null during cleanup after the associated run is deleted. */ + runId: string | null; + agentId: string | null; + /** Company-validated projects to prepare on submit; empty during cleanup. */ + projectIds: string[]; + operation: PluginEnvironmentTaskOperation; +} + +/** An accepted stop is not proof of process termination. */ +export function parseEnvironmentTaskResult(operation: PluginEnvironmentTaskOperation, taskId: string, value: unknown): PluginEnvironmentTaskResult { + const result = environmentTaskResultSchema.parse(value); + const expected = operation.kind === "describe" ? "description" : ["status", "connection", "credentials"].includes(operation.kind) ? operation.kind : "accepted"; + if (result.taskId !== taskId || result.kind !== expected) throw new Error("Invalid plugin-provided Runner execution response"); + return result; +} diff --git a/packages/plugins/sdk/src/index.ts b/packages/plugins/sdk/src/index.ts index 8698d6ba4bf..9b74f554d8c 100644 --- a/packages/plugins/sdk/src/index.ts +++ b/packages/plugins/sdk/src/index.ts @@ -463,4 +463,6 @@ export type { PluginEnvironmentSyncErrorDiagnostic } from "./environment-sync-er export type { AiConnectionPool, AiConnectionPoolConfig, AiConnectionPoolMember, AiConnectionRouterRequest, AiConnectionRouterResult, AiConnectionRouterSelection } from "@paperclipai/shared"; +export { environmentTaskOperationSchema, environmentTaskResultSchema, environmentTaskCredentialSchema, parseEnvironmentTaskResult } from "./environment-tasks.js"; +export type { PluginEnvironmentTaskOperation, PluginEnvironmentTaskParams, PluginEnvironmentTaskResult, PluginEnvironmentTaskCredential } from "./environment-tasks.js"; export type { AgentLifecycleRequest, AgentLifecycleResult } from "@paperclipai/shared"; diff --git a/packages/plugins/sdk/src/protocol.ts b/packages/plugins/sdk/src/protocol.ts index 83e8d88b097..b6cb01a679c 100644 --- a/packages/plugins/sdk/src/protocol.ts +++ b/packages/plugins/sdk/src/protocol.ts @@ -1,3 +1,4 @@ +import type { PluginEnvironmentTaskParams, PluginEnvironmentTaskResult } from "./environment-tasks.js"; import type { PluginAgentRuntimeState } from "./types.js"; import type { AgentLifecycleRequest, AgentLifecycleResult } from "@paperclipai/shared"; import type { AiConnectionRouterRequest, AiConnectionRouterResult } from "@paperclipai/shared"; @@ -1393,6 +1394,7 @@ export interface HostToWorkerMethods { params: PluginEnvironmentValidateConfigParams, result: PluginEnvironmentValidationResult, ]; + environmentTask: [params: PluginEnvironmentTaskParams, result: PluginEnvironmentTaskResult]; environmentProbe: [ params: PluginEnvironmentProbeParams, result: PluginEnvironmentProbeResult, @@ -1521,6 +1523,7 @@ export const HOST_TO_WORKER_OPTIONAL_METHODS: readonly HostToWorkerMethodName[] "resolveExternalObject", "refreshExternalObjects", "environmentValidateConfig", + "environmentTask", "environmentProbe", "environmentAcquireLease", "environmentResumeLease", diff --git a/packages/plugins/sdk/src/worker-rpc-host.ts b/packages/plugins/sdk/src/worker-rpc-host.ts index 3de92ac5216..8b98a83e64a 100644 --- a/packages/plugins/sdk/src/worker-rpc-host.ts +++ b/packages/plugins/sdk/src/worker-rpc-host.ts @@ -1,3 +1,4 @@ +import { environmentTaskOperationSchema, parseEnvironmentTaskResult, type PluginEnvironmentTaskParams } from "./environment-tasks.js"; import type { AgentLifecycleRequest, AgentLifecycleResult } from "@paperclipai/shared"; import type { AiConnectionRouterRequest, AiConnectionRouterResult } from "@paperclipai/shared"; import { environmentCreationCleanupErrorData } from "./environment-creation-cleanup.js"; @@ -1682,6 +1683,13 @@ export function startWorkerRpcHost(options: WorkerRpcHostOptions): WorkerRpcHost case "environmentValidateConfig": return handleEnvironmentValidateConfig(params as PluginEnvironmentValidateConfigParams); + case "environmentTask": { + if (!plugin.definition.onEnvironmentTask) throw methodNotImplemented("environmentTask"); + const input = params as PluginEnvironmentTaskParams; + const operation = environmentTaskOperationSchema.parse(input.operation); + return parseEnvironmentTaskResult(operation, input.taskId, await plugin.definition.onEnvironmentTask({ ...input, operation })); + } + case "environmentProbe": return handleEnvironmentProbe(params as PluginEnvironmentProbeParams); @@ -1801,6 +1809,7 @@ export function startWorkerRpcHost(options: WorkerRpcHostOptions): WorkerRpcHost if (plugin.definition.onResolveExternalObject) supportedMethods.push("resolveExternalObject"); if (plugin.definition.onRefreshExternalObjects) supportedMethods.push("refreshExternalObjects"); if (plugin.definition.onEnvironmentValidateConfig) supportedMethods.push("environmentValidateConfig"); + if (plugin.definition.onEnvironmentTask) supportedMethods.push("environmentTask"); if (plugin.definition.onEnvironmentProbe) supportedMethods.push("environmentProbe"); if (plugin.definition.onEnvironmentAcquireLease) supportedMethods.push("environmentAcquireLease"); if (plugin.definition.onEnvironmentResumeLease) supportedMethods.push("environmentResumeLease"); diff --git a/packages/shared/src/types/plugin.ts b/packages/shared/src/types/plugin.ts index 384f9e9b4c9..bbfb2c97a07 100644 --- a/packages/shared/src/types/plugin.ts +++ b/packages/shared/src/types/plugin.ts @@ -191,6 +191,8 @@ export interface SandboxProviderCapabilities { } export interface PluginEnvironmentDriverDeclaration { + /** Provides Runner execution, resource preparation, and PRP connections via environmentTask. */ + supportsTasks?: boolean; /** Stable driver key, unique within the plugin. Namespaced by plugin ID at runtime. */ driverKey: string; /** diff --git a/packages/shared/src/validators/plugin.ts b/packages/shared/src/validators/plugin.ts index ca6fa079630..5cc7d174c45 100644 --- a/packages/shared/src/validators/plugin.ts +++ b/packages/shared/src/validators/plugin.ts @@ -177,6 +177,7 @@ export type SandboxProviderCapabilitiesInput = z.infer { + let db: ReturnType; + let stopDb: (() => Promise) | undefined; + beforeAll(async () => { + const started = await startEmbeddedPostgresTestDatabase("environment-task-authority"); + stopDb = started.stop; db = createDb(started.connectionString); + }); + afterAll(async () => { await stopDb?.(); }); + + it("authorizes a new task lease only after scoped prior execution has stopped", async () => { + const companyId = randomUUID(), agentId = randomUUID(), environmentId = randomUUID(), pluginId = randomUUID(); + const oldRunId = randomUUID(), nextRunId = randomUUID(), oldLeaseId = randomUUID(), nextLeaseId = randomUUID(); + await db.insert(companies).values({ id: companyId, name: "Lease transition", issuePrefix: `T${companyId.slice(0, 7).replace(/-/g, "").toUpperCase()}` }); + await db.insert(agents).values({ id: agentId, companyId, name: "Fixture", role: "engineer", adapterType: "paperclip_runner" }); + await db.insert(environments).values({ id: environmentId, companyId, name: `Fixture ${environmentId}`, driver: "plugin", config: {} }); + await db.insert(plugins).values({ id: pluginId, pluginKey: pluginId, packageName: "fixture-continuation", version: "1.0.0", status: "ready", manifestJson: { + capabilities: ["environment.drivers.register"], environmentDrivers: [{ driverKey: "tasks", supportsTasks: true }], + } as never }); + await db.insert(heartbeatRuns).values([ + { id: oldRunId, companyId, agentId, invocationSource: "manual", status: "succeeded" }, + { id: nextRunId, companyId, agentId, invocationSource: "manual", status: "running" }, + ]); + const pin = { driver: "plugin", pluginId, driverKey: "tasks", environmentId, runnerTaskLease: true }; + const prior = { id: oldLeaseId, companyId, environmentId, heartbeatRunId: oldRunId, provider: "plugin:fixture:tasks", providerLeaseId: "old-task" }; + await db.insert(environmentLeases).values([ + { ...prior, status: "released", leasePolicy: "ephemeral", cleanupStatus: "success", releasedAt: new Date(), + metadata: { ...pin, runnerTaskSessionRetained: true, runnerTaskIdentity: { runnerId: "runner", leaseId: oldLeaseId, runId: oldRunId, sessionId: "session" }, remoteExecutionTermination: remoteTerminationReceipt(prior, { providerLeaseId: "old-task", state: "stopped" }) } }, + { id: nextLeaseId, companyId, environmentId, heartbeatRunId: nextRunId, provider: prior.provider, providerLeaseId: "new-task", + status: "active", leasePolicy: "ephemeral", metadata: pin }, + ]); + const priorIdentity = { runnerInstanceId: "runner", normalizedSessionId: "session", environmentLeaseId: oldLeaseId, runId: oldRunId }; + const authorize = () => resolvePluginTaskResume(db, { companyId, leaseId: nextLeaseId, priorIdentity }); + await expect(authorize()).resolves.toBe("old-task"); + for (const field of ["runnerInstanceId", "environmentLeaseId", "normalizedSessionId", "runId"]) { + await expect(resolvePluginTaskResume(db, { companyId, leaseId: nextLeaseId, + priorIdentity: { ...priorIdentity, [field]: randomUUID() } })).rejects.toThrow("not authorized"); + } + await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, oldRunId)); + await expect(authorize()).rejects.toThrow("not authorized"); + await db.update(heartbeatRuns).set({ status: "succeeded" }).where(eq(heartbeatRuns.id, oldRunId)); + await expect(resolvePluginTaskResume(db, { companyId: randomUUID(), leaseId: nextLeaseId, priorIdentity })).rejects.toThrow("not authorized"); + await db.update(environmentLeases).set({ metadata: { ...pin, pluginId: randomUUID() } }).where(eq(environmentLeases.id, nextLeaseId)); + await expect(authorize()).rejects.toThrow("not authorized"); + await db.update(environmentLeases).set({ metadata: pin }).where(eq(environmentLeases.id, nextLeaseId)); + await db.update(environmentLeases).set({ cleanupStatus: "failed" }).where(eq(environmentLeases.id, oldLeaseId)); + await expect(authorize()).rejects.toThrow("not authorized"); + await db.update(environmentLeases).set({ cleanupStatus: "success" }).where(eq(environmentLeases.id, oldLeaseId)); + const workers = { getWorker: () => ({ supportedMethods: ["environmentTask"] }), + call: async (_id: string, _method: string, params: any) => ({ kind: "accepted", taskId: params.taskId }), + } as unknown as PluginWorkerManager; + const runtime = environmentRuntimeService(db, { pluginWorkerManager: workers }); + const runner = { runnerId: "runner", leaseId: oldLeaseId, sessionId: "session", runId: nextRunId, + turnId: "turn", itemId: "item", harness: "codex", protocolMin: 1, protocolMax: 3 }; + await expect(runtime.task({ companyId, leaseId: nextLeaseId, operation: { + kind: "submit", projectIds: [], bootstrapTicket: "fixture", runner, resumeFrom: "old-task", + } })).resolves.toMatchObject({ kind: "accepted" }); + const [second] = await db.select().from(environmentLeases).where(eq(environmentLeases.id, nextLeaseId)); + await db.update(environmentLeases).set({ status: "released", cleanupStatus: "success", releasedAt: new Date(), + metadata: { ...second.metadata, runnerTaskSessionRetained: true, + remoteExecutionTermination: remoteTerminationReceipt(second, { providerLeaseId: "new-task", state: "stopped" }) }, + }).where(eq(environmentLeases.id, nextLeaseId)); + await db.update(heartbeatRuns).set({ status: "succeeded" }).where(eq(heartbeatRuns.id, nextRunId)); + const thirdRunId = randomUUID(), thirdLeaseId = randomUUID(); + await db.insert(heartbeatRuns).values({ id: thirdRunId, companyId, agentId, invocationSource: "manual", status: "running" }); + await db.insert(environmentLeases).values({ id: thirdLeaseId, companyId, environmentId, heartbeatRunId: thirdRunId, + status: "active", provider: prior.provider, providerLeaseId: "third-task", metadata: pin }); + await expect(resolvePluginTaskResume(db, { companyId, leaseId: thirdLeaseId, + priorIdentity: { ...priorIdentity, runId: nextRunId } })).resolves.toBe("new-task"); + const submit = { kind: "submit", projectIds: [], bootstrapTicket: "fixture", runner: { ...runner, runId: thirdRunId }, resumeFrom: "new-task" } as const; + await expect(runtime.task({ companyId, leaseId: thirdLeaseId, operation: { ...submit, projectIds: [] } })).resolves.toMatchObject({ kind: "accepted" }); + await expect(runtime.task({ companyId, leaseId: thirdLeaseId, operation: { ...submit, projectIds: [], runner } })).rejects.toThrow("identity mismatch"); + // Same-run recovery must choose the latest task in the saved lineage, not an earlier physical lease. + const [third] = await db.select().from(environmentLeases).where(eq(environmentLeases.id, thirdLeaseId)); + await db.update(environmentLeases).set({ status: "released", cleanupStatus: "success", releasedAt: new Date(), + metadata: { ...third.metadata, runnerTaskSessionRetained: true, + remoteExecutionTermination: remoteTerminationReceipt(third, { providerLeaseId: "third-task", state: "stopped" }) }, + }).where(eq(environmentLeases.id, thirdLeaseId)); + const retryLeaseId = randomUUID(); + await db.insert(environmentLeases).values({ id: retryLeaseId, companyId, environmentId, heartbeatRunId: thirdRunId, + status: "active", provider: prior.provider, providerLeaseId: "retry-task", metadata: pin }); + await expect(runtime.task({ companyId, leaseId: retryLeaseId, operation: { ...submit, projectIds: [], resumeFrom: "third-task" } })).resolves.toMatchObject({ kind: "accepted" }); + const [retry] = await db.select().from(environmentLeases).where(eq(environmentLeases.id, retryLeaseId)); + await db.update(environmentLeases).set({ status: "released", cleanupStatus: "success", releasedAt: new Date(), + metadata: { ...retry.metadata, runnerTaskSessionRetained: true, + remoteExecutionTermination: remoteTerminationReceipt(retry, { providerLeaseId: "retry-task", state: "stopped" }) }, + }).where(eq(environmentLeases.id, retryLeaseId)); + const finalLeaseId = randomUUID(); + await db.insert(environmentLeases).values({ id: finalLeaseId, companyId, environmentId, heartbeatRunId: thirdRunId, + status: "active", provider: prior.provider, providerLeaseId: "final-task", metadata: pin }); + await expect(resolvePluginTaskResume(db, { companyId, leaseId: finalLeaseId, + priorIdentity: { ...priorIdentity, runId: thirdRunId } })).resolves.toBe("retry-task"); + + }); + +}); diff --git a/server/src/__tests__/environment-task-runtime.test.ts b/server/src/__tests__/environment-task-runtime.test.ts new file mode 100644 index 00000000000..fec9d2c2d76 --- /dev/null +++ b/server/src/__tests__/environment-task-runtime.test.ts @@ -0,0 +1,199 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { PgDialect } from "drizzle-orm/pg-core"; +import type { Db } from "@paperclipai/db"; +import { environmentTaskOperationSchema, parseEnvironmentTaskResult } from "@paperclipai/plugin-sdk"; +import { executeEnvironmentTask } from "../services/environment-task-runtime.js"; + +const state = vi.hoisted(() => ({ plugin: {} as any })); +vi.mock("../services/plugin-registry.js", () => ({ pluginRegistryService: () => ({ getById: vi.fn(async () => state.plugin) }) })); +const leaseId = "10000000-0000-4000-8000-000000000001"; +const projectId = "10000000-0000-4000-8000-000000000002"; +const secondProjectId = "10000000-0000-4000-8000-000000000003"; +const row = () => ({ + lease: { id: leaseId, companyId: "company", environmentId: "environment", providerLeaseId: "attempt-1", heartbeatRunId: "run", issueId: "issue", status: "active", expiresAt: null as Date | null, + metadata: { driver: "plugin", pluginId: "original", driverKey: "tasks" } }, + environment: { id: "environment", config: { pluginKey: "test.provider", driverKey: "tasks", driverConfig: {} } } as { id: string; config: Record } | null, + run: { id: "run", agentId: "agent", status: "running" } as { id: string; agentId: string; status: string } | null, +}); +function database(value: ReturnType | null = row(), projectRows = [{ id: projectId }]) { + const results = [value ? [value] : [], [{ id: "issue" }]]; + const query: any = { from: () => query, leftJoin: () => query, where: vi.fn(() => query), limit: () => Promise.resolve(results.shift()), then: (resolve: (rows: typeof projectRows) => unknown) => Promise.resolve(projectRows).then(resolve) }; + const update: any = { set: () => update, where: () => update, + returning: async () => [{ id: leaseId }], then: (resolve: (value: unknown) => unknown) => Promise.resolve([]).then(resolve) }; + return { db: { select: () => query, update: () => update } as unknown as Db, query }; +} +function worker(result: unknown = { kind: "accepted", taskId: "attempt-1" }) { + return { getWorker: vi.fn(() => ({ supportedMethods: ["environmentTask"] })), call: vi.fn(async () => result) }; +} +const submit = { kind: "submit" as const, projectIds: [projectId], runner: { protocolMin: 1, protocolMax: 2, harness: "codex", runnerId: "runner", leaseId, runId: "run", sessionId: "session", turnId: "turn", itemId: "item" }, bootstrapTicket: "transient-test-ticket" }; + +beforeEach(() => { + state.plugin = { id: "original", pluginKey: "test.provider", status: "ready", manifestJson: { capabilities: ["environment.drivers.register"], environmentDrivers: [{ driverKey: "tasks", supportsTasks: true }] } }; +}); +describe("plugin-provided Runner execution", () => { + it("accepts artifact identity and excludes filesystem transfer from the contract", () => { + const operation = environmentTaskOperationSchema.parse({ kind: "describe" }); + const description = { kind: "description", taskId: "attempt-1", runner: { version: "fixture", digest: `sha256:${"a".repeat(64)}` } }; + expect(parseEnvironmentTaskResult(operation, "attempt-1", description)).toEqual(description); + for (const field of ["resources", "environment", "connection", "codexCommand"]) { + expect(() => parseEnvironmentTaskResult(operation, "attempt-1", { ...description, [field]: {} })).toThrow(); + } + for (const kind of ["checkpoint", "read_file", "check_project_auth"]) { + expect(environmentTaskOperationSchema.safeParse({ kind }).success).toBe(false); + } + }); + it("allows only the private, unexpired setup identity on an ad-hoc lease", async () => { + const value = row(); + Object.assign(value.lease, { heartbeatRunId: null, issueId: null }); value.run = null; + Object.assign(value.lease.metadata, { runnerTaskProbe: { runId: "probe-run", sessionId: "probe-session", expiresAtMs: Date.now() + 60_000 } }); + const operation = { ...submit, projectIds: [], runner: { ...submit.runner, runId: "probe-run", sessionId: "probe-session" }, + launch: { runnerVersion: "fixture", runnerDigest: `sha256:${"a".repeat(64)}`, + maxOutboxBytes: 262144, p0ReserveBytes: 65536, maxRuntimeMs: 0, maxLifetimeMs: 0, + lifecyclePolicy: { mode: "per_turn" as const, idleTimeoutMs: null }, + } }; + const workers = worker(); + await expect(executeEnvironmentTask(database(value).db, workers as never, { companyId: "company", leaseId, operation })).resolves.toMatchObject({ kind: "accepted" }); + expect(workers.call).toHaveBeenCalledWith("original", "environmentTask", expect.objectContaining({ runId: "probe-run", projectIds: [] }), 15_000); + const lifetime = (workers.call.mock.calls[0] as unknown as [string, string, { operation: typeof operation }])[2].operation.launch.maxRuntimeMs; + expect(lifetime).toBeGreaterThan(0); expect(lifetime).toBeLessThanOrEqual(60_000); + expect((workers.call.mock.calls[0] as unknown as [string, string, { operation: typeof operation }])[2].operation.launch).not.toHaveProperty("maxLifetimeMs"); + for (const invalid of [{ ...operation, projectIds: [projectId] }, { ...operation, runner: { ...operation.runner, sessionId: "other" } }, submit]) { + await expect(executeEnvironmentTask(database(value).db, workers as never, { companyId: "company", leaseId, operation: invalid })).rejects.toThrow("identity mismatch"); + } + Object.assign(value.lease.metadata, { runnerTaskProbe: { runId: "probe-run", sessionId: "probe-session", expiresAtMs: Date.now() - 1 } }); + await expect(executeEnvironmentTask(database(value).db, workers as never, { companyId: "company", leaseId, operation })).rejects.toThrow("not active"); + await expect(executeEnvironmentTask(database(value).db, workers as never, { companyId: "company", leaseId, operation: { kind: "stop" } })).resolves.toMatchObject({ kind: "accepted" }); + }); + it("dispatches with host-derived scope and a persisted attempt identity", async () => { + const { db, query } = database(); const workers = worker(); + await expect(executeEnvironmentTask(db, workers as never, { companyId: "company", leaseId, operation: submit })).resolves.toEqual({ kind: "accepted", taskId: "attempt-1" }); + const scope = new PgDialect().sqlToQuery(query.where.mock.calls[0][0]); + expect(scope.sql).toContain('"environment_leases"."company_id" ='); + expect(scope.sql).toContain('"environment_leases"."id" ='); + expect(scope.params).toEqual([leaseId, "company"]); + expect(workers.call).toHaveBeenCalledWith("original", "environmentTask", expect.objectContaining({ + taskId: "attempt-1", companyId: "company", agentId: "agent", projectIds: [projectId], runId: "run", operation: submit, + }), 15_000); + }); + it("validates every requested project in the task company before dispatch", async () => { + const operation = { ...submit, projectIds: [projectId, secondProjectId] }; + const { db, query } = database(row(), [{ id: projectId }, { id: secondProjectId }]); + const workers = worker(); + await executeEnvironmentTask(db, workers as never, { companyId: "company", leaseId, operation }); + expect(workers.call).toHaveBeenCalledWith("original", "environmentTask", expect.objectContaining({ projectIds: operation.projectIds }), 15_000); + const scope = new PgDialect().sqlToQuery(query.where.mock.calls[2][0]); + expect(scope.sql).toContain('"projects"."company_id" ='); + expect(scope.params).toEqual(["company", projectId, secondProjectId]); + const rejected = worker(); + await expect(executeEnvironmentTask(database(row(), [{ id: projectId }]).db, rejected as never, { companyId: "company", leaseId, operation })).rejects.toThrow("project unavailable"); + expect(rejected.call).not.toHaveBeenCalled(); + }); + it("validates unique project IDs and allows tasks with no project mounts", async () => { + for (const projectIds of [[projectId, projectId], ["invalid"], Array(65).fill(projectId)]) { + expect(environmentTaskOperationSchema.safeParse({ ...submit, projectIds }).success).toBe(false); + } + const workers = worker(); + await executeEnvironmentTask(database().db, workers as never, { companyId: "company", leaseId, operation: { ...submit, projectIds: [] } }); + expect(workers.call).toHaveBeenCalledWith("original", "environmentTask", expect.objectContaining({ projectIds: [] }), 15_000); + }); + it("rejects an absent or cross-company lease before calling a worker", async () => { + const workers = worker(); + await expect(executeEnvironmentTask(database(null).db, workers as never, { companyId: "other", leaseId, operation: submit })).rejects.toThrow("lease unavailable"); + expect(workers.call).not.toHaveBeenCalled(); + }); + it("rejects mismatched Runner binding and inactive submission", async () => { + const workers = worker(); + await expect(executeEnvironmentTask(database().db, workers as never, { companyId: "company", leaseId, operation: { ...submit, runner: { ...submit.runner, runId: "other" } } })).rejects.toThrow("identity mismatch"); + const value = row(); value.lease.status = "released"; + await expect(executeEnvironmentTask(database(value).db, workers as never, { companyId: "company", leaseId, operation: submit })).rejects.toThrow("not active"); + expect(workers.call).not.toHaveBeenCalled(); + }); + it("uses the pinned provider for cleanup after environment edits", async () => { + const value = row(); value.lease.status = "released"; value.environment!.config.pluginKey = "replacement"; + const workers = worker(); + await executeEnvironmentTask(database(value).db, workers as never, { companyId: "company", leaseId, operation: { kind: "stop" } }); + expect(workers.call).toHaveBeenCalledWith("original", "environmentTask", expect.objectContaining({ config: {}, operation: { kind: "stop" } }), 15_000); + }); + it("keeps cleanup available after environment, run, and issue deletion", async () => { + const value = row(); value.environment = null; value.run = null; + for (const kind of ["stop", "complete", "status"] as const) { + const workers = worker(kind === "status" ? { kind: "status", taskId: "attempt-1", phase: "cancelled" } : undefined); + await executeEnvironmentTask(database(value).db, workers as never, { companyId: "company", leaseId, operation: { kind } }); + expect(workers.call).toHaveBeenCalledWith("original", "environmentTask", expect.objectContaining({ + config: {}, environmentId: null, runId: null, agentId: null, projectIds: [], + }), 15_000); + } + await expect(executeEnvironmentTask(database(value).db, worker() as never, { companyId: "company", leaseId, operation: submit })).rejects.toThrow("not active"); + }); + it("preserves opaque provider IDs in requests and receipts", async () => { + const value = row(); value.lease.providerLeaseId = `provider:attempt.${"x".repeat(120)}`; + const workers = worker({ kind: "accepted", taskId: value.lease.providerLeaseId }); + await expect(executeEnvironmentTask(database(value).db, workers as never, { companyId: "company", leaseId, operation: { kind: "stop" } })).resolves.toMatchObject({ taskId: value.lease.providerLeaseId }); + }); + it("returns a validated connection only while the lease and run are active", async () => { + const connection = { kind: "connection", taskId: "attempt-1", endpoint: { + kind: "authenticated_websocket", websocketUrl: "wss://runner.example.test/connect", generation: "attempt-1", + secretHeaders: [{ name: "Authorization", value: "Bearer fixture" }], + } }; + const input = { companyId: "company", leaseId, operation: { kind: "connection" as const } }; + await expect(executeEnvironmentTask(database().db, worker(connection) as never, input)).resolves.toEqual(connection); + for (const state of ["expired", "finished", "deleted"] as const) { + const value = row(); + if (state === "expired") value.lease.expiresAt = new Date(0); + if (state === "finished") value.run!.status = "succeeded"; + if (state === "deleted") value.environment = null; + const workers = worker(connection); + await expect(executeEnvironmentTask(database(value).db, workers as never, input)).rejects.toThrow("not active"); + expect(workers.call).not.toHaveBeenCalled(); + } + for (const endpoint of [ + { ...connection.endpoint, websocketUrl: "ws://runner.example.test/connect" }, + { ...connection.endpoint, websocketUrl: "wss://runner.example.test/connect?token=fixture" }, + { ...connection.endpoint, secretHeaders: [{ name: "Authorization", value: "fixture\r\nInjected: true" }] }, + { ...connection.endpoint, secretHeaders: [{ name: "Invalid Header", value: "fixture" }] }, + ]) { + await expect(executeEnvironmentTask(database().db, worker({ ...connection, endpoint }) as never, input)).rejects.toThrow("reconcile the same task"); + } + }); + it.each(["manifest", "worker"])("requires live %s support", async kind => { + const workers = worker(); + if (kind === "manifest") state.plugin.manifestJson.environmentDrivers[0].supportsTasks = false; + else workers.getWorker.mockReturnValue({ supportedMethods: [] }); + await expect(executeEnvironmentTask(database().db, workers as never, { companyId: "company", leaseId, operation: submit })).rejects.toThrow("provider unavailable"); + expect(workers.call).not.toHaveBeenCalled(); + }); + it("rejects wrong task receipts and sanitizes worker errors", async () => { + await expect(executeEnvironmentTask(database().db, worker({ kind: "accepted", taskId: "other" }) as never, { companyId: "company", leaseId, operation: submit })).rejects.toThrow("reconcile the same task"); + const workers = worker(); workers.call.mockRejectedValue(new Error("private-credential")); + await expect(executeEnvironmentTask(database().db, workers as never, { companyId: "company", leaseId, operation: submit })).rejects.toThrow(/^Plugin-provided Runner execution operation unavailable; reconcile the same task before retrying$/); + }); + it("accepts PRP identity characters and length limits", () => { + for (const field of ["runnerId", "leaseId", "runId", "sessionId", "turnId", "itemId"]) { + for (const value of ["session:part.1", "a".repeat(160)]) { + expect(environmentTaskOperationSchema.safeParse({ ...submit, runner: { ...submit.runner, [field]: value } }).success).toBe(true); + } + for (const value of ["a".repeat(161), "../escape", "a/b", ""]) { + expect(environmentTaskOperationSchema.safeParse({ ...submit, runner: { ...submit.runner, [field]: value } }).success).toBe(false); + } + } + }); + it("validates the client's inclusive PRP version range", () => { + for (const range of [{ protocolMin: 1, protocolMax: 1 }, { protocolMin: 1, protocolMax: 2 }, { protocolMin: 3, protocolMax: 5 }]) { + expect(environmentTaskOperationSchema.safeParse({ ...submit, runner: { ...submit.runner, ...range } }).success).toBe(true); + } + for (const range of [ + { protocolMin: 0, protocolMax: 2 }, { protocolMin: 2, protocolMax: 1 }, + { protocolMin: 1.5, protocolMax: 2 }, { protocolMin: 1, protocolMax: 2.5 }, + { protocolMin: 1, protocolMax: Number.MAX_SAFE_INTEGER + 1 }, + { protocolMin: undefined, protocolMax: 2 }, { protocolMin: 1, protocolMax: undefined }, + ]) { + expect(environmentTaskOperationSchema.safeParse({ ...submit, runner: { ...submit.runner, ...range } }).success).toBe(false); + } + }); + it("validates operation and result shape without making acceptance mean readiness", () => { + expect(environmentTaskOperationSchema.safeParse({ ...submit, surprise: true }).success).toBe(false); + expect(environmentTaskOperationSchema.safeParse({ ...submit, runner: { ...submit.runner, unknownField: true } }).success).toBe(false); + expect(() => parseEnvironmentTaskResult({ kind: "status" }, "attempt-1", { kind: "accepted", taskId: "attempt-1" })).toThrow(); + expect(parseEnvironmentTaskResult(submit, "attempt-1", { kind: "accepted", taskId: "attempt-1" }).kind).toBe("accepted"); + }); +}); diff --git a/server/src/__tests__/plugin-environment-driver-seam.test.ts b/server/src/__tests__/plugin-environment-driver-seam.test.ts index 56c13d091ea..1e691fd8a04 100644 --- a/server/src/__tests__/plugin-environment-driver-seam.test.ts +++ b/server/src/__tests__/plugin-environment-driver-seam.test.ts @@ -8,6 +8,7 @@ import { parseMessage, serializeMessage, } from "../../../packages/plugins/sdk/src/protocol.js"; +import type { PluginEnvironmentTaskParams } from "../../../packages/plugins/sdk/src/environment-tasks.js"; import { definePlugin } from "../../../packages/plugins/sdk/src/define-plugin.js"; import { startWorkerRpcHost } from "../../../packages/plugins/sdk/src/worker-rpc-host.js"; import { pluginManifestV1Schema, type PaperclipPluginManifestV1 } from "@paperclipai/shared"; @@ -94,6 +95,54 @@ describe("plugin environment driver seam", () => { }); }); + it("negotiates and validates plugin-provided Runner execution RPCs", async () => { + const calls: unknown[] = []; + const plugin = definePlugin({ + async setup() {}, + async onEnvironmentTask(input) { + calls.push(input); + return { kind: "accepted", taskId: input.taskId }; + }, + }); + const stdin = new PassThrough(); + const stdout = new PassThrough(); + const host = startWorkerRpcHost({ plugin, stdin, stdout }); + const responses: unknown[] = []; + stdout.on("data", chunk => { + for (const line of String(chunk).split("\n").filter(Boolean)) responses.push(parseMessage(line)); + }); + try { + const manifest = { ...baseManifest, environmentDrivers: [{ ...baseManifest.environmentDrivers![0], supportsTasks: true }] }; + expect(pluginManifestV1Schema.parse(manifest).environmentDrivers![0].supportsTasks).toBe(true); + stdin.write(serializeMessage(createRequest("initialize", { + manifest, config: {}, instanceInfo: { instanceId: "instance-1", hostVersion: "1.0.0" }, apiVersion: 1, + }, 1))); + await waitForResponses(responses, 1); + const initialized = responses[0]; + expect(isJsonRpcSuccessResponse(initialized)).toBe(true); + if (!isJsonRpcSuccessResponse(initialized)) return; + expect(initialized.result.supportedMethods).toContain("environmentTask"); + const projectIds = ["10000000-0000-4000-8000-000000000001", "10000000-0000-4000-8000-000000000002"]; + const params: PluginEnvironmentTaskParams = { driverKey: "fake-plugin", companyId: "company", environmentId: "environment", config: {}, + taskId: "attempt", runId: "run", agentId: "agent", projectIds, lease: { providerLeaseId: "attempt" }, + operation: { kind: "submit", projectIds, bootstrapTicket: "transient-test-ticket", runner: { + protocolMin: 1, protocolMax: 2, harness: "codex", runnerId: "runner", leaseId: "lease", runId: "run", sessionId: "session:part.1", turnId: "turn", itemId: "item", + } } }; + + stdin.write(serializeMessage(createRequest("environmentTask", params, 2))); + await waitForResponses(responses, 2); + expect(responses[1]).toMatchObject({ result: { kind: "accepted", taskId: "attempt" } }); + expect(calls[0]).toMatchObject({ projectIds, operation: { kind: "submit", projectIds } }); + stdin.write(serializeMessage(createRequest("environmentTask", { ...params, operation: { kind: "status" } }, 3))); + await waitForResponses(responses, 3); + expect(isJsonRpcErrorResponse(responses[2])).toBe(true); // Wrong response kind. + stdin.write(serializeMessage(createRequest("environmentTask", { ...params, operation: { kind: "shell", command: "echo test" } }, 4))); + await waitForResponses(responses, 4); + expect(isJsonRpcErrorResponse(responses[3])).toBe(true); + expect(calls).toHaveLength(2); // Invalid operation did not reach the plugin. + } finally { host.stop(); } + }); + it("dispatches environment driver worker hooks and reports support", async () => { const plugin = definePlugin({ async setup() {}, diff --git a/server/src/__tests__/plugin-task-credential-writeback.test.ts b/server/src/__tests__/plugin-task-credential-writeback.test.ts new file mode 100644 index 00000000000..8ef383854c9 --- /dev/null +++ b/server/src/__tests__/plugin-task-credential-writeback.test.ts @@ -0,0 +1,117 @@ +import { randomUUID } from "node:crypto"; +import { mkdtemp, realpath, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { and, eq } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import { agents, companies, companyMemberships, createDb, environmentLeases, environments, heartbeatRuns, plugins } from "@paperclipai/db"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; +import { aiConnectionService } from "../services/ai-connections.js"; +import { prepareManagedAiRuntime } from "../services/ai-connection-runtime.js"; +import { registerPluginTaskCredentialHandoff } from "../services/plugin-task-credential-writeback.js"; +import { executeEnvironmentTask, settleDispatchedEnvironmentTask } from "../services/environment-task-runtime.js"; +import { secretService } from "../services/secrets.js"; +import type { PluginWorkerManager } from "../services/plugin-worker-manager.js"; + +vi.mock("../services/subscription-refresh.js", () => ({ refreshSubscriptionConnection: async () => {} })); +const support = await getEmbeddedPostgresTestSupport(); +const suite = support.supported ? describe : describe.skip; +const auth = (token: string, year = 2020) => JSON.stringify({ tokens: { account_id: "fixture-account", access_token: "fixture-access", refresh_token: token }, last_refresh: `${year}-01-01T00:00:00Z` }); + +suite("durable plugin task credential handoff", () => { + let database: Awaited>, db: ReturnType, home: string; + beforeAll(async () => { + home = await realpath(await mkdtemp(join(tmpdir(), "plugin-task-secret-"))); + vi.stubEnv("PAPERCLIP_HOME", home); vi.stubEnv("PAPERCLIP_INSTANCE_ID", "plugin-task-fixture"); + database = await startEmbeddedPostgresTestDatabase("plugin-task-credential-"); db = createDb(database.connectionString); + }, 90_000); + afterAll(async () => { await database?.cleanup(); vi.unstubAllEnvs(); if (home) await rm(home, { recursive: true, force: true }); }); + async function fixture(dispatched = true) { + const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID(), pluginId = randomUUID(), environmentId = randomUUID(), leaseId = randomUUID(); + const owner = `fixture-${randomUUID()}`; + await db.insert(companies).values({ id: companyId, name: "Task credential fixture", issuePrefix: `C${companyId.slice(0, 6)}` }); + await db.insert(agents).values({ id: agentId, companyId, name: "Fixture", adapterType: "paperclip_runner" }); + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const service = aiConnectionService(db); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Fixture Codex", loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("fixture-original")); + const [selection] = await service.quotaAccounts(companyId, owner); + const managed = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "paperclip_runner", + binding: { mode: "responsible_user", provider: "openai", method: "subscription" }, config: { provider: "codex" } }); + expect(managed.credentialWriteback).toBeDefined(); + await db.insert(environments).values({ id: environmentId, name: `Fixture ${environmentId}`, driver: "plugin", status: "active", config: { pluginKey: pluginId, driverKey: "tasks", driverConfig: {} } }); + await db.insert(plugins).values({ id: pluginId, pluginKey: pluginId, packageName: "fixture-tasks", version: "1.0.0", status: "ready", manifestJson: { + capabilities: ["environment.drivers.register"], environmentDrivers: [{ driverKey: "tasks", supportsTasks: true }], + } as never }); + await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, invocationSource: "manual", status: "running" }); + await db.insert(environmentLeases).values({ id: leaseId, companyId, environmentId, heartbeatRunId: runId, provider: `plugin:${pluginId}:tasks`, providerLeaseId: "fixture-task", status: "active", + metadata: { driver: "plugin", pluginId, driverKey: "tasks", runnerTaskLease: true, runnerTaskDispatched: dispatched, runnerTaskExecutionStopped: false } }); + await registerPluginTaskCredentialHandoff(db, { companyId, leaseId }, managed.sessionIdentity, managed.credentialWriteback!); + // Recovery must not depend on the original controller or its private home. + await rm(managed.home!, { recursive: true, force: true }); + const calls: string[] = []; + let failExport = false; + const workers = { getWorker: () => ({ supportedMethods: ["environmentTask"] }), + call: async (_plugin: string, _method: string, params: any) => { + calls.push(params.operation.kind); + expect(params.companyId).toBe(companyId); expect(params.taskId).toBe("fixture-task"); + if (params.operation.kind === "status") return { kind: "status", taskId: "fixture-task", phase: "cancelled", executionStopped: true }; + if (params.operation.kind === "credentials") { + expect(params.operation.accountIdentity).toBe(managed.sessionIdentity); + if (failExport) throw new Error("fixture export unavailable"); + return { kind: "credentials", taskId: "fixture-task", credential: { type: "codex-auth-json", authJson: auth("fixture-refreshed", 2021) } }; + } + return { kind: "accepted", taskId: "fixture-task" }; + } } as unknown as PluginWorkerManager; + const lease = () => db.select().from(environmentLeases).where(and(eq(environmentLeases.id, leaseId), eq(environmentLeases.companyId, companyId))).then(rows => rows[0]!); + return { companyId, agentId, runId, environmentId, leaseId, owner, selection: selection!, managed, workers, calls, lease, + failExport(value: boolean) { failExport = value; }, credential: () => service.credential(selection!), + settle: () => settleDispatchedEnvironmentTask(db, workers, { companyId, id: leaseId }) }; + } + it("persists refreshed credentials before completing a stopped task, after controller files disappear", async () => { + const f = await fixture(); + expect(await f.settle()).toBe(true); + expect(await f.credential()).toBe(auth("fixture-refreshed", 2021)); + expect(f.calls.indexOf("credentials")).toBeLessThan(f.calls.indexOf("complete")); + const saved = await f.lease(); + expect(saved.metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete" }); + expect(JSON.stringify(saved.metadata)).not.toContain("fixture-refreshed"); + expect(JSON.stringify(saved.metadata)).not.toContain(f.managed.home!); + }); + it("keeps a failed transfer pending and recovers through a fresh cleanup after run and environment deletion", async () => { + const f = await fixture(); f.failExport(true); + await expect(f.settle()).rejects.toThrow("termination is unconfirmed"); + expect(f.calls).not.toContain("complete"); + expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "pending" }); + expect(await f.credential()).toBe(auth("fixture-original")); + await db.delete(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + await db.delete(environments).where(eq(environments.id, f.environmentId)); + f.failExport(false); + expect(await f.settle()).toBe(true); + expect(await f.credential()).toBe(auth("fixture-refreshed", 2021)); + expect(f.calls.at(-1)).toBe("complete"); + }); + it("does not overwrite an explicit account replacement", async () => { + const f = await fixture(); + await secretService(db).rotate(f.managed.credentialWriteback!.secretId, { value: auth("fixture-reconnected", 2022) }, { userId: f.owner }); + expect(await f.settle()).toBe(true); + expect(f.calls).not.toContain("credentials"); + expect(await f.credential()).toBe(auth("fixture-reconnected", 2022)); + }); + it("allows concurrent cleanup to complete the same durable handoff", async () => { + const f = await fixture(); + expect(await Promise.all([f.settle(), f.settle()])).toEqual([true, true]); + expect(await f.credential()).toBe(auth("fixture-refreshed", 2021)); + expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete" }); + }); + it("retires an undispatched task without exporting credentials and blocks a late submit", async () => { + const f = await fixture(false); + expect(await f.settle()).toBe(false); + expect(f.calls).toEqual([]); + expect((await f.lease()).metadata?.runnerTaskCredentialHandoff).toMatchObject({ state: "complete" }); + await expect(executeEnvironmentTask(db, f.workers, { companyId: f.companyId, leaseId: f.leaseId, operation: { + kind: "submit", projectIds: [], bootstrapTicket: "fixture-ticket", runner: { harness: "codex", protocolMin: 1, protocolMax: 3, + runnerId: "runner", leaseId: f.leaseId, runId: f.runId, sessionId: "session", turnId: "turn", itemId: "item" }, + } })).rejects.toThrow("not active"); + expect(f.calls).toEqual([]); + }); +}); diff --git a/server/src/services/ai-connection-runtime.ts b/server/src/services/ai-connection-runtime.ts index c201f2451fd..74a4173cb7c 100644 --- a/server/src/services/ai-connection-runtime.ts +++ b/server/src/services/ai-connection-runtime.ts @@ -1,3 +1,4 @@ +import type { PluginTaskCredentialWritebackBinding } from "./plugin-task-credential-writeback.js"; import { OPENROUTER_ATTRIBUTION_HEADERS, openRouterClaudeHeaders } from "@paperclipai/shared/openrouter-attribution"; import { createHash } from "node:crypto"; import { HttpError, unprocessable } from "../errors.js"; @@ -385,6 +386,10 @@ export async function prepareManagedAiRuntime( const sessionIdentity = `${selection.grant.id}:${input.responsibleUserId ?? "shared"}:${noAuth ? "no-auth" : `${credentialRef!.secretId}:${freshness!.epoch}`}`; return { sessionIdentity, + ...(subscriptionFile && input.binding.provider === "openai" && credentialRef && freshness ? { + credentialWriteback: { grantId: selection.grant.id, secretId: credentialRef.secretId, sessionEpoch: freshness.epoch, + versionSelector: credentialRef.versionSelector ?? "latest" } satisfies PluginTaskCredentialWritebackBinding, + } : {}), config: { ...input.config, ...projected?.config, diff --git a/server/src/services/environment-runtime.ts b/server/src/services/environment-runtime.ts index 35a57175092..7d14dbf7c81 100644 --- a/server/src/services/environment-runtime.ts +++ b/server/src/services/environment-runtime.ts @@ -1,3 +1,5 @@ +import { executeEnvironmentTask, pluginEnvironmentSupportsTasks, settleDispatchedEnvironmentTask } from "./environment-task-runtime.js"; +import type { PluginEnvironmentTaskOperation } from "@paperclipai/plugin-sdk"; import { allowLegacyShutdownWorkspaceCleanup, legacyShutdownWorkspaceResourceProtected } from "./legacy-shutdown-workspace-settlement.js"; import { nativeWorkspaceExportPendingForLease } from "./native-runtime/native-workspace-finalization-state.js"; import { beginIdleTrackedWork } from "./task-admission.js"; @@ -3594,16 +3596,16 @@ function createPluginEnvironmentDriver( return { plugin, driver }; } - async function resolvePluginDriverForRelease(input: EnvironmentDriverReleaseInput) { + async function resolvePluginDriverForRelease(input: { environment: Environment | null; lease: EnvironmentDriverReleaseInput["lease"] }) { const metadata = input.lease.metadata ?? {}; const metadataPluginId = readString(metadata.pluginId); const metadataPluginKey = readString(metadata.pluginKey); const metadataDriverKey = readString(metadata.driverKey); - const currentConfig = tryParseCurrentPluginConfig(input.environment); + const currentConfig = input.environment ? tryParseCurrentPluginConfig(input.environment) : null; if (!metadataPluginId && !metadataPluginKey && !metadataDriverKey) { if (!currentConfig) { - throw new Error(`Expected plugin environment config for driver "${input.environment.driver}".`); + throw new Error(`Expected pinned plugin environment config.`); } const { plugin } = await resolvePluginDriver(currentConfig); return { @@ -3698,6 +3700,8 @@ function createPluginEnvironmentDriver( metadata: { ...(input.agentId ? { agentId: input.agentId } : {}), providerMetadata: providerLease.metadata ?? {}, + environmentId: input.environment.id, + ...(input.adapterType === "paperclip_runner" && driver.supportsTasks === true ? { runnerTaskLease: true } : {}), driver: input.environment.driver, executionWorkspaceMode: input.executionWorkspaceMode, pluginId: plugin.id, @@ -3708,6 +3712,9 @@ function createPluginEnvironmentDriver( }, async releaseRunLease(input) { + await db.update(environmentLeases).set({ metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || '{"pluginLeaseCleanupMethod":"release"}'::jsonb` }).where(and(eq(environmentLeases.id, input.lease.id), eq(environmentLeases.companyId, input.lease.companyId))); + const taskStopped = await settleDispatchedEnvironmentTask(db, workerManager, input.lease) || + (input.lease.metadata?.runnerTaskLease === true && input.lease.metadata?.runnerTaskDispatched !== true); const { plugin, driverKey, driverConfig } = await resolvePluginDriverForRelease(input); await workerManager.call(plugin.id, "environmentReleaseLease", { driverKey, @@ -3718,17 +3725,18 @@ function createPluginEnvironmentDriver( providerLeaseId: input.lease.providerLeaseId, leaseMetadata: input.lease.metadata ?? undefined, }); - return await environmentsSvc.releaseLease(input.lease.id, input.status, { cleanupStatus: "success" }); + return await environmentsSvc.releaseLease(input.lease.id, input.status, { + cleanupStatus: "success", + ...(taskStopped ? { remoteExecutionTermination: remoteTerminationReceipt(input.lease, + { providerLeaseId: input.lease.providerLeaseId, state: "stopped" }) } : {}), + }); }, async resumeRunLease(input) { if (!input.lease.providerLeaseId) { throw new Error(`Plugin environment lease "${input.lease.id}" does not have a provider lease id to resume.`); } - const { pluginKey, driverKey, driverConfig } = await resolvePluginDriverForRelease({ - ...input, - status: "released", - }); + const { pluginKey, driverKey, driverConfig } = await resolvePluginDriverForRelease(input); return await resumePluginEnvironmentLease({ db, workerManager, @@ -3746,10 +3754,10 @@ function createPluginEnvironmentDriver( }, async destroyRunLease(input) { - const { pluginKey, driverKey, driverConfig } = await resolvePluginDriverForRelease({ - ...input, - status: "failed", - }); + await db.update(environmentLeases).set({ metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || '{"pluginLeaseCleanupMethod":"destroy"}'::jsonb` }).where(and(eq(environmentLeases.id, input.lease.id), eq(environmentLeases.companyId, input.lease.companyId))); + const taskStopped = await settleDispatchedEnvironmentTask(db, workerManager, input.lease) || + (input.lease.metadata?.runnerTaskLease === true && input.lease.metadata?.runnerTaskDispatched !== true); + const { pluginKey, driverKey, driverConfig } = await resolvePluginDriverForRelease(input); await destroyPluginEnvironmentLease({ db, workerManager, @@ -3769,14 +3777,37 @@ function createPluginEnvironmentDriver( cleanupStatus: "success", ...(input.lease.status === "pending_cleanup" && typeof input.lease.metadata?.pendingCleanupAttemptId === "string" ? { expectedPendingCleanupAttemptId: input.lease.metadata.pendingCleanupAttemptId } : {}), + ...(taskStopped ? { remoteExecutionTermination: remoteTerminationReceipt(input.lease, + { providerLeaseId: input.lease.providerLeaseId, state: "stopped" }) } : {}), }); }, + async isPendingCleanupWorkerReady(input) { + try { + const { plugin } = await resolvePluginDriverForRelease(input); + const methods = workerManager.getWorker(plugin.id)?.supportedMethods ?? []; + return methods.includes(input.lease.metadata?.pluginLeaseCleanupMethod === "destroy" ? "environmentDestroyLease" : "environmentReleaseLease") && + (input.lease.metadata?.runnerTaskDispatched !== true || methods.includes("environmentTask")); + } catch { return false; } + }, + + async retryPendingSandboxTeardown(input) { + const taskStopped = await settleDispatchedEnvironmentTask(db, workerManager, input.lease) || + (input.lease.metadata?.runnerTaskLease === true && input.lease.metadata?.runnerTaskDispatched !== true); + const { plugin, driverKey, driverConfig } = await resolvePluginDriverForRelease(input); + const environmentId = input.environment?.id ?? readString(input.lease.metadata?.environmentId); + if (!environmentId) throw new Error("Plugin lease cleanup has no saved environment identity"); + const receipt = await workerManager.call(plugin.id, input.lease.metadata?.pluginLeaseCleanupMethod === "destroy" ? "environmentDestroyLease" : "environmentReleaseLease", { + driverKey, companyId: input.lease.companyId, environmentId, issueId: input.lease.issueId, + config: driverConfig, providerLeaseId: input.lease.providerLeaseId, leaseMetadata: input.lease.metadata ?? undefined, + }); + return taskStopped ? { providerLeaseId: input.lease.providerLeaseId, state: "stopped" } : receipt; + }, + async realizeWorkspace(input) { const { plugin, pluginKey, driverKey, driverConfig } = await resolvePluginDriverForRelease({ environment: input.environment, lease: input.lease, - status: "released", }); return await realizePluginEnvironmentWorkspace({ db, @@ -3807,7 +3838,6 @@ function createPluginEnvironmentDriver( const { plugin, pluginKey, driverKey, driverConfig } = await resolvePluginDriverForRelease({ environment: input.environment, lease: input.lease, - status: "released", }); return await executePluginEnvironmentCommand({ db, @@ -4051,6 +4081,17 @@ export function environmentRuntimeService( return resolveSandboxDuplexBridgeInput(experimental); }, + async supportsTasks(input: { lease: EnvironmentLease }) { + return options.pluginWorkerManager + ? pluginEnvironmentSupportsTasks(db, options.pluginWorkerManager, input.lease) : false; + }, + + /** Dispatch plugin-provided Runner execution, including resource preparation and task lifecycle. */ + async task(input: { companyId: string; leaseId: string; operation: PluginEnvironmentTaskOperation }) { + if (!options.pluginWorkerManager) throw new Error("Plugin-provided Runner execution worker manager unavailable"); + return executeEnvironmentTask(db, options.pluginWorkerManager, input); + }, + async acquireRunLease(input: { companyId: string; environment: Environment; @@ -4198,7 +4239,15 @@ export function environmentRuntimeService( } continue; } - if (!environment) continue; + if (!environment) { + if (leaseSnapshot.metadata?.driver === "plugin" && options.pluginWorkerManager) { + const receipt = await requireDriverKey("plugin").retryPendingSandboxTeardown!({ environment: null, lease: leaseSnapshot }); + await environmentsSvc.releaseLease(leaseSnapshot.id, status, { cleanupStatus: "success", + remoteExecutionTermination: remoteTerminationReceipt(leaseSnapshot, receipt), + }); + } + continue; + } if ( providerResourceDisposition === "keep_running" && leaseSnapshot.leasePolicy === "reuse_by_environment" diff --git a/server/src/services/environment-task-runtime.ts b/server/src/services/environment-task-runtime.ts new file mode 100644 index 00000000000..a8b080a37b5 --- /dev/null +++ b/server/src/services/environment-task-runtime.ts @@ -0,0 +1,202 @@ +import { settlePluginTaskCredentialHandoff } from "./plugin-task-credential-writeback.js"; +import { and, eq, inArray, or, sql } from "drizzle-orm"; +import { environmentLeases, environments, heartbeatRuns, issues, projects, type Db } from "@paperclipai/db"; +import { environmentTaskOperationSchema, parseEnvironmentTaskResult, type PluginEnvironmentTaskOperation } from "@paperclipai/plugin-sdk"; +import { pluginRegistryService } from "./plugin-registry.js"; +import { hasRemoteTerminationReceipt, remoteExecutionHasStopped } from "./remote-execution-termination.js"; +import type { PluginWorkerManager } from "./plugin-worker-manager.js"; + +export async function pluginEnvironmentSupportsTasks(db: Db, workers: PluginWorkerManager, lease: { companyId: string; metadata: Record | null }) { + const metadata = lease.metadata ?? {}; + if (metadata.driver !== "plugin" || typeof metadata.pluginId !== "string" || typeof metadata.driverKey !== "string") return false; + const plugin = await pluginRegistryService(db).getById(metadata.pluginId); + const driver = plugin?.manifestJson.environmentDrivers?.find(value => value.driverKey === metadata.driverKey); + const methods = workers.getWorker(metadata.pluginId)?.supportedMethods ?? []; + return plugin?.status === "ready" && driver?.supportsTasks === true && + plugin.manifestJson.capabilities.includes("environment.drivers.register") && + methods.includes("environmentTask"); +} + +/** Server-only dispatch for plugin-provided Runner execution. + * The caller supplies an authorized company and a persisted lease, + * never a plugin identity, resource endpoint, or agent/project authority. + * A timeout is ambiguous: retry the same operation against the same lease. + */ +export async function executeEnvironmentTask(db: Db, workers: PluginWorkerManager, input: { + companyId: string; + leaseId: string; + operation: PluginEnvironmentTaskOperation; +}) { + const operation = environmentTaskOperationSchema.parse(input.operation); + const [row] = await db.select({ lease: environmentLeases, environment: environments, run: heartbeatRuns }) + .from(environmentLeases) + .leftJoin(environments, eq(environments.id, environmentLeases.environmentId)) + .leftJoin(heartbeatRuns, and(eq(heartbeatRuns.id, environmentLeases.heartbeatRunId), eq(heartbeatRuns.companyId, environmentLeases.companyId))) + .where(and(eq(environmentLeases.id, input.leaseId), eq(environmentLeases.companyId, input.companyId))).limit(1); + if (!row || !row.lease.providerLeaseId) throw new Error("Plugin-provided Runner execution lease unavailable"); + const { lease, environment, run } = row; + const taskId = row.lease.providerLeaseId; + const metadata = lease.metadata ?? {}; + const probe = metadata.runnerTaskProbe as { runId?: unknown; sessionId?: unknown; expiresAtMs?: unknown } | undefined; + const activeProbe = lease.heartbeatRunId === null && lease.issueId === null && probe + && typeof probe.runId === "string" && typeof probe.sessionId === "string" + && typeof probe.expiresAtMs === "number" && probe.expiresAtMs > Date.now(); + if ((operation.kind === "describe" || operation.kind === "submit" || operation.kind === "connection") && + (metadata.runnerTaskRetiring === true || !environment || (!activeProbe && run?.status !== "running") || lease.status !== "active" || (lease.expiresAt && lease.expiresAt.getTime() <= Date.now()))) throw new Error("Plugin-provided Runner execution lease is not active"); + if (operation.kind === "submit" && (operation.runner.runId !== (activeProbe ? probe.runId : run?.id) || (!operation.resumeFrom && operation.runner.leaseId !== lease.id) + || (activeProbe && (operation.runner.sessionId !== probe.sessionId || operation.projectIds.length > 0)))) { + throw new Error("Plugin-provided Runner execution identity mismatch"); + } + if (operation.kind === "submit" && operation.resumeFrom) { + const sources = await db.select({ runId: environmentLeases.heartbeatRunId }).from(environmentLeases).where(and( + eq(environmentLeases.companyId, input.companyId), eq(environmentLeases.provider, lease.provider!), + eq(environmentLeases.providerLeaseId, operation.resumeFrom), + )).limit(2); + const priorRunId = sources.length === 1 ? sources[0]!.runId : null; + if (activeProbe || !priorRunId || operation.resumeFrom !== await resolvePluginTaskResume(db, { + companyId: input.companyId, leaseId: lease.id, + priorIdentity: { runnerInstanceId: operation.runner.runnerId, environmentLeaseId: operation.runner.leaseId, + runId: priorRunId, normalizedSessionId: operation.runner.sessionId }, + })) throw new Error("Plugin Runner resume identity mismatch"); + } + if (metadata.driver !== "plugin" || typeof metadata.pluginId !== "string" || typeof metadata.driverKey !== "string") { + throw new Error("Plugin-provided Runner execution lease has no pinned plugin driver"); + } + const plugin = await pluginRegistryService(db).getById(metadata.pluginId); + const driver = plugin?.manifestJson.environmentDrivers?.find(value => value.driverKey === metadata.driverKey); + if (!plugin || plugin.status !== "ready" || !driver?.supportsTasks || + !plugin.manifestJson.capabilities.includes("environment.drivers.register") || + !workers.getWorker(plugin.id)?.supportedMethods.includes("environmentTask")) { + throw new Error("Plugin-provided Runner execution provider unavailable"); + } + const issue = operation.kind === "submit" && lease.issueId ? await db.select({ id: issues.id }).from(issues) + .where(and(eq(issues.id, lease.issueId), eq(issues.companyId, input.companyId))).limit(1).then(rows => rows[0]) : null; + if (operation.kind === "submit" && lease.issueId && !issue) throw new Error("Plugin-provided Runner execution issue unavailable"); + const projectIds = operation.kind === "submit" ? operation.projectIds : []; + if (projectIds.length > 0) { + const authorizedProjects = await db.select({ id: projects.id }).from(projects) + .where(and(eq(projects.companyId, input.companyId), inArray(projects.id, projectIds))); + if (authorizedProjects.length !== projectIds.length) throw new Error("Plugin-provided Runner execution project unavailable"); + } + // Provider identity comes from the lease. Editing the environment must never + // redirect status or cleanup to a replacement plugin. + const config = (environment?.config ?? {}) as Record; + const driverConfig = config.pluginKey === plugin.pluginKey && config.driverKey === metadata.driverKey + ? (config.driverConfig as Record | undefined) ?? {} : {}; + if (operation.kind === "submit" && activeProbe) { + const remainingMs = Number(probe.expiresAtMs) - Date.now(); + if (remainingMs <= 0 || !operation.launch) throw new Error("Plugin-provided Runner probe lease is not active"); + // Runner itself must stop if Core disappears during a setup probe. + const requestedLimits = [operation.launch.maxRuntimeMs, operation.launch.maxLifetimeMs] + .filter((value): value is number => typeof value === "number" && value > 0); + operation.launch.maxRuntimeMs = Math.min(remainingMs, ...requestedLimits); + delete operation.launch.maxLifetimeMs; + } + if (operation.kind === "submit") { + // No ticket, headers, environment values, or launch paths belong in metadata. + const marked = await db.update(environmentLeases).set({ + metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || ${JSON.stringify({ runnerTaskLease: true, runnerTaskDispatched: true, runnerTaskExecutionStopped: false, runnerTaskSessionRetained: false, runnerTaskIdentity: operation.runner, runnerTaskResumeFrom: operation.resumeFrom ?? null })}::jsonb`, + }).where(and(eq(environmentLeases.id, lease.id), eq(environmentLeases.companyId, input.companyId), + eq(environmentLeases.status, "active"), sql`coalesce(${environmentLeases.metadata}->>'runnerTaskRetiring', 'false') != 'true'`)).returning({ id: environmentLeases.id }); + if (marked.length !== 1) throw new Error("Plugin-provided Runner execution lease is not active"); + } + try { + if (operation.kind === "complete") { + await settlePluginTaskCredentialHandoff(db, { companyId: input.companyId, leaseId: lease.id }, metadata, async accountIdentity => { + const snapshot = await executeEnvironmentTask(db, workers, { ...input, operation: { kind: "credentials", accountIdentity } }); + if (snapshot.kind !== "credentials") throw new Error("Plugin Runner credential snapshot unavailable"); + return snapshot.credential; + }); + } + const result = await workers.call(plugin.id, "environmentTask", { + driverKey: metadata.driverKey, companyId: input.companyId, environmentId: environment?.id ?? null, + issueId: lease.issueId, config: driverConfig, + taskId, runId: run?.id ?? (activeProbe ? String(probe.runId) : null), agentId: run?.agentId ?? null, projectIds, + lease: { providerLeaseId: lease.providerLeaseId, metadata: lease.metadata ?? undefined }, operation, + }, 15_000); + const parsed = parseEnvironmentTaskResult(operation, taskId, result); + if (parsed.kind === "status" && parsed.executionStopped === true) { + await db.update(environmentLeases).set({ + metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || ${JSON.stringify({ runnerTaskExecutionStopped: true, runnerTaskSessionRetained: parsed.sessionStateAvailable === true })}::jsonb`, + }).where(and(eq(environmentLeases.id, lease.id), eq(environmentLeases.companyId, input.companyId))); + } + if (operation.kind === "complete") { + await db.update(environmentLeases).set({ + metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || '{"runnerTaskDispatched":false}'::jsonb`, + }).where(and(eq(environmentLeases.id, lease.id), eq(environmentLeases.companyId, input.companyId), + sql`${environmentLeases.metadata}->>'runnerTaskExecutionStopped' = 'true'`)); + } + return parsed; + } catch { + // Worker errors and schema diagnostics can contain request credentials. + throw new Error("Plugin-provided Runner execution operation unavailable; reconcile the same task before retrying"); + } +} + +/** Failure leaves the dispatch marker and provider handle available to cleanup retries. */ +export async function settleDispatchedEnvironmentTask(db: Db, workers: PluginWorkerManager, lease: { id: string; companyId: string }) { + const [current] = await db.update(environmentLeases).set({ + metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || '{"runnerTaskRetiring":true}'::jsonb`, + }).where(and(eq(environmentLeases.id, lease.id), eq(environmentLeases.companyId, lease.companyId))) + .returning({ metadata: environmentLeases.metadata }); + if (current?.metadata?.runnerTaskDispatched !== true) { + if (current?.metadata) await settlePluginTaskCredentialHandoff(db, { companyId: lease.companyId, leaseId: lease.id }, current.metadata, + async () => { throw new Error("Undispatched task must not request credentials"); }); + return current?.metadata?.runnerTaskExecutionStopped === true; + } + const scope = { companyId: lease.companyId, leaseId: lease.id }; + try { + await executeEnvironmentTask(db, workers, { ...scope, operation: { kind: "stop" } }); + const state = await executeEnvironmentTask(db, workers, { ...scope, operation: { kind: "status" } }); + if (state.kind !== "status" || state.executionStopped !== true) throw new Error("Plugin Runner termination is unconfirmed"); + await executeEnvironmentTask(db, workers, { ...scope, operation: { kind: "complete" } }); + return true; + } catch { + await db.update(environmentLeases).set({ status: "pending_cleanup", cleanupStatus: "failed", + failureReason: "plugin_runner_task_cleanup_unconfirmed", releasedAt: new Date(), updatedAt: new Date(), + }).where(and(eq(environmentLeases.id, lease.id), eq(environmentLeases.companyId, lease.companyId))); + throw new Error("Plugin Runner termination is unconfirmed; the saved lease requires cleanup"); + } +} + +/** A saved provider task can resume only within the same authorized agent workspace. */ +export async function resolvePluginTaskResume(db: Db, input: { + companyId: string; leaseId: string; + priorIdentity: { runnerInstanceId: string; environmentLeaseId: string; runId: string; normalizedSessionId: string }; +}): Promise { + const rows = await db.select({ lease: environmentLeases, run: heartbeatRuns }).from(environmentLeases) + .innerJoin(heartbeatRuns, and(eq(heartbeatRuns.id, environmentLeases.heartbeatRunId), eq(heartbeatRuns.companyId, input.companyId))) + .where(and(eq(environmentLeases.companyId, input.companyId), + or(eq(environmentLeases.id, input.leaseId), eq(heartbeatRuns.id, input.priorIdentity.runId)))); + const next = rows.find(row => row.lease.id === input.leaseId); + // A retained PRP lease is logical; later provider tasks have different physical lease IDs. + const candidates = rows.filter(row => { + const identity = row.lease.metadata?.runnerTaskIdentity as Record | undefined; + return row.lease.id !== input.leaseId && row.run.id === input.priorIdentity.runId && row.lease.metadata?.runnerTaskSessionRetained === true && + identity?.runnerId === input.priorIdentity.runnerInstanceId && identity.leaseId === input.priorIdentity.environmentLeaseId && + identity.sessionId === input.priorIdentity.normalizedSessionId && identity.runId === input.priorIdentity.runId && + !rows.some(successor => successor.lease.id !== input.leaseId && + successor.lease.metadata?.runnerTaskResumeFrom === row.lease.providerLeaseId); + }); + const old = candidates.length === 1 ? candidates[0] : undefined; + const identity = old?.lease.metadata?.runnerTaskIdentity as Record | undefined; + if (!old || !next || !old.lease.providerLeaseId || !identity || + identity.runnerId !== input.priorIdentity.runnerInstanceId || identity.leaseId !== input.priorIdentity.environmentLeaseId || + identity.sessionId !== input.priorIdentity.normalizedSessionId || identity.runId !== input.priorIdentity.runId || + old.lease.metadata?.driver !== "plugin" || next.lease.metadata?.driver !== "plugin" || + !next.lease.providerLeaseId || next.lease.providerLeaseId === old.lease.providerLeaseId || + (old.run.id !== next.run.id && !["succeeded", "interrupted", "failed", "cancelled", "timed_out"].includes(old.run.status)) || + old.run.agentId !== next.run.agentId || old.lease.issueId !== next.lease.issueId || + old.lease.executionWorkspaceId !== next.lease.executionWorkspaceId || + !old.lease.environmentId || old.lease.environmentId !== next.lease.environmentId || + !next.lease.provider?.startsWith("plugin:") || old.lease.provider !== next.lease.provider || + old.lease.metadata?.pluginId !== next.lease.metadata?.pluginId || old.lease.metadata?.driverKey !== next.lease.metadata?.driverKey || + next.lease.status !== "active" || next.lease.releasedAt || next.run.status !== "running" || + (next.lease.expiresAt && next.lease.expiresAt.getTime() <= Date.now()) || + !(old.run.id === next.run.id + ? rows.filter(row => row.lease.id !== next.lease.id).every(row => hasRemoteTerminationReceipt(row.lease)) + : await remoteExecutionHasStopped(db, input.companyId, old.run.id))) { + throw new Error("Plugin Runner retained session is not authorized"); + } + return old.lease.providerLeaseId; +} diff --git a/server/src/services/plugin-task-credential-writeback.ts b/server/src/services/plugin-task-credential-writeback.ts new file mode 100644 index 00000000000..e13630b5f58 --- /dev/null +++ b/server/src/services/plugin-task-credential-writeback.ts @@ -0,0 +1,143 @@ +import { aiConnectionMetadataSchema } from "@paperclipai/shared"; +import { and, eq, sql } from "drizzle-orm"; +import { companySecrets, connectionGrants, environmentLeases, toolConnections, type Db } from "@paperclipai/db"; +import { z, type PluginEnvironmentTaskCredential } from "@paperclipai/plugin-sdk"; +import { mkdtemp, writeFile, rm, realpath } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { isDeepStrictEqual } from "node:util"; +import { decideCodexAuthMerge, withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server"; +import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge"; +import { aiConnectionService } from "./ai-connections.js"; +import { secretService } from "./secrets.js"; + +const bindingSchema = z.object({ + grantId: z.string().uuid(), + secretId: z.string().uuid(), + sessionEpoch: z.number().int().nonnegative(), + versionSelector: z.union([z.literal("latest"), z.number().int().positive()]), +}).strict(); +export type PluginTaskCredentialWritebackBinding = z.infer; +const handoffSchema = z.object({ + accountIdentity: z.string().min(1), + binding: bindingSchema, + state: z.enum(["pending", "complete"]), +}).strict(); +const key = "runnerTaskCredentialHandoff"; +type Scope = { companyId: string; leaseId: string }; + +/** Save authority, never credential bytes or a controller-local temporary path. */ +export async function registerPluginTaskCredentialHandoff( + db: Db, + scope: Scope, + accountIdentity: string, + binding: PluginTaskCredentialWritebackBinding, +) { + const handoff = handoffSchema.parse({ accountIdentity, binding, state: "pending" }); + await db.transaction(async tx => { + const [lease] = await tx.select().from(environmentLeases).where(and( + eq(environmentLeases.companyId, scope.companyId), eq(environmentLeases.id, scope.leaseId), + )).for("update"); + if (!lease || lease.status !== "active" || lease.metadata?.runnerTaskLease !== true || + lease.metadata.runnerTaskRetiring === true) { + throw new Error("Plugin task credential lease unavailable"); + } + const saved = lease.metadata[key]; + if (saved) { + const prior = handoffSchema.parse(saved); + if (prior.accountIdentity !== accountIdentity || !isDeepStrictEqual(prior.binding, handoff.binding)) { + throw new Error("Plugin task credential authority changed"); + } + return; + } + await tx.update(environmentLeases).set({ metadata: { ...lease.metadata, [key]: handoff } }) + .where(eq(environmentLeases.id, lease.id)); + }); +} + +/** Runs for live completion and the cleanup sweep, including after Core restarts. */ +export async function settlePluginTaskCredentialHandoff( + db: Db, + scope: Scope, + metadata: Record, + fetchCredential: (accountIdentity: string) => Promise, +) { + if (!metadata[key]) return; + const handoff = handoffSchema.parse(metadata[key]); + if (handoff.state === "complete") return; + if (metadata.runnerTaskDispatched === true && metadata.runnerTaskExecutionStopped !== true) { + throw new Error("Plugin task credential process termination is unconfirmed"); + } + const binding = handoff.binding; + const home = await mkdtemp(join(await realpath(tmpdir()), "paperclip-task-credential-")); + try { + // Revocation or an explicit account replacement intentionally discards an old refresh. + const account = async (database: Db) => { + const [grant] = await database.select().from(connectionGrants).where(and( + eq(connectionGrants.companyId, scope.companyId), eq(connectionGrants.id, binding.grantId), + )).for("update"); + if (!grant || grant.status !== "active") return null; + const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential"); + if (!ref || ref.secretId !== binding.secretId || + (ref.versionSelector ?? "latest") !== binding.versionSelector) return null; + const [secret] = await database.select().from(companySecrets).where(and( + eq(companySecrets.companyId, scope.companyId), eq(companySecrets.id, binding.secretId), + )).for("update"); + if (!secret || secret.aiSessionEpoch !== binding.sessionEpoch) return null; + const [connection] = await database.select().from(toolConnections).where(and( + eq(toolConnections.companyId, scope.companyId), eq(toolConnections.id, grant.connectionId), + )); + const ai = aiConnectionMetadataSchema.safeParse(connection?.config.ai); + if (!connection || !ai.success || ai.data.provider !== "openai" || ai.data.method !== "subscription") return null; + return { connection, grant }; + }; + const mutate = (operation: (database: Db) => Promise) => withAccountHomeSecretMutationLock( + undefined, scope.companyId, () => db.transaction(tx => operation(tx as unknown as Db)), + ); + const current = metadata.runnerTaskDispatched === true && await mutate(account); + if (current) { + // Do not hold database or account locks while the provider exports its credential. + const credential = await fetchCredential(handoff.accountIdentity); + if (credential !== null) { + if (credential.type !== "codex-auth-json") throw new Error("Plugin task credential provider mismatch"); + const incoming = join(home, "incoming.json"), destination = join(home, "current.json"); + await writeFile(incoming, credential.authJson, { mode: 0o600 }); + const writeBack = () => mutate(async database => { + const selected = await account(database); + if (!selected) return; + const value = await aiConnectionService(database).credential(selected); + await writeFile(destination, value, { mode: 0o600 }); + if (await decideCodexAuthMerge(incoming, destination, { errorLabel: "AI account refresh" }) !== 10) return; + await secretService(database).rotate(binding.secretId, { + value: credential.authJson, preserveAiSessionEpoch: true, + }, { userId: selected.grant.subjectUserId }); + }); + for (let attempt = 0; ; attempt++) { + try { + await writeBack(); + break; + } catch (error) { + if (attempt >= 2 || (error as NodeJS.ErrnoException)?.code !== WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE) throw error; + } + } + } + } + const completed = await db.update(environmentLeases).set({ + metadata: sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) || ${JSON.stringify({ [key]: { ...handoff, state: "complete" } })}::jsonb`, + }).where(and( + eq(environmentLeases.companyId, scope.companyId), eq(environmentLeases.id, scope.leaseId), + sql`${environmentLeases.metadata}->${key} = ${JSON.stringify(handoff)}::jsonb`, + )).returning({ id: environmentLeases.id }); + if (completed.length !== 1) { + const [current] = await db.select({ metadata: environmentLeases.metadata }).from(environmentLeases).where(and( + eq(environmentLeases.companyId, scope.companyId), eq(environmentLeases.id, scope.leaseId), + )); + if (!isDeepStrictEqual(handoffSchema.safeParse(current?.metadata?.[key]).data, { ...handoff, state: "complete" })) { + throw new Error("Plugin task credential handoff changed"); + } + } + } finally { + // The provider retains the recoverable credential until completion succeeds. + await rm(home, { recursive: true, force: true }); + } +}