Fix agent loop reload and panic cleanup stability

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)
 		}
 	}
 

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)
+	}
+}

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)
 }

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() {

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 {