krow_backend/go-api/internal/runtime/store_test.go

package runtime_test

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

	"github.com/krow/krow-backend/go-api/internal/runtime"
	"github.com/krow/krow-backend/go-api/internal/testutil"
)

// trajectory builds a saveable run for the given harness org.
func trajectory(orgID, runID string) *runtime.Trajectory {
	started := time.Now().Add(-2 * time.Second).UTC()
	return &runtime.Trajectory{
		RunID:        runID,
		OrgID:        orgID,
		AgentID:      "activity-agent",
		AgentVersion: 3,
		Tier:         "balanced",
		Model:        "llama-3.3-70b-versatile",
		StartedAt:    started,
		EndedAt:      started.Add(1200 * time.Millisecond),
		Termination:  runtime.TerminationCompleted,
		Entries: []runtime.Entry{
			{Seq: 1, At: started, Kind: runtime.EntryMessage, Role: "user", Text: "what happened?"},
			{Seq: 2, At: started, Kind: runtime.EntryBudget, Budget: &runtime.Snapshot{StepsLeft: 8, TokensLeft: 120000}},
			{Seq: 3, At: started, Kind: runtime.EntryMessage, Role: "assistant", Text: "Twelve events."},
		},
		Usage: runtime.RunUsage{InputTokens: 900, OutputTokens: 120, TotalTokens: 1020, ModelCalls: 1},
	}
}

func TestPostgresSinkSavesAndReadsBack(t *testing.T) {
	h := testutil.New(t)
	ctx := context.Background()
	sink := runtime.NewPostgresSink(h.Pool)

	traj := trajectory(h.OrgID, "run_store_basic")
	if err := sink.Save(ctx, traj); err != nil {
		t.Fatalf("save: %v", err)
	}

	var (
		agentID, tier, model, termination string
		version, modelCalls               int
		total                             int64
		entries                           []byte
	)
	err := h.Pool.QueryRow(ctx, `
		SELECT agent_id, agent_version, tier, model, termination, total_tokens, model_calls, entries
		  FROM agent_runs WHERE run_id = $1`, traj.RunID,
	).Scan(&agentID, &version, &tier, &model, &termination, &total, &modelCalls, &entries)
	if err != nil {
		t.Fatalf("read back: %v", err)
	}

	if agentID != "activity-agent" || version != 3 {
		t.Errorf("stored %s v%d, want activity-agent v3", agentID, version)
	}
	if termination != string(runtime.TerminationCompleted) {
		t.Errorf("termination = %q, want Completed", termination)
	}
	// Both the tier asked for and the model that answered, so a trajectory read
	// a year later does not require knowing that week's routing.
	if tier != "balanced" || model != "llama-3.3-70b-versatile" {
		t.Errorf("tier/model = %q/%q, want balanced/llama-3.3-70b-versatile", tier, model)
	}
	if total != 1020 || modelCalls != 1 {
		t.Errorf("usage = %d tokens over %d calls, want 1020/1", total, modelCalls)
	}

	var round []runtime.Entry
	if err := json.Unmarshal(entries, &round); err != nil {
		t.Fatalf("entries did not round-trip: %v", err)
	}
	if len(round) != 3 || round[0].Role != "user" || round[2].Text != "Twelve events." {
		t.Errorf("entries round-tripped as %+v", round)
	}
}

func TestPostgresSinkIsIdempotentPerRun(t *testing.T) {
	// A run id is generated once and written once. A second save is a retry of
	// one that already landed — failing it would turn a harmless duplicate
	// into a reported error on a run that succeeded.
	h := testutil.New(t)
	ctx := context.Background()
	sink := runtime.NewPostgresSink(h.Pool)

	traj := trajectory(h.OrgID, "run_store_twice")
	if err := sink.Save(ctx, traj); err != nil {
		t.Fatalf("first save: %v", err)
	}
	if err := sink.Save(ctx, traj); err != nil {
		t.Fatalf("second save should be a no-op, got: %v", err)
	}

	var n int
	if err := h.Pool.QueryRow(ctx,
		`SELECT count(*) FROM agent_runs WHERE run_id = $1`, traj.RunID).Scan(&n); err != nil {
		t.Fatalf("count: %v", err)
	}
	if n != 1 {
		t.Errorf("%d rows for one run id, want 1", n)
	}
}

func TestPostgresSinkRefusesRunsItCannotAttribute(t *testing.T) {
	h := testutil.New(t)
	ctx := context.Background()
	sink := runtime.NewPostgresSink(h.Pool)

	cases := map[string]*runtime.Trajectory{
		"no run id": func() *runtime.Trajectory {
			tr := trajectory(h.OrgID, "")
			return tr
		}(),
		// I5: tenancy is not optional. Caught here so the failure names the
		// cause rather than surfacing as a NOT NULL violation.
		"no org": func() *runtime.Trajectory {
			tr := trajectory("", "run_no_org")
			return tr
		}(),
		// An invented termination must never reach the column that evals and
		// dashboards group by.
		"invented termination": func() *runtime.Trajectory {
			tr := trajectory(h.OrgID, "run_bad_term")
			tr.Termination = "Finished"
			return tr
		}(),
	}

	for name, tr := range cases {
		if err := sink.Save(ctx, tr); err == nil {
			t.Errorf("%s: save should have been refused", name)
		}
	}
}

func TestPostgresSinkStoresAnEmptyTrajectory(t *testing.T) {
	// A run that recorded nothing is still a run worth keeping, and a nil
	// slice marshals to "null", which the jsonb_typeof CHECK refuses.
	h := testutil.New(t)
	ctx := context.Background()
	sink := runtime.NewPostgresSink(h.Pool)

	traj := trajectory(h.OrgID, "run_store_empty")
	traj.Entries = nil
	traj.Termination = runtime.TerminationBudgetExceeded

	if err := sink.Save(ctx, traj); err != nil {
		t.Fatalf("save: %v", err)
	}

	var kind string
	if err := h.Pool.QueryRow(ctx,
		`SELECT jsonb_typeof(entries) FROM agent_runs WHERE run_id = $1`, traj.RunID).Scan(&kind); err != nil {
		t.Fatalf("read back: %v", err)
	}
	if kind != "array" {
		t.Errorf("entries stored as %q, want array", kind)
	}
}