From a49695e906cc2cc218d2d50c724ef13e01b4f793 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 21 Nov 2025 20:52:10 +0800 Subject: [PATCH] Refactor interrupt handling and enhance logging in Assistant and Context components - Updated the interrupt controller to accept context IDs during initialization, improving traceability of interrupt signals. - Enhanced logging in the Assistant's Stream method to provide detailed trace information for stream lifecycle events. - Refactored context release logic to improve cleanup processes and added logging for context management. - Adjusted test cases to align with the new interrupt handling and logging structure, ensuring robust functionality and traceability. --- agent/assistant/agent.go | 112 ++++++++++- agent/assistant/agent_interrupt_test.go | 3 +- agent/context/context.go | 43 ++++- agent/context/interrupt.go | 18 +- agent/context/interrupt_test.go | 3 +- agent/context/openapi.go | 7 +- agent/llm/providers/openai/gpt5_test.go | 2 +- agent/llm/providers/openai/openai.go | 157 ++++++--------- agent/llm/providers/openai/openai_test.go | 69 ++----- openapi/chat/completions.go | 8 +- trace/manager.go | 29 ++- trace/pubsub/pubsub.go | 142 ++++++++++++++ trace/pubsub/subscriber.go | 62 ++++++ trace/state.go | 131 ++++--------- trace/subscription.go | 70 ++----- trace/trace.go | 224 +++++++++++++++++++++- 16 files changed, 722 insertions(+), 358 deletions(-) create mode 100644 trace/pubsub/pubsub.go create mode 100644 trace/pubsub/subscriber.go diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 91bacb4c..b9250167 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -2,8 +2,11 @@ package assistant import ( "fmt" + "time" + jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/connector" + "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" "github.com/yaoapp/yao/agent/llm" @@ -16,9 +19,15 @@ import ( // handler is optional, if not provided, a default handler will be used func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Message, handler ...context.StreamFunc) (*context.Response, error) { + log.Trace("[AGENT] Stream started: assistant=%s, contextID=%s", ast.ID, ctx.ID) + defer log.Trace("[AGENT] Stream ended: assistant=%s, contextID=%s", ast.ID, ctx.ID) + var err error + streamStartTime := time.Now() // Set up interrupt handler if interrupt controller is available + // InterruptController handles user interrupt signals (stop button) for appending messages + // HTTP context cancellation is handled naturally by LLM/Agent layers if ctx.Interrupt != nil { ctx.Interrupt.SetHandler(func(c *context.Context, signal *context.InterruptSignal) error { return ast.handleInterrupt(c, signal) @@ -31,6 +40,12 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa _ = traceID // traceID is available for trace logging + // Determine stream handler + streamHandler := ast.getStreamHandler(ctx, handler...) + + // Send ChunkStreamStart only for root stack (agent-level stream start) + ast.sendAgentStreamStart(ctx, streamHandler, streamStartTime) + // Trace Add trace, _ := ctx.Trace() var agentNode types.Node = nil @@ -49,6 +64,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa if agentNode != nil { agentNode.Fail(err) } + // Send error stream_end for root stack + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } @@ -66,6 +83,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa if agentNode != nil { agentNode.Fail(err) } + // Send error stream_end for root stack + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } @@ -87,6 +106,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa if agentNode != nil { agentNode.Fail(err) } + // Send error stream_end for root stack + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } @@ -96,6 +117,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa if agentNode != nil { agentNode.Fail(err) } + // Send error stream_end for root stack + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } @@ -125,18 +148,20 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // Create LLM instance with connector and options llmInstance, err := llm.New(conn, completionOptions) if err != nil { + // Send error stream_end for root stack + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } - // Use provided handler or default handler - streamHandler := llm.DefaultStreamHandler(ctx) - if len(handler) > 0 && handler[0] != nil { - streamHandler = handler[0] - } - - // Call the LLM Completion Stream + // Call the LLM Completion Stream (streamHandler was set earlier) + log.Trace("[AGENT] Calling LLM Stream: assistant=%s", ast.ID) completionResponse, err = llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler) + log.Trace("[AGENT] LLM Stream returned: assistant=%s, err=%v", ast.ID, err) if err != nil { + // Send error stream_end for root stack + log.Trace("[AGENT] Calling sendStreamEndOnError") + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) + log.Trace("[AGENT] sendStreamEndOnError returned") return nil, err } @@ -160,6 +185,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa var err error doneResponse, err = ast.Script.Done(ctx, fullMessages, completionResponse, mcpResponse) if err != nil { + // Send error stream_end for root stack + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } } @@ -171,8 +198,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa agentNode.SetOutput(context.Response{Create: createResponse, Done: doneResponse, Completion: completionResponse}) } - // Only close output if this is the root call (entry point) - // Nested calls (from MCP, hooks, etc.) should not close the output + // Only close output and send stream_end if this is the root call (entry point) + // Nested calls (from MCP, hooks, etc.) should not close the output or send stream_end // Note: Flush is already handled by the stream handler (handleStreamEnd) if ctx.Stack != nil && ctx.Stack.IsRoot() { // Log closing output for root call @@ -184,6 +211,9 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa }) } + // Send ChunkStreamEnd (agent-level stream completion) + ast.sendAgentStreamEnd(ctx, streamHandler, streamStartTime, "completed", nil, completionResponse) + // Close the output writer to send [DONE] marker if err := output.Close(ctx); err != nil { if trace, _ := ctx.Trace(); trace != nil { @@ -597,6 +627,70 @@ func (ast *Assistant) WithHistory(ctx *context.Context, messages []context.Messa return messages, nil } +// getStreamHandler returns the stream handler from the provided handlers or a default one +func (ast *Assistant) getStreamHandler(ctx *context.Context, handler ...context.StreamFunc) context.StreamFunc { + if len(handler) > 0 && handler[0] != nil { + return handler[0] + } + return llm.DefaultStreamHandler(ctx) +} + +// sendAgentStreamStart sends ChunkStreamStart for root stack only (agent-level stream start) +// This ensures only one stream_start per agent execution, even with multiple LLM calls +func (ast *Assistant) sendAgentStreamStart(ctx *context.Context, handler context.StreamFunc, startTime time.Time) { + if ctx.Stack == nil || !ctx.Stack.IsRoot() || handler == nil { + return + } + + requestID := fmt.Sprintf("agent_req_%d", startTime.UnixNano()) + startData := &context.StreamStartData{ + RequestID: requestID, + Timestamp: startTime.UnixMilli(), + Model: ast.ID, // Use assistant ID as the "model" for agent-level stream + } + + if startJSON, err := jsoniter.Marshal(startData); err == nil { + handler(context.ChunkStreamStart, startJSON) + } +} + +// sendAgentStreamEnd sends ChunkStreamEnd for root stack only (agent-level stream completion) +func (ast *Assistant) sendAgentStreamEnd(ctx *context.Context, handler context.StreamFunc, startTime time.Time, status string, err error, response *context.CompletionResponse) { + if ctx.Stack == nil || !ctx.Stack.IsRoot() || handler == nil { + return + } + + // Check if context is cancelled - if so, skip handler call to avoid blocking + if ctx.Context != nil && ctx.Context.Err() != nil { + log.Trace("[AGENT] Context cancelled, skipping sendAgentStreamEnd handler call") + return + } + + endData := &context.StreamEndData{ + RequestID: fmt.Sprintf("agent_req_%d", startTime.UnixNano()), + Timestamp: time.Now().UnixMilli(), + DurationMs: time.Since(startTime).Milliseconds(), + Status: status, + } + + if err != nil { + endData.Error = err.Error() + } + + if response != nil && response.Usage != nil { + endData.Usage = response.Usage + } + + if endJSON, marshalErr := jsoniter.Marshal(endData); marshalErr == nil { + handler(context.ChunkStreamEnd, endJSON) + } +} + +// sendStreamEndOnError sends ChunkStreamEnd with error status for root stack only +func (ast *Assistant) sendStreamEndOnError(ctx *context.Context, handler context.StreamFunc, startTime time.Time, err error) { + ast.sendAgentStreamEnd(ctx, handler, startTime, "error", err, nil) +} + // handleInterrupt handles the interrupt signal // This is called by the interrupt listener when a signal is received func (ast *Assistant) handleInterrupt(ctx *context.Context, signal *context.InterruptSignal) error { diff --git a/agent/assistant/agent_interrupt_test.go b/agent/assistant/agent_interrupt_test.go index f12cbd5f..d2bb9768 100644 --- a/agent/assistant/agent_interrupt_test.go +++ b/agent/assistant/agent_interrupt_test.go @@ -47,7 +47,6 @@ func newTestContextWithInterrupt(chatID, assistantID string) *context.Context { // Initialize interrupt controller ctx.Interrupt = context.NewInterruptController() - ctx.Interrupt.SetContextID(ctx.ID) // Register context globally if err := context.Register(ctx); err != nil { @@ -55,7 +54,7 @@ func newTestContextWithInterrupt(chatID, assistantID string) *context.Context { } // Start interrupt listener - ctx.Interrupt.Start() + ctx.Interrupt.Start(ctx.ID) return ctx } diff --git a/agent/context/context.go b/agent/context/context.go index d4de0693..65e9ef21 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -75,6 +75,8 @@ func WithTimeout(parent Context, timeout time.Duration) (Context, context.Cancel // Release the context and clean up all resources including stacks and trace func (ctx *Context) Release() { + log.Trace("[RELEASE] Context cleanup started: contextID=%s, assistantID=%s", ctx.ID, ctx.AssistantID) + // Unregister from global registry if ctx.ID != "" { Unregister(ctx.ID) @@ -82,29 +84,61 @@ func (ctx *Context) Release() { // Stop interrupt controller if ctx.Interrupt != nil { + log.Trace("[RELEASE] Stopping interrupt controller") ctx.Interrupt.Stop() ctx.Interrupt = nil } // Complete and release trace if exists if ctx.trace != nil && ctx.Stack != nil && ctx.Stack.TraceID != "" { - // Mark trace as complete (sends final event) - _ = ctx.trace.MarkComplete() + log.Trace("[RELEASE] Releasing trace: traceID=%s", ctx.Stack.TraceID) - // Release from global registry (removes from registry and closes resources) - _ = trace.Release(ctx.Stack.TraceID) + // Check if context is cancelled - if so, mark as cancelled instead of complete + if ctx.Context != nil && ctx.Context.Err() != nil { + log.Trace("[RELEASE] Context cancelled, marking trace as cancelled: err=%v", ctx.Context.Err()) + + // Mark trace as cancelled (saves to disk and broadcasts to subscribers) + log.Trace("[RELEASE] Calling trace.MarkCancelled: traceID=%s", ctx.Stack.TraceID) + if err := trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error()); err != nil { + log.Trace("[RELEASE] Failed to mark trace as cancelled: %v", err) + } else { + log.Trace("[RELEASE] Successfully marked trace as cancelled") + } + + // Release trace from registry + // Subscribers will be notified via channel close and will cleanup automatically + log.Trace("[RELEASE] Calling trace.Release: traceID=%s", ctx.Stack.TraceID) + if err := trace.Release(ctx.Stack.TraceID); err != nil { + log.Trace("[RELEASE] Failed to release trace from registry: %v", err) + } else { + log.Trace("[RELEASE] Successfully released trace from registry") + } + } else { + // Normal case: mark complete then release + if err := ctx.trace.MarkComplete(); err != nil { + log.Trace("[RELEASE] Failed to mark trace complete: %v", err) + } + + if err := trace.Release(ctx.Stack.TraceID); err != nil { + log.Trace("[RELEASE] Failed to release trace: %v", err) + } + } ctx.trace = nil + } else { + log.Trace("[RELEASE] No trace to release (trace=%v, stack=%v)", ctx.trace != nil, ctx.Stack != nil) } // Clear space if ctx.Space != nil { + log.Trace("[RELEASE] Clearing space") ctx.Space.Clear() ctx.Space = nil } // Clear stacks if ctx.Stacks != nil { + log.Trace("[RELEASE] Clearing %d stacks", len(ctx.Stacks)) for k := range ctx.Stacks { delete(ctx.Stacks, k) } @@ -117,6 +151,7 @@ func (ctx *Context) Release() { // Clear writer reference ctx.Writer = nil + log.Trace("[RELEASE] Context cleanup completed: contextID=%s", ctx.ID) ctx = nil } diff --git a/agent/context/interrupt.go b/agent/context/interrupt.go index c4419298..94dab6af 100644 --- a/agent/context/interrupt.go +++ b/agent/context/interrupt.go @@ -4,6 +4,8 @@ import ( "context" "fmt" "time" + + "github.com/yaoapp/kun/log" ) // NewInterruptController creates a new interrupt controller @@ -17,13 +19,14 @@ func NewInterruptController() *InterruptController { } // Start starts the interrupt listener goroutine -func (ic *InterruptController) Start() { +func (ic *InterruptController) Start(contextID string) { if ic.listenerStarted { return } ic.mutex.Lock() ic.listenerStarted = true + ic.contextID = contextID ic.mutex.Unlock() go ic.listen() @@ -37,23 +40,16 @@ func (ic *InterruptController) SetHandler(handler InterruptHandler) { ic.handler = handler } -// SetContextID sets the context ID for retrieving the parent context -func (ic *InterruptController) SetContextID(contextID string) { - if ic == nil { - return - } - ic.contextID = contextID -} - // listen is the main listener goroutine that processes interrupt signals func (ic *InterruptController) listen() { for { select { case signal := <-ic.queue: + // Handle user interrupt signal (stop button, for appending messages) ic.handleSignal(signal) case <-ic.ctx.Done(): - // Context cancelled, stop listening + // Internal context cancelled, stop listening return } } @@ -65,6 +61,8 @@ func (ic *InterruptController) handleSignal(signal *InterruptSignal) { return } + log.Trace("[INTERRUPT] Signal received: type=%s, messages=%d, timestamp=%d", signal.Type, len(signal.Messages), signal.Timestamp) + ic.mutex.Lock() // If no current interrupt, set it as current diff --git a/agent/context/interrupt_test.go b/agent/context/interrupt_test.go index dac3a83b..92df85bb 100644 --- a/agent/context/interrupt_test.go +++ b/agent/context/interrupt_test.go @@ -44,7 +44,6 @@ func newTestContextWithInterrupt(chatID, assistantID string) *Context { // Initialize interrupt controller ctx.Interrupt = NewInterruptController() - ctx.Interrupt.SetContextID(ctx.ID) // Register context globally if err := Register(ctx); err != nil { @@ -52,7 +51,7 @@ func newTestContextWithInterrupt(chatID, assistantID string) *Context { } // Start interrupt listener - ctx.Interrupt.Start() + ctx.Interrupt.Start(ctx.ID) return ctx } diff --git a/agent/context/openapi.go b/agent/context/openapi.go index cc5bb765..76bfb092 100644 --- a/agent/context/openapi.go +++ b/agent/context/openapi.go @@ -69,15 +69,16 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest // Initialize interrupt controller ctx.Interrupt = NewInterruptController() - ctx.Interrupt.SetContextID(ctx.ID) - // Register context to global registry first + // Register context to global registry first (required for interrupt handler callback) if err := Register(ctx); err != nil { return nil, nil, fmt.Errorf("failed to register context: %w", err) } // Start interrupt listener after registration - ctx.Interrupt.Start() + // Only monitors interrupt signals (user stop button for appending messages) + // HTTP context cancellation is handled by LLM/Agent layers naturally + ctx.Interrupt.Start(ctx.ID) return completionReq, ctx, nil } diff --git a/agent/llm/providers/openai/gpt5_test.go b/agent/llm/providers/openai/gpt5_test.go index ff0f55e5..a4bc4d90 100644 --- a/agent/llm/providers/openai/gpt5_test.go +++ b/agent/llm/providers/openai/gpt5_test.go @@ -284,7 +284,7 @@ func TestGPT5Vision(t *testing.T) { { Type: context.ContentImageURL, ImageURL: &context.ImageURL{ - URL: "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/320px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", + URL: "https://raw.githubusercontent.com/YaoApp/yao/refs/heads/main/yao/data/icons/icon.png", }, }, }, diff --git a/agent/llm/providers/openai/openai.go b/agent/llm/providers/openai/openai.go index 9d1b7dc9..c3a0dae5 100644 --- a/agent/llm/providers/openai/openai.go +++ b/agent/llm/providers/openai/openai.go @@ -9,6 +9,7 @@ import ( jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/http" + "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" "github.com/yaoapp/yao/agent/llm/adapters" @@ -265,13 +266,21 @@ func (p *Provider) Stream(ctx *context.Context, messages []context.Message, opti } response, err := p.streamWithRetry(ctx, currentMessages, options, handler) + log.Trace("[LLM] streamWithRetry returned: err=%v", err) if err == nil { - if trace != nil { + if trace != nil && goCtx.Err() == nil { trace.Debug("OpenAI Stream: Request completed successfully") } return response, nil } lastErr = err + log.Trace("[LLM] Checking context after error: goCtx.Err()=%v", goCtx.Err()) + + // Check for context cancellation before logging (trace calls may block if context is cancelled) + if goCtx.Err() != nil { + log.Trace("[LLM] Context cancelled in retry loop, returning") + return nil, fmt.Errorf("context cancelled: %w", goCtx.Err()) + } if trace != nil { trace.Debug("OpenAI Stream: Request failed", map[string]any{ @@ -361,18 +370,8 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess } } - // Send stream_start event - if handler != nil { - model, _ := p.GetModel() - startData := &context.StreamStartData{ - RequestID: requestID, - Timestamp: streamStartTime.UnixMilli(), - Model: model, - } - if startJSON, err := jsoniter.Marshal(startData); err == nil { - handler(context.ChunkStreamStart, startJSON) - } - } + // Note: ChunkStreamStart/End are now sent at Agent level, not LLM level + // This is because an agent may make multiple LLM calls in one stream // Preprocess messages and options through adapters processedMessages := messages @@ -403,19 +402,6 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess // Build request body requestBody, err := p.buildRequestBody(processedMessages, processedOptions, true) if err != nil { - // Send stream_end with error - if handler != nil { - endData := &context.StreamEndData{ - RequestID: requestID, - Timestamp: time.Now().UnixMilli(), - DurationMs: time.Since(streamStartTime).Milliseconds(), - Status: "error", - Error: err.Error(), - } - if endJSON, err := jsoniter.Marshal(endData); err == nil { - handler(context.ChunkStreamEnd, endJSON) - } - } return nil, fmt.Errorf("failed to build request body: %w", err) } @@ -681,7 +667,9 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess } // Make streaming request (goCtx already set at function start) + log.Trace("[LLM] Starting HTTP Stream request: url=%s", url) err = req.Stream(goCtx, "POST", requestBody, wrappedHandler) + log.Trace("[LLM] HTTP Stream request returned: err=%v", err) // Check if we captured an error response if errorDetected && errorBuffer.Len() > 0 { @@ -708,30 +696,19 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess } } - // Log any error from streaming - if err != nil && trace != nil { - trace.Error(i18n.T(ctx.Locale, "llm.openai.stream.error"), map[string]any{"error": err.Error()}) // "OpenAI Stream Error" + // Check if error is due to context cancellation FIRST (before logging) + // This prevents blocking on trace operations when context is cancelled + if err != nil && goCtx.Err() != nil { + log.Trace("[LLM] Context cancelled detected, skipping handler calls and returning") + // NOTE: Do NOT call handler or groupTracker.endGroup here + // The connection is already closed, calling handler may block indefinitely + // Just return the error immediately + return nil, fmt.Errorf("stream cancelled: %w", goCtx.Err()) } - // Check if error is due to context cancellation - if err != nil && goCtx.Err() != nil { - // End current group if active - groupTracker.endGroup(handler) - - // Send stream_end with cancellation status - if handler != nil { - endData := &context.StreamEndData{ - RequestID: requestID, - Timestamp: time.Now().UnixMilli(), - DurationMs: time.Since(streamStartTime).Milliseconds(), - Status: "cancelled", - Error: goCtx.Err().Error(), - } - if endJSON, err := jsoniter.Marshal(endData); err == nil { - handler(context.ChunkStreamEnd, endJSON) - } - } - return nil, fmt.Errorf("stream cancelled: %w", goCtx.Err()) + // Log any error from streaming (only if not cancelled) + if err != nil && trace != nil { + trace.Error(i18n.T(ctx.Locale, "llm.openai.stream.error"), map[string]any{"error": err.Error()}) // "OpenAI Stream Error" } if err != nil { @@ -742,18 +719,6 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess if handler != nil { errData := []byte(err.Error()) handler(context.ChunkError, errData) - - // Send stream_end with error - endData := &context.StreamEndData{ - RequestID: requestID, - Timestamp: time.Now().UnixMilli(), - DurationMs: time.Since(streamStartTime).Milliseconds(), - Status: "error", - Error: err.Error(), - } - if endJSON, err := jsoniter.Marshal(endData); err == nil { - handler(context.ChunkStreamEnd, endJSON) - } } return nil, fmt.Errorf("streaming request failed: %w", err) } @@ -783,18 +748,6 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess if handler != nil { errData := []byte(err.Error()) handler(context.ChunkError, errData) - - // Send stream_end with error - endData := &context.StreamEndData{ - RequestID: requestID, - Timestamp: time.Now().UnixMilli(), - DurationMs: time.Since(streamStartTime).Milliseconds(), - Status: "error", - Error: err.Error(), - } - if endJSON, err := jsoniter.Marshal(endData); err == nil { - handler(context.ChunkStreamEnd, endJSON) - } } return nil, err } @@ -835,20 +788,6 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess // End current group groupTracker.endGroup(handler) - // Send stream_end with validation error - if handler != nil { - endData := &context.StreamEndData{ - RequestID: requestID, - Timestamp: time.Now().UnixMilli(), - DurationMs: time.Since(streamStartTime).Milliseconds(), - Status: "error", - Error: "tool call validation failed", - } - if endJSON, err := jsoniter.Marshal(endData); err == nil { - handler(context.ChunkStreamEnd, endJSON) - } - } - // Tool call validation failed, need to retry with error feedback return nil, fmt.Errorf("tool call validation failed: %w", err) } @@ -857,20 +796,6 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess // End final group if still active groupTracker.endGroup(handler) - // Send stream_end event (success) - if handler != nil { - endData := &context.StreamEndData{ - RequestID: requestID, - Timestamp: time.Now().UnixMilli(), - DurationMs: time.Since(streamStartTime).Milliseconds(), - Status: "completed", - Usage: response.Usage, - } - if endJSON, err := jsoniter.Marshal(endData); err == nil { - handler(context.ChunkStreamEnd, endJSON) - } - } - return response, nil } @@ -1140,7 +1065,37 @@ func (p *Provider) buildRequestBody(messages []context.Message, options *context } if msg.Content != nil { - apiMsg["content"] = msg.Content + // Check if Content is []context.ContentPart and convert to API format + if parts, ok := msg.Content.([]context.ContentPart); ok { + apiParts := make([]map[string]interface{}, 0, len(parts)) + for _, part := range parts { + apiPart := map[string]interface{}{ + "type": string(part.Type), + } + switch part.Type { + case context.ContentText: + apiPart["text"] = part.Text + case context.ContentImageURL: + if part.ImageURL != nil { + apiPart["image_url"] = map[string]interface{}{ + "url": part.ImageURL.URL, + } + if part.ImageURL.Detail != "" { + apiPart["image_url"].(map[string]interface{})["detail"] = part.ImageURL.Detail + } + } + case context.ContentInputAudio: + if part.InputAudio != nil { + apiPart["input_audio"] = part.InputAudio + } + } + apiParts = append(apiParts, apiPart) + } + apiMsg["content"] = apiParts + } else { + // Content is string or already in map format, use as is + apiMsg["content"] = msg.Content + } } if msg.Name != nil { diff --git a/agent/llm/providers/openai/openai_test.go b/agent/llm/providers/openai/openai_test.go index ca03e02c..f437cac9 100644 --- a/agent/llm/providers/openai/openai_test.go +++ b/agent/llm/providers/openai/openai_test.go @@ -1250,7 +1250,8 @@ func TestOpenAIProxySupport(t *testing.T) { t.Log("HTTP proxy support is implemented via http.GetTransport using environment variables") } -// TestOpenAIStreamLifecycleEvents tests that lifecycle events are sent correctly +// TestOpenAIStreamLifecycleEvents tests that LLM-level lifecycle events (group_start/end) are sent correctly +// Note: stream_start/end are now sent at Agent level, not LLM level func TestOpenAIStreamLifecycleEvents(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() @@ -1282,9 +1283,8 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) { ctx := newTestContext("test-lifecycle", "openai.gpt-4o") - // Track lifecycle events + // Track lifecycle events (only group-level events at LLM layer) var events []string - var streamStartReceived, streamEndReceived bool var groupStartReceived, groupEndReceived bool handler := func(chunkType context.StreamChunkType, data []byte) int { @@ -1292,31 +1292,10 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) { switch chunkType { case context.ChunkStreamStart: - streamStartReceived = true - var startData context.StreamStartData - if err := json.Unmarshal(data, &startData); err == nil { - t.Logf("✓ stream_start: request_id=%s, model=%s", startData.RequestID, startData.Model) - if startData.RequestID == "" { - t.Error("stream_start missing request_id") - } - } else { - t.Errorf("Failed to parse stream_start data: %v", err) - } + t.Error("❌ LLM layer should NOT send stream_start (now sent at Agent level)") case context.ChunkStreamEnd: - streamEndReceived = true - var endData context.StreamEndData - if err := json.Unmarshal(data, &endData); err == nil { - t.Logf("✓ stream_end: status=%s, duration=%dms", endData.Status, endData.DurationMs) - if endData.Status != "completed" { - t.Errorf("stream_end status should be 'completed', got '%s'", endData.Status) - } - if endData.DurationMs <= 0 { - t.Error("stream_end duration should be > 0") - } - } else { - t.Errorf("Failed to parse stream_end data: %v", err) - } + t.Error("❌ LLM layer should NOT send stream_end (now sent at Agent level)") case context.ChunkGroupStart: groupStartReceived = true @@ -1359,13 +1338,7 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) { t.Fatal("Response is nil") } - // Validate that all lifecycle events were received - if !streamStartReceived { - t.Error("stream_start event was not received") - } - if !streamEndReceived { - t.Error("stream_end event was not received") - } + // Validate that LLM-level lifecycle events were received if !groupStartReceived { t.Error("group_start event was not received") } @@ -1373,20 +1346,14 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) { t.Error("group_end event was not received") } - // Validate event order: stream_start should be first, stream_end should be last - if len(events) < 4 { - t.Errorf("Expected at least 4 events, got %d", len(events)) - } else { - if events[0] != "stream_start" { - t.Errorf("First event should be stream_start, got %s", events[0]) - } - if events[len(events)-1] != "stream_end" { - t.Errorf("Last event should be stream_end, got %s", events[len(events)-1]) - } + // Validate event order: group_start should come before group_end + if len(events) < 2 { + t.Errorf("Expected at least 2 events (group_start, group_end), got %d", len(events)) } t.Logf("Total events received: %d", len(events)) - t.Log("Lifecycle events test completed successfully") + t.Log("LLM lifecycle events test completed successfully") + t.Log("Note: stream_start/end are now tested at Agent level, not LLM level") } // TestOpenAIStreamContextCancellation tests that stream respects context cancellation @@ -1426,18 +1393,14 @@ func TestOpenAIStreamContextCancellation(t *testing.T) { ctx.Context = goCtx var receivedChunks int - var receivedStreamEnd bool handler := func(chunkType context.StreamChunkType, data []byte) int { if chunkType == context.ChunkText || chunkType == context.ChunkToolCall { receivedChunks++ } + // Note: stream_end is now sent at Agent level, not LLM level if chunkType == context.ChunkStreamEnd { - receivedStreamEnd = true - var endData context.StreamEndData - if err := json.Unmarshal(data, &endData); err == nil { - t.Logf("stream_end status: %s, error: %s", endData.Status, endData.Error) - } + t.Error("❌ LLM layer should NOT send stream_end (now sent at Agent level)") } return 0 } @@ -1462,13 +1425,9 @@ func TestOpenAIStreamContextCancellation(t *testing.T) { t.Logf("Warning: Response is not nil despite cancellation (partial response)") } - // Should have received stream_end event (even for cancellation) - if !receivedStreamEnd { - t.Error("Expected stream_end event even for cancelled stream") - } - t.Logf("Received %d chunks before cancellation", receivedChunks) t.Log("Context cancellation test completed successfully") + t.Log("Note: stream_end for cancellation is now sent at Agent level") } // TestOpenAIStreamWithTemperature tests different temperature settings diff --git a/openapi/chat/completions.go b/openapi/chat/completions.go index 78952261..1ee42ded 100644 --- a/openapi/chat/completions.go +++ b/openapi/chat/completions.go @@ -4,6 +4,7 @@ import ( "fmt" "github.com/gin-gonic/gin" + "github.com/yaoapp/kun/log" "github.com/yaoapp/kun/utils" "github.com/yaoapp/yao/agent" "github.com/yaoapp/yao/agent/assistant" @@ -37,7 +38,10 @@ func GinCreateCompletions(c *gin.Context) { return } - defer ctx.Release() // Release the context after the request is complete + defer func() { + log.Trace("[HTTP] Handler defer: calling ctx.Release()") + ctx.Release() + }() // Print request info for debugging fmt.Println("-----------------------------------------------") @@ -73,7 +77,9 @@ func GinCreateCompletions(c *gin.Context) { // Stream the completion (uses default handler which sends to ctx.Writer) // The Stream method will automatically close the writer and send [DONE] marker + log.Trace("[HTTP] Calling ast.Stream()") res, err := ast.Stream(ctx, completionReq.Messages) + log.Trace("[HTTP] ast.Stream() returned, err=%v", err) if err != nil { fmt.Println("-----------------------------------------------") fmt.Println("Error: ", err.Error()) diff --git a/trace/manager.go b/trace/manager.go index e78639a1..ee249903 100644 --- a/trace/manager.go +++ b/trace/manager.go @@ -6,6 +6,8 @@ import ( "time" gonanoid "github.com/matoous/go-nanoid/v2" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/trace/pubsub" "github.com/yaoapp/yao/trace/types" ) @@ -17,10 +19,12 @@ type manager struct { driver types.Driver stateCmdChan chan stateCommand // Single channel for all state mutations autoArchive bool // Auto-archive on complete/fail + pubsub *pubsub.PubSub // Reference to independent pubsub service (for publishing only, doesn't own it) } // NewManager creates a new trace manager instance -func NewManager(ctx context.Context, traceID string, driver types.Driver, option *types.TraceOption) (types.Manager, error) { +// pubsubService: reference to independent pubsub service (manager doesn't own it, just publishes to it) +func NewManager(ctx context.Context, traceID string, driver types.Driver, pubsubService *pubsub.PubSub, option *types.TraceOption) (types.Manager, error) { // Create a cancellable context for the manager managerCtx, cancel := context.WithCancel(ctx) @@ -37,6 +41,7 @@ func NewManager(ctx context.Context, traceID string, driver types.Driver, option driver: driver, stateCmdChan: make(chan stateCommand, 100), // Buffered channel for performance autoArchive: autoArchive, + pubsub: pubsubService, // Reference only, doesn't manage lifecycle } // Start state worker goroutine @@ -44,18 +49,26 @@ func NewManager(ctx context.Context, traceID string, driver types.Driver, option // Try to load existing updates from driver (for resumed traces) if existingUpdates, err := driver.LoadUpdates(ctx, traceID, 0); err == nil && len(existingUpdates) > 0 { + log.Trace("[MANAGER] NewManager: loaded %d existing updates from driver for trace %s", len(existingUpdates), traceID) m.stateSetUpdates(existingUpdates) // Check if trace was already completed for _, update := range existingUpdates { if update.Type == types.UpdateTypeComplete { + log.Trace("[MANAGER] NewManager: trace %s was already completed, marking as completed", traceID) m.stateMarkCompleted() if data, ok := update.Data.(*types.TraceCompleteData); ok { + log.Trace("[MANAGER] NewManager: setting trace status to %s", data.Status) m.stateSetTraceStatus(data.Status) } break } } } else { + if err != nil { + log.Trace("[MANAGER] NewManager: failed to load existing updates for trace %s: %v", traceID, err) + } else { + log.Trace("[MANAGER] NewManager: no existing updates found for trace %s, creating new trace", traceID) + } // New trace - create and broadcast init event now := time.Now().UnixMilli() m.addUpdateAndBroadcast(&types.TraceUpdate{ @@ -75,16 +88,22 @@ func genNodeID() string { return id } -// addUpdateAndBroadcast persists, adds to history, and broadcasts an update +// addUpdateAndBroadcast persists, adds to history, and publishes an update func (m *manager) addUpdateAndBroadcast(update *types.TraceUpdate) { // Persist to driver (synchronous - no race) - _ = m.driver.SaveUpdate(context.Background(), m.traceID, update) + if err := m.driver.SaveUpdate(context.Background(), m.traceID, update); err != nil { + log.Trace("[MANAGER] addUpdateAndBroadcast: failed to save update type=%s for trace %s: %v", update.Type, m.traceID, err) + } else { + log.Trace("[MANAGER] addUpdateAndBroadcast: successfully saved update type=%s for trace %s", update.Type, m.traceID) + } // Add to in-memory history m.stateAddUpdate(update) - // Broadcast to subscribers - m.stateBroadcast(update) + // Publish to independent PubSub service (manager just publishes, doesn't manage pubsub lifecycle) + if m.pubsub != nil { + m.pubsub.Publish(update) + } } // checkContext checks if context is cancelled diff --git a/trace/pubsub/pubsub.go b/trace/pubsub/pubsub.go new file mode 100644 index 00000000..9f8ee309 --- /dev/null +++ b/trace/pubsub/pubsub.go @@ -0,0 +1,142 @@ +package pubsub + +import ( + "sync" + + gonanoid "github.com/matoous/go-nanoid/v2" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/trace/types" +) + +// PubSub is an independent publish-subscribe service for trace updates +// It acts as a message broker between trace writers and readers +type PubSub struct { + eventBus chan *types.TraceUpdate // Event bus for incoming events + subscribers map[string]chan *types.TraceUpdate // Active subscribers + mu sync.RWMutex // Protects subscribers map + stopCh chan struct{} // Signal to stop the service + stopped bool // Whether service is stopped +} + +// New creates a new PubSub service +func New() *PubSub { + ps := &PubSub{ + eventBus: make(chan *types.TraceUpdate, 1000), // Buffered event bus + subscribers: make(map[string]chan *types.TraceUpdate), + stopCh: make(chan struct{}), + stopped: false, + } + + // Start forwarding service + go ps.forward() + + return ps +} + +// forward continuously forwards events from eventBus to all subscribers +// This runs in a dedicated goroutine +func (ps *PubSub) forward() { + for { + select { + case event := <-ps.eventBus: + ps.mu.RLock() + subscriberCount := len(ps.subscribers) + + if subscriberCount == 0 { + // No subscribers, discard event + ps.mu.RUnlock() + continue + } + + // Forward to all subscribers (non-blocking) + for subID, ch := range ps.subscribers { + select { + case ch <- event: + // Sent successfully + default: + // Subscriber is slow or channel full, skip + log.Trace("[PUBSUB] Subscriber %s is slow, skipping event type=%s", subID, event.Type) + } + } + ps.mu.RUnlock() + + case <-ps.stopCh: + return + } + } +} + +// Publish sends an event to the event bus +// This is called by trace writers (e.g., manager.addUpdateAndBroadcast) +func (ps *PubSub) Publish(event *types.TraceUpdate) { + if ps.stopped { + return + } + + select { + case ps.eventBus <- event: + // Event published successfully + default: + // Event bus full, this shouldn't happen with large buffer (log as warning) + log.Warn("[PUBSUB] Event bus full, discarding event type=%s", event.Type) + } +} + +// Subscribe creates a new subscription and returns a channel for receiving updates +// The caller is responsible for reading from the channel and closing it when done +func (ps *PubSub) Subscribe(bufferSize int) (<-chan *types.TraceUpdate, string) { + // Generate unique subscriber ID + subID, _ := gonanoid.Generate("0123456789abcdefghijklmnopqrstuvwxyz", 12) + + // Create subscriber channel + ch := make(chan *types.TraceUpdate, bufferSize) + + // Register subscriber + ps.mu.Lock() + ps.subscribers[subID] = ch + ps.mu.Unlock() + + return ch, subID +} + +// Unsubscribe removes a subscriber and closes its channel +func (ps *PubSub) Unsubscribe(subID string) { + ps.mu.Lock() + defer ps.mu.Unlock() + + ch, exists := ps.subscribers[subID] + if !exists { + return + } + + // Remove from map + delete(ps.subscribers, subID) + + // Close channel + close(ch) +} + +// SubscriberCount returns the number of active subscribers +func (ps *PubSub) SubscriberCount() int { + ps.mu.RLock() + defer ps.mu.RUnlock() + return len(ps.subscribers) +} + +// Stop stops the forwarding service and closes all subscriber channels +func (ps *PubSub) Stop() { + if ps.stopped { + return + } + + ps.stopped = true + close(ps.stopCh) + + // Close all subscriber channels + ps.mu.Lock() + for _, ch := range ps.subscribers { + close(ch) + } + ps.subscribers = make(map[string]chan *types.TraceUpdate) + ps.mu.Unlock() +} diff --git a/trace/pubsub/subscriber.go b/trace/pubsub/subscriber.go new file mode 100644 index 00000000..26ada520 --- /dev/null +++ b/trace/pubsub/subscriber.go @@ -0,0 +1,62 @@ +package pubsub + +import ( + "time" + + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/trace/types" +) + +// Subscriber represents a subscription to trace updates +type Subscriber struct { + ID string + Channel <-chan *types.TraceUpdate + pubsub *PubSub +} + +// Unsubscribe removes the subscription and closes the channel +func (s *Subscriber) Unsubscribe() { + s.pubsub.Unsubscribe(s.ID) +} + +// SubscribeWithHistory creates a subscription and replays historical updates first +// historicalUpdates: updates to replay before starting live stream +// bufferSize: size of the subscription channel buffer +func (ps *PubSub) SubscribeWithHistory(historicalUpdates []*types.TraceUpdate, bufferSize int) *Subscriber { + // Create subscription + ch, subID := ps.Subscribe(bufferSize) + + // Create subscriber + sub := &Subscriber{ + ID: subID, + Channel: ch, + pubsub: ps, + } + + // Replay historical updates in a goroutine + // This allows the subscription to start immediately + go func() { + // Get writable channel for replay + ps.mu.RLock() + writeCh, exists := ps.subscribers[subID] + ps.mu.RUnlock() + + if !exists { + return + } + + // Replay all historical updates (blocking send to ensure delivery) + for i, update := range historicalUpdates { + select { + case writeCh <- update: + // Sent successfully + case <-time.After(5 * time.Second): + // Timeout - subscriber is too slow or disconnected + log.Trace("[PUBSUB] Subscriber %s timed out during replay at update %d/%d", subID, i, len(historicalUpdates)) + return + } + } + }() + + return sub +} diff --git a/trace/state.go b/trace/state.go index a4b26519..398cd749 100644 --- a/trace/state.go +++ b/trace/state.go @@ -3,6 +3,7 @@ package trace import ( "time" + "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/trace/types" ) @@ -17,7 +18,7 @@ type managerState struct { traceStatus types.TraceStatus completed bool updates []*types.TraceUpdate - subscribers map[string]chan *types.TraceUpdate + // Note: subscribers moved to SubscriptionManager (no longer in state) } // State command interface - all commands are processed serially @@ -190,62 +191,8 @@ func (c *cmdSetUpdates) execute(s *managerState) { s.updates = c.updates } -// --- Subscriber Commands --- - -type cmdAddSubscriber struct { - id string - ch chan *types.TraceUpdate -} - -func (c *cmdAddSubscriber) execute(s *managerState) { - s.subscribers[c.id] = c.ch -} - -type cmdRemoveSubscriber struct { - id string -} - -func (c *cmdRemoveSubscriber) execute(s *managerState) { - delete(s.subscribers, c.id) -} - -type cmdGetSubscribers struct { - resp chan map[string]chan *types.TraceUpdate -} - -func (c *cmdGetSubscribers) execute(s *managerState) { - // Return a copy of the map - subs := make(map[string]chan *types.TraceUpdate, len(s.subscribers)) - for id, ch := range s.subscribers { - subs[id] = ch - } - c.resp <- subs -} - -// --- Broadcast Command (special - sends to all subscribers) --- - -type cmdBroadcast struct { - update *types.TraceUpdate -} - -func (c *cmdBroadcast) execute(s *managerState) { - // Send to all subscribers (non-blocking, with panic recovery) - for _, ch := range s.subscribers { - func(channel chan *types.TraceUpdate) { - defer func() { - // Recover from panic if channel is closed - if r := recover(); r != nil { - // Channel was closed, ignore - } - }() - select { - case channel <- c.update: - default: - // Subscriber is slow, skip (non-blocking) - } - }(ch) - } -} +// --- Subscriber Commands (REMOVED - now handled by SubscriptionManager) --- +// Subscriber management has been decoupled from state machine for better separation of concerns // --- Space KV Commands (for concurrent safety) --- // These ensure all operations on a space are serialized through state worker @@ -273,39 +220,39 @@ func (m *manager) startStateWorker() { traceStatus: types.TraceStatusPending, completed: false, updates: make([]*types.TraceUpdate, 0, 100), - subscribers: make(map[string]chan *types.TraceUpdate), + // subscribers removed - now managed by SubscriptionManager } - // Process commands until context is cancelled or trace is completed + // Process commands until channel is closed (on Release) + // Note: We don't exit on context cancellation anymore - state machine should continue + // running until Release() is called, which closes the channel for { - select { - case cmd, ok := <-m.stateCmdChan: - if !ok { - // Channel closed - return - } - cmd.execute(state) + cmd, ok := <-m.stateCmdChan + if !ok { + // Channel closed by Release(), exit cleanly + return + } + cmd.execute(state) - // Exit after processing completion - if state.completed { - // Drain remaining commands with timeout - drainTimer := time.NewTimer(100 * time.Millisecond) - defer drainTimer.Stop() - drainLoop: - for { - select { - case cmd := <-m.stateCmdChan: - cmd.execute(state) - case <-drainTimer.C: + // Optional: Exit after processing completion (but only after draining) + // This is mainly for optimization - the channel will be closed by Release() anyway + if state.completed { + // Drain remaining commands with timeout + drainTimer := time.NewTimer(100 * time.Millisecond) + defer drainTimer.Stop() + drainLoop: + for { + select { + case cmd, ok := <-m.stateCmdChan: + if !ok { + // Channel closed during drain break drainLoop } + cmd.execute(state) + case <-drainTimer.C: + break drainLoop } - return } - case <-m.ctx.Done(): - // Context cancelled - continue processing for a short time to handle cancellation - // Then exit to prevent deadlock - time.Sleep(10 * time.Millisecond) return } } @@ -391,26 +338,12 @@ func (m *manager) stateGetUpdates(since int64) []*types.TraceUpdate { } 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} } -func (m *manager) stateAddSubscriber(id string, ch chan *types.TraceUpdate) { - m.stateCmdChan <- &cmdAddSubscriber{id: id, ch: ch} -} - -func (m *manager) stateRemoveSubscriber(id string) { - m.stateCmdChan <- &cmdRemoveSubscriber{id: id} -} - -func (m *manager) stateGetSubscribers() map[string]chan *types.TraceUpdate { - resp := make(chan map[string]chan *types.TraceUpdate, 1) - m.stateCmdChan <- &cmdGetSubscribers{resp: resp} - return <-resp -} - -func (m *manager) stateBroadcast(update *types.TraceUpdate) { - m.stateCmdChan <- &cmdBroadcast{update: update} -} +// Subscription management methods removed - now handled by SubscriptionManager +// See subscription_manager.go and subscription.go for the new implementation // stateExecuteSpaceOp executes a space operation serially through state worker func (m *manager) stateExecuteSpaceOp(spaceID string, fn func() error) error { diff --git a/trace/subscription.go b/trace/subscription.go index dfe97fae..7537e36b 100644 --- a/trace/subscription.go +++ b/trace/subscription.go @@ -1,9 +1,8 @@ package trace import ( - "time" + "fmt" - gonanoid "github.com/matoous/go-nanoid/v2" "github.com/yaoapp/yao/trace/types" ) @@ -19,67 +18,22 @@ func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error) // subscribe is the internal implementation for subscriptions func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) { - // Generate unique subscriber ID - subID, _ := gonanoid.Generate("0123456789abcdefghijklmnopqrstuvwxyz", 12) - - // Create update channel - updateCh := make(chan *types.TraceUpdate, 100) - - // Register subscriber - m.stateAddSubscriber(subID, updateCh) - - // Start replay and stream goroutine (will auto-cleanup on completion) - go m.replayAndStream(subID, updateCh, since) - - return updateCh, nil -} - -// replayAndStream replays historical updates and streams new ones -func (m *manager) replayAndStream(subID string, ch chan *types.TraceUpdate, since int64) { - // Auto-cleanup on exit - MUST remove from map before closing channel - defer func() { - // Remove from subscribers map first to prevent new broadcasts - m.stateRemoveSubscriber(subID) - // Close channel (any in-flight broadcasts will be caught by recover) - close(ch) - }() - // Get historical updates updates := m.stateGetUpdates(since) - // Replay historical updates and check if trace was already completed - traceWasCompleted := false - for _, update := range updates { - select { - case ch <- update: - // Check if this is a trace complete event - if update.Type == types.UpdateTypeComplete { - traceWasCompleted = true - } - case <-m.ctx.Done(): - return - } + // Use manager's pubsub reference (always available) + if m.pubsub == nil { + return nil, fmt.Errorf("pubsub service not initialized for trace: %s", m.traceID) } - // If trace was already completed in historical events, exit immediately - if traceWasCompleted { - return + // Create subscription with historical replay + // Buffer size should be large enough to hold historical updates plus some live updates + // Using max of 1000 or len(updates)+100 to handle large traces + bufferSize := 1000 + if len(updates)+100 > bufferSize { + bufferSize = len(updates) + 100 } + sub := m.pubsub.SubscribeWithHistory(updates, bufferSize) - // Continue streaming new updates - // The channel will receive updates via broadcast from addUpdate - // Monitor completion to know when to exit - ticker := time.NewTicker(100 * time.Millisecond) - defer ticker.Stop() - - for { - select { - case <-ticker.C: - if m.stateIsCompleted() { - return - } - case <-m.ctx.Done(): - return - } - } + return sub.Channel, nil } diff --git a/trace/trace.go b/trace/trace.go index 2a3bc7fc..b2011d62 100644 --- a/trace/trace.go +++ b/trace/trace.go @@ -7,7 +7,9 @@ import ( "time" gonanoid "github.com/matoous/go-nanoid/v2" + "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/trace/local" + "github.com/yaoapp/yao/trace/pubsub" "github.com/yaoapp/yao/trace/store" "github.com/yaoapp/yao/trace/types" ) @@ -18,10 +20,14 @@ const ( Store = "store" // Gou store storage ) -// Global trace registry +// Global trace registry and pubsub services var ( registry = make(map[string]*types.TraceInfo) registryMu sync.RWMutex + + // Each trace has its own independent pubsub service + pubsubRegistry = make(map[string]*pubsub.PubSub) + pubsubRegistryMu sync.RWMutex ) // getDriver creates a driver instance based on driver type and options @@ -137,9 +143,23 @@ func New(ctx context.Context, driver string, option *types.TraceOption, driverOp return LoadFromStorage(ctx, driver, traceID, driverOptions...) } - // Create Manager instance with the driver - manager, err := NewManager(ctx, traceID, drv, option) + // Create independent PubSub service for this trace + pubsubService := pubsub.New() + + // Register pubsub service + pubsubRegistryMu.Lock() + pubsubRegistry[traceID] = pubsubService + pubsubRegistryMu.Unlock() + + // Create Manager instance with the driver and pubsub reference + // Manager uses pubsub only for publishing, doesn't manage its lifecycle + manager, err := NewManager(ctx, traceID, drv, pubsubService, option) if err != nil { + // Clean up pubsub if manager creation fails + pubsubRegistryMu.Lock() + delete(pubsubRegistry, traceID) + pubsubRegistryMu.Unlock() + pubsubService.Stop() return "", nil, fmt.Errorf("failed to create manager: %w", err) } @@ -176,6 +196,13 @@ func New(ctx context.Context, driver string, option *types.TraceOption, driverOp return traceID, manager, nil } +// GetPubSub returns the pubsub service for a trace +func GetPubSub(traceID string) *pubsub.PubSub { + pubsubRegistryMu.RLock() + defer pubsubRegistryMu.RUnlock() + return pubsubRegistry[traceID] +} + // Load loads an existing trace by ID from the registry // Returns: manager, error // traceID: the trace ID to load @@ -224,13 +251,29 @@ func LoadFromStorage(ctx context.Context, driver string, traceID string, driverO return "", nil, fmt.Errorf("trace not found in storage: %s", traceID) } - // Create Manager instance with the driver + // Create or reuse PubSub service for this trace + pubsubRegistryMu.Lock() + pubsubService, exists := pubsubRegistry[traceID] + if !exists { + pubsubService = pubsub.New() + pubsubRegistry[traceID] = pubsubService + } + pubsubRegistryMu.Unlock() + + // Create Manager instance with the driver and pubsub reference // Note: We need to reconstruct the manager from stored data // TODO: Implement proper restoration of manager state from storage // For loaded traces, we don't have the original option, so pass nil - manager, err := NewManager(ctx, traceID, drv, nil) + manager, err := NewManager(ctx, traceID, drv, pubsubService, nil) if err != nil { drv.Close() + if !exists { + // Clean up pubsub if we just created it + pubsubRegistryMu.Lock() + delete(pubsubRegistry, traceID) + pubsubRegistryMu.Unlock() + pubsubService.Stop() + } return "", nil, fmt.Errorf("failed to create manager: %w", err) } @@ -287,9 +330,150 @@ func GetInfo(ctx context.Context, driver string, traceID string, options ...any) return storedInfo, nil } +// MarkCancelled marks the trace as cancelled without using the context +// This is useful when the HTTP context has been cancelled and we can't use it for trace operations +// traceID: the trace ID to mark as cancelled +// reason: the cancellation reason +func MarkCancelled(traceID string, reason string) error { + log.Trace("[TRACE] MarkCancelled called: traceID=%s, reason=%s", traceID, reason) + + registryMu.RLock() + info, exists := registry[traceID] + registryMu.RUnlock() + + if !exists { + log.Trace("[TRACE] MarkCancelled: trace not found in registry") + return fmt.Errorf("trace not found in registry: %s", traceID) + } + + mgr, ok := info.Manager.(*manager) + if !ok { + log.Trace("[TRACE] MarkCancelled: invalid manager type") + return fmt.Errorf("invalid manager type for trace: %s", traceID) + } + + log.Trace("[TRACE] MarkCancelled: starting to mark nodes and trace as cancelled") + + // Get independent pubsub service + ps := GetPubSub(traceID) + if ps != nil { + log.Trace("[TRACE] MarkCancelled: current subscriber count: %d", ps.SubscriberCount()) + } + + now := time.Now().UnixMilli() + + // Use background context since the original context is cancelled + bgCtx := context.Background() + + // Load trace tree from driver (disk) + log.Trace("[TRACE] MarkCancelled: loading trace tree from driver") + rootNode, err := mgr.driver.LoadTrace(bgCtx, traceID) + if err != nil { + log.Trace("[TRACE] MarkCancelled: failed to load trace tree: %v", err) + return fmt.Errorf("failed to load trace tree: %w", err) + } + + if rootNode == nil { + log.Trace("[TRACE] MarkCancelled: no root node found") + return fmt.Errorf("no root node found for trace: %s", traceID) + } + + // Mark incomplete nodes as failed (recursively walk tree) + log.Trace("[TRACE] MarkCancelled: marking incomplete nodes as failed") + var markNodesFailed func(node *types.TraceNode) + markNodesFailed = func(node *types.TraceNode) { + if node.Status != types.StatusCompleted && node.Status != types.StatusFailed { + log.Trace("[TRACE] MarkCancelled: marking node %s as failed", node.ID) + node.Status = types.StatusFailed + node.EndTime = now + node.UpdatedAt = now + + // Save node to driver + if err := mgr.driver.SaveNode(bgCtx, traceID, node); err != nil { + log.Trace("[TRACE] MarkCancelled: failed to save node %s: %v", node.ID, err) + } + + // Broadcast node failed event (also saves to disk) + subscriberCount := 0 + if ps := GetPubSub(traceID); ps != nil { + subscriberCount = ps.SubscriberCount() + } + log.Trace("[TRACE] MarkCancelled: publishing node failed event for node %s to %d subscribers", node.ID, subscriberCount) + mgr.addUpdateAndBroadcast(&types.TraceUpdate{ + Type: types.UpdateTypeNodeFailed, + TraceID: traceID, + NodeID: node.ID, + Timestamp: now, + Data: &types.NodeFailedData{ + NodeID: node.ID, + Status: types.CompleteStatusFailed, + EndTime: now, + Duration: now - node.StartTime, + Error: reason, + }, + }) + log.Trace("[TRACE] MarkCancelled: node failed event broadcasted for node %s", node.ID) + } + + // Process children + for _, child := range node.Children { + markNodesFailed(child) + } + } + markNodesFailed(rootNode) + + // Load trace info + log.Trace("[TRACE] MarkCancelled: loading trace info from driver") + traceInfo, err := mgr.driver.LoadTraceInfo(bgCtx, traceID) + if err != nil { + log.Trace("[TRACE] MarkCancelled: failed to load trace info: %v", err) + return fmt.Errorf("failed to load trace info: %w", err) + } + + // Update trace status to cancelled + log.Trace("[TRACE] MarkCancelled: updating trace status to cancelled") + traceInfo.Status = types.TraceStatusCancelled + traceInfo.UpdatedAt = now + if err := mgr.driver.SaveTraceInfo(bgCtx, traceInfo); err != nil { + log.Trace("[TRACE] MarkCancelled: failed to save trace info: %v", err) + return fmt.Errorf("failed to save trace info: %w", err) + } + + // Set trace status in state machine + mgr.stateSetTraceStatus(types.TraceStatusCancelled) + mgr.stateMarkCompleted() + + // Broadcast completion update (saves to disk and publishes to subscribers) + subscriberCount := 0 + if ps := GetPubSub(traceID); ps != nil { + subscriberCount = ps.SubscriberCount() + } + log.Trace("[TRACE] MarkCancelled: publishing completion update to %d subscribers", subscriberCount) + totalDuration := int64(0) + if rootNode.CreatedAt > 0 { + totalDuration = now - rootNode.CreatedAt + } + + mgr.addUpdateAndBroadcast(&types.TraceUpdate{ + Type: types.UpdateTypeComplete, + TraceID: traceID, + Timestamp: now, + Data: &types.TraceCompleteData{ + TraceID: traceID, + Status: types.TraceStatusCancelled, + TotalDuration: totalDuration, + }, + }) + + log.Trace("[TRACE] MarkCancelled: completed successfully") + return nil +} + // Release releases a trace from the registry and closes its resources // traceID: the trace ID to release func Release(traceID string) error { + log.Trace("[TRACE] Release called: traceID=%s", traceID) + registryMu.Lock() info, exists := registry[traceID] if exists { @@ -298,14 +482,38 @@ func Release(traceID string) error { registryMu.Unlock() if !exists { + log.Trace("[TRACE] Release: trace not found in registry") return fmt.Errorf("trace not found in registry: %s", traceID) } - // Cancel the manager's context to stop background goroutines - if mgr, ok := info.Manager.(*manager); ok && mgr.cancel != nil { - mgr.cancel() + // Stop manager + if mgr, ok := info.Manager.(*manager); ok { + // Close state machine channel to stop state worker goroutine + log.Trace("[TRACE] Release: closing state command channel") + close(mgr.stateCmdChan) + + // Cancel the manager's context to stop other background operations + if mgr.cancel != nil { + log.Trace("[TRACE] Release: cancelling manager context") + mgr.cancel() + } } + // Stop independent PubSub service + pubsubRegistryMu.Lock() + ps, psExists := pubsubRegistry[traceID] + if psExists { + delete(pubsubRegistry, traceID) + } + pubsubRegistryMu.Unlock() + + if psExists && ps != nil { + subscriberCount := ps.SubscriberCount() + log.Trace("[TRACE] Release: stopping pubsub service with %d active subscribers", subscriberCount) + ps.Stop() + } + + log.Trace("[TRACE] Release: completed") return nil }