oikos/cmd/nomos/store_test.go

package main

// Integration tests against a real Postgres, mirroring
// internal/db/integration_test.go's pattern: guarded by
// OIKOS_TEST_DATABASE_URL (skipped when unset), throwaway database per run,
// full migrations applied, dropped on cleanup. Run with:
//
//   docker compose up -d postgres
//   OIKOS_TEST_DATABASE_URL="postgres://oikos:oikos_dev@localhost:5432/oikos?sslmode=disable" go test ./cmd/nomos/

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"os"
	"strings"
	"testing"

	"github.com/dtoro/oikos/internal/db"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

// newTestStore creates a throwaway, fully-migrated database and returns a
// *store connected to it, cleaned up (including a matching task:<session>
// entity type in the ontology, needed by createTaskEntity/proposePlan tests)
// via t.Cleanup.
func newTestStore(t *testing.T) *store {
	t.Helper()
	baseURL := os.Getenv("OIKOS_TEST_DATABASE_URL")
	if baseURL == "" {
		t.Skip("OIKOS_TEST_DATABASE_URL not set — skipping integration test")
	}
	ctx := context.Background()

	admin, err := pgx.Connect(ctx, baseURL)
	if err != nil {
		t.Fatalf("connect admin: %v", err)
	}
	dbName := fmt.Sprintf("oikos_test_nomos_%08x", rand.Int63())
	if _, err := admin.Exec(ctx, "CREATE DATABASE "+dbName); err != nil {
		admin.Close(ctx)
		t.Fatalf("create test db: %v", err)
	}
	admin.Close(ctx)

	testURL := swapTestDatabase(baseURL, dbName)
	pool, err := db.New(ctx, testURL)
	if err != nil {
		t.Fatalf("connect test db: %v", err)
	}
	t.Cleanup(func() {
		pool.Close()
		admin, err := pgx.Connect(ctx, baseURL)
		if err == nil {
			admin.Exec(ctx, "DROP DATABASE IF EXISTS "+dbName+" WITH (FORCE)")
			admin.Close(ctx)
		}
	})
	if err := pool.Migrate(ctx); err != nil {
		t.Fatalf("migrate: %v", err)
	}

	// session_plan_steps/session_questions tests don't need the ontology
	// seed, but createTaskEntity's INSERT INTO entities (type='task') has an
	// FK to entity_types — seed the minimal rows it needs directly rather
	// than pulling in the full seeds/ontology.yaml ingest path.
	if _, err := pool.Exec(ctx, `
		INSERT INTO entity_types (name, domain, layer) VALUES ('entity', 'meta', 'meta')
		ON CONFLICT DO NOTHING;
		INSERT INTO entity_types (name, parent_type, domain, layer) VALUES ('task', 'entity', 'cognition', 'cognition')
		ON CONFLICT DO NOTHING;`); err != nil {
		t.Fatalf("seed minimal ontology: %v", err)
	}

	return &store{pool: pool.Pool}
}

func swapTestDatabase(url, dbName string) string {
	qi := strings.Index(url, "?")
	params, base := "", url
	if qi >= 0 {
		params = url[qi:]
		base = url[:qi]
	}
	si := strings.LastIndex(base, "/")
	return base[:si+1] + dbName + params
}

// TestGetRecentMessages_Truncation is the concrete proof for fix A2 of
// plans/2026-07-11-nomos-agent-code-review.md: chatWith used to replay a
// session's ENTIRE history on every turn with no bound. getRecentMessages
// caps that; this test checks both sides — under the limit, nothing is
// dropped and truncated=false; over it, only the most recent `limit` come
// back, in chronological order, with truncated=true.
func TestGetRecentMessages_Truncation(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()

	sess, err := s.createSession(ctx, "history window test")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}

	const total = 35
	const limit = 30
	for i := 0; i < total; i++ {
		role := "user"
		if i%2 == 1 {
			role = "assistant"
		}
		body := fmt.Appendf(nil, `{"role":%q,"text":"msg-%d"}`, role, i)
		if err := s.saveMessage(ctx, sess.ID, role, body); err != nil {
			t.Fatalf("saveMessage %d: %v", i, err)
		}
	}

	msgs, truncated, err := s.getRecentMessages(ctx, sess.ID, limit)
	if err != nil {
		t.Fatalf("getRecentMessages: %v", err)
	}
	if !truncated {
		t.Errorf("truncated = false, want true (%d messages > limit %d)", total, limit)
	}
	if len(msgs) != limit {
		t.Fatalf("got %d messages, want %d", len(msgs), limit)
	}
	// Chronological order: the oldest of the RETAINED messages should be the
	// (total-limit)-th one saved (msg-5, since msg-0..4 were dropped), and
	// the last should be the most recently saved (msg-34).
	wantFirst := fmt.Sprintf("msg-%d", total-limit)
	wantLast := fmt.Sprintf("msg-%d", total-1)
	if got := extractText(msgs[0].Content); got != wantFirst {
		t.Errorf("first retained message = %q, want %q", got, wantFirst)
	}
	if got := extractText(msgs[len(msgs)-1].Content); got != wantLast {
		t.Errorf("last retained message = %q, want %q", got, wantLast)
	}

	// Under the limit: nothing dropped.
	sess2, err := s.createSession(ctx, "small session")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}
	for i := 0; i < 5; i++ {
		body := fmt.Appendf(nil, `{"role":"user","text":"msg-%d"}`, i)
		if err := s.saveMessage(ctx, sess2.ID, "user", body); err != nil {
			t.Fatalf("saveMessage: %v", err)
		}
	}
	msgs2, truncated2, err := s.getRecentMessages(ctx, sess2.ID, limit)
	if err != nil {
		t.Fatalf("getRecentMessages (small): %v", err)
	}
	if truncated2 {
		t.Errorf("truncated = true for a 5-message session under a %d limit, want false", limit)
	}
	if len(msgs2) != 5 {
		t.Errorf("got %d messages, want 5", len(msgs2))
	}
}

// TestProposePlan_RefuseInFlight is the concrete proof for the plan-drift
// fix (2026-07-14, "plan added twice in the sidebar"): proposePlan must
// REPLACE the step list only while every existing step is still 'pending'
// (a genuine pre-execution revision), and REFUSE the call once any step has
// started. The prior append-mode safety net (commit 5384499) preserved
// history but duplicated the plan in the sidebar when the agent re-proposed
// on "proceed". Refusing is the correct default — the agent must advance
// with update_plan_step + run.
func TestProposePlan_RefuseInFlight(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()

	sess, err := s.createSession(ctx, "plan refuse test")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}

	// First call: no steps exist yet — must persist as-is (replace mode,
	// trivially: nothing to replace).
	out1, err := s.proposePlan(ctx, sess.ID, []planStepInput{{Title: "Step A"}})
	if err != nil {
		t.Fatalf("proposePlan #1: %v", err)
	}
	if len(out1) != 1 || out1[0]["seq"] != 1 {
		t.Fatalf("proposePlan #1 = %+v, want one step at seq 1", out1)
	}
	if out1[0]["generation"] != 1 {
		t.Fatalf("proposePlan #1 generation = %v, want 1", out1[0]["generation"])
	}

	// Mark step 1 as started.
	if err := s.updatePlanStep(ctx, sess.ID, 1, "running", ""); err != nil {
		t.Fatalf("updatePlanStep: %v", err)
	}

	// Second call, simulating a model that re-proposes mid-flight (the
	// operator-reported "proceed" bug): since step 1 has left 'pending',
	// this MUST refuse with errPlanInFlight, not append or replace.
	_, err = s.proposePlan(ctx, sess.ID, []planStepInput{{Title: "Step B"}})
	if !errors.Is(err, errPlanInFlight) {
		t.Fatalf("proposePlan #2: err = %v, want errPlanInFlight (refuse mid-flight re-proposal)", err)
	}

	// The original step 1 must be untouched — not erased, not appended to.
	steps, err := s.getPlanSteps(ctx, sess.ID)
	if err != nil {
		t.Fatalf("getPlanSteps: %v", err)
	}
	if len(steps) != 1 {
		t.Fatalf("got %d persisted steps, want 1 (refused call must not mutate the plan)", len(steps))
	}
	if steps[0].Title != "Step A" || steps[0].Status != "running" {
		t.Errorf("step 1 = %+v, want Step A still running (refused call must not touch it)", steps[0])
	}

	// Third call BEFORE anything runs on a fresh session: every step is
	// still pending, so this must REPLACE, not refuse.
	sess2, err := s.createSession(ctx, "plan replace test")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}
	if _, err := s.proposePlan(ctx, sess2.ID, []planStepInput{{Title: "Original"}}); err != nil {
		t.Fatalf("proposePlan (initial): %v", err)
	}
	if _, err := s.proposePlan(ctx, sess2.ID, []planStepInput{{Title: "Revised"}}); err != nil {
		t.Fatalf("proposePlan (revise before execution): %v", err)
	}
	revisedSteps, err := s.getPlanSteps(ctx, sess2.ID)
	if err != nil {
		t.Fatalf("getPlanSteps: %v", err)
	}
	if len(revisedSteps) != 1 || revisedSteps[0].Title != "Revised" {
		t.Fatalf("got %+v, want a single 'Revised' step (pre-execution revise must replace, not refuse)", revisedSteps)
	}
	if revisedSteps[0].Generation != 1 {
		t.Fatalf("revised step generation = %d, want 1 (fresh-start after DELETE resets generation)", revisedSteps[0].Generation)
	}
}

// TestHadDiscoveryAndWriteback is the store-level proof for D.1 (refuse
// complete_task when discovery ran without writeback). hadDiscovery must
// report true only after a successful `run` call; hadEntityWriteback must
// report true only after a successful update_entity_attributes or
// create_relationship call. The D.1 gate in tasks.go combines these: refuse
// success when hadDiscovery && !hadEntityWriteback.
func TestHadDiscoveryAndWriteback(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()

	sess, err := s.createSession(ctx, "discovery test")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}

	// Before any tool calls: no discovery, no writeback.
	if s.hadDiscovery(ctx, sess.ID) {
		t.Fatal("hadDiscovery = true before any tool calls, want false")
	}
	if s.hadEntityWriteback(ctx, sess.ID) {
		t.Fatal("hadEntityWriteback = true before any tool calls, want false")
	}

	// A `run` call (discovery) — should set hadDiscovery, not hadEntityWriteback.
	agentID := uuid.New()
	s.logActivity(ctx, agentID, sess.ID, "run", nil, "", "uptime output", 100, true, "corr-1")
	if !s.hadDiscovery(ctx, sess.ID) {
		t.Fatal("hadDiscovery = false after a successful run call, want true")
	}
	if s.hadEntityWriteback(ctx, sess.ID) {
		t.Fatal("hadEntityWriteback = true after only a run call, want false")
	}

	// A failed run call should NOT count as discovery (no facts learned).
	sess2, err := s.createSession(ctx, "failed discovery test")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}
	s.logActivity(ctx, agentID, sess2.ID, "run", nil, "", "ssh timeout", 100, false, "corr-2")
	if s.hadDiscovery(ctx, sess2.ID) {
		t.Fatal("hadDiscovery = true after a failed run call, want false (no facts learned)")
	}

	// A get_entity call should NOT count as discovery (DB lookup, not live state).
	sess3, err := s.createSession(ctx, "lookup test")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}
	s.logActivity(ctx, agentID, sess3.ID, "get_entity", nil, "", "entity row", 10, true, "corr-3")
	if s.hadDiscovery(ctx, sess3.ID) {
		t.Fatal("hadDiscovery = true after get_entity, want false (DB lookups are not discovery)")
	}

	// update_entity_attributes sets hadEntityWriteback.
	sess4, err := s.createSession(ctx, "writeback test")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}
	s.logActivity(ctx, agentID, sess4.ID, "update_entity_attributes", nil, "", "ok", 10, true, "corr-4")
	if !s.hadEntityWriteback(ctx, sess4.ID) {
		t.Fatal("hadEntityWriteback = false after update_entity_attributes, want true")
	}
	// And the discovery+writeback combination (the conv3 scenario).
	s.logActivity(ctx, agentID, sess4.ID, "run", nil, "", "apt-get update output", 100, true, "corr-5")
	if !s.hadDiscovery(ctx, sess4.ID) {
		t.Fatal("hadDiscovery = false after run+writeback, want true")
	}
	if !s.hadEntityWriteback(ctx, sess4.ID) {
		t.Fatal("hadEntityWriteback = false after run+writeback, want true")
	}
}