YiemAgent/src/agent_loop.jl

"""
    agent_loop.jl - Low-level agent loop implementation

This module implements the core agentLoop functionality that works with AgentMessage
throughout, transforming to Message[] only at the LLM call boundary.
"""

module AgentLoop

using ..Types: *
using ..StreamFn: *

# ============================================================================
# Event sink type
# ============================================================================

const AgentEventSink = Function

# ============================================================================
# Main agent loop function
# ============================================================================

function agentLoop(
    prompts::Vector{AgentMessage},
    context::AgentContext,
    config::AgentLoopConfig,
    signal::Union{Nothing, AbortSignal},
    stream_fn::StreamFn,
)::EventStream
    stream = createAgentStream()

    Threads.@spawn begin
        messages = runAgentLoop(
            prompts,
            context,
            config,
            (event) -> push!(stream, event),
            signal,
            stream_fn,
        )
        end(stream, messages)
    end

    return stream
end

# ============================================================================
# Continue agent loop function
# ============================================================================

function agentLoopContinue(
    context::AgentContext,
    config::AgentLoopConfig,
    signal::Union{Nothing, AbortSignal},
    stream_fn::StreamFn,
)::EventStream
    if isempty(context.messages)
        throw(ErrorException("Cannot continue: no messages in context"))
    end

    if context.messages[end].role == "assistant"
        throw(ErrorException("Cannot continue from message role: assistant"))
    end

    stream = createAgentStream()

    Threads.@spawn begin
        messages = runAgentLoopContinue(
            context,
            config,
            (event) -> push!(stream, event),
            signal,
            stream_fn,
        )
        end(stream, messages)
    end

    return stream
end

# ============================================================================
# Run agent loop function
# ============================================================================

function runAgentLoop(
    prompts::Vector{AgentMessage},
    context::AgentContext,
    config::AgentLoopConfig,
    emit::AgentEventSink,
    signal::Union{Nothing, AbortSignal},
    stream_fn::StreamFn,
)::Vector{AgentMessage}
    new_messages::Vector{AgentMessage} = copy(prompts)
    current_context::AgentContext = AgentContext(
        context.system_prompt,
        vcat(context.messages, copy(prompts)),
        context.tools,
    )

    emit(AgentStartEvent())
    emit(TurnStartEvent())
    for prompt in prompts
        emit(MessageStartEvent(prompt))
        emit(MessageEndEvent(prompt))
    end

    runLoop(
        current_context,
        new_messages,
        config,
        signal,
        emit,
        stream_fn,
    )
    return new_messages
end

# ============================================================================
# Run agent loop continue function
# ============================================================================

function runAgentLoopContinue(
    context::AgentContext,
    config::AgentLoopConfig,
    emit::AgentEventSink,
    signal::Union{Nothing, AbortSignal},
    stream_fn::StreamFn,
)::Vector{AgentMessage}
    if isempty(context.messages)
        throw(ErrorException("Cannot continue: no messages in context"))
    end

    if context.messages[end].role == "assistant"
        throw(ErrorException("Cannot continue from message role: assistant"))
    end

    new_messages::Vector{AgentMessage} = []
    current_context::AgentContext = context

    emit(AgentStartEvent())
    emit(TurnStartEvent())

    runLoop(
        current_context,
        new_messages,
        config,
        signal,
        emit,
        stream_fn,
    )
    return new_messages
end

# ============================================================================
# Create agent stream function
# ============================================================================

function createAgentStream()::EventStream
    return EventStream(
        (event::AgentEvent) -> event isa AgentEndEvent,
        (event::AgentEvent) -> event isa AgentEndEvent ? event.messages : AgentMessage[],
    )
end

# ============================================================================
# Main loop logic shared by agentLoop and agentLoopContinue
# ============================================================================

function runLoop(
    initial_context::AgentContext,
    new_messages::Vector{AgentMessage},
    initial_config::AgentLoopConfig,
    signal::Union{Nothing, AbortSignal},
    emit::AgentEventSink,
    stream_function::StreamFn,
)::Nothing
    current_context::AgentContext = initial_context
    config::AgentLoopConfig = initial_config
    first_turn::Bool = true
    pending_messages::Vector{AgentMessage} = getSteeringMessages(config) do
        get_steering_messages(config)
    end

    while true
        has_more_tool_calls::Bool = true

        while has_more_tool_calls || !isempty(pending_messages)
            if !first_turn
                emit(TurnStartEvent())
            else
                first_turn = false
            end

            if !isempty(pending_messages)
                for message in pending_messages
                    emit(MessageStartEvent(message))
                    emit(MessageEndEvent(message))
                    push!(current_context.messages, message)
                    push!(new_messages, message)
                end
                pending_messages = AgentMessage[]
            end

            message = streamAssistantResponse(
                current_context,
                config,
                signal,
                emit,
                stream_function,
            )
            push!(new_messages, message)

            if message.stop_reason in ("error", "aborted")
                emit(TurnEndEvent(message, ToolResultMessage[]))
                emit(AgentEndEvent(new_messages))
                return
            end

            tool_calls = filter(
                (c) -> c isa ToolCall,
                message.content,
            )

            tool_results::Vector{ToolResultMessage} = []
            has_more_tool_calls = false
            if !isempty(tool_calls)
                executed_tool_batch =
                    message.stop_reason == "length"
                        ? failToolCallsFromTruncatedMessage(tool_calls, emit)
                        : executeToolCalls(
                            current_context,
                            message,
                            config,
                            signal,
                            emit,
                        )
                append!(tool_results, executed_tool_batch.messages)
                has_more_tool_calls = !executed_tool_batch.terminate

                for result in tool_results
                    push!(current_context.messages, result)
                    push!(new_messages, result)
                end
            end

            emit(TurnEndEvent(message, tool_results))

            next_turn_context = PrepareNextTurnContext(
                message,
                tool_results,
                current_context,
                new_messages,
            )
            next_turn_snapshot = prepare_next_turn(config, next_turn_context)

            if !isnothing(next_turn_snapshot)
                current_context = next_turn_snapshot.context
                config = AgentLoopConfig(
                    model = next_turn_snapshot.model,
                    reasoning = next_turn_snapshot.thinking_level,
                    convert_to_llm = config.convert_to_llm,
                    transform_context = config.transform_context,
                    get_api_key = config.get_api_key,
                    should_stop_after_turn = config.should_stop_after_turn,
                    prepare_next_turn = config.prepare_next_turn,
                    get_steering_messages = config.get_steering_messages,
                    get_follow_up_messages = config.get_follow_up_messages,
                    tool_execution = config.tool_execution,
                    before_tool_call = config.before_tool_call,
                    after_tool_call = config.after_tool_call,
                    max_tokens = config.max_tokens,
                    temperature = config.temperature,
                    reasoning = config.reasoning,
                    cache_retention = config.cache_retention,
                    session_id = config.session_id,
                    headers = config.headers,
                    metadata = config.metadata,
                    transport = config.transport,
                    signal = signal,
                    api_key = config.api_key,
                    on_payload = config.on_payload,
                    on_response = config.on_response,
                    max_retry_delay_ms = config.max_retry_delay_ms,
                )
            end

            if should_stop_after_turn(config, next_turn_context)
                emit(AgentEndEvent(new_messages))
                return
            end

            pending_messages = getSteeringMessages(config) do
                get_steering_messages(config)
            end
        end

        follow_up_messages = getFollowUpMessages(config) do
            get_follow_up_messages(config)
        end

        if !isempty(follow_up_messages)
            pending_messages = follow_up_messages
            continue
        end

        break
    end

    emit(AgentEndEvent(new_messages))
end

# ============================================================================
# Helper types
# ============================================================================

struct PrepareNextTurnContext
    message::AssistantMessage
    tool_results::Vector{ToolResultMessage}
    context::AgentContext
    new_messages::Vector{AgentMessage}
end

struct AgentLoopTurnUpdate
    context::Union{AgentContext, Nothing}
    model::Union{Model, Nothing}
    thinking_level::Union{ThinkingLevel, Nothing}
end

# ============================================================================
# Helper functions for getting messages from queues
# ============================================================================

macro getSteeringMessages(config)
    :(get_steering_messages($(esc(config))))
end

macro getFollowUpMessages(config)
    :(get_follow_up_messages($(esc(config))))
end

function get_steering_messages(config::AgentLoopConfig)::Vector{AgentMessage}
    return isnothing(config.get_steering_messages) ? AgentMessage[] : config.get_steering_messages()
end

function get_follow_up_messages(config::AgentLoopConfig)::Vector{AgentMessage}
    return isnothing(config.get_follow_up_messages) ? AgentMessage[] : config.get_follow_up_messages()
end

function prepare_next_turn(config::AgentLoopConfig, context::PrepareNextTurnContext)::Union{AgentLoopTurnUpdate, Nothing}
    return isnothing(config.prepare_next_turn) ? nothing : config.prepare_next_turn(context)
end

function should_stop_after_turn(config::AgentLoopConfig, context::PrepareNextTurnContext)::Bool
    return isnothing(config.should_stop_after_turn) ? false : config.should_stop_after_turn(context)
end

# ============================================================================
# Stream assistant response function
# ============================================================================

function streamAssistantResponse(
    context::AgentContext,
    config::AgentLoopConfig,
    signal::Union{Nothing, AbortSignal},
    emit::AgentEventSink,
    stream_function::StreamFn,
)::AssistantMessage
    messages::Vector{AgentMessage} = context.messages

    if !isnothing(config.transform_context)
        messages = config.transform_context(messages, signal)
    end

    llm_messages::Vector{Message} = config.convert_to_llm(messages)

    llm_context::Context = Context(
        context.system_prompt,
        llm_messages,
        context.tools,
    )

    resolved_api_key::Union{String, Nothing} =
        !isnothing(config.get_api_key)
            ? config.get_api_key(config.model.provider)
            : nothing

    response = stream_function(
        config.model,
        llm_context,
        merge(
            config,
            Dict(:apiKey => resolved_api_key, :signal => signal),
        ),
    )

    partial_message::Union{AssistantMessage, Nothing} = nothing
    added_partial::Bool = false

    for event in response
        if event.type == "start"
            partial_message = event.partial
            push!(context.messages, partial_message)
            added_partial = true
            emit(MessageStartEvent(copy(partial_message)))
        elseif event.type in ("text_start", "text_delta", "text_end", "thinking_start", "thinking_delta", "thinking_end", "toolcall_start", "toolcall_delta", "toolcall_end")
            if !isnothing(partial_message)
                partial_message = event.partial
                context.messages[end] = partial_message
                emit(MessageUpdateEvent(copy(partial_message), event))
            end
        elseif event.type in ("done", "error")
            final_message = response.result()
            if added_partial
                context.messages[end] = final_message
            else
                push!(context.messages, final_message)
            end
            if !added_partial
                emit(MessageStartEvent(copy(final_message)))
            end
            emit(MessageEndEvent(final_message))
            return final_message
        end
    end

    final_message = response.result()
    if added_partial
        context.messages[end] = final_message
    else
        push!(context.messages, final_message)
        emit(MessageStartEvent(copy(final_message)))
    end
    emit(MessageEndEvent(final_message))
    return final_message
end

# ============================================================================
# Fail tool calls from truncated message
# ============================================================================

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

function failToolCallsFromTruncatedMessage(
    tool_calls::Vector{ToolCall},
    emit::AgentEventSink,
)::ExecutedToolCallBatch
    messages::Vector{ToolResultMessage} = []

    for tool_call in tool_calls
        emit(ToolExecutionStartEvent(tool_call.id, tool_call.name, tool_call.arguments))

        finalized = FinalizedToolCallOutcome(
            tool_call,
            createErrorToolResult(
                "Tool call \"$(tool_call.name)\" was not executed: the response hit the output token limit, so its arguments may be truncated. Re-issue the tool call with complete arguments.",
            ),
            true,
        )

        emitToolExecutionEnd(finalized, emit)
        tool_result_message = createToolResultMessage(finalized)
        emitToolResultMessage(tool_result_message, emit)
        push!(messages, tool_result_message)
    end

    return ExecutedToolCallBatch(messages, false)
end

# ============================================================================
# Execute tool calls
# ============================================================================

function executeToolCalls(
    current_context::AgentContext,
    assistant_message::AssistantMessage,
    config::AgentLoopConfig,
    signal::Union{Nothing, AbortSignal},
    emit::AgentEventSink,
)::ExecutedToolCallBatch
    tool_calls = filter(
        (c) -> c isa ToolCall,
        assistant_message.content,
    )

    has_sequential_tool_call = any(
        (tc) -> begin
            tool = findfirst((t) -> t.name == tc.name, current_context.tools)
            !isnothing(tool) && tool.execution_mode == EXECUTION_SEQUENTIAL
        end,
        tool_calls,
    )

    if config.tool_execution == EXECUTION_SEQUENTIAL || has_sequential_tool_call
        return executeToolCallsSequential(
            current_context,
            assistant_message,
            tool_calls,
            config,
            signal,
            emit,
        )
    end
    return executeToolCallsParallel(
        current_context,
        assistant_message,
        tool_calls,
        config,
        signal,
        emit,
    )
end

# ============================================================================
# Execute tool calls sequentially
# ============================================================================

function executeToolCallsSequential(
    current_context::AgentContext,
    assistant_message::AssistantMessage,
    tool_calls::Vector{ToolCall},
    config::AgentLoopConfig,
    signal::Union{Nothing, AbortSignal},
    emit::AgentEventSink,
)::ExecutedToolCallBatch
    finalized_calls::Vector{FinalizedToolCallOutcome} = []
    messages::Vector{ToolResultMessage} = []

    for tool_call in tool_calls
        emit(ToolExecutionStartEvent(tool_call.id, tool_call.name, tool_call.arguments))

        preparation = prepareToolCall(current_context, assistant_message, tool_call, config, signal)

        finalized = if preparation.kind == "immediate"
            FinalizedToolCallOutcome(tool_call, preparation.result, preparation.is_error)
        else
            executed = executePreparedToolCall(preparation, signal, emit)
            finalizeExecutedToolCall(
                current_context,
                assistant_message,
                preparation,
                executed,
                config,
                signal,
            )
        end

        emitToolExecutionEnd(finalized, emit)
        tool_result_message = createToolResultMessage(finalized)
        emitToolResultMessage(tool_result_message, emit)
        push!(finalized_calls, finalized)
        push!(messages, tool_result_message)

        if !isnothing(signal) && signal.aborted
            break
        end
    end

    return ExecutedToolCallBatch(messages, shouldTerminateToolBatch(finalized_calls))
end

# ============================================================================
# Execute tool calls in parallel
# ============================================================================

function executeToolCallsParallel(
    current_context::AgentContext,
    assistant_message::AssistantMessage,
    tool_calls::Vector{ToolCall},
    config::AgentLoopConfig,
    signal::Union{Nothing, AbortSignal},
    emit::AgentEventSink,
)::ExecutedToolCallBatch
    finalized_calls::Vector{Union{FinalizedToolCallOutcome, Function}} = []

    for tool_call in tool_calls
        emit(ToolExecutionStartEvent(tool_call.id, tool_call.name, tool_call.arguments))

        preparation = prepareToolCall(current_context, assistant_message, tool_call, config, signal)

        if preparation.kind == "immediate"
            finalized = FinalizedToolCallOutcome(
                tool_call,
                preparation.result,
                preparation.is_error,
            )
            emitToolExecutionEnd(finalized, emit)
            push!(finalized_calls, finalized)
            if !isnothing(signal) && signal.aborted
                break
            end
            continue
        end

        push!(finalized_calls, () -> begin
            executed = executePreparedToolCall(preparation, signal, emit)
            finalized = finalizeExecutedToolCall(
                current_context,
                assistant_message,
                preparation,
                executed,
                config,
                signal,
            )
            emitToolExecutionEnd(finalized, emit)
            return finalized
        end)

        if !isnothing(signal) && signal.aborted
            break
        end
    end

    ordered_finalized_calls = map(
        (entry) -> if entry isa Function
            entry()
        else
            entry
        end,
        finalized_calls,
    )

    messages::Vector{ToolResultMessage} = []
    for finalized in ordered_finalized_calls
        tool_result_message = createToolResultMessage(finalized)
        emitToolResultMessage(tool_result_message, emit)
        push!(messages, tool_result_message)
    end

    return ExecutedToolCallBatch(messages, shouldTerminateToolBatch(ordered_finalized_calls))
end

# ============================================================================
# Prepared tool call types
# ============================================================================

struct PreparedToolCall
    kind::String
    tool_call::ToolCall
    tool::AgentTool
    args::Any
end

struct ImmediateToolCallOutcome
    kind::String
    result::AgentToolResultMutable
    is_error::Bool
end

struct ExecutedToolCallOutcome
    result::AgentToolResultMutable
    is_error::Bool
end

struct FinalizedToolCallOutcome
    tool_call::ToolCall
    result::AgentToolResultMutable
    is_error::Bool
end

# ============================================================================
# Helper functions
# ============================================================================

function shouldTerminateToolBatch(finalized_calls::Vector{FinalizedToolCallOutcome})::Bool
    return !isempty(finalized_calls) && all(
        (finalized) -> finalized.result.terminate === true,
        finalized_calls,
    )
end

function prepareToolCallArguments(tool::AgentTool, tool_call::ToolCall)::ToolCall
    if isnothing(tool.prepare_arguments)
        return tool_call
    end
    prepared_arguments = tool.prepare_arguments(tool_call.arguments)
    if prepared_arguments === tool_call.arguments
        return tool_call
    end
    return ToolCall(
        tool_call.type,
        tool_call.id,
        tool_call.name,
        prepared_arguments,
        tool_call.partial_json,
    )
end

function prepareToolCall(
    current_context::AgentContext,
    assistant_message::AssistantMessage,
    tool_call::ToolCall,
    config::AgentLoopConfig,
    signal::Union{Nothing, AbortSignal},
)::Union{PreparedToolCall, ImmediateToolCallOutcome}
    tool = findfirst((t) -> t.name == tool_call.name, current_context.tools)
    if isnothing(tool)
        return ImmediateToolCallOutcome("immediate", createErrorToolResult("Tool $(tool_call.name) not found"), true)
    end

    try
        prepared_tool_call = prepareToolCallArguments(tool, tool_call)
        validated_args = validateToolArguments(tool, prepared_tool_call)

        if !isnothing(config.before_tool_call)
            before_result = config.before_tool_call(
                BeforeToolCallContext(assistant_message, tool_call, validated_args, current_context),
                signal,
            )
            if !isnothing(signal) && signal.aborted
                return ImmediateToolCallOutcome("immediate", createErrorToolResult("Operation aborted"), true)
            end
            if !isnothing(before_result) && before_result.block
                reason = isnothing(before_result.reason) ? "Tool execution was blocked" : before_result.reason
                return ImmediateToolCallOutcome("immediate", createErrorToolResult(reason), true)
            end
        end

        if !isnothing(signal) && signal.aborted
            return ImmediateToolCallOutcome("immediate", createErrorToolResult("Operation aborted"), true)
        end

        return PreparedToolCall("prepared", tool_call, tool, validated_args)
    catch error
        return ImmediateToolCallOutcome("immediate", createErrorToolResult(string(error)), true)
    end
end

function executePreparedToolCall(
    prepared::PreparedToolCall,
    signal::Union{Nothing, AbortSignal},
    emit::AgentEventSink,
)::ExecutedToolCallOutcome
    update_events::Vector{Future} = []
    accepting_updates::Bool = true

    try
        result = prepared.tool.execute(
            prepared.tool_call.id,
            prepared.args,
            signal,
            (partial_result) -> begin
                if !accepting_updates
                    return
                end
                push!(
                    update_events,
                    Threads.@spawn begin
                        emit(
                            ToolExecutionUpdateEvent(
                                prepared.tool_call.id,
                                prepared.tool_call.name,
                                prepared.tool_call.arguments,
                                partial_result,
                            ),
                        )
                    end,
                )
            end,
        )
        accepting_updates = false
        wait.(update_events)
        return ExecutedToolCallOutcome(result, false)
    catch error
        accepting_updates = false
        wait.(update_events)
        return ExecutedToolCallOutcome(createErrorToolResult(string(error)), true)
    finally
        accepting_updates = false
    end
end

function finalizeExecutedToolCall(
    current_context::AgentContext,
    assistant_message::AssistantMessage,
    prepared::PreparedToolCall,
    executed::ExecutedToolCallOutcome,
    config::AgentLoopConfig,
    signal::Union{Nothing, AbortSignal},
)::FinalizedToolCallOutcome
    result = executed.result
    is_error = executed.is_error

    if !isnothing(config.after_tool_call)
        try
            after_result = config.after_tool_call(
                AfterToolCallContext(
                    assistant_message,
                    prepared.tool_call,
                    prepared.args,
                    result,
                    is_error,
                    current_context,
                ),
                signal,
            )
            if !isnothing(after_result)
                result = AgentToolResultMutable(
                    isnothing(after_result.content) ? result.content : after_result.content,
                    isnothing(after_result.details) ? result.details : after_result.details,
                    isnothing(after_result.usage) ? result.usage : after_result.usage,
                    result.added_tool_names,
                    isnothing(after_result.terminate) ? result.terminate : after_result.terminate,
                )
                is_error = isnothing(after_result.is_error) ? is_error : after_result.is_error
            end
        catch error
            result = createErrorToolResult(string(error))
            is_error = true
        end
    end

    return FinalizedToolCallOutcome(prepared.tool_call, result, is_error)
end

function createErrorToolResult(message::String)::AgentToolResultMutable
    return AgentToolResultMutable([TextContent(message)], Dict{String, Any}(), nothing, nothing, nothing)
end

function emitToolExecutionEnd(finalized::FinalizedToolCallOutcome, emit::AgentEventSink)::Nothing
    emit(ToolExecutionEndEvent(
        finalized.tool_call.id,
        finalized.tool_call.name,
        finalized.result,
        finalized.is_error,
    ))
    return nothing
end

function createToolResultMessage(finalized::FinalizedToolCallOutcome)::ToolResultMessage
    return ToolResultMessage(
        "toolResult",
        finalized.tool_call.id,
        finalized.tool_call.name,
        isnothing(finalized.result.content) ? MessageContent[] : finalized.result.content,
        finalized.result.details,
        finalized.result.usage,
        finalized.result.added_tool_names,
        finalized.is_error,
        Dates.now(Dates.UTC).datetime,
    )
end

function emitToolResultMessage(tool_result_message::ToolResultMessage, emit::AgentEventSink)::Nothing
    emit(MessageStartEvent(tool_result_message))
    emit(MessageEndEvent(tool_result_message))
    return nothing
end

# ============================================================================
# Validation helper
# ============================================================================

function validateToolArguments(tool::AgentTool, tool_call::ToolCall)::Any
    # Simplified validation - in a full implementation, this would use TypeBox-like validation
    return tool_call.arguments
end

end