zesdex/apps/interfaces/tui/src/turn.rs

//! Agent turn engine — runs LLM + tool execution on a background thread.
//!
//! Flow: push user message → spawn OS thread → loop: call blocking LLM
//! client → execute tool calls → push TurnEvents → repeat until done.

use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::collections::VecDeque;

use tracing::{debug, info, warn};

use zesdex_domain::core::tool_call::sanitize_tool_arguments;
use zesdex_domain::core::ChatMessage;
use zesdex_infrastructure::llm::provider::LlmClient;
use zesdex_infrastructure::tools::{all_tools, tool_defs, ToolCtx};
use zesdex_infrastructure::TurnEvent;
use crate::state::AppStateRest;

/// Spawn an agent turn on a background OS thread.
pub fn spawn_agent_turn(state: &mut AppStateRest, text: String) {
    // compare_exchange: only mark in-flight if not already running
    if state.turn_in_flight_flag
        .compare_exchange(false, true, Ordering::SeqCst, Ordering::Relaxed)
        .is_err()
    {
        return; // already running
    }

    let turn_events = state.turn_events.clone();
    let in_flight = state.turn_in_flight_flag.clone();
    let abort = state.abort_flag.clone();
    let session_dir = state.session_dir.clone();
    let workspace_roots = state.workspace_roots.clone();

    // Resolve LLM provider configuration from settings
    let provider_name = &state.settings.provider;
    let provider_cfg = state.app_config.providers.get(provider_name).cloned();
    
    let mut api_key = String::new();
    if let Some(key) = state.settings.api_keys.get(provider_name) {
        api_key = key.clone();
    } else if let Some(ref cfg) = provider_cfg {
        if let Some(ref default_key) = cfg.default_api_key {
            api_key = default_key.clone();
        }
        if api_key.is_empty() {
            if let Some(ref env_name) = cfg.api_key_env {
                if let Ok(val) = std::env::var(env_name) {
                    api_key = val;
                }
            }
        }
    }

    let model = state.settings.model.clone();
    let api_base = provider_cfg.map(|cfg| cfg.api_base.clone());

    let mut messages: Vec<ChatMessage> = state
        .session_runtime
        .as_ref()
        .map(|rt| rt.messages.clone())
        .unwrap_or_default();
    messages.push(ChatMessage::user(text));

    if let Some(ref mut rt) = state.session_runtime {
        rt.messages = messages.clone();
    }

    info!("spawning agent turn with {} messages (model: {})", messages.len(), model);

    std::thread::spawn(move || {
        run_turn(
            &mut messages,
            &session_dir,
            &workspace_roots,
            &turn_events,
            &in_flight,
            &abort,
            api_key,
            model,
            api_base,
        );
    });
}

/// The core agent turn — LLM call → tool execution → repeat.
fn run_turn(
    messages: &mut Vec<ChatMessage>,
    session_dir: &Path,
    workspace_roots: &[PathBuf],
    turn_events: &Arc<Mutex<VecDeque<TurnEvent>>>,
    in_flight: &Arc<AtomicBool>,
    abort: &Arc<AtomicBool>,
    api_key: String,
    model: String,
    api_base: Option<String>,
) {
    let client = LlmClient::new(api_key, model, api_base);

    let tools = all_tools();
    let defs = tool_defs(&tools);

    // Prepend system message
    let sys_msg = ChatMessage::system(
        "You are Zesdex, an AI coding assistant. You have access to various tools via native function calling to help the user. \
         For any non-trivial tasks, you MUST prioritize creating a structured plan (using `plan_enter`) and a list of TODOs (using `todowrite`) BEFORE executing any other tools or modifying files. \
         Use other tools when necessary. If a tool returns an error, fix the issue before retrying. \
         Respond conversationally, concisely, and helpfully."
            .to_string(),
    );
    messages.insert(0, sys_msg);

    let tool_ctx = ToolCtx::builder()
        .session_dir(session_dir.to_path_buf())
        .workspaces(workspace_roots.to_vec())
        .build();

    for iteration in 0..50 {
        if abort.load(Ordering::SeqCst) {
            abort.store(false, Ordering::SeqCst);
            push_event(turn_events, TurnEvent::SystemNote {
                kind: "info".into(),
                message: "Turn aborted by user".into(),
            });
            break;
        }

        debug!("agent turn iteration {iteration}");

        // Stream the LLM response
        push_event(turn_events, TurnEvent::StreamStart);

        let result = client.chat_with_tools_streaming(
            messages,
            Some(defs.clone()),
            Some(0.7),
            Some(4096),
            |event| {
                if abort.load(Ordering::SeqCst) {
                    return false;
                }
                match event {
                    zesdex_domain::core::StreamEvent::Token(s) | zesdex_domain::core::StreamEvent::Reasoning(s) => {
                        push_event(turn_events, TurnEvent::StreamToken(s.clone()));
                    }
                    _ => {}
                }
                true
            },
            Some(&abort),
        );

        match result {
            Ok((assistant_msg, usage)) => {
                let content = assistant_msg.content.clone().unwrap_or_default();
                let tool_calls = assistant_msg.tool_calls.clone().unwrap_or_default();

                push_event(turn_events, TurnEvent::StreamDone(assistant_msg.clone()));

                if let Some((tokens_in, tokens_out)) = usage {
                    push_event(turn_events, TurnEvent::Usage {
                        tokens_in,
                        tokens_out,
                    });
                }

                if tool_calls.is_empty() {
                    messages.push(ChatMessage::assistant(Some(content)));
                    break;
                }

                messages.push(assistant_msg);

                for tc in &tool_calls {
                    let name = &tc.function.name;
                    let args = sanitize_tool_arguments(&tc.function.arguments);

                    debug!("executing tool: {name}");

                    let output = if let Some(tool) = tools.iter().find(|t| t.name() == name) {
                        match tool.run(&tool_ctx, &args) {
                            Ok(o) => o,
                            Err(e) => format!("Error: {e}"),
                        }
                    } else {
                        format!("Unknown tool: {name}")
                    };

                    let is_error = output.starts_with("Error:");

                    push_event(turn_events, TurnEvent::ToolResult {
                        tool_call_id: tc.id.clone(),
                        tool_name: name.clone(),
                        output: output.clone(),
                        is_error,
                        path: None,
                    });

                    messages.push(ChatMessage::tool(tc.id.clone(), output.clone()));
                }
            }
            Err(e) => {
                warn!("LLM call failed: {e}");
                push_event(turn_events, TurnEvent::Error(format!("LLM error: {e}")));
                break;
            }
        }
    }

    // Propagate accumulated messages back to session_runtime so the next
    // turn starts with the full history (assistant replies + tool results).
    // TurnEvent::Compacted already exists on the enum and is handled in
    // action.rs to write back to state.session_runtime.messages.
    push_event(turn_events, TurnEvent::Compacted(messages.clone()));
    push_event(turn_events, TurnEvent::Done);
    mark_done(in_flight);
}

fn push_event(queue: &Arc<Mutex<VecDeque<TurnEvent>>>, event: TurnEvent) {
    if let Ok(mut q) = queue.lock() {
        q.push_back(event);
    }
}

/// Mark the turn as done using lock-free atomic store.
fn mark_done(flag: &Arc<AtomicBool>) {
    flag.store(false, Ordering::SeqCst);
}