Update dependencies and enhance logging in the Assistant module
- Add new indirect dependencies including various Charmbracelet packages for improved UI handling. - Enhance logging in the Assistant module by adding tool completion and start logging for better traceability of tool calls. - Modify context handling in the RequestLogger to support a stack-based assistant ID management, improving the logging structure for agent requests. - Implement event service integration for better trace management and debugging capabilities.
This commit is contained in:
Max
parent
fd24e31912
commit
af981dc0b0
24 changed files with 1678 additions and 864 deletions
|
|
@ -103,6 +103,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
|
||||||
|
|
||||||
// Log end of request
|
// Log end of request
|
||||||
ctx.Logger.End(finalStatus == context.StepStatusCompleted, finalError)
|
ctx.Logger.End(finalStatus == context.StepStatusCompleted, finalError)
|
||||||
|
ctx.Logger.RestoreAssistantID()
|
||||||
}()
|
}()
|
||||||
|
|
||||||
// Determine stream handler
|
// Determine stream handler
|
||||||
|
|
|
||||||
|
|
@ -319,6 +319,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall
|
||||||
result.Content = result.Error.Error()
|
result.Content = result.Error.Error()
|
||||||
result.IsRetryableError = true // Argument parsing error is retryable by LLM
|
result.IsRetryableError = true // Argument parsing error is retryable by LLM
|
||||||
ctx.Logger.Error("Failed to parse arguments: %v", err)
|
ctx.Logger.Error("Failed to parse arguments: %v", err)
|
||||||
|
ctx.Logger.ToolComplete(toolCall.Function.Name, false)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Fail(result.Error)
|
toolNode.Fail(result.Error)
|
||||||
}
|
}
|
||||||
|
|
@ -333,6 +334,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall
|
||||||
result.Content = result.Error.Error()
|
result.Content = result.Error.Error()
|
||||||
result.IsRetryableError = true // Type error is retryable by LLM
|
result.IsRetryableError = true // Type error is retryable by LLM
|
||||||
ctx.Logger.Error("Arguments must be an object, got %T", parsed)
|
ctx.Logger.Error("Arguments must be an object, got %T", parsed)
|
||||||
|
ctx.Logger.ToolComplete(toolCall.Function.Name, false)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Fail(result.Error)
|
toolNode.Fail(result.Error)
|
||||||
}
|
}
|
||||||
|
|
@ -346,6 +348,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall
|
||||||
result.Content = result.Error.Error()
|
result.Content = result.Error.Error()
|
||||||
result.IsRetryableError = true // Validation error is retryable by LLM
|
result.IsRetryableError = true // Validation error is retryable by LLM
|
||||||
ctx.Logger.Error("Argument validation failed: %v", err)
|
ctx.Logger.Error("Argument validation failed: %v", err)
|
||||||
|
ctx.Logger.ToolComplete(toolCall.Function.Name, false)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Fail(result.Error)
|
toolNode.Fail(result.Error)
|
||||||
}
|
}
|
||||||
|
|
@ -393,7 +396,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall
|
||||||
}
|
}
|
||||||
|
|
||||||
result.Content = string(contentBytes)
|
result.Content = string(contentBytes)
|
||||||
ctx.Logger.ToolComplete(toolName, true)
|
ctx.Logger.ToolComplete(toolCall.Function.Name, true)
|
||||||
|
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Complete(map[string]any{
|
toolNode.Complete(map[string]any{
|
||||||
|
|
@ -563,6 +566,7 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context
|
||||||
Arguments: args,
|
Arguments: args,
|
||||||
})
|
})
|
||||||
callMap[toolName] = tc
|
callMap[toolName] = tc
|
||||||
|
ctx.Logger.ToolStart(tc.Function.Name)
|
||||||
|
|
||||||
// Add trace input for this tool
|
// Add trace input for this tool
|
||||||
parallelInputs = append(parallelInputs, types.TraceParallelInput{
|
parallelInputs = append(parallelInputs, types.TraceParallelInput{
|
||||||
|
|
@ -598,11 +602,15 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context
|
||||||
mcpResponse, err := client.CallToolsParallel(mcpCtx, mcpCalls, ctx)
|
mcpResponse, err := client.CallToolsParallel(mcpCtx, mcpCalls, ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
ctx.Logger.Error("Parallel call failed: %v", err)
|
ctx.Logger.Error("Parallel call failed: %v", err)
|
||||||
// Mark all trace nodes as failed
|
for i, node := range toolNodes {
|
||||||
for _, node := range toolNodes {
|
|
||||||
if node != nil {
|
if node != nil {
|
||||||
node.Fail(err)
|
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
|
return nil, true
|
||||||
}
|
}
|
||||||
|
|
@ -631,6 +639,7 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context
|
||||||
result.Content = result.Error.Error()
|
result.Content = result.Error.Error()
|
||||||
result.IsRetryableError = false // Serialization error is not retryable
|
result.IsRetryableError = false // Serialization error is not retryable
|
||||||
hasErrors = true
|
hasErrors = true
|
||||||
|
ctx.Logger.ToolComplete(originalCall.Function.Name, false)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Fail(result.Error)
|
toolNode.Fail(result.Error)
|
||||||
}
|
}
|
||||||
|
|
@ -643,11 +652,12 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context
|
||||||
result.IsRetryableError = isRetryableToolError(result.Error)
|
result.IsRetryableError = isRetryableToolError(result.Error)
|
||||||
hasErrors = true
|
hasErrors = true
|
||||||
ctx.Logger.Error("Tool call failed: %s - %s (retryable: %v)", toolName, result.Content, result.IsRetryableError)
|
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 {
|
if toolNode != nil {
|
||||||
toolNode.Fail(result.Error)
|
toolNode.Fail(result.Error)
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
// Success
|
ctx.Logger.ToolComplete(originalCall.Function.Name, true)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Complete(map[string]any{
|
toolNode.Complete(map[string]any{
|
||||||
"result": mcpResult.Content,
|
"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)
|
ctx.Logger.Debug("Calling %d tools sequentially on server '%s'", len(toolCalls), serverID)
|
||||||
|
|
||||||
for _, tc := range toolCalls {
|
for _, tc := range toolCalls {
|
||||||
|
ctx.Logger.ToolStart(tc.Function.Name)
|
||||||
|
|
||||||
_, toolName, ok := ParseMCPToolName(tc.Function.Name)
|
_, toolName, ok := ParseMCPToolName(tc.Function.Name)
|
||||||
if !ok {
|
if !ok {
|
||||||
results = append(results, ToolCallResult{
|
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),
|
Content: fmt.Sprintf("Invalid tool name format: %s", tc.Function.Name),
|
||||||
Error: fmt.Errorf("invalid tool name format"),
|
Error: fmt.Errorf("invalid tool name format"),
|
||||||
})
|
})
|
||||||
|
ctx.Logger.ToolComplete(tc.Function.Name, false)
|
||||||
hasErrors = true
|
hasErrors = true
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
@ -727,6 +740,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte
|
||||||
}
|
}
|
||||||
results = append(results, result)
|
results = append(results, result)
|
||||||
hasErrors = true
|
hasErrors = true
|
||||||
|
ctx.Logger.ToolComplete(tc.Function.Name, false)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Fail(err)
|
toolNode.Fail(err)
|
||||||
}
|
}
|
||||||
|
|
@ -747,6 +761,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte
|
||||||
}
|
}
|
||||||
results = append(results, result)
|
results = append(results, result)
|
||||||
hasErrors = true
|
hasErrors = true
|
||||||
|
ctx.Logger.ToolComplete(tc.Function.Name, false)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Fail(err)
|
toolNode.Fail(err)
|
||||||
}
|
}
|
||||||
|
|
@ -765,6 +780,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte
|
||||||
}
|
}
|
||||||
results = append(results, result)
|
results = append(results, result)
|
||||||
hasErrors = true
|
hasErrors = true
|
||||||
|
ctx.Logger.ToolComplete(tc.Function.Name, false)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Fail(err)
|
toolNode.Fail(err)
|
||||||
}
|
}
|
||||||
|
|
@ -788,6 +804,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte
|
||||||
result.IsRetryableError = isRetryableToolError(err)
|
result.IsRetryableError = isRetryableToolError(err)
|
||||||
hasErrors = true
|
hasErrors = true
|
||||||
ctx.Logger.Error("Tool call failed: %s - %v (retryable: %v)", toolName, err, result.IsRetryableError)
|
ctx.Logger.Error("Tool call failed: %s - %v (retryable: %v)", toolName, err, result.IsRetryableError)
|
||||||
|
ctx.Logger.ToolComplete(tc.Function.Name, false)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Fail(err)
|
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.Content = fmt.Sprintf("Failed to serialize result: %v", err)
|
||||||
result.IsRetryableError = false // Serialization error is not retryable
|
result.IsRetryableError = false // Serialization error is not retryable
|
||||||
hasErrors = true
|
hasErrors = true
|
||||||
|
ctx.Logger.ToolComplete(tc.Function.Name, false)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Fail(err)
|
toolNode.Fail(err)
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
result.Content = string(contentBytes)
|
result.Content = string(contentBytes)
|
||||||
|
ctx.Logger.ToolComplete(tc.Function.Name, !mcpResult.IsError)
|
||||||
if toolNode != nil {
|
if toolNode != nil {
|
||||||
toolNode.Complete(map[string]any{
|
toolNode.Complete(map[string]any{
|
||||||
"result": mcpResult.Content,
|
"result": mcpResult.Content,
|
||||||
|
|
|
||||||
|
|
@ -5,6 +5,7 @@ import (
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
agentContext "github.com/yaoapp/yao/agent/context"
|
agentContext "github.com/yaoapp/yao/agent/context"
|
||||||
|
"github.com/yaoapp/yao/trace/types"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Orchestrator handles parallel agent calls with different concurrency patterns
|
// 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
|
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
|
// Execute the agent call with the provided context
|
||||||
// The agent.Stream method will use the context's Writer for output
|
// The agent.Stream method will use the context's Writer for output
|
||||||
resp, err := agent.Stream(ctx, req.Messages, ctxOpts)
|
resp, err := agent.Stream(ctx, req.Messages, ctxOpts)
|
||||||
if err != nil {
|
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))
|
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)
|
return NewResult(req.AgentID, resp, nil)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -71,19 +71,21 @@ func (ctx *Context) Release() {
|
||||||
ctx.Interrupt = nil
|
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.trace != nil && ctx.Stack != nil && ctx.Stack.TraceID != "" {
|
||||||
if ctx.Logger != nil {
|
if ctx.ForkParent == nil {
|
||||||
ctx.Logger.Cleanup("Trace: " + ctx.Stack.TraceID)
|
if ctx.Logger != nil {
|
||||||
}
|
ctx.Logger.Cleanup("Trace: " + ctx.Stack.TraceID)
|
||||||
|
}
|
||||||
// Check if context is cancelled - if so, mark as cancelled instead of complete
|
if ctx.Context != nil && ctx.Context.Err() != nil {
|
||||||
if ctx.Context != nil && ctx.Context.Err() != nil {
|
trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error())
|
||||||
trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error())
|
trace.Release(ctx.Stack.TraceID)
|
||||||
trace.Release(ctx.Stack.TraceID)
|
} else {
|
||||||
} else {
|
ctx.trace.MarkComplete()
|
||||||
ctx.trace.MarkComplete()
|
trace.Release(ctx.Stack.TraceID)
|
||||||
trace.Release(ctx.Stack.TraceID)
|
}
|
||||||
}
|
}
|
||||||
ctx.trace = nil
|
ctx.trace = nil
|
||||||
}
|
}
|
||||||
|
|
@ -191,7 +193,7 @@ func (ctx *Context) Fork() *Context {
|
||||||
|
|
||||||
// Create independent resources to avoid race conditions
|
// Create independent resources to avoid race conditions
|
||||||
IDGenerator: message.NewIDGenerator(),
|
IDGenerator: message.NewIDGenerator(),
|
||||||
Logger: NewRequestLogger(ctx.AssistantID, ctx.ChatID, childID),
|
Logger: NewRequestLogger(ctx.AssistantID, ctx.ChatID, childID, WithParentID(ctx.ID)),
|
||||||
messageMetadata: newMessageMetadataStore(),
|
messageMetadata: newMessageMetadataStore(),
|
||||||
|
|
||||||
// Inherit context metadata
|
// Inherit context metadata
|
||||||
|
|
|
||||||
|
|
@ -72,11 +72,12 @@ type LogEntry struct {
|
||||||
|
|
||||||
// RequestLogger provides request-scoped async logging
|
// RequestLogger provides request-scoped async logging
|
||||||
type RequestLogger struct {
|
type RequestLogger struct {
|
||||||
assistantID string
|
assistantIDStack []string // Stack-based: delegate calls push, pop on exit; top = current
|
||||||
chatID string
|
chatID string
|
||||||
requestID string
|
requestID string
|
||||||
shortID string // Short version of requestID for display
|
shortID string // Short version of requestID for display
|
||||||
startTime time.Time
|
parentID string // Parent request ID for A2A tree structure
|
||||||
|
startTime time.Time
|
||||||
|
|
||||||
ch chan LogEntry
|
ch chan LogEntry
|
||||||
done chan struct{}
|
done chan struct{}
|
||||||
|
|
@ -86,19 +87,33 @@ type RequestLogger struct {
|
||||||
mu sync.RWMutex
|
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
|
// noopLogger is a shared no-op logger instance
|
||||||
var noopLogger = &RequestLogger{noop: true}
|
var noopLogger = &RequestLogger{noop: true}
|
||||||
|
|
||||||
// NewRequestLogger creates a new request-scoped logger with async processing
|
// 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{
|
l := &RequestLogger{
|
||||||
assistantID: assistantID,
|
assistantIDStack: []string{assistantID},
|
||||||
chatID: chatID,
|
chatID: chatID,
|
||||||
requestID: requestID,
|
requestID: requestID,
|
||||||
shortID: shortID(requestID),
|
shortID: shortID(requestID),
|
||||||
startTime: time.Now(),
|
startTime: time.Now(),
|
||||||
ch: make(chan LogEntry, 100), // Buffered channel
|
ch: make(chan LogEntry, 100), // Buffered channel
|
||||||
done: make(chan struct{}),
|
done: make(chan struct{}),
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, opt := range opts {
|
||||||
|
opt(l)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Start consumer goroutine
|
// Start consumer goroutine
|
||||||
|
|
@ -112,12 +127,35 @@ func Noop() *RequestLogger {
|
||||||
return noopLogger
|
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) {
|
func (l *RequestLogger) SetAssistantID(id string) {
|
||||||
if l.noop {
|
if l.noop {
|
||||||
return
|
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
|
// 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) {
|
func (l *RequestLogger) processEntry(entry LogEntry) {
|
||||||
if config.IsDevelopment() {
|
if config.IsDevelopment() {
|
||||||
l.printDev(entry)
|
l.printDev(entry)
|
||||||
|
l.writeLog(entry, true)
|
||||||
} else {
|
} 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) {
|
func (l *RequestLogger) printDev(entry LogEntry) {
|
||||||
|
if GetTUIProgram() != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
switch entry.Level {
|
switch entry.Level {
|
||||||
case LogLevelTrace:
|
case LogLevelTrace:
|
||||||
fmt.Printf("%s → %s%s\n", colorGray, entry.Message, colorReset)
|
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
|
// writeLog writes structured events to kun/log
|
||||||
func (l *RequestLogger) printProd(entry LogEntry) {
|
func (l *RequestLogger) writeLog(entry LogEntry, devMode bool) {
|
||||||
prefix := fmt.Sprintf("[AGENT] %s ", l.shortID)
|
prefix := fmt.Sprintf("[AGENT] %s ", l.shortID)
|
||||||
|
if devMode {
|
||||||
|
kunlog.Trace("%s%s", prefix, entry.Message)
|
||||||
|
return
|
||||||
|
}
|
||||||
switch entry.Level {
|
switch entry.Level {
|
||||||
case LogLevelTrace:
|
case LogLevelTrace:
|
||||||
kunlog.Trace("%s%s", prefix, entry.Message)
|
kunlog.Trace("%s%s", prefix, entry.Message)
|
||||||
|
|
@ -263,18 +308,29 @@ func (l *RequestLogger) Start() {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
kunlog.Trace("[AGENT] Request %s started: assistant=%s, chat=%s, request=%s",
|
||||||
|
l.shortID, l.currentAssistantID(), shortID(l.chatID), shortID(l.requestID))
|
||||||
|
|
||||||
if !config.IsDevelopment() {
|
if !config.IsDevelopment() {
|
||||||
kunlog.Trace("[AGENT] Request %s started: assistant=%s, chat=%s, request=%s",
|
|
||||||
l.shortID, l.assistantID, shortID(l.chatID), shortID(l.requestID))
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// Development: colorful output (direct print, not through channel for immediate display)
|
SendTUI(AgentEventMsg{
|
||||||
|
RequestID: l.requestID,
|
||||||
|
ParentID: l.parentID,
|
||||||
|
AssistantID: l.currentAssistantID(),
|
||||||
|
Event: EventRequestStart,
|
||||||
|
})
|
||||||
|
|
||||||
|
if GetTUIProgram() != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
fmt.Println()
|
fmt.Println()
|
||||||
fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("═", 60), colorReset)
|
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%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 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 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)
|
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)
|
duration := time.Since(l.startTime)
|
||||||
|
|
||||||
|
if success {
|
||||||
|
kunlog.Trace("[AGENT] Request %s completed: assistant=%s, duration=%v",
|
||||||
|
l.shortID, l.currentAssistantID(), duration.Round(time.Millisecond))
|
||||||
|
} else {
|
||||||
|
kunlog.Error("[AGENT] Request %s failed: assistant=%s, duration=%v, error=%v",
|
||||||
|
l.shortID, l.currentAssistantID(), duration.Round(time.Millisecond), err)
|
||||||
|
}
|
||||||
|
|
||||||
if !config.IsDevelopment() {
|
if !config.IsDevelopment() {
|
||||||
if success {
|
|
||||||
kunlog.Trace("[AGENT] Request %s completed: assistant=%s, duration=%v",
|
|
||||||
l.shortID, l.assistantID, duration.Round(time.Millisecond))
|
|
||||||
} else {
|
|
||||||
kunlog.Trace("[AGENT] Request %s failed: assistant=%s, duration=%v, error=%v",
|
|
||||||
l.shortID, l.assistantID, duration.Round(time.Millisecond), err)
|
|
||||||
}
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// Development: colorful output (direct print for immediate display)
|
data := map[string]interface{}{"duration": duration.Round(time.Millisecond)}
|
||||||
|
if err != nil {
|
||||||
|
data["error"] = err.Error()
|
||||||
|
}
|
||||||
|
SendTUI(AgentEventMsg{
|
||||||
|
RequestID: l.requestID,
|
||||||
|
AssistantID: l.currentAssistantID(),
|
||||||
|
Event: EventRequestEnd,
|
||||||
|
Data: data,
|
||||||
|
})
|
||||||
|
|
||||||
|
if GetTUIProgram() != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset)
|
fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset)
|
||||||
if success {
|
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 {
|
} 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 {
|
if err != nil {
|
||||||
fmt.Printf("%s Error: %s%v%s\n", colorGray, colorRed, err, colorReset)
|
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 Duration: %s%v%s\n", colorGray, colorWhite, duration.Round(time.Millisecond), colorReset)
|
||||||
fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset)
|
fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset)
|
||||||
fmt.Println()
|
fmt.Println()
|
||||||
|
|
@ -323,12 +394,22 @@ func (l *RequestLogger) Phase(name string) {
|
||||||
}
|
}
|
||||||
|
|
||||||
elapsed := time.Since(l.startTime).Round(time.Millisecond)
|
elapsed := time.Since(l.startTime).Round(time.Millisecond)
|
||||||
|
kunlog.Trace("[AGENT] %s Phase: %s (+%v)", l.shortID, name, elapsed)
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
if !config.IsDevelopment() {
|
||||||
fmt.Printf("%s ▶ %s%s %s[+%v]%s\n", colorBoldBlue, name, colorReset, colorGray, elapsed, colorReset)
|
return
|
||||||
} else {
|
|
||||||
kunlog.Trace("[AGENT] %s Phase: %s (+%v)", l.shortID, name, elapsed)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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
|
// 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)
|
elapsed := time.Since(l.startTime).Round(time.Millisecond)
|
||||||
|
kunlog.Trace("[AGENT] %s Phase completed: %s (+%v)", l.shortID, name, elapsed)
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
if !config.IsDevelopment() {
|
||||||
fmt.Printf("%s ✓ %s%s %s[+%v]%s\n", colorGreen, name, colorReset, colorGray, elapsed, colorReset)
|
return
|
||||||
} else {
|
|
||||||
kunlog.Trace("[AGENT] %s Phase completed: %s (+%v)", l.shortID, name, elapsed)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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)
|
// PhaseSkip logs a skipped phase (development only)
|
||||||
|
|
@ -351,9 +442,21 @@ func (l *RequestLogger) PhaseSkip(name, reason string) {
|
||||||
if l.noop {
|
if l.noop {
|
||||||
return
|
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
|
// 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)
|
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() {
|
if !config.IsDevelopment() {
|
||||||
fmt.Printf("%s 🤖 LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset)
|
return
|
||||||
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)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
SendTUI(AgentEventMsg{
|
||||||
|
RequestID: l.requestID,
|
||||||
|
Event: EventLLMCall,
|
||||||
|
Data: map[string]interface{}{
|
||||||
|
"connector": connector,
|
||||||
|
"model": model,
|
||||||
|
"messages": messageCount,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
if GetTUIProgram() != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
fmt.Printf("%s LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset)
|
||||||
|
fmt.Printf("%s Connector: %s%s%s\n", colorGray, colorWhite, connector, colorReset)
|
||||||
|
if model != "" {
|
||||||
|
fmt.Printf("%s Model: %s%s%s\n", colorGray, colorWhite, model, colorReset)
|
||||||
|
}
|
||||||
|
fmt.Printf("%s Messages: %s%d%s\n", colorGray, colorWhite, messageCount, colorReset)
|
||||||
}
|
}
|
||||||
|
|
||||||
// LLMComplete logs the completion of an LLM call
|
// LLMComplete logs the completion of an LLM call
|
||||||
|
|
@ -388,15 +505,30 @@ func (l *RequestLogger) LLMComplete(tokens int, hasToolCalls bool) {
|
||||||
status = "tool_calls"
|
status = "tool_calls"
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
kunlog.Trace("[AGENT] %s LLM response: status=%s, tokens=%d (+%v)", l.shortID, status, tokens, elapsed)
|
||||||
fmt.Printf("%s ✓ LLM Response (%s)%s", colorGreen, status, colorReset)
|
|
||||||
if tokens > 0 {
|
if !config.IsDevelopment() {
|
||||||
fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset)
|
return
|
||||||
}
|
|
||||||
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)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
SendTUI(AgentEventMsg{
|
||||||
|
RequestID: l.requestID,
|
||||||
|
Event: EventLLMDone,
|
||||||
|
Data: map[string]interface{}{
|
||||||
|
"detail": fmt.Sprintf("%s [tokens:%d, %v]", status, tokens, elapsed),
|
||||||
|
"tokens": tokens,
|
||||||
|
"status": status,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
if GetTUIProgram() != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
fmt.Printf("%s + LLM Response (%s)%s", colorGreen, status, colorReset)
|
||||||
|
if tokens > 0 {
|
||||||
|
fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset)
|
||||||
|
}
|
||||||
|
fmt.Printf(" %s[+%v]%s\n", colorGray, elapsed, colorReset)
|
||||||
}
|
}
|
||||||
|
|
||||||
// ToolStart logs the start of tool execution
|
// ToolStart logs the start of tool execution
|
||||||
|
|
@ -405,11 +537,22 @@ func (l *RequestLogger) ToolStart(toolName string) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
kunlog.Trace("[AGENT] %s Tool call: %s", l.shortID, toolName)
|
||||||
fmt.Printf("%s 🔧 Tool: %s%s\n", colorYellow, toolName, colorReset)
|
|
||||||
} else {
|
if !config.IsDevelopment() {
|
||||||
kunlog.Trace("[AGENT] %s Tool call: %s", l.shortID, toolName)
|
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
|
// ToolComplete logs the completion of tool execution
|
||||||
|
|
@ -418,18 +561,29 @@ func (l *RequestLogger) ToolComplete(toolName string, success bool) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
if success {
|
||||||
if success {
|
kunlog.Trace("[AGENT] %s Tool completed: %s", l.shortID, toolName)
|
||||||
fmt.Printf("%s ✓ %s completed%s\n", colorGreen, toolName, colorReset)
|
|
||||||
} else {
|
|
||||||
fmt.Printf("%s ✗ %s failed%s\n", colorRed, toolName, colorReset)
|
|
||||||
}
|
|
||||||
} else {
|
} else {
|
||||||
if success {
|
kunlog.Error("[AGENT] %s Tool failed: %s", l.shortID, toolName)
|
||||||
kunlog.Trace("[AGENT] %s Tool completed: %s", l.shortID, toolName)
|
}
|
||||||
} else {
|
|
||||||
kunlog.Trace("[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)
|
elapsed := time.Since(l.startTime).Round(time.Millisecond)
|
||||||
|
kunlog.Trace("[AGENT] %s Hook: %s (+%v)", l.shortID, hookName, elapsed)
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
if !config.IsDevelopment() {
|
||||||
fmt.Printf("%s 🪝 Hook: %s%s %s[+%v]%s\n", colorMagenta, hookName, colorReset, colorGray, elapsed, colorReset)
|
return
|
||||||
} else {
|
|
||||||
kunlog.Trace("[AGENT] %s Hook: %s (+%v)", l.shortID, hookName, elapsed)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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
|
// HookComplete logs the completion of a hook
|
||||||
|
|
@ -454,11 +618,22 @@ func (l *RequestLogger) HookComplete(hookName string) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
kunlog.Trace("[AGENT] %s Hook completed: %s", l.shortID, hookName)
|
||||||
fmt.Printf("%s ✓ %s done%s\n", colorGreen, hookName, colorReset)
|
|
||||||
} else {
|
if !config.IsDevelopment() {
|
||||||
kunlog.Trace("[AGENT] %s Hook completed: %s", l.shortID, hookName)
|
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
|
// Cleanup logs resource cleanup
|
||||||
|
|
@ -467,11 +642,12 @@ func (l *RequestLogger) Cleanup(resource string) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource)
|
||||||
fmt.Printf("%s ✓ %s%s\n", colorGray, resource, colorReset)
|
|
||||||
} else {
|
if !config.IsDevelopment() || GetTUIProgram() != nil {
|
||||||
kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource)
|
return
|
||||||
}
|
}
|
||||||
|
fmt.Printf("%s + %s%s\n", colorGray, resource, colorReset)
|
||||||
}
|
}
|
||||||
|
|
||||||
// HistoryLoad logs history loading
|
// HistoryLoad logs history loading
|
||||||
|
|
@ -480,11 +656,12 @@ func (l *RequestLogger) HistoryLoad(count, maxSize int) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize)
|
||||||
fmt.Printf("%s Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset)
|
|
||||||
} else {
|
if !config.IsDevelopment() || GetTUIProgram() != nil {
|
||||||
kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize)
|
return
|
||||||
}
|
}
|
||||||
|
fmt.Printf("%s Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset)
|
||||||
}
|
}
|
||||||
|
|
||||||
// HistoryOverlap logs overlap detection
|
// HistoryOverlap logs overlap detection
|
||||||
|
|
@ -494,11 +671,12 @@ func (l *RequestLogger) HistoryOverlap(overlapCount int) {
|
||||||
}
|
}
|
||||||
|
|
||||||
if overlapCount > 0 {
|
if overlapCount > 0 {
|
||||||
if config.IsDevelopment() {
|
kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount)
|
||||||
fmt.Printf("%s Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset)
|
|
||||||
} else {
|
if !config.IsDevelopment() || GetTUIProgram() != nil {
|
||||||
kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount)
|
return
|
||||||
}
|
}
|
||||||
|
fmt.Printf("%s Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -508,11 +686,22 @@ func (l *RequestLogger) Release() {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.IsDevelopment() {
|
kunlog.Trace("[AGENT] %s Release started", l.shortID)
|
||||||
fmt.Printf("%s 🧹 RELEASE %s%s %s(%s)%s\n", colorBoldYellow, l.shortID, colorReset, colorGray, l.assistantID, colorReset)
|
|
||||||
} else {
|
if !config.IsDevelopment() {
|
||||||
kunlog.Trace("[AGENT] %s Release started", l.shortID)
|
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
799
agent/context/tui.go
Normal 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
90
agent/context/tui_msg.go
Normal 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
42
cmd/start.go
|
|
@ -8,11 +8,14 @@ import (
|
||||||
"strings"
|
"strings"
|
||||||
"syscall"
|
"syscall"
|
||||||
|
|
||||||
|
tea "github.com/charmbracelet/bubbletea"
|
||||||
"github.com/fatih/color"
|
"github.com/fatih/color"
|
||||||
|
"github.com/mattn/go-isatty"
|
||||||
"github.com/spf13/cobra"
|
"github.com/spf13/cobra"
|
||||||
"github.com/yaoapp/gou/api"
|
"github.com/yaoapp/gou/api"
|
||||||
"github.com/yaoapp/gou/connector"
|
"github.com/yaoapp/gou/connector"
|
||||||
"github.com/yaoapp/gou/fs"
|
"github.com/yaoapp/gou/fs"
|
||||||
|
"github.com/yaoapp/gou/helper"
|
||||||
"github.com/yaoapp/gou/mcp"
|
"github.com/yaoapp/gou/mcp"
|
||||||
"github.com/yaoapp/gou/plugin"
|
"github.com/yaoapp/gou/plugin"
|
||||||
"github.com/yaoapp/gou/schedule"
|
"github.com/yaoapp/gou/schedule"
|
||||||
|
|
@ -20,7 +23,9 @@ import (
|
||||||
"github.com/yaoapp/gou/store"
|
"github.com/yaoapp/gou/store"
|
||||||
"github.com/yaoapp/gou/task"
|
"github.com/yaoapp/gou/task"
|
||||||
"github.com/yaoapp/gou/websocket"
|
"github.com/yaoapp/gou/websocket"
|
||||||
|
"github.com/yaoapp/kun/exception"
|
||||||
"github.com/yaoapp/kun/log"
|
"github.com/yaoapp/kun/log"
|
||||||
|
agentcontext "github.com/yaoapp/yao/agent/context"
|
||||||
"github.com/yaoapp/yao/config"
|
"github.com/yaoapp/yao/config"
|
||||||
"github.com/yaoapp/yao/engine"
|
"github.com/yaoapp/yao/engine"
|
||||||
"github.com/yaoapp/yao/openapi"
|
"github.com/yaoapp/yao/openapi"
|
||||||
|
|
@ -33,6 +38,7 @@ import (
|
||||||
|
|
||||||
var startDebug = false
|
var startDebug = false
|
||||||
var startDisableWatching = false
|
var startDisableWatching = false
|
||||||
|
var startTUI = false
|
||||||
|
|
||||||
var startCmd = &cobra.Command{
|
var startCmd = &cobra.Command{
|
||||||
Use: "start",
|
Use: "start",
|
||||||
|
|
@ -265,8 +271,9 @@ var startCmd = &cobra.Command{
|
||||||
|
|
||||||
switch v {
|
switch v {
|
||||||
case http.READY:
|
case http.READY:
|
||||||
fmt.Println(color.GreenString(L("✨Server is up and running...")))
|
fmt.Println(color.GreenString(L("Server is up and running...")))
|
||||||
fmt.Println(color.GreenString("✨Ctrl+C to stop"))
|
fmt.Println(color.GreenString("Ctrl+C to stop"))
|
||||||
|
initAgentTUI()
|
||||||
break
|
break
|
||||||
|
|
||||||
case http.CLOSED:
|
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() {
|
func init() {
|
||||||
startCmd.PersistentFlags().BoolVarP(&startDebug, "debug", "", false, L("Development mode"))
|
startCmd.PersistentFlags().BoolVarP(&startDebug, "debug", "", false, L("Development mode"))
|
||||||
startCmd.PersistentFlags().BoolVarP(&startDisableWatching, "disable-watching", "", false, L("Disable watching"))
|
startCmd.PersistentFlags().BoolVarP(&startDisableWatching, "disable-watching", "", false, L("Disable watching"))
|
||||||
|
startCmd.PersistentFlags().BoolVarP(&startTUI, "tui", "", false, L("Enable TUI for agent request visualization"))
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
package engine
|
package engine
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"log"
|
"log"
|
||||||
"os"
|
"os"
|
||||||
|
|
@ -24,6 +25,7 @@ import (
|
||||||
"github.com/yaoapp/yao/config"
|
"github.com/yaoapp/yao/config"
|
||||||
"github.com/yaoapp/yao/connector"
|
"github.com/yaoapp/yao/connector"
|
||||||
"github.com/yaoapp/yao/data"
|
"github.com/yaoapp/yao/data"
|
||||||
|
"github.com/yaoapp/yao/event"
|
||||||
"github.com/yaoapp/yao/flow"
|
"github.com/yaoapp/yao/flow"
|
||||||
"github.com/yaoapp/yao/fs"
|
"github.com/yaoapp/yao/fs"
|
||||||
"github.com/yaoapp/yao/i18n"
|
"github.com/yaoapp/yao/i18n"
|
||||||
|
|
@ -48,6 +50,8 @@ import (
|
||||||
"github.com/yaoapp/yao/websocket"
|
"github.com/yaoapp/yao/websocket"
|
||||||
"github.com/yaoapp/yao/widget"
|
"github.com/yaoapp/yao/widget"
|
||||||
"github.com/yaoapp/yao/widgets"
|
"github.com/yaoapp/yao/widgets"
|
||||||
|
|
||||||
|
_ "github.com/yaoapp/yao/trace" // register trace handler/listener via init()
|
||||||
)
|
)
|
||||||
|
|
||||||
// LoadHooks used to load custom widgets/processes
|
// 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})
|
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
|
// Load Uploaders
|
||||||
err = loadStep("Uploader", func() error {
|
err = loadStep("Uploader", func() error {
|
||||||
return attachment.Load(cfg)
|
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
|
// Stop Runtime
|
||||||
err = runtime.Stop()
|
err = runtime.Stop()
|
||||||
|
|
||||||
|
|
|
||||||
14
go.mod
14
go.mod
|
|
@ -64,11 +64,18 @@ require (
|
||||||
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 // indirect
|
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 // indirect
|
||||||
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.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/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/blang/semver/v4 v4.0.0 // indirect
|
||||||
github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect
|
github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect
|
||||||
github.com/bytedance/sonic v1.13.2 // indirect
|
github.com/bytedance/sonic v1.13.2 // indirect
|
||||||
github.com/bytedance/sonic/loader v0.2.4 // indirect
|
github.com/bytedance/sonic/loader v0.2.4 // indirect
|
||||||
github.com/cespare/xxhash/v2 v2.3.0 // 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/cloudwego/base64x v0.1.5 // indirect
|
||||||
github.com/containerd/errdefs v1.0.0 // indirect
|
github.com/containerd/errdefs v1.0.0 // indirect
|
||||||
github.com/containerd/errdefs/pkg v0.3.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/dlclark/regexp2 v1.11.5 // indirect
|
||||||
github.com/docker/go-units v0.5.0 // indirect
|
github.com/docker/go-units v0.5.0 // indirect
|
||||||
github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6 // 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/felixge/httpsnoop v1.0.4 // indirect
|
||||||
github.com/gabriel-vasile/mimetype v1.4.9 // indirect
|
github.com/gabriel-vasile/mimetype v1.4.9 // indirect
|
||||||
github.com/gin-contrib/sse v1.1.0 // 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/klauspost/cpuid/v2 v2.2.10 // indirect
|
||||||
github.com/leodido/go-urn v1.4.0 // indirect
|
github.com/leodido/go-urn v1.4.0 // indirect
|
||||||
github.com/lib/pq v1.10.9 // 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/mark3labs/mcp-go v0.32.0 // indirect
|
||||||
github.com/mattn/go-colorable v0.1.14 // indirect
|
github.com/mattn/go-colorable v0.1.14 // indirect
|
||||||
github.com/mattn/go-isatty v0.0.20 // 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-runewidth v0.0.16 // indirect
|
||||||
github.com/mattn/go-sqlite3 v1.14.28 // indirect
|
github.com/mattn/go-sqlite3 v1.14.28 // indirect
|
||||||
github.com/miekg/dns v1.1.66 // 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/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
|
||||||
github.com/modern-go/reflect2 v1.0.2 // indirect
|
github.com/modern-go/reflect2 v1.0.2 // indirect
|
||||||
github.com/montanaflynn/stats v0.7.1 // 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/neo4j/neo4j-go-driver/v5 v5.28.1 // indirect
|
||||||
github.com/oklog/run v1.1.0 // indirect
|
github.com/oklog/run v1.1.0 // indirect
|
||||||
github.com/opencontainers/go-digest v1.0.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/pbkdf2 v1.0.0 // indirect
|
||||||
github.com/xdg-go/scram v1.1.2 // indirect
|
github.com/xdg-go/scram v1.1.2 // indirect
|
||||||
github.com/xdg-go/stringprep v1.0.4 // 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/efp v0.0.1 // indirect
|
||||||
github.com/xuri/nfp v0.0.1 // indirect
|
github.com/xuri/nfp v0.0.1 // indirect
|
||||||
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
|
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
|
||||||
|
|
|
||||||
29
go.sum
29
go.sum
|
|
@ -39,6 +39,8 @@ github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3 h1:BRXS0U76Z8wfF+bnkilA2QwpIch6U
|
||||||
github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3/go.mod h1:bNXKFFyaiVvWuR6O16h/I1724+aXe/tAkA9/QS01t5k=
|
github.com/aws/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 h1:Z//5NuZCSW6R4PhQ93hShNbyBbn8BWCmCVCt+Q8Io5k=
|
||||||
github.com/aws/smithy-go v1.22.3/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI=
|
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 h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ=
|
||||||
github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk=
|
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=
|
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/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 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
||||||
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
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 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4=
|
||||||
github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
|
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=
|
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 h1:oP4q0fw+fOSWn3DfFi4EXdT+B+gTtzx8GC9xsc26Znk=
|
||||||
github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6/go.mod h1:iL2twTeMvZnrg54ZoPDNfJaJaqy0xIQFuBdrLsmspwQ=
|
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/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 h1:k1bTSim+usBG27w7BfOCorhgx3tO+6bAfMj5pR+6SKg=
|
||||||
github.com/evanw/esbuild v0.25.4/go.mod h1:D2vIQZqV/vIf/VRHtViaUtViZmG7o+kKmlBfVQuRi48=
|
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=
|
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/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 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
|
||||||
github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
|
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 h1:fgwmbfL2gbd67obg57OfV2Dnrhs1HtSdlY/i5fn7MU8=
|
||||||
github.com/mark3labs/mcp-go v0.32.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
|
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=
|
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.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 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
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 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
|
||||||
github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
|
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=
|
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/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 h1:BtR3DsxpApHfKReaPO1fCqF4pThRwH9uwvXzm+GnMFQ=
|
||||||
github.com/mozillazg/go-pinyin v0.20.0/go.mod h1:iR4EnMMRXkfpFVV5FMi4FNB6wGq9NV6uDWbUuPhP4Yc=
|
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 h1:RKWQW7wTgYAY2fU9S+9LaJ9OwRPbRc0I17tlT7nDmAY=
|
||||||
github.com/neo4j/neo4j-go-driver/v5 v5.28.1/go.mod h1:Vff8OwT7QpLm7L2yYr85XNWe9Rbqlbeb9asNXJTHO4k=
|
github.com/neo4j/neo4j-go-driver/v5 v5.28.1/go.mod h1:Vff8OwT7QpLm7L2yYr85XNWe9Rbqlbeb9asNXJTHO4k=
|
||||||
github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
|
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/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 h1:XLI/Ng3O1Atzq0oBs3TWm+5ZVgkq2aqdlvP9JtoZ6c8=
|
||||||
github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM=
|
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 h1:fws5Rv3myXyYni8uwj2qKjVaRP30PdjeYe2Y6FDsCL8=
|
||||||
github.com/xuri/efp v0.0.1/go.mod h1:ybY/Jr0T0GTCnYjKqmdwxyxn2BQf2RcQIIvex5QldPI=
|
github.com/xuri/efp v0.0.1/go.mod h1:ybY/Jr0T0GTCnYjKqmdwxyxn2BQf2RcQIIvex5QldPI=
|
||||||
github.com/xuri/excelize/v2 v2.9.1 h1:VdSGk+rraGmgLHGFaGG9/9IWu1nj4ufjJ7uwMDtj8Qw=
|
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-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-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-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-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-20220503163025-988cb79eb6c6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||||
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||||
|
|
|
||||||
|
|
@ -66,12 +66,14 @@
|
||||||
// 7. Starts V8 JavaScript runtime
|
// 7. Starts V8 JavaScript runtime
|
||||||
// 8. Registers query engines for database operations
|
// 8. Registers query engines for database operations
|
||||||
// 9. Creates temporary data directories for test isolation
|
// 9. Creates temporary data directories for test isolation
|
||||||
|
// 10. Starts the Event Service (handlers registered via init(), e.g. trace)
|
||||||
//
|
//
|
||||||
// WHAT test.Clean() DOES:
|
// WHAT test.Clean() DOES:
|
||||||
// 1. Stops V8 runtime and releases resources
|
// 1. Stops the Event Service (drains in-flight events)
|
||||||
// 2. Closes all database connections
|
// 2. Stops V8 runtime and releases resources
|
||||||
// 3. Removes temporary test data stores
|
// 3. Closes all database connections
|
||||||
// 4. Resets global state to prevent test interference
|
// 4. Removes temporary test data stores
|
||||||
|
// 5. Resets global state to prevent test interference
|
||||||
//
|
//
|
||||||
// WHAT test.Start() DOES:
|
// WHAT test.Start() DOES:
|
||||||
// 1. Creates Gin HTTP server with API routes
|
// 1. Creates Gin HTTP server with API routes
|
||||||
|
|
@ -159,6 +161,7 @@
|
||||||
package test
|
package test
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
|
@ -184,11 +187,14 @@ import (
|
||||||
"github.com/yaoapp/xun/capsule"
|
"github.com/yaoapp/xun/capsule"
|
||||||
"github.com/yaoapp/yao/config"
|
"github.com/yaoapp/yao/config"
|
||||||
"github.com/yaoapp/yao/data"
|
"github.com/yaoapp/yao/data"
|
||||||
|
"github.com/yaoapp/yao/event"
|
||||||
"github.com/yaoapp/yao/fs"
|
"github.com/yaoapp/yao/fs"
|
||||||
"github.com/yaoapp/yao/helper"
|
"github.com/yaoapp/yao/helper"
|
||||||
"github.com/yaoapp/yao/runtime"
|
"github.com/yaoapp/yao/runtime"
|
||||||
"github.com/yaoapp/yao/share"
|
"github.com/yaoapp/yao/share"
|
||||||
"github.com/yaoapp/yao/utils"
|
"github.com/yaoapp/yao/utils"
|
||||||
|
|
||||||
|
_ "github.com/yaoapp/yao/trace" // register trace event handler via init()
|
||||||
)
|
)
|
||||||
|
|
||||||
var testServer *http.Server = nil
|
var testServer *http.Server = nil
|
||||||
|
|
@ -488,10 +494,15 @@ func Prepare(t *testing.T, cfg config.Config, opts ...interface{}) {
|
||||||
load(t, cfg)
|
load(t, cfg)
|
||||||
startRuntime(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
|
// Clean the test environment
|
||||||
func Clean() {
|
func Clean() {
|
||||||
|
event.Stop(context.Background())
|
||||||
dbclose()
|
dbclose()
|
||||||
runtime.Stop()
|
runtime.Stop()
|
||||||
|
|
||||||
|
|
|
||||||
25
trace/event_listener.go
Normal file
25
trace/event_listener.go
Normal 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
28
trace/handler.go
Normal 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),
|
||||||
|
)
|
||||||
|
}
|
||||||
154
trace/manager.go
154
trace/manager.go
|
|
@ -3,63 +3,56 @@ package trace
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
gonanoid "github.com/matoous/go-nanoid/v2"
|
gonanoid "github.com/matoous/go-nanoid/v2"
|
||||||
"github.com/yaoapp/kun/log"
|
"github.com/yaoapp/kun/log"
|
||||||
"github.com/yaoapp/yao/trace/pubsub"
|
"github.com/yaoapp/yao/event"
|
||||||
"github.com/yaoapp/yao/trace/types"
|
"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 {
|
type manager struct {
|
||||||
ctx context.Context
|
mu sync.Mutex
|
||||||
cancel context.CancelFunc
|
traceID string
|
||||||
traceID string
|
driver types.Driver
|
||||||
driver types.Driver
|
state *managerState
|
||||||
stateCmdChan chan stateCommand // Single channel for all state mutations
|
autoArchive bool
|
||||||
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)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewManager creates a new trace manager instance
|
// 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, option *types.TraceOption) (types.Manager, error) {
|
||||||
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
|
|
||||||
autoArchive := false
|
autoArchive := false
|
||||||
if option != nil {
|
if option != nil {
|
||||||
autoArchive = option.AutoArchive
|
autoArchive = option.AutoArchive
|
||||||
}
|
}
|
||||||
|
|
||||||
m := &manager{
|
m := &manager{
|
||||||
ctx: managerCtx,
|
traceID: traceID,
|
||||||
cancel: cancel,
|
driver: driver,
|
||||||
traceID: traceID,
|
autoArchive: autoArchive,
|
||||||
driver: driver,
|
state: &managerState{
|
||||||
stateCmdChan: make(chan stateCommand, 100), // Buffered channel for performance
|
spaces: make(map[string]*types.TraceSpace),
|
||||||
autoArchive: autoArchive,
|
traceStatus: types.TraceStatusPending,
|
||||||
pubsub: pubsubService, // Reference only, doesn't manage lifecycle
|
updates: make([]*types.TraceUpdate, 0, 100),
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
// Start state worker goroutine
|
// Load existing updates from driver (for resumed traces).
|
||||||
go m.startStateWorker()
|
// Safe to access m.state directly here — no Queue yet, single goroutine.
|
||||||
|
|
||||||
// Try to load existing updates from driver (for resumed traces)
|
|
||||||
if existingUpdates, err := driver.LoadUpdates(ctx, traceID, 0); err == nil && len(existingUpdates) > 0 {
|
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)
|
log.Trace("[MANAGER] NewManager: loaded %d existing updates from driver for trace %s", len(existingUpdates), traceID)
|
||||||
m.stateSetUpdates(existingUpdates)
|
m.state.updates = existingUpdates
|
||||||
// Check if trace was already completed
|
|
||||||
for _, update := range existingUpdates {
|
for _, update := range existingUpdates {
|
||||||
if update.Type == types.UpdateTypeComplete {
|
if update.Type == types.UpdateTypeComplete {
|
||||||
log.Trace("[MANAGER] NewManager: trace %s was already completed, marking as completed", traceID)
|
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 {
|
if data, ok := update.Data.(*types.TraceCompleteData); ok {
|
||||||
log.Trace("[MANAGER] NewManager: setting trace status to %s", data.Status)
|
log.Trace("[MANAGER] NewManager: setting trace status to %s", data.Status)
|
||||||
m.stateSetTraceStatus(data.Status)
|
m.state.traceStatus = data.Status
|
||||||
}
|
}
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
|
@ -70,7 +63,6 @@ func NewManager(ctx context.Context, traceID string, driver types.Driver, pubsub
|
||||||
} else {
|
} else {
|
||||||
log.Trace("[MANAGER] NewManager: no existing updates found for trace %s, creating new trace", traceID)
|
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()
|
now := time.Now().UnixMilli()
|
||||||
m.addUpdateAndBroadcast(&types.TraceUpdate{
|
m.addUpdateAndBroadcast(&types.TraceUpdate{
|
||||||
Type: types.UpdateTypeInit,
|
Type: types.UpdateTypeInit,
|
||||||
|
|
@ -89,36 +81,24 @@ func genNodeID() string {
|
||||||
return id
|
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) {
|
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 {
|
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)
|
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)
|
m.stateAddUpdate(update)
|
||||||
|
|
||||||
// Publish to independent PubSub service (manager just publishes, doesn't manage pubsub lifecycle)
|
// Broadcast to subscribers via event service (fire-and-forget, non-blocking).
|
||||||
if m.pubsub != nil {
|
// Uses Push with the update as payload so event.Subscribe filters can match by traceID.
|
||||||
m.pubsub.Publish(update)
|
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 {
|
func (m *manager) checkContext() error {
|
||||||
select {
|
return nil
|
||||||
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
|
// 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
|
// 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)
|
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
|
// Add to each parent's children
|
||||||
for _, parent := range currentNodes {
|
for _, parent := range currentNodes {
|
||||||
parent.Children = append(parent.Children, newNodeData)
|
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
|
// Log error but continue
|
||||||
m.Error("Failed to update parent node %s: %v", parent.ID, err)
|
m.Error("Failed to update parent node %s: %v", parent.ID, err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Save new node
|
// 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
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -310,7 +290,7 @@ func (m *manager) Parallel(parallelInputs []types.TraceParallelInput) ([]types.N
|
||||||
// Save all nodes in batch - collect errors
|
// Save all nodes in batch - collect errors
|
||||||
var saveErrors []error
|
var saveErrors []error
|
||||||
for _, data := range nodeData {
|
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))
|
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
|
// 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)
|
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,
|
NodeID: node.ID,
|
||||||
}
|
}
|
||||||
// Save log (ignore errors for non-critical logging)
|
// 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
|
// Broadcast log event
|
||||||
m.addUpdateAndBroadcast(&types.TraceUpdate{
|
m.addUpdateAndBroadcast(&types.TraceUpdate{
|
||||||
|
|
@ -404,7 +384,7 @@ func (m *manager) SetOutput(output types.TraceOutput) error {
|
||||||
for _, node := range nodes {
|
for _, node := range nodes {
|
||||||
node.Output = output
|
node.Output = output
|
||||||
node.UpdatedAt = 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
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -434,7 +414,7 @@ func (m *manager) SetMetadata(key string, value any) error {
|
||||||
}
|
}
|
||||||
node.Metadata[key] = value
|
node.Metadata[key] = value
|
||||||
node.UpdatedAt = 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
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -485,7 +465,7 @@ func (m *manager) Complete(output ...types.TraceOutput) error {
|
||||||
node.Status = types.StatusCompleted
|
node.Status = types.StatusCompleted
|
||||||
node.EndTime = now
|
node.EndTime = now
|
||||||
node.UpdatedAt = 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
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -516,7 +496,7 @@ func (m *manager) Fail(err error) error {
|
||||||
node.Status = types.StatusFailed
|
node.Status = types.StatusFailed
|
||||||
node.EndTime = now
|
node.EndTime = now
|
||||||
node.UpdatedAt = 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
|
return saveErr
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -545,7 +525,7 @@ func (m *manager) GetRootNode() (*types.TraceNode, error) {
|
||||||
|
|
||||||
// GetNode returns a node by ID
|
// GetNode returns a node by ID
|
||||||
func (m *manager) GetNode(id string) (*types.TraceNode, error) {
|
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
|
// GetCurrentNodes returns current active nodes
|
||||||
|
|
@ -581,7 +561,7 @@ func (m *manager) MarkComplete() error {
|
||||||
|
|
||||||
// Auto-archive if enabled
|
// Auto-archive if enabled
|
||||||
if m.autoArchive {
|
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
|
// Log error but don't fail the complete operation
|
||||||
m.Debug("Failed to auto-archive trace", map[string]any{
|
m.Debug("Failed to auto-archive trace", map[string]any{
|
||||||
"trace_id": m.traceID,
|
"trace_id": m.traceID,
|
||||||
|
|
@ -610,7 +590,7 @@ func (m *manager) CreateSpace(option types.TraceSpaceOption) (*types.TraceSpace,
|
||||||
}
|
}
|
||||||
|
|
||||||
// Save to driver
|
// 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
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -637,7 +617,7 @@ func (m *manager) GetSpace(id string) (*types.TraceSpace, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Load from driver
|
// 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 {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
@ -658,7 +638,7 @@ func (m *manager) HasSpace(id string) bool {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check in driver
|
// Check in driver
|
||||||
space, _ := m.driver.LoadSpace(m.ctx, m.traceID, id)
|
space, _ := m.driver.LoadSpace(context.Background(), m.traceID, id)
|
||||||
return space != nil
|
return space != nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -674,7 +654,7 @@ func (m *manager) DeleteSpace(id string) error {
|
||||||
m.stateDeleteSpace(id)
|
m.stateDeleteSpace(id)
|
||||||
|
|
||||||
// Delete from driver
|
// 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
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -693,7 +673,7 @@ func (m *manager) DeleteSpace(id string) error {
|
||||||
// ListSpaces returns all spaces
|
// ListSpaces returns all spaces
|
||||||
func (m *manager) ListSpaces() []*types.TraceSpace {
|
func (m *manager) ListSpaces() []*types.TraceSpace {
|
||||||
// Load from driver to ensure we have all spaces
|
// 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 {
|
if err != nil {
|
||||||
// Fallback to cached spaces
|
// Fallback to cached spaces
|
||||||
return m.stateGetAllSpaces()
|
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)
|
// Set value in driver (through state worker for concurrent safety)
|
||||||
err = m.stateExecuteSpaceOp(spaceID, func() error {
|
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
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
// Update space timestamp
|
// Update space timestamp
|
||||||
space.UpdatedAt = now
|
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
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -761,7 +741,7 @@ func (m *manager) GetSpaceValue(spaceID, key string) (any, error) {
|
||||||
var result any
|
var result any
|
||||||
err := m.stateExecuteSpaceOp(spaceID, func() error {
|
err := m.stateExecuteSpaceOp(spaceID, func() error {
|
||||||
var err 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 err
|
||||||
})
|
})
|
||||||
return result, 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 {
|
func (m *manager) HasSpaceValue(spaceID, key string) bool {
|
||||||
var result bool
|
var result bool
|
||||||
_ = m.stateExecuteSpaceOp(spaceID, func() error {
|
_ = 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 nil
|
||||||
})
|
})
|
||||||
return result
|
return result
|
||||||
|
|
@ -787,7 +767,7 @@ func (m *manager) DeleteSpaceValue(spaceID, key string) error {
|
||||||
|
|
||||||
// Delete value from driver (through state worker for concurrent safety)
|
// Delete value from driver (through state worker for concurrent safety)
|
||||||
err := m.stateExecuteSpaceOp(spaceID, func() error {
|
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 {
|
if err != nil {
|
||||||
|
|
@ -816,7 +796,7 @@ func (m *manager) ClearSpaceValues(spaceID string) error {
|
||||||
|
|
||||||
// Clear values from driver (through state worker for concurrent safety)
|
// Clear values from driver (through state worker for concurrent safety)
|
||||||
err := m.stateExecuteSpaceOp(spaceID, func() error {
|
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 {
|
if err != nil {
|
||||||
|
|
@ -840,7 +820,7 @@ func (m *manager) ListSpaceKeys(spaceID string) []string {
|
||||||
var keys []string
|
var keys []string
|
||||||
_ = m.stateExecuteSpaceOp(spaceID, func() error {
|
_ = m.stateExecuteSpaceOp(spaceID, func() error {
|
||||||
var err 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 err
|
||||||
})
|
})
|
||||||
return keys
|
return keys
|
||||||
|
|
@ -865,7 +845,7 @@ func (m *manager) GetTraceInfo() (*types.TraceInfo, error) {
|
||||||
if err := m.checkContext(); err != nil {
|
if err := m.checkContext(); err != nil {
|
||||||
return nil, err
|
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
|
// GetAllNodes retrieves all nodes from storage
|
||||||
|
|
@ -875,7 +855,7 @@ func (m *manager) GetAllNodes() ([]*types.TraceNode, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Load the root node tree from storage
|
// 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 {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
@ -906,7 +886,7 @@ func (m *manager) GetNodeByID(nodeID string) (*types.TraceNode, error) {
|
||||||
if err := m.checkContext(); err != nil {
|
if err := m.checkContext(); err != nil {
|
||||||
return nil, err
|
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
|
// GetAllLogs retrieves all logs from storage
|
||||||
|
|
@ -914,7 +894,7 @@ func (m *manager) GetAllLogs() ([]*types.TraceLog, error) {
|
||||||
if err := m.checkContext(); err != nil {
|
if err := m.checkContext(); err != nil {
|
||||||
return nil, err
|
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
|
// 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 {
|
if err := m.checkContext(); err != nil {
|
||||||
return nil, err
|
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
|
// GetAllSpaces retrieves all spaces from storage
|
||||||
|
|
@ -932,7 +912,7 @@ func (m *manager) GetAllSpaces() ([]*types.TraceSpace, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Get all space IDs from driver
|
// 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 {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
@ -940,7 +920,7 @@ func (m *manager) GetAllSpaces() ([]*types.TraceSpace, error) {
|
||||||
// Load all spaces
|
// Load all spaces
|
||||||
spaces := make([]*types.TraceSpace, 0, len(spaceIDs))
|
spaces := make([]*types.TraceSpace, 0, len(spaceIDs))
|
||||||
for _, spaceID := range 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 {
|
if err != nil {
|
||||||
continue // Skip spaces that fail to load
|
continue // Skip spaces that fail to load
|
||||||
}
|
}
|
||||||
|
|
@ -959,7 +939,7 @@ func (m *manager) GetSpaceByID(spaceID string) (*types.TraceSpaceData, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Load space metadata
|
// 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 {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
@ -968,7 +948,7 @@ func (m *manager) GetSpaceByID(spaceID string) (*types.TraceSpaceData, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Load all keys in the space
|
// 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 {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
@ -976,7 +956,7 @@ func (m *manager) GetSpaceByID(spaceID string) (*types.TraceSpaceData, error) {
|
||||||
// Load all key-value pairs
|
// Load all key-value pairs
|
||||||
data := make(map[string]any)
|
data := make(map[string]any)
|
||||||
for _, key := range keys {
|
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 {
|
if err != nil {
|
||||||
continue // Skip keys that fail to load
|
continue // Skip keys that fail to load
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
package trace
|
package trace
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/yaoapp/yao/trace/types"
|
"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,
|
NodeID: n.data.ID,
|
||||||
}
|
}
|
||||||
// Save log (ignore errors for non-critical logging)
|
// 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
|
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)
|
n.data.Children = append(n.data.Children, childNodeData)
|
||||||
|
|
||||||
// Save both nodes
|
// 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
|
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
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -120,7 +121,7 @@ func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node
|
||||||
n.data.Children = append(n.data.Children, childNodeData)
|
n.data.Children = append(n.data.Children, childNodeData)
|
||||||
|
|
||||||
// Save node
|
// 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
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -132,7 +133,7 @@ func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node
|
||||||
}
|
}
|
||||||
|
|
||||||
// Save parent 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
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -165,7 +166,7 @@ func (n *node) Join(nodes []*types.TraceNode, input types.TraceInput, option typ
|
||||||
}
|
}
|
||||||
|
|
||||||
// Save join node
|
// 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
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -185,7 +186,7 @@ func (n *node) ID() string {
|
||||||
func (n *node) SetOutput(output types.TraceOutput) error {
|
func (n *node) SetOutput(output types.TraceOutput) error {
|
||||||
n.data.Output = output
|
n.data.Output = output
|
||||||
n.data.UpdatedAt = time.Now().UnixMilli()
|
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
|
// SetMetadata sets node metadata
|
||||||
|
|
@ -195,14 +196,14 @@ func (n *node) SetMetadata(key string, value any) error {
|
||||||
}
|
}
|
||||||
n.data.Metadata[key] = value
|
n.data.Metadata[key] = value
|
||||||
n.data.UpdatedAt = time.Now().UnixMilli()
|
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
|
// SetStatus sets the node status
|
||||||
func (n *node) SetStatus(status string) error {
|
func (n *node) SetStatus(status string) error {
|
||||||
n.data.Status = types.NodeStatus(status)
|
n.data.Status = types.NodeStatus(status)
|
||||||
n.data.UpdatedAt = time.Now().UnixMilli()
|
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)
|
// 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.Status = types.StatusCompleted
|
||||||
n.data.EndTime = now
|
n.data.EndTime = now
|
||||||
n.data.UpdatedAt = 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)
|
// 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.EndTime = now
|
||||||
n.data.UpdatedAt = 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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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()
|
|
||||||
}
|
|
||||||
|
|
@ -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
|
|
||||||
}
|
|
||||||
|
|
@ -6,18 +6,18 @@ import (
|
||||||
"github.com/yaoapp/yao/trace/types"
|
"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 {
|
type space struct {
|
||||||
ctx context.Context
|
|
||||||
traceID string
|
traceID string
|
||||||
data *types.TraceSpace
|
data *types.TraceSpace
|
||||||
driver types.Driver
|
driver types.Driver
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewSpace creates a new space instance
|
// 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{
|
return &space{
|
||||||
ctx: ctx,
|
|
||||||
traceID: traceID,
|
traceID: traceID,
|
||||||
data: data,
|
data: data,
|
||||||
driver: driver,
|
driver: driver,
|
||||||
|
|
@ -31,32 +31,32 @@ func (s *space) ID() string {
|
||||||
|
|
||||||
// Set stores a value by key
|
// Set stores a value by key
|
||||||
func (s *space) Set(key string, value any) error {
|
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
|
// Get retrieves a value by key
|
||||||
func (s *space) Get(key string) (any, error) {
|
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
|
// Has checks if a key exists
|
||||||
func (s *space) Has(key string) bool {
|
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
|
// Delete removes a key-value pair
|
||||||
func (s *space) Delete(key string) error {
|
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
|
// Clear removes all key-value pairs
|
||||||
func (s *space) Clear() error {
|
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
|
// Keys returns all keys in the space
|
||||||
func (s *space) Keys() []string {
|
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 {
|
if err != nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
|
||||||
362
trace/state.go
362
trace/state.go
|
|
@ -2,16 +2,13 @@ package trace
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
"sync/atomic"
|
|
||||||
|
|
||||||
"github.com/yaoapp/kun/log"
|
"github.com/yaoapp/kun/log"
|
||||||
"github.com/yaoapp/yao/trace/types"
|
"github.com/yaoapp/yao/trace/types"
|
||||||
)
|
)
|
||||||
|
|
||||||
// State management using channel-based serialization (no locks needed)
|
// managerState holds all mutable state for a trace.
|
||||||
// All state mutations go through a single worker goroutine
|
// Protected by manager.mu — all access goes through state* methods which acquire the lock.
|
||||||
|
|
||||||
// managerState holds all mutable state (accessed only by state worker)
|
|
||||||
type managerState struct {
|
type managerState struct {
|
||||||
rootNode *types.TraceNode
|
rootNode *types.TraceNode
|
||||||
currentNodes []*types.TraceNode
|
currentNodes []*types.TraceNode
|
||||||
|
|
@ -19,344 +16,133 @@ type managerState struct {
|
||||||
traceStatus types.TraceStatus
|
traceStatus types.TraceStatus
|
||||||
completed bool
|
completed bool
|
||||||
updates []*types.TraceUpdate
|
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) {
|
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 {
|
func (m *manager) stateGetRoot() *types.TraceNode {
|
||||||
resp := make(chan *types.TraceNode, 1)
|
m.mu.Lock()
|
||||||
if !m.safeSend(&cmdGetRoot{resp: resp}) {
|
defer m.mu.Unlock()
|
||||||
return nil // Context cancelled
|
return m.state.rootNode
|
||||||
}
|
|
||||||
return <-resp
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *manager) stateSetCurrentNodes(nodes []*types.TraceNode) {
|
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 {
|
func (m *manager) stateGetCurrentNodes() []*types.TraceNode {
|
||||||
resp := make(chan []*types.TraceNode, 1)
|
m.mu.Lock()
|
||||||
if !m.safeSend(&cmdGetCurrentNodes{resp: resp}) {
|
defer m.mu.Unlock()
|
||||||
return nil // Context cancelled
|
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) {
|
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) {
|
func (m *manager) stateGetSpace(id string) (*types.TraceSpace, bool) {
|
||||||
resp := make(chan *types.TraceSpace, 1)
|
m.mu.Lock()
|
||||||
if !m.safeSend(&cmdGetSpace{id: id, resp: resp}) {
|
defer m.mu.Unlock()
|
||||||
return nil, false // Context cancelled
|
space, ok := m.state.spaces[id]
|
||||||
}
|
return space, ok
|
||||||
space := <-resp
|
|
||||||
return space, space != nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *manager) stateSetSpace(id string, space *types.TraceSpace) {
|
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) {
|
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 {
|
func (m *manager) stateGetAllSpaces() []*types.TraceSpace {
|
||||||
resp := make(chan []*types.TraceSpace, 1)
|
m.mu.Lock()
|
||||||
if !m.safeSend(&cmdGetAllSpaces{resp: resp}) {
|
defer m.mu.Unlock()
|
||||||
return nil // Context cancelled
|
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) {
|
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 {
|
func (m *manager) stateGetTraceStatus() types.TraceStatus {
|
||||||
resp := make(chan types.TraceStatus, 1)
|
m.mu.Lock()
|
||||||
if !m.safeSend(&cmdGetTraceStatus{resp: resp}) {
|
defer m.mu.Unlock()
|
||||||
return types.TraceStatusCancelled // Context cancelled
|
return m.state.traceStatus
|
||||||
}
|
|
||||||
return <-resp
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *manager) stateMarkCompleted() bool {
|
func (m *manager) stateMarkCompleted() bool {
|
||||||
resp := make(chan bool, 1)
|
m.mu.Lock()
|
||||||
if !m.safeSend(&cmdMarkCompleted{resp: resp}) {
|
defer m.mu.Unlock()
|
||||||
return true // Context cancelled, treat as completed
|
if m.state.completed {
|
||||||
|
return false
|
||||||
}
|
}
|
||||||
return <-resp
|
m.state.completed = true
|
||||||
|
return true
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *manager) stateIsCompleted() bool {
|
func (m *manager) stateIsCompleted() bool {
|
||||||
resp := make(chan bool, 1)
|
m.mu.Lock()
|
||||||
if !m.safeSend(&cmdIsCompleted{resp: resp}) {
|
defer m.mu.Unlock()
|
||||||
return true // Context cancelled, treat as completed
|
return m.state.completed
|
||||||
}
|
|
||||||
return <-resp
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *manager) stateAddUpdate(update *types.TraceUpdate) {
|
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 {
|
func (m *manager) stateGetUpdates(since int64) []*types.TraceUpdate {
|
||||||
resp := make(chan []*types.TraceUpdate, 1)
|
m.mu.Lock()
|
||||||
if !m.safeSend(&cmdGetUpdates{since: since, resp: resp}) {
|
defer m.mu.Unlock()
|
||||||
return nil // Context cancelled
|
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) {
|
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)
|
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
|
// stateExecuteSpaceOp executes a space operation while holding the lock.
|
||||||
// See subscription_manager.go and subscription.go for the new implementation
|
|
||||||
|
|
||||||
// stateExecuteSpaceOp executes a space operation serially through state worker
|
|
||||||
func (m *manager) stateExecuteSpaceOp(spaceID string, fn func() error) error {
|
func (m *manager) stateExecuteSpaceOp(spaceID string, fn func() error) error {
|
||||||
resp := make(chan error, 1)
|
m.mu.Lock()
|
||||||
if !m.safeSend(&cmdSpaceKVOp{spaceID: spaceID, fn: fn, resp: resp}) {
|
defer m.mu.Unlock()
|
||||||
return fmt.Errorf("trace %s: state worker stopped", m.traceID)
|
err := fn()
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("trace %s: space op failed: %w", m.traceID, err)
|
||||||
}
|
}
|
||||||
return <-resp
|
return nil
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -3,12 +3,18 @@ package trace
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/yaoapp/yao/event"
|
||||||
|
eventTypes "github.com/yaoapp/yao/event/types"
|
||||||
"github.com/yaoapp/yao/trace/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)
|
// Subscribe creates a new subscription for trace updates (replays all historical events from the beginning)
|
||||||
func (m *manager) Subscribe() (<-chan *types.TraceUpdate, error) {
|
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
|
// 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)
|
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) {
|
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
|
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
105
trace/trace.go
|
|
@ -4,13 +4,11 @@ import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"sync"
|
"sync"
|
||||||
"sync/atomic"
|
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
gonanoid "github.com/matoous/go-nanoid/v2"
|
gonanoid "github.com/matoous/go-nanoid/v2"
|
||||||
"github.com/yaoapp/kun/log"
|
"github.com/yaoapp/kun/log"
|
||||||
"github.com/yaoapp/yao/trace/local"
|
"github.com/yaoapp/yao/trace/local"
|
||||||
"github.com/yaoapp/yao/trace/pubsub"
|
|
||||||
"github.com/yaoapp/yao/trace/store"
|
"github.com/yaoapp/yao/trace/store"
|
||||||
"github.com/yaoapp/yao/trace/types"
|
"github.com/yaoapp/yao/trace/types"
|
||||||
)
|
)
|
||||||
|
|
@ -21,14 +19,10 @@ const (
|
||||||
Store = "store" // Gou store storage
|
Store = "store" // Gou store storage
|
||||||
)
|
)
|
||||||
|
|
||||||
// Global trace registry and pubsub services
|
// Global trace registry
|
||||||
var (
|
var (
|
||||||
registry = make(map[string]*types.TraceInfo)
|
registry = make(map[string]*types.TraceInfo)
|
||||||
registryMu sync.RWMutex
|
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
|
// 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...)
|
return LoadFromStorage(ctx, driver, traceID, driverOptions...)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Create independent PubSub service for this trace
|
// Create Manager instance with the driver
|
||||||
pubsubService := pubsub.New()
|
manager, err := NewManager(ctx, traceID, drv, option)
|
||||||
|
|
||||||
// Register pubsub service
|
|
||||||
pubsubRegistryMu.Lock()
|
|
||||||
pubsubRegistry[traceID] = pubsubService
|
|
||||||
pubsubRegistryMu.Unlock()
|
|
||||||
|
|
||||||
// Create Manager instance with the driver and pubsub reference
|
|
||||||
// Manager uses pubsub only for publishing, doesn't manage its lifecycle
|
|
||||||
manager, err := NewManager(ctx, traceID, drv, pubsubService, option)
|
|
||||||
if err != nil {
|
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)
|
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
|
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
|
// Load loads an existing trace by ID from the registry
|
||||||
// Returns: manager, error
|
// Returns: manager, error
|
||||||
// traceID: the trace ID to load
|
// 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)
|
return "", nil, fmt.Errorf("trace not found in storage: %s", traceID)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Create or reuse PubSub service for this trace
|
manager, err := NewManager(ctx, traceID, drv, nil)
|
||||||
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)
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
drv.Close()
|
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)
|
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")
|
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()
|
now := time.Now().UnixMilli()
|
||||||
|
|
||||||
// Use background context since the original context is cancelled
|
// 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)
|
log.Trace("[TRACE] MarkCancelled: failed to save node %s: %v", node.ID, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Broadcast node failed event (also saves to disk)
|
log.Trace("[TRACE] MarkCancelled: broadcasting node failed event for node %s", node.ID)
|
||||||
subscriberCount := 0
|
|
||||||
if ps := GetPubSub(traceID); ps != nil {
|
|
||||||
subscriberCount = ps.SubscriberCount()
|
|
||||||
}
|
|
||||||
log.Trace("[TRACE] MarkCancelled: publishing node failed event for node %s to %d subscribers", node.ID, subscriberCount)
|
|
||||||
mgr.addUpdateAndBroadcast(&types.TraceUpdate{
|
mgr.addUpdateAndBroadcast(&types.TraceUpdate{
|
||||||
Type: types.UpdateTypeNodeFailed,
|
Type: types.UpdateTypeNodeFailed,
|
||||||
TraceID: traceID,
|
TraceID: traceID,
|
||||||
|
|
@ -445,12 +387,7 @@ func MarkCancelled(traceID string, reason string) error {
|
||||||
mgr.stateSetTraceStatus(types.TraceStatusCancelled)
|
mgr.stateSetTraceStatus(types.TraceStatusCancelled)
|
||||||
mgr.stateMarkCompleted()
|
mgr.stateMarkCompleted()
|
||||||
|
|
||||||
// Broadcast completion update (saves to disk and publishes to subscribers)
|
log.Trace("[TRACE] MarkCancelled: broadcasting completion update")
|
||||||
subscriberCount := 0
|
|
||||||
if ps := GetPubSub(traceID); ps != nil {
|
|
||||||
subscriberCount = ps.SubscriberCount()
|
|
||||||
}
|
|
||||||
log.Trace("[TRACE] MarkCancelled: publishing completion update to %d subscribers", subscriberCount)
|
|
||||||
totalDuration := int64(0)
|
totalDuration := int64(0)
|
||||||
if rootNode.CreatedAt > 0 {
|
if rootNode.CreatedAt > 0 {
|
||||||
totalDuration = now - rootNode.CreatedAt
|
totalDuration = now - rootNode.CreatedAt
|
||||||
|
|
@ -488,35 +425,7 @@ func Release(traceID string) error {
|
||||||
return fmt.Errorf("trace not found in registry: %s", traceID)
|
return fmt.Errorf("trace not found in registry: %s", traceID)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Stop manager with safe three-step shutdown sequence.
|
_ = info.Manager
|
||||||
// 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()
|
|
||||||
}
|
|
||||||
|
|
||||||
log.Trace("[TRACE] Release: completed")
|
log.Trace("[TRACE] Release: completed")
|
||||||
return nil
|
return nil
|
||||||
|
|
|
||||||
|
|
@ -13,11 +13,8 @@ import (
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestMain(m *testing.M) {
|
func TestMain(m *testing.M) {
|
||||||
// Prepare test environment (initializes stores, models, etc.)
|
|
||||||
test.Prepare(&testing.T{}, config.Conf)
|
test.Prepare(&testing.T{}, config.Conf)
|
||||||
defer test.Clean()
|
defer test.Clean()
|
||||||
|
|
||||||
// Run tests
|
|
||||||
os.Exit(m.Run())
|
os.Exit(m.Run())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -234,15 +231,17 @@ func TestContextCancellation(t *testing.T) {
|
||||||
defer trace.Release(traceID)
|
defer trace.Release(traceID)
|
||||||
defer trace.Remove(context.Background(), d.DriverType, traceID, d.DriverOptions...)
|
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()
|
cancel()
|
||||||
|
|
||||||
// Operations should fail with context error
|
node, err := manager.Add("test", types.TraceNodeOption{
|
||||||
_, err = manager.Add("test", types.TraceNodeOption{
|
|
||||||
Label: "Test",
|
Label: "Test",
|
||||||
Type: "test",
|
Type: "test",
|
||||||
})
|
})
|
||||||
assert.Error(t, err)
|
assert.NoError(t, err)
|
||||||
|
assert.NotNil(t, node)
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -186,9 +186,10 @@ func TestConcurrentReleaseAndMarkCancelled(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// TestSafeSendAfterClosed verifies that operations using safeSend after
|
// TestOperationsAfterRelease verifies that using a manager reference after
|
||||||
// Release return gracefully instead of panicking.
|
// Release does not panic. The manager is removed from registry but its
|
||||||
func TestSafeSendAfterClosed(t *testing.T) {
|
// in-memory state remains valid (no channel close or context cancel).
|
||||||
|
func TestOperationsAfterRelease(t *testing.T) {
|
||||||
drivers := trace.GetTestDrivers()
|
drivers := trace.GetTestDrivers()
|
||||||
|
|
||||||
for _, d := range drivers {
|
for _, d := range drivers {
|
||||||
|
|
@ -205,21 +206,13 @@ func TestSafeSendAfterClosed(t *testing.T) {
|
||||||
err = trace.Release(traceID)
|
err = trace.Release(traceID)
|
||||||
assert.NoError(t, err)
|
assert.NoError(t, err)
|
||||||
|
|
||||||
// All of these internally use safeSend. After Release they should
|
// After Release, the manager object is still usable (state in memory).
|
||||||
// return nil/error/zero-value, never panic.
|
// These calls should not panic.
|
||||||
manager.Info("post-close info")
|
manager.Info("post-release info")
|
||||||
manager.Debug("post-close debug")
|
manager.Debug("post-release debug")
|
||||||
manager.Error("post-close error")
|
|
||||||
manager.Warn("post-close warn")
|
|
||||||
|
|
||||||
root, _ := manager.GetRootNode()
|
root, _ := manager.GetRootNode()
|
||||||
assert.Nil(t, root)
|
assert.NotNil(t, root)
|
||||||
|
|
||||||
nodes, _ := manager.GetCurrentNodes()
|
|
||||||
assert.Nil(t, nodes)
|
|
||||||
|
|
||||||
status := manager.IsComplete()
|
|
||||||
assert.True(t, status)
|
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue