fix(hooks): fail-open void dispatch and preserve tool-call audit args
This commit is contained in:
parent
0d2c4f9368
commit
c86904cd61
5 changed files with 246 additions and 5 deletions
|
|
@ -803,9 +803,14 @@ func (al *AgentLoop) runLLMIteration(
|
|||
|
||||
// Save assistant message with tool calls to session
|
||||
agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg)
|
||||
assistantMsgIndex := len(messages) - 1
|
||||
assistantSessionIndex := -1
|
||||
if history := agent.Sessions.GetHistory(opts.SessionKey); len(history) > 0 {
|
||||
assistantSessionIndex = len(history) - 1
|
||||
}
|
||||
|
||||
// Execute tool calls
|
||||
for _, tc := range normalizedToolCalls {
|
||||
for tcIdx, tc := range normalizedToolCalls {
|
||||
argsJSON, _ := json.Marshal(tc.Arguments)
|
||||
argsPreview := utils.Truncate(string(argsJSON), 200)
|
||||
logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview),
|
||||
|
|
@ -858,12 +863,29 @@ func (al *AgentLoop) runLLMIteration(
|
|||
if tc.Arguments == nil {
|
||||
tc.Arguments = make(map[string]any)
|
||||
}
|
||||
|
||||
// Keep persisted assistant tool-call arguments aligned with rewritten execution args.
|
||||
updateToolCallArguments(&messages[assistantMsgIndex], tcIdx, tc.Arguments)
|
||||
if assistantSessionIndex >= 0 {
|
||||
history := agent.Sessions.GetHistory(opts.SessionKey)
|
||||
if assistantSessionIndex < len(history) {
|
||||
updateToolCallArguments(&history[assistantSessionIndex], tcIdx, tc.Arguments)
|
||||
agent.Sessions.SetHistory(opts.SessionKey, history)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var toolDuration time.Duration
|
||||
if !toolCanceled {
|
||||
toolStart := time.Now()
|
||||
toolResult = agent.Tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, opts.Channel, opts.ChatID, asyncCallback)
|
||||
toolResult = agent.Tools.ExecuteWithContext(
|
||||
ctx,
|
||||
tc.Name,
|
||||
tc.Arguments,
|
||||
opts.Channel,
|
||||
opts.ChatID,
|
||||
asyncCallback,
|
||||
)
|
||||
toolDuration = time.Since(toolStart)
|
||||
}
|
||||
|
||||
|
|
@ -1040,6 +1062,19 @@ func (al *AgentLoop) GetStartupInfo() map[string]any {
|
|||
return info
|
||||
}
|
||||
|
||||
// updateToolCallArguments patches the serialized arguments for a tool call in-place.
|
||||
func updateToolCallArguments(msg *providers.Message, toolCallIndex int, args map[string]any) {
|
||||
if msg == nil || toolCallIndex < 0 || toolCallIndex >= len(msg.ToolCalls) {
|
||||
return
|
||||
}
|
||||
toolCall := &msg.ToolCalls[toolCallIndex]
|
||||
if toolCall.Function == nil {
|
||||
return
|
||||
}
|
||||
argumentsJSON, _ := json.Marshal(args)
|
||||
toolCall.Function.Arguments = string(argumentsJSON)
|
||||
}
|
||||
|
||||
// formatMessagesForLog formats messages for logging
|
||||
func formatMessagesForLog(messages []providers.Message) string {
|
||||
if len(messages) == 0 {
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ package agent
|
|||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
|
@ -340,3 +341,76 @@ func TestSetHooksNilRestoresDirectMessageCallback(t *testing.T) {
|
|||
t.Fatalf("unexpected outbound message: %#v", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBeforeToolCallArgRewriteUpdatesAssistantTranscript(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "agent-plugin-test-*")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
Model: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
msgBus := bus.NewMessageBus()
|
||||
provider := &nilArgsProvider{}
|
||||
al := NewAgentLoop(cfg, msgBus, provider)
|
||||
al.RegisterTool(&nilArgsCaptureTool{})
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
t.Fatal("expected default agent")
|
||||
}
|
||||
sessionKey := "agent:" + defaultAgent.ID + ":s2"
|
||||
|
||||
reg := hooks.NewHookRegistry()
|
||||
reg.OnBeforeToolCall("rewrite-args", 0, func(_ context.Context, e *hooks.BeforeToolCallEvent) error {
|
||||
e.Args["rewritten"] = "yes"
|
||||
return nil
|
||||
})
|
||||
if err := al.SetHooks(reg); err != nil {
|
||||
t.Fatalf("SetHooks: %v", err)
|
||||
}
|
||||
|
||||
if _, err := al.ProcessDirectWithChannel(
|
||||
context.Background(),
|
||||
"run rewrite test",
|
||||
sessionKey,
|
||||
"cli",
|
||||
"direct",
|
||||
); err != nil {
|
||||
t.Fatalf("ProcessDirectWithChannel: %v", err)
|
||||
}
|
||||
|
||||
history := defaultAgent.Sessions.GetHistory(sessionKey)
|
||||
|
||||
foundToolCall := false
|
||||
for _, msg := range history {
|
||||
if msg.Role != "assistant" || len(msg.ToolCalls) == 0 {
|
||||
continue
|
||||
}
|
||||
if msg.ToolCalls[0].Function == nil {
|
||||
t.Fatal("expected tool call function payload")
|
||||
}
|
||||
var args map[string]any
|
||||
if err := json.Unmarshal([]byte(msg.ToolCalls[0].Function.Arguments), &args); err != nil {
|
||||
t.Fatalf("failed to decode persisted tool call args: %v", err)
|
||||
}
|
||||
if got := args["rewritten"]; got != "yes" {
|
||||
t.Fatalf("expected rewritten arg to be persisted, got %#v", got)
|
||||
}
|
||||
foundToolCall = true
|
||||
break
|
||||
}
|
||||
|
||||
if !foundToolCall {
|
||||
t.Fatal("expected assistant tool call message in session history")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,11 +11,14 @@ import (
|
|||
"fmt"
|
||||
"reflect"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
const voidHookWaitBudget = 50 * time.Millisecond
|
||||
|
||||
// HookHandler is the callback signature for all hooks.
|
||||
type HookHandler[T any] func(ctx context.Context, event *T) error
|
||||
|
||||
|
|
@ -204,6 +207,17 @@ func cloneReflectValue(v reflect.Value) reflect.Value {
|
|||
out.Index(i).Set(cloneReflectValue(v.Index(i)))
|
||||
}
|
||||
return out
|
||||
case reflect.Struct:
|
||||
out := reflect.New(v.Type()).Elem()
|
||||
for i := range v.NumField() {
|
||||
field := out.Field(i)
|
||||
if !field.CanSet() {
|
||||
// Preserve original value for structs with non-settable fields.
|
||||
return v
|
||||
}
|
||||
field.Set(cloneReflectValue(v.Field(i)))
|
||||
}
|
||||
return out
|
||||
default:
|
||||
return v
|
||||
}
|
||||
|
|
@ -310,7 +324,9 @@ func cloneVoidEvent[T any](event *T) *T {
|
|||
}
|
||||
}
|
||||
|
||||
// triggerVoid runs all handlers concurrently and waits for completion.
|
||||
// triggerVoid runs all handlers concurrently.
|
||||
// It waits up to a small budget to collect immediate completions, then
|
||||
// continues fail-open to avoid blocking the core agent pipeline.
|
||||
// Each handler receives a cloned event to avoid shared-state mutation races.
|
||||
// Errors are logged but do not propagate to the caller.
|
||||
func triggerVoid[T any](ctx context.Context, hooks []HookRegistration[T], event *T, hookName string) {
|
||||
|
|
@ -343,7 +359,27 @@ func triggerVoid[T any](ctx context.Context, hooks []HookRegistration[T], event
|
|||
}
|
||||
}(h)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-ctx.Done():
|
||||
logger.WarnCF("hooks", "Void hook dispatch interrupted by context",
|
||||
map[string]any{
|
||||
"hook": hookName,
|
||||
})
|
||||
case <-time.After(voidHookWaitBudget):
|
||||
logger.WarnCF("hooks", "Void hook dispatch exceeded wait budget; continuing",
|
||||
map[string]any{
|
||||
"hook": hookName,
|
||||
"wait_budget_ms": voidHookWaitBudget.Milliseconds(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// triggerModifying runs handlers sequentially by priority, stopping if Cancel is set.
|
||||
|
|
|
|||
|
|
@ -226,6 +226,96 @@ func TestVoidHooksReceiveIsolatedLLMInputToolSchema(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestVoidHooksReceiveIsolatedStructValuesInMap(t *testing.T) {
|
||||
type schemaSpec struct {
|
||||
Required []string
|
||||
Meta map[string]string
|
||||
}
|
||||
|
||||
r := NewHookRegistry()
|
||||
ctx := context.Background()
|
||||
|
||||
r.OnLLMInput("struct-mutator", 0, func(_ context.Context, e *LLMInputEvent) error {
|
||||
spec, ok := e.Tools[0].Function.Parameters["schema"].(schemaSpec)
|
||||
if !ok {
|
||||
t.Fatal("schema should be schemaSpec")
|
||||
}
|
||||
spec.Required[0] = "mutated"
|
||||
spec.Meta["k"] = "changed"
|
||||
e.Tools[0].Function.Parameters["schema"] = spec
|
||||
return nil
|
||||
})
|
||||
|
||||
event := &LLMInputEvent{
|
||||
AgentID: "a1",
|
||||
Model: "m1",
|
||||
Tools: []providers.ToolDefinition{
|
||||
{
|
||||
Type: "function",
|
||||
Function: providers.ToolFunctionDefinition{
|
||||
Name: "message",
|
||||
Parameters: map[string]any{
|
||||
"schema": schemaSpec{
|
||||
Required: []string{"content"},
|
||||
Meta: map[string]string{"k": "v"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
r.TriggerLLMInput(ctx, event)
|
||||
|
||||
spec, ok := event.Tools[0].Function.Parameters["schema"].(schemaSpec)
|
||||
if !ok {
|
||||
t.Fatal("schema should remain schemaSpec")
|
||||
}
|
||||
if len(spec.Required) != 1 || spec.Required[0] != "content" {
|
||||
t.Fatalf("expected required to remain unchanged, got %#v", spec.Required)
|
||||
}
|
||||
if got := spec.Meta["k"]; got != "v" {
|
||||
t.Fatalf("expected meta[k] to remain v, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVoidHooksFailOpenOnSlowHandler(t *testing.T) {
|
||||
r := NewHookRegistry()
|
||||
ctx := context.Background()
|
||||
|
||||
started := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
done := make(chan struct{})
|
||||
|
||||
r.OnLLMInput("slow", 0, func(_ context.Context, _ *LLMInputEvent) error {
|
||||
close(started)
|
||||
<-release
|
||||
close(done)
|
||||
return nil
|
||||
})
|
||||
|
||||
begin := time.Now()
|
||||
r.TriggerLLMInput(ctx, &LLMInputEvent{AgentID: "a1"})
|
||||
elapsed := time.Since(begin)
|
||||
|
||||
if elapsed > voidHookWaitBudget*3 {
|
||||
t.Fatalf("expected fail-open dispatch within budget, got %s", elapsed)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(1 * time.Second):
|
||||
t.Fatal("timeout waiting for slow handler to start")
|
||||
}
|
||||
|
||||
close(release)
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(1 * time.Second):
|
||||
t.Fatal("timeout waiting for slow handler to finish after release")
|
||||
}
|
||||
}
|
||||
|
||||
func TestModifyingHookPriority(t *testing.T) {
|
||||
r := NewHookRegistry()
|
||||
ctx := context.Background()
|
||||
|
|
|
|||
|
|
@ -141,7 +141,13 @@ func TestPolicyDemoPluginAuditHooks(t *testing.T) {
|
|||
}
|
||||
|
||||
pm.HookRegistry().TriggerSessionStart(context.Background(), &hooks.SessionEvent{AgentID: "a1", SessionKey: "s1"})
|
||||
pm.HookRegistry().TriggerAfterToolCall(context.Background(), &hooks.AfterToolCallEvent{ToolName: "web_search", Duration: 45 * time.Millisecond})
|
||||
pm.HookRegistry().TriggerAfterToolCall(
|
||||
context.Background(),
|
||||
&hooks.AfterToolCallEvent{
|
||||
ToolName: "web_search",
|
||||
Duration: 45 * time.Millisecond,
|
||||
},
|
||||
)
|
||||
pm.HookRegistry().TriggerSessionEnd(context.Background(), &hooks.SessionEvent{AgentID: "a1", SessionKey: "s1"})
|
||||
|
||||
stats := p.Snapshot()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue