From 55f7d03458e6da90e34afc55ad5694b6408117a4 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 26 Nov 2025 19:45:22 +0800 Subject: [PATCH] Enhance state management with safeSend method to prevent panics on closed channels - Introduced safeSend method to handle context cancellation and prevent panics when sending commands to state command channel. - Updated existing state management methods to utilize safeSend, ensuring robust handling of command sending in various states. - Improved overall stability of state management by addressing potential race conditions and channel closure issues. --- agent/assistant/agent_interrupt_test.go | 8 +-- trace/state.go | 67 +++++++++++++++++++------ 2 files changed, 56 insertions(+), 19 deletions(-) diff --git a/agent/assistant/agent_interrupt_test.go b/agent/assistant/agent_interrupt_test.go index d2bb9768..716fdd6e 100644 --- a/agent/assistant/agent_interrupt_test.go +++ b/agent/assistant/agent_interrupt_test.go @@ -9,6 +9,7 @@ import ( "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/output/message" "github.com/yaoapp/yao/agent/testutils" "github.com/yaoapp/yao/openapi/oauth/types" ) @@ -29,9 +30,10 @@ func newTestContextWithInterrupt(chatID, assistantID string) *context.Context { UserAgent: "TestAgent/1.0", IP: "127.0.0.1", }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "/test/route", + Referer: context.RefererAPI, + Accept: context.AcceptWebCUI, + Route: "/test/route", + IDGenerator: message.NewIDGenerator(), // Initialize context-scoped ID generator Metadata: map[string]interface{}{ "test": "interrupt_test", }, diff --git a/trace/state.go b/trace/state.go index 398cd749..c16af3a0 100644 --- a/trace/state.go +++ b/trace/state.go @@ -260,86 +260,121 @@ func (m *manager) startStateWorker() { // Helper methods for manager to send commands +// safeSend checks if context is cancelled before sending to avoid panic on closed channel +func (m *manager) safeSend(cmd stateCommand) (ok bool) { + // Use defer/recover to handle the case where channel is closed mid-send + defer func() { + if r := recover(); r != nil { + // Channel was closed, silently return false + ok = false + } + }() + + select { + case <-m.ctx.Done(): + // Context cancelled, channel may be closed + return false + case m.stateCmdChan <- cmd: + return true + } +} + func (m *manager) stateSetRoot(node *types.TraceNode) { - m.stateCmdChan <- &cmdSetRoot{node: node} + m.safeSend(&cmdSetRoot{node: node}) } func (m *manager) stateGetRoot() *types.TraceNode { resp := make(chan *types.TraceNode, 1) - m.stateCmdChan <- &cmdGetRoot{resp: resp} + if !m.safeSend(&cmdGetRoot{resp: resp}) { + return nil // Context cancelled + } return <-resp } func (m *manager) stateSetCurrentNodes(nodes []*types.TraceNode) { - m.stateCmdChan <- &cmdSetCurrentNodes{nodes: nodes} + m.safeSend(&cmdSetCurrentNodes{nodes: nodes}) } func (m *manager) stateGetCurrentNodes() []*types.TraceNode { resp := make(chan []*types.TraceNode, 1) - m.stateCmdChan <- &cmdGetCurrentNodes{resp: resp} + if !m.safeSend(&cmdGetCurrentNodes{resp: resp}) { + return nil // Context cancelled + } return <-resp } func (m *manager) stateUpdateRootAndCurrent(root *types.TraceNode, current []*types.TraceNode) { - m.stateCmdChan <- &cmdUpdateRootAndCurrent{root: root, current: current} + m.safeSend(&cmdUpdateRootAndCurrent{root: root, current: current}) } func (m *manager) stateGetSpace(id string) (*types.TraceSpace, bool) { resp := make(chan *types.TraceSpace, 1) - m.stateCmdChan <- &cmdGetSpace{id: id, resp: resp} + if !m.safeSend(&cmdGetSpace{id: id, resp: resp}) { + return nil, false // Context cancelled + } space := <-resp return space, space != nil } func (m *manager) stateSetSpace(id string, space *types.TraceSpace) { - m.stateCmdChan <- &cmdSetSpace{id: id, space: space} + m.safeSend(&cmdSetSpace{id: id, space: space}) } func (m *manager) stateDeleteSpace(id string) { - m.stateCmdChan <- &cmdDeleteSpace{id: id} + m.safeSend(&cmdDeleteSpace{id: id}) } func (m *manager) stateGetAllSpaces() []*types.TraceSpace { resp := make(chan []*types.TraceSpace, 1) - m.stateCmdChan <- &cmdGetAllSpaces{resp: resp} + if !m.safeSend(&cmdGetAllSpaces{resp: resp}) { + return nil // Context cancelled + } return <-resp } func (m *manager) stateSetTraceStatus(status types.TraceStatus) { - m.stateCmdChan <- &cmdSetTraceStatus{status: status} + m.safeSend(&cmdSetTraceStatus{status: status}) } func (m *manager) stateGetTraceStatus() types.TraceStatus { resp := make(chan types.TraceStatus, 1) - m.stateCmdChan <- &cmdGetTraceStatus{resp: resp} + if !m.safeSend(&cmdGetTraceStatus{resp: resp}) { + return types.TraceStatusCancelled // Context cancelled + } return <-resp } func (m *manager) stateMarkCompleted() bool { resp := make(chan bool, 1) - m.stateCmdChan <- &cmdMarkCompleted{resp: resp} + if !m.safeSend(&cmdMarkCompleted{resp: resp}) { + return true // Context cancelled, treat as completed + } return <-resp } func (m *manager) stateIsCompleted() bool { resp := make(chan bool, 1) - m.stateCmdChan <- &cmdIsCompleted{resp: resp} + if !m.safeSend(&cmdIsCompleted{resp: resp}) { + return true // Context cancelled, treat as completed + } return <-resp } func (m *manager) stateAddUpdate(update *types.TraceUpdate) { - m.stateCmdChan <- &cmdAddUpdate{update: update} + m.safeSend(&cmdAddUpdate{update: update}) } func (m *manager) stateGetUpdates(since int64) []*types.TraceUpdate { resp := make(chan []*types.TraceUpdate, 1) - m.stateCmdChan <- &cmdGetUpdates{since: since, resp: resp} + if !m.safeSend(&cmdGetUpdates{since: since, resp: resp}) { + return nil // Context cancelled + } return <-resp } func (m *manager) stateSetUpdates(updates []*types.TraceUpdate) { log.Trace("[STATE] stateSetUpdates: setting %d updates for trace %s", len(updates), m.traceID) - m.stateCmdChan <- &cmdSetUpdates{updates: updates} + m.safeSend(&cmdSetUpdates{updates: updates}) } // Subscription management methods removed - now handled by SubscriptionManager