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.
///
/// Flow: compare-exchange the in-flight flag → snapshot state fields →
/// clone session runtime messages → push user message → spawn OS thread
/// that runs `run_turn`.
#[tracing::instrument(skip(state))]
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 || {
        let params = TurnParams {
            messages: &mut messages,
            session_dir: &session_dir,
            workspace_roots: &workspace_roots,
            turn_events: &turn_events,
            in_flight: &in_flight,
            abort: &abort,
            api_key,
            model,
            api_base,
        };
        run_turn(params);
    });
}

/// Parameters for a turn, grouped to avoid too-many-arguments lint.
///
/// Holds all the references and owned values that `run_turn` needs:
/// message history, turn-event queue, abort/in-flight flags, API credentials,
/// and environment paths.
struct TurnParams<'a> {
    messages: &'a mut Vec<ChatMessage>,
    session_dir: &'a Path,
    workspace_roots: &'a [PathBuf],
    turn_events: &'a Arc<Mutex<VecDeque<TurnEvent>>>,
    in_flight: &'a Arc<AtomicBool>,
    abort: &'a Arc<AtomicBool>,
    api_key: String,
    model: String,
    api_base: Option<String>,
}

/// The core agent turn — LLM call → tool execution → repeat.
///
/// Flow: build `LlmClient` → compile tools → prepend system message →
/// loop (max 50 iterations): abort check → stream LLM response →
/// push events → execute tool calls → push results → break on
/// no tool calls or error → emit final `Compacted` + `Done`.
#[tracing::instrument(skip(params))]
fn run_turn(params: TurnParams) {
    let client = LlmClient::new(params.api_key, params.model, params.api_base);

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

    let mut sys_prompt = "You are Zesdex, an AI coding assistant. You have access to various tools via native function calling to help the user. \
         For complex problems, use `seq_think` to reason step-by-step. \
         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. \
         For large multi-step operations or delegating tasks, you MUST prioritize using `workflow_run` (to run a yaml workflow script) or `hive_mind` (to orchestrate multiple agents) to complete the task efficiently. \
         Use other tools when necessary. If a tool returns an error, fix the issue before retrying. \
         Respond conversationally, concisely, and helpfully.".to_string();

    if let Some(root) = params.workspace_roots.first() {
        let tree = zesdex_infrastructure::utils::build_workspace_tree(root, 800);
        let rich_ctx = zesdex_infrastructure::utils::build_rich_context(root);
        
        sys_prompt.push_str("\n\nWorkspace structure:\n```\n");
        sys_prompt.push_str(&tree);
        sys_prompt.push_str("\n```\n\n");
        
        sys_prompt.push_str(&rich_ctx);
    }

    let sys_msg = ChatMessage::system(sys_prompt);
    params.messages.insert(0, sys_msg);

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

    for iteration in 0..50 {
        if params.abort.load(Ordering::SeqCst) {
            params.abort.store(false, Ordering::SeqCst);
            push_event(params.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(params.turn_events, TurnEvent::StreamStart);

        let result = client.chat_with_tools_streaming(
            params.messages,
            Some(defs.clone()),
            Some(0.7),
            Some(4096),
            |event| {
                if params.abort.load(Ordering::SeqCst) {
                    return false;
                }
                match event {
                    zesdex_domain::core::StreamEvent::Token(s) => {
                        push_event(params.turn_events, TurnEvent::StreamToken(s.clone()));
                    }
                    zesdex_domain::core::StreamEvent::Reasoning(s) => {
                        push_event(params.turn_events, TurnEvent::StreamReasoning(s.clone()));
                    }
                    _ => {}
                }
                true
            },
            Some(params.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(params.turn_events, TurnEvent::StreamDone(assistant_msg.clone()));

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

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

                params.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(params.turn_events, TurnEvent::ToolResult {
                        tool_call_id: tc.id.clone(),
                        tool_name: name.clone(),
                        output: output.clone(),
                        is_error,
                        path: None,
                    });

                    params.messages.push(ChatMessage::tool(tc.id.clone(), output.clone()));
                }
            }
            Err(e) => {
                warn!("LLM call failed: {e}");
                push_event(params.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).
    push_event(params.turn_events, TurnEvent::Compacted(params.messages.clone()));
    push_event(params.turn_events, TurnEvent::Done);
    mark_done(params.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);
}