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 hel::hel_config::HelConfig;
12use hel::hel_elicitation::ElicitationResponse;
13use hel::hel_state::{ManagedSessionSnapshot, SessionRecord};
14use hel::hel_worker::{
15    AnalyzeDeltaRepository, RelayCommand, RelayCursor, RelayEvent, RelayOperationalState, RepoDelta,
16};
17use hel::hel_worker_launch::ReviewerLaunchConfig;
18
19pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
20
21#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
22#[serde(tag = "kind", content = "detail", rename_all = "snake_case")]
23pub enum ViewError {
24    Unreachable(String),
25    TargetMissing(String),
26    ProjectionIntegrity(String),
27}
28
29impl ViewError {
30    pub fn detail(&self) -> &str {
31        match self {
32            Self::Unreachable(detail)
33            | Self::TargetMissing(detail)
34            | Self::ProjectionIntegrity(detail) => detail,
35        }
36    }
37}
38
39#[derive(Debug, Clone, PartialEq, Default)]
40pub struct ManagedSessionView {
41    pub snapshot: Option<ManagedSessionSnapshot>,
42    pub connected: bool,
43    pub error: Option<ViewError>,
44}
45
46#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
47pub struct RelayAttachment {
48    pub state: RelayOperationalState,
49    pub events: Vec<RelayEvent>,
50    pub through_ordinal: u64,
51    pub through_digest: String,
52}
53
54#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
55pub struct StartedReviewer {
56    pub native_session_id: Option<String>,
57    pub config_options: Vec<SessionConfigOption>,
58    pub reused: bool,
59    pub state: RelayOperationalState,
60}
61
62#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
63#[serde(rename_all = "snake_case")]
64pub enum ReviewerAction {
65    Start {
66        config: Box<ReviewerLaunchConfig>,
67    },
68    Submit {
69        command_id: String,
70        command: RelayCommand,
71    },
72    Attach {
73        after_ordinal: u64,
74        after_digest: String,
75    },
76    Acknowledge {
77        through_ordinal: u64,
78        through_digest: String,
79    },
80    Status,
81    RespondElicitation {
82        elicitation_id: String,
83        response: ElicitationResponse,
84    },
85    Pause,
86    CaptureDelta {
87        baselines: std::collections::BTreeMap<std::path::PathBuf, String>,
88    },
89    AdvanceBaseline {
90        trees: std::collections::BTreeMap<std::path::PathBuf, String>,
91    },
92    AnalyzeDelta {
93        repositories: Vec<AnalyzeDeltaRepository>,
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 => "reviewer_pause",
108            Self::CaptureDelta { .. } => "reviewer_capture_delta",
109            Self::AdvanceBaseline { .. } => "reviewer_advance_baseline",
110            Self::AnalyzeDelta { .. } => "reviewer_analyze_delta",
111            Self::TakeLaneDispatches => "reviewer_take_lane_dispatches",
112        }
113    }
114}
115
116#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
117#[serde(rename_all = "snake_case")]
118pub enum ReviewerOutcome {
119    Started(Box<StartedReviewer>),
120    Accepted {
121        ordinal: u64,
122    },
123    Attached(Box<RelayAttachment>),
124    Acknowledged(RelayCursor),
125    Status(Box<RelayOperationalState>),
126    ElicitationResolved,
127    Paused,
128    Delta {
129        repositories: Vec<RepoDelta>,
130    },
131    BaselineAdvanced,
132    ChangedFunctions {
133        packet: String,
134    },
135    LaneDispatches {
136        requests: Vec<hel::hel_review::lanes::ReviewSubagentRequest>,
137    },
138}
139
140pub struct PendingRelaySubmit {
141    completion: BoxFuture<'static, Result<u64>>,
142}
143
144impl PendingRelaySubmit {
145    pub fn new(completion: BoxFuture<'static, Result<u64>>) -> Self {
146        Self { completion }
147    }
148
149    pub async fn wait(self) -> Result<u64> {
150        self.completion.await
151    }
152}
153
154pub struct PendingRelaySync {
155    completion: BoxFuture<'static, Result<()>>,
156}
157
158impl PendingRelaySync {
159    pub fn new(completion: BoxFuture<'static, Result<()>>) -> Self {
160        Self { completion }
161    }
162
163    pub async fn wait(self) -> Result<()> {
164        self.completion.await
165    }
166}
167
168pub trait SessionHandleBackend: Send + Sync {
169    fn clone_box(&self) -> Box<dyn SessionHandleBackend>;
170    fn session_id(&self) -> &str;
171    fn view(&self) -> ManagedSessionView;
172    fn is_stopped(&self) -> bool;
173    fn has_changed(&self) -> Result<bool>;
174    fn changed(&mut self) -> BoxFuture<'_, Result<ManagedSessionView>>;
175    fn enqueue_submit(
176        &self,
177        command_id: String,
178        command: RelayCommand,
179    ) -> BoxFuture<'_, Result<PendingRelaySubmit>>;
180    fn enqueue_sync(&self) -> BoxFuture<'_, Result<PendingRelaySync>>;
181    fn respond_elicitation(
182        &self,
183        elicitation_id: String,
184        response: ElicitationResponse,
185    ) -> BoxFuture<'_, Result<()>>;
186    fn stop_background_task(&self, background_task_id: String) -> BoxFuture<'_, Result<()>>;
187    fn reviewer(
188        &self,
189        role: Option<String>,
190        action: ReviewerAction,
191    ) -> BoxFuture<'_, Result<ReviewerOutcome>>;
192}
193
194pub struct SessionHandle {
195    backend: Box<dyn SessionHandleBackend>,
196}
197
198impl SessionHandle {
199    pub fn new(backend: impl SessionHandleBackend + 'static) -> Self {
200        Self {
201            backend: Box::new(backend),
202        }
203    }
204
205    pub fn session_id(&self) -> &str {
206        self.backend.session_id()
207    }
208
209    pub fn view(&self) -> ManagedSessionView {
210        self.backend.view()
211    }
212
213    pub fn is_stopped(&self) -> bool {
214        self.backend.is_stopped()
215    }
216
217    pub fn has_changed(&self) -> Result<bool> {
218        self.backend.has_changed()
219    }
220
221    pub async fn changed(&mut self) -> Result<ManagedSessionView> {
222        self.backend.changed().await
223    }
224
225    pub async fn submit(&self, command_id: String, command: RelayCommand) -> Result<u64> {
226        self.enqueue_submit(command_id, command).await?.wait().await
227    }
228
229    pub async fn enqueue_submit(
230        &self,
231        command_id: String,
232        command: RelayCommand,
233    ) -> Result<PendingRelaySubmit> {
234        self.backend.enqueue_submit(command_id, command).await
235    }
236
237    pub async fn sync_now(&self) -> Result<()> {
238        self.enqueue_sync().await?.wait().await
239    }
240
241    pub async fn enqueue_sync(&self) -> Result<PendingRelaySync> {
242        self.backend.enqueue_sync().await
243    }
244
245    pub async fn respond_elicitation(
246        &self,
247        elicitation_id: String,
248        response: ElicitationResponse,
249    ) -> Result<()> {
250        self.backend
251            .respond_elicitation(elicitation_id, response)
252            .await
253    }
254
255    pub async fn stop_background_task(&self, background_task_id: String) -> Result<()> {
256        self.backend.stop_background_task(background_task_id).await
257    }
258
259    pub async fn reviewer(&self, action: ReviewerAction) -> Result<ReviewerOutcome> {
260        self.reviewer_as(None, action).await
261    }
262
263    pub async fn reviewer_as(
264        &self,
265        role: Option<String>,
266        action: ReviewerAction,
267    ) -> Result<ReviewerOutcome> {
268        self.backend.reviewer(role, action).await
269    }
270}
271
272impl Clone for SessionHandle {
273    fn clone(&self) -> Self {
274        Self {
275            backend: self.backend.clone_box(),
276        }
277    }
278}
279
280impl fmt::Debug for SessionHandle {
281    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
282        formatter
283            .debug_struct("SessionHandle")
284            .field("session_id", &self.session_id())
285            .finish_non_exhaustive()
286    }
287}
288
289pub trait SessionControlBackend: Send + Sync {
290    fn session(&self, session_id: String) -> BoxFuture<'_, Result<SessionHandle>>;
291}
292
293#[derive(Clone)]
294pub struct SessionControl {
295    backend: Arc<dyn SessionControlBackend>,
296}
297
298impl SessionControl {
299    pub fn new(backend: impl SessionControlBackend + 'static) -> Self {
300        Self {
301            backend: Arc::new(backend),
302        }
303    }
304
305    pub async fn session(&self, session_id: impl Into<String>) -> Result<SessionHandle> {
306        self.backend.session(session_id.into()).await
307    }
308
309    pub async fn wait_for_session(
310        &self,
311        session_id: &str,
312        timeout: Duration,
313    ) -> Result<SessionHandle> {
314        tokio::time::timeout(timeout, async {
315            loop {
316                match self.session(session_id.to_owned()).await {
317                    Ok(handle) => return Ok(handle),
318                    Err(error) => {
319                        tracing::trace!(session_id, "waiting for session actor: {error:#}");
320                        tokio::time::sleep(Duration::from_millis(25)).await;
321                    }
322                }
323            }
324        })
325        .await
326        .with_context(|| {
327            format!(
328                "session {session_id} did not become available within {} seconds",
329                timeout.as_secs()
330            )
331        })?
332    }
333}
334
335impl fmt::Debug for SessionControl {
336    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
337        formatter.write_str("SessionControl(..)")
338    }
339}
340
341pub trait ReviewerStagerBackend: Send + Sync {
342    fn stage(
343        &self,
344        config: HelConfig,
345        session: SessionRecord,
346        profile_id: String,
347        generation: u64,
348    ) -> Result<ReviewerLaunchConfig>;
349}
350
351#[derive(Clone)]
352pub struct ReviewerStager {
353    backend: Arc<dyn ReviewerStagerBackend>,
354}
355
356impl ReviewerStager {
357    pub fn new(backend: impl ReviewerStagerBackend + 'static) -> Self {
358        Self {
359            backend: Arc::new(backend),
360        }
361    }
362
363    pub fn stage(
364        &self,
365        config: HelConfig,
366        session: SessionRecord,
367        profile_id: String,
368        generation: u64,
369    ) -> Result<ReviewerLaunchConfig> {
370        self.backend.stage(config, session, profile_id, generation)
371    }
372
373    #[doc(hidden)]
374    pub fn unavailable(message: impl Into<String>) -> Self {
375        Self::new(UnavailableReviewerStager(message.into()))
376    }
377}
378
379impl fmt::Debug for ReviewerStager {
380    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
381        formatter.write_str("ReviewerStager(..)")
382    }
383}
384
385struct UnavailableReviewerStager(String);
386
387impl ReviewerStagerBackend for UnavailableReviewerStager {
388    fn stage(
389        &self,
390        _config: HelConfig,
391        _session: SessionRecord,
392        _profile_id: String,
393        _generation: u64,
394    ) -> Result<ReviewerLaunchConfig> {
395        anyhow::bail!(self.0.clone())
396    }
397}
398
399pub fn new_command_id(prefix: &str) -> Result<String> {
400    ensure!(!prefix.trim().is_empty(), "command ID prefix is required");
401    let mut random = [0_u8; 16];
402    getrandom::fill(&mut random)
403        .map_err(|error| anyhow::anyhow!("generate command ID: {error}"))?;
404    Ok(format!("{prefix}-{}", hex(&random)))
405}
406
407fn hex(bytes: &[u8]) -> String {
408    const DIGITS: &[u8; 16] = b"0123456789abcdef";
409    let mut output = String::with_capacity(bytes.len() * 2);
410    for byte in bytes {
411        output.push(char::from(DIGITS[usize::from(byte >> 4)]));
412        output.push(char::from(DIGITS[usize::from(byte & 0x0f)]));
413    }
414    output
415}
416
417/// A stopped session and a manager that resolves its live replacement.
418///
419/// Chat's cross-crate tests use this hand-written client fake to verify actor
420/// replacement without depending on the controller implementation crate.
421#[doc(hidden)]
422pub struct ReplacementSessionTestFixture {
423    pub stopped: SessionHandle,
424    pub control: SessionControl,
425    pub submitted: tokio::sync::mpsc::UnboundedReceiver<RelayCommand>,
426}
427
428#[derive(Clone)]
429struct ReplacementTestSession {
430    session_id: String,
431    stopped: bool,
432    accepted_ordinal: u64,
433    submitted: Option<tokio::sync::mpsc::UnboundedSender<RelayCommand>>,
434    view: tokio::sync::watch::Receiver<ManagedSessionView>,
435    _view_guard: Option<Arc<tokio::sync::watch::Sender<ManagedSessionView>>>,
436}
437
438impl SessionHandleBackend for ReplacementTestSession {
439    fn clone_box(&self) -> Box<dyn SessionHandleBackend> {
440        Box::new(self.clone())
441    }
442
443    fn session_id(&self) -> &str {
444        &self.session_id
445    }
446
447    fn view(&self) -> ManagedSessionView {
448        self.view.borrow().clone()
449    }
450
451    fn is_stopped(&self) -> bool {
452        self.stopped
453    }
454
455    fn has_changed(&self) -> Result<bool> {
456        self.view.has_changed().context("session manager stopped")
457    }
458
459    fn changed(&mut self) -> BoxFuture<'_, Result<ManagedSessionView>> {
460        Box::pin(async move {
461            self.view
462                .changed()
463                .await
464                .context("session manager stopped")?;
465            Ok(self.view())
466        })
467    }
468
469    fn enqueue_submit(
470        &self,
471        _command_id: String,
472        command: RelayCommand,
473    ) -> BoxFuture<'_, Result<PendingRelaySubmit>> {
474        let submitted = self.submitted.clone();
475        let stopped = self.stopped;
476        let accepted_ordinal = self.accepted_ordinal;
477        Box::pin(async move {
478            ensure!(!stopped, "session manager stopped");
479            let submitted = submitted.context("unsupported test operation")?;
480            submitted
481                .send(command)
482                .context("test submit observer stopped")?;
483            Ok(PendingRelaySubmit::new(Box::pin(async move {
484                Ok(accepted_ordinal)
485            })))
486        })
487    }
488
489    fn enqueue_sync(&self) -> BoxFuture<'_, Result<PendingRelaySync>> {
490        let stopped = self.stopped;
491        Box::pin(async move {
492            ensure!(!stopped, "session manager stopped");
493            Ok(PendingRelaySync::new(Box::pin(async { Ok(()) })))
494        })
495    }
496
497    fn respond_elicitation(
498        &self,
499        _elicitation_id: String,
500        _response: ElicitationResponse,
501    ) -> BoxFuture<'_, Result<()>> {
502        Box::pin(async { anyhow::bail!("unsupported test operation") })
503    }
504
505    fn stop_background_task(&self, _background_task_id: String) -> BoxFuture<'_, Result<()>> {
506        Box::pin(async { anyhow::bail!("unsupported test operation") })
507    }
508
509    fn reviewer(
510        &self,
511        _role: Option<String>,
512        _action: ReviewerAction,
513    ) -> BoxFuture<'_, Result<ReviewerOutcome>> {
514        Box::pin(async { anyhow::bail!("unsupported test operation") })
515    }
516}
517
518struct ReplacementTestControl {
519    session_id: String,
520    replacement: SessionHandle,
521}
522
523impl SessionControlBackend for ReplacementTestControl {
524    fn session(&self, session_id: String) -> BoxFuture<'_, Result<SessionHandle>> {
525        Box::pin(async move {
526            ensure!(
527                session_id == self.session_id,
528                "session {session_id} is not managed"
529            );
530            Ok(self.replacement.clone())
531        })
532    }
533}
534
535#[doc(hidden)]
536pub fn replacement_session_test_fixture(
537    session_id: &str,
538    accepted_ordinal: u64,
539) -> ReplacementSessionTestFixture {
540    let (stopped_view_tx, stopped_view) =
541        tokio::sync::watch::channel(ManagedSessionView::default());
542    drop(stopped_view_tx);
543    let stopped = SessionHandle::new(ReplacementTestSession {
544        session_id: session_id.to_owned(),
545        stopped: true,
546        accepted_ordinal,
547        submitted: None,
548        view: stopped_view,
549        _view_guard: None,
550    });
551
552    let (view_tx, view) = tokio::sync::watch::channel(ManagedSessionView::default());
553    let (submitted_tx, submitted) = tokio::sync::mpsc::unbounded_channel();
554    let replacement = SessionHandle::new(ReplacementTestSession {
555        session_id: session_id.to_owned(),
556        stopped: false,
557        accepted_ordinal,
558        submitted: Some(submitted_tx),
559        view,
560        _view_guard: Some(Arc::new(view_tx)),
561    });
562    let control = SessionControl::new(ReplacementTestControl {
563        session_id: session_id.to_owned(),
564        replacement,
565    });
566    ReplacementSessionTestFixture {
567        stopped,
568        control,
569        submitted,
570    }
571}
mj_client/session.rs

mj_client/
session.rs
