diff options
Diffstat (limited to 'crates/atuin-ai')
| -rw-r--r-- | crates/atuin-ai/src/stream.rs | 20 | ||||
| -rw-r--r-- | crates/atuin-ai/src/tui/state.rs | 44 | ||||
| -rw-r--r-- | crates/atuin-ai/src/tui/view/turn.rs | 1 |
3 files changed, 54 insertions, 11 deletions
diff --git a/crates/atuin-ai/src/stream.rs b/crates/atuin-ai/src/stream.rs index 4673f2cd..9c21fc05 100644 --- a/crates/atuin-ai/src/stream.rs +++ b/crates/atuin-ai/src/stream.rs @@ -40,6 +40,8 @@ pub(crate) enum StreamContent { tool_use_id: String, content: String, is_error: bool, + remote: bool, + content_length: Option<usize>, }, } @@ -179,7 +181,9 @@ fn create_chat_stream( let tool_use_id = json.get("tool_use_id").and_then(|v| v.as_str()).unwrap_or("").to_string(); let content = json.get("content").and_then(|v| v.as_str()).unwrap_or("").to_string(); let is_error = json.get("is_error").and_then(|v| v.as_bool()).unwrap_or(false); - yield Ok(StreamFrame::Content(StreamContent::ToolResult { tool_use_id, content, is_error })); + let remote = json.get("remote").and_then(|v| v.as_bool()).unwrap_or(false); + let content_length = json.get("content_length").and_then(|v| v.as_u64()).map(|v| v as usize); + yield Ok(StreamFrame::Content(StreamContent::ToolResult { tool_use_id, content, is_error, remote, content_length })); } } "status" => { @@ -299,6 +303,8 @@ fn apply_content_frame( id, format!("Tool not enabled: capability '{required_cap}' was not advertised by this client"), true, + false, + None, ); }); return; @@ -321,11 +327,17 @@ fn apply_content_frame( tool_use_id, content, is_error, + remote, + content_length, } => { handle.update(move |state| { - state - .conversation - .add_tool_result(tool_use_id, content, is_error); + state.conversation.add_tool_result( + tool_use_id, + content, + is_error, + remote, + content_length, + ); }); } } diff --git a/crates/atuin-ai/src/tui/state.rs b/crates/atuin-ai/src/tui/state.rs index 69b35909..37200025 100644 --- a/crates/atuin-ai/src/tui/state.rs +++ b/crates/atuin-ai/src/tui/state.rs @@ -40,11 +40,16 @@ pub(crate) enum ConversationEvent { name: String, input: serde_json::Value, }, - /// Tool result (usually from server-side execution) + /// Tool result (from server-side or client-side execution) ToolResult { tool_use_id: String, content: String, is_error: bool, + /// Server-side results are stored in the DB; the client sends an opaque + /// reference (`remote: true`) instead of the full content. + remote: bool, + /// Approximate content length for token estimation of remote results. + content_length: Option<usize>, }, /// Out-of-band output from the system - not sent to the server OutOfBandOutput { @@ -193,15 +198,31 @@ impl Conversation { tool_use_id, content, is_error, + remote, + content_length, } => { - messages.push(serde_json::json!({ - "role": "user", - "content": [{ + let tool_result = if *remote { + let mut obj = serde_json::json!({ + "type": "tool_result", + "tool_use_id": tool_use_id, + "remote": true, + "is_error": is_error + }); + if let Some(len) = content_length { + obj["content_length"] = serde_json::json!(len); + } + obj + } else { + serde_json::json!({ "type": "tool_result", "tool_use_id": tool_use_id, "content": content, "is_error": is_error - }] + }) + }; + messages.push(serde_json::json!({ + "role": "user", + "content": [tool_result] })); i += 1; } @@ -299,11 +320,20 @@ impl Conversation { } /// Add a tool result event during streaming - pub fn add_tool_result(&mut self, tool_use_id: String, content: String, is_error: bool) { + pub fn add_tool_result( + &mut self, + tool_use_id: String, + content: String, + is_error: bool, + remote: bool, + content_length: Option<usize>, + ) { self.events.push(ConversationEvent::ToolResult { tool_use_id, content, is_error, + remote, + content_length, }); } @@ -563,7 +593,7 @@ impl Session { let content = outcome.format_for_llm(); let is_error = outcome.is_error(); self.conversation - .add_tool_result(tool_id.to_string(), content, is_error); + .add_tool_result(tool_id.to_string(), content, is_error, false, None); } /// Get the footer text for current mode diff --git a/crates/atuin-ai/src/tui/view/turn.rs b/crates/atuin-ai/src/tui/view/turn.rs index 6949236c..7369f151 100644 --- a/crates/atuin-ai/src/tui/view/turn.rs +++ b/crates/atuin-ai/src/tui/view/turn.rs @@ -159,6 +159,7 @@ impl<'a> TurnBuilder<'a> { tool_use_id, content, is_error, + .. } => { self.add_tool_result(tool_use_id, content, *is_error); } |