feat: add streaming to markdown, fix raycast mcps exposing

src/beaver_gateway/backends/raycast.py

@@ -17,12 +17,24 @@ Two halves live here:
   ``arguments``) — Anthropic wants one ``content_block_start`` → deltas
   → ``content_block_stop`` per block, so we de-duplicate the final
   summary against the per-delta increments already emitted.
+
+MCP wiring: Raycast has no native MCP concept, so when an agent declares
+``expose_mcps`` we splice each MCP's tools into the wire request as
+``Tool.local(name=f"{mcp}__{tool}", ...)`` and run a gateway-internal
+loop. Every time the model emits a tool_call for one of those local
+tools we route it back to the underlying MCP in-process, append the
+result as a ``tool`` message, and re-issue the stream — all inside one
+Anthropic envelope (one ``message_start`` … one ``message_stop``). The
+tool_use blocks DO surface to the caller (mirrors what
+``ClaudeCodeBackendAdapter`` does), but tool_results stay internal.
 """
 
 from __future__ import annotations
 
 import json
+import logging
 import uuid
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, cast
 
 from raycast_api import Message as RaycastMessage
@@ -47,6 +59,8 @@ if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
 
     from anthropic.types import MessageParam
+    from fastmcp import FastMCP
+    from fastmcp.tools.base import Tool as FastMCPTool
     from raycast_api import ChatStreamChunk, Client
 
     from beaver_gateway.agents.base import BaseAgent
@@ -58,6 +72,16 @@ else:
 __all__ = ["RaycastBackend"]
 
 
+_log = logging.getLogger("beaver_gateway.backends.raycast")
+
+
+# Cap on consecutive tool-call turns inside one Anthropic envelope.
+# Real conversations rarely chain more than a handful; the limit only
+# fires on a model that loops, and surfaces as a clean ``end_turn``
+# with an error tool_result rather than a hang.
+_MAX_TOOL_TURNS = 20
+
+
 _RAYCAST_TO_ANTHROPIC_STOP: dict[str, StopReason] = {
     "stop": "end_turn",
     "STOP": "end_turn",
@@ -132,9 +156,7 @@ def _extract_tool_result_text(content: object) -> str:
     return json.dumps(content, separators=(",", ":"), ensure_ascii=False)
 
 
-def _to_raycast_messages(
-    messages: Iterable[MessageParam],
-) -> list[RaycastMessage]:
+def _to_raycast_messages(messages: Iterable[MessageParam]) -> list[RaycastMessage]:
     """Convert an Anthropic message history into a Raycast one.
 
     ``tool_use_id → name`` is tracked across the iteration so that
@@ -206,27 +228,143 @@ def _to_raycast_messages(
                 )
         out.append(
             RaycastMessage.assistant(
-                text="\n".join(text_parts),
-                tool_calls=tool_calls or None,
+                text="\n".join(text_parts), tool_calls=tool_calls or None
             )
         )
 
     return out
 
 
-def _build_tool_list(
-    agent: RaycastAgent,
-) -> list[Tool | RemoteTool | str] | None:
-    """Wrap each native-tool name as a Raycast remote tool.
+@dataclass(frozen=True, slots=True)
+class _AgentToolCatalog:
+    """Per-agent MCP tool catalog + routing map.
 
-    Phase 1.2 scope: only the three model-agnostic remote tools
-    (`web_search`, `search_images`, `read_page`). Client-defined tools
-    coming through the Anthropic body's ``tools`` field stay out — that's
-    Phase 1.3+ alongside ``accept_client_tools``.
+    Cached on :class:`RaycastBackend` keyed by agent name. ``tools`` is
+    the fully-rendered ``tools`` argument for ``chat.stream`` (local MCP
+    tools first, native remote tools last). ``mcp_routing`` maps each
+    wire-name we registered back to the underlying ``(mcp, original)``
+    pair so we can dispatch a tool_call without re-parsing the prefix.
     """
-    if not agent.available_native_tools:
-        return None
-    return [Tool.remote(name) for name in agent.available_native_tools]
+
+    tools: tuple[Tool, ...]
+    mcp_routing: Mapping[str, tuple[str, str]]
+
+
+_EMPTY_SCHEMA: dict[str, Any] = {"type": "object", "properties": {}}
+
+
+def _sanitize_schema(schema: Mapping[str, Any] | None) -> dict[str, Any]:
+    """Strip JSON-Schema meta keys that Raycast/Gemini's tool API rejects.
+
+    MCP servers ship draft-07 schemas including ``$schema``/``$id``/etc.
+    OpenAI-style function-calling expects a leaner subset — leaving the
+    meta keys in causes the upstream provider to reject the entire
+    request with an opaque ``unknown_api_error``. We drop the
+    document-level meta keys and keep the structural keys
+    (``type``/``properties``/``required``/...).
+    """
+    if not schema:
+        return dict(_EMPTY_SCHEMA)
+    out: dict[str, Any] = {}
+    for key, value in schema.items():
+        if key.startswith("$"):
+            continue
+        out[key] = value
+    if "type" not in out:
+        out["type"] = "object"
+    if out.get("type") == "object" and "properties" not in out:
+        out["properties"] = {}
+    return out
+
+
+def _build_agent_catalog(
+    agent: RaycastAgent, mcp_tools: Mapping[str, list[FastMCPTool]]
+) -> _AgentToolCatalog:
+    """Resolve ``agent.expose_mcps`` against the prefetched MCP tool lists.
+
+    Each exposed MCP contributes one ``Tool.local`` per tool, named
+    ``{mcp}__{tool}`` (mirrors claude-code's wire convention, so a model
+    that has seen one style sees a familiar one here). ``ExposedMcp.tools``
+    optionally filters to a subset by original (unprefixed) name.
+
+    A missing MCP entry (broken at startup, see
+    ``cli._prefetch_mcp_tools``) silently contributes nothing — surfaces
+    in logs at start, doesn't crash request-time.
+    """
+    routing: dict[str, tuple[str, str]] = {}
+    local_tools: list[Tool] = []
+    for em in agent.expose_mcps:
+        for mt in mcp_tools.get(em.name, []):
+            if em.tools is not None and mt.name not in em.tools:
+                continue
+            wire_name = f"{em.name}__{mt.name}"
+            routing[wire_name] = (em.name, mt.name)
+            local_tools.append(
+                Tool.local(
+                    name=wire_name,
+                    description=mt.description or "",
+                    parameters=_sanitize_schema(mt.parameters),
+                )
+            )
+    remote_tools = [Tool.remote(n) for n in agent.available_native_tools]
+    return _AgentToolCatalog(
+        tools=tuple(local_tools + remote_tools), mcp_routing=routing
+    )
+
+
+def _render_mcp_tool_result(result: Any) -> str:
+    """Flatten a ``fastmcp.ToolResult`` into the text Raycast carries.
+
+    Raycast ``tool`` messages only hold a string. MCP results can carry
+    text, images, structured payloads, etc. — we keep text blocks
+    verbatim, JSON-encode anything else, and fall back to
+    ``structured_content`` if the content list is empty.
+    """
+    parts: list[str] = []
+    for block in getattr(result, "content", None) or []:
+        text = getattr(block, "text", None)
+        if isinstance(text, str):
+            parts.append(text)
+            continue
+        dump = getattr(block, "model_dump", None)
+        if callable(dump):
+            parts.append(json.dumps(dump(), ensure_ascii=False))
+        else:
+            parts.append(str(block))
+    if not parts:
+        structured = getattr(result, "structured_content", None)
+        if structured is not None:
+            parts.append(json.dumps(structured, ensure_ascii=False))
+    return "\n".join(parts)
+
+
+@dataclass(slots=True)
+class _TurnAccumulator:
+    """Pieces collected from one Raycast stream needed to feed the next.
+
+    ``text_parts`` is the assistant's text portion (joined and replayed
+    on the next ``chat.stream`` call so the model sees its own prior
+    reply). ``tool_*`` mirrors the same data the wire-state already
+    emitted, but kept in raw form so we can build ``ToolCall`` objects
+    and dispatch them through the gateway.
+    """
+
+    text_parts: list[str] = field(default_factory=list)
+    tool_order: list[str] = field(default_factory=list)
+    tool_names: dict[str, str] = field(default_factory=dict)
+    tool_args: dict[str, list[str]] = field(default_factory=dict)
+    finish_reason: str | None = None
+    usage: dict[str, int] | None = None
+
+    def add_tool_phase1(self, tool_id: str, name: str) -> None:
+        if tool_id in self.tool_names:
+            return
+        self.tool_order.append(tool_id)
+        self.tool_names[tool_id] = name
+        self.tool_args.setdefault(tool_id, [])
+
+    def add_tool_args(self, tool_id: str, fragment: str) -> None:
+        self.tool_args.setdefault(tool_id, []).append(fragment)
 
 
 class _BlockState:
@@ -241,6 +379,10 @@ class _BlockState:
     in phase 1 only and ``index`` in every chunk: ``tool_id_to_block``
     keys by Raycast tool-call id, ``tool_idx_to_id`` resolves chunk
     indices back to that id for the no-id delta chunks.
+
+    State spans the entire Anthropic envelope (all Raycast turns) so
+    block indices grow monotonically across the tool-call loop. The
+    per-turn ``_TurnAccumulator`` carries the volatile bits.
     """
 
     __slots__ = ("index", "kind", "tool_id_to_block", "tool_idx_to_id")
@@ -251,18 +393,49 @@ class _BlockState:
         self.tool_id_to_block: dict[str, int] = {}
         self.tool_idx_to_id: dict[int, str] = {}
 
+    def reset_turn_indexing(self) -> None:
+        """Drop chunk-index → id routing between turns.
+
+        Raycast chunk ``index`` fields restart at 0 inside each fresh
+        ``chat.stream``; reusing the previous turn's table would alias
+        new tool_calls onto old ids. ``tool_id_to_block`` we keep — ids
+        are stream-unique and new ones won't collide.
+        """
+        self.tool_idx_to_id = {}
+
 
 class RaycastBackend:
     """Adapter from ``raycast-api`` chat streams to Anthropic stream events.
 
     Construction takes a long-lived :class:`raycast_api.Client` (one per
-    gateway — bearer + device_id are process-wide). Each
-    :meth:`complete` call opens one ``chat.stream`` and yields a fully
-    Anthropic-shaped event sequence.
+    gateway — bearer + device_id are process-wide) plus the in-process
+    MCP server map and the prefetched per-MCP tool list. Each
+    :meth:`complete` call opens one or more ``chat.stream`` calls — one
+    per agent "turn" inside the gateway-internal tool-call loop — and
+    yields a single Anthropic stream envelope spanning all of them.
     """
 
-    def __init__(self, client: Client) -> None:
+    def __init__(
+        self,
+        client: Client,
+        *,
+        mcp_servers: Mapping[str, FastMCP] | None = None,
+        mcp_tools: Mapping[str, list[FastMCPTool]] | None = None,
+    ) -> None:
         self._client = client
+        self._mcp_servers: Mapping[str, FastMCP] = mcp_servers or {}
+        self._mcp_tools: Mapping[str, list[FastMCPTool]] = mcp_tools or {}
+        # Cached per-agent catalog — agents are immutable, so a single
+        # render at first use covers the gateway's lifetime.
+        self._agent_catalog: dict[str, _AgentToolCatalog] = {}
+
+    def _catalog_for(self, agent: RaycastAgent) -> _AgentToolCatalog:
+        cached = self._agent_catalog.get(agent.name)
+        if cached is not None:
+            return cached
+        cat = _build_agent_catalog(agent, self._mcp_tools)
+        self._agent_catalog[agent.name] = cat
+        return cat
 
     async def complete(
         self,
@@ -277,7 +450,13 @@ class RaycastBackend:
             raise TypeError(msg)
 
         raycast_messages = _to_raycast_messages(messages)
-        tools = _build_tool_list(agent)
+        catalog = self._catalog_for(agent)
+        # Native remote tools + spliced MCP locals. ``None`` keeps the
+        # SDK from sending a ``tools`` field at all when neither is
+        # declared.
+        tools_arg: list[Tool | RemoteTool | str] | None = (
+            list(catalog.tools) if catalog.tools else None
+        )
 
         # On the wire Raycast uses ``system_instructions`` as a format
         # marker (``"markdown"`` for AI_CHAT, ``"plain"`` otherwise —
@@ -294,7 +473,8 @@ class RaycastBackend:
         async for event in self._stream(
             agent=agent,
             raycast_messages=raycast_messages,
-            tools=tools,
+            tools=tools_arg,
+            mcp_routing=catalog.mcp_routing,
             prompt_content=prompt_content,
             temperature=_first_set(options.get("temperature"), agent.temperature),
             reasoning_effort=_first_set(
@@ -310,6 +490,7 @@ class RaycastBackend:
         agent: RaycastAgent,
         raycast_messages: Sequence[RaycastMessage],
         tools: list[Tool | RemoteTool | str] | None,
+        mcp_routing: Mapping[str, tuple[str, str]],
         prompt_content: str | None,
         temperature: float | None,
         reasoning_effort: str | None,
@@ -319,45 +500,153 @@ class RaycastBackend:
         yield build_message_start(message_id=message_id, model=agent.model)
 
         state = _BlockState()
-        final_finish: str | None = None
-        final_usage: dict[str, int] | None = None
+        # ``working_messages`` is the rolling history fed back to
+        # Raycast as we resolve MCP tool_calls turn by turn.
+        working_messages = list(raycast_messages)
+        last_usage: dict[str, int] | None = None
+        last_finish: str | None = None
 
-        stream = self._client.chat.stream(
-            model=agent.model,
-            messages=list(raycast_messages),
-            source=agent.source,
-            # ``system_instructions=None`` → SDK substitutes the source
-            # default (``"markdown"`` / ``"plain"``). Real prompt goes
-            # into ``additional_system_instructions``.
-            additional_system_instructions=prompt_content,
-            user_preferences=agent.user_preferences,
-            tools=tools,
-            tool_choice=tool_choice,
-            temperature=temperature,
-            reasoning_effort=reasoning_effort,
-        )
+        for _turn in range(_MAX_TOOL_TURNS):
+            acc = _TurnAccumulator()
+            state.reset_turn_indexing()
 
-        async for chunk in stream:
-            for event in self._handle_chunk(chunk, state):
-                yield event
-            if chunk.finish_reason:
-                final_finish = chunk.finish_reason
-            if chunk.usage:
-                final_usage = chunk.usage
+            stream = self._client.chat.stream(
+                model=agent.model,
+                messages=working_messages,
+                source=agent.source,
+                # ``system_instructions=None`` → SDK substitutes the
+                # source default (``"markdown"`` / ``"plain"``). Real
+                # prompt goes into ``additional_system_instructions``.
+                additional_system_instructions=prompt_content,
+                user_preferences=agent.user_preferences,
+                tools=tools,
+                tool_choice=tool_choice,
+                temperature=temperature,
+                reasoning_effort=reasoning_effort,
+            )
 
-        # Close whatever block is still open before the message delta.
+            async for chunk in stream:
+                for event in self._handle_chunk(chunk, state, acc):
+                    yield event
+                if chunk.finish_reason:
+                    acc.finish_reason = chunk.finish_reason
+                if chunk.usage:
+                    acc.usage = chunk.usage
+
+            last_finish = acc.finish_reason
+            if acc.usage:
+                last_usage = acc.usage
+
+            # Figure out which (if any) tool_calls land on our MCP
+            # routing table. Anything not in the table is left to bubble
+            # out of the envelope as a regular tool_use block — the
+            # Anthropic caller can then respond with a tool_result the
+            # usual way, and the next ``complete`` invocation will
+            # carry it back in.
+            pending_mcp_calls = [
+                tid for tid in acc.tool_order if acc.tool_names.get(tid) in mcp_routing
+            ]
+            if not pending_mcp_calls:
+                break
+
+            # Close whatever block is still open before we step into
+            # tool execution — the next turn's chunks start a fresh
+            # block sequence.
+            if state.kind is not None:
+                yield build_content_block_stop(state.index)
+                state.kind = None
+
+            # Echo the assistant turn so Raycast sees its own reply +
+            # tool_calls in subsequent context.
+            assistant_text = "".join(acc.text_parts)
+            assistant_tool_calls = [
+                ToolCall(
+                    id=tid,
+                    name=acc.tool_names.get(tid, ""),
+                    arguments="".join(acc.tool_args.get(tid, [])) or "{}",
+                )
+                for tid in acc.tool_order
+            ]
+            working_messages.append(
+                RaycastMessage.assistant(
+                    text=assistant_text, tool_calls=assistant_tool_calls or None
+                )
+            )
+
+            # Dispatch each MCP-routed call and append a ``tool``
+            # message. Calls not in the routing table get a placeholder
+            # error so the model can correct itself rather than the
+            # gateway hanging the conversation.
+            for tid in acc.tool_order:
+                tool_name = acc.tool_names.get(tid, "")
+                args_str = "".join(acc.tool_args.get(tid, [])) or "{}"
+                if tool_name in mcp_routing:
+                    result_text = await self._dispatch_mcp_call(
+                        tool_name=tool_name, args_json=args_str, mcp_routing=mcp_routing
+                    )
+                else:
+                    # Non-MCP tool — shouldn't really happen because we
+                    # haven't surfaced any other locals, but defend.
+                    result_text = f"Tool {tool_name!r} is not handled by the gateway."
+                working_messages.append(
+                    RaycastMessage.tool(
+                        tool_call_id=tid, name=tool_name, result=result_text
+                    )
+                )
+
+            # Loop: re-stream with the new history.
+        else:
+            _log.warning(
+                "raycast tool-call loop hit %d-turn cap for agent %r; "
+                "closing the envelope",
+                _MAX_TOOL_TURNS,
+                agent.name,
+            )
+
+        # Close whatever block is still open before the final delta.
         if state.kind is not None:
             yield build_content_block_stop(state.index)
             state.kind = None
 
         yield build_message_delta(
-            stop_reason=_map_stop_reason(final_finish),
-            usage=final_usage,
+            stop_reason=_map_stop_reason(last_finish), usage=last_usage
         )
         yield build_message_stop()
 
+    async def _dispatch_mcp_call(
+        self,
+        *,
+        tool_name: str,
+        args_json: str,
+        mcp_routing: Mapping[str, tuple[str, str]],
+    ) -> str:
+        """Route one tool_call to the matching MCP and render its result.
+
+        Errors (missing server, JSON parse, tool exception) surface as
+        text in the ``tool`` message rather than crashing the stream —
+        the model gets a chance to recover or apologize.
+        """
+        mcp_name, original_name = mcp_routing[tool_name]
+        server = self._mcp_servers.get(mcp_name)
+        if server is None:
+            return f"Error: MCP {mcp_name!r} is not registered."
+        try:
+            args = json.loads(args_json) if args_json else {}
+        except json.JSONDecodeError as exc:
+            return f"Error: tool arguments are not valid JSON ({exc})."
+        if not isinstance(args, dict):
+            return "Error: tool arguments must be a JSON object."
+        try:
+            result = await server.call_tool(original_name, args)
+        except Exception as exc:  # noqa: BLE001
+            _log.exception(
+                "MCP call failed: %s.%s args=%r", mcp_name, original_name, args
+            )
+            return f"Error calling {tool_name}: {exc}"
+        return _render_mcp_tool_result(result)
+
     def _handle_chunk(
-        self, chunk: ChatStreamChunk, state: _BlockState
+        self, chunk: ChatStreamChunk, state: _BlockState, acc: _TurnAccumulator
     ) -> Iterable[MessageStreamEvent]:
         """Translate one Raycast chunk into zero or more Anthropic events.
 
@@ -373,6 +662,7 @@ class RaycastBackend:
         if chunk.text:
             events.extend(self._ensure_kind(state, "text"))
             events.append(build_text_delta(state.index, chunk.text))
+            acc.text_parts.append(chunk.text)
 
         if chunk.reasoning:
             events.extend(self._ensure_kind(state, "thinking"))
@@ -380,7 +670,9 @@ class RaycastBackend:
 
         if chunk.tool_calls:
             events.extend(
-                self._handle_tool_calls(chunk, state, is_final_summary=is_final_summary)
+                self._handle_tool_calls(
+                    chunk, state, acc, is_final_summary=is_final_summary
+                )
             )
 
         return events
@@ -409,6 +701,7 @@ class RaycastBackend:
         self,
         chunk: ChatStreamChunk,
         state: _BlockState,
+        acc: _TurnAccumulator,
         *,
         is_final_summary: bool,
     ) -> Iterable[MessageStreamEvent]:
@@ -419,6 +712,11 @@ class RaycastBackend:
         by id when present, otherwise by their wire ``index`` field, and
         the final-summary chunk's arguments string is dropped because the
         deltas already streamed it.
+
+        Side-effects on ``acc`` mirror what gets emitted to the wire so
+        the gateway can rebuild a full ``ToolCall`` for the next Raycast
+        turn (it needs the joined ``arguments`` JSON string, which the
+        wire never gives us as a whole — only in fragments).
         """
         events: list[MessageStreamEvent] = []
         raw_tcs = chunk.raw.get("tool_calls") or []
@@ -449,20 +747,35 @@ class RaycastBackend:
                 state.kind = "tool_use"
                 block_idx = state.index
                 state.tool_id_to_block[tool_id] = block_idx
+                acc.add_tool_phase1(tool_id, tc.name or "")
                 events.append(
                     build_tool_use_block_start(
                         block_idx, tool_use_id=tool_id, name=tc.name or ""
                     )
                 )
-                if tc.arguments and not is_final_summary:
+                # Streaming providers (GPT) deliver arguments as deltas
+                # across chunks and re-state the full string in a final
+                # summary chunk; non-streaming-args providers (Gemini)
+                # only send args in the final summary. Emit args in
+                # both cases when this is the first appearance — final
+                # summary then IS the full string.
+                if tc.arguments:
                     events.append(build_input_json_delta(block_idx, tc.arguments))
+                    acc.add_tool_args(tool_id, tc.arguments)
                 continue
 
-            # Existing block. Skip args on the final summary — they're a
-            # full restatement of what's already been delta'd.
+            # Existing block. Final summary chunks restate the full args
+            # string; if we already streamed deltas, that restatement is
+            # a duplicate (skip). If we streamed nothing (the streaming
+            # provider didn't send mid-arg deltas — Gemini path again),
+            # the summary IS the args — emit it once.
             if is_final_summary:
+                if tc.arguments and not acc.tool_args.get(tool_id):
+                    events.append(build_input_json_delta(block_idx, tc.arguments))
+                    acc.add_tool_args(tool_id, tc.arguments)
                 continue
             if tc.arguments:
                 events.append(build_input_json_delta(block_idx, tc.arguments))
+                acc.add_tool_args(tool_id, tc.arguments)
 
         return events