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