Skip to main content

claude_code_cli_acp/session/
manager.rs

1use std::{
2    future::Future,
3    path::{Path, PathBuf},
4    sync::Mutex,
5    time::Duration,
6};
7
8use agent_client_protocol::schema::{
9    McpServer, PromptRequest, SessionConfigOption, SessionConfigValueId, SessionId,
10    SessionModeState, SessionModelState, SessionUpdate,
11};
12
13use crate::{
14    compat::claude_probe::ClaudeCli,
15    config::{
16        session::SessionConfigState,
17        settings::{SettingsPaths, load_merged_settings},
18    },
19    pty::session::{ClaudePtyConfig, ClaudePtySession},
20    terminal::recognizers::{self, PermissionDecision, PermissionDialog},
21    transcript::{
22        events::{TranscriptEvent, TranscriptEventKind},
23        tailer::{TranscriptLocator, TranscriptTailer},
24    },
25};
26
27#[derive(Clone)]
28pub struct SessionManager {
29    claude: ClaudeCli,
30}
31
32impl SessionManager {
33    pub fn new() -> Self {
34        Self {
35            claude: ClaudeCli::from_env(),
36        }
37    }
38
39    pub fn create_session(
40        &self,
41        session_id: SessionId,
42        cwd: PathBuf,
43        mcp_servers: Vec<McpServer>,
44    ) -> anyhow::Result<std::sync::Arc<ManagedSession>> {
45        Ok(std::sync::Arc::new(ManagedSession::new(
46            self.claude.clone(),
47            session_id,
48            cwd,
49            mcp_servers,
50            None,
51        )))
52    }
53
54    pub fn load_session(
55        &self,
56        session_id: SessionId,
57        cwd: PathBuf,
58        mcp_servers: Vec<McpServer>,
59    ) -> anyhow::Result<std::sync::Arc<ManagedSession>> {
60        self.create_session(session_id, cwd, mcp_servers)
61    }
62
63    pub fn create_print_session(
64        &self,
65        session_id: String,
66        cwd: PathBuf,
67        model: Option<String>,
68    ) -> anyhow::Result<ManagedSession> {
69        Ok(ManagedSession::new(
70            self.claude.clone(),
71            SessionId::new(session_id),
72            cwd,
73            Vec::new(),
74            model,
75        ))
76    }
77}
78
79impl Default for SessionManager {
80    fn default() -> Self {
81        Self::new()
82    }
83}
84
85#[derive(Debug, Clone)]
86pub struct TurnOptions {
87    pub timeout: Duration,
88    pub model: Option<String>,
89    pub permission_mode: Option<String>,
90    pub resume: Option<String>,
91    pub continue_last: bool,
92    pub initial_prompt_argument: bool,
93    pub attach_on_timeout: bool,
94    pub attach_on_permission: bool,
95}
96
97impl TurnOptions {
98    pub fn from_prompt_request(_request: &PromptRequest) -> Self {
99        Self {
100            timeout: Duration::from_secs(120),
101            model: None,
102            permission_mode: None,
103            resume: None,
104            continue_last: false,
105            initial_prompt_argument: false,
106            attach_on_timeout: false,
107            attach_on_permission: false,
108        }
109    }
110}
111
112pub struct ManagedSession {
113    claude: ClaudeCli,
114    session_id: SessionId,
115    cwd: PathBuf,
116    mcp_servers: Vec<McpServer>,
117    model: Mutex<Option<String>>,
118    permission_mode: Mutex<Option<String>>,
119    config: Mutex<SessionConfigState>,
120    pty: Mutex<Option<ClaudePtySession>>,
121    prompt_lock: tokio::sync::Mutex<()>,
122}
123
124impl ManagedSession {
125    fn new(
126        claude: ClaudeCli,
127        session_id: SessionId,
128        cwd: PathBuf,
129        mcp_servers: Vec<McpServer>,
130        model: Option<String>,
131    ) -> Self {
132        let settings = SettingsPaths::for_cwd(&cwd)
133            .map(|paths| load_merged_settings(&paths).settings)
134            .unwrap_or_default();
135        let mut config = SessionConfigState::from_settings(&settings);
136        if let Some(model) = model.as_deref()
137            && let Ok(resolved) = config.set_model(model)
138        {
139            drop(resolved);
140        }
141        Self {
142            claude,
143            session_id,
144            cwd,
145            mcp_servers,
146            model: Mutex::new(model),
147            permission_mode: Mutex::new(None),
148            config: Mutex::new(config),
149            pty: Mutex::new(None),
150            prompt_lock: tokio::sync::Mutex::new(()),
151        }
152    }
153
154    pub fn session_id(&self) -> &SessionId {
155        &self.session_id
156    }
157
158    pub fn cwd(&self) -> &Path {
159        &self.cwd
160    }
161
162    pub fn set_model(&self, model: Option<String>) -> anyhow::Result<()> {
163        if let Some(model) = model.as_deref() {
164            let resolved = self.config.lock().unwrap().set_model(model)?;
165            *self.model.lock().unwrap() = Some(resolved);
166        } else {
167            *self.model.lock().unwrap() = None;
168        }
169        Ok(())
170    }
171
172    pub fn set_permission_mode(&self, permission_mode: Option<String>) -> anyhow::Result<()> {
173        if let Some(permission_mode) = permission_mode.as_deref() {
174            self.config.lock().unwrap().set_mode(permission_mode)?;
175            *self.permission_mode.lock().unwrap() =
176                Some(self.config.lock().unwrap().mode().to_string());
177        } else {
178            *self.permission_mode.lock().unwrap() = None;
179        }
180        Ok(())
181    }
182
183    pub fn modes(&self) -> SessionModeState {
184        self.config.lock().unwrap().modes()
185    }
186
187    pub fn models(&self) -> SessionModelState {
188        self.config.lock().unwrap().models()
189    }
190
191    pub fn config_options(&self) -> Vec<SessionConfigOption> {
192        self.config.lock().unwrap().config_options()
193    }
194
195    pub fn set_config_option(
196        &self,
197        config_id: &str,
198        value: &SessionConfigValueId,
199    ) -> anyhow::Result<Option<SessionUpdate>> {
200        let update = self.config.lock().unwrap().set_option(config_id, value)?;
201        match config_id {
202            "mode" => {
203                *self.permission_mode.lock().unwrap() =
204                    Some(self.config.lock().unwrap().mode().to_string());
205            }
206            "model" => {
207                *self.model.lock().unwrap() = Some(self.config.lock().unwrap().model().to_string());
208            }
209            _ => {}
210        }
211        Ok(update)
212    }
213
214    pub async fn prompt(&self, prompt: String, options: TurnOptions) -> anyhow::Result<TurnOutput> {
215        self.prompt_with_permission_handler(prompt, options, |request| async move {
216            anyhow::bail!(
217                "Claude requested permission before transcript completion for session {}: {}",
218                request.session_id,
219                request.dialog.title
220            )
221        })
222        .await
223    }
224
225    pub async fn prompt_with_permission_handler<F, Fut>(
226        &self,
227        prompt: String,
228        options: TurnOptions,
229        mut permission_handler: F,
230    ) -> anyhow::Result<TurnOutput>
231    where
232        F: FnMut(PendingPermission) -> Fut + Send,
233        Fut: Future<Output = anyhow::Result<PermissionDecision>> + Send,
234    {
235        let _prompt_guard = self.prompt_lock.lock().await;
236        let (mut pty, reused_pty) = self.ensure_pty(&options, &prompt)?;
237        let locator = TranscriptLocator::default_home()?;
238        let mut tailer =
239            TranscriptTailer::from_locator_at_end(self.session_id.0.to_string(), &locator)?;
240        if !options.initial_prompt_argument || reused_pty {
241            wait_for_idle_prompt(&mut pty, options.timeout)?;
242            pty.submit_prompt(&prompt)?;
243        }
244
245        let deadline = tokio::time::Instant::now() + options.timeout;
246        let mut events = Vec::new();
247        let mut active_permission_fingerprint: Option<String> = None;
248        loop {
249            if tailer.is_none() {
250                tailer = TranscriptTailer::from_locator(self.session_id.0.to_string(), &locator)?;
251            }
252            if let Some(tailer) = tailer.as_mut() {
253                events.extend(tailer.poll()?);
254            }
255            if events.iter().any(is_assistant_terminal_event) && pty.is_idle() {
256                break;
257            }
258            if let Some(dialog) = pty.permission_dialog()? {
259                let fingerprint = permission_fingerprint(&dialog);
260                if active_permission_fingerprint.as_deref() == Some(&fingerprint) {
261                    tokio::time::sleep(Duration::from_millis(150)).await;
262                    continue;
263                }
264                if options.attach_on_permission {
265                    pty.detach_for_user()?;
266                    anyhow::bail!(
267                        "attached user to Claude session {} for permission request",
268                        self.session_id.0
269                    );
270                }
271                let decision = permission_handler(PendingPermission {
272                    session_id: self.session_id.clone(),
273                    dialog,
274                })
275                .await?;
276                if !pty.select_permission(decision)? {
277                    anyhow::bail!(
278                        "unable to select Claude permission option {:?} for session {}",
279                        decision,
280                        self.session_id.0
281                    );
282                }
283                active_permission_fingerprint = Some(fingerprint);
284            } else {
285                active_permission_fingerprint = None;
286            }
287            if tokio::time::Instant::now() >= deadline {
288                if options.attach_on_timeout {
289                    pty.detach_for_user()?;
290                }
291                let screen_status = pty
292                    .screen_snapshot()
293                    .map(|text| recognizers::recognize_screen(&text))
294                    .unwrap_or(recognizers::ScreenStatus::Unknown);
295                anyhow::bail!(
296                    "timed out waiting for Claude transcript completion for session {} (screen status: {:?})",
297                    self.session_id.0,
298                    screen_status
299                );
300            }
301            tokio::time::sleep(Duration::from_millis(150)).await;
302        }
303
304        let screen_text = pty.screen_snapshot().ok();
305        *self.pty.lock().unwrap() = Some(pty);
306        Ok(TurnOutput {
307            events,
308            screen_text,
309        })
310    }
311
312    pub async fn cancel(&self) -> anyhow::Result<()> {
313        if let Some(pty) = self.pty.lock().unwrap().as_mut() {
314            pty.send_interrupt()?;
315        }
316        Ok(())
317    }
318
319    pub async fn shutdown(&self) -> anyhow::Result<()> {
320        if let Some(mut pty) = self.pty.lock().unwrap().take() {
321            pty.send_exit()?;
322            pty.terminate()?;
323        }
324        Ok(())
325    }
326
327    fn ensure_pty(
328        &self,
329        options: &TurnOptions,
330        prompt: &str,
331    ) -> anyhow::Result<(ClaudePtySession, bool)> {
332        if let Some(pty) = self.pty.lock().unwrap().take() {
333            return Ok((pty, true));
334        }
335        let mut model = self
336            .model
337            .lock()
338            .unwrap()
339            .clone()
340            .or_else(|| Some(self.config.lock().unwrap().model().to_string()));
341        if model.as_deref() == Some("default") {
342            model = None;
343        }
344        if options.model.is_some() {
345            model = options.model.clone();
346        }
347        let permission_mode = options
348            .permission_mode
349            .clone()
350            .or_else(|| self.permission_mode.lock().unwrap().clone());
351        let config = ClaudePtyConfig {
352            executable: self.claude.executable().to_path_buf(),
353            cwd: self.cwd.clone(),
354            session_id: self.session_id.0.to_string(),
355            model,
356            permission_mode,
357            setting_sources: std::env::var("CLAUDE_CODE_ACP_SETTING_SOURCES")
358                .ok()
359                .filter(|sources| !sources.trim().is_empty()),
360            resume: options.resume.clone(),
361            continue_last: options.continue_last,
362            mcp_servers: self.mcp_servers.clone(),
363            extra_args: if options.initial_prompt_argument {
364                vec![prompt.into()]
365            } else {
366                Vec::new()
367            },
368            rows: 24,
369            cols: 80,
370        };
371        Ok((ClaudePtySession::spawn(config)?, false))
372    }
373}
374
375#[derive(Clone, Debug)]
376pub struct PendingPermission {
377    pub session_id: SessionId,
378    pub dialog: PermissionDialog,
379}
380
381fn is_assistant_terminal_event(event: &TranscriptEvent) -> bool {
382    match event.kind {
383        TranscriptEventKind::AssistantMessage => event
384            .text
385            .as_deref()
386            .is_some_and(|text| !text.trim().is_empty()),
387        TranscriptEventKind::ToolResult => {
388            event.session_id.is_some()
389                && event
390                    .text
391                    .as_deref()
392                    .is_some_and(|text| !text.trim().is_empty())
393        }
394        _ => false,
395    }
396}
397
398fn permission_fingerprint(dialog: &PermissionDialog) -> String {
399    format!(
400        "{}::{:?}",
401        dialog.title,
402        dialog
403            .options
404            .iter()
405            .map(|option| (&option.accelerator, &option.label, option.decision))
406            .collect::<Vec<_>>()
407    )
408}
409
410fn wait_for_idle_prompt(pty: &mut ClaudePtySession, timeout: Duration) -> anyhow::Result<()> {
411    let startup_timeout = timeout.min(Duration::from_secs(20));
412    let deadline = std::time::Instant::now() + startup_timeout;
413    let mut confirmed_workspace_trust = false;
414    loop {
415        let screen_status = pty
416            .screen_snapshot()
417            .map(|text| recognizers::recognize_screen(&text))
418            .unwrap_or(recognizers::ScreenStatus::Unknown);
419        match screen_status {
420            recognizers::ScreenStatus::Idle => return Ok(()),
421            recognizers::ScreenStatus::WorkspaceTrust if !confirmed_workspace_trust => {
422                pty.write_bytes(b"\r")?;
423                confirmed_workspace_trust = true;
424            }
425            _ => {}
426        }
427        if std::time::Instant::now() >= deadline {
428            anyhow::bail!("timed out waiting for Claude interactive prompt");
429        }
430        std::thread::sleep(Duration::from_millis(100));
431    }
432}
433
434pub struct TurnOutput {
435    pub events: Vec<TranscriptEvent>,
436    pub screen_text: Option<String>,
437}
438
439impl TurnOutput {
440    pub fn final_text(&self) -> String {
441        self.events
442            .iter()
443            .filter(|event| matches!(event.kind, TranscriptEventKind::AssistantMessage))
444            .filter_map(|event| event.text.as_deref())
445            .collect::<Vec<_>>()
446            .join("")
447            .trim()
448            .to_string()
449            .or_else_screen(self.screen_text.as_deref())
450    }
451
452    pub fn model(&self) -> Option<String> {
453        self.events.iter().find_map(|event| event.model.clone())
454    }
455}
456
457trait ScreenFallback {
458    fn or_else_screen(self, screen: Option<&str>) -> String;
459}
460
461impl ScreenFallback for String {
462    fn or_else_screen(self, screen: Option<&str>) -> String {
463        if self.is_empty() {
464            screen.unwrap_or_default().to_string()
465        } else {
466            self
467        }
468    }
469}