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

Implement trace management system with local and store drivers

Browse source 
- Introduced a global trace registry to manage active traces with thread-safe access.
- Added support for two storage drivers: Local (file-based) and Store (Gou store).
- Enhanced trace ID generation to include a date prefix and improved uniqueness.
- Implemented functions for creating, loading, and releasing traces, along with checking their existence in both the registry and persistent storage.
- Added functionality to retrieve trace metadata and list active traces, improving trace management capabilities.
This commit is contained in:
Max 2025-11-18 11:50:05 +08:00
parent a824670def
commit 652035311d
12 changed files with 2855 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

572
agent/trace/README.md Normal file
View file

@ -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<br/>Status: Running]
Input[Input Processing<br/>✓ Completed<br/>Output: validated=true, items=3]
Fork{Parallel Fork}
Space1[(Worker Results<br/>Space)]
WorkerA[Worker A<br/>✓ Completed<br/>Output: worker=1]
WorkerB[Worker B<br/>✓ Completed<br/>Output: worker=2]
WorkerC[Worker C<br/>✓ Completed<br/>Output: worker=3]
Join((Auto Join))
Agg[Aggregation<br/>✓ Completed<br/>Output: total=3, success=true]
Space2[(Session Data<br/>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

143
agent/trace/local/driver.go Normal file
View file

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

742
agent/trace/manager.go Normal file
View file

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

266
agent/trace/node.go Normal file
View file

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

64
agent/trace/space.go Normal file
View file

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

147
agent/trace/store/driver.go Normal file
View file

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

106
agent/trace/subscription.go Normal file
View file

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

358
agent/trace/trace.go
View file

@ -1,27 +1,367 @@
package trace package trace
import ( import (
"context"
"fmt" "fmt"
"sync"
"time" "time"
gonanoid "github.com/matoous/go-nanoid/v2" 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) // safe: optional parameter, reserved for future safe mode implementation (collision detection)
func GenTraceID(safe ...bool) string { func GenTraceID(safe ...bool) string {
// TODO: Implement safe mode with collision detection when needed // 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) now := time.Now()
const alphabet = "1234567890" // Date prefix: YYYYMMDD (8 digits)
const length = 8 // 8 characters provides good balance of uniqueness and readability 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 { if err != nil {
// Fallback to timestamp-based ID if NanoID generation fails // Fallback to nanoseconds if NanoID generation fails
return fmt.Sprintf("%d", time.Now().UnixNano()) 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)
} }

73
agent/trace/types/driver.go Normal file
View file

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

97
agent/trace/types/events.go Normal file
View file

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

106
agent/trace/types/interfaces.go Normal file
View file

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

190
agent/trace/types/types.go Normal file
View file

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