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 retries_consumed: u32,
277 ) -> Result<(reqwest::Response, RetryMetadata)> {
278 let last_error: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
279 let mut retry_config = self.retry_config.clone();
280 retry_config.max_retries = retry_config.max_retries.saturating_sub(retries_consumed);
281
282 retry_request(
283 &retry_config,
284 "OpenAIProtocolDriver",
285 || async {
286 let request_builder = self.client.post(&self.api_url);
290 let request_builder = match &self.auth_provider {
291 Some(provider) => {
292 let (name, value) =
293 provider.auth_header().await.map_err(SendOutcome::Fatal)?;
294 request_builder.header(name, value)
295 }
296 None => apply_openai_api_auth(request_builder, &self.api_url, &self.api_key),
297 };
298 request_builder
299 .header("Content-Type", "application/json")
300 .json(request)
301 .send()
302 .await
303 .map_err(SendOutcome::Send)
304 },
305 |response, attempts, can_retry| {
306 let last_error = Arc::clone(&last_error);
307 let model = model.to_string();
308 async move {
309 let status = response.status();
310
311 if can_retry {
312 let rate_limit_info = if is_rate_limit_status(status) {
314 Some(RateLimitInfo::from_openai_headers(response.headers()))
315 } else {
316 None
317 };
318
319 let error_text = response.text().await.unwrap_or_default();
320
321 if is_openai_request_too_large(status, &error_text) {
323 return RetryDecision::Terminal(AgentLoopError::request_too_large(
324 format!("OpenAI API error ({}): {}", status, error_text),
325 ));
326 }
327
328 if is_provider_quota_message(&error_text) {
331 return RetryDecision::Terminal(AgentLoopError::llm_kind(
332 LlmErrorKind::QuotaExhausted,
333 format!("OpenAI API error ({}): {}", status, error_text),
334 ));
335 }
336
337 let wait = rate_limit_info
338 .as_ref()
339 .map(|info| info.recommended_wait(&self.retry_config, attempts))
340 .unwrap_or_else(|| self.retry_config.calculate_backoff(attempts));
341
342 *last_error.lock().unwrap() = Some(error_text);
343 return RetryDecision::Retry {
344 wait,
345 rate_limit_info,
346 };
347 }
348
349 let error_text = response.text().await.unwrap_or_default();
351 let error_msg = format!("OpenAI API error ({}): {}", status, error_text);
352
353 if is_openai_model_not_found(status, &error_text) {
355 return RetryDecision::Terminal(AgentLoopError::model_not_available(model));
356 }
357
358 if is_openai_request_too_large(status, &error_text) {
360 return RetryDecision::Terminal(AgentLoopError::request_too_large(
361 error_msg,
362 ));
363 }
364
365 let kind = LlmErrorKind::from_provider_status(status.as_u16(), &error_text);
368
369 if attempts > 0 {
370 return RetryDecision::Terminal(AgentLoopError::llm_kind(
371 kind,
372 format!(
373 "{} (after {} retries, last error: {})",
374 error_msg,
375 attempts,
376 last_error.lock().unwrap().take().unwrap_or_default()
377 ),
378 ));
379 }
380
381 RetryDecision::Terminal(AgentLoopError::llm_kind(kind, error_msg))
382 }
383 },
384 |e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
385 )
386 .await
387 }
388
389 fn convert_role(role: &LlmMessageRole) -> &'static str {
390 match role {
391 LlmMessageRole::System => "system",
392 LlmMessageRole::User => "user",
393 LlmMessageRole::Assistant => "assistant",
394 LlmMessageRole::Tool => "tool",
395 }
396 }
397
398 fn convert_message(msg: &LlmMessage) -> OpenAiMessage {
399 let content = match &msg.content {
400 LlmMessageContent::Text(text) => OpenAiContent::Text(text.clone()),
401 LlmMessageContent::Parts(parts) => {
402 let openai_parts: Vec<OpenAiContentPart> = parts
403 .iter()
404 .map(|part| match part {
405 LlmContentPart::Text { text } => OpenAiContentPart::Text {
406 r#type: "text".to_string(),
407 text: text.clone(),
408 },
409 LlmContentPart::Image { url } => OpenAiContentPart::ImageUrl {
410 r#type: "image_url".to_string(),
411 image_url: OpenAiImageUrl { url: url.clone() },
412 },
413 LlmContentPart::Audio { url } => OpenAiContentPart::InputAudio {
414 r#type: "input_audio".to_string(),
415 input_audio: OpenAiInputAudio {
416 data: url.clone(),
417 format: "wav".to_string(),
418 },
419 },
420 })
421 .collect();
422 OpenAiContent::Parts(openai_parts)
423 }
424 };
425
426 let tool_calls = if msg.role == LlmMessageRole::Assistant {
428 msg.tool_calls.as_ref().map(|calls| {
429 calls
430 .iter()
431 .map(|tc| OpenAiToolCall {
432 id: tc.id.clone(),
433 r#type: "function".to_string(),
434 function: OpenAiFunctionCall {
435 name: tc.name.clone(),
436 arguments: serde_json::to_string(&tc.arguments).unwrap_or_default(),
437 },
438 })
439 .collect()
440 })
441 } else {
442 None
443 };
444
445 OpenAiMessage {
446 role: Self::convert_role(&msg.role).to_string(),
447 content: Some(content),
448 tool_calls,
449 tool_call_id: msg.tool_call_id.clone(),
450 }
451 }
452
453 fn convert_tools(tools: &[ToolDefinition]) -> Vec<OpenAiTool> {
454 tools
455 .iter()
456 .map(|tool| OpenAiTool {
457 r#type: "function".to_string(),
458 function: OpenAiFunction {
459 name: tool.name().to_string(),
460 description: tool.description().to_string(),
461 parameters: crate::tool_schema_compat::sanitize_openai_tool_schema(
462 tool.parameters(),
463 ),
464 },
465 })
466 .collect()
467 }
468}
469
470fn drop_orphaned_tool_messages(messages: &[LlmMessage]) -> Vec<LlmMessage> {
474 use std::collections::HashSet;
475
476 let visible_call_ids: HashSet<&str> = messages
477 .iter()
478 .filter(|m| m.role == LlmMessageRole::Assistant)
479 .flat_map(|m| m.tool_calls.iter().flatten())
480 .map(|tc| tc.id.as_str())
481 .collect();
482
483 if visible_call_ids.is_empty() {
484 return messages
485 .iter()
486 .filter(|m| m.role != LlmMessageRole::Tool)
487 .cloned()
488 .collect();
489 }
490
491 messages
492 .iter()
493 .filter(|m| {
494 if m.role == LlmMessageRole::Tool {
495 return m
496 .tool_call_id
497 .as_deref()
498 .is_none_or(|id| visible_call_ids.contains(id));
499 }
500 true
501 })
502 .cloned()
503 .collect()
504}
505
506#[async_trait]
507impl ChatDriver for OpenAIProtocolChatDriver {
508 async fn chat_completion_stream(
509 &self,
510 messages: Vec<LlmMessage>,
511 config: &LlmCallConfig,
512 ) -> Result<LlmResponseStream> {
513 let messages = drop_orphaned_tool_messages(&messages);
517 let openai_messages: Vec<OpenAiMessage> =
518 messages.iter().map(Self::convert_message).collect();
519
520 let tools = if config.tools.is_empty() {
521 None
522 } else {
523 Some(Self::convert_tools(&config.tools))
524 };
525
526 let metadata = if config.metadata.is_empty() {
528 None
529 } else {
530 Some(config.metadata.clone())
531 };
532
533 let request = OpenAiRequest {
534 model: config.model.clone(),
535 messages: openai_messages,
536 temperature: config.temperature,
537 max_tokens: config.max_tokens,
538 stream: true,
539 stream_options: Some(OpenAiStreamOptions {
540 include_usage: true,
541 }),
542 tools,
543 parallel_tool_calls: config
544 .resolved_parallel_tool_calls(self.supports_parallel_tool_calls(&config.model)),
545 reasoning_effort: config
547 .reasoning_effort
548 .as_ref()
549 .filter(|e| !e.eq_ignore_ascii_case("none"))
550 .cloned(),
551 service_tier: config.speed.clone(),
552 verbosity: config.verbosity.clone(),
553 metadata,
554 };
555
556 let (event_stream, retry_metadata) =
562 connect_sse_with_reconnect(&self.retry_config, "OpenAIProtocolDriver", |attempts| {
563 self.send_chat_completion_request(&request, &config.model, attempts)
564 })
565 .await?;
566
567 let model = config.model.clone();
568 let total_tokens = Arc::new(Mutex::new(0u32));
569 let prompt_tokens = Arc::new(Mutex::new(0u32));
570 let cache_read_tokens = Arc::new(Mutex::new(Option::<u32>::None));
571 let provider_cost_usd = Arc::new(Mutex::new(Option::<f64>::None));
574 let accumulated_tool_calls = Arc::new(Mutex::new(StreamToolCallAccumulator::new()));
575 let finish_reason = Arc::new(Mutex::new(Option::<String>::None));
576 let response_id = Arc::new(Mutex::new(Option::<String>::None));
579 let shared_retry_metadata = if retry_metadata.had_retries() {
581 Some(Arc::new(retry_metadata))
582 } else {
583 None
584 };
585
586 let converted_stream: LlmResponseStream = Box::pin(
590 event_stream
591 .then(move |result| {
592 let model = model.clone();
593 let total_tokens = Arc::clone(&total_tokens);
594 let prompt_tokens = Arc::clone(&prompt_tokens);
595 let cache_read_tokens = Arc::clone(&cache_read_tokens);
596 let provider_cost_usd = Arc::clone(&provider_cost_usd);
597 let accumulated_tool_calls = Arc::clone(&accumulated_tool_calls);
598 let finish_reason = Arc::clone(&finish_reason);
599 let response_id = Arc::clone(&response_id);
600 let retry_metadata_for_done = shared_retry_metadata.clone();
601
602 async move {
603 let event = match result {
604 Ok(event) => event,
605 Err(e) => {
606 return vec![Ok(LlmStreamEvent::Error(
607 format!("Stream error: {}", e).into(),
608 ))];
609 }
610 };
611
612 if event.data == "[DONE]" {
613 let output_tokens = *total_tokens.lock().unwrap();
614 let input_tokens = *prompt_tokens.lock().unwrap();
615 let cached = *cache_read_tokens.lock().unwrap();
616 let cost = *provider_cost_usd.lock().unwrap();
617 let resp_id = response_id.lock().unwrap().clone();
618 let mut reason = finish_reason.lock().unwrap().clone();
619
620 let mut events = Vec::new();
621
622 {
631 let mut acc = accumulated_tool_calls.lock().unwrap();
632 if let Some(event) =
633 take_pending_tool_calls(&mut acc, reason.as_deref())
634 {
635 events.push(Ok(event));
636 reason.get_or_insert_with(|| "tool_calls".to_string());
637 }
638 }
639
640 events.push(Ok(LlmStreamEvent::Done(Box::new(
641 LlmCompletionMetadata {
642 total_tokens: Some(input_tokens + output_tokens),
645 prompt_tokens: Some(disjoint_prompt_tokens(
646 input_tokens,
647 cached,
648 )),
649 completion_tokens: Some(output_tokens),
650 cache_read_tokens: cached,
651 cache_creation_tokens: None,
652 provider_cost_usd: cost,
653 model: Some(model),
654 finish_reason: reason.or_else(|| Some("stop".to_string())),
655 retry_metadata: retry_metadata_for_done
656 .map(|arc| (*arc).clone()),
657 response_id: resp_id,
658 phase: None,
659 },
660 ))));
661
662 return events;
663 }
664
665 match serde_json::from_str::<OpenAiStreamChunk>(&event.data) {
666 Ok(chunk) => {
667 if let Some(id) = &chunk.id {
672 let mut rid = response_id.lock().unwrap();
673 if rid.is_none() {
674 *rid = Some(id.clone());
675 }
676 }
677
678 if let Some(usage) = &chunk.usage {
680 if let Some(pt) = usage.prompt_tokens {
681 *prompt_tokens.lock().unwrap() = pt;
682 }
683 if let Some(ct) = usage.completion_tokens {
684 *total_tokens.lock().unwrap() = ct;
685 }
686 if let Some(details) = &usage.prompt_tokens_details
688 && details.cached_tokens.is_some()
689 {
690 *cache_read_tokens.lock().unwrap() = details.cached_tokens;
691 }
692 if usage.cost.is_some() {
695 *provider_cost_usd.lock().unwrap() = usage.cost;
696 }
697 }
698
699 if let Some(choice) = chunk.choices.first() {
700 let mut tt = total_tokens.lock().unwrap();
701 let mut acc = accumulated_tool_calls.lock().unwrap();
702 let mut fr = finish_reason.lock().unwrap();
703 let stream_event =
704 process_stream_choice(choice, &mut tt, &mut acc, &mut fr);
705 return vec![Ok(stream_event)];
706 }
707 vec![Ok(LlmStreamEvent::TextDelta(String::new()))]
708 }
709 Err(e) => vec![Ok(LlmStreamEvent::Error(
710 format!("Failed to parse chunk: {}", e).into(),
711 ))],
712 }
713 }
714 })
715 .flat_map(futures::stream::iter),
716 );
717
718 Ok(converted_stream)
719 }
720
721 fn supports_parallel_tool_calls(&self, _model: &str) -> bool {
725 true
726 }
727}
728
729impl std::fmt::Debug for OpenAIProtocolChatDriver {
730 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
731 f.debug_struct("OpenAIProtocolChatDriver")
732 .field("api_url", &self.api_url)
733 .field("api_key", &"[REDACTED]")
734 .finish()
735 }
736}
737
738pub fn is_openai_model_not_found(status: reqwest::StatusCode, error_text: &str) -> bool {
749 let error_lower = error_text.to_lowercase();
750
751 if status == reqwest::StatusCode::NOT_FOUND
753 || status == reqwest::StatusCode::BAD_REQUEST
754 || status == reqwest::StatusCode::FORBIDDEN
755 {
756 if error_lower.contains("model_not_found") {
758 return true;
759 }
760 }
761
762 if status == reqwest::StatusCode::NOT_FOUND {
764 if error_lower.contains("does not exist") {
765 return true;
766 }
767 if error_lower.contains("model") && error_lower.contains("not found") {
768 return true;
769 }
770 }
771
772 false
773}
774
775pub fn is_openai_request_too_large(status: reqwest::StatusCode, error_text: &str) -> bool {
782 let error_lower = error_text.to_lowercase();
783
784 if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
786 if error_lower.contains("request too large") {
788 return true;
789 }
790 if error_lower.contains("tokens") && error_lower.contains("limit") {
792 return true;
793 }
794 }
795
796 if status == reqwest::StatusCode::BAD_REQUEST {
798 if error_lower.contains("context_length_exceeded") {
800 return true;
801 }
802 if error_lower.contains("maximum context length") {
804 return true;
805 }
806 }
807
808 if error_lower.contains("tokens must be reduced")
810 || error_lower.contains("reduce the length")
811 || error_lower.contains("input is too long")
812 {
813 return true;
814 }
815
816 false
817}
818
819#[derive(Debug, Serialize)]
824struct OpenAiRequest {
825 model: String,
826 messages: Vec<OpenAiMessage>,
827 #[serde(skip_serializing_if = "Option::is_none")]
828 temperature: Option<f32>,
829 #[serde(skip_serializing_if = "Option::is_none")]
830 max_tokens: Option<u32>,
831 stream: bool,
832 #[serde(skip_serializing_if = "Option::is_none")]
834 stream_options: Option<OpenAiStreamOptions>,
835 #[serde(skip_serializing_if = "Option::is_none")]
836 tools: Option<Vec<OpenAiTool>>,
837 #[serde(skip_serializing_if = "Option::is_none")]
840 parallel_tool_calls: Option<bool>,
841 #[serde(skip_serializing_if = "Option::is_none")]
842 reasoning_effort: Option<String>,
843 #[serde(skip_serializing_if = "Option::is_none")]
846 service_tier: Option<String>,
847 #[serde(skip_serializing_if = "Option::is_none")]
851 verbosity: Option<String>,
852 #[serde(skip_serializing_if = "Option::is_none")]
855 metadata: Option<std::collections::HashMap<String, String>>,
856}
857
858#[derive(Debug, Serialize)]
859struct OpenAiStreamOptions {
860 include_usage: bool,
861}
862
863#[derive(Debug, Serialize, Deserialize)]
864#[serde(untagged)]
865enum OpenAiContent {
866 Text(String),
867 Parts(Vec<OpenAiContentPart>),
868}
869
870#[derive(Debug, Serialize, Deserialize)]
871#[serde(untagged)]
872enum OpenAiContentPart {
873 Text {
874 r#type: String,
875 text: String,
876 },
877 ImageUrl {
878 r#type: String,
879 image_url: OpenAiImageUrl,
880 },
881 InputAudio {
882 r#type: String,
883 input_audio: OpenAiInputAudio,
884 },
885}
886
887#[derive(Debug, Serialize, Deserialize)]
888struct OpenAiImageUrl {
889 url: String,
890}
891
892#[derive(Debug, Serialize, Deserialize)]
893struct OpenAiInputAudio {
894 data: String,
895 format: String,
896}
897
898#[derive(Debug, Serialize, Deserialize)]
899struct OpenAiMessage {
900 role: String,
901 #[serde(skip_serializing_if = "Option::is_none")]
902 content: Option<OpenAiContent>,
903 #[serde(skip_serializing_if = "Option::is_none")]
904 tool_calls: Option<Vec<OpenAiToolCall>>,
905 #[serde(skip_serializing_if = "Option::is_none")]
906 tool_call_id: Option<String>,
907}
908
909#[derive(Debug, Serialize, Deserialize)]
910struct OpenAiTool {
911 r#type: String,
912 function: OpenAiFunction,
913}
914
915#[derive(Debug, Serialize, Deserialize)]
916struct OpenAiFunction {
917 name: String,
918 description: String,
919 parameters: Value,
920}
921
922#[derive(Debug, Serialize, Deserialize)]
923struct OpenAiToolCall {
924 id: String,
925 r#type: String,
926 function: OpenAiFunctionCall,
927}
928
929#[derive(Debug, Serialize, Deserialize)]
930struct OpenAiFunctionCall {
931 name: String,
932 arguments: String,
933}
934
935#[derive(Debug, Deserialize)]
936#[allow(dead_code)] struct OpenAiStreamChunk {
938 #[serde(default)]
940 id: Option<String>,
941 #[serde(default)]
943 model: Option<String>,
944 choices: Vec<OpenAiStreamChoice>,
945 #[serde(default)]
946 usage: Option<OpenAiUsage>,
947}
948
949#[derive(Debug, Deserialize)]
950struct OpenAiUsage {
951 prompt_tokens: Option<u32>,
952 completion_tokens: Option<u32>,
953 #[serde(default)]
955 prompt_tokens_details: Option<OpenAiPromptTokensDetails>,
956 #[serde(default)]
959 cost: Option<f64>,
960}
961
962#[derive(Debug, Deserialize, Default)]
963struct OpenAiPromptTokensDetails {
964 #[serde(default)]
966 cached_tokens: Option<u32>,
967}
968
969#[derive(Debug, Deserialize)]
970struct OpenAiStreamChoice {
971 delta: OpenAiDelta,
972 #[serde(default)]
973 finish_reason: Option<String>,
974}
975
976#[derive(Debug, Deserialize)]
977struct OpenAiDelta {
978 #[serde(default)]
979 content: Option<String>,
980 #[serde(default)]
981 tool_calls: Option<Vec<OpenAiStreamToolCall>>,
982}
983
984#[derive(Debug, Deserialize)]
985struct OpenAiStreamToolCall {
986 index: u32,
987 id: Option<String>,
988 function: Option<OpenAiStreamFunction>,
989}
990
991#[derive(Debug, Deserialize)]
992struct OpenAiStreamFunction {
993 name: Option<String>,
994 arguments: Option<String>,
995}
996
997fn take_pending_tool_calls(
1008 accumulated_tool_calls: &mut StreamToolCallAccumulator,
1009 finish_reason: Option<&str>,
1010) -> Option<LlmStreamEvent> {
1011 if accumulated_tool_calls.is_empty() {
1012 return None;
1013 }
1014
1015 if !matches!(finish_reason, None | Some("tool_calls")) {
1018 let _ = accumulated_tool_calls.take_finalized();
1019 return None;
1020 }
1021
1022 let calls = accumulated_tool_calls.take_pending_strict();
1023 if calls.is_empty() {
1024 None
1025 } else {
1026 Some(LlmStreamEvent::ToolCalls(calls))
1027 }
1028}
1029
1030fn process_stream_choice(
1040 choice: &OpenAiStreamChoice,
1041 total_tokens: &mut u32,
1042 accumulated_tool_calls: &mut StreamToolCallAccumulator,
1043 finish_reason: &mut Option<String>,
1044) -> LlmStreamEvent {
1045 if let Some(tool_calls) = &choice.delta.tool_calls {
1049 for tc in tool_calls {
1050 accumulated_tool_calls.apply_indexed_delta(
1051 tc.index,
1052 tc.id.as_deref(),
1053 tc.function.as_ref().and_then(|f| f.name.as_deref()),
1054 tc.function.as_ref().and_then(|f| f.arguments.as_deref()),
1055 );
1056 }
1057 return LlmStreamEvent::TextDelta(String::new());
1058 }
1059
1060 if let Some(content) = &choice.delta.content
1063 && !content.is_empty()
1064 {
1065 *total_tokens += 1;
1066 return LlmStreamEvent::TextDelta(content.clone());
1067 }
1068
1069 if let Some(fr) = &choice.finish_reason {
1073 *finish_reason = Some(fr.clone());
1074
1075 if fr == "tool_calls" && !accumulated_tool_calls.is_empty() {
1076 return LlmStreamEvent::ToolCalls(accumulated_tool_calls.take_finalized());
1077 }
1078 }
1079
1080 LlmStreamEvent::TextDelta(String::new())
1081}
1082
1083#[cfg(test)]
1088mod tests {
1089 use super::*;
1090 use serde_json::json;
1091
1092 #[test]
1093 fn test_convert_message_preserves_multiple_system_messages() {
1094 let messages = [
1099 LlmMessage::text(LlmMessageRole::System, "A"),
1100 LlmMessage::text(LlmMessageRole::User, "hi"),
1101 LlmMessage::text(LlmMessageRole::System, "B"),
1102 ];
1103 let converted: Vec<OpenAiMessage> = messages
1104 .iter()
1105 .map(OpenAIProtocolChatDriver::convert_message)
1106 .collect();
1107 let system_texts: Vec<String> = converted
1108 .iter()
1109 .filter(|m| m.role == "system")
1110 .filter_map(|m| match &m.content {
1111 Some(OpenAiContent::Text(t)) => Some(t.clone()),
1112 _ => None,
1113 })
1114 .collect();
1115 assert_eq!(system_texts, vec!["A".to_string(), "B".to_string()]);
1116 }
1117
1118 #[test]
1119 fn test_driver_with_api_key() {
1120 let driver = OpenAIProtocolChatDriver::new("test-key");
1121 assert!(format!("{:?}", driver).contains("OpenAIProtocolChatDriver"));
1122 }
1123
1124 #[test]
1125 fn test_driver_with_base_url() {
1126 let driver = OpenAIProtocolChatDriver::with_base_url(
1127 "test-key",
1128 "https://custom.api.com/v1/completions",
1129 );
1130 assert!(format!("{:?}", driver).contains("OpenAIProtocolChatDriver"));
1131 assert_eq!(driver.api_url(), "https://custom.api.com/v1/completions");
1132 }
1133
1134 #[test]
1135 fn test_is_azure_openai_api_url() {
1136 assert!(is_azure_openai_api_url(
1137 "https://example.openai.azure.com/openai/v1/chat/completions"
1138 ));
1139 assert!(is_azure_openai_api_url(
1140 "https://example.services.ai.azure.com/openai/v1/responses"
1141 ));
1142 assert!(!is_azure_openai_api_url(
1143 "https://api.openai.com/v1/chat/completions"
1144 ));
1145 }
1146
1147 #[test]
1148 fn test_request_includes_stream_options_for_usage() {
1149 let request = OpenAiRequest {
1152 verbosity: None,
1153 service_tier: None,
1154 model: "gpt-4o".to_string(),
1155 messages: vec![OpenAiMessage {
1156 role: "user".to_string(),
1157 content: Some(OpenAiContent::Text("Hello".to_string())),
1158 tool_calls: None,
1159 tool_call_id: None,
1160 }],
1161 temperature: None,
1162 max_tokens: None,
1163 stream: true,
1164 stream_options: Some(OpenAiStreamOptions {
1165 include_usage: true,
1166 }),
1167 tools: None,
1168 parallel_tool_calls: None,
1169 reasoning_effort: None,
1170 metadata: None,
1171 };
1172
1173 let json = serde_json::to_value(&request).unwrap();
1174 assert_eq!(json["stream"], true);
1175 assert_eq!(json["stream_options"]["include_usage"], true);
1176 }
1177
1178 #[test]
1179 fn test_request_includes_metadata() {
1180 let mut metadata = std::collections::HashMap::new();
1182 metadata.insert("session_id".to_string(), "session_abc123".to_string());
1183 metadata.insert("agent_id".to_string(), "agent_xyz789".to_string());
1184
1185 let request = OpenAiRequest {
1186 verbosity: None,
1187 service_tier: None,
1188 model: "gpt-4o".to_string(),
1189 messages: vec![OpenAiMessage {
1190 role: "user".to_string(),
1191 content: Some(OpenAiContent::Text("Hello".to_string())),
1192 tool_calls: None,
1193 tool_call_id: None,
1194 }],
1195 temperature: None,
1196 max_tokens: None,
1197 stream: true,
1198 stream_options: None,
1199 tools: None,
1200 parallel_tool_calls: None,
1201 reasoning_effort: None,
1202 metadata: Some(metadata),
1203 };
1204
1205 let json = serde_json::to_value(&request).unwrap();
1206 assert_eq!(json["metadata"]["session_id"], "session_abc123");
1207 assert_eq!(json["metadata"]["agent_id"], "agent_xyz789");
1208 }
1209
1210 #[test]
1211 fn test_usage_chunk_parsing() {
1212 let usage_chunk = r#"{
1215 "id": "chatcmpl-123",
1216 "object": "chat.completion.chunk",
1217 "created": 1234567890,
1218 "model": "gpt-4o",
1219 "choices": [],
1220 "usage": {
1221 "prompt_tokens": 150,
1222 "completion_tokens": 42,
1223 "total_tokens": 192
1224 }
1225 }"#;
1226
1227 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1228 assert!(chunk.usage.is_some());
1229 let usage = chunk.usage.unwrap();
1230 assert_eq!(usage.prompt_tokens, Some(150));
1231 assert_eq!(usage.completion_tokens, Some(42));
1232 }
1233
1234 #[test]
1235 fn test_usage_chunk_with_cached_tokens() {
1236 let usage_chunk = r#"{
1238 "id": "chatcmpl-123",
1239 "choices": [],
1240 "usage": {
1241 "prompt_tokens": 150,
1242 "completion_tokens": 42,
1243 "prompt_tokens_details": {
1244 "cached_tokens": 100
1245 }
1246 }
1247 }"#;
1248
1249 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1250 let usage = chunk.usage.unwrap();
1251 assert_eq!(usage.prompt_tokens, Some(150));
1252 assert_eq!(usage.completion_tokens, Some(42));
1253 assert!(usage.prompt_tokens_details.is_some());
1254 assert_eq!(
1255 usage.prompt_tokens_details.unwrap().cached_tokens,
1256 Some(100)
1257 );
1258 }
1259
1260 #[test]
1261 fn test_usage_chunk_with_openrouter_cost() {
1262 let usage_chunk = r#"{
1264 "id": "gen-123",
1265 "choices": [],
1266 "usage": {
1267 "prompt_tokens": 194,
1268 "completion_tokens": 2,
1269 "total_tokens": 196,
1270 "cost": 0.00095
1271 }
1272 }"#;
1273
1274 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1275 let usage = chunk.usage.unwrap();
1276 assert_eq!(usage.cost, Some(0.00095));
1277 }
1278
1279 #[test]
1280 fn test_usage_chunk_without_cost_defaults_none() {
1281 let usage_chunk = r#"{
1283 "id": "chatcmpl-123",
1284 "choices": [],
1285 "usage": { "prompt_tokens": 10, "completion_tokens": 5 }
1286 }"#;
1287
1288 let chunk: OpenAiStreamChunk = serde_json::from_str(usage_chunk).unwrap();
1289 assert_eq!(chunk.usage.unwrap().cost, None);
1290 }
1291
1292 #[test]
1293 fn test_chunk_id_is_captured() {
1294 let chunk_with_id: OpenAiStreamChunk =
1295 serde_json::from_str(r#"{"id":"gen-abc123","choices":[]}"#).unwrap();
1296 assert_eq!(chunk_with_id.id.as_deref(), Some("gen-abc123"));
1297
1298 let chunk_no_id: OpenAiStreamChunk = serde_json::from_str(r#"{"choices":[]}"#).unwrap();
1299 assert!(chunk_no_id.id.is_none());
1300 }
1301
1302 #[test]
1303 fn test_finish_reason_chunk_parsing() {
1304 let finish_chunk = r#"{
1306 "id": "chatcmpl-123",
1307 "choices": [{
1308 "index": 0,
1309 "delta": {},
1310 "finish_reason": "stop"
1311 }]
1312 }"#;
1313
1314 let chunk: OpenAiStreamChunk = serde_json::from_str(finish_chunk).unwrap();
1315 assert!(chunk.usage.is_none()); assert_eq!(chunk.choices.len(), 1);
1317 assert_eq!(chunk.choices[0].finish_reason, Some("stop".to_string()));
1318 }
1319
1320 #[test]
1325 fn test_is_openai_request_too_large_429_request_too_large() {
1326 let error = r#"{"error":{"message":"Request too large for gpt-4o in organization org-xxx on tokens per min (TPM): Limit 500000, Requested 538772."}}"#;
1327 assert!(is_openai_request_too_large(
1328 reqwest::StatusCode::TOO_MANY_REQUESTS,
1329 error
1330 ));
1331 }
1332
1333 #[test]
1334 fn test_is_openai_request_too_large_429_token_limit() {
1335 let error =
1336 r#"{"error":{"message":"tokens per min (TPM): Limit 500000, Requested 600000"}}"#;
1337 assert!(is_openai_request_too_large(
1338 reqwest::StatusCode::TOO_MANY_REQUESTS,
1339 error
1340 ));
1341 }
1342
1343 #[test]
1344 fn test_is_openai_request_too_large_400_context_length() {
1345 let error = r#"{"error":{"code":"context_length_exceeded","message":"This model's maximum context length is 128000 tokens."}}"#;
1346 assert!(is_openai_request_too_large(
1347 reqwest::StatusCode::BAD_REQUEST,
1348 error
1349 ));
1350 }
1351
1352 #[test]
1353 fn test_is_openai_request_too_large_400_max_context() {
1354 let error =
1355 r#"{"error":{"message":"This model's maximum context length is 128000 tokens"}}"#;
1356 assert!(is_openai_request_too_large(
1357 reqwest::StatusCode::BAD_REQUEST,
1358 error
1359 ));
1360 }
1361
1362 #[test]
1363 fn test_is_openai_request_too_large_tokens_must_be_reduced() {
1364 let error = r#"{"error":{"message":"The input or output tokens must be reduced"}}"#;
1365 assert!(is_openai_request_too_large(
1366 reqwest::StatusCode::BAD_REQUEST,
1367 error
1368 ));
1369 }
1370
1371 #[test]
1372 fn test_is_openai_request_too_large_false_for_other_errors() {
1373 let error = r#"{"error":{"message":"Rate limit exceeded: too many requests per minute"}}"#;
1375 assert!(!is_openai_request_too_large(
1376 reqwest::StatusCode::TOO_MANY_REQUESTS,
1377 error
1378 ));
1379
1380 let error = r#"{"error":{"message":"Internal server error"}}"#;
1382 assert!(!is_openai_request_too_large(
1383 reqwest::StatusCode::INTERNAL_SERVER_ERROR,
1384 error
1385 ));
1386
1387 let error = r#"{"error":{"message":"Invalid request"}}"#;
1389 assert!(!is_openai_request_too_large(
1390 reqwest::StatusCode::BAD_REQUEST,
1391 error
1392 ));
1393 }
1394
1395 #[test]
1400 fn test_is_openai_model_not_found_real_error() {
1401 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}}"#;
1403 assert!(is_openai_model_not_found(
1404 reqwest::StatusCode::NOT_FOUND,
1405 error
1406 ));
1407 }
1408
1409 #[test]
1410 fn test_is_openai_model_not_found_does_not_exist() {
1411 let error = r#"{"error":{"message":"The model 'fake-model' does not exist"}}"#;
1412 assert!(is_openai_model_not_found(
1413 reqwest::StatusCode::NOT_FOUND,
1414 error
1415 ));
1416 }
1417
1418 #[test]
1419 fn test_is_openai_model_not_found_generic_not_found() {
1420 let error = r#"{"error":{"message":"Model not found"}}"#;
1421 assert!(is_openai_model_not_found(
1422 reqwest::StatusCode::NOT_FOUND,
1423 error
1424 ));
1425 }
1426
1427 #[test]
1428 fn test_is_openai_model_not_found_400_with_model_not_found_code() {
1429 let error = r#"{"error":{"code":"model_not_found","message":"The requested model 'gpt-99' does not exist.","type":"invalid_request_error","param":"model"}}"#;
1431 assert!(is_openai_model_not_found(
1432 reqwest::StatusCode::BAD_REQUEST,
1433 error
1434 ));
1435 }
1436
1437 #[test]
1438 fn test_is_openai_model_not_found_false_for_non_model_error() {
1439 let error = r#"{"error":{"code":"invalid_request","message":"Some other error"}}"#;
1441 assert!(!is_openai_model_not_found(
1442 reqwest::StatusCode::BAD_REQUEST,
1443 error
1444 ));
1445 }
1446
1447 #[test]
1448 fn test_is_openai_model_not_found_false_for_other_404() {
1449 let error = r#"{"error":{"message":"Endpoint not found"}}"#;
1451 assert!(!is_openai_model_not_found(
1452 reqwest::StatusCode::NOT_FOUND,
1453 error
1454 ));
1455 }
1456
1457 #[test]
1458 fn test_is_openai_model_not_found_403_tier_gated_model() {
1459 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}}"#;
1462 assert!(is_openai_model_not_found(
1463 reqwest::StatusCode::FORBIDDEN,
1464 error
1465 ));
1466 }
1467
1468 #[test]
1469 fn test_is_openai_model_not_found_403_plain_auth_error_is_not_model_not_found() {
1470 let error = r#"{"error":{"message":"Invalid authentication credentials","type":"authentication_error"}}"#;
1473 assert!(!is_openai_model_not_found(
1474 reqwest::StatusCode::FORBIDDEN,
1475 error
1476 ));
1477 }
1478
1479 #[test]
1484 fn test_reasoning_effort_none_is_omitted() {
1485 let request = OpenAiRequest {
1488 verbosity: None,
1489 service_tier: None,
1490 model: "gpt-4o-mini".to_string(),
1491 messages: vec![OpenAiMessage {
1492 role: "user".to_string(),
1493 content: Some(OpenAiContent::Text("Hello".to_string())),
1494 tool_calls: None,
1495 tool_call_id: None,
1496 }],
1497 temperature: None,
1498 max_tokens: None,
1499 stream: true,
1500 stream_options: None,
1501 tools: None,
1502 parallel_tool_calls: None,
1503 reasoning_effort: Some("none".to_string())
1504 .as_ref()
1505 .filter(|e| !e.eq_ignore_ascii_case("none"))
1506 .cloned(),
1507 metadata: None,
1508 };
1509
1510 let json = serde_json::to_value(&request).unwrap();
1511 assert!(
1512 json.get("reasoning_effort").is_none(),
1513 "reasoning_effort should be omitted when effort is 'none'"
1514 );
1515 }
1516
1517 #[test]
1518 fn test_reasoning_effort_high_is_included() {
1519 let request = OpenAiRequest {
1520 verbosity: None,
1521 service_tier: None,
1522 model: "o3-mini".to_string(),
1523 messages: vec![OpenAiMessage {
1524 role: "user".to_string(),
1525 content: Some(OpenAiContent::Text("Hello".to_string())),
1526 tool_calls: None,
1527 tool_call_id: None,
1528 }],
1529 temperature: None,
1530 max_tokens: None,
1531 stream: true,
1532 stream_options: None,
1533 tools: None,
1534 parallel_tool_calls: None,
1535 reasoning_effort: Some("high".to_string())
1536 .as_ref()
1537 .filter(|e| !e.eq_ignore_ascii_case("none"))
1538 .cloned(),
1539 metadata: None,
1540 };
1541
1542 let json = serde_json::to_value(&request).unwrap();
1543 assert_eq!(json["reasoning_effort"], "high");
1544 }
1545
1546 #[test]
1550 fn test_request_serializes_parallel_tool_calls() {
1551 fn build(flag: Option<bool>) -> serde_json::Value {
1552 let request = OpenAiRequest {
1553 verbosity: None,
1554 service_tier: None,
1555 model: "gpt-4o-mini".to_string(),
1556 messages: vec![OpenAiMessage {
1557 role: "user".to_string(),
1558 content: Some(OpenAiContent::Text("Hello".to_string())),
1559 tool_calls: None,
1560 tool_call_id: None,
1561 }],
1562 temperature: None,
1563 max_tokens: None,
1564 stream: true,
1565 stream_options: None,
1566 tools: None,
1567 parallel_tool_calls: flag,
1568 reasoning_effort: None,
1569 metadata: None,
1570 };
1571 serde_json::to_value(&request).unwrap()
1572 }
1573
1574 assert!(build(None).get("parallel_tool_calls").is_none());
1576 assert_eq!(build(Some(true))["parallel_tool_calls"], true);
1578 assert_eq!(build(Some(false))["parallel_tool_calls"], false);
1579 }
1580
1581 #[test]
1584 fn test_request_serializes_service_tier() {
1585 fn build(tier: Option<&str>) -> serde_json::Value {
1586 let request = OpenAiRequest {
1587 service_tier: tier.map(str::to_string),
1588 verbosity: None,
1589 model: "gpt-4o-mini".to_string(),
1590 messages: vec![OpenAiMessage {
1591 role: "user".to_string(),
1592 content: Some(OpenAiContent::Text("Hello".to_string())),
1593 tool_calls: None,
1594 tool_call_id: None,
1595 }],
1596 temperature: None,
1597 max_tokens: None,
1598 stream: true,
1599 stream_options: None,
1600 tools: None,
1601 parallel_tool_calls: None,
1602 reasoning_effort: None,
1603 metadata: None,
1604 };
1605 serde_json::to_value(&request).unwrap()
1606 }
1607
1608 assert!(build(None).get("service_tier").is_none());
1609 assert_eq!(build(Some("flex"))["service_tier"], "flex");
1610 assert_eq!(build(Some("priority"))["service_tier"], "priority");
1611 }
1612
1613 #[test]
1616 fn test_request_serializes_verbosity() {
1617 fn build(verbosity: Option<&str>) -> serde_json::Value {
1618 let request = OpenAiRequest {
1619 service_tier: None,
1620 verbosity: verbosity.map(str::to_string),
1621 model: "gpt-5.6-sol".to_string(),
1622 messages: vec![OpenAiMessage {
1623 role: "user".to_string(),
1624 content: Some(OpenAiContent::Text("Hello".to_string())),
1625 tool_calls: None,
1626 tool_call_id: None,
1627 }],
1628 temperature: None,
1629 max_tokens: None,
1630 stream: true,
1631 stream_options: None,
1632 tools: None,
1633 parallel_tool_calls: None,
1634 reasoning_effort: None,
1635 metadata: None,
1636 };
1637 serde_json::to_value(&request).unwrap()
1638 }
1639
1640 assert!(build(None).get("verbosity").is_none());
1641 assert_eq!(build(Some("low"))["verbosity"], "low");
1642 assert_eq!(build(Some("high"))["verbosity"], "high");
1643 }
1644
1645 fn choice(json_str: &str) -> OpenAiStreamChoice {
1650 serde_json::from_str(json_str).unwrap()
1651 }
1652
1653 #[test]
1657 fn test_empty_content_finish_chunk_still_emits_tool_calls() {
1658 let mut total_tokens = 0u32;
1659 let mut acc = StreamToolCallAccumulator::new();
1660 let mut finish_reason: Option<String> = None;
1661
1662 let e = process_stream_choice(
1664 &choice(
1665 r#"{"delta":{"content":null,"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":""}}]},"finish_reason":null}"#,
1666 ),
1667 &mut total_tokens,
1668 &mut acc,
1669 &mut finish_reason,
1670 );
1671 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()));
1672
1673 let e = process_stream_choice(
1675 &choice(
1676 r#"{"delta":{"content":null,"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":\"Cargo.toml\"}"}}]},"finish_reason":null}"#,
1677 ),
1678 &mut total_tokens,
1679 &mut acc,
1680 &mut finish_reason,
1681 );
1682 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()));
1683
1684 let e = process_stream_choice(
1687 &choice(r#"{"delta":{"content":""},"finish_reason":"tool_calls"}"#),
1688 &mut total_tokens,
1689 &mut acc,
1690 &mut finish_reason,
1691 );
1692 match e {
1693 LlmStreamEvent::ToolCalls(calls) => {
1694 assert_eq!(calls.len(), 1);
1695 assert_eq!(calls[0].id, "call_1");
1696 assert_eq!(calls[0].name, "read_file");
1697 assert_eq!(calls[0].arguments, json!({"path": "Cargo.toml"}));
1698 }
1699 other => panic!("expected ToolCalls, got {:?}", other),
1700 }
1701 assert_eq!(finish_reason.as_deref(), Some("tool_calls"));
1702
1703 let e = process_stream_choice(
1706 &choice(r#"{"delta":{"content":""},"finish_reason":"tool_calls"}"#),
1707 &mut total_tokens,
1708 &mut acc,
1709 &mut finish_reason,
1710 );
1711 assert!(
1712 matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()),
1713 "tool calls must only be emitted once"
1714 );
1715 }
1716
1717 #[test]
1719 fn test_non_empty_content_is_emitted() {
1720 let mut total_tokens = 0u32;
1721 let mut acc = StreamToolCallAccumulator::new();
1722 let mut finish_reason: Option<String> = None;
1723
1724 let e = process_stream_choice(
1725 &choice(r#"{"delta":{"content":"hello"},"finish_reason":null}"#),
1726 &mut total_tokens,
1727 &mut acc,
1728 &mut finish_reason,
1729 );
1730 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s == "hello"));
1731 assert_eq!(total_tokens, 1);
1732 }
1733
1734 #[test]
1738 fn test_tool_call_arguments_accumulate_across_many_chunks() {
1739 let mut total_tokens = 0u32;
1740 let mut acc = StreamToolCallAccumulator::new();
1741 let mut finish_reason: Option<String> = None;
1742
1743 process_stream_choice(
1745 &choice(
1746 r#"{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"write_file","arguments":""}}]},"finish_reason":null}"#,
1747 ),
1748 &mut total_tokens,
1749 &mut acc,
1750 &mut finish_reason,
1751 );
1752
1753 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}"#;
1754 assert!(payload.chars().count() > 100);
1755
1756 let mut expected = String::new();
1758 for ch in payload.chars() {
1759 let frag = ch.to_string();
1760 let chunk = json!({
1761 "delta": {"tool_calls": [{"index": 0, "function": {"arguments": frag}}]},
1762 "finish_reason": null
1763 })
1764 .to_string();
1765 process_stream_choice(
1766 &choice(&chunk),
1767 &mut total_tokens,
1768 &mut acc,
1769 &mut finish_reason,
1770 );
1771 expected.push_str(&frag);
1772 }
1773
1774 let e = process_stream_choice(
1780 &choice(r#"{"delta":{},"finish_reason":"tool_calls"}"#),
1781 &mut total_tokens,
1782 &mut acc,
1783 &mut finish_reason,
1784 );
1785 match e {
1786 LlmStreamEvent::ToolCalls(calls) => {
1787 assert_eq!(calls.len(), 1);
1788 assert_eq!(calls[0].id, "call_1");
1789 assert_eq!(
1790 calls[0].arguments,
1791 serde_json::from_str::<serde_json::Value>(payload).unwrap()
1792 );
1793 }
1794 other => panic!("expected ToolCalls, got {:?}", other),
1795 }
1796 }
1797
1798 #[test]
1801 fn test_finish_chunk_without_content_emits_tool_calls() {
1802 let mut total_tokens = 0u32;
1803 let mut acc = StreamToolCallAccumulator::new();
1804 let mut finish_reason: Option<String> = None;
1805
1806 process_stream_choice(
1807 &choice(
1808 r#"{"delta":{"tool_calls":[{"index":0,"id":"call_9","function":{"name":"list_dir","arguments":"{}"}}]},"finish_reason":null}"#,
1809 ),
1810 &mut total_tokens,
1811 &mut acc,
1812 &mut finish_reason,
1813 );
1814
1815 let e = process_stream_choice(
1816 &choice(r#"{"delta":{},"finish_reason":"tool_calls"}"#),
1817 &mut total_tokens,
1818 &mut acc,
1819 &mut finish_reason,
1820 );
1821 match e {
1822 LlmStreamEvent::ToolCalls(calls) => {
1823 assert_eq!(calls.len(), 1);
1824 assert_eq!(calls[0].name, "list_dir");
1825 }
1826 other => panic!("expected ToolCalls, got {:?}", other),
1827 }
1828 }
1829
1830 fn seeded_acc(id: &str, name: &str, arguments: &str) -> StreamToolCallAccumulator {
1834 let mut acc = StreamToolCallAccumulator::new();
1835 acc.apply_indexed_delta(0, Some(id), Some(name), Some(arguments));
1836 acc
1837 }
1838
1839 #[test]
1843 fn test_take_pending_tool_calls_flushes_then_drains_without_finish_reason() {
1844 let mut acc = seeded_acc("call_1", "read_file", r#"{"path":"Cargo.toml"}"#);
1845
1846 match take_pending_tool_calls(&mut acc, None) {
1847 Some(LlmStreamEvent::ToolCalls(calls)) => {
1848 assert_eq!(calls.len(), 1);
1849 assert_eq!(calls[0].name, "read_file");
1850 assert_eq!(calls[0].arguments, json!({"path": "Cargo.toml"}));
1851 }
1852 other => panic!("expected ToolCalls, got {:?}", other),
1853 }
1854 assert!(acc.is_empty(), "accumulator must be drained after flush");
1855 assert!(take_pending_tool_calls(&mut acc, None).is_none());
1856 }
1857
1858 #[test]
1859 fn test_take_pending_tool_calls_discards_non_tool_finish_reason() {
1860 let mut acc = seeded_acc("call_cut", "read_file", r#"{"path":"#);
1861
1862 assert!(take_pending_tool_calls(&mut acc, Some("length")).is_none());
1863 assert!(
1864 acc.is_empty(),
1865 "discarded unsafe fallback calls must still drain the accumulator"
1866 );
1867 }
1868
1869 #[test]
1870 fn test_take_pending_tool_calls_rejects_malformed_fallback_arguments() {
1871 let mut acc = seeded_acc("call_cut", "read_file", r#"{"path":"#);
1872
1873 assert!(take_pending_tool_calls(&mut acc, None).is_none());
1874 assert!(
1875 acc.is_empty(),
1876 "malformed fallback calls must be drained instead of re-emitted"
1877 );
1878 }
1879
1880 #[test]
1881 fn test_non_tool_finish_reason_leaves_pending_calls_for_done_discard() {
1882 let mut total_tokens = 0u32;
1883 let mut acc = StreamToolCallAccumulator::new();
1884 let mut finish_reason: Option<String> = None;
1885
1886 process_stream_choice(
1887 &choice(
1888 r#"{"delta":{"tool_calls":[{"index":0,"id":"call_cut","function":{"name":"read_file","arguments":"{\"path\":"}}]},"finish_reason":null}"#,
1889 ),
1890 &mut total_tokens,
1891 &mut acc,
1892 &mut finish_reason,
1893 );
1894
1895 let e = process_stream_choice(
1896 &choice(r#"{"delta":{},"finish_reason":"length"}"#),
1897 &mut total_tokens,
1898 &mut acc,
1899 &mut finish_reason,
1900 );
1901
1902 assert!(matches!(e, LlmStreamEvent::TextDelta(s) if s.is_empty()));
1903 assert_eq!(finish_reason.as_deref(), Some("length"));
1904 assert!(take_pending_tool_calls(&mut acc, finish_reason.as_deref()).is_none());
1905 assert!(acc.is_empty());
1906 }
1907
1908 #[test]
1909 fn drop_orphaned_tool_messages_removes_unmatched_tool_results() {
1910 use crate::driver_registry::LlmMessageContent;
1911
1912 let messages = vec![
1913 LlmMessage::text(LlmMessageRole::User, "hello"),
1914 LlmMessage {
1915 role: LlmMessageRole::Tool,
1916 content: LlmMessageContent::Text("result".to_string()),
1917 tool_calls: None,
1918 tool_call_id: Some("call_trimmed".to_string()),
1919 phase: None,
1920 thinking: None,
1921 thinking_signature: None,
1922 },
1923 ];
1924 let filtered = drop_orphaned_tool_messages(&messages);
1925 assert_eq!(filtered.len(), 1);
1926 assert_eq!(filtered[0].role, LlmMessageRole::User);
1927 }
1928
1929 #[test]
1930 fn drop_orphaned_tool_messages_keeps_matched_tool_results() {
1931 use crate::driver_registry::LlmMessageContent;
1932 use crate::tool_types::ToolCall;
1933
1934 let messages = vec![
1935 LlmMessage {
1936 role: LlmMessageRole::Assistant,
1937 content: LlmMessageContent::Text(String::new()),
1938 tool_calls: Some(vec![ToolCall {
1939 id: "call_1".to_string(),
1940 name: "read_file".to_string(),
1941 arguments: json!({}),
1942 }]),
1943 tool_call_id: None,
1944 phase: None,
1945 thinking: None,
1946 thinking_signature: None,
1947 },
1948 LlmMessage {
1949 role: LlmMessageRole::Tool,
1950 content: LlmMessageContent::Text("file content".to_string()),
1951 tool_calls: None,
1952 tool_call_id: Some("call_1".to_string()),
1953 phase: None,
1954 thinking: None,
1955 thinking_signature: None,
1956 },
1957 ];
1958 let filtered = drop_orphaned_tool_messages(&messages);
1959 assert_eq!(filtered.len(), 2);
1960 }
1961}