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

mj_client/
session.rs
