diff --git a/coworker/engine.py b/coworker/engine.py index ae34d4a0..7ef75aa9 100644 --- a/coworker/engine.py +++ b/coworker/engine.py @@ -1018,8 +1018,10 @@ def _outbound_messages(self) -> list[dict[str, Any]]: `ts` — (providers reject unknown keys), unconditionally — whether or not a `` block is added. When a context provider yields a non-empty string, an ephemeral `` block is appended to the - last user message. Never mutates `self.messages`, so neither the strip nor the block is - persisted/replayed. + last user message. A partial assistant turn abandoned by a turn that died mid-stream — + identified by the error/interrupted notice that follows it — is dropped too: display-only + like the notices, and rejected outright by several providers as the final message. + Never mutates `self.messages`, so none of this is persisted/replayed. """ # Strip the display-only sidecars — `source` (connector cards), `_display` # (e.g. filter-hidden counts), `ts` (append-time timestamps), `reasoning` @@ -1033,14 +1035,43 @@ def _outbound_messages(self) -> list[dict[str, Any]]: source_messages = _compaction.apply_to_outbound( self.messages, self.compaction_state ) + # A turn that dies mid-stream (provider error, interrupt) leaves its partial assistant + # message in the transcript — so the user can still see what streamed — followed by the + # notice saying why it stopped. That partial is display-only in the same sense the + # notices are: not a turn the model finished, and several providers reject it as the + # FINAL message. Gemini 400s ("Requests ending with a model turn are not supported"), + # Claude 4.6+ 400s ("does not support assistant message prefill"), and one ending in + # whitespace 400s on older Claude ("final assistant content cannot end with trailing + # whitespace"). `retry()` replays the same tail, so the session wedges for good. + # The trailing notice is what identifies it: a COMPLETED turn's assistant message looks + # identical but has no error/interrupted notice after it, and must be kept (its `extras` + # sidecar is replayed back to the owning provider). Trailing `tool_calls` mean a turn + # suspended awaiting durable resume (see `resume`), never an abandoned one. Consecutive + # failed retries stack one partial each, so walk back through the whole run. + abandoned: set[int] = set() + failed = False + for i in range(len(source_messages) - 1, -1, -1): + message = source_messages[i] + if message.get("role") == "notice": + failed = failed or message.get("kind") in ("error", "interrupted") + continue + if ( + failed + and message.get("role") == "assistant" + and not message.get("tool_calls") + ): + abandoned.add(i) + failed = False # that notice belongs to this partial + continue + break out = [ ( {k: v for k, v in msg.items() if k not in _SIDECARS} if any(s in msg for s in _SIDECARS) else msg ) - for msg in source_messages - if msg.get("role") != "notice" + for i, msg in enumerate(source_messages) + if msg.get("role") != "notice" and i not in abandoned ] # PDF attachments (stored as `file` parts) are adapted to the ACTIVE model right # here — never in the persisted history — so a mid-session model switch always diff --git a/tests/test_durable_resume.py b/tests/test_durable_resume.py index fee2bcad..264f0466 100644 --- a/tests/test_durable_resume.py +++ b/tests/test_durable_resume.py @@ -94,6 +94,47 @@ async def scenario(): assert mgr.inbox.pending(sid) == [] # nothing left pending +def test_resume_payload_keeps_the_suspended_tool_call(tmp_path): + # `_outbound_messages` drops a TRAILING assistant turn (the partial a mid-stream death + # leaves behind). A suspended turn is also stored ending on an assistant turn, so this + # pins that resume is unaffected: the pending calls are answered first, making the tail a + # `tool` message, and the assistant turn that made the call still reaches the model. + class Recording(ScriptedProvider): + def __init__(self, turns): + super().__init__(turns) + self.payloads = [] + + def complete(self, *, model, messages, tools=None, **settings): + self.payloads.append(messages) + return super().complete( + model=model, messages=messages, tools=tools, **settings + ) + + provider = Recording( + [ + _tool("ask_user", {"question": "Which region?", "options": ["a", "b"]}, "call_q"), + _text("You chose b."), + ] + ) + mgr = SessionManager(workspace=tmp_path, provider=provider) + sid = "dur-payload" + + async def scenario(): + engine = mgr.get_engine(sid, agent="cowork", workspace=str(tmp_path)) + item = await _run_until_pending(mgr, sid, engine) + await mgr.resolve_inbox(item.id, "b") + + asyncio.run(scenario()) + + resumed = next( + p for p in provider.payloads if any(m.get("role") == "tool" for m in p) + ) + assert [m.get("role") for m in resumed][-2:] == ["assistant", "tool"] + assert any( + m.get("role") == "assistant" and m.get("tool_calls") for m in resumed + ), "the assistant turn that made the call must still reach the model" + + def test_durable_resume_approval_executes_tool(tmp_path): # The model wants a write (needs approval); on durable resume "allow" must RE-EXECUTE the tool. target = tmp_path / "scratch_marker.txt" diff --git a/tests/test_engine.py b/tests/test_engine.py index 90e38fdb..7e6cb25c 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -319,6 +319,88 @@ def test_streaming_emits_deltas(tmp_path): assert events[-1].type == EventType.TURN_END +class DyingStreamProvider(ProviderClient): + """Streams a delta, then dies — the shape of a provider 503 landing mid-turn.""" + + def __init__(self): + self.payloads = [] + + def complete(self, **kwargs): # pragma: no cover - streamed instead + raise NotImplementedError + + def capabilities(self, model): + return ModelCapabilities() + + def stream(self, *, model, messages, tools=None, **settings): + self.payloads.append(messages) + yield StreamChunk(text_delta="I will request permission to ") + raise RuntimeError("Error code: 503 - {'error': {'type': 'overloaded_error'}}") + + +def test_retry_after_mid_stream_death_does_not_replay_the_partial(tmp_path): + # The partial stays in the transcript so the user can see what streamed, but several + # providers reject a payload whose FINAL message is an assistant turn — Gemini ("Requests + # ending with a model turn are not supported"), Claude 4.6+ ("does not support assistant + # message prefill"), older Claude when it ends in whitespace. Replaying it wedges retry. + provider = DyingStreamProvider() + engine = TurnEngine( + provider=provider, + registry=ToolRegistry(), + permissions=PermissionEngine(workspace_root=tmp_path), + model="gpt-5.5", + ) + _collect(engine, "arrange my visa documents") + + # kept in history (display), followed by the error notice + assert engine.messages[-2]["role"] == "assistant" + assert engine.messages[-2]["content"] == "I will request permission to " + assert engine.messages[-1]["role"] == "notice" + + async def _retry(): + return [ev async for ev in engine.retry()] + + asyncio.run(_retry()) + + sent = provider.payloads[-1] + assert sent, "retry never reached the provider" + assert sent[-1]["role"] != "assistant" + assert all(m.get("content") != "I will request permission to " for m in sent) + + +def test_completed_assistant_turn_is_never_dropped(tmp_path): + # The counterpart to the test above: a finished turn's assistant message is structurally + # identical to an abandoned partial (trailing, no tool_calls). Only the error notice tells + # them apart, so a completed turn with no notice after it must survive intact. + engine, _ = _engine(tmp_path, [_text_turn("all done")]) + _collect(engine, "hi") + + assert engine.messages[-1]["role"] == "assistant" + outbound = engine._outbound_messages() + assert outbound[-1]["role"] == "assistant" + assert outbound[-1]["content"] == "all done" + + +def test_consecutive_deaths_do_not_stack_partials_in_the_payload(tmp_path): + # Each failed retry appends its own partial, so one drop is not enough. + provider = DyingStreamProvider() + engine = TurnEngine( + provider=provider, + registry=ToolRegistry(), + permissions=PermissionEngine(workspace_root=tmp_path), + model="gpt-5.5", + ) + _collect(engine, "go") + + async def _retry(): + return [ev async for ev in engine.retry()] + + asyncio.run(_retry()) + asyncio.run(_retry()) + + assert sum(1 for m in engine.messages if m.get("role") == "assistant") == 3 + assert provider.payloads[-1][-1]["role"] != "assistant" + + def _pdf_file_part(): import base64 import io