From ff9b955494178fdc19df04e48e56f0078bbd62ee Mon Sep 17 00:00:00 2001 From: Devika Verma Date: Mon, 2 Nov 2026 10:50:08 +0530 Subject: [PATCH] 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 | 37 +++++++++++++++++++++---- coworker/observability/engine_hooks.py | 22 +++++++++++++-- coworker/observability/provider_view.py | 11 +++++--- tests/test_observability_engine.py | 8 +++++- tests/test_observability_model_steps.py | 21 ++++++++++++++ 5 files changed, 87 insertions(+), 12 deletions(-) diff --git a/coworker/engine.py b/coworker/engine.py index 59f29301..b7984c47 100644 --- a/coworker/engine.py +++ b/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 diff --git a/coworker/observability/engine_hooks.py b/coworker/observability/engine_hooks.py index 28fc261e..492a727d 100644 --- a/coworker/observability/engine_hooks.py +++ b/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 diff --git a/coworker/observability/provider_view.py b/coworker/observability/provider_view.py index f9e1d6f4..7b429cd8 100644 --- a/coworker/observability/provider_view.py +++ b/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 ---------------------------------------------------------------------------- diff --git a/tests/test_observability_engine.py b/tests/test_observability_engine.py index 1975d92d..8ca62406 100644 --- a/tests/test_observability_engine.py +++ b/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" diff --git a/tests/test_observability_model_steps.py b/tests/test_observability_model_steps.py index fd17ac8d..02bc2764 100644 --- a/tests/test_observability_model_steps.py +++ b/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