nordisk/nornir

src/warehouse/generator/claude.rs

//! `claude` backend — a [`Generator`] over the shared `crate::viz::claude` client
//! (Anthropic Messages API, default model `claude-opus-4-8`, `$ANTHROPIC_API_KEY`).
//!
//! This is the funnel **executor's** live generator: `run_subtask_loop`/`drain_plan`
//! ask a `Generator` for a patch/file body, and the compiler (`AnswerJudge`) is the
//! oracle. Here that generator is Claude — a DISTINCT call from the funnel *planner*
//! (`funnel::decompose` decompose), keeping planner and executor separate.
//!
//! It reuses the existing `viz::claude` transport seam (Law #1 — no new HTTP client):
//! the real path is `HttpTransport` (feature `claude`), tests inject a `FakeTransport`.
//! Gated on `any(feature = "viz", feature = "server")` because that is where the
//! `viz::claude` client compiles (mirrors `nornir.rs::llm_decomposition`).

use std::sync::Arc;
use std::time::Instant;

use anyhow::Result;

use super::{Backend, GenAnswer, GenRequest, Generator};
use crate::viz::claude::{
    ClaudeMessage, ClaudeRequest, ClaudeTransport, DEFAULT_MODEL, api_key_present,
    default_transport,
};

/// A [`Generator`] that completes via Claude. The transport is injectable so the
/// backend is unit-testable offline with a `FakeTransport` (no network, no key).
pub struct ClaudeGenerator {
    id: String,
    model: String,
    transport: Arc<dyn ClaudeTransport>,
}

impl ClaudeGenerator {
    /// Build the backend for `model` (empty ⇒ [`DEFAULT_MODEL`]) using the runtime
    /// transport (`HttpTransport` under the `claude` feature, a disabled stub
    /// otherwise — so `complete` returns a clean error, never a panic).
    pub fn new(model: &str) -> Self {
        Self::with_transport_impl(model, default_transport())
    }

    /// Build the backend with an explicit transport (tests pass a `FakeTransport`).
    pub fn with_transport(model: &str, transport: Arc<dyn ClaudeTransport>) -> Self {
        Self::with_transport_impl(model, transport)
    }

    fn with_transport_impl(model: &str, transport: Arc<dyn ClaudeTransport>) -> Self {
        let model = if model.trim().is_empty() {
            DEFAULT_MODEL.to_string()
        } else {
            model.trim().to_string()
        };
        Self {
            id: format!("claude:{model}"),
            model,
            transport,
        }
    }
}

impl Backend for ClaudeGenerator {
    fn id(&self) -> &str {
        &self.id
    }
    /// Usable when an API key is resolvable (`$ANTHROPIC_API_KEY` or the tab token).
    /// Without the `claude` feature the transport is disabled and `complete` errors,
    /// so callers still degrade gracefully; `available()` here reports key presence.
    fn available(&self) -> bool {
        api_key_present()
    }
}

impl Generator for ClaudeGenerator {
    fn complete(&self, req: &GenRequest) -> Result<GenAnswer> {
        // The `viz::claude` ClaudeRequest has no system field (opus-4-8 rejects some
        // legacy knobs); fold the system prompt into the user turn. `temperature` is
        // intentionally ignored — the client does not send it (unsupported on 4.8).
        let prompt = match &req.system {
            Some(s) if !s.trim().is_empty() => format!("{s}\n\n{}", req.prompt),
            _ => req.prompt.clone(),
        };
        let mut creq = ClaudeRequest::new(self.model.as_str(), vec![ClaudeMessage::user(prompt)]);
        // Respect the caller's token cap (clamp into u32); a 0 cap keeps the client default.
        if req.max_tokens > 0 {
            creq.max_tokens = req.max_tokens.min(u32::MAX as usize) as u32;
        }

        let t0 = Instant::now();
        let text = self
            .transport
            .send(&creq)
            .map_err(|e| anyhow::anyhow!("claude generate: {e}"))?;
        let latency_ms = t0.elapsed().as_secs_f64() * 1000.0;

        // The transport returns assembled text only (no usage block), so token
        // counts are a coarse chars/4 estimate — enough for the bake-off row.
        Ok(GenAnswer {
            text: text.clone(),
            tokens_in: (req.prompt.len() / 4) as i64,
            tokens_out: (text.len() / 4) as i64,
            tokens_per_s: 0.0,
            latency_ms,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::viz::claude::{ClaudeError, FakeTransport};

    #[test]
    fn completes_via_injected_transport() {
        let t = Arc::new(FakeTransport::replying(
            "fn add(a: i32, b: i32) -> i32 { a + b }",
        ));
        let g = ClaudeGenerator::with_transport("claude-opus-4-8", t);
        assert_eq!(g.id(), "claude:claude-opus-4-8");
        let ans = g.complete(&GenRequest::new("implement add")).unwrap();
        assert!(
            ans.text.contains("fn add"),
            "returns the transport's completion"
        );
        assert!(
            ans.tokens_out > 0,
            "token economics estimated from the text"
        );
    }

    #[test]
    fn empty_model_defaults_and_system_folds_into_the_prompt() {
        let t = Arc::new(FakeTransport::replying("ok"));
        let g = ClaudeGenerator::with_transport("", t.clone());
        assert_eq!(
            g.id(),
            format!("claude:{DEFAULT_MODEL}"),
            "empty model ⇒ default"
        );
        let req = GenRequest::new("do the thing").with_system("You are a Rust expert.");
        g.complete(&req).unwrap();
        let seen = t.requests();
        assert_eq!(seen.len(), 1);
        let sent = &seen[0].messages[0].content;
        assert!(sent.contains("Rust expert"), "system prompt folded in");
        assert!(sent.contains("do the thing"), "user prompt present");
    }

    #[test]
    fn transport_error_surfaces_not_panics() {
        let t = Arc::new(FakeTransport::failing(ClaudeError::NoCredential));
        let g = ClaudeGenerator::with_transport("m", t);
        assert!(
            g.complete(&GenRequest::new("x")).is_err(),
            "no key ⇒ Err, never a panic"
        );
    }
}