dci-tool 0.1.0

//! `dci-mcp` — serve a Direct Corpus Interaction agent as a stateful MCP tool.
//!
//! Exposes one corpus as the `dci_investigate` Model Context Protocol tool over
//! stdio (JSON-RPC on stdin/stdout). Any MCP client can then ask questions
//! about the corpus and pass a `session_id` to continue an investigation.
//! Bring your own model via the relevant provider environment variable.
//!
//! Logs go to stderr so they never corrupt the JSON-RPC stream on stdout.

use std::path::PathBuf;
use std::process::ExitCode;

use anyhow::{Context, Result};
use clap::{Parser, ValueEnum};
use dci_tool::{CorpusRoot, DciAgent, DciMcpService, Limits, SessionConfig};
use rig_core::client::{CompletionClient, ProviderClient};
use rig_core::completion::CompletionModel;
use rig_core::providers::{anthropic, ollama, openai};

/// Supported model providers (selected at runtime; bring your own key).
#[derive(Debug, Clone, Copy, ValueEnum)]
enum Provider {
    /// OpenAI (`OPENAI_API_KEY`).
    Openai,
    /// Anthropic (`ANTHROPIC_API_KEY`).
    Anthropic,
    /// Local Ollama (`OLLAMA_*`, no key required).
    Ollama,
}

/// Serve a corpus as a stateful MCP tool over stdio.
#[derive(Debug, Parser)]
#[command(name = "dci-mcp", version, about)]
struct Cli {
    /// Path to the corpus directory to expose.
    #[arg(short, long, default_value = ".")]
    corpus: PathBuf,

    /// Model provider to use.
    #[arg(short, long, value_enum, default_value_t = Provider::Openai)]
    provider: Provider,

    /// Model name (defaults to a sensible model per provider).
    #[arg(short, long)]
    model: Option<String>,

    /// Maximum tool-calling turns per investigation.
    #[arg(long, default_value_t = dci_tool::DEFAULT_MAX_TURNS)]
    max_turns: usize,

    /// Do not honor `.gitignore` rules while searching (useful for log corpora).
    #[arg(long)]
    no_gitignore: bool,

    /// Maximum number of concurrent sessions retained (0 disables the cap).
    /// Least-recently-active sessions are evicted past this.
    #[arg(long, default_value_t = SessionConfig::default().max_sessions)]
    max_sessions: usize,

    /// Maximum turns kept per session (0 disables the cap). Oldest turns are
    /// dropped past this.
    #[arg(long, default_value_t = SessionConfig::default().max_turns_per_session)]
    max_turns_per_session: usize,

    /// Idle seconds before a session is pruned (0 disables expiry).
    #[arg(long, default_value_t = SessionConfig::default().ttl.as_secs())]
    session_ttl_secs: u64,
}

#[tokio::main]
async fn main() -> ExitCode {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .with_writer(std::io::stderr)
        .init();

    match run(Cli::parse()).await {
        Ok(()) => ExitCode::SUCCESS,
        Err(err) => {
            eprintln!("error: {err:#}");
            ExitCode::FAILURE
        }
    }
}

async fn run(cli: Cli) -> Result<()> {
    let limits = Limits {
        respect_gitignore: !cli.no_gitignore,
        ..Limits::default()
    };
    let corpus = CorpusRoot::with_limits(&cli.corpus, limits)
        .with_context(|| format!("opening corpus at {}", cli.corpus.display()))?;

    let session_config = SessionConfig {
        max_sessions: cli.max_sessions,
        max_turns_per_session: cli.max_turns_per_session,
        ttl: std::time::Duration::from_secs(cli.session_ttl_secs),
    };

    let service = match cli.provider {
        Provider::Openai => {
            let client = openai::Client::from_env().context("initializing OpenAI client")?;
            let model_id = cli.model.as_deref().unwrap_or(openai::GPT_4O).to_string();
            let model = client.completion_model(&model_id);
            build_service(model, corpus, cli.max_turns, None, &model_id, session_config)
        }
        Provider::Anthropic => {
            let client = anthropic::Client::from_env().context("initializing Anthropic client")?;
            let model_id = cli
                .model
                .as_deref()
                .unwrap_or(anthropic::completion::CLAUDE_SONNET_4_6)
                .to_string();
            let model = client.completion_model(&model_id);
            // Anthropic requires an explicit output token budget.
            build_service(model, corpus, cli.max_turns, Some(4096), &model_id, session_config)
        }
        Provider::Ollama => {
            let client = ollama::Client::from_env().context("initializing Ollama client")?;
            let model_id = cli.model.as_deref().unwrap_or("llama3.1").to_string();
            let model = client.completion_model(&model_id);
            build_service(model, corpus, cli.max_turns, None, &model_id, session_config)
        }
    };

    tracing::info!(
        tool = service.tool_name(),
        "serving DCI corpus over MCP stdio"
    );
    service
        .serve_stdio()
        .await
        .map_err(|e| anyhow::anyhow!("MCP server error: {e}"))
}

fn build_service<M: CompletionModel + 'static>(
    model: M,
    corpus: CorpusRoot,
    max_turns: usize,
    max_tokens: Option<u64>,
    model_label: &str,
    session_config: SessionConfig,
) -> DciMcpService {
    let mut builder = DciAgent::builder(model, corpus)
        .max_turns(max_turns)
        .model_label(model_label);
    if let Some(max_tokens) = max_tokens {
        builder = builder.max_tokens(max_tokens);
    }
    DciMcpService::new_with_config(builder.build(), session_config)
}