isokit/crates/isokit-cli/src/agent.rs

//! The Studio's agent window talks to a language model through these routes,
//! so the endpoint key stays on this machine and the browser needs no CORS
//! permission from the endpoint (any OpenAI-compatible `/v1` API: vLLM,
//! llama.cpp, Ollama, LM Studio, a hosted one):
//!
//! - `GET  /api/agent/config`  endpoint, model, whether a key is set (never the key)
//! - `POST /api/agent/config`  save `{ url, model, think, key }` (key "" forgets it)
//! - `GET  /api/agent/models`  the endpoint's `/models`
//! - `POST /api/agent/chat`    the endpoint's `/chat/completions`, streamed through
//!
//! Settings come from `ISOKIT_LLM_URL`, `ISOKIT_LLM_KEY` and `ISOKIT_LLM_MODEL`,
//! else from `agent.json` in the user's config folder (`%APPDATA%\isokit`,
//! `~/.config/isokit`; `ISOKIT_CONFIG_DIR` moves it), which the window writes.

use serde_json::{json, Map, Value};
use std::io::{Read, Write};
use std::path::PathBuf;
use std::time::Duration;
use tiny_http::Request;

/// B4L's inference server (vLLM, Qwen); any OpenAI-compatible base URL works.
pub const DEFAULT_URL: &str = "https://inference.p0.contact/v1";

#[derive(Debug, Clone, PartialEq)]
pub struct Settings {
    pub url: String,
    pub key: String,
    pub model: String,
    pub think: bool,
    /// where url / key came from: "env", "file" or "default" / ""
    pub url_from: &'static str,
    pub key_from: &'static str,
}

pub fn config_file() -> Option<PathBuf> {
    if let Some(d) = var("ISOKIT_CONFIG_DIR") {
        return Some(PathBuf::from(d).join("agent.json"));
    }
    let base = if cfg!(windows) {
        var("APPDATA").map(PathBuf::from)
    } else {
        var("XDG_CONFIG_HOME")
            .map(PathBuf::from)
            .or_else(|| var("HOME").map(|h| PathBuf::from(h).join(".config")))
    };
    base.map(|b| b.join("isokit").join("agent.json"))
}

fn var(k: &str) -> Option<String> {
    std::env::var(k)
        .ok()
        .map(|v| v.trim().to_string())
        .filter(|v| !v.is_empty())
}

fn read_file() -> Map<String, Value> {
    config_file()
        .and_then(|p| std::fs::read_to_string(p).ok())
        .and_then(|s| serde_json::from_str::<Map<String, Value>>(&s).ok())
        .unwrap_or_default()
}

pub fn load() -> Settings {
    let f = read_file();
    let fs = |k: &str| {
        f.get(k)
            .and_then(|v| v.as_str())
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty())
    };
    let (url, url_from) = match (var("ISOKIT_LLM_URL"), fs("url")) {
        (Some(u), _) => (u, "env"),
        (None, Some(u)) => (u, "file"),
        _ => (DEFAULT_URL.to_string(), "default"),
    };
    let (key, key_from) = match (var("ISOKIT_LLM_KEY"), fs("key")) {
        (Some(k), _) => (k, "env"),
        (None, Some(k)) => (k, "file"),
        _ => (String::new(), ""),
    };
    Settings {
        url,
        key,
        model: var("ISOKIT_LLM_MODEL")
            .or_else(|| fs("model"))
            .unwrap_or_default(),
        think: f.get("think").and_then(|v| v.as_bool()).unwrap_or(true),
        url_from,
        key_from,
    }
}

/// What the page may see: everything but the key.
pub fn public(s: &Settings) -> Value {
    json!({
        "url": s.url, "url_from": s.url_from, "model": s.model, "think": s.think,
        "has_key": !s.key.is_empty(), "key_from": s.key_from,
        "file": config_file().map(|p| p.display().to_string()),
    })
}

/// Merge the window's settings into the config file. A missing field stays
/// as it is; `key: ""` forgets the key.
pub fn save(body: &Value) -> Result<Value, String> {
    let path = config_file().ok_or("no config folder (set ISOKIT_CONFIG_DIR)")?;
    let mut f = read_file();
    for k in ["url", "model", "key"] {
        if let Some(v) = body.get(k).and_then(|v| v.as_str()) {
            let v = v.trim();
            if v.is_empty() {
                f.remove(k);
            } else {
                f.insert(k.into(), json!(v));
            }
        }
    }
    if let Some(b) = body.get("think").and_then(|v| v.as_bool()) {
        f.insert("think".into(), json!(b));
    }
    if let Some(dir) = path.parent() {
        std::fs::create_dir_all(dir).map_err(|e| format!("{}: {e}", dir.display()))?;
    }
    let text = serde_json::to_string_pretty(&Value::Object(f)).unwrap_or_default();
    std::fs::write(&path, text + "\n").map_err(|e| format!("{}: {e}", path.display()))?;
    Ok(public(&load()))
}

/// Forward a request to `<url><path>` on its own thread (an answer can take
/// minutes; the Studio keeps serving meanwhile) and stream the answer back.
pub fn forward(req: Request, path: &'static str, body: Option<String>) {
    std::thread::spawn(move || {
        let s = load();
        let url = format!("{}{}", s.url.trim_end_matches('/'), path);
        let agent: ureq::Agent = ureq::Agent::config_builder()
            .http_status_as_error(false)
            .timeout_connect(Some(Duration::from_secs(15)))
            .timeout_recv_response(Some(Duration::from_secs(600)))
            .user_agent(concat!("isokit-studio/", env!("CARGO_PKG_VERSION")))
            .build()
            .into();
        let auth = format!("Bearer {}", s.key);
        let res = match body {
            Some(b) => {
                let b = with_model(&b, &s.model);
                let mut r = agent.post(&url).header("Content-Type", "application/json");
                if !s.key.is_empty() {
                    r = r.header("Authorization", &auth);
                }
                r.send(b.as_bytes())
            }
            None => {
                let mut r = agent.get(&url);
                if !s.key.is_empty() {
                    r = r.header("Authorization", &auth);
                }
                r.call()
            }
        };
        match res {
            Ok(resp) => stream_back(req, resp),
            Err(e) => {
                let msg = format!("{url}: {e}");
                let body = json!({ "error": { "message": msg } }).to_string();
                let r = tiny_http::Response::from_string(body)
                    .with_status_code(502)
                    .with_header(
                        tiny_http::Header::from_bytes("Content-Type", "application/json")
                            .expect("valid header"),
                    );
                let _ = req.respond(r);
            }
        }
    });
}

/// The configured model when the page left it open.
fn with_model(body: &str, model: &str) -> String {
    match serde_json::from_str::<Value>(body) {
        Ok(Value::Object(mut m)) => {
            let open = m
                .get("model")
                .and_then(|v| v.as_str())
                .is_none_or(|s| s.is_empty());
            if open && !model.is_empty() {
                m.insert("model".into(), json!(model));
            }
            Value::Object(m).to_string()
        }
        _ => body.to_string(),
    }
}

/// Copy the endpoint's answer to the page as it arrives (chunked, flushed per
/// read: tiny_http's own chunked writer would hold back 8 kB at a time).
fn stream_back(req: Request, resp: ureq::http::Response<ureq::Body>) {
    let status = resp.status();
    let ctype = resp
        .headers()
        .get("content-type")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("application/json")
        .to_string();
    let sse = ctype.starts_with("text/event-stream");
    let mut body = resp.into_body().into_reader();
    let mut w = req.into_writer();
    let head = format!(
        "HTTP/1.1 {} {}\r\nContent-Type: {}\r\nCache-Control: no-store\r\nX-Accel-Buffering: no\r\nTransfer-Encoding: chunked\r\n\r\n",
        status.as_u16(),
        status.canonical_reason().unwrap_or(""),
        ctype
    );
    if w.write_all(head.as_bytes())
        .and_then(|_| w.flush())
        .is_err()
    {
        return;
    }
    let mut buf = vec![0u8; 16 << 10];
    loop {
        let chunk: Vec<u8> = match body.read(&mut buf) {
            Ok(0) => break,
            Ok(n) => buf[..n].to_vec(),
            // the endpoint broke off: tell a streaming page why
            Err(e) if sse => {
                let ev = json!({ "error": { "message": format!("the endpoint broke off: {e}") } });
                format!("data: {ev}\n\n").into_bytes()
            }
            Err(_) => break,
        };
        let last = sse && chunk.starts_with(b"data: {\"error\"");
        let ok = write!(w, "{:x}\r\n", chunk.len())
            .and_then(|_| w.write_all(&chunk))
            .and_then(|_| w.write_all(b"\r\n"))
            .and_then(|_| w.flush());
        if ok.is_err() {
            return; // the page stopped reading (Stop): dropping the reader ends the request upstream
        }
        if last {
            break;
        }
    }
    let _ = w.write_all(b"0\r\n\r\n").and_then(|_| w.flush());
}

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

    #[test]
    fn model_fills_in_only_when_open() {
        let b = with_model(r#"{"model":"","messages":[]}"#, "qwen");
        assert_eq!(serde_json::from_str::<Value>(&b).unwrap()["model"], "qwen");
        let b = with_model(r#"{"model":"other"}"#, "qwen");
        assert_eq!(serde_json::from_str::<Value>(&b).unwrap()["model"], "other");
        let b = with_model(r#"{"messages":[]}"#, "");
        assert!(serde_json::from_str::<Value>(&b)
            .unwrap()
            .get("model")
            .is_none());
    }
}