feat(agent-sdk): run assigned tasks in pre-warmed sandboxes

[captainsafia]

Description

A pre-warmed standby sandbox can now receive a run from the server, install it in place, and run it. Before this change, a standby sandbox prepared, heartbeated, and exited at its deadline, so no run ever used it. This is the client side of the server's assignment contract, where the server commits a run to a sandbox before it tells the sandbox.

The protocol has no accept or reject step. The sandbox receives assignment or retired on a control socket and sends only heartbeat. After it installs the run's credential, it reconnects with that credential. The server admits the connection, revokes the standby credential, and replays the assignment; the replay confirms delivery. Only then does the run start, so a cancelled or superseded run never starts.

Commits

Each commit is one review unit, in reading order. The branch is linear on master at #16538 (ae60375ef). Only the branch tip is validated; intermediate commits do not all build.

#	Commit	Changes
1	52b0c69 feat(terminal): stop a local shell's event loop for replacement	Adds Message::ShutdownForReplacement, so the PTY event loop can stop a shell without ending the terminal.
2	09a451a feat(terminal): replace an idle local shell with a new environment	Adds TerminalView::replace_shell, which restarts the shell behind a terminal with a new directory and environment. It accepts only an idle, writable local shell, and its error names the condition that refused the replacement. The pane routes the request to the local terminal manager of that view, which starts the replacement on the cloud shell-recovery resources. The replacement environment becomes the baseline for later shell recoveries. Planned replacement is not supported on Windows, because pre-warmed sandboxes run only on Linux.
3	d03c91b feat(mcp): restart templatable MCP servers with a new runtime environment	Adds a runtime environment that overrides the environment of spawned CLI MCP servers without being persisted. replace_runtime_environment stops every server and returns the ones to restart. The runtime environment is native-only.
4	830a748 feat(server-api): talk to pre-warmed sandboxes over the control socket	Replaces the HTTP heartbeat with the /ws control socket and its assignment, retired, and heartbeat messages. 410 Gone maps to "unavailable". A message type the client does not know decodes as Unknown, so the server can add types without breaking a waiting sandbox. The decode test uses the server's own assignment fixture. The configuration query keeps #16538's fingerprint/spec contract.
5	cecdda2 refactor(agent-sdk): resolve agent driver options in one place	split_run_options is the only destructure of AgentDriverOptions. create and the pre-warmed path both apply the result with apply_run_configuration, so a new option cannot reach one path and miss the other.
6	ab51390 refactor(agent-sdk): let launches without CLI arguments reuse run setup	Execution-config parsing, Git credential bootstrap, team metadata refresh, the Factory-repository decision, and Bedrock OIDC setup now accept a launch with no RunAgentArgs, so the assignment path calls them instead of copying them.
7	63f521e feat(agent-sdk): let a run reuse a workspace prepared in standby	WorkspaceState records whether the workspace is pending or prepared. A prepared workspace skips cloning, setup commands, and the file-based MCP scan. same_preparation lets the assignment path refuse a run that asks for a different workspace.
8	fbcc0c5 feat(agent-sdk): install a pre-warmed assignment on the standby driver	prepare_prewarmed_assignment installs the run credential first, then installs the run: options, secrets, attachments, cloud providers, Bedrock, the shell and MCP environment, and the task binding. MCP servers get only the assignment's runtime environment, which is what a worker gives them in a normal run, and never the terminal's managed secrets. The team metadata refresh uses the same retry as a normal cloud run. An installation failure is reported against the run, because the sandbox cannot hand the run back.
9	133f95b feat(agent-sdk): run assigned tasks from a pre-warmed standby sandbox	run_lifecycle waits in standby for the assignment, installs it, confirms delivery, and runs it. It holds an assignment that arrives before preparation finishes, ignores replays of the same assignment, and ignores unknown messages on the same connection.
10	96319d8 refactor(agent-sdk): identify a pre-warmed assignment by its execution ID	The standby driver dedupes and confirms replayed assignments on execution_config.execution_id, and the client no longer reads an assignment ID. A sandbox gets at most one assignment and an execution gets at most one sandbox, so the separate ID was redundant.
11	0298a80 fix(server-api): skip non-text frames on the pre-warmed control socket	Review fix. tungstenite gives server pings to the reader, so a ping dropped the socket and the sandbox reconnected about once a minute. The reader now skips frames that are not text.
12	66af38f fix(agent-sdk): continue a pre-warmed run when MCP servers fail to restart	Review fix. The MCP restart result now goes through handle_mcp_startup_result, so a server that fails to restart degrades the run as it does in a cold run, unless strict MCP startup is set.
13	cbf6838 test(terminal): run the planned replacement baseline test only on Unix	CI fix. The baseline test used the app context only on Unix, so Windows Clippy rejected the unused closure parameter. Windows refuses planned shell replacement, so the test now runs only on Unix.

Decisions and behavior changes

• The run waits for delivery confirmation. Starting immediately and confirming in the background would start sooner, but it could run a task the server already cancelled. Confirmation retries transient failures for 5 minutes, then fails, so an unreachable server does not hold a run until its deadline.
• Cloud providers are accepted. The server assigns runs that use them, and refusing an assignment now fails the run outright.
• Secrets for agent run (commit 5). A typed auth secret is now skipped as a whole when one of its variables is set in environment_variables. Before, only the colliding variable was replaced.
• No planned shell replacement on Windows. Pre-warmed sandboxes run only on Linux, so the Windows child exit watcher stays unchanged.
• No standby shutdown report. The server learns that a sandbox stopped from its Namespace instance callback, and it has no sandbox shutdown endpoint. warp-agent-docker#281 calls harness-support report-shutdown --standby-id, which this PR does not add.

Review guide

Start with commits 5 and 8; together they hold most of the risk. In apply_run_configuration, check which runtime state survives an in-place assignment: the harness, the output format, and a prepared workspace. In prepare_prewarmed_assignment, check that nothing can fail before the run credential is installed, apart from parsing the task ID. In commit 2, the replacement refusal conditions in TerminalView::replace_shell decide when an assignment can install.

Linked Issue

Part of REMOTE-3002. This PR does not add the typed assignment context or the standby telemetry correlation that the issue describes. Plans: Align the pre-warmed sandbox client with the new server assignment contract · Address the automated review of PR #16539 · Trim the PR #16539 diff · Conversation

• The linked issue is labeled ready-to-spec or ready-to-implement. The work is tracked in Linear, not in a GitHub issue.
• Where appropriate, screenshots or a short video of the implementation are included below (especially for user-visible or UI changes). No user-visible UI changed.

Testing

On the branch tip:

• cargo nextest run -p warp --lib -E 'test(prewarmed) | test(agent_sdk::) | test(control_reader)': 1018 passed.
• cargo clippy --locked -p warp --all-targets --tests -- -D warnings: passed.
• CI's wasm Clippy command, cargo clippy --locked --target wasm32-unknown-unknown --profile release-wasm-debug_assertions -- -D warnings: passed.
• ./script/format: no changes.

The wider filter that adds test(shell_replacement) | test(shell_recovery) | test(terminal_manager) | test(templatable_manager) | test(local_tty::event_loop_tests) passed 1046 tests on the previous tip. The terminal and MCP manager code has not changed since then.

The standby tests cover an early assignment, replays, unknown messages, non-text frames, retirement before and after installation, Gone and transient failures during confirmation, and both deadlines. A mock-server test checks that an installation failure is reported against the run with the run's own credential. The standby bootstrap test runs oz agent run-cloud-standby against a mock server through readiness.

• I have manually tested my changes locally with ./script/run. No live Namespace assignment ran, because the server side of the socket contract is not deployed.

Agent Mode

• Warp Agent Mode - This PR was created via Warp's AI Agent Mode

CHANGELOG-NONE

Co-Authored-By: Warp agent@warp.dev

[warp-factories]

@captainsafia, Warp CI failed on this PR's current head. May windows runner test diagnose it, commit fixes, and push to this same PR branch? Your consent also opts this PR into automatic handling of future CI failures. Please reply to this request mentioning @warp with explicit permission. We will not change the branch without it.

Responding as windows runner test: Open session · View in factory

[warp-for-oss]

@captainsafia

I'm starting a first review of this pull request.

You can view the conversation on Warp.

I completed the review and no human review was requested for this pull request.

Comment /warp-agent-review on this pull request to retrigger a review (up to 3 times on the same pull request).

Powered by Oz

[warp-for-oss]

Overview

This PR adds the client-side pre-warmed sandbox assignment flow, including control-socket assignment handling, in-place driver/shell/MCP environment replacement, and supporting terminal/standby tests.

Concerns

• The planned shell replacement failure path in TerminalManager still invokes the old cloud-shell recovery failure callback before branching to the new replacement callback, which consumes the anyhow::Error and prevents the file from compiling.

Verdict

Found: 1 critical, 0 important, 0 suggestions

Request changes

Comment /warp-agent-review on this pull request to retrigger a review (up to 3 times on the same pull request).

Powered by Oz

[vorporeal]

read through all of the commits, and had an agent read the full PR, leave some comments from its own findings, and leave some more from me asking it questions as i read.

approving to unblock, but i think there's a bunch of stuff in here that's worth doing (whether it's in this PR or in a follow-up is up to you).

[vorporeal]

pending review comment

[vorporeal]

oops created this so there was a pending review that the agent could attach additional comments to and then i forgot to delete it

[vorporeal]

i assume this is more of a TODO than a "we cannot support this"? (otherwise, we're saying that we can't do pre-warmed sandboxes on windows, right?)

[vorporeal]

this isn't associated with any sort of "request", right? should this be planned_replacement_id?

[vorporeal]

might be worth a quick rename of this, given we're now using it for more than just recovery. perhaps PendingShellReplacement? (or PendingShellRespawn?)

[vorporeal]

[agent] The server pings this socket every 54s (socketPingPeriod in warp-server logic/ai/ambient_agents/prewarmed/socket.go, gorilla/websocket). tungstenite 0.24 queues the pong and also returns Message::Ping to the stream consumer, so this text() check turns every server ping into an error. maintain_standby_connection then drops the socket and reconnects after 1s. So every standby sandbox tears down and reopens its socket about once a minute for its whole life, re-running auth each time, with nothing logged.

It also stretches the heartbeat cadence to ~55s, because the 30s heartbeat timer restarts on each new connection. The server only treats a sandbox as an assignment candidate if it heartbeated within 90s (candidateHeartbeatWindow in assignment.go), so a single missed heartbeat at that cadence makes the sandbox temporarily unassignable.

Delivery confirmation is not affected, since the server replays the assignment immediately on connect, well before the first ping.

Could this skip non-text frames instead, e.g. loop with let Some(text) = message.text() else { continue };, the way the shared-session code does in app/src/terminal/shared_session/sharer/network.rs:1521? A Close frame would then end the stream normally on the next read. A test with a fake stream that yields a Ping and then a Text message would cover it.

[vorporeal]

[agent] FYI, not blocking: this branch does no refresh after the assignment arrives (no fetch/checkout of the prepared repos, no Factory definition checkout, and global skills are only resolved, not cloned), so a pre-warmed run works on branch heads as of standby preparation. I understand a post-assignment refresh is planned; two things to keep in mind when it lands:

• Global skill repos are separate checkouts from the workspace repos, so the refresh needs to cover them too, or a global skill repo added/updated after standby will be missing or stale.
• Repos can carry expected_sha, and same_preparation compares clone_requests, which include it. If the server ever sends an assignment whose pinned SHA differs from what standby prepared, the run (already committed) fails at same_preparation. A refresh that checks out the expected SHA, with same_preparation ignoring the SHA, would be the more forgiving behavior.

[vorporeal]

[agent] Question about the overall failure model, since cloud agent reliability is a top priority right now: because the sandbox can't hand a committed run back, every failure between the assignment and the run actually starting turns into a failed run, even though nothing has executed yet. That includes installation failures here (a retained MCP server failing to restart, same_preparation rejecting the workspace, the shell replacement being refused or timing out, secrets or Bedrock setup failing), plus the post-install exits in run_lifecycle (deadline during installation, delivery not confirmed within the window). In all of those cases a fresh cold sandbox would very likely have succeeded.

Could these pre-start failures fall back to a cold start instead of failing the run? This doesn't need an accept/decline step in the protocol. Roughly:

• the sandbox reports "installation failed before the run started" as a distinct outcome (e.g. a specific status or error class on the run, or a socket/HTTP call), using the run key it already holds;
• the server treats that as "release this execution back to pending and retire the sandbox", then dispatches the run cold;
• retries are capped (e.g. at most one pre-warmed attempt per execution, then cold only) so a bad config can't loop.

The task never reaching IN_PROGRESS before these failures seems like a reasonable server-side signal that requeuing is safe. Is this something you've considered, or planned as a follow-up? It would need matching server work in warp-server's prewarmed assignment code.

[vorporeal]

nit: you can leverage warpui::r#async::BoxFuture to simplify the types here.

[vorporeal]

should should_share be overwritten at (or at least closer to) the source based on this ngrok check? helps ensure that all consumers get it right.

[vorporeal]

[agent] Installation currently runs every step one after another, but most steps don't depend on each other. Since installation latency is the whole point of pre-warming, it seems worth restructuring this into a few clearly ordered groups, where each group runs in parallel and steps are serialized only where a real dependency exists.

Dependencies (after the run credential and task ID are installed, which should stay first):

Step	Kind	Needs	Needed by
Validate payload; derive team scope, share requests, skill specs, Bedrock role/region, runtime_env (+ key, deadline), Factory contract	pure/sync	payload	everything
Fetch task secrets	network	task ID + key	shell (via build_secret_env_vars)
Fetch + download attachments (incl. handoff)	network	task ID + working_dir	task prompt only
Bootstrap Git credentials	network + git config write	task ID	the run itself (nothing in installation)
Refresh team metadata (up to 30s of retries)	network	key	building options (conservatively), bind, run
Mint Bedrock OIDC credentials	network	task ID + team scope (both known from the payload)	the run
Resolve configured skills	local only (clone_repositories = false)	prepared checkouts	building options
Build options/task, same_preparation, cloud providers, split_run_options + apply_run_configuration	sync, main thread	skills, secrets, attachments, Bedrock config	shell, bind
Replace shell	up to 75s	resolved_env_vars (depends on secrets)	sharing
Restart retained MCP servers (stop up to 20s, start up to 20s)	process spawn	only runtime_env	nothing in installation
Bind task + start sharing + wait for share	network, up to 20s	team scope; should follow the shell replacement	the run

The notable ones:

• MCP depends only on runtime_env, which is known right after parsing the payload. Its one small link is wait_for_mcp_servers_started reading me.mcp_startup_timeout, which is the default either way and could be passed explicitly.
• All network steps need only the credential. Attachments and Bedrock currently write into options/task as they go; they'd need to return results instead.
• Sharing should follow the shell replacement, but waiting for the share doesn't need to block the MCP restart.
• The cold path already does this kind of thing: fetch_secrets_and_attachments fetches secrets, attachments, metadata, and the handoff snapshot with one futures::join!.

Suggested shape:

async fn install_assignment(...) -> Result<Task, AgentDriverError> {
    // 0. Pure: validate and derive everything from the payload. No I/O.
    let plan = AssignmentPlan::from_payload(config, runtime_env, api_key, expires_at)?;

    // The one step with no inputs beyond the payload, and the slowest: start it now, join it last.
    let mcp_restart = Self::replace_prewarmed_mcp_environment(spawner, plan.mcp_environment.clone());

    // 1. Network: everything that only needs the run credential.
    let (secrets, attachments, (), (), bedrock) = futures::try_join!(
        fetch_task_secrets(spawner, task_id),
        fetch_and_download_attachments(ai_client, server_api, task_id, &attachments_dir),
        AgentDriverRunner::bootstrap_git_credentials_for_task(spawner, ..),
        AgentDriverRunner::refresh_team_metadata(spawner, true),
        mint_bedrock_credentials(spawner, task_id, &plan.team_scope, plan.bedrock),
    )?;

    // 2. Sync, on the main thread: resolve skills (local), build options/task, check
    //    same_preparation, load cloud providers, attach secrets/attachments/bedrock, then
    //    split_run_options + apply_run_configuration.
    let (task, should_share) = spawner.spawn(move |me, ctx| me.install_run_configuration(..)).await??;

    // 3. Shell replacement (needs resolved_env_vars from step 2).
    Self::replace_shell_with_resolved_env(spawner).await?;

    // 4. Sharing (after the shell) overlaps with the rest of the MCP restart.
    futures::try_join!(Self::bind_and_share(spawner, task_id, plan.share, should_share), mcp_restart)?;
    Ok(task)
}

If you'd rather keep this strictly "groups, each a join", the MCP restart can go in step 3 next to the shell replacement. That's simpler, but gives up overlapping MCP with the network group. Starting it early seems worth it: it's the only step with no inputs and usually the slowest, and one future started early and joined at the end stays easy to follow.

Roughly, the critical path goes from the sum of everything (five network round trips + shell + MCP stop/start + share wait) to the slowest network call + the shell replacement + whichever is longer of the share wait and the remaining MCP time. That's an estimate from reading the code, not a measurement.

Notes:

• Use try_join! so the first error cancels the rest; dropping a half-finished MCP restart or shell replacement is fine, since a failed installation ends the process.
• I didn't confirm whether building options actually reads team metadata. The shape above doesn't depend on it, because step 2 waits for all of step 1. If it doesn't, the refresh could move off the critical path entirely (joined before bind).
• This pairs with the MCP comment on mcp_startup.rs: if MCP restart failures degrade instead of failing the run, joining the MCP restart last makes a degraded result just a warning.
• Nothing else waits on the MCP manager during installation (the next wait is in run()), so the overlapping-waits warning on wait_for_mcp_servers_started isn't an issue, as long as that stays true.

[captainsafia]

@vorporeal Thanks! I'll take advantage of the approval to unblock being able to test the end-to-end in the next nightly. I've addressed a few of the impact impact items in the PR here and will circle back with the others in follow-ups.