ClawdDotNet/tests/ClawdDotNet.Core.Tests/Infrastructure/FakeTool.cs

using System.Collections.Concurrent;
using System.Text.Json;
using ClawdDotNet.Core.Tools;

namespace ClawdDotNet.Core.Tests.Infrastructure;

/// <summary>
/// Configurable tool for engine tests: delivers fixed results, can delay
/// (to provoke concurrency) or throw.
/// </summary>
internal sealed class FakeTool : IAgentTool
{
    private readonly Func<JsonElement, AgentToolContext, CancellationToken, Task<ToolResult>> _handler;

    public string Name { get; }
    public string Description => $"Test-Tool {Name}";

    public JsonElement InputSchema { get; } = JsonDocument.Parse("""
        { "type": "object", "properties": { "action": { "type": "string" } } }
        """).RootElement.Clone();

    /// <summary>Every call is recorded — including the calling agent ID.</summary>
    public ConcurrentBag<string> Invocations { get; } = new();

    /// <summary>Whether the tool declares itself fit for simultaneous execution (#551).</summary>
    public bool IsParallelSafe { get; private init; }

    /// <summary>
    /// Highest number of simultaneously running calls this tool has seen.
    ///
    /// That is the only proof that counts: a shorter total runtime could also
    /// come from a fast machine, but a peak of 2 can only
    /// arise if two calls actually overlapped each other.
    /// </summary>
    public int MaxConcurrent => _maxConcurrent;

    private int _current;
    private int _maxConcurrent;

    private FakeTool(string name, Func<JsonElement, AgentToolContext, CancellationToken, Task<ToolResult>> handler)
    {
        Name = name;
        _handler = handler;
    }

    /// <summary>
    /// A slow tool that declares itself fit for concurrency and counts
    /// how many calls overlapped.
    /// </summary>
    public static FakeTool SlowParallelSafe(TimeSpan delay, string result = "ok", string name = "ParallelTool")
        => new(name, async (_, _, ct) =>
        {
            await Task.Delay(delay, ct);
            return ToolResult.Ok(result);
        })
        { IsParallelSafe = true };

    /// <summary>The same, but without declaring itself fit — must stay serial.</summary>
    public static FakeTool SlowSerial(TimeSpan delay, string result = "ok", string name = "SerialTool")
        => new(name, async (_, _, ct) =>
        {
            await Task.Delay(delay, ct);
            return ToolResult.Ok(result);
        });

    public static FakeTool Returning(string result, string name = "TestTool")
        => new(name, (_, _, _) => Task.FromResult(ToolResult.Ok(result)));

    /// <summary>Delays execution — enlarges the time window for races.</summary>
    public static FakeTool Slow(TimeSpan delay, string result = "ok", string name = "TestTool")
        => new(name, async (_, _, ct) =>
        {
            await Task.Delay(delay, ct);
            return ToolResult.Ok(result);
        });

    /// <summary>Fails every time with the same message - for the repetition guard (#581).</summary>
    public static FakeTool Failing(string error, string name = "TestTool")
        => new(name, (_, _, _) => Task.FromResult(ToolResult.Fail(error)));

    /// <summary>#720: Refuses with a write conflict.</summary>
    public static FakeTool Conflicting(string error, string name = "TestTool")
        => new(name, (_, _, _) => Task.FromResult(ToolResult.Conflict(error)));

    /// <summary>Hands the context it was called with to <paramref name="seen"/>.</summary>
    public static FakeTool Observing(Action<AgentToolContext> seen, string name = "TestTool")
        => new(name, (_, context, _) => { seen(context); return Task.FromResult(ToolResult.Ok("ok")); });

    public static FakeTool Throwing(Exception ex, string name = "TestTool")
        => new(name, (_, _, _) => throw ex);

    public async Task<ToolResult> ExecuteAsync(JsonElement input, AgentToolContext context, CancellationToken ct)
    {
        Invocations.Add(context.AgentId);

        var now = Interlocked.Increment(ref _current);

        // Advance the peak without ever lowering it: compare and swap
        // until the stored value is no longer smaller.
        int seen;
        while (now > (seen = Volatile.Read(ref _maxConcurrent)))
            Interlocked.CompareExchange(ref _maxConcurrent, now, seen);

        try
        {
            return await _handler(input, context, ct);
        }
        finally
        {
            Interlocked.Decrement(ref _current);
        }
    }
}

internal sealed class InMemoryStateStore : ClawdDotNet.Core.State.IStateStore
{
    private readonly ConcurrentDictionary<string, string> _values = new();

    public Task<string?> GetAsync(string key, CancellationToken ct)
        => Task.FromResult(_values.GetValueOrDefault(key));

    public Task SetAsync(string key, string value, CancellationToken ct)
    {
        _values[key] = value;
        return Task.CompletedTask;
    }

    public Task DeleteAsync(string key, CancellationToken ct)
    {
        _values.TryRemove(key, out _);
        return Task.CompletedTask;
    }
}