Skip to main content

winx_code_agent/runtime/
session_store.rs

1//! Embedded session index used by the shell runtime.
2
3use std::collections::HashMap;
4use std::sync::{Arc, LazyLock, Mutex as StdMutex, Weak};
5
6use crate::errors::Result;
7use crate::state::pty::{PtyShell, SharedPtyShell};
8use crate::tools::background_shell::{BackgroundShellManager, ExitedShellInfo};
9
10/// Stable address of a terminal within one logical session.
11#[derive(Clone, Debug, Eq, PartialEq)]
12pub enum ShellTarget {
13    Main,
14    Background(String),
15}
16
17/// Single index for main and background terminals in the embedded runtime.
18///
19/// Main shells remain owned by `BashState`; weak references keep this index from
20/// extending their lifetime. Background ownership stays in the existing manager
21/// until the daemon runtime takes ownership in a later migration step.
22#[derive(Debug, Default)]
23pub struct SessionStore {
24    mains: HashMap<String, Weak<tokio::sync::Mutex<Option<PtyShell>>>>,
25    operation_barriers: HashMap<String, Weak<tokio::sync::RwLock<()>>>,
26    backgrounds: BackgroundShellManager,
27}
28
29impl SessionStore {
30    pub fn new() -> Self {
31        Self::default()
32    }
33
34    pub fn bind_main(&mut self, thread_id: &str, shell: &SharedPtyShell) {
35        self.mains.retain(|_, shell| shell.strong_count() > 0);
36        self.operation_barriers.retain(|_, barrier| barrier.strong_count() > 0);
37        self.mains.insert(thread_id.to_string(), Arc::downgrade(shell));
38    }
39
40    /// Session-wide ordering barrier. Shell actions and interrupts hold a read
41    /// guard while reset/configure holds a write guard, so an incarnation can
42    /// never change between execution-token validation and the PTY operation.
43    pub(crate) fn operation_barrier(&mut self, thread_id: &str) -> Arc<tokio::sync::RwLock<()>> {
44        if let Some(barrier) =
45            self.operation_barriers.get(thread_id).and_then(std::sync::Weak::upgrade)
46        {
47            return barrier;
48        }
49        let barrier = Arc::new(tokio::sync::RwLock::new(()));
50        self.operation_barriers.insert(thread_id.to_string(), Arc::downgrade(&barrier));
51        barrier
52    }
53
54    pub fn resolve(&self, thread_id: &str, target: &ShellTarget) -> Option<SharedPtyShell> {
55        match target {
56            ShellTarget::Main => self.mains.get(thread_id).and_then(Weak::upgrade),
57            ShellTarget::Background(id) => self.backgrounds.get_shell(thread_id, id),
58        }
59    }
60
61    pub(crate) fn register_shell(&mut self, thread_id: &str, shell: PtyShell) -> Result<String> {
62        self.backgrounds.register_shell(thread_id, shell)
63    }
64
65    pub(crate) fn get_shell(&self, thread_id: &str, bg_command_id: &str) -> Option<SharedPtyShell> {
66        self.backgrounds.get_shell(thread_id, bg_command_id)
67    }
68
69    pub(crate) fn remove_shell(&mut self, bg_command_id: &str) -> bool {
70        self.backgrounds.remove_shell(bg_command_id)
71    }
72
73    pub(crate) fn prune_finished_shells(&mut self) {
74        self.backgrounds.prune_finished_shells();
75    }
76
77    pub(crate) fn peek_tombstone(
78        &self,
79        thread_id: &str,
80        bg_command_id: &str,
81    ) -> Option<ExitedShellInfo> {
82        self.backgrounds.peek_tombstone(thread_id, bg_command_id)
83    }
84
85    /// Live background shell ids owned by `thread_id`.
86    pub(crate) fn background_ids_for(&self, thread_id: &str) -> Vec<String> {
87        self.backgrounds.ids_for(thread_id)
88    }
89
90    pub(crate) fn get_running_info(&mut self, thread_id: &str) -> String {
91        self.backgrounds.get_running_info(thread_id)
92    }
93}
94
95static EMBEDDED_SESSION_STORE: LazyLock<StdMutex<SessionStore>> =
96    LazyLock::new(|| StdMutex::new(SessionStore::new()));
97
98/// Lock the embedded session index, recovering from poisoning like the legacy
99/// background manager did.
100pub(crate) fn lock_session_store() -> std::sync::MutexGuard<'static, SessionStore> {
101    EMBEDDED_SESSION_STORE.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
102}