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

using AxCopilot.Models;

namespace AxCopilot.Services.Agent;

public sealed class AgentQueryContextWindowResult
{
    public required List<ChatMessage> Messages { get; init; }
    public int SourceMessageCount { get; init; }
    public int ViewMessageCount { get; init; }
    public int WindowStartIndex { get; init; }
    public bool BoundaryApplied { get; init; }
    public bool ToolPairExpanded { get; init; }
    public int PreservedToolPairCount { get; init; }
    public int TruncatedToolResultCount { get; init; }
    public int ReusedToolResultPreviewCount { get; init; }
    public int TokensBeforeBudget { get; init; }
    public int TokensAfterBudget { get; init; }
}

/// <summary>
/// claude-code의 messagesForQuery처럼, 저장된 전체 대화와 실제 LLM에 보내는 query view를 분리합니다.
/// </summary>
public static class AgentQueryContextBuilder
{
    private const int ProtectedRecentNonSystemMessages = 8;
    private const string PostCompactContextMetaKind = "post_compact_context";

    public static AgentQueryContextWindowResult Build(IReadOnlyList<ChatMessage> sourceMessages)
    {
        if (sourceMessages is IList<ChatMessage> mutableSourceMessages)
            AgentMessageInvariantHelper.PopulateMissingToolResultPreviews(mutableSourceMessages);

        if (sourceMessages.Count == 0)
        {
            return new AgentQueryContextWindowResult
            {
                Messages = new List<ChatMessage>(),
                SourceMessageCount = 0,
                ViewMessageCount = 0,
                WindowStartIndex = 0,
                BoundaryApplied = false,
                ToolPairExpanded = false,
                PreservedToolPairCount = 0,
                TruncatedToolResultCount = 0,
                ReusedToolResultPreviewCount = 0,
                TokensBeforeBudget = 0,
                TokensAfterBudget = 0,
            };
        }

        var startIndex = FindWindowStartIndex(sourceMessages, out var boundaryApplied);
        var adjustedStartIndex = AgentMessageInvariantHelper.AdjustStartIndexForToolPairs(
            sourceMessages,
            startIndex,
            out var preservedToolPairs);
        var toolPairExpanded = adjustedStartIndex < startIndex;

        var windowMessages = new List<ChatMessage>(sourceMessages.Count);
        for (var i = 0; i < sourceMessages.Count; i++)
        {
            var include = string.Equals(sourceMessages[i].Role, "system", StringComparison.OrdinalIgnoreCase)
                || i >= adjustedStartIndex;
            if (!include)
                continue;

            windowMessages.Add(CloneMessage(sourceMessages[i]));
        }

        if (boundaryApplied)
            InjectPostCompactContextMessage(windowMessages);

        var tokensBeforeBudget = TokenEstimator.EstimateMessages(windowMessages);
        var budgetResult = AgentToolResultBudget.Apply(windowMessages, ProtectedRecentNonSystemMessages, sourceMessages: sourceMessages);
        var tokensAfterBudget = TokenEstimator.EstimateMessages(windowMessages);

        return new AgentQueryContextWindowResult
        {
            Messages = windowMessages,
            SourceMessageCount = sourceMessages.Count,
            ViewMessageCount = windowMessages.Count,
            WindowStartIndex = adjustedStartIndex,
            BoundaryApplied = boundaryApplied,
            ToolPairExpanded = toolPairExpanded,
            PreservedToolPairCount = preservedToolPairs,
            TruncatedToolResultCount = budgetResult.TruncatedCount,
            ReusedToolResultPreviewCount = budgetResult.ReusedPreviewCount,
            TokensBeforeBudget = tokensBeforeBudget,
            TokensAfterBudget = tokensAfterBudget,
        };
    }

    private static int FindWindowStartIndex(IReadOnlyList<ChatMessage> sourceMessages, out bool boundaryApplied)
    {
        for (var i = sourceMessages.Count - 1; i >= 0; i--)
        {
            if (IsQueryBoundaryMarker(sourceMessages[i]))
            {
                boundaryApplied = true;
                return i;
            }
        }

        boundaryApplied = false;
        return 0;
    }

    private static bool IsQueryBoundaryMarker(ChatMessage message)
    {
        var metaKind = message.MetaKind ?? "";
        if (metaKind.Equals("microcompact_boundary", StringComparison.OrdinalIgnoreCase)
            || metaKind.Equals("session_memory_compaction", StringComparison.OrdinalIgnoreCase)
            || metaKind.Equals("collapsed_boundary", StringComparison.OrdinalIgnoreCase))
        {
            return true;
        }

        var content = message.Content ?? "";
        return content.StartsWith("[이전 대화 요약", StringComparison.Ordinal)
               || content.StartsWith("[세션 메모리 압축", StringComparison.Ordinal)
               || content.StartsWith("[이전 실행 묶음 압축", StringComparison.Ordinal)
               || content.StartsWith("[이전 압축 경계 병합", StringComparison.Ordinal);
    }

    private static ChatMessage CloneMessage(ChatMessage source, string? contentOverride = null)
    {
        return new ChatMessage
        {
            MsgId = source.MsgId,
            Role = source.Role,
            Content = contentOverride ?? source.Content,
            Timestamp = source.Timestamp,
            MetaKind = source.MetaKind,
            MetaRunId = source.MetaRunId,
            Feedback = source.Feedback,
            ResponseElapsedMs = source.ResponseElapsedMs,
            PromptTokens = source.PromptTokens,
            CompletionTokens = source.CompletionTokens,
            AttachedFiles = source.AttachedFiles?.ToList(),
            QueryPreviewContent = source.QueryPreviewContent,
            Images = source.Images?.Select(image => new ImageAttachment
            {
                Base64 = image.Base64,
                MimeType = image.MimeType,
                FileName = image.FileName,
            }).ToList(),
        };
    }

    private static void InjectPostCompactContextMessage(List<ChatMessage> messages)
    {
        if (messages.Count == 0)
            return;

        if (messages.Any(m => string.Equals(m.MetaKind, PostCompactContextMetaKind, StringComparison.OrdinalIgnoreCase)))
            return;

        var attachedFiles = messages
            .SelectMany(m => m.AttachedFiles ?? Enumerable.Empty<string>())
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .Take(5)
            .ToList();
        var imageCount = messages.Sum(m => m.Images?.Count ?? 0);
        var compactSummaryCount = messages.Count(m =>
            string.Equals(m.MetaKind, "microcompact_boundary", StringComparison.OrdinalIgnoreCase)
            || string.Equals(m.MetaKind, "session_memory_compaction", StringComparison.OrdinalIgnoreCase)
            || string.Equals(m.MetaKind, "collapsed_boundary", StringComparison.OrdinalIgnoreCase));
        var branchContextCount = messages.Count(m =>
            string.Equals(m.MetaKind, "branch_context", StringComparison.OrdinalIgnoreCase));
        var structuredToolHistoryCount = messages.Count(m =>
        {
            var content = m.Content ?? "";
            return content.StartsWith("{\"_tool_use_blocks\"", StringComparison.Ordinal)
                   || content.StartsWith("{\"type\":\"tool_result\"", StringComparison.Ordinal);
        });
        var recentToolNames = messages
            .Where(m => m.Role == "user")
            .Select(m => TryExtractToolResultToolName(m, out var toolName) ? toolName : "")
            .Where(toolName => !string.IsNullOrWhiteSpace(toolName))
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .Take(4)
            .ToList();

        if (attachedFiles.Count == 0
            && imageCount == 0
            && compactSummaryCount == 0
            && structuredToolHistoryCount == 0
            && branchContextCount == 0
            && recentToolNames.Count == 0)
            return;

        var lines = new List<string> { "[post-compact context]" };
        if (compactSummaryCount > 0)
            lines.Add($"restored compact summaries: {compactSummaryCount}");
        if (branchContextCount > 0)
            lines.Add($"restored branch context: {branchContextCount}");
        if (structuredToolHistoryCount > 0)
            lines.Add($"restored tool history blocks: {structuredToolHistoryCount}");
        if (recentToolNames.Count > 0)
            lines.Add("restored recent tools: " + string.Join(", ", recentToolNames));
        if (attachedFiles.Count > 0)
            lines.Add("restored file refs: " + string.Join(", ", attachedFiles));
        if (imageCount > 0)
            lines.Add($"restored image refs: {imageCount}");

        var insertIndex = 0;
        while (insertIndex < messages.Count && string.Equals(messages[insertIndex].Role, "system", StringComparison.OrdinalIgnoreCase))
            insertIndex++;

        messages.Insert(insertIndex, new ChatMessage
        {
            Role = "system",
            MetaKind = PostCompactContextMetaKind,
            Content = string.Join("\n", lines),
            Timestamp = DateTime.Now,
            AttachedFiles = attachedFiles.Count > 0 ? attachedFiles : null,
        });
    }

    private static bool TryExtractToolResultToolName(ChatMessage message, out string toolName)
    {
        toolName = "";
        if (!string.Equals(message.Role, "user", StringComparison.OrdinalIgnoreCase)
            || string.IsNullOrWhiteSpace(message.Content)
            || !message.Content.StartsWith("{\"type\":\"tool_result\"", StringComparison.Ordinal))
            return false;

        try
        {
            using var doc = System.Text.Json.JsonDocument.Parse(message.Content);
            if (doc.RootElement.TryGetProperty("tool_name", out var toolNameEl))
            {
                toolName = toolNameEl.GetString() ?? "";
                return !string.IsNullOrWhiteSpace(toolName);
            }
        }
        catch
        {
        }

        return false;
    }
}