[fix] trajectory agent step (#359)

fix trajectory agent step

Author: JqRrt

backend/modules/observability/domain/trace/entity/loop_span/trajectory.go

@@ -251,14 +251,10 @@ func buildAgentSteps(agentSpan *Span, spanMap map[string]*Span) []*Step {
 	childSpans := getDirectChildren(agentSpan, spanMap)
 
 	for _, childSpan := range childSpans {
-		// 深度遍历每个分支收集所有普通子节点，每个分支直到遇到agent节点为止
-		branchSteps := collectOtherSteps(childSpan, spanMap)
-		steps = append(steps, branchSteps...)
-
-		// 对每个直接子节点，向下深度遍历找到每个分支的第一个agent/model/tool节点
-		agentModelToolSteps := findAgentModelToolNode(childSpan, spanMap)
-		if len(agentModelToolSteps) > 0 {
-			steps = append(steps, agentModelToolSteps...)
+		// 深度遍历每个分支收集所有子节点，每个分支直到遇到agent节点为止
+		branchSteps := collectSubSteps(childSpan, spanMap)
+		if len(branchSteps) > 0 {
+			steps = append(steps, branchSteps...)
 		}
 	}
 
@@ -315,34 +311,8 @@ func buildStep(span *Span) *Step {
 	return step
 }
 
-// findAgentModelToolNode 向下深度遍历，找到每个分支的第一个agent/model/tool节点
-func findAgentModelToolNode(startSpan *Span, spanMap map[string]*Span) []*Step {
-	if startSpan == nil {
-		return nil
-	}
-
-	steps := make([]*Step, 0)
-	stepType := getStepType(startSpan)
-
-	// 如果当前节点就是agent/model/tool，直接返回
-	if stepType == StepTypeAgent || stepType == StepTypeModel || stepType == StepTypeTool {
-		steps = append(steps, buildStep(startSpan))
-		return steps
-	}
-
-	// 如果是other节点，继续向下遍历
-	children := getDirectChildren(startSpan, spanMap)
-	for _, child := range children {
-		if result := findAgentModelToolNode(child, spanMap); len(result) > 0 {
-			steps = append(steps, result...)
-		}
-	}
-
-	return steps
-}
-
-// collectOtherSteps 深度遍历分支，收集任意层级的普通子节点，直到遇到agent节点为止
-func collectOtherSteps(startSpan *Span, spanMap map[string]*Span) []*Step {
+// collectSubSteps 深度遍历分支，收集任意层级的普通子节点，直到遇到agent节点为止
+func collectSubSteps(startSpan *Span, spanMap map[string]*Span) []*Step {
 	if startSpan == nil {
 		return nil
 	}
@@ -351,19 +321,15 @@ func collectOtherSteps(startSpan *Span, spanMap map[string]*Span) []*Step {
 	stepType := getStepType(startSpan)
 
 	// 如果当前节点是agent节点，停止遍历
+	steps = append(steps, buildStep(startSpan))
 	if stepType == StepTypeAgent {
 		return steps
 	}
 
-	// 如果是普通节点，添加到结果中，然后继续向下遍历
-	if stepType != StepTypeModel && stepType != StepTypeTool {
-		steps = append(steps, buildStep(startSpan))
-	}
-
 	// 获取当前节点的子节点，继续深度遍历
 	children := getDirectChildren(startSpan, spanMap)
 	for _, child := range children {
-		childSteps := collectOtherSteps(child, spanMap)
+		childSteps := collectSubSteps(child, spanMap)
 		if len(childSteps) > 0 {
 			steps = append(steps, childSteps...)
 		}

backend/modules/observability/domain/trace/entity/loop_span/trajectory_test.go

@@ -220,6 +220,70 @@ func TestBuildTrajectoryFromSpans_ComplexTree(t *testing.T) {
 	assert.NotEmpty(t, s)
 }
 
+func TestBuildTrajectoryFromSpans_ModelIsParentOfTool(t *testing.T) {
+	t.Parallel()
+	traceID := "trace1"
+
+	// 根 agent
+	root := &Span{
+		SpanID:         "r",
+		ParentID:       "", // 作为root
+		TraceID:        traceID,
+		SpanName:       "root-agent",
+		SpanType:       "agent",
+		StartTime:      0,       // us
+		DurationMicros: 3000000, // 3s
+		Input:          "root-in",
+		Output:         "root-out",
+		TagsString:     map[string]string{},
+		TagsLong:       map[string]int64{},
+	}
+
+	// 分支1：model->tool
+	m1 := &Span{
+		SpanID:    "m1",
+		ParentID:  "r",
+		TraceID:   traceID,
+		SpanName:  "parser-1",
+		SpanType:  "model",
+		StartTime: 100, // us
+	}
+	t1 := &Span{
+		SpanID:   "t1",
+		ParentID: "m1",
+		TraceID:  traceID,
+		SpanName: "tool-1",
+		SpanType: "tool",
+	}
+
+	spans := SpanList{root, m1, t1}
+	traj := BuildTrajectoryFromSpans(spans)
+	assert.NotNil(t, traj)
+	assert.NotNil(t, traj.RootStep)
+	assert.NotNil(t, traj.ID)
+	assert.Equal(t, traceID, *traj.ID)
+
+	assert.Equal(t, 1, len(traj.AgentSteps))
+
+	// AgentSteps 包含 root 和 a1
+	// 找到 root 对应的 AgentStep
+	var m1Step, t1Step *Step
+	for _, s := range traj.AgentSteps[0].Steps {
+		if s != nil && s.ID != nil {
+			switch *s.ID {
+			case "m1":
+				m1Step = s
+			case "t1":
+				t1Step = s
+			default:
+			}
+		}
+	}
+	assert.NotNil(t, m1Step)
+	assert.NotNil(t, t1Step)
+	assert.Equal(t, *t1Step.ParentID, "m1")
+}
+
 func TestGetDirectChildren_Sorting(t *testing.T) {
 	t.Parallel()
 	parent := &Span{SpanID: "p"}