1use async_trait::async_trait;
18use futures::StreamExt;
19use reqwest::{Client, Url};
20use serde::{Deserialize, Serialize};
21use serde_json::Value;
22use std::sync::{Arc, Mutex};
23
24use crate::driver_registry::{
25 ChatDriver, LlmCallConfig, LlmCompletionMetadata, LlmContentPart, LlmMessage,
26 LlmMessageContent, LlmMessageRole, LlmResponseStream, LlmStreamEvent, disjoint_prompt_tokens,
27};
28use crate::error::{AgentLoopError, LlmErrorKind, Result};
29use crate::llm_retry::{
30 LlmRetryConfig, RateLimitInfo, RetryDecision, RetryMetadata, SendOutcome, is_rate_limit_status,
31 retry_request, send_error_message,
32};
33use crate::runtime_provider::ProviderEndpoint;
34use crate::stream_accumulator::StreamToolCallAccumulator;
35use crate::stream_reconnect::connect_sse_with_reconnect;
36use crate::tool_types::ToolDefinition;
37use crate::user_facing_error::is_provider_quota_message;
38
39pub fn is_azure_openai_api_url(api_url: &str) -> bool {
40 Url::parse(api_url)
41 .ok()
42 .and_then(|url| url.host_str().map(|host| host.to_ascii_lowercase()))
43 .is_some_and(|host| {
44 host.ends_with(".openai.azure.com") || host.ends_with(".services.ai.azure.com")
45 })
46}
47
48pub fn is_openai_api_url(api_url: &str) -> bool {
52 Url::parse(api_url)
53 .ok()
54 .and_then(|url| url.host_str().map(|host| host.to_ascii_lowercase()))
55 .is_some_and(|host| host == "api.openai.com")
56}
57
58const OPENAI_MODELS_URL: &str = "https://api.openai.com/v1/models";
68
69pub fn url_host_eq(api_url: &str, host: &str) -> bool {
71 Url::parse(api_url)
72 .ok()
73 .and_then(|url| url.host_str().map(str::to_owned))
74 .is_some_and(|h| h.eq_ignore_ascii_case(host))
75}
76
77pub fn normalize_api_url(base_url: &str, endpoint_suffix: &str) -> String {
80 let trimmed = base_url.trim_end_matches('/');
81 if trimmed.ends_with(endpoint_suffix) {
82 trimmed.to_string()
83 } else {
84 format!("{trimmed}{endpoint_suffix}")
85 }
86}
87
88pub fn models_url_for_api_url(api_url: &str) -> String {
90 let trimmed = api_url.trim_end_matches('/');
91
92 if let Some(prefix) = trimmed.strip_suffix("/responses") {
93 return format!("{prefix}/models");
94 }
95 if let Some(prefix) = trimmed.strip_suffix("/chat/completions") {
96 return format!("{prefix}/models");
97 }
98 if trimmed.ends_with("/models") {
99 return trimmed.to_string();
100 }
101 if trimmed.ends_with("/v1") || trimmed.ends_with("/openai/v1") {
102 return format!("{trimmed}/models");
103 }
104
105 OPENAI_MODELS_URL.to_string()
106}
107
108pub fn models_api_status_error(status: reqwest::StatusCode) -> AgentLoopError {
111 AgentLoopError::llm(format!("Models API returned status {status}"))
112}
113
114#[derive(Clone)]
137pub struct OpenAIProtocolChatDriver {
138 client: Client,
139 retry_config: LlmRetryConfig,
141}
142
143impl OpenAIProtocolChatDriver {
144 pub fn new() -> Self {
146 Self {
147 client: crate::driver_helpers::shared_streaming_http_client(),
148 retry_config: LlmRetryConfig::default(),
149 }
150 }
151
152 pub fn with_retry_config(mut self, config: LlmRetryConfig) -> Self {
154 self.retry_config = config;
155 self
156 }
157
158 pub fn client(&self) -> &Client {
160 &self.client
161 }
162
163 async fn send_chat_completion_request(
172 &self,
173 endpoint: &ProviderEndpoint,
174 api_url: &str,
175 request: &OpenAiRequest,
176 model: &str,
177 retries_consumed: u32,
178 ) -> Result<(reqwest::Response, RetryMetadata)> {
179 let last_error: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
180 let mut retry_config = self.retry_config.clone();
181 retry_config.max_retries = retry_config.max_retries.saturating_sub(retries_consumed);
182
183 let body = serde_json::to_vec(request)
184 .map_err(|e| AgentLoopError::llm(format!("failed to serialize request: {e}")))?;
185 retry_request(
186 &retry_config,
187 "OpenAIProtocolDriver",
188 || async {
189 let resolved = endpoint
190 .resolve("POST", api_url, &body)
191 .await
192 .map_err(SendOutcome::Fatal)?;
193 let mut request_builder = self.client.post(&resolved.url);
194 for (name, value) in resolved.headers {
195 request_builder = request_builder.header(name, value);
196 }
197 request_builder
198 .header("Content-Type", "application/json")
199 .body(body.clone())
200 .send()
201 .await
202 .map_err(SendOutcome::Send)
203 },
204 |response, attempts, can_retry| {
205 let last_error = Arc::clone(&last_error);
206 let model = model.to_string();
207 async move {
208 let status = response.status();
209
210 if can_retry {
211 let rate_limit_info = if is_rate_limit_status(status) {
213 Some(RateLimitInfo::from_openai_headers(response.headers()))
214 } else {
215 None
216 };
217
218 let error_text = response.text().await.unwrap_or_default();
219
220 if is_openai_request_too_large(status, &error_text) {
222 return RetryDecision::Terminal(AgentLoopError::request_too_large(
223 format!("OpenAI API error ({}): {}", status, error_text),
224 ));
225 }
226
227 if is_provider_quota_message(&error_text) {
230 return RetryDecision::Terminal(AgentLoopError::llm_kind(
231 LlmErrorKind::QuotaExhausted,
232 format!("OpenAI API error ({}): {}", status, error_text),
233 ));
234 }
235
236 let wait = rate_limit_info
237 .as_ref()
238 .map(|info| info.recommended_wait(&self.retry_config, attempts))
239 .unwrap_or_else(|| self.retry_config.calculate_backoff(attempts));
240
241 *last_error.lock().unwrap() = Some(error_text);
242 return RetryDecision::Retry {
243 wait,
244 rate_limit_info,
245 };
246 }
247
248 let error_text = response.text().await.unwrap_or_default();
250 let error_msg = format!("OpenAI API error ({}): {}", status, error_text);
251
252 if is_openai_model_not_found(status, &error_text) {
254 return RetryDecision::Terminal(AgentLoopError::model_not_available(model));
255 }
256
257 if is_openai_request_too_large(status, &error_text) {
259 return RetryDecision::Terminal(AgentLoopError::request_too_large(
260 error_msg,
261 ));
262 }
263
264 let kind = LlmErrorKind::from_provider_status(status.as_u16(), &error_text);
267
268 if attempts > 0 {
269 return RetryDecision::Terminal(AgentLoopError::llm_kind(
270 kind,
271 format!(
272 "{} (after {} retries, last error: {})",
273 error_msg,
274 attempts,
275 last_error.lock().unwrap().take().unwrap_or_default()
276 ),
277 ));
278 }
279
280 RetryDecision::Terminal(AgentLoopError::llm_kind(kind, error_msg))
281 }
282 },
283 |e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
284 )
285 .await
286 }
287
288 fn convert_role(role: &LlmMessageRole) -> &'static str {
289 match role {
290 LlmMessageRole::System => "system",
291 LlmMessageRole::User => "user",
292 LlmMessageRole::Assistant => "assistant",
293 LlmMessageRole::Tool => "tool",
294 }
295 }
296
297 fn convert_message(msg: &LlmMessage) -> OpenAiMessage {
298 let content = match &msg.content {
299 LlmMessageContent::Text(text) => OpenAiContent::Text(text.clone()),
300 LlmMessageContent::Parts(parts) => {
301 let openai_parts: Vec<OpenAiContentPart> = parts
302 .iter()
303 .map(|part| match part {
304 LlmContentPart::Text { text } => OpenAiContentPart::Text {
305 r#type: "text".to_string(),
306 text: text.clone(),
307 },
308 LlmContentPart::Image { url } => OpenAiContentPart::ImageUrl {
309 r#type: "image_url".to_string(),
310 image_url: OpenAiImageUrl { url: url.clone() },
311 },
312 LlmContentPart::Audio { url } => OpenAiContentPart::InputAudio {
313 r#type: "input_audio".to_string(),
314 input_audio: OpenAiInputAudio {
315 data: url.clone(),
316 format: "wav".to_string(),
317 },
318 },
319 })
320 .collect();
321 OpenAiContent::Parts(openai_parts)
322 }
323 };
324
325 let tool_calls = if msg.role == LlmMessageRole::Assistant {
327 msg.tool_calls.as_ref().map(|calls| {
328 calls
329 .iter()
330 .map(|tc| OpenAiToolCall {
331 id: tc.id.clone(),
332 r#type: "function".to_string(),
333 function: OpenAiFunctionCall {
334 name: tc.name.clone(),
335 arguments: serde_json::to_string(&tc.arguments).unwrap_or_default(),
336 },
337 })
338 .collect()
339 })
340 } else {
341 None
342 };
343
344 OpenAiMessage {
345 role: Self::convert_role(&msg.role).to_string(),
346 content: Some(content),
347 tool_calls,
348 tool_call_id: msg.tool_call_id.clone(),
349 }
350 }
351
352 fn convert_tools(tools: &[ToolDefinition]) -> Vec<OpenAiTool> {
353 tools
354 .iter()
355 .map(|tool| OpenAiTool {
356 r#type: "function".to_string(),
357 function: OpenAiFunction {
358 name: tool.name().to_string(),
359 description: tool.description().to_string(),
360 parameters: crate::tool_schema_compat::sanitize_openai_tool_schema(
361 tool.parameters(),
362 ),
363 },
364 })
365 .collect()
366 }
367}
368
369impl Default for OpenAIProtocolChatDriver {
370 fn default() -> Self {
371 Self::new()
372 }
373}
374
375fn drop_orphaned_tool_messages(messages: &[LlmMessage]) -> Vec<LlmMessage> {
379 use std::collections::HashSet;
380
381 let visible_call_ids: HashSet<&str> = messages
382 .iter()
383 .filter(|m| m.role == LlmMessageRole::Assistant)
384 .flat_map(|m| m.tool_calls.iter().flatten())
385 .map(|tc| tc.id.as_str())
386 .collect();
387
388 if visible_call_ids.is_empty() {
389 return messages
390 .iter()
391 .filter(|m| m.role != LlmMessageRole::Tool)
392 .cloned()
393 .collect();
394 }
395
396 messages
397 .iter()
398 .filter(|m| {
399 if m.role == LlmMessageRole::Tool {
400 return m
401 .tool_call_id
402 .as_deref()
403 .is_none_or(|id| visible_call_ids.contains(id));
404 }
405 true
406 })
407 .cloned()
408 .collect()
409}
410
411#[async_trait]
412impl ChatDriver for OpenAIProtocolChatDriver {
413 async fn chat_completion_stream(
414 &self,
415 endpoint: &ProviderEndpoint,
416 messages: Vec<LlmMessage>,
417 config: &LlmCallConfig,
418 ) -> Result<LlmResponseStream> {
419 let messages = drop_orphaned_tool_messages(&messages);
423 let openai_messages: Vec<OpenAiMessage> =
424 messages.iter().map(Self::convert_message).collect();
425
426 let tools = if config.tools.is_empty() {
427 None
428 } else {
429 Some(Self::convert_tools(&config.tools))
430 };
431
432 let metadata = if config.metadata.is_empty() {
434 None
435 } else {
436 Some(config.metadata.clone())
437 };
438
439 let request = OpenAiRequest {
440 model: config.model.clone(),
441 messages: openai_messages,
442 temperature: config.temperature,
443 max_tokens: config.max_tokens,
444 stream: true,
445 stream_options: Some(OpenAiStreamOptions {
446 include_usage: true,
447 }),
448 tools,
449 parallel_tool_calls: config
450 .resolved_parallel_tool_calls(self.supports_parallel_tool_calls(&config.model)),
451 reasoning_effort: config
453 .reasoning_effort
454 .as_ref()
455 .filter(|e| !e.eq_ignore_ascii_case("none"))
456 .cloned(),
457 service_tier: config.speed.clone(),
458 verbosity: config.verbosity.clone(),
459 metadata,
460 };
461
462 let api_url = endpoint.url("chat/completions").ok_or_else(|| {
468 AgentLoopError::Configuration(
469 "OpenAI Chat Completions provider has no base URL".to_string(),
470 )
471 })?;
472 let (event_stream, retry_metadata) =
473 connect_sse_with_reconnect(&self.retry_config, "OpenAIProtocolDriver", |attempts| {
474 self.send_chat_completion_request(
475 endpoint,
476 &api_url,
477 &request,
478 &config.model,
479 attempts,
480 )
481 })
482 .await?;
483
484 let model = config.model.clone();
485 let total_tokens = Arc::new(Mutex::new(0u32));
486 let prompt_tokens = Arc::new(Mutex::new(0u32));
487 let cache_read_tokens = Arc::new(Mutex::new(Option::<u32>::None));
488 let provider_cost_usd = Arc::new(Mutex::new(Option::<f64>::None));
491 let accumulated_tool_calls = Arc::new(Mutex::new(StreamToolCallAccumulator::new()));
492 let finish_reason = Arc::new(Mutex::new(Option::<String>::None));
493 let response_id = Arc::new(Mutex::new(Option::<String>::None));
496 let shared_retry_metadata = if retry_metadata.had_retries() {
498 Some(Arc::new(retry_metadata))
499 } else {
500 None
501 };
502
503 let converted_stream: LlmResponseStream = Box::pin(
507 event_stream
508 .then(move |result| {
509 let model = model.clone();
510 let total_tokens = Arc::clone(&total_tokens);
511 let prompt_tokens = Arc::clone(&prompt_tokens);
512 let cache_read_tokens = Arc::clone(&cache_read_tokens);
513 let provider_cost_usd = Arc::clone(&provider_cost_usd);
514 let accumulated_tool_calls = Arc::clone(&accumulated_tool_calls);
515 let finish_reason = Arc::clone(&finish_reason);
516 let response_id = Arc::clone(&response_id);
517 let retry_metadata_for_done = shared_retry_metadata.clone();
518
519 async move {
520 let event = match result {
521 Ok(event) => event,
522 Err(e) => {
523 return vec![Ok(LlmStreamEvent::Error(
524 format!("Stream error: {}", e).into(),
525 ))];
526 }
527 };
528
529 if event.data == "[DONE]" {
530 let output_tokens = *total_tokens.lock().unwrap();
531 let input_tokens = *prompt_tokens.lock().unwrap();
532 let cached = *cache_read_tokens.lock().unwrap();
533 let cost = *provider_cost_usd.lock().unwrap();
534 let resp_id = response_id.lock().unwrap().clone();
535 let mut reason = finish_reason.lock().unwrap().clone();
536
537 let mut events = Vec::new();
538
539 {
548 let mut acc = accumulated_tool_calls.lock().unwrap();
549 if let Some(event) =
550 take_pending_tool_calls(&mut acc, reason.as_deref())
551 {
552 events.push(Ok(event));
553 reason.get_or_insert_with(|| "tool_calls".to_string());
554 }
555 }
556
557 events.push(Ok(LlmStreamEvent::Done(Box::new(
558 LlmCompletionMetadata {
559 total_tokens: Some(input_tokens + output_tokens),
562 prompt_tokens: Some(disjoint_prompt_tokens(
563 input_tokens,
564 cached,
565 )),
566 completion_tokens: Some(output_tokens),
567 cache_read_tokens: cached,
568 cache_creation_tokens: None,
569 provider_cost_usd: cost,
570 model: Some(model),
571 finish_reason: reason.or_else(|| Some("stop".to_string())),
572 retry_metadata: retry_metadata_for_done
573 .map(|arc| (*arc).clone()),
574 response_id: resp_id,
575 phase: None,
576 },
577 ))));
578
579 return events;
580 }
581
582 match serde_json::from_str::<OpenAiStreamChunk>(&event.data) {
583 Ok(chunk) => {
584 if let Some(id) = &chunk.id {
589 let mut rid = response_id.lock().unwrap();
590 if rid.is_none() {
591 *rid = Some(id.clone());
592 }
593 }
594
595 if let Some(usage) = &chunk.usage {
597 if let Some(pt) = usage.prompt_tokens {
598 *prompt_tokens.lock().unwrap() = pt;
599 }
600 if let Some(ct) = usage.completion_tokens {
601 *total_tokens.lock().unwrap() = ct;
602 }
603 if let Some(details) = &usage.prompt_tokens_details
605 && details.cached_tokens.is_some()
606 {
607 *cache_read_tokens.lock().unwrap() = details.cached_tokens;
608 }
609 if usage.cost.is_some() {
612 *provider_cost_usd.lock().unwrap() = usage.cost;
613 }
614 }
615
616 if let Some(choice) = chunk.choices.first() {
617 let mut tt = total_tokens.lock().unwrap();
618 let mut acc = accumulated_tool_calls.lock().unwrap();
619 let mut fr = finish_reason.lock().unwrap();
620 let stream_event =
621 process_stream_choice(choice, &mut tt, &mut acc, &mut fr);
622 return vec![Ok(stream_event)];
623 }
624 vec![Ok(LlmStreamEvent::TextDelta(String::new()))]
625 }
626 Err(e) => vec![Ok(LlmStreamEvent::Error(
627 format!("Failed to parse chunk: {}", e).into(),
628 ))],
629 }
630 }
631 })
632 .flat_map(futures::stream::iter),
633 );
634
635 Ok(converted_stream)
636 }
637
638 fn supports_parallel_tool_calls(&self, _model: &str) -> bool {
642 true
643 }
644}
645
646impl std::fmt::Debug for OpenAIProtocolChatDriver {
647 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
648 f.debug_struct("OpenAIProtocolChatDriver")
649 .field("protocol", &"openai_chat_completions")
650 .finish()
651 }
652}
653
654pub fn is_openai_model_not_found(status: reqwest::StatusCode, error_text: &str) -> bool {
665 let error_lower = error_text.to_lowercase();
666
667 if status == reqwest::StatusCode::NOT_FOUND
669 || status == reqwest::StatusCode::BAD_REQUEST
670 || status == reqwest::StatusCode::FORBIDDEN
671 {
672 if error_lower.contains("model_not_found") {
674 return true;
675 }
676 }
677
678 if status == reqwest::StatusCode::NOT_FOUND {
680 if error_lower.contains("does not exist") {
681 return true;
682 }
683 if error_lower.contains("model") && error_lower.contains("not found") {
684 return true;
685 }
686 }
687
688 false
689}
690
691pub fn is_openai_request_too_large(status: reqwest::StatusCode, error_text: &str) -> bool {
698 let error_lower = error_text.to_lowercase();
699
700 if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
702 if error_lower.contains("request too large") {
704 return true;
705 }
706 if error_lower.contains("tokens") && error_lower.contains("limit") {
708 return true;
709 }
710 }
711
712 if status == reqwest::StatusCode::BAD_REQUEST {
714 if error_lower.contains("context_length_exceeded") {
716 return true;
717 }
718 if error_lower.contains("maximum context length") {
720 return true;
721 }
722 }
723
724 if error_lower.contains("tokens must be reduced")
726 || error_lower.contains("reduce the length")
727 || error_lower.contains("input is too long")
728 {
729 return true;
730 }
731
732 false
733}
734
735#[derive(Debug, Serialize)]
740struct OpenAiRequest {
741 model: String,
742 messages: Vec<OpenAiMessage>,
743 #[serde(skip_serializing_if = "Option::is_none")]
744 temperature: Option<f32>,
745 #[serde(skip_serializing_if = "Option::is_none")]
746 max_tokens: Option<u32>,
747 stream: bool,
748 #[serde(skip_serializing_if = "Option::is_none")]
750 stream_options: Option<OpenAiStreamOptions>,
751 #[serde(skip_serializing_if = "Option::is_none")]
752 tools: Option<Vec<OpenAiTool>>,
753 #[serde(skip_serializing_if = "Option::is_none")]
756 parallel_tool_calls: Option<bool>,
757 #[serde(skip_serializing_if = "Option::is_none")]
758 reasoning_effort: Option<String>,
759 #[serde(skip_serializing_if = "Option::is_none")]
762 service_tier: Option<String>,
763 #[serde(skip_serializing_if = "Option::is_none")]
767 verbosity: Option<String>,
768 #[serde(skip_serializing_if = "Option::is_none")]
771 metadata: Option<std::collections::HashMap<String, String>>,
772}
773
774#[derive(Debug, Serialize)]
775struct OpenAiStreamOptions {
776 include_usage: bool,
777}
778
779#[derive(Debug, Serialize, Deserialize)]
780#[serde(untagged)]
781enum OpenAiContent {
782 Text(String),
783 Parts(Vec<OpenAiContentPart>),
784}
785
786#[derive(Debug, Serialize, Deserialize)]
787#[serde(untagged)]
788enum OpenAiContentPart {
789 Text {
790 r#type: String,
791 text: String,
792 },
793 ImageUrl {
794 r#type: String,
795 image_url: OpenAiImageUrl,
796 },
797 InputAudio {
798 r#type: String,
799 input_audio: OpenAiInputAudio,
800 },
801}
802
803#[derive(Debug, Serialize, Deserialize)]
804struct OpenAiImageUrl {
805 url: String,
806}
807
808#[derive(Debug, Serialize, Deserialize)]
809struct OpenAiInputAudio {
810 data: String,
811 format: String,
812}
813
814#[derive(Debug, Serialize, Deserialize)]
815struct OpenAiMessage {
816 role: String,
817 #[serde(skip_serializing_if = "Option::is_none")]
818 content: Option<OpenAiContent>,
819 #[serde(skip_serializing_if = "Option::is_none")]
820 tool_calls: Option<Vec<OpenAiToolCall>>,
821 #[serde(skip_serializing_if = "Option::is_none")]
822 tool_call_id: Option<String>,
823}
824
825#[derive(Debug, Serialize, Deserialize)]
826struct OpenAiTool {
827 r#type: String,
828 function: OpenAiFunction,
829}
830
831#[derive(Debug, Serialize, Deserialize)]
832struct OpenAiFunction {
833 name: String,
834 description: String,
835 parameters: Value,
836}
837
838#[derive(Debug, Serialize, Deserialize)]
839struct OpenAiToolCall {
840 id: String,
841 r#type: String,
842 function: OpenAiFunctionCall,
843}
844
845#[derive(Debug, Serialize, Deserialize)]
846struct OpenAiFunctionCall {
847 name: String,
848 arguments: String,
849}
850
851#[derive(Debug, Deserialize)]
852#[allow(dead_code)] struct OpenAiStreamChunk {
854 #[serde(default)]
856 id: Option<String>,
857 #[serde(default)]
859 model: Option<String>,
860 choices: Vec<OpenAiStreamChoice>,
861 #[serde(default)]
862 usage: Option<OpenAiUsage>,
863}
864
865#[derive(Debug, Deserialize)]
866struct OpenAiUsage {
867 prompt_tokens: Option<u32>,
868 completion_tokens: Option<u32>,
869 #[serde(default)]
871 prompt_tokens_details: Option<OpenAiPromptTokensDetails>,
872 #[serde(default)]
875 cost: Option<f64>,
876}
877
878#[derive(Debug, Deserialize, Default)]
879struct OpenAiPromptTokensDetails {
880 #[serde(default)]
882 cached_tokens: Option<u32>,
883}
884
885#[derive(Debug, Deserialize)]
886struct OpenAiStreamChoice {
887 delta: OpenAiDelta,
888 #[serde(default)]
889 finish_reason: Option<String>,
890}
891
892#[derive(Debug, Deserialize)]
893struct OpenAiDelta {
894 #[serde(default)]
895 content: Option<String>,
896 #[serde(default)]
897 tool_calls: Option<Vec<OpenAiStreamToolCall>>,
898}
899
900#[derive(Debug, Deserialize)]
901struct OpenAiStreamToolCall {
902 index: u32,
903 id: Option<String>,
904 function: Option<OpenAiStreamFunction>,
905}
906
907#[derive(Debug, Deserialize)]
908struct OpenAiStreamFunction {
909 name: Option<String>,
910 arguments: Option<String>,
911}
912
913fn take_pending_tool_calls(
924 accumulated_tool_calls: &mut StreamToolCallAccumulator,
925 finish_reason: Option<&str>,
926) -> Option<LlmStreamEvent> {
927 if accumulated_tool_calls.is_empty() {
928 return None;
929 }
930
931 if !matches!(finish_reason, None | Some("tool_calls")) {
934 let _ = accumulated_tool_calls.take_finalized();
935 return None;
936 }
937
938 let calls = accumulated_tool_calls.take_pending_strict();
939 if calls.is_empty() {
940 None
941 } else {
942 Some(LlmStreamEvent::ToolCalls(calls))
943 }
944}
945
946fn process_stream_choice(
956 choice: &OpenAiStreamChoice,
957 total_tokens: &mut u32,
958 accumulated_tool_calls: &mut StreamToolCallAccumulator,
959 finish_reason: &mut Option<String>,
960) -> LlmStreamEvent {
961 if let Some(tool_calls) = &choice.delta.tool_calls {
965 for tc in tool_calls {
966 accumulated_tool_calls.apply_indexed_delta(
967 tc.index,
968 tc.id.as_deref(),
969 tc.function.as_ref().and_then(|f| f.name.as_deref()),
970 tc.function.as_ref().and_then(|f| f.arguments.as_deref()),
971 );
972 }
973 return LlmStreamEvent::TextDelta(String::new());
974 }
975
976 if let Some(content) = &choice.delta.content
979 && !content.is_empty()
980 {
981 *total_tokens += 1;
982 return LlmStreamEvent::TextDelta(content.clone());
983 }
984
985 if let Some(fr) = &choice.finish_reason {
989 *finish_reason = Some(fr.clone());
990
991 if fr == "tool_calls" && !accumulated_tool_calls.is_empty() {
992 return LlmStreamEvent::ToolCalls(accumulated_tool_calls.take_finalized());
993 }
994 }
995
996 LlmStreamEvent::TextDelta(String::new())
997}
998
999#[cfg(test)]
1004mod tests {
1005 use super::*;
1006 use serde_json::json;
1007
1008 #[test]
1009 fn test_convert_message_preserves_multiple_system_messages() {
1010 let messages = [
1015 LlmMessage::text(LlmMessageRole::System, "A"),
1016 LlmMessage::text(LlmMessageRole::User, "hi"),
1017 LlmMessage::text(LlmMessageRole::System, "B"),
1018 ];
1019 let converted: Vec<OpenAiMessage> = messages
1020 .iter()
1021 .map(OpenAIProtocolChatDriver::convert_message)
1022 .collect();
1023 let system_texts: Vec<String> = converted
1024 .iter()
1025 .filter(|m| m.role == "system")
1026 .filter_map(|m| match &m.content {
1027 Some(OpenAiContent::Text(t)) => Some(t.clone()),
1028 _ => None,
1029 })
1030 .collect();
1031 assert_eq!(system_texts, vec!["A".to_string(), "B".to_string()]);
1032 }
1033
1034 #[test]
1035 fn test_driver_is_wire_only() {
1036 let driver = OpenAIProtocolChatDriver::new();
1037 assert!(format!("{:?}", driver).contains("OpenAIProtocolChatDriver"));
1038 }
1039
1040 #[test]
1041 fn test_is_azure_openai_api_url() {
1042 assert!(is_azure_openai_api_url(
1043 "https://example.openai.azure.com/openai/v1/chat/completions"
1044 ));
1045 assert!(is_azure_openai_api_url(
1046 "https://example.services.ai.azure.com/openai/v1/responses"
1047 ));
1048 assert!(!is_azure_openai_api_url(
1049 "https://api.openai.com/v1/chat/completions"
1050 ));
1051 }
1052
1053 #[test]
1054 fn test_request_includes_stream_options_for_usage() {
1055 let request = OpenAiRequest {
1058 verbosity: None,
1059 service_tier: None,
1060 model: "gpt-4o".to_string(),
1061 messages: vec![OpenAiMessage {
1062 role: "user".to_string(),
1063 content: Some(OpenAiContent::Text("Hello".to_string())),
1064 tool_calls: None,
1065 tool_call_id: None,
1066 }],
1067 temperature: None,
1068 max_tokens: None,
1069 stream: true,
1070 stream_options: Some(OpenAiStreamOptions {
1071 include_usage: true,
1072 }),
1073 tools: None,
1074 parallel_tool_calls: None,
1075 reasoning_effort: None,
1076 metadata: None,
1077 };
1078
1079 let json = serde_json::to_value(&request).unwrap();
1080 assert_eq!(json["stream"], true);
1081 assert_eq!(json["stream_options"]["include_usage"], true);
1082 }
1083
1084 #[test]
1085 fn test_request_includes_metadata() {
1086 let mut metadata = std::collections::HashMap::new();
1088 metadata.insert("session_id".to_string(), "session_abc123".to_string());
1089 metadata.insert("agent_id".to_string(), "agent_xyz789".to_string());
1090
1091 let request = OpenAiRequest {
1092 verbosity: None,
1093 service_tier: None,
1094 model: "gpt-4o".to_string(),
1095 messages: vec![OpenAiMessage {
1096 role: "user".to_string(),
1097 content: Some(OpenAiContent::Text("Hello".to_string())),
1098 tool_calls: None,
1099 tool_call_id: None,
1100 }],
1101 temperature: None,
1102 max_tokens: None,
1103 stream: true,
1104 stream_options: None,
1105 tools: None,
1106 parallel_tool_calls: None,
1107 reasoning_effort: None,
1108 metadata: Some(metadata),
1109 };
1110
1111 let json = serde_json::to_value(&request).unwrap();
1112 assert_eq!(json["metadata"]["session_id"], "session_abc123");
1113 assert_eq!(json["metadata"]["agent_id"], "agent_xyz789");
1114 }
1115
1116 #[test]
1117 fn test_usage_chunk_parsing() {
1118 let usage_chunk = r#"{
1121 "id": "chatcmpl-123",
1122 "object": "chat.completion.chunk",
1123 "created": 1234567890,
1124 "model": "gpt-4o",
1125 "choices": [],
1126 "usage": {
1127 "prompt_tokens": 150,
1128 "completion_tokens": 42,
1129 "total_tokens": 192
1130 }
1131 }"#;
1132
1133 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1134 assert!(chunk.usage.is_some());
1135 let usage = chunk.usage.unwrap();
1136 assert_eq!(usage.prompt_tokens, Some(150));
1137 assert_eq!(usage.completion_tokens, Some(42));
1138 }
1139
1140 #[test]
1141 fn test_usage_chunk_with_cached_tokens() {
1142 let usage_chunk = r#"{
1144 "id": "chatcmpl-123",
1145 "choices": [],
1146 "usage": {
1147 "prompt_tokens": 150,
1148 "completion_tokens": 42,
1149 "prompt_tokens_details": {
1150 "cached_tokens": 100
1151 }
1152 }
1153 }"#;
1154
1155 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1156 let usage = chunk.usage.unwrap();
1157 assert_eq!(usage.prompt_tokens, Some(150));
1158 assert_eq!(usage.completion_tokens, Some(42));
1159 assert!(usage.prompt_tokens_details.is_some());
1160 assert_eq!(
1161 usage.prompt_tokens_details.unwrap().cached_tokens,
1162 Some(100)
1163 );
1164 }
1165
1166 #[test]
1167 fn test_usage_chunk_with_openrouter_cost() {
1168 let usage_chunk = r#"{
1170 "id": "gen-123",
1171 "choices": [],
1172 "usage": {
1173 "prompt_tokens": 194,
1174 "completion_tokens": 2,
1175 "total_tokens": 196,
1176 "cost": 0.00095
1177 }
1178 }"#;
1179
1180 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1181 let usage = chunk.usage.unwrap();
1182 assert_eq!(usage.cost, Some(0.00095));
1183 }
1184
1185 #[test]
1186 fn test_usage_chunk_without_cost_defaults_none() {
1187 let usage_chunk = r#"{
1189 "id": "chatcmpl-123",
1190 "choices": [],
1191 "usage": { "prompt_tokens": 10, "completion_tokens": 5 }
1192 }"#;
1193
1194 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1195 assert_eq!(chunk.usage.unwrap().cost, None);
1196 }
1197
1198 #[test]
1199 fn test_chunk_id_is_captured() {
1200 let chunk_with_id: OpenAiStreamChunk =
1201 serde_json::from_str(r#"{"id":"gen-abc123","choices":[]}"#).unwrap();
1202 assert_eq!(chunk_with_id.id.as_deref(), Some("gen-abc123"));
1203
1204 let chunk_no_id: OpenAiStreamChunk = serde_json::from_str(r#"{"choices":[]}"#).unwrap();
1205 assert!(chunk_no_id.id.is_none());
1206 }
1207
1208 #[test]
1209 fn test_finish_reason_chunk_parsing() {
1210 let finish_chunk = r#"{
1212 "id": "chatcmpl-123",
1213 "choices": [{
1214 "index": 0,
1215 "delta": {},
1216 "finish_reason": "stop"
1217 }]
1218 }"#;
1219
1220 let chunk: OpenAiStreamChunk = serde_json::from_str(finish_chunk).unwrap();
1221 assert!(chunk.usage.is_none()); assert_eq!(chunk.choices.len(), 1);
1223 assert_eq!(chunk.choices[0].finish_reason, Some("stop".to_string()));
1224 }
1225
1226 #[test]
1231 fn test_is_openai_request_too_large_429_request_too_large() {
1232 let error = r#"{"error":{"message":"Request too large for gpt-4o in organization org-xxx on tokens per min (TPM): Limit 500000, Requested 538772."}}"#;
1233 assert!(is_openai_request_too_large(
1234 reqwest::StatusCode::TOO_MANY_REQUESTS,
1235 error
1236 ));
1237 }
1238
1239 #[test]
1240 fn test_is_openai_request_too_large_429_token_limit() {
1241 let error =
1242 r#"{"error":{"message":"tokens per min (TPM): Limit 500000, Requested 600000"}}"#;
1243 assert!(is_openai_request_too_large(
1244 reqwest::StatusCode::TOO_MANY_REQUESTS,
1245 error
1246 ));
1247 }
1248
1249 #[test]
1250 fn test_is_openai_request_too_large_400_context_length() {
1251 let error = r#"{"error":{"code":"context_length_exceeded","message":"This model's maximum context length is 128000 tokens."}}"#;
1252 assert!(is_openai_request_too_large(
1253 reqwest::StatusCode::BAD_REQUEST,
1254 error
1255 ));
1256 }
1257
1258 #[test]
1259 fn test_is_openai_request_too_large_400_max_context() {
1260 let error =
1261 r#"{"error":{"message":"This model's maximum context length is 128000 tokens"}}"#;
1262 assert!(is_openai_request_too_large(
1263 reqwest::StatusCode::BAD_REQUEST,
1264 error
1265 ));
1266 }
1267
1268 #[test]
1269 fn test_is_openai_request_too_large_tokens_must_be_reduced() {
1270 let error = r#"{"error":{"message":"The input or output tokens must be reduced"}}"#;
1271 assert!(is_openai_request_too_large(
1272 reqwest::StatusCode::BAD_REQUEST,
1273 error
1274 ));
1275 }
1276
1277 #[test]
1278 fn test_is_openai_request_too_large_false_for_other_errors() {
1279 let error = r#"{"error":{"message":"Rate limit exceeded: too many requests per minute"}}"#;
1281 assert!(!is_openai_request_too_large(
1282 reqwest::StatusCode::TOO_MANY_REQUESTS,
1283 error
1284 ));
1285
1286 let error = r#"{"error":{"message":"Internal server error"}}"#;
1288 assert!(!is_openai_request_too_large(
1289 reqwest::StatusCode::INTERNAL_SERVER_ERROR,
1290 error
1291 ));
1292
1293 let error = r#"{"error":{"message":"Invalid request"}}"#;
1295 assert!(!is_openai_request_too_large(
1296 reqwest::StatusCode::BAD_REQUEST,
1297 error
1298 ));
1299 }
1300
1301 #[test]
1306 fn test_is_openai_model_not_found_real_error() {
1307 let error = r#"{"error":{"code":"model_not_found","message":"The model 'gpt-99' does not exist or you do not have access to it.","type":"invalid_request_error","param":null}}"#;
1309 assert!(is_openai_model_not_found(
1310 reqwest::StatusCode::NOT_FOUND,
1311 error
1312 ));
1313 }
1314
1315 #[test]
1316 fn test_is_openai_model_not_found_does_not_exist() {
1317 let error = r#"{"error":{"message":"The model 'fake-model' does not exist"}}"#;
1318 assert!(is_openai_model_not_found(
1319 reqwest::StatusCode::NOT_FOUND,
1320 error
1321 ));
1322 }
1323
1324 #[test]
1325 fn test_is_openai_model_not_found_generic_not_found() {
1326 let error = r#"{"error":{"message":"Model not found"}}"#;
1327 assert!(is_openai_model_not_found(
1328 reqwest::StatusCode::NOT_FOUND,
1329 error
1330 ));
1331 }
1332
1333 #[test]
1334 fn test_is_openai_model_not_found_400_with_model_not_found_code() {
1335 let error = r#"{"error":{"code":"model_not_found","message":"The requested model 'gpt-99' does not exist.","type":"invalid_request_error","param":"model"}}"#;
1337 assert!(is_openai_model_not_found(
1338 reqwest::StatusCode::BAD_REQUEST,
1339 error
1340 ));
1341 }
1342
1343 #[test]
1344 fn test_is_openai_model_not_found_false_for_non_model_error() {
1345 let error = r#"{"error":{"code":"invalid_request","message":"Some other error"}}"#;
1347 assert!(!is_openai_model_not_found(
1348 reqwest::StatusCode::BAD_REQUEST,
1349 error
1350 ));
1351 }
1352
1353 #[test]
1354 fn test_is_openai_model_not_found_false_for_other_404() {
1355 let error = r#"{"error":{"message":"Endpoint not found"}}"#;
1357 assert!(!is_openai_model_not_found(
1358 reqwest::StatusCode::NOT_FOUND,
1359 error
1360 ));
1361 }
1362
1363 #[test]
1364 fn test_is_openai_model_not_found_403_tier_gated_model() {
1365 let error = r#"{"error":{"code":"model_not_found","message":"The model 'gpt-5.4-mini' does not exist or you do not have access to it.","type":"invalid_request_error","param":null}}"#;
1368 assert!(is_openai_model_not_found(
1369 reqwest::StatusCode::FORBIDDEN,
1370 error
1371 ));
1372 }
1373
1374 #[test]
1375 fn test_is_openai_model_not_found_403_plain_auth_error_is_not_model_not_found() {
1376 let error = r#"{"error":{"message":"Invalid authentication credentials","type":"authentication_error"}}"#;
1379 assert!(!is_openai_model_not_found(
1380 reqwest::StatusCode::FORBIDDEN,
1381 error
1382 ));
1383 }
1384
1385 #[test]
1390 fn test_reasoning_effort_none_is_omitted() {
1391 let request = OpenAiRequest {
1394 verbosity: None,
1395 service_tier: None,
1396 model: "gpt-4o-mini".to_string(),
1397 messages: vec![OpenAiMessage {
1398 role: "user".to_string(),
1399 content: Some(OpenAiContent::Text("Hello".to_string())),
1400 tool_calls: None,
1401 tool_call_id: None,
1402 }],
1403 temperature: None,
1404 max_tokens: None,
1405 stream: true,
1406 stream_options: None,
1407 tools: None,
1408 parallel_tool_calls: None,
1409 reasoning_effort: Some("none".to_string())
1410 .as_ref()
1411 .filter(|e| !e.eq_ignore_ascii_case("none"))
1412 .cloned(),
1413 metadata: None,
1414 };
1415
1416 let json = serde_json::to_value(&request).unwrap();
1417 assert!(
1418 json.get("reasoning_effort").is_none(),
1419 "reasoning_effort should be omitted when effort is 'none'"
1420 );
1421 }
1422
1423 #[test]
1424 fn test_reasoning_effort_high_is_included() {
1425 let request = OpenAiRequest {
1426 verbosity: None,
1427 service_tier: None,
1428 model: "o3-mini".to_string(),
1429 messages: vec![OpenAiMessage {
1430 role: "user".to_string(),
1431 content: Some(OpenAiContent::Text("Hello".to_string())),
1432 tool_calls: None,
1433 tool_call_id: None,
1434 }],
1435 temperature: None,
1436 max_tokens: None,
1437 stream: true,
1438 stream_options: None,
1439 tools: None,
1440 parallel_tool_calls: None,
1441 reasoning_effort: Some("high".to_string())
1442 .as_ref()
1443 .filter(|e| !e.eq_ignore_ascii_case("none"))
1444 .cloned(),
1445 metadata: None,
1446 };
1447
1448 let json = serde_json::to_value(&request).unwrap();
1449 assert_eq!(json["reasoning_effort"], "high");
1450 }
1451
1452 #[test]
1456 fn test_request_serializes_parallel_tool_calls() {
1457 fn build(flag: Option<bool>) -> serde_json::Value {
1458 let request = OpenAiRequest {
1459 verbosity: None,
1460 service_tier: None,
1461 model: "gpt-4o-mini".to_string(),
1462 messages: vec![OpenAiMessage {
1463 role: "user".to_string(),
1464 content: Some(OpenAiContent::Text("Hello".to_string())),
1465 tool_calls: None,
1466 tool_call_id: None,
1467 }],
1468 temperature: None,
1469 max_tokens: None,
1470 stream: true,
1471 stream_options: None,
1472 tools: None,
1473 parallel_tool_calls: flag,
1474 reasoning_effort: None,
1475 metadata: None,
1476 };
1477 serde_json::to_value(&request).unwrap()
1478 }
1479
1480 assert!(build(None).get("parallel_tool_calls").is_none());
1482 assert_eq!(build(Some(true))["parallel_tool_calls"], true);
1484 assert_eq!(build(Some(false))["parallel_tool_calls"], false);
1485 }
1486
1487 #[test]
1490 fn test_request_serializes_service_tier() {
1491 fn build(tier: Option<&str>) -> serde_json::Value {
1492 let request = OpenAiRequest {
1493 service_tier: tier.map(str::to_string),
1494 verbosity: None,
1495 model: "gpt-4o-mini".to_string(),
1496 messages: vec![OpenAiMessage {
1497 role: "user".to_string(),
1498 content: Some(OpenAiContent::Text("Hello".to_string())),
1499 tool_calls: None,
1500 tool_call_id: None,
1501 }],
1502 temperature: None,
1503 max_tokens: None,
1504 stream: true,
1505 stream_options: None,
1506 tools: None,
1507 parallel_tool_calls: None,
1508 reasoning_effort: None,
1509 metadata: None,
1510 };
1511 serde_json::to_value(&request).unwrap()
1512 }
1513
1514 assert!(build(None).get("service_tier").is_none());
1515 assert_eq!(build(Some("flex"))["service_tier"], "flex");
1516 assert_eq!(build(Some("priority"))["service_tier"], "priority");
1517 }
1518
1519 #[test]
1522 fn test_request_serializes_verbosity() {
1523 fn build(verbosity: Option<&str>) -> serde_json::Value {
1524 let request = OpenAiRequest {
1525 service_tier: None,
1526 verbosity: verbosity.map(str::to_string),
1527 model: "gpt-5.6-sol".to_string(),
1528 messages: vec![OpenAiMessage {
1529 role: "user".to_string(),
1530 content: Some(OpenAiContent::Text("Hello".to_string())),
1531 tool_calls: None,
1532 tool_call_id: None,
1533 }],
1534 temperature: None,
1535 max_tokens: None,
1536 stream: true,
1537 stream_options: None,
1538 tools: None,
1539 parallel_tool_calls: None,
1540 reasoning_effort: None,
1541 metadata: None,
1542 };
1543 serde_json::to_value(&request).unwrap()
1544 }
1545
1546 assert!(build(None).get("verbosity").is_none());
1547 assert_eq!(build(Some("low"))["verbosity"], "low");
1548 assert_eq!(build(Some("high"))["verbosity"], "high");
1549 }
1550
1551 fn choice(json_str: &str) -> OpenAiStreamChoice {
1556 serde_json::from_str(json_str).unwrap()
1557 }
1558
1559 #[test]
1563 fn test_empty_content_finish_chunk_still_emits_tool_calls() {
1564 let mut total_tokens = 0u32;
1565 let mut acc = StreamToolCallAccumulator::new();
1566 let mut finish_reason: Option<String> = None;
1567
1568 let e = process_stream_choice(
1570 &choice(
1571 r#"{"delta":{"content":null,"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":""}}]},"finish_reason":null}"#,
1572 ),
1573 &mut total_tokens,
1574 &mut acc,
1575 &mut finish_reason,
1576 );
1577 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()));
1578
1579 let e = process_stream_choice(
1581 &choice(
1582 r#"{"delta":{"content":null,"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":\"Cargo.toml\"}"}}]},"finish_reason":null}"#,
1583 ),
1584 &mut total_tokens,
1585 &mut acc,
1586 &mut finish_reason,
1587 );
1588 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()));
1589
1590 let e = process_stream_choice(
1593 &choice(r#"{"delta":{"content":""},"finish_reason":"tool_calls"}"#),
1594 &mut total_tokens,
1595 &mut acc,
1596 &mut finish_reason,
1597 );
1598 match e {
1599 LlmStreamEvent::ToolCalls(calls) => {
1600 assert_eq!(calls.len(), 1);
1601 assert_eq!(calls[0].id, "call_1");
1602 assert_eq!(calls[0].name, "read_file");
1603 assert_eq!(calls[0].arguments, json!({"path": "Cargo.toml"}));
1604 }
1605 other => panic!("expected ToolCalls, got {:?}", other),
1606 }
1607 assert_eq!(finish_reason.as_deref(), Some("tool_calls"));
1608
1609 let e = process_stream_choice(
1612 &choice(r#"{"delta":{"content":""},"finish_reason":"tool_calls"}"#),
1613 &mut total_tokens,
1614 &mut acc,
1615 &mut finish_reason,
1616 );
1617 assert!(
1618 matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()),
1619 "tool calls must only be emitted once"
1620 );
1621 }
1622
1623 #[test]
1625 fn test_non_empty_content_is_emitted() {
1626 let mut total_tokens = 0u32;
1627 let mut acc = StreamToolCallAccumulator::new();
1628 let mut finish_reason: Option<String> = None;
1629
1630 let e = process_stream_choice(
1631 &choice(r#"{"delta":{"content":"hello"},"finish_reason":null}"#),
1632 &mut total_tokens,
1633 &mut acc,
1634 &mut finish_reason,
1635 );
1636 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s == "hello"));
1637 assert_eq!(total_tokens, 1);
1638 }
1639
1640 #[test]
1644 fn test_tool_call_arguments_accumulate_across_many_chunks() {
1645 let mut total_tokens = 0u32;
1646 let mut acc = StreamToolCallAccumulator::new();
1647 let mut finish_reason: Option<String> = None;
1648
1649 process_stream_choice(
1651 &choice(
1652 r#"{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"write_file","arguments":""}}]},"finish_reason":null}"#,
1653 ),
1654 &mut total_tokens,
1655 &mut acc,
1656 &mut finish_reason,
1657 );
1658
1659 let payload = r#"{"path":"a.rs","contents":"a fairly long contents value streamed one character at a time to exceed one hundred chunks","n":987654321}"#;
1660 assert!(payload.chars().count() > 100);
1661
1662 let mut expected = String::new();
1664 for ch in payload.chars() {
1665 let frag = ch.to_string();
1666 let chunk = json!({
1667 "delta": {"tool_calls": [{"index": 0, "function": {"arguments": frag}}]},
1668 "finish_reason": null
1669 })
1670 .to_string();
1671 process_stream_choice(
1672 &choice(&chunk),
1673 &mut total_tokens,
1674 &mut acc,
1675 &mut finish_reason,
1676 );
1677 expected.push_str(&frag);
1678 }
1679
1680 let e = process_stream_choice(
1686 &choice(r#"{"delta":{},"finish_reason":"tool_calls"}"#),
1687 &mut total_tokens,
1688 &mut acc,
1689 &mut finish_reason,
1690 );
1691 match e {
1692 LlmStreamEvent::ToolCalls(calls) => {
1693 assert_eq!(calls.len(), 1);
1694 assert_eq!(calls[0].id, "call_1");
1695 assert_eq!(
1696 calls[0].arguments,
1697 serde_json::from_str::<serde_json::Value>(payload).unwrap()
1698 );
1699 }
1700 other => panic!("expected ToolCalls, got {:?}", other),
1701 }
1702 }
1703
1704 #[test]
1707 fn test_finish_chunk_without_content_emits_tool_calls() {
1708 let mut total_tokens = 0u32;
1709 let mut acc = StreamToolCallAccumulator::new();
1710 let mut finish_reason: Option<String> = None;
1711
1712 process_stream_choice(
1713 &choice(
1714 r#"{"delta":{"tool_calls":[{"index":0,"id":"call_9","function":{"name":"list_dir","arguments":"{}"}}]},"finish_reason":null}"#,
1715 ),
1716 &mut total_tokens,
1717 &mut acc,
1718 &mut finish_reason,
1719 );
1720
1721 let e = process_stream_choice(
1722 &choice(r#"{"delta":{},"finish_reason":"tool_calls"}"#),
1723 &mut total_tokens,
1724 &mut acc,
1725 &mut finish_reason,
1726 );
1727 match e {
1728 LlmStreamEvent::ToolCalls(calls) => {
1729 assert_eq!(calls.len(), 1);
1730 assert_eq!(calls[0].name, "list_dir");
1731 }
1732 other => panic!("expected ToolCalls, got {:?}", other),
1733 }
1734 }
1735
1736 fn seeded_acc(id: &str, name: &str, arguments: &str) -> StreamToolCallAccumulator {
1740 let mut acc = StreamToolCallAccumulator::new();
1741 acc.apply_indexed_delta(0, Some(id), Some(name), Some(arguments));
1742 acc
1743 }
1744
1745 #[test]
1749 fn test_take_pending_tool_calls_flushes_then_drains_without_finish_reason() {
1750 let mut acc = seeded_acc("call_1", "read_file", r#"{"path":"Cargo.toml"}"#);
1751
1752 match take_pending_tool_calls(&mut acc, None) {
1753 Some(LlmStreamEvent::ToolCalls(calls)) => {
1754 assert_eq!(calls.len(), 1);
1755 assert_eq!(calls[0].name, "read_file");
1756 assert_eq!(calls[0].arguments, json!({"path": "Cargo.toml"}));
1757 }
1758 other => panic!("expected ToolCalls, got {:?}", other),
1759 }
1760 assert!(acc.is_empty(), "accumulator must be drained after flush");
1761 assert!(take_pending_tool_calls(&mut acc, None).is_none());
1762 }
1763
1764 #[test]
1765 fn test_take_pending_tool_calls_discards_non_tool_finish_reason() {
1766 let mut acc = seeded_acc("call_cut", "read_file", r#"{"path":"#);
1767
1768 assert!(take_pending_tool_calls(&mut acc, Some("length")).is_none());
1769 assert!(
1770 acc.is_empty(),
1771 "discarded unsafe fallback calls must still drain the accumulator"
1772 );
1773 }
1774
1775 #[test]
1776 fn test_take_pending_tool_calls_rejects_malformed_fallback_arguments() {
1777 let mut acc = seeded_acc("call_cut", "read_file", r#"{"path":"#);
1778
1779 assert!(take_pending_tool_calls(&mut acc, None).is_none());
1780 assert!(
1781 acc.is_empty(),
1782 "malformed fallback calls must be drained instead of re-emitted"
1783 );
1784 }
1785
1786 #[test]
1787 fn test_non_tool_finish_reason_leaves_pending_calls_for_done_discard() {
1788 let mut total_tokens = 0u32;
1789 let mut acc = StreamToolCallAccumulator::new();
1790 let mut finish_reason: Option<String> = None;
1791
1792 process_stream_choice(
1793 &choice(
1794 r#"{"delta":{"tool_calls":[{"index":0,"id":"call_cut","function":{"name":"read_file","arguments":"{\"path\":"}}]},"finish_reason":null}"#,
1795 ),
1796 &mut total_tokens,
1797 &mut acc,
1798 &mut finish_reason,
1799 );
1800
1801 let e = process_stream_choice(
1802 &choice(r#"{"delta":{},"finish_reason":"length"}"#),
1803 &mut total_tokens,
1804 &mut acc,
1805 &mut finish_reason,
1806 );
1807
1808 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()));
1809 assert_eq!(finish_reason.as_deref(), Some("length"));
1810 assert!(take_pending_tool_calls(&mut acc, finish_reason.as_deref()).is_none());
1811 assert!(acc.is_empty());
1812 }
1813
1814 #[test]
1815 fn drop_orphaned_tool_messages_removes_unmatched_tool_results() {
1816 use crate::driver_registry::LlmMessageContent;
1817
1818 let messages = vec![
1819 LlmMessage::text(LlmMessageRole::User, "hello"),
1820 LlmMessage {
1821 role: LlmMessageRole::Tool,
1822 content: LlmMessageContent::Text("result".to_string()),
1823 tool_calls: None,
1824 tool_call_id: Some("call_trimmed".to_string()),
1825 phase: None,
1826 thinking: None,
1827 thinking_signature: None,
1828 },
1829 ];
1830 let filtered = drop_orphaned_tool_messages(&messages);
1831 assert_eq!(filtered.len(), 1);
1832 assert_eq!(filtered[0].role, LlmMessageRole::User);
1833 }
1834
1835 #[test]
1836 fn drop_orphaned_tool_messages_keeps_matched_tool_results() {
1837 use crate::driver_registry::LlmMessageContent;
1838 use crate::tool_types::ToolCall;
1839
1840 let messages = vec![
1841 LlmMessage {
1842 role: LlmMessageRole::Assistant,
1843 content: LlmMessageContent::Text(String::new()),
1844 tool_calls: Some(vec![ToolCall {
1845 id: "call_1".to_string(),
1846 name: "read_file".to_string(),
1847 arguments: json!({}),
1848 }]),
1849 tool_call_id: None,
1850 phase: None,
1851 thinking: None,
1852 thinking_signature: None,
1853 },
1854 LlmMessage {
1855 role: LlmMessageRole::Tool,
1856 content: LlmMessageContent::Text("file content".to_string()),
1857 tool_calls: None,
1858 tool_call_id: Some("call_1".to_string()),
1859 phase: None,
1860 thinking: None,
1861 thinking_signature: None,
1862 },
1863 ];
1864 let filtered = drop_orphaned_tool_messages(&messages);
1865 assert_eq!(filtered.len(), 2);
1866 }
1867}