feat: add Kimi Code provider with OAuth device flow authentication (#8466)

Signed-off-by: DaeHee Lee <lee111dae11@proton.me>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

crates/goose/src/providers/declarative/kimi.json

@@ -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
-}

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<ProviderRegistry> {
         registry.register::<GeminiOAuthProvider>(true);
         registry.register::<GithubCopilotProvider>(false);
         registry.register::<GoogleProvider>(true);
+        registry.register::<KimiCodeProvider>(true);
         registry.register::<LiteLLMProvider>(false);
         registry.register::<NanoGptProvider>(true);
         registry.register::<OllamaProvider>(true);
@@ -94,6 +96,10 @@ async fn init_registry() -> RwLock<ProviderRegistry> {
         "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);

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<Utc>,
+}
+
+#[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<KimiToken> {
+        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<Option<KimiToken>>,
+    #[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<Self> {
+        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<String> {
+        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<String> {
+        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<KimiToken> {
+        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<KimiToken> {
+        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<KimiToken> {
+        #[derive(Serialize)]
+        struct DeviceAuthReq<'a> {
+            client_id: &'a str,
+        }
+        #[derive(Deserialize)]
+        struct DeviceAuthResp {
+            device_code: String,
+            user_code: String,
+            verification_uri_complete: Option<String>,
+            verification_uri: String,
+            interval: Option<u64>,
+            expires_in: Option<u64>,
+        }
+
+        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<KimiToken> {
+        #[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<String>,
+            refresh_token: Option<String>,
+            expires_in: Option<i64>,
+            error: Option<String>,
+        }
+
+        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<KimiToken> {
+        #[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<String>,
+            expires_in: Option<i64>,
+        }
+
+        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<reqwest::Response, ProviderError> {
+        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<crate::config::ExtensionConfig>,
+    ) -> BoxFuture<'static, Result<Self::Provider>> {
+        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<MessageStream, ProviderError> {
+        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<Vec<String>, 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
+        );
+    }
+}

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;