AX-Copilot-Codex/src/AxCopilot/Services/Agent/StreamingToolExecutionCoordinator.cs

using System.Collections.Generic;
using System.Diagnostics;
using System.Text;
using System.Text.Json;
using AxCopilot.Models;
using AxCopilot.Services;

namespace AxCopilot.Services.Agent;

internal sealed class StreamingToolExecutionCoordinator : IToolExecutionCoordinator
{
    private static readonly HashSet<string> PrefetchableReadOnlyTools = new(StringComparer.OrdinalIgnoreCase)
    {
        "file_read", "document_read",
        "env_tool", "datetime_tool",
        "dev_env_detect", "memory", "json_tool", "regex_tool", "base64_tool",
        "hash_tool", "image_analyze"
    };

    private readonly ILlmService _llm;
    private readonly Func<string, IReadOnlyCollection<string>, string> _resolveRequestedToolName;
    private readonly Func<string, JsonElement, AgentContext, List<ChatMessage>?, CancellationToken, Task<ToolResult>> _executeToolAsync;
    private readonly Action<AgentEventType, string, string> _emitEvent;
    private readonly Func<string?, bool> _isContextOverflowError;
    private readonly Func<List<ChatMessage>, bool> _forceContextRecovery;
    private readonly Func<Exception, bool> _isTransientLlmError;
    private readonly Func<int, Exception, int> _computeTransientBackoffDelayMs;

    public StreamingToolExecutionCoordinator(
        ILlmService llm,
        Func<string, IReadOnlyCollection<string>, string> resolveRequestedToolName,
        Func<string, JsonElement, AgentContext, List<ChatMessage>?, CancellationToken, Task<ToolResult>> executeToolAsync,
        Action<AgentEventType, string, string> emitEvent,
        Func<string?, bool> isContextOverflowError,
        Func<List<ChatMessage>, bool> forceContextRecovery,
        Func<Exception, bool> isTransientLlmError,
        Func<int, Exception, int> computeTransientBackoffDelayMs)
    {
        _llm = llm;
        _resolveRequestedToolName = resolveRequestedToolName;
        _executeToolAsync = executeToolAsync;
        _emitEvent = emitEvent;
        _isContextOverflowError = isContextOverflowError;
        _forceContextRecovery = forceContextRecovery;
        _isTransientLlmError = isTransientLlmError;
        _computeTransientBackoffDelayMs = computeTransientBackoffDelayMs;
    }

    public async Task<ToolPrefetchResult?> TryPrefetchReadOnlyToolAsync(
        ContentBlock block,
        IReadOnlyCollection<IAgentTool> tools,
        AgentContext context,
        CancellationToken ct)
    {
        var activeToolNames = tools.Select(t => t.Name).Distinct(StringComparer.OrdinalIgnoreCase).ToList();
        var resolvedToolName = _resolveRequestedToolName(block.ToolName, activeToolNames);
        block.ResolvedToolName = resolvedToolName;

        if (!PrefetchableReadOnlyTools.Contains(resolvedToolName))
            return null;

        _emitEvent(
            AgentEventType.Thinking,
            resolvedToolName,
            $"읽기 도구 조기 실행 준비: {resolvedToolName}");

        var sw = Stopwatch.StartNew();
        try
        {
            var input = block.ToolInput ?? JsonDocument.Parse("{}").RootElement;
            var result = await _executeToolAsync(resolvedToolName, input, context, null, ct);
            sw.Stop();
            return new ToolPrefetchResult(result, sw.ElapsedMilliseconds, resolvedToolName);
        }
        catch (Exception ex)
        {
            sw.Stop();
            return new ToolPrefetchResult(
                ToolResult.Fail($"조기 실행 오류: {ex.Message}"),
                sw.ElapsedMilliseconds,
                resolvedToolName);
        }
    }

    public async Task<List<ContentBlock>> SendWithToolsWithRecoveryAsync(
        List<ChatMessage> messages,
        IReadOnlyCollection<IAgentTool> tools,
        CancellationToken ct,
        string phaseLabel,
        AgentLoopService.RunState? runState = null,
        bool forceToolCall = false,
        Func<ContentBlock, Task<ToolPrefetchResult?>>? prefetchToolCallAsync = null,
        Func<ToolStreamEvent, Task>? onStreamEventAsync = null)
    {
        var transientRetries = runState?.TransientLlmErrorRetries ?? 0;
        var contextRecoveryRetries = runState?.ContextRecoveryAttempts ?? 0;

        while (true)
        {
            var streamedAnyPartialState = false;
            try
            {
                if (onStreamEventAsync == null)
                    return await _llm.SendWithToolsAsync(messages, tools, ct, forceToolCall, prefetchToolCallAsync);

                var blocks = new List<ContentBlock>();
                var textBuilder = new StringBuilder();
                await foreach (var evt in _llm.StreamWithToolsAsync(messages, tools, forceToolCall, prefetchToolCallAsync, ct).WithCancellation(ct))
                {
                    await onStreamEventAsync(evt);
                    if (evt.Kind == ToolStreamEventKind.TextDelta && !string.IsNullOrWhiteSpace(evt.Text))
                    {
                        streamedAnyPartialState = true;
                        textBuilder.Append(evt.Text);
                    }
                    else if (evt.Kind == ToolStreamEventKind.ToolCallReady && evt.ToolCall != null)
                    {
                        streamedAnyPartialState = true;
                        blocks.Add(evt.ToolCall);
                    }
                }

                var result = new List<ContentBlock>();
                var text = textBuilder.ToString().Trim();
                if (!string.IsNullOrWhiteSpace(text))
                    result.Add(new ContentBlock { Type = "text", Text = text });
                result.AddRange(blocks);
                return result;
            }
            catch (Exception ex)
            {
                if (_isContextOverflowError(ex.Message)
                    && contextRecoveryRetries < 2
                    && _forceContextRecovery(messages))
                {
                    if (onStreamEventAsync != null && streamedAnyPartialState)
                        await onStreamEventAsync(new ToolStreamEvent(ToolStreamEventKind.RetryReset, $"{phaseLabel}:retry"));
                    contextRecoveryRetries++;
                    if (runState != null)
                        runState.ContextRecoveryAttempts = contextRecoveryRetries;

                    _emitEvent(
                        AgentEventType.Thinking,
                        "",
                        $"{phaseLabel}: 컨텍스트 한도 초과로 대화를 압축한 후 재시도합니다 ({contextRecoveryRetries}/2)");
                    continue;
                }

                if (ct.IsCancellationRequested)
                    throw;

                if (_isTransientLlmError(ex) && transientRetries < 3)
                {
                    if (onStreamEventAsync != null && streamedAnyPartialState)
                        await onStreamEventAsync(new ToolStreamEvent(ToolStreamEventKind.RetryReset, $"{phaseLabel}:retry"));
                    transientRetries++;
                    if (runState != null)
                        runState.TransientLlmErrorRetries = transientRetries;

                    var delayMs = _computeTransientBackoffDelayMs(transientRetries, ex);
                    _emitEvent(
                        AgentEventType.Thinking,
                        "",
                        $"{phaseLabel}: 일시적 LLM 오류로 {delayMs}ms 후 재시도합니다 ({transientRetries}/3)");
                    await Task.Delay(delayMs, ct);
                    continue;
                }

                throw;
            }
        }
    }
}