Support Codex warm plugin task reuse and recovery (#17390)

server/src/services/native-runtime/plugin-runner-task.ts

@@ -24,17 +24,20 @@ export function createPluginRunnerTaskOwner(input: {
   onSessionRetained?: (state: Record<string, unknown>) => Promise<void>;
   onDispatched?: () => Promise<void> | void;
   pollMs?: number;
+  warmTaskVersion?: 1;
+  adoptIdentity?: DurableRecoveryIdentity;
 }) {
   const pollMs = input.pollMs ?? 500;
   let registrationGeneration = 0;
   let activeOutbound: ReturnType<typeof connectRunnerPrpIngress> | null = null;
   let preparedLaunchId: string | undefined;
-  let preparedIdentity: DurableRecoveryIdentity | undefined;
-  let launched = false;
+  let preparedIdentity: DurableRecoveryIdentity | undefined = input.adoptIdentity;
+  let launched = !!input.adoptIdentity;
   let stopRequested = false;
   let released = false;
-  let submitted = false;
+  let submitted = !!input.adoptIdentity;
   let detached = false;
+  let adoptionVerified = false;
   let markReady!: () => void;
   let failReady!: (error: unknown) => void;
   let confirmStopped: (code: number) => void = () => {};
@@ -44,6 +47,7 @@ export function createPluginRunnerTaskOwner(input: {
   void failure.catch(() => undefined);
   const running = new Promise<void>((resolve, reject) => { markReady = resolve; failReady = reject; });
   void running.catch(() => undefined);
+  if (input.adoptIdentity) markReady();
   const call = async (operation: PluginEnvironmentTaskOperation) => {
     const phase = ({ prepare: "execution.plugin_prepare", check_project_auth: "execution.plugin_project_auth",
       submit: "execution.plugin_submit", connection: "execution.plugin_connection" } as const)[operation.kind as "prepare" | "check_project_auth" | "submit" | "connection"];
@@ -54,23 +58,32 @@ export function createPluginRunnerTaskOwner(input: {
     }
     return result;
   };
+  const checkProjectAuth = async (configurationId:string) => {
+    if (input.projectAuthCheckVersion !== 1) throw new Error("Plugin Runner managed project authentication validation unavailable");
+    const provider = input.harness === "codex" ? "openai" : input.harness === "claude" ? "anthropic" : input.harness;
+    const validation = await call({kind:"check_project_auth",configurationId,provider,projectIds:[...input.projectIds]});
+    if (released || stopRequested) throw new Error("Plugin Runner preparation cancelled");
+    if (validation.kind !== "project_auth" || !validation.allowed || validation.configurationId !== configurationId
+      || validation.provider !== provider) throw new Error("Plugin Runner project authentication conflicts with the selected managed account");
+  };
   const prepareLaunch: RunnerLaunchPreparer = async (requirements, recovery) => {
     if (released || stopRequested) throw new Error("Plugin Runner preparation cancelled");
+    // Reconnecting the previous authority must not prepare it through the new run's lease.
+    if (recovery && input.adoptIdentity && Object.entries(input.adoptIdentity).every(([key,value]) => requirements.identity[key as keyof DurableRecoveryIdentity] === value)) {
+      await checkProjectAuth(recovery.configurationId);
+      preparedLaunchId = recovery.configurationId;
+      return recovery;
+    }
     const result = await call(environmentTaskOperationSchema.parse({ kind: "prepare", requirements,
       projectIds: [...input.projectIds], ...(recovery ? { recovery } : {}),
       ...(input.resumeFrom ? { resumeFrom: input.resumeFrom } : {}),
+      ...(launched && preparedIdentity && requirements.identity.runId !== preparedIdentity.runId
+        ? { priorIdentity: preparedIdentity, accountIdentity: input.accountIdentity } : {}),
     }));
     if (released || stopRequested) throw new Error("Plugin Runner preparation cancelled");
     if (result.kind !== "prepared") throw new Error("Plugin Runner preparation unavailable");
     if (input.projectAuthCheckVersion !== 1) throw new Error("Plugin Runner managed project authentication validation unavailable");
-    if (!result.projectAuthCheckAtSubmit) {
-      const provider = input.harness === "codex" ? "openai" : input.harness === "claude" ? "anthropic" : input.harness;
-      const validation = await call({ kind: "check_project_auth", configurationId: result.launch.configurationId,
-        provider, projectIds: [...input.projectIds] });
-      if (released || stopRequested) throw new Error("Plugin Runner preparation cancelled");
-      if (validation.kind !== "project_auth" || !validation.allowed || validation.configurationId !== result.launch.configurationId
-        || validation.provider !== provider) throw new Error("Plugin Runner project authentication conflicts with the selected managed account");
-    }
+    if (launched || !result.projectAuthCheckAtSubmit) await checkProjectAuth(result.launch.configurationId);
     preparedLaunchId = result.launch.configurationId;
     preparedIdentity = requirements.identity;
     return result.launch;
@@ -115,7 +128,7 @@ export function createPluginRunnerTaskOwner(input: {
     if (launched || released) throw new Error("Plugin Runner replacement requires a new task lease");
     const native = spec.nativeRunner;
     if (!preparedLaunchId) throw new Error("Plugin Runner launch requires preparation");
-    if (!native || native.lifecyclePolicy?.mode !== "per_turn"
+    if (!native || (native.lifecyclePolicy?.mode !== "per_turn" && !(native.lifecyclePolicy?.mode === "warm" && input.harness === "codex" && input.accountIdentity && input.warmTaskVersion === 1))
       || !preparedIdentity || Object.entries(preparedIdentity).some(([key, value]) => native.identity[key as keyof DurableRecoveryIdentity] !== value)) {
       throw new Error("Plugin Runner requires a structured per-turn native launch");
     }
@@ -138,7 +151,11 @@ export function createPluginRunnerTaskOwner(input: {
     });
     launched = true;
     const child: RunnerProcessHandle["child"] = { exitCode: null, signalCode: null,
-      kill: () => { stopRequested = true; return true; } };
+      kill: () => {
+        stopRequested = true;
+        if (native.lifecyclePolicy?.mode === "warm" && submitted) void stopAndConfirm().catch(rejectFailure);
+        return true;
+      } };
     const completion = new Promise<Awaited<RunnerProcessHandle["completion"]>>(resolve => {
       confirmStopped = code => {
         if (child.exitCode !== null) return;
@@ -172,7 +189,11 @@ export function createPluginRunnerTaskOwner(input: {
         if (detached) return;
         if (stopRequested && !stopSent) { await call({ kind: "stop" }); stopSent = true; }
         const state = await readStatus();
-        if (state.phase === "running" && !stopRequested) markReady();
+        if (state.phase === "running" && !stopRequested) {
+          markReady();
+          // PRP owns live failure detection; warm idle tasks must not keep polling Launcher.
+          if (native.lifecyclePolicy?.mode === "warm") return;
+        }
         if (state.executionStopped === true) return;
         await delay(pollMs);
       }
@@ -251,6 +272,16 @@ export function createPluginRunnerTaskOwner(input: {
       return result.state;
     },
     controlPlaneRegistration: (authority, identity, signal) => registration(authority, identity?.runId ?? "", signal),
+    ...(input.adoptIdentity ? { adoptExistingRunner: {
+      // Provider ownership is attested by Launcher, never by a controller-local PID.
+      pid: 0, processGroupId: null, startedAt: new Date().toISOString(),
+      isAlive: async () => {
+        if (!adoptionVerified) { if ((await readStatus()).executionStopped === true) return false; adoptionVerified = true; }
+        // Authenticated PRP detects disconnects; this cache never asserts termination.
+        return stoppedState === undefined;
+      },
+      signal: async () => { stopRequested = true; await stopAndConfirm(); return true; },
+    } } : {}),
   };
   const readWorkspaceFile = async (path: string, maxBytes: number): Promise<Buffer> => {
     if (!submitted || !preparedIdentity || released || stopRequested) throw new Error("Plugin Runner file access unavailable");
@@ -262,5 +293,14 @@ export function createPluginRunnerTaskOwner(input: {
         createHash("sha256").update(bytes).digest("hex") !== result.sha256) throw new Error("Plugin Runner file integrity check failed");
     return bytes;
   };
-  return { launcher, registration, driver, readWorkspaceFile };
+  return { launcher, registration, driver, readWorkspaceFile,
+    rebind(next: typeof input) {
+      if (!launched || stopRequested || stoppedState || detached || next.harness !== "codex" ||
+          input.accountIdentity !== next.accountIdentity || input.credentialIdentity !== next.credentialIdentity ||
+          JSON.stringify(input.projectIds) !== JSON.stringify(next.projectIds)) throw new Error("Plugin warm task cannot be rebound");
+      input = next;
+      released = false;
+    },
+  };
+
 }