1pub mod actions;
3
4pub use actions::SqsClientProtocol;
5
6use anyhow::{Context, Result};
7use aws_config::BehaviorVersion;
8use aws_sdk_sqs::types::{MessageAttributeValue, QueueAttributeName};
9use std::net::SocketAddr;
10use std::sync::Arc;
11use tokio::sync::mpsc;
12use tracing::{debug, error, info, warn};
13
14use crate::client::llm_budget::call_llm_for_client;
15use crate::client::sqs::actions::{
16 SQS_CLIENT_CONNECTED_EVENT, SQS_MESSAGE_RECEIVED_EVENT, SQS_MESSAGE_SENT_EVENT,
17};
18use crate::llm::actions::client_trait::{Client, ClientActionResult};
19use crate::llm::actions::protocol_trait::Protocol;
20use crate::llm::ollama_client::OllamaClient;
21use crate::llm::ClientLlmResult;
22use crate::protocol::Event;
23use crate::state::app_state::AppState;
24use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
25use crate::state::{AccessLogOwner, ClientId, ClientStatus};
26use crate::utils::truncate::truncate_for_log;
27
28fn endpoint_url_from_remote_addr(remote_addr: &str) -> Option<String> {
39 let trimmed = remote_addr.trim();
40 if trimmed.is_empty() {
41 return None;
42 }
43 if trimmed.starts_with("http://") || trimmed.starts_with("https://") {
44 Some(trimmed.to_string())
45 } else {
46 Some(format!("http://{trimmed}"))
47 }
48}
49
50pub struct SqsClient;
52
53const MAX_FOLLOWUP_DEPTH: u8 = 4;
63
64impl SqsClient {
65 pub async fn connect_with_llm_actions(
67 remote_addr: String,
68 llm_client: OllamaClient,
69 app_state: Arc<AppState>,
70 status_tx: mpsc::UnboundedSender<String>,
71 client_id: ClientId,
72 startup_params: Option<crate::protocol::StartupParams>,
73 ) -> Result<SocketAddr> {
74 let queue_url = startup_params
76 .as_ref()
77 .map(|p| p.get_string("queue_url"))
78 .transpose()?
79 .context("Missing required 'queue_url' parameter")?;
80
81 let region = startup_params
82 .as_ref()
83 .map(|p| p.get_optional_string("region"))
84 .transpose()?
85 .flatten();
86 if let Some(region) = ®ion {
88 crate::client::aws_support::validate_region("SQS", region)?;
89 }
90
91 let endpoint_url = startup_params
92 .as_ref()
93 .map(|p| p.get_optional_string("endpoint_url"))
94 .transpose()?
95 .flatten();
96
97 let endpoint_url = endpoint_url.or_else(|| endpoint_url_from_remote_addr(&remote_addr));
100
101 let access_key_id = startup_params
102 .as_ref()
103 .map(|p| p.get_optional_string("access_key_id"))
104 .transpose()?
105 .flatten();
106
107 let secret_access_key = startup_params
108 .as_ref()
109 .map(|p| p.get_optional_string("secret_access_key"))
110 .transpose()?
111 .flatten();
112
113 crate::client::aws_support::refuse_ambient_credentials(
116 "SQS",
117 endpoint_url.as_deref(),
118 access_key_id.is_some() && secret_access_key.is_some(),
119 )?;
120
121 info!(
122 "SQS client {} connecting to queue: {}",
123 client_id, queue_url
124 );
125
126 let mut config_loader = aws_config::defaults(BehaviorVersion::latest());
128
129 if let Some(reg) = ®ion {
130 config_loader = config_loader.region(aws_config::Region::new(reg.clone()));
131 }
132
133 let config = config_loader.load().await;
139
140 let mut sqs_builder = aws_sdk_sqs::config::Builder::from(&config);
142
143 if let Some(endpoint) = &endpoint_url {
144 sqs_builder = sqs_builder.endpoint_url(endpoint);
145 }
146
147 if let (Some(key_id), Some(secret)) = (&access_key_id, &secret_access_key) {
152 sqs_builder = sqs_builder.credentials_provider(aws_sdk_sqs::config::Credentials::new(
153 key_id.clone(),
154 secret.clone(),
155 None, None, "netget_sqs_client",
158 ));
159 }
160
161 let sqs_config = sqs_builder.build();
162 let sqs = aws_sdk_sqs::Client::from_conf(sqs_config);
163
164 app_state
166 .update_client_status(client_id, ClientStatus::Connected)
167 .await;
168 let _ = status_tx.send(format!(
169 "[CLIENT] SQS client {} connected to {}",
170 client_id, queue_url
171 ));
172 let _ = status_tx.send("__UPDATE_UI__".to_string());
173
174 info!("SQS client {} connected", client_id);
175
176 let protocol = Arc::new(SqsClientProtocol::new());
177
178 let command_rx =
185 crate::client::command_support::register_command_channel(&app_state, client_id).await;
186 let cmd_task = tokio::spawn(Self::command_loop(
187 command_rx,
188 sqs.clone(),
189 queue_url.clone(),
190 protocol.clone(),
191 client_id,
192 app_state.clone(),
193 llm_client.clone(),
194 status_tx.clone(),
195 ));
196 app_state.register_client_task(client_id, cmd_task).await;
197
198 let conn_sqs = sqs.clone();
207 let conn_queue = queue_url.clone();
208 let conn_protocol = protocol.clone();
209 let conn_llm = llm_client.clone();
210 let conn_state = app_state.clone();
211 let conn_status = status_tx.clone();
212 let conn_task = tokio::spawn(async move {
213 let event = Event::new(
214 &SQS_CLIENT_CONNECTED_EVENT,
215 serde_json::json!({
216 "queue_url": conn_queue.clone(),
217 }),
218 );
219
220 let Some(instruction) = conn_state.get_instruction_for_client(client_id).await else {
221 return;
222 };
223 let memory = conn_state
224 .get_memory_for_client(client_id)
225 .await
226 .unwrap_or_default();
227
228 match call_llm_for_client(
229 &conn_llm,
230 &conn_state,
231 client_id.to_string(),
232 &instruction,
233 &memory,
234 Some(&event),
235 conn_protocol.as_ref(),
236 &conn_status,
237 )
238 .await
239 {
240 Ok(ClientLlmResult {
241 actions,
242 memory_updates,
243 }) => {
244 if let Some(mem) = memory_updates {
245 conn_state.set_memory_for_client(client_id, mem).await;
246 }
247
248 if let Err(e) = Self::execute_actions(
249 actions,
250 &conn_sqs,
251 &conn_queue,
252 conn_protocol.clone(),
253 &conn_llm,
254 &conn_state,
255 &conn_status,
256 client_id,
257 0,
258 )
259 .await
260 {
261 error!(
262 "SQS client {} connect-time actions failed: {}",
263 client_id, e
264 );
265 }
266 }
267 Err(e) => {
268 error!("LLM error for SQS client {}: {}", client_id, e);
269 }
270 }
271 });
272 app_state.register_client_task(client_id, conn_task).await;
273
274 let dummy_port = 10_000u16.wrapping_add((client_id.as_u32() % 50_000) as u16);
281 let dummy_addr: SocketAddr = format!("127.0.0.1:{dummy_port}")
282 .parse()
283 .context("Failed to create dummy socket address")?;
284
285 Ok(dummy_addr)
286 }
287
288 #[allow(clippy::too_many_arguments)]
290 fn execute_actions<'a>(
291 actions: Vec<serde_json::Value>,
292 sqs: &'a aws_sdk_sqs::Client,
293 queue_url: &'a str,
294 protocol: Arc<SqsClientProtocol>,
295 llm_client: &'a OllamaClient,
296 app_state: &'a Arc<AppState>,
297 status_tx: &'a mpsc::UnboundedSender<String>,
298 client_id: ClientId,
299 depth: u8,
300 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + 'a>> {
301 Box::pin(async move {
302 for action in actions {
303 let result = match protocol.execute_action(action) {
304 Ok(result) => result,
305 Err(e) => {
306 error!("SQS client {} rejected action: {}", client_id, e);
307 continue;
308 }
309 };
310 let outcome = Self::apply_action(
311 result,
312 &sqs,
313 queue_url,
314 protocol.clone(),
315 llm_client,
316 app_state,
317 status_tx,
318 client_id,
319 depth,
320 )
321 .await;
322 if matches!(outcome, ClientSendOutcome::Disconnected) {
323 info!("SQS client {} disconnecting", client_id);
324 app_state
325 .update_client_status(client_id, ClientStatus::Disconnected)
326 .await;
327 app_state.remove_client_handle(client_id).await;
328 let _ =
329 status_tx.send(format!("[CLIENT] SQS client {} disconnected", client_id));
330 let _ = status_tx.send("__UPDATE_UI__".to_string());
331 break;
332 }
333 debug!("SQS client {} action outcome: {:?}", client_id, outcome);
334 }
335 Ok(())
336 })
337 }
338
339 #[allow(clippy::too_many_arguments)]
347 fn apply_action<'a>(
348 result: ClientActionResult,
349 sqs: &'a aws_sdk_sqs::Client,
350 queue_url: &'a str,
351 protocol: Arc<SqsClientProtocol>,
352 llm_client: &'a OllamaClient,
353 app_state: &'a Arc<AppState>,
354 status_tx: &'a mpsc::UnboundedSender<String>,
355 client_id: ClientId,
356 depth: u8,
357 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ClientSendOutcome> + Send + 'a>> {
358 Box::pin(async move {
359 match result {
360 ClientActionResult::Custom { name, data } => {
361 let outcome = match name.as_str() {
362 "send_message" => {
363 Self::send_message(
364 sqs,
365 queue_url,
366 &data,
367 protocol.clone(),
368 llm_client,
369 app_state,
370 status_tx,
371 client_id,
372 depth,
373 )
374 .await
375 }
376 "receive_messages" => {
377 Self::receive_messages(
378 sqs,
379 queue_url,
380 &data,
381 protocol.clone(),
382 llm_client,
383 app_state,
384 status_tx,
385 client_id,
386 depth,
387 )
388 .await
389 }
390 "delete_message" => {
391 Self::delete_message(sqs, queue_url, &data, client_id).await
392 }
393 "purge_queue" => Self::purge_queue(sqs, queue_url, client_id).await,
394 "get_queue_attributes" => {
395 Self::get_queue_attributes(sqs, queue_url, &data, client_id).await
396 }
397 other => Err(anyhow::anyhow!("Unknown SQS operation: {}", other)),
398 };
399 match outcome {
400 Ok(value) => ClientSendOutcome::Executed {
401 detail: format!(
402 "{} completed: {}",
403 name,
404 truncate_for_log(&value.to_string(), 200)
405 ),
406 },
407 Err(e) => {
408 error!("SQS client {} operation {} failed: {}", client_id, name, e);
409 let _ = status_tx
410 .send(format!("[ERROR] SQS operation {} failed: {}", name, e));
411 ClientSendOutcome::Rejected {
415 error: format!(
416 "{} failed: {}",
417 name,
418 truncate_for_log(&e.to_string(), 200)
419 ),
420 }
421 }
422 }
423 }
424 ClientActionResult::Disconnect => ClientSendOutcome::Disconnected,
425 ClientActionResult::WaitForMore => ClientSendOutcome::Executed {
426 detail: "wait_for_more".to_string(),
427 },
428 ClientActionResult::NoAction => ClientSendOutcome::Executed {
429 detail: "no_action".to_string(),
430 },
431 ClientActionResult::SendData(_) => ClientSendOutcome::Executed {
432 detail: "send_data has no meaning for the SQS client: it speaks the SQS API \
433 through the AWS SDK, not a socket this client owns"
434 .to_string(),
435 },
436 ClientActionResult::Multiple(_) => ClientSendOutcome::Executed {
437 detail:
438 "the SQS client's own verbs never produce Multiple; nothing was applied"
439 .to_string(),
440 },
441 }
442 })
443 }
444
445 #[allow(clippy::too_many_arguments)]
454 async fn command_loop(
455 mut command_rx: mpsc::Receiver<ClientCommand>,
456 sqs: aws_sdk_sqs::Client,
457 queue_url: String,
458 protocol: Arc<SqsClientProtocol>,
459 client_id: ClientId,
460 app_state: Arc<AppState>,
461 llm_client: OllamaClient,
462 status_tx: mpsc::UnboundedSender<String>,
463 ) {
464 while let Some(command) = command_rx.recv().await {
465 let action = command.action.clone();
466 let outcome = match protocol.execute_action(action.clone()) {
467 Err(e) => ClientSendOutcome::Rejected {
468 error: e.to_string(),
469 },
470 Ok(result) => {
471 Self::apply_action(
474 result,
475 &sqs,
476 &queue_url,
477 protocol.clone(),
478 &llm_client,
479 &app_state,
480 &status_tx,
481 client_id,
482 0,
483 )
484 .await
485 }
486 };
487 let disconnected = matches!(outcome, ClientSendOutcome::Disconnected);
488
489 app_state
490 .record_access_log(
491 AccessLogOwner::Client(client_id.as_u32()),
492 protocol.protocol_name(),
493 None,
494 "injected_action",
495 action,
496 vec![serde_json::to_value(&outcome).unwrap_or(serde_json::Value::Null)],
497 )
498 .await;
499 let _ = status_tx.send("__UPDATE_UI__".to_string());
500 crate::client::command_support::reply(command, Ok(outcome));
501
502 if disconnected {
503 info!("SQS client {} disconnecting on injected action", client_id);
504 app_state
505 .update_client_status(client_id, ClientStatus::Disconnected)
506 .await;
507 break;
508 }
509 }
510
511 info!("SQS client {} command loop stopped", client_id);
512 app_state.remove_client_handle(client_id).await;
514 let _ = status_tx.send("__UPDATE_UI__".to_string());
515 }
516
517 #[allow(clippy::too_many_arguments)]
519 async fn send_message(
520 sqs: &aws_sdk_sqs::Client,
521 queue_url: &str,
522 data: &serde_json::Value,
523 protocol: Arc<SqsClientProtocol>,
524 llm_client: &OllamaClient,
525 app_state: &Arc<AppState>,
526 status_tx: &mpsc::UnboundedSender<String>,
527 client_id: ClientId,
528 depth: u8,
529 ) -> Result<serde_json::Value> {
530 let message_body = data
531 .get("message_body")
532 .and_then(|v| v.as_str())
533 .context("Missing message_body")?;
534
535 let mut request = sqs
536 .send_message()
537 .queue_url(queue_url)
538 .message_body(message_body);
539
540 if let Some(delay) = data.get("delay_seconds").and_then(|v| v.as_i64()) {
542 request = request.delay_seconds(delay.clamp(0, 900) as i32);
544 }
545
546 if let Some(attrs) = data.get("message_attributes").and_then(|v| v.as_object()) {
548 for (key, value) in attrs {
549 if let Some(value_str) = value.as_str() {
550 let msg_attr = MessageAttributeValue::builder()
551 .data_type("String")
552 .string_value(value_str)
553 .build()?;
554 request = request.message_attributes(key.clone(), msg_attr);
555 }
556 }
557 }
558
559 let response = request
560 .send()
561 .await
562 .context("Failed to send message to SQS")?;
563
564 let message_id = response.message_id().unwrap_or("unknown").to_string();
565 info!("SQS client {} sent message: {}", client_id, message_id);
566
567 let event = Event::new(
573 &SQS_MESSAGE_SENT_EVENT,
574 serde_json::json!({
575 "message_id": message_id,
576 }),
577 );
578 Self::spawn_event_notification(
579 event,
580 sqs.clone(),
581 queue_url.to_string(),
582 protocol,
583 llm_client.clone(),
584 app_state.clone(),
585 status_tx.clone(),
586 client_id,
587 depth,
588 )
589 .await;
590
591 Ok(serde_json::json!({ "message_id": message_id }))
592 }
593
594 #[allow(clippy::too_many_arguments)]
596 async fn receive_messages(
597 sqs: &aws_sdk_sqs::Client,
598 queue_url: &str,
599 data: &serde_json::Value,
600 protocol: Arc<SqsClientProtocol>,
601 llm_client: &OllamaClient,
602 app_state: &Arc<AppState>,
603 status_tx: &mpsc::UnboundedSender<String>,
604 client_id: ClientId,
605 depth: u8,
606 ) -> Result<serde_json::Value> {
607 let mut request = sqs.receive_message().queue_url(queue_url);
608
609 if let Some(max) = data.get("max_messages").and_then(|v| v.as_i64()) {
611 request = request.max_number_of_messages(max.min(10).max(1) as i32);
612 }
613
614 if let Some(wait) = data.get("wait_time_seconds").and_then(|v| v.as_i64()) {
616 request = request.wait_time_seconds(wait.min(20).max(0) as i32);
617 }
618
619 if let Some(timeout) = data.get("visibility_timeout").and_then(|v| v.as_i64()) {
621 request = request.visibility_timeout(timeout.clamp(0, 43_200) as i32);
624 }
625
626 let response = request
627 .send()
628 .await
629 .context("Failed to receive messages from SQS")?;
630
631 let messages = response.messages();
632 let message_count = messages.len();
633 let received_ids: Vec<String> = messages
634 .iter()
635 .map(|m| m.message_id().unwrap_or("").to_string())
636 .collect();
637 info!(
638 "SQS client {} received {} messages",
639 client_id, message_count
640 );
641
642 {
647 let messages_json: Vec<serde_json::Value> = messages
649 .iter()
650 .map(|msg| {
651 let attributes: std::collections::HashMap<String, String> = msg
653 .attributes()
654 .map(|attrs| {
655 attrs
656 .iter()
657 .map(|(k, v)| (k.as_str().to_string(), v.clone()))
658 .collect()
659 })
660 .unwrap_or_default();
661
662 let message_attributes: std::collections::HashMap<String, String> = msg
663 .message_attributes()
664 .map(|attrs| {
665 attrs
666 .iter()
667 .map(|(k, v)| {
668 (k.clone(), v.string_value().unwrap_or("").to_string())
669 })
670 .collect()
671 })
672 .unwrap_or_default();
673
674 serde_json::json!({
675 "message_id": msg.message_id().unwrap_or(""),
676 "receipt_handle": msg.receipt_handle().unwrap_or(""),
677 "body": msg.body().unwrap_or(""),
678 "attributes": attributes,
679 "message_attributes": message_attributes,
680 })
681 })
682 .collect();
683
684 let event = Event::new(
686 &SQS_MESSAGE_RECEIVED_EVENT,
687 serde_json::json!({
688 "messages": messages_json,
689 }),
690 );
691
692 Self::spawn_event_notification(
695 event,
696 sqs.clone(),
697 queue_url.to_string(),
698 protocol,
699 llm_client.clone(),
700 app_state.clone(),
701 status_tx.clone(),
702 client_id,
703 depth,
704 )
705 .await;
706 }
707
708 Ok(serde_json::json!({
709 "count": message_count,
710 "message_ids": received_ids,
711 }))
712 }
713
714 #[allow(clippy::too_many_arguments)]
718 async fn spawn_event_notification(
719 event: Event,
720 sqs: aws_sdk_sqs::Client,
721 queue_url: String,
722 protocol: Arc<SqsClientProtocol>,
723 llm_client: OllamaClient,
724 app_state: Arc<AppState>,
725 status_tx: mpsc::UnboundedSender<String>,
726 client_id: ClientId,
727 depth: u8,
728 ) {
729 if depth >= MAX_FOLLOWUP_DEPTH {
734 warn!(
735 "SQS client {} stopping the follow-up chain at depth {} on {}: the model's \
736 answers have provoked {} rounds of events",
737 client_id, depth, event.event_type.id, MAX_FOLLOWUP_DEPTH
738 );
739 let _ = status_tx.send(format!(
740 "[CLIENT] SQS client {client_id} follow-up chain capped at \
741 {MAX_FOLLOWUP_DEPTH} rounds"
742 ));
743 return;
744 }
745
746 let registrar = app_state.clone();
747 let handle = tokio::spawn(async move {
748 let Some(instruction) = app_state.get_instruction_for_client(client_id).await else {
751 debug!(
752 "SQS client {} has no instruction; not asking the model about {}",
753 client_id, event.event_type.id
754 );
755 return;
756 };
757 let memory = app_state
758 .get_memory_for_client(client_id)
759 .await
760 .unwrap_or_default();
761
762 let (actions, memory_updates) = match call_llm_for_client(
767 &llm_client,
768 &app_state,
769 client_id.to_string(),
770 &instruction,
771 &memory,
772 Some(&event),
773 protocol.as_ref(),
774 &status_tx,
775 )
776 .await
777 {
778 Ok(ClientLlmResult {
779 actions,
780 memory_updates,
781 }) => (actions, memory_updates),
782 Err(e) => {
783 error!(
784 "SQS client {} LLM error on {}: {}",
785 client_id, event.event_type.id, e
786 );
787 return;
788 }
789 };
790
791 if let Some(mem) = memory_updates {
792 app_state.set_memory_for_client(client_id, mem).await;
793 }
794
795 if let Err(e) = Self::execute_actions(
796 actions,
797 &sqs,
798 &queue_url,
799 protocol,
800 &llm_client,
801 &app_state,
802 &status_tx,
803 client_id,
804 depth + 1,
805 )
806 .await
807 {
808 error!("SQS client {} follow-up actions failed: {}", client_id, e);
809 }
810 });
811 registrar.register_client_task(client_id, handle).await;
812 }
813
814 async fn delete_message(
816 sqs: &aws_sdk_sqs::Client,
817 queue_url: &str,
818 data: &serde_json::Value,
819 client_id: ClientId,
820 ) -> Result<serde_json::Value> {
821 let receipt_handle = data
822 .get("receipt_handle")
823 .and_then(|v| v.as_str())
824 .context("Missing receipt_handle")?;
825
826 sqs.delete_message()
827 .queue_url(queue_url)
828 .receipt_handle(receipt_handle)
829 .send()
830 .await
831 .context("Failed to delete message from SQS")?;
832
833 info!("SQS client {} deleted message", client_id);
834 Ok(serde_json::json!({ "deleted": true }))
835 }
836
837 async fn purge_queue(
839 sqs: &aws_sdk_sqs::Client,
840 queue_url: &str,
841 client_id: ClientId,
842 ) -> Result<serde_json::Value> {
843 sqs.purge_queue()
844 .queue_url(queue_url)
845 .send()
846 .await
847 .context("Failed to purge SQS queue")?;
848
849 info!("SQS client {} purged queue", client_id);
850 Ok(serde_json::json!({ "purged": true }))
851 }
852
853 async fn get_queue_attributes(
855 sqs: &aws_sdk_sqs::Client,
856 queue_url: &str,
857 data: &serde_json::Value,
858 client_id: ClientId,
859 ) -> Result<serde_json::Value> {
860 let mut request = sqs.get_queue_attributes().queue_url(queue_url);
861
862 if let Some(attr_names) = data.get("attribute_names").and_then(|v| v.as_array()) {
864 for attr in attr_names {
865 if let Some(attr_str) = attr.as_str() {
866 match attr_str {
868 "ApproximateNumberOfMessages" => {
869 request = request
870 .attribute_names(QueueAttributeName::ApproximateNumberOfMessages);
871 }
872 "QueueArn" => {
873 request = request.attribute_names(QueueAttributeName::QueueArn);
874 }
875 _ => {
876 request = request.attribute_names(QueueAttributeName::from(attr_str));
877 }
878 }
879 }
880 }
881 } else {
882 request = request.attribute_names(QueueAttributeName::All);
884 }
885
886 let response = request
887 .send()
888 .await
889 .context("Failed to get queue attributes from SQS")?;
890
891 let attributes: std::collections::HashMap<String, String> = response
892 .attributes()
893 .map(|attrs| {
894 attrs
895 .iter()
896 .map(|(k, v)| (k.as_str().to_string(), v.clone()))
897 .collect()
898 })
899 .unwrap_or_default();
900 info!(
901 "SQS client {} got queue attributes: {:?}",
902 client_id, attributes
903 );
904 Ok(serde_json::json!({ "attributes": attributes }))
905 }
906}