feat: present options to user when context length is exceeded (#2207)

Co-authored-by: Yingjie He <yingjiehe@squareup.com>

crates/goose-server/src/routes/context.rs

@@ -0,0 +1,64 @@
+use super::utils::verify_secret_key;
+use crate::state::AppState;
+use axum::{
+    extract::State,
+    http::{HeaderMap, StatusCode},
+    routing::post,
+    Json, Router,
+};
+use goose::message::Message;
+use serde::{Deserialize, Serialize};
+
+// Direct message serialization for context mgmt request
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ContextManageRequest {
+    messages: Vec<Message>,
+    manage_action: String,
+}
+
+// Direct message serialization for context mgmt request
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ContextManageResponse {
+    messages: Vec<Message>,
+    token_counts: Vec<usize>,
+}
+
+async fn manage_context(
+    State(state): State<AppState>,
+    headers: HeaderMap,
+    Json(request): Json<ContextManageRequest>,
+) -> Result<Json<ContextManageResponse>, StatusCode> {
+    verify_secret_key(&headers, &state)?;
+
+    // Get a lock on the shared agent
+    let agent = state.agent.read().await;
+    let agent = agent.as_ref().ok_or(StatusCode::PRECONDITION_REQUIRED)?;
+
+    let mut processed_messages: Vec<Message> = vec![];
+    let mut token_counts: Vec<usize> = vec![];
+    if request.manage_action == "trunction" {
+        (processed_messages, token_counts) = agent
+            .truncate_context(&request.messages)
+            .await
+            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    } else if request.manage_action == "summarize" {
+        (processed_messages, token_counts) = agent
+            .summarize_context(&request.messages)
+            .await
+            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    }
+
+    Ok(Json(ContextManageResponse {
+        messages: processed_messages,
+        token_counts,
+    }))
+}
+
+// Configure routes for this module
+pub fn routes(state: AppState) -> Router {
+    Router::new()
+        .route("/context/manage", post(manage_context))
+        .with_state(state)
+}

crates/goose-server/src/routes/mod.rs

@@ -2,6 +2,7 @@
 pub mod agent;
 pub mod config_management;
 pub mod configs;
+pub mod context;
 pub mod extension;
 pub mod health;
 pub mod recipe;
@@ -16,6 +17,7 @@ pub fn configure(state: crate::state::AppState) -> Router {
         .merge(health::routes())
         .merge(reply::routes(state.clone()))
         .merge(agent::routes(state.clone()))
+        .merge(context::routes(state.clone()))
         .merge(extension::routes(state.clone()))
         .merge(configs::routes(state.clone()))
         .merge(config_management::routes(state.clone()))