feat: add auto compaction, fix/refactor conversation state handling (#1275)

Author: brandonskiser

Cargo.lock

@@ -98,7 +98,7 @@ impl StreamingClient {
     }
 
     pub async fn send_message(&self, conversation_state: ConversationState) -> Result<SendMessageOutput, Error> {
-        debug!("Sending conversation: {:?}", conversation_state);
+        debug!("Sending conversation: {:#?}", conversation_state);
         let ConversationState {
             conversation_id,
             user_input_message,
@@ -132,8 +132,15 @@ impl StreamingClient {
                     Ok(resp) => Ok(SendMessageOutput::Codewhisperer(resp)),
                     Err(e) => {
                         let is_quota_breach = e.raw_response().is_some_and(|resp| resp.status().as_u16() == 429);
+                        let is_context_window_overflow = e.as_service_error().is_some_and(|err| {
+                            matches!(err, err if err.meta().code() == Some("ValidationException")
+                                && err.meta().message() == Some("Input is too long."))
+                        });
+
                         if is_quota_breach {
                             Err(Error::QuotaBreach("quota has reached its limit"))
+                        } else if is_context_window_overflow {
+                            Err(Error::ContextWindowOverflow)
                         } else {
                             Err(e.into())
                         }

crates/fig_api_client/src/clients/streaming_client.rs

@@ -1,6 +1,6 @@
 use amzn_codewhisperer_client::operation::generate_completions::GenerateCompletionsError;
 use amzn_codewhisperer_client::operation::list_available_customizations::ListAvailableCustomizationsError;
-use amzn_codewhisperer_streaming_client::operation::generate_assistant_response::GenerateAssistantResponseError;
+pub use amzn_codewhisperer_streaming_client::operation::generate_assistant_response::GenerateAssistantResponseError;
 // use amzn_codewhisperer_streaming_client::operation::send_message::SendMessageError as
 // CodewhispererSendMessageError;
 use amzn_codewhisperer_streaming_client::types::error::ChatResponseStreamError as CodewhispererChatResponseStreamError;
@@ -10,7 +10,7 @@ use amzn_qdeveloper_streaming_client::operation::send_message::SendMessageError
 use amzn_qdeveloper_streaming_client::types::error::ChatResponseStreamError as QDeveloperChatResponseStreamError;
 use aws_credential_types::provider::error::CredentialsError;
 use aws_smithy_runtime_api::client::orchestrator::HttpResponse;
-use aws_smithy_runtime_api::client::result::SdkError;
+pub use aws_smithy_runtime_api::client::result::SdkError;
 use aws_smithy_types::event_stream::RawMessage;
 use fig_aws_common::SdkErrorDisplay;
 use thiserror::Error;
@@ -48,6 +48,14 @@ pub enum Error {
     #[error("quota has reached its limit")]
     QuotaBreach(&'static str),
 
+    /// Returned from the backend when the user input is too large to fit within the model context
+    /// window.
+    ///
+    /// Note that we currently do not receive token usage information regarding how large the
+    /// context window is.
+    #[error("the context window has overflowed")]
+    ContextWindowOverflow,
+
     #[error(transparent)]
     SmithyBuild(#[from] aws_smithy_types::error::operation::BuildError),
 
@@ -71,6 +79,7 @@ impl Error {
             | Error::QDeveloperChatResponseStream(_)
             | Error::SmithyBuild(_)
             | Error::UnsupportedConsolas(_)
+            | Error::ContextWindowOverflow
             | Error::QuotaBreach(_) => false,
         }
     }
@@ -84,6 +93,7 @@ impl Error {
             Error::ListAvailableServices(e) => e.as_service_error().is_some(),
             Error::CodewhispererGenerateAssistantResponse(e) => e.as_service_error().is_some(),
             Error::QDeveloperSendMessage(e) => e.as_service_error().is_some(),
+            Error::ContextWindowOverflow => true,
             Error::CodewhispererChatResponseStream(_)
             | Error::QDeveloperChatResponseStream(_)
             | Error::SmithyBuild(_)

crates/fig_api_client/src/error.rs

@@ -3,7 +3,7 @@ pub(crate) mod consts;
 pub(crate) mod credentials;
 mod customization;
 mod endpoints;
-mod error;
+pub mod error;
 pub(crate) mod interceptor;
 pub mod model;
 

crates/fig_api_client/src/lib.rs

@@ -1,4 +1,8 @@
 use aws_smithy_types::Document;
+use serde::{
+    Deserialize,
+    Serialize,
+};
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 #[serde(rename_all = "camelCase")]
@@ -284,7 +288,7 @@ impl From<ToolResultContentBlock> for amzn_qdeveloper_streaming_client::types::T
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum ToolResultStatus {
     Error,
     Success,
@@ -497,7 +501,7 @@ impl From<amzn_qdeveloper_streaming_client::types::ChatResponseStream> for ChatR
     }
 }
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
 pub struct EnvState {
     pub operating_system: Option<String>,
     pub current_working_directory: Option<String>,
@@ -534,7 +538,7 @@ impl From<EnvState> for amzn_qdeveloper_streaming_client::types::EnvState {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct EnvironmentVariable {
     pub key: String,
     pub value: String,
@@ -598,9 +602,10 @@ impl From<ShellHistoryEntry> for amzn_qdeveloper_streaming_client::types::ShellH
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ShellState {
     pub shell_name: String,
+    #[serde(skip)]
     pub shell_history: Option<Vec<ShellHistoryEntry>>,
 }
 

crates/fig_api_client/src/model.rs

@@ -0,0 +1,19 @@
+use super::token_counter::TokenCounter;
+
+// These limits are the internal undocumented values from the service for each item
+
+pub const MAX_CURRENT_WORKING_DIRECTORY_LEN: usize = 256;
+
+/// Limit to send the number of messages as part of chat.
+pub const MAX_CONVERSATION_STATE_HISTORY_LEN: usize = 250;
+
+pub const MAX_TOOL_RESPONSE_SIZE: usize = 800_000;
+
+/// TODO: Use this to gracefully handle user message sizes.
+#[allow(dead_code)]
+pub const MAX_USER_MESSAGE_SIZE: usize = 600_000;
+
+/// In tokens
+pub const CONTEXT_WINDOW_SIZE: usize = 200_000;
+
+pub const MAX_CHARS: usize = TokenCounter::token_to_chars(CONTEXT_WINDOW_SIZE); // Character-based warning threshold

crates/q_cli/src/cli/chat/consts.rs

@@ -18,6 +18,7 @@ use serde::{
     Deserialize,
     Serialize,
 };
+use tracing::debug;
 
 use super::hooks::{
     Hook,
@@ -553,7 +554,7 @@ fn profile_dir_path(ctx: &Context, profile_name: &str) -> Result<PathBuf> {
 }
 
 /// Path to the context config file for `profile_name`.
-fn profile_context_path(ctx: &Context, profile_name: &str) -> Result<PathBuf> {
+pub fn profile_context_path(ctx: &Context, profile_name: &str) -> Result<PathBuf> {
     Ok(directories::chat_profiles_dir(ctx)?
         .join(profile_name)
         .join("context.json"))
@@ -564,6 +565,7 @@ fn profile_context_path(ctx: &Context, profile_name: &str) -> Result<PathBuf> {
 /// If the global configuration file doesn't exist, returns a default configuration.
 async fn load_global_config(ctx: &Context) -> Result<ContextConfig> {
     let global_path = directories::chat_global_context_path(&ctx)?;
+    debug!(?global_path, "loading profile config");
     if ctx.fs().exists(&global_path) {
         let contents = ctx.fs().read_to_string(&global_path).await?;
         let config: ContextConfig =
@@ -587,6 +589,7 @@ async fn load_global_config(ctx: &Context) -> Result<ContextConfig> {
 /// If the profile configuration file doesn't exist, creates a default configuration.
 async fn load_profile_config(ctx: &Context, profile_name: &str) -> Result<ContextConfig> {
     let profile_path = profile_context_path(ctx, profile_name)?;
+    debug!(?profile_path, "loading profile config");
     if ctx.fs().exists(&profile_path) {
         let contents = ctx.fs().read_to_string(&profile_path).await?;
         let config: ContextConfig =