lakehouse/crates/gateway/src/v1/kimi.rs

//! Kimi For Coding adapter — direct provider for `kimi-for-coding`
//! (kimi-k2.6 underneath). Used when Ollama Cloud's `kimi-k2:1t` is
//! returning sustained 5xx (broken upstream) and OpenRouter's
//! `moonshotai/kimi-k2.6` is rate-limited.
//!
//! Endpoint per `kimi.com/code/docs` and `moonshotai.github.io/kimi-cli`:
//!   base_url:  https://api.kimi.com/coding/v1
//!   model id:  kimi-for-coding
//!   auth:      Bearer sk-kimi-…
//!   protocol:  OpenAI Chat Completions compatible
//!
//! IMPORTANT: `api.kimi.com` is a separate account system from
//! `api.moonshot.ai` and `api.moonshot.cn`. Keys are NOT interchangeable.
//! This adapter is for `sk-kimi-*` keys provisioned via the Kimi
//! membership console only.
//!
//! Key sourcing priority:
//!   1. Env var `KIMI_API_KEY` (loaded from /etc/lakehouse/kimi.env via
//!      systemd EnvironmentFile=)
//!   2. /etc/lakehouse/kimi.env directly (rescue path if env not loaded)
//!
//! First hit wins. Resolved once at gateway startup, stored on
//! `V1State.kimi_key`.

use std::time::Duration;
use serde::{Deserialize, Serialize};

use super::{ChatRequest, ChatResponse, Choice, Message, UsageBlock};

const KIMI_BASE_URL: &str = "https://api.kimi.com/coding/v1";
// Default 600s — kimi-for-coding is a reasoning model; on large
// code-audit prompts (~50KB+ input + 8K output) it routinely needs
// 3-8 min to think + emit. Override with KIMI_TIMEOUT_SECS env var.
const KIMI_TIMEOUT_SECS_DEFAULT: u64 = 600;

fn kimi_timeout_secs() -> u64 {
    std::env::var("KIMI_TIMEOUT_SECS")
        .ok()
        .and_then(|s| s.trim().parse::<u64>().ok())
        .filter(|&n| n > 0)
        .unwrap_or(KIMI_TIMEOUT_SECS_DEFAULT)
}

pub fn resolve_kimi_key() -> Option<String> {
    if let Ok(k) = std::env::var("KIMI_API_KEY") {
        if !k.trim().is_empty() { return Some(k.trim().to_string()); }
    }
    if let Ok(raw) = std::fs::read_to_string("/etc/lakehouse/kimi.env") {
        for line in raw.lines() {
            if let Some(rest) = line.strip_prefix("KIMI_API_KEY=") {
                let k = rest.trim().trim_matches('"').trim_matches('\'');
                if !k.is_empty() { return Some(k.to_string()); }
            }
        }
    }
    None
}

pub async fn chat(
    key: &str,
    req: &ChatRequest,
) -> Result<ChatResponse, String> {
    // Strip the "kimi/" namespace prefix if the caller used it so the
    // upstream API sees the bare model id (e.g. "kimi-for-coding").
    let model = req.model.strip_prefix("kimi/").unwrap_or(&req.model).to_string();

    // Flatten content to a plain String. api.kimi.com is text-only on
    // the coding endpoint; the OpenAI multimodal array shape
    // ([{type:"text",text:"..."},{type:"image_url",...}]) returns 400.
    // Message::text() concats text-parts and drops non-text. Caught
    // 2026-04-27 by Kimi's self-audit (kimi.rs:137 — content as raw
    // serde_json::Value risked upstream rejection).
    let body = KimiChatBody {
        model: model.clone(),
        messages: req.messages.iter().map(|m| KimiMessage {
            role: m.role.clone(),
            content: serde_json::Value::String(m.text()),
        }).collect(),
        max_tokens: req.max_tokens.unwrap_or(800),
        temperature: req.temperature.unwrap_or(0.3),
        stream: false,
    };

    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(kimi_timeout_secs()))
        .build()
        .map_err(|e| format!("build client: {e}"))?;

    let t0 = std::time::Instant::now();
    let resp = client
        .post(format!("{}/chat/completions", KIMI_BASE_URL))
        .bearer_auth(key)
        // api.kimi.com gates this endpoint by User-Agent — only sanctioned
        // coding agents (Claude Code, Kimi CLI, Roo Code, Kilo Code) get
        // through. Generic clients receive 403 access_terminated_error.
        // J accepted the TOS risk on 2026-04-27; revisit if Moonshot
        // tightens enforcement.
        .header("User-Agent", "claude-code/1.0.0")
        .json(&body)
        .send()
        .await
        .map_err(|e| format!("api.kimi.com unreachable: {e}"))?;

    let status = resp.status();
    if !status.is_success() {
        let body = resp.text().await.unwrap_or_else(|_| "?".into());
        return Err(format!("api.kimi.com {}: {}", status, body));
    }

    let parsed: KimiChatResponse = resp.json().await
        .map_err(|e| format!("invalid kimi response: {e}"))?;

    let latency_ms = t0.elapsed().as_millis();
    let choice = parsed.choices.into_iter().next()
        .ok_or_else(|| "kimi returned no choices".to_string())?;
    let text = choice.message.content;

    let prompt_tokens = parsed.usage.as_ref().map(|u| u.prompt_tokens).unwrap_or_else(|| {
        let chars: usize = req.messages.iter().map(|m| m.text().chars().count()).sum();
        ((chars + 3) / 4) as u32
    });
    let completion_tokens = parsed.usage.as_ref().map(|u| u.completion_tokens).unwrap_or_else(|| {
        ((text.chars().count() + 3) / 4) as u32
    });

    tracing::info!(
        target: "v1.chat",
        provider = "kimi",
        model = %model,
        prompt_tokens,
        completion_tokens,
        latency_ms = latency_ms as u64,
        "kimi chat completed",
    );

    Ok(ChatResponse {
        id: format!("chatcmpl-{}", chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)),
        object: "chat.completion",
        created: chrono::Utc::now().timestamp(),
        model,
        choices: vec![Choice {
            index: 0,
            message: Message { role: "assistant".into(), content: serde_json::Value::String(text) },
            finish_reason: choice.finish_reason.unwrap_or_else(|| "stop".into()),
        }],
        usage: UsageBlock {
            prompt_tokens,
            completion_tokens,
            total_tokens: prompt_tokens + completion_tokens,
        },
    })
}

// -- Kimi wire shapes (OpenAI-compatible) --

#[derive(Serialize)]
struct KimiChatBody {
    model: String,
    messages: Vec<KimiMessage>,
    max_tokens: u32,
    temperature: f64,
    stream: bool,
}

#[derive(Serialize)]
struct KimiMessage { role: String, content: serde_json::Value }

#[derive(Deserialize)]
struct KimiChatResponse {
    choices: Vec<KimiChoice>,
    #[serde(default)]
    usage: Option<KimiUsage>,
}

#[derive(Deserialize)]
struct KimiChoice {
    message: KimiMessageResp,
    #[serde(default)]
    finish_reason: Option<String>,
}

#[derive(Deserialize)]
struct KimiMessageResp { content: String }

#[derive(Deserialize)]
struct KimiUsage { prompt_tokens: u32, completion_tokens: u32 }

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn resolve_kimi_key_does_not_panic() {
        let _ = resolve_kimi_key();
    }

    #[test]
    fn chat_body_serializes_to_openai_shape() {
        let body = KimiChatBody {
            model: "kimi-for-coding".into(),
            messages: vec![
                KimiMessage { role: "user".into(), content: "review this".into() },
            ],
            max_tokens: 800,
            temperature: 0.3,
            stream: false,
        };
        let json = serde_json::to_string(&body).unwrap();
        assert!(json.contains("\"model\":\"kimi-for-coding\""));
        assert!(json.contains("\"messages\""));
        assert!(json.contains("\"max_tokens\":800"));
        assert!(json.contains("\"stream\":false"));
    }

    #[test]
    fn model_prefix_strip() {
        let cases = [
            ("kimi/kimi-for-coding", "kimi-for-coding"),
            ("kimi-for-coding", "kimi-for-coding"),
            ("kimi/kimi-k2.6", "kimi-k2.6"),
        ];
        for (input, expected) in cases {
            let out = input.strip_prefix("kimi/").unwrap_or(input);
            assert_eq!(out, expected, "{input} should become {expected}");
        }
    }
}