Enhance trace logging and error handling across components

- Integrated trace management into the Assistant's Stream method, allowing for detailed logging of processing steps and errors.
- Improved error handling in the OpenAI stream and post methods, adding trace logs for retries and failures.
- Enhanced CUI and OpenAI writers to log message adaptation and sending errors, improving debugging capabilities.
- Updated trace manager to use milliseconds for timestamps, ensuring consistency across the trace system.
- Refactored log methods in the trace node to streamline logging and broadcasting of events.

agent/assistant/agent.go

@@ -7,6 +7,7 @@ import (
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/llm"
 	"github.com/yaoapp/yao/agent/output"
+	"github.com/yaoapp/yao/trace/types"
 	"github.com/yaoapp/yao/utils/jsonschema"
 )
 
@@ -29,20 +30,51 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 
 	_ = traceID // traceID is available for trace logging
 
+	// Trace Add
+	trace, _ := ctx.Trace()
+	var agentNode types.Node = nil
+	if trace != nil {
+		agentNode, _ = trace.Add(inputMessages, types.TraceNodeOption{
+			Label:       fmt.Sprintf("Assistant %s", ast.Name),
+			Icon:        "assistant",
+			Description: fmt.Sprintf("Assistant %s is processing the request", ast.Name),
+		})
+		if agentNode != nil {
+			// Mark the node as complete when the function returns
+			defer agentNode.Complete()
+		}
+	}
+
 	// Full input messages with chat history
 	fullMessages, err := ast.WithHistory(ctx, inputMessages)
 	if err != nil {
+		if agentNode != nil {
+			agentNode.Fail(err)
+		}
 		return nil, err
 	}
 
+	// Log the chat history
+	if agentNode != nil {
+		agentNode.Info("Get Chat History", map[string]any{"messages": fullMessages})
+	}
+
 	// Request Create hook ( Optional )
 	var createResponse *context.HookCreateResponse
 	if ast.Script != nil {
 		var err error
 		createResponse, err = ast.Script.Create(ctx, fullMessages)
 		if err != nil {
+			if agentNode != nil {
+				agentNode.Fail(err)
+			}
 			return nil, err
 		}
+
+		// Log the create response
+		if agentNode != nil {
+			agentNode.Debug("Call Create Hook", map[string]any{"response": createResponse})
+		}
 	}
 
 	var completionOptions *context.CompletionOptions // default is nil
@@ -54,12 +86,18 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 		// Build the LLM request first
 		completionMessages, completionOptions, err = ast.BuildRequest(ctx, inputMessages, createResponse)
 		if err != nil {
+			if agentNode != nil {
+				agentNode.Fail(err)
+			}
 			return nil, err
 		}
 
 		// Get connector object and capabilities
 		conn, capabilities, err := ast.GetConnector(ctx)
 		if err != nil {
+			if agentNode != nil {
+				agentNode.Fail(err)
+			}
 			return nil, err
 		}
 
@@ -68,6 +106,22 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 			completionOptions.Capabilities = capabilities
 		}
 
+		// Log the capabilities
+		if agentNode != nil {
+			agentNode.Debug("Get Connector Capabilities", map[string]any{"capabilities": capabilities})
+		}
+
+		// Trace Add
+		if trace != nil {
+			trace.Add(
+				map[string]any{"messages": completionMessages, "options": completionOptions},
+				types.TraceNodeOption{
+					Label: fmt.Sprintf("LLM %s", conn.ID()), Icon: "llm", Description: fmt.Sprintf("LLM %s is processing the request", conn.ID()),
+				},
+			)
+
+		}
+
 		// Create LLM instance with connector and options
 		llmInstance, err := llm.New(conn, completionOptions)
 		if err != nil {
@@ -86,6 +140,10 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 			return nil, err
 		}
 
+		// Mark LLM Request Complete
+		if trace != nil {
+			trace.Complete(completionResponse)
+		}
 	}
 
 	// Request MCP hook ( Optional )

agent/llm/handlers/stream.go

@@ -1,8 +1,6 @@
 package handlers
 
 import (
-	"fmt"
-
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/output"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -20,10 +18,11 @@ func DefaultStreamHandler(ctx *context.Context) context.StreamFunc {
 	}
 
 	return func(chunkType context.StreamChunkType, data []byte) int {
-		fmt.Println("-----------------------------------------------")
+		trace, _ := ctx.Trace()
-		fmt.Println("Chunk Type: ", string(chunkType))
+		if trace != nil {
-		fmt.Println("Data: ", string(data))
+			trace.Info("LLM Stream", map[string]any{"data": string(data)})
-		fmt.Println("-----------------------------------------------")
+		}
+
 		// Handle different chunk types
 		switch chunkType {
 		case context.ChunkStreamStart:

agent/llm/providers/openai/openai.go

@@ -9,7 +9,6 @@ import (
 	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/http"
-	"github.com/yaoapp/kun/log"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/llm/adapters"
 	"github.com/yaoapp/yao/agent/llm/providers/base"
@@ -186,6 +185,14 @@ func detectReasoningFormat(cap *context.ModelCapabilities) adapters.ReasoningFor
 
 // Stream stream completion from OpenAI API
 func (p *Provider) Stream(ctx *context.Context, messages []context.Message, options *context.CompletionOptions, handler context.StreamFunc) (*context.CompletionResponse, error) {
+	// Add debug log
+	trace, _ := ctx.Trace()
+	if trace != nil {
+		trace.Debug("OpenAI Stream: Starting stream request", map[string]any{
+			"message_count": len(messages),
+		})
+	}
+
 	maxRetries := 3
 	maxValidationRetries := 3
 	var lastErr error
@@ -219,7 +226,16 @@ func (p *Provider) Stream(ctx *context.Context, messages []context.Message, opti
 		if attempt > 0 {
 			// Exponential backoff: 1s, 2s, 4s
 			backoff := time.Duration(1<<uint(attempt-1)) * time.Second
-			log.Warn("OpenAI stream request failed, retrying in %v (attempt %d/%d): %v", backoff, attempt+1, maxRetries, lastErr)
+
+			// Add debug log to trace
+			if trace != nil {
+				trace.Warn("OpenAI stream request failed, retrying", map[string]any{
+					"backoff":     backoff.String(),
+					"attempt":     attempt + 1,
+					"max_retries": maxRetries,
+					"error":       lastErr.Error(),
+				})
+			}
 
 			// Sleep with context cancellation support
 			timer := time.NewTimer(backoff)
@@ -249,16 +265,31 @@ func (p *Provider) Stream(ctx *context.Context, messages []context.Message, opti
 
 		response, err := p.streamWithRetry(ctx, currentMessages, options, handler)
 		if err == nil {
+			if trace != nil {
+				trace.Debug("OpenAI Stream: Request completed successfully")
+			}
 			return response, nil
 		}
 		lastErr = err
 
+		if trace != nil {
+			trace.Debug("OpenAI Stream: Request failed", map[string]any{
+				"error": err.Error(),
+			})
+		}
+
 		// Check if error is tool call validation failure
 		if isToolCallValidationError(err) {
 			// Handle tool call validation retry with feedback to LLM
 			validationRetryMessages := currentMessages
 			for validationAttempt := 0; validationAttempt < maxValidationRetries; validationAttempt++ {
-				log.Warn("Tool call validation failed (attempt %d/%d): %v", validationAttempt+1, maxValidationRetries, err)
+				if trace != nil {
+					trace.Warn("Tool call validation failed", map[string]any{
+						"attempt":     validationAttempt + 1,
+						"max_retries": maxValidationRetries,
+						"error":       err.Error(),
+					})
+				}
 
 				// Add error feedback to conversation history
 				validationRetryMessages = append(validationRetryMessages, context.Message{
@@ -301,6 +332,14 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 	streamStartTime := time.Now()
 	requestID := fmt.Sprintf("req_%d", streamStartTime.UnixNano())
 
+	// Add debug log
+	trace, _ := ctx.Trace()
+	if trace != nil {
+		trace.Debug("OpenAI Stream: streamWithRetry starting", map[string]any{
+			"request_id": requestID,
+		})
+	}
+
 	// Get Go context for cancellation support
 	goCtx := ctx.Context
 	if goCtx == nil {
@@ -394,6 +433,12 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 	// Build URL
 	url := buildAPIURL(host, "/chat/completions")
 
+	if trace != nil {
+		trace.Debug("OpenAI Stream: Sending request", map[string]any{
+			"url": url,
+		})
+	}
+
 	// Create HTTP request with proxy support
 	req := http.New(url).
 		SetHeader("Content-Type", "application/json").
@@ -413,7 +458,9 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 		// Check for context cancellation
 		select {
 		case <-goCtx.Done():
-			log.Warn("Stream cancelled by context")
+			if trace != nil {
+				trace.Warn("Stream cancelled by context")
+			}
 			return http.HandlerReturnBreak
 		default:
 		}
@@ -421,7 +468,9 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 		// Check for force interrupt signal
 		if ctx.Interrupt != nil {
 			if signal := ctx.Interrupt.Peek(); signal != nil && signal.Type == context.InterruptForce {
-				log.Warn("Stream cancelled by force interrupt")
+				if trace != nil {
+					trace.Warn("Stream cancelled by force interrupt")
+				}
 				return http.HandlerReturnBreak
 			}
 		}
@@ -430,13 +479,16 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 			return http.HandlerReturnOk
 		}
 
-		// Log raw stream data for debugging
+		// Record LLM raw output to trace
-		log.Trace("OpenAI Stream Raw Data: %s", string(data))
+		if trace != nil {
+			trace.Debug("LLM Raw Output", map[string]any{
+				"data": string(data),
+			})
+		}
 
 		// Parse SSE data
 		dataStr := string(data)
 		if !strings.HasPrefix(dataStr, "data: ") {
-			log.Trace("Skipping non-SSE line: %s", dataStr)
 			return http.HandlerReturnOk
 		}
 
@@ -451,7 +503,11 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 		// Parse JSON chunk
 		var chunk StreamChunk
 		if err := jsoniter.UnmarshalFromString(dataStr, &chunk); err != nil {
-			log.Warn("Failed to parse stream chunk: %v", err)
+			if trace != nil {
+				trace.Warn("Failed to parse stream chunk", map[string]any{
+					"error": err.Error(),
+				})
+			}
 			return http.HandlerReturnOk
 		}
 
@@ -590,8 +646,13 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 	}
 
 	// Log request for debugging
-	if requestBodyJSON, marshalErr := jsoniter.Marshal(requestBody); marshalErr == nil {
+	if trace != nil {
-		log.Debug("OpenAI Stream Request - URL: %s, Body: %s", url, string(requestBodyJSON))
+		if requestBodyJSON, marshalErr := jsoniter.Marshal(requestBody); marshalErr == nil {
+			trace.Debug("OpenAI Stream Request", map[string]any{
+				"url":  url,
+				"body": string(requestBodyJSON),
+			})
+		}
 	}
 
 	// Buffer to capture non-SSE error responses
@@ -624,7 +685,11 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 	// Check if we captured an error response
 	if errorDetected && errorBuffer.Len() > 0 {
 		errorJSON := errorBuffer.String()
-		log.Error("OpenAI API returned error response: %s", errorJSON)
+		if trace != nil {
+			trace.Error("OpenAI API returned error response", map[string]any{
+				"response": errorJSON,
+			})
+		}
 
 		// Try to parse error
 		var apiError struct {
@@ -645,8 +710,10 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 	}
 
 	// Log any error from streaming
-	if err != nil {
+	if err != nil && trace != nil {
-		log.Error("OpenAI Stream Error: %v", err)
+		trace.Error("OpenAI Stream Error", map[string]any{
+			"error": err.Error(),
+		})
 	}
 
 	// Check if error is due to context cancellation
@@ -696,14 +763,21 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 
 	// Check if we received any data
 	if accumulator.id == "" {
-		log.Warn("OpenAI stream completed but no data was received (accumulator.id is empty)")
+		if trace != nil {
+			trace.Warn("OpenAI stream completed but no data was received")
 
-		// Log request details for debugging
+			// Log request details for debugging
-		if requestBodyJSON, err := jsoniter.Marshal(requestBody); err == nil {
+			if requestBodyJSON, err := jsoniter.Marshal(requestBody); err == nil {
-			log.Error("Request body that caused empty response: %s", string(requestBodyJSON))
+				trace.Error("Request body that caused empty response", map[string]any{
+					"body": string(requestBodyJSON),
+				})
+			}
+			trace.Error("Request details", map[string]any{
+				"url":     url,
+				"model":   accumulator.model,
+				"created": accumulator.created,
+			})
 		}
-		log.Error("Request URL: %s", url)
-		log.Error("Model in accumulator: %s, Created: %d", accumulator.model, accumulator.created)
 
 		err := fmt.Errorf("no data received from OpenAI API")
 
@@ -807,6 +881,14 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 
 // Post post completion request to OpenAI API
 func (p *Provider) Post(ctx *context.Context, messages []context.Message, options *context.CompletionOptions) (*context.CompletionResponse, error) {
+	// Add debug log
+	trace, _ := ctx.Trace()
+	if trace != nil {
+		trace.Debug("OpenAI Post: Starting non-stream request", map[string]any{
+			"message_count": len(messages),
+		})
+	}
+
 	maxRetries := 3
 	maxValidationRetries := 3
 	var lastErr error
@@ -833,7 +915,14 @@ func (p *Provider) Post(ctx *context.Context, messages []context.Message, option
 		if attempt > 0 {
 			// Exponential backoff
 			backoff := time.Duration(1<<uint(attempt-1)) * time.Second
-			log.Warn("OpenAI post request failed, retrying in %v (attempt %d/%d): %v", backoff, attempt+1, maxRetries, lastErr)
+			if trace != nil {
+				trace.Warn("OpenAI post request failed, retrying", map[string]any{
+					"backoff":     backoff.String(),
+					"attempt":     attempt + 1,
+					"max_retries": maxRetries,
+					"error":       lastErr.Error(),
+				})
+			}
 
 			// Sleep with context cancellation support
 			timer := time.NewTimer(backoff)
@@ -857,7 +946,13 @@ func (p *Provider) Post(ctx *context.Context, messages []context.Message, option
 			// Handle tool call validation retry with feedback to LLM
 			validationRetryMessages := currentMessages
 			for validationAttempt := 0; validationAttempt < maxValidationRetries; validationAttempt++ {
-				log.Warn("Tool call validation failed (attempt %d/%d): %v", validationAttempt+1, maxValidationRetries, err)
+				if trace != nil {
+					trace.Warn("Tool call validation failed", map[string]any{
+						"attempt":     validationAttempt + 1,
+						"max_retries": maxValidationRetries,
+						"error":       err.Error(),
+					})
+				}
 
 				// Add error feedback to conversation history
 				validationRetryMessages = append(validationRetryMessages, context.Message{
@@ -897,6 +992,9 @@ func (p *Provider) Post(ctx *context.Context, messages []context.Message, option
 
 // postWithRetry performs a single POST request attempt
 func (p *Provider) postWithRetry(ctx *context.Context, messages []context.Message, options *context.CompletionOptions) (*context.CompletionResponse, error) {
+	// Get trace from context
+	trace, _ := ctx.Trace()
+
 	// Preprocess messages and options through adapters
 	processedMessages := messages
 	processedOptions := options
@@ -958,8 +1056,12 @@ func (p *Provider) postWithRetry(ctx *context.Context, messages []context.Messag
 				}
 			}
 			// Log full response data for debugging
-			if respJSON, err := jsoniter.Marshal(resp.Data); err == nil {
+			if trace != nil {
-				log.Error("OpenAI API error response: %s", string(respJSON))
+				if respJSON, err := jsoniter.Marshal(resp.Data); err == nil {
+					trace.Error("OpenAI API error response", map[string]any{
+						"response": string(respJSON),
+					})
+				}
 			}
 		}
 		return nil, fmt.Errorf("HTTP %d: %s", resp.Code, errorMsg)

agent/output/adapters/cui/writer.go

@@ -26,12 +26,23 @@ func (w *Writer) Write(msg *message.Message) error {
 	// CUI adapter passes messages through as-is
 	chunks, err := w.adapter.Adapt(msg)
 	if err != nil {
+		if trace, _ := w.ctx.Trace(); trace != nil {
+			trace.Error("CUI Writer: Failed to adapt message", map[string]any{
+				"error":        err.Error(),
+				"message_type": msg.Type,
+			})
+		}
 		return err
 	}
 
 	// Send each chunk
 	for _, chunk := range chunks {
 		if err := w.sendChunk(chunk); err != nil {
+			if trace, _ := w.ctx.Trace(); trace != nil {
+				trace.Error("CUI Writer: Failed to send chunk", map[string]any{
+					"error": err.Error(),
+				})
+			}
 			return err
 		}
 	}
@@ -46,6 +57,12 @@ func (w *Writer) WriteGroup(group *message.MessageGroup) error {
 
 	// Send the group
 	if err := w.sendChunk(group); err != nil {
+		if trace, _ := w.ctx.Trace(); trace != nil {
+			trace.Error("CUI Writer: Failed to send message group", map[string]any{
+				"error":    err.Error(),
+				"group_id": group.ID,
+			})
+		}
 		return err
 	}
 
@@ -70,10 +87,31 @@ func (w *Writer) sendChunk(chunk interface{}) error {
 	// Convert chunk to JSON
 	data, err := json.Marshal(chunk)
 	if err != nil {
+		if trace, _ := w.ctx.Trace(); trace != nil {
+			trace.Error("CUI Writer: Failed to marshal chunk", map[string]any{
+				"error": err.Error(),
+			})
+		}
 		return err
 	}
 
+	// Log outgoing data to trace for debugging
+	if trace, _ := w.ctx.Trace(); trace != nil {
+		trace.Debug("CUI Writer: Sending chunk to client", map[string]any{
+			"data": string(data),
+		})
+	}
+
 	// Send via context's writer
 	// The context knows how to send data based on the connection type (SSE, WebSocket, etc.)
-	return w.ctx.Send(data)
+	if err := w.ctx.Send(data); err != nil {
+		if trace, _ := w.ctx.Trace(); trace != nil {
+			trace.Error("CUI Writer: Failed to send data to client", map[string]any{
+				"error": err.Error(),
+			})
+		}
+		return err
+	}
+
+	return nil
 }

agent/output/adapters/openai/writer.go

@@ -2,7 +2,6 @@ package openai
 
 import (
 	"encoding/json"
-	"fmt"
 
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -32,6 +31,12 @@ func (w *Writer) Write(msg *message.Message) error {
 	// Convert message to OpenAI format using adapter
 	chunks, err := w.adapter.Adapt(msg)
 	if err != nil {
+		if trace, _ := w.ctx.Trace(); trace != nil {
+			trace.Error("OpenAI Writer: Failed to adapt message", map[string]any{
+				"error":        err.Error(),
+				"message_type": msg.Type,
+			})
+		}
 		return err
 	}
 
@@ -50,6 +55,11 @@ func (w *Writer) Write(msg *message.Message) error {
 		}
 
 		if err := w.sendChunk(chunk); err != nil {
+			if trace, _ := w.ctx.Trace(); trace != nil {
+				trace.Error("OpenAI Writer: Failed to send chunk", map[string]any{
+					"error": err.Error(),
+				})
+			}
 			return err
 		}
 	}
@@ -63,6 +73,13 @@ func (w *Writer) WriteGroup(group *message.MessageGroup) error {
 	// Just send each message individually
 	for _, msg := range group.Messages {
 		if err := w.Write(msg); err != nil {
+			if trace, _ := w.ctx.Trace(); trace != nil {
+				trace.Error("OpenAI Writer: Failed to write message in group", map[string]any{
+					"error":        err.Error(),
+					"group_id":     group.ID,
+					"message_type": msg.Type,
+				})
+			}
 			return err
 		}
 	}
@@ -88,25 +105,55 @@ func (w *Writer) sendChunk(chunk interface{}) error {
 	// Convert chunk to JSON
 	data, err := json.Marshal(chunk)
 	if err != nil {
+		if trace, _ := w.ctx.Trace(); trace != nil {
+			trace.Error("OpenAI Writer: Failed to marshal chunk", map[string]any{
+				"error": err.Error(),
+			})
+		}
 		return err
 	}
 
-	// Debug: print the chunk being sent
-	fmt.Println("-----------------------------------------------")
-	fmt.Println("Sending SSE chunk: ", string(data))
-	fmt.Println("-----------------------------------------------")
-
 	// Format as SSE: "data: {json}\n\n"
 	sseData := append([]byte("data: "), data...)
 	sseData = append(sseData, []byte("\n\n")...)
 
+	// Log outgoing data to trace for debugging
+	if trace, _ := w.ctx.Trace(); trace != nil {
+		trace.Debug("OpenAI Writer: Sending chunk to client", map[string]any{
+			"data": string(data),
+		})
+	}
+
 	// Send via context's writer
-	return w.ctx.Send(sseData)
+	if err := w.ctx.Send(sseData); err != nil {
+		if trace, _ := w.ctx.Trace(); trace != nil {
+			trace.Error("OpenAI Writer: Failed to send data to client", map[string]any{
+				"error": err.Error(),
+			})
+		}
+		return err
+	}
+
+	return nil
 }
 
 // sendDone sends the final [DONE] message
 func (w *Writer) sendDone() error {
+	// Log completion to trace
+	if trace, _ := w.ctx.Trace(); trace != nil {
+		trace.Debug("OpenAI Writer: Sending [DONE] to client")
+	}
+
 	// OpenAI SSE format uses "data: [DONE]\n\n" to signal completion
 	doneData := []byte("data: [DONE]\n\n")
-	return w.ctx.Send(doneData)
+	if err := w.ctx.Send(doneData); err != nil {
+		if trace, _ := w.ctx.Trace(); trace != nil {
+			trace.Error("OpenAI Writer: Failed to send [DONE] to client", map[string]any{
+				"error": err.Error(),
+			})
+		}
+		return err
+	}
+
+	return nil
 }

trace/README.md

@@ -99,7 +99,7 @@ func main() {
             manager.SetSpaceValue(space.ID, fmt.Sprintf("worker_%d", idx+1), map[string]any{
                 "id":     idx + 1,
                 "status": "done",
-                "time":   time.Now().Unix(),
+                "time":   time.Now().UnixMilli(),
             })
 
             // Each worker completes itself
@@ -124,7 +124,7 @@ func main() {
         Icon:  "database",
     })
     manager.SetSpaceValue(sessionSpace.ID, "total_processed", 3)
-    manager.SetSpaceValue(sessionSpace.ID, "timestamp", time.Now().Unix())
+    manager.SetSpaceValue(sessionSpace.ID, "timestamp", time.Now().UnixMilli())
 
     manager.Complete(map[string]any{"total": 3, "success": true})
 

trace/manager.go

@@ -49,7 +49,7 @@ func NewManager(ctx context.Context, traceID string, driver types.Driver) (types
 		}
 	} else {
 		// New trace - create and broadcast init event
-		now := time.Now().Unix()
+		now := time.Now().UnixMilli()
 		m.addUpdateAndBroadcast(&types.TraceUpdate{
 			Type:      types.UpdateTypeInit,
 			TraceID:   traceID,
@@ -100,7 +100,7 @@ func (m *manager) handleCancellation() {
 		return // Already completed
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Get current nodes (state worker will process this before exiting)
 	nodes := m.stateGetCurrentNodes()
@@ -125,7 +125,7 @@ func (m *manager) handleCancellation() {
 					NodeID:   node.ID,
 					Status:   types.CompleteStatusCancelled,
 					EndTime:  now,
-					Duration: (now - node.StartTime) * 1000,
+					Duration: now - node.StartTime, // Already in milliseconds
 					Error:    "context cancelled",
 				},
 			})
@@ -139,7 +139,7 @@ func (m *manager) handleCancellation() {
 	totalDuration := int64(0)
 	rootNode := m.stateGetRoot()
 	if rootNode != nil && rootNode.CreatedAt > 0 {
-		totalDuration = (now - rootNode.CreatedAt) * 1000
+		totalDuration = now - rootNode.CreatedAt // Already in milliseconds
 	}
 
 	// Broadcast trace cancelled event (this will be processed even after state worker starts draining)
@@ -169,7 +169,7 @@ func (m *manager) Add(input types.TraceInput, option types.TraceNodeOption) (typ
 		return nil, err
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Check if root exists
 	rootNode := m.stateGetRoot()
@@ -291,7 +291,7 @@ func (m *manager) Parallel(parallelInputs []types.TraceParallelInput) ([]types.N
 		return nil, fmt.Errorf("root node does not exist, please call Add first before using Parallel")
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Get current nodes
 	currentNodes := m.stateGetCurrentNodes()
@@ -356,33 +356,32 @@ func (m *manager) Parallel(parallelInputs []types.TraceParallelInput) ([]types.N
 }
 
 // Info logs info message to current node(s)
-func (m *manager) Info(format string, args ...any) types.Manager {
+func (m *manager) Info(message string, args ...any) types.Manager {
-	m.log("info", format, args...)
+	m.log("info", message, args...)
 	return m
 }
 
 // Debug logs debug message to current node(s)
-func (m *manager) Debug(format string, args ...any) types.Manager {
+func (m *manager) Debug(message string, args ...any) types.Manager {
-	m.log("debug", format, args...)
+	m.log("debug", message, args...)
 	return m
 }
 
 // Error logs error message to current node(s)
-func (m *manager) Error(format string, args ...any) types.Manager {
+func (m *manager) Error(message string, args ...any) types.Manager {
-	m.log("error", format, args...)
+	m.log("error", message, args...)
 	return m
 }
 
 // Warn logs warning message to current node(s)
-func (m *manager) Warn(format string, args ...any) types.Manager {
+func (m *manager) Warn(message string, args ...any) types.Manager {
-	m.log("warn", format, args...)
+	m.log("warn", message, args...)
 	return m
 }
 
 // log helper method to log messages
-func (m *manager) log(level string, format string, args ...any) {
+func (m *manager) log(level string, message string, args ...any) {
-	message := fmt.Sprintf(format, args...)
+	now := time.Now().UnixMilli()
-	now := time.Now().Unix()
 
 	// Get current nodes
 	nodes := m.stateGetCurrentNodes()
@@ -393,6 +392,7 @@ func (m *manager) log(level string, format string, args ...any) {
 			Timestamp: now,
 			Level:     level,
 			Message:   message,
+			Data:      args,
 			NodeID:    node.ID,
 		}
 		// Save log (ignore errors for non-critical logging)
@@ -415,7 +415,7 @@ func (m *manager) SetOutput(output types.TraceOutput) error {
 		return err
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 	nodes := m.stateGetCurrentNodes()
 	for _, node := range nodes {
 		node.Output = output
@@ -442,7 +442,7 @@ func (m *manager) SetMetadata(key string, value any) error {
 		return err
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 	nodes := m.stateGetCurrentNodes()
 	for _, node := range nodes {
 		if node.Metadata == nil {
@@ -473,7 +473,7 @@ func (m *manager) Complete(output ...types.TraceOutput) error {
 		return err
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 	nodes := m.stateGetCurrentNodes()
 
 	// Determine output value once
@@ -493,7 +493,7 @@ func (m *manager) Complete(output ...types.TraceOutput) error {
 			NodeID:   node.ID,
 			Status:   types.CompleteStatusSuccess,
 			EndTime:  now,
-			Duration: (now - node.StartTime) * 1000,
+			Duration: now - node.StartTime, // Already in milliseconds
 			Output:   node.Output,
 		}
 
@@ -523,7 +523,7 @@ func (m *manager) Fail(err error) error {
 		return ctxErr
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 	// Log error first
 	m.Error("Node failed: %v", err)
 
@@ -546,7 +546,7 @@ func (m *manager) Fail(err error) error {
 				NodeID:   node.ID,
 				Status:   types.CompleteStatusFailed,
 				EndTime:  now,
-				Duration: (node.EndTime - node.StartTime) * 1000, // Convert to milliseconds
+				Duration: node.EndTime - node.StartTime, // Already in milliseconds
 				Error:    err.Error(),
 			},
 		})
@@ -580,11 +580,11 @@ func (m *manager) MarkComplete() error {
 	m.stateSetTraceStatus(types.TraceStatusCompleted)
 
 	// Calculate total duration from root node
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 	totalDuration := int64(0)
 	rootNode := m.stateGetRoot()
 	if rootNode != nil && rootNode.CreatedAt > 0 {
-		totalDuration = (now - rootNode.CreatedAt) * 1000 // Convert to milliseconds
+		totalDuration = now - rootNode.CreatedAt // Already in milliseconds
 	}
 
 	// Broadcast completion event
@@ -604,7 +604,7 @@ func (m *manager) CreateSpace(option types.TraceSpaceOption) (*types.TraceSpace,
 		return nil, err
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Create space instance
 	space := &types.TraceSpace{
@@ -673,7 +673,7 @@ func (m *manager) DeleteSpace(id string) error {
 		return err
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Remove from cache
 	m.stateDeleteSpace(id)
@@ -722,7 +722,7 @@ func (m *manager) SetSpaceValue(spaceID, key string, value any) error {
 		return err
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Get space
 	space, err := m.GetSpace(spaceID)
@@ -788,7 +788,7 @@ func (m *manager) DeleteSpaceValue(spaceID, key string) error {
 		return err
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Delete value from driver (through state worker for concurrent safety)
 	err := m.stateExecuteSpaceOp(spaceID, func() error {
@@ -817,7 +817,7 @@ func (m *manager) ClearSpaceValues(spaceID string) error {
 		return err
 	}
 
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Clear values from driver (through state worker for concurrent safety)
 	err := m.stateExecuteSpaceOp(spaceID, func() error {

trace/node.go

@@ -1,7 +1,6 @@
 package trace
 
 import (
-	"fmt"
 	"time"
 
 	"github.com/yaoapp/yao/trace/types"
@@ -14,32 +13,32 @@ type node struct {
 }
 
 // Info logs info message (public method, broadcasts event)
-func (n *node) Info(format string, args ...any) types.Node {
+func (n *node) Info(message string, args ...any) types.Node {
-	n.logWithBroadcast("info", format, args...)
+	n.logWithBroadcast("info", message, args...)
 	return n
 }
 
 // Debug logs debug message (public method, broadcasts event)
-func (n *node) Debug(format string, args ...any) types.Node {
+func (n *node) Debug(message string, args ...any) types.Node {
-	n.logWithBroadcast("debug", format, args...)
+	n.logWithBroadcast("debug", message, args...)
 	return n
 }
 
 // Error logs error message (public method, broadcasts event)
-func (n *node) Error(format string, args ...any) types.Node {
+func (n *node) Error(message string, args ...any) types.Node {
-	n.logWithBroadcast("error", format, args...)
+	n.logWithBroadcast("error", message, args...)
 	return n
 }
 
 // Warn logs warning message (public method, broadcasts event)
-func (n *node) Warn(format string, args ...any) types.Node {
+func (n *node) Warn(message string, args ...any) types.Node {
-	n.logWithBroadcast("warn", format, args...)
+	n.logWithBroadcast("warn", message, args...)
 	return n
 }
 
 // logWithBroadcast logs and broadcasts event (for external calls)
-func (n *node) logWithBroadcast(level string, format string, args ...any) {
+func (n *node) logWithBroadcast(level string, message string, args ...any) {
-	log := n.log(level, format, args...)
+	log := n.log(level, message, args...)
 
 	// Broadcast event
 	n.manager.addUpdateAndBroadcast(&types.TraceUpdate{
@@ -52,12 +51,12 @@ func (n *node) logWithBroadcast(level string, format string, args ...any) {
 }
 
 // log logs without broadcasting (for internal Manager calls)
-func (n *node) log(level string, format string, args ...any) *types.TraceLog {
+func (n *node) log(level string, message string, args ...any) *types.TraceLog {
-	message := fmt.Sprintf(format, args...)
 	log := &types.TraceLog{
-		Timestamp: time.Now().Unix(),
+		Timestamp: time.Now().UnixMilli(),
 		Level:     level,
 		Message:   message,
+		Data:      args,
 		NodeID:    n.data.ID,
 	}
 	// Save log (ignore errors for non-critical logging)
@@ -67,7 +66,7 @@ func (n *node) log(level string, format string, args ...any) *types.TraceLog {
 
 // Add creates next sequential node
 func (n *node) Add(input types.TraceInput, option types.TraceNodeOption) (types.Node, error) {
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Create child node data
 	childNodeData := &types.TraceNode{
@@ -102,7 +101,7 @@ func (n *node) Add(input types.TraceInput, option types.TraceNodeOption) (types.
 
 // Parallel creates multiple concurrent child nodes
 func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node, error) {
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 	nodeInterfaces := make([]types.Node, 0, len(parallelInputs))
 
 	// Create multiple child nodes
@@ -142,7 +141,7 @@ func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node
 
 // Join joins multiple nodes into one
 func (n *node) Join(nodes []*types.TraceNode, input types.TraceInput, option types.TraceNodeOption) (types.Node, error) {
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Create join node data
 	joinNodeData := &types.TraceNode{
@@ -177,7 +176,7 @@ func (n *node) ID() string {
 // SetOutput sets the node output
 func (n *node) SetOutput(output types.TraceOutput) error {
 	n.data.Output = output
-	n.data.UpdatedAt = time.Now().Unix()
+	n.data.UpdatedAt = time.Now().UnixMilli()
 	return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data)
 }
 
@@ -187,14 +186,14 @@ func (n *node) SetMetadata(key string, value any) error {
 		n.data.Metadata = make(map[string]any)
 	}
 	n.data.Metadata[key] = value
-	n.data.UpdatedAt = time.Now().Unix()
+	n.data.UpdatedAt = time.Now().UnixMilli()
 	return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data)
 }
 
 // SetStatus sets the node status
 func (n *node) SetStatus(status string) error {
 	n.data.Status = types.NodeStatus(status)
-	n.data.UpdatedAt = time.Now().Unix()
+	n.data.UpdatedAt = time.Now().UnixMilli()
 	return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data)
 }
 
@@ -219,7 +218,7 @@ func (n *node) Complete(output ...types.TraceOutput) error {
 
 // complete marks as completed without broadcasting (for Manager calls)
 func (n *node) complete(output ...types.TraceOutput) error {
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Set output if provided
 	if len(output) > 0 {
@@ -255,7 +254,7 @@ func (n *node) Fail(err error) error {
 
 // fail marks as failed without broadcasting (for Manager calls)
 func (n *node) fail(err error) error {
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Update status
 	n.data.Status = types.StatusFailed

trace/subscription.go

@@ -7,9 +7,9 @@ import (
 	"github.com/yaoapp/yao/trace/types"
 )
 
-// Subscribe creates a new subscription for trace updates (real-time from now)
+// Subscribe creates a new subscription for trace updates (replays all historical events from the beginning)
 func (m *manager) Subscribe() (<-chan *types.TraceUpdate, error) {
-	return m.subscribe(time.Now().Unix())
+	return m.subscribe(0) // Subscribe from beginning to get all historical events
 }
 
 // SubscribeFrom creates a subscription starting from a specific timestamp
@@ -47,15 +47,25 @@ func (m *manager) replayAndStream(subID string, ch chan *types.TraceUpdate, sinc
 	// Get historical updates
 	updates := m.stateGetUpdates(since)
 
-	// Replay historical updates
+	// Replay historical updates and check if trace was already completed
+	traceWasCompleted := false
 	for _, update := range updates {
 		select {
 		case ch <- update:
+			// Check if this is a trace complete event
+			if update.Type == types.UpdateTypeComplete {
+				traceWasCompleted = true
+			}
 		case <-m.ctx.Done():
 			return
 		}
 	}
 
+	// If trace was already completed in historical events, exit immediately
+	if traceWasCompleted {
+		return
+	}
+
 	// Continue streaming new updates
 	// The channel will receive updates via broadcast from addUpdate
 	// Monitor completion to know when to exit

trace/trace.go

@@ -103,7 +103,7 @@ func GenTraceID(safe ...bool) string {
 // option: trace options (optional)
 // driverOptions: driver-specific options (e.g., base path for local, store name for store)
 func New(ctx context.Context, driver string, option *types.TraceOption, driverOptions ...any) (string, types.Manager, error) {
-	now := time.Now().Unix()
+	now := time.Now().UnixMilli()
 
 	// Handle nil option
 	if option == nil {
@@ -235,7 +235,7 @@ func LoadFromStorage(ctx context.Context, driver string, traceID string, driverO
 
 	// Update stored info with new manager
 	storedInfo.Manager = manager
-	storedInfo.UpdatedAt = time.Now().Unix()
+	storedInfo.UpdatedAt = time.Now().UnixMilli()
 
 	// Register in global registry
 	registryMu.Lock()

trace/trace_subscription_test.go

@@ -124,7 +124,7 @@ func TestSubscribeFrom(t *testing.T) {
 			time.Sleep(1100 * time.Millisecond)
 
 			// Record timestamp (simulate user noting current time before refresh)
-			resumeTimestamp := time.Now().Unix()
+			resumeTimestamp := time.Now().UnixMilli()
 
 			// Wait again to ensure next operations are after resumeTimestamp
 			time.Sleep(100 * time.Millisecond)

trace/types/driver.go

@@ -4,10 +4,11 @@ import "context"
 
 // TraceLog represents a log entry
 type TraceLog struct {
-	Timestamp int64  // Log timestamp
+	Timestamp int64    // Log timestamp (milliseconds since epoch)
-	Level     string // Log level (info, debug, error, warn)
+	Level     string   // Log level (info, debug, error, warn)
-	Message   string // Log message
+	Message   string   // Log message
-	NodeID    string // Node ID this log belongs to
+	Data      []any    // Additional data arguments
+	NodeID    string   // Node ID this log belongs to
 }
 
 // Driver defines the storage driver interface that providers must implement
@@ -71,7 +72,7 @@ type Driver interface {
 	// SaveUpdate persists a trace update event to storage
 	SaveUpdate(ctx context.Context, traceID string, update *TraceUpdate) error
 
-	// LoadUpdates loads trace update events from storage (filtering by timestamp)
+	// LoadUpdates loads trace update events from storage (filtering by timestamp in milliseconds)
 	LoadUpdates(ctx context.Context, traceID string, since int64) ([]*TraceUpdate, error)
 
 	// Close closes the driver and releases resources

trace/types/events.go

@@ -18,7 +18,7 @@ func (n *TraceNode) ToCompleteData() *NodeCompleteData {
 		NodeID:   n.ID,
 		Status:   CompleteStatusSuccess,
 		EndTime:  n.EndTime,
-		Duration: (n.EndTime - n.StartTime) * 1000, // Convert to milliseconds
+		Duration: n.EndTime - n.StartTime, // Already in milliseconds
 		Output:   n.Output,
 	}
 }
@@ -29,7 +29,7 @@ func (n *TraceNode) ToFailedData(err error) *NodeFailedData {
 		NodeID:   n.ID,
 		Status:   CompleteStatusFailed,
 		EndTime:  n.EndTime,
-		Duration: (n.EndTime - n.StartTime) * 1000, // Convert to milliseconds
+		Duration: n.EndTime - n.StartTime, // Already in milliseconds
 		Error:    err.Error(),
 	}
 }

trace/types/interfaces.go

@@ -12,10 +12,10 @@ type Manager interface {
 	Parallel(parallelInputs []TraceParallelInput) ([]Node, error)
 
 	// Log Operations - log to current node(s) with chainable interface
-	Info(format string, args ...any) Manager
+	Info(message string, args ...any) Manager
-	Debug(format string, args ...any) Manager
+	Debug(message string, args ...any) Manager
-	Error(format string, args ...any) Manager
+	Error(message string, args ...any) Manager
-	Warn(format string, args ...any) Manager
+	Warn(message string, args ...any) Manager
 
 	// Node Status Operations - operate on current node(s)
 	SetOutput(output TraceOutput) error
@@ -60,10 +60,10 @@ type Manager interface {
 // Context is bound to Node at creation time
 type Node interface {
 	// Log Operations - chainable interface
-	Info(format string, args ...any) Node
+	Info(message string, args ...any) Node
-	Debug(format string, args ...any) Node
+	Debug(message string, args ...any) Node
-	Error(format string, args ...any) Node
+	Error(message string, args ...any) Node
-	Warn(format string, args ...any) Node
+	Warn(message string, args ...any) Node
 
 	// Node Tree Operations
 	Add(input TraceInput, option TraceNodeOption) (Node, error)

trace/types/types.go

@@ -61,10 +61,10 @@ type TraceNode struct {
 	Status          NodeStatus   // Node status (pending, running, completed, failed, skipped)
 	Input           TraceInput   // Node input data
 	Output          TraceOutput  // Node output data
-	CreatedAt       int64        // Creation timestamp
+	CreatedAt       int64        // Creation timestamp (milliseconds since epoch)
-	StartTime       int64        // Start timestamp
+	StartTime       int64        // Start timestamp (milliseconds since epoch)
-	EndTime         int64        // End timestamp
+	EndTime         int64        // End timestamp (milliseconds since epoch)
-	UpdatedAt       int64        // Last update timestamp
+	UpdatedAt       int64        // Last update timestamp (milliseconds since epoch)
 	// Other fields will be added during implementation
 }
 
@@ -72,8 +72,8 @@ type TraceNode struct {
 type TraceSpace struct {
 	ID               string // Space ID
 	TraceSpaceOption        // Embedded option fields (Label, Icon, Description, TTL, Metadata)
-	CreatedAt        int64  // Creation timestamp
+	CreatedAt        int64  // Creation timestamp (milliseconds since epoch)
-	UpdatedAt        int64  // Last update timestamp
+	UpdatedAt        int64  // Last update timestamp (milliseconds since epoch)
 	// Internal data storage will be managed by implementation
 }
 
@@ -118,7 +118,7 @@ type TraceUpdate struct {
 	TraceID   string // Trace ID
 	NodeID    string // Node ID (optional, for node/log updates)
 	SpaceID   string // Space ID (optional, for space updates)
-	Timestamp int64  // Update timestamp
+	Timestamp int64  // Update timestamp (milliseconds since epoch)
 	Data      any    // Update data (payload structures below)
 }
 
@@ -141,18 +141,18 @@ type NodeStartData struct {
 // NodeCompleteData payload for "node_complete" event
 type NodeCompleteData struct {
 	NodeID   string         `json:"nodeId"`
-	Status   CompleteStatus `json:"status"` // "success" or "failed"
+	Status   CompleteStatus `json:"status"`   // "success" or "failed"
-	EndTime  int64          `json:"endTime"`
+	EndTime  int64          `json:"endTime"`  // milliseconds since epoch
-	Duration int64          `json:"duration"` // in milliseconds
+	Duration int64          `json:"duration"` // duration in milliseconds
 	Output   TraceOutput    `json:"output,omitempty"`
 }
 
 // NodeFailedData payload for "node_failed" event (same as NodeCompleteData but with error)
 type NodeFailedData struct {
 	NodeID   string         `json:"nodeId"`
-	Status   CompleteStatus `json:"status"` // "failed"
+	Status   CompleteStatus `json:"status"`   // "failed"
-	EndTime  int64          `json:"endTime"`
+	EndTime  int64          `json:"endTime"`  // milliseconds since epoch
-	Duration int64          `json:"duration"`
+	Duration int64          `json:"duration"` // duration in milliseconds
 	Error    string         `json:"error"`
 }
 
@@ -168,15 +168,15 @@ type MemoryItem struct {
 	Type       string `json:"type"`
 	Title      string `json:"title,omitempty"`
 	Content    any    `json:"content"`
-	Timestamp  int64  `json:"timestamp"`
+	Timestamp  int64  `json:"timestamp"`            // milliseconds since epoch
 	Importance string `json:"importance,omitempty"` // "high", "medium", "low"
 }
 
 // TraceCompleteData payload for "complete" event
 type TraceCompleteData struct {
 	TraceID       string      `json:"traceId"`
-	Status        TraceStatus `json:"status"` // "completed"
+	Status        TraceStatus `json:"status"`        // "completed"
-	TotalDuration int64       `json:"totalDuration"`
+	TotalDuration int64       `json:"totalDuration"` // duration in milliseconds
 }
 
 // SpaceDeletedData payload for "space_deleted" event
@@ -197,9 +197,9 @@ type TraceInfo struct {
 	Driver    string         `json:"driver"`
 	Status    TraceStatus    `json:"status"` // Trace status
 	Options   []any          `json:"options,omitempty"`
-	Manager   Manager        `json:"-"` // Not persisted
+	Manager   Manager        `json:"-"`          // Not persisted
-	CreatedAt int64          `json:"created_at"`
+	CreatedAt int64          `json:"created_at"` // milliseconds since epoch
-	UpdatedAt int64          `json:"updated_at"`
+	UpdatedAt int64          `json:"updated_at"` // milliseconds since epoch
 	CreatedBy string         `json:"__yao_created_by,omitempty"`
 	UpdatedBy string         `json:"__yao_updated_by,omitempty"`
 	TeamID    string         `json:"__yao_team_id,omitempty"`