path: root/src/Agent.Core/Runtime/AgentRuntime.cs

using System.Text.Json;
using Agent.Core.Ollama;
using Agent.Core.Prompts;
using Agent.Core.Sessions;
using Agent.Core.Skills;
using Agent.Core.Tasks;
using Agent.Core.Tools;

namespace Agent.Core.Runtime;

public sealed class AgentRuntime(TaskRegistry registry, IOllamaClient ollama, AgentRuntimeSettings settings, InternalToolDispatcher? toolDispatcher = null, IAgentRuntimeObserver? observer = null)
{
    private readonly InternalToolDispatcher _tools = toolDispatcher ?? new(
        new DenyByDefaultPolicy(InternalToolCatalog.All.Select(static item => item.Name).ToHashSet(StringComparer.Ordinal)),
        new FakeInternalToolBackend());
    private readonly IAgentRuntimeObserver _observer = observer ?? NullAgentRuntimeObserver.Instance;

    public async Task ExecuteAsync(Guid taskId, CancellationToken cancellationToken)
    {
        if (!registry.TryGet(taskId, null, out TaskSnapshot? task) || task is null) return;
        if (!registry.Sessions.TryGet(task.SessionId, out AgentSession? session) || session is null)
        {
            registry.Transition(taskId, AgentTaskState.Failed, result: "Session not found.");
            return;
        }

        try
        {
            await _tools.BeginTaskAsync(taskId, cancellationToken).ConfigureAwait(false);
            if (task.Images.Count > 0 && await ollama.SupportsVisionAsync(cancellationToken).ConfigureAwait(false) is false)
                throw new InvalidOperationException(settings.VisionUnsupportedMessage);

            // Progressive Disclosure Stage 2: LLM Skill Selection.
            IReadOnlyList<AgentSkill> selectedSkills = await SelectSkillsAsync(task.Request, cancellationToken).ConfigureAwait(false);

            var messages = PromptComposer.Compose(new PromptContext(
                FixedPrompt.Load(),
                settings.UserSystemPrompt,
                settings.Locale,
                settings.Skills,
                selectedSkills,  // Stage 3: selected skill instructions injected by PromptComposer.
                settings.SandboxCapabilitySummary,
                "Current task: " + task.Request,
                string.Empty,
                session.Conversation,
                task.Request,
                task.Images)).ToList();
            var guard = new ToolUseGuard(settings.SameFailureLimit);
            bool requiresExternalEffect = RequiresExternalEffect(task.Request);
            bool externalEffectEvidence = false;
            var unresolvedToolFailures = new Dictionary<string, string>(StringComparer.Ordinal);
            IReadOnlyDictionary<string, int> requiredToolCalls = RequiredToolCalls(task.Request);
            var successfulToolCalls = new Dictionary<string, int>(StringComparer.Ordinal);
            string? finalContent = null;
            for (int iteration = 0; iteration < settings.MaxIterations; iteration++)
            {
                _observer.LlmRequest(taskId, iteration, messages);
                OllamaResult result = await ollama.ChatWithToolsAsync(messages, task.Temperature, InternalToolCatalog.All, cancellationToken).ConfigureAwait(false);
                _observer.LlmResponse(taskId, iteration, result);
                registry.AddUsage(taskId, result.InputTokens, result.OutputTokens);
                IReadOnlyList<OllamaToolCall> calls = result.ToolCalls ?? [];
                if (calls.Count > 0)
                {
                    messages.Add(new OllamaMessage("assistant", result.Content, ToolCalls: calls));
                    string? failureFingerprint = null;
                    foreach (OllamaToolCall call in calls)
                    {
                        ToolExecutionResult toolResult;
                        try
                        {
                            toolResult = await _tools.DispatchAsync(taskId, call.Function.Name, call.Function.Arguments, cancellationToken).ConfigureAwait(false);
                        }
                        catch (UnauthorizedAccessException ex)
                        {
                            toolResult = new(false, ex.Message, false, $"{call.Function.Name}:{ex.GetType().Name}:{ex.Message}", RetryAllowed: false);
                        }
                        catch (Exception ex) when (ex is ToolValidationException or ArgumentException)
                        {
                            toolResult = new(false, ex.Message, false, $"{call.Function.Name}:{ex.GetType().Name}:{ex.Message}");
                        }
                        externalEffectEvidence |= toolResult.Success && toolResult.ExternalEffect;
                        if (toolResult.Success) unresolvedToolFailures.Remove(call.Function.Name);
                        else unresolvedToolFailures[call.Function.Name] = toolResult.FailureFingerprint ?? toolResult.Output;
                        if (toolResult.Success) successfulToolCalls[call.Function.Name] = successfulToolCalls.GetValueOrDefault(call.Function.Name) + 1;
                        _observer.ToolResult(taskId, call.Function.Name, toolResult);
                        if (!toolResult.Success && !toolResult.RetryAllowed)
                            throw new InvalidOperationException($"Tool '{call.Function.Name}' was denied by sandbox policy and will not be retried: {toolResult.Output}");
                        failureFingerprint = toolResult.FailureFingerprint ?? failureFingerprint;
                        messages.Add(new OllamaMessage("tool", toolResult.Output, ToolName: call.Function.Name));
                    }
                    ToolUseGuardResult callGuard = guard.Evaluate(requiresExternalEffect, true, externalEffectEvidence, failureFingerprint);
                    if (callGuard.Decision == ToolUseGuardDecision.Fail) throw new InvalidOperationException(callGuard.Observation);
                    continue;
                }

                if (unresolvedToolFailures.Count > 0)
                {
                    messages.Add(new OllamaMessage("system", "Completion is not allowed while tool failures remain unresolved. Retry or correct: " +
                        string.Join("; ", unresolvedToolFailures.Select(static item => $"{item.Key}={item.Value}"))));
                    continue;
                }
                string[] missingCalls = requiredToolCalls
                    .Where(requirement => successfulToolCalls.GetValueOrDefault(requirement.Key) < requirement.Value)
                    .Select(requirement => $"{requirement.Key}: {successfulToolCalls.GetValueOrDefault(requirement.Key)}/{requirement.Value}")
                    .ToArray();
                if (missingCalls.Length > 0)
                {
                    messages.Add(new OllamaMessage("system", "Completion is not allowed because the request explicitly requires successful tool calls that are still missing: " + string.Join("; ", missingCalls)));
                    continue;
                }
                ToolUseGuardResult completion = guard.Evaluate(requiresExternalEffect, false, externalEffectEvidence);
                if (completion.Decision == ToolUseGuardDecision.Retry)
                {
                    messages.Add(new OllamaMessage("system", completion.Observation));
                    continue;
                }
                if (completion.Decision == ToolUseGuardDecision.Fail) throw new InvalidOperationException(completion.Observation);
                finalContent = result.Content;
                break;
            }
            if (finalContent is null) throw new InvalidOperationException($"Agent exceeded the iteration limit of {settings.MaxIterations}.");
            string? cleanupFailure = await EndTaskSafelyAsync(taskId, true, CancellationToken.None).ConfigureAwait(false);
            if (cleanupFailure is not null) throw new InvalidOperationException(cleanupFailure);
            session.Append(new ConversationMessage("user", task.Request));
            session.Append(new ConversationMessage("assistant", finalContent));
            registry.Transition(taskId, AgentTaskState.Completed, result: finalContent);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            string? cleanupFailure = await EndTaskSafelyAsync(taskId, false, CancellationToken.None).ConfigureAwait(false);
            string result = cleanupFailure is null ? "Cancelled." : $"Cancelled. {cleanupFailure}";
            registry.Transition(taskId, AgentTaskState.Cancelled, result: result);
        }
        catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
        {
            string? cleanupFailure = await EndTaskSafelyAsync(taskId, false, CancellationToken.None).ConfigureAwait(false);
            string result = cleanupFailure is null ? ex.Message : $"{ex.Message} {cleanupFailure}";
            registry.Transition(taskId, AgentTaskState.Failed, result: result);
        }
    }

    private async Task<IReadOnlyList<AgentSkill>> SelectSkillsAsync(string request, CancellationToken cancellationToken)
    {
        if (settings.Skills.Count == 0) return [];

        // Build a minimal prompt with skill catalog for the LLM to choose from.
        string catalog = string.Join("\n", settings.Skills.Select(static s => $"- {s.Name}: {s.Description}"));
        var selectionMessages = new List<OllamaMessage>
        {
            new("system", $"Select skills relevant to the user request. Available skills:\n{catalog}\n\nReturn a JSON object with a 'selected_skills' array containing skill names. Select only relevant skills. If none are relevant, return an empty array."),
            new("user", request)
        };

        // JSON Schema for structured output.
        string schema = """{"type":"object","properties":{"selected_skills":{"type":"array","items":{"type":"string"}}},"required":["selected_skills"]}""";
        using JsonDocument formatDoc = JsonDocument.Parse(schema);

        OllamaResult result;
        try
        {
            result = await ollama.ChatStructuredAsync(selectionMessages, 0.0, formatDoc.RootElement.Clone(), cancellationToken).ConfigureAwait(false);
        }
        catch
        {
            // If structured output fails (e.g., model doesn't support it), fall back to no skills.
            return [];
        }

        return ParseSelectedSkills(result.Content);
    }

    private IReadOnlyList<AgentSkill> ParseSelectedSkills(string content)
    {
        try
        {
            using JsonDocument doc = JsonDocument.Parse(content);
            if (!doc.RootElement.TryGetProperty("selected_skills", out JsonElement skillsElement)) return [];
            var selectedNames = new HashSet<string>(StringComparer.Ordinal);
            foreach (JsonElement item in skillsElement.EnumerateArray())
            {
                string? name = item.GetString();
                if (!string.IsNullOrWhiteSpace(name)) selectedNames.Add(name);
            }
            // Validate: only known skills, no duplicates, max 10.
            var valid = new List<AgentSkill>();
            foreach (AgentSkill skill in settings.Skills)
            {
                if (selectedNames.Contains(skill.Name))
                {
                    valid.Add(skill);
                    if (valid.Count >= 10) break;
                }
            }
            return valid;
        }
        catch
        {
            return [];
        }
    }

    private async Task<string?> EndTaskSafelyAsync(Guid taskId, bool success, CancellationToken cancellationToken)
    {
        try
        {
            await _tools.EndTaskAsync(taskId, success, cancellationToken).ConfigureAwait(false);
            return null;
        }
        catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
        {
            return $"Task resource cleanup failed: {ex.Message}";
        }
    }

    private static bool RequiresExternalEffect(string request)
    {
        string normalized = request.ToLowerInvariant();
        string[] indicators = ["edit", "modify", "change", "fix", "write", "delete", "create", "build", "test", "run", "修正", "変更", "編集", "作成", "削除", "ビルド", "テスト", "実行"];
        return indicators.Any(normalized.Contains);
    }

    private static IReadOnlyDictionary<string, int> RequiredToolCalls(string request)
    {
        var result = new Dictionary<string, int>(StringComparer.Ordinal);
        foreach (InternalToolDefinition tool in InternalToolCatalog.All)
        {
            int count = 0;
            int offset = 0;
            while ((offset = request.IndexOf(tool.Name, offset, StringComparison.OrdinalIgnoreCase)) >= 0)
            {
                count++;
                offset += tool.Name.Length;
            }
            if (count > 0) result[tool.Name] = count;
        }
        return result;
    }
}