From 3dd5d04215094ce1aca9b1d6fb76b32840014823 Mon Sep 17 00:00:00 2001 From: Sriharsha Etigowni Date: Fri, 30 Oct 2026 20:43:07 +0000 Subject: [PATCH 1/5] feat(copilot_sdk): add support for Copilot SDK 1.x Use public SDK 1.x APIs, preserve taskflow MCP confirmation policies, and cover authentication, model selection, streaming and cleanup with offline regressions. BREAKING CHANGE: Python 3.11 or newer is required. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- README.md | 17 +- pyproject.toml | 5 +- src/seclab_taskflow_agent/runner.py | 4 +- src/seclab_taskflow_agent/sdk/base.py | 7 +- .../sdk/copilot_sdk/backend.py | 428 +++++++--- .../sdk/copilot_sdk/mcp.py | 163 ++-- .../sdk/copilot_sdk/permissions.py | 61 +- tests/test_sdk_copilot_adapter.py | 802 ++++++++++++++---- tests/test_sdk_copilot_mcp.py | 307 ++++--- tests/test_sdk_copilot_permissions.py | 134 ++- 10 files changed, 1397 insertions(+), 531 deletions(-) diff --git a/README.md b/README.md index 8a0aacc2..045a5bf8 100644 --- a/README.md +++ b/README.md @@ -89,9 +89,19 @@ The runner can drive three SDKs behind a common interface: multi-personality handoffs, both `chat_completions` and `responses` `api_type`, `temperature`, `parallel_tool_calls`, `exclude_from_context`, and MCP over stdio, SSE, and streamable HTTP. -- **`copilot_sdk`** — the GitHub Copilot Python SDK. Supports streaming, +- **`copilot_sdk`** — the GitHub Copilot Python SDK (1.x). Supports streaming, `reasoning_effort`, MCP over stdio/SSE/HTTP, and per-tool permission - gating. The SDK selects its own wire protocol per model, so the YAML + gating. Authentication uses the env var named by `token` in + `model_settings` (default `COPILOT_GITHUB_TOKEN`); a named but empty + variable is an error. Without a token override, the SDK's default + authentication chain applies. The requested model must be available to + the account; runtime switches are rejected unless `auto` was explicitly + selected. Only taskflow-configured MCP tools are exposed, through the + existing SecLab wrappers: `blocked_tools`, toolbox confirmations, and + the explicit `headless` opt-in retain their usual behavior. Other SDK + tools and permission kinds are denied. `exclude_from_context` uses + terminal tools so successful tool results are not sent back to the model. + The SDK selects its own wire protocol per model, so the YAML `api_type` field is not honoured; multi-personality handoffs, `temperature`, and `parallel_tool_calls` are likewise not available. Taskflows that use unsupported fields fail at load time with a @@ -261,7 +271,8 @@ The two models are defined in [multi_model.yaml](examples/model_configs/multi_mo ## Requirements -Python >= 3.10 or Docker +Python >= 3.11 or Docker. Copilot SDK 1.x requires Python 3.11+, so this +upgrade explicitly ends Python 3.10 support for the package. ## Configuration diff --git a/pyproject.toml b/pyproject.toml index 84873602..7271e7b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "seclab-taskflow-agent" dynamic = ["version"] description = "A taskflow agent for the SecLab project, enabling secure and automated workflow execution." readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.11" license = "MIT" keywords = [ "agentic-workflows", @@ -24,7 +24,6 @@ authors = [ classifiers = [ "Development Status :: 4 - Beta", "Programming Language :: Python", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -56,7 +55,7 @@ dependencies = [ "email-validator==2.3.0", "exceptiongroup==1.3.0", "fastmcp==3.2.0", - "github-copilot-sdk>=0.2.2,<0.3", + "github-copilot-sdk==1.0.28", "griffelib==2.1.0", "h11==0.16.0", "httpcore==1.0.9", diff --git a/src/seclab_taskflow_agent/runner.py b/src/seclab_taskflow_agent/runner.py index a2892d1b..6da6bdd5 100644 --- a/src/seclab_taskflow_agent/runner.py +++ b/src/seclab_taskflow_agent/runner.py @@ -723,9 +723,7 @@ async def deploy_task_agents( mcp_params = mcp_client_params(available_tools, toolboxes) server_prompts = [sp for _, (_, _, sp, _) in mcp_params.items()] - # Wrap each built MCP server in a neutral spec. The openai - # adapter unwraps ``params["_native"]``; the Copilot adapter reads - # ``kind`` plus the raw transport keys. + # Share the connected, policy-enforcing wrapper with each adapter. mcp_specs: list[MCPServerSpec] = [] for entry in entries: raw_params, *_ = mcp_params.get(entry.server.name, ({}, [], "", 0.0)) diff --git a/src/seclab_taskflow_agent/sdk/base.py b/src/seclab_taskflow_agent/sdk/base.py index 639f3b1e..6ee59775 100644 --- a/src/seclab_taskflow_agent/sdk/base.py +++ b/src/seclab_taskflow_agent/sdk/base.py @@ -71,10 +71,9 @@ class TokenUsage: class MCPServerSpec: """Backend-neutral MCP server descriptor. - The openai-agents adapter only needs the built native server - handle that the runner stores under ``params["_native"]``. The - Copilot SDK adapter consumes ``kind`` plus the transport keys - (``command``/``args``/``env``/``url``/``headers``). + The runner stores the connected, policy-enforcing server wrapper under + ``params["_native"]``. Adapters invoke tools through that wrapper so + transport lifecycle and toolbox confirmations remain runner-owned. """ name: str diff --git a/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py b/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py index 31ef437e..6b20d732 100644 --- a/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py +++ b/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py @@ -1,14 +1,18 @@ # SPDX-FileCopyrightText: GitHub, Inc. # SPDX-License-Identifier: MIT -"""GitHub Copilot SDK backend adapter. +"""GitHub Copilot SDK (1.x) backend adapter. -* :meth:`build` starts a :class:`CopilotClient`, creates a session bound - to the agent's model and MCP servers, and wires a queue that captures +* :meth:`build` starts a :class:`copilot.CopilotClient` authenticated with + the token from ``spec.token_env`` (default ``COPILOT_GITHUB_TOKEN``), + verifies the requested model is available, creates a session bound to + that model and the agent's policy-enforcing MCP tools, and wires a queue that captures every ``SessionEvent``. -* :meth:`run_streamed` calls ``session.send`` and drains the queue - yielding :class:`TextDelta` until ``session.idle`` (success) or - ``session.error`` (translated to :class:`BackendUnexpectedError`). +* :meth:`run_streamed` calls ``session.send`` and drains the queue, + yielding neutral stream events until the ``session.idle`` event for the + sent message (success) or a ``session.error`` (translated to a neutral + backend error). Any model switch away from the requested model is an + error, never a silent fallback. * :meth:`aclose` disconnects the session and stops the client. """ @@ -22,14 +26,33 @@ import os from collections.abc import AsyncIterator from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any from ..base import AgentSpec, StreamEvent, TextDelta, TokenUsage, ToolEnd -from ..errors import BackendBadRequestError, BackendCapabilityError, BackendUnexpectedError -from .mcp import build_mcp_config +from ..errors import ( + BackendBadRequestError, + BackendCapabilityError, + BackendError, + BackendRateLimitError, + BackendTimeoutError, + BackendUnexpectedError, +) +from .mcp import build_mcp_tools from .permissions import build_permission_handler -_VALID_REASONING = ("low", "medium", "high", "xhigh") +if TYPE_CHECKING: + from copilot.session import ReasoningEffort + +_VALID_REASONING = ("low", "medium", "high", "xhigh", "max") + +#: Env var consulted when the taskflow does not name one via ``token:``. +DEFAULT_TOKEN_ENV = "COPILOT_GITHUB_TOKEN" # noqa: S105 - environment variable name + +#: Copilot's router alias. Requesting it explicitly opts out of model pinning. +AUTO_MODEL = "auto" + +_BAD_REQUEST_STATUS = frozenset({400, 401, 403, 404, 422}) +_TIMEOUT_STATUS = frozenset({408, 504}) @dataclass @@ -38,18 +61,28 @@ class _CopilotHandle: session: Any event_queue: asyncio.Queue[Any] exclude_from_context: bool = False + model: str = "" def _resolve_token(token_env: str | None) -> str | None: - """Resolve a token from an environment variable name. + """Resolve the GitHub token from an environment variable *name*. - Returns the env var's value, or ``None`` if the var is unset/empty. - Never returns the env var *name* itself — that would send a literal - string like ``"AI_API_TOKEN"`` as a bearer token. + * ``token_env`` set: the variable must be non-empty, otherwise + :class:`BackendBadRequestError` is raised (fail closed rather than + silently using whatever identity the Copilot CLI is logged in as). + * ``token_env`` unset: :data:`DEFAULT_TOKEN_ENV` is used when present; + otherwise ``None`` is returned and the SDK's default auth chain + applies. + + Never returns the env var *name* itself and never includes the token + value in error messages. """ - if not token_env: - return None - return os.getenv(token_env) or None + if token_env: + value = os.getenv(token_env) + if not value: + raise BackendBadRequestError(f"copilot_sdk: token environment variable {token_env!r} is unset or empty") + return value + return os.getenv(DEFAULT_TOKEN_ENV) or None def _normalize_model(model: str) -> str: @@ -65,20 +98,19 @@ def _normalize_model(model: str) -> str: invalidate any reproducibility guarantee a taskflow makes about the model under test. """ + model = model.split("/", 1)[-1].strip() if not model: - raise BackendBadRequestError( - "copilot_sdk: model is required (the SDK would otherwise pick a default)" - ) - return model.split("/", 1)[1] if "/" in model else model + raise BackendBadRequestError("copilot_sdk: model is required (the SDK would otherwise pick a default)") + return model def _tool_result_text(data: Any) -> str | None: - """Extract the textual payload from a ``TOOL_EXECUTION_COMPLETE`` event. + """Extract the textual payload from a ``tool.execution_complete`` event. The Copilot SDK delivers MCP tool results as a ``Result`` object with a - list of ``ContentElement`` entries. We concatenate the ``text`` fields - of the text elements, mirroring how openai-agents stringifies an MCP - text content list before handing it to the runner's tool-end hook. + list of content entries. We concatenate the ``text`` fields of the text + entries, mirroring how openai-agents stringifies an MCP text content + list before handing it to the runner's tool-end hook. """ result = getattr(data, "result", None) if result is None: @@ -91,18 +123,94 @@ def _tool_result_text(data: Any) -> str | None: return content if isinstance(content, str) else None -def _reasoning_effort(model_settings: dict[str, Any]) -> str | None: +def _reasoning_effort(model_settings: dict[str, Any]) -> ReasoningEffort | None: raw = model_settings.get("reasoning_effort") if raw is None: return None if raw not in _VALID_REASONING: raise BackendBadRequestError( - f"copilot_sdk: invalid reasoning_effort {raw!r} " - f"(expected one of {_VALID_REASONING})" + f"copilot_sdk: invalid reasoning_effort {raw!r} (expected one of {_VALID_REASONING})" ) return raw +def _translate_exception(exc: Exception, action: str) -> Exception: + """Map an exception raised by an SDK call to a neutral backend error.""" + if isinstance(exc, BackendError): + return exc + if isinstance(exc, TimeoutError): + return BackendTimeoutError(f"copilot_sdk: {action} timed out") + return BackendUnexpectedError(f"copilot_sdk: {action} failed: {exc}") + + +def _session_error(data: Any) -> Exception: + """Translate ``session.error`` event data into a neutral backend error.""" + message = getattr(data, "message", None) or "session error" + error_type = str(getattr(data, "error_type", "") or "").lower() + status = getattr(data, "status_code", None) + detail = f"copilot_sdk: {message}" + if error_type: + detail += f" (error_type={error_type}" + detail += f", status={status})" if status is not None else ")" + if error_type == "rate_limit" or status == 429: + return BackendRateLimitError(detail) + if error_type == "timeout" or status in _TIMEOUT_STATUS: + return BackendTimeoutError(detail) + if error_type in ("authentication", "authorization") or status in _BAD_REQUEST_STATUS: + return BackendBadRequestError(detail) + return BackendUnexpectedError(detail) + + +def _model_mismatch(expected: str, actual: str | None) -> bool: + return bool(actual) and actual != expected + + +def _check_model(event: Any, expected: str) -> None: + """Raise if *event* shows the session running a model other than *expected*.""" + from copilot.session_events import SessionEventType + + etype = getattr(event, "type", None) + data = getattr(event, "data", None) + if etype == SessionEventType.SESSION_MODEL_CHANGE: + actual = getattr(data, "new_model", None) + source = "session.model_change" + elif etype == SessionEventType.SESSION_START: + actual = getattr(data, "selected_model", None) + source = "session.start" + elif etype in (SessionEventType.ASSISTANT_MESSAGE, SessionEventType.ASSISTANT_USAGE): + actual = getattr(data, "model", None) + source = str(etype.value) + else: + return + if _model_mismatch(expected, actual): + raise BackendUnexpectedError( + f"copilot_sdk: requested model {expected!r} but {source} reports " + f"{actual!r}; refusing to continue with a different model" + ) + + +async def _verify_model(client: Any, model: str) -> None: + """Ensure *model* is in the authenticated user's catalogue and enabled.""" + models = await client.list_models() + for info in models: + if getattr(info, "id", None) != model: + continue + policy = getattr(info, "policy", None) + if getattr(policy, "state", None) == "disabled": + raise BackendBadRequestError(f"copilot_sdk: model {model!r} is disabled by policy for this account") + return + raise BackendBadRequestError( + f"copilot_sdk: model {model!r} is not available to this account ({len(models)} models listed)" + ) + + +async def _stop_client(client: Any) -> None: + try: + await client.stop() + except Exception: + logging.exception("copilot_sdk: client.stop failed") + + class CopilotSDKBackend: """Adapter that drives the GitHub Copilot SDK.""" @@ -132,14 +240,10 @@ def validate(self, spec: AgentSpec) -> None: wire protocol per model. """ if spec.handoffs or spec.in_handoff_graph: - raise BackendCapabilityError( - "copilot_sdk: agent handoffs are not supported" - ) + raise BackendCapabilityError("copilot_sdk: agent handoffs are not supported") for unsupported in ("temperature", "parallel_tool_calls"): if unsupported in spec.model_settings: - raise BackendCapabilityError( - f"copilot_sdk: model_settings.{unsupported} is not supported" - ) + raise BackendCapabilityError(f"copilot_sdk: model_settings.{unsupported} is not supported") async def build( self, @@ -149,13 +253,13 @@ async def build( agent_hooks: Any = None, ) -> _CopilotHandle: del run_hooks, agent_hooks # no parity yet - from copilot import CopilotClient, SubprocessConfig - from copilot.session import SystemMessageReplaceConfig + from copilot import CopilotClient, ToolSet + model = _normalize_model(spec.model) + reasoning_effort = _reasoning_effort(spec.model_settings) token = _resolve_token(spec.token_env) - config = SubprocessConfig(github_token=token) if token else SubprocessConfig() - client = CopilotClient(config, auto_start=False) - await client.start() + + client = CopilotClient(github_token=token, use_logged_in_user=False) if token else CopilotClient() loop = asyncio.get_running_loop() event_queue: asyncio.Queue[Any] = asyncio.Queue() @@ -166,39 +270,57 @@ def _on_event(event: Any) -> None: loop.call_soon_threadsafe(event_queue.put_nowait, event) try: + await client.start() + auth = await client.get_auth_status() + if not auth.isAuthenticated: + raise BackendBadRequestError( + "copilot_sdk: authentication failed; provide a Copilot-enabled " + "GitHub token or sign in with the Copilot CLI" + ) + await _verify_model(client, model) + tools = await build_mcp_tools( + spec.mcp_servers, + blocked_tools=spec.blocked_tools, + exclude_from_context=spec.exclude_from_context, + ) + available_tools = ToolSet() + for tool in tools: + available_tools.add_custom(tool.name) session = await client.create_session( - model=_normalize_model(spec.model), - on_permission_request=build_permission_handler( - spec.blocked_tools, headless=spec.headless - ), + model=model, + on_permission_request=build_permission_handler([tool.name for tool in tools]), + tools=tools, + available_tools=available_tools, streaming=True, - mcp_servers=build_mcp_config(spec.mcp_servers) or None, - excluded_tools=spec.blocked_tools or None, - reasoning_effort=_reasoning_effort(spec.model_settings), - system_message=( - SystemMessageReplaceConfig(content=spec.instructions) - if spec.instructions - else None - ), - # Pin everything that could otherwise be silently changed - # by user-level configuration: no on-disk config discovery, - # no built-in custom agents, no skill auto-loading. The - # taskflow YAML is the single source of truth. + include_sub_agent_streaming_events=False, + reasoning_effort=reasoning_effort, + system_message=({"mode": "replace", "content": spec.instructions} if spec.instructions else None), + # Preserve SDK authentication defaults without discovering + # ambient instructions, hooks, tools or cross-session memory. enable_config_discovery=False, - custom_agents=None, - agent=None, - skill_directories=[], + skip_custom_instructions=True, + enable_on_demand_instruction_discovery=False, + enable_file_hooks=False, + enable_host_git_operations=False, + enable_session_store=False, + enable_skills=False, + skip_embedding_retrieval=True, + memory={"enabled": False}, on_event=_on_event, ) - except Exception: - await client.stop() - raise + except BaseException as exc: + await _stop_client(client) + translated = _translate_exception(exc, "session setup") if isinstance(exc, Exception) else exc + if translated is exc: + raise + raise translated from exc return _CopilotHandle( client=client, session=session, event_queue=event_queue, exclude_from_context=spec.exclude_from_context, + model=model, ) async def run_streamed( @@ -209,62 +331,144 @@ async def run_streamed( max_turns: int, ) -> AsyncIterator[StreamEvent]: del max_turns # the SDK manages turn limits internally - from copilot.generated.session_events import SessionEventType - - await agent.session.send(prompt) - - aborted = False + from copilot.session_events import SessionEventType + + expected_model = getattr(agent, "model", "") or "" + routed_models: set[str] = set() + + def _enforce_model(event: Any) -> None: + if not expected_model: + return + if expected_model == AUTO_MODEL: + # The taskflow explicitly asked Copilot to route; record + # which concrete model served it instead of pinning. + actual = getattr(getattr(event, "data", None), "model", None) + if actual and actual not in routed_models: + routed_models.add(actual) + logging.info("copilot_sdk: model 'auto' routed to %r", actual) + return + _check_model(event, expected_model) + + started = False + completed = False + tool_names: dict[str, str] = {} + streamed_message_ids: set[str] = set() try: + # Discard an earlier turn's events, retaining model-pin checks. + while True: + try: + stale = agent.event_queue.get_nowait() + except asyncio.QueueEmpty: + break + if not getattr(stale, "agent_id", None): + _enforce_model(stale) + + try: + message_id = await agent.session.send(prompt) + except Exception as exc: + translated = _translate_exception(exc, "session.send") + if translated is exc: + raise + raise translated from exc + while True: event = await agent.event_queue.get() etype = getattr(event, "type", None) data = getattr(event, "data", None) + if getattr(event, "agent_id", None): + # Sub-agent events: account for usage only. + if etype == SessionEventType.ASSISTANT_USAGE: + yield _token_usage(data) + continue + + _enforce_model(event) + + if ( + etype == SessionEventType.SESSION_WARNING + and getattr(data, "warning_type", None) == "policy_blocked" + and getattr(data, "originating_message_id", None) == message_id + ): + raise BackendBadRequestError( + f"copilot_sdk: {getattr(data, 'message', None) or 'Request blocked by policy'}" + ) + + if not started: + if etype == SessionEventType.USER_MESSAGE: + started = getattr(data, "message_id", None) == message_id + continue + if etype == SessionEventType.ASSISTANT_MESSAGE and ( + getattr(data, "originating_message_id", None) == message_id + ): + started = True + elif etype == SessionEventType.SESSION_ERROR and getattr(data, "originating_message_id", None) in ( + None, + message_id, + ): + raise _session_error(data) + else: + continue + if etype == SessionEventType.ASSISTANT_MESSAGE_DELTA: + if getattr(data, "parent_tool_call_id", None): + continue text = getattr(data, "delta_content", None) or "" if text: + streamed_message_ids.add(getattr(data, "message_id", None) or "") yield TextDelta(text=text) + elif etype == SessionEventType.ASSISTANT_MESSAGE: + # Fallback for messages that were not streamed as deltas. + if getattr(data, "parent_tool_call_id", None): + continue + if (getattr(data, "message_id", None) or "") not in streamed_message_ids: + content = getattr(data, "content", None) or "" + if content: + yield TextDelta(text=content) elif etype == SessionEventType.ASSISTANT_USAGE: - # Copilot reports usage from the chat_completions surface, - # so (like OpenAI) ``input_tokens`` INCLUDES cached tokens - # and ``cache_read_tokens`` is a subset already counted in - # it, not an additional amount. - yield TokenUsage( - model=getattr(data, "model", "") or "", - input_tokens=int(getattr(data, "input_tokens", 0) or 0), - output_tokens=int(getattr(data, "output_tokens", 0) or 0), - cache_read_tokens=int(getattr(data, "cache_read_tokens", 0) or 0), - cache_write_tokens=int(getattr(data, "cache_write_tokens", 0) or 0), - ) + yield _token_usage(data) + elif etype == SessionEventType.TOOL_EXECUTION_START: + call_id = getattr(data, "tool_call_id", None) + if call_id: + tool_names[call_id] = ( + getattr(data, "mcp_tool_name", None) or getattr(data, "tool_name", None) or "" + ) elif etype == SessionEventType.TOOL_EXECUTION_COMPLETE: - if getattr(data, "success", False): - text = _tool_result_text(data) - if text is not None: - yield ToolEnd( - tool_name=getattr(data, "tool_name", "") or "", - text=text, - ) - # exclude_from_context parity with openai-agents' - # ToolsToFinalOutputFunction: stop the session after - # the first successful tool call so the model never - # sees the result fed back into its context. - if agent.exclude_from_context: - with contextlib.suppress(Exception): - await agent.session.abort() - aborted = True - return + call_id = getattr(data, "tool_call_id", None) + tool_name = tool_names.pop(call_id, "") if call_id else "" + if not getattr(data, "success", False): + error = getattr(data, "error", None) + logging.debug( + "copilot_sdk: tool %r failed: %s", + tool_name, + getattr(error, "message", None) or "unknown error", + ) + continue + text = _tool_result_text(data) + if text is not None: + yield ToolEnd(tool_name=tool_name, text=text) + # The registered terminal tool prevents a follow-up + # model call; abort also releases any remaining work. + if agent.exclude_from_context: + with contextlib.suppress(Exception): + await agent.session.abort() + completed = True + return elif etype == SessionEventType.SESSION_ERROR: - raise BackendUnexpectedError( - getattr(data, "message", None) or "copilot_sdk session error" - ) + if getattr(data, "originating_message_id", None) in (None, message_id): + raise _session_error(data) elif etype == SessionEventType.SESSION_IDLE: + mode = getattr(data, "mode", None) + if getattr(mode, "value", mode) == "autopilot": + continue + if getattr(data, "aborted", False): + raise BackendUnexpectedError("copilot_sdk: session aborted before the response completed") + completed = True return finally: # If the consumer abandoned the iterator (idle-timeout, error, # caller break) the session may still be mid-turn. Abort so the - # background event task drains and aclose() can run cleanly. - # Skipped when we already aborted above for exclude_from_context. - if not aborted: + # runtime stops generating and aclose() can run cleanly. + if not completed: with contextlib.suppress(Exception): await agent.session.abort() @@ -272,10 +476,22 @@ async def aclose(self, agent: _CopilotHandle | None) -> None: if agent is None: return try: - await agent.session.disconnect() - except Exception: - logging.exception("copilot_sdk: session.disconnect failed") - try: - await agent.client.stop() - except Exception: - logging.exception("copilot_sdk: client.stop failed") + try: + await agent.session.disconnect() + except Exception: + logging.exception("copilot_sdk: session.disconnect failed") + finally: + await _stop_client(agent.client) + + +def _token_usage(data: Any) -> TokenUsage: + # Copilot reports usage from the chat_completions surface, so (like + # OpenAI) ``input_tokens`` INCLUDES cached tokens and + # ``cache_read_tokens`` is a subset already counted in it. + return TokenUsage( + model=getattr(data, "model", "") or "", + input_tokens=int(getattr(data, "input_tokens", 0) or 0), + output_tokens=int(getattr(data, "output_tokens", 0) or 0), + cache_read_tokens=int(getattr(data, "cache_read_tokens", 0) or 0), + cache_write_tokens=int(getattr(data, "cache_write_tokens", 0) or 0), + ) diff --git a/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py b/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py index 8cb23cd3..e654a2bf 100644 --- a/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py +++ b/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py @@ -1,85 +1,102 @@ # SPDX-FileCopyrightText: GitHub, Inc. # SPDX-License-Identifier: MIT -"""Translate neutral :class:`MCPServerSpec` values into Copilot SDK MCP configs.""" +"""Expose runner-owned MCP tools through the public Copilot SDK tool API.""" from __future__ import annotations -__all__ = ["build_mcp_config"] +__all__ = ["build_mcp_tools"] -import os -import shutil -from typing import Any +import json +from typing import TYPE_CHECKING, Any from ..base import MCPServerSpec - - -def _resolve_command(command: str) -> str: - """Return an absolute path for *command*. - - The Copilot SDK spawns MCP servers in a fresh subprocess that does not - inherit our active virtualenv, so a bare ``python`` resolves to the - system interpreter and the toolbox import fails. Resolving against the - runner's ``PATH`` (which has the venv prepended) gives the SDK an - absolute path that points at the correct interpreter. +from ..errors import BackendCapabilityError + +if TYPE_CHECKING: + from copilot import Tool, ToolInvocation, ToolResult + + +def _tool_handler(server: Any, name: str): + async def _invoke(invocation: ToolInvocation) -> ToolResult: + from copilot import ToolBinaryResult, ToolResult + from mcp.types import BlobResourceContents, EmbeddedResource, ImageContent, TextContent + + if invocation.tool_name != name: + return ToolResult(result_type="denied", text_result_for_llm="Unknown taskflow tool") + try: + result = await server.call_tool(name, invocation.arguments or {}) + except (EOFError, OSError): + return ToolResult( + result_type="denied", + text_result_for_llm="Required toolbox confirmation is unavailable", + ) + + texts: list[str] = [] + binaries: list[ToolBinaryResult] = [] + for content in result.content: + if isinstance(content, TextContent): + texts.append(content.text) + elif isinstance(content, ImageContent): + binaries.append(ToolBinaryResult(data=content.data, mime_type=content.mimeType, type="image")) + elif isinstance(content, EmbeddedResource): + resource = content.resource + if isinstance(resource, BlobResourceContents): + binaries.append( + ToolBinaryResult( + data=resource.blob, + mime_type=resource.mimeType or "application/octet-stream", + type="resource", + description=str(resource.uri), + ) + ) + else: + texts.append(resource.text) + else: + texts.append(content.model_dump_json(exclude_none=True)) + if not texts and result.structuredContent is not None: + texts.append(json.dumps(result.structuredContent)) + return ToolResult( + result_type="failure" if result.isError else "success", + text_result_for_llm="".join(texts), + binary_results_for_llm=binaries or None, + ) + + return _invoke + + +async def build_mcp_tools( + specs: list[MCPServerSpec], + *, + blocked_tools: list[str], + exclude_from_context: bool, +) -> list[Tool]: + """Keep MCP execution behind SecLab's confirmation and transport wrappers. + + Registering the raw servers with Copilot as well would bypass those + wrappers and start duplicate MCP connections. Terminal tools enforce + ``exclude_from_context`` before the runtime can make another model call. """ - if not command or os.path.isabs(command): - return command - return shutil.which(command) or command - - -def _timeout_ms(params: dict[str, Any]) -> int | None: - timeout_s = params.get("timeout") - return int(timeout_s * 1000) if timeout_s else None - + from copilot import Tool -def _stdio_config(spec: MCPServerSpec) -> dict[str, Any]: - p = spec.params - config: dict[str, Any] = { - "type": "stdio", - "command": _resolve_command(p.get("command", "")), - "args": list(p.get("args") or []), - "tools": ["*"], # session-level excluded_tools handles filtering - } - if p.get("env"): - config["env"] = dict(p["env"]) - # The Copilot SDK does not inherit the runner's cwd by default, which - # breaks toolboxes that resolve state directories from relative paths - # (e.g. ``MEMCACHE_STATE_DIR=example_taskflow/``). Pin to the current - # process directory unless the toolbox specifies one explicitly. - config["cwd"] = p.get("cwd") or os.getcwd() - timeout = _timeout_ms(p) - if timeout: - config["timeout"] = timeout - return config - - -def _remote_config(spec: MCPServerSpec) -> dict[str, Any]: - p = spec.params - config: dict[str, Any] = { - "type": "sse" if spec.kind == "sse" else "http", - "url": p.get("url", ""), - "tools": ["*"], - } - if p.get("headers"): - config["headers"] = dict(p["headers"]) - timeout = _timeout_ms(p) - if timeout: - config["timeout"] = timeout - return config - - -def build_mcp_config(specs: list[MCPServerSpec]) -> dict[str, dict[str, Any]]: - """Return a ``mcp_servers`` mapping suitable for ``create_session``. - - Remote specs without a ``url`` (the runner is hosting the transport - in-process) are skipped — the SDK cannot consume an in-process - transport. - """ - out: dict[str, dict[str, Any]] = {} + tools: list[Tool] = [] + blocked = set(blocked_tools) for spec in specs: - if spec.kind == "stdio": - out[spec.name] = _stdio_config(spec) - elif spec.kind in ("sse", "streamable") and spec.params.get("url"): - out[spec.name] = _remote_config(spec) - return out + server = spec.params.get("_native") + if server is None: + raise BackendCapabilityError( + f"copilot_sdk: MCP server {spec.name!r} must be connected by the taskflow runner" + ) + for tool in await server.list_tools_unfiltered(): + if tool.name in blocked or tool.name.removeprefix(server.namespace) in blocked: + continue + tools.append( + Tool( + name=tool.name, + description=tool.description or "", + parameters=tool.inputSchema, + handler=_tool_handler(server, tool.name), + is_terminal=exclude_from_context, + ) + ) + return tools diff --git a/src/seclab_taskflow_agent/sdk/copilot_sdk/permissions.py b/src/seclab_taskflow_agent/sdk/copilot_sdk/permissions.py index e918a72c..ac38c89e 100644 --- a/src/seclab_taskflow_agent/sdk/copilot_sdk/permissions.py +++ b/src/seclab_taskflow_agent/sdk/copilot_sdk/permissions.py @@ -1,53 +1,38 @@ # SPDX-FileCopyrightText: GitHub, Inc. # SPDX-License-Identifier: MIT -"""Permission handling for the Copilot SDK backend. - -The Copilot SDK requires every session to provide an -``on_permission_request`` callback. The taskflow YAML exposes a single -policy knob, ``blocked_tools``: any tool whose name or raw -command/path/url appears in that list is denied (exact string match); -everything else is approved. The ``headless`` flag is accepted for -signature parity with the openai-agents path but is unused — the -runner never has a TTY, so there is no interactive fallback to prompt -for. +"""Allow only registered, policy-enforcing taskflow tool handlers. + +An SDK approval permits dispatch to the wrapper, not execution without its +confirmation checks. Raw MCP requests, built-in tools, unknown tools and +managed-approval requests are denied, including reads and URL fetches. """ from __future__ import annotations __all__ = ["build_permission_handler"] -from collections.abc import Callable -from typing import Any - +from collections.abc import Callable, Collection +from typing import TYPE_CHECKING -def _request_field(request: Any, *names: str) -> str | None: - for name in names: - value = getattr(request, name, None) - if value: - return str(value) - return None +if TYPE_CHECKING: + from copilot import PermissionRequest, PermissionRequestResult + from copilot.session import PermissionInvocation def build_permission_handler( - blocked_tools: list[str], - *, - headless: bool, -) -> Callable[[Any, dict[str, str]], Any]: - # Imported lazily so test code that does not exercise the handler - # can import this module without the SDK installed. - from copilot.session import PermissionRequestResult - - del headless # see module docstring - blocked = set(blocked_tools) - - def _handler(request: Any, _invocation: dict[str, str]) -> Any: - tool_id = _request_field(request, "tool_name", "full_command_text", "path", "url") - if tool_id and tool_id in blocked: - return PermissionRequestResult( - kind="denied-by-rules", - message=f"Tool {tool_id!r} is blocked by taskflow configuration", - ) - return PermissionRequestResult(kind="approved") + allowed_tools: Collection[str], +) -> Callable[[PermissionRequest, PermissionInvocation], PermissionRequestResult]: + from copilot.rpc import PermissionDecisionApproveOnce, PermissionDecisionReject + from copilot.session_events import PermissionRequestCustomTool + + allowed = frozenset(allowed_tools) + + def _handler(request: PermissionRequest, _invocation: PermissionInvocation) -> PermissionRequestResult: + if getattr(request, "managed_approval_required", False): + return PermissionDecisionReject(feedback="Managed approval is required and cannot be granted by a taskflow") + if isinstance(request, PermissionRequestCustomTool) and request.tool_name in allowed: + return PermissionDecisionApproveOnce() + return PermissionDecisionReject(feedback="Only configured taskflow tools may be invoked") return _handler diff --git a/tests/test_sdk_copilot_adapter.py b/tests/test_sdk_copilot_adapter.py index 29ba1144..7ade5127 100644 --- a/tests/test_sdk_copilot_adapter.py +++ b/tests/test_sdk_copilot_adapter.py @@ -6,18 +6,33 @@ from __future__ import annotations import asyncio +import inspect +from collections.abc import Sequence from dataclasses import dataclass +from datetime import datetime, timezone from types import SimpleNamespace -from typing import Any +from typing import Any, ClassVar +from unittest.mock import AsyncMock +from uuid import uuid4 import pytest pytest.importorskip("copilot") -from copilot.generated.session_events import SessionEventType +import copilot +from copilot.session_events import ( + AssistantMessageData, + AssistantMessageDeltaData, + AssistantUsageData, + SessionEvent, + SessionEventType, + SessionIdleData, + UserMessageData, +) +from mcp.types import Tool from seclab_taskflow_agent.sdk import get_backend -from seclab_taskflow_agent.sdk.base import AgentSpec, MCPServerSpec, TextDelta, TokenUsage +from seclab_taskflow_agent.sdk.base import AgentSpec, MCPServerSpec, TextDelta, TokenUsage, ToolEnd from seclab_taskflow_agent.sdk.copilot_sdk.backend import ( CopilotSDKBackend, _reasoning_effort, @@ -26,6 +41,8 @@ from seclab_taskflow_agent.sdk.errors import ( BackendBadRequestError, BackendCapabilityError, + BackendRateLimitError, + BackendTimeoutError, BackendUnexpectedError, ) @@ -37,7 +54,7 @@ def test_get_backend_returns_copilot_sdk_instance(): def _spec(**overrides) -> AgentSpec: - base = { + base: dict[str, Any] = { "name": "a", "instructions": "", "model": "gpt-5-mini", @@ -74,202 +91,672 @@ def test_validate_rejects_handoff_targets(): backend.validate(_spec(handoffs=[_spec(name="b")])) +MODEL = "gpt-5.6-luna" +MSG_ID = "req-1" +SECRET = "offline-test-token" # noqa: S105 - offline fixture, not a credential +_SDK_CLIENT = copilot.CopilotClient +_CREATE_SIGNATURE = inspect.signature(copilot.CopilotClient.create_session) + + @dataclass class _Event: type: SessionEventType data: Any = None + agent_id: str | None = None + + +def _ev(etype: SessionEventType, agent_id: str | None = None, **data: Any) -> _Event: + return _Event(etype, SimpleNamespace(**data), agent_id) + + +def _started(message_id: str = MSG_ID) -> _Event: + return _ev(SessionEventType.USER_MESSAGE, message_id=message_id) + + +def _idle(aborted: bool = False) -> _Event: + return _ev(SessionEventType.SESSION_IDLE, aborted=aborted) + + +def _delta(text: str, message_id: str = "m1") -> _Event: + return _ev(SessionEventType.ASSISTANT_MESSAGE_DELTA, delta_content=text, message_id=message_id) class _FakeSession: - def __init__(self) -> None: + def __init__(self, events: list[Any] | None = None, on_event: Any = None) -> None: self.sent: list[str] = [] + self.aborted = 0 + self.disconnected = 0 + self._events = events or [] + self._on_event = on_event + self.send_error: BaseException | None = None async def send(self, prompt: str) -> str: + if self.send_error is not None: + raise self.send_error self.sent.append(prompt) - return "req-1" + # Simulate the SDK dispatching events from its reader thread. + if self._on_event is not None: + for ev in self._events: + self._on_event(ev) + return MSG_ID + + async def abort(self) -> None: + self.aborted += 1 + + async def disconnect(self) -> None: + self.disconnected += 1 + + +def _handle( + events: list[Any], + *, + stale: Sequence[Any] = (), + session: Any = None, + exclude: bool = False, + model: str = MODEL, +): + """Return a handle whose session emits *events* when ``send`` is called. + + *stale* events are queued before ``send`` (i.e. left over from a + previous turn). + """ + queue: asyncio.Queue[Any] = asyncio.Queue() + for ev in stale: + queue.put_nowait(ev) + session = session or _FakeSession() + session._events = events + session._on_event = queue.put_nowait + return SimpleNamespace( + client=None, + session=session, + event_queue=queue, + exclude_from_context=exclude, + model=model, + ) + + +async def _collect(handle: Any, prompt: str = "go") -> list[Any]: + async with asyncio.timeout(5): + return [ev async for ev in CopilotSDKBackend().run_streamed(handle, prompt, max_turns=10)] + + +def _stream(events: list[Any], **kw) -> tuple[Any, list[Any]]: + handle = _handle(events, **kw) + return handle, asyncio.run(_collect(handle)) + + +# --- run_streamed ----------------------------------------------------------- def test_run_streamed_translates_deltas_until_idle(): - async def _run(): - session = _FakeSession() - queue: asyncio.Queue[Any] = asyncio.Queue() - for ev in ( - _Event(SessionEventType.ASSISTANT_MESSAGE_DELTA, SimpleNamespace(delta_content="hi ")), - _Event(SessionEventType.ASSISTANT_MESSAGE_DELTA, SimpleNamespace(delta_content="there")), - _Event(SessionEventType.SESSION_IDLE), - ): - queue.put_nowait(ev) - handle = SimpleNamespace(client=None, session=session, event_queue=queue) - backend = CopilotSDKBackend() - return session, [ - event async for event in backend.run_streamed(handle, "go", max_turns=10) + prompt = "Explain CWE-476 in one sentence." + handle = _handle([_started(), _delta("hi "), _delta("there"), _idle(), _delta("late")]) + out = asyncio.run(_collect(handle, prompt)) + assert handle.session.sent == [prompt] + assert out == [TextDelta(text="hi "), TextDelta(text="there")] + assert handle.session.aborted == 0 + + +def test_run_streamed_ignores_events_before_own_user_message(): + # A late idle/delta from an earlier (aborted) turn must not end this one. + _, out = _stream([_delta("stale"), _idle(aborted=True), _started("other"), _started(), _delta("fresh"), _idle()]) + assert out == [TextDelta(text="fresh")] + + +def test_run_streamed_drains_stale_queue_before_send(): + _, out = _stream([_started(), _delta("new"), _idle()], stale=[_started(), _delta("old"), _idle()]) + assert out == [TextDelta(text="new")] + + +def test_run_streamed_falls_back_to_full_assistant_message(): + _, out = _stream( + [ + _started(), + _ev(SessionEventType.ASSISTANT_MESSAGE, content="whole", message_id="m1", model=MODEL), + _idle(), ] + ) + assert out == [TextDelta(text="whole")] - session, out = asyncio.run(_run()) - assert session.sent == ["go"] - assert out == [TextDelta(text="hi "), TextDelta(text="there")] +def test_run_streamed_does_not_duplicate_streamed_message(): + _, out = _stream( + [ + _started(), + _delta("part", message_id="m1"), + _ev(SessionEventType.ASSISTANT_MESSAGE, content="part", message_id="m1", model=MODEL), + _idle(), + ] + ) + assert out == [TextDelta(text="part")] -def test_run_streamed_emits_token_usage(): - """An assistant.usage event is translated into a neutral TokenUsage.""" - async def _run(): - queue: asyncio.Queue[Any] = asyncio.Queue() - for ev in ( - _Event( +def test_streaming_with_public_sdk_event_types(): + def event(etype, data): + return SessionEvent(type=etype, data=data, id=uuid4(), timestamp=datetime.now(timezone.utc)) + + _, out = _stream( + [ + event(SessionEventType.USER_MESSAGE, UserMessageData(content="go", message_id=MSG_ID)), + event( + SessionEventType.ASSISTANT_MESSAGE_DELTA, + AssistantMessageDeltaData( + delta_content="hello", + message_id="m1", + ), + ), + event( + SessionEventType.ASSISTANT_MESSAGE, + AssistantMessageData( + content="hello", + message_id="m1", + originating_message_id=MSG_ID, + model=MODEL, + ), + ), + event( SessionEventType.ASSISTANT_USAGE, - SimpleNamespace( - model="claude-x", - input_tokens=100, - output_tokens=20, - cache_read_tokens=64, - cache_write_tokens=8, + AssistantUsageData( + model=MODEL, + input_tokens=10, + output_tokens=1, + cache_read_tokens=4, ), ), - _Event(SessionEventType.SESSION_IDLE), - ): - queue.put_nowait(ev) - handle = SimpleNamespace(client=None, session=_FakeSession(), event_queue=queue) - return [ - event async for event in CopilotSDKBackend().run_streamed(handle, "go", max_turns=10) + event(SessionEventType.SESSION_IDLE, SessionIdleData()), + ] + ) + assert out == [ + TextDelta(text="hello"), + TokenUsage(model=MODEL, input_tokens=10, output_tokens=1, cache_read_tokens=4), + ] + + +def test_full_response_correlates_without_user_message_event(): + _, out = _stream( + [ + _ev( + SessionEventType.ASSISTANT_MESSAGE, + content="reply", + message_id="m1", + originating_message_id=MSG_ID, + model=MODEL, + ), + _idle(), ] + ) + assert out == [TextDelta(text="reply")] + - out = asyncio.run(_run()) +def test_run_streamed_emits_token_usage(): + _, out = _stream( + [ + _started(), + _ev( + SessionEventType.ASSISTANT_USAGE, + model=MODEL, + input_tokens=100, + output_tokens=20, + cache_read_tokens=64, + cache_write_tokens=None, + ), + _idle(), + ] + ) assert out == [ - TokenUsage( - model="claude-x", - input_tokens=100, - output_tokens=20, - cache_read_tokens=64, - cache_write_tokens=8, - ) + TokenUsage(model=MODEL, input_tokens=100, output_tokens=20, cache_read_tokens=64, cache_write_tokens=0) ] -def test_run_streamed_session_error_raises(): - async def _run(): - queue: asyncio.Queue[Any] = asyncio.Queue() - queue.put_nowait(_Event(SessionEventType.SESSION_ERROR, SimpleNamespace(message="boom"))) - handle = SimpleNamespace(client=None, session=_FakeSession(), event_queue=queue) - async for _ in CopilotSDKBackend().run_streamed(handle, "go", max_turns=10): - pass +def test_run_streamed_ignores_sub_agent_output_but_counts_usage(): + _, out = _stream( + [ + _started(), + _ev(SessionEventType.ASSISTANT_MESSAGE_DELTA, "sub", delta_content="sub", message_id="s1"), + _ev(SessionEventType.SESSION_IDLE, "sub", aborted=False), + _ev(SessionEventType.ASSISTANT_USAGE, "sub", model="other-model", input_tokens=5, output_tokens=1), + _delta("main"), + _idle(), + ] + ) + assert out == [TokenUsage(model="other-model", input_tokens=5, output_tokens=1), TextDelta(text="main")] + + +def _tool_events(name: str = "echo", success: bool = True) -> list[_Event]: + result = SimpleNamespace(contents=[SimpleNamespace(text="tool-out")], content="ignored") + return [ + _ev(SessionEventType.TOOL_EXECUTION_START, tool_call_id="c1", tool_name=f"tb-{name}", mcp_tool_name=name), + _ev( + SessionEventType.TOOL_EXECUTION_COMPLETE, + tool_call_id="c1", + success=success, + result=result if success else None, + error=None if success else SimpleNamespace(message="denied", code="denied"), + ), + ] + + +def test_run_streamed_correlates_tool_name_from_start_event(): + handle, out = _stream([_started(), *_tool_events(), _delta("post-tool"), _idle()]) + assert out == [ToolEnd(tool_name="echo", text="tool-out"), TextDelta(text="post-tool")] + assert handle.session.aborted == 0 - with pytest.raises(BackendUnexpectedError, match="boom"): - asyncio.run(_run()) + +def test_run_streamed_skips_failed_tool(): + _, out = _stream([_started(), *_tool_events(success=False), _idle()]) + assert out == [] + + +def test_run_streamed_tool_falls_back_to_content_string(): + complete = _ev( + SessionEventType.TOOL_EXECUTION_COMPLETE, + tool_call_id="c9", + success=True, + result=SimpleNamespace(contents=None, content="plain"), + ) + _, out = _stream([_started(), complete, _idle()]) + assert out == [ToolEnd(tool_name="", text="plain")] def test_run_streamed_aborts_after_tool_when_excluded(): - aborted: list[bool] = [] + handle, out = _stream([_started(), *_tool_events(), _delta("should-not-see"), _idle()], exclude=True) + assert handle.session.aborted == 1 + assert out == [ToolEnd(tool_name="echo", text="tool-out")] - class _AbortableSession(_FakeSession): - async def abort(self) -> None: - aborted.append(True) - async def _run(): - session = _AbortableSession() - queue: asyncio.Queue[Any] = asyncio.Queue() - result = SimpleNamespace( - contents=[SimpleNamespace(text="tool-out")], content=None - ) - for ev in ( - _Event( - SessionEventType.TOOL_EXECUTION_COMPLETE, - SimpleNamespace(success=True, result=result, tool_name="echo"), +@pytest.mark.parametrize( + ("data", "exc"), + [ + ({"error_type": "rate_limit", "message": "slow down", "status_code": 429}, BackendRateLimitError), + ({"error_type": "query", "message": "slow down", "status_code": 429}, BackendRateLimitError), + ({"error_type": "authentication", "message": "bad creds", "status_code": 401}, BackendBadRequestError), + ({"error_type": "query", "message": "no such model", "status_code": 404}, BackendBadRequestError), + ({"error_type": "query", "message": "gateway", "status_code": 504}, BackendTimeoutError), + ({"error_type": "quota", "message": "boom", "status_code": None}, BackendUnexpectedError), + ], +) +def test_run_streamed_session_error_mapping(data, exc): + handle = _handle([_started(), _ev(SessionEventType.SESSION_ERROR, originating_message_id=None, **data)]) + with pytest.raises(exc, match=data["message"]): + asyncio.run(_collect(handle)) + assert handle.session.aborted == 1 + + +def test_run_streamed_session_error_before_start_is_raised(): + handle = _handle( + [ + _ev( + SessionEventType.SESSION_ERROR, + error_type="authentication", + message="denied", + status_code=None, + originating_message_id=MSG_ID, + ) + ] + ) + with pytest.raises(BackendBadRequestError, match="denied"): + asyncio.run(_collect(handle)) + + +def test_policy_block_without_user_message_fails_instead_of_waiting(): + handle = _handle( + [ + _ev( + SessionEventType.SESSION_WARNING, + warning_type="policy_blocked", + message="Request blocked by policy", + originating_message_id=MSG_ID, ), - # idle should NOT be reached because exclude_from_context aborts - _Event( - SessionEventType.ASSISTANT_MESSAGE_DELTA, - SimpleNamespace(delta_content="should-not-see"), + _idle(), + ] + ) + with pytest.raises(BackendBadRequestError, match="blocked by policy"): + asyncio.run(_collect(handle)) + assert handle.session.aborted == 1 + + +def test_run_streamed_ignores_error_for_other_message(): + _, out = _stream( + [ + _started(), + _ev( + SessionEventType.SESSION_ERROR, + error_type="query", + message="x", + status_code=None, + originating_message_id="other", ), - _Event(SessionEventType.SESSION_IDLE), - ): - queue.put_nowait(ev) - handle = SimpleNamespace( - client=None, - session=session, - event_queue=queue, - exclude_from_context=True, - ) - return [ - ev async for ev in CopilotSDKBackend().run_streamed(handle, "go", max_turns=10) + _delta("ok"), + _idle(), ] + ) + assert out == [TextDelta(text="ok")] - out = asyncio.run(_run()) - assert aborted == [True] - # Only the ToolEnd is emitted; the trailing delta is not consumed. - assert len(out) == 1 - assert out[0].text == "tool-out" +def test_run_streamed_unexpected_abort_raises(): + with pytest.raises(BackendUnexpectedError, match="aborted"): + _stream([_started(), _delta("partial"), _idle(aborted=True)]) -def test_run_streamed_keeps_running_when_exclude_disabled(): - async def _run(): - queue: asyncio.Queue[Any] = asyncio.Queue() - result = SimpleNamespace( - contents=[SimpleNamespace(text="tool-out")], content=None - ) - for ev in ( - _Event( - SessionEventType.TOOL_EXECUTION_COMPLETE, - SimpleNamespace(success=True, result=result, tool_name="echo"), - ), - _Event( - SessionEventType.ASSISTANT_MESSAGE_DELTA, - SimpleNamespace(delta_content="post-tool"), - ), - _Event(SessionEventType.SESSION_IDLE), - ): - queue.put_nowait(ev) - handle = SimpleNamespace( - client=None, - session=_FakeSession(), - event_queue=queue, - exclude_from_context=False, - ) - return [ - ev async for ev in CopilotSDKBackend().run_streamed(handle, "go", max_turns=10) + +def test_autopilot_idle_does_not_complete_response(): + _, out = _stream( + [ + _started(), + _ev(SessionEventType.SESSION_IDLE, mode="autopilot"), + _delta("final"), + _idle(), ] + ) + assert out == [TextDelta(text="final")] - out = asyncio.run(_run()) - assert [type(e).__name__ for e in out] == ["ToolEnd", "TextDelta"] +@pytest.mark.parametrize( + "event", + [ + _ev(SessionEventType.SESSION_MODEL_CHANGE, new_model="gpt-4.1", previous_model=MODEL, cause="fallback"), + _ev(SessionEventType.ASSISTANT_USAGE, model="gpt-4.1", input_tokens=1, output_tokens=1), + _ev(SessionEventType.ASSISTANT_MESSAGE, content="hi", message_id="m1", model="gpt-4.1"), + ], + ids=lambda e: e.type.value, +) +def test_run_streamed_rejects_model_switch(event): + handle = _handle([_started(), event, _idle()]) + with pytest.raises(BackendUnexpectedError, match="refusing"): + asyncio.run(_collect(handle)) + assert handle.session.aborted == 1 + + +def test_run_streamed_auto_model_accepts_routed_model(caplog): + caplog.set_level("INFO") + _, out = _stream( + [ + _started(), + _ev(SessionEventType.SESSION_MODEL_CHANGE, new_model="gpt-6-luna", previous_model="auto"), + _delta("routed"), + _ev(SessionEventType.ASSISTANT_USAGE, model="gpt-6-luna", input_tokens=1, output_tokens=1), + _idle(), + ], + model="auto", + ) + assert out == [TextDelta(text="routed"), TokenUsage(model="gpt-6-luna", input_tokens=1, output_tokens=1)] + assert "'auto' routed to 'gpt-6-luna'" in caplog.text + + +def test_run_streamed_rejects_model_switch_in_stale_events(): + handle = _handle([], stale=[_ev(SessionEventType.SESSION_START, selected_model="gpt-4.1")]) + with pytest.raises(BackendUnexpectedError, match=r"gpt-4\.1"): + asyncio.run(_collect(handle)) + assert handle.session.sent == [] -def test_run_streamed_aborts_session_on_consumer_break(): - """If the consumer abandons the iterator mid-turn, the adapter must - still call session.abort() so the SDK background task can drain. - """ - aborted: list[bool] = [] - class _AbortableSession(_FakeSession): - async def abort(self) -> None: - aborted.append(True) +@pytest.mark.parametrize( + ("raised", "exc"), + [(TimeoutError("t"), BackendTimeoutError), (RuntimeError("closed"), BackendUnexpectedError)], +) +def test_run_streamed_send_errors_are_translated(raised, exc): + session = _FakeSession() + session.send_error = raised + with pytest.raises(exc) as info: + _stream([], session=session) + assert info.value.__cause__ is raised + assert session.aborted == 1 + + +def test_cancellation_during_send_aborts_session(): + session = _FakeSession() + session.send_error = asyncio.CancelledError() + with pytest.raises(asyncio.CancelledError): + _stream([], session=session) + assert session.aborted == 1 + +def test_cancellation_while_waiting_for_events_aborts_session(): async def _run(): - queue: asyncio.Queue[Any] = asyncio.Queue() - for ev in ( - _Event( - SessionEventType.ASSISTANT_MESSAGE_DELTA, - SimpleNamespace(delta_content="hello"), - ), - _Event( - SessionEventType.ASSISTANT_MESSAGE_DELTA, - SimpleNamespace(delta_content="world"), - ), - _Event(SessionEventType.SESSION_IDLE), - ): - queue.put_nowait(ev) - handle = SimpleNamespace( - client=None, - session=_AbortableSession(), - event_queue=queue, - exclude_from_context=False, - ) - # Take only the first event then break — exercises the finally - # path of the async generator. - async for _ev in CopilotSDKBackend().run_streamed(handle, "go", max_turns=10): + sent = asyncio.Event() + + class _WaitingSession(_FakeSession): + async def send(self, _prompt): + sent.set() + return MSG_ID + + handle = _handle([], session=_WaitingSession()) + task = asyncio.create_task(_collect(handle)) + await sent.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert handle.session.aborted == 1 + + asyncio.run(_run()) + + +def test_run_streamed_aborts_session_on_consumer_break(): + handle = _handle([_started(), _delta("hello"), _delta("world"), _idle()]) + + async def _run(): + async for _ev_ in CopilotSDKBackend().run_streamed(handle, "go", max_turns=10): break asyncio.run(_run()) - assert aborted == [True] + assert handle.session.aborted == 1 + + +# --- build ------------------------------------------------------------------ + + +class _FakeClient: + instances: ClassVar[list[_FakeClient]] = [] + models: ClassVar[list[Any]] = [] + start_error: ClassVar[BaseException | None] = None + create_error: ClassVar[BaseException | None] = None + authenticated: ClassVar[bool] = True + session_events: ClassVar[list[Any]] = [] + + def __init__(self, **kwargs: Any) -> None: + # The real public constructor validates options without starting a runtime. + _SDK_CLIENT(**kwargs) + self.kwargs = kwargs + self.started = False + self.stopped = False + self.create_kwargs: dict[str, Any] | None = None + self.session: _FakeSession | None = None + _FakeClient.instances.append(self) + + async def start(self) -> None: + if self.start_error is not None: + raise self.start_error + self.started = True + + async def list_models(self) -> list[Any]: + return self.models + + async def get_auth_status(self): + return copilot.GetAuthStatusResponse(isAuthenticated=self.authenticated) + + async def create_session(self, **kwargs: Any) -> _FakeSession: + _CREATE_SIGNATURE.bind(self, **kwargs) + if self.create_error is not None: + raise self.create_error + self.create_kwargs = kwargs + self.session = _FakeSession(self.session_events, kwargs["on_event"]) + return self.session + + async def stop(self) -> None: + self.stopped = True + + +def _model(model_id: str, state: str | None = "enabled") -> Any: + return SimpleNamespace(id=model_id, policy=SimpleNamespace(state=state) if state else None) + + +@pytest.fixture +def fake_client(monkeypatch): + monkeypatch.setattr(_FakeClient, "instances", []) + monkeypatch.setattr(_FakeClient, "models", [_model(MODEL), _model("off", "disabled")]) + monkeypatch.setattr(_FakeClient, "start_error", None) + monkeypatch.setattr(_FakeClient, "create_error", None) + monkeypatch.setattr(_FakeClient, "authenticated", True) + monkeypatch.setattr(_FakeClient, "session_events", []) + monkeypatch.setattr(copilot, "CopilotClient", _FakeClient) + monkeypatch.setenv("COPILOT_GITHUB_TOKEN", SECRET) + return _FakeClient + + +def _build(**overrides): + overrides.setdefault("model", MODEL) + overrides.setdefault("token_env", "COPILOT_GITHUB_TOKEN") + spec = _spec(**overrides) + return asyncio.run(CopilotSDKBackend().build(spec)) + + +def test_build_uses_explicit_token_and_pins_session(fake_client, caplog): + caplog.set_level("DEBUG") + native = SimpleNamespace( + namespace="ns", + list_tools_unfiltered=AsyncMock( + return_value=[ + Tool(name="nsecho", inputSchema={"type": "object"}), + Tool(name="nsbad_tool", inputSchema={"type": "object"}), + ] + ), + ) + handle = _build( + instructions="You are a security analyst.", + blocked_tools=["bad_tool"], + model_settings={"reasoning_effort": "max"}, + mcp_servers=[MCPServerSpec(name="tb", kind="stdio", params={"_native": native})], + ) + (client,) = fake_client.instances + assert client.kwargs == {"github_token": SECRET, "use_logged_in_user": False} + assert client.started + assert not client.stopped + kw = client.create_kwargs + assert kw["model"] == MODEL + assert kw["system_message"] == {"mode": "replace", "content": "You are a security analyst."} + assert kw["streaming"] is True + assert kw["reasoning_effort"] == "max" + assert kw["available_tools"].to_list() == ["custom:nsecho"] + assert [tool.name for tool in kw["tools"]] == ["nsecho"] + assert "mcp_servers" not in kw + assert kw["enable_config_discovery"] is False + assert kw["skip_custom_instructions"] is True + assert kw["include_sub_agent_streaming_events"] is False + assert kw["enable_skills"] is False + assert kw["enable_file_hooks"] is False + assert kw["enable_on_demand_instruction_discovery"] is False + assert kw["memory"] == {"enabled": False} + assert callable(kw["on_permission_request"]) + # No auto-mode handler: the SDK then declines any automatic model switch. + assert "on_auto_mode_switch_request" not in kw + assert handle.model == MODEL + assert SECRET not in caplog.text + + +def test_build_strips_provider_prefix(fake_client): + _build(model="openai/" + MODEL) + assert fake_client.instances[0].create_kwargs["model"] == MODEL + + +def test_build_without_instructions_sends_no_system_message(fake_client): + _build() + assert fake_client.instances[0].create_kwargs["system_message"] is None + + +def test_build_missing_token_env_fails_closed(fake_client, monkeypatch): + monkeypatch.delenv("SECLAB_TEST_MISSING_TOKEN", raising=False) + spec = _spec(model=MODEL, token_env="SECLAB_TEST_MISSING_TOKEN") # noqa: S106 - env var name + with pytest.raises(BackendBadRequestError, match="SECLAB_TEST_MISSING_TOKEN"): + asyncio.run(CopilotSDKBackend().build(spec)) + assert fake_client.instances == [] + + +def test_build_defaults_to_copilot_github_token(fake_client): + asyncio.run(CopilotSDKBackend().build(_spec(model=MODEL))) + assert fake_client.instances[0].kwargs["github_token"] == SECRET + + +def test_build_explicit_token_overrides_default_identity(fake_client, monkeypatch): + value = "different-offline-token" + monkeypatch.setenv("SECLAB_TEST_TOKEN", value) + _build(token_env="SECLAB_TEST_TOKEN") # noqa: S106 - env var name + assert fake_client.instances[0].kwargs["github_token"] == value + assert fake_client.instances[0].kwargs["use_logged_in_user"] is False + + +def test_build_without_any_token_uses_sdk_default_auth(fake_client, monkeypatch): + monkeypatch.delenv("COPILOT_GITHUB_TOKEN") + asyncio.run(CopilotSDKBackend().build(_spec(model=MODEL))) + assert fake_client.instances[0].kwargs == {} + + +def test_build_authentication_failure_is_actionable_and_stops_client(fake_client): + fake_client.authenticated = False + with pytest.raises(BackendBadRequestError, match="authentication failed"): + _build() + (client,) = fake_client.instances + assert client.create_kwargs is None + assert client.stopped + + +@pytest.mark.parametrize(("model", "match"), [("nope", "not available"), ("off", "disabled")]) +def test_build_rejects_unavailable_model_without_fallback(fake_client, model, match): + with pytest.raises(BackendBadRequestError, match=match): + _build(model=model) + (client,) = fake_client.instances + assert client.create_kwargs is None + assert client.stopped + + +@pytest.mark.parametrize( + ("attr", "raised", "exc"), + [ + ("start_error", RuntimeError("auth failed"), BackendUnexpectedError), + ("create_error", RuntimeError("bad mcp"), BackendUnexpectedError), + ("create_error", TimeoutError(), BackendTimeoutError), + ], +) +def test_build_failures_are_translated_and_stop_client(fake_client, monkeypatch, attr, raised, exc): + monkeypatch.setattr(_FakeClient, attr, raised) + with pytest.raises(exc) as info: + _build() + assert info.value.__cause__ is raised + assert fake_client.instances[0].stopped + assert SECRET not in str(info.value) + + +@pytest.mark.parametrize("attr", ["start_error", "create_error"]) +def test_build_cancellation_stops_client(fake_client, monkeypatch, attr): + monkeypatch.setattr(fake_client, attr, asyncio.CancelledError()) + with pytest.raises(asyncio.CancelledError): + _build() + assert fake_client.instances[0].stopped + + +def test_build_and_stream_end_to_end(fake_client): + fake_client.session_events = [ + _ev(SessionEventType.SESSION_START, selected_model=MODEL), + _started(), + _delta("CWE-476 is a NULL pointer dereference."), + _ev(SessionEventType.ASSISTANT_USAGE, model=MODEL, input_tokens=10, output_tokens=8), + _idle(), + ] + + async def _run(): + backend = CopilotSDKBackend() + handle = await backend.build(_spec(model=MODEL, token_env="COPILOT_GITHUB_TOKEN")) # noqa: S106 - env var name + try: + return handle, [ev async for ev in backend.run_streamed(handle, "Explain", max_turns=1)] + finally: + await backend.aclose(handle) + + handle, out = asyncio.run(_run()) + assert out == [ + TextDelta(text="CWE-476 is a NULL pointer dereference."), + TokenUsage(model=MODEL, input_tokens=10, output_tokens=8), + ] + assert handle.session.sent == ["Explain"] + assert handle.session.disconnected == 1 + assert handle.client.stopped def test_invalid_reasoning_effort_rejected(): @@ -277,12 +764,13 @@ def test_invalid_reasoning_effort_rejected(): _reasoning_effort({"reasoning_effort": "ludicrous"}) -def test_normalize_model_rejects_empty(): +@pytest.mark.parametrize("model", ["", " ", "openai/", "openai/ "]) +def test_normalize_model_rejects_empty(model): # Passing an empty model would let the SDK silently fall back to its # built-in default, which would invalidate any reproducibility # guarantee a taskflow makes about the model under test. with pytest.raises(BackendBadRequestError, match="model is required"): - _normalize_model("") + _normalize_model(model) def test_normalize_model_strips_provider_prefix(): @@ -307,7 +795,9 @@ async def stop(self) -> None: asyncio.run(CopilotSDKBackend().aclose(handle)) -def test_mcp_specs_are_used(): - # Sanity: MCPServerSpec is what build() consumes via build_mcp_config. - spec = _spec(mcp_servers=[MCPServerSpec(name="m", kind="stdio", params={"command": "x"})]) - CopilotSDKBackend().validate(spec) +def test_aclose_stops_client_even_if_disconnect_is_cancelled(): + session = SimpleNamespace(disconnect=AsyncMock(side_effect=asyncio.CancelledError())) + client = SimpleNamespace(stop=AsyncMock()) + with pytest.raises(asyncio.CancelledError): + asyncio.run(CopilotSDKBackend().aclose(SimpleNamespace(client=client, session=session))) + client.stop.assert_awaited_once() diff --git a/tests/test_sdk_copilot_mcp.py b/tests/test_sdk_copilot_mcp.py index 54f6fec1..c6877b74 100644 --- a/tests/test_sdk_copilot_mcp.py +++ b/tests/test_sdk_copilot_mcp.py @@ -1,118 +1,221 @@ # SPDX-FileCopyrightText: GitHub, Inc. # SPDX-License-Identifier: MIT -"""Tests for ``sdk.copilot_sdk.mcp``.""" +"""Offline tests for policy-enforcing Copilot MCP tools.""" from __future__ import annotations +import asyncio +import json import sys +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock import pytest +pytest.importorskip("copilot") + +from copilot import ToolInvocation +from mcp.types import CallToolResult, ImageContent, ListToolsResult, TextContent, Tool + +from seclab_taskflow_agent.mcp_utils import MCPNamespaceWrap from seclab_taskflow_agent.sdk.base import MCPServerSpec -from seclab_taskflow_agent.sdk.copilot_sdk.mcp import build_mcp_config +from seclab_taskflow_agent.sdk.copilot_sdk.mcp import build_mcp_tools +from seclab_taskflow_agent.sdk.errors import BackendCapabilityError + + +def _server(*, confirms=()): + transport = SimpleNamespace( + name="toolbox", + session=SimpleNamespace( + list_tools=AsyncMock( + return_value=ListToolsResult( + tools=[Tool(name="echo", description="Echo input", inputSchema={"type": "object"})], + ) + ) + ), + call_tool=AsyncMock( + return_value=CallToolResult( + content=[TextContent(type="text", text="tool-output")], + ) + ), + ) + return MCPNamespaceWrap(list(confirms), transport), transport -def test_stdio_spec_translated_to_subprocess_config(): - specs = [ - MCPServerSpec( - name="files", - kind="stdio", - params={ - "command": "/usr/bin/mcp", - "args": ["--root", "/tmp"], - "env": {"FOO": "1"}, - "cwd": "/work", - "timeout": 2.5, - }, +async def _tools(server, *, blocked=(), exclude=False, kind="stdio"): + return await build_mcp_tools( + [MCPServerSpec(name="toolbox", kind=kind, params={"_native": server})], + blocked_tools=list(blocked), + exclude_from_context=exclude, + ) + + +@pytest.mark.parametrize("kind", ["stdio", "sse", "streamable"]) +def test_registers_and_invokes_existing_mcp_connection(kind): + server, transport = _server() + + async def _run(): + (tool,) = await _tools(server, kind=kind) + assert tool.name == server.namespace + "echo" + assert tool.description == "Echo input" + assert tool.parameters == {"type": "object"} + assert not tool.skip_permission + assert not tool.is_terminal + return await tool.handler(ToolInvocation(tool_name=tool.name, arguments={"value": 42})) + + result = asyncio.run(_run()) + assert result.result_type == "success" + assert result.text_result_for_llm == "tool-output" + transport.call_tool.assert_awaited_once_with("echo", {"value": 42}) + + +@pytest.mark.parametrize("namespaced", [False, True]) +def test_blocked_tools_are_not_registered(namespaced): + server, transport = _server() + blocked = server.namespace + "echo" if namespaced else "echo" + assert asyncio.run(_tools(server, blocked=[blocked])) == [] + transport.call_tool.assert_not_awaited() + + +@pytest.mark.parametrize("approved", [False, True]) +def test_toolbox_confirmation_cannot_be_bypassed(approved): + server, transport = _server(confirms=["echo"]) + server.confirm_tool = Mock(return_value=approved) + + async def _run(): + (tool,) = await _tools(server) + return await tool.handler(ToolInvocation(tool_name=tool.name, arguments={"value": 42})) + + result = asyncio.run(_run()) + server.confirm_tool.assert_called_once_with("echo", [{"value": 42}]) + if approved: + transport.call_tool.assert_awaited_once() + assert result.text_result_for_llm == "tool-output" + else: + transport.call_tool.assert_not_awaited() + assert "not allowed" in result.text_result_for_llm + + +def test_headless_wrapper_retains_explicit_confirmation_opt_out(): + server, transport = _server(confirms=[]) + server.confirm_tool = Mock(side_effect=AssertionError("unexpected confirmation")) + + async def _run(): + (tool,) = await _tools(server) + await tool.handler(ToolInvocation(tool_name=tool.name)) + + asyncio.run(_run()) + server.confirm_tool.assert_not_called() + transport.call_tool.assert_awaited_once() + + +def test_missing_confirmation_fails_closed(): + server, transport = _server(confirms=["echo"]) + server.confirm_tool = Mock(side_effect=EOFError) + + async def _run(): + (tool,) = await _tools(server) + return await tool.handler(ToolInvocation(tool_name=tool.name)) + + result = asyncio.run(_run()) + assert result.result_type == "denied" + transport.call_tool.assert_not_awaited() + + +def test_invocation_cannot_select_another_tool(): + server, transport = _server() + + async def _run(): + (tool,) = await _tools(server) + return await tool.handler(ToolInvocation(tool_name="unregistered")) + + assert asyncio.run(_run()).result_type == "denied" + transport.call_tool.assert_not_awaited() + + +def test_exclude_from_context_registers_terminal_tools(): + server, _ = _server() + (tool,) = asyncio.run(_tools(server, exclude=True)) + assert tool.is_terminal + + +def test_mcp_errors_are_reported_as_tool_failures(): + server, transport = _server() + transport.call_tool.return_value = CallToolResult( + isError=True, + content=[TextContent(type="text", text="invalid key")], + ) + + async def _run(): + (tool,) = await _tools(server) + return await tool.handler(ToolInvocation(tool_name=tool.name)) + + result = asyncio.run(_run()) + assert result.result_type == "failure" + assert result.text_result_for_llm == "invalid key" + + +def test_preserves_structured_results_and_images(): + server, transport = _server() + transport.call_tool.return_value = CallToolResult( + content=[ImageContent(type="image", data="aW1hZ2U=", mimeType="image/png")], + structuredContent={"value": [1, 2, 3]}, + ) + + async def _run(): + (tool,) = await _tools(server) + return await tool.handler(ToolInvocation(tool_name=tool.name)) + + result = asyncio.run(_run()) + assert json.loads(result.text_result_for_llm) == {"value": [1, 2, 3]} + assert result.binary_results_for_llm[0].data == "aW1hZ2U=" + assert result.binary_results_for_llm[0].mime_type == "image/png" + + +def test_raw_mcp_specs_cannot_bypass_runner_policy(): + with pytest.raises(BackendCapabilityError, match="connected by the taskflow runner"): + asyncio.run( + build_mcp_tools( + [MCPServerSpec(name="unmanaged", kind="stdio", params={"command": "unused"})], + blocked_tools=[], + exclude_from_context=False, + ) ) - ] - assert build_mcp_config(specs) == { - "files": { - "type": "stdio", - "command": "/usr/bin/mcp", - "args": ["--root", "/tmp"], - "env": {"FOO": "1"}, - "cwd": "/work", - "tools": ["*"], - "timeout": 2500, - } - } - - -def test_sse_spec_translated_to_remote_config(): - specs = [ - MCPServerSpec( - name="docs", - kind="sse", + + +def test_memcache_round_trip_without_copilot_credentials(tmp_path): + from agents.mcp import MCPServerStdio + + async def _run(): + async with MCPServerStdio( params={ - "url": "https://example.com/mcp", - "headers": {"Authorization": "Bearer x"}, - "timeout": 10.0, + "command": sys.executable, + "args": ["-m", "seclab_taskflow_agent.mcp_servers.memcache.memcache"], + "env": { + "MEMCACHE_STATE_DIR": str(tmp_path), + "MEMCACHE_BACKEND": "dictionary_file", + "LOG_DIR": str(tmp_path), + }, }, - ) - ] - assert build_mcp_config(specs) == { - "docs": { - "type": "sse", - "url": "https://example.com/mcp", - "headers": {"Authorization": "Bearer x"}, - "tools": ["*"], - "timeout": 10000, - } - } - - -def test_streamable_spec_uses_http_type(): - specs = [ - MCPServerSpec( - name="gh", - kind="streamable", - params={"url": "https://api.example.com/mcp"}, - ) - ] - out = build_mcp_config(specs) - assert out["gh"]["type"] == "http" - assert out["gh"]["url"] == "https://api.example.com/mcp" - - -def test_streamable_spec_without_url_is_skipped(): - assert build_mcp_config([MCPServerSpec(name="x", kind="streamable", params={})]) == {} - - -def test_stdio_minimal_defaults(): - import os - - specs = [MCPServerSpec(name="m", kind="stdio", params={"command": "x"})] - assert build_mcp_config(specs) == { - "m": { - "type": "stdio", - "command": "x", - "args": [], - "tools": ["*"], - "cwd": os.getcwd(), - } - } - - -def test_stdio_resolves_bare_command_via_path(tmp_path, monkeypatch): - if sys.platform == "win32": - pytest.skip("PATHEXT-based shutil.which lookup is POSIX-specific in this test") - bin_dir = tmp_path / "bin" - bin_dir.mkdir() - fake = bin_dir / "fakecmd" - fake.write_text("#!/bin/sh\n") - fake.chmod(0o755) - monkeypatch.setenv("PATH", str(bin_dir)) - specs = [MCPServerSpec(name="m", kind="stdio", params={"command": "fakecmd"})] - out = build_mcp_config(specs) - assert out["m"]["command"] == str(fake) - - -def test_stdio_explicit_cwd_preserved(): - specs = [ - MCPServerSpec( - name="m", kind="stdio", params={"command": "/usr/bin/x", "cwd": "/tmp"} - ) - ] - assert build_mcp_config(specs)["m"]["cwd"] == "/tmp" + name="memcache", + ) as transport: + server = MCPNamespaceWrap([], transport) + tools = {t.name.removeprefix(server.namespace): t for t in await _tools(server)} + setter = tools["memcache_set_state"] + result = await setter.handler( + ToolInvocation( + tool_name=setter.name, + arguments={"key": "acc", "value": [1, 2, 3]}, + ) + ) + assert result.result_type == "success" + getter = tools["memcache_get_state"] + return await getter.handler(ToolInvocation(tool_name=getter.name, arguments={"key": "acc"})) + + async def _bounded(): + return await asyncio.wait_for(_run(), timeout=60) + + result = asyncio.run(_bounded()) + assert result.result_type == "success" + assert json.loads(result.text_result_for_llm) == [1, 2, 3] diff --git a/tests/test_sdk_copilot_permissions.py b/tests/test_sdk_copilot_permissions.py index c9179b71..6c352b1b 100644 --- a/tests/test_sdk_copilot_permissions.py +++ b/tests/test_sdk_copilot_permissions.py @@ -9,50 +9,98 @@ pytest.importorskip("copilot") -from dataclasses import dataclass +from copilot.rpc import PermissionDecisionApproveOnce, PermissionDecisionReject +from copilot.session_events import ( + PermissionRequestCustomTool, + PermissionRequestHook, + PermissionRequestMcp, + PermissionRequestMemory, + PermissionRequestRead, + PermissionRequestShell, + PermissionRequestUrl, + PermissionRequestWrite, +) from seclab_taskflow_agent.sdk.copilot_sdk.permissions import build_permission_handler -@dataclass -class _Req: - tool_name: str | None = None - full_command_text: str | None = None - path: str | None = None - url: str | None = None - - -def test_blocked_tool_is_denied_by_rules(): - handler = build_permission_handler(["dangerous"], headless=True) - result = handler(_Req(tool_name="dangerous"), {}) - assert result.kind == "denied-by-rules" - assert "dangerous" in (result.message or "") - - -def test_headless_approves_unblocked_tool(): - handler = build_permission_handler(["other"], headless=True) - result = handler(_Req(tool_name="safe"), {}) - assert result.kind == "approved" - - -def test_non_headless_still_approves_unblocked_tool(): - # The Copilot SDK requires a permission verdict; the runner has no - # interactive TTY, so any non-blocked tool is approved regardless of - # the YAML ``headless`` flag (the openai-agents path has the same - # effective behaviour). - handler = build_permission_handler([], headless=False) - result = handler(_Req(tool_name="safe"), {}) - assert result.kind == "approved" - - -def test_block_matches_command_text_when_tool_name_missing(): - handler = build_permission_handler(["rm -rf /"], headless=True) - result = handler(_Req(full_command_text="rm -rf /"), {}) - assert result.kind == "denied-by-rules" - - -def test_handler_signature_matches_sdk_callback() -> None: - # Sanity check: the SDK calls handler(request, invocation); make sure - # we accept exactly two positional arguments. - handler = build_permission_handler([], headless=True) - handler(_Req(tool_name="ok"), {"invocation": "x"}) # type: ignore[arg-type] +def _mcp(tool_name: str = "safe", **kw) -> PermissionRequestMcp: + return PermissionRequestMcp( + read_only=False, + server_name="toolbox", + tool_name=tool_name, + tool_title=tool_name, + **kw, + ) + + +def _shell(command: str = "ls") -> PermissionRequestShell: + return PermissionRequestShell( + can_offer_session_approval=False, + commands=[], + full_command_text=command, + has_write_file_redirection=False, + intention="run", + possible_paths=[], + possible_urls=[], + ) + + +def _write(path: str = "/tmp/x") -> PermissionRequestWrite: + return PermissionRequestWrite(can_offer_session_approval=False, diff="+x", file_name=path, intention="edit") + + +@pytest.mark.parametrize( + "request_", + [ + _mcp(), + PermissionRequestRead(intention="read", path="/etc/hostname"), + PermissionRequestUrl(intention="fetch", url="https://example.com"), + _shell(), + _write(), + PermissionRequestMemory(fact="remember me"), + PermissionRequestHook(tool_name="h"), + ], + ids=lambda r: r.kind, +) +def test_raw_mcp_and_builtin_permissions_are_denied(request_): + result = build_permission_handler(["safe"])(request_, {"session_id": "s"}) + assert isinstance(result, PermissionDecisionReject) + + +def test_only_registered_policy_enforcing_tools_are_approved(): + handler = build_permission_handler(["safe"]) + result = handler( + PermissionRequestCustomTool(tool_description="d", tool_name="safe"), + {"session_id": "s", "managed_settings_enabled": False}, + ) + assert isinstance(result, PermissionDecisionApproveOnce) + assert result.to_dict() == {"kind": "approve-once"} + + +def test_unknown_kind_is_denied(): + class _Future: + kind = "something-new" + + result = build_permission_handler([])(_Future(), {"session_id": "s"}) + assert isinstance(result, PermissionDecisionReject) + + +def test_unregistered_custom_tool_is_rejected(): + result = build_permission_handler(["safe"])( + PermissionRequestCustomTool(tool_description="d", tool_name="dangerous"), + {"session_id": "s"}, + ) + assert isinstance(result, PermissionDecisionReject) + + +def test_managed_approval_required_is_rejected(): + result = build_permission_handler(["safe"])( + PermissionRequestCustomTool( + tool_description="d", + tool_name="safe", + managed_approval_required=True, + ), + {"session_id": "s"}, + ) + assert isinstance(result, PermissionDecisionReject) From 0af51257f8a4b7b9b376592ad4914c0029e8f7be Mon Sep 17 00:00:00 2001 From: Sriharsha Etigowni Date: Fri, 30 Oct 2026 21:42:03 +0000 Subject: [PATCH 2/5] fix(copilot_sdk): address SDK and MCP review findings Allow authenticated auto routing, keep configured tools and large outputs accessible, and reject foreign assistant messages before model checks. Preserve structured MCP payloads and declared FastMCP result envelopes, paginate tool discovery through the policy wrapper, and add focused regression coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/seclab_taskflow_agent/mcp_utils.py | 25 ++++--- .../sdk/copilot_sdk/backend.py | 11 ++- .../sdk/copilot_sdk/mcp.py | 16 +++-- tests/test_mcp_utils.py | 30 +++++++- tests/test_sdk_copilot_adapter.py | 26 ++++++- tests/test_sdk_copilot_mcp.py | 70 +++++++++++++++++++ 6 files changed, 162 insertions(+), 16 deletions(-) diff --git a/src/seclab_taskflow_agent/mcp_utils.py b/src/seclab_taskflow_agent/mcp_utils.py index 84cc3134..ee6d7b55 100644 --- a/src/seclab_taskflow_agent/mcp_utils.py +++ b/src/seclab_taskflow_agent/mcp_utils.py @@ -109,7 +109,7 @@ async def list_tools(self, *args: Any, **kwargs: Any) -> list[Any]: return namespaced_tools async def list_tools_unfiltered(self) -> list[Any]: - """List tools directly from the MCP session, namespace-prefixed. + """List every page of tools from the MCP session, namespace-prefixed. Bypasses any tool_filter configured on the wrapped openai-agents server (which would require ``run_context`` and ``agent`` arguments @@ -125,7 +125,7 @@ async def list_tools_unfiltered(self) -> list[Any]: Raises ``RuntimeError`` if the underlying server has no active MCP session yet (caller should ensure the server is connected - before calling this). + before calling this), or if a pagination cursor repeats. """ session = getattr(self._obj, "session", None) if session is None: @@ -135,12 +135,21 @@ async def list_tools_unfiltered(self) -> list[Any]: ) result = await session.list_tools() namespaced_tools: list[Any] = [] - for tool in result.tools: - tool_copy = tool.copy() if hasattr(tool, "copy") else tool - # Idempotent: strip existing prefix before re-applying - base_name = tool_copy.name.removeprefix(self.namespace) - tool_copy.name = f"{self.namespace}{base_name}" - namespaced_tools.append(tool_copy) + seen_cursors: set[str] = set() + while True: + for tool in result.tools: + tool_copy = tool.copy() if hasattr(tool, "copy") else tool + # Idempotent: strip existing prefix before re-applying + base_name = tool_copy.name.removeprefix(self.namespace) + tool_copy.name = f"{self.namespace}{base_name}" + namespaced_tools.append(tool_copy) + cursor = getattr(result, "nextCursor", None) + if cursor is None: + break + if cursor in seen_cursors: + raise RuntimeError("MCP server returned a repeated tool-list cursor") + seen_cursors.add(cursor) + result = await session.list_tools(cursor=cursor) return namespaced_tools def confirm_tool(self, tool_name: str, args: list[Any]) -> bool: diff --git a/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py b/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py index 6b20d732..665bc6c5 100644 --- a/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py +++ b/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py @@ -190,7 +190,9 @@ def _check_model(event: Any, expected: str) -> None: async def _verify_model(client: Any, model: str) -> None: - """Ensure *model* is in the authenticated user's catalogue and enabled.""" + """Validate concrete models; ``auto`` is a router alias, not a catalogue ID.""" + if model == AUTO_MODEL: + return models = await client.list_models() for info in models: if getattr(info, "id", None) != model: @@ -291,6 +293,8 @@ def _on_event(event: Any) -> None: on_permission_request=build_permission_handler([tool.name for tool in tools]), tools=tools, available_tools=available_tools, + # The custom-only policy cannot retrieve SDK-offloaded files. + large_output={"enabled": False}, streaming=True, include_sub_agent_streaming_events=False, reasoning_effort=reasoning_effort, @@ -376,6 +380,11 @@ def _enforce_model(event: Any) -> None: etype = getattr(event, "type", None) data = getattr(event, "data", None) + if etype == SessionEventType.ASSISTANT_MESSAGE and getattr( + data, "originating_message_id", None + ) not in (None, message_id): + continue + if getattr(event, "agent_id", None): # Sub-agent events: account for usage only. if etype == SessionEventType.ASSISTANT_USAGE: diff --git a/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py b/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py index e654a2bf..59281389 100644 --- a/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py +++ b/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py @@ -17,7 +17,7 @@ from copilot import Tool, ToolInvocation, ToolResult -def _tool_handler(server: Any, name: str): +def _tool_handler(server: Any, name: str, *, unwrap_result: bool): async def _invoke(invocation: ToolInvocation) -> ToolResult: from copilot import ToolBinaryResult, ToolResult from mcp.types import BlobResourceContents, EmbeddedResource, ImageContent, TextContent @@ -54,8 +54,12 @@ async def _invoke(invocation: ToolInvocation) -> ToolResult: texts.append(resource.text) else: texts.append(content.model_dump_json(exclude_none=True)) - if not texts and result.structuredContent is not None: - texts.append(json.dumps(result.structuredContent)) + if result.structuredContent is not None: + structured: Any = result.structuredContent + # FastMCP explicitly marks envelopes for non-object return values. + if unwrap_result and list(structured) == ["result"]: + structured = structured["result"] + texts = [structured if isinstance(structured, str) else json.dumps(structured)] return ToolResult( result_type="failure" if result.isError else "success", text_result_for_llm="".join(texts), @@ -90,12 +94,16 @@ async def build_mcp_tools( for tool in await server.list_tools_unfiltered(): if tool.name in blocked or tool.name.removeprefix(server.namespace) in blocked: continue + output_schema = getattr(tool, "outputSchema", None) or {} tools.append( Tool( name=tool.name, description=tool.description or "", parameters=tool.inputSchema, - handler=_tool_handler(server, tool.name), + handler=_tool_handler( + server, tool.name, unwrap_result=output_schema.get("x-fastmcp-wrap-result") is True + ), + defer="never", is_terminal=exclude_from_context, ) ) diff --git a/tests/test_mcp_utils.py b/tests/test_mcp_utils.py index fd9ee489..6f5dffa0 100644 --- a/tests/test_mcp_utils.py +++ b/tests/test_mcp_utils.py @@ -8,7 +8,7 @@ import asyncio import logging from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, call import pytest @@ -58,6 +58,34 @@ def test_list_tools_unfiltered_prefixes_names_from_session(): assert result[1].name == f"{ns}write_file" +def test_list_tools_unfiltered_follows_cursors_through_empty_pages(): + session = SimpleNamespace( + list_tools=AsyncMock( + side_effect=[ + SimpleNamespace(tools=[_FakeTool("first")], nextCursor="page2"), + SimpleNamespace(tools=[], nextCursor="page3"), + SimpleNamespace(tools=[_FakeTool("last")], nextCursor=None), + ] + ) + ) + wrapper = _make_wrapper("RepoContext", session=session) + + result = asyncio.run(wrapper.list_tools_unfiltered()) + + assert [tool.name for tool in result] == [wrapper.namespace + "first", wrapper.namespace + "last"] + assert session.list_tools.await_args_list == [call(), call(cursor="page2"), call(cursor="page3")] + + +def test_list_tools_unfiltered_rejects_repeated_cursors(): + page = SimpleNamespace(tools=[], nextCursor="repeated") + session = SimpleNamespace(list_tools=AsyncMock(side_effect=[page, page])) + wrapper = _make_wrapper("RepoContext", session=session) + + with pytest.raises(RuntimeError, match=r"repeated.*cursor"): + asyncio.run(wrapper.list_tools_unfiltered()) + assert session.list_tools.await_count == 2 + + def test_list_tools_unfiltered_no_double_prefix_when_called_twice(): """Calling list_tools_unfiltered twice should not double-prefix names.""" session = MagicMock() diff --git a/tests/test_sdk_copilot_adapter.py b/tests/test_sdk_copilot_adapter.py index 7ade5127..8b4d12e9 100644 --- a/tests/test_sdk_copilot_adapter.py +++ b/tests/test_sdk_copilot_adapter.py @@ -231,6 +231,19 @@ def test_run_streamed_does_not_duplicate_streamed_message(): assert out == [TextDelta(text="part")] +@pytest.mark.parametrize("foreign_model", [MODEL, "other-model"]) +def test_run_streamed_ignores_full_messages_from_other_requests(foreign_model): + foreign = _ev( + SessionEventType.ASSISTANT_MESSAGE, + content="old response", + message_id="old", + originating_message_id="other-request", + model=foreign_model, + ) + _, out = _stream([_started(), foreign, _delta("current response"), _idle()]) + assert out == [TextDelta(text="current response")] + + def test_streaming_with_public_sdk_event_types(): def event(etype, data): return SessionEvent(type=etype, data=data, id=uuid4(), timestamp=datetime.now(timezone.utc)) @@ -636,6 +649,7 @@ def test_build_uses_explicit_token_and_pins_session(fake_client, caplog): assert kw["streaming"] is True assert kw["reasoning_effort"] == "max" assert kw["available_tools"].to_list() == ["custom:nsecho"] + assert kw["large_output"] == {"enabled": False} assert [tool.name for tool in kw["tools"]] == ["nsecho"] assert "mcp_servers" not in kw assert kw["enable_config_discovery"] is False @@ -657,6 +671,13 @@ def test_build_strips_provider_prefix(fake_client): assert fake_client.instances[0].create_kwargs["model"] == MODEL +def test_build_auto_accepts_a_concrete_only_model_catalogue(fake_client): + fake_client.models = [_model(MODEL)] + handle = _build(model="auto") + assert handle.model == "auto" + assert fake_client.instances[0].create_kwargs["model"] == "auto" + + def test_build_without_instructions_sends_no_system_message(fake_client): _build() assert fake_client.instances[0].create_kwargs["system_message"] is None @@ -689,10 +710,11 @@ def test_build_without_any_token_uses_sdk_default_auth(fake_client, monkeypatch) assert fake_client.instances[0].kwargs == {} -def test_build_authentication_failure_is_actionable_and_stops_client(fake_client): +@pytest.mark.parametrize("model", [MODEL, "auto"]) +def test_build_authentication_failure_is_actionable_and_stops_client(fake_client, model): fake_client.authenticated = False with pytest.raises(BackendBadRequestError, match="authentication failed"): - _build() + _build(model=model) (client,) = fake_client.instances assert client.create_kwargs is None assert client.stopped diff --git a/tests/test_sdk_copilot_mcp.py b/tests/test_sdk_copilot_mcp.py index c6877b74..a2dd374e 100644 --- a/tests/test_sdk_copilot_mcp.py +++ b/tests/test_sdk_copilot_mcp.py @@ -62,6 +62,7 @@ async def _run(): assert tool.parameters == {"type": "object"} assert not tool.skip_permission assert not tool.is_terminal + assert tool.defer == "never" return await tool.handler(ToolInvocation(tool_name=tool.name, arguments={"value": 42})) result = asyncio.run(_run()) @@ -78,6 +79,30 @@ def test_blocked_tools_are_not_registered(namespaced): transport.call_tool.assert_not_awaited() +@pytest.mark.parametrize("namespaced", [False, True]) +def test_paginated_tools_are_filtered_and_invoked_through_the_wrapper(namespaced): + server, transport = _server(confirms=["later"]) + server.confirm_tool = Mock(return_value=True) + transport.session.list_tools.side_effect = [ + ListToolsResult(tools=[Tool(name="echo", inputSchema={})], nextCursor="next"), + ListToolsResult( + tools=[Tool(name="blocked", inputSchema={}), Tool(name="later", inputSchema={})], + ), + ] + blocked = server.namespace + "blocked" if namespaced else "blocked" + + async def _run(): + tools = {tool.name.removeprefix(server.namespace): tool for tool in await _tools(server, blocked=[blocked])} + assert set(tools) == {"echo", "later"} + tool = tools["later"] + return await tool.handler(ToolInvocation(tool_name=tool.name, arguments={"value": 42})) + + result = asyncio.run(_run()) + assert result.result_type == "success" + server.confirm_tool.assert_called_once_with("later", [{"value": 42}]) + transport.call_tool.assert_awaited_once_with("later", {"value": 42}) + + @pytest.mark.parametrize("approved", [False, True]) def test_toolbox_confirmation_cannot_be_bypassed(approved): server, transport = _server(confirms=["echo"]) @@ -173,6 +198,51 @@ async def _run(): assert result.binary_results_for_llm[0].mime_type == "image/png" +@pytest.mark.parametrize("structured", [{"items": [1, 2, 3]}, {}, {"result": [1, 2, 3]}]) +def test_structured_payload_takes_precedence_over_text(structured): + server, transport = _server() + transport.call_tool.return_value = CallToolResult( + content=[TextContent(type="text", text="Found 3 items")], + structuredContent=structured, + ) + + async def _run(): + (tool,) = await _tools(server) + return await tool.handler(ToolInvocation(tool_name=tool.name)) + + result = asyncio.run(_run()) + assert json.loads(result.text_result_for_llm) == structured + + +@pytest.mark.parametrize( + ("structured", "expected"), + [ + ({"result": "[1, 2, 3]"}, [1, 2, 3]), + ({"result": [1, 2, 3]}, [1, 2, 3]), + ({"result": [1, 2, 3], "extra": True}, {"result": [1, 2, 3], "extra": True}), + ], +) +def test_structured_fastmcp_envelopes_preserve_original_payloads(structured, expected): + server, transport = _server() + transport.session.list_tools.return_value.tools[0].outputSchema = { + "type": "object", + "properties": {"result": {}}, + "required": ["result"], + "x-fastmcp-wrap-result": True, + } + transport.call_tool.return_value = CallToolResult( + content=[TextContent(type="text", text="Found 3 items")], + structuredContent=structured, + ) + + async def _run(): + (tool,) = await _tools(server) + return await tool.handler(ToolInvocation(tool_name=tool.name)) + + result = asyncio.run(_run()) + assert json.loads(result.text_result_for_llm) == expected + + def test_raw_mcp_specs_cannot_bypass_runner_policy(): with pytest.raises(BackendCapabilityError, match="connected by the taskflow runner"): asyncio.run( From e00214c0256996f410d7ccad84d1ae270067dae1 Mon Sep 17 00:00:00 2001 From: Sriharsha Etigowni Date: Fri, 30 Oct 2026 22:07:58 +0000 Subject: [PATCH 3/5] fix(copilot_sdk): handle MCP refusals and reconnecting servers Report declined confirmations as MCP errors rather than successful tool calls. Keep unfiltered tool pagination inside one serialized per-call connection for reconnecting transports, with cleanup on errors and cancellation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/seclab_taskflow_agent/mcp_transport.py | 27 ++++----- src/seclab_taskflow_agent/mcp_utils.py | 14 ++++- tests/test_mcp_utils.py | 64 ++++++++++++++++++++++ tests/test_sdk_copilot_mcp.py | 14 +++-- 4 files changed, 99 insertions(+), 20 deletions(-) diff --git a/src/seclab_taskflow_agent/mcp_transport.py b/src/seclab_taskflow_agent/mcp_transport.py index 8632fd8d..195cbe08 100644 --- a/src/seclab_taskflow_agent/mcp_transport.py +++ b/src/seclab_taskflow_agent/mcp_transport.py @@ -28,7 +28,8 @@ import socket import subprocess import time -from collections.abc import Callable +from collections.abc import AsyncIterator, Callable +from contextlib import asynccontextmanager from threading import Event, Thread from typing import IO, Any from urllib.parse import urlparse @@ -291,23 +292,23 @@ async def cleanup(self) -> None: """No-op — cleanup happens per-call.""" logging.debug("Ignoring mcp cleanup request on purpose") - async def list_tools(self, *args: Any, **kwargs: Any) -> Any: - """Connect, list tools, then disconnect.""" + @asynccontextmanager + async def tool_session(self) -> AsyncIterator[Any]: + """Hold one serialized connection for an entire tool operation.""" async with self.reconnecting_lock: - await super().connect() try: - result = await super().list_tools(*args, **kwargs) + await super().connect() + yield self.session finally: await super().cleanup() - return result + + async def list_tools(self, *args: Any, **kwargs: Any) -> Any: + """Connect, list tools, then disconnect.""" + async with self.tool_session(): + return await super().list_tools(*args, **kwargs) async def call_tool(self, *args: Any, **kwargs: Any) -> Any: """Connect, call tool, then disconnect.""" logging.debug("Using reconnecting call_tool for stdio mcp") - async with self.reconnecting_lock: - await super().connect() - try: - result = await super().call_tool(*args, **kwargs) - finally: - await super().cleanup() - return result + async with self.tool_session(): + return await super().call_tool(*args, **kwargs) diff --git a/src/seclab_taskflow_agent/mcp_utils.py b/src/seclab_taskflow_agent/mcp_utils.py index ee6d7b55..7def9f81 100644 --- a/src/seclab_taskflow_agent/mcp_utils.py +++ b/src/seclab_taskflow_agent/mcp_utils.py @@ -117,17 +117,24 @@ async def list_tools_unfiltered(self) -> list[Any]: run loop -- e.g. when handing tools to a different SDK at build time). + Reconnecting transports hold one connection open for all pages. + Prefixing is idempotent: if a tool's name already starts with this wrapper's namespace (e.g. because the underlying session returned a previously-namespaced object), the existing prefix is stripped before re-applying so calling this method multiple times never yields ``name``. - Raises ``RuntimeError`` if the underlying server has no active + Raises ``RuntimeError`` if a persistent server has no active MCP session yet (caller should ensure the server is connected before calling this), or if a pagination cursor repeats. """ - session = getattr(self._obj, "session", None) + if isinstance(self._obj, ReconnectingMCPServerStdio): + async with self._obj.tool_session() as session: + return await self._list_session_tools(session) + return await self._list_session_tools(getattr(self._obj, "session", None)) + + async def _list_session_tools(self, session: Any) -> list[Any]: if session is None: raise RuntimeError( f"MCPNamespaceWrap({self._obj!r}): underlying server has no " @@ -180,7 +187,8 @@ async def call_tool(self, *args: Any, **kwargs: Any) -> Any: if self.confirms and tool_name in self.confirms: if not self.confirm_tool(tool_name, _args[1:]): result = CallToolResult( - content=[TextContent(type="text", text="Tool call not allowed.", annotations=None, meta=None)] + content=[TextContent(type="text", text="Tool call not allowed.", annotations=None, meta=None)], + isError=True, ) return result _args[0] = tool_name diff --git a/tests/test_mcp_utils.py b/tests/test_mcp_utils.py index 6f5dffa0..2793ca2c 100644 --- a/tests/test_mcp_utils.py +++ b/tests/test_mcp_utils.py @@ -11,7 +11,9 @@ from unittest.mock import AsyncMock, MagicMock, call import pytest +from agents.mcp import MCPServerStdio +from seclab_taskflow_agent.mcp_transport import ReconnectingMCPServerStdio from seclab_taskflow_agent.mcp_utils import ( MCPNamespaceWrap, _env_names, @@ -131,6 +133,56 @@ def test_list_tools_unfiltered_raises_when_session_missing(): asyncio.run(wrapper.list_tools_unfiltered()) +@pytest.mark.parametrize("outcome", ["success", "listing-error", "cancelled", "connect-error"]) +def test_unfiltered_reconnecting_listing_owns_its_session(monkeypatch, outcome): + session = SimpleNamespace( + list_tools=AsyncMock( + side_effect=[ + SimpleNamespace(tools=[_FakeTool("first")], nextCursor="next"), + SimpleNamespace(tools=[_FakeTool("last")], nextCursor=None), + ] + ) + ) + if outcome in ("listing-error", "cancelled"): + error = asyncio.CancelledError() if outcome == "cancelled" else RuntimeError("listing failed") + session.list_tools.side_effect = [ + SimpleNamespace(tools=[_FakeTool("first")], nextCursor="next"), + error, + ] + events = [] + + async def connect(server): + events.append("connect") + if outcome == "connect-error": + raise RuntimeError("connection failed") + server.session = session + + async def cleanup(server): + events.append("cleanup") + server.session = None + + monkeypatch.setattr(MCPServerStdio, "connect", connect) + monkeypatch.setattr(MCPServerStdio, "cleanup", cleanup) + transport = ReconnectingMCPServerStdio(params={"command": "unused"}, name="Repo") + transport.list_tools = AsyncMock(side_effect=AssertionError("filtered listing must not be used")) + wrapper = MCPNamespaceWrap([], transport) + + if outcome == "success": + tools = asyncio.run(wrapper.list_tools_unfiltered()) + assert [tool.name for tool in tools] == [wrapper.namespace + "first", wrapper.namespace + "last"] + else: + exception = asyncio.CancelledError if outcome == "cancelled" else RuntimeError + with pytest.raises(exception): + asyncio.run(wrapper.list_tools_unfiltered()) + + assert events == ["connect", "cleanup"] + assert transport.session is None + assert not transport.reconnecting_lock.locked() + transport.list_tools.assert_not_awaited() + expected_calls = [] if outcome == "connect-error" else [call(), call(cursor="next")] + assert session.list_tools.await_args_list == expected_calls + + def test_list_tools_unfiltered_does_not_share_state_with_caller(): """Mutating returned tool names must not affect the underlying tools.""" original = _FakeTool("read_file") @@ -211,6 +263,18 @@ def test_list_tools_existing_behaviour_unchanged(): assert result[0].name == f"{ns}read_file" +def test_call_tool_reports_confirmation_refusal_as_an_error(monkeypatch): + transport = SimpleNamespace(name="Repo", call_tool=AsyncMock()) + wrapper = MCPNamespaceWrap(["write_file"], transport) + monkeypatch.setattr("builtins.input", lambda _prompt: "no") + + result = asyncio.run(wrapper.call_tool(wrapper.namespace + "write_file", {"path": "example.txt"})) + + assert result.isError is True + assert result.content[0].text == "Tool call not allowed." + transport.call_tool.assert_not_awaited() + + # -- secret redaction in debug logs -- diff --git a/tests/test_sdk_copilot_mcp.py b/tests/test_sdk_copilot_mcp.py index a2dd374e..3f8c85e6 100644 --- a/tests/test_sdk_copilot_mcp.py +++ b/tests/test_sdk_copilot_mcp.py @@ -18,6 +18,7 @@ from copilot import ToolInvocation from mcp.types import CallToolResult, ImageContent, ListToolsResult, TextContent, Tool +from seclab_taskflow_agent.mcp_transport import ReconnectingMCPServerStdio from seclab_taskflow_agent.mcp_utils import MCPNamespaceWrap from seclab_taskflow_agent.sdk.base import MCPServerSpec from seclab_taskflow_agent.sdk.copilot_sdk.mcp import build_mcp_tools @@ -104,15 +105,17 @@ async def _run(): @pytest.mark.parametrize("approved", [False, True]) -def test_toolbox_confirmation_cannot_be_bypassed(approved): +@pytest.mark.parametrize("exclude", [False, True]) +def test_toolbox_confirmation_cannot_be_bypassed(approved, exclude): server, transport = _server(confirms=["echo"]) server.confirm_tool = Mock(return_value=approved) async def _run(): - (tool,) = await _tools(server) + (tool,) = await _tools(server, exclude=exclude) return await tool.handler(ToolInvocation(tool_name=tool.name, arguments={"value": 42})) result = asyncio.run(_run()) + assert result.result_type == ("success" if approved else "failure") server.confirm_tool.assert_called_once_with("echo", [{"value": 42}]) if approved: transport.call_tool.assert_awaited_once() @@ -254,11 +257,14 @@ def test_raw_mcp_specs_cannot_bypass_runner_policy(): ) -def test_memcache_round_trip_without_copilot_credentials(tmp_path): +@pytest.mark.parametrize("reconnecting", [False, True]) +def test_memcache_round_trip_without_copilot_credentials(tmp_path, reconnecting): from agents.mcp import MCPServerStdio + server_class = ReconnectingMCPServerStdio if reconnecting else MCPServerStdio + async def _run(): - async with MCPServerStdio( + async with server_class( params={ "command": sys.executable, "args": ["-m", "seclab_taskflow_agent.mcp_servers.memcache.memcache"], From 5c1aaf50c777a9fa9f68b9879840117223cb4ee3 Mon Sep 17 00:00:00 2001 From: Sriharsha Etigowni Date: Fri, 30 Oct 2026 22:27:09 +0000 Subject: [PATCH 4/5] fix(copilot_sdk): ignore stale pre-send responses Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../sdk/copilot_sdk/backend.py | 7 +++- tests/test_sdk_copilot_adapter.py | 37 ++++++++++++++++++- 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py b/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py index 665bc6c5..8f671092 100644 --- a/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py +++ b/src/seclab_taskflow_agent/sdk/copilot_sdk/backend.py @@ -358,13 +358,16 @@ def _enforce_model(event: Any) -> None: tool_names: dict[str, str] = {} streamed_message_ids: set[str] = set() try: - # Discard an earlier turn's events, retaining model-pin checks. + # Only session-level events describe the model for the next request. while True: try: stale = agent.event_queue.get_nowait() except asyncio.QueueEmpty: break - if not getattr(stale, "agent_id", None): + if getattr(stale, "type", None) in ( + SessionEventType.SESSION_START, + SessionEventType.SESSION_MODEL_CHANGE, + ) and not getattr(stale, "agent_id", None): _enforce_model(stale) try: diff --git a/tests/test_sdk_copilot_adapter.py b/tests/test_sdk_copilot_adapter.py index 8b4d12e9..3c503920 100644 --- a/tests/test_sdk_copilot_adapter.py +++ b/tests/test_sdk_copilot_adapter.py @@ -208,6 +208,31 @@ def test_run_streamed_drains_stale_queue_before_send(): assert out == [TextDelta(text="new")] +def test_run_streamed_discards_stale_foreign_response_before_model_checks(): + stale = [ + _ev( + SessionEventType.ASSISTANT_MESSAGE, + content="old response", + message_id="old", + originating_message_id="other-request", + model="gpt-4.1", + ), + _ev(SessionEventType.ASSISTANT_USAGE, model="gpt-4.1", input_tokens=99, output_tokens=99), + ] + handle, out = _stream( + [ + _started(), + _delta("new"), + _ev(SessionEventType.ASSISTANT_USAGE, model=MODEL, input_tokens=1, output_tokens=2), + _idle(), + ], + stale=stale, + ) + assert handle.session.sent == ["go"] + assert handle.session.aborted == 0 + assert out == [TextDelta(text="new"), TokenUsage(model=MODEL, input_tokens=1, output_tokens=2)] + + def test_run_streamed_falls_back_to_full_assistant_message(): _, out = _stream( [ @@ -495,8 +520,16 @@ def test_run_streamed_auto_model_accepts_routed_model(caplog): assert "'auto' routed to 'gpt-6-luna'" in caplog.text -def test_run_streamed_rejects_model_switch_in_stale_events(): - handle = _handle([], stale=[_ev(SessionEventType.SESSION_START, selected_model="gpt-4.1")]) +@pytest.mark.parametrize( + "event", + [ + _ev(SessionEventType.SESSION_START, selected_model="gpt-4.1"), + _ev(SessionEventType.SESSION_MODEL_CHANGE, new_model="gpt-4.1", previous_model=MODEL), + ], + ids=lambda e: e.type.value, +) +def test_run_streamed_rejects_model_switch_in_stale_events(event): + handle = _handle([], stale=[event]) with pytest.raises(BackendUnexpectedError, match=r"gpt-4\.1"): asyncio.run(_collect(handle)) assert handle.session.sent == [] From 3ad4fc2a3e0052ad082a3a117ae88727d0e23528 Mon Sep 17 00:00:00 2001 From: Sriharsha Etigowni Date: Fri, 30 Oct 2026 23:48:52 +0000 Subject: [PATCH 5/5] fix(copilot_sdk): distinguish confirmation and transport failures Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/seclab_taskflow_agent/mcp_utils.py | 11 +++++++- .../sdk/copilot_sdk/mcp.py | 6 ++-- tests/test_sdk_copilot_mcp.py | 28 +++++++++++++++++-- 3 files changed, 40 insertions(+), 5 deletions(-) diff --git a/src/seclab_taskflow_agent/mcp_utils.py b/src/seclab_taskflow_agent/mcp_utils.py index 7def9f81..5778e46b 100644 --- a/src/seclab_taskflow_agent/mcp_utils.py +++ b/src/seclab_taskflow_agent/mcp_utils.py @@ -12,6 +12,7 @@ __all__ = [ "COMPRESSED_NAME_LENGTH", "DEFAULT_MCP_CLIENT_SESSION_TIMEOUT", + "MCPConfirmationUnavailableError", "MCPNamespaceWrap", "compress_name", "mcp_client_params", @@ -70,6 +71,10 @@ def compress_name(name: str) -> str: return COMPRESSED_NAME_PREFIX + m.hexdigest()[:COMPRESSED_NAME_LENGTH] +class MCPConfirmationUnavailableError(RuntimeError): + """A required tool confirmation could not be obtained.""" + + class MCPNamespaceWrap: """MCP client wrapper that prefixes tool names with a namespace hash. @@ -185,7 +190,11 @@ async def call_tool(self, *args: Any, **kwargs: Any) -> Any: tool_name = tool_name.removeprefix(self.namespace) # to run headless, just make confirms an empty list if self.confirms and tool_name in self.confirms: - if not self.confirm_tool(tool_name, _args[1:]): + try: + approved = self.confirm_tool(tool_name, _args[1:]) + except (EOFError, OSError) as exc: + raise MCPConfirmationUnavailableError("Required toolbox confirmation is unavailable") from exc + if not approved: result = CallToolResult( content=[TextContent(type="text", text="Tool call not allowed.", annotations=None, meta=None)], isError=True, diff --git a/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py b/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py index 59281389..47e59e3a 100644 --- a/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py +++ b/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py @@ -22,14 +22,16 @@ async def _invoke(invocation: ToolInvocation) -> ToolResult: from copilot import ToolBinaryResult, ToolResult from mcp.types import BlobResourceContents, EmbeddedResource, ImageContent, TextContent + from ...mcp_utils import MCPConfirmationUnavailableError + if invocation.tool_name != name: return ToolResult(result_type="denied", text_result_for_llm="Unknown taskflow tool") try: result = await server.call_tool(name, invocation.arguments or {}) - except (EOFError, OSError): + except MCPConfirmationUnavailableError as exc: return ToolResult( result_type="denied", - text_result_for_llm="Required toolbox confirmation is unavailable", + text_result_for_llm=str(exc), ) texts: list[str] = [] diff --git a/tests/test_sdk_copilot_mcp.py b/tests/test_sdk_copilot_mcp.py index 3f8c85e6..5ce48113 100644 --- a/tests/test_sdk_copilot_mcp.py +++ b/tests/test_sdk_copilot_mcp.py @@ -138,9 +138,10 @@ async def _run(): transport.call_tool.assert_awaited_once() -def test_missing_confirmation_fails_closed(): +@pytest.mark.parametrize("error_type", [EOFError, OSError]) +def test_missing_confirmation_fails_closed(monkeypatch, error_type): server, transport = _server(confirms=["echo"]) - server.confirm_tool = Mock(side_effect=EOFError) + monkeypatch.setattr("builtins.input", Mock(side_effect=error_type)) async def _run(): (tool,) = await _tools(server) @@ -151,6 +152,29 @@ async def _run(): transport.call_tool.assert_not_awaited() +@pytest.mark.parametrize("error_type", [EOFError, OSError]) +@pytest.mark.parametrize("requires_confirmation", [False, True]) +def test_transport_errors_are_not_reported_as_confirmation_denials(error_type, requires_confirmation): + server, transport = _server(confirms=["echo"] if requires_confirmation else []) + server.confirm_tool = Mock(return_value=True) + error = error_type("MCP connection lost") + transport.call_tool.side_effect = error + + async def _run(): + (tool,) = await _tools(server) + return await tool.handler(ToolInvocation(tool_name=tool.name, arguments={"value": 42})) + + with pytest.raises(error_type, match="MCP connection lost") as raised: + asyncio.run(_run()) + + assert raised.value is error + transport.call_tool.assert_awaited_once_with("echo", {"value": 42}) + if requires_confirmation: + server.confirm_tool.assert_called_once_with("echo", [{"value": 42}]) + else: + server.confirm_tool.assert_not_called() + + def test_invocation_cannot_select_another_tool(): server, transport = _server()