Skip to main content

vtcode_llm/providers/
custom_provider.rs

1use async_trait::async_trait;
2use vtcode_commons::llm::BackendKind;
3use vtcode_config::core::AnthropicConfig;
4use vtcode_config::core::{CustomProviderApiFormat, CustomProviderConfig, ResolvedCustomProviderProfile};
5use vtcode_config::{ModelConfig, OpenAIConfig, PromptCachingConfig, TimeoutsConfig};
6
7use crate::provider::{
8    LLMError, LLMProvider, LLMRequest, LLMResponse, LLMStream, Message, ResponsesCompactionOptions, SamplingOverrides,
9};
10use crate::providers::anthropic::{self, AnthropicProvider};
11use crate::providers::common::{resolve_model, validate_request_common};
12use crate::providers::openai::{CustomProviderAuthHandle, OpenAIProvider};
13
14/// Profile-aware backend router for custom providers.
15pub struct CustomProviderBackendRouter {
16    provider_key: String,
17    display_name: String,
18    default_model: String,
19    supported_models: Vec<String>,
20    custom_config: CustomProviderConfig,
21    anthropic_config: AnthropicConfig,
22    openai_auto: OpenAIProvider,
23    openai_chat: OpenAIProvider,
24    openai_responses: OpenAIProvider,
25    anthropic: AnthropicProvider,
26}
27
28impl CustomProviderBackendRouter {
29    #[allow(
30        clippy::too_many_arguments,
31        reason = "Intentional compatibility, platform, or test-only suppression."
32    )]
33    pub fn from_config(
34        custom_config: CustomProviderConfig,
35        api_key: Option<String>,
36        model: Option<String>,
37        base_url: String,
38        prompt_cache: Option<PromptCachingConfig>,
39        timeouts: Option<TimeoutsConfig>,
40        openai: Option<OpenAIConfig>,
41        anthropic: Option<AnthropicConfig>,
42        model_behavior: Option<ModelConfig>,
43        custom_provider_auth: Option<CustomProviderAuthHandle>,
44    ) -> Self {
45        let provider_key = custom_config.name.clone();
46        let display_name = custom_config.display_name.clone();
47        let supported_models = custom_config.effective_models();
48        let configured_default = resolve_model(model, &custom_config.model);
49        let default_model = if configured_default.trim().is_empty() {
50            supported_models.first().cloned().unwrap_or_default()
51        } else {
52            configured_default
53        };
54        let anthropic_config = anthropic.unwrap_or_default();
55
56        let openai_auto = OpenAIProvider::from_custom_config(
57            provider_key.clone(),
58            display_name.clone(),
59            api_key.clone(),
60            Some(default_model.clone()),
61            Some(base_url.clone()),
62            prompt_cache.clone(),
63            timeouts.clone(),
64            openai.clone(),
65            model_behavior.clone(),
66            custom_provider_auth.clone(),
67            Some(supported_models.clone()),
68        );
69        let openai_chat = OpenAIProvider::from_custom_config(
70            provider_key.clone(),
71            display_name.clone(),
72            api_key.clone(),
73            Some(default_model.clone()),
74            Some(base_url.clone()),
75            prompt_cache.clone(),
76            timeouts.clone(),
77            openai.clone(),
78            model_behavior.clone(),
79            custom_provider_auth.clone(),
80            Some(supported_models.clone()),
81        )
82        .with_api_format_override(Some(CustomProviderApiFormat::OpenAIChat));
83        let openai_responses = OpenAIProvider::from_custom_config(
84            provider_key.clone(),
85            display_name.clone(),
86            api_key.clone(),
87            Some(default_model.clone()),
88            Some(base_url.clone()),
89            prompt_cache.clone(),
90            timeouts.clone(),
91            openai.clone(),
92            model_behavior.clone(),
93            custom_provider_auth.clone(),
94            Some(supported_models.clone()),
95        )
96        .with_api_format_override(Some(CustomProviderApiFormat::OpenAIResponses));
97        let anthropic = AnthropicProvider::from_config(
98            api_key,
99            Some(default_model.clone()),
100            Some(base_url),
101            prompt_cache,
102            timeouts,
103            Some(anthropic_config.clone()),
104            model_behavior,
105        )
106        .with_custom_auth(custom_provider_auth);
107
108        Self {
109            provider_key,
110            display_name,
111            default_model,
112            supported_models,
113            custom_config,
114            anthropic_config,
115            openai_auto,
116            openai_chat,
117            openai_responses,
118            anthropic,
119        }
120    }
121
122    fn resolved_model<'a>(&'a self, model: &'a str) -> &'a str {
123        if model.trim().is_empty() {
124            &self.default_model
125        } else {
126            model
127        }
128    }
129
130    fn profile_for_model(&self, model: &str) -> ResolvedCustomProviderProfile {
131        self.custom_config.resolved_profile(model)
132    }
133
134    fn backend_for_model(&self, model: &str) -> &dyn LLMProvider {
135        match self.profile_for_model(self.resolved_model(model)).api_format {
136            Some(CustomProviderApiFormat::AnthropicMessages) => &self.anthropic,
137            Some(CustomProviderApiFormat::OpenAIChat) => &self.openai_chat,
138            Some(CustomProviderApiFormat::OpenAIResponses) => &self.openai_responses,
139            Some(CustomProviderApiFormat::Auto) | None => &self.openai_auto,
140        }
141    }
142
143    fn override_bool(value: Option<bool>, default: bool) -> bool {
144        value.unwrap_or(default)
145    }
146}
147
148#[async_trait]
149impl LLMProvider for CustomProviderBackendRouter {
150    fn name(&self) -> &str {
151        &self.provider_key
152    }
153
154    fn backend_kind(&self) -> BackendKind {
155        self.backend_for_model(&self.default_model).backend_kind()
156    }
157
158    fn supports_streaming(&self) -> bool {
159        true
160    }
161
162    fn supports_non_streaming(&self, model: &str) -> bool {
163        self.backend_for_model(model).supports_non_streaming(self.resolved_model(model))
164    }
165
166    fn supports_reasoning(&self, model: &str) -> bool {
167        let resolved = self.resolved_model(model);
168        let profile = self.profile_for_model(resolved);
169        Self::override_bool(profile.supports_reasoning, self.backend_for_model(model).supports_reasoning(resolved))
170    }
171
172    fn supports_reasoning_effort(&self, model: &str) -> bool {
173        let resolved = self.resolved_model(model);
174        let profile = self.profile_for_model(resolved);
175        // A pin in THIS model's own profile implies effort support unless the
176        // same profile explicitly disables it. Provider-level defaults — both
177        // the effort value and the capability flag — must not flip support for
178        // every model on the endpoint.
179        let own = self.custom_config.profiles.get(resolved);
180        let own_pin = own.and_then(|p| p.reasoning_effort).is_some();
181        if own_pin {
182            return own.and_then(|p| p.supports_reasoning_effort) != Some(false);
183        }
184        Self::override_bool(
185            profile.supports_reasoning_effort,
186            self.backend_for_model(model).supports_reasoning_effort(resolved),
187        )
188    }
189
190    fn sampling_overrides(&self, model: &str) -> SamplingOverrides {
191        let resolved = self.resolved_model(model);
192        let profile = self.profile_for_model(resolved);
193        SamplingOverrides {
194            temperature: profile.temperature,
195            top_p: profile.top_p,
196            top_k: profile.top_k,
197            presence_penalty: profile.presence_penalty,
198            frequency_penalty: profile.frequency_penalty,
199            max_tokens: profile.max_tokens,
200            reasoning_effort: profile.reasoning_effort,
201            suppresses_sampling_with_reasoning: profile.api_format == Some(CustomProviderApiFormat::AnthropicMessages),
202            profile_aware: true,
203        }
204    }
205
206    fn supports_tools(&self, model: &str) -> bool {
207        let resolved = self.resolved_model(model);
208        let profile = self.profile_for_model(resolved);
209        Self::override_bool(profile.supports_tools, self.backend_for_model(model).supports_tools(resolved))
210    }
211
212    fn supports_parallel_tool_config(&self, model: &str) -> bool {
213        let resolved = self.resolved_model(model);
214        let profile = self.profile_for_model(resolved);
215        Self::override_bool(
216            profile.supports_parallel_tool_calls,
217            self.backend_for_model(model).supports_parallel_tool_config(resolved),
218        )
219    }
220
221    fn supports_structured_output(&self, model: &str) -> bool {
222        let resolved = self.resolved_model(model);
223        let profile = self.profile_for_model(resolved);
224        Self::override_bool(
225            profile.supports_structured_output,
226            self.backend_for_model(model).supports_structured_output(resolved),
227        )
228    }
229
230    fn supports_context_caching(&self, model: &str) -> bool {
231        let resolved = self.resolved_model(model);
232        let profile = self.profile_for_model(resolved);
233        Self::override_bool(
234            profile.supports_context_caching,
235            self.backend_for_model(model).supports_context_caching(resolved),
236        )
237    }
238
239    fn supports_vision(&self, model: &str) -> bool {
240        let resolved = self.resolved_model(model);
241        let profile = self.profile_for_model(resolved);
242        Self::override_bool(profile.supports_vision, self.backend_for_model(model).supports_vision(resolved))
243    }
244
245    fn supports_responses_compaction(&self, model: &str) -> bool {
246        let resolved = self.resolved_model(model);
247        let profile = self.profile_for_model(resolved);
248        let backend_default = match self.profile_for_model(resolved).api_format {
249            Some(CustomProviderApiFormat::AnthropicMessages) => false,
250            _ => self.backend_for_model(model).supports_responses_compaction(resolved),
251        };
252        Self::override_bool(profile.supports_responses_compaction, backend_default)
253    }
254
255    fn supports_native_allowed_tools(&self, model: &str) -> bool {
256        self.backend_for_model(model)
257            .supports_native_allowed_tools(self.resolved_model(model))
258    }
259
260    fn supports_context_edits(&self, model: &str) -> bool {
261        let resolved = self.resolved_model(model);
262        let profile = self.profile_for_model(resolved);
263        Self::override_bool(
264            profile.supports_context_edits,
265            self.backend_for_model(model).supports_context_edits(resolved),
266        )
267    }
268
269    fn supports_turn_scoped_system_messages(&self, model: &str) -> bool {
270        let resolved = self.resolved_model(model);
271        matches!(self.profile_for_model(resolved).api_format, Some(CustomProviderApiFormat::AnthropicMessages))
272            && self.backend_for_model(model).supports_turn_scoped_system_messages(resolved)
273    }
274
275    fn supports_manual_openai_compaction(&self, model: &str) -> bool {
276        let resolved = self.resolved_model(model);
277        let profile = self.profile_for_model(resolved);
278        let backend_default = match self.profile_for_model(resolved).api_format {
279            Some(CustomProviderApiFormat::AnthropicMessages) => false,
280            _ => self.backend_for_model(model).supports_manual_openai_compaction(resolved),
281        };
282        Self::override_bool(profile.supports_responses_compaction, backend_default)
283    }
284
285    fn supports_native_inline_compaction(&self, model: &str) -> bool {
286        let resolved = self.resolved_model(model);
287        let profile = self.profile_for_model(resolved);
288        Self::override_bool(
289            profile.supports_context_edits,
290            self.backend_for_model(model).supports_native_inline_compaction(resolved),
291        )
292    }
293
294    fn manual_openai_compaction_unavailable_message(&self, model: &str) -> String {
295        self.backend_for_model(model)
296            .manual_openai_compaction_unavailable_message(self.resolved_model(model))
297    }
298
299    fn effective_context_size(&self, model: &str) -> usize {
300        let resolved = self.resolved_model(model);
301        let profile = self.profile_for_model(resolved);
302        profile
303            .context_window
304            .unwrap_or_else(|| self.backend_for_model(model).effective_context_size(resolved))
305    }
306
307    async fn compact_history(&self, model: &str, history: &[Message]) -> Result<Vec<Message>, LLMError> {
308        self.backend_for_model(model)
309            .compact_history(self.resolved_model(model), history)
310            .await
311    }
312
313    async fn compact_history_with_options(
314        &self,
315        model: &str,
316        history: &[Message],
317        options: &ResponsesCompactionOptions,
318    ) -> Result<Vec<Message>, LLMError> {
319        self.backend_for_model(model)
320            .compact_history_with_options(self.resolved_model(model), history, options)
321            .await
322    }
323
324    async fn stream(&self, request: LLMRequest) -> Result<LLMStream, LLMError> {
325        self.backend_for_model(&request.model).stream(request).await
326    }
327
328    async fn stream_normalized(&self, request: LLMRequest) -> Result<crate::provider::LLMNormalizedStream, LLMError> {
329        self.backend_for_model(&request.model).stream_normalized(request).await
330    }
331
332    async fn generate(&self, request: LLMRequest) -> Result<LLMResponse, LLMError> {
333        self.backend_for_model(&request.model).generate(request).await
334    }
335
336    async fn generate_auxiliary(&self, request: LLMRequest) -> Result<LLMResponse, LLMError> {
337        self.backend_for_model(&request.model).generate_auxiliary(request).await
338    }
339
340    fn supported_models(&self) -> Vec<String> {
341        self.supported_models.clone()
342    }
343
344    fn validate_request(&self, request: &LLMRequest) -> Result<(), LLMError> {
345        let resolved_model = self.resolved_model(&request.model);
346        match self.profile_for_model(resolved_model).api_format {
347            Some(CustomProviderApiFormat::AnthropicMessages) => {
348                // Anthropic-shaped custom profiles use the Anthropic message
349                // contract, including its turn-scoped system marker. Keep the
350                // custom display name for diagnostics, but validate message roles
351                // against the backend contract so the marker is not rejected.
352                validate_request_common(request, &self.display_name, "anthropic", Some(&self.supported_models))?;
353                anthropic::validation::validate_request(
354                    request,
355                    &self.default_model,
356                    &self.anthropic_config,
357                    &self.display_name,
358                )
359            }
360            Some(CustomProviderApiFormat::OpenAIChat)
361            | Some(CustomProviderApiFormat::OpenAIResponses)
362            | Some(CustomProviderApiFormat::Auto)
363            | None => self.backend_for_model(resolved_model).validate_request(request),
364        }
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use super::CustomProviderBackendRouter;
371    use crate::provider::{LLMProvider, LLMRequest, Message, SamplingOverrides};
372    use crate::providers::openai::CustomProviderAuthHandle;
373    use serde_json::json;
374    use std::path::Path;
375    use tempfile::TempDir;
376    use vtcode_config::constants::models;
377    use vtcode_config::core::{
378        AnthropicConfig, CustomProviderApiFormat, CustomProviderCommandAuthConfig, CustomProviderConfig,
379        CustomProviderProfileConfig,
380    };
381    use vtcode_config::types::ReasoningEffortLevel;
382    use wiremock::matchers::{method, path};
383    use wiremock::{Mock, MockServer, ResponseTemplate};
384
385    fn sampling_router() -> CustomProviderBackendRouter {
386        let mut profiles = std::collections::BTreeMap::new();
387        profiles.insert(
388            "cold-model".to_string(),
389            CustomProviderProfileConfig {
390                temperature: Some(0.0),
391                top_p: Some(0.9),
392                reasoning_effort: Some(ReasoningEffortLevel::Low),
393                ..Default::default()
394            },
395        );
396
397        CustomProviderBackendRouter::from_config(
398            CustomProviderConfig {
399                name: "sampling-test".to_string(),
400                display_name: "Sampling Test".to_string(),
401                base_url: "https://llm.example/v1".to_string(),
402                temperature: Some(0.5),
403                models: vec!["cold-model".to_string(), "warm-model".to_string()],
404                profiles,
405                ..Default::default()
406            },
407            None,
408            None,
409            "https://llm.example/v1".to_string(),
410            None,
411            None,
412            None,
413            None,
414            None,
415            None,
416        )
417    }
418
419    #[tokio::test]
420    async fn auxiliary_generation_forwards_isolation_to_custom_openai_backend() {
421        let server = MockServer::start().await;
422        Mock::given(method("POST"))
423            .and(path("/responses"))
424            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
425                "id": "resp_custom_auxiliary",
426                "output": [{"type": "message", "content": [{"type": "output_text", "text": "SAFE"}]}]
427            })))
428            .expect(1)
429            .mount(&server)
430            .await;
431        let router = CustomProviderBackendRouter::from_config(
432            CustomProviderConfig {
433                name: "auxiliary-test".into(),
434                display_name: "Auxiliary Test".into(),
435                model: "gpt-6-luna".into(),
436                models: vec!["gpt-6-luna".into()],
437                api_format: CustomProviderApiFormat::OpenAIResponses,
438                ..Default::default()
439            },
440            Some("fixture-key".into()),
441            None,
442            server.uri(),
443            None,
444            None,
445            Some(vtcode_config::OpenAIConfig { websocket_mode: true, ..Default::default() }),
446            None,
447            None,
448            None,
449        );
450        let response = router
451            .generate_auxiliary(LLMRequest {
452                model: "gpt-6-luna".into(),
453                messages: vec![Message::user("independent custom probe".into())].into(),
454                ..Default::default()
455            })
456            .await
457            .expect("HTTP auxiliary response");
458        assert_eq!(response.content.as_deref(), Some("SAFE"));
459        let requests = server.received_requests().await.expect("recorded requests");
460        assert_eq!(requests.len(), 1, "auxiliary must not attempt a WebSocket upgrade or warmup");
461        assert_eq!(requests[0].method.as_str(), "POST");
462        assert_eq!(requests[0].url.path(), "/responses");
463        let payload: serde_json::Value = requests[0].body_json().expect("request payload");
464        assert_eq!(payload["model"], "gpt-6-luna");
465    }
466
467    #[test]
468    fn sampling_overrides_resolve_per_model_with_provider_defaults() {
469        let router = sampling_router();
470
471        let cold = router.sampling_overrides("cold-model");
472        assert_eq!(
473            cold,
474            SamplingOverrides {
475                temperature: Some(0.0),
476                top_p: Some(0.9),
477                top_k: None,
478                presence_penalty: None,
479                frequency_penalty: None,
480                max_tokens: None,
481                reasoning_effort: Some(ReasoningEffortLevel::Low),
482                suppresses_sampling_with_reasoning: false,
483                profile_aware: true,
484            }
485        );
486        assert!(router.supports_reasoning_effort("cold-model"));
487
488        // Models without a profile fall back to provider-level defaults.
489        let warm = router.sampling_overrides("warm-model");
490        assert_eq!(warm.temperature, Some(0.5));
491        assert_eq!(warm.top_p, None);
492    }
493
494    #[test]
495    fn sampling_overrides_suppression_matrix_matches_native_semantics() {
496        let profile_openai = SamplingOverrides { profile_aware: true, ..Default::default() };
497        let profile_anthropic = SamplingOverrides {
498            suppresses_sampling_with_reasoning: true,
499            profile_aware: true,
500            ..Default::default()
501        };
502        let builtin_default = SamplingOverrides::default();
503
504        // Custom openai-shaped profile keeps pinned values during reasoning.
505        assert!(!profile_openai.suppresses_sampling(false, true));
506
507        // Custom anthropic-messages profile drops them.
508        assert!(profile_anthropic.suppresses_sampling(false, true));
509        assert!(!profile_anthropic.suppresses_sampling(false, false));
510
511        // Built-in Anthropic/MiniMax shape suppresses only via native match;
512        // built-in OpenAI shape never does through overrides alone.
513        assert!(builtin_default.suppresses_sampling(true, true));
514        assert!(!builtin_default.suppresses_sampling(true, false));
515        assert!(!builtin_default.suppresses_sampling(false, true));
516
517        // Level "none"/"unknown" is not active reasoning anywhere.
518        assert!(!profile_anthropic.suppresses_sampling(false, false));
519    }
520
521    #[test]
522    fn turn_scoped_system_capability_follows_selected_api_profile() {
523        let mut profiles = std::collections::BTreeMap::new();
524        profiles.insert(
525            models::anthropic::CLAUDE_FABLE_5.to_string(),
526            CustomProviderProfileConfig {
527                api_format: CustomProviderApiFormat::AnthropicMessages,
528                ..Default::default()
529            },
530        );
531
532        let router = CustomProviderBackendRouter::from_config(
533            CustomProviderConfig {
534                name: "profiled-transport-test".to_string(),
535                display_name: "Profiled Transport Test".to_string(),
536                base_url: "https://llm.example/v1".to_string(),
537                model: models::anthropic::CLAUDE_FABLE_5.to_string(),
538                models: vec![
539                    models::anthropic::CLAUDE_FABLE_5.to_string(),
540                    "openai-shaped".to_string(),
541                ],
542                profiles,
543                ..Default::default()
544            },
545            Some("test-key".to_string()),
546            None,
547            "https://llm.example/v1".to_string(),
548            None,
549            None,
550            None,
551            Some(AnthropicConfig::default()),
552            None,
553            None,
554        );
555
556        assert!(router.supports_turn_scoped_system_messages(models::anthropic::CLAUDE_FABLE_5));
557        assert!(!router.supports_turn_scoped_system_messages("openai-shaped"));
558    }
559
560    #[test]
561    fn anthropic_profile_accepts_turn_scoped_system_marker() {
562        let mut profiles = std::collections::BTreeMap::new();
563        profiles.insert(
564            models::anthropic::CLAUDE_FABLE_5.to_string(),
565            CustomProviderProfileConfig {
566                api_format: CustomProviderApiFormat::AnthropicMessages,
567                ..Default::default()
568            },
569        );
570        let router = CustomProviderBackendRouter::from_config(
571            CustomProviderConfig {
572                name: "profiled-validation-test".to_string(),
573                display_name: "Profiled Validation Test".to_string(),
574                base_url: "https://llm.example/v1".to_string(),
575                model: models::anthropic::CLAUDE_FABLE_5.to_string(),
576                models: vec![models::anthropic::CLAUDE_FABLE_5.to_string()],
577                profiles,
578                ..Default::default()
579            },
580            Some("test-key".to_string()),
581            None,
582            "https://llm.example/v1".to_string(),
583            None,
584            None,
585            None,
586            Some(AnthropicConfig::default()),
587            None,
588            None,
589        );
590        let request = LLMRequest {
591            model: models::anthropic::CLAUDE_FABLE_5.to_string(),
592            messages: vec![Message::turn_scoped_system("visible-output notice".to_string())].into(),
593            ..Default::default()
594        };
595
596        router
597            .validate_request(&request)
598            .expect("custom Anthropic profile should accept native marker");
599    }
600
601    fn write_tokens_file(dir: &Path, tokens: &[&str]) {
602        std::fs::write(dir.join("tokens.txt"), tokens.join("\n")).expect("write tokens file");
603    }
604
605    #[cfg(unix)]
606    fn auth_fixture(dir: &TempDir, tokens: &[&str]) -> CustomProviderCommandAuthConfig {
607        use std::os::unix::fs::PermissionsExt;
608
609        write_tokens_file(dir.path(), tokens);
610        let script_path = dir.path().join("print-token.sh");
611        std::fs::write(
612            &script_path,
613            r#"#!/bin/sh
614    first_line=$(sed -n '1p' tokens.txt)
615    printf '%s\n' "$first_line"
616    tail -n +2 tokens.txt > tokens.next
617    mv tokens.next tokens.txt
618    "#,
619        )
620        .expect("write script");
621        let mut perms = std::fs::metadata(&script_path).expect("metadata").permissions();
622        perms.set_mode(0o755);
623        std::fs::set_permissions(&script_path, perms).expect("set permissions");
624        CustomProviderCommandAuthConfig {
625            command: "./print-token.sh".to_string(),
626            args: Vec::new(),
627            cwd: Some(dir.path().to_path_buf()),
628            timeout_ms: 5_000,
629            refresh_interval_ms: 60_000,
630        }
631    }
632
633    #[cfg(windows)]
634    fn auth_fixture(dir: &TempDir, tokens: &[&str]) -> CustomProviderCommandAuthConfig {
635        write_tokens_file(dir.path(), tokens);
636        let script_path = dir.path().join("print-token.ps1");
637        std::fs::write(
638            &script_path,
639            r#"$lines = Get-Content -Path tokens.txt
640    if ($lines.Count -eq 0) { exit 1 }
641    Write-Output $lines[0]
642    $lines | Select-Object -Skip 1 | Set-Content -Path tokens.txt
643    "#,
644        )
645        .expect("write script");
646        CustomProviderCommandAuthConfig {
647            command: "powershell".to_string(),
648            args: vec![
649                "-NoProfile".to_string(),
650                "-ExecutionPolicy".to_string(),
651                "Bypass".to_string(),
652                "-File".to_string(),
653                script_path.to_string_lossy().into_owned(),
654            ],
655            cwd: Some(dir.path().to_path_buf()),
656            timeout_ms: 5_000,
657            refresh_interval_ms: 60_000,
658        }
659    }
660
661    fn anthropic_response_body(text: &str) -> serde_json::Value {
662        json!({
663            "content": [{
664                "type": "text",
665                "text": text,
666            }],
667            "stop_reason": "end_turn",
668            "usage": {
669                "input_tokens": 1,
670                "output_tokens": 1,
671                "cache_creation_input_tokens": 0,
672                "cache_read_input_tokens": 0
673            }
674        })
675    }
676
677    #[tokio::test]
678    async fn anthropic_command_auth_refreshes_and_uses_protocol_headers() {
679        let server = MockServer::start().await;
680        let tempdir = TempDir::new().expect("tempdir");
681        let seen_keys = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
682        let auth_config = auth_fixture(&tempdir, &["first-token", "second-token"]);
683
684        let seen_for_mock = std::sync::Arc::clone(&seen_keys);
685        Mock::given(method("POST"))
686            .and(path("/messages"))
687            .respond_with(move |request: &wiremock::Request| {
688                let api_key = request
689                    .headers
690                    .get("x-api-key")
691                    .and_then(|value| value.to_str().ok())
692                    .unwrap_or_default()
693                    .to_string();
694                seen_for_mock.lock().expect("mutex").push(api_key);
695                let count = seen_for_mock.lock().expect("mutex").len();
696                if count == 1 {
697                    ResponseTemplate::new(401)
698                } else {
699                    ResponseTemplate::new(200).set_body_json(anthropic_response_body("anthropic command auth"))
700                }
701            })
702            .expect(2)
703            .mount(&server)
704            .await;
705
706        let config = CustomProviderConfig {
707            name: "anthropic-custom".to_string(),
708            display_name: "Anthropic Custom".to_string(),
709            base_url: server.uri(),
710            api_format: CustomProviderApiFormat::AnthropicMessages,
711            context_window: None,
712            temperature: None,
713            top_p: None,
714            top_k: None,
715            presence_penalty: None,
716            frequency_penalty: None,
717            max_tokens: None,
718            reasoning_effort: None,
719            supports_tools: None,
720            supports_reasoning: None,
721            supports_reasoning_effort: None,
722            supports_vision: None,
723            supports_structured_output: None,
724            supports_parallel_tool_calls: None,
725            supports_context_caching: None,
726            supports_responses_compaction: None,
727            supports_context_edits: Some(true),
728            api_key_env: "ANTHROPIC_CUSTOM_API_KEY".to_string(),
729            auth: None,
730            model: models::anthropic::DEFAULT_MODEL.to_string(),
731            models: vec![models::anthropic::DEFAULT_MODEL.to_string()],
732            profiles: std::collections::BTreeMap::new(),
733        };
734
735        let router = CustomProviderBackendRouter::from_config(
736            config,
737            None,
738            Some(models::anthropic::DEFAULT_MODEL.to_string()),
739            server.uri(),
740            None,
741            None,
742            None,
743            Some(AnthropicConfig::default()),
744            None,
745            Some(CustomProviderAuthHandle::new(auth_config, None)),
746        );
747
748        let response = router
749            .generate(LLMRequest {
750                model: models::anthropic::DEFAULT_MODEL.to_string(),
751                messages: vec![Message::user("hello".to_string())].into(),
752                ..Default::default()
753            })
754            .await
755            .expect("anthropic auth should refresh and succeed");
756
757        assert_eq!(response.content.as_deref(), Some("anthropic command auth"));
758        assert_eq!(
759            seen_keys.lock().expect("mutex").as_slice(),
760            &["first-token".to_string(), "second-token".to_string()]
761        );
762    }
763}