fix(sandbox): accept extra stdout in kernel warmup probe (#52)

The warmup probe required stdout to be byte-identical to its marker,
so any extra interpreter output alongside the marker discarded a
healthy session and re-probed. A liveness probe only needs to prove
the marker round-tripped: check containment instead. Cold/stale
sessions still lack the marker, so the reset_session() retry path is
preserved.

Robustness cleanup per the re-scope of #51: the originally reported
latency figures were withdrawn (measured against a pre-#95 sandbox
image), so no performance claim is made. The docstring is also
rewritten stack-agnostically; whether the cold-start race exists
against the current Go sandbox agent is unverified.

Closes #51

Author: Korel

src/prokube/sandbox/sandbox.py

@@ -405,18 +405,23 @@ def wait_until_ready(self, timeout: int = 120) -> None:
         )
 
     def _warmup_kernel(self, deadline: float) -> None:
-        """Probe the Jupyter kernel until it echoes a unique marker back.
+        """Probe the sandbox interpreter until it echoes a unique marker back.
 
-        After the pod reaches Running, ipykernel still needs ~1–2s before its
-        first ``execute_request`` produces visible iopub stdout. During that
-        window, execd's ``exec`` sub-resource returns successfully but the
-        stdout never reaches the SSE stream, so the first user ``run_code()``
-        call races the cold kernel pipeline and silently returns empty output.
+        The first execution after a pod reaches Running can race a cold
+        interpreter: a probe may return successfully with empty stdout before
+        the execution pipeline is fully live, in which case the first user
+        ``run_code()`` call would silently return empty output. (Whether the
+        current sandbox agent still exhibits this race is unverified; the
+        probe is kept as cheap insurance until warmup is re-measured against
+        it.)
 
         This method hides that race by running a tiny ``print(<marker>)``
-        probe in a loop until the stdout matches, proving the kernel pipeline
-        is end-to-end live. The probe is bounded by ``deadline`` (the same
-        deadline used by :meth:`wait_until_ready`), so it can never exceed the
+        probe in a loop until the marker appears in stdout, proving the
+        execution pipeline is end-to-end live. The check is containment, not
+        equality: the interpreter may append unrelated text (e.g. warnings)
+        to the same stream, and extra output does not make the session any
+        less live. The probe is bounded by ``deadline`` (the same deadline
+        used by :meth:`wait_until_ready`), so it can never exceed the
         caller's overall timeout budget. If the deadline is reached without
         success, a warning is logged and the method returns without raising —
         the user may still get useful results, and we don't want to block
@@ -425,10 +430,10 @@ def _warmup_kernel(self, deadline: float) -> None:
         Notes:
             * The marker is per-call (``uuid4().hex``) to avoid collisions
               with any user code that happens to print a similar literal.
-            * If a probe returns successfully but without the marker, discard
-              that session before retrying. Otherwise a cold/stale Jupyter
-              session can be reused forever and every probe keeps returning
-              empty stdout.
+            * If a probe returns successfully but stdout does not contain the
+              marker, discard that session before retrying. Otherwise a
+              cold/stale session can be reused forever and every probe keeps
+              returning empty stdout.
 
         Args:
             deadline: ``time.monotonic()`` value after which the probe gives
@@ -471,7 +476,7 @@ def _warmup_kernel(self, deadline: float) -> None:
                 sleep_for = min(0.5, max(0.0, deadline - time.monotonic()))
                 time.sleep(sleep_for)
                 continue
-            if result.stdout.strip() == marker:
+            if marker in result.stdout:
                 return
             self._code.reset_session()
             # Loop top will recompute remaining and exit if deadline passed.

tests/test_pause_resume.py

@@ -742,6 +742,61 @@ def _probe_callback(request: httpx.Request) -> httpx.Response:
 
         sbx._client.close()
 
+    def test_wait_until_ready_warmup_accepts_extra_stdout(
+        self, mock_env, monkeypatch, httpx_mock: HTTPXMock
+    ):
+        """A probe whose stdout contains the marker plus extra text succeeds in one attempt.
+
+        Regression test for issue #51: the kernel may append unrelated
+        warnings (e.g. IPython's history-thread SQLite error) to the same
+        stdout as the marker. That session is live and must not be discarded.
+        """
+        monkeypatch.setattr("time.sleep", lambda _: None)
+
+        _mock_version(httpx_mock)
+        _mock_claim(httpx_mock)
+        httpx_mock.add_response(
+            method="GET",
+            url=f"{BASE}/_platform/sandbox/test-ws/sandboxes/sandbox-test",
+            json={"name": "sandbox-test", "phase": "Running"},
+        )
+
+        call_counter = {"n": 0}
+
+        def _probe_callback(request: httpx.Request) -> httpx.Response:
+            call_counter["n"] += 1
+            marker = _extract_marker(request) or ""
+            return httpx.Response(
+                200,
+                json={
+                    "stdout": (
+                        f"{marker}\n"
+                        "The history saving thread hit an unexpected error "
+                        "(OperationalError('attempt to write a readonly "
+                        "database')). History will not be written to the "
+                        "database.\n"
+                    ),
+                    "stderr": "",
+                    "success": True,
+                    "execution_time_ms": 1,
+                },
+            )
+
+        httpx_mock.add_callback(
+            _probe_callback,
+            method="POST",
+            url=f"{BASE}/_platform/sandbox/test-ws/sandboxes/sandbox-test/exec",
+            is_reusable=True,
+        )
+
+        sbx = Sandbox.from_pool("python-pool")
+        sbx.wait_until_ready(timeout=5)
+
+        assert sbx.status == "Running"
+        assert call_counter["n"] == 1
+
+        sbx._client.close()
+
     def test_wait_until_ready_warmup_caps_per_probe_timeout(
         self, mock_env, monkeypatch, httpx_mock: HTTPXMock
     ):