Skip to main content

ag_store/
orchestration.rs

1//! Orchestration and orchestration-task persistence adapters.
2//!
3//! Task rows are written when the controller proposes a plan, before any child
4//! session exists. That ordering is what makes fan-out idempotent: a restart or
5//! a retry reuses the `(session_orchestration_id, task_key)` unique key instead
6//! of creating a second child for the same subtask.
7
8use std::sync::Arc;
9
10use ag_session::IntegrationApproach;
11use async_trait::async_trait;
12use sqlx::SqlitePool;
13
14use super::status;
15use crate::timestamp::TimestampSource;
16use crate::{DbError, DbResultExt};
17
18/// Row returned when loading one `session_orchestration`.
19#[derive(Clone, Debug, Eq, PartialEq)]
20pub struct SessionOrchestrationRow {
21    /// Project containing the controller and all child sessions.
22    pub controller_project_id: i64,
23    /// Controller session that owns this orchestration.
24    pub controller_session_id: String,
25    /// Canonical single-goal statement approved for this campaign.
26    pub goal_statement: String,
27    /// Stable database identifier.
28    pub id: i64,
29    /// Maximum number of children allowed to run at once.
30    pub max_parallelism: i64,
31    /// Exact managed task whose questions are mirrored onto the controller.
32    pub relayed_question_task_id: Option<i64>,
33    /// Persisted orchestration status string.
34    pub status: String,
35    /// Monotonic identity for durable verification turns.
36    pub verification_generation: i64,
37}
38
39/// Row returned when loading one `session_orchestration_task`.
40#[derive(Clone, Debug, Eq, PartialEq)]
41pub struct SessionOrchestrationTaskRow {
42    /// Serialized acceptance criteria checked during verification.
43    pub acceptance_criteria: String,
44    /// Serialized changed paths outside the task's expected planning areas.
45    pub area_violations: String,
46    /// Whether the latest child diff stayed within its expected planning areas.
47    pub areas_compliant: Option<bool>,
48    /// Number of child sessions created for this task so far.
49    pub attempt_count: i64,
50    /// Persisted added-line count from the latest child diff refresh.
51    pub child_added_lines: i64,
52    /// Latest assistant answer emitted by the linked child.
53    pub child_answer: Option<String>,
54    /// Persisted deleted-line count from the latest child diff refresh.
55    pub child_deleted_lines: i64,
56    /// Durable focused-review generation state observed on the child.
57    pub child_focused_review_status: Option<String>,
58    /// Latest focused-review markdown observed on the child.
59    pub child_focused_review_text: Option<String>,
60    /// Whether the latest child diff refresh found any content.
61    pub child_has_diff: Option<bool>,
62    /// Total input tokens observed on the linked child session.
63    pub child_input_tokens: i64,
64    /// Total output tokens observed on the linked child session.
65    pub child_output_tokens: i64,
66    /// Persisted clarification questions on the linked child.
67    pub child_questions: Option<String>,
68    /// Child session created for this task, when one exists.
69    pub child_session_id: Option<String>,
70    /// Persisted lifecycle status observed on the linked child session.
71    pub child_status: Option<String>,
72    /// Cumulative summary observed on the linked child session.
73    pub child_summary: Option<String>,
74    /// Monotonic identity for durable feedback delivery attempts.
75    pub continuation_generation: i64,
76    /// Feedback prompt waiting to resume the existing managed child.
77    pub continuation_prompt: Option<String>,
78    /// Stable database identifier.
79    pub id: i64,
80    /// Number of bounded automatic spawn retries already consumed.
81    pub infrastructure_retry_count: i64,
82    /// Persisted execution behavior string.
83    pub kind: String,
84    /// Most recent failure detail, when the task failed.
85    pub last_error: Option<String>,
86    /// Stable integration order selected on the approval board.
87    pub merge_position: i64,
88    /// Standalone prompt handed to the child session.
89    pub prompt: String,
90    /// Bounded full report captured from a temporary research child.
91    pub research_report: Option<String>,
92    /// Bounded child-reported result summary used for fan-in, when present.
93    pub result_summary: Option<String>,
94    /// Number of automatic focused-review remediation turns already consumed.
95    pub review_iteration: i64,
96    /// Persisted task status string.
97    pub status: String,
98    /// Stable subtask key unique within the owning orchestration.
99    pub task_key: String,
100    /// Serialized repository areas this task expects to touch.
101    pub touched_areas: String,
102    /// Short human-readable task title.
103    pub title: String,
104    /// Controller explanation for the latest verification verdict.
105    pub verification_reason: Option<String>,
106    /// Latest controller verdict for this task.
107    pub verification_verdict: Option<String>,
108}
109
110/// Task scope and child base needed to compute verification evidence.
111#[derive(Clone, Debug, Eq, PartialEq)]
112pub struct OrchestrationTaskScopeRow {
113    /// Base branch used by the managed child worktree.
114    pub base_branch: String,
115    /// Stable orchestration-task identifier.
116    pub id: i64,
117    /// Serialized repository areas assigned to this task.
118    pub touched_areas: String,
119}
120
121/// Bulk-loaded controller progress and child adjacency for one session row.
122#[derive(Clone, Debug, Eq, PartialEq)]
123pub struct SessionOrchestrationMetadataRow {
124    /// Owning controller for an orchestration child, when this is a child row.
125    pub controller_session_id: Option<String>,
126    /// Latest orchestration status for a controller row, when this is a
127    /// controller.
128    pub orchestration_status: Option<String>,
129    /// Number of child tasks currently creating or running.
130    pub running_task_count: i64,
131    /// Session receiving the derived metadata.
132    pub session_id: String,
133    /// Number of child tasks currently waiting for user input.
134    pub waiting_task_count: i64,
135}
136
137/// Values used to persist one planned orchestration task.
138///
139/// Owns its fields so the persistence trait method stays lifetime-free. A
140/// borrowed variant forced the trait to carry a generic lifetime, which
141/// `mockall::automock` drops in the generated mock. Owning the data is
142/// allocation-cheap on this once-per-plan path.
143pub struct PersistedOrchestrationTask {
144    /// Serialized acceptance criteria checked during verification.
145    pub acceptance_criteria: String,
146    /// Persisted execution behavior.
147    pub kind: String,
148    /// Stable integration order selected on the approval board.
149    pub merge_position: i64,
150    /// Standalone prompt handed to the child session.
151    pub prompt: String,
152    /// Owning orchestration identifier.
153    pub session_orchestration_id: i64,
154    /// Stable subtask key unique within the owning orchestration.
155    pub task_key: String,
156    /// Short human-readable task title.
157    pub title: String,
158    /// Serialized repository areas this task expects to touch.
159    pub touched_areas: String,
160}
161
162/// Orchestration persistence boundary used by the coordinator and tests.
163///
164/// The coordinator owns its own pool through this trait so reconciliation reads
165/// never contend with the foreground session-runtime mailbox.
166#[cfg_attr(any(test, feature = "test-utils"), mockall::automock)]
167#[async_trait]
168pub trait OrchestrationRepository: Send + Sync {
169    /// Inserts one orchestration and returns its stable identifier.
170    async fn insert_orchestration(
171        &self,
172        controller_session_id: &str,
173        status: &str,
174        max_parallelism: i64,
175    ) -> Result<i64, DbError>;
176
177    /// Inserts one planned task, replacing any previous attempt that used the
178    /// same `task_key` within the same orchestration.
179    ///
180    /// Re-proposing a task key preserves the existing row identity and its
181    /// `attempt_count`, so a retry updates the plan in place instead of fanning
182    /// out a duplicate child. The retry transaction detaches both persisted
183    /// directions of any prior child link before replacement creation.
184    async fn upsert_orchestration_task(
185        &self,
186        task: PersistedOrchestrationTask,
187    ) -> Result<i64, DbError>;
188
189    /// Loads the most recent orchestration owned by one controller session.
190    async fn load_orchestration_for_controller(
191        &self,
192        controller_session_id: &str,
193    ) -> Result<Option<SessionOrchestrationRow>, DbError>;
194
195    /// Loads every orchestration whose persisted status is still active.
196    async fn load_active_orchestrations(&self) -> Result<Vec<SessionOrchestrationRow>, DbError>;
197
198    /// Loads managed `Reviewing` sessions whose incomplete focused-review
199    /// state must be regenerated during app startup.
200    async fn load_recoverable_focused_review_session_ids(
201        &self,
202        project_id: i64,
203    ) -> Result<Vec<String>, DbError>;
204
205    /// Loads controller progress and child adjacency for one project's
206    /// sessions in a single query.
207    async fn load_session_metadata_for_project(
208        &self,
209        project_id: i64,
210    ) -> Result<Vec<SessionOrchestrationMetadataRow>, DbError>;
211
212    /// Loads all tasks belonging to one orchestration in stable plan order.
213    async fn load_orchestration_tasks(
214        &self,
215        session_orchestration_id: i64,
216    ) -> Result<Vec<SessionOrchestrationTaskRow>, DbError>;
217
218    /// Loads the destination selected for verified child branches.
219    async fn load_orchestration_integration_approach(&self, id: i64) -> Result<String, DbError>;
220
221    /// Loads the declared task scope linked to one managed child.
222    async fn load_orchestration_task_scope_for_child(
223        &self,
224        child_session_id: &str,
225    ) -> Result<Option<OrchestrationTaskScopeRow>, DbError>;
226
227    /// Loads a child session already persisted for one orchestration task.
228    async fn load_child_session_id_for_task(&self, task_id: i64)
229    -> Result<Option<String>, DbError>;
230
231    /// Atomically blocks new fan-out before cascade cancellation begins.
232    async fn begin_orchestration_cancellation(&self, id: i64) -> Result<bool, DbError>;
233
234    /// Atomically claims one planned task while its orchestration is running.
235    async fn claim_orchestration_task(&self, id: i64) -> Result<bool, DbError>;
236
237    /// Atomically claims one focused-review remediation turn and clears the
238    /// consumed child review cache.
239    async fn claim_orchestration_review_application(
240        &self,
241        id: i64,
242        prompt: &str,
243        iteration_limit: i64,
244    ) -> Result<bool, DbError>;
245
246    /// Atomically claims roll-up submission for one running orchestration.
247    async fn claim_orchestration_rollup(&self, id: i64) -> Result<bool, DbError>;
248
249    /// Completes a submitted roll-up unless cancellation won the state race.
250    async fn complete_orchestration_rollup(&self, id: i64) -> Result<bool, DbError>;
251
252    /// Persists one explicit controller verdict for a settled task.
253    async fn record_orchestration_verdict(
254        &self,
255        id: i64,
256        task_key: &str,
257        is_pass: bool,
258        reason: &str,
259    ) -> Result<bool, DbError>;
260
261    /// Archives a finalized campaign and makes its controller terminal.
262    async fn complete_orchestration_campaign(&self, id: i64) -> Result<bool, DbError>;
263
264    /// Loads the durable worker-operation status for one roll-up delivery.
265    async fn load_rollup_operation_status(
266        &self,
267        operation_id: &str,
268    ) -> Result<Option<String>, DbError>;
269
270    /// Updates one orchestration's persisted status.
271    async fn update_orchestration_status(&self, id: i64, status: &str) -> Result<(), DbError>;
272
273    /// Approves parked proposed tasks and resumes campaign execution.
274    async fn approve_orchestration_plan(&self, id: i64) -> Result<bool, DbError>;
275
276    /// Persists the selected integration destination and starts integration.
277    async fn approve_orchestration_integration(
278        &self,
279        id: i64,
280        approach: IntegrationApproach,
281    ) -> Result<bool, DbError>;
282
283    /// Updates plan metadata before fan-out begins.
284    async fn update_orchestration_plan(
285        &self,
286        id: i64,
287        goal_statement: &str,
288        max_parallelism: i64,
289    ) -> Result<(), DbError>;
290
291    /// Routes feedback to a live managed child without replacing its branch.
292    async fn queue_orchestration_continuation(
293        &self,
294        id: i64,
295        prompt: &str,
296        acceptance_criteria: &str,
297        touched_areas: &str,
298    ) -> Result<bool, DbError>;
299
300    /// Returns previously verified tasks to fan-in before a follow-up wave.
301    async fn reset_orchestration_verification(&self, id: i64) -> Result<(), DbError>;
302
303    /// Records one infrastructure failure and returns the next task status.
304    async fn record_orchestration_spawn_failure(
305        &self,
306        id: i64,
307        error: &str,
308        retry_limit: i64,
309    ) -> Result<String, DbError>;
310
311    /// Permanently transfers one managed child to ordinary user ownership.
312    async fn detach_orchestration_child(&self, child_session_id: &str) -> Result<bool, DbError>;
313
314    /// Claims one managed task's questions and mirrors them onto its
315    /// controller.
316    async fn surface_orchestration_questions(
317        &self,
318        session_orchestration_id: i64,
319        task_id: i64,
320        questions: &str,
321    ) -> Result<bool, DbError>;
322
323    /// Clears the exact question proxy claimed by one orchestration.
324    async fn clear_orchestration_questions(
325        &self,
326        session_orchestration_id: i64,
327    ) -> Result<(), DbError>;
328
329    /// Links one child and counts the attempt only while fan-out still owns it.
330    async fn link_orchestration_task_child(
331        &self,
332        id: i64,
333        child_session_id: &str,
334    ) -> Result<bool, DbError>;
335
336    /// Updates one task's status and failure detail.
337    async fn update_orchestration_task_status(
338        &self,
339        id: i64,
340        status: &str,
341        last_error: Option<String>,
342    ) -> Result<(), DbError>;
343
344    /// Records one child's bounded result summary for fan-in.
345    async fn update_orchestration_task_result_summary(
346        &self,
347        id: i64,
348        result_summary: &str,
349    ) -> Result<(), DbError>;
350
351    /// Records one research child's bounded report before its temporary
352    /// worktree is discarded.
353    async fn update_orchestration_task_research_report(
354        &self,
355        id: i64,
356        research_report: &str,
357    ) -> Result<(), DbError>;
358
359    /// Records a mechanically computed touched-area planning comparison.
360    async fn update_orchestration_task_area_compliance(
361        &self,
362        id: i64,
363        areas_compliant: Option<bool>,
364        area_violations: &str,
365    ) -> Result<(), DbError>;
366}
367
368/// `SQLite` implementation of [`OrchestrationRepository`].
369#[derive(Clone)]
370pub(crate) struct SqliteOrchestrationRepository(SqlitePool, Arc<dyn TimestampSource>);
371
372struct OrchestrationTransition {
373    context: &'static str,
374    from_orchestration_status: &'static str,
375    from_task_status: &'static str,
376    require_pass_verdict: bool,
377    to_orchestration_status: &'static str,
378    to_task_status: &'static str,
379}
380
381impl SqliteOrchestrationRepository {
382    /// Creates an orchestration repository backed by the provided pool.
383    pub(crate) fn new(pool: SqlitePool, timestamp_source: Arc<dyn TimestampSource>) -> Self {
384        Self(pool, timestamp_source)
385    }
386
387    /// Returns the shared persistence timestamp in Unix seconds.
388    fn now(&self) -> i64 {
389        self.1.now_timestamp_seconds()
390    }
391
392    async fn transition_orchestration_and_tasks(
393        &self,
394        id: i64,
395        transition: OrchestrationTransition,
396    ) -> Result<bool, DbError> {
397        let now = self.now();
398        let mut transaction = self.0.begin().await.db_context(transition.context)?;
399        let result = sqlx::query(
400            r"
401UPDATE session_orchestration
402SET status = ?,
403    updated_at = ?
404WHERE id = ?
405  AND status = ?
406",
407        )
408        .bind(transition.to_orchestration_status)
409        .bind(now)
410        .bind(id)
411        .bind(transition.from_orchestration_status)
412        .execute(&mut *transaction)
413        .await
414        .db_context(transition.context)?;
415        if result.rows_affected() == 1 {
416            sqlx::query(
417                r"
418UPDATE session_orchestration_task
419SET status = ?,
420    updated_at = ?
421WHERE session_orchestration_id = ?
422  AND status = ?
423  AND (? = 0 OR verification_verdict = 'Pass')
424",
425            )
426            .bind(transition.to_task_status)
427            .bind(now)
428            .bind(id)
429            .bind(transition.from_task_status)
430            .bind(i64::from(transition.require_pass_verdict))
431            .execute(&mut *transaction)
432            .await
433            .db_context(transition.context)?;
434        }
435        transaction.commit().await.db_context(transition.context)?;
436
437        Ok(result.rows_affected() == 1)
438    }
439}
440
441#[async_trait]
442impl OrchestrationRepository for SqliteOrchestrationRepository {
443    async fn insert_orchestration(
444        &self,
445        controller_session_id: &str,
446        status: &str,
447        max_parallelism: i64,
448    ) -> Result<i64, DbError> {
449        status::validate_orchestration(status)?;
450        let now = self.now();
451
452        let row = sqlx::query!(
453            r#"
454INSERT INTO session_orchestration (
455    controller_session_id,
456    goal_statement,
457    status,
458    max_parallelism,
459    created_at,
460    updated_at
461)
462VALUES (?, '', ?, ?, ?, ?)
463RETURNING id AS "id!: i64"
464"#,
465            controller_session_id,
466            status,
467            max_parallelism,
468            now,
469            now
470        )
471        .fetch_one(&self.0)
472        .await?;
473
474        Ok(row.id)
475    }
476
477    async fn upsert_orchestration_task(
478        &self,
479        task: PersistedOrchestrationTask,
480    ) -> Result<i64, DbError> {
481        let PersistedOrchestrationTask {
482            acceptance_criteria,
483            kind,
484            merge_position,
485            prompt,
486            session_orchestration_id,
487            task_key,
488            title,
489            touched_areas,
490        } = task;
491        kind.parse::<ag_session::OrchestrationTaskKind>()
492            .map_err(|_| DbError::InvalidData {
493                entity: "orchestration task kind",
494                reason: format!("unknown persisted kind `{kind}`"),
495            })?;
496        let now = self.now();
497        let mut transaction = self
498            .0
499            .begin()
500            .await
501            .db_context("upsert orchestration task")?;
502
503        let row = sqlx::query!(
504            r#"
505INSERT INTO session_orchestration_task (
506    session_orchestration_id,
507    task_key,
508    title,
509    prompt,
510    touched_areas,
511    acceptance_criteria,
512    kind,
513    merge_position,
514    status,
515    created_at,
516    updated_at
517)
518VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'Planned', ?, ?)
519ON CONFLICT(session_orchestration_id, task_key) DO UPDATE
520SET title = excluded.title,
521    prompt = excluded.prompt,
522    kind = excluded.kind,
523    touched_areas = excluded.touched_areas,
524    acceptance_criteria = excluded.acceptance_criteria,
525    merge_position = excluded.merge_position,
526    status = 'Planned',
527    child_session_id = NULL,
528    continuation_prompt = NULL,
529    review_iteration = 0,
530    research_report = NULL,
531    result_summary = NULL,
532    verification_reason = NULL,
533    verification_verdict = NULL,
534    last_error = NULL,
535    updated_at = excluded.updated_at
536RETURNING id AS "id!: i64"
537"#,
538            session_orchestration_id,
539            task_key,
540            title,
541            prompt,
542            touched_areas,
543            acceptance_criteria,
544            kind,
545            merge_position,
546            now,
547            now
548        )
549        .fetch_one(&mut *transaction)
550        .await
551        .db_context("upsert orchestration task")?;
552
553        sqlx::query!(
554            r"
555UPDATE session
556SET orchestration_task_id = NULL,
557    updated_at = ?
558WHERE orchestration_task_id = ?
559",
560            now,
561            row.id
562        )
563        .execute(&mut *transaction)
564        .await
565        .db_context("upsert orchestration task")?;
566
567        transaction
568            .commit()
569            .await
570            .db_context("upsert orchestration task")?;
571
572        Ok(row.id)
573    }
574
575    async fn load_orchestration_for_controller(
576        &self,
577        controller_session_id: &str,
578    ) -> Result<Option<SessionOrchestrationRow>, DbError> {
579        let row = sqlx::query_as!(
580            SessionOrchestrationRow,
581            r#"
582SELECT orchestration.id AS "id!: i64",
583       session.project_id AS "controller_project_id!: i64",
584       orchestration.controller_session_id,
585       orchestration.goal_statement,
586       orchestration.relayed_question_task_id,
587       orchestration.status,
588       orchestration.max_parallelism,
589       orchestration.verification_generation
590FROM session_orchestration AS orchestration
591INNER JOIN session
592ON session.id = orchestration.controller_session_id
593WHERE orchestration.controller_session_id = ?
594ORDER BY orchestration.id DESC
595LIMIT 1
596"#,
597            controller_session_id
598        )
599        .fetch_optional(&self.0)
600        .await?;
601        if let Some(row) = &row {
602            status::validate_orchestration(&row.status)?;
603        }
604
605        Ok(row)
606    }
607
608    async fn load_active_orchestrations(&self) -> Result<Vec<SessionOrchestrationRow>, DbError> {
609        let rows = sqlx::query_as!(
610            SessionOrchestrationRow,
611            r#"
612SELECT orchestration.id AS "id!: i64",
613       session.project_id AS "controller_project_id!: i64",
614       orchestration.controller_session_id,
615       orchestration.goal_statement,
616       orchestration.relayed_question_task_id,
617       orchestration.status,
618       orchestration.max_parallelism,
619       orchestration.verification_generation
620FROM session_orchestration AS orchestration
621INNER JOIN session
622ON session.id = orchestration.controller_session_id
623WHERE orchestration.status IN (
624    'AwaitingApproval',
625    'Running',
626    'Verifying',
627    'AwaitingIntegration',
628    'Integrating',
629    'Canceling'
630)
631ORDER BY orchestration.id
632"#
633        )
634        .fetch_all(&self.0)
635        .await?;
636        for row in &rows {
637            status::validate_orchestration(&row.status)?;
638        }
639
640        Ok(rows)
641    }
642
643    async fn load_recoverable_focused_review_session_ids(
644        &self,
645        project_id: i64,
646    ) -> Result<Vec<String>, DbError> {
647        let session_ids = sqlx::query_scalar!(
648            r#"
649SELECT child.id AS "id!: String"
650FROM session_orchestration_task AS task
651INNER JOIN session_orchestration AS orchestration
652ON orchestration.id = task.session_orchestration_id
653INNER JOIN session AS child
654ON child.id = task.child_session_id
655WHERE child.project_id = ?
656  AND orchestration.status IN ('AwaitingApproval', 'Running')
657  AND task.status = 'Reviewing'
658  AND child.status IN ('Review', 'AgentReview')
659  AND (
660      child.focused_review_status IS NULL
661      OR child.focused_review_status = 'Pending'
662  )
663ORDER BY task.id
664"#,
665            project_id
666        )
667        .fetch_all(&self.0)
668        .await?;
669
670        Ok(session_ids)
671    }
672
673    async fn load_session_metadata_for_project(
674        &self,
675        project_id: i64,
676    ) -> Result<Vec<SessionOrchestrationMetadataRow>, DbError> {
677        let rows = sqlx::query_as!(
678            SessionOrchestrationMetadataRow,
679            r#"
680WITH latest_orchestration_id AS (
681    SELECT controller_session_id,
682           MAX(id) AS orchestration_id
683    FROM session_orchestration
684    GROUP BY controller_session_id
685),
686latest_orchestration AS (
687    SELECT orchestration.id,
688           orchestration.controller_session_id,
689           orchestration.status
690    FROM session_orchestration AS orchestration
691    INNER JOIN latest_orchestration_id AS latest
692    ON latest.orchestration_id = orchestration.id
693),
694controller_metadata AS (
695    SELECT orchestration.controller_session_id AS session_id,
696           orchestration.status AS orchestration_status,
697           COALESCE(SUM(
698               CASE WHEN task.status IN (
699                        'Creating',
700                        'Running',
701                        'Reviewing',
702                        'ReviewApplying',
703                        'ContinuationPending'
704                    )
705                    THEN 1 ELSE 0 END
706           ), 0) AS running_task_count,
707           COALESCE(SUM(
708               CASE WHEN task.status = 'WaitingForInput' THEN 1 ELSE 0 END
709           ), 0) AS waiting_task_count
710    FROM latest_orchestration AS orchestration
711    LEFT JOIN session_orchestration_task AS task
712    ON task.session_orchestration_id = orchestration.id
713    GROUP BY orchestration.id
714),
715child_metadata AS (
716    SELECT task.child_session_id AS session_id,
717           orchestration.controller_session_id
718    FROM session_orchestration_task AS task
719    INNER JOIN session_orchestration AS orchestration
720    ON orchestration.id = task.session_orchestration_id
721    WHERE task.child_session_id IS NOT NULL
722)
723SELECT session.id AS "session_id!: String",
724       child_metadata.controller_session_id,
725       controller_metadata.orchestration_status,
726       COALESCE(controller_metadata.running_task_count, 0) AS "running_task_count!: i64",
727       COALESCE(controller_metadata.waiting_task_count, 0) AS "waiting_task_count!: i64"
728FROM session
729LEFT JOIN controller_metadata
730ON controller_metadata.session_id = session.id
731LEFT JOIN child_metadata
732ON child_metadata.session_id = session.id
733WHERE session.project_id = ?
734  AND (
735      controller_metadata.session_id IS NOT NULL
736      OR child_metadata.session_id IS NOT NULL
737  )
738ORDER BY session.id
739"#,
740            project_id
741        )
742        .fetch_all(&self.0)
743        .await?;
744        for row in &rows {
745            if let Some(orchestration_status) = &row.orchestration_status {
746                status::validate_orchestration(orchestration_status)?;
747            }
748        }
749
750        Ok(rows)
751    }
752
753    async fn load_orchestration_tasks(
754        &self,
755        session_orchestration_id: i64,
756    ) -> Result<Vec<SessionOrchestrationTaskRow>, DbError> {
757        let rows = sqlx::query_as!(
758            SessionOrchestrationTaskRow,
759            r#"
760SELECT task.id AS "id!: i64",
761       task.acceptance_criteria,
762       task.area_violations,
763       task.areas_compliant AS "areas_compliant?: bool",
764       task.attempt_count,
765       COALESCE(child.added_lines, 0) AS "child_added_lines!: i64",
766       (
767           SELECT message.content
768           FROM session_message AS message
769           WHERE message.session_id = task.child_session_id
770             AND message.kind = 'assistant_answer'
771           ORDER BY message.position DESC
772           LIMIT 1
773       ) AS child_answer,
774       COALESCE(child.deleted_lines, 0) AS "child_deleted_lines!: i64",
775       child.focused_review_status AS child_focused_review_status,
776       child.focused_review_text AS child_focused_review_text,
777       child.has_diff AS "child_has_diff?: bool",
778       COALESCE(child.input_tokens, 0) AS "child_input_tokens!: i64",
779       COALESCE(child.output_tokens, 0) AS "child_output_tokens!: i64",
780       task.child_session_id,
781       child.status AS child_status,
782       child.questions AS child_questions,
783       child.summary AS child_summary,
784       task.continuation_generation,
785       task.continuation_prompt,
786       task.infrastructure_retry_count,
787       task.kind,
788       task.last_error,
789       task.merge_position,
790       task.prompt,
791       task.research_report,
792       task.result_summary,
793       task.review_iteration,
794       task.status,
795       task.task_key,
796       task.touched_areas,
797       task.title,
798       task.verification_reason,
799       task.verification_verdict
800FROM session_orchestration_task AS task
801LEFT JOIN session AS child
802ON child.id = task.child_session_id
803WHERE task.session_orchestration_id = ?
804ORDER BY task.merge_position, task.id
805"#,
806            session_orchestration_id
807        )
808        .fetch_all(&self.0)
809        .await?;
810        for row in &rows {
811            status::validate_orchestration_task(&row.status)?;
812            row.kind
813                .parse::<ag_session::OrchestrationTaskKind>()
814                .map_err(|_| DbError::InvalidData {
815                    entity: "orchestration task kind",
816                    reason: format!("unknown persisted kind `{}`", row.kind),
817                })?;
818            if let Some(child_status) = &row.child_status {
819                status::validate_session(child_status)?;
820            }
821        }
822
823        Ok(rows)
824    }
825
826    async fn load_orchestration_integration_approach(&self, id: i64) -> Result<String, DbError> {
827        sqlx::query_scalar::<_, String>(
828            "SELECT integration_approach FROM session_orchestration WHERE id = ?",
829        )
830        .bind(id)
831        .fetch_one(&self.0)
832        .await
833        .db_context("load orchestration integration approach")
834    }
835
836    async fn load_orchestration_task_scope_for_child(
837        &self,
838        child_session_id: &str,
839    ) -> Result<Option<OrchestrationTaskScopeRow>, DbError> {
840        sqlx::query_as!(
841            OrchestrationTaskScopeRow,
842            r#"
843SELECT child.base_branch,
844       task.id AS "id!: i64",
845       task.touched_areas
846FROM session_orchestration_task AS task
847INNER JOIN session AS child
848ON child.id = task.child_session_id
849WHERE child.id = ?
850  AND child.role = 'OrchestrationWorker'
851  AND task.kind = 'Implementation'
852"#,
853            child_session_id
854        )
855        .fetch_optional(&self.0)
856        .await
857        .db_context("load orchestration task scope for child")
858    }
859
860    async fn load_child_session_id_for_task(
861        &self,
862        task_id: i64,
863    ) -> Result<Option<String>, DbError> {
864        let row = sqlx::query!(
865            r#"
866SELECT id AS "id!: String"
867FROM session
868WHERE orchestration_task_id = ?
869"#,
870            task_id
871        )
872        .fetch_optional(&self.0)
873        .await?;
874
875        Ok(row.map(|row| row.id))
876    }
877
878    async fn begin_orchestration_cancellation(&self, id: i64) -> Result<bool, DbError> {
879        let now = self.now();
880
881        let result = sqlx::query!(
882            r"
883UPDATE session_orchestration
884SET status = 'Canceling',
885    updated_at = ?
886WHERE id = ?
887  AND status IN (
888      'AwaitingApproval',
889      'Running',
890      'Verifying',
891      'AwaitingIntegration',
892      'Integrating',
893      'Canceling'
894  )
895",
896            now,
897            id
898        )
899        .execute(&self.0)
900        .await
901        .db_context("begin orchestration cancellation")?;
902
903        Ok(result.rows_affected() == 1)
904    }
905
906    async fn claim_orchestration_task(&self, id: i64) -> Result<bool, DbError> {
907        let now = self.now();
908
909        let result = sqlx::query!(
910            r"
911UPDATE session_orchestration_task
912SET status = 'Creating',
913    last_error = NULL,
914    updated_at = ?
915WHERE id = ?
916  AND status = 'Planned'
917  AND EXISTS (
918      SELECT 1
919      FROM session_orchestration
920      WHERE session_orchestration.id = session_orchestration_task.session_orchestration_id
921        AND session_orchestration.status = 'Running'
922  )
923",
924            now,
925            id
926        )
927        .execute(&self.0)
928        .await
929        .db_context("claim orchestration task")?;
930
931        Ok(result.rows_affected() == 1)
932    }
933
934    async fn claim_orchestration_review_application(
935        &self,
936        id: i64,
937        prompt: &str,
938        iteration_limit: i64,
939    ) -> Result<bool, DbError> {
940        let now = self.now();
941        let mut transaction = self
942            .0
943            .begin()
944            .await
945            .db_context("claim orchestration review application")?;
946        let claim = sqlx::query!(
947            r"
948UPDATE session_orchestration_task
949SET continuation_generation = continuation_generation + 1,
950    continuation_prompt = ?,
951    review_iteration = review_iteration + 1,
952    status = 'ReviewApplying',
953    updated_at = ?
954WHERE id = ?
955  AND status = 'Reviewing'
956  AND review_iteration < ?
957  AND child_session_id IS NOT NULL
958",
959            prompt,
960            now,
961            id,
962            iteration_limit
963        )
964        .execute(&mut *transaction)
965        .await
966        .db_context("claim orchestration review application")?;
967        if claim.rows_affected() == 0 {
968            transaction
969                .rollback()
970                .await
971                .db_context("claim orchestration review application")?;
972
973            return Ok(false);
974        }
975
976        sqlx::query!(
977            r"
978UPDATE session
979SET focused_review_status = NULL,
980    focused_review_diff_hash = NULL,
981    focused_review_text = NULL,
982    updated_at = ?
983WHERE id = (
984    SELECT child_session_id
985    FROM session_orchestration_task
986    WHERE id = ?
987)
988",
989            now,
990            id
991        )
992        .execute(&mut *transaction)
993        .await
994        .db_context("claim orchestration review application")?;
995        transaction
996            .commit()
997            .await
998            .db_context("claim orchestration review application")?;
999
1000        Ok(true)
1001    }
1002
1003    async fn claim_orchestration_rollup(&self, id: i64) -> Result<bool, DbError> {
1004        let now = self.now();
1005
1006        let result = sqlx::query!(
1007            r"
1008UPDATE session_orchestration
1009SET status = 'Verifying',
1010    verification_generation = verification_generation + 1,
1011    updated_at = ?
1012WHERE id = ?
1013  AND status = 'Running'
1014",
1015            now,
1016            id
1017        )
1018        .execute(&self.0)
1019        .await
1020        .db_context("claim orchestration rollup")?;
1021
1022        Ok(result.rows_affected() == 1)
1023    }
1024
1025    async fn complete_orchestration_rollup(&self, id: i64) -> Result<bool, DbError> {
1026        self.transition_orchestration_and_tasks(
1027            id,
1028            OrchestrationTransition {
1029                context: "complete orchestration rollup",
1030                from_orchestration_status: "Verifying",
1031                from_task_status: "Ready",
1032                require_pass_verdict: true,
1033                to_orchestration_status: "AwaitingIntegration",
1034                to_task_status: "AwaitingIntegration",
1035            },
1036        )
1037        .await
1038    }
1039
1040    async fn record_orchestration_verdict(
1041        &self,
1042        id: i64,
1043        task_key: &str,
1044        is_pass: bool,
1045        reason: &str,
1046    ) -> Result<bool, DbError> {
1047        let now = self.now();
1048        let verdict = if is_pass { "Pass" } else { "Flag" };
1049        let result = sqlx::query!(
1050            r"
1051UPDATE session_orchestration_task
1052SET verification_reason = ?,
1053    verification_verdict = ?,
1054    updated_at = ?
1055WHERE session_orchestration_id = ?
1056  AND task_key = ?
1057  AND status IN ('Ready', 'Reported')
1058  AND EXISTS (
1059      SELECT 1
1060      FROM session_orchestration
1061      WHERE id = ?
1062        AND status = 'Verifying'
1063  )
1064",
1065            reason,
1066            verdict,
1067            now,
1068            id,
1069            task_key,
1070            id
1071        )
1072        .execute(&self.0)
1073        .await
1074        .db_context("record orchestration verdict")?;
1075
1076        Ok(result.rows_affected() == 1)
1077    }
1078
1079    async fn complete_orchestration_campaign(&self, id: i64) -> Result<bool, DbError> {
1080        let now = self.now();
1081        let mut transaction = self
1082            .0
1083            .begin()
1084            .await
1085            .db_context("complete orchestration campaign")?;
1086        let orchestration = sqlx::query!(
1087            r#"
1088UPDATE session_orchestration
1089SET status = 'Done',
1090    updated_at = ?
1091WHERE id = ?
1092  AND status IN ('AwaitingIntegration', 'Integrating')
1093RETURNING controller_session_id AS "controller_session_id!: String"
1094"#,
1095            now,
1096            id
1097        )
1098        .fetch_optional(&mut *transaction)
1099        .await
1100        .db_context("complete orchestration campaign")?;
1101        if let Some(orchestration) = &orchestration {
1102            sqlx::query!(
1103                r"
1104UPDATE session
1105SET status = 'Done',
1106    questions = '',
1107    updated_at = ?
1108WHERE id = ?
1109  AND role = 'Orchestrator'
1110  AND status IN ('Review', 'Question')
1111",
1112                now,
1113                orchestration.controller_session_id
1114            )
1115            .execute(&mut *transaction)
1116            .await
1117            .db_context("complete orchestration campaign")?;
1118        }
1119        transaction
1120            .commit()
1121            .await
1122            .db_context("complete orchestration campaign")?;
1123
1124        Ok(orchestration.is_some())
1125    }
1126
1127    async fn load_rollup_operation_status(
1128        &self,
1129        operation_id: &str,
1130    ) -> Result<Option<String>, DbError> {
1131        let row = sqlx::query!(
1132            r#"
1133SELECT status AS "status!: String"
1134FROM session_operation
1135WHERE id = ?
1136"#,
1137            operation_id
1138        )
1139        .fetch_optional(&self.0)
1140        .await?;
1141        if let Some(row) = &row {
1142            status::validate_operation(&row.status)?;
1143        }
1144
1145        Ok(row.map(|row| row.status))
1146    }
1147
1148    async fn update_orchestration_status(&self, id: i64, status: &str) -> Result<(), DbError> {
1149        status::validate_orchestration(status)?;
1150        let now = self.now();
1151
1152        sqlx::query!(
1153            r"
1154UPDATE session_orchestration
1155SET status = ?,
1156    updated_at = ?
1157WHERE id = ?
1158",
1159            status,
1160            now,
1161            id
1162        )
1163        .execute(&self.0)
1164        .await?;
1165
1166        Ok(())
1167    }
1168
1169    async fn approve_orchestration_plan(&self, id: i64) -> Result<bool, DbError> {
1170        self.transition_orchestration_and_tasks(
1171            id,
1172            OrchestrationTransition {
1173                context: "approve orchestration plan",
1174                from_orchestration_status: "AwaitingApproval",
1175                from_task_status: "Proposed",
1176                require_pass_verdict: false,
1177                to_orchestration_status: "Running",
1178                to_task_status: "Planned",
1179            },
1180        )
1181        .await
1182    }
1183
1184    async fn approve_orchestration_integration(
1185        &self,
1186        id: i64,
1187        approach: IntegrationApproach,
1188    ) -> Result<bool, DbError> {
1189        let approach = approach.to_string();
1190        let now = self.now();
1191        let result = sqlx::query(
1192            r"
1193UPDATE session_orchestration
1194SET integration_approach = ?,
1195    status = 'Integrating',
1196    updated_at = ?
1197WHERE id = ?
1198  AND status = 'AwaitingIntegration'
1199",
1200        )
1201        .bind(approach)
1202        .bind(now)
1203        .bind(id)
1204        .execute(&self.0)
1205        .await
1206        .db_context("approve orchestration integration")?;
1207
1208        Ok(result.rows_affected() == 1)
1209    }
1210
1211    async fn update_orchestration_plan(
1212        &self,
1213        id: i64,
1214        goal_statement: &str,
1215        max_parallelism: i64,
1216    ) -> Result<(), DbError> {
1217        let now = self.now();
1218
1219        sqlx::query!(
1220            r"
1221UPDATE session_orchestration
1222SET goal_statement = ?,
1223    max_parallelism = ?,
1224    updated_at = ?
1225WHERE id = ?
1226  AND status = 'AwaitingApproval'
1227",
1228            goal_statement,
1229            max_parallelism,
1230            now,
1231            id
1232        )
1233        .execute(&self.0)
1234        .await?;
1235
1236        Ok(())
1237    }
1238
1239    async fn queue_orchestration_continuation(
1240        &self,
1241        id: i64,
1242        prompt: &str,
1243        acceptance_criteria: &str,
1244        touched_areas: &str,
1245    ) -> Result<bool, DbError> {
1246        let now = self.now();
1247        let mut transaction = self
1248            .0
1249            .begin()
1250            .await
1251            .db_context("queue orchestration continuation")?;
1252        let result = sqlx::query!(
1253            r"
1254UPDATE session_orchestration_task
1255SET acceptance_criteria = ?,
1256    area_violations = '[]',
1257    areas_compliant = NULL,
1258    continuation_generation = continuation_generation + 1,
1259    continuation_prompt = ?,
1260    review_iteration = 0,
1261    status = 'ContinuationPending',
1262    touched_areas = ?,
1263    result_summary = NULL,
1264    verification_verdict = NULL,
1265    verification_reason = NULL,
1266    last_error = NULL,
1267    updated_at = ?
1268WHERE id = ?
1269  AND child_session_id IS NOT NULL
1270  AND status IN ('Ready', 'AwaitingIntegration', 'IntegrationFailed')
1271",
1272            acceptance_criteria,
1273            prompt,
1274            touched_areas,
1275            now,
1276            id
1277        )
1278        .execute(&mut *transaction)
1279        .await?;
1280
1281        if result.rows_affected() == 1 {
1282            sqlx::query!(
1283                r"
1284UPDATE session
1285SET focused_review_status = NULL,
1286    focused_review_diff_hash = NULL,
1287    focused_review_text = NULL,
1288    updated_at = ?
1289WHERE id = (
1290    SELECT child_session_id
1291    FROM session_orchestration_task
1292    WHERE id = ?
1293)
1294",
1295                now,
1296                id
1297            )
1298            .execute(&mut *transaction)
1299            .await?;
1300        }
1301        transaction.commit().await?;
1302
1303        Ok(result.rows_affected() == 1)
1304    }
1305
1306    async fn reset_orchestration_verification(&self, id: i64) -> Result<(), DbError> {
1307        let now = self.now();
1308
1309        sqlx::query!(
1310            r"
1311UPDATE session_orchestration_task
1312SET status = 'Ready',
1313    verification_reason = NULL,
1314    verification_verdict = NULL,
1315    updated_at = ?
1316WHERE session_orchestration_id = ?
1317  AND status = 'AwaitingIntegration'
1318",
1319            now,
1320            id
1321        )
1322        .execute(&self.0)
1323        .await
1324        .db_context("reset orchestration verification")?;
1325
1326        Ok(())
1327    }
1328
1329    async fn record_orchestration_spawn_failure(
1330        &self,
1331        id: i64,
1332        error: &str,
1333        retry_limit: i64,
1334    ) -> Result<String, DbError> {
1335        let now = self.now();
1336        let mut transaction = self
1337            .0
1338            .begin()
1339            .await
1340            .db_context("record orchestration spawn failure")?;
1341        sqlx::query!(
1342            r"
1343UPDATE session
1344SET orchestration_task_id = NULL,
1345    updated_at = ?
1346WHERE orchestration_task_id = ?
1347",
1348            now,
1349            id
1350        )
1351        .execute(&mut *transaction)
1352        .await
1353        .db_context("record orchestration spawn failure")?;
1354        let row = sqlx::query!(
1355            r#"
1356UPDATE session_orchestration_task
1357SET child_session_id = NULL,
1358    infrastructure_retry_count = infrastructure_retry_count + 1,
1359    status = CASE
1360        WHEN infrastructure_retry_count < ? THEN 'Planned'
1361        ELSE 'Failed'
1362    END,
1363    last_error = ?,
1364    updated_at = ?
1365WHERE id = ?
1366RETURNING status AS "status!: String"
1367"#,
1368            retry_limit,
1369            error,
1370            now,
1371            id
1372        )
1373        .fetch_one(&mut *transaction)
1374        .await
1375        .db_context("record orchestration spawn failure")?;
1376        transaction
1377            .commit()
1378            .await
1379            .db_context("record orchestration spawn failure")?;
1380        status::validate_orchestration_task(&row.status)?;
1381
1382        Ok(row.status)
1383    }
1384
1385    async fn detach_orchestration_child(&self, child_session_id: &str) -> Result<bool, DbError> {
1386        let now = self.now();
1387        let mut transaction = self
1388            .0
1389            .begin()
1390            .await
1391            .db_context("detach orchestration child")?;
1392        let task = sqlx::query!(
1393            r#"
1394SELECT orchestration_task_id AS "task_id!: i64"
1395FROM session
1396WHERE id = ?
1397  AND role = 'OrchestrationWorker'
1398  AND orchestration_task_id IS NOT NULL
1399"#,
1400            child_session_id
1401        )
1402        .fetch_optional(&mut *transaction)
1403        .await
1404        .db_context("detach orchestration child")?;
1405        let Some(task) = task else {
1406            transaction
1407                .commit()
1408                .await
1409                .db_context("detach orchestration child")?;
1410
1411            return Ok(false);
1412        };
1413
1414        sqlx::query!(
1415            r"
1416UPDATE session_orchestration_task
1417SET child_session_id = NULL,
1418    status = 'Detached',
1419    updated_at = ?
1420WHERE id = ?
1421",
1422            now,
1423            task.task_id
1424        )
1425        .execute(&mut *transaction)
1426        .await
1427        .db_context("detach orchestration child")?;
1428        sqlx::query!(
1429            r"
1430UPDATE session
1431SET orchestration_task_id = NULL,
1432    role = 'Worker',
1433    updated_at = ?
1434WHERE id = ?
1435",
1436            now,
1437            child_session_id
1438        )
1439        .execute(&mut *transaction)
1440        .await
1441        .db_context("detach orchestration child")?;
1442        transaction
1443            .commit()
1444            .await
1445            .db_context("detach orchestration child")?;
1446
1447        Ok(true)
1448    }
1449
1450    async fn surface_orchestration_questions(
1451        &self,
1452        session_orchestration_id: i64,
1453        task_id: i64,
1454        questions: &str,
1455    ) -> Result<bool, DbError> {
1456        let now = self.now();
1457        let mut transaction = self
1458            .0
1459            .begin()
1460            .await
1461            .db_context("surface orchestration questions")?;
1462        let claim = sqlx::query!(
1463            r"
1464UPDATE session_orchestration
1465SET relayed_question_task_id = ?,
1466    updated_at = ?
1467WHERE id = ?
1468  AND relayed_question_task_id IS NULL
1469  AND EXISTS (
1470      SELECT 1
1471      FROM session_orchestration_task AS task
1472      WHERE task.id = ?
1473        AND task.session_orchestration_id = session_orchestration.id
1474        AND task.status = 'WaitingForInput'
1475        AND task.child_session_id IS NOT NULL
1476  )
1477  AND EXISTS (
1478      SELECT 1
1479      FROM session AS controller
1480      WHERE controller.id = session_orchestration.controller_session_id
1481        AND controller.role = 'Orchestrator'
1482        AND controller.status IN ('Review', 'Question')
1483        AND COALESCE(controller.questions, '') = ''
1484  )
1485",
1486            task_id,
1487            now,
1488            session_orchestration_id,
1489            task_id
1490        )
1491        .execute(&mut *transaction)
1492        .await
1493        .db_context("surface orchestration questions")?;
1494        if claim.rows_affected() == 0 {
1495            transaction
1496                .rollback()
1497                .await
1498                .db_context("surface orchestration questions")?;
1499
1500            return Ok(false);
1501        }
1502        sqlx::query!(
1503            r"
1504UPDATE session
1505SET questions = ?,
1506    status = 'Question',
1507    updated_at = ?
1508WHERE id = (
1509    SELECT controller_session_id
1510    FROM session_orchestration
1511    WHERE id = ?
1512)
1513",
1514            questions,
1515            now,
1516            session_orchestration_id
1517        )
1518        .execute(&mut *transaction)
1519        .await
1520        .db_context("surface orchestration questions")?;
1521        transaction
1522            .commit()
1523            .await
1524            .db_context("surface orchestration questions")?;
1525
1526        Ok(true)
1527    }
1528
1529    async fn clear_orchestration_questions(
1530        &self,
1531        session_orchestration_id: i64,
1532    ) -> Result<(), DbError> {
1533        let now = self.now();
1534        let mut transaction = self
1535            .0
1536            .begin()
1537            .await
1538            .db_context("clear orchestration questions")?;
1539        sqlx::query!(
1540            r"
1541UPDATE session
1542SET questions = '',
1543    status = 'Review',
1544    updated_at = ?
1545WHERE id = (
1546    SELECT controller_session_id
1547    FROM session_orchestration
1548    WHERE id = ?
1549      AND relayed_question_task_id IS NOT NULL
1550)
1551  AND role = 'Orchestrator'
1552  AND status = 'Question'
1553",
1554            now,
1555            session_orchestration_id
1556        )
1557        .execute(&mut *transaction)
1558        .await
1559        .db_context("clear orchestration questions")?;
1560        sqlx::query!(
1561            r"
1562UPDATE session_orchestration
1563SET relayed_question_task_id = NULL,
1564    updated_at = ?
1565WHERE id = ?
1566  AND relayed_question_task_id IS NOT NULL
1567",
1568            now,
1569            session_orchestration_id
1570        )
1571        .execute(&mut *transaction)
1572        .await
1573        .db_context("clear orchestration questions")?;
1574        transaction
1575            .commit()
1576            .await
1577            .db_context("clear orchestration questions")?;
1578
1579        Ok(())
1580    }
1581
1582    async fn link_orchestration_task_child(
1583        &self,
1584        id: i64,
1585        child_session_id: &str,
1586    ) -> Result<bool, DbError> {
1587        let now = self.now();
1588
1589        let result = sqlx::query!(
1590            r"
1591UPDATE session_orchestration_task
1592SET child_session_id = ?,
1593    status = 'Running',
1594    attempt_count = attempt_count + 1,
1595    updated_at = ?
1596WHERE id = ?
1597  AND status = 'Creating'
1598  AND EXISTS (
1599      SELECT 1
1600      FROM session_orchestration
1601      WHERE session_orchestration.id = session_orchestration_task.session_orchestration_id
1602        AND session_orchestration.status = 'Running'
1603  )
1604",
1605            child_session_id,
1606            now,
1607            id
1608        )
1609        .execute(&self.0)
1610        .await?;
1611
1612        Ok(result.rows_affected() == 1)
1613    }
1614
1615    async fn update_orchestration_task_status(
1616        &self,
1617        id: i64,
1618        status: &str,
1619        last_error: Option<String>,
1620    ) -> Result<(), DbError> {
1621        status::validate_orchestration_task(status)?;
1622        let now = self.now();
1623
1624        sqlx::query!(
1625            r"
1626UPDATE session_orchestration_task
1627SET status = ?,
1628    last_error = ?,
1629    updated_at = ?
1630WHERE id = ?
1631",
1632            status,
1633            last_error,
1634            now,
1635            id
1636        )
1637        .execute(&self.0)
1638        .await?;
1639
1640        Ok(())
1641    }
1642
1643    async fn update_orchestration_task_result_summary(
1644        &self,
1645        id: i64,
1646        result_summary: &str,
1647    ) -> Result<(), DbError> {
1648        let now = self.now();
1649
1650        sqlx::query!(
1651            r"
1652UPDATE session_orchestration_task
1653SET result_summary = ?,
1654    updated_at = ?
1655WHERE id = ?
1656",
1657            result_summary,
1658            now,
1659            id
1660        )
1661        .execute(&self.0)
1662        .await?;
1663
1664        Ok(())
1665    }
1666
1667    async fn update_orchestration_task_research_report(
1668        &self,
1669        id: i64,
1670        research_report: &str,
1671    ) -> Result<(), DbError> {
1672        let now = self.now();
1673
1674        sqlx::query!(
1675            r"
1676UPDATE session_orchestration_task
1677SET research_report = ?,
1678    updated_at = ?
1679WHERE id = ?
1680  AND kind = 'Research'
1681",
1682            research_report,
1683            now,
1684            id
1685        )
1686        .execute(&self.0)
1687        .await?;
1688
1689        Ok(())
1690    }
1691
1692    async fn update_orchestration_task_area_compliance(
1693        &self,
1694        id: i64,
1695        areas_compliant: Option<bool>,
1696        area_violations: &str,
1697    ) -> Result<(), DbError> {
1698        let now = self.now();
1699
1700        sqlx::query!(
1701            r"
1702UPDATE session_orchestration_task
1703SET area_violations = ?,
1704    areas_compliant = ?,
1705    updated_at = ?
1706WHERE id = ?
1707",
1708            area_violations,
1709            areas_compliant,
1710            now,
1711            id
1712        )
1713        .execute(&self.0)
1714        .await?;
1715
1716        Ok(())
1717    }
1718}
1719
1720#[cfg(test)]
1721mod tests {
1722    use ag_agent::{AgentKind, ReasoningLevel, SpeedMode};
1723    use ag_session::{
1724        FocusedReviewStatus, OrchestrationStatus, OrchestrationTaskKind, OrchestrationTaskStatus,
1725        SessionMessageKind,
1726    };
1727
1728    use super::*;
1729    use crate::{AppRepositories, PersistedSessionCreation, SessionRow};
1730
1731    /// Inserts one project plus one controller session and returns the
1732    /// repository bundle ready for orchestration persistence assertions.
1733    async fn controller_fixture() -> AppRepositories {
1734        controller_fixture_with_pool().await.0
1735    }
1736
1737    async fn controller_fixture_with_pool() -> (AppRepositories, SqlitePool) {
1738        let (database, pool) = AppRepositories::in_memory_with_pool()
1739            .await
1740            .expect("db should open");
1741        let project_id = database
1742            .projects()
1743            .upsert_project("/tmp/project", None)
1744            .await
1745            .expect("failed to upsert project");
1746        database
1747            .sessions()
1748            .insert_session_with_agent(PersistedSessionCreation {
1749                agent: "codex",
1750                base_branch: "main",
1751                id: "controller",
1752                is_draft: false,
1753                model: AgentKind::Codex.default_model().as_str(),
1754                orchestration_task_id: None,
1755                parent_session_id: None,
1756                personality_id: None,
1757                project_id,
1758                reasoning_level: ReasoningLevel::default(),
1759                role: Some("Orchestrator"),
1760                speed_mode: SpeedMode::Normal,
1761                status: "Review",
1762            })
1763            .await
1764            .expect("failed to insert controller session");
1765
1766        (database, pool)
1767    }
1768
1769    #[tokio::test]
1770    async fn hydration_rejects_unknown_orchestration_status() {
1771        // Arrange
1772        let (database, pool) = AppRepositories::in_memory_with_pool()
1773            .await
1774            .expect("db should open");
1775        let project_id = database
1776            .projects()
1777            .upsert_project("/tmp/invalid-orchestration", None)
1778            .await
1779            .expect("failed to upsert project");
1780        database
1781            .sessions()
1782            .insert_session("controller", "gpt-5.6-sol", "main", "Draft", project_id)
1783            .await
1784            .expect("failed to insert controller session");
1785        let orchestration_id = database
1786            .orchestrations()
1787            .insert_orchestration("controller", "Running", 1)
1788            .await
1789            .expect("failed to insert orchestration");
1790        sqlx::query("UPDATE session_orchestration SET status = 'Unknown' WHERE id = ?")
1791            .bind(orchestration_id)
1792            .execute(&pool)
1793            .await
1794            .expect("failed to corrupt orchestration status");
1795
1796        // Act
1797        let error = database
1798            .orchestrations()
1799            .load_orchestration_for_controller("controller")
1800            .await
1801            .expect_err("invalid status should fail hydration");
1802
1803        // Assert
1804        assert!(matches!(
1805            error,
1806            DbError::InvalidStatus {
1807                entity: "orchestration",
1808                value,
1809            } if value == "Unknown"
1810        ));
1811    }
1812
1813    #[tokio::test]
1814    async fn rollup_recovery_failures_report_semantic_operation_context() {
1815        // Arrange
1816        let (database, pool) = AppRepositories::in_memory_with_pool()
1817            .await
1818            .expect("db should open");
1819        sqlx::query("DROP TABLE session_orchestration")
1820            .execute(&pool)
1821            .await
1822            .expect("failed to drop orchestration table");
1823
1824        // Act
1825        let claim_error = database
1826            .orchestrations()
1827            .claim_orchestration_rollup(1)
1828            .await
1829            .expect_err("claim should fail");
1830        let completion_error = database
1831            .orchestrations()
1832            .complete_orchestration_rollup(1)
1833            .await
1834            .expect_err("completion should fail");
1835
1836        // Assert
1837        assert!(matches!(
1838            claim_error,
1839            DbError::QueryContext {
1840                operation: "claim orchestration rollup",
1841                ..
1842            }
1843        ));
1844        assert!(matches!(
1845            completion_error,
1846            DbError::QueryContext {
1847                operation: "complete orchestration rollup",
1848                ..
1849            }
1850        ));
1851    }
1852
1853    /// Builds one planned task payload for `orchestration_id`.
1854    fn planned_task(session_orchestration_id: i64, task_key: &str) -> PersistedOrchestrationTask {
1855        PersistedOrchestrationTask {
1856            acceptance_criteria: format!(r#"["Complete {task_key}"]"#),
1857            kind: "Implementation".to_string(),
1858            merge_position: 0,
1859            prompt: format!("Complete {task_key}"),
1860            session_orchestration_id,
1861            task_key: task_key.to_string(),
1862            title: format!("Task {task_key}"),
1863            touched_areas: format!(r#"["crates/{task_key}/"]"#),
1864        }
1865    }
1866
1867    /// Persists one worker session carrying the durable reverse task link.
1868    async fn insert_orchestration_child(
1869        database: &AppRepositories,
1870        project_id: i64,
1871        session_id: &str,
1872        task_id: i64,
1873    ) {
1874        database
1875            .sessions()
1876            .insert_session_with_agent(PersistedSessionCreation {
1877                agent: "codex",
1878                base_branch: "main",
1879                id: session_id,
1880                is_draft: false,
1881                model: AgentKind::Codex.default_model().as_str(),
1882                orchestration_task_id: Some(task_id),
1883                parent_session_id: None,
1884                personality_id: None,
1885                project_id,
1886                reasoning_level: ReasoningLevel::default(),
1887                role: Some("OrchestrationWorker"),
1888                speed_mode: SpeedMode::Normal,
1889                status: "Review",
1890            })
1891            .await
1892            .expect("failed to insert orchestration child");
1893    }
1894
1895    async fn insert_waiting_orchestration_task(
1896        database: &AppRepositories,
1897        project_id: i64,
1898        session_orchestration_id: i64,
1899        task_key: &str,
1900        child_session_id: &str,
1901    ) -> i64 {
1902        let task_id = database
1903            .orchestrations()
1904            .upsert_orchestration_task(planned_task(session_orchestration_id, task_key))
1905            .await
1906            .expect("failed to insert waiting task");
1907        assert!(
1908            database
1909                .orchestrations()
1910                .claim_orchestration_task(task_id)
1911                .await
1912                .expect("failed to claim waiting task")
1913        );
1914        insert_orchestration_child(database, project_id, child_session_id, task_id).await;
1915        assert!(
1916            database
1917                .orchestrations()
1918                .link_orchestration_task_child(task_id, child_session_id)
1919                .await
1920                .expect("failed to link waiting child")
1921        );
1922        database
1923            .orchestrations()
1924            .update_orchestration_task_status(
1925                task_id,
1926                &OrchestrationTaskStatus::WaitingForInput.to_string(),
1927                None,
1928            )
1929            .await
1930            .expect("failed to wait for child question");
1931
1932        task_id
1933    }
1934
1935    async fn surface_and_clear_orchestration_questions(
1936        database: &AppRepositories,
1937        session_orchestration_id: i64,
1938        task_id: i64,
1939        questions: &str,
1940    ) -> bool {
1941        let surfaced = database
1942            .orchestrations()
1943            .surface_orchestration_questions(session_orchestration_id, task_id, questions)
1944            .await
1945            .expect("failed to surface child question");
1946        database
1947            .orchestrations()
1948            .clear_orchestration_questions(session_orchestration_id)
1949            .await
1950            .expect("failed to clear child question");
1951
1952        surfaced
1953    }
1954
1955    async fn load_detached_campaign_state(
1956        database: &AppRepositories,
1957        orchestration_id: i64,
1958    ) -> (SessionOrchestrationTaskRow, SessionRow, SessionRow) {
1959        let task = database
1960            .orchestrations()
1961            .load_orchestration_tasks(orchestration_id)
1962            .await
1963            .expect("failed to load detached task")
1964            .remove(0);
1965        let child = database
1966            .sessions()
1967            .load_session("child-alpha")
1968            .await
1969            .expect("failed to load detached child")
1970            .expect("detached child should exist");
1971        let controller = database
1972            .sessions()
1973            .load_session("controller")
1974            .await
1975            .expect("failed to load controller")
1976            .expect("controller should exist");
1977
1978        (task, child, controller)
1979    }
1980
1981    #[tokio::test]
1982    /// Persists one plan before any child exists and loads it back for the
1983    /// owning controller session.
1984    async fn test_planned_orchestration_round_trips_before_any_child_exists() {
1985        // Arrange
1986        let database = controller_fixture().await;
1987
1988        // Act
1989        let orchestration_id = database
1990            .orchestrations()
1991            .insert_orchestration(
1992                "controller",
1993                &OrchestrationStatus::AwaitingApproval.to_string(),
1994                3,
1995            )
1996            .await
1997            .expect("failed to insert orchestration");
1998        database
1999            .orchestrations()
2000            .upsert_orchestration_task(planned_task(orchestration_id, "alpha"))
2001            .await
2002            .expect("failed to insert task");
2003        let orchestration = database
2004            .orchestrations()
2005            .load_orchestration_for_controller("controller")
2006            .await
2007            .expect("failed to load orchestration")
2008            .expect("orchestration should exist");
2009        let tasks = database
2010            .orchestrations()
2011            .load_orchestration_tasks(orchestration_id)
2012            .await
2013            .expect("failed to load tasks");
2014
2015        // Assert
2016        assert_eq!(orchestration.controller_project_id, 1);
2017        assert_eq!(orchestration.controller_session_id, "controller");
2018        assert_eq!(orchestration.max_parallelism, 3);
2019        assert_eq!(
2020            orchestration.status,
2021            OrchestrationStatus::AwaitingApproval.to_string()
2022        );
2023        assert_eq!(tasks.len(), 1);
2024        assert_eq!(tasks[0].task_key, "alpha");
2025        assert_eq!(tasks[0].child_session_id, None);
2026        assert_eq!(tasks[0].attempt_count, 0);
2027        assert_eq!(
2028            tasks[0].kind,
2029            OrchestrationTaskKind::Implementation.to_string()
2030        );
2031        assert_eq!(tasks[0].touched_areas, r#"["crates/alpha/"]"#);
2032        assert!(tasks[0].child_answer.is_none());
2033        assert!(tasks[0].research_report.is_none());
2034    }
2035
2036    #[tokio::test]
2037    async fn research_task_round_trips_report_and_latest_child_answer_without_scope() {
2038        // Arrange
2039        let database = controller_fixture().await;
2040        let orchestration_id = database
2041            .orchestrations()
2042            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2043            .await
2044            .expect("failed to insert orchestration");
2045        let task_id = database
2046            .orchestrations()
2047            .upsert_orchestration_task(PersistedOrchestrationTask {
2048                kind: OrchestrationTaskKind::Research.to_string(),
2049                touched_areas: "[]".to_string(),
2050                ..planned_task(orchestration_id, "architecture")
2051            })
2052            .await
2053            .expect("failed to insert research task");
2054        assert!(
2055            database
2056                .orchestrations()
2057                .claim_orchestration_task(task_id)
2058                .await
2059                .expect("failed to claim research task")
2060        );
2061        database
2062            .sessions()
2063            .insert_session_with_agent(PersistedSessionCreation {
2064                agent: "codex",
2065                base_branch: "main",
2066                id: "research-child",
2067                is_draft: false,
2068                model: AgentKind::Codex.default_model().as_str(),
2069                orchestration_task_id: Some(task_id),
2070                parent_session_id: None,
2071                personality_id: None,
2072                project_id: 1,
2073                reasoning_level: ReasoningLevel::default(),
2074                role: Some("OrchestrationResearcher"),
2075                speed_mode: SpeedMode::Normal,
2076                status: "Review",
2077            })
2078            .await
2079            .expect("failed to insert research child");
2080        assert!(
2081            database
2082                .orchestrations()
2083                .link_orchestration_task_child(task_id, "research-child")
2084                .await
2085                .expect("failed to link research child")
2086        );
2087        database
2088            .sessions()
2089            .append_session_message(
2090                "research-child",
2091                SessionMessageKind::AssistantAnswer,
2092                "Full architecture report",
2093            )
2094            .await
2095            .expect("failed to persist research answer");
2096
2097        // Act
2098        database
2099            .orchestrations()
2100            .update_orchestration_task_research_report(task_id, "Full architecture report")
2101            .await
2102            .expect("failed to persist research report");
2103        let task = database
2104            .orchestrations()
2105            .load_orchestration_tasks(orchestration_id)
2106            .await
2107            .expect("failed to load research task")
2108            .remove(0);
2109        let scope = database
2110            .orchestrations()
2111            .load_orchestration_task_scope_for_child("research-child")
2112            .await
2113            .expect("failed to inspect research child scope");
2114
2115        // Assert
2116        assert_eq!(task.kind, OrchestrationTaskKind::Research.to_string());
2117        assert_eq!(
2118            task.child_answer.as_deref(),
2119            Some("Full architecture report")
2120        );
2121        assert_eq!(
2122            task.research_report.as_deref(),
2123            Some("Full architecture report")
2124        );
2125        assert!(scope.is_none());
2126    }
2127
2128    #[tokio::test]
2129    async fn orchestration_task_kind_validation_rejects_unknown_writes_and_hydration() {
2130        // Arrange
2131        let (database, pool) = controller_fixture_with_pool().await;
2132        let orchestration_id = database
2133            .orchestrations()
2134            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2135            .await
2136            .expect("failed to insert orchestration");
2137        let invalid_task = PersistedOrchestrationTask {
2138            kind: "Unknown".to_string(),
2139            ..planned_task(orchestration_id, "invalid-write")
2140        };
2141
2142        // Act
2143        let write_error = database
2144            .orchestrations()
2145            .upsert_orchestration_task(invalid_task)
2146            .await
2147            .expect_err("unknown task kind should be rejected before persistence");
2148        sqlx::query(
2149            "INSERT INTO session_orchestration_task (session_orchestration_id, task_key, title, \
2150             prompt, status, kind) VALUES (?, 'invalid-read', 'Invalid', 'Inspect', 'Planned', \
2151             'Unknown')",
2152        )
2153        .bind(orchestration_id)
2154        .execute(&pool)
2155        .await
2156        .expect("failed to seed invalid persisted kind");
2157        let read_error = database
2158            .orchestrations()
2159            .load_orchestration_tasks(orchestration_id)
2160            .await
2161            .expect_err("unknown persisted task kind should fail hydration");
2162
2163        // Assert
2164        assert!(matches!(
2165            write_error,
2166            DbError::InvalidData {
2167                entity: "orchestration task kind",
2168                reason,
2169            } if reason == "unknown persisted kind `Unknown`"
2170        ));
2171        assert!(matches!(
2172            read_error,
2173            DbError::InvalidData {
2174                entity: "orchestration task kind",
2175                reason,
2176            } if reason == "unknown persisted kind `Unknown`"
2177        ));
2178    }
2179
2180    #[tokio::test]
2181    /// Reuses the same row when a retry re-proposes an existing task key, so a
2182    /// respawn cannot fan out a duplicate child for one subtask.
2183    async fn test_retry_with_same_task_key_updates_the_existing_row() {
2184        // Arrange
2185        let database = controller_fixture().await;
2186        let project_id = database
2187            .projects()
2188            .upsert_project("/tmp/project", None)
2189            .await
2190            .expect("failed to load project");
2191        let orchestration_id = database
2192            .orchestrations()
2193            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2194            .await
2195            .expect("failed to insert orchestration");
2196        let first_id = database
2197            .orchestrations()
2198            .upsert_orchestration_task(planned_task(orchestration_id, "alpha"))
2199            .await
2200            .expect("failed to insert task");
2201        let first_claim = database
2202            .orchestrations()
2203            .claim_orchestration_task(first_id)
2204            .await
2205            .expect("failed to claim first attempt");
2206        insert_orchestration_child(&database, project_id, "child-old", first_id).await;
2207        let first_link = database
2208            .orchestrations()
2209            .link_orchestration_task_child(first_id, "child-old")
2210            .await
2211            .expect("failed to link old child");
2212        database
2213            .orchestrations()
2214            .update_orchestration_task_status(
2215                first_id,
2216                &OrchestrationTaskStatus::Failed.to_string(),
2217                Some("agent crashed".to_string()),
2218            )
2219            .await
2220            .expect("failed to fail task");
2221
2222        // Act
2223        let retried_id = database
2224            .orchestrations()
2225            .upsert_orchestration_task(PersistedOrchestrationTask {
2226                title: "Task alpha, retried".to_string(),
2227                ..planned_task(orchestration_id, "alpha")
2228            })
2229            .await
2230            .expect("failed to retry task");
2231        let detached_child = database
2232            .orchestrations()
2233            .load_child_session_id_for_task(first_id)
2234            .await
2235            .expect("failed to check detached child");
2236        let retry_claim = database
2237            .orchestrations()
2238            .claim_orchestration_task(retried_id)
2239            .await
2240            .expect("failed to claim replacement attempt");
2241        insert_orchestration_child(&database, project_id, "child-replacement", retried_id).await;
2242        let replacement_link = database
2243            .orchestrations()
2244            .link_orchestration_task_child(retried_id, "child-replacement")
2245            .await
2246            .expect("failed to link replacement child");
2247        let linked_child = database
2248            .orchestrations()
2249            .load_child_session_id_for_task(first_id)
2250            .await
2251            .expect("failed to check replacement child");
2252        let tasks = database
2253            .orchestrations()
2254            .load_orchestration_tasks(orchestration_id)
2255            .await
2256            .expect("failed to load tasks");
2257
2258        // Assert
2259        assert!(first_claim);
2260        assert!(first_link);
2261        assert!(retry_claim);
2262        assert!(replacement_link);
2263        assert_eq!(retried_id, first_id);
2264        assert_eq!(tasks.len(), 1);
2265        assert_eq!(tasks[0].title, "Task alpha, retried");
2266        assert_eq!(
2267            tasks[0].status,
2268            OrchestrationTaskStatus::Running.to_string()
2269        );
2270        assert_eq!(tasks[0].attempt_count, 2);
2271        assert_eq!(
2272            tasks[0].child_session_id.as_deref(),
2273            Some("child-replacement")
2274        );
2275        assert_eq!(tasks[0].last_error, None);
2276        assert_eq!(detached_child, None);
2277        assert_eq!(linked_child.as_deref(), Some("child-replacement"));
2278    }
2279
2280    #[tokio::test]
2281    /// Links a created child, counts the attempt, and exposes its observed
2282    /// session state through the task snapshot.
2283    async fn test_child_linkage_counts_attempts_and_loads_observed_state() {
2284        // Arrange
2285        let database = controller_fixture().await;
2286        let project_id = database
2287            .projects()
2288            .upsert_project("/tmp/project", None)
2289            .await
2290            .expect("failed to upsert project");
2291        database
2292            .sessions()
2293            .insert_session("child-a", "gpt-5.6-sol", "main", "Draft", project_id)
2294            .await
2295            .expect("failed to insert child session");
2296        let orchestration_id = database
2297            .orchestrations()
2298            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2299            .await
2300            .expect("failed to insert orchestration");
2301        let task_id = database
2302            .orchestrations()
2303            .upsert_orchestration_task(planned_task(orchestration_id, "alpha"))
2304            .await
2305            .expect("failed to insert task");
2306        let claimed = database
2307            .orchestrations()
2308            .claim_orchestration_task(task_id)
2309            .await
2310            .expect("failed to claim task");
2311
2312        // Act
2313        let linked = database
2314            .orchestrations()
2315            .link_orchestration_task_child(task_id, "child-a")
2316            .await
2317            .expect("failed to link child");
2318        database
2319            .orchestrations()
2320            .update_orchestration_task_result_summary(task_id, "Added the parser")
2321            .await
2322            .expect("failed to record summary");
2323        let task = database
2324            .orchestrations()
2325            .load_orchestration_tasks(orchestration_id)
2326            .await
2327            .expect("failed to load task")
2328            .remove(0);
2329
2330        // Assert
2331        assert!(claimed);
2332        assert!(linked);
2333        assert_eq!(task.id, task_id);
2334        assert_eq!(task.attempt_count, 1);
2335        assert_eq!(task.child_session_id.as_deref(), Some("child-a"));
2336        assert_eq!(task.result_summary.as_deref(), Some("Added the parser"));
2337        assert_eq!(task.status, OrchestrationTaskStatus::Running.to_string());
2338        assert_eq!(task.child_status.as_deref(), Some("Draft"));
2339        assert_eq!(task.child_summary, None);
2340        assert_eq!(task.child_input_tokens, 0);
2341        assert_eq!(task.child_output_tokens, 0);
2342    }
2343
2344    #[tokio::test]
2345    /// Loads running, submitting, and canceling orchestrations while excluding
2346    /// terminal rows.
2347    async fn test_active_orchestration_load_includes_recoverable_states() {
2348        // Arrange
2349        let database = controller_fixture().await;
2350        let running_id = database
2351            .orchestrations()
2352            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2353            .await
2354            .expect("failed to insert running orchestration");
2355        let submitting_id = database
2356            .orchestrations()
2357            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2358            .await
2359            .expect("failed to insert submitting orchestration");
2360        let canceling_id = database
2361            .orchestrations()
2362            .insert_orchestration("controller", &OrchestrationStatus::Canceling.to_string(), 2)
2363            .await
2364            .expect("failed to insert canceling orchestration");
2365        let settled_id = database
2366            .orchestrations()
2367            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2368            .await
2369            .expect("failed to insert settled orchestration");
2370
2371        // Act
2372        database
2373            .orchestrations()
2374            .update_orchestration_status(settled_id, &OrchestrationStatus::Done.to_string())
2375            .await
2376            .expect("failed to settle orchestration");
2377        let first_claim = database
2378            .orchestrations()
2379            .claim_orchestration_rollup(submitting_id)
2380            .await
2381            .expect("failed to claim roll-up");
2382        let duplicate_claim = database
2383            .orchestrations()
2384            .claim_orchestration_rollup(submitting_id)
2385            .await
2386            .expect("failed to repeat roll-up claim");
2387        let active = database
2388            .orchestrations()
2389            .load_active_orchestrations()
2390            .await
2391            .expect("failed to load active orchestrations");
2392
2393        // Assert
2394        assert!(first_claim);
2395        assert!(!duplicate_claim);
2396        assert_eq!(
2397            active.iter().map(|row| row.id).collect::<Vec<_>>(),
2398            vec![running_id, submitting_id, canceling_id]
2399        );
2400        assert_eq!(active[1].status, OrchestrationStatus::Verifying.to_string());
2401        assert_eq!(active[1].verification_generation, 1);
2402        assert_eq!(active[2].status, OrchestrationStatus::Canceling.to_string());
2403    }
2404
2405    #[tokio::test]
2406    /// Uses the cancellation status as a durable barrier against both a late
2407    /// task claim and a late child link.
2408    async fn test_cancellation_barrier_blocks_fan_out_claims_and_links() {
2409        // Arrange
2410        let database = controller_fixture().await;
2411        let project_id = database
2412            .projects()
2413            .upsert_project("/tmp/project", None)
2414            .await
2415            .expect("failed to reload project");
2416        let orchestration_id = database
2417            .orchestrations()
2418            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2419            .await
2420            .expect("failed to insert orchestration");
2421        let late_task_id = database
2422            .orchestrations()
2423            .upsert_orchestration_task(planned_task(orchestration_id, "late"))
2424            .await
2425            .expect("failed to insert late task");
2426        let claimed_task_id = database
2427            .orchestrations()
2428            .upsert_orchestration_task(planned_task(orchestration_id, "claimed"))
2429            .await
2430            .expect("failed to insert claimed task");
2431        let initial_claim = database
2432            .orchestrations()
2433            .claim_orchestration_task(claimed_task_id)
2434            .await
2435            .expect("failed to claim task before cancellation");
2436        insert_orchestration_child(&database, project_id, "child-claimed", claimed_task_id).await;
2437
2438        // Act
2439        let cancellation_started = database
2440            .orchestrations()
2441            .begin_orchestration_cancellation(orchestration_id)
2442            .await
2443            .expect("failed to begin cancellation");
2444        let cancellation_retried = database
2445            .orchestrations()
2446            .begin_orchestration_cancellation(orchestration_id)
2447            .await
2448            .expect("failed to retry cancellation");
2449        let late_claim = database
2450            .orchestrations()
2451            .claim_orchestration_task(late_task_id)
2452            .await
2453            .expect("failed to inspect late claim");
2454        let late_link = database
2455            .orchestrations()
2456            .link_orchestration_task_child(claimed_task_id, "child-claimed")
2457            .await
2458            .expect("failed to inspect late link");
2459        let rollup_completion = database
2460            .orchestrations()
2461            .complete_orchestration_rollup(orchestration_id)
2462            .await
2463            .expect("failed to inspect late roll-up completion");
2464        let orchestration = database
2465            .orchestrations()
2466            .load_orchestration_for_controller("controller")
2467            .await
2468            .expect("failed to load orchestration")
2469            .expect("orchestration should exist");
2470        // Assert
2471        assert!(initial_claim);
2472        assert!(cancellation_started);
2473        assert!(cancellation_retried);
2474        assert!(!late_claim);
2475        assert!(!late_link);
2476        assert!(!rollup_completion);
2477        assert_eq!(
2478            orchestration.status,
2479            OrchestrationStatus::Canceling.to_string()
2480        );
2481    }
2482
2483    #[tokio::test]
2484    /// Observes the durable worker operation until successful completion and
2485    /// only then settles its submitting orchestration.
2486    async fn test_rollup_operation_status_controls_orchestration_completion() {
2487        // Arrange
2488        let database = controller_fixture().await;
2489        let orchestration_id = database
2490            .orchestrations()
2491            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2492            .await
2493            .expect("failed to insert orchestration");
2494        let claimed = database
2495            .orchestrations()
2496            .claim_orchestration_rollup(orchestration_id)
2497            .await
2498            .expect("failed to claim roll-up");
2499        assert!(claimed);
2500        let operation_id = format!("orchestration-rollup-{orchestration_id}-1");
2501
2502        // Act
2503        let missing_status = database
2504            .orchestrations()
2505            .load_rollup_operation_status(&operation_id)
2506            .await
2507            .expect("failed to load missing operation");
2508        database
2509            .operations()
2510            .claim_session_operation(&operation_id, "controller", "reply")
2511            .await
2512            .expect("failed to claim operation");
2513        let queued_status = database
2514            .orchestrations()
2515            .load_rollup_operation_status(&operation_id)
2516            .await
2517            .expect("failed to load queued operation");
2518        database
2519            .operations()
2520            .mark_session_operation_running(&operation_id)
2521            .await
2522            .expect("failed to run operation");
2523        let running_status = database
2524            .orchestrations()
2525            .load_rollup_operation_status(&operation_id)
2526            .await
2527            .expect("failed to load running operation");
2528        database
2529            .operations()
2530            .mark_session_operation_failed(&operation_id, "turn failed")
2531            .await
2532            .expect("failed to fail operation");
2533        let failed_status = database
2534            .orchestrations()
2535            .load_rollup_operation_status(&operation_id)
2536            .await
2537            .expect("failed to load failed operation");
2538        database
2539            .operations()
2540            .claim_session_operation(&operation_id, "controller", "reply")
2541            .await
2542            .expect("failed to reclaim operation");
2543        database
2544            .operations()
2545            .mark_session_operation_done(&operation_id)
2546            .await
2547            .expect("failed to complete operation");
2548        let done_status = database
2549            .orchestrations()
2550            .load_rollup_operation_status(&operation_id)
2551            .await
2552            .expect("failed to load completed operation");
2553        let completed = database
2554            .orchestrations()
2555            .complete_orchestration_rollup(orchestration_id)
2556            .await
2557            .expect("failed to complete orchestration");
2558        let duplicate_completion = database
2559            .orchestrations()
2560            .complete_orchestration_rollup(orchestration_id)
2561            .await
2562            .expect("failed to inspect duplicate completion");
2563        let orchestration = database
2564            .orchestrations()
2565            .load_orchestration_for_controller("controller")
2566            .await
2567            .expect("failed to load orchestration")
2568            .expect("orchestration should exist");
2569
2570        // Assert
2571        assert_eq!(missing_status, None);
2572        assert_eq!(queued_status.as_deref(), Some("queued"));
2573        assert_eq!(running_status.as_deref(), Some("running"));
2574        assert_eq!(failed_status.as_deref(), Some("failed"));
2575        assert_eq!(done_status.as_deref(), Some("done"));
2576        assert!(completed);
2577        assert!(!duplicate_completion);
2578        assert_eq!(
2579            orchestration.status,
2580            OrchestrationStatus::AwaitingIntegration.to_string()
2581        );
2582    }
2583
2584    #[tokio::test]
2585    /// Bulk-loads controller progress and child adjacency without per-session
2586    /// orchestration queries.
2587    async fn test_session_metadata_for_project_loads_controller_and_children_together() {
2588        // Arrange
2589        let database = controller_fixture().await;
2590        let project_id = database
2591            .projects()
2592            .upsert_project("/tmp/project", None)
2593            .await
2594            .expect("failed to reload project");
2595        let orchestration_id = database
2596            .orchestrations()
2597            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2598            .await
2599            .expect("failed to insert orchestration");
2600        let running_task_id = database
2601            .orchestrations()
2602            .upsert_orchestration_task(planned_task(orchestration_id, "running"))
2603            .await
2604            .expect("failed to insert running task");
2605        let waiting_task_id = database
2606            .orchestrations()
2607            .upsert_orchestration_task(planned_task(orchestration_id, "waiting"))
2608            .await
2609            .expect("failed to insert waiting task");
2610        insert_orchestration_child(&database, project_id, "child-running", running_task_id).await;
2611        insert_orchestration_child(&database, project_id, "child-waiting", waiting_task_id).await;
2612        let running_claim = database
2613            .orchestrations()
2614            .claim_orchestration_task(running_task_id)
2615            .await
2616            .expect("failed to claim running task");
2617        let waiting_claim = database
2618            .orchestrations()
2619            .claim_orchestration_task(waiting_task_id)
2620            .await
2621            .expect("failed to claim waiting task");
2622        let running_link = database
2623            .orchestrations()
2624            .link_orchestration_task_child(running_task_id, "child-running")
2625            .await
2626            .expect("failed to link running child");
2627        let waiting_link = database
2628            .orchestrations()
2629            .link_orchestration_task_child(waiting_task_id, "child-waiting")
2630            .await
2631            .expect("failed to link waiting child");
2632        database
2633            .orchestrations()
2634            .update_orchestration_task_status(
2635                waiting_task_id,
2636                &OrchestrationTaskStatus::WaitingForInput.to_string(),
2637                None,
2638            )
2639            .await
2640            .expect("failed to mark waiting child");
2641
2642        // Act
2643        let metadata = database
2644            .orchestrations()
2645            .load_session_metadata_for_project(project_id)
2646            .await
2647            .expect("failed to load bulk orchestration metadata");
2648
2649        // Assert
2650        assert!(running_claim);
2651        assert!(waiting_claim);
2652        assert!(running_link);
2653        assert!(waiting_link);
2654        assert_eq!(
2655            metadata,
2656            vec![
2657                SessionOrchestrationMetadataRow {
2658                    controller_session_id: Some("controller".to_string()),
2659                    orchestration_status: None,
2660                    running_task_count: 0,
2661                    session_id: "child-running".to_string(),
2662                    waiting_task_count: 0,
2663                },
2664                SessionOrchestrationMetadataRow {
2665                    controller_session_id: Some("controller".to_string()),
2666                    orchestration_status: None,
2667                    running_task_count: 0,
2668                    session_id: "child-waiting".to_string(),
2669                    waiting_task_count: 0,
2670                },
2671                SessionOrchestrationMetadataRow {
2672                    controller_session_id: None,
2673                    orchestration_status: Some(OrchestrationStatus::Running.to_string()),
2674                    running_task_count: 1,
2675                    session_id: "controller".to_string(),
2676                    waiting_task_count: 1,
2677                },
2678            ]
2679        );
2680    }
2681
2682    #[tokio::test]
2683    async fn approval_persists_plan_and_releases_proposed_tasks() {
2684        // Arrange
2685        let database = controller_fixture().await;
2686        let orchestration_id = database
2687            .orchestrations()
2688            .insert_orchestration(
2689                "controller",
2690                &OrchestrationStatus::AwaitingApproval.to_string(),
2691                2,
2692            )
2693            .await
2694            .expect("failed to insert orchestration");
2695        let task_id = database
2696            .orchestrations()
2697            .upsert_orchestration_task(planned_task(orchestration_id, "follow-up"))
2698            .await
2699            .expect("failed to insert proposed task");
2700        database
2701            .orchestrations()
2702            .update_orchestration_task_status(
2703                task_id,
2704                &OrchestrationTaskStatus::Proposed.to_string(),
2705                None,
2706            )
2707            .await
2708            .expect("failed to propose task");
2709
2710        // Act
2711        database
2712            .orchestrations()
2713            .update_orchestration_plan(orchestration_id, "Ship the campaign", 4)
2714            .await
2715            .expect("failed to update campaign plan");
2716        let approved = database
2717            .orchestrations()
2718            .approve_orchestration_plan(orchestration_id)
2719            .await
2720            .expect("failed to approve campaign");
2721        let duplicate_approval = database
2722            .orchestrations()
2723            .approve_orchestration_plan(orchestration_id)
2724            .await
2725            .expect("failed to inspect duplicate approval");
2726        let orchestration = database
2727            .orchestrations()
2728            .load_orchestration_for_controller("controller")
2729            .await
2730            .expect("failed to load campaign")
2731            .expect("campaign should exist");
2732        let task = database
2733            .orchestrations()
2734            .load_orchestration_tasks(orchestration_id)
2735            .await
2736            .expect("failed to load campaign tasks")
2737            .remove(0);
2738
2739        // Assert
2740        assert!(approved);
2741        assert!(!duplicate_approval);
2742        assert_eq!(orchestration.goal_statement, "Ship the campaign");
2743        assert_eq!(orchestration.max_parallelism, 4);
2744        assert_eq!(
2745            orchestration.status,
2746            OrchestrationStatus::Running.to_string()
2747        );
2748        assert_eq!(task.status, OrchestrationTaskStatus::Planned.to_string());
2749    }
2750
2751    #[tokio::test]
2752    async fn integration_approval_persists_selected_approach_atomically() {
2753        // Arrange
2754        let database = controller_fixture().await;
2755        let orchestration_id = database
2756            .orchestrations()
2757            .insert_orchestration(
2758                "controller",
2759                &OrchestrationStatus::AwaitingIntegration.to_string(),
2760                2,
2761            )
2762            .await
2763            .expect("failed to insert orchestration");
2764
2765        // Act
2766        let approved = database
2767            .orchestrations()
2768            .approve_orchestration_integration(orchestration_id, IntegrationApproach::ReviewRequest)
2769            .await
2770            .expect("failed to approve integration");
2771        let duplicate_approval = database
2772            .orchestrations()
2773            .approve_orchestration_integration(orchestration_id, IntegrationApproach::LocalMerge)
2774            .await
2775            .expect("failed to inspect duplicate approval");
2776        let approach = database
2777            .orchestrations()
2778            .load_orchestration_integration_approach(orchestration_id)
2779            .await
2780            .expect("failed to load integration approach");
2781        let orchestration = database
2782            .orchestrations()
2783            .load_orchestration_for_controller("controller")
2784            .await
2785            .expect("failed to load orchestration")
2786            .expect("orchestration should exist");
2787
2788        // Assert
2789        assert!(approved);
2790        assert!(!duplicate_approval);
2791        assert_eq!(approach, IntegrationApproach::ReviewRequest.to_string());
2792        assert_eq!(
2793            orchestration.status,
2794            OrchestrationStatus::Integrating.to_string()
2795        );
2796    }
2797
2798    #[tokio::test]
2799    async fn recoverable_focused_reviews_exclude_outstanding_continuations() {
2800        // Arrange
2801        let database = controller_fixture().await;
2802        let project_id = database
2803            .projects()
2804            .upsert_project("/tmp/project", None)
2805            .await
2806            .expect("failed to reload project");
2807        let orchestration_id = database
2808            .orchestrations()
2809            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2810            .await
2811            .expect("failed to insert orchestration");
2812        let task_id = database
2813            .orchestrations()
2814            .upsert_orchestration_task(planned_task(orchestration_id, "recoverable"))
2815            .await
2816            .expect("failed to insert recoverable task");
2817        assert!(
2818            database
2819                .orchestrations()
2820                .claim_orchestration_task(task_id)
2821                .await
2822                .expect("failed to claim recoverable task")
2823        );
2824        insert_orchestration_child(&database, project_id, "child-recoverable", task_id).await;
2825        assert!(
2826            database
2827                .orchestrations()
2828                .link_orchestration_task_child(task_id, "child-recoverable")
2829                .await
2830                .expect("failed to link recoverable child")
2831        );
2832        database
2833            .orchestrations()
2834            .update_orchestration_task_status(
2835                task_id,
2836                &OrchestrationTaskStatus::Reviewing.to_string(),
2837                None,
2838            )
2839            .await
2840            .expect("failed to mark child reviewing");
2841
2842        // Act
2843        let incomplete = database
2844            .orchestrations()
2845            .load_recoverable_focused_review_session_ids(project_id)
2846            .await
2847            .expect("failed to load incomplete review");
2848        database
2849            .orchestrations()
2850            .update_orchestration_task_status(
2851                task_id,
2852                &OrchestrationTaskStatus::ContinuationPending.to_string(),
2853                None,
2854            )
2855            .await
2856            .expect("failed to mark continuation pending");
2857        let pending = database
2858            .orchestrations()
2859            .load_recoverable_focused_review_session_ids(project_id)
2860            .await
2861            .expect("failed to inspect pending continuation recovery");
2862        database
2863            .orchestrations()
2864            .update_orchestration_task_status(
2865                task_id,
2866                &OrchestrationTaskStatus::ReviewApplying.to_string(),
2867                None,
2868            )
2869            .await
2870            .expect("failed to mark review application pending");
2871        let applying = database
2872            .orchestrations()
2873            .load_recoverable_focused_review_session_ids(project_id)
2874            .await
2875            .expect("failed to inspect review application recovery");
2876        database
2877            .orchestrations()
2878            .update_orchestration_task_status(
2879                task_id,
2880                &OrchestrationTaskStatus::Reviewing.to_string(),
2881                None,
2882            )
2883            .await
2884            .expect("failed to restore reviewing state");
2885        database
2886            .sessions()
2887            .update_session_focused_review(
2888                "child-recoverable",
2889                Some(FocusedReviewStatus::Ready),
2890                Some("42".to_string()),
2891                Some("### Suggestions\n\n- None".to_string()),
2892            )
2893            .await
2894            .expect("failed to complete focused review");
2895        let completed = database
2896            .orchestrations()
2897            .load_recoverable_focused_review_session_ids(project_id)
2898            .await
2899            .expect("failed to reload completed review");
2900
2901        // Assert
2902        assert_eq!(incomplete, ["child-recoverable"]);
2903        assert!(pending.is_empty() && applying.is_empty() && completed.is_empty());
2904    }
2905
2906    #[tokio::test]
2907    async fn review_application_claim_is_bounded_and_clears_consumed_review() {
2908        // Arrange
2909        let database = controller_fixture().await;
2910        let project_id = database
2911            .projects()
2912            .upsert_project("/tmp/project", None)
2913            .await
2914            .expect("failed to reload project");
2915        let orchestration_id = database
2916            .orchestrations()
2917            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
2918            .await
2919            .expect("failed to insert orchestration");
2920        let task_id = database
2921            .orchestrations()
2922            .upsert_orchestration_task(planned_task(orchestration_id, "reviewed"))
2923            .await
2924            .expect("failed to insert reviewed task");
2925        assert!(
2926            database
2927                .orchestrations()
2928                .claim_orchestration_task(task_id)
2929                .await
2930                .expect("failed to claim reviewed task")
2931        );
2932        insert_orchestration_child(&database, project_id, "child-reviewed", task_id).await;
2933        assert!(
2934            database
2935                .orchestrations()
2936                .link_orchestration_task_child(task_id, "child-reviewed")
2937                .await
2938                .expect("failed to link reviewed child")
2939        );
2940        database
2941            .orchestrations()
2942            .update_orchestration_task_status(
2943                task_id,
2944                &OrchestrationTaskStatus::Reviewing.to_string(),
2945                None,
2946            )
2947            .await
2948            .expect("failed to mark child reviewing");
2949        database
2950            .sessions()
2951            .update_session_focused_review(
2952                "child-reviewed",
2953                Some(FocusedReviewStatus::Ready),
2954                Some("42".to_string()),
2955                Some("### Suggestions\n\n- Fix it".to_string()),
2956            )
2957            .await
2958            .expect("failed to seed focused review");
2959
2960        // Act
2961        let claimed = database
2962            .orchestrations()
2963            .claim_orchestration_review_application(task_id, "Verify then apply", 3)
2964            .await
2965            .expect("failed to claim review application");
2966        let duplicate_claim = database
2967            .orchestrations()
2968            .claim_orchestration_review_application(task_id, "Duplicate", 3)
2969            .await
2970            .expect("failed to inspect duplicate review application");
2971        let task = database
2972            .orchestrations()
2973            .load_orchestration_tasks(orchestration_id)
2974            .await
2975            .expect("failed to load reviewed task")
2976            .remove(0);
2977        let review_cache = database
2978            .sessions()
2979            .load_session_focused_reviews_for_project(project_id)
2980            .await
2981            .expect("failed to load consumed review cache");
2982
2983        // Assert
2984        assert!(claimed);
2985        assert!(!duplicate_claim);
2986        assert_eq!(
2987            task.status,
2988            OrchestrationTaskStatus::ReviewApplying.to_string()
2989        );
2990        assert_eq!(task.continuation_generation, 1);
2991        assert_eq!(
2992            task.continuation_prompt.as_deref(),
2993            Some("Verify then apply")
2994        );
2995        assert_eq!(task.review_iteration, 1);
2996        assert_eq!(review_cache, [] as [crate::SessionFocusedReviewRow; 0]);
2997        assert_eq!(task.child_focused_review_status, None);
2998        assert_eq!(task.child_focused_review_text, None);
2999    }
3000
3001    #[tokio::test]
3002    async fn managed_child_continuation_questions_and_detach_are_durable() {
3003        // Arrange
3004        let database = controller_fixture().await;
3005        let project_id = database
3006            .projects()
3007            .upsert_project("/tmp/project", None)
3008            .await
3009            .expect("failed to reload project");
3010        let orchestration_id = database
3011            .orchestrations()
3012            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
3013            .await
3014            .expect("failed to insert orchestration");
3015        let task_id = database
3016            .orchestrations()
3017            .upsert_orchestration_task(planned_task(orchestration_id, "alpha"))
3018            .await
3019            .expect("failed to insert task");
3020        assert!(
3021            database
3022                .orchestrations()
3023                .claim_orchestration_task(task_id)
3024                .await
3025                .expect("failed to claim task")
3026        );
3027        insert_orchestration_child(&database, project_id, "child-alpha", task_id).await;
3028        assert!(
3029            database
3030                .orchestrations()
3031                .link_orchestration_task_child(task_id, "child-alpha")
3032                .await
3033                .expect("failed to link child")
3034        );
3035        database
3036            .orchestrations()
3037            .update_orchestration_task_status(
3038                task_id,
3039                &OrchestrationTaskStatus::WaitingForInput.to_string(),
3040                None,
3041            )
3042            .await
3043            .expect("failed to wait for child question");
3044
3045        // Act
3046        let surfaced = surface_and_clear_orchestration_questions(
3047            &database,
3048            orchestration_id,
3049            task_id,
3050            r#"[{"text":"Choose one"}]"#,
3051        )
3052        .await;
3053        database
3054            .orchestrations()
3055            .update_orchestration_task_status(
3056                task_id,
3057                &OrchestrationTaskStatus::Ready.to_string(),
3058                None,
3059            )
3060            .await
3061            .expect("failed to settle child after its question");
3062        let queued = database
3063            .orchestrations()
3064            .queue_orchestration_continuation(
3065                task_id,
3066                "Add the missing edge case",
3067                r#"["The edge case is tested"]"#,
3068                r#"["docs/"]"#,
3069            )
3070            .await
3071            .expect("failed to queue continuation");
3072        let duplicate_queue = database
3073            .orchestrations()
3074            .queue_orchestration_continuation(
3075                task_id,
3076                "Duplicate",
3077                r#"["Duplicate"]"#,
3078                r#"["ignored/"]"#,
3079            )
3080            .await
3081            .expect("failed to inspect duplicate continuation");
3082        let detached = database
3083            .orchestrations()
3084            .detach_orchestration_child("child-alpha")
3085            .await
3086            .expect("failed to detach child");
3087        let duplicate_detach = database
3088            .orchestrations()
3089            .detach_orchestration_child("child-alpha")
3090            .await
3091            .expect("failed to inspect duplicate detach");
3092        let (task, child, controller) =
3093            load_detached_campaign_state(&database, orchestration_id).await;
3094        let continuation_prompt = task.continuation_prompt.as_deref();
3095
3096        // Assert
3097        assert!(queued && surfaced && detached);
3098        assert!(!duplicate_queue && !duplicate_detach);
3099        assert_eq!(task.status, OrchestrationTaskStatus::Detached.to_string());
3100        assert_eq!(task.child_session_id, None);
3101        assert_eq!(task.continuation_generation, 1);
3102        assert_eq!(continuation_prompt, Some("Add the missing edge case"));
3103        assert_eq!(task.touched_areas, r#"["docs/"]"#);
3104        assert_eq!(child.role.as_deref(), Some("Worker"));
3105        assert_eq!(controller.status, "Review");
3106        assert_eq!(controller.questions.as_deref(), Some(""));
3107    }
3108
3109    #[tokio::test]
3110    async fn question_relay_preserves_controller_questions() {
3111        // Arrange
3112        let database = controller_fixture().await;
3113        let project_id = database
3114            .projects()
3115            .upsert_project("/tmp/project", None)
3116            .await
3117            .expect("failed to reload project");
3118        let orchestration_id = database
3119            .orchestrations()
3120            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
3121            .await
3122            .expect("failed to insert orchestration");
3123        let task_id = insert_waiting_orchestration_task(
3124            &database,
3125            project_id,
3126            orchestration_id,
3127            "worker",
3128            "child-worker",
3129        )
3130        .await;
3131        let controller_question = r#"[{"text":"Controller question"}]"#;
3132        database
3133            .sessions()
3134            .update_session_questions("controller", controller_question)
3135            .await
3136            .expect("failed to seed controller question");
3137
3138        // Act
3139        let blocked_by_controller = database
3140            .orchestrations()
3141            .surface_orchestration_questions(
3142                orchestration_id,
3143                task_id,
3144                r#"[{"text":"Child question"}]"#,
3145            )
3146            .await
3147            .expect("failed to inspect controller question");
3148        let orchestration = database
3149            .orchestrations()
3150            .load_orchestration_for_controller("controller")
3151            .await
3152            .expect("failed to load orchestration")
3153            .expect("orchestration should exist");
3154        let controller = database
3155            .sessions()
3156            .load_session("controller")
3157            .await
3158            .expect("failed to load controller")
3159            .expect("controller should exist");
3160
3161        // Assert
3162        assert!(!blocked_by_controller);
3163        assert_eq!(orchestration.relayed_question_task_id, None);
3164        assert_eq!(controller.questions.as_deref(), Some(controller_question));
3165    }
3166
3167    #[tokio::test]
3168    async fn question_relay_claims_one_exact_task_at_a_time() {
3169        // Arrange
3170        let database = controller_fixture().await;
3171        let project_id = database
3172            .projects()
3173            .upsert_project("/tmp/project", None)
3174            .await
3175            .expect("failed to reload project");
3176        let orchestration_id = database
3177            .orchestrations()
3178            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
3179            .await
3180            .expect("failed to insert orchestration");
3181        let first_task_id = insert_waiting_orchestration_task(
3182            &database,
3183            project_id,
3184            orchestration_id,
3185            "first",
3186            "child-first",
3187        )
3188        .await;
3189        let second_task_id = insert_waiting_orchestration_task(
3190            &database,
3191            project_id,
3192            orchestration_id,
3193            "second",
3194            "child-second",
3195        )
3196        .await;
3197
3198        // Act
3199        let first_surfaced = database
3200            .orchestrations()
3201            .surface_orchestration_questions(
3202                orchestration_id,
3203                first_task_id,
3204                r#"[{"text":"First child question"}]"#,
3205            )
3206            .await
3207            .expect("failed to surface first child question");
3208        let second_blocked = database
3209            .orchestrations()
3210            .surface_orchestration_questions(
3211                orchestration_id,
3212                second_task_id,
3213                r#"[{"text":"Second child question"}]"#,
3214            )
3215            .await
3216            .expect("failed to inspect occupied relay");
3217        let claimed_orchestration = database
3218            .orchestrations()
3219            .load_orchestration_for_controller("controller")
3220            .await
3221            .expect("failed to load claimed orchestration")
3222            .expect("orchestration should exist");
3223        database
3224            .orchestrations()
3225            .clear_orchestration_questions(orchestration_id)
3226            .await
3227            .expect("failed to release first relay");
3228        let second_surfaced = database
3229            .orchestrations()
3230            .surface_orchestration_questions(
3231                orchestration_id,
3232                second_task_id,
3233                r#"[{"text":"Second child question"}]"#,
3234            )
3235            .await
3236            .expect("failed to surface second child question");
3237        let second_claimed_orchestration = database
3238            .orchestrations()
3239            .load_orchestration_for_controller("controller")
3240            .await
3241            .expect("failed to load second claimed orchestration")
3242            .expect("orchestration should exist");
3243
3244        // Assert
3245        assert!(first_surfaced);
3246        assert!(!second_blocked);
3247        assert_eq!(
3248            claimed_orchestration.relayed_question_task_id,
3249            Some(first_task_id)
3250        );
3251        assert!(second_surfaced);
3252        assert_eq!(
3253            second_claimed_orchestration.relayed_question_task_id,
3254            Some(second_task_id)
3255        );
3256    }
3257
3258    #[tokio::test]
3259    async fn infrastructure_retries_are_bounded_and_completion_is_idempotent() {
3260        // Arrange
3261        let database = controller_fixture().await;
3262        let orchestration_id = database
3263            .orchestrations()
3264            .insert_orchestration("controller", &OrchestrationStatus::Running.to_string(), 2)
3265            .await
3266            .expect("failed to insert orchestration");
3267        let task_id = database
3268            .orchestrations()
3269            .upsert_orchestration_task(planned_task(orchestration_id, "alpha"))
3270            .await
3271            .expect("failed to insert task");
3272
3273        // Act
3274        let first_status = database
3275            .orchestrations()
3276            .record_orchestration_spawn_failure(task_id, "provider unavailable", 2)
3277            .await
3278            .expect("failed to record first retry");
3279        let second_status = database
3280            .orchestrations()
3281            .record_orchestration_spawn_failure(task_id, "provider unavailable", 2)
3282            .await
3283            .expect("failed to record second retry");
3284        let final_status = database
3285            .orchestrations()
3286            .record_orchestration_spawn_failure(task_id, "provider unavailable", 2)
3287            .await
3288            .expect("failed to exhaust retries");
3289        database
3290            .orchestrations()
3291            .update_orchestration_status(
3292                orchestration_id,
3293                &OrchestrationStatus::Integrating.to_string(),
3294            )
3295            .await
3296            .expect("failed to begin integration completion");
3297        let completed = database
3298            .orchestrations()
3299            .complete_orchestration_campaign(orchestration_id)
3300            .await
3301            .expect("failed to complete campaign");
3302        let duplicate_completion = database
3303            .orchestrations()
3304            .complete_orchestration_campaign(orchestration_id)
3305            .await
3306            .expect("failed to inspect duplicate completion");
3307        let orchestration = database
3308            .orchestrations()
3309            .load_orchestration_for_controller("controller")
3310            .await
3311            .expect("failed to load completed campaign")
3312            .expect("completed campaign should exist");
3313        let task = database
3314            .orchestrations()
3315            .load_orchestration_tasks(orchestration_id)
3316            .await
3317            .expect("failed to load exhausted task")
3318            .remove(0);
3319        let controller = database
3320            .sessions()
3321            .load_session("controller")
3322            .await
3323            .expect("failed to load completed controller")
3324            .expect("controller should exist");
3325
3326        // Assert
3327        assert_eq!(first_status, OrchestrationTaskStatus::Planned.to_string());
3328        assert_eq!(second_status, OrchestrationTaskStatus::Planned.to_string());
3329        assert_eq!(final_status, OrchestrationTaskStatus::Failed.to_string());
3330        assert!(completed);
3331        assert!(!duplicate_completion);
3332        assert_eq!(task.infrastructure_retry_count, 3);
3333        assert_eq!(task.status, OrchestrationTaskStatus::Failed.to_string());
3334        assert_eq!(orchestration.status, OrchestrationStatus::Done.to_string());
3335        assert_eq!(controller.status, "Done");
3336    }
3337
3338    #[tokio::test]
3339    /// Returns no orchestration for a controller session that never planned.
3340    async fn test_missing_orchestration_returns_none() {
3341        // Arrange
3342        let database = controller_fixture().await;
3343
3344        // Act
3345        let orchestration = database
3346            .orchestrations()
3347            .load_orchestration_for_controller("controller")
3348            .await
3349            .expect("failed to load orchestration");
3350        // Assert
3351        assert_eq!(orchestration, None);
3352    }
3353}