YiemAgent/etc.jl

struct preparedToolCall
    tool::AgentTool
    toolCall::AgentToolCall
    args::Any
end

struct immediateOutcome
    result::AgentToolResult
    isError::Bool
end

struct executedOutcome
    result::AgentToolResult
    isError::Bool
end

struct finalizedOutcome
    toolCall::AgentToolCall
    result::AgentToolResult
    isError::Bool
end

struct toolCallBatch
    messages::Vector{ToolResultMessage}
    terminate::Bool
end

# ── helpers ─────────────────────────────────────────────────────

function createErrorToolResult(msg::String)::AgentToolResult
    return AgentToolResult([TextContent("text", msg)], Dict{Any,Any}())
end

function createToolResultMessage(f::finalizedOutcome)::ToolResultMessage
    return ToolResultMessage(
        "toolResult", f.toolCall.id, f.toolCall.name,
        f.result.content, f.result.details, f.result.usage,
        get(f.result, :addedToolNames, String[]), f.isError, now_millis()
    )
end

function shouldTerminate(batches::Vector{finalizedOutcome})::Bool
    return !isempty(batches) && all(b -> b.result.terminate, batches)
end

# ── per-call preparation ────────────────────────────────────────

function prepareToolCall(
    context::AgentContext,
    assistantMsg::AssistantMessage,
    toolCall::AgentToolCall,
    config::AgentLoopConfig,
    signal::Union{Nothing,AbortSignal},
)::Union{preparedToolCall,immediateOutcome}

    tool = find(t -> t.name == toolCall.name, context.tools)
    if tool === nothing
        return immediateOutcome(createErrorToolResult("Tool $toolCall.name not found"), true)
    end

    try
        # 1. prepare arguments (tool-specific transform)
        preparedArgs = prepareToolCallArguments(tool, toolCall)
        validatedArgs = validateToolArguments(tool, preparedArgs)

        # 2. beforeToolCall hook
        if config.before_tool_call !== nothing
            before = config.before_tool_call(
                AssistantMsgCtx(assistantMsg, toolCall, validatedArgs, context), signal
            )
            if signal !== nothing && signal.aborted
                return immediateOutcome(createErrorToolResult("Operation aborted"), true)
            end
            if before !== nothing && before.block
                return immediateOutcome(
                    createErrorToolResult(get(before, :reason, "Tool execution was blocked")), true)
            end
        end

        return preparedToolCall(tool, toolCall, validatedArgs)
    catch err
        return immediateOutcome(createErrorToolResult(sprint(showerror, err)), true)
    end
end

# ── per-call execution ──────────────────────────────────────────

function executePreparedToolCall(
    prep::preparedToolCall,
    signal::Union{Nothing,AbortSignal},
    emit::AgentEventSink,
)::executedOutcome

    updateEvents = Promise[]
    accepting = true

    try
        result = prep.tool.execute(
            prep.toolCall.id, prep.args, signal,
            partialResult -> begin
                if accepting
                    push!(updateEvents,
                          emit(ToolExecUpdateEvent(prep.toolCall.id, prep.toolCall.name,
                                                   prep.toolCall.arguments, partialResult)))
                end
            end
        )
        accepting = false
        wait.(updateEvents)
        return executedOutcome(result, false)
    catch err
        accepting = false
        wait.(updateEvents)
        return executedOutcome(createErrorToolResult(sprint(showerror, err)), true)
    end
end

# ── per-call finalization ───────────────────────────────────────

function finalizeExecutedToolCall(
    context::AgentContext,
    assistantMsg::AssistantMessage,
    prep::preparedToolCall,
    executed::executedOutcome,
    config::AgentLoopConfig,
    signal::Union{Nothing,AbortSignal},
)::finalizedOutcome

    result = executed.result
    isError = executed.isError

    if config.afterToolCalls !== nothing
        try
            after = config.afterToolCalls(
                AfterCtx(assistantMsg, prep.toolCall, prep.args, result, isError, context), signal
            )
            if after !== nothing
                result = merge(result, Dict(:content=>get(after,:content,result.content),
                                            :details=>get(after,:details,result.details),
                                            :usage=>get(after,:usage,result.usage),
                                            :terminate=>get(after,:terminate,result.terminate)))
                isError = get(after, :is_error, isError)
            end
        catch err
            result = createErrorToolResult(sprint(showerror, err))
            isError = true
        end
    end

    return finalizedOutcome(prep.toolCall, result, isError)
end

function emitToolExecutionEnd(finalized::finalizedOutcome, emit::AgentEventSink)
    emit(ToolExecEndEvent(finalized.toolCall.id, finalized.toolCall.name,
                          finalized.result, finalized.isError))
end

# ── sequential execution ────────────────────────────────────────

function executeToolCallsSequential(
    context::AgentContext,
    assistantMsg::AssistantMessage,
    toolCalls::Vector{AgentToolCall},
    config::AgentLoopConfig,
    signal::Union{Nothing,AbortSignal},
    emit::AgentEventSink,
)::toolCallBatch

    finalizedCalls = finalizedOutcome[]
    messages = ToolResultMessage[]

    for tc in toolCalls
        emit(ToolExecStartEvent(tc.id, tc.name, tc.arguments))

        prep = prepareToolCall(context, assistantMsg, tc, config, signal)

        if prep isa immediateOutcome
            finalized = finalizedOutcome(tc, prep.result, prep.isError)
        else
            executed = executePreparedToolCall(prep, signal, emit)
            finalized = finalizeExecutedToolCall(context, assistantMsg, prep, executed, config, signal)
        end

        emitToolExecutionEnd(finalized, emit)
        push!(messages, createToolResultMessage(finalized))
        push!(finalizedCalls, finalized)

        if signal !== nothing && signal.aborted
            break
        end
    end

    return toolCallBatch(messages, shouldTerminate(finalizedCalls))
end

# ── parallel execution ──────────────────────────────────────────

function executeToolCallsParallel(
    context::AgentContext,
    assistantMsg::AssistantMessage,
    toolCalls::Vector{AgentToolCall},
    config::AgentLoopConfig,
    signal::Union{Nothing,AbortSignal},
    emit::AgentEventSink,
)::toolCallBatch

    # Each entry: finalizedOutcome (already done) or Task → finalizedOutcome (pending)
    entries = Union{finalizedOutcome,Task{finalizedOutcome}}[]

    for tc in toolCalls
        emit(ToolExecStartEvent(tc.id, tc.name, tc.arguments))

        prep = prepareToolCall(context, assistantMsg, tc, config, signal)

        if prep isa immediateOutcome
            finalized = finalizedOutcome(tc, prep.result, prep.isError)
            emitToolExecutionEnd(finalized, emit)
            push!(entries, finalized)
        else
            # spawn lazy computation task
            task = Task() do
                executed = executePreparedToolCall(prep, signal, emit)
                finalized = finalizeExecutedToolCall(context, assistantMsg, prep, executed, config, signal)
                emitToolExecutionEnd(finalized, emit)
                return finalized
            end
            schedule(task)
            push!(entries, task)
        end

        if signal !== nothing && signal.aborted
            break
        end
    end

    # Wait for all tasks, collect in order
    finalizedCalls = finalizedOutcome[]
    for entry in entries
        outcome = entry isa Task ? fetch(entry) : entry
        push!(finalizedCalls, outcome)
    end

    messages = ToolResultMessage[]
    for f in finalizedCalls
        push!(messages, createToolResultMessage(f))
    end

    return toolCallBatch(messages, shouldTerminate(finalizedCalls))
end

# ── dispatcher ──────────────────────────────────────────────────

function executeToolCalls(
    context::AgentContext,
    assistantMsg::AssistantMessage,
    toolCalls::Vector{AgentToolCall},
    config::AgentLoopConfig,
    signal::Union{Nothing,AbortSignal},
    emit::AgentEventSink,
)::toolCallBatch

    # Check if any tool is marked sequential, or config forces sequential
    hasSequential = any(tc ->
        any(t -> t.name == tc.name && get(t.executionMode, "parallel") == "sequential",
            context.tools), toolCalls)

    if config.tool_execution == "sequential" || hasSequential
        return executeToolCallsSequential(context, assistantMsg, toolCalls, config, signal, emit)
    else
        return executeToolCallsParallel(context, assistantMsg, toolCalls, config, signal, emit)
    end
end