Skip to main content

rpi_ai/
types.rs

1//! Mirrors `packages/ai/src/types.ts` — the core LLM type contract every other
2//! crate consumes.
3//!
4//! These types are the boundary between the provider-agnostic agent loop and the
5//! provider wire formats. Everything the loop talks in (`Content`, `Message`,
6//! `AssistantMessage`, `AssistantMessageEvent`) is defined here; provider code
7//! only produces/consumes these.
8
9use serde::{Deserialize, Serialize};
10use std::collections::BTreeMap;
11
12/// A newtype wrapping a JSON Schema value. Tools carry their parameters schema here
13/// (produced by `schemars::JsonSchema` derive, stored as a `serde_json::Value` so it
14/// can be sent over the wire to providers verbatim). Mirrors the TypeBox `TSchema`
15/// parameter shape on the TS `Tool`.
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
17#[serde(transparent)]
18pub struct Schema(pub serde_json::Value);
19
20impl Schema {
21    pub fn new(value: serde_json::Value) -> Self {
22        Self(value)
23    }
24
25    pub fn empty_object() -> Self {
26        Self(serde_json::Value::Object(serde_json::Map::new()))
27    }
28
29    pub fn as_value(&self) -> &serde_json::Value {
30        &self.0
31    }
32
33    pub fn into_value(self) -> serde_json::Value {
34        self.0
35    }
36
37    pub fn is_empty(&self) -> bool {
38        self.0.is_null()
39            || (self.0.is_object() && self.0.as_object().map(|m| m.is_empty()).unwrap_or(true))
40    }
41}
42
43impl From<serde_json::Value> for Schema {
44    fn from(v: serde_json::Value) -> Self {
45        Self(v)
46    }
47}
48
49// ----------------------------------------------------------------------------
50// Content blocks — mirror TS `TextContent | ImageContent | ThinkingContent | ToolCall`
51// ----------------------------------------------------------------------------
52
53/// A text content block.
54#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
55#[serde(rename_all = "camelCase")]
56pub struct TextContent {
57    // `kind` renames to "type" so direct serialization carries the block
58    // discriminator. `default`: `Content` is `#[serde(tag="type")]`, and the
59    // tagged path consumes the `type` field as the tag before the payload is
60    // deserialized — without a default, re-reading a persisted `Content`
61    // failed with "missing field `type`" (sessions could be written but never
62    // restored). The unit-marker default is the correct value per variant.
63    #[serde(default, rename = "type")]
64    pub kind: TextContentType,
65    pub text: String,
66    /// Provider message-metadata signature (OpenAI Responses legacy id string or
67    /// `TextSignatureV1` JSON). Optional.
68    #[serde(default, skip_serializing_if = "Option::is_none")]
69    pub text_signature: Option<String>,
70}
71
72/// Marker so the `type` field serializes as the literal `"text"`.
73#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
74pub struct TextContentType;
75impl Serialize for TextContentType {
76    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
77        s.serialize_str("text")
78    }
79}
80impl<'de> Deserialize<'de> for TextContentType {
81    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
82        let s = String::deserialize(d)?;
83        if s != "text" {
84            return Err(serde::de::Error::custom(format!(
85                "expected \"text\", got {s:?}"
86            )));
87        }
88        Ok(Self)
89    }
90}
91
92/// Thinking / reasoning content.
93#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
94#[serde(rename_all = "camelCase")]
95pub struct ThinkingContent {
96    #[serde(default, rename = "type")]
97    pub kind: ThinkingContentType,
98    pub thinking: String,
99    #[serde(default, skip_serializing_if = "Option::is_none")]
100    pub thinking_signature: Option<String>,
101    /// When true, safety filters redacted the thinking; the opaque payload is in
102    /// `thinking_signature`.
103    #[serde(default, skip_serializing_if = "is_false")]
104    pub redacted: bool,
105}
106
107#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
108pub struct ThinkingContentType;
109impl Serialize for ThinkingContentType {
110    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
111        s.serialize_str("thinking")
112    }
113}
114impl<'de> Deserialize<'de> for ThinkingContentType {
115    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
116        let s = String::deserialize(d)?;
117        if s != "thinking" {
118            return Err(serde::de::Error::custom(format!(
119                "expected \"thinking\", got {s:?}"
120            )));
121        }
122        Ok(Self)
123    }
124}
125
126fn is_false(b: &bool) -> bool {
127    !*b
128}
129
130/// An image content block (base64-encoded).
131#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
132#[serde(rename_all = "camelCase")]
133pub struct ImageContent {
134    // `Content` is internally tagged and consumes `type` before deserializing
135    // this struct. Default the marker just like text/thinking content so
136    // persisted image blocks round-trip through JSONL.
137    #[serde(default, rename = "type")]
138    pub kind: ImageContentType,
139    pub data: String,
140    pub mime_type: String,
141}
142
143#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
144pub struct ImageContentType;
145impl Serialize for ImageContentType {
146    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
147        s.serialize_str("image")
148    }
149}
150impl<'de> Deserialize<'de> for ImageContentType {
151    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
152        let s = String::deserialize(d)?;
153        if s != "image" {
154            return Err(serde::de::Error::custom(format!(
155                "expected \"image\", got {s:?}"
156            )));
157        }
158        Ok(Self)
159    }
160}
161
162/// A tool call request from the model.
163#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
164#[serde(rename_all = "camelCase")]
165pub struct ToolCall {
166    #[serde(default, rename = "type")]
167    pub kind: ToolCallType,
168    pub id: String,
169    pub name: String,
170    /// LLM-supplied arguments. `Value::Object` in the happy path; we permit any
171    /// `Value` so partial-parse recovery has somewhere to park malformed input.
172    pub arguments: serde_json::Value,
173    /// Google-specific opaque signature for reusing thought context.
174    #[serde(default, skip_serializing_if = "Option::is_none")]
175    pub thought_signature: Option<String>,
176    /// OpenAI Responses namespace for dynamic/namespaced tools.
177    #[serde(default, skip_serializing_if = "Option::is_none")]
178    pub namespace: Option<String>,
179}
180
181#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
182pub struct ToolCallType;
183impl Serialize for ToolCallType {
184    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
185        s.serialize_str("toolCall")
186    }
187}
188impl<'de> Deserialize<'de> for ToolCallType {
189    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
190        let s = String::deserialize(d)?;
191        if s != "toolCall" {
192            return Err(serde::de::Error::custom(format!(
193                "expected \"toolCall\", got {s:?}"
194            )));
195        }
196        Ok(Self)
197    }
198}
199
200/// `TextContent | ImageContent | ThinkingContent | ToolCall` — assistant-side
201/// content blocks. User messages use only text/image; assistant messages may carry
202/// all four. Tagged via the `type` field so it round-trips the TS union exactly.
203#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
204#[serde(tag = "type", rename_all = "camelCase")]
205pub enum Content {
206    Text(TextContent),
207    Thinking(ThinkingContent),
208    Image(ImageContent),
209    ToolCall(ToolCall),
210}
211
212// Each concrete block already serializes its own `type` marker. Deriving an
213// internally tagged serializer here would emit that field a second time, which
214// raw JSON readers reject (a Value-based round-trip silently hides the bug).
215impl Serialize for Content {
216    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
217        match self {
218            Self::Text(block) => block.serialize(serializer),
219            Self::Thinking(block) => block.serialize(serializer),
220            Self::Image(block) => block.serialize(serializer),
221            Self::ToolCall(block) => block.serialize(serializer),
222        }
223    }
224}
225
226impl Content {
227    pub fn text<S: Into<String>>(s: S) -> Self {
228        Content::Text(TextContent {
229            kind: TextContentType,
230            text: s.into(),
231            text_signature: None,
232        })
233    }
234
235    pub fn thinking<S: Into<String>>(s: S) -> Self {
236        Content::Thinking(ThinkingContent {
237            kind: ThinkingContentType,
238            thinking: s.into(),
239            thinking_signature: None,
240            redacted: false,
241        })
242    }
243
244    pub fn tool_call<I: Into<String>, N: Into<String>>(
245        id: I,
246        name: N,
247        arguments: serde_json::Value,
248    ) -> Self {
249        Content::ToolCall(ToolCall {
250            kind: ToolCallType,
251            id: id.into(),
252            name: name.into(),
253            arguments,
254            thought_signature: None,
255            namespace: None,
256        })
257    }
258
259    /// Plain-text extraction (mirrors `contentText`): joins all `Text` blocks with
260    /// `\n`, dropping thinking/tool-call/image blocks.
261    pub fn text_only(content: &[Content], sep: &str) -> String {
262        content
263            .iter()
264            .filter_map(|c| match c {
265                Content::Text(t) => Some(t.text.as_str()),
266                _ => None,
267            })
268            .collect::<Vec<_>>()
269            .join(sep)
270    }
271}
272
273// ----------------------------------------------------------------------------
274// Usage + costs
275// ----------------------------------------------------------------------------
276
277#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
278#[serde(rename_all = "camelCase")]
279pub struct ModelCostRates {
280    pub input: f64,
281    pub output: f64,
282    pub cache_read: f64,
283    pub cache_write: f64,
284}
285
286#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
287#[serde(rename_all = "camelCase")]
288pub struct ModelCostTier {
289    #[serde(flatten)]
290    pub rates: ModelCostRates,
291    pub input_tokens_above: i64,
292}
293
294#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
295#[serde(rename_all = "camelCase")]
296pub struct ModelCost {
297    #[serde(flatten)]
298    pub rates: ModelCostRates,
299    #[serde(default, skip_serializing_if = "Vec::is_empty")]
300    pub tiers: Vec<ModelCostTier>,
301}
302
303impl Default for ModelCost {
304    fn default() -> Self {
305        Self {
306            rates: ModelCostRates::default(),
307            tiers: Vec::new(),
308        }
309    }
310}
311
312/// Per-request token usage + cost. Mirrors TS `Usage`.
313#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
314#[serde(rename_all = "camelCase")]
315pub struct Usage {
316    pub input: i64,
317    pub output: i64,
318    pub cache_read: i64,
319    pub cache_write: i64,
320    /// Subset of `cache_write` written with 1h retention (Anthropic only).
321    #[serde(default, skip_serializing_if = "Option::is_none")]
322    pub cache_write_1h: Option<i64>,
323    /// Reasoning tokens, when reported. Subset of `output`.
324    #[serde(default, skip_serializing_if = "Option::is_none")]
325    pub reasoning: Option<i64>,
326    pub total_tokens: i64,
327    pub cost: UsageCost,
328}
329
330#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
331#[serde(rename_all = "camelCase")]
332pub struct UsageCost {
333    pub input: f64,
334    pub output: f64,
335    pub cache_read: f64,
336    pub cache_write: f64,
337    pub total: f64,
338}
339
340impl Usage {
341    pub fn zero() -> Self {
342        Self {
343            input: 0,
344            output: 0,
345            cache_read: 0,
346            cache_write: 0,
347            cache_write_1h: None,
348            reasoning: None,
349            total_tokens: 0,
350            cost: UsageCost::default(),
351        }
352    }
353
354    pub fn add(&mut self, other: &Usage) {
355        self.input += other.input;
356        self.output += other.output;
357        self.cache_read += other.cache_read;
358        self.cache_write += other.cache_write;
359        if other.cache_write_1h.is_some() {
360            self.cache_write_1h =
361                Some(self.cache_write_1h.unwrap_or(0) + other.cache_write_1h.unwrap());
362        }
363        if other.reasoning.is_some() {
364            self.reasoning = Some(self.reasoning.unwrap_or(0) + other.reasoning.unwrap());
365        }
366        self.total_tokens += other.total_tokens;
367        self.cost.input += other.cost.input;
368        self.cost.output += other.cost.output;
369        self.cost.cache_read += other.cost.cache_read;
370        self.cost.cache_write += other.cost.cache_write;
371        self.cost.total += other.cost.total;
372    }
373
374    /// `usage.totalTokens || (input + output + cacheRead + cacheWrite)` — mirrors
375    /// `calculateContextTokens`.
376    pub fn context_tokens(&self) -> i64 {
377        let sum = self.total_tokens;
378        if sum > 0 {
379            sum
380        } else {
381            self.input + self.output + self.cache_read + self.cache_write
382        }
383    }
384}
385
386// ----------------------------------------------------------------------------
387// Stop reason + provider/api identifiers
388// ----------------------------------------------------------------------------
389
390#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
391#[serde(rename_all = "lowercase")]
392pub enum StopReason {
393    Pending,
394    Stop,
395    Length,
396    ToolUse,
397    Error,
398    Aborted,
399    Deferred,
400}
401
402impl StopReason {
403    pub fn as_str(self) -> &'static str {
404        match self {
405            StopReason::Pending => "pending",
406            StopReason::Stop => "stop",
407            StopReason::Length => "length",
408            StopReason::ToolUse => "toolUse",
409            StopReason::Error => "error",
410            StopReason::Aborted => "aborted",
411            StopReason::Deferred => "deferred",
412        }
413    }
414}
415
416/// Pi known APIs. The string fallback covers custom APIs; we keep a typed enum for
417/// the common ones and expose `Api::Other(String)` for the rest.
418#[derive(Debug, Clone, PartialEq, Eq, Hash)]
419pub enum Api {
420    OpenaiCompletions,
421    MistralConversations,
422    OpenaiResponses,
423    AzureOpenaiResponses,
424    OpenaiCodexResponses,
425    AnthropicMessages,
426    BedrockConverseStream,
427    GoogleGenerativeAi,
428    GoogleVertex,
429    PiMessages,
430    Faux,
431    Other(String),
432}
433
434impl Api {
435    pub fn as_str(&self) -> &str {
436        match self {
437            Api::OpenaiCompletions => "openai-completions",
438            Api::MistralConversations => "mistral-conversations",
439            Api::OpenaiResponses => "openai-responses",
440            Api::AzureOpenaiResponses => "azure-openai-responses",
441            Api::OpenaiCodexResponses => "openai-codex-responses",
442            Api::AnthropicMessages => "anthropic-messages",
443            Api::BedrockConverseStream => "bedrock-converse-stream",
444            Api::GoogleGenerativeAi => "google-generative-ai",
445            Api::GoogleVertex => "google-vertex",
446            Api::PiMessages => "pi-messages",
447            Api::Faux => "faux",
448            Api::Other(s) => s.as_str(),
449        }
450    }
451}
452
453impl Serialize for Api {
454    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
455        s.serialize_str(self.as_str())
456    }
457}
458impl<'de> Deserialize<'de> for Api {
459    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
460        let s = String::deserialize(d)?;
461        Ok(match s.as_str() {
462            "openai-completions" => Api::OpenaiCompletions,
463            "mistral-conversations" => Api::MistralConversations,
464            "openai-responses" => Api::OpenaiResponses,
465            "azure-openai-responses" => Api::AzureOpenaiResponses,
466            "openai-codex-responses" => Api::OpenaiCodexResponses,
467            "anthropic-messages" => Api::AnthropicMessages,
468            "bedrock-converse-stream" => Api::BedrockConverseStream,
469            "google-generative-ai" => Api::GoogleGenerativeAi,
470            "google-vertex" => Api::GoogleVertex,
471            "pi-messages" => Api::PiMessages,
472            "faux" => Api::Faux,
473            other => Api::Other(other.to_string()),
474        })
475    }
476}
477
478/// Provider identifier (string; keeps the TS `KnownProvider | string` open enum).
479pub type ProviderId = String;
480
481/// Input modalities a model accepts.
482#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
483#[serde(rename_all = "lowercase")]
484pub enum InputModality {
485    Text,
486    Image,
487}
488
489/// Thinking levels, mirroring TS `ThinkingLevel | "off"`.
490#[derive(
491    Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
492)]
493#[serde(rename_all = "lowercase")]
494pub enum ThinkingLevel {
495    #[default]
496    Off,
497    Minimal,
498    Low,
499    Medium,
500    High,
501    Xhigh,
502    Max,
503}
504
505/// Maps pi thinking levels to provider-specific values; `None` marks a level as
506/// unsupported. Mirrors `ThinkingLevelMap = Partial<Record<...>>`.
507pub type ThinkingLevelMap = BTreeMap<ThinkingLevel, Option<String>>;
508
509/// Custom token budgets per thinking level for token-budgeted reasoning models.
510/// Mirrors TS `ThinkingBudgets`; all fields optional (defaults live in
511/// `simple-options.ts::adjustMaxTokensForThinking`).
512#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
513#[serde(rename_all = "camelCase")]
514pub struct ThinkingBudgets {
515    #[serde(default, skip_serializing_if = "Option::is_none")]
516    pub minimal: Option<u64>,
517    #[serde(default, skip_serializing_if = "Option::is_none")]
518    pub low: Option<u64>,
519    #[serde(default, skip_serializing_if = "Option::is_none")]
520    pub medium: Option<u64>,
521    #[serde(default, skip_serializing_if = "Option::is_none")]
522    pub high: Option<u64>,
523}
524
525// ----------------------------------------------------------------------------
526// Deferred responses
527// ----------------------------------------------------------------------------
528
529/// A durable handle to a deferred provider response (long-poll APIs).
530#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
531#[serde(rename_all = "camelCase")]
532pub struct DeferredHandle {
533    pub provider: String,
534    pub model_id: String,
535    pub api: String,
536    pub id: String,
537    #[serde(default, skip_serializing_if = "Option::is_none")]
538    pub expires_at: Option<i64>,
539    #[serde(default, skip_serializing_if = "Option::is_none")]
540    pub poll_after_ms: Option<i64>,
541    #[serde(default, skip_serializing_if = "Option::is_none")]
542    pub data: Option<serde_json::Value>,
543}
544
545// ----------------------------------------------------------------------------
546// Messages — mirror TS `Message = UserMessage | AssistantMessage | ToolResultMessage`
547// ----------------------------------------------------------------------------
548
549/// User content may be a plain string or a list of text/image blocks.
550#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
551#[serde(untagged)]
552pub enum UserContent {
553    Text(String),
554    Blocks(Vec<Content>),
555}
556
557impl UserContent {
558    pub fn as_text(&self) -> Option<&str> {
559        match self {
560            UserContent::Text(s) => Some(s),
561            UserContent::Blocks(_) => None,
562        }
563    }
564}
565
566impl From<String> for UserContent {
567    fn from(s: String) -> Self {
568        UserContent::Text(s)
569    }
570}
571
572impl From<&str> for UserContent {
573    fn from(s: &str) -> Self {
574        UserContent::Text(s.to_string())
575    }
576}
577
578#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
579#[serde(rename_all = "camelCase")]
580pub struct UserMessage {
581    #[serde(rename = "role")]
582    pub role: UserRole,
583    pub content: UserContent,
584    pub timestamp: i64,
585}
586
587#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
588pub struct UserRole;
589impl Serialize for UserRole {
590    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
591        s.serialize_str("user")
592    }
593}
594impl<'de> Deserialize<'de> for UserRole {
595    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
596        let s = String::deserialize(d)?;
597        if s != "user" {
598            return Err(serde::de::Error::custom(format!(
599                "expected \"user\", got {s:?}"
600            )));
601        }
602        Ok(Self)
603    }
604}
605
606impl UserMessage {
607    pub fn new(content: impl Into<UserContent>, timestamp: i64) -> Self {
608        Self {
609            role: UserRole,
610            content: content.into(),
611            timestamp,
612        }
613    }
614}
615
616#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
617#[serde(rename_all = "camelCase")]
618pub struct ToolResultMessage {
619    #[serde(rename = "role")]
620    pub role: ToolResultRole,
621    pub tool_call_id: String,
622    pub tool_name: String,
623    /// Result content (text/image blocks).
624    pub content: Vec<Content>,
625    #[serde(default, skip_serializing_if = "Option::is_none")]
626    pub details: Option<serde_json::Value>,
627    #[serde(default, skip_serializing_if = "Option::is_none")]
628    pub usage: Option<Usage>,
629    #[serde(default, skip_serializing_if = "Vec::is_empty")]
630    pub added_tool_names: Vec<String>,
631    pub is_error: bool,
632    pub timestamp: i64,
633}
634
635#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
636pub struct ToolResultRole;
637impl Serialize for ToolResultRole {
638    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
639        s.serialize_str("toolResult")
640    }
641}
642impl<'de> Deserialize<'de> for ToolResultRole {
643    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
644        let s = String::deserialize(d)?;
645        if s != "toolResult" {
646            return Err(serde::de::Error::custom(format!(
647                "expected \"toolResult\", got {s:?}"
648            )));
649        }
650        Ok(Self)
651    }
652}
653
654/// Concrete assistant message. Mirrors TS `AssistantMessage`.
655#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
656#[serde(rename_all = "camelCase")]
657pub struct AssistantMessage {
658    #[serde(rename = "role")]
659    pub role: AssistantRole,
660    pub content: Vec<Content>,
661    pub api: Api,
662    pub provider: ProviderId,
663    pub model: String,
664    #[serde(default, skip_serializing_if = "Option::is_none")]
665    pub response_model: Option<String>,
666    #[serde(default, skip_serializing_if = "Option::is_none")]
667    pub response_id: Option<String>,
668    pub usage: Usage,
669    pub stop_reason: StopReason,
670    #[serde(default, skip_serializing_if = "Option::is_none")]
671    pub deferred: Option<DeferredHandle>,
672    #[serde(default, skip_serializing_if = "Option::is_none")]
673    pub error_message: Option<String>,
674    #[serde(default, skip_serializing_if = "Option::is_none")]
675    pub raw_stop_reason: Option<String>,
676    #[serde(default, skip_serializing_if = "Option::is_none")]
677    pub end_turn: Option<bool>,
678    pub timestamp: i64,
679}
680
681#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
682pub struct AssistantRole;
683impl Serialize for AssistantRole {
684    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
685        s.serialize_str("assistant")
686    }
687}
688impl<'de> Deserialize<'de> for AssistantRole {
689    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
690        let s = String::deserialize(d)?;
691        if s != "assistant" {
692            return Err(serde::de::Error::custom(format!(
693                "expected \"assistant\", got {s:?}"
694            )));
695        }
696        Ok(Self)
697    }
698}
699
700impl AssistantMessage {
701    /// Construct a new assistant message with zeroed usage and `Pending` stop reason,
702    /// ready for the stream mapper to mutate per-event.
703    pub fn empty(
704        api: Api,
705        provider: impl Into<String>,
706        model: impl Into<String>,
707        timestamp: i64,
708    ) -> Self {
709        Self {
710            role: AssistantRole,
711            content: Vec::new(),
712            api,
713            provider: provider.into(),
714            model: model.into(),
715            response_model: None,
716            response_id: None,
717            usage: Usage::zero(),
718            stop_reason: StopReason::Pending,
719            deferred: None,
720            error_message: None,
721            raw_stop_reason: None,
722            end_turn: None,
723            timestamp,
724        }
725    }
726
727    /// Convenience: an error/aborted terminal message (used by providers + faux).
728    pub fn terminal(
729        api: Api,
730        provider: impl Into<String>,
731        model: impl Into<String>,
732        stop_reason: StopReason,
733        error_message: impl Into<String>,
734        timestamp: i64,
735    ) -> Self {
736        Self {
737            role: AssistantRole,
738            content: Vec::new(),
739            api,
740            provider: provider.into(),
741            model: model.into(),
742            response_model: None,
743            response_id: None,
744            usage: Usage::zero(),
745            stop_reason,
746            deferred: None,
747            error_message: Some(error_message.into()),
748            raw_stop_reason: None,
749            end_turn: None,
750            timestamp,
751        }
752    }
753}
754
755/// `Message = UserMessage | AssistantMessage | ToolResultMessage`. The base
756/// provider-facing message type.
757#[derive(Debug, Clone, PartialEq, Serialize)]
758#[serde(tag = "role", rename_all = "camelCase")]
759pub enum Message {
760    User(UserMessage),
761    Assistant(Box<AssistantMessage>),
762    ToolResult(Box<ToolResultMessage>),
763}
764
765// The concrete message structs intentionally retain their `role` field so
766// they serialize exactly like Pi's TypeScript messages. An internally tagged
767// enum would consume that field before deserializing the struct variant,
768// causing `missing field role` for otherwise valid messages. Decode the tag
769// explicitly and pass the complete object through to the concrete type.
770impl<'de> Deserialize<'de> for Message {
771    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
772        let value = serde_json::Value::deserialize(deserializer)?;
773        let role = value
774            .get("role")
775            .and_then(serde_json::Value::as_str)
776            .ok_or_else(|| serde::de::Error::missing_field("role"))?;
777        match role {
778            "user" => serde_json::from_value(value)
779                .map(Message::User)
780                .map_err(serde::de::Error::custom),
781            "assistant" => serde_json::from_value(value)
782                .map(|message| Message::Assistant(Box::new(message)))
783                .map_err(serde::de::Error::custom),
784            "toolResult" => serde_json::from_value(value)
785                .map(|message| Message::ToolResult(Box::new(message)))
786                .map_err(serde::de::Error::custom),
787            other => Err(serde::de::Error::unknown_variant(
788                other,
789                &["user", "assistant", "toolResult"],
790            )),
791        }
792    }
793}
794
795impl From<UserMessage> for Message {
796    fn from(m: UserMessage) -> Self {
797        Message::User(m)
798    }
799}
800impl From<AssistantMessage> for Message {
801    fn from(m: AssistantMessage) -> Self {
802        Message::Assistant(Box::new(m))
803    }
804}
805impl From<ToolResultMessage> for Message {
806    fn from(m: ToolResultMessage) -> Self {
807        Message::ToolResult(Box::new(m))
808    }
809}
810
811// ----------------------------------------------------------------------------
812// Tool + Context
813// ----------------------------------------------------------------------------
814
815/// OpenAI grammar variants — mirror TS `GrammarFormat` (used by constrained sampling).
816#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
817pub enum GrammarFormat {
818    #[serde(rename = "openai_lark")]
819    OpenaiLark,
820    #[serde(rename = "openai_regex")]
821    OpenaiRegex,
822}
823
824pub type GrammarVariants = BTreeMap<GrammarFormat, String>;
825
826/// Provider-side constrained sampling config (json_schema strict, or grammar).
827#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
828#[serde(tag = "type", rename_all = "snake_case")]
829pub enum ConstrainedSamplingConfig {
830    JsonSchema {
831        #[serde(rename = "strict")]
832        strict: ConstrainedStrictness,
833    },
834    Grammar {
835        variants: GrammarVariants,
836    },
837}
838
839#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
840#[serde(rename_all = "lowercase")]
841pub enum ConstrainedStrictness {
842    Prefer,
843    Require,
844}
845
846/// A tool definition. Mirrors TS `Tool`.
847#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
848#[serde(rename_all = "camelCase")]
849pub struct Tool {
850    pub name: String,
851    pub description: String,
852    pub parameters: Schema,
853    #[serde(default, skip_serializing_if = "Option::is_none")]
854    pub constrained_sampling: Option<ConstrainedSamplingConfig>,
855}
856
857/// The context handed to a provider stream call. Mirrors TS `Context`.
858#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
859#[serde(rename_all = "camelCase")]
860pub struct Context {
861    #[serde(default, skip_serializing_if = "Option::is_none")]
862    pub system_prompt: Option<String>,
863    pub messages: Vec<Message>,
864    #[serde(default, skip_serializing_if = "Vec::is_empty")]
865    pub tools: Vec<Tool>,
866}
867
868impl Context {
869    pub fn new(messages: Vec<Message>) -> Self {
870        Self {
871            system_prompt: None,
872            messages,
873            tools: Vec::new(),
874        }
875    }
876}
877
878// ----------------------------------------------------------------------------
879// AssistantMessageEvent — the streaming protocol union
880// ----------------------------------------------------------------------------
881
882/// `AssistantMessageEvent` tagged union, mirrors TS exactly.
883///
884/// Providers emit events in this order: `Start` is emitted once before any per-block
885/// events, then a `{Start,Delta*,End}` triple per content block (text / thinking /
886/// toolcall), then a terminal `Done` (carrying the final successful message) or
887/// `Error` (carrying a terminal message with `StopReason::Error`/`Aborted`).
888///
889/// Every non-terminal event references the in-progress `AssistantMessage` via
890/// `partial` so subscribers can render a live view. We wrap the partial in `Arc`
891/// so the event is cheap to clone across broadcast subscribers.
892#[derive(Debug, Clone, PartialEq)]
893pub enum AssistantMessageEvent {
894    Start {
895        partial: std::sync::Arc<AssistantMessage>,
896    },
897    TextStart {
898        content_index: usize,
899        partial: std::sync::Arc<AssistantMessage>,
900    },
901    TextDelta {
902        content_index: usize,
903        delta: String,
904        partial: std::sync::Arc<AssistantMessage>,
905    },
906    TextEnd {
907        content_index: usize,
908        content: String,
909        partial: std::sync::Arc<AssistantMessage>,
910    },
911    ThinkingStart {
912        content_index: usize,
913        partial: std::sync::Arc<AssistantMessage>,
914    },
915    ThinkingDelta {
916        content_index: usize,
917        delta: String,
918        partial: std::sync::Arc<AssistantMessage>,
919    },
920    ThinkingEnd {
921        content_index: usize,
922        content: String,
923        partial: std::sync::Arc<AssistantMessage>,
924    },
925    ToolCallStart {
926        content_index: usize,
927        partial: std::sync::Arc<AssistantMessage>,
928    },
929    ToolCallDelta {
930        content_index: usize,
931        delta: String,
932        partial: std::sync::Arc<AssistantMessage>,
933    },
934    ToolCallEnd {
935        content_index: usize,
936        tool_call: ToolCall,
937        partial: std::sync::Arc<AssistantMessage>,
938    },
939    Done {
940        reason: DoneReason,
941        message: AssistantMessage,
942    },
943    Error {
944        reason: ErrorReason,
945        error: AssistantMessage,
946    },
947}
948
949/// The terminal reason on a `Done` event — the success subset of `StopReason`.
950#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
951pub enum DoneReason {
952    Stop,
953    Length,
954    ToolUse,
955    Deferred,
956}
957
958impl DoneReason {
959    pub fn as_str(self) -> &'static str {
960        match self {
961            DoneReason::Stop => "stop",
962            DoneReason::Length => "length",
963            DoneReason::ToolUse => "toolUse",
964            DoneReason::Deferred => "deferred",
965        }
966    }
967}
968
969impl From<DoneReason> for StopReason {
970    fn from(r: DoneReason) -> Self {
971        match r {
972            DoneReason::Stop => StopReason::Stop,
973            DoneReason::Length => StopReason::Length,
974            DoneReason::ToolUse => StopReason::ToolUse,
975            DoneReason::Deferred => StopReason::Deferred,
976        }
977    }
978}
979
980/// The terminal reason on an `Error` event — the failure subset of `StopReason`.
981#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
982pub enum ErrorReason {
983    Aborted,
984    Error,
985}
986
987impl ErrorReason {
988    pub fn as_str(self) -> &'static str {
989        match self {
990            ErrorReason::Aborted => "aborted",
991            ErrorReason::Error => "error",
992        }
993    }
994}
995
996impl From<ErrorReason> for StopReason {
997    fn from(r: ErrorReason) -> Self {
998        match r {
999            ErrorReason::Aborted => StopReason::Aborted,
1000            ErrorReason::Error => StopReason::Error,
1001        }
1002    }
1003}
1004
1005impl AssistantMessageEvent {
1006    /// Returns the `type` tag, mirroring the TS discriminator.
1007    pub fn type_tag(&self) -> &'static str {
1008        match self {
1009            AssistantMessageEvent::Start { .. } => "start",
1010            AssistantMessageEvent::TextStart { .. } => "text_start",
1011            AssistantMessageEvent::TextDelta { .. } => "text_delta",
1012            AssistantMessageEvent::TextEnd { .. } => "text_end",
1013            AssistantMessageEvent::ThinkingStart { .. } => "thinking_start",
1014            AssistantMessageEvent::ThinkingDelta { .. } => "thinking_delta",
1015            AssistantMessageEvent::ThinkingEnd { .. } => "thinking_end",
1016            AssistantMessageEvent::ToolCallStart { .. } => "toolcall_start",
1017            AssistantMessageEvent::ToolCallDelta { .. } => "toolcall_delta",
1018            AssistantMessageEvent::ToolCallEnd { .. } => "toolcall_end",
1019            AssistantMessageEvent::Done { .. } => "done",
1020            AssistantMessageEvent::Error { .. } => "error",
1021        }
1022    }
1023
1024    /// True for the two terminal variants.
1025    pub fn is_terminal(&self) -> bool {
1026        matches!(
1027            self,
1028            AssistantMessageEvent::Done { .. } | AssistantMessageEvent::Error { .. }
1029        )
1030    }
1031
1032    /// Snapshot of the partial message for live rendering. Terminal events do not
1033    /// carry a `partial`; returns the terminal message itself for `Done`/`Error`.
1034    pub fn partial(&self) -> &AssistantMessage {
1035        match self {
1036            AssistantMessageEvent::Start { partial }
1037            | AssistantMessageEvent::TextStart { partial, .. }
1038            | AssistantMessageEvent::TextDelta { partial, .. }
1039            | AssistantMessageEvent::TextEnd { partial, .. }
1040            | AssistantMessageEvent::ThinkingStart { partial, .. }
1041            | AssistantMessageEvent::ThinkingDelta { partial, .. }
1042            | AssistantMessageEvent::ThinkingEnd { partial, .. }
1043            | AssistantMessageEvent::ToolCallStart { partial, .. }
1044            | AssistantMessageEvent::ToolCallDelta { partial, .. }
1045            | AssistantMessageEvent::ToolCallEnd { partial, .. } => partial,
1046            AssistantMessageEvent::Done { message, .. } => message,
1047            AssistantMessageEvent::Error { error, .. } => error,
1048        }
1049    }
1050}
1051
1052// serde for AssistantMessageEvent — tagged on "type" matching the wire protocol.
1053impl Serialize for AssistantMessageEvent {
1054    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
1055        use serde::ser::SerializeMap;
1056        let tag = self.type_tag();
1057        match self {
1058            AssistantMessageEvent::Start { partial } => {
1059                let mut m = s.serialize_map(Some(2))?;
1060                m.serialize_entry("type", tag)?;
1061                m.serialize_entry("partial", &**partial)?;
1062                m.end()
1063            }
1064            AssistantMessageEvent::TextStart {
1065                content_index,
1066                partial,
1067            }
1068            | AssistantMessageEvent::ThinkingStart {
1069                content_index,
1070                partial,
1071            }
1072            | AssistantMessageEvent::ToolCallStart {
1073                content_index,
1074                partial,
1075            } => {
1076                let mut m = s.serialize_map(Some(3))?;
1077                m.serialize_entry("type", tag)?;
1078                m.serialize_entry("contentIndex", content_index)?;
1079                m.serialize_entry("partial", &**partial)?;
1080                m.end()
1081            }
1082            AssistantMessageEvent::TextDelta {
1083                content_index,
1084                delta,
1085                partial,
1086            }
1087            | AssistantMessageEvent::ThinkingDelta {
1088                content_index,
1089                delta,
1090                partial,
1091            } => {
1092                let mut m = s.serialize_map(Some(4))?;
1093                m.serialize_entry("type", tag)?;
1094                m.serialize_entry("contentIndex", content_index)?;
1095                m.serialize_entry("delta", delta)?;
1096                m.serialize_entry("partial", &**partial)?;
1097                m.end()
1098            }
1099            AssistantMessageEvent::ToolCallDelta {
1100                content_index,
1101                delta,
1102                partial,
1103            } => {
1104                let mut m = s.serialize_map(Some(4))?;
1105                m.serialize_entry("type", tag)?;
1106                m.serialize_entry("contentIndex", content_index)?;
1107                m.serialize_entry("delta", delta)?;
1108                m.serialize_entry("partial", &**partial)?;
1109                m.end()
1110            }
1111            AssistantMessageEvent::TextEnd {
1112                content_index,
1113                content,
1114                partial,
1115            }
1116            | AssistantMessageEvent::ThinkingEnd {
1117                content_index,
1118                content,
1119                partial,
1120            } => {
1121                let mut m = s.serialize_map(Some(4))?;
1122                m.serialize_entry("type", tag)?;
1123                m.serialize_entry("contentIndex", content_index)?;
1124                m.serialize_entry("content", content)?;
1125                m.serialize_entry("partial", &**partial)?;
1126                m.end()
1127            }
1128            AssistantMessageEvent::ToolCallEnd {
1129                content_index,
1130                tool_call,
1131                partial,
1132            } => {
1133                let mut m = s.serialize_map(Some(4))?;
1134                m.serialize_entry("type", tag)?;
1135                m.serialize_entry("contentIndex", content_index)?;
1136                m.serialize_entry("toolCall", tool_call)?;
1137                m.serialize_entry("partial", &**partial)?;
1138                m.end()
1139            }
1140            AssistantMessageEvent::Done { reason, message } => {
1141                let mut m = s.serialize_map(Some(3))?;
1142                m.serialize_entry("type", tag)?;
1143                m.serialize_entry("reason", reason.as_str())?;
1144                m.serialize_entry("message", message)?;
1145                m.end()
1146            }
1147            AssistantMessageEvent::Error { reason, error } => {
1148                let mut m = s.serialize_map(Some(3))?;
1149                m.serialize_entry("type", tag)?;
1150                m.serialize_entry("reason", reason.as_str())?;
1151                m.serialize_entry("error", error)?;
1152                m.end()
1153            }
1154        }
1155    }
1156}