winx_code_agent/runtime/
session_store.rs1use 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#[derive(Clone, Debug, Eq, PartialEq)]
12pub enum ShellTarget {
13 Main,
14 Background(String),
15}
16
17#[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 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 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
98pub(crate) fn lock_session_store() -> std::sync::MutexGuard<'static, SessionStore> {
101 EMBEDDED_SESSION_STORE.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
102}