"""Shared A2A request-context utilities. How it works ------------ pydantic-ai's ``AgentWorker.run_task`` calls ``agent.run(message_history=...)`` without any ``deps=`` argument. To give ``@tool`` functions a properly-typed ``RunContext[RequestContext]``, we need deps to be injected at that call site. We do this with two cooperating pieces: 1. **``MetadataAwareWorker``** — subclasses ``AgentWorker``, overrides ``run_task`` to extract ``Message.metadata`` from the A2A params and store a ``RequestContext`` in a task-local ``ContextVar`` before delegating to the parent ``run_task``. 2. **``_DepsInjectingAgent``** — a thin wrapper around the raw pydantic-ai ``Agent`` that intercepts every ``agent.run()`` call and injects ``deps=_request_context_var.get()`` automatically. This is what ``AgentWorker`` holds as ``self.agent``. The net effect: ``AgentWorker.run_task`` calls ``wrapped_agent.run(...)``, which injects the ``RequestContext`` that ``MetadataAwareWorker`` set in the ContextVar — so every ``@tool`` function receives a fully-populated ``RunContext[RequestContext]`` containing the caller's metadata. Thread / concurrency safety --------------------------- fasta2a's ``Worker._loop`` is sequential — it awaits each task to completion before pulling the next. There is never true concurrency at the worker level. Even if fasta2a moved to per-task concurrency, ``ContextVar`` is asyncio-task-local: each spawned task inherits an **independent copy** of the context snapshot at creation time. We also call ``reset(token)`` after each task for explicit cleanup. Usage ----- 1. Replace ``agent.to_a2a(...)`` with ``make_metadata_aware_app(agent, ...)``. 2. Set ``deps_type=RequestContext`` on the ``EnterpriseAgent``:: agent = EnterpriseAgent(..., deps_type=RequestContext) 3. Use ``@agent.tool`` (not ``tool_plain``) with ``RunContext`` as first arg:: from pydantic_ai import RunContext from backend.chanakya.a2a_context import RequestContext @my_agent.tool def my_tool(ctx: RunContext[RequestContext], query: str) -> str: user = ctx.deps.user_email # typed, no .get() needed ... 4. When Chanakya forwards context to a sub-agent:: @chanakya.tool async def ask_jira(ctx: RunContext[RequestContext], question: str) -> str: return await _ask_agent(JIRA_URL, question, ctx=ctx.deps.to_metadata()) """ from __future__ import annotations import asyncio from contextvars import ContextVar from dataclasses import dataclass, field, replace from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any _IMPERSONATION_FILE = Path("/tmp/eb_dev_impersonation") def _read_impersonation() -> str: try: return _IMPERSONATION_FILE.read_text().strip() except FileNotFoundError: return "" from fasta2a.broker import InMemoryBroker from fasta2a.storage import InMemoryStorage from pydantic_ai._a2a import AgentWorker, worker_lifespan if TYPE_CHECKING: from fasta2a.schema import TaskSendParams # --------------------------------------------------------------------------- # Public dataclass — the typed deps object your tools receive # --------------------------------------------------------------------------- @dataclass class RequestContext: """Per-request metadata propagated from the A2A ``Message.metadata`` dict. Accessible inside any ``@agent.tool`` via ``ctx.deps``. The *runtime channel* fields below (``conversation_id``, ``tatasteel_result``, ``ui_result``, ``progress_queue``) are **not** part of the A2A metadata wire format; they are populated by tools and HTTP handlers during ``agent.run()`` and read back after it completes. They replace the previous stringly-keyed ``raw["_*"]`` side-channels. """ user_email: str = "" tenant: str = "" role: str = "" # ----------------------------------------------------------------------- # Runtime channels — populated by HTTP handlers, not from A2A metadata # ----------------------------------------------------------------------- #: Active conversation for history loading and persistence. conversation_id: int | None = field(default=None) #: Payload written by ask_tatasteel; read by ask_ui_agent + server handler. tatasteel_result: list[dict[str, Any]] | None = field(default=None) #: First/primary tatasteel payload for the user's main question. tatasteel_primary_result: dict[str, Any] | None = field(default=None) #: Additional tatasteel payloads gathered for decision support. tatasteel_decision_support_results: list[dict[str, Any]] = field(default_factory=list) #: Generated SQL from the most recent tatasteel query; injected into done event for display. tatasteel_generated_sql: str | None = field(default=None) #: Plain-English explanation of what the SQL does and any translation assumptions. tatasteel_sql_explanation: str | None = field(default=None) #: Full log of every query run by tatasteel: [{sql, explanation}, ...]. tatasteel_query_log: list[dict[str, str]] = field(default_factory=list) #: Structured web claims from the most recent web_search call; injected into done event. web_claims: list[dict[str, str]] = field(default_factory=list) #: Bullet-point market context text from the most recent web search enrichment. market_context_text: str = "" #: done event from UI agent; stored by ask_ui_agent for DB save and sync handler response. ui_result: dict[str, Any] | None = field(default=None) #: SSE progress queue injected by the stream handler. progress_queue: asyncio.Queue | None = field(default=None) #: Chronological list of tool execution events — consumed by _emit_flow_log. flow_log: list[dict[str, Any]] = field(default_factory=list) #: Pre-loaded conversation history serialized as [{role, content}] — set by #: the HTTP handler so ask_tatasteel and ask_ui_agent can pass it directly #: in the sub-agent HTTP request instead of each process loading from DB. message_history_wire: list[dict[str, str]] | None = field(default=None) #: Full original A2A metadata dict — for any keys not listed above. raw: dict[str, Any] = field(default_factory=dict) @classmethod def from_metadata(cls, metadata: dict[str, Any]) -> "RequestContext": """Build a ``RequestContext`` from a raw A2A ``Message.metadata`` dict.""" return cls( user_email=metadata.get("user_email", ""), tenant=metadata.get("tenant", ""), role=metadata.get("role", ""), raw=metadata, ) def __copy__(self) -> "RequestContext": """Inject dev impersonation whenever pylogue copies deps for a new request. Uses ``dataclasses.replace`` so any new runtime-channel field is automatically preserved without manually listing every attribute. """ return replace(self, user_email=_read_impersonation() or self.user_email, raw=dict(self.raw)) def to_metadata(self) -> dict[str, Any]: """Serialise back to a plain dict suitable for ``Message.metadata``.""" d: dict[str, Any] = dict(self.raw) # preserve any extra keys if self.user_email: d["user_email"] = self.user_email if self.tenant: d["tenant"] = self.tenant if self.role: d["role"] = self.role return d # --------------------------------------------------------------------------- # Internal ContextVar — transport between MetadataAwareWorker and # _DepsInjectingAgent. Not part of the public API; use RunContext in tools. # --------------------------------------------------------------------------- _request_context_var: ContextVar[RequestContext] = ContextVar( "_a2a_request_context", default=RequestContext() ) # --------------------------------------------------------------------------- # _DepsInjectingAgent — wraps the raw pydantic-ai Agent # --------------------------------------------------------------------------- class _DepsInjectingAgent: """Thin wrapper that injects ``deps`` into every ``agent.run()`` call. ``AgentWorker.run_task`` calls ``self.agent.run(message_history=...)`` without a ``deps`` argument. This wrapper intercepts that call and adds ``deps=_request_context_var.get()``, which ``MetadataAwareWorker.run_task`` has already populated with the correct ``RequestContext`` for this task. """ def __init__(self, agent: Any) -> None: self._agent = agent async def run(self, *args: Any, **kwargs: Any) -> Any: ctx = _request_context_var.get() impersonated = _read_impersonation() if impersonated: ctx = replace(ctx, user_email=impersonated) kwargs.setdefault("deps", ctx) return await self._agent.run(*args, **kwargs) # Forward async context-manager protocol used by worker_lifespan async def __aenter__(self) -> "_DepsInjectingAgent": await self._agent.__aenter__() return self async def __aexit__(self, *args: Any) -> None: await self._agent.__aexit__(*args) def __getattr__(self, name: str) -> Any: return getattr(self._agent, name) # --------------------------------------------------------------------------- # MetadataAwareWorker # --------------------------------------------------------------------------- @dataclass class MetadataAwareWorker(AgentWorker): """``AgentWorker`` that populates ``_request_context_var`` before each run. Flow per task ------------- 1. Extract ``params["message"]["metadata"]`` (travels intact through the broker queue from the original ``send_message`` call). 2. Build ``RequestContext.from_metadata(metadata)`` and set the ContextVar. 3. Call ``super().run_task(params)`` — which eventually calls ``_DepsInjectingAgent.run()``, which reads the ContextVar and passes the ``RequestContext`` as ``deps`` to pydantic-ai. 4. Reset the ContextVar token after the task finishes (or raises). """ async def run_task(self, params: "TaskSendParams") -> None: metadata: dict[str, Any] = ( params.get("message", {}).get("metadata") or {} ) token = _request_context_var.set(RequestContext.from_metadata(metadata)) try: await super().run_task(params) finally: _request_context_var.reset(token) # --------------------------------------------------------------------------- # make_metadata_aware_app — drop-in for agent.to_a2a(...) # --------------------------------------------------------------------------- def make_metadata_aware_app(agent_wrapper: Any, **to_a2a_kwargs: Any) -> Any: """Create a ``FastA2A`` app where every ``@tool`` receives a populated ``RunContext[RequestContext]`` derived from the incoming ``Message.metadata``. Drop-in replacement for ``agent.to_a2a(**kwargs)``. """ storage = InMemoryStorage() broker = InMemoryBroker() pydantic_agent = agent_wrapper.agent # the underlying pydantic_ai.Agent # Wrap the agent so AgentWorker.run() calls get deps= injected wrapped_agent = _DepsInjectingAgent(pydantic_agent) # type: ignore[arg-type] worker = MetadataAwareWorker( agent=wrapped_agent, # type: ignore[arg-type] broker=broker, storage=storage, ) lifespan = partial(worker_lifespan, worker=worker, agent=pydantic_agent) return agent_wrapper.to_a2a( storage=storage, broker=broker, lifespan=lifespan, **to_a2a_kwargs, )