what am i doing
This commit is contained in:
		
							parent
							
								
									e3611ee15a
								
							
						
					
					
						commit
						b7a29405f0
					
				
					 6 changed files with 133 additions and 87 deletions
				
			
		
							
								
								
									
										1
									
								
								Cargo.lock
									
									
									
										generated
									
									
									
								
							
							
						
						
									
										1
									
								
								Cargo.lock
									
									
									
										generated
									
									
									
								
							|  | @ -709,6 +709,7 @@ dependencies = [ | ||||||
|  "bytes", |  "bytes", | ||||||
|  "chrono", |  "chrono", | ||||||
|  "futures", |  "futures", | ||||||
|  |  "once_cell", | ||||||
|  "rustix", |  "rustix", | ||||||
|  "serde", |  "serde", | ||||||
|  "serde_json", |  "serde_json", | ||||||
|  |  | ||||||
|  | @ -9,6 +9,7 @@ async-nats = "0.33.0" | ||||||
| bytes = "1.5.0" | bytes = "1.5.0" | ||||||
| chrono = { version = "0.4.33", default-features = false, features = ["now", "serde"] } | chrono = { version = "0.4.33", default-features = false, features = ["now", "serde"] } | ||||||
| futures = { version = "0.3.30", default-features = false, features = ["std"] } | futures = { version = "0.3.30", default-features = false, features = ["std"] } | ||||||
|  | once_cell = "1.19.0" | ||||||
| rustix = { version = "0.38.30", features = ["termios", "stdio", "pty", "process"] } | rustix = { version = "0.38.30", features = ["termios", "stdio", "pty", "process"] } | ||||||
| serde = { version = "1.0.195", features = ["derive"] } | serde = { version = "1.0.195", features = ["derive"] } | ||||||
| serde_json = "1.0.111" | serde_json = "1.0.111" | ||||||
|  |  | ||||||
|  | @ -5,7 +5,7 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; | ||||||
| use serde::{Deserialize, Serialize}; | use serde::{Deserialize, Serialize}; | ||||||
| use tokio::sync::watch; | use tokio::sync::watch; | ||||||
| 
 | 
 | ||||||
| use crate::messaging::{Client, Message}; | use crate::messaging::Message; | ||||||
| 
 | 
 | ||||||
| const MEMORY_USAGE_CRITICAL_THRESHOLD: f64 = 90.0; | const MEMORY_USAGE_CRITICAL_THRESHOLD: f64 = 90.0; | ||||||
| const CPU_USAGE_CRITICAL_THRESHOLD: f32 = 90.0; | const CPU_USAGE_CRITICAL_THRESHOLD: f32 = 90.0; | ||||||
|  | @ -171,7 +171,7 @@ impl Default for HealthMonitor { | ||||||
|     } |     } | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| pub async fn init_health_subsystem(client: Client) -> HealthMonitor { | pub async fn init_health_subsystem(client: crate::messaging::Client) -> HealthMonitor { | ||||||
|     let health_monitor = HealthMonitor::new(); |     let health_monitor = HealthMonitor::new(); | ||||||
|     let health_monitor_clone = health_monitor.clone(); |     let health_monitor_clone = health_monitor.clone(); | ||||||
|     let health_monitor_ret = health_monitor.clone(); |     let health_monitor_ret = health_monitor.clone(); | ||||||
|  |  | ||||||
|  | @ -1,4 +1,7 @@ | ||||||
| use std::fmt::{Debug, Display}; | use std::{ | ||||||
|  |     fmt::{Debug, Display}, | ||||||
|  |     sync::Arc, | ||||||
|  | }; | ||||||
| 
 | 
 | ||||||
| use anyhow::{anyhow, Result}; | use anyhow::{anyhow, Result}; | ||||||
| use bytes::Bytes; | use bytes::Bytes; | ||||||
|  | @ -6,11 +9,7 @@ use futures::Stream; | ||||||
| use tokio::sync::mpsc; | use tokio::sync::mpsc; | ||||||
| use tokio_stream::StreamExt; | use tokio_stream::StreamExt; | ||||||
| 
 | 
 | ||||||
| #[derive(Clone)] | const PREFIX: &'static str = "agents.v1"; | ||||||
| pub struct Client { |  | ||||||
|     id: String, |  | ||||||
|     nats: async_nats::Client, |  | ||||||
| } |  | ||||||
| 
 | 
 | ||||||
| #[derive(Debug)] | #[derive(Debug)] | ||||||
| pub enum Subject { | pub enum Subject { | ||||||
|  | @ -46,18 +45,19 @@ impl TryFrom<&str> for Subject { | ||||||
| pub struct Message { | pub struct Message { | ||||||
|     subject: Subject, |     subject: Subject, | ||||||
|     payload: Bytes, |     payload: Bytes, | ||||||
|     reply: Option<async_nats::Subject>, |     // reply: Option<async_nats::Subject>,
 | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| impl Message { | impl Message { | ||||||
|     fn from_transport(msg: async_nats::Message) -> Result<Self> { |     fn from_transport(subject: &str, payload: Bytes) -> Result<Self> { | ||||||
|         let suffix = msg.subject.split_terminator('.').last().unwrap_or_default(); |         let subject = subject.try_into()?; | ||||||
|         let subject = suffix.try_into()?; | 
 | ||||||
|  |         // msg.headers.unwrap().get("x-idempotent-key");
 | ||||||
| 
 | 
 | ||||||
|         Ok(Message { |         Ok(Message { | ||||||
|             subject, |             subject, | ||||||
|             payload: msg.payload, |             payload, | ||||||
|             reply: msg.reply, |             // reply: None,
 | ||||||
|         }) |         }) | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|  | @ -73,20 +73,16 @@ impl Message { | ||||||
|         Ok(Message { |         Ok(Message { | ||||||
|             subject: Subject::Health, |             subject: Subject::Health, | ||||||
|             payload: Bytes::from_iter(serde_json::to_vec(&health)?), |             payload: Bytes::from_iter(serde_json::to_vec(&health)?), | ||||||
|             reply: None, |             // reply: None,
 | ||||||
|         }) |         }) | ||||||
|     } |     } | ||||||
|  | } | ||||||
| 
 | 
 | ||||||
|     async fn open_subchannels(&self, client: &Client, sender: mpsc::Sender<Self>) { | #[derive(Clone)] | ||||||
|         match self.subject { | pub struct Client { | ||||||
|             Subject::OpenTerminal => { |     id: Arc<String>, | ||||||
|                 // let subject = format!("agents.v1.{}.terminal.{}.input");
 |     nats: async_nats::Client, | ||||||
|                 // client.nats.subscribe(subject).await;
 |     subj_prefix: Arc<String>, | ||||||
|                 // sender.send(message).await.unwrap();
 |  | ||||||
|             } |  | ||||||
|             _ => {} |  | ||||||
|         } |  | ||||||
|     } |  | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| impl Client { | impl Client { | ||||||
|  | @ -101,8 +97,9 @@ impl Client { | ||||||
|         tracing::debug!("connected to NATS"); |         tracing::debug!("connected to NATS"); | ||||||
| 
 | 
 | ||||||
|         Ok(Self { |         Ok(Self { | ||||||
|             id: id.to_owned(), |  | ||||||
|             nats, |             nats, | ||||||
|  |             id: Arc::new(id.to_owned()), | ||||||
|  |             subj_prefix: Arc::new(format!("{}.{}", PREFIX, id)), | ||||||
|         }) |         }) | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|  | @ -111,26 +108,29 @@ impl Client { | ||||||
|         self.nats.publish(subject, msg.payload).await.unwrap(); |         self.nats.publish(subject, msg.payload).await.unwrap(); | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|     pub async fn reply(&self, msg: Message, mut stream: impl Stream<Item = Bytes> + Unpin) { |     // pub async fn reply(&self, msg: Message, mut stream: impl Stream<Item = Bytes> + Unpin) {
 | ||||||
|         match msg.reply { |     //     match msg.reply {
 | ||||||
|             Some(reply) => { |     //         Some(reply) => {
 | ||||||
|                 while let Some(data) = stream.next().await { |     //             while let Some(data) = stream.next().await {
 | ||||||
|                     self.nats.publish(reply.clone(), data.into()).await.unwrap(); |     //                 self.nats.publish(reply.clone(), data.into()).await.unwrap();
 | ||||||
|                 } |     //             }
 | ||||||
|             } |     //         }
 | ||||||
|             None => tracing::warn!(?msg, "tried to reply to message without a reply subject"), |     //         None => tracing::warn!(?msg, "tried to reply to message without a reply subject"),
 | ||||||
|         } |     //     }
 | ||||||
|     } |     // }
 | ||||||
| 
 | 
 | ||||||
|     pub async fn messages_channel(&self) -> Result<mpsc::Receiver<Message>> { |     pub async fn messages_channel(&self) -> Result<mpsc::Receiver<Message>> { | ||||||
|         let subject = format!("agents.v1.{}.*", self.id); |  | ||||||
|         let (sender, receiver) = mpsc::channel(100); |         let (sender, receiver) = mpsc::channel(100); | ||||||
|         let mut stream = self.subscribe(subject).await?; |         let mut stream = self.subscribe("*").await?; | ||||||
| 
 | 
 | ||||||
|         let self_clone = self.clone(); |         let self_clone = self.clone(); | ||||||
|         tokio::spawn(async move { |         tokio::spawn(async move { | ||||||
|             while let Some(msg) = stream.next().await { |             while let Some(msg) = stream.next().await { | ||||||
|                 self_clone.clone().open_subchannels(&msg).await; |                 self_clone | ||||||
|  |                     .clone() | ||||||
|  |                     .open_subchannels(&msg, sender.clone()) | ||||||
|  |                     .await; | ||||||
|  | 
 | ||||||
|                 sender.send(msg).await.unwrap(); |                 sender.send(msg).await.unwrap(); | ||||||
|             } |             } | ||||||
|         }); |         }); | ||||||
|  | @ -138,17 +138,23 @@ impl Client { | ||||||
|         Ok(receiver) |         Ok(receiver) | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|     async fn subscribe(&self, subject: String) -> Result<impl Stream<Item = Message>> { |     async fn subscribe(&self, subject: &str) -> Result<impl Stream<Item = Message>> { | ||||||
|  |         let prefix = Arc::clone(&self.subj_prefix); | ||||||
|  | 
 | ||||||
|         let stream = self |         let stream = self | ||||||
|             .nats |             .nats | ||||||
|             .subscribe(subject.clone()) |             .subscribe(format!("{}.{}", self.subj_prefix, subject)) | ||||||
|             .await? |             .await? | ||||||
|             .filter_map(|data| match Message::from_transport(data) { |             .filter_map(move |data| { | ||||||
|  |                 let subject = data.subject[..].trim_start_matches(&*prefix); | ||||||
|  | 
 | ||||||
|  |                 match Message::from_transport(subject, data.payload) { | ||||||
|                     Ok(msg) => Some(msg), |                     Ok(msg) => Some(msg), | ||||||
|                     Err(err) => { |                     Err(err) => { | ||||||
|                         tracing::warn!("{}", err); |                         tracing::warn!("{}", err); | ||||||
|                         None |                         None | ||||||
|                     } |                     } | ||||||
|  |                 } | ||||||
|             }); |             }); | ||||||
| 
 | 
 | ||||||
|         tracing::debug!("subscribed on {}", subject); |         tracing::debug!("subscribed on {}", subject); | ||||||
|  | @ -156,29 +162,29 @@ impl Client { | ||||||
|         Ok(stream) |         Ok(stream) | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|     async fn open_subchannels(self, message: &Message, sender: mpsc::Sender<Message>) { |     async fn open_subchannels(&self, message: &Message, sender: mpsc::Sender<Message>) { | ||||||
|  |         fn send_messages_from_stream( | ||||||
|  |             mut stream: impl Stream<Item = Message> + Send + Unpin + 'static, | ||||||
|  |             sender: mpsc::Sender<Message>, | ||||||
|  |         ) { | ||||||
|  |             tokio::spawn(async move { | ||||||
|  |                 while let Some(message) = stream.next().await { | ||||||
|  |                     sender.send(message).await.unwrap(); | ||||||
|  |                 } | ||||||
|  |             }); | ||||||
|  |         } | ||||||
|  | 
 | ||||||
|         match message.subject { |         match message.subject { | ||||||
|             Subject::OpenTerminal => { |             Subject::OpenTerminal => { | ||||||
|                 let terminal_id = "test"; |                 let input_stream = self.subscribe("terminal.key.input").await.unwrap(); | ||||||
|                 let stream = self |                 let resize_stream = self.subscribe("terminal.key.resize").await.unwrap(); | ||||||
|                     .subscribe(format!( | 
 | ||||||
|                         "agents.v1.{}.terminal.{}.input", |                 send_messages_from_stream(input_stream, sender.clone()); | ||||||
|                         self.id, terminal_id |                 send_messages_from_stream(resize_stream, sender); | ||||||
|                     )) |  | ||||||
|                     .await |  | ||||||
|                     .unwrap(); |  | ||||||
|             } |             } | ||||||
|             _ => {} |             _ => {} | ||||||
|         } |         } | ||||||
|     } |     } | ||||||
| 
 |  | ||||||
|     fn send_messages_from_stream(self, stream: impl Stream<Item = Message> + Send + Unpin) { |  | ||||||
|         tokio::spawn(async { |  | ||||||
|             // while let Some(msg) = stream.next().await {
 |  | ||||||
|             //
 |  | ||||||
|             // }
 |  | ||||||
|         }); |  | ||||||
|     } |  | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| impl Debug for Client { | impl Debug for Client { | ||||||
|  |  | ||||||
|  | @ -1,7 +1,7 @@ | ||||||
| use anyhow::Context; | use std::collections::HashMap; | ||||||
|  | 
 | ||||||
| use bytes::Bytes; | use bytes::Bytes; | ||||||
| use thiserror::Error; | use thiserror::Error; | ||||||
| use tokio_stream::StreamExt; |  | ||||||
| 
 | 
 | ||||||
| mod exec; | mod exec; | ||||||
| mod terminal; | mod terminal; | ||||||
|  | @ -12,6 +12,20 @@ enum ServiceError { | ||||||
|     BodyFormatError, |     BodyFormatError, | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
|  | struct Service { | ||||||
|  |     terminals: HashMap<String, ()>, | ||||||
|  | } | ||||||
|  | 
 | ||||||
|  | impl Service { | ||||||
|  |     fn serve(&mut self, message: crate::messaging::Message) { | ||||||
|  |         match message.subject() { | ||||||
|  |             crate::messaging::Subject::OpenTerminal => {} | ||||||
|  |             crate::messaging::Subject::Exec => todo!(), | ||||||
|  |             crate::messaging::Subject::Health => todo!(), | ||||||
|  |         } | ||||||
|  |     } | ||||||
|  | } | ||||||
|  | 
 | ||||||
| async fn route_message(message: crate::messaging::Message) -> Result<(), ServiceError> { | async fn route_message(message: crate::messaging::Message) -> Result<(), ServiceError> { | ||||||
|     match message.subject() { |     match message.subject() { | ||||||
|         crate::messaging::Subject::Health => {} |         crate::messaging::Subject::Health => {} | ||||||
|  | @ -30,7 +44,7 @@ async fn route_message(message: crate::messaging::Message) -> Result<(), Service | ||||||
| 
 | 
 | ||||||
| struct Ctx<T> { | struct Ctx<T> { | ||||||
|     body: T, |     body: T, | ||||||
|     // input_streams:
 |     terminals: HashMap<String, (crate::pty::Pty, tokio::process::Child)>, | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| impl<T> Ctx<T> | impl<T> Ctx<T> | ||||||
|  | @ -38,10 +52,13 @@ where | ||||||
|     T: TryFrom<Bytes>, |     T: TryFrom<Bytes>, | ||||||
| { | { | ||||||
|     fn with_body(body: Bytes) -> Result<Self, ServiceError> { |     fn with_body(body: Bytes) -> Result<Self, ServiceError> { | ||||||
|         Ok(Self { |         let body = body | ||||||
|             body: body |  | ||||||
|             .try_into() |             .try_into() | ||||||
|                 .map_err(|_err| ServiceError::BodyFormatError)?, |             .map_err(|_err| ServiceError::BodyFormatError)?; | ||||||
|  | 
 | ||||||
|  |         Ok(Self { | ||||||
|  |             body, | ||||||
|  |             terminals: HashMap::default(), | ||||||
|         }) |         }) | ||||||
|     } |     } | ||||||
| } | } | ||||||
|  |  | ||||||
|  | @ -1,38 +1,44 @@ | ||||||
|  | use std::convert::Infallible; | ||||||
|  | 
 | ||||||
| use bytes::Bytes; | use bytes::Bytes; | ||||||
| use futures::Stream; |  | ||||||
| use serde::Deserialize; | use serde::Deserialize; | ||||||
| use tokio::io::AsyncWriteExt; | use tokio::io::AsyncWriteExt; | ||||||
| use tokio_stream::StreamExt; | use tokio_stream::StreamExt; | ||||||
| use tokio_util::codec::{BytesCodec, FramedRead}; | use tokio_util::codec::{BytesCodec, FramedRead}; | ||||||
| 
 | 
 | ||||||
|  | use crate::pty::open_shell; | ||||||
|  | 
 | ||||||
| use super::Ctx; | use super::Ctx; | ||||||
| 
 | 
 | ||||||
| #[derive(Debug, Deserialize)] | #[derive(Debug, Deserialize)] | ||||||
| pub struct OpenTerminalMessage { | pub struct OpenTerminalMessage(Bytes); | ||||||
|     id: String, |  | ||||||
| } |  | ||||||
| 
 | 
 | ||||||
| impl TryFrom<Bytes> for OpenTerminalMessage { | impl TryFrom<Bytes> for OpenTerminalMessage { | ||||||
|     type Error = serde_json::Error; |     type Error = Infallible; | ||||||
| 
 | 
 | ||||||
|     fn try_from(value: Bytes) -> Result<Self, Self::Error> { |     fn try_from(value: Bytes) -> Result<Self, Self::Error> { | ||||||
|         serde_json::from_slice(&value[..]) |         Ok(Self(value)) | ||||||
|     } |     } | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| pub async fn open_terminal(ctx: Ctx<OpenTerminalMessage>) -> anyhow::Result<()> { | #[derive(Debug, Deserialize)] | ||||||
|  | pub struct TerminalInput(Bytes); | ||||||
|  | 
 | ||||||
|  | impl TryFrom<Bytes> for TerminalInput { | ||||||
|  |     type Error = Infallible; | ||||||
|  | 
 | ||||||
|  |     fn try_from(value: Bytes) -> Result<Self, Self::Error> { | ||||||
|  |         Ok(Self(value)) | ||||||
|  |     } | ||||||
|  | } | ||||||
|  | 
 | ||||||
|  | pub async fn open_terminal( | ||||||
|  |     mut ctx: Ctx<OpenTerminalMessage>, | ||||||
|  | ) -> anyhow::Result<Ctx<OpenTerminalMessage>> { | ||||||
|     let pty = crate::pty::Pty::open()?; |     let pty = crate::pty::Pty::open()?; | ||||||
|     let mut pty_clone = pty.try_clone()?; |     let shell = open_shell(pty.child()?, "/bin/bash")?; | ||||||
| 
 | 
 | ||||||
|     tokio::spawn(async move { |     let _out_stream = FramedRead::new(pty.try_clone()?, BytesCodec::new()).filter_map(|inner| { | ||||||
|         while let Some(data) = tokio_stream::once(b"foo").next().await { |  | ||||||
|             if let Err(err) = pty_clone.write_all(&data[..]).await { |  | ||||||
|                 tracing::warn!(%err, "pseudoterminal write error"); |  | ||||||
|             } |  | ||||||
|         } |  | ||||||
|     }); |  | ||||||
| 
 |  | ||||||
|     let _out_stream = FramedRead::new(pty, BytesCodec::new()).filter_map(|inner| { |  | ||||||
|         inner |         inner | ||||||
|             .map(|bytes| bytes.freeze()) |             .map(|bytes| bytes.freeze()) | ||||||
|             .map_err(|err| { |             .map_err(|err| { | ||||||
|  | @ -41,5 +47,20 @@ pub async fn open_terminal(ctx: Ctx<OpenTerminalMessage>) -> anyhow::Result<()> | ||||||
|             .ok() |             .ok() | ||||||
|     }); |     }); | ||||||
| 
 | 
 | ||||||
|     Ok(()) |     ctx.terminals.insert(String::from("test"), (pty, shell)); | ||||||
|  | 
 | ||||||
|  |     Ok(ctx) | ||||||
|  | } | ||||||
|  | 
 | ||||||
|  | pub async fn terminal_input( | ||||||
|  |     terminal_id: &str, | ||||||
|  |     mut ctx: Ctx<TerminalInput>, | ||||||
|  | ) -> anyhow::Result<Ctx<TerminalInput>> { | ||||||
|  |     let (pty, _) = ctx.terminals.get_mut(terminal_id).unwrap(); | ||||||
|  | 
 | ||||||
|  |     if let Err(err) = pty.write_all(&ctx.body.0[..]).await { | ||||||
|  |         tracing::warn!(%err, "pseudoterminal write error"); | ||||||
|  |     } | ||||||
|  | 
 | ||||||
|  |     Ok(ctx) | ||||||
| } | } | ||||||
|  |  | ||||||
		Loading…
	
		Reference in a new issue