microsoft--agent-framework/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/FileMemory/FileMemoryProviderTests.cs

// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Moq;

namespace Microsoft.Agents.AI.UnitTests.Harness.FileMemory;

public class FileMemoryProviderTests
{
    #region Constructor Validation

    [Fact]
    public void Constructor_NullFileStore_Throws()
    {
        Assert.Throws<ArgumentNullException>(() => new FileMemoryProvider(null!));
    }

    [Fact]
    public void Constructor_WithDefaults_Succeeds()
    {
        // Act
        var provider = new FileMemoryProvider(new InMemoryAgentFileStore());

        // Assert
        Assert.NotNull(provider);
    }

    [Fact]
    public void Constructor_WithStateInitializer_Succeeds()
    {
        // Act
        var provider = new FileMemoryProvider(
            new InMemoryAgentFileStore(),
            _ => new FileMemoryState { WorkingFolder = "custom" });

        // Assert
        Assert.NotNull(provider);
    }

    #endregion

    #region ProvideAIContextAsync Tests

    [Fact]
    public async Task ProvideAIContextAsync_ReturnsToolsAsync()
    {
        // Arrange
        var (tools, _, session) = await CreateToolsAsync();

        // Assert - 7 tools: Write, Read, Delete, Ls, Grep, Replace, ReplaceLines
        Assert.Equal(7, tools.Count());
    }

    [Fact]
    public async Task ProvideAIContextAsync_ReturnsInstructionsAsync()
    {
        // Arrange
        var provider = new FileMemoryProvider(new InMemoryAgentFileStore());
        var agent = new Mock<AIAgent>().Object;
        var session = new ChatClientAgentSession();
#pragma warning disable MAAI001
        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
#pragma warning restore MAAI001

        // Act
        AIContext result = await provider.InvokingAsync(context);

        // Assert
        Assert.NotNull(result.Instructions);
        Assert.Contains("file-based memory", result.Instructions);
        Assert.Contains("compacted", result.Instructions);
    }

    #endregion

    #region SaveFile Tests

    [Fact]
    public async Task SaveFile_CreatesFileAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");

        // Act
        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "Test content",
            ["description"] = "",
        }, session);

        // Assert
        var content = await store.ReadAsync("notes.md");
        Assert.Equal("Test content", content);
    }

    [Fact]
    public async Task SaveFile_WithDescription_CreatesBothFilesAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");

        // Act
        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "research.md",
            ["content"] = "Long research content...",
            ["description"] = "Summary of research findings",
        }, session);

        // Assert
        var content = await store.ReadAsync("research.md");
        Assert.Equal("Long research content...", content);
        var desc = await store.ReadAsync("research_description.md");
        Assert.Equal("Summary of research findings", desc);
    }

    [Fact]
    public async Task SaveFile_WithoutDescription_DeletesStaleDescriptionAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");

        // Save with description first.
        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "Original",
            ["description"] = "Old description",
        }, session);
        Assert.NotNull(await store.ReadAsync("notes_description.md"));

        // Act — overwrite without description.
        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "Updated",
        }, session);

        // Assert — stale description file is removed.
        Assert.Equal("Updated", await store.ReadAsync("notes.md"));
        Assert.Null(await store.ReadAsync("notes_description.md"));
    }

    [Fact]
    public async Task SaveFile_WithCustomState_CreatesInSubfolderAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, state, session) = await CreateToolsAsync(store, _ => new FileMemoryState { WorkingFolder = "session123" });
        var saveFile = GetTool(tools, "file_memory_write");

        // Act
        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "Session content",
            ["description"] = "",
        }, session);

        // Assert
        Assert.Equal("session123", state.WorkingFolder);
        var content = await store.ReadAsync("session123/notes.md");
        Assert.Equal("Session content", content);
    }

    #endregion

    #region ReadFile Tests

    [Fact]
    public async Task ReadFile_ExistingFile_ReturnsContentAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        await store.WriteAsync("notes.md", "Stored content");
        var (tools, _, session) = await CreateToolsAsync(store);
        var readFile = GetTool(tools, "file_memory_read");

        // Act
        var result = await InvokeWithRunContextAsync(readFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
        }, session);

        // Assert
        var text = Assert.IsType<JsonElement>(result).GetString();
        Assert.Equal("Stored content", text);
    }

    [Fact]
    public async Task ReadFile_NonExistent_ReturnsNotFoundMessageAsync()
    {
        // Arrange
        var (tools, _, session) = await CreateToolsAsync();
        var readFile = GetTool(tools, "file_memory_read");

        // Act
        var result = await InvokeWithRunContextAsync(readFile, new AIFunctionArguments
        {
            ["fileName"] = "nonexistent.md",
        }, session);

        // Assert
        var text = Assert.IsType<JsonElement>(result).GetString();
        Assert.Contains("not found", text);
    }

    [Fact]
    public async Task ReadFile_InternalFile_ThrowsAsync()
    {
        // Arrange — the memory index file name is reserved for internal use.
        var store = new InMemoryAgentFileStore();
        await store.WriteAsync("memories.md", "internal");
        var (tools, _, session) = await CreateToolsAsync(store);
        var readFile = GetTool(tools, "file_memory_read");

        // Act & Assert — reserved internal files are not reachable through the read tool.
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(readFile, new AIFunctionArguments
            {
                ["fileName"] = "memories.md",
            }, session));
    }

    #endregion

    #region DeleteFile Tests

    [Fact]
    public async Task DeleteFile_ExistingFile_DeletesAndReturnsConfirmationAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        await store.WriteAsync("notes.md", "Content");
        var (tools, _, session) = await CreateToolsAsync(store);
        var deleteFile = GetTool(tools, "file_memory_delete");

        // Act
        var result = await InvokeWithRunContextAsync(deleteFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
        }, session);

        // Assert
        var text = Assert.IsType<JsonElement>(result).GetString();
        Assert.Contains("deleted", text);
        Assert.False(await store.FileExistsAsync("notes.md"));
    }

    [Fact]
    public async Task DeleteFile_AlsoDeletesDescriptionFileAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        await store.WriteAsync("notes.md", "Content");
        await store.WriteAsync("notes_description.md", "Description");
        var (tools, _, session) = await CreateToolsAsync(store);
        var deleteFile = GetTool(tools, "file_memory_delete");

        // Act
        await InvokeWithRunContextAsync(deleteFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
        }, session);

        // Assert
        Assert.False(await store.FileExistsAsync("notes.md"));
        Assert.False(await store.FileExistsAsync("notes_description.md"));
    }

    [Fact]
    public async Task DeleteFile_InternalFile_ThrowsAsync()
    {
        // Arrange — the memory index file name is reserved for internal use.
        var store = new InMemoryAgentFileStore();
        await store.WriteAsync("memories.md", "internal");
        var (tools, _, session) = await CreateToolsAsync(store);
        var deleteFile = GetTool(tools, "file_memory_delete");

        // Act & Assert — reserved internal files cannot be deleted through the delete tool.
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(deleteFile, new AIFunctionArguments
            {
                ["fileName"] = "memories.md",
            }, session));
        Assert.True(await store.FileExistsAsync("memories.md"));
    }

    #endregion

    #region ListFiles Tests

    [Fact]
    public async Task ListFiles_ReturnsFilesWithDescriptionsAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        await store.WriteAsync("notes.md", "Content");
        await store.WriteAsync("notes_description.md", "A description");
        await store.WriteAsync("other.md", "Other content");
        var (tools, _, session) = await CreateToolsAsync(store);
        var listFiles = GetTool(tools, "file_memory_ls");

        // Act
        var result = await InvokeWithRunContextAsync(listFiles, new AIFunctionArguments(), session);

        // Assert
        var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
        Assert.Equal(2, entries.Count);

        var notesEntry = entries.First(e => e.GetProperty("name").GetString() == "notes.md");
        Assert.Equal("A description", notesEntry.GetProperty("description").GetString());

        var otherEntry = entries.First(e => e.GetProperty("name").GetString() == "other.md");
        Assert.False(otherEntry.TryGetProperty("description", out _));
    }

    [Fact]
    public async Task ListFiles_HidesDescriptionFilesAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        await store.WriteAsync("notes.md", "Content");
        await store.WriteAsync("notes_description.md", "Desc");
        var (tools, _, session) = await CreateToolsAsync(store);
        var listFiles = GetTool(tools, "file_memory_ls");

        // Act
        var result = await InvokeWithRunContextAsync(listFiles, new AIFunctionArguments(), session);

        // Assert
        var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
        Assert.Single(entries);
        Assert.Equal("notes.md", entries[0].GetProperty("name").GetString());
    }

    #endregion

    #region SearchFiles Tests

    [Fact]
    public async Task SearchFiles_FindsMatchingContentAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        await store.WriteAsync("notes.md", "Important research findings about AI");
        var (tools, _, session) = await CreateToolsAsync(store);
        var searchFiles = GetTool(tools, "file_memory_grep");

        // Act
        var result = await InvokeWithRunContextAsync(searchFiles, new AIFunctionArguments
        {
            ["regexPattern"] = "research findings",
            ["globPattern"] = "",
        }, session);

        // Assert
        var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
        Assert.Single(entries);
        Assert.Equal("notes.md", entries[0].GetProperty("fileName").GetString());
        Assert.True(entries[0].TryGetProperty("matchingLines", out var matchingLines));
        Assert.True(matchingLines.GetArrayLength() > 0);
    }

    [Fact]
    public async Task SearchFiles_WithFilePattern_FiltersResultsAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        await store.WriteAsync("notes.md", "Important data");
        await store.WriteAsync("data.txt", "Important data");
        var (tools, _, session) = await CreateToolsAsync(store);
        var searchFiles = GetTool(tools, "file_memory_grep");

        // Act
        var result = await InvokeWithRunContextAsync(searchFiles, new AIFunctionArguments
        {
            ["regexPattern"] = "Important",
            ["globPattern"] = "*.md",
        }, session);

        // Assert
        var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
        Assert.Single(entries);
        Assert.Equal("notes.md", entries[0].GetProperty("fileName").GetString());
    }

    #endregion

    #region State Initializer Tests

    [Fact]
    public async Task CustomStateInitializer_SetsWorkingFolderAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (_, state, _) = await CreateToolsAsync(store, _ => new FileMemoryState { WorkingFolder = "user42" });

        // Assert
        Assert.Equal("user42", state.WorkingFolder);
    }

    [Fact]
    public async Task DefaultStateInitializer_UsesEmptyWorkingFolderAsync()
    {
        // Arrange
        var (_, state, _) = await CreateToolsAsync();

        // Assert
        Assert.Equal(string.Empty, state.WorkingFolder);
    }

    [Fact]
    public async Task State_PersistsAcrossInvocationsAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var provider = new FileMemoryProvider(store, _ => new FileMemoryState { WorkingFolder = "persistent" });
        var agent = new Mock<AIAgent>().Object;
        var session = new ChatClientAgentSession();
#pragma warning disable MAAI001
        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
#pragma warning restore MAAI001

        // Act - first invocation initializes state
        await provider.InvokingAsync(context);
        session.StateBag.TryGetValue<FileMemoryState>("FileMemoryProvider", out var state1, AgentJsonUtilities.DefaultOptions);

        // Second invocation should reuse the same folder
        await provider.InvokingAsync(context);
        session.StateBag.TryGetValue<FileMemoryState>("FileMemoryProvider", out var state2, AgentJsonUtilities.DefaultOptions);

        // Assert
        Assert.NotNull(state1);
        Assert.NotNull(state2);
        Assert.Equal(state1!.WorkingFolder, state2!.WorkingFolder);
    }

    #endregion

    #region Path Traversal Protection

    [Fact]
    public async Task SaveFile_PathTraversal_ThrowsAsync()
    {
        // Arrange
        var (tools, _, session) = await CreateToolsAsync();
        var saveFile = GetTool(tools, "file_memory_write");

        // Act & Assert — the normalization error bubbles up.
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
            {
                ["fileName"] = "../escape.md",
                ["content"] = "Content",
                ["description"] = "",
            }, session));
    }

    [Fact]
    public async Task SaveFile_AbsolutePath_ThrowsAsync()
    {
        // Arrange
        var (tools, _, session) = await CreateToolsAsync();
        var saveFile = GetTool(tools, "file_memory_write");

        // Act & Assert
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
            {
                ["fileName"] = "/etc/passwd",
                ["content"] = "Content",
                ["description"] = "",
            }, session));
    }

    [Fact]
    public async Task SaveFile_DriveRootedPath_ThrowsAsync()
    {
        // Arrange
        var (tools, _, session) = await CreateToolsAsync();
        var saveFile = GetTool(tools, "file_memory_write");

        // Act & Assert
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
            {
                ["fileName"] = "C:\\temp\\file.md",
                ["content"] = "Content",
            }, session));
    }

    [Fact]
    public async Task SaveFile_DoubleDotsInFileName_AllowedAsync()
    {
        // Arrange — "notes..md" is not a path traversal attempt.
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");

        // Act
        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "notes..md",
            ["content"] = "Content",
        }, session);

        // Assert
        Assert.Equal("Content", await store.ReadAsync("notes..md"));
    }

    #endregion

    #region Memory Index Tests

    [Fact]
    public async Task SaveFile_CreatesMemoryIndexAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");

        // Act
        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "Test content",
        }, session);

        // Assert — memories.md should exist and contain the file entry.
        string? index = await store.ReadAsync("memories.md");
        Assert.NotNull(index);
        Assert.Contains("**notes.md**", index);
    }

    [Fact]
    public async Task SaveFile_WithDescription_IndexIncludesDescriptionAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");

        // Act
        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "research.md",
            ["content"] = "Research data",
            ["description"] = "Key findings",
        }, session);

        // Assert
        string? index = await store.ReadAsync("memories.md");
        Assert.NotNull(index);
        Assert.Contains("**research.md**: Key findings", index);
    }

    [Fact]
    public async Task DeleteFile_UpdatesMemoryIndexAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");
        var deleteFile = GetTool(tools, "file_memory_delete");

        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "Content",
        }, session);

        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "other.md",
            ["content"] = "Other",
        }, session);

        // Act
        await InvokeWithRunContextAsync(deleteFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
        }, session);

        // Assert — index should only contain other.md
        string? index = await store.ReadAsync("memories.md");
        Assert.NotNull(index);
        Assert.DoesNotContain("notes.md", index);
        Assert.Contains("**other.md**", index);
    }

    [Fact]
    public async Task MemoryIndex_CappedAt50EntriesAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");

        // Act — save 55 files
        for (int i = 0; i < 55; i++)
        {
            await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
            {
                ["fileName"] = $"file{i:D3}.md",
                ["content"] = $"Content {i}",
            }, session);
        }

        // Assert — index should have at most 50 entries
        string? index = await store.ReadAsync("memories.md");
        Assert.NotNull(index);

        int entryCount = 0;
        foreach (string line in index!.Split('\n'))
        {
            if (line.StartsWith("- **", StringComparison.Ordinal))
            {
                entryCount++;
            }
        }

        Assert.Equal(50, entryCount);
    }

    [Fact]
    public async Task ListFiles_HidesMemoryIndexAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");
        var listFiles = GetTool(tools, "file_memory_ls");

        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "Content",
        }, session);

        // Act
        var result = await InvokeWithRunContextAsync(listFiles, new AIFunctionArguments(), session);

        // Assert — memories.md should not appear in the listing
        var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
        Assert.Single(entries);
        Assert.Equal("notes.md", entries[0].GetProperty("name").GetString());
    }

    [Fact]
    public async Task ProvideAIContextAsync_InjectsMemoryIndexMessageAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var provider = new FileMemoryProvider(store);
        var agent = new Mock<AIAgent>().Object;
        var session = new ChatClientAgentSession();

        // First, save a file via tool invocation to create the index.
#pragma warning disable MAAI001
        var initContext = new AIContextProvider.InvokingContext(agent, session, new AIContext());
#pragma warning restore MAAI001
        AIContext initResult = await provider.InvokingAsync(initContext);
        var saveFile = GetTool(initResult.Tools!, "file_memory_write");
        await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
        {
            ["fileName"] = "research.md",
            ["content"] = "Data",
            ["description"] = "Research summary",
        }, session);

        // Act — invoke the provider again; it should now inject the memory index.
#pragma warning disable MAAI001
        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
#pragma warning restore MAAI001
        AIContext result = await provider.InvokingAsync(context);

        // Assert
        Assert.NotNull(result.Messages);
        var messages = result.Messages!.ToList();
        Assert.Single(messages);
        Assert.Equal(ChatRole.User, messages[0].Role);
        Assert.Contains("memory index", messages[0].Text, StringComparison.OrdinalIgnoreCase);
        Assert.Contains("research.md", messages[0].Text);
    }

    [Fact]
    public async Task ProvideAIContextAsync_NoFiles_NoMessageInjectedAsync()
    {
        // Arrange
        var provider = new FileMemoryProvider(new InMemoryAgentFileStore());
        var agent = new Mock<AIAgent>().Object;
        var session = new ChatClientAgentSession();
#pragma warning disable MAAI001
        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
#pragma warning restore MAAI001

        // Act
        AIContext result = await provider.InvokingAsync(context);

        // Assert — no memories.md exists, so no message should be injected
        Assert.Null(result.Messages);
    }

    #endregion

    #region Replace Tests

    [Fact]
    public async Task Replace_SingleOccurrence_ReplacesAndReturnsCountAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var write = GetTool(tools, "file_memory_write");
        var replace = GetTool(tools, "file_memory_replace");
        await InvokeWithRunContextAsync(write, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "Hello world",
        }, session);

        // Act
        var result = await InvokeWithRunContextAsync(replace, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["oldString"] = "world",
            ["newString"] = "there",
        }, session);

        // Assert
        var text = Assert.IsType<JsonElement>(result).GetString();
        Assert.Contains("Replaced 1 occurrence(s)", text);
        Assert.Equal("Hello there", await store.ReadAsync("notes.md"));
    }

    [Fact]
    public async Task Replace_InternalFile_ThrowsAsync()
    {
        // Arrange — the memory index file name is reserved for internal use.
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var replace = GetTool(tools, "file_memory_replace");

        // Act & Assert
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(replace, new AIFunctionArguments
            {
                ["fileName"] = "memories.md",
                ["oldString"] = "x",
                ["newString"] = "y",
            }, session));
    }

    [Fact]
    public async Task Replace_ReplaceAll_ReplacesEveryOccurrenceAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var write = GetTool(tools, "file_memory_write");
        var replace = GetTool(tools, "file_memory_replace");
        await InvokeWithRunContextAsync(write, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "a a a",
        }, session);

        // Act
        var result = await InvokeWithRunContextAsync(replace, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["oldString"] = "a",
            ["newString"] = "b",
            ["replaceAll"] = true,
        }, session);

        // Assert
        var text = Assert.IsType<JsonElement>(result).GetString();
        Assert.Contains("Replaced 3 occurrence(s)", text);
        Assert.Equal("b b b", await store.ReadAsync("notes.md"));
    }

    [Fact]
    public async Task Replace_MultipleOccurrences_WithoutReplaceAll_ThrowsAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var write = GetTool(tools, "file_memory_write");
        var replace = GetTool(tools, "file_memory_replace");
        await InvokeWithRunContextAsync(write, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "a a a",
        }, session);

        // Act & Assert — exception bubbles, content unchanged.
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(replace, new AIFunctionArguments
            {
                ["fileName"] = "notes.md",
                ["oldString"] = "a",
                ["newString"] = "b",
            }, session));
        Assert.Equal("a a a", await store.ReadAsync("notes.md"));
    }

    [Fact]
    public async Task Replace_NonExistentFile_ReturnsNotFoundAsync()
    {
        // Arrange
        var (tools, _, session) = await CreateToolsAsync();
        var replace = GetTool(tools, "file_memory_replace");

        // Act
        var result = await InvokeWithRunContextAsync(replace, new AIFunctionArguments
        {
            ["fileName"] = "missing.md",
            ["oldString"] = "x",
            ["newString"] = "y",
        }, session);

        // Assert
        var text = Assert.IsType<JsonElement>(result).GetString();
        Assert.Contains("not found", text);
    }

    #endregion

    #region ReplaceLines Tests

    [Fact]
    public async Task ReplaceLines_ReplacesSpecifiedLinesAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var write = GetTool(tools, "file_memory_write");
        var replaceLines = GetTool(tools, "file_memory_replace_lines");
        await InvokeWithRunContextAsync(write, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "line1\nline2\nline3",
        }, session);

        // Act
        var result = await InvokeWithRunContextAsync(replaceLines, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["edits"] = new List<FileLineEdit> { new() { LineNumber = 2, NewLine = "CHANGED\n" } },
        }, session);

        // Assert
        var text = Assert.IsType<JsonElement>(result).GetString();
        Assert.Contains("Replaced 1 line(s)", text);
        Assert.Equal("line1\nCHANGED\nline3", await store.ReadAsync("notes.md"));
    }

    [Fact]
    public async Task ReplaceLines_InternalFile_ThrowsAsync()
    {
        // Arrange — the memory index file name is reserved for internal use.
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var replaceLines = GetTool(tools, "file_memory_replace_lines");

        // Act & Assert
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(replaceLines, new AIFunctionArguments
            {
                ["fileName"] = "memories.md",
                ["edits"] = new List<FileLineEdit> { new() { LineNumber = 1, NewLine = "X" } },
            }, session));
    }

    [Fact]
    public async Task ReplaceLines_OutOfRange_ThrowsAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var write = GetTool(tools, "file_memory_write");
        var replaceLines = GetTool(tools, "file_memory_replace_lines");
        await InvokeWithRunContextAsync(write, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "line1\nline2",
        }, session);

        // Act & Assert — exception bubbles, content unchanged.
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(replaceLines, new AIFunctionArguments
            {
                ["fileName"] = "notes.md",
                ["edits"] = new List<FileLineEdit> { new() { LineNumber = 5, NewLine = "X" } },
            }, session));
        Assert.Equal("line1\nline2", await store.ReadAsync("notes.md"));
    }

    #endregion

    #region Write Guards

    [Fact]
    public async Task Write_ReservedInternalFileName_ThrowsAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var write = GetTool(tools, "file_memory_write");

        // Act & Assert — the file is not written and an ArgumentException is thrown.
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(write, new AIFunctionArguments
            {
                ["fileName"] = "memories.md",
                ["content"] = "Content",
            }, session));
        Assert.False(await store.FileExistsAsync("memories.md"));
    }

    [Fact]
    public async Task Write_NestedPath_ThrowsAsync()
    {
        // Arrange — memory files must be written with a flat name.
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var write = GetTool(tools, "file_memory_write");

        // Act & Assert
        await Assert.ThrowsAsync<ArgumentException>(async () =>
            await InvokeWithRunContextAsync(write, new AIFunctionArguments
            {
                ["fileName"] = "sub/notes.md",
                ["content"] = "Content",
            }, session));
    }

    [Fact]
    public async Task Write_WithDescription_ReturnsWrittenWithDescriptionAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var write = GetTool(tools, "file_memory_write");

        // Act
        var result = await InvokeWithRunContextAsync(write, new AIFunctionArguments
        {
            ["fileName"] = "notes.md",
            ["content"] = "Content",
            ["description"] = "A summary",
        }, session);

        // Assert
        var text = Assert.IsType<JsonElement>(result).GetString();
        Assert.Contains("written with description", text);
    }

    #endregion

    #region Helper Methods

    private static FileMemoryProvider CreateProvider(InMemoryAgentFileStore? store = null, Func<AgentSession?, FileMemoryState>? stateInitializer = null)
    {
        return new FileMemoryProvider(store ?? new InMemoryAgentFileStore(), stateInitializer);
    }

    private static async Task<(IEnumerable<AITool> Tools, FileMemoryState State, AgentSession Session)> CreateToolsAsync(InMemoryAgentFileStore? store = null, Func<AgentSession?, FileMemoryState>? stateInitializer = null)
    {
        var provider = CreateProvider(store, stateInitializer);
        var agent = new Mock<AIAgent>().Object;
        var session = new ChatClientAgentSession();
#pragma warning disable MAAI001
        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
#pragma warning restore MAAI001

        AIContext result = await provider.InvokingAsync(context);

        session.StateBag.TryGetValue<FileMemoryState>("FileMemoryProvider", out var state, AgentJsonUtilities.DefaultOptions);

        return (result.Tools!, state!, session);
    }

    private static AIFunction GetTool(IEnumerable<AITool> tools, string name)
    {
        return (AIFunction)tools.First(t => t is AIFunction f && f.Name == name);
    }

    /// <summary>
    /// Invokes a tool within a mock <see cref="AIAgent.CurrentRunContext"/> so that
    /// the tool methods can access the session via <c>AIAgent.CurrentRunContext?.Session</c>.
    /// </summary>
    /// <param name="tool">The tool to invoke.</param>
    /// <param name="arguments">The arguments to pass to the tool.</param>
    /// <param name="session">
    /// An optional session to use in the run context. When provided, ensures the tool executes
    /// against the same session whose state was initialized during <see cref="CreateToolsAsync"/>.
    /// When <see langword="null"/>, a new session is created.
    /// </param>
    private static async Task<object?> InvokeWithRunContextAsync(AIFunction tool, AIFunctionArguments arguments, AgentSession? session = null)
    {
        var agent = new Mock<AIAgent>().Object;
        session ??= new ChatClientAgentSession();
        var messages = new List<ChatMessage>();

        // Set up the ambient run context so tool methods can access the session.
        var runContext = new AgentRunContext(agent, session, messages, null);

        // Use reflection to set the protected static CurrentRunContext property.
        var property = typeof(AIAgent).GetProperty("CurrentRunContext", System.Reflection.BindingFlags.Public | System.Reflection.BindingFlags.Static);
        var setter = property!.GetSetMethod(true)!;
        var previousContext = AIAgent.CurrentRunContext;
        try
        {
            setter.Invoke(null, [runContext]);
            return await tool.InvokeAsync(arguments);
        }
        finally
        {
            setter.Invoke(null, [previousContext]);
        }
    }

    #endregion

    #region Options Tests

    /// <summary>
    /// Verify that custom instructions override the default.
    /// </summary>
    [Fact]
    public async Task Options_CustomInstructions_OverridesDefaultAsync()
    {
        // Arrange
        var options = new FileMemoryProviderOptions { Instructions = "Custom file memory instructions." };
        var provider = new FileMemoryProvider(new InMemoryAgentFileStore(), options: options);
        var agent = new Mock<AIAgent>().Object;
        var session = new ChatClientAgentSession();
#pragma warning disable MAAI001
        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
#pragma warning restore MAAI001

        // Act
        AIContext result = await provider.InvokingAsync(context);

        // Assert
        Assert.Equal("Custom file memory instructions.", result.Instructions);
    }

    /// <summary>
    /// Verify that null options uses default instructions.
    /// </summary>
    [Fact]
    public async Task Options_Null_UsesDefaultInstructionsAsync()
    {
        // Arrange
        var provider = new FileMemoryProvider(new InMemoryAgentFileStore());
        var agent = new Mock<AIAgent>().Object;
        var session = new ChatClientAgentSession();
#pragma warning disable MAAI001
        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
#pragma warning restore MAAI001

        // Act
        AIContext result = await provider.InvokingAsync(context);

        // Assert
        Assert.Contains("file-based memory", result.Instructions);
    }

    #endregion

    #region Thread Safety Tests

    [Fact]
    public async Task ConcurrentSaves_ProduceConsistentIndexAsync()
    {
        // Arrange
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");
        const int FileCount = 20;

        // Act — save multiple files in parallel.
        var tasks = new Task[FileCount];
        for (int i = 0; i < FileCount; i++)
        {
            int index = i;
            tasks[i] = InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
            {
                ["fileName"] = $"file{index}.md",
                ["content"] = $"Content {index}",
                ["description"] = $"Description {index}",
            }, session);
        }

        await Task.WhenAll(tasks);

        // Assert — the memory index should contain all files.
        string? indexContent = await store.ReadAsync("memories.md");
        Assert.NotNull(indexContent);
        for (int i = 0; i < FileCount; i++)
        {
            Assert.Contains($"**file{i}.md**", indexContent);
        }
    }

    [Fact]
    public async Task ConcurrentSaveAndDelete_ProduceConsistentIndexAsync()
    {
        // Arrange — pre-populate files that will be deleted.
        var store = new InMemoryAgentFileStore();
        var (tools, _, session) = await CreateToolsAsync(store);
        var saveFile = GetTool(tools, "file_memory_write");
        var deleteFile = GetTool(tools, "file_memory_delete");

        for (int i = 0; i < 5; i++)
        {
            await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
            {
                ["fileName"] = $"delete{i}.md",
                ["content"] = $"To be deleted {i}",
            }, session);
        }

        // Act — concurrently save new files and delete existing ones.
        var tasks = new List<Task>();
        for (int i = 0; i < 5; i++)
        {
            int index = i;
            tasks.Add(InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
            {
                ["fileName"] = $"keep{index}.md",
                ["content"] = $"Kept {index}",
            }, session));
            tasks.Add(InvokeWithRunContextAsync(deleteFile, new AIFunctionArguments
            {
                ["fileName"] = $"delete{index}.md",
            }, session));
        }

        await Task.WhenAll(tasks);

        // Assert — index should contain only the kept files.
        string? indexContent = await store.ReadAsync("memories.md");
        Assert.NotNull(indexContent);
        for (int i = 0; i < 5; i++)
        {
            Assert.Contains($"**keep{i}.md**", indexContent);
            Assert.DoesNotContain($"**delete{i}.md**", indexContent);
        }
    }

    [Fact]
    public void Dispose_ReleasesResources()
    {
        // Arrange
        var provider = new FileMemoryProvider(new InMemoryAgentFileStore());

        // Act
        provider.Dispose();

        // Assert — calling Dispose again should not throw (idempotent SemaphoreSlim.Dispose).
        provider.Dispose();
    }

    [Fact]
    public async Task SaveFile_AfterDispose_ThrowsAsync()
    {
        // Arrange — create tools from a provider, then dispose the provider.
        var store = new InMemoryAgentFileStore();
        var provider = CreateProvider(store);
        var agent = new Mock<AIAgent>().Object;
        var session = new ChatClientAgentSession();
#pragma warning disable MAAI001
        var context = new AIContextProvider.InvokingContext(agent, session, new AIContext());
#pragma warning restore MAAI001
        AIContext result = await provider.InvokingAsync(context);
        var saveFile = GetTool(result.Tools!, "file_memory_write");
        provider.Dispose();

        // Act & Assert — the disposed SemaphoreSlim should throw ObjectDisposedException.
        await Assert.ThrowsAsync<ObjectDisposedException>(async () =>
            await InvokeWithRunContextAsync(saveFile, new AIFunctionArguments
            {
                ["fileName"] = "notes.md",
                ["content"] = "Should fail",
            }, session));
    }

    #endregion
}