oikos/cmd/nomos/store.go

package main

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

	"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()
	}
}

type session struct {
	ID           string    `json:"id"`
	Title        string    `json:"title"`
	Actor        string    `json:"actor"`
	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"}, 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
	}
	return &session{ID: id, Title: title, Actor: "agent:nomos", CreatedAt: time.Now(), LastActiveAt: time.Now()}, nil
}

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, 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.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)
}

// 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)
}