fix(observability): a streamed model step ends as its reply arrives; a turn's end waits for shadow checks (OPE-244)

Two races found by running the schema test repeatedly:

- The streamed view ended its model step only after the engine had read
  the final piece. The engine could add the reply and end the turn first,
  and the turn's end then closed the step as interrupted, without its
  usage or finish reason. The step now ends when the final piece
  arrives, before the engine sees it.
- A shadow check runs in the background and can outlast the turn's last
  event. The turn's end is now recorded with the time of its ending
  event but sent once running shadow checks finish (at most 30 s), so
  their steps are not closed as interrupted while they run.

Also: a server test's fake provider answers the session title apart, so
the title call can no longer take the turn's only reply.

coworker/engine.py

@@ -234,7 +234,7 @@ def _records_turn(kind: str) -> Callable[[Callable[..., AsyncIterator[Event]]],
             finally:
                 await inner.aclose()
                 if started:
-                    obs.end_turn("interrupted" if self._cancel.is_set() else "unknown")
+                    self._end_recorded_turn("interrupted" if self._cancel.is_set() else "unknown")
 
         return wrapper
 
@@ -567,6 +567,7 @@ class TurnEngine:
         # Auto-approver check steps by tool call (E12), and shadow checks waiting to be
         # compared with your answer: call id -> (step id, verdict or None).
         self._review_steps: dict[str, str] = {}
+        self._turn_end_tasks: set[asyncio.Task] = set()
         # The coworker fingerprint and the session-start context facts, set by the session
         # build (observability E16, E5); and this turn's context block, to spot a change.
         self.persona_fp: Optional[str] = None
@@ -797,11 +798,34 @@ class TurnEngine:
             if status in ("max_iterations_exceeded", "truncated"):
                 obs.instant_step(kind="notice", name="step_limit" if status == "max_iterations_exceeded" else "truncated",
                                  status="ok", attrs={})
-            obs.end_turn(status)
+            self._end_recorded_turn(status)
         elif event.type is EventType.INTERRUPTED:
-            obs.end_turn("interrupted")
+            self._end_recorded_turn("interrupted")
         elif event.type is EventType.ERROR:
-            obs.end_turn("error")
+            self._end_recorded_turn("error")
+
+    def _end_recorded_turn(self, status: str) -> None:
+        """`turn_ended` with this moment's time. If shadow checks are still running in the
+        background, it is sent once they end (at most 30 s later), so their steps are not
+        closed as interrupted while they run; the turn's recorded end stays this moment."""
+        pending = self._obs.close_turn(status)
+        if pending is None:
+            return
+        running = [t for t in self._shadow_tasks if not t.done()]
+        if not running:
+            self._obs.emit_turn_ended(pending)
+            return
+
+        async def after_shadows() -> None:
+            try:
+                await asyncio.wait_for(asyncio.gather(*running, return_exceptions=True), 30)
+            except Exception:  # noqa: BLE001 - a slow shadow check must not lose the turn's end
+                pass
+            self._obs.emit_turn_ended(pending)
+
+        task = asyncio.ensure_future(after_shadows())
+        self._turn_end_tasks.add(task)
+        task.add_done_callback(self._turn_end_tasks.discard)
 
     def _risk_class(self, tool_name: str) -> Optional[str]:
         try:
@@ -2229,9 +2253,12 @@ class TurnEngine:
         task.add_done_callback(self._shadow_tasks.discard)
 
     async def drain_shadow_reviews(self) -> None:
-        """Await in-flight shadow verdicts (tests and orderly shutdown; never the hot path)."""
+        """Await in-flight shadow verdicts (tests and orderly shutdown; never the hot path),
+        and the recorded turn ends that waited for them."""
         if self._shadow_tasks:
             await asyncio.gather(*list(self._shadow_tasks), return_exceptions=True)
+        if self._turn_end_tasks:
+            await asyncio.gather(*list(self._turn_end_tasks), return_exceptions=True)
 
     async def _authorize(self, tool_call: ToolCall) -> "AsyncIterator[Event | bool]":
         """Permission flow for one call (TOOL_PROPOSED is emitted by the caller). Yields

coworker/observability/engine_hooks.py

@@ -105,9 +105,18 @@ class TurnRecording:
             return self.turn_id
 
     def end_turn(self, status: str) -> None:
+        pending = self.close_turn(status)
+        if pending is not None:
+            self.emit_turn_ended(pending)
+
+    def close_turn(self, status: str) -> Optional[dict[str, Any]]:
+        """The turn is over, now: what it still held is sent, and its `turn_ended` is built
+        with this moment's clocks. `emit_turn_ended` sends it, possibly later, once the
+        turn's background work (shadow checks) has ended, so the turn's end does not close
+        those steps as interrupted while they still run."""
         with self._lock:
             if self._ended or self.turn_id is None:
-                return
+                return None
             self._ended = True
             self._release(None)  # a retry that added no message at all
             for ev in self._tool_ends.values():  # a tool whose result message never came
@@ -116,7 +125,10 @@ class TurnRecording:
             self._flush_reply(None)
             ev = self._base("turn_ended", self.turn_id, self._turn_session)
             ev.update({"status": status, "last_message_id": self._last_mid})
-            self.recorder.emit(ev)
+            return ev
+
+    def emit_turn_ended(self, event: dict[str, Any]) -> None:
+        self.recorder.emit(event)
 
     @property
     def in_turn(self) -> bool:
@@ -423,6 +435,12 @@ class HelperRecording:
     def end_turn(self, status: str) -> None:
         pass
 
+    def close_turn(self, status: str) -> None:
+        return None
+
+    def emit_turn_ended(self, event: dict[str, Any]) -> None:
+        pass
+
     def note_message(self, message: dict[str, Any]) -> None:
         pass
 

coworker/observability/provider_view.py

@@ -142,16 +142,19 @@ class ProviderView(ProviderClient):
                     except StopIteration:
                         break
                 call.saw(chunk)
-                if getattr(chunk, "turn", None) is not None:
+                if getattr(chunk, "turn", None) is not None and turn is None:
+                    # The step ends as the reply arrives, before the engine sees it: the
+                    # engine may add the reply and end the turn before this generator runs
+                    # again, and the turn's end must find this step already closed.
                     turn = chunk.turn
+                    call.finished(turn)
                 yield chunk
         except GeneratorExit:
             raise
         except BaseException as exc:
-            call.failed(exc)
+            if turn is None:
+                call.failed(exc)
             raise
-        if turn is not None:
-            call.finished(turn)
 
 
 # -- one recorded call ----------------------------------------------------------------------------

tests/test_observability_engine.py

@@ -172,7 +172,13 @@ def test_a_turn_through_the_server_lands_in_recordings_db(tmp_path):
 
     from coworker.server import SessionManager, create_app
 
-    manager = SessionManager(workspace=tmp_path, provider=ScriptedProvider([_text("hi there")]))
+    class WithTitle(ScriptedProvider):
+        def complete(self, *, model, messages, tools=None, **settings):
+            if str(messages[0].get("content") or "").startswith("You title"):
+                return _text("Greeting")  # the session title, which can run before the turn's call
+            return super().complete(model=model, messages=messages, tools=tools, **settings)
+
+    manager = SessionManager(workspace=tmp_path, provider=WithTitle([_text("hi there")]))
     client = TestClient(create_app(manager))
     with client.websocket_connect("/ws/session/obs1") as ws:
         assert ws.receive_json()["type"] == "ready"

tests/test_observability_model_steps.py

@@ -265,3 +265,24 @@ def test_decision_reviewer_usage_vectors(vector):
 
     entries = [{"body": json.dumps({"usage": vector["raw"]})}]
     assert usage_attrs(_reply_usage(entries)) == {"reported": True, **vector["expected"]}
+
+
+def test_a_streamed_model_step_is_closed_before_the_engine_sees_the_reply():
+    """The engine may add the reply and end the turn before the stream generator runs again;
+    the turn's end must then find the model step closed, not close it as interrupted."""
+    from coworker.providers.base import StreamChunk
+
+    class Streaming(ScriptedProvider):
+        def stream(self, *, model, messages, tools=None, **settings):
+            yield StreamChunk(text_delta="hi")
+            yield StreamChunk(turn=AssistantTurn(text="hi", finish_reason="stop", usage=TokenUsage(input=3, output=1)))
+
+    rec = Recorder()
+    sink = TurnRecording(rec, lambda: "sess_s")
+    sink.begin_turn(kind="run", trigger="user", attrs={}, stamp={}, first_message_id=None)
+    view = RecordingProvider(Streaming([])).tagged(purpose="helper", sink=sink)
+    for chunk in view.stream(model="openai:model-a", messages=[]):
+        if chunk.turn is not None:
+            [step] = [s for s in rec.as_dict()["steps"] if s["kind"] == "model"]
+            assert step["status"] == "ok" and step["attrs"]["finish_reason"] == "stop"
+            assert step["attrs"]["ttft_ms"] is not None