diff --git a/src/agent_loop.py b/src/agent_loop.py index 1565d08e5..4949667db 100644 --- a/src/agent_loop.py +++ b/src/agent_loop.py @@ -5521,6 +5521,15 @@ async def stream_agent_loop( except (json.JSONDecodeError, Exception): pass + # Only a successful, authorized document execution may affect the + # editor. Start the authorized stream before any completed-document + # event: handleDocUpdate finalizes that stream, while sending a + # doc_update first can enter diff mode and make the later stream + # discard/save the stale pre-update document. + if tool_result_is_successful(result): + for doc_event in _document_stream_events(block): + yield f'data: {json.dumps(doc_event)}\n\n' + # Emit doc-specific event for document tools — the frontend # document panel handles this; no need to show content in chat. if is_doc_tool and "action" in result: @@ -5759,13 +5768,6 @@ async def stream_agent_loop( f'data: {json.dumps({"type": "ask_user", "data": _pending_ask_user_event})}\n\n' ) - # Only a successful, authorized document execution may affect the - # editor. Model deltas and raw fences are proposals and can be - # invalidated by an earlier result in the same tool batch. - if tool_result_is_successful(result): - for doc_event in _document_stream_events(block): - yield f'data: {json.dumps(doc_event)}\n\n' - # Native document tools open in the editor + carry the REAL doc id. # Emit a doc_update so the frontend opens/activates it and sends it # back as active_doc_id next turn (otherwise the agent can't "see" diff --git a/tests/test_external_context_tool_gate.py b/tests/test_external_context_tool_gate.py index d65269176..ed9355a06 100644 --- a/tests/test_external_context_tool_gate.py +++ b/tests/test_external_context_tool_gate.py @@ -745,3 +745,54 @@ def test_document_stream_events_are_derived_from_authorized_block(): {"type": "doc_stream_open", "title": "Title", "language": "markdown"}, {"type": "doc_stream_delta", "content": "Body"}, ] + + +def test_authorized_document_stream_precedes_completed_update(monkeypatch): + import src.agent_loop as agent_loop + + monkeypatch.setattr( + agent_loop, + "get_setting", + lambda key, default=None: default, + raising=False, + ) + monkeypatch.setattr(agent_loop, "get_mcp_manager", lambda: None, raising=False) + monkeypatch.setattr(agent_loop, "estimate_tokens", lambda *args, **kwargs: 10) + + async def fake_stream(*args, **kwargs): + yield "data: " + json.dumps( + {"delta": "```update_document\nNew body\n```"} + ) + "\n\n" + yield "data: [DONE]\n\n" + + async def fake_execute(block, *args, **kwargs): + assert block.tool_type == "update_document" + return ( + block.tool_type, + { + "action": "update", + "doc_id": "doc-1", + "title": "Existing", + "language": "markdown", + "content": "New body", + "version": 2, + }, + ) + + monkeypatch.setattr(agent_loop, "stream_llm_with_fallback", fake_stream) + monkeypatch.setattr(agent_loop, "execute_tool_block", fake_execute) + + events = _collect_agent_events( + agent_loop.stream_agent_loop( + "http://local.test/v1", + "small-local-model", + [{"role": "user", "content": "update the active document"}], + max_rounds=1, + relevant_tools={"update_document"}, + ) + ) + event_types = [event.get("type") for event in events] + + assert event_types.index("doc_stream_open") < event_types.index("doc_update") + assert event_types.index("doc_stream_delta") < event_types.index("doc_update") + assert event_types.index("doc_update") < event_types.index("tool_output")