Python: Access agent context in as_tool scenarios

[giles17]

Motivation and Context

Adds get_current_agent_run_context() to enable tools and sub-agents to access their parent agent's run context during execution.

Closes #3411

When using agents as tools within other agents, there was no way to access the parent agent's context (conversation_id, messages, thread, etc.) from deeply nested
code. This is needed for:

• Correlating multi-agent conversations in storage systems
• Sharing session data between parent and child agents
• Accessing parent's messages or options from tools

Changes

• New _agent_context.py module:

  • get_current_agent_run_context() - Returns the current AgentContext or None if outside an agent run
  • agent_run_scope() - Context manager for setting ambient context (internal use)

• _middleware.py: Sets ambient context in AgentMiddlewareLayer.run() for both streaming and non-streaming

• _agents.py: Updates context with resolved thread after _prepare_run_context()

• _tools.py: Updates thread's service_thread_id when conversation_id becomes available

Usage

from agent_framework import get_current_agent_run_context, tool

@tool
def my_tool(query: str) -> str:
    parent_ctx = get_current_agent_run_context()
    if parent_ctx:
        # Access parent's conversation_id, messages, agent, etc.
        conv_id = parent_ctx.thread.service_thread_id if parent_ctx.thread else None
    return "result"

Description

Contribution Checklist

• The code builds clean without any errors or warnings
• The PR follows the Contribution Guidelines
• All unit tests pass, and I have added new tests where possible
• Is this a breaking change? If yes, add "[BREAKING]" prefix to the title of the PR.

[Copilot]

Pull request overview

Adds an “ambient” agent run context API so code executing during an agent run (notably tools and agent-as-tool wrappers) can access the current AgentContext (thread/messages/options/etc.) without having it explicitly passed through every call layer.

Changes:

• Introduces get_current_agent_run_context() (and an internal agent_run_scope() context manager) backed by contextvars.
• Sets the ambient context in AgentMiddlewareLayer.run() for non-streaming runs and streaming runs (via a wrapped ResponseStream in the no-middleware path).
• Updates thread propagation so the ambient context can be updated with the resolved thread and conversation id as they become available, plus adds tests for basic scoping behavior.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
python/packages/core/agent_framework/_agent_context.py	New ContextVar-based ambient context getter + scope manager.
python/packages/core/agent_framework/_middleware.py	Establishes ambient context during agent runs; adds stream wrapper helper.
python/packages/core/agent_framework/_agents.py	Updates ambient context with the resolved thread after _prepare_run_context().
python/packages/core/agent_framework/_tools.py	Updates ambient context thread’s service_thread_id when conversation_id becomes known.
python/packages/core/agent_framework/init.py	Re-exports new ambient context API from the package root.
python/packages/core/tests/core/test_agent_context.py	New tests for context scoping/isolation and middleware visibility.

Comments suppressed due to low confidence (2)

python/packages/core/agent_framework/_middleware.py:1240

• In the streaming + middleware case, the ambient agent context is only set while pipeline.execute(...) runs. The returned ResponseStream is iterated after _execute() returns, so get_current_agent_run_context() will be unset during stream iteration/tool calls. Wrap the ResponseStream produced by _execute_stream() with _wrap_stream_with_context(..., context) (similar to the no-middleware streaming path) so context remains available throughout iteration and finalization.

        if stream:
            # For streaming, wrap execution in ResponseStream.from_awaitable
            async def _execute_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse]:
                result = await _execute()
                if result is None:
                    # Create empty stream if middleware terminated without setting result
                    return ResponseStream(_empty_async_iterable())
                if isinstance(result, ResponseStream):
                    return result
                # If result is AgentResponse (shouldn't happen for streaming), convert to stream
                raise ValueError("Expected ResponseStream for streaming, got AgentResponse")

            return ResponseStream.from_awaitable(_execute_stream())

python/packages/core/agent_framework/_agents.py:474

• The as_tool() docstring no longer documents the exceptions it can raise, but the implementation still raises TypeError (when self isn’t an AgentProtocol) and ValueError (when the tool name can’t be determined). Please restore/retain the Raises: section so the docstring matches the actual behavior.

        """Create a FunctionTool that wraps this agent.

        Keyword Args:
            name: The name for the tool. If None, uses the agent's name.
            description: The description for the tool. If None, uses the agent's description or empty string.
            arg_name: The name of the function argument (default: "task").
            arg_description: The description for the function argument.
                If None, defaults to "Task for {tool_name}".
            stream_callback: Optional callback for streaming responses. If provided, uses run(..., stream=True).

        Returns:
            A FunctionTool that can be used as a tool by other agents.

        Examples:
            .. code-block:: python

                from agent_framework import ChatAgent

                # Create an agent
                agent = ChatAgent(chat_client=client, name="research-agent", description="Performs research tasks")

                # Convert the agent to a tool
                research_tool = agent.as_tool()

                # Use the tool with another agent
                coordinator = ChatAgent(chat_client=client, name="coordinator", tools=research_tool)
        """
        # Verify that self implements AgentProtocol
        if not isinstance(self, AgentProtocol):
            raise TypeError(f"Agent {self.__class__.__name__} must implement AgentProtocol to be used as a tool")

        tool_name = name or _sanitize_agent_name(self.name)
        if tool_name is None:
            raise ValueError("Agent tool name cannot be None. Either provide a name parameter or set the agent's name.")
        tool_description = description or self.description or ""

[markwallace-microsoft]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/core/agent_framework
_agent_context.py	15	0	100%
_agents.py	328	35	89%	473, 893, 929, 1036–1038, 1151, 1192, 1194, 1203–1208, 1214, 1216, 1226–1227, 1234, 1236–1237, 1245–1249, 1257–1258, 1260, 1265, 1267, 1301, 1341, 1361
_middleware.py	346	16	95%	81, 84, 89, 798, 800, 802, 923, 950, 952, 977, 1058, 1062, 1233, 1237, 1298, 1372
_tools.py	793	85	89%	232, 278, 329, 331, 359, 529, 564–565, 667, 669, 689, 707, 721, 733, 738, 740, 747, 780, 851–853, 894, 916–944, 979, 987, 1228, 1433, 1490, 1494, 1573–1577, 1595, 1597–1598, 1703, 1707, 1760, 1762, 1785, 1787, 1851, 1878, 1931, 1999, 2178–2179, 2206, 2214, 2227, 2237–2238, 2273, 2329, 2361
TOTAL	16627	2022	87%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
3918	225 💤	0 ❌	0 🔥	1m 8s ⏱️

[eavanvalkenburg]

This scares me because it looks very complex, we have a mechanism in function tools (add **kwargs) to your definition and then you get additional context, can't we use that?

[eavanvalkenburg]

We shouldn't consume a stream unless we absolutely have to.

[dmytrostruk]

Looks like _wrap_stream_with_context function is used only when the pipeline doesn't have middleware registered, but how about the case where middleware exists, should we add wrapping there as well?

[dmytrostruk]

We will also need to add an example that will demonstrate how to use this functionality and how it resolves the problem described in the issue associated with this PR.

[dmytrostruk]

New unit tests don't test the case where context is retrieved within the tool, as described in PR description.