mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-15 20:50:54 +08:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c2d1f557c3 | ||
|
|
e636408f66 |
+50
-26
@@ -1486,11 +1486,13 @@ type TaskListMetrics struct {
|
||||
Findings FindingSeverityCounts // this task's vulnerability count in the findings table (bucketed by severity)
|
||||
}
|
||||
|
||||
// TaskListMetricsAll returns list aggregates for every live task in one query.
|
||||
// The lateral lookups use the per-exploration indexes instead of grouping the
|
||||
// complete activity history on every poll.
|
||||
func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
rows, err := d.Query(`
|
||||
// taskListMetricsQuery is the select shared by TaskListMetricsAll and
|
||||
// TaskListMetrics; each appends its own WHERE on task. The lateral lookups use
|
||||
// the per-exploration indexes instead of grouping the complete activity history
|
||||
// on every poll. latest_activity orders by the column activity.created_at, so
|
||||
// idx_act_latest (exploration_id, created_at DESC) serves the LIMIT 1; ordering
|
||||
// by the output alias would sort every activity row of the task.
|
||||
const taskListMetricsQuery = `
|
||||
SELECT task.exploration_id,
|
||||
COALESCE(token_metrics.input_tokens,0),
|
||||
COALESCE(token_metrics.output_tokens,0),
|
||||
@@ -1514,10 +1516,10 @@ func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
WHERE exploration_id=task.exploration_id AND kind='result'
|
||||
) token_metrics ON true
|
||||
LEFT JOIN LATERAL (
|
||||
SELECT EXTRACT(EPOCH FROM created_at)::bigint AS created_at
|
||||
SELECT EXTRACT(EPOCH FROM activity.created_at)::bigint AS created_at
|
||||
FROM activity
|
||||
WHERE exploration_id=task.exploration_id
|
||||
ORDER BY created_at DESC
|
||||
WHERE activity.exploration_id=task.exploration_id
|
||||
ORDER BY activity.created_at DESC
|
||||
LIMIT 1
|
||||
) latest_activity ON true
|
||||
LEFT JOIN LATERAL (
|
||||
@@ -1539,30 +1541,41 @@ func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
FROM findings
|
||||
WHERE task_id=task.id
|
||||
) finding_metrics ON true
|
||||
WHERE task.deleted_at IS NULL`)
|
||||
WHERE task.deleted_at IS NULL`
|
||||
|
||||
// scanTaskListMetrics reads one row of taskListMetricsQuery.
|
||||
func scanTaskListMetrics(row interface{ Scan(...any) error }) (int64, TaskListMetrics, error) {
|
||||
var explorationID int64
|
||||
var metrics TaskListMetrics
|
||||
err := row.Scan(
|
||||
&explorationID,
|
||||
&metrics.Tokens.InputTokens,
|
||||
&metrics.Tokens.OutputTokens,
|
||||
&metrics.Tokens.CacheReadTokens,
|
||||
&metrics.Tokens.CacheWriteTokens,
|
||||
&metrics.LastActivity,
|
||||
&metrics.Goals.Total,
|
||||
&metrics.Goals.Met,
|
||||
&metrics.RunningIntents,
|
||||
&metrics.Findings.Critical,
|
||||
&metrics.Findings.High,
|
||||
&metrics.Findings.Medium,
|
||||
&metrics.Findings.Low,
|
||||
)
|
||||
return explorationID, metrics, err
|
||||
}
|
||||
|
||||
// TaskListMetricsAll returns list aggregates for every live task in one query.
|
||||
func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
rows, err := d.Query(taskListMetricsQuery)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := map[int64]TaskListMetrics{}
|
||||
for rows.Next() {
|
||||
var explorationID int64
|
||||
var metrics TaskListMetrics
|
||||
if err := rows.Scan(
|
||||
&explorationID,
|
||||
&metrics.Tokens.InputTokens,
|
||||
&metrics.Tokens.OutputTokens,
|
||||
&metrics.Tokens.CacheReadTokens,
|
||||
&metrics.Tokens.CacheWriteTokens,
|
||||
&metrics.LastActivity,
|
||||
&metrics.Goals.Total,
|
||||
&metrics.Goals.Met,
|
||||
&metrics.RunningIntents,
|
||||
&metrics.Findings.Critical,
|
||||
&metrics.Findings.High,
|
||||
&metrics.Findings.Medium,
|
||||
&metrics.Findings.Low,
|
||||
); err != nil {
|
||||
explorationID, metrics, err := scanTaskListMetrics(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[explorationID] = metrics
|
||||
@@ -1570,6 +1583,17 @@ func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// TaskListMetrics returns the same aggregates for the one live task of
|
||||
// exploration expID, with the same laterals as TaskListMetricsAll. A missing or
|
||||
// deleted task gives zero metrics, as a missing key of TaskListMetricsAll does.
|
||||
func (d *DB) TaskListMetrics(expID int64) (TaskListMetrics, error) {
|
||||
_, metrics, err := scanTaskListMetrics(d.QueryRow(taskListMetricsQuery+` AND task.exploration_id=$1`, expID))
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return TaskListMetrics{}, nil
|
||||
}
|
||||
return metrics, err
|
||||
}
|
||||
|
||||
// GoalCountsAll returns goal totals (total / met) per exploration id, one query
|
||||
// for all tasks — used by listTasks so the task list shows progress for every task.
|
||||
func (d *DB) GoalCountsAll() (map[int64]GoalCounts, error) {
|
||||
|
||||
@@ -117,3 +117,133 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExplorationID); er
|
||||
t.Fatalf("unexpected goal metrics: %+v", metrics.Goals)
|
||||
}
|
||||
}
|
||||
|
||||
// TaskListMetrics(expID) is exactly the row TaskListMetricsAll returns for that
|
||||
// exploration, on every counter, and TaskListMetricsAll still lists every task.
|
||||
func TestTaskListMetricsMatchesAll(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
newest := time.Now().Add(-time.Minute).Truncate(time.Second)
|
||||
seed := func(name string, n int) *Task {
|
||||
t.Helper()
|
||||
task, err := d.CreateTask(name, "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, err := d.DeleteTaskCascade(task.ID, false, true)
|
||||
cleanupErr(t, "task "+name, err)
|
||||
})
|
||||
exp := task.ExplorationID
|
||||
// n result rows with tokens; the newest activity first, then an older row
|
||||
// with a higher id, so ordering by id instead of created_at would fail.
|
||||
for _, q := range []struct {
|
||||
sql string
|
||||
args []any
|
||||
}{
|
||||
{`INSERT INTO activity(exploration_id,kind,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,created_at)
|
||||
SELECT $1,'result',10,5,3,1,$3::timestamptz - interval '1 hour' FROM generate_series(1,$2::int)`, []any{exp, n, newest}},
|
||||
{`INSERT INTO activity(exploration_id,kind,created_at) VALUES ($1,'text',$2::timestamptz)`, []any{exp, newest}},
|
||||
{`INSERT INTO activity(exploration_id,kind,created_at) SELECT $1,'text',$3::timestamptz - interval '2 hours' FROM generate_series(1,$2::int)`, []any{exp, n, newest}},
|
||||
{`INSERT INTO exploration_nodes(exploration_id,kind,payload,state)
|
||||
SELECT $1,'goal','{}',CASE WHEN g<=$2::int THEN 'met' ELSE 'open' END FROM generate_series(1,$2::int+1) g`, []any{exp, n}},
|
||||
{`INSERT INTO exploration_nodes(exploration_id,kind,payload,state)
|
||||
SELECT $1,'intent','{}',CASE WHEN g<=$2::int THEN 'running' ELSE 'done' END FROM generate_series(1,$2::int+2) g`, []any{exp, n}},
|
||||
} {
|
||||
if _, err := d.Exec(q.sql, q.args...); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if _, err := d.Exec(`INSERT INTO findings(task_id,severity)
|
||||
SELECT $1, s FROM unnest(ARRAY['critical','high','high','medium','low','info']) s, generate_series(1,$2)`, task.ID, n); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return task
|
||||
}
|
||||
a := seed("task-list-metrics-a", 2)
|
||||
b := seed("task-list-metrics-b", 3)
|
||||
|
||||
all, err := d.TaskListMetricsAll()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, c := range []struct {
|
||||
task *Task
|
||||
n int
|
||||
}{{a, 2}, {b, 3}} {
|
||||
row, ok := all[c.task.ExplorationID]
|
||||
if !ok {
|
||||
t.Fatalf("TaskListMetricsAll misses task %d", c.task.ID)
|
||||
}
|
||||
one, err := d.TaskListMetrics(c.task.ExplorationID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if one != row {
|
||||
t.Fatalf("task %d: TaskListMetrics=%+v, TaskListMetricsAll=%+v", c.task.ID, one, row)
|
||||
}
|
||||
n := c.n
|
||||
want := TaskListMetrics{
|
||||
Tokens: TokenUsage{InputTokens: 10 * n, OutputTokens: 5 * n, CacheReadTokens: 3 * n, CacheWriteTokens: n},
|
||||
LastActivity: newest.Unix(),
|
||||
Goals: GoalCounts{Total: n + 1, Met: n},
|
||||
RunningIntents: n,
|
||||
Findings: FindingSeverityCounts{Critical: n, High: 2 * n, Medium: n, Low: n},
|
||||
}
|
||||
if one != want {
|
||||
t.Fatalf("task %d metrics=%+v, want %+v", c.task.ID, one, want)
|
||||
}
|
||||
}
|
||||
if none, err := d.TaskListMetrics(-1); err != nil || none != (TaskListMetrics{}) {
|
||||
t.Fatalf("missing exploration: %+v, %v; want zero metrics", none, err)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkTaskListMetrics compares the per-exploration metrics query with the
|
||||
// all-tasks one on three tasks of 40,000 activity rows each. No time assertion.
|
||||
func BenchmarkTaskListMetrics(b *testing.B) {
|
||||
dsn, _, err := DSN()
|
||||
if err != nil {
|
||||
b.Skipf("no database config (%v) - skipping", err)
|
||||
}
|
||||
d, err := Open(dsn)
|
||||
if err != nil {
|
||||
b.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
b.Cleanup(func() { d.Close() })
|
||||
var exp int64
|
||||
for i := range 3 {
|
||||
task, err := d.CreateTask(fmt.Sprintf("bench-task-list-metrics-%d", i), "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.Cleanup(func() { _, _ = d.DeleteTaskCascade(task.ID, false, true) })
|
||||
if _, err := d.Exec(`INSERT INTO activity(exploration_id,kind,input_tokens,created_at)
|
||||
SELECT $1, CASE WHEN g%20=0 THEN 'result' ELSE 'tool_use' END, 100, now() - (g||' seconds')::interval
|
||||
FROM generate_series(1,40000) g`, task.ExplorationID); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
exp = task.ExplorationID
|
||||
}
|
||||
if _, err := d.Exec(`ANALYZE activity`); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.Run("one", func(b *testing.B) {
|
||||
for range b.N {
|
||||
if _, err := d.TaskListMetrics(exp); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
b.Run("all", func(b *testing.B) {
|
||||
for range b.N {
|
||||
if _, err := d.TaskListMetricsAll(); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
+2
-2
@@ -1927,11 +1927,11 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
dto := taskDTO(t, s.resolvedTaskStatus(t))
|
||||
metrics, err := s.m.PG().TaskListMetricsAll()
|
||||
metrics, err := s.m.PG().TaskListMetrics(t.ExpID)
|
||||
if err != nil {
|
||||
log.Printf("[tasks] detail metrics query failed for %s — metrics zeroed: %v", t.ID, err)
|
||||
}
|
||||
s.applyTaskMetrics(&dto, t, metrics[t.ExpID])
|
||||
s.applyTaskMetrics(&dto, t, metrics)
|
||||
archiveBlockers, _ := s.m.PG().TaskArchiveBlockers()
|
||||
applyTaskArchiveBlocker(&dto, archiveBlockers)
|
||||
writeJSON(w, 200, dto)
|
||||
|
||||
@@ -38,6 +38,21 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// A second task, so the detail must pick its own row and the list must keep
|
||||
// serving every task.
|
||||
other, err := m.CreateTask("detail metrics other", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
otherID, _ := strconv.ParseInt(other.ID, 10, 64)
|
||||
defer func() { _, _ = m.pg.DeleteTaskCascade(otherID, false, true) }()
|
||||
defer func() { _, _ = m.pg.Exec(`DELETE FROM findings WHERE task_id=$1`, taskID) }()
|
||||
if _, err := m.pg.Exec(`
|
||||
INSERT INTO findings(task_id, severity) VALUES ($1,'critical'), ($1,'high'), ($1,'high'), ($1,'low'), ($2,'medium')`,
|
||||
taskID, otherID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
|
||||
defer stopTestServer(t, s, cancel)
|
||||
@@ -45,7 +60,19 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// After New: restoring the task resets running intents to open.
|
||||
if _, err := m.pg.Exec(`
|
||||
INSERT INTO exploration_nodes(exploration_id, kind, payload, state)
|
||||
VALUES ($1,'intent','{}','running'), ($1,'intent','{}','done')`, task.ExpID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
type findingCounts struct {
|
||||
Critical int `json:"critical"`
|
||||
High int `json:"high"`
|
||||
Medium int `json:"medium"`
|
||||
Low int `json:"low"`
|
||||
}
|
||||
var probe struct {
|
||||
Tokens struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
@@ -53,8 +80,10 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
CacheReadTokens int `json:"cache_read_tokens"`
|
||||
CacheWriteTokens int `json:"cache_write_tokens"`
|
||||
} `json:"tokens"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
InFlight int `json:"in_flight"`
|
||||
Findings findingCounts `json:"findings"`
|
||||
}
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/tasks/"+task.ID, nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
@@ -73,6 +102,9 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
if probe.GoalsMet != 1 || probe.GoalsTotal != 2 {
|
||||
t.Fatalf("detail goal progress not filled from metrics: %d/%d", probe.GoalsMet, probe.GoalsTotal)
|
||||
}
|
||||
if probe.InFlight != 1 || probe.Findings != (findingCounts{Critical: 1, High: 2, Low: 1}) {
|
||||
t.Fatalf("detail in-flight/findings not filled from this task's metrics: %d %+v", probe.InFlight, probe.Findings)
|
||||
}
|
||||
|
||||
// The list endpoint must report the same numbers for the same task.
|
||||
listReq := httptest.NewRequest(http.MethodGet, "/api/tasks", nil)
|
||||
@@ -89,27 +121,32 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
} `json:"tokens"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
InFlight int `json:"in_flight"`
|
||||
Findings findingCounts `json:"findings"`
|
||||
} `json:"tasks"`
|
||||
}
|
||||
if err := json.Unmarshal(listRec.Body.Bytes(), &listed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
found, foundOther := false, false
|
||||
for _, lt := range listed.Tasks {
|
||||
if lt.ID == other.ID {
|
||||
foundOther = lt.Findings == (findingCounts{Medium: 1})
|
||||
}
|
||||
if lt.ID != task.ID {
|
||||
continue
|
||||
}
|
||||
found = true
|
||||
if lt.Tokens.InputTokens != probe.Tokens.InputTokens ||
|
||||
lt.Tokens.OutputTokens != probe.Tokens.OutputTokens ||
|
||||
lt.GoalsMet != probe.GoalsMet || lt.GoalsTotal != probe.GoalsTotal {
|
||||
t.Fatalf("list and detail metrics disagree: list=%+v/%d detail=%+v/%d",
|
||||
lt.Tokens, lt.GoalsMet, probe.Tokens, probe.GoalsTotal)
|
||||
lt.GoalsMet != probe.GoalsMet || lt.GoalsTotal != probe.GoalsTotal ||
|
||||
lt.InFlight != probe.InFlight || lt.Findings != probe.Findings {
|
||||
t.Fatalf("list and detail metrics disagree: list=%+v detail=%+v", lt, probe)
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("task %s missing from list response", task.ID)
|
||||
if !found || !foundOther {
|
||||
t.Fatalf("list misses a task or its metrics: task %s found=%v, task %s with its findings=%v", task.ID, found, other.ID, foundOther)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user