feat(tasks): add runtime controls and inherited context
This commit is contained in:
1 parent
18926c129f
commit
4392b4fee1
50 files changed
+8291
-663
No files matched your search
@@ -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
@@ -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).
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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 = ©
|
||||
}
|
||||
}
|
||||
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, ¤tRevision); 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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user