diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index d62b9268f..3546feef7 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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 { diff --git a/pkg/agent/plugin_test.go b/pkg/agent/plugin_test.go index 115923476..c0b6f0625 100644 --- a/pkg/agent/plugin_test.go +++ b/pkg/agent/plugin_test.go @@ -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") + } +} diff --git a/pkg/hooks/hooks.go b/pkg/hooks/hooks.go index 3a6a0a675..9865c86fe 100644 --- a/pkg/hooks/hooks.go +++ b/pkg/hooks/hooks.go @@ -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. diff --git a/pkg/hooks/hooks_test.go b/pkg/hooks/hooks_test.go index 8ca3e4c2c..d21467a55 100644 --- a/pkg/hooks/hooks_test.go +++ b/pkg/hooks/hooks_test.go @@ -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() diff --git a/pkg/plugin/demoplugin/policy_demo_test.go b/pkg/plugin/demoplugin/policy_demo_test.go index 825c673f0..4d41084f1 100644 --- a/pkg/plugin/demoplugin/policy_demo_test.go +++ b/pkg/plugin/demoplugin/policy_demo_test.go @@ -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()