Publish harness and TUI open-source
initial sync from the monorepo
This commit is contained in:
grokkybara[bot]
commit
c68e39f604
2734 changed files with 1437016 additions and 0 deletions
336
crates/common/xai-tool-runtime/src/context.rs
Normal file
336
crates/common/xai-tool-runtime/src/context.rs
Normal file
|
|
@ -0,0 +1,336 @@
|
|||
//! Context types and the typed-extension store they share.
|
||||
|
||||
use std::any::{Any, TypeId};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use xai_tool_protocol::ToolCallId;
|
||||
|
||||
/// Open typed-extension store keyed by `TypeId`.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct TypedExtensions {
|
||||
map: HashMap<TypeId, Arc<dyn Any + Send + Sync>>,
|
||||
}
|
||||
|
||||
impl TypedExtensions {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn insert<T: Send + Sync + 'static>(&mut self, value: T) -> &mut Self {
|
||||
self.map.insert(TypeId::of::<T>(), Arc::new(value));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn insert_arc<T: Send + Sync + 'static>(&mut self, value: Arc<T>) -> &mut Self {
|
||||
self.map.insert(TypeId::of::<T>(), value);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn get<T: Send + Sync + 'static>(&self) -> Option<Arc<T>> {
|
||||
self.map
|
||||
.get(&TypeId::of::<T>())
|
||||
.cloned()
|
||||
.and_then(|arc| Arc::downcast::<T>(arc).ok())
|
||||
}
|
||||
|
||||
pub fn contains<T: Send + Sync + 'static>(&self) -> bool {
|
||||
self.map.contains_key(&TypeId::of::<T>())
|
||||
}
|
||||
|
||||
pub fn remove<T: Send + Sync + 'static>(&mut self) -> Option<Arc<T>> {
|
||||
self.map
|
||||
.remove(&TypeId::of::<T>())
|
||||
.and_then(|arc| Arc::downcast::<T>(arc).ok())
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.map.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.map.is_empty()
|
||||
}
|
||||
|
||||
/// Copy entries from `defaults` that are not already present in `self`.
|
||||
pub fn merge_defaults(&mut self, defaults: &TypedExtensions) {
|
||||
for (key, value) in &defaults.map {
|
||||
self.map.entry(*key).or_insert_with(|| value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-call context.
|
||||
#[derive(Clone)]
|
||||
pub struct ToolCallContext {
|
||||
pub call_id: ToolCallId,
|
||||
pub extensions: TypedExtensions,
|
||||
}
|
||||
|
||||
impl Default for ToolCallContext {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
call_id: ToolCallId::new_v7(),
|
||||
extensions: TypedExtensions::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolCallContext {
|
||||
pub fn new(call_id: ToolCallId) -> Self {
|
||||
Self {
|
||||
call_id,
|
||||
extensions: TypedExtensions::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Delegate to `self.extensions.insert()`.
|
||||
pub fn insert<T: Send + Sync + 'static>(&mut self, value: T) -> &mut Self {
|
||||
self.extensions.insert(value);
|
||||
self
|
||||
}
|
||||
|
||||
/// Delegate to `self.extensions.get()`.
|
||||
pub fn get<T: Send + Sync + 'static>(&self) -> Option<Arc<T>> {
|
||||
self.extensions.get::<T>()
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-turn context consumed by [`crate::Tool::should_list`].
|
||||
#[derive(Clone, Default)]
|
||||
pub struct ListToolsContext {
|
||||
pub extensions: TypedExtensions,
|
||||
}
|
||||
|
||||
impl ListToolsContext {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
// Runtime-blessed per-concept extensions. One type per concept so
|
||||
// dispatchers install exactly what they have and tools depend on
|
||||
// exactly what they need.
|
||||
|
||||
/// Working directory for relative path resolution.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Cwd(pub PathBuf);
|
||||
|
||||
/// Opaque behaviour version. Tools that branch on this MUST treat
|
||||
/// unknown values as a hard error.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct BehaviorVersion(pub String);
|
||||
|
||||
/// Distributed-trace correlation context (e.g. W3C `traceparent`).
|
||||
///
|
||||
/// Receive-side carrier only: stamped from the inbound wire value for
|
||||
/// tool impls to read, never serialized back out.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct TraceContext(pub String);
|
||||
|
||||
/// Session ID context — identifies which hub session this call belongs to.
|
||||
/// Used by multi-session tool servers to dispatch to the correct
|
||||
/// per-session state.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SessionContext(pub String);
|
||||
|
||||
/// Cooperative-cancellation handle for the current tool call. Tools MAY
|
||||
/// poll/await this for graceful shutdown; the dispatcher also hard-cancels
|
||||
/// by dropping the call future when it fires.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Cancellation(pub tokio_util::sync::CancellationToken);
|
||||
|
||||
/// Per-user feature-flag bag attached as a [`ToolCallContext`] extension.
|
||||
/// Dispatcher resolves; tools read. Default = "off" for every field so an
|
||||
/// absent extension never accidentally opts a feature in. Extend by
|
||||
/// adding fields with safe defaults; new fields need `#[serde(default)]`
|
||||
/// so older `session.bind` payloads stay deserializable.
|
||||
#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct WorkspaceViewerContext {
|
||||
/// When `true`, `BashTool` emits `bash_output_chunk` Progress frames.
|
||||
#[serde(default)]
|
||||
pub stream_tool_progress: bool,
|
||||
}
|
||||
|
||||
/// Wire shape of the Computer Hub `session.bind` metadata — one definition
|
||||
/// shared by the emitter (serializes) and the workspace consumer
|
||||
/// (deserializes), so the two can't drift on field names/types.
|
||||
///
|
||||
/// Excludes anything not meant for the workspace (cached tool definitions,
|
||||
/// and terminal-provisioning inputs like image/fuse/isolation) so they can
|
||||
/// never reach the wire. Every field tolerates a missing/malformed value
|
||||
/// (drops to default) to keep valid siblings and mixed-version compatibility.
|
||||
#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct WorkspaceBindMetadata {
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub preset: Option<String>,
|
||||
/// Raw string; the workspace maps it to its own capability enum.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub capability_mode: Option<String>,
|
||||
/// Explicit toolset in the grok-tools gRPC wire shape. Empty = unset.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Vec::is_empty"
|
||||
)]
|
||||
pub tools: Vec<xai_grok_tools_api::ToolConfigEntry>,
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub viewer_ctx: Option<WorkspaceViewerContext>,
|
||||
/// Initial auto-approve (YOLO) state for the bound session. Omitted when
|
||||
/// unset (legacy emitters / wire compat with older workspace servers);
|
||||
/// consumers fail closed on `None`.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub yolo_mode: Option<bool>,
|
||||
/// Optional/additive: omitted by emitters that don't yet write it.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub manifest_version: Option<String>,
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub manifest_hash: Option<String>,
|
||||
/// Opt-in: forward SystemNotifications produced in this session to the gateway.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub system_notifications: Option<bool>,
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "std::ops::Not::not"
|
||||
)]
|
||||
pub rpc_only: bool,
|
||||
}
|
||||
|
||||
/// Deserialize a field, falling back to its default on a malformed value
|
||||
/// instead of failing the whole struct.
|
||||
fn ok_or_default<'de, D, T>(deserializer: D) -> Result<T, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
T: serde::de::DeserializeOwned + Default,
|
||||
{
|
||||
let value = <serde_json::Value as serde::Deserialize>::deserialize(deserializer)?;
|
||||
Ok(serde_json::from_value(value).unwrap_or_default())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod bind_metadata_tests {
|
||||
use super::WorkspaceBindMetadata;
|
||||
|
||||
#[test]
|
||||
fn serialize_omits_empty_fields() {
|
||||
let md = WorkspaceBindMetadata::default();
|
||||
assert_eq!(serde_json::to_value(&md).unwrap(), serde_json::json!({}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_trips_populated() {
|
||||
let md = WorkspaceBindMetadata {
|
||||
preset: Some("explore".to_owned()),
|
||||
capability_mode: Some("read_only".to_owned()),
|
||||
tools: vec![xai_grok_tools_api::ToolConfigEntry {
|
||||
id: "GrokBuild:grep".to_owned(),
|
||||
..Default::default()
|
||||
}],
|
||||
viewer_ctx: Some(super::WorkspaceViewerContext {
|
||||
stream_tool_progress: true,
|
||||
}),
|
||||
yolo_mode: Some(true),
|
||||
manifest_version: Some("v1".to_owned()),
|
||||
manifest_hash: Some("abc123".to_owned()),
|
||||
system_notifications: Some(true),
|
||||
rpc_only: true,
|
||||
};
|
||||
let value = serde_json::to_value(&md).unwrap();
|
||||
let back: WorkspaceBindMetadata = serde_json::from_value(value).unwrap();
|
||||
assert_eq!(back.preset.as_deref(), Some("explore"));
|
||||
assert_eq!(back.capability_mode.as_deref(), Some("read_only"));
|
||||
assert_eq!(back.tools.len(), 1);
|
||||
assert!(back.viewer_ctx.unwrap().stream_tool_progress);
|
||||
assert_eq!(back.yolo_mode, Some(true));
|
||||
assert_eq!(back.manifest_version.as_deref(), Some("v1"));
|
||||
assert_eq!(back.manifest_hash.as_deref(), Some("abc123"));
|
||||
assert_eq!(back.system_notifications, Some(true));
|
||||
assert!(back.rpc_only);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rpc_only_omitted_when_false_wire_compatible() {
|
||||
let md = WorkspaceBindMetadata::default();
|
||||
let value = serde_json::to_value(&md).unwrap();
|
||||
assert!(value.get("rpc_only").is_none());
|
||||
|
||||
let md: WorkspaceBindMetadata =
|
||||
serde_json::from_value(serde_json::json!({"preset": "explore"})).unwrap();
|
||||
assert!(!md.rpc_only);
|
||||
|
||||
let md: WorkspaceBindMetadata =
|
||||
serde_json::from_value(serde_json::json!({"rpc_only": true})).unwrap();
|
||||
assert!(md.rpc_only);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn system_notifications_is_wire_compatible() {
|
||||
let md = WorkspaceBindMetadata::default();
|
||||
let value = serde_json::to_value(&md).unwrap();
|
||||
assert!(value.get("system_notifications").is_none());
|
||||
|
||||
let md = WorkspaceBindMetadata {
|
||||
system_notifications: Some(true),
|
||||
..Default::default()
|
||||
};
|
||||
let value = serde_json::to_value(&md).unwrap();
|
||||
let back: WorkspaceBindMetadata = serde_json::from_value(value).unwrap();
|
||||
assert_eq!(back.system_notifications, Some(true));
|
||||
|
||||
let md: WorkspaceBindMetadata =
|
||||
serde_json::from_value(serde_json::json!({"preset": "explore"})).unwrap();
|
||||
assert!(md.system_notifications.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_field_falls_back_to_default_keeping_siblings() {
|
||||
// `tools` is the wrong type and `capability_mode` is fine: the bad
|
||||
// field drops to default, the good sibling survives.
|
||||
let value = serde_json::json!({
|
||||
"preset": "explore",
|
||||
"capability_mode": "read_only",
|
||||
"tools": "not-a-list",
|
||||
});
|
||||
let md: WorkspaceBindMetadata = serde_json::from_value(value).unwrap();
|
||||
assert_eq!(md.preset.as_deref(), Some("explore"));
|
||||
assert_eq!(md.capability_mode.as_deref(), Some("read_only"));
|
||||
assert!(md.tools.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn legacy_payload_without_viewer_ctx_parses() {
|
||||
let md: WorkspaceBindMetadata =
|
||||
serde_json::from_value(serde_json::json!({"preset": "explore"})).unwrap();
|
||||
assert!(md.viewer_ctx.is_none());
|
||||
}
|
||||
}
|
||||
68
crates/common/xai-tool-runtime/src/dispatch.rs
Normal file
68
crates/common/xai-tool-runtime/src/dispatch.rs
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
//! Object-safe `ToolDispatch` trait — the runtime contract for handling tool calls.
|
||||
//!
|
||||
//! `Tool` itself is not object-safe (it carries associated `Args` /
|
||||
//! `Output` types), so implementations expose a JSON-typed surface and rely on
|
||||
//! per-tool adapters to encode/decode at the boundary. The default
|
||||
//! `call_terminal` impl drains the stream so the common "I just want the
|
||||
//! result" path doesn't have to depend on `futures` internals.
|
||||
//!
|
||||
//! This crate is upstream of every concrete impl. Doc-comments here describe
|
||||
//! trait semantics in terms of "the runtime" or "the implementation" —
|
||||
//! concrete dispatch routers live downstream and are intentionally not named
|
||||
//! here.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use serde_json::Value;
|
||||
|
||||
use xai_tool_protocol::ToolId;
|
||||
|
||||
use crate::context::ToolCallContext;
|
||||
use crate::error::ToolError;
|
||||
use crate::tool::{ToolStream, ToolStreamItem, TypedToolOutput};
|
||||
|
||||
/// Object-safe tool dispatch interface.
|
||||
///
|
||||
/// Implementations route the `tool_id` to the correct tool, decode `args`
|
||||
/// against the tool's typed `Args`, and return the streaming result as
|
||||
/// [`TypedToolOutput`] — preserving model-facing content blocks and
|
||||
/// optional chat-completion metadata end-to-end. Raw `Value` only appears
|
||||
/// at JSON-RPC wire encode/decode boundaries.
|
||||
#[async_trait]
|
||||
pub trait ToolDispatch: Send + Sync {
|
||||
/// Streaming dispatch. The returned stream MUST end with exactly one
|
||||
/// `Terminal` item per the [`ToolStream`] invariant.
|
||||
async fn call(
|
||||
&self,
|
||||
tool_id: ToolId,
|
||||
args: Value,
|
||||
ctx: ToolCallContext,
|
||||
) -> ToolStream<TypedToolOutput>;
|
||||
|
||||
/// Drain the stream and return only the terminal result. Useful for
|
||||
/// callers that don't care about progress chunks.
|
||||
///
|
||||
/// Default impl pulls items off the stream and discards `Progress`
|
||||
/// items; the first `Terminal` short-circuits. A stream that ends
|
||||
/// without a `Terminal` is a protocol violation by the implementation;
|
||||
/// the default surfaces this as `ToolError::Custom { code:
|
||||
/// "stream_no_terminal", ... }`.
|
||||
async fn call_terminal(
|
||||
&self,
|
||||
tool_id: ToolId,
|
||||
args: Value,
|
||||
ctx: ToolCallContext,
|
||||
) -> Result<TypedToolOutput, ToolError> {
|
||||
let mut stream = self.call(tool_id, args, ctx).await;
|
||||
while let Some(item) = stream.next().await {
|
||||
match item {
|
||||
ToolStreamItem::Progress(_) => continue,
|
||||
ToolStreamItem::Terminal(result) => return result,
|
||||
}
|
||||
}
|
||||
Err(ToolError::custom(
|
||||
"stream_no_terminal",
|
||||
"dispatch stream ended without a terminal item",
|
||||
))
|
||||
}
|
||||
}
|
||||
554
crates/common/xai-tool-runtime/src/error.rs
Normal file
554
crates/common/xai-tool-runtime/src/error.rs
Normal file
|
|
@ -0,0 +1,554 @@
|
|||
//! Cross-ecosystem error type for tool execution.
|
||||
//!
|
||||
//! `ToolError` is a struct with a `kind` discriminator and a tool-provided
|
||||
//! `detail` string. The `detail` is the model-facing message — tools MUST
|
||||
//! provide a human-readable explanation of what went wrong, since this text
|
||||
//! is sent back to the model to inform its next action.
|
||||
//!
|
||||
//! The wire boundary is bridged by `From<ToolError> for ToolErrorWire`.
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use xai_tool_protocol::{ToolErrorWire, ToolId};
|
||||
|
||||
/// Discriminator for tool errors.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub enum ToolErrorKind {
|
||||
/// The tool has no implementation for the requested operation.
|
||||
NotImplemented,
|
||||
/// Inputs failed validation.
|
||||
InvalidArguments,
|
||||
/// No tool registered under the given id.
|
||||
NotFound,
|
||||
/// Caller lacks required permissions (403-shaped).
|
||||
PermissionDenied,
|
||||
/// Authentication failed (401-shaped).
|
||||
Unauthorized,
|
||||
/// The tool ran past its time budget.
|
||||
Timeout,
|
||||
/// The caller cancelled the tool call.
|
||||
Cancelled,
|
||||
/// Rate limit exceeded.
|
||||
RateLimited,
|
||||
/// The caller's usage pool / billing balance is exhausted (out
|
||||
/// of credits). Payment-required-shaped; distinct from
|
||||
/// `RateLimited` so the surface can show "out of credits"
|
||||
/// rather than "try again later".
|
||||
UsagePoolExhausted,
|
||||
/// The caller hit a usage limit with no balance verdict behind it
|
||||
/// (the balance gate was skipped/dormant and the non-billable
|
||||
/// allowance ran out). Payment-required-shaped, but distinct from
|
||||
/// `UsagePoolExhausted` (an explicit out-of-balance verdict) so the
|
||||
/// surface can show a "usage limit reached" message.
|
||||
UsageLimitReached,
|
||||
/// The billing global rate limiter shed this request (transient
|
||||
/// load shed). Distinct from `RateLimited` (per-user / per-message
|
||||
/// quota) so the surface can render a billing-specific
|
||||
/// "try again later" with a retry hint; the `retry_after_secs`
|
||||
/// hint, when known, rides in `ToolError::details`. Named to match
|
||||
/// the chat surface's `global_rate_limit` typed error.
|
||||
GlobalRateLimit,
|
||||
/// The caller hit their per-user concurrency cap (too many media
|
||||
/// generations already in flight). Transient — retry once one
|
||||
/// finishes. Distinct from `GlobalRateLimit` (a shared-backend load
|
||||
/// shed) so the surface can tailor a "too many in progress" message.
|
||||
/// Named to match the chat surface's `concurrency_limit` typed error.
|
||||
ConcurrencyLimit,
|
||||
/// Upstream service unavailable.
|
||||
ServiceUnavailable,
|
||||
/// Network-level failure.
|
||||
NetworkError,
|
||||
/// Tool body returned an error.
|
||||
Execution,
|
||||
/// Requested behavior version not supported.
|
||||
BehaviorVersionUnsupported,
|
||||
/// Render-card budget exceeded.
|
||||
RenderLimited,
|
||||
/// Terminal subprocess failure.
|
||||
TerminalError,
|
||||
/// Forward-compat catch-all.
|
||||
Custom,
|
||||
}
|
||||
|
||||
impl ToolErrorKind {
|
||||
/// Snake-case identifier for metrics / logs.
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::NotImplemented => "not_implemented",
|
||||
Self::InvalidArguments => "invalid_arguments",
|
||||
Self::NotFound => "not_found",
|
||||
Self::PermissionDenied => "permission_denied",
|
||||
Self::Unauthorized => "unauthorized",
|
||||
Self::Timeout => "timeout",
|
||||
Self::Cancelled => "cancelled",
|
||||
Self::RateLimited => "rate_limited",
|
||||
Self::UsagePoolExhausted => "usage_pool_exhausted",
|
||||
Self::UsageLimitReached => "usage_limit_reached",
|
||||
Self::GlobalRateLimit => "global_rate_limit",
|
||||
Self::ConcurrencyLimit => "concurrency_limit",
|
||||
Self::ServiceUnavailable => "service_unavailable",
|
||||
Self::NetworkError => "network_error",
|
||||
Self::Execution => "execution",
|
||||
Self::BehaviorVersionUnsupported => "behavior_version_unsupported",
|
||||
Self::RenderLimited => "render_limited",
|
||||
Self::TerminalError => "terminal_error",
|
||||
Self::Custom => "custom",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Cross-ecosystem error type for tool execution.
|
||||
///
|
||||
/// Every error carries:
|
||||
/// - `kind` — the machine-readable discriminator
|
||||
/// - `detail` — the model/user-facing message that tools MUST provide
|
||||
/// - `source` — optional causal chain for debugging (not sent to the model)
|
||||
/// - `details` — optional structured metadata (JSON Schema validation
|
||||
/// report, retry_after hints, etc.)
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub struct ToolError {
|
||||
pub kind: ToolErrorKind,
|
||||
/// Human-readable message provided by the tool. This is sent back to
|
||||
/// the model so it can understand what went wrong and adjust its next
|
||||
/// action. Tools MUST make this specific and actionable.
|
||||
pub detail: String,
|
||||
/// Optional causal chain for developer debugging. NOT sent to the model.
|
||||
#[serde(skip)]
|
||||
source: Option<anyhow::Error>,
|
||||
/// Optional structured metadata (e.g. per-field validation errors,
|
||||
/// `retry_after` hints, `tool_id`, `card_id`).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub details: Option<Value>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for ToolError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let mut d = f.debug_struct("ToolError");
|
||||
d.field("kind", &self.kind);
|
||||
d.field("detail", &self.detail);
|
||||
if let Some(ref source) = self.source {
|
||||
d.field("source", &format!("{source:#}"));
|
||||
}
|
||||
if let Some(ref details) = self.details {
|
||||
d.field("details", details);
|
||||
}
|
||||
d.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for ToolError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}", self.detail)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ToolError {
|
||||
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
|
||||
self.source.as_ref().map(|e| e.as_ref() as &_)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Constructors — one per kind for ergonomic tool code
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
impl ToolError {
|
||||
/// Core constructor. All other constructors delegate here.
|
||||
pub fn new(kind: ToolErrorKind, detail: impl Into<String>) -> Self {
|
||||
Self {
|
||||
kind,
|
||||
detail: detail.into(),
|
||||
source: None,
|
||||
details: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach structured metadata.
|
||||
pub fn with_details(mut self, details: Value) -> Self {
|
||||
self.details = Some(details);
|
||||
self
|
||||
}
|
||||
|
||||
/// Attach a causal error chain (for developer logs, not sent to model).
|
||||
pub fn with_source(mut self, source: impl Into<anyhow::Error>) -> Self {
|
||||
self.source = Some(source.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn not_implemented(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::NotImplemented, detail)
|
||||
}
|
||||
|
||||
pub fn invalid_arguments(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::InvalidArguments, detail)
|
||||
}
|
||||
|
||||
pub fn not_found(tool_id: ToolId, detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::NotFound, detail)
|
||||
.with_details(serde_json::json!({ "tool_id": tool_id.as_str() }))
|
||||
}
|
||||
|
||||
pub fn permission_denied(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::PermissionDenied, detail)
|
||||
}
|
||||
|
||||
pub fn unauthorized(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::Unauthorized, detail)
|
||||
}
|
||||
|
||||
pub fn timeout(tool_id: ToolId, detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::Timeout, detail)
|
||||
.with_details(serde_json::json!({ "tool_id": tool_id.as_str() }))
|
||||
}
|
||||
|
||||
pub fn cancelled(tool_id: ToolId, detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::Cancelled, detail)
|
||||
.with_details(serde_json::json!({ "tool_id": tool_id.as_str() }))
|
||||
}
|
||||
|
||||
pub fn rate_limited(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::RateLimited, detail)
|
||||
}
|
||||
|
||||
pub fn usage_pool_exhausted(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::UsagePoolExhausted, detail)
|
||||
}
|
||||
|
||||
pub fn usage_limit_reached(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::UsageLimitReached, detail)
|
||||
}
|
||||
|
||||
pub fn global_rate_limit(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::GlobalRateLimit, detail)
|
||||
}
|
||||
|
||||
pub fn concurrency_limit(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::ConcurrencyLimit, detail)
|
||||
}
|
||||
|
||||
pub fn service_unavailable(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::ServiceUnavailable, detail)
|
||||
}
|
||||
|
||||
pub fn network_error(detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::NetworkError, detail)
|
||||
}
|
||||
|
||||
pub fn execution(tool_id: ToolId, detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::Execution, detail)
|
||||
.with_details(serde_json::json!({ "tool_id": tool_id.as_str() }))
|
||||
}
|
||||
|
||||
pub fn terminal_error(tool_id: ToolId, detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::TerminalError, detail)
|
||||
.with_details(serde_json::json!({ "tool_id": tool_id.as_str() }))
|
||||
}
|
||||
|
||||
pub fn custom(code: impl Into<String>, detail: impl Into<String>) -> Self {
|
||||
Self::new(ToolErrorKind::Custom, detail)
|
||||
.with_details(serde_json::json!({ "code": code.into() }))
|
||||
}
|
||||
|
||||
/// Snake-case identifier for the kind. Delegates to
|
||||
/// [`ToolErrorKind::as_str`].
|
||||
pub fn variant_name(&self) -> &'static str {
|
||||
self.kind.as_str()
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// From impls
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
impl From<serde_json::Error> for ToolError {
|
||||
fn from(value: serde_json::Error) -> Self {
|
||||
Self::invalid_arguments(value.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Wire bridge
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Carry a [`ToolError`]'s structured `details` onto a `Custom` wire variant
|
||||
/// while keeping the round-trip recognizable: the decoder
|
||||
/// (`tool_error_from_wire`) replaces the `{"code": <subcode>}` object that
|
||||
/// `ToolError::custom` installs with the wire `details` verbatim, so the
|
||||
/// subcode is merged into object-shaped details (without clobbering an
|
||||
/// existing `code` key). Non-object details pass through unchanged.
|
||||
fn custom_details_with_code(details: Option<Value>, code: &str) -> Option<Value> {
|
||||
match details {
|
||||
Some(Value::Object(mut map)) => {
|
||||
map.entry("code")
|
||||
.or_insert_with(|| Value::String(code.to_owned()));
|
||||
Some(Value::Object(map))
|
||||
}
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ToolError> for ToolErrorWire {
|
||||
fn from(err: ToolError) -> Self {
|
||||
// Extract structured fields from `details` when the wire shape needs
|
||||
// them. The `detail` string is always the model-facing message.
|
||||
let details_val = err.details.as_ref();
|
||||
|
||||
match err.kind {
|
||||
ToolErrorKind::NotImplemented => Self::Custom {
|
||||
subcode: "not_implemented".to_owned(),
|
||||
message: err.detail,
|
||||
details: custom_details_with_code(err.details, "not_implemented"),
|
||||
},
|
||||
ToolErrorKind::InvalidArguments => Self::InvalidArguments {
|
||||
message: err.detail,
|
||||
details: err.details,
|
||||
},
|
||||
ToolErrorKind::NotFound => {
|
||||
let tool_id = details_val
|
||||
.and_then(|d| d.get("tool_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|s| ToolId::new(s).ok())
|
||||
.unwrap_or_else(|| ToolId::new("unknown").unwrap());
|
||||
Self::ToolNotFound { tool_id }
|
||||
}
|
||||
ToolErrorKind::PermissionDenied => Self::PermissionDenied { reason: err.detail },
|
||||
ToolErrorKind::Unauthorized => Self::Custom {
|
||||
subcode: "unauthorized".to_owned(),
|
||||
message: err.detail,
|
||||
details: custom_details_with_code(err.details, "unauthorized"),
|
||||
},
|
||||
ToolErrorKind::Timeout => {
|
||||
let tool_id = details_val
|
||||
.and_then(|d| d.get("tool_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|s| ToolId::new(s).ok())
|
||||
.unwrap_or_else(|| ToolId::new("unknown").unwrap());
|
||||
let elapsed_ms = details_val
|
||||
.and_then(|d| d.get("elapsed_ms"))
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0);
|
||||
Self::Timeout {
|
||||
tool_id,
|
||||
elapsed_ms,
|
||||
}
|
||||
}
|
||||
ToolErrorKind::Cancelled => {
|
||||
let tool_id = details_val
|
||||
.and_then(|d| d.get("tool_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|s| ToolId::new(s).ok())
|
||||
.unwrap_or_else(|| ToolId::new("unknown").unwrap());
|
||||
Self::Cancelled { tool_id }
|
||||
}
|
||||
ToolErrorKind::RateLimited => Self::Custom {
|
||||
subcode: "rate_limited".to_owned(),
|
||||
message: err.detail,
|
||||
details: custom_details_with_code(err.details, "rate_limited"),
|
||||
},
|
||||
ToolErrorKind::UsagePoolExhausted => Self::Custom {
|
||||
subcode: "usage_pool_exhausted".to_owned(),
|
||||
message: err.detail,
|
||||
details: custom_details_with_code(err.details, "usage_pool_exhausted"),
|
||||
},
|
||||
ToolErrorKind::UsageLimitReached => Self::Custom {
|
||||
subcode: "usage_limit_reached".to_owned(),
|
||||
message: err.detail,
|
||||
details: custom_details_with_code(err.details, "usage_limit_reached"),
|
||||
},
|
||||
ToolErrorKind::GlobalRateLimit => Self::Custom {
|
||||
subcode: "global_rate_limit".to_owned(),
|
||||
message: err.detail,
|
||||
details: custom_details_with_code(err.details, "global_rate_limit"),
|
||||
},
|
||||
ToolErrorKind::ConcurrencyLimit => Self::Custom {
|
||||
subcode: "concurrency_limit".to_owned(),
|
||||
message: err.detail,
|
||||
details: custom_details_with_code(err.details, "concurrency_limit"),
|
||||
},
|
||||
ToolErrorKind::ServiceUnavailable => Self::Custom {
|
||||
subcode: "service_unavailable".to_owned(),
|
||||
message: err.detail,
|
||||
details: custom_details_with_code(err.details, "service_unavailable"),
|
||||
},
|
||||
ToolErrorKind::NetworkError => Self::Custom {
|
||||
subcode: "network_error".to_owned(),
|
||||
message: err.detail,
|
||||
details: custom_details_with_code(err.details, "network_error"),
|
||||
},
|
||||
ToolErrorKind::Execution => {
|
||||
let tool_id = details_val
|
||||
.and_then(|d| d.get("tool_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|s| ToolId::new(s).ok())
|
||||
.unwrap_or_else(|| ToolId::new("unknown").unwrap());
|
||||
Self::Execution {
|
||||
tool_id,
|
||||
message: err.detail,
|
||||
}
|
||||
}
|
||||
ToolErrorKind::BehaviorVersionUnsupported => {
|
||||
let tool_id = details_val
|
||||
.and_then(|d| d.get("tool_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|s| ToolId::new(s).ok())
|
||||
.unwrap_or_else(|| ToolId::new("unknown").unwrap());
|
||||
let requested = details_val
|
||||
.and_then(|d| d.get("requested"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_owned();
|
||||
Self::BehaviorVersionUnsupported { tool_id, requested }
|
||||
}
|
||||
ToolErrorKind::RenderLimited => {
|
||||
let tool_id = details_val
|
||||
.and_then(|d| d.get("tool_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|s| ToolId::new(s).ok())
|
||||
.unwrap_or_else(|| ToolId::new("unknown").unwrap());
|
||||
let card_id = details_val
|
||||
.and_then(|d| d.get("card_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(str::to_owned);
|
||||
Self::RenderLimited {
|
||||
tool_id,
|
||||
card_id,
|
||||
reason: err.detail,
|
||||
}
|
||||
}
|
||||
ToolErrorKind::TerminalError => {
|
||||
let tool_id = details_val
|
||||
.and_then(|d| d.get("tool_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|s| ToolId::new(s).ok())
|
||||
.unwrap_or_else(|| ToolId::new("unknown").unwrap());
|
||||
Self::TerminalError {
|
||||
tool_id,
|
||||
message: err.detail,
|
||||
}
|
||||
}
|
||||
ToolErrorKind::Custom => {
|
||||
let subcode = details_val
|
||||
.and_then(|d| d.get("code"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("custom")
|
||||
.to_owned();
|
||||
Self::Custom {
|
||||
subcode,
|
||||
message: err.detail,
|
||||
details: err.details,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod wire_bridge_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn service_unavailable_details_survive_wire_projection() {
|
||||
// Structured details used to be dropped (`details: None`) for the
|
||||
// Custom-mapped kinds; they must now ride the wire with the subcode
|
||||
// merged in so recognizers keying on `details.code` keep working.
|
||||
let err = ToolError::service_unavailable("sandbox not ready")
|
||||
.with_details(serde_json::json!({ "retry_after_ms": 1500 }));
|
||||
let wire = ToolErrorWire::from(err);
|
||||
let ToolErrorWire::Custom {
|
||||
subcode,
|
||||
message,
|
||||
details,
|
||||
} = wire
|
||||
else {
|
||||
panic!("expected Custom");
|
||||
};
|
||||
assert_eq!(subcode, "service_unavailable");
|
||||
assert_eq!(message, "sandbox not ready");
|
||||
let details = details.expect("details preserved");
|
||||
assert_eq!(
|
||||
details.get("retry_after_ms").and_then(|v| v.as_u64()),
|
||||
Some(1500)
|
||||
);
|
||||
assert_eq!(
|
||||
details.get("code").and_then(|v| v.as_str()),
|
||||
Some("service_unavailable"),
|
||||
"subcode merged into details for round-trip recognizability",
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rate_limit_and_usage_kinds_merge_subcode_uniformly() {
|
||||
// Same property as service_unavailable, applied to every
|
||||
// Custom-mapped kind: object details without a `code` key gain the
|
||||
// subcode, so decode-side recognizers keying on `details.code` can
|
||||
// still classify the error.
|
||||
let cases: [(ToolError, &str); 5] = [
|
||||
(ToolError::rate_limited("slow down"), "rate_limited"),
|
||||
(
|
||||
ToolError::usage_pool_exhausted("pool empty"),
|
||||
"usage_pool_exhausted",
|
||||
),
|
||||
(
|
||||
ToolError::usage_limit_reached("limit hit"),
|
||||
"usage_limit_reached",
|
||||
),
|
||||
(
|
||||
ToolError::global_rate_limit("global limit"),
|
||||
"global_rate_limit",
|
||||
),
|
||||
(
|
||||
ToolError::concurrency_limit("too many in flight"),
|
||||
"concurrency_limit",
|
||||
),
|
||||
];
|
||||
for (err, subcode) in cases {
|
||||
let err = err.with_details(serde_json::json!({ "retry_after_ms": 250 }));
|
||||
let ToolErrorWire::Custom {
|
||||
subcode: got,
|
||||
details,
|
||||
..
|
||||
} = ToolErrorWire::from(err)
|
||||
else {
|
||||
panic!("expected Custom for {subcode}");
|
||||
};
|
||||
assert_eq!(got, subcode);
|
||||
let details = details.expect("details preserved");
|
||||
assert_eq!(
|
||||
details.get("code").and_then(|v| v.as_str()),
|
||||
Some(subcode),
|
||||
"subcode merged for {subcode}",
|
||||
);
|
||||
assert_eq!(
|
||||
details.get("retry_after_ms").and_then(|v| v.as_u64()),
|
||||
Some(250),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_details_with_code_does_not_clobber_existing_code() {
|
||||
let merged = custom_details_with_code(
|
||||
Some(serde_json::json!({ "code": "workspace_unavailable", "retryable": true })),
|
||||
"service_unavailable",
|
||||
)
|
||||
.expect("details kept");
|
||||
assert_eq!(
|
||||
merged.get("code").and_then(|v| v.as_str()),
|
||||
Some("workspace_unavailable"),
|
||||
"an existing code key must win",
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_details_with_code_passes_none_and_non_objects_through() {
|
||||
assert_eq!(custom_details_with_code(None, "network_error"), None);
|
||||
let arr = serde_json::json!([1, 2, 3]);
|
||||
assert_eq!(
|
||||
custom_details_with_code(Some(arr.clone()), "network_error"),
|
||||
Some(arr),
|
||||
);
|
||||
}
|
||||
}
|
||||
45
crates/common/xai-tool-runtime/src/lib.rs
Normal file
45
crates/common/xai-tool-runtime/src/lib.rs
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
//! xAI Computer Hub — unified runtime contract.
|
||||
//!
|
||||
//! Single home for the `Tool` trait, `ToolDispatch`, `ToolError`,
|
||||
//! `ToolNotification`, `ToolSearchIndex`, `ToolCallContext`, `ToolStream`,
|
||||
//! and the helper constructors that build well-formed streams. Adapters
|
||||
//! for individual tool sources re-export from here so every tool author
|
||||
//! sees the same surface.
|
||||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
pub mod context;
|
||||
pub mod dispatch;
|
||||
pub mod error;
|
||||
pub mod notification;
|
||||
pub mod render;
|
||||
pub mod search;
|
||||
pub mod streaming;
|
||||
pub mod tool;
|
||||
|
||||
pub use context::{
|
||||
BehaviorVersion, Cancellation, Cwd, ListToolsContext, SessionContext, ToolCallContext,
|
||||
TraceContext, TypedExtensions, WorkspaceBindMetadata, WorkspaceViewerContext,
|
||||
};
|
||||
pub use dispatch::ToolDispatch;
|
||||
pub use error::{ToolError, ToolErrorKind};
|
||||
pub use notification::{
|
||||
BashExecutionBackgrounded, BashExecutionComplete, BashExecutionFailed, BashExecutionTimeout,
|
||||
BashNotificationBase, BashOutputChunk, FileRead, FileWritten, LspServerCrashed,
|
||||
LspServerFailed, LspServerReady, LspServerRetrying, LspServerStarting, MonitorEvent,
|
||||
PlanModeEntered, PlanModeExited, ScheduledTaskCreated, ScheduledTaskFired,
|
||||
ScheduledTaskRemoved, TaskKind, TaskSnapshot, ToolNotification, ToolNotificationHandle,
|
||||
UserQuestionAsked,
|
||||
};
|
||||
pub use render::{
|
||||
ModelOutputExtractor, ToolChatCompletion, ToolChatCompletionResponse, ToolCodeExecutionResult,
|
||||
ToolOutput, ToolStreamError, extract_content_blocks, extractor_for,
|
||||
};
|
||||
pub use search::{SearchSnapshot, ServerSummary, ToolIndex, ToolSearchIndex, ToolSearchResult};
|
||||
pub use streaming::{PartialResultPayload, stream_chunk};
|
||||
pub use tool::{
|
||||
ArcTool, ArcToolFamily, ContentBlock, Tool, ToolDyn, ToolFamily, ToolProgress, ToolStream,
|
||||
ToolStreamItem, ToolVariant, TypedToolOutput, terminal_only, with_progress,
|
||||
};
|
||||
|
||||
pub use xai_tool_protocol::{StreamingSpec, ToolCallId, ToolCapabilities, ToolId, ToolScope};
|
||||
530
crates/common/xai-tool-runtime/src/notification.rs
Normal file
530
crates/common/xai-tool-runtime/src/notification.rs
Normal file
|
|
@ -0,0 +1,530 @@
|
|||
//! Tool notifications — typed messages a running tool emits to subscribers
|
||||
//! (TUI, gateway, audit log, ...) for live visibility into execution.
|
||||
//!
|
||||
//! The enum and its payload structs use unconditional serde derives so wire
|
||||
//! adapters can serialise them without enabling additional features.
|
||||
//!
|
||||
//! Each `ToolNotification` variant has a parallel `send_*` convenience on
|
||||
//! [`ToolNotificationHandle`]. The two surfaces are kept in lockstep — when
|
||||
//! adding a variant here, add the `send_*` constructor too.
|
||||
//!
|
||||
//! The handle is built on `futures::channel::mpsc` so it is runtime-neutral:
|
||||
//! the trait crate doesn't pin a particular async executor on its
|
||||
//! consumers.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use futures::channel::mpsc;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Common fields shared by every bash notification variant. Hoisted into a
|
||||
/// dedicated struct so the variants stay in lockstep on tool_call_id /
|
||||
/// command / output / cwd, and so payload-shape changes only need to be
|
||||
/// made once.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BashNotificationBase {
|
||||
/// Tool call id, used to correlate with the originating tool call.
|
||||
pub tool_call_id: String,
|
||||
|
||||
/// The command being executed.
|
||||
pub command: String,
|
||||
|
||||
/// Captured output bytes. May be truncated; use `output_lossy` for a
|
||||
/// `String` rendering that handles invalid UTF-8.
|
||||
pub output: Vec<u8>,
|
||||
|
||||
/// Total bytes received before any truncation.
|
||||
pub total_bytes: usize,
|
||||
|
||||
/// Whether `output` was truncated to fit a size cap.
|
||||
pub truncated: bool,
|
||||
|
||||
/// Working directory the command ran in.
|
||||
pub cwd: PathBuf,
|
||||
}
|
||||
|
||||
impl BashNotificationBase {
|
||||
/// Lossy UTF-8 rendering of `output`. Invalid bytes become U+FFFD.
|
||||
pub fn output_lossy(&self) -> std::borrow::Cow<'_, str> {
|
||||
String::from_utf8_lossy(&self.output)
|
||||
}
|
||||
}
|
||||
|
||||
/// Incremental output chunk streamed during a bash command. Sent
|
||||
/// periodically while the process is still running.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BashOutputChunk {
|
||||
#[serde(flatten)]
|
||||
pub base: BashNotificationBase,
|
||||
}
|
||||
|
||||
/// Sent when a bash process exits. Carries the exit status (or the killing
|
||||
/// signal name when the process didn't exit normally).
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BashExecutionComplete {
|
||||
#[serde(flatten)]
|
||||
pub base: BashNotificationBase,
|
||||
|
||||
/// `Some(code)` for a normal exit; `None` when the process was killed
|
||||
/// by a signal before reaching `exit(2)`.
|
||||
pub exit_code: Option<i32>,
|
||||
|
||||
/// Signal that terminated the process (e.g. `"SIGKILL"`). `None` when
|
||||
/// the process exited normally.
|
||||
pub signal: Option<String>,
|
||||
}
|
||||
|
||||
impl BashExecutionComplete {
|
||||
/// `true` when termination was triggered by a signal.
|
||||
pub fn was_signaled(&self) -> bool {
|
||||
self.signal.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
/// Sent when a bash command exceeded its configured timeout and was
|
||||
/// killed.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BashExecutionTimeout {
|
||||
#[serde(flatten)]
|
||||
pub base: BashNotificationBase,
|
||||
|
||||
/// Wall time the command ran for before being killed.
|
||||
pub elapsed: Duration,
|
||||
|
||||
/// Configured timeout that was exceeded.
|
||||
pub timeout: Duration,
|
||||
}
|
||||
|
||||
/// Sent when a foreground bash command was moved to the background. The
|
||||
/// process keeps running; a downstream task monitor emits the eventual
|
||||
/// [`BashExecutionComplete`].
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BashExecutionBackgrounded {
|
||||
#[serde(flatten)]
|
||||
pub base: BashNotificationBase,
|
||||
|
||||
/// File the full output stream is being written to. Background tasks
|
||||
/// always tee to disk so consumers can fetch the rest later.
|
||||
pub output_file: PathBuf,
|
||||
|
||||
/// Background task registry id. Distinct from `base.tool_call_id`:
|
||||
/// the task id is generated when backgrounding, the tool call id was
|
||||
/// assigned when the originating tool was invoked.
|
||||
pub task_id: String,
|
||||
}
|
||||
|
||||
/// Sent when a bash command failed to spawn. Distinct from
|
||||
/// [`BashExecutionComplete`] with a non-zero `exit_code` because the
|
||||
/// process never started.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BashExecutionFailed {
|
||||
pub tool_call_id: String,
|
||||
pub command: String,
|
||||
pub cwd: PathBuf,
|
||||
/// Error message describing the spawn / IO failure.
|
||||
pub error: String,
|
||||
}
|
||||
|
||||
/// Emitted when a tool reads a file. Subscribers use this for state
|
||||
/// snapshotting (rewind, audit) of accessed files.
|
||||
///
|
||||
/// **Reserved for a future `ToolNotification::FileRead` variant.** The
|
||||
/// struct is kept in the public API so adapters can construct it ahead of
|
||||
/// time, but it is not currently dispatched by any
|
||||
/// [`ToolNotificationHandle`] helper. Adding the enum variant here is a
|
||||
/// breaking change for exhaustive `match` consumers, so the variant is
|
||||
/// deferred until a downstream crate has a real consumer wired up.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct FileRead {
|
||||
pub tool_call_id: String,
|
||||
/// Absolute filesystem path of the file that was read.
|
||||
pub absolute_path: PathBuf,
|
||||
}
|
||||
|
||||
/// Emitted when a tool writes a file. Carries the full pre- and post-edit
|
||||
/// content so subscribers can rewind without re-reading the disk.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct FileWritten {
|
||||
pub tool_call_id: String,
|
||||
/// Absolute filesystem path of the file that was written.
|
||||
pub absolute_path: PathBuf,
|
||||
/// Full file content after the write.
|
||||
pub content: String,
|
||||
/// Full file content before the write. `None` for a fresh file.
|
||||
pub previous_content: Option<String>,
|
||||
/// Whether the write created a new file.
|
||||
pub is_new_file: bool,
|
||||
}
|
||||
|
||||
/// Sent when the agent transitions into plan mode. Subscribers use this to
|
||||
/// enforce the read-only constraint and switch UI affordances.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct PlanModeEntered {
|
||||
pub tool_call_id: String,
|
||||
}
|
||||
|
||||
/// Sent when the agent transitions out of plan mode. Carries the plan
|
||||
/// document so subscribers can present it for approval without an extra
|
||||
/// file read.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct PlanModeExited {
|
||||
pub tool_call_id: String,
|
||||
/// Plan content as captured at exit time. `None` when the plan file
|
||||
/// did not exist or was empty.
|
||||
pub plan_content: Option<String>,
|
||||
/// Path the plan file lives at.
|
||||
pub plan_file_path: String,
|
||||
}
|
||||
|
||||
/// Sent when the agent issues a structured question to the user.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UserQuestionAsked {
|
||||
pub tool_call_id: String,
|
||||
/// Serialised question payload. Subscribers render it directly; the
|
||||
/// runtime does not introspect its shape.
|
||||
pub questions_json: serde_json::Value,
|
||||
}
|
||||
|
||||
/// LSP server is being spawned and is waiting for the initialise
|
||||
/// handshake.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LspServerStarting {
|
||||
pub server_name: String,
|
||||
pub command: String,
|
||||
}
|
||||
|
||||
/// LSP server completed initialisation and is ready to serve requests.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LspServerReady {
|
||||
pub server_name: String,
|
||||
}
|
||||
|
||||
/// LSP server process died unexpectedly.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LspServerCrashed {
|
||||
pub server_name: String,
|
||||
}
|
||||
|
||||
/// LSP server is being retried after a crash. Carries the retry attempt
|
||||
/// count and computed backoff so subscribers can render progress.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LspServerRetrying {
|
||||
pub server_name: String,
|
||||
pub attempt: u32,
|
||||
pub max_restarts: u32,
|
||||
pub backoff_ms: u64,
|
||||
}
|
||||
|
||||
/// LSP server is permanently dead. Either init failed (`attempts == 0`)
|
||||
/// or the configured retry budget was exhausted.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LspServerFailed {
|
||||
pub server_name: String,
|
||||
pub error: String,
|
||||
/// `0` for init failure, `> 0` when the retry budget was exhausted.
|
||||
pub attempts: u32,
|
||||
}
|
||||
|
||||
/// Sent when a scheduled task fired and its prompt should be executed.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ScheduledTaskFired {
|
||||
pub task_id: String,
|
||||
pub prompt: String,
|
||||
pub human_schedule: String,
|
||||
/// RFC 3339 timestamp of the next fire, when the task is recurring.
|
||||
pub next_fire_at: Option<String>,
|
||||
}
|
||||
|
||||
/// Sent when a scheduled task is removed (deleted, expired, or one-shot
|
||||
/// completed).
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ScheduledTaskRemoved {
|
||||
pub task_id: String,
|
||||
}
|
||||
|
||||
/// Sent when a scheduled task is created.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ScheduledTaskCreated {
|
||||
pub task_id: String,
|
||||
pub prompt: String,
|
||||
pub human_schedule: String,
|
||||
/// RFC 3339 timestamp of the upcoming first fire.
|
||||
pub next_fire_at: Option<String>,
|
||||
}
|
||||
|
||||
/// Streaming event from a Monitor tool background process. Each event is
|
||||
/// already XML-wrapped for direct injection into the conversation; the
|
||||
/// raw text is preserved for plain-text consumers.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct MonitorEvent {
|
||||
pub task_id: String,
|
||||
pub description: String,
|
||||
/// XML-wrapped event text, ready for conversation injection.
|
||||
pub event_text: String,
|
||||
/// Raw text without XML wrapping.
|
||||
pub raw_text: String,
|
||||
}
|
||||
|
||||
/// Snapshot of a background task's state. Identical shape to the Grok
|
||||
/// Build `TaskSnapshot` so subscribers can decode without per-source
|
||||
/// adapters.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct TaskSnapshot {
|
||||
pub task_id: String,
|
||||
/// Actual command that was executed (may be wrapped by an isolation
|
||||
/// harness).
|
||||
pub command: String,
|
||||
/// Original user-provided command before isolation wrapping. When
|
||||
/// present, model- and user-facing surfaces should prefer it over
|
||||
/// `command`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub display_command: Option<String>,
|
||||
pub cwd: String,
|
||||
pub start_time: SystemTime,
|
||||
pub end_time: Option<SystemTime>,
|
||||
pub output: String,
|
||||
pub output_file: PathBuf,
|
||||
pub truncated: bool,
|
||||
pub exit_code: Option<i32>,
|
||||
pub signal: Option<String>,
|
||||
pub completed: bool,
|
||||
/// Distinguishes monitor tasks from regular bash tasks.
|
||||
#[serde(default)]
|
||||
pub kind: TaskKind,
|
||||
}
|
||||
|
||||
impl TaskSnapshot {
|
||||
/// Wall-time duration in seconds. Falls back to `now` for tasks that
|
||||
/// haven't completed.
|
||||
pub fn duration_secs(&self) -> f64 {
|
||||
let end = self.end_time.unwrap_or_else(SystemTime::now);
|
||||
end.duration_since(self.start_time)
|
||||
.map(|d| d.as_secs_f64())
|
||||
.unwrap_or(0.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Distinguishes background-task kinds.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TaskKind {
|
||||
/// Regular bash command.
|
||||
#[default]
|
||||
Bash,
|
||||
/// Monitor tool — streams stdout events with rate limiting.
|
||||
Monitor,
|
||||
}
|
||||
|
||||
/// A typed notification a tool emits during or after execution.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum ToolNotification {
|
||||
BashOutputChunk(BashOutputChunk),
|
||||
BashExecutionComplete(BashExecutionComplete),
|
||||
BashExecutionTimeout(BashExecutionTimeout),
|
||||
BashExecutionBackgrounded(BashExecutionBackgrounded),
|
||||
BashExecutionFailed(BashExecutionFailed),
|
||||
FileWritten(FileWritten),
|
||||
TaskCompleted(TaskSnapshot),
|
||||
PlanModeEntered(PlanModeEntered),
|
||||
PlanModeExited(PlanModeExited),
|
||||
UserQuestionAsked(UserQuestionAsked),
|
||||
LspServerStarting(LspServerStarting),
|
||||
LspServerReady(LspServerReady),
|
||||
LspServerCrashed(LspServerCrashed),
|
||||
LspServerRetrying(LspServerRetrying),
|
||||
LspServerFailed(LspServerFailed),
|
||||
ScheduledTaskFired(ScheduledTaskFired),
|
||||
ScheduledTaskRemoved(ScheduledTaskRemoved),
|
||||
ScheduledTaskCreated(ScheduledTaskCreated),
|
||||
MonitorEvent(MonitorEvent),
|
||||
}
|
||||
|
||||
impl ToolNotification {
|
||||
/// Stable `PascalCase` name of the active variant. Mirrors the serde
|
||||
/// `tag = "type"` discriminator used on the wire.
|
||||
pub fn variant_name(&self) -> &'static str {
|
||||
match self {
|
||||
Self::BashOutputChunk(_) => "BashOutputChunk",
|
||||
Self::BashExecutionComplete(_) => "BashExecutionComplete",
|
||||
Self::BashExecutionTimeout(_) => "BashExecutionTimeout",
|
||||
Self::BashExecutionBackgrounded(_) => "BashExecutionBackgrounded",
|
||||
Self::BashExecutionFailed(_) => "BashExecutionFailed",
|
||||
Self::FileWritten(_) => "FileWritten",
|
||||
Self::TaskCompleted(_) => "TaskCompleted",
|
||||
Self::PlanModeEntered(_) => "PlanModeEntered",
|
||||
Self::PlanModeExited(_) => "PlanModeExited",
|
||||
Self::UserQuestionAsked(_) => "UserQuestionAsked",
|
||||
Self::LspServerStarting(_) => "LspServerStarting",
|
||||
Self::LspServerReady(_) => "LspServerReady",
|
||||
Self::LspServerCrashed(_) => "LspServerCrashed",
|
||||
Self::LspServerRetrying(_) => "LspServerRetrying",
|
||||
Self::LspServerFailed(_) => "LspServerFailed",
|
||||
Self::ScheduledTaskFired(_) => "ScheduledTaskFired",
|
||||
Self::ScheduledTaskRemoved(_) => "ScheduledTaskRemoved",
|
||||
Self::ScheduledTaskCreated(_) => "ScheduledTaskCreated",
|
||||
Self::MonitorEvent(_) => "MonitorEvent",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Cloneable handle for emitting [`ToolNotification`]s.
|
||||
///
|
||||
/// Built on `futures::channel::mpsc::UnboundedSender` so the sender side
|
||||
/// is runtime-neutral — the trait crate does not pin tokio (or any other
|
||||
/// executor) on its callers. Sends are non-blocking and best-effort:
|
||||
/// errors (a closed receiver) are silently dropped, matching the
|
||||
/// established convention for fire-and-forget notification streams.
|
||||
#[derive(Clone)]
|
||||
pub struct ToolNotificationHandle {
|
||||
sender: mpsc::UnboundedSender<ToolNotification>,
|
||||
}
|
||||
|
||||
impl ToolNotificationHandle {
|
||||
/// Wrap a sender obtained elsewhere.
|
||||
pub fn new(sender: mpsc::UnboundedSender<ToolNotification>) -> Self {
|
||||
Self { sender }
|
||||
}
|
||||
|
||||
/// Alias for [`Self::new`]. Tests and consumers that own the receiver
|
||||
/// half use this for symmetry.
|
||||
pub fn from_sender(sender: mpsc::UnboundedSender<ToolNotification>) -> Self {
|
||||
Self { sender }
|
||||
}
|
||||
|
||||
/// Build both halves of a fresh channel and return them paired.
|
||||
pub fn channel() -> (Self, mpsc::UnboundedReceiver<ToolNotification>) {
|
||||
let (sender, receiver) = mpsc::unbounded();
|
||||
(Self { sender }, receiver)
|
||||
}
|
||||
|
||||
/// Build a handle whose sends are silently dropped. Use for callers
|
||||
/// that don't care about notifications (smoke tests, dry-run
|
||||
/// utilities). NOT a sensible default for production paths — the
|
||||
/// silent-drop behaviour makes notification bugs invisible.
|
||||
pub fn noop() -> Self {
|
||||
let (sender, _receiver) = mpsc::unbounded();
|
||||
Self { sender }
|
||||
}
|
||||
|
||||
/// Send a fully-built notification. Errors are deliberately swallowed;
|
||||
/// notifications are best-effort.
|
||||
pub fn send(&self, notification: ToolNotification) {
|
||||
let _ = self.sender.unbounded_send(notification);
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::BashOutputChunk`]: an incremental
|
||||
/// stdout/stderr chunk while a bash command is still running.
|
||||
pub fn send_bash_output_chunk(&self, chunk: BashOutputChunk) {
|
||||
self.send(ToolNotification::BashOutputChunk(chunk));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::BashExecutionComplete`]: a bash command
|
||||
/// exited (normally or via signal).
|
||||
pub fn send_bash_complete(&self, complete: BashExecutionComplete) {
|
||||
self.send(ToolNotification::BashExecutionComplete(complete));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::BashExecutionTimeout`]: a bash command
|
||||
/// exceeded its configured timeout and was killed.
|
||||
pub fn send_bash_timeout(&self, timeout: BashExecutionTimeout) {
|
||||
self.send(ToolNotification::BashExecutionTimeout(timeout));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::BashExecutionBackgrounded`]: a
|
||||
/// foreground bash command was moved to the background.
|
||||
pub fn send_bash_backgrounded(&self, backgrounded: BashExecutionBackgrounded) {
|
||||
self.send(ToolNotification::BashExecutionBackgrounded(backgrounded));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::BashExecutionFailed`]: a bash command
|
||||
/// could not be spawned.
|
||||
pub fn send_bash_failed(&self, failed: BashExecutionFailed) {
|
||||
self.send(ToolNotification::BashExecutionFailed(failed));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::FileWritten`]: a tool wrote to a file
|
||||
/// on disk.
|
||||
pub fn send_file_written(&self, written: FileWritten) {
|
||||
self.send(ToolNotification::FileWritten(written));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::TaskCompleted`]: a background task
|
||||
/// transitioned to a terminal state.
|
||||
pub fn send_task_complete(&self, task_completed: TaskSnapshot) {
|
||||
self.send(ToolNotification::TaskCompleted(task_completed));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::PlanModeEntered`]: the agent
|
||||
/// transitioned into plan mode.
|
||||
pub fn send_plan_mode_entered(&self, entered: PlanModeEntered) {
|
||||
self.send(ToolNotification::PlanModeEntered(entered));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::PlanModeExited`]: the agent transitioned
|
||||
/// out of plan mode and the captured plan is attached.
|
||||
pub fn send_plan_mode_exited(&self, exited: PlanModeExited) {
|
||||
self.send(ToolNotification::PlanModeExited(exited));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::UserQuestionAsked`]: the agent issued a
|
||||
/// structured question payload to the user.
|
||||
pub fn send_user_question_asked(&self, asked: UserQuestionAsked) {
|
||||
self.send(ToolNotification::UserQuestionAsked(asked));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::LspServerStarting`]: an LSP server is
|
||||
/// being spawned.
|
||||
pub fn send_lsp_starting(&self, starting: LspServerStarting) {
|
||||
self.send(ToolNotification::LspServerStarting(starting));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::LspServerReady`]: an LSP server
|
||||
/// finished its initialise handshake.
|
||||
pub fn send_lsp_ready(&self, ready: LspServerReady) {
|
||||
self.send(ToolNotification::LspServerReady(ready));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::LspServerCrashed`]: an LSP server
|
||||
/// process died unexpectedly.
|
||||
pub fn send_lsp_crashed(&self, crashed: LspServerCrashed) {
|
||||
self.send(ToolNotification::LspServerCrashed(crashed));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::LspServerRetrying`]: an LSP server is
|
||||
/// being restarted after a crash.
|
||||
pub fn send_lsp_retrying(&self, retrying: LspServerRetrying) {
|
||||
self.send(ToolNotification::LspServerRetrying(retrying));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::LspServerFailed`]: an LSP server is
|
||||
/// permanently dead (init failure or retry budget exhausted).
|
||||
pub fn send_lsp_failed(&self, failed: LspServerFailed) {
|
||||
self.send(ToolNotification::LspServerFailed(failed));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::ScheduledTaskFired`]: a recurring or
|
||||
/// one-shot scheduled task fired and its prompt should be executed.
|
||||
pub fn send_scheduled_task_fired(&self, fired: ScheduledTaskFired) {
|
||||
self.send(ToolNotification::ScheduledTaskFired(fired));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::ScheduledTaskRemoved`]: a scheduled task
|
||||
/// was deleted, expired, or a one-shot variant completed.
|
||||
pub fn send_scheduled_task_removed(&self, removed: ScheduledTaskRemoved) {
|
||||
self.send(ToolNotification::ScheduledTaskRemoved(removed));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::ScheduledTaskCreated`]: a new scheduled
|
||||
/// task was registered and should appear in subscriber views.
|
||||
pub fn send_scheduled_task_created(&self, created: ScheduledTaskCreated) {
|
||||
self.send(ToolNotification::ScheduledTaskCreated(created));
|
||||
}
|
||||
|
||||
/// Send a [`ToolNotification::MonitorEvent`]: a streaming event from
|
||||
/// a Monitor background process, ready for conversation injection.
|
||||
pub fn send_monitor_event(&self, event: MonitorEvent) {
|
||||
self.send(ToolNotification::MonitorEvent(event));
|
||||
}
|
||||
}
|
||||
731
crates/common/xai-tool-runtime/src/render.rs
Normal file
731
crates/common/xai-tool-runtime/src/render.rs
Normal file
|
|
@ -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()
|
||||
}]
|
||||
);
|
||||
}
|
||||
}
|
||||
83
crates/common/xai-tool-runtime/src/search.rs
Normal file
83
crates/common/xai-tool-runtime/src/search.rs
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
//! Backend-agnostic tool search interface.
|
||||
//!
|
||||
//! `ToolSearchIndex` is a `Send + Sync` trait so concrete implementations
|
||||
//! can live in different crates (BM25, OpenSearch, in-memory linear) and
|
||||
//! be stored as `Arc<dyn ToolSearchIndex>` for shared access across tasks.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
/// A single tool search hit.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ToolSearchResult {
|
||||
/// Qualified tool name (e.g. `"linear__save_issue"`).
|
||||
pub tool_name: String,
|
||||
/// Origin server name (e.g. `"linear"`).
|
||||
pub server_name: String,
|
||||
/// Tool description.
|
||||
pub description: String,
|
||||
/// Backend-defined relevance score; comparable within a single
|
||||
/// snapshot but not across snapshots.
|
||||
pub score: f32,
|
||||
/// Parameter names from the tool's input schema, in declaration order.
|
||||
pub parameters: Vec<String>,
|
||||
/// Full JSON Schema for the tool's input. Included so callers can
|
||||
/// construct dispatched tool calls without a separate schema fetch.
|
||||
pub input_schema: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Snapshot of a search query — results plus index metadata captured from
|
||||
/// the same point-in-time view.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct SearchSnapshot {
|
||||
pub results: Vec<ToolSearchResult>,
|
||||
/// Number of indexed tools that did not appear in `results`.
|
||||
pub total_hidden_tools: usize,
|
||||
/// `true` when the index reflects all available tools. `false` while
|
||||
/// the index source is still warming up.
|
||||
pub is_ready: bool,
|
||||
}
|
||||
|
||||
/// Summary of an MCP server (or other tool source) available for search.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ServerSummary {
|
||||
/// Server name (e.g. `"linear"`, `"slack"`).
|
||||
pub name: String,
|
||||
/// Optional short description of the server's surface area.
|
||||
pub description: Option<String>,
|
||||
/// Unqualified tool names, sorted alphabetically. Use
|
||||
/// [`Self::tool_count`] for a count without indirection.
|
||||
pub tool_names: Vec<String>,
|
||||
}
|
||||
|
||||
impl ServerSummary {
|
||||
/// Number of tools the server exposes.
|
||||
pub fn tool_count(&self) -> usize {
|
||||
self.tool_names.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Backend-agnostic search interface.
|
||||
///
|
||||
/// Implementations must be `Send + Sync` so they can be wrapped in
|
||||
/// `Arc<dyn ToolSearchIndex>` and shared across concurrent tasks.
|
||||
pub trait ToolSearchIndex: Send + Sync {
|
||||
/// Run a query against a single consistent index snapshot. Returning
|
||||
/// the metadata alongside the results lets the caller render an
|
||||
/// accurate "N results out of M" line without a second call.
|
||||
fn search_snapshot(&self, query: &str, limit: usize) -> SearchSnapshot;
|
||||
|
||||
/// Enumerate the unique servers in the index. Used to render the
|
||||
/// system-reminder listing connected integrations.
|
||||
fn list_server_summaries(&self) -> Vec<ServerSummary>;
|
||||
}
|
||||
|
||||
/// Resource wrapper for storing a `ToolSearchIndex` behind an `Arc` in
|
||||
/// shared resource maps.
|
||||
#[derive(Clone)]
|
||||
pub struct ToolIndex(pub Arc<dyn ToolSearchIndex>);
|
||||
|
||||
impl std::fmt::Debug for ToolIndex {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ToolIndex").finish()
|
||||
}
|
||||
}
|
||||
439
crates/common/xai-tool-runtime/src/streaming.rs
Normal file
439
crates/common/xai-tool-runtime/src/streaming.rs
Normal file
|
|
@ -0,0 +1,439 @@
|
|||
//! Canonical partial-result streaming contract shared by every streaming tool.
|
||||
//!
|
||||
//! A tool declares a [`StreamingSpec`] in its [`ToolCapabilities`] and emits
|
||||
//! deltas from `execute` via [`stream_chunk`], which materializes the spec into
|
||||
//! a [`PartialResultPayload`] carried by [`ToolProgress::Custom`]. Downstream
|
||||
//! layers dispatch on the envelope's `subkind` rather than on the tool's
|
||||
//! identity.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use xai_tool_protocol::StreamingSpec;
|
||||
|
||||
use crate::tool::ToolProgress;
|
||||
|
||||
/// Per-frame `delta` byte cap used when [`StreamingSpec::max_delta_bytes`] is
|
||||
/// unset. Guards against a single oversized tick flooding the harness in one
|
||||
/// frame. Deliberately independent of `ToolCapabilities::max_frame_bytes`,
|
||||
/// which caps whole frames (16 MiB ceiling), not deltas.
|
||||
const DEFAULT_MAX_DELTA_BYTES: usize = 16 * 1024;
|
||||
|
||||
/// Canonical payload carried by a streaming tool's [`ToolProgress::Custom`].
|
||||
///
|
||||
/// Downstream layers dispatch on the envelope's `subkind`. Deltas are
|
||||
/// append-only and lossless (see [`stream_chunk`]).
|
||||
///
|
||||
/// Parsed strictly (`deny_unknown_fields`): an unexpected field is a hard
|
||||
/// deserialize error rather than being silently ignored, so producer/consumer
|
||||
/// schema drift (e.g. a stale field from an un-updated producer) is caught and
|
||||
/// the frame is dropped with a warning instead of misinterpreted.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct PartialResultPayload {
|
||||
/// Content produced since the previous tick (the delta).
|
||||
pub delta: String,
|
||||
|
||||
/// Monotonic total bytes produced so far (NOT the current buffer length).
|
||||
pub total_bytes: u64,
|
||||
|
||||
/// Cumulative content was lost upstream and will never be delivered
|
||||
/// (distinct from a single-tick `gap`).
|
||||
#[serde(default)]
|
||||
pub truncated: bool,
|
||||
|
||||
/// This delta has a gap: a single oversized tick overflowed the tail
|
||||
/// buffer and its middle was dropped.
|
||||
#[serde(default)]
|
||||
pub gap: bool,
|
||||
}
|
||||
|
||||
/// Byte count of an incomplete (still-arriving) UTF-8 sequence at the very end
|
||||
/// of `bytes`, or 0 when the slice ends on a complete sequence or in invalid
|
||||
/// bytes that can never become valid (those are surfaced lossily instead of
|
||||
/// held forever).
|
||||
fn incomplete_utf8_suffix_len(bytes: &[u8]) -> usize {
|
||||
match std::str::from_utf8(bytes) {
|
||||
Ok(_) => 0,
|
||||
// `error_len() == None` means the error is an incomplete sequence at
|
||||
// the end of the input — the only case worth holding back.
|
||||
Err(e) if e.error_len().is_none() => bytes.len() - e.valid_up_to(),
|
||||
Err(_) => 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build at most one [`ToolProgress::Custom`] delta from a monotonic byte
|
||||
/// source, with UTF-8-safe slicing at both the tick boundary and the per-frame
|
||||
/// cap.
|
||||
///
|
||||
/// `tail` is the source's (possibly truncated) tail buffer — the newest bytes
|
||||
/// are always at its end. `total` is the monotonic count of bytes produced so
|
||||
/// far; `last_total` records how much has already been surfaced and is advanced
|
||||
/// in place. Returns `None` when `total` has not advanced (no new bytes).
|
||||
///
|
||||
/// Deltas are **append-only and lossless**: when a delta would end mid-way
|
||||
/// through a multi-byte UTF-8 sequence, or exceeds the per-frame cap
|
||||
/// ([`StreamingSpec::max_delta_bytes`], default 16 KiB), the excess bytes are
|
||||
/// *held back* — `last_total` advances only past the emitted bytes, so the next
|
||||
/// call re-slices the remainder from the (still-growing) tail. Concatenated
|
||||
/// deltas are therefore always valid UTF-8 and lossless.
|
||||
///
|
||||
/// `truncated` is the caller's cumulative upstream-truncation flag (e.g. a
|
||||
/// source that hit a hard output cap and will never deliver the elided bytes).
|
||||
/// It is copied into the payload verbatim and is intentionally distinct from
|
||||
/// the per-tick `gap` (a single oversized tick overflowed the tail buffer and
|
||||
/// its middle was dropped upstream). Sources with no cumulative-truncation
|
||||
/// notion pass `false`.
|
||||
pub fn stream_chunk(
|
||||
spec: &StreamingSpec,
|
||||
tail: &[u8],
|
||||
total: u64,
|
||||
last_total: &mut u64,
|
||||
truncated: bool,
|
||||
) -> Option<ToolProgress> {
|
||||
if total <= *last_total {
|
||||
return None;
|
||||
}
|
||||
let new = total - *last_total;
|
||||
let tail_len = tail.len() as u64;
|
||||
// Deltas are keyed off the monotonic `total`, not the buffer length: when
|
||||
// all genuinely-new bytes still fit in the tail we slice its suffix; when a
|
||||
// single tick's burst exceeded the buffer the middle was dropped upstream,
|
||||
// so we emit what survived plus a `gap` marker.
|
||||
let (delta_bytes, gap) = if new <= tail_len {
|
||||
(&tail[(tail_len - new) as usize..], false)
|
||||
} else {
|
||||
(tail, true)
|
||||
};
|
||||
|
||||
let cap = spec
|
||||
.max_delta_bytes
|
||||
.map_or(DEFAULT_MAX_DELTA_BYTES, |c| c as usize);
|
||||
|
||||
// Defer: emit the longest prefix that fits the cap AND ends on a complete
|
||||
// UTF-8 sequence; hold the rest back for the next call (the tail still
|
||||
// contains it, since `last_total` only advances past the emitted bytes).
|
||||
// Nothing is dropped.
|
||||
let mut cut = delta_bytes.len().min(cap);
|
||||
while cut > 0 && incomplete_utf8_suffix_len(&delta_bytes[..cut]) > 0 {
|
||||
cut -= 1;
|
||||
}
|
||||
// A cap smaller than one multi-byte char would deadlock at cut == 0 while
|
||||
// bytes remain; emit the full first char in that pathological case rather
|
||||
// than stalling forever.
|
||||
if cut == 0 && !delta_bytes.is_empty() {
|
||||
cut = delta_bytes.len().min(4);
|
||||
while cut < delta_bytes.len() && incomplete_utf8_suffix_len(&delta_bytes[..cut]) > 0 {
|
||||
cut += 1;
|
||||
}
|
||||
}
|
||||
if cut == 0 {
|
||||
return None;
|
||||
}
|
||||
let delta = String::from_utf8_lossy(&delta_bytes[..cut]).into_owned();
|
||||
let consumed = cut as u64;
|
||||
|
||||
// Advance only past what was emitted (gap case: the upstream-dropped
|
||||
// middle counts as consumed — those bytes can never be re-sliced).
|
||||
*last_total = if gap {
|
||||
total - (delta_bytes.len() as u64 - consumed.min(delta_bytes.len() as u64))
|
||||
} else {
|
||||
*last_total + consumed
|
||||
};
|
||||
|
||||
let payload = PartialResultPayload {
|
||||
delta,
|
||||
total_bytes: total,
|
||||
// The caller's cumulative upstream-loss flag passes through verbatim;
|
||||
// the per-tick `gap` is reported separately.
|
||||
truncated,
|
||||
gap,
|
||||
};
|
||||
Some(ToolProgress::Custom {
|
||||
subkind: spec.subkind.clone(),
|
||||
// Infallible: a struct of String/u64/bool/Copy-enums has no map keys
|
||||
// or floats that could make `to_value` fail.
|
||||
payload: serde_json::to_value(&payload).expect("PartialResultPayload always serializes"),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn spec_with(max_delta_bytes: Option<u32>) -> StreamingSpec {
|
||||
StreamingSpec {
|
||||
subkind: "test_chunk".to_owned(),
|
||||
max_delta_bytes,
|
||||
}
|
||||
}
|
||||
|
||||
/// Run `stream_chunk`, assert it produced a frame, and decode the payload.
|
||||
fn run(
|
||||
spec: &StreamingSpec,
|
||||
tail: &[u8],
|
||||
total: u64,
|
||||
last_total: &mut u64,
|
||||
truncated: bool,
|
||||
) -> PartialResultPayload {
|
||||
let progress = stream_chunk(spec, tail, total, last_total, truncated)
|
||||
.expect("expected a progress frame");
|
||||
let ToolProgress::Custom { subkind, payload } = progress else {
|
||||
panic!("expected ToolProgress::Custom");
|
||||
};
|
||||
assert_eq!(subkind, "test_chunk");
|
||||
serde_json::from_value(payload).expect("payload decodes")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_new_bytes_returns_none_and_leaves_last_total() {
|
||||
let spec = spec_with(None);
|
||||
let mut last = 10;
|
||||
assert!(stream_chunk(&spec, b"abc", 10, &mut last, false).is_none());
|
||||
assert!(stream_chunk(&spec, b"abc", 5, &mut last, false).is_none());
|
||||
assert_eq!(
|
||||
last, 10,
|
||||
"last_total is untouched when total does not advance"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn emits_suffix_delta_and_advances_last_total() {
|
||||
let spec = spec_with(None);
|
||||
let mut last = 2;
|
||||
// total 2 -> 5: 3 genuinely-new bytes, all present in the tail suffix.
|
||||
let p = run(&spec, b"abcde", 5, &mut last, false);
|
||||
assert_eq!(p.delta, "cde");
|
||||
assert_eq!(p.total_bytes, 5);
|
||||
assert!(!p.gap);
|
||||
assert!(!p.truncated);
|
||||
assert_eq!(last, 5, "last_total advanced in place");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gap_set_when_new_exceeds_surviving_tail() {
|
||||
// 100 new bytes but only a 4-byte tail survived upstream: the middle
|
||||
// was dropped, so the whole tail is emitted with gap = true.
|
||||
let spec = spec_with(None);
|
||||
let mut last = 0;
|
||||
let p = run(&spec, b"tail", 100, &mut last, false);
|
||||
assert_eq!(p.delta, "tail");
|
||||
assert_eq!(p.total_bytes, 100);
|
||||
assert!(p.gap);
|
||||
assert!(
|
||||
!p.truncated,
|
||||
"a per-tick gap must not set cumulative truncated"
|
||||
);
|
||||
assert_eq!(last, 100, "fully-emitted gap delta consumes the total");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncated_is_caller_supplied_and_distinct_from_gap() {
|
||||
let spec = spec_with(None);
|
||||
let mut last = 0;
|
||||
// Caller reports cumulative upstream truncation; no per-tick gap here.
|
||||
let p = run(&spec, b"abc", 3, &mut last, true);
|
||||
assert!(
|
||||
p.truncated,
|
||||
"caller's cumulative flag passes through verbatim"
|
||||
);
|
||||
assert!(!p.gap, "no tail overflow this tick");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_multibyte_split_across_ticks_is_held_back_and_reassembled() {
|
||||
// Tick 1 delivers "aé" cut mid-'é' (0xC3 without 0xA9). The lone lead
|
||||
// byte is held back, NOT emitted as U+FFFD.
|
||||
let spec = spec_with(None);
|
||||
let mut last = 0;
|
||||
let p = run(&spec, b"a\xC3", 2, &mut last, false);
|
||||
assert_eq!(p.delta, "a", "incomplete UTF-8 suffix held back");
|
||||
assert_eq!(last, 1, "last_total advances only past emitted bytes");
|
||||
|
||||
// Tick 2: the continuation byte arrives; the held bytes re-slice from
|
||||
// the tail and the char comes out whole.
|
||||
let p = run(&spec, "aé".as_bytes(), 3, &mut last, false);
|
||||
assert_eq!(p.delta, "é", "held bytes reassemble into a whole char");
|
||||
assert_eq!(last, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_over_cap_defers_remainder_to_next_call_without_loss() {
|
||||
// Cap 4: a 9-byte burst is paced out over capped frames; nothing is
|
||||
// dropped and the concatenation is lossless.
|
||||
let spec = spec_with(Some(4));
|
||||
let mut last = 0;
|
||||
let mut out = String::new();
|
||||
while last < 9 {
|
||||
let p = run(&spec, b"abcdefghi", 9, &mut last, false);
|
||||
assert!(p.delta.len() <= 4, "every frame respects the cap");
|
||||
out.push_str(&p.delta);
|
||||
}
|
||||
assert_eq!(out, "abcdefghi", "deferred remainders are all emitted");
|
||||
assert_eq!(last, 9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_cap_cut_respects_utf8_boundaries() {
|
||||
// 7 ASCII bytes + 'é' (2 bytes) = 9 bytes. A cap of 8 would split the
|
||||
// 'é'; the cut backs off and the 'é' is deferred whole.
|
||||
let tail = "aaaaaaaé".as_bytes();
|
||||
let spec = spec_with(Some(8));
|
||||
let mut last = 0;
|
||||
let p = run(&spec, tail, tail.len() as u64, &mut last, false);
|
||||
assert_eq!(p.delta, "aaaaaaa", "backed off the split multibyte char");
|
||||
let p = run(&spec, tail, tail.len() as u64, &mut last, false);
|
||||
assert_eq!(p.delta, "é");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn payload_decodes_with_optional_flags_absent() {
|
||||
// Wire tolerance: a payload missing the bool/count fields still
|
||||
// decodes (serde defaults), matching the ToolCapabilities convention.
|
||||
let p: PartialResultPayload = serde_json::from_value(serde_json::json!({
|
||||
"delta": "x",
|
||||
"total_bytes": 1,
|
||||
}))
|
||||
.expect("payload with absent flags decodes");
|
||||
assert!(!p.truncated);
|
||||
assert!(!p.gap);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn payload_rejects_unknown_field() {
|
||||
// Strict (`deny_unknown_fields`): a stale/typo'd field — e.g. a removed
|
||||
// `accumulation` from an un-updated producer — is a hard error, not
|
||||
// silently ignored, so schema drift never decodes into a partial frame.
|
||||
let decoded = serde_json::from_value::<PartialResultPayload>(serde_json::json!({
|
||||
"delta": "x",
|
||||
"total_bytes": 1,
|
||||
"truncated": false,
|
||||
"gap": false,
|
||||
"accumulation": "append",
|
||||
}));
|
||||
assert!(
|
||||
decoded.is_err(),
|
||||
"an unknown field must be rejected under deny_unknown_fields"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Limit / latency invariants ──────────────────────────────────────────
|
||||
|
||||
/// A backlog drains in exactly `ceil(new / cap)` calls — no extra round-trips.
|
||||
#[test]
|
||||
fn drains_backlog_in_minimum_ticks() {
|
||||
let cap = 4usize;
|
||||
let spec = spec_with(Some(cap as u32));
|
||||
let data = b"abcdefghij";
|
||||
let total = data.len() as u64;
|
||||
let mut last = 0;
|
||||
let mut ticks = 0usize;
|
||||
let mut out = String::new();
|
||||
while last < total {
|
||||
out.push_str(&run(&spec, data, total, &mut last, false).delta);
|
||||
ticks += 1;
|
||||
assert!(ticks <= 100, "must terminate");
|
||||
}
|
||||
assert_eq!(out, "abcdefghij", "lossless");
|
||||
assert_eq!(
|
||||
ticks,
|
||||
data.len().div_ceil(cap),
|
||||
"no extra ticks beyond ceil(new / cap)"
|
||||
);
|
||||
}
|
||||
|
||||
/// ASCII frames fill to the cap — no under-fill, no empty trailing frame.
|
||||
#[test]
|
||||
fn ascii_frames_fill_to_cap() {
|
||||
let spec = spec_with(Some(4));
|
||||
let data = b"abcdefgh";
|
||||
let mut last = 0;
|
||||
assert_eq!(run(&spec, data, 8, &mut last, false).delta, "abcd");
|
||||
assert_eq!(run(&spec, data, 8, &mut last, false).delta, "efgh");
|
||||
assert!(
|
||||
stream_chunk(&spec, data, 8, &mut last, false).is_none(),
|
||||
"no spurious empty trailing frame"
|
||||
);
|
||||
}
|
||||
|
||||
/// A delta exactly at the cap is one frame, no gap, nothing deferred.
|
||||
#[test]
|
||||
fn exact_cap_emits_single_frame() {
|
||||
let cap = 8u64;
|
||||
let spec = spec_with(Some(cap as u32));
|
||||
let mut last = 0;
|
||||
let p = run(&spec, b"abcdefgh", cap, &mut last, false);
|
||||
assert_eq!(p.delta, "abcdefgh");
|
||||
assert!(!p.gap);
|
||||
assert_eq!(last, cap);
|
||||
assert!(stream_chunk(&spec, b"abcdefgh", cap, &mut last, false).is_none());
|
||||
}
|
||||
|
||||
/// A cap smaller than one char still emits a whole char — never stalls.
|
||||
#[test]
|
||||
fn tiny_cap_below_char_still_makes_progress() {
|
||||
let spec = spec_with(Some(1));
|
||||
let mut last = 0;
|
||||
let p = run(&spec, "é".as_bytes(), 2, &mut last, false);
|
||||
assert_eq!(
|
||||
p.delta, "é",
|
||||
"emits the whole first char despite cap < charlen"
|
||||
);
|
||||
assert_eq!(last, 2, "and makes forward progress");
|
||||
}
|
||||
|
||||
/// UTF-8 backoff loses at most 3 bytes, so frames stay within 3 of the cap.
|
||||
#[test]
|
||||
fn utf8_backoff_stays_within_three_bytes_of_cap() {
|
||||
let cap = 7usize; // splits a 4-byte char -> backs off to 4 (cap - 3)
|
||||
let spec = spec_with(Some(cap as u32));
|
||||
let data = "😀😀😀😀".as_bytes();
|
||||
let total = data.len() as u64;
|
||||
let mut last = 0;
|
||||
while last < total {
|
||||
let n = run(&spec, data, total, &mut last, false).delta.len();
|
||||
assert!(
|
||||
n.is_multiple_of(4) && n >= 4,
|
||||
"emits whole 4-byte chars, got {n}"
|
||||
);
|
||||
assert!(
|
||||
n >= cap - 3,
|
||||
"frame stays within 3 bytes of the cap, got {n}"
|
||||
);
|
||||
assert!(n <= cap, "frame respects the cap, got {n}");
|
||||
}
|
||||
assert_eq!(last, total, "drains losslessly");
|
||||
}
|
||||
|
||||
/// A gap drains only the surviving tail; never re-scans the dropped middle.
|
||||
#[test]
|
||||
fn gap_with_cap_paces_surviving_tail_and_terminates() {
|
||||
let spec = spec_with(Some(4));
|
||||
let tail = b"abcdefgh";
|
||||
let total = 1000u64; // only 8 of 1000 bytes survived in the tail
|
||||
let mut last = 0;
|
||||
let mut ticks = 0usize;
|
||||
let mut emitted = 0usize;
|
||||
let mut saw_gap = false;
|
||||
while last < total {
|
||||
let Some(progress) = stream_chunk(&spec, tail, total, &mut last, false) else {
|
||||
break;
|
||||
};
|
||||
let ToolProgress::Custom { payload, .. } = progress else {
|
||||
panic!("expected ToolProgress::Custom");
|
||||
};
|
||||
let p: PartialResultPayload = serde_json::from_value(payload).unwrap();
|
||||
saw_gap |= p.gap;
|
||||
emitted += p.delta.len();
|
||||
ticks += 1;
|
||||
assert!(
|
||||
ticks <= 4,
|
||||
"gap pacing must drain only the surviving tail, not re-scan the dropped middle"
|
||||
);
|
||||
}
|
||||
assert!(saw_gap, "first frame reports the gap");
|
||||
assert_eq!(
|
||||
emitted,
|
||||
tail.len(),
|
||||
"emits exactly the surviving tail bytes — no replay of dropped middle"
|
||||
);
|
||||
}
|
||||
}
|
||||
440
crates/common/xai-tool-runtime/src/tool.rs
Normal file
440
crates/common/xai-tool-runtime/src/tool.rs
Normal file
|
|
@ -0,0 +1,440 @@
|
|||
//! The unified `Tool` trait, the streaming primitives it produces, and the
|
||||
//! helper constructors tool authors use to build well-formed streams.
|
||||
//!
|
||||
//! `Tool::execute` is the canonical streaming entry point; the runtime
|
||||
//! always calls it. The default impl wraps `Tool::run` (the simpler
|
||||
//! convenience hook) into a single-item terminal stream so blocking tools
|
||||
//! don't have to think about streaming. A tool that overrides neither gets
|
||||
//! a `NotImplemented` terminal at runtime.
|
||||
//!
|
||||
//! `ToolStream<T>` is a type alias for an opaque pinned stream; the helper
|
||||
//! free functions [`terminal_only`] and [`with_progress`] are the supported
|
||||
//! ways to build one. Stream invariant: at most arbitrarily many `Progress`
|
||||
//! items, ending in exactly one `Terminal`.
|
||||
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::stream::{self, Stream, StreamExt};
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use xai_tool_protocol::{ToolCapabilities, ToolId};
|
||||
use xai_tool_types::ToolDescription;
|
||||
|
||||
use crate::context::{ListToolsContext, ToolCallContext};
|
||||
use crate::error::ToolError;
|
||||
use crate::render::{ToolChatCompletionResponse, ToolOutput};
|
||||
|
||||
/// The unified tool trait used by every tool source.
|
||||
///
|
||||
/// Implement either `run` (blocking) or `execute` (streaming). The
|
||||
/// runtime only ever invokes `execute`.
|
||||
pub trait Tool: Send + Sync {
|
||||
/// Typed input. Must be deserialisable from JSON for wire dispatch.
|
||||
type Args: for<'de> Deserialize<'de> + JsonSchema + Send + 'static;
|
||||
|
||||
/// Typed output. Must implement [`ToolOutput`] which provides
|
||||
/// model-facing content blocks and optional chat-completion
|
||||
/// responses. All methods have defaults, so most output types
|
||||
/// just need:
|
||||
/// ```rust,ignore
|
||||
/// impl ToolOutput for MyOutput {}
|
||||
/// ```
|
||||
type Output: Serialize + ToolOutput + Send + 'static;
|
||||
|
||||
/// Stable identity used by the runtime to route to this tool.
|
||||
fn id(&self) -> ToolId;
|
||||
|
||||
/// Model-facing description and argument schema.
|
||||
///
|
||||
/// Receives the per-turn [`ListToolsContext`] (viewer context,
|
||||
/// attachments, etc.) — the same context [`Tool::should_list`] consumes
|
||||
/// — so descriptions can be context-aware. Most tools ignore `_ctx` and
|
||||
/// return a static description. Callers outside a listing turn pass
|
||||
/// [`ListToolsContext::default`].
|
||||
fn description(&self, _ctx: &ListToolsContext) -> ToolDescription;
|
||||
|
||||
/// Per-tool capability flags (concurrency, scope, frame caps, ...).
|
||||
fn capabilities(&self) -> ToolCapabilities {
|
||||
ToolCapabilities::default()
|
||||
}
|
||||
|
||||
/// Whether [`Tool::description`] varies with the per-turn
|
||||
/// [`ListToolsContext`]. `false` for the common case of a static
|
||||
/// description.
|
||||
fn has_dynamic_description(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Per-turn listing predicate. Return `false` to exclude this tool
|
||||
/// from the model-facing manifest for a given turn.
|
||||
fn should_list(&self, _ctx: &ListToolsContext) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Streaming entry point. Default impl wraps `run` into a single-item
|
||||
/// stream so blocking tools just override `run`.
|
||||
///
|
||||
/// Uses a native `async fn` in trait (RPITIT) with an explicit `Send`
|
||||
/// bound rather than `#[async_trait]`, so the returned future is not
|
||||
/// boxed. The `Tool` trait is only ever consumed generically (type
|
||||
/// erasure goes through [`ToolDyn`]), so it does not need to be
|
||||
/// dyn-compatible.
|
||||
fn execute(
|
||||
&self,
|
||||
ctx: ToolCallContext,
|
||||
args: Self::Args,
|
||||
) -> impl Future<Output = ToolStream<Self::Output>> + Send {
|
||||
async move {
|
||||
let result = self.run(ctx, args).await;
|
||||
terminal_only(result)
|
||||
}
|
||||
}
|
||||
|
||||
/// Blocking convenience entry point. Default returns
|
||||
/// `Err(ToolError::not_implemented(...))` so a tool that overrides
|
||||
/// neither method fails loudly at the first call.
|
||||
fn run(
|
||||
&self,
|
||||
_ctx: ToolCallContext,
|
||||
_args: Self::Args,
|
||||
) -> impl Future<Output = Result<Self::Output, ToolError>> + Send {
|
||||
async move {
|
||||
Err(ToolError::not_implemented(
|
||||
"Tool must implement either `run` or `execute`",
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Stream of items a tool produces during a single call. Shape:
|
||||
/// `[Progress(_)*, Terminal(Result<T, ToolError>)]`.
|
||||
pub type ToolStream<T> = Pin<Box<dyn Stream<Item = ToolStreamItem<T>> + Send>>;
|
||||
|
||||
/// One item in a [`ToolStream`].
|
||||
#[derive(Debug)]
|
||||
pub enum ToolStreamItem<T> {
|
||||
/// Intermediate progress. Zero or more per stream.
|
||||
Progress(ToolProgress),
|
||||
/// Terminal result. Exactly one per stream, always last.
|
||||
Terminal(Result<T, ToolError>),
|
||||
}
|
||||
|
||||
impl<T> ToolStreamItem<T> {
|
||||
/// `true` for the `Terminal` variant. Stream consumers use this to
|
||||
/// short-circuit once the final item has been seen.
|
||||
pub fn is_terminal(&self) -> bool {
|
||||
matches!(self, Self::Terminal(_))
|
||||
}
|
||||
}
|
||||
|
||||
/// Open-ended progress payload. The `Custom` arm is the escape hatch for
|
||||
/// tool-specific shapes that don't map onto `Text` or `Content`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum ToolProgress {
|
||||
/// Free-form text chunk (terminal stdout, log line, partial response).
|
||||
Text { text: String },
|
||||
/// Rich content blocks.
|
||||
Content { blocks: Vec<ContentBlock> },
|
||||
/// Tool-defined progress payload. `subkind` is a stable snake-case
|
||||
/// discriminator owned by the tool. The outer `"kind"` serde tag is
|
||||
/// always `"custom"` for this variant; `subkind` is the producer's
|
||||
/// own discriminator and lives one level deeper to avoid colliding
|
||||
/// with the tag.
|
||||
Custom {
|
||||
subkind: String,
|
||||
payload: serde_json::Value,
|
||||
},
|
||||
}
|
||||
|
||||
/// Rich content block for `ToolProgress::Content`. Mirrors the wire-side
|
||||
/// `McpBlock` shape on the protocol crate so adapters can move blocks
|
||||
/// across the wire boundary without re-encoding.
|
||||
///
|
||||
/// The `Image` variant carries optional metadata fields (`media_id`,
|
||||
/// `filename`, `path`, `metadata`) that the Grok SLOP converter uses to
|
||||
/// build `Media` objects with the right identifiers and file paths. Tool
|
||||
/// authors populate these in their `ToolOutput::model_output()` impls so
|
||||
/// downstream consumers don't need to know the concrete tool type.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ContentBlock {
|
||||
/// Plain text. Equivalent to `ToolProgress::Text` but allowed inside a
|
||||
/// content list so a tool can interleave images and text.
|
||||
Text { text: String },
|
||||
/// Image. `mime_type` is e.g. `"image/png"`; `data` is base64.
|
||||
///
|
||||
/// Optional metadata fields are used by the Grok SLOP converter to
|
||||
/// produce `Media` objects with the correct identifiers:
|
||||
/// - `media_id`: unique image ID for referencing in subsequent tool calls
|
||||
/// - `filename`: human-readable filename (e.g. `"xK29f.png"`)
|
||||
/// - `path`: file path on the Grok Computer filesystem
|
||||
/// - `metadata`: arbitrary key-value pairs (e.g. `title`, `webpage_url`)
|
||||
Image {
|
||||
#[serde(alias = "mimeType")]
|
||||
mime_type: String,
|
||||
data: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
media_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
filename: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
path: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "std::collections::HashMap::is_empty")]
|
||||
metadata: std::collections::HashMap<String, String>,
|
||||
},
|
||||
/// Resource pointer. `uri` is required; `mime_type` and `text` are
|
||||
/// optional preview data.
|
||||
Resource {
|
||||
uri: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none", alias = "mimeType")]
|
||||
mime_type: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
text: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Build a single-item stream containing only the terminal result.
|
||||
///
|
||||
/// The most common shape — a blocking tool's `Tool::run` is wrapped this
|
||||
/// way by the default `execute` impl.
|
||||
pub fn terminal_only<T: Send + 'static>(result: Result<T, ToolError>) -> ToolStream<T> {
|
||||
Box::pin(stream::iter(std::iter::once(ToolStreamItem::Terminal(
|
||||
result,
|
||||
))))
|
||||
}
|
||||
|
||||
/// Build a stream that emits each progress item from `progress` then
|
||||
/// resolves `terminal` and emits its result as the final `Terminal` item.
|
||||
///
|
||||
/// `terminal` is awaited only after `progress` has fully drained, so a
|
||||
/// progress producer that pulls from the same upstream can drive the
|
||||
/// terminal value without conflicts.
|
||||
pub fn with_progress<T, P, F>(progress: P, terminal: F) -> ToolStream<T>
|
||||
where
|
||||
T: Send + 'static,
|
||||
P: Stream<Item = ToolProgress> + Send + 'static,
|
||||
F: Future<Output = Result<T, ToolError>> + Send + 'static,
|
||||
{
|
||||
let progress = progress.map(ToolStreamItem::Progress);
|
||||
let tail = stream::once(async move { ToolStreamItem::Terminal(terminal.await) });
|
||||
Box::pin(progress.chain(tail))
|
||||
}
|
||||
|
||||
/// Type-erased tool output that bundles the serialised JSON value with
|
||||
/// model-facing content blocks extracted at serialisation time.
|
||||
///
|
||||
/// When the `ToolDyn` blanket impl serialises a typed `Tool::Output` to
|
||||
/// JSON it also calls [`ToolOutput::model_output`] and
|
||||
/// [`ToolOutput::chat_completion_output`], capturing both
|
||||
/// here so downstream consumers never need to deserialise the `Value`
|
||||
/// back into the concrete type.
|
||||
///
|
||||
/// # MCP invariant
|
||||
///
|
||||
/// `model_output` is **always non-empty**. The [`ToolOutput`]
|
||||
/// default serialises the typed output as a JSON text block, so even
|
||||
/// tools that never override `model_output()` produce MCP-compliant
|
||||
/// content. Tools that override the method are expected to return at
|
||||
/// least one block.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TypedToolOutput {
|
||||
/// Identity of the tool that produced this output.
|
||||
pub tool_id: ToolId,
|
||||
/// Serialised JSON representation of the tool output.
|
||||
pub value: Value,
|
||||
/// Model-facing content blocks. Always contains at least one block
|
||||
/// (MCP compliance — see struct-level docs).
|
||||
pub model_output: Vec<ContentBlock>,
|
||||
/// Optional chat-completion response frame, extracted from the typed
|
||||
/// output via [`ToolOutput`] at type-erasure time.
|
||||
///
|
||||
/// `None` for the vast majority of tools. Present when a tool
|
||||
/// produces a frontend-ready chat response (render cards, progress
|
||||
/// reports, etc.).
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub chat_completion_output: Option<ToolChatCompletionResponse>,
|
||||
}
|
||||
|
||||
impl TypedToolOutput {
|
||||
/// Convenience constructor for building a `TypedToolOutput` from an
|
||||
/// already-serialised `Value` (e.g. at wire decode boundaries).
|
||||
///
|
||||
/// `model_output` is derived from the value via
|
||||
/// [`extract_content_blocks`](crate::render::extract_content_blocks);
|
||||
/// `chat_completion_output` is always `None` since the wire format
|
||||
/// does not carry it.
|
||||
pub fn from_value(tool_id: ToolId, value: Value) -> Self {
|
||||
let model_output = crate::render::extract_content_blocks(&value);
|
||||
Self {
|
||||
tool_id,
|
||||
value,
|
||||
model_output,
|
||||
chat_completion_output: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Override the `chat_completion_output` that [`Self::from_value`]
|
||||
/// leaves `None` — used by wire-decode boundaries that recover the
|
||||
/// frame the plain `from_value` path cannot carry.
|
||||
pub fn with_chat_completion_output(
|
||||
mut self,
|
||||
chat_completion_output: Option<ToolChatCompletionResponse>,
|
||||
) -> Self {
|
||||
self.chat_completion_output = chat_completion_output;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolOutput for TypedToolOutput {
|
||||
fn model_output(&self) -> Vec<ContentBlock> {
|
||||
self.model_output.clone()
|
||||
}
|
||||
|
||||
fn chat_completion_output(&self) -> Option<ToolChatCompletionResponse> {
|
||||
self.chat_completion_output.clone()
|
||||
}
|
||||
}
|
||||
|
||||
/// Type erased tool trait. Auto-generated for every typed Tool implementation.
|
||||
#[async_trait]
|
||||
pub trait ToolDyn: Send + Sync {
|
||||
/// Stable identity. Same value as [`Tool::id`].
|
||||
fn id(&self) -> ToolId;
|
||||
|
||||
/// Model-facing description. Same value as [`Tool::description`].
|
||||
fn description(&self, ctx: &ListToolsContext) -> ToolDescription;
|
||||
|
||||
/// Per-tool capability flags. Same value as [`Tool::capabilities`].
|
||||
fn capabilities(&self) -> ToolCapabilities {
|
||||
ToolCapabilities::default()
|
||||
}
|
||||
|
||||
/// Same value as [`Tool::has_dynamic_description`].
|
||||
fn has_dynamic_description(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn should_list(&self, _ctx: &ListToolsContext) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// JSON-typed streaming entry point. The returned stream MUST honour
|
||||
/// the same `[Progress*, Terminal]` invariant as [`ToolStream`].
|
||||
///
|
||||
/// Terminal items carry [`TypedToolOutput`] which bundles both the
|
||||
/// serialised JSON `Value` and the model-facing content blocks
|
||||
/// extracted from the typed output.
|
||||
async fn execute(&self, ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<T: Tool> ToolDyn for T {
|
||||
fn id(&self) -> ToolId {
|
||||
Tool::id(self)
|
||||
}
|
||||
|
||||
fn description(&self, ctx: &ListToolsContext) -> ToolDescription {
|
||||
Tool::description(self, ctx)
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ToolCapabilities {
|
||||
Tool::capabilities(self)
|
||||
}
|
||||
|
||||
fn has_dynamic_description(&self) -> bool {
|
||||
Tool::has_dynamic_description(self)
|
||||
}
|
||||
|
||||
fn should_list(&self, ctx: &ListToolsContext) -> bool {
|
||||
Tool::should_list(self, ctx)
|
||||
}
|
||||
|
||||
async fn execute(&self, ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput> {
|
||||
let typed_args: T::Args = match serde_json::from_value(args) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return terminal_only(Err(ToolError::invalid_arguments(e.to_string()))),
|
||||
};
|
||||
|
||||
let tool_id = Tool::id(self);
|
||||
let typed_stream = Tool::execute(self, ctx, typed_args).await;
|
||||
|
||||
Box::pin(typed_stream.map(move |item| {
|
||||
match item {
|
||||
ToolStreamItem::Progress(p) => ToolStreamItem::Progress(p),
|
||||
ToolStreamItem::Terminal(Ok(out)) => {
|
||||
// TypedToolOutput::value and the model_output
|
||||
// fallback.
|
||||
match serde_json::to_value(&out) {
|
||||
Ok(value) => {
|
||||
let custom = out.model_output();
|
||||
let model_output = if custom.is_empty() {
|
||||
// Default path: extract blocks from the
|
||||
// already-serialised Value.
|
||||
crate::render::extract_content_blocks(&value)
|
||||
} else {
|
||||
custom
|
||||
};
|
||||
let chat_completion_output = out.chat_completion_output();
|
||||
ToolStreamItem::Terminal(Ok(TypedToolOutput {
|
||||
tool_id: tool_id.clone(),
|
||||
value,
|
||||
model_output,
|
||||
chat_completion_output,
|
||||
}))
|
||||
}
|
||||
Err(e) => ToolStreamItem::Terminal(Err(ToolError::execution(
|
||||
tool_id.clone(),
|
||||
format!("serializing tool output to JSON: {e}"),
|
||||
)
|
||||
.with_source(e))),
|
||||
}
|
||||
}
|
||||
ToolStreamItem::Terminal(Err(e)) => ToolStreamItem::Terminal(Err(e)),
|
||||
}
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
/// Convenience alias for the most common [`ToolDyn`] handle shape.
|
||||
pub type ArcTool = Arc<dyn ToolDyn>;
|
||||
|
||||
/// Variant identifier for tools that ship multiple implementations under
|
||||
/// one stable [`ToolId`].
|
||||
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
|
||||
pub enum ToolVariant {
|
||||
/// The implicit fallback variant.
|
||||
Default,
|
||||
/// A named variant. The string is treated opaquely by the registry.
|
||||
Variant(String),
|
||||
}
|
||||
|
||||
/// Group of related tools that share one [`ToolId`] but route to different
|
||||
/// implementations chosen by a [`ToolVariant`].
|
||||
pub trait ToolFamily: Send + Sync {
|
||||
/// Identity shared by every variant in this family.
|
||||
fn id(&self) -> ToolId;
|
||||
|
||||
/// Resolve a `variant` to its concrete tool. Returns `None` when the
|
||||
/// family does not expose the requested variant.
|
||||
fn get_tool(&self, variant: &ToolVariant) -> Option<ArcTool>;
|
||||
|
||||
/// Every variant the family exposes. Registries iterate this once at
|
||||
/// startup and cache the results, so allocating here is fine.
|
||||
fn variants(&self) -> Vec<ToolVariant>;
|
||||
|
||||
/// Variant name the default falls back to when the family's `Default`
|
||||
/// arm is itself a named variant. Returns `None` when the default is
|
||||
/// the [`ToolVariant::Default`] sentinel.
|
||||
fn default_variant_name(&self) -> Option<&'static str> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Convenience alias for the most common [`ToolFamily`] handle shape.
|
||||
pub type ArcToolFamily = Arc<dyn ToolFamily>;
|
||||
Loading…
Reference in a new issue