1use async_trait::async_trait;
18use futures::StreamExt;
19use reqwest::{Client, RequestBuilder, Url};
20use serde::{Deserialize, Serialize};
21use serde_json::Value;
22use std::borrow::Cow;
23use std::sync::{Arc, Mutex};
24
25use crate::driver_registry::{
26 ChatDriver, LlmCallConfig, LlmCompletionMetadata, LlmContentPart, LlmMessage,
27 LlmMessageContent, LlmMessageRole, LlmResponseStream, LlmStreamEvent, disjoint_prompt_tokens,
28};
29use crate::error::{AgentLoopError, LlmErrorKind, Result};
30use crate::llm_retry::{
31 LlmRetryConfig, RateLimitInfo, RetryDecision, RetryMetadata, SendOutcome, is_rate_limit_status,
32 retry_request, send_error_message,
33};
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
39const DEFAULT_API_URL: &str = "https://api.openai.com/v1/chat/completions";
40
41pub(crate) fn openai_auth_header_pair<'a>(
49 api_url: &str,
50 api_key: &'a str,
51) -> (&'static str, Cow<'a, str>) {
52 if is_azure_openai_api_url(api_url) {
53 ("api-key", Cow::Borrowed(api_key))
54 } else {
55 ("Authorization", Cow::Owned(format!("Bearer {}", api_key)))
56 }
57}
58
59pub(crate) fn apply_openai_api_auth(
60 request: RequestBuilder,
61 api_url: &str,
62 api_key: &str,
63) -> RequestBuilder {
64 let (name, value) = openai_auth_header_pair(api_url, api_key);
65 request.header(name, value.as_ref())
66}
67
68#[async_trait]
82pub trait AuthHeaderProvider: Send + Sync {
83 async fn auth_header(&self) -> Result<(String, String)>;
87}
88
89pub fn is_azure_openai_api_url(api_url: &str) -> bool {
90 Url::parse(api_url)
91 .ok()
92 .and_then(|url| url.host_str().map(|host| host.to_ascii_lowercase()))
93 .is_some_and(|host| {
94 host.ends_with(".openai.azure.com") || host.ends_with(".services.ai.azure.com")
95 })
96}
97
98pub fn is_openai_api_url(api_url: &str) -> bool {
102 Url::parse(api_url)
103 .ok()
104 .and_then(|url| url.host_str().map(|host| host.to_ascii_lowercase()))
105 .is_some_and(|host| host == "api.openai.com")
106}
107
108const OPENAI_MODELS_URL: &str = "https://api.openai.com/v1/models";
118
119pub fn url_host_eq(api_url: &str, host: &str) -> bool {
121 Url::parse(api_url)
122 .ok()
123 .and_then(|url| url.host_str().map(str::to_owned))
124 .is_some_and(|h| h.eq_ignore_ascii_case(host))
125}
126
127pub fn normalize_api_url(base_url: &str, endpoint_suffix: &str) -> String {
130 let trimmed = base_url.trim_end_matches('/');
131 if trimmed.ends_with(endpoint_suffix) {
132 trimmed.to_string()
133 } else {
134 format!("{trimmed}{endpoint_suffix}")
135 }
136}
137
138pub fn models_url_for_api_url(api_url: &str) -> String {
140 let trimmed = api_url.trim_end_matches('/');
141
142 if let Some(prefix) = trimmed.strip_suffix("/responses") {
143 return format!("{prefix}/models");
144 }
145 if let Some(prefix) = trimmed.strip_suffix("/chat/completions") {
146 return format!("{prefix}/models");
147 }
148 if trimmed.ends_with("/models") {
149 return trimmed.to_string();
150 }
151 if trimmed.ends_with("/v1") || trimmed.ends_with("/openai/v1") {
152 return format!("{trimmed}/models");
153 }
154
155 OPENAI_MODELS_URL.to_string()
156}
157
158pub fn apply_models_api_auth(
161 request: RequestBuilder,
162 api_url: &str,
163 api_key: &str,
164) -> RequestBuilder {
165 if is_azure_openai_api_url(api_url) {
166 request.header("api-key", api_key)
167 } else {
168 request.bearer_auth(api_key)
169 }
170}
171
172pub fn models_api_status_error(status: reqwest::StatusCode) -> AgentLoopError {
175 AgentLoopError::llm(format!("Models API returned status {status}"))
176}
177
178#[derive(Clone)]
202pub struct OpenAIProtocolChatDriver {
203 client: Client,
204 api_key: String,
205 api_url: String,
206 retry_config: LlmRetryConfig,
208 auth_provider: Option<Arc<dyn AuthHeaderProvider>>,
211}
212
213impl OpenAIProtocolChatDriver {
214 pub fn new(api_key: impl Into<String>) -> Self {
216 Self {
217 client: crate::driver_helpers::shared_streaming_http_client(),
218 api_key: api_key.into(),
219 api_url: DEFAULT_API_URL.to_string(),
220 retry_config: LlmRetryConfig::default(),
221 auth_provider: None,
222 }
223 }
224
225 pub fn with_base_url(api_key: impl Into<String>, api_url: impl Into<String>) -> Self {
227 Self {
228 client: crate::driver_helpers::shared_streaming_http_client(),
229 api_key: api_key.into(),
230 api_url: api_url.into(),
231 retry_config: LlmRetryConfig::default(),
232 auth_provider: None,
233 }
234 }
235
236 pub fn with_retry_config(mut self, config: LlmRetryConfig) -> Self {
238 self.retry_config = config;
239 self
240 }
241
242 pub fn with_auth_provider(mut self, provider: Arc<dyn AuthHeaderProvider>) -> Self {
245 self.auth_provider = Some(provider);
246 self
247 }
248
249 pub fn api_url(&self) -> &str {
251 &self.api_url
252 }
253
254 pub fn api_key(&self) -> &str {
256 &self.api_key
257 }
258
259 pub fn client(&self) -> &Client {
261 &self.client
262 }
263
264 async fn send_chat_completion_request(
273 &self,
274 request: &OpenAiRequest,
275 model: &str,
276 ) -> Result<(reqwest::Response, RetryMetadata)> {
277 let last_error: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
278
279 retry_request(
280 &self.retry_config,
281 "OpenAIProtocolDriver",
282 || async {
283 let request_builder = self.client.post(&self.api_url);
287 let request_builder = match &self.auth_provider {
288 Some(provider) => {
289 let (name, value) =
290 provider.auth_header().await.map_err(SendOutcome::Fatal)?;
291 request_builder.header(name, value)
292 }
293 None => apply_openai_api_auth(request_builder, &self.api_url, &self.api_key),
294 };
295 request_builder
296 .header("Content-Type", "application/json")
297 .json(request)
298 .send()
299 .await
300 .map_err(SendOutcome::Send)
301 },
302 |response, attempts, can_retry| {
303 let last_error = Arc::clone(&last_error);
304 let model = model.to_string();
305 async move {
306 let status = response.status();
307
308 if can_retry {
309 let rate_limit_info = if is_rate_limit_status(status) {
311 Some(RateLimitInfo::from_openai_headers(response.headers()))
312 } else {
313 None
314 };
315
316 let error_text = response.text().await.unwrap_or_default();
317
318 if is_openai_request_too_large(status, &error_text) {
320 return RetryDecision::Terminal(AgentLoopError::request_too_large(
321 format!("OpenAI API error ({}): {}", status, error_text),
322 ));
323 }
324
325 if is_provider_quota_message(&error_text) {
328 return RetryDecision::Terminal(AgentLoopError::llm_kind(
329 LlmErrorKind::QuotaExhausted,
330 format!("OpenAI API error ({}): {}", status, error_text),
331 ));
332 }
333
334 let wait = rate_limit_info
335 .as_ref()
336 .map(|info| info.recommended_wait(&self.retry_config, attempts))
337 .unwrap_or_else(|| self.retry_config.calculate_backoff(attempts));
338
339 *last_error.lock().unwrap() = Some(error_text);
340 return RetryDecision::Retry {
341 wait,
342 rate_limit_info,
343 };
344 }
345
346 let error_text = response.text().await.unwrap_or_default();
348 let error_msg = format!("OpenAI API error ({}): {}", status, error_text);
349
350 if is_openai_model_not_found(status, &error_text) {
352 return RetryDecision::Terminal(AgentLoopError::model_not_available(model));
353 }
354
355 if is_openai_request_too_large(status, &error_text) {
357 return RetryDecision::Terminal(AgentLoopError::request_too_large(
358 error_msg,
359 ));
360 }
361
362 let kind = LlmErrorKind::from_provider_status(status.as_u16(), &error_text);
365
366 if attempts > 0 {
367 return RetryDecision::Terminal(AgentLoopError::llm_kind(
368 kind,
369 format!(
370 "{} (after {} retries, last error: {})",
371 error_msg,
372 attempts,
373 last_error.lock().unwrap().take().unwrap_or_default()
374 ),
375 ));
376 }
377
378 RetryDecision::Terminal(AgentLoopError::llm_kind(kind, error_msg))
379 }
380 },
381 |e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
382 )
383 .await
384 }
385
386 fn convert_role(role: &LlmMessageRole) -> &'static str {
387 match role {
388 LlmMessageRole::System => "system",
389 LlmMessageRole::User => "user",
390 LlmMessageRole::Assistant => "assistant",
391 LlmMessageRole::Tool => "tool",
392 }
393 }
394
395 fn convert_message(msg: &LlmMessage) -> OpenAiMessage {
396 let content = match &msg.content {
397 LlmMessageContent::Text(text) => OpenAiContent::Text(text.clone()),
398 LlmMessageContent::Parts(parts) => {
399 let openai_parts: Vec<OpenAiContentPart> = parts
400 .iter()
401 .map(|part| match part {
402 LlmContentPart::Text { text } => OpenAiContentPart::Text {
403 r#type: "text".to_string(),
404 text: text.clone(),
405 },
406 LlmContentPart::Image { url } => OpenAiContentPart::ImageUrl {
407 r#type: "image_url".to_string(),
408 image_url: OpenAiImageUrl { url: url.clone() },
409 },
410 LlmContentPart::Audio { url } => OpenAiContentPart::InputAudio {
411 r#type: "input_audio".to_string(),
412 input_audio: OpenAiInputAudio {
413 data: url.clone(),
414 format: "wav".to_string(),
415 },
416 },
417 })
418 .collect();
419 OpenAiContent::Parts(openai_parts)
420 }
421 };
422
423 let tool_calls = if msg.role == LlmMessageRole::Assistant {
425 msg.tool_calls.as_ref().map(|calls| {
426 calls
427 .iter()
428 .map(|tc| OpenAiToolCall {
429 id: tc.id.clone(),
430 r#type: "function".to_string(),
431 function: OpenAiFunctionCall {
432 name: tc.name.clone(),
433 arguments: serde_json::to_string(&tc.arguments).unwrap_or_default(),
434 },
435 })
436 .collect()
437 })
438 } else {
439 None
440 };
441
442 OpenAiMessage {
443 role: Self::convert_role(&msg.role).to_string(),
444 content: Some(content),
445 tool_calls,
446 tool_call_id: msg.tool_call_id.clone(),
447 }
448 }
449
450 fn convert_tools(tools: &[ToolDefinition]) -> Vec<OpenAiTool> {
451 tools
452 .iter()
453 .map(|tool| OpenAiTool {
454 r#type: "function".to_string(),
455 function: OpenAiFunction {
456 name: tool.name().to_string(),
457 description: tool.description().to_string(),
458 parameters: tool.parameters().clone(),
459 },
460 })
461 .collect()
462 }
463}
464
465fn drop_orphaned_tool_messages(messages: &[LlmMessage]) -> Vec<LlmMessage> {
469 use std::collections::HashSet;
470
471 let visible_call_ids: HashSet<&str> = messages
472 .iter()
473 .filter(|m| m.role == LlmMessageRole::Assistant)
474 .flat_map(|m| m.tool_calls.iter().flatten())
475 .map(|tc| tc.id.as_str())
476 .collect();
477
478 if visible_call_ids.is_empty() {
479 return messages
480 .iter()
481 .filter(|m| m.role != LlmMessageRole::Tool)
482 .cloned()
483 .collect();
484 }
485
486 messages
487 .iter()
488 .filter(|m| {
489 if m.role == LlmMessageRole::Tool {
490 return m
491 .tool_call_id
492 .as_deref()
493 .is_none_or(|id| visible_call_ids.contains(id));
494 }
495 true
496 })
497 .cloned()
498 .collect()
499}
500
501#[async_trait]
502impl ChatDriver for OpenAIProtocolChatDriver {
503 async fn chat_completion_stream(
504 &self,
505 messages: Vec<LlmMessage>,
506 config: &LlmCallConfig,
507 ) -> Result<LlmResponseStream> {
508 let messages = drop_orphaned_tool_messages(&messages);
512 let openai_messages: Vec<OpenAiMessage> =
513 messages.iter().map(Self::convert_message).collect();
514
515 let tools = if config.tools.is_empty() {
516 None
517 } else {
518 Some(Self::convert_tools(&config.tools))
519 };
520
521 let metadata = if config.metadata.is_empty() {
523 None
524 } else {
525 Some(config.metadata.clone())
526 };
527
528 let request = OpenAiRequest {
529 model: config.model.clone(),
530 messages: openai_messages,
531 temperature: config.temperature,
532 max_tokens: config.max_tokens,
533 stream: true,
534 stream_options: Some(OpenAiStreamOptions {
535 include_usage: true,
536 }),
537 tools,
538 parallel_tool_calls: config
539 .resolved_parallel_tool_calls(self.supports_parallel_tool_calls(&config.model)),
540 reasoning_effort: config
542 .reasoning_effort
543 .as_ref()
544 .filter(|e| !e.eq_ignore_ascii_case("none"))
545 .cloned(),
546 service_tier: config.speed.clone(),
547 verbosity: config.verbosity.clone(),
548 metadata,
549 };
550
551 let (event_stream, retry_metadata) =
557 connect_sse_with_reconnect(&self.retry_config, "OpenAIProtocolDriver", |_attempt| {
558 self.send_chat_completion_request(&request, &config.model)
559 })
560 .await?;
561
562 let model = config.model.clone();
563 let total_tokens = Arc::new(Mutex::new(0u32));
564 let prompt_tokens = Arc::new(Mutex::new(0u32));
565 let cache_read_tokens = Arc::new(Mutex::new(Option::<u32>::None));
566 let provider_cost_usd = Arc::new(Mutex::new(Option::<f64>::None));
569 let accumulated_tool_calls = Arc::new(Mutex::new(StreamToolCallAccumulator::new()));
570 let finish_reason = Arc::new(Mutex::new(Option::<String>::None));
571 let response_id = Arc::new(Mutex::new(Option::<String>::None));
574 let shared_retry_metadata = if retry_metadata.had_retries() {
576 Some(Arc::new(retry_metadata))
577 } else {
578 None
579 };
580
581 let converted_stream: LlmResponseStream = Box::pin(
585 event_stream
586 .then(move |result| {
587 let model = model.clone();
588 let total_tokens = Arc::clone(&total_tokens);
589 let prompt_tokens = Arc::clone(&prompt_tokens);
590 let cache_read_tokens = Arc::clone(&cache_read_tokens);
591 let provider_cost_usd = Arc::clone(&provider_cost_usd);
592 let accumulated_tool_calls = Arc::clone(&accumulated_tool_calls);
593 let finish_reason = Arc::clone(&finish_reason);
594 let response_id = Arc::clone(&response_id);
595 let retry_metadata_for_done = shared_retry_metadata.clone();
596
597 async move {
598 let event = match result {
599 Ok(event) => event,
600 Err(e) => {
601 return vec![Ok(LlmStreamEvent::Error(
602 format!("Stream error: {}", e).into(),
603 ))];
604 }
605 };
606
607 if event.data == "[DONE]" {
608 let output_tokens = *total_tokens.lock().unwrap();
609 let input_tokens = *prompt_tokens.lock().unwrap();
610 let cached = *cache_read_tokens.lock().unwrap();
611 let cost = *provider_cost_usd.lock().unwrap();
612 let resp_id = response_id.lock().unwrap().clone();
613 let mut reason = finish_reason.lock().unwrap().clone();
614
615 let mut events = Vec::new();
616
617 {
626 let mut acc = accumulated_tool_calls.lock().unwrap();
627 if let Some(event) =
628 take_pending_tool_calls(&mut acc, reason.as_deref())
629 {
630 events.push(Ok(event));
631 reason.get_or_insert_with(|| "tool_calls".to_string());
632 }
633 }
634
635 events.push(Ok(LlmStreamEvent::Done(Box::new(
636 LlmCompletionMetadata {
637 total_tokens: Some(input_tokens + output_tokens),
640 prompt_tokens: Some(disjoint_prompt_tokens(
641 input_tokens,
642 cached,
643 )),
644 completion_tokens: Some(output_tokens),
645 cache_read_tokens: cached,
646 cache_creation_tokens: None,
647 provider_cost_usd: cost,
648 model: Some(model),
649 finish_reason: reason.or_else(|| Some("stop".to_string())),
650 retry_metadata: retry_metadata_for_done
651 .map(|arc| (*arc).clone()),
652 response_id: resp_id,
653 phase: None,
654 },
655 ))));
656
657 return events;
658 }
659
660 match serde_json::from_str::<OpenAiStreamChunk>(&event.data) {
661 Ok(chunk) => {
662 if let Some(id) = &chunk.id {
667 let mut rid = response_id.lock().unwrap();
668 if rid.is_none() {
669 *rid = Some(id.clone());
670 }
671 }
672
673 if let Some(usage) = &chunk.usage {
675 if let Some(pt) = usage.prompt_tokens {
676 *prompt_tokens.lock().unwrap() = pt;
677 }
678 if let Some(ct) = usage.completion_tokens {
679 *total_tokens.lock().unwrap() = ct;
680 }
681 if let Some(details) = &usage.prompt_tokens_details
683 && details.cached_tokens.is_some()
684 {
685 *cache_read_tokens.lock().unwrap() = details.cached_tokens;
686 }
687 if usage.cost.is_some() {
690 *provider_cost_usd.lock().unwrap() = usage.cost;
691 }
692 }
693
694 if let Some(choice) = chunk.choices.first() {
695 let mut tt = total_tokens.lock().unwrap();
696 let mut acc = accumulated_tool_calls.lock().unwrap();
697 let mut fr = finish_reason.lock().unwrap();
698 let stream_event =
699 process_stream_choice(choice, &mut tt, &mut acc, &mut fr);
700 return vec![Ok(stream_event)];
701 }
702 vec![Ok(LlmStreamEvent::TextDelta(String::new()))]
703 }
704 Err(e) => vec![Ok(LlmStreamEvent::Error(
705 format!("Failed to parse chunk: {}", e).into(),
706 ))],
707 }
708 }
709 })
710 .flat_map(futures::stream::iter),
711 );
712
713 Ok(converted_stream)
714 }
715
716 fn supports_parallel_tool_calls(&self, _model: &str) -> bool {
720 true
721 }
722}
723
724impl std::fmt::Debug for OpenAIProtocolChatDriver {
725 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
726 f.debug_struct("OpenAIProtocolChatDriver")
727 .field("api_url", &self.api_url)
728 .field("api_key", &"[REDACTED]")
729 .finish()
730 }
731}
732
733pub fn is_openai_model_not_found(status: reqwest::StatusCode, error_text: &str) -> bool {
744 let error_lower = error_text.to_lowercase();
745
746 if status == reqwest::StatusCode::NOT_FOUND
748 || status == reqwest::StatusCode::BAD_REQUEST
749 || status == reqwest::StatusCode::FORBIDDEN
750 {
751 if error_lower.contains("model_not_found") {
753 return true;
754 }
755 }
756
757 if status == reqwest::StatusCode::NOT_FOUND {
759 if error_lower.contains("does not exist") {
760 return true;
761 }
762 if error_lower.contains("model") && error_lower.contains("not found") {
763 return true;
764 }
765 }
766
767 false
768}
769
770pub fn is_openai_request_too_large(status: reqwest::StatusCode, error_text: &str) -> bool {
777 let error_lower = error_text.to_lowercase();
778
779 if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
781 if error_lower.contains("request too large") {
783 return true;
784 }
785 if error_lower.contains("tokens") && error_lower.contains("limit") {
787 return true;
788 }
789 }
790
791 if status == reqwest::StatusCode::BAD_REQUEST {
793 if error_lower.contains("context_length_exceeded") {
795 return true;
796 }
797 if error_lower.contains("maximum context length") {
799 return true;
800 }
801 }
802
803 if error_lower.contains("tokens must be reduced")
805 || error_lower.contains("reduce the length")
806 || error_lower.contains("input is too long")
807 {
808 return true;
809 }
810
811 false
812}
813
814#[derive(Debug, Serialize)]
819struct OpenAiRequest {
820 model: String,
821 messages: Vec<OpenAiMessage>,
822 #[serde(skip_serializing_if = "Option::is_none")]
823 temperature: Option<f32>,
824 #[serde(skip_serializing_if = "Option::is_none")]
825 max_tokens: Option<u32>,
826 stream: bool,
827 #[serde(skip_serializing_if = "Option::is_none")]
829 stream_options: Option<OpenAiStreamOptions>,
830 #[serde(skip_serializing_if = "Option::is_none")]
831 tools: Option<Vec<OpenAiTool>>,
832 #[serde(skip_serializing_if = "Option::is_none")]
835 parallel_tool_calls: Option<bool>,
836 #[serde(skip_serializing_if = "Option::is_none")]
837 reasoning_effort: Option<String>,
838 #[serde(skip_serializing_if = "Option::is_none")]
841 service_tier: Option<String>,
842 #[serde(skip_serializing_if = "Option::is_none")]
846 verbosity: Option<String>,
847 #[serde(skip_serializing_if = "Option::is_none")]
850 metadata: Option<std::collections::HashMap<String, String>>,
851}
852
853#[derive(Debug, Serialize)]
854struct OpenAiStreamOptions {
855 include_usage: bool,
856}
857
858#[derive(Debug, Serialize, Deserialize)]
859#[serde(untagged)]
860enum OpenAiContent {
861 Text(String),
862 Parts(Vec<OpenAiContentPart>),
863}
864
865#[derive(Debug, Serialize, Deserialize)]
866#[serde(untagged)]
867enum OpenAiContentPart {
868 Text {
869 r#type: String,
870 text: String,
871 },
872 ImageUrl {
873 r#type: String,
874 image_url: OpenAiImageUrl,
875 },
876 InputAudio {
877 r#type: String,
878 input_audio: OpenAiInputAudio,
879 },
880}
881
882#[derive(Debug, Serialize, Deserialize)]
883struct OpenAiImageUrl {
884 url: String,
885}
886
887#[derive(Debug, Serialize, Deserialize)]
888struct OpenAiInputAudio {
889 data: String,
890 format: String,
891}
892
893#[derive(Debug, Serialize, Deserialize)]
894struct OpenAiMessage {
895 role: String,
896 #[serde(skip_serializing_if = "Option::is_none")]
897 content: Option<OpenAiContent>,
898 #[serde(skip_serializing_if = "Option::is_none")]
899 tool_calls: Option<Vec<OpenAiToolCall>>,
900 #[serde(skip_serializing_if = "Option::is_none")]
901 tool_call_id: Option<String>,
902}
903
904#[derive(Debug, Serialize, Deserialize)]
905struct OpenAiTool {
906 r#type: String,
907 function: OpenAiFunction,
908}
909
910#[derive(Debug, Serialize, Deserialize)]
911struct OpenAiFunction {
912 name: String,
913 description: String,
914 parameters: Value,
915}
916
917#[derive(Debug, Serialize, Deserialize)]
918struct OpenAiToolCall {
919 id: String,
920 r#type: String,
921 function: OpenAiFunctionCall,
922}
923
924#[derive(Debug, Serialize, Deserialize)]
925struct OpenAiFunctionCall {
926 name: String,
927 arguments: String,
928}
929
930#[derive(Debug, Deserialize)]
931#[allow(dead_code)] struct OpenAiStreamChunk {
933 #[serde(default)]
935 id: Option<String>,
936 #[serde(default)]
938 model: Option<String>,
939 choices: Vec<OpenAiStreamChoice>,
940 #[serde(default)]
941 usage: Option<OpenAiUsage>,
942}
943
944#[derive(Debug, Deserialize)]
945struct OpenAiUsage {
946 prompt_tokens: Option<u32>,
947 completion_tokens: Option<u32>,
948 #[serde(default)]
950 prompt_tokens_details: Option<OpenAiPromptTokensDetails>,
951 #[serde(default)]
954 cost: Option<f64>,
955}
956
957#[derive(Debug, Deserialize, Default)]
958struct OpenAiPromptTokensDetails {
959 #[serde(default)]
961 cached_tokens: Option<u32>,
962}
963
964#[derive(Debug, Deserialize)]
965struct OpenAiStreamChoice {
966 delta: OpenAiDelta,
967 #[serde(default)]
968 finish_reason: Option<String>,
969}
970
971#[derive(Debug, Deserialize)]
972struct OpenAiDelta {
973 #[serde(default)]
974 content: Option<String>,
975 #[serde(default)]
976 tool_calls: Option<Vec<OpenAiStreamToolCall>>,
977}
978
979#[derive(Debug, Deserialize)]
980struct OpenAiStreamToolCall {
981 index: u32,
982 id: Option<String>,
983 function: Option<OpenAiStreamFunction>,
984}
985
986#[derive(Debug, Deserialize)]
987struct OpenAiStreamFunction {
988 name: Option<String>,
989 arguments: Option<String>,
990}
991
992fn take_pending_tool_calls(
1003 accumulated_tool_calls: &mut StreamToolCallAccumulator,
1004 finish_reason: Option<&str>,
1005) -> Option<LlmStreamEvent> {
1006 if accumulated_tool_calls.is_empty() {
1007 return None;
1008 }
1009
1010 if !matches!(finish_reason, None | Some("tool_calls")) {
1013 let _ = accumulated_tool_calls.take_finalized();
1014 return None;
1015 }
1016
1017 let calls = accumulated_tool_calls.take_pending_strict();
1018 if calls.is_empty() {
1019 None
1020 } else {
1021 Some(LlmStreamEvent::ToolCalls(calls))
1022 }
1023}
1024
1025fn process_stream_choice(
1035 choice: &OpenAiStreamChoice,
1036 total_tokens: &mut u32,
1037 accumulated_tool_calls: &mut StreamToolCallAccumulator,
1038 finish_reason: &mut Option<String>,
1039) -> LlmStreamEvent {
1040 if let Some(tool_calls) = &choice.delta.tool_calls {
1044 for tc in tool_calls {
1045 accumulated_tool_calls.apply_indexed_delta(
1046 tc.index,
1047 tc.id.as_deref(),
1048 tc.function.as_ref().and_then(|f| f.name.as_deref()),
1049 tc.function.as_ref().and_then(|f| f.arguments.as_deref()),
1050 );
1051 }
1052 return LlmStreamEvent::TextDelta(String::new());
1053 }
1054
1055 if let Some(content) = &choice.delta.content
1058 && !content.is_empty()
1059 {
1060 *total_tokens += 1;
1061 return LlmStreamEvent::TextDelta(content.clone());
1062 }
1063
1064 if let Some(fr) = &choice.finish_reason {
1068 *finish_reason = Some(fr.clone());
1069
1070 if fr == "tool_calls" && !accumulated_tool_calls.is_empty() {
1071 return LlmStreamEvent::ToolCalls(accumulated_tool_calls.take_finalized());
1072 }
1073 }
1074
1075 LlmStreamEvent::TextDelta(String::new())
1076}
1077
1078#[cfg(test)]
1083mod tests {
1084 use super::*;
1085 use serde_json::json;
1086
1087 #[test]
1088 fn test_convert_message_preserves_multiple_system_messages() {
1089 let messages = [
1094 LlmMessage::text(LlmMessageRole::System, "A"),
1095 LlmMessage::text(LlmMessageRole::User, "hi"),
1096 LlmMessage::text(LlmMessageRole::System, "B"),
1097 ];
1098 let converted: Vec<OpenAiMessage> = messages
1099 .iter()
1100 .map(OpenAIProtocolChatDriver::convert_message)
1101 .collect();
1102 let system_texts: Vec<String> = converted
1103 .iter()
1104 .filter(|m| m.role == "system")
1105 .filter_map(|m| match &m.content {
1106 Some(OpenAiContent::Text(t)) => Some(t.clone()),
1107 _ => None,
1108 })
1109 .collect();
1110 assert_eq!(system_texts, vec!["A".to_string(), "B".to_string()]);
1111 }
1112
1113 #[test]
1114 fn test_driver_with_api_key() {
1115 let driver = OpenAIProtocolChatDriver::new("test-key");
1116 assert!(format!("{:?}", driver).contains("OpenAIProtocolChatDriver"));
1117 }
1118
1119 #[test]
1120 fn test_driver_with_base_url() {
1121 let driver = OpenAIProtocolChatDriver::with_base_url(
1122 "test-key",
1123 "https://custom.api.com/v1/completions",
1124 );
1125 assert!(format!("{:?}", driver).contains("OpenAIProtocolChatDriver"));
1126 assert_eq!(driver.api_url(), "https://custom.api.com/v1/completions");
1127 }
1128
1129 #[test]
1130 fn test_is_azure_openai_api_url() {
1131 assert!(is_azure_openai_api_url(
1132 "https://example.openai.azure.com/openai/v1/chat/completions"
1133 ));
1134 assert!(is_azure_openai_api_url(
1135 "https://example.services.ai.azure.com/openai/v1/responses"
1136 ));
1137 assert!(!is_azure_openai_api_url(
1138 "https://api.openai.com/v1/chat/completions"
1139 ));
1140 }
1141
1142 #[test]
1143 fn test_request_includes_stream_options_for_usage() {
1144 let request = OpenAiRequest {
1147 verbosity: None,
1148 service_tier: None,
1149 model: "gpt-4o".to_string(),
1150 messages: vec![OpenAiMessage {
1151 role: "user".to_string(),
1152 content: Some(OpenAiContent::Text("Hello".to_string())),
1153 tool_calls: None,
1154 tool_call_id: None,
1155 }],
1156 temperature: None,
1157 max_tokens: None,
1158 stream: true,
1159 stream_options: Some(OpenAiStreamOptions {
1160 include_usage: true,
1161 }),
1162 tools: None,
1163 parallel_tool_calls: None,
1164 reasoning_effort: None,
1165 metadata: None,
1166 };
1167
1168 let json = serde_json::to_value(&request).unwrap();
1169 assert_eq!(json["stream"], true);
1170 assert_eq!(json["stream_options"]["include_usage"], true);
1171 }
1172
1173 #[test]
1174 fn test_request_includes_metadata() {
1175 let mut metadata = std::collections::HashMap::new();
1177 metadata.insert("session_id".to_string(), "session_abc123".to_string());
1178 metadata.insert("agent_id".to_string(), "agent_xyz789".to_string());
1179
1180 let request = OpenAiRequest {
1181 verbosity: None,
1182 service_tier: None,
1183 model: "gpt-4o".to_string(),
1184 messages: vec![OpenAiMessage {
1185 role: "user".to_string(),
1186 content: Some(OpenAiContent::Text("Hello".to_string())),
1187 tool_calls: None,
1188 tool_call_id: None,
1189 }],
1190 temperature: None,
1191 max_tokens: None,
1192 stream: true,
1193 stream_options: None,
1194 tools: None,
1195 parallel_tool_calls: None,
1196 reasoning_effort: None,
1197 metadata: Some(metadata),
1198 };
1199
1200 let json = serde_json::to_value(&request).unwrap();
1201 assert_eq!(json["metadata"]["session_id"], "session_abc123");
1202 assert_eq!(json["metadata"]["agent_id"], "agent_xyz789");
1203 }
1204
1205 #[test]
1206 fn test_usage_chunk_parsing() {
1207 let usage_chunk = r#"{
1210 "id": "chatcmpl-123",
1211 "object": "chat.completion.chunk",
1212 "created": 1234567890,
1213 "model": "gpt-4o",
1214 "choices": [],
1215 "usage": {
1216 "prompt_tokens": 150,
1217 "completion_tokens": 42,
1218 "total_tokens": 192
1219 }
1220 }"#;
1221
1222 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1223 assert!(chunk.usage.is_some());
1224 let usage = chunk.usage.unwrap();
1225 assert_eq!(usage.prompt_tokens, Some(150));
1226 assert_eq!(usage.completion_tokens, Some(42));
1227 }
1228
1229 #[test]
1230 fn test_usage_chunk_with_cached_tokens() {
1231 let usage_chunk = r#"{
1233 "id": "chatcmpl-123",
1234 "choices": [],
1235 "usage": {
1236 "prompt_tokens": 150,
1237 "completion_tokens": 42,
1238 "prompt_tokens_details": {
1239 "cached_tokens": 100
1240 }
1241 }
1242 }"#;
1243
1244 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1245 let usage = chunk.usage.unwrap();
1246 assert_eq!(usage.prompt_tokens, Some(150));
1247 assert_eq!(usage.completion_tokens, Some(42));
1248 assert!(usage.prompt_tokens_details.is_some());
1249 assert_eq!(
1250 usage.prompt_tokens_details.unwrap().cached_tokens,
1251 Some(100)
1252 );
1253 }
1254
1255 #[test]
1256 fn test_usage_chunk_with_openrouter_cost() {
1257 let usage_chunk = r#"{
1259 "id": "gen-123",
1260 "choices": [],
1261 "usage": {
1262 "prompt_tokens": 194,
1263 "completion_tokens": 2,
1264 "total_tokens": 196,
1265 "cost": 0.00095
1266 }
1267 }"#;
1268
1269 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1270 let usage = chunk.usage.unwrap();
1271 assert_eq!(usage.cost, Some(0.00095));
1272 }
1273
1274 #[test]
1275 fn test_usage_chunk_without_cost_defaults_none() {
1276 let usage_chunk = r#"{
1278 "id": "chatcmpl-123",
1279 "choices": [],
1280 "usage": { "prompt_tokens": 10, "completion_tokens": 5 }
1281 }"#;
1282
1283 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1284 assert_eq!(chunk.usage.unwrap().cost, None);
1285 }
1286
1287 #[test]
1288 fn test_chunk_id_is_captured() {
1289 let chunk_with_id: OpenAiStreamChunk =
1290 serde_json::from_str(r#"{"id":"gen-abc123","choices":[]}"#).unwrap();
1291 assert_eq!(chunk_with_id.id.as_deref(), Some("gen-abc123"));
1292
1293 let chunk_no_id: OpenAiStreamChunk = serde_json::from_str(r#"{"choices":[]}"#).unwrap();
1294 assert!(chunk_no_id.id.is_none());
1295 }
1296
1297 #[test]
1298 fn test_finish_reason_chunk_parsing() {
1299 let finish_chunk = r#"{
1301 "id": "chatcmpl-123",
1302 "choices": [{
1303 "index": 0,
1304 "delta": {},
1305 "finish_reason": "stop"
1306 }]
1307 }"#;
1308
1309 let chunk: OpenAiStreamChunk = serde_json::from_str(finish_chunk).unwrap();
1310 assert!(chunk.usage.is_none()); assert_eq!(chunk.choices.len(), 1);
1312 assert_eq!(chunk.choices[0].finish_reason, Some("stop".to_string()));
1313 }
1314
1315 #[test]
1320 fn test_is_openai_request_too_large_429_request_too_large() {
1321 let error = r#"{"error":{"message":"Request too large for gpt-4o in organization org-xxx on tokens per min (TPM): Limit 500000, Requested 538772."}}"#;
1322 assert!(is_openai_request_too_large(
1323 reqwest::StatusCode::TOO_MANY_REQUESTS,
1324 error
1325 ));
1326 }
1327
1328 #[test]
1329 fn test_is_openai_request_too_large_429_token_limit() {
1330 let error =
1331 r#"{"error":{"message":"tokens per min (TPM): Limit 500000, Requested 600000"}}"#;
1332 assert!(is_openai_request_too_large(
1333 reqwest::StatusCode::TOO_MANY_REQUESTS,
1334 error
1335 ));
1336 }
1337
1338 #[test]
1339 fn test_is_openai_request_too_large_400_context_length() {
1340 let error = r#"{"error":{"code":"context_length_exceeded","message":"This model's maximum context length is 128000 tokens."}}"#;
1341 assert!(is_openai_request_too_large(
1342 reqwest::StatusCode::BAD_REQUEST,
1343 error
1344 ));
1345 }
1346
1347 #[test]
1348 fn test_is_openai_request_too_large_400_max_context() {
1349 let error =
1350 r#"{"error":{"message":"This model's maximum context length is 128000 tokens"}}"#;
1351 assert!(is_openai_request_too_large(
1352 reqwest::StatusCode::BAD_REQUEST,
1353 error
1354 ));
1355 }
1356
1357 #[test]
1358 fn test_is_openai_request_too_large_tokens_must_be_reduced() {
1359 let error = r#"{"error":{"message":"The input or output tokens must be reduced"}}"#;
1360 assert!(is_openai_request_too_large(
1361 reqwest::StatusCode::BAD_REQUEST,
1362 error
1363 ));
1364 }
1365
1366 #[test]
1367 fn test_is_openai_request_too_large_false_for_other_errors() {
1368 let error = r#"{"error":{"message":"Rate limit exceeded: too many requests per minute"}}"#;
1370 assert!(!is_openai_request_too_large(
1371 reqwest::StatusCode::TOO_MANY_REQUESTS,
1372 error
1373 ));
1374
1375 let error = r#"{"error":{"message":"Internal server error"}}"#;
1377 assert!(!is_openai_request_too_large(
1378 reqwest::StatusCode::INTERNAL_SERVER_ERROR,
1379 error
1380 ));
1381
1382 let error = r#"{"error":{"message":"Invalid request"}}"#;
1384 assert!(!is_openai_request_too_large(
1385 reqwest::StatusCode::BAD_REQUEST,
1386 error
1387 ));
1388 }
1389
1390 #[test]
1395 fn test_is_openai_model_not_found_real_error() {
1396 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}}"#;
1398 assert!(is_openai_model_not_found(
1399 reqwest::StatusCode::NOT_FOUND,
1400 error
1401 ));
1402 }
1403
1404 #[test]
1405 fn test_is_openai_model_not_found_does_not_exist() {
1406 let error = r#"{"error":{"message":"The model 'fake-model' does not exist"}}"#;
1407 assert!(is_openai_model_not_found(
1408 reqwest::StatusCode::NOT_FOUND,
1409 error
1410 ));
1411 }
1412
1413 #[test]
1414 fn test_is_openai_model_not_found_generic_not_found() {
1415 let error = r#"{"error":{"message":"Model not found"}}"#;
1416 assert!(is_openai_model_not_found(
1417 reqwest::StatusCode::NOT_FOUND,
1418 error
1419 ));
1420 }
1421
1422 #[test]
1423 fn test_is_openai_model_not_found_400_with_model_not_found_code() {
1424 let error = r#"{"error":{"code":"model_not_found","message":"The requested model 'gpt-99' does not exist.","type":"invalid_request_error","param":"model"}}"#;
1426 assert!(is_openai_model_not_found(
1427 reqwest::StatusCode::BAD_REQUEST,
1428 error
1429 ));
1430 }
1431
1432 #[test]
1433 fn test_is_openai_model_not_found_false_for_non_model_error() {
1434 let error = r#"{"error":{"code":"invalid_request","message":"Some other error"}}"#;
1436 assert!(!is_openai_model_not_found(
1437 reqwest::StatusCode::BAD_REQUEST,
1438 error
1439 ));
1440 }
1441
1442 #[test]
1443 fn test_is_openai_model_not_found_false_for_other_404() {
1444 let error = r#"{"error":{"message":"Endpoint not found"}}"#;
1446 assert!(!is_openai_model_not_found(
1447 reqwest::StatusCode::NOT_FOUND,
1448 error
1449 ));
1450 }
1451
1452 #[test]
1453 fn test_is_openai_model_not_found_403_tier_gated_model() {
1454 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}}"#;
1457 assert!(is_openai_model_not_found(
1458 reqwest::StatusCode::FORBIDDEN,
1459 error
1460 ));
1461 }
1462
1463 #[test]
1464 fn test_is_openai_model_not_found_403_plain_auth_error_is_not_model_not_found() {
1465 let error = r#"{"error":{"message":"Invalid authentication credentials","type":"authentication_error"}}"#;
1468 assert!(!is_openai_model_not_found(
1469 reqwest::StatusCode::FORBIDDEN,
1470 error
1471 ));
1472 }
1473
1474 #[test]
1479 fn test_reasoning_effort_none_is_omitted() {
1480 let request = OpenAiRequest {
1483 verbosity: None,
1484 service_tier: None,
1485 model: "gpt-4o-mini".to_string(),
1486 messages: vec![OpenAiMessage {
1487 role: "user".to_string(),
1488 content: Some(OpenAiContent::Text("Hello".to_string())),
1489 tool_calls: None,
1490 tool_call_id: None,
1491 }],
1492 temperature: None,
1493 max_tokens: None,
1494 stream: true,
1495 stream_options: None,
1496 tools: None,
1497 parallel_tool_calls: None,
1498 reasoning_effort: Some("none".to_string())
1499 .as_ref()
1500 .filter(|e| !e.eq_ignore_ascii_case("none"))
1501 .cloned(),
1502 metadata: None,
1503 };
1504
1505 let json = serde_json::to_value(&request).unwrap();
1506 assert!(
1507 json.get("reasoning_effort").is_none(),
1508 "reasoning_effort should be omitted when effort is 'none'"
1509 );
1510 }
1511
1512 #[test]
1513 fn test_reasoning_effort_high_is_included() {
1514 let request = OpenAiRequest {
1515 verbosity: None,
1516 service_tier: None,
1517 model: "o3-mini".to_string(),
1518 messages: vec![OpenAiMessage {
1519 role: "user".to_string(),
1520 content: Some(OpenAiContent::Text("Hello".to_string())),
1521 tool_calls: None,
1522 tool_call_id: None,
1523 }],
1524 temperature: None,
1525 max_tokens: None,
1526 stream: true,
1527 stream_options: None,
1528 tools: None,
1529 parallel_tool_calls: None,
1530 reasoning_effort: Some("high".to_string())
1531 .as_ref()
1532 .filter(|e| !e.eq_ignore_ascii_case("none"))
1533 .cloned(),
1534 metadata: None,
1535 };
1536
1537 let json = serde_json::to_value(&request).unwrap();
1538 assert_eq!(json["reasoning_effort"], "high");
1539 }
1540
1541 #[test]
1545 fn test_request_serializes_parallel_tool_calls() {
1546 fn build(flag: Option<bool>) -> serde_json::Value {
1547 let request = OpenAiRequest {
1548 verbosity: None,
1549 service_tier: None,
1550 model: "gpt-4o-mini".to_string(),
1551 messages: vec![OpenAiMessage {
1552 role: "user".to_string(),
1553 content: Some(OpenAiContent::Text("Hello".to_string())),
1554 tool_calls: None,
1555 tool_call_id: None,
1556 }],
1557 temperature: None,
1558 max_tokens: None,
1559 stream: true,
1560 stream_options: None,
1561 tools: None,
1562 parallel_tool_calls: flag,
1563 reasoning_effort: None,
1564 metadata: None,
1565 };
1566 serde_json::to_value(&request).unwrap()
1567 }
1568
1569 assert!(build(None).get("parallel_tool_calls").is_none());
1571 assert_eq!(build(Some(true))["parallel_tool_calls"], true);
1573 assert_eq!(build(Some(false))["parallel_tool_calls"], false);
1574 }
1575
1576 #[test]
1579 fn test_request_serializes_service_tier() {
1580 fn build(tier: Option<&str>) -> serde_json::Value {
1581 let request = OpenAiRequest {
1582 service_tier: tier.map(str::to_string),
1583 verbosity: None,
1584 model: "gpt-4o-mini".to_string(),
1585 messages: vec![OpenAiMessage {
1586 role: "user".to_string(),
1587 content: Some(OpenAiContent::Text("Hello".to_string())),
1588 tool_calls: None,
1589 tool_call_id: None,
1590 }],
1591 temperature: None,
1592 max_tokens: None,
1593 stream: true,
1594 stream_options: None,
1595 tools: None,
1596 parallel_tool_calls: None,
1597 reasoning_effort: None,
1598 metadata: None,
1599 };
1600 serde_json::to_value(&request).unwrap()
1601 }
1602
1603 assert!(build(None).get("service_tier").is_none());
1604 assert_eq!(build(Some("flex"))["service_tier"], "flex");
1605 assert_eq!(build(Some("priority"))["service_tier"], "priority");
1606 }
1607
1608 #[test]
1611 fn test_request_serializes_verbosity() {
1612 fn build(verbosity: Option<&str>) -> serde_json::Value {
1613 let request = OpenAiRequest {
1614 service_tier: None,
1615 verbosity: verbosity.map(str::to_string),
1616 model: "gpt-5.6-sol".to_string(),
1617 messages: vec![OpenAiMessage {
1618 role: "user".to_string(),
1619 content: Some(OpenAiContent::Text("Hello".to_string())),
1620 tool_calls: None,
1621 tool_call_id: None,
1622 }],
1623 temperature: None,
1624 max_tokens: None,
1625 stream: true,
1626 stream_options: None,
1627 tools: None,
1628 parallel_tool_calls: None,
1629 reasoning_effort: None,
1630 metadata: None,
1631 };
1632 serde_json::to_value(&request).unwrap()
1633 }
1634
1635 assert!(build(None).get("verbosity").is_none());
1636 assert_eq!(build(Some("low"))["verbosity"], "low");
1637 assert_eq!(build(Some("high"))["verbosity"], "high");
1638 }
1639
1640 fn choice(json_str: &str) -> OpenAiStreamChoice {
1645 serde_json::from_str(json_str).unwrap()
1646 }
1647
1648 #[test]
1652 fn test_empty_content_finish_chunk_still_emits_tool_calls() {
1653 let mut total_tokens = 0u32;
1654 let mut acc = StreamToolCallAccumulator::new();
1655 let mut finish_reason: Option<String> = None;
1656
1657 let e = process_stream_choice(
1659 &choice(
1660 r#"{"delta":{"content":null,"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":""}}]},"finish_reason":null}"#,
1661 ),
1662 &mut total_tokens,
1663 &mut acc,
1664 &mut finish_reason,
1665 );
1666 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()));
1667
1668 let e = process_stream_choice(
1670 &choice(
1671 r#"{"delta":{"content":null,"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":\"Cargo.toml\"}"}}]},"finish_reason":null}"#,
1672 ),
1673 &mut total_tokens,
1674 &mut acc,
1675 &mut finish_reason,
1676 );
1677 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()));
1678
1679 let e = process_stream_choice(
1682 &choice(r#"{"delta":{"content":""},"finish_reason":"tool_calls"}"#),
1683 &mut total_tokens,
1684 &mut acc,
1685 &mut finish_reason,
1686 );
1687 match e {
1688 LlmStreamEvent::ToolCalls(calls) => {
1689 assert_eq!(calls.len(), 1);
1690 assert_eq!(calls[0].id, "call_1");
1691 assert_eq!(calls[0].name, "read_file");
1692 assert_eq!(calls[0].arguments, json!({"path": "Cargo.toml"}));
1693 }
1694 other => panic!("expected ToolCalls, got {:?}", other),
1695 }
1696 assert_eq!(finish_reason.as_deref(), Some("tool_calls"));
1697
1698 let e = process_stream_choice(
1701 &choice(r#"{"delta":{"content":""},"finish_reason":"tool_calls"}"#),
1702 &mut total_tokens,
1703 &mut acc,
1704 &mut finish_reason,
1705 );
1706 assert!(
1707 matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()),
1708 "tool calls must only be emitted once"
1709 );
1710 }
1711
1712 #[test]
1714 fn test_non_empty_content_is_emitted() {
1715 let mut total_tokens = 0u32;
1716 let mut acc = StreamToolCallAccumulator::new();
1717 let mut finish_reason: Option<String> = None;
1718
1719 let e = process_stream_choice(
1720 &choice(r#"{"delta":{"content":"hello"},"finish_reason":null}"#),
1721 &mut total_tokens,
1722 &mut acc,
1723 &mut finish_reason,
1724 );
1725 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s == "hello"));
1726 assert_eq!(total_tokens, 1);
1727 }
1728
1729 #[test]
1733 fn test_tool_call_arguments_accumulate_across_many_chunks() {
1734 let mut total_tokens = 0u32;
1735 let mut acc = StreamToolCallAccumulator::new();
1736 let mut finish_reason: Option<String> = None;
1737
1738 process_stream_choice(
1740 &choice(
1741 r#"{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"write_file","arguments":""}}]},"finish_reason":null}"#,
1742 ),
1743 &mut total_tokens,
1744 &mut acc,
1745 &mut finish_reason,
1746 );
1747
1748 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}"#;
1749 assert!(payload.chars().count() > 100);
1750
1751 let mut expected = String::new();
1753 for ch in payload.chars() {
1754 let frag = ch.to_string();
1755 let chunk = json!({
1756 "delta": {"tool_calls": [{"index": 0, "function": {"arguments": frag}}]},
1757 "finish_reason": null
1758 })
1759 .to_string();
1760 process_stream_choice(
1761 &choice(&chunk),
1762 &mut total_tokens,
1763 &mut acc,
1764 &mut finish_reason,
1765 );
1766 expected.push_str(&frag);
1767 }
1768
1769 let e = process_stream_choice(
1775 &choice(r#"{"delta":{},"finish_reason":"tool_calls"}"#),
1776 &mut total_tokens,
1777 &mut acc,
1778 &mut finish_reason,
1779 );
1780 match e {
1781 LlmStreamEvent::ToolCalls(calls) => {
1782 assert_eq!(calls.len(), 1);
1783 assert_eq!(calls[0].id, "call_1");
1784 assert_eq!(
1785 calls[0].arguments,
1786 serde_json::from_str::<serde_json::Value>(payload).unwrap()
1787 );
1788 }
1789 other => panic!("expected ToolCalls, got {:?}", other),
1790 }
1791 }
1792
1793 #[test]
1796 fn test_finish_chunk_without_content_emits_tool_calls() {
1797 let mut total_tokens = 0u32;
1798 let mut acc = StreamToolCallAccumulator::new();
1799 let mut finish_reason: Option<String> = None;
1800
1801 process_stream_choice(
1802 &choice(
1803 r#"{"delta":{"tool_calls":[{"index":0,"id":"call_9","function":{"name":"list_dir","arguments":"{}"}}]},"finish_reason":null}"#,
1804 ),
1805 &mut total_tokens,
1806 &mut acc,
1807 &mut finish_reason,
1808 );
1809
1810 let e = process_stream_choice(
1811 &choice(r#"{"delta":{},"finish_reason":"tool_calls"}"#),
1812 &mut total_tokens,
1813 &mut acc,
1814 &mut finish_reason,
1815 );
1816 match e {
1817 LlmStreamEvent::ToolCalls(calls) => {
1818 assert_eq!(calls.len(), 1);
1819 assert_eq!(calls[0].name, "list_dir");
1820 }
1821 other => panic!("expected ToolCalls, got {:?}", other),
1822 }
1823 }
1824
1825 fn seeded_acc(id: &str, name: &str, arguments: &str) -> StreamToolCallAccumulator {
1829 let mut acc = StreamToolCallAccumulator::new();
1830 acc.apply_indexed_delta(0, Some(id), Some(name), Some(arguments));
1831 acc
1832 }
1833
1834 #[test]
1838 fn test_take_pending_tool_calls_flushes_then_drains_without_finish_reason() {
1839 let mut acc = seeded_acc("call_1", "read_file", r#"{"path":"Cargo.toml"}"#);
1840
1841 match take_pending_tool_calls(&mut acc, None) {
1842 Some(LlmStreamEvent::ToolCalls(calls)) => {
1843 assert_eq!(calls.len(), 1);
1844 assert_eq!(calls[0].name, "read_file");
1845 assert_eq!(calls[0].arguments, json!({"path": "Cargo.toml"}));
1846 }
1847 other => panic!("expected ToolCalls, got {:?}", other),
1848 }
1849 assert!(acc.is_empty(), "accumulator must be drained after flush");
1850 assert!(take_pending_tool_calls(&mut acc, None).is_none());
1851 }
1852
1853 #[test]
1854 fn test_take_pending_tool_calls_discards_non_tool_finish_reason() {
1855 let mut acc = seeded_acc("call_cut", "read_file", r#"{"path":"#);
1856
1857 assert!(take_pending_tool_calls(&mut acc, Some("length")).is_none());
1858 assert!(
1859 acc.is_empty(),
1860 "discarded unsafe fallback calls must still drain the accumulator"
1861 );
1862 }
1863
1864 #[test]
1865 fn test_take_pending_tool_calls_rejects_malformed_fallback_arguments() {
1866 let mut acc = seeded_acc("call_cut", "read_file", r#"{"path":"#);
1867
1868 assert!(take_pending_tool_calls(&mut acc, None).is_none());
1869 assert!(
1870 acc.is_empty(),
1871 "malformed fallback calls must be drained instead of re-emitted"
1872 );
1873 }
1874
1875 #[test]
1876 fn test_non_tool_finish_reason_leaves_pending_calls_for_done_discard() {
1877 let mut total_tokens = 0u32;
1878 let mut acc = StreamToolCallAccumulator::new();
1879 let mut finish_reason: Option<String> = None;
1880
1881 process_stream_choice(
1882 &choice(
1883 r#"{"delta":{"tool_calls":[{"index":0,"id":"call_cut","function":{"name":"read_file","arguments":"{\"path\":"}}]},"finish_reason":null}"#,
1884 ),
1885 &mut total_tokens,
1886 &mut acc,
1887 &mut finish_reason,
1888 );
1889
1890 let e = process_stream_choice(
1891 &choice(r#"{"delta":{},"finish_reason":"length"}"#),
1892 &mut total_tokens,
1893 &mut acc,
1894 &mut finish_reason,
1895 );
1896
1897 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()));
1898 assert_eq!(finish_reason.as_deref(), Some("length"));
1899 assert!(take_pending_tool_calls(&mut acc, finish_reason.as_deref()).is_none());
1900 assert!(acc.is_empty());
1901 }
1902
1903 #[test]
1904 fn drop_orphaned_tool_messages_removes_unmatched_tool_results() {
1905 use crate::driver_registry::LlmMessageContent;
1906
1907 let messages = vec![
1908 LlmMessage::text(LlmMessageRole::User, "hello"),
1909 LlmMessage {
1910 role: LlmMessageRole::Tool,
1911 content: LlmMessageContent::Text("result".to_string()),
1912 tool_calls: None,
1913 tool_call_id: Some("call_trimmed".to_string()),
1914 phase: None,
1915 thinking: None,
1916 thinking_signature: None,
1917 },
1918 ];
1919 let filtered = drop_orphaned_tool_messages(&messages);
1920 assert_eq!(filtered.len(), 1);
1921 assert_eq!(filtered[0].role, LlmMessageRole::User);
1922 }
1923
1924 #[test]
1925 fn drop_orphaned_tool_messages_keeps_matched_tool_results() {
1926 use crate::driver_registry::LlmMessageContent;
1927 use crate::tool_types::ToolCall;
1928
1929 let messages = vec![
1930 LlmMessage {
1931 role: LlmMessageRole::Assistant,
1932 content: LlmMessageContent::Text(String::new()),
1933 tool_calls: Some(vec![ToolCall {
1934 id: "call_1".to_string(),
1935 name: "read_file".to_string(),
1936 arguments: json!({}),
1937 }]),
1938 tool_call_id: None,
1939 phase: None,
1940 thinking: None,
1941 thinking_signature: None,
1942 },
1943 LlmMessage {
1944 role: LlmMessageRole::Tool,
1945 content: LlmMessageContent::Text("file content".to_string()),
1946 tool_calls: None,
1947 tool_call_id: Some("call_1".to_string()),
1948 phase: None,
1949 thinking: None,
1950 thinking_signature: None,
1951 },
1952 ];
1953 let filtered = drop_orphaned_tool_messages(&messages);
1954 assert_eq!(filtered.len(), 2);
1955 }
1956}