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