Skip to main content

rpi_agent/
types.rs

1//! Mirrors `packages/agent/src/types.ts` β€” the public type contract of the
2//! agent layer: tool results, hook payloads, agent state, events.
3//!
4//! The TS source declares `AgentMessage = Message | Custom` via declaration
5//! merging; the Rust port models it as an open [`AgentMessage`] enum (see
6//! [`crate::message`]). Everything else here is a straight port of the TS
7//! interfaces, adapted to Rust ownership/async idioms.
8
9use rpi_ai::types::{AssistantMessage, ImageContent, TextContent, ToolResultMessage, Usage};
10use std::collections::HashSet;
11use std::sync::Arc;
12
13use crate::message::AgentMessage;
14
15/// How a batch of tool calls from one assistant message are executed.
16///
17/// - `Sequential`: each call is prepared, executed, finalized before the next starts.
18/// - `Parallel`: calls are prepared sequentially, then allowed tools execute
19///   concurrently. `tool_execution_end` fires in completion order; tool-result
20///   `MessageEnd` fires later in assistant source order.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
22pub enum ToolExecutionMode {
23    Sequential,
24    #[default]
25    Parallel,
26}
27
28/// Controls how many queued user messages are injected at a drain point.
29/// `All` drains every queued message; `OneAtATime` drains only the oldest.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
31pub enum QueueMode {
32    All,
33    #[default]
34    OneAtATime,
35}
36
37/// Result returned from a tool `execute`. Mirrors TS `AgentToolResult<T>`:
38/// `content` goes back to the model, `details` are structured log/UI payload,
39/// `usage` is optionally reported, and `terminate` hints the batch should stop.
40#[derive(Debug, Clone, Default, serde::Serialize)]
41#[serde(rename_all = "camelCase")]
42pub struct AgentToolResult {
43    pub content: Vec<TextContentOrImage>,
44    pub details: serde_json::Value,
45    pub usage: Option<Usage>,
46    pub added_tool_names: Vec<String>,
47    pub terminate: bool,
48}
49/// an enum so tool authors stay within the provider-content language without
50/// pulling in the full `Content` union (which adds Thinking/ToolCall).
51#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
52#[serde(untagged)]
53pub enum TextContentOrImage {
54    Text(TextContent),
55    Image(ImageContent),
56}
57
58impl TextContentOrImage {
59    pub fn text<S: Into<String>>(s: S) -> Self {
60        TextContentOrImage::Text(TextContent {
61            kind: rpi_ai::types::TextContentType,
62            text: s.into(),
63            text_signature: None,
64        })
65    }
66}
67
68impl AgentToolResult {
69    /// Convenience: a single text block, empty details.
70    pub fn text(message: impl Into<String>) -> Self {
71        Self {
72            content: vec![TextContentOrImage::text(message)],
73            details: serde_json::Value::Null,
74            usage: None,
75            added_tool_names: Vec::new(),
76            terminate: false,
77        }
78    }
79
80    /// Convenience: an error text block. `is_error` is carried on the
81    /// `ToolResultMessage`, not the result; this just builds the content.
82    pub fn error_text(message: impl Into<String>) -> Self {
83        Self::text(message)
84    }
85
86    pub fn into_content(self) -> Vec<rpi_ai::types::Content> {
87        self.content
88            .into_iter()
89            .map(|c| match c {
90                TextContentOrImage::Text(t) => rpi_ai::types::Content::Text(t),
91                TextContentOrImage::Image(i) => rpi_ai::types::Content::Image(i),
92            })
93            .collect()
94    }
95}
96
97impl From<AgentToolResult> for Result<AgentToolResult, crate::AgentError> {
98    fn from(r: AgentToolResult) -> Self {
99        Ok(r)
100    }
101}
102
103/// Partial result pushed by a tool's `on_update` callback during execution.
104/// Mirrors TS `AgentToolUpdateCallback<T>` payload.
105pub type ToolResultPartial = AgentToolResult;
106
107/// Result of a `before_tool_call` hook. `block` prevents execution; the loop
108/// emits an error tool result with `reason` (or a default) instead. `terminate`
109/// participates in the batch early-termination rule (all results must set it).
110///
111/// `args` is the Rust equivalent of TS `beforeToolCall` mutating the validated
112/// args object in place: JS callbacks receive `args` by reference and write to
113/// it; Rust hands the hook an immutable `&serde_json::Value`, so to rewrite the
114/// args the hook returns them here. Replacement args are applied **without
115/// re-validation** β€” mirroring TS, where the mutation happens after
116/// `validateToolArguments` and is never re-checked. `None` keeps the validated
117/// args.
118#[derive(Debug, Clone, Default)]
119pub struct BeforeToolCallResult {
120    pub block: bool,
121    pub reason: Option<String>,
122    pub terminate: bool,
123    pub args: Option<serde_json::Value>,
124}
125
126/// Partial override returned from `after_tool_call`. Field-by-field merge:
127/// provided values replace the executed result's fields; omitted fields keep
128/// the original. No deep merge.
129#[derive(Debug, Clone, Default)]
130pub struct AfterToolCallResult {
131    pub content: Option<Vec<TextContentOrImage>>,
132    pub details: Option<serde_json::Value>,
133    pub is_error: Option<bool>,
134    pub usage: Option<Usage>,
135    pub terminate: Option<bool>,
136}
137
138/// Context passed to `before_tool_call`. Mirrors TS `BeforeToolCallContext`.
139pub struct BeforeToolCallContext<'a> {
140    pub assistant_message: &'a AssistantMessage,
141    pub tool_call: &'a rpi_ai::types::ToolCall,
142    pub args: &'a serde_json::Value,
143    pub context: &'a AgentContext,
144}
145
146/// Context passed to `after_tool_call`. Mirrors TS `AfterToolCallContext`.
147pub struct AfterToolCallContext<'a> {
148    pub assistant_message: &'a AssistantMessage,
149    pub tool_call: &'a rpi_ai::types::ToolCall,
150    pub args: &'a serde_json::Value,
151    pub result: &'a AgentToolResult,
152    pub is_error: bool,
153    pub context: &'a AgentContext,
154}
155
156/// Context passed to `should_stop_after_turn` / `prepare_next_turn`.
157pub struct ShouldStopAfterTurnContext<'a> {
158    pub message: &'a AssistantMessage,
159    pub tool_results: &'a [ToolResultMessage],
160    pub context: &'a AgentContext,
161    pub new_messages: &'a [AgentMessage],
162}
163
164/// A context snapshot handed to the low-level loop. Mirrors TS `AgentContext`.
165#[derive(Clone, Default)]
166pub struct AgentContext {
167    pub system_prompt: String,
168    pub messages: Vec<AgentMessage>,
169    pub tools: Vec<Arc<dyn crate::agent_tool::AgentTool>>,
170}
171
172impl std::fmt::Debug for AgentContext {
173    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174        f.debug_struct("AgentContext")
175            .field("system_prompt", &self.system_prompt)
176            .field("messages", &self.messages)
177            .field(
178                "tools",
179                &self
180                    .tools
181                    .iter()
182                    .map(|t| t.schema().name.as_str())
183                    .collect::<Vec<_>>(),
184            )
185            .finish()
186    }
187}
188
189impl AgentContext {
190    pub fn new(messages: Vec<AgentMessage>) -> Self {
191        Self {
192            system_prompt: String::new(),
193            messages,
194            tools: Vec::new(),
195        }
196    }
197}
198
199/// Replacement runtime state returned by `prepare_next_turn`. `None` on any
200/// field means "keep current".
201#[derive(Debug, Clone, Default)]
202pub struct AgentLoopTurnUpdate {
203    pub context: Option<AgentContext>,
204    pub model: Option<rpi_ai::model::Model>,
205    pub thinking_level: Option<rpi_ai::types::ThinkingLevel>,
206}
207
208/// Public agent state snapshot. Mirrors TS `AgentState` (the readable subset).
209/// `tools`/`messages` clone on read so callers can't mutate internal state.
210#[derive(Clone)]
211pub struct AgentState {
212    pub system_prompt: String,
213    pub model: rpi_ai::model::Model,
214    pub thinking_level: rpi_ai::types::ThinkingLevel,
215    pub tools: Vec<Arc<dyn crate::agent_tool::AgentTool>>,
216    pub messages: Vec<AgentMessage>,
217    pub is_streaming: bool,
218    /// The in-flight assistant message, shared so the per-delta `MessageUpdate`
219    /// does not deep-clone it. `None` unless an assistant message is streaming.
220    pub streaming_message: Option<std::sync::Arc<AssistantMessage>>,
221    pub pending_tool_calls: HashSet<String>,
222    pub error_message: Option<String>,
223}
224
225impl std::fmt::Debug for AgentState {
226    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
227        f.debug_struct("AgentState")
228            .field("system_prompt", &self.system_prompt)
229            .field("model", &self.model)
230            .field("thinking_level", &self.thinking_level)
231            .field(
232                "tools",
233                &self
234                    .tools
235                    .iter()
236                    .map(|t| t.schema().name.as_str())
237                    .collect::<Vec<_>>(),
238            )
239            .field("messages", &self.messages)
240            .field("is_streaming", &self.is_streaming)
241            .field("streaming_message", &self.streaming_message)
242            .field("pending_tool_calls", &self.pending_tool_calls)
243            .field("error_message", &self.error_message)
244            .finish()
245    }
246}
247
248impl Default for AgentState {
249    fn default() -> Self {
250        Self {
251            system_prompt: String::new(),
252            model: default_model(),
253            thinking_level: rpi_ai::types::ThinkingLevel::Off,
254            tools: Vec::new(),
255            messages: Vec::new(),
256            is_streaming: false,
257            streaming_message: None,
258            pending_tool_calls: HashSet::new(),
259            error_message: None,
260        }
261    }
262}
263
264/// The placeholder model used when none is configured. Mirrors TS `DEFAULT_MODEL`.
265pub(crate) fn default_model() -> rpi_ai::model::Model {
266    rpi_ai::model::Model::new(
267        "unknown",
268        "unknown",
269        rpi_ai::types::Api::Other("unknown".into()),
270        "unknown",
271        "",
272    )
273}
274
275#[cfg(test)]
276mod result_codec_tests {
277    use super::*;
278    use serde_json::json;
279    use std::io::{self, Write};
280
281    #[test]
282    fn canonical_result_shape_keeps_null_usage_and_separate_error_authority() {
283        let value = serde_json::to_value(AgentToolResult::error_text("failed")).unwrap();
284        assert_eq!(
285            value,
286            json!({
287                "content": [{"type":"text", "text":"failed"}],
288                "details": null, "usage": null, "addedToolNames": [], "terminate": false,
289            })
290        );
291    }
292
293    #[test]
294    fn canonical_text_signatures_preserve_native_strings_exactly() {
295        // The native field is Option<String>, including JSON-formatted signatures.
296        for signature in [
297            "primitive-签名\"\\\n",
298            "123",
299            r#"{ "version": 1, "id": "sig-πŸ˜€", "nested": [true, null] }"#,
300        ] {
301            let result = AgentToolResult {
302                content: vec![TextContentOrImage::Text(TextContent {
303                    kind: rpi_ai::types::TextContentType,
304                    text: "η•ŒπŸ˜€\"\\\n".into(),
305                    text_signature: Some(signature.into()),
306                })],
307                ..Default::default()
308            };
309            let encoded = serde_json::to_string(&result).unwrap();
310            let value: serde_json::Value = serde_json::from_str(&encoded).unwrap();
311            assert_eq!(
312                value["content"][0],
313                json!({
314                    "type":"text", "text":"η•ŒπŸ˜€\"\\\n", "textSignature":signature,
315                })
316            );
317            assert_eq!(
318                value["content"][0]["textSignature"]
319                    .as_str()
320                    .unwrap()
321                    .as_bytes(),
322                signature.as_bytes()
323            );
324        }
325    }
326
327    #[test]
328    fn canonical_result_preserves_images_business_details_usage_and_batch_hints() {
329        let details = json!({"runtimeKey":"customer", "runtimeHandle":{"id":7},
330            "__rpi":{"owner":"business"}, "__rpiRuntimeKey":["nested", {"runtimeKey":true}]});
331        let mut usage = Usage::zero();
332        usage.input = 7;
333        usage.output = 3;
334        usage.cache_read = 2;
335        usage.cache_write = 4;
336        usage.cache_write_1h = Some(1);
337        usage.reasoning = Some(2);
338        usage.total_tokens = 16;
339        usage.cost.input = 0.5;
340        usage.cost.total = 0.75;
341        let result = AgentToolResult {
342            content: vec![TextContentOrImage::Image(ImageContent {
343                kind: rpi_ai::types::ImageContentType,
344                data: "AAEC+/==".into(),
345                mime_type: "image/png".into(),
346            })],
347            details: details.clone(),
348            usage: Some(usage),
349            added_tool_names: vec!["read".into(), "read".into(), "写".into()],
350            terminate: true,
351        };
352        assert_eq!(
353            serde_json::to_value(&result).unwrap(),
354            json!({
355                "content":[{"type":"image", "data":"AAEC+/==", "mimeType":"image/png"}],
356                "details":details,
357                "usage":{"input":7,"output":3,"cacheRead":2,"cacheWrite":4,
358                    "cacheWrite1h":1,"reasoning":2,"totalTokens":16,
359                    "cost":{"input":0.5,"output":0.0,"cacheRead":0.0,"cacheWrite":0.0,"total":0.75}},
360                "addedToolNames":["read","read","写"], "terminate":true,
361            })
362        );
363    }
364
365    #[test]
366    fn borrowed_writer_counts_large_unicode_and_image_without_a_payload_value() {
367        #[derive(Default)]
368        struct ByteCounter {
369            bytes: usize,
370            chunks: usize,
371        }
372        impl Write for ByteCounter {
373            fn write(&mut self, chunk: &[u8]) -> io::Result<usize> {
374                self.bytes += chunk.len();
375                self.chunks += 1;
376                Ok(chunk.len())
377            }
378            fn flush(&mut self) -> io::Result<()> {
379                Ok(())
380            }
381        }
382        let mut result = AgentToolResult::text("η•ŒπŸ˜€Γ©\"\\\n".repeat(150_000));
383        result.content.push(TextContentOrImage::Image(ImageContent {
384            kind: rpi_ai::types::ImageContentType,
385            data: "aGVsbG8=".repeat(150_000).into(),
386            mime_type: "image/webp".into(),
387        }));
388        let mut counter = ByteCounter::default();
389        serde_json::to_writer(&mut counter, &result).unwrap();
390        let encoded = serde_json::to_string(&result).unwrap();
391        assert!(counter.bytes > 1024 * 1024);
392        assert!(counter.chunks > 10);
393        assert_eq!(counter.bytes, encoded.len());
394        assert!(encoded.len() > encoded.chars().count());
395        assert_eq!(result.content.len(), 2);
396    }
397
398    #[test]
399    fn borrowed_writer_propagates_io_failure() {
400        struct FailingWriter;
401        impl Write for FailingWriter {
402            fn write(&mut self, _: &[u8]) -> io::Result<usize> {
403                Err(io::Error::other("journal unavailable"))
404            }
405            fn flush(&mut self) -> io::Result<()> {
406                Ok(())
407            }
408        }
409        let result = AgentToolResult::text("still owned");
410        let error = serde_json::to_writer(FailingWriter, &result).unwrap_err();
411        assert!(error.is_io());
412        assert!(error.to_string().contains("journal unavailable"));
413        assert_eq!(
414            result.content,
415            vec![TextContentOrImage::text("still owned")]
416        );
417    }
418}