feat(agent): add Playwright browser MCP integration

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(

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()

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)

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]:

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)}

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",