oikos/cmd/nomos/store.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"regexp"
	"strings"
	"time"

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

const maxToolResultSize = 4096

type store struct {
	pool *pgxpool.Pool
}

func newStore(ctx context.Context, databaseURL string) (*store, error) {
	if databaseURL == "" {
		return nil, nil
	}
	pool, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		return nil, fmt.Errorf("connect db: %w", err)
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("ping db: %w", err)
	}
	return &store{pool: pool}, nil
}

func (s *store) close() {
	if s.pool != nil {
		s.pool.Close()
	}
}

// session is a chat session elevated to a task: goal-structured work with a
// lifecycle status and an outcome (see migration 018 / the task-board plan).
// Outcome/Summary/EntityID are empty until set, hence omitempty.
type session struct {
	ID           string    `json:"id"`
	Title        string    `json:"title"`
	Actor        string    `json:"actor"`
	Goal         string    `json:"goal"`
	Status       string    `json:"status"`
	Outcome      string    `json:"outcome,omitempty"`
	Summary      string    `json:"summary,omitempty"`
	EntityID     string    `json:"entity_id,omitempty"`
	CreatedAt    time.Time `json:"created_at"`
	LastActiveAt time.Time `json:"last_active_at"`
}

type message struct {
	ID        string          `json:"id"`
	SessionID string          `json:"session_id"`
	Role      string          `json:"role"`
	Content   json.RawMessage `json:"content"`
	CreatedAt time.Time       `json:"created_at"`
}

func (s *store) createSession(ctx context.Context, title string) (*session, error) {
	if s == nil {
		return &session{ID: "ephemeral", Title: title, Actor: "agent:nomos", Status: "active"}, nil
	}
	var id string
	err := s.pool.QueryRow(ctx,
		`INSERT INTO agent_sessions (title, actor) VALUES ($1, 'agent:nomos') RETURNING id`,
		title).Scan(&id)
	if err != nil {
		return nil, err
	}
	// Give the task its own entity so knowledge and involved-entity edges hang
	// off the existing relationships graph. Best-effort: a failure here must not
	// block the chat — the session is usable without a graph anchor.
	entityID := s.createTaskEntity(ctx, id, title)
	return &session{ID: id, Title: title, Actor: "agent:nomos", Status: "active",
		EntityID: entityID, CreatedAt: time.Now(), LastActiveAt: time.Now()}, nil
}

// createTaskEntity creates (or reuses) the task:<session-id> entity that
// anchors this task's knowledge and involved-entity relationships, and records
// it on the session. Returns the entity id, or "" on failure — non-fatal, see
// caller. Requires the 'task' entity type (seeds/ontology.yaml).
func (s *store) createTaskEntity(ctx context.Context, sessionID, title string) string {
	entityID, _ := uuid.NewV7()
	slug := "task:" + sessionID
	// name is UNIQUE(type,name) and chat titles collide ("hi" ×6), so key the
	// name on the session id and keep the human title in attributes for display.
	name := "task " + sessionID
	attrs, _ := json.Marshal(map[string]any{"title": title})
	if err := s.pool.QueryRow(ctx, `
		INSERT INTO entities (id, slug, type, name, attributes)
		VALUES ($1, $2, 'task', $3, $4)
		ON CONFLICT (slug) DO UPDATE SET updated_at = now()
		RETURNING id`, entityID, slug, name, string(attrs)).Scan(&entityID); err != nil {
		slog.Warn("nomos: could not create task entity", "session", sessionID, "error", err)
		return ""
	}
	if _, err := s.pool.Exec(ctx,
		`UPDATE agent_sessions SET entity_id = $1 WHERE id = $2`, entityID, sessionID); err != nil {
		slog.Warn("nomos: could not link task entity", "session", sessionID, "error", err)
	}
	return entityID.String()
}

func (s *store) saveMessage(ctx context.Context, sessionID, role string, content json.RawMessage) error {
	if s == nil {
		return nil
	}
	_, err := s.pool.Exec(ctx,
		`INSERT INTO agent_messages (session_id, role, content) VALUES ($1, $2, $3)`,
		sessionID, role, truncateToolResults(content))
	return err
}

// insertMessageReturningID and updateMessage exist for the auto-continuation
// worker's live-progress persistence (see continue.go): rather than saving
// one message only once the whole continuation finishes — which could be
// several minutes of silence in the UI even though frontend polling exists —
// the worker inserts a placeholder immediately and updates the SAME row as
// each tool call completes, so a poller sees individual steps land, not just
// a final rolled-up summary.
func (s *store) insertMessageReturningID(ctx context.Context, sessionID, role string, content json.RawMessage) (uuid.UUID, error) {
	if s == nil {
		return uuid.Nil, nil
	}
	var id uuid.UUID
	err := s.pool.QueryRow(ctx,
		`INSERT INTO agent_messages (session_id, role, content) VALUES ($1, $2, $3) RETURNING id`,
		sessionID, role, truncateToolResults(content)).Scan(&id)
	return id, err
}

func (s *store) updateMessage(ctx context.Context, id uuid.UUID, content json.RawMessage) error {
	if s == nil || id == uuid.Nil {
		return nil
	}
	_, err := s.pool.Exec(ctx,
		`UPDATE agent_messages SET content = $2 WHERE id = $1`,
		id, truncateToolResults(content))
	return err
}

func truncateToolResults(content json.RawMessage) json.RawMessage {
	var m map[string]any
	if err := json.Unmarshal(content, &m); err != nil {
		return content
	}
	toolCalls, ok := m["tool_calls"].([]any)
	if !ok || len(toolCalls) == 0 {
		return content
	}
	changed := false
	for i, raw := range toolCalls {
		tc, ok := raw.(map[string]any)
		if !ok {
			continue
		}
		if result, ok := tc["result"]; ok {
			resultJSON, _ := json.Marshal(result)
			if len(resultJSON) > maxToolResultSize {
				tc["result"] = string(resultJSON[:maxToolResultSize]) + fmt.Sprintf("...truncated (%d bytes total)", len(resultJSON))
				toolCalls[i] = tc
				changed = true
			}
		}
	}
	if !changed {
		return content
	}
	m["tool_calls"] = toolCalls
	out, err := json.Marshal(m)
	if err != nil {
		return content
	}
	return out
}

func (s *store) touchSession(ctx context.Context, id string) {
	if s != nil {
		s.pool.Exec(ctx, `UPDATE agent_sessions SET last_active_at=now() WHERE id=$1`, id)
	}
}

func (s *store) listSessions(ctx context.Context) ([]session, error) {
	if s == nil {
		return nil, nil
	}
	rows, err := s.pool.Query(ctx,
		`SELECT id, title, actor, goal, status, COALESCE(outcome, ''), summary,
		        COALESCE(entity_id::text, ''), created_at, last_active_at
		 FROM agent_sessions ORDER BY last_active_at DESC LIMIT 50`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []session
	for rows.Next() {
		var sess session
		if err := rows.Scan(&sess.ID, &sess.Title, &sess.Actor, &sess.Goal, &sess.Status,
			&sess.Outcome, &sess.Summary, &sess.EntityID, &sess.CreatedAt, &sess.LastActiveAt); err != nil {
			return nil, err
		}
		out = append(out, sess)
	}
	return out, rows.Err()
}

func (s *store) getMessages(ctx context.Context, sessionID string) ([]message, error) {
	if s == nil {
		return nil, nil
	}
	rows, err := s.pool.Query(ctx,
		`SELECT id, session_id, role, content, created_at FROM agent_messages WHERE session_id=$1 ORDER BY created_at ASC`,
		sessionID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []message
	for rows.Next() {
		var m message
		if err := rows.Scan(&m.ID, &m.SessionID, &m.Role, &m.Content, &m.CreatedAt); err != nil {
			return nil, err
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

func (s *store) deleteSession(ctx context.Context, id string) error {
	if s == nil {
		return nil
	}
	_, err := s.pool.Exec(ctx, `DELETE FROM agent_messages WHERE session_id = $1`, id)
	if err != nil {
		return err
	}
	_, err = s.pool.Exec(ctx, `DELETE FROM agent_sessions WHERE id = $1`, id)
	return err
}

func (s *store) updateSessionTitle(ctx context.Context, id, title string) error {
	if s == nil {
		return nil
	}
	_, err := s.pool.Exec(ctx, `UPDATE agent_sessions SET title = $1 WHERE id = $2`, title, id)
	return err
}

// resolveAgentID looks up the UUID of the agent entity (e.g. "agent:nomos").
// Returns uuid.Nil if the store is absent or the slug is unknown.
func (s *store) resolveAgentID(ctx context.Context, slug string) uuid.UUID {
	if s == nil {
		return uuid.Nil
	}
	var id uuid.UUID
	if err := s.pool.QueryRow(ctx, `SELECT id FROM entities WHERE slug = $1`, slug).Scan(&id); err != nil {
		return uuid.Nil
	}
	return id
}

// linkExecution records that a gated execution was initiated by a chat
// session, so the auto-continuation worker can feed its result back to that
// session when it finishes. Idempotent — the same execution may appear in
// several tool results across a turn.
func (s *store) linkExecution(ctx context.Context, execID uuid.UUID, sessionID string) {
	if s == nil || execID == uuid.Nil || sessionID == "" || sessionID == "ephemeral" {
		return
	}
	s.pool.Exec(ctx, `
		INSERT INTO nomos_plan_executions (execution_id, session_id)
		VALUES ($1, $2) ON CONFLICT (execution_id) DO NOTHING`, execID, sessionID)
}

// taskSlugRe matches an entity slug: a lowercase type prefix then colon-
// separated segments (host:hubris, lxc:caddy, check:ping:8cf). Mirrors the
// frontend SessionGraph regex so the panel and the involves-graph agree on
// what counts as an entity reference.
var taskSlugRe = regexp.MustCompile(`[a-z][a-z-]*:[a-z0-9][a-z0-9._/-]*(?::[a-z0-9._/-]+)*`)

// touchExcludedTypes are entity types too noisy to record as task involvement:
// a health question names dozens of check:… slugs, executions/tasks are
// bookkeeping, not things the task "worked on".
var touchExcludedTypes = map[string]bool{"check": true, "execution": true, "task": true}

// recordTouched links the task to every entity referenced in a tool call's
// args (task —involves→ entity) and publishes one entity.touched event per
// entity so the live context panel can pulse it. Best-effort: it never blocks
// or fails the tool call. Only args are inspected — what the agent chose to act
// on — never results, since a single bulk query result would otherwise pull the
// whole fleet into the task's graph.
func (s *store) recordTouched(ctx context.Context, sessionID, toolName string, args map[string]any) {
	if s == nil || sessionID == "" || sessionID == "ephemeral" || len(args) == 0 {
		return
	}
	slugs := map[string]struct{}{}
	collectTaskSlugs(args, slugs)
	if len(slugs) == 0 {
		return
	}
	var taskEntityID uuid.UUID
	if err := s.pool.QueryRow(ctx,
		`SELECT entity_id FROM agent_sessions WHERE id = $1`, sessionID).Scan(&taskEntityID); err != nil || taskEntityID == uuid.Nil {
		return // no task entity to anchor edges on
	}
	q := sqlcgen.New(s.pool)
	for slug := range slugs {
		var entID uuid.UUID
		var etype string
		if err := s.pool.QueryRow(ctx,
			`SELECT id, type FROM entities WHERE slug = $1`, slug).Scan(&entID, &etype); err != nil {
			continue // unknown slug — skip
		}
		if touchExcludedTypes[etype] || entID == taskEntityID {
			continue
		}
		// Idempotent involves edge (task → entity), same guard as upsert_knowledge.
		s.pool.Exec(ctx, `
			INSERT INTO relationships (source_id, target_id, type, attributes, valid_from)
			SELECT $1, $2, 'involves', '{"by":"nomos"}'::jsonb, now()
			WHERE NOT EXISTS (
				SELECT 1 FROM relationships
				WHERE source_id = $1 AND target_id = $2 AND type = 'involves' AND valid_to IS NULL)`,
			taskEntityID, entID)
		// Live pulse for the panel. correlation_id = sessionID lets the frontend
		// filter to the active task.
		_ = observability.Event(ctx, q, "entity.touched", &entID, "info", "nomos", sessionID,
			map[string]any{"slug": slug, "tool": toolName})
	}
}

// collectTaskSlugs recursively pulls entity slugs out of tool-call args,
// mirroring the frontend's collectSlugs so both sides see the same references.
func collectTaskSlugs(v any, out map[string]struct{}) {
	switch t := v.(type) {
	case string:
		for _, m := range taskSlugRe.FindAllString(t, -1) {
			out[strings.TrimRight(m, ".,;)]")] = struct{}{}
		}
	case []any:
		for _, e := range t {
			collectTaskSlugs(e, out)
		}
	case map[string]any:
		for _, e := range t {
			collectTaskSlugs(e, out)
		}
	}
}

// pendingContinuation is one finished execution whose result hasn't yet been
// fed back to its originating session.
type pendingContinuation struct {
	ExecID    uuid.UUID
	SessionID string
	Status    string
	Result    string
	Action    string
}

// pendingContinuations returns executions that have reached a terminal state
// but haven't been continued yet — the worker's work list. Bounded so one
// tick can't fan out unboundedly.
func (s *store) pendingContinuations(ctx context.Context, limit int) []pendingContinuation {
	if s == nil {
		return nil
	}
	rows, err := s.pool.Query(ctx, `
		SELECT l.execution_id, l.session_id, e.status,
		       COALESCE(e.result::text, ''), COALESCE(e.action, '')
		FROM nomos_plan_executions l
		JOIN executions e ON e.entity_id = l.execution_id
		WHERE l.continued_at IS NULL
		  AND e.status IN ('completed', 'failed', 'cancelled', 'denied', 'revoked')
		ORDER BY l.created_at
		LIMIT $1`, limit)
	if err != nil {
		return nil
	}
	defer rows.Close()
	var out []pendingContinuation
	for rows.Next() {
		var p pendingContinuation
		if err := rows.Scan(&p.ExecID, &p.SessionID, &p.Status, &p.Result, &p.Action); err == nil {
			out = append(out, p)
		}
	}
	return out
}

// markContinued stamps an execution as fed-back so the worker won't process it
// again (prevents an auto-continuation loop).
func (s *store) markContinued(ctx context.Context, execID uuid.UUID) {
	if s == nil {
		return
	}
	s.pool.Exec(ctx, `UPDATE nomos_plan_executions SET continued_at = now() WHERE execution_id = $1`, execID)
}

// assentWindowActive reports whether this agent currently has an open assent
// window — the scope gate for auto-continuation. We only auto-continue
// executions that are part of an approved plan, never stray one-off actions.
func (s *store) assentWindowActive(ctx context.Context, agentID uuid.UUID) bool {
	if s == nil || agentID == uuid.Nil {
		return false
	}
	var expires time.Time
	key := "assent_window.agent:" + agentID.String()
	if err := s.pool.QueryRow(ctx, `SELECT value::timestamptz FROM autonomy_settings WHERE key = $1`, key).Scan(&expires); err != nil {
		return false
	}
	return time.Now().Before(expires)
}

// destructiveWindowDuration is intentionally shorter than the general assent
// window (30 min): it's a narrow, scoped grant for a multi-step DESTRUCTIVE
// recovery (e.g. "stop then destroy this specific half-provisioned
// container"), not a standing license to destroy things.
const destructiveWindowDuration = 15 * time.Minute

// destructiveWindowKey scopes the grant to one agent AND one target entity —
// an explicit typed confirmation ("I confirm") for a destructive action on
// target X must never be read as authorizing a destructive action on target Y.
func destructiveWindowKey(agentID uuid.UUID, targetSlug string) string {
	return "destructive_window.agent:" + agentID.String() + ".target:" + targetSlug
}

// openDestructiveWindow records a short, target-scoped grant after an
// operator's EXPLICIT typed confirmation (never loose assent) authorized a
// destructive action. Real case this exists for: recovering a failed destroy
// took "stop" (destructive) then "destroy" (destructive) — same container,
// two separate typed-confirmation round trips, because each was gated
// independently. One explicit confirmation on a target should cover the
// short follow-up sequence needed to finish what was just confirmed.
func (s *store) openDestructiveWindow(ctx context.Context, agentID uuid.UUID, targetSlug string) {
	if s == nil || agentID == uuid.Nil || targetSlug == "" {
		return
	}
	expires := time.Now().Add(destructiveWindowDuration).UTC().Format(time.RFC3339)
	s.pool.Exec(ctx, `INSERT INTO autonomy_settings (key, value) VALUES ($1, $2)
		ON CONFLICT (key) DO UPDATE SET value = $2`, destructiveWindowKey(agentID, targetSlug), expires)
}

// destructiveWindowActive reports whether target has a live, explicitly-
// confirmed destructive grant for this agent.
func (s *store) destructiveWindowActive(ctx context.Context, agentID uuid.UUID, targetSlug string) bool {
	if s == nil || agentID == uuid.Nil || targetSlug == "" {
		return false
	}
	var expires time.Time
	if err := s.pool.QueryRow(ctx, `SELECT value::timestamptz FROM autonomy_settings WHERE key = $1`,
		destructiveWindowKey(agentID, targetSlug)).Scan(&expires); err != nil {
		return false
	}
	return time.Now().Before(expires)
}

// executionTarget resolves the target entity slug for an execution — used to
// scope the destructive window to the right entity when a chat-assent typed
// confirmation grants a destructive execution.
func (s *store) executionTarget(ctx context.Context, execID uuid.UUID) string {
	if s == nil {
		return ""
	}
	var slug string
	s.pool.QueryRow(ctx, `
		SELECT e.slug FROM executions ex JOIN entities e ON e.id = ex.target_entity_id
		WHERE ex.entity_id = $1`, execID).Scan(&slug)
	return slug
}

// logActivity records a tool call. agent_id is the agent entity UUID and is
// NOT NULL in the schema, so we skip logging when it can't be resolved.
// The (nullable) session_id column carries the conversation id.
func (s *store) logActivity(ctx context.Context, agentID uuid.UUID, sessionID, toolName, inputSummary, outputSummary string, durationMs int, success bool, correlationID string) {
	if s == nil || agentID == uuid.Nil {
		return
	}
	s.pool.Exec(ctx, `
		INSERT INTO agent_activity
		    (agent_id, session_id, activity_type, tool_name, input_summary, output_summary,
		     duration_ms, success, correlation_id)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`,
		agentID, sessionID, "tool_call", toolName, inputSummary, outputSummary,
		durationMs, success, correlationID)
}