Skip to main content

llm_codegen/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use proc_macro2::TokenStream;
4use quote::{ToTokens, format_ident, quote};
5use serde::Deserialize;
6use std::collections::{BTreeMap, HashMap};
7use std::fmt::Write;
8use std::path::Path;
9
10type ModelsDevData = HashMap<String, ProviderData>;
11type ContextWindowOverride = fn(&str, u32) -> u32;
12
13#[derive(Debug, Deserialize)]
14struct ProviderData {
15    #[allow(dead_code)]
16    id: String,
17    #[allow(dead_code)]
18    name: String,
19    #[serde(default)]
20    #[allow(dead_code)]
21    env: Vec<String>,
22    #[serde(default)]
23    models: HashMap<String, ModelData>,
24}
25
26#[derive(Debug, Deserialize)]
27struct ModelData {
28    id: String,
29    name: String,
30    #[serde(default)]
31    tool_call: Option<bool>,
32    #[serde(default)]
33    reasoning: Option<bool>,
34    #[serde(default)]
35    #[allow(dead_code)]
36    cost: Option<CostData>,
37    #[serde(default)]
38    limit: Option<LimitData>,
39    #[serde(default)]
40    modalities: Option<ModalitiesData>,
41}
42
43#[derive(Debug, Deserialize, Default)]
44struct ModalitiesData {
45    #[serde(default)]
46    input: Vec<String>,
47}
48
49#[derive(Debug, Deserialize)]
50#[allow(dead_code)]
51struct CostData {
52    #[serde(default)]
53    input: f64,
54    #[serde(default)]
55    output: f64,
56    #[serde(default)]
57    cache_read: Option<f64>,
58    #[serde(default)]
59    cache_write: Option<f64>,
60}
61
62#[derive(Debug, Deserialize)]
63struct LimitData {
64    #[serde(default)]
65    context: u32,
66    #[serde(default)]
67    #[allow(dead_code)]
68    output: u32,
69}
70
71impl CostData {
72    fn has_prompt_caching(&self) -> bool {
73        self.cache_read.is_some() || self.cache_write.is_some()
74    }
75}
76
77/// Provider configuration for codegen (catalog providers with known model lists)
78struct ProviderConfig {
79    /// Unique provider key used in `provider_models` map (e.g. "codex")
80    dev_id: &'static str,
81    /// models.dev provider ID to read models from (defaults to `dev_id` when `None`)
82    source_dev_id: Option<&'static str>,
83    /// Additional models.dev keys whose models are merged into this provider
84    extra_source_ids: &'static [&'static str],
85    /// Only include models whose ID passes this filter (None = include all)
86    model_filter: Option<fn(&str) -> bool>,
87    /// Provider-specific generated context window override
88    context_window_override: Option<ContextWindowOverride>,
89    /// Our Rust enum name (e.g. "Gemini")
90    enum_name: &'static str,
91    /// Our internal provider name used for parsing (e.g. "gemini")
92    parser_name: &'static str,
93    /// Human-readable provider name (e.g. "AWS Bedrock")
94    display_name: &'static str,
95    /// Env var our code actually checks (None for providers with complex credential chains)
96    env_var: Option<&'static str>,
97    /// OAuth provider ID for providers that require OAuth login (e.g. "codex")
98    oauth_provider_id: Option<&'static str>,
99    /// Default reasoning levels for models that support reasoning (empty = use standard 3)
100    default_reasoning_levels: &'static [&'static str],
101    /// When true, the inner catalog enum is named `{Enum}FoundationModel` and
102    /// `LlmModel::{Enum}` carries a hand-written `{Enum}Model` wrapper (defined
103    /// outside of codegen) that adds a `Profile(String)` fall-through plus any
104    /// provider-specific parsing policy. Used for Bedrock to accept arbitrary
105    /// inference profile IDs at runtime while keeping ARNs out of model identity.
106    is_hybrid_dynamic: bool,
107}
108
109impl ProviderConfig {
110    /// Shorthand for providers with default `source_dev_id`, `model_filter`, and `oauth_provider_id`.
111    const fn standard(
112        dev_id: &'static str,
113        enum_name: &'static str,
114        parser_name: &'static str,
115        display_name: &'static str,
116        env_var: Option<&'static str>,
117    ) -> Self {
118        Self {
119            dev_id,
120            source_dev_id: None,
121            extra_source_ids: &[],
122            model_filter: None,
123            context_window_override: None,
124            enum_name,
125            parser_name,
126            display_name,
127            env_var,
128            oauth_provider_id: None,
129            default_reasoning_levels: &["low", "medium", "high"],
130            is_hybrid_dynamic: false,
131        }
132    }
133
134    /// Inner catalog-enum name. For hybrid providers the outer `{enum_name}Model`
135    /// is a wrapper; the catalog enum is `{enum_name}FoundationModel`.
136    fn inner_enum_name(&self) -> String {
137        if self.is_hybrid_dynamic {
138            format!("{}FoundationModel", self.enum_name)
139        } else {
140            format!("{}Model", self.enum_name)
141        }
142    }
143
144    /// Outer enum name as referenced by `LlmModel::{enum_name}(...)`.
145    fn outer_enum_name(&self) -> String {
146        format!("{}Model", self.enum_name)
147    }
148
149    /// The models.dev key to look up in the JSON data.
150    fn json_key(&self) -> &'static str {
151        self.source_dev_id.unwrap_or(self.dev_id)
152    }
153}
154
155/// Dynamic provider — model name is user-supplied at runtime, no fixed enum
156#[allow(clippy::struct_field_names)]
157struct DynamicProviderConfig {
158    /// Rust variant name in `LlmModel` (e.g. "Ollama")
159    enum_name: &'static str,
160    /// Parser name used in "provider:model" strings (e.g. "ollama")
161    parser_name: &'static str,
162    /// Human-readable provider name (e.g. "Ollama")
163    display_name: &'static str,
164}
165
166const PROVIDERS: &[ProviderConfig] = &[
167    ProviderConfig::standard("anthropic", "Anthropic", "anthropic", "Anthropic", Some("ANTHROPIC_API_KEY")),
168    ProviderConfig {
169        dev_id: "codex",
170        source_dev_id: Some("openai"),
171        extra_source_ids: &[],
172        model_filter: Some(|id| id.contains("codex") || id.starts_with("gpt-5.") || id == "gpt-5"),
173        context_window_override: Some(codex_subscription_context_window),
174        enum_name: "Codex",
175        parser_name: "codex",
176        display_name: "Codex",
177        env_var: None,
178        oauth_provider_id: Some("codex"),
179        default_reasoning_levels: &["low", "medium", "high", "xhigh"],
180        is_hybrid_dynamic: false,
181    },
182    ProviderConfig::standard("deepseek", "DeepSeek", "deepseek", "DeepSeek", Some("DEEPSEEK_API_KEY")),
183    ProviderConfig::standard("google", "Gemini", "gemini", "Gemini", Some("GEMINI_API_KEY")),
184    ProviderConfig::standard("moonshotai", "Moonshot", "moonshot", "Moonshot", Some("MOONSHOT_API_KEY")),
185    ProviderConfig::standard("openai", "Openai", "openai", "OpenAI", Some("OPENAI_API_KEY")),
186    ProviderConfig::standard("openrouter", "OpenRouter", "openrouter", "OpenRouter", Some("OPENROUTER_API_KEY")),
187    ProviderConfig {
188        extra_source_ids: &["zai-coding-plan"],
189        ..ProviderConfig::standard("zai", "ZAi", "zai", "ZAI", Some("ZAI_API_KEY"))
190    },
191    ProviderConfig {
192        is_hybrid_dynamic: true,
193        ..ProviderConfig::standard("amazon-bedrock", "Bedrock", "bedrock", "AWS Bedrock", None)
194    },
195];
196
197const DYNAMIC_PROVIDERS: &[DynamicProviderConfig] = &[
198    DynamicProviderConfig { enum_name: "Ollama", parser_name: "ollama", display_name: "Ollama" },
199    DynamicProviderConfig { enum_name: "LlamaCpp", parser_name: "llamacpp", display_name: "LlamaCpp" },
200];
201
202const CODEX_SUBSCRIPTION_CONTEXT_WINDOW: u32 = 272_000;
203
204fn codex_subscription_context_window(model_id: &str, default_context_window: u32) -> u32 {
205    match model_id {
206        "gpt-5.5" | "gpt-5.4" | "gpt-5.4-mini" | "gpt-5.3-codex" | "gpt-5.2" | "codex-auto-review" => {
207            CODEX_SUBSCRIPTION_CONTEXT_WINDOW
208        }
209        _ => default_context_window,
210    }
211}
212
213#[derive(Debug, Clone)]
214struct ModelInfo {
215    variant_name: String,
216    model_id: String,
217    display_name: String,
218    context_window: u32,
219    reasoning_levels: Vec<String>,
220    input_modalities: Vec<String>,
221    supports_prompt_caching: bool,
222}
223
224type ProviderModels = BTreeMap<&'static str, Vec<ModelInfo>>;
225
226struct CodegenCtx {
227    provider_models: ProviderModels,
228}
229
230/// Output of the code generator.
231pub struct GeneratedOutput {
232    /// The generated Rust source (for `generated.rs`).
233    pub rust_source: String,
234    /// Per-provider markdown documentation keyed by provider identifier.
235    ///
236    /// Keys are provider `dev_ids` (e.g. `"anthropic"`, `"ollama"`) and values
237    /// are markdown strings suitable for `#![doc = include_str!(...)]`.
238    pub provider_docs: HashMap<String, String>,
239}
240
241#[derive(Debug, thiserror::Error)]
242pub enum CodegenError {
243    #[error("read: {0}")]
244    Read(#[from] std::io::Error),
245    #[error("parse: {0}")]
246    Parse(#[from] serde_json::Error),
247    #[error("Provider '{0}' not found in models.dev data")]
248    ProviderNotFound(String),
249}
250
251/// Run the codegen, returning the generated Rust source and per-provider docs.
252pub fn generate(models_json_path: &Path) -> Result<GeneratedOutput, CodegenError> {
253    let json_bytes = std::fs::read_to_string(models_json_path)?;
254    let data: ModelsDevData = serde_json::from_str(&json_bytes)?;
255
256    let provider_models = build_provider_models(&data)?;
257    let ctx = CodegenCtx { provider_models };
258    Ok(GeneratedOutput { rust_source: emit_generated_source(&ctx), provider_docs: emit_provider_docs(&ctx) })
259}
260
261fn build_provider_models(data: &ModelsDevData) -> Result<ProviderModels, CodegenError> {
262    let mut provider_models = ProviderModels::new();
263
264    for cfg in PROVIDERS {
265        let json_key = cfg.json_key();
266        let provider_data = data.get(json_key).ok_or_else(|| CodegenError::ProviderNotFound(json_key.to_string()))?;
267
268        let mut models: Vec<ModelInfo> = collect_models_from(cfg, &provider_data.models);
269
270        for &extra_key in cfg.extra_source_ids {
271            if let Some(extra_data) = data.get(extra_key) {
272                let extra = collect_models_from(cfg, &extra_data.models);
273                let existing_ids: std::collections::HashSet<String> =
274                    models.iter().map(|m| m.model_id.clone()).collect();
275                models.extend(extra.into_iter().filter(|m| !existing_ids.contains(&m.model_id)));
276            }
277        }
278
279        models.sort_by(|a, b| a.model_id.cmp(&b.model_id));
280        provider_models.insert(cfg.dev_id, models);
281    }
282
283    Ok(provider_models)
284}
285
286fn collect_models_from(cfg: &ProviderConfig, models: &HashMap<String, ModelData>) -> Vec<ModelInfo> {
287    models
288        .values()
289        .filter(|m| m.tool_call == Some(true))
290        .filter(|m| !is_alias(&m.id))
291        .filter(|m| cfg.model_filter.is_none_or(|f| f(&m.id)))
292        .map(|m| {
293            let reasoning_levels = if m.reasoning.unwrap_or(false) {
294                cfg.default_reasoning_levels.iter().map(|s| (*s).to_string()).collect()
295            } else {
296                Vec::new()
297            };
298            let input_modalities =
299                m.modalities.as_ref().map_or_else(|| vec!["text".to_string()], |md| md.input.clone());
300            let source_context_window = m.limit.as_ref().map_or(0, |l| l.context);
301            let context_window = cfg.context_window_override.map_or(source_context_window, |override_context_window| {
302                override_context_window(&m.id, source_context_window)
303            });
304            ModelInfo {
305                variant_name: model_id_to_variant(&m.id),
306                model_id: m.id.clone(),
307                display_name: m.name.clone(),
308                context_window,
309                reasoning_levels,
310                input_modalities,
311                supports_prompt_caching: m.cost.as_ref().is_some_and(CostData::has_prompt_caching),
312            }
313        })
314        .collect()
315}
316
317/// Returns true for "latest" alias IDs that just point to another model
318fn is_alias(id: &str) -> bool {
319    id.ends_with("-latest")
320}
321
322/// Convert a model ID like "claude-sonnet-4-5-20250929" into a `PascalCase` variant name.
323/// Treats `-`, `.`, `/`, and `:` as word separators.
324fn model_id_to_variant(id: &str) -> String {
325    let mut result = String::new();
326    let mut capitalize_next = true;
327
328    for ch in id.chars() {
329        if ch == '-' || ch == '.' || ch == '/' || ch == ':' {
330            capitalize_next = true;
331        } else if capitalize_next {
332            result.push(ch.to_ascii_uppercase());
333            capitalize_next = false;
334        } else {
335            result.push(ch);
336        }
337    }
338
339    if result.starts_with(|c: char| c.is_ascii_digit()) {
340        result.insert(0, '_');
341    }
342
343    result
344}
345
346fn emit_generated_source(ctx: &CodegenCtx) -> String {
347    let provider_enum = emit_provider_enum();
348    let provider_enum_impl = emit_provider_enum_impl();
349    let provider_enum_display = emit_provider_enum_display();
350    let provider_enum_fromstr = emit_provider_enum_fromstr();
351    let provider_enums = emit_provider_enums(&ctx.provider_models);
352    let provider_impls = emit_provider_impls(&ctx.provider_models);
353    let llm_model_enum = emit_llm_model_enum();
354    let from_impls = emit_from_impls();
355    let llm_model_impl = emit_llm_model_impl();
356    let display_impl = emit_display_impl();
357    let fromstr_impl = emit_fromstr_impl();
358
359    let file_tokens = quote! {
360        use std::borrow::Cow;
361        use std::sync::LazyLock;
362        use crate::ReasoningEffort;
363
364        #provider_enum
365        #provider_enum_impl
366        #provider_enum_display
367        #provider_enum_fromstr
368        #provider_enums
369        #provider_impls
370        #llm_model_enum
371        #from_impls
372        #llm_model_impl
373        #display_impl
374        #fromstr_impl
375    };
376
377    let file: syn::File = syn::parse2(file_tokens).expect("generated tokens parse as Rust");
378    let formatted = prettyplease::unparse(&file);
379    format!(
380        "// Auto-generated from models.dev — do not edit manually\n// Regenerated automatically by build.rs\n\n{formatted}"
381    )
382}
383
384fn emit_provider_enum() -> TokenStream {
385    let catalog_variants = PROVIDERS.iter().map(|cfg| format_ident!("{}", cfg.enum_name));
386    let dynamic_variants = DYNAMIC_PROVIDERS.iter().map(|d| format_ident!("{}", d.enum_name));
387    quote! {
388        /// Typed provider identifier — covers both catalog providers
389        /// (`Anthropic`, `Codex`, …) and dynamic providers whose model name is
390        /// user-supplied (`Ollama`, `LlamaCpp`).
391        #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
392        pub enum Provider {
393            #(#catalog_variants,)*
394            #(#dynamic_variants,)*
395        }
396    }
397}
398
399fn emit_provider_enum_impl() -> TokenStream {
400    let parser_arms = provider_match_arms(|cfg| cfg.parser_name, |d| d.parser_name);
401    let display_arms = provider_match_arms(|cfg| cfg.display_name, |d| d.display_name);
402
403    let env_var_some = PROVIDERS.iter().filter_map(|cfg| {
404        cfg.env_var.map(|var| {
405            let v = format_ident!("{}", cfg.enum_name);
406            quote! { Self::#v => Some(#var), }
407        })
408    });
409
410    let env_var_none = provider_or_pats(|cfg| cfg.env_var.is_none(), |_| true);
411    let oauth_some = PROVIDERS.iter().filter_map(|cfg| {
412        cfg.oauth_provider_id.map(|id| {
413            let v = format_ident!("{}", cfg.enum_name);
414            quote! { Self::#v => Some(#id), }
415        })
416    });
417    let oauth_none = provider_or_pats(|cfg| cfg.oauth_provider_id.is_none(), |_| true);
418
419    let is_local_true = provider_or_pats(|_| false, |_| true);
420    let is_local_false = provider_or_pats(|_| true, |_| false);
421
422    let all_variants = PROVIDERS
423        .iter()
424        .map(|cfg| format_ident!("{}", cfg.enum_name))
425        .chain(DYNAMIC_PROVIDERS.iter().map(|d| format_ident!("{}", d.enum_name)));
426
427    quote! {
428        impl Provider {
429            /// All providers — catalog and dynamic — in declaration order.
430            pub const ALL: &[Provider] = &[#(Self::#all_variants),*];
431
432            /// Parser name used in `provider:model` strings (e.g. `"anthropic"`).
433            pub fn parser_name(self) -> &'static str {
434                match self { #parser_arms }
435            }
436
437            /// Human-readable provider name (e.g. `"AWS Bedrock"`).
438            pub fn display_name(self) -> &'static str {
439                match self { #display_arms }
440            }
441
442            /// API-key env var the provider requires, if any.
443            pub fn required_env_var(self) -> Option<&'static str> {
444                match self {
445                    #(#env_var_some)*
446                    #env_var_none => None,
447                }
448            }
449
450            /// OAuth provider ID if this provider authenticates via OAuth.
451            pub fn oauth_provider_id(self) -> Option<&'static str> {
452                match self {
453                    #(#oauth_some)*
454                    #oauth_none => None,
455                }
456            }
457
458            /// Local providers run models on the user's machine — there's no
459            /// remote API to call and no env var to satisfy.
460            pub fn is_local(self) -> bool {
461                match self {
462                    #is_local_true => true,
463                    #is_local_false => false,
464                }
465            }
466        }
467    }
468}
469
470fn emit_provider_enum_display() -> TokenStream {
471    quote! {
472        impl std::fmt::Display for Provider {
473            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
474                f.write_str(self.parser_name())
475            }
476        }
477    }
478}
479
480fn emit_provider_enum_fromstr() -> TokenStream {
481    let catalog_arms = PROVIDERS.iter().map(|cfg| {
482        let v = format_ident!("{}", cfg.enum_name);
483        let name = cfg.parser_name;
484        quote! { #name => Ok(Self::#v), }
485    });
486
487    let dynamic_arms = DYNAMIC_PROVIDERS.iter().map(|d| {
488        let v = format_ident!("{}", d.enum_name);
489        let name = d.parser_name;
490        quote! { #name => Ok(Self::#v), }
491    });
492
493    quote! {
494        impl std::str::FromStr for Provider {
495            type Err = String;
496            fn from_str(s: &str) -> Result<Self, Self::Err> {
497                match s {
498                    #(#catalog_arms)*
499                    #(#dynamic_arms)*
500                    other => Err(format!("Unknown provider: '{other}'")),
501                }
502            }
503        }
504    }
505}
506
507fn emit_provider_enums(provider_models: &ProviderModels) -> TokenStream {
508    let enums = PROVIDERS.iter().map(|cfg| {
509        let inner = format_ident!("{}", cfg.inner_enum_name());
510        let variants = provider_models[cfg.dev_id].iter().map(|m| format_ident!("{}", m.variant_name));
511        quote! {
512            #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
513            pub enum #inner {
514                #(#variants,)*
515            }
516        }
517    });
518    quote! { #(#enums)* }
519}
520
521fn emit_provider_impls(provider_models: &ProviderModels) -> TokenStream {
522    let impls = PROVIDERS.iter().map(|cfg| {
523        let models = &provider_models[cfg.dev_id];
524        let enum_ident = format_ident!("{}", cfg.inner_enum_name());
525
526        let model_id_arms = models.iter().map(|m| {
527            let v = format_ident!("{}", m.variant_name);
528            let id = &m.model_id;
529            quote! { Self::#v => #id, }
530        });
531
532        let display_name_arms = grouped_arms(
533            models,
534            |m| m.display_name.clone(),
535            |m| {
536                let s = &m.display_name;
537                quote! { #s }
538            },
539        );
540
541        let context_window_arms =
542            grouped_arms(models, |m| m.context_window, |m| num_lit_with_underscores(m.context_window));
543
544        let reasoning_levels_arms = emit_reasoning_levels_arms(models);
545
546        let prompt_caching_arms = grouped_arms(
547            models,
548            |m| m.supports_prompt_caching,
549            |m| {
550                let b = m.supports_prompt_caching;
551                quote! { #b }
552            },
553        );
554
555        let modality_methods = ["image", "audio"].iter().map(|modality| {
556            let method = format_ident!("supports_{}", modality);
557            let mod_owned = (*modality).to_string();
558            let arms = grouped_arms(models, move |m| m.input_modalities.contains(&mod_owned), {
559                let mod_owned = (*modality).to_string();
560                move |m| {
561                    let b = m.input_modalities.contains(&mod_owned);
562                    quote! { #b }
563                }
564            });
565            quote! {
566                #[allow(clippy::too_many_lines)]
567                pub fn #method(self) -> bool {
568                    match self { #arms }
569                }
570            }
571        });
572
573        let all_variants = models.iter().map(|m| format_ident!("{}", m.variant_name));
574
575        let from_str_impl = emit_from_str_impl(&enum_ident, cfg.parser_name, models);
576
577        quote! {
578            impl #enum_ident {
579                #[allow(clippy::too_many_lines)]
580                fn model_id(self) -> &'static str {
581                    match self { #(#model_id_arms)* }
582                }
583
584                #[allow(clippy::too_many_lines)]
585                fn display_name(self) -> &'static str {
586                    match self { #display_name_arms }
587                }
588
589                #[allow(clippy::too_many_lines)]
590                fn context_window(self) -> u32 {
591                    match self { #context_window_arms }
592                }
593
594                #[allow(clippy::too_many_lines)]
595                pub fn reasoning_levels(self) -> &'static [ReasoningEffort] {
596                    match self { #reasoning_levels_arms }
597                }
598
599                pub fn supports_reasoning(self) -> bool {
600                    !self.reasoning_levels().is_empty()
601                }
602
603                #[allow(clippy::too_many_lines)]
604                pub fn supports_prompt_caching(self) -> bool {
605                    match self { #prompt_caching_arms }
606                }
607
608                #(#modality_methods)*
609
610                const ALL: &[#enum_ident] = &[#(Self::#all_variants),*];
611            }
612
613            #from_str_impl
614        }
615    });
616    quote! { #(#impls)* }
617}
618
619fn emit_from_str_impl(enum_ident: &proc_macro2::Ident, parser_name: &str, models: &[ModelInfo]) -> TokenStream {
620    let arms = models.iter().map(|m| {
621        let id = &m.model_id;
622        let v = format_ident!("{}", m.variant_name);
623        quote! { #id => Ok(Self::#v), }
624    });
625    let err_msg = format!("Unknown {parser_name} model: '{{s}}'");
626    quote! {
627        impl std::str::FromStr for #enum_ident {
628            type Err = String;
629
630            #[allow(clippy::too_many_lines)]
631            fn from_str(s: &str) -> Result<Self, Self::Err> {
632                match s {
633                    #(#arms)*
634                    _ => Err(format!(#err_msg)),
635                }
636            }
637        }
638    }
639}
640
641/// Emit match arms grouped by value to avoid clippy `match_same_arms`.
642fn grouped_arms<K, R>(
643    models: &[ModelInfo],
644    key_fn: impl Fn(&ModelInfo) -> K,
645    rhs_fn: impl Fn(&ModelInfo) -> R,
646) -> TokenStream
647where
648    K: Eq + Ord,
649    R: ToTokens,
650{
651    let mut groups: BTreeMap<K, Vec<&ModelInfo>> = BTreeMap::new();
652    for m in models {
653        groups.entry(key_fn(m)).or_default().push(m);
654    }
655    let arms = groups.values().map(|members| {
656        let pats = members.iter().map(|m| {
657            let v = format_ident!("{}", m.variant_name);
658            quote! { Self::#v }
659        });
660        let rhs = rhs_fn(members[0]);
661        quote! { #(#pats)|* => #rhs, }
662    });
663    quote! { #(#arms)* }
664}
665
666fn emit_reasoning_levels_arms(models: &[ModelInfo]) -> TokenStream {
667    grouped_arms(
668        models,
669        |m| m.reasoning_levels.clone(),
670        |m| {
671            if m.reasoning_levels.is_empty() {
672                quote! { &[] }
673            } else {
674                let items = m.reasoning_levels.iter().map(|l| {
675                    let variant = format_ident!("{}", level_str_to_variant(l));
676                    quote! { ReasoningEffort::#variant }
677                });
678                quote! { &[#(#items),*] }
679            }
680        },
681    )
682}
683
684/// Map a reasoning level string to its `ReasoningEffort` variant name.
685fn level_str_to_variant(level: &str) -> &'static str {
686    match level {
687        "low" => "Low",
688        "medium" => "Medium",
689        "high" => "High",
690        "xhigh" => "Xhigh",
691        other => panic!("Unknown reasoning level: {other}"),
692    }
693}
694
695fn emit_llm_model_enum() -> TokenStream {
696    let catalog_variants = PROVIDERS.iter().map(|cfg| {
697        let v = format_ident!("{}", cfg.enum_name);
698        let inner = format_ident!("{}Model", cfg.enum_name);
699        quote! { #v(#inner) }
700    });
701    let dynamic_variants = DYNAMIC_PROVIDERS.iter().map(|d| {
702        let v = format_ident!("{}", d.enum_name);
703        quote! { #v(String) }
704    });
705    quote! {
706        /// A model from a specific provider
707        #[derive(Debug, Clone, PartialEq, Eq, Hash)]
708        pub enum LlmModel {
709            #(#catalog_variants,)*
710            #(#dynamic_variants,)*
711        }
712    }
713}
714
715fn emit_from_impls() -> TokenStream {
716    let impls = PROVIDERS.iter().map(|cfg| {
717        let outer = format_ident!("{}Model", cfg.enum_name);
718        let v = format_ident!("{}", cfg.enum_name);
719        quote! {
720            impl From<#outer> for LlmModel {
721                fn from(m: #outer) -> Self {
722                    LlmModel::#v(m)
723                }
724            }
725        }
726    });
727    quote! { #(#impls)* }
728}
729
730fn emit_llm_model_impl() -> TokenStream {
731    let model_id = emit_llm_model_id();
732    let display_name = emit_llm_display_name();
733    let provider = emit_llm_provider();
734    let provider_enum = emit_llm_provider_enum();
735    let provider_display_name = emit_llm_provider_display_name();
736    let context_window = emit_llm_context_window();
737    let required_env_var = emit_llm_required_env_var();
738    let all_required_env_vars = emit_llm_all_required_env_vars();
739    let oauth_provider_id = emit_llm_oauth_provider_id();
740    let reasoning_levels = emit_llm_reasoning_levels();
741    let supports_reasoning = emit_llm_supports_reasoning();
742    let supports_prompt_caching = emit_llm_supports_prompt_caching();
743    let modality_methods = ["image", "audio"].iter().map(|m| emit_llm_supports_modality(m));
744    let all = emit_llm_all();
745
746    quote! {
747        impl LlmModel {
748            #model_id
749            #display_name
750            #provider
751            #provider_enum
752            #provider_display_name
753            #context_window
754            #required_env_var
755            #all_required_env_vars
756            #oauth_provider_id
757            #reasoning_levels
758            #supports_reasoning
759            #supports_prompt_caching
760            #(#modality_methods)*
761            #all
762        }
763    }
764}
765
766fn emit_llm_model_id() -> TokenStream {
767    let catalog_arms = PROVIDERS.iter().map(|cfg| {
768        let v = format_ident!("{}", cfg.enum_name);
769        if cfg.is_hybrid_dynamic {
770            quote! { Self::#v(m) => m.model_id(), }
771        } else {
772            quote! { Self::#v(m) => Cow::Borrowed(m.model_id()), }
773        }
774    });
775    let dyn_pats = dynamic_pattern_with_binding("s");
776    quote! {
777        /// Raw model ID (e.g. `claude-opus-4-6`, `llama3.2`)
778        pub fn model_id(&self) -> Cow<'static, str> {
779            match self {
780                #(#catalog_arms)*
781                #dyn_pats => Cow::Owned(s.clone()),
782            }
783        }
784    }
785}
786
787fn emit_llm_display_name() -> TokenStream {
788    let catalog_arms = PROVIDERS.iter().map(|cfg| {
789        let v = format_ident!("{}", cfg.enum_name);
790        if cfg.is_hybrid_dynamic {
791            quote! { Self::#v(m) => m.display_name(), }
792        } else {
793            quote! { Self::#v(m) => Cow::Borrowed(m.display_name()), }
794        }
795    });
796    let dyn_arms = DYNAMIC_PROVIDERS.iter().map(|d| {
797        let v = format_ident!("{}", d.enum_name);
798        let fmt = format!("{} {{s}}", d.enum_name);
799        quote! { Self::#v(s) => Cow::Owned(format!(#fmt)), }
800    });
801    quote! {
802        /// Human-readable display name (e.g. `Claude Opus 4.6`)
803        pub fn display_name(&self) -> Cow<'static, str> {
804            match self {
805                #(#catalog_arms)*
806                #(#dyn_arms)*
807            }
808        }
809    }
810}
811
812fn emit_llm_provider() -> TokenStream {
813    let arms = llm_match_arms_ignored(|cfg| cfg.parser_name, |d| d.parser_name);
814    quote! {
815        /// Provider identifier (e.g. `anthropic`)
816        pub fn provider(&self) -> &'static str {
817            match self { #arms }
818        }
819    }
820}
821
822fn emit_llm_provider_enum() -> TokenStream {
823    let arms = llm_match_arms_ignored(
824        |cfg| {
825            let v = format_ident!("{}", cfg.enum_name);
826            quote! { Provider::#v }
827        },
828        |d| {
829            let v = format_ident!("{}", d.enum_name);
830            quote! { Provider::#v }
831        },
832    );
833    quote! {
834        /// Typed provider identifier.
835        pub fn provider_enum(&self) -> Provider {
836            match self { #arms }
837        }
838    }
839}
840
841fn emit_llm_provider_display_name() -> TokenStream {
842    let arms = llm_match_arms_ignored(|cfg| cfg.display_name, |d| d.display_name);
843    quote! {
844        /// Human-readable provider name (e.g. `AWS Bedrock`)
845        pub fn provider_display_name(&self) -> &'static str {
846            match self { #arms }
847        }
848    }
849}
850
851fn emit_llm_context_window() -> TokenStream {
852    let catalog_arms = PROVIDERS.iter().map(|cfg| {
853        let v = format_ident!("{}", cfg.enum_name);
854        if cfg.is_hybrid_dynamic {
855            quote! { Self::#v(m) => m.context_window(), }
856        } else {
857            quote! { Self::#v(m) => Some(m.context_window()), }
858        }
859    });
860    let dyn_pats = dynamic_pattern_with_binding("_");
861    quote! {
862        /// Context window size in tokens (None for dynamic providers)
863        pub fn context_window(&self) -> Option<u32> {
864            match self {
865                #(#catalog_arms)*
866                #dyn_pats => None,
867            }
868        }
869    }
870}
871
872fn emit_llm_required_env_var() -> TokenStream {
873    let some_arms = PROVIDERS.iter().filter_map(|cfg| {
874        cfg.env_var.map(|var| {
875            let v = format_ident!("{}", cfg.enum_name);
876            quote! { Self::#v(_) => Some(#var), }
877        })
878    });
879    let none_pats = llm_or_pats(|cfg| cfg.env_var.is_none(), |_| true);
880    quote! {
881        /// Required env var for this model's provider (None for local providers)
882        pub fn required_env_var(&self) -> Option<&'static str> {
883            match self {
884                #(#some_arms)*
885                #none_pats => None,
886            }
887        }
888    }
889}
890
891fn emit_llm_all_required_env_vars() -> TokenStream {
892    let vars = PROVIDERS.iter().filter_map(|cfg| cfg.env_var);
893    quote! {
894        /// All provider API key env var names (deduplicated, static)
895        pub const ALL_REQUIRED_ENV_VARS: &[&str] = &[#(#vars),*];
896    }
897}
898
899fn emit_llm_oauth_provider_id() -> TokenStream {
900    let some_arms = PROVIDERS.iter().filter_map(|cfg| {
901        cfg.oauth_provider_id.map(|id| {
902            let v = format_ident!("{}", cfg.enum_name);
903            quote! { Self::#v(_) => Some(#id), }
904        })
905    });
906    let none_pats = llm_or_pats(|cfg| cfg.oauth_provider_id.is_none(), |_| true);
907    quote! {
908        /// OAuth provider ID if this model requires OAuth login (e.g. `"codex"`)
909        pub fn oauth_provider_id(&self) -> Option<&'static str> {
910            match self {
911                #(#some_arms)*
912                #none_pats => None,
913            }
914        }
915    }
916}
917
918fn emit_llm_reasoning_levels() -> TokenStream {
919    let body = llm_delegate_with_dynamic_default("reasoning_levels", &quote! { &[] });
920    quote! {
921        /// Reasoning levels supported by this model (empty if not a reasoning model)
922        pub fn reasoning_levels(&self) -> &'static [ReasoningEffort] {
923            #body
924        }
925    }
926}
927
928fn emit_llm_supports_reasoning() -> TokenStream {
929    quote! {
930        /// Whether this model supports reasoning/extended thinking
931        pub fn supports_reasoning(&self) -> bool {
932            !self.reasoning_levels().is_empty()
933        }
934    }
935}
936
937fn emit_llm_supports_prompt_caching() -> TokenStream {
938    let body = llm_delegate_with_dynamic_default("supports_prompt_caching", &quote! { false });
939    quote! {
940        /// Whether this model supports provider-side prompt caching
941        pub fn supports_prompt_caching(&self) -> bool {
942            #body
943        }
944    }
945}
946
947fn emit_llm_supports_modality(modality: &str) -> TokenStream {
948    let method = format!("supports_{modality}");
949    let method_ident = format_ident!("{}", method);
950    let doc = format!(" Whether this model supports {modality} input");
951    let body = llm_delegate_with_dynamic_default(&method, &quote! { false });
952    quote! {
953        #[doc = #doc]
954        pub fn #method_ident(&self) -> bool {
955            #body
956        }
957    }
958}
959
960fn emit_llm_all() -> TokenStream {
961    let pushes = PROVIDERS.iter().map(|cfg| {
962        let inner = format_ident!("{}", cfg.inner_enum_name());
963        let outer = format_ident!("{}", cfg.outer_enum_name());
964        let v = format_ident!("{}", cfg.enum_name);
965        if cfg.is_hybrid_dynamic {
966            quote! {
967                v.extend(#inner::ALL.iter().copied().map(#outer::Foundation).map(LlmModel::#v));
968            }
969        } else {
970            quote! {
971                v.extend(#inner::ALL.iter().copied().map(LlmModel::#v));
972            }
973        }
974    });
975    quote! {
976        /// All catalog models (excludes dynamic providers)
977        pub fn all() -> &'static [LlmModel] {
978            static ALL: LazyLock<Vec<LlmModel>> = LazyLock::new(|| {
979                let mut v = Vec::new();
980                #(#pushes)*
981                v
982            });
983            &ALL
984        }
985    }
986}
987
988fn emit_display_impl() -> TokenStream {
989    quote! {
990        impl std::fmt::Display for LlmModel {
991            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
992                write!(f, "{}:{}", self.provider(), self.model_id())
993            }
994        }
995    }
996}
997
998fn emit_fromstr_impl() -> TokenStream {
999    let catalog_arms = PROVIDERS.iter().map(|cfg| {
1000        let name = cfg.parser_name;
1001        let outer = format_ident!("{}Model", cfg.enum_name);
1002        let v = format_ident!("{}", cfg.enum_name);
1003        quote! { #name => model_str.parse::<#outer>().map(Self::#v), }
1004    });
1005    let dyn_arms = DYNAMIC_PROVIDERS.iter().map(|d| {
1006        let name = d.parser_name;
1007        let v = format_ident!("{}", d.enum_name);
1008        quote! { #name => Ok(Self::#v(model_str.to_string())), }
1009    });
1010    quote! {
1011        impl std::str::FromStr for LlmModel {
1012            type Err = String;
1013
1014            /// Parse a `provider:model` string into an `LlmModel`
1015            fn from_str(s: &str) -> Result<Self, Self::Err> {
1016                let (provider_str, model_str) = s.split_once(':').unwrap_or((s, ""));
1017                match provider_str {
1018                    #(#catalog_arms)*
1019                    #(#dyn_arms)*
1020                    _ => Err(format!("Unknown provider: '{provider_str}'")),
1021                }
1022            }
1023        }
1024    }
1025}
1026
1027/// Build a `Self::Ollama(b) | Self::LlamaCpp(b)` pattern for all dynamic providers.
1028fn dynamic_pattern_with_binding(binding: &str) -> TokenStream {
1029    let binding_ident = if binding == "_" {
1030        quote! { _ }
1031    } else {
1032        let b = format_ident!("{}", binding);
1033        quote! { #b }
1034    };
1035    let pats = DYNAMIC_PROVIDERS.iter().map(|d| {
1036        let v = format_ident!("{}", d.enum_name);
1037        quote! { Self::#v(#binding_ident) }
1038    });
1039    quote! { #(#pats)|* }
1040}
1041
1042/// Build `Self::A => va, Self::B => vb, ...` arms for every `Provider` variant
1043/// (catalog + dynamic). The `Provider` enum carries no inner data so there is
1044/// no binding.
1045fn provider_match_arms<V: ToTokens>(
1046    catalog_value: impl Fn(&ProviderConfig) -> V,
1047    dynamic_value: impl Fn(&DynamicProviderConfig) -> V,
1048) -> TokenStream {
1049    let catalog = PROVIDERS.iter().map(|cfg| {
1050        let v = format_ident!("{}", cfg.enum_name);
1051        let val = catalog_value(cfg);
1052        quote! { Self::#v => #val, }
1053    });
1054    let dynamic = DYNAMIC_PROVIDERS.iter().map(|d| {
1055        let v = format_ident!("{}", d.enum_name);
1056        let val = dynamic_value(d);
1057        quote! { Self::#v => #val, }
1058    });
1059    quote! { #(#catalog)* #(#dynamic)* }
1060}
1061
1062/// Build `Self::A | Self::B | ...` patterns selecting `Provider` variants by
1063/// predicate, across catalog + dynamic.
1064fn provider_or_pats(
1065    include_catalog: impl Fn(&ProviderConfig) -> bool,
1066    include_dynamic: impl Fn(&DynamicProviderConfig) -> bool,
1067) -> TokenStream {
1068    let catalog = PROVIDERS.iter().filter(|cfg| include_catalog(cfg)).map(|cfg| {
1069        let v = format_ident!("{}", cfg.enum_name);
1070        quote! { Self::#v }
1071    });
1072    let dynamic = DYNAMIC_PROVIDERS.iter().filter(|d| include_dynamic(d)).map(|d| {
1073        let v = format_ident!("{}", d.enum_name);
1074        quote! { Self::#v }
1075    });
1076    let pats = catalog.chain(dynamic);
1077    quote! { #(#pats)|* }
1078}
1079
1080/// Build `Self::A(_) => va, ...` arms for every `LlmModel` variant — the
1081/// wrapped inner value is ignored.
1082fn llm_match_arms_ignored<V: ToTokens>(
1083    catalog_value: impl Fn(&ProviderConfig) -> V,
1084    dynamic_value: impl Fn(&DynamicProviderConfig) -> V,
1085) -> TokenStream {
1086    let catalog = PROVIDERS.iter().map(|cfg| {
1087        let v = format_ident!("{}", cfg.enum_name);
1088        let val = catalog_value(cfg);
1089        quote! { Self::#v(_) => #val, }
1090    });
1091    let dynamic = DYNAMIC_PROVIDERS.iter().map(|d| {
1092        let v = format_ident!("{}", d.enum_name);
1093        let val = dynamic_value(d);
1094        quote! { Self::#v(_) => #val, }
1095    });
1096    quote! { #(#catalog)* #(#dynamic)* }
1097}
1098
1099/// Build `Self::A(_) | Self::B(_) | ...` patterns selecting `LlmModel`
1100/// variants by predicate, across catalog + dynamic.
1101fn llm_or_pats(
1102    include_catalog: impl Fn(&ProviderConfig) -> bool,
1103    include_dynamic: impl Fn(&DynamicProviderConfig) -> bool,
1104) -> TokenStream {
1105    let catalog = PROVIDERS.iter().filter(|cfg| include_catalog(cfg)).map(|cfg| {
1106        let v = format_ident!("{}", cfg.enum_name);
1107        quote! { Self::#v(_) }
1108    });
1109    let dynamic = DYNAMIC_PROVIDERS.iter().filter(|d| include_dynamic(d)).map(|d| {
1110        let v = format_ident!("{}", d.enum_name);
1111        quote! { Self::#v(_) }
1112    });
1113    let pats = catalog.chain(dynamic);
1114    quote! { #(#pats)|* }
1115}
1116
1117/// Build the body of an `LlmModel` method that delegates to a same-named
1118/// method on the inner catalog enum, with a single combined arm for all
1119/// dynamic providers.
1120fn llm_delegate_with_dynamic_default(method: &str, dynamic_value: &TokenStream) -> TokenStream {
1121    let method_ident = format_ident!("{}", method);
1122    let catalog_arms = PROVIDERS.iter().map(|cfg| {
1123        let v = format_ident!("{}", cfg.enum_name);
1124        quote! { Self::#v(m) => m.#method_ident(), }
1125    });
1126    let dyn_pat = dynamic_pattern_with_binding("_");
1127    quote! {
1128        match self {
1129            #(#catalog_arms)*
1130            #dyn_pat => #dynamic_value,
1131        }
1132    }
1133}
1134
1135/// Emit a `u32` literal with underscore separators (e.g. `200_000`).
1136fn num_lit_with_underscores(n: u32) -> TokenStream {
1137    format_number(n).parse().expect("formatted number parses as a token")
1138}
1139
1140/// Format a number with underscore separators (e.g. `200000` → `200_000`).
1141fn format_number(n: u32) -> String {
1142    let s = n.to_string();
1143    if s.len() <= 4 {
1144        return s;
1145    }
1146    let mut result = String::with_capacity(s.len() + s.len() / 3);
1147    for (i, ch) in s.chars().enumerate() {
1148        if i > 0 && (s.len() - i).is_multiple_of(3) {
1149            result.push('_');
1150        }
1151        result.push(ch);
1152    }
1153    result
1154}
1155
1156fn emit_provider_docs(ctx: &CodegenCtx) -> HashMap<String, String> {
1157    let mut docs = HashMap::new();
1158
1159    for cfg in PROVIDERS {
1160        let models = &ctx.provider_models[cfg.dev_id];
1161        let mut doc = String::new();
1162
1163        pushln(&mut doc, format!("`{}` LLM provider.", cfg.display_name));
1164        blank(&mut doc);
1165
1166        pushln(&mut doc, "# Authentication");
1167        blank(&mut doc);
1168        match cfg.env_var {
1169            Some(var) => pushln(&mut doc, format!("Set the `{var}` environment variable.")),
1170            None if cfg.oauth_provider_id.is_some() => {
1171                pushln(&mut doc, "This provider uses OAuth authentication.");
1172            }
1173            None => {
1174                pushln(
1175                    &mut doc,
1176                    "Uses the default AWS credential chain (environment variables, config files, IAM roles).",
1177                );
1178            }
1179        }
1180        blank(&mut doc);
1181
1182        pushln(&mut doc, "# Supported models");
1183        blank(&mut doc);
1184        pushln(&mut doc, "| Model ID | Name | Context | Reasoning | Image | Audio |");
1185        pushln(&mut doc, "|----------|------|---------|-----------|-------|-------|");
1186        for model in models {
1187            let ctx_str = format_context_window(model.context_window);
1188            let reasoning = if model.reasoning_levels.is_empty() { "" } else { "yes" };
1189            let image = if model.input_modalities.contains(&"image".to_string()) { "yes" } else { "" };
1190            let audio = if model.input_modalities.contains(&"audio".to_string()) { "yes" } else { "" };
1191            pushln(
1192                &mut doc,
1193                format!(
1194                    "| `{}` | `{}` | `{}` | {} | {} | {} |",
1195                    model.model_id, model.display_name, ctx_str, reasoning, image, audio
1196                ),
1197            );
1198        }
1199
1200        docs.insert(cfg.dev_id.to_string(), doc);
1201    }
1202
1203    for dyn_cfg in DYNAMIC_PROVIDERS {
1204        let mut doc = String::new();
1205        pushln(&mut doc, format!("`{}` LLM provider.", dyn_cfg.display_name));
1206        blank(&mut doc);
1207        pushln(
1208            &mut doc,
1209            format!("This provider accepts any model name at runtime (e.g. `{}:my-model`).", dyn_cfg.parser_name),
1210        );
1211        pushln(&mut doc, "No API key is required.");
1212        docs.insert(dyn_cfg.parser_name.to_string(), doc);
1213    }
1214
1215    docs
1216}
1217
1218/// Format a token count as human-readable (e.g. `1_000_000` → `1M`, `200_000` → `200k`).
1219fn format_context_window(tokens: u32) -> String {
1220    if tokens == 0 {
1221        return "unknown".to_string();
1222    }
1223    if tokens >= 1_000_000 && tokens.is_multiple_of(1_000_000) {
1224        format!("{}M", tokens / 1_000_000)
1225    } else if tokens >= 1_000 && tokens.is_multiple_of(1_000) {
1226        format!("{}k", tokens / 1_000)
1227    } else {
1228        format_number(tokens)
1229    }
1230}
1231
1232fn pushln(out: &mut String, line: impl AsRef<str>) {
1233    writeln!(out, "{}", line.as_ref()).expect("writing to String should not fail");
1234}
1235
1236fn blank(out: &mut String) {
1237    pushln(out, "");
1238}
1239
1240#[cfg(test)]
1241mod tests {
1242    use super::*;
1243    use serde_json::Value;
1244    use serde_json::json;
1245    use tempfile::NamedTempFile;
1246
1247    // ── Helper unit tests ────────────────────────────────────────────────────
1248
1249    #[test]
1250    fn model_id_to_variant_pascal_cases_segments() {
1251        assert_eq!(model_id_to_variant("claude-sonnet-4-5-20250929"), "ClaudeSonnet4520250929");
1252        assert_eq!(model_id_to_variant("gemini-2.5-flash"), "Gemini25Flash");
1253        assert_eq!(model_id_to_variant("deepseek-chat"), "DeepseekChat");
1254        assert_eq!(model_id_to_variant("glm-4.5"), "Glm45");
1255    }
1256
1257    #[test]
1258    fn model_id_to_variant_handles_slash_and_colon() {
1259        assert_eq!(model_id_to_variant("anthropic/claude-opus-4.6"), "AnthropicClaudeOpus46");
1260        assert_eq!(model_id_to_variant("openai/gpt-5.1-codex-max"), "OpenaiGpt51CodexMax");
1261        assert_eq!(model_id_to_variant("deepseek/deepseek-r1:free"), "DeepseekDeepseekR1Free");
1262    }
1263
1264    #[test]
1265    fn is_alias_detects_latest_suffix() {
1266        assert!(is_alias("claude-sonnet-4-5-latest"));
1267        assert!(is_alias("claude-3-7-sonnet-latest"));
1268        assert!(!is_alias("claude-sonnet-4-5-20250929"));
1269    }
1270
1271    #[test]
1272    fn codex_subscription_context_window_overrides_known_codex_models() {
1273        for model_id in ["gpt-5.5", "gpt-5.4", "gpt-5.4-mini", "gpt-5.3-codex", "gpt-5.2", "codex-auto-review"] {
1274            assert_eq!(codex_subscription_context_window(model_id, 1_050_000), 272_000);
1275        }
1276    }
1277
1278    #[test]
1279    fn codex_subscription_context_window_leaves_unknown_models_unchanged() {
1280        assert_eq!(codex_subscription_context_window("gpt-5.3-codex-spark", 128_000), 128_000);
1281        assert_eq!(codex_subscription_context_window("some-future-model", 400_000), 400_000);
1282    }
1283
1284    #[test]
1285    fn format_context_window_formats_correctly() {
1286        assert_eq!(format_context_window(1_000_000), "1M");
1287        assert_eq!(format_context_window(200_000), "200k");
1288        assert_eq!(format_context_window(8_000), "8k");
1289        assert_eq!(format_context_window(0), "unknown");
1290    }
1291
1292    #[test]
1293    fn level_str_to_variant_covers_all_reasoning_efforts() {
1294        for effort in utils::ReasoningEffort::all() {
1295            let _ = level_str_to_variant(effort.as_str());
1296        }
1297    }
1298
1299    #[test]
1300    fn build_sorts_models_and_filters_aliases_and_non_tool_call() {
1301        let mut data = minimal_models_dev_json();
1302        anthropic_models(
1303            &mut data,
1304            json!({
1305                "b-model": {"id": "b-model", "name": "B Model", "tool_call": true, "limit": {"context": 2000, "output": 0}},
1306                "a-model": {"id": "a-model", "name": "A Model", "tool_call": true, "limit": {"context": 1000, "output": 0}},
1307                "alpha-latest": {"id": "alpha-latest", "name": "Alias", "tool_call": true, "limit": {"context": 500, "output": 0}},
1308                "no-tools": {"id": "no-tools", "name": "No Tools", "tool_call": false, "limit": {"context": 500, "output": 0}}
1309            }),
1310        );
1311
1312        let models = build_from_value(&data);
1313        let ids: Vec<&str> = models["anthropic"].iter().map(|m| m.model_id.as_str()).collect();
1314        assert_eq!(ids, vec!["a-model", "b-model"]);
1315    }
1316
1317    #[test]
1318    fn build_extra_source_ids_merges_unique_models_into_provider() {
1319        let mut data = minimal_models_dev_json();
1320        zai_extra_models(
1321            &mut data,
1322            json!({
1323                "extra-model": {"id": "extra-model", "name": "Extra Model", "tool_call": true, "limit": {"context": 4000, "output": 0}}
1324            }),
1325        );
1326
1327        let models = build_from_value(&data);
1328        assert!(models["zai"].iter().any(|m| m.model_id == "extra-model"));
1329    }
1330
1331    #[test]
1332    fn build_extra_source_ids_does_not_duplicate_existing_models() {
1333        let mut data = minimal_models_dev_json();
1334        let shared = json!({
1335            "shared-model": {"id": "shared-model", "name": "Shared Model", "tool_call": true, "limit": {"context": 1000, "output": 0}}
1336        });
1337        insert_models(&mut data, "zai", shared.clone());
1338        insert_models(&mut data, "zai-coding-plan", shared);
1339
1340        let models = build_from_value(&data);
1341        let count = models["zai"].iter().filter(|m| m.model_id == "shared-model").count();
1342        assert_eq!(count, 1);
1343    }
1344
1345    #[test]
1346    fn build_derives_prompt_caching_from_cost_fields() {
1347        let mut data = minimal_models_dev_json();
1348        insert_models(
1349            &mut data,
1350            "amazon-bedrock",
1351            json!({
1352                "cached": {
1353                    "id": "cached", "name": "Cached", "tool_call": true,
1354                    "limit": {"context": 200_000, "output": 0},
1355                    "cost": {"input": 3.0, "output": 15.0, "cache_read": 0.3, "cache_write": 3.75}
1356                },
1357                "uncached": {
1358                    "id": "uncached", "name": "Uncached", "tool_call": true,
1359                    "limit": {"context": 200_000, "output": 0},
1360                    "cost": {"input": 3.0, "output": 15.0}
1361                }
1362            }),
1363        );
1364
1365        let models = build_from_value(&data);
1366        let bedrock = &models["amazon-bedrock"];
1367        let cached = bedrock.iter().find(|m| m.model_id == "cached").unwrap();
1368        let uncached = bedrock.iter().find(|m| m.model_id == "uncached").unwrap();
1369        assert!(cached.supports_prompt_caching);
1370        assert!(!uncached.supports_prompt_caching);
1371    }
1372
1373    #[test]
1374    fn build_assigns_codex_four_reasoning_levels() {
1375        let mut data = minimal_models_dev_json();
1376        insert_models(
1377            &mut data,
1378            "openai",
1379            json!({
1380                "gpt-5.4-codex": {
1381                    "id": "gpt-5.4-codex", "name": "GPT-5.4 Codex", "tool_call": true, "reasoning": true,
1382                    "limit": {"context": 200_000, "output": 0}
1383                }
1384            }),
1385        );
1386
1387        let models = build_from_value(&data);
1388        let codex_model = models["codex"].iter().find(|m| m.model_id == "gpt-5.4-codex").unwrap();
1389        assert_eq!(codex_model.reasoning_levels, vec!["low", "medium", "high", "xhigh"]);
1390    }
1391
1392    #[test]
1393    fn build_applies_codex_subscription_context_window_override() {
1394        let mut data = minimal_models_dev_json();
1395        insert_models(
1396            &mut data,
1397            "openai",
1398            json!({
1399                "gpt-5.5": {
1400                    "id": "gpt-5.5", "name": "GPT-5.5", "tool_call": true, "reasoning": true,
1401                    "limit": {"context": 1_050_000, "output": 128_000}
1402                }
1403            }),
1404        );
1405
1406        let models = build_from_value(&data);
1407        let codex = models["codex"].iter().find(|m| m.model_id == "gpt-5.5").unwrap();
1408        let openai = models["openai"].iter().find(|m| m.model_id == "gpt-5.5").unwrap();
1409        assert_eq!(codex.context_window, 272_000);
1410        assert_eq!(openai.context_window, 1_050_000);
1411    }
1412
1413    // ── Markdown docs ────────────────────────────────────────────────────────
1414
1415    #[test]
1416    fn generate_emits_provider_docs() {
1417        let mut data = minimal_models_dev_json();
1418        anthropic_models(
1419            &mut data,
1420            json!({
1421                "claude-test": {
1422                    "id": "claude-test", "name": "Claude Test", "tool_call": true, "reasoning": true,
1423                    "limit": {"context": 200_000, "output": 0},
1424                    "modalities": {"input": ["text", "image"]}
1425                }
1426            }),
1427        );
1428
1429        let tmp = NamedTempFile::new().unwrap();
1430        std::fs::write(tmp.path(), serde_json::to_string(&data).unwrap()).unwrap();
1431        let output = generate(tmp.path()).unwrap();
1432
1433        let anthropic_doc = &output.provider_docs["anthropic"];
1434        assert!(anthropic_doc.contains("`Anthropic` LLM provider."));
1435        assert!(anthropic_doc.contains("`ANTHROPIC_API_KEY`"));
1436        assert!(anthropic_doc.contains("| `claude-test` | `Claude Test` | `200k` | yes | yes |  |"));
1437
1438        let ollama_doc = &output.provider_docs["ollama"];
1439        assert!(ollama_doc.contains("`Ollama` LLM provider."));
1440        assert!(ollama_doc.contains("any model name at runtime"));
1441    }
1442
1443    fn build_from_value(data: &Value) -> ProviderModels {
1444        let parsed: ModelsDevData = serde_json::from_value(data.clone()).expect("parse fixture");
1445        build_provider_models(&parsed).expect("build provider models")
1446    }
1447
1448    fn anthropic_models(data: &mut Value, models: Value) {
1449        insert_models(data, "anthropic", models);
1450    }
1451
1452    fn zai_extra_models(data: &mut Value, models: Value) {
1453        insert_models(data, "zai-coding-plan", models);
1454    }
1455
1456    fn insert_models(data: &mut Value, provider_key: &str, models: Value) {
1457        let provider = data.as_object_mut().unwrap().get_mut(provider_key).unwrap().as_object_mut().unwrap();
1458        provider.insert("models".to_string(), models);
1459    }
1460
1461    fn minimal_models_dev_json() -> Value {
1462        let mut root = serde_json::Map::new();
1463        for cfg in PROVIDERS {
1464            let json_key = cfg.json_key();
1465            root.entry(json_key.to_string())
1466                .or_insert_with(|| json!({"id": json_key, "name": json_key, "env": [], "models": {}}));
1467            for &extra in cfg.extra_source_ids {
1468                root.entry(extra.to_string())
1469                    .or_insert_with(|| json!({"id": extra, "name": extra, "env": [], "models": {}}));
1470            }
1471        }
1472        Value::Object(root)
1473    }
1474}