Fix Claude model limits and provider labels in model pickers

crates/galaxy_agent_rig/src/native.rs

@@ -1,4 +1,5 @@
-use std::sync::Arc;
+use std::collections::HashSet;
+use std::sync::{Arc, OnceLock};
 
 use async_trait::async_trait;
 use futures::StreamExt;
@@ -9,6 +10,7 @@ use galaxy_agent_core::{
 };
 use rig_core::client::{CompletionClient, ModelListingClient};
 use rig_core::providers::gemini;
+use rig_core::http_client::{self, HttpClientExt};
 
 use crate::anthropic_auth::{anthropic_client, prepare_subscription_request, restore_tool_names};
 use crate::request::build_completion_request;
@@ -19,22 +21,67 @@ pub struct RigModelInfo {
     pub id: String,
     pub display_name: String,
     pub context_size: Option<u32>,
+    pub max_output_tokens: Option<u32>,
 }
 
 pub async fn discover_anthropic_models(api_key: &str) -> Result<Vec<RigModelInfo>, String> {
     let client = anthropic_client(api_key)?;
-    let models = client
-        .list_models()
-        .await
-        .map_err(|error| error.to_string())?;
-    Ok(models
-        .into_iter()
-        .map(|model| RigModelInfo {
-            display_name: model.display_name().to_string(),
-            id: model.id,
-            context_size: model.context_length,
-        })
-        .collect())
+    let mut models = Vec::new();
+    let mut cursors = HashSet::new();
+    let mut path = "/v1/models?limit=1000".to_string();
+    loop {
+        let request = client.get(&path).and_then(|request| Ok(request.body(http_client::NoBody)?))
+            .map_err(|_| "Could not build the Anthropic models request.".to_string())?;
+        let response = client.send::<_, Vec<u8>>(request).await
+            .map_err(|_| "Could not retrieve Anthropic models.".to_string())?;
+        if !response.status().is_success() {
+            return Err(format!("Anthropic model discovery returned HTTP {}.", response.status()));
+        }
+        let body = response.into_body().await
+            .map_err(|_| "Could not read the Anthropic models response.".to_string())?;
+        let page: serde_json::Value = serde_json::from_slice(&body)
+            .map_err(|_| "Invalid Anthropic models response.".to_string())?;
+        let entries = page["data"].as_array()
+            .ok_or_else(|| "Anthropic models response is missing its model list.".to_string())?;
+        for entry in entries {
+            models.push(anthropic_model_info(entry)?);
+        }
+        match page["has_more"].as_bool() {
+            Some(false) => return Ok(models),
+            Some(true) => {},
+            None => return Err("Anthropic models response is missing pagination state.".to_string()),
+        }
+        let cursor = page["last_id"].as_str().filter(|cursor| {
+            !cursor.is_empty() && cursor.bytes().all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
+        }).ok_or_else(|| "Invalid Anthropic model pagination cursor.".to_string())?;
+        if !cursors.insert(cursor.to_string()) {
+            return Err("Anthropic model pagination did not advance.".to_string());
+        }
+        path = format!("/v1/models?limit=1000&after_id={cursor}");
+    }
+}
+
+fn positive_token_limit(value: &serde_json::Value) -> Option<u32> {
+    value.as_u64().and_then(|tokens| u32::try_from(tokens).ok()).filter(|tokens| *tokens > 0)
+}
+
+fn anthropic_model_info(entry: &serde_json::Value) -> Result<RigModelInfo, String> {
+    let id = entry["id"].as_str().filter(|id| !id.is_empty())
+        .ok_or_else(|| "Anthropic returned a model without an ID.".to_string())?;
+    // Read live limits first. Rig's generic model lister drops these fields.
+    // The models.dev snapshot is metadata only, never an availability allowlist.
+    static LIMITS: OnceLock<serde_json::Value> = OnceLock::new();
+    let limits = LIMITS.get_or_init(|| serde_json::from_str(include_str!("../data/anthropic_limits.json"))
+        .expect("bundled Anthropic limits must be valid JSON"));
+    let fallback = &limits["models"][id];
+    Ok(RigModelInfo {
+        id: id.to_string(),
+        display_name: entry["display_name"].as_str().unwrap_or(id).to_string(),
+        context_size: positive_token_limit(&entry["max_input_tokens"])
+            .or_else(|| positive_token_limit(&fallback["context"])),
+        max_output_tokens: positive_token_limit(&entry["max_tokens"])
+            .or_else(|| positive_token_limit(&fallback["output"])),
+    })
 }
 
 pub async fn discover_gemini_models(api_key: &str) -> Result<Vec<RigModelInfo>, String> {
@@ -50,6 +97,7 @@ pub async fn discover_gemini_models(api_key: &str) -> Result<Vec<RigModelInfo>,
             display_name: model.display_name().to_string(),
             id: model.id,
             context_size: model.context_length,
+            max_output_tokens: None,
         })
         .collect())
 }
@@ -76,6 +124,7 @@ pub fn vertex_ai_model_catalog() -> Vec<RigModelInfo> {
         id: id.to_string(),
         display_name: display_name.to_string(),
         context_size: None,
+        max_output_tokens: None,
     })
     .collect()
 }