Skip to main content

ai_session/core/
lifecycle.rs

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