1use std::collections::{HashMap, HashSet};
2use std::time::Duration;
3
4use serde_json::{json, Value};
5
6use crate::app_server::outgoing::{OutgoingEnvelope, OutgoingMessageSender};
7use crate::app_server::protocol::{
8 generate_app_server_json_schema_bundle, generate_app_server_typescript_bundle, AppClientInfo,
9 AppServerCapabilities, AppServerError, AppServerErrorCode, ApprovalDecision,
10 ApprovalResolveParams, InitializeParams, InitializeResponse, JsonRpcMessage,
11 JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, ModelListParams, ModelListResponse,
12 SchemaExportResponse, ServerNotification, ThreadArchiveParams, ThreadArchiveResponse,
13 ThreadArchivedParams, ThreadListParams, ThreadListResponse, ThreadReadParams,
14 ThreadReadResponse, ThreadResumeParams, ThreadResumeResponse, ThreadStartParams,
15 ThreadStartResponse, ThreadStartedParams, TurnInterruptParams, TurnInterruptResponse,
16 TurnStartParams, TurnStartResponse,
17};
18use crate::app_server::request_serialization::{
19 RequestSerializationQueue, RequestSerializationScope,
20};
21use crate::app_server::run_adapter::AppServerRunAdapter;
22use crate::app_server::thread_state::ThreadStateManager;
23use crate::app_server::thread_store::{SqliteThreadStore, ThreadStoreError};
24use crate::app_server::transport::ConnectionId;
25use crate::{Agent, Runner};
26
27pub struct MessageProcessor {
28 outgoing: OutgoingMessageSender,
29 connections: HashMap<ConnectionId, ConnectionSessionState>,
30 run_adapter: Option<AppServerRunAdapter>,
31 thread_state: ThreadStateManager,
32 request_queue: RequestSerializationQueue,
33}
34
35#[derive(Debug, Clone, Default)]
36pub struct ConnectionSessionState {
37 initialized: bool,
38 ready_for_notifications: bool,
39 client_info: Option<AppClientInfo>,
40 experimental_api: bool,
41 opt_out_notification_methods: HashSet<String>,
42}
43
44impl ConnectionSessionState {
45 pub fn initialized(&self) -> bool {
46 self.initialized
47 }
48
49 pub fn ready_for_notifications(&self) -> bool {
50 self.ready_for_notifications
51 }
52
53 pub fn client_info(&self) -> Option<&AppClientInfo> {
54 self.client_info.as_ref()
55 }
56
57 pub fn experimental_api(&self) -> bool {
58 self.experimental_api
59 }
60
61 pub fn opt_out_notification_methods(&self) -> &HashSet<String> {
62 &self.opt_out_notification_methods
63 }
64}
65
66impl MessageProcessor {
67 pub fn new_for_tests(
68 outgoing_capacity: usize,
69 ) -> (Self, tokio::sync::mpsc::Receiver<OutgoingEnvelope>) {
70 let (outgoing, rx) = OutgoingMessageSender::channel(outgoing_capacity);
71 (
72 Self {
73 outgoing,
74 connections: HashMap::new(),
75 run_adapter: None,
76 thread_state: ThreadStateManager::default(),
77 request_queue: RequestSerializationQueue::default(),
78 },
79 rx,
80 )
81 }
82
83 pub fn new_for_tests_with_runtime(
84 outgoing_capacity: usize,
85 runner: Runner,
86 agent: Agent,
87 store: SqliteThreadStore,
88 ) -> (Self, tokio::sync::mpsc::Receiver<OutgoingEnvelope>) {
89 Self::new_for_tests_with_runtime_and_approval_timeout(
90 outgoing_capacity,
91 runner,
92 agent,
93 store,
94 Duration::from_secs(30),
95 )
96 }
97
98 pub fn new_for_tests_with_runtime_and_approval_timeout(
99 outgoing_capacity: usize,
100 runner: Runner,
101 agent: Agent,
102 store: SqliteThreadStore,
103 approval_request_timeout: Duration,
104 ) -> (Self, tokio::sync::mpsc::Receiver<OutgoingEnvelope>) {
105 let (outgoing, rx) = OutgoingMessageSender::channel(outgoing_capacity);
106 let thread_state = ThreadStateManager::default();
107 let run_adapter =
108 AppServerRunAdapter::new(runner, agent, store, thread_state.clone(), outgoing.clone())
109 .with_approval_request_timeout(approval_request_timeout);
110 (
111 Self {
112 outgoing,
113 connections: HashMap::new(),
114 run_adapter: Some(run_adapter),
115 thread_state,
116 request_queue: RequestSerializationQueue::default(),
117 },
118 rx,
119 )
120 }
121
122 pub fn outgoing(&self) -> &OutgoingMessageSender {
123 &self.outgoing
124 }
125
126 pub fn connection_state(&self, connection_id: ConnectionId) -> Option<&ConnectionSessionState> {
127 self.connections.get(&connection_id)
128 }
129
130 pub async fn process_message(&mut self, connection_id: ConnectionId, message: JsonRpcMessage) {
131 self.outgoing.register_connection(connection_id).await;
132 match message {
133 JsonRpcMessage::Request(request) => {
134 self.process_request(connection_id, request).await;
135 }
136 JsonRpcMessage::Notification(notification) => {
137 self.process_notification(connection_id, notification).await;
138 }
139 JsonRpcMessage::Response(response) => {
140 let _ = self.outgoing.resolve_server_response(response).await;
141 }
142 JsonRpcMessage::Error(error) => {
143 let _ = self.outgoing.resolve_server_error(error).await;
144 }
145 }
146 }
147
148 async fn process_request(&mut self, connection_id: ConnectionId, request: JsonRpcRequest) {
149 if request.method == "initialize" {
150 self.process_initialize(connection_id, request).await;
151 return;
152 }
153
154 if !self
155 .connections
156 .get(&connection_id)
157 .is_some_and(ConnectionSessionState::initialized)
158 {
159 let _ = self
160 .outgoing
161 .send_error(connection_id, request.id, AppServerError::not_initialized())
162 .await;
163 return;
164 }
165
166 if let Some(scope) =
167 RequestSerializationScope::for_method(&request.method, request.params.as_ref())
168 {
169 let queue = self.request_queue.clone();
170 queue
171 .run(scope, async move {
172 self.process_initialized_request(connection_id, request)
173 .await;
174 })
175 .await;
176 return;
177 }
178
179 self.process_initialized_request(connection_id, request)
180 .await;
181 }
182
183 async fn process_initialized_request(
184 &mut self,
185 connection_id: ConnectionId,
186 request: JsonRpcRequest,
187 ) {
188 match request.method.as_str() {
189 "thread/start" => self.process_thread_start(connection_id, request).await,
190 "thread/resume" => self.process_thread_resume(connection_id, request).await,
191 "thread/read" => self.process_thread_read(connection_id, request).await,
192 "thread/list" => self.process_thread_list(connection_id, request).await,
193 "thread/archive" => self.process_thread_archive(connection_id, request).await,
194 "turn/start" => self.process_turn_start(connection_id, request).await,
195 "turn/interrupt" => self.process_turn_interrupt(connection_id, request).await,
196 "turn/steer" => self.process_unsupported(connection_id, request).await,
197 "approval/resolve" => self.process_approval_resolve(connection_id, request).await,
198 "model/list" => self.process_model_list(connection_id, request).await,
199 "schema/export" => self.process_schema_export(connection_id, request).await,
200 _ => {
201 let _ = self
202 .outgoing
203 .send_error(
204 connection_id,
205 request.id,
206 AppServerError::new(
207 AppServerErrorCode::MethodNotFound,
208 format!("Method not found: {}", request.method),
209 )
210 .with_data(serde_json::json!({ "method": request.method })),
211 )
212 .await;
213 }
214 }
215 }
216
217 async fn process_schema_export(
218 &mut self,
219 connection_id: ConnectionId,
220 request: JsonRpcRequest,
221 ) {
222 let json_schema = match generate_app_server_json_schema_bundle() {
223 Ok(bundle) => bundle,
224 Err(error) => {
225 let _ = self
226 .outgoing
227 .send_error(
228 connection_id,
229 request.id,
230 AppServerError::internal(error.to_string()),
231 )
232 .await;
233 return;
234 }
235 };
236 let typescript = match generate_app_server_typescript_bundle() {
237 Ok(bundle) => bundle,
238 Err(error) => {
239 let _ = self
240 .outgoing
241 .send_error(
242 connection_id,
243 request.id,
244 AppServerError::internal(error.to_string()),
245 )
246 .await;
247 return;
248 }
249 };
250 let result = serde_json::to_value(SchemaExportResponse {
251 json_schema,
252 typescript,
253 })
254 .expect("schema export response serializes");
255 let _ = self
256 .outgoing
257 .send_response(connection_id, request.id, result)
258 .await;
259 }
260
261 async fn process_model_list(&mut self, connection_id: ConnectionId, request: JsonRpcRequest) {
262 let _params = match parse_params::<ModelListParams>(request.params) {
263 Ok(params) => params,
264 Err(error) => {
265 let _ = self
266 .outgoing
267 .send_error(connection_id, request.id, error)
268 .await;
269 return;
270 }
271 };
272 let result = serde_json::to_value(ModelListResponse { models: Vec::new() })
273 .expect("model list response serializes");
274 let _ = self
275 .outgoing
276 .send_response(connection_id, request.id, result)
277 .await;
278 }
279
280 async fn process_approval_resolve(
281 &mut self,
282 connection_id: ConnectionId,
283 request: JsonRpcRequest,
284 ) {
285 let params = match parse_params::<ApprovalResolveParams>(request.params) {
286 Ok(params) => params,
287 Err(error) => {
288 let _ = self
289 .outgoing
290 .send_error(connection_id, request.id, error)
291 .await;
292 return;
293 }
294 };
295 let decision = match params.decision {
296 ApprovalDecision::Allow => "allow",
297 ApprovalDecision::Deny => "deny",
298 };
299 let resolved = self
300 .outgoing
301 .resolve_server_response(JsonRpcResponse {
302 id: crate::app_server::protocol::RequestId::String(params.request_id.clone()),
303 result: json!({ "decision": decision }),
304 })
305 .await;
306 if !resolved {
307 let _ = self
308 .outgoing
309 .send_error(
310 connection_id,
311 request.id,
312 AppServerError::invalid_params("Unknown approval request"),
313 )
314 .await;
315 return;
316 }
317 let _ = self
318 .outgoing
319 .send_response(connection_id, request.id, json!({}))
320 .await;
321 }
322
323 async fn process_unsupported(&mut self, connection_id: ConnectionId, request: JsonRpcRequest) {
324 let method = request.method.clone();
325 let _ = self
326 .outgoing
327 .send_error(
328 connection_id,
329 request.id,
330 AppServerError::unsupported_method(method),
331 )
332 .await;
333 }
334
335 async fn process_thread_resume(
336 &mut self,
337 connection_id: ConnectionId,
338 request: JsonRpcRequest,
339 ) {
340 let params = match parse_params::<ThreadResumeParams>(request.params) {
341 Ok(params) => params,
342 Err(error) => {
343 let _ = self
344 .outgoing
345 .send_error(connection_id, request.id, error)
346 .await;
347 return;
348 }
349 };
350 let Some(adapter) = self.run_adapter.clone() else {
351 let _ = self
352 .outgoing
353 .send_error(
354 connection_id,
355 request.id,
356 AppServerError::internal("App Server runtime is not configured"),
357 )
358 .await;
359 return;
360 };
361 let Some(thread) = (match adapter.store().get_thread(¶ms.thread_id) {
362 Ok(thread) => thread,
363 Err(error) => {
364 let _ = self
365 .outgoing
366 .send_error(connection_id, request.id, store_error(error))
367 .await;
368 return;
369 }
370 }) else {
371 let _ = self
372 .outgoing
373 .send_error(
374 connection_id,
375 request.id,
376 AppServerError::invalid_params("Unknown thread"),
377 )
378 .await;
379 return;
380 };
381 let items = match adapter.store().replay_items(¶ms.thread_id) {
382 Ok(items) => items,
383 Err(error) => {
384 let _ = self
385 .outgoing
386 .send_error(connection_id, request.id, store_error(error))
387 .await;
388 return;
389 }
390 };
391 if params.subscribe {
392 self.thread_state
393 .subscribe(params.thread_id.clone(), connection_id)
394 .await;
395 }
396 let active_turn = adapter.active_turn(¶ms.thread_id).await;
397 let result = serde_json::to_value(ThreadResumeResponse {
398 thread,
399 items,
400 active_turn,
401 })
402 .expect("thread resume response serializes");
403 let _ = self
404 .outgoing
405 .send_response(connection_id, request.id, result)
406 .await;
407 }
408
409 async fn process_initialize(&mut self, connection_id: ConnectionId, request: JsonRpcRequest) {
410 let state = self.connections.entry(connection_id).or_default();
411 if state.initialized {
412 let _ = self
413 .outgoing
414 .send_error(
415 connection_id,
416 request.id,
417 AppServerError::already_initialized(),
418 )
419 .await;
420 return;
421 }
422
423 let params = match parse_params::<InitializeParams>(request.params) {
424 Ok(params) => params,
425 Err(error) => {
426 let _ = self
427 .outgoing
428 .send_error(connection_id, request.id, error)
429 .await;
430 return;
431 }
432 };
433 state.initialized = true;
434 state.client_info = Some(params.client_info);
435 state.experimental_api = params.capabilities.experimental_api;
436 state.opt_out_notification_methods = params
437 .capabilities
438 .opt_out_notification_methods
439 .into_iter()
440 .collect();
441 self.outgoing
442 .configure_connection(connection_id, state.opt_out_notification_methods.clone())
443 .await;
444
445 let result = serde_json::to_value(InitializeResponse::new(
446 "vv-agent-rs",
447 env!("CARGO_PKG_VERSION"),
448 AppServerCapabilities::mvp(),
449 ))
450 .expect("initialize response serializes");
451 let _ = self
452 .outgoing
453 .send_response(connection_id, request.id, result)
454 .await;
455 }
456
457 async fn process_thread_list(&mut self, connection_id: ConnectionId, request: JsonRpcRequest) {
458 let params = match parse_params::<ThreadListParams>(request.params) {
459 Ok(params) => params,
460 Err(error) => {
461 let _ = self
462 .outgoing
463 .send_error(connection_id, request.id, error)
464 .await;
465 return;
466 }
467 };
468 let Some(adapter) = self.run_adapter.clone() else {
469 let _ = self
470 .outgoing
471 .send_error(
472 connection_id,
473 request.id,
474 AppServerError::internal("App Server runtime is not configured"),
475 )
476 .await;
477 return;
478 };
479 let include_archived = params.include_archived || params.archived.unwrap_or(false);
480 let mut threads = match adapter.store().list_threads(include_archived) {
481 Ok(threads) => threads,
482 Err(error) => {
483 let _ = self
484 .outgoing
485 .send_error(connection_id, request.id, store_error(error))
486 .await;
487 return;
488 }
489 };
490 if let Some(archived) = params.archived {
491 threads.retain(|thread| thread.archived == archived);
492 }
493 let offset = params.offset.unwrap_or_default();
494 let limit = params.limit.unwrap_or(threads.len());
495 let threads = threads.into_iter().skip(offset).take(limit).collect();
496 let result =
497 serde_json::to_value(ThreadListResponse { threads }).expect("thread list serializes");
498 let _ = self
499 .outgoing
500 .send_response(connection_id, request.id, result)
501 .await;
502 }
503
504 async fn process_thread_archive(
505 &mut self,
506 connection_id: ConnectionId,
507 request: JsonRpcRequest,
508 ) {
509 let params = match parse_params::<ThreadArchiveParams>(request.params) {
510 Ok(params) => params,
511 Err(error) => {
512 let _ = self
513 .outgoing
514 .send_error(connection_id, request.id, error)
515 .await;
516 return;
517 }
518 };
519 let Some(adapter) = self.run_adapter.clone() else {
520 let _ = self
521 .outgoing
522 .send_error(
523 connection_id,
524 request.id,
525 AppServerError::internal("App Server runtime is not configured"),
526 )
527 .await;
528 return;
529 };
530 match adapter.store().get_thread(¶ms.thread_id) {
531 Ok(Some(_)) => {}
532 Ok(None) => {
533 let _ = self
534 .outgoing
535 .send_error(
536 connection_id,
537 request.id,
538 AppServerError::invalid_params("Unknown thread"),
539 )
540 .await;
541 return;
542 }
543 Err(error) => {
544 let _ = self
545 .outgoing
546 .send_error(connection_id, request.id, store_error(error))
547 .await;
548 return;
549 }
550 }
551 if let Err(error) = adapter.store().archive_thread(¶ms.thread_id) {
552 let _ = self
553 .outgoing
554 .send_error(connection_id, request.id, store_error(error))
555 .await;
556 return;
557 }
558 self.thread_state
559 .subscribe(params.thread_id.clone(), connection_id)
560 .await;
561 let result = serde_json::to_value(ThreadArchiveResponse {})
562 .expect("thread archive response serializes");
563 let _ = self
564 .outgoing
565 .send_response(connection_id, request.id, result)
566 .await;
567 let subscribers = self.thread_state.subscribers(¶ms.thread_id).await;
568 for subscriber in subscribers {
569 let _ = self
570 .outgoing
571 .send_notification(
572 subscriber,
573 ServerNotification::ThreadArchived(ThreadArchivedParams {
574 thread_id: params.thread_id.clone(),
575 }),
576 )
577 .await;
578 }
579 }
580
581 async fn process_notification(
582 &mut self,
583 connection_id: ConnectionId,
584 notification: JsonRpcNotification,
585 ) {
586 if notification.method != "initialized" {
587 return;
588 }
589 let Some(state) = self.connections.get_mut(&connection_id) else {
590 return;
591 };
592 if !state.initialized {
593 return;
594 }
595 state.ready_for_notifications = true;
596 self.outgoing
597 .mark_ready_for_notifications(connection_id)
598 .await;
599 }
600
601 async fn process_thread_start(&mut self, connection_id: ConnectionId, request: JsonRpcRequest) {
602 let params = match parse_params::<ThreadStartParams>(request.params) {
603 Ok(params) => params,
604 Err(error) => {
605 let _ = self
606 .outgoing
607 .send_error(connection_id, request.id, error)
608 .await;
609 return;
610 }
611 };
612 let Some(adapter) = self.run_adapter.clone() else {
613 let _ = self
614 .outgoing
615 .send_error(
616 connection_id,
617 request.id,
618 AppServerError::internal("App Server runtime is not configured"),
619 )
620 .await;
621 return;
622 };
623 let thread = match adapter.store().create_thread(params) {
624 Ok(thread) => thread,
625 Err(error) => {
626 let _ = self
627 .outgoing
628 .send_error(connection_id, request.id, store_error(error))
629 .await;
630 return;
631 }
632 };
633 self.thread_state
634 .subscribe(thread.id.clone(), connection_id)
635 .await;
636 let result = serde_json::to_value(ThreadStartResponse {
637 thread: thread.clone(),
638 })
639 .expect("thread start response serializes");
640 let _ = self
641 .outgoing
642 .send_response(connection_id, request.id, result)
643 .await;
644 let _ = self
645 .outgoing
646 .send_notification(
647 connection_id,
648 ServerNotification::ThreadStarted(ThreadStartedParams { thread }),
649 )
650 .await;
651 }
652
653 async fn process_thread_read(&mut self, connection_id: ConnectionId, request: JsonRpcRequest) {
654 let params = match parse_params::<ThreadReadParams>(request.params) {
655 Ok(params) => params,
656 Err(error) => {
657 let _ = self
658 .outgoing
659 .send_error(connection_id, request.id, error)
660 .await;
661 return;
662 }
663 };
664 let Some(adapter) = self.run_adapter.clone() else {
665 let _ = self
666 .outgoing
667 .send_error(
668 connection_id,
669 request.id,
670 AppServerError::internal("App Server runtime is not configured"),
671 )
672 .await;
673 return;
674 };
675 let Some(thread) = (match adapter.store().get_thread(¶ms.thread_id) {
676 Ok(thread) => thread,
677 Err(error) => {
678 let _ = self
679 .outgoing
680 .send_error(connection_id, request.id, store_error(error))
681 .await;
682 return;
683 }
684 }) else {
685 let _ = self
686 .outgoing
687 .send_error(
688 connection_id,
689 request.id,
690 AppServerError::invalid_params("Unknown thread"),
691 )
692 .await;
693 return;
694 };
695 let mut items = match adapter.store().replay_items(¶ms.thread_id) {
696 Ok(items) => items,
697 Err(error) => {
698 let _ = self
699 .outgoing
700 .send_error(connection_id, request.id, store_error(error))
701 .await;
702 return;
703 }
704 };
705 if let Some(after_item_id) = params.after_item_id {
706 if let Some(index) = items.iter().position(|item| item.id == after_item_id) {
707 items = items.into_iter().skip(index + 1).collect();
708 }
709 }
710 let active_turn = adapter.active_turn(¶ms.thread_id).await;
711 let result = serde_json::to_value(ThreadReadResponse {
712 thread,
713 items,
714 active_turn,
715 })
716 .expect("thread read response serializes");
717 let _ = self
718 .outgoing
719 .send_response(connection_id, request.id, result)
720 .await;
721 }
722
723 async fn process_turn_start(&mut self, connection_id: ConnectionId, request: JsonRpcRequest) {
724 let params = match parse_params::<TurnStartParams>(request.params) {
725 Ok(params) => params,
726 Err(error) => {
727 let _ = self
728 .outgoing
729 .send_error(connection_id, request.id, error)
730 .await;
731 return;
732 }
733 };
734 let Some(adapter) = self.run_adapter.clone() else {
735 let _ = self
736 .outgoing
737 .send_error(
738 connection_id,
739 request.id,
740 AppServerError::internal("App Server runtime is not configured"),
741 )
742 .await;
743 return;
744 };
745 self.thread_state
746 .subscribe(params.thread_id.clone(), connection_id)
747 .await;
748 let turn = match adapter.start_turn(params).await {
749 Ok(turn) => turn,
750 Err(error) => {
751 let _ = self
752 .outgoing
753 .send_error(connection_id, request.id, error)
754 .await;
755 return;
756 }
757 };
758 let result = serde_json::to_value(TurnStartResponse { turn: turn.clone() })
759 .expect("turn start response serializes");
760 let _ = self
761 .outgoing
762 .send_response(connection_id, request.id, result)
763 .await;
764 let _ = adapter.notify_turn_started(&turn).await;
765 adapter
766 .spawn_event_forwarding(turn.thread_id.clone(), turn.id.clone())
767 .await;
768 }
769
770 async fn process_turn_interrupt(
771 &mut self,
772 connection_id: ConnectionId,
773 request: JsonRpcRequest,
774 ) {
775 let params = match parse_params::<TurnInterruptParams>(request.params) {
776 Ok(params) => params,
777 Err(error) => {
778 let _ = self
779 .outgoing
780 .send_error(connection_id, request.id, error)
781 .await;
782 return;
783 }
784 };
785 let Some(adapter) = self.run_adapter.clone() else {
786 let _ = self
787 .outgoing
788 .send_error(
789 connection_id,
790 request.id,
791 AppServerError::internal("App Server runtime is not configured"),
792 )
793 .await;
794 return;
795 };
796 match adapter
797 .interrupt_turn(¶ms.thread_id, ¶ms.turn_id)
798 .await
799 {
800 Ok(outcome) => {
801 let result = serde_json::to_value(TurnInterruptResponse {})
802 .expect("turn interrupt response serializes");
803 let _ = self
804 .outgoing
805 .send_response(connection_id, request.id, result)
806 .await;
807 if let Some(resolved_approval) = outcome.approval_resolved {
808 let _ = adapter.notify_approval_resolved(resolved_approval).await;
809 }
810 if let Some(completed_turn) = outcome.completed_turn {
811 let _ = adapter.notify_turn_completed(completed_turn).await;
812 }
813 }
814 Err(error) => {
815 let _ = self
816 .outgoing
817 .send_error(connection_id, request.id, error)
818 .await;
819 }
820 }
821 }
822}
823
824fn parse_params<T: serde::de::DeserializeOwned>(
825 params: Option<Value>,
826) -> Result<T, AppServerError> {
827 let params = params.ok_or_else(|| AppServerError::invalid_params("Missing params"))?;
828 serde_json::from_value(params)
829 .map_err(|error| AppServerError::invalid_params(error.to_string()))
830}
831
832fn store_error(error: ThreadStoreError) -> AppServerError {
833 AppServerError::internal(error.to_string())
834}