1use std::future::Future;
5use std::pin::Pin;
6use std::{
7 any::TypeId,
8 collections::HashMap,
9 sync::{LazyLock, Mutex},
10};
11
12use futures_core::Stream;
13use serde::{Deserialize, Serialize};
14
15use zeph_common::ToolName;
16
17pub use zeph_common::ToolDefinition;
18
19use crate::embed::owned_strs;
20use crate::error::LlmError;
21
22static SCHEMA_CACHE: LazyLock<Mutex<HashMap<TypeId, (serde_json::Value, String)>>> =
23 LazyLock::new(|| Mutex::new(HashMap::new()));
24
25pub(crate) fn cached_schema<T: schemars::JsonSchema + 'static>()
31-> Result<(serde_json::Value, String), crate::LlmError> {
32 let type_id = TypeId::of::<T>();
33 if let Ok(cache) = SCHEMA_CACHE.lock()
34 && let Some(entry) = cache.get(&type_id)
35 {
36 return Ok(entry.clone());
37 }
38 let schema = schemars::schema_for!(T);
39 let value = serde_json::to_value(&schema)
40 .map_err(|e| crate::LlmError::StructuredParse(e.to_string()))?;
41 let pretty = serde_json::to_string_pretty(&schema)
42 .map_err(|e| crate::LlmError::StructuredParse(e.to_string()))?;
43 if let Ok(mut cache) = SCHEMA_CACHE.lock() {
44 cache.insert(type_id, (value.clone(), pretty.clone()));
45 }
46 Ok((value, pretty))
47}
48
49pub(crate) fn short_type_name<T: ?Sized>() -> &'static str {
63 std::any::type_name::<T>()
64 .rsplit("::")
65 .next()
66 .unwrap_or("Output")
67}
68
69#[non_exhaustive]
78#[derive(Debug, Clone, Default)]
79pub struct ChatExtras {
80 pub entropy: Option<f64>,
85}
86
87impl ChatExtras {
88 #[must_use]
101 pub fn with_entropy(entropy: f64) -> Self {
102 Self {
103 entropy: Some(entropy),
104 }
105 }
106}
107
108#[non_exhaustive]
113#[derive(Debug, Clone)]
114pub enum StreamChunk {
115 Content(String),
117 Thinking(String),
119 Compaction(String),
122 ToolUse(Vec<ToolUseRequest>),
124}
125
126pub type ChatStream = Pin<Box<dyn Stream<Item = Result<StreamChunk, LlmError>> + Send>>;
131
132#[derive(Debug, Clone, Serialize, Deserialize)]
138pub struct ToolUseRequest {
139 pub id: String,
141 pub name: ToolName,
143 pub input: serde_json::Value,
145}
146
147#[non_exhaustive]
153#[derive(Debug, Clone)]
154pub enum ThinkingBlock {
155 Thinking { thinking: String, signature: String },
157 Redacted { data: String },
159}
160
161pub const MAX_TOKENS_TRUNCATION_MARKER: &str = "max_tokens limit reached";
164
165#[non_exhaustive]
173#[derive(Debug, Clone)]
174pub enum ChatResponse {
175 Text(String),
177 ToolUse {
179 text: Option<String>,
181 tool_calls: Vec<ToolUseRequest>,
182 thinking_blocks: Vec<ThinkingBlock>,
185 },
186}
187
188pub type EmbedFuture = Pin<Box<dyn Future<Output = Result<Vec<f32>, LlmError>> + Send>>;
190
191pub type EmbedFn = Box<dyn Fn(&str) -> EmbedFuture + Send + Sync>;
196
197pub type StatusTx = tokio::sync::mpsc::UnboundedSender<String>;
203
204#[must_use]
207pub fn default_debug_request_json(
208 messages: &[Message],
209 tools: &[ToolDefinition],
210) -> serde_json::Value {
211 serde_json::json!({
212 "model": serde_json::Value::Null,
213 "max_tokens": serde_json::Value::Null,
214 "messages": serde_json::to_value(messages).unwrap_or(serde_json::Value::Array(vec![])),
215 "tools": serde_json::to_value(tools).unwrap_or(serde_json::Value::Array(vec![])),
216 "temperature": serde_json::Value::Null,
217 "cache_control": serde_json::Value::Null,
218 })
219}
220
221#[derive(Debug, Clone, Default)]
230pub struct GenerationOverrides {
231 pub temperature: Option<f64>,
233 pub top_p: Option<f64>,
235 pub top_k: Option<usize>,
237 pub frequency_penalty: Option<f64>,
239 pub presence_penalty: Option<f64>,
241}
242
243#[non_exhaustive]
250#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
251#[serde(rename_all = "lowercase")]
252pub enum Role {
253 System,
254 User,
255 Assistant,
256}
257
258#[non_exhaustive]
273#[derive(Clone, Debug, Serialize, Deserialize)]
274#[serde(tag = "kind", rename_all = "snake_case")]
275pub enum MessagePart {
276 Text { text: String },
278 ToolOutput {
280 tool_name: zeph_common::ToolName,
281 body: String,
282 #[serde(default, skip_serializing_if = "Option::is_none")]
283 compacted_at: Option<i64>,
284 },
285 Recall { text: String },
287 CodeContext { text: String },
289 Summary { text: String },
291 CrossSession { text: String },
293 ToolUse {
295 id: String,
296 name: String,
297 input: serde_json::Value,
298 },
299 ToolResult {
301 tool_use_id: String,
302 content: String,
303 #[serde(default)]
304 is_error: bool,
305 },
306 Image(Box<ImageData>),
308 ThinkingBlock { thinking: String, signature: String },
310 RedactedThinkingBlock { data: String },
312 Compaction { summary: String },
315}
316
317impl MessagePart {
318 #[must_use]
321 pub fn as_plain_text(&self) -> Option<&str> {
322 match self {
323 Self::Text { text }
324 | Self::Recall { text }
325 | Self::CodeContext { text }
326 | Self::Summary { text }
327 | Self::CrossSession { text } => Some(text.as_str()),
328 _ => None,
329 }
330 }
331
332 #[must_use]
334 pub fn as_image(&self) -> Option<&ImageData> {
335 if let Self::Image(img) = self {
336 Some(img)
337 } else {
338 None
339 }
340 }
341}
342
343#[derive(Clone, Debug, Serialize, Deserialize)]
344pub struct ImageData {
349 #[serde(with = "serde_bytes_base64")]
350 pub data: Vec<u8>,
351 pub mime_type: String,
352}
353
354mod serde_bytes_base64 {
355 use base64::{Engine, engine::general_purpose::STANDARD};
356 use serde::{Deserialize, Deserializer, Serializer};
357
358 pub fn serialize<S>(bytes: &[u8], s: S) -> Result<S::Ok, S::Error>
359 where
360 S: Serializer,
361 {
362 s.serialize_str(&STANDARD.encode(bytes))
363 }
364
365 pub fn deserialize<'de, D>(d: D) -> Result<Vec<u8>, D::Error>
366 where
367 D: Deserializer<'de>,
368 {
369 let s = String::deserialize(d)?;
370 STANDARD.decode(&s).map_err(serde::de::Error::custom)
371 }
372}
373
374#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
390#[serde(rename_all = "snake_case")]
391#[non_exhaustive]
392pub enum MessageVisibility {
393 Both,
395 AgentOnly,
397 UserOnly,
399}
400
401impl MessageVisibility {
402 #[must_use]
404 pub fn is_agent_visible(self) -> bool {
405 matches!(self, MessageVisibility::Both | MessageVisibility::AgentOnly)
406 }
407
408 #[must_use]
410 pub fn is_user_visible(self) -> bool {
411 matches!(self, MessageVisibility::Both | MessageVisibility::UserOnly)
412 }
413}
414
415impl Default for MessageVisibility {
416 fn default() -> Self {
418 MessageVisibility::Both
419 }
420}
421
422impl MessageVisibility {
423 #[must_use]
425 pub fn as_db_str(self) -> &'static str {
426 match self {
427 MessageVisibility::Both => "both",
428 MessageVisibility::AgentOnly => "agent_only",
429 MessageVisibility::UserOnly => "user_only",
430 }
431 }
432
433 #[must_use]
437 pub fn from_db_str(s: &str) -> Self {
438 match s {
439 "agent_only" => MessageVisibility::AgentOnly,
440 "user_only" => MessageVisibility::UserOnly,
441 _ => MessageVisibility::Both,
442 }
443 }
444}
445
446#[derive(Clone, Debug, Serialize, Deserialize)]
451pub struct MessageMetadata {
452 pub visibility: MessageVisibility,
454 #[serde(default, skip_serializing_if = "Option::is_none")]
456 pub compacted_at: Option<i64>,
457 #[serde(default, skip_serializing_if = "Option::is_none")]
460 pub deferred_summary: Option<String>,
461 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
464 pub focus_pinned: bool,
465 #[serde(default, skip_serializing_if = "Option::is_none")]
468 pub focus_marker_id: Option<uuid::Uuid>,
469 #[serde(skip)]
472 pub db_id: Option<i64>,
473 #[serde(default, skip_serializing_if = "Option::is_none")]
478 pub fidelity_tag: Option<zeph_common::ContextFidelity>,
479 #[serde(skip)]
483 pub embedding: Option<Vec<f32>>,
484}
485
486impl Default for MessageMetadata {
487 fn default() -> Self {
488 Self {
489 visibility: MessageVisibility::Both,
490 compacted_at: None,
491 deferred_summary: None,
492 focus_pinned: false,
493 focus_marker_id: None,
494 db_id: None,
495 fidelity_tag: None,
496 embedding: None,
497 }
498 }
499}
500
501impl MessageMetadata {
502 #[must_use]
504 pub fn agent_only() -> Self {
505 Self {
506 visibility: MessageVisibility::AgentOnly,
507 compacted_at: None,
508 deferred_summary: None,
509 focus_pinned: false,
510 focus_marker_id: None,
511 db_id: None,
512 fidelity_tag: None,
513 embedding: None,
514 }
515 }
516
517 #[must_use]
519 pub fn user_only() -> Self {
520 Self {
521 visibility: MessageVisibility::UserOnly,
522 compacted_at: None,
523 deferred_summary: None,
524 focus_pinned: false,
525 focus_marker_id: None,
526 db_id: None,
527 fidelity_tag: None,
528 embedding: None,
529 }
530 }
531
532 #[must_use]
534 pub fn focus_pinned() -> Self {
535 Self {
536 visibility: MessageVisibility::AgentOnly,
537 compacted_at: None,
538 deferred_summary: None,
539 focus_pinned: true,
540 focus_marker_id: None,
541 db_id: None,
542 fidelity_tag: None,
543 embedding: None,
544 }
545 }
546}
547
548#[derive(Clone, Debug, Serialize, Deserialize)]
575pub struct Message {
576 pub role: Role,
577 pub content: String,
579 #[serde(default)]
580 pub parts: Vec<MessagePart>,
581 #[serde(default)]
582 pub metadata: MessageMetadata,
583}
584
585impl Default for Message {
586 fn default() -> Self {
587 Self {
588 role: Role::User,
589 content: String::new(),
590 parts: vec![],
591 metadata: MessageMetadata::default(),
592 }
593 }
594}
595
596impl Message {
597 #[must_use]
602 pub fn from_legacy(role: Role, content: impl Into<String>) -> Self {
603 Self {
604 role,
605 content: content.into(),
606 parts: vec![],
607 metadata: MessageMetadata::default(),
608 }
609 }
610
611 #[must_use]
616 pub fn from_parts(role: Role, parts: Vec<MessagePart>) -> Self {
617 let content = Self::flatten_parts(&parts);
618 Self {
619 role,
620 content,
621 parts,
622 metadata: MessageMetadata::default(),
623 }
624 }
625
626 #[must_use]
629 pub fn to_llm_content(&self) -> &str {
630 &self.content
631 }
632
633 pub fn rebuild_content(&mut self) {
635 if !self.parts.is_empty() {
636 self.content = Self::flatten_parts(&self.parts);
637 }
638 }
639
640 fn flatten_parts(parts: &[MessagePart]) -> String {
641 use std::fmt::Write;
642 let mut out = String::new();
643 for part in parts {
644 match part {
645 MessagePart::Text { text }
646 | MessagePart::Recall { text }
647 | MessagePart::CodeContext { text }
648 | MessagePart::Summary { text }
649 | MessagePart::CrossSession { text } => out.push_str(text),
650 MessagePart::ToolOutput {
651 tool_name,
652 body,
653 compacted_at,
654 } => {
655 if compacted_at.is_some() {
656 if body.is_empty() {
657 let _ = write!(out, "[tool output: {tool_name}] (pruned)");
658 } else {
659 let _ = write!(out, "[tool output: {tool_name}] {body}");
660 }
661 } else {
662 let _ = write!(out, "[tool output: {tool_name}]\n```\n{body}\n```");
663 }
664 }
665 MessagePart::ToolUse { id, name, .. } => {
666 let _ = write!(out, "[tool_use: {name}({id})]");
667 }
668 MessagePart::ToolResult {
669 tool_use_id,
670 content,
671 ..
672 } => {
673 let _ = write!(out, "[tool_result: {tool_use_id}]\n{content}");
674 }
675 MessagePart::Image(img) => {
676 let _ = write!(out, "[image: {}, {} bytes]", img.mime_type, img.data.len());
677 }
678 MessagePart::ThinkingBlock { .. }
680 | MessagePart::RedactedThinkingBlock { .. }
681 | MessagePart::Compaction { .. } => {}
682 }
683 }
684 out
685 }
686}
687
688pub trait LlmProvider: Send + Sync {
758 fn context_window(&self) -> Option<usize> {
762 None
763 }
764
765 fn chat(&self, messages: &[Message]) -> impl Future<Output = Result<String, LlmError>> + Send;
771
772 fn chat_stream(
778 &self,
779 messages: &[Message],
780 ) -> impl Future<Output = Result<ChatStream, LlmError>> + Send;
781
782 fn supports_streaming(&self) -> bool;
784
785 fn embed(&self, text: &str) -> impl Future<Output = Result<Vec<f32>, LlmError>> + Send;
791
792 fn embed_batch(
802 &self,
803 texts: &[&str],
804 ) -> impl Future<Output = Result<Vec<Vec<f32>>, LlmError>> + Send {
805 let owned = owned_strs(texts);
806 async move {
807 let mut results = Vec::with_capacity(owned.len());
808 for text in &owned {
809 results.push(self.embed(text).await?);
810 }
811 Ok(results)
812 }
813 }
814
815 fn supports_embeddings(&self) -> bool;
817
818 fn name(&self) -> &str;
820
821 #[allow(clippy::unnecessary_literal_bound)]
824 fn model_identifier(&self) -> &str {
825 ""
826 }
827
828 fn supports_vision(&self) -> bool {
830 false
831 }
832
833 fn supports_tool_use(&self) -> bool {
840 false
841 }
842
843 fn chat_with_tools(
851 &self,
852 messages: &[Message],
853 _tools: &[ToolDefinition],
854 ) -> impl std::future::Future<Output = Result<ChatResponse, LlmError>> + Send {
855 let msgs = messages.to_vec();
856 async move { Ok(ChatResponse::Text(self.chat(&msgs).await?)) }
857 }
858
859 fn last_cache_usage(&self) -> Option<(u64, u64)> {
862 None
863 }
864
865 fn last_usage(&self) -> Option<(u64, u64)> {
868 None
869 }
870
871 fn last_reasoning_tokens(&self) -> Option<u64> {
876 None
877 }
878
879 fn take_compaction_summary(&self) -> Option<String> {
882 None
883 }
884
885 fn chat_with_extras(
900 &self,
901 messages: &[Message],
902 ) -> impl Future<Output = Result<(String, ChatExtras), LlmError>> + Send {
903 let msgs = messages.to_vec();
904 async move { Ok((self.chat(&msgs).await?, ChatExtras::default())) }
905 }
906
907 #[must_use]
911 fn debug_request_json(
912 &self,
913 messages: &[Message],
914 tools: &[ToolDefinition],
915 _stream: bool,
916 ) -> serde_json::Value {
917 default_debug_request_json(messages, tools)
918 }
919
920 fn list_models(&self) -> Vec<String> {
923 vec![]
924 }
925
926 fn supports_structured_output(&self) -> bool {
928 false
929 }
930
931 #[allow(async_fn_in_trait)]
942 async fn chat_typed<T>(&self, messages: &[Message]) -> Result<T, LlmError>
943 where
944 T: serde::de::DeserializeOwned + schemars::JsonSchema + 'static,
945 Self: Sized,
946 {
947 let (_, schema_json) = cached_schema::<T>()?;
948 let type_name = short_type_name::<T>();
949
950 let mut augmented = messages.to_vec();
951 let instruction = format!(
952 "Respond with a valid JSON object matching this schema. \
953 Output ONLY the JSON, no markdown fences or extra text.\n\n\
954 Type: {type_name}\nSchema:\n```json\n{schema_json}\n```"
955 );
956 augmented.insert(0, Message::from_legacy(Role::System, instruction));
957
958 let raw = self.chat(&augmented).await?;
959 let cleaned = strip_json_fences(&raw);
960 match serde_json::from_str::<T>(cleaned) {
961 Ok(val) => Ok(val),
962 Err(first_err) => {
963 augmented.push(Message::from_legacy(Role::Assistant, &raw));
964 augmented.push(Message::from_legacy(
965 Role::User,
966 format!(
967 "Your response was not valid JSON. Error: {first_err}. \
968 Please output ONLY valid JSON matching the schema."
969 ),
970 ));
971 let retry_raw = self.chat(&augmented).await?;
972 let retry_cleaned = strip_json_fences(&retry_raw);
973 serde_json::from_str::<T>(retry_cleaned).map_err(|e| {
974 LlmError::StructuredParse(format!("parse failed after retry: {e}"))
975 })
976 }
977 }
978 }
979}
980
981fn strip_json_fences(s: &str) -> &str {
985 s.trim()
986 .trim_start_matches("```json")
987 .trim_start_matches("```")
988 .trim_end_matches("```")
989 .trim()
990}
991
992#[cfg(test)]
993mod tests {
994 use std::assert_matches;
995 use tokio_stream::StreamExt;
996
997 use super::*;
998
999 struct StubProvider {
1000 response: String,
1001 }
1002
1003 impl LlmProvider for StubProvider {
1004 async fn chat(&self, _messages: &[Message]) -> Result<String, LlmError> {
1005 Ok(self.response.clone())
1006 }
1007
1008 async fn chat_stream(&self, messages: &[Message]) -> Result<ChatStream, LlmError> {
1009 let response = self.chat(messages).await?;
1010 Ok(Box::pin(tokio_stream::once(Ok(StreamChunk::Content(
1011 response,
1012 )))))
1013 }
1014
1015 fn supports_streaming(&self) -> bool {
1016 false
1017 }
1018
1019 async fn embed(&self, _text: &str) -> Result<Vec<f32>, LlmError> {
1020 Ok(vec![0.1, 0.2, 0.3])
1021 }
1022
1023 fn supports_embeddings(&self) -> bool {
1024 false
1025 }
1026
1027 fn name(&self) -> &'static str {
1028 "stub"
1029 }
1030 }
1031
1032 #[test]
1033 fn context_window_default_returns_none() {
1034 let provider = StubProvider {
1035 response: String::new(),
1036 };
1037 assert!(provider.context_window().is_none());
1038 }
1039
1040 #[test]
1041 fn supports_streaming_default_returns_false() {
1042 let provider = StubProvider {
1043 response: String::new(),
1044 };
1045 assert!(!provider.supports_streaming());
1046 }
1047
1048 #[test]
1049 fn supports_tool_use_default_returns_false() {
1050 let provider = StubProvider {
1056 response: String::new(),
1057 };
1058 assert!(!provider.supports_tool_use());
1059 }
1060
1061 #[tokio::test]
1062 async fn chat_stream_default_yields_single_chunk() {
1063 let provider = StubProvider {
1064 response: "hello world".into(),
1065 };
1066 let messages = vec![Message {
1067 role: Role::User,
1068 content: "test".into(),
1069 parts: vec![],
1070 metadata: MessageMetadata::default(),
1071 }];
1072
1073 let mut stream = provider.chat_stream(&messages).await.unwrap();
1074 let chunk = stream.next().await.unwrap().unwrap();
1075 assert_matches!(chunk, StreamChunk::Content(s) if s == "hello world");
1076 assert!(stream.next().await.is_none());
1077 }
1078
1079 #[tokio::test]
1080 async fn chat_stream_default_propagates_chat_error() {
1081 struct FailProvider;
1082
1083 impl LlmProvider for FailProvider {
1084 async fn chat(&self, _messages: &[Message]) -> Result<String, LlmError> {
1085 Err(LlmError::Unavailable)
1086 }
1087
1088 async fn chat_stream(&self, messages: &[Message]) -> Result<ChatStream, LlmError> {
1089 let response = self.chat(messages).await?;
1090 Ok(Box::pin(tokio_stream::once(Ok(StreamChunk::Content(
1091 response,
1092 )))))
1093 }
1094
1095 fn supports_streaming(&self) -> bool {
1096 false
1097 }
1098
1099 async fn embed(&self, _text: &str) -> Result<Vec<f32>, LlmError> {
1100 Err(LlmError::Unavailable)
1101 }
1102
1103 fn supports_embeddings(&self) -> bool {
1104 false
1105 }
1106
1107 fn name(&self) -> &'static str {
1108 "fail"
1109 }
1110 }
1111
1112 let provider = FailProvider;
1113 let messages = vec![Message {
1114 role: Role::User,
1115 content: "test".into(),
1116 parts: vec![],
1117 metadata: MessageMetadata::default(),
1118 }];
1119
1120 let result = provider.chat_stream(&messages).await;
1121 assert!(result.is_err());
1122 if let Err(e) = result {
1123 assert!(e.to_string().contains("provider unavailable"));
1124 }
1125 }
1126
1127 #[tokio::test]
1128 async fn stub_provider_embed_returns_vector() {
1129 let provider = StubProvider {
1130 response: String::new(),
1131 };
1132 let embedding = provider.embed("test").await.unwrap();
1133 assert_eq!(embedding, vec![0.1, 0.2, 0.3]);
1134 }
1135
1136 #[tokio::test]
1137 async fn fail_provider_embed_propagates_error() {
1138 struct FailProvider;
1139
1140 impl LlmProvider for FailProvider {
1141 async fn chat(&self, _messages: &[Message]) -> Result<String, LlmError> {
1142 Err(LlmError::Unavailable)
1143 }
1144
1145 async fn chat_stream(&self, messages: &[Message]) -> Result<ChatStream, LlmError> {
1146 let response = self.chat(messages).await?;
1147 Ok(Box::pin(tokio_stream::once(Ok(StreamChunk::Content(
1148 response,
1149 )))))
1150 }
1151
1152 fn supports_streaming(&self) -> bool {
1153 false
1154 }
1155
1156 async fn embed(&self, _text: &str) -> Result<Vec<f32>, LlmError> {
1157 Err(LlmError::EmbedUnsupported {
1158 provider: "fail".into(),
1159 })
1160 }
1161
1162 fn supports_embeddings(&self) -> bool {
1163 false
1164 }
1165
1166 fn name(&self) -> &'static str {
1167 "fail"
1168 }
1169 }
1170
1171 let provider = FailProvider;
1172 let result = provider.embed("test").await;
1173 assert!(result.is_err());
1174 assert!(
1175 result
1176 .unwrap_err()
1177 .to_string()
1178 .contains("embedding not supported")
1179 );
1180 }
1181
1182 #[test]
1183 fn role_serialization() {
1184 let system = Role::System;
1185 let user = Role::User;
1186 let assistant = Role::Assistant;
1187
1188 assert_eq!(serde_json::to_string(&system).unwrap(), "\"system\"");
1189 assert_eq!(serde_json::to_string(&user).unwrap(), "\"user\"");
1190 assert_eq!(serde_json::to_string(&assistant).unwrap(), "\"assistant\"");
1191 }
1192
1193 #[test]
1194 fn role_deserialization() {
1195 let system: Role = serde_json::from_str("\"system\"").unwrap();
1196 let user: Role = serde_json::from_str("\"user\"").unwrap();
1197 let assistant: Role = serde_json::from_str("\"assistant\"").unwrap();
1198
1199 assert_eq!(system, Role::System);
1200 assert_eq!(user, Role::User);
1201 assert_eq!(assistant, Role::Assistant);
1202 }
1203
1204 #[test]
1205 fn message_clone() {
1206 let msg = Message {
1207 role: Role::User,
1208 content: "test".into(),
1209 parts: vec![],
1210 metadata: MessageMetadata::default(),
1211 };
1212 let cloned = msg.clone();
1213 assert_eq!(cloned.role, msg.role);
1214 assert_eq!(cloned.content, msg.content);
1215 }
1216
1217 #[test]
1218 fn message_debug() {
1219 let msg = Message {
1220 role: Role::Assistant,
1221 content: "response".into(),
1222 parts: vec![],
1223 metadata: MessageMetadata::default(),
1224 };
1225 let debug = format!("{msg:?}");
1226 assert!(debug.contains("Assistant"));
1227 assert!(debug.contains("response"));
1228 }
1229
1230 #[test]
1231 fn message_serialization() {
1232 let msg = Message {
1233 role: Role::User,
1234 content: "hello".into(),
1235 parts: vec![],
1236 metadata: MessageMetadata::default(),
1237 };
1238 let json = serde_json::to_string(&msg).unwrap();
1239 assert!(json.contains("\"role\":\"user\""));
1240 assert!(json.contains("\"content\":\"hello\""));
1241 }
1242
1243 #[test]
1244 fn message_part_serde_round_trip() {
1245 let parts = vec![
1246 MessagePart::Text {
1247 text: "hello".into(),
1248 },
1249 MessagePart::ToolOutput {
1250 tool_name: "bash".into(),
1251 body: "output".into(),
1252 compacted_at: None,
1253 },
1254 MessagePart::Recall {
1255 text: "recall".into(),
1256 },
1257 MessagePart::CodeContext {
1258 text: "code".into(),
1259 },
1260 MessagePart::Summary {
1261 text: "summary".into(),
1262 },
1263 ];
1264 let json = serde_json::to_string(&parts).unwrap();
1265 let deserialized: Vec<MessagePart> = serde_json::from_str(&json).unwrap();
1266 assert_eq!(deserialized.len(), 5);
1267 }
1268
1269 #[test]
1270 fn from_legacy_creates_empty_parts() {
1271 let msg = Message::from_legacy(Role::User, "hello");
1272 assert_eq!(msg.role, Role::User);
1273 assert_eq!(msg.content, "hello");
1274 assert!(msg.parts.is_empty());
1275 assert_eq!(msg.to_llm_content(), "hello");
1276 }
1277
1278 #[test]
1279 fn from_parts_flattens_content() {
1280 let msg = Message::from_parts(
1281 Role::System,
1282 vec![MessagePart::Recall {
1283 text: "recalled data".into(),
1284 }],
1285 );
1286 assert_eq!(msg.content, "recalled data");
1287 assert_eq!(msg.to_llm_content(), "recalled data");
1288 assert_eq!(msg.parts.len(), 1);
1289 }
1290
1291 #[test]
1292 fn from_parts_tool_output_format() {
1293 let msg = Message::from_parts(
1294 Role::User,
1295 vec![MessagePart::ToolOutput {
1296 tool_name: "bash".into(),
1297 body: "hello world".into(),
1298 compacted_at: None,
1299 }],
1300 );
1301 assert!(msg.content.contains("[tool output: bash]"));
1302 assert!(msg.content.contains("hello world"));
1303 }
1304
1305 #[test]
1306 fn message_deserializes_without_parts() {
1307 let json = r#"{"role":"user","content":"hello"}"#;
1308 let msg: Message = serde_json::from_str(json).unwrap();
1309 assert_eq!(msg.content, "hello");
1310 assert!(msg.parts.is_empty());
1311 }
1312
1313 #[test]
1314 fn flatten_skips_compacted_tool_output_empty_body() {
1315 let msg = Message::from_parts(
1317 Role::User,
1318 vec![
1319 MessagePart::Text {
1320 text: "prefix ".into(),
1321 },
1322 MessagePart::ToolOutput {
1323 tool_name: "bash".into(),
1324 body: String::new(),
1325 compacted_at: Some(1234),
1326 },
1327 MessagePart::Text {
1328 text: " suffix".into(),
1329 },
1330 ],
1331 );
1332 assert!(msg.content.contains("(pruned)"));
1333 assert!(msg.content.contains("prefix "));
1334 assert!(msg.content.contains(" suffix"));
1335 }
1336
1337 #[test]
1338 fn flatten_compacted_tool_output_with_reference_renders_body() {
1339 let ref_notice = "[tool output pruned; full content at /tmp/overflow/big.txt]";
1341 let msg = Message::from_parts(
1342 Role::User,
1343 vec![MessagePart::ToolOutput {
1344 tool_name: "bash".into(),
1345 body: ref_notice.into(),
1346 compacted_at: Some(1234),
1347 }],
1348 );
1349 assert!(msg.content.contains(ref_notice));
1350 assert!(!msg.content.contains("(pruned)"));
1351 }
1352
1353 #[test]
1354 fn rebuild_content_syncs_after_mutation() {
1355 let mut msg = Message::from_parts(
1356 Role::User,
1357 vec![MessagePart::ToolOutput {
1358 tool_name: "bash".into(),
1359 body: "original".into(),
1360 compacted_at: None,
1361 }],
1362 );
1363 assert!(msg.content.contains("original"));
1364
1365 if let MessagePart::ToolOutput {
1366 ref mut compacted_at,
1367 ref mut body,
1368 ..
1369 } = msg.parts[0]
1370 {
1371 *compacted_at = Some(999);
1372 body.clear(); }
1374 msg.rebuild_content();
1375
1376 assert!(msg.content.contains("(pruned)"));
1377 assert!(!msg.content.contains("original"));
1378 }
1379
1380 #[test]
1381 fn message_part_tool_use_serde_round_trip() {
1382 let part = MessagePart::ToolUse {
1383 id: "toolu_123".into(),
1384 name: "bash".into(),
1385 input: serde_json::json!({"command": "ls"}),
1386 };
1387 let json = serde_json::to_string(&part).unwrap();
1388 let deserialized: MessagePart = serde_json::from_str(&json).unwrap();
1389 if let MessagePart::ToolUse { id, name, input } = deserialized {
1390 assert_eq!(id, "toolu_123");
1391 assert_eq!(name, "bash");
1392 assert_eq!(input["command"], "ls");
1393 } else {
1394 panic!("expected ToolUse");
1395 }
1396 }
1397
1398 #[test]
1399 fn message_part_tool_result_serde_round_trip() {
1400 let part = MessagePart::ToolResult {
1401 tool_use_id: "toolu_123".into(),
1402 content: "file1.rs\nfile2.rs".into(),
1403 is_error: false,
1404 };
1405 let json = serde_json::to_string(&part).unwrap();
1406 let deserialized: MessagePart = serde_json::from_str(&json).unwrap();
1407 if let MessagePart::ToolResult {
1408 tool_use_id,
1409 content,
1410 is_error,
1411 } = deserialized
1412 {
1413 assert_eq!(tool_use_id, "toolu_123");
1414 assert_eq!(content, "file1.rs\nfile2.rs");
1415 assert!(!is_error);
1416 } else {
1417 panic!("expected ToolResult");
1418 }
1419 }
1420
1421 #[test]
1422 fn message_part_tool_result_is_error_default() {
1423 let json = r#"{"kind":"tool_result","tool_use_id":"id","content":"err"}"#;
1424 let part: MessagePart = serde_json::from_str(json).unwrap();
1425 if let MessagePart::ToolResult { is_error, .. } = part {
1426 assert!(!is_error);
1427 } else {
1428 panic!("expected ToolResult");
1429 }
1430 }
1431
1432 #[test]
1433 fn chat_response_construction() {
1434 let text = ChatResponse::Text("hello".into());
1435 assert_matches!(text, ChatResponse::Text(s) if s == "hello");
1436
1437 let tool_use = ChatResponse::ToolUse {
1438 text: Some("I'll run that".into()),
1439 tool_calls: vec![ToolUseRequest {
1440 id: "1".into(),
1441 name: "bash".into(),
1442 input: serde_json::json!({}),
1443 }],
1444 thinking_blocks: vec![],
1445 };
1446 assert_matches!(tool_use, ChatResponse::ToolUse { .. });
1447 }
1448
1449 #[test]
1450 fn flatten_parts_tool_use() {
1451 let msg = Message::from_parts(
1452 Role::Assistant,
1453 vec![MessagePart::ToolUse {
1454 id: "t1".into(),
1455 name: "bash".into(),
1456 input: serde_json::json!({"command": "ls"}),
1457 }],
1458 );
1459 assert!(msg.content.contains("[tool_use: bash(t1)]"));
1460 }
1461
1462 #[test]
1463 fn flatten_parts_tool_result() {
1464 let msg = Message::from_parts(
1465 Role::User,
1466 vec![MessagePart::ToolResult {
1467 tool_use_id: "t1".into(),
1468 content: "output here".into(),
1469 is_error: false,
1470 }],
1471 );
1472 assert!(msg.content.contains("[tool_result: t1]"));
1473 assert!(msg.content.contains("output here"));
1474 }
1475
1476 #[test]
1477 fn tool_definition_serde_round_trip() {
1478 let def = ToolDefinition {
1479 name: "bash".into(),
1480 description: "Execute a shell command".into(),
1481 parameters: serde_json::json!({"type": "object"}),
1482 output_schema: None,
1483 };
1484 let json = serde_json::to_string(&def).unwrap();
1485 let deserialized: ToolDefinition = serde_json::from_str(&json).unwrap();
1486 assert_eq!(deserialized.name, "bash");
1487 assert_eq!(deserialized.description, "Execute a shell command");
1488 }
1489
1490 #[tokio::test]
1491 async fn chat_with_tools_default_delegates_to_chat() {
1492 let provider = StubProvider {
1493 response: "hello".into(),
1494 };
1495 let messages = vec![Message::from_legacy(Role::User, "test")];
1496 let result = provider.chat_with_tools(&messages, &[]).await.unwrap();
1497 assert_matches!(result, ChatResponse::Text(s) if s == "hello");
1498 }
1499
1500 #[test]
1501 fn tool_output_compacted_at_serde_default() {
1502 let json = r#"{"kind":"tool_output","tool_name":"bash","body":"out"}"#;
1503 let part: MessagePart = serde_json::from_str(json).unwrap();
1504 if let MessagePart::ToolOutput { compacted_at, .. } = part {
1505 assert!(compacted_at.is_none());
1506 } else {
1507 panic!("expected ToolOutput");
1508 }
1509 }
1510
1511 #[test]
1514 fn strip_json_fences_plain_json() {
1515 assert_eq!(strip_json_fences(r#"{"a": 1}"#), r#"{"a": 1}"#);
1516 }
1517
1518 #[test]
1519 fn strip_json_fences_with_json_fence() {
1520 assert_eq!(strip_json_fences("```json\n{\"a\": 1}\n```"), r#"{"a": 1}"#);
1521 }
1522
1523 #[test]
1524 fn strip_json_fences_with_plain_fence() {
1525 assert_eq!(strip_json_fences("```\n{\"a\": 1}\n```"), r#"{"a": 1}"#);
1526 }
1527
1528 #[test]
1529 fn strip_json_fences_whitespace() {
1530 assert_eq!(strip_json_fences(" \n "), "");
1531 }
1532
1533 #[test]
1534 fn strip_json_fences_empty() {
1535 assert_eq!(strip_json_fences(""), "");
1536 }
1537
1538 #[test]
1539 fn strip_json_fences_outer_whitespace() {
1540 assert_eq!(
1541 strip_json_fences(" ```json\n{\"a\": 1}\n``` "),
1542 r#"{"a": 1}"#
1543 );
1544 }
1545
1546 #[test]
1547 fn strip_json_fences_only_opening_fence() {
1548 assert_eq!(strip_json_fences("```json\n{\"a\": 1}"), r#"{"a": 1}"#);
1549 }
1550
1551 #[derive(Debug, serde::Deserialize, schemars::JsonSchema, PartialEq)]
1554 struct TestOutput {
1555 value: String,
1556 }
1557
1558 struct SequentialStub {
1559 responses: std::sync::Mutex<Vec<Result<String, LlmError>>>,
1560 }
1561
1562 impl SequentialStub {
1563 fn new(responses: Vec<Result<String, LlmError>>) -> Self {
1564 Self {
1565 responses: std::sync::Mutex::new(responses),
1566 }
1567 }
1568 }
1569
1570 impl LlmProvider for SequentialStub {
1571 async fn chat(&self, _messages: &[Message]) -> Result<String, LlmError> {
1572 let mut responses = self.responses.lock().unwrap();
1573 if responses.is_empty() {
1574 return Err(LlmError::Other("no more responses".into()));
1575 }
1576 responses.remove(0)
1577 }
1578
1579 async fn chat_stream(&self, messages: &[Message]) -> Result<ChatStream, LlmError> {
1580 let response = self.chat(messages).await?;
1581 Ok(Box::pin(tokio_stream::once(Ok(StreamChunk::Content(
1582 response,
1583 )))))
1584 }
1585
1586 fn supports_streaming(&self) -> bool {
1587 false
1588 }
1589
1590 async fn embed(&self, _text: &str) -> Result<Vec<f32>, LlmError> {
1591 Err(LlmError::EmbedUnsupported {
1592 provider: "sequential-stub".into(),
1593 })
1594 }
1595
1596 fn supports_embeddings(&self) -> bool {
1597 false
1598 }
1599
1600 fn name(&self) -> &'static str {
1601 "sequential-stub"
1602 }
1603 }
1604
1605 #[tokio::test]
1606 async fn chat_typed_happy_path() {
1607 let provider = StubProvider {
1608 response: r#"{"value": "hello"}"#.into(),
1609 };
1610 let messages = vec![Message::from_legacy(Role::User, "test")];
1611 let result: TestOutput = provider.chat_typed(&messages).await.unwrap();
1612 assert_eq!(
1613 result,
1614 TestOutput {
1615 value: "hello".into()
1616 }
1617 );
1618 }
1619
1620 #[tokio::test]
1621 async fn chat_typed_retry_succeeds() {
1622 let provider = SequentialStub::new(vec![
1623 Ok("not valid json".into()),
1624 Ok(r#"{"value": "ok"}"#.into()),
1625 ]);
1626 let messages = vec![Message::from_legacy(Role::User, "test")];
1627 let result: TestOutput = provider.chat_typed(&messages).await.unwrap();
1628 assert_eq!(result, TestOutput { value: "ok".into() });
1629 }
1630
1631 #[tokio::test]
1632 async fn chat_typed_both_fail() {
1633 let provider = SequentialStub::new(vec![Ok("bad json".into()), Ok("still bad".into())]);
1634 let messages = vec![Message::from_legacy(Role::User, "test")];
1635 let result = provider.chat_typed::<TestOutput>(&messages).await;
1636 let err = result.unwrap_err();
1637 assert!(err.to_string().contains("parse failed after retry"));
1638 }
1639
1640 #[tokio::test]
1641 async fn chat_typed_chat_error_propagates() {
1642 let provider = SequentialStub::new(vec![Err(LlmError::Unavailable)]);
1643 let messages = vec![Message::from_legacy(Role::User, "test")];
1644 let result = provider.chat_typed::<TestOutput>(&messages).await;
1645 assert_matches!(result, Err(LlmError::Unavailable));
1646 }
1647
1648 #[tokio::test]
1649 async fn chat_typed_strips_fences() {
1650 let provider = StubProvider {
1651 response: "```json\n{\"value\": \"fenced\"}\n```".into(),
1652 };
1653 let messages = vec![Message::from_legacy(Role::User, "test")];
1654 let result: TestOutput = provider.chat_typed(&messages).await.unwrap();
1655 assert_eq!(
1656 result,
1657 TestOutput {
1658 value: "fenced".into()
1659 }
1660 );
1661 }
1662
1663 #[test]
1664 fn supports_structured_output_default_false() {
1665 let provider = StubProvider {
1666 response: String::new(),
1667 };
1668 assert!(!provider.supports_structured_output());
1669 }
1670
1671 #[test]
1672 fn structured_parse_error_display() {
1673 let err = LlmError::StructuredParse("test error".into());
1674 assert_eq!(
1675 err.to_string(),
1676 "structured output parse failed: test error"
1677 );
1678 }
1679
1680 #[test]
1681 fn message_part_image_roundtrip_json() {
1682 let part = MessagePart::Image(Box::new(ImageData {
1683 data: vec![1, 2, 3, 4],
1684 mime_type: "image/jpeg".into(),
1685 }));
1686 let json = serde_json::to_string(&part).unwrap();
1687 let decoded: MessagePart = serde_json::from_str(&json).unwrap();
1688 match decoded {
1689 MessagePart::Image(img) => {
1690 assert_eq!(img.data, vec![1, 2, 3, 4]);
1691 assert_eq!(img.mime_type, "image/jpeg");
1692 }
1693 _ => panic!("expected Image variant"),
1694 }
1695 }
1696
1697 #[test]
1698 fn flatten_parts_includes_image_placeholder() {
1699 let msg = Message::from_parts(
1700 Role::User,
1701 vec![
1702 MessagePart::Text {
1703 text: "see this".into(),
1704 },
1705 MessagePart::Image(Box::new(ImageData {
1706 data: vec![0u8; 100],
1707 mime_type: "image/png".into(),
1708 })),
1709 ],
1710 );
1711 let content = msg.to_llm_content();
1712 assert!(content.contains("see this"));
1713 assert!(content.contains("[image: image/png"));
1714 }
1715
1716 #[test]
1717 fn supports_vision_default_false() {
1718 let provider = StubProvider {
1719 response: String::new(),
1720 };
1721 assert!(!provider.supports_vision());
1722 }
1723
1724 #[test]
1725 fn message_metadata_default_both_visible() {
1726 let m = MessageMetadata::default();
1727 assert!(m.visibility.is_agent_visible());
1728 assert!(m.visibility.is_user_visible());
1729 assert_eq!(m.visibility, MessageVisibility::Both);
1730 assert!(m.compacted_at.is_none());
1731 }
1732
1733 #[test]
1734 fn message_metadata_agent_only() {
1735 let m = MessageMetadata::agent_only();
1736 assert!(m.visibility.is_agent_visible());
1737 assert!(!m.visibility.is_user_visible());
1738 assert_eq!(m.visibility, MessageVisibility::AgentOnly);
1739 }
1740
1741 #[test]
1742 fn message_metadata_user_only() {
1743 let m = MessageMetadata::user_only();
1744 assert!(!m.visibility.is_agent_visible());
1745 assert!(m.visibility.is_user_visible());
1746 assert_eq!(m.visibility, MessageVisibility::UserOnly);
1747 }
1748
1749 #[test]
1750 fn message_metadata_serde_default() {
1751 let json = r#"{"role":"user","content":"hello"}"#;
1752 let msg: Message = serde_json::from_str(json).unwrap();
1753 assert!(msg.metadata.visibility.is_agent_visible());
1754 assert!(msg.metadata.visibility.is_user_visible());
1755 }
1756
1757 #[test]
1758 fn message_metadata_round_trip() {
1759 let msg = Message {
1760 role: Role::User,
1761 content: "test".into(),
1762 parts: vec![],
1763 metadata: MessageMetadata::agent_only(),
1764 };
1765 let json = serde_json::to_string(&msg).unwrap();
1766 let decoded: Message = serde_json::from_str(&json).unwrap();
1767 assert!(decoded.metadata.visibility.is_agent_visible());
1768 assert!(!decoded.metadata.visibility.is_user_visible());
1769 assert_eq!(decoded.metadata.visibility, MessageVisibility::AgentOnly);
1770 }
1771
1772 #[test]
1773 fn message_part_compaction_round_trip() {
1774 let part = MessagePart::Compaction {
1775 summary: "Context was summarized.".to_owned(),
1776 };
1777 let json = serde_json::to_string(&part).unwrap();
1778 let decoded: MessagePart = serde_json::from_str(&json).unwrap();
1779 assert!(
1780 matches!(decoded, MessagePart::Compaction { summary } if summary == "Context was summarized.")
1781 );
1782 }
1783
1784 #[test]
1785 fn flatten_parts_compaction_contributes_no_text() {
1786 let parts = vec![
1789 MessagePart::Text {
1790 text: "Hello".to_owned(),
1791 },
1792 MessagePart::Compaction {
1793 summary: "Summary".to_owned(),
1794 },
1795 ];
1796 let msg = Message::from_parts(Role::Assistant, parts);
1797 assert_eq!(msg.content.trim(), "Hello");
1799 }
1800
1801 #[test]
1802 fn stream_chunk_compaction_variant() {
1803 let chunk = StreamChunk::Compaction("A summary".to_owned());
1804 assert_matches!(chunk, StreamChunk::Compaction(s) if s == "A summary");
1805 }
1806
1807 #[test]
1808 fn short_type_name_extracts_last_segment() {
1809 struct MyOutput;
1810 assert_eq!(short_type_name::<MyOutput>(), "MyOutput");
1811 }
1812
1813 #[test]
1814 fn short_type_name_primitive_returns_full_name() {
1815 assert_eq!(short_type_name::<u32>(), "u32");
1817 assert_eq!(short_type_name::<bool>(), "bool");
1818 }
1819
1820 #[test]
1821 fn short_type_name_nested_path_returns_last() {
1822 assert_eq!(
1824 short_type_name::<std::collections::HashMap<u32, u32>>(),
1825 "HashMap<u32, u32>"
1826 );
1827 }
1828
1829 #[test]
1832 fn summary_roundtrip() {
1833 let part = MessagePart::Summary {
1834 text: "hello".to_string(),
1835 };
1836 let json = serde_json::to_string(&part).expect("serialization must not fail");
1837 assert!(
1838 json.contains("\"kind\":\"summary\""),
1839 "must use internally-tagged format, got: {json}"
1840 );
1841 assert!(
1842 !json.contains("\"Summary\""),
1843 "must not use externally-tagged format, got: {json}"
1844 );
1845 let decoded: MessagePart =
1846 serde_json::from_str(&json).expect("deserialization must not fail");
1847 match decoded {
1848 MessagePart::Summary { text } => assert_eq!(text, "hello"),
1849 other => panic!("expected MessagePart::Summary, got {other:?}"),
1850 }
1851 }
1852
1853 #[tokio::test]
1854 async fn embed_batch_default_empty_returns_empty() {
1855 let provider = StubProvider {
1856 response: String::new(),
1857 };
1858 let result = provider.embed_batch(&[]).await.unwrap();
1859 assert!(result.is_empty());
1860 }
1861
1862 #[tokio::test]
1863 async fn embed_batch_default_calls_embed_sequentially() {
1864 let provider = StubProvider {
1865 response: String::new(),
1866 };
1867 let texts = ["hello", "world", "foo"];
1868 let result = provider.embed_batch(&texts).await.unwrap();
1869 assert_eq!(result.len(), 3);
1870 for vec in &result {
1872 assert_eq!(vec, &[0.1_f32, 0.2, 0.3]);
1873 }
1874 }
1875
1876 #[test]
1877 fn message_visibility_db_roundtrip_both() {
1878 assert_eq!(MessageVisibility::Both.as_db_str(), "both");
1879 assert_eq!(
1880 MessageVisibility::from_db_str("both"),
1881 MessageVisibility::Both
1882 );
1883 }
1884
1885 #[test]
1886 fn message_visibility_db_roundtrip_agent_only() {
1887 assert_eq!(MessageVisibility::AgentOnly.as_db_str(), "agent_only");
1888 assert_eq!(
1889 MessageVisibility::from_db_str("agent_only"),
1890 MessageVisibility::AgentOnly
1891 );
1892 }
1893
1894 #[test]
1895 fn message_visibility_db_roundtrip_user_only() {
1896 assert_eq!(MessageVisibility::UserOnly.as_db_str(), "user_only");
1897 assert_eq!(
1898 MessageVisibility::from_db_str("user_only"),
1899 MessageVisibility::UserOnly
1900 );
1901 }
1902
1903 #[test]
1904 fn message_visibility_from_db_str_unknown_defaults_to_both() {
1905 assert_eq!(
1906 MessageVisibility::from_db_str("unknown_future_value"),
1907 MessageVisibility::Both
1908 );
1909 assert_eq!(MessageVisibility::from_db_str(""), MessageVisibility::Both);
1910 }
1911}