mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-13 11:40:53 +08:00
- resetSettledRun clears settling and the budget ledger; a done/failed task revived to running now plans and claims again (was gated by the old wrap-up). resetTimeoutRevival calls it. Test: TestResetSettledRunRestartsBudget. - docker-compose.yml passes the four ARTEX_*_BUDGET_* variables through. - .env.example and README document the variables, the wrap-up, the in-memory ledger and the practical floor. - llmrec: the usage sink runs even without a PostgreSQL recorder; comment fixed. - BudgetWrapupBlocks allows update_finding_report.
414 lines
16 KiB
Go
414 lines
16 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/Autumn-27/artex/agent"
|
|
"github.com/Autumn-27/artex/db"
|
|
)
|
|
|
|
type goalSpec struct {
|
|
Text string
|
|
VulnClass string
|
|
}
|
|
|
|
// launchTask runs the shared post-creation sequence for a task created via ANY
|
|
// path (HTTP createTask or orchestration spawn_task), avoiding copy-paste in both places:
|
|
// 1. seed the root assets to feed the event-driven loop;
|
|
// 2. an optional seed intent so the worker starts without waiting for the first planner round;
|
|
// 3. decompose goals asynchronously in the background (emit a "round 0 goal decomposition" round + LLM decomposition steps + each goal, visible in the UI),
|
|
// then engine.Run once done — the engine starts only after the goal nodes are ready, avoiding a race where the planner runs before the goals exist.
|
|
//
|
|
// Asynchronous (goroutine) so the caller returns immediately; both paths behave the same: the task exists in seconds and goals decompose in the background.
|
|
func (s *Server) launchTask(t *Task, seedText string, seedFirstIntent bool) {
|
|
if !s.engine.beginTaskOperation(t.ID) {
|
|
return
|
|
}
|
|
s.seed(t, seedText)
|
|
if seedFirstIntent {
|
|
s.seedFirstIntent(t)
|
|
}
|
|
s.engine.decInflight(t.ID)
|
|
if _, err := s.admitTask(t, "bootstrap"); err != nil {
|
|
log.Printf("[concurrency] task %s failed to start: %v", t.ID, err)
|
|
}
|
|
}
|
|
|
|
func (s *Server) startTaskEngine(t *Task) {
|
|
ctx := s.engine.execContextFor(s.ctx, t.ID)
|
|
if ctx.Err() != nil || s.engine.IsDeleting(t.ID) {
|
|
return
|
|
}
|
|
s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
|
Summary: "round 0 goal decomposition"})
|
|
goals := s.createGoals(ctx, t, func(r db.Activity) {
|
|
s.engine.emitActivity(t, r)
|
|
})
|
|
if ctx.Err() != nil || s.engine.IsDeleting(t.ID) {
|
|
return
|
|
}
|
|
for _, g := range goals {
|
|
summary := g.Text
|
|
if g.VulnClass != "" {
|
|
summary = fmt.Sprintf("[%s] %s", g.VulnClass, g.Text)
|
|
}
|
|
s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "text", Summary: summary})
|
|
}
|
|
if ctx.Err() != nil || s.engine.IsDeleting(t.ID) {
|
|
return
|
|
}
|
|
s.engine.Run(s.ctx, t)
|
|
}
|
|
|
|
func (s *Server) occupiesConcurrencySlot(t *Task) bool {
|
|
if t == nil {
|
|
return false
|
|
}
|
|
lifecycle := t.lifecycleSnapshot()
|
|
if lifecycle.Queued || lifecycle.Paused || isTerminalStatus(lifecycle.Status) {
|
|
return false
|
|
}
|
|
// Deletion temporarily pauses the Engine but has not committed yet. Preserve
|
|
// the task's slot until PostgreSQL deletion succeeds; otherwise FIFO promotion
|
|
// during the drain window could over-admit if deletion later aborts and the
|
|
// persisted running task is restored.
|
|
if s.engine.IsDeleting(t.ID) {
|
|
return true
|
|
}
|
|
return !s.engine.IsPaused(t.ID) && (s.engine.ReadyFor(t) || s.engine.ActiveLLMCalls(t.ID) > 0)
|
|
}
|
|
|
|
func (s *Server) runningTaskCount(excludeID string) int {
|
|
count := 0
|
|
for _, task := range s.m.List() {
|
|
if task.ID != excludeID && s.occupiesConcurrencySlot(task) {
|
|
count++
|
|
}
|
|
}
|
|
return count
|
|
}
|
|
|
|
// admitTask is the single admission path for new, resumed, rerun and follow-up
|
|
// work. It atomically either starts the task or appends it to the persistent FIFO
|
|
// queue. mode is bootstrap for a freshly-created task and resume otherwise.
|
|
func (s *Server) admitTask(t *Task, mode string) (queued bool, err error) {
|
|
return s.admitTaskWhen(t, mode, false)
|
|
}
|
|
|
|
// admitPausedTask is the task-control resume path. The paused precondition is
|
|
// checked under the same scheduler lock as admission so a concurrent pause,
|
|
// dequeue or FIFO promotion cannot leave database and Engine state divergent.
|
|
func (s *Server) admitPausedTask(t *Task) (queued bool, err error) {
|
|
return s.admitTaskWhen(t, "resume", true)
|
|
}
|
|
|
|
func (s *Server) admitTaskWhen(t *Task, mode string, requirePaused bool) (queued bool, err error) {
|
|
if t == nil {
|
|
return false, fmt.Errorf("task not found")
|
|
}
|
|
if mode != "bootstrap" {
|
|
mode = "resume"
|
|
}
|
|
s.concMu.Lock()
|
|
defer s.concMu.Unlock()
|
|
// Delete installs its barrier under concMu as well. Re-resolve after acquiring
|
|
// the lock so a request that captured a task pointer before successful deletion
|
|
// cannot revive that stale handle after StopTask clears its Engine maps.
|
|
current, exists := s.m.Task(t.ID)
|
|
if !exists || current != t || s.engine.IsDeleting(t.ID) {
|
|
return false, fmt.Errorf("task is being deleted")
|
|
}
|
|
if !s.engine.beginTaskOperation(t.ID) {
|
|
return false, fmt.Errorf("task is being deleted")
|
|
}
|
|
defer s.engine.decInflight(t.ID)
|
|
lifecycle := t.lifecycleSnapshot()
|
|
if requirePaused {
|
|
if isTerminalStatus(lifecycle.Status) {
|
|
return false, fmt.Errorf("a terminal task cannot be resumed")
|
|
}
|
|
if !lifecycle.Paused {
|
|
return false, fmt.Errorf("only a paused task can be resumed")
|
|
}
|
|
mode = s.resumeAdmissionMode(t)
|
|
}
|
|
wasTerminal := isTerminalStatus(lifecycle.Status)
|
|
wasPaused := lifecycle.Paused || s.engine.IsPaused(t.ID)
|
|
wasQueued := lifecycle.Queued
|
|
engineWasPaused := s.engine.IsPaused(t.ID)
|
|
enabled, limit := s.m.ConcurrencyLimit()
|
|
ready := s.engine.ReadyFor(t)
|
|
|
|
// Work added to a task that is already admitted only needs a wake-up. This
|
|
// matters when the configured limit was lowered below the current running
|
|
// count: an existing task must not suddenly mark itself queued while its
|
|
// planner/workers are still live. Still compare-and-commit the persisted
|
|
// status: a concurrent terminal transition must win instead of being silently
|
|
// reported as a successful follow-up admission.
|
|
if !wasTerminal && !wasPaused && !wasQueued && s.engine.Started(t.ID) && (!enabled || ready) {
|
|
if err := s.m.ApplyTaskAdmission(t.ID, lifecycle.Status, lifecycle.Status, false, "resume", false); err != nil {
|
|
return false, err
|
|
}
|
|
s.startAdmittedTask(t, "resume")
|
|
return false, nil
|
|
}
|
|
|
|
// Preserve the original first-run mode across repeated admissions. Legacy
|
|
// queued rows have an empty queue_mode, so infer bootstrap from their graph.
|
|
if wasQueued {
|
|
switch lifecycle.QueueMode {
|
|
case "bootstrap":
|
|
mode = "bootstrap"
|
|
case "":
|
|
mode = s.resumeAdmissionMode(t)
|
|
}
|
|
}
|
|
|
|
readyBacklog := enabled && s.hasReadyQueuedTask(t.ID)
|
|
atCapacity := enabled && s.runningTaskCount(t.ID) >= limit
|
|
shouldQueue := enabled && (!ready || readyBacklog || atCapacity)
|
|
|
|
// Install the execution barrier before reviving a terminal/paused task. Without
|
|
// this ordering, its already-running worker loops can claim the newly-opened
|
|
// intent in the gap between status=running and queued=true.
|
|
if shouldQueue || wasTerminal || wasPaused || wasQueued {
|
|
s.engine.Pause(t.ID, agent.Causef("queued_for_admission", "task waiting for admission to run",
|
|
"the task is waiting on the concurrency queue or for its admission status to commit; this execution has stopped, and intents are claimed again only once it holds a run slot"))
|
|
}
|
|
|
|
status := lifecycle.Status
|
|
if wasTerminal {
|
|
status = "running"
|
|
}
|
|
if err := s.m.ApplyTaskAdmission(t.ID, lifecycle.Status, status, shouldQueue, mode, wasQueued); err != nil {
|
|
if !engineWasPaused && !wasPaused && !wasQueued {
|
|
s.engine.Resume(t)
|
|
}
|
|
return false, err
|
|
}
|
|
if lifecycle.Status == "timeout" && status == "running" {
|
|
// ApplyTaskAdmission reset this timed-out run's persisted clock. Clear the
|
|
// matching Engine gates before either parking it in FIFO or starting work;
|
|
// otherwise the old settling flag would make every worker skip forever.
|
|
s.engine.resetTimeoutRevival(t.ID)
|
|
}
|
|
if wasTerminal && status == "running" {
|
|
// done/failed -> running: the previous run's settle state must not gate the new one.
|
|
s.engine.resetSettledRun(t.ID)
|
|
}
|
|
if shouldQueue {
|
|
if !wasQueued {
|
|
summary := fmt.Sprintf("Queued: concurrency limit %d reached; starts automatically once a slot frees up", limit)
|
|
switch {
|
|
case !ready:
|
|
summary = "Queued: no runnable LLM configuration right now; starts automatically once one is available"
|
|
case readyBacklog:
|
|
summary = "Queued: earlier tasks are already waiting; starts automatically in FIFO order"
|
|
}
|
|
s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "text", Summary: summary})
|
|
}
|
|
return true, nil
|
|
}
|
|
s.startAdmittedTask(t, mode)
|
|
return false, nil
|
|
}
|
|
|
|
func (s *Server) hasReadyQueuedTask(excludeID string) bool {
|
|
for _, task := range s.m.List() {
|
|
lifecycle := task.lifecycleSnapshot()
|
|
if task.ID != excludeID && lifecycle.Queued && !lifecycle.Paused && !isTerminalStatus(lifecycle.Status) && s.engine.ReadyFor(task) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (s *Server) startAdmittedTask(t *Task, mode string) {
|
|
if mode == "bootstrap" {
|
|
if !s.engine.beginTaskOperation(t.ID) {
|
|
return
|
|
}
|
|
// A first-run task may have kept the Engine pause barrier while waiting
|
|
// in the concurrency queue. Clear it only after operation admission, or
|
|
// startTaskEngine's first execContextFor call would return a cancelled
|
|
// context and silently strand the dequeued task.
|
|
if s.engine.IsPaused(t.ID) {
|
|
s.engine.Resume(t)
|
|
}
|
|
go func() {
|
|
defer s.engine.decInflight(t.ID)
|
|
s.startTaskEngine(t)
|
|
}()
|
|
return
|
|
}
|
|
s.engine.Run(s.ctx, t)
|
|
s.engine.Resume(t)
|
|
// Run returns early for an already-started task. Explicitly ensure a timeout
|
|
// coordinator exists after resume; resetTimeoutRevival cleared the completed
|
|
// coordinator's marker and the next real call will stamp a fresh deadline.
|
|
s.engine.startDeadlineCoordinator(s.ctx, t)
|
|
}
|
|
|
|
func (s *Server) reconcileConcurrency() {
|
|
s.concMu.Lock()
|
|
defer s.concMu.Unlock()
|
|
enabled, limit := s.m.ConcurrencyLimit()
|
|
|
|
// A task whose provider chain becomes unavailable cannot do useful work and
|
|
// must not reserve a limited running slot forever. Park it persistently so a
|
|
// later profile edit/recovery can re-enter through the same FIFO path.
|
|
if enabled {
|
|
for _, task := range s.m.List() {
|
|
lifecycle := task.lifecycleSnapshot()
|
|
if lifecycle.Queued || lifecycle.Paused || s.engine.IsPaused(task.ID) || s.engine.IsDeleting(task.ID) ||
|
|
isTerminalStatus(lifecycle.Status) || s.engine.ReadyFor(task) || s.engine.ActiveLLMCalls(task.ID) > 0 {
|
|
continue
|
|
}
|
|
mode := s.resumeAdmissionMode(task)
|
|
s.engine.Pause(task.ID, agent.Causef("llm_unavailable_queued", "LLM unavailable, task moved to the wait queue",
|
|
"the task currently cannot resolve a runnable Planner/Worker LLM and has released its concurrency slot; it continues in queue order once the configuration recovers"))
|
|
if err := s.m.EnqueueTask(task.ID, mode); err != nil {
|
|
s.engine.Resume(task)
|
|
log.Printf("[concurrency] task %s failed to enqueue due to LLM unavailability: %v", task.ID, err)
|
|
continue
|
|
}
|
|
s.engine.emitActivity(task, db.Activity{Worker: "system", Kind: "text",
|
|
Summary: "Queued: no runnable LLM configuration right now; starts automatically once one is available"})
|
|
}
|
|
}
|
|
|
|
type queuedTask struct {
|
|
task *Task
|
|
lifecycle taskLifecycleState
|
|
}
|
|
queued := []queuedTask{}
|
|
for _, task := range s.m.List() {
|
|
lifecycle := task.lifecycleSnapshot()
|
|
if lifecycle.Queued && !isTerminalStatus(lifecycle.Status) {
|
|
queued = append(queued, queuedTask{task: task, lifecycle: lifecycle})
|
|
}
|
|
}
|
|
sort.SliceStable(queued, func(i, j int) bool {
|
|
left, right := queued[i].lifecycle.QueuedAt, queued[j].lifecycle.QueuedAt
|
|
if left == 0 {
|
|
left = queued[i].task.CreatedAt * int64(1e9)
|
|
}
|
|
if right == 0 {
|
|
right = queued[j].task.CreatedAt * int64(1e9)
|
|
}
|
|
if left != right {
|
|
return left < right
|
|
}
|
|
// Old rows may not have queued_at and task creation timestamps are only
|
|
// kept to second precision in memory. Task ids are monotonic, providing a
|
|
// deterministic oldest-first fallback for those ties.
|
|
leftID, leftErr := strconv.ParseInt(queued[i].task.ID, 10, 64)
|
|
rightID, rightErr := strconv.ParseInt(queued[j].task.ID, 10, 64)
|
|
if leftErr == nil && rightErr == nil {
|
|
return leftID < rightID
|
|
}
|
|
return queued[i].task.ID < queued[j].task.ID
|
|
})
|
|
slots := len(queued)
|
|
if enabled {
|
|
slots = limit - s.runningTaskCount("")
|
|
}
|
|
for _, entry := range queued {
|
|
if slots <= 0 {
|
|
break
|
|
}
|
|
task := entry.task
|
|
lifecycle := task.lifecycleSnapshot()
|
|
// Keep FIFO order, but do not consume a concurrency slot for a task
|
|
// whose explicit chain/global provider is unavailable. It will be retried
|
|
// after the user configures or resets its LLM chain.
|
|
if lifecycle.Paused || (enabled && !s.engine.ReadyFor(task)) {
|
|
continue
|
|
}
|
|
mode := lifecycle.QueueMode
|
|
if mode != "bootstrap" && mode != "resume" {
|
|
mode = s.resumeAdmissionMode(task)
|
|
}
|
|
if mode != "bootstrap" {
|
|
mode = "resume"
|
|
}
|
|
if !s.engine.beginTaskOperation(task.ID) {
|
|
continue
|
|
}
|
|
if err := s.m.ApplyTaskAdmission(task.ID, lifecycle.Status, lifecycle.Status, false, mode, false); err != nil {
|
|
s.engine.decInflight(task.ID)
|
|
continue
|
|
}
|
|
s.startAdmittedTask(task, mode)
|
|
s.engine.decInflight(task.ID)
|
|
slots--
|
|
}
|
|
}
|
|
|
|
// reviveTask brings a stopped task back to running: pull a terminal state (done/failed/timeout) back to running,
|
|
// lift the pause, and (re)start the engine loop + wake it. A task already running and not paused: only one Notify
|
|
// inside Run, nearly side-effect free. Used by "main agent set_goals adds a goal" and "rerun a blocked intent".
|
|
//
|
|
// Why revival must be explicit: the terminal-state gate of the planner/worker loops (engine.go)
|
|
// swallows ordinary notifies — graph edits + Notify cannot wake a task judged complete; after a restart its goroutine may be gone, hence Run.
|
|
func (s *Server) reviveTask(t *Task) {
|
|
if t == nil {
|
|
return
|
|
}
|
|
if _, err := s.admitTask(t, "resume"); err != nil {
|
|
log.Printf("[revive] task %s failed to revive: %v", t.ID, err)
|
|
}
|
|
}
|
|
|
|
// createGoals materializes the goal node(s) under the task root (rel objective).
|
|
// Decomposition is done ENTIRELY by the LLM (the project requires an LLM). There is
|
|
// no rule-based fallback splitter — it only ever produced garbage (shredded URLs,
|
|
// meaningless 2-way splits). If the LLM yields nothing (an error), the raw task goal
|
|
// is used verbatim as a single goal so the task still has something to judge against.
|
|
// Returns the seeded specs so callers can emit activity records for them.
|
|
// emit, when non-nil, is forwarded to DecomposeGoals so LLM steps are visible in the UI.
|
|
func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity)) []goalSpec {
|
|
if t == nil {
|
|
return nil
|
|
}
|
|
// Use the SAME LLM the task runs on (its pinned profile, else the active profile),
|
|
// NOT agent.FromEnv() — the LLM is configured via the UI (DB profile), not env vars,
|
|
// so FromEnv returned empty and every task silently fell back to the crude rule
|
|
// splitter (which shredded URLs / made meaningless 2-way splits).
|
|
var specs []goalSpec
|
|
var as *db.AssetStore
|
|
if s.m != nil {
|
|
as = s.m.Assets()
|
|
}
|
|
taskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
|
s.engine.BeginLLMCall(t.ID)
|
|
goalRuntime := s.agentsForTask(t).runtime
|
|
decomposed := agent.DecomposeGoalsWithProvider(ctx, goalRuntime, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, goalRuntime.nonStreaming(), goalRuntime.maxTokens(), emit)
|
|
s.engine.EndLLMCall(t.ID)
|
|
for _, g := range decomposed {
|
|
if strings.TrimSpace(g.Text) != "" {
|
|
specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass})
|
|
}
|
|
}
|
|
if len(specs) == 0 {
|
|
// No decomposed goals (LLM error / no provider): use the raw task goal verbatim
|
|
// as a single goal so the task still has something to judge against. This is the
|
|
// only path that writes here — decomposed goals are already persisted by the tool.
|
|
if g := strings.TrimSpace(t.Goal); g != "" {
|
|
log.Printf("[goals] task %s: LLM goal decomposition produced nothing, falling back to the raw goal as a single goal", t.ID)
|
|
origin, _ := t.Store.OriginFactID()
|
|
id, _ := t.Store.AddNode(db.KindGoal, map[string]any{"text": g}, 0, "open", "system", nil)
|
|
if origin > 0 && id > 0 {
|
|
_ = t.Store.Link(origin, db.RelSpawns, id) // goal descends from the task root (origin fact)
|
|
}
|
|
specs = []goalSpec{{Text: g}}
|
|
}
|
|
}
|
|
return specs
|
|
}
|