magi-code 0.100.0

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()
        && run.hooks.is_none_or(|hooks| !hooks.has_configured_hooks())
        && pending.rejection.is_none()
        && matches!(
            ToolCapability::from_dispatch_name(&pending.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 next = AtomicUsize::new(0);
    let dispatches = std::thread::scope(|scope| {
        let mut handles = Vec::new();
        for _ in 0..prepared.len().min(MAX_WORKERS) {
            match std::thread::Builder::new()
                .name("inspection-tool".to_string())
                .spawn_scoped(scope, || dispatch_pending(tools, &prepared, &next))
            {
                Ok(handle) => handles.push(handle),
                // Existing workers drain the queue. If none started, use this thread.
                Err(_) => break,
            }
        }
        if handles.is_empty() {
            return dispatch_pending(tools, &prepared, &next);
        }
        let mut results = Vec::with_capacity(prepared.len());
        for handle in handles {
            if let Ok(completed) = handle.join() {
                results.extend(completed);
            }
        }
        results
    });
    let mut ordered: Vec<_> = (0..prepared.len()).map(|_| None).collect();
    for (index, runtime, outcome) in dispatches {
        ordered[index] = Some((runtime, outcome));
    }
    let mut first_error = None;
    for (prepared, dispatch) in prepared.into_iter().zip(ordered) {
        let PreparedToolCall {
            call,
            activity,
            context,
            ..
        } = prepared;
        run.turn_state.append_function_call_if_missing(&call);
        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(),
            },
        };
        // 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(),
        );
        let completion = completion.and_then(|outcome| {
            finish_tool_call(run, call, activity, outcome, run.cancellation.is_canceled())
        });
        if let Err(error) = completion {
            first_error.get_or_insert(error);
        }
    }
    run.cancellation.check()?;
    first_error.map_or(Ok(()), Err)
}

fn dispatch_pending(
    tools: &ToolRuntime,
    prepared: &[PreparedToolCall],
    next: &AtomicUsize,
) -> DispatchResults {
    let mut results = Vec::new();
    loop {
        let index = next.fetch_add(1, Ordering::Relaxed);
        let Some(prepared) = prepared.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(
            &prepared.call.name,
            prepared.call.arguments.clone(),
            prepared.context.clone(),
        );
        results.push((index, runtime, outcome));
    }
    results
}