From 1ad8104e52de6c2fcdbd9cf7a2750fdc5cef0173 Mon Sep 17 00:00:00 2001 From: droideronline Date: Wed, 30 Sep 2026 00:58:06 +0530 Subject: [PATCH 1/3] Python: Forward function_invocation_kwargs to GitHubCopilotAgent tools RawGitHubCopilotAgent.run swallowed function_invocation_kwargs in **kwargs and the Copilot SDK tool handler invoked FunctionTool without a FunctionInvocationContext, so tools never saw per-run kwargs or the AgentSession. Thread both through session create/resume into the tool handler and pass them via FunctionInvocationContext. Fixes #8865 --- .../agent_framework_github_copilot/_agent.py | 151 ++++++++++++++++-- .../tests/test_github_copilot_agent.py | 112 +++++++++++++ 2 files changed, 250 insertions(+), 13 deletions(-) diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index b0287e734c1..472d5150a65 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -23,6 +23,7 @@ BaseAgent, Content, ContextProvider, + FunctionInvocationContext, HistoryProvider, Message, ResponseStream, @@ -813,6 +814,7 @@ def run( session: AgentSession | None = None, middleware: Sequence[AgentMiddlewareTypes] | None = None, options: OptionsT | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse]: ... @@ -825,6 +827,7 @@ def run( session: AgentSession | None = None, middleware: Sequence[AgentMiddlewareTypes] | None = None, options: OptionsT | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, **kwargs: Any, ) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ... @@ -836,6 +839,7 @@ def run( session: AgentSession | None = None, middleware: Sequence[AgentMiddlewareTypes] | None = None, options: OptionsT | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]: """Get a response from the agent. @@ -854,6 +858,9 @@ def run( compatibility; pass middleware via :class:`GitHubCopilotAgent` which forwards it through :class:`AgentTelemetryLayer`. options: Runtime options (model, timeout, etc.). + function_invocation_kwargs: Keyword arguments forwarded to tool invocations. + Tools receive them through ``FunctionInvocationContext.kwargs`` when they + declare a ``FunctionInvocationContext`` parameter. kwargs: Additional keyword arguments for compatibility with the shared agent interface (e.g. compaction_strategy, tokenizer). Not used by this agent. @@ -887,11 +894,22 @@ def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse: return AgentResponse.from_updates(updates) return ResponseStream( - self._stream_updates(messages=messages, session=session, options=options, _ctx_holder=ctx_holder), + self._stream_updates( + messages=messages, + session=session, + options=options, + function_invocation_kwargs=function_invocation_kwargs, + _ctx_holder=ctx_holder, + ), finalizer=_finalize, result_hooks=[_after_run_hook], ) - return self._run_impl(messages=messages, session=session, options=options) + return self._run_impl( + messages=messages, + session=session, + options=options, + function_invocation_kwargs=function_invocation_kwargs, + ) @staticmethod def _parse_usage_details_from_copilot(data: AssistantUsageData) -> UsageDetails | None: @@ -920,6 +938,7 @@ async def _run_impl( *, session: AgentSession | None = None, options: OptionsT | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> AgentResponse: """Non-streaming implementation of run.""" if not self._started: @@ -952,7 +971,12 @@ async def _run_impl( existing = list(opts.get("tools") or []) opts["tools"] = existing + list(session_context.tools) - copilot_session = await self._get_or_create_session(session, streaming=False, runtime_options=opts) + copilot_session = await self._get_or_create_session( + session, + streaming=False, + runtime_options=opts, + function_invocation_kwargs=function_invocation_kwargs, + ) usage_details: UsageDetails | None = None finish_reason: str | None = None model: str | None = None @@ -1032,6 +1056,7 @@ async def _stream_updates( *, session: AgentSession | None = None, options: OptionsT | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, _ctx_holder: dict[str, Any] | None = None, ) -> AsyncIterable[AgentResponseUpdate]: """Internal method to stream updates from GitHub Copilot. @@ -1042,6 +1067,7 @@ async def _stream_updates( Keyword Args: session: The conversation session associated with the message(s). options: Runtime options (model, timeout, etc.). + function_invocation_kwargs: Keyword arguments forwarded to tool invocations. _ctx_holder: Internal dict populated with session_context and session so that the caller (via a ResponseStream result_hook) can run after_run providers without duplicating the updates buffer. @@ -1081,7 +1107,12 @@ async def _stream_updates( existing = list(opts.get("tools") or []) opts["tools"] = existing + list(session_context.tools) - copilot_session = await self._get_or_create_session(session, streaming=True, runtime_options=opts) + copilot_session = await self._get_or_create_session( + session, + streaming=True, + runtime_options=opts, + function_invocation_kwargs=function_invocation_kwargs, + ) if _ctx_holder is not None: _ctx_holder["session_context"] = session_context @@ -1301,12 +1332,20 @@ def _prepare_attachments_for_copilot(messages: Sequence[Message]) -> list[Attach def _prepare_tools( self, tools: Sequence[ToolTypes | CopilotTool], + *, + agent_session: AgentSession | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> list[CopilotTool]: """Convert Agent Framework tools to Copilot SDK tools. Args: tools: List of Agent Framework tools. + Keyword Args: + agent_session: The conversation session, exposed to tools through their + ``FunctionInvocationContext``. + function_invocation_kwargs: Keyword arguments forwarded to tool invocations. + Returns: List of Copilot SDK tools. """ @@ -1316,16 +1355,33 @@ def _prepare_tools( if isinstance(tool, CopilotTool): copilot_tools.append(tool) elif isinstance(tool, FunctionTool): - copilot_tools.append(self._tool_to_copilot_tool(tool)) + copilot_tools.append( + self._tool_to_copilot_tool( + tool, + agent_session=agent_session, + function_invocation_kwargs=function_invocation_kwargs, + ) + ) elif isinstance(tool, MutableMapping): copilot_tools.append(tool) # type: ignore[arg-type] # Note: Other tool types (e.g., dict-based hosted tools) are skipped return copilot_tools - def _tool_to_copilot_tool(self, ai_func: FunctionTool) -> CopilotTool: + def _tool_to_copilot_tool( + self, + ai_func: FunctionTool, + *, + agent_session: AgentSession | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, + ) -> CopilotTool: """Convert an FunctionTool to a Copilot SDK tool. + Each invocation receives a ``FunctionInvocationContext`` carrying the run's + ``function_invocation_kwargs`` and ``agent_session``, so tools that declare a + ``FunctionInvocationContext`` parameter can read them the same way they do + with a chat-client-backed agent. + Approval for tools declared with ``approval_mode="always_require"`` is normally enforced by the Copilot SDK's native ``on_pre_tool_use`` hook (see :meth:`_build_session_hooks`). When the deprecated ``on_function_approval`` @@ -1335,6 +1391,7 @@ def _tool_to_copilot_tool(self, ai_func: FunctionTool) -> CopilotTool: """ approval_handler = self._function_approval_handler enforce = approval_handler is not None and ai_func.approval_mode == "always_require" + runtime_kwargs = dict(function_invocation_kwargs) if function_invocation_kwargs is not None else {} async def handler(invocation: ToolInvocation) -> ToolResult: args: dict[str, Any] = invocation.arguments or {} @@ -1353,11 +1410,21 @@ async def handler(invocation: ToolInvocation) -> ToolResult: result_type="failure", error="approval_denied", ) + context = FunctionInvocationContext( + function=ai_func, + arguments=args, + session=agent_session, + kwargs=runtime_kwargs, + ) if ai_func.input_model: args_instance = ai_func.input_model(**args) - result = await ai_func.invoke(arguments=args_instance) + result = await ai_func.invoke( + arguments=args_instance, context=context, tool_call_id=invocation.tool_call_id or None + ) else: - result = await ai_func.invoke(arguments=args) + result = await ai_func.invoke( + arguments=args, context=context, tool_call_id=invocation.tool_call_id or None + ) rich = [c for c in result if c.type in ("data", "uri")] if rich: logger.warning( @@ -1476,6 +1543,7 @@ async def _get_or_create_session( agent_session: AgentSession, streaming: bool = False, runtime_options: dict[str, Any] | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> CopilotSession: """Get an existing session or create a new one for the session. @@ -1483,6 +1551,7 @@ async def _get_or_create_session( agent_session: The conversation session. streaming: Whether to enable streaming for the session. runtime_options: Runtime options from run that take precedence. + function_invocation_kwargs: Keyword arguments forwarded to tool invocations. Returns: A CopilotSession instance. @@ -1500,9 +1569,20 @@ async def _get_or_create_session( raise AgentException( "GitHubCopilotAgent expects a string service_session_id for session resumption." ) - return await self._resume_session(service_session_id, streaming, runtime_options) + return await self._resume_session( + service_session_id, + streaming, + runtime_options, + agent_session=agent_session, + function_invocation_kwargs=function_invocation_kwargs, + ) - session = await self._create_session(streaming, runtime_options) + session = await self._create_session( + streaming, + runtime_options, + agent_session=agent_session, + function_invocation_kwargs=function_invocation_kwargs, + ) agent_session.service_session_id = session.session_id return session except Exception as ex: @@ -1533,6 +1613,9 @@ def _build_session_kwargs( self, streaming: bool, runtime_options: dict[str, Any] | None, + *, + agent_session: AgentSession | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> dict[str, Any]: """Assemble keyword arguments for ``create_session`` / ``resume_session``. @@ -1551,6 +1634,11 @@ def _build_session_kwargs( streaming: Whether to enable streaming for the session. runtime_options: Runtime options that take precedence over default_options. + Keyword Args: + agent_session: The conversation session, exposed to tools through their + ``FunctionInvocationContext``. + function_invocation_kwargs: Keyword arguments forwarded to tool invocations. + Returns: The keyword arguments to splat into the SDK session factory. """ @@ -1565,7 +1653,15 @@ def _build_session_kwargs( # and flattens tool-collection wrappers, which can otherwise hide an MCPTool. all_tools = normalize_tools(list(self._tools or []) + list(kwargs.get("tools") or [])) _reject_mcp_tools(all_tools) - kwargs["tools"] = self._prepare_tools(all_tools) if all_tools else None + kwargs["tools"] = ( + self._prepare_tools( + all_tools, + agent_session=agent_session, + function_invocation_kwargs=function_invocation_kwargs, + ) + if all_tools + else None + ) kwargs["streaming"] = streaming # model may already be present from per-run options (merged above); otherwise fall @@ -1606,23 +1702,40 @@ async def _create_session( self, streaming: bool, runtime_options: dict[str, Any] | None = None, + *, + agent_session: AgentSession | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> CopilotSession: """Create a new Copilot session. Args: streaming: Whether to enable streaming for the session. runtime_options: Runtime options that take precedence over default_options. + + Keyword Args: + agent_session: The conversation session, exposed to tools. + function_invocation_kwargs: Keyword arguments forwarded to tool invocations. """ if not self._client: raise RuntimeError("GitHub Copilot client not initialized. Call start() first.") - return await self._client.create_session(**self._build_session_kwargs(streaming, runtime_options)) + return await self._client.create_session( + **self._build_session_kwargs( + streaming, + runtime_options, + agent_session=agent_session, + function_invocation_kwargs=function_invocation_kwargs, + ) + ) async def _resume_session( self, session_id: str, streaming: bool, runtime_options: dict[str, Any] | None = None, + *, + agent_session: AgentSession | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> CopilotSession: """Resume an existing Copilot session by ID. @@ -1630,11 +1743,23 @@ async def _resume_session( session_id: The session ID to resume. streaming: Whether to enable streaming for the session. runtime_options: Runtime options that take precedence over default_options. + + Keyword Args: + agent_session: The conversation session, exposed to tools. + function_invocation_kwargs: Keyword arguments forwarded to tool invocations. """ if not self._client: raise RuntimeError("GitHub Copilot client not initialized. Call start() first.") - return await self._client.resume_session(session_id, **self._build_session_kwargs(streaming, runtime_options)) + return await self._client.resume_session( + session_id, + **self._build_session_kwargs( + streaming, + runtime_options, + agent_session=agent_session, + function_invocation_kwargs=function_invocation_kwargs, + ), + ) class GitHubCopilotAgent( # type: ignore[misc] diff --git a/python/packages/github_copilot/tests/test_github_copilot_agent.py b/python/packages/github_copilot/tests/test_github_copilot_agent.py index b0c9bdcad19..73218ca2180 100644 --- a/python/packages/github_copilot/tests/test_github_copilot_agent.py +++ b/python/packages/github_copilot/tests/test_github_copilot_agent.py @@ -26,6 +26,7 @@ AgentSession, Content, ContextProvider, + FunctionInvocationContext, HistoryProvider, MCPStdioTool, Message, @@ -2422,6 +2423,117 @@ def no_args_tool() -> str: assert result.result_type == "success" assert result.text_result_for_llm == "no args result" + async def test_run_forwards_function_invocation_kwargs_to_tools( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """Test that function_invocation_kwargs passed to run() reach tools via FunctionInvocationContext.""" + captured: dict[str, Any] = {} + + def lookup(query: str, ctx: FunctionInvocationContext) -> str: + """Look something up for the current user.""" + captured["kwargs"] = dict(ctx.kwargs) + captured["session"] = ctx.session + return f"{ctx.kwargs['user_id']}:{query}" + + mock_session.send_and_wait.return_value = assistant_message_event + agent = GitHubCopilotAgent(client=mock_client, tools=[lookup]) + session = agent.create_session() + + await agent.run("Hello", session=session, function_invocation_kwargs={"user_id": "u-1"}) + + copilot_tool = mock_client.create_session.call_args.kwargs["tools"][0] + result = await copilot_tool.handler(ToolInvocation(arguments={"query": "orders"}, tool_call_id="call-1")) + + assert result.result_type == "success" + assert result.text_result_for_llm == "u-1:orders" + assert captured["kwargs"] == {"user_id": "u-1"} + assert captured["session"] is session + # The tool's JSON schema must not expose the context parameter to the model. + assert "ctx" not in copilot_tool.parameters["properties"] + + async def test_run_streaming_forwards_function_invocation_kwargs_to_tools( + self, + mock_client: MagicMock, + mock_session: MagicMock, + session_idle_event: SessionEvent, + ) -> None: + """Test that streaming runs forward function_invocation_kwargs to tools.""" + captured: dict[str, Any] = {} + + def lookup(query: str, ctx: FunctionInvocationContext) -> str: + """Look something up for the current user.""" + captured.update(ctx.kwargs) + return query + + def mock_on(handler: Any) -> Any: + handler(session_idle_event) + return lambda: None + + mock_session.on = mock_on + agent = GitHubCopilotAgent(client=mock_client, tools=[lookup]) + + async for _ in agent.run("Hello", stream=True, function_invocation_kwargs={"tenant": "contoso"}): + pass + + copilot_tool = mock_client.create_session.call_args.kwargs["tools"][0] + result = await copilot_tool.handler(ToolInvocation(arguments={"query": "q"})) + + assert result.result_type == "success" + assert captured == {"tenant": "contoso"} + + async def test_resumed_session_forwards_function_invocation_kwargs_to_tools( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """Test that tools on a resumed session get the current run's function_invocation_kwargs.""" + captured: dict[str, Any] = {} + + def lookup(query: str, ctx: FunctionInvocationContext) -> str: + """Look something up for the current user.""" + captured.update(ctx.kwargs) + return query + + mock_session.send_and_wait.return_value = assistant_message_event + agent = GitHubCopilotAgent(client=mock_client, tools=[lookup]) + session = agent.create_session() + session.service_session_id = "existing-session-id" + + await agent.run("Hello", session=session, function_invocation_kwargs={"user_id": "u-2"}) + + mock_client.create_session.assert_not_called() + copilot_tool = mock_client.resume_session.call_args.kwargs["tools"][0] + await copilot_tool.handler(ToolInvocation(arguments={"query": "q"})) + + assert captured == {"user_id": "u-2"} + + async def test_tool_without_context_parameter_ignores_function_invocation_kwargs( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """Test that tools without a context parameter still run when function_invocation_kwargs are set.""" + + def echo(text: str) -> str: + """Echo the text back.""" + return text + + mock_session.send_and_wait.return_value = assistant_message_event + agent = GitHubCopilotAgent(client=mock_client, tools=[echo]) + + await agent.run("Hello", function_invocation_kwargs={"user_id": "u-3"}) + + copilot_tool = mock_client.create_session.call_args.kwargs["tools"][0] + result = await copilot_tool.handler(ToolInvocation(arguments={"text": "hi"})) + + assert result.result_type == "success" + assert result.text_result_for_llm == "hi" + def test_copilot_tool_passthrough( self, mock_client: MagicMock, From d401cdc88dc90169bebcfc160e145038d65a6941 Mon Sep 17 00:00:00 2001 From: droideronline Date: Wed, 30 Sep 2026 11:30:47 +0530 Subject: [PATCH 2/3] Python: Serialize overlapping GitHubCopilotAgent runs on the same session The Copilot SDK keeps one session object per service session ID and re-registers tools on every resume, dispatching tool calls by session ID. Since tool handlers now capture the run's function_invocation_kwargs and session, an overlapping resume could let one run's tool calls execute with another run's context. Hold a per-session run lock from create/resume until the run completes. --- .../agent_framework_github_copilot/_agent.py | 137 +++++++++++----- .../tests/test_github_copilot_agent.py | 149 ++++++++++++++++++ 2 files changed, 250 insertions(+), 36 deletions(-) diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index 472d5150a65..5ccb9032ae9 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -9,7 +9,8 @@ import logging import sys import warnings -from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence +import weakref +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence from pathlib import Path from typing import Any, ClassVar, Generic, Literal, TypedDict, cast, overload from urllib.parse import urlparse @@ -182,6 +183,10 @@ async def _resolve_function_approval( logger = logging.getLogger("agent_framework.github_copilot") +# One lock per (CopilotClient, service session ID), shared across agents that share a client. +# Entries disappear once no run holds or awaits the lock. +_SESSION_RUN_LOCKS: weakref.WeakValueDictionary[tuple[int, str], asyncio.Lock] = weakref.WeakValueDictionary() + _MCP_TOOL_MESSAGE = ( "MCP server '{name}' cannot be passed to GitHubCopilotAgent as a tool: the Copilot SDK " "connects to MCP servers itself, so a framework-managed MCPTool would keep none of its " @@ -971,12 +976,6 @@ async def _run_impl( existing = list(opts.get("tools") or []) opts["tools"] = existing + list(session_context.tools) - copilot_session = await self._get_or_create_session( - session, - streaming=False, - runtime_options=opts, - function_invocation_kwargs=function_invocation_kwargs, - ) usage_details: UsageDetails | None = None finish_reason: str | None = None model: str | None = None @@ -1010,14 +1009,20 @@ def usage_event_handler(event: SessionEvent) -> None: prompt = "\n".join(session_context.instructions) + "\n" + prompt attachments = self._prepare_attachments_for_copilot(context_messages) - unsubscribe = copilot_session.on(usage_event_handler) - try: - mark_feature_used(FeatureIndex.GITHUB_COPILOT) - response_event = await copilot_session.send_and_wait(prompt, attachments=attachments, timeout=timeout) - except Exception as ex: - raise AgentException(f"GitHub Copilot request failed: {ex}") from ex - finally: - unsubscribe() + async with self._session_run_scope( + session, + streaming=False, + runtime_options=opts, + function_invocation_kwargs=function_invocation_kwargs, + ) as copilot_session: + unsubscribe = copilot_session.on(usage_event_handler) + try: + mark_feature_used(FeatureIndex.GITHUB_COPILOT) + response_event = await copilot_session.send_and_wait(prompt, attachments=attachments, timeout=timeout) + except Exception as ex: + raise AgentException(f"GitHub Copilot request failed: {ex}") from ex + finally: + unsubscribe() response_messages: list[Message] = [] response_id: str | None = None @@ -1107,17 +1112,6 @@ async def _stream_updates( existing = list(opts.get("tools") or []) opts["tools"] = existing + list(session_context.tools) - copilot_session = await self._get_or_create_session( - session, - streaming=True, - runtime_options=opts, - function_invocation_kwargs=function_invocation_kwargs, - ) - - if _ctx_holder is not None: - _ctx_holder["session_context"] = session_context - _ctx_holder["session"] = session - # Build the prompt from the full session context so provider-injected messages are included. context_messages = session_context.get_messages(include_input=True) prompt = "\n".join([message.text for message in context_messages]) @@ -1203,18 +1197,28 @@ def event_handler(event: SessionEvent) -> None: error_msg = error_data.message or "Unknown error" queue.put_nowait(AgentException(f"GitHub Copilot session error: {error_msg}")) - unsubscribe = copilot_session.on(event_handler) + async with self._session_run_scope( + session, + streaming=True, + runtime_options=opts, + function_invocation_kwargs=function_invocation_kwargs, + ) as copilot_session: + if _ctx_holder is not None: + _ctx_holder["session_context"] = session_context + _ctx_holder["session"] = session + + unsubscribe = copilot_session.on(event_handler) - try: - mark_feature_used(FeatureIndex.GITHUB_COPILOT) - await copilot_session.send(prompt, attachments=attachments) + try: + mark_feature_used(FeatureIndex.GITHUB_COPILOT) + await copilot_session.send(prompt, attachments=attachments) - while (item := await queue.get()) is not None: - if isinstance(item, Exception): - raise item - yield item - finally: - unsubscribe() + while (item := await queue.get()) is not None: + if isinstance(item, Exception): + raise item + yield item + finally: + unsubscribe() async def _run_before_providers( self, @@ -1538,6 +1542,67 @@ def default_pre_tool_use( return {"on_pre_tool_use": default_pre_tool_use} + def _session_run_lock(self, service_session_id: str) -> asyncio.Lock: + """Return the lock that serializes runs on one Copilot service session.""" + key = (id(self._client), service_session_id) + lock = _SESSION_RUN_LOCKS.get(key) + if lock is None: + lock = asyncio.Lock() + _SESSION_RUN_LOCKS[key] = lock + return lock + + @contextlib.asynccontextmanager + async def _session_run_scope( + self, + agent_session: AgentSession, + *, + streaming: bool, + runtime_options: dict[str, Any], + function_invocation_kwargs: Mapping[str, Any] | None, + ) -> AsyncGenerator[CopilotSession]: + """Open the Copilot session for one run and hold it exclusively until the run ends. + + The Copilot SDK keeps a single session object per service session ID and + re-registers tools on it on every resume, dispatching tool calls by session ID. + Tool handlers capture the run's ``function_invocation_kwargs`` and session, so + an overlapping resume would let one run's tool calls execute with another run's + context. Runs on the same service session are therefore serialized. + + Args: + agent_session: The conversation session. + + Keyword Args: + streaming: Whether to enable streaming for the session. + runtime_options: Runtime options from run that take precedence. + function_invocation_kwargs: Keyword arguments forwarded to tool invocations. + + Yields: + The created or resumed CopilotSession. + """ + lock: asyncio.Lock | None = None + service_session_id = agent_session.service_session_id + if isinstance(service_session_id, str): + resume_lock = self._session_run_lock(service_session_id) + await resume_lock.acquire() + lock = resume_lock + try: + copilot_session = await self._get_or_create_session( + agent_session, + streaming=streaming, + runtime_options=runtime_options, + function_invocation_kwargs=function_invocation_kwargs, + ) + if lock is None: + # A newly created service session ID is not known to any other run yet, + # so this acquire never waits. + create_lock = self._session_run_lock(copilot_session.session_id) + await create_lock.acquire() + lock = create_lock + yield copilot_session + finally: + if lock is not None: + lock.release() + async def _get_or_create_session( self, agent_session: AgentSession, diff --git a/python/packages/github_copilot/tests/test_github_copilot_agent.py b/python/packages/github_copilot/tests/test_github_copilot_agent.py index 73218ca2180..50b71199ef6 100644 --- a/python/packages/github_copilot/tests/test_github_copilot_agent.py +++ b/python/packages/github_copilot/tests/test_github_copilot_agent.py @@ -2534,6 +2534,155 @@ def echo(text: str) -> str: assert result.result_type == "success" assert result.text_result_for_llm == "hi" + @staticmethod + def _last_registration_wins_client(mock_client: MagicMock, mock_session: MagicMock) -> dict[str, Any]: + """Make create/resume mimic the SDK: the latest call's tools replace earlier ones for the session.""" + registered: dict[str, Any] = {} + + async def register(*args: Any, **kwargs: Any) -> MagicMock: + registered["tools"] = kwargs["tools"] + return mock_session + + mock_client.create_session = AsyncMock(side_effect=register) + mock_client.resume_session = AsyncMock(side_effect=register) + return registered + + async def test_overlapping_runs_on_same_session_keep_their_own_tool_context( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """Test that a concurrent run cannot replace the tool handlers of a run still in progress.""" + import asyncio + + registered = self._last_registration_wins_client(mock_client, mock_session) + seen: list[str] = [] + + def whoami(ctx: FunctionInvocationContext) -> str: + """Return the calling tenant.""" + seen.append(ctx.kwargs["tenant"]) + return ctx.kwargs["tenant"] + + first_run_sending = asyncio.Event() + release_first_run = asyncio.Event() + calls = 0 + + async def send_and_wait(*args: Any, **kwargs: Any) -> SessionEvent: + nonlocal calls + calls += 1 + if calls == 1: + first_run_sending.set() + await release_first_run.wait() + # The SDK dispatches the tool call to whatever handler is registered now. + await registered["tools"][0].handler(ToolInvocation(arguments={})) + return assistant_message_event + + mock_session.send_and_wait = send_and_wait + agent = GitHubCopilotAgent(client=mock_client, tools=[whoami]) + await agent.start() + session = agent.create_session() + session.service_session_id = "shared-session-id" + + run_a = asyncio.create_task(agent.run("Hi", session=session, function_invocation_kwargs={"tenant": "a"})) + await first_run_sending.wait() + run_b = asyncio.create_task(agent.run("Hi", session=session, function_invocation_kwargs={"tenant": "b"})) + await asyncio.sleep(0.05) + + # Run B must wait for run A instead of resuming the session and replacing A's handlers. + assert mock_client.resume_session.await_count == 1 + + release_first_run.set() + await asyncio.gather(run_a, run_b) + + assert seen == ["a", "b"] + assert mock_client.resume_session.await_count == 2 + + async def test_overlapping_streaming_runs_on_same_session_are_serialized( + self, + mock_client: MagicMock, + mock_session: MagicMock, + session_idle_event: SessionEvent, + ) -> None: + """Test that streaming runs on the same session do not overlap.""" + import asyncio + + self._last_registration_wins_client(mock_client, mock_session) + handlers: list[Any] = [] + mock_session.on = lambda handler: handlers.append(handler) or (lambda: None) + mock_session.send = AsyncMock() + + agent = GitHubCopilotAgent(client=mock_client) + await agent.start() + session = agent.create_session() + session.service_session_id = "shared-session-id" + + async def consume(tenant: str) -> None: + async for _ in agent.run("Hi", session=session, stream=True, function_invocation_kwargs={"t": tenant}): + pass + + run_a = asyncio.create_task(consume("a")) + run_b = asyncio.create_task(consume("b")) + await asyncio.sleep(0.05) + + assert mock_client.resume_session.await_count == 1 + + handlers[0](session_idle_event) + await asyncio.sleep(0.05) + assert mock_client.resume_session.await_count == 2 + + handlers[1](session_idle_event) + await asyncio.gather(run_a, run_b) + + async def test_session_run_lock_released_after_failed_run( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """Test that a failed run does not leave the session locked for later runs.""" + import asyncio + + mock_session.send_and_wait = AsyncMock(side_effect=[RuntimeError("boom"), assistant_message_event]) + agent = GitHubCopilotAgent(client=mock_client) + session = agent.create_session() + + with pytest.raises(AgentException): + await agent.run("Hi", session=session) + + response = await asyncio.wait_for(agent.run("Hi", session=session), timeout=1) + assert response.text == "Test response" + + async def test_runs_on_different_sessions_are_not_serialized( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """Test that the run lock only applies to runs sharing a Copilot service session.""" + import asyncio + + in_flight = 0 + max_in_flight = 0 + + async def send_and_wait(*args: Any, **kwargs: Any) -> SessionEvent: + nonlocal in_flight, max_in_flight + in_flight += 1 + max_in_flight = max(max_in_flight, in_flight) + await asyncio.sleep(0.05) + in_flight -= 1 + return assistant_message_event + + mock_session.send_and_wait = send_and_wait + agent = GitHubCopilotAgent(client=mock_client) + first, second = agent.create_session(), agent.create_session() + first.service_session_id = "session-1" + second.service_session_id = "session-2" + + await asyncio.gather(agent.run("Hi", session=first), agent.run("Hi", session=second)) + + assert max_in_flight == 2 + def test_copilot_tool_passthrough( self, mock_client: MagicMock, From 0eb1e19605732ec34b4d0241697ac27cfd8962c3 Mon Sep 17 00:00:00 2001 From: droideronline Date: Thu, 1 Oct 2026 12:49:23 +0530 Subject: [PATCH 3/3] Python: Fix typing of overlapping-run Copilot tests --- .../github_copilot/tests/test_github_copilot_agent.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/python/packages/github_copilot/tests/test_github_copilot_agent.py b/python/packages/github_copilot/tests/test_github_copilot_agent.py index 50b71199ef6..bbcf35d114a 100644 --- a/python/packages/github_copilot/tests/test_github_copilot_agent.py +++ b/python/packages/github_copilot/tests/test_github_copilot_agent.py @@ -2584,9 +2584,9 @@ async def send_and_wait(*args: Any, **kwargs: Any) -> SessionEvent: session = agent.create_session() session.service_session_id = "shared-session-id" - run_a = asyncio.create_task(agent.run("Hi", session=session, function_invocation_kwargs={"tenant": "a"})) + run_a = asyncio.ensure_future(agent.run("Hi", session=session, function_invocation_kwargs={"tenant": "a"})) await first_run_sending.wait() - run_b = asyncio.create_task(agent.run("Hi", session=session, function_invocation_kwargs={"tenant": "b"})) + run_b = asyncio.ensure_future(agent.run("Hi", session=session, function_invocation_kwargs={"tenant": "b"})) await asyncio.sleep(0.05) # Run B must wait for run A instead of resuming the session and replacing A's handlers. @@ -2609,7 +2609,12 @@ async def test_overlapping_streaming_runs_on_same_session_are_serialized( self._last_registration_wins_client(mock_client, mock_session) handlers: list[Any] = [] - mock_session.on = lambda handler: handlers.append(handler) or (lambda: None) + + def mock_on(handler: Any) -> Any: + handlers.append(handler) + return lambda: None + + mock_session.on = mock_on mock_session.send = AsyncMock() agent = GitHubCopilotAgent(client=mock_client)