context_management: handle summarization in UI (#2377)

crates/goose-server/src/openapi.rs

@@ -3,8 +3,16 @@ use goose::agents::extension::ToolInfo;
 use goose::agents::ExtensionConfig;
 use goose::config::permission::PermissionLevel;
 use goose::config::ExtensionEntry;
+use goose::message::{
+    ContextLengthExceeded, FrontendToolRequest, Message, MessageContent, RedactedThinkingContent,
+    ThinkingContent, ToolConfirmationRequest, ToolRequest, ToolResponse,
+};
 use goose::permission::permission_confirmation::PrincipalType;
 use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata};
+use mcp_core::content::{Annotations, Content, EmbeddedResource, ImageContent, TextContent};
+use mcp_core::handler::ToolResultSchema;
+use mcp_core::resource::ResourceContents;
+use mcp_core::role::Role;
 use mcp_core::tool::{Tool, ToolAnnotations};
 use utoipa::OpenApi;
 
@@ -25,6 +33,7 @@ use utoipa::OpenApi;
         super::routes::config_management::upsert_permissions,
         super::routes::agent::get_tools,
         super::routes::reply::confirm_permission,
+        super::routes::context::manage_context, // Added this path
     ),
     components(schemas(
         super::routes::config_management::UpsertConfigQuery,
@@ -37,6 +46,25 @@ use utoipa::OpenApi;
         super::routes::config_management::ToolPermission,
         super::routes::config_management::UpsertPermissionsQuery,
         super::routes::reply::PermissionConfirmationRequest,
+        super::routes::context::ContextManageRequest,
+        super::routes::context::ContextManageResponse,
+        Message,
+        MessageContent,
+        Content,
+        EmbeddedResource,
+        ImageContent,
+        Annotations,
+        TextContent,
+        ToolResponse,
+        ToolRequest,
+        ToolResultSchema,
+        ToolConfirmationRequest,
+        ThinkingContent,
+        RedactedThinkingContent,
+        FrontendToolRequest,
+        ResourceContents,
+        ContextLengthExceeded,
+        Role,
         ProviderMetadata,
         ExtensionEntry,
         ExtensionConfig,

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

@@ -9,23 +9,43 @@ use axum::{
 use goose::message::Message;
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
+use utoipa::ToSchema;
 
-// Direct message serialization for context mgmt request
-#[derive(Debug, Deserialize)]
+/// Request payload for context management operations
+#[derive(Debug, Deserialize, ToSchema)]
 #[serde(rename_all = "camelCase")]
 pub struct ContextManageRequest {
-    messages: Vec<Message>,
-    manage_action: String,
+    /// Collection of messages to be managed
+    pub messages: Vec<Message>,
+    /// Operation to perform: "truncation" or "summarize"
+    pub manage_action: String,
 }
 
-// Direct message serialization for context mgmt request
-#[derive(Debug, Serialize)]
+/// Response from context management operations
+#[derive(Debug, Serialize, ToSchema)]
 #[serde(rename_all = "camelCase")]
 pub struct ContextManageResponse {
-    messages: Vec<Message>,
-    token_counts: Vec<usize>,
+    /// Processed messages after the operation
+    pub messages: Vec<Message>,
+    /// Token counts for each processed message
+    pub token_counts: Vec<usize>,
 }
 
+#[utoipa::path(
+    post,
+    path = "/context/manage",
+    request_body = ContextManageRequest,
+    responses(
+        (status = 200, description = "Context managed successfully", body = ContextManageResponse),
+        (status = 401, description = "Unauthorized - Invalid or missing API key"),
+        (status = 412, description = "Precondition failed - Agent not available"),
+        (status = 500, description = "Internal server error")
+    ),
+    security(
+        ("api_key" = [])
+    ),
+    tag = "Context Management"
+)]
 async fn manage_context(
     State(state): State<Arc<AppState>>,
     headers: HeaderMap,
@@ -40,7 +60,8 @@ async fn manage_context(
 
     let mut processed_messages: Vec<Message> = vec![];
     let mut token_counts: Vec<usize> = vec![];
-    if request.manage_action == "trunction" {
+
+    if request.manage_action == "truncation" {
         (processed_messages, token_counts) = agent
             .truncate_context(&request.messages)
             .await

crates/goose/src/message.rs

@@ -16,14 +16,17 @@ use mcp_core::role::Role;
 use mcp_core::tool::ToolCall;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
+use utoipa::ToSchema;
 
 mod tool_result_serde;
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
+#[derive(ToSchema)]
 pub struct ToolRequest {
     pub id: String,
     #[serde(with = "tool_result_serde")]
+    #[schema(value_type = Object)]
     pub tool_call: ToolResult<ToolCall>,
 }
 
@@ -45,14 +48,17 @@ impl ToolRequest {
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
+#[derive(ToSchema)]
 pub struct ToolResponse {
     pub id: String,
     #[serde(with = "tool_result_serde")]
+    #[schema(value_type = Object)]
     pub tool_result: ToolResult<Vec<Content>>,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
+#[derive(ToSchema)]
 pub struct ToolConfirmationRequest {
     pub id: String,
     pub tool_name: String,
@@ -60,31 +66,33 @@ pub struct ToolConfirmationRequest {
     pub prompt: Option<String>,
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
 pub struct ThinkingContent {
     pub thinking: String,
     pub signature: String,
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
 pub struct RedactedThinkingContent {
     pub data: String,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
+#[derive(ToSchema)]
 pub struct FrontendToolRequest {
     pub id: String,
     #[serde(with = "tool_result_serde")]
+    #[schema(value_type = Object)]
     pub tool_call: ToolResult<ToolCall>,
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
 pub struct ContextLengthExceeded {
     pub msg: String,
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
 /// Content passed inside a message, which can be both simple content and tool content
 #[serde(tag = "type", rename_all = "camelCase")]
 pub enum MessageContent {
@@ -273,7 +281,7 @@ impl From<PromptMessage> for Message {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)]
 /// A message to or from an LLM
 #[serde(rename_all = "camelCase")]
 pub struct Message {

crates/mcp-core/src/content.rs

@@ -5,8 +5,9 @@ use super::role::Role;
 use crate::resource::ResourceContents;
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct Annotations {
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -14,6 +15,8 @@ pub struct Annotations {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub priority: Option<f32>,
     #[serde(skip_serializing_if = "Option::is_none")]
+    #[schema(value_type = String, format = "date-time", example = "2023-01-01T00:00:00Z")]
+    // for openapi
     pub timestamp: Option<DateTime<Utc>>,
 }
 
@@ -33,7 +36,7 @@ impl Annotations {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TextContent {
     pub text: String,
@@ -41,7 +44,7 @@ pub struct TextContent {
     pub annotations: Option<Annotations>,
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ImageContent {
     pub data: String,
@@ -50,7 +53,7 @@ pub struct ImageContent {
     pub annotations: Option<Annotations>,
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct EmbeddedResource {
     pub resource: ResourceContents,
@@ -67,7 +70,7 @@ impl EmbeddedResource {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(tag = "type", rename_all = "camelCase")]
 pub enum Content {
     Text(TextContent),

crates/mcp-core/src/handler.rs

@@ -1,8 +1,11 @@
 use async_trait::async_trait;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
+#[allow(unused_imports)] // this is used in schema below
+use serde_json::json;
 use serde_json::Value;
 use thiserror::Error;
+use utoipa::ToSchema;
 
 #[non_exhaustive]
 #[derive(Error, Debug, Clone, Deserialize, Serialize, PartialEq)]
@@ -19,6 +22,18 @@ pub enum ToolError {
 
 pub type ToolResult<T> = std::result::Result<T, ToolError>;
 
+// Define schema manually without generics issues
+#[derive(ToSchema)]
+#[schema(example = json!({"success": true, "data": {}}))]
+pub struct ToolResultSchema {
+    #[schema(example = "Operation completed successfully")]
+    pub message: Option<String>,
+    #[schema(example = true)]
+    pub success: bool,
+    #[schema(value_type = Object)]
+    pub data: Option<serde_json::Value>,
+}
+
 #[derive(Error, Debug)]
 pub enum ResourceError {
     #[error("Execution failed: {0}")]

crates/mcp-core/src/resource.rs

@@ -1,10 +1,10 @@
+use crate::content::Annotations;
 /// Resources that servers provide to clients
 use anyhow::{anyhow, Result};
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
 use url::Url;
-
-use crate::content::Annotations;
+use utoipa::ToSchema;
 
 const EPSILON: f32 = 1e-6; // Tolerance for floating point comparison
 
@@ -28,6 +28,7 @@ pub struct Resource {
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 #[serde(rename_all = "camelCase", untagged)]
+#[derive(ToSchema)]
 pub enum ResourceContents {
     TextResourceContents {
         uri: String,

crates/mcp-core/src/role.rs

@@ -1,7 +1,8 @@
 /// Roles to describe the origin/ownership of content
 use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
 #[serde(rename_all = "lowercase")]
 pub enum Role {
     User,