fix(dotnet): coordinate mixed client and server approvals - #2450

using System.Runtime.CompilerServices;
using System.Text.Json;
using AGUI.Abstractions;
using Microsoft.Extensions.AI;
using Xunit;

namespace AGUI.Server.UnitTests;

/// <summary>
/// Regression coverage for ag-ui#3271: on a continuation, a client tool that already produced a
/// result must stay declared exactly as the client declared it. Resolving one call must not
/// change the schema or the execution owner of later calls to the same tool. Each new call has
/// to reach the client with its own arguments, and must never be answered on the server with
/// the earlier call's result.
/// </summary>
public sealed class ContinuationClientToolDeclarationTest
{
    private static readonly JsonSerializerOptions SerializerOptions = AIJsonUtilities.DefaultOptions;

    private const string ThreadId = "thread-1";
    private const string RunId = "run-1";

    private static readonly JsonElement ColourSchema = JsonDocument.Parse(
        """{"type":"object","properties":{"colour":{"type":"string"}},"required":["colour"]}""")
        .RootElement.Clone();

    [Fact]
    public void Continuation_ClientToolWithEarlierResult_KeepsItsParameterSchema()
    {
        var context = BuildContinuationInput().ToChatRequestContext(SerializerOptions);

        var declaration = Assert.IsAssignableFrom<AIFunctionDeclaration>(Assert.Single(context.ChatOptions.Tools!));
        Assert.Equal("set_colour", declaration.Name);
        Assert.Equal("Change the UI colour.", declaration.Description);
        Assert.True(declaration.JsonSchema.GetProperty("properties").TryGetProperty("colour", out _));
        Assert.Equal("colour", declaration.JsonSchema.GetProperty("required")[0].GetString());
    }

    /// <summary>
    /// Whatever the adapter installs for the client tool, invoking it for a call other than the
    /// resolved one must not produce the earlier call's result: a new call belongs to the client.
    /// </summary>
    [Fact]
    public async Task Continuation_ClientToolWithEarlierResult_DoesNotAnswerNewCallsWithTheEarlierResult()
    {
        var tool = Assert.Single(BuildContinuationInput().ToChatRequestContext(SerializerOptions).ChatOptions.Tools!);
        if (tool is not AIFunction function)
        {
            return; // Declaration only: nothing on the server can answer the call.
        }

        object? result;
        try
        {
            result = await function.InvokeAsync(new AIFunctionArguments(
                new Dictionary<string, object?> { ["colour"] = "blue" }));
        }
        catch (InvalidOperationException)
        {
            return; // Refusing a call it has no result for is also correct.
        }

        Assert.NotEqual("Applied red", result?.ToString());
    }

    /// <summary>
    /// The reported scenario, driven through FunctionInvokingChatClient: after "Applied red", the
    /// model asks for blue. The model must see the colour parameter, and the new call must reach
    /// the client as a plain TOOL_CALL carrying its own arguments, with no server-side result.
    /// </summary>
    [Fact]
    public async Task Continuation_NewCallToResolvedClientTool_ReachesTheClientWithItsArguments()
    {
        var llm = new ScriptedChatClient(
            new FunctionCallContent("call-new", "set_colour", new Dictionary<string, object?> { ["colour"] = "blue" }));

        var events = await RunTurn(llm, serverTools: [], BuildContinuationInput());

        var advertised = Assert.IsAssignableFrom<AIFunctionDeclaration>(Assert.Single(llm.ToolsSeen[0]));
        Assert.True(advertised.JsonSchema.GetProperty("properties").TryGetProperty("colour", out _));

        AssertSurfacedToClient(events, "call-new", "blue");
        Assert.IsType<RunFinishedSuccessOutcome>(events.OfType<RunFinishedEvent>().Single().Outcome);
        Assert.Equal(1, llm.Calls);
    }

    /// <summary>
    /// Several UI actions in one request: the continuation resolves the earlier client call and
    /// runs a pending server tool, then the model calls the client tool again in the same run.
    /// That second call must still go to the client, not be answered with "Applied red".
    /// </summary>
    [Fact]
    public async Task MixedContinuation_ClientToolCalledAgainAfterServerTool_ReachesTheClient()
    {
        var serverToolInvoked = false;
        var serverTools = new List<AITool>
        {
            AIFunctionFactory.Create(() => { serverToolInvoked = true; return "table shown"; }, "show_table"),
        };

        var llm = new ScriptedChatClient(
            new FunctionCallContent("call-dock", "set_colour", new Dictionary<string, object?> { ["colour"] = "blue" }));

        var input = BuildContinuationInput();
        var assistant = (AGUIAssistantMessage)input.Messages[1];
        assistant.ToolCalls!.Add(new AGUIToolCall
        {
            Id = "call-server",
            Function = new AGUIToolCallFunction { Name = "show_table", Arguments = "{}" }
        });
        input.Messages.RemoveAt(input.Messages.Count - 1);

        var events = await RunTurn(llm, serverTools, input);

        Assert.True(serverToolInvoked, "the pending server tool should run on the continuation");
        Assert.Contains(events.OfType<ToolCallResultEvent>(), e => e.ToolCallId == "call-server");
        AssertSurfacedToClient(events, "call-dock", "blue");
        Assert.IsType<RunFinishedSuccessOutcome>(events.OfType<RunFinishedEvent>().Single().Outcome);
    }

    private static void AssertSurfacedToClient(List<BaseEvent> events, string callId, string colour)
    {
        var start = Assert.Single(events.OfType<ToolCallStartEvent>(), e => e.ToolCallId == callId);
        Assert.Equal("set_colour", start.ToolCallName);

        var args = Assert.Single(events.OfType<ToolCallArgsEvent>(), e => e.ToolCallId == callId);
        using var parsed = JsonDocument.Parse(args.Delta);
        Assert.Equal(colour, parsed.RootElement.GetProperty("colour").GetString());

        Assert.DoesNotContain(events.OfType<ToolCallResultEvent>(), e => e.ToolCallId == callId);
    }

    private static RunAgentInput BuildFirstTurnInput() =>
        new()
        {
            ThreadId = ThreadId,
            RunId = RunId,
            Tools = [new AGUITool { Name = "set_colour", Description = "Change the UI colour.", Parameters = ColourSchema }],
            Messages = [new AGUIUserMessage { Content = "Make it red." }]
        };

    /// <summary>
    /// The continuation from the issue: an earlier set_colour call answered by the client with
    /// "Applied red", then a new user request for blue.
    /// </summary>
    private static RunAgentInput BuildContinuationInput()
    {
        var input = BuildFirstTurnInput();
        input.Messages =
        [
            .. input.Messages,
            new AGUIAssistantMessage
            {
                ToolCalls =
                [
                    new AGUIToolCall
                    {
                        Id = "call-old",
                        Function = new AGUIToolCallFunction { Name = "set_colour", Arguments = """{"colour":"red"}""" }
                    }
                ]
            },
            new AGUIToolMessage { ToolCallId = "call-old", Content = "Applied red" },
            new AGUIUserMessage { Content = "Now change it to blue." }
        ];
        return input;
    }

    private static async Task<List<BaseEvent>> RunTurn(
        ScriptedChatClient llm, IList<AITool> serverTools, RunAgentInput input)
    {
        // Server tools are registered through the adapter so a continuation may resume them.
        var context = input.ToChatRequestContext(SerializerOptions, null, serverTools);

        var ficc = new ChatClientBuilder(llm).UseFunctionInvocation().Build();
        var events = new List<BaseEvent>();
        await foreach (var evt in ficc.GetStreamingResponseAsync(context.Messages, context.ChatOptions)
            .AsAGUIEventStreamAsync(context).ConfigureAwait(false))
        {
            events.Add(evt);
        }

        return events;
    }

    /// <summary>
    /// A fake model: the first call proposes the given tool calls, every later call answers with
    /// text. Records the client tools it was offered on each call.
    /// </summary>
    private sealed class ScriptedChatClient(params AIContent[] firstCallContents) : IChatClient
    {
        public int Calls { get; private set; }

        public List<List<AITool>> ToolsSeen { get; } = [];

        public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
            IEnumerable<ChatMessage> messages,
            ChatOptions? options = null,
            [EnumeratorCancellation] CancellationToken cancellationToken = default)
        {
            await Task.Yield();
            ToolsSeen.Add(options?.Tools?.Where(t => t.Name == "set_colour").ToList() ?? []);
            if (Calls++ == 0)
            {
                yield return new ChatResponseUpdate
                {
                    Role = ChatRole.Assistant,
                    Contents = [.. firstCallContents],
                    FinishReason = ChatFinishReason.ToolCalls,
                };
                yield break;
            }

            yield return new ChatResponseUpdate { Role = ChatRole.Assistant, Contents = [new TextContent("done")] };
        }

        public Task<ChatResponse> GetResponseAsync(
            IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default) =>
            throw new NotSupportedException();

        public object? GetService(Type serviceType, object? serviceKey = null) => null;

        public void Dispose()
        {
        }
    }
}