Merge remote-tracking branch 'origin/main' into agent_crate_tests

Author: ekang7

.github/workflows/rust.yml

@@ -75,46 +75,46 @@ jobs:
       - name: Run tests
         run: cargo test --locked --workspace --lib --bins --test '*' --exclude fig_desktop-fuzz
 
-  cargo-clippy-windows-chat-cli:
-    name: Clippy Windows (chat_cli)
-    runs-on: windows-latest
-    timeout-minutes: 60
-    steps:
-      - uses: actions/checkout@v4
-      - uses: dtolnay/[email protected]
-        id: toolchain
-        with:
-          components: clippy
-      - uses: actions/cache@v4
-        with:
-          path: |
-            ~/.cargo/registry/index/
-            ~/.cargo/registry/cache/
-            ~/.cargo/git/db/
-            target/
-          key: cargo-clippy-windows-chat-cli-${{ hashFiles('**/Cargo.lock') }}-${{ steps.toolchain.outputs.cachekey }}
-      - run: cargo clippy --locked -p chat_cli --color always -- -D warnings
+  # cargo-clippy-windows-chat-cli:
+  #   name: Clippy Windows (chat_cli)
+  #   runs-on: windows-latest
+  #   timeout-minutes: 60
+  #   steps:
+  #     - uses: actions/checkout@v4
+  #     - uses: dtolnay/[email protected]
+  #       id: toolchain
+  #       with:
+  #         components: clippy
+  #     - uses: actions/cache@v4
+  #       with:
+  #         path: |
+  #           ~/.cargo/registry/index/
+  #           ~/.cargo/registry/cache/
+  #           ~/.cargo/git/db/
+  #           target/
+  #         key: cargo-clippy-windows-chat-cli-${{ hashFiles('**/Cargo.lock') }}-${{ steps.toolchain.outputs.cachekey }}
+  #     - run: cargo clippy --locked -p chat_cli --color always -- -D warnings
 
-  cargo-test-windows-chat-cli:
-    name: Test Windows (chat_cli)
-    runs-on: windows-latest
-    timeout-minutes: 60
-    steps:
-      - uses: actions/checkout@v4
-      - uses: dtolnay/rust-toolchain@nightly
-        id: toolchain
-        with:
-          components: llvm-tools-preview
-      - uses: actions/cache@v4
-        with:
-          path: |
-            ~/.cargo/registry/index/
-            ~/.cargo/registry/cache/
-            ~/.cargo/git/db/
-            target/
-          key: cargo-test-windows-chat-cli-${{ hashFiles('**/Cargo.lock') }}-${{ steps.toolchain.outputs.cachekey }}
-      - name: Run tests
-        run: cargo test --locked -p chat_cli
+  # cargo-test-windows-chat-cli:
+  #   name: Test Windows (chat_cli)
+  #   runs-on: windows-latest
+  #   timeout-minutes: 60
+  #   steps:
+  #     - uses: actions/checkout@v4
+  #     - uses: dtolnay/rust-toolchain@nightly
+  #       id: toolchain
+  #       with:
+  #         components: llvm-tools-preview
+  #     - uses: actions/cache@v4
+  #       with:
+  #         path: |
+  #           ~/.cargo/registry/index/
+  #           ~/.cargo/registry/cache/
+  #           ~/.cargo/git/db/
+  #           target/
+  #         key: cargo-test-windows-chat-cli-${{ hashFiles('**/Cargo.lock') }}-${{ steps.toolchain.outputs.cachekey }}
+  #     - name: Run tests
+  #       run: cargo test --locked -p chat_cli
 
   cargo-fmt:
     name: Fmt

Cargo.lock

@@ -8,10 +8,11 @@ authors = ["Amazon Q CLI Team ([email protected])", "Chay Nabors (nabochay@amazon
 edition = "2024"
 homepage = "https://aws.amazon.com/q/"
 publish = false
-version = "1.19.3"
+version = "1.19.5"
 license = "MIT OR Apache-2.0"
 
 [workspace.dependencies]
+agent = { path = "crates/agent" }
 amzn-codewhisperer-client = { path = "crates/amzn-codewhisperer-client" }
 amzn-codewhisperer-streaming-client = { path = "crates/amzn-codewhisperer-streaming-client" }
 amzn-consolas-client = { path = "crates/amzn-consolas-client" }
@@ -112,6 +113,7 @@ tempfile = "3.18.0"
 thiserror = "2.0.12"
 time = { version = "0.3.39", features = ["parsing", "formatting", "local-offset", "macros", "serde"] }
 tokio = { version = "1.45.0", features = ["full"] }
+tokio-stream = { version = "0.1.17", features = ["io-util"] }
 tokio-tungstenite = "0.26.2"
 tokio-util = { version = "0.7.15", features = ["codec", "compat"] }
 toml = "0.8.12"

Cargo.toml

@@ -69,7 +69,7 @@ tempfile.workspace = true
 thiserror.workspace = true
 time.workspace = true
 tokio.workspace = true
-tokio-stream = { version = "0.1.17", features = ["io-util"] }
+tokio-stream.workspace = true
 tokio-util.workspace = true
 tracing.workspace = true
 tracing-appender = "0.2.3"

crates/agent/Cargo.toml

@@ -6,10 +6,6 @@ use std::sync::{
 use std::time::Duration;
 
 use futures::Stream;
-use serde::{
-    Deserialize,
-    Serialize,
-};
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::sync::CancellationToken;
@@ -27,7 +23,6 @@ use super::types::{
     Message,
     ToolSpec,
 };
-use crate::agent::rts::RtsModel;
 
 /// Represents a backend implementation for a converse stream compatible API.
 ///
@@ -51,59 +46,6 @@ pub trait Model: std::fmt::Debug + Send + Sync + 'static {
     }
 }
 
-/// The supported backends
-#[derive(Debug, Clone)]
-pub enum Models {
-    Rts(RtsModel),
-    Test(MockModel),
-}
-
-impl Models {
-    pub fn state(&self) -> ModelsState {
-        match self {
-            Models::Rts(v) => ModelsState::Rts {
-                conversation_id: Some(v.conversation_id().to_string()),
-                model_id: v.model_id().map(String::from),
-            },
-            Models::Test(_) => ModelsState::Test,
-        }
-    }
-}
-
-/// A serializable representation of the state contained within [Models].
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub enum ModelsState {
-    Rts {
-        conversation_id: Option<String>,
-        model_id: Option<String>,
-    },
-    Test,
-}
-
-impl Default for ModelsState {
-    fn default() -> Self {
-        Self::Rts {
-            conversation_id: None,
-            model_id: None,
-        }
-    }
-}
-
-impl Model for Models {
-    fn stream(
-        &self,
-        messages: Vec<Message>,
-        tool_specs: Option<Vec<ToolSpec>>,
-        system_prompt: Option<String>,
-        cancel_token: CancellationToken,
-    ) -> Pin<Box<dyn Stream<Item = StreamResult> + Send + 'static>> {
-        match self {
-            Models::Rts(rts_model) => rts_model.stream(messages, tool_specs, system_prompt, cancel_token),
-            Models::Test(test_model) => test_model.stream(messages, tool_specs, system_prompt, cancel_token),
-        }
-    }
-}
-
 #[derive(Debug, Clone)]
 pub struct MockModel {
     inner: Arc<Mutex<mock::Inner>>,

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

@@ -14,11 +14,6 @@ use serde_json::Map;
 use tracing::error;
 use uuid::Uuid;
 
-use crate::api_client::error::{
-    ApiClientError,
-    ConverseStreamError,
-};
-
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub enum StreamEvent {
@@ -74,12 +69,10 @@ impl StreamError {
         self
     }
 
-    /// Helper for downcasting the error source to [ConverseStreamError].
-    ///
-    /// Just defining this here for simplicity
-    pub fn as_rts_error(&self) -> Option<&ConverseStreamError> {
+    /// Helper for downcasting a [StreamErrorSource] to a concrete type.
+    pub fn as_concrete_error<T: StreamErrorSource>(&self) -> Option<&T> {
         if let Some(source) = &self.source {
-            (*source).as_any().downcast_ref::<ConverseStreamError>()
+            (*source).as_any().downcast_ref::<T>()
         } else {
             None
         }
@@ -164,18 +157,6 @@ pub trait StreamErrorSource: std::any::Any + std::error::Error + Send + Sync {
     fn as_any(&self) -> &dyn std::any::Any;
 }
 
-impl StreamErrorSource for ConverseStreamError {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-}
-
-impl StreamErrorSource for ApiClientError {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-}
-
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct Message {
@@ -573,7 +554,6 @@ mod tests {
     use std::str::FromStr;
 
     use super::*;
-    use crate::api_client::error::ConverseStreamErrorKind;
 
     macro_rules! test_ser_deser {
         ($ty:ident, $variant:expr, $text:expr) => {
@@ -585,19 +565,6 @@ mod tests {
         };
     }
 
-    #[test]
-    fn test_other_stream_err_downcasting() {
-        let err = StreamError::new(StreamErrorKind::Interrupted).with_source(Arc::new(ConverseStreamError::new(
-            ConverseStreamErrorKind::ModelOverloadedError,
-            None::<aws_smithy_types::error::operation::BuildError>, /* annoying type inference
-                                                                     * required */
-        )));
-        assert!(
-            err.as_rts_error()
-                .is_some_and(|r| matches!(r.kind, ConverseStreamErrorKind::ModelOverloadedError))
-        );
-    }
-
     #[test]
     fn test_image_format_ser_deser() {
         test_ser_deser!(ImageFormat, ImageFormat::Gif, "gif");

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

@@ -4,7 +4,6 @@ pub mod consts;
 pub mod mcp;
 mod permissions;
 pub mod protocol;
-pub mod rts;
 pub mod task_executor;
 mod tool_utils;
 pub mod tools;