feat(tasks): add runtime controls and inherited context

This commit is contained in:
pphua committed 2026-08-24 19:29:00 +08:00
1 parent 18926c129f
commit 4392b4fee1
50 files changed
+8291 -663

No files matched your search

+2 -2
View File
@@ -103,7 +103,7 @@ type builtinAgent struct {
key, name, role, desc string
vars []promptVar
interactiveShell bool // 建行时的默认交互式 shell 开关;ON CONFLICT 不覆盖用户后续手动开关
runSeconds *int // 建行时的单次 run 墙钟上限(秒);nil=用表默认(600),0=不限时
runSeconds *int // 建行时的单次 run 墙钟上限(秒);nil=用种子默认(1200),0=不限时
}
type promptVar struct{ name, desc, example, source string }
@@ -145,7 +145,7 @@ func (d *DB) seedBuiltins() error {
var agentID int64
err := d.QueryRow(`
INSERT INTO agents(key, name, description, role, builtin, enabled, interactive_shell, run_seconds)
VALUES ($1, $2, NULLIF($3,''), $4, true, true, $5, COALESCE($6, 600))
VALUES ($1, $2, NULLIF($3,''), $4, true, true, $5, COALESCE($6, 1200))
ON CONFLICT (key) DO UPDATE SET name = EXCLUDED.name, description = EXCLUDED.description
RETURNING id`, a.key, a.name, a.desc, a.role, a.interactiveShell, a.runSeconds).Scan(&agentID)
if err != nil {
+314 -52
View File
@@ -3,11 +3,17 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"sort"
"strings"
"time"
)
// ErrIntentStateConflict marks a failed compare-and-set transition. Callers use
// it to distinguish an expected control race (HTTP 409) from a storage failure.
var ErrIntentStateConflict = errors.New("intent state changed concurrently")
// utf8Clean makes a string safe for a PostgreSQL text column. It (1) replaces
// invalid UTF-8 byte sequences with U+FFFD and (2) strips NUL (0x00) bytes.
// Tool output (nmap/curl/… stdout) can carry raw/truncated bytes that PostgreSQL's
@@ -25,34 +31,40 @@ func utf8Clean(s string) string {
// Node is a typed reasoning node (= old task_nodes). kind ∈ goal|intent|finding|hint.
type Node struct {
ID int64 `json:"id"`
Kind string `json:"kind"`
Payload json.RawMessage `json:"payload"`
Priority int `json:"priority"`
State string `json:"state"`
Origin string `json:"origin,omitempty"`
Owner string `json:"owner,omitempty"`
Anchors []int64 `json:"anchors,omitempty"`
CreatedAt time.Time `json:"created_at"`
ID int64 `json:"id"`
Kind string `json:"kind"`
Payload json.RawMessage `json:"payload"`
Priority int `json:"priority"`
State string `json:"state"`
Origin string `json:"origin,omitempty"`
Owner string `json:"owner,omitempty"`
BlockedReason string `json:"blocked_reason,omitempty"`
Anchors []int64 `json:"anchors,omitempty"`
CreatedAt time.Time `json:"created_at"`
SourceTaskID int64 `json:"source_task_id,omitempty"`
Inherited bool `json:"inherited,omitempty"`
}
// Activity is one worker execution step (= old task_activity).
type Activity struct {
ID int64 `json:"id"`
NodeID *int64 `json:"node_id,omitempty"`
Worker string `json:"worker,omitempty"`
Kind string `json:"kind,omitempty"`
Tool string `json:"tool,omitempty"`
ToolUseID string `json:"tool_use_id,omitempty"`
IsError bool `json:"is_error"`
Summary string `json:"summary,omitempty"`
Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads
CreatedAt time.Time `json:"created_at"`
ID int64 `json:"id"`
NodeID *int64 `json:"node_id,omitempty"`
Worker string `json:"worker,omitempty"`
Kind string `json:"kind,omitempty"`
Tool string `json:"tool,omitempty"`
ToolUseID string `json:"tool_use_id,omitempty"`
IsError bool `json:"is_error"`
Summary string `json:"summary,omitempty"`
Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads
Metadata json.RawMessage `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
// Token usage — set only on the terminal kind='result' record; nil elsewhere.
InputTokens *int `json:"input_tokens,omitempty"`
OutputTokens *int `json:"output_tokens,omitempty"`
CacheReadTokens *int `json:"cache_read_tokens,omitempty"`
CacheWriteTokens *int `json:"cache_write_tokens,omitempty"`
InputTokens *int `json:"input_tokens,omitempty"`
OutputTokens *int `json:"output_tokens,omitempty"`
CacheReadTokens *int `json:"cache_read_tokens,omitempty"`
CacheWriteTokens *int `json:"cache_write_tokens,omitempty"`
SourceTaskID int64 `json:"source_task_id,omitempty"`
Inherited bool `json:"inherited,omitempty"`
}
// TokenUsage is a per-worker token aggregate (TokenStatsByWorker).
@@ -64,6 +76,17 @@ type TokenUsage struct {
CacheWriteTokens int `json:"cache_write_tokens"`
}
// SessionTokenUsage is the authoritative completed-run total for one task UI
// session. Worker sessions are keyed by intent rather than the reusable work#N
// executor name, so retries and reassignment remain attached to the same row.
type SessionTokenUsage struct {
Session string `json:"session"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheReadTokens int `json:"cache_read_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
}
// DailyTokenBucket is one day's global token aggregate across all tasks.
type DailyTokenBucket struct {
Day string `json:"day"` // "YYYY-MM-DD"
@@ -201,13 +224,10 @@ ON CONFLICT (exploration_id, src_id, rel, dst_id) DO NOTHING`, s.expID, from, re
// SetNodeState updates any node's state (never deletes).
func (s *ExplorationStore) SetNodeState(id int64, state string) error {
_, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1 WHERE id=$2 AND exploration_id=$3`, state, id, s.expID)
_, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1, blocked_reason=NULL WHERE id=$2 AND exploration_id=$3`, state, id, s.expID)
return err
}
<<<<<<< Updated upstream
// SetIntentState updates an intent node's state (done/blocked/exhausted/open).
=======
// UpdateGoalPayload rewrites a goal node's payload text (and optional vulnclass);
// scoped to kind='goal' so it can never mutate an intent/fact by id. Returns an
// error if no such goal exists in this exploration.
@@ -243,13 +263,12 @@ func (s *ExplorationStore) DeleteGoal(id int64) error {
}
// SetIntentState updates an intent node's state.
>>>>>>> Stashed changes
// ResetRunningIntents returns any intent left in 'running' back to 'open' so it
// is re-claimed. Called on startup: a 'running' intent with no live worker (a
// backend restart or crashed worker goroutine left it stuck) would otherwise spin
// forever in the UI. Workers re-run an intent from scratch, so reopening is safe.
// forever in the UI. Workers resume from their transcript, so reopening is safe.
func (s *ExplorationStore) ResetRunningIntents() (int64, error) {
res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL
res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL
WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID)
if err != nil {
return 0, err
@@ -259,12 +278,12 @@ WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID)
}
// ReopenIntent flips ONE not-successfully-finished intent (blocked/exhausted/stopped)
// back to 'open' so a worker re-claims and re-runs it from scratch (graph writes it
// already produced stay). done/open/running are left untouched. Returns whether a row
// changed (false = intent absent or not in a rerunnable state). Used by the "重跑" button.
// back to 'open' so a worker re-claims it (graph writes it already produced stay).
// The worker resumes from its prior LLM transcript instead of restarting from scratch.
// done/open/running are left untouched. Returns whether a row changed. Used by "重跑".
func (s *ExplorationStore) ReopenIntent(id int64) (bool, error) {
res, err := s.db.Exec(`UPDATE exploration_nodes
SET state='open', completed_at=NULL
SET state='open', completed_at=NULL, blocked_reason=NULL
WHERE id=$1 AND exploration_id=$2 AND kind='intent'
AND state IN ('blocked','exhausted','stopped')`, id, s.expID)
if err != nil {
@@ -276,9 +295,9 @@ WHERE id=$1 AND exploration_id=$2 AND kind='intent'
// ReopenBlockedIntents flips EVERY 'blocked' intent in this exploration back to 'open'
// (batch rerun after e.g. an LLM/network outage that blocked many at once). Returns the
// number reopened. Workers re-run each from scratch; kept graph writes remain.
// number reopened. Workers resume from their transcripts; kept graph writes remain.
func (s *ExplorationStore) ReopenBlockedIntents() (int64, error) {
res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL
res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL
WHERE exploration_id=$1 AND kind='intent' AND state='blocked'`, s.expID)
if err != nil {
return 0, err
@@ -291,14 +310,11 @@ func (s *ExplorationStore) SetIntentState(id int64, state string) error {
// terminal states stamp completed_at; reopening (back to open/running) clears it.
terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped"
_, err := s.db.Exec(`UPDATE exploration_nodes
SET state=$1, completed_at = CASE WHEN $4 THEN now() ELSE NULL END
SET state=$1, blocked_reason=NULL, completed_at = CASE WHEN $4 THEN now() ELSE NULL END
WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, terminal)
return err
}
<<<<<<< Updated upstream
const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), created_at`
=======
// CompareAndSetIntentState transitions one local intent only when its current
// state still matches expected. It is the state boundary used by pause/resume and
// worker settlement, so a stale API read cannot overwrite a concurrent finish.
@@ -593,12 +609,11 @@ func (u *TokenUsage) add(other TokenUsage) {
}
const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), COALESCE(blocked_reason,''), created_at`
>>>>>>> Stashed changes
func scanNode(sc interface{ Scan(...any) error }) (*Node, error) {
var n Node
var payload []byte
if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.CreatedAt); err != nil {
if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.BlockedReason, &n.CreatedAt); err != nil {
return nil, err
}
n.Payload = json.RawMessage(payload)
@@ -790,6 +805,31 @@ WHERE e.exploration_id=$1 AND e.rel=$2 AND n.kind='finding'`, s.expID, RelYields
return out, rows.Err()
}
// FindingIntentsTerminal is the inherited-history variant: it returns finding
// lineage only when the producing intent has reached an immutable terminal
// state. This keeps a source task's live worker topology private.
func (s *ExplorationStore) FindingIntentsTerminal() (map[int64]int64, error) {
rows, err := s.db.Query(`SELECT e.dst_id, e.src_id
FROM exploration_edges e
JOIN exploration_nodes finding ON finding.id=e.dst_id AND finding.exploration_id=e.exploration_id
JOIN exploration_nodes intent ON intent.id=e.src_id AND intent.exploration_id=e.exploration_id
WHERE e.exploration_id=$1 AND e.rel=$2 AND finding.kind='finding'
AND intent.kind='intent' AND intent.state IN ('done','blocked','exhausted','stopped')`, s.expID, RelYields)
if err != nil {
return nil, err
}
defer rows.Close()
out := map[int64]int64{}
for rows.Next() {
var findingID, intentID int64
if err := rows.Scan(&findingID, &intentID); err != nil {
return nil, err
}
out[findingID] = intentID
}
return out, rows.Err()
}
// Frontier returns open intents ordered by priority desc then id (FIFO tiebreak).
func (s *ExplorationStore) Frontier(limit int) ([]*Node, error) {
if limit <= 0 {
@@ -851,12 +891,16 @@ func (s *ExplorationStore) Stats() (map[string]int, error) {
// AppendActivity records one worker step and returns its id.
func (s *ExplorationStore) AppendActivity(a Activity) (int64, error) {
metadata := a.Metadata
if len(metadata) == 0 {
metadata = json.RawMessage(`{}`)
}
var id int64
err := s.db.QueryRow(`
INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13)
INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, metadata, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13,$14)
RETURNING id`, s.expID, a.NodeID, utf8Clean(a.Worker), utf8Clean(a.Kind), utf8Clean(a.Tool), utf8Clean(a.ToolUseID), a.IsError, utf8Clean(a.Summary), utf8Clean(a.Detail),
a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id)
metadata, a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id)
return id, err
}
@@ -986,6 +1030,41 @@ func (s *ExplorationStore) TokenStatsByWorker() ([]TokenUsage, error) {
return out, rows.Err()
}
// TokenStatsBySession aggregates every persisted completed run, independent of
// activity-history pagination. Main/planner use fixed keys; Worker executions use
// their intent node id, even when a different work#N process handled a retry.
func (s *ExplorationStore) TokenStatsBySession() ([]SessionTokenUsage, error) {
rows, err := s.db.Query(`SELECT
CASE
WHEN COALESCE(a.worker,'')='mainagent' THEN 'main'
WHEN COALESCE(a.worker,'')='planner' THEN 'plan'
WHEN n.id IS NOT NULL THEN 'intent:' || n.id::text
ELSE ''
END AS session_key,
COALESCE(SUM(a.input_tokens),0), COALESCE(SUM(a.output_tokens),0),
COALESCE(SUM(a.cache_read_tokens),0), COALESCE(SUM(a.cache_write_tokens),0)
FROM activity a
LEFT JOIN exploration_nodes n
ON n.id=a.node_id AND n.exploration_id=a.exploration_id AND n.kind='intent'
WHERE a.exploration_id=$1 AND a.kind='result'
AND (COALESCE(a.worker,'') IN ('mainagent','planner') OR n.id IS NOT NULL)
GROUP BY session_key
ORDER BY session_key`, s.expID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []SessionTokenUsage{}
for rows.Next() {
var u SessionTokenUsage
if err := rows.Scan(&u.Session, &u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens); err != nil {
return nil, err
}
out = append(out, u)
}
return out, rows.Err()
}
// ActivityList returns steps after sinceID (exclusive), optionally filtered by node.
// Returns items and the new cursor (max id seen).
func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) ([]Activity, int64, error) {
@@ -994,7 +1073,7 @@ func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int)
}
var rows *sql.Rows
var err error
const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens`
const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), metadata, created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens`
if nodeID != nil {
rows, err = s.db.Query(`SELECT `+cols+`
FROM activity WHERE exploration_id=$1 AND node_id=$2 AND id>$3 ORDER BY id LIMIT $4`, s.expID, *nodeID, sinceID, limit)
@@ -1008,6 +1087,43 @@ FROM activity WHERE exploration_id=$1 AND id>$2 ORDER BY id LIMIT $3`, s.expID,
defer rows.Close()
out := []Activity{}
cursor := sinceID
for rows.Next() {
var a Activity
if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.Metadata, &a.CreatedAt,
&a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil {
return nil, sinceID, err
}
if a.ID > cursor {
cursor = a.ID
}
out = append(out, a)
}
return out, cursor, rows.Err()
}
// ActivityListForTerminalIntent is the inherited, incremental-read boundary.
// The intent state check and activity read happen in one statement so a source
// intent reopened between API-level checks cannot expose its new in-flight trace.
// Model reasoning and accounting rows are never inherited.
func (s *ExplorationStore) ActivityListForTerminalIntent(nodeID, sinceID int64, limit int) ([]Activity, int64, error) {
if limit <= 0 {
limit = 300
}
rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''),
COALESCE(a.tool,''), COALESCE(a.tool_use_id,''), a.is_error, COALESCE(a.summary,''),
a.created_at, a.input_tokens, a.output_tokens, a.cache_read_tokens, a.cache_write_tokens
FROM activity a
JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id
WHERE a.exploration_id=$1 AND a.node_id=$2 AND a.id>$3
AND n.kind='intent' AND n.state IN ('done','blocked','exhausted','stopped')
AND a.kind NOT IN ('thinking','usage')
ORDER BY a.id LIMIT $4`, s.expID, nodeID, sinceID, limit)
if err != nil {
return nil, sinceID, err
}
defer rows.Close()
out := []Activity{}
cursor := sinceID
for rows.Next() {
var a Activity
if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt,
@@ -1054,7 +1170,7 @@ func (s *ExplorationStore) ActivityPage(f ActivitySessionFilter, before int64, l
if limit <= 0 {
limit = 200
}
const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens`
const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), metadata, created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens`
args := []any{s.expID}
cond, cargs := f.cond(len(args) + 1)
args = append(args, cargs...)
@@ -1071,6 +1187,51 @@ ORDER BY id DESC LIMIT $%d`, cond, beforeArg, beforeArg, limitArg)
}
defer rows.Close()
desc := []Activity{}
for rows.Next() {
var a Activity
if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.Metadata, &a.CreatedAt,
&a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil {
return nil, false, err
}
desc = append(desc, a)
}
if err := rows.Err(); err != nil {
return nil, false, err
}
hasMore := len(desc) > limit
if hasMore {
desc = desc[:limit]
}
// reverse the newest-first window into ascending id order for display.
out := make([]Activity, len(desc))
for i, a := range desc {
out[len(desc)-1-i] = a
}
return out, hasMore, nil
}
// ActivityPageForTerminalIntent is the inherited session-history boundary. It
// combines ownership, terminal-state and row-kind checks in one query, closing
// the race where a previously completed source intent is reopened before paging.
func (s *ExplorationStore) ActivityPageForTerminalIntent(nodeID, before int64, limit int) ([]Activity, bool, error) {
if limit <= 0 {
limit = 200
}
rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''),
COALESCE(a.tool,''), COALESCE(a.tool_use_id,''), a.is_error, COALESCE(a.summary,''),
a.created_at, a.input_tokens, a.output_tokens, a.cache_read_tokens, a.cache_write_tokens
FROM activity a
JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id
WHERE a.exploration_id=$1 AND a.node_id=$2
AND n.kind='intent' AND n.state IN ('done','blocked','exhausted','stopped')
AND a.kind NOT IN ('thinking','usage')
AND ($3 <= 0 OR a.id < $3)
ORDER BY a.id DESC LIMIT $4`, s.expID, nodeID, before, limit+1)
if err != nil {
return nil, false, err
}
defer rows.Close()
desc := []Activity{}
for rows.Next() {
var a Activity
if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt,
@@ -1086,7 +1247,6 @@ ORDER BY id DESC LIMIT $%d`, cond, beforeArg, beforeArg, limitArg)
if hasMore {
desc = desc[:limit]
}
// reverse the newest-first window into ascending id order for display.
out := make([]Activity, len(desc))
for i, a := range desc {
out[len(desc)-1-i] = a
@@ -1152,6 +1312,26 @@ ORDER BY id LIMIT $3`, s.expID, nodeID, limit)
return scanTrace(rows)
}
// ActivityTraceForTerminalIntent returns one inherited work trace only while its
// source intent is still terminal at query time.
func (s *ExplorationStore) ActivityTraceForTerminalIntent(nodeID int64, limit int) ([]Activity, error) {
if limit <= 0 {
limit = 500
}
rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'')
FROM activity a
JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id
WHERE a.exploration_id=$1 AND a.node_id=$2 AND n.kind='intent'
AND n.state IN ('done','blocked','exhausted','stopped')
AND a.kind NOT IN ('thinking','usage')
ORDER BY a.id LIMIT $3`, s.expID, nodeID, limit)
if err != nil {
return nil, err
}
defer rows.Close()
return scanTrace(rows)
}
// ActivityTraceSearch finds steps whose summary OR detail matches q
// (case-insensitive), returning summaries only. nodeID != nil scopes to one work;
// nil searches across ALL works (node_id IS NOT NULL — planner/main steps have no
@@ -1179,13 +1359,59 @@ AND (summary ILIKE $2 OR detail ILIKE $2) ORDER BY id LIMIT $3`, s.expID, like,
return scanTrace(rows)
}
// ActivityTraceSearchForTerminalIntent is the scoped inherited search variant;
// terminal eligibility is evaluated by the same statement that reads the rows.
func (s *ExplorationStore) ActivityTraceSearchForTerminalIntent(nodeID int64, q string, limit int) ([]Activity, error) {
if limit <= 0 {
limit = 100
}
like := "%" + q + "%"
rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'')
FROM activity a
JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id
WHERE a.exploration_id=$1 AND a.node_id=$2 AND n.kind='intent'
AND n.state IN ('done','blocked','exhausted','stopped')
AND a.kind NOT IN ('thinking','usage')
AND (a.summary ILIKE $3 OR a.detail ILIKE $3)
ORDER BY a.id LIMIT $4`, s.expID, nodeID, like, limit)
if err != nil {
return nil, err
}
defer rows.Close()
return scanTrace(rows)
}
// ActivityTraceSearchTerminalIntents searches inherited worker history without
// exposing planner/main rows or work that is still open/running in the source.
func (s *ExplorationStore) ActivityTraceSearchTerminalIntents(q string, limit int) ([]Activity, error) {
if limit <= 0 {
limit = 100
}
like := "%" + q + "%"
rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'')
FROM activity a
JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id
WHERE a.exploration_id=$1 AND n.kind='intent'
AND n.state IN ('done','blocked','exhausted','stopped')
AND a.kind NOT IN ('thinking','usage')
AND (a.summary ILIKE $2 OR a.detail ILIKE $2)
ORDER BY a.id LIMIT $3`, s.expID, like, limit)
if err != nil {
return nil, err
}
defer rows.Close()
return scanTrace(rows)
}
// AssetRef is a compact exploration node (intent/fact/finding) that references an
// asset via an anchor — for the coverage graph's node drawer.
type AssetRef struct {
ID int64 `json:"id"`
Kind string `json:"kind"`
State string `json:"state"`
Summary string `json:"summary"`
ID int64 `json:"id"`
Kind string `json:"kind"`
State string `json:"state"`
Summary string `json:"summary"`
SourceTaskID int64 `json:"source_task_id,omitempty"`
Inherited bool `json:"inherited,omitempty"`
}
// AssetRefs returns the intents / facts / findings in THIS exploration anchored to
@@ -1278,6 +1504,42 @@ FROM activity WHERE exploration_id=$1 AND kind<>'thinking' AND id IN (`+strings.
return out, rows.Err()
}
// ActivityByIDsForTerminalIntents is the inherited-detail read boundary. It
// deliberately excludes node-less planner/main rows, non-intent activities, live
// source work, and thinking. Local callers that need legacy behavior use
// ActivityByIDs instead.
func (s *ExplorationStore) ActivityByIDsForTerminalIntents(ids []int64) ([]Activity, error) {
if len(ids) == 0 {
return nil, nil
}
ph := make([]string, len(ids))
args := make([]any, len(ids)+1)
args[0] = s.expID
for i, id := range ids {
ph[i] = fmt.Sprintf("$%d", i+2)
args[i+1] = id
}
rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.detail,'')
FROM activity a
JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id
WHERE a.exploration_id=$1 AND a.kind NOT IN ('thinking','usage') AND n.kind='intent'
AND n.state IN ('done','blocked','exhausted','stopped')
AND a.id IN (`+strings.Join(ph, ",")+`) ORDER BY a.id`, args...)
if err != nil {
return nil, err
}
defer rows.Close()
out := []Activity{}
for rows.Next() {
var a Activity
if err := rows.Scan(&a.ID, &a.NodeID, &a.Kind, &a.Tool, &a.IsError, &a.Detail); err != nil {
return nil, err
}
out = append(out, a)
}
return out, rows.Err()
}
// AddStandaloneFinding writes a finding to the standalone findings table, which
// persists across task deletion (task_id / node_id become NULL when the task or
// exploration node is deleted). taskID and nodeID may be 0 (stored as NULL).
+350
View File
@@ -0,0 +1,350 @@
package db
import (
"database/sql"
"sort"
)
// DirectSourceStore binds one directly related task to its exploration store.
// It is intentionally a read-side helper: callers keep using the receiver store
// for every graph mutation, frontier lookup, and intent claim.
type DirectSourceStore struct {
Task TaskSource
Store *ExplorationStore
}
// DirectSourceStores resolves the live, direct task relations for this
// exploration. It deliberately queries on every call: relations disappear when
// a source task is deleted, and inherited context must not retain stale rows.
// Sources of a source are never expanded.
func (s *ExplorationStore) DirectSourceStores() ([]DirectSourceStore, error) {
rows, err := s.db.Query(`
SELECT source.id, source.exploration_id, source.description, source.goal, source.status
FROM tasks owner
JOIN task_relations relation ON relation.task_id=owner.id
JOIN tasks source ON source.id=relation.source_task_id AND source.deleted_at IS NULL
WHERE owner.exploration_id=$1 AND owner.deleted_at IS NULL
ORDER BY relation.created_at, source.id`, s.expID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []DirectSourceStore{}
for rows.Next() {
var source TaskSource
if err := rows.Scan(&source.TaskID, &source.ExplorationID, &source.Description, &source.Goal, &source.Status); err != nil {
return nil, err
}
out = append(out, DirectSourceStore{Task: source, Store: s.db.Exploration(source.ExplorationID)})
}
return out, rows.Err()
}
// TaskID returns the live task bound to this exploration. Explorations created
// directly in tests or maintenance code have no task and return zero.
func (s *ExplorationStore) TaskID() (int64, error) {
var id int64
err := s.db.QueryRow(`SELECT id FROM tasks WHERE exploration_id=$1 AND deleted_at IS NULL`, s.expID).Scan(&id)
if err != nil {
if err == sql.ErrNoRows {
return 0, nil
}
return 0, err
}
return id, nil
}
func markInheritedNode(n *Node, taskID int64) *Node {
if n == nil {
return nil
}
n.SourceTaskID = taskID
n.Inherited = true
return n
}
func markInheritedActivities(in []Activity, taskID int64) []Activity {
for i := range in {
in[i].SourceTaskID = taskID
in[i].Inherited = true
}
return in
}
func inheritedIntentTerminal(state string) bool {
switch state {
case "done", "blocked", "exhausted", "stopped":
return true
default:
return false
}
}
// ListByKindWithSources returns local nodes followed by nodes from each direct
// source in relation order. The limit remains per exploration, matching the
// existing ListByKind contract while ensuring one large task cannot hide all
// inherited context from another source.
func (s *ExplorationStore) ListByKindWithSources(kind string, limit int) ([]*Node, error) {
out, err := s.ListByKind(kind, limit)
if err != nil {
return nil, err
}
sources, err := s.DirectSourceStores()
if err != nil {
return nil, err
}
for _, source := range sources {
nodes, err := source.Store.ListByKind(kind, limit)
if err != nil {
return nil, err
}
for _, node := range nodes {
if kind == KindIntent && !inheritedIntentTerminal(node.State) {
continue
}
out = append(out, markInheritedNode(node, source.Task.TaskID))
}
}
return out, nil
}
// GetNodeWithSources reads a node only when it belongs to this exploration or
// one of its direct sources. Inherited nodes are tagged so tool callers can keep
// them read-only and show their provenance.
func (s *ExplorationStore) GetNodeWithSources(id int64) (*Node, error) {
node, err := s.GetNode(id)
if err != nil || node != nil {
return node, err
}
sources, err := s.DirectSourceStores()
if err != nil {
return nil, err
}
for _, source := range sources {
node, err = source.Store.GetNode(id)
if err != nil {
return nil, err
}
if node != nil {
if node.Kind == KindIntent && !inheritedIntentTerminal(node.State) {
continue
}
return markInheritedNode(node, source.Task.TaskID), nil
}
}
return nil, nil
}
// FindingIntentsWithSources combines finding lineage for the current
// exploration and each direct source. Node ids are globally unique.
func (s *ExplorationStore) FindingIntentsWithSources() (map[int64]int64, error) {
out, err := s.FindingIntents()
if err != nil {
return nil, err
}
sources, err := s.DirectSourceStores()
if err != nil {
return nil, err
}
for _, source := range sources {
items, err := source.Store.FindingIntentsTerminal()
if err != nil {
return nil, err
}
for findingID, intentID := range items {
out[findingID] = intentID
}
}
return out, nil
}
// ActivityTraceWithSources returns a work trace when its intent belongs to the
// current exploration or a direct source. It never searches indirect sources.
func (s *ExplorationStore) ActivityTraceWithSources(nodeID int64, limit int) ([]Activity, error) {
acts, err := s.ActivityTrace(nodeID, limit)
if err != nil || len(acts) > 0 {
return acts, err
}
sources, err := s.DirectSourceStores()
if err != nil {
return nil, err
}
for _, source := range sources {
node, nodeErr := source.Store.GetNode(nodeID)
if nodeErr != nil {
return nil, nodeErr
}
if node == nil || node.Kind != KindIntent {
continue
}
acts, err = source.Store.ActivityTraceForTerminalIntent(nodeID, limit)
if err != nil {
return nil, err
}
if len(acts) > 0 {
return markInheritedActivities(acts, source.Task.TaskID), nil
}
}
return []Activity{}, nil
}
// ActivityListWithSources is the source-aware equivalent used by
// get_worker_output. Node ids are global, so the first owning exploration is
// unambiguous even when the work has not emitted any activity yet.
func (s *ExplorationStore) ActivityListWithSources(nodeID, sinceID int64, limit int) ([]Activity, int64, error) {
node, err := s.GetNode(nodeID)
if err != nil {
return nil, sinceID, err
}
if node != nil {
return s.ActivityList(&nodeID, sinceID, limit)
}
sources, err := s.DirectSourceStores()
if err != nil {
return nil, sinceID, err
}
for _, source := range sources {
node, err = source.Store.GetNode(nodeID)
if err != nil {
return nil, sinceID, err
}
if node == nil || node.Kind != KindIntent {
continue
}
acts, cursor, err := source.Store.ActivityListForTerminalIntent(nodeID, sinceID, limit)
if err != nil {
return nil, sinceID, err
}
return markInheritedActivities(acts, source.Task.TaskID), cursor, nil
}
return []Activity{}, sinceID, nil
}
// ActivityDetailWithSources keeps the legacy local-task lookup (including local
// thinking rows), while inherited details are restricted to terminal worker
// intents. Source planner/main rows have no node id and must never become part of
// inherited context.
func (s *ExplorationStore) ActivityDetailWithSources(id int64) (string, error) {
detail, err := s.ActivityDetail(id)
if err != nil || detail != "" {
return detail, err
}
acts, err := s.ActivityByIDsWithSources([]int64{id})
if err != nil || len(acts) == 0 {
return "", err
}
return acts[0].Detail, nil
}
// ActivityTraceSearchWithSources performs the scoped worker-trace search used
// by get_worker_trace. A node id resolves to at most one exploration because
// exploration node ids are global.
func (s *ExplorationStore) ActivityTraceSearchWithSources(nodeID int64, q string, limit int) ([]Activity, error) {
acts, err := s.ActivityTraceSearch(&nodeID, q, limit)
if err != nil || len(acts) > 0 {
return acts, err
}
sources, err := s.DirectSourceStores()
if err != nil {
return nil, err
}
for _, source := range sources {
node, nodeErr := source.Store.GetNode(nodeID)
if nodeErr != nil {
return nil, nodeErr
}
if node == nil || node.Kind != KindIntent {
continue
}
acts, err = source.Store.ActivityTraceSearchForTerminalIntent(nodeID, q, limit)
if err != nil {
return nil, err
}
if len(acts) > 0 {
return markInheritedActivities(acts, source.Task.TaskID), nil
}
}
return []Activity{}, nil
}
// ActivityTraceSearchAllWithSources searches local worker traces plus every
// direct source. The local owner is excluded only from the current exploration;
// inherited traces are immutable historical context.
func (s *ExplorationStore) ActivityTraceSearchAllWithSources(excludeNodeID int64, q string, limit int) ([]Activity, error) {
if limit <= 0 {
limit = 100
}
out, err := s.ActivityTraceSearchExcluding(excludeNodeID, q, limit)
if err != nil {
return nil, err
}
sources, err := s.DirectSourceStores()
if err != nil {
return nil, err
}
for _, source := range sources {
acts, err := source.Store.ActivityTraceSearchTerminalIntents(q, limit)
if err != nil {
return nil, err
}
out = append(out, markInheritedActivities(acts, source.Task.TaskID)...)
}
sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
if len(out) > limit {
out = out[:limit]
}
return out, nil
}
// ActivityByIDsWithSources loads local step details with the legacy behavior. For
// direct sources it only returns rows attached to terminal intents, preventing an
// arbitrary global activity id from exposing source planner/main transcripts.
func (s *ExplorationStore) ActivityByIDsWithSources(ids []int64) ([]Activity, error) {
out, err := s.ActivityByIDs(ids)
if err != nil {
return nil, err
}
sources, err := s.DirectSourceStores()
if err != nil {
return nil, err
}
for _, source := range sources {
acts, err := source.Store.ActivityByIDsForTerminalIntents(ids)
if err != nil {
return nil, err
}
out = append(out, markInheritedActivities(acts, source.Task.TaskID)...)
}
sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
return out, nil
}
// AssetRefsWithSources returns anchored nodes from this exploration and each
// direct source. Inherited entries retain their owning task id so API/UI callers
// can present them as immutable context.
func (s *ExplorationStore) AssetRefsWithSources(assetID int64) ([]AssetRef, error) {
out, err := s.AssetRefs(assetID)
if err != nil {
return nil, err
}
sources, err := s.DirectSourceStores()
if err != nil {
return nil, err
}
for _, source := range sources {
refs, err := source.Store.AssetRefs(assetID)
if err != nil {
return nil, err
}
for i := range refs {
if refs[i].Kind == KindIntent && !inheritedIntentTerminal(refs[i].State) {
continue
}
refs[i].SourceTaskID = source.Task.TaskID
refs[i].Inherited = true
out = append(out, refs[i])
}
}
sort.Slice(out, func(i, j int) bool { return out[i].ID > out[j].ID })
return out, nil
}
+431
View File
@@ -0,0 +1,431 @@
package db
import (
"fmt"
"testing"
)
func TestInheritedActivityReadsRequireTerminalIntent(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("source activity boundary", "source activity boundary")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) })
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "source work"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
appendStep := func(kind, summary string) int64 {
t.Helper()
id, appendErr := store.AppendActivity(Activity{
NodeID: &intentID, Worker: "worker", Kind: kind, Summary: summary, Detail: summary + " detail",
})
if appendErr != nil {
t.Fatal(appendErr)
}
return id
}
textID := appendStep("text", "shared text")
thinkingID := appendStep("thinking", "private reasoning")
usageID := appendStep("usage", "private accounting")
resultID := appendStep("result", "shared result")
if err := store.SetIntentState(intentID, "done"); err != nil {
t.Fatal(err)
}
page, more, err := store.ActivityPageForTerminalIntent(intentID, 0, 10)
if err != nil || more || len(page) != 2 || page[0].ID != textID || page[1].ID != resultID {
t.Fatalf("terminal page boundary: page=%+v more=%v err=%v", page, more, err)
}
list, cursor, err := store.ActivityListForTerminalIntent(intentID, 0, 10)
if err != nil || len(list) != 2 || cursor != resultID {
t.Fatalf("terminal list boundary: list=%+v cursor=%d err=%v", list, cursor, err)
}
trace, err := store.ActivityTraceForTerminalIntent(intentID, 10)
if err != nil || len(trace) != 2 {
t.Fatalf("terminal trace boundary: trace=%+v err=%v", trace, err)
}
hits, err := store.ActivityTraceSearchForTerminalIntent(intentID, "shared", 10)
if err != nil || len(hits) != 2 {
t.Fatalf("terminal search boundary: hits=%+v err=%v", hits, err)
}
details, err := store.ActivityByIDsForTerminalIntents([]int64{textID, thinkingID, usageID, resultID})
if err != nil || len(details) != 2 || details[0].ID != textID || details[1].ID != resultID {
t.Fatalf("terminal detail boundary: details=%+v err=%v", details, err)
}
// Reopening a source intent must close every inherited activity read even if
// callers still hold a stale terminal-state snapshot.
if err := store.SetIntentState(intentID, "running"); err != nil {
t.Fatal(err)
}
page, _, err = store.ActivityPageForTerminalIntent(intentID, 0, 10)
if err != nil || len(page) != 0 {
t.Fatalf("reopened intent leaked through page: page=%+v err=%v", page, err)
}
list, _, err = store.ActivityListForTerminalIntent(intentID, 0, 10)
if err != nil || len(list) != 0 {
t.Fatalf("reopened intent leaked through list: list=%+v err=%v", list, err)
}
trace, err = store.ActivityTraceForTerminalIntent(intentID, 10)
if err != nil || len(trace) != 0 {
t.Fatalf("reopened intent leaked through trace: trace=%+v err=%v", trace, err)
}
hits, err = store.ActivityTraceSearchForTerminalIntent(intentID, "shared", 10)
if err != nil || len(hits) != 0 {
t.Fatalf("reopened intent leaked through search: hits=%+v err=%v", hits, err)
}
details, err = store.ActivityByIDsForTerminalIntents([]int64{textID, resultID})
if err != nil || len(details) != 0 {
t.Fatalf("reopened intent leaked through detail: details=%+v err=%v", details, err)
}
}
func TestExplorationDirectSourceReadView(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
grand, err := d.CreateTask("grand source", "grand goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
source, err := d.CreateTaskWithOptions("direct source", "source goal", TaskCreateOptions{
SourceTaskIDs: []int64{grand.ID},
})
if err != nil {
t.Fatal(err)
}
current, err := d.CreateTaskWithOptions("current", "current goal", TaskCreateOptions{
SourceTaskIDs: []int64{source.ID},
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteTask(current.ID)
_ = d.DeleteTask(source.ID)
_ = d.DeleteTask(grand.ID)
})
grandStore := d.Exploration(grand.ExplorationID)
sourceStore := d.Exploration(source.ExplorationID)
currentStore := d.Exploration(current.ExplorationID)
grandFact, err := grandStore.AddNode(KindFact, map[string]any{"summary": "grand-only fact"}, 0, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
sourceFact, err := sourceStore.AddNode(KindFact, map[string]any{"summary": "direct source fact"}, 0, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
sourceFinding, err := sourceStore.AddNode(KindFinding, map[string]any{"summary": "direct source finding"}, 0, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
sourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "source work"}, 9, nil, "planner")
if err != nil {
t.Fatal(err)
}
if err := sourceStore.Link(sourceIntent, RelYields, sourceFinding); err != nil {
t.Fatal(err)
}
currentIntent, err := currentStore.AddIntent(map[string]any{"summary": "current work"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
currentFact, err := currentStore.AddNode(KindFact, map[string]any{"summary": "current fact"}, 0, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
sources, err := currentStore.DirectSourceStores()
if err != nil {
t.Fatal(err)
}
if len(sources) != 1 || sources[0].Task.TaskID != source.ID {
t.Fatalf("want only direct source %d, got %+v", source.ID, sources)
}
frontier, err := currentStore.Frontier(20)
if err != nil {
t.Fatal(err)
}
if len(frontier) != 1 || frontier[0].ID != currentIntent {
t.Fatalf("source intent leaked into local frontier: %+v", frontier)
}
if claimed, err := currentStore.ClaimIntent(sourceIntent, "wrong-task-worker"); err != nil || claimed {
t.Fatalf("source intent must not be claimable: claimed=%v err=%v", claimed, err)
}
if got, err := currentStore.GetNodeWithSources(sourceIntent); err != nil || got != nil {
t.Fatalf("live source intent must not be readable as inherited history: node=%+v err=%v", got, err)
}
if got, err := currentStore.ListByKindWithSources(KindIntent, 100); err != nil || len(got) != 1 || got[0].ID != currentIntent {
t.Fatalf("live source intent leaked through inherited list: nodes=%+v err=%v", got, err)
}
if got, err := currentStore.FindingIntentsWithSources(); err != nil || got[sourceFinding] != 0 {
t.Fatalf("live source intent leaked through finding lineage: lineage=%+v err=%v", got, err)
}
facts, err := currentStore.ListByKindWithSources(KindFact, 100)
if err != nil {
t.Fatal(err)
}
seenCurrent, seenSource, seenGrand := false, false, false
for _, fact := range facts {
switch fact.ID {
case currentFact:
seenCurrent = !fact.Inherited && fact.SourceTaskID == 0
case sourceFact:
seenSource = fact.Inherited && fact.SourceTaskID == source.ID
case grandFact:
seenGrand = true
}
}
if !seenCurrent || !seenSource || seenGrand {
t.Fatalf("unexpected direct-source fact view: current=%v source=%v grand=%v", seenCurrent, seenSource, seenGrand)
}
gotSource, err := currentStore.GetNodeWithSources(sourceFact)
if err != nil || gotSource == nil || !gotSource.Inherited || gotSource.SourceTaskID != source.ID {
t.Fatalf("source node lookup: node=%+v err=%v", gotSource, err)
}
gotGrand, err := currentStore.GetNodeWithSources(grandFact)
if err != nil || gotGrand != nil {
t.Fatalf("indirect source must be invisible: node=%+v err=%v", gotGrand, err)
}
// Existing write methods remain bound to currentStore.expID. Passing an
// inherited id is a no-op and cannot mutate the source blackboard.
if err := currentStore.SetNodeState(sourceFact, "dismissed"); err != nil {
t.Fatal(err)
}
unchanged, err := sourceStore.GetNode(sourceFact)
if err != nil || unchanged == nil || unchanged.State != "confirmed" {
t.Fatalf("inherited fact was mutated: node=%+v err=%v", unchanged, err)
}
if err := sourceStore.SetIntentState(sourceIntent, "done"); err != nil {
t.Fatal(err)
}
if got, err := currentStore.GetNodeWithSources(sourceIntent); err != nil || got == nil || !got.Inherited {
t.Fatalf("terminal source intent should become readable history: node=%+v err=%v", got, err)
}
if got, err := currentStore.FindingIntentsWithSources(); err != nil || got[sourceFinding] != sourceIntent {
t.Fatalf("terminal source finding lineage missing: lineage=%+v err=%v", got, err)
}
sourcePlannerStep, err := sourceStore.AppendActivity(Activity{
Worker: "planner", Kind: "text", Summary: "private source planner", Detail: "private source planner transcript",
})
if err != nil {
t.Fatal(err)
}
openSourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "source work still running"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if err := sourceStore.SetIntentState(openSourceIntent, "running"); err != nil {
t.Fatal(err)
}
openSourceStep, err := sourceStore.AppendActivity(Activity{
NodeID: &openSourceIntent, Worker: "source-live-worker", Kind: "text",
Summary: "private live source work", Detail: "private live source transcript",
})
if err != nil {
t.Fatal(err)
}
sourceStep, err := sourceStore.AppendActivity(Activity{
NodeID: &sourceIntent, Worker: "source-worker", Kind: "tool_result", Tool: "HTTP",
Summary: "source-secret response", Detail: "source-secret full detail",
})
if err != nil {
t.Fatal(err)
}
currentPlannerStep, err := currentStore.AppendActivity(Activity{
Worker: "planner", Kind: "text", Summary: "current planner", Detail: "current planner detail",
})
if err != nil {
t.Fatal(err)
}
currentThinkingStep, err := currentStore.AppendActivity(Activity{
Worker: "planner", Kind: "thinking", Summary: "current thinking", Detail: "current thinking detail",
})
if err != nil {
t.Fatal(err)
}
grandStep, err := grandStore.AppendActivity(Activity{
NodeID: &grandFact, Worker: "grand-worker", Kind: "tool_result", Tool: "HTTP",
Summary: "grand-only response", Detail: "grand-only detail",
})
if err != nil {
t.Fatal(err)
}
trace, err := currentStore.ActivityTraceWithSources(sourceIntent, 20)
if err != nil || len(trace) != 1 || trace[0].ID != sourceStep || !trace[0].Inherited || trace[0].SourceTaskID != source.ID {
t.Fatalf("source trace lookup: trace=%+v err=%v", trace, err)
}
hits, err := currentStore.ActivityTraceSearchAllWithSources(0, "source-secret", 20)
if err != nil || len(hits) != 1 || hits[0].ID != sourceStep || !hits[0].Inherited {
t.Fatalf("source trace search: hits=%+v err=%v", hits, err)
}
grandHits, err := currentStore.ActivityTraceSearchAllWithSources(0, "grand-only", 20)
if err != nil || len(grandHits) != 0 {
t.Fatalf("indirect source trace leaked: hits=%+v err=%v", grandHits, err)
}
liveHits, err := currentStore.ActivityTraceSearchAllWithSources(0, "private live source", 20)
if err != nil || len(liveHits) != 0 {
t.Fatalf("live source trace leaked: hits=%+v err=%v", liveHits, err)
}
if liveTrace, err := currentStore.ActivityTraceWithSources(openSourceIntent, 20); err != nil || len(liveTrace) != 0 {
t.Fatalf("live inherited intent trace leaked: trace=%+v err=%v", liveTrace, err)
}
details, err := currentStore.ActivityByIDsWithSources([]int64{
sourcePlannerStep, openSourceStep, sourceStep, currentPlannerStep, grandStep,
})
if err != nil || len(details) != 2 || details[0].ID != sourceStep || !details[0].Inherited || details[1].ID != currentPlannerStep || details[1].Inherited {
t.Fatalf("source trace detail allow-list: details=%+v err=%v", details, err)
}
if detail, err := currentStore.ActivityDetailWithSources(sourcePlannerStep); err != nil || detail != "" {
t.Fatalf("source planner transcript leaked: detail=%q err=%v", detail, err)
}
if detail, err := currentStore.ActivityDetailWithSources(openSourceStep); err != nil || detail != "" {
t.Fatalf("live source intent transcript leaked: detail=%q err=%v", detail, err)
}
if detail, err := currentStore.ActivityDetailWithSources(currentPlannerStep); err != nil || detail != "current planner detail" {
t.Fatalf("local activity detail behavior changed: detail=%q err=%v", detail, err)
}
if detail, err := currentStore.ActivityDetailWithSources(currentThinkingStep); err != nil || detail != "current thinking detail" {
t.Fatalf("local thinking detail behavior changed: detail=%q err=%v", detail, err)
}
}
func TestTaskAssetContextUsesDirectSourceScopeAndAnchors(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
grand, err := d.CreateTask("grand assets", "grand goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
source, err := d.CreateTaskWithOptions("source assets", "source goal", TaskCreateOptions{SourceTaskIDs: []int64{grand.ID}})
if err != nil {
t.Fatal(err)
}
current, err := d.CreateTaskWithOptions("current assets", "current goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID}})
if err != nil {
t.Fatal(err)
}
assets := d.Assets()
sourceTestedDomain := fmt.Sprintf("context-tested-%d.invalid", source.ID)
sourceUntestedDomain := fmt.Sprintf("context-untested-%d.invalid", source.ID)
sourceAnchoredOnlyDomain := fmt.Sprintf("context-anchor-only-%d.invalid", source.ID)
grandDomain := fmt.Sprintf("context-grand-%d.invalid", grand.ID)
sourceTestedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceTestedDomain, TaskID: source.ID})
if err != nil {
t.Fatal(err)
}
sourceUntestedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceUntestedDomain, TaskID: source.ID})
if err != nil {
t.Fatal(err)
}
sourceAnchoredOnlyID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceAnchoredOnlyDomain, TaskID: source.ID})
if err != nil {
t.Fatal(err)
}
grandID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: grandDomain, TaskID: grand.ID})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteTask(current.ID)
_ = d.DeleteTask(source.ID)
_ = d.DeleteTask(grand.ID)
_, _ = assets.DeleteByIDs([]int64{sourceTestedID, sourceUntestedID, sourceAnchoredOnlyID, grandID})
})
if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceTestedDomain, "test", "agent"); err != nil {
t.Fatal(err)
}
if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceUntestedDomain, "test", "agent"); err != nil {
t.Fatal(err)
}
if _, err := assets.AddAgentScope(grand.ID, "root_domain", grandDomain, "test", "agent"); err != nil {
t.Fatal(err)
}
if _, err := d.Exploration(source.ExplorationID).AddNode(KindFact, map[string]any{"summary": "tested source asset"}, 0, "confirmed", "worker", []int64{sourceTestedID}); err != nil {
t.Fatal(err)
}
if _, err := d.Exploration(source.ExplorationID).AddNode(KindFact, map[string]any{"summary": "anchored-only source asset"}, 0, "confirmed", "worker", []int64{sourceAnchoredOnlyID}); err != nil {
t.Fatal(err)
}
if _, err := d.Exploration(grand.ExplorationID).AddNode(KindFact, map[string]any{"summary": "indirect tested asset"}, 0, "confirmed", "worker", []int64{grandID}); err != nil {
t.Fatal(err)
}
scopes, err := assets.ListTaskScopeWithSources(current.ID)
if err != nil {
t.Fatal(err)
}
if len(scopes) != 2 || scopes[0].TaskID != source.ID || scopes[1].TaskID != source.ID {
t.Fatalf("direct source scopes only: %+v", scopes)
}
cov, err := assets.TaskCoverageWithSources(current.ID)
if err != nil {
t.Fatal(err)
}
if cov.Denominator != 3 || cov.Tested != 2 {
t.Fatalf("combined coverage: %+v", cov)
}
untested, total, err := assets.ListUntestedAssetsWithSources(current.ID, "", 10, 0)
if err != nil {
t.Fatal(err)
}
if total != 1 || len(untested) != 1 || untested[0].ID != sourceUntestedID {
t.Fatalf("combined untested assets: total=%d assets=%+v", total, untested)
}
hosts, err := assets.HostsByTaskWithSources(current.ID)
if err != nil {
t.Fatal(err)
}
hostSet := map[string]bool{}
for _, host := range hosts {
hostSet[host] = true
}
if !hostSet[sourceTestedDomain] || !hostSet[sourceUntestedDomain] || !hostSet[sourceAnchoredOnlyDomain] || hostSet[grandDomain] {
t.Fatalf("direct source hosts only: %v", hosts)
}
refs, err := d.Exploration(current.ExplorationID).AssetRefsWithSources(sourceAnchoredOnlyID)
if err != nil || len(refs) != 1 || !refs[0].Inherited || refs[0].SourceTaskID != source.ID || refs[0].Kind != KindFact {
t.Fatalf("source asset refs provenance: refs=%+v err=%v", refs, err)
}
grandRefs, err := d.Exploration(current.ExplorationID).AssetRefsWithSources(grandID)
if err != nil || len(grandRefs) != 0 {
t.Fatalf("indirect source asset refs leaked: refs=%+v err=%v", grandRefs, err)
}
graph, err := assets.BuildCoverageGraph(current.ID, current.ExplorationID)
if err != nil {
t.Fatal(err)
}
seenAnchoredOnly := false
for _, node := range graph.Nodes {
if node.AssetID == sourceAnchoredOnlyID {
seenAnchoredOnly = node.Tested && node.InScope
}
}
if !seenAnchoredOnly {
t.Fatalf("anchored-only source asset missing from coverage graph: %+v", graph.Nodes)
}
}
+135 -1
View File
@@ -1,6 +1,9 @@
package db
import "testing"
import (
"fmt"
"testing"
)
func TestExplorationFlow(t *testing.T) {
d, err := Open(testDSN(t))
@@ -112,3 +115,134 @@ func TestExplorationFlow(t *testing.T) {
t.Fatalf("stats: %+v", st)
}
}
func TestIntentPauseResumeAndCancelCleanup(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("test", "worker control cleanup")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
es := d.Exploration(expID)
assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{
Domain: fmt.Sprintf("cancel-intent-%d.invalid", expID),
})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, assetID)
otherIntent, err := es.AddIntent(map[string]any{"summary": "keep intent"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
intentID, err := es.AddIntent(map[string]any{"summary": "cancel intent"}, 10, nil, "planner")
if err != nil {
t.Fatal(err)
}
if claimed, err := es.ClaimIntent(intentID, "worker-control-test"); err != nil || !claimed {
t.Fatalf("initial claim: claimed=%v err=%v", claimed, err)
}
if err := es.SetIntentState(intentID, "paused"); err != nil {
t.Fatal(err)
}
frontier, err := es.Frontier(10)
if err != nil {
t.Fatal(err)
}
for _, node := range frontier {
if node.ID == intentID {
t.Fatalf("paused intent %d must not enter frontier", intentID)
}
}
if claimed, err := es.ClaimIntent(intentID, "worker-while-paused"); err != nil || claimed {
t.Fatalf("paused intent claim: claimed=%v err=%v", claimed, err)
}
if err := es.SetIntentState(intentID, "open"); err != nil {
t.Fatal(err)
}
if claimed, err := es.ClaimIntent(intentID, "worker-after-resume"); err != nil || !claimed {
t.Fatalf("resumed intent claim: claimed=%v err=%v", claimed, err)
}
if err := es.SetIntentState(intentID, "paused"); err != nil {
t.Fatal(err)
}
directFact, err := es.AddNode("fact", map[string]any{"summary": "remove fact"}, 0, "confirmed", "worker", []int64{assetID})
if err != nil {
t.Fatal(err)
}
directFinding, err := es.AddNode("finding", map[string]any{"summary": "remove finding"}, 0, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
keptFact, err := es.AddNode("fact", map[string]any{"summary": "keep fact"}, 0, "confirmed", "other-worker", nil)
if err != nil {
t.Fatal(err)
}
if err := es.Link(intentID, RelYields, directFact); err != nil {
t.Fatal(err)
}
if err := es.Link(intentID, RelYields, directFinding); err != nil {
t.Fatal(err)
}
findingRowID, err := es.AddStandaloneFinding(0, directFinding, "test", "cancelled finding", SeverityHigh, "summary", "evidence", "worker", []int64{assetID})
if err != nil {
t.Fatal(err)
}
activityID, err := es.AppendActivity(Activity{NodeID: &intentID, Worker: "worker", Kind: "result", Summary: "remove activity"})
if err != nil {
t.Fatal(err)
}
keptActivityID, err := es.AppendActivity(Activity{NodeID: &otherIntent, Worker: "other-worker", Kind: "result", Summary: "keep activity"})
if err != nil {
t.Fatal(err)
}
cleanup, err := es.CancelIntent(intentID)
if err != nil {
t.Fatal(err)
}
if cleanup.Intents != 1 || cleanup.Facts != 1 || cleanup.Findings != 1 || cleanup.Activities != 1 {
t.Fatalf("unexpected cleanup counts: %+v", cleanup)
}
for _, nodeID := range []int64{intentID, directFact, directFinding} {
node, err := es.GetNode(nodeID)
if err != nil {
t.Fatal(err)
}
if node != nil {
t.Fatalf("node %d survived intent cancellation", nodeID)
}
}
for _, nodeID := range []int64{otherIntent, keptFact} {
node, err := es.GetNode(nodeID)
if err != nil || node == nil {
t.Fatalf("unrelated node %d removed: node=%v err=%v", nodeID, node, err)
}
}
assertCount := func(query string, want int, args ...any) {
t.Helper()
var got int
if err := d.QueryRow(query, args...).Scan(&got); err != nil {
t.Fatal(err)
}
if got != want {
t.Fatalf("query count=%d, want %d: %s", got, want, query)
}
}
assertCount(`SELECT COUNT(*) FROM findings WHERE id=$1`, 0, findingRowID)
assertCount(`SELECT COUNT(*) FROM activity WHERE id=$1`, 0, activityID)
assertCount(`SELECT COUNT(*) FROM activity WHERE id=$1`, 1, keptActivityID)
assertCount(`SELECT COUNT(*) FROM exploration_edges WHERE exploration_id=$1 AND (src_id=$2 OR dst_id=$2)`, 0, expID, intentID)
assertCount(`SELECT COUNT(*) FROM exploration_anchors WHERE node_id=$1`, 0, directFact)
assertCount(`SELECT COUNT(*) FROM assets WHERE id=$1`, 1, assetID)
}
+130
View File
@@ -0,0 +1,130 @@
package db
import (
"fmt"
"testing"
)
func TestTokenStatsBySessionUsesCompleteIntentHistory(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("token sessions", "aggregate complete session history")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "worker session"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
factID, err := store.AddNode(KindFact, map[string]any{"summary": "not a worker"}, 0, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
otherExpID, err := d.CreateExploration("other token sessions", "foreign intent must not match")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, otherExpID)
foreignIntentID, err := d.Exploration(otherExpID).AddIntent(map[string]any{"summary": "foreign worker"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
appendResult := func(nodeID *int64, worker string, input, output, read, write int) {
t.Helper()
if _, appendErr := store.AppendActivity(Activity{
NodeID: nodeID,
Worker: worker,
Kind: "result",
InputTokens: &input,
OutputTokens: &output,
CacheReadTokens: &read,
CacheWriteTokens: &write,
}); appendErr != nil {
t.Fatal(appendErr)
}
}
appendResult(nil, "mainagent", 11, 12, 13, 14)
appendResult(nil, "planner", 21, 22, 23, 24)
// More than ActivityPage's default limit proves this aggregate reads the
// persisted history directly instead of summing the currently loaded page.
const completedRuns = 205
for i := 0; i < completedRuns; i++ {
appendResult(&intentID, fmt.Sprintf("work#%d", i%3+1), 1, 2, 3, 4)
}
// All three rows remain part of the legacy worker/whole-task totals, but none
// represents a local Worker session and therefore none may create intent:*.
appendResult(&factID, "work#fact", 31, 32, 33, 34)
appendResult(nil, "work#missing-node", 41, 42, 43, 44)
appendResult(&foreignIntentID, "work#foreign", 51, 52, 53, 54)
ignoredInput := 1000
if _, err := store.AppendActivity(Activity{
NodeID: &intentID,
Worker: "work#1",
Kind: "usage",
InputTokens: &ignoredInput,
}); err != nil {
t.Fatal(err)
}
sessions, err := store.TokenStatsBySession()
if err != nil {
t.Fatal(err)
}
got := make(map[string]SessionTokenUsage, len(sessions))
for _, usage := range sessions {
got[usage.Session] = usage
}
if len(got) != 3 {
t.Fatalf("sessions = %+v, want only main, plan, and the local intent", sessions)
}
assertSessionTokens(t, got["main"], "main", 11, 12, 13, 14)
assertSessionTokens(t, got["plan"], "plan", 21, 22, 23, 24)
assertSessionTokens(t, got[fmt.Sprintf("intent:%d", intentID)], fmt.Sprintf("intent:%d", intentID),
completedRuns, completedRuns*2, completedRuns*3, completedRuns*4)
workers, err := store.TokenStatsByWorker()
if err != nil {
t.Fatal(err)
}
if len(workers) != 8 { // main, planner, 3 executors, and the 3 deliberately invalid session rows
t.Fatalf("legacy workers changed: got %d entries: %+v", len(workers), workers)
}
total, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, total,
11+21+completedRuns+31+41+51,
12+22+completedRuns*2+32+42+52,
13+23+completedRuns*3+33+43+53,
14+24+completedRuns*4+34+44+54)
}
func assertSessionTokens(t *testing.T, got SessionTokenUsage, session string, input, output, read, write int) {
t.Helper()
if got.Session != session || got.InputTokens != input || got.OutputTokens != output ||
got.CacheReadTokens != read || got.CacheWriteTokens != write {
t.Fatalf("session %q = %+v, want input=%d output=%d cache-read=%d cache-write=%d",
session, got, input, output, read, write)
}
}
func assertTokenUsage(t *testing.T, got TokenUsage, input, output, read, write int) {
t.Helper()
if got.InputTokens != input || got.OutputTokens != output ||
got.CacheReadTokens != read || got.CacheWriteTokens != write {
t.Fatalf("token total = %+v, want input=%d output=%d cache-read=%d cache-write=%d",
got, input, output, read, write)
}
}
+32
View File
@@ -0,0 +1,32 @@
package db
import "fmt"
// DiscardOpenIntent compensates a follow-up creation when task admission fails.
// The task execution gate must still be held by the caller, so the intent cannot
// be claimed between this check and deletion. Edges and anchors cascade with the
// node; activity is deleted explicitly because its node FK otherwise becomes NULL.
func (s *ExplorationStore) DiscardOpenIntent(id int64) error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if _, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id); err != nil {
return err
}
res, err := tx.Exec(`DELETE FROM exploration_nodes
WHERE id=$1 AND exploration_id=$2 AND kind='intent' AND state='open'`, id, s.expID)
if err != nil {
return err
}
removed, err := res.RowsAffected()
if err != nil {
return err
}
if removed != 1 {
return fmt.Errorf("open intent %d was not available for admission rollback", id)
}
return tx.Commit()
}
+285
View File
@@ -0,0 +1,285 @@
package db
import (
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestCompareAndSetIntentStateAllowsSingleWinner(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("intent CAS", "only one controller wins")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "controlled"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if claimed, err := store.ClaimIntent(intentID, "worker"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
const contenders = 12
var winners atomic.Int32
var wg sync.WaitGroup
start := make(chan struct{})
for range contenders {
wg.Add(1)
go func() {
defer wg.Done()
<-start
changed, transitionErr := store.CompareAndSetIntentState(intentID, "running", "paused")
if transitionErr != nil {
t.Errorf("transition: %v", transitionErr)
return
}
if changed {
winners.Add(1)
}
}()
}
close(start)
wg.Wait()
if got := winners.Load(); got != 1 {
t.Fatalf("CAS winners=%d, want 1", got)
}
node, err := store.GetNode(intentID)
if err != nil || node == nil || node.State != "paused" {
t.Fatalf("node=%+v err=%v, want paused", node, err)
}
}
func TestCancelIntentPreservesTokenMeteringWithoutDoubleCount(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("cancel token rollup", "preserve consumed tokens")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "cancelled"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if claimed, err := store.ClaimIntent(intentID, "work#1"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
baselineDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
baseline := dailyTokenBuckets(baselineDaily)
appendUsage := func(kind string, at time.Time, input, output, read, write *int) {
t.Helper()
activityID, appendErr := store.AppendActivity(Activity{
NodeID: &intentID, Worker: "work#1", Kind: kind,
InputTokens: input, OutputTokens: output,
CacheReadTokens: read, CacheWriteTokens: write,
})
if appendErr != nil {
t.Fatal(appendErr)
}
if _, updateErr := d.Exec(`UPDATE activity SET created_at=$1 WHERE id=$2`, at, activityID); updateErr != nil {
t.Fatal(updateErr)
}
}
values := func(input, output, read, write int) (*int, *int, *int, *int) {
return &input, &output, &read, &write
}
now := time.Now().UTC()
atUTCNoon := func(daysAgo int) time.Time {
at := now.AddDate(0, 0, -daysAgo)
return time.Date(at.Year(), at.Month(), at.Day(), 12, 0, 0, 0, time.UTC)
}
authoritativeDay := atUTCNoon(6)
fallbackDay := atUTCNoon(4)
trailingDay := atUTCNoon(2)
// Run 1 has an authoritative result; its preceding cumulative frame must not
// be counted a second time or move usage to the frame's earlier date.
i, o, r, w := values(100, 100, 100, 100)
appendUsage("usage", atUTCNoon(7), i, o, r, w)
i, o, r, w = values(10, 11, 12, 13)
appendUsage("result", authoritativeDay, i, o, r, w)
// Run 2 models a legacy failed result that omitted usage. Its fallback belongs
// to the result's date, not the cumulative frame's date.
i, o, r, w = values(20, 21, 22, 23)
appendUsage("usage", atUTCNoon(5), i, o, r, w)
appendUsage("result", fallbackDay, nil, nil, nil, nil)
// Run 3 was interrupted before a terminal result was persisted.
i, o, r, w = values(30, 31, 32, 33)
appendUsage("usage", trailingDay, i, o, r, w)
beforeCancelDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
beforeCancel := dailyTokenBuckets(beforeCancelDaily)
assertDailyTokenDelta(t, beforeCancel, baseline, authoritativeDay, 10, 11, 12)
assertDailyTokenDelta(t, beforeCancel, baseline, fallbackDay, 0, 0, 0)
assertDailyTokenDelta(t, beforeCancel, baseline, trailingDay, 0, 0, 0)
if err := store.SetIntentState(intentID, "paused"); err != nil {
t.Fatal(err)
}
cleanup, err := store.CancelIntent(intentID)
if err != nil {
t.Fatal(err)
}
if cleanup.Activities != 5 {
t.Fatalf("deleted activities=%d, want 5", cleanup.Activities)
}
total, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, total, 60, 63, 66, 69)
afterCancelDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
afterCancel := dailyTokenBuckets(afterCancelDaily)
assertDailyTokenDelta(t, afterCancel, baseline, authoritativeDay, 10, 11, 12)
assertDailyTokenDelta(t, afterCancel, baseline, fallbackDay, 20, 21, 22)
assertDailyTokenDelta(t, afterCancel, baseline, trailingDay, 30, 31, 32)
assertDailyTokenDelta(t, afterCancel, baseline, now, 0, 0, 0)
sessions, err := store.TokenStatsBySession()
if err != nil {
t.Fatal(err)
}
if len(sessions) != 0 {
t.Fatalf("cancelled intent leaked into sessions: %+v", sessions)
}
var rollups, datedRollups int
if err := d.QueryRow(`SELECT COUNT(*), COUNT(*) FILTER (
WHERE metadata->>'token_day'=TO_CHAR(created_at AT TIME ZONE 'UTC', 'YYYY-MM-DD')
) FROM activity
WHERE exploration_id=$1 AND worker='token-ledger' AND kind='result'
AND metadata->>'cancelled_intent_id'=$2`, expID, fmt.Sprint(intentID)).Scan(&rollups, &datedRollups); err != nil {
t.Fatal(err)
}
if rollups != 3 || datedRollups != rollups {
t.Fatalf("token rollups=%d dated=%d, want three correctly dated rows", rollups, datedRollups)
}
if _, err := store.CancelIntent(intentID); err == nil {
t.Fatal("second cancellation unexpectedly succeeded")
}
afterRetry, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, afterRetry, 60, 63, 66, 69)
afterRetryDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
assertDailyTokenBucketsEqual(t, dailyTokenBuckets(afterRetryDaily), afterCancel)
}
func dailyTokenBuckets(items []DailyTokenBucket) map[string]DailyTokenBucket {
out := make(map[string]DailyTokenBucket, len(items))
for _, item := range items {
out[item.Day] = item
}
return out
}
func assertDailyTokenDelta(t *testing.T, got, baseline map[string]DailyTokenBucket, at time.Time, input, output, read int) {
t.Helper()
day := at.UTC().Format(time.DateOnly)
actual, initial := got[day], baseline[day]
if actual.InputTokens-initial.InputTokens != input ||
actual.OutputTokens-initial.OutputTokens != output ||
actual.CacheReadTokens-initial.CacheReadTokens != read {
t.Fatalf("token delta for %s = input:%d output:%d cache-read:%d, want %d/%d/%d",
day, actual.InputTokens-initial.InputTokens, actual.OutputTokens-initial.OutputTokens,
actual.CacheReadTokens-initial.CacheReadTokens, input, output, read)
}
}
func assertDailyTokenBucketsEqual(t *testing.T, got, want map[string]DailyTokenBucket) {
t.Helper()
if len(got) != len(want) {
t.Fatalf("daily token buckets changed after retry: got=%+v want=%+v", got, want)
}
for day, expected := range want {
if actual, ok := got[day]; !ok || actual != expected {
t.Fatalf("daily token bucket %s changed after retry: got=%+v want=%+v", day, actual, expected)
}
}
}
func TestCancelIntentPreservesOutputsYieldedByAnotherIntent(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("shared intent output", "preserve shared facts and findings")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
first, err := store.AddIntent(map[string]any{"summary": "first"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
second, err := store.AddIntent(map[string]any{"summary": "second"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
fact, err := store.AddNode(KindFact, map[string]any{"text": "shared fact"}, 1, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
finding, err := store.AddNode(KindFinding, map[string]any{"summary": "shared finding"}, 1, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
for _, intentID := range []int64{first, second} {
if err := store.Link(intentID, RelYields, fact); err != nil {
t.Fatal(err)
}
if err := store.Link(intentID, RelYields, finding); err != nil {
t.Fatal(err)
}
}
if claimed, err := store.ClaimIntent(first, "work#1"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
if err := store.SetIntentState(first, "paused"); err != nil {
t.Fatal(err)
}
cleanup, err := store.CancelIntent(first)
if err != nil {
t.Fatal(err)
}
if cleanup.Intents != 1 || cleanup.Facts != 0 || cleanup.Findings != 0 {
t.Fatalf("cleanup=%+v, want only the cancelled intent", cleanup)
}
for _, nodeID := range []int64{fact, finding} {
node, getErr := store.GetNode(nodeID)
if getErr != nil || node == nil {
t.Fatalf("shared node %d was removed: node=%+v err=%v", nodeID, node, getErr)
}
}
}
+14 -6
View File
@@ -29,8 +29,12 @@ target AS (
JOIN task_scope ts ON (
(ts.kind='company' AND a.company_id = ts.company_id)
OR (ts.kind='root_domain' AND a.root_domain = ts.domain)
OR (ts.kind='subdomain' AND a.domain = ts.domain)
OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet)
OR (ts.kind='subdomain' AND a.domain = ts.domain)
OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet)
OR (ts.kind='icp' AND (
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value
))
)
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
UNION
@@ -53,7 +57,7 @@ tested AS (
func (s *AssetStore) ListTaskScopeWithSources(taskID int64) ([]TaskScope, error) {
rows, err := s.db.Query(`WITH `+directTaskContextCTE+`
SELECT ts.id, ts.task_id, ts.kind, COALESCE(ts.company_id,0), COALESCE(c.name,''), COALESCE(ts.domain,''),
COALESCE(ts.net::text,''), ts.source, COALESCE(ts.reason,'')
COALESCE(ts.net::text,''), COALESCE(ts.value,''), ts.source, COALESCE(ts.reason,'')
FROM task_scope ts
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
LEFT JOIN companies c ON c.id=ts.company_id
@@ -66,7 +70,7 @@ ORDER BY CASE WHEN ts.task_id=$1 THEN 0 ELSE 1 END, ts.id`, taskID)
for rows.Next() {
var scope TaskScope
var companyID int64
if err := rows.Scan(&scope.ID, &scope.TaskID, &scope.Kind, &companyID, &scope.CompanyName, &scope.Domain, &scope.Net, &scope.Source, &scope.Reason); err != nil {
if err := rows.Scan(&scope.ID, &scope.TaskID, &scope.Kind, &companyID, &scope.CompanyName, &scope.Domain, &scope.Net, &scope.Value, &scope.Source, &scope.Reason); err != nil {
return nil, err
}
if companyID > 0 {
@@ -175,8 +179,12 @@ context_assets AS (
JOIN task_scope ts ON (
(ts.kind='company' AND a.company_id=ts.company_id)
OR (ts.kind='root_domain' AND a.root_domain=ts.domain)
OR (ts.kind='subdomain' AND a.domain=ts.domain)
OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet)
OR (ts.kind='subdomain' AND a.domain=ts.domain)
OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet)
OR (ts.kind='icp' AND (
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value
))
)
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
)
+419
View File
@@ -0,0 +1,419 @@
package db
import (
"database/sql"
"fmt"
"strings"
"time"
"unicode/utf8"
)
const IntentBlockedLLMQuota = "llm_quota_exhausted"
// TaskLLMProfile is one ordered entry in a task's explicit failover chain.
type TaskLLMProfile struct {
ProfileID int64 `json:"profile_id"`
Position int `json:"position"`
Status string `json:"status"`
LastError string `json:"last_error,omitempty"`
ExhaustedAt *time.Time `json:"exhausted_at,omitempty"`
}
// TaskSource identifies one directly related task and its exploration.
type TaskSource struct {
TaskID int64
ExplorationID int64
Description string
Goal string
Status string
}
// TaskLLMTransition reports the shared task-level result of marking one profile
// quota-exhausted. NextProfileID is nil when the explicit chain is exhausted.
type TaskLLMTransition struct {
PreviousProfileID int64
NextProfileID *int64
ChainExhausted bool
Advanced bool
// Stale means the chain was replaced after the failing call selected its
// provider. The caller may retry a pre-stream request against the new chain,
// but must not report or persist a transition for this result.
Stale bool
}
func (d *DB) hydrateTaskContext(t *Task) error {
if t == nil {
return nil
}
// A nil cursor on a non-empty chain is a persisted end-of-chain marker. Do not
// repair it from an earlier ready entry: the user may have manually selected a
// profile in the middle and legitimately exhausted every candidate after it.
legacyProfileID, activeProfileID, revision, chain, err := d.taskLLMContext(t.ID)
if err != nil {
return err
}
t.LLMProfileID = legacyProfileID
t.ActiveLLMProfileID = activeProfileID
t.LLMChainRevision = revision
sources, err := d.TaskSourceIDs(t.ID)
if err != nil {
return err
}
t.SourceTaskIDs = sources
companies, err := d.TaskCompanyIDs(t.ID)
if err != nil {
return err
}
t.CompanyIDs = companies
t.LLMProfileIDs = make([]int64, 0, len(chain))
t.LLMFailoverState = "default"
var latest *TaskLLMProfile
activeReady := false
for i := range chain {
entry := chain[i]
t.LLMProfileIDs = append(t.LLMProfileIDs, entry.ProfileID)
if t.ActiveLLMProfileID != nil && entry.ProfileID == *t.ActiveLLMProfileID && entry.Status == "ready" {
activeReady = true
}
if entry.ExhaustedAt != nil && (latest == nil || entry.ExhaustedAt.After(*latest.ExhaustedAt)) {
copy := entry
latest = &copy
}
}
if len(chain) > 0 {
if activeReady {
t.LLMFailoverState = "ready"
} else {
t.LLMFailoverState = "chain_exhausted"
}
}
if latest != nil {
t.LLMFailoverReason = latest.LastError
}
return nil
}
// taskLLMContext reads the compatibility profile, current cursor, and ordered
// chain in one statement. Runtime chain edits commit atomically, and one SQL
// statement gives hydration one matching snapshot instead of a transient mix of
// an old task cursor and a newly replaced chain (or vice versa).
func (d *DB) taskLLMContext(taskID int64) (*int64, *int64, int64, []TaskLLMProfile, error) {
rows, err := d.Query(`
SELECT t.llm_profile_id, t.active_llm_profile_id, t.llm_chain_revision,
p.profile_id, p.position, p.status, COALESCE(p.last_error,''), p.exhausted_at
FROM tasks t
LEFT JOIN task_llm_profiles p ON p.task_id=t.id
WHERE t.id=$1
ORDER BY p.position NULLS LAST`, taskID)
if err != nil {
return nil, nil, 0, nil, err
}
defer rows.Close()
var (
legacyProfileID *int64
activeProfileID *int64
revision int64
chain []TaskLLMProfile
found bool
)
for rows.Next() {
var legacy, active, profileID, position sql.NullInt64
var rowRevision int64
var status, lastError sql.NullString
var exhaustedAt sql.NullTime
if err := rows.Scan(&legacy, &active, &rowRevision, &profileID, &position, &status, &lastError, &exhaustedAt); err != nil {
return nil, nil, 0, nil, err
}
if !found {
found = true
revision = rowRevision
if legacy.Valid {
id := legacy.Int64
legacyProfileID = &id
}
if active.Valid {
id := active.Int64
activeProfileID = &id
}
}
if !profileID.Valid {
continue
}
entry := TaskLLMProfile{
ProfileID: profileID.Int64,
Position: int(position.Int64),
Status: status.String,
LastError: lastError.String,
}
if exhaustedAt.Valid {
ts := exhaustedAt.Time
entry.ExhaustedAt = &ts
}
chain = append(chain, entry)
}
if err := rows.Err(); err != nil {
return nil, nil, 0, nil, err
}
if !found {
return nil, nil, 0, nil, sql.ErrNoRows
}
return legacyProfileID, activeProfileID, revision, chain, nil
}
func (d *DB) TaskSourceIDs(taskID int64) ([]int64, error) {
rows, err := d.Query(`SELECT source_task_id FROM task_relations WHERE task_id=$1 ORDER BY created_at, source_task_id`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
out = append(out, id)
}
return out, rows.Err()
}
// TaskCompanyIDs returns the companies whose asset scopes are available to the
// task. The task_scope rows remain the single source of truth.
func (d *DB) TaskCompanyIDs(taskID int64) ([]int64, error) {
rows, err := d.Query(`
SELECT company_id FROM task_scope
WHERE task_id=$1 AND kind='company' AND company_id IS NOT NULL
ORDER BY id`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
out = append(out, id)
}
return out, rows.Err()
}
func (d *DB) TaskSources(taskID int64) ([]TaskSource, error) {
rows, err := d.Query(`
SELECT t.id, t.exploration_id, t.description, t.goal, t.status
FROM task_relations r
JOIN tasks t ON t.id=r.source_task_id AND t.deleted_at IS NULL
WHERE r.task_id=$1
ORDER BY r.created_at, r.source_task_id`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []TaskSource
for rows.Next() {
var source TaskSource
if err := rows.Scan(&source.TaskID, &source.ExplorationID, &source.Description, &source.Goal, &source.Status); err != nil {
return nil, err
}
out = append(out, source)
}
return out, rows.Err()
}
func (d *DB) TaskLLMProfiles(taskID int64) ([]TaskLLMProfile, error) {
rows, err := d.Query(`SELECT profile_id, position, status, COALESCE(last_error,''), exhausted_at
FROM task_llm_profiles WHERE task_id=$1 ORDER BY position`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []TaskLLMProfile
for rows.Next() {
var entry TaskLLMProfile
if err := rows.Scan(&entry.ProfileID, &entry.Position, &entry.Status, &entry.LastError, &entry.ExhaustedAt); err != nil {
return nil, err
}
out = append(out, entry)
}
return out, rows.Err()
}
// ReplaceTaskLLMProfiles atomically replaces and resets the explicit task chain.
// activeProfileID=0 selects the first entry. An empty list restores the existing
// agent-binding/global fallback behavior.
func (d *DB) ReplaceTaskLLMProfiles(taskID int64, profileIDs []int64, activeProfileID int64) error {
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
// Keep the task -> profile lock order shared by all task LLM mutations and
// DeleteProfile. Inserts below may take KEY SHARE locks on llm_profiles for
// their foreign keys, so the task row must be locked before any of them.
var status string
if err := tx.QueryRow(`SELECT status FROM tasks WHERE id=$1 AND deleted_at IS NULL FOR UPDATE`, taskID).Scan(&status); err != nil {
if err == sql.ErrNoRows {
return fmt.Errorf("task %d not found", taskID)
}
return err
}
if IsTerminal(status) {
return fmt.Errorf("terminal task cannot change LLM profiles")
}
if _, err := tx.Exec(`DELETE FROM task_llm_profiles WHERE task_id=$1`, taskID); err != nil {
return err
}
if err := insertTaskLLMProfiles(tx, taskID, profileIDs); err != nil {
return err
}
if len(profileIDs) == 0 {
activeProfileID = 0
} else if activeProfileID == 0 {
activeProfileID = profileIDs[0]
}
if activeProfileID > 0 {
found := false
for _, id := range profileIDs {
found = found || id == activeProfileID
}
if !found {
return fmt.Errorf("active LLM profile %d is not in the task chain", activeProfileID)
}
}
var active any
if activeProfileID > 0 {
active = activeProfileID
}
if _, err := tx.Exec(`UPDATE tasks
SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, taskID, active); err != nil {
return err
}
return tx.Commit()
}
// MarkTaskLLMProfileQuotaExhausted advances the shared task cursor once. A late
// in-flight error from an older profile records that entry as exhausted but does
// not advance past the profile another call already selected.
func (d *DB) MarkTaskLLMProfileQuotaExhausted(taskID, profileID int64, reason string) (TaskLLMTransition, error) {
return d.markTaskLLMProfileQuotaExhausted(taskID, profileID, 0, false, reason)
}
// MarkTaskLLMProfileQuotaExhaustedAtRevision applies a provider failure only
// when it belongs to the chain snapshot used to start that call.
func (d *DB) MarkTaskLLMProfileQuotaExhaustedAtRevision(taskID, profileID, revision int64, reason string) (TaskLLMTransition, error) {
return d.markTaskLLMProfileQuotaExhausted(taskID, profileID, revision, true, reason)
}
func (d *DB) markTaskLLMProfileQuotaExhausted(taskID, profileID, revision int64, checkRevision bool, reason string) (TaskLLMTransition, error) {
var out TaskLLMTransition
tx, err := d.Begin()
if err != nil {
return out, err
}
defer tx.Rollback()
var (
active sql.NullInt64
currentRevision int64
)
if err := tx.QueryRow(`SELECT active_llm_profile_id, llm_chain_revision FROM tasks WHERE id=$1 FOR UPDATE`, taskID).Scan(&active, &currentRevision); err != nil {
return out, err
}
out.PreviousProfileID = profileID
if checkRevision && revision != currentRevision {
out.Stale = true
return out, tx.Commit()
}
var (
position int
entryStatus string
)
if err := tx.QueryRow(`SELECT position, status FROM task_llm_profiles WHERE task_id=$1 AND profile_id=$2`, taskID, profileID).
Scan(&position, &entryStatus); err != nil {
if err == sql.ErrNoRows {
// The chain was edited while this request was in flight. Its provider
// error remains valid for the caller, but it must not mutate the new chain.
return out, tx.Commit()
}
return out, err
}
// Once the cursor reached the end, late failures must be idempotent. This also
// prevents an older in-flight request from reviving a ready entry before a
// manually selected starting position.
if !active.Valid {
out.ChainExhausted = true
return out, tx.Commit()
}
reason = truncateUTF8(strings.TrimSpace(reason), 1000)
if entryStatus != "quota_exhausted" {
if _, err := tx.Exec(`UPDATE task_llm_profiles
SET status='quota_exhausted', last_error=$3, exhausted_at=now()
WHERE task_id=$1 AND profile_id=$2`, taskID, profileID, reason); err != nil {
return out, err
}
}
if active.Valid && active.Int64 != profileID {
next := active.Int64
out.NextProfileID = &next
return out, tx.Commit()
}
var next int64
err = tx.QueryRow(`SELECT profile_id FROM task_llm_profiles
WHERE task_id=$1 AND position>$2 AND status='ready'
ORDER BY position LIMIT 1`, taskID, position).Scan(&next)
switch err {
case nil:
out.Advanced = true
out.NextProfileID = &next
if _, err := tx.Exec(`UPDATE tasks
SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, taskID, next); err != nil {
return out, err
}
case sql.ErrNoRows:
out.Advanced = true
out.ChainExhausted = true
if _, err := tx.Exec(`UPDATE tasks
SET active_llm_profile_id=NULL, llm_profile_id=NULL, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, taskID); err != nil {
return out, err
}
default:
return out, err
}
return out, tx.Commit()
}
func truncateUTF8(value string, maxBytes int) string {
value = strings.ToValidUTF8(value, "\uFFFD")
if maxBytes <= 0 {
return ""
}
if len(value) <= maxBytes {
return value
}
end := maxBytes
for end > 0 && !utf8.ValidString(value[:end]) {
end--
}
return value[:end]
}
func (s *ExplorationStore) SetIntentBlockedReason(id int64, reason string) error {
_, err := s.db.Exec(`UPDATE exploration_nodes
SET state='blocked', blocked_reason=NULLIF($1,''), completed_at=now()
WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, reason, id, s.expID)
return err
}
func (s *ExplorationStore) ReopenIntentsByBlockedReason(reason string) (int64, error) {
res, err := s.db.Exec(`UPDATE exploration_nodes
SET state='open', blocked_reason=NULL, completed_at=NULL
WHERE exploration_id=$1 AND kind='intent' AND state='blocked' AND blocked_reason=$2`, s.expID, reason)
if err != nil {
return 0, err
}
n, _ := res.RowsAffected()
return n, nil
}
+527
View File
@@ -0,0 +1,527 @@
package db
import (
"database/sql"
"fmt"
"testing"
"time"
)
func TestTaskLLMProfileMutationsLockTaskBeforeProfile(t *testing.T) {
dsn := testDSN(t)
d, err := Open(dsn)
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
first, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("lock-order-first-%d", suffix), Format: "openai", Model: "first", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
second, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("lock-order-second-%d", suffix), Format: "openai", Model: "second", APIKey: "test-key",
})
if err != nil {
_ = d.DeleteProfile(first)
t.Fatal(err)
}
task, err := d.CreateTaskWithOptions("LLM lock order", "verify concurrent mutation locks", TaskCreateOptions{
LLMProfileIDs: []int64{first, second},
})
if err != nil {
_ = d.DeleteProfile(first)
_ = d.DeleteProfile(second)
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteTask(task.ID)
_ = d.DeleteProfile(first)
_ = d.DeleteProfile(second)
})
replaceDB, replacePID := openSingleConnectionTestDB(t, dsn)
replaceBlocker, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer replaceBlocker.Rollback()
if _, err := replaceBlocker.Exec(`SELECT id FROM tasks WHERE id=$1 FOR UPDATE`, task.ID); err != nil {
replaceBlocker.Rollback()
t.Fatal(err)
}
replaceDone := make(chan error, 1)
go func() {
replaceDone <- replaceDB.ReplaceTaskLLMProfiles(task.ID, []int64{second, first}, second)
}()
if err := waitForBackendBlock(d, replacePID, replaceDone); err != nil {
replaceBlocker.Rollback()
t.Fatal(err)
}
assertProfilesUnlocked(t, d, first, second)
if err := replaceBlocker.Rollback(); err != nil {
t.Fatal(err)
}
if err := waitForMutationResult(replaceDone); err != nil {
t.Fatalf("replace chain: %v", err)
}
deleteDB, deletePID := openSingleConnectionTestDB(t, dsn)
deleteBlocker, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer deleteBlocker.Rollback()
if _, err := deleteBlocker.Exec(`SELECT id FROM tasks WHERE id=$1 FOR UPDATE`, task.ID); err != nil {
deleteBlocker.Rollback()
t.Fatal(err)
}
deleteDone := make(chan error, 1)
go func() {
deleteDone <- deleteDB.DeleteProfile(second)
}()
if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil {
deleteBlocker.Rollback()
t.Fatal(err)
}
assertProfilesUnlocked(t, d, second)
if err := deleteBlocker.Rollback(); err != nil {
t.Fatal(err)
}
if err := waitForMutationResult(deleteDone); err != nil {
t.Fatalf("delete profile: %v", err)
}
got, err := d.GetTask(task.ID)
if err != nil {
t.Fatal(err)
}
if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != first {
t.Fatalf("deleting the active profile did not select its successor: %+v", got)
}
}
func TestDeleteProfileLocksNonTaskReferencesBeforeProfile(t *testing.T) {
dsn := testDSN(t)
t.Run("agent", func(t *testing.T) {
d, err := Open(dsn)
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
profileID, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("agent-lock-profile-%d", suffix), Format: "openai", Model: "agent-lock", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
agent, err := d.CreateAgent(fmt.Sprintf("lock_agent_%d", suffix), "lock agent", "")
if err != nil {
_ = d.DeleteProfile(profileID)
t.Fatal(err)
}
if err := d.SetAgentLLMProfile(agent.Key, &profileID); err != nil {
_ = d.DeleteAgent(agent.Key)
_ = d.DeleteProfile(profileID)
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteAgent(agent.Key)
_ = d.DeleteProfile(profileID)
})
blocker, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer blocker.Rollback()
if _, err := blocker.Exec(`SELECT id FROM agents WHERE id=$1 FOR UPDATE`, agent.ID); err != nil {
t.Fatal(err)
}
deleteDB, deletePID := openSingleConnectionTestDB(t, dsn)
deleteDone := make(chan error, 1)
go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }()
if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil {
blocker.Rollback()
t.Fatal(err)
}
assertProfilesUnlocked(t, d, profileID)
if err := blocker.Rollback(); err != nil {
t.Fatal(err)
}
if err := waitForMutationResult(deleteDone); err != nil {
t.Fatalf("delete agent profile: %v", err)
}
got, err := d.GetAgentByKey(agent.Key)
if err != nil || got == nil || got.LLMProfileID != nil {
t.Fatalf("agent binding was not cleared: agent=%+v err=%v", got, err)
}
})
t.Run("conversation", func(t *testing.T) {
d, err := Open(dsn)
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
profileID, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("conversation-lock-profile-%d", suffix), Format: "openai", Model: "conversation-lock", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
conversation, err := d.CreateConversation("planner", "lock conversation", &profileID)
if err != nil {
_ = d.DeleteProfile(profileID)
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteConversation(conversation.ID)
_ = d.DeleteProfile(profileID)
})
blocker, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer blocker.Rollback()
if _, err := blocker.Exec(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, conversation.ID); err != nil {
t.Fatal(err)
}
deleteDB, deletePID := openSingleConnectionTestDB(t, dsn)
deleteDone := make(chan error, 1)
go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }()
if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil {
blocker.Rollback()
t.Fatal(err)
}
assertProfilesUnlocked(t, d, profileID)
if err := blocker.Rollback(); err != nil {
t.Fatal(err)
}
if err := waitForMutationResult(deleteDone); err != nil {
t.Fatalf("delete conversation profile: %v", err)
}
got, err := d.GetConversation(conversation.ID)
if err != nil || got == nil || got.LLMProfileID != nil {
t.Fatalf("conversation binding was not cleared: conversation=%+v err=%v", got, err)
}
})
}
func TestNonTaskProfileMutationsLockReferenceBeforeProfile(t *testing.T) {
dsn := testDSN(t)
t.Run("agent", func(t *testing.T) {
d, err := Open(dsn)
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
profileID, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("agent-write-lock-profile-%d", suffix), Format: "openai", Model: "agent-write-lock", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
agent, err := d.CreateAgent(fmt.Sprintf("write_lock_agent_%d", suffix), "write lock agent", "")
if err != nil {
_ = d.DeleteProfile(profileID)
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteAgent(agent.Key)
_ = d.DeleteProfile(profileID)
})
blocker, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer blocker.Rollback()
if _, err := blocker.Exec(`SELECT id FROM agents WHERE id=$1 FOR UPDATE`, agent.ID); err != nil {
t.Fatal(err)
}
mutationDB, mutationPID := openSingleConnectionTestDB(t, dsn)
mutationDone := make(chan error, 1)
go func() { mutationDone <- mutationDB.SetAgentLLMProfile(agent.Key, &profileID) }()
if err := waitForBackendBlock(d, mutationPID, mutationDone); err != nil {
blocker.Rollback()
t.Fatal(err)
}
assertProfilesUnlocked(t, d, profileID)
if err := blocker.Rollback(); err != nil {
t.Fatal(err)
}
if err := waitForMutationResult(mutationDone); err != nil {
t.Fatalf("bind agent profile: %v", err)
}
})
t.Run("conversation", func(t *testing.T) {
d, err := Open(dsn)
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
profileID, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("conversation-write-lock-profile-%d", suffix), Format: "openai", Model: "conversation-write-lock", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
conversation, err := d.CreateConversation("planner", "write lock conversation", nil)
if err != nil {
_ = d.DeleteProfile(profileID)
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteConversation(conversation.ID)
_ = d.DeleteProfile(profileID)
})
blocker, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer blocker.Rollback()
if _, err := blocker.Exec(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, conversation.ID); err != nil {
t.Fatal(err)
}
mutationDB, mutationPID := openSingleConnectionTestDB(t, dsn)
mutationDone := make(chan error, 1)
go func() { mutationDone <- mutationDB.UpdateConversationProfile(conversation.ID, &profileID) }()
if err := waitForBackendBlock(d, mutationPID, mutationDone); err != nil {
blocker.Rollback()
t.Fatal(err)
}
assertProfilesUnlocked(t, d, profileID)
if err := blocker.Rollback(); err != nil {
t.Fatal(err)
}
if err := waitForMutationResult(mutationDone); err != nil {
t.Fatalf("bind conversation profile: %v", err)
}
})
}
func TestCreateConversationAndDeleteProfileDoNotDeadlock(t *testing.T) {
dsn := testDSN(t)
d, err := Open(dsn)
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
profileID, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("conversation-create-race-profile-%d", suffix), Format: "openai", Model: "conversation-create-race", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
title := fmt.Sprintf("conversation create race %d", suffix)
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM conversations WHERE title=$1`, title)
_ = d.DeleteProfile(profileID)
})
profileBlocker, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer profileBlocker.Rollback()
if _, err := profileBlocker.Exec(`SELECT id FROM llm_profiles WHERE id=$1 FOR UPDATE`, profileID); err != nil {
t.Fatal(err)
}
createDB, createPID := openSingleConnectionTestDB(t, dsn)
type createResult struct {
conversation *Conversation
err error
}
createDone := make(chan createResult, 1)
createStatus := make(chan error, 1)
go func() {
conversation, err := createDB.CreateConversation("planner", title, &profileID)
createDone <- createResult{conversation: conversation, err: err}
createStatus <- err
}()
if err := waitForBackendBlock(d, createPID, createStatus); err != nil {
profileBlocker.Rollback()
t.Fatal(err)
}
deleteDB, deletePID := openSingleConnectionTestDB(t, dsn)
deleteDone := make(chan error, 1)
go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }()
if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil {
profileBlocker.Rollback()
t.Fatal(err)
}
if err := profileBlocker.Rollback(); err != nil {
t.Fatal(err)
}
var created createResult
select {
case created = <-createDone:
case <-time.After(12 * time.Second):
t.Fatal("timed out waiting for conversation creation")
}
if err := waitForMutationResult(deleteDone); err != nil {
t.Fatalf("delete profile during conversation creation: %v", err)
}
var persisted int
if err := d.QueryRow(`SELECT count(*) FROM conversations WHERE title=$1`, title).Scan(&persisted); err != nil {
t.Fatal(err)
}
if created.err == nil {
if created.conversation == nil || persisted != 1 {
t.Fatalf("successful creation was not committed atomically: conversation=%+v count=%d", created.conversation, persisted)
}
} else if persisted != 0 {
t.Fatalf("failed creation left a partial conversation row: err=%v count=%d", created.err, persisted)
}
}
func TestDeleteProfileRetriesReferenceCommittedAfterInitialScan(t *testing.T) {
dsn := testDSN(t)
d, err := Open(dsn)
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
profileID, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("late-reference-profile-%d", suffix), Format: "openai", Model: "late-reference", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
task, err := d.CreateTask("late profile reference", "exercise delete retry", nil, 0, 0)
if err != nil {
_ = d.DeleteProfile(profileID)
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteTask(task.ID)
_ = d.DeleteProfile(profileID)
})
// Keep the new reference uncommitted while deletion takes its initial
// READ COMMITTED snapshot. The task is therefore absent from the first lock
// set, while its FK KEY SHARE lock makes deletion wait at the profile row.
referenceTx, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer referenceTx.Rollback()
if _, err := referenceTx.Exec(`UPDATE tasks
SET llm_profile_id=$2, active_llm_profile_id=$2
WHERE id=$1`, task.ID, profileID); err != nil {
t.Fatal(err)
}
if _, err := referenceTx.Exec(`INSERT INTO task_llm_profiles(task_id, profile_id, position)
VALUES ($1,$2,0)`, task.ID, profileID); err != nil {
t.Fatal(err)
}
deleteDB, deletePID := openSingleConnectionTestDB(t, dsn)
deleteDone := make(chan error, 1)
go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }()
if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil {
referenceTx.Rollback()
t.Fatal(err)
}
if err := referenceTx.Commit(); err != nil {
t.Fatal(err)
}
if err := waitForMutationResult(deleteDone); err != nil {
t.Fatalf("delete profile after late reference: %v", err)
}
got, err := d.GetTask(task.ID)
if err != nil {
t.Fatal(err)
}
if got.ActiveLLMProfileID != nil || len(got.LLMProfileIDs) != 0 || got.LLMChainRevision != 1 {
t.Fatalf("late task reference was not handled by a locked retry: %+v", got)
}
}
func openSingleConnectionTestDB(t *testing.T, dsn string) (*DB, int) {
t.Helper()
sqlDB, err := sql.Open("pgx", dsn)
if err != nil {
t.Fatal(err)
}
sqlDB.SetMaxOpenConns(1)
sqlDB.SetMaxIdleConns(1)
t.Cleanup(func() { _ = sqlDB.Close() })
if _, err := sqlDB.Exec(`SET statement_timeout='10s'`); err != nil {
t.Fatal(err)
}
var pid int
if err := sqlDB.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil {
t.Fatal(err)
}
return &DB{sqlDB}, pid
}
func waitForBackendBlock(observer *DB, pid int, done <-chan error) error {
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
select {
case err := <-done:
return fmt.Errorf("mutation returned before reaching the expected reference-row lock: %v", err)
default:
}
var blockers int
if err := observer.QueryRow(`SELECT cardinality(pg_blocking_pids($1))`, pid).Scan(&blockers); err != nil {
return err
}
if blockers > 0 {
return nil
}
time.Sleep(10 * time.Millisecond)
}
return fmt.Errorf("backend %d did not block within 5s", pid)
}
func assertProfilesUnlocked(t *testing.T, d *DB, profileIDs ...int64) {
t.Helper()
probe, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer probe.Rollback()
for _, profileID := range profileIDs {
var lockedID int64
if err := probe.QueryRow(`SELECT id FROM llm_profiles WHERE id=$1 FOR UPDATE NOWAIT`, profileID).Scan(&lockedID); err != nil {
t.Fatalf("profile %d was locked before the task row: %v", profileID, err)
}
}
}
func waitForMutationResult(done <-chan error) error {
select {
case err := <-done:
return err
case <-time.After(12 * time.Second):
return fmt.Errorf("timed out waiting for task LLM mutation")
}
}
+27
View File
@@ -0,0 +1,27 @@
package db
import (
"strings"
"testing"
"unicode/utf8"
)
func TestTruncateUTF8PreservesValidEncoding(t *testing.T) {
t.Parallel()
input := strings.Repeat("额度不足", 400)
got := truncateUTF8(input, 1000)
if len(got) > 1000 {
t.Fatalf("truncated value has %d bytes, want at most 1000", len(got))
}
if !utf8.ValidString(got) {
t.Fatalf("truncated value is not valid UTF-8: %q", got[len(got)-8:])
}
}
func TestTruncateUTF8RepairsInvalidInput(t *testing.T) {
t.Parallel()
got := truncateUTF8("bad\xffvalue", 1000)
if !utf8.ValidString(got) {
t.Fatalf("repaired value is not valid UTF-8: %q", got)
}
}
+210
View File
@@ -0,0 +1,210 @@
package db
import (
"database/sql"
"fmt"
"testing"
"time"
_ "github.com/jackc/pgx/v5/stdlib"
)
func TestDeleteTaskTrafficHostsUseOneLockedTransaction(t *testing.T) {
dsn := testDSN(t)
d, err := Open(dsn)
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
t.Run("sharing committed before deletion lock is observed", func(t *testing.T) {
first, second, host, rootAssetID := createTaskDeleteRaceFixture(t, d)
serviceURL := "https://" + host + "/concurrent-owner"
t.Cleanup(func() {
_ = d.DeleteTask(first.ID)
_ = d.DeleteTask(second.ID)
_, _ = d.Exec(`DELETE FROM assets WHERE id=$1 OR url=$2`, rootAssetID, serviceURL)
})
writer, _ := openTaskDeleteTestDB(t, dsn)
writerTx, err := writer.Begin()
if err != nil {
t.Fatal(err)
}
defer writerTx.Rollback()
if _, err := writerTx.Exec(`
INSERT INTO assets(type, url, service_type, domain, task_ids)
VALUES ('service', $1, 'http', $2, ARRAY[$3]::bigint[])`, serviceURL, host, second.ID); err != nil {
t.Fatal(err)
}
deleter, deleterPID := openTaskDeleteTestDB(t, dsn)
var preparedHosts []string
deleteDone := make(chan error, 1)
go func() {
_, deleteErr := deleter.DeleteTaskCascadePrepared(first.ID, true, false, false, func(p TaskDeletePreparation) error {
preparedHosts = append([]string(nil), p.TrafficHosts...)
return nil
})
deleteDone <- deleteErr
}()
if err := waitForTaskDeleteBlock(d, deleterPID, deleteDone); err != nil {
t.Fatal(err)
}
if err := writerTx.Commit(); err != nil {
t.Fatal(err)
}
if err := waitForTaskDeleteResult(deleteDone); err != nil {
t.Fatal(err)
}
if containsDeleteHost(preparedHosts, host) {
t.Fatalf("newly shared host %q was selected for traffic deletion: %v", host, preparedHosts)
}
var remaining int
if err := d.QueryRow(`SELECT count(*) FROM assets WHERE url=$1 AND $2=ANY(task_ids)`, serviceURL, second.ID).Scan(&remaining); err != nil {
t.Fatal(err)
}
if remaining != 1 {
t.Fatalf("concurrent owner's asset was not preserved: count=%d", remaining)
}
})
t.Run("sharing cannot enter between host resolution and commit", func(t *testing.T) {
first, second, host, rootAssetID := createTaskDeleteRaceFixture(t, d)
serviceURL := "https://" + host + "/late-owner"
t.Cleanup(func() {
_ = d.DeleteTask(first.ID)
_ = d.DeleteTask(second.ID)
_, _ = d.Exec(`DELETE FROM assets WHERE id=$1 OR url=$2`, rootAssetID, serviceURL)
})
deleter, _ := openTaskDeleteTestDB(t, dsn)
prepared := make(chan TaskDeletePreparation, 1)
releasePrepare := make(chan struct{})
deleteDone := make(chan error, 1)
go func() {
_, deleteErr := deleter.DeleteTaskCascadePrepared(first.ID, true, false, false, func(p TaskDeletePreparation) error {
prepared <- p
<-releasePrepare
return nil
})
deleteDone <- deleteErr
}()
var plan TaskDeletePreparation
select {
case plan = <-prepared:
case err := <-deleteDone:
close(releasePrepare)
t.Fatalf("delete returned before preparation: %v", err)
case <-time.After(5 * time.Second):
close(releasePrepare)
t.Fatal("timed out waiting for task delete preparation")
}
if !containsDeleteHost(plan.TrafficHosts, host) {
close(releasePrepare)
t.Fatalf("exclusive host %q missing from preparation: %v", host, plan.TrafficHosts)
}
writer, writerPID := openTaskDeleteTestDB(t, dsn)
writerDone := make(chan error, 1)
go func() {
_, writeErr := writer.Exec(`
INSERT INTO assets(type, url, service_type, domain, task_ids)
VALUES ('service', $1, 'http', $2, ARRAY[$3]::bigint[])`, serviceURL, host, second.ID)
writerDone <- writeErr
}()
if err := waitForTaskDeleteBlock(d, writerPID, writerDone); err != nil {
close(releasePrepare)
t.Fatal(err)
}
close(releasePrepare)
if err := waitForTaskDeleteResult(deleteDone); err != nil {
t.Fatal(err)
}
if err := waitForTaskDeleteResult(writerDone); err != nil {
t.Fatal(err)
}
})
}
func createTaskDeleteRaceFixture(t *testing.T, d *DB) (first, second *Task, host string, rootAssetID int64) {
t.Helper()
var err error
first, err = d.CreateTask("delete race owner", "delete safely", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
second, err = d.CreateTask("delete race sharer", "preserve shared host", nil, 0, 0)
if err != nil {
_ = d.DeleteTask(first.ID)
t.Fatal(err)
}
host = fmt.Sprintf("task-delete-race-%d.example.test", first.ID)
rootAssetID, err = d.Assets().UpsertRootDomain(UpsertRootDomainReq{Domain: host, TaskID: first.ID})
if err != nil {
_ = d.DeleteTask(first.ID)
_ = d.DeleteTask(second.ID)
t.Fatal(err)
}
return first, second, host, rootAssetID
}
func openTaskDeleteTestDB(t *testing.T, dsn string) (*DB, int) {
t.Helper()
sqlDB, err := sql.Open("pgx", dsn)
if err != nil {
t.Fatal(err)
}
sqlDB.SetMaxOpenConns(1)
sqlDB.SetMaxIdleConns(1)
t.Cleanup(func() { _ = sqlDB.Close() })
if _, err := sqlDB.Exec(`SET statement_timeout='10s'`); err != nil {
t.Fatal(err)
}
var pid int
if err := sqlDB.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil {
t.Fatal(err)
}
return &DB{sqlDB}, pid
}
func waitForTaskDeleteBlock(observer *DB, pid int, done <-chan error) error {
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
select {
case err := <-done:
return fmt.Errorf("operation returned before reaching the task deletion lock: %v", err)
default:
}
var blockers int
if err := observer.QueryRow(`SELECT cardinality(pg_blocking_pids($1))`, pid).Scan(&blockers); err != nil {
return err
}
if blockers > 0 {
return nil
}
time.Sleep(10 * time.Millisecond)
}
return fmt.Errorf("backend %d did not block within 5s", pid)
}
func waitForTaskDeleteResult(done <-chan error) error {
select {
case err := <-done:
return err
case <-time.After(12 * time.Second):
return fmt.Errorf("timed out waiting for concurrent task deletion operation")
}
}
func containsDeleteHost(hosts []string, want string) bool {
for _, host := range hosts {
if host == want {
return true
}
}
return false
}
+66
View File
@@ -0,0 +1,66 @@
package db
import (
"testing"
"time"
)
func TestTaskQueuePreservesBootstrapAndFIFOPosition(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
task, err := d.CreateTask("queue metadata", "keep first-run mode", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(task.ID)
if err := d.Enqueue(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
first, err := d.GetTask(task.ID)
if err != nil || first == nil || first.QueuedAt == nil {
t.Fatalf("first enqueue: task=%+v err=%v", first, err)
}
// A follow-up/rerun can try to admit an already queued task as resume. The
// original bootstrap mode and FIFO timestamp must remain authoritative.
time.Sleep(time.Millisecond)
if err := d.Enqueue(task.ID, "resume"); err != nil {
t.Fatal(err)
}
second, err := d.GetTask(task.ID)
if err != nil || second == nil {
t.Fatalf("second enqueue: task=%+v err=%v", second, err)
}
if second.QueueMode != "bootstrap" {
t.Fatalf("queue mode=%q, want bootstrap", second.QueueMode)
}
if second.QueuedAt == nil || !second.QueuedAt.Equal(*first.QueuedAt) {
t.Fatalf("repeated enqueue moved FIFO position: first=%v second=%v", first.QueuedAt, second.QueuedAt)
}
// Pausing a queued task removes it from the queue but keeps its required
// startup mode. A later requeue receives a new tail position.
if err := d.Dequeue(task.ID, false); err != nil {
t.Fatal(err)
}
paused, err := d.GetTask(task.ID)
if err != nil || paused == nil || paused.Queued || paused.QueueMode != "bootstrap" {
t.Fatalf("paused queue metadata: task=%+v err=%v", paused, err)
}
time.Sleep(time.Millisecond)
if err := d.Enqueue(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
requeued, err := d.GetTask(task.ID)
if err != nil || requeued == nil || requeued.QueuedAt == nil {
t.Fatalf("requeue: task=%+v err=%v", requeued, err)
}
if !requeued.QueuedAt.After(*first.QueuedAt) {
t.Fatalf("requeue did not move to FIFO tail: first=%v requeued=%v", first.QueuedAt, requeued.QueuedAt)
}
}
+47
View File
@@ -0,0 +1,47 @@
package db
import (
"errors"
"reflect"
"strings"
"testing"
)
func TestCreateTaskRejectsTooManySourcesBeforeOpeningTransaction(t *testing.T) {
sourceIDs := make([]int64, MaxTaskSourceCount+1)
for i := range sourceIDs {
sourceIDs[i] = int64(i + 1)
}
// No database handle is needed: validation must run before Begin so an
// oversized request cannot consume a connection or create partial rows.
_, err := (&DB{}).CreateTaskWithOptions("child", "goal", TaskCreateOptions{SourceTaskIDs: sourceIDs})
if err == nil || !strings.Contains(err.Error(), "too many source tasks") {
t.Fatalf("expected source-count validation error, got %v", err)
}
}
func TestNormalizeTaskCompanyIDs(t *testing.T) {
got, err := NormalizeTaskCompanyIDs([]int64{4, 2, 4, 7, 2})
if err != nil {
t.Fatal(err)
}
if want := []int64{4, 2, 7}; !reflect.DeepEqual(got, want) {
t.Fatalf("NormalizeTaskCompanyIDs=%v, want %v", got, want)
}
if _, err := NormalizeTaskCompanyIDs([]int64{1, 0}); !errors.Is(err, ErrTaskCompanyIDsInvalid) {
t.Fatalf("invalid company id error=%v", err)
}
}
func TestCreateTaskRejectsTooManyCompaniesBeforeOpeningTransaction(t *testing.T) {
companyIDs := make([]int64, MaxTaskCompanyCount+1)
for i := range companyIDs {
companyIDs[i] = int64(i + 1)
}
_, err := (&DB{}).CreateTaskWithOptions("child", "goal", TaskCreateOptions{CompanyIDs: companyIDs})
if !errors.Is(err, ErrTaskCompanyIDsInvalid) {
t.Fatalf("expected company-count validation error, got %v", err)
}
}
+281 -45
View File
@@ -1,7 +1,11 @@
package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"strconv"
"time"
)
@@ -9,15 +13,29 @@ import (
type Task struct {
ID int64 `json:"id"`
Name string `json:"name"` // 可选任务名称;空=未命名
CategoryID *int64 `json:"category_id,omitempty"`
CategoryName string `json:"category_name,omitempty"`
Description string `json:"description"`
Goal string `json:"goal"`
ExplorationID int64 `json:"exploration_id"`
Status string `json:"status"`
Paused bool `json:"paused"`
Queued bool `json:"queued"`
QueuedAt *time.Time `json:"queued_at,omitempty"`
QueueMode string `json:"queue_mode,omitempty"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层)
CreatedAt time.Time `json:"created_at"`
CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil
// Task-level ordered LLM chain. LLMProfileID remains the compatibility alias
// for ActiveLLMProfileID while older API clients still send one profile id.
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"`
LLMChainRevision int64 `json:"-"`
LLMFailoverState string `json:"llm_failover_state,omitempty"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
CompanyIDs []int64 `json:"company_ids,omitempty"`
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层)
CreatedAt time.Time `json:"created_at"`
CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil
// 任务级超时(见 docs/任务级超时与收尾设计.md)。
TimeoutSeconds int `json:"timeout_seconds"` // 0=不限时
FirstRunAt *time.Time `json:"first_run_at,omitempty"` // 首次真正开始运行的时刻(非 created_at);nil=尚未运行
@@ -32,6 +50,23 @@ type Task struct {
CoverageEnabled bool `json:"coverage_enabled"`
}
// TaskDeleteResult reports optional related-data cleanup performed in the same
// transaction as the task/exploration delete.
type TaskDeleteResult struct {
AssetsDeleted int64
AssetsDetached int64
FindingsDeleted int64
LLMRecordsDeleted int64
}
// TaskDeletePreparation is produced inside the PostgreSQL deletion transaction
// after asset and anchor writers have been excluded. Prepare callbacks may use
// TrafficHosts to stage an external traffic deletion before PostgreSQL commits.
type TaskDeletePreparation struct {
ExplorationID int64
TrafficHosts []string
}
// IsTerminal reports whether a task status is a terminal (finished) state.
// 单一真源,替换散落各处的 done/failed 硬编码判定。
func IsTerminal(status string) bool {
@@ -45,6 +80,46 @@ func IsTerminal(status string) bool {
// 低于它(含缺省 0 / 负值 / 误配的小值)一律抬到 10min,防止把 planner 打爆。
const MinPlanHeartbeatSeconds = 600
// MaxTaskSourceCount bounds the amount of live inherited context one task can
// pull into every planner/main-agent prompt. Inheritance is intentionally direct
// only; keeping the fan-in bounded also prevents a single create request from
// multiplying graph and asset-context queries without limit.
const MaxTaskSourceCount = 8
// MaxTaskCompanyCount bounds the number of company asset scopes attached to a
// task. Company scopes are prompt context, so an unbounded list would inflate
// every planner and worker call.
const MaxTaskCompanyCount = 32
var (
ErrTaskCompanyIDsInvalid = errors.New("invalid task company ids")
ErrTaskCompanyNotFound = errors.New("task company not found")
)
// NormalizeTaskCompanyIDs validates IDs and removes duplicates while retaining
// the user's first-seen order.
func NormalizeTaskCompanyIDs(ids []int64) ([]int64, error) {
if len(ids) == 0 {
return nil, nil
}
seen := make(map[int64]struct{}, min(len(ids), MaxTaskCompanyCount))
normalized := make([]int64, 0, min(len(ids), MaxTaskCompanyCount))
for _, id := range ids {
if id <= 0 {
return nil, fmt.Errorf("%w: company id must be positive", ErrTaskCompanyIDsInvalid)
}
if _, exists := seen[id]; exists {
continue
}
seen[id] = struct{}{}
normalized = append(normalized, id)
if len(normalized) > MaxTaskCompanyCount {
return nil, fmt.Errorf("%w: got more than %d unique companies", ErrTaskCompanyIDsInvalid, MaxTaskCompanyCount)
}
}
return normalized, nil
}
func normalizeHeartbeat(sec int) int {
if sec < MinPlanHeartbeatSeconds {
return MinPlanHeartbeatSeconds
@@ -53,8 +128,6 @@ func normalizeHeartbeat(sec int) int {
}
func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) {
<<<<<<< Updated upstream
=======
var ids []int64
if llmProfileID != nil {
ids = []int64{*llmProfileID}
@@ -68,6 +141,7 @@ func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSe
// with the task/exploration row.
type TaskCreateOptions struct {
Name string // 可选任务名称;空=未命名
CategoryID *int64
SourceTaskIDs []int64
CompanyIDs []int64
LLMProfileIDs []int64
@@ -89,7 +163,6 @@ func (d *DB) CreateTaskWithOptions(description, goal string, opts TaskCreateOpti
return nil, err
}
opts.CompanyIDs = companyIDs
>>>>>>> Stashed changes
tx, err := d.Begin()
if err != nil {
return nil, err
@@ -114,22 +187,30 @@ INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, or
VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); err != nil {
return nil, err
}
<<<<<<< Updated upstream
if timeoutSeconds < 0 {
timeoutSeconds = 0
=======
if opts.TimeoutSeconds < 0 {
opts.TimeoutSeconds = 0
}
opts.PlanHeartbeatSeconds = normalizeHeartbeat(opts.PlanHeartbeatSeconds)
coverageEnabled := opts.CoverageEnabled == nil || *opts.CoverageEnabled
var categoryName string
if opts.CategoryID != nil {
if *opts.CategoryID <= 0 {
return nil, fmt.Errorf("%w: category id must be positive", ErrTaskCategoryInvalid)
}
if err := tx.QueryRow(`SELECT name FROM task_categories WHERE id=$1`, *opts.CategoryID).Scan(&categoryName); err != nil {
if err == sql.ErrNoRows {
return nil, ErrTaskCategoryNotFound
}
return nil, err
}
}
var active *int64
if len(opts.LLMProfileIDs) > 0 {
id := opts.LLMProfileIDs[0]
active = &id
}
t := &Task{
Name: opts.Name,
Name: opts.Name, CategoryID: opts.CategoryID, CategoryName: categoryName,
Description: description, Goal: goal, ExplorationID: expID,
LLMProfileID: active, ActiveLLMProfileID: active,
LLMProfileIDs: append([]int64(nil), opts.LLMProfileIDs...),
@@ -137,31 +218,30 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload));
CompanyIDs: append([]int64(nil), opts.CompanyIDs...),
TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds,
CoverageEnabled: coverageEnabled,
>>>>>>> Stashed changes
}
planHeartbeatSeconds = normalizeHeartbeat(planHeartbeatSeconds) // <=0→300,<30→30
t := &Task{Description: description, Goal: goal, ExplorationID: expID, LLMProfileID: llmProfileID, TimeoutSeconds: timeoutSeconds, PlanHeartbeatSeconds: planHeartbeatSeconds}
if err := tx.QueryRow(`
<<<<<<< Updated upstream
INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, timeout_seconds, plan_heartbeat_seconds) VALUES ($1,$2,$3,$4,$5,$6)
RETURNING id, status, paused, created_at`, description, goal, expID, llmProfileID, timeoutSeconds, planHeartbeatSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
=======
INSERT INTO tasks(name, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled)
VALUES ($1,$2,$3,$4,$5,$5,$6,$7,$8)
RETURNING id, status, paused, created_at`, opts.Name, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, coverageEnabled).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
>>>>>>> Stashed changes
INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled)
VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9)
RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, coverageEnabled).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
return nil, err
}
if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil {
return nil, err
}
if err := insertTaskCompanies(tx, t.ID, opts.CompanyIDs); err != nil {
return nil, err
}
if err := insertTaskLLMProfiles(tx, t.ID, opts.LLMProfileIDs); err != nil {
return nil, err
}
if len(opts.LLMProfileIDs) == 0 {
t.LLMFailoverState = "default"
} else {
t.LLMFailoverState = "ready"
}
return t, tx.Commit()
}
<<<<<<< Updated upstream
const taskCols = `id, description, goal, exploration_id, status, paused, llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at`
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
var t Task
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.LLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
=======
func insertTaskCompanies(tx *sql.Tx, taskID int64, companyIDs []int64) error {
if len(companyIDs) == 0 {
return nil
@@ -222,12 +302,13 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error {
return nil
}
const taskCols = `id, COALESCE(name,''), description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(coverage_enabled,true), first_run_at, deadline_at`
const taskCols = `id, COALESCE(name,''), category_id,
COALESCE((SELECT category.name FROM task_categories category WHERE category.id=tasks.category_id),''),
description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(coverage_enabled,true), first_run_at, deadline_at`
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
var t Task
if err := sc.Scan(&t.ID, &t.Name, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.CoverageEnabled, &t.FirstRunAt, &t.DeadlineAt); err != nil {
>>>>>>> Stashed changes
if err := sc.Scan(&t.ID, &t.Name, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.CoverageEnabled, &t.FirstRunAt, &t.DeadlineAt); err != nil {
return nil, err
}
return &t, nil
@@ -254,7 +335,16 @@ func (d *DB) ListTasks() ([]*Task, error) {
}
out = append(out, t)
}
return out, rows.Err()
if err := rows.Err(); err != nil {
return nil, err
}
rows.Close()
for _, task := range out {
if err := d.hydrateTaskContext(task); err != nil {
return nil, err
}
}
return out, nil
}
// GetTask returns one alive task (nil if not found/deleted).
@@ -266,6 +356,9 @@ func (d *DB) GetTask(id int64) (*Task, error) {
}
return nil, err
}
if err := d.hydrateTaskContext(t); err != nil {
return nil, err
}
return t, nil
}
@@ -275,6 +368,43 @@ func (d *DB) SetPaused(id int64, paused bool) error {
return err
}
// Enqueue places a task at the tail of the persistent admission queue. Repeating
// the operation while it is already queued keeps its original FIFO position.
func (d *DB) Enqueue(id int64, mode string) error {
if mode != "bootstrap" && mode != "resume" {
return fmt.Errorf("invalid queue mode %q", mode)
}
_, err := d.Exec(`UPDATE tasks
SET queued=true,
queued_at=CASE WHEN queued THEN COALESCE(queued_at, now()) ELSE now() END,
queue_mode=CASE
WHEN queue_mode='bootstrap' OR $2='bootstrap' THEN 'bootstrap'
ELSE 'resume'
END
WHERE id=$1`, id, mode)
return err
}
// Dequeue removes the concurrency hold. clearMode is false when a user pauses a
// queued task, preserving whether its next admission must bootstrap or resume.
func (d *DB) Dequeue(id int64, clearMode bool) error {
_, err := d.Exec(`UPDATE tasks
SET queued=false,
queued_at=NULL,
queue_mode=CASE WHEN $2 THEN '' ELSE queue_mode END
WHERE id=$1`, id, clearMode)
return err
}
// SetQueued is the compatibility helper used by older callers and tests. New
// scheduling code should use Enqueue/Dequeue so FIFO metadata is explicit.
func (d *DB) SetQueued(id int64, queued bool) error {
if queued {
return d.Enqueue(id, "bootstrap")
}
return d.Dequeue(id, true)
}
// SetStatus updates a task's lifecycle status. Entering a terminal state
// (done/failed/timeout) stamps completed_at once (COALESCE keeps the first stamp
// stable); moving back to a non-terminal state clears it, so a re-run has no stale
@@ -323,28 +453,134 @@ UPDATE tasks
return deadline, err
}
// DeleteTask hard-deletes a task and its exploration subgraph (nodes/edges/
// activity via ON DELETE CASCADE). The task row is removed first so the
// ON DELETE RESTRICT on tasks.exploration_id does not block the exploration delete.
// Global assets are NOT touched.
// DeleteTask preserves the historical behavior: remove the task and its
// exploration subgraph while retaining global assets.
func (d *DB) DeleteTask(id int64) error {
_, err := d.DeleteTaskCascade(id, false, false)
return err
}
// DeleteTaskCascade hard-deletes a task and its exploration subgraph
// (nodes/edges/activity via ON DELETE CASCADE). Optional standalone findings and
// assets owned only by this task are deleted in the same transaction; shared
// assets are retained and only have this task id detached. deleteLLMRecords is
// optional to preserve callers of the original two-option API.
func (d *DB) DeleteTaskCascade(id int64, deleteAssets, deleteFindings bool, deleteLLMRecords ...bool) (TaskDeleteResult, error) {
deleteRecords := len(deleteLLMRecords) > 0 && deleteLLMRecords[0]
return d.DeleteTaskCascadePrepared(id, deleteAssets, deleteFindings, deleteRecords, nil)
}
// DeleteTaskCascadePrepared coordinates reversible external deletion with the
// PostgreSQL cascade. When prepare is non-nil, host ownership is resolved and
// prepare is invoked inside this transaction while asset and anchor writers are
// excluded. The locks remain held through commit, closing the window where a
// host could become shared after its traffic had already been staged.
//
// prepare must only stage reversible work. Any returned error or later database
// error rolls PostgreSQL back; the caller remains responsible for rolling back
// external work that its callback staged successfully.
func (d *DB) DeleteTaskCascadePrepared(
id int64,
deleteAssets, deleteFindings, deleteLLMRecords bool,
prepare func(TaskDeletePreparation) error,
) (TaskDeleteResult, error) {
var result TaskDeleteResult
tx, err := d.Begin()
if err != nil {
return err
return result, err
}
defer tx.Rollback()
var expID int64
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, id).Scan(&expID); err != nil {
if err.Error() == "sql: no rows in result set" {
return nil // already gone
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1 FOR UPDATE`, id).Scan(&expID); err != nil {
if err == sql.ErrNoRows {
return result, nil // already gone
}
return err
return result, err
}
// SHARE ROW EXCLUSIVE conflicts with every INSERT/UPDATE/DELETE on these
// tables and with another coordinated deletion. Taking both in one fixed order
// prevents phantoms (a distinct asset row for the same host) as well as new
// ownership/anchor references until the deletion transaction commits.
if deleteAssets || prepare != nil {
if _, err := tx.Exec(`LOCK TABLE assets, exploration_anchors IN SHARE ROW EXCLUSIVE MODE`); err != nil {
return result, err
}
}
if prepare != nil {
hosts, err := hostsForTaskDeletion(tx, id, expID)
if err != nil {
return result, err
}
if err := prepare(TaskDeletePreparation{
ExplorationID: expID,
TrafficHosts: hosts,
}); err != nil {
return result, err
}
}
if deleteAssets {
// Ownership is the union of explicit task_ids and this exploration's anchors.
// An asset is deletable only when no other live task references it through
// either mechanism. This covers legacy anchor-only seeds without destroying
// evidence anchored by another task.
res, err := tx.Exec(`
WITH candidate_assets AS (
SELECT id FROM assets WHERE $1 = ANY(task_ids)
UNION
SELECT ea.asset_id
FROM exploration_anchors ea
JOIN exploration_nodes n ON n.id=ea.node_id
WHERE n.exploration_id=$2
),
deletable AS (
SELECT a.id
FROM assets a
JOIN candidate_assets c ON c.id=a.id
WHERE NOT EXISTS (
SELECT 1 FROM tasks t
WHERE t.id<>$1 AND t.deleted_at IS NULL AND t.id=ANY(a.task_ids)
) AND NOT EXISTS (
SELECT 1
FROM exploration_anchors ea
JOIN exploration_nodes n ON n.id=ea.node_id
JOIN tasks t ON t.exploration_id=n.exploration_id
WHERE ea.asset_id=a.id AND t.id<>$1 AND t.deleted_at IS NULL
)
)
DELETE FROM assets a USING deletable d WHERE a.id=d.id`, id, expID)
if err != nil {
return result, err
}
result.AssetsDeleted, _ = res.RowsAffected()
res, err = tx.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, id)
if err != nil {
return result, err
}
result.AssetsDetached, _ = res.RowsAffected()
}
if deleteFindings {
res, err := tx.Exec(`DELETE FROM findings WHERE task_id = $1`, id)
if err != nil {
return result, err
}
result.FindingsDeleted, _ = res.RowsAffected()
}
if deleteLLMRecords {
res, err := tx.Exec(`DELETE FROM llm_records WHERE COALESCE(task_id,'')=$1`, strconv.FormatInt(id, 10))
if err != nil {
return result, err
}
result.LLMRecordsDeleted, _ = res.RowsAffected()
}
// llm_usage (the token metering ledger) is intentionally NOT deleted with the
// task — it is kept as historical accounting even after the task is gone.
if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1`, id); err != nil {
return err
return result, err
}
if _, err := tx.Exec(`DELETE FROM explorations WHERE id=$1`, expID); err != nil {
return err
return result, err
}
return tx.Commit()
return result, tx.Commit()
}
+506 -1
View File
@@ -1,6 +1,12 @@
package db
import "testing"
import (
"errors"
"fmt"
"sync"
"testing"
"time"
)
func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
d, err := Open(testDSN(t))
@@ -27,6 +33,23 @@ func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
if err != nil || got == nil || !got.Paused {
t.Fatalf("paused not persisted: %+v err=%v", got, err)
}
if got.Queued {
t.Fatalf("new task should not be queued: %+v", got)
}
// queued (concurrency-hold) flag round-trips independently of paused
if err := d.SetQueued(tk.ID, true); err != nil {
t.Fatal(err)
}
if g, _ := d.GetTask(tk.ID); g == nil || !g.Queued {
t.Fatalf("queued not persisted: %+v", g)
}
if err := d.SetQueued(tk.ID, false); err != nil {
t.Fatal(err)
}
if g, _ := d.GetTask(tk.ID); g == nil || g.Queued {
t.Fatalf("queued not cleared: %+v", g)
}
// list contains it
list, _ := d.ListTasks()
@@ -58,3 +81,485 @@ func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
t.Fatalf("exploration should be deleted, got %d", exps)
}
}
func TestTaskDeleteCascadeAssets(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
first, err := d.CreateTask("级联删除测试", "目标A", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
second, err := d.CreateTask("共享资产保留测试", "目标B", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
assets := d.Assets()
exclusiveID, err := assets.UpsertRootDomain(UpsertRootDomainReq{
Domain: fmt.Sprintf("delete-%d.example.test", first.ID),
TaskID: first.ID,
})
if err != nil {
t.Fatal(err)
}
sharedDomain := fmt.Sprintf("shared-%d.example.test", first.ID)
sharedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sharedDomain, TaskID: first.ID})
if err != nil {
t.Fatal(err)
}
if _, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sharedDomain, TaskID: second.ID}); err != nil {
t.Fatal(err)
}
anchorOnlyDomain := fmt.Sprintf("anchor-only-%d.example.test", first.ID)
anchorOnlyID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: anchorOnlyDomain})
if err != nil {
t.Fatal(err)
}
firstOrigin, err := d.Exploration(first.ExplorationID).OriginFactID()
if err != nil || firstOrigin == 0 {
t.Fatalf("first origin: id=%d err=%v", firstOrigin, err)
}
if err := d.Exploration(first.ExplorationID).Anchor(firstOrigin, anchorOnlyID); err != nil {
t.Fatal(err)
}
protectedDomain := fmt.Sprintf("other-anchor-%d.example.test", first.ID)
protectedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: protectedDomain, TaskID: first.ID})
if err != nil {
t.Fatal(err)
}
secondOrigin, err := d.Exploration(second.ExplorationID).OriginFactID()
if err != nil || secondOrigin == 0 {
t.Fatalf("second origin: id=%d err=%v", secondOrigin, err)
}
if err := d.Exploration(second.ExplorationID).Anchor(secondOrigin, protectedID); err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteTask(first.ID)
_ = d.DeleteTask(second.ID)
_, _ = assets.DeleteByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID})
})
hosts, err := assets.HostsByTask(first.ID)
if err != nil {
t.Fatal(err)
}
hostSet := make(map[string]bool, len(hosts))
for _, host := range hosts {
hostSet[host] = true
}
if !hostSet[fmt.Sprintf("delete-%d.example.test", first.ID)] || !hostSet[sharedDomain] {
t.Fatalf("task hosts missing cascade fixtures: %v", hosts)
}
deletableHosts, err := assets.HostsForTaskDeletion(first.ID, first.ExplorationID)
if err != nil {
t.Fatal(err)
}
deletableSet := make(map[string]bool, len(deletableHosts))
for _, host := range deletableHosts {
deletableSet[host] = true
}
if !deletableSet[fmt.Sprintf("delete-%d.example.test", first.ID)] || !deletableSet[anchorOnlyDomain] {
t.Fatalf("exclusive or anchor-only cleanup host missing: %v", deletableHosts)
}
if deletableSet[sharedDomain] || deletableSet[protectedDomain] {
t.Fatalf("shared host was not protected from traffic cleanup: %v", deletableHosts)
}
findingID, err := d.AddFinding(first.ID, 0, "__task_delete_cascade__", "", SeverityHigh, "summary", "evidence", "tester", []int64{exclusiveID})
if err != nil {
t.Fatal(err)
}
result, err := d.DeleteTaskCascade(first.ID, true, true)
if err != nil {
t.Fatal(err)
}
if result.AssetsDeleted < 2 || result.AssetsDetached < 2 {
t.Fatalf("unexpected asset cleanup result: %+v", result)
}
if result.FindingsDeleted != 1 {
t.Fatalf("unexpected finding cleanup result: %+v", result)
}
if finding, err := d.GetFinding(findingID); err != nil || finding != nil {
t.Fatalf("finding should be deleted, got finding=%+v err=%v", finding, err)
}
remaining, err := assets.GetByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID})
if err != nil {
t.Fatal(err)
}
if len(remaining) != 2 {
t.Fatalf("only shared and other-anchored assets should remain, got %+v", remaining)
}
byID := make(map[int64]*Asset, len(remaining))
for _, asset := range remaining {
byID[asset.ID] = asset
}
if shared := byID[sharedID]; shared == nil || len(shared.TaskIDs) != 1 || shared.TaskIDs[0] != second.ID {
t.Fatalf("deleted task should be detached from shared asset: %+v", shared)
}
if protected := byID[protectedID]; protected == nil || len(protected.TaskIDs) != 0 {
t.Fatalf("other task's anchor should preserve the asset after detaching ownership: %+v", protected)
}
}
func TestTaskRelationsAndLLMFailoverChain(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
profileIDs := make([]int64, 0, 3)
for i := 0; i < 3; i++ {
id, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("task-chain-%d-%d", suffix, i), Format: "openai",
Model: fmt.Sprintf("model-%d", i), APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
profileIDs = append(profileIDs, id)
}
sourceA, err := d.CreateTask("source A", "goal A", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
sourceB, err := d.CreateTask("source B", "goal B", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
child, err := d.CreateTaskWithOptions("child", "new goal", TaskCreateOptions{
SourceTaskIDs: []int64{sourceA.ID, sourceB.ID},
LLMProfileIDs: profileIDs,
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteTask(child.ID)
_ = d.DeleteTask(sourceA.ID)
_ = d.DeleteTask(sourceB.ID)
for _, id := range profileIDs {
_ = d.DeleteProfile(id)
}
})
got, err := d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(got.SourceTaskIDs) != fmt.Sprint([]int64{sourceA.ID, sourceB.ID}) {
t.Fatalf("unexpected direct sources: %v", got.SourceTaskIDs)
}
if fmt.Sprint(got.LLMProfileIDs) != fmt.Sprint(profileIDs) || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] {
t.Fatalf("unexpected initial chain: %+v", got)
}
initialRevision := got.LLMChainRevision
if err := d.ReplaceTaskLLMProfiles(child.ID, profileIDs, profileIDs[0]); err != nil {
t.Fatal(err)
}
staleRevision, err := d.MarkTaskLLMProfileQuotaExhaustedAtRevision(child.ID, profileIDs[0], initialRevision, "late quota from prior generation")
if err != nil {
t.Fatal(err)
}
if !staleRevision.Stale || staleRevision.Advanced || staleRevision.ChainExhausted {
t.Fatalf("old chain revision changed replacement chain: %+v", staleRevision)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.LLMChainRevision <= initialRevision || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] {
t.Fatalf("replacement revision/cursor not preserved: %+v", got)
}
replacementRevision := got.LLMChainRevision
transition, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "insufficient_quota")
if err != nil {
t.Fatal(err)
}
if !transition.Advanced || transition.ChainExhausted || transition.NextProfileID == nil || *transition.NextProfileID != profileIDs[1] {
t.Fatalf("unexpected first transition: %+v", transition)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.LLMChainRevision != replacementRevision+1 {
t.Fatalf("automatic cursor advance did not increment revision: before=%d after=%d", replacementRevision, got.LLMChainRevision)
}
late, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "late duplicate")
if err != nil {
t.Fatal(err)
}
if late.Advanced || late.NextProfileID == nil || *late.NextProfileID != profileIDs[1] {
t.Fatalf("late failure advanced past current profile: %+v", late)
}
if _, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[1], "quota_exceeded"); err != nil {
t.Fatal(err)
}
last, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[2], "余额不足")
if err != nil {
t.Fatal(err)
}
if !last.Advanced || !last.ChainExhausted || last.NextProfileID != nil {
t.Fatalf("unexpected exhausted transition: %+v", last)
}
duplicateLast, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[2], "late final duplicate")
if err != nil {
t.Fatal(err)
}
if duplicateLast.Advanced || !duplicateLast.ChainExhausted || duplicateLast.NextProfileID != nil {
t.Fatalf("duplicate final failure must be idempotent: %+v", duplicateLast)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.ActiveLLMProfileID != nil || got.LLMFailoverState != "chain_exhausted" {
t.Fatalf("chain exhaustion was not persisted: %+v", got)
}
// Two requests can observe the same last active profile. Only the transaction
// that clears the cursor may report an advance; the late one is idempotent.
if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2]}, profileIDs[2]); err != nil {
t.Fatal(err)
}
transitions := make(chan TaskLLMTransition, 2)
errs := make(chan error, 2)
var wg sync.WaitGroup
for i := 0; i < 2; i++ {
wg.Add(1)
go func() {
defer wg.Done()
transition, markErr := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[2], "concurrent final quota")
if markErr != nil {
errs <- markErr
return
}
transitions <- transition
}()
}
wg.Wait()
close(errs)
close(transitions)
for markErr := range errs {
t.Fatal(markErr)
}
advanced := 0
for transition := range transitions {
if transition.Advanced {
advanced++
}
if !transition.ChainExhausted {
t.Fatalf("concurrent final transition must report exhausted chain: %+v", transition)
}
}
if advanced != 1 {
t.Fatalf("last profile advanced %d times, want exactly once", advanced)
}
// Starting manually from the middle consumes only candidates after that cursor.
// Earlier ready profiles must not be revived by hydration or unrelated deletes.
if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2], profileIDs[0], profileIDs[1]}, profileIDs[0]); err != nil {
t.Fatal(err)
}
if step, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "middle quota"); err != nil || step.NextProfileID == nil || *step.NextProfileID != profileIDs[1] {
t.Fatalf("middle cursor did not advance to its successor: transition=%+v err=%v", step, err)
}
if end, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[1], "tail quota"); err != nil || !end.ChainExhausted {
t.Fatalf("tail did not exhaust manual chain: transition=%+v err=%v", end, err)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.ActiveLLMProfileID != nil || got.LLMFailoverState != "chain_exhausted" {
t.Fatalf("hydration revived a profile before the manual cursor: %+v", got)
}
unrelatedID, err := d.SaveProfile(&LLMProfile{
Name: fmt.Sprintf("task-chain-unrelated-%d", suffix), Format: "openai", Model: "other", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
if err := d.DeleteProfile(unrelatedID); err != nil {
t.Fatal(err)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.ActiveLLMProfileID != nil || got.LLMFailoverState != "chain_exhausted" {
t.Fatalf("deleting an unrelated profile revived an exhausted chain: %+v", got)
}
// A call that finishes after its profile was removed from the chain is stale:
// preserve the new cursor and surface the original provider error upstream.
if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2], profileIDs[1]}, profileIDs[2]); err != nil {
t.Fatal(err)
}
stale, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "obsolete in-flight quota")
if err != nil {
t.Fatalf("removed in-flight profile returned an internal error: %v", err)
}
if stale.Advanced || stale.ChainExhausted || stale.NextProfileID != nil {
t.Fatalf("removed in-flight profile changed the replacement chain: %+v", stale)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[2] {
t.Fatalf("stale failure changed active replacement profile: %+v", got)
}
if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2], profileIDs[0], profileIDs[1]}, profileIDs[0]); err != nil {
t.Fatal(err)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] || got.LLMFailoverState != "ready" || got.LLMFailoverReason != "" {
t.Fatalf("chain reset did not clear failure state: %+v", got)
}
if err := d.DeleteProfile(profileIDs[0]); err != nil {
t.Fatal(err)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[1] {
t.Fatalf("deleting active profile did not select next ready entry: %+v", got)
}
if err := d.DeleteProfile(profileIDs[1]); err != nil {
t.Fatal(err)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.ActiveLLMProfileID != nil || len(got.LLMProfileIDs) != 0 || got.LLMFailoverState != "default" {
t.Fatalf("deleting active chain tail must fall back instead of wrapping: %+v", got)
}
if err := d.DeleteTask(sourceA.ID); err != nil {
t.Fatal(err)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(got.SourceTaskIDs) != fmt.Sprint([]int64{sourceB.ID}) {
t.Fatalf("source delete did not cascade only its relation: %v", got.SourceTaskIDs)
}
}
func TestTaskContextRejectsDuplicatesAndTerminalEdits(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
source, err := d.CreateTask("source", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(source.ID) })
if _, err := d.CreateTaskWithOptions("bad", "goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID, source.ID}}); err == nil {
t.Fatal("duplicate source task ids should be rejected")
}
task, err := d.CreateTask("terminal", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
if err := d.SetStatus(task.ID, "done"); err != nil {
t.Fatal(err)
}
if err := d.ReplaceTaskLLMProfiles(task.ID, nil, 0); err == nil {
t.Fatal("terminal task LLM edit should be rejected")
}
}
func TestCreateTaskWithCompanyScopes(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
companyA, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company A %d", suffix), "")
if err != nil {
t.Fatal(err)
}
companyB, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company B %d", suffix), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, companyA, companyB)
})
task, err := d.CreateTaskWithOptions("company-scoped task", "use company scope", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyB},
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
if fmt.Sprint(task.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("creation result company IDs=%v", task.CompanyIDs)
}
got, err := d.GetTask(task.ID)
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(got.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("hydrated company IDs=%v", got.CompanyIDs)
}
var scopeCount int
if err := d.QueryRow(`SELECT count(*) FROM task_scope WHERE task_id=$1 AND kind='company'`, task.ID).Scan(&scopeCount); err != nil {
t.Fatal(err)
}
if scopeCount != 2 {
t.Fatalf("company task scope rows=%d want 2", scopeCount)
}
duplicateTask, err := d.CreateTaskWithOptions("duplicate company scope", "deduplicate", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyA, companyB, companyA},
})
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(duplicateTask.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("company IDs were not normalized: %v", duplicateTask.CompanyIDs)
}
if err := d.DeleteTask(duplicateTask.ID); err != nil {
t.Fatal(err)
}
badDescription := fmt.Sprintf("invalid-company-%d", suffix)
if _, err := d.CreateTaskWithOptions(badDescription, "rollback", TaskCreateOptions{CompanyIDs: []int64{companyA, 1 << 62}}); !errors.Is(err, ErrTaskCompanyNotFound) {
t.Fatalf("missing company error=%v, want %v", err, ErrTaskCompanyNotFound)
}
var leaked int
if err := d.QueryRow(`SELECT count(*) FROM explorations WHERE description=$1`, badDescription).Scan(&leaked); err != nil {
t.Fatal(err)
}
if leaked != 0 {
t.Fatalf("failed company association leaked %d exploration rows", leaked)
}
}