| 1 | //! Moonshot/Kimi native search adapters. |
| 2 | |
| 3 | use anyhow::{Context, Result, bail}; |
| 4 | use reqwest::header::{HeaderName, HeaderValue}; |
| 5 | use serde_json::{Map, Value, json}; |
| 6 | use uuid::Uuid; |
| 7 | |
| 8 | use super::{ |
| 9 | ProviderNativeSearchClient, ProviderNativeSearchRequest, ProviderNativeSearchResponse, |
| 10 | bounded_answer, citation_from_url, citations_from_text, push_citation, |
| 11 | }; |
| 12 | use crate::{ |
| 13 | client::api_url, |
| 14 | config::{MOONSHOT_KIMI_K3_MODEL, moonshot_base_url_is_exact_kimi_code}, |
| 15 | }; |
| 16 | |
| 17 | const MAX_NATIVE_SEARCH_ROUNDS: usize = 4; |
| 18 | const MAX_NATIVE_SEARCH_TOOL_CALLS: usize = 8; |
| 19 | const NATIVE_SEARCH_MAX_COMPLETION_TOKENS: u32 = 4_096; |
| 20 | const WEB_SEARCH_FORMULA_URI: &str = "moonshot/web-search:latest"; |
| 21 | const WEB_SEARCH_FORMULA_FUNCTION: &str = "web_search"; |
| 22 | |
| 23 | pub(super) async fn search( |
| 24 | client: &ProviderNativeSearchClient, |
| 25 | request: &ProviderNativeSearchRequest, |
| 26 | ) -> Result<ProviderNativeSearchResponse> { |
| 27 | if moonshot_base_url_is_exact_kimi_code(&client.inner.base_url) { |
| 28 | // Exact Kimi Code membership endpoint only: the structured `/search` |
| 29 | // service is a first-party contract, so differently-cased or adjacent |
| 30 | // Kimi-hosted paths must not inherit it. |
| 31 | search_kimi_code(client, request).await |
| 32 | } else if client |
| 33 | .inner |
| 34 | .default_model |
| 35 | .trim() |
| 36 | .eq_ignore_ascii_case(MOONSHOT_KIMI_K3_MODEL) |
| 37 | { |
| 38 | search_formula(client, request).await |
| 39 | } else { |
| 40 | search_builtin(client, request).await |
| 41 | } |
| 42 | } |
| 43 | |
| 44 | async fn search_kimi_code( |
| 45 | client: &ProviderNativeSearchClient, |
| 46 | request: &ProviderNativeSearchRequest, |
| 47 | ) -> Result<ProviderNativeSearchResponse> { |
| 48 | let call_id = HeaderValue::from_str(&Uuid::new_v4().to_string()) |
| 49 | .context("failed to build Kimi search call id")?; |
| 50 | let url = format!("{}/search", client.inner.base_url.trim_end_matches('/')); |
| 51 | let payload = client |
| 52 | .post_json( |
| 53 | &url, |
| 54 | &json!({ "text_query": request.query }), |
| 55 | &[(HeaderName::from_static("x-msh-tool-call-id"), call_id)], |
| 56 | ) |
| 57 | .await?; |
| 58 | Ok(parse_kimi_code(&payload)) |
| 59 | } |
| 60 | |
| 61 | async fn search_builtin( |
| 62 | client: &ProviderNativeSearchClient, |
| 63 | request: &ProviderNativeSearchRequest, |
| 64 | ) -> Result<ProviderNativeSearchResponse> { |
| 65 | let tools = builtin_search_tools(); |
| 66 | let mut messages = vec![json!({ |
| 67 | "role": "user", |
| 68 | "content": super::search_prompt(request), |
| 69 | })]; |
| 70 | let mut tool_calls_executed = 0; |
| 71 | let url = api_url(&client.inner.base_url, "chat/completions"); |
| 72 | |
| 73 | for _ in 0..MAX_NATIVE_SEARCH_ROUNDS { |
| 74 | let body = json!({ |
| 75 | "model": client.inner.default_model, |
| 76 | "messages": &messages, |
| 77 | "tools": &tools, |
| 78 | "max_completion_tokens": NATIVE_SEARCH_MAX_COMPLETION_TOKENS, |
| 79 | "stream": false, |
| 80 | "thinking": { "type": "disabled" }, |
| 81 | }); |
| 82 | let payload = client.post_json(&url, &body, &[]).await?; |
| 83 | let choice = payload |
| 84 | .pointer("/choices/0") |
| 85 | .context("Kimi web search response omitted choices[0]")?; |
| 86 | let message = choice |
| 87 | .get("message") |
| 88 | .and_then(Value::as_object) |
| 89 | .context("Kimi web search response omitted assistant message")?; |
| 90 | if choice.get("finish_reason").and_then(Value::as_str) != Some("tool_calls") { |
| 91 | return Ok(parse_final_message(message)); |
| 92 | } |
| 93 | |
| 94 | messages.push(Value::Object(message.clone())); |
| 95 | let tool_calls = message |
| 96 | .get("tool_calls") |
| 97 | .and_then(Value::as_array) |
| 98 | .context("Kimi returned tool_calls finish reason without tool calls")?; |
| 99 | if tool_calls.is_empty() { |
| 100 | bail!("Kimi returned an empty native web-search tool call list"); |
| 101 | } |
| 102 | reserve_native_search_tool_calls(&mut tool_calls_executed, tool_calls.len())?; |
| 103 | for tool_call in tool_calls { |
| 104 | if tool_call.pointer("/function/name").and_then(Value::as_str) != Some("$web_search") { |
| 105 | bail!("Kimi native search requested an unexpected tool"); |
| 106 | } |
| 107 | let id = tool_call |
| 108 | .get("id") |
| 109 | .and_then(Value::as_str) |
| 110 | .context("Kimi native web-search call omitted id")?; |
| 111 | let arguments = tool_call |
| 112 | .pointer("/function/arguments") |
| 113 | .and_then(Value::as_str) |
| 114 | .context("Kimi native web-search call omitted arguments")?; |
| 115 | let _: Value = serde_json::from_str(arguments) |
| 116 | .context("Kimi native web-search arguments were not valid JSON")?; |
| 117 | messages.push(json!({ |
| 118 | "role": "tool", |
| 119 | "tool_call_id": id, |
| 120 | "name": "$web_search", |
| 121 | "content": arguments, |
| 122 | })); |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | bail!("Kimi native web search exceeded the bounded tool-call loop") |
| 127 | } |
| 128 | |
| 129 | async fn search_formula( |
| 130 | client: &ProviderNativeSearchClient, |
| 131 | request: &ProviderNativeSearchRequest, |
| 132 | ) -> Result<ProviderNativeSearchResponse> { |
| 133 | let formula_path = format!("formulas/{WEB_SEARCH_FORMULA_URI}"); |
| 134 | let tools_payload = client |
| 135 | .get_json(&api_url( |
| 136 | &client.inner.base_url, |
| 137 | &format!("{formula_path}/tools"), |
| 138 | )) |
| 139 | .await?; |
| 140 | let tools = formula_web_search_tools(&tools_payload)?; |
| 141 | let mut messages = vec![json!({ |
| 142 | "role": "user", |
| 143 | "content": super::search_prompt(request), |
| 144 | })]; |
| 145 | let mut tool_calls_executed = 0; |
| 146 | let chat_url = api_url(&client.inner.base_url, "chat/completions"); |
| 147 | let fiber_url = api_url(&client.inner.base_url, &format!("{formula_path}/fibers")); |
| 148 | |
| 149 | for _ in 0..MAX_NATIVE_SEARCH_ROUNDS { |
| 150 | let body = json!({ |
| 151 | "model": client.inner.default_model, |
| 152 | "messages": &messages, |
| 153 | "tools": &tools, |
| 154 | "max_completion_tokens": NATIVE_SEARCH_MAX_COMPLETION_TOKENS, |
| 155 | "stream": false, |
| 156 | }); |
| 157 | let payload = client.post_json(&chat_url, &body, &[]).await?; |
| 158 | let choice = payload |
| 159 | .pointer("/choices/0") |
| 160 | .context("Kimi Formula web search response omitted choices[0]")?; |
| 161 | let message = choice |
| 162 | .get("message") |
| 163 | .and_then(Value::as_object) |
| 164 | .context("Kimi Formula web search response omitted assistant message")?; |
| 165 | let Some(tool_calls) = message |
| 166 | .get("tool_calls") |
| 167 | .and_then(Value::as_array) |
| 168 | .filter(|calls| !calls.is_empty()) |
| 169 | else { |
| 170 | return Ok(parse_final_message(message)); |
| 171 | }; |
| 172 | |
| 173 | reserve_native_search_tool_calls(&mut tool_calls_executed, tool_calls.len())?; |
| 174 | messages.push(Value::Object(message.clone())); |
| 175 | for tool_call in tool_calls { |
| 176 | let id = tool_call |
| 177 | .get("id") |
| 178 | .and_then(Value::as_str) |
| 179 | .context("Kimi Formula web-search call omitted id")?; |
| 180 | let function = tool_call |
| 181 | .get("function") |
| 182 | .and_then(Value::as_object) |
| 183 | .context("Kimi Formula web-search call omitted function")?; |
| 184 | let name = function |
| 185 | .get("name") |
| 186 | .and_then(Value::as_str) |
| 187 | .context("Kimi Formula web-search call omitted function name")?; |
| 188 | if name != WEB_SEARCH_FORMULA_FUNCTION { |
| 189 | bail!("Kimi Formula web search requested an unexpected tool"); |
| 190 | } |
| 191 | let arguments = function |
| 192 | .get("arguments") |
| 193 | .and_then(Value::as_str) |
| 194 | .context("Kimi Formula web-search call omitted arguments")?; |
| 195 | let _: Value = serde_json::from_str(arguments) |
| 196 | .context("Kimi Formula web-search arguments were not valid JSON")?; |
| 197 | let fiber = client |
| 198 | .post_json( |
| 199 | &fiber_url, |
| 200 | &json!({ "name": name, "arguments": arguments }), |
| 201 | &[], |
| 202 | ) |
| 203 | .await?; |
| 204 | messages.push(json!({ |
| 205 | "role": "tool", |
| 206 | "tool_call_id": id, |
| 207 | "content": formula_fiber_result(&fiber)?, |
| 208 | })); |
| 209 | } |
| 210 | } |
| 211 | |
| 212 | bail!("Kimi Formula web search exceeded the bounded tool-call loop") |
| 213 | } |
| 214 | |
| 215 | fn reserve_native_search_tool_calls(executed: &mut usize, additional: usize) -> Result<()> { |
| 216 | let total = executed |
| 217 | .checked_add(additional) |
| 218 | .context("Kimi native web search tool-call count overflowed")?; |
| 219 | if total > MAX_NATIVE_SEARCH_TOOL_CALLS { |
| 220 | bail!( |
| 221 | "Kimi native web search exceeded the {MAX_NATIVE_SEARCH_TOOL_CALLS}-call safety limit" |
| 222 | ); |
| 223 | } |
| 224 | *executed = total; |
| 225 | Ok(()) |
| 226 | } |
| 227 | |
| 228 | fn formula_web_search_tools(payload: &Value) -> Result<Value> { |
| 229 | let tools = payload |
| 230 | .get("tools") |
| 231 | .and_then(Value::as_array) |
| 232 | .context("Kimi web-search Formula omitted tools")?; |
| 233 | if tools.len() != 1 |
| 234 | || tools[0].get("type").and_then(Value::as_str) != Some("function") |
| 235 | || tools[0].pointer("/function/name").and_then(Value::as_str) |
| 236 | != Some(WEB_SEARCH_FORMULA_FUNCTION) |
| 237 | { |
| 238 | bail!("Kimi web-search Formula returned an unexpected tool declaration"); |
| 239 | } |
| 240 | Ok(Value::Array(tools.clone())) |
| 241 | } |
| 242 | |
| 243 | fn formula_fiber_result(payload: &Value) -> Result<&str> { |
| 244 | if payload.get("status").and_then(Value::as_str) != Some("succeeded") { |
| 245 | bail!("Kimi web-search Formula fiber did not succeed"); |
| 246 | } |
| 247 | payload |
| 248 | .pointer("/context/output") |
| 249 | .or_else(|| payload.pointer("/context/encrypted_output")) |
| 250 | .and_then(Value::as_str) |
| 251 | .map(str::trim) |
| 252 | .filter(|result| !result.is_empty()) |
| 253 | .context("Kimi web-search Formula fiber omitted its result") |
| 254 | } |
| 255 | |
| 256 | fn builtin_search_tools() -> Value { |
| 257 | json!([{ |
| 258 | "type": "builtin_function", |
| 259 | "function": { "name": "$web_search" } |
| 260 | }]) |
| 261 | } |
| 262 | |
| 263 | fn parse_kimi_code(payload: &Value) -> ProviderNativeSearchResponse { |
| 264 | let mut citations = Vec::new(); |
| 265 | if let Some(results) = payload.get("search_results").and_then(Value::as_array) { |
| 266 | for result in results { |
| 267 | let Some(url) = result.get("url").and_then(Value::as_str) else { |
| 268 | continue; |
| 269 | }; |
| 270 | let title = result |
| 271 | .get("title") |
| 272 | .and_then(Value::as_str) |
| 273 | .map(str::to_string); |
| 274 | let snippet = result |
| 275 | .get("snippet") |
| 276 | .and_then(Value::as_str) |
| 277 | .map(str::to_string); |
| 278 | let published = result |
| 279 | .get("date") |
| 280 | .and_then(Value::as_str) |
| 281 | .filter(|value| !value.is_empty()) |
| 282 | .map(str::to_string); |
| 283 | push_citation( |
| 284 | &mut citations, |
| 285 | citation_from_url(url, title, snippet, published), |
| 286 | ); |
| 287 | } |
| 288 | } |
| 289 | ProviderNativeSearchResponse { |
| 290 | answer: None, |
| 291 | citations, |
| 292 | } |
| 293 | } |
| 294 | |
| 295 | fn parse_final_message(message: &Map<String, Value>) -> ProviderNativeSearchResponse { |
| 296 | let answer = message |
| 297 | .get("content") |
| 298 | .and_then(Value::as_str) |
| 299 | .map(str::trim) |
| 300 | .filter(|text| !text.is_empty()) |
| 301 | .map(str::to_string); |
| 302 | let citations = answer |
| 303 | .as_deref() |
| 304 | .map(citations_from_text) |
| 305 | .unwrap_or_default(); |
| 306 | ProviderNativeSearchResponse { |
| 307 | answer: bounded_answer(answer.into_iter().collect()), |
| 308 | citations, |
| 309 | } |
| 310 | } |
| 311 | |
| 312 | #[cfg(test)] |
| 313 | mod tests { |
| 314 | use super::*; |
| 315 | use crate::config::{Config, ProviderConfig, ProvidersConfig}; |
| 316 | use wiremock::matchers::{body_partial_json, body_string_contains, header, method, path}; |
| 317 | use wiremock::{Mock, MockServer, ResponseTemplate}; |
| 318 | |
| 319 | fn request() -> ProviderNativeSearchRequest { |
| 320 | ProviderNativeSearchRequest { |
| 321 | query: "current release".to_string(), |
| 322 | max_results: 3, |
| 323 | domains: Vec::new(), |
| 324 | } |
| 325 | } |
| 326 | |
| 327 | #[test] |
| 328 | fn kimi_code_request_and_structured_response_contract() { |
| 329 | let body = json!({ "text_query": request().query }); |
| 330 | assert_eq!(body["text_query"], "current release"); |
| 331 | assert_eq!(body.as_object().map(serde_json::Map::len), Some(1)); |
| 332 | |
| 333 | let parsed = parse_kimi_code(&json!({ |
| 334 | "search_results": [{ |
| 335 | "title": "Kimi", |
| 336 | "url": "https://example.com/kimi", |
| 337 | "snippet": "Summary", |
| 338 | "date": "2026-08-28" |
| 339 | }] |
| 340 | })); |
| 341 | assert_eq!(parsed.citations.len(), 1); |
| 342 | assert_eq!(parsed.citations[0].snippet.as_deref(), Some("Summary")); |
| 343 | assert_eq!(parsed.citations[0].published.as_deref(), Some("2026-08-28")); |
| 344 | } |
| 345 | |
| 346 | #[test] |
| 347 | fn direct_search_contracts_are_bounded() { |
| 348 | let tools = builtin_search_tools(); |
| 349 | assert_eq!(tools[0]["function"]["name"], "$web_search"); |
| 350 | assert_eq!(NATIVE_SEARCH_MAX_COMPLETION_TOKENS, 4_096); |
| 351 | |
| 352 | let formula_tools = formula_web_search_tools(&json!({ |
| 353 | "tools": [{ |
| 354 | "type": "function", |
| 355 | "function": { "name": "web_search" } |
| 356 | }] |
| 357 | })) |
| 358 | .expect("formula tools"); |
| 359 | assert_eq!(formula_tools[0]["function"]["name"], "web_search"); |
| 360 | assert_eq!( |
| 361 | formula_fiber_result(&json!({ |
| 362 | "status": "succeeded", |
| 363 | "context": { "encrypted_output": "encrypted result" } |
| 364 | })) |
| 365 | .expect("formula result"), |
| 366 | "encrypted result" |
| 367 | ); |
| 368 | } |
| 369 | |
| 370 | #[test] |
| 371 | fn native_search_tool_call_limit_is_total_not_per_round() { |
| 372 | let mut executed = 0; |
| 373 | reserve_native_search_tool_calls(&mut executed, 4).expect("first rounds"); |
| 374 | reserve_native_search_tool_calls(&mut executed, 4).expect("final allowed round"); |
| 375 | assert_eq!(executed, MAX_NATIVE_SEARCH_TOOL_CALLS); |
| 376 | assert!(reserve_native_search_tool_calls(&mut executed, 1).is_err()); |
| 377 | } |
| 378 | |
| 379 | #[test] |
| 380 | fn kimi_code_dispatch_reuses_the_exact_route_matcher() { |
| 381 | for route in [ |
| 382 | "https://api.kimi.com/coding/v1", |
| 383 | "https://api.kimi.com/coding/v1/", |
| 384 | "HTTPS://API.KIMI.COM/coding/v1", |
| 385 | ] { |
| 386 | assert!( |
| 387 | moonshot_base_url_is_exact_kimi_code(route), |
| 388 | "{route} is the membership endpoint" |
| 389 | ); |
| 390 | } |
| 391 | for neighboring_route in [ |
| 392 | // A case-variant path is a different route, not the official one. |
| 393 | "https://API.KIMI.COM/CODING/V1", |
| 394 | "https://api.kimi.com/coding/v2", |
| 395 | "https://api.kimi.com/coding", |
| 396 | "http://api.kimi.com/coding/v1", |
| 397 | "https://api.moonshot.ai/v1", |
| 398 | ] { |
| 399 | assert!( |
| 400 | !moonshot_base_url_is_exact_kimi_code(neighboring_route), |
| 401 | "{neighboring_route} must not reach the Kimi Code /search service" |
| 402 | ); |
| 403 | } |
| 404 | } |
| 405 | |
| 406 | #[tokio::test] |
| 407 | async fn k3_formula_executes_tool_fiber_and_returns_citations() { |
| 408 | let server = MockServer::start().await; |
| 409 | Mock::given(method("GET")) |
| 410 | .and(path("/v1/formulas/moonshot/web-search:latest/tools")) |
| 411 | .and(header("authorization", "Bearer moonshot-test-key")) |
| 412 | .respond_with(ResponseTemplate::new(200).set_body_json(json!({ |
| 413 | "tools": [{ |
| 414 | "type": "function", |
| 415 | "function": { |
| 416 | "name": "web_search", |
| 417 | "description": "Search the web", |
| 418 | "parameters": { |
| 419 | "type": "object", |
| 420 | "properties": { "query": { "type": "string" } }, |
| 421 | "required": ["query"] |
| 422 | } |
| 423 | } |
| 424 | }] |
| 425 | }))) |
| 426 | .expect(1) |
| 427 | .mount(&server) |
| 428 | .await; |
| 429 | Mock::given(method("POST")) |
| 430 | .and(path("/v1/chat/completions")) |
| 431 | .and(header("authorization", "Bearer moonshot-test-key")) |
| 432 | .and(body_partial_json(json!({ |
| 433 | "model": "kimi-k3", |
| 434 | "max_completion_tokens": 4096, |
| 435 | "tools": [{ |
| 436 | "type": "function", |
| 437 | "function": { "name": "web_search" } |
| 438 | }] |
| 439 | }))) |
| 440 | .respond_with(ResponseTemplate::new(200).set_body_json(json!({ |
| 441 | "choices": [{ |
| 442 | "finish_reason": "tool_calls", |
| 443 | "message": { |
| 444 | "role": "assistant", |
| 445 | "content": "", |
| 446 | "tool_calls": [{ |
| 447 | "id": "web_search:0", |
| 448 | "type": "function", |
| 449 | "function": { |
| 450 | "name": "web_search", |
| 451 | "arguments": "{\"query\":\"current release\"}" |
| 452 | } |
| 453 | }] |
| 454 | } |
| 455 | }] |
| 456 | }))) |
| 457 | .up_to_n_times(1) |
| 458 | .expect(1) |
| 459 | .mount(&server) |
| 460 | .await; |
| 461 | Mock::given(method("POST")) |
| 462 | .and(path("/v1/formulas/moonshot/web-search:latest/fibers")) |
| 463 | .and(header("authorization", "Bearer moonshot-test-key")) |
| 464 | .and(body_partial_json(json!({ |
| 465 | "name": "web_search", |
| 466 | "arguments": "{\"query\":\"current release\"}" |
| 467 | }))) |
| 468 | .respond_with(ResponseTemplate::new(200).set_body_json(json!({ |
| 469 | "status": "succeeded", |
| 470 | "context": { "encrypted_output": "encrypted-search-result" } |
| 471 | }))) |
| 472 | .expect(1) |
| 473 | .mount(&server) |
| 474 | .await; |
| 475 | Mock::given(method("POST")) |
| 476 | .and(path("/v1/chat/completions")) |
| 477 | .and(header("authorization", "Bearer moonshot-test-key")) |
| 478 | .and(body_string_contains("encrypted-search-result")) |
| 479 | .respond_with(ResponseTemplate::new(200).set_body_json(json!({ |
| 480 | "choices": [{ |
| 481 | "finish_reason": "stop", |
| 482 | "message": { |
| 483 | "role": "assistant", |
| 484 | "content": "See https://example.com/kimi for the current result." |
| 485 | } |
| 486 | }] |
| 487 | }))) |
| 488 | .expect(1) |
| 489 | .mount(&server) |
| 490 | .await; |
| 491 | |
| 492 | let config = Config { |
| 493 | provider: Some("moonshot".to_string()), |
| 494 | providers: Some(ProvidersConfig { |
| 495 | moonshot: ProviderConfig { |
| 496 | api_key: Some("moonshot-test-key".to_string()), |
| 497 | base_url: Some(format!("{}/v1", server.uri())), |
| 498 | model: Some("kimi-k3".to_string()), |
| 499 | ..ProviderConfig::default() |
| 500 | }, |
| 501 | ..ProvidersConfig::default() |
| 502 | }), |
| 503 | ..Config::default() |
| 504 | }; |
| 505 | let client = ProviderNativeSearchClient::new( |
| 506 | crate::client::CodewhaleClient::new(&config).expect("test Moonshot client"), |
| 507 | ) |
| 508 | .expect("Moonshot native adapter"); |
| 509 | |
| 510 | let response = search_formula(&client, &request()) |
| 511 | .await |
| 512 | .expect("K3 Formula search"); |
| 513 | |
| 514 | assert_eq!(response.citations.len(), 1); |
| 515 | assert_eq!(response.citations[0].url, "https://example.com/kimi"); |
| 516 | } |
| 517 | } |
| 518 |