Implement trace management system with local and store drivers
- 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
parent
a824670def
commit
652035311d
12 changed files with 2855 additions and 9 deletions
572
agent/trace/README.md
Normal file
572
agent/trace/README.md
Normal 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
143
agent/trace/local/driver.go
Normal 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
742
agent/trace/manager.go
Normal 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
266
agent/trace/node.go
Normal 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
64
agent/trace/space.go
Normal 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
147
agent/trace/store/driver.go
Normal 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
106
agent/trace/subscription.go
Normal 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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
73
agent/trace/types/driver.go
Normal file
73
agent/trace/types/driver.go
Normal 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
97
agent/trace/types/events.go
Normal 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
106
agent/trace/types/interfaces.go
Normal 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
190
agent/trace/types/types.go
Normal 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
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue