#include <atomic>
#include <chrono>
#include <thread>
#include <agent_smith/detail/tool_dispatch.hpp>
#include <agent_smith/detail/timeout.hpp>
#include <agent_smith/sync_wait.hpp>
#include <gtest/gtest.h>
namespace agent_smith
{
namespace
{
using namespace std::chrono_literals;
ToolSchema testSchema(std::string name)
{
return {std::move(name), "dispatch test",
{{"type", "object"}, {"properties", nlohmann::json::object()}}};
}
mw::Task<mw::E<ToolResult>> rendezvous(
std::atomic<int>& started, nlohmann::json args, ToolContext context)
{
++started;
while(started.load() < 2 && !context.cancel.cancelled())
{
co_await context.runtime->sleepFor(1ms);
}
if(context.cancel.cancelled())
{
co_return std::unexpected(mw::runtimeError("not concurrent"));
}
co_await context.runtime->sleepFor(
std::chrono::milliseconds(args.at("delay").get<int>()));
co_return ToolResult{.content = context.call_id};
}
TEST(ToolDispatchTest, StartsAllCallsAndPreservesInputOrder)
{
Runtime runtime;
ToolRegistry registry;
std::atomic<int> started{0};
registry.add(makeTool(testSchema("parallel"),
[&started](nlohmann::json args, ToolContext context)
{ return rendezvous(started, std::move(args), std::move(context)); }));
auto result = syncWait(internal::dispatchTools(runtime, registry,
{{"first", "parallel", {{"delay", 30}}},
{"second", "parallel", {{"delay", 0}}}},
{.default_tool_timeout = 1s}));
ASSERT_TRUE(result);
ASSERT_EQ(result->size(), 2U);
EXPECT_FALSE((*result)[0].is_error);
EXPECT_FALSE((*result)[1].is_error);
EXPECT_EQ((*result)[0].content, "first");
EXPECT_EQ((*result)[1].content, "second");
}
TEST(ToolDispatchTest, ConvertsFailuresAndTruncatesResults)
{
Runtime runtime;
ToolRegistry registry;
registry.add(makeTool(testSchema("failure"),
[](nlohmann::json, ToolContext) -> mw::E<ToolResult>
{ return std::unexpected(mw::runtimeError("failed")); }));
registry.add(makeTool(testSchema("reported"),
[](nlohmann::json, ToolContext)
{ return ToolResult{.content = "reported failure", .is_error = true}; }));
registry.add(makeTool(testSchema("large"),
[](nlohmann::json, ToolContext) { return std::string(100, 'x'); }));
auto result = syncWait(internal::dispatchTools(runtime, registry,
{{"a", "missing", {}}, {"b", "failure", {}},
{"c", "reported", {}}, {"d", "large", {}}},
{.tool_result_max_bytes = 64}));
ASSERT_TRUE(result);
ASSERT_EQ(result->size(), 4U);
for(int index = 0; index < 3; ++index)
{
EXPECT_TRUE((*result)[index].is_error);
EXPECT_TRUE(nlohmann::json::parse((*result)[index].content).at("error"));
}
EXPECT_EQ(nlohmann::json::parse((*result)[0].content).at("message"),
"unknown tool: missing");
EXPECT_EQ((*result)[3].content,
std::string(64, 'x') + "\n…[truncated, original size 100 bytes]");
}
mw::Task<mw::E<ToolResult>> waitForCancellation(
std::atomic<bool>& cancelled, nlohmann::json, ToolContext context)
{
while(!context.cancel.cancelled())
{
co_await context.runtime->sleepFor(1ms);
}
cancelled = true;
co_return ToolResult{.content = "late"};
}
TEST(ToolDispatchTest, AppliesOverrideAndSignalsCancellation)
{
std::atomic<bool> cancelled{false};
{
Runtime runtime;
ToolRegistry registry;
registry.add(makeTool(testSchema("slow"),
[&cancelled](nlohmann::json args, ToolContext context)
{ return waitForCancellation(cancelled, std::move(args),
std::move(context)); }),
{.timeout = 20ms});
const auto start = std::chrono::steady_clock::now();
auto result = syncWait(internal::dispatchTools(runtime, registry,
{{"call", "slow", {}}}, {.default_tool_timeout = 5s}));
ASSERT_TRUE(result);
ASSERT_EQ(result->size(), 1U);
EXPECT_TRUE(result->front().is_error);
EXPECT_EQ(result->front().tool_call_id, "call");
EXPECT_LT(std::chrono::steady_clock::now() - start, 1s);
}
EXPECT_TRUE(cancelled);
}
mw::Task<mw::E<int>> delayedValue(Runtime& runtime)
{
co_await runtime.sleepFor(40ms);
co_return 42;
}
TEST(TimeoutTest, ReturnsTypedErrorAndRetainsLateOperation)
{
Runtime runtime;
CancellationSource cancel;
auto result = syncWait(internal::withTimeout(
runtime, delayedValue(runtime), 5ms, &cancel));
ASSERT_FALSE(result);
ASSERT_NE(result.error().as<TimeoutError>(), nullptr);
EXPECT_EQ(result.error().as<TimeoutError>()->budget, 5ms);
EXPECT_TRUE(cancel.token().cancelled());
}
TEST(TimeoutTest, CancelsLosingTimerPromptly)
{
const auto start = std::chrono::steady_clock::now();
{
Runtime runtime;
auto result = syncWait(internal::withTimeout(
runtime, delayedValue(runtime), 30s));
ASSERT_TRUE(result);
EXPECT_EQ(*result, 42);
}
EXPECT_LT(std::chrono::steady_clock::now() - start, 2s);
}
TEST(ToolDispatchTest, SyncCallsRunConcurrentlyAndUseDefaultBudget)
{
Runtime runtime;
ToolRegistry registry;
std::atomic<int> started{0};
registry.add(makeTool(testSchema("sync"),
[&started](nlohmann::json, ToolContext context) -> ToolResult
{
++started;
while(started.load() < 2 && !context.cancel.cancelled())
{
std::this_thread::yield();
}
return {.content = std::to_string(context.budget.count()),
.is_error = context.cancel.cancelled()};
}));
auto result = syncWait(internal::dispatchTools(runtime, registry,
{{"a", "sync", {}}, {"b", "sync", {}}},
{.default_tool_timeout = 500ms}));
ASSERT_TRUE(result);
ASSERT_EQ(result->size(), 2U);
EXPECT_FALSE((*result)[0].is_error);
EXPECT_FALSE((*result)[1].is_error);
EXPECT_EQ((*result)[0].content, "500");
}
class LateTool final : public Tool
{
public:
LateTool(std::atomic<bool>& destroyed, std::atomic<bool>& finished)
: destroyed_(destroyed), finished_(finished), schema_(testSchema("late"))
{}
~LateTool() override { destroyed_ = true; }
const ToolSchema& schema() const override { return schema_; }
mw::Task<mw::E<ToolResult>> invoke(nlohmann::json,
ToolContext context) override
{
co_await context.runtime->sleepFor(80ms);
finished_ = true;
co_return ToolResult{.content = schema_.name};
}
private:
std::atomic<bool>& destroyed_;
std::atomic<bool>& finished_;
ToolSchema schema_;
};
TEST(ToolDispatchTest, RetainsTimedOutToolAfterRegistryDestruction)
{
std::atomic<bool> destroyed{false};
std::atomic<bool> finished{false};
{
Runtime runtime;
{
ToolRegistry registry;
registry.add(std::make_unique<LateTool>(destroyed, finished));
auto result = syncWait(internal::dispatchTools(runtime, registry,
{{"a", "late", {}}}, {.default_tool_timeout = 5ms}));
ASSERT_TRUE(result);
EXPECT_TRUE(result->front().is_error);
}
EXPECT_FALSE(destroyed);
}
EXPECT_TRUE(finished);
EXPECT_TRUE(destroyed);
}
TEST(ToolDispatchTest, HandlesEmptyBatchAndZeroBudget)
{
Runtime runtime;
ToolRegistry registry;
auto empty = syncWait(internal::dispatchTools(runtime, registry, {}));
ASSERT_TRUE(empty);
EXPECT_TRUE(empty->empty());
std::atomic<bool> invoked{false};
registry.add(makeTool(testSchema("zero"),
[&invoked](nlohmann::json, ToolContext)
{ invoked = true; return "unexpected"; }));
auto result = syncWait(internal::dispatchTools(runtime, registry,
{{"a", "zero", {}}}, {.default_tool_timeout = 0ms}));
ASSERT_TRUE(result);
EXPECT_TRUE(result->front().is_error);
EXPECT_FALSE(invoked);
}
TEST(ToolDispatchTest, TruncationKeepsUtf8Serializable)
{
Runtime runtime;
ToolRegistry registry;
registry.add(makeTool(testSchema("unicode"),
[](nlohmann::json, ToolContext) { return "北京天气"; }));
auto result = syncWait(internal::dispatchTools(runtime, registry,
{{"a", "unicode", {}}}, {.tool_result_max_bytes = 4}));
ASSERT_TRUE(result);
EXPECT_EQ(result->front().content,
"北\n…[truncated, original size 12 bytes]");
EXPECT_NO_THROW(toOpenAIJson(Message(result->front())).dump());
}
TEST(ToolDispatchTest, SyncTimeoutReturnsBeforeNonCooperativeWorkFinishes)
{
std::atomic<bool> finished{false};
{
Runtime runtime;
ToolRegistry registry;
registry.add(makeTool(testSchema("blocking"),
[&finished](nlohmann::json, ToolContext)
{
std::this_thread::sleep_for(200ms);
finished = true;
return "late";
}));
auto result = syncWait(internal::dispatchTools(runtime, registry,
{{"a", "blocking", {}}}, {.default_tool_timeout = 10ms}));
ASSERT_TRUE(result);
EXPECT_TRUE(result->front().is_error);
EXPECT_FALSE(finished);
}
EXPECT_TRUE(finished);
}
} // namespace
} // namespace agent_smith