test: add tests for board wiring, recursion guards, allowlist, and hooks
Cover SetBoard switching, BoardAware interface compliance, depth limit, cycle detection, depth propagation, allowlist block/permit/default-open, AllowlistCheckerFunc adapter, and ToolHook before/after/block/chain behavior.
This commit is contained in:
parent
42d73bde4d
commit
284f17e05f
3 changed files with 409 additions and 0 deletions
|
|
@ -286,6 +286,57 @@ func TestBlackboardTool_MissingKey(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestBlackboardTool_SetBoard(t *testing.T) {
|
||||||
|
bb1 := NewBlackboard()
|
||||||
|
bb2 := NewBlackboard()
|
||||||
|
bb2.Set("from_session", "session_data", "system")
|
||||||
|
|
||||||
|
tool := NewBlackboardTool(bb1, "agent1")
|
||||||
|
|
||||||
|
// Initially reads from bb1 (empty)
|
||||||
|
result := tool.Execute(context.Background(), map[string]any{
|
||||||
|
"action": "read",
|
||||||
|
"key": "from_session",
|
||||||
|
})
|
||||||
|
if !contains(result.ForLLM, "No entry") {
|
||||||
|
t.Errorf("expected 'No entry' before SetBoard, got %q", result.ForLLM)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Switch to session board
|
||||||
|
tool.SetBoard(bb2)
|
||||||
|
|
||||||
|
// Now reads from bb2
|
||||||
|
result = tool.Execute(context.Background(), map[string]any{
|
||||||
|
"action": "read",
|
||||||
|
"key": "from_session",
|
||||||
|
})
|
||||||
|
if !contains(result.ForLLM, "session_data") {
|
||||||
|
t.Errorf("expected 'session_data' after SetBoard, got %q", result.ForLLM)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Writes go to bb2, not bb1
|
||||||
|
tool.Execute(context.Background(), map[string]any{
|
||||||
|
"action": "write",
|
||||||
|
"key": "new_key",
|
||||||
|
"value": "new_val",
|
||||||
|
})
|
||||||
|
if bb1.Get("new_key") != "" {
|
||||||
|
t.Error("write went to old board after SetBoard")
|
||||||
|
}
|
||||||
|
if bb2.Get("new_key") != "new_val" {
|
||||||
|
t.Error("write didn't go to new board after SetBoard")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBoardAware_Interface(t *testing.T) {
|
||||||
|
// Verify both tools implement BoardAware
|
||||||
|
bb := NewBlackboard()
|
||||||
|
var _ BoardAware = NewBlackboardTool(bb, "test")
|
||||||
|
|
||||||
|
resolver := newMockResolver()
|
||||||
|
var _ BoardAware = NewHandoffTool(resolver, bb, "test")
|
||||||
|
}
|
||||||
|
|
||||||
func contains(s, sub string) bool {
|
func contains(s, sub string) bool {
|
||||||
return len(s) >= len(sub) && (s == sub || len(s) > 0 && containsStr(s, sub))
|
return len(s) >= len(sub) && (s == sub || len(s) > 0 && containsStr(s, sub))
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -355,6 +355,251 @@ func TestAgentInfo_Capabilities(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestExecuteHandoff_DepthLimit(t *testing.T) {
|
||||||
|
provider := &mockProvider{response: "done"}
|
||||||
|
resolver := newMockResolver(&AgentInfo{
|
||||||
|
ID: "target",
|
||||||
|
Name: "Target",
|
||||||
|
Model: "test",
|
||||||
|
Provider: provider,
|
||||||
|
Tools: tools.NewToolRegistry(),
|
||||||
|
MaxIter: 5,
|
||||||
|
})
|
||||||
|
|
||||||
|
bb := NewBlackboard()
|
||||||
|
result := ExecuteHandoff(context.Background(), resolver, bb, HandoffRequest{
|
||||||
|
FromAgentID: "main",
|
||||||
|
ToAgentID: "target",
|
||||||
|
Task: "do something",
|
||||||
|
Depth: 3, // at max depth
|
||||||
|
MaxDepth: 3,
|
||||||
|
Visited: []string{"main", "agent-a", "agent-b"},
|
||||||
|
}, "cli", "direct")
|
||||||
|
|
||||||
|
if result.Success {
|
||||||
|
t.Error("expected failure at max depth")
|
||||||
|
}
|
||||||
|
if !strings.Contains(result.Error, "depth limit") {
|
||||||
|
t.Errorf("Error = %q, expected 'depth limit'", result.Error)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestExecuteHandoff_CycleDetection(t *testing.T) {
|
||||||
|
provider := &mockProvider{response: "done"}
|
||||||
|
resolver := newMockResolver(
|
||||||
|
&AgentInfo{ID: "main", Name: "Main", Model: "test", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
&AgentInfo{ID: "coder", Name: "Coder", Model: "test", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
)
|
||||||
|
|
||||||
|
bb := NewBlackboard()
|
||||||
|
|
||||||
|
// Try to hand off to "main" which is already in the visited chain
|
||||||
|
result := ExecuteHandoff(context.Background(), resolver, bb, HandoffRequest{
|
||||||
|
FromAgentID: "coder",
|
||||||
|
ToAgentID: "main",
|
||||||
|
Task: "some task",
|
||||||
|
Depth: 1,
|
||||||
|
Visited: []string{"main", "coder"},
|
||||||
|
}, "cli", "direct")
|
||||||
|
|
||||||
|
if result.Success {
|
||||||
|
t.Error("expected failure due to cycle detection")
|
||||||
|
}
|
||||||
|
if !strings.Contains(result.Error, "cycle detected") {
|
||||||
|
t.Errorf("Error = %q, expected 'cycle detected'", result.Error)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestExecuteHandoff_DefaultMaxDepth(t *testing.T) {
|
||||||
|
provider := &mockProvider{response: "done"}
|
||||||
|
resolver := newMockResolver(&AgentInfo{
|
||||||
|
ID: "target", Name: "Target", Model: "test",
|
||||||
|
Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5,
|
||||||
|
})
|
||||||
|
|
||||||
|
bb := NewBlackboard()
|
||||||
|
|
||||||
|
// Depth 2 with default max (3) should succeed
|
||||||
|
result := ExecuteHandoff(context.Background(), resolver, bb, HandoffRequest{
|
||||||
|
FromAgentID: "main",
|
||||||
|
ToAgentID: "target",
|
||||||
|
Task: "do something",
|
||||||
|
Depth: 2,
|
||||||
|
Visited: []string{"main", "middle"},
|
||||||
|
}, "cli", "direct")
|
||||||
|
if !result.Success {
|
||||||
|
t.Fatalf("expected success at depth 2 (max 3), got error: %s", result.Error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Depth 3 with default max should fail
|
||||||
|
result = ExecuteHandoff(context.Background(), resolver, bb, HandoffRequest{
|
||||||
|
FromAgentID: "main",
|
||||||
|
ToAgentID: "target",
|
||||||
|
Task: "do something",
|
||||||
|
Depth: 3,
|
||||||
|
Visited: []string{"main", "a", "b"},
|
||||||
|
}, "cli", "direct")
|
||||||
|
if result.Success {
|
||||||
|
t.Error("expected failure at depth 3 with default max 3")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestExecuteHandoff_PropagatesDepthToTarget(t *testing.T) {
|
||||||
|
provider := &mockProvider{response: "done"}
|
||||||
|
targetRegistry := tools.NewToolRegistry()
|
||||||
|
innerResolver := newMockResolver()
|
||||||
|
targetHandoff := NewHandoffTool(innerResolver, NewBlackboard(), "target")
|
||||||
|
targetRegistry.Register(targetHandoff)
|
||||||
|
|
||||||
|
resolver := newMockResolver(&AgentInfo{
|
||||||
|
ID: "target", Name: "Target", Model: "test",
|
||||||
|
Provider: provider, Tools: targetRegistry, MaxIter: 5,
|
||||||
|
})
|
||||||
|
|
||||||
|
bb := NewBlackboard()
|
||||||
|
result := ExecuteHandoff(context.Background(), resolver, bb, HandoffRequest{
|
||||||
|
FromAgentID: "main",
|
||||||
|
ToAgentID: "target",
|
||||||
|
Task: "do something",
|
||||||
|
Depth: 1,
|
||||||
|
Visited: []string{"main"},
|
||||||
|
MaxDepth: 5,
|
||||||
|
}, "cli", "direct")
|
||||||
|
|
||||||
|
if !result.Success {
|
||||||
|
t.Fatalf("expected success, got error: %s", result.Error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify the target's handoff tool got the propagated depth
|
||||||
|
if targetHandoff.depth != 2 {
|
||||||
|
t.Errorf("target handoff depth = %d, want 2", targetHandoff.depth)
|
||||||
|
}
|
||||||
|
if len(targetHandoff.visited) != 2 || targetHandoff.visited[0] != "main" || targetHandoff.visited[1] != "target" {
|
||||||
|
t.Errorf("target handoff visited = %v, want [main target]", targetHandoff.visited)
|
||||||
|
}
|
||||||
|
if targetHandoff.maxDepth != 5 {
|
||||||
|
t.Errorf("target handoff maxDepth = %d, want 5", targetHandoff.maxDepth)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandoffTool_AllowlistBlocks(t *testing.T) {
|
||||||
|
provider := &mockProvider{response: "done"}
|
||||||
|
resolver := newMockResolver(
|
||||||
|
&AgentInfo{ID: "main", Name: "Main", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
&AgentInfo{ID: "restricted", Name: "Restricted", Model: "test", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
)
|
||||||
|
|
||||||
|
bb := NewBlackboard()
|
||||||
|
tool := NewHandoffTool(resolver, bb, "main")
|
||||||
|
tool.SetAllowlistChecker(AllowlistCheckerFunc(func(from, to string) bool {
|
||||||
|
return to == "allowed-agent" // only allow "allowed-agent"
|
||||||
|
}))
|
||||||
|
|
||||||
|
result := tool.Execute(context.Background(), map[string]any{
|
||||||
|
"agent_id": "restricted",
|
||||||
|
"task": "do something",
|
||||||
|
})
|
||||||
|
if !result.IsError {
|
||||||
|
t.Error("expected error for blocked handoff")
|
||||||
|
}
|
||||||
|
if !strings.Contains(result.ForLLM, "not allowed") {
|
||||||
|
t.Errorf("ForLLM = %q, expected 'not allowed'", result.ForLLM)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandoffTool_AllowlistPermits(t *testing.T) {
|
||||||
|
provider := &mockProvider{response: "allowed result"}
|
||||||
|
resolver := newMockResolver(
|
||||||
|
&AgentInfo{ID: "main", Name: "Main", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
&AgentInfo{ID: "coder", Name: "Coder", Model: "test", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
)
|
||||||
|
|
||||||
|
bb := NewBlackboard()
|
||||||
|
tool := NewHandoffTool(resolver, bb, "main")
|
||||||
|
tool.SetAllowlistChecker(AllowlistCheckerFunc(func(from, to string) bool {
|
||||||
|
return to == "coder" // allow coder
|
||||||
|
}))
|
||||||
|
|
||||||
|
result := tool.Execute(context.Background(), map[string]any{
|
||||||
|
"agent_id": "coder",
|
||||||
|
"task": "write code",
|
||||||
|
})
|
||||||
|
if result.IsError {
|
||||||
|
t.Fatalf("expected success, got error: %s", result.ForLLM)
|
||||||
|
}
|
||||||
|
if !strings.Contains(result.ForLLM, "allowed result") {
|
||||||
|
t.Errorf("ForLLM = %q, expected 'allowed result'", result.ForLLM)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandoffTool_NoAllowlistAllowsAll(t *testing.T) {
|
||||||
|
provider := &mockProvider{response: "ok"}
|
||||||
|
resolver := newMockResolver(
|
||||||
|
&AgentInfo{ID: "main", Name: "Main", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
&AgentInfo{ID: "any", Name: "Any", Model: "test", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
)
|
||||||
|
|
||||||
|
bb := NewBlackboard()
|
||||||
|
tool := NewHandoffTool(resolver, bb, "main")
|
||||||
|
// No allowlist checker set
|
||||||
|
|
||||||
|
result := tool.Execute(context.Background(), map[string]any{
|
||||||
|
"agent_id": "any",
|
||||||
|
"task": "anything",
|
||||||
|
})
|
||||||
|
if result.IsError {
|
||||||
|
t.Fatalf("expected success with no allowlist, got: %s", result.ForLLM)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandoffTool_SetBoard(t *testing.T) {
|
||||||
|
provider := &mockProvider{response: "done"}
|
||||||
|
resolver := newMockResolver(
|
||||||
|
&AgentInfo{ID: "main", Name: "Main", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
&AgentInfo{ID: "coder", Name: "Coder", Model: "test", Provider: provider, Tools: tools.NewToolRegistry(), MaxIter: 5},
|
||||||
|
)
|
||||||
|
|
||||||
|
bb1 := NewBlackboard()
|
||||||
|
bb2 := NewBlackboard()
|
||||||
|
bb2.Set("session_data", "hello", "system")
|
||||||
|
|
||||||
|
tool := NewHandoffTool(resolver, bb1, "main")
|
||||||
|
|
||||||
|
// Switch to session board
|
||||||
|
tool.SetBoard(bb2)
|
||||||
|
|
||||||
|
// Execute with context that writes to blackboard
|
||||||
|
tool.Execute(context.Background(), map[string]any{
|
||||||
|
"agent_id": "coder",
|
||||||
|
"task": "write code",
|
||||||
|
"context": map[string]any{"language": "Go"},
|
||||||
|
})
|
||||||
|
|
||||||
|
// Context should have been written to bb2 (session board), not bb1
|
||||||
|
if bb1.Get("language") != "" {
|
||||||
|
t.Error("context was written to old board")
|
||||||
|
}
|
||||||
|
if bb2.Get("language") != "Go" {
|
||||||
|
t.Errorf("context not written to session board: %q", bb2.Get("language"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAllowlistCheckerFunc(t *testing.T) {
|
||||||
|
checker := AllowlistCheckerFunc(func(from, to string) bool {
|
||||||
|
return from == "main" && to == "coder"
|
||||||
|
})
|
||||||
|
|
||||||
|
if !checker.CanHandoff("main", "coder") {
|
||||||
|
t.Error("expected main->coder to be allowed")
|
||||||
|
}
|
||||||
|
if checker.CanHandoff("main", "other") {
|
||||||
|
t.Error("expected main->other to be blocked")
|
||||||
|
}
|
||||||
|
if checker.CanHandoff("other", "coder") {
|
||||||
|
t.Error("expected other->coder to be blocked")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestBuildHandoffSystemPrompt(t *testing.T) {
|
func TestBuildHandoffSystemPrompt(t *testing.T) {
|
||||||
agent := &AgentInfo{
|
agent := &AgentInfo{
|
||||||
Name: "Code Agent",
|
Name: "Code Agent",
|
||||||
|
|
|
||||||
113
pkg/tools/hooks_test.go
Normal file
113
pkg/tools/hooks_test.go
Normal file
|
|
@ -0,0 +1,113 @@
|
||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
// testHook records calls and optionally blocks execution.
|
||||||
|
type testHook struct {
|
||||||
|
beforeCalls []string
|
||||||
|
afterCalls []string
|
||||||
|
blockTool string // if non-empty, block this tool name
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *testHook) BeforeExecute(_ context.Context, toolName string, _ map[string]interface{}) error {
|
||||||
|
h.beforeCalls = append(h.beforeCalls, toolName)
|
||||||
|
if h.blockTool != "" && toolName == h.blockTool {
|
||||||
|
return errors.New("blocked by test hook")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *testHook) AfterExecute(_ context.Context, toolName string, _ map[string]interface{}, _ *ToolResult) {
|
||||||
|
h.afterCalls = append(h.afterCalls, toolName)
|
||||||
|
}
|
||||||
|
|
||||||
|
// dummyTool is a minimal tool for hook testing.
|
||||||
|
type dummyTool struct {
|
||||||
|
name string
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *dummyTool) Name() string { return d.name }
|
||||||
|
func (d *dummyTool) Description() string { return "test tool" }
|
||||||
|
func (d *dummyTool) Parameters() map[string]interface{} { return nil }
|
||||||
|
func (d *dummyTool) Execute(_ context.Context, _ map[string]interface{}) *ToolResult { return NewToolResult("ok") }
|
||||||
|
|
||||||
|
func TestToolHook_BeforeAndAfterCalled(t *testing.T) {
|
||||||
|
reg := NewToolRegistry()
|
||||||
|
reg.Register(&dummyTool{name: "test_tool"})
|
||||||
|
|
||||||
|
hook := &testHook{}
|
||||||
|
reg.AddHook(hook)
|
||||||
|
|
||||||
|
reg.Execute(context.Background(), "test_tool", nil)
|
||||||
|
|
||||||
|
if len(hook.beforeCalls) != 1 || hook.beforeCalls[0] != "test_tool" {
|
||||||
|
t.Errorf("beforeCalls = %v, want [test_tool]", hook.beforeCalls)
|
||||||
|
}
|
||||||
|
if len(hook.afterCalls) != 1 || hook.afterCalls[0] != "test_tool" {
|
||||||
|
t.Errorf("afterCalls = %v, want [test_tool]", hook.afterCalls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolHook_BlocksExecution(t *testing.T) {
|
||||||
|
reg := NewToolRegistry()
|
||||||
|
reg.Register(&dummyTool{name: "blocked_tool"})
|
||||||
|
|
||||||
|
hook := &testHook{blockTool: "blocked_tool"}
|
||||||
|
reg.AddHook(hook)
|
||||||
|
|
||||||
|
result := reg.Execute(context.Background(), "blocked_tool", nil)
|
||||||
|
|
||||||
|
if !result.IsError {
|
||||||
|
t.Error("expected error result when hook blocks")
|
||||||
|
}
|
||||||
|
if len(hook.beforeCalls) != 1 {
|
||||||
|
t.Errorf("beforeCalls count = %d, want 1", len(hook.beforeCalls))
|
||||||
|
}
|
||||||
|
// AfterExecute should still be called (for observability)
|
||||||
|
if len(hook.afterCalls) != 1 {
|
||||||
|
t.Errorf("afterCalls count = %d, want 1 (observability)", len(hook.afterCalls))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolHook_MultipleHooks(t *testing.T) {
|
||||||
|
reg := NewToolRegistry()
|
||||||
|
reg.Register(&dummyTool{name: "multi"})
|
||||||
|
|
||||||
|
hook1 := &testHook{}
|
||||||
|
hook2 := &testHook{}
|
||||||
|
reg.AddHook(hook1)
|
||||||
|
reg.AddHook(hook2)
|
||||||
|
|
||||||
|
reg.Execute(context.Background(), "multi", nil)
|
||||||
|
|
||||||
|
if len(hook1.beforeCalls) != 1 || len(hook2.beforeCalls) != 1 {
|
||||||
|
t.Error("expected both hooks to be called")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolHook_FirstBlockStopsChain(t *testing.T) {
|
||||||
|
reg := NewToolRegistry()
|
||||||
|
reg.Register(&dummyTool{name: "chain_test"})
|
||||||
|
|
||||||
|
hook1 := &testHook{blockTool: "chain_test"}
|
||||||
|
hook2 := &testHook{}
|
||||||
|
reg.AddHook(hook1)
|
||||||
|
reg.AddHook(hook2)
|
||||||
|
|
||||||
|
result := reg.Execute(context.Background(), "chain_test", nil)
|
||||||
|
|
||||||
|
if !result.IsError {
|
||||||
|
t.Error("expected error when first hook blocks")
|
||||||
|
}
|
||||||
|
// hook1 should have been called, hook2's Before should NOT
|
||||||
|
if len(hook1.beforeCalls) != 1 {
|
||||||
|
t.Error("hook1 before should have been called")
|
||||||
|
}
|
||||||
|
if len(hook2.beforeCalls) != 0 {
|
||||||
|
t.Error("hook2 before should NOT have been called (chain stopped)")
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue