feat(sprites): lazy-warm named-attach sessions

When attaching to an existing sprite (created_by_us=False),
SpritesSandboxSession now skips the eager wait-for-running poll
and lets the first I/O operation drive the wake-up via a new
_ensure_warm() guard. The platform auto-wakes a paused sprite when
traffic arrives, so this is essentially free — and avoids paying
1–10s of polling latency just to hand back a session handle.

The created-by-us path is unchanged: a fresh sprite still needs a
provisioning poll, and we set _warmth_verified=True after.

A new _invalidate_warmth() hook lets recovery flows force a re-poll
after a transport error.

Live timing of named-attach:
- create():       3s → 0.01s
- first exec:     adds the wake-up roundtrip (formerly paid eagerly)
- subsequent I/O: unchanged (cached warm flag)

3 new lazy-warm tests; the _attach test helper now also marks
_warmth_verified=True so existing tests that bypass startup keep
working.

Author: vaurdan

src/agents/extensions/sandbox/sprites/sandbox.py

@@ -257,6 +257,7 @@ class SpritesSandboxSession(BaseSandboxSession):
     _pty_lock: asyncio.Lock
     _pty_processes: dict[int, _SpritePtyProcessEntry]
     _reserved_pty_process_ids: set[int]
+    _warmth_verified: bool
 
     def __init__(
         self,
@@ -276,6 +277,12 @@ def __init__(
         self._pty_lock = asyncio.Lock()
         self._pty_processes = {}
         self._reserved_pty_process_ids = set()
+        # ``_warmth_verified`` tracks whether we've already confirmed the sprite
+        # is warm/running. Set when a fresh sprite is provisioned (we have to
+        # poll anyway), or when the first I/O operation drives a successful
+        # control-plane connect. Resetting via ``_invalidate_warmth`` after a
+        # transport error forces a re-check on the next operation.
+        self._warmth_verified = False
 
     @classmethod
     def from_state(
@@ -336,19 +343,47 @@ async def _ensure_sprite(self) -> Sprite:
                     cause=exc,
                 ) from exc
             await self._maybe_update_url_settings(sprite)
-        else:
-            sprite = await asyncio.to_thread(client.sprite, self.state.sprite_name)
-
+            self._sprite = sprite
+            # Fresh sprite: poll until it reaches a ready status, then refresh
+            # its info so ``url`` / ``organization_name`` are populated. We
+            # have to poll regardless because the create call doesn't tell us
+            # when provisioning finishes.
+            await self._wait_for_sprite_running()
+            self._warmth_verified = True
+            refreshed: Sprite
+            try:
+                refreshed = await asyncio.to_thread(client.get_sprite, self.state.sprite_name)
+            except SpriteError:
+                refreshed = sprite
+            self._sprite = refreshed
+            return refreshed
+
+        # Named-attach: skip the readiness poll on attach. The sprite may be
+        # cold (paused) and we don't want to pay wake-up latency just to hand
+        # the agent a session handle. The first I/O operation drives the
+        # wake-up via ``_ensure_warm`` — the platform auto-wakes a paused
+        # sprite when traffic arrives, so this is essentially free.
+        sprite = await asyncio.to_thread(client.sprite, self.state.sprite_name)
         self._sprite = sprite
+        return sprite
+
+    async def _ensure_warm(self) -> None:
+        """Block until the sprite is ready to accept I/O, but only on first use.
+
+        ``_warmth_verified`` is sticky for the life of the session; cached
+        until a transport error invalidates it (e.g., the sprite was deleted
+        out from under us and we have to re-attach in a recovery flow).
+        """
+
+        if self._warmth_verified:
+            return
         await self._wait_for_sprite_running()
-        # Refresh sprite info so url / organization_name / status are populated.
-        refreshed: Sprite
-        try:
-            refreshed = await asyncio.to_thread(client.get_sprite, self.state.sprite_name)
-        except SpriteError:
-            refreshed = sprite
-        self._sprite = refreshed
-        return refreshed
+        self._warmth_verified = True
+
+    def _invalidate_warmth(self) -> None:
+        """Force the next I/O operation to re-poll the sprite's status."""
+
+        self._warmth_verified = False
 
     def _build_sprite_config(self) -> sprites.SpriteConfig | None:
         if (
@@ -586,6 +621,8 @@ async def _exec_with_cwd(
         if not normalized:
             return ExecResult(stdout=b"", stderr=b"", exit_code=0)
 
+        await self._ensure_warm()
+
         control: ControlConnection | None = None
         op_conn: OpConn | None = None
         try:
@@ -656,6 +693,7 @@ async def pty_exec_start(
     ) -> PtyExecUpdate:
         sanitized_command = self._prepare_exec_command(*command, shell=shell, user=user)
         # ``_ensure_control`` will lazily call ``_ensure_sprite``; no extra await here.
+        await self._ensure_warm()
 
         cc: ControlConnection | None = None
         op: OpConn | None = None
@@ -972,6 +1010,7 @@ async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase
 
         normalized_path = await self._validate_path_access(path)
         sprite = await self._ensure_sprite()
+        await self._ensure_warm()
         try:
             payload = await asyncio.to_thread(
                 lambda: (sprite.filesystem("/") / sandbox_path_str(normalized_path)).read_bytes()
@@ -1003,6 +1042,7 @@ async def write(
             )
 
         sprite = await self._ensure_sprite()
+        await self._ensure_warm()
         try:
             await asyncio.to_thread(
                 lambda: (sprite.filesystem("/") / sandbox_path_str(normalized_path)).write_bytes(

tests/extensions/test_sandbox_sprites.py

@@ -47,13 +47,20 @@ def _attach(inner: SpritesSandboxSession, *, client: Any, sprite: Any = None) ->
     ``setattr`` is used to sidestep mypy's invariant attribute typing — the fakes
     duck-type the real ``SpritesClient``/``Sprite`` interface only as far as the
     tests exercise.
+
+    Also marks ``_warmth_verified=True`` so I/O paths skip the lazy
+    wake-up poll — the test has already set up the fake sprite directly,
+    so we can trust it's "warm enough" for the assertions under test.
+    Tests that specifically exercise the wait-for-running poll override
+    this back to False.
     """
 
     # ``setattr`` (instead of plain assignment) silences mypy's invariant attribute
     # check; the fakes only duck-type the parts we exercise.
     setattr(inner, "_client", client)  # noqa: B010
     if sprite is not None:
         setattr(inner, "_sprite", sprite)  # noqa: B010
+        setattr(inner, "_warmth_verified", True)  # noqa: B010
 
 
 # ---------- Fakes ----------
@@ -1072,3 +1079,76 @@ def test_sprites_checkpoints_tool_count_depends_on_allow_restore() -> None:
     cap_yes = SpritesCheckpoints(allow_restore=True)
     assert len(cap_no.tools()) == 2  # create + list
     assert len(cap_yes.tools()) == 3  # + restore
+
+
+# ---------- 17. Lazy wake-up ----------
+
+
+@pytest.mark.asyncio
+async def test_named_attach_create_does_not_poll_for_running(
+    patched_sprites: dict[str, Any],
+) -> None:
+    """create() with sprite_name should NOT call get_sprite during attach.
+
+    The platform auto-wakes the sprite on first traffic; polling here would
+    pay wake-up latency just to hand back a session handle. The first I/O
+    operation drives the wake-up via _ensure_warm.
+    """
+
+    fake_client = patched_sprites["client"]
+    fake_client._sprites_by_name["existing"] = _FakeSprite(name="existing")
+    client = SpritesSandboxClient(token="tok")
+    options = SpritesSandboxClientOptions(sprite_name="existing")
+    session = await client.create(options=options)
+    inner = session._inner
+    assert isinstance(inner, SpritesSandboxSession)
+    # No get_sprite calls because we did not poll for warmth.
+    assert fake_client.get_sprite_calls == []
+    # And the warmth flag stays False, so the next I/O will trigger the poll.
+    assert inner._warmth_verified is False
+
+
+@pytest.mark.asyncio
+async def test_lazy_warm_polls_on_first_exec(patched_sprites: dict[str, Any]) -> None:
+    fake_client = patched_sprites["client"]
+    fake_client._sprites_by_name["existing"] = _FakeSprite(name="existing")
+    fake_control = patched_sprites["control"]
+    fake_control.next_ops.append(_FakeOpConn(stdout=b"", exit_code=0))
+
+    client = SpritesSandboxClient(token="tok")
+    session = await client.create(options=SpritesSandboxClientOptions(sprite_name="existing"))
+    inner = session._inner
+    assert isinstance(inner, SpritesSandboxSession)
+    assert inner._warmth_verified is False
+    assert fake_client.get_sprite_calls == []
+
+    # First exec drives the wake-up poll.
+    await inner._exec_internal("echo", "hi")
+    assert fake_client.get_sprite_calls == ["existing"]
+    assert inner._warmth_verified is True
+
+    # Subsequent exec does NOT re-poll.
+    fake_control.next_ops.append(_FakeOpConn(stdout=b"", exit_code=0))
+    await inner._exec_internal("echo", "hi2")
+    # Still just the one poll from the first call.
+    assert fake_client.get_sprite_calls == ["existing"]
+
+
+@pytest.mark.asyncio
+async def test_lazy_warm_invalidate_forces_repoll(patched_sprites: dict[str, Any]) -> None:
+    fake_client = patched_sprites["client"]
+    fake_client._sprites_by_name["existing"] = _FakeSprite(name="existing")
+    fake_control = patched_sprites["control"]
+    fake_control.next_ops.extend([_FakeOpConn(exit_code=0), _FakeOpConn(exit_code=0)])
+
+    client = SpritesSandboxClient(token="tok")
+    session = await client.create(options=SpritesSandboxClientOptions(sprite_name="existing"))
+    inner = session._inner
+    assert isinstance(inner, SpritesSandboxSession)
+
+    await inner._exec_internal("echo", "1")
+    assert len(fake_client.get_sprite_calls) == 1
+
+    inner._invalidate_warmth()
+    await inner._exec_internal("echo", "2")
+    assert len(fake_client.get_sprite_calls) == 2