Skip to main content

oxicode_ai/providers/
azure.rs

1//! Azure OpenAI provider implementation
2
3use bytes::Bytes;
4use futures::{Stream, StreamExt};
5use reqwest::Client;
6use serde::Deserialize;
7use serde_json::Value as JsonValue;
8use serde_json::json;
9use std::future::Future;
10use std::pin::Pin;
11use std::sync::Arc;
12
13use crate::{
14    Api, AssistantMessage, ContentBlock, Context, Model, Provider, ProviderEvent, StopReason,
15    StreamOptions, StreamResult, Usage, error::ProviderError,
16};
17
18use super::shared_client;
19
20/// Azure OpenAI provider
21///
22/// Uses Azure-specific endpoint format:
23/// https://{resource}.openai.azure.com/openai/deployments/{deployment}/chat/completions
24///
25/// Supports the following environment variables:
26///   - AZURE_OPENAI_API_KEY: API key for authentication
27///
28/// Azure OpenAI provider
29///
30/// Configuration is resolved at runtime from:
31/// - auth.json (for api_key)
32/// - settings.toml (for resource_name, deployment_name via custom_provider)
33/// - StreamOptions (for request-time override)
34#[derive(Clone)]
35pub struct AzureProvider {
36    client: &'static Client,
37    api_key: Option<String>,
38    resource_name: Option<String>,
39    deployment_name: Option<String>,
40}
41
42impl AzureProvider {
43    /// Create a new Azure provider without configuration.
44    ///
45    /// Configuration is resolved at request time via auth.json or StreamOptions.
46    pub fn new() -> Self {
47        Self {
48            client: shared_client(),
49            api_key: None,
50            resource_name: None,
51            deployment_name: None,
52        }
53    }
54
55    /// Create with explicit configuration (public API for external consumers)
56    #[cfg(test)]
57    pub fn with_config(
58        api_key: impl Into<String>,
59        resource_name: impl Into<String>,
60        deployment_name: impl Into<String>,
61    ) -> Self {
62        Self {
63            client: shared_client(),
64            api_key: Some(api_key.into()),
65            resource_name: Some(resource_name.into()),
66            deployment_name: Some(deployment_name.into()),
67        }
68    }
69
70    /// Build the Azure endpoint URL
71    fn build_url(&self, model: &Model) -> Result<String, ProviderError> {
72        // Priority: model.base_url > resource_name env > fallback
73        if !model.base_url.is_empty() && model.base_url != "https://api.openai.com" {
74            // Use the provided base URL directly (already includes deployment)
75            return Ok(format!(
76                "{}/chat/completions?api-version=2024-02-15-preview",
77                model.base_url.trim_end_matches('/')
78            ));
79        }
80
81        // Fallback to constructing from environment variables
82        let resource = self.resource_name.as_ref().ok_or_else(|| {
83            ProviderError::InvalidResponse("AZURE_OPENAI_RESOURCE_NAME not set".into())
84        })?;
85
86        let deployment = self.deployment_name.as_ref().ok_or_else(|| {
87            ProviderError::InvalidResponse("AZURE_OPENAI_DEPLOYMENT_NAME not set".into())
88        })?;
89
90        let url = format!(
91            "https://{}.openai.azure.com/openai/deployments/{}/chat/completions?api-version=2024-02-15-preview",
92            resource, deployment
93        );
94
95        Ok(url)
96    }
97
98    /// Get the API key (from options or self)
99    fn get_api_key(&self, options: &Option<StreamOptions>) -> Result<String, ProviderError> {
100        options
101            .as_ref()
102            .and_then(|o| o.api_key.as_ref())
103            .or(self.api_key.as_ref())
104            .cloned()
105            .ok_or_else(|| ProviderError::MissingApiKey)
106    }
107
108    /// Build request headers with Azure-specific api-key authentication
109    fn build_headers(
110        &self,
111        api_key: &str,
112        options: &Option<StreamOptions>,
113    ) -> Result<reqwest::header::HeaderMap, ProviderError> {
114        let mut headers = reqwest::header::HeaderMap::new();
115
116        // Azure uses api-key header instead of Bearer token
117        headers.insert(
118            "api-key",
119            api_key.parse().map_err(|e| {
120                ProviderError::InvalidResponse(format!("invalid header value: {e}"))
121            })?,
122        );
123        headers.insert(
124            reqwest::header::CONTENT_TYPE,
125            "application/json".parse().map_err(|e| {
126                ProviderError::InvalidResponse(format!("invalid header value: {e}"))
127            })?,
128        );
129
130        // Add custom headers from options
131        if let Some(opts) = options {
132            for (k, v) in &opts.headers {
133                if let (Ok(name), Ok(value)) = (
134                    k.parse::<reqwest::header::HeaderName>(),
135                    v.parse::<reqwest::header::HeaderValue>(),
136                ) {
137                    headers.insert(name, value);
138                }
139            }
140        }
141
142        Ok(headers)
143    }
144}
145
146impl Default for AzureProvider {
147    fn default() -> Self {
148        Self::new()
149    }
150}
151
152impl Provider for AzureProvider {
153    fn stream<'a>(
154        &'a self,
155        model: &'a Model,
156        context: &'a Context,
157        options: Option<StreamOptions>,
158    ) -> Pin<Box<dyn Future<Output = StreamResult> + Send + 'a>> {
159        Box::pin(async move {
160            // Build URL
161            let url = self.build_url(model)?;
162
163            // Get API key
164            let api_key = self.get_api_key(&options)?;
165
166            // Build messages
167            let messages = build_messages(context)?;
168
169            // Build request body
170            let mut body = serde_json::json!({
171                "messages": messages,
172                "stream": true,
173            });
174
175            // Add model if not already in URL (some deployments use it)
176            if model.id != "default" && model.id != "azure" {
177                body["model"] = serde_json::json!(model.id);
178            }
179
180            // Add optional parameters
181            if let Some(ref opts) = options {
182                if let Some(temp) = opts.temperature {
183                    body["temperature"] = serde_json::json!(temp);
184                }
185
186                if let Some(max) = opts.max_tokens {
187                    body["max_tokens"] = serde_json::json!(max);
188                }
189            }
190
191            // Add tools if present
192            if !context.tools.is_empty() {
193                body["tools"] = build_tools(&context.tools)?;
194            }
195
196            // Force the tool choice when a Named choice is set (and tools exist).
197            if let Some(choice) = options
198                .as_ref()
199                .and_then(|o| build_tool_choice(o.tool_choice.as_ref()))
200            {
201                body["tool_choice"] = choice;
202            }
203
204            // Build headers
205            let headers = self.build_headers(&api_key, &options)?;
206
207            // Make request
208            let response = self
209                .client
210                .post(&url)
211                .headers(headers)
212                .json(&body)
213                .send()
214                .await
215                .map_err(ProviderError::RequestFailed)?;
216
217            if !response.status().is_success() {
218                let status = response.status();
219                let body: String = response.text().await.unwrap_or_default();
220                return Err(ProviderError::HttpError(
221                    crate::error::HttpErrorDetail::new(status.as_u16(), body),
222                ));
223            }
224
225            // Create event stream
226            let provider_name = model.provider.clone();
227            let model_id = model.id.clone();
228
229            let stream = response
230                .bytes_stream()
231                .scan(
232                    Vec::<u8>::new(),
233                    move |pending_bytes, chunk: Result<Bytes, reqwest::Error>| {
234                        let pn = provider_name.clone();
235                        let mid = model_id.clone();
236                        // Synchronous computation — split_complete_lines and
237                        // parse_sse_events are sync, so no async block needed.
238                        // Using ready() avoids the lifetime error where an async
239                        // block would borrow `pending_bytes` beyond its scope.
240                        let events = match chunk {
241                            Ok(bytes) => {
242                                // Accumulate across HTTP chunk boundaries so SSE
243                                // lines split mid-stream are not silently dropped
244                                // (F-10, code audit 2026-07-25). Same pattern as
245                                // the openai/anthropic/google/vertex providers.
246                                let mut combined =
247                                    Vec::with_capacity(pending_bytes.len() + bytes.len());
248                                combined.extend_from_slice(pending_bytes);
249                                combined.extend_from_slice(&bytes);
250                                let (text, trailing) = super::sse::split_complete_lines(&combined);
251                                *pending_bytes = trailing;
252                                parse_sse_events(&text, &pn, &mid)
253                            }
254                            Err(e) => vec![ProviderEvent::Error {
255                                reason: StopReason::Error,
256                                error: create_error_message(&e.to_string(), &pn, &mid),
257                            }],
258                        };
259                        std::future::ready(Some(futures::stream::iter(events)))
260                    },
261                )
262                .flatten();
263
264            Ok(Box::pin(stream) as Pin<Box<dyn Stream<Item = ProviderEvent> + Send>>)
265        })
266    }
267}
268
269/// Build messages array from context
270fn build_messages(context: &Context) -> Result<Vec<JsonValue>, ProviderError> {
271    let mut messages = Vec::new();
272
273    // System prompt
274    if let Some(ref prompt) = context.system_prompt {
275        messages.push(serde_json::json!({
276            "role": "system",
277            "content": prompt,
278        }));
279    }
280
281    // Conversation messages
282    for msg in &context.messages {
283        match msg {
284            crate::Message::User(u) => {
285                let content: String = match &u.content {
286                    crate::MessageContent::Text(s) => s.clone(),
287                    crate::MessageContent::Blocks(blocks) => blocks_to_content(blocks)?.to_string(),
288                };
289                messages.push(serde_json::json!({
290                    "role": "user",
291                    "content": content,
292                }));
293            }
294            crate::Message::Assistant(a) => {
295                let content = blocks_to_content(&a.content)?.to_string();
296                messages.push(serde_json::json!({
297                    "role": "assistant",
298                    "content": content,
299                }));
300            }
301            crate::Message::ToolResult(t) => {
302                let content = blocks_to_content(&t.content)?.to_string();
303                messages.push(serde_json::json!({
304                    "role": "tool",
305                    "tool_call_id": t.tool_call_id,
306                    "tool_name": t.tool_name,
307                    "content": content,
308                }));
309            }
310        }
311    }
312
313    Ok(messages)
314}
315
316/// Convert content blocks to a string representation
317fn blocks_to_content(blocks: &[ContentBlock]) -> Result<JsonValue, ProviderError> {
318    if blocks.len() == 1
319        && let Some(text) = blocks[0].as_text()
320    {
321        return Ok(JsonValue::String(text.to_string()));
322    }
323
324    let items: Result<Vec<_>, _> = blocks
325        .iter()
326        .map(|block| match block {
327            ContentBlock::Text(t) => Ok(serde_json::json!({
328                "type": "text",
329                "text": t.text,
330            })),
331            ContentBlock::ToolCall(tc) => Ok(serde_json::json!({
332                "type": "function",
333                "id": tc.id,
334                "function": {
335                    "name": tc.name,
336                    "arguments": tc.arguments.to_string(),
337                },
338            })),
339            ContentBlock::Thinking(th) => Ok(serde_json::json!({
340                "type": "thinking",
341                "thinking": th.thinking,
342            })),
343            ContentBlock::Image(img) => Ok(serde_json::json!({
344                "type": "image_url",
345                "image_url": {
346                    "url": format!("data:{};base64,{}", img.mime_type, img.data),
347                },
348            })),
349            ContentBlock::Unknown(_) => Err(ProviderError::InvalidResponse(
350                "Unknown content block type".into(),
351            )),
352        })
353        .collect();
354
355    Ok(serde_json::json!(items?))
356}
357
358/// Map a `ToolChoice` to Azure OpenAI's forced-tool-choice shape.
359fn build_tool_choice(tool_choice: Option<&crate::tools::ToolChoice>) -> Option<JsonValue> {
360    match tool_choice {
361        None | Some(crate::tools::ToolChoice::Auto) => None,
362        Some(crate::tools::ToolChoice::Named(name)) => {
363            Some(json!({"type": "function", "function": {"name": name}}))
364        }
365    }
366}
367
368/// Build tools array
369fn build_tools(tools: &[crate::Tool]) -> Result<JsonValue, ProviderError> {
370    let items: Vec<_> = tools
371        .iter()
372        .map(|tool| {
373            serde_json::json!({
374                "type": "function",
375                "function": {
376                    "name": tool.name,
377                    "description": tool.description,
378                    "parameters": tool.parameters,
379                },
380            })
381        })
382        .collect();
383
384    Ok(serde_json::json!(items))
385}
386
387/// Parse SSE event stream from a byte buffer.
388///
389/// This is identical to the OpenAI provider's SSE parsing logic.
390fn parse_sse_events(text: &str, provider: &str, model_id: &str) -> Vec<ProviderEvent> {
391    let mut events = Vec::with_capacity(text.len() / 80);
392    let mut partial_message = AssistantMessage::new(Api::OpenAiCompletions, provider, model_id);
393
394    let mut accumulated_usage = Usage::default();
395
396    for line in text.split('\n') {
397        let line = line.trim_end_matches('\r');
398        if line.is_empty() {
399            continue;
400        }
401
402        // Fast rejection for non-data lines (comments, event tags, etc.)
403        if !line.starts_with("data: ") {
404            continue;
405        }
406
407        let data = &line[6..]; // skip "data: "
408
409        // Early exit on stream end
410        if data == "[DONE]" {
411            break;
412        }
413
414        if data.is_empty() {
415            continue;
416        }
417
418        let chunk = match serde_json::from_str::<SSEChunk>(data) {
419            Ok(c) => c,
420            Err(_) => continue,
421        };
422
423        // Get this chunk's usage for setting on Done events
424        let this_chunk_usage = chunk.usage.as_ref();
425
426        for choice in &chunk.choices {
427            if let Some(delta) = &choice.delta {
428                if let Some(content) = &delta.content {
429                    // pi-mono: accumulate into partial_message so the TUI can
430                    // diff against its snapshot tracker.
431                    let last_text_idx = partial_message
432                        .content
433                        .iter()
434                        .rposition(|b| matches!(b, ContentBlock::Text(_)));
435                    if let Some(idx) = last_text_idx
436                        && let ContentBlock::Text(t) = &mut partial_message.content[idx]
437                    {
438                        t.text.push_str(content);
439                    } else {
440                        partial_message
441                            .content
442                            .push(ContentBlock::Text(crate::TextContent::new(content.clone())));
443                    }
444                    events.push(ProviderEvent::TextDelta {
445                        content_index: choice.index,
446                        delta: content.clone(),
447                        partial: Arc::new(partial_message.clone()),
448                    });
449                }
450
451                if let Some(tool_calls) = &delta.tool_calls {
452                    for tc in tool_calls {
453                        if let Some(func) = &tc.function {
454                            events.push(ProviderEvent::ToolCallDelta {
455                                content_index: choice.index,
456                                delta: func.arguments.clone().unwrap_or_default(),
457                                partial: Arc::new(partial_message.clone()),
458                            });
459                        }
460                    }
461                }
462            }
463
464            if choice.finish_reason.is_some() {
465                // For Done events: prefer current chunk's usage if available,
466                // otherwise fall back to accumulated usage
467                let reason = match choice.finish_reason.as_deref() {
468                    Some("stop") => StopReason::Stop,
469                    Some("length") => StopReason::Length,
470                    Some("tool_calls") => StopReason::ToolUse,
471                    _ => StopReason::Stop,
472                };
473
474                let mut done_msg = partial_message.clone();
475
476                // Use current chunk's usage if present, otherwise accumulated
477                if let Some(usage) = this_chunk_usage {
478                    done_msg.usage.input = usage.prompt_tokens;
479                    done_msg.usage.output = usage.completion_tokens;
480                    done_msg.usage.cache_read = usage
481                        .prompt_tokens_details
482                        .as_ref()
483                        .map(|d| d.cached_tokens)
484                        .unwrap_or(0);
485                    done_msg.usage.total_tokens = usage.total_tokens;
486                } else {
487                    done_msg.usage = accumulated_usage.clone();
488                }
489
490                events.push(ProviderEvent::Done {
491                    reason,
492                    message: done_msg,
493                });
494            }
495        }
496
497        // Update accumulated usage for next chunks
498        if let Some(usage) = this_chunk_usage {
499            accumulated_usage.input = usage.prompt_tokens;
500            accumulated_usage.output = usage.completion_tokens;
501            accumulated_usage.cache_read = usage
502                .prompt_tokens_details
503                .as_ref()
504                .map(|d| d.cached_tokens)
505                .unwrap_or(0);
506            accumulated_usage.total_tokens = usage.total_tokens;
507        }
508    }
509
510    events
511}
512
513/// Create error assistant message
514fn create_error_message(msg: &str, provider: &str, model_id: &str) -> AssistantMessage {
515    let mut message = AssistantMessage::new(Api::OpenAiCompletions, provider, model_id);
516    message.stop_reason = StopReason::Error;
517    message.error_message = Some(msg.to_string());
518    message
519}
520
521// SSE chunk structure (same as OpenAI)
522#[derive(Debug, Deserialize)]
523struct SSEChunk {
524    _id: Option<String>,
525    #[serde(rename = "model")]
526    _model: Option<String>,
527    choices: Vec<Choice>,
528    usage: Option<UsageInfo>,
529}
530
531#[derive(Debug, Deserialize)]
532struct Choice {
533    index: usize,
534    delta: Option<Delta>,
535    finish_reason: Option<String>,
536}
537
538#[derive(Debug, Deserialize)]
539struct Delta {
540    content: Option<String>,
541    tool_calls: Option<Vec<ToolCallDelta>>,
542}
543
544#[derive(Debug, Deserialize)]
545struct ToolCallDelta {
546    _index: Option<usize>,
547    _id: Option<String>,
548    #[serde(rename = "type")]
549    _type_: Option<String>,
550    function: Option<FunctionDelta>,
551}
552
553#[derive(Debug, Deserialize)]
554struct FunctionDelta {
555    _name: Option<String>,
556    arguments: Option<String>,
557}
558
559#[derive(Debug, Deserialize, Clone)]
560struct UsageInfo {
561    prompt_tokens: usize,
562    completion_tokens: usize,
563    total_tokens: usize,
564    #[serde(rename = "prompt_tokens_details")]
565    prompt_tokens_details: Option<PromptTokensDetails>,
566}
567
568#[derive(Debug, Deserialize, Clone)]
569struct PromptTokensDetails {
570    #[serde(rename = "cached_tokens")]
571    cached_tokens: usize,
572}
573
574#[cfg(test)]
575mod tests {
576    use super::*;
577
578    #[test]
579    fn build_tool_choice_maps_named_to_azure_shape() {
580        assert!(build_tool_choice(None).is_none());
581        assert!(build_tool_choice(Some(&crate::tools::ToolChoice::Auto)).is_none());
582        assert_eq!(
583            build_tool_choice(Some(&crate::tools::ToolChoice::Named("todo".into()))),
584            Some(serde_json::json!({"type": "function", "function": {"name": "todo"}}))
585        );
586    }
587
588    fn make_test_model(id: &str, base_url: &str) -> Model {
589        Model::new(id, id, Api::OpenAiCompletions, "azure", base_url)
590    }
591
592    #[test]
593    fn test_build_url_from_base_url() {
594        let provider = AzureProvider::new();
595        let model = make_test_model(
596            "gpt-4o",
597            "https://my-resource.openai.azure.com/openai/deployments/gpt-4o",
598        );
599
600        let url = provider.build_url(&model).unwrap();
601        assert!(url.contains("api-version=2024-02-15-preview"));
602        assert!(url.contains("my-resource"));
603        assert!(url.contains("gpt-4o"));
604    }
605
606    #[test]
607    fn test_build_url_missing_resource() {
608        let provider = AzureProvider {
609            client: shared_client(),
610            api_key: Some("test-key".to_string()),
611            resource_name: None,
612            deployment_name: Some("gpt-4o".to_string()),
613        };
614
615        let model = make_test_model("default", "");
616
617        let result = provider.build_url(&model);
618        assert!(result.is_err());
619        match result.unwrap_err() {
620            ProviderError::InvalidResponse(msg) => {
621                assert!(msg.contains("AZURE_OPENAI_RESOURCE_NAME"));
622            }
623            _ => panic!("Expected InvalidResponse"),
624        }
625    }
626
627    #[test]
628    fn test_build_url_missing_deployment() {
629        let provider = AzureProvider {
630            client: shared_client(),
631            api_key: Some("test-key".to_string()),
632            resource_name: Some("my-resource".to_string()),
633            deployment_name: None,
634        };
635
636        let model = make_test_model("default", "");
637
638        let result = provider.build_url(&model);
639        assert!(result.is_err());
640        match result.unwrap_err() {
641            ProviderError::InvalidResponse(msg) => {
642                assert!(msg.contains("AZURE_OPENAI_DEPLOYMENT_NAME"));
643            }
644            _ => panic!("Expected InvalidResponse"),
645        }
646    }
647
648    #[test]
649    fn test_build_url_from_env_vars() {
650        let provider = AzureProvider {
651            client: shared_client(),
652            api_key: Some("test-key".to_string()),
653            resource_name: Some("my-resource".to_string()),
654            deployment_name: Some("gpt-4o".to_string()),
655        };
656
657        let model = make_test_model("default", "");
658
659        let url = provider.build_url(&model).unwrap();
660        assert_eq!(
661            url,
662            "https://my-resource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-02-15-preview"
663        );
664    }
665
666    #[test]
667    fn test_parse_sse_events_text() {
668        let sse_data = r#"data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}
669
670data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":"stop"}]}
671
672data: [DONE]"#;
673
674        let events = parse_sse_events(sse_data, "azure", "gpt-4o");
675
676        // Should have text delta events and a done event
677        assert!(events.len() >= 3);
678
679        // Check first text delta
680        match &events[0] {
681            ProviderEvent::TextDelta { delta, .. } => assert_eq!(delta, "Hello"),
682            _ => panic!("Expected TextDelta event"),
683        }
684
685        // Check done event
686        match &events[events.len() - 1] {
687            ProviderEvent::Done { reason, .. } => assert_eq!(*reason, StopReason::Stop),
688            _ => panic!("Expected Done event"),
689        }
690    }
691
692    #[test]
693    fn test_parse_sse_events_with_tool_calls() {
694        let sse_data = r#"data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"tool_calls":[{"id":"call_abc123","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":null}]}
695
696data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"location\":"}}]},"finish_reason":null}]}
697
698data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"Boston\"}"}}]},"finish_reason":"tool_calls"}]}
699
700data: [DONE]"#;
701
702        let events = parse_sse_events(sse_data, "azure", "gpt-4o");
703
704        // Should have tool call delta events and a done event
705        assert!(events.len() >= 4);
706
707        // Check for tool call delta
708        let has_tool_call = events
709            .iter()
710            .any(|e| matches!(e, ProviderEvent::ToolCallDelta { .. }));
711        assert!(
712            has_tool_call,
713            "Should have at least one ToolCallDelta event"
714        );
715
716        // Check done event
717        match &events[events.len() - 1] {
718            ProviderEvent::Done { reason, .. } => assert_eq!(*reason, StopReason::ToolUse),
719            _ => panic!("Expected Done event with ToolUse reason"),
720        }
721    }
722
723    #[test]
724    fn test_parse_sse_events_usage() {
725        let sse_data = r#"data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,"prompt_tokens_details":{"cached_tokens":0}}}
726
727data: [DONE]"#;
728
729        let events = parse_sse_events(sse_data, "azure", "gpt-4o");
730
731        // Find the done event and check usage
732        let done_event = events
733            .iter()
734            .find(|e| matches!(e, ProviderEvent::Done { .. }));
735        assert!(done_event.is_some());
736
737        if let ProviderEvent::Done { message, .. } = done_event.unwrap() {
738            assert_eq!(message.usage.input, 10);
739            assert_eq!(message.usage.output, 5);
740            assert_eq!(message.usage.total_tokens, 15);
741        }
742    }
743
744    #[test]
745    fn test_build_headers_includes_api_key() {
746        let provider = AzureProvider::new();
747        let api_key = "test-api-key-12345";
748
749        let headers = provider.build_headers(api_key, &None).unwrap();
750
751        // Check api-key header is present
752        let api_key_header = headers.get("api-key");
753        assert!(api_key_header.is_some());
754        assert_eq!(api_key_header.unwrap().to_str().unwrap(), api_key);
755
756        // Check content-type is present
757        let content_type = headers.get(reqwest::header::CONTENT_TYPE);
758        assert!(content_type.is_some());
759    }
760
761    #[test]
762    fn test_build_headers_no_bearer_token() {
763        let provider = AzureProvider::new();
764        let api_key = "test-api-key-12345";
765
766        let headers = provider.build_headers(api_key, &None).unwrap();
767
768        // Azure should NOT use Authorization header with Bearer token
769        let auth_header = headers.get(reqwest::header::AUTHORIZATION);
770        assert!(
771            auth_header.is_none(),
772            "Azure should not use Bearer token authentication"
773        );
774    }
775
776    #[test]
777    fn test_with_config_constructor() {
778        let provider = AzureProvider::with_config("my-api-key", "my-resource", "gpt-4o");
779
780        // Verify the internal state through build_url
781        let model = make_test_model("default", "");
782
783        let url = provider.build_url(&model).unwrap();
784        assert!(url.contains("my-resource"));
785        assert!(url.contains("gpt-4o"));
786    }
787
788    #[test]
789    fn test_azure_endpoint_format() {
790        let provider = AzureProvider {
791            client: shared_client(),
792            api_key: Some("key".to_string()),
793            resource_name: Some("my-resource".to_string()),
794            deployment_name: Some("gpt-4-turbo".to_string()),
795        };
796
797        let model = make_test_model("default", "");
798        let url = provider.build_url(&model).unwrap();
799
800        // Verify the complete Azure endpoint format
801        assert!(url.starts_with("https://"));
802        assert!(url.contains(".openai.azure.com"));
803        assert!(url.contains("/openai/deployments/"));
804        assert!(url.contains("gpt-4-turbo"));
805        assert!(url.contains("chat/completions"));
806        assert!(url.contains("api-version=2024-02-15-preview"));
807    }
808}