Compaction overhaul (#5186)

Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: David Katz <dkatz@squareup.com>
Co-authored-by: Alex Hancock <alexhancock@block.xyz>

crates/goose-server/src/openapi.rs

@@ -18,10 +18,10 @@ use goose::config::declarative_providers::{
     DeclarativeProviderConfig, LoadedProvider, ProviderEngine,
 };
 use goose::conversation::message::{
-    ContextLengthExceeded, FrontendToolRequest, Message, MessageContent, MessageMetadata,
-    RedactedThinkingContent, SummarizationRequested, ThinkingContent, ToolConfirmationRequest,
-    ToolRequest, ToolResponse,
+    ConversationCompacted, FrontendToolRequest, Message, MessageContent, MessageMetadata,
+    RedactedThinkingContent, ThinkingContent, ToolConfirmationRequest, ToolRequest, ToolResponse,
 };
+
 use utoipa::openapi::schema::{
     AdditionalProperties, AnyOfBuilder, ArrayBuilder, ObjectBuilder, OneOfBuilder, Schema,
     SchemaFormat, SchemaType,
@@ -420,8 +420,7 @@ derive_utoipa!(Icon as IconSchema);
         RedactedThinkingContent,
         FrontendToolRequest,
         ResourceContentsSchema,
-        ContextLengthExceeded,
-        SummarizationRequested,
+        ConversationCompacted,
         JsonObjectSchema,
         RoleSchema,
         ProviderMetadata,

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

@@ -1,6 +1,6 @@
 use crate::state::AppState;
 use axum::{extract::State, http::StatusCode, routing::post, Json, Router};
-use goose::conversation::{message::Message, Conversation};
+use goose::conversation::message::Message;
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 use utoipa::ToSchema;
@@ -11,8 +11,6 @@ use utoipa::ToSchema;
 pub struct ContextManageRequest {
     /// Collection of messages to be managed
     pub messages: Vec<Message>,
-    /// Operation to perform: "truncation" or "summarize"
-    pub manage_action: String,
     /// Optional session ID for session-specific agent
     pub session_id: String,
 }
@@ -48,28 +46,20 @@ async fn manage_context(
 ) -> Result<Json<ContextManageResponse>, StatusCode> {
     let agent = state.get_agent_for_route(request.session_id).await?;
 
-    let mut processed_messages = Conversation::new_unvalidated(vec![]);
-    let mut token_counts: Vec<usize> = vec![];
-
-    if request.manage_action == "truncation" {
-        (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)?;
-    }
+    let (_, processed_messages, token_counts, _) = goose::context_mgmt::check_and_compact_messages(
+        &agent,
+        &request.messages,
+        true,
+        false,
+        None,
+        None,
+    )
+    .await
+    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    // TODO(Douwe): store into db
 
     Ok(Json(ContextManageResponse {
-        messages: processed_messages
-            .messages()
-            .iter()
-            .filter(|m| m.is_user_visible())
-            .cloned()
-            .collect(),
+        messages: processed_messages.messages().to_vec(),
         token_counts,
     }))
 }

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

@@ -121,9 +121,9 @@ impl IntoResponse for SseResponse {
     }
 }
 
-#[derive(Debug, Serialize)]
+#[derive(Debug, Serialize, utoipa::ToSchema)]
 #[serde(tag = "type")]
-enum MessageEvent {
+pub enum MessageEvent {
     Message {
         message: Message,
     },
@@ -141,6 +141,9 @@ enum MessageEvent {
         request_id: String,
         message: ServerNotification,
     },
+    UpdateConversation {
+        conversation: Conversation,
+    },
     Ping,
 }
 
@@ -306,10 +309,9 @@ pub async fn reply(
                             }
                         }
                         Ok(Some(Ok(AgentEvent::HistoryReplaced(new_messages)))) => {
-                            // Replace the message history with the compacted messages
-                            all_messages = Conversation::new_unvalidated(new_messages);
-                            // Note: We don't send this as a stream event since it's an internal operation
-                            // The client will see the compaction notification message that was sent before this event
+                            all_messages = new_messages.clone();
+                            stream_event(MessageEvent::UpdateConversation {conversation: new_messages}, &tx, &cancel_token).await;
+
                         }
                         Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => {
                             stream_event(MessageEvent::ModelChange { model, mode }, &tx, &cancel_token).await;