tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Session manager (#4648)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2025-09-26 14:56:45 -04:00
committed by GitHub
parent d0aa14a1ea
commit dc292883b7
61 changed files with 2895 additions and 5572 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/openapi.rs
+13 -26
View File
@@ -3,10 +3,11 @@ use goose::agents::extension::ToolInfo;
use goose::agents::ExtensionConfig;
use goose::config::permission::PermissionLevel;
use goose::config::ExtensionEntry;
use goose::conversation::Conversation;
use goose::permission::permission_confirmation::PrincipalType;
use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata};
use goose::session::info::SessionInfo;
use goose::session::SessionMetadata;
use goose::session::{Session, SessionInsights};
use rmcp::model::{
Annotations, Content, EmbeddedResource, ImageContent, RawEmbeddedResource, RawImageContent,
RawResource, RawTextContent, ResourceContents, Role, TextContent, Tool, ToolAnnotations,
@@ -45,8 +46,6 @@ macro_rules! derive_utoipa {
}
fn convert_schemars_to_utoipa(schema: rmcp::schemars::Schema) -> RefOr<Schema> {
// For schemars 1.0+, we need to work with the public API
// The schema is now a wrapper around a JSON Value that can be either an object or bool
if let Some(true) = schema.as_bool() {
return RefOr::T(Schema::Object(ObjectBuilder::new().build()));
}
@@ -55,12 +54,10 @@ fn convert_schemars_to_utoipa(schema: rmcp::schemars::Schema) -> RefOr<Schema> {
return RefOr::T(Schema::Object(ObjectBuilder::new().build()));
}
// For object schemas, we'll need to work with the JSON Value directly
if let Some(obj) = schema.as_object() {
return convert_json_object_to_utoipa(obj);
}
// Fallback
RefOr::T(Schema::Object(ObjectBuilder::new().build()))
}
@@ -69,12 +66,10 @@ fn convert_json_object_to_utoipa(
) -> RefOr<Schema> {
use serde_json::Value;
// Handle $ref
if let Some(Value::String(reference)) = obj.get("$ref") {
return RefOr::Ref(Ref::new(reference.clone()));
}
// Handle oneOf, allOf, anyOf
if let Some(Value::Array(one_of)) = obj.get("oneOf") {
let mut builder = OneOfBuilder::new();
for item in one_of {
@@ -105,11 +100,9 @@ fn convert_json_object_to_utoipa(
return RefOr::T(Schema::AnyOf(builder.build()));
}
// Handle type-based schemas
match obj.get("type") {
Some(Value::String(type_str)) => convert_typed_schema(type_str, obj),
Some(Value::Array(types)) => {
// Multiple types - use AnyOf
let mut builder = AnyOfBuilder::new();
for type_val in types {
if let Value::String(type_str) = type_val {
@@ -119,7 +112,7 @@ fn convert_json_object_to_utoipa(
RefOr::T(Schema::AnyOf(builder.build()))
}
None => RefOr::T(Schema::Object(ObjectBuilder::new().build())),
_ => RefOr::T(Schema::Object(ObjectBuilder::new().build())), // Handle other value types
_ => RefOr::T(Schema::Object(ObjectBuilder::new().build())),
}
}
@@ -133,7 +126,6 @@ fn convert_typed_schema(
"object" => {
let mut object_builder = ObjectBuilder::new();
// Add properties
if let Some(Value::Object(properties)) = obj.get("properties") {
for (name, prop_value) in properties {
if let Ok(prop_schema) = rmcp::schemars::Schema::try_from(prop_value.clone()) {
@@ -143,7 +135,6 @@ fn convert_typed_schema(
}
}
// Add required fields
if let Some(Value::Array(required)) = obj.get("required") {
for req in required {
if let Value::String(field_name) = req {
@@ -152,7 +143,6 @@ fn convert_typed_schema(
}
}
// Handle additional properties
if let Some(additional) = obj.get("additionalProperties") {
match additional {
Value::Bool(false) => {
@@ -178,7 +168,6 @@ fn convert_typed_schema(
"array" => {
let mut array_builder = ArrayBuilder::new();
// Add items schema
if let Some(items) = obj.get("items") {
match items {
Value::Object(_) | Value::Bool(_) => {
@@ -188,7 +177,6 @@ fn convert_typed_schema(
}
}
Value::Array(item_schemas) => {
// Multiple item types - use AnyOf
let mut any_of = AnyOfBuilder::new();
for item in item_schemas {
if let Ok(schema) = rmcp::schemars::Schema::try_from(item.clone()) {
@@ -202,7 +190,6 @@ fn convert_typed_schema(
}
}
// Add constraints
if let Some(Value::Number(min_items)) = obj.get("minItems") {
if let Some(min) = min_items.as_u64() {
array_builder = array_builder.min_items(Some(min as usize));
@@ -333,8 +320,6 @@ struct AnnotatedSchema {}
impl<'__s> ToSchema<'__s> for AnnotatedSchema {
fn schema() -> (&'__s str, utoipa::openapi::RefOr<utoipa::openapi::Schema>) {
// Create a oneOf schema with only the variants we actually use in the API
// This avoids the circular reference from RawContent::Audio(AudioContent)
let schema = Schema::OneOf(
OneOfBuilder::new()
.item(RefOr::Ref(Ref::new("#/components/schemas/RawTextContent")))
@@ -352,7 +337,6 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
}
}
#[allow(dead_code)] // Used by utoipa for OpenAPI generation
#[derive(OpenApi)]
#[openapi(
paths(
@@ -384,7 +368,10 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
super::routes::reply::confirm_permission,
super::routes::context::manage_context,
super::routes::session::list_sessions,
super::routes::session::get_session_history,
super::routes::session::get_session,
super::routes::session::get_session_insights,
super::routes::session::update_session_description,
super::routes::session::delete_session,
super::routes::schedule::create_schedule,
super::routes::schedule::list_schedules,
super::routes::schedule::delete_schedule,
@@ -419,7 +406,7 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
super::routes::context::ContextManageRequest,
super::routes::context::ContextManageResponse,
super::routes::session::SessionListResponse,
super::routes::session::SessionHistoryResponse,
super::routes::session::UpdateSessionDescriptionRequest,
Message,
MessageContent,
MessageMetadata,
@@ -454,9 +441,10 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
PermissionLevel,
PrincipalType,
ModelInfo,
SessionInfo,
SessionMetadata,
goose::session::ExtensionData,
Session,
SessionInsights,
Conversation,
goose::session::extension_data::ExtensionData,
super::routes::schedule::CreateScheduleRequest,
super::routes::schedule::UpdateScheduleRequest,
super::routes::schedule::KillJobResponse,
@@ -498,7 +486,6 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
super::routes::agent::UpdateRouterToolSelectorRequest,
super::routes::agent::StartAgentRequest,
super::routes::agent::ResumeAgentRequest,
super::routes::agent::StartAgentResponse,
super::routes::agent::ErrorResponse,
super::routes::setup::SetupResponse,
))
crates/goose-server/src/routes/agent.rs
+26 -64
View File
@@ -6,13 +6,11 @@ use axum::{
Json, Router,
};
use goose::config::PermissionManager;
use goose::conversation::message::Message;
use goose::conversation::Conversation;
use goose::model::ModelConfig;
use goose::providers::create;
use goose::recipe::{Recipe, Response};
use goose::session;
use goose::session::SessionMetadata;
use goose::session::{Session, SessionManager};
use goose::{
agents::{extension::ToolInfo, extension_manager::get_parameter_names},
config::permission::PermissionLevel,
@@ -22,7 +20,6 @@ use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use tracing::error;
#[derive(Deserialize, utoipa::ToSchema)]
pub struct ExtendPromptRequest {
@@ -81,14 +78,6 @@ pub struct ResumeAgentRequest {
session_id: String,
}
// This is the same as SessionHistoryResponse
#[derive(Serialize, utoipa::ToSchema)]
pub struct StartAgentResponse {
session_id: String,
metadata: SessionMetadata,
messages: Vec<Message>,
}
#[derive(Serialize, utoipa::ToSchema)]
pub struct ErrorResponse {
error: String,
@@ -99,7 +88,7 @@ pub struct ErrorResponse {
path = "/agent/start",
request_body = StartAgentRequest,
responses(
(status = 200, description = "Agent started successfully", body = StartAgentResponse),
(status = 200, description = "Agent started successfully", body = Session),
(status = 400, description = "Bad request - invalid working directory"),
(status = 401, description = "Unauthorized - invalid secret key"),
(status = 500, description = "Internal server error")
@@ -108,39 +97,28 @@ pub struct ErrorResponse {
async fn start_agent(
State(state): State<Arc<AppState>>,
Json(payload): Json<StartAgentRequest>,
) -> Result<Json<StartAgentResponse>, StatusCode> {
let session_id = session::generate_session_id();
) -> Result<Json<Session>, StatusCode> {
let counter = state.session_counter.fetch_add(1, Ordering::SeqCst) + 1;
let description = format!("New session {}", counter);
let metadata = SessionMetadata {
working_dir: PathBuf::from(&payload.working_dir),
description: format!("New session {}", counter),
schedule_id: None,
message_count: 0,
total_tokens: Some(0),
input_tokens: Some(0),
output_tokens: Some(0),
accumulated_total_tokens: Some(0),
accumulated_input_tokens: Some(0),
accumulated_output_tokens: Some(0),
extension_data: Default::default(),
recipe: payload.recipe,
};
let mut session =
SessionManager::create_session(PathBuf::from(&payload.working_dir), description)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
Ok(path) => path,
Err(_) => return Err(StatusCode::BAD_REQUEST),
};
if let Some(recipe) = payload.recipe {
SessionManager::update_session(&session.id)
.recipe(Some(recipe))
.apply()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let conversation = Conversation::empty();
session::storage::save_messages_with_metadata(&session_path, &metadata, &conversation)
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
session = SessionManager::get_session(&session.id, false)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
}
Ok(Json(StartAgentResponse {
session_id,
metadata,
messages: conversation.messages().clone(),
}))
Ok(Json(session))
}
#[utoipa::path(
@@ -148,7 +126,7 @@ async fn start_agent(
path = "/agent/resume",
request_body = ResumeAgentRequest,
responses(
(status = 200, description = "Agent started successfully", body = StartAgentResponse),
(status = 200, description = "Agent started successfully", body = Session),
(status = 400, description = "Bad request - invalid working directory"),
(status = 401, description = "Unauthorized - invalid secret key"),
(status = 500, description = "Internal server error")
@@ -156,28 +134,12 @@ async fn start_agent(
)]
async fn resume_agent(
Json(payload): Json<ResumeAgentRequest>,
) -> Result<Json<StartAgentResponse>, StatusCode> {
let session_path =
match session::get_path(session::Identifier::Name(payload.session_id.clone())) {
Ok(path) => path,
Err(_) => return Err(StatusCode::BAD_REQUEST),
};
) -> Result<Json<Session>, StatusCode> {
let session = SessionManager::get_session(&payload.session_id, true)
.await
.map_err(|_| StatusCode::NOT_FOUND)?;
let metadata = session::read_metadata(&session_path).map_err(|_| StatusCode::NOT_FOUND)?;
let conversation = match session::read_messages(&session_path) {
Ok(messages) => messages,
Err(e) => {
error!("Failed to read session messages: {:?}", e);
return Err(StatusCode::NOT_FOUND);
}
};
Ok(Json(StartAgentResponse {
session_id: payload.session_id.clone(),
metadata,
messages: conversation.messages().clone(),
}))
Ok(Json(session))
}
#[utoipa::path(
crates/goose-server/src/routes/config_management.rs
+2 -2
View File
@@ -822,8 +822,8 @@ mod tests {
let status_code = result.unwrap_err();
assert!(status_code == StatusCode::BAD_REQUEST,
"Expected BAD_REQUEST (authentication error) or INTERNAL_SERVER_ERROR (other errors), got: {}",
status_code
"Expected BAD_REQUEST (authentication error) or INTERNAL_SERVER_ERROR (other errors), got: {}",
status_code
);
std::env::remove_var("OPENAI_API_KEY");
crates/goose-server/src/routes/reply.rs
+12 -69
View File
@@ -11,14 +11,12 @@ use futures::{stream::StreamExt, Stream};
use goose::conversation::message::{Message, MessageContent};
use goose::conversation::Conversation;
use goose::execution::SessionExecutionMode;
use goose::permission::{Permission, PermissionConfirmation};
use goose::session::SessionManager;
use goose::{
agents::{AgentEvent, SessionConfig},
permission::permission_confirmation::PrincipalType,
};
use goose::{
permission::{Permission, PermissionConfirmation},
session,
};
use mcp_core::ToolResult;
use rmcp::model::{Content, ServerNotification};
use serde::{Deserialize, Serialize};
@@ -228,30 +226,13 @@ async fn reply_handler(
}
};
// Load session metadata to get the working directory and other config
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
Ok(path) => path,
Err(e) => {
tracing::error!("Failed to get session path for {}: {}", session_id, e);
let _ = stream_event(
MessageEvent::Error {
error: format!("Failed to get session path: {}", e),
},
&task_tx,
&cancel_token,
)
.await;
return;
}
};
let session_metadata = match session::read_metadata(&session_path) {
let session = match SessionManager::get_session(&session_id, false).await {
Ok(metadata) => metadata,
Err(e) => {
tracing::error!("Failed to read session metadata for {}: {}", session_id, e);
tracing::error!("Failed to read session for {}: {}", session_id, e);
let _ = stream_event(
MessageEvent::Error {
error: format!("Failed to read session metadata: {}", e),
error: format!("Failed to read session: {}", e),
},
&task_tx,
&cancel_token,
@@ -262,9 +243,9 @@ async fn reply_handler(
};
let session_config = SessionConfig {
id: session::Identifier::Name(session_id.clone()),
working_dir: session_metadata.working_dir.clone(),
schedule_id: session_metadata.schedule_id.clone(),
id: session_id.clone(),
working_dir: session.working_dir.clone(),
schedule_id: session.schedule_id.clone(),
execution_mode: None,
max_turns: None,
retry_config: None,
@@ -294,22 +275,6 @@ async fn reply_handler(
};
let mut all_messages = messages.clone();
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
Ok(path) => path,
Err(e) => {
tracing::error!("Failed to get session path: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: format!("Failed to get session path: {}", e),
},
&task_tx,
&cancel_token,
)
.await;
return;
}
};
let saved_message_count = all_messages.len();
let mut heartbeat_interval = tokio::time::interval(Duration::from_millis(500));
loop {
@@ -376,40 +341,18 @@ async fn reply_handler(
}
}
if all_messages.len() > saved_message_count {
if let Ok(provider) = agent.provider().await {
let provider = Arc::clone(&provider);
let session_path_clone = session_path.to_path_buf();
let all_messages_clone = all_messages.clone();
let working_dir = session_config.working_dir.clone();
tokio::spawn(async move {
if let Err(e) = session::persist_messages(
&session_path_clone,
&all_messages_clone,
Some(provider),
Some(working_dir),
)
.await
{
tracing::error!("Failed to store session history: {:?}", e);
}
});
}
}
let session_duration = session_start.elapsed();
if let Ok(metadata) = session::read_metadata(&session_path) {
let total_tokens = metadata.total_tokens.unwrap_or(0);
let message_count = metadata.message_count;
if let Ok(session) = SessionManager::get_session(&session_id, true).await {
let total_tokens = session.total_tokens.unwrap_or(0);
tracing::info!(
counter.goose.session_completions = 1,
session_type = "app",
interface = "ui",
exit_type = "normal",
duration_ms = session_duration.as_millis() as u64,
total_tokens,
message_count,
total_tokens = total_tokens,
message_count = session.message_count,
"Session completed"
);
crates/goose-server/src/routes/schedule.rs
+15 -16
View File
@@ -319,24 +319,23 @@ async fn sessions_handler(
.await
{
Ok(session_tuples) => {
// Expecting Vec<(String, goose::session::storage::SessionMetadata)>
let display_infos: Vec<SessionDisplayInfo> = session_tuples
.into_iter()
.map(|(session_name, metadata)| SessionDisplayInfo {
let mut display_infos = Vec::new();
for (session_name, session) in session_tuples {
display_infos.push(SessionDisplayInfo {
id: session_name.clone(),
name: metadata.description, // Use description as name
name: session.description,
created_at: parse_session_name_to_iso(&session_name),
working_dir: metadata.working_dir.to_string_lossy().into_owned(),
schedule_id: metadata.schedule_id, // This is the ID of the schedule itself
message_count: metadata.message_count,
total_tokens: metadata.total_tokens,
input_tokens: metadata.input_tokens,
output_tokens: metadata.output_tokens,
accumulated_total_tokens: metadata.accumulated_total_tokens,
accumulated_input_tokens: metadata.accumulated_input_tokens,
accumulated_output_tokens: metadata.accumulated_output_tokens,
})
.collect();
working_dir: session.working_dir.to_string_lossy().into_owned(),
schedule_id: session.schedule_id,
message_count: session.message_count,
total_tokens: session.total_tokens,
input_tokens: session.input_tokens,
output_tokens: session.output_tokens,
accumulated_total_tokens: session.accumulated_total_tokens,
accumulated_input_tokens: session.accumulated_input_tokens,
accumulated_output_tokens: session.accumulated_output_tokens,
});
}
Ok(Json(display_infos))
}
Err(e) => {
crates/goose-server/src/routes/session.rs
+38 -274
View File
@@ -1,7 +1,3 @@
use chrono::DateTime;
use std::collections::HashMap;
use std::sync::Arc;
use crate::state::AppState;
use axum::{
extract::Path,
@@ -9,65 +5,28 @@ use axum::{
routing::{delete, get, put},
Json, Router,
};
use goose::conversation::message::Message;
use goose::session;
use goose::session::info::{get_valid_sorted_sessions, SessionInfo, SortOrder};
use goose::session::SessionMetadata;
use goose::session::session_manager::SessionInsights;
use goose::session::{Session, SessionManager};
use serde::{Deserialize, Serialize};
use tracing::{error, info};
use std::sync::Arc;
use utoipa::ToSchema;
#[derive(Serialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct SessionListResponse {
/// List of available session information objects
sessions: Vec<SessionInfo>,
}
#[derive(Serialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct SessionHistoryResponse {
/// Unique identifier for the session
session_id: String,
/// Session metadata containing creation time and other details
metadata: SessionMetadata,
/// List of messages in the session conversation
messages: Vec<Message>,
sessions: Vec<Session>,
}
#[derive(Deserialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct UpdateSessionMetadataRequest {
pub struct UpdateSessionDescriptionRequest {
/// Updated description (name) for the session (max 200 characters)
description: String,
}
const MAX_DESCRIPTION_LENGTH: usize = 200;
#[derive(Serialize, ToSchema, Debug)]
#[serde(rename_all = "camelCase")]
pub struct SessionInsights {
/// Total number of sessions
total_sessions: usize,
/// Most active working directories with session counts
most_active_dirs: Vec<(String, usize)>,
/// Average session duration in minutes
avg_session_duration: f64,
/// Total tokens used across all sessions
total_tokens: i64,
/// Activity trend for the last 7 days
recent_activity: Vec<(String, usize)>,
}
#[derive(Serialize, ToSchema, Debug)]
#[allow(dead_code)]
#[serde(rename_all = "camelCase")]
pub struct ActivityHeatmapCell {
pub week: usize,
pub day: usize,
pub count: usize,
}
#[utoipa::path(
get,
path = "/sessions",
@@ -81,9 +40,9 @@ pub struct ActivityHeatmapCell {
),
tag = "Session Management"
)]
// List all available sessions
async fn list_sessions() -> Result<Json<SessionListResponse>, StatusCode> {
let sessions = get_valid_sorted_sessions(SortOrder::Descending)
let sessions = SessionManager::list_sessions()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
Ok(Json(SessionListResponse { sessions }))
@@ -96,7 +55,7 @@ async fn list_sessions() -> Result<Json<SessionListResponse>, StatusCode> {
("session_id" = String, Path, description = "Unique identifier for the session")
),
responses(
(status = 200, description = "Session history retrieved successfully", body = SessionHistoryResponse),
(status = 200, description = "Session history retrieved successfully", body = Session),
(status = 401, description = "Unauthorized - Invalid or missing API key"),
(status = 404, description = "Session not found"),
(status = 500, description = "Internal server error")
@@ -106,40 +65,13 @@ async fn list_sessions() -> Result<Json<SessionListResponse>, StatusCode> {
),
tag = "Session Management"
)]
// Get a specific session's history
async fn get_session_history(
Path(session_id): Path<String>,
) -> Result<Json<SessionHistoryResponse>, StatusCode> {
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
Ok(path) => path,
Err(_) => return Err(StatusCode::BAD_REQUEST),
};
async fn get_session(Path(session_id): Path<String>) -> Result<Json<Session>, StatusCode> {
let session = SessionManager::get_session(&session_id, true)
.await
.map_err(|_| StatusCode::NOT_FOUND)?;
let metadata = session::read_metadata(&session_path).map_err(|_| StatusCode::NOT_FOUND)?;
let messages = match session::read_messages(&session_path) {
Ok(messages) => messages,
Err(e) => {
error!("Failed to read session messages: {:?}", e);
return Err(StatusCode::NOT_FOUND);
}
};
// Filter messages to only include user_visible ones
let user_visible_messages: Vec<Message> = messages
.messages()
.iter()
.filter(|m| m.is_user_visible())
.cloned()
.collect();
Ok(Json(SessionHistoryResponse {
session_id,
metadata,
messages: user_visible_messages,
}))
Ok(Json(session))
}
#[utoipa::path(
get,
path = "/sessions/insights",
@@ -154,115 +86,21 @@ async fn get_session_history(
tag = "Session Management"
)]
async fn get_session_insights() -> Result<Json<SessionInsights>, StatusCode> {
info!("Received request for session insights");
let sessions = get_valid_sorted_sessions(SortOrder::Descending).map_err(|e| {
error!("Failed to get session info: {:?}", e);
StatusCode::INTERNAL_SERVER_ERROR
})?;
// Filter out sessions without descriptions
let sessions: Vec<SessionInfo> = sessions
.into_iter()
.filter(|session| !session.metadata.description.is_empty())
.collect();
info!("Found {} sessions with descriptions", sessions.len());
// Calculate insights
let total_sessions = sessions.len();
// Debug: Log if we have very few sessions, which might indicate filtering issues
if total_sessions == 0 {
info!("Warning: No sessions found with descriptions");
}
// Track directory usage
let mut dir_counts: HashMap<String, usize> = HashMap::new();
let mut total_duration = 0.0;
let mut total_tokens = 0;
let mut activity_by_date: HashMap<String, usize> = HashMap::new();
for session in &sessions {
// Track directory usage
let dir = session.metadata.working_dir.to_string_lossy().to_string();
*dir_counts.entry(dir).or_insert(0) += 1;
// Track tokens - only add positive values to prevent negative totals
if let Some(tokens) = session.metadata.accumulated_total_tokens {
match tokens.cmp(&0) {
std::cmp::Ordering::Greater => {
total_tokens += tokens as i64;
}
std::cmp::Ordering::Less => {
// Log negative token values for debugging
info!(
"Warning: Session {} has negative accumulated_total_tokens: {}",
session.id, tokens
);
}
std::cmp::Ordering::Equal => {
// Zero tokens, no action needed
}
}
}
// Track activity by date
if let Ok(date) = DateTime::parse_from_str(&session.modified, "%Y-%m-%d %H:%M:%S UTC") {
let date_str = date.format("%Y-%m-%d").to_string();
*activity_by_date.entry(date_str).or_insert(0) += 1;
}
// Calculate session duration from messages
let session_path = session::get_path(session::Identifier::Name(session.id.clone()));
if let Ok(session_path) = session_path {
if let Ok(messages) = session::read_messages(&session_path) {
if let (Some(first), Some(last)) = (messages.first(), messages.last()) {
let duration = (last.created - first.created) as f64 / 60.0; // Convert to minutes
total_duration += duration;
}
}
}
}
// Get top 3 most active directories
let mut dir_vec: Vec<(String, usize)> = dir_counts.into_iter().collect();
dir_vec.sort_by(|a, b| b.1.cmp(&a.1));
let most_active_dirs = dir_vec.into_iter().take(3).collect();
// Calculate average session duration
let avg_session_duration = if total_sessions > 0 {
total_duration / total_sessions as f64
} else {
0.0
};
// Get last 7 days of activity
let mut activity_vec: Vec<(String, usize)> = activity_by_date.into_iter().collect();
activity_vec.sort_by(|a, b| b.0.cmp(&a.0)); // Sort by date descending
let recent_activity = activity_vec.into_iter().take(7).collect();
let insights = SessionInsights {
total_sessions,
most_active_dirs,
avg_session_duration,
total_tokens,
recent_activity,
};
info!("Returning insights: {:?}", insights);
let insights = SessionManager::get_insights()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
Ok(Json(insights))
}
#[utoipa::path(
put,
path = "/sessions/{session_id}/metadata",
request_body = UpdateSessionMetadataRequest,
path = "/sessions/{session_id}/description",
request_body = UpdateSessionDescriptionRequest,
params(
("session_id" = String, Path, description = "Unique identifier for the session")
),
responses(
(status = 200, description = "Session metadata updated successfully"),
(status = 200, description = "Session description updated successfully"),
(status = 400, description = "Bad request - Description too long (max 200 characters)"),
(status = 401, description = "Unauthorized - Invalid or missing API key"),
(status = 404, description = "Session not found"),
@@ -273,27 +111,17 @@ async fn get_session_insights() -> Result<Json<SessionInsights>, StatusCode> {
),
tag = "Session Management"
)]
// Update session metadata
async fn update_session_metadata(
async fn update_session_description(
Path(session_id): Path<String>,
Json(request): Json<UpdateSessionMetadataRequest>,
Json(request): Json<UpdateSessionDescriptionRequest>,
) -> Result<StatusCode, StatusCode> {
// Validate description length
if request.description.len() > MAX_DESCRIPTION_LENGTH {
return Err(StatusCode::BAD_REQUEST);
}
let session_path = session::get_path(session::Identifier::Name(session_id.clone()))
.map_err(|_| StatusCode::BAD_REQUEST)?;
// Read current metadata
let mut metadata = session::read_metadata(&session_path).map_err(|_| StatusCode::NOT_FOUND)?;
// Update description
metadata.description = request.description;
// Save updated metadata
session::update_metadata(&session_path, &metadata)
SessionManager::update_session(&session_id)
.description(request.description)
.apply()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
@@ -302,7 +130,7 @@ async fn update_session_metadata(
#[utoipa::path(
delete,
path = "/sessions/{session_id}/delete",
path = "/sessions/{session_id}",
params(
("session_id" = String, Path, description = "Unique identifier for the session")
),
@@ -317,93 +145,29 @@ async fn update_session_metadata(
),
tag = "Session Management"
)]
// Delete a session
async fn delete_session(Path(session_id): Path<String>) -> Result<StatusCode, StatusCode> {
// Get the session path
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
Ok(path) => path,
Err(_) => return Err(StatusCode::BAD_REQUEST),
};
// Check if session file exists
if !session_path.exists() {
return Err(StatusCode::NOT_FOUND);
}
// Delete the session file
std::fs::remove_file(&session_path).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
SessionManager::delete_session(&session_id)
.await
.map_err(|e| {
if e.to_string().contains("not found") {
StatusCode::NOT_FOUND
} else {
StatusCode::INTERNAL_SERVER_ERROR
}
})?;
Ok(StatusCode::OK)
}
// Configure routes for this module
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/sessions", get(list_sessions))
.route("/sessions/{session_id}", get(get_session_history))
.route("/sessions/{session_id}/delete", delete(delete_session))
.route("/sessions/{session_id}", get(get_session))
.route("/sessions/{session_id}", delete(delete_session))
.route("/sessions/insights", get(get_session_insights))
.route(
"/sessions/{session_id}/metadata",
put(update_session_metadata),
"/sessions/{session_id}/description",
put(update_session_description),
)
.with_state(state)
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_update_session_metadata_request_deserialization() {
// Test that our request struct can be deserialized properly
let json = r#"{"description": "test description"}"#;
let request: UpdateSessionMetadataRequest = serde_json::from_str(json).unwrap();
assert_eq!(request.description, "test description");
}
#[tokio::test]
async fn test_update_session_metadata_request_validation() {
// Test empty description
let empty_request = UpdateSessionMetadataRequest {
description: "".to_string(),
};
assert_eq!(empty_request.description, "");
// Test normal description
let normal_request = UpdateSessionMetadataRequest {
description: "My Session Name".to_string(),
};
assert_eq!(normal_request.description, "My Session Name");
// Test description at max length (should be valid)
let max_length_description = "A".repeat(MAX_DESCRIPTION_LENGTH);
let max_request = UpdateSessionMetadataRequest {
description: max_length_description.clone(),
};
assert_eq!(max_request.description, max_length_description);
assert_eq!(max_request.description.len(), MAX_DESCRIPTION_LENGTH);
// Test description over max length
let over_max_description = "A".repeat(MAX_DESCRIPTION_LENGTH + 1);
let over_max_request = UpdateSessionMetadataRequest {
description: over_max_description.clone(),
};
assert_eq!(over_max_request.description, over_max_description);
assert!(over_max_request.description.len() > MAX_DESCRIPTION_LENGTH);
}
#[tokio::test]
async fn test_description_length_validation() {
// Test the validation logic used in the endpoint
let valid_description = "A".repeat(MAX_DESCRIPTION_LENGTH);
assert!(valid_description.len() <= MAX_DESCRIPTION_LENGTH);
let invalid_description = "A".repeat(MAX_DESCRIPTION_LENGTH + 1);
assert!(invalid_description.len() > MAX_DESCRIPTION_LENGTH);
// Test edge cases
assert!(String::new().len() <= MAX_DESCRIPTION_LENGTH); // Empty string
assert!("Short".len() <= MAX_DESCRIPTION_LENGTH); // Short string
}
}