Cyrene_v2/internal/runtime/trace/graph.go

package trace

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"
)

// Flow graph wire types (DevDocs/design/19 §1). The JSON field names are
// camelCase so the frontend's FlowGraph type and parseGraph read them directly.
// Node status strings mirror the frontend AgentStatus keys so parseGraph and
// FlowGraphView pass them through unchanged.

// Node status labels kept as named constants so no bare string appears in the
// mapping code (project rule §10).
const (
	// FlowStatusIdle is a node that has not started.
	FlowStatusIdle = "idle"
	// FlowStatusRunning is a node whose span is in flight.
	FlowStatusRunning = "running"
	// FlowStatusDone is a successfully completed node.
	FlowStatusDone = "done"
	// FlowStatusError is a node whose span finished with an error.
	FlowStatusError = "error"
)

// maxSummaryRunes bounds how many runes of a user string / tool result are
// rendered into a node label or prop summary, so sensitive or long content is
// collapsed to a safe, dashboard-friendly prefix.
const maxSummaryRunes = 80

// FlowNode is one node in a flow graph. Status/SpanKind mirror the frontier.
type FlowNode struct {
	ID        string         `json:"id"`
	Label     string         `json:"label"`
	Kind      string         `json:"kind,omitempty"`
	Status    string         `json:"status,omitempty"`
	StartedAt time.Time      `json:"startedAt,omitempty"`
	EndedAt   time.Time      `json:"endedAt,omitempty"`
	Props     map[string]any `json:"props,omitempty"`
	Log       []Event        `json:"log,omitempty"`
}

// FlowEdge is a directed edge between two nodes.
type FlowEdge struct {
	From string `json:"from"`
	To   string `json:"to"`
}

// FlowGraph is a live, single-round flow graph published as ws.graph_update.
type FlowGraph struct {
	ID        string     `json:"id"`
	Nodes     []FlowNode `json:"nodes"`
	Edges     []FlowEdge `json:"edges"`
	SessionID string     `json:"sessionId,omitempty"`
	StartedAt time.Time  `json:"startedAt,omitempty"`
	EndedAt   time.Time  `json:"endedAt,omitempty"`
}

// NodeSpec is the input to RunGraph.Upsert: the full description of a node plus
// the parent id used to wire its incoming edge.
type NodeSpec struct {
	ID        string
	Label     string
	Kind      string
	Status    string
	StartedAt time.Time
	EndedAt   time.Time
	Props     map[string]any
	Log       []Event
	ParentID  string
}

// RunGraph accumulates the spans of one round into a FlowGraph. It is a pure
// structure: it knows nothing about the bus, sessions or the loop, so it is
// unit-testable in isolation. A caller serializes access (see loop.FlowPublisher).
type RunGraph struct {
	session   string
	round     int
	nodes     map[string]*FlowNode
	order     []string
	edges     map[string]bool
	edgeOrder []FlowEdge
	last      string
	started   time.Time
	ended     time.Time
}

// NewRunGraph allocates an empty graph for one session round.
func NewRunGraph(session string) *RunGraph {
	return &RunGraph{
		session: session,
		nodes:   make(map[string]*FlowNode),
		edges:   make(map[string]bool),
	}
}

// Start marks the round began at now and resets the node/edge sets so a reused
// graph starts fresh for a new round.
func (g *RunGraph) Start(round int, now time.Time) {
	g.round = round
	g.started = now
	g.ended = time.Time{}
	g.nodes = make(map[string]*FlowNode)
	g.order = nil
	g.edges = make(map[string]bool)
	g.edgeOrder = nil
	g.last = ""
}

// Upsert creates or merges a node and wires its incoming edge. The edge source
// is the node's ParentID when that parent is already a node, otherwise the most
// recently added node — so a round assembles as a left-to-right trunk even when
// the parent container (a state_transition span) is omitted from the graph.
func (g *RunGraph) Upsert(spec NodeSpec) {
	if spec.ID == "" {
		return
	}
	if existing := g.nodes[spec.ID]; existing != nil {
		g.merge(existing, spec)
		return
	}
	n := &FlowNode{
		ID:        spec.ID,
		Label:     spec.Label,
		Kind:      spec.Kind,
		Status:    spec.Status,
		StartedAt: spec.StartedAt,
		EndedAt:   spec.EndedAt,
	}
	if spec.Props != nil {
		n.Props = spec.Props
	}
	n.Log = append(n.Log, spec.Log...)
	g.nodes[spec.ID] = n
	g.order = append(g.order, spec.ID)

	from := g.last
	if spec.ParentID != "" {
		if _, ok := g.nodes[spec.ParentID]; ok {
			from = spec.ParentID
		}
	}
	if from != "" && from != spec.ID {
		g.addEdge(from, spec.ID)
	}
	g.last = spec.ID
}

// merge updates an existing node in place (a running placeholder being finalized
// by its finished span), keeping label/kind and appending props/log.
func (g *RunGraph) merge(existing *FlowNode, spec NodeSpec) {
	if spec.Label != "" {
		existing.Label = spec.Label
	}
	if spec.Kind != "" {
		existing.Kind = spec.Kind
	}
	if spec.Status != "" {
		existing.Status = spec.Status
	}
	if !spec.EndedAt.IsZero() {
		existing.EndedAt = spec.EndedAt
	}
	if len(spec.Props) > 0 {
		if existing.Props == nil {
			existing.Props = make(map[string]any, len(spec.Props))
		}
		for k, v := range spec.Props {
			existing.Props[k] = v
		}
	}
	existing.Log = append(existing.Log, spec.Log...)
}

// addEdge records a directed edge once.
func (g *RunGraph) addEdge(from, to string) {
	key := from + "\x00" + to
	if g.edges[key] {
		return
	}
	g.edges[key] = true
	g.edgeOrder = append(g.edgeOrder, FlowEdge{From: from, To: to})
}

// SeedInput creates the "input" node for the round. It returns the node id so a
// caller can reference it.
func (g *RunGraph) SeedInput(text string) string {
	g.Upsert(NodeSpec{
		ID:        "input",
		Label:     truncate(text, maxSummaryRunes),
		Kind:      "input",
		Status:    FlowStatusDone,
		StartedAt: g.started,
	})
	return "input"
}

// SeedOutput creates the "output" node for the round, wired to the last node so
// the trunk ends at the answer.
func (g *RunGraph) SeedOutput(text string) string {
	g.ended = time.Now()
	g.Upsert(NodeSpec{
		ID:        "output",
		Label:     truncate(text, maxSummaryRunes),
		Kind:      "output",
		Status:    FlowStatusDone,
		StartedAt: g.ended,
		EndedAt:   g.ended,
	})
	return "output"
}

// AddSpan folds a finished span record into the graph as a node with its final
// status. State-transition spans are the step container and are omitted.
func (g *RunGraph) AddSpan(rec SpanRecord) {
	kind := kindForSpan(rec.Kind)
	if kind == "" {
		return
	}
	g.Upsert(NodeSpec{
		ID:        rec.ID,
		Label:     labelForSpan(rec),
		Kind:      kind,
		Status:    statusForFinished(rec.Status),
		StartedAt: rec.StartedAt,
		EndedAt:   rec.EndedAt,
		Props:     safeSpanProps(rec),
		Log:       rec.Events,
		ParentID:  rec.ParentID,
	})
}

// AddSpanStart registers a running placeholder for a span that just began, so a
// node appears as "running" before it finishes. State-transition spans are
// omitted, matching AddSpan. It reports whether a node was created.
func (g *RunGraph) AddSpanStart(id, name string, kind SpanKind, parent string, at time.Time) bool {
	nodeKind := kindForSpan(kind)
	if nodeKind == "" {
		return false
	}
	g.Upsert(NodeSpec{
		ID:        id,
		Label:     labelForStart(name, kind),
		Kind:      nodeKind,
		Status:    FlowStatusRunning,
		StartedAt: at,
		ParentID:  parent,
	})
	return true
}

// Snapshot returns a copy of the graph safe to encode and publish.
func (g *RunGraph) Snapshot() *FlowGraph {
	nodes := make([]FlowNode, 0, len(g.order))
	for _, id := range g.order {
		nodes = append(nodes, *g.nodes[id])
	}
	return &FlowGraph{
		ID:        fmt.Sprintf("%s-%d", g.session, g.round),
		Nodes:     nodes,
		Edges:     append([]FlowEdge(nil), g.edgeOrder...),
		SessionID: g.session,
		StartedAt: g.started,
		EndedAt:   g.ended,
	}
}

// kindForSpan maps a span kind to a flow node kind. state_transition is the step
// container and maps to "" so it is never rendered as a node.
func kindForSpan(k SpanKind) string {
	switch k {
	case SpanKindLLM:
		return "llm"
	case SpanKindTool:
		return "tool"
	case SpanKindSubagent:
		return "subagent"
	case SpanKindMemory:
		return "memory"
	case SpanKindGuardrail:
		return "guardrail"
	case SpanKindOrchestration:
		return "orchestration"
	default:
		return ""
	}
}

// labelForSpan renders a node label from a finished span.
func labelForSpan(rec SpanRecord) string {
	if kind := kindForSpan(rec.Kind); kind == "tool" {
		return strings.TrimPrefix(rec.Name, "tool.")
	}
	if rec.Kind == SpanKindLLM {
		return "思考"
	}
	if rec.Name != "" {
		return rec.Name
	}
	return string(rec.Kind)
}

// labelForStart renders a node label from a just-started span.
func labelForStart(name string, kind SpanKind) string {
	if kind == SpanKindLLM {
		return "思考"
	}
	if kind == SpanKindTool {
		return strings.TrimPrefix(name, "tool.")
	}
	if name != "" {
		return name
	}
	return string(kind)
}

// safeSpanProps carries only a whitelist of span props to the frontend, plus a
// desensitized tool result summary. It never forwards raw arguments or secrets,
// and truncates string text so long content cannot leak a full sensitive payload.
func safeSpanProps(rec SpanRecord) map[string]any {
	var props map[string]any
	for _, k := range safePropKeys {
		v, ok := rec.Props[k]
		if !ok {
			continue
		}
		if props == nil {
			props = make(map[string]any)
		}
		props[k] = redactValue(v)
	}
	if rec.Kind == SpanKindTool {
		if v, ok := rec.Props["result"]; ok {
			if props == nil {
				props = make(map[string]any)
			}
			props["result_summary"] = truncate(summarize(v), maxSummaryRunes)
		}
	}
	return props
}

// redactValue shortens string props so a long error/step value cannot leak
// sensitive content through the dashboard; non-string values pass through.
func redactValue(v any) any {
	if s, ok := v.(string); ok {
		return truncate(s, maxSummaryRunes)
	}
	return v
}

// safePropKeys are the span prop keys forwarded to a node's props.
var safePropKeys = []string{"stop_reason", "step", "error", "error_kind", "tool"}

// statusForFinished maps a span terminal status to a node status.
func statusForFinished(s SpanStatus) string {
	switch s {
	case SpanStatusError:
		return FlowStatusError
	case SpanStatusOK:
		return FlowStatusDone
	default:
		return FlowStatusIdle
	}
}

// summarize renders an arbitrary value as a short text string for a prop summary.
func summarize(v any) string {
	switch x := v.(type) {
	case nil:
		return ""
	case string:
		return x
	default:
		b, err := json.Marshal(x)
		if err != nil {
			return fmt.Sprintf("%v", x)
		}
		return string(b)
	}
}

// truncate cuts s to at most max runes, appending an ellipsis when it overflows.
func truncate(s string, max int) string {
	if max <= 0 {
		return s
	}
	r := []rune(s)
	if len(r) <= max {
		return s
	}
	return string(r[:max]) + "…"
}