Rich ToolResult: follow-up content, metadata, control hints, replay class (tinyagents#171)

crates/tinytools/README.md

@@ -35,7 +35,77 @@ impl Tool for Echo {
 `Tool` describes a callable capability. Its result is a `ToolResult` block list
 with a reported-error flag and optional markdown rendering. `ToolSpec` is the
 model-visible declaration. `ToolRunContext` exposes only tool-relevant run
-facts: workspace, thread id, and output cap.
+facts: workspace, thread id, and output cap — plus `host_extension()`, the
+same type-erased escape hatch `Tool::host_extension` offers, so a tool written
+against one specific harness can downcast to that harness's full context.
+
+## Rich tool returns
+
+`ToolContent` has four block kinds: `Text`, `Json`, `Image`, and `File`.
+`Image` carries a MIME `media_type` plus `ImageData::{Base64, Url}`; `File`
+carries a display `name`, a `media_type`, and `FileData::{Base64, Url, Path}`.
+`ToolResult::text()`, `output()`, and `output_for_llm()` render `Text`
+verbatim, pretty-print `Json`, skip `Json` from `text()` specifically (as
+before), and render a short placeholder for `Image`/`File` — `[image
+image/png]`, `[file report.pdf (application/pdf)]` — so a model still gets a
+sensible turn even when a renderer does not special-case the new block kinds.
+
+`ToolResult` carries three more optional surfaces beyond `content` and
+`markdown_formatted`:
+
+- `follow_up: Vec<ToolContent>` — content the caller should present to the
+  model as a *separate* user message after the tool result (a screenshot, a
+  generated document). It is never included in `text()`, `output()`, or
+  `output_for_llm()`; a host that wants to honour it reads the field directly.
+  Attach it with `with_follow_up(..)`.
+- `metadata: Option<serde_json::Value>` — host-only data (trace ids, raw
+  provider payloads) that is never shown to the model. Attach it with
+  `with_metadata(..)`.
+- `control: Option<ToolControl>` — loop-control hints a harness may honour:
+  `return_direct: Option<bool>`, `terminate`, `goto: Option<String>`, and
+  `state_update: Option<serde_json::Value>`. Set them with the builders
+  `return_direct()`, `dont_return_direct()`, `terminate()`, `with_goto(..)`,
+  and `with_state_update(..)`, which lazily create the `ToolControl`.
+  `return_direct` is tri-state, not a defaulted `bool`: a call that only used
+  `with_goto(..)`, `with_state_update(..)`, or `terminate()` leaves it `None`
+  rather than an implicit `false`, so it cannot silently suppress a tool's
+  static `true` default — see "Static and per-call return-direct" below.
+
+`ToolResult::retry(message)` and `ToolResult::failed(message)` both set
+`is_error`, same as `error(message)`, but additionally tag
+`error_kind: Option<ToolErrorKind>` as `Retry` or `Failed` — Pydantic AI's
+`ModelRetry` versus a permanent tool failure — so a harness can decide whether
+to loop the model back in or surface the failure as final.
+
+Every new field is `#[serde(default)]` and, where it can be empty or absent,
+`skip_serializing_if`, so a `ToolResult` persisted before these fields existed
+still decodes, and a plain result's wire shape is unchanged.
+
+## Static and per-call return-direct
+
+`Tool::return_direct()` is a static, per-tool default (`false`) for a tool
+whose entire purpose is to hand the model's answer straight back — a
+final-answer or handoff tool overrides it to `true`. `ToolResult::control`'s
+`return_direct: Option<bool>` is the per-*call* override on `ToolControl`; a
+harness should prefer `Some(..)` on the result it just received over the
+tool's static declaration, and fall back to the static declaration when it is
+`None`. `None` is the outcome of a call that never touched `return_direct` —
+including one that only used `with_goto(..)`, `with_state_update(..)`, or
+`terminate()` — so it must not be read as an explicit override. Call
+`return_direct()` for `Some(true)`, or `dont_return_direct()` for `Some(false)`
+to force the call to *not* return directly even when the tool's static
+declaration is `true`.
+
+## Replay after a crash
+
+`ToolPolicy`'s `ToolRuntime` carries `replay: ToolReplay`, mirroring pi's
+`replay` classification: whether an orphaned in-flight call for a tool may be
+safely re-executed after a crash. It defaults to `ToolReplay::Never`; a tool
+that is idempotent or otherwise safe to repeat declares `ToolReplay::Safe`
+through its `ToolPolicy`. This lives on the existing declarative policy
+surface rather than as a new `Tool` trait method, consistent with how every
+other runtime requirement (timeout, retries, cancellation, sandboxing) is
+already expressed there.
 
 `ToolPolicy` is the complete host-readable declaration around a call:
 

crates/tinytools/src/context/test.rs

@@ -56,3 +56,35 @@ fn an_implementor_is_readable_through_the_trait_object() {
         Some(PathBuf::from("/tmp/worktree"))
     );
 }
+
+/// A context carrying a host-owned payload behind the erased hook.
+struct Hosted {
+    tag: HostTag,
+}
+
+#[derive(Debug, PartialEq)]
+struct HostTag(&'static str);
+
+impl ToolRunContext for Hosted {
+    fn host_extension(&self) -> Option<&(dyn std::any::Any + Send + Sync)> {
+        Some(&self.tag)
+    }
+}
+
+#[test]
+fn the_default_host_extension_is_absent() {
+    let erased: &dyn ToolRunContext = &Bare;
+    assert!(erased.host_extension().is_none());
+}
+
+#[test]
+fn a_host_recovers_its_own_context_by_downcasting() {
+    let hosted = Hosted {
+        tag: HostTag("call-7"),
+    };
+    let erased: &dyn ToolRunContext = &hosted;
+    let tag = erased
+        .host_extension()
+        .and_then(|any| any.downcast_ref::<HostTag>());
+    assert_eq!(tag, Some(&HostTag("call-7")));
+}

crates/tinytools/src/context/types.rs

@@ -61,4 +61,17 @@ pub trait ToolRunContext: Send + Sync {
     fn workspace_policy_id(&self) -> Option<&str> {
         self.workspace().map(|w| w.policy_id.as_str())
     }
+
+    /// The host's own context object, erased, for a tool written against a
+    /// specific harness that needs more than the portable facts above.
+    ///
+    /// The same escape hatch as [`Tool::host_extension`][crate::Tool::host_extension]:
+    /// this crate has no business naming the harness's context type, so a
+    /// host returns `Some(self)` and a tool that knows which host it runs
+    /// under downcasts. Every other implementor returns `None` and pays
+    /// nothing. A tool that only needs the workspace, thread id, or output
+    /// cap should keep using the typed methods.
+    fn host_extension(&self) -> Option<&(dyn std::any::Any + Send + Sync)> {
+        None
+    }
 }

crates/tinytools/src/lib.rs

@@ -122,9 +122,9 @@ pub use naming::{
 };
 pub use permission::PermissionLevel;
 pub use policy::{
-    ToolAccess, ToolDisplay, ToolPolicy, ToolRuntime, ToolSideEffects, WorkspaceAccess,
+    ToolAccess, ToolDisplay, ToolPolicy, ToolReplay, ToolRuntime, ToolSideEffects, WorkspaceAccess,
 };
-pub use result::{ToolContent, ToolResult};
+pub use result::{FileData, ImageData, ToolContent, ToolControl, ToolErrorKind, ToolResult};
 pub use spec::ToolSpec;
 pub use tool::{Tool, ToolExposure};
 pub use workspace::{SandboxMode, WorkspaceDescriptor};

crates/tinytools/src/policy/mod.rs

@@ -7,7 +7,7 @@
 mod types;
 
 pub use types::{
-    ToolAccess, ToolDisplay, ToolPolicy, ToolRuntime, ToolSideEffects, WorkspaceAccess,
+    ToolAccess, ToolDisplay, ToolPolicy, ToolReplay, ToolRuntime, ToolSideEffects, WorkspaceAccess,
 };
 
 #[cfg(test)]

crates/tinytools/src/policy/test.rs

@@ -2,7 +2,9 @@
 
 #![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)]
 
-use super::{ToolAccess, ToolDisplay, ToolPolicy, ToolRuntime, ToolSideEffects, WorkspaceAccess};
+use super::{
+    ToolAccess, ToolDisplay, ToolPolicy, ToolReplay, ToolRuntime, ToolSideEffects, WorkspaceAccess,
+};
 use crate::{SandboxMode, ToolTimeout};
 
 #[test]
@@ -106,6 +108,7 @@ fn policy_round_trips_through_its_stable_json_shape() {
                 "cancelable": false,
                 "sandbox": "required",
                 "streaming": false,
+                "replay": "never",
             },
             "access": {
                 "workspace": "any",
@@ -140,6 +143,7 @@ fn fully_populated_policy_has_a_pinned_json_wire_shape() {
             sandbox: SandboxMode::Required,
             max_result_bytes: Some(8_192),
             streaming: true,
+            replay: ToolReplay::Safe,
         },
         access: ToolAccess {
             workspace: WorkspaceAccess::Scoped,
@@ -170,6 +174,7 @@ fn fully_populated_policy_has_a_pinned_json_wire_shape() {
             "sandbox": "required",
             "max_result_bytes": 8192,
             "streaming": true,
+            "replay": "safe",
         },
         "access": {
             "workspace": "scoped",
@@ -190,3 +195,34 @@ fn fully_populated_policy_has_a_pinned_json_wire_shape() {
         policy
     );
 }
+
+#[test]
+fn replay_defaults_to_never() {
+    assert_eq!(ToolRuntime::default().replay, ToolReplay::Never);
+}
+
+#[test]
+fn replay_round_trips_through_json() {
+    let runtime = ToolRuntime {
+        replay: ToolReplay::Safe,
+        ..ToolRuntime::default()
+    };
+    let encoded = serde_json::to_string(&runtime).expect("serializable");
+    assert!(encoded.contains("\"replay\":\"safe\""));
+    let back: ToolRuntime = serde_json::from_str(&encoded).expect("deserializable");
+    assert_eq!(back.replay, ToolReplay::Safe);
+}
+
+#[test]
+fn legacy_runtime_json_without_replay_defaults_to_never() {
+    // A `ToolRuntime` persisted before `replay` existed should still decode.
+    let literal = r#"{
+        "timeout": { "mode": "inherit" },
+        "idempotent": false,
+        "cancelable": false,
+        "sandbox": "disabled",
+        "streaming": false
+    }"#;
+    let decoded: ToolRuntime = serde_json::from_str(literal).expect("deserializable");
+    assert_eq!(decoded.replay, ToolReplay::Never);
+}

crates/tinytools/src/policy/types.rs

@@ -90,6 +90,24 @@ impl ToolDisplay {
     }
 }
 
+/// Whether an orphaned in-flight call may be safely re-executed after a
+/// crash.
+///
+/// A host that persists an in-flight call and recovers after a crash has to
+/// decide whether to replay it. Pi's `replay` classification is the reference
+/// design: most tools are not safe to blindly re-run (a payment, a send), so
+/// [`Self::Never`] is the default and a tool must opt into [`Self::Safe`].
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum ToolReplay {
+    /// An orphaned call must not be re-executed after a crash.
+    #[default]
+    Never,
+    /// An orphaned call may be safely re-executed after a crash — the tool is
+    /// idempotent or otherwise safe to repeat.
+    Safe,
+}
+
 /// Runtime requirements a tool declares for safe execution.
 ///
 /// A host decides how to apply these requirements. In particular, this type
@@ -117,6 +135,10 @@ pub struct ToolRuntime {
     pub max_result_bytes: Option<usize>,
     /// Whether the tool can emit streaming result fragments.
     pub streaming: bool,
+    /// Whether an orphaned in-flight call for this tool may be safely
+    /// re-executed after a crash. See [`ToolReplay`].
+    #[serde(default)]
+    pub replay: ToolReplay,
 }
 
 /// Access requirements a tool declares before a host exposes or runs it.

crates/tinytools/src/result/mod.rs

@@ -2,7 +2,7 @@
 
 mod types;
 
-pub use types::{ToolContent, ToolResult};
+pub use types::{FileData, ImageData, ToolContent, ToolControl, ToolErrorKind, ToolResult};
 
 #[cfg(test)]
 mod test;