path: root/tests/Agent.Core.Tests/PromptAndRuntimeTests.cs

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

namespace Agent.Core.Tests;

public sealed class PromptAndRuntimeTests
{
    [Fact]
    public void Prompt_order_is_fixed_and_request_is_last()
    {
        IReadOnlyList<OllamaMessage> messages = PromptComposer.Compose(new PromptContext(
            "fixed", "user system", "ja", ["skill"], "capability", "graph", "summary",
            [new ConversationMessage("user", "old"), new ConversationMessage("assistant", "reply")],
            "new request", ["image"]));

        Assert.Equal(["fixed", "user system", "Respond in locale: ja.", "skill", "capability", "graph", "summary", "old", "reply", "new request"], messages.Select(x => x.Content));
        Assert.Equal("system", messages[0].Role);
        Assert.Equal("user", messages[^1].Role);
        Assert.Equal(["image"], messages[^1].Images);
    }

    [Fact]
    public async Task Runtime_continues_session_and_accumulates_usage()
    {
        var sessions = new SessionStore();
        var registry = new TaskRegistry(sessions);
        var ollama = new RecordingOllamaClient();
        var runtime = new AgentRuntime(registry, ollama, new AgentRuntimeSettings("custom", "ja", [], "sandbox"));
        TaskSnapshot first = registry.Enqueue("first", null, 0.5);
        Assert.True(registry.Transition(first.TaskId, AgentTaskState.Planning));
        Assert.True(registry.Transition(first.TaskId, AgentTaskState.Running));

        await runtime.ExecuteAsync(first.TaskId, TestContext.Current.CancellationToken);
        Assert.True(registry.TryGet(first.TaskId, null, out TaskSnapshot? completed));
        Assert.True(completed!.State == AgentTaskState.Completed, completed.Result);
        Assert.Equal(3, completed.InputTokens);
        Assert.Equal(2, completed.OutputTokens);

        TaskSnapshot second = registry.Enqueue("second", first.SessionId, null);
        Assert.True(registry.Transition(second.TaskId, AgentTaskState.Planning));
        Assert.True(registry.Transition(second.TaskId, AgentTaskState.Running));
        await runtime.ExecuteAsync(second.TaskId, TestContext.Current.CancellationToken);

        Assert.Equal(2, ollama.Calls.Count);
        Assert.Contains(ollama.Calls[1], message => message.Role == "user" && message.Content == "first");
        Assert.Contains(ollama.Calls[1], message => message.Role == "assistant" && message.Content == "answer-1");
        Assert.Equal(0.5, ollama.Temperatures[0]);
        Assert.Null(ollama.Temperatures[1]);
    }

    [Fact]
    public void Session_store_is_process_local()
    {
        SessionStore first = new();
        Guid id = first.Create().SessionId;
        SessionStore restarted = new();

        Assert.False(restarted.TryGet(id, out _));
    }

    [Fact]
    public async Task Non_vision_model_rejects_images_with_selected_locale_message_before_chat()
    {
        TaskRegistry registry = new();
        var ollama = new NonVisionOllamaClient();
        const string localized = "選択したモデルは画像入力に対応していません。";
        var runtime = new AgentRuntime(registry, ollama,
            new AgentRuntimeSettings("", "ja", [], "sandbox", VisionUnsupportedMessage: localized));
        TaskSnapshot task = registry.Enqueue("画像を確認", null, null, ["base64-image"]);
        registry.Transition(task.TaskId, AgentTaskState.Planning);
        registry.Transition(task.TaskId, AgentTaskState.Running);

        await runtime.ExecuteAsync(task.TaskId, TestContext.Current.CancellationToken);

        Assert.True(registry.TryGet(task.TaskId, null, out TaskSnapshot? failed));
        Assert.Equal(AgentTaskState.Failed, failed!.State);
        Assert.Equal(localized, failed.Result);
        Assert.Equal(0, ollama.ChatCalls);
    }

    [Fact]
    public async Task Runtime_dispatches_native_tool_call_then_completes_only_after_effect_evidence()
    {
        TaskRegistry registry = new();
        FakeInternalToolBackend backend = new();
        var dispatcher = new InternalToolDispatcher(
            new DenyByDefaultPolicy(InternalToolCatalog.All.Select(static item => item.Name).ToHashSet(StringComparer.Ordinal)), backend);
        var ollama = new ToolCallingOllamaClient();
        var runtime = new AgentRuntime(registry, ollama, new AgentRuntimeSettings("", "ja", [], "sandbox"), dispatcher);
        TaskSnapshot task = registry.Enqueue("fileを修正", null, null);
        registry.Transition(task.TaskId, AgentTaskState.Planning);
        registry.Transition(task.TaskId, AgentTaskState.Running);

        await runtime.ExecuteAsync(task.TaskId, TestContext.Current.CancellationToken);

        Assert.True(registry.TryGet(task.TaskId, null, out TaskSnapshot? completed));
        Assert.True(completed!.State == AgentTaskState.Completed, $"{completed.Result}; backend={backend.Calls.Count}:{string.Join(',', backend.Calls.Select(x => x.Name))}; ollama={ollama.CallCount}");
        Assert.Equal("修正しました。", completed.Result);
        Assert.Single(backend.Calls);
        Assert.Equal("file_edit", backend.Calls[0].Name);
        Assert.Equal(2, ollama.CallCount);
        Assert.Equal(6, ollama.ToolCount);
    }

    [Fact]
    public async Task Runtime_reaches_terminal_failed_state_when_resource_cleanup_fails()
    {
        TaskRegistry registry = new();
        var backend = new CleanupFailingBackend();
        var dispatcher = new InternalToolDispatcher(new DenyByDefaultPolicy(), backend);
        var runtime = new AgentRuntime(registry, new RecordingOllamaClient(), new AgentRuntimeSettings("", "en", [], "sandbox"), dispatcher);
        TaskSnapshot task = registry.Enqueue("analyze", null, null);
        registry.Transition(task.TaskId, AgentTaskState.Planning);
        registry.Transition(task.TaskId, AgentTaskState.Running);

        await runtime.ExecuteAsync(task.TaskId, TestContext.Current.CancellationToken);

        Assert.True(registry.TryGet(task.TaskId, null, out TaskSnapshot? failed));
        Assert.Equal(AgentTaskState.Failed, failed!.State);
        Assert.Contains("cleanup failed", failed.Result, StringComparison.OrdinalIgnoreCase);
        Assert.Equal(2, backend.EndCalls);
    }

    [Fact]
    public async Task Runtime_cannot_complete_with_an_unresolved_tool_failure()
    {
        TaskRegistry registry = new();
        var dispatcher = new InternalToolDispatcher(
            new DenyByDefaultPolicy(new HashSet<string>(["process_run"], StringComparer.Ordinal)), new FailingProcessBackend());
        var runtime = new AgentRuntime(registry, new FailingThenFinalOllamaClient(), new AgentRuntimeSettings("", "en", [], "sandbox", MaxIterations: 2), dispatcher);
        TaskSnapshot task = registry.Enqueue("run build", null, null);
        registry.Transition(task.TaskId, AgentTaskState.Planning);
        registry.Transition(task.TaskId, AgentTaskState.Running);

        await runtime.ExecuteAsync(task.TaskId, TestContext.Current.CancellationToken);

        Assert.True(registry.TryGet(task.TaskId, null, out TaskSnapshot? failed));
        Assert.Equal(AgentTaskState.Failed, failed!.State);
        Assert.Contains("iteration limit", failed.Result, StringComparison.OrdinalIgnoreCase);
    }

    [Fact]
    public async Task Runtime_does_not_retry_a_sandbox_policy_denial()
    {
        TaskRegistry registry = new();
        var backend = new UnauthorizedBackend();
        var dispatcher = new InternalToolDispatcher(
            new DenyByDefaultPolicy(new HashSet<string>(["file_search"], StringComparer.Ordinal)), backend);
        var ollama = new FileSearchOllamaClient();
        var runtime = new AgentRuntime(registry, ollama, new AgentRuntimeSettings("", "en", [], "sandbox"), dispatcher);
        TaskSnapshot task = registry.Enqueue("research regional curry", null, null);
        registry.Transition(task.TaskId, AgentTaskState.Planning);
        registry.Transition(task.TaskId, AgentTaskState.Running);

        await runtime.ExecuteAsync(task.TaskId, TestContext.Current.CancellationToken);

        Assert.True(registry.TryGet(task.TaskId, null, out TaskSnapshot? failed));
        Assert.Equal(AgentTaskState.Failed, failed!.State);
        Assert.Contains("denied by sandbox policy", failed.Result, StringComparison.OrdinalIgnoreCase);
        Assert.Contains("not approved for Read", failed.Result, StringComparison.OrdinalIgnoreCase);
        Assert.Equal(1, ollama.CallCount);
        Assert.Equal(1, backend.CallCount);
    }

    private sealed class RecordingOllamaClient : IOllamaClient
    {
        public List<IReadOnlyList<OllamaMessage>> Calls { get; } = [];
        public List<double?> Temperatures { get; } = [];

        public Task<OllamaResult> ChatAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, CancellationToken cancellationToken)
        {
            Calls.Add(messages);
            Temperatures.Add(temperatureOverride);
            return Task.FromResult(new OllamaResult($"answer-{Calls.Count}", 3, 2));
        }
    }

    private sealed class ToolCallingOllamaClient : IOllamaClient
    {
        public int CallCount { get; private set; }
        public int ToolCount { get; private set; }
        public Task<OllamaResult> ChatAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, CancellationToken cancellationToken) =>
            Task.FromResult(new OllamaResult("修正しました。", 1, 1));
        public Task<OllamaResult> ChatWithToolsAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, IReadOnlyList<InternalToolDefinition> tools, CancellationToken cancellationToken)
        {
            CallCount++;
            ToolCount = tools.Count;
            if (CallCount == 1)
            {
                using JsonDocument document = JsonDocument.Parse("""{"path":"x","operation":"overwrite","content":"new"}""");
                return Task.FromResult(new OllamaResult("", 2, 1, [new OllamaToolCall(new OllamaFunctionCall("file_edit", document.RootElement.Clone()))]));
            }
            return Task.FromResult(new OllamaResult("修正しました。", 2, 2));
        }
    }

    private sealed class NonVisionOllamaClient : IOllamaClient
    {
        public int ChatCalls { get; private set; }
        public Task<bool?> SupportsVisionAsync(CancellationToken cancellationToken) => Task.FromResult<bool?>(false);
        public Task<OllamaResult> ChatAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, CancellationToken cancellationToken)
        {
            ChatCalls++;
            return Task.FromResult(new OllamaResult("unexpected", 0, 0));
        }
    }

    private sealed class CleanupFailingBackend : IInternalToolBackend
    {
        public int EndCalls { get; private set; }
        public Task<ToolExecutionResult> ExecuteAsync(Guid taskId, string toolName, JsonElement arguments, CancellationToken cancellationToken) =>
            Task.FromResult(new ToolExecutionResult(true, "unused", false));
        public Task EndTaskAsync(Guid taskId, bool success, CancellationToken cancellationToken)
        {
            EndCalls++;
            throw new IOException("cleanup failed");
        }
    }

    private sealed class FailingProcessBackend : IInternalToolBackend
    {
        public Task<ToolExecutionResult> ExecuteAsync(Guid taskId, string toolName, JsonElement arguments, CancellationToken cancellationToken) =>
            Task.FromResult(new ToolExecutionResult(false, "exit 1", true, "process:1"));
    }

    private sealed class FailingThenFinalOllamaClient : IOllamaClient
    {
        private int _calls;
        public Task<OllamaResult> ChatAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, CancellationToken cancellationToken) =>
            Task.FromResult(new OllamaResult("done", 1, 1));
        public Task<OllamaResult> ChatWithToolsAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, IReadOnlyList<InternalToolDefinition> tools, CancellationToken cancellationToken)
        {
            if (_calls++ > 0) return Task.FromResult(new OllamaResult("done", 1, 1));
            using JsonDocument arguments = JsonDocument.Parse("""{"executable":"x","args":[],"cwd":"."}""");
            return Task.FromResult(new OllamaResult("", 1, 1, [new OllamaToolCall(new OllamaFunctionCall("process_run", arguments.RootElement.Clone()))]));
        }
    }

    private sealed class UnauthorizedBackend : IInternalToolBackend
    {
        public int CallCount { get; private set; }

        public Task<ToolExecutionResult> ExecuteAsync(Guid taskId, string toolName, JsonElement arguments, CancellationToken cancellationToken)
        {
            CallCount++;
            throw new UnauthorizedAccessException("'C:\\Users' is not approved for Read.");
        }
    }

    private sealed class FileSearchOllamaClient : IOllamaClient
    {
        public int CallCount { get; private set; }

        public Task<OllamaResult> ChatAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, CancellationToken cancellationToken) =>
            Task.FromResult(new OllamaResult("unexpected", 1, 1));

        public Task<OllamaResult> ChatWithToolsAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, IReadOnlyList<InternalToolDefinition> tools, CancellationToken cancellationToken)
        {
            CallCount++;
            using JsonDocument arguments = JsonDocument.Parse("""{"root":"C:\\Users","mode":"name","query":"*.txt"}""");
            return Task.FromResult(new OllamaResult("", 1, 1, [new OllamaToolCall(new OllamaFunctionCall("file_search", arguments.RootElement.Clone()))]));
        }
    }
}