diff --git a/agent/trace/README.md b/agent/trace/README.md
new file mode 100644
index 00000000..9f512183
--- /dev/null
+++ b/agent/trace/README.md
@@ -0,0 +1,572 @@
+# Trace Package
+
+A trace system for logging and visualizing execution flow with real-time event streaming support.
+
+## Features
+
+- **Node Tree Structure**: Build execution trees with sequential and parallel operations
+- **Real-time Events**: Subscribe to trace updates with history replay and SSE support
+- **Memory Spaces**: Key-value storage for session data and context
+- **Dual Storage**: Local disk and Gou store backends
+- **Concurrent Safe**: Thread-safe operations with context cancellation
+- **Auto-join**: Automatic handling of parallel to sequential transitions
+
+## Quick Start
+
+Complete example program:
+
+```go
+package main
+
+import (
+ "context"
+ "fmt"
+ "sync"
+ "time"
+
+ "github.com/yaoapp/yao/agent/trace"
+ "github.com/yaoapp/yao/agent/trace/types"
+)
+
+func main() {
+ // Create a new trace with custom ID
+ ctx := context.Background()
+ traceID := trace.GenTraceID()
+
+ option := &types.TraceOption{
+ ID: traceID,
+ CreatedBy: "user@example.com",
+ Metadata: map[string]any{"task": "demo"},
+ }
+
+ _, manager, err := trace.New(ctx, trace.Local, option)
+ if err != nil {
+ panic(err)
+ }
+ defer trace.Release(traceID)
+
+ fmt.Printf("Trace ID: %s\n", traceID)
+
+ // Step 1: Input processing
+ manager.Info("Starting input processing")
+ _, err = manager.Add("user input data", types.TraceNodeOption{
+ Label: "Input Processing",
+ Icon: "processor",
+ })
+ if err != nil {
+ panic(err)
+ }
+ manager.Complete(map[string]any{"validated": true, "items": 3})
+
+ // Step 2: Parallel processing - each worker completes independently
+ manager.Info("Starting parallel tasks")
+
+ // Create a shared space for workers to store results
+ space, _ := manager.CreateSpace(types.TraceSpaceOption{
+ Label: "Worker Results",
+ Icon: "storage",
+ })
+
+ nodes, _ := manager.Parallel([]types.TraceParallelInput{
+ {
+ Input: "Processing task A",
+ Option: types.TraceNodeOption{Label: "Worker A", Icon: "cpu"},
+ },
+ {
+ Input: "Processing task B",
+ Option: types.TraceNodeOption{Label: "Worker B", Icon: "cpu"},
+ },
+ {
+ Input: "Processing task C",
+ Option: types.TraceNodeOption{Label: "Worker C", Icon: "cpu"},
+ },
+ })
+
+ // Each parallel node completes itself
+ var wg sync.WaitGroup
+ for i, node := range nodes {
+ wg.Add(1)
+ go func(idx int, n types.Node) {
+ defer wg.Done()
+
+ // Worker performs its task
+ n.Info("Worker %d processing", idx+1)
+
+ // Simulate different completion times
+ time.Sleep(time.Duration(50+idx*20) * time.Millisecond)
+
+ // Store result in shared space
+ manager.SetSpaceValue(space.ID, fmt.Sprintf("worker_%d", idx+1), map[string]any{
+ "id": idx + 1,
+ "status": "done",
+ "time": time.Now().Unix(),
+ })
+
+ // Each worker completes itself
+ n.Complete(map[string]any{"worker": idx + 1, "status": "done"})
+ }(i, node)
+ }
+ wg.Wait()
+
+ // Step 3: Aggregation (auto-joins parallel branches)
+ manager.Info("Aggregating results")
+ _, err = manager.Add("Merging outputs", types.TraceNodeOption{
+ Label: "Aggregation",
+ Icon: "merge",
+ })
+ if err != nil {
+ panic(err)
+ }
+
+ // Create another space for session data
+ sessionSpace, _ := manager.CreateSpace(types.TraceSpaceOption{
+ Label: "Session Data",
+ Icon: "database",
+ })
+ manager.SetSpaceValue(sessionSpace.ID, "total_processed", 3)
+ manager.SetSpaceValue(sessionSpace.ID, "timestamp", time.Now().Unix())
+
+ manager.Complete(map[string]any{"total": 3, "success": true})
+
+ // Mark trace as completed
+ manager.MarkComplete()
+
+ fmt.Println("Trace completed successfully!")
+}
+```
+
+### Execution Flow
+
+The program creates the following node tree:
+
+```mermaid
+graph TD
+ Root[Root Node
Status: Running]
+ Input[Input Processing
✓ Completed
Output: validated=true, items=3]
+
+ Fork{Parallel Fork}
+ Space1[(Worker Results
Space)]
+ WorkerA[Worker A
✓ Completed
Output: worker=1]
+ WorkerB[Worker B
✓ Completed
Output: worker=2]
+ WorkerC[Worker C
✓ Completed
Output: worker=3]
+
+ Join((Auto Join))
+ Agg[Aggregation
✓ Completed
Output: total=3, success=true]
+ Space2[(Session Data
Space)]
+
+ Root --> Input
+ Input --> Fork
+ Fork --> WorkerA
+ Fork --> WorkerB
+ Fork --> WorkerC
+ WorkerA -.-> Space1
+ WorkerB -.-> Space1
+ WorkerC -.-> Space1
+ WorkerA --> Join
+ WorkerB --> Join
+ WorkerC --> Join
+ Join --> Agg
+ Agg -.-> Space2
+
+ style Root fill:#e1f5ff
+ style Input fill:#c8e6c9
+ style WorkerA fill:#c8e6c9
+ style WorkerB fill:#c8e6c9
+ style WorkerC fill:#c8e6c9
+ style Agg fill:#c8e6c9
+ style Space1 fill:#fff9c4
+ style Space2 fill:#fff9c4
+```
+
+## Subscribe to Events
+
+Real-time event subscription example:
+
+```go
+package main
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "time"
+
+ "github.com/yaoapp/yao/agent/trace"
+ "github.com/yaoapp/yao/agent/trace/types"
+)
+
+func main() {
+ ctx := context.Background()
+ traceID := trace.GenTraceID()
+
+ option := &types.TraceOption{
+ ID: traceID,
+ CreatedBy: "user@example.com",
+ }
+
+ _, manager, _ := trace.New(ctx, trace.Local, option)
+ defer trace.Release(traceID)
+
+ // Subscribe to all events (history + real-time)
+ updates, _ := manager.Subscribe()
+
+ // Start event listener in goroutine
+ go func() {
+ for update := range updates {
+ // Convert to JSON for display
+ data, _ := json.MarshalIndent(update, "", " ")
+ fmt.Printf("\n[Event] %s at %d\n%s\n",
+ update.Type, update.Timestamp, string(data))
+
+ // Handle specific events
+ switch update.Type {
+ case types.UpdateTypeNodeStart:
+ fmt.Println("→ Node started")
+
+ case types.UpdateTypeNodeComplete:
+ fmt.Println("✓ Node completed")
+
+ case types.UpdateTypeMemoryAdd:
+ fmt.Println("💾 Memory updated")
+
+ case types.UpdateTypeComplete:
+ fmt.Println("🎉 Trace completed!")
+ return
+ }
+ }
+ }()
+
+ // Execute trace operations
+ manager.Info("Processing started")
+ manager.Add("Task 1", types.TraceNodeOption{Label: "Task 1"})
+ manager.Complete(map[string]any{"status": "ok"})
+
+ manager.Add("Task 2", types.TraceNodeOption{Label: "Task 2"})
+ manager.Complete(map[string]any{"status": "ok"})
+
+ manager.MarkComplete()
+
+ // Wait for events to be processed
+ time.Sleep(1 * time.Second)
+}
+```
+
+### Event Output Example
+
+```json
+[Event] init at 1700123456
+{
+ "Type": "init",
+ "TraceID": "20251118123456789012",
+ "Timestamp": 1700123456,
+ "Data": {
+ "traceId": "20251118123456789012",
+ "agentName": "",
+ "rootNode": {...}
+ }
+}
+→ Node started
+
+[Event] node_start at 1700123457
+{
+ "Type": "node_start",
+ "TraceID": "20251118123456789012",
+ "NodeID": "abc123def456",
+ "Timestamp": 1700123457,
+ "Data": {
+ "node": {
+ "ID": "abc123def456",
+ "Label": "Task 1",
+ "Status": "running"
+ }
+ }
+}
+→ Node started
+
+[Event] node_complete at 1700123458
+{
+ "Type": "node_complete",
+ "TraceID": "20251118123456789012",
+ "NodeID": "abc123def456",
+ "Timestamp": 1700123458,
+ "Data": {
+ "nodeId": "abc123def456",
+ "status": "success",
+ "endTime": 1700123458,
+ "duration": 1000,
+ "output": {"status": "ok"}
+ }
+}
+✓ Node completed
+
+[Event] complete at 1700123460
+{
+ "Type": "complete",
+ "TraceID": "20251118123456789012",
+ "Timestamp": 1700123460,
+ "Data": {
+ "traceId": "20251118123456789012",
+ "status": "completed",
+ "totalDuration": 4000
+ }
+}
+🎉 Trace completed!
+```
+
+## API Reference
+
+### Trace Management
+
+#### `New(ctx, driver, option, driverOptions...) (traceID, Manager, error)`
+
+Create a new trace or load existing one from storage.
+
+**Drivers:**
+
+- `trace.Local` - Local disk storage (default path: `./traces`)
+- `trace.Store` - Gou store backend
+
+**Example:**
+
+```go
+// Local with custom path
+traceID, manager, _ := trace.New(ctx, trace.Local, nil, "/data/traces")
+
+// Store with custom name
+traceID, manager, _ := trace.New(ctx, trace.Store, nil, "my_traces")
+
+// With trace options
+option := &types.TraceOption{
+ ID: "custom-id",
+ CreatedBy: "user@example.com",
+ TeamID: "team-001",
+ Metadata: map[string]any{"version": "1.0"},
+}
+traceID, manager, _ := trace.New(ctx, trace.Local, option)
+```
+
+#### `LoadFromStorage(ctx, driver, traceID, options...) (traceID, Manager, error)`
+
+Load an existing trace from persistent storage.
+
+#### `Load(traceID) (Manager, error)`
+
+Get an active trace from registry.
+
+#### `GetInfo(ctx, driver, traceID, options...) (*TraceInfo, error)`
+
+Retrieve trace metadata from storage.
+
+#### `IsLoaded(traceID) bool`
+
+Check if trace is active in registry.
+
+#### `Exists(ctx, driver, traceID, options...) (bool, error)`
+
+Check if trace exists in persistent storage.
+
+#### `Release(traceID) error`
+
+Remove trace from registry and release resources.
+
+#### `Remove(ctx, driver, traceID, options...) error`
+
+Delete trace and all associated data permanently.
+
+#### `List() []string`
+
+List all active trace IDs in registry.
+
+### Manager Interface
+
+#### Node Operations
+
+- `Add(input, option) Node` - Create sequential node, returns Node interface (auto-joins if parallel)
+- `Parallel(inputs) []Node` - Create concurrent child nodes, returns Node interfaces for direct control
+- `GetRootNode() *TraceNode` - Get root node data
+- `GetNode(id) *TraceNode` - Get node data by ID
+- `GetCurrentNodes() []*TraceNode` - Get active node data
+
+#### Logging (Chainable)
+
+- `Info(format, args...)` - Log info message
+- `Debug(format, args...)` - Log debug message
+- `Error(format, args...)` - Log error message
+- `Warn(format, args...)` - Log warning message
+
+#### Node Status
+
+- `SetOutput(output)` - Set output for current nodes
+- `SetMetadata(key, value)` - Set metadata
+- `Complete(output...)` - Mark nodes as completed
+- `Fail(err)` - Mark nodes as failed
+- `MarkComplete()` - Mark entire trace as completed
+
+#### Memory Spaces
+
+- `CreateSpace(option)` - Create new space
+- `GetSpace(id)` - Get space by ID
+- `HasSpace(id)` - Check if space exists
+- `DeleteSpace(id)` - Delete space
+- `ListSpaces()` - List all spaces
+
+#### Space Key-Value
+
+- `SetSpaceValue(spaceID, key, value)` - Set value (broadcasts event)
+- `GetSpaceValue(spaceID, key)` - Get value
+- `HasSpaceValue(spaceID, key)` - Check key existence
+- `DeleteSpaceValue(spaceID, key)` - Delete key
+- `ClearSpaceValues(spaceID)` - Clear all keys
+- `ListSpaceKeys(spaceID)` - List all keys
+
+#### Subscription
+
+- `Subscribe()` - Subscribe to events (history + real-time)
+- `SubscribeFrom(since)` - Subscribe from timestamp
+- `IsComplete()` - Check if trace is completed
+
+### Node Interface
+
+The Node interface is returned by `Manager.Add()` and `Manager.Parallel()` for direct control of individual nodes, typically used in parallel operations.
+
+#### Node Operations
+
+- `Add(input, option) Node` - Create child node
+- `Parallel(inputs) []Node` - Create parallel child nodes
+- `Join(nodes, input, option) Node` - Join multiple nodes into one
+- `ID() string` - Get node ID
+
+#### Logging (Chainable)
+
+- `Info(format, args...)` - Log info message
+- `Debug(format, args...)` - Log debug message
+- `Error(format, args...)` - Log error message
+- `Warn(format, args...)` - Log warning message
+
+#### Node Status
+
+- `SetOutput(output)` - Set node output
+- `SetMetadata(key, value)` - Set node metadata
+- `SetStatus(status)` - Set node status
+- `Complete(output...)` - Mark node as completed (broadcasts event)
+- `Fail(err)` - Mark node as failed (broadcasts event)
+
+### Event Types
+
+- `init` - Trace initialization
+- `node_start` - Node created
+- `node_complete` - Node completed
+- `node_failed` - Node failed
+- `node_updated` - Node data updated
+- `log_added` - Log entry added
+- `memory_add` - Space value added
+- `memory_update` - Space value updated
+- `memory_delete` - Space value deleted
+- `space_created` - Space created
+- `space_deleted` - Space deleted
+- `complete` - Trace completed
+
+## Storage Format
+
+### TraceID Format
+
+`YYYYMMDDnnnnnnnnnnnn` (20 digits)
+
+- First 8 digits: Date (YYYYMMDD)
+- Last 12 digits: Unique identifier
+
+Example: `20251118123456789012`
+
+### Local Driver Structure
+
+```
+./traces/
+ └── 20251118/
+ └── {traceID}/
+ ├── trace_info.json
+ ├── nodes/
+ │ ├── {nodeID}.json
+ │ └── ...
+ ├── spaces/
+ │ ├── {spaceID}.json
+ │ └── {spaceID}/
+ │ └── data.json
+ └── logs/
+ └── {nodeID}.jsonl
+```
+
+## Advanced Usage
+
+### Custom Node Operations
+
+For fine-grained control in parallel operations:
+
+```go
+nodes, _ := manager.Parallel(parallelInputs)
+
+// Each goroutine controls its own node
+var wg sync.WaitGroup
+for i, node := range nodes {
+ wg.Add(1)
+ go func(idx int, n types.Node) {
+ defer wg.Done()
+
+ n.Info("Worker %d started", idx+1)
+
+ // Simulate different processing times
+ time.Sleep(time.Duration(100+idx*50) * time.Millisecond)
+
+ n.Complete(result)
+ }(i, node)
+}
+wg.Wait()
+```
+
+### Context Cancellation
+
+The manager respects context cancellation:
+
+```go
+ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+defer cancel()
+
+traceID, manager, _ := trace.New(ctx, trace.Local, nil)
+
+// All operations will check context
+manager.Add(input, option) // Returns error if context cancelled
+```
+
+### Server-Sent Events (SSE)
+
+```go
+func traceSSEHandler(w http.ResponseWriter, r *http.Request) {
+ traceID := r.URL.Query().Get("traceId")
+ manager, _ := trace.Load(traceID)
+
+ updates, _ := manager.Subscribe()
+
+ w.Header().Set("Content-Type", "text/event-stream")
+ w.Header().Set("Cache-Control", "no-cache")
+
+ for update := range updates {
+ json.NewEncoder(w).Encode(update)
+ if f, ok := w.(http.Flusher); ok {
+ f.Flush()
+ }
+ }
+}
+```
+
+## Best Practices
+
+1. **Always handle errors**: Check errors from all operations
+2. **Release resources**: Call `Release()` when done or use defer
+3. **Complete traces**: Always call `MarkComplete()` when finished
+4. **Use context**: Pass context with timeout for long-running operations
+5. **Buffer channels**: Subscription channels are buffered (100), handle updates promptly
+6. **Unique IDs**: Let the system generate trace IDs for uniqueness
+7. **Metadata**: Use metadata for custom fields and debugging info
+
+## License
+
+Copyright (c) 2025 YaoApp
diff --git a/agent/trace/local/driver.go b/agent/trace/local/driver.go
new file mode 100644
index 00000000..f25704d5
--- /dev/null
+++ b/agent/trace/local/driver.go
@@ -0,0 +1,143 @@
+package local
+
+import (
+ "context"
+
+ "github.com/yaoapp/yao/agent/trace/types"
+)
+
+// Driver the local disk storage driver implementation
+type Driver struct {
+ basePath string // Base directory for storing trace files
+}
+
+// New creates a new local driver
+func New(basePath string) (*Driver, error) {
+ // TODO: Implement initialization (create directories, etc.)
+ return &Driver{
+ basePath: basePath,
+ }, nil
+}
+
+// SaveNode persists a node to disk
+func (d *Driver) SaveNode(ctx context.Context, traceID string, node *types.TraceNode) error {
+ // TODO: Implement disk save
+ // File path: {basePath}/{YYYYMMDD}/{traceID}/nodes/{nodeID}.json
+ return nil
+}
+
+// LoadNode loads a node from disk
+func (d *Driver) LoadNode(ctx context.Context, traceID string, nodeID string) (*types.TraceNode, error) {
+ // TODO: Implement disk load
+ return nil, nil
+}
+
+// LoadTrace loads the entire trace tree from disk
+func (d *Driver) LoadTrace(ctx context.Context, traceID string) (*types.TraceNode, error) {
+ // TODO: Implement disk load trace
+ // File path: {basePath}/{YYYYMMDD}/{traceID}/trace.json
+ return nil, nil
+}
+
+// SaveSpace persists a space to disk
+func (d *Driver) SaveSpace(ctx context.Context, traceID string, space *types.TraceSpace) error {
+ // TODO: Implement disk save space
+ // File path: {basePath}/{YYYYMMDD}/{traceID}/spaces/{spaceID}.json
+ return nil
+}
+
+// LoadSpace loads a space from disk
+func (d *Driver) LoadSpace(ctx context.Context, traceID string, spaceID string) (*types.TraceSpace, error) {
+ // TODO: Implement disk load space
+ return nil, nil
+}
+
+// DeleteSpace removes a space from disk
+func (d *Driver) DeleteSpace(ctx context.Context, traceID string, spaceID string) error {
+ // TODO: Implement disk delete space
+ return nil
+}
+
+// ListSpaces lists all space IDs for a trace from disk
+func (d *Driver) ListSpaces(ctx context.Context, traceID string) ([]string, error) {
+ // TODO: Implement disk list spaces
+ return nil, nil
+}
+
+// SetSpaceKey stores a value by key in a space
+func (d *Driver) SetSpaceKey(ctx context.Context, traceID, spaceID, key string, value any) error {
+ // TODO: Implement disk set space key
+ // File path: {basePath}/{YYYYMMDD}/{traceID}/spaces/{spaceID}/data.json
+ return nil
+}
+
+// GetSpaceKey retrieves a value by key from a space
+func (d *Driver) GetSpaceKey(ctx context.Context, traceID, spaceID, key string) (any, error) {
+ // TODO: Implement disk get space key
+ return nil, nil
+}
+
+// HasSpaceKey checks if a key exists in a space
+func (d *Driver) HasSpaceKey(ctx context.Context, traceID, spaceID, key string) bool {
+ // TODO: Implement disk has space key
+ return false
+}
+
+// DeleteSpaceKey removes a key-value pair from a space
+func (d *Driver) DeleteSpaceKey(ctx context.Context, traceID, spaceID, key string) error {
+ // TODO: Implement disk delete space key
+ return nil
+}
+
+// ClearSpaceKeys removes all key-value pairs from a space
+func (d *Driver) ClearSpaceKeys(ctx context.Context, traceID, spaceID string) error {
+ // TODO: Implement disk clear space keys
+ return nil
+}
+
+// ListSpaceKeys returns all keys in a space
+func (d *Driver) ListSpaceKeys(ctx context.Context, traceID, spaceID string) ([]string, error) {
+ // TODO: Implement disk list space keys
+ return nil, nil
+}
+
+// SaveLog appends a log entry to disk
+func (d *Driver) SaveLog(ctx context.Context, traceID string, log *types.TraceLog) error {
+ // TODO: Implement disk save log
+ // File path: {basePath}/{YYYYMMDD}/{traceID}/logs/{nodeID}.jsonl (append mode)
+ return nil
+}
+
+// LoadLogs loads all logs for a trace or specific node from disk
+func (d *Driver) LoadLogs(ctx context.Context, traceID string, nodeID string) ([]*types.TraceLog, error) {
+ // TODO: Implement disk load logs
+ // If nodeID is empty, load all logs
+ // If nodeID provided, load logs for that node only
+ return nil, nil
+}
+
+// SaveTraceInfo persists trace metadata to disk
+func (d *Driver) SaveTraceInfo(ctx context.Context, info *types.TraceInfo) error {
+ // TODO: Implement disk save trace info
+ // File path: {basePath}/{YYYYMMDD}/{traceID}/trace_info.json
+ return nil
+}
+
+// LoadTraceInfo loads trace metadata from disk
+func (d *Driver) LoadTraceInfo(ctx context.Context, traceID string) (*types.TraceInfo, error) {
+ // TODO: Implement disk load trace info
+ return nil, nil
+}
+
+// DeleteTrace removes entire trace from disk
+func (d *Driver) DeleteTrace(ctx context.Context, traceID string) error {
+ // TODO: Implement disk delete trace
+ // Delete directory: {basePath}/{YYYYMMDD}/{traceID}/
+ return nil
+}
+
+// Close closes the local driver
+func (d *Driver) Close() error {
+ // TODO: Implement cleanup if needed
+ return nil
+}
diff --git a/agent/trace/manager.go b/agent/trace/manager.go
new file mode 100644
index 00000000..0a4341c7
--- /dev/null
+++ b/agent/trace/manager.go
@@ -0,0 +1,742 @@
+package trace
+
+import (
+ "context"
+ "fmt"
+ "sync"
+ "time"
+
+ gonanoid "github.com/matoous/go-nanoid/v2"
+ "github.com/yaoapp/yao/agent/trace/types"
+)
+
+// manager implements the Manager interface with unified business logic
+type manager struct {
+ ctx context.Context
+ traceID string
+ driver types.Driver
+ rootNode *types.TraceNode
+ currentNodes []*types.TraceNode
+ spaces map[string]*types.TraceSpace
+ mu sync.RWMutex // Protects currentNodes and spaces
+
+ // Subscription mechanism
+ updates []*types.TraceUpdate // Update history (all events)
+ updatesMu sync.RWMutex // Protects updates
+ subscribers map[string]chan *types.TraceUpdate // Active subscribers
+ subMu sync.RWMutex // Protects subscribers
+ completed bool // Trace completion status
+}
+
+// NewManager creates a new trace manager instance
+func NewManager(ctx context.Context, traceID string, driver types.Driver) (types.Manager, error) {
+ // Create root node
+ now := time.Now().Unix()
+ rootNode := &types.TraceNode{
+ ID: genNodeID(),
+ ParentID: "",
+ Children: []*types.TraceNode{},
+ Status: types.StatusRunning,
+ CreatedAt: now,
+ StartTime: now,
+ UpdatedAt: now,
+ TraceNodeOption: types.TraceNodeOption{
+ Label: "Root",
+ Icon: "root",
+ },
+ }
+
+ // Save root node
+ if err := driver.SaveNode(ctx, traceID, rootNode); err != nil {
+ return nil, fmt.Errorf("failed to save root node: %w", err)
+ }
+
+ m := &manager{
+ ctx: ctx,
+ traceID: traceID,
+ driver: driver,
+ rootNode: rootNode,
+ currentNodes: []*types.TraceNode{rootNode},
+ spaces: make(map[string]*types.TraceSpace),
+ updates: make([]*types.TraceUpdate, 0, 100),
+ subscribers: make(map[string]chan *types.TraceUpdate),
+ completed: false,
+ }
+
+ // Broadcast init event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeInit,
+ TraceID: traceID,
+ Timestamp: now,
+ Data: types.NewTraceInitData(traceID, rootNode),
+ })
+
+ // Broadcast root node start event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeNodeStart,
+ TraceID: traceID,
+ NodeID: rootNode.ID,
+ Timestamp: now,
+ Data: rootNode.ToStartData(),
+ })
+
+ return m, nil
+}
+
+// genNodeID generates a unique node ID
+func genNodeID() string {
+ id, _ := gonanoid.Generate("0123456789abcdefghijklmnopqrstuvwxyz", 12)
+ return id
+}
+
+// checkContext checks if context is cancelled
+func (m *manager) checkContext() error {
+ select {
+ case <-m.ctx.Done():
+ return m.ctx.Err()
+ default:
+ return nil
+ }
+}
+
+// newNode creates a node instance that broadcasts events (for external use)
+func (m *manager) newNode(data *types.TraceNode) types.Node {
+ return &node{
+ manager: m,
+ data: data,
+ }
+}
+
+// Helper functions for thread-safe access
+
+// getCurrentNodes returns a copy of current nodes (thread-safe read)
+func (m *manager) getCurrentNodes() []*types.TraceNode {
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+ nodes := make([]*types.TraceNode, len(m.currentNodes))
+ copy(nodes, m.currentNodes)
+ return nodes
+}
+
+// getSpace returns a space by ID (thread-safe read)
+func (m *manager) getSpace(id string) (*types.TraceSpace, bool) {
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+ space, ok := m.spaces[id]
+ return space, ok
+}
+
+// setSpace stores a space (thread-safe write)
+func (m *manager) setSpace(id string, space *types.TraceSpace) {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ m.spaces[id] = space
+}
+
+// deleteSpace removes a space (thread-safe write)
+func (m *manager) deleteSpace(id string) {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ delete(m.spaces, id)
+}
+
+// getAllSpaces returns all spaces (thread-safe read)
+func (m *manager) getAllSpaces() []*types.TraceSpace {
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+ spaces := make([]*types.TraceSpace, 0, len(m.spaces))
+ for _, space := range m.spaces {
+ spaces = append(spaces, space)
+ }
+ return spaces
+}
+
+// Add creates next sequential node - auto-joins if currently in parallel state
+func (m *manager) Add(input types.TraceInput, option types.TraceNodeOption) (types.Node, error) {
+ if err := m.checkContext(); err != nil {
+ return nil, err
+ }
+
+ m.mu.Lock()
+ defer m.mu.Unlock()
+
+ now := time.Now().Unix()
+
+ // If in parallel state (multiple current nodes), auto-join first
+ var parentNode *types.TraceNode
+ if len(m.currentNodes) > 1 {
+ // Auto-join: create join node
+ parentNode = &types.TraceNode{
+ ID: genNodeID(),
+ ParentID: m.currentNodes[0].ParentID, // Same parent as parallel nodes
+ Children: []*types.TraceNode{},
+ Status: types.StatusCompleted,
+ CreatedAt: now,
+ StartTime: now,
+ EndTime: now,
+ UpdatedAt: now,
+ TraceNodeOption: types.TraceNodeOption{Label: "Join", Icon: "join"},
+ }
+ // Save join node
+ if err := m.driver.SaveNode(m.ctx, m.traceID, parentNode); err != nil {
+ return nil, err
+ }
+ } else {
+ parentNode = m.currentNodes[0]
+ }
+
+ // Create new node data
+ newNodeData := &types.TraceNode{
+ ID: genNodeID(),
+ ParentID: parentNode.ID,
+ Children: []*types.TraceNode{},
+ TraceNodeOption: option,
+ Status: types.StatusRunning,
+ Input: input,
+ CreatedAt: now,
+ StartTime: now,
+ UpdatedAt: now,
+ }
+
+ // Add to parent's children
+ parentNode.Children = append(parentNode.Children, newNodeData)
+
+ // Save nodes
+ if err := m.driver.SaveNode(m.ctx, m.traceID, newNodeData); err != nil {
+ return nil, err
+ }
+ if err := m.driver.SaveNode(m.ctx, m.traceID, parentNode); err != nil {
+ return nil, err
+ }
+
+ // Set as current node
+ m.currentNodes = []*types.TraceNode{newNodeData}
+
+ // Broadcast node start event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeNodeStart,
+ TraceID: m.traceID,
+ NodeID: newNodeData.ID,
+ Timestamp: now,
+ Data: newNodeData.ToStartData(),
+ })
+
+ // Return Node interface
+ return &node{
+ manager: m,
+ data: newNodeData,
+ }, nil
+}
+
+// Parallel creates multiple concurrent child nodes, returns Node interfaces for direct control
+func (m *manager) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node, error) {
+ if err := m.checkContext(); err != nil {
+ return nil, err
+ }
+
+ m.mu.Lock()
+ defer m.mu.Unlock()
+
+ now := time.Now().Unix()
+ parentNode := m.currentNodes[0]
+ nodeData := make([]*types.TraceNode, 0, len(parallelInputs))
+ nodeInterfaces := make([]types.Node, 0, len(parallelInputs))
+
+ // Create multiple child nodes
+ for _, input := range parallelInputs {
+ data := &types.TraceNode{
+ ID: genNodeID(),
+ ParentID: parentNode.ID,
+ Children: []*types.TraceNode{},
+ TraceNodeOption: input.Option,
+ Status: types.StatusRunning,
+ Input: input.Input,
+ CreatedAt: now,
+ StartTime: now,
+ UpdatedAt: now,
+ }
+ nodeData = append(nodeData, data)
+ parentNode.Children = append(parentNode.Children, data)
+
+ // Save node
+ if err := m.driver.SaveNode(m.ctx, m.traceID, data); err != nil {
+ return nil, err
+ }
+
+ // Create Node interface wrapper
+ nodeInterfaces = append(nodeInterfaces, &node{
+ manager: m,
+ data: data,
+ })
+ }
+
+ // Save parent node
+ if err := m.driver.SaveNode(m.ctx, m.traceID, parentNode); err != nil {
+ return nil, err
+ }
+
+ // Set all as current nodes (parallel state)
+ m.currentNodes = nodeData
+
+ // Broadcast parallel nodes as batch (frontend supports data.nodes[])
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeNodeStart,
+ TraceID: m.traceID,
+ Timestamp: now,
+ Data: types.NodesToStartData(nodeData),
+ })
+
+ return nodeInterfaces, nil
+}
+
+// Info logs info message to current node(s)
+func (m *manager) Info(format string, args ...any) types.Manager {
+ m.log("info", format, args...)
+ return m
+}
+
+// Debug logs debug message to current node(s)
+func (m *manager) Debug(format string, args ...any) types.Manager {
+ m.log("debug", format, args...)
+ return m
+}
+
+// Error logs error message to current node(s)
+func (m *manager) Error(format string, args ...any) types.Manager {
+ m.log("error", format, args...)
+ return m
+}
+
+// Warn logs warning message to current node(s)
+func (m *manager) Warn(format string, args ...any) types.Manager {
+ m.log("warn", format, args...)
+ return m
+}
+
+// log helper method to log messages
+func (m *manager) log(level string, format string, args ...any) {
+ message := fmt.Sprintf(format, args...)
+ now := time.Now().Unix()
+
+ // Get current nodes safely
+ nodes := m.getCurrentNodes()
+
+ // Log to all current nodes
+ for _, node := range nodes {
+ log := &types.TraceLog{
+ Timestamp: now,
+ Level: level,
+ Message: message,
+ NodeID: node.ID,
+ }
+ // Save log (ignore errors for non-critical logging)
+ _ = m.driver.SaveLog(m.ctx, m.traceID, log)
+
+ // Broadcast log event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeLogAdded,
+ TraceID: m.traceID,
+ NodeID: node.ID,
+ Timestamp: now,
+ Data: log,
+ })
+ }
+}
+
+// SetOutput sets output for current node(s)
+func (m *manager) SetOutput(output types.TraceOutput) error {
+ if err := m.checkContext(); err != nil {
+ return err
+ }
+
+ now := time.Now().Unix()
+ nodes := m.getCurrentNodes()
+ for _, node := range nodes {
+ node.Output = output
+ node.UpdatedAt = now
+ if err := m.driver.SaveNode(m.ctx, m.traceID, node); err != nil {
+ return err
+ }
+
+ // Broadcast node update event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeNodeUpdated,
+ TraceID: m.traceID,
+ NodeID: node.ID,
+ Timestamp: now,
+ Data: node,
+ })
+ }
+ return nil
+}
+
+// SetMetadata sets metadata for current node(s)
+func (m *manager) SetMetadata(key string, value any) error {
+ if err := m.checkContext(); err != nil {
+ return err
+ }
+
+ now := time.Now().Unix()
+ nodes := m.getCurrentNodes()
+ for _, node := range nodes {
+ if node.Metadata == nil {
+ node.Metadata = make(map[string]any)
+ }
+ node.Metadata[key] = value
+ node.UpdatedAt = now
+ if err := m.driver.SaveNode(m.ctx, m.traceID, node); err != nil {
+ return err
+ }
+
+ // Broadcast node update event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeNodeUpdated,
+ TraceID: m.traceID,
+ NodeID: node.ID,
+ Timestamp: now,
+ Data: node.ToStartData(), // Send updated node
+ })
+ }
+ return nil
+}
+
+// Complete marks current node(s) as completed
+// Optional output parameter: if provided, sets the output before completing
+func (m *manager) Complete(output ...types.TraceOutput) error {
+ if err := m.checkContext(); err != nil {
+ return err
+ }
+
+ now := time.Now().Unix()
+ nodes := m.getCurrentNodes()
+
+ // Set output if provided
+ if len(output) > 0 {
+ for _, node := range nodes {
+ node.Output = output[0]
+ }
+ }
+
+ for _, node := range nodes {
+ node.Status = types.StatusCompleted
+ node.EndTime = now
+ node.UpdatedAt = now
+ if err := m.driver.SaveNode(m.ctx, m.traceID, node); err != nil {
+ return err
+ }
+
+ // Broadcast node complete event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeNodeComplete,
+ TraceID: m.traceID,
+ NodeID: node.ID,
+ Timestamp: now,
+ Data: node.ToCompleteData(),
+ })
+ }
+ return nil
+}
+
+// Fail marks current node(s) as failed
+func (m *manager) Fail(err error) error {
+ if ctxErr := m.checkContext(); ctxErr != nil {
+ return ctxErr
+ }
+
+ now := time.Now().Unix()
+ // Log error first
+ m.Error("Node failed: %v", err)
+
+ nodes := m.getCurrentNodes()
+ for _, node := range nodes {
+ node.Status = types.StatusFailed
+ node.EndTime = now
+ node.UpdatedAt = now
+ if saveErr := m.driver.SaveNode(m.ctx, m.traceID, node); saveErr != nil {
+ return saveErr
+ }
+
+ // Broadcast node failed event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeNodeFailed,
+ TraceID: m.traceID,
+ NodeID: node.ID,
+ Timestamp: now,
+ Data: &types.NodeFailedData{
+ NodeID: node.ID,
+ Status: "failed",
+ EndTime: now,
+ Duration: (node.EndTime - node.StartTime) * 1000, // Convert to milliseconds
+ Error: err.Error(),
+ },
+ })
+ }
+ return nil
+}
+
+// GetRootNode returns the root node
+func (m *manager) GetRootNode() (*types.TraceNode, error) {
+ return m.rootNode, nil
+}
+
+// GetNode returns a node by ID
+func (m *manager) GetNode(id string) (*types.TraceNode, error) {
+ return m.driver.LoadNode(m.ctx, m.traceID, id)
+}
+
+// GetCurrentNodes returns current active nodes
+func (m *manager) GetCurrentNodes() ([]*types.TraceNode, error) {
+ return m.getCurrentNodes(), nil
+}
+
+// MarkComplete marks the entire trace as completed
+func (m *manager) MarkComplete() error {
+ m.updatesMu.Lock()
+ if m.completed {
+ m.updatesMu.Unlock()
+ return nil // Already completed
+ }
+ m.completed = true
+ m.updatesMu.Unlock()
+
+ // Calculate total duration from root node
+ now := time.Now().Unix()
+ totalDuration := int64(0)
+ if m.rootNode != nil && m.rootNode.CreatedAt > 0 {
+ totalDuration = (now - m.rootNode.CreatedAt) * 1000 // Convert to milliseconds
+ }
+
+ // Broadcast completion event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeComplete,
+ TraceID: m.traceID,
+ Timestamp: now,
+ Data: types.NewTraceCompleteData(m.traceID, totalDuration),
+ })
+
+ return nil
+}
+
+// CreateSpace creates a new memory space
+func (m *manager) CreateSpace(option types.TraceSpaceOption) (*types.TraceSpace, error) {
+ if err := m.checkContext(); err != nil {
+ return nil, err
+ }
+
+ now := time.Now().Unix()
+
+ // Create space instance
+ space := &types.TraceSpace{
+ ID: genNodeID(), // Reuse node ID generator
+ TraceSpaceOption: option,
+ CreatedAt: now,
+ UpdatedAt: now,
+ }
+
+ // Save to driver
+ if err := m.driver.SaveSpace(m.ctx, m.traceID, space); err != nil {
+ return nil, err
+ }
+
+ // Cache in memory (thread-safe)
+ m.setSpace(space.ID, space)
+
+ // Broadcast space created event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeSpaceCreated,
+ TraceID: m.traceID,
+ SpaceID: space.ID,
+ Timestamp: now,
+ Data: space,
+ })
+
+ return space, nil
+}
+
+// GetSpace returns a space by ID
+func (m *manager) GetSpace(id string) (*types.TraceSpace, error) {
+ // Check cache first (thread-safe)
+ if space, ok := m.getSpace(id); ok {
+ return space, nil
+ }
+
+ // Load from driver
+ space, err := m.driver.LoadSpace(m.ctx, m.traceID, id)
+ if err != nil {
+ return nil, err
+ }
+
+ // Cache it (thread-safe)
+ if space != nil {
+ m.setSpace(id, space)
+ }
+
+ return space, nil
+}
+
+// HasSpace checks if a space exists
+func (m *manager) HasSpace(id string) bool {
+ // Check cache (thread-safe)
+ if _, ok := m.getSpace(id); ok {
+ return true
+ }
+
+ // Check in driver
+ space, _ := m.driver.LoadSpace(m.ctx, m.traceID, id)
+ return space != nil
+}
+
+// DeleteSpace deletes a space
+func (m *manager) DeleteSpace(id string) error {
+ if err := m.checkContext(); err != nil {
+ return err
+ }
+
+ now := time.Now().Unix()
+
+ // Remove from cache (thread-safe)
+ m.deleteSpace(id)
+
+ // Delete from driver
+ if err := m.driver.DeleteSpace(m.ctx, m.traceID, id); err != nil {
+ return err
+ }
+
+ // Broadcast space deleted event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeSpaceDeleted,
+ TraceID: m.traceID,
+ SpaceID: id,
+ Timestamp: now,
+ Data: types.NewSpaceDeletedData(id),
+ })
+
+ return nil
+}
+
+// ListSpaces returns all spaces
+func (m *manager) ListSpaces() []*types.TraceSpace {
+ // Load from driver to ensure we have all spaces
+ spaceIDs, err := m.driver.ListSpaces(m.ctx, m.traceID)
+ if err != nil {
+ // Fallback to cached spaces (thread-safe)
+ return m.getAllSpaces()
+ }
+
+ // Load all spaces
+ spaces := make([]*types.TraceSpace, 0, len(spaceIDs))
+ for _, id := range spaceIDs {
+ space, err := m.GetSpace(id) // Use GetSpace to leverage cache
+ if err == nil && space != nil {
+ spaces = append(spaces, space)
+ }
+ }
+
+ return spaces
+}
+
+// SetSpaceValue sets a value in a space and broadcasts memory_add event
+func (m *manager) SetSpaceValue(spaceID, key string, value any) error {
+ if err := m.checkContext(); err != nil {
+ return err
+ }
+
+ now := time.Now().Unix()
+
+ // Get space
+ space, err := m.GetSpace(spaceID)
+ if err != nil || space == nil {
+ return fmt.Errorf("space not found: %s", spaceID)
+ }
+
+ // Set value in driver
+ if err := m.driver.SetSpaceKey(m.ctx, m.traceID, spaceID, key, value); err != nil {
+ return err
+ }
+
+ // Update space timestamp
+ space.UpdatedAt = now
+ if err := m.driver.SaveSpace(m.ctx, m.traceID, space); err != nil {
+ return err
+ }
+
+ // Broadcast memory_add event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeMemoryAdd,
+ TraceID: m.traceID,
+ SpaceID: spaceID,
+ Timestamp: now,
+ Data: space.ToMemoryAddData(key, value, now),
+ })
+
+ return nil
+}
+
+// GetSpaceValue gets a value from a space
+func (m *manager) GetSpaceValue(spaceID, key string) (any, error) {
+ return m.driver.GetSpaceKey(m.ctx, m.traceID, spaceID, key)
+}
+
+// HasSpaceValue checks if a key exists in a space
+func (m *manager) HasSpaceValue(spaceID, key string) bool {
+ return m.driver.HasSpaceKey(m.ctx, m.traceID, spaceID, key)
+}
+
+// DeleteSpaceValue deletes a value from a space and broadcasts memory_delete event
+func (m *manager) DeleteSpaceValue(spaceID, key string) error {
+ if err := m.checkContext(); err != nil {
+ return err
+ }
+
+ now := time.Now().Unix()
+
+ // Delete value from driver
+ if err := m.driver.DeleteSpaceKey(m.ctx, m.traceID, spaceID, key); err != nil {
+ return err
+ }
+
+ // Broadcast memory_delete event
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeMemoryDelete,
+ TraceID: m.traceID,
+ SpaceID: spaceID,
+ Timestamp: now,
+ Data: types.NewMemoryDeleteData(spaceID, key),
+ })
+
+ return nil
+}
+
+// ClearSpaceValues clears all values from a space
+func (m *manager) ClearSpaceValues(spaceID string) error {
+ if err := m.checkContext(); err != nil {
+ return err
+ }
+
+ now := time.Now().Unix()
+
+ // Clear values from driver
+ if err := m.driver.ClearSpaceKeys(m.ctx, m.traceID, spaceID); err != nil {
+ return err
+ }
+
+ // Broadcast memory_delete event (for all keys)
+ m.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeMemoryDelete,
+ TraceID: m.traceID,
+ SpaceID: spaceID,
+ Timestamp: now,
+ Data: types.NewMemoryDeleteAllData(spaceID),
+ })
+
+ return nil
+}
+
+// ListSpaceKeys returns all keys in a space
+func (m *manager) ListSpaceKeys(spaceID string) []string {
+ keys, err := m.driver.ListSpaceKeys(m.ctx, m.traceID, spaceID)
+ if err != nil {
+ return nil
+ }
+ return keys
+}
diff --git a/agent/trace/node.go b/agent/trace/node.go
new file mode 100644
index 00000000..066ed612
--- /dev/null
+++ b/agent/trace/node.go
@@ -0,0 +1,266 @@
+package trace
+
+import (
+ "fmt"
+ "time"
+
+ "github.com/yaoapp/yao/agent/trace/types"
+)
+
+// node implements the Node interface for custom node operations
+type node struct {
+ manager *manager
+ data *types.TraceNode
+}
+
+// Info logs info message (public method, broadcasts event)
+func (n *node) Info(format string, args ...any) types.Node {
+ n.logWithBroadcast("info", format, args...)
+ return n
+}
+
+// Debug logs debug message (public method, broadcasts event)
+func (n *node) Debug(format string, args ...any) types.Node {
+ n.logWithBroadcast("debug", format, args...)
+ return n
+}
+
+// Error logs error message (public method, broadcasts event)
+func (n *node) Error(format string, args ...any) types.Node {
+ n.logWithBroadcast("error", format, args...)
+ return n
+}
+
+// Warn logs warning message (public method, broadcasts event)
+func (n *node) Warn(format string, args ...any) types.Node {
+ n.logWithBroadcast("warn", format, args...)
+ return n
+}
+
+// logWithBroadcast logs and broadcasts event (for external calls)
+func (n *node) logWithBroadcast(level string, format string, args ...any) {
+ log := n.log(level, format, args...)
+
+ // Broadcast event
+ n.manager.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeLogAdded,
+ TraceID: n.manager.traceID,
+ NodeID: n.data.ID,
+ Timestamp: log.Timestamp,
+ Data: log,
+ })
+}
+
+// log logs without broadcasting (for internal Manager calls)
+func (n *node) log(level string, format string, args ...any) *types.TraceLog {
+ message := fmt.Sprintf(format, args...)
+ log := &types.TraceLog{
+ Timestamp: time.Now().Unix(),
+ Level: level,
+ Message: message,
+ NodeID: n.data.ID,
+ }
+ // Save log (ignore errors for non-critical logging)
+ _ = n.manager.driver.SaveLog(n.manager.ctx, n.manager.traceID, log)
+ return log
+}
+
+// Add creates next sequential node
+func (n *node) Add(input types.TraceInput, option types.TraceNodeOption) (types.Node, error) {
+ now := time.Now().Unix()
+
+ // Create child node data
+ childNodeData := &types.TraceNode{
+ ID: genNodeID(),
+ ParentID: n.data.ID,
+ Children: []*types.TraceNode{},
+ TraceNodeOption: option,
+ Status: types.StatusRunning,
+ Input: input,
+ CreatedAt: now,
+ StartTime: now,
+ UpdatedAt: now,
+ }
+
+ // Add to parent's children
+ n.data.Children = append(n.data.Children, childNodeData)
+
+ // Save both nodes
+ if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, childNodeData); err != nil {
+ return nil, err
+ }
+ if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data); err != nil {
+ return nil, err
+ }
+
+ // Return Node interface
+ return &node{
+ manager: n.manager,
+ data: childNodeData,
+ }, nil
+}
+
+// Parallel creates multiple concurrent child nodes
+func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node, error) {
+ now := time.Now().Unix()
+ nodeInterfaces := make([]types.Node, 0, len(parallelInputs))
+
+ // Create multiple child nodes
+ for _, input := range parallelInputs {
+ childNodeData := &types.TraceNode{
+ ID: genNodeID(),
+ ParentID: n.data.ID,
+ Children: []*types.TraceNode{},
+ TraceNodeOption: input.Option,
+ Status: types.StatusRunning,
+ Input: input.Input,
+ CreatedAt: now,
+ StartTime: now,
+ UpdatedAt: now,
+ }
+ n.data.Children = append(n.data.Children, childNodeData)
+
+ // Save node
+ if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, childNodeData); err != nil {
+ return nil, err
+ }
+
+ // Create Node interface wrapper
+ nodeInterfaces = append(nodeInterfaces, &node{
+ manager: n.manager,
+ data: childNodeData,
+ })
+ }
+
+ // Save parent node
+ if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data); err != nil {
+ return nil, err
+ }
+
+ return nodeInterfaces, nil
+}
+
+// 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()
+
+ // Create join node data
+ joinNodeData := &types.TraceNode{
+ ID: genNodeID(),
+ ParentID: n.data.ID,
+ Children: []*types.TraceNode{},
+ TraceNodeOption: option,
+ Status: types.StatusRunning,
+ Input: input,
+ CreatedAt: now,
+ StartTime: now,
+ UpdatedAt: now,
+ }
+
+ // Save join node
+ if err := n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, joinNodeData); err != nil {
+ return nil, err
+ }
+
+ // Return Node interface
+ return &node{
+ manager: n.manager,
+ data: joinNodeData,
+ }, nil
+}
+
+// ID returns the node ID
+func (n *node) ID() string {
+ return n.data.ID
+}
+
+// SetOutput sets the node output
+func (n *node) SetOutput(output types.TraceOutput) error {
+ n.data.Output = output
+ n.data.UpdatedAt = time.Now().Unix()
+ return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data)
+}
+
+// SetMetadata sets node metadata
+func (n *node) SetMetadata(key string, value any) error {
+ if n.data.Metadata == nil {
+ n.data.Metadata = make(map[string]any)
+ }
+ n.data.Metadata[key] = value
+ n.data.UpdatedAt = time.Now().Unix()
+ 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 = status
+ n.data.UpdatedAt = time.Now().Unix()
+ return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data)
+}
+
+// Complete marks the node as completed (public method, broadcasts event)
+// Optional output parameter: if provided, sets the output before completing
+func (n *node) Complete(output ...types.TraceOutput) error {
+ if err := n.complete(output...); err != nil {
+ return err
+ }
+
+ // Broadcast event
+ n.manager.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeNodeComplete,
+ TraceID: n.manager.traceID,
+ NodeID: n.data.ID,
+ Timestamp: n.data.EndTime,
+ Data: n.data.ToCompleteData(),
+ })
+
+ return nil
+}
+
+// complete marks as completed without broadcasting (for Manager calls)
+func (n *node) complete(output ...types.TraceOutput) error {
+ now := time.Now().Unix()
+
+ // Set output if provided
+ if len(output) > 0 {
+ n.data.Output = output[0]
+ }
+
+ n.data.Status = types.StatusCompleted
+ n.data.EndTime = now
+ n.data.UpdatedAt = now
+ return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data)
+}
+
+// Fail marks the node as failed (public method, broadcasts event)
+func (n *node) Fail(err error) error {
+ // Log error first
+ n.Error("Node failed: %v", err)
+
+ if saveErr := n.fail(err); saveErr != nil {
+ return saveErr
+ }
+
+ // Broadcast event
+ n.manager.addUpdate(&types.TraceUpdate{
+ Type: types.UpdateTypeNodeFailed,
+ TraceID: n.manager.traceID,
+ NodeID: n.data.ID,
+ Timestamp: n.data.EndTime,
+ Data: n.data.ToFailedData(err),
+ })
+
+ return nil
+}
+
+// fail marks as failed without broadcasting (for Manager calls)
+func (n *node) fail(err error) error {
+ now := time.Now().Unix()
+
+ // Update status
+ n.data.Status = types.StatusFailed
+ n.data.EndTime = now
+ n.data.UpdatedAt = now
+
+ return n.manager.driver.SaveNode(n.manager.ctx, n.manager.traceID, n.data)
+}
diff --git a/agent/trace/space.go b/agent/trace/space.go
new file mode 100644
index 00000000..f03fecf8
--- /dev/null
+++ b/agent/trace/space.go
@@ -0,0 +1,64 @@
+package trace
+
+import (
+ "context"
+
+ "github.com/yaoapp/yao/agent/trace/types"
+)
+
+// space implements the Space interface for custom space operations
+type space struct {
+ ctx context.Context
+ traceID string
+ data *types.TraceSpace
+ driver types.Driver
+}
+
+// NewSpace creates a new space instance
+func NewSpace(ctx context.Context, traceID string, data *types.TraceSpace, driver types.Driver) types.Space {
+ return &space{
+ ctx: ctx,
+ traceID: traceID,
+ data: data,
+ driver: driver,
+ }
+}
+
+// ID returns the space identifier
+func (s *space) ID() string {
+ return s.data.ID
+}
+
+// Set stores a value by key
+func (s *space) Set(key string, value any) error {
+ return s.driver.SetSpaceKey(s.ctx, s.traceID, s.data.ID, key, value)
+}
+
+// Get retrieves a value by key
+func (s *space) Get(key string) (any, error) {
+ return s.driver.GetSpaceKey(s.ctx, s.traceID, s.data.ID, key)
+}
+
+// Has checks if a key exists
+func (s *space) Has(key string) bool {
+ return s.driver.HasSpaceKey(s.ctx, s.traceID, s.data.ID, key)
+}
+
+// Delete removes a key-value pair
+func (s *space) Delete(key string) error {
+ return s.driver.DeleteSpaceKey(s.ctx, s.traceID, s.data.ID, key)
+}
+
+// Clear removes all key-value pairs
+func (s *space) Clear() error {
+ return s.driver.ClearSpaceKeys(s.ctx, s.traceID, s.data.ID)
+}
+
+// Keys returns all keys in the space
+func (s *space) Keys() []string {
+ keys, err := s.driver.ListSpaceKeys(s.ctx, s.traceID, s.data.ID)
+ if err != nil {
+ return nil
+ }
+ return keys
+}
diff --git a/agent/trace/store/driver.go b/agent/trace/store/driver.go
new file mode 100644
index 00000000..e5a30f58
--- /dev/null
+++ b/agent/trace/store/driver.go
@@ -0,0 +1,147 @@
+package store
+
+import (
+ "context"
+
+ "github.com/yaoapp/yao/agent/trace/types"
+)
+
+// Driver the gou store storage driver implementation
+type Driver struct {
+ storeName string // Store name in gou
+}
+
+// New creates a new store driver
+func New(storeName string) (*Driver, error) {
+ // TODO: Implement initialization (connect to gou store, etc.)
+ return &Driver{
+ storeName: storeName,
+ }, nil
+}
+
+// SaveNode persists a node to store
+func (d *Driver) SaveNode(ctx context.Context, traceID string, node *types.TraceNode) error {
+ // TODO: Implement store save
+ // Key: trace:{traceID}:node:{nodeID}
+ return nil
+}
+
+// LoadNode loads a node from store
+func (d *Driver) LoadNode(ctx context.Context, traceID string, nodeID string) (*types.TraceNode, error) {
+ // TODO: Implement store load
+ return nil, nil
+}
+
+// LoadTrace loads the entire trace tree from store
+func (d *Driver) LoadTrace(ctx context.Context, traceID string) (*types.TraceNode, error) {
+ // TODO: Implement store load trace
+ // Key: trace:{traceID}
+ return nil, nil
+}
+
+// SaveSpace persists a space to store
+func (d *Driver) SaveSpace(ctx context.Context, traceID string, space *types.TraceSpace) error {
+ // TODO: Implement store save space
+ // Key: trace:{traceID}:space:{spaceID}
+ return nil
+}
+
+// LoadSpace loads a space from store
+func (d *Driver) LoadSpace(ctx context.Context, traceID string, spaceID string) (*types.TraceSpace, error) {
+ // TODO: Implement store load space
+ return nil, nil
+}
+
+// DeleteSpace removes a space from store
+func (d *Driver) DeleteSpace(ctx context.Context, traceID string, spaceID string) error {
+ // TODO: Implement store delete space
+ return nil
+}
+
+// ListSpaces lists all space IDs for a trace from store
+func (d *Driver) ListSpaces(ctx context.Context, traceID string) ([]string, error) {
+ // TODO: Implement store list spaces
+ // Use pattern matching: trace:{traceID}:space:*
+ return nil, nil
+}
+
+// SetSpaceKey stores a value by key in a space
+func (d *Driver) SetSpaceKey(ctx context.Context, traceID, spaceID, key string, value any) error {
+ // TODO: Implement store set space key
+ // Key: trace:{traceID}:space:{spaceID}:key:{key}
+ return nil
+}
+
+// GetSpaceKey retrieves a value by key from a space
+func (d *Driver) GetSpaceKey(ctx context.Context, traceID, spaceID, key string) (any, error) {
+ // TODO: Implement store get space key
+ return nil, nil
+}
+
+// HasSpaceKey checks if a key exists in a space
+func (d *Driver) HasSpaceKey(ctx context.Context, traceID, spaceID, key string) bool {
+ // TODO: Implement store has space key
+ return false
+}
+
+// DeleteSpaceKey removes a key-value pair from a space
+func (d *Driver) DeleteSpaceKey(ctx context.Context, traceID, spaceID, key string) error {
+ // TODO: Implement store delete space key
+ return nil
+}
+
+// ClearSpaceKeys removes all key-value pairs from a space
+func (d *Driver) ClearSpaceKeys(ctx context.Context, traceID, spaceID string) error {
+ // TODO: Implement store clear space keys
+ // Delete keys: trace:{traceID}:space:{spaceID}:key:*
+ return nil
+}
+
+// ListSpaceKeys returns all keys in a space
+func (d *Driver) ListSpaceKeys(ctx context.Context, traceID, spaceID string) ([]string, error) {
+ // TODO: Implement store list space keys
+ // Use pattern matching: trace:{traceID}:space:{spaceID}:key:*
+ return nil, nil
+}
+
+// SaveLog appends a log entry to store
+func (d *Driver) SaveLog(ctx context.Context, traceID string, log *types.TraceLog) error {
+ // TODO: Implement store save log
+ // Key: trace:{traceID}:logs:{nodeID} (list type, append)
+ return nil
+}
+
+// LoadLogs loads all logs for a trace or specific node from store
+func (d *Driver) LoadLogs(ctx context.Context, traceID string, nodeID string) ([]*types.TraceLog, error) {
+ // TODO: Implement store load logs
+ // If nodeID is empty, load all logs from trace:{traceID}:logs:*
+ // If nodeID provided, load from trace:{traceID}:logs:{nodeID}
+ return nil, nil
+}
+
+// SaveTraceInfo persists trace metadata to store
+func (d *Driver) SaveTraceInfo(ctx context.Context, info *types.TraceInfo) error {
+ // TODO: Implement store save trace info
+ // Key: trace:{traceID}:info
+ return nil
+}
+
+// LoadTraceInfo loads trace metadata from store
+func (d *Driver) LoadTraceInfo(ctx context.Context, traceID string) (*types.TraceInfo, error) {
+ // TODO: Implement store load trace info
+ return nil, nil
+}
+
+// DeleteTrace removes entire trace from store
+func (d *Driver) DeleteTrace(ctx context.Context, traceID string) error {
+ // TODO: Implement store delete trace
+ // Delete keys: trace:{traceID}* (including all spaces, nodes, and logs)
+ return nil
+}
+
+// Close closes the store driver
+func (d *Driver) Close() error {
+ // TODO: Implement cleanup if needed
+ return nil
+}
+
diff --git a/agent/trace/subscription.go b/agent/trace/subscription.go
new file mode 100644
index 00000000..8d4c6ec2
--- /dev/null
+++ b/agent/trace/subscription.go
@@ -0,0 +1,106 @@
+package trace
+
+import (
+ "github.com/yaoapp/yao/agent/trace/types"
+)
+
+// Subscription Operations
+
+// addUpdate adds an update to history and broadcasts to subscribers
+func (m *manager) addUpdate(update *types.TraceUpdate) {
+ // Add to history
+ m.updatesMu.Lock()
+ m.updates = append(m.updates, update)
+ m.updatesMu.Unlock()
+
+ // Broadcast to real-time subscribers (non-blocking, in goroutine)
+ go m.broadcast(update)
+}
+
+// broadcast sends update to all active subscribers (non-blocking)
+func (m *manager) broadcast(update *types.TraceUpdate) {
+ m.subMu.RLock()
+ defer m.subMu.RUnlock()
+
+ for _, ch := range m.subscribers {
+ select {
+ case ch <- update:
+ // Sent successfully
+ default:
+ // Channel full, skip (or could log warning)
+ }
+ }
+}
+
+// Subscribe subscribes to all trace updates (replay history + real-time)
+func (m *manager) Subscribe() (<-chan *types.TraceUpdate, error) {
+ return m.SubscribeFrom(0)
+}
+
+// SubscribeFrom subscribes from a specific timestamp (for resume)
+func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error) {
+ // Create subscriber channel with buffer
+ ch := make(chan *types.TraceUpdate, 100)
+ subID := genNodeID()
+
+ // Register subscriber
+ m.subMu.Lock()
+ m.subscribers[subID] = ch
+ m.subMu.Unlock()
+
+ // Start replay and streaming goroutine
+ go m.replayAndStream(ch, subID, since)
+
+ return ch, nil
+}
+
+// replayAndStream replays history then streams real-time updates
+func (m *manager) replayAndStream(ch chan *types.TraceUpdate, subID string, since int64) {
+ defer func() {
+ // Close channel and cleanup subscriber
+ close(ch)
+ m.subMu.Lock()
+ delete(m.subscribers, subID)
+ m.subMu.Unlock()
+ }()
+
+ // Step 1: Replay history
+ m.updatesMu.RLock()
+ history := make([]*types.TraceUpdate, 0)
+ for _, update := range m.updates {
+ if update.Timestamp > since {
+ history = append(history, update)
+ }
+ }
+ isCompleted := m.completed
+ m.updatesMu.RUnlock()
+
+ // Send history in order
+ for _, update := range history {
+ select {
+ case ch <- update:
+ // Sent successfully
+ // Optional: add small delay to control replay speed
+ // time.Sleep(10 * time.Millisecond)
+ case <-m.ctx.Done():
+ // Context cancelled, stop
+ return
+ }
+ }
+
+ // Step 2: If already completed, exit
+ if isCompleted {
+ return
+ }
+
+ // Step 3: Wait for completion or context cancellation
+ // Real-time updates are sent by broadcast() method
+ <-m.ctx.Done()
+}
+
+// IsComplete checks if the trace is completed
+func (m *manager) IsComplete() bool {
+ m.updatesMu.RLock()
+ defer m.updatesMu.RUnlock()
+ return m.completed
+}
diff --git a/agent/trace/trace.go b/agent/trace/trace.go
index ea39374f..413e3f07 100644
--- a/agent/trace/trace.go
+++ b/agent/trace/trace.go
@@ -1,27 +1,367 @@
package trace
import (
+ "context"
"fmt"
+ "sync"
"time"
gonanoid "github.com/matoous/go-nanoid/v2"
+ "github.com/yaoapp/yao/agent/trace/local"
+ "github.com/yaoapp/yao/agent/trace/store"
+ "github.com/yaoapp/yao/agent/trace/types"
)
-// GenTraceID generate a new trace ID using NanoID algorithm
+// Driver types
+const (
+ Local = "local" // Local disk storage
+ Store = "store" // Gou store storage
+)
+
+// Global trace registry
+var (
+ registry = make(map[string]*types.TraceInfo)
+ registryMu sync.RWMutex
+)
+
+// getDriver creates a driver instance based on driver type and options
+func getDriver(driver string, options ...any) (types.Driver, error) {
+ var drv types.Driver
+ var err error
+
+ switch driver {
+ case Local:
+ basePath := "./traces" // default
+ if len(options) > 0 {
+ if path, ok := options[0].(string); ok {
+ basePath = path
+ }
+ }
+ drv, err = local.New(basePath)
+ if err != nil {
+ return nil, fmt.Errorf("failed to create local driver: %w", err)
+ }
+
+ case Store:
+ storeName := "trace" // default
+ if len(options) > 0 {
+ if name, ok := options[0].(string); ok {
+ storeName = name
+ }
+ }
+ drv, err = store.New(storeName)
+ if err != nil {
+ return nil, fmt.Errorf("failed to create store driver: %w", err)
+ }
+
+ default:
+ return nil, fmt.Errorf("unknown driver: %s", driver)
+ }
+
+ return drv, nil
+}
+
+// GenTraceID generate a new trace ID with date prefix format: YYYYMMDDnnnnnnnnnnnn
+// Format: 20251118123456789012 (8-digit date + 12-digit unique ID)
+// The date prefix enables directory-based storage organization (e.g., traces/20251118/)
// safe: optional parameter, reserved for future safe mode implementation (collision detection)
func GenTraceID(safe ...bool) string {
// TODO: Implement safe mode with collision detection when needed
- // For now, NanoID provides sufficient uniqueness without collision checking
- // URL-safe alphabet (no ambiguous characters like 0/O, 1/l/I)
- const alphabet = "1234567890"
- const length = 8 // 8 characters provides good balance of uniqueness and readability
+ now := time.Now()
+ // Date prefix: YYYYMMDD (8 digits)
+ prefix := now.Format("20060102")
- id, err := gonanoid.Generate(alphabet, length)
+ // Generate 12-digit unique suffix (timestamp in microseconds + random)
+ // Using timestamp ensures uniqueness within the same day
+ timestamp := fmt.Sprintf("%06d", now.Unix()%1000000) // 6 digits from timestamp
+
+ const alphabet = "0123456789"
+ const length = 6
+
+ random, err := gonanoid.Generate(alphabet, length)
if err != nil {
- // Fallback to timestamp-based ID if NanoID generation fails
- return fmt.Sprintf("%d", time.Now().UnixNano())
+ // Fallback to nanoseconds if NanoID generation fails
+ random = fmt.Sprintf("%06d", now.Nanosecond()%1000000)
}
- return id
+ return prefix + timestamp + random
+}
+
+// New creates a new trace manager with specified driver
+// Returns: traceID, manager, error
+// ctx: context for the trace manager
+// driver: Local or Store
+// 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()
+
+ // Handle nil option
+ if option == nil {
+ option = &types.TraceOption{}
+ }
+
+ // Generate ID if not provided
+ traceID := option.ID
+ if traceID == "" {
+ traceID = GenTraceID()
+ }
+
+ // Check if ID already exists (in registry or storage)
+ if IsLoaded(traceID) {
+ return "", nil, fmt.Errorf("trace ID already loaded in registry: %s", traceID)
+ }
+
+ // Create driver instance
+ drv, err := getDriver(driver, driverOptions...)
+ if err != nil {
+ return "", nil, err
+ }
+
+ // Check if exists in storage - if so, load it instead
+ exists, err := Exists(ctx, driver, traceID, driverOptions...)
+ if err != nil {
+ return "", nil, fmt.Errorf("failed to check trace existence: %w", err)
+ }
+ if exists {
+ // Trace exists in storage, load it
+ return LoadFromStorage(ctx, driver, traceID, driverOptions...)
+ }
+
+ // Create Manager instance with the driver
+ manager, err := NewManager(ctx, traceID, drv)
+ if err != nil {
+ return "", nil, fmt.Errorf("failed to create manager: %w", err)
+ }
+
+ // Create trace info
+ info := &types.TraceInfo{
+ ID: traceID,
+ Driver: driver,
+ Options: driverOptions,
+ Manager: manager,
+ CreatedAt: now,
+ UpdatedAt: now,
+ CreatedBy: option.CreatedBy,
+ UpdatedBy: option.CreatedBy, // Initially same as CreatedBy
+ TeamID: option.TeamID,
+ TenantID: option.TenantID,
+ Metadata: option.Metadata,
+ }
+
+ // Register trace in global registry
+ registryMu.Lock()
+ registry[traceID] = info
+ registryMu.Unlock()
+
+ // Persist trace info to driver
+ if err := drv.SaveTraceInfo(ctx, info); err != nil {
+ // If save fails, remove from registry and return error
+ registryMu.Lock()
+ delete(registry, traceID)
+ registryMu.Unlock()
+ return "", nil, fmt.Errorf("failed to save trace info: %w", err)
+ }
+
+ return traceID, manager, nil
+}
+
+// Load loads an existing trace by ID from the registry
+// Returns: manager, error
+// traceID: the trace ID to load
+func Load(traceID string) (types.Manager, error) {
+ registryMu.RLock()
+ info, exists := registry[traceID]
+ registryMu.RUnlock()
+
+ if !exists {
+ return nil, fmt.Errorf("trace not found in registry: %s (use LoadFromStorage to load from persistent storage)", traceID)
+ }
+
+ return info.Manager, nil
+}
+
+// LoadFromStorage loads a trace from persistent storage and activates it in registry
+// This is used to resume a trace that was previously created but not currently loaded
+// Returns: traceID, manager, error
+// ctx: context for the trace manager
+// driver: Local or Store (must match the driver used to create the trace)
+// traceID: the trace ID to load
+// driverOptions: driver-specific options (e.g., base path for local, store name for store)
+func LoadFromStorage(ctx context.Context, driver string, traceID string, driverOptions ...any) (string, types.Manager, error) {
+ // Check if already loaded
+ if IsLoaded(traceID) {
+ registryMu.RLock()
+ info := registry[traceID]
+ registryMu.RUnlock()
+ return traceID, info.Manager, nil
+ }
+
+ // Create driver instance
+ drv, err := getDriver(driver, driverOptions...)
+ if err != nil {
+ return "", nil, err
+ }
+
+ // Load trace info from storage
+ storedInfo, err := drv.LoadTraceInfo(ctx, traceID)
+ if err != nil {
+ drv.Close()
+ return "", nil, fmt.Errorf("failed to load trace info: %w", err)
+ }
+ if storedInfo == nil {
+ drv.Close()
+ return "", nil, fmt.Errorf("trace not found in storage: %s", traceID)
+ }
+
+ // Create Manager instance with the driver
+ // Note: We need to reconstruct the manager from stored data
+ // TODO: Implement proper restoration of manager state from storage
+ manager, err := NewManager(ctx, traceID, drv)
+ if err != nil {
+ drv.Close()
+ return "", nil, fmt.Errorf("failed to create manager: %w", err)
+ }
+
+ // Update stored info with new manager
+ storedInfo.Manager = manager
+ storedInfo.UpdatedAt = time.Now().Unix()
+
+ // Register in global registry
+ registryMu.Lock()
+ registry[traceID] = storedInfo
+ registryMu.Unlock()
+
+ return traceID, manager, nil
+}
+
+// GetInfo returns the trace metadata from storage
+// This function reads from persistent storage and can be used even if the trace is not in registry
+// ctx: context for the operation
+// driver: Local or Store (must match the driver used to create the trace)
+// traceID: the trace ID
+// options: driver-specific options (e.g., base path for local, store name for store)
+func GetInfo(ctx context.Context, driver string, traceID string, options ...any) (*types.TraceInfo, error) {
+ // Try registry first (if trace is active)
+ registryMu.RLock()
+ info, exists := registry[traceID]
+ registryMu.RUnlock()
+
+ if exists {
+ // Return a copy to prevent external modification
+ infoCopy := *info
+ infoCopy.Manager = nil // Don't expose manager in info
+ return &infoCopy, nil
+ }
+
+ // Not in registry, load from driver
+ drv, err := getDriver(driver, options...)
+ if err != nil {
+ return nil, err
+ }
+ defer drv.Close()
+
+ // Load from storage
+ storedInfo, err := drv.LoadTraceInfo(ctx, traceID)
+ if err != nil {
+ return nil, fmt.Errorf("failed to load trace info: %w", err)
+ }
+
+ if storedInfo == nil {
+ return nil, fmt.Errorf("trace not found: %s", traceID)
+ }
+
+ // Don't expose manager for stored info (manager is only available for active traces)
+ storedInfo.Manager = nil
+ return storedInfo, nil
+}
+
+// Release releases a trace from the registry and closes its resources
+// traceID: the trace ID to release
+func Release(traceID string) error {
+ registryMu.Lock()
+ _, exists := registry[traceID]
+ if exists {
+ delete(registry, traceID)
+ }
+ registryMu.Unlock()
+
+ if !exists {
+ return fmt.Errorf("trace not found in registry: %s", traceID)
+ }
+
+ // Close driver resources if manager has a close method
+ // (Currently manager doesn't expose driver, but driver has Close method)
+ // This is handled when the context is cancelled or program exits
+
+ return nil
+}
+
+// IsLoaded checks if a trace is loaded in the registry (active state)
+func IsLoaded(traceID string) bool {
+ registryMu.RLock()
+ defer registryMu.RUnlock()
+ _, exists := registry[traceID]
+ return exists
+}
+
+// Exists checks if a trace exists in persistent storage
+// ctx: context for the operation
+// driver: Local or Store (must match the driver used to create the trace)
+// traceID: the trace ID
+// options: driver-specific options (e.g., base path for local, store name for store)
+func Exists(ctx context.Context, driver string, traceID string, options ...any) (bool, error) {
+ // Check registry first (if loaded)
+ if IsLoaded(traceID) {
+ return true, nil
+ }
+
+ // Check persistent storage
+ drv, err := getDriver(driver, options...)
+ if err != nil {
+ return false, err
+ }
+ defer drv.Close()
+
+ // Try to load trace info from storage
+ info, err := drv.LoadTraceInfo(ctx, traceID)
+ if err != nil {
+ // If error is not found, return false; otherwise return error
+ return false, nil
+ }
+
+ return info != nil, nil
+}
+
+// List returns all active trace IDs in the registry
+func List() []string {
+ registryMu.RLock()
+ defer registryMu.RUnlock()
+
+ ids := make([]string, 0, len(registry))
+ for id := range registry {
+ ids = append(ids, id)
+ }
+ return ids
+}
+
+// Remove deletes a trace and all its associated data (nodes, spaces)
+// This is a destructive operation and cannot be undone
+// Automatically releases the trace from registry if it exists
+// driver: Local or Store (must match the driver used to create the trace)
+// traceID: the trace ID to remove
+// options: driver-specific options
+func Remove(ctx context.Context, driver string, traceID string, options ...any) error {
+ // Release from registry first (if exists)
+ _ = Release(traceID) // Ignore error if not in registry
+
+ // Create driver instance
+ drv, err := getDriver(driver, options...)
+ if err != nil {
+ return err
+ }
+ defer drv.Close()
+
+ return drv.DeleteTrace(ctx, traceID)
}
diff --git a/agent/trace/types/driver.go b/agent/trace/types/driver.go
new file mode 100644
index 00000000..2015cd7a
--- /dev/null
+++ b/agent/trace/types/driver.go
@@ -0,0 +1,73 @@
+package types
+
+import "context"
+
+// TraceLog represents a log entry
+type TraceLog struct {
+ Timestamp int64 // Log timestamp
+ Level string // Log level (info, debug, error, warn)
+ Message string // Log message
+ NodeID string // Node ID this log belongs to
+}
+
+// Driver defines the storage driver interface that providers must implement
+// Driver is only responsible for persistence operations, not business logic
+type Driver interface {
+ // SaveNode persists a node to storage
+ SaveNode(ctx context.Context, traceID string, node *TraceNode) error
+
+ // LoadNode loads a node from storage
+ LoadNode(ctx context.Context, traceID string, nodeID string) (*TraceNode, error)
+
+ // LoadTrace loads the entire trace tree from storage
+ LoadTrace(ctx context.Context, traceID string) (*TraceNode, error)
+
+ // SaveSpace persists a space to storage
+ SaveSpace(ctx context.Context, traceID string, space *TraceSpace) error
+
+ // LoadSpace loads a space from storage
+ LoadSpace(ctx context.Context, traceID string, spaceID string) (*TraceSpace, error)
+
+ // DeleteSpace removes a space from storage
+ DeleteSpace(ctx context.Context, traceID string, spaceID string) error
+
+ // ListSpaces lists all space IDs for a trace
+ ListSpaces(ctx context.Context, traceID string) ([]string, error)
+
+ // Space KV Operations
+ // SetSpaceKey stores a value by key in a space
+ SetSpaceKey(ctx context.Context, traceID, spaceID, key string, value any) error
+
+ // GetSpaceKey retrieves a value by key from a space
+ GetSpaceKey(ctx context.Context, traceID, spaceID, key string) (any, error)
+
+ // HasSpaceKey checks if a key exists in a space
+ HasSpaceKey(ctx context.Context, traceID, spaceID, key string) bool
+
+ // DeleteSpaceKey removes a key-value pair from a space
+ DeleteSpaceKey(ctx context.Context, traceID, spaceID, key string) error
+
+ // ClearSpaceKeys removes all key-value pairs from a space
+ ClearSpaceKeys(ctx context.Context, traceID, spaceID string) error
+
+ // ListSpaceKeys returns all keys in a space
+ ListSpaceKeys(ctx context.Context, traceID, spaceID string) ([]string, error)
+
+ // SaveLog appends a log entry to storage
+ SaveLog(ctx context.Context, traceID string, log *TraceLog) error
+
+ // LoadLogs loads all logs for a trace or specific node
+ LoadLogs(ctx context.Context, traceID string, nodeID string) ([]*TraceLog, error)
+
+ // SaveTraceInfo persists trace metadata to storage
+ SaveTraceInfo(ctx context.Context, info *TraceInfo) error
+
+ // LoadTraceInfo loads trace metadata from storage
+ LoadTraceInfo(ctx context.Context, traceID string) (*TraceInfo, error)
+
+ // DeleteTrace removes entire trace and all its data
+ DeleteTrace(ctx context.Context, traceID string) error
+
+ // Close closes the driver and releases resources
+ Close() error
+}
diff --git a/agent/trace/types/events.go b/agent/trace/types/events.go
new file mode 100644
index 00000000..f34b507a
--- /dev/null
+++ b/agent/trace/types/events.go
@@ -0,0 +1,97 @@
+package types
+
+// Helper functions and methods to create event data
+
+// ToStartData converts TraceNode to NodeStartData (single node)
+func (n *TraceNode) ToStartData() *NodeStartData {
+ return &NodeStartData{Node: n}
+}
+
+// NodesToStartData creates NodeStartData for multiple nodes (parallel operations)
+func NodesToStartData(nodes []*TraceNode) *NodeStartData {
+ return &NodeStartData{Nodes: nodes}
+}
+
+// ToCompleteData converts TraceNode to NodeCompleteData
+func (n *TraceNode) ToCompleteData() *NodeCompleteData {
+ return &NodeCompleteData{
+ NodeID: n.ID,
+ Status: "success",
+ EndTime: n.EndTime,
+ Duration: (n.EndTime - n.StartTime) * 1000, // Convert to milliseconds
+ Output: n.Output,
+ }
+}
+
+// ToFailedData converts TraceNode to NodeFailedData
+func (n *TraceNode) ToFailedData(err error) *NodeFailedData {
+ return &NodeFailedData{
+ NodeID: n.ID,
+ Status: "failed",
+ EndTime: n.EndTime,
+ Duration: (n.EndTime - n.StartTime) * 1000, // Convert to milliseconds
+ Error: err.Error(),
+ }
+}
+
+// ToMemoryAddData creates MemoryAddData for a space key-value operation
+func (s *TraceSpace) ToMemoryAddData(key string, value any, timestamp int64) *MemoryAddData {
+ item := MemoryItem{
+ ID: key,
+ Type: s.ID, // Space ID as type
+ Content: value,
+ Timestamp: timestamp,
+ }
+ // Use Label as title if available
+ if s.Label != "" {
+ item.Title = s.Label
+ }
+ return &MemoryAddData{
+ Type: s.ID,
+ Item: item,
+ }
+}
+
+// NewTraceInitData creates init event data
+func NewTraceInitData(traceID string, rootNode *TraceNode, agentName ...string) *TraceInitData {
+ data := &TraceInitData{
+ TraceID: traceID,
+ RootNode: rootNode,
+ }
+ if len(agentName) > 0 {
+ data.AgentName = agentName[0]
+ }
+ return data
+}
+
+// NewTraceCompleteData creates trace complete event data
+func NewTraceCompleteData(traceID string, totalDuration int64) *TraceCompleteData {
+ return &TraceCompleteData{
+ TraceID: traceID,
+ Status: "completed",
+ TotalDuration: totalDuration,
+ }
+}
+
+// NewSpaceDeletedData creates space deleted event data
+func NewSpaceDeletedData(spaceID string) *SpaceDeletedData {
+ return &SpaceDeletedData{
+ SpaceID: spaceID,
+ }
+}
+
+// NewMemoryDeleteData creates memory delete event data (single key)
+func NewMemoryDeleteData(spaceID, key string) *MemoryDeleteData {
+ return &MemoryDeleteData{
+ SpaceID: spaceID,
+ Key: key,
+ }
+}
+
+// NewMemoryDeleteAllData creates memory delete event data (all keys cleared)
+func NewMemoryDeleteAllData(spaceID string) *MemoryDeleteData {
+ return &MemoryDeleteData{
+ SpaceID: spaceID,
+ Cleared: true,
+ }
+}
diff --git a/agent/trace/types/interfaces.go b/agent/trace/types/interfaces.go
new file mode 100644
index 00000000..7bd0a763
--- /dev/null
+++ b/agent/trace/types/interfaces.go
@@ -0,0 +1,106 @@
+package types
+
+// Manager the trace manager interface
+// Manager automatically tracks current node(s) state, users don't need to manage nodes manually
+// Context is bound to Manager at creation time
+type Manager interface {
+
+ // Node Tree Operations - work on current node(s)
+ // Add creates next sequential node - auto-joins if currently in parallel state
+ Add(input TraceInput, option TraceNodeOption) (Node, error)
+ // Parallel creates multiple concurrent child nodes, returns Node interfaces for direct control
+ Parallel(parallelInputs []TraceParallelInput) ([]Node, error)
+
+ // Log Operations - log to current node(s) with chainable interface
+ Info(format string, args ...any) Manager
+ Debug(format string, args ...any) Manager
+ Error(format string, args ...any) Manager
+ Warn(format string, args ...any) Manager
+
+ // Node Status Operations - operate on current node(s)
+ SetOutput(output TraceOutput) error
+ SetMetadata(key string, value any) error
+ Complete(output ...TraceOutput) error // Optional output parameter
+ Fail(err error) error
+
+ // Query Operations
+ GetRootNode() (*TraceNode, error)
+ GetNode(id string) (*TraceNode, error)
+ GetCurrentNodes() ([]*TraceNode, error)
+
+ // Memory Space Operations
+ CreateSpace(option TraceSpaceOption) (*TraceSpace, error)
+ GetSpace(id string) (*TraceSpace, error)
+ HasSpace(id string) bool
+ DeleteSpace(id string) error
+ ListSpaces() []*TraceSpace
+
+ // Space Key-Value Operations (with automatic event broadcasting)
+ SetSpaceValue(spaceID, key string, value any) error
+ GetSpaceValue(spaceID, key string) (any, error)
+ HasSpaceValue(spaceID, key string) bool
+ DeleteSpaceValue(spaceID, key string) error
+ ClearSpaceValues(spaceID string) error
+ ListSpaceKeys(spaceID string) []string
+
+ // Trace Control Operations
+ // MarkComplete marks the entire trace as completed (sends trace_complete event)
+ MarkComplete() error
+
+ // Subscription Operations
+ // Subscribe subscribes to trace updates (replay history + real-time)
+ Subscribe() (<-chan *TraceUpdate, error)
+ // SubscribeFrom subscribes from a specific timestamp (for resume)
+ SubscribeFrom(since int64) (<-chan *TraceUpdate, error)
+ // IsComplete checks if the trace is completed
+ IsComplete() bool
+}
+
+// Node represents a trace node with operations for tree building and logging
+// Context is bound to Node at creation time
+type Node interface {
+ // Log Operations - chainable interface
+ Info(format string, args ...any) Node
+ Debug(format string, args ...any) Node
+ Error(format string, args ...any) Node
+ Warn(format string, args ...any) Node
+
+ // Node Tree Operations
+ Add(input TraceInput, option TraceNodeOption) (Node, error)
+ Parallel(parallelInputs []TraceParallelInput) ([]Node, error)
+ Join(nodes []*TraceNode, input TraceInput, option TraceNodeOption) (Node, error)
+
+ // Node Data Operations
+ ID() string
+ SetOutput(output TraceOutput) error
+ SetMetadata(key string, value any) error
+
+ // Node Status Operations
+ SetStatus(status string) error
+ Complete(output ...TraceOutput) error // Optional output parameter
+ Fail(err error) error
+}
+
+// Space represents a key-value storage space
+type Space interface {
+ // ID returns the space identifier
+ ID() string
+
+ // Set stores a value by key
+ Set(key string, value any) error
+
+ // Get retrieves a value by key
+ Get(key string) (any, error)
+
+ // Has checks if a key exists
+ Has(key string) bool
+
+ // Delete removes a key-value pair
+ Delete(key string) error
+
+ // Clear removes all key-value pairs
+ Clear() error
+
+ // Keys returns all keys in the space
+ Keys() []string
+}
diff --git a/agent/trace/types/types.go b/agent/trace/types/types.go
new file mode 100644
index 00000000..e06366b9
--- /dev/null
+++ b/agent/trace/types/types.go
@@ -0,0 +1,190 @@
+package types
+
+// Node status constants
+const (
+ StatusPending = "pending" // Node created but not started
+ StatusRunning = "running" // Node is currently executing
+ StatusCompleted = "completed" // Node finished successfully
+ StatusFailed = "failed" // Node failed with error
+ StatusSkipped = "skipped" // Node was skipped
+)
+
+// TraceNodeOption defines options for creating a node
+type TraceNodeOption struct {
+ Label string // Display label in UI
+ Icon string // Icon identifier
+ Description string // Node description
+ Metadata map[string]any // Additional metadata
+}
+
+// TraceSpaceOption defines options for creating a space
+type TraceSpaceOption struct {
+ Label string // Display label in UI
+ Icon string // Icon identifier
+ Description string // Space description
+ TTL int64 // Time to live in seconds (0 = no expiration) - for display/record only
+ Metadata map[string]any // Additional metadata
+}
+
+// TraceNode the trace node implementation
+type TraceNode struct {
+ ID string // Node ID
+ ParentID string // Parent node ID
+ Children []*TraceNode // Child nodes (for tree structure)
+ TraceNodeOption // Embedded option fields (Label, Icon, Description, Metadata)
+ Status string // Node status (pending, running, completed, failed, skipped)
+ Input TraceInput // Node input data
+ Output TraceOutput // Node output data
+ CreatedAt int64 // Creation timestamp
+ StartTime int64 // Start timestamp
+ EndTime int64 // End timestamp
+ UpdatedAt int64 // Last update timestamp
+ // Other fields will be added during implementation
+}
+
+// TraceSpace the trace memory space implementation (can add methods for serialization)
+type TraceSpace struct {
+ ID string // Space ID
+ TraceSpaceOption // Embedded option fields (Label, Icon, Description, TTL, Metadata)
+ CreatedAt int64 // Creation timestamp
+ UpdatedAt int64 // Last update timestamp
+ // Internal data storage will be managed by implementation
+}
+
+// TraceParallelInput defines input and options for a parallel node
+type TraceParallelInput struct {
+ Input TraceInput // Input data for the node
+ Option TraceNodeOption // Display options (label, icon, etc.)
+}
+
+// TraceInput the trace input (can add methods for validation)
+type TraceInput = any
+
+// TraceOutput the trace output (can add methods for formatting)
+type TraceOutput = any
+
+// Update event type constants (matching frontend SSE events)
+const (
+ // Trace lifecycle events
+ UpdateTypeInit = "init" // Trace initialization
+ UpdateTypeComplete = "complete" // Entire trace completed
+
+ // Node lifecycle events
+ UpdateTypeNodeStart = "node_start" // Node started (created)
+ UpdateTypeNodeComplete = "node_complete" // Node completed successfully
+ UpdateTypeNodeFailed = "node_failed" // Node failed with error
+ UpdateTypeNodeUpdated = "node_updated" // Node data updated (output, metadata, status)
+
+ // Log events
+ UpdateTypeLogAdded = "log_added" // Log entry added to node
+
+ // Memory/Space events
+ UpdateTypeMemoryAdd = "memory_add" // Memory space item added (key-value added)
+ UpdateTypeMemoryUpdate = "memory_update" // Memory space item updated
+ UpdateTypeMemoryDelete = "memory_delete" // Memory space item deleted
+ UpdateTypeSpaceCreated = "space_created" // Space was created
+ UpdateTypeSpaceDeleted = "space_deleted" // Space was deleted
+)
+
+// TraceUpdate represents a trace update event for subscriptions
+type TraceUpdate struct {
+ Type string // Update type (see UpdateType* constants)
+ 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
+ Data any // Update data (payload structures below)
+}
+
+// Event payload structures (matching frontend SSE format)
+
+// TraceInitData payload for "init" event
+type TraceInitData struct {
+ TraceID string `json:"traceId"`
+ AgentName string `json:"agentName,omitempty"`
+ RootNode *TraceNode `json:"rootNode,omitempty"`
+}
+
+// NodeStartData payload for "node_start" event
+// Supports both single node and multiple nodes (for parallel operations)
+type NodeStartData struct {
+ Node *TraceNode `json:"node,omitempty"` // Single node
+ Nodes []*TraceNode `json:"nodes,omitempty"` // Multiple nodes (for parallel)
+}
+
+// NodeCompleteData payload for "node_complete" event
+type NodeCompleteData struct {
+ NodeID string `json:"nodeId"`
+ Status string `json:"status"` // "success" or "failed"
+ EndTime int64 `json:"endTime"`
+ Duration int64 `json:"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 string `json:"status"` // "failed"
+ EndTime int64 `json:"endTime"`
+ Duration int64 `json:"duration"`
+ Error string `json:"error"`
+}
+
+// MemoryAddData payload for "memory_add" event
+type MemoryAddData struct {
+ Type string `json:"type"` // Space type/ID (e.g., "context", "intent", "knowledge")
+ Item MemoryItem `json:"item"`
+}
+
+// MemoryItem represents an item in memory space
+type MemoryItem struct {
+ ID string `json:"id"`
+ Type string `json:"type"`
+ Title string `json:"title,omitempty"`
+ Content any `json:"content"`
+ Timestamp int64 `json:"timestamp"`
+ Importance string `json:"importance,omitempty"` // "high", "medium", "low"
+}
+
+// TraceCompleteData payload for "complete" event
+type TraceCompleteData struct {
+ TraceID string `json:"traceId"`
+ Status string `json:"status"` // "completed"
+ TotalDuration int64 `json:"totalDuration"`
+}
+
+// SpaceDeletedData payload for "space_deleted" event
+type SpaceDeletedData struct {
+ SpaceID string `json:"spaceId"`
+}
+
+// MemoryDeleteData payload for "memory_delete" event
+type MemoryDeleteData struct {
+ SpaceID string `json:"spaceId"`
+ Key string `json:"key,omitempty"` // Empty when clearing all
+ Cleared bool `json:"cleared,omitempty"` // True when clearing all keys
+}
+
+// TraceInfo stores trace metadata and manager instance
+type TraceInfo struct {
+ ID string `json:"id"`
+ Driver string `json:"driver"`
+ Options []any `json:"options,omitempty"`
+ Manager Manager `json:"-"` // Not persisted
+ CreatedAt int64 `json:"created_at"`
+ UpdatedAt int64 `json:"updated_at"`
+ CreatedBy string `json:"__yao_created_by,omitempty"`
+ UpdatedBy string `json:"__yao_updated_by,omitempty"`
+ TeamID string `json:"__yao_team_id,omitempty"`
+ TenantID string `json:"__yao_tenant_id,omitempty"`
+ Metadata map[string]any `json:"metadata,omitempty"`
+}
+
+// TraceOption defines options for creating a trace
+type TraceOption struct {
+ ID string // Optional trace ID (if empty, generates new ID)
+ CreatedBy string // User who created the trace
+ TeamID string // Team ID
+ TenantID string // Tenant ID
+ Metadata map[string]any // Additional metadata
+}