from __future__ import annotations import functools import json import locale import logging import os import re import urllib.error import urllib.request import uuid from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple from api_clients.A2A.constants import ( ARTIFACT_DESCRIPTORS, AcceptedOutputMode, ) from api_clients.base_agent_client import BaseAgentClient from cli_logging.cli_logger import emit_structured_log as emit_cli_log from cli_logging.console_diagnostics import emit_structured_log from cli_logging.logging_utils import Operation, build_run_context from retry_policy import is_timeout_error from api_clients.user_agent import build_user_agent from api_clients.A2A.protocol import ( A2A_VERSION_HEADER, DEFAULT_PROTOCOL_VERSION, PROTOCOL_VERSION_ERROR_CODES, ProtocolProfile, get_profile, join_text_parts, normalize_task_state, part_data, protocol_version_error, unwrap_send_result, ) from error_messages import ( AgentRequestError, ERROR_CODE_AGENT_REQUEST_FAILED, ERROR_CODE_AGENT_AUTH, ERROR_CODE_AGENT_RATE_LIMITED, ERROR_CODE_AGENT_TIMEOUT, ERROR_CODE_AGENT_RESPONSE_UNPARSEABLE, ) def _agent_code_for_status(status: "int | None") -> str: """Map an HTTP status to the matching agent-request error code.""" if status == 429: return ERROR_CODE_AGENT_RATE_LIMITED if status in (401, 403): return ERROR_CODE_AGENT_AUTH return ERROR_CODE_AGENT_REQUEST_FAILED # Feature flag required by the experimental A2A surface. _A2A_FEATURE_FLAG = "feature.EnableA2AServer" # Default HTTP request timeout (seconds). The env-configurable value (resolved # via ``_resolve_request_timeout_secs`` from ``WORKIQ_REQUEST_TIMEOUT_SECS``) # applies only to the retryable message-send path. The agent-discovery and # agent-card-resolution calls use this fixed default directly because they are # not yet covered by a retry mechanism (see issue #406). _DEFAULT_REQUEST_TIMEOUT_SECS = 300 _ENV_REQUEST_TIMEOUT_SECS = "WORKIQ_REQUEST_TIMEOUT_SECS" def _resolve_request_timeout_secs() -> float: """Resolve the effective request timeout in seconds. Honors the ``WORKIQ_REQUEST_TIMEOUT_SECS`` environment variable when it is set to a positive number; otherwise falls back to ``_DEFAULT_REQUEST_TIMEOUT_SECS``. """ raw = os.environ.get(_ENV_REQUEST_TIMEOUT_SECS) if raw is None or not raw.strip(): return _DEFAULT_REQUEST_TIMEOUT_SECS try: value = float(raw) except (TypeError, ValueError): emit_cli_log( "warning", f"Ignoring invalid {_ENV_REQUEST_TIMEOUT_SECS}={raw!r}; " f"expected a positive number of seconds. " f"Using default {_DEFAULT_REQUEST_TIMEOUT_SECS:g}s.", Operation.SETUP, ) return _DEFAULT_REQUEST_TIMEOUT_SECS if value <= 0: emit_cli_log( "warning", f"Ignoring non-positive {_ENV_REQUEST_TIMEOUT_SECS}={raw!r}; " f"expected a positive number of seconds. " f"Using default {_DEFAULT_REQUEST_TIMEOUT_SECS:g}s.", Operation.SETUP, ) return _DEFAULT_REQUEST_TIMEOUT_SECS return value # Lookup derived from ARTIFACT_DESCRIPTORS for use in the response parser: # server returns artifacts by name, we expose them under a stable key. _ARTIFACT_NAME_TO_RESULT_KEY: Dict[str, str] = dict(ARTIFACT_DESCRIPTORS.values()) # OAI citation marker patterns — compiled once at module level. # Marker format: \ue200cite(\ue202turn{X}search{Y})+\ue201 _CITATION_REF_PAT = re.compile(r"\ue202turn\d+search(\d+)") _CITATION_BLOCK_PAT = re.compile(r"\ue200cite(?:\ue202turn\d+search\d+)+\ue201") # WorkIQ response header carrying the server-side correlation id. _REQUEST_ID_HEADER = "request-id" def _extract_request_id(headers: Any) -> Optional[str]: """Return the ``request-id`` header value from an HTTP headers-like object. Accepts anything exposing a case-insensitive ``get`` (e.g. an ``http.client.HTTPMessage``) and returns the stripped ``request-id`` value, or ``None`` when it is absent, empty, or the object cannot be read. """ getter = getattr(headers, "get", None) if not callable(getter): return None try: value = getter(_REQUEST_ID_HEADER) except Exception: return None if not value: return None return str(value).strip() or None class A2AClient(BaseAgentClient): """A2A (Agent-to-Agent) JSON-RPC 2.0 client for Work IQ agents.""" def __init__( self, *, a2a_endpoint: str, access_token: str, logger: Optional[logging.Logger] = None, diagnostic_records: Optional[List[Dict[str, Any]]] = None, token_refresh_fn: Optional[Callable[[], str]] = None, protocol_version: str = DEFAULT_PROTOCOL_VERSION, ) -> None: """ Args: a2a_endpoint: Base URL of the A2A endpoint. access_token: Bearer token for A2A authentication. logger: Logger to use. Defaults to a module-level logger if not provided. diagnostic_records: List to accumulate structured log entries. token_refresh_fn: Optional callable that returns a fresh access token string. When provided, a single HTTP 401 response will trigger a token refresh and one automatic retry, making the refresh invisible to the caller. protocol_version: A2A wire revision to speak. Defaults to :data:`~api_clients.A2A.protocol.DEFAULT_PROTOCOL_VERSION`. ``"0.3"`` selects the legacy dialect as an emergency rollback lever; it never engages on its own. """ self._endpoint = a2a_endpoint.rstrip("/") self._access_token = access_token self._logger = logger or logging.getLogger(__name__) self._diagnostic_records = diagnostic_records self._token_refresh_fn = token_refresh_fn self._protocol: ProtocolProfile = get_profile(protocol_version) self._resolved_agent_url: Optional[str] = None self._missing_artifact_warned: set[Tuple[Optional[str], str]] = set() self._request_timeout_secs = _resolve_request_timeout_secs() if not self._protocol.is_default: emit_structured_log( "warning", f"[A2A] Using legacy protocol version {self._protocol.version}. " f"This is a rollback lever and is expected to be removed once " f"{DEFAULT_PROTOCOL_VERSION} is proven.", Operation.SETUP, logger=self._logger, diagnostic_records=self._diagnostic_records, ) emit_structured_log( "debug", f"[A2A] Speaking protocol version {self._protocol.version}.", Operation.SETUP, logger=self._logger, diagnostic_records=self._diagnostic_records, ) emit_structured_log( "info", f"Request timeout set to {self._request_timeout_secs:g}s.", Operation.SETUP, logger=self._logger, diagnostic_records=self._diagnostic_records, ) # ------------------------------------------------------------------ # # BaseAgentClient implementation # # ------------------------------------------------------------------ # def resolve_agent(self, agent_id: str) -> None: """Pre-resolve agent URL from agent card and cache for the session.""" self._resolved_agent_url = self._resolve_agent_url(agent_id) def fetch_available_agents(self) -> List[Dict[str, Any]]: """Fetch agents from the A2A discovery endpoint. Calls GET {endpoint}/.agents. Each A2A agent card is normalized to include 'gptId', 'name', and 'provider' so it is compatible with the shared select_agent_interactively selector. Returns an empty list if the endpoint is unreachable or returns an error. """ try: agents_url = f"{self._endpoint}/.agents" headers = self._build_request_headers() emit_structured_log( "debug", f"[A2A] Fetching available agents from: {agents_url}", Operation.FETCH_AGENTS, logger=self._logger, diagnostic_records=self._diagnostic_records, ) req = urllib.request.Request(agents_url, headers=headers, method="GET") # Fixed timeout (not env-configurable) and not retried; see issue #406. with urllib.request.urlopen(req, timeout=_DEFAULT_REQUEST_TIMEOUT_SECS) as resp: agents = json.loads(resp.read().decode("utf-8")) emit_structured_log( "debug", f"[A2A] Available agents response: {json.dumps(agents)}", Operation.FETCH_AGENTS, logger=self._logger, diagnostic_records=self._diagnostic_records ) return [self._normalize_agent_card(a) for a in agents] except urllib.error.HTTPError as e: emit_structured_log( "warning", f"[A2A] Unable to fetch agents list (HTTP {e.code}).", Operation.FETCH_AGENTS, logger=self._logger, diagnostic_records=self._diagnostic_records, ) return [] except Exception as e: emit_structured_log( "warning", f"[A2A] Error fetching agents: {e}", Operation.FETCH_AGENTS, logger=self._logger, diagnostic_records=self._diagnostic_records, ) return [] @staticmethod def _normalize_agent_card(agent: Dict[str, Any]) -> Dict[str, Any]: """Normalize an A2A agent card to the shape expected by the selector. """ return { "gptId": agent.get("agentId"), "name": agent.get("name"), "provider": agent.get("provider") } def send_prompt( self, prompt: str, agent_id: str | None = None, conversation_context: Optional[Dict[str, Any]] = None, *, accepted_output_modes: Optional[Iterable[AcceptedOutputMode]] = None, ) -> Tuple[Dict[str, Any], Optional[Dict[str, Any]]]: """Send a single prompt to the A2A endpoint and return the response with context. Args: prompt: The prompt string to send. agent_id: Target agent ID. Required for A2A — no fallback discovery. conversation_context: Context from a previous turn (contains context_id), or None for the first turn / single-turn usage. accepted_output_modes: A2A media types for ``configuration.acceptedOutputModes`` (use constants from :mod:`api_clients.A2A.constants`). ``None`` or empty omits the configuration block; callers compose this list per-prompt so non-retrieval prompts stay on the lighter payload path. Returns: Tuple of (enhanced_response_dict, conversation_context). """ agent_id = (agent_id or "").strip() if not agent_id: raise ValueError("agent_id is required for A2A requests.") headers = self._build_request_headers(include_content_type=True) agent_url = self._resolved_agent_url or self._resolve_agent_url(agent_id) context_id = conversation_context.get("context_id") if conversation_context else None emit_structured_log( "debug", "[A2A] Sending prompt to agent.", Operation.SEND_PROMPT, logger=self._logger, diagnostic_records=self._diagnostic_records, ) # Materialize once: needed for both the payload and the # post-response missing-artifact check. requested_modes = list(accepted_output_modes) if accepted_output_modes else [] payload = self._build_chat_payload( prompt, context_id, accepted_output_modes=requested_modes, ) emit_structured_log( "debug", f"[A2A] Sending to {agent_url}: {payload.decode('utf-8')}", Operation.SEND_PROMPT, logger=self._logger, diagnostic_records=self._diagnostic_records, ) result_dict, raw_result = self._send_and_parse_message(agent_url, payload, headers) # Surface missing opted-in artifacts. First occurrence per # (agent_id, mode) is a warning; subsequent occurrences drop to # debug to avoid a flood when an agent doesn't yet support the # requested mode across a large prompt set. for mode in requested_modes: descriptor = ARTIFACT_DESCRIPTORS.get(mode) if descriptor is None: continue _, result_key = descriptor if result_key not in result_dict: seen_key = (agent_id, mode) level = "debug" if seen_key in self._missing_artifact_warned else "warning" self._missing_artifact_warned.add(seen_key) emit_structured_log( level, f"[A2A] Requested output mode '{mode}' but no matching artifact was returned.", Operation.SEND_PROMPT, logger=self._logger, diagnostic_records=self._diagnostic_records, ) # Build updated context for subsequent turns new_context_id = context_id if raw_result: new_context_id = raw_result.get("contextId") or context_id updated_context = {"context_id": new_context_id} if new_context_id else None return result_dict, updated_context # ------------------------------------------------------------------ # # Private helpers # # ------------------------------------------------------------------ # def _build_request_headers(self, *, include_content_type: bool = False) -> Dict[str, str]: headers: Dict[str, str] = { "Authorization": f"Bearer {self._access_token}", "X-variants": _A2A_FEATURE_FLAG, "User-Agent": build_user_agent(), # Declares the protocol revision this client speaks so a server # doing version negotiation answers in the same revision. A2A_VERSION_HEADER: self._protocol.version, } if include_content_type: headers["Content-Type"] = "application/json" return headers def _build_chat_payload( self, prompt: str, context_id: str | None = None, *, accepted_output_modes: Optional[Iterable[AcceptedOutputMode]] = None, ) -> bytes: message: Dict[str, Any] = self._protocol.message_envelope( [self._protocol.text_part(prompt)] ) message["messageId"] = str(uuid.uuid4()) message["metadata"] = {"location": self._get_a2a_location()} if context_id: message["contextId"] = context_id params: Dict[str, Any] = {"message": message} # Per-prompt opt-in: omit the configuration block on default prompts # to keep response payloads small. # # ``dict.fromkeys`` dedupes while # preserving order, matching A2A's preference-ordered # ``acceptedOutputModes`` semantics. modes = list(dict.fromkeys(accepted_output_modes or ())) if modes: params["configuration"] = self._protocol.send_configuration(modes) return json.dumps({ "jsonrpc": "2.0", "method": self._protocol.send_method, "params": params, "id": str(uuid.uuid4()), }).encode("utf-8") @staticmethod @functools.lru_cache(maxsize=1) def _get_a2a_location() -> Dict[str, Any]: locale_str = locale.getlocale()[0] or "" country = locale_str.split("_")[-1] if "_" in locale_str else "" return { "countryOrRegion": country, "countryOrRegionConfidence": 1.0, "timeZone": BaseAgentClient._get_iana_timezone_name(), } def _resolve_agent_url(self, agent_id: str) -> str: """Resolve the agent URL from the agent card, falling back to base URL. The card supplies the URL only. It is deliberately *not* consulted for protocol-version compatibility: an endpoint can advertise a stale version while already serving a newer one, so a card-based veto would block endpoints that work. A genuine mismatch surfaces on the first send as JSON-RPC ``-32009`` or ``-32601``; see :func:`~api_clients.A2A.protocol.protocol_version_error`. """ headers = self._build_request_headers(include_content_type=True) base_agent_url = f"{self._endpoint}/{agent_id}" card_url = f"{base_agent_url}/.well-known/agent-card.json" agent_url = base_agent_url emit_structured_log( "debug", f"[A2A] Fetching agent card from: {card_url}", Operation.FETCH_AGENTS, logger=self._logger, diagnostic_records=self._diagnostic_records, ) try: card_req = urllib.request.Request(card_url, headers=headers) # Fixed timeout (not env-configurable) and not retried; see issue #406. with urllib.request.urlopen(card_req, timeout=_DEFAULT_REQUEST_TIMEOUT_SECS) as resp: raw_card = resp.read().decode("utf-8") if raw_card.strip(): card = json.loads(raw_card) card_url_value = card.get("url") if isinstance(card_url_value, str) and card_url_value.strip(): agent_url = card_url_value.strip() except (urllib.error.HTTPError, urllib.error.URLError, json.JSONDecodeError, UnicodeDecodeError) as e: emit_structured_log( "debug", f"[A2A] Agent card fetch failed ({e}); using base URL: {base_agent_url}", Operation.FETCH_AGENTS, logger=self._logger, diagnostic_records=self._diagnostic_records, ) emit_structured_log( "debug", f"[A2A] Resolved agent URL: {agent_url}", Operation.SEND_PROMPT, logger=self._logger, diagnostic_records=self._diagnostic_records, ) return agent_url def _send_and_parse_message( self, agent_url: str, payload: bytes, headers: Dict[str, str], ) -> tuple[Dict[str, Any], Dict[str, Any]]: """Send the prompt, parse the response, and always log A2A correlation. Delegates to :meth:`_send_and_parse_core`, which fills ``correlation`` with the ``request-id`` (response header), ``conversation-id`` (the A2A ``contextId``), and ``response-timestamp`` (``status.timestamp``) as they become available. A ``finally`` block then emits a single info line for every outcome — success or failure — so each A2A request can be correlated to a server-side log even when it fails. """ correlation: Dict[str, Optional[str]] = { "request_id": None, "conversation_id": None, "response_timestamp": None, } try: return self._send_and_parse_core(agent_url, payload, headers, correlation) finally: self._emit_correlation_log(correlation) def _emit_correlation_log(self, correlation: Dict[str, Optional[str]]) -> None: """Emit a single ``info`` correlation record for the A2A request. The ``request-id``, ``conversation-id``, and ``response-timestamp`` are threaded through ``run_context`` so they populate the structured fields on the diagnostic record. ``info`` level keeps the ids visible at the default verbosity so failures can be correlated to server-side logs. """ emit_structured_log( "info", "[A2A] Response correlation", Operation.SEND_PROMPT, logger=self._logger, diagnostic_records=self._diagnostic_records, run_context=build_run_context( operation=Operation.SEND_PROMPT, request_id=correlation["request_id"], conversation_id=correlation["conversation_id"], response_timestamp=correlation["response_timestamp"], ), ) def _send_and_parse_core( self, agent_url: str, payload: bytes, headers: Dict[str, str], correlation: Dict[str, Optional[str]], ) -> tuple[Dict[str, Any], Dict[str, Any]]: """Send a JSON-RPC message to the agent and parse the response. When a ``token_refresh_fn`` was supplied at construction time and the server responds with HTTP 401 (Unauthorized), the token is refreshed automatically and the request is retried exactly once. This keeps long-running eval sessions alive beyond the initial token lifetime without requiring any user interaction. Transient connection errors (including socket timeouts) are surfaced as an :class:`~error_messages.AgentRequestError` carrying an ``agent*`` ``code``, with the original exception chained via ``from`` so the higher-level evaluation runner can detect timeouts and retry the send. As ids become available they are written into ``correlation`` so the caller can log them regardless of whether this method returns or raises. Returns: A tuple of (result_dict, raw_result) where result_dict is the normalized response dict (raw_response_text, display_response_text, a2a_attributions, and a metadata block carrying ``conversation_id``) and raw_result is the parsed JSON result object. Raises: AgentRequestError: On HTTP errors, connection errors, request timeouts, JSON parse errors, or A2A protocol errors. The ``code`` attribute categorizes the failure (``agentTimeout`` / ``agentRateLimited`` / ``agentAuthenticationError`` / ``agentResponseUnparseable`` / ``agentRequestFailed``). """ timeout = self._request_timeout_secs req = urllib.request.Request(agent_url, data=payload, headers=headers, method="POST") try: with urllib.request.urlopen(req, timeout=timeout) as resp: correlation["request_id"] = _extract_request_id(resp.headers) raw = resp.read().decode("utf-8", errors="replace") except urllib.error.HTTPError as e: correlation["request_id"] = _extract_request_id(e.headers) if e.code == 401 and self._token_refresh_fn is not None: emit_structured_log( "info", "[A2A] Access token expired (HTTP 401); refreshing token and retrying.", Operation.AUTHENTICATE, logger=self._logger, diagnostic_records=self._diagnostic_records, ) new_token = self._token_refresh_fn() if not new_token: raise AgentRequestError( "A2A request failed (HTTP 401 Unauthorized) and token refresh returned no token.", code=ERROR_CODE_AGENT_AUTH, ) from e self._access_token = new_token headers["Authorization"] = f"Bearer {self._access_token}" retry_req = urllib.request.Request( agent_url, data=payload, headers=headers, method="POST" ) try: with urllib.request.urlopen(retry_req, timeout=timeout) as resp: correlation["request_id"] = _extract_request_id(resp.headers) raw = resp.read().decode("utf-8", errors="replace") except urllib.error.HTTPError as retry_e: correlation["request_id"] = ( _extract_request_id(retry_e.headers) or correlation["request_id"] ) body = "" try: body = retry_e.read().decode("utf-8", errors="replace") except Exception: pass raise AgentRequestError( f"A2A request failed (HTTP {retry_e.code} {retry_e.reason}) after token refresh." + (f" Body: {body}" if body else ""), code=_agent_code_for_status(retry_e.code), ) from retry_e except (urllib.error.URLError, TimeoutError) as retry_e: if is_timeout_error(retry_e): raise AgentRequestError( f"A2A request timed out after {timeout:g}s waiting for the " f"agent's response, even after refreshing the token " f"(set {_ENV_REQUEST_TIMEOUT_SECS} to allow more time).", code=ERROR_CODE_AGENT_TIMEOUT, ) from retry_e raise AgentRequestError( f"A2A connection error after token refresh: {getattr(retry_e, 'reason', str(retry_e))}" ) from retry_e else: body = "" try: body = e.read().decode("utf-8", errors="replace") except Exception: pass raise AgentRequestError( f"A2A request failed (HTTP {e.code} {e.reason})." + (f" Body: {body}" if body else ""), code=_agent_code_for_status(e.code), ) from e except (urllib.error.URLError, TimeoutError) as e: if is_timeout_error(e): raise AgentRequestError( f"A2A request timed out after {timeout:g}s waiting for the " f"agent's response (set {_ENV_REQUEST_TIMEOUT_SECS} to allow " f"more time).", code=ERROR_CODE_AGENT_TIMEOUT, ) from e raise AgentRequestError( f"A2A connection error: {getattr(e, 'reason', str(e))}" ) from e emit_structured_log( "debug", f"[A2A] Raw response: {raw}", Operation.SEND_PROMPT, logger=self._logger, diagnostic_records=self._diagnostic_records, ) try: data = json.loads(raw) except json.JSONDecodeError as e: raise AgentRequestError( f"A2A response is not valid JSON: {e}", code=ERROR_CODE_AGENT_RESPONSE_UNPARSEABLE, ) from e if "error" in data: err = data["error"] # A server that rejects the version outright (-32009), or does not # know the method belonging to the version we asked for (-32601), # is telling us it will not serve this revision. That deserves a # message naming the mismatch rather than a generic protocol error. if err.get("code") in PROTOCOL_VERSION_ERROR_CODES: raise protocol_version_error( agent_url, self._protocol, code=err["code"], server_message=err.get("message"), ) raise AgentRequestError( f"A2A JSON-RPC error {err.get('code', 'unknown')}: {err.get('message', 'no message')}", code=ERROR_CODE_AGENT_RESPONSE_UNPARSEABLE, ) if "result" not in data: raise AgentRequestError( f"A2A response missing 'result' key. Keys present: {list(data.keys())}", code=ERROR_CODE_AGENT_RESPONSE_UNPARSEABLE, ) result = data["result"] # 1.0 wraps the SendMessage result in a ``task`` / ``message`` member # instead of tagging it with a ``kind`` field. payload_kind, payload = unwrap_send_result(result) # Correlation ids for downstream logging: the conversation id is the A2A # ``contextId`` and the server timestamp lives under ``status.timestamp`` # for task responses. Written to ``correlation`` so the caller logs them # on every outcome. correlation["conversation_id"] = payload.get("contextId") or result.get("contextId") correlation["response_timestamp"] = (payload.get("status") or {}).get("timestamp") text = "" attributions: List[Dict[str, Any]] = [] # Map of result-dict key → extracted data dict for any A2A artifact # listed in :data:`constants.ARTIFACT_NAME_TO_RESULT_KEY`. Extensible: # add a new entry to that map (and a new media-type constant if the # artifact requires opt-in) and this loop picks it up without further # changes. extracted_artifacts: Dict[str, Dict[str, Any]] = {} if payload_kind == "message": text = join_text_parts(payload.get("parts", [])) attributions = payload.get("metadata", {}).get("attributions", []) elif payload_kind == "task": raw_state = payload.get("status", {}).get("state") # Normalizes the 1.0 ``TASK_STATE_`` prefix and the hyphenated # spelling into a single comparable form. state = normalize_task_state(raw_state) if state == "completed": msg = payload.get("status", {}).get("message") or {} all_parts = list(msg.get("parts", [])) for artifact in payload.get("artifacts", []): artifact_name = artifact.get("name") if artifact_name in _ARTIFACT_NAME_TO_RESULT_KEY: data = self._extract_artifact_part_data( artifact, artifact_name ) if data is not None: result_key = _ARTIFACT_NAME_TO_RESULT_KEY[artifact_name] extracted_artifacts[result_key] = data all_parts.extend(artifact.get("parts", [])) text = join_text_parts(all_parts) attributions = msg.get("metadata", {}).get("attributions", []) elif state in ("failed", "canceled", "rejected"): status_msg = payload.get("status", {}).get("message") or {} detail = join_text_parts(status_msg.get("parts", [])).strip() suffix = f" Detail: {detail}" if detail else "" raise AgentRequestError( f"A2A task {state}. Task id: {payload.get('id')}{suffix}" ) elif state in ("input_required", "auth_required"): requirement = { "input_required": "user input", "auth_required": "authentication", }.get(state, state.replace("_", " ")) raise AgentRequestError( f"A2A task requires {requirement} and cannot proceed automatically." f" Task id: {payload.get('id')}", code=( ERROR_CODE_AGENT_AUTH if state == "auth_required" else ERROR_CODE_AGENT_REQUEST_FAILED ), ) elif state in ("submitted", "working"): raise AgentRequestError( f"A2A task is still {state}; synchronous send returned before completion." f" Task id: {payload.get('id')}" ) else: raise AgentRequestError( f"A2A task in unexpected state: {raw_state!r}. Task id: {payload.get('id')}" ) else: raise AgentRequestError( f"Unexpected A2A result payload; expected a 'task' or 'message' member. " f"Keys present: {sorted(result.keys()) if isinstance(result, dict) else type(result).__name__}", code=ERROR_CODE_AGENT_RESPONSE_UNPARSEABLE, ) display_text = self._replace_citation_markers(text, attributions) result_dict: Dict[str, Any] = { "raw_response_text": text, "display_response_text": display_text, "a2a_attributions": attributions, "metadata": { "conversation_id": correlation["conversation_id"], }, } for result_key, data in extracted_artifacts.items(): result_dict[result_key] = data # Return the unwrapped Task/Message, not the 1.0 oneof wrapper, so # callers keep reading ``contextId`` straight off it. return result_dict, payload @staticmethod def _extract_artifact_part_data( artifact: Dict[str, Any], artifact_name: str, ) -> Optional[Dict[str, Any]]: """Extract the first data payload from an artifact matching the given name. Returns the ``data`` field of the first data part when ``artifact["name"] == artifact_name`` and the data is a dict, or None otherwise. Parts are discriminated by member presence (1.0) rather than by a ``kind`` field. """ if artifact.get("name") != artifact_name: return None for part in artifact.get("parts", []): data = part_data(part) if data is not None: return data return None @staticmethod def _replace_citation_markers( text: str, attributions: List[Dict[str, Any]], ) -> str: """Replace OAI Unicode citation markers with markdown links. Marker format: \\ue200cite(\\ue202turn{X}search{Y})+\\ue201 Compound markers (multiple turn/search refs between a single pair of bookend characters) are also handled. The search{Y} number is NOT a direct array index — it's a grounding result number. Mapping: unique search numbers in first-appearance order → citation_attrs[0, 1, ...]. Args: text: Response text that may contain OAI citation markers. attributions: Attribution objects from A2A response metadata. Returns: Text with citation markers replaced by markdown links, or the original text unchanged if there are no citation attributions. """ citation_attrs = [a for a in attributions if a.get("attributionType") == "citation"] if not text: return text if not citation_attrs: return text # Build ordered map: search-number-string → 0-based index into citation_attrs seen: Dict[str, int] = {} for m in _CITATION_REF_PAT.finditer(text): k = m.group(1) if k not in seen: seen[k] = len(seen) def replace_citation(m: re.Match) -> str: links = [] for idx_str in _CITATION_REF_PAT.findall(m.group(0)): pos = seen.get(idx_str) if pos is not None and pos < len(citation_attrs): attr = citation_attrs[pos] url = attr.get("seeMoreWebUrl") or "" label = attr.get("providerDisplayName") or url or idx_str if url: links.append(f"[{label}]({url})") return " ".join(links) return _CITATION_BLOCK_PAT.sub(replace_citation, text)