1use std::sync::atomic::{AtomicU64, Ordering};
2use std::sync::Arc;
3use std::time::{Duration, SystemTime, UNIX_EPOCH};
4
5use serde_json::json;
6
7use crate::app_server::outgoing::OutgoingMessageSender;
8use crate::app_server::protocol::{
9 map_run_event_to_notifications, AgentMessageDeltaParams, AppItem, AppItemKind, AppItemStatus,
10 AppServerError, AppServerErrorCode, AppTurn, ApprovalDecision, ApprovalRequestParams,
11 ApprovalResolveParams, ItemCompletedParams, ItemStartedParams, JsonRpcError, JsonRpcErrorBody,
12 RequestId, ServerNotification, ServerRequest, ThreadStatus, TurnCompletedParams,
13 TurnStartParams, TurnStartedParams, TurnStatus, UserInput,
14};
15use crate::app_server::thread_state::{ActiveTurn, ThreadStateManager};
16use crate::app_server::thread_store::{SqliteThreadStore, ThreadStoreError};
17use crate::tools::ApprovalDecision as ToolApprovalDecision;
18use crate::{
19 Agent, ApprovalBroker, ApprovalFuture, ApprovalProvider, ApprovalRequest, ModelRef, RunConfig,
20 RunHandle, Runner,
21};
22
23#[derive(Clone)]
24pub struct AppServerRunAdapter {
25 runner: Runner,
26 agent: Agent,
27 store: SqliteThreadStore,
28 state: ThreadStateManager,
29 outgoing: OutgoingMessageSender,
30 next_turn_id: Arc<AtomicU64>,
31 approval_request_timeout: Duration,
32}
33
34impl AppServerRunAdapter {
35 pub fn new(
36 runner: Runner,
37 agent: Agent,
38 store: SqliteThreadStore,
39 state: ThreadStateManager,
40 outgoing: OutgoingMessageSender,
41 ) -> Self {
42 Self {
43 runner,
44 agent,
45 store,
46 state,
47 outgoing,
48 next_turn_id: Arc::new(AtomicU64::new(1)),
49 approval_request_timeout: Duration::from_secs(30),
50 }
51 }
52
53 pub fn with_approval_request_timeout(mut self, timeout: Duration) -> Self {
54 self.approval_request_timeout = timeout;
55 self
56 }
57
58 pub fn store(&self) -> &SqliteThreadStore {
59 &self.store
60 }
61
62 pub fn state(&self) -> &ThreadStateManager {
63 &self.state
64 }
65
66 pub async fn start_turn(&self, params: TurnStartParams) -> Result<AppTurn, AppServerError> {
67 let thread = self
68 .store
69 .get_thread(¶ms.thread_id)
70 .map_err(store_error)?
71 .ok_or_else(|| AppServerError::invalid_params("Unknown thread"))?;
72 if thread.archived {
73 return Err(AppServerError::invalid_params("Thread is archived"));
74 }
75 if self.state.active_turn(¶ms.thread_id).await.is_some() {
76 return Err(AppServerError::invalid_params(
77 "Thread already has an active turn",
78 ));
79 }
80
81 let turn_id = format!("turn_{}", self.next_turn_id.fetch_add(1, Ordering::Relaxed));
82 let input = params.input.clone();
83 let input_text = input_text(&input);
84 let run_id = format!("{}_run", self.agent.name());
85 let mut config = RunConfig::default();
86 if let Some(model) = params.model {
87 config.model = Some(ModelRef::named(model));
88 }
89 let approval_broker = ApprovalBroker::default();
90 config.approval_provider = Some(Arc::new(AppServerApprovalProvider));
91 config.approval_broker = Some(approval_broker);
92 config
93 .metadata
94 .insert("thread_id".to_string(), json!(params.thread_id));
95 config
96 .metadata
97 .insert("turn_id".to_string(), json!(turn_id.clone()));
98
99 let handle = self
100 .runner
101 .start(&self.agent, input_text, config)
102 .await
103 .map_err(AppServerError::internal)?;
104 let turn = AppTurn {
105 id: turn_id.clone(),
106 thread_id: params.thread_id.clone(),
107 run_id,
108 status: TurnStatus::Running,
109 input,
110 started_at_ms: Some(timestamp_millis()),
111 completed_at_ms: None,
112 token_usage: None,
113 };
114 self.store
115 .set_active_turn(¶ms.thread_id, Some(&turn_id), ThreadStatus::Running)
116 .map_err(store_error)?;
117 self.state
118 .set_active_turn(
119 params.thread_id,
120 ActiveTurn {
121 turn: turn.clone(),
122 handle,
123 },
124 )
125 .await;
126 Ok(turn)
127 }
128
129 pub async fn notify_turn_started(&self, turn: &AppTurn) -> Result<(), AppServerError> {
130 self.broadcast_to_thread(
131 &turn.thread_id,
132 ServerNotification::TurnStarted(TurnStartedParams { turn: turn.clone() }),
133 )
134 .await
135 }
136
137 pub async fn spawn_event_forwarding(&self, thread_id: String, turn_id: String) {
138 let Some(active) = self.state.active_turn(&thread_id).await else {
139 return;
140 };
141 let adapter = self.clone();
142 tokio::spawn(async move {
143 let mut events = active.handle.events();
144 while let Some(event) = events.next().await {
145 match event {
146 Ok(event) => {
147 let notifications =
148 map_run_event_to_notifications(&thread_id, &turn_id, &event);
149 for notification in notifications {
150 if let Some(item) = item_from_notification(¬ification) {
151 let _ = adapter.store.append_item(&thread_id, &turn_id, item);
152 }
153 let _ = adapter
154 .broadcast_to_thread(&thread_id, notification.clone())
155 .await;
156 if let ServerNotification::ApprovalRequested(approval) = ¬ification {
157 adapter
158 .route_approval_request(
159 &thread_id,
160 &turn_id,
161 &active.handle,
162 approval.clone(),
163 )
164 .await;
165 }
166 if is_terminal_turn_notification(¬ification) {
167 adapter.state.clear_active_turn(&thread_id, &turn_id).await;
168 let _ = adapter.store.set_active_turn(
169 &thread_id,
170 None,
171 ThreadStatus::Idle,
172 );
173 }
174 }
175 }
176 Err(error) => {
177 let _ = adapter
178 .broadcast_to_thread(
179 &thread_id,
180 ServerNotification::ErrorWarning(
181 crate::app_server::protocol::WarningParams {
182 message: error,
183 code: Some("event_stream".to_string()),
184 },
185 ),
186 )
187 .await;
188 break;
189 }
190 }
191 }
192 });
193 }
194
195 pub async fn interrupt_turn(
196 &self,
197 thread_id: &str,
198 turn_id: &str,
199 ) -> Result<InterruptTurnOutcome, AppServerError> {
200 if let Some(pending) = self.state.pending_approval(thread_id).await {
201 if pending.turn_id == turn_id {
202 let mut completed_turn = None;
203 if let Some(active) = self.state.active_turn(thread_id).await {
204 active.handle.cancel();
205 let _ = active
206 .handle
207 .approve(
208 &pending.request_id,
209 ToolApprovalDecision::timeout("turn interrupted"),
210 )
211 .await;
212 completed_turn = Some(interrupted_turn(active.turn));
213 }
214 let _ = self
215 .outgoing
216 .resolve_server_error(JsonRpcError {
217 id: RequestId::String(pending.request_id.clone()),
218 error: JsonRpcErrorBody {
219 code: AppServerErrorCode::InternalError.code(),
220 message: "turn interrupted".to_string(),
221 data: None,
222 },
223 })
224 .await;
225 self.state
226 .clear_pending_approval(thread_id, &pending.request_id)
227 .await;
228 self.store
229 .set_active_turn(thread_id, None, ThreadStatus::Idle)
230 .map_err(store_error)?;
231 return Ok(InterruptTurnOutcome {
232 approval_resolved: Some(ApprovalResolveParams {
233 thread_id: thread_id.to_string(),
234 turn_id: turn_id.to_string(),
235 request_id: pending.request_id,
236 decision: ApprovalDecision::Deny,
237 }),
238 completed_turn,
239 });
240 }
241 }
242 if let Some(active) = self.state.active_turn(thread_id).await {
243 if active.turn.id != turn_id {
244 return Err(AppServerError::invalid_params("No matching active turn"));
245 }
246 active.handle.cancel();
247 self.store
248 .set_active_turn(thread_id, None, ThreadStatus::Idle)
249 .map_err(store_error)?;
250 return Ok(InterruptTurnOutcome {
251 approval_resolved: None,
252 completed_turn: Some(interrupted_turn(active.turn)),
253 });
254 }
255 Err(AppServerError::invalid_params("No matching active turn"))
256 }
257
258 pub async fn notify_approval_resolved(
259 &self,
260 params: ApprovalResolveParams,
261 ) -> Result<(), AppServerError> {
262 let thread_id = params.thread_id.clone();
263 self.broadcast_to_thread(&thread_id, ServerNotification::ApprovalResolved(params))
264 .await
265 }
266
267 pub async fn notify_turn_completed(&self, turn: AppTurn) -> Result<(), AppServerError> {
268 let thread_id = turn.thread_id.clone();
269 self.broadcast_to_thread(
270 &thread_id,
271 ServerNotification::TurnCompleted(TurnCompletedParams { turn }),
272 )
273 .await
274 }
275
276 pub async fn active_turn(&self, thread_id: &str) -> Option<AppTurn> {
277 self.state
278 .active_turn(thread_id)
279 .await
280 .map(|active| active.turn)
281 }
282
283 async fn broadcast_to_thread(
284 &self,
285 thread_id: &str,
286 notification: ServerNotification,
287 ) -> Result<(), AppServerError> {
288 let subscribers = self.state.subscribers(thread_id).await;
289 for connection_id in subscribers {
290 self.outgoing
291 .send_notification(connection_id, notification.clone())
292 .await?;
293 }
294 Ok(())
295 }
296
297 async fn route_approval_request(
298 &self,
299 thread_id: &str,
300 turn_id: &str,
301 handle: &RunHandle,
302 approval: ApprovalRequestParams,
303 ) {
304 let subscribers = self.state.subscribers(thread_id).await;
305 let Some(connection_id) = subscribers.first().copied() else {
306 let _ = handle
307 .approve(
308 &approval.request_id,
309 ToolApprovalDecision::timeout("approval client disconnected"),
310 )
311 .await;
312 return;
313 };
314 self.state
315 .set_pending_approval(thread_id, turn_id, approval.request_id.clone())
316 .await;
317 let response = self
318 .outgoing
319 .send_server_request_with_id_and_timeout(
320 connection_id,
321 RequestId::String(approval.request_id.clone()),
322 ServerRequest::ApprovalRequest(approval.clone()),
323 self.approval_request_timeout,
324 )
325 .await;
326 let (decision, protocol_decision) = match response {
327 Ok(value) => tool_approval_decision_from_response(value),
328 Err(error) => (
329 ToolApprovalDecision::timeout(error.message().to_string()),
330 ApprovalDecision::Deny,
331 ),
332 };
333 if self
334 .state
335 .pending_approval(thread_id)
336 .await
337 .is_none_or(|pending| pending.request_id != approval.request_id)
338 {
339 return;
340 }
341 self.state
342 .clear_pending_approval(thread_id, &approval.request_id)
343 .await;
344 let _ = handle.approve(&approval.request_id, decision).await;
345 let _ = self
346 .broadcast_to_thread(
347 thread_id,
348 ServerNotification::ApprovalResolved(
349 crate::app_server::protocol::ApprovalResolveParams {
350 thread_id: thread_id.to_string(),
351 turn_id: turn_id.to_string(),
352 request_id: approval.request_id,
353 decision: protocol_decision,
354 },
355 ),
356 )
357 .await;
358 }
359}
360
361pub struct InterruptTurnOutcome {
362 pub approval_resolved: Option<ApprovalResolveParams>,
363 pub completed_turn: Option<AppTurn>,
364}
365
366struct AppServerApprovalProvider;
367
368impl ApprovalProvider for AppServerApprovalProvider {
369 fn should_request(&self, _request: &ApprovalRequest) -> bool {
370 true
371 }
372
373 fn decide(&self, _request: &ApprovalRequest) -> ApprovalFuture<Option<ToolApprovalDecision>> {
374 Box::pin(async { Ok(None) })
375 }
376}
377
378fn item_from_notification(notification: &ServerNotification) -> Option<AppItem> {
379 match notification {
380 ServerNotification::AgentMessageDelta(AgentMessageDeltaParams {
381 item_id, delta, ..
382 }) => Some(AppItem {
383 id: item_id.clone(),
384 run_event_id: item_id.clone(),
385 kind: AppItemKind::AgentMessage,
386 status: AppItemStatus::Completed,
387 created_at_ms: timestamp_millis(),
388 completed_at_ms: Some(timestamp_millis()),
389 content: Some(json!({ "text": delta })),
390 }),
391 ServerNotification::ItemStarted(ItemStartedParams { item, .. })
392 | ServerNotification::ItemCompleted(ItemCompletedParams { item, .. }) => Some(item.clone()),
393 ServerNotification::ApprovalRequested(ApprovalRequestParams {
394 request_id,
395 tool_name,
396 preview,
397 ..
398 }) => Some(AppItem {
399 id: request_id.clone(),
400 run_event_id: request_id.clone(),
401 kind: AppItemKind::ApprovalRequest,
402 status: AppItemStatus::InProgress,
403 created_at_ms: timestamp_millis(),
404 completed_at_ms: None,
405 content: Some(json!({
406 "toolName": tool_name,
407 "preview": preview,
408 "choices": [ApprovalDecision::Allow, ApprovalDecision::Deny],
409 })),
410 }),
411 _ => None,
412 }
413}
414
415fn is_terminal_turn_notification(notification: &ServerNotification) -> bool {
416 matches!(notification, ServerNotification::TurnCompleted(_))
417}
418
419fn interrupted_turn(mut turn: AppTurn) -> AppTurn {
420 turn.status = TurnStatus::Interrupted;
421 turn.completed_at_ms = Some(timestamp_millis());
422 turn
423}
424
425fn input_text(input: &[UserInput]) -> String {
426 input
427 .iter()
428 .map(|item| item.text.as_str())
429 .collect::<Vec<_>>()
430 .join("\n")
431}
432
433fn store_error(error: ThreadStoreError) -> AppServerError {
434 AppServerError::internal(error.to_string())
435}
436
437fn tool_approval_decision_from_response(
438 value: serde_json::Value,
439) -> (ToolApprovalDecision, ApprovalDecision) {
440 match value
441 .get("decision")
442 .and_then(serde_json::Value::as_str)
443 .unwrap_or_default()
444 {
445 "allow" | "approved" => (ToolApprovalDecision::allow(), ApprovalDecision::Allow),
446 "deny" | "denied" => {
447 let reason = value
448 .get("message")
449 .and_then(serde_json::Value::as_str)
450 .unwrap_or("approval denied");
451 (ToolApprovalDecision::deny(reason), ApprovalDecision::Deny)
452 }
453 _ => (
454 ToolApprovalDecision::deny("invalid approval response"),
455 ApprovalDecision::Deny,
456 ),
457 }
458}
459
460fn timestamp_millis() -> u128 {
461 SystemTime::now()
462 .duration_since(UNIX_EPOCH)
463 .map(|duration| duration.as_millis())
464 .unwrap_or_default()
465}