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, Serialize, Deserialize)]
344pub struct ImageData {
349 #[serde(with = "serde_bytes_base64")]
350 pub data: Vec<u8>,
351 pub mime_type: String,
352}
353
354impl std::fmt::Debug for ImageData {
355 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
360 write!(f, "[image: {}, {} bytes]", self.mime_type, self.data.len())
361 }
362}
363
364mod serde_bytes_base64 {
365 use base64::{Engine, engine::general_purpose::STANDARD};
366 use serde::{Deserialize, Deserializer, Serializer};
367
368 pub fn serialize<S>(bytes: &[u8], s: S) -> Result<S::Ok, S::Error>
369 where
370 S: Serializer,
371 {
372 s.serialize_str(&STANDARD.encode(bytes))
373 }
374
375 pub fn deserialize<'de, D>(d: D) -> Result<Vec<u8>, D::Error>
376 where
377 D: Deserializer<'de>,
378 {
379 let s = String::deserialize(d)?;
380 STANDARD.decode(&s).map_err(serde::de::Error::custom)
381 }
382}
383
384#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
400#[serde(rename_all = "snake_case")]
401#[non_exhaustive]
402pub enum MessageVisibility {
403 Both,
405 AgentOnly,
407 UserOnly,
409}
410
411impl MessageVisibility {
412 #[must_use]
414 pub fn is_agent_visible(self) -> bool {
415 matches!(self, MessageVisibility::Both | MessageVisibility::AgentOnly)
416 }
417
418 #[must_use]
420 pub fn is_user_visible(self) -> bool {
421 matches!(self, MessageVisibility::Both | MessageVisibility::UserOnly)
422 }
423}
424
425impl Default for MessageVisibility {
426 fn default() -> Self {
428 MessageVisibility::Both
429 }
430}
431
432impl MessageVisibility {
433 #[must_use]
435 pub fn as_db_str(self) -> &'static str {
436 match self {
437 MessageVisibility::Both => "both",
438 MessageVisibility::AgentOnly => "agent_only",
439 MessageVisibility::UserOnly => "user_only",
440 }
441 }
442
443 #[must_use]
447 pub fn from_db_str(s: &str) -> Self {
448 match s {
449 "agent_only" => MessageVisibility::AgentOnly,
450 "user_only" => MessageVisibility::UserOnly,
451 _ => MessageVisibility::Both,
452 }
453 }
454}
455
456#[derive(Clone, Debug, Serialize, Deserialize)]
461pub struct MessageMetadata {
462 pub visibility: MessageVisibility,
464 #[serde(default, skip_serializing_if = "Option::is_none")]
466 pub compacted_at: Option<i64>,
467 #[serde(default, skip_serializing_if = "Option::is_none")]
470 pub deferred_summary: Option<String>,
471 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
474 pub focus_pinned: bool,
475 #[serde(default, skip_serializing_if = "Option::is_none")]
478 pub focus_marker_id: Option<uuid::Uuid>,
479 #[serde(skip)]
482 pub db_id: Option<i64>,
483 #[serde(default, skip_serializing_if = "Option::is_none")]
488 pub fidelity_tag: Option<zeph_common::ContextFidelity>,
489 #[serde(skip)]
493 pub embedding: Option<Vec<f32>>,
494}
495
496impl Default for MessageMetadata {
497 fn default() -> Self {
498 Self {
499 visibility: MessageVisibility::Both,
500 compacted_at: None,
501 deferred_summary: None,
502 focus_pinned: false,
503 focus_marker_id: None,
504 db_id: None,
505 fidelity_tag: None,
506 embedding: None,
507 }
508 }
509}
510
511impl MessageMetadata {
512 #[must_use]
514 pub fn agent_only() -> Self {
515 Self {
516 visibility: MessageVisibility::AgentOnly,
517 compacted_at: None,
518 deferred_summary: None,
519 focus_pinned: false,
520 focus_marker_id: None,
521 db_id: None,
522 fidelity_tag: None,
523 embedding: None,
524 }
525 }
526
527 #[must_use]
529 pub fn user_only() -> Self {
530 Self {
531 visibility: MessageVisibility::UserOnly,
532 compacted_at: None,
533 deferred_summary: None,
534 focus_pinned: false,
535 focus_marker_id: None,
536 db_id: None,
537 fidelity_tag: None,
538 embedding: None,
539 }
540 }
541
542 #[must_use]
544 pub fn focus_pinned() -> Self {
545 Self {
546 visibility: MessageVisibility::AgentOnly,
547 compacted_at: None,
548 deferred_summary: None,
549 focus_pinned: true,
550 focus_marker_id: None,
551 db_id: None,
552 fidelity_tag: None,
553 embedding: None,
554 }
555 }
556}
557
558#[derive(Clone, Debug, Serialize, Deserialize)]
585pub struct Message {
586 pub role: Role,
587 pub content: String,
589 #[serde(default)]
590 pub parts: Vec<MessagePart>,
591 #[serde(default)]
592 pub metadata: MessageMetadata,
593}
594
595impl Default for Message {
596 fn default() -> Self {
597 Self {
598 role: Role::User,
599 content: String::new(),
600 parts: vec![],
601 metadata: MessageMetadata::default(),
602 }
603 }
604}
605
606impl Message {
607 #[must_use]
612 pub fn from_legacy(role: Role, content: impl Into<String>) -> Self {
613 Self {
614 role,
615 content: content.into(),
616 parts: vec![],
617 metadata: MessageMetadata::default(),
618 }
619 }
620
621 #[must_use]
626 pub fn from_parts(role: Role, parts: Vec<MessagePart>) -> Self {
627 let content = Self::flatten_parts(&parts);
628 Self {
629 role,
630 content,
631 parts,
632 metadata: MessageMetadata::default(),
633 }
634 }
635
636 #[must_use]
639 pub fn to_llm_content(&self) -> &str {
640 &self.content
641 }
642
643 pub fn rebuild_content(&mut self) {
645 if !self.parts.is_empty() {
646 self.content = Self::flatten_parts(&self.parts);
647 }
648 }
649
650 fn flatten_parts(parts: &[MessagePart]) -> String {
651 use std::fmt::Write;
652 let mut out = String::new();
653 for part in parts {
654 match part {
655 MessagePart::Text { text }
656 | MessagePart::Recall { text }
657 | MessagePart::CodeContext { text }
658 | MessagePart::Summary { text }
659 | MessagePart::CrossSession { text } => out.push_str(text),
660 MessagePart::ToolOutput {
661 tool_name,
662 body,
663 compacted_at,
664 } => {
665 if compacted_at.is_some() {
666 if body.is_empty() {
667 let _ = write!(out, "[tool output: {tool_name}] (pruned)");
668 } else {
669 let _ = write!(out, "[tool output: {tool_name}] {body}");
670 }
671 } else {
672 let _ = write!(out, "[tool output: {tool_name}]\n```\n{body}\n```");
673 }
674 }
675 MessagePart::ToolUse { id, name, .. } => {
676 let _ = write!(out, "[tool_use: {name}({id})]");
677 }
678 MessagePart::ToolResult {
679 tool_use_id,
680 content,
681 ..
682 } => {
683 let _ = write!(out, "[tool_result: {tool_use_id}]\n{content}");
684 }
685 MessagePart::Image(img) => {
686 let _ = write!(out, "[image: {}, {} bytes]", img.mime_type, img.data.len());
687 }
688 MessagePart::ThinkingBlock { .. }
690 | MessagePart::RedactedThinkingBlock { .. }
691 | MessagePart::Compaction { .. } => {}
692 }
693 }
694 out
695 }
696}
697
698pub trait LlmProvider: Send + Sync {
768 fn context_window(&self) -> Option<usize> {
772 None
773 }
774
775 fn chat(&self, messages: &[Message]) -> impl Future<Output = Result<String, LlmError>> + Send;
781
782 fn chat_stream(
788 &self,
789 messages: &[Message],
790 ) -> impl Future<Output = Result<ChatStream, LlmError>> + Send;
791
792 fn supports_streaming(&self) -> bool;
794
795 fn embed(&self, text: &str) -> impl Future<Output = Result<Vec<f32>, LlmError>> + Send;
801
802 fn embed_batch(
812 &self,
813 texts: &[&str],
814 ) -> impl Future<Output = Result<Vec<Vec<f32>>, LlmError>> + Send {
815 let owned = owned_strs(texts);
816 async move {
817 let mut results = Vec::with_capacity(owned.len());
818 for text in &owned {
819 results.push(self.embed(text).await?);
820 }
821 Ok(results)
822 }
823 }
824
825 fn supports_embeddings(&self) -> bool;
827
828 fn name(&self) -> &str;
830
831 #[allow(clippy::unnecessary_literal_bound)]
834 fn model_identifier(&self) -> &str {
835 ""
836 }
837
838 fn effective_model_identifier(&self) -> &str {
850 self.model_identifier()
851 }
852
853 fn supports_vision(&self) -> bool {
855 false
856 }
857
858 fn supports_tool_use(&self) -> bool {
865 false
866 }
867
868 fn chat_with_tools(
876 &self,
877 messages: &[Message],
878 _tools: &[ToolDefinition],
879 ) -> impl std::future::Future<Output = Result<ChatResponse, LlmError>> + Send {
880 let msgs = messages.to_vec();
881 async move { Ok(ChatResponse::Text(self.chat(&msgs).await?)) }
882 }
883
884 fn last_cache_usage(&self) -> Option<(u64, u64)> {
887 None
888 }
889
890 fn last_usage(&self) -> Option<(u64, u64)> {
893 None
894 }
895
896 fn last_reasoning_tokens(&self) -> Option<u64> {
901 None
902 }
903
904 fn take_compaction_summary(&self) -> Option<String> {
907 None
908 }
909
910 fn chat_with_extras(
925 &self,
926 messages: &[Message],
927 ) -> impl Future<Output = Result<(String, ChatExtras), LlmError>> + Send {
928 let msgs = messages.to_vec();
929 async move { Ok((self.chat(&msgs).await?, ChatExtras::default())) }
930 }
931
932 #[must_use]
936 fn debug_request_json(
937 &self,
938 messages: &[Message],
939 tools: &[ToolDefinition],
940 _stream: bool,
941 ) -> serde_json::Value {
942 default_debug_request_json(messages, tools)
943 }
944
945 fn list_models(&self) -> Vec<String> {
948 vec![]
949 }
950
951 fn supports_structured_output(&self) -> bool {
953 false
954 }
955
956 #[allow(async_fn_in_trait)]
967 async fn chat_typed<T>(&self, messages: &[Message]) -> Result<T, LlmError>
968 where
969 T: serde::de::DeserializeOwned + schemars::JsonSchema + 'static,
970 Self: Sized,
971 {
972 let (_, schema_json) = cached_schema::<T>()?;
973 let type_name = short_type_name::<T>();
974
975 let mut augmented = messages.to_vec();
976 let instruction = format!(
977 "Respond with a valid JSON object matching this schema. \
978 Output ONLY the JSON, no markdown fences or extra text.\n\n\
979 Type: {type_name}\nSchema:\n```json\n{schema_json}\n```"
980 );
981 augmented.insert(0, Message::from_legacy(Role::System, instruction));
982
983 let raw = self.chat(&augmented).await?;
984 let cleaned = strip_json_fences(&raw);
985 match serde_json::from_str::<T>(cleaned) {
986 Ok(val) => Ok(val),
987 Err(first_err) => {
988 augmented.push(Message::from_legacy(Role::Assistant, &raw));
989 augmented.push(Message::from_legacy(
990 Role::User,
991 format!(
992 "Your response was not valid JSON. Error: {first_err}. \
993 Please output ONLY valid JSON matching the schema."
994 ),
995 ));
996 let retry_raw = self.chat(&augmented).await?;
997 let retry_cleaned = strip_json_fences(&retry_raw);
998 serde_json::from_str::<T>(retry_cleaned).map_err(|e| {
999 LlmError::StructuredParse(format!("parse failed after retry: {e}"))
1000 })
1001 }
1002 }
1003 }
1004}
1005
1006fn strip_json_fences(s: &str) -> &str {
1010 s.trim()
1011 .trim_start_matches("```json")
1012 .trim_start_matches("```")
1013 .trim_end_matches("```")
1014 .trim()
1015}
1016
1017#[cfg(test)]
1018mod tests {
1019 use std::assert_matches;
1020 use tokio_stream::StreamExt;
1021
1022 use super::*;
1023
1024 struct StubProvider {
1025 response: String,
1026 }
1027
1028 impl LlmProvider for StubProvider {
1029 async fn chat(&self, _messages: &[Message]) -> Result<String, LlmError> {
1030 Ok(self.response.clone())
1031 }
1032
1033 async fn chat_stream(&self, messages: &[Message]) -> Result<ChatStream, LlmError> {
1034 let response = self.chat(messages).await?;
1035 Ok(Box::pin(tokio_stream::once(Ok(StreamChunk::Content(
1036 response,
1037 )))))
1038 }
1039
1040 fn supports_streaming(&self) -> bool {
1041 false
1042 }
1043
1044 async fn embed(&self, _text: &str) -> Result<Vec<f32>, LlmError> {
1045 Ok(vec![0.1, 0.2, 0.3])
1046 }
1047
1048 fn supports_embeddings(&self) -> bool {
1049 false
1050 }
1051
1052 fn name(&self) -> &'static str {
1053 "stub"
1054 }
1055 }
1056
1057 #[test]
1058 fn test_image_data_debug_redacts_bytes() {
1059 let img = ImageData {
1060 data: vec![0xAB, 0xCD, 0xEF],
1061 mime_type: "image/png".to_owned(),
1062 };
1063 let debug = format!("{img:?}");
1064 assert_eq!(debug, "[image: image/png, 3 bytes]");
1065 assert!(!debug.contains("171") && !debug.contains("205") && !debug.contains("239"));
1066 }
1067
1068 #[test]
1069 fn context_window_default_returns_none() {
1070 let provider = StubProvider {
1071 response: String::new(),
1072 };
1073 assert!(provider.context_window().is_none());
1074 }
1075
1076 #[test]
1077 fn supports_streaming_default_returns_false() {
1078 let provider = StubProvider {
1079 response: String::new(),
1080 };
1081 assert!(!provider.supports_streaming());
1082 }
1083
1084 #[test]
1085 fn supports_tool_use_default_returns_false() {
1086 let provider = StubProvider {
1092 response: String::new(),
1093 };
1094 assert!(!provider.supports_tool_use());
1095 }
1096
1097 #[tokio::test]
1098 async fn chat_stream_default_yields_single_chunk() {
1099 let provider = StubProvider {
1100 response: "hello world".into(),
1101 };
1102 let messages = vec![Message {
1103 role: Role::User,
1104 content: "test".into(),
1105 parts: vec![],
1106 metadata: MessageMetadata::default(),
1107 }];
1108
1109 let mut stream = provider.chat_stream(&messages).await.unwrap();
1110 let chunk = stream.next().await.unwrap().unwrap();
1111 assert_matches!(chunk, StreamChunk::Content(s) if s == "hello world");
1112 assert!(stream.next().await.is_none());
1113 }
1114
1115 #[tokio::test]
1116 async fn chat_stream_default_propagates_chat_error() {
1117 struct FailProvider;
1118
1119 impl LlmProvider for FailProvider {
1120 async fn chat(&self, _messages: &[Message]) -> Result<String, LlmError> {
1121 Err(LlmError::Unavailable)
1122 }
1123
1124 async fn chat_stream(&self, messages: &[Message]) -> Result<ChatStream, LlmError> {
1125 let response = self.chat(messages).await?;
1126 Ok(Box::pin(tokio_stream::once(Ok(StreamChunk::Content(
1127 response,
1128 )))))
1129 }
1130
1131 fn supports_streaming(&self) -> bool {
1132 false
1133 }
1134
1135 async fn embed(&self, _text: &str) -> Result<Vec<f32>, LlmError> {
1136 Err(LlmError::Unavailable)
1137 }
1138
1139 fn supports_embeddings(&self) -> bool {
1140 false
1141 }
1142
1143 fn name(&self) -> &'static str {
1144 "fail"
1145 }
1146 }
1147
1148 let provider = FailProvider;
1149 let messages = vec![Message {
1150 role: Role::User,
1151 content: "test".into(),
1152 parts: vec![],
1153 metadata: MessageMetadata::default(),
1154 }];
1155
1156 let result = provider.chat_stream(&messages).await;
1157 assert!(result.is_err());
1158 if let Err(e) = result {
1159 assert!(e.to_string().contains("provider unavailable"));
1160 }
1161 }
1162
1163 #[tokio::test]
1164 async fn stub_provider_embed_returns_vector() {
1165 let provider = StubProvider {
1166 response: String::new(),
1167 };
1168 let embedding = provider.embed("test").await.unwrap();
1169 assert_eq!(embedding, vec![0.1, 0.2, 0.3]);
1170 }
1171
1172 #[tokio::test]
1173 async fn fail_provider_embed_propagates_error() {
1174 struct FailProvider;
1175
1176 impl LlmProvider for FailProvider {
1177 async fn chat(&self, _messages: &[Message]) -> Result<String, LlmError> {
1178 Err(LlmError::Unavailable)
1179 }
1180
1181 async fn chat_stream(&self, messages: &[Message]) -> Result<ChatStream, LlmError> {
1182 let response = self.chat(messages).await?;
1183 Ok(Box::pin(tokio_stream::once(Ok(StreamChunk::Content(
1184 response,
1185 )))))
1186 }
1187
1188 fn supports_streaming(&self) -> bool {
1189 false
1190 }
1191
1192 async fn embed(&self, _text: &str) -> Result<Vec<f32>, LlmError> {
1193 Err(LlmError::EmbedUnsupported {
1194 provider: "fail".into(),
1195 })
1196 }
1197
1198 fn supports_embeddings(&self) -> bool {
1199 false
1200 }
1201
1202 fn name(&self) -> &'static str {
1203 "fail"
1204 }
1205 }
1206
1207 let provider = FailProvider;
1208 let result = provider.embed("test").await;
1209 assert!(result.is_err());
1210 assert!(
1211 result
1212 .unwrap_err()
1213 .to_string()
1214 .contains("embedding not supported")
1215 );
1216 }
1217
1218 #[test]
1219 fn role_serialization() {
1220 let system = Role::System;
1221 let user = Role::User;
1222 let assistant = Role::Assistant;
1223
1224 assert_eq!(serde_json::to_string(&system).unwrap(), "\"system\"");
1225 assert_eq!(serde_json::to_string(&user).unwrap(), "\"user\"");
1226 assert_eq!(serde_json::to_string(&assistant).unwrap(), "\"assistant\"");
1227 }
1228
1229 #[test]
1230 fn role_deserialization() {
1231 let system: Role = serde_json::from_str("\"system\"").unwrap();
1232 let user: Role = serde_json::from_str("\"user\"").unwrap();
1233 let assistant: Role = serde_json::from_str("\"assistant\"").unwrap();
1234
1235 assert_eq!(system, Role::System);
1236 assert_eq!(user, Role::User);
1237 assert_eq!(assistant, Role::Assistant);
1238 }
1239
1240 #[test]
1241 fn message_clone() {
1242 let msg = Message {
1243 role: Role::User,
1244 content: "test".into(),
1245 parts: vec![],
1246 metadata: MessageMetadata::default(),
1247 };
1248 let cloned = msg.clone();
1249 assert_eq!(cloned.role, msg.role);
1250 assert_eq!(cloned.content, msg.content);
1251 }
1252
1253 #[test]
1254 fn message_debug() {
1255 let msg = Message {
1256 role: Role::Assistant,
1257 content: "response".into(),
1258 parts: vec![],
1259 metadata: MessageMetadata::default(),
1260 };
1261 let debug = format!("{msg:?}");
1262 assert!(debug.contains("Assistant"));
1263 assert!(debug.contains("response"));
1264 }
1265
1266 #[test]
1267 fn message_serialization() {
1268 let msg = Message {
1269 role: Role::User,
1270 content: "hello".into(),
1271 parts: vec![],
1272 metadata: MessageMetadata::default(),
1273 };
1274 let json = serde_json::to_string(&msg).unwrap();
1275 assert!(json.contains("\"role\":\"user\""));
1276 assert!(json.contains("\"content\":\"hello\""));
1277 }
1278
1279 #[test]
1280 fn message_part_serde_round_trip() {
1281 let parts = vec![
1282 MessagePart::Text {
1283 text: "hello".into(),
1284 },
1285 MessagePart::ToolOutput {
1286 tool_name: "bash".into(),
1287 body: "output".into(),
1288 compacted_at: None,
1289 },
1290 MessagePart::Recall {
1291 text: "recall".into(),
1292 },
1293 MessagePart::CodeContext {
1294 text: "code".into(),
1295 },
1296 MessagePart::Summary {
1297 text: "summary".into(),
1298 },
1299 ];
1300 let json = serde_json::to_string(&parts).unwrap();
1301 let deserialized: Vec<MessagePart> = serde_json::from_str(&json).unwrap();
1302 assert_eq!(deserialized.len(), 5);
1303 }
1304
1305 #[test]
1306 fn from_legacy_creates_empty_parts() {
1307 let msg = Message::from_legacy(Role::User, "hello");
1308 assert_eq!(msg.role, Role::User);
1309 assert_eq!(msg.content, "hello");
1310 assert!(msg.parts.is_empty());
1311 assert_eq!(msg.to_llm_content(), "hello");
1312 }
1313
1314 #[test]
1315 fn from_parts_flattens_content() {
1316 let msg = Message::from_parts(
1317 Role::System,
1318 vec![MessagePart::Recall {
1319 text: "recalled data".into(),
1320 }],
1321 );
1322 assert_eq!(msg.content, "recalled data");
1323 assert_eq!(msg.to_llm_content(), "recalled data");
1324 assert_eq!(msg.parts.len(), 1);
1325 }
1326
1327 #[test]
1328 fn from_parts_tool_output_format() {
1329 let msg = Message::from_parts(
1330 Role::User,
1331 vec![MessagePart::ToolOutput {
1332 tool_name: "bash".into(),
1333 body: "hello world".into(),
1334 compacted_at: None,
1335 }],
1336 );
1337 assert!(msg.content.contains("[tool output: bash]"));
1338 assert!(msg.content.contains("hello world"));
1339 }
1340
1341 #[test]
1342 fn message_deserializes_without_parts() {
1343 let json = r#"{"role":"user","content":"hello"}"#;
1344 let msg: Message = serde_json::from_str(json).unwrap();
1345 assert_eq!(msg.content, "hello");
1346 assert!(msg.parts.is_empty());
1347 }
1348
1349 #[test]
1350 fn flatten_skips_compacted_tool_output_empty_body() {
1351 let msg = Message::from_parts(
1353 Role::User,
1354 vec![
1355 MessagePart::Text {
1356 text: "prefix ".into(),
1357 },
1358 MessagePart::ToolOutput {
1359 tool_name: "bash".into(),
1360 body: String::new(),
1361 compacted_at: Some(1234),
1362 },
1363 MessagePart::Text {
1364 text: " suffix".into(),
1365 },
1366 ],
1367 );
1368 assert!(msg.content.contains("(pruned)"));
1369 assert!(msg.content.contains("prefix "));
1370 assert!(msg.content.contains(" suffix"));
1371 }
1372
1373 #[test]
1374 fn flatten_compacted_tool_output_with_reference_renders_body() {
1375 let ref_notice = "[tool output pruned; full content at /tmp/overflow/big.txt]";
1377 let msg = Message::from_parts(
1378 Role::User,
1379 vec![MessagePart::ToolOutput {
1380 tool_name: "bash".into(),
1381 body: ref_notice.into(),
1382 compacted_at: Some(1234),
1383 }],
1384 );
1385 assert!(msg.content.contains(ref_notice));
1386 assert!(!msg.content.contains("(pruned)"));
1387 }
1388
1389 #[test]
1390 fn rebuild_content_syncs_after_mutation() {
1391 let mut msg = Message::from_parts(
1392 Role::User,
1393 vec![MessagePart::ToolOutput {
1394 tool_name: "bash".into(),
1395 body: "original".into(),
1396 compacted_at: None,
1397 }],
1398 );
1399 assert!(msg.content.contains("original"));
1400
1401 if let MessagePart::ToolOutput {
1402 ref mut compacted_at,
1403 ref mut body,
1404 ..
1405 } = msg.parts[0]
1406 {
1407 *compacted_at = Some(999);
1408 body.clear(); }
1410 msg.rebuild_content();
1411
1412 assert!(msg.content.contains("(pruned)"));
1413 assert!(!msg.content.contains("original"));
1414 }
1415
1416 #[test]
1417 fn message_part_tool_use_serde_round_trip() {
1418 let part = MessagePart::ToolUse {
1419 id: "toolu_123".into(),
1420 name: "bash".into(),
1421 input: serde_json::json!({"command": "ls"}),
1422 };
1423 let json = serde_json::to_string(&part).unwrap();
1424 let deserialized: MessagePart = serde_json::from_str(&json).unwrap();
1425 if let MessagePart::ToolUse { id, name, input } = deserialized {
1426 assert_eq!(id, "toolu_123");
1427 assert_eq!(name, "bash");
1428 assert_eq!(input["command"], "ls");
1429 } else {
1430 panic!("expected ToolUse");
1431 }
1432 }
1433
1434 #[test]
1435 fn message_part_tool_result_serde_round_trip() {
1436 let part = MessagePart::ToolResult {
1437 tool_use_id: "toolu_123".into(),
1438 content: "file1.rs\nfile2.rs".into(),
1439 is_error: false,
1440 };
1441 let json = serde_json::to_string(&part).unwrap();
1442 let deserialized: MessagePart = serde_json::from_str(&json).unwrap();
1443 if let MessagePart::ToolResult {
1444 tool_use_id,
1445 content,
1446 is_error,
1447 } = deserialized
1448 {
1449 assert_eq!(tool_use_id, "toolu_123");
1450 assert_eq!(content, "file1.rs\nfile2.rs");
1451 assert!(!is_error);
1452 } else {
1453 panic!("expected ToolResult");
1454 }
1455 }
1456
1457 #[test]
1458 fn message_part_tool_result_is_error_default() {
1459 let json = r#"{"kind":"tool_result","tool_use_id":"id","content":"err"}"#;
1460 let part: MessagePart = serde_json::from_str(json).unwrap();
1461 if let MessagePart::ToolResult { is_error, .. } = part {
1462 assert!(!is_error);
1463 } else {
1464 panic!("expected ToolResult");
1465 }
1466 }
1467
1468 #[test]
1469 fn chat_response_construction() {
1470 let text = ChatResponse::Text("hello".into());
1471 assert_matches!(text, ChatResponse::Text(s) if s == "hello");
1472
1473 let tool_use = ChatResponse::ToolUse {
1474 text: Some("I'll run that".into()),
1475 tool_calls: vec![ToolUseRequest {
1476 id: "1".into(),
1477 name: "bash".into(),
1478 input: serde_json::json!({}),
1479 }],
1480 thinking_blocks: vec![],
1481 };
1482 assert_matches!(tool_use, ChatResponse::ToolUse { .. });
1483 }
1484
1485 #[test]
1486 fn flatten_parts_tool_use() {
1487 let msg = Message::from_parts(
1488 Role::Assistant,
1489 vec![MessagePart::ToolUse {
1490 id: "t1".into(),
1491 name: "bash".into(),
1492 input: serde_json::json!({"command": "ls"}),
1493 }],
1494 );
1495 assert!(msg.content.contains("[tool_use: bash(t1)]"));
1496 }
1497
1498 #[test]
1499 fn flatten_parts_tool_result() {
1500 let msg = Message::from_parts(
1501 Role::User,
1502 vec![MessagePart::ToolResult {
1503 tool_use_id: "t1".into(),
1504 content: "output here".into(),
1505 is_error: false,
1506 }],
1507 );
1508 assert!(msg.content.contains("[tool_result: t1]"));
1509 assert!(msg.content.contains("output here"));
1510 }
1511
1512 #[test]
1513 fn tool_definition_serde_round_trip() {
1514 let def = ToolDefinition {
1515 name: "bash".into(),
1516 description: "Execute a shell command".into(),
1517 parameters: serde_json::json!({"type": "object"}),
1518 output_schema: None,
1519 };
1520 let json = serde_json::to_string(&def).unwrap();
1521 let deserialized: ToolDefinition = serde_json::from_str(&json).unwrap();
1522 assert_eq!(deserialized.name, "bash");
1523 assert_eq!(deserialized.description, "Execute a shell command");
1524 }
1525
1526 #[tokio::test]
1527 async fn chat_with_tools_default_delegates_to_chat() {
1528 let provider = StubProvider {
1529 response: "hello".into(),
1530 };
1531 let messages = vec![Message::from_legacy(Role::User, "test")];
1532 let result = provider.chat_with_tools(&messages, &[]).await.unwrap();
1533 assert_matches!(result, ChatResponse::Text(s) if s == "hello");
1534 }
1535
1536 #[test]
1537 fn tool_output_compacted_at_serde_default() {
1538 let json = r#"{"kind":"tool_output","tool_name":"bash","body":"out"}"#;
1539 let part: MessagePart = serde_json::from_str(json).unwrap();
1540 if let MessagePart::ToolOutput { compacted_at, .. } = part {
1541 assert!(compacted_at.is_none());
1542 } else {
1543 panic!("expected ToolOutput");
1544 }
1545 }
1546
1547 #[test]
1550 fn strip_json_fences_plain_json() {
1551 assert_eq!(strip_json_fences(r#"{"a": 1}"#), r#"{"a": 1}"#);
1552 }
1553
1554 #[test]
1555 fn strip_json_fences_with_json_fence() {
1556 assert_eq!(strip_json_fences("```json\n{\"a\": 1}\n```"), r#"{"a": 1}"#);
1557 }
1558
1559 #[test]
1560 fn strip_json_fences_with_plain_fence() {
1561 assert_eq!(strip_json_fences("```\n{\"a\": 1}\n```"), r#"{"a": 1}"#);
1562 }
1563
1564 #[test]
1565 fn strip_json_fences_whitespace() {
1566 assert_eq!(strip_json_fences(" \n "), "");
1567 }
1568
1569 #[test]
1570 fn strip_json_fences_empty() {
1571 assert_eq!(strip_json_fences(""), "");
1572 }
1573
1574 #[test]
1575 fn strip_json_fences_outer_whitespace() {
1576 assert_eq!(
1577 strip_json_fences(" ```json\n{\"a\": 1}\n``` "),
1578 r#"{"a": 1}"#
1579 );
1580 }
1581
1582 #[test]
1583 fn strip_json_fences_only_opening_fence() {
1584 assert_eq!(strip_json_fences("```json\n{\"a\": 1}"), r#"{"a": 1}"#);
1585 }
1586
1587 #[derive(Debug, serde::Deserialize, schemars::JsonSchema, PartialEq)]
1590 struct TestOutput {
1591 value: String,
1592 }
1593
1594 struct SequentialStub {
1595 responses: std::sync::Mutex<Vec<Result<String, LlmError>>>,
1596 }
1597
1598 impl SequentialStub {
1599 fn new(responses: Vec<Result<String, LlmError>>) -> Self {
1600 Self {
1601 responses: std::sync::Mutex::new(responses),
1602 }
1603 }
1604 }
1605
1606 impl LlmProvider for SequentialStub {
1607 async fn chat(&self, _messages: &[Message]) -> Result<String, LlmError> {
1608 let mut responses = self.responses.lock().unwrap();
1609 if responses.is_empty() {
1610 return Err(LlmError::Other("no more responses".into()));
1611 }
1612 responses.remove(0)
1613 }
1614
1615 async fn chat_stream(&self, messages: &[Message]) -> Result<ChatStream, LlmError> {
1616 let response = self.chat(messages).await?;
1617 Ok(Box::pin(tokio_stream::once(Ok(StreamChunk::Content(
1618 response,
1619 )))))
1620 }
1621
1622 fn supports_streaming(&self) -> bool {
1623 false
1624 }
1625
1626 async fn embed(&self, _text: &str) -> Result<Vec<f32>, LlmError> {
1627 Err(LlmError::EmbedUnsupported {
1628 provider: "sequential-stub".into(),
1629 })
1630 }
1631
1632 fn supports_embeddings(&self) -> bool {
1633 false
1634 }
1635
1636 fn name(&self) -> &'static str {
1637 "sequential-stub"
1638 }
1639 }
1640
1641 #[tokio::test]
1642 async fn chat_typed_happy_path() {
1643 let provider = StubProvider {
1644 response: r#"{"value": "hello"}"#.into(),
1645 };
1646 let messages = vec![Message::from_legacy(Role::User, "test")];
1647 let result: TestOutput = provider.chat_typed(&messages).await.unwrap();
1648 assert_eq!(
1649 result,
1650 TestOutput {
1651 value: "hello".into()
1652 }
1653 );
1654 }
1655
1656 #[tokio::test]
1657 async fn chat_typed_retry_succeeds() {
1658 let provider = SequentialStub::new(vec![
1659 Ok("not valid json".into()),
1660 Ok(r#"{"value": "ok"}"#.into()),
1661 ]);
1662 let messages = vec![Message::from_legacy(Role::User, "test")];
1663 let result: TestOutput = provider.chat_typed(&messages).await.unwrap();
1664 assert_eq!(result, TestOutput { value: "ok".into() });
1665 }
1666
1667 #[tokio::test]
1668 async fn chat_typed_both_fail() {
1669 let provider = SequentialStub::new(vec![Ok("bad json".into()), Ok("still bad".into())]);
1670 let messages = vec![Message::from_legacy(Role::User, "test")];
1671 let result = provider.chat_typed::<TestOutput>(&messages).await;
1672 let err = result.unwrap_err();
1673 assert!(err.to_string().contains("parse failed after retry"));
1674 }
1675
1676 #[tokio::test]
1677 async fn chat_typed_chat_error_propagates() {
1678 let provider = SequentialStub::new(vec![Err(LlmError::Unavailable)]);
1679 let messages = vec![Message::from_legacy(Role::User, "test")];
1680 let result = provider.chat_typed::<TestOutput>(&messages).await;
1681 assert_matches!(result, Err(LlmError::Unavailable));
1682 }
1683
1684 #[tokio::test]
1685 async fn chat_typed_strips_fences() {
1686 let provider = StubProvider {
1687 response: "```json\n{\"value\": \"fenced\"}\n```".into(),
1688 };
1689 let messages = vec![Message::from_legacy(Role::User, "test")];
1690 let result: TestOutput = provider.chat_typed(&messages).await.unwrap();
1691 assert_eq!(
1692 result,
1693 TestOutput {
1694 value: "fenced".into()
1695 }
1696 );
1697 }
1698
1699 #[test]
1700 fn supports_structured_output_default_false() {
1701 let provider = StubProvider {
1702 response: String::new(),
1703 };
1704 assert!(!provider.supports_structured_output());
1705 }
1706
1707 #[test]
1708 fn structured_parse_error_display() {
1709 let err = LlmError::StructuredParse("test error".into());
1710 assert_eq!(
1711 err.to_string(),
1712 "structured output parse failed: test error"
1713 );
1714 }
1715
1716 #[test]
1717 fn message_part_image_roundtrip_json() {
1718 let part = MessagePart::Image(Box::new(ImageData {
1719 data: vec![1, 2, 3, 4],
1720 mime_type: "image/jpeg".into(),
1721 }));
1722 let json = serde_json::to_string(&part).unwrap();
1723 let decoded: MessagePart = serde_json::from_str(&json).unwrap();
1724 match decoded {
1725 MessagePart::Image(img) => {
1726 assert_eq!(img.data, vec![1, 2, 3, 4]);
1727 assert_eq!(img.mime_type, "image/jpeg");
1728 }
1729 _ => panic!("expected Image variant"),
1730 }
1731 }
1732
1733 #[test]
1734 fn flatten_parts_includes_image_placeholder() {
1735 let msg = Message::from_parts(
1736 Role::User,
1737 vec![
1738 MessagePart::Text {
1739 text: "see this".into(),
1740 },
1741 MessagePart::Image(Box::new(ImageData {
1742 data: vec![0u8; 100],
1743 mime_type: "image/png".into(),
1744 })),
1745 ],
1746 );
1747 let content = msg.to_llm_content();
1748 assert!(content.contains("see this"));
1749 assert!(content.contains("[image: image/png"));
1750 }
1751
1752 #[test]
1753 fn supports_vision_default_false() {
1754 let provider = StubProvider {
1755 response: String::new(),
1756 };
1757 assert!(!provider.supports_vision());
1758 }
1759
1760 #[test]
1761 fn message_metadata_default_both_visible() {
1762 let m = MessageMetadata::default();
1763 assert!(m.visibility.is_agent_visible());
1764 assert!(m.visibility.is_user_visible());
1765 assert_eq!(m.visibility, MessageVisibility::Both);
1766 assert!(m.compacted_at.is_none());
1767 }
1768
1769 #[test]
1770 fn message_metadata_agent_only() {
1771 let m = MessageMetadata::agent_only();
1772 assert!(m.visibility.is_agent_visible());
1773 assert!(!m.visibility.is_user_visible());
1774 assert_eq!(m.visibility, MessageVisibility::AgentOnly);
1775 }
1776
1777 #[test]
1778 fn message_metadata_user_only() {
1779 let m = MessageMetadata::user_only();
1780 assert!(!m.visibility.is_agent_visible());
1781 assert!(m.visibility.is_user_visible());
1782 assert_eq!(m.visibility, MessageVisibility::UserOnly);
1783 }
1784
1785 #[test]
1786 fn message_metadata_serde_default() {
1787 let json = r#"{"role":"user","content":"hello"}"#;
1788 let msg: Message = serde_json::from_str(json).unwrap();
1789 assert!(msg.metadata.visibility.is_agent_visible());
1790 assert!(msg.metadata.visibility.is_user_visible());
1791 }
1792
1793 #[test]
1794 fn message_metadata_round_trip() {
1795 let msg = Message {
1796 role: Role::User,
1797 content: "test".into(),
1798 parts: vec![],
1799 metadata: MessageMetadata::agent_only(),
1800 };
1801 let json = serde_json::to_string(&msg).unwrap();
1802 let decoded: Message = serde_json::from_str(&json).unwrap();
1803 assert!(decoded.metadata.visibility.is_agent_visible());
1804 assert!(!decoded.metadata.visibility.is_user_visible());
1805 assert_eq!(decoded.metadata.visibility, MessageVisibility::AgentOnly);
1806 }
1807
1808 #[test]
1809 fn message_part_compaction_round_trip() {
1810 let part = MessagePart::Compaction {
1811 summary: "Context was summarized.".to_owned(),
1812 };
1813 let json = serde_json::to_string(&part).unwrap();
1814 let decoded: MessagePart = serde_json::from_str(&json).unwrap();
1815 assert!(
1816 matches!(decoded, MessagePart::Compaction { summary } if summary == "Context was summarized.")
1817 );
1818 }
1819
1820 #[test]
1821 fn flatten_parts_compaction_contributes_no_text() {
1822 let parts = vec![
1825 MessagePart::Text {
1826 text: "Hello".to_owned(),
1827 },
1828 MessagePart::Compaction {
1829 summary: "Summary".to_owned(),
1830 },
1831 ];
1832 let msg = Message::from_parts(Role::Assistant, parts);
1833 assert_eq!(msg.content.trim(), "Hello");
1835 }
1836
1837 #[test]
1838 fn stream_chunk_compaction_variant() {
1839 let chunk = StreamChunk::Compaction("A summary".to_owned());
1840 assert_matches!(chunk, StreamChunk::Compaction(s) if s == "A summary");
1841 }
1842
1843 #[test]
1844 fn short_type_name_extracts_last_segment() {
1845 struct MyOutput;
1846 assert_eq!(short_type_name::<MyOutput>(), "MyOutput");
1847 }
1848
1849 #[test]
1850 fn short_type_name_primitive_returns_full_name() {
1851 assert_eq!(short_type_name::<u32>(), "u32");
1853 assert_eq!(short_type_name::<bool>(), "bool");
1854 }
1855
1856 #[test]
1857 fn short_type_name_nested_path_returns_last() {
1858 assert_eq!(
1860 short_type_name::<std::collections::HashMap<u32, u32>>(),
1861 "HashMap<u32, u32>"
1862 );
1863 }
1864
1865 #[test]
1868 fn summary_roundtrip() {
1869 let part = MessagePart::Summary {
1870 text: "hello".to_string(),
1871 };
1872 let json = serde_json::to_string(&part).expect("serialization must not fail");
1873 assert!(
1874 json.contains("\"kind\":\"summary\""),
1875 "must use internally-tagged format, got: {json}"
1876 );
1877 assert!(
1878 !json.contains("\"Summary\""),
1879 "must not use externally-tagged format, got: {json}"
1880 );
1881 let decoded: MessagePart =
1882 serde_json::from_str(&json).expect("deserialization must not fail");
1883 match decoded {
1884 MessagePart::Summary { text } => assert_eq!(text, "hello"),
1885 other => panic!("expected MessagePart::Summary, got {other:?}"),
1886 }
1887 }
1888
1889 #[tokio::test]
1890 async fn embed_batch_default_empty_returns_empty() {
1891 let provider = StubProvider {
1892 response: String::new(),
1893 };
1894 let result = provider.embed_batch(&[]).await.unwrap();
1895 assert!(result.is_empty());
1896 }
1897
1898 #[tokio::test]
1899 async fn embed_batch_default_calls_embed_sequentially() {
1900 let provider = StubProvider {
1901 response: String::new(),
1902 };
1903 let texts = ["hello", "world", "foo"];
1904 let result = provider.embed_batch(&texts).await.unwrap();
1905 assert_eq!(result.len(), 3);
1906 for vec in &result {
1908 assert_eq!(vec, &[0.1_f32, 0.2, 0.3]);
1909 }
1910 }
1911
1912 #[test]
1913 fn message_visibility_db_roundtrip_both() {
1914 assert_eq!(MessageVisibility::Both.as_db_str(), "both");
1915 assert_eq!(
1916 MessageVisibility::from_db_str("both"),
1917 MessageVisibility::Both
1918 );
1919 }
1920
1921 #[test]
1922 fn message_visibility_db_roundtrip_agent_only() {
1923 assert_eq!(MessageVisibility::AgentOnly.as_db_str(), "agent_only");
1924 assert_eq!(
1925 MessageVisibility::from_db_str("agent_only"),
1926 MessageVisibility::AgentOnly
1927 );
1928 }
1929
1930 #[test]
1931 fn message_visibility_db_roundtrip_user_only() {
1932 assert_eq!(MessageVisibility::UserOnly.as_db_str(), "user_only");
1933 assert_eq!(
1934 MessageVisibility::from_db_str("user_only"),
1935 MessageVisibility::UserOnly
1936 );
1937 }
1938
1939 #[test]
1940 fn message_visibility_from_db_str_unknown_defaults_to_both() {
1941 assert_eq!(
1942 MessageVisibility::from_db_str("unknown_future_value"),
1943 MessageVisibility::Both
1944 );
1945 assert_eq!(MessageVisibility::from_db_str(""), MessageVisibility::Both);
1946 }
1947}