YiemAgent/test/toolTest.jl

using Test
using Dates
using YiemAgent
using YiemAgent.toolRegistry
using YiemAgent.type
using YiemAgent.agentCore

@testset "register_all_tools with toolStore" begin

    # ------------------------------------------------------------------ #
    # 1. register_all_tools registers all static tools + listTools       #
    # ------------------------------------------------------------------ #
    store = toolStore(name="test1")
    loaded = register_all_tools(store)
    @test !isempty(loaded)
    @test length(loaded) == 4  # getWeather + getTime + writeTool + listTools

    names = [k for k in keys(loaded)]
    @test "getTime" in names
    @test "getWeather" in names
    @test "writeTool" in names
    @test "listTools" in names

    # ------------------------------------------------------------------ #
    # 2. register_all_tools returns tools in registration order          #
    # ------------------------------------------------------------------ #
    @test collect(keys(loaded))[1] == "getWeather"
    @test collect(keys(loaded))[2] == "getTime"
    @test collect(keys(loaded))[3] == "writeTool"
    @test collect(keys(loaded))[4] == "listTools"

    # ------------------------------------------------------------------ #
    # 3. Verify loaded tool fields are correct                           #
    # ------------------------------------------------------------------ #
    # getTime
    time_tool = loaded["getTime"]
    @test time_tool.name == "getTime"
    @test time_tool.label == "Time Lookup"
    @test time_tool.validateRequiredArgs !== nothing
    @test time_tool.parallelToolExecute == false
    @test time_tool.inputSchema["required"] == Any[]

    # getWeather
    weather = loaded["getWeather"]
    @test weather.name == "getWeather"
    @test weather.label == "Weather Lookup"
    @test weather.execute !== nothing
    @test weather.parallelToolExecute == false
    @test weather.inputSchema["required"] == ["city"]

    # writeTool
    wt = loaded["writeTool"]
    @test wt.name == "writeTool"
    @test wt.label == "Create Tool"
    @test wt.execute !== nothing
    @test "name" in wt.inputSchema["required"]
    @test "executeCode" in wt.inputSchema["required"]

    # ------------------------------------------------------------------ #
    # 4. Tool execution returns valid results                            #
    # ------------------------------------------------------------------ #
    sig = nothing
    op = x -> x  # no-op partial result callback

    # execute getTime
    result_t = time_tool.execute("call-1", Dict{String,Any}("city" => "Tokyo"), sig, op)
    @test result_t isa agentToolResult
    @test result_t.content[1] isa textContent
    @test occursin("Tokyo", result_t.content[1].text)

    # execute getTime with timezone
    result_tz = time_tool.execute("call-2", Dict{String,Any}("timezone" => "America/New_York"), sig, op)
    @test result_tz isa agentToolResult
    @test occursin("America/New_York", result_tz.content[1].text)

    # execute getWeather
    result_w = weather.execute("call-3", Dict{String,Any}("city" => "Bangkok"), sig, op)
    @test result_w isa agentToolResult
    @test result_w.content[1] isa textContent
    @test occursin("Bangkok", result_w.content[1].text)

    # execute getWeather with units
    result_w2 = weather.execute("call-4", Dict{String,Any}("city" => "London", "units" => "fahrenheit"), sig, op)
    @test occursin("72\u00b0F", result_w2.content[1].text)

    # ------------------------------------------------------------------ #
    # 5. getTools / registerTool / clearTools (per-store isolation)      #
    # ------------------------------------------------------------------ #
    store3 = toolStore(name="test3")
    registry_tools = getTools(store3)
    @test isempty(registry_tools)

    # Register tools manually
    registerTool(store3, loaded["getTime"])
    registerTool(store3, loaded["getWeather"])
    registerTool(store3, loaded["writeTool"])

    reg = getTools(store3)
    @test !isempty(reg)
    @test "getTime" in keys(reg)
    @test "getWeather" in keys(reg)
    @test "writeTool" in keys(reg)
    @test collect(keys(reg))[1] == "getTime"
    @test collect(keys(reg))[2] == "getWeather"
    @test collect(keys(reg))[3] == "writeTool"

    clearTools(store3)
    @test isempty(getTools(store3))

    test_tool = agentTool(
        name = "manualTool",
        label = "Manual Tool",
        description = "Registered manually",
        inputSchema = Dict{String,Any}("type" => "object", "properties" => Dict{String,Any}(), "required" => Any[]),
        execute = (toolCallId, args, signal, onPartialResult) ->
            agentToolResult([textContent("manual")], Dict{Any,Any}(), nothing, false),
        prepareArguments = nothing,
        validateRequiredArgs = nothing,
        parallelToolExecute = true
    )
    registerTool(store3, test_tool)
    reg = getTools(store3)
    @test haskey(reg, "manualTool")
    @test length(reg) == 1
    @test reg["manualTool"].parallelToolExecute == true

    # ------------------------------------------------------------------ #
    # 6. getTools returns direct reference (mutations affect registry)   #
    # ------------------------------------------------------------------ #
    copy1 = getTools(store3)
    copy2 = getTools(store3)
    @test copy1 === copy2  # same reference, not a deep copy
    empty!(copy1)
    @test isempty(getTools(store3))  # mutation propagates

    # ------------------------------------------------------------------ #
    # 7. Per-store isolation — two stores don't share tools              #
    # ------------------------------------------------------------------ #
    storeA = toolStore(name="isolationA")
    storeB = toolStore(name="isolationB")

    registerTool(storeA, loaded["getTime"])
    registerTool(storeB, loaded["getWeather"])

    regA = getTools(storeA)
    regB = getTools(storeB)

    @test "getTime" in keys(regA)
    @test "getWeather" ∉ keys(regA)
    @test "getWeather" in keys(regB)
    @test "getTime" ∉ keys(regB)

    clearTools(storeA)
    @test isempty(getTools(storeA))
    @test !isempty(getTools(storeB))  # storeB unaffected
end

@testset "listTool" begin
    store = toolStore(name="test_list")
    register_all_tools(store)  # auto-registers getWeather, getTime, writeTool + listTools

    # register_all_tools auto-registers listTool
    @test "listTools" in keys(store.tools)

    # listTool returns an agentTool, not a string or array
    list_t = listTool(store)
    @test list_t isa agentTool
    @test list_t.name == "listTools"
    @test list_t.label == "List Tools"
    @test isempty(list_t.inputSchema["required"])

    # Verify all tools appear (3 loaded + listTools = 4)
    result = list_t.execute("call-1", Dict{String,Any}(), nothing, x -> x)
    @test result isa agentToolResult
    @test result.content[1] isa textContent
    @test occursin("listTools", result.content[1].text)
    @test occursin("getWeather", result.content[1].text)
    @test occursin("getTime", result.content[1].text)
    @test occursin("writeTool", result.content[1].text)
    @test result.details["count"] == 4

    # Each listTool call creates an independent closure
    storeB = toolStore(name="test_listB")
    registerTool(storeB, store.tools["getWeather"])
    list_tB = listTool(storeB)

    resultA = list_t.execute("call-3", Dict{String,Any}(), nothing, x -> x)
    resultB = list_tB.execute("call-4", Dict{String,Any}(), nothing, x -> x)

    @test occursin("getWeather", resultA.content[1].text)
    @test occursin("getWeather", resultB.content[1].text)
    @test occursin("getTime", resultA.content[1].text)
    @test occursin("getTime", resultB.content[1].text) == false  # storeB only has getWeather
end

@testset "executePreparedToolCall with static tools" begin
    # Tests executePreparedToolCall with statically loaded tools.
    # The world-age issue is resolved because tool.execute comes from
    # a statically included module, not a dynamically created one.

    store = toolStore(name="test_static")
    register_all_tools(store)

    weather_tool = store.tools["getWeather"]

    # Create a preparedToolCall that mimics what prepareToolCall() returns
    tool_call = agentToolCall(
        "function", "call-static-1", "getWeather",
        Dict{String,Any}("city" => "San Francisco")
    )
    prep = preparedToolCall(
        weather_tool, tool_call, Dict{String,Any}("city" => "San Francisco")
    )
    sig = abortSignal(false)

    # This call goes through: executePreparedToolCall -> prep.tool.execute(...)
    result = executePreparedToolCall(
        prep, sig, x -> nothing
    )

    @test result isa executedOutcome
    @test result.isError == false
    @test result.result.content[1] isa textContent
    @test occursin("San Francisco", result.result.content[1].text)
end

@testset "executePreparedToolCall with validation (static tools)" begin
    # Tests executePreparedToolCall with a tool that has custom validation hooks.
    # This exercises the full tool execution path including validation.

    store = toolStore(name="test_static_validate")
    register_all_tools(store)

    time_tool = store.tools["getTime"]

    tool_call = agentToolCall(
        "function", "call-static-2", "getTime",
        Dict{String,Any}("timezone" => "America/New_York")
    )
    prep = preparedToolCall(
        time_tool, tool_call, Dict{String,Any}("timezone" => "America/New_York")
    )
    sig = abortSignal(false)

    result = executePreparedToolCall(
        prep, sig, x -> nothing
    )

    @test result isa executedOutcome
    @test result.isError == false
    @test result.result.content[1] isa textContent
    @test occursin("America/New_York", result.result.content[1].text)
end

@testset "executeToolCallsSequential with static tools (full pipeline)" begin
    # Tests the full tool execution pipeline: executeToolCallsSequential
    # which calls prepareToolCall -> executePreparedToolCall -> finalizeExecutedToolCall
    # with statically loaded tools.

    store = toolStore(name="test_full_pipeline")
    register_all_tools(store)

    # Build agentContext from the store's tools
    tools = getTools(store)
    ctx = agentContext(
        "test system prompt",
        agentMessage[],
        tools
    )

    # Create an assistant message containing tool calls
    assistant_msg = assistantMessage(
        role="assistant",
        content=Vector{messageContent}(),
        api="openai",
        provider="test",
        model="test-model",
        usage=llmUsage(0, 0),
        stopReason="tool_calls",
        errorMessage=nothing,
        timestamp=now()
    )

    # Create tool calls for multiple statically loaded tools
    tool_calls = [
        agentToolCall(
            "function", "call-seq-1", "getWeather",
            Dict{String,Any}("city" => "Tokyo")
        ),
        agentToolCall(
            "function", "call-seq-2", "getTime",
            Dict{String,Any}("timezone" => "Europe/London")
        ),
    ]

    config = agentLoopConfig(
        nothing, nothing, "sequential"
    )
    sig = abortSignal(false)

    # Execute the full pipeline
    batch = executeToolCallsSequential(
        ctx, assistant_msg, tool_calls, config, sig, x -> nothing
    )

    @test batch.messages isa Vector{toolResultMessage}
    @test length(batch.messages) == 2
    @test batch.messages[1].toolName == "getWeather"
    @test batch.messages[1].isError == false
    @test occursin("Tokyo", batch.messages[1].content[1].text)
    @test batch.messages[2].toolName == "getTime"
    @test batch.messages[2].isError == false
    @test occursin("Europe/London", batch.messages[2].content[1].text)
end

@testset "executeToolCallsParallel with static tools (full pipeline)" begin
    # Same as above but tests parallel execution path.

    store = toolStore(name="test_parallel")
    register_all_tools(store)

    tools = getTools(store)
    ctx = agentContext(
        "test system prompt",
        agentMessage[],
        tools
    )

    assistant_msg = assistantMessage(
        role="assistant",
        content=Vector{messageContent}(),
        api="openai",
        provider="test",
        model="test-model",
        usage=llmUsage(0, 0),
        stopReason="tool_calls",
        errorMessage=nothing,
        timestamp=now()
    )

    tool_calls = [
        agentToolCall(
            "function", "call-par-1", "getWeather",
            Dict{String,Any}("city" => "Paris")
        ),
        agentToolCall(
            "function", "call-par-2", "getTime",
            Dict{String,Any}("city" => "Sydney")
        ),
    ]

    config = agentLoopConfig(
        nothing, nothing, "parallel"
    )
    sig = abortSignal(false)

    batch = executeToolCallsParallel(
        ctx, assistant_msg, tool_calls, config, sig, x -> nothing
    )

    @test batch.messages isa Vector{toolResultMessage}
    @test length(batch.messages) == 2
    @test batch.messages[1].toolName == "getWeather"
    @test batch.messages[1].isError == false
    @test batch.messages[2].toolName == "getTime"
    @test batch.messages[2].isError == false
end