Merge pull request 'feat: forward tool-result images to the MCP caller in send_message results' (#5) from feat/tool-result-image-passthrough into main

Reviewed-on: #5

docs/pallas.md

@@ -111,7 +111,7 @@ server.py main()
 |---|---|
 | `pallas.server` | CLI entry point, configuration loading, agent lifecycle orchestration, model registration |
 | `pallas.registry` | Starlette app serving `GET /.well-known/mcp/server.json` — builds the agent catalogue from `agents.yaml` + `fastagent.config.yaml` |
-| `pallas.multimodal_server` | `MultimodalAgentMCPServer` — `AgentMCPServer` subclass adding image attachment support and conversation history prompts |
+| `pallas.multimodal_server` | `MultimodalAgentMCPServer` — `AgentMCPServer` subclass adding image attachment support, tool-result image passthrough, and conversation history prompts |
 | `pallas.health` | Two-layer health: startup LLM preflight validation + runtime `get_health` MCP tool with downstream server probing |
 
 ---
@@ -195,6 +195,7 @@ agents:
 | `agents.<name>.depends_on` | no | List of agent names that must start and become ready before this agent |
 | `agents.<name>.max_iterations` | no | Hard cap on agentic-loop turns per `send_message`. Default: `15`. fast-agent returns a partial answer once exceeded |
 | `agents.<name>.loop_repeat_threshold` | no | Halt the loop after this many consecutive identical `(tool, args) → result` rounds. Default: `3`. `0` disables the guard |
+| `agents.<name>.max_result_images` | no | Cap on tool-result images forwarded in the final `send_message` result (most recent kept). Default: `8`. `0` disables image passthrough |
 
 ### `fastagent.config.yaml` Extensions
 
@@ -447,6 +448,14 @@ Each agent's MCP tool accepts:
 
 When `images` is provided, the message is sent as a `PromptMessageExtended` containing both `TextContent` and `ImageContent` parts — the agent's underlying model must support vision.
 
+### Tool-Result Image Passthrough
+
+Images work in both directions. fast-agent's `agent.send()` returns only the final assistant text, so images produced by downstream tools during the agentic loop (playwright screenshots, rommie desktop captures) would otherwise reach the agent's own vision model but never the MCP caller. A per-request `after_tool_call` hook (`pallas.image_passthrough`) collects every `ImageContent` block from the turn's tool results; at end of turn `send_message` returns a `CallToolResult` whose content is the assistant's text block followed by the collected images. Turns that produce no images return the plain string, unchanged from previous releases.
+
+Only the most recent `max_result_images` (default `8`) are forwarded — screenshots are usually taken to show final state, and an unbounded loop of full-desktop captures would balloon the HTTP response. Forwarded images are counted in `pallas_result_images_total`; drops are logged.
+
+Because a lead agent receives a sub-agent's images as ordinary tool-result content, the same hook on the lead's own turn forwards them again — images cascade hop-by-hop up the delegation chain (playwright → dolores → lead → Daedalus) with no extra wiring.
+
 ### Conversation History Prompt
 
 For agents with `instance_scope != "request"`, a `{agent}_history` prompt is registered that returns the full conversation history as FastMCP `Message` objects. This allows clients to retrieve the stored context.
@@ -606,6 +615,7 @@ scrape_configs:
 | `pallas_llm_provider_up` | gauge | `provider` | `1` when the active LLM provider passed its last preflight or runtime re-probe |
 | `pallas_agent_health_status` | gauge | `agent` | Aggregate from the last `get_health`: `1`=ok, `0.5`=degraded, `0`=error |
 | `pallas_agent_loop_aborted_total` | counter | `agent`, `reason` | Agentic loops force-stopped by a runtime guard. `reason` ∈ `repeat` (identical-tool-call loop detected) |
+| `pallas_result_images_total` | counter | `agent` | Tool-result images forwarded to the MCP caller in `send_message` results |
 
 Standard process metrics (RSS, CPU, GC, open FDs) are emitted by `prometheus-client`'s default collectors on the same endpoint.
 
@@ -683,6 +693,7 @@ This avoids the brittle pattern of inferring capabilities from model name substr
 | `pallas.multimodal_server` | `multimodal_server.py` | `MultimodalAgentMCPServer` — extends `AgentMCPServer` with image support, conversation history prompts, bearer token propagation |
 | `pallas.health` | `health.py` | LLM provider preflight validation, downstream MCP server probing, `get_health` tool registration |
 | `pallas.loop_guard` | `loop_guard.py` | Per-request `ToolRunnerHooks` that halt the agentic loop on repeated-identical tool calls |
+| `pallas.image_passthrough` | `image_passthrough.py` | Per-request `ToolRunnerHooks` that collect tool-result images so `send_message` can return them in the final `CallToolResult` |
 | `pallas.log` | `log.py` | JSON log configuration, third-party traceback capture, Rich-TUI-safe handler attachment |
 | `pallas._fastagent_patch` | `_fastagent_patch.py` | Monkey-patches fast-agent at import time: per-request bearer forwarding via `httpx.Auth`, diagnostic trace-capture wrappers around `send_request` / `session.call_tool` / `_execute_on_server` |
 

pallas/image_passthrough.py

@@ -0,0 +1,146 @@
+"""Tool-result image passthrough for ``send_message``.
+
+fast-agent's ``agent.send()`` returns only the final assistant text, so
+``ImageContent`` produced by downstream tools during the agentic loop —
+playwright screenshots, rommie desktop captures — reaches the agent's own
+vision model but never crosses the MCP boundary. The caller (Daedalus, or a
+lead agent driving a sub-agent) sees a text-only ``CallToolResult`` and the
+live stream reduces result images to a ``"+N images"`` preview.
+
+This module fixes that by installing a per-request ``after_tool_call`` hook
+(same composition pattern as ``assistant_stream`` / ``loop_guard``) that
+collects every ``ImageContent`` block from the turn's tool results. At end of
+turn ``multimodal_server.send_message`` appends the collected images to the
+final ``CallToolResult`` via FastMCP's ``ToolResult`` passthrough, after the
+assistant's text block. Daedalus already renders image blocks in final
+results; a lead agent calling a sub-agent receives them as ordinary
+tool-result content, so images cascade hop-by-hop up the delegation chain.
+
+Only the most recent ``max_images`` are forwarded — a long loop that grabs a
+screenshot per iteration would otherwise balloon the HTTP response — and the
+final state is what the screenshots are usually for. Dropped images are
+logged; forwarded ones are counted in ``pallas_result_images_total``.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from fast_agent.agents.tool_runner import ToolRunnerHooks
+from fast_agent.types import PromptMessageExtended
+from fastmcp.tools import ToolResult
+from mcp.types import ImageContent, TextContent
+
+from pallas.assistant_stream import _merge_hooks
+
+logger = logging.getLogger("pallas.image_passthrough")
+
+DEFAULT_MAX_IMAGES = 8
+
+
+class ImageCollector:
+    """Per-request accumulator of tool-result ``ImageContent`` blocks."""
+
+    def __init__(
+        self,
+        *,
+        agent_name: str,
+        conversation_id: str | None,
+        max_images: int = DEFAULT_MAX_IMAGES,
+    ) -> None:
+        self._agent_name = agent_name
+        self._conversation_id = conversation_id
+        self._max_images = max_images
+        self._images: list[ImageContent] = []
+
+    def as_after_tool_call_hook(self):
+        async def _hook(_runner: Any, message: PromptMessageExtended) -> None:
+            try:
+                for result in (message.tool_results or {}).values():
+                    for block in getattr(result, "content", None) or []:
+                        if isinstance(block, ImageContent):
+                            self._images.append(block)
+            except Exception:  # never let collection break a live turn
+                logger.warning(
+                    "image_passthrough collection failed",
+                    exc_info=True,
+                    extra={
+                        "agent": self._agent_name,
+                        "conversation_id": self._conversation_id,
+                    },
+                )
+
+        return _hook
+
+    def collected(self) -> list[ImageContent]:
+        """Return the images to forward — the most recent ``max_images``.
+
+        When the loop produced more than the cap, the oldest are dropped:
+        screenshots are usually taken to show the *final* state, and the
+        intermediate ones were already summarized on the live stream.
+        """
+        if len(self._images) <= self._max_images:
+            return list(self._images)
+        dropped = len(self._images) - self._max_images
+        logger.warning(
+            "image_passthrough dropping oldest tool-result images",
+            extra={
+                "agent": self._agent_name,
+                "conversation_id": self._conversation_id,
+                "collected": len(self._images),
+                "forwarded": self._max_images,
+                "dropped": dropped,
+            },
+        )
+        return self._images[-self._max_images :]
+
+
+def build_result(text: str, images: list[ImageContent]) -> str | ToolResult:
+    """Assemble ``send_message``'s return value.
+
+    With no images the plain string is returned so the wire shape is
+    byte-identical to previous releases. With images, an explicit FastMCP
+    ``ToolResult`` carries ``[text, *images]`` content blocks — FastMCP
+    passes it through verbatim, bypassing return-annotation conversion.
+    """
+    if not images:
+        return text
+    content: list[TextContent | ImageContent] = []
+    if text:
+        content.append(TextContent(type="text", text=text))
+    content.extend(images)
+    return ToolResult(content=content)
+
+
+def install_for_request(
+    agent: Any,
+    *,
+    agent_name: str,
+    conversation_id: str | None,
+    max_images: int = DEFAULT_MAX_IMAGES,
+):
+    """Install the image collector on a request-scoped agent instance.
+
+    Returns ``(collector, restore)``. Call ``restore`` in a ``finally`` like
+    the other per-request hooks. A non-positive ``max_images`` disables
+    passthrough entirely (returns ``(None, no-op)``).
+    """
+    if max_images is None or max_images < 1:
+        return None, lambda: None
+
+    collector = ImageCollector(
+        agent_name=agent_name,
+        conversation_id=conversation_id,
+        max_images=max_images,
+    )
+
+    extra = ToolRunnerHooks(after_tool_call=collector.as_after_tool_call_hook())
+
+    previous = getattr(agent, "tool_runner_hooks", None)
+    agent.tool_runner_hooks = _merge_hooks(previous, extra)
+
+    def restore() -> None:
+        agent.tool_runner_hooks = previous
+
+    return collector, restore

pallas/metrics.py

@@ -138,6 +138,13 @@ agent_loop_aborted_total = Counter(
     registry=REGISTRY,
 )
 
+result_images_total = Counter(
+    "pallas_result_images_total",
+    "Tool-result images forwarded to the MCP caller in send_message results",
+    labelnames=["agent"],
+    registry=REGISTRY,
+)
+
 
 # ── Helpers ──────────────────────────────────────────────────────────────────
 
@@ -147,6 +154,11 @@ def record_loop_abort(agent: str, reason: str) -> None:
     agent_loop_aborted_total.labels(agent=agent, reason=reason).inc()
 
 
+def record_result_images(agent: str, count: int) -> None:
+    """Record tool-result images forwarded in a send_message result."""
+    result_images_total.labels(agent=agent).inc(count)
+
+
 def set_agent_info(agents: dict[str, dict]) -> None:
     """Record the deployment's configured agents (called once at startup)."""
     for name, agent in agents.items():

pallas/multimodal_server.py

@@ -9,7 +9,10 @@ Overrides register_agent_tools to:
     so callers own conversation state and seed it on every turn,
   * accept an optional ``conversation_id`` string that is recorded in
     structured logs and progress notification metadata for end-to-end
-    trace correlation.
+    trace correlation,
+  * append images produced by downstream tools during the turn (playwright
+    screenshots, rommie desktop captures) to the final ``CallToolResult``
+    so they reach the MCP caller — see ``pallas.image_passthrough``.
 
 Drop-in replacement for AgentMCPServer.  When combined with
 ``instance_scope="request"`` (the Pallas default), this gives a fully
@@ -29,11 +32,17 @@ from fast_agent.mcp.server import AgentMCPServer
 from fast_agent.types import PromptMessageExtended, RequestParams
 
 from pallas.assistant_stream import install_for_request as _install_assistant_stream
+from pallas.image_passthrough import (
+    DEFAULT_MAX_IMAGES,
+    build_result as _build_image_result,
+    install_for_request as _install_image_passthrough,
+)
 from pallas.loop_guard import install_for_request as _install_loop_guard
 from pallas.progress import EnrichedMCPToolProgressManager
 from pallas import metrics as _pallas_metrics
 from fastmcp import Context as MCPContext
 from fastmcp.prompts import Message
+from fastmcp.tools import ToolResult
 from mcp.types import ImageContent, TextContent
 from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
 from starlette.responses import JSONResponse, Response
@@ -188,7 +197,7 @@ class MultimodalAgentMCPServer(AgentMCPServer):
             images: list[dict] | None = None,
             history: list[dict] | None = None,
             conversation_id: str | None = None,
-        ) -> str:
+        ) -> str | ToolResult:
             """Send a single turn to the agent.
 
             Parameters
@@ -209,6 +218,12 @@ class MultimodalAgentMCPServer(AgentMCPServer):
             conversation_id:
                 Optional opaque identifier, logged for trace correlation.
                 Pallas does not interpret it.
+
+            Returns the assistant's final text. When downstream tools
+            produced images this turn (screenshots etc.), returns a
+            ``ToolResult`` whose content is the text block followed by the
+            most recent ``max_result_images`` image blocks, so they reach
+            the MCP caller instead of dying inside ``message_history``.
             """
             report_progress = self._build_progress_reporter(ctx)
             request_params = RequestParams(
@@ -245,8 +260,20 @@ class MultimodalAgentMCPServer(AgentMCPServer):
                 conversation_id=conversation_id,
                 threshold=self._request_limits.get("loop_repeat_threshold", 3),
             )
+            # Collect tool-result images (screenshots) so the final
+            # CallToolResult can carry them to the caller — fast-agent's
+            # send() return value is text-only.
+            image_collector, restore_images = _install_image_passthrough(
+                agent,
+                agent_name=agent_name,
+                conversation_id=conversation_id,
+                max_images=self._request_limits.get(
+                    "max_result_images", DEFAULT_MAX_IMAGES
+                ),
+            )
 
             def restore_hooks() -> None:
+                restore_images()
                 restore_guard()
                 restore_stream()
             try:
@@ -314,7 +341,25 @@ class MultimodalAgentMCPServer(AgentMCPServer):
                     return await execute_send()
 
                 try:
-                    return await asyncio.wait_for(_dispatch(), timeout=turn_timeout)
+                    response = await asyncio.wait_for(
+                        _dispatch(), timeout=turn_timeout
+                    )
+                    images = (
+                        image_collector.collected() if image_collector else []
+                    )
+                    if images:
+                        _pallas_metrics.record_result_images(
+                            agent_name, len(images)
+                        )
+                        logger.debug(
+                            f"Forwarding {len(images)} tool-result image(s) "
+                            f"from agent '{agent_name}'",
+                            name="result_images_forwarded",
+                            agent=agent_name,
+                            image_count=len(images),
+                            conversation_id=conversation_id,
+                        )
+                    return _build_image_result(response, images)
                 except asyncio.TimeoutError:
                     logger.warning(
                         f"Agent '{agent_name}' turn exceeded {turn_timeout}s wall-clock limit",

pallas/server.py

@@ -67,6 +67,7 @@ def _build_agents_table(config: dict) -> dict[str, dict]:
             "streaming_timeout": agent.get("streaming_timeout"),
             "turn_timeout": agent.get("turn_timeout"),
             "loop_repeat_threshold": agent.get("loop_repeat_threshold"),
+            "max_result_images": agent.get("max_result_images"),
         }
         for name, agent in config["agents"].items()
     }
@@ -271,6 +272,7 @@ async def _start_agent(name: str, agents: dict[str, dict]) -> None:
                 "streaming_timeout",
                 "turn_timeout",
                 "loop_repeat_threshold",
+                "max_result_images",
             )
             if entry.get(k) is not None
         }

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pallas-mcp"
-version = "0.6.0"
+version = "0.7.0"
 description = "FastAgent MCP Bridge — generic runtime for serving FastAgent agents over StreamableHTTP"
 requires-python = ">=3.13.5"
 dependencies = [

tests/test_image_passthrough.py

@@ -0,0 +1,138 @@
+"""Tests for ``pallas.image_passthrough``.
+
+Drives the ``after_tool_call`` hook with handcrafted
+``PromptMessageExtended`` objects and asserts the collection, cap, and
+result-assembly behaviour. No fast-agent runtime is involved — the hook is
+a pure async function. Uses ``asyncio.run`` directly to match the
+convention in the other test modules (pallas has no pytest-asyncio
+dependency).
+"""
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+
+from fast_agent.types import PromptMessageExtended
+from fastmcp.tools import ToolResult
+from mcp.types import CallToolResult, ImageContent, TextContent
+
+from pallas.image_passthrough import (
+    ImageCollector,
+    build_result,
+    install_for_request,
+)
+
+
+def _run(coro):
+    return asyncio.run(coro)
+
+
+def _image(tag: str) -> ImageContent:
+    return ImageContent(type="image", data=f"b64-{tag}", mimeType="image/png")
+
+
+def _result_message(*blocks, call_id: str = "toolu_1") -> PromptMessageExtended:
+    return PromptMessageExtended(
+        role="user",
+        content=[],
+        tool_results={call_id: CallToolResult(content=list(blocks))},
+    )
+
+
+def _collector(max_images: int = 8) -> ImageCollector:
+    return ImageCollector(
+        agent_name="dolores", conversation_id="c1", max_images=max_images
+    )
+
+
+def test_collects_images_and_ignores_text():
+    collector = _collector()
+    hook = collector.as_after_tool_call_hook()
+
+    async def go():
+        await hook(None, _result_message(TextContent(type="text", text="snapshot")))
+        await hook(None, _result_message(_image("a")))
+        await hook(
+            None,
+            _result_message(TextContent(type="text", text="took it"), _image("b")),
+        )
+
+    _run(go())
+    assert [i.data for i in collector.collected()] == ["b64-a", "b64-b"]
+
+
+def test_collects_across_multiple_results_in_one_message():
+    collector = _collector()
+    hook = collector.as_after_tool_call_hook()
+    message = PromptMessageExtended(
+        role="user",
+        content=[],
+        tool_results={
+            "toolu_1": CallToolResult(content=[_image("a")]),
+            "toolu_2": CallToolResult(content=[_image("b")]),
+        },
+    )
+
+    _run(hook(None, message))
+    assert len(collector.collected()) == 2
+
+
+def test_cap_keeps_most_recent():
+    collector = _collector(max_images=2)
+    hook = collector.as_after_tool_call_hook()
+
+    async def go():
+        for tag in ("a", "b", "c", "d"):
+            await hook(None, _result_message(_image(tag)))
+
+    _run(go())
+    assert [i.data for i in collector.collected()] == ["b64-c", "b64-d"]
+
+
+def test_empty_tool_results_are_harmless():
+    collector = _collector()
+    hook = collector.as_after_tool_call_hook()
+    message = PromptMessageExtended(role="user", content=[], tool_results=None)
+
+    _run(hook(None, message))
+    assert collector.collected() == []
+
+
+def test_build_result_plain_text_when_no_images():
+    assert build_result("answer", []) == "answer"
+
+
+def test_build_result_wraps_text_and_images():
+    result = build_result("answer", [_image("a"), _image("b")])
+    assert isinstance(result, ToolResult)
+    assert [b.type for b in result.content] == ["text", "image", "image"]
+    assert result.content[0].text == "answer"
+    assert result.content[1].data == "b64-a"
+
+
+def test_build_result_omits_empty_text_block():
+    result = build_result("", [_image("a")])
+    assert isinstance(result, ToolResult)
+    assert [b.type for b in result.content] == ["image"]
+
+
+def test_install_disabled_with_nonpositive_max():
+    agent = SimpleNamespace(tool_runner_hooks="sentinel")
+    collector, restore = install_for_request(
+        agent, agent_name="a", conversation_id=None, max_images=0
+    )
+    assert collector is None
+    assert agent.tool_runner_hooks == "sentinel"  # untouched
+    restore()  # no-op, must not raise
+
+
+def test_install_merges_and_restores():
+    agent = SimpleNamespace(tool_runner_hooks=None)
+    collector, restore = install_for_request(
+        agent, agent_name="a", conversation_id=None
+    )
+    assert collector is not None
+    assert agent.tool_runner_hooks is not None
+    assert agent.tool_runner_hooks.after_tool_call is not None
+    restore()
+    assert agent.tool_runner_hooks is None