diff --git a/README.md b/README.md index 8a0aacc..045a5bf 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 8487360..7271e7b 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/mcp_transport.py b/src/seclab_taskflow_agent/mcp_transport.py index 8632fd8..195cbe0 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 84cc313..5778e46 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. @@ -109,7 +114,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 @@ -117,17 +122,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). + 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 " @@ -135,12 +147,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: @@ -169,9 +190,14 @@ 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)] + content=[TextContent(type="text", text="Tool call not allowed.", annotations=None, meta=None)], + isError=True, ) return result _args[0] = tool_name diff --git a/src/seclab_taskflow_agent/runner.py b/src/seclab_taskflow_agent/runner.py index a2892d1..6da6bdd 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 639f3b1..6ee5977 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 31ef437..8f67109 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,96 @@ 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: + """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: + 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 +242,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 +255,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 +272,59 @@ 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, + # The custom-only policy cannot retrieve SDK-offloaded files. + large_output={"enabled": False}, 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 +335,152 @@ 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: + # 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 getattr(stale, "type", None) in ( + SessionEventType.SESSION_START, + SessionEventType.SESSION_MODEL_CHANGE, + ) and 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 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: + 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 +488,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 8cb23cd..47e59e3 100644 --- a/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py +++ b/src/seclab_taskflow_agent/sdk/copilot_sdk/mcp.py @@ -1,85 +1,112 @@ # 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, *, unwrap_result: bool): + 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 MCPConfirmationUnavailableError as exc: + return ToolResult( + result_type="denied", + text_result_for_llm=str(exc), + ) + + 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 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), + 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 - - -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]] = {} + from copilot import Tool + + 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 + 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, unwrap_result=output_schema.get("x-fastmcp-wrap-result") is True + ), + defer="never", + 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 e918a72..ac38c89 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_mcp_utils.py b/tests/test_mcp_utils.py index fd9ee48..2793ca2 100644 --- a/tests/test_mcp_utils.py +++ b/tests/test_mcp_utils.py @@ -8,10 +8,12 @@ import asyncio import logging from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +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, @@ -58,6 +60,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() @@ -103,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") @@ -183,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_adapter.py b/tests/test_sdk_copilot_adapter.py index 29ba114..3c50392 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,727 @@ 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, + ) -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) - ] +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 ----------------------------------------------------------- + - session, out = asyncio.run(_run()) - assert session.sent == ["go"] +def test_run_streamed_translates_deltas_until_idle(): + 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_emits_token_usage(): - """An assistant.usage event is translated into a neutral TokenUsage.""" +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")] - async def _run(): - queue: asyncio.Queue[Any] = asyncio.Queue() - for ev in ( - _Event( + +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_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( + [ + _started(), + _ev(SessionEventType.ASSISTANT_MESSAGE, content="whole", message_id="m1", model=MODEL), + _idle(), + ] + ) + assert out == [TextDelta(text="whole")] + + +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")] + + +@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)) + + _, 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 + + +def test_run_streamed_skips_failed_tool(): + _, out = _stream([_started(), *_tool_events(success=False), _idle()]) + assert out == [] + - with pytest.raises(BackendUnexpectedError, match="boom"): - asyncio.run(_run()) +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_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] = [] +@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 == [] + + +@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 - class _AbortableSession(_FakeSession): - async def abort(self) -> None: - aborted.append(True) +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 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 + 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_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 + + +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 == {} + + +@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(model=model) + (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 +819,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 +850,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 54f6fec..5ce4811 100644 --- a/tests/test_sdk_copilot_mcp.py +++ b/tests/test_sdk_copilot_mcp.py @@ -1,118 +1,321 @@ # 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_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_config +from seclab_taskflow_agent.sdk.copilot_sdk.mcp import build_mcp_tools +from seclab_taskflow_agent.sdk.errors import BackendCapabilityError -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, - }, - ) - ] - assert build_mcp_config(specs) == { - "files": { - "type": "stdio", - "command": "/usr/bin/mcp", - "args": ["--root", "/tmp"], - "env": {"FOO": "1"}, - "cwd": "/work", - "tools": ["*"], - "timeout": 2500, - } - } +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_sse_spec_translated_to_remote_config(): - specs = [ - MCPServerSpec( - name="docs", - kind="sse", - params={ - "url": "https://example.com/mcp", - "headers": {"Authorization": "Bearer x"}, - "timeout": 10.0, - }, - ) - ] - assert build_mcp_config(specs) == { - "docs": { - "type": "sse", - "url": "https://example.com/mcp", - "headers": {"Authorization": "Bearer x"}, - "tools": ["*"], - "timeout": 10000, - } - } +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, + ) -def test_streamable_spec_uses_http_type(): - specs = [ - MCPServerSpec( - name="gh", - kind="streamable", - params={"url": "https://api.example.com/mcp"}, - ) +@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 + assert tool.defer == "never" + 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("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={})], + ), ] - out = build_mcp_config(specs) - assert out["gh"]["type"] == "http" - assert out["gh"]["url"] == "https://api.example.com/mcp" + 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]) +@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, 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() + 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() + + +@pytest.mark.parametrize("error_type", [EOFError, OSError]) +def test_missing_confirmation_fails_closed(monkeypatch, error_type): + server, transport = _server(confirms=["echo"]) + monkeypatch.setattr("builtins.input", Mock(side_effect=error_type)) + + 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_streamable_spec_without_url_is_skipped(): - assert build_mcp_config([MCPServerSpec(name="x", kind="streamable", params={})]) == {} +@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 -def test_stdio_minimal_defaults(): - import os + async def _run(): + (tool,) = await _tools(server) + return await tool.handler(ToolInvocation(tool_name=tool.name, arguments={"value": 42})) - specs = [MCPServerSpec(name="m", kind="stdio", params={"command": "x"})] - assert build_mcp_config(specs) == { - "m": { - "type": "stdio", - "command": "x", - "args": [], - "tools": ["*"], - "cwd": os.getcwd(), - } + 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() + + 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" + + +@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_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"} + +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)["m"]["cwd"] == "/tmp" + + +@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 server_class( + params={ + "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), + }, + }, + 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 c9179b7..6c352b1 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)