clanker-saddle/src/Services/AllocationService.vala

using Invercargill;
using Invercargill.DataStructures;
using Invercargill.Expressions;
using InvercargillSql;
using InvercargillSql.Orm;
using Inversion;
using ClankerSaddle.Data;

namespace ClankerSaddle.Services {

    /**
     * One offerable task for the wake computation: a fresh queued task
     * (dependencies met, {@link affine_saddle_id} 0) or a parked task
     * ready to resume — blocked/review-requested with its newest block
     * answered, or suspended with a terminal subtask event since it
     * parked — pinned to the saddle holding its open session.
     */
    public class OfferableTask : Object {
        public TaskEntity task { get; set; }
        /** The saddle whose parked workspace holds the task (0 = fresh). */
        public int64 affine_saddle_id { get; set; }
        /** True when resuming a parked workspace (vs. a fresh dispatch). */
        public bool parked { get; set; }
        /** True when the parked task is a suspended parent's resume. */
        public bool suspended_resume { get; set; }
    }

    /**
     * A claimed allocation: the transactional half of
     * `POST /api/saddle/next-task` succeeded — the task is `assigned`
     * to the asking saddle, its session row exists and the saddle row
     * is busy. The assignment payload is assembled from it afterwards.
     */
    public class AllocationClaim : Object {
        public int64 task_id { get; set; }
        public int64 session_id { get; set; }
        public bool parked { get; set; }
        public bool suspended_resume { get; set; }
    }

    /**
     * The ask/allocation half of the plan-10 work-acquisition flip:
     * daemons ask (`POST /api/saddle/next-task`), the allocation picks
     * the best offerable task for the asking saddle and claims it in a
     * single write transaction, and the dispatcher's wake computation
     * uses the same offerable set to decide which sleeping daemons to
     * wake.
     *
     * Two filters, in priority order:
     *
     * 1. **Workspace affinity** — a task parked (blocked, awaiting
     *    review or suspended) on this saddle is offered ONLY to this
     *    saddle: its workspace lives here, and capabilities do not
     *    apply (the parked session pins the pairing, exactly like the
     *    pre-flip resume paths).
     * 2. **Capability match** — fresh queued tasks (dependencies met)
     *    go to the asking saddle only when its capability set is a
     *    superset of the task's required tags.
     *
     * The claim itself runs inside one `BEGIN IMMEDIATE` transaction
     * with no awaits, so concurrent asks cannot interleave between the
     * decision and the writes: exactly one ask wins a task even when
     * several fire simultaneously. Every candidate is re-validated
     * inside the transaction (a racing ask may have claimed it after
     * the candidate list was built), and a lost race simply falls
     * through to the next candidate.
     */
    public class AllocationService : Object {

        private OrmSession db = inject<OrmSession>();
        private Connection connection = inject<Connection>();
        private TaskService tasks = inject<TaskService>();
        private CapabilityService capabilities = inject<CapabilityService>();
        private SubtaskEventQueue subtask_events = inject<SubtaskEventQueue>();
        private StatePulse pulse = inject<StatePulse>();
        private SaddlePollState poll_state = inject<SaddlePollState>();

        /**
         * Allocates the best offerable task to `saddle`, or null when
         * nothing is offerable (the daemon sleeps). The transactional
         * claim runs first; the assignment payload (answers, resume
         * notice, effective rules, project notes, git identity, llm
         * config, session id) is assembled from the claimed rows after
         * the commit.
         */
        public async Json.Object? allocate_next(SaddleService saddles, ProjectService projects,
                Saddle saddle) throws Error {
            if (LlmEndpoint.from_environment() == null) {
                // No LLM endpoint configured: assignment stays paused
                // (the dispatcher logs it once); the daemon sleeps.
                return null;
            }

            var candidates = yield candidates_for(saddle);
            AllocationClaim? claim = null;
            foreach (var candidate in candidates) {
                claim = claim_candidate(saddle, candidate);
                if (claim != null) {
                    break;
                }
            }
            if (claim == null) {
                return null;
            }
            return yield assemble_assignment(saddles, projects, saddle, (!)claim);
        }

        /**
         * Every offerable task (see {@link OfferableTask}) — the
         * dispatcher's wake computation walks this list: parked tasks
         * target their affine saddle first, fresh tasks any capable
         * sleeping one. A parked task without an open session has no
         * workspace to resume and is not offerable.
         */
        public async ImmutableLot<OfferableTask> offerable_tasks() throws Error {
            var rows = new Vector<OfferableTask>();
            foreach (var status in new string[] { STATUS_BLOCKED, STATUS_REVIEW_REQUESTED }) {
                foreach (var task in yield tasks.list_tasks(status, null)) {
                    if (yield tasks.newest_block_answered(task.id)) {
                        var parked = yield parked_offerable(task, false);
                        if (parked != null) {
                            rows.add((!)parked);
                        }
                    }
                }
            }
            foreach (var parent in yield tasks.list_tasks(STATUS_SUSPENDED, null)) {
                if (yield tasks.has_terminal_subtask_since(parent)) {
                    var parked = yield parked_offerable(parent, true);
                    if (parked != null) {
                        rows.add((!)parked);
                    }
                }
            }
            foreach (var task in yield tasks.get_ready_tasks()) {
                rows.add(new OfferableTask() {
                    task = task,
                    affine_saddle_id = 0,
                    parked = false,
                    suspended_resume = false
                });
            }
            return rows.to_immutable_buffer();
        }

        /**
         * The offerable tasks `saddle` may be allocated: affinity
         * candidates first (parked on this saddle, priority order),
         * then fresh queued tasks whose required capabilities are
         * covered (ready order — priority, then creation).
         */
        private async ImmutableLot<OfferableTask> candidates_for(Saddle saddle) throws Error {
            var rows = new Vector<OfferableTask>();
            foreach (var offerable in yield offerable_tasks()) {
                if (offerable.parked) {
                    if (offerable.affine_saddle_id == saddle.id) {
                        rows.add(offerable);
                    }
                } else if (yield capabilities.saddle_covers_task(saddle.id,
                        offerable.task.id)) {
                    rows.add(offerable);
                }
            }
            return rows.to_immutable_buffer();
        }

        /**
         * A parked offerable bound to its session's saddle, or null when
         * the task has no open (parked) session — nothing to resume.
         */
        private async OfferableTask? parked_offerable(TaskEntity task, bool suspended_resume)
                throws Error {
            var session = yield tasks.open_session_for_task(task.id);
            if (session == null) {
                return null;
            }
            return new OfferableTask() {
                task = task,
                affine_saddle_id = ((!)session).saddle_id,
                parked = true,
                suspended_resume = suspended_resume
            };
        }

        /**
         * The single-transaction claim. `BEGIN IMMEDIATE` takes the
         * SQLite write lock up front; the section is fully synchronous
         * (no awaits), so no other ask, action or tick can interleave
         * between the re-validation reads and the claim writes —
         * concurrent asks race to exactly one winner. Returns null when
         * the candidate was claimed by a racing ask or its saddle
         * became unavailable (the caller falls through to the next
         * candidate).
         */
        private AllocationClaim? claim_candidate(Saddle saddle, OfferableTask candidate) throws Error {
            connection.execute("BEGIN IMMEDIATE");
            try {
                var task = db.query<TaskEntity>()
                    .where(expr("id == $0", elem<int64?>(candidate.task.id)))
                    .first();
                if (task == null || !still_claimable((!)task, candidate)) {
                    connection.execute("ROLLBACK");
                    return null;
                }

                var row = db.query<Saddle>()
                    .where(expr("id == $0", elem<int64?>(saddle.id)))
                    .first();
                if (row == null || (((!)row).status != "idle"
                        && (!(candidate.parked && ((!)row).current_task_id == candidate.task.id)))) {
                    connection.execute("ROLLBACK");
                    return null;
                }

                int64 session_id = 0;
                if (candidate.parked) {
                    var session = db.query<Session>()
                        .where(expr("task_id == $0 && saddle_id == $1",
                            elem<int64?>(candidate.task.id), elem<int64?>(saddle.id)))
                        .order_by(expr("id"))
                        .materialise();
                    Session? newest_open = null;
                    foreach (var open in session) {
                        if (open.ended_at == null) {
                            newest_open = open;
                        }
                    }
                    if (newest_open == null) {
                        // The parked session closed underneath us (a
                        // concurrent requeue); nothing to resume onto.
                        connection.execute("ROLLBACK");
                        return null;
                    }
                    // Plan 11: the resuming assignment carries a FRESH
                    // session id — the parked row closes with outcome
                    // 'resumed' (the park record) and the daemon keys
                    // every telemetry push of the resumed round on the
                    // new row.
                    ((!)newest_open).ended_at = new DateTime.now_local();
                    ((!)newest_open).outcome = "resumed";
                    db.update<Session>((!)newest_open);
                    var resumed = new Session() {
                        task_id = candidate.task.id,
                        saddle_id = saddle.id,
                        started_at = new DateTime.now_local()
                    };
                    db.insert<Session>(resumed);
                    session_id = resumed.id;
                } else {
                    var session = new Session() {
                        task_id = candidate.task.id,
                        saddle_id = saddle.id,
                        started_at = new DateTime.now_local()
                    };
                    db.insert<Session>(session);
                    session_id = session.id;
                }

                ((!)task).status = STATUS_ASSIGNED;
                ((!)task).started_at = new DateTime.now_local();
                db.update<TaskEntity>((!)task);

                ((!)row).status = "busy";
                ((!)row).current_task_id = candidate.task.id;
                ((!)row).allocated_at = new DateTime.now_local();
                db.update<Saddle>((!)row);

                connection.execute("COMMIT");
                return new AllocationClaim() {
                    task_id = candidate.task.id,
                    session_id = session_id,
                    parked = candidate.parked,
                    suspended_resume = candidate.suspended_resume
                };
            } catch (Error e) {
                try {
                    connection.execute("ROLLBACK");
                } catch (Error rollback_error) {
                    warning("[ClankerSaddle] Allocation rollback failed: %s",
                        rollback_error.message);
                }
                throw e;
            }
        }

        /**
         * The in-transaction re-validation: the task must still be in
         * its candidate state (a racing ask that claimed it flipped the
         * status), and a fresh candidate's dependencies must still be
         * met.
         */
        private bool still_claimable(TaskEntity task, OfferableTask candidate) throws Error {
            if (candidate.parked) {
                return task.status == STATUS_BLOCKED
                    || task.status == STATUS_REVIEW_REQUESTED
                    || (candidate.suspended_resume && task.status == STATUS_SUSPENDED);
            }
            if (task.status != STATUS_QUEUED) {
                return false;
            }
            foreach (var dependency in db.query<TaskDependency>()
                    .where(expr("task_id == $0", elem<int64?>(task.id)))
                    .materialise()) {
                var other = db.query<TaskEntity>()
                    .where(expr("id == $0", elem<int64?>(dependency.depends_on_task_id)))
                    .first();
                if (other == null || ((!)other).status != STATUS_DONE) {
                    return false;
                }
            }
            return true;
        }

        /**
         * Assembles the claimed task's assignment payload (today's
         * /start shape): pending block answers ride the task object,
         * and every parked resume carries a `resume_notice` — a
         * suspended parent's the coalesced subtask report, an answered
         * block's the verdict text in the stable wording ({@link
         * TaskService.verdict_notice_for}) — plus the
         * `branch_advanced` verdict. The hand-over itself starts the
         * session's push-liveness clock (the daemon cannot push before
         * it owns the assignment), so the staleness probe's window
         * opens here.
         */
        private async Json.Object assemble_assignment(SaddleService saddles,
                ProjectService projects, Saddle saddle, AllocationClaim claim) throws Error {
            var detail = yield tasks.get_task(claim.task_id);
            if (detail == null) {
                throw new Error(Quark.from_string("clanker-saddle"), 48,
                    @"Task $(claim.task_id) vanished after its allocation claim");
            }
            var task = ((!)detail).task;

            string? resume_notice = null;
            bool? branch_advanced = null;
            if (claim.parked) {
                branch_advanced = yield tasks.branch_advanced_since(task.id);
                if (claim.suspended_resume) {
                    resume_notice = yield tasks.subtask_status_report(task.id);
                    // The parking window is over: a later suspension
                    // starts a fresh event window.
                    yield tasks.clear_suspension(task.id);
                } else {
                    // Answering = allocating (plan 11): the verdict rides
                    // the resuming assignment's resume_notice, injected as
                    // the restored conversation's opening turn.
                    var newest = yield tasks.newest_block(task.id);
                    if (newest != null && ((!)newest).answer != null) {
                        resume_notice = TaskService.verdict_notice_for((!)newest);
                    }
                }
            }

            poll_state.note_push(claim.session_id);
            var payload = yield saddles.build_assignment_payload(saddle, task,
                claim.session_id, projects, resume_notice, branch_advanced);
            payload.set_int_member("session_id", claim.session_id);

            yield saddles.mark_task_comments_delivered(task.id);
            subtask_events.record_dispatch(task);
            yield pulse.bump(task.id, "task dispatched");
            print(@"[ClankerSaddle] Allocated task $(task.id) to saddle '$(saddle.name)' on ask (session $(claim.session_id))\n");
            return payload;
        }
    }
}