From 3a68d268371d5328f008bc695c41585100e5549a Mon Sep 17 00:00:00 2001 From: SiYue-ZO <2835601846@qq.com> Date: Thu, 21 May 2026 01:14:42 +0800 Subject: [PATCH] Fix agent loop reload and panic cleanup stability --- pkg/agent/agent.go | 51 +++----- pkg/agent/agent_test.go | 106 +++++++++++++++++ pkg/agent/agent_utils.go | 50 ++++++++ pkg/agent/runtime_event_logger_test.go | 159 +++++++++++++++++++++++++ pkg/agent/turn_state.go | 12 +- 5 files changed, 343 insertions(+), 35 deletions(-) diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 0392829f..2723831f 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -117,6 +117,7 @@ const ( handledToolResponseSummary = "Requested output delivered via tool attachment." sessionKeyAgentPrefix = "agent:" pendingTurnPrefix = "pending-" + providerReloadGracePeriod = 30 * time.Second metadataKeyMessageKind = "message_kind" metadataKeyToolCalls = "tool_calls" metadataKeyOutboundKind = "outbound_kind" @@ -208,6 +209,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { // slot. The goroutine is spawned immediately so the main loop keeps // draining the inbound channel. The goroutine blocks on the semaphore. go func(m bus.InboundMessage) { + var releaseSession bool // Acquire semaphore slot (blocks if at capacity) select { case al.workerSem <- struct{}{}: @@ -215,7 +217,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { case <-ctx.Done(): // Context canceled while waiting for a slot — clean up the // placeholder to prevent session-level deadlock. - al.activeTurnStates.Delete(sessionKey) + al.releaseSessionTurnState(sessionKey, nil) return } @@ -224,16 +226,21 @@ func (al *AgentLoop) Run(ctx context.Context) error { // completes normally, clearActiveTurn deletes the real turnState and // this becomes a no-op (the key is already gone). defer func() { + if releaseSession { + al.releaseSessionTurnState(sessionKey, nil) + return + } if actual, ok := al.activeTurnStates.Load(sessionKey); ok { if ts, ok := actual.(*turnState); ok && strings.HasPrefix(ts.turnID, pendingTurnPrefix) { // Placeholder still present — runTurn never replaced it. - al.activeTurnStates.Delete(sessionKey) + al.releaseSessionTurnState(sessionKey, ts) } } }() defer func() { if r := recover(); r != nil { + releaseSession = true logger.RecoverPanicNoExit(r) logger.ErrorCF("agent", "Worker goroutine panicked", map[string]any{ @@ -251,7 +258,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { } if al.takePendingStop(sessionKey) { - al.activeTurnStates.Delete(sessionKey) + al.releaseSessionTurnState(sessionKey, nil) target := &continuationTarget{ SessionKey: sessionKey, Channel: m.Channel, @@ -365,37 +372,23 @@ func (al *AgentLoop) ReloadProviderAndConfig( return fmt.Errorf("config cannot be nil") } - // Create new registry with updated config and provider - // Wrap in defer/recover to handle any panics gracefully var registry *AgentRegistry - var panicErr error - done := make(chan struct{}, 1) - - go func() { + func() { defer func() { if r := recover(); r != nil { logger.RecoverPanicNoExit(r) - panicErr = fmt.Errorf("panic during registry creation: %v", r) logger.ErrorCF("agent", "Panic during registry creation", map[string]any{"panic": r}) + registry = nil } - close(done) }() - registry = NewAgentRegistry(cfg, provider) }() - - // Wait for completion or context cancellation - select { - case <-done: - if registry == nil { - if panicErr != nil { - return fmt.Errorf("registry creation failed: %w", panicErr) - } - return fmt.Errorf("registry creation failed (nil result)") + if registry == nil { + if err := ctx.Err(); err != nil { + return fmt.Errorf("context canceled during registry creation: %w", err) } - case <-ctx.Done(): - return fmt.Errorf("context canceled during registry creation: %w", ctx.Err()) + return fmt.Errorf("registry creation failed") } // Check context again before proceeding @@ -471,17 +464,7 @@ func (al *AgentLoop) ReloadProviderAndConfig( // This prevents blocking readers while closing if oldProvider, ok := extractProvider(oldRegistry); ok { if stateful, ok := oldProvider.(providers.StatefulProvider); ok { - // Give in-flight requests a moment to complete - // Use a reasonable timeout that balances cleanup vs resource usage - select { - case <-time.After(100 * time.Millisecond): - stateful.Close() - case <-ctx.Done(): - // Context canceled, close immediately but log warning - logger.WarnCF("agent", "Context canceled during provider cleanup, forcing close", - map[string]any{"error": ctx.Err()}) - stateful.Close() - } + al.closeReloadedProvider(ctx, stateful) } } diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index aaf3d1a8..58541140 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -13,6 +13,7 @@ import ( "slices" "strings" "sync" + "sync/atomic" "testing" "time" @@ -119,6 +120,31 @@ type recordingProvider struct { lastModel string } +type panicAfterStartProvider struct { + started chan struct{} + calls atomic.Int32 +} + +func (p *panicAfterStartProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + p.calls.Add(1) + select { + case <-p.started: + default: + close(p.started) + } + panic("provider panic after turn registration") +} + +func (p *panicAfterStartProvider) GetDefaultModel() string { + return "panic-after-start" +} + func (r *recordingProvider) Chat( ctx context.Context, messages []providers.Message, @@ -5710,3 +5736,83 @@ func (p *concurrentMockProvider) Chat( func (p *concurrentMockProvider) GetDefaultModel() string { return "test-model" } + +func TestRunWorkerPanicReleasesSessionTurnState(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = t.TempDir() + cfg.Agents.Defaults.MaxParallelTurns = 1 + + msgBus := bus.NewMessageBus() + provider := &panicAfterStartProvider{started: make(chan struct{})} + al := NewAgentLoop(cfg, msgBus, provider) + defer al.Close() + + runCtx, cancelRun := context.WithCancel(context.Background()) + runDone := make(chan error, 1) + go func() { + runDone <- al.Run(runCtx) + }() + defer func() { + cancelRun() + select { + case err := <-runDone: + if err != nil { + t.Fatalf("Run() error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for Run() to exit") + } + }() + + msg := bus.InboundMessage{ + Context: bus.InboundContext{ + Channel: "test", + ChatID: "panic-chat", + ChatType: "direct", + SenderID: "user1", + }, + Content: "trigger panic", + SessionKey: "panic-session", + } + route, _, err := al.resolveMessageRoute(msg) + if err != nil { + t.Fatalf("resolveMessageRoute() error = %v", err) + } + scopeKey := resolveScopeKey(al.allocateRouteSession(route, msg).SessionKey, msg.SessionKey) + + if err := msgBus.PublishInbound(context.Background(), msg); err != nil { + t.Fatalf("PublishInbound(first) error = %v", err) + } + + select { + case <-provider.started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for first turn to start") + } + + deadline := time.Now().Add(2 * time.Second) + for { + if al.getActiveTurnState(scopeKey) == nil { + break + } + if time.Now().After(deadline) { + t.Fatal("session turn state remained stuck after worker panic") + } + time.Sleep(10 * time.Millisecond) + } + + if err := msgBus.PublishInbound(context.Background(), msg); err != nil { + t.Fatalf("PublishInbound(second) error = %v", err) + } + + deadline = time.Now().Add(2 * time.Second) + for { + if provider.calls.Load() >= 2 { + break + } + if time.Now().After(deadline) { + t.Fatal("second message did not start a new turn after panic cleanup") + } + time.Sleep(10 * time.Millisecond) + } +} diff --git a/pkg/agent/agent_utils.go b/pkg/agent/agent_utils.go index 432a9f24..e1da3bf6 100644 --- a/pkg/agent/agent_utils.go +++ b/pkg/agent/agent_utils.go @@ -13,6 +13,7 @@ import ( "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/commands" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/session" "github.com/sipeed/picoclaw/pkg/utils" @@ -584,6 +585,55 @@ func closeProviderIfStateful(provider providers.LLMProvider) { } } +func (al *AgentLoop) waitForActiveRequests(ctx context.Context, timeout time.Duration) bool { + done := make(chan struct{}) + go func() { + al.activeRequests.Wait() + close(done) + }() + + if timeout <= 0 { + select { + case <-done: + return true + case <-ctx.Done(): + return false + } + } + + timer := time.NewTimer(timeout) + defer timer.Stop() + + select { + case <-done: + return true + case <-timer.C: + return false + case <-ctx.Done(): + return false + } +} + +func (al *AgentLoop) closeReloadedProvider(ctx context.Context, provider providers.StatefulProvider) { + waitCtx := ctx + if waitCtx == nil { + waitCtx = context.Background() + } + + drained := al.waitForActiveRequests(waitCtx, providerReloadGracePeriod) + if !drained { + fields := map[string]any{"grace_period": providerReloadGracePeriod.String()} + if err := waitCtx.Err(); err != nil { + fields["error"] = err.Error() + logger.WarnCF("agent", "Provider reload interrupted while waiting for in-flight requests", fields) + } else { + logger.WarnCF("agent", "Provider reload grace period expired with in-flight requests still running", fields) + } + } + + provider.Close() +} + func makePendingTurnID(sessionKey string, seq uint64) string { return pendingTurnPrefix + sessionKey + "-" + fmt.Sprintf("%d", seq) } diff --git a/pkg/agent/runtime_event_logger_test.go b/pkg/agent/runtime_event_logger_test.go index 1c95b365..1da9803c 100644 --- a/pkg/agent/runtime_event_logger_test.go +++ b/pkg/agent/runtime_event_logger_test.go @@ -2,6 +2,7 @@ package agent import ( "context" + "errors" "sync/atomic" "testing" "time" @@ -9,6 +10,7 @@ import ( "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/config" runtimeevents "github.com/sipeed/picoclaw/pkg/events" + "github.com/sipeed/picoclaw/pkg/providers" ) func TestRuntimeEventLoggerFiltering(t *testing.T) { @@ -191,6 +193,163 @@ func TestReloadProviderAndConfigRefreshesRuntimeEventLogger(t *testing.T) { } } +type reloadBlockingProvider struct { + chatStarted chan struct{} + releaseChat chan struct{} + closeCalled chan struct{} +} + +func (p *reloadBlockingProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + select { + case <-p.chatStarted: + default: + close(p.chatStarted) + } + + select { + case <-p.releaseChat: + return &providers.LLMResponse{Content: "done"}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (p *reloadBlockingProvider) GetDefaultModel() string { + return "reload-blocking" +} + +func (p *reloadBlockingProvider) Close() { + select { + case <-p.closeCalled: + default: + close(p.closeCalled) + } +} + +func TestReloadProviderAndConfigWaitsForInFlightRequestsBeforeClosingOldProvider(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = t.TempDir() + + oldProvider := &reloadBlockingProvider{ + chatStarted: make(chan struct{}), + releaseChat: make(chan struct{}), + closeCalled: make(chan struct{}), + } + al := NewAgentLoop(cfg, bus.NewMessageBus(), oldProvider) + defer al.Close() + + msg := testInboundMessage(bus.InboundMessage{ + Channel: "test", + ChatID: "reload-chat", + SenderID: "user-1", + Content: "hold request open", + }) + + reqDone := make(chan error, 1) + go func() { + _, err := al.processMessage(context.Background(), msg) + reqDone <- err + }() + + select { + case <-oldProvider.chatStarted: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for in-flight provider request") + } + + reloadDone := make(chan error, 1) + go func() { + reloaded := config.DefaultConfig() + reloaded.Agents.Defaults.Workspace = cfg.Agents.Defaults.Workspace + reloadDone <- al.ReloadProviderAndConfig(context.Background(), &mockProvider{}, reloaded) + }() + + select { + case <-oldProvider.closeCalled: + t.Fatal("old provider closed before in-flight request completed") + case err := <-reloadDone: + t.Fatalf("reload returned early: %v", err) + case <-time.After(150 * time.Millisecond): + } + + close(oldProvider.releaseChat) + + select { + case err := <-reqDone: + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for in-flight request to complete") + } + + select { + case err := <-reloadDone: + if err != nil { + t.Fatalf("ReloadProviderAndConfig() error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for reload to finish") + } + + select { + case <-oldProvider.closeCalled: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for old provider close") + } +} + +func TestWaitForActiveRequestsHonorsContextCancellation(t *testing.T) { + al := &AgentLoop{} + al.activeRequests.Add(1) + defer al.activeRequests.Done() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if al.waitForActiveRequests(ctx, time.Second) { + t.Fatal("waitForActiveRequests() = true, want false on canceled context") + } +} + +func TestReloadProviderAndConfigReturnsCanceledErrorWhenRegistryCreationPanics(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = t.TempDir() + + al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) + defer al.Close() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := al.ReloadProviderAndConfig(ctx, &panicProviderForReloadTest{}, cfg) + if !errors.Is(err, context.Canceled) { + t.Fatalf("ReloadProviderAndConfig() error = %v, want context canceled", err) + } +} + +type panicProviderForReloadTest struct{} + +func (p *panicProviderForReloadTest) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + return &providers.LLMResponse{Content: "unused"}, nil +} + +func (p *panicProviderForReloadTest) GetDefaultModel() string { + panic("boom") +} + func TestCloseRuntimeEventLoggerSubscriptionWaitsForDrain(t *testing.T) { eventBus := runtimeevents.NewBus() defer func() { diff --git a/pkg/agent/turn_state.go b/pkg/agent/turn_state.go index ddd1eb89..2161490e 100644 --- a/pkg/agent/turn_state.go +++ b/pkg/agent/turn_state.go @@ -285,7 +285,17 @@ func (al *AgentLoop) registerActiveTurn(ts *turnState) { } func (al *AgentLoop) clearActiveTurn(ts *turnState) { - al.activeTurnStates.Delete(ts.sessionKey) + al.releaseSessionTurnState(ts.sessionKey, ts) +} + +func (al *AgentLoop) releaseSessionTurnState(sessionKey string, expected *turnState) { + if expected == nil { + al.activeTurnStates.Delete(sessionKey) + return + } + if actual, ok := al.activeTurnStates.Load(sessionKey); ok && actual == expected { + al.activeTurnStates.Delete(sessionKey) + } } func (al *AgentLoop) getActiveTurnState(sessionKey string) *turnState {