diff --git a/daiv/automation/agent/graph.py b/daiv/automation/agent/graph.py index f3bc3a8f9..b6724f935 100644 --- a/daiv/automation/agent/graph.py +++ b/daiv/automation/agent/graph.py @@ -25,7 +25,6 @@ ModelName, ) from automation.agent.deferred.conf import settings as deferred_settings -from automation.agent.mcp.toolkits import MCPToolkit from automation.agent.middlewares.deferred_tools import DeferredToolsMiddleware from automation.agent.middlewares.ensure_response import ensure_non_empty_response from automation.agent.middlewares.file_system import ( @@ -52,6 +51,7 @@ if TYPE_CHECKING: from collections.abc import Sequence + from langchain_core.tools.base import BaseTool from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.store.base import BaseStore @@ -147,6 +147,7 @@ async def create_daiv_agent( thinking_level: ThinkingLevel | None | type[_Unset] = _Unset, *, ctx: RuntimeCtx, + mcp_tools: Sequence[BaseTool], auto_commit_changes: bool = True, checkpointer: BaseCheckpointSaver | None = None, store: BaseStore | None = None, @@ -165,6 +166,10 @@ async def create_daiv_agent( model_names: The model names to use for the agent. thinking_level: The thinking level to use for the agent. ctx: The runtime context. + mcp_tools: MCP tools bound to persistent sessions for the agent run's + lifetime. Callers must obtain these from ``MCPToolkit.aopen()`` so + stateful servers (e.g. Playwright) share a single session across + tool calls; passing freshly-discovered tools here breaks that. auto_commit_changes: Whether to commit the changes to the repository when the agent finishes. checkpointer: The checkpointer to use for the agent. store: The store to use for the agent. @@ -229,8 +234,6 @@ async def create_daiv_agent( ) subagents.extend(custom_subagents) - mcp_tools = await MCPToolkit.get_tools() - user_middleware: list[AgentMiddleware[Any, Any, Any]] = [ TodoListMiddleware(system_prompt=dynamic_write_todos_system_prompt(bash_tool_enabled=_sandbox_enabled)), SkillsMiddleware( diff --git a/daiv/automation/agent/mcp/conf.py b/daiv/automation/agent/mcp/conf.py index cd27dcb87..9adbf1169 100644 --- a/daiv/automation/agent/mcp/conf.py +++ b/daiv/automation/agent/mcp/conf.py @@ -14,7 +14,7 @@ class MCPSettings(BaseSettings): ), ) - # Built-in MCP server URLs (supergateway containers). Set to None to disable. + # Built-in MCP server URLs (each runs in a dedicated container). Set to None to disable. # Note: Docker Swarm normalizes service names with dashes to underscores. SENTRY_URL: str | None = Field( default="http://mcp_sentry:8000/mcp", description="The streamable HTTP URL of the Sentry supergateway container" @@ -23,6 +23,10 @@ class MCPSettings(BaseSettings): default="http://mcp_context7:8000/mcp", description="The streamable HTTP URL of the Context7 supergateway container", ) + PLAYWRIGHT_URL: str | None = Field( + default="http://mcp_playwright:8931/mcp", + description="The streamable HTTP URL of the Playwright browser MCP container", + ) settings = MCPSettings() diff --git a/daiv/automation/agent/mcp/servers.py b/daiv/automation/agent/mcp/servers.py index d98552f03..9d1189dde 100644 --- a/daiv/automation/agent/mcp/servers.py +++ b/daiv/automation/agent/mcp/servers.py @@ -51,3 +51,22 @@ def is_enabled(self) -> bool: def get_connection(self) -> StreamableHttpConnection: assert settings.CONTEXT7_URL is not None # guaranteed by is_enabled() return StreamableHttpConnection(transport="streamable_http", url=settings.CONTEXT7_URL) + + +@mcp_server +class PlaywrightMCPServer(MCPServer): + """ + Full Playwright MCP tool surface exposed (no tool_filter). The container + runs headless Chromium with --isolated, so each MCP session gets a fresh + browser context. See docs/customization/mcp-tools.md for the trust-boundary + discussion. + """ + + name = "playwright" + + def is_enabled(self) -> bool: + return settings.PLAYWRIGHT_URL is not None + + def get_connection(self) -> StreamableHttpConnection: + assert settings.PLAYWRIGHT_URL is not None # guaranteed by is_enabled() + return StreamableHttpConnection(transport="streamable_http", url=settings.PLAYWRIGHT_URL) diff --git a/daiv/automation/agent/mcp/toolkits.py b/daiv/automation/agent/mcp/toolkits.py index 85b63983c..5327ff9ac 100644 --- a/daiv/automation/agent/mcp/toolkits.py +++ b/daiv/automation/agent/mcp/toolkits.py @@ -1,13 +1,18 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from contextlib import AsyncExitStack, asynccontextmanager +from typing import TYPE_CHECKING, Any -from langchain_mcp_adapters.client import MultiServerMCPClient +from langchain_mcp_adapters.tools import load_mcp_tools +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client from automation.agent.toolkits import BaseToolkit if TYPE_CHECKING: + from collections.abc import AsyncIterator, Callable + from langchain_core.tools.base import BaseTool from automation.agent.mcp.schemas import ToolFilter @@ -19,38 +24,141 @@ def _get_connection_url(conn) -> str: return getattr(conn, "url", "unknown") +@asynccontextmanager +async def _open_streamable_mcp_session( + *, url: str, headers: dict[str, Any] | None, terminate_on_close: bool, initialize: bool +) -> AsyncIterator[tuple[ClientSession, Callable[[], str | None]]]: + """Open a single Streamable-HTTP MCP session and expose the server-issued session id. + + ``langchain_mcp_adapters``' helper discards the ``get_session_id`` callback; + we re-implement it so callers can persist the id and resume on a later turn. + Set ``initialize=False`` when reconnecting to a session the server already + knows — the Playwright server rejects re-init with -32600 "Server already + initialized" otherwise. + """ + async with ( + streamablehttp_client(url, headers=headers, terminate_on_close=terminate_on_close) as (read, write, get_id), + ClientSession(read, write) as session, + ): + if initialize: + await session.initialize() + yield session, get_id + + class MCPToolkit(BaseToolkit): @classmethod - async def get_tools(cls) -> list[BaseTool]: + @asynccontextmanager + async def aopen( + cls, *, session_ids: dict[str, str] | None = None + ) -> AsyncIterator[tuple[list[BaseTool], dict[str, str]]]: + """Open one persistent MCP session per registered server for the block's scope. + + Tools yielded inside the ``async with`` reuse their server's session on + every invocation. Stateful MCP servers (Playwright above all) need this: + each MCP session gets its own browser context, so without session reuse + ``navigate`` and ``snapshot`` land in different browsers and the second + sees ``about:blank``. + + ``session_ids`` controls cross-turn behaviour: + + - ``None`` (default): sessions terminate on context exit; each call gets + fresh state. Right for one-shot runs (jobs, webhooks). + - ``dict``: persistent mode. For each server with an id in the dict, we + reconnect to that server-side session (skipping init); for servers + without one, we open fresh and capture the id back into the dict. The + MCP session stays alive past close so the next call with the same dict + can resume it. Right for chat threads. + + On a stale id (404 "Session not found"), we transparently fall back to + a fresh session and overwrite the id in ``session_ids`` so the caller's + next persistence write doesn't preserve the dead value. A server that + fails to open is logged and skipped; the rest still connect. + """ from automation.agent.mcp.registry import mcp_registry connections, tool_filters = mcp_registry.get_connections_and_filters() - if not connections: - return [] + yield [], {} + return server_urls = {name: _get_connection_url(conn) for name, conn in connections.items()} - logger.debug("Connecting to MCP servers: %s", server_urls) - - client = MultiServerMCPClient(connections, tool_name_prefix=True) + logger.debug("Opening MCP sessions: %s", server_urls) + + persist = session_ids is not None + ids_out: dict[str, str] = dict(session_ids) if session_ids else {} + + async with AsyncExitStack() as stack: + tools: list[BaseTool] = [] + for name, conn in connections.items(): + if conn.get("transport") != "streamable_http": + logger.warning("MCP %s: non-streamable transports are not supported by aopen", name) + continue + url = conn["url"] + base_headers = dict(conn.get("headers") or {}) + existing_id = ids_out.get(name) + server_tools, captured_id = await _open_server(stack, url, base_headers, persist, existing_id, name) + if server_tools is None: + if existing_id and name in ids_out: + # Drop the stale id so the caller doesn't persist a dead one. + del ids_out[name] + continue + tools.extend(server_tools) + if persist: + new_id = captured_id or existing_id + if new_id: + ids_out[name] = new_id + elif name in ids_out: + # Server doesn't issue session ids (stateless); drop the empty slot. + del ids_out[name] + + if tool_filters: + tools = _apply_tool_filters(tools, tool_filters) + + for tool in tools: + tool.handle_tool_error = True + tool.handle_validation_error = True + tool.tags = ["mcp_server"] + tool.metadata = {"mcp_server": tool.name} + + yield tools, ids_out + + +async def _open_server( + stack: AsyncExitStack, url: str, base_headers: dict[str, Any], persist: bool, existing_id: str | None, name: str +) -> tuple[list[BaseTool] | None, str | None]: + """Open one MCP server, attempting resume when an existing id is supplied. + + Returns ``(tools, captured_id)`` on success, ``(None, None)`` on failure. + A stale ``existing_id`` is recovered from automatically by opening fresh. + """ + terminate_on_close = not persist + if existing_id: try: - tools = await client.get_tools() + session, get_id = await stack.enter_async_context( + _open_streamable_mcp_session( + url=url, + headers={**base_headers, "Mcp-Session-Id": existing_id}, + terminate_on_close=terminate_on_close, + initialize=False, + ) + ) + tools = await load_mcp_tools(session, server_name=name, tool_name_prefix=True) + return tools, get_id() or existing_id except Exception: - logger.warning("Error getting tools from MCP servers: %s", server_urls, exc_info=True) - tools = [] - - if tool_filters: - tools = _apply_tool_filters(tools, tool_filters) - - # Handle tool errors and validation errors gracefully to allow the agent to continue - for tool in tools: - tool.handle_tool_error = True - tool.handle_validation_error = True - tool.tags = ["mcp_server"] - tool.metadata = {"mcp_server": tool.name} - - return tools + logger.warning("MCP %s: resume of session %s failed; opening fresh", name, existing_id, exc_info=True) + + try: + session, get_id = await stack.enter_async_context( + _open_streamable_mcp_session( + url=url, headers=base_headers or None, terminate_on_close=terminate_on_close, initialize=True + ) + ) + tools = await load_mcp_tools(session, server_name=name, tool_name_prefix=True) + except Exception: + logger.warning("MCP %s: open failed", name, exc_info=True) + return None, None + return tools, get_id() def _apply_tool_filters(tools: list[BaseTool], filters: dict[str, ToolFilter]) -> list[BaseTool]: diff --git a/daiv/automation/agent/middlewares/mcp_session.py b/daiv/automation/agent/middlewares/mcp_session.py new file mode 100644 index 000000000..ee95bef2a --- /dev/null +++ b/daiv/automation/agent/middlewares/mcp_session.py @@ -0,0 +1,44 @@ +"""Persist MCP session ids in graph state across chat turns. + +The actual session opening is owned by ``MCPToolkit.aopen`` at the agent-run +scope; this middleware exists only to push the resulting ids dict into state so +the LangGraph checkpointer carries them to the next turn, where the chat +streaming wrapper reads them back and asks ``aopen`` to resume. + +Without this, the dict captured during ``aopen`` would die with the request and +each turn would still get a fresh browser context. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NotRequired + +from langchain.agents.middleware import AgentMiddleware, AgentState + +if TYPE_CHECKING: + from typing import Any + + +class MCPSessionState(AgentState): + """State extension carrying MCP session ids keyed by server name.""" + + mcp_session_ids: NotRequired[dict[str, str]] + + +class MCPSessionStateMiddleware(AgentMiddleware): + """Writes the current MCP session ids dict into state on each turn. + + The dict is shared by reference with ``MCPToolkit.aopen`` — it's mutated + in-place when a stale id is recovered, so by the time ``before_agent`` + fires the dict already reflects what actually got connected this turn. + """ + + state_schema = MCPSessionState + + def __init__(self, session_ids: dict[str, str]) -> None: + self._session_ids = session_ids + + async def abefore_agent(self, state: MCPSessionState, runtime: Any) -> dict[str, Any] | None: # noqa: ARG002 + if not self._session_ids: + return None + return {"mcp_session_ids": dict(self._session_ids)} diff --git a/daiv/chat/api/streaming.py b/daiv/chat/api/streaming.py index ae5c3b3f7..3364735b5 100644 --- a/daiv/chat/api/streaming.py +++ b/daiv/chat/api/streaming.py @@ -10,6 +10,8 @@ from langgraph.store.memory import InMemoryStore from automation.agent.graph import create_daiv_agent +from automation.agent.mcp.toolkits import MCPToolkit +from automation.agent.middlewares.mcp_session import MCPSessionStateMiddleware from automation.agent.utils import build_langsmith_config, get_daiv_agent_kwargs from codebase.base import Scope from codebase.context import set_runtime_ctx @@ -37,6 +39,23 @@ HEARTBEAT_INTERVAL_S = 5.0 +async def _aread_mcp_session_ids(checkpointer: Any, thread_id: str) -> dict[str, str]: + """Return the previous turn's MCP session ids from the checkpointer, or empty. + + Loaded once at stream start so ``MCPToolkit.aopen`` can resume stateful MCP + sessions (Playwright browser context) instead of opening fresh ones. + """ + try: + tup = await checkpointer.aget_tuple({"configurable": {"thread_id": thread_id}}) + except Exception: + logger.warning("chat: failed to load checkpoint for thread_id=%s", thread_id, exc_info=True) + return {} + if tup is None: + return {} + ids = (tup.checkpoint or {}).get("channel_values", {}).get("mcp_session_ids") or {} + return {str(k): str(v) for k, v in ids.items()} if isinstance(ids, dict) else {} + + class RuntimeContextLangGraphAGUIAgent(LangGraphAGUIAgent): """Inject the daiv RuntimeCtx dataclass into upstream's stream kwargs. @@ -120,14 +139,25 @@ async def events(self) -> AsyncIterator[str]: repo_id=self.repo_id, scope=Scope.GLOBAL, ref=self.ref, sandbox_env_id=self.sandbox_environment_id ) as runtime_ctx, ): - agent_kwargs = get_daiv_agent_kwargs( - model_config=runtime_ctx.config.models.agent, - agent_model=self.agent_model, - agent_thinking_level=self.agent_thinking_level, - ) - agent = await create_daiv_agent( - ctx=runtime_ctx, checkpointer=checkpointer, store=InMemoryStore(), **agent_kwargs - ) + # Load MCP session ids from the previous turn's checkpoint so + # stateful servers (Playwright above all) resume their browser + # context instead of starting fresh; the middleware below writes + # the current ids back to state for the next turn. + prev_session_ids = await _aread_mcp_session_ids(checkpointer, self.thread_id) + async with MCPToolkit.aopen(session_ids=prev_session_ids) as (mcp_tools, current_session_ids): + agent_kwargs = get_daiv_agent_kwargs( + model_config=runtime_ctx.config.models.agent, + agent_model=self.agent_model, + agent_thinking_level=self.agent_thinking_level, + ) + agent = await create_daiv_agent( + ctx=runtime_ctx, + mcp_tools=mcp_tools, + checkpointer=checkpointer, + store=InMemoryStore(), + middleware=[MCPSessionStateMiddleware(current_session_ids)], + **agent_kwargs, + ) langsmith_config = build_langsmith_config( runtime_ctx, trigger="chat", diff --git a/daiv/codebase/managers/issue_addressor.py b/daiv/codebase/managers/issue_addressor.py index cee547b5f..932b788de 100644 --- a/daiv/codebase/managers/issue_addressor.py +++ b/daiv/codebase/managers/issue_addressor.py @@ -6,6 +6,7 @@ from langchain_core.messages import HumanMessage from automation.agent.graph import create_daiv_agent +from automation.agent.mcp.toolkits import MCPToolkit from automation.agent.usage_tracking import build_usage_summary, track_usage_metadata from automation.agent.utils import build_langsmith_config, extract_text_content, get_daiv_agent_kwargs from automation.agent.validators import AgentConfigurationError @@ -106,7 +107,7 @@ async def _address_issue(self) -> AgentResult: HumanMessage(name=self.issue.author.username, id=str(self.issue.iid), content=message_content) ) - async with open_checkpointer() as checkpointer: + async with open_checkpointer() as checkpointer, MCPToolkit.aopen() as mcp_tools: try: agent_kwargs = get_daiv_agent_kwargs( model_config=self.ctx.config.models.agent, use_max=self.issue.has_max_label() @@ -119,7 +120,7 @@ async def _address_issue(self) -> AgentResult: ) return daiv_agent = await create_daiv_agent( - ctx=self.ctx, checkpointer=checkpointer, store=self.store, **agent_kwargs + ctx=self.ctx, mcp_tools=mcp_tools, checkpointer=checkpointer, store=self.store, **agent_kwargs ) agent_config = build_langsmith_config( self.ctx, diff --git a/daiv/codebase/managers/review_addressor.py b/daiv/codebase/managers/review_addressor.py index 78c474afd..cdb6c0d8b 100644 --- a/daiv/codebase/managers/review_addressor.py +++ b/daiv/codebase/managers/review_addressor.py @@ -12,6 +12,7 @@ from unidiff.patch import Line from automation.agent.graph import create_daiv_agent +from automation.agent.mcp.toolkits import MCPToolkit from automation.agent.usage_tracking import build_usage_summary, track_usage_metadata from automation.agent.utils import build_langsmith_config, extract_text_content, get_daiv_agent_kwargs from automation.agent.validators import AgentConfigurationError @@ -269,10 +270,10 @@ async def _address_comments(self) -> AgentResult: self.ctx.repository.slug, self.merge_request.merge_request_id, self.mention_comment_id ) - async with open_checkpointer() as checkpointer: + async with open_checkpointer() as checkpointer, MCPToolkit.aopen() as mcp_tools: agent_kwargs = get_daiv_agent_kwargs(model_config=self.ctx.config.models.agent) daiv_agent = await create_daiv_agent( - ctx=self.ctx, checkpointer=checkpointer, store=self.store, **agent_kwargs + ctx=self.ctx, mcp_tools=mcp_tools, checkpointer=checkpointer, store=self.store, **agent_kwargs ) agent_config = build_langsmith_config( self.ctx, diff --git a/daiv/jobs/tasks.py b/daiv/jobs/tasks.py index a83a2fa87..602d0f3fe 100644 --- a/daiv/jobs/tasks.py +++ b/daiv/jobs/tasks.py @@ -4,6 +4,7 @@ from langchain_core.messages import HumanMessage from automation.agent.graph import create_daiv_agent +from automation.agent.mcp.toolkits import MCPToolkit from automation.agent.results import AgentResult, build_agent_result from automation.agent.usage_tracking import build_usage_summary, track_usage_metadata from automation.agent.utils import build_langsmith_config, extract_text_content, get_daiv_agent_kwargs @@ -56,6 +57,7 @@ async def run_job_task( repo_id=repo_id, scope=Scope.GLOBAL, ref=ref, sandbox_env_id=sandbox_environment_id ) as runtime_ctx, open_checkpointer() as checkpointer, + MCPToolkit.aopen() as mcp_tools, ): agent_kwargs = get_daiv_agent_kwargs( model_config=runtime_ctx.config.models.agent, @@ -70,7 +72,9 @@ async def run_job_task( extra_metadata={"ref": ref, "override_source": "explicit" if agent_model else None}, configurable={"thread_id": thread_id}, ) - daiv_agent = await create_daiv_agent(ctx=runtime_ctx, checkpointer=checkpointer, **agent_kwargs) + daiv_agent = await create_daiv_agent( + ctx=runtime_ctx, mcp_tools=mcp_tools, checkpointer=checkpointer, **agent_kwargs + ) with track_usage_metadata() as usage_handler: result = await daiv_agent.ainvoke(input_data, config=config, context=runtime_ctx) except Exception: diff --git a/docker-compose.yml b/docker-compose.yml index 06153b7bd..3b22ceea7 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -223,6 +223,37 @@ services: retries: 3 start_period: 30s + mcp_playwright: + image: mcr.microsoft.com/playwright/mcp:latest + profiles: + - mcp + - full + restart: unless-stopped + container_name: daiv-mcp-playwright + environment: + # Semicolon-separated origin lists. Empty defaults = no restriction. + # Per upstream docs these flags are a usability/scope guardrail, NOT a + # security boundary — they do not affect redirects. + MCP_PLAYWRIGHT_ALLOWED_ORIGINS: ${MCP_PLAYWRIGHT_ALLOWED_ORIGINS:-} + MCP_PLAYWRIGHT_BLOCKED_ORIGINS: ${MCP_PLAYWRIGHT_BLOCKED_ORIGINS:-} + entrypoint: ["sh", "-c"] + # --allowed-hosts is DNS-rebinding protection (Host header allowlist, exact host:port match). + # Default is just the bound host, which excludes the docker service name `mcp_playwright` + # used by the daiv app — must enumerate it here. Includes `localhost:8931` for the healthcheck. + command: + - >- + set -eu; + exec node /app/cli.js --headless --isolated --browser chromium --host 0.0.0.0 --port 8931 + --allowed-hosts mcp_playwright:8931,localhost:8931 + $${MCP_PLAYWRIGHT_ALLOWED_ORIGINS:+--allowed-origins "$$MCP_PLAYWRIGHT_ALLOWED_ORIGINS"} + $${MCP_PLAYWRIGHT_BLOCKED_ORIGINS:+--blocked-origins "$$MCP_PLAYWRIGHT_BLOCKED_ORIGINS"} + healthcheck: + test: ["CMD-SHELL", "node -e \"require('http').get('http://localhost:8931/mcp', r => process.exit(0)).on('error', () => process.exit(1))\""] + interval: 30s + timeout: 5s + retries: 3 + start_period: 30s + smee: build: context: ./ diff --git a/docs/customization/mcp-tools.md b/docs/customization/mcp-tools.md index 200d86923..f883b10da 100644 --- a/docs/customization/mcp-tools.md +++ b/docs/customization/mcp-tools.md @@ -1,6 +1,6 @@ # MCP Tools -[Model Context Protocol (MCP)](https://modelcontextprotocol.io/) tools extend DAIV's agent with access to external services. Each MCP server runs in its own isolated container via [supergateway](https://github.com/supercorp-ai/supergateway), ensuring that a compromised or misbehaving server cannot affect others. +[Model Context Protocol (MCP)](https://modelcontextprotocol.io/) tools extend DAIV's agent with access to external services. Each MCP server runs in its own isolated container — most are wrapped by [supergateway](https://github.com/supercorp-ai/supergateway) to bridge stdio servers to streamable HTTP; servers that speak HTTP natively (like Playwright) run directly. Either way, a compromised or misbehaving server cannot affect others. ## Built-in MCP servers @@ -49,6 +49,40 @@ MCP_CONTEXT7_URL=http://mcp-context7:8000/mcp # Default; set to None to disable Context7 credentials (`CONTEXT7_API_KEY`) are consumed by the `mcp-context7` container. Add them to your secrets env file. +### Playwright + +The [Playwright MCP Server](https://github.com/microsoft/playwright-mcp) gives DAIV a headless Chromium browser for inspecting web pages, capturing screenshots, reading the console, and driving UI interactions. Each agent session gets a fresh, isolated browser context. + +**Available tools** (selected — the full surface is exposed, ~23 tools). Names below are what the agent sees: the upstream server emits `browser_*` tool names; DAIV prefixes them with the server name (`playwright`) so the agent calls them as `playwright_browser_*`. + +| Tool | Description | +|------|-------------| +| `playwright_browser_navigate` | Open a URL in the browser | +| `playwright_browser_take_screenshot` | Capture a screenshot of the current page | +| `playwright_browser_snapshot` | Get an accessibility-tree snapshot of the page | +| `playwright_browser_console_messages` | Read messages logged to the JavaScript console | +| `playwright_browser_network_requests` | Read network requests made by the page | +| `playwright_browser_click` / `playwright_browser_type` / `playwright_browser_fill_form` | Interact with form elements and controls | +| `playwright_browser_run_code_unsafe` | Evaluate arbitrary JavaScript in the page context | + +**Use cases:** Verifying a frontend change by loading the running app and screenshotting the result; reproducing UI bug reports against a live URL to inspect console and network state; exercising end-to-end flows the agent is building. + +**Configuration:** + +```bash +MCP_PLAYWRIGHT_URL=http://mcp_playwright:8931/mcp # Default; set to None to disable + +# Optional: restrict where the browser may navigate. Semicolon-separated. +# Empty = no restriction. The blocklist is evaluated before the allowlist. +MCP_PLAYWRIGHT_ALLOWED_ORIGINS=https://app.example.com;https://docs.example.com +MCP_PLAYWRIGHT_BLOCKED_ORIGINS=https://ads.example.com +``` + +No credentials required. The container runs `mcr.microsoft.com/playwright/mcp` with `--headless --isolated --browser chromium` and exposes the MCP endpoint on port `8931`. + +!!! warning + Unlike Sentry and Context7, the Playwright server exposes its **full tool surface** with no filter — including `browser_run_code_unsafe` (arbitrary JavaScript execution). The agent can navigate to any URL reachable from the container. The origin allow/block lists above are a usability guardrail (the upstream docs explicitly call out that they are *not* a security boundary and *do not* affect redirects); if you need real network isolation, deploy the container behind an egress proxy. Enable the `mcp` compose profile only when you accept this trust boundary. + ## User-defined MCP servers DAIV can connect to any MCP server that implements the [Model Context Protocol](https://modelcontextprotocol.io/). If you need to build a custom server, see the [MCP server documentation](https://modelcontextprotocol.io/docs/concepts/servers). This section covers how to connect an existing server to DAIV. diff --git a/docs/reference/env-variables.md b/docs/reference/env-variables.md index c0168a894..f60274223 100644 --- a/docs/reference/env-variables.md +++ b/docs/reference/env-variables.md @@ -269,6 +269,14 @@ MCP (Model Context Protocol) tools extend agent capabilities by providing access | `MCP_SERVERS_CONFIG_FILE` | Path to user-defined MCP servers JSON config file | *(none)* | `/path/to/mcp.json` | | `MCP_SENTRY_URL` | Streamable HTTP URL for the Sentry supergateway container (set to `None` to disable) | `http://mcp-sentry:8000/mcp` | `http://localhost:8001/mcp` | | `MCP_CONTEXT7_URL` | Streamable HTTP URL for the Context7 supergateway container (set to `None` to disable) | `http://mcp-context7:8000/mcp` | `http://localhost:8002/mcp` | +| `MCP_PLAYWRIGHT_URL` | Streamable HTTP URL for the Playwright browser MCP container (set to `None` to disable) | `http://mcp_playwright:8931/mcp` | `http://localhost:8931/mcp` | + +The Playwright container itself reads two additional environment variables — set them on the `mcp_playwright` service, **not** on the DAIV app: + +| Variable | Description | Default | Example | +|---------------------------------|----------------------------------------------------------------|:------------------------------:|---------| +| `MCP_PLAYWRIGHT_ALLOWED_ORIGINS` | Semicolon-separated origins the browser is allowed to navigate to. Empty = allow all. Not a security boundary; redirects bypass it. | *(empty)* | `https://app.example.com;https://docs.example.com` | +| `MCP_PLAYWRIGHT_BLOCKED_ORIGINS` | Semicolon-separated origins the browser is blocked from. Evaluated before the allowlist. | *(empty)* | `https://ads.example.com;https://tracker.test` | !!! info For detailed MCP server configuration including user-defined servers, see [MCP Tools](../customization/mcp-tools.md). diff --git a/tests/integration_tests/test_sandbox.py b/tests/integration_tests/test_sandbox.py index 7566b971a..e51b1937e 100644 --- a/tests/integration_tests/test_sandbox.py +++ b/tests/integration_tests/test_sandbox.py @@ -85,6 +85,7 @@ async def test_sandbox_bash_tool_activated(model_name, inputs, runtime_ctx): agent = await create_daiv_agent( ctx=runtime_ctx, + mcp_tools=[], model_names=[model_name], auto_commit_changes=False, interrupt_on=INTERRUPT_ALL_TOOLS_CONFIG, @@ -147,6 +148,7 @@ async def test_sandbox_policy_blocks_forbidden_commands(model_name, user_message agent = await create_daiv_agent( ctx=runtime_ctx, + mcp_tools=[], model_names=[model_name], auto_commit_changes=False, interrupt_on=INTERRUPT_ALL_TOOLS_CONFIG, diff --git a/tests/integration_tests/test_skills.py b/tests/integration_tests/test_skills.py index 1368b6b37..38dc7bea4 100644 --- a/tests/integration_tests/test_skills.py +++ b/tests/integration_tests/test_skills.py @@ -35,6 +35,7 @@ async def test_skill_activated(model_name, user_message, skill): async with set_runtime_ctx(repo_id="srtab/daiv", scope=Scope.GLOBAL, ref="main") as ctx: agent = await create_daiv_agent( ctx=ctx, + mcp_tools=[], model_names=[model_name], auto_commit_changes=False, checkpointer=InMemorySaver(), diff --git a/tests/unit_tests/automation/agent/mcp/test_servers.py b/tests/unit_tests/automation/agent/mcp/test_servers.py index 57a849153..6bf83dcbd 100644 --- a/tests/unit_tests/automation/agent/mcp/test_servers.py +++ b/tests/unit_tests/automation/agent/mcp/test_servers.py @@ -1,6 +1,6 @@ from unittest.mock import patch -from automation.agent.mcp.servers import Context7MCPServer, SentryMCPServer +from automation.agent.mcp.servers import Context7MCPServer, PlaywrightMCPServer, SentryMCPServer class TestSentryMCPServer: @@ -75,3 +75,42 @@ def test_context7_get_connection_returns_correct_url(self, mock_settings): assert connection["transport"] == "streamable_http" assert connection["url"] == "http://mcp-context7:8000/mcp" + + +class TestPlaywrightMCPServer: + def test_playwright_server_has_correct_name(self): + server = PlaywrightMCPServer() + assert server.name == "playwright" + + @patch("automation.agent.mcp.servers.settings") + def test_playwright_server_is_enabled_when_url_set(self, mock_settings): + mock_settings.PLAYWRIGHT_URL = "http://mcp_playwright:8931/mcp" + server = PlaywrightMCPServer() + + assert server.is_enabled() is True + + @patch("automation.agent.mcp.servers.settings") + def test_playwright_server_is_disabled_when_url_none(self, mock_settings): + mock_settings.PLAYWRIGHT_URL = None + server = PlaywrightMCPServer() + + assert server.is_enabled() is False + + def test_playwright_server_has_no_tool_filter(self): + """Documents the 'full Playwright surface' decision as a test assertion. + + A future change that adds a filter has to update this test, making the + deviation visible in the diff. + """ + server = PlaywrightMCPServer() + + assert server.tool_filter is None + + @patch("automation.agent.mcp.servers.settings") + def test_playwright_get_connection_returns_correct_url(self, mock_settings): + mock_settings.PLAYWRIGHT_URL = "http://mcp_playwright:8931/mcp" + server = PlaywrightMCPServer() + connection = server.get_connection() + + assert connection["transport"] == "streamable_http" + assert connection["url"] == "http://mcp_playwright:8931/mcp" diff --git a/tests/unit_tests/automation/agent/mcp/test_toolkits.py b/tests/unit_tests/automation/agent/mcp/test_toolkits.py index cca180012..dd4451c80 100644 --- a/tests/unit_tests/automation/agent/mcp/test_toolkits.py +++ b/tests/unit_tests/automation/agent/mcp/test_toolkits.py @@ -1,7 +1,10 @@ -from unittest.mock import MagicMock +from contextlib import asynccontextmanager +from unittest.mock import MagicMock, patch + +import pytest from automation.agent.mcp.schemas import ToolFilter -from automation.agent.mcp.toolkits import _apply_tool_filters +from automation.agent.mcp.toolkits import MCPToolkit, _apply_tool_filters def _make_tool(name: str) -> MagicMock: @@ -87,3 +90,122 @@ def test_multiple_servers_filtered_independently(self): assert "sentry_delete_project" not in names assert "context7_query-docs" in names assert "context7_resolve-library-id" not in names + + +class TestAopenSessionIdHandling: + """The cross-turn resume logic: ``MCPToolkit.aopen`` reconnects to existing + server-side sessions when a previous id is supplied, recovers from a stale + id by opening fresh, and drops empty slots from the dict it yields back. + + We stub the low-level session opener so the test runs without a real MCP + server but still exercises every branch of ``_open_server``. + """ + + @pytest.fixture(autouse=True) + def mock_mcp_toolkit_aopen(self): + """Override the conftest autouse fixture that stubs ``aopen`` so this + class exercises the real implementation.""" + yield + + async def _run_aopen(self, *, connections, session_open_results, session_ids=None): + """Patch the registry + per-server opener and drive ``aopen`` once. + + ``session_open_results`` is a dict mapping ``(server_name, mode)`` — + where ``mode`` is "resume" if a session id was passed in or "fresh" + otherwise — to ``(tools, captured_id)`` or to an ``Exception`` to raise. + """ + captured_calls: list[tuple[str, str]] = [] + + @asynccontextmanager + async def _fake_opener(*, url, headers, terminate_on_close, initialize): + # `initialize=False` only happens on the resume branch (per _open_server), + # so we can use it to dispatch the right scripted result. + mode = "fresh" if initialize else "resume" + server_name = next(name for name, conn in connections.items() if conn["url"] == url) + captured_calls.append((server_name, mode)) + outcome = session_open_results[(server_name, mode)] + if isinstance(outcome, Exception): + raise outcome + + session = MagicMock(name=f"session-{server_name}") + get_id = lambda captured=outcome[1]: captured # noqa: E731 + yield session, get_id + + async def _fake_load_tools(session, *, server_name, tool_name_prefix): + # Dispatch on the most recently opened (server, mode) — this matters + # in the recovery path where the same server is opened twice. + outcome = session_open_results[captured_calls[-1]] + return list(outcome[0]) + + registry = MagicMock() + registry.get_connections_and_filters.return_value = (connections, {}) + + # ``aopen`` imports the registry at runtime via + # ``from automation.agent.mcp.registry import mcp_registry`` — patch the + # symbol at its source so the local rebinding inside the function sees the mock. + with ( + patch("automation.agent.mcp.toolkits._open_streamable_mcp_session", _fake_opener), + patch("automation.agent.mcp.toolkits.load_mcp_tools", side_effect=_fake_load_tools), + patch("automation.agent.mcp.registry.mcp_registry", registry), + ): + async with MCPToolkit.aopen(session_ids=session_ids) as (tools, ids): + return tools, ids, captured_calls + + async def test_fresh_open_captures_ids_when_persist_requested(self): + connections = {"playwright": {"transport": "streamable_http", "url": "http://pw/mcp", "headers": None}} + tools, ids, calls = await self._run_aopen( + connections=connections, + session_open_results={("playwright", "fresh"): ([_make_tool("playwright_browser_navigate")], "abc-123")}, + session_ids={}, + ) + + assert calls == [("playwright", "fresh")] + assert ids == {"playwright": "abc-123"} + assert [t.name for t in tools] == ["playwright_browser_navigate"] + + async def test_resume_skips_initialize_and_keeps_id(self): + connections = {"playwright": {"transport": "streamable_http", "url": "http://pw/mcp", "headers": None}} + tools, ids, calls = await self._run_aopen( + connections=connections, + session_open_results={("playwright", "resume"): ([_make_tool("playwright_browser_snapshot")], "abc-123")}, + session_ids={"playwright": "abc-123"}, + ) + + assert calls == [("playwright", "resume")] + assert ids == {"playwright": "abc-123"} + + async def test_resume_failure_falls_back_to_fresh_and_overwrites_id(self): + connections = {"playwright": {"transport": "streamable_http", "url": "http://pw/mcp", "headers": None}} + tools, ids, calls = await self._run_aopen( + connections=connections, + session_open_results={ + ("playwright", "resume"): RuntimeError("404 Session not found"), + ("playwright", "fresh"): ([_make_tool("playwright_browser_navigate")], "new-456"), + }, + session_ids={"playwright": "stale-id"}, + ) + + assert calls == [("playwright", "resume"), ("playwright", "fresh")] + assert ids == {"playwright": "new-456"}, "stale id must be replaced, not preserved" + + async def test_stateless_server_does_not_pollute_ids_dict(self): + connections = {"sentry": {"transport": "streamable_http", "url": "http://sn/mcp", "headers": None}} + tools, ids, calls = await self._run_aopen( + connections=connections, + # Stateless servers (supergateway-fronted) return no Mcp-Session-Id header, + # so the get_session_id callback yields None. + session_open_results={("sentry", "fresh"): ([_make_tool("sentry_whoami")], None)}, + session_ids={}, + ) + + assert ids == {}, "an empty id slot would make the next turn attempt a doomed resume" + + async def test_no_session_ids_means_no_persistence_mode(self): + connections = {"playwright": {"transport": "streamable_http", "url": "http://pw/mcp", "headers": None}} + tools, ids, calls = await self._run_aopen( + connections=connections, + session_open_results={("playwright", "fresh"): ([_make_tool("playwright_browser_navigate")], "abc-123")}, + session_ids=None, + ) + + assert ids == {}, "session_ids=None signals one-shot mode; returned dict is empty" diff --git a/tests/unit_tests/automation/agent/test_graph_deferred.py b/tests/unit_tests/automation/agent/test_graph_deferred.py index fc9372d2d..b52ea5b2e 100644 --- a/tests/unit_tests/automation/agent/test_graph_deferred.py +++ b/tests/unit_tests/automation/agent/test_graph_deferred.py @@ -10,7 +10,6 @@ def _common_patches(): patch("automation.agent.graph.create_explore_subagent"), patch("automation.agent.graph.load_custom_subagents", AsyncMock(return_value=[])), patch("automation.agent.graph.create_deep_agent"), - patch("automation.agent.graph.MCPToolkit"), patch("automation.agent.graph.deferred_settings"), patch("automation.agent.graph.BaseAgent"), patch("automation.agent.graph.site_settings"), @@ -39,7 +38,6 @@ async def _run(self, *, flag_on: bool): mock_create_explore, mock_load_custom, mock_create_deep_agent, - mock_toolkit, mock_deferred_settings, mock_base_agent, mock_site_settings, @@ -51,7 +49,6 @@ async def _run(self, *, flag_on: bool): mock_deferred_settings.TOP_K_MAX = 10 mcp_tool = StructuredTool.from_function(func=lambda **k: "x", name="t1", description="d") - mock_toolkit.get_tools = AsyncMock(return_value=[mcp_tool]) mock_base_agent.get_model.return_value = MagicMock() @@ -70,25 +67,23 @@ async def _run(self, *, flag_on: bool): ctx.config.context_file_name = "AGENTS.md" ctx.git_platform = MagicMock() - await create_daiv_agent(ctx=ctx, auto_commit_changes=False) - return mock_create_deep_agent, mock_toolkit, mcp_tool + await create_daiv_agent(ctx=ctx, mcp_tools=[mcp_tool], auto_commit_changes=False) + return mock_create_deep_agent, mcp_tool finally: for p in patches: p.stop() async def test_flag_off_passes_eager_tools(self): - mock_create_deep_agent, mock_toolkit, mcp_tool = await self._run(flag_on=False) + mock_create_deep_agent, mcp_tool = await self._run(flag_on=False) - mock_toolkit.get_tools.assert_awaited() kwargs = mock_create_deep_agent.call_args.kwargs assert kwargs["tools"] == [mcp_tool] middleware_types = [type(m).__name__ for m in kwargs["middleware"]] assert "DeferredToolsMiddleware" not in middleware_types async def test_flag_on_passes_empty_tools_and_adds_middleware(self): - mock_create_deep_agent, mock_toolkit, _ = await self._run(flag_on=True) + mock_create_deep_agent, _ = await self._run(flag_on=True) - mock_toolkit.get_tools.assert_awaited() kwargs = mock_create_deep_agent.call_args.kwargs assert kwargs["tools"] == [] middleware_types = [type(m).__name__ for m in kwargs["middleware"]] diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index 35d25e7cd..7da7897aa 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -1,4 +1,4 @@ -from contextlib import contextmanager +from contextlib import asynccontextmanager, contextmanager from pathlib import Path from tempfile import TemporaryDirectory from unittest.mock import AsyncMock, Mock, patch @@ -83,6 +83,24 @@ def mock_settings(monkeypatch): yield codebase_settings +@pytest.fixture(autouse=True) +def mock_mcp_toolkit_aopen(): + """Stub ``MCPToolkit.aopen`` so unit tests never reach real MCP servers. + + Production wraps every agent run with ``async with MCPToolkit.aopen() as mcp_tools:``; + without this fixture the real classmethod would try to open sessions against + ``mcp_sentry`` / ``mcp_context7`` / ``mcp_playwright`` and hang or error out. + Tests that specifically exercise the MCP pool override this fixture locally. + """ + + @asynccontextmanager + async def _aopen_empty(*, session_ids=None): + yield [], dict(session_ids or {}) + + with patch("automation.agent.mcp.toolkits.MCPToolkit.aopen", _aopen_empty): + yield + + @pytest.fixture(autouse=True) def mock_generate_title_task(): """Stub the titling tasks so the ImmediateBackend doesn't fire real LLM calls.