hud-python 0.5.28__tar.gz → 0.5.30__tar.gz

{hud_python-0.5.28 → hud_python-0.5.30}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hud-python
-Version: 0.5.28
+Version: 0.5.30
 Summary: SDK for the HUD platform.
 Project-URL: Homepage, https://github.com/hud-evals/hud-python
 Project-URL: Bug Tracker, https://github.com/hud-evals/hud-python/issues
@@ -34,11 +34,12 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Python: <3.13,>=3.11
+Requires-Dist: a2a-sdk>=0.3.24
 Requires-Dist: blessed>=1.20.0
-Requires-Dist: fastmcp==2.13.3
+Requires-Dist: fastmcp==3.0.2
 Requires-Dist: httpx<1,>=0.23.0
-Requires-Dist: mcp<1.23,>1.21.1
-Requires-Dist: openai>=2.8.1
+Requires-Dist: mcp<2.0,>=1.24.0
+Requires-Dist: openai>=2.26.0
 Requires-Dist: packaging>=21.0
 Requires-Dist: prompt-toolkit==3.0.51
 Requires-Dist: pydantic-settings<3,>=2.2
@@ -51,7 +52,6 @@ Requires-Dist: typer>=0.9.0
 Requires-Dist: watchfiles>=0.21.0
 Provides-Extra: agent
 Requires-Dist: anthropic>=0.78.0; extra == 'agent'
-Requires-Dist: datasets>=2.14.0; extra == 'agent'
 Requires-Dist: google-genai; extra == 'agent'
 Requires-Dist: langchain>=1.1.0; extra == 'agent'
 Requires-Dist: mcp-use==1.5.0; extra == 'agent'
@@ -60,7 +60,6 @@ Requires-Dist: pillow>=11.1.0; extra == 'agent'
 Requires-Dist: tornado>=6.5.2; extra == 'agent'
 Provides-Extra: agents
 Requires-Dist: anthropic>=0.78.0; extra == 'agents'
-Requires-Dist: datasets>=2.14.0; extra == 'agents'
 Requires-Dist: google-genai; extra == 'agents'
 Requires-Dist: langchain>=1.1.0; extra == 'agents'
 Requires-Dist: mcp-use==1.5.0; extra == 'agents'
@@ -71,7 +70,6 @@ Provides-Extra: bedrock
 Requires-Dist: anthropic[bedrock]>=0.78.0; extra == 'bedrock'
 Provides-Extra: dev
 Requires-Dist: anthropic>=0.78.0; extra == 'dev'
-Requires-Dist: datasets>=2.14.0; extra == 'dev'
 Requires-Dist: dotenv>=0.9.9; extra == 'dev'
 Requires-Dist: google-adk; extra == 'dev'
 Requires-Dist: google-genai; extra == 'dev'

{hud_python-0.5.28 → hud_python-0.5.30}/hud/__init__.py

@@ -12,6 +12,7 @@ from . import patches as _patches  # noqa: F401
 from .environment import Environment
 from .eval import EvalContext
 from .eval import run_eval as eval
+from .services import Chat
 from .telemetry.instrument import instrument
 
 
@@ -30,6 +31,7 @@ def trace(*args: object, **kwargs: object) -> EvalContext:
 
 
 __all__ = [
+    "Chat",
     "Environment",
     "EvalContext",
     "eval",

{hud_python-0.5.28 → hud_python-0.5.30}/hud/agents/base.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import asyncio
 import json
 import logging
+import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, ClassVar, Literal
@@ -12,7 +13,8 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal
 import mcp.types as types
 
 from hud.tools.native_types import NativeToolSpec
-from hud.types import AgentResponse, AgentType, BaseAgentConfig, MCPToolCall, MCPToolResult, Trace
+from hud.tools.types import Citation
+from hud.types import AgentType, BaseAgentConfig, InferenceResult, MCPToolCall, MCPToolResult, Trace
 from hud.utils.hud_console import HUDConsole
 
 from .types import BaseCreateParams
@@ -417,21 +419,25 @@ class MCPAgent(ABC):
             await self._initialize_from_ctx(ctx)
 
         try:
-            # Build initial context - optionally append setup tool output
-            # Check ctx first (task-level override), then fall back to agent config
-            append_setup = getattr(ctx, "append_setup_output", False) or getattr(
-                self.config, "append_setup_output", False
-            )
-            initial_prompt = ctx.prompt
-            if append_setup:
-                setup_output = getattr(ctx, "setup_output", None)
-                if setup_output:
-                    initial_prompt = f"{initial_prompt}\n\n{setup_output}"
+            # Build initial context
+            conversation: list[dict[str, str]] | None = getattr(ctx, "conversation", None)
 
-            # Build initial blocks (text prompt + optional screenshot)
-            initial_blocks = text_to_blocks(initial_prompt)
+            if conversation:
+                # Multi-turn: build alternating role messages
+                initial_messages = await self._build_conversation_messages(conversation)
+            else:
+                # Single-turn: single user message from prompt
+                append_setup = getattr(ctx, "append_setup_output", False) or getattr(
+                    self.config, "append_setup_output", False
+                )
+                initial_prompt = ctx.prompt
+                if append_setup:
+                    setup_output = getattr(ctx, "setup_output", None)
+                    if setup_output:
+                        initial_prompt = f"{initial_prompt}\n\n{setup_output}"
+                initial_messages = await self.format_message(initial_prompt)
 
-            result = await self._run_context(initial_blocks, max_steps=max_steps)
+            result = await self._run_context(initial_messages, max_steps=max_steps)
 
             # Propagate error state to context for platform visibility
             if result.isError and hasattr(ctx, "error"):
@@ -440,7 +446,15 @@ class MCPAgent(ABC):
 
             # Submit final answer to context (only if scenario is running)
             if result.content and ctx.has_scenario:
-                await ctx.submit(result.content)
+                if result.citations:
+                    await ctx.submit(
+                        {
+                            "content": result.content,
+                            "citations": result.citations,
+                        }
+                    )
+                else:
+                    await ctx.submit(result.content)
 
             return result
 
@@ -460,30 +474,48 @@ class MCPAgent(ABC):
             # Cleanup auto-created resources
             await self._cleanup()
 
-    async def _run_context(
-        self, context: list[types.ContentBlock], *, max_steps: int = 10
-    ) -> Trace:
+    def _map_role(self, role: str) -> str:
+        """Map a canonical role name to the provider-specific role.
+
+        Override in subclasses where the provider uses different role names.
+        Default passes through (works for OpenAI and Claude which use "assistant").
         """
-        Run the agent with the given context messages. This is the core agent loop.
+        return role
+
+    async def _build_conversation_messages(self, conversation: list[dict[str, str]]) -> list[Any]:
+        """Build provider-formatted messages from a conversation history."""
+        result: list[Any] = []
+        for msg in conversation:
+            role = self._map_role(msg.get("role", "user"))
+            content = msg.get("content", "")
+            formatted = await self.format_message(content)
+            for fm in formatted:
+                if isinstance(fm, dict):
+                    fm["role"] = role
+                elif hasattr(fm, "role"):
+                    fm.role = role  # type: ignore[attr-defined]
+            result.extend(formatted)
+        return result
+
+    async def _run_context(self, initial_messages: list[Any], *, max_steps: int = 10) -> Trace:
+        """
+        Run the agent with pre-built messages. This is the core agent loop.
 
         Args:
-            context: The context to complete
+            initial_messages: Provider-formatted messages (from format_message or conversation)
             max_steps: Maximum number of steps (-1 for infinite)
 
         Returns:
             Trace with reward, done, content fields and trace steps
         """
-        final_response = None
+        final_response: InferenceResult | None = None
         error = None
 
         messages: list[Any] = []
 
         try:
-            # Start with system messages
             messages = await self.get_system_messages()
-
-            # Add initial context
-            messages.extend(await self.format_message(context))
+            messages.extend(initial_messages)
             self.console.debug(f"Messages: {messages}")
 
             step_count = 0
@@ -513,6 +545,19 @@ class MCPAgent(ABC):
                             except Exception as e:
                                 self.console.warning_log(f"Auto-respond failed: {e}")
                         if decision == "STOP":
+                            if (
+                                getattr(self.ctx, "scenario_enable_citations", False)
+                                and not response.citations
+                            ):
+                                recovered = self._recover_citations_from_content(response)
+                                if recovered:
+                                    self.console.info_log(
+                                        "Recovered citations from JSON answer payload"
+                                    )
+                                else:
+                                    self.console.warning_log(
+                                        "Citations required by scenario but missing in final response"  # noqa: E501
+                                    )
                             self.console.debug("Stopping execution")
                             final_response = response
                             break
@@ -564,7 +609,6 @@ class MCPAgent(ABC):
         else:
             is_error = False
 
-        # Ensure all parameters are the correct type
         # Use ctx.reward if already set (e.g., from scenario evaluate), otherwise 0.0
         # Note: For v4 tasks with evaluate_tool, reward is set in __aexit__ after this returns,
         # so callers should prefer ctx.reward over Trace.reward for the final result.
@@ -574,17 +618,81 @@ class MCPAgent(ABC):
             if ctx_reward is not None:
                 reward = ctx_reward
 
-        trace_params = {
-            "reward": reward,
-            "done": True,
-            "messages": messages,
-            "content": final_response.content if final_response else error,
-            "isError": is_error,
-            "info": {"error": error} if error else {},
-        }
-        trace_result = Trace(**trace_params)
+        return Trace(
+            reward=reward,
+            done=True,
+            messages=messages,
+            content=final_response.content if final_response else error,
+            isError=is_error,
+            citations=final_response.citations if final_response else [],
+            info={"error": error} if error else {},
+        )
 
-        return trace_result
+    def _recover_citations_from_content(self, response: InferenceResult) -> bool:
+        """Try to extract citations from model content when native citations are missing.
+
+        Handles two cases: raw JSON content and fenced ```json blocks.
+        """
+        raw = response.content or ""
+        if not raw:
+            return False
+
+        # Try raw content first, then try extracting from fenced block.
+        for text in dict.fromkeys([raw, self._extract_fenced_json(raw) or ""]):
+            if not text:
+                continue
+            try:
+                parsed = json.loads(text)
+            except (json.JSONDecodeError, TypeError):
+                continue
+            if not isinstance(parsed, dict):
+                continue
+
+            raw_citations = parsed.get("citations")
+            if not isinstance(raw_citations, list) or not raw_citations:
+                continue
+
+            normalized: list[Citation] = [
+                c
+                for cit in raw_citations
+                if isinstance(cit, dict) and (c := self._normalize_citation(cit)) is not None
+            ]
+            if not normalized:
+                continue
+
+            content = parsed.get("content")
+            if isinstance(content, str) and content.strip():
+                response.content = content
+            response.citations = [c.model_dump(exclude={"provider_data"}) for c in normalized]
+            return True
+
+        return False
+
+    @staticmethod
+    def _extract_fenced_json(value: str) -> str | None:
+        """Extract JSON content from a fenced code block."""
+        match = re.search(r"```(?:json)?\s*\n(.*?)```", value, re.DOTALL)
+        return match.group(1).strip() if match else None
+
+    @staticmethod
+    def _normalize_citation(cit: dict[str, Any]) -> Citation | None:
+        """Normalize a citation dict to canonical Citation shape.
+
+        Maps common key aliases to canonical names and validates via Citation.
+        Returns None only if construction fails (e.g. extra-forbid violation).
+        """
+        source = cit.get("source") or cit.get("document") or ""
+        try:
+            return Citation(
+                type=cit.get("type", "document_citation"),
+                text=cit.get("text") or cit.get("cited_text", ""),
+                source=str(source),
+                title=cit.get("title") or cit.get("document_title"),
+                start_index=cit.get("start_index", cit.get("start_char_index")),
+                end_index=cit.get("end_index", cit.get("end_char_index")),
+            )
+        except Exception:
+            return None
 
     async def call_tools(
         self, tool_call: MCPToolCall | list[MCPToolCall] | None = None
@@ -629,7 +737,7 @@ class MCPAgent(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def get_response(self, messages: list[Any]) -> AgentResponse:
+    async def get_response(self, messages: list[Any]) -> InferenceResult:
         """
         Get response from the model including any tool calls.
 
@@ -756,8 +864,10 @@ def _parse_spec_dict(spec_dict: dict[str, Any]) -> NativeToolSpec | None:
     """Parse a dict (from MCP meta) into a NativeToolSpec."""
     if not spec_dict:
         return None
-    known_fields = {"api_type", "api_name", "beta", "hosted", "role", "supported_models"}
+    known_fields = {"api_type", "api_name", "beta", "hosted", "role", "supported_models", "extra"}
     extra = {k: v for k, v in spec_dict.items() if k not in known_fields}
+    if isinstance(spec_dict.get("extra"), dict):
+        extra.update(spec_dict["extra"])
     supported_models_raw = spec_dict.get("supported_models")
     supported_models: tuple[str, ...] | None = None
     if supported_models_raw:

{hud_python-0.5.28 → hud_python-0.5.30}/hud/agents/claude.py

@@ -17,6 +17,7 @@ from anthropic.types.beta import (
     BetaContentBlockParam,
     BetaImageBlockParam,
     BetaMessageParam,
+    BetaPlainTextSourceParam,
     BetaRequestDocumentBlockParam,
     BetaTextBlockParam,
     BetaToolBash20250124Param,
@@ -31,7 +32,7 @@ from anthropic.types.beta import (
 from hud.settings import settings
 from hud.tools.computer.settings import computer_settings
 from hud.tools.native_types import NativeToolSpec
-from hud.types import AgentResponse, AgentType, BaseAgentConfig, MCPToolCall, MCPToolResult
+from hud.types import AgentType, BaseAgentConfig, InferenceResult, MCPToolCall, MCPToolResult
 from hud.utils.hud_console import HUDConsole
 from hud.utils.types import with_signature
 
@@ -155,9 +156,11 @@ class ClaudeAgent(MCPAgent):
 
         # these will be initialized in _convert_tools_for_claude
         self.has_computer_tool = False
-        self.tool_mapping = {}
-        self.claude_tools = []
-        self._required_betas = set()
+        self.tool_mapping: dict[str, str] = {}
+        self.claude_tools: list[BetaToolUnionParam] = []
+        self._required_betas: set[str] = set()
+        self._tool_search_threshold: int | None = None
+        self._gated_screenshot_tools: set[str] = set()
 
     def _on_tools_ready(self) -> None:
         """Build Claude-specific tool mappings after tools are discovered."""
@@ -167,6 +170,67 @@ class ClaudeAgent(MCPAgent):
         """No system messages for Claude because applied in get_response"""
         return []
 
+    def _result_from_response_blocks(self, response_blocks: list[Any]) -> InferenceResult:
+        """Extract text/tool calls/citations from Anthropic response blocks."""
+        result = InferenceResult(content="", tool_calls=[], done=True)
+        text_content = ""
+        thinking_content = ""
+        citations: list[dict[str, Any]] = []
+
+        for block in response_blocks:
+            block_type = getattr(block, "type", None)
+            if block_type == "tool_use":
+                block_input = getattr(block, "input", {})
+                mcp_name = self.tool_mapping.get(
+                    getattr(block, "name", ""),
+                    getattr(block, "name", ""),
+                )
+                arguments = block_input if isinstance(block_input, dict) else block_input.__dict__
+                if mcp_name in self._gated_screenshot_tools:
+                    arguments = {**arguments, "take_screenshot_on_click": False}
+                    logger.debug(
+                        "Injected take_screenshot_on_click=False for gated tool %s", mcp_name
+                    )
+                tool_call = MCPToolCall(
+                    id=getattr(block, "id", ""),
+                    name=mcp_name,
+                    arguments=arguments,
+                )
+                result.tool_calls.append(tool_call)
+                result.done = False
+            elif block_type == "text":
+                text = getattr(block, "text", "") or ""
+                text_content += text
+                block_citations = getattr(block, "citations", None) or []
+                for cit in block_citations:
+                    cit_dict = {
+                        "type": "document_citation",
+                        "text": getattr(cit, "cited_text", "") or "",
+                        "source": (
+                            str(idx)
+                            if (idx := getattr(cit, "document_index", None)) is not None
+                            else getattr(cit, "document_title", "") or ""
+                        ),
+                        "title": getattr(cit, "document_title", None),
+                        "start_index": getattr(cit, "start_char_index", None),
+                        "end_index": getattr(cit, "end_char_index", None),
+                    }
+                    normalized = self._normalize_citation(cit_dict)
+                    if normalized is not None:
+                        citations.append(normalized.model_dump(exclude={"provider_data"}))
+            elif block_type == "thinking":
+                thinking = getattr(block, "thinking", "") or ""
+                if thinking:
+                    if thinking_content:
+                        thinking_content += "\n"
+                    thinking_content += thinking
+
+        result.content = text_content
+        result.citations = citations
+        if thinking_content:
+            result.reasoning = thinking_content
+        return result
+
     async def format_blocks(self, blocks: list[types.ContentBlock]) -> list[BetaMessageParam]:
         """Format messages for Claude."""
         # Convert MCP content types to Anthropic content types
@@ -234,7 +298,7 @@ class ClaudeAgent(MCPAgent):
             content=[text_to_content_block(retry_text)],
         )
 
-    async def get_response(self, messages: list[BetaMessageParam]) -> AgentResponse:
+    async def get_response(self, messages: list[BetaMessageParam]) -> InferenceResult:
         """Get response from Claude including any tool calls."""
         messages_cached = self._add_prompt_caching(messages)
         # betas to use - collected during tool conversion based on native specs
@@ -242,6 +306,24 @@ class ClaudeAgent(MCPAgent):
         # anthropic-beta header which the API rejects.
         betas: list[str] | Omit = list(self._required_betas) if self._required_betas else Omit()
 
+        effective_tools: list[BetaToolUnionParam] = list(self.claude_tools)
+        if self._tool_search_threshold is not None:
+            generic_count = sum(
+                1 for t in effective_tools if isinstance(t, dict) and "input_schema" in t
+            )
+            if generic_count > self._tool_search_threshold:
+                logger.debug(
+                    "tool_search: %d generic tools > threshold %d, applying defer_loading",
+                    generic_count,
+                    self._tool_search_threshold,
+                )
+                effective_tools = [
+                    {**t, "defer_loading": True}
+                    if isinstance(t, dict) and "input_schema" in t
+                    else t
+                    for t in effective_tools
+                ]
+
         # Bedrock doesn't support .stream() - use create(stream=True) instead
         if isinstance(self.anthropic_client, AsyncAnthropicBedrock):
             try:
@@ -250,7 +332,7 @@ class ClaudeAgent(MCPAgent):
                     system=self.system_prompt if self.system_prompt is not None else Omit(),
                     max_tokens=self.max_tokens,
                     messages=messages_cached,
-                    tools=self.claude_tools,
+                    tools=effective_tools,
                     tool_choice={"type": "auto", "disable_parallel_tool_use": True},
                     betas=betas,
                 )
@@ -271,7 +353,7 @@ class ClaudeAgent(MCPAgent):
                         system=self.system_prompt if self.system_prompt is not None else Omit(),
                         max_tokens=self.max_tokens,
                         messages=messages_cached,
-                        tools=self.claude_tools,
+                        tools=effective_tools,
                         tool_choice={"type": "auto", "disable_parallel_tool_use": True},
                         betas=betas,
                     ) as stream:
@@ -315,34 +397,7 @@ class ClaudeAgent(MCPAgent):
                 raise ValueError("Claude response missing after stream retries")
 
         # Process response
-        result = AgentResponse(content="", tool_calls=[], done=True)
-
-        # Extract text content and reasoning
-        text_content = ""
-        thinking_content = ""
-
-        for block in response.content:
-            if block.type == "tool_use":
-                tool_call = MCPToolCall(
-                    id=block.id,
-                    # look up name in tool_mapping if available, otherwise use block name
-                    name=self.tool_mapping.get(block.name, block.name),
-                    arguments=block.input
-                    if isinstance(block.input, dict)
-                    else block.input.__dict__,
-                )
-                result.tool_calls.append(tool_call)
-                result.done = False
-            elif block.type == "text":
-                text_content += block.text
-            elif hasattr(block, "type") and block.type == "thinking":
-                if thinking_content:
-                    thinking_content += "\n"
-                thinking_content += block.thinking
-
-        result.content = text_content
-        if thinking_content:
-            result.reasoning = thinking_content
+        result = self._result_from_response_blocks(list(response.content))
 
         return result
 
@@ -353,23 +408,28 @@ class ClaudeAgent(MCPAgent):
 
         Handles EmbeddedResource (PDFs), images, and text content.
         """
+        citations_enabled = bool(
+            getattr(self.ctx, "scenario_enable_citations", False) if self.ctx else False
+        )
+
         # Process each tool result
-        user_content = []
+        user_content: list[BetaToolResultBlockParam | BetaRequestDocumentBlockParam] = []
 
         for tool_call, result in zip(tool_calls, tool_results, strict=True):
-            # Extract Claude-specific metadata from extra fields
             tool_use_id = tool_call.id
             if not tool_use_id:
                 self.hud_console.warning(f"No tool_use_id found for {tool_call.name}")
                 continue
 
-            # Convert MCP tool results to Claude format
+            # Blocks placed inside the tool_result (text, images)
             claude_blocks: list[
                 BetaTextBlockParam | BetaImageBlockParam | BetaRequestDocumentBlockParam
             ] = []
+            # Citable document blocks placed as siblings after the tool_result
+            # so Claude's citation system indexes them properly.
+            sibling_docs: list[BetaRequestDocumentBlockParam] = []
 
             if result.isError:
-                # Extract error message from content
                 error_msg = "Tool execution failed"
                 for content in result.content:
                     if isinstance(content, types.TextContent):
@@ -377,27 +437,37 @@ class ClaudeAgent(MCPAgent):
                         break
                 claude_blocks.append(text_to_content_block(f"Error: {error_msg}"))
             else:
-                # Process success content
                 for content in result.content:
                     if isinstance(content, types.TextContent):
                         claude_blocks.append(text_to_content_block(content.text))
+                        if citations_enabled:
+                            sibling_docs.append(
+                                text_document_block(content.text, title=tool_call.name)
+                            )
                     elif isinstance(content, types.ImageContent):
                         claude_blocks.append(base64_to_content_block(content.data))
                     elif isinstance(content, types.EmbeddedResource):
-                        # Handle embedded resources (PDFs)
                         resource = content.resource
                         if (
                             isinstance(resource, types.BlobResourceContents)
                             and resource.mimeType == "application/pdf"
                         ):
                             claude_blocks.append(
-                                document_to_content_block(base64_data=resource.blob)
+                                document_to_content_block(
+                                    base64_data=resource.blob,
+                                )
                             )
+                            if citations_enabled:
+                                sibling_docs.append(
+                                    document_to_content_block(
+                                        base64_data=resource.blob,
+                                        enable_citations=True,
+                                    )
+                                )
 
-            # Add tool result
             user_content.append(tool_use_content_block(tool_use_id, claude_blocks))
+            user_content.extend(sibling_docs)
 
-        # Return as a user message containing all tool results
         return [
             BetaMessageParam(
                 role="user",
@@ -418,12 +488,28 @@ class ClaudeAgent(MCPAgent):
         self.tool_mapping: dict[str, str] = {}
         self.claude_tools: list[BetaToolUnionParam] = []
         self._required_betas: set[str] = set()
+        self._tool_search_threshold = None
+        self._gated_screenshot_tools: set[str] = set()
 
         categorized = self._categorized_tools
 
-        # Log skipped hosted tools (Claude doesn't support hosted tools currently)
-        for tool, _spec in categorized.hosted:
-            logger.debug("Skipping hosted tool %s for Claude", tool.name)
+        # Process hosted tools
+        for tool, spec in categorized.hosted:
+            if not spec.api_type:
+                logger.debug("Skipping hosted tool %s: no api_type", tool.name)
+                continue
+            tool_def: dict[str, Any] = {
+                "type": spec.api_type,
+                "name": spec.api_name or tool.name,
+            }
+            api_extra = {k: v for k, v in spec.extra.items() if k != "threshold"}
+            tool_def.update(api_extra)
+            if spec.beta:
+                self._required_betas.add(spec.beta)
+            if "threshold" in spec.extra:
+                self._tool_search_threshold = spec.extra["threshold"]
+            self.claude_tools.append(tool_def)  # type: ignore[arg-type]
+            logger.debug("Added hosted tool %s (%s) for Claude", tool.name, spec.api_type)
 
         # Process native tools
         for tool, spec in categorized.native:
@@ -437,6 +523,9 @@ class ClaudeAgent(MCPAgent):
 
             if spec.api_type and spec.api_type.startswith("computer"):
                 self.has_computer_tool = True
+            if spec.api_type == "computer_20251124":
+                self._gated_screenshot_tools.add(tool.name)
+                logger.debug("Screenshot gating enabled for tool %s (computer_20251124)", tool.name)
 
         # Process generic tools
         for tool in categorized.generic:
@@ -611,9 +700,27 @@ def text_to_content_block(text: str) -> BetaTextBlockParam:
     return {"type": "text", "text": text}
 
 
-def document_to_content_block(base64_data: str) -> BetaRequestDocumentBlockParam:
+def text_document_block(text: str, *, title: str | None = None) -> BetaRequestDocumentBlockParam:
+    """Wrap plain text as a citable document block."""
+    block = BetaRequestDocumentBlockParam(
+        type="document",
+        source=BetaPlainTextSourceParam(
+            type="text",
+            media_type="text/plain",
+            data=text,
+        ),
+        citations={"enabled": True},
+    )
+    if title:
+        block["title"] = title
+    return block
+
+
+def document_to_content_block(
+    base64_data: str, *, enable_citations: bool = False
+) -> BetaRequestDocumentBlockParam:
     """Convert base64 PDF to Claude document content block."""
-    return BetaRequestDocumentBlockParam(
+    block = BetaRequestDocumentBlockParam(
         type="document",
         source=BetaBase64PDFSourceParam(
             type="base64",
@@ -621,6 +728,9 @@ def document_to_content_block(base64_data: str) -> BetaRequestDocumentBlockParam
             data=base64_data,
         ),
     )
+    if enable_citations:
+        block["citations"] = {"enabled": True}
+    return block
 
 
 def tool_use_content_block(