Add live agent views and improve Codex monitoring

internal/store/postgres/runs.go

@@ -15,6 +15,7 @@ type RunRow struct {
 	SpanCount int        `json:"span_count"`
 	ToolCount int        `json:"tool_count"`
 	Model     string     `json:"model,omitempty"`
+	Spans     []SpanRow  `json:"spans,omitempty"`
 }
 
 type SessionDetail struct {
@@ -84,10 +85,20 @@ func (d *DB) GetSessionWithRuns(ctx context.Context, sessionID string) (*Session
 		runs = append(runs, r)
 	}
 
-	return &session, runs, rows.Err()
+	if err := rows.Err(); err != nil {
+		return nil, nil, err
+	}
+
+	runs, err = d.attachSpansToRuns(ctx, sessionID, runs)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return &session, runs, nil
 }
 
 type SpanRow struct {
+	RunID     string          `json:"run_id,omitempty"`
 	SpanID    string          `json:"span_id"`
 	Name      string          `json:"name"`
 	Kind      string          `json:"kind"`
@@ -129,9 +140,18 @@ func (d *DB) GetRunWithSpans(ctx context.Context, runID string) (*RunDetail, []S
 		return nil, nil, err
 	}
 
-	// Get spans
-	spansQuery := `
+	spans, err := d.listSpansForRun(ctx, runID)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return &run, spans, nil
+}
+
+func (d *DB) listSpansForRun(ctx context.Context, runID string) ([]SpanRow, error) {
+	rows, err := d.sql.QueryContext(ctx, `
 		SELECT
+			run_id,
 			span_id,
 			COALESCE(payload->'attributes'->>'name', payload->'event'->>'type', type) as name,
 			COALESCE(payload->'attributes'->>'span_kind', 'unknown') as kind,
@@ -145,21 +165,84 @@ func (d *DB) GetRunWithSpans(ctx context.Context, runID string) (*RunDetail, []S
 		FROM events
 		WHERE run_id = $1 AND span_id IS NOT NULL
 		ORDER BY ts ASC
-	`
-	rows, err := d.sql.QueryContext(ctx, spansQuery, runID)
+	`, runID)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 	defer rows.Close()
 
-	var spans []SpanRow
+	spansByID := make(map[string]*SpanRow)
+	var spanOrder []string
 	for rows.Next() {
 		var s SpanRow
-		if err := rows.Scan(&s.SpanID, &s.Name, &s.Kind, &s.StartedAt, &s.Duration, &s.Status, &s.Payload); err != nil {
-			return nil, nil, err
+		if err := rows.Scan(&s.RunID, &s.SpanID, &s.Name, &s.Kind, &s.StartedAt, &s.Duration, &s.Status, &s.Payload); err != nil {
+			return nil, err
 		}
-		spans = append(spans, s)
+
+		existing := spansByID[s.SpanID]
+		if existing == nil {
+			copy := s
+			spansByID[s.SpanID] = &copy
+			spanOrder = append(spanOrder, s.SpanID)
+			continue
+		}
+
+		if existing.Name == "" && s.Name != "" {
+			existing.Name = s.Name
+		}
+		if existing.Kind == "" || existing.Kind == "unknown" {
+			existing.Kind = s.Kind
+		}
+		if s.Duration != nil {
+			existing.Duration = s.Duration
+		}
+		if s.Status == "error" {
+			existing.Status = "error"
+		}
+		existing.Payload = s.Payload
 	}
 
-	return &run, spans, rows.Err()
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	spans := make([]SpanRow, 0, len(spanOrder))
+	for _, spanID := range spanOrder {
+		spans = append(spans, *spansByID[spanID])
+	}
+	return spans, nil
+}
+
+func (d *DB) attachSpansToRuns(ctx context.Context, sessionID string, runs []RunRow) ([]RunRow, error) {
+	rows, err := d.sql.QueryContext(ctx, `
+		SELECT DISTINCT run_id
+		FROM events
+		WHERE session_id = $1 AND run_id IS NOT NULL
+		ORDER BY run_id
+	`, sessionID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	spansByRun := make(map[string][]SpanRow)
+	for rows.Next() {
+		var runID string
+		if err := rows.Scan(&runID); err != nil {
+			return nil, err
+		}
+		spans, err := d.listSpansForRun(ctx, runID)
+		if err != nil {
+			return nil, err
+		}
+		spansByRun[runID] = spans
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	for i := range runs {
+		runs[i].Spans = spansByRun[runs[i].RunID]
+	}
+	return runs, nil
 }