diff --git a/tests/cli/test_frontend_runtime_proxy.py b/tests/cli/test_frontend_runtime_proxy.py index e0aaf52a9..08c5f2494 100644 --- a/tests/cli/test_frontend_runtime_proxy.py +++ b/tests/cli/test_frontend_runtime_proxy.py @@ -4657,6 +4657,304 @@ async def aclose(self) -> None: ) +def test_runtime_proxy_reports_incomplete_a2a_stream_after_partial_eof( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + app = _create_frontend_app(monkeypatch, tmp_path) + endpoint = "https://runtime.example" + rpc_url = endpoint + "/a2a/jsonrpc" + + class _FakeRuntimeClient: + def __init__(self, **kwargs: Any) -> None: + del kwargs + + def get_runtime(self, request: Any) -> SimpleNamespace: + del request + return SimpleNamespace( + runtime_id="runtime-1", + name="a2a-agent", + project_name="default", + network_configurations=[ + SimpleNamespace(endpoint=endpoint, network_type="public") + ], + authorizer_configuration=SimpleNamespace( + key_auth=SimpleNamespace(api_key="runtime-api-key"), + custom_jwt_authorizer=None, + ), + tags=[], + ) + + class _FakeUpstreamResponse: + def __init__( + self, + *, + status_code: int, + body: bytes, + content_type: str = "application/json", + ) -> None: + self.status_code = status_code + self._body = body + self.headers = {"content-type": content_type} + + async def aiter_raw(self): + yield self._body + + async def aclose(self) -> None: + pass + + class _FakeAsyncClient: + def __init__(self, **kwargs: Any) -> None: + del kwargs + + def build_request( + self, + method: str, + url: str, + *, + params: dict[str, str], + headers: dict[str, str], + content: bytes, + ) -> dict[str, Any]: + return { + "method": method, + "url": url, + "params": params, + "headers": headers, + "content": content, + } + + async def send(self, request: dict[str, Any], *, stream: bool): + del stream + url = request["url"] + if url == endpoint + "/.well-known/agent-card.json": + return _FakeUpstreamResponse( + status_code=200, + body=json.dumps( + { + "name": "default", + "description": "a2a-agent", + "url": rpc_url, + "version": "0.0.1", + "capabilities": {"streaming": True}, + } + ).encode(), + ) + if url == endpoint + "/list-apps": + return _FakeUpstreamResponse( + status_code=404, + body=b'{"detail":"Not Found"}', + ) + if url == rpc_url: + payload = json.loads(request["content"]) + assert payload["method"] == "message/stream" + frame = { + "jsonrpc": "2.0", + "id": payload["id"], + "result": { + "kind": "status-update", + "taskId": "task-1", + "status": { + "state": "working", + "message": { + "role": "agent", + "parts": [{"kind": "text", "text": "partial"}], + }, + }, + }, + } + return _FakeUpstreamResponse( + status_code=200, + body=b"data: " + json.dumps(frame).encode() + b"\n\n", + content_type="text/event-stream", + ) + raise AssertionError(f"unexpected upstream URL: {url}") + + async def aclose(self) -> None: + pass + + monkeypatch.setattr( + "agentkit.sdk.runtime.client.AgentkitRuntimeClient", + _FakeRuntimeClient, + ) + monkeypatch.setattr("httpx.AsyncClient", _FakeAsyncClient) + + with TestClient(app) as client: + response = client.post( + "/web/runtime-proxy/runtime-1/run_sse?_runtime_region=cn-beijing", + json={ + "app_name": "a2a-default", + "user_id": "user", + "session_id": "sid", + "new_message": {"role": "user", "parts": [{"text": "hello"}]}, + }, + ) + + assert response.status_code == 200 + assert '"text": "partial"' in response.text + assert ( + '"error": "A2A stream ended before a terminal task status ' + '(last state: working)."' in response.text + ) + + +def test_runtime_proxy_reuses_a2a_stream_http_client( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + app = _create_frontend_app(monkeypatch, tmp_path) + endpoint = "https://runtime.example" + rpc_url = endpoint + "/a2a/jsonrpc" + stream_client_ids: list[int] = [] + closed_stream_clients: list[int] = [] + rpc_requests = 0 + + class _FakeRuntimeClient: + def __init__(self, **kwargs: Any) -> None: + del kwargs + + def get_runtime(self, request: Any) -> SimpleNamespace: + del request + return SimpleNamespace( + runtime_id="runtime-1", + name="a2a-agent", + project_name="default", + network_configurations=[ + SimpleNamespace(endpoint=endpoint, network_type="public") + ], + authorizer_configuration=SimpleNamespace( + key_auth=SimpleNamespace(api_key="runtime-api-key"), + custom_jwt_authorizer=None, + ), + tags=[], + ) + + class _FakeUpstreamResponse: + def __init__( + self, + *, + status_code: int, + body: bytes, + content_type: str = "application/json", + ) -> None: + self.status_code = status_code + self._body = body + self.headers = {"content-type": content_type} + + async def aiter_raw(self): + yield self._body + + async def aclose(self) -> None: + pass + + class _FakeAsyncClient: + def __init__(self, **kwargs: Any) -> None: + self.is_stream_client = kwargs.get("limits") is not None + if self.is_stream_client: + stream_client_ids.append(id(self)) + + @property + def is_closed(self) -> bool: + return id(self) in closed_stream_clients + + def build_request( + self, + method: str, + url: str, + *, + params: dict[str, str], + headers: dict[str, str], + content: bytes, + ) -> dict[str, Any]: + return { + "method": method, + "url": url, + "params": params, + "headers": headers, + "content": content, + } + + async def send(self, request: dict[str, Any], *, stream: bool): + nonlocal rpc_requests + del stream + url = request["url"] + if url == endpoint + "/.well-known/agent-card.json": + return _FakeUpstreamResponse( + status_code=200, + body=json.dumps( + { + "name": "default", + "description": "a2a-agent", + "url": rpc_url, + "version": "0.0.1", + "capabilities": {"streaming": True}, + } + ).encode(), + ) + if url == endpoint + "/list-apps": + return _FakeUpstreamResponse( + status_code=404, + body=b'{"detail":"Not Found"}', + ) + if url == rpc_url: + assert self.is_stream_client + rpc_requests += 1 + payload = json.loads(request["content"]) + frame = { + "jsonrpc": "2.0", + "id": payload["id"], + "result": { + "kind": "status-update", + "taskId": "task-1", + "final": True, + "status": { + "state": "completed", + "message": { + "role": "agent", + "parts": [{"kind": "text", "text": "done"}], + }, + }, + }, + } + return _FakeUpstreamResponse( + status_code=200, + body=b"data: " + json.dumps(frame).encode() + b"\n\n", + content_type="text/event-stream", + ) + raise AssertionError(f"unexpected upstream URL: {url}") + + async def aclose(self) -> None: + if self.is_stream_client: + closed_stream_clients.append(id(self)) + + monkeypatch.setattr( + "agentkit.sdk.runtime.client.AgentkitRuntimeClient", + _FakeRuntimeClient, + ) + monkeypatch.setattr("httpx.AsyncClient", _FakeAsyncClient) + + with TestClient(app) as client: + for _ in range(2): + response = client.post( + "/web/runtime-proxy/runtime-1/run_sse?_runtime_region=cn-beijing", + json={ + "app_name": "a2a-default", + "user_id": "user", + "session_id": "sid", + "new_message": { + "role": "user", + "parts": [{"text": "hello"}], + }, + }, + ) + assert response.status_code == 200 + assert '"text": "done"' in response.text + + assert rpc_requests == 2 + assert len(stream_client_ids) == 1 + assert closed_stream_clients == stream_client_ids + + def test_runtime_proxy_resolves_studio_media_before_forwarding( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/tests/cli/test_runtime_a2a_stream.py b/tests/cli/test_runtime_a2a_stream.py index 1f5ca657c..67699d3d3 100644 --- a/tests/cli/test_runtime_a2a_stream.py +++ b/tests/cli/test_runtime_a2a_stream.py @@ -911,7 +911,7 @@ def test_decoder_suppresses_repeated_transport_heartbeats(): assert decoder.project(working, author="default") == [] -def test_does_not_map_submitted_user_echo_or_terminal_status_message(): +def test_does_not_map_submitted_user_echo(): submitted = { "kind": "status-update", "status": { @@ -919,8 +919,16 @@ def test_does_not_map_submitted_user_echo_or_terminal_status_message(): "message": {"role": "user", "parts": [{"kind": "text", "text": "hello"}]}, }, } + + assert a2a_event_to_studio_events(submitted, author="default") == [] + + +@pytest.mark.parametrize("mpa_a2a", [False, True]) +def test_decoder_projects_final_completed_status_message(mpa_a2a): + decoder = A2AStreamDecoder(mpa_a2a=mpa_a2a) completed = { "kind": "status-update", + "taskId": "task-1", "final": True, "status": { "state": "completed", @@ -928,8 +936,96 @@ def test_does_not_map_submitted_user_echo_or_terminal_status_message(): }, } - assert a2a_event_to_studio_events(submitted, author="default") == [] - assert a2a_event_to_studio_events(completed, author="default") == [] + projected = decoder.project(completed, author="default") + + assert projected == [ + { + "id": "a2a-task-1-final-0", + "author": "default", + "partial": False, + "turnComplete": True, + "content": {"role": "model", "parts": [{"text": "done"}]}, + "customMetadata": {"projectionSource": "a2a-terminal-status"}, + } + ] + assert decoder.finalize_projection(author="default") == [] + assert decoder.stream_incomplete_error() == "" + + +@pytest.mark.parametrize("mpa_a2a", [False, True]) +def test_decoder_projects_final_failed_status_message_as_error(mpa_a2a): + decoder = A2AStreamDecoder(mpa_a2a=mpa_a2a) + failed = { + "kind": "status-update", + "taskId": "task-1", + "final": True, + "status": { + "state": "failed", + "message": {"role": "agent", "parts": [{"kind": "text", "text": "boom"}]}, + }, + } + + projected = decoder.project(failed, author="default") + + assert projected == [ + { + "id": "a2a-task-1-failed-error", + "author": "default", + "partial": False, + "content": {"role": "model", "parts": []}, + "error": "A2A task failed: boom", + } + ] + assert decoder.finalize_projection(author="default") == [] + assert decoder.stream_incomplete_error() == "" + + +def test_decoder_suppresses_duplicate_terminal_status_after_final_artifact(): + decoder = A2AStreamDecoder() + assert decoder.project(_outer_result("done"), author="default") + + duplicate = { + "kind": "status-update", + "taskId": "task-1", + "final": True, + "status": { + "state": "completed", + "message": {"role": "agent", "parts": [{"kind": "text", "text": "done"}]}, + }, + } + + assert decoder.project(duplicate, author="default") == [] + assert decoder.stream_incomplete_error() == "" + + +def test_decoder_reports_incomplete_stream_after_partial_without_terminal(): + decoder = A2AStreamDecoder() + decoder.project( + { + "kind": "status-update", + "taskId": "task-1", + "status": { + "state": "working", + "message": { + "role": "agent", + "parts": [{"kind": "text", "text": "partial"}], + }, + }, + }, + author="default", + ) + + assert ( + decoder.stream_incomplete_error() + == "A2A stream ended before a terminal task status (last state: working)." + ) + + +def test_decoder_accepts_completed_projection_without_terminal_status(): + decoder = A2AStreamDecoder() + assert decoder.project(_outer_result("done"), author="default") + + assert decoder.stream_incomplete_error() == "" def test_projection_suppresses_cumulative_working_message_after_deltas(): diff --git a/veadk/cli/cli_frontend.py b/veadk/cli/cli_frontend.py index aabc58afb..a39a252ac 100644 --- a/veadk/cli/cli_frontend.py +++ b/veadk/cli/cli_frontend.py @@ -11552,6 +11552,47 @@ async def _list_region_window(reg: str) -> tuple[list[dict], bool]: _rt_conn_cache: dict[tuple[str, str], tuple[str, str, str, str, float]] = {} _a2a_task_by_session: OrderedDict[tuple[str, str, str, str], str] = OrderedDict() _a2a_task_cache_limit = 1_000 + _runtime_a2a_stream_client: httpx.AsyncClient | None = None + + def _shared_runtime_a2a_stream_client() -> httpx.AsyncClient: + nonlocal _runtime_a2a_stream_client + if _runtime_a2a_stream_client is None or getattr( + _runtime_a2a_stream_client, + "is_closed", + False, + ): + _runtime_a2a_stream_client = httpx.AsyncClient( + timeout=None, + limits=httpx.Limits( + max_connections=100, + max_keepalive_connections=20, + keepalive_expiry=30, + ), + ) + return _runtime_a2a_stream_client + + async def _close_runtime_a2a_stream_client() -> None: + nonlocal _runtime_a2a_stream_client + client, _runtime_a2a_stream_client = _runtime_a2a_stream_client, None + close = getattr(client, "aclose", None) + if close is not None: + await close() + + runtime_router = getattr(app, "router", None) + add_runtime_shutdown = getattr(runtime_router, "add_event_handler", None) + if callable(add_runtime_shutdown): + add_runtime_shutdown("shutdown", _close_runtime_a2a_stream_client) + runtime_a2a_lifespan = app.router.lifespan_context + + @asynccontextmanager + async def _runtime_a2a_stream_client_lifespan(current_app: Any): + async with runtime_a2a_lifespan(current_app): + try: + yield + finally: + await _close_runtime_a2a_stream_client() + + app.router.lifespan_context = _runtime_a2a_stream_client_lifespan def _remember_a2a_task( *, @@ -11906,7 +11947,7 @@ async def _runtime_a2a_run_sse( ) + "\n\n" ).encode("utf-8") - client = httpx.AsyncClient(timeout=None) + client = _shared_runtime_a2a_stream_client() upstream = None try: request = client.build_request( @@ -12010,6 +12051,16 @@ async def _runtime_a2a_run_sse( yield ( "data: " + json.dumps(event, ensure_ascii=False) + "\n\n" ).encode("utf-8") + incomplete_error = decoder.stream_incomplete_error() + if incomplete_error: + yield ( + "data: " + + json.dumps( + {"error": incomplete_error}, + ensure_ascii=False, + ) + + "\n\n" + ).encode("utf-8") return body = await _runtime_proxy_buffer(upstream) if upstream.status_code >= 400: @@ -12080,7 +12131,6 @@ async def _runtime_a2a_run_sse( finally: if upstream is not None: await upstream.aclose() - await client.aclose() def _resolve_runtime_conn( runtime_id: str, diff --git a/veadk/cli/runtime_a2a_stream.py b/veadk/cli/runtime_a2a_stream.py index 336a667d7..fdcc2d66d 100644 --- a/veadk/cli/runtime_a2a_stream.py +++ b/veadk/cli/runtime_a2a_stream.py @@ -24,6 +24,9 @@ from uuid import uuid4 _MAX_EVENT_TEXT_CHARS = 64_000 +_TERMINAL_TASK_STATES = frozenset( + {"completed", "failed", "canceled", "cancelled", "rejected"} +) @dataclass @@ -46,7 +49,10 @@ def __init__(self, *, mpa_a2a: bool = False) -> None: self._heartbeat_states: set[tuple[str, str]] = set() self._usage_snapshots: dict[str, dict[str, int]] = {} self._terminal_state = "" + self._latest_state = "" + self._completed_projection_seen = False self._sandbox_final_answers: dict[str, set[str]] = {} + self._final_answer_texts: dict[str, set[str]] = {} self._projected_sandbox_ids: set[tuple[str, str, str, str]] = set() def feed(self, chunk: str | bytes) -> list[dict[str, Any]]: @@ -96,16 +102,20 @@ def project(self, event: Any, *, author: str) -> list[dict[str, Any]]: else "unknown" ) metadata = event.get("metadata") if isinstance(event, Mapping) else None - if isinstance(event, Mapping) and ( - (event.get("kind") == "status-update" and event.get("final") is True) - or event.get("kind") == "task" - ): + if isinstance(event, Mapping): status = event.get("status") - self._terminal_state = ( + state = ( str(status.get("state") or "").lower() if isinstance(status, Mapping) else "" ) + if state: + self._latest_state = state + if state and ( + (event.get("kind") == "status-update" and event.get("final") is True) + or event.get("kind") == "task" + ): + self._terminal_state = state cumulative_snapshot = ( isinstance(event, Mapping) and event.get("kind") == "status-update" @@ -224,15 +234,27 @@ def project(self, event: Any, *, author: str) -> list[dict[str, Any]]: ): output.append(item) continue + item_metadata = item.get("customMetadata") + answer_texts = _answer_texts_from_event(item) + if ( + item.get("partial") is False + and isinstance(item_metadata, Mapping) + and item_metadata.get("projectionSource") == "a2a-terminal-status" + and answer_texts + and all( + text in self._final_answer_texts.get(task_id, set()) + for text in answer_texts + ) + ): + continue if item.get("partial") is not True: - parts = item.get("content", {}).get("parts", []) - if any( - isinstance(part, Mapping) - and part.get("thought") is not True - and str(part.get("text") or "") - for part in parts - ): + if answer_texts: self._partial_text["answer"] = "" + self._final_answer_texts.setdefault(task_id, set()).update( + answer_texts + ) + if item.get("turnComplete") is True: + self._completed_projection_seen = True output.append(item) continue parts = item.get("content", {}).get("parts", []) @@ -315,6 +337,15 @@ def project(self, event: Any, *, author: str) -> list[dict[str, Any]]: output.append(item) return output + def stream_incomplete_error(self) -> str: + if ( + self._terminal_state in _TERMINAL_TASK_STATES + or self._completed_projection_seen + ): + return "" + state = self._latest_state or "unknown" + return f"A2A stream ended before a terminal task status (last state: {state})." + def _normalize_mpa_reasoning( self, item: dict[str, Any], task_id: str, cumulative: bool ) -> dict[str, Any] | None: @@ -477,23 +508,45 @@ def a2a_event_to_studio_events(event: Any, *, author: str) -> list[dict[str, Any return [] if event.get("kind") == "status-update": status = event.get("status") - if not isinstance(status, Mapping) or event.get("final") is True: + if not isinstance(status, Mapping): + return [] + state = str(status.get("state") or "").lower() + task_id = str(event.get("taskId") or event.get("id") or "unknown") + message = status.get("message") + if event.get("final") is True: + if state == "completed" and isinstance(message, Mapping): + if message.get("role") == "user": + return [] + return _message_to_studio_events( + message, + author=author, + partial=False, + turn_complete=True, + event_id_prefix=f"a2a-{task_id}-final", + projection_source="a2a-terminal-status", + ) + if state in _TERMINAL_TASK_STATES - {"completed"}: + return [ + _terminal_status_error_event( + state=state, + message=message, + author=author, + task_id=task_id, + ) + ] return [] - state = str(status.get("state") or "") if state not in {"submitted", "working"}: return [] - message = status.get("message") if isinstance(message, Mapping) and message.get("role") == "user": return [] if state == "working" and isinstance(message, Mapping): projected = ( - _message_to_partial_events(message, author=author) + _message_to_studio_events(message, author=author, partial=True) if message.get("role") == "agent" else [] ) if projected: return projected - task_id = str(event.get("taskId") or event.get("id") or "unknown") return [ { "id": f"a2a-{task_id}-{state}", @@ -513,6 +566,31 @@ def a2a_event_to_studio_events(event: Any, *, author: str) -> list[dict[str, Any turn_complete=True, ) ) + status = event.get("status") + if isinstance(status, Mapping): + state = str(status.get("state") or "").lower() + task_id = str(event.get("id") or event.get("taskId") or "unknown") + message = status.get("message") + if state == "completed" and isinstance(message, Mapping): + events.extend( + _message_to_studio_events( + message, + author=author, + partial=False, + turn_complete=True, + event_id_prefix=f"a2a-{task_id}-final", + projection_source="a2a-terminal-status", + ) + ) + elif state in _TERMINAL_TASK_STATES - {"completed"}: + events.append( + _terminal_status_error_event( + state=state, + message=message, + author=author, + task_id=task_id, + ) + ) return events if event.get("kind") != "artifact-update": return [] @@ -595,11 +673,47 @@ def coalesce(container: Any, *, snapshot: bool | None) -> Any: return event -def _message_to_partial_events( - message: Mapping[str, Any], *, author: str +def _answer_texts_from_event(event: Mapping[str, Any]) -> list[str]: + parts = event.get("content", {}).get("parts", []) + if not isinstance(parts, list): + return [] + texts: list[str] = [] + for part in parts: + if not isinstance(part, Mapping) or part.get("thought") is True: + continue + text = str(part.get("text") or "").strip() + if text: + texts.append(text) + return texts + + +def _message_visible_text(message: Any) -> str: + if not isinstance(message, Mapping) or message.get("role") == "user": + return "" + texts: list[str] = [] + for part in message.get("parts") or []: + if not isinstance(part, Mapping): + continue + metadata = part.get("metadata") + if isinstance(metadata, Mapping) and metadata.get("adk_thought") is True: + continue + text = str(part.get("text") or "").strip() + if text: + texts.append(text[:_MAX_EVENT_TEXT_CHARS]) + return "\n".join(texts)[:_MAX_EVENT_TEXT_CHARS] + + +def _message_to_studio_events( + message: Mapping[str, Any], + *, + author: str, + partial: bool, + turn_complete: bool = False, + event_id_prefix: str = "", + projection_source: str = "a2a-status", ) -> list[dict[str, Any]]: events: list[dict[str, Any]] = [] - message_id = str(message.get("messageId") or uuid4()) + message_id = event_id_prefix or str(message.get("messageId") or uuid4()) for index, part in enumerate(message.get("parts") or []): if not isinstance(part, Mapping): continue @@ -608,30 +722,47 @@ def _message_to_partial_events( text = str(part.get("text") or "")[:_MAX_EVENT_TEXT_CHARS] if not text: continue + custom_metadata: dict[str, Any] = {} + if thought or projection_source != "a2a-status": + custom_metadata["projectionSource"] = projection_source + if thought: + custom_metadata["thoughtKind"] = "reasoning" + if ( + isinstance(metadata, Mapping) + and metadata.get("reasoningSnapshot") is True + ): + custom_metadata["reasoningSnapshot"] = True events.append( _text_event( text, author=author, event_id=f"{message_id}-{index}", - partial=True, + partial=partial, + turn_complete=turn_complete, thought=thought, - custom_metadata={ - "thoughtKind": "reasoning", - "projectionSource": "a2a-status", - **( - {"reasoningSnapshot": True} - if isinstance(metadata, Mapping) - and metadata.get("reasoningSnapshot") is True - else {} - ), - } - if thought - else None, + custom_metadata=custom_metadata or None, ) ) return events +def _terminal_status_error_event( + *, + state: str, + message: Any, + author: str, + task_id: str, +) -> dict[str, Any]: + detail = _message_visible_text(message) + return { + "id": f"a2a-{task_id}-{state}-error", + "author": author, + "partial": False, + "content": {"role": "model", "parts": []}, + "error": f"A2A task {state}: {detail}" if detail else f"A2A task {state}.", + } + + def _artifact_to_studio_events( artifact: Any, *, author: str, turn_complete: bool ) -> list[dict[str, Any]]: