1use std::path::PathBuf;
4use std::sync::Arc;
5
6use axum::{
7 extract::{Path, State},
8 http::StatusCode,
9 Json,
10};
11
12use super::types::{
13 CreateSessionRequest, CreateSessionResponse, ErrorResponse, ExecuteCommandRequest,
14 ExecuteCommandResponse, ListSessionsResponse, SessionStatusResponse, SessionSummary,
15};
16use crate::execution::{Command, CommandExecutor};
17use crate::session::{SessionConfig, SessionId, SessionState, SessionStore};
18
19#[derive(Clone)]
21pub struct AppState {
22 pub store: Arc<SessionStore>,
23 pub executor: Arc<CommandExecutor>,
24}
25
26impl AppState {
27 pub fn new() -> Self {
28 let store = Arc::new(SessionStore::new());
29 let executor = Arc::new(CommandExecutor::new(Arc::clone(&store)));
30 Self { store, executor }
31 }
32}
33
34impl Default for AppState {
35 fn default() -> Self {
36 Self::new()
37 }
38}
39
40pub async fn health() -> &'static str {
42 "OK"
43}
44
45pub async fn api_info() -> Json<serde_json::Value> {
47 Json(serde_json::json!({
48 "name": "shell-tunnel",
49 "version": env!("CARGO_PKG_VERSION"),
50 "status": "running"
51 }))
52}
53
54pub async fn list_sessions(
56 State(state): State<AppState>,
57) -> Result<Json<ListSessionsResponse>, (StatusCode, Json<ErrorResponse>)> {
58 let ids = state.store.list_ids().map_err(|e| {
59 (
60 StatusCode::INTERNAL_SERVER_ERROR,
61 Json(ErrorResponse::internal_error(e.to_string())),
62 )
63 })?;
64
65 let mut sessions = Vec::with_capacity(ids.len());
66 for id in ids {
67 if let Ok(Some(session)) = state.store.get(&id) {
68 sessions.push(SessionSummary {
69 session_id: session.id.as_u64(),
70 state: format!("{:?}", session.state),
71 idle_seconds: session.idle_duration().as_secs_f64(),
72 });
73 }
74 }
75
76 Ok(Json(ListSessionsResponse {
77 count: sessions.len(),
78 sessions,
79 }))
80}
81
82pub async fn create_session(
84 State(state): State<AppState>,
85 Json(req): Json<CreateSessionRequest>,
86) -> Result<(StatusCode, Json<CreateSessionResponse>), (StatusCode, Json<ErrorResponse>)> {
87 let config = SessionConfig {
88 shell: req.shell,
89 working_dir: req.working_dir,
90 env: req.env,
91 };
92
93 let session_id = state.store.create(config).map_err(|e| {
94 (
95 StatusCode::INTERNAL_SERVER_ERROR,
96 Json(ErrorResponse::internal_error(e.to_string())),
97 )
98 })?;
99
100 state
102 .store
103 .update(&session_id, |s| {
104 let _ = s.state.transition_to(SessionState::Idle);
105 })
106 .ok();
107
108 Ok((
109 StatusCode::CREATED,
110 Json(CreateSessionResponse::new(session_id)),
111 ))
112}
113
114pub async fn get_session(
116 State(state): State<AppState>,
117 Path(session_id): Path<u64>,
118) -> Result<Json<SessionStatusResponse>, (StatusCode, Json<ErrorResponse>)> {
119 let id = SessionId::from_raw(session_id);
120
121 let session = state
122 .store
123 .get(&id)
124 .map_err(|e| {
125 (
126 StatusCode::INTERNAL_SERVER_ERROR,
127 Json(ErrorResponse::internal_error(e.to_string())),
128 )
129 })?
130 .ok_or_else(|| {
131 (
132 StatusCode::NOT_FOUND,
133 Json(ErrorResponse::session_not_found(&session_id.to_string())),
134 )
135 })?;
136
137 Ok(Json(SessionStatusResponse::from_session(&session)))
138}
139
140pub async fn delete_session(
142 State(state): State<AppState>,
143 Path(session_id): Path<u64>,
144) -> Result<StatusCode, (StatusCode, Json<ErrorResponse>)> {
145 let id = SessionId::from_raw(session_id);
146
147 state
149 .store
150 .update(&id, |s| {
151 let _ = s.state.transition_to(SessionState::Terminated);
152 })
153 .map_err(|_| {
154 (
155 StatusCode::NOT_FOUND,
156 Json(ErrorResponse::session_not_found(&session_id.to_string())),
157 )
158 })?;
159
160 state.store.remove(&id).map_err(|e| {
162 (
163 StatusCode::INTERNAL_SERVER_ERROR,
164 Json(ErrorResponse::internal_error(e.to_string())),
165 )
166 })?;
167
168 Ok(StatusCode::NO_CONTENT)
169}
170
171pub async fn execute_command(
173 State(state): State<AppState>,
174 Path(session_id): Path<u64>,
175 Json(req): Json<ExecuteCommandRequest>,
176) -> Result<Json<ExecuteCommandResponse>, (StatusCode, Json<ErrorResponse>)> {
177 let id = SessionId::from_raw(session_id);
178
179 let session = state
181 .store
182 .get(&id)
183 .map_err(|e| {
184 (
185 StatusCode::INTERNAL_SERVER_ERROR,
186 Json(ErrorResponse::internal_error(e.to_string())),
187 )
188 })?
189 .ok_or_else(|| {
190 (
191 StatusCode::NOT_FOUND,
192 Json(ErrorResponse::session_not_found(&session_id.to_string())),
193 )
194 })?;
195
196 if !session.state.can_execute() {
197 return Err((
198 StatusCode::CONFLICT,
199 Json(ErrorResponse::invalid_state(session.state)),
200 ));
201 }
202
203 let mut cmd = Command::new(&req.command);
205 if let Some(dir) = &req.working_dir {
206 cmd = cmd.working_dir(PathBuf::from(dir));
207 }
208 if let Some(timeout) = req.timeout() {
209 cmd = cmd.timeout(timeout);
210 }
211 for (key, value) in &req.env {
212 cmd = cmd.env(key, value);
213 }
214
215 let result = state
217 .executor
218 .execute_in_session(&id, &cmd)
219 .await
220 .map_err(|e| {
221 (
222 StatusCode::INTERNAL_SERVER_ERROR,
223 Json(ErrorResponse::internal_error(e.to_string())),
224 )
225 })?;
226
227 state
229 .store
230 .update(&id, |s| {
231 s.context.record_execution(&req.command, result.exit_code);
232 })
233 .ok();
234
235 Ok(Json(ExecuteCommandResponse::from_result(&result)))
236}
237
238pub async fn execute_oneshot(
240 State(state): State<AppState>,
241 Json(req): Json<ExecuteCommandRequest>,
242) -> Result<Json<ExecuteCommandResponse>, (StatusCode, Json<ErrorResponse>)> {
243 let mut cmd = Command::new(&req.command);
245 if let Some(dir) = &req.working_dir {
246 cmd = cmd.working_dir(PathBuf::from(dir));
247 }
248 if let Some(timeout) = req.timeout() {
249 cmd = cmd.timeout(timeout);
250 }
251 for (key, value) in &req.env {
252 cmd = cmd.env(key, value);
253 }
254
255 let result = state.executor.execute(&cmd).await.map_err(|e| {
257 (
258 StatusCode::INTERNAL_SERVER_ERROR,
259 Json(ErrorResponse::internal_error(e.to_string())),
260 )
261 })?;
262
263 Ok(Json(ExecuteCommandResponse::from_result(&result)))
264}
265
266#[cfg(test)]
267mod tests {
268 use super::*;
269
270 #[test]
271 fn test_app_state_new() {
272 let state = AppState::new();
273 assert_eq!(state.store.count(), 0);
274 }
275
276 #[tokio::test]
277 async fn test_health_endpoint() {
278 let response = health().await;
279 assert_eq!(response, "OK");
280 }
281
282 #[tokio::test]
283 async fn test_api_info_endpoint() {
284 let response = api_info().await;
285 let json = response.0;
286 assert_eq!(json["name"], "shell-tunnel");
287 assert_eq!(json["status"], "running");
288 }
289}