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    backgrounds: BackgroundShellManager,
26}
27
28impl SessionStore {
29    pub fn new() -> Self {
30        Self::default()
31    }
32
33    pub fn bind_main(&mut self, thread_id: &str, shell: &SharedPtyShell) {
34        self.mains.retain(|_, shell| shell.strong_count() > 0);
35        self.mains.insert(thread_id.to_string(), Arc::downgrade(shell));
36    }
37
38    pub fn resolve(&self, thread_id: &str, target: &ShellTarget) -> Option<SharedPtyShell> {
39        match target {
40            ShellTarget::Main => self.mains.get(thread_id).and_then(Weak::upgrade),
41            ShellTarget::Background(id) => self.backgrounds.get_shell(thread_id, id),
42        }
43    }
44
45    pub(crate) fn register_shell(&mut self, thread_id: &str, shell: PtyShell) -> Result<String> {
46        self.backgrounds.register_shell(thread_id, shell)
47    }
48
49    pub(crate) fn get_shell(&self, thread_id: &str, bg_command_id: &str) -> Option<SharedPtyShell> {
50        self.backgrounds.get_shell(thread_id, bg_command_id)
51    }
52
53    pub(crate) fn remove_shell(&mut self, bg_command_id: &str) -> bool {
54        self.backgrounds.remove_shell(bg_command_id)
55    }
56
57    pub(crate) fn prune_finished_shells(&mut self) {
58        self.backgrounds.prune_finished_shells();
59    }
60
61    pub(crate) fn peek_tombstone(
62        &self,
63        thread_id: &str,
64        bg_command_id: &str,
65    ) -> Option<ExitedShellInfo> {
66        self.backgrounds.peek_tombstone(thread_id, bg_command_id)
67    }
68
69    pub(crate) fn get_running_info(&mut self, thread_id: &str) -> String {
70        self.backgrounds.get_running_info(thread_id)
71    }
72}
73
74static EMBEDDED_SESSION_STORE: LazyLock<StdMutex<SessionStore>> =
75    LazyLock::new(|| StdMutex::new(SessionStore::new()));
76
77/// Lock the embedded session index, recovering from poisoning like the legacy
78/// background manager did.
79pub(crate) fn lock_session_store() -> std::sync::MutexGuard<'static, SessionStore> {
80    EMBEDDED_SESSION_STORE.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
81}