
    'j-                    j   U d Z ddlmZ ddlZddlmZ ddlmZmZm	Z	 ddl
mZ ddlmZ ddlmZmZ  ed	          ZddZddlmZ ddlmZ ddlmZmZ erddlmZ e G d d                      Z ed e                      Zded<    G d d          Ze G d de                      Z d dZ!dS )!u	  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())
    )annotationsN)
ContextVar)	dataclassfieldreplace)partial)Path)TYPE_CHECKINGAnyz/tmp/eb_dev_impersonationreturnstrc                 |    	 t                                                                           S # t          $ r Y dS w xY w)N )_IMPERSONATION_FILE	read_textstripFileNotFoundError     A/var/www/html/ai-enterprise-brain/backend/chanakya/a2a_context.py_read_impersonationr   H   sG    ",,..44666   rrs   *- 
;;)InMemoryBroker)InMemoryStorage)AgentWorkerworker_lifespan)TaskSendParamsc                  .   e Zd ZU dZdZded<   dZded<   dZded<    ed          Z	d	ed
<    ed          Z
ded<    ed          Zded<    ee          Zded<    ed          Zded<    ed          Zded<    ee          Zded<    ee          Zded<   dZded<    ed          Zded<    ed          Zded<    ee          Zded<    ed          Zded<    ee          Zded <   ed'd$            Zd(d%Zd)d&ZdS )*RequestContexta  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.
    r   r   
user_emailtenantroleNdefaultz
int | Noneconversation_idzlist[dict[str, Any]] | Nonetatasteel_resultzdict[str, Any] | Nonetatasteel_primary_result)default_factoryzlist[dict[str, Any]]"tatasteel_decision_support_resultsz
str | Nonetatasteel_generated_sqltatasteel_sql_explanationzlist[dict[str, str]]tatasteel_query_log
web_claimsmarket_context_text	ui_resultzasyncio.Queue | Noneprogress_queueflow_logzlist[dict[str, str]] | Nonemessage_history_wiredict[str, Any]rawmetadatar   'RequestContext'c                     | |                     dd          |                     dd          |                     dd          |          S )zDBuild a ``RequestContext`` from a raw A2A ``Message.metadata`` dict.r   r   r    r!   )r   r    r!   r3   )get)clsr4   s     r   from_metadatazRequestContext.from_metadata   sS     s||L"55<<"--fb))	
 
 
 	
r   c                p    t          | t                      p| j        t          | j                            S )zInject 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.
        )r   r3   )r   r   r   dictr3   selfs    r   __copy__zRequestContext.__copy__   s4     t(;(=(=(PVZ[_[cVdVdeeeer   c                    t          | j                  }| j        r
| j        |d<   | j        r
| j        |d<   | j        r
| j        |d<   |S )zASerialise back to a plain dict suitable for ``Message.metadata``.r   r    r!   )r;   r3   r   r    r!   )r=   ds     r   to_metadatazRequestContext.to_metadata   sS     NN? 	."oAlO; 	&+AhK9 	"	AfIr   )r4   r2   r   r5   )r   r5   )r   r2   )__name__
__module____qualname____doc__r   __annotations__r    r!   r   r$   r%   r&   listr(   r)   r*   r+   r,   r-   r.   r/   r0   r1   r;   r3   classmethodr9   r>   rA   r   r   r   r   r   Z   sC        	 	 JFDNNNN #(%"5"5"5O555549E$4G4G4GGGGG6;eD6I6I6IIIII?DuUY?Z?Z?Z&ZZZZ*/%*=*=*=====,1E$,?,?,?????05d0K0K0KKKKK',uT'B'B'BJBBBB!!!!!',uT':':':I::::+05+>+>+>N>>>> &+U4%@%@%@H@@@@ 9>d8K8K8KKKKK  %555C5555
 
 
 [
f f f f	 	 	 	 	 	r   r   _a2a_request_contextr"   zContextVar[RequestContext]_request_context_varc                  :    e Zd ZdZddZdd	ZddZddZddZdS )_DepsInjectingAgenta~  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.
    agentr   r   Nonec                    || _         d S N)_agent)r=   rM   s     r   __init__z_DepsInjectingAgent.__init__   s    r   argskwargsc                   K   t                                           }t                      }|rt          ||          }|                    d|            | j        j        |i | d {V S )N)r   deps)rJ   r7   r   r   
setdefaultrQ   run)r=   rS   rT   ctximpersonateds        r   rX   z_DepsInjectingAgent.run   s}      "&&((*,, 	8#,777C&#&&&$T[_d5f555555555r   '_DepsInjectingAgent'c                H   K   | j                                          d {V  | S rP   )rQ   
__aenter__r<   s    r   r]   z_DepsInjectingAgent.__aenter__   s1      k$$&&&&&&&&&r   c                4   K    | j         j        |  d {V  d S rP   )rQ   	__aexit__)r=   rS   s     r   r_   z_DepsInjectingAgent.__aexit__   s0      #dk#T**********r   namer   c                ,    t          | j        |          S rP   )getattrrQ   )r=   r`   s     r   __getattr__z_DepsInjectingAgent.__getattr__   s    t{D)))r   N)rM   r   r   rN   )rS   r   rT   r   r   r   )r   r[   )rS   r   r   rN   )r`   r   r   r   )	rB   rC   rD   rE   rR   rX   r]   r_   rc   r   r   r   rL   rL      s            6 6 6 6   + + + +* * * * * *r   rL   c                  $     e Zd ZdZd fdZ xZS )MetadataAwareWorkeru^  ``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).
    params'TaskSendParams'r   rN   c                  K   |                     di                                d          pi }t                              t                              |                    }	 t                                          |           d {V  t                              |           d S # t                              |           w xY w)Nmessager4   )r7   rJ   setr   r9   superrun_taskreset)r=   rf   r4   token	__class__s       r   rl   zMetadataAwareWorker.run_task   s      JJy"%%))*55; 	 %(()E)Eh)O)OPP	.''""6********* &&u----- &&u----s   "'B% %C)rf   rg   r   rN   )rB   rC   rD   rE   rl   __classcell__)ro   s   @r   re   re      sG         . . . . . . . . . .r   re   agent_wrapperr   to_a2a_kwargsc                    t                      }t                      }| j        }t          |          }t	          |||          }t          t          ||          } | j        d|||d|S )zCreate 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)``.
    )rM   brokerstorage)workerrM   )ru   rt   lifespanr   )r   r   rM   rL   re   r   r   to_a2a)rq   rr   ru   rt   pydantic_agentwrapped_agentrv   rw   s           r   make_metadata_aware_appr{      s     GF"(N (77M   F
 v^LLLH=   	  r   )r   r   )rq   r   rr   r   r   r   )"rE   
__future__r   asynciocontextvarsr   dataclassesr   r   r   	functoolsr   pathlibr	   typingr
   r   r   r   fasta2a.brokerr   fasta2a.storager   pydantic_ai._a2ar   r   fasta2a.schemar   r   rJ   rF   rL   re   r{   r   r   r   <module>r      s  : : :v # " " " " "  " " " " " " 1 1 1 1 1 1 1 1 1 1             % % % % % % % %d677     * ) ) ) ) ) + + + + + + 9 9 9 9 9 9 9 9 .------ N N N N N N N Nl 4>:NN$4$44 4 4     * * * * * * * *H . . . . .+ . . .:     r   