Publish harness and TUI open-source

initial sync from the monorepo

crates/common/xai-tool-runtime/src/render.rs

@@ -0,0 +1,731 @@
+//! Model-facing output extraction.
+//!
+//! Tool outputs carry both structured data (for agent/client logic) and
+//! a model-facing representation as MCP content blocks. [`ToolOutput`]
+//! is the trait that the runtime uses to extract the model-facing part.
+//!
+//! The default [`ToolOutput`] implementation serialises the output
+//! to JSON, then walks the structure looking for embedded
+//! [`ContentBlock`]-shaped values (images, resources). These are
+//! promoted to proper content block types; everything else becomes
+//! [`ContentBlock::Text`].
+//!
+//! Use [`extract_content_blocks`] directly when you need the same
+//! conversion on an arbitrary [`serde_json::Value`].
+//!
+//! # Example — custom model output
+//!
+//! ```rust
+//! use serde::Serialize;
+//! use xai_tool_runtime::render::ToolOutput;
+//! use xai_tool_runtime::ContentBlock;
+//!
+//! #[derive(Serialize)]
+//! struct BashOutput {
+//!     stdout: String,
+//!     exit_code: i32,
+//!     model_output: Vec<ContentBlock>,
+//! }
+//!
+//! impl ToolOutput for BashOutput {
+//!     fn model_output(&self) -> Vec<ContentBlock> {
+//!         self.model_output.clone()
+//!     }
+//! }
+//! ```
+//!
+//! # Example — default (automatic MCP extraction)
+//!
+//! Types that don't override `model_output()` get automatic extraction.
+//! Embedded images and resources are promoted; the rest is JSON text:
+//!
+//! ```rust
+//! use serde::Serialize;
+//! use xai_tool_runtime::render::ToolOutput;
+//!
+//! #[derive(Serialize)]
+//! struct SimpleOutput { answer: String }
+//! impl ToolOutput for SimpleOutput {}
+//! // model sees: ContentBlock::Text { text: r#"{"answer":"..."}"# }
+//! ```
+
+use std::sync::Arc;
+
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+
+use crate::tool::ContentBlock;
+
+/// Unified trait for typed tool outputs.
+///
+/// Combines model-facing content extraction and optional
+/// chat-completion response generation into a single trait.
+pub trait ToolOutput: Serialize {
+    /// Returns the model-facing content blocks for this output.
+    ///
+    /// Return an empty `Vec` to signal "use automatic extraction" — the
+    /// runtime will call [`extract_content_blocks`] on the serialised
+    /// JSON value instead.
+    fn model_output(&self) -> Vec<ContentBlock> {
+        Vec::new()
+    }
+
+    /// Build a chat-completion response frame from this tool output (sent to client),
+    /// if applicable.  Returns `None` by default.
+    fn chat_completion_output(&self) -> Option<ToolChatCompletionResponse> {
+        None
+    }
+}
+
+/// Blanket impl so `serde_json::Value` can be used directly as a
+/// `Tool::Output` (handy for stub/test tools and pass-through proxies).
+impl ToolOutput for Value {}
+
+/// `String` is a common output type for simple tools.
+impl ToolOutput for String {}
+
+/// Lets `xai_tool_types::TaskOutputOutput` be used directly as a `Tool::Output`
+/// (handy for stub/test tools and pass-through proxies).
+impl ToolOutput for xai_tool_types::TaskOutputOutput {}
+
+/// Lets `xai_tool_types::SubagentCompletedOutput` be used directly as a
+/// `Tool::Output` (the `task` tool's structured completion output).
+impl ToolOutput for xai_tool_types::SubagentCompletedOutput {}
+
+/// Lets `xai_tool_types::KillTaskOutput` be used directly as a `Tool::Output`
+/// (the `kill_task` tool's typed result / not-found output).
+impl ToolOutput for xai_tool_types::KillTaskOutput {}
+
+/// Delegate through `Box<T>` so boxed outputs (e.g. large response
+/// structs) work without a manual impl.
+impl<T: ToolOutput + Serialize + ?Sized> ToolOutput for Box<T> {
+    fn model_output(&self) -> Vec<ContentBlock> {
+        (**self).model_output()
+    }
+
+    fn chat_completion_output(&self) -> Option<ToolChatCompletionResponse> {
+        (**self).chat_completion_output()
+    }
+}
+
+/// Minimal representation of a chat-completion response streamed to the
+/// frontend.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ToolChatCompletionResponse {
+    /// The main completion payload.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub result: Option<ToolChatCompletion>,
+    /// Structured stream error (e.g. rate-limit, tool failure).
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stream_error: Option<ToolStreamError>,
+}
+
+/// Minimal chat completion response for tool result (sent to client).
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ToolChatCompletion {
+    /// Always `"assistant"`.
+    #[serde(default)]
+    pub sender: String,
+    /// Text body of the response.
+    #[serde(default)]
+    pub message: String,
+    /// Tag discriminator: `"final"`, `"raw_function_result"`,
+    /// `"tool_usage_card"`, `"tool_partial_output"`, etc.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub message_tag: Option<String>,
+    /// Identifies the tool-usage card this result belongs to.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_usage_card_id: Option<String>,
+    /// JSON-encoded card attachment (images, render cards, files).
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub card_attachment: Option<String>,
+    /// Media generation type: `"image_gen"`, `"video_gen"`, etc.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub media_gen_type: Option<String>,
+    /// Code execution result.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub code_execution_result: Option<ToolCodeExecutionResult>,
+    /// Catch-all for additional fields the tool wants to set. Merged
+    /// into the proto `ChatCompletion` by the downstream converter.
+    #[serde(flatten)]
+    pub extra: serde_json::Map<String, Value>,
+}
+
+/// Lightweight code-execution result carried on the completion.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ToolCodeExecutionResult {
+    #[serde(default)]
+    pub stdout: String,
+    #[serde(default)]
+    pub stderr: String,
+    #[serde(default)]
+    pub exit_code: i32,
+    #[serde(default)]
+    pub command_timed_out: bool,
+}
+
+/// Structured stream error returned alongside the completion.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ToolStreamError {
+    pub message: String,
+    /// Opaque typed-error payload. The downstream chat layer
+    /// deserialises this into the concrete proto enum variant.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub typed_error: Option<Value>,
+}
+
+/// Extract MCP-compatible [`ContentBlock`]s from a serialised JSON value.
+///
+/// Strategies are tried in order — first match wins:
+///
+/// | # | Shape | Result |
+/// |---|-------|--------|
+/// | 1 | Value is itself a `ContentBlock` (`{"type":"text",…}`) | `vec![block]` |
+/// | 2 | Array containing ≥ 1 `ContentBlock` | each element: block or text |
+/// | 3 | Object with `"content": [...]` (MCP `CallToolResult`) | `structuredContent` (if any) as JSON text, followed by the content array |
+/// | 4 | Object with mixed fields | block-shaped fields extracted, rest as JSON text |
+/// | 5 | Anything else | `ContentBlock::Text` with the stringified value |
+pub fn extract_content_blocks(value: &Value) -> Vec<ContentBlock> {
+    // 1. Value IS a single ContentBlock.
+    if let Some(block) = try_parse_block(value) {
+        return vec![block];
+    }
+
+    // 2. Array: convert each element (block-shaped -> block, else -> text).
+    //    Only enter this path when at least one element looks like a
+    //    ContentBlock so plain arrays like [1,2,3] fall through to text.
+    if let Some(arr) = value.as_array()
+        && !arr.is_empty()
+        && arr.iter().any(looks_like_content_block)
+    {
+        return arr.iter().map(value_to_block).collect();
+    }
+
+    if let Some(obj) = value.as_object() {
+        // grok-build `ToolRunResult`: the model sees `prompt_text` (reminders
+        // appended), never a JSON dump of the structured result.
+        if let Some(Value::String(prompt_text)) = obj.get("prompt_text")
+            && obj.contains_key("output")
+            && obj.contains_key("effective_tool_name")
+        {
+            return vec![ContentBlock::Text {
+                text: prompt_text.clone(),
+            }];
+        }
+
+        // 3. Object with a `"content"` array -> the standard MCP
+        //    CallToolResult shape.
+        if let Some(Value::Array(arr)) = obj.get("content")
+            && !arr.is_empty()
+            && arr.iter().any(looks_like_content_block)
+        {
+            // Surface `structuredContent` so IDs/handles the server
+            // expects the model to round-trip aren't dropped.
+            let structured = obj
+                .get("structuredContent")
+                .filter(|v| !v.is_null())
+                .map(|v| ContentBlock::Text {
+                    text: v.to_string(),
+                });
+            let mut blocks = Vec::with_capacity(arr.len() + structured.is_some() as usize);
+            blocks.extend(structured);
+            blocks.extend(arr.iter().map(value_to_block));
+            return blocks;
+        }
+
+        // 4. Mixed object -> pull block-shaped field values out; collect
+        //    the remaining fields into a single JSON text block.
+        let mut extracted = Vec::new();
+        let mut remainder = serde_json::Map::new();
+
+        for (key, val) in obj {
+            match classify_field(val) {
+                FieldShape::Block(block) => extracted.push(block),
+                FieldShape::Blocks(blocks) => extracted.extend(blocks),
+                FieldShape::Other => {
+                    remainder.insert(key.clone(), val.clone());
+                }
+            }
+        }
+
+        if !extracted.is_empty() {
+            let mut result = Vec::new();
+            if !remainder.is_empty() {
+                result.push(ContentBlock::Text {
+                    text: Value::Object(remainder).to_string(),
+                });
+            }
+            result.extend(extracted);
+            return result;
+        }
+    }
+
+    // 5. Fallback -> render the whole value as text.
+    vec![value_to_block(value)]
+}
+
+// ---------------------------------------------------------------------------
+// Internal helpers
+// ---------------------------------------------------------------------------
+
+/// The `ContentBlock` enum is `#[serde(tag = "type", rename_all =
+/// "snake_case")]`, so a JSON object can only be a content block when
+/// it has `"type"` set to one of these three values.
+const CONTENT_BLOCK_TYPES: &[&str] = &["text", "image", "resource"];
+
+/// Cheap check: could `value` plausibly deserialise as a
+/// [`ContentBlock`]?  Only objects with a `"type"` field whose value
+/// is one of the known discriminators pass.  This avoids a full
+/// `from_value(clone())` on the vast majority of values.
+fn looks_like_content_block(value: &Value) -> bool {
+    value
+        .as_object()
+        .and_then(|obj| obj.get("type"))
+        .and_then(Value::as_str)
+        .is_some_and(|t| CONTENT_BLOCK_TYPES.contains(&t))
+}
+
+/// Try to parse `value` as a [`ContentBlock`].  Returns `None`
+/// immediately when the value doesn't pass the cheap
+/// [`looks_like_content_block`] check, avoiding clone + full
+/// deserialisation for non-matching shapes.
+fn try_parse_block(value: &Value) -> Option<ContentBlock> {
+    if !looks_like_content_block(value) {
+        return None;
+    }
+    serde_json::from_value::<ContentBlock>(value.clone()).ok()
+}
+
+/// Result of inspecting a single object field value.
+enum FieldShape {
+    /// The field value IS a single `ContentBlock`.
+    Block(ContentBlock),
+    /// The field value is an array where *every* element is a `ContentBlock`.
+    Blocks(Vec<ContentBlock>),
+    /// The field value does not look like block content.
+    Other,
+}
+
+/// Classify a field value as block content.
+///
+/// Conservative for arrays: all elements must deserialise as
+/// `ContentBlock`; mixed arrays go to `Other` so ambiguous data
+/// (e.g. `"scores": [0.9, 0.8]`) is not silently dropped.
+fn classify_field(value: &Value) -> FieldShape {
+    // Single block.
+    if let Some(block) = try_parse_block(value) {
+        return FieldShape::Block(block);
+    }
+    // Array of blocks — strict: every element must parse.
+    if let Some(arr) = value.as_array()
+        && !arr.is_empty()
+        && arr.iter().all(looks_like_content_block)
+    {
+        let blocks: Result<Vec<ContentBlock>, _> = arr
+            .iter()
+            .map(|v| serde_json::from_value::<ContentBlock>(v.clone()))
+            .collect();
+        if let Ok(blocks) = blocks {
+            return FieldShape::Blocks(blocks);
+        }
+    }
+    FieldShape::Other
+}
+
+/// Convert a single `Value` to a `ContentBlock`.
+///
+/// Uses the cheap [`try_parse_block`] check first; on failure wraps
+/// the value as `ContentBlock::Text`.  Strings are used verbatim (no
+/// extra JSON quoting); all other types go through `Value::to_string`.
+fn value_to_block(value: &Value) -> ContentBlock {
+    try_parse_block(value).unwrap_or_else(|| ContentBlock::Text {
+        text: match value {
+            Value::String(s) => s.clone(),
+            other => other.to_string(),
+        },
+    })
+}
+
+// ---------------------------------------------------------------------------
+// Type-erased extractor (used by the toolbox registry)
+// ---------------------------------------------------------------------------
+
+/// Type-erased model output extractor.
+pub type ModelOutputExtractor = Arc<dyn Fn(&Value) -> Option<Vec<ContentBlock>> + Send + Sync>;
+
+/// Build a [`ModelOutputExtractor`] for a concrete output type.
+pub fn extractor_for<T>() -> ModelOutputExtractor
+where
+    T: ToolOutput + serde::de::DeserializeOwned + 'static,
+{
+    Arc::new(|value: &Value| {
+        serde_json::from_value::<T>(value.clone())
+            .ok()
+            .map(|output| {
+                let blocks = output.model_output();
+                if blocks.is_empty() {
+                    extract_content_blocks(value)
+                } else {
+                    blocks
+                }
+            })
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    // ── ToolOutput with custom override ─────────────────────────────
+
+    #[derive(Serialize)]
+    struct FakeOutput {
+        blocks: Vec<ContentBlock>,
+    }
+
+    impl ToolOutput for FakeOutput {
+        fn model_output(&self) -> Vec<ContentBlock> {
+            self.blocks.clone()
+        }
+    }
+
+    #[test]
+    fn custom_text_block() {
+        let o = FakeOutput {
+            blocks: vec![ContentBlock::Text {
+                text: "hello".into(),
+            }],
+        };
+        assert_eq!(o.model_output().len(), 1);
+        assert_eq!(
+            o.model_output()[0],
+            ContentBlock::Text {
+                text: "hello".into()
+            }
+        );
+    }
+
+    #[test]
+    fn custom_multimodal() {
+        let o = FakeOutput {
+            blocks: vec![
+                ContentBlock::Text {
+                    text: "result:".into(),
+                },
+                ContentBlock::Image {
+                    mime_type: "image/png".into(),
+                    data: "iVBOR...".into(),
+                    media_id: None,
+                    filename: None,
+                    path: None,
+                    metadata: Default::default(),
+                },
+            ],
+        };
+        assert_eq!(o.model_output().len(), 2);
+    }
+
+    // ── ToolOutput default → empty (runtime fills via extract) ──────
+
+    #[test]
+    fn default_model_output_returns_empty() {
+        #[derive(Serialize)]
+        struct Plain {
+            value: u32,
+        }
+        impl ToolOutput for Plain {}
+
+        // Default signals "use automatic extraction" by returning empty.
+        assert!(Plain { value: 42 }.model_output().is_empty());
+    }
+
+    #[test]
+    fn runtime_fills_empty_model_output_via_extract() {
+        // Simulates what the ToolDyn blanket does: serialise once,
+        // then extract_content_blocks on the Value.
+        #[derive(Serialize)]
+        struct Plain {
+            value: u32,
+        }
+        impl ToolOutput for Plain {}
+
+        let p = Plain { value: 42 };
+        let value = serde_json::to_value(&p).unwrap();
+        let custom = p.model_output();
+        let blocks = if custom.is_empty() {
+            extract_content_blocks(&value)
+        } else {
+            custom
+        };
+        assert_eq!(blocks.len(), 1);
+        assert_eq!(
+            blocks[0],
+            ContentBlock::Text {
+                text: r#"{"value":42}"#.into(),
+            }
+        );
+    }
+
+    // ── extract_content_blocks unit tests ──────────────────────────
+
+    // Strategy 1: single ContentBlock
+    #[test]
+    fn extract_single_text_block() {
+        let v = json!({"type": "text", "text": "hi"});
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(blocks, vec![ContentBlock::Text { text: "hi".into() }]);
+    }
+
+    #[test]
+    fn extract_single_image_block() {
+        let v = json!({"type": "image", "mime_type": "image/png", "data": "abc"});
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(
+            blocks,
+            vec![ContentBlock::Image {
+                mime_type: "image/png".into(),
+                data: "abc".into(),
+                media_id: None,
+                filename: None,
+                path: None,
+                metadata: Default::default(),
+            }]
+        );
+    }
+
+    #[test]
+    fn extract_single_resource_block() {
+        let v = json!({"type": "resource", "uri": "file:///x"});
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(
+            blocks,
+            vec![ContentBlock::Resource {
+                uri: "file:///x".into(),
+                mime_type: None,
+                text: None,
+            }]
+        );
+    }
+
+    #[test]
+    fn extract_mcp_image_with_camel_case() {
+        let v = json!({"type": "image", "mimeType": "image/png", "data": "abc"});
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(
+            blocks,
+            vec![ContentBlock::Image {
+                mime_type: "image/png".into(),
+                data: "abc".into(),
+                media_id: None,
+                filename: None,
+                path: None,
+                metadata: Default::default(),
+            }]
+        );
+    }
+
+    // Strategy 2: array of blocks
+    #[test]
+    fn extract_array_of_blocks() {
+        let v = json!([
+            {"type": "text", "text": "a"},
+            {"type": "image", "mime_type": "image/png", "data": "b"},
+        ]);
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(blocks.len(), 2);
+        assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "a"));
+        assert!(matches!(&blocks[1], ContentBlock::Image { .. }));
+    }
+
+    #[test]
+    fn extract_mixed_array_promotes_non_blocks_to_text() {
+        let v = json!([
+            {"type": "text", "text": "a"},
+            42,
+            "raw string",
+        ]);
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(blocks.len(), 3);
+        assert_eq!(blocks[0], ContentBlock::Text { text: "a".into() });
+        assert_eq!(blocks[1], ContentBlock::Text { text: "42".into() });
+        assert_eq!(
+            blocks[2],
+            ContentBlock::Text {
+                text: "raw string".into()
+            }
+        );
+    }
+
+    #[test]
+    fn plain_array_falls_through_to_text() {
+        // No block-shaped elements → single text fallback.
+        let v = json!([1, 2, 3]);
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(blocks.len(), 1);
+        assert_eq!(
+            blocks[0],
+            ContentBlock::Text {
+                text: "[1,2,3]".into()
+            }
+        );
+    }
+
+    // Strategy 3: object with "content" key
+    #[test]
+    fn extract_content_field() {
+        let v = json!({
+            "is_error": false,
+            "content": [
+                {"type": "text", "text": "summary"},
+                {"type": "image", "mime_type": "image/png", "data": "b64"},
+            ],
+        });
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(blocks.len(), 2);
+        assert_eq!(
+            blocks[0],
+            ContentBlock::Text {
+                text: "summary".into()
+            }
+        );
+        assert!(matches!(blocks[1], ContentBlock::Image { .. }));
+    }
+
+    // Strategy 3 + structuredContent (MCP CallToolResult)
+    #[test]
+    fn extract_content_with_structured_content_surfaces_id() {
+        let v = json!({
+            "content": [
+                {"type": "resource", "uri": "ui://tldraw/canvas",
+                 "mimeType": "text/html", "text": "<html>…</html>"},
+            ],
+            "structuredContent": {"drawing_id": "abc123", "title": "sketch"},
+            "isError": false,
+            "_meta": {"ui": {"resourceUri": "ui://tldraw/canvas"}},
+        });
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(blocks.len(), 2, "expected structured + resource");
+        // structuredContent rendered as JSON text, ahead of content.
+        let ContentBlock::Text { text } = &blocks[0] else {
+            panic!("expected first block to be Text, got {:?}", blocks[0]);
+        };
+        assert!(
+            text.contains("abc123"),
+            "structuredContent must surface drawing_id, got: {text}"
+        );
+        assert!(matches!(&blocks[1], ContentBlock::Resource { .. }));
+    }
+
+    #[test]
+    fn extract_content_with_null_structured_content_omits_it() {
+        let v = json!({
+            "content": [{"type": "text", "text": "ok"}],
+            "structuredContent": null,
+        });
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(blocks.len(), 1);
+        assert_eq!(blocks[0], ContentBlock::Text { text: "ok".into() });
+    }
+
+    #[test]
+    fn extract_content_without_structured_content_unchanged() {
+        let v = json!({
+            "content": [{"type": "text", "text": "ok"}],
+            "isError": false,
+        });
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(blocks.len(), 1);
+        assert_eq!(blocks[0], ContentBlock::Text { text: "ok".into() });
+    }
+
+    // Strategy 4: mixed object with block-shaped field values
+    #[test]
+    fn extract_mixed_object_separates_blocks_and_remainder() {
+        let v = json!({
+            "summary": "found 3 results",
+            "count": 3,
+            "screenshot": {"type": "image", "mime_type": "image/png", "data": "b64"},
+        });
+        let blocks = extract_content_blocks(&v);
+        // Remainder (summary + count) as JSON text, then the image.
+        assert_eq!(blocks.len(), 2);
+        // First block is the remainder text (field order in JSON objects
+        // is not guaranteed, so just check it's Text and non-empty).
+        assert!(matches!(&blocks[0], ContentBlock::Text { text } if text.contains("summary")));
+        assert!(matches!(&blocks[1], ContentBlock::Image { .. }));
+    }
+
+    #[test]
+    fn extract_object_with_block_array_field() {
+        let v = json!({
+            "metadata": "info",
+            "results": [
+                {"type": "text", "text": "a"},
+                {"type": "text", "text": "b"},
+            ],
+        });
+        let blocks = extract_content_blocks(&v);
+        // results field → 2 blocks extracted, metadata → remainder text.
+        assert_eq!(blocks.len(), 3);
+        assert!(matches!(&blocks[0], ContentBlock::Text { text } if text.contains("metadata")));
+        assert_eq!(blocks[1], ContentBlock::Text { text: "a".into() });
+        assert_eq!(blocks[2], ContentBlock::Text { text: "b".into() });
+    }
+
+    // Strategy 5: fallback
+    #[test]
+    fn extract_plain_string() {
+        let blocks = extract_content_blocks(&json!("hello world"));
+        assert_eq!(
+            blocks,
+            vec![ContentBlock::Text {
+                text: "hello world".into()
+            }]
+        );
+    }
+
+    #[test]
+    fn extract_number() {
+        let blocks = extract_content_blocks(&json!(42));
+        assert_eq!(blocks, vec![ContentBlock::Text { text: "42".into() }]);
+    }
+
+    #[test]
+    fn extract_null() {
+        let blocks = extract_content_blocks(&json!(null));
+        assert_eq!(
+            blocks,
+            vec![ContentBlock::Text {
+                text: "null".into()
+            }]
+        );
+    }
+
+    #[test]
+    fn extract_plain_object_no_blocks() {
+        let v = json!({"a": 1, "b": "two"});
+        let blocks = extract_content_blocks(&v);
+        assert_eq!(blocks.len(), 1);
+        assert!(matches!(&blocks[0], ContentBlock::Text { .. }));
+    }
+
+    /// `extract_content_blocks` surfaces a `ToolRunResult`'s `prompt_text` as the
+    /// model content, not a JSON dump of the struct.
+    #[test]
+    fn tool_run_result_shape_extracts_prompt_text() {
+        let prompt = "1: a.txt\n2: b.txt\n<system-reminder>\nThe todo_write tool \
+                      hasn't been used recently.\n</system-reminder>";
+        let v = json!({
+            "output": {"list_dir": {"entries": ["a.txt", "b.txt"]}},
+            "prompt_text": prompt,
+            "effective_tool_name": null,
+        });
+        assert_eq!(
+            extract_content_blocks(&v),
+            vec![ContentBlock::Text {
+                text: prompt.to_owned()
+            }]
+        );
+    }
+}