From 61562e947d1ce938e0a2fbe51bd02d97b26991ff Mon Sep 17 00:00:00 2001
From: h <h@kotikot.com>
Date: Fri, 28 Aug 2026 19:38:47 +0200
Subject: [PATCH] fix(backends): remember the session id from the first frame
 so a cut turn stays resumable

---
 src/beaver_gateway/backends/claude_sdk.py | 18 +++++++++++++++---
 tests/test_conversations.py               | 13 +++++++++----
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/beaver_gateway/backends/claude_sdk.py b/src/beaver_gateway/backends/claude_sdk.py
index 66125f6..ddbed32 100644
--- a/src/beaver_gateway/backends/claude_sdk.py
+++ b/src/beaver_gateway/backends/claude_sdk.py
@@ -325,7 +325,9 @@ class ClaudeSdkBackend:
             try:
                 # Events go out as the CLI produces them: the frontends
                 # stream text and thinking live, the turn is not buffered.
-                async for event in self._run_turn(live, prompt, turn, observer):
+                async for event in self._run_turn(
+                    live, prompt, turn, observer, capture
+                ):
                     yield event
             except Exception:
                 live.running_turn = None
@@ -343,7 +345,9 @@ class ClaudeSdkBackend:
                 turn = _Turn()
                 async with live.lock:
                     live.running_turn = turn_id or message_id
-                    async for event in self._run_turn(live, prompt, turn, observer):
+                    async for event in self._run_turn(
+                        live, prompt, turn, observer, capture
+                    ):
                         yield event
             live.turns += 1
             live.last_used = time.monotonic()
@@ -403,12 +407,15 @@ class ClaudeSdkBackend:
         prompt: str,
         turn: _Turn,
         observer: Callable[[Any], None] | None = None,
+        capture: TurnCapture | None = None,
     ) -> AsyncIterator[MessageStreamEvent]:
         """Run one prompt, yielding wire events as they arrive.
 
         ``turn`` is filled in place (result, synthesized history, count of
         events already yielded) so the caller can finish bookkeeping - and
-        decide whether a retry is still possible - after a failure.
+        decide whether a retry is still possible - after a failure. The
+        session id lands in ``capture`` with the first frame, so a turn cut
+        by a restart still leaves a resumable session behind.
         """
         streaming = self._agent.options.include_partial_messages
         raw: list[Any] = []
@@ -418,6 +425,11 @@ class ClaudeSdkBackend:
         async for message in live.client.receive_response():
             if observer is not None:
                 observer(message)
+            session_id = getattr(message, "session_id", None)
+            if isinstance(session_id, str) and session_id and live.session_id is None:
+                live.session_id = session_id
+                if capture is not None:
+                    capture.session_id = session_id
             if isinstance(message, MirrorErrorMessage):
                 live.dirty = True
                 _log.error(
diff --git a/tests/test_conversations.py b/tests/test_conversations.py
index 01331b8..2874f72 100644
--- a/tests/test_conversations.py
+++ b/tests/test_conversations.py
@@ -96,11 +96,13 @@ class ScriptedClient:
     async def receive_response(self):
         prompt = self.prompts[-1]
         yield StreamEvent(
-            uuid="u", session_id="s", event={"type": "message_start", "message": {}}
+            uuid="u",
+            session_id=self.session_id,
+            event={"type": "message_start", "message": {}},
         )
         yield StreamEvent(
             uuid="u",
-            session_id="s",
+            session_id=self.session_id,
             event={
                 "type": "content_block_start",
                 "index": 0,
@@ -109,7 +111,7 @@ class ScriptedClient:
         )
         yield StreamEvent(
             uuid="u",
-            session_id="s",
+            session_id=self.session_id,
             event={
                 "type": "content_block_delta",
                 "index": 0,
@@ -117,7 +119,9 @@ class ScriptedClient:
             },
         )
         yield StreamEvent(
-            uuid="u", session_id="s", event={"type": "content_block_stop", "index": 0}
+            uuid="u",
+            session_id=self.session_id,
+            event={"type": "content_block_stop", "index": 0},
         )
         yield AssistantMessage(content=[TextBlock(text=f"ok:{prompt}")], model="m")
         hold = ScriptedClient.hold
@@ -505,6 +509,7 @@ async def test_graceful_stop_keeps_running_turn_for_recover(world: World) -> Non
     await world.conversations.stop()
     row = await world.conversations.get(conv.external_id)
     assert row.running_turn is not None
+    assert row.session_id == ScriptedClient.instances[0].session_id
     cut = await world.conversations.recover()
     assert [c.external_id for c in cut] == [conv.external_id]
     assert (await world.conversations.get(conv.external_id)).running_turn is None