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

Refactor trace node management to support multiple parents and enhance type handling

Browse source 
- Updated the TraceNode structure to allow multiple parent IDs, facilitating implicit joins in trace management.
- Enhanced the TraceNodeOption to include a type field for better categorization of nodes.
- Refactored related methods in the trace manager and node handling to accommodate the new parent structure and type field.
- Improved test coverage for trace node creation and retrieval, ensuring robust handling of the new multi-parent functionality.
This commit is contained in:
Max 2025-11-21 17:33:00 +08:00
parent 30778bbf2f
commit e781e8f591
18 changed files with 733 additions and 155 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
agent/assistant/agent.go
View file

@ -37,14 +37,10 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
if trace != nil {
agentNode, _ = trace.Add(inputMessages, types.TraceNodeOption{
Label: i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.label"), // "Assistant {{name}}"
Type: "agent",
Icon: "assistant",
Description: i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.description"), // "Assistant {{name}} is processing the request"
})
if agentNode != nil {
// Mark the node as complete when the function returns
defer agentNode.Complete()
}
}
// Full input messages with chat history
@ -119,11 +115,11 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
map[string]any{"messages": completionMessages, "options": completionOptions},
types.TraceNodeOption{
Label: fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.label"), conn.ID()), // "LLM %s"
Icon: "llm",
Type: "llm",
Icon: "psychology",
Description: fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.description"), conn.ID()), // "LLM %s is processing the request"
},
)
}
// Create LLM instance with connector and options

33
agent/i18n/i18n.go
View file

@ -331,20 +331,35 @@ func TranslateGlobal(locale string, input any) any {
i18ns = map[string]I18n{}
}
// Try the exact locale first
i18n, has := i18ns[locale]
if !has {
parts := strings.Split(locale, "-")
if len(parts) > 1 {
i18n, has = i18ns[parts[1]]
}
if !has {
i18n, has = i18ns[parts[0]]
if has {
result := i18n.Parse(input)
// If the result is the same as input (not translated), try fallback
if result != input {
return result
}
}
if has {
return i18n.Parse(input)
// Fallback logic: for "en-us", try "en"
parts := strings.Split(locale, "-")
if len(parts) > 1 {
// Try the language code (e.g., "en" for "en-us")
if fallbackI18n, hasFallback := i18ns[parts[0]]; hasFallback {
result := fallbackI18n.Parse(input)
if result != input {
return result
}
}
// Try the country code (e.g., "us" for "en-us")
if fallbackI18n, hasFallback := i18ns[parts[1]]; hasFallback {
result := fallbackI18n.Parse(input)
if result != input {
return result
}
}
}
return input
}

42
agent/i18n/i18n_test.go
View file

@ -733,6 +733,48 @@ func TestTranslateGlobal(t *testing.T) {
// Restore
Locales["__global__"] = temp
})
t.Run("TranslateGlobal fallback from en-us to en", func(t *testing.T) {
// Create a scenario where en-us has limited messages, but en has more
// This simulates the real-world case: en-us locale exists with 3 messages,
// but en has 41 messages including llm.handlers.stream.info
// Create en-us with only a few messages
enUSMessages := map[string]any{
"button.ok": "OK (US)", // en-us specific
"app.name": "My App",
"app.version": "1.0",
}
Locales["__global__"]["en-us"] = I18n{
Locale: "en-us",
Messages: enUSMessages,
}
// Test 1: Key exists in en-us - should use en-us
result := TranslateGlobal("en-us", "button.ok")
if result != "OK (US)" {
t.Errorf("Expected 'OK (US)' from en-us, got %v", result)
}
// Test 2: Key does NOT exist in en-us but exists in en - should fallback to en
result = TranslateGlobal("en-us", "button.cancel")
if result != "Cancel" {
t.Errorf("Expected 'Cancel' (fallback from en-us to en), got %v", result)
}
// Test 3: Built-in key that exists in en but not in en-us
result = TranslateGlobal("en-us", "llm.handlers.stream.info")
expected := "LLM Stream"
if result != expected {
t.Errorf("Expected '%s' (fallback from en-us to en), got %v", expected, result)
}
// Test 4: Direct key (not template) also should fallback
result = TranslateGlobal("en-us", "{{llm.handlers.stream.info}}")
if result != expected {
t.Errorf("Expected '%s' (fallback from en-us to en with template), got %v", expected, result)
}
})
}
// TestGetLocalesIntegration tests GetLocales with real assistant data

2
agent/llm/handlers/stream.go
View file

@ -21,7 +21,7 @@ func DefaultStreamHandler(ctx *context.Context) context.StreamFunc {
return func(chunkType context.StreamChunkType, data []byte) int {
trace, _ := ctx.Trace()
if trace != nil {
trace.Info(i18n.T(ctx.Locale, "llm.handlers.stream.info"), map[string]any{"data": string(data)}) // "LLM Stream"
trace.Info(i18n.T(ctx.Locale, "llm.handlers.stream.info"), map[string]any{"data": string(data)})
}
// Handle different chunk types

5
openapi/tests/trace/common_test.go
View file

@ -63,6 +63,8 @@ func prepareTestTrace(t *testing.T) *testTraceData {
// Create a memory space
space, err := manager.CreateSpace(types.TraceSpaceOption{
Label: "Test Space",
Type: "memory",
Icon: "database",
Description: "A test memory space",
})
assert.NoError(t, err)
@ -78,6 +80,7 @@ func prepareTestTrace(t *testing.T) *testTraceData {
// Add first node
node1, err := manager.Add("test input 1", types.TraceNodeOption{
Label: "First Node",
Type: "agent",
Icon: "icon1",
Description: "First test node",
Metadata: map[string]any{"node_order": 1},
@ -94,6 +97,7 @@ func prepareTestTrace(t *testing.T) *testTraceData {
// Add second node
node2, err := manager.Add("test input 2", types.TraceNodeOption{
Label: "Second Node",
Type: "tool",
Icon: "icon2",
Description: "Second test node",
Metadata: map[string]any{"node_order": 2},
@ -110,6 +114,7 @@ func prepareTestTrace(t *testing.T) *testTraceData {
// Add third node
node3, err := manager.Add("test input 3", types.TraceNodeOption{
Label: "Third Node",
Type: "custom",
Icon: "icon3",
Description: "Third test node",
Metadata: map[string]any{"node_order": 3},

27
openapi/tests/trace/nodes_test.go
View file

@ -53,6 +53,7 @@ func TestGetNodes(t *testing.T) {
// Verify node structure and metadata
metadataFound := 0
typeFound := 0
for _, n := range nodes {
node, ok := n.(map[string]interface{})
assert.True(t, ok, "Each node should be an object")
@ -61,6 +62,20 @@ func TestGetNodes(t *testing.T) {
assert.NotNil(t, node["status"], "Node should have status")
assert.NotNil(t, node["created_at"], "Node should have created_at")
// Verify type field is present
if node["type"] != nil {
nodeType, ok := node["type"].(string)
assert.True(t, ok, "Type should be a string")
assert.NotEmpty(t, nodeType, "Type should not be empty")
typeFound++
}
// Verify parent_ids field (should be array or null)
if node["parent_ids"] != nil {
_, ok := node["parent_ids"].([]interface{})
assert.True(t, ok, "parent_ids should be an array")
}
// Check if metadata is present (should be for all our test nodes)
if node["metadata"] != nil {
metadata, ok := node["metadata"].(map[string]interface{})
@ -72,6 +87,8 @@ func TestGetNodes(t *testing.T) {
}
}
assert.Equal(t, 3, typeFound, "All 3 nodes should have type field")
assert.Equal(t, 3, metadataFound, "All 3 nodes should have metadata with node_order")
t.Logf("Retrieved %d nodes for trace %s (all with metadata)", count, data.TraceID)
@ -107,9 +124,17 @@ func TestGetNodeByID(t *testing.T) {
// Verify response structure
assert.Equal(t, data.Node1ID, responseData["id"], "Node ID should match")
assert.Equal(t, "First Node", responseData["label"], "Node label should match")
assert.Equal(t, "agent", responseData["type"], "Node type should match")
assert.Equal(t, "icon1", responseData["icon"], "Node icon should match")
assert.Equal(t, "First test node", responseData["description"], "Node description should match")
// Verify parent_ids field (should be array or null)
if responseData["parent_ids"] != nil {
parentIDs, ok := responseData["parent_ids"].([]interface{})
assert.True(t, ok, "parent_ids should be an array")
t.Logf("Node has %d parent(s)", len(parentIDs))
}
// Verify metadata is present and correct
assert.NotNil(t, responseData["metadata"], "Metadata should be present")
metadata, ok := responseData["metadata"].(map[string]interface{})
@ -120,7 +145,7 @@ func TestGetNodeByID(t *testing.T) {
assert.NotNil(t, responseData["input"], "Input should be present")
assert.NotNil(t, responseData["output"], "Output should be present")
t.Logf("Retrieved node %s from trace %s with metadata: %+v", data.Node1ID, data.TraceID, metadata)
t.Logf("Retrieved node %s (type: %s) from trace %s with metadata: %+v", data.Node1ID, responseData["type"], data.TraceID, metadata)
}
// TestGetNodeByIDNotFound tests getting a non-existent node

20
openapi/tests/trace/spaces_test.go
View file

@ -18,6 +18,7 @@ func TestGetSpaces(t *testing.T) {
// Create additional spaces for this test
space1, err := data.Manager.CreateSpace(types.TraceSpaceOption{
Label: "Memory Space",
Type: "memory",
Icon: "memory",
Description: "Test memory space",
})
@ -25,6 +26,7 @@ func TestGetSpaces(t *testing.T) {
space2, err := data.Manager.CreateSpace(types.TraceSpaceOption{
Label: "Cache Space",
Type: "cache",
Icon: "cache",
Description: "Test cache space",
})
@ -64,6 +66,7 @@ func TestGetSpaces(t *testing.T) {
// Verify space metadata (should not include data field)
spaceLabels := make(map[string]bool)
typeFound := 0
for _, s := range spaces {
space := s.(map[string]any)
assert.NotNil(t, space["id"])
@ -71,9 +74,20 @@ func TestGetSpaces(t *testing.T) {
assert.NotNil(t, space["created_at"])
assert.NotNil(t, space["updated_at"])
assert.Nil(t, space["data"]) // Should NOT include key-value data
// Verify type field is present
if space["type"] != nil {
spaceType, ok := space["type"].(string)
assert.True(t, ok, "Type should be a string")
assert.NotEmpty(t, spaceType, "Type should not be empty")
typeFound++
}
spaceLabels[space["label"].(string)] = true
}
assert.GreaterOrEqual(t, typeFound, 2, "At least 2 spaces should have type field")
assert.True(t, spaceLabels["Memory Space"])
assert.True(t, spaceLabels["Cache Space"])
@ -88,9 +102,10 @@ func TestGetSpaceByID(t *testing.T) {
// Create a space with specific data
space, err := data.Manager.CreateSpace(types.TraceSpaceOption{
Label: "Detailed Space",
Type: "detailed",
Icon: "memory",
Description: "Space with detailed data",
Metadata: map[string]any{"type": "cache"},
Metadata: map[string]any{"cache_enabled": true},
})
assert.NoError(t, err)
@ -122,6 +137,7 @@ func TestGetSpaceByID(t *testing.T) {
// Verify space metadata
assert.Equal(t, space.ID, result["id"])
assert.Equal(t, "Detailed Space", result["label"])
assert.Equal(t, "detailed", result["type"])
assert.Equal(t, "memory", result["icon"])
assert.Equal(t, "Space with detailed data", result["description"])
assert.NotNil(t, result["created_at"])
@ -129,7 +145,7 @@ func TestGetSpaceByID(t *testing.T) {
// Verify metadata
metadata := result["metadata"].(map[string]any)
assert.Equal(t, "cache", metadata["type"])
assert.Equal(t, true, metadata["cache_enabled"])
// Verify key-value data
spaceData := result["data"].(map[string]any)

6
openapi/trace/nodes.go
View file

@ -48,8 +48,9 @@ func GetNodes(c *gin.Context) {
for _, node := range nodes {
nodeInfo := gin.H{
"id": node.ID,
"parent_id": node.ParentID,
"parent_ids": node.ParentIDs,
"label": node.Label,
"type": node.Type,
"icon": node.Icon,
"description": node.Description,
"status": node.Status,
@ -124,8 +125,9 @@ func GetNode(c *gin.Context) {
// Prepare detailed node response
nodeData := gin.H{
"id": node.ID,
"parent_id": node.ParentID,
"parent_ids": node.ParentIDs,
"label": node.Label,
"type": node.Type,
"icon": node.Icon,
"description": node.Description,
"status": node.Status,

2
openapi/trace/spaces.go
View file

@ -49,6 +49,7 @@ func GetSpaces(c *gin.Context) {
spaceInfo := gin.H{
"id": space.ID,
"label": space.Label,
"type": space.Type,
"icon": space.Icon,
"description": space.Description,
"ttl": space.TTL,
@ -122,6 +123,7 @@ func GetSpace(c *gin.Context) {
responseData := gin.H{
"id": spaceData.ID,
"label": spaceData.Label,
"type": spaceData.Type,
"icon": spaceData.Icon,
"description": spaceData.Description,
"ttl": spaceData.TTL,

7
trace/jsapi/node.go
View file

@ -23,12 +23,19 @@ func parseTraceNodeOption(obj *v8go.Object) types.TraceNodeOption {
if labelVal, err := obj.Get("label"); err == nil && !labelVal.IsNullOrUndefined() {
option.Label = labelVal.String()
}
if typeVal, err := obj.Get("type"); err == nil && !typeVal.IsNullOrUndefined() {
option.Type = typeVal.String()
}
if iconVal, err := obj.Get("icon"); err == nil && !iconVal.IsNullOrUndefined() {
option.Icon = iconVal.String()
}
if descVal, err := obj.Get("description"); err == nil && !descVal.IsNullOrUndefined() {
option.Description = descVal.String()
}
if autoCompleteVal, err := obj.Get("autoCompleteParent"); err == nil && !autoCompleteVal.IsNullOrUndefined() {
boolVal := autoCompleteVal.Boolean()
option.AutoCompleteParent = &boolVal
}
return option
}

3
trace/jsapi/trace.go
View file

@ -443,6 +443,9 @@ func traceCreateSpaceMethod(iso *v8go.Isolate, manager types.Manager) *v8go.Func
if labelVal, err := optionObj.Get("label"); err == nil && !labelVal.IsNullOrUndefined() {
option.Label = labelVal.String()
}
if typeVal, err := optionObj.Get("type"); err == nil && !typeVal.IsNullOrUndefined() {
option.Type = typeVal.String()
}
if iconVal, err := optionObj.Get("icon"); err == nil && !iconVal.IsNullOrUndefined() {
option.Icon = iconVal.String()
}

25
trace/local/driver.go
View file

@ -20,9 +20,10 @@ import (
// Only stores IDs of children instead of full child nodes
type persistNode struct {
ID string `json:"ID"`
ParentID string `json:"ParentID"`
ParentIDs []string `json:"ParentIDs,omitempty"`
ChildrenIDs []string `json:"ChildrenIDs,omitempty"`
Label string `json:"Label,omitempty"`
Type string `json:"Type,omitempty"`
Icon string `json:"Icon,omitempty"`
Description string `json:"Description,omitempty"`
Metadata map[string]any `json:"Metadata,omitempty"`
@ -51,9 +52,10 @@ func toPersistNode(node *types.TraceNode) *persistNode {
return &persistNode{
ID: node.ID,
ParentID: node.ParentID,
ParentIDs: node.ParentIDs,
ChildrenIDs: childrenIDs,
Label: node.Label,
Type: node.Type,
Icon: node.Icon,
Description: node.Description,
Metadata: node.Metadata,
@ -74,11 +76,12 @@ func fromPersistNode(pn *persistNode) *types.TraceNode {
}
return &types.TraceNode{
ID: pn.ID,
ParentID: pn.ParentID,
Children: nil, // Children will be loaded separately if needed
ID: pn.ID,
ParentIDs: pn.ParentIDs,
Children: nil, // Children will be loaded separately if needed
TraceNodeOption: types.TraceNodeOption{
Label: pn.Label,
Type: pn.Type,
Icon: pn.Icon,
Description: pn.Description,
Metadata: pn.Metadata,
@ -125,8 +128,12 @@ func New(basePath string) (*Driver, error) {
// Format: {basePath}/{YYYYMMDD}/{traceID}/
func (d *Driver) getTracePath(traceID string) string {
// Extract date prefix from traceID (first 8 digits)
datePrefix := traceID[:8]
return filepath.Join(d.basePath, datePrefix, traceID)
// If traceID is short, use "others" as prefix
prefix := "others"
if len(traceID) >= 8 {
prefix = traceID[:8]
}
return filepath.Join(d.basePath, prefix, traceID)
}
// ensureTraceDir creates the trace directory if it doesn't exist
@ -250,7 +257,7 @@ func (d *Driver) LoadTrace(ctx context.Context, traceID string) (*types.TraceNod
return nil, nil
}
// Find root node ID (node with empty ParentID) by checking each file
// Find root node ID (node with empty ParentIDs) by checking each file
var rootNodeID string
for _, entry := range entries {
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
@ -269,7 +276,7 @@ func (d *Driver) LoadTrace(ctx context.Context, traceID string) (*types.TraceNod
continue
}
if pn.ParentID == "" {
if len(pn.ParentIDs) == 0 {
rootNodeID = nodeID
break
}

159
trace/manager.go
View file

@ -100,77 +100,6 @@ func (m *manager) checkContext() error {
}
}
// handleCancellation marks nodes and trace as cancelled (called when context is done)
func (m *manager) handleCancellation() {
// Mark as completed first - this will trigger state worker to exit
// IMPORTANT: Must mark completed before any state queries to prevent deadlock
if !m.stateMarkCompleted() {
return // Already completed
}
now := time.Now().UnixMilli()
// Get current nodes (state worker will process this before exiting)
nodes := m.stateGetCurrentNodes()
// Mark only running/pending nodes as cancelled
for _, node := range nodes {
if node.Status == types.StatusRunning || node.Status == types.StatusPending {
node.Status = types.StatusCancelled
node.EndTime = now
node.UpdatedAt = now
// Save node with background context (ignore errors)
_ = m.driver.SaveNode(context.Background(), m.traceID, node)
// Broadcast cancelled event
m.addUpdateAndBroadcast(&types.TraceUpdate{
Type: types.UpdateTypeNodeFailed,
TraceID: m.traceID,
NodeID: node.ID,
Timestamp: now,
Data: &types.NodeFailedData{
NodeID: node.ID,
Status: types.CompleteStatusCancelled,
EndTime: now,
Duration: now - node.StartTime, // Already in milliseconds
Error: "context cancelled",
},
})
}
}
// Update trace status
m.stateSetTraceStatus(types.TraceStatusCancelled)
// Calculate total duration
totalDuration := int64(0)
rootNode := m.stateGetRoot()
if rootNode != nil && rootNode.CreatedAt > 0 {
totalDuration = now - rootNode.CreatedAt // Already in milliseconds
}
// Broadcast trace cancelled event (this will be processed even after state worker starts draining)
m.addUpdateAndBroadcast(&types.TraceUpdate{
Type: types.UpdateTypeComplete,
TraceID: m.traceID,
Timestamp: now,
Data: &types.TraceCompleteData{
TraceID: m.traceID,
Status: types.TraceStatusCancelled,
TotalDuration: totalDuration,
},
})
}
// 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,
}
}
// 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 {
@ -186,7 +115,7 @@ func (m *manager) Add(input types.TraceInput, option types.TraceNodeOption) (typ
// Create root node
rootNode = &types.TraceNode{
ID: genNodeID(),
ParentID: "",
ParentIDs: []string{}, // Root has no parents
Children: []*types.TraceNode{},
TraceNodeOption: option,
Status: types.StatusRunning,
@ -222,33 +151,32 @@ func (m *manager) Add(input types.TraceInput, option types.TraceNodeOption) (typ
// Get current nodes
currentNodes := m.stateGetCurrentNodes()
var parentNode *types.TraceNode
if len(currentNodes) > 1 {
// Auto-join: create join node
parentNode = &types.TraceNode{
ID: genNodeID(),
ParentID: currentNodes[0].ParentID,
Children: []*types.TraceNode{},
Status: types.StatusCompleted,
CreatedAt: now,
StartTime: now,
EndTime: now,
UpdatedAt: now,
TraceNodeOption: types.TraceNodeOption{Label: "Join", Icon: "join"},
// Auto-complete parent nodes if enabled (default: true when nil)
autoCompleteParent := option.AutoCompleteParent == nil || *option.AutoCompleteParent
if autoCompleteParent {
for _, current := range currentNodes {
// Only auto-complete running or pending nodes
if current.Status == types.StatusRunning || current.Status == types.StatusPending {
// Use node.Complete() method to properly complete the node with broadcast
currentNodeInterface := &node{manager: m, data: current}
if err := currentNodeInterface.Complete(); err != nil {
// Log error but don't fail the operation
m.Error("Failed to auto-complete parent node %s: %v", current.ID, err)
}
}
}
}
// Save join node
if err := m.driver.SaveNode(m.ctx, m.traceID, parentNode); err != nil {
return nil, err
}
} else {
parentNode = currentNodes[0]
// Collect parent IDs from all current nodes (supports implicit join)
parentIDs := make([]string, 0, len(currentNodes))
for _, current := range currentNodes {
parentIDs = append(parentIDs, current.ID)
}
// Create new node
// Create new node with multiple parents (implicit join)
newNodeData := &types.TraceNode{
ID: genNodeID(),
ParentID: parentNode.ID,
ParentIDs: parentIDs, // Multiple parents for implicit join
Children: []*types.TraceNode{},
TraceNodeOption: option,
Status: types.StatusRunning,
@ -258,14 +186,17 @@ func (m *manager) Add(input types.TraceInput, option types.TraceNodeOption) (typ
UpdatedAt: now,
}
// Update 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
// Add to each parent's children
for _, parent := range currentNodes {
parent.Children = append(parent.Children, newNodeData)
if err := m.driver.SaveNode(m.ctx, m.traceID, parent); err != nil {
// Log error but continue
m.Error("Failed to update parent node %s: %v", parent.ID, err)
}
}
if err := m.driver.SaveNode(m.ctx, m.traceID, parentNode); err != nil {
// Save new node
if err := m.driver.SaveNode(m.ctx, m.traceID, newNodeData); err != nil {
return nil, err
}
@ -303,16 +234,40 @@ func (m *manager) Parallel(parallelInputs []types.TraceParallelInput) ([]types.N
// Get current nodes
currentNodes := m.stateGetCurrentNodes()
// Auto-complete parent node if any option has AutoCompleteParent enabled (default: true when nil)
shouldAutoComplete := false
for _, input := range parallelInputs {
if input.Option.AutoCompleteParent == nil || *input.Option.AutoCompleteParent {
shouldAutoComplete = true
break
}
}
if shouldAutoComplete {
for _, current := range currentNodes {
// Only auto-complete running or pending nodes
if current.Status == types.StatusRunning || current.Status == types.StatusPending {
// Use node.Complete() method to properly complete the node with broadcast
currentNodeInterface := &node{manager: m, data: current}
if err := currentNodeInterface.Complete(); err != nil {
// Log error but don't fail the operation
m.Error("Failed to auto-complete parent node %s: %v", current.ID, err)
}
}
}
}
parentNode := currentNodes[0]
nodeData := make([]*types.TraceNode, 0, len(parallelInputs))
nodeInterfaces := make([]types.Node, 0, len(parallelInputs))
// Create multiple child nodes
// Create multiple child nodes with single parent
for _, input := range parallelInputs {
data := &types.TraceNode{
ID: genNodeID(),
ParentID: parentNode.ID,
ParentIDs: []string{parentNode.ID}, // Single parent for parallel branches
Children: []*types.TraceNode{},
TraceNodeOption: input.Option,
Status: types.StatusRunning,

16
trace/node.go
View file

@ -71,7 +71,7 @@ func (n *node) Add(input types.TraceInput, option types.TraceNodeOption) (types.
// Create child node data
childNodeData := &types.TraceNode{
ID: genNodeID(),
ParentID: n.data.ID,
ParentIDs: []string{n.data.ID}, // Single parent
Children: []*types.TraceNode{},
TraceNodeOption: option,
Status: types.StatusRunning,
@ -108,7 +108,7 @@ func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node
for _, input := range parallelInputs {
childNodeData := &types.TraceNode{
ID: genNodeID(),
ParentID: n.data.ID,
ParentIDs: []string{n.data.ID}, // Single parent for parallel branches
Children: []*types.TraceNode{},
TraceNodeOption: input.Option,
Status: types.StatusRunning,
@ -143,10 +143,18 @@ func (n *node) Parallel(parallelInputs []types.TraceParallelInput) ([]types.Node
func (n *node) Join(nodes []*types.TraceNode, input types.TraceInput, option types.TraceNodeOption) (types.Node, error) {
now := time.Now().UnixMilli()
// Create join node data
// Collect parent IDs from all nodes
parentIDs := make([]string, 0, len(nodes))
for _, node := range nodes {
if node != nil {
parentIDs = append(parentIDs, node.ID)
}
}
// Create join node data with multiple parents
joinNodeData := &types.TraceNode{
ID: genNodeID(),
ParentID: n.data.ID,
ParentIDs: parentIDs, // Multiple parents for explicit join
Children: []*types.TraceNode{},
TraceNodeOption: option,
Status: types.StatusRunning,

17
trace/store/driver.go
View file

@ -19,9 +19,10 @@ import (
// Only stores IDs of children instead of full child nodes
type persistNode struct {
ID string `json:"ID"`
ParentID string `json:"ParentID"`
ParentIDs []string `json:"ParentIDs,omitempty"`
ChildrenIDs []string `json:"ChildrenIDs,omitempty"`
Label string `json:"Label,omitempty"`
Type string `json:"Type,omitempty"`
Icon string `json:"Icon,omitempty"`
Description string `json:"Description,omitempty"`
Metadata map[string]any `json:"Metadata,omitempty"`
@ -50,9 +51,10 @@ func toPersistNode(node *types.TraceNode) *persistNode {
return &persistNode{
ID: node.ID,
ParentID: node.ParentID,
ParentIDs: node.ParentIDs,
ChildrenIDs: childrenIDs,
Label: node.Label,
Type: node.Type,
Icon: node.Icon,
Description: node.Description,
Metadata: node.Metadata,
@ -73,11 +75,12 @@ func fromPersistNode(pn *persistNode) *types.TraceNode {
}
return &types.TraceNode{
ID: pn.ID,
ParentID: pn.ParentID,
Children: nil, // Children will be loaded separately if needed
ID: pn.ID,
ParentIDs: pn.ParentIDs,
Children: nil, // Children will be loaded separately if needed
TraceNodeOption: types.TraceNodeOption{
Label: pn.Label,
Type: pn.Type,
Icon: pn.Icon,
Description: pn.Description,
Metadata: pn.Metadata,
@ -253,7 +256,7 @@ func (d *Driver) LoadTrace(ctx context.Context, traceID string) (*types.TraceNod
return nil, nil
}
// Find root node ID (node with empty ParentID) by checking each node
// Find root node ID (node with empty ParentIDs) by checking each node
var rootNodeID string
for _, key := range nodeKeys {
// Extract node ID from key (format: prefix:traceID:nodes:nodeID)
@ -279,7 +282,7 @@ func (d *Driver) LoadTrace(ctx context.Context, traceID string) (*types.TraceNod
continue
}
if pn.ParentID == "" {
if len(pn.ParentIDs) == 0 {
rootNodeID = nodeID
break
}

473
trace/trace_autocomplete_test.go Normal file
View file

@ -0,0 +1,473 @@
package trace_test
import (
"context"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/trace"
"github.com/yaoapp/yao/trace/types"
)
// TestAutoCompleteParentDefault tests the default behavior where parent nodes are auto-completed
func TestAutoCompleteParentDefault(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
defer trace.Release(traceID)
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Add first node
node1, err := manager.Add("step1", types.TraceNodeOption{
Label: "Step 1",
})
assert.NoError(t, err)
// Verify node1 is running
currentNodes, err := manager.GetCurrentNodes()
assert.NoError(t, err)
assert.Len(t, currentNodes, 1)
assert.Equal(t, types.StatusRunning, currentNodes[0].Status)
// Add second node - node1 should be auto-completed
node2, err := manager.Add("step2", types.TraceNodeOption{
Label: "Step 2",
})
assert.NoError(t, err)
// Verify node1 is now completed
node1Data, err := manager.GetNodeByID(node1.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusCompleted, node1Data.Status)
// Verify node2 is running
currentNodes, err = manager.GetCurrentNodes()
assert.NoError(t, err)
assert.Len(t, currentNodes, 1)
assert.Equal(t, node2.ID(), currentNodes[0].ID)
assert.Equal(t, types.StatusRunning, currentNodes[0].Status)
})
}
}
// TestAutoCompleteParentDisabled tests disabling auto-complete behavior
func TestAutoCompleteParentDisabled(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
defer trace.Release(traceID)
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Add first node
node1, err := manager.Add("step1", types.TraceNodeOption{
Label: "Step 1",
})
assert.NoError(t, err)
// Add second node with auto-complete disabled
falseVal := false
node2, err := manager.Add("step2", types.TraceNodeOption{
Label: "Step 2",
AutoCompleteParent: &falseVal,
})
assert.NoError(t, err)
// Verify node1 is still running (not auto-completed)
node1Data, err := manager.GetNodeByID(node1.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusRunning, node1Data.Status)
// Verify node2 is running
node2Data, err := manager.GetNodeByID(node2.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusRunning, node2Data.Status)
})
}
}
// TestAutoCompleteParentParallel tests auto-complete with parallel nodes
func TestAutoCompleteParentParallel(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
defer trace.Release(traceID)
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Add root node
rootNode, err := manager.Add("root", types.TraceNodeOption{
Label: "Root",
})
assert.NoError(t, err)
// Create parallel nodes
parallelNodes, err := manager.Parallel([]types.TraceParallelInput{
{Input: "task1", Option: types.TraceNodeOption{Label: "Task 1"}},
{Input: "task2", Option: types.TraceNodeOption{Label: "Task 2"}},
{Input: "task3", Option: types.TraceNodeOption{Label: "Task 3"}},
})
assert.NoError(t, err)
assert.Len(t, parallelNodes, 3)
// Root node should be auto-completed
rootData, err := manager.GetNodeByID(rootNode.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusCompleted, rootData.Status)
// All parallel nodes should be running
for _, node := range parallelNodes {
nodeData, err := manager.GetNodeByID(node.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusRunning, nodeData.Status)
}
// Add a merge node - all parallel nodes should be auto-completed
mergeNode, err := manager.Add("merge", types.TraceNodeOption{
Label: "Merge",
})
assert.NoError(t, err)
// All parallel nodes should now be completed
for _, node := range parallelNodes {
nodeData, err := manager.GetNodeByID(node.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusCompleted, nodeData.Status)
}
// Merge node should be running
mergeData, err := manager.GetNodeByID(mergeNode.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusRunning, mergeData.Status)
})
}
}
// TestAutoCompleteParentConcurrent tests auto-complete with concurrent operations
func TestAutoCompleteParentConcurrent(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
defer trace.Release(traceID)
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Add root node
_, err = manager.Add("root", types.TraceNodeOption{Label: "Root"})
assert.NoError(t, err)
// Create parallel nodes
parallelNodes, err := manager.Parallel([]types.TraceParallelInput{
{Input: "task1", Option: types.TraceNodeOption{Label: "Task 1"}},
{Input: "task2", Option: types.TraceNodeOption{Label: "Task 2"}},
{Input: "task3", Option: types.TraceNodeOption{Label: "Task 3"}},
})
assert.NoError(t, err)
// Complete all parallel nodes (simulating concurrent work)
var wg sync.WaitGroup
for _, pNode := range parallelNodes {
wg.Add(1)
go func(node types.Node) {
defer wg.Done()
// Small delay to simulate real work
time.Sleep(10 * time.Millisecond)
err := node.Complete(map[string]any{"done": true})
assert.NoError(t, err)
}(pNode)
}
wg.Wait()
// All parallel nodes should be completed
for _, node := range parallelNodes {
nodeData, err := manager.GetNodeByID(node.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusCompleted, nodeData.Status)
}
// Add a merge node - should work correctly with all parents completed
mergeNode, err := manager.Add("merge", types.TraceNodeOption{
Label: "Merge Results",
})
assert.NoError(t, err)
// Verify merge node is running
mergeData, err := manager.GetNodeByID(mergeNode.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusRunning, mergeData.Status)
// Verify merge node has all parallel nodes as parents
assert.Len(t, mergeData.ParentIDs, 3)
})
}
}
// TestAutoCompleteParentWithFailedNode tests that failed nodes are not auto-completed
func TestAutoCompleteParentWithFailedNode(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
defer trace.Release(traceID)
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Add first node
node1, err := manager.Add("step1", types.TraceNodeOption{
Label: "Step 1",
})
assert.NoError(t, err)
// Fail node1
err = node1.Fail(assert.AnError)
assert.NoError(t, err)
// Verify node1 is failed
node1Data, err := manager.GetNodeByID(node1.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusFailed, node1Data.Status)
// Add second node with auto-complete enabled
node2, err := manager.Add("step2", types.TraceNodeOption{
Label: "Step 2",
})
assert.NoError(t, err)
// node1 should still be failed (not auto-completed to success)
node1Data, err = manager.GetNodeByID(node1.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusFailed, node1Data.Status)
// node2 should be running
node2Data, err := manager.GetNodeByID(node2.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusRunning, node2Data.Status)
})
}
}
// TestAutoCompleteParentWithCompletedNode tests that already completed nodes are not affected
func TestAutoCompleteParentWithCompletedNode(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
defer trace.Release(traceID)
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Add first node
node1, err := manager.Add("step1", types.TraceNodeOption{
Label: "Step 1",
})
assert.NoError(t, err)
// Complete node1 manually with output
err = node1.Complete(map[string]any{"result": "manual"})
assert.NoError(t, err)
// Verify node1 is completed with output
node1Data, err := manager.GetNodeByID(node1.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusCompleted, node1Data.Status)
assert.Equal(t, map[string]any{"result": "manual"}, node1Data.Output)
// Add second node with auto-complete enabled
node2, err := manager.Add("step2", types.TraceNodeOption{
Label: "Step 2",
})
assert.NoError(t, err)
// node1 should still be completed with same output
node1Data, err = manager.GetNodeByID(node1.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusCompleted, node1Data.Status)
assert.Equal(t, map[string]any{"result": "manual"}, node1Data.Output)
// node2 should be running
node2Data, err := manager.GetNodeByID(node2.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusRunning, node2Data.Status)
})
}
}
// TestAutoCompleteParentEvents tests that auto-complete generates proper events
func TestAutoCompleteParentEvents(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
defer trace.Release(traceID)
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Subscribe to updates
updates, err := manager.Subscribe()
assert.NoError(t, err)
// Collect updates in background
var receivedUpdates []*types.TraceUpdate
var updatesMu sync.Mutex
done := make(chan bool)
go func() {
timeout := time.After(2 * time.Second)
for {
select {
case update, ok := <-updates:
if !ok {
done <- true
return
}
updatesMu.Lock()
receivedUpdates = append(receivedUpdates, update)
updatesMu.Unlock()
case <-timeout:
done <- true
return
}
}
}()
// Add first node
node1, err := manager.Add("step1", types.TraceNodeOption{
Label: "Step 1",
})
assert.NoError(t, err)
// Small delay
time.Sleep(50 * time.Millisecond)
// Add second node - should trigger auto-complete of node1
_, err = manager.Add("step2", types.TraceNodeOption{
Label: "Step 2",
})
assert.NoError(t, err)
// Wait a bit for events to be processed
time.Sleep(100 * time.Millisecond)
// Wait for goroutine to finish (with timeout)
<-done
// Verify we received auto-complete event for node1
updatesMu.Lock()
defer updatesMu.Unlock()
var foundNode1Start bool
var foundNode1Complete bool
var foundNode2Start bool
for _, update := range receivedUpdates {
switch update.Type {
case types.UpdateTypeNodeStart:
if data, ok := update.Data.(*types.NodeStartData); ok && data.Node != nil {
if data.Node.ID == node1.ID() {
foundNode1Start = true
} else if data.Node.Label == "Step 2" {
foundNode2Start = true
}
}
case types.UpdateTypeNodeComplete:
if data, ok := update.Data.(*types.NodeCompleteData); ok {
if data.NodeID == node1.ID() {
foundNode1Complete = true
}
}
}
}
assert.True(t, foundNode1Start, "Should receive node_start event for node1")
assert.True(t, foundNode1Complete, "Should receive node_complete event for node1 (auto-completed)")
assert.True(t, foundNode2Start, "Should receive node_start event for node2")
})
}
}
// TestAutoCompleteParentSequentialChain tests a long sequential chain
func TestAutoCompleteParentSequentialChain(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
defer trace.Release(traceID)
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Create a chain of 10 nodes
nodeCount := 10
nodes := make([]types.Node, nodeCount)
for i := 0; i < nodeCount; i++ {
node, err := manager.Add("input", types.TraceNodeOption{
Label: "Step " + string(rune('0'+i)),
})
assert.NoError(t, err)
nodes[i] = node
// All previous nodes should be completed
for j := 0; j < i; j++ {
nodeData, err := manager.GetNodeByID(nodes[j].ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusCompleted, nodeData.Status,
"Node %d should be completed when node %d is added", j, i)
}
// Current node should be running
currentData, err := manager.GetNodeByID(node.ID())
assert.NoError(t, err)
assert.Equal(t, types.StatusRunning, currentData.Status)
}
// Get all nodes
allNodes, err := manager.GetAllNodes()
assert.NoError(t, err)
assert.Len(t, allNodes, nodeCount)
// All except the last should be completed
completedCount := 0
runningCount := 0
for _, nodeData := range allNodes {
if nodeData.Status == types.StatusCompleted {
completedCount++
} else if nodeData.Status == types.StatusRunning {
runningCount++
}
}
assert.Equal(t, nodeCount-1, completedCount, "Should have %d completed nodes", nodeCount-1)
assert.Equal(t, 1, runningCount, "Should have 1 running node")
})
}
}

24
trace/trace_basic_test.go
View file

@ -47,7 +47,11 @@ func TestTraceNew(t *testing.T) {
assert.Nil(t, root)
// Add first node - this should become the root
node, err := manager.Add("test input", types.TraceNodeOption{Label: "First Node", Icon: "test"})
node, err := manager.Add("test input", types.TraceNodeOption{
Label: "First Node",
Type: "test",
Icon: "test",
})
assert.NoError(t, err)
assert.NotNil(t, node)
@ -56,6 +60,7 @@ func TestTraceNew(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, root)
assert.Equal(t, "First Node", root.Label)
assert.Equal(t, "test", root.Type)
assert.Equal(t, "test", root.Icon)
})
}
@ -109,10 +114,18 @@ func TestTraceLoadFromStorage(t *testing.T) {
assert.NoError(t, err)
// Add some data
_, err = manager.Add("test", types.TraceNodeOption{Label: "Test"})
_, err = manager.Add("test", types.TraceNodeOption{
Label: "Test",
Type: "test_node",
Icon: "node",
})
assert.NoError(t, err)
space, err := manager.CreateSpace(types.TraceSpaceOption{Label: "Test Space"})
space, err := manager.CreateSpace(types.TraceSpaceOption{
Label: "Test Space",
Type: "test_space",
Icon: "space",
})
assert.NoError(t, err)
err = manager.SetSpaceValue(space.ID, "key", "value")
@ -225,7 +238,10 @@ func TestContextCancellation(t *testing.T) {
cancel()
// Operations should fail with context error
_, err = manager.Add("test", types.TraceNodeOption{Label: "Test"})
_, err = manager.Add("test", types.TraceNodeOption{
Label: "Test",
Type: "test",
})
assert.Error(t, err)
})
}

17
trace/types/types.go
View file

@ -37,15 +37,18 @@ const (
// TraceNodeOption defines options for creating a node
type TraceNodeOption struct {
Label string `json:"label"` // Display label in UI
Icon string `json:"icon"` // Icon identifier
Description string `json:"description"` // Node description
Metadata map[string]any `json:"metadata,omitempty"` // Additional metadata
Label string `json:"label"` // Display label in UI
Type string `json:"type"` // Node type identifier
Icon string `json:"icon"` // Icon identifier
Description string `json:"description"` // Node description
Metadata map[string]any `json:"metadata,omitempty"` // Additional metadata
AutoCompleteParent *bool `json:"auto_complete_parent,omitempty"` // Auto-complete parent node(s) when this node is created (nil = default true)
}
// TraceSpaceOption defines options for creating a space
type TraceSpaceOption struct {
Label string `json:"label"` // Display label in UI
Type string `json:"type"` // Space type identifier
Icon string `json:"icon"` // Icon identifier
Description string `json:"description"` // Space description
TTL int64 `json:"ttl"` // Time to live in seconds (0 = no expiration) - for display/record only
@ -54,9 +57,9 @@ type TraceSpaceOption struct {
// TraceNode the trace node implementation
type TraceNode struct {
ID string `json:"id"` // Node ID
ParentID string `json:"parent_id"` // Parent node ID
Children []*TraceNode `json:"children"` // Child nodes (for tree structure)
ID string `json:"id"` // Node ID
ParentIDs []string `json:"parent_ids"` // Parent node IDs (supports multiple parents for implicit join)
Children []*TraceNode `json:"children"` // Child nodes (for tree structure)
TraceNodeOption `json:",inline"` // Embedded option fields (Label, Icon, Description, Metadata)
Status NodeStatus `json:"status"` // Node status (pending, running, completed, failed, skipped)
Input TraceInput `json:"input,omitempty"` // Node input data