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    /// Replace the worker's ephemeral GitHub token before a controller-owned
249    /// operation that needs the newest installation credential.
250    fn install_github_token(&self, _token: String) -> BoxFuture<'_, Result<()>> {
251        Box::pin(async { anyhow::bail!("GitHub token refresh is unavailable") })
252    }
253    fn reviewer(
254        &self,
255        role: Option<String>,
256        action: ReviewerAction,
257    ) -> BoxFuture<'_, Result<ReviewerOutcome>>;
258}
259
260pub struct SessionHandle {
261    backend: Box<dyn SessionHandleBackend>,
262}
263
264impl SessionHandle {
265    pub async fn transcript_history(
266        &self,
267        before: Option<mj_core::storage::TranscriptCursor>,
268    ) -> Result<mj_core::storage::TranscriptHistoryPage> {
269        self.backend.transcript_history(before).await
270    }
271
272    pub async fn search_prompts(
273        &self,
274        bundle_id: String,
275        scope: mj_core::storage::HistoryScope,
276        query: String,
277    ) -> Result<Vec<mj_core::storage::PromptHistoryEntry>> {
278        self.backend.search_prompts(bundle_id, scope, query).await
279    }
280    pub async fn review_state(&self) -> Result<ReviewState> {
281        self.backend.review_state().await
282    }
283
284    pub async fn resolve_review_settings(
285        &self,
286        cancelled: Arc<std::sync::atomic::AtomicBool>,
287    ) -> Result<mj_core::review::settings::ResolvedReviewSettings> {
288        self.backend.resolve_review_settings(cancelled).await
289    }
290
291    pub fn new(backend: impl SessionHandleBackend + 'static) -> Self {
292        Self {
293            backend: Box::new(backend),
294        }
295    }
296
297    pub fn session_id(&self) -> &str {
298        self.backend.session_id()
299    }
300
301    pub fn view(&self) -> ManagedSessionView {
302        self.backend.view()
303    }
304
305    pub fn is_stopped(&self) -> bool {
306        self.backend.is_stopped()
307    }
308
309    pub fn has_changed(&self) -> Result<bool> {
310        self.backend.has_changed()
311    }
312
313    pub async fn changed(&mut self) -> Result<ManagedSessionView> {
314        self.backend.changed().await
315    }
316
317    pub async fn submit(&self, command_id: String, command: RelayCommand) -> Result<u64> {
318        self.enqueue_submit(command_id, command).await?.wait().await
319    }
320
321    /// Apply a setting and wait for its durable success or rejection.
322    pub async fn set_config(&self, key: String, value: String) -> Result<()> {
323        self.set_config_with_id(new_command_id("set-config")?, key, value)
324            .await
325            .map(|_| ())
326    }
327
328    /// Apply a setting under a caller-chosen command id and return the
329    /// ordinal at which the relay accepted it. A caller that retries with the
330    /// same id after a lost reply gets the original ordinal back.
331    pub async fn set_config_with_id(
332        &self,
333        command_id: String,
334        key: String,
335        value: String,
336    ) -> Result<u64> {
337        self.apply_configuration_command(command_id, RelayCommand::SetConfig { key, value })
338            .await
339    }
340
341    async fn apply_configuration_command(
342        &self,
343        command_id: String,
344        command: RelayCommand,
345    ) -> Result<u64> {
346        let ordinal = self.submit(command_id.clone(), command).await?;
347        tokio::time::timeout(Duration::from_secs(60), async {
348            loop {
349                if let Some(error) = self.backend.config_result(command_id.clone()).await? {
350                    if let Some(error) = error {
351                        anyhow::bail!("{error}");
352                    }
353                    self.sync_now().await?;
354                    return Ok(ordinal);
355                }
356                ensure!(
357                    !self.is_stopped(),
358                    "session stopped while applying configuration"
359                );
360                if let Some(error) = self.view().error {
361                    anyhow::bail!("configuration connection failed: {}", error.detail());
362                }
363                tokio::time::sleep(Duration::from_millis(50)).await;
364            }
365        })
366        .await
367        .context("configuration command did not complete within 60 seconds")?
368    }
369
370    /// A follow-up prompt must wait for the mode change, not just admission.
371    pub async fn apply_plan_control(
372        &self,
373        command_id: String,
374        control: mj_core::acp::PlanControl,
375    ) -> Result<()> {
376        match control {
377            mj_core::acp::PlanControl::RestoreExecutionMode => self
378                .apply_configuration_command(command_id, RelayCommand::RestoreExecutionMode)
379                .await
380                .map(|_| ()),
381            mj_core::acp::PlanControl::SetConfig { key, value } => self
382                .set_config_with_id(command_id, key, value)
383                .await
384                .map(|_| ()),
385            mj_core::acp::PlanControl::SetSessionMode { mode_id } => {
386                self.submit(
387                    command_id,
388                    RelayCommand::SetSessionMode {
389                        mode_id: mode_id.clone(),
390                    },
391                )
392                .await?;
393                tokio::time::timeout(Duration::from_secs(60), async {
394                    loop {
395                        self.sync_now().await?;
396                        if self
397                            .view()
398                            .snapshot
399                            .as_ref()
400                            .and_then(|snapshot| snapshot.operational.modes.as_ref())
401                            .is_some_and(|modes| modes.current_mode_id.to_string() == mode_id)
402                        {
403                            return Ok(());
404                        }
405                        ensure!(!self.is_stopped(), "session stopped while changing mode");
406                        tokio::time::sleep(Duration::from_millis(50)).await;
407                    }
408                })
409                .await
410                .context("session mode change did not complete within 60 seconds")?
411            }
412        }
413    }
414
415    pub async fn enqueue_submit(
416        &self,
417        command_id: String,
418        command: RelayCommand,
419    ) -> Result<PendingRelaySubmit> {
420        self.backend.enqueue_submit(command_id, command).await
421    }
422
423    pub async fn sync_now(&self) -> Result<()> {
424        self.enqueue_sync().await?.wait().await
425    }
426
427    pub async fn enqueue_sync(&self) -> Result<PendingRelaySync> {
428        self.backend.enqueue_sync().await
429    }
430
431    pub async fn respond_elicitation(
432        &self,
433        elicitation_id: String,
434        response: ElicitationResponse,
435    ) -> Result<()> {
436        self.backend
437            .respond_elicitation(elicitation_id, response)
438            .await
439    }
440
441    pub async fn stop_background_task(&self, background_task_id: String) -> Result<()> {
442        self.backend.stop_background_task(background_task_id).await
443    }
444
445    pub async fn install_github_token(&self, token: String) -> Result<()> {
446        self.backend.install_github_token(token).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
