claude_codex/providers/kimi/
mod.rs

1pub mod auth;
2pub mod client;
3pub mod count_tokens;
4pub mod translate;
5
6use async_trait::async_trait;
7use axum::Json;
8use axum::response::{IntoResponse, Response};
9use http::StatusCode;
10use std::time::{SystemTime, UNIX_EPOCH};
11
12use crate::anthropic::error::json_error;
13use crate::anthropic::schema::{CountTokensResponse, MessagesRequest};
14use crate::monitor::usage_from_anthropic_sse;
15use crate::provider::{CliHandlers, Provider, RequestContext};
16use crate::providers::kimi::auth::token_store::file_store;
17use crate::providers::kimi::translate::accumulate::accumulate_response;
18use crate::providers::kimi::translate::model_allowlist::{assert_allowed_model, resolve_model};
19use crate::providers::kimi::translate::request::{TranslateOptions, translate_request};
20use crate::providers::kimi::translate::stream::translate_stream_bytes;
21use crate::registry::KIMI_MODELS;
22
23fn now_ms() -> u64 {
24    SystemTime::now()
25        .duration_since(UNIX_EPOCH)
26        .unwrap_or_default()
27        .as_millis() as u64
28}
29
30pub struct KimiProvider;
31
32impl Default for KimiProvider {
33    fn default() -> Self {
34        Self::new()
35    }
36}
37
38impl KimiProvider {
39    pub fn new() -> Self {
40        Self
41    }
42}
43
44#[async_trait]
45impl Provider for KimiProvider {
46    fn name(&self) -> &'static str {
47        "kimi"
48    }
49
50    fn supported_models(&self) -> Vec<String> {
51        KIMI_MODELS.iter().map(|s| s.to_string()).collect()
52    }
53
54    fn cli(&self) -> &'static dyn CliHandlers {
55        &KIMI_CLI
56    }
57
58    async fn handle_messages(&self, body: MessagesRequest, ctx: RequestContext) -> Response {
59        let message_id = format!("msg_{}", uuid::Uuid::new_v4().to_string().replace('-', ""));
60        let want_stream = body.stream;
61        let model = body.model.as_deref().unwrap_or("kimi-for-coding");
62        let resolved = resolve_model(model);
63
64        if let Err(e) = assert_allowed_model(&resolved) {
65            return json_error(
66                StatusCode::BAD_REQUEST,
67                "invalid_request_error",
68                format!(
69                    "Model \"{model}\" resolves to unsupported model \"{}\"",
70                    e.model
71                ),
72            );
73        }
74        if let Some(monitor) = ctx.monitor.as_ref() {
75            monitor.model_resolved(&ctx.req_id, &resolved);
76        }
77
78        let translated = match translate_request(
79            &body,
80            TranslateOptions {
81                session_id: ctx.session_id.clone(),
82            },
83        ) {
84            Ok(t) => t,
85            Err(e) => {
86                return json_error(
87                    StatusCode::BAD_REQUEST,
88                    "invalid_request_error",
89                    e.to_string(),
90                );
91            }
92        };
93
94        // KimiHttpClient uses a blocking client whose lifecycle belongs on a
95        // blocking thread.
96        if let Some(monitor) = ctx.monitor.as_ref() {
97            monitor.upstream_started(&ctx.req_id);
98        }
99        let upstream = match tokio::task::spawn_blocking(move || {
100            let client = client::KimiHttpClient::new();
101            let result = client.post_kimi(&translated);
102            drop(client);
103            result
104        })
105        .await
106        {
107            Ok(Ok(r)) => r,
108            Ok(Err(e)) => {
109                return map_kimi_error_to_response(&e);
110            }
111            Err(join_err) => {
112                return json_error(
113                    StatusCode::BAD_GATEWAY,
114                    "api_error",
115                    format!("Blocking task join error: {join_err}"),
116                );
117            }
118        };
119
120        if want_stream {
121            let sse_bytes = match translate_stream_bytes(&upstream.body, &message_id, model) {
122                Ok(b) => b,
123                Err(e) => {
124                    return json_error(
125                        StatusCode::BAD_GATEWAY,
126                        "api_error",
127                        format!("Stream translation error: {e}"),
128                    );
129                }
130            };
131            if let Some(monitor) = ctx.monitor.as_ref() {
132                let (input_tokens, output_tokens) = usage_from_anthropic_sse(&sse_bytes);
133                monitor.stream_progress(
134                    &ctx.req_id,
135                    sse_bytes.len() as u64,
136                    count_sse_events(&sse_bytes),
137                    input_tokens,
138                    output_tokens,
139                );
140            }
141
142            let headers = [
143                (http::header::CONTENT_TYPE, "text/event-stream"),
144                (http::header::CACHE_CONTROL, "no-cache"),
145                (http::header::CONNECTION, "keep-alive"),
146            ];
147            (headers, sse_bytes).into_response()
148        } else {
149            match accumulate_response(&upstream.body, &message_id, model) {
150                Ok(json) => {
151                    if let Some(monitor) = ctx.monitor.as_ref() {
152                        monitor.usage_updated(
153                            &ctx.req_id,
154                            json.pointer("/usage/input_tokens").and_then(|v| v.as_u64()),
155                            json.pointer("/usage/output_tokens")
156                                .and_then(|v| v.as_u64()),
157                        );
158                    }
159                    (StatusCode::OK, Json(json)).into_response()
160                }
161                Err(e) => json_error(
162                    StatusCode::BAD_GATEWAY,
163                    "api_error",
164                    format!("Accumulation error: {e}"),
165                ),
166            }
167        }
168    }
169
170    async fn handle_count_tokens(&self, body: MessagesRequest, ctx: RequestContext) -> Response {
171        let model = body.model.as_deref().unwrap_or("kimi-for-coding");
172        let resolved = resolve_model(model);
173        if let Some(monitor) = ctx.monitor.as_ref() {
174            monitor.model_resolved(&ctx.req_id, &resolved);
175        }
176        let tokens = count_tokens::count_tokens(&body);
177        if let Some(monitor) = ctx.monitor.as_ref() {
178            monitor.usage_updated(&ctx.req_id, Some(tokens), None);
179        }
180        (
181            StatusCode::OK,
182            Json(CountTokensResponse {
183                input_tokens: tokens,
184            }),
185        )
186            .into_response()
187    }
188}
189
190fn count_sse_events(bytes: &[u8]) -> u64 {
191    String::from_utf8_lossy(bytes).matches("event:").count() as u64
192}
193
194fn map_kimi_error_to_response(err: &client::KimiError) -> Response {
195    match err.status {
196        401 | 403 => json_error(
197            StatusCode::UNAUTHORIZED,
198            "authentication_error",
199            err.detail.as_deref().unwrap_or("Authentication failed"),
200        ),
201        429 => {
202            let retry_after = err.retry_after.as_deref().unwrap_or("5");
203            let resp = json_error(
204                StatusCode::TOO_MANY_REQUESTS,
205                "rate_limit_error",
206                &err.message,
207            );
208            // Forward retry-after header
209            let headers = [(http::header::RETRY_AFTER, retry_after)];
210            (headers, resp).into_response()
211        }
212        _ => json_error(
213            StatusCode::BAD_GATEWAY,
214            "api_error",
215            err.detail.as_deref().unwrap_or("Upstream error"),
216        ),
217    }
218}
219
220// ---------------------------------------------------------------------------
221// CLI
222// ---------------------------------------------------------------------------
223
224pub(crate) struct KimiCli;
225
226impl CliHandlers for KimiCli {
227    fn login(&self) -> Result<(), anyhow::Error> {
228        let tokens = auth::login::run_device_login()?;
229        let store = file_store();
230        let manager = auth::manager::KimiAuthManager::new(store);
231        let saved = manager.persist_initial_tokens(&tokens)?;
232        println!("Auth saved in {}", manager.store.auth_path());
233        if let Some(ref uid) = saved.user_id {
234            println!("User: {uid}");
235        }
236        println!("Authentication complete");
237        Ok(())
238    }
239
240    fn device(&self) -> Result<(), anyhow::Error> {
241        self.login()
242    }
243
244    fn status(&self) -> Result<(), anyhow::Error> {
245        let store = file_store();
246        let stored = store.load_auth()?;
247        match stored {
248            Some(auth) => {
249                println!("Auth path: {}", store.auth_path());
250                println!("Authenticated: true");
251                if let Some(ref uid) = auth.user_id {
252                    println!("User: {uid}");
253                }
254                if let Some(ref scope) = auth.scope {
255                    println!("Scope: {scope}");
256                }
257                let remaining = auth.expires.saturating_sub(now_ms()) / 1000;
258                println!("Expires in {remaining}s");
259                Ok(())
260            }
261            None => {
262                anyhow::bail!("Not authenticated");
263            }
264        }
265    }
266
267    fn logout(&self) -> Result<(), anyhow::Error> {
268        let store = file_store();
269        store.clear_auth()?;
270        println!("Logged out");
271        Ok(())
272    }
273}
274
275pub(crate) static KIMI_CLI: KimiCli = KimiCli;
claude_codex/providers/kimi/mod.rs

claude_codex/providers/kimi/
mod.rs
