diff --git a/crates/goose/src/providers/declarative/kimi.json b/crates/goose/src/providers/declarative/kimi.json deleted file mode 100644 index e1ccf9ad5a..0000000000 --- a/crates/goose/src/providers/declarative/kimi.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "name": "kimi", - "engine": "openai", - "display_name": "Kimi", - "description": "Kimi Code AI models optimized for coding tasks", - "api_key_env": "MOONSHOT_API_KEY", - "base_url": "https://api.moonshot.cn/v1/chat/completions", - "models": [ - {"name": "kimi-for-coding", "context_limit": 262144}, - {"name": "kimi-code", "context_limit": 262144} - ], - "supports_streaming": true -} diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index afb9db277e..c0c228bbcf 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -25,6 +25,7 @@ use super::{ gemini_oauth::GeminiOAuthProvider, githubcopilot::GithubCopilotProvider, google::GoogleProvider, + kimicode::KimiCodeProvider, litellm::LiteLLMProvider, nanogpt::NanoGptProvider, ollama::OllamaProvider, @@ -72,6 +73,7 @@ async fn init_registry() -> RwLock { registry.register::(true); registry.register::(false); registry.register::(true); + registry.register::(true); registry.register::(false); registry.register::(true); registry.register::(true); @@ -94,6 +96,10 @@ async fn init_registry() -> RwLock { "databricks", Arc::new(|| Box::pin(DatabricksProvider::cleanup())), ); + registry.set_cleanup( + "kimi_code", + Arc::new(|| Box::pin(KimiCodeProvider::cleanup())), + ); if let Err(e) = load_custom_providers_into_registry(&mut registry) { tracing::warn!("Failed to load custom providers: {}", e); diff --git a/crates/goose/src/providers/kimicode.rs b/crates/goose/src/providers/kimicode.rs new file mode 100644 index 0000000000..d0b721382b --- /dev/null +++ b/crates/goose/src/providers/kimicode.rs @@ -0,0 +1,881 @@ +use crate::config::paths::Paths; +use crate::config::Config; +use crate::session_context::SESSION_ID_HEADER; +use anyhow::{anyhow, Context, Result}; +use async_stream::try_stream; +use async_trait::async_trait; +use chrono::{DateTime, Duration, Utc}; +use futures::TryStreamExt; +use reqwest::header::{HeaderMap, HeaderValue}; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::io; +use std::time::Duration as StdDuration; +use tokio::pin; +use tokio_util::io::StreamReader; +use uuid::Uuid; + +use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata}; +use super::errors::ProviderError; +use super::formats::anthropic::{create_request, response_to_streaming_message}; +use super::openai_compatible::handle_status_openai_compat; +use super::retry::ProviderRetry; +use super::utils::RequestLog; +use crate::conversation::message::Message; +use crate::model::ModelConfig; +use futures::future::BoxFuture; +use rmcp::model::Tool; + +const KIMI_CODE_PROVIDER_NAME: &str = "kimi_code"; +pub const KIMI_CODE_DEFAULT_MODEL: &str = "kimi-k2.5"; +pub const KIMI_CODE_DEFAULT_FAST_MODEL: &str = "kimi-k2.5"; +pub const KIMI_CODE_KNOWN_MODELS: &[&str] = &["kimi-k2.5", "kimi-k2-thinking"]; + +const KIMI_CODE_DOC_URL: &str = "https://www.kimi.com/code/docs/en/"; +const KIMI_CODE_CLIENT_ID: &str = "17e5f671-d194-4dfb-9706-5516cb48c098"; +const KIMI_AUTH_HOST: &str = "https://auth.kimi.com"; +const KIMI_API_BASE: &str = "https://api.kimi.com/coding"; +const KIMI_MSH_PLATFORM: &str = "kimi_cli"; +const KIMI_MSH_VERSION: &str = "0.1.0"; + +/// Refresh the access token if it expires within this many seconds. +const REFRESH_THRESHOLD_SECS: i64 = 300; + +/// Fallback access-token lifetime when the server omits `expires_in`. +const DEFAULT_TOKEN_LIFETIME_SECS: i64 = 3600; + +/// Fallback device-code window when the server omits `expires_in` +/// from `device_authorization`. +const DEFAULT_DEVICE_CODE_LIFETIME_SECS: u64 = 300; + +/// Fallback poll interval when the server omits `interval`. +const DEFAULT_POLL_INTERVAL_SECS: u64 = 5; + +/// Extra seconds added to the poll interval after an RFC 8628 `slow_down`. +const SLOW_DOWN_BACKOFF_SECS: u64 = 5; + +/// Marker key written to the user config when OAuth completes successfully. +/// `check_provider_configured` (server) keys off this when an OAuth-flow +/// provider has no required secret env var. +const KIMI_CONFIGURED_MARKER: &str = "kimi_code_configured"; + +// ── Token persistence ──────────────────────────────────────────────────────── + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct KimiToken { + access_token: String, + refresh_token: String, + expires_at: DateTime, +} + +#[derive(Debug)] +struct TokenCache { + path: std::path::PathBuf, +} + +impl TokenCache { + fn new() -> Self { + Self { + path: Paths::in_config_dir("kimicode/token.json"), + } + } + + async fn load(&self) -> Option { + let raw = tokio::fs::read_to_string(&self.path).await.ok()?; + match serde_json::from_str(&raw) { + Ok(token) => Some(token), + Err(e) => { + tracing::warn!( + "kimicode token cache at {:?} is corrupted ({}); ignoring and re-authenticating", + self.path, + e + ); + None + } + } + } + + async fn save(&self, token: &KimiToken) -> Result<()> { + if let Some(parent) = self.path.parent() { + tokio::fs::create_dir_all(parent).await?; + } + tokio::fs::write(&self.path, serde_json::to_string(token)?).await?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + tokio::fs::set_permissions(&self.path, std::fs::Permissions::from_mode(0o600)).await?; + } + Ok(()) + } + + async fn clear(&self) -> Result<()> { + match tokio::fs::remove_file(&self.path).await { + Ok(()) => Ok(()), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(e) => Err(e.into()), + } + } +} + +// ── Provider ───────────────────────────────────────────────────────────────── + +#[derive(Debug, serde::Serialize)] +pub struct KimiCodeProvider { + #[serde(skip)] + client: Client, + #[serde(skip)] + token_cache: TokenCache, + #[serde(skip)] + cached_token: tokio::sync::Mutex>, + #[serde(skip)] + device_id: String, + #[serde(skip)] + auth_host: String, + #[serde(skip)] + api_base: String, + model: ModelConfig, + #[serde(skip)] + name: String, +} + +impl KimiCodeProvider { + pub async fn cleanup() -> Result<()> { + TokenCache::new().clear().await + } + + pub async fn from_env(model: ModelConfig) -> Result { + let model = model.with_fast(KIMI_CODE_DEFAULT_FAST_MODEL, KIMI_CODE_PROVIDER_NAME)?; + let client = Client::builder() + .timeout(StdDuration::from_secs(600)) + .build()?; + let device_id = Self::get_or_create_device_id().await?; + Ok(Self { + client, + token_cache: TokenCache::new(), + cached_token: tokio::sync::Mutex::new(None), + device_id, + auth_host: KIMI_AUTH_HOST.to_string(), + api_base: KIMI_API_BASE.to_string(), + model, + name: KIMI_CODE_PROVIDER_NAME.to_string(), + }) + } + + fn is_valid_device_id(id: &str) -> bool { + !id.is_empty() && HeaderValue::from_str(id).is_ok() + } + + async fn get_or_create_device_id() -> Result { + let path = Paths::in_config_dir("kimicode/device_id"); + if let Ok(raw) = tokio::fs::read_to_string(&path).await { + let id = raw.trim().to_string(); + if Self::is_valid_device_id(&id) { + return Ok(id); + } + tracing::warn!("kimicode device_id at {:?} is invalid; regenerating", path); + } + let id = Uuid::new_v4().to_string().replace('-', ""); + if let Some(parent) = path.parent() { + tokio::fs::create_dir_all(parent).await?; + } + tokio::fs::write(&path, &id).await?; + Ok(id) + } + + fn kimi_headers(&self) -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert( + "X-Msh-Platform", + HeaderValue::from_static(KIMI_MSH_PLATFORM), + ); + headers.insert("X-Msh-Version", HeaderValue::from_static(KIMI_MSH_VERSION)); + // Normally validated in `get_or_create_device_id`; skip the header if + // that validation was bypassed (e.g. test-constructed provider). + if let Ok(value) = HeaderValue::from_str(&self.device_id) { + headers.insert("X-Msh-Device-Id", value); + } + headers + } + + // ── Token management ───────────────────────────────────────────────────── + + /// Returns a valid access token, refreshing or re-authenticating as needed. + async fn get_access_token(&self) -> Result { + Ok(self.ensure_token().await?.access_token) + } + + /// Ensures we have a usable token, walking the cache → refresh → device-flow ladder. + async fn ensure_token(&self) -> Result { + let mut guard = self.cached_token.lock().await; + + if let Some(token) = guard.clone() { + if let Some(usable) = self.use_or_refresh(token).await { + *guard = Some(usable.clone()); + return Ok(usable); + } + } + + if let Some(token) = self.token_cache.load().await { + if let Some(usable) = self.use_or_refresh(token).await { + *guard = Some(usable.clone()); + return Ok(usable); + } + } + + tracing::info!("kimicode: starting OAuth device-flow login"); + let token = self.device_flow_login().await?; + self.token_cache.save(&token).await?; + *guard = Some(token.clone()); + Ok(token) + } + + /// Returns a usable token derived from `token`, or `None` if it is unusable. + /// On a successful refresh, the new token is also persisted to disk. + async fn use_or_refresh(&self, token: KimiToken) -> Option { + if token.expires_at - Utc::now() > Duration::seconds(REFRESH_THRESHOLD_SECS) { + return Some(token); + } + match self.do_refresh_token(&token.refresh_token).await { + Ok(refreshed) => { + tracing::debug!("kimicode: token refreshed"); + if let Err(e) = self.token_cache.save(&refreshed).await { + tracing::warn!("failed to persist refreshed kimicode token: {}", e); + } + Some(refreshed) + } + Err(e) => { + tracing::debug!("kimicode: token refresh failed: {}", e); + if token.expires_at > Utc::now() { + tracing::debug!("kimicode: falling back to still-unexpired token"); + Some(token) + } else { + None + } + } + } + } + + async fn device_flow_login(&self) -> Result { + #[derive(Serialize)] + struct DeviceAuthReq<'a> { + client_id: &'a str, + } + #[derive(Deserialize)] + struct DeviceAuthResp { + device_code: String, + user_code: String, + verification_uri_complete: Option, + verification_uri: String, + interval: Option, + expires_in: Option, + } + + let resp: DeviceAuthResp = self + .client + .post(format!("{}/api/oauth/device_authorization", self.auth_host)) + .headers(self.kimi_headers()) + .form(&DeviceAuthReq { + client_id: KIMI_CODE_CLIENT_ID, + }) + .send() + .await + .context("failed to request device authorization")? + .error_for_status() + .context("device authorization request failed")? + .json() + .await + .context("failed to parse device authorization response")?; + + let verify_url = resp + .verification_uri_complete + .as_deref() + .unwrap_or(&resp.verification_uri); + let interval = resp.interval.unwrap_or(DEFAULT_POLL_INTERVAL_SECS); + + if let Ok(mut clipboard) = arboard::Clipboard::new() { + let _ = clipboard.set_text(&resp.user_code); + } + if let Err(e) = webbrowser::open(verify_url) { + tracing::warn!("Failed to open browser: {}", e); + } + + // stderr so CLI workflows parsing stdout aren't interfered with. + eprintln!( + "Please visit {} and enter code {}", + verify_url, resp.user_code + ); + + let expires_in = resp.expires_in.unwrap_or(DEFAULT_DEVICE_CODE_LIFETIME_SECS); + self.poll_for_token(&resp.device_code, interval, expires_in) + .await + } + + async fn poll_for_token( + &self, + device_code: &str, + interval_secs: u64, + expires_in_secs: u64, + ) -> Result { + #[derive(Serialize)] + struct PollReq<'a> { + client_id: &'a str, + device_code: &'a str, + grant_type: &'static str, + } + #[derive(Deserialize, Debug)] + struct PollResp { + access_token: Option, + refresh_token: Option, + expires_in: Option, + error: Option, + } + + let deadline = + tokio::time::Instant::now() + tokio::time::Duration::from_secs(expires_in_secs); + let mut effective_interval = interval_secs; + loop { + if tokio::time::Instant::now() >= deadline { + return Err(anyhow!("timed out waiting for user authorization")); + } + tokio::time::sleep(tokio::time::Duration::from_secs(effective_interval)).await; + + let response = self + .client + .post(format!("{}/api/oauth/token", self.auth_host)) + .headers(self.kimi_headers()) + .form(&PollReq { + client_id: KIMI_CODE_CLIENT_ID, + device_code, + grant_type: "urn:ietf:params:oauth:grant-type:device_code", + }) + .send() + .await + .context("failed to poll for token")?; + + // RFC 8628 returns pending/slow_down as 4xx with a JSON error payload, + // so don't `error_for_status()` before parsing — but if the body is + // unparseable AND the status is non-2xx, surface the HTTP status. + let status = response.status(); + let bytes = response + .bytes() + .await + .context("failed to read token poll response")?; + let resp: PollResp = match serde_json::from_slice(&bytes) { + Ok(p) => p, + Err(e) => { + if !status.is_success() { + return Err(anyhow!( + "token poll HTTP {}: {}", + status, + String::from_utf8_lossy(&bytes) + )); + } + return Err( + anyhow::Error::new(e).context("failed to parse token poll response") + ); + } + }; + + if let Some(access_token) = resp.access_token { + // RFC 6749: refresh_token is optional in token responses. + // Kimi currently returns one, but be defensive for servers/ + // versions that do not. + let refresh_token = resp.refresh_token.unwrap_or_default(); + let expires_in = resp.expires_in.unwrap_or(DEFAULT_TOKEN_LIFETIME_SECS); + return Ok(KimiToken { + access_token, + refresh_token, + expires_at: Utc::now() + Duration::seconds(expires_in), + }); + } + + match resp.error.as_deref() { + Some("authorization_pending") => { + tracing::debug!("authorization pending, continuing to poll"); + } + // RFC 8628: client MUST increase polling interval by 5 seconds + Some("slow_down") => { + tracing::debug!("slow_down received, increasing poll interval"); + effective_interval += SLOW_DOWN_BACKOFF_SECS; + } + Some(err) => { + return Err(anyhow!("authorization failed: {}", err)); + } + None => { + tracing::debug!("unexpected poll response: no token and no error"); + } + } + } + } + + async fn do_refresh_token(&self, refresh_token: &str) -> Result { + #[derive(Serialize)] + struct RefreshReq<'a> { + client_id: &'a str, + grant_type: &'static str, + refresh_token: &'a str, + } + #[derive(Deserialize)] + struct RefreshResp { + access_token: String, + refresh_token: Option, + expires_in: Option, + } + + let resp: RefreshResp = self + .client + .post(format!("{}/api/oauth/token", self.auth_host)) + .headers(self.kimi_headers()) + .form(&RefreshReq { + client_id: KIMI_CODE_CLIENT_ID, + grant_type: "refresh_token", + refresh_token, + }) + .send() + .await + .context("failed to refresh token")? + .error_for_status() + .context("token refresh failed")? + .json() + .await + .context("failed to parse token refresh response")?; + + // RFC 6749 §6: the server MAY omit `refresh_token` from a refresh + // response, in which case the client should keep reusing the prior one. + let next_refresh_token = resp + .refresh_token + .unwrap_or_else(|| refresh_token.to_string()); + let expires_in = resp.expires_in.unwrap_or(DEFAULT_TOKEN_LIFETIME_SECS); + Ok(KimiToken { + access_token: resp.access_token, + refresh_token: next_refresh_token, + expires_at: Utc::now() + Duration::seconds(expires_in), + }) + } + + // ── HTTP ───────────────────────────────────────────────────────────────── + + async fn post( + &self, + session_id: Option<&str>, + payload: &Value, + ) -> Result { + let access_token = self.get_access_token().await.map_err(|e| { + ProviderError::Authentication(format!("Failed to get Kimi access token: {}", e)) + })?; + + let mut builder = self + .client + .post(format!("{}/v1/messages", self.api_base)) + .bearer_auth(access_token) + .headers(self.kimi_headers()) + .json(payload); + + if let Some(sid) = session_id { + builder = builder.header(SESSION_ID_HEADER, sid); + } + + builder + .send() + .await + .map_err(|e| ProviderError::RequestFailed(e.to_string())) + } +} + +// ── ProviderDef ─────────────────────────────────────────────────────────────── + +impl ProviderDef for KimiCodeProvider { + type Provider = Self; + + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + KIMI_CODE_PROVIDER_NAME, + "Kimi Code", + "Kimi Code AI models optimized for coding tasks", + KIMI_CODE_DEFAULT_MODEL, + KIMI_CODE_KNOWN_MODELS.to_vec(), + KIMI_CODE_DOC_URL, + // Marker key — the actual token lives in ~/.config/goose/kimicode/token.json. + // `oauth_flow=true` routes config through `configure_oauth`; + // readiness is tracked via the `kimi_code_configured` param. + vec![ConfigKey::new_oauth_device_code( + "KIMI_CODE_TOKEN", + true, + true, + None, + false, + )], + ) + .with_setup_steps(vec![ + "Run `goose configure` and select 'Kimi Code'", + "A browser window will open — log in to kimi.com and enter the displayed code", + "Once authorized, Goose will save your token automatically", + ]) + } + + fn from_env( + model: ModelConfig, + _extensions: Vec, + ) -> BoxFuture<'static, Result> { + Box::pin(Self::from_env(model)) + } +} + +// ── Provider trait ──────────────────────────────────────────────────────────── + +#[async_trait] +impl Provider for KimiCodeProvider { + fn get_name(&self) -> &str { + &self.name + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() + } + + async fn stream( + &self, + model_config: &ModelConfig, + session_id: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + let mut payload = create_request(model_config, system, messages, tools) + .map_err(|e| ProviderError::RequestFailed(e.to_string()))?; + payload + .as_object_mut() + .unwrap() + .insert("stream".to_string(), Value::Bool(true)); + + let mut log = RequestLog::start(model_config, &payload) + .map_err(|e| ProviderError::RequestFailed(e.to_string()))?; + + let response = self + .with_retry(|| async { + let resp = self.post(Some(session_id), &payload).await?; + handle_status_openai_compat(resp).await + }) + .await + .inspect_err(|e| { + let _ = log.error(e); + })?; + + let stream = response.bytes_stream().map_err(io::Error::other); + + Ok(Box::pin(try_stream! { + let stream_reader = StreamReader::new(stream); + let framed = tokio_util::codec::FramedRead::new( + stream_reader, + tokio_util::codec::LinesCodec::new(), + ) + .map_err(anyhow::Error::from); + + let message_stream = response_to_streaming_message(framed); + pin!(message_stream); + while let Some(message) = futures::StreamExt::next(&mut message_stream).await { + let (message, usage) = message.map_err(|e| { + ProviderError::RequestFailed(format!("Stream decode error: {}", e)) + })?; + log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?; + yield (message, usage); + } + })) + } + + async fn fetch_supported_models(&self) -> Result, ProviderError> { + Ok(KIMI_CODE_KNOWN_MODELS + .iter() + .map(|s| s.to_string()) + .collect()) + } + + async fn configure_oauth(&self) -> Result<(), ProviderError> { + self.ensure_token() + .await + .map_err(|e| ProviderError::Authentication(format!("OAuth flow failed: {}", e)))?; + + Config::global() + .set_param(KIMI_CONFIGURED_MARKER, Value::Bool(true)) + .map_err(|e| { + ProviderError::ExecutionError(format!( + "Failed to record kimi_code configured state: {}", + e + )) + })?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + use serde_json::json; + use wiremock::matchers::{body_string_contains, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + fn test_provider(server_uri: &str, device_id: &str) -> KimiCodeProvider { + KimiCodeProvider { + client: Client::new(), + token_cache: TokenCache { + path: std::env::temp_dir() + .join(format!("goose-kimicode-test-{}.json", Uuid::new_v4())), + }, + cached_token: tokio::sync::Mutex::new(None), + device_id: device_id.to_string(), + auth_host: server_uri.to_string(), + api_base: server_uri.to_string(), + model: ModelConfig::new(KIMI_CODE_DEFAULT_MODEL).unwrap(), + name: KIMI_CODE_PROVIDER_NAME.to_string(), + } + } + + // ── KimiToken serde ─────────────────────────────────────────────────────── + + #[test] + fn kimi_token_roundtrip() { + let token = KimiToken { + access_token: "acc_test".to_string(), + refresh_token: "ref_test".to_string(), + expires_at: Utc::now() + Duration::seconds(3600), + }; + let json = serde_json::to_string(&token).expect("serialize"); + let decoded: KimiToken = serde_json::from_str(&json).expect("deserialize"); + assert_eq!(decoded.access_token, token.access_token); + assert_eq!(decoded.refresh_token, token.refresh_token); + assert_eq!(decoded.expires_at.timestamp(), token.expires_at.timestamp()); + } + + // ── Headers ─────────────────────────────────────────────────────────────── + + #[tokio::test] + async fn kimi_headers_contains_required_fields() { + let provider = test_provider("http://localhost", "testdeviceid"); + let headers = provider.kimi_headers(); + assert_eq!( + headers.get("X-Msh-Platform").and_then(|v| v.to_str().ok()), + Some(KIMI_MSH_PLATFORM) + ); + assert_eq!( + headers.get("X-Msh-Version").and_then(|v| v.to_str().ok()), + Some(KIMI_MSH_VERSION) + ); + assert_eq!( + headers.get("X-Msh-Device-Id").and_then(|v| v.to_str().ok()), + Some("testdeviceid") + ); + } + + #[tokio::test] + async fn kimi_headers_skips_invalid_device_id_without_panic() { + // U+0000 is an invalid header byte; must not panic. + let provider = test_provider("http://localhost", "bad\u{0000}id"); + let headers = provider.kimi_headers(); + assert!(headers.get("X-Msh-Device-Id").is_none()); + assert!(headers.contains_key("X-Msh-Platform")); + } + + #[test] + fn validates_device_id_rejects_invalid_bytes() { + assert!(KimiCodeProvider::is_valid_device_id("abc123")); + assert!(!KimiCodeProvider::is_valid_device_id("")); + assert!(!KimiCodeProvider::is_valid_device_id("bad\u{0000}id")); + assert!(!KimiCodeProvider::is_valid_device_id("bad\nid")); + } + + // ── Metadata ────────────────────────────────────────────────────────────── + + #[test] + fn metadata_has_oauth_device_code_key() { + let meta = KimiCodeProvider::metadata(); + let key = meta + .config_keys + .iter() + .find(|k| k.name == "KIMI_CODE_TOKEN") + .expect("KIMI_CODE_TOKEN config key should exist"); + assert!(key.oauth_flow, "should be an OAuth flow key"); + assert!(key.device_code_flow, "should use device code flow"); + assert!(key.secret, "token should be stored securely"); + } + + #[test] + fn metadata_has_setup_steps() { + let meta = KimiCodeProvider::metadata(); + assert!( + !meta.setup_steps.is_empty(), + "setup_steps should be populated" + ); + } + + // ── Refresh / poll behavior ────────────────────────────────────────────── + + #[tokio::test] + async fn use_or_refresh_returns_fresh_token_without_calling_endpoint() { + let server = MockServer::start().await; + // Refusing all requests proves no network call was made. + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + + let provider = test_provider(&server.uri(), "abc"); + let fresh = KimiToken { + access_token: "acc".to_string(), + refresh_token: "ref".to_string(), + expires_at: Utc::now() + Duration::seconds(REFRESH_THRESHOLD_SECS + 600), + }; + + let usable = provider.use_or_refresh(fresh.clone()).await.unwrap(); + assert_eq!(usable.access_token, "acc"); + } + + #[tokio::test] + async fn use_or_refresh_falls_back_to_existing_token_when_refresh_fails() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(503)) + .mount(&server) + .await; + + let provider = test_provider(&server.uri(), "abc"); + // Inside refresh threshold but not yet expired. + let near_stale = KimiToken { + access_token: "still_good".to_string(), + refresh_token: "ref".to_string(), + expires_at: Utc::now() + Duration::seconds(60), + }; + + let usable = provider.use_or_refresh(near_stale).await.unwrap(); + assert_eq!(usable.access_token, "still_good"); + } + + #[tokio::test] + async fn use_or_refresh_returns_new_token_on_successful_refresh() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "new_access", + "refresh_token": "new_refresh", + "expires_in": 3600, + }))) + .mount(&server) + .await; + + let provider = test_provider(&server.uri(), "abc"); + let near_stale = KimiToken { + access_token: "old".to_string(), + refresh_token: "old_refresh".to_string(), + expires_at: Utc::now() + Duration::seconds(60), + }; + + let usable = provider.use_or_refresh(near_stale).await.unwrap(); + assert_eq!(usable.access_token, "new_access"); + assert_eq!(usable.refresh_token, "new_refresh"); + } + + #[tokio::test] + async fn poll_for_token_handles_authorization_pending_then_success() { + let server = MockServer::start().await; + + // First call: authorization_pending (returned as 400 per RFC 8628). + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": "authorization_pending", + }))) + .up_to_n_times(1) + .mount(&server) + .await; + + // Subsequent call: token issued. + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "the_token", + "refresh_token": "the_refresh", + "expires_in": 1800, + }))) + .mount(&server) + .await; + + let provider = test_provider(&server.uri(), "abc"); + let token = provider.poll_for_token("device-abc", 0, 30).await.unwrap(); + assert_eq!(token.access_token, "the_token"); + assert_eq!(token.refresh_token, "the_refresh"); + } + + #[tokio::test] + async fn poll_for_token_accepts_response_without_refresh_token() { + // RFC 6749: refresh_token is optional in token responses. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "access_only", + "expires_in": 1800, + }))) + .mount(&server) + .await; + + let provider = test_provider(&server.uri(), "abc"); + let token = provider.poll_for_token("device-abc", 0, 5).await.unwrap(); + assert_eq!(token.access_token, "access_only"); + assert_eq!(token.refresh_token, ""); + } + + #[tokio::test] + async fn use_or_refresh_preserves_refresh_token_when_server_omits_it() { + // RFC 6749 §6: if the refresh response omits `refresh_token`, the + // client should keep reusing the prior one. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "new_access", + "expires_in": 3600, + }))) + .mount(&server) + .await; + + let provider = test_provider(&server.uri(), "abc"); + let near_stale = KimiToken { + access_token: "old".to_string(), + refresh_token: "original_refresh".to_string(), + expires_at: Utc::now() + Duration::seconds(60), + }; + + let usable = provider.use_or_refresh(near_stale).await.unwrap(); + assert_eq!(usable.access_token, "new_access"); + assert_eq!(usable.refresh_token, "original_refresh"); + } + + #[tokio::test] + async fn poll_for_token_surfaces_http_error_on_unparseable_body() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(502).set_body_string("Bad Gateway")) + .mount(&server) + .await; + + let provider = test_provider(&server.uri(), "abc"); + let err = provider + .poll_for_token("device-abc", 0, 5) + .await + .unwrap_err(); + let msg = format!("{:#}", err); + assert!(msg.contains("502"), "expected status in error: {}", msg); + assert!( + msg.contains("Bad Gateway"), + "expected body in error: {}", + msg + ); + } +} diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index a49b5a483a..616e7c49ff 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -28,6 +28,7 @@ pub mod gemini_oauth; pub mod githubcopilot; pub mod google; mod init; +pub mod kimicode; pub mod litellm; #[cfg(feature = "local-inference")] pub mod local_inference;