1use 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#[derive(Clone, Debug, Eq, PartialEq)]
20pub struct SessionOrchestrationRow {
21 pub controller_project_id: i64,
23 pub controller_session_id: String,
25 pub goal_statement: String,
27 pub id: i64,
29 pub max_parallelism: i64,
31 pub relayed_question_task_id: Option<i64>,
33 pub status: String,
35 pub verification_generation: i64,
37}
38
39#[derive(Clone, Debug, Eq, PartialEq)]
41pub struct SessionOrchestrationTaskRow {
42 pub acceptance_criteria: String,
44 pub area_violations: String,
46 pub areas_compliant: Option<bool>,
48 pub attempt_count: i64,
50 pub child_added_lines: i64,
52 pub child_answer: Option<String>,
54 pub child_deleted_lines: i64,
56 pub child_focused_review_status: Option<String>,
58 pub child_focused_review_text: Option<String>,
60 pub child_has_diff: Option<bool>,
62 pub child_input_tokens: i64,
64 pub child_output_tokens: i64,
66 pub child_questions: Option<String>,
68 pub child_session_id: Option<String>,
70 pub child_status: Option<String>,
72 pub child_summary: Option<String>,
74 pub continuation_generation: i64,
76 pub continuation_prompt: Option<String>,
78 pub id: i64,
80 pub infrastructure_retry_count: i64,
82 pub kind: String,
84 pub last_error: Option<String>,
86 pub merge_position: i64,
88 pub prompt: String,
90 pub research_report: Option<String>,
92 pub result_summary: Option<String>,
94 pub review_iteration: i64,
96 pub status: String,
98 pub task_key: String,
100 pub touched_areas: String,
102 pub title: String,
104 pub verification_reason: Option<String>,
106 pub verification_verdict: Option<String>,
108}
109
110#[derive(Clone, Debug, Eq, PartialEq)]
112pub struct OrchestrationTaskScopeRow {
113 pub base_branch: String,
115 pub id: i64,
117 pub touched_areas: String,
119}
120
121#[derive(Clone, Debug, Eq, PartialEq)]
123pub struct SessionOrchestrationMetadataRow {
124 pub controller_session_id: Option<String>,
126 pub orchestration_status: Option<String>,
129 pub running_task_count: i64,
131 pub session_id: String,
133 pub waiting_task_count: i64,
135}
136
137pub struct PersistedOrchestrationTask {
144 pub acceptance_criteria: String,
146 pub kind: String,
148 pub merge_position: i64,
150 pub prompt: String,
152 pub session_orchestration_id: i64,
154 pub task_key: String,
156 pub title: String,
158 pub touched_areas: String,
160}
161
162#[cfg_attr(any(test, feature = "test-utils"), mockall::automock)]
167#[async_trait]
168pub trait OrchestrationRepository: Send + Sync {
169 async fn insert_orchestration(
171 &self,
172 controller_session_id: &str,
173 status: &str,
174 max_parallelism: i64,
175 ) -> Result<i64, DbError>;
176
177 async fn upsert_orchestration_task(
185 &self,
186 task: PersistedOrchestrationTask,
187 ) -> Result<i64, DbError>;
188
189 async fn load_orchestration_for_controller(
191 &self,
192 controller_session_id: &str,
193 ) -> Result<Option<SessionOrchestrationRow>, DbError>;
194
195 async fn load_active_orchestrations(&self) -> Result<Vec<SessionOrchestrationRow>, DbError>;
197
198 async fn load_recoverable_focused_review_session_ids(
201 &self,
202 project_id: i64,
203 ) -> Result<Vec<String>, DbError>;
204
205 async fn load_session_metadata_for_project(
208 &self,
209 project_id: i64,
210 ) -> Result<Vec<SessionOrchestrationMetadataRow>, DbError>;
211
212 async fn load_orchestration_tasks(
214 &self,
215 session_orchestration_id: i64,
216 ) -> Result<Vec<SessionOrchestrationTaskRow>, DbError>;
217
218 async fn load_orchestration_integration_approach(&self, id: i64) -> Result<String, DbError>;
220
221 async fn load_orchestration_task_scope_for_child(
223 &self,
224 child_session_id: &str,
225 ) -> Result<Option<OrchestrationTaskScopeRow>, DbError>;
226
227 async fn load_child_session_id_for_task(&self, task_id: i64)
229 -> Result<Option<String>, DbError>;
230
231 async fn begin_orchestration_cancellation(&self, id: i64) -> Result<bool, DbError>;
233
234 async fn claim_orchestration_task(&self, id: i64) -> Result<bool, DbError>;
236
237 async fn claim_orchestration_review_application(
240 &self,
241 id: i64,
242 prompt: &str,
243 iteration_limit: i64,
244 ) -> Result<bool, DbError>;
245
246 async fn claim_orchestration_rollup(&self, id: i64) -> Result<bool, DbError>;
248
249 async fn complete_orchestration_rollup(&self, id: i64) -> Result<bool, DbError>;
251
252 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 async fn complete_orchestration_campaign(&self, id: i64) -> Result<bool, DbError>;
263
264 async fn load_rollup_operation_status(
266 &self,
267 operation_id: &str,
268 ) -> Result<Option<String>, DbError>;
269
270 async fn update_orchestration_status(&self, id: i64, status: &str) -> Result<(), DbError>;
272
273 async fn approve_orchestration_plan(&self, id: i64) -> Result<bool, DbError>;
275
276 async fn approve_orchestration_integration(
278 &self,
279 id: i64,
280 approach: IntegrationApproach,
281 ) -> Result<bool, DbError>;
282
283 async fn update_orchestration_plan(
285 &self,
286 id: i64,
287 goal_statement: &str,
288 max_parallelism: i64,
289 ) -> Result<(), DbError>;
290
291 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 async fn reset_orchestration_verification(&self, id: i64) -> Result<(), DbError>;
302
303 async fn record_orchestration_spawn_failure(
305 &self,
306 id: i64,
307 error: &str,
308 retry_limit: i64,
309 ) -> Result<String, DbError>;
310
311 async fn detach_orchestration_child(&self, child_session_id: &str) -> Result<bool, DbError>;
313
314 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 async fn clear_orchestration_questions(
325 &self,
326 session_orchestration_id: i64,
327 ) -> Result<(), DbError>;
328
329 async fn link_orchestration_task_child(
331 &self,
332 id: i64,
333 child_session_id: &str,
334 ) -> Result<bool, DbError>;
335
336 async fn update_orchestration_task_status(
338 &self,
339 id: i64,
340 status: &str,
341 last_error: Option<String>,
342 ) -> Result<(), DbError>;
343
344 async fn update_orchestration_task_result_summary(
346 &self,
347 id: i64,
348 result_summary: &str,
349 ) -> Result<(), DbError>;
350
351 async fn update_orchestration_task_research_report(
354 &self,
355 id: i64,
356 research_report: &str,
357 ) -> Result<(), DbError>;
358
359 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#[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 pub(crate) fn new(pool: SqlitePool, timestamp_source: Arc<dyn TimestampSource>) -> Self {
384 Self(pool, timestamp_source)
385 }
386
387 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 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 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 let error = database
1798 .orchestrations()
1799 .load_orchestration_for_controller("controller")
1800 .await
1801 .expect_err("invalid status should fail hydration");
1802
1803 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 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 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!(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 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 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 async fn test_planned_orchestration_round_trips_before_any_child_exists() {
1985 let database = controller_fixture().await;
1987
1988 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_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 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 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_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 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 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!(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 async fn test_retry_with_same_task_key_updates_the_existing_row() {
2184 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 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!(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 async fn test_child_linkage_counts_attempts_and_loads_observed_state() {
2284 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 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!(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 async fn test_active_orchestration_load_includes_recoverable_states() {
2348 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 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!(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 async fn test_cancellation_barrier_blocks_fan_out_claims_and_links() {
2409 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 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!(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 async fn test_rollup_operation_status_controls_orchestration_completion() {
2487 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 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_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 async fn test_session_metadata_for_project_loads_controller_and_children_together() {
2588 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 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!(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 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 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!(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 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 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!(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 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 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_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 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 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!(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 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 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!(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 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 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!(!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 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 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!(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 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 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_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 async fn test_missing_orchestration_returns_none() {
3341 let database = controller_fixture().await;
3343
3344 let orchestration = database
3346 .orchestrations()
3347 .load_orchestration_for_controller("controller")
3348 .await
3349 .expect("failed to load orchestration");
3350 assert_eq!(orchestration, None);
3352 }
3353}