lsp-llm 0.2.1

//! Anthropic-backed advisor implementation.
//!
//! Compiled only with `--features llm`. The deterministic stack
//! never touches anything in this module — it's the sole carrier
//! of `reqwest`, prompt strings, and any other "AI" dependency in
//! the workspace. `cargo build --release` (default features) keeps
//! it out of the binary.

use std::fmt::Write as _;
use std::sync::Arc;
use std::time::Duration;

use lsp_types::{Position, Range, Url};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tracing::{debug, warn};

use crate::{
    AskRequest, AskResponse, Citation, ENV_ANTHROPIC_KEY, ENV_BACKEND, ENV_ENABLED, ENV_MODEL,
    is_runtime_enabled,
};

/// Default model when `AXON_LLM_MODEL` isn't set. Picks the cheapest
/// Claude tier — sufficient for the short, grounded responses
/// `askAdvisor` produces. Override by exporting `AXON_LLM_MODEL`.
const DEFAULT_MODEL: &str = "claude-haiku-4-5-20251001";

/// Network deadline. The advisor is interactive — anything past 30 s
/// is dead time the user could spend reading docs instead.
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);

/// Errors the advisor can surface to the Backend. Each is mapped to
/// a JSON-RPC error code at the call site.
#[derive(Debug)]
pub enum AdvisorError {
    /// Compile-time on, runtime off. Caller should surface
    /// `MethodNotFound` with the canonical message.
    Disabled,
    /// `AXON_LLM_BACKEND` was set to something other than `claude`,
    /// or `ANTHROPIC_API_KEY` is missing.
    Misconfigured(String),
    /// Network or transport failure. Includes the upstream error
    /// for tracing.
    Transport(String),
    /// HTTP response rejected by the upstream API (non-2xx).
    Upstream { status: u16, body: String },
    /// Response body didn't match the expected Anthropic shape.
    Parse(String),
}

impl std::fmt::Display for AdvisorError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Disabled => write!(f, "advisor disabled (set {ENV_ENABLED}=1)"),
            Self::Misconfigured(msg) => write!(f, "advisor misconfigured: {msg}"),
            Self::Transport(msg) => write!(f, "advisor transport error: {msg}"),
            Self::Upstream { status, body } => {
                write!(f, "advisor upstream error {status}: {body}")
            }
            Self::Parse(msg) => write!(f, "advisor response parse error: {msg}"),
        }
    }
}

impl std::error::Error for AdvisorError {}

/// Anthropic-backed advisor. Holds a reusable HTTP client and the
/// resolved configuration; each `ask` call is a single Messages API
/// round-trip.
pub struct Advisor {
    client: Client,
    api_key: String,
    model: String,
}

impl Advisor {
    /// Build an [`Advisor`] from environment variables. Returns
    /// `None` when the runtime flag is off; returns `Err` when the
    /// flag is on but configuration is incomplete (so the Backend
    /// can log a clear diagnostic instead of silently misbehaving).
    pub fn from_env() -> Result<Self, AdvisorError> {
        if !is_runtime_enabled() {
            return Err(AdvisorError::Disabled);
        }
        let backend = std::env::var(ENV_BACKEND)
            .ok()
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| "claude".into());
        if backend != "claude" {
            return Err(AdvisorError::Misconfigured(format!(
                "unsupported backend {backend:?} (only `claude` is wired today)"
            )));
        }
        let api_key = std::env::var(ENV_ANTHROPIC_KEY).map_err(|_| {
            AdvisorError::Misconfigured(format!(
                "{ENV_ANTHROPIC_KEY} is not set; required for backend `claude`"
            ))
        })?;
        if api_key.is_empty() {
            return Err(AdvisorError::Misconfigured(format!(
                "{ENV_ANTHROPIC_KEY} is empty"
            )));
        }
        let model = std::env::var(ENV_MODEL)
            .ok()
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| DEFAULT_MODEL.into());

        let client = Client::builder()
            .timeout(REQUEST_TIMEOUT)
            .user_agent(concat!("axon-lsp/", env!("CARGO_PKG_VERSION")))
            .build()
            .map_err(|e| AdvisorError::Transport(e.to_string()))?;

        Ok(Self {
            client,
            api_key,
            model,
        })
    }

    /// Run the advisor: build a prompt grounded in the embedded
    /// docs corpus + the optional code snippet, post to Anthropic,
    /// extract the answer, and synthesise citations for every doc
    /// entry name mentioned in the answer.
    pub async fn ask(&self, request: AskRequest) -> Result<AskResponse, AdvisorError> {
        let docs = relevant_docs(&request);
        let prompt = build_prompt(&request, &docs);
        debug!(
            chars = prompt.chars().count(),
            docs = docs.len(),
            "advisor prompt built"
        );

        let payload = MessagesRequest {
            model: &self.model,
            max_tokens: 1024,
            messages: vec![Message {
                role: "user",
                content: &prompt,
            }],
        };

        let response = self
            .client
            .post("https://api.anthropic.com/v1/messages")
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", "2023-06-01")
            .json(&payload)
            .send()
            .await
            .map_err(|e| AdvisorError::Transport(e.to_string()))?;

        let status = response.status();
        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            warn!(status = %status, "advisor upstream error");
            return Err(AdvisorError::Upstream {
                status: status.as_u16(),
                body,
            });
        }

        let parsed: MessagesResponse = response
            .json()
            .await
            .map_err(|e| AdvisorError::Parse(e.to_string()))?;
        let answer = parsed
            .content
            .into_iter()
            .filter(|b| b.kind == "text")
            .map(|b| b.text)
            .collect::<Vec<_>>()
            .join("\n");

        let mut citations = synthesise_citations(&answer, &docs);
        if let Some(ctx) = &request.context {
            citations.push(Citation {
                uri: ctx.uri.clone(),
                range: ctx.range,
            });
        }

        Ok(AskResponse { answer, citations })
    }
}

/// Convenience wrapper used by `axon-lsp::server::Backend` at boot.
/// Wraps the raw constructor so the Backend can store
/// `Option<Arc<Advisor>>` without unwrapping `Result` plumbing.
#[must_use]
pub fn try_advisor_from_env() -> Option<Arc<Advisor>> {
    match Advisor::from_env() {
        Ok(a) => Some(Arc::new(a)),
        Err(AdvisorError::Disabled) => None,
        Err(other) => {
            warn!(error = %other, "advisor disabled — misconfigured");
            None
        }
    }
}

// ── Prompt construction ────────────────────────────────────────────

fn build_prompt(request: &AskRequest, docs: &[&'static lsp_docs::DocEntry]) -> String {
    let mut out = String::with_capacity(2048);
    out.push_str(
        "You are an assistant for the Axon programming language. \
         Answer the user's question concisely in Markdown. \
         Cite documentation entries by their `name` when relevant. \
         If you don't know the answer, say so — don't invent symbols.\n\n",
    );

    if !docs.is_empty() {
        out.push_str("## Grounding documentation\n\n");
        for d in docs {
            let kind = match d.kind {
                lsp_docs::DocKind::Type => "type",
                lsp_docs::DocKind::Syntax => "syntax",
                lsp_docs::DocKind::Handler => "handler",
            };
            // `write!` into a `String` is infallible — the helper
            // can't err. The trait method just requires the
            // unwrap.
            let _ = write!(out, "### {kind} `{}`\n\n", d.name);
            out.push_str(d.body);
            out.push_str("\n\n");
        }
    }

    if let Some(ctx) = &request.context
        && !ctx.text.is_empty()
    {
        out.push_str("## User's selected code\n\n```axon\n");
        out.push_str(&ctx.text);
        if !ctx.text.ends_with('\n') {
            out.push('\n');
        }
        out.push_str("```\n\n");
    }

    out.push_str("## Question\n\n");
    out.push_str(&request.question);
    out.push('\n');
    out
}

/// Pick the docs to send as grounding. The corpus is small (~18
/// entries, ~5 KB total), so for v0.1.0 we send the full set.
/// Future expansions can switch to BM25 / embedding-based
/// retrieval if the corpus grows past a few hundred entries.
fn relevant_docs(_request: &AskRequest) -> Vec<&'static lsp_docs::DocEntry> {
    lsp_docs::iter_entries().collect()
}

/// Find the doc-entry names that appear in `answer` and synthesise
/// citation entries for each. URI scheme is `axon-lsp://docs/...`
/// — editors may not navigate it natively, but the structure is
/// there for future tooling.
fn synthesise_citations(answer: &str, docs: &[&'static lsp_docs::DocEntry]) -> Vec<Citation> {
    let mut out = Vec::new();
    for d in docs {
        // Word-bounded contains: avoid matching `String` inside
        // `Stringify`, etc. Cheap pass — corpus is tiny.
        if mentions_word(answer, d.name) {
            let kind = match d.kind {
                lsp_docs::DocKind::Type => "types",
                lsp_docs::DocKind::Syntax => "syntax",
                lsp_docs::DocKind::Handler => "handlers",
            };
            // SAFETY: the URI is built from `&'static str`s under
            // our control — never fails.
            if let Ok(uri) = Url::parse(&format!("axon-lsp://docs/{kind}/{name}.md", name = d.name))
            {
                out.push(Citation {
                    uri,
                    range: Range {
                        start: Position::new(0, 0),
                        end: Position::new(0, 0),
                    },
                });
            }
        }
    }
    out
}

fn mentions_word(haystack: &str, needle: &str) -> bool {
    if needle.is_empty() {
        return false;
    }
    let mut i = 0;
    while let Some(found) = haystack[i..].find(needle) {
        let abs = i + found;
        let before = haystack[..abs].chars().last();
        let after = haystack[abs + needle.len()..].chars().next();
        let is_word = |c: char| c.is_ascii_alphanumeric() || c == '_';
        let left_ok = before.is_none_or(|c| !is_word(c));
        let right_ok = after.is_none_or(|c| !is_word(c));
        if left_ok && right_ok {
            return true;
        }
        i = abs + needle.len();
    }
    false
}

// ── Anthropic Messages API wire types ──────────────────────────────

#[derive(Serialize)]
struct MessagesRequest<'a> {
    model: &'a str,
    max_tokens: u32,
    messages: Vec<Message<'a>>,
}

#[derive(Serialize)]
struct Message<'a> {
    role: &'a str,
    content: &'a str,
}

#[derive(Deserialize)]
struct MessagesResponse {
    content: Vec<ContentBlock>,
}

#[derive(Deserialize)]
struct ContentBlock {
    #[serde(rename = "type")]
    kind: String,
    #[serde(default)]
    text: String,
}

#[cfg(test)]
mod tests {
    use super::{build_prompt, mentions_word, synthesise_citations};
    use crate::AskRequest;

    #[test]
    fn build_prompt_includes_grounding_and_question() {
        let req = AskRequest {
            question: "What is a Trusted<T>?".into(),
            context: None,
        };
        let docs: Vec<&lsp_docs::DocEntry> = lsp_docs::iter_entries().collect();
        let prompt = build_prompt(&req, &docs);
        assert!(prompt.contains("## Grounding documentation"));
        assert!(prompt.contains("type `Trusted`"));
        assert!(prompt.contains("## Question"));
        assert!(prompt.contains("What is a Trusted<T>?"));
    }

    #[test]
    fn mentions_word_is_word_bounded() {
        assert!(mentions_word("the String type", "String"));
        assert!(mentions_word("`String`.", "String"));
        assert!(!mentions_word("Stringify is unrelated", "String"));
        assert!(!mentions_word("xStringy", "String"));
    }

    #[test]
    fn synthesise_citations_skips_unmentioned() {
        let docs: Vec<&lsp_docs::DocEntry> = lsp_docs::iter_entries().collect();
        let answer = "Use a `Trusted<T>` to carry compliance evidence.";
        let cites = synthesise_citations(answer, &docs);
        assert!(cites.iter().any(|c| c.uri.as_str().contains("Trusted")));
        assert!(!cites.iter().any(|c| c.uri.as_str().contains("String")));
    }
}