Land release blocker fixes and workspace cleanup

crates/llm/src/lib.rs

@@ -18,6 +18,8 @@ pub mod prompts;
 pub use model_manager::{recommend_tier, LlmModelId, LlmModelInfo};
 
 const DEFAULT_CONTEXT_TOKENS: u32 = 4096;
+const MAX_CONTEXT_TOKENS: u32 = 8192;
+const CONTEXT_RESERVE_TOKENS: u32 = 64;
 const GENERATION_SEED: u32 = 0;
 
 #[derive(Debug, thiserror::Error)]
@@ -26,6 +28,15 @@ pub enum EngineError {
     NotLoaded,
     #[error("LLM load failed: {0}")]
     LoadFailed(String),
+    #[error(
+        "prompt too long: {prompt_tokens} prompt tokens exceed the {available_prompt_tokens}-token prompt budget for an {context_window}-token context with {max_tokens} reserved response tokens"
+    )]
+    PromptTooLong {
+        prompt_tokens: usize,
+        max_tokens: u32,
+        available_prompt_tokens: u32,
+        context_window: u32,
+    },
     #[error("inference failed: {0}")]
     Inference(String),
     #[error("model output not valid JSON: {0}")]
@@ -149,7 +160,7 @@ impl LlmEngine {
             return Ok(String::new());
         }
 
-        let n_ctx = context_window_size(prompt_tokens.len(), config.max_tokens);
+        let n_ctx = preflight_context_window(prompt_tokens.len(), config.max_tokens)?;
         let thread_count = i32::try_from(num_cpus::get().max(1)).unwrap_or(4);
         let ctx_params = LlamaContextParams::default()
             .with_n_ctx(Some(
@@ -317,8 +328,26 @@ impl LlmEngine {
 fn context_window_size(prompt_tokens: usize, max_tokens: u32) -> u32 {
     let required = prompt_tokens
         .saturating_add(max_tokens as usize)
-        .saturating_add(64);
-    DEFAULT_CONTEXT_TOKENS.max(required.min(8192) as u32)
+        .saturating_add(CONTEXT_RESERVE_TOKENS as usize);
+    DEFAULT_CONTEXT_TOKENS.max(required.min(MAX_CONTEXT_TOKENS as usize) as u32)
+}
+
+fn preflight_context_window(prompt_tokens: usize, max_tokens: u32) -> Result<u32, EngineError> {
+    let required = prompt_tokens
+        .saturating_add(max_tokens as usize)
+        .saturating_add(CONTEXT_RESERVE_TOKENS as usize);
+    if required > MAX_CONTEXT_TOKENS as usize {
+        let available_prompt_tokens =
+            MAX_CONTEXT_TOKENS.saturating_sub(max_tokens.saturating_add(CONTEXT_RESERVE_TOKENS));
+        return Err(EngineError::PromptTooLong {
+            prompt_tokens,
+            max_tokens,
+            available_prompt_tokens,
+            context_window: MAX_CONTEXT_TOKENS,
+        });
+    }
+
+    Ok(context_window_size(prompt_tokens, max_tokens))
 }
 
 fn first_stop_index(text: &str, stop_sequences: &[String]) -> Option<usize> {
@@ -417,4 +446,24 @@ mod tests {
         let index = first_stop_index(text, &["<|im_end|>".into(), "zzz".into()]);
         assert_eq!(index, Some(5));
     }
+
+    #[test]
+    fn prompt_preflight_rejects_oversized_prompt_tokens() {
+        let err = preflight_context_window(7_105, 1_024).unwrap_err();
+        assert!(matches!(
+            err,
+            EngineError::PromptTooLong {
+                prompt_tokens: 7_105,
+                max_tokens: 1_024,
+                available_prompt_tokens: 7_104,
+                context_window: MAX_CONTEXT_TOKENS,
+            }
+        ));
+    }
+
+    #[test]
+    fn prompt_preflight_keeps_prompts_within_budget() {
+        let n_ctx = preflight_context_window(7_104, 1_024).unwrap();
+        assert_eq!(n_ctx, MAX_CONTEXT_TOKENS);
+    }
 }