magi-code 0.102.2

use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};

const MAX_WORKERS: usize = 4;

type DispatchResults = Vec<(usize, ToolRuntime, ToolDispatchOutcome)>;

pub(super) fn eligible(run: &ToolLifecycleRun<'_, '_, '_, '_>, pending: &PendingToolCall) -> bool {
    run.tools.is_some() && eligible_call(run.hooks, &pending.call, pending.rejection.is_some())
}

pub(in crate::agent) fn eligible_call(
    hooks: Option<&HookRuntime>,
    call: &ToolCall,
    refused: bool,
) -> bool {
    hooks.is_none_or(|hooks| !hooks.has_configured_hooks())
        && !refused
        && matches!(
            ToolCapability::from_dispatch_name(&call.name),
            Some(
                ToolCapability::Read
                    | ToolCapability::Find
                    | ToolCapability::ListFiles
                    | ToolCapability::Grep
                    | ToolCapability::AstGrep
                    | ToolCapability::Web
                    | ToolCapability::ViewImage
            )
        )
}

pub(super) fn run(
    run: &mut ToolLifecycleRun<'_, '_, '_, '_>,
    group: Vec<(usize, PendingToolCall)>,
) -> anyhow::Result<()> {
    // Nothing can dispatch until every call in this run has its durable record.
    let prepared = group
        .into_iter()
        .map(|(index, pending)| prepare_tool_call(run, index, pending))
        .collect::<anyhow::Result<Vec<_>>>()?;
    let tools = run.tools.expect("eligible calls have a tool runtime");
    let calls: Vec<_> = prepared
        .iter()
        .map(|prepared| (prepared.call.clone(), prepared.context.clone()))
        .collect();
    let mut prepared = prepared.into_iter();
    let result = dispatch_ordered(tools, &calls, |_, outcome| {
        let prepared = prepared.next().expect("one completion per prepared call");
        let PreparedToolCall {
            call,
            activity,
            context,
            ..
        } = prepared;
        run.turn_state.append_function_call_if_missing(&call);
        // Record even canceled/queued calls. Never leave durable calls dangling merely
        // because cancellation arrived while the scoped workers were joining.
        let completion = record_tool_dispatch(
            Some(tools),
            None,
            run.session_persistence,
            run.output_sink,
            run.subdir_instruction_state.as_deref_mut(),
            call.clone(),
            context,
            outcome.into(),
        );
        completion.and_then(|outcome| {
            finish_tool_call(run, call, activity, outcome, run.cancellation.is_canceled())
        })
    });
    run.cancellation.check()?;
    result
}

/// One bounded worker path for direct and Code Mode inspection calls. All workers
/// join before snapshots and completions are delivered on the caller in input order.
pub(in crate::agent) fn dispatch_ordered(
    tools: &ToolRuntime,
    calls: &[(ToolCall, ToolDispatchContext)],
    mut complete: impl FnMut(usize, ToolDispatchOutcome) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
    let next = AtomicUsize::new(0);
    let dispatches = std::thread::scope(|scope| {
        let mut handles = Vec::new();
        for _ in 0..calls.len().min(MAX_WORKERS) {
            match std::thread::Builder::new()
                .name("inspection-tool".to_string())
                .spawn_scoped(scope, || dispatch_pending(tools, calls, &next))
            {
                Ok(handle) => handles.push(handle),
                Err(_) => break,
            }
        }
        if handles.is_empty() {
            return dispatch_pending(tools, calls, &next);
        }
        let mut results = Vec::with_capacity(calls.len());
        for handle in handles {
            if let Ok(completed) = handle.join() {
                results.extend(completed);
            }
        }
        results
    });
    let mut ordered: Vec<_> = (0..calls.len()).map(|_| None).collect();
    for (index, runtime, outcome) in dispatches {
        ordered[index] = Some((runtime, outcome));
    }
    let mut first_error = None;
    for (index, ((call, _), dispatch)) in calls.iter().zip(ordered).enumerate() {
        let outcome = match dispatch {
            Some((runtime, outcome)) => {
                if let Err(error) = tools.record_parallel_read_snapshots(&runtime) {
                    first_error.get_or_insert(error);
                }
                outcome
            }
            None => ToolDispatchOutcome {
                result: ToolResult {
                    tool_name: call.name.clone(),
                    success: false,
                    content: "inspection tool worker panicked".to_string(),
                    metadata: json!({}),
                    display: ToolResultDisplay::default(),
                },
                touched_paths: Vec::new(),
                changed_paths: Vec::new(),
            },
        };
        if let Err(error) = complete(index, outcome) {
            first_error.get_or_insert(error);
        }
    }
    first_error.map_or(Ok(()), Err)
}

fn dispatch_pending(
    tools: &ToolRuntime,
    calls: &[(ToolCall, ToolDispatchContext)],
    next: &AtomicUsize,
) -> DispatchResults {
    let mut results = Vec::new();
    loop {
        let index = next.fetch_add(1, Ordering::Relaxed);
        let Some((call, context)) = calls.get(index) else {
            break;
        };
        let runtime = tools.for_parallel_call();
        // Normal dispatch retains cancellation, inspection and per-surface ceilings.
        let outcome = runtime.dispatch_with_context_outcome(
            &call.name,
            call.arguments.clone(),
            context.clone(),
        );
        results.push((index, runtime, outcome));
    }
    results
}