agentmon/internal/store/postgres/runs_test.go

package postgres

import (
	"encoding/json"
	"testing"
	"time"
)

func TestMergeEnvelopeJSON_MergesNestedPayloads(t *testing.T) {
	existing := json.RawMessage(`{
		"event":{"type":"span.start"},
		"attributes":{"name":"tool.call","span_kind":"tool"},
		"payload":{"input":{"query":"status"},"prompt_preview":"summarize"}
	}`)
	next := json.RawMessage(`{
		"event":{"type":"span.end"},
		"payload":{"result_preview":"ok","duration_ms":42}
	}`)

	merged := mergeEnvelopeJSON(existing, next)

	var got map[string]any
	if err := json.Unmarshal(merged, &got); err != nil {
		t.Fatalf("unmarshal merged payload: %v", err)
	}

	payload, ok := got["payload"].(map[string]any)
	if !ok {
		t.Fatal("expected merged payload object")
	}
	if _, ok := payload["input"].(map[string]any); !ok {
		t.Fatal("expected input from first event to be preserved")
	}
	if payload["result_preview"] != "ok" {
		t.Fatalf("expected result_preview to be merged, got %#v", payload["result_preview"])
	}
	if payload["duration_ms"] != float64(42) {
		t.Fatalf("expected duration_ms to be merged, got %#v", payload["duration_ms"])
	}

	attrs, ok := got["attributes"].(map[string]any)
	if !ok || attrs["name"] != "tool.call" {
		t.Fatalf("expected attributes to be preserved, got %#v", got["attributes"])
	}

	event, ok := got["event"].(map[string]any)
	if !ok || event["type"] != "span.end" {
		t.Fatalf("expected later event metadata to win, got %#v", got["event"])
	}
}

func TestMergeSpanEvent_PreservesStartPayloadDetails(t *testing.T) {
	existing := &SpanRow{
		Name:   "tool.call",
		Kind:   "tool",
		Status: "success",
		Payload: json.RawMessage(`{
			"attributes":{"name":"tool.call","span_kind":"tool"},
			"payload":{"input":{"command":"ls"}}
		}`),
	}

	mergeSpanEvent(existing, SpanRow{
		Status: "success",
		Payload: json.RawMessage(`{
			"payload":{"result_preview":"done","duration_ms":12}
		}`),
	})

	var got map[string]any
	if err := json.Unmarshal(existing.Payload, &got); err != nil {
		t.Fatalf("unmarshal merged span payload: %v", err)
	}

	payload := got["payload"].(map[string]any)
	if _, ok := payload["input"].(map[string]any); !ok {
		t.Fatal("expected input to remain after merge")
	}
	if payload["result_preview"] != "done" {
		t.Fatalf("expected result_preview after merge, got %#v", payload["result_preview"])
	}
}

func TestFindRunIndexForSpan_MatchesContainingRunWindow(t *testing.T) {
	start := time.Date(2026, 3, 23, 10, 0, 0, 0, time.UTC)
	run1End := start.Add(2 * time.Minute)
	run2Start := start.Add(3 * time.Minute)
	run2End := start.Add(7 * time.Minute)

	runs := []RunRow{
		{
			RunID:     "run-1",
			StartedAt: start,
			EndedAt:   &run1End,
		},
		{
			RunID:     "run-2",
			StartedAt: run2Start,
			EndedAt:   &run2End,
		},
	}

	idx := findRunIndexForSpan(runs, start.Add(4*time.Minute))
	if idx != 1 {
		t.Fatalf("expected span to attach to run-2, got index %d", idx)
	}
}

func TestFindRunIndexForSpan_MatchesOpenRun(t *testing.T) {
	start := time.Date(2026, 3, 23, 10, 0, 0, 0, time.UTC)
	run1End := start.Add(2 * time.Minute)
	run2Start := start.Add(3 * time.Minute)

	runs := []RunRow{
		{
			RunID:     "run-1",
			StartedAt: start,
			EndedAt:   &run1End,
		},
		{
			RunID:     "run-2",
			StartedAt: run2Start,
		},
	}

	idx := findRunIndexForSpan(runs, start.Add(5*time.Minute))
	if idx != 1 {
		t.Fatalf("expected span to attach to open run-2, got index %d", idx)
	}
}

func TestAddSpanToRunSpans_MergesDuplicateSpanIDs(t *testing.T) {
	rs := &runSpans{}

	addSpanToRunSpans(rs, SpanRow{
		SpanID: "span-1",
		Name:   "tool.call",
		Kind:   "tool",
		Payload: json.RawMessage(`{
			"payload":{"input":{"command":"ls"}}
		}`),
	})

	addSpanToRunSpans(rs, SpanRow{
		SpanID:  "span-1",
		Status:  "error",
		Duration: func() *int64 { v := int64(42); return &v }(),
		Payload: json.RawMessage(`{
			"payload":{"result_preview":"failed"}
		}`),
	})

	if len(rs.order) != 1 {
		t.Fatalf("expected one span order entry, got %d", len(rs.order))
	}

	span := rs.byID["span-1"]
	if span == nil {
		t.Fatal("expected merged span to be present")
	}
	if span.Name != "tool.call" {
		t.Fatalf("expected name to be preserved, got %q", span.Name)
	}
	if span.Status != "error" {
		t.Fatalf("expected error status to win, got %q", span.Status)
	}
	if span.Duration == nil || *span.Duration != 42 {
		t.Fatalf("expected duration to be merged, got %#v", span.Duration)
	}

	var got map[string]any
	if err := json.Unmarshal(span.Payload, &got); err != nil {
		t.Fatalf("unmarshal merged payload: %v", err)
	}
	payload, ok := got["payload"].(map[string]any)
	if !ok {
		t.Fatal("expected merged payload object")
	}
	if _, ok := payload["input"].(map[string]any); !ok {
		t.Fatal("expected input to remain after merge")
	}
	if payload["result_preview"] != "failed" {
		t.Fatalf("expected result_preview to be merged, got %#v", payload["result_preview"])
	}
}