From f6991f16820f82baa6366341a705c60d90b4c2e7 Mon Sep 17 00:00:00 2001 From: zoujiawei Date: Thu, 30 Jul 2026 10:42:46 +0800 Subject: [PATCH] fix(a2a): preserve authoritative text replacements Propagate runner replace markers through the RuntimeEvent bridge and A2A artifacts, while leaving ADK partial chunks on their documented incremental semantics. Add direct executor and public route regression coverage. --- ksadk/a2a/executor.py | 12 +++-- ksadk/runtime/runner_adapter.py | 14 +++-- tests/test_a2a_text_replacement.py | 85 ++++++++++++++++++++++++++++++ 3 files changed, 102 insertions(+), 9 deletions(-) create mode 100644 tests/test_a2a_text_replacement.py diff --git a/ksadk/a2a/executor.py b/ksadk/a2a/executor.py index 0ea605d..61eb722 100644 --- a/ksadk/a2a/executor.py +++ b/ksadk/a2a/executor.py @@ -342,8 +342,9 @@ async def _run_streaming( if self.include_reasoning: await artifacts.push("thinking", text) continue - output_text += text - await artifacts.push("text", text) + replace = bool(isinstance(chunk, dict) and chunk.get("replace")) + output_text = text if replace else output_text + text + await artifacts.push("text", text, replace_snapshot=replace) await artifacts.close() return output_text @@ -444,7 +445,8 @@ async def _run_runtime( reasoning_text += text await artifacts.push("thinking", text) continue - # TEXT_COMPLETED 是累计全文,去重只发新增 suffix;TEXT_DELTA 是增量直接透传。 + # TEXT_COMPLETED 是累计全文,去重只发新增 suffix;TEXT_DELTA 默认是增量, + # 但 runner 显式标记 replace 时是权威快照。 if event.event_type == EventType.TEXT_COMPLETED: if not output_text: delta = text @@ -460,8 +462,8 @@ async def _run_runtime( replace_snapshot = True else: delta = text - output_text += text - replace_snapshot = False + replace_snapshot = bool(event.payload.get("replace")) + output_text = text if replace_snapshot else output_text + text if not delta: continue await artifacts.push("text", delta, replace_snapshot=replace_snapshot) diff --git a/ksadk/runtime/runner_adapter.py b/ksadk/runtime/runner_adapter.py index b7e9ecd..cc7eaab 100644 --- a/ksadk/runtime/runner_adapter.py +++ b/ksadk/runtime/runner_adapter.py @@ -597,7 +597,7 @@ async def _map_runner_stream( else nullcontext() ) runner_name = _runner_name(self._runner) - accumulated_output: list[str] = [] + accumulated_output = "" usage: dict[str, Any] = {} runner_gen: Optional[AsyncIterator[Any]] = None async with _conversation_span_scope(runner_name) as span: @@ -664,7 +664,10 @@ async def _map_runner_stream( chunk.get("delta") or chunk.get("output") or chunk.get("data") ) if text: - accumulated_output.append(text) + if chunk_type == "final" or chunk.get("replace"): + accumulated_output = text + else: + accumulated_output += text raw_usage = chunk.get("usage") if isinstance(raw_usage, dict): usage.update(raw_usage) @@ -677,7 +680,7 @@ async def _map_runner_stream( yield self._event(handle, event_type, payload) finally: if accumulated_output: - _set_conversation_output_attributes(span, "".join(accumulated_output)) + _set_conversation_output_attributes(span, accumulated_output) _set_conversation_usage_attributes(span, usage) if runner_gen is not None: aclose = getattr(runner_gen, "aclose", None) @@ -840,7 +843,10 @@ def _chunk_to_event( text = self._coerce(chunk.get("delta") or chunk.get("output") or chunk.get("data")) if not text: return None - return self._event(handle, EventType.TEXT_DELTA, {"text": text}, phase="commentary") + payload: dict[str, Any] = {"text": text} + if chunk.get("replace"): + payload["replace"] = True + return self._event(handle, EventType.TEXT_DELTA, payload, phase="commentary") def _event( self, diff --git a/tests/test_a2a_text_replacement.py b/tests/test_a2a_text_replacement.py new file mode 100644 index 0000000..e3be76a --- /dev/null +++ b/tests/test_a2a_text_replacement.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from types import SimpleNamespace +from typing import Any + +import httpx +import pytest +from fastapi import FastAPI + +from ksadk.a2a import A2ARuntimeTaskAdapter, add_a2a_protocol_routes +from ksadk.a2a.executor import A2ARuntimeExecutor +from ksadk.a2a.langgraph import stream_a2a_agent_to_writer +from ksadk.a2a.routes import A2AConfig +from ksadk.runtime.runner_adapter import RunnerRuntimeAdapter + + +class _RecordingUpdater: + def __init__(self) -> None: + self.artifacts: list[dict[str, Any]] = [] + + async def add_artifact(self, **kwargs: Any) -> None: + self.artifacts.append(kwargs) + + +class _ReplacingRunner: + async def stream(self, _runner_input: dict[str, Any]) -> AsyncIterator[dict[str, Any]]: + yield {"type": "text", "delta": "旧答"} + yield {"type": "text", "delta": "新答", "replace": True} + yield {"type": "final", "output": "新答"} + + +@pytest.mark.asyncio +async def test_runner_text_replacement_is_marked_as_authoritative_snapshot() -> None: + runner = _ReplacingRunner() + executor = A2ARuntimeExecutor(runner=runner) + updater = _RecordingUpdater() + + output = await executor._run_streaming( # noqa: SLF001 + SimpleNamespace(task_id="task-1"), + updater, # type: ignore[arg-type] + runner.stream, + {}, + ) + + assert output == "新答" + assert [item["parts"][0].text for item in updater.artifacts] == ["旧答", "新答"] + assert updater.artifacts[-1]["append"] is False + assert updater.artifacts[-1]["parts"][0].metadata["ksadk_output_snapshot"] is True + + +@pytest.mark.asyncio +async def test_a2a_server_round_trip_preserves_text_replacement(tmp_path: Any) -> None: + runner = _ReplacingRunner() + app = FastAPI() + add_a2a_protocol_routes( + app, + runner, + A2AConfig( + enabled=True, + base_url="http://testserver", + agent_name="replacing-agent", + task_store_dsn=f"sqlite+aiosqlite:///{tmp_path}/tasks.db", + ), + task_adapter=A2ARuntimeTaskAdapter( + RunnerRuntimeAdapter(runner, runtime_type="test"), # type: ignore[arg-type] + runtime_type="test", + ), + ) + written: list[dict[str, Any]] = [] + transport = httpx.ASGITransport(app=app) + + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + output = await stream_a2a_agent_to_writer( + "http://testserver", + "问题", + writer=written.append, + httpx_client=client, + ) + + assert output == "新答" + assert written == [ + {"type": "text", "delta": "旧答", "replace": False}, + {"type": "text", "delta": "新答", "replace": True}, + ]