fix: truncate FileContext to limits (#572)

* fix: validate FileContext limits

* truncate

* fmt

* clean up

Author: grant0417

crates/fig_api_client/src/clients/client.rs

@@ -43,6 +43,11 @@ use crate::{
     Error,
 };
 
+// Limits for FileContext
+pub const FILE_CONTEXT_LEFT_FILE_CONTENT_MAX_LEN: usize = 10240;
+pub const FILE_CONTEXT_RIGHT_FILE_CONTENT_MAX_LEN: usize = 10240;
+pub const FILE_CONTEXT_FILE_NAME_MAX_LEN: usize = 1024;
+
 mod inner {
     use amzn_codewhisperer_client::Client as CodewhispererClient;
     use amzn_consolas_client::Client as ConsolasClient;
@@ -96,7 +101,36 @@ impl Client {
         Ok(Self(inner::Inner::Consolas(ConsolasClient::from_conf(conf))))
     }
 
-    pub async fn generate_recommendations(&self, input: RecommendationsInput) -> Result<RecommendationsOutput, Error> {
+    pub async fn generate_recommendations(
+        &self,
+        mut input: RecommendationsInput,
+    ) -> Result<RecommendationsOutput, Error> {
+        let truncate_left = |s: String, max_len: usize| {
+            if s.len() > max_len {
+                s[(s.len() - max_len)..].into()
+            } else {
+                s
+            }
+        };
+
+        let truncate_right = |s: String, max_len: usize| {
+            if s.len() > max_len { s[..max_len].into() } else { s }
+        };
+
+        let filename = truncate_right(input.file_context.filename, FILE_CONTEXT_FILE_NAME_MAX_LEN);
+        let left_content = truncate_left(
+            input.file_context.left_file_content,
+            FILE_CONTEXT_LEFT_FILE_CONTENT_MAX_LEN,
+        );
+        let right_content = truncate_right(
+            input.file_context.right_file_content,
+            FILE_CONTEXT_RIGHT_FILE_CONTENT_MAX_LEN,
+        );
+
+        input.file_context.filename = filename;
+        input.file_context.left_file_content = left_content;
+        input.file_context.right_file_content = right_content;
+
         match &self.0 {
             inner::Inner::Codewhisperer(client) => Ok(codewhisperer_generate_recommendation(client, input).await?),
             inner::Inner::Consolas(client) => Ok(consolas_generate_recommendation(client, input).await?),

crates/fig_api_client/src/clients/mod.rs

@@ -2,5 +2,10 @@ mod client;
 mod shared;
 mod streaming_client;
 
-pub use client::Client;
+pub use client::{
+    Client,
+    FILE_CONTEXT_FILE_NAME_MAX_LEN,
+    FILE_CONTEXT_LEFT_FILE_CONTENT_MAX_LEN,
+    FILE_CONTEXT_RIGHT_FILE_CONTENT_MAX_LEN,
+};
 pub use streaming_client::StreamingClient;

crates/fig_api_client/src/lib.rs

@@ -1,12 +1,11 @@
-pub(crate) mod clients;
+pub mod clients;
 pub(crate) mod consts;
 pub(crate) mod credentials;
 mod customization;
 mod endpoints;
 mod error;
 pub(crate) mod interceptor;
 pub mod model;
-// mod stage;
 
 pub use clients::{
     Client,

crates/figterm/src/inline/mod.rs

@@ -10,6 +10,7 @@ use std::time::{
 };
 
 use fig_api_client::Client;
+use fig_api_client::clients::FILE_CONTEXT_LEFT_FILE_CONTENT_MAX_LEN;
 use fig_api_client::model::{
     FileContext,
     LanguageName,
@@ -232,7 +233,7 @@ pub async fn handle_request(
             },
         };
 
-        let prompt = prompt(&history, buffer);
+        let prompt = prompt(&history, buffer).unwrap();
 
         let input = RecommendationsInput {
             file_context: FileContext {
@@ -362,20 +363,28 @@ pub async fn handle_set_enabled(figterm_request: InlineShellCompletionSetEnabled
     *INLINE_ENABLED.lock().await = figterm_request.enabled;
 }
 
-fn prompt(history: &[CommandInfo], buffer: &str) -> String {
-    history
-        .iter()
-        .rev()
-        .filter_map(|c| c.command.clone())
-        .chain([buffer.into()])
-        .enumerate()
-        .fold(String::new(), |mut acc, (i, c)| {
-            if i > 0 {
-                acc.push('\n');
-            }
-            let _ = write!(acc, "{:>5}  {c}", i + 1);
-            acc
-        })
+fn prompt(history: &[CommandInfo], buffer: &str) -> Option<String> {
+    for i in (0..history.len()).rev() {
+        let formatted_prompt = history
+            .iter()
+            .rev()
+            .take(i + 1)
+            .filter_map(|c| c.command.clone())
+            .chain([buffer.into()])
+            .enumerate()
+            .fold(String::new(), |mut acc, (i, c)| {
+                if i > 0 {
+                    acc.push('\n');
+                }
+                let _ = write!(acc, "{:>5}  {c}", i + 1);
+                acc
+            });
+
+        if formatted_prompt.len() < FILE_CONTEXT_LEFT_FILE_CONTENT_MAX_LEN {
+            return Some(formatted_prompt);
+        }
+    }
+    None
 }
 
 static RE_1: LazyLock<Regex> = LazyLock::new(|| Regex::new(&format!("{}\\s+.*", *HISTORY_COUNT + 1)).unwrap());
@@ -420,7 +429,7 @@ mod tests {
             },
         ];
 
-        let prompt = prompt(&history, "echo ");
+        let prompt = prompt(&history, "echo ").unwrap();
         println!("{prompt}");
 
         assert_eq!(prompt, "    1  echo hello\n    2  echo world\n    3  echo ");
@@ -446,6 +455,16 @@ mod tests {
         );
     }
 
+    #[test]
+    fn too_long_prompt() {
+        let history = vec![CommandInfo {
+            command: Some("a".repeat(FILE_CONTEXT_LEFT_FILE_CONTENT_MAX_LEN + 1)),
+            ..Default::default()
+        }];
+
+        assert!(prompt(&history, "echo ").is_none());
+    }
+
     #[ignore = "not in CI"]
     #[tokio::test]
     async fn test_inline_suggestion_prompt() {
@@ -458,7 +477,7 @@ mod tests {
                 0,
             )
             .unwrap();
-        let prompt = prompt(&commands, "cd ");
+        let prompt = prompt(&commands, "cd ").unwrap();
 
         let client = fig_api_client::Client::new().await.unwrap();
         let out = client