opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor interrupt handling and enhance logging in Assistant and Context components

Browse source 
- 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.
This commit is contained in:
Max 2025-11-21 20:52:10 +08:00
parent c4df6e51c1
commit a49695e906
16 changed files with 722 additions and 358 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

112
agent/assistant/agent.go
View file

@ -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 {

3
agent/assistant/agent_interrupt_test.go
View file

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

43
agent/context/context.go
View file

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

18
agent/context/interrupt.go
View file

@ -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

3
agent/context/interrupt_test.go
View file

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

7
agent/context/openapi.go
View file

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

2
agent/llm/providers/openai/gpt5_test.go
View file

@ -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",
},
},
},

157
agent/llm/providers/openai/openai.go
View file

@ -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 {

69
agent/llm/providers/openai/openai_test.go
View file

@ -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

8
openapi/chat/completions.go
View file

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

29
trace/manager.go
View file

@ -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

142
trace/pubsub/pubsub.go Normal file
View file

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

62
trace/pubsub/subscriber.go Normal file
View file

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

131
trace/state.go
View file

@ -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 {

70
trace/subscription.go
View file

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

224
trace/trace.go
View file

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