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

Update dependencies and enhance logging in the Assistant module

Browse source 
- 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.
This commit is contained in:
Max 2026-02-23 19:22:31 +08:00
parent fd24e31912
commit af981dc0b0
24 changed files with 1678 additions and 864 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
agent/assistant/agent.go
View file

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

27
agent/assistant/mcp.go
View file

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

48
agent/caller/orchestrator.go
View file

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

10
agent/context/context.go
View file

@ -71,13 +71,14 @@ 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.ForkParent == nil {
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)
@ -85,6 +86,7 @@ func (ctx *Context) Release() {
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

335
agent/context/log.go
View file

@ -72,10 +72,11 @@ type LogEntry struct {
// RequestLogger provides request-scoped async logging
type RequestLogger struct {
assistantID string
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
@ -86,13 +87,23 @@ 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,
assistantIDStack: []string{assistantID},
chatID: chatID,
requestID: requestID,
shortID: shortID(requestID),
@ -101,6 +112,10 @@ func NewRequestLogger(assistantID, chatID, requestID string) *RequestLogger {
done: make(chan struct{}),
}
for _, opt := range opts {
opt(l)
}
// Start consumer goroutine
go l.consume()
@ -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
}
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))
l.shortID, l.currentAssistantID(), shortID(l.chatID), shortID(l.requestID))
if !config.IsDevelopment() {
return
}
SendTUI(AgentEventMsg{
RequestID: l.requestID,
ParentID: l.parentID,
AssistantID: l.currentAssistantID(),
Event: EventRequestStart,
})
if GetTUIProgram() != nil {
return
}
// Development: colorful output (direct print, not through channel for immediate display)
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 !config.IsDevelopment() {
if success {
kunlog.Trace("[AGENT] Request %s completed: assistant=%s, duration=%v",
l.shortID, l.assistantID, duration.Round(time.Millisecond))
l.shortID, l.currentAssistantID(), 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)
kunlog.Error("[AGENT] Request %s failed: assistant=%s, duration=%v, error=%v",
l.shortID, l.currentAssistantID(), duration.Round(time.Millisecond), err)
}
if !config.IsDevelopment() {
return
}
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
}
// Development: colorful output (direct print for immediate display)
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)
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)
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)
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)
} else {
kunlog.Trace("[AGENT] %s LLM call: connector=%s, model=%s, messages=%d (+%v)", l.shortID, connector, model, messageCount, elapsed)
}
}
// 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)
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)
} else {
kunlog.Trace("[AGENT] %s LLM response: status=%s, tokens=%d (+%v)", l.shortID, status, tokens, elapsed)
}
}
// 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)
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)
}
} 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)
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)
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)
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)
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)
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)
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)
}
// =============================================================================

799
agent/context/tui.go Normal file
View file

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

90
agent/context/tui_msg.go Normal file
View file

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

42
cmd/start.go
View file

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

15
engine/load.go
View file

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

14
go.mod
View file

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

29
go.sum
View file

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

19
test/utils.go
View file

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

25
trace/event_listener.go Normal file
View file

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

28
trace/handler.go Normal file
View file

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

142
trace/manager.go
View file

@ -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
mu sync.Mutex
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)
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
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
}
}
// 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
}

23
trace/node.go
View file

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

142
trace/pubsub/pubsub.go
View file

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

62
trace/pubsub/subscriber.go
View file

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

20
trace/space.go
View file

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

362
trace/state.go
View file

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

78
trace/subscription.go
View file

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

105
trace/trace.go
View file

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

13
trace/trace_basic_test.go
View file

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

25
trace/trace_lifecycle_test.go
View file

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