From af981dc0b070456bb1783f8d6ff98c189f253bd3 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 23 Feb 2026 19:22:31 +0800 Subject: [PATCH] Update dependencies and enhance logging in the Assistant module - Add new indirect dependencies including various Charmbracelet packages for improved UI handling. - Enhance logging in the Assistant module by adding tool completion and start logging for better traceability of tool calls. - Modify context handling in the RequestLogger to support a stack-based assistant ID management, improving the logging structure for agent requests. - Implement event service integration for better trace management and debugging capabilities. --- agent/assistant/agent.go | 1 + agent/assistant/mcp.go | 27 +- agent/caller/orchestrator.go | 48 ++ agent/context/context.go | 28 +- agent/context/log.go | 393 ++++++++++++----- agent/context/tui.go | 799 ++++++++++++++++++++++++++++++++++ agent/context/tui_msg.go | 90 ++++ cmd/start.go | 42 +- engine/load.go | 15 + go.mod | 14 + go.sum | 29 ++ test/utils.go | 19 +- trace/event_listener.go | 25 ++ trace/handler.go | 28 ++ trace/manager.go | 154 +++---- trace/node.go | 23 +- trace/pubsub/pubsub.go | 142 ------ trace/pubsub/subscriber.go | 62 --- trace/space.go | 20 +- trace/state.go | 362 ++++----------- trace/subscription.go | 78 +++- trace/trace.go | 105 +---- trace/trace_basic_test.go | 13 +- trace/trace_lifecycle_test.go | 25 +- 24 files changed, 1678 insertions(+), 864 deletions(-) create mode 100644 agent/context/tui.go create mode 100644 agent/context/tui_msg.go create mode 100644 trace/event_listener.go create mode 100644 trace/handler.go delete mode 100644 trace/pubsub/pubsub.go delete mode 100644 trace/pubsub/subscriber.go diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 5d9f070d..e1fcdda0 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -103,6 +103,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // Log end of request ctx.Logger.End(finalStatus == context.StepStatusCompleted, finalError) + ctx.Logger.RestoreAssistantID() }() // Determine stream handler diff --git a/agent/assistant/mcp.go b/agent/assistant/mcp.go index 1e544e00..4696f230 100644 --- a/agent/assistant/mcp.go +++ b/agent/assistant/mcp.go @@ -319,6 +319,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Content = result.Error.Error() result.IsRetryableError = true // Argument parsing error is retryable by LLM ctx.Logger.Error("Failed to parse arguments: %v", err) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } @@ -333,6 +334,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Content = result.Error.Error() result.IsRetryableError = true // Type error is retryable by LLM ctx.Logger.Error("Arguments must be an object, got %T", parsed) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } @@ -346,6 +348,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Content = result.Error.Error() result.IsRetryableError = true // Validation error is retryable by LLM ctx.Logger.Error("Argument validation failed: %v", err) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } @@ -393,7 +396,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall } result.Content = string(contentBytes) - ctx.Logger.ToolComplete(toolName, true) + ctx.Logger.ToolComplete(toolCall.Function.Name, true) if toolNode != nil { toolNode.Complete(map[string]any{ @@ -563,6 +566,7 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context Arguments: args, }) callMap[toolName] = tc + ctx.Logger.ToolStart(tc.Function.Name) // Add trace input for this tool parallelInputs = append(parallelInputs, types.TraceParallelInput{ @@ -598,11 +602,15 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context mcpResponse, err := client.CallToolsParallel(mcpCtx, mcpCalls, ctx) if err != nil { ctx.Logger.Error("Parallel call failed: %v", err) - // Mark all trace nodes as failed - for _, node := range toolNodes { + for i, node := range toolNodes { if node != nil { node.Fail(err) } + if i < len(mcpCalls) { + if tc, ok := callMap[mcpCalls[i].Name]; ok { + ctx.Logger.ToolComplete(tc.Function.Name, false) + } + } } return nil, true } @@ -631,6 +639,7 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context result.Content = result.Error.Error() result.IsRetryableError = false // Serialization error is not retryable hasErrors = true + ctx.Logger.ToolComplete(originalCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } @@ -643,11 +652,12 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context result.IsRetryableError = isRetryableToolError(result.Error) hasErrors = true ctx.Logger.Error("Tool call failed: %s - %s (retryable: %v)", toolName, result.Content, result.IsRetryableError) + ctx.Logger.ToolComplete(originalCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } } else { - // Success + ctx.Logger.ToolComplete(originalCall.Function.Name, true) if toolNode != nil { toolNode.Complete(map[string]any{ "result": mcpResult.Content, @@ -670,6 +680,8 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte ctx.Logger.Debug("Calling %d tools sequentially on server '%s'", len(toolCalls), serverID) for _, tc := range toolCalls { + ctx.Logger.ToolStart(tc.Function.Name) + _, toolName, ok := ParseMCPToolName(tc.Function.Name) if !ok { results = append(results, ToolCallResult{ @@ -678,6 +690,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte Content: fmt.Sprintf("Invalid tool name format: %s", tc.Function.Name), Error: fmt.Errorf("invalid tool name format"), }) + ctx.Logger.ToolComplete(tc.Function.Name, false) hasErrors = true continue } @@ -727,6 +740,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte } results = append(results, result) hasErrors = true + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } @@ -747,6 +761,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte } results = append(results, result) hasErrors = true + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } @@ -765,6 +780,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte } results = append(results, result) hasErrors = true + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } @@ -788,6 +804,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte result.IsRetryableError = isRetryableToolError(err) hasErrors = true ctx.Logger.Error("Tool call failed: %s - %v (retryable: %v)", toolName, err, result.IsRetryableError) + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } @@ -806,11 +823,13 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte result.Content = fmt.Sprintf("Failed to serialize result: %v", err) result.IsRetryableError = false // Serialization error is not retryable hasErrors = true + ctx.Logger.ToolComplete(tc.Function.Name, false) if toolNode != nil { toolNode.Fail(err) } } else { result.Content = string(contentBytes) + ctx.Logger.ToolComplete(tc.Function.Name, !mcpResult.IsError) if toolNode != nil { toolNode.Complete(map[string]any{ "result": mcpResult.Content, diff --git a/agent/caller/orchestrator.go b/agent/caller/orchestrator.go index 65589c0a..fdf78655 100644 --- a/agent/caller/orchestrator.go +++ b/agent/caller/orchestrator.go @@ -5,6 +5,7 @@ import ( "sync" agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/trace/types" ) // Orchestrator handles parallel agent calls with different concurrency patterns @@ -257,13 +258,60 @@ func (o *Orchestrator) callAgentWithContext(ctx *agentContext.Context, req *Requ ctxOpts.OnMessage = req.Handler } + // Add trace node for A2A call using the ORIGINAL parent context's trace + // (forked contexts have nil trace and nil Stack, so ctx.Trace() would create a new orphan trace) + parentTrace, _ := o.ctx.Trace() + var a2aNode types.Node + if parentTrace != nil { + a2aNode, _ = parentTrace.Add( + map[string]any{ + "agent_id": req.AgentID, + "referer": string(ctx.Referer), + }, + types.TraceNodeOption{ + Label: fmt.Sprintf("Agent: %s", req.AgentID), + Type: "agent_call", + Icon: "smart_toy", + Description: fmt.Sprintf("A2A call to '%s'", req.AgentID), + }, + ) + } + + // Notify TUI of A2A call start (use parent requestID so it appears in parent panel) + parentRequestID := o.ctx.RequestID() + agentContext.SendTUI(agentContext.AgentEventMsg{ + RequestID: parentRequestID, + Event: agentContext.EventA2AStart, + Data: map[string]interface{}{"target": req.AgentID}, + }) + // Execute the agent call with the provided context // The agent.Stream method will use the context's Writer for output resp, err := agent.Stream(ctx, req.Messages, ctxOpts) if err != nil { + if a2aNode != nil { + a2aNode.Fail(err) + } + agentContext.SendTUI(agentContext.AgentEventMsg{ + RequestID: parentRequestID, + Event: agentContext.EventA2ADone, + Data: map[string]interface{}{"target": req.AgentID, "error": err.Error()}, + }) return NewResult(req.AgentID, nil, fmt.Errorf("agent call failed: %w", err)) } + if a2aNode != nil { + a2aNode.Complete(map[string]any{ + "agent_id": req.AgentID, + "status": "completed", + }) + } + agentContext.SendTUI(agentContext.AgentEventMsg{ + RequestID: parentRequestID, + Event: agentContext.EventA2ADone, + Data: map[string]interface{}{"target": req.AgentID}, + }) + return NewResult(req.AgentID, resp, nil) } diff --git a/agent/context/context.go b/agent/context/context.go index b419cf49..05582029 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -71,19 +71,21 @@ func (ctx *Context) Release() { ctx.Interrupt = nil } - // Complete and release trace if exists + // Complete and release trace if exists. + // Only the root context (non-forked) owns the trace lifecycle. + // Forked contexts share the same trace manager but must not release it. if ctx.trace != nil && ctx.Stack != nil && ctx.Stack.TraceID != "" { - if ctx.Logger != nil { - ctx.Logger.Cleanup("Trace: " + ctx.Stack.TraceID) - } - - // Check if context is cancelled - if so, mark as cancelled instead of complete - if ctx.Context != nil && ctx.Context.Err() != nil { - trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error()) - trace.Release(ctx.Stack.TraceID) - } else { - ctx.trace.MarkComplete() - trace.Release(ctx.Stack.TraceID) + if ctx.ForkParent == nil { + if ctx.Logger != nil { + ctx.Logger.Cleanup("Trace: " + ctx.Stack.TraceID) + } + if ctx.Context != nil && ctx.Context.Err() != nil { + trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error()) + trace.Release(ctx.Stack.TraceID) + } else { + ctx.trace.MarkComplete() + trace.Release(ctx.Stack.TraceID) + } } ctx.trace = nil } @@ -191,7 +193,7 @@ func (ctx *Context) Fork() *Context { // Create independent resources to avoid race conditions IDGenerator: message.NewIDGenerator(), - Logger: NewRequestLogger(ctx.AssistantID, ctx.ChatID, childID), + Logger: NewRequestLogger(ctx.AssistantID, ctx.ChatID, childID, WithParentID(ctx.ID)), messageMetadata: newMessageMetadataStore(), // Inherit context metadata diff --git a/agent/context/log.go b/agent/context/log.go index 16658e42..503fc3fe 100644 --- a/agent/context/log.go +++ b/agent/context/log.go @@ -72,11 +72,12 @@ type LogEntry struct { // RequestLogger provides request-scoped async logging type RequestLogger struct { - assistantID string - chatID string - requestID string - shortID string // Short version of requestID for display - startTime time.Time + assistantIDStack []string // Stack-based: delegate calls push, pop on exit; top = current + chatID string + requestID string + shortID string // Short version of requestID for display + parentID string // Parent request ID for A2A tree structure + startTime time.Time ch chan LogEntry done chan struct{} @@ -86,19 +87,33 @@ type RequestLogger struct { mu sync.RWMutex } +// LoggerOption configures a RequestLogger +type LoggerOption func(*RequestLogger) + +// WithParentID sets the parent request ID for A2A tree structure +func WithParentID(parentID string) LoggerOption { + return func(l *RequestLogger) { + l.parentID = parentID + } +} + // noopLogger is a shared no-op logger instance var noopLogger = &RequestLogger{noop: true} // NewRequestLogger creates a new request-scoped logger with async processing -func NewRequestLogger(assistantID, chatID, requestID string) *RequestLogger { +func NewRequestLogger(assistantID, chatID, requestID string, opts ...LoggerOption) *RequestLogger { l := &RequestLogger{ - assistantID: assistantID, - chatID: chatID, - requestID: requestID, - shortID: shortID(requestID), - startTime: time.Now(), - ch: make(chan LogEntry, 100), // Buffered channel - done: make(chan struct{}), + assistantIDStack: []string{assistantID}, + chatID: chatID, + requestID: requestID, + shortID: shortID(requestID), + startTime: time.Now(), + ch: make(chan LogEntry, 100), // Buffered channel + done: make(chan struct{}), + } + + for _, opt := range opts { + opt(l) } // Start consumer goroutine @@ -112,12 +127,35 @@ func Noop() *RequestLogger { return noopLogger } -// SetAssistantID sets the assistant ID (called when entering Stream) +// SetAssistantID pushes a new assistant ID onto the stack (called when entering Stream). +// Each SetAssistantID must be paired with a RestoreAssistantID on exit. func (l *RequestLogger) SetAssistantID(id string) { if l.noop { return } - l.assistantID = id + l.mu.Lock() + l.assistantIDStack = append(l.assistantIDStack, id) + l.mu.Unlock() +} + +// RestoreAssistantID pops the current assistant ID, reverting to the previous one. +// Safe to call even if the stack has only one entry (the initial ID is never removed). +func (l *RequestLogger) RestoreAssistantID() { + if l.noop { + return + } + l.mu.Lock() + if len(l.assistantIDStack) > 1 { + l.assistantIDStack = l.assistantIDStack[:len(l.assistantIDStack)-1] + } + l.mu.Unlock() +} + +func (l *RequestLogger) currentAssistantID() string { + if len(l.assistantIDStack) == 0 { + return "" + } + return l.assistantIDStack[len(l.assistantIDStack)-1] } // Close closes the logger and waits for all entries to be processed @@ -148,13 +186,17 @@ func (l *RequestLogger) consume() { func (l *RequestLogger) processEntry(entry LogEntry) { if config.IsDevelopment() { l.printDev(entry) + l.writeLog(entry, true) } else { - l.printProd(entry) + l.writeLog(entry, false) } } -// printDev prints colorful output for development mode +// printDev sends to TUI if available, otherwise prints colored output to stdout func (l *RequestLogger) printDev(entry LogEntry) { + if GetTUIProgram() != nil { + return + } switch entry.Level { case LogLevelTrace: fmt.Printf("%s β†’ %s%s\n", colorGray, entry.Message, colorReset) @@ -169,10 +211,13 @@ func (l *RequestLogger) printDev(entry LogEntry) { } } -// printProd logs to kun/log for production mode -func (l *RequestLogger) printProd(entry LogEntry) { +// writeLog writes structured events to kun/log +func (l *RequestLogger) writeLog(entry LogEntry, devMode bool) { prefix := fmt.Sprintf("[AGENT] %s ", l.shortID) - + if devMode { + kunlog.Trace("%s%s", prefix, entry.Message) + return + } switch entry.Level { case LogLevelTrace: kunlog.Trace("%s%s", prefix, entry.Message) @@ -263,18 +308,29 @@ func (l *RequestLogger) Start() { return } + kunlog.Trace("[AGENT] Request %s started: assistant=%s, chat=%s, request=%s", + l.shortID, l.currentAssistantID(), shortID(l.chatID), shortID(l.requestID)) + if !config.IsDevelopment() { - kunlog.Trace("[AGENT] Request %s started: assistant=%s, chat=%s, request=%s", - l.shortID, l.assistantID, shortID(l.chatID), shortID(l.requestID)) return } - // Development: colorful output (direct print, not through channel for immediate display) + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + ParentID: l.parentID, + AssistantID: l.currentAssistantID(), + Event: EventRequestStart, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Println() fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("═", 60), colorReset) - fmt.Printf("%s πŸš€ AGENT REQUEST %s%s\n", colorBoldCyan, l.shortID, colorReset) + fmt.Printf("%s AGENT REQUEST %s%s\n", colorBoldCyan, l.shortID, colorReset) fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("─", 60), colorReset) - fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.assistantID, colorReset) + fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.currentAssistantID(), colorReset) fmt.Printf("%s Chat ID: %s%s%s\n", colorGray, colorWhite, l.chatID, colorReset) fmt.Printf("%s Request: %s%s%s\n", colorGray, colorWhite, l.requestID, colorReset) fmt.Printf("%s Time: %s%s%s\n", colorGray, colorWhite, l.startTime.Format("15:04:05.000"), colorReset) @@ -289,28 +345,43 @@ func (l *RequestLogger) End(success bool, err error) { duration := time.Since(l.startTime) + if success { + kunlog.Trace("[AGENT] Request %s completed: assistant=%s, duration=%v", + l.shortID, l.currentAssistantID(), duration.Round(time.Millisecond)) + } else { + kunlog.Error("[AGENT] Request %s failed: assistant=%s, duration=%v, error=%v", + l.shortID, l.currentAssistantID(), duration.Round(time.Millisecond), err) + } + if !config.IsDevelopment() { - if success { - kunlog.Trace("[AGENT] Request %s completed: assistant=%s, duration=%v", - l.shortID, l.assistantID, duration.Round(time.Millisecond)) - } else { - kunlog.Trace("[AGENT] Request %s failed: assistant=%s, duration=%v, error=%v", - l.shortID, l.assistantID, duration.Round(time.Millisecond), err) - } return } - // Development: colorful output (direct print for immediate display) + data := map[string]interface{}{"duration": duration.Round(time.Millisecond)} + if err != nil { + data["error"] = err.Error() + } + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + AssistantID: l.currentAssistantID(), + Event: EventRequestEnd, + Data: data, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset) if success { - fmt.Printf("%s βœ… REQUEST %s COMPLETED%s\n", colorBoldGreen, l.shortID, colorReset) + fmt.Printf("%s REQUEST %s COMPLETED%s\n", colorBoldGreen, l.shortID, colorReset) } else { - fmt.Printf("%s ❌ REQUEST %s FAILED%s\n", colorBoldRed, l.shortID, colorReset) + fmt.Printf("%s REQUEST %s FAILED%s\n", colorBoldRed, l.shortID, colorReset) if err != nil { fmt.Printf("%s Error: %s%v%s\n", colorGray, colorRed, err, colorReset) } } - fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.assistantID, colorReset) + fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.currentAssistantID(), colorReset) fmt.Printf("%s Duration: %s%v%s\n", colorGray, colorWhite, duration.Round(time.Millisecond), colorReset) fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset) fmt.Println() @@ -323,12 +394,22 @@ func (l *RequestLogger) Phase(name string) { } elapsed := time.Since(l.startTime).Round(time.Millisecond) + kunlog.Trace("[AGENT] %s Phase: %s (+%v)", l.shortID, name, elapsed) - if config.IsDevelopment() { - fmt.Printf("%s β–Ά %s%s %s[+%v]%s\n", colorBoldBlue, name, colorReset, colorGray, elapsed, colorReset) - } else { - kunlog.Trace("[AGENT] %s Phase: %s (+%v)", l.shortID, name, elapsed) + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventPhase, + Data: map[string]interface{}{"name": name, "elapsed": elapsed}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s > %s%s %s[+%v]%s\n", colorBoldBlue, name, colorReset, colorGray, elapsed, colorReset) } // PhaseComplete logs the completion of a phase @@ -338,12 +419,22 @@ func (l *RequestLogger) PhaseComplete(name string) { } elapsed := time.Since(l.startTime).Round(time.Millisecond) + kunlog.Trace("[AGENT] %s Phase completed: %s (+%v)", l.shortID, name, elapsed) - if config.IsDevelopment() { - fmt.Printf("%s βœ“ %s%s %s[+%v]%s\n", colorGreen, name, colorReset, colorGray, elapsed, colorReset) - } else { - kunlog.Trace("[AGENT] %s Phase completed: %s (+%v)", l.shortID, name, elapsed) + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventPhaseDone, + Data: map[string]interface{}{"name": name, "elapsed": elapsed}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s + %s%s %s[+%v]%s\n", colorGreen, name, colorReset, colorGray, elapsed, colorReset) } // PhaseSkip logs a skipped phase (development only) @@ -351,9 +442,21 @@ func (l *RequestLogger) PhaseSkip(name, reason string) { if l.noop { return } - if config.IsDevelopment() { - fmt.Printf("%s ⊘ %s (%s)%s\n", colorGray, name, reason, colorReset) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventPhaseSkip, + Data: map[string]interface{}{"name": name, "reason": reason}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s - %s (%s)%s\n", colorGray, name, reason, colorReset) } // LLMStart logs the start of an LLM call @@ -363,17 +466,31 @@ func (l *RequestLogger) LLMStart(connector, model string, messageCount int) { } elapsed := time.Since(l.startTime).Round(time.Millisecond) + kunlog.Trace("[AGENT] %s LLM call: connector=%s, model=%s, messages=%d (+%v)", l.shortID, connector, model, messageCount, elapsed) - if config.IsDevelopment() { - fmt.Printf("%s πŸ€– LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset) - fmt.Printf("%s Connector: %s%s%s\n", colorGray, colorWhite, connector, colorReset) - if model != "" { - fmt.Printf("%s Model: %s%s%s\n", colorGray, colorWhite, model, colorReset) - } - fmt.Printf("%s Messages: %s%d%s\n", colorGray, colorWhite, messageCount, colorReset) - } else { - kunlog.Trace("[AGENT] %s LLM call: connector=%s, model=%s, messages=%d (+%v)", l.shortID, connector, model, messageCount, elapsed) + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventLLMCall, + Data: map[string]interface{}{ + "connector": connector, + "model": model, + "messages": messageCount, + }, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset) + fmt.Printf("%s Connector: %s%s%s\n", colorGray, colorWhite, connector, colorReset) + if model != "" { + fmt.Printf("%s Model: %s%s%s\n", colorGray, colorWhite, model, colorReset) + } + fmt.Printf("%s Messages: %s%d%s\n", colorGray, colorWhite, messageCount, colorReset) } // LLMComplete logs the completion of an LLM call @@ -388,15 +505,30 @@ func (l *RequestLogger) LLMComplete(tokens int, hasToolCalls bool) { status = "tool_calls" } - if config.IsDevelopment() { - fmt.Printf("%s βœ“ LLM Response (%s)%s", colorGreen, status, colorReset) - if tokens > 0 { - fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset) - } - fmt.Printf(" %s[+%v]%s\n", colorGray, elapsed, colorReset) - } else { - kunlog.Trace("[AGENT] %s LLM response: status=%s, tokens=%d (+%v)", l.shortID, status, tokens, elapsed) + kunlog.Trace("[AGENT] %s LLM response: status=%s, tokens=%d (+%v)", l.shortID, status, tokens, elapsed) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventLLMDone, + Data: map[string]interface{}{ + "detail": fmt.Sprintf("%s [tokens:%d, %v]", status, tokens, elapsed), + "tokens": tokens, + "status": status, + }, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s + LLM Response (%s)%s", colorGreen, status, colorReset) + if tokens > 0 { + fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset) + } + fmt.Printf(" %s[+%v]%s\n", colorGray, elapsed, colorReset) } // ToolStart logs the start of tool execution @@ -405,11 +537,22 @@ func (l *RequestLogger) ToolStart(toolName string) { return } - if config.IsDevelopment() { - fmt.Printf("%s πŸ”§ Tool: %s%s\n", colorYellow, toolName, colorReset) - } else { - kunlog.Trace("[AGENT] %s Tool call: %s", l.shortID, toolName) + kunlog.Trace("[AGENT] %s Tool call: %s", l.shortID, toolName) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventToolCall, + Data: map[string]interface{}{"name": toolName}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s Tool: %s%s\n", colorYellow, toolName, colorReset) } // ToolComplete logs the completion of tool execution @@ -418,18 +561,29 @@ func (l *RequestLogger) ToolComplete(toolName string, success bool) { return } - if config.IsDevelopment() { - if success { - fmt.Printf("%s βœ“ %s completed%s\n", colorGreen, toolName, colorReset) - } else { - fmt.Printf("%s βœ— %s failed%s\n", colorRed, toolName, colorReset) - } + if success { + kunlog.Trace("[AGENT] %s Tool completed: %s", l.shortID, toolName) } else { - if success { - kunlog.Trace("[AGENT] %s Tool completed: %s", l.shortID, toolName) - } else { - kunlog.Trace("[AGENT] %s Tool failed: %s", l.shortID, toolName) - } + kunlog.Error("[AGENT] %s Tool failed: %s", l.shortID, toolName) + } + + if !config.IsDevelopment() { + return + } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventToolDone, + Data: map[string]interface{}{"name": toolName, "success": success}, + }) + + if GetTUIProgram() != nil { + return + } + if success { + fmt.Printf("%s + %s completed%s\n", colorGreen, toolName, colorReset) + } else { + fmt.Printf("%s x %s failed%s\n", colorRed, toolName, colorReset) } } @@ -440,12 +594,22 @@ func (l *RequestLogger) HookStart(hookName string) { } elapsed := time.Since(l.startTime).Round(time.Millisecond) + kunlog.Trace("[AGENT] %s Hook: %s (+%v)", l.shortID, hookName, elapsed) - if config.IsDevelopment() { - fmt.Printf("%s πŸͺ Hook: %s%s %s[+%v]%s\n", colorMagenta, hookName, colorReset, colorGray, elapsed, colorReset) - } else { - kunlog.Trace("[AGENT] %s Hook: %s (+%v)", l.shortID, hookName, elapsed) + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventHook, + Data: map[string]interface{}{"name": hookName}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s Hook: %s%s %s[+%v]%s\n", colorMagenta, hookName, colorReset, colorGray, elapsed, colorReset) } // HookComplete logs the completion of a hook @@ -454,11 +618,22 @@ func (l *RequestLogger) HookComplete(hookName string) { return } - if config.IsDevelopment() { - fmt.Printf("%s βœ“ %s done%s\n", colorGreen, hookName, colorReset) - } else { - kunlog.Trace("[AGENT] %s Hook completed: %s", l.shortID, hookName) + kunlog.Trace("[AGENT] %s Hook completed: %s", l.shortID, hookName) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventHookDone, + Data: map[string]interface{}{"name": hookName}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s + %s done%s\n", colorGreen, hookName, colorReset) } // Cleanup logs resource cleanup @@ -467,11 +642,12 @@ func (l *RequestLogger) Cleanup(resource string) { return } - if config.IsDevelopment() { - fmt.Printf("%s βœ“ %s%s\n", colorGray, resource, colorReset) - } else { - kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource) + kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource) + + if !config.IsDevelopment() || GetTUIProgram() != nil { + return } + fmt.Printf("%s + %s%s\n", colorGray, resource, colorReset) } // HistoryLoad logs history loading @@ -480,11 +656,12 @@ func (l *RequestLogger) HistoryLoad(count, maxSize int) { return } - if config.IsDevelopment() { - fmt.Printf("%s Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset) - } else { - kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize) + kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize) + + if !config.IsDevelopment() || GetTUIProgram() != nil { + return } + fmt.Printf("%s Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset) } // HistoryOverlap logs overlap detection @@ -494,11 +671,12 @@ func (l *RequestLogger) HistoryOverlap(overlapCount int) { } if overlapCount > 0 { - if config.IsDevelopment() { - fmt.Printf("%s Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset) - } else { - kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount) + kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount) + + if !config.IsDevelopment() || GetTUIProgram() != nil { + return } + fmt.Printf("%s Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset) } } @@ -508,11 +686,22 @@ func (l *RequestLogger) Release() { return } - if config.IsDevelopment() { - fmt.Printf("%s 🧹 RELEASE %s%s %s(%s)%s\n", colorBoldYellow, l.shortID, colorReset, colorGray, l.assistantID, colorReset) - } else { - kunlog.Trace("[AGENT] %s Release started", l.shortID) + kunlog.Trace("[AGENT] %s Release started", l.shortID) + + if !config.IsDevelopment() { + return } + + SendTUI(AgentEventMsg{ + RequestID: l.requestID, + Event: EventContextRelease, + Data: map[string]interface{}{"assistant": l.currentAssistantID()}, + }) + + if GetTUIProgram() != nil { + return + } + fmt.Printf("%s RELEASE %s%s %s(%s)%s\n", colorBoldYellow, l.shortID, colorReset, colorGray, l.currentAssistantID(), colorReset) } // ============================================================================= diff --git a/agent/context/tui.go b/agent/context/tui.go new file mode 100644 index 00000000..5062bcc9 --- /dev/null +++ b/agent/context/tui.go @@ -0,0 +1,799 @@ +package context + +import ( + "fmt" + "strings" + "sync" + "time" + + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" +) + +var ( + tuiProgram *tea.Program + tuiProgramMu sync.RWMutex +) + +// SetTUIProgram sets the global TUI program (called from start.go after HTTP READY) +func SetTUIProgram(p *tea.Program) { + tuiProgramMu.Lock() + tuiProgram = p + tuiProgramMu.Unlock() +} + +// GetTUIProgram returns the global TUI program (nil if not in TUI mode) +func GetTUIProgram() *tea.Program { + tuiProgramMu.RLock() + defer tuiProgramMu.RUnlock() + return tuiProgram +} + +// SendTUI sends a message to the TUI program if available +func SendTUI(msg tea.Msg) { + if p := GetTUIProgram(); p != nil { + p.Send(msg) + } +} + +// TUILogWriter implements io.Writer to bridge gou DevWriter -> TUI AppLogMsg +type TUILogWriter struct { + Program *tea.Program +} + +func (w *TUILogWriter) Write(p []byte) (n int, err error) { + content := strings.TrimRight(string(p), "\n") + if content == "" { + return len(p), nil + } + w.Program.Send(AppLogMsg{Content: content}) + return len(p), nil +} + +// ─── Styles ─────────────────────────────────────────────────────────────────── + +var ( + boxRunning = lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(lipgloss.Color("33")). + PaddingLeft(1).PaddingRight(1) + + boxDone = lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(lipgloss.Color("240")). + PaddingLeft(1).PaddingRight(1) + + boxFailed = lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(lipgloss.Color("31")). + PaddingLeft(1).PaddingRight(1) + + boxAppLog = lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(lipgloss.Color("240")). + PaddingLeft(1).PaddingRight(1) + + sRunning = lipgloss.NewStyle().Foreground(lipgloss.Color("33")) + sDone = lipgloss.NewStyle().Foreground(lipgloss.Color("34")) + sFailed = lipgloss.NewStyle().Foreground(lipgloss.Color("31")) + sDim = lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + sBold = lipgloss.NewStyle().Bold(true) + sYellow = lipgloss.NewStyle().Foreground(lipgloss.Color("33")) + sRed = lipgloss.NewStyle().Foreground(lipgloss.Color("31")) + sBlue = lipgloss.NewStyle().Foreground(lipgloss.Color("34")) + sMagenta = lipgloss.NewStyle().Foreground(lipgloss.Color("35")) + sTree = lipgloss.NewStyle().Foreground(lipgloss.Color("240")) +) + +// ─── Data ───────────────────────────────────────────────────────────────────── + +// RequestPanel represents a single top-level agent request +type RequestPanel struct { + RequestID string + ShortID string + AssistantID string + StartTime time.Time + EndTime time.Time // set when done/failed, freezes elapsed display + Status PanelStatus + Nodes []TreeNode + ParentID string + Collapsed bool + viewRow int // Y offset of the header line (for mouse click) +} + +// TreeNode represents a step within a request panel +type TreeNode struct { + Kind NodeKind + Label string + Status NodeStatus + Detail string + Children []*TreeNode + StartTime time.Time + EndTime time.Time + Collapsed bool +} + +// AgentTUIModel is the bubbletea Model for agent request visualization +type AgentTUIModel struct { + panels []*RequestPanel + panelIndex map[string]int // requestID -> index in panels (first registration wins) + appLogs []AppLogEntry + appLogExpand bool + appLogRow int // Y offset of app log header + cursor int + width int + height int + scrollOffset int + autoFollow bool // auto-scroll to bottom when new content arrives + mouseOn bool + quitting bool +} + +// NewAgentTUIModel creates a new TUI model +func NewAgentTUIModel() AgentTUIModel { + return AgentTUIModel{ + panels: []*RequestPanel{}, + panelIndex: map[string]int{}, + appLogs: []AppLogEntry{}, + width: 80, + height: 24, + autoFollow: true, + } +} + +func (m AgentTUIModel) Init() tea.Cmd { + return tickCmd() +} + +func tickCmd() tea.Cmd { + return tea.Tick(500*time.Millisecond, func(t time.Time) tea.Msg { + return TickMsg(t) + }) +} + +// ─── Update ─────────────────────────────────────────────────────────────────── + +func (m AgentTUIModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.width = msg.Width + m.height = msg.Height + return m, nil + + case tea.KeyMsg: + return m.handleKey(msg) + + case tea.MouseMsg: + return m.handleMouse(msg) + + case AgentEventMsg: + return m.handleAgentEvent(msg), nil + + case AppLogMsg: + m.appLogs = append(m.appLogs, AppLogEntry{ + Content: msg.Content, + Time: time.Now(), + }) + return m, nil + + case TickMsg: + return m, tickCmd() + } + return m, nil +} + +func (m AgentTUIModel) handleKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { + topPanels := m.topLevelPanels() + total := len(topPanels) + 1 // +1 for app log + viewH := m.viewHeight() + + switch msg.String() { + case "q", "ctrl+c": + m.quitting = true + return m, tea.Quit + + // Scrolling + case "j", "down": + m.scrollOffset++ + m.autoFollow = false + case "k", "up": + if m.scrollOffset > 0 { + m.scrollOffset-- + } + m.autoFollow = false + case "pgdown", "ctrl+d": + m.scrollOffset += viewH / 2 + m.autoFollow = false + case "pgup", "ctrl+u": + m.scrollOffset -= viewH / 2 + if m.scrollOffset < 0 { + m.scrollOffset = 0 + } + m.autoFollow = false + case "G", "end": + m.autoFollow = true + case "g", "home": + m.scrollOffset = 0 + m.autoFollow = false + + // Cursor navigation for panel selection (wraps around) + case "tab": + m.cursor = (m.cursor + 1) % total + m.scrollToCursor(topPanels) + case "shift+tab": + m.cursor = (m.cursor - 1 + total) % total + m.scrollToCursor(topPanels) + + case "enter", " ": + if m.cursor < len(topPanels) { + topPanels[m.cursor].Collapsed = !topPanels[m.cursor].Collapsed + } else { + m.appLogExpand = !m.appLogExpand + } + case "c": + m.appLogExpand = !m.appLogExpand + case "a": + for _, p := range m.panels { + p.Collapsed = false + } + m.appLogExpand = true + case "A": + for _, p := range m.panels { + p.Collapsed = true + } + m.appLogExpand = false + case "m": + m.mouseOn = !m.mouseOn + if m.mouseOn { + return m, tea.EnableMouseCellMotion + } + return m, tea.DisableMouse + } + return m, nil +} + +func (m AgentTUIModel) viewHeight() int { + h := m.height - 2 // reserve for status bar + if h < 4 { + h = 4 + } + return h +} + +func (m *AgentTUIModel) scrollToCursor(topPanels []*RequestPanel) { + targetRow := 0 + if m.cursor < len(topPanels) { + targetRow = topPanels[m.cursor].viewRow + } else { + targetRow = m.appLogRow + } + viewH := m.viewHeight() + if targetRow < m.scrollOffset { + m.scrollOffset = targetRow + } else if targetRow >= m.scrollOffset+viewH { + m.scrollOffset = targetRow - viewH + 3 + } +} + +func (m AgentTUIModel) handleMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) { + switch { + case msg.Button == tea.MouseButtonWheelUp: + m.scrollOffset -= 3 + if m.scrollOffset < 0 { + m.scrollOffset = 0 + } + m.autoFollow = false + return m, nil + case msg.Button == tea.MouseButtonWheelDown: + m.scrollOffset += 3 + m.autoFollow = false + return m, nil + } + + if msg.Button != tea.MouseButtonLeft || msg.Action != tea.MouseActionRelease { + return m, nil + } + y := msg.Y + m.scrollOffset + + // Check app log header + if y == m.appLogRow { + m.appLogExpand = !m.appLogExpand + return m, nil + } + + // Check panel headers + for _, p := range m.panels { + if p.ParentID != "" { + continue + } + if y == p.viewRow { + p.Collapsed = !p.Collapsed + return m, nil + } + } + return m, nil +} + +// ─── Agent Events ───────────────────────────────────────────────────────────── + +func (m *AgentTUIModel) handleAgentEvent(msg AgentEventMsg) tea.Model { + switch msg.Event { + case EventRequestStart: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + // Delegate sub-call: same requestID, different assistantID. + // Add as a tree node inside the existing panel instead of creating a new one. + p := m.panels[idx] + p.Nodes = append(p.Nodes, TreeNode{ + Kind: NodeA2A, + Label: msg.AssistantID, + Status: NodeRunning, + StartTime: time.Now(), + }) + return m + } + + panel := &RequestPanel{ + RequestID: msg.RequestID, + ShortID: shortID(msg.RequestID), + AssistantID: msg.AssistantID, + StartTime: time.Now(), + Status: PanelRunning, + ParentID: msg.ParentID, + } + m.panelIndex[msg.RequestID] = len(m.panels) + m.panels = append(m.panels, panel) + + case EventRequestEnd: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + p := m.panels[idx] + + // Only mark panel done if the ending assistantID matches the panel's original assistantID + // (delegate sub-calls End with a different assistantID, they update their tree node instead) + if msg.AssistantID == p.AssistantID || msg.AssistantID == "" { + if errVal, has := msg.Data["error"]; has && errVal != nil { + p.Status = PanelFailed + } else { + p.Status = PanelSuccess + } + p.EndTime = time.Now() + p.Collapsed = true + + // Finalize any still-running child nodes (e.g. hook interrupted mid-execution) + finalStatus := NodeDone + if p.Status == PanelFailed { + finalStatus = NodeFailed + } + for i := range p.Nodes { + if p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = finalStatus + p.Nodes[i].EndTime = p.EndTime + } + } + } else { + // Delegate sub-call finished: mark its tree node as done + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeA2A && p.Nodes[i].Label == msg.AssistantID && p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + break + } + } + } + } + + case EventLLMCall: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{ + Kind: NodeLLM, Label: "LLM", Status: NodeRunning, StartTime: time.Now(), + }) + } + + case EventLLMDone: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + p := m.panels[idx] + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeLLM && p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + if d, has := msg.Data["detail"]; has { + p.Nodes[i].Detail = fmt.Sprintf("%v", d) + } + break + } + } + } + + case EventToolCall: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + name := dataStr(msg.Data, "name") + m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{ + Kind: NodeTool, Label: name, Status: NodeRunning, StartTime: time.Now(), + }) + } + + case EventToolDone: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + name := dataStr(msg.Data, "name") + p := m.panels[idx] + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeTool && p.Nodes[i].Label == name && p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + break + } + } + } + + case EventHook: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + name := dataStr(msg.Data, "name") + m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{ + Kind: NodeHook, Label: name, Status: NodeRunning, StartTime: time.Now(), + }) + } + + case EventHookDone: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + p := m.panels[idx] + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeHook && p.Nodes[i].Status == NodeRunning { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + break + } + } + } + + case EventA2AStart: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + target := dataStr(msg.Data, "target") + m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{ + Kind: NodeA2A, Label: target, Status: NodeRunning, StartTime: time.Now(), + }) + } + + case EventA2ADone: + if idx, ok := m.panelIndex[msg.RequestID]; ok { + target := dataStr(msg.Data, "target") + p := m.panels[idx] + for i := len(p.Nodes) - 1; i >= 0; i-- { + if p.Nodes[i].Kind == NodeA2A && p.Nodes[i].Status == NodeRunning && (target == "" || p.Nodes[i].Label == target) { + p.Nodes[i].Status = NodeDone + p.Nodes[i].EndTime = time.Now() + break + } + } + } + } + return m +} + +// ─── View ───────────────────────────────────────────────────────────────────── + +func (m AgentTUIModel) View() string { + if m.quitting { + return "" + } + + boxW := m.width - 2 + if boxW < 40 { + boxW = 40 + } + + // Render full content + var sb strings.Builder + row := 0 + topIdx := 0 + + for _, panel := range m.panels { + if panel.ParentID != "" { + continue + } + selected := (topIdx == m.cursor) + rendered := m.renderPanelBox(panel, boxW, selected, &row) + sb.WriteString(rendered) + sb.WriteString("\n") + row++ + topIdx++ + } + + // App Log + m.appLogRow = row + sb.WriteString(m.renderAppLogBox(boxW, topIdx == m.cursor, &row)) + + fullContent := sb.String() + lines := strings.Split(fullContent, "\n") + totalLines := len(lines) + viewH := m.viewHeight() + + // Auto-follow: snap to bottom + if m.autoFollow { + m.scrollOffset = totalLines - viewH + } + + // Clamp scroll offset + maxScroll := totalLines - viewH + if maxScroll < 0 { + maxScroll = 0 + } + if m.scrollOffset > maxScroll { + m.scrollOffset = maxScroll + } + if m.scrollOffset < 0 { + m.scrollOffset = 0 + } + + // Slice visible lines + end := m.scrollOffset + viewH + if end > totalLines { + end = totalLines + } + visible := lines[m.scrollOffset:end] + + // Build output + var out strings.Builder + out.WriteString(strings.Join(visible, "\n")) + + // Status bar with scroll indicator + mouseLabel := "off" + if m.mouseOn { + mouseLabel = "on" + } + scrollInfo := "" + if totalLines > viewH { + pct := 100 + if maxScroll > 0 { + pct = m.scrollOffset * 100 / maxScroll + } + scrollInfo = fmt.Sprintf(" [%d%%]", pct) + } + followLabel := "" + if m.autoFollow { + followLabel = " AUTO" + } + hint := sDim.Render(fmt.Sprintf(" j/k:scroll tab:select space:toggle a/A:all G:bottom g:top m:mouse(%s)%s%s q:quit", + mouseLabel, scrollInfo, followLabel)) + out.WriteString("\n" + hint) + + return out.String() +} + +func (m AgentTUIModel) renderPanelBox(panel *RequestPanel, boxW int, selected bool, row *int) string { + // Record header row for mouse + panel.viewRow = *row + + elapsed := m.panelElapsed(panel) + icon, statusText, style := panelStatusDisplay(panel.Status, elapsed) + + // Title line + collapser := "β–Ύ" + if panel.Collapsed { + collapser = "β–Έ" + } + cursor := " " + if selected { + cursor = "β€Ί" + } + title := fmt.Sprintf("%s %s %s %s %s", + sDim.Render(cursor), + sDim.Render(collapser), + sBold.Render(panel.ShortID), + panel.AssistantID, + style.Render(icon+" "+statusText), + ) + + if panel.Collapsed { + box := boxForStatus(panel.Status).Width(boxW) + result := box.Render(title) + *row += strings.Count(result, "\n") + 1 + return result + } + + // Build body + var body strings.Builder + body.WriteString(title + "\n") + + for _, node := range panel.Nodes { + body.WriteString(m.renderTreeNode(node, " ", false, panel)) + } + + // Render fork children (different requestID, parentID matches) + children := m.childPanels(panel.RequestID) + for i, child := range children { + isLast := (i == len(children)-1) + body.WriteString(m.renderChildSummary(child, " ", isLast)) + } + + box := boxForStatus(panel.Status).Width(boxW) + result := box.Render(body.String()) + *row += strings.Count(result, "\n") + 1 + return result +} + +func (m AgentTUIModel) renderTreeNode(node TreeNode, prefix string, isChild bool, panel *RequestPanel) string { + panelEnded := panel != nil && panel.Status != PanelRunning + displayNode := node + if panelEnded && displayNode.Status == NodeRunning { + displayNode.Status = NodeFailed + } + icon, statusText := nodeStatusDisplay(displayNode) + elapsed := m.nodeElapsed(node, panelEnded, panel.EndTime) + + label := "" + switch node.Kind { + case NodeHook: + label = sMagenta.Render("Hook: "+node.Label) + " " + statusText + case NodeLLM: + detail := "" + if node.Detail != "" { + detail = " " + sDim.Render("["+node.Detail+"]") + } + label = sBlue.Render("LLM") + " " + statusText + detail + case NodeTool: + label = sTree.Render("β”œ ") + sYellow.Render(node.Label) + " " + statusText + case NodeA2A: + label = sTree.Render("β€· ") + sBold.Render(node.Label) + " " + statusText + case NodePhase: + label = node.Label + " " + statusText + default: + label = node.Label + " " + statusText + } + + _ = icon + line := prefix + label + if elapsed != "" { + line += " " + sDim.Render(elapsed) + } + return line + "\n" +} + +func (m AgentTUIModel) renderChildSummary(panel *RequestPanel, prefix string, isLast bool) string { + elapsed := m.panelElapsed(panel) + icon, statusText, style := panelStatusDisplay(panel.Status, elapsed) + + branch := sTree.Render("β”œβ”€ ") + if isLast { + branch = sTree.Render("└─ ") + } + return fmt.Sprintf("%s%s%s %s %s\n", + prefix, branch, + sBold.Render(panel.ShortID+" "+panel.AssistantID), + style.Render(icon+" "+statusText), + sDim.Render(elapsed), + ) +} + +func (m AgentTUIModel) renderAppLogBox(boxW int, selected bool, row *int) string { + cursor := " " + if selected { + cursor = "β€Ί" + } + collapser := "β–Έ" + if m.appLogExpand { + collapser = "β–Ύ" + } + + count := len(m.appLogs) + title := fmt.Sprintf("%s %s %s (%d)", + sDim.Render(cursor), + sDim.Render(collapser), + sBold.Render("App Output"), + count, + ) + + if !m.appLogExpand || count == 0 { + result := boxAppLog.Width(boxW).Render(title) + *row += strings.Count(result, "\n") + 1 + return result + } + + var body strings.Builder + body.WriteString(title + "\n") + + start := 0 + if count > 50 { + start = count - 50 + } + for _, entry := range m.appLogs[start:] { + body.WriteString(" " + entry.Content + "\n") + } + + result := boxAppLog.Width(boxW).Render(body.String()) + *row += strings.Count(result, "\n") + 1 + return result +} + +// ─── Helpers ────────────────────────────────────────────────────────────────── + +func (m AgentTUIModel) topLevelPanels() []*RequestPanel { + var result []*RequestPanel + for _, p := range m.panels { + if p.ParentID == "" { + result = append(result, p) + } + } + return result +} + +func (m AgentTUIModel) childPanels(parentRequestID string) []*RequestPanel { + var result []*RequestPanel + for _, p := range m.panels { + if p.ParentID == parentRequestID { + result = append(result, p) + } + } + return result +} + +func (m AgentTUIModel) panelElapsed(p *RequestPanel) string { + if p.Status != PanelRunning && !p.EndTime.IsZero() { + return fmtDuration(p.EndTime.Sub(p.StartTime)) + } + return fmtDuration(time.Since(p.StartTime)) +} + +func (m AgentTUIModel) nodeElapsed(n TreeNode, panelEnded bool, panelEndTime time.Time) string { + if n.Status == NodeDone || n.Status == NodeFailed { + if !n.EndTime.IsZero() { + return fmtDuration(n.EndTime.Sub(n.StartTime)) + } + } + if n.Status == NodeRunning { + if panelEnded && !panelEndTime.IsZero() { + return fmtDuration(panelEndTime.Sub(n.StartTime)) + } + return fmtDuration(time.Since(n.StartTime)) + } + return "" +} + +func panelStatusDisplay(status PanelStatus, elapsed string) (icon string, text string, style lipgloss.Style) { + switch status { + case PanelRunning: + return "⟳", "running " + elapsed, sRunning + case PanelSuccess: + return "βœ“", "done " + elapsed, sDone + case PanelFailed: + return "βœ—", "failed " + elapsed, sFailed + } + return "", "", sDim +} + +func nodeStatusDisplay(n TreeNode) (icon string, text string) { + switch n.Status { + case NodePending: + return "…", sDim.Render("…") + case NodeRunning: + return "⟳", sRunning.Render("⟳") + case NodeDone: + return "βœ“", sDone.Render("βœ“") + case NodeFailed: + return "βœ—", sFailed.Render("βœ—") + } + return "", "" +} + +func boxForStatus(status PanelStatus) lipgloss.Style { + switch status { + case PanelRunning: + return boxRunning + case PanelFailed: + return boxFailed + default: + return boxDone + } +} + +func dataStr(data map[string]interface{}, key string) string { + if v, ok := data[key]; ok { + return fmt.Sprintf("%v", v) + } + return "" +} + +func fmtDuration(d time.Duration) string { + if d < time.Second { + return fmt.Sprintf("%dms", d.Milliseconds()) + } + return fmt.Sprintf("%.1fs", d.Seconds()) +} diff --git a/agent/context/tui_msg.go b/agent/context/tui_msg.go new file mode 100644 index 00000000..f910df44 --- /dev/null +++ b/agent/context/tui_msg.go @@ -0,0 +1,90 @@ +package context + +import "time" + +// EventType represents the type of agent lifecycle event +type EventType int + +const ( + EventRequestStart EventType = iota + EventPhase + EventPhaseDone + EventPhaseSkip + EventLLMCall + EventLLMDone + EventToolCall + EventToolDone + EventHook + EventHookDone + EventA2AStart + EventA2ADone + EventRequestEnd + EventContextFork + EventContextRelease +) + +// AgentEventMsg is sent from RequestLogger to the TUI Program +type AgentEventMsg struct { + RequestID string + ParentID string + AssistantID string + Event EventType + Data map[string]interface{} +} + +// AppLogLevel represents the severity of application-side output +type AppLogLevel int + +const ( + AppLogLevelLog AppLogLevel = iota + AppLogLevelInfo + AppLogLevelWarn + AppLogLevelError + AppLogLevelException +) + +// AppLogMsg is sent from the DevWriter (gou layer) to the TUI Program +type AppLogMsg struct { + Level AppLogLevel + Content string +} + +// AppLogEntry stores a single application output entry +type AppLogEntry struct { + Level AppLogLevel + Content string + Time time.Time +} + +// PanelStatus represents the lifecycle state of a request panel +type PanelStatus int + +const ( + PanelRunning PanelStatus = iota + PanelSuccess + PanelFailed +) + +// NodeKind represents the type of a tree node within a request panel +type NodeKind int + +const ( + NodePhase NodeKind = iota + NodeLLM + NodeTool + NodeHook + NodeA2A +) + +// NodeStatus represents the state of a tree node +type NodeStatus int + +const ( + NodePending NodeStatus = iota + NodeRunning + NodeDone + NodeFailed +) + +// TickMsg triggers periodic UI refresh for elapsed time display +type TickMsg time.Time diff --git a/cmd/start.go b/cmd/start.go index 9b221474..d663505b 100644 --- a/cmd/start.go +++ b/cmd/start.go @@ -8,11 +8,14 @@ import ( "strings" "syscall" + tea "github.com/charmbracelet/bubbletea" "github.com/fatih/color" + "github.com/mattn/go-isatty" "github.com/spf13/cobra" "github.com/yaoapp/gou/api" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/fs" + "github.com/yaoapp/gou/helper" "github.com/yaoapp/gou/mcp" "github.com/yaoapp/gou/plugin" "github.com/yaoapp/gou/schedule" @@ -20,7 +23,9 @@ import ( "github.com/yaoapp/gou/store" "github.com/yaoapp/gou/task" "github.com/yaoapp/gou/websocket" + "github.com/yaoapp/kun/exception" "github.com/yaoapp/kun/log" + agentcontext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/engine" "github.com/yaoapp/yao/openapi" @@ -33,6 +38,7 @@ import ( var startDebug = false var startDisableWatching = false +var startTUI = false var startCmd = &cobra.Command{ Use: "start", @@ -265,8 +271,9 @@ var startCmd = &cobra.Command{ switch v { case http.READY: - fmt.Println(color.GreenString(L("✨Server is up and running..."))) - fmt.Println(color.GreenString("✨Ctrl+C to stop")) + fmt.Println(color.GreenString(L("Server is up and running..."))) + fmt.Println(color.GreenString("Ctrl+C to stop")) + initAgentTUI() break case http.CLOSED: @@ -627,7 +634,38 @@ func colorMehtod(method string) string { } } +// initAgentTUI initializes the TUI for agent request visualization in dev mode. +// Must be called after HTTP READY to avoid interfering with startup messages. +func initAgentTUI() { + if !config.IsDevelopment() { + return + } + + if !startTUI && os.Getenv("YAO_TUI") != "on" { + return + } + + if !isatty.IsTerminal(os.Stdout.Fd()) { + return + } + + model := agentcontext.NewAgentTUIModel() + p := tea.NewProgram(model, tea.WithoutSignalHandler()) + + agentcontext.SetTUIProgram(p) + tuiWriter := &agentcontext.TUILogWriter{Program: p} + helper.SetDevWriter(tuiWriter) + exception.SetWriter(tuiWriter) + + go func() { + if _, err := p.Run(); err != nil { + log.Error("TUI error: %s", err.Error()) + } + }() +} + func init() { startCmd.PersistentFlags().BoolVarP(&startDebug, "debug", "", false, L("Development mode")) startCmd.PersistentFlags().BoolVarP(&startDisableWatching, "disable-watching", "", false, L("Disable watching")) + startCmd.PersistentFlags().BoolVarP(&startTUI, "tui", "", false, L("Enable TUI for agent request visualization")) } diff --git a/engine/load.go b/engine/load.go index 411d19a5..f96f2090 100644 --- a/engine/load.go +++ b/engine/load.go @@ -1,6 +1,7 @@ package engine import ( + "context" "fmt" "log" "os" @@ -24,6 +25,7 @@ import ( "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/connector" "github.com/yaoapp/yao/data" + "github.com/yaoapp/yao/event" "github.com/yaoapp/yao/flow" "github.com/yaoapp/yao/fs" "github.com/yaoapp/yao/i18n" @@ -48,6 +50,8 @@ import ( "github.com/yaoapp/yao/websocket" "github.com/yaoapp/yao/widget" "github.com/yaoapp/yao/widgets" + + _ "github.com/yaoapp/yao/trace" // register trace handler/listener via init() ) // LoadHooks used to load custom widgets/processes @@ -212,6 +216,14 @@ func Load(cfg config.Config, options LoadOption, progressCallback ...func(string warnings = append(warnings, Warning{Widget: "Store", Error: err}) } + // Start Event Service (handlers registered via init(), e.g. trace) + err = loadStep("Event", func() error { + return event.Start() + }, callback) + if err != nil { + warnings = append(warnings, Warning{Widget: "Event", Error: err}) + } + // Load Uploaders err = loadStep("Uploader", func() error { return attachment.Load(cfg) @@ -425,6 +437,9 @@ func Unload() (err error) { } } + // Stop Event Service (before runtime, so in-flight handlers can still use V8) + event.Stop(context.Background()) + // Stop Runtime err = runtime.Stop() diff --git a/go.mod b/go.mod index 4b25c0cf..55089598 100644 --- a/go.mod +++ b/go.mod @@ -64,11 +64,18 @@ require ( github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 // indirect github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.15 // indirect github.com/aws/smithy-go v1.22.3 // indirect + github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/blang/semver/v4 v4.0.0 // indirect github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect github.com/bytedance/sonic v1.13.2 // indirect github.com/bytedance/sonic/loader v0.2.4 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/charmbracelet/bubbletea v1.3.10 // indirect + github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect + github.com/charmbracelet/lipgloss v1.1.0 // indirect + github.com/charmbracelet/x/ansi v0.10.1 // indirect + github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect + github.com/charmbracelet/x/term v0.2.1 // indirect github.com/cloudwego/base64x v0.1.5 // indirect github.com/containerd/errdefs v1.0.0 // indirect github.com/containerd/errdefs/pkg v0.3.0 // indirect @@ -78,6 +85,7 @@ require ( github.com/dlclark/regexp2 v1.11.5 // indirect github.com/docker/go-units v0.5.0 // indirect github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6 // indirect + github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/gabriel-vasile/mimetype v1.4.9 // indirect github.com/gin-contrib/sse v1.1.0 // indirect @@ -115,9 +123,11 @@ require ( github.com/klauspost/cpuid/v2 v2.2.10 // indirect github.com/leodido/go-urn v1.4.0 // indirect github.com/lib/pq v1.10.9 // indirect + github.com/lucasb-eyer/go-colorful v1.3.0 // indirect github.com/mark3labs/mcp-go v0.32.0 // indirect github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-localereader v0.0.1 // indirect github.com/mattn/go-runewidth v0.0.16 // indirect github.com/mattn/go-sqlite3 v1.14.28 // indirect github.com/miekg/dns v1.1.66 // indirect @@ -126,6 +136,9 @@ require ( github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/montanaflynn/stats v0.7.1 // indirect + github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect + github.com/muesli/cancelreader v0.2.2 // indirect + github.com/muesli/termenv v0.16.0 // indirect github.com/neo4j/neo4j-go-driver/v5 v5.28.1 // indirect github.com/oklog/run v1.1.0 // indirect github.com/opencontainers/go-digest v1.0.0 // indirect @@ -159,6 +172,7 @@ require ( github.com/xdg-go/pbkdf2 v1.0.0 // indirect github.com/xdg-go/scram v1.1.2 // indirect github.com/xdg-go/stringprep v1.0.4 // indirect + github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/xuri/efp v0.0.1 // indirect github.com/xuri/nfp v0.0.1 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/go.sum b/go.sum index 52182b71..0ff37529 100644 --- a/go.sum +++ b/go.sum @@ -39,6 +39,8 @@ github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3 h1:BRXS0U76Z8wfF+bnkilA2QwpIch6U github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3/go.mod h1:bNXKFFyaiVvWuR6O16h/I1724+aXe/tAkA9/QS01t5k= github.com/aws/smithy-go v1.22.3 h1:Z//5NuZCSW6R4PhQ93hShNbyBbn8BWCmCVCt+Q8Io5k= github.com/aws/smithy-go v1.22.3/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI= +github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= +github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/blang/semver v3.5.1+incompatible h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ= github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk= github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM= @@ -62,6 +64,18 @@ github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqy github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw= +github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4= +github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs= +github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk= +github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY= +github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30= +github.com/charmbracelet/x/ansi v0.10.1 h1:rL3Koar5XvX0pHGfovN03f5cxLbCF2YvLeyz7D2jVDQ= +github.com/charmbracelet/x/ansi v0.10.1/go.mod h1:3RQDQ6lDnROptfpWuUVIUG64bD2g2BgntdxH0Ya5TeE= +github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd h1:vy0GVL4jeHEwG5YOXDmi86oYw2yuYUGqz6a8sLwg0X8= +github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs= +github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ= +github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg= github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4= github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= @@ -98,6 +112,8 @@ github.com/emersion/go-sasl v0.0.0-20200509203442-7bfe0ed36a21/go.mod h1:iL2twTe github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6 h1:oP4q0fw+fOSWn3DfFi4EXdT+B+gTtzx8GC9xsc26Znk= github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6/go.mod h1:iL2twTeMvZnrg54ZoPDNfJaJaqy0xIQFuBdrLsmspwQ= github.com/emersion/go-textwrapper v0.0.0-20200911093747-65d896831594/go.mod h1:aqO8z8wPrjkscevZJFVE1wXJrLpC5LtJG7fqLOsPb2U= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= github.com/evanw/esbuild v0.25.4 h1:k1bTSim+usBG27w7BfOCorhgx3tO+6bAfMj5pR+6SKg= github.com/evanw/esbuild v0.25.4/go.mod h1:D2vIQZqV/vIf/VRHtViaUtViZmG7o+kKmlBfVQuRi48= github.com/expr-lang/expr v1.17.7 h1:Q0xY/e/2aCIp8g9s/LGvMDCC5PxYlvHgDZRQ4y16JX8= @@ -231,6 +247,8 @@ github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= +github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag= +github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mark3labs/mcp-go v0.32.0 h1:fgwmbfL2gbd67obg57OfV2Dnrhs1HtSdlY/i5fn7MU8= github.com/mark3labs/mcp-go v0.32.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4= github.com/matoous/go-nanoid/v2 v2.1.0 h1:P64+dmq21hhWdtvZfEAofnvJULaRR1Yib0+PnU669bE= @@ -243,6 +261,8 @@ github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Ky github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4= +github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= @@ -269,6 +289,12 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= github.com/mozillazg/go-pinyin v0.20.0 h1:BtR3DsxpApHfKReaPO1fCqF4pThRwH9uwvXzm+GnMFQ= github.com/mozillazg/go-pinyin v0.20.0/go.mod h1:iR4EnMMRXkfpFVV5FMi4FNB6wGq9NV6uDWbUuPhP4Yc= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= +github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= +github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= +github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= +github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= github.com/neo4j/neo4j-go-driver/v5 v5.28.1 h1:RKWQW7wTgYAY2fU9S+9LaJ9OwRPbRc0I17tlT7nDmAY= github.com/neo4j/neo4j-go-driver/v5 v5.28.1/go.mod h1:Vff8OwT7QpLm7L2yYr85XNWe9Rbqlbeb9asNXJTHO4k= github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE= @@ -379,6 +405,8 @@ github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY= github.com/xdg-go/scram v1.1.2/go.mod h1:RT/sEzTbU5y00aCK8UOx6R7YryM0iF1N2MOmC3kKLN4= github.com/xdg-go/stringprep v1.0.4 h1:XLI/Ng3O1Atzq0oBs3TWm+5ZVgkq2aqdlvP9JtoZ6c8= github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/xuri/efp v0.0.1 h1:fws5Rv3myXyYni8uwj2qKjVaRP30PdjeYe2Y6FDsCL8= github.com/xuri/efp v0.0.1/go.mod h1:ybY/Jr0T0GTCnYjKqmdwxyxn2BQf2RcQIIvex5QldPI= github.com/xuri/excelize/v2 v2.9.1 h1:VdSGk+rraGmgLHGFaGG9/9IWu1nj4ufjJ7uwMDtj8Qw= @@ -471,6 +499,7 @@ golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= diff --git a/test/utils.go b/test/utils.go index 11bf75f5..8b3efb80 100644 --- a/test/utils.go +++ b/test/utils.go @@ -66,12 +66,14 @@ // 7. Starts V8 JavaScript runtime // 8. Registers query engines for database operations // 9. Creates temporary data directories for test isolation +// 10. Starts the Event Service (handlers registered via init(), e.g. trace) // // WHAT test.Clean() DOES: -// 1. Stops V8 runtime and releases resources -// 2. Closes all database connections -// 3. Removes temporary test data stores -// 4. Resets global state to prevent test interference +// 1. Stops the Event Service (drains in-flight events) +// 2. Stops V8 runtime and releases resources +// 3. Closes all database connections +// 4. Removes temporary test data stores +// 5. Resets global state to prevent test interference // // WHAT test.Start() DOES: // 1. Creates Gin HTTP server with API routes @@ -159,6 +161,7 @@ package test import ( + "context" "fmt" "os" "path/filepath" @@ -184,11 +187,14 @@ import ( "github.com/yaoapp/xun/capsule" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/data" + "github.com/yaoapp/yao/event" "github.com/yaoapp/yao/fs" "github.com/yaoapp/yao/helper" "github.com/yaoapp/yao/runtime" "github.com/yaoapp/yao/share" "github.com/yaoapp/yao/utils" + + _ "github.com/yaoapp/yao/trace" // register trace event handler via init() ) var testServer *http.Server = nil @@ -488,10 +494,15 @@ func Prepare(t *testing.T, cfg config.Config, opts ...interface{}) { load(t, cfg) startRuntime(t, cfg) + // Start event service (trace handler registered via blank import above) + if err := event.Start(); err != nil { + t.Fatalf("Failed to start event service: %v", err) + } } // Clean the test environment func Clean() { + event.Stop(context.Background()) dbclose() runtime.Stop() diff --git a/trace/event_listener.go b/trace/event_listener.go new file mode 100644 index 00000000..f699595d --- /dev/null +++ b/trace/event_listener.go @@ -0,0 +1,25 @@ +package trace + +import ( + "context" + + "github.com/yaoapp/yao/event" + eventTypes "github.com/yaoapp/yao/event/types" +) + +// traceUpdateListener receives trace update events for cross-cutting concerns +// (e.g., audit logging, metrics). Trace updates are broadcast via event.Push +// and delivered to this listener and any dynamic subscribers. +type traceUpdateListener struct{} + +func (l *traceUpdateListener) OnEvent(ev *eventTypes.Event) {} + +func (l *traceUpdateListener) Shutdown(ctx context.Context) error { + return nil +} + +func init() { + event.Listen("trace.*", &traceUpdateListener{}, + event.BufferSize(4096), + ) +} diff --git a/trace/handler.go b/trace/handler.go new file mode 100644 index 00000000..e73feafd --- /dev/null +++ b/trace/handler.go @@ -0,0 +1,28 @@ +package trace + +import ( + "context" + + "github.com/yaoapp/yao/event" + eventTypes "github.com/yaoapp/yao/event/types" +) + +// traceHandler processes trace events dispatched through the event service. +// It enables event.Push routing for trace.* events (used by addUpdateAndBroadcast). +type traceHandler struct{} + +func (h *traceHandler) Handle(ctx context.Context, ev *eventTypes.Event, resp chan<- eventTypes.Result) { + resp <- eventTypes.Result{} +} + +func (h *traceHandler) Shutdown(ctx context.Context) error { + return nil +} + +func init() { + event.Register("trace", &traceHandler{}, + event.MaxWorkers(256), + event.ReservedWorkers(32), + event.QueueSize(4096), + ) +} diff --git a/trace/manager.go b/trace/manager.go index 8aaf6b25..2b4989e4 100644 --- a/trace/manager.go +++ b/trace/manager.go @@ -3,63 +3,56 @@ package trace import ( "context" "fmt" + "sync" "time" gonanoid "github.com/matoous/go-nanoid/v2" "github.com/yaoapp/kun/log" - "github.com/yaoapp/yao/trace/pubsub" + "github.com/yaoapp/yao/event" "github.com/yaoapp/yao/trace/types" ) -// manager implements the Manager interface with channel-based state management +// manager implements the Manager interface. +// State is protected by a mutex, replacing the old channel-based state worker. +// This eliminates the context-cancel bug while maintaining thread safety. type manager struct { - ctx context.Context - cancel context.CancelFunc - traceID string - driver types.Driver - stateCmdChan chan stateCommand // Single channel for all state mutations - closed int32 // Atomic flag: 1 = closed, safeSend rejects new commands - autoArchive bool // Auto-archive on complete/fail - pubsub *pubsub.PubSub // Reference to independent pubsub service (for publishing only, doesn't own it) + mu sync.Mutex + traceID string + driver types.Driver + state *managerState + autoArchive bool } -// NewManager creates a new trace manager instance -// 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) - - // Determine auto-archive setting +// NewManager creates a new trace manager instance. +func NewManager(ctx context.Context, traceID string, driver types.Driver, option *types.TraceOption) (types.Manager, error) { autoArchive := false if option != nil { autoArchive = option.AutoArchive } m := &manager{ - ctx: managerCtx, - cancel: cancel, - traceID: traceID, - driver: driver, - stateCmdChan: make(chan stateCommand, 100), // Buffered channel for performance - autoArchive: autoArchive, - pubsub: pubsubService, // Reference only, doesn't manage lifecycle + traceID: traceID, + driver: driver, + autoArchive: autoArchive, + state: &managerState{ + spaces: make(map[string]*types.TraceSpace), + traceStatus: types.TraceStatusPending, + updates: make([]*types.TraceUpdate, 0, 100), + }, } - // Start state worker goroutine - go m.startStateWorker() - - // Try to load existing updates from driver (for resumed traces) + // Load existing updates from driver (for resumed traces). + // Safe to access m.state directly here β€” no Queue yet, single goroutine. 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 + m.state.updates = existingUpdates for _, update := range existingUpdates { if update.Type == types.UpdateTypeComplete { log.Trace("[MANAGER] NewManager: trace %s was already completed, marking as completed", traceID) - m.stateMarkCompleted() + m.state.completed = true if data, ok := update.Data.(*types.TraceCompleteData); ok { log.Trace("[MANAGER] NewManager: setting trace status to %s", data.Status) - m.stateSetTraceStatus(data.Status) + m.state.traceStatus = data.Status } break } @@ -70,7 +63,6 @@ func NewManager(ctx context.Context, traceID string, driver types.Driver, pubsub } 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{ Type: types.UpdateTypeInit, @@ -89,36 +81,24 @@ func genNodeID() string { return id } -// addUpdateAndBroadcast persists, adds to history, and publishes an update +// addUpdateAndBroadcast persists, adds to history, and broadcasts via event service. func (m *manager) addUpdateAndBroadcast(update *types.TraceUpdate) { - // Persist to driver (synchronous - no race) 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) - // Publish to independent PubSub service (manager just publishes, doesn't manage pubsub lifecycle) - if m.pubsub != nil { - m.pubsub.Publish(update) - } + // Broadcast to subscribers via event service (fire-and-forget, non-blocking). + // Uses Push with the update as payload so event.Subscribe filters can match by traceID. + event.Push(context.Background(), "trace.update", update) } -// checkContext checks if context is cancelled +// checkContext checks if the trace has been completed/released. +// With event-based state management, the manager no longer binds a context. +// Lifecycle is controlled by QueueCreate/QueueRelease. func (m *manager) checkContext() error { - select { - case <-m.ctx.Done(): - // Context cancelled - just return the error - // Don't call handleCancellation here to avoid deadlock - // handleCancellation should be called explicitly when needed - return m.ctx.Err() - default: - return nil - } + return nil } // Add creates next sequential node - auto-joins if currently in parallel state @@ -147,7 +127,7 @@ func (m *manager) Add(input types.TraceInput, option types.TraceNodeOption) (typ } // Save root node - if err := m.driver.SaveNode(m.ctx, m.traceID, rootNode); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, rootNode); err != nil { return nil, fmt.Errorf("failed to save root node: %w", err) } @@ -210,14 +190,14 @@ func (m *manager) Add(input types.TraceInput, option types.TraceNodeOption) (typ // Add to each parent's children for _, parent := range currentNodes { parent.Children = append(parent.Children, newNodeData) - if err := m.driver.SaveNode(m.ctx, m.traceID, parent); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, parent); err != nil { // Log error but continue m.Error("Failed to update parent node %s: %v", parent.ID, err) } } // Save new node - if err := m.driver.SaveNode(m.ctx, m.traceID, newNodeData); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, newNodeData); err != nil { return nil, err } @@ -310,7 +290,7 @@ func (m *manager) Parallel(parallelInputs []types.TraceParallelInput) ([]types.N // Save all nodes in batch - collect errors var saveErrors []error for _, data := range nodeData { - if err := m.driver.SaveNode(m.ctx, m.traceID, data); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, data); err != nil { saveErrors = append(saveErrors, fmt.Errorf("failed to save node %s: %w", data.ID, err)) } } @@ -321,7 +301,7 @@ func (m *manager) Parallel(parallelInputs []types.TraceParallelInput) ([]types.N } // Save parent node - if err := m.driver.SaveNode(m.ctx, m.traceID, parentNode); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, parentNode); err != nil { return nil, fmt.Errorf("failed to save parent node: %w", err) } @@ -380,7 +360,7 @@ func (m *manager) log(level string, message string, args ...any) { NodeID: node.ID, } // Save log (ignore errors for non-critical logging) - _ = m.driver.SaveLog(m.ctx, m.traceID, log) + _ = m.driver.SaveLog(context.Background(), m.traceID, log) // Broadcast log event m.addUpdateAndBroadcast(&types.TraceUpdate{ @@ -404,7 +384,7 @@ func (m *manager) SetOutput(output types.TraceOutput) error { for _, node := range nodes { node.Output = output node.UpdatedAt = now - if err := m.driver.SaveNode(m.ctx, m.traceID, node); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, node); err != nil { return err } @@ -434,7 +414,7 @@ func (m *manager) SetMetadata(key string, value any) error { } node.Metadata[key] = value node.UpdatedAt = now - if err := m.driver.SaveNode(m.ctx, m.traceID, node); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, node); err != nil { return err } @@ -485,7 +465,7 @@ func (m *manager) Complete(output ...types.TraceOutput) error { node.Status = types.StatusCompleted node.EndTime = now node.UpdatedAt = now - if err := m.driver.SaveNode(m.ctx, m.traceID, node); err != nil { + if err := m.driver.SaveNode(context.Background(), m.traceID, node); err != nil { return err } @@ -516,7 +496,7 @@ func (m *manager) Fail(err error) error { node.Status = types.StatusFailed node.EndTime = now node.UpdatedAt = now - if saveErr := m.driver.SaveNode(m.ctx, m.traceID, node); saveErr != nil { + if saveErr := m.driver.SaveNode(context.Background(), m.traceID, node); saveErr != nil { return saveErr } @@ -545,7 +525,7 @@ func (m *manager) GetRootNode() (*types.TraceNode, error) { // GetNode returns a node by ID func (m *manager) GetNode(id string) (*types.TraceNode, error) { - return m.driver.LoadNode(m.ctx, m.traceID, id) + return m.driver.LoadNode(context.Background(), m.traceID, id) } // GetCurrentNodes returns current active nodes @@ -581,7 +561,7 @@ func (m *manager) MarkComplete() error { // Auto-archive if enabled if m.autoArchive { - if err := m.driver.Archive(m.ctx, m.traceID); err != nil { + if err := m.driver.Archive(context.Background(), m.traceID); err != nil { // Log error but don't fail the complete operation m.Debug("Failed to auto-archive trace", map[string]any{ "trace_id": m.traceID, @@ -610,7 +590,7 @@ func (m *manager) CreateSpace(option types.TraceSpaceOption) (*types.TraceSpace, } // Save to driver - if err := m.driver.SaveSpace(m.ctx, m.traceID, space); err != nil { + if err := m.driver.SaveSpace(context.Background(), m.traceID, space); err != nil { return nil, err } @@ -637,7 +617,7 @@ func (m *manager) GetSpace(id string) (*types.TraceSpace, error) { } // Load from driver - space, err := m.driver.LoadSpace(m.ctx, m.traceID, id) + space, err := m.driver.LoadSpace(context.Background(), m.traceID, id) if err != nil { return nil, err } @@ -658,7 +638,7 @@ func (m *manager) HasSpace(id string) bool { } // Check in driver - space, _ := m.driver.LoadSpace(m.ctx, m.traceID, id) + space, _ := m.driver.LoadSpace(context.Background(), m.traceID, id) return space != nil } @@ -674,7 +654,7 @@ func (m *manager) DeleteSpace(id string) error { m.stateDeleteSpace(id) // Delete from driver - if err := m.driver.DeleteSpace(m.ctx, m.traceID, id); err != nil { + if err := m.driver.DeleteSpace(context.Background(), m.traceID, id); err != nil { return err } @@ -693,7 +673,7 @@ func (m *manager) DeleteSpace(id string) error { // ListSpaces returns all spaces func (m *manager) ListSpaces() []*types.TraceSpace { // Load from driver to ensure we have all spaces - spaceIDs, err := m.driver.ListSpaces(m.ctx, m.traceID) + spaceIDs, err := m.driver.ListSpaces(context.Background(), m.traceID) if err != nil { // Fallback to cached spaces return m.stateGetAllSpaces() @@ -727,13 +707,13 @@ func (m *manager) SetSpaceValue(spaceID, key string, value any) error { // Set value in driver (through state worker for concurrent safety) err = m.stateExecuteSpaceOp(spaceID, func() error { - if err := m.driver.SetSpaceKey(m.ctx, m.traceID, spaceID, key, value); err != nil { + if err := m.driver.SetSpaceKey(context.Background(), m.traceID, spaceID, key, value); err != nil { return err } // Update space timestamp space.UpdatedAt = now - if err := m.driver.SaveSpace(m.ctx, m.traceID, space); err != nil { + if err := m.driver.SaveSpace(context.Background(), m.traceID, space); err != nil { return err } @@ -761,7 +741,7 @@ func (m *manager) GetSpaceValue(spaceID, key string) (any, error) { var result any err := m.stateExecuteSpaceOp(spaceID, func() error { var err error - result, err = m.driver.GetSpaceKey(m.ctx, m.traceID, spaceID, key) + result, err = m.driver.GetSpaceKey(context.Background(), m.traceID, spaceID, key) return err }) return result, err @@ -771,7 +751,7 @@ func (m *manager) GetSpaceValue(spaceID, key string) (any, error) { func (m *manager) HasSpaceValue(spaceID, key string) bool { var result bool _ = m.stateExecuteSpaceOp(spaceID, func() error { - result = m.driver.HasSpaceKey(m.ctx, m.traceID, spaceID, key) + result = m.driver.HasSpaceKey(context.Background(), m.traceID, spaceID, key) return nil }) return result @@ -787,7 +767,7 @@ func (m *manager) DeleteSpaceValue(spaceID, key string) error { // Delete value from driver (through state worker for concurrent safety) err := m.stateExecuteSpaceOp(spaceID, func() error { - return m.driver.DeleteSpaceKey(m.ctx, m.traceID, spaceID, key) + return m.driver.DeleteSpaceKey(context.Background(), m.traceID, spaceID, key) }) if err != nil { @@ -816,7 +796,7 @@ func (m *manager) ClearSpaceValues(spaceID string) error { // Clear values from driver (through state worker for concurrent safety) err := m.stateExecuteSpaceOp(spaceID, func() error { - return m.driver.ClearSpaceKeys(m.ctx, m.traceID, spaceID) + return m.driver.ClearSpaceKeys(context.Background(), m.traceID, spaceID) }) if err != nil { @@ -840,7 +820,7 @@ func (m *manager) ListSpaceKeys(spaceID string) []string { var keys []string _ = m.stateExecuteSpaceOp(spaceID, func() error { var err error - keys, err = m.driver.ListSpaceKeys(m.ctx, m.traceID, spaceID) + keys, err = m.driver.ListSpaceKeys(context.Background(), m.traceID, spaceID) return err }) return keys @@ -865,7 +845,7 @@ func (m *manager) GetTraceInfo() (*types.TraceInfo, error) { if err := m.checkContext(); err != nil { return nil, err } - return m.driver.LoadTraceInfo(m.ctx, m.traceID) + return m.driver.LoadTraceInfo(context.Background(), m.traceID) } // GetAllNodes retrieves all nodes from storage @@ -875,7 +855,7 @@ func (m *manager) GetAllNodes() ([]*types.TraceNode, error) { } // Load the root node tree from storage - rootNode, err := m.driver.LoadTrace(m.ctx, m.traceID) + rootNode, err := m.driver.LoadTrace(context.Background(), m.traceID) if err != nil { return nil, err } @@ -906,7 +886,7 @@ func (m *manager) GetNodeByID(nodeID string) (*types.TraceNode, error) { if err := m.checkContext(); err != nil { return nil, err } - return m.driver.LoadNode(m.ctx, m.traceID, nodeID) + return m.driver.LoadNode(context.Background(), m.traceID, nodeID) } // GetAllLogs retrieves all logs from storage @@ -914,7 +894,7 @@ func (m *manager) GetAllLogs() ([]*types.TraceLog, error) { if err := m.checkContext(); err != nil { return nil, err } - return m.driver.LoadLogs(m.ctx, m.traceID, "") + return m.driver.LoadLogs(context.Background(), m.traceID, "") } // GetLogsByNode retrieves logs for a specific node from storage @@ -922,7 +902,7 @@ func (m *manager) GetLogsByNode(nodeID string) ([]*types.TraceLog, error) { if err := m.checkContext(); err != nil { return nil, err } - return m.driver.LoadLogs(m.ctx, m.traceID, nodeID) + return m.driver.LoadLogs(context.Background(), m.traceID, nodeID) } // GetAllSpaces retrieves all spaces from storage @@ -932,7 +912,7 @@ func (m *manager) GetAllSpaces() ([]*types.TraceSpace, error) { } // Get all space IDs from driver - spaceIDs, err := m.driver.ListSpaces(m.ctx, m.traceID) + spaceIDs, err := m.driver.ListSpaces(context.Background(), m.traceID) if err != nil { return nil, err } @@ -940,7 +920,7 @@ func (m *manager) GetAllSpaces() ([]*types.TraceSpace, error) { // Load all spaces spaces := make([]*types.TraceSpace, 0, len(spaceIDs)) for _, spaceID := range spaceIDs { - space, err := m.driver.LoadSpace(m.ctx, m.traceID, spaceID) + space, err := m.driver.LoadSpace(context.Background(), m.traceID, spaceID) if err != nil { continue // Skip spaces that fail to load } @@ -959,7 +939,7 @@ func (m *manager) GetSpaceByID(spaceID string) (*types.TraceSpaceData, error) { } // Load space metadata - space, err := m.driver.LoadSpace(m.ctx, m.traceID, spaceID) + space, err := m.driver.LoadSpace(context.Background(), m.traceID, spaceID) if err != nil { return nil, err } @@ -968,7 +948,7 @@ func (m *manager) GetSpaceByID(spaceID string) (*types.TraceSpaceData, error) { } // Load all keys in the space - keys, err := m.driver.ListSpaceKeys(m.ctx, m.traceID, spaceID) + keys, err := m.driver.ListSpaceKeys(context.Background(), m.traceID, spaceID) if err != nil { return nil, err } @@ -976,7 +956,7 @@ func (m *manager) GetSpaceByID(spaceID string) (*types.TraceSpaceData, error) { // Load all key-value pairs data := make(map[string]any) for _, key := range keys { - value, err := m.driver.GetSpaceKey(m.ctx, m.traceID, spaceID, key) + value, err := m.driver.GetSpaceKey(context.Background(), m.traceID, spaceID, key) if err != nil { continue // Skip keys that fail to load } diff --git a/trace/node.go b/trace/node.go index cd09f14e..aca80e6c 100644 --- a/trace/node.go +++ b/trace/node.go @@ -1,6 +1,7 @@ package trace import ( + "context" "time" "github.com/yaoapp/yao/trace/types" @@ -60,7 +61,7 @@ func (n *node) log(level string, message string, args ...any) *types.TraceLog { NodeID: n.data.ID, } // Save log (ignore errors for non-critical logging) - _ = n.manager.driver.SaveLog(n.manager.ctx, n.manager.traceID, log) + _ = n.manager.driver.SaveLog(context.Background(), n.manager.traceID, log) return log } @@ -85,10 +86,10 @@ func (n *node) Add(input types.TraceInput, option types.TraceNodeOption) (types. n.data.Children = append(n.data.Children, childNodeData) // Save both nodes - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, childNodeData); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, childNodeData); err != nil { return nil, err } - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data); err != nil { return nil, err } @@ -120,7 +121,7 @@ func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node n.data.Children = append(n.data.Children, childNodeData) // Save node - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, childNodeData); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, childNodeData); err != nil { return nil, err } @@ -132,7 +133,7 @@ func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node } // Save parent node - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data); err != nil { return nil, err } @@ -165,7 +166,7 @@ func (n *node) Join(nodes []*types.TraceNode, input types.TraceInput, option typ } // Save join node - if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, joinNodeData); err != nil { + if err := n.manager.driver.SaveNode(context.Background(), n.manager.traceID, joinNodeData); err != nil { return nil, err } @@ -185,7 +186,7 @@ func (n *node) ID() string { func (n *node) SetOutput(output types.TraceOutput) error { n.data.Output = output n.data.UpdatedAt = time.Now().UnixMilli() - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } // SetMetadata sets node metadata @@ -195,14 +196,14 @@ func (n *node) SetMetadata(key string, value any) error { } n.data.Metadata[key] = value n.data.UpdatedAt = time.Now().UnixMilli() - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } // SetStatus sets the node status func (n *node) SetStatus(status string) error { n.data.Status = types.NodeStatus(status) n.data.UpdatedAt = time.Now().UnixMilli() - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } // Complete marks the node as completed (public method, broadcasts event) @@ -236,7 +237,7 @@ func (n *node) complete(output ...types.TraceOutput) error { n.data.Status = types.StatusCompleted n.data.EndTime = now n.data.UpdatedAt = now - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } // Fail marks the node as failed (public method, broadcasts event) @@ -269,5 +270,5 @@ func (n *node) fail(err error) error { n.data.EndTime = now n.data.UpdatedAt = now - return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data) + return n.manager.driver.SaveNode(context.Background(), n.manager.traceID, n.data) } diff --git a/trace/pubsub/pubsub.go b/trace/pubsub/pubsub.go deleted file mode 100644 index 9f8ee309..00000000 --- a/trace/pubsub/pubsub.go +++ /dev/null @@ -1,142 +0,0 @@ -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 deleted file mode 100644 index 26ada520..00000000 --- a/trace/pubsub/subscriber.go +++ /dev/null @@ -1,62 +0,0 @@ -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/space.go b/trace/space.go index 2806e7e4..1b3131cc 100644 --- a/trace/space.go +++ b/trace/space.go @@ -6,18 +6,18 @@ import ( "github.com/yaoapp/yao/trace/types" ) -// space implements the Space interface for custom space operations +// space implements the Space interface for custom space operations. +// Uses context.Background() for driver calls to decouple from caller context +// (fixes the context-fork bug where parent cancellation breaks child ops). type space struct { - ctx context.Context traceID string data *types.TraceSpace driver types.Driver } // NewSpace creates a new space instance -func NewSpace(ctx context.Context, traceID string, data *types.TraceSpace, driver types.Driver) types.Space { +func NewSpace(traceID string, data *types.TraceSpace, driver types.Driver) types.Space { return &space{ - ctx: ctx, traceID: traceID, data: data, driver: driver, @@ -31,32 +31,32 @@ func (s *space) ID() string { // Set stores a value by key func (s *space) Set(key string, value any) error { - return s.driver.SetSpaceKey(s.ctx, s.traceID, s.data.ID, key, value) + return s.driver.SetSpaceKey(context.Background(), s.traceID, s.data.ID, key, value) } // Get retrieves a value by key func (s *space) Get(key string) (any, error) { - return s.driver.GetSpaceKey(s.ctx, s.traceID, s.data.ID, key) + return s.driver.GetSpaceKey(context.Background(), s.traceID, s.data.ID, key) } // Has checks if a key exists func (s *space) Has(key string) bool { - return s.driver.HasSpaceKey(s.ctx, s.traceID, s.data.ID, key) + return s.driver.HasSpaceKey(context.Background(), s.traceID, s.data.ID, key) } // Delete removes a key-value pair func (s *space) Delete(key string) error { - return s.driver.DeleteSpaceKey(s.ctx, s.traceID, s.data.ID, key) + return s.driver.DeleteSpaceKey(context.Background(), s.traceID, s.data.ID, key) } // Clear removes all key-value pairs func (s *space) Clear() error { - return s.driver.ClearSpaceKeys(s.ctx, s.traceID, s.data.ID) + return s.driver.ClearSpaceKeys(context.Background(), s.traceID, s.data.ID) } // Keys returns all keys in the space func (s *space) Keys() []string { - keys, err := s.driver.ListSpaceKeys(s.ctx, s.traceID, s.data.ID) + keys, err := s.driver.ListSpaceKeys(context.Background(), s.traceID, s.data.ID) if err != nil { return nil } diff --git a/trace/state.go b/trace/state.go index 63e4915e..e49404e2 100644 --- a/trace/state.go +++ b/trace/state.go @@ -2,16 +2,13 @@ package trace import ( "fmt" - "sync/atomic" "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/trace/types" ) -// State management using channel-based serialization (no locks needed) -// All state mutations go through a single worker goroutine - -// managerState holds all mutable state (accessed only by state worker) +// managerState holds all mutable state for a trace. +// Protected by manager.mu β€” all access goes through state* methods which acquire the lock. type managerState struct { rootNode *types.TraceNode currentNodes []*types.TraceNode @@ -19,344 +16,133 @@ type managerState struct { traceStatus types.TraceStatus completed bool updates []*types.TraceUpdate - // Note: subscribers moved to SubscriptionManager (no longer in state) -} - -// State command interface - all commands are processed serially -type stateCommand interface { - execute(s *managerState) -} - -// Commands with response channels for synchronous operations - -// --- Root Node Commands --- - -type cmdSetRoot struct { - node *types.TraceNode -} - -func (c *cmdSetRoot) execute(s *managerState) { - s.rootNode = c.node -} - -type cmdGetRoot struct { - resp chan *types.TraceNode -} - -func (c *cmdGetRoot) execute(s *managerState) { - c.resp <- s.rootNode -} - -// --- Current Nodes Commands --- - -type cmdSetCurrentNodes struct { - nodes []*types.TraceNode -} - -func (c *cmdSetCurrentNodes) execute(s *managerState) { - s.currentNodes = c.nodes -} - -type cmdGetCurrentNodes struct { - resp chan []*types.TraceNode -} - -func (c *cmdGetCurrentNodes) execute(s *managerState) { - // Return a copy to prevent external mutation - nodes := make([]*types.TraceNode, len(s.currentNodes)) - copy(nodes, s.currentNodes) - c.resp <- nodes -} - -type cmdUpdateRootAndCurrent struct { - root *types.TraceNode - current []*types.TraceNode -} - -func (c *cmdUpdateRootAndCurrent) execute(s *managerState) { - s.rootNode = c.root - s.currentNodes = c.current -} - -// --- Space Commands --- - -type cmdGetSpace struct { - id string - resp chan *types.TraceSpace -} - -func (c *cmdGetSpace) execute(s *managerState) { - c.resp <- s.spaces[c.id] -} - -type cmdSetSpace struct { - id string - space *types.TraceSpace -} - -func (c *cmdSetSpace) execute(s *managerState) { - s.spaces[c.id] = c.space -} - -type cmdDeleteSpace struct { - id string -} - -func (c *cmdDeleteSpace) execute(s *managerState) { - delete(s.spaces, c.id) -} - -type cmdGetAllSpaces struct { - resp chan []*types.TraceSpace -} - -func (c *cmdGetAllSpaces) execute(s *managerState) { - spaces := make([]*types.TraceSpace, 0, len(s.spaces)) - for _, space := range s.spaces { - spaces = append(spaces, space) - } - c.resp <- spaces -} - -// --- Trace Status Commands --- - -type cmdSetTraceStatus struct { - status types.TraceStatus -} - -func (c *cmdSetTraceStatus) execute(s *managerState) { - s.traceStatus = c.status -} - -type cmdGetTraceStatus struct { - resp chan types.TraceStatus -} - -func (c *cmdGetTraceStatus) execute(s *managerState) { - c.resp <- s.traceStatus -} - -// --- Completion Commands --- - -type cmdMarkCompleted struct { - resp chan bool // Returns true if marked, false if already completed -} - -func (c *cmdMarkCompleted) execute(s *managerState) { - if s.completed { - c.resp <- false - } else { - s.completed = true - c.resp <- true - } -} - -type cmdIsCompleted struct { - resp chan bool -} - -func (c *cmdIsCompleted) execute(s *managerState) { - c.resp <- s.completed -} - -// --- Update Commands --- - -type cmdAddUpdate struct { - update *types.TraceUpdate -} - -func (c *cmdAddUpdate) execute(s *managerState) { - s.updates = append(s.updates, c.update) -} - -type cmdGetUpdates struct { - since int64 - resp chan []*types.TraceUpdate -} - -func (c *cmdGetUpdates) execute(s *managerState) { - filtered := make([]*types.TraceUpdate, 0) - for _, update := range s.updates { - if update.Timestamp >= c.since { - filtered = append(filtered, update) - } - } - c.resp <- filtered -} - -type cmdSetUpdates struct { - updates []*types.TraceUpdate -} - -func (c *cmdSetUpdates) execute(s *managerState) { - s.updates = c.updates -} - -// --- 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 - -type cmdSpaceKVOp struct { - spaceID string - fn func() error - resp chan error -} - -func (c *cmdSpaceKVOp) execute(s *managerState) { - // Execute the operation (typically a driver call) - // The function is provided by caller and executed serially here - err := c.fn() - c.resp <- err -} - -// State worker - processes all commands serially in a single goroutine. -// Exits only when stateCmdChan is closed by Release(). The for-range loop -// automatically drains any buffered commands before returning (Go spec guarantee). -func (m *manager) startStateWorker() { - state := &managerState{ - rootNode: nil, - currentNodes: []*types.TraceNode{}, - spaces: make(map[string]*types.TraceSpace), - traceStatus: types.TraceStatusPending, - completed: false, - updates: make([]*types.TraceUpdate, 0, 100), - } - for cmd := range m.stateCmdChan { - cmd.execute(state) - } -} - -// Helper methods for manager to send commands - -// safeSend sends a command to the state worker channel. Returns false if the -// manager is closed, context is cancelled, or the channel was closed mid-send. -// The atomic closed flag provides a fast-path rejection before touching the channel, -// which is critical in CGO callback stacks where recover() may not work. -func (m *manager) safeSend(cmd stateCommand) (ok bool) { - if atomic.LoadInt32(&m.closed) == 1 { - return false - } - defer func() { - if r := recover(); r != nil { - ok = false - } - }() - select { - case <-m.ctx.Done(): - return false - case m.stateCmdChan <- cmd: - return true - } } func (m *manager) stateSetRoot(node *types.TraceNode) { - m.safeSend(&cmdSetRoot{node: node}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.rootNode = node } func (m *manager) stateGetRoot() *types.TraceNode { - resp := make(chan *types.TraceNode, 1) - if !m.safeSend(&cmdGetRoot{resp: resp}) { - return nil // Context cancelled - } - return <-resp + m.mu.Lock() + defer m.mu.Unlock() + return m.state.rootNode } func (m *manager) stateSetCurrentNodes(nodes []*types.TraceNode) { - m.safeSend(&cmdSetCurrentNodes{nodes: nodes}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.currentNodes = nodes } func (m *manager) stateGetCurrentNodes() []*types.TraceNode { - resp := make(chan []*types.TraceNode, 1) - if !m.safeSend(&cmdGetCurrentNodes{resp: resp}) { - return nil // Context cancelled + m.mu.Lock() + defer m.mu.Unlock() + if m.state.currentNodes == nil { + return nil } - return <-resp + nodes := make([]*types.TraceNode, len(m.state.currentNodes)) + copy(nodes, m.state.currentNodes) + return nodes } func (m *manager) stateUpdateRootAndCurrent(root *types.TraceNode, current []*types.TraceNode) { - m.safeSend(&cmdUpdateRootAndCurrent{root: root, current: current}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.rootNode = root + m.state.currentNodes = current } func (m *manager) stateGetSpace(id string) (*types.TraceSpace, bool) { - resp := make(chan *types.TraceSpace, 1) - if !m.safeSend(&cmdGetSpace{id: id, resp: resp}) { - return nil, false // Context cancelled - } - space := <-resp - return space, space != nil + m.mu.Lock() + defer m.mu.Unlock() + space, ok := m.state.spaces[id] + return space, ok } func (m *manager) stateSetSpace(id string, space *types.TraceSpace) { - m.safeSend(&cmdSetSpace{id: id, space: space}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.spaces[id] = space } func (m *manager) stateDeleteSpace(id string) { - m.safeSend(&cmdDeleteSpace{id: id}) + m.mu.Lock() + defer m.mu.Unlock() + delete(m.state.spaces, id) } func (m *manager) stateGetAllSpaces() []*types.TraceSpace { - resp := make(chan []*types.TraceSpace, 1) - if !m.safeSend(&cmdGetAllSpaces{resp: resp}) { - return nil // Context cancelled + m.mu.Lock() + defer m.mu.Unlock() + spaces := make([]*types.TraceSpace, 0, len(m.state.spaces)) + for _, space := range m.state.spaces { + spaces = append(spaces, space) } - return <-resp + return spaces } func (m *manager) stateSetTraceStatus(status types.TraceStatus) { - m.safeSend(&cmdSetTraceStatus{status: status}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.traceStatus = status } func (m *manager) stateGetTraceStatus() types.TraceStatus { - resp := make(chan types.TraceStatus, 1) - if !m.safeSend(&cmdGetTraceStatus{resp: resp}) { - return types.TraceStatusCancelled // Context cancelled - } - return <-resp + m.mu.Lock() + defer m.mu.Unlock() + return m.state.traceStatus } func (m *manager) stateMarkCompleted() bool { - resp := make(chan bool, 1) - if !m.safeSend(&cmdMarkCompleted{resp: resp}) { - return true // Context cancelled, treat as completed + m.mu.Lock() + defer m.mu.Unlock() + if m.state.completed { + return false } - return <-resp + m.state.completed = true + return true } func (m *manager) stateIsCompleted() bool { - resp := make(chan bool, 1) - if !m.safeSend(&cmdIsCompleted{resp: resp}) { - return true // Context cancelled, treat as completed - } - return <-resp + m.mu.Lock() + defer m.mu.Unlock() + return m.state.completed } func (m *manager) stateAddUpdate(update *types.TraceUpdate) { - m.safeSend(&cmdAddUpdate{update: update}) + m.mu.Lock() + defer m.mu.Unlock() + m.state.updates = append(m.state.updates, update) } func (m *manager) stateGetUpdates(since int64) []*types.TraceUpdate { - resp := make(chan []*types.TraceUpdate, 1) - if !m.safeSend(&cmdGetUpdates{since: since, resp: resp}) { - return nil // Context cancelled + m.mu.Lock() + defer m.mu.Unlock() + filtered := make([]*types.TraceUpdate, 0) + for _, update := range m.state.updates { + if update.Timestamp >= since { + filtered = append(filtered, update) + } } - return <-resp + return filtered } func (m *manager) stateSetUpdates(updates []*types.TraceUpdate) { + m.mu.Lock() + defer m.mu.Unlock() log.Trace("[STATE] stateSetUpdates: setting %d updates for trace %s", len(updates), m.traceID) - m.safeSend(&cmdSetUpdates{updates: updates}) + m.state.updates = updates } -// 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 +// stateExecuteSpaceOp executes a space operation while holding the lock. func (m *manager) stateExecuteSpaceOp(spaceID string, fn func() error) error { - resp := make(chan error, 1) - if !m.safeSend(&cmdSpaceKVOp{spaceID: spaceID, fn: fn, resp: resp}) { - return fmt.Errorf("trace %s: state worker stopped", m.traceID) + m.mu.Lock() + defer m.mu.Unlock() + err := fn() + if err != nil { + return fmt.Errorf("trace %s: space op failed: %w", m.traceID, err) } - return <-resp + return nil } diff --git a/trace/subscription.go b/trace/subscription.go index 7537e36b..ac76cf12 100644 --- a/trace/subscription.go +++ b/trace/subscription.go @@ -3,12 +3,18 @@ package trace import ( "fmt" + "github.com/yaoapp/yao/event" + eventTypes "github.com/yaoapp/yao/event/types" "github.com/yaoapp/yao/trace/types" ) +func dedupKey(u *types.TraceUpdate) string { + return fmt.Sprintf("%s:%s:%d", u.Type, u.NodeID, u.Timestamp) +} + // Subscribe creates a new subscription for trace updates (replays all historical events from the beginning) func (m *manager) Subscribe() (<-chan *types.TraceUpdate, error) { - return m.subscribe(0) // Subscribe from beginning to get all historical events + return m.subscribe(0) } // SubscribeFrom creates a subscription starting from a specific timestamp @@ -16,24 +22,60 @@ func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error) return m.subscribe(since) } -// subscribe is the internal implementation for subscriptions +// subscribe creates a subscription channel that first replays historical +// updates, then streams live events via the event service's Subscriber. +// The subscriber is registered BEFORE reading historical state to prevent +// missing events that occur between the state snapshot and subscriber setup. func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) { - // Get historical updates - updates := m.stateGetUpdates(since) - - // Use manager's pubsub reference (always available) - if m.pubsub == nil { - return nil, fmt.Errorf("pubsub service not initialized for trace: %s", m.traceID) - } - - // 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) - return sub.Channel, nil + out := make(chan *types.TraceUpdate, bufferSize) + + // Register live subscriber FIRST to avoid missing events between snapshot and subscribe. + liveCh := make(chan *eventTypes.Event, bufferSize) + traceID := m.traceID + subID := event.Subscribe("trace.*", liveCh, event.Filter(func(ev *eventTypes.Event) bool { + update, ok := ev.Payload.(*types.TraceUpdate) + if !ok { + return false + } + return update.TraceID == traceID + })) + + // THEN snapshot historical updates (may overlap with live events). + historical := m.stateGetUpdates(since) + + // Build a set of historical event identifiers for dedup. + // Key: "type:nodeID:timestamp" is unique enough for trace events. + histSeen := make(map[string]struct{}, len(historical)) + for _, u := range historical { + histSeen[dedupKey(u)] = struct{}{} + } + + go func() { + defer close(out) + defer event.Unsubscribe(subID) + + for _, update := range historical { + out <- update + } + + for ev := range liveCh { + update, ok := ev.Payload.(*types.TraceUpdate) + if !ok { + continue + } + key := dedupKey(update) + if _, dup := histSeen[key]; dup { + delete(histSeen, key) + continue + } + out <- update + if update.Type == types.UpdateTypeComplete { + return + } + } + }() + + return out, nil } diff --git a/trace/trace.go b/trace/trace.go index ebdb837b..3450567c 100644 --- a/trace/trace.go +++ b/trace/trace.go @@ -4,13 +4,11 @@ import ( "context" "fmt" "sync" - "sync/atomic" "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" ) @@ -21,14 +19,10 @@ const ( Store = "store" // Gou store storage ) -// Global trace registry and pubsub services +// Global trace registry 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 @@ -145,23 +139,9 @@ func New(ctx context.Context, driver string, option *types.TraceOption, driverOp return LoadFromStorage(ctx, driver, traceID, driverOptions...) } - // 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) + // Create Manager instance with the driver + manager, err := NewManager(ctx, traceID, drv, 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) } @@ -198,13 +178,6 @@ 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 @@ -253,29 +226,9 @@ func LoadFromStorage(ctx context.Context, driver string, traceID string, driverO return "", nil, fmt.Errorf("trace not found in storage: %s", traceID) } - // 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, pubsubService, nil) + manager, err := NewManager(ctx, traceID, drv, 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) } @@ -356,12 +309,6 @@ func MarkCancelled(traceID string, reason string) error { 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 @@ -395,12 +342,7 @@ func MarkCancelled(traceID string, reason string) error { 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) + log.Trace("[TRACE] MarkCancelled: broadcasting node failed event for node %s", node.ID) mgr.addUpdateAndBroadcast(&types.TraceUpdate{ Type: types.UpdateTypeNodeFailed, TraceID: traceID, @@ -445,12 +387,7 @@ func MarkCancelled(traceID string, reason string) error { 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) + log.Trace("[TRACE] MarkCancelled: broadcasting completion update") totalDuration := int64(0) if rootNode.CreatedAt > 0 { totalDuration = now - rootNode.CreatedAt @@ -488,35 +425,7 @@ func Release(traceID string) error { return fmt.Errorf("trace not found in registry: %s", traceID) } - // Stop manager with safe three-step shutdown sequence. - // Order matters: flag blocks new writes -> cancel unblocks in-flight safeSend -> - // close terminates state worker (which drains remaining buffer first). - if mgr, ok := info.Manager.(*manager); ok { - // Step 1: Set closed flag β€” new safeSend calls return false immediately - atomic.StoreInt32(&mgr.closed, 1) - - // Step 2: Cancel context β€” unblocks any safeSend blocked in select on ctx.Done - if mgr.cancel != nil { - mgr.cancel() - } - - // Step 3: Close channel β€” state worker for-range exits after draining buffer - close(mgr.stateCmdChan) - } - - // 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() - } + _ = info.Manager log.Trace("[TRACE] Release: completed") return nil diff --git a/trace/trace_basic_test.go b/trace/trace_basic_test.go index 77564a51..e4a2512b 100644 --- a/trace/trace_basic_test.go +++ b/trace/trace_basic_test.go @@ -13,11 +13,8 @@ import ( ) func TestMain(m *testing.M) { - // Prepare test environment (initializes stores, models, etc.) test.Prepare(&testing.T{}, config.Conf) defer test.Clean() - - // Run tests os.Exit(m.Run()) } @@ -234,15 +231,17 @@ func TestContextCancellation(t *testing.T) { defer trace.Release(traceID) defer trace.Remove(context.Background(), d.DriverType, traceID, d.DriverOptions...) - // Cancel context + // Cancel the creation context β€” trace operations should still work. + // This is the core context-fork fix: trace managers no longer bind + // to the caller's context, so parent cancellation cannot break child ops. cancel() - // Operations should fail with context error - _, err = manager.Add("test", types.TraceNodeOption{ + node, err := manager.Add("test", types.TraceNodeOption{ Label: "Test", Type: "test", }) - assert.Error(t, err) + assert.NoError(t, err) + assert.NotNil(t, node) }) } } diff --git a/trace/trace_lifecycle_test.go b/trace/trace_lifecycle_test.go index c28f4452..88fa1a24 100644 --- a/trace/trace_lifecycle_test.go +++ b/trace/trace_lifecycle_test.go @@ -186,9 +186,10 @@ func TestConcurrentReleaseAndMarkCancelled(t *testing.T) { } } -// TestSafeSendAfterClosed verifies that operations using safeSend after -// Release return gracefully instead of panicking. -func TestSafeSendAfterClosed(t *testing.T) { +// TestOperationsAfterRelease verifies that using a manager reference after +// Release does not panic. The manager is removed from registry but its +// in-memory state remains valid (no channel close or context cancel). +func TestOperationsAfterRelease(t *testing.T) { drivers := trace.GetTestDrivers() for _, d := range drivers { @@ -205,21 +206,13 @@ func TestSafeSendAfterClosed(t *testing.T) { err = trace.Release(traceID) assert.NoError(t, err) - // All of these internally use safeSend. After Release they should - // return nil/error/zero-value, never panic. - manager.Info("post-close info") - manager.Debug("post-close debug") - manager.Error("post-close error") - manager.Warn("post-close warn") + // After Release, the manager object is still usable (state in memory). + // These calls should not panic. + manager.Info("post-release info") + manager.Debug("post-release debug") root, _ := manager.GetRootNode() - assert.Nil(t, root) - - nodes, _ := manager.GetCurrentNodes() - assert.Nil(t, nodes) - - status := manager.IsComplete() - assert.True(t, status) + assert.NotNil(t, root) }) } }