1use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
6pub enum BackendKind {
7 Gemini,
8 OpenAI,
9 Anthropic,
10 DeepSeek,
11 Mistral,
12 OpenRouter,
13 Ollama,
14 LlamaCpp,
15 ZAI,
16 Moonshot,
17 HuggingFace,
18 Minimax,
19 MiMo,
20 OpenCodeZen,
21 OpenCodeGo,
22 Qwen,
23 StepFun,
24 Evolink,
25 Poolside,
26}
27
28#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
29pub struct Usage {
30 pub prompt_tokens: u32,
31 pub completion_tokens: u32,
32 pub total_tokens: u32,
33 pub cached_prompt_tokens: Option<u32>,
34 pub cache_creation_tokens: Option<u32>,
35 pub cache_read_tokens: Option<u32>,
36 #[serde(default, skip_serializing_if = "Option::is_none")]
39 pub iterations: Option<Vec<serde_json::Value>>,
40}
41
42impl Usage {
43 #[inline]
44 fn has_cache_read_metric(&self) -> bool {
45 self.cache_read_tokens.is_some() || self.cached_prompt_tokens.is_some()
46 }
47
48 #[inline]
49 fn has_any_cache_metrics(&self) -> bool {
50 self.has_cache_read_metric() || self.cache_creation_tokens.is_some()
51 }
52
53 #[inline]
54 pub fn cache_read_tokens_or_fallback(&self) -> u32 {
55 self.cache_read_tokens.or(self.cached_prompt_tokens).unwrap_or(0)
56 }
57
58 #[inline]
59 pub fn cache_creation_tokens_or_zero(&self) -> u32 {
60 self.cache_creation_tokens.unwrap_or(0)
61 }
62
63 #[inline]
64 pub fn cache_hit_rate(&self) -> Option<f64> {
65 if !self.has_any_cache_metrics() {
66 return None;
67 }
68 let read = self.cache_read_tokens_or_fallback() as f64;
69 let creation = self.cache_creation_tokens_or_zero() as f64;
70 let total = read + creation;
71 if total > 0.0 {
72 Some((read / total) * 100.0)
73 } else {
74 None
75 }
76 }
77
78 #[inline]
79 pub fn is_cache_hit(&self) -> Option<bool> {
80 self.has_any_cache_metrics().then(|| self.cache_read_tokens_or_fallback() > 0)
81 }
82
83 #[inline]
84 pub fn is_cache_miss(&self) -> Option<bool> {
85 self.has_any_cache_metrics().then(|| {
86 self.cache_creation_tokens_or_zero() > 0 && self.cache_read_tokens_or_fallback() == 0
87 })
88 }
89
90 #[inline]
91 pub fn total_cache_tokens(&self) -> u32 {
92 let read = self.cache_read_tokens_or_fallback();
93 let creation = self.cache_creation_tokens_or_zero();
94 read + creation
95 }
96
97 #[inline]
98 pub fn cache_savings_ratio(&self) -> Option<f64> {
99 if !self.has_cache_read_metric() {
100 return None;
101 }
102 let read = self.cache_read_tokens_or_fallback() as f64;
103 let prompt = self.prompt_tokens as f64;
104 if prompt > 0.0 {
105 Some(read / prompt)
106 } else {
107 None
108 }
109 }
110}
111
112#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
114pub struct BalanceInfo {
115 pub display: String,
117 pub is_available: bool,
119}
120
121#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct DeepSeekBalanceResponse {
124 pub is_available: bool,
125 pub balance_infos: Vec<DeepSeekCurrencyBalance>,
126}
127
128#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct DeepSeekCurrencyBalance {
130 pub currency: String,
131 pub total_balance: String,
132 #[serde(default)]
133 pub granted_balance: String,
134 #[serde(default)]
135 pub topped_up_balance: String,
136}
137
138impl From<DeepSeekBalanceResponse> for BalanceInfo {
139 fn from(resp: DeepSeekBalanceResponse) -> Self {
140 let display = resp
141 .balance_infos
142 .first()
143 .map(|b| {
144 let symbol = match b.currency.as_str() {
145 "CNY" => "¥",
146 "USD" => "$",
147 _ => &b.currency,
148 };
149 format!("{}{}", b.total_balance, symbol)
150 })
151 .unwrap_or_else(|| "N/A".to_string());
152 BalanceInfo { display, is_available: resp.is_available }
153 }
154}
155
156#[cfg(test)]
157mod usage_tests {
158 use super::Usage;
159
160 #[test]
161 fn cache_helpers_fall_back_to_cached_prompt_tokens() {
162 let usage = Usage {
163 prompt_tokens: 1_000,
164 completion_tokens: 200,
165 total_tokens: 1_200,
166 cached_prompt_tokens: Some(600),
167 cache_creation_tokens: Some(150),
168 cache_read_tokens: None,
169 iterations: None,
170 };
171
172 assert_eq!(usage.cache_read_tokens_or_fallback(), 600);
173 assert_eq!(usage.cache_creation_tokens_or_zero(), 150);
174 assert_eq!(usage.total_cache_tokens(), 750);
175 assert_eq!(usage.is_cache_hit(), Some(true));
176 assert_eq!(usage.is_cache_miss(), Some(false));
177 assert_eq!(usage.cache_savings_ratio(), Some(0.6));
178 assert_eq!(usage.cache_hit_rate(), Some(80.0));
179 }
180
181 #[test]
182 fn cache_helpers_preserve_unknown_without_metrics() {
183 let usage = Usage {
184 prompt_tokens: 1_000,
185 completion_tokens: 200,
186 total_tokens: 1_200,
187 cached_prompt_tokens: None,
188 cache_creation_tokens: None,
189 cache_read_tokens: None,
190 iterations: None,
191 };
192
193 assert_eq!(usage.total_cache_tokens(), 0);
194 assert_eq!(usage.is_cache_hit(), None);
195 assert_eq!(usage.is_cache_miss(), None);
196 assert_eq!(usage.cache_savings_ratio(), None);
197 assert_eq!(usage.cache_hit_rate(), None);
198 }
199}
200
201#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
202pub enum FinishReason {
203 #[default]
204 Stop,
205 Length,
206 ToolCalls,
207 ContentFilter,
208 Pause,
209 Refusal,
210 Error(String),
211}
212
213#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
215pub struct ToolCall {
216 pub id: String,
218
219 #[serde(rename = "type")]
221 pub call_type: String,
222
223 #[serde(skip_serializing_if = "Option::is_none")]
225 pub function: Option<FunctionCall>,
226
227 #[serde(skip_serializing_if = "Option::is_none")]
229 pub text: Option<String>,
230
231 #[serde(skip_serializing_if = "Option::is_none")]
233 pub thought_signature: Option<String>,
234}
235
236#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
238pub struct FunctionCall {
239 #[serde(default, skip_serializing_if = "Option::is_none")]
241 pub namespace: Option<String>,
242
243 pub name: String,
245
246 pub arguments: String,
248}
249
250impl ToolCall {
251 pub fn function(id: String, name: String, arguments: String) -> Self {
253 Self::function_with_namespace(id, None, name, arguments)
254 }
255
256 pub fn function_with_namespace(
258 id: String,
259 namespace: Option<String>,
260 name: String,
261 arguments: String,
262 ) -> Self {
263 Self {
264 id,
265 call_type: "function".to_owned(),
266 function: Some(FunctionCall { namespace, name, arguments }),
267 text: None,
268 thought_signature: None,
269 }
270 }
271
272 pub fn custom(id: String, name: String, text: String) -> Self {
274 Self {
275 id,
276 call_type: "custom".to_owned(),
277 function: Some(FunctionCall { namespace: None, name, arguments: text.clone() }),
278 text: Some(text),
279 thought_signature: None,
280 }
281 }
282
283 pub fn is_custom(&self) -> bool {
285 self.call_type == "custom"
286 }
287
288 pub fn tool_name(&self) -> Option<&str> {
290 self.function.as_ref().map(|function| function.name.as_str())
291 }
292
293 pub fn raw_input(&self) -> Option<&str> {
295 self.text
296 .as_deref()
297 .or_else(|| self.function.as_ref().map(|function| function.arguments.as_str()))
298 }
299
300 pub fn parsed_arguments(&self) -> Result<serde_json::Value, serde_json::Error> {
302 if let Some(ref func) = self.function {
303 parse_tool_arguments(&func.arguments)
304 } else {
305 serde_json::from_str("")
307 }
308 }
309
310 pub fn execution_arguments(&self) -> Result<serde_json::Value, serde_json::Error> {
316 if self.is_custom() {
317 return Ok(serde_json::Value::String(self.raw_input().unwrap_or_default().to_string()));
318 }
319
320 self.parsed_arguments()
321 }
322
323 pub fn validate(&self) -> Result<(), String> {
325 if self.id.is_empty() {
326 return Err("Tool call ID cannot be empty".to_owned());
327 }
328
329 match self.call_type.as_str() {
330 "function" => {
331 if let Some(func) = &self.function {
332 if func.name.is_empty() {
333 return Err("Function name cannot be empty".to_owned());
334 }
335 if let Err(e) = self.parsed_arguments() {
337 return Err(format!("Invalid JSON in function arguments: {e}"));
338 }
339 } else {
340 return Err("Function tool call missing function details".to_owned());
341 }
342 }
343 "custom" => {
344 if let Some(func) = &self.function {
346 if func.name.is_empty() {
347 return Err("Custom tool name cannot be empty".to_owned());
348 }
349 } else {
350 return Err("Custom tool call missing function details".to_owned());
351 }
352 }
353 _ => return Err(format!("Unsupported tool call type: {}", self.call_type)),
354 }
355
356 Ok(())
357 }
358}
359
360fn parse_tool_arguments(raw_arguments: &str) -> Result<serde_json::Value, serde_json::Error> {
361 let trimmed = raw_arguments.trim();
362 match serde_json::from_str(trimmed) {
363 Ok(parsed) => Ok(parsed),
364 Err(primary_error) => {
365 if let Some(candidate) = extract_balanced_json(trimmed)
366 && let Ok(parsed) = serde_json::from_str(candidate)
367 {
368 return Ok(parsed);
369 }
370 if let Some(candidate) = repair_tag_polluted_json(trimmed)
371 && let Ok(parsed) = serde_json::from_str(&candidate)
372 {
373 return Ok(parsed);
374 }
375 if let Some(repaired) = close_incomplete_json_prefix(trimmed)
376 && let Ok(parsed) = serde_json::from_str(&repaired)
377 {
378 return Ok(parsed);
379 }
380 Err(primary_error)
381 }
382 }
383}
384
385fn extract_balanced_json(input: &str) -> Option<&str> {
386 let start = input.find(['{', '['])?;
387 let opening = input.as_bytes().get(start).copied()?;
388 let closing = match opening {
389 b'{' => b'}',
390 b'[' => b']',
391 _ => return None,
392 };
393
394 let mut depth = 0usize;
395 let mut in_string = false;
396 let mut escaped = false;
397
398 for (offset, ch) in input[start..].char_indices() {
399 if in_string {
400 if escaped {
401 escaped = false;
402 continue;
403 }
404 if ch == '\\' {
405 escaped = true;
406 continue;
407 }
408 if ch == '"' {
409 in_string = false;
410 }
411 continue;
412 }
413
414 match ch {
415 '"' => in_string = true,
416 _ if ch as u32 == opening as u32 => depth += 1,
417 _ if ch as u32 == closing as u32 => {
418 depth = depth.saturating_sub(1);
419 if depth == 0 {
420 let end = start + offset + ch.len_utf8();
421 return input.get(start..end);
422 }
423 }
424 _ => {}
425 }
426 }
427
428 None
429}
430
431fn repair_tag_polluted_json(input: &str) -> Option<String> {
432 let start = input.find(['{', '['])?;
433 let candidate = input.get(start..)?;
434 let boundary = find_provider_markup_boundary(candidate)?;
435 if boundary == 0 {
436 return None;
437 }
438
439 close_incomplete_json_prefix(candidate[..boundary].trim_end())
440}
441
442fn find_provider_markup_boundary(input: &str) -> Option<usize> {
443 const PROVIDER_MARKERS: &[&str] = &[
444 "<</",
445 "</parameter>",
446 "</invoke>",
447 "</minimax:tool_call>",
448 "<minimax:tool_call>",
449 "<parameter name=\"",
450 "<invoke name=\"",
451 "<tool_call>",
452 "</tool_call>",
453 ];
454
455 input.char_indices().find_map(|(offset, _)| {
456 let rest = input.get(offset..)?;
457 PROVIDER_MARKERS.iter().any(|marker| rest.starts_with(marker)).then_some(offset)
458 })
459}
460
461fn close_incomplete_json_prefix(prefix: &str) -> Option<String> {
462 if prefix.is_empty() {
463 return None;
464 }
465
466 let mut repaired = String::with_capacity(prefix.len() + 8);
467 let mut expected_closers = Vec::new();
468 let mut in_string = false;
469 let mut escaped = false;
470
471 for ch in prefix.chars() {
472 repaired.push(ch);
473
474 if in_string {
475 if escaped {
476 escaped = false;
477 continue;
478 }
479
480 match ch {
481 '\\' => escaped = true,
482 '"' => in_string = false,
483 _ => {}
484 }
485 continue;
486 }
487
488 match ch {
489 '"' => in_string = true,
490 '{' => expected_closers.push('}'),
491 '[' => expected_closers.push(']'),
492 '}' | ']' if expected_closers.pop() != Some(ch) => return None,
493 '}' | ']' => {}
494 _ => {}
495 }
496 }
497
498 if in_string {
499 repaired.push('"');
500 }
501 for closer in expected_closers.drain(..) {
502 repaired.push(closer);
503 }
504
505 Some(repaired)
506}
507
508#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
510pub struct LLMResponse {
511 pub content: Option<String>,
513
514 pub tool_calls: Option<Vec<ToolCall>>,
516
517 pub model: String,
519
520 pub usage: Option<Usage>,
522
523 pub finish_reason: FinishReason,
525
526 pub reasoning: Option<String>,
528
529 pub reasoning_details: Option<Vec<String>>,
531
532 pub tool_references: Vec<String>,
534
535 pub request_id: Option<String>,
537
538 pub organization_id: Option<String>,
540
541 pub compaction: Option<String>,
546}
547
548impl LLMResponse {
549 pub fn new(model: impl Into<String>, content: impl Into<String>) -> Self {
551 Self {
552 content: Some(content.into()),
553 tool_calls: None,
554 model: model.into(),
555 usage: None,
556 finish_reason: FinishReason::Stop,
557 reasoning: None,
558 reasoning_details: None,
559 tool_references: Vec::new(),
560 request_id: None,
561 organization_id: None,
562 compaction: None,
563 }
564 }
565
566 pub fn content_text(&self) -> &str {
568 self.content.as_deref().unwrap_or("")
569 }
570
571 pub fn content_string(&self) -> String {
573 self.content.clone().unwrap_or_default()
574 }
575}
576
577#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
578pub struct LLMErrorMetadata {
579 pub provider: Option<String>,
580 pub status: Option<u16>,
581 pub code: Option<String>,
582 pub request_id: Option<String>,
583 pub organization_id: Option<String>,
584 pub retry_after: Option<String>,
585 pub message: Option<String>,
586}
587
588impl LLMErrorMetadata {
589 #[must_use]
592 pub fn new(
593 provider: impl Into<String>,
594 status: Option<u16>,
595 code: Option<String>,
596 request_id: Option<String>,
597 organization_id: Option<String>,
598 retry_after: Option<String>,
599 message: Option<String>,
600 ) -> Box<Self> {
601 Box::new(Self {
602 provider: Some(provider.into()),
603 status,
604 code,
605 request_id,
606 organization_id,
607 retry_after,
608 message,
609 })
610 }
611}
612
613#[derive(Debug, thiserror::Error, Serialize, Deserialize, Clone)]
615#[serde(tag = "type", rename_all = "snake_case")]
616pub enum LLMError {
617 #[error("Authentication failed: {message}")]
618 Authentication {
619 message: String,
620 metadata: Option<Box<LLMErrorMetadata>>,
621 },
622 #[error("Rate limit exceeded")]
623 RateLimit {
624 metadata: Option<Box<LLMErrorMetadata>>,
625 },
626 #[error("Invalid request: {message}")]
627 InvalidRequest {
628 message: String,
629 metadata: Option<Box<LLMErrorMetadata>>,
630 },
631 #[error("Network error: {message}")]
632 Network {
633 message: String,
634 metadata: Option<Box<LLMErrorMetadata>>,
635 },
636 #[error("Provider error: {message}")]
637 Provider {
638 message: String,
639 metadata: Option<Box<LLMErrorMetadata>>,
640 },
641}
642
643#[cfg(test)]
644mod tests {
645 use super::ToolCall;
646 use serde_json::json;
647
648 #[test]
649 fn parsed_arguments_accepts_trailing_characters() {
650 let call = ToolCall::function(
651 "call_read".to_string(),
652 "exec_command".to_string(),
653 r#"{"path":"src/main.rs"} trailing text"#.to_string(),
654 );
655
656 let parsed = call.parsed_arguments().expect("arguments with trailing text should recover");
657 assert_eq!(parsed, json!({"path":"src/main.rs"}));
658 }
659
660 #[test]
661 fn parsed_arguments_accepts_code_fenced_json() {
662 let call = ToolCall::function(
663 "call_read".to_string(),
664 "exec_command".to_string(),
665 "```json\n{\"path\":\"src/lib.rs\",\"limit\":25}\n```".to_string(),
666 );
667
668 let parsed = call.parsed_arguments().expect("code-fenced arguments should recover");
669 assert_eq!(parsed, json!({"path":"src/lib.rs","limit":25}));
670 }
671
672 #[test]
673 fn parsed_arguments_recovers_truncated_json_missing_closing_brace() {
674 let call = ToolCall::function(
675 "call_search".to_string(),
676 "code_search".to_string(),
677 r#"{"query":"context","path":".","file_types":["rust"],"result_types":["definition"],"max_results":20"#
678 .to_string(),
679 );
680
681 let parsed = call
682 .parsed_arguments()
683 .expect("truncated JSON missing closing brace should recover");
684 assert_eq!(
685 parsed,
686 json!({
687 "query": "context",
688 "path": ".",
689 "file_types": ["rust"],
690 "result_types": ["definition"],
691 "max_results": 20
692 })
693 );
694 }
695
696 #[test]
697 fn parsed_arguments_rejects_incomplete_json() {
698 let call = ToolCall::function(
699 "call_read".to_string(),
700 "exec_command".to_string(),
701 r#"{"path":"src/main.rs","limit""#.to_string(),
702 );
703
704 assert!(call.parsed_arguments().is_err());
705 }
706
707 #[test]
708 fn parsed_arguments_recovers_truncated_minimax_markup() {
709 let call = ToolCall::function(
710 "call_search".to_string(),
711 "code_search".to_string(),
712 "{\"query\":\"persistent_memory\",\"file_types\":[\"rust\"],\"result_types\":[\"text\"],\"max_results\":20,\"path\":\"crates/codegen/vtcode-core/src</parameter>\n<</invoke>\n</minimax:tool_call>".to_string(),
713 );
714
715 let parsed = call.parsed_arguments().expect("minimax markup spillover should recover");
716 assert_eq!(
717 parsed,
718 json!({
719 "query": "persistent_memory",
720 "path": "crates/codegen/vtcode-core/src",
721 "file_types": ["rust"],
722 "result_types": ["text"],
723 "max_results": 20
724 })
725 );
726 }
727
728 #[test]
729 fn function_call_serializes_optional_namespace() {
730 let call = ToolCall::function_with_namespace(
731 "call_read".to_string(),
732 Some("workspace".to_string()),
733 "exec_command".to_string(),
734 r#"{"path":"src/main.rs"}"#.to_string(),
735 );
736
737 let json = serde_json::to_value(&call).expect("tool call should serialize");
738 assert_eq!(json["function"]["namespace"], "workspace");
739 assert_eq!(json["function"]["name"], "exec_command");
740 }
741
742 #[test]
743 fn custom_tool_call_exposes_raw_execution_arguments() {
744 let patch = "*** Begin Patch\n*** End Patch\n".to_string();
745 let call =
746 ToolCall::custom("call_patch".to_string(), "apply_patch".to_string(), patch.clone());
747
748 assert!(call.is_custom());
749 assert_eq!(call.tool_name(), Some("apply_patch"));
750 assert_eq!(call.raw_input(), Some(patch.as_str()));
751 assert_eq!(call.execution_arguments().expect("custom arguments"), json!(patch));
752 assert!(
753 call.parsed_arguments().is_err(),
754 "custom tool payload should stay freeform rather than JSON"
755 );
756 }
757}