1use std::time::Duration;
20
21use serde_json::{json, Value};
22
23use super::{
24 wake_prompt, ProviderKind, ProviderSession, SessionKind, WakeAdapter, WakeError, WakeLetter,
25 WakeOutcome,
26};
27
28pub const KIMI_SERVER_URL_ENV: &str = "M4A_KIMI_SERVER_URL";
30pub const KIMI_SERVER_TOKEN_ENV: &str = "M4A_KIMI_SERVER_TOKEN";
32
33pub struct KimiServerAdapter {
35 origin: Option<String>,
36 bearer: Option<String>,
37}
38
39impl KimiServerAdapter {
40 pub fn new(origin: Option<String>, bearer: Option<String>) -> Self {
42 let origin = origin
43 .map(|o| o.trim().trim_end_matches('/').to_string())
44 .filter(|o| is_loopback_http(o));
45 Self {
46 origin,
47 bearer: bearer.filter(|b| !b.is_empty()),
48 }
49 }
50
51 pub fn from_env() -> Self {
53 Self::new(
54 std::env::var(KIMI_SERVER_URL_ENV).ok(),
55 std::env::var(KIMI_SERVER_TOKEN_ENV).ok(),
56 )
57 }
58}
59
60impl KimiServerAdapter {
61 pub fn is_configured(&self) -> bool {
63 self.origin.is_some() && self.bearer.is_some()
64 }
65
66 pub fn into_parts(self) -> (Option<String>, Option<String>) {
68 (self.origin, self.bearer)
69 }
70}
71
72pub fn discover_local_server(home: Option<&std::path::Path>) -> Option<KimiServerAdapter> {
77 let home = home.map(std::path::Path::to_path_buf).or_else(|| {
78 std::env::var_os("KIMI_CODE_HOME")
79 .filter(|v| !v.is_empty())
80 .map(std::path::PathBuf::from)
81 .or_else(|| {
82 std::env::var_os(if cfg!(windows) { "USERPROFILE" } else { "HOME" })
83 .map(|h| std::path::PathBuf::from(h).join(".kimi-code"))
84 })
85 })?;
86 let mut instances: Vec<(std::time::SystemTime, Value)> =
87 std::fs::read_dir(home.join("server").join("instances"))
88 .ok()?
89 .filter_map(Result::ok)
90 .filter_map(|item| {
91 let modified = item.metadata().and_then(|m| m.modified()).ok()?;
92 let value: Value =
93 serde_json::from_slice(&std::fs::read(item.path()).ok()?).ok()?;
94 Some((modified, value))
95 })
96 .collect();
97 instances.sort_by(|a, b| b.0.cmp(&a.0));
98 let token = std::fs::read_to_string(home.join("server.token")).ok()?;
99 let token = token.trim().to_string();
100 for (_, instance) in instances {
101 if let Some(pid) = instance["pid"].as_u64() {
102 if cfg!(target_os = "linux") && !std::path::Path::new(&format!("/proc/{pid}")).exists()
103 {
104 continue;
105 }
106 }
107 let Some(port) = instance["port"].as_u64() else {
108 continue;
109 };
110 let host = instance["host"].as_str().unwrap_or("127.0.0.1");
111 let host = if host == "0.0.0.0" || host.is_empty() {
112 "127.0.0.1"
113 } else {
114 host
115 };
116 let adapter =
117 KimiServerAdapter::new(Some(format!("http://{host}:{port}")), Some(token.clone()));
118 if adapter.is_configured() {
119 return Some(adapter);
120 }
121 }
122 None
123}
124
125pub fn prompt_body(text: &str) -> Value {
127 json!({"content": [{"type": "text", "text": text}]})
128}
129
130fn is_loopback_http(origin: &str) -> bool {
131 let Some(rest) = origin.strip_prefix("http://") else {
132 return false;
133 };
134 let host = rest.split('/').next().unwrap_or("");
135 let host = host.rsplit_once(':').map(|(h, _)| h).unwrap_or(host);
136 matches!(host, "127.0.0.1" | "localhost" | "[::1]")
137}
138
139fn session_path_ok(id: &str) -> bool {
140 !id.is_empty()
141 && id
142 .chars()
143 .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
144}
145
146impl WakeAdapter for KimiServerAdapter {
147 fn kind(&self) -> SessionKind {
148 SessionKind::local(ProviderKind::KimiCode)
149 }
150
151 fn probe(&self, session: &ProviderSession) -> Result<(), WakeError> {
152 if self.origin.is_none() {
153 return Err(WakeError::Unavailable(
154 "kimi server url unset or not loopback".into(),
155 ));
156 }
157 if self.bearer.is_none() {
158 return Err(WakeError::Unavailable("kimi server token unset".into()));
159 }
160 if !session_path_ok(&session.session_id) {
161 return Err(WakeError::Unavailable(
162 "kimi session id is not a plain id".into(),
163 ));
164 }
165 Ok(())
166 }
167
168 fn wake(
169 &mut self,
170 session: &ProviderSession,
171 letter: &WakeLetter<'_>,
172 ) -> Result<WakeOutcome, WakeError> {
173 self.probe(session)?;
174 let origin = self.origin.as_deref().unwrap_or_default();
175 let bearer = self.bearer.as_deref().unwrap_or_default();
176 let url = format!("{origin}/api/v1/sessions/{}/prompts", session.session_id);
177 let client = reqwest::blocking::Client::builder()
178 .timeout(Duration::from_secs(20))
179 .build()
180 .map_err(|err| WakeError::Transport(err.to_string()))?;
181 let response = client
182 .post(url)
183 .bearer_auth(bearer)
184 .header(reqwest::header::CONTENT_TYPE, "application/json")
185 .body(prompt_body(&wake_prompt(session, letter)).to_string())
186 .send()
187 .map_err(|err| WakeError::Unavailable(format!("kimi server: {}", err.without_url())))?;
188 let status = response.status();
189 let body: Value = response
190 .bytes()
191 .ok()
192 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
193 .unwrap_or(Value::Null);
194 if !status.is_success() || body["code"] != json!(0) {
195 let msg = body["msg"]
196 .as_str()
197 .unwrap_or("")
198 .chars()
199 .take(120)
200 .collect::<String>();
201 return Err(WakeError::Transport(format!(
202 "kimi prompt: http {} {msg}",
203 status.as_u16()
204 )));
205 }
206 match body["data"]["status"].as_str() {
207 Some("running" | "queued") => Ok(WakeOutcome::Delivered),
208 Some(other) => Err(WakeError::Transport(format!("kimi prompt status {other}"))),
209 None => Err(WakeError::Transport("kimi prompt: no status".into())),
210 }
211 }
212}
213
214#[cfg(test)]
215mod tests {
216 use super::*;
217 use crate::tests::{letter, session};
218 use std::io::{BufRead, BufReader, Read, Write};
219 use std::net::TcpListener;
220
221 fn fake(status: &'static str) -> (String, std::thread::JoinHandle<(String, String)>) {
222 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
223 let origin = format!("http://{}", listener.local_addr().unwrap());
224 let handle = std::thread::spawn(move || {
225 let (stream, _) = listener.accept().unwrap();
226 let mut reader = BufReader::new(stream.try_clone().unwrap());
227 let mut request_line = String::new();
228 reader.read_line(&mut request_line).unwrap();
229 let mut auth = String::new();
230 let mut len = 0usize;
231 loop {
232 let mut line = String::new();
233 reader.read_line(&mut line).unwrap();
234 if line == "\r\n" {
235 break;
236 }
237 let lower = line.to_ascii_lowercase();
238 if lower.starts_with("authorization:") {
239 auth = line.trim().to_string();
240 }
241 if let Some(v) = lower.strip_prefix("content-length:") {
242 len = v.trim().parse().unwrap();
243 }
244 }
245 let mut body = vec![0; len];
246 reader.read_exact(&mut body).unwrap();
247 let reply = format!(
248 "{{\"code\":0,\"msg\":\"success\",\"data\":{{\"prompt_id\":\"p\",\"status\":\"{status}\"}}}}"
249 );
250 let mut stream = stream;
251 write!(
252 stream,
253 "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{reply}",
254 reply.len()
255 )
256 .unwrap();
257 (
258 format!("{} {auth}", request_line.trim()),
259 String::from_utf8(body).unwrap(),
260 )
261 });
262 (origin, handle)
263 }
264
265 #[test]
266 fn posts_prompt_with_bearer_and_accepts_running() {
267 let (origin, server) = fake("running");
268 let mut adapter = KimiServerAdapter::new(Some(origin), Some("tok".into()));
269 let mut s = session(SessionKind::local(ProviderKind::KimiCode));
270 s.session_id = "session_abc-1".into();
271 assert_eq!(
272 adapter.wake(&s, &letter("hi")).unwrap(),
273 WakeOutcome::Delivered
274 );
275 let (head, body) = server.join().unwrap();
276 assert!(head.starts_with("POST /api/v1/sessions/session_abc-1/prompts HTTP/1.1"));
277 assert!(head.ends_with("Bearer tok"));
278 let body: Value = serde_json::from_str(&body).unwrap();
279 assert_eq!(body["content"][0]["type"], "text");
280 assert!(body["content"][0]["text"].as_str().unwrap().ends_with("hi"));
281 }
282
283 #[test]
284 fn blocked_is_an_error_and_remote_origin_is_refused() {
285 let (origin, server) = fake("blocked");
286 let mut adapter = KimiServerAdapter::new(Some(origin), Some("tok".into()));
287 let s = session(SessionKind::local(ProviderKind::KimiCode));
288 assert!(adapter.wake(&s, &letter("hi")).is_err());
289 server.join().unwrap();
290 let remote = KimiServerAdapter::new(Some("http://10.0.0.2:58627".into()), Some("t".into()));
291 assert!(matches!(remote.probe(&s), Err(WakeError::Unavailable(_))));
292 let mut bad = session(SessionKind::local(ProviderKind::KimiCode));
293 bad.session_id = "../x".into();
294 let ok = KimiServerAdapter::new(Some("http://127.0.0.1:1".into()), Some("t".into()));
295 assert!(ok.probe(&bad).is_err());
296 }
297
298 #[test]
300 #[ignore]
301 fn live_prompt() {
302 let mut adapter = KimiServerAdapter::new(
303 std::env::var("M4A_LIVE_KIMI_URL").ok(),
304 std::env::var("M4A_LIVE_KIMI_TOKEN").ok(),
305 );
306 let mut s = session(SessionKind::local(ProviderKind::KimiCode));
307 s.session_id = std::env::var("M4A_LIVE_KIMI_SESSION").unwrap();
308 assert_eq!(
309 adapter.wake(&s, &letter("live probe")).unwrap(),
310 WakeOutcome::Delivered
311 );
312 }
313
314 #[test]
315 fn discovers_live_loopback_instance_and_skips_dead_pids() {
316 let home = std::env::temp_dir().join(format!("m4a-kimi-home-{}", std::process::id()));
317 let _ = std::fs::remove_dir_all(&home);
318 let inst = home.join("server").join("instances");
319 std::fs::create_dir_all(&inst).unwrap();
320 assert!(discover_local_server(Some(&home)).is_none());
321 std::fs::write(home.join("server.token"), "tok\n").unwrap();
322 std::fs::write(
323 inst.join("dead.json"),
324 r#"{"pid":4000000000,"host":"127.0.0.1","port":1}"#,
325 )
326 .unwrap();
327 if cfg!(target_os = "linux") {
328 assert!(discover_local_server(Some(&home)).is_none());
329 }
330 std::fs::write(
331 inst.join("live.json"),
332 format!(
333 r#"{{"pid":{},"host":"127.0.0.1","port":58627}}"#,
334 std::process::id()
335 ),
336 )
337 .unwrap();
338 let (url, token) = discover_local_server(Some(&home)).unwrap().into_parts();
339 assert_eq!(url.as_deref(), Some("http://127.0.0.1:58627"));
340 assert_eq!(token.as_deref(), Some("tok"));
341 let _ = std::fs::remove_dir_all(&home);
342 }
343}