Recover setup after environment and policy changes

Resume preparation after a stop completes or a policy hold clears. Keep probes blocked when the selected environment could not be saved, preserve the account, and offer a save retry.

Verified 64 lifecycle/policy tests, the focused policy recovery case, 110 wizard/helper tests, and server/UI TypeScript checks.

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/modules/agent-lifecycle/__tests__/lifecycle.test.ts

@@ -171,8 +171,8 @@ const support = await getEmbeddedPostgresTestSupport();
     expect(await current(id)).toMatchObject({ lifecycleState: "paused", lifecycleHolds: ["setup", "budget"] });
     await lifecycle.updateConfiguration(id, { adapterConfig: { provider: "codex", model: "changed" } });
     expect((await current(id)).lifecycleState).toBe("paused");
-    // Restore only the setup hold to exercise environment changes under the wizard's authority.
-    await db.update(agents).set({ lifecycleHolds: ["setup"], pauseReason: "setup" }).where(eq(agents.id, id));
+    await store.change(id, "reconcile"); await work.process(id);
+    expect(await current(id)).toMatchObject({ lifecycleState: "verifying", lifecycleHolds: ["setup"] });
     await lifecycle.updateConfiguration(id, { adapterConfig: { provider: "codex", model: "prepared" } });
     await work.process(id);
     expect((await current(id)).lifecycleState).toBe("verifying");
@@ -181,6 +181,11 @@ const support = await getEmbeddedPostgresTestSupport();
     expect((await current(id)).lifecycleState).toBe("paused");
     await lifecycle.updateConfiguration(id, { defaultEnvironmentId: environmentId }); await work.process(id);
     expect((await current(id)).lifecycleState).toBe("verifying");
+    await lifecycle.updateConfiguration(id, { defaultEnvironmentId: local.id });
+    expect((await current(id)).lifecycleState).toBe("pausing");
+    await lifecycle.updateConfiguration(id, { defaultEnvironmentId: environmentId });
+    await work.process(id);
+    expect((await current(id)).lifecycleState).toBe("verifying");
     const finished = await lifecycle.requestHire(companyId, input, { setupAgentId: id, runnerResolved: true });
     expect(finished.id).toBe(id);
     await work.process(id);

server/src/modules/agent-lifecycle/adapters/postgres.ts

@@ -62,6 +62,8 @@ async function changeLockedAgent(tx: Db, effects: LifecycleEffects, agent: Lifec
   if (restartPause && next === "paused") next = "pausing";
   if (command === "complete" && next === "ready" && holds.length) next = "pausing";
   if (command === "complete" && next === "paused" && !holds.length) next = "resuming";
+  if (["complete", "reconcile"].includes(command) && next === "paused" && holds.length === 1 && holds[0] === "setup"
+    && await effects.setupRequiresPreparation(tx, agent.companyId, agent.defaultEnvironmentId)) next = "preparing";
   if (next === agent.lifecycleState) {
     if (JSON.stringify(holds) !== JSON.stringify(agent.lifecycleHolds) || agent.pauseReason !== (holds[0] ?? null)) {
       const [updated] = await tx.update(agents).set({ lifecycleHolds: holds, pauseReason: holds[0] ?? null }).where(eq(agents.id, agent.id)).returning();

ui/src/components/new-agent/NewAgentSetup.tsx

@@ -413,7 +413,9 @@ function Setup({
   );
   const needsCurrentSetup = ["codex_local", "claude_local", "opencode_local", "pi_local"].includes(adapterType);
   const preparingComputer = created?.pauseReason === "setup" && ["preparing", "pausing"].includes(created.lifecycleState ?? "");
-  const ready = Boolean(created) && !environmentSaving && !preparingComputer && connectionSettingsReady && (!needsCurrentSetup || (
+  const environmentSelectionSaved = (created?.defaultEnvironmentId ?? null)
+    === (environmentOverride || (forced.forced || managedOnly ? environmentId : null));
+  const ready = Boolean(created) && environmentSelectionSaved && !environmentSaving && !preparingComputer && connectionSettingsReady && (!needsCurrentSetup || (
     envs.isSuccess && !envs.isFetching
     && settings.isSuccess && !settings.isFetching
     && experimental.isSuccess && !experimental.isFetching
@@ -657,6 +659,15 @@ function Setup({
     finally { savingRef.current = false; setSaving(false); }
   }
 
+  async function changeEnvironment(selected = environmentOverride) {
+    setEnvironmentOverride(selected);
+    resetTest();
+    setEnvironmentSaving(true);
+    try { await saveDraft(connection, false, undefined, selected); }
+    catch (cause) { setError(cause instanceof Error ? cause.message : "Could not save the environment."); }
+    finally { setEnvironmentSaving(false); }
+  }
+
   async function runTest(nextConnection = connection): Promise<Record<string, unknown> | null> {
     if (!ready) return null;
     const run = ++generation.current;
@@ -1004,7 +1015,7 @@ function Setup({
                           ) ??
                           result?.checks.find(
                             (check) =>
-                              check.code.includes("hello_probe") &&
+                              (check.code.includes("hello_probe") || check.code.startsWith("environment_agent_")) &&
                               check.level === "warn",
                           )
                         )?.message
@@ -1390,13 +1401,7 @@ function Setup({
                             value={environmentOverride}
                             disabled={forced.forced}
                             onChange={(event) => {
-                              const selected = event.target.value;
-                              setEnvironmentOverride(selected);
-                              resetTest();
-                              setEnvironmentSaving(true);
-                              void saveDraft(connection, false, undefined, selected)
-                                .catch(cause => setError(cause instanceof Error ? cause.message : "Could not save the environment."))
-                                .finally(() => setEnvironmentSaving(false));
+                              void changeEnvironment(event.target.value);
                             }}
                           >
                             <option value="">
@@ -1427,6 +1432,7 @@ function Setup({
                         {error}
                       </p>
                     )}
+                    {error && !environmentSelectionSaved && <Button disabled={busy} onClick={() => void changeEnvironment()}>Retry saving environment</Button>}
                     <div className="flex justify-between border-t border-border py-5">
                       <div className="flex items-center gap-3">
                         <Button type="button" variant="ghost" disabled={busy} onClick={() => void saveAndExit()}>Save & exit</Button>

ui/src/pages/NewAgent.test.tsx

@@ -10,6 +10,32 @@ import {
 } from "./NewAgent.test-support";
 
 describe("New agent setup", () => {
+  it("keeps the account on an environment save failure and blocks probes until the selection is saved", async () => {
+    envApi.list.mockResolvedValue([
+      { id: "local-1", name: "Local", status: "active", driver: "local", config: {} },
+      { id: "computer", name: "Agent computer", status: "active", driver: "plugin", config: {} },
+    ]);
+    await render("codex_local");
+    await click("OpenAIAPI key");
+    await fill("API key", "setup-test-key");
+    await click("Connect");
+    const originalBinding = api.testEnvironment.mock.calls[0][2].aiConnection;
+    const probes = api.testEnvironment.mock.calls.length;
+    api.update.mockRejectedValueOnce(new Error("Environment save unavailable"));
+    const select = container.querySelector<HTMLSelectElement>('[aria-label="Environment"]')!;
+    await act(async () => { select.value = "computer"; select.dispatchEvent(new Event("change", { bubbles: true })); });
+    await settle();
+    expect(container.textContent).toContain("Configure your agent");
+    expect(container.textContent).toContain("Environment save unavailable");
+    await click("Run test");
+    expect(api.testEnvironment).toHaveBeenCalledTimes(probes);
+    await click("Retry saving environment");
+    await click("Run test");
+    expect(api.testEnvironment).toHaveBeenLastCalledWith("company-1", expect.any(String), expect.objectContaining({
+      environmentId: "computer", aiConnection: originalBinding,
+    }));
+  });
+
   it("waits for background preparation and tests the same saved computer after it completes", async () => {
     const id = "00000000-0000-4000-8000-000000000092";
     const draft = { id, companyId: "company-1", name: "Atlas", adapterType: "paperclip_runner", adapterConfig: { provider: "codex" },