wip

Author: brandonskiser

crates/agent/src/agent/agent_loop/model.rs

@@ -28,13 +28,13 @@ pub trait Model {
     ) -> Pin<Box<dyn Stream<Item = Result<StreamEvent, StreamError>> + Send + 'static>>;
 }
 
-/// Required for defining [Model] with a [Box<dyn Model>] for [AgentLoopRequest].
+/// Required for defining [Model] with a [Box<dyn Model>] for [super::AgentLoopRequest].
 pub trait AgentLoopModel: Model + std::fmt::Debug + Send + Sync + 'static {}
 
 // Helper blanket impl
 impl<T> AgentLoopModel for T where T: Model + std::fmt::Debug + Send + Sync + 'static {}
 
-/// The supporte
+/// The supported backends
 #[derive(Debug, Clone)]
 pub enum Models {
     Rts(RtsModel),

crates/agent/src/agent/agent_loop/types.rs

@@ -272,6 +272,7 @@ pub struct ImageBlock {
 pub enum ImageFormat {
     Gif,
     #[serde(alias = "jpg")]
+    #[strum(serialize = "jpeg", serialize = "jpg")]
     Jpeg,
     Png,
     Webp,
@@ -422,6 +423,8 @@ pub struct MetadataEvent {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct MetadataMetrics {
+    pub request_start_time: DateTime<Utc>,
+    pub request_end_time: DateTime<Utc>,
     pub time_to_first_chunk: Option<Duration>,
     pub time_between_chunks: Option<Vec<Duration>>,
     pub response_stream_len: u32,
@@ -479,6 +482,11 @@ mod tests {
         test_ser_deser!(ImageFormat, ImageFormat::Png, "png");
         test_ser_deser!(ImageFormat, ImageFormat::Webp, "webp");
         test_ser_deser!(ImageFormat, ImageFormat::Jpeg, "jpeg");
-        assert_eq!(ImageFormat::from_str("jpg").unwrap(), ImageFormat::Jpeg);
+        assert_eq!(
+            ImageFormat::from_str("jpg").unwrap(),
+            ImageFormat::Jpeg,
+            "expected 'jpg' to parse to {}",
+            ImageFormat::Jpeg
+        );
     }
 }

crates/agent/src/agent/mod.rs

@@ -2104,38 +2104,3 @@ pub enum HookStage {
     /// Hooks after executing tool uses
     PostToolUse { tool_results: Vec<ToolExecutorResult> },
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[tokio::test]
-    async fn test_collect_resources() {
-        let r = collect_resources(vec!["file://AGENTS.md"]).await;
-        println!("{:?}", r);
-    }
-
-    #[tokio::test]
-    async fn test_agent() {
-        let _ = tracing_subscriber::fmt::try_init();
-
-        let path = "/Users/bskiser/.aws/amazonq/cli-agents/idk.json";
-        let contents = tokio::fs::read_to_string(path).await.unwrap();
-        let cfg: Config = serde_json::from_str(&contents).unwrap();
-        let mut agent = Agent::from_config(cfg).await.unwrap().spawn();
-        let init_res = agent.recv().await.unwrap();
-        println!("Init res: {:?}", init_res);
-
-        agent
-            .send_prompt(SendPromptArgs {
-                content: vec![InputItem::Text("what tools do you have?".to_string())],
-            })
-            .await
-            .unwrap();
-
-        loop {
-            let res = agent.recv().await.unwrap();
-            println!("res: {:?}", res);
-        }
-    }
-}

crates/agent/src/agent/rts/mod.rs

@@ -6,9 +6,12 @@ use std::sync::Arc;
 use std::time::{
     Duration,
     Instant,
-    SystemTime,
 };
 
+use chrono::{
+    DateTime,
+    Utc,
+};
 use eyre::Result;
 use futures::Stream;
 use tokio::sync::mpsc;
@@ -111,7 +114,7 @@ impl RtsModel {
         };
 
         let request_start_time = Instant::now();
-        let request_start_time_sys = SystemTime::now();
+        let request_start_time_sys = Utc::now();
         let token_clone = cancel_token.clone();
         let result = tokio::select! {
             _ = token_clone.cancelled() => {
@@ -144,7 +147,7 @@ impl RtsModel {
         tx: mpsc::Sender<std::result::Result<StreamEvent, StreamError>>,
         token: CancellationToken,
         request_start_time: Instant,
-        request_start_time_sys: SystemTime,
+        request_start_time_sys: DateTime<Utc>,
     ) {
         match res {
             Ok(output) => {
@@ -335,35 +338,7 @@ impl Model for RtsModel {
                 .await;
         });
 
-        Box::pin(RtsDropWrapper {
-            receiver_stream: ReceiverStream::new(rx),
-            cancel_token,
-        })
-    }
-}
-
-#[derive(Debug)]
-struct RtsDropWrapper {
-    receiver_stream: ReceiverStream<Result<StreamEvent, StreamError>>,
-    cancel_token: CancellationToken,
-}
-
-impl Stream for RtsDropWrapper {
-    type Item = Result<StreamEvent, StreamError>;
-
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Option<Self::Item>> {
-        Pin::new(&mut self.receiver_stream).poll_next(cx)
-    }
-}
-
-impl Drop for RtsDropWrapper {
-    fn drop(&mut self) {
-        // TODO - I don't think RtsDropWrapper is really required here.
-        //
-        // Cancelling is already handled by agent_loop correctly (when AgentLoop is dropped, the
-        // cancel token will call cancel)
-        // debug!("rts stream dropped, cancelling");
-        // self.cancel_token.cancel();
+        Box::pin(ReceiverStream::new(rx))
     }
 }
 
@@ -393,7 +368,7 @@ struct ResponseParser {
     /// Time immediately before sending the request.
     request_start_time: Instant,
     /// Time immediately before sending the request, as a [SystemTime].
-    request_start_time_sys: SystemTime,
+    request_start_time_sys: DateTime<Utc>,
     time_to_first_chunk: Option<Duration>,
     time_between_chunks: Vec<Duration>,
     /// Total size (in bytes) of the response received so far.
@@ -407,7 +382,7 @@ impl ResponseParser {
         cancel_token: CancellationToken,
         request_id: Option<String>,
         request_start_time: Instant,
-        request_start_time_sys: SystemTime,
+        request_start_time_sys: DateTime<Utc>,
     ) -> Self {
         Self {
             response,
@@ -621,6 +596,8 @@ impl ResponseParser {
     fn make_metadata(&self) -> StreamEvent {
         StreamEvent::Metadata(MetadataEvent {
             metrics: Some(MetadataMetrics {
+                request_start_time: self.request_start_time_sys,
+                request_end_time: Utc::now(),
                 time_to_first_chunk: self.time_to_first_chunk,
                 time_between_chunks: if self.time_between_chunks.is_empty() {
                     None

crates/agent/src/agent/rts/types.rs

@@ -40,25 +40,6 @@ impl From<ToolUseBlock> for model::ToolUse {
     }
 }
 
-// impl From<ToolResultBlock> for model::ToolResult {
-//     fn from(v: ToolResultBlock) -> Self {
-//         Self {
-//             tool_use_id: v.tool_use_id,
-//             content: v.content.into_iter().map(Into::into).collect(),
-//             status: v.status.into(),
-//         }
-//     }
-// }
-
-// impl From<ToolResultContentBlock> for model::ToolResultContentBlock {
-//     fn from(v: ToolResultContentBlock) -> Self {
-//         match v {
-//             ToolResultContentBlock::Text(t) => Self::Text(t),
-//             ToolResultContentBlock::Json(v) => Self::Json(serde_value_to_document(v)),
-//         }
-//     }
-// }
-
 impl From<ToolResultStatus> for model::ToolResultStatus {
     fn from(value: ToolResultStatus) -> Self {
         match value {

crates/agent/src/agent/task_executor/mod.rs

@@ -216,7 +216,7 @@ impl TaskExecutor {
                     }
                 });
             },
-            HookConfig::Tool(tool) => (),
+            HookConfig::Tool(_) => (),
         };
 
         let start_time = Utc::now();
@@ -282,6 +282,12 @@ impl TaskExecutor {
     }
 }
 
+impl Default for TaskExecutor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 #[derive(Debug)]
 pub enum ExecuteRequest {
     Tool(StartToolExecution),
@@ -336,6 +342,7 @@ struct ExecutingHook {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
+#[allow(clippy::large_enum_variant)]
 pub enum TaskExecutorEvent {
     /// A tool has started executing
     ToolExecutionStart(ToolExecutionStartEvent),
@@ -406,6 +413,7 @@ impl ToolExecutionId {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
+#[allow(clippy::large_enum_variant)]
 pub enum ExecutorResult {
     Tool(ToolExecutorResult),
     Hook(HookExecutorResult),
@@ -505,7 +513,7 @@ impl HookResult {
     pub fn is_success(&self) -> bool {
         match self {
             HookResult::Command(res) => res.as_ref().is_ok_and(|r| r.exit_code == 0),
-            HookResult::Tool { .. } => todo!(),
+            HookResult::Tool { .. } => panic!("unimplemented"),
         }
     }
 
@@ -516,7 +524,6 @@ impl HookResult {
     pub fn output(&self) -> Option<&str> {
         match self {
             HookResult::Command(Ok(CommandResult { output, .. })) => Some(output),
-            HookResult::Tool { output } => todo!(),
             _ => None,
         }
     }
@@ -668,44 +675,42 @@ fn sanitize_user_prompt(input: &str) -> String {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::agent::types::AgentId;
-
-    const TEST_AGENT_NAME: &str = "test_agent";
 
     const TEST_COMMAND_HOOK: &str = r#"
 {
     "command": "echo hello world"
 }
 "#;
 
-    async fn run_with_timeout<T: Future>(fut: T) {
-        match tokio::time::timeout(std::time::Duration::from_millis(500), fut).await {
+    async fn run_with_timeout<T: Future>(timeout: Duration, fut: T) {
+        match tokio::time::timeout(timeout, fut).await {
             Ok(_) => (),
             Err(e) => panic!("Future failed to resolve within timeout: {}", e),
         }
     }
 
     #[tokio::test]
     async fn test_hook_execution() {
-        let mut bg = TaskExecutor::new();
-
-        let agent_id = AgentId::new(TEST_AGENT_NAME.to_string());
-        bg.start_hook_execution(StartHookExecution {
-            id: HookExecutionId {
-                hook: Hook {
-                    trigger: HookTrigger::UserPromptSubmit,
-                    config: serde_json::from_str(TEST_COMMAND_HOOK).unwrap(),
+        let mut executor = TaskExecutor::new();
+
+        executor
+            .start_hook_execution(StartHookExecution {
+                id: HookExecutionId {
+                    hook: Hook {
+                        trigger: HookTrigger::UserPromptSubmit,
+                        config: serde_json::from_str(TEST_COMMAND_HOOK).unwrap(),
+                    },
+                    tool_context: None,
                 },
-                tool_context: None,
-            },
-            prompt: None,
-        })
-        .await;
+                prompt: None,
+            })
+            .await;
 
-        run_with_timeout(async move {
+        run_with_timeout(Duration::from_millis(100), async move {
             let mut event_buf = Vec::new();
             loop {
-                bg.recv_next(&mut event_buf).await;
+                executor.recv_next(&mut event_buf).await;
+                // Check if we get a "hello world" successful hook execution.
                 if event_buf.iter().any(|ev| match ev {
                     TaskExecutorEvent::HookExecutionEnd(HookExecutionEndEvent { result, .. }) => {
                         let HookExecutorResult::Completed { result, .. } = result else {
@@ -720,9 +725,9 @@ mod tests {
                     },
                     _ => false,
                 }) {
+                    // Hook succeeded with expected output, break.
                     break;
                 }
-                println!("{:?}", event_buf);
                 event_buf.drain(..);
             }
         })

crates/agent/src/agent/tool_utils.rs

@@ -18,6 +18,7 @@ use super::tools::BuiltInTool;
 /// Categorizes different types of tool name validation failures according to the requirements by
 /// the RTS API.
 #[derive(Debug, Clone)]
+#[allow(dead_code)] // TODO
 pub struct ToolValidationError {
     mcp_server_name: String,
     tool_spec: ToolSpec,
@@ -34,14 +35,18 @@ impl ToolValidationError {
     }
 }
 
+// TODO - remove dead code. Keeping for debug purposes
 #[derive(Debug, Clone)]
 pub enum ToolValidationErrorKind {
-    OutOfSpecName { transformed_name: String },
+    OutOfSpecName {
+        #[allow(dead_code)]
+        transformed_name: String,
+    },
     EmptyName,
     NameTooLong,
     EmptyDescription,
     DescriptionTooLong,
-    NameCollision(CanonicalToolName),
+    NameCollision(#[allow(dead_code)] CanonicalToolName),
 }
 
 /// Represents a set of tool specs that conforms to backend validations.
@@ -113,9 +118,9 @@ impl SanitizedToolSpec {
 ///
 /// - `canonical_names` - List of tool names to include in the generated tool specs
 /// - `mcp_tool_specs` - Map from an MCP server name to a list of tool specs as returned by the
-/// server
+///   server
 /// - `aliases` - Map from a canonical tool name to an aliased name. This refers to the `aliases`
-/// field in the agent config
+///   field in the agent config
 pub fn sanitize_tool_specs(
     canonical_names: Vec<CanonicalToolName>,
     mcp_tool_specs: HashMap<String, Vec<ToolSpec>>,

crates/agent/src/agent/tools/execute_cmd.rs

@@ -5,8 +5,6 @@ use std::collections::HashMap;
 use std::process::Stdio;
 
 use bstr::ByteSlice as _;
-use futures::StreamExt;
-use rand::seq::IndexedRandom;
 use schemars::{
     JsonSchema,
     schema_for,