Skip to main content

ai_session/core/
lifecycle.rs

1//! Session lifecycle management
2
3use super::attention::AttentionState;
4use super::headless::HeadlessHandle;
5use super::pty::PtyHandle;
6use super::terminal::TerminalHandle;
7use super::{AISession, SessionConfig, SessionStatus};
8use anyhow::Result;
9use portable_pty::CommandBuilder;
10use std::io::ErrorKind;
11
12/// Start a session
13pub async fn start_session(session: &AISession) -> Result<()> {
14    // Update status
15    {
16        let mut status = session.status.write().await;
17        if *status != SessionStatus::Initializing {
18            return Err(anyhow::anyhow!("Session already started"));
19        }
20        *status = SessionStatus::Running;
21    }
22
23    let shell_env = std::env::var("SHELL").ok();
24    let shell = session
25        .config
26        .shell
27        .as_deref()
28        .or(shell_env.as_deref())
29        .unwrap_or("/bin/bash");
30
31    let terminal = if session.config.force_headless {
32        TerminalHandle::Headless(
33            HeadlessHandle::spawn_shell(
34                shell,
35                &session.config.working_directory,
36                session.config.environment.iter(),
37            )
38            .await?,
39        )
40    } else {
41        match spawn_pty(&session.config, shell).await {
42            Ok(pty) => TerminalHandle::Pty(pty),
43            Err(err) => {
44                if session.config.allow_headless_fallback && is_permission_denied(&err) {
45                    tracing::warn!(
46                        "PTY unavailable ({}). Falling back to headless shell for session {}",
47                        err,
48                        session.id
49                    );
50                    TerminalHandle::Headless(
51                        HeadlessHandle::spawn_shell(
52                            shell,
53                            &session.config.working_directory,
54                            session.config.environment.iter(),
55                        )
56                        .await?,
57                    )
58                } else {
59                    return Err(err);
60                }
61            }
62        }
63    };
64
65    // Store terminal handle
66    {
67        let mut terminal_lock = session.terminal.write().await;
68        *terminal_lock = Some(terminal);
69    }
70
71    // Update last activity
72    *session.last_activity.write().await = chrono::Utc::now();
73
74    // Surface to attention rail: session is now actively running.
75    session.set_attention(AttentionState::Running);
76
77    Ok(())
78}
79
80/// Stop a session
81pub async fn stop_session(session: &AISession) -> Result<()> {
82    // Update status
83    {
84        let mut status = session.status.write().await;
85        if *status != SessionStatus::Running && *status != SessionStatus::Paused {
86            return Ok(()); // Already stopped
87        }
88        *status = SessionStatus::Terminating;
89    }
90
91    // Clear terminal handle (this will close the underlying IO)
92    {
93        let mut terminal_lock = session.terminal.write().await;
94        if let Some(terminal) = terminal_lock.take() {
95            terminal.shutdown().await?;
96        }
97    }
98
99    // Clear process handle
100    {
101        let mut process_lock = session.process.write().await;
102        if let Some(mut process) = process_lock.take() {
103            let _ = process.kill().await;
104        }
105    }
106
107    // Update status
108    {
109        let mut status = session.status.write().await;
110        *status = SessionStatus::Terminated;
111    }
112
113    // Preserve a terminal Error/Done verdict if the session already produced
114    // one; otherwise mark a clean shutdown as Done.
115    if matches!(
116        session.attention(),
117        AttentionState::Idle | AttentionState::Running | AttentionState::Waiting
118    ) {
119        session.set_attention(AttentionState::Done);
120    }
121
122    Ok(())
123}
124
125async fn spawn_pty(config: &SessionConfig, shell: &str) -> Result<PtyHandle> {
126    let pty = PtyHandle::new(config.pty_size.0, config.pty_size.1)?;
127    let mut cmd = CommandBuilder::new(shell);
128    cmd.cwd(&config.working_directory);
129
130    for (key, value) in &config.environment {
131        cmd.env(key, value);
132    }
133
134    pty.spawn_command(cmd).await?;
135    Ok(pty)
136}
137
138fn is_permission_denied(err: &anyhow::Error) -> bool {
139    err.chain().any(|cause| {
140        if let Some(io_err) = cause.downcast_ref::<std::io::Error>() {
141            io_err.kind() == ErrorKind::PermissionDenied
142        } else {
143            let msg = cause.to_string();
144            msg.contains("PermissionDenied") || msg.contains("Operation not permitted")
145        }
146    })
147}
148
149/// Pause a session
150pub async fn pause_session(session: &AISession) -> Result<()> {
151    let mut status = session.status.write().await;
152    if *status != SessionStatus::Running {
153        return Err(anyhow::anyhow!("Session not running"));
154    }
155    *status = SessionStatus::Paused;
156    Ok(())
157}
158
159/// Resume a session
160pub async fn resume_session(session: &AISession) -> Result<()> {
161    let mut status = session.status.write().await;
162    if *status != SessionStatus::Paused {
163        return Err(anyhow::anyhow!("Session not paused"));
164    }
165    *status = SessionStatus::Running;
166    *session.last_activity.write().await = chrono::Utc::now();
167    Ok(())
168}