storkit/server/src/llm/chat.rs

use crate::llm::prompts::SYSTEM_PROMPT;
use crate::llm::types::{Message, Role, ToolCall, ToolDefinition, ToolFunctionDefinition};
use crate::state::SessionState;
use crate::store::StoreOps;
use serde::Deserialize;
use serde_json::json;

const MAX_TURNS: usize = 30;
const KEY_ANTHROPIC_API_KEY: &str = "anthropic_api_key";

#[derive(Deserialize, Clone)]
pub struct ProviderConfig {
    pub provider: String,
    pub model: String,
    pub base_url: Option<String>,
    pub enable_tools: Option<bool>,
}

fn get_anthropic_api_key_exists_impl(store: &dyn StoreOps) -> bool {
    match store.get(KEY_ANTHROPIC_API_KEY) {
        Some(value) => value.as_str().map(|k| !k.is_empty()).unwrap_or(false),
        None => false,
    }
}

fn set_anthropic_api_key_impl(store: &dyn StoreOps, api_key: &str) -> Result<(), String> {
    store.set(KEY_ANTHROPIC_API_KEY, json!(api_key));
    store.save()?;

    match store.get(KEY_ANTHROPIC_API_KEY) {
        Some(value) => {
            if let Some(retrieved) = value.as_str() {
                if retrieved != api_key {
                    return Err("Retrieved key does not match saved key".to_string());
                }
            } else {
                return Err("Stored value is not a string".to_string());
            }
        }
        None => {
            return Err("API key was saved but cannot be retrieved".to_string());
        }
    }

    Ok(())
}

fn get_anthropic_api_key_impl(store: &dyn StoreOps) -> Result<String, String> {
    match store.get(KEY_ANTHROPIC_API_KEY) {
        Some(value) => {
            if let Some(key) = value.as_str() {
                if key.is_empty() {
                    Err("Anthropic API key is empty. Please set your API key.".to_string())
                } else {
                    Ok(key.to_string())
                }
            } else {
                Err("Stored API key is not a string".to_string())
            }
        }
        None => Err("Anthropic API key not found. Please set your API key.".to_string()),
    }
}

fn parse_tool_arguments(args_str: &str) -> Result<serde_json::Value, String> {
    serde_json::from_str(args_str).map_err(|e| format!("Error parsing arguments: {e}"))
}

pub fn get_tool_definitions() -> Vec<ToolDefinition> {
    vec![
        ToolDefinition {
            kind: "function".to_string(),
            function: ToolFunctionDefinition {
                name: "read_file".to_string(),
                description: "Reads the complete content of a file from the project. Use this to understand existing code before making changes.".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "path": { "type": "string", "description": "Relative path to the file from project root" }
                    },
                    "required": ["path"]
                }),
            },
        },
        ToolDefinition {
            kind: "function".to_string(),
            function: ToolFunctionDefinition {
                name: "write_file".to_string(),
                description: "Creates or completely overwrites a file with new content. YOU MUST USE THIS to implement code changes - do not suggest code to the user. The content parameter must contain the COMPLETE file including all imports, functions, and unchanged code.".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "path": { "type": "string", "description": "Relative path to the file from project root" },
                        "content": { "type": "string", "description": "The complete file content to write (not a diff or partial code)" }
                    },
                    "required": ["path", "content"]
                }),
            },
        },
        ToolDefinition {
            kind: "function".to_string(),
            function: ToolFunctionDefinition {
                name: "list_directory".to_string(),
                description: "Lists all files and directories at a given path. Use this to explore the project structure.".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "path": { "type": "string", "description": "Relative path to list (use '.' for project root)" }
                    },
                    "required": ["path"]
                }),
            },
        },
        ToolDefinition {
            kind: "function".to_string(),
            function: ToolFunctionDefinition {
                name: "search_files".to_string(),
                description: "Searches for text patterns across all files in the project. Use this to find functions, variables, or code patterns when you don't know which file they're in."
                    .to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "query": { "type": "string", "description": "The text pattern to search for across all files" }
                    },
                    "required": ["query"]
                }),
            },
        },
        ToolDefinition {
            kind: "function".to_string(),
            function: ToolFunctionDefinition {
                name: "exec_shell".to_string(),
                description: "Executes a shell command in the project root directory. Use this to run tests, build commands, git operations, or any command-line tool. Examples: cargo check, npm test, git status.".to_string(),
                parameters: json!({
                    "type": "object",
                    "properties": {
                        "command": {
                            "type": "string",
                            "description": "The command binary to execute (e.g., 'git', 'cargo', 'npm', 'ls')"
                        },
                        "args": {
                            "type": "array",
                            "items": { "type": "string" },
                            "description": "Array of arguments to pass to the command (e.g., ['status'] for git status)"
                        }
                    },
                    "required": ["command", "args"]
                }),
            },
        },
    ]
}

pub async fn get_ollama_models(base_url: Option<String>) -> Result<Vec<String>, String> {
    use crate::llm::providers::ollama::OllamaProvider;
    let url = base_url.unwrap_or_else(|| "http://localhost:11434".to_string());
    OllamaProvider::get_models(&url).await
}

pub fn get_anthropic_api_key_exists(store: &dyn StoreOps) -> Result<bool, String> {
    Ok(get_anthropic_api_key_exists_impl(store))
}

pub fn set_anthropic_api_key(store: &dyn StoreOps, api_key: String) -> Result<(), String> {
    set_anthropic_api_key_impl(store, &api_key)
}

pub async fn chat<F, U>(
    messages: Vec<Message>,
    config: ProviderConfig,
    state: &SessionState,
    store: &dyn StoreOps,
    mut on_update: F,
    mut on_token: U,
) -> Result<Vec<Message>, String>
where
    F: FnMut(&[Message]) + Send,
    U: FnMut(&str) + Send,
{
    use crate::llm::providers::anthropic::AnthropicProvider;
    use crate::llm::providers::ollama::OllamaProvider;

    let _ = state.cancel_tx.send(false);
    let mut cancel_rx = state.cancel_rx.clone();
    cancel_rx.borrow_and_update();

    let base_url = config
        .base_url
        .clone()
        .unwrap_or_else(|| "http://localhost:11434".to_string());

    eprintln!("[chat] provider={} model={}", config.provider, config.model);
    let is_claude_code = config.provider == "claude-code";
    let is_claude = !is_claude_code && config.model.starts_with("claude-");

    if !is_claude_code && !is_claude && config.provider.as_str() != "ollama" {
        return Err(format!("Unsupported provider: {}", config.provider));
    }

    // Claude Code provider: bypasses our tool loop entirely.
    // Claude Code has its own agent loop, tools, and context management.
    // We just pipe the user message in and stream raw output back.
    if is_claude_code {
        use crate::llm::providers::claude_code::ClaudeCodeProvider;

        let user_message = messages
            .iter()
            .rev()
            .find(|m| m.role == Role::User)
            .map(|m| m.content.clone())
            .ok_or_else(|| "No user message found".to_string())?;

        let project_root = state
            .get_project_root()
            .unwrap_or_else(|_| std::path::PathBuf::from("."));

        let provider = ClaudeCodeProvider::new();
        let response = provider
            .chat_stream(
                &user_message,
                &project_root.to_string_lossy(),
                &mut cancel_rx,
                |token| on_token(token),
            )
            .await
            .map_err(|e| format!("Claude Code Error: {e}"))?;

        let assistant_msg = Message {
            role: Role::Assistant,
            content: response.content.unwrap_or_default(),
            tool_calls: None,
            tool_call_id: None,
        };

        let mut result = messages.clone();
        result.push(assistant_msg);
        on_update(&result);
        return Ok(result);
    }

    let tool_defs = get_tool_definitions();
    let tools = if config.enable_tools.unwrap_or(true) {
        tool_defs.as_slice()
    } else {
        &[]
    };

    let mut current_history = messages.clone();

    current_history.insert(
        0,
        Message {
            role: Role::System,
            content: SYSTEM_PROMPT.to_string(),
            tool_calls: None,
            tool_call_id: None,
        },
    );

    current_history.insert(
        1,
        Message {
            role: Role::System,
            content: "REMINDER: Distinguish between showing examples (use code blocks in chat) vs implementing changes (use write_file tool). Keywords like 'show me', 'example', 'how does' = chat response. Keywords like 'create', 'add', 'implement', 'fix' = use tools."
                .to_string(),
            tool_calls: None,
            tool_call_id: None,
        },
    );

    let mut new_messages: Vec<Message> = Vec::new();
    let mut turn_count = 0;

    loop {
        if *cancel_rx.borrow() {
            return Err("Chat cancelled by user".to_string());
        }

        if turn_count >= MAX_TURNS {
            return Err("Max conversation turns reached.".to_string());
        }
        turn_count += 1;

        let response = if is_claude {
            let api_key = get_anthropic_api_key_impl(store)?;
            let anthropic_provider = AnthropicProvider::new(api_key);
            anthropic_provider
                .chat_stream(
                    &config.model,
                    &current_history,
                    tools,
                    &mut cancel_rx,
                    |token| on_token(token),
                )
                .await
                .map_err(|e| format!("Anthropic Error: {e}"))?
        } else {
            let ollama_provider = OllamaProvider::new(base_url.clone());
            ollama_provider
                .chat_stream(
                    &config.model,
                    &current_history,
                    tools,
                    &mut cancel_rx,
                    |token| on_token(token),
                )
                .await
                .map_err(|e| format!("Ollama Error: {e}"))?
        };

        if let Some(tool_calls) = response.tool_calls {
            let assistant_msg = Message {
                role: Role::Assistant,
                content: response.content.unwrap_or_default(),
                tool_calls: Some(tool_calls.clone()),
                tool_call_id: None,
            };

            current_history.push(assistant_msg.clone());
            new_messages.push(assistant_msg);
            on_update(&current_history[2..]);

            for call in tool_calls {
                if *cancel_rx.borrow() {
                    return Err("Chat cancelled before tool execution".to_string());
                }

                let output = execute_tool(&call, state).await;

                let tool_msg = Message {
                    role: Role::Tool,
                    content: output,
                    tool_calls: None,
                    tool_call_id: call.id,
                };

                current_history.push(tool_msg.clone());
                new_messages.push(tool_msg);
                on_update(&current_history[2..]);
            }
        } else {
            let assistant_msg = Message {
                role: Role::Assistant,
                content: response.content.unwrap_or_default(),
                tool_calls: None,
                tool_call_id: None,
            };

            new_messages.push(assistant_msg.clone());
            current_history.push(assistant_msg);
            on_update(&current_history[2..]);
            break;
        }
    }

    Ok(new_messages)
}

async fn execute_tool(call: &ToolCall, state: &SessionState) -> String {
    use crate::io::{fs, search, shell};

    let name = call.function.name.as_str();
    let args: serde_json::Value = match parse_tool_arguments(&call.function.arguments) {
        Ok(v) => v,
        Err(e) => return e,
    };

    match name {
        "read_file" => {
            let path = args["path"].as_str().unwrap_or("").to_string();
            match fs::read_file(path, state).await {
                Ok(content) => content,
                Err(e) => format!("Error: {e}"),
            }
        }
        "write_file" => {
            let path = args["path"].as_str().unwrap_or("").to_string();
            let content = args["content"].as_str().unwrap_or("").to_string();
            match fs::write_file(path, content, state).await {
                Ok(()) => "File written successfully.".to_string(),
                Err(e) => format!("Error: {e}"),
            }
        }
        "list_directory" => {
            let path = args["path"].as_str().unwrap_or("").to_string();
            match fs::list_directory(path, state).await {
                Ok(entries) => serde_json::to_string(&entries).unwrap_or_default(),
                Err(e) => format!("Error: {e}"),
            }
        }
        "search_files" => {
            let query = args["query"].as_str().unwrap_or("").to_string();
            match search::search_files(query, state).await {
                Ok(results) => serde_json::to_string(&results).unwrap_or_default(),
                Err(e) => format!("Error: {e}"),
            }
        }
        "exec_shell" => {
            let command = args["command"].as_str().unwrap_or("").to_string();
            let args_vec: Vec<String> = args["args"]
                .as_array()
                .map(|arr| {
                    arr.iter()
                        .map(|v| v.as_str().unwrap_or("").to_string())
                        .collect()
                })
                .unwrap_or_default();

            match shell::exec_shell(command, args_vec, state).await {
                Ok(output) => serde_json::to_string(&output).unwrap_or_default(),
                Err(e) => format!("Error: {e}"),
            }
        }
        _ => format!("Unknown tool: {name}"),
    }
}

pub fn cancel_chat(state: &SessionState) -> Result<(), String> {
    state.cancel_tx.send(true).map_err(|e| e.to_string())?;
    Ok(())
}