From 956a26715b02ec2d60c6a6ffcf9c6146633a4318 Mon Sep 17 00:00:00 2001
From: Ahmed Allam <ahmed39652003@gmail.com>
Date: Wed, 14 Oct 2026 00:09:35 +0000
Subject: [PATCH] feat(runtime): accept an awaitable extra_files in
 create_or_reuse so the sandbox create does not wait for it

---
 strix/core/runner.py             |  4 ++--
 strix/runtime/session_manager.py | 31 +++++++++++++++++++-----------
 tests/test_deferred_session.py   | 33 ++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/strix/core/runner.py b/strix/core/runner.py
index f9bb4940..5c5182a6 100644
--- a/strix/core/runner.py
+++ b/strix/core/runner.py
@@ -8,7 +8,7 @@ import io
 import json
 import logging
 import uuid
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -203,7 +203,7 @@ async def run_strix_scan(
     scan_id: str | None = None,
     image: str,
     local_sources: list[dict[str, Any]] | None = None,
-    extra_files: list[dict[str, Any]] | None = None,
+    extra_files: list[dict[str, Any]] | Awaitable[list[dict[str, Any]] | None] | None = None,
     coordinator: AgentCoordinator | None = None,
     interactive: bool = False,
     max_turns: int = DEFAULT_MAX_TURNS,
diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py
index cfa5bc4c..d92b1fdb 100644
--- a/strix/runtime/session_manager.py
+++ b/strix/runtime/session_manager.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
+import inspect
 import io
 import logging
 import os
@@ -259,12 +260,20 @@ def _gitdir_from_pointer(git_file: Path) -> Path | None:
     return None
 
 
+async def _resolve_extra_file_archive(
+    extra_files: list[dict[str, Any]] | Awaitable[list[dict[str, Any]] | None] | None,
+    local_sources: list[dict[str, Any]],
+) -> bytes | None:
+    files = await extra_files if inspect.isawaitable(extra_files) else extra_files
+    return build_extra_file_archive(files, local_sources) if files else None
+
+
 async def create_or_reuse(
     scan_id: str,
     *,
     image: str,
     local_sources: list[dict[str, Any]],
-    extra_files: list[dict[str, Any]] | None = None,
+    extra_files: list[dict[str, Any]] | Awaitable[list[dict[str, Any]] | None] | None = None,
     status_sink: StatusSink | None = None,
     defer: bool = False,
 ) -> dict[str, Any]:
@@ -276,7 +285,9 @@ async def create_or_reuse(
     Each ``extra_files`` entry (``{"workspace_path": "/workspace/<rel>",
     "content": bytes | str}``) lands as a regular, agent-writable file at its
     ``workspace_path`` on every backend: the files are uploaded as one archive
-    and unpacked inside the sandbox right after bring-up.
+    and unpacked inside the sandbox right after bring-up. ``extra_files`` may
+    also be an awaitable resolving to that list: provisioning starts at once and
+    the files are awaited only when the sandbox is ready to receive them.
     """
 
     def report(phase: str) -> None:
@@ -297,9 +308,6 @@ async def create_or_reuse(
     else:
         bind_mounts = []
         entries = build_manifest_entries(local_sources)
-    extra_file_archive = (
-        build_extra_file_archive(extra_files, local_sources) if extra_files else None
-    )
 
     # Caido runs as an in-container sidecar; HTTP(S) traffic from any
     # process started via ``session.exec`` (the SDK's Shell tool, etc.)
@@ -339,13 +347,14 @@ async def create_or_reuse(
                 bind_mounts=bind_mounts,
             )
         )
-        if extra_file_archive is not None:
-            report("Placing workspace files")
-            try:
+        try:
+            extra_file_archive = await _resolve_extra_file_archive(extra_files, local_sources)
+            if extra_file_archive is not None:
+                report("Placing workspace files")
                 await stage_extra_files(session, extra_file_archive)
-            except BaseException:
-                await _discard_session(client, session)
-                raise
+        except BaseException:
+            await _discard_session(client, session)
+            raise
         return client, session
 
     async def _caido_url(session: Any) -> str:
diff --git a/tests/test_deferred_session.py b/tests/test_deferred_session.py
index aac5d226..774426fb 100644
--- a/tests/test_deferred_session.py
+++ b/tests/test_deferred_session.py
@@ -242,6 +242,39 @@ async def test_extra_files_are_staged_inside_the_deferred_bring_up(backend: _Bac
         await session_manager.cleanup(scan_id)
 
 
+@pytest.mark.asyncio
+async def test_awaitable_extra_files_do_not_delay_the_backend_call(backend: _Backend) -> None:
+    scan_id = "deferred-extra-awaitable"
+    files: asyncio.Future[list[dict[str, Any]] | None] = asyncio.get_running_loop().create_future()
+    try:
+        bundle = await _create(scan_id, extra_files=files)
+        backend.gate.set()
+        await _settle()
+        assert backend.calls == 1
+        assert backend.session.writes == []
+        files.set_result([{"workspace_path": "/workspace/.strix/notes.md", "content": "hi"}])
+        await asyncio.wait_for(bundle["session"].resolve(), timeout=1)
+        assert len(backend.session.writes) == 1
+    finally:
+        await session_manager.cleanup(scan_id)
+
+
+@pytest.mark.asyncio
+async def test_awaitable_extra_files_resolving_to_none_stage_nothing(backend: _Backend) -> None:
+    scan_id = "deferred-extra-none"
+
+    async def _nothing() -> list[dict[str, Any]] | None:
+        return None
+
+    try:
+        bundle = await _create(scan_id, extra_files=_nothing())
+        backend.gate.set()
+        await asyncio.wait_for(bundle["session"].resolve(), timeout=1)
+        assert backend.session.writes == []
+    finally:
+        await session_manager.cleanup(scan_id)
+
+
 @pytest.mark.asyncio
 async def test_default_path_still_blocks_until_the_sandbox_is_up(backend: _Backend) -> None:
     scan_id = "deferred-off"