awaken-server 0.6.0

Multi-protocol HTTP server with SSE, mailbox, and protocol adapters for Awaken
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
//! Mailbox service: persistent run queue, dispatch execution, leasing, and lifecycle management.

use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use async_trait::async_trait;
use parking_lot::Mutex as SyncMutex;
use thiserror::Error;
use tokio::sync::{Mutex, RwLock};
use tokio::task::JoinHandle;

use awaken_runtime::{ResolveError, RunActivation};
use awaken_server_contract::contract::commit_coordinator::CommitCoordinator;
use awaken_server_contract::contract::event::AgentEvent;
use awaken_server_contract::contract::event_sink::EventSink;
use awaken_server_contract::contract::mailbox::{MailboxStore, RunDispatchStatus};
use awaken_server_contract::contract::message::Message;
use awaken_server_contract::contract::run::{
    RunActivationSnapshot, RunInputSnapshot, RunIntent, RunKind, RunOptions, RunTraceContext,
};
use awaken_server_contract::contract::staged_commit::{
    OutboxServerEventPublisher, StagedCommitCoordinator,
};
use awaken_server_contract::contract::storage::{RunRecord, StorageError, ThreadRunStore};
use awaken_server_contract::contract::suspension::{ToolCallOutcome, ToolCallResume};
use awaken_server_contract::contract::tool_intercept::{AdapterKind, RunMode};

use crate::transport::channel_sink::ReconnectableEventSink;

/// Guard window for inline-claimed dispatches: if the process crashes between
/// enqueue and claim, the sweep will reclaim the dispatch after this period.
const INLINE_CLAIM_GUARD_MS: u64 = 60_000;
#[cfg(not(test))]
const REMOTE_CANCEL_WAIT_MS: u64 = 5_000;
#[cfg(test)]
const REMOTE_CANCEL_WAIT_MS: u64 = 250;
const REMOTE_CANCEL_POLL_MS: u64 = 25;
const DISPATCH_SIGNAL_BATCH_DEFAULT: usize = 32;
const DISPATCH_SIGNAL_EXPIRES_DEFAULT: Duration = Duration::from_millis(500);
const DISPATCH_SIGNAL_ERROR_DELAY: Duration = Duration::from_millis(250);
const DISPATCH_SIGNAL_BLOCKED_NACK_BASE_DELAY_DEFAULT: Duration = Duration::from_millis(500);
const DISPATCH_SIGNAL_BLOCKED_NACK_MAX_DELAY_DEFAULT: Duration = Duration::from_secs(30);
const DISPATCH_SIGNAL_BATCH_ENV: &str = "AWAKEN_DISPATCH_SIGNAL_BATCH_SIZE";
const DISPATCH_SIGNAL_EXPIRES_ENV: &str = "AWAKEN_DISPATCH_SIGNAL_FETCH_EXPIRES_MS";
const DISPATCH_SIGNAL_NACK_BASE_DELAY_ENV: &str = "AWAKEN_DISPATCH_SIGNAL_NACK_BASE_DELAY_MS";
const DISPATCH_SIGNAL_NACK_MAX_DELAY_ENV: &str = "AWAKEN_DISPATCH_SIGNAL_NACK_MAX_DELAY_MS";
const DISPATCH_SIGNAL_MAX_CONCURRENT_HANDLERS_DEFAULT: usize = 32;
const DISPATCH_SIGNAL_MAX_CONCURRENT_HANDLERS_ENV: &str =
    "AWAKEN_DISPATCH_SIGNAL_MAX_CONCURRENT_HANDLERS";
const TERMINAL_RECONCILE_BATCH: usize = 100;
const MAILBOX_DEPTH_STATUSES: [RunDispatchStatus; 6] = [
    RunDispatchStatus::Queued,
    RunDispatchStatus::Claimed,
    RunDispatchStatus::Acked,
    RunDispatchStatus::Cancelled,
    RunDispatchStatus::Superseded,
    RunDispatchStatus::DeadLetter,
];

/// Validation message returned when an inline submit loses the active-run race.
pub(crate) const ACTIVE_RUN_CONFLICT_MESSAGE: &str =
    "thread has an active run; cannot claim inline";

pub(super) fn run_activation_snapshot(
    request: &RunActivation,
    persisted_input: RunInputSnapshot,
    resolution_id: Option<String>,
) -> RunActivationSnapshot {
    RunActivationSnapshot {
        intent: request.intent.clone(),
        input: persisted_input,
        options: request.options.clone(),
        trace: request.trace.clone(),
        seeded_decisions: request.control.seeded_decisions.clone(),
        resolution_id,
    }
}

// ── Legacy run snapshot conversion ───────────────────────────────

/// Typed envelope for RunActivation fields that Mailbox stores opaquely.
/// Centralizes the legacy run snapshot round-trip round-trip.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub(super) struct LegacyRunSnapshotExtras {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    overrides: Option<awaken_server_contract::contract::inference::InferenceOverride>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    decisions: Vec<(
        String,
        awaken_server_contract::contract::suspension::ToolCallResume,
    )>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    frontend_tools: Vec<awaken_server_contract::contract::tool::ToolDescriptor>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    continue_run_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    run_id_hint: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    dispatch_id_hint: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    parent_thread_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    transport_request_id: Option<String>,
    #[serde(default)]
    run_mode: RunMode,
    #[serde(default)]
    adapter: AdapterKind,
}

impl LegacyRunSnapshotExtras {
    #[cfg(test)]
    fn from_request(request: &awaken_runtime::RunActivation) -> Self {
        Self {
            overrides: request.options.overrides.clone(),
            decisions: request.control.seeded_decisions.clone(),
            frontend_tools: request.options.frontend_tools.clone(),
            continue_run_id: request.resume_run_id().map(str::to_owned),
            run_id_hint: request.persistence.run_id_hint.clone(),
            dispatch_id_hint: request.persistence.dispatch_id_hint.clone(),
            parent_thread_id: request.trace.parent_thread_id.clone(),
            transport_request_id: request.trace.transport_request_id.clone(),
            run_mode: request.trace.run_mode,
            adapter: request.trace.adapter,
        }
    }

    #[cfg(test)]
    fn to_value(&self) -> Result<Option<serde_json::Value>, serde_json::Error> {
        if self.overrides.is_none()
            && self.decisions.is_empty()
            && self.frontend_tools.is_empty()
            && self.continue_run_id.is_none()
            && self.run_id_hint.is_none()
            && self.dispatch_id_hint.is_none()
            && self.parent_thread_id.is_none()
            && self.transport_request_id.is_none()
            && self.run_mode == RunMode::Foreground
            && self.adapter == AdapterKind::Internal
        {
            Ok(None)
        } else {
            serde_json::to_value(self).map(Some)
        }
    }

    fn from_value(value: &serde_json::Value) -> Result<Self, serde_json::Error> {
        serde_json::from_value(value.clone())
    }

    #[cfg(test)]
    fn apply_to(self, mut request: awaken_runtime::RunActivation) -> awaken_runtime::RunActivation {
        if let Some(ov) = self.overrides {
            request = request.with_overrides(ov);
        }
        if !self.decisions.is_empty() {
            request = request.with_decisions(self.decisions);
        }
        if !self.frontend_tools.is_empty() {
            request = request.with_frontend_tools(self.frontend_tools);
        }
        if let Some(crid) = self.continue_run_id {
            request = request.with_continue_run_id(crid);
        }
        if let Some(run_id_hint) = self.run_id_hint {
            request = request.with_run_id_hint(run_id_hint);
        }
        if let Some(dispatch_id_hint) = self.dispatch_id_hint {
            request = request.with_dispatch_id_hint(dispatch_id_hint);
        }
        if let Some(parent_thread_id) = self.parent_thread_id {
            request = request.with_parent_thread_id(parent_thread_id);
        }
        if let Some(transport_request_id) = self.transport_request_id {
            request = request.with_transport_request_id(transport_request_id);
        }
        request
            .with_run_mode(self.run_mode)
            .with_adapter(self.adapter)
    }
}

pub(super) struct LegacyRunRequestSnapshotAdapter {
    pub snapshot: awaken_server_contract::contract::storage::RunRequestSnapshot,
    pub input: RunInputSnapshot,
    pub resolution_id: Option<String>,
    pub thread_id: String,
    pub agent_id: Option<String>,
    pub parent_run_id: Option<String>,
    pub extras: Option<LegacyRunSnapshotExtras>,
}

impl TryFrom<LegacyRunRequestSnapshotAdapter> for RunActivationSnapshot {
    type Error = String;

    fn try_from(value: LegacyRunRequestSnapshotAdapter) -> Result<Self, Self::Error> {
        let mut options = RunOptions {
            overrides: None,
            frontend_tools: value.snapshot.frontend_tools,
        };
        let mut trace = RunTraceContext {
            parent_run_id: value.parent_run_id,
            parent_thread_id: value.snapshot.parent_thread_id,
            origin: value.snapshot.origin.into(),
            adapter: AdapterKind::Internal,
            run_mode: RunMode::Resume,
            dispatch_id: None,
            session_id: None,
            transport_request_id: value.snapshot.transport_request_id,
            correlation_id: None,
        };
        let mut kind = RunKind::NewIntent;
        let mut seeded_decisions = value
            .snapshot
            .decisions
            .into_iter()
            .map(|decision| (decision.call_id, decision.resume))
            .collect::<Vec<_>>();

        if let Some(extras) = value.extras {
            if let Some(overrides) = extras.overrides {
                options.overrides = Some(overrides);
            }
            if !extras.frontend_tools.is_empty() {
                options.frontend_tools = extras.frontend_tools;
            }
            if !extras.decisions.is_empty() {
                seeded_decisions = extras.decisions;
            }
            if let Some(run_id) = extras.continue_run_id {
                kind = RunKind::HitlResume { run_id };
            }
            if let Some(parent_thread_id) = extras.parent_thread_id {
                trace.parent_thread_id = Some(parent_thread_id);
            }
            if let Some(transport_request_id) = extras.transport_request_id {
                trace.transport_request_id = Some(transport_request_id);
            }
            trace.run_mode = extras.run_mode;
            trace.adapter = extras.adapter;
        }

        Ok(RunActivationSnapshot {
            intent: RunIntent {
                agent_id: value.agent_id,
                thread_id: value.thread_id,
                kind,
            },
            input: value.input,
            options,
            trace,
            seeded_decisions,
            resolution_id: value.resolution_id,
        })
    }
}

pub(super) fn legacy_input_snapshot(
    run: &RunRecord,
    snapshot: &awaken_server_contract::contract::storage::RunRequestSnapshot,
) -> RunInputSnapshot {
    if let Some(input) = run.input.as_ref() {
        return RunInputSnapshot {
            thread_id: input.thread_id.clone(),
            range: input.range,
            trigger_message_ids: input.trigger_message_ids.clone(),
            selected_message_ids: input.selected_message_ids.clone(),
            context_policy: input.context_policy.clone(),
            compacted_snapshot_id: input.compacted_snapshot_id.clone(),
        };
    }
    RunInputSnapshot {
        thread_id: run.thread_id.clone(),
        range: None,
        trigger_message_ids: snapshot.input_message_ids.clone(),
        selected_message_ids: Vec::new(),
        context_policy: None,
        compacted_snapshot_id: None,
    }
}

// ── TaskDoneMailboxNotify ────────────────────────────────────────────

/// Fallback for inbox delivery when the agent's run has ended.
///
/// Implements [`OnInboxClosed`](awaken_runtime::inbox::OnInboxClosed) — when an `InboxSender::send()` fails
/// because the receiver was dropped (agent run returned with AwaitingTasks),
/// this enqueues a mailbox wake dispatch so the thread gets a continuation run.
pub struct TaskDoneMailboxNotify {
    mailbox: Arc<Mailbox>,
    thread_id: String,
    continue_run_id: Option<String>,
}

impl TaskDoneMailboxNotify {
    pub fn new(mailbox: Arc<Mailbox>, thread_id: String, continue_run_id: Option<String>) -> Self {
        Self {
            mailbox,
            thread_id,
            continue_run_id,
        }
    }
}

impl awaken_runtime::inbox::OnInboxClosed for TaskDoneMailboxNotify {
    fn closed(&self, message: &serde_json::Value) {
        let mailbox = self.mailbox.clone();
        let thread_id = self.thread_id.clone();
        let continue_run_id = self.continue_run_id.clone();
        let wake_message = awaken_runtime::inbox::inbox_event_message(message);

        // Spawn because OnInboxClosed::closed is sync but enqueue+dispatch is async
        tokio::spawn(async move {
            let mut request = RunActivation::new(thread_id.clone(), vec![wake_message])
                .with_origin(awaken_server_contract::contract::storage::RunRequestOrigin::Internal)
                .with_run_mode(RunMode::InternalWake)
                .with_adapter(AdapterKind::Internal);
            if let Some(run_id) = continue_run_id {
                request = request.with_continue_run_id(run_id);
            }
            if let Err(e) = mailbox.submit_background(request).await {
                tracing::warn!(thread_id, error = %e, "failed to enqueue background task wake dispatch");
            }
        });
    }
}

// ── Public types ─────────────────────────────────────────────────────

/// Result returned by submit/submit_background.
#[derive(Debug, Clone)]
pub struct MailboxSubmitResult {
    pub dispatch_id: String,
    pub run_id: String,
    pub thread_id: String,
    pub status: MailboxDispatchStatus,
}

/// Dispatch status for a submitted run activation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MailboxDispatchStatus {
    /// Job was claimed and is executing now.
    Running,
    /// Job is queued, waiting for the current run to finish.
    Queued,
}

/// Mailbox service errors.
#[derive(Debug, Error)]
pub enum MailboxError {
    #[error("validation error: {0}")]
    Validation(String),
    #[error("store error: {0}")]
    Store(#[from] StorageError),
    #[error("resolution error while {context}: {source}")]
    Resolution {
        context: &'static str,
        source: ResolveError,
    },
    #[error("internal error: {0}")]
    Internal(String),
    /// A foreground submit cannot be consumed because a barrier pending entry
    /// ahead of it must be consumed first. Returned by the submit preflight
    /// *before* any interrupt/cancel side effect, so a blocked submit never
    /// cancels the active run (ADR-0042 D6: barriers are never skipped).
    #[error("delivery blocked by barrier: pending '{blocking_pending_id}' must be consumed first")]
    DeliveryBlockedByBarrier { blocking_pending_id: String },
}

/// Outcome classification for runtime run results.
#[derive(Debug)]
pub enum MailboxRunOutcome {
    /// Run completed successfully.
    Completed,
    /// Transient infrastructure failure -- retry.
    TransientError(String),
    /// Permanent failure -- do not retry.
    PermanentError(String),
}

impl MailboxRunOutcome {
    fn metric_label(&self) -> &'static str {
        match self {
            Self::Completed => "completed",
            Self::TransientError(_) => "transient_error",
            Self::PermanentError(_) => "permanent_error",
        }
    }
}

/// Configuration for the Mailbox service.
#[derive(Debug, Clone)]
pub struct MailboxConfig {
    /// Lease duration in milliseconds (default 30_000).
    pub lease_ms: u64,
    /// Lease duration in milliseconds when the run is suspended/waiting
    /// for human input (default 600_000 = 10 minutes).
    pub suspended_lease_ms: u64,
    /// How often to renew leases (default 10s).
    pub lease_renewal_interval: Duration,
    /// How often to sweep for expired leases (default 30s).
    pub sweep_interval: Duration,
    /// How often to run GC for terminal dispatches (default 60s).
    pub gc_interval: Duration,
    /// How long to keep terminal dispatches before purging (default 24h).
    pub gc_ttl: Duration,
    /// Default max attempts before dead-lettering (default 5).
    pub default_max_attempts: u32,
    /// Default retry delay in milliseconds (default 250).
    pub default_retry_delay_ms: u64,
    /// Maximum retry delay in milliseconds for exponential backoff (default 30_000).
    pub max_retry_delay_ms: u64,
}

impl Default for MailboxConfig {
    fn default() -> Self {
        Self {
            lease_ms: 30_000,
            suspended_lease_ms: 600_000,
            lease_renewal_interval: Duration::from_secs(10),
            sweep_interval: Duration::from_secs(30),
            gc_interval: Duration::from_secs(60),
            gc_ttl: Duration::from_secs(24 * 60 * 60),
            default_max_attempts: 5,
            default_retry_delay_ms: 250,
            max_retry_delay_ms: 30_000,
        }
    }
}

/// Callback invoked during mailbox maintenance GC ticks.
pub type MailboxMaintenanceCallback = Arc<dyn Fn() + Send + Sync + 'static>;

/// Startup recovery retry settings used by lifecycle startup.
#[derive(Clone)]
pub struct MailboxStartupRecoveryConfig {
    /// Maximum recovery attempts before giving up. Values below 1 are treated
    /// as one attempt.
    pub max_attempts: u32,
    /// Delay between failed recovery attempts.
    pub retry_delay: Duration,
}

impl Default for MailboxStartupRecoveryConfig {
    fn default() -> Self {
        Self {
            max_attempts: 1,
            retry_delay: Duration::from_millis(250),
        }
    }
}

/// Configuration for framework-managed mailbox lifecycle tasks.
#[derive(Clone)]
pub struct MailboxLifecycleConfig {
    /// Delay before startup recovery and maintenance begin.
    pub startup_delay: Duration,
    /// Retry policy for startup recovery.
    pub startup_recovery: MailboxStartupRecoveryConfig,
    /// Optional cleanup hook for application-owned resources.
    pub maintenance_callback: Option<MailboxMaintenanceCallback>,
}

impl Default for MailboxLifecycleConfig {
    fn default() -> Self {
        Self {
            startup_delay: Duration::ZERO,
            startup_recovery: MailboxStartupRecoveryConfig::default(),
            maintenance_callback: None,
        }
    }
}

/// Handle for framework-managed mailbox lifecycle tasks.
///
/// Dropping the handle does not stop lifecycle tasks. Call [`shutdown`](Self::shutdown)
/// for quiescent shutdown or [`abort`](Self::abort) for fire-and-forget stop.
#[derive(Clone)]
pub struct MailboxLifecycleHandle {
    tasks: Arc<StdMutex<Option<MailboxLifecycleTasks>>>,
    transition_lock: Arc<Mutex<()>>,
}

impl MailboxLifecycleHandle {
    /// Abort lifecycle tasks. This is idempotent.
    pub fn abort(&self) {
        if let Some(tasks) = self.tasks.lock().expect("lifecycle lock poisoned").take() {
            tasks.abort();
        }
    }

    /// Abort lifecycle tasks and wait until they have fully exited.
    ///
    /// This is the quiescent shutdown path. Use it when a caller needs a hard
    /// guarantee that a subsequent lifecycle start cannot overlap old recovery
    /// or maintenance tasks.
    pub async fn shutdown(&self) -> Result<(), MailboxError> {
        let _transition_guard = self.transition_lock.lock().await;
        let tasks = self.tasks.lock().expect("lifecycle lock poisoned").take();
        if let Some(tasks) = tasks {
            tasks.shutdown().await?;
        }
        Ok(())
    }

    /// Returns true while lifecycle tasks are registered for this mailbox.
    pub fn is_running(&self) -> bool {
        self.tasks
            .lock()
            .expect("lifecycle lock poisoned")
            .is_some()
    }
}

struct MailboxLifecycleTasks {
    recover_task: Option<JoinHandle<()>>,
    dispatch_signal_task: Option<JoinHandle<()>>,
    maintenance_task: JoinHandle<()>,
}

impl MailboxLifecycleTasks {
    fn abort(self) {
        if let Some(task) = self.recover_task {
            task.abort();
        }
        if let Some(task) = self.dispatch_signal_task {
            task.abort();
        }
        self.maintenance_task.abort();
    }

    async fn shutdown(self) -> Result<(), MailboxError> {
        if let Some(task) = self.recover_task {
            task.abort();
            await_lifecycle_task("mailbox startup recovery", task).await?;
        }
        if let Some(task) = self.dispatch_signal_task {
            task.abort();
            await_lifecycle_task("mailbox dispatch signal loop", task).await?;
        }
        self.maintenance_task.abort();
        await_lifecycle_task("mailbox maintenance", self.maintenance_task).await
    }
}

async fn await_lifecycle_task(name: &str, task: JoinHandle<()>) -> Result<(), MailboxError> {
    match task.await {
        Ok(()) => Ok(()),
        Err(error) if error.is_cancelled() => Ok(()),
        Err(error) if error.is_panic() => Err(MailboxError::Internal(format!("{name} panicked"))),
        Err(error) => Err(MailboxError::Internal(format!("{name} failed: {error}"))),
    }
}

// ── Internal types ───────────────────────────────────────────────────

/// Per-thread worker status.
enum MailboxWorkerStatus {
    Idle,
    /// Transitional: claim in progress. Prevents TOCTOU race where two
    /// concurrent dispatches both see Idle and both try to claim.
    Claiming,
    Running {
        dispatch_id: String,
        run_id: String,
        lease_handle: JoinHandle<()>,
        sink: Arc<ReconnectableEventSink>,
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DispatchAttempt {
    Claimed,
    Busy,
    NoEligible,
    TransientError,
}

impl DispatchAttempt {
    fn started_execution(self) -> bool {
        matches!(self, DispatchAttempt::Claimed)
    }
}

/// Cached thread state, valid for the duration of a lease.
struct ThreadContext {
    messages: Vec<Message>,
    latest_run: Option<RunRecord>,
    run_cache: HashMap<String, RunRecord>,
}

impl ThreadContext {
    async fn load(run_store: &dyn ThreadRunStore, thread_id: &str) -> Result<Self, MailboxError> {
        let messages = run_store
            .load_messages(thread_id)
            .await?
            .unwrap_or_default();
        let latest_run = run_store.latest_run(thread_id).await?;
        let mut run_cache = HashMap::new();
        if let Some(ref run) = latest_run {
            run_cache.insert(run.run_id.clone(), run.clone());
        }
        Ok(Self {
            messages,
            latest_run,
            run_cache,
        })
    }

    fn get_run(&self, run_id: &str) -> Option<&RunRecord> {
        self.run_cache.get(run_id)
    }

    fn apply_checkpoint(&mut self, messages: &[Message], run: &RunRecord) {
        self.messages = messages.to_vec();
        self.latest_run = Some(run.clone());
        self.run_cache.insert(run.run_id.clone(), run.clone());
    }
}

/// Per-thread worker. Store is the sole queue authority.
struct MailboxWorker {
    status: MailboxWorkerStatus,
    thread_ctx: Option<ThreadContext>,
}

impl Default for MailboxWorker {
    fn default() -> Self {
        Self {
            status: MailboxWorkerStatus::Idle,
            thread_ctx: None,
        }
    }
}

// ── Suspension-aware event sink ──────────────────────────────────────

/// Wraps an inner `EventSink` and sets a shared flag when the run
/// enters a suspended (waiting) state, detected by a `ToolCallDone`
/// event with `ToolCallOutcome::Suspended`.
struct SuspensionAwareSink {
    inner: Arc<dyn EventSink>,
    suspended: Arc<AtomicBool>,
}

#[async_trait]
impl EventSink for SuspensionAwareSink {
    async fn emit(&self, event: AgentEvent) {
        if matches!(
            &event,
            AgentEvent::ToolCallDone {
                outcome: ToolCallOutcome::Suspended,
                ..
            }
        ) {
            self.suspended.store(true, Ordering::Release);
        }
        // Reset the flag when the run resumes from suspension.
        if matches!(&event, AgentEvent::ToolCallResumed { .. }) {
            self.suspended.store(false, Ordering::Release);
        }
        self.inner.emit(event).await;
    }

    async fn close(&self) {
        self.inner.close().await;
    }
}

/// RAII guard that decrements the active-runs gauge on drop.
struct ActiveRunGuard;

impl Drop for ActiveRunGuard {
    fn drop(&mut self) {
        crate::metrics::dec_active_runs();
    }
}

// ── Mailbox service ──────────────────────────────────────────────────

/// Unified persistent run queue.
///
/// Orchestrates `MailboxStore` (dispatch persistence) + `ThreadRunStore`
/// (run/message truth) + `RunDispatchExecutor` (execution)
/// with lease-based distributed claim, per-thread serialization, sweep,
/// and garbage collection.
pub struct Mailbox {
    executor: Arc<dyn RunDispatchExecutor>,
    store: Arc<dyn MailboxStore>,
    /// Single durable-write boundary from `executor.commit_coordinator()` (see `try_new`).
    coordinator: Arc<dyn CommitCoordinator>,
    /// Staged variant of `coordinator`, when the executor exposes one. Present
    /// only when the executor can commit canonical-event/outbox writes
    /// atomically with the checkpoint; required for runtime event capture.
    staged_coordinator: Option<Arc<dyn StagedCommitCoordinator>>,
    /// Full thread/run store for server-side reads and queries. Supplied by
    /// the caller, which builds the coordinator from this same store.
    run_store: Arc<dyn ThreadRunStore>,
    pending_thread_run_store: Option<Arc<dyn awaken_stores::PendingThreadRunStore>>,
    consumer_id: String,
    workers: RwLock<HashMap<String, Arc<SyncMutex<MailboxWorker>>>>,
    config: MailboxConfig,
    runtime_event_capture: Option<RuntimeEventCaptureConfig>,
    server_event_publisher: Option<Arc<dyn OutboxServerEventPublisher>>,
    server_event_origin: String,
    lifecycle_tasks: Arc<StdMutex<Option<MailboxLifecycleTasks>>>,
    lifecycle_start_lock: Arc<Mutex<()>>,
    /// Striped per-thread locks serializing the message-append
    /// read-modify-write in `prepare_run_for_dispatch` (see `lock_thread_append`).
    thread_append_locks: Box<[Mutex<()>]>,
    /// Retry queue for checkpoint events that failed to publish (see `checkpoint_repair`).
    checkpoint_repair_queue: Arc<StdMutex<VecDeque<checkpoint_repair::CheckpointRepairTask>>>,
    /// Optional source for materializing a run's pinned `RegistrySet` from its
    /// persisted resolution_id (publication snapshot_version) at execution
    /// time. Lets a durable run resolve against a pinned publication — e.g. an
    /// unsaved admin-sandbox draft agent — by handing the runtime a resolved
    /// `RegistrySet` via `with_pinned_registry_set`. Pin/publication logic stays
    /// server-side; the runtime never learns what a "pin" is. `None` keeps the
    /// live-registry resolution used by all other deployments.
    pinned_registry: Option<PinnedRegistrySource>,
}

/// Server-side handle for re-materializing a run's pinned publication.
#[derive(Clone)]
struct PinnedRegistrySource {
    store: Arc<dyn awaken_server_contract::VersionedRegistryStore>,
    scope: awaken_server_contract::ScopeId,
}

impl Mailbox {
    /// Create a new Mailbox service.
    pub fn new(
        executor: impl IntoDispatchExecutor,
        store: Arc<dyn MailboxStore>,
        run_store: Arc<dyn ThreadRunStore>,
        consumer_id: String,
        config: MailboxConfig,
    ) -> Self {
        Self::try_new(executor, store, run_store, consumer_id, config)
            .expect("Mailbox requires a CommitCoordinator outside unit-test fallback")
    }

    /// Fallible constructor for production wiring that must fail closed.
    ///
    /// ADR-0038 D7: the mailbox adopts `executor.commit_coordinator()` and
    /// derives its `ThreadRunStore` from that coordinator. Unit tests retain
    /// the legacy implicit run-store coordinator for small harnesses; every
    /// non-test build returns an error instead of silently taking a partial
    /// durable path with no canonical events/outbox writes.
    pub fn try_new(
        executor: impl IntoDispatchExecutor,
        store: Arc<dyn MailboxStore>,
        run_store: Arc<dyn ThreadRunStore>,
        consumer_id: String,
        config: MailboxConfig,
    ) -> Result<Self, MailboxError> {
        let executor = executor.into_dispatch_executor();
        let staged_coordinator = executor.staged_commit_coordinator();
        let coordinator = if let Some(coordinator) = executor.commit_coordinator() {
            coordinator
        } else if cfg!(test) {
            tracing::warn!(
                "using unit-test MailboxRunStoreCoordinator fallback; non-test executors must \
                 provide a CommitCoordinator"
            );
            Arc::new(MailboxRunStoreCoordinator::new(Arc::clone(&run_store)))
                as Arc<dyn CommitCoordinator>
        } else {
            return Err(MailboxError::Internal(
                "Mailbox requires a CommitCoordinator; wire a durable \
                 Memory/File/Pg coordinator through the runtime"
                    .to_string(),
            ));
        };
        if let (Some(coordinator_identity), Some(run_store_identity)) = (
            coordinator.thread_run_storage_identity(),
            run_store.thread_run_storage_identity(),
        ) && coordinator_identity != run_store_identity
        {
            return Err(MailboxError::Validation(format!(
                "mailbox run_store must match executor CommitCoordinator thread/run store \
                 (coordinator={coordinator_identity}, run_store={run_store_identity})"
            )));
        }
        Ok(Self {
            executor,
            store,
            coordinator,
            staged_coordinator,
            run_store,
            pending_thread_run_store: None,
            consumer_id,
            workers: RwLock::new(HashMap::new()),
            config,
            runtime_event_capture: None,
            server_event_publisher: None,
            server_event_origin: "mailbox".to_string(),
            lifecycle_tasks: Arc::new(StdMutex::new(None)),
            lifecycle_start_lock: Arc::new(Mutex::new(())),
            thread_append_locks: (0..Self::THREAD_APPEND_STRIPES)
                .map(|_| Mutex::new(()))
                .collect(),
            checkpoint_repair_queue: Arc::new(StdMutex::new(VecDeque::new())),
            pinned_registry: None,
        })
    }

    /// Wire a versioned-registry source so durable runs can resolve against a
    /// pinned publication (by their persisted resolution_id) at execution
    /// time. Used to run unsaved draft agents (admin sandbox) durably. The
    /// runtime stays pin-agnostic — it only receives the materialized
    /// `RegistrySet` via `with_pinned_registry_set`.
    #[must_use]
    pub fn with_pinned_registry(
        self,
        store: Arc<dyn awaken_server_contract::VersionedRegistryStore>,
        scope: impl Into<String>,
    ) -> Self {
        self.try_with_pinned_registry(store, scope)
            .expect("pinned registry scope_id must be valid")
    }

    pub fn try_with_pinned_registry(
        mut self,
        store: Arc<dyn awaken_server_contract::VersionedRegistryStore>,
        scope: impl Into<String>,
    ) -> Result<Self, awaken_server_contract::ScopeError> {
        let scope = awaken_server_contract::ScopeId::new(scope.into())?;
        self.pinned_registry = Some(PinnedRegistrySource { store, scope });
        Ok(self)
    }

    /// Materialize the pinned `RegistrySet` for a run's resolution_id
    /// (publication snapshot_version), overlaying live runtime objects. Returns
    /// `Ok(None)` only when no pinned-registry source is wired.
    async fn materialize_pinned_registry_set(
        &self,
        resolution_id: &str,
    ) -> Result<Option<awaken_runtime::registry::RegistrySet>, MailboxError> {
        let Some(source) = self.pinned_registry.as_ref() else {
            return Ok(None);
        };
        let snapshot_version = resolution_id.parse::<u64>().map_err(|error| {
            MailboxError::Internal(format!(
                "invalid pinned registry resolution id '{resolution_id}': {error}"
            ))
        })?;
        let live = self.executor.live_registry_set().ok_or_else(|| {
            MailboxError::Internal(
                "pinned registry materialization requires a live registry snapshot".to_string(),
            )
        })?;
        let materializer = crate::services::frozen_registry::FrozenAgentRegistryMaterializer::new(
            source.store.clone(),
        );
        let frozen = materializer
            .materialize(awaken_server_contract::VersionSelector::Publication {
                scope_id: source.scope.as_str().to_string(),
                snapshot_version,
            })
            .await
            .map_err(|error| {
                MailboxError::Internal(format!(
                    "failed to materialize pinned registry publication {snapshot_version}: {error}"
                ))
            })?;
        Ok(Some(frozen.to_registry_set(&live)))
    }

    /// Number of stripes for `lock_thread_append` (defined in `submit`).
    const THREAD_APPEND_STRIPES: usize = 256;

    /// Default bounded channel capacity for the runtime->SSE relay.
    const EVENT_CHANNEL_CAPACITY: usize = 256;

    /// Single source of truth for the thread/run store. Always returns the
    /// store wrapped by the mailbox's [`CommitCoordinator`] — callers that
    /// previously held a parallel `Arc<dyn ThreadRunStore>` should reach
    /// it through this accessor instead.
    pub fn thread_run_store(&self) -> &Arc<dyn ThreadRunStore> {
        &self.run_store
    }

    /// Borrow the durable-write coordinator.
    pub fn coordinator(&self) -> &Arc<dyn CommitCoordinator> {
        &self.coordinator
    }

    /// Forward a tool-call decision to an active run in this process only.
    ///
    /// Distributed callers should use [`Self::send_decision_live`] so remote
    /// active runs can receive the decision through targeted live delivery.
    pub fn send_decision(&self, id: &str, tool_call_id: String, resume: ToolCallResume) -> bool {
        self.executor.send_decision(id, tool_call_id, resume)
    }
}

mod cancel;
mod checkpoint;
mod checkpoint_repair;
mod coordinator_facade;
mod decision;
mod dispatch_execution;
mod helpers;
mod lifecycle;
mod live_delivery;
mod pending_delivery;
mod runtime_event_capture;
mod server_event_capture;
mod signal_loop;
mod staging_coordinator;
mod submit;

use self::coordinator_facade::MailboxRunStoreCoordinator;
use self::{helpers::*, runtime_event_capture::RuntimeEventCaptureConfig};
pub use crate::run_dispatch::RunDispatchExecutor;
pub use coordinator_facade::IntoDispatchExecutor;
#[cfg(test)]
mod tests;