Skip to main content

mj_client/
session.rs

1//! The session operations used by interactive control surfaces.
2
3use std::fmt;
4use std::future::Future;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::time::Duration;
8
9use agent_client_protocol::schema::v1::SessionConfigOption;
10use anyhow::{Context, Result, ensure};
11use mj_core::config::Config;
12use mj_core::elicitation::ElicitationResponse;
13use mj_core::state::{ManagedSessionSnapshot, SessionRecord};
14
15use mj_core::relay::{
16    AnalyzeDeltaRepository, RelayCommand, RelayCursor, RelayEvent, RelayOperationalState, RepoDelta,
17};
18use mj_core::worker_launch::ReviewerLaunchConfig;
19
20pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
21
22#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
23#[serde(tag = "kind", content = "detail", rename_all = "snake_case")]
24pub enum ViewError {
25    Unreachable(String),
26    TargetMissing(String),
27    ProjectionIntegrity(String),
28}
29
30impl ViewError {
31    pub fn detail(&self) -> &str {
32        match self {
33            Self::Unreachable(detail)
34            | Self::TargetMissing(detail)
35            | Self::ProjectionIntegrity(detail) => detail,
36        }
37    }
38}
39
40#[derive(Debug, Clone, PartialEq, Default)]
41pub struct ManagedSessionView {
42    pub snapshot: Option<ManagedSessionSnapshot>,
43    pub connected: bool,
44    pub error: Option<ViewError>,
45}
46
47#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
48pub struct RelayAttachment {
49    pub state: RelayOperationalState,
50    pub events: Vec<RelayEvent>,
51    pub through_ordinal: u64,
52    pub through_digest: String,
53}
54
55#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
56pub struct StartedReviewer {
57    pub native_session_id: Option<String>,
58    pub config_options: Vec<SessionConfigOption>,
59    pub reused: bool,
60    pub state: RelayOperationalState,
61}
62
63#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
64#[serde(rename_all = "snake_case")]
65pub enum ReviewerAction {
66    Start {
67        config: Box<ReviewerLaunchConfig>,
68    },
69    Submit {
70        command_id: String,
71        command: RelayCommand,
72    },
73    Attach {
74        after_ordinal: u64,
75        after_digest: String,
76    },
77    Acknowledge {
78        through_ordinal: u64,
79        through_digest: String,
80    },
81    Status,
82    RespondElicitation {
83        elicitation_id: String,
84        response: ElicitationResponse,
85    },
86    Pause,
87    /// Stop only the preparation that owns this generation; stale cleanup must not stop its replacement.
88    PauseGeneration {
89        generation: u64,
90    },
91    CaptureDelta {
92        baselines: std::collections::BTreeMap<std::path::PathBuf, String>,
93    },
94    AdvanceBaseline {
95        trees: std::collections::BTreeMap<std::path::PathBuf, String>,
96    },
97    AnalyzeDelta {
98        repositories: Vec<AnalyzeDeltaRepository>,
99    },
100    TakeLaneDispatches,
101}
102
103impl ReviewerAction {
104    pub const fn operation_name(&self) -> &'static str {
105        match self {
106            Self::Start { .. } => "reviewer_start",
107            Self::Submit { .. } => "reviewer_submit",
108            Self::Attach { .. } => "reviewer_attach",
109            Self::Acknowledge { .. } => "reviewer_acknowledge",
110            Self::Status => "reviewer_status",
111            Self::RespondElicitation { .. } => "reviewer_respond_elicitation",
112            Self::Pause | Self::PauseGeneration { .. } => "reviewer_pause",
113            Self::CaptureDelta { .. } => "reviewer_capture_delta",
114            Self::AdvanceBaseline { .. } => "reviewer_advance_baseline",
115            Self::AnalyzeDelta { .. } => "reviewer_analyze_delta",
116            Self::TakeLaneDispatches => "reviewer_take_lane_dispatches",
117        }
118    }
119}
120
121#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
122#[serde(rename_all = "snake_case")]
123pub enum ReviewerOutcome {
124    Started(Box<StartedReviewer>),
125    Accepted {
126        ordinal: u64,
127    },
128    Attached(Box<RelayAttachment>),
129    Acknowledged(RelayCursor),
130    Status(Box<RelayOperationalState>),
131    ElicitationResolved,
132    Paused,
133    Delta {
134        repositories: Vec<RepoDelta>,
135    },
136    BaselineAdvanced,
137    ChangedFunctions {
138        packet: String,
139    },
140    LaneDispatches {
141        requests: Vec<mj_core::review::lanes::ReviewSubagentRequest>,
142    },
143}
144
145/// Preserve whether a submit failed before delivery or lost its acknowledgement.
146#[derive(Debug)]
147pub struct SubmitFailure {
148    pub message: String,
149    pub unconfirmed: bool,
150    /// The worker itself read the command and rejected it. Other definite
151    /// failures (a full queue, a worker that could not be reached) were
152    /// never looked at and may succeed on a later attempt.
153    pub refused: bool,
154}
155impl From<String> for SubmitFailure {
156    fn from(message: String) -> Self {
157        Self {
158            message,
159            unconfirmed: false,
160            refused: false,
161        }
162    }
163}
164impl From<&str> for SubmitFailure {
165    fn from(message: &str) -> Self {
166        message.to_owned().into()
167    }
168}
169
170/// The worker read the command and rejected it; resending it unchanged will
171/// be rejected again.
172#[derive(Debug)]
173pub struct Refused;
174impl std::fmt::Display for Refused {
175    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
176        f.write_str("the worker refused the command")
177    }
178}
179impl std::error::Error for Refused {}
180
181/// Submission lost its acknowledgement; callers must reconcile before retrying.
182#[derive(Debug)]
183pub struct DeliveryUnconfirmed;
184impl std::fmt::Display for DeliveryUnconfirmed {
185    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
186        f.write_str("delivery unconfirmed")
187    }
188}
189impl std::error::Error for DeliveryUnconfirmed {}
190
191pub struct PendingRelaySubmit {
192    completion: BoxFuture<'static, Result<u64>>,
193}
194
195impl PendingRelaySubmit {
196    pub fn new(completion: BoxFuture<'static, Result<u64>>) -> Self {
197        Self { completion }
198    }
199
200    pub async fn wait(self) -> Result<u64> {
201        self.completion.await
202    }
203}
204
205pub struct PendingRelaySync {
206    completion: BoxFuture<'static, Result<()>>,
207}
208
209impl PendingRelaySync {
210    pub fn new(completion: BoxFuture<'static, Result<()>>) -> Self {
211        Self { completion }
212    }
213
214    pub async fn wait(self) -> Result<()> {
215        self.completion.await
216    }
217}
218
219#[derive(Debug, Default)]
220pub struct ReviewState {
221    pub review: Option<mj_core::storage::StoredReview>,
222}
223
224pub trait SessionHandleBackend: Send + Sync {
225    fn transcript_history(
226        &self,
227        _before: Option<mj_core::storage::TranscriptCursor>,
228    ) -> BoxFuture<'_, Result<mj_core::storage::TranscriptHistoryPage>> {
229        Box::pin(async { anyhow::bail!("earlier conversation history is unavailable") })
230    }
231    fn search_prompts(
232        &self,
233        bundle_id: String,
234        scope: mj_core::storage::HistoryScope,
235        query: String,
236    ) -> BoxFuture<'_, Result<Vec<mj_core::storage::PromptHistoryEntry>>>;
237    fn review_state(&self) -> BoxFuture<'_, Result<ReviewState>>;
238    fn resolve_review_settings(
239        &self,
240        cancelled: Arc<std::sync::atomic::AtomicBool>,
241    ) -> BoxFuture<'_, Result<mj_core::review::settings::ResolvedReviewSettings>> {
242        let _ = cancelled;
243        Box::pin(async { anyhow::bail!("review settings resolution is unavailable") })
244    }
245
246    fn config_result(&self, command_id: String) -> BoxFuture<'_, Result<Option<Option<String>>>>;
247
248    fn clone_box(&self) -> Box<dyn SessionHandleBackend>;
249    fn session_id(&self) -> &str;
250    fn view(&self) -> ManagedSessionView;
251    fn is_stopped(&self) -> bool;
252    fn has_changed(&self) -> Result<bool>;
253    fn changed(&mut self) -> BoxFuture<'_, Result<ManagedSessionView>>;
254    fn enqueue_submit(
255        &self,
256        command_id: String,
257        command: RelayCommand,
258    ) -> BoxFuture<'_, Result<PendingRelaySubmit>>;
259    fn enqueue_sync(&self) -> BoxFuture<'_, Result<PendingRelaySync>>;
260    fn respond_elicitation(
261        &self,
262        elicitation_id: String,
263        response: ElicitationResponse,
264    ) -> BoxFuture<'_, Result<()>>;
265    fn stop_background_task(&self, background_task_id: String) -> BoxFuture<'_, Result<()>>;
266    fn reviewer(
267        &self,
268        role: Option<String>,
269        action: ReviewerAction,
270    ) -> BoxFuture<'_, Result<ReviewerOutcome>>;
271}
272
273pub struct SessionHandle {
274    backend: Box<dyn SessionHandleBackend>,
275}
276
277impl SessionHandle {
278    pub async fn transcript_history(
279        &self,
280        before: Option<mj_core::storage::TranscriptCursor>,
281    ) -> Result<mj_core::storage::TranscriptHistoryPage> {
282        self.backend.transcript_history(before).await
283    }
284
285    pub async fn search_prompts(
286        &self,
287        bundle_id: String,
288        scope: mj_core::storage::HistoryScope,
289        query: String,
290    ) -> Result<Vec<mj_core::storage::PromptHistoryEntry>> {
291        self.backend.search_prompts(bundle_id, scope, query).await
292    }
293    pub async fn review_state(&self) -> Result<ReviewState> {
294        self.backend.review_state().await
295    }
296
297    pub async fn resolve_review_settings(
298        &self,
299        cancelled: Arc<std::sync::atomic::AtomicBool>,
300    ) -> Result<mj_core::review::settings::ResolvedReviewSettings> {
301        self.backend.resolve_review_settings(cancelled).await
302    }
303
304    pub fn new(backend: impl SessionHandleBackend + 'static) -> Self {
305        Self {
306            backend: Box::new(backend),
307        }
308    }
309
310    pub fn session_id(&self) -> &str {
311        self.backend.session_id()
312    }
313
314    pub fn view(&self) -> ManagedSessionView {
315        self.backend.view()
316    }
317
318    pub fn is_stopped(&self) -> bool {
319        self.backend.is_stopped()
320    }
321
322    pub fn has_changed(&self) -> Result<bool> {
323        self.backend.has_changed()
324    }
325
326    pub async fn changed(&mut self) -> Result<ManagedSessionView> {
327        self.backend.changed().await
328    }
329
330    pub async fn submit(&self, command_id: String, command: RelayCommand) -> Result<u64> {
331        self.enqueue_submit(command_id, command).await?.wait().await
332    }
333
334    /// Apply a setting and wait for its durable success or rejection.
335    pub async fn set_config(&self, key: String, value: String) -> Result<()> {
336        self.set_config_with_id(new_command_id("set-config")?, key, value)
337            .await
338            .map(|_| ())
339    }
340
341    /// Apply a setting under a caller-chosen command id and return the
342    /// ordinal at which the relay accepted it. A caller that retries with the
343    /// same id after a lost reply gets the original ordinal back.
344    pub async fn set_config_with_id(
345        &self,
346        command_id: String,
347        key: String,
348        value: String,
349    ) -> Result<u64> {
350        self.apply_configuration_command(command_id, RelayCommand::SetConfig { key, value })
351            .await
352    }
353
354    async fn apply_configuration_command(
355        &self,
356        command_id: String,
357        command: RelayCommand,
358    ) -> Result<u64> {
359        let ordinal = self.submit(command_id.clone(), command).await?;
360        tokio::time::timeout(Duration::from_secs(60), async {
361            loop {
362                if let Some(error) = self.backend.config_result(command_id.clone()).await? {
363                    if let Some(error) = error {
364                        anyhow::bail!("{error}");
365                    }
366                    self.sync_now().await?;
367                    return Ok(ordinal);
368                }
369                ensure!(
370                    !self.is_stopped(),
371                    "session stopped while applying configuration"
372                );
373                if let Some(error) = self.view().error {
374                    anyhow::bail!("configuration connection failed: {}", error.detail());
375                }
376                tokio::time::sleep(Duration::from_millis(50)).await;
377            }
378        })
379        .await
380        .context("configuration command did not complete within 60 seconds")?
381    }
382
383    /// A follow-up prompt must wait for the mode change, not just admission.
384    pub async fn apply_plan_control(
385        &self,
386        command_id: String,
387        control: mj_core::acp::PlanControl,
388    ) -> Result<()> {
389        match control {
390            mj_core::acp::PlanControl::RestoreExecutionMode => self
391                .apply_configuration_command(command_id, RelayCommand::RestoreExecutionMode)
392                .await
393                .map(|_| ()),
394            mj_core::acp::PlanControl::SetConfig { key, value } => self
395                .set_config_with_id(command_id, key, value)
396                .await
397                .map(|_| ()),
398            mj_core::acp::PlanControl::SetSessionMode { mode_id } => {
399                self.submit(
400                    command_id,
401                    RelayCommand::SetSessionMode {
402                        mode_id: mode_id.clone(),
403                    },
404                )
405                .await?;
406                tokio::time::timeout(Duration::from_secs(60), async {
407                    loop {
408                        self.sync_now().await?;
409                        if self
410                            .view()
411                            .snapshot
412                            .as_ref()
413                            .and_then(|snapshot| snapshot.operational.modes.as_ref())
414                            .is_some_and(|modes| modes.current_mode_id.to_string() == mode_id)
415                        {
416                            return Ok(());
417                        }
418                        ensure!(!self.is_stopped(), "session stopped while changing mode");
419                        tokio::time::sleep(Duration::from_millis(50)).await;
420                    }
421                })
422                .await
423                .context("session mode change did not complete within 60 seconds")?
424            }
425        }
426    }
427
428    pub async fn enqueue_submit(
429        &self,
430        command_id: String,
431        command: RelayCommand,
432    ) -> Result<PendingRelaySubmit> {
433        self.backend.enqueue_submit(command_id, command).await
434    }
435
436    pub async fn sync_now(&self) -> Result<()> {
437        self.enqueue_sync().await?.wait().await
438    }
439
440    pub async fn enqueue_sync(&self) -> Result<PendingRelaySync> {
441        self.backend.enqueue_sync().await
442    }
443
444    pub async fn respond_elicitation(
445        &self,
446        elicitation_id: String,
447        response: ElicitationResponse,
448    ) -> Result<()> {
449        self.backend
450            .respond_elicitation(elicitation_id, response)
451            .await
452    }
453
454    pub async fn stop_background_task(&self, background_task_id: String) -> Result<()> {
455        self.backend.stop_background_task(background_task_id).await
456    }
457
458    pub async fn reviewer(&self, action: ReviewerAction) -> Result<ReviewerOutcome> {
459        self.reviewer_as(None, action).await
460    }
461
462    pub async fn reviewer_as(
463        &self,
464        role: Option<String>,
465        action: ReviewerAction,
466    ) -> Result<ReviewerOutcome> {
467        self.backend.reviewer(role, action).await
468    }
469}
470
471impl Clone for SessionHandle {
472    fn clone(&self) -> Self {
473        Self {
474            backend: self.backend.clone_box(),
475        }
476    }
477}
478
479impl fmt::Debug for SessionHandle {
480    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
481        formatter
482            .debug_struct("SessionHandle")
483            .field("session_id", &self.session_id())
484            .finish_non_exhaustive()
485    }
486}
487
488pub trait SessionControlBackend: Send + Sync {
489    fn session(&self, session_id: String) -> BoxFuture<'_, Result<SessionHandle>>;
490}
491
492#[derive(Clone)]
493pub struct SessionControl {
494    backend: Arc<dyn SessionControlBackend>,
495}
496
497impl SessionControl {
498    pub fn new(backend: impl SessionControlBackend + 'static) -> Self {
499        Self {
500            backend: Arc::new(backend),
501        }
502    }
503
504    pub async fn session(&self, session_id: impl Into<String>) -> Result<SessionHandle> {
505        self.backend.session(session_id.into()).await
506    }
507
508    pub async fn wait_for_session(
509        &self,
510        session_id: &str,
511        timeout: Duration,
512    ) -> Result<SessionHandle> {
513        tokio::time::timeout(timeout, async {
514            loop {
515                match self.session(session_id.to_owned()).await {
516                    Ok(handle) => return Ok(handle),
517                    Err(error) => {
518                        tracing::trace!(session_id, "waiting for session actor: {error:#}");
519                        tokio::time::sleep(Duration::from_millis(25)).await;
520                    }
521                }
522            }
523        })
524        .await
525        .with_context(|| {
526            format!(
527                "session {session_id} did not become available within {} seconds",
528                timeout.as_secs()
529            )
530        })?
531    }
532}
533
534impl fmt::Debug for SessionControl {
535    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
536        formatter.write_str("SessionControl(..)")
537    }
538}
539
540pub trait ReviewerStagerBackend: Send + Sync {
541    fn stage(
542        &self,
543        config: Config,
544        session: SessionRecord,
545        profile_id: String,
546        generation: u64,
547        cancelled: Arc<std::sync::atomic::AtomicBool>,
548    ) -> Result<ReviewerLaunchConfig>;
549}
550
551#[derive(Clone)]
552pub struct ReviewerStager {
553    backend: Arc<dyn ReviewerStagerBackend>,
554}
555
556impl ReviewerStager {
557    pub fn new(backend: impl ReviewerStagerBackend + 'static) -> Self {
558        Self {
559            backend: Arc::new(backend),
560        }
561    }
562
563    pub fn stage(
564        &self,
565        config: Config,
566        session: SessionRecord,
567        profile_id: String,
568        generation: u64,
569        cancelled: Arc<std::sync::atomic::AtomicBool>,
570    ) -> Result<ReviewerLaunchConfig> {
571        self.backend
572            .stage(config, session, profile_id, generation, cancelled)
573    }
574
575    #[doc(hidden)]
576    pub fn unavailable(message: impl Into<String>) -> Self {
577        Self::new(UnavailableReviewerStager(message.into()))
578    }
579}
580
581impl fmt::Debug for ReviewerStager {
582    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
583        formatter.write_str("ReviewerStager(..)")
584    }
585}
586
587struct UnavailableReviewerStager(String);
588
589impl ReviewerStagerBackend for UnavailableReviewerStager {
590    fn stage(
591        &self,
592        _config: Config,
593        _session: SessionRecord,
594        _profile_id: String,
595        _generation: u64,
596        _cancelled: Arc<std::sync::atomic::AtomicBool>,
597    ) -> Result<ReviewerLaunchConfig> {
598        anyhow::bail!(self.0.clone())
599    }
600}
601
602pub fn new_command_id(prefix: &str) -> Result<String> {
603    ensure!(!prefix.trim().is_empty(), "command ID prefix is required");
604    let mut random = [0_u8; 16];
605    getrandom::fill(&mut random)
606        .map_err(|error| anyhow::anyhow!("generate command ID: {error}"))?;
607    Ok(format!("{prefix}-{}", mj_core::hex::lower_hex(random)))
608}
609
610/// A stopped session and a manager that resolves its live replacement.
611///
612/// Chat's cross-crate tests use this hand-written client fake to verify actor
613/// replacement without depending on the controller implementation crate.
614#[doc(hidden)]
615pub struct ReplacementSessionTestFixture {
616    pub stopped: SessionHandle,
617    pub control: SessionControl,
618    pub submitted: tokio::sync::mpsc::UnboundedReceiver<RelayCommand>,
619    pub replacement_view: Arc<tokio::sync::watch::Sender<ManagedSessionView>>,
620}
621
622#[derive(Clone)]
623struct ReplacementTestSession {
624    #[cfg(test)]
625    history: Option<tokio::sync::mpsc::UnboundedSender<HistoryTestRequest>>,
626    session_id: String,
627    stopped: bool,
628    accepted_ordinal: u64,
629    submitted: Option<tokio::sync::mpsc::UnboundedSender<RelayCommand>>,
630    view: tokio::sync::watch::Receiver<ManagedSessionView>,
631    _view_guard: Option<Arc<tokio::sync::watch::Sender<ManagedSessionView>>>,
632}
633
634impl SessionHandleBackend for ReplacementTestSession {
635    fn search_prompts(
636        &self,
637        _bundle_id: String,
638        _scope: mj_core::storage::HistoryScope,
639        _query: String,
640    ) -> BoxFuture<'_, Result<Vec<mj_core::storage::PromptHistoryEntry>>> {
641        #[cfg(test)]
642        if let Some(history) = &self.history {
643            let (response, result) = tokio::sync::oneshot::channel();
644            let sent = history.send(HistoryTestRequest {
645                bundle_id: _bundle_id,
646                scope: _scope,
647                query: _query,
648                response,
649            });
650            return Box::pin(async move {
651                sent.map_err(|_| anyhow::anyhow!("history backend closed"))?;
652                result
653                    .await
654                    .map_err(|_| anyhow::anyhow!("history response dropped"))?
655            });
656        }
657        Box::pin(async { Ok(Vec::new()) })
658    }
659    fn review_state(&self) -> BoxFuture<'_, Result<ReviewState>> {
660        Box::pin(async { Ok(ReviewState::default()) })
661    }
662
663    fn config_result(&self, _command_id: String) -> BoxFuture<'_, Result<Option<Option<String>>>> {
664        Box::pin(async { Ok(None) })
665    }
666    fn clone_box(&self) -> Box<dyn SessionHandleBackend> {
667        Box::new(self.clone())
668    }
669
670    fn session_id(&self) -> &str {
671        &self.session_id
672    }
673
674    fn view(&self) -> ManagedSessionView {
675        self.view.borrow().clone()
676    }
677
678    fn is_stopped(&self) -> bool {
679        self.stopped
680    }
681
682    fn has_changed(&self) -> Result<bool> {
683        self.view.has_changed().context("session manager stopped")
684    }
685
686    fn changed(&mut self) -> BoxFuture<'_, Result<ManagedSessionView>> {
687        Box::pin(async move {
688            self.view
689                .changed()
690                .await
691                .context("session manager stopped")?;
692            Ok(self.view())
693        })
694    }
695
696    fn enqueue_submit(
697        &self,
698        _command_id: String,
699        command: RelayCommand,
700    ) -> BoxFuture<'_, Result<PendingRelaySubmit>> {
701        let submitted = self.submitted.clone();
702        let stopped = self.stopped;
703        let accepted_ordinal = self.accepted_ordinal;
704        Box::pin(async move {
705            ensure!(!stopped, "session manager stopped");
706            let submitted = submitted.context("unsupported test operation")?;
707            submitted
708                .send(command)
709                .context("test submit observer stopped")?;
710            Ok(PendingRelaySubmit::new(Box::pin(async move {
711                Ok(accepted_ordinal)
712            })))
713        })
714    }
715
716    fn enqueue_sync(&self) -> BoxFuture<'_, Result<PendingRelaySync>> {
717        let stopped = self.stopped;
718        Box::pin(async move {
719            ensure!(!stopped, "session manager stopped");
720            Ok(PendingRelaySync::new(Box::pin(async { Ok(()) })))
721        })
722    }
723
724    fn respond_elicitation(
725        &self,
726        _elicitation_id: String,
727        _response: ElicitationResponse,
728    ) -> BoxFuture<'_, Result<()>> {
729        Box::pin(async { anyhow::bail!("unsupported test operation") })
730    }
731
732    fn stop_background_task(&self, _background_task_id: String) -> BoxFuture<'_, Result<()>> {
733        Box::pin(async { anyhow::bail!("unsupported test operation") })
734    }
735
736    fn reviewer(
737        &self,
738        _role: Option<String>,
739        _action: ReviewerAction,
740    ) -> BoxFuture<'_, Result<ReviewerOutcome>> {
741        Box::pin(async { anyhow::bail!("unsupported test operation") })
742    }
743}
744
745struct ReplacementTestControl {
746    session_id: String,
747    replacement: SessionHandle,
748}
749
750impl SessionControlBackend for ReplacementTestControl {
751    fn session(&self, session_id: String) -> BoxFuture<'_, Result<SessionHandle>> {
752        Box::pin(async move {
753            ensure!(
754                session_id == self.session_id,
755                "session {session_id} is not managed"
756            );
757            Ok(self.replacement.clone())
758        })
759    }
760}
761
762#[doc(hidden)]
763pub fn replacement_session_test_fixture(
764    session_id: &str,
765    accepted_ordinal: u64,
766) -> ReplacementSessionTestFixture {
767    let (stopped_view_tx, stopped_view) =
768        tokio::sync::watch::channel(ManagedSessionView::default());
769    drop(stopped_view_tx);
770    let stopped = SessionHandle::new(ReplacementTestSession {
771        #[cfg(test)]
772        history: None,
773        session_id: session_id.to_owned(),
774        stopped: true,
775        accepted_ordinal,
776        submitted: None,
777        view: stopped_view,
778        _view_guard: None,
779    });
780
781    let (view_tx, view) = tokio::sync::watch::channel(ManagedSessionView::default());
782    let view_tx = Arc::new(view_tx);
783    let (submitted_tx, submitted) = tokio::sync::mpsc::unbounded_channel();
784    let replacement = SessionHandle::new(ReplacementTestSession {
785        #[cfg(test)]
786        history: None,
787        session_id: session_id.to_owned(),
788        stopped: false,
789        accepted_ordinal,
790        submitted: Some(submitted_tx),
791        view,
792        _view_guard: Some(view_tx.clone()),
793    });
794    let control = SessionControl::new(ReplacementTestControl {
795        session_id: session_id.to_owned(),
796        replacement,
797    });
798    ReplacementSessionTestFixture {
799        stopped,
800        control,
801        submitted,
802        replacement_view: view_tx,
803    }
804}
805
806#[cfg(test)]
807struct HistoryTestRequest {
808    bundle_id: String,
809    scope: mj_core::storage::HistoryScope,
810    query: String,
811    response: tokio::sync::oneshot::Sender<Result<Vec<mj_core::storage::PromptHistoryEntry>>>,
812}
813
814#[cfg(test)]
815mod storage_tests {
816    use super::*;
817    use mj_core::storage::{HistoryScope, PromptHistoryEntry};
818
819    #[tokio::test]
820    async fn history_search_yields_until_backend_responds_and_propagates_failures() {
821        let (history, mut requests) = tokio::sync::mpsc::unbounded_channel();
822        let (view_guard, view) = tokio::sync::watch::channel(ManagedSessionView::default());
823        let session = SessionHandle::new(ReplacementTestSession {
824            history: Some(history),
825            session_id: "session".into(),
826            stopped: false,
827            accepted_ordinal: 0,
828            submitted: None,
829            view,
830            _view_guard: Some(Arc::new(view_guard)),
831        });
832        let search =
833            session.search_prompts("bundle".into(), HistoryScope::Project, "needle".into());
834        tokio::pin!(search);
835        let request = tokio::select! {
836            biased;
837            result = &mut search => panic!("search completed before storage replied: {result:?}"),
838            request = requests.recv() => request.unwrap(),
839        };
840        assert_eq!(request.bundle_id, "bundle");
841        assert_eq!(request.scope, HistoryScope::Project);
842        assert_eq!(request.query, "needle");
843        request
844            .response
845            .send(Err(anyhow::anyhow!("storage unavailable")))
846            .unwrap();
847        assert!(
848            search
849                .await
850                .unwrap_err()
851                .to_string()
852                .contains("storage unavailable")
853        );
854
855        let search = session.search_prompts("bundle".into(), HistoryScope::Project, "retry".into());
856        tokio::pin!(search);
857        let request = tokio::select! {
858            biased;
859            result = &mut search => panic!("retry completed before storage replied: {result:?}"),
860            request = requests.recv() => request.unwrap(),
861        };
862        request
863            .response
864            .send(Ok(vec![PromptHistoryEntry {
865                id: 1,
866                session_id: "session".into(),
867                text: "retry works".into(),
868            }]))
869            .unwrap();
870        assert_eq!(search.await.unwrap()[0].text, "retry works");
871    }
872}