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 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
77pub(crate) fn lock_session_store() -> std::sync::MutexGuard<'static, SessionStore> {
80 EMBEDDED_SESSION_STORE.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
81}