YiemAgent/etc.jl

# ── executeToolCalls() Julia pseudo code ──────────────────────────
# Full call stack from runLoop → executeToolCalls → prepare → execute → finalize → emit

"""
    preparedToolCall(tool, toolCall, args)

Intermediate state between tool call validation and execution.
This struct is created after `prepareToolCall` succeeds and serves
as the bridge to the execution phase. Keeping the resolved tool,
original call metadata, and validated args together avoids repeated
lookups and allows the execution phase to access all necessary data
without carrying the full context through the call chain.
"""
struct preparedToolCall
    tool::agentTool               # The resolved tool definition from the context
    toolCall::agentToolCall       # The original tool call from the assistant
    args::any                      # Validated (and coerced) argument values
end

"""
    immediateOutcome(result, isError)

A tool call that was resolved without actual execution — either
because the tool was not found, validation failed, or a
`beforeToolCall` hook blocked the call. The result is produced
immediately and emitted as a tool result message.

Returning an outcome instead of throwing an exception is intentional:
it lets the agent feed the error back to the LLM as a tool result so
the model can recover — for example, by re-issuing a tool call with
corrected arguments after a validation failure.
"""
struct immediateOutcome
    result::agentToolResult         # The pre-computed tool result
    isError::bool                  # Whether this outcome represents an error
end

"""
    executedOutcome(result, isError)

A tool call that has been executed by `tool.execute()` but has not
yet been through the `afterToolCall` hook. This intermediate state
is necessary because the hook may mutate the result (content, usage,
termination, error status). Keeping execution and finalization separate
allows the hook to inspect the raw result and decide whether to
transform it or replace it entirely.
"""
struct executedOutcome
    result::agentToolResult         # The tool's execution result
    isError::bool                  # Whether execution raised an error
end

"""
    finalizedOutcome(toolCall, result, isError)

The complete outcome of a tool call after both execution and the
`afterToolCall` hook. This is the final form that is used to
construct the `toolResultMessage` emitted to the agent loop.

The three-phase design (prepare → execute → finalize) exists so that
each phase has a single responsibility: preparation handles validation
and gating, execution performs the actual work, and finalization
applies post-processing hooks. This separation allows the agent loop
to emit `tool_execution_end` events with the finalized data while
keeping each phase independently testable and swappable.
"""
struct finalizedOutcome
    toolCall::agentToolCall        # The original tool call reference
    result::agentToolResult         # The final tool result (post-afterToolCall)
    isError::bool                  # Whether the call failed or was blocked
end

"""
    ToolCallBatch(messages, terminate)

A batch of tool result messages from executing one or more tool calls.
The `terminate` flag indicates whether all tools in the batch requested
termination, which causes the agent loop to stop processing further turns.

This flag is set by the tool implementation (not the end user) to signal
that the agent should not call the LLM again. Typical use cases:

- Task completion: a tool like `deploy` or `submit` finishes its work and
  returns `terminate: true` so the agent stops instead of asking the LLM
  what to do next.
- Unrecoverable error: a tool hits a fatal condition (e.g. database
  connection lost, auth token expired) and returns `terminate: true` so
  the agent stops with an error message rather than retrying.
- Async handoff: a tool triggers a long-running external operation and
  wants the agent to stop now; the external system will later resume the
  agent via `continue()`.

If `terminate` is `false` (default), the agent loop feeds the tool results
back to the LLM for another turn.
"""
struct toolCallBatch
    messages::vector{toolResultMessage}  # Tool result messages for this batch
    terminate::bool                      # Whether the batch should terminate the loop
end

"""
    createErrorToolResult(msg)

Builds an `agentToolResult` containing a single text content item
with the provided error message and an empty details dictionary.
Used when a tool call cannot be executed due to errors.

Returning a result instead of throwing ensures that errors at any
point in the tool call pipeline are fed back to the LLM as a tool
result message. This allows the model to see the error and decide
whether to retry, re-issue the call with different arguments, or
report failure to the user.
"""
function createErrorToolResult(msg::String)::agentToolResult
    return agentToolResult([textContent("text", msg)], dict{any,any}())
end

"""
    createToolResultMessage(f)

Constructs a `toolResultMessage` from a `finalizedOutcome`.
Normalizes missing content to an empty array and includes
the `addedToolNames` field only when the tool dynamically
registered new tools during execution.

This conversion is necessary because the tool result is an
`agentToolResult` used by tool implementations, while the agent
loop consumes `toolResultMessage` objects that become part of the
conversation history. The message format includes metadata like
timestamp and tool call ID that the raw result does not carry,
and it is the object emitted via `message_start`/`message_end`
events so the LLM receives the result as a proper assistant/user
message in the context window.
"""
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, nowMillis()
    )
end

"""
    shouldTerminate(finalizedCalls)

Returns `true` only when every finalized call in the batch has
`result.terminate == true`. All tools must agree — if any tool
did not request termination, the agent continues. This prevents
a single tool that happens to set `terminate: true` (e.g. for
metadata purposes) from accidentally stopping the agent when
other tools in the batch did not intend to terminate.
"""
function shouldTerminate(batches::vector{finalizedOutcome})::bool
    return !isempty(batches) && all(b -> b.result.terminate, batches)
end

"""
    prepareToolCallArguments(tool, toolCall)

Calls the tool's optional `prepareArguments` hook to transform
the raw argument values from the LLM before schema validation.
If the tool has no hook or the hook returns the same object
reference, the original call is returned unchanged.

This hook allows tools to normalize arguments that the LLM may
have produced in a non-standard format — for example, converting
a date string to a timestamp, expanding a short file path to an
absolute path, or normalizing casing. It runs before schema
validation so the validator sees the normalized form rather than
raw LLM output.
"""
function prepareToolCallArguments(tool::agentTool, toolCall::agentToolCall)::agentToolCall
    if tool.prepareArguments === nothing
        return toolCall
    end
    prepared = tool.prepareArguments(toolCall.arguments)
    if prepared == toolCall.arguments
        return toolCall
    end
    return merge(toolCall, dict(:arguments => prepared))
end

"""
    prepareToolCall(context, assistantMsg, toolCall, config, signal)

Resolves the tool by name, prepares and validates its arguments,
and runs the `beforeToolCall` hook. Returns a `preparedToolCall`
if successful or an `immediateOutcome` if the tool is not found,
validation fails, the hook blocks execution, or the signal is
aborted. Errors during preparation are caught and returned as
immediate error outcomes so the agent loop can feed them back
to the model.

The key design decision here is that preparation never throws.
Every failure path returns an `immediateOutcome` with an error
result. This ensures the agent loop always receives a valid tool
result message for every tool call the assistant requested,
regardless of whether preparation succeeded. The LLM can then
use the error message to decide whether to retry with different
arguments or acknowledge the failure.
"""
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)
        prepared = prepareToolCallArguments(tool, toolCall)
        validatedArgs = validateToolArguments(tool, prepared)

        # 2. beforeToolCall hook — can block
        if config.beforeToolCall !== nothing
            before = config.beforeToolCall(
                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 ──────────────────────────────────────────

"""
    executePreparedToolCall(prep, signal, emit)

Executes the tool by calling `tool.execute()` with the validated
arguments, the abort signal, and a callback for streaming partial
results. Emits `toolExecutionUpdate` events for each partial
result batch. Waits for all pending update events to settle before
returning. Catches execution errors and returns them as an error
outcome. The `accepting` guard prevents emitting updates after
the call has finished.

Long-running tools (e.g. file uploads, model training, web scraping)
may take seconds or minutes. The streaming update mechanism allows
UI listeners and other consumers to show progress in real time rather
than waiting for the entire call to complete. The `accepting` guard
ensures that if the tool's execute function yields after emitting
updates but before returning, no duplicate or stale updates are
emitted after the result has already been captured.
"""
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 ───────────────────────────────────────

"""
    finalizeExecutedToolCall(context, assistantMsg, prep, executed, config, signal)

Runs the `afterToolCall` hook on the executed result, allowing
the consumer to mutate the result content, details, usage,
termination flag, or error status. Catches errors from the
hook and converts them to error outcomes. Returns a
`finalizedOutcome` that is used to construct the tool result
message.

The `afterToolCall` hook exists as a post-processing step that
runs after every tool call regardless of success or failure.
Common use cases include:

- Masking sensitive data from result content before the LLM
  sees it (e.g. removing API keys from error messages).
- Normalizing usage tracking data into a consistent format.
- Inspecting the result and deciding to flip `terminate: true`
  based on business logic (e.g. "if deployment failed, stop
  the agent rather than retrying").
- Wrapping an error result in a friendlier message for the LLM
  to understand.

If the hook itself throws, the error is caught and the result
becomes an error outcome. This ensures the tool pipeline never
breaks due to a buggy hook.
"""
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.afterToolCall !== nothing
        try
            after = config.afterToolCall(
                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, :isError, isError)
            end
        catch err
            result = createErrorToolResult(sprint(showerror, err))
            isError = true
        end
    end

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

"""
    emitToolExecutionEnd(finalized, emit)

Emits the `toolExecutionEnd` event with the finalized outcome,
signalling to listeners that the tool call has completed.

This event is part of the tool execution lifecycle:
`tool_execution_start` → (zero or more `tool_execution_update` events) →
`tool_execution_end`. Listeners (such as the TUI or logging systems)
use this lifecycle to track individual tool calls. The event carries
the final result so listeners have all the data they need without
requiring external state lookups.
"""
function emitToolExecutionEnd(finalized::finalizedOutcome, emit::agentEventSink)
    emit(toolExecEndEvent(finalized.toolCall.id, finalized.toolCall.name,
                          finalized.result, finalized.isError))
end

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

"""
    executeToolCallsSequential(context, assistantMsg, toolCalls, config, signal, emit)

Executes tool calls one at a time in the order they appear. For each
call: emits `toolExecutionStart`, runs `prepareToolCall`,
then either resolves the immediate outcome or executes/finalizes
the prepared call. Emits `toolExecutionEnd` and creates the tool
result message before proceeding to the next call. Respects the
abort signal — if aborted, remaining calls are skipped. Returns
a batch with `terminate` determined by whether all results set
the termination flag.

Sequential execution is required when tool calls have implicit
dependencies — for example, a `create_database` tool must complete
before `create_table` can reference it. It is also the safer
default because it prevents race conditions when multiple tools
share state (e.g. writing to the same file or API rate limits).
Use parallel only when you are confident the tools are independent.
"""
function executeToolCallsSequential(
    context::agentContext,
    assistantMsg::assistantMessage,
    toolCalls::vector{agentToolCall},
    config::agentLoopConfig,
    signal::union{nothing,abortSignal},
    emit::agentEventSink,
)::agentToolCallBatch

    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 ──────────────────────────────────────────

"""
    executeToolCallsParallel(context, assistantMsg, toolCalls, config, signal, emit)

Prepares all tool calls concurrently and spawns a task for each
prepared call. Immediate outcomes are resolved instantly. Task
entries are collected in order, then `fetch`ed to await all
concurrent executions. Tool result messages are created from
finalized outcomes in order and returned as a batch. Respects
the abort signal — if aborted during preparation, remaining
calls are skipped. Finalization order preserves the original
call order.

Parallel execution is appropriate when the assistant requests
independent tools — for example, reading multiple files, querying
separate databases, or making independent API calls. It reduces
wall-clock time compared to sequential execution. The tradeoff is
that parallel calls can overwhelm external resources (rate limits,
connection pools, disk I/O). Finalization preserves the original
call order so tool result messages appear in the same order the
assistant requested them, regardless of which call finishes first.
"""
function executeToolCallsParallel(
    context::agentContext,
    assistantMsg::assistantMessage,
    toolCalls::vector{agentToolCall},
    config::agentLoopConfig,
    signal::union{nothing,abortSignal},
    emit::agentEventSink,
)::agentToolCallBatch

    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
            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

    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


"""
    executeToolCalls(context, assistantMsg, toolCalls, config, signal, emit)

Dispatches to sequential or parallel execution. Uses sequential mode
when `config.toolExecution == "sequential"` or when any of the
tool calls reference a tool with `executionMode: "sequential"`.
Otherwise uses parallel execution. This is the entry point called
from `streamAssistantResponse` in the agent loop.

The sequential mode takes priority over parallel because it is the
safe default. If even one tool in a batch is marked sequential, all
tools execute sequentially — this prevents a single dependent tool
from racing with an otherwise independent one. The per-tool
`executionMode` allows fine-grained control (e.g. most tools are
parallel but a specific write tool is sequential), while the config-level
`toolExecution` provides a global override.
"""
function executeToolCalls(
    context::agentContext,
    assistantMsg::assistantMessage,
    toolCalls::vector{agentToolCall},
    config::agentLoopConfig,
    signal::union{nothing,abortSignal},
    emit::agentEventSink,
)::agentToolCallBatch

    hasSequential = any(tc ->
        any(t -> t.name == tc.name && get(t.executionMode, "parallel") == "sequential",
            context.tools), toolCalls)

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