Session manager (#4648)

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
Generated
+477 -324
View File
File diff suppressed because it is too large Load Diff
+4 -4
View File
@@ -3,7 +3,6 @@ use chrono::{DateTime, Utc};
use goose::conversation::Conversation;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Mutex;
@@ -18,9 +17,9 @@ pub struct BenchAgentError {
#[async_trait]
pub trait BenchBaseSession: Send + Sync {
async fn headless(&mut self, message: String) -> anyhow::Result<()>;
fn session_file(&self) -> Option<PathBuf>;
fn message_history(&self) -> Conversation;
fn get_total_token_usage(&self) -> anyhow::Result<Option<i32>>;
fn get_session_id(&self) -> anyhow::Result<String>;
}
// struct for managing agent-session-access. to be passed to evals for benchmarking
pub struct BenchAgent {
@@ -52,7 +51,8 @@ impl BenchAgent {
pub(crate) async fn get_token_usage(&self) -> Option<i32> {
self.session.get_total_token_usage().ok().flatten()
}
pub(crate) fn session_file(&self) -> Option<PathBuf> {
self.session.session_file()
pub(crate) fn get_session_id(&self) -> anyhow::Result<String> {
self.session.get_session_id()
}
}
@@ -5,6 +5,7 @@ use crate::eval_suites::{EvaluationSuite, ExtensionRequirements};
use crate::reporting::EvaluationResult;
use crate::utilities::await_process_exits;
use anyhow::{bail, Context, Result};
use goose::session::SessionManager;
use std::env;
use std::fs;
use std::future::Future;
@@ -155,15 +156,14 @@ impl EvalRunner {
.canonicalize()
.context("Failed to canonicalize current directory path")?;
BenchmarkWorkDir::deep_copy(
agent
.session_file()
.expect("Failed to get session file")
.as_path(),
here.as_path(),
false,
)
.context("Failed to copy session file to evaluation directory")?;
let session_id = agent.get_session_id()?.to_string();
let session = SessionManager::get_session(&session_id, true).await?;
let session_json = serde_json::to_string_pretty(&session)
.context("Failed to serialize session to JSON")?;
fs::write(here.join("session.json"), session_json)
.context("Failed to write session JSON to evaluation directory")?;
tracing::info!("Evaluation completed successfully");
} else {
+69 -19
View File
@@ -18,8 +18,8 @@ use crate::commands::schedule::{
use crate::commands::session::{handle_session_list, handle_session_remove};
use crate::recipes::extract_from_cli::extract_recipe_info_from_cli;
use crate::recipes::recipe::{explain_recipe, render_recipe_as_yaml};
use crate::session;
use crate::session::{build_session, SessionBuilderConfig, SessionSettings};
use goose::session::SessionManager;
use goose_bench::bench_config::BenchRunConfig;
use goose_bench::runners::bench_runner::BenchRunner;
use goose_bench::runners::eval_runner::EvalRunner;
@@ -48,26 +48,45 @@ struct Identifier {
)]
name: Option<String>,
#[arg(
long = "session-id",
value_name = "SESSION_ID",
help = "Session ID (e.g., '20250921_143022')",
long_help = "Specify a session ID directly. When used with --resume, will resume this specific session if it exists."
)]
session_id: Option<String>,
#[arg(
short,
long,
value_name = "PATH",
help = "Path for the chat session (e.g., './playground.jsonl')",
long_help = "Specify a path for your chat session. When used with --resume, will resume this specific session if it exists."
help = "Legacy: Path for the chat session",
long_help = "Legacy parameter for backward compatibility. Extracts session ID from the file path (e.g., '/path/to/20250325_200615.
jsonl' -> '20250325_200615')."
)]
path: Option<PathBuf>,
}
fn extract_identifier(identifier: Identifier) -> session::Identifier {
if let Some(name) = identifier.name {
session::Identifier::Name(name)
async fn get_session_id(identifier: Identifier) -> Result<String> {
if let Some(session_id) = identifier.session_id {
Ok(session_id)
} else if let Some(name) = identifier.name {
let sessions = SessionManager::list_sessions().await?;
sessions
.into_iter()
.find(|s| s.description == name)
.map(|s| s.id)
.ok_or_else(|| anyhow::anyhow!("No session found with name '{}'", name))
} else if let Some(path) = identifier.path {
session::Identifier::Path(path)
path.file_stem()
.and_then(|s| s.to_str())
.map(|s| s.to_string())
.ok_or_else(|| anyhow::anyhow!("Could not extract session ID from path: {:?}", path))
} else {
unreachable!()
}
}
fn parse_key_val(s: &str) -> Result<(String, String), String> {
match s.split_once('=') {
Some((key, value)) => Ok((key.to_string(), value.to_string())),
@@ -121,6 +140,14 @@ enum SessionCommand {
long_help = "Path to save the exported Markdown. If not provided, output will be sent to stdout"
)]
output: Option<PathBuf>,
#[arg(
long = "format",
value_name = "FORMAT",
help = "Output format (markdown, json, yaml)",
default_value = "markdown"
)]
format: String,
},
}
@@ -768,19 +795,24 @@ pub async fn cli() -> Result<()> {
format,
ascending,
}) => {
handle_session_list(verbose, format, ascending)?;
handle_session_list(verbose, format, ascending).await?;
Ok(())
}
Some(SessionCommand::Remove { id, regex }) => {
handle_session_remove(id, regex)?;
handle_session_remove(id, regex).await?;
return Ok(());
}
Some(SessionCommand::Export { identifier, output }) => {
Some(SessionCommand::Export {
identifier,
output,
format,
}) => {
let session_identifier = if let Some(id) = identifier {
extract_identifier(id)
get_session_id(id).await?
} else {
// If no identifier is provided, prompt for interactive selection
match crate::commands::session::prompt_interactive_session_selection() {
match crate::commands::session::prompt_interactive_session_selection().await
{
Ok(id) => id,
Err(e) => {
eprintln!("Error: {}", e);
@@ -789,7 +821,12 @@ pub async fn cli() -> Result<()> {
}
};
crate::commands::session::handle_session_export(session_identifier, output)?;
crate::commands::session::handle_session_export(
session_identifier,
output,
format,
)
.await?;
Ok(())
}
None => {
@@ -803,9 +840,15 @@ pub async fn cli() -> Result<()> {
"Session started"
);
let session_id = if let Some(id) = identifier {
Some(get_session_id(id).await?)
} else {
None
};
// Run session command by default
let mut session: crate::Session = build_session(SessionBuilderConfig {
identifier: identifier.map(extract_identifier),
let mut session: crate::CliSession = build_session(SessionBuilderConfig {
session_id,
resume,
no_session: false,
extensions,
@@ -841,6 +884,7 @@ pub async fn cli() -> Result<()> {
let (total_tokens, message_count) = session
.get_metadata()
.await
.map(|m| (m.total_tokens.unwrap_or(0), m.message_count))
.unwrap_or((0, 0));
@@ -994,9 +1038,14 @@ pub async fn cli() -> Result<()> {
std::process::exit(1);
}
};
let session_id = if let Some(id) = identifier {
Some(get_session_id(id).await?)
} else {
None
};
let mut session = build_session(SessionBuilderConfig {
identifier: identifier.map(extract_identifier),
session_id,
resume,
no_session,
extensions,
@@ -1048,6 +1097,7 @@ pub async fn cli() -> Result<()> {
let (total_tokens, message_count) = session
.get_metadata()
.await
.map(|m| (m.total_tokens.unwrap_or(0), m.message_count))
.unwrap_or((0, 0));
@@ -1172,7 +1222,7 @@ pub async fn cli() -> Result<()> {
} else {
// Run session command by default
let mut session = build_session(SessionBuilderConfig {
identifier: None,
session_id: None,
resume: false,
no_session: false,
extensions: Vec::new(),
@@ -1188,7 +1238,7 @@ pub async fn cli() -> Result<()> {
max_tool_repetitions: None,
max_turns: None,
scheduled_job_id: None,
interactive: true, // Default case is always interactive
interactive: true,
quiet: false,
sub_recipes: None,
final_output_response: None,
+14 -12
View File
@@ -1,38 +1,42 @@
use crate::session::build_session;
use crate::session::SessionBuilderConfig;
use crate::{logging, session, Session};
use crate::{logging, CliSession};
use async_trait::async_trait;
use goose::conversation::Conversation;
use goose_bench::bench_session::{BenchAgent, BenchBaseSession};
use goose_bench::eval_suites::ExtensionRequirements;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Mutex;
// allow session obj to be used in benchmarking
#[async_trait]
impl BenchBaseSession for Session {
impl BenchBaseSession for CliSession {
async fn headless(&mut self, message: String) -> anyhow::Result<()> {
self.headless(message).await
}
fn session_file(&self) -> Option<PathBuf> {
self.session_file()
}
fn message_history(&self) -> Conversation {
self.message_history()
}
fn get_total_token_usage(&self) -> anyhow::Result<Option<i32>> {
self.get_total_token_usage()
// Since the trait requires sync but the session method is async,
// we need to block on the async call
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(self.get_total_token_usage())
})
}
fn get_session_id(&self) -> anyhow::Result<String> {
self.session_id()
.cloned()
.ok_or_else(|| anyhow::anyhow!("No session ID available"))
}
}
pub async fn agent_generator(
requirements: ExtensionRequirements,
session_id: String,
) -> BenchAgent {
let identifier = Some(session::Identifier::Name(session_id));
let base_session = build_session(SessionBuilderConfig {
identifier,
session_id: Some(session_id),
resume: false,
no_session: false,
extensions: requirements.external,
@@ -56,10 +60,8 @@ pub async fn agent_generator(
})
.await;
// package session obj into benchmark-compatible struct
let bench_agent = BenchAgent::new(Box::new(base_session));
// Initialize logging with error capture
let errors = Some(Arc::new(Mutex::new(bench_agent.get_errors().await)));
logging::setup_logging(Some("bench"), errors).expect("Failed to initialize logging");
+1 -2
View File
@@ -219,11 +219,10 @@ pub async fn handle_schedule_sessions(id: String, limit: Option<u32>) -> Result<
// sessions is now Vec<(String, SessionMetadata)>
for (session_name, metadata) in sessions {
println!(
" - Session ID: {}, Working Dir: {}, Description: \"{}\", Messages: {}, Schedule ID: {:?}",
" - Session ID: {}, Working Dir: {}, Description: \"{}\", Schedule ID: {:?}",
session_name, // Display the session_name as Session ID
metadata.working_dir.display(),
metadata.description,
metadata.message_count,
metadata.schedule_id.as_deref().unwrap_or("N/A")
);
}
+70 -109
View File
@@ -1,19 +1,19 @@
use crate::session::message_to_markdown;
use anyhow::{Context, Result};
use cliclack::{confirm, multiselect, select};
use goose::session::info::{get_valid_sorted_sessions, SessionInfo, SortOrder};
use goose::session::{self, Identifier};
use goose::session::{Session, SessionManager};
use goose::utils::safe_truncate;
use regex::Regex;
use std::fs;
use std::path::{Path, PathBuf};
use std::path::PathBuf;
const TRUNCATED_DESC_LENGTH: usize = 60;
pub fn remove_sessions(sessions: Vec<SessionInfo>) -> Result<()> {
pub async fn remove_sessions(sessions: Vec<Session>) -> Result<()> {
println!("The following sessions will be removed:");
for session in &sessions {
println!("- {}", session.id);
println!("- {} {}", session.id, session.description);
}
let should_delete = confirm("Are you sure you want to delete these sessions?")
@@ -22,8 +22,7 @@ pub fn remove_sessions(sessions: Vec<SessionInfo>) -> Result<()> {
if should_delete {
for session in sessions {
fs::remove_file(session.path.clone())
.with_context(|| format!("Failed to remove session file '{}'", session.path))?;
SessionManager::delete_session(&session.id).await?;
println!("Session `{}` removed.", session.id);
}
} else {
@@ -33,7 +32,7 @@ pub fn remove_sessions(sessions: Vec<SessionInfo>) -> Result<()> {
Ok(())
}
fn prompt_interactive_session_removal(sessions: &[SessionInfo]) -> Result<Vec<SessionInfo>> {
fn prompt_interactive_session_removal(sessions: &[Session]) -> Result<Vec<Session>> {
if sessions.is_empty() {
println!("No sessions to delete.");
return Ok(vec![]);
@@ -43,16 +42,16 @@ fn prompt_interactive_session_removal(sessions: &[SessionInfo]) -> Result<Vec<Se
"Select sessions to delete (use spacebar, Enter to confirm, Ctrl+C to cancel):",
);
let display_map: std::collections::HashMap<String, SessionInfo> = sessions
let display_map: std::collections::HashMap<String, Session> = sessions
.iter()
.map(|s| {
let desc = if s.metadata.description.is_empty() {
let desc = if s.description.is_empty() {
"(no description)"
} else {
&s.metadata.description
&s.description
};
let truncated_desc = safe_truncate(desc, TRUNCATED_DESC_LENGTH);
let display_text = format!("{} - {} ({})", s.modified, truncated_desc, s.id);
let display_text = format!("{} - {} ({})", s.updated_at, truncated_desc, s.id);
(display_text, s.clone())
})
.collect();
@@ -63,7 +62,7 @@ fn prompt_interactive_session_removal(sessions: &[SessionInfo]) -> Result<Vec<Se
let selected_display_texts: Vec<String> = selector.interact()?;
let selected_sessions: Vec<SessionInfo> = selected_display_texts
let selected_sessions: Vec<Session> = selected_display_texts
.into_iter()
.filter_map(|text| display_map.get(&text).cloned())
.collect();
@@ -71,8 +70,8 @@ fn prompt_interactive_session_removal(sessions: &[SessionInfo]) -> Result<Vec<Se
Ok(selected_sessions)
}
pub fn handle_session_remove(id: Option<String>, regex_string: Option<String>) -> Result<()> {
let all_sessions = match get_valid_sorted_sessions(SortOrder::Descending) {
pub async fn handle_session_remove(id: Option<String>, regex_string: Option<String>) -> Result<()> {
let all_sessions = match SessionManager::list_sessions().await {
Ok(sessions) => sessions,
Err(e) => {
tracing::error!("Failed to retrieve sessions: {:?}", e);
@@ -80,7 +79,7 @@ pub fn handle_session_remove(id: Option<String>, regex_string: Option<String>) -
}
};
let matched_sessions: Vec<SessionInfo>;
let matched_sessions: Vec<Session>;
if let Some(id_val) = id {
if let Some(session) = all_sessions.iter().find(|s| s.id == id_val) {
@@ -112,23 +111,16 @@ pub fn handle_session_remove(id: Option<String>, regex_string: Option<String>) -
return Ok(());
}
remove_sessions(matched_sessions)
remove_sessions(matched_sessions).await
}
pub fn handle_session_list(verbose: bool, format: String, ascending: bool) -> Result<()> {
let sort_order = if ascending {
SortOrder::Ascending
pub async fn handle_session_list(verbose: bool, format: String, ascending: bool) -> Result<()> {
let mut sessions = SessionManager::list_sessions().await?;
if ascending {
sessions.sort_by(|a, b| a.updated_at.cmp(&b.updated_at));
} else {
SortOrder::Descending
};
let sessions = match get_valid_sorted_sessions(sort_order) {
Ok(sessions) => sessions,
Err(e) => {
tracing::error!("Failed to list sessions: {:?}", e);
return Err(anyhow::anyhow!("Failed to list sessions"));
}
};
sessions.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));
}
match format.as_str() {
"json" => {
@@ -138,27 +130,18 @@ pub fn handle_session_list(verbose: bool, format: String, ascending: bool) -> Re
if sessions.is_empty() {
println!("No sessions found");
return Ok(());
} else {
println!("Available sessions:");
for SessionInfo {
id,
path,
metadata,
modified,
} in sessions
{
let description = if metadata.description.is_empty() {
"(none)"
} else {
&metadata.description
};
let output = format!("{} - {} - {}", id, description, modified);
if verbose {
println!(" {}", output);
println!(" Path: {}", path);
} else {
println!("{}", output);
}
}
println!("Available sessions:");
for session in sessions {
let output = format!(
"{} - {} - {}",
session.id, session.description, session.updated_at
);
if verbose {
println!(" {}", output);
} else {
println!("{}", output);
}
}
}
@@ -166,68 +149,55 @@ pub fn handle_session_list(verbose: bool, format: String, ascending: bool) -> Re
Ok(())
}
/// Export a session to Markdown without creating a full Session object
///
/// This function directly reads messages from the session file and converts them to Markdown
/// without creating an Agent or prompting about working directories.
pub fn handle_session_export(identifier: Identifier, output_path: Option<PathBuf>) -> Result<()> {
// Get the session file path
let session_file_path = match goose::session::get_path(identifier.clone()) {
Ok(path) => path,
pub async fn handle_session_export(
session_id: String,
output_path: Option<PathBuf>,
format: String,
) -> Result<()> {
let session = match SessionManager::get_session(&session_id, true).await {
Ok(session) => session,
Err(e) => {
return Err(anyhow::anyhow!("Invalid session identifier: {}", e));
return Err(anyhow::anyhow!(
"Session '{}' not found or failed to read: {}",
session_id,
e
));
}
};
if !session_file_path.exists() {
return Err(anyhow::anyhow!(
"Session file not found (expected path: {})",
session_file_path.display()
));
}
// Read messages directly without using Session
let messages = match goose::session::read_messages(&session_file_path) {
Ok(msgs) => msgs,
Err(e) => {
return Err(anyhow::anyhow!("Failed to read session messages: {}", e));
let output = match format.as_str() {
"json" => serde_json::to_string_pretty(&session)?,
"yaml" => serde_yaml::to_string(&session)?,
"markdown" => {
let conversation = session
.conversation
.ok_or_else(|| anyhow::anyhow!("Session has no messages"))?;
export_session_to_markdown(conversation.messages().to_vec(), &session.description)
}
_ => return Err(anyhow::anyhow!("Unsupported format: {}", format)),
};
// Generate the markdown content using the export functionality
let markdown =
export_session_to_markdown(messages.messages().clone(), &session_file_path, None);
// Output the markdown
if let Some(output) = output_path {
fs::write(&output, markdown)
.with_context(|| format!("Failed to write to output file: {}", output.display()))?;
println!("Session exported to {}", output.display());
if let Some(output_path) = output_path {
fs::write(&output_path, output).with_context(|| {
format!("Failed to write to output file: {}", output_path.display())
})?;
println!("Session exported to {}", output_path.display());
} else {
println!("{}", markdown);
println!("{}", output);
}
Ok(())
}
/// Convert a list of messages to markdown format for session export
///
/// This function handles the formatting of a complete session including headers,
/// message organization, and proper tool request/response pairing.
fn export_session_to_markdown(
messages: Vec<goose::conversation::message::Message>,
session_file: &Path,
session_name_override: Option<&str>,
session_name: &String,
) -> String {
let mut markdown_output = String::new();
let session_name = session_name_override.unwrap_or_else(|| {
session_file
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("Unnamed Session")
});
markdown_output.push_str(&format!("# Session Export: {}\n\n", session_name));
if messages.is_empty() {
@@ -293,15 +263,8 @@ fn export_session_to_markdown(
/// Prompt the user to interactively select a session
///
/// Shows a list of available sessions and lets the user select one
pub fn prompt_interactive_session_selection() -> Result<session::Identifier> {
// Get sessions sorted by modification date (newest first)
let sessions = match get_valid_sorted_sessions(SortOrder::Descending) {
Ok(sessions) => sessions,
Err(e) => {
tracing::error!("Failed to list sessions: {:?}", e);
return Err(anyhow::anyhow!("Failed to list sessions"));
}
};
pub async fn prompt_interactive_session_selection() -> Result<String> {
let sessions = SessionManager::list_sessions().await?;
if sessions.is_empty() {
return Err(anyhow::anyhow!("No sessions found"));
@@ -311,19 +274,17 @@ pub fn prompt_interactive_session_selection() -> Result<session::Identifier> {
let mut selector = select("Select a session to export:");
// Map to display text
let display_map: std::collections::HashMap<String, SessionInfo> = sessions
let display_map: std::collections::HashMap<String, Session> = sessions
.iter()
.map(|s| {
let desc = if s.metadata.description.is_empty() {
let desc = if s.description.is_empty() {
"(no description)"
} else {
&s.metadata.description
&s.description
};
let truncated_desc = safe_truncate(desc, TRUNCATED_DESC_LENGTH);
// Truncate description if too long
let truncated_desc = safe_truncate(desc, 40);
let display_text = format!("{} - {} ({})", s.modified, truncated_desc, s.id);
let display_text = format!("{} - {} ({})", s.updated_at, truncated_desc, s.id);
(display_text, s.clone())
})
.collect();
@@ -346,7 +307,7 @@ pub fn prompt_interactive_session_selection() -> Result<session::Identifier> {
// Retrieve the selected session
if let Some(session) = display_map.get(&selected_display_text) {
Ok(goose::session::Identifier::Name(session.id.clone()))
Ok(session.id.clone())
} else {
Err(anyhow::anyhow!("Invalid selection"))
}
+54 -183
View File
@@ -8,24 +8,25 @@ use axum::{
routing::get,
Json, Router,
};
use goose::session::SessionManager;
use webbrowser;
use futures::{sink::SinkExt, stream::StreamExt};
use goose::agents::{Agent, AgentEvent};
use goose::conversation::message::Message as GooseMessage;
use goose::conversation::Conversation;
use goose::session;
use axum::response::Redirect;
use serde::{Deserialize, Serialize};
use std::{net::SocketAddr, sync::Arc};
use tokio::sync::{Mutex, RwLock};
use tower_http::cors::{Any, CorsLayer};
use tracing::error;
type SessionStore = Arc<RwLock<std::collections::HashMap<String, Arc<Mutex<Conversation>>>>>;
type CancellationStore = Arc<RwLock<std::collections::HashMap<String, tokio::task::AbortHandle>>>;
#[derive(Clone)]
struct AppState {
agent: Arc<Agent>,
sessions: SessionStore,
cancellations: CancellationStore,
}
@@ -123,7 +124,6 @@ pub async fn handle_web(port: u16, host: String, open: bool) -> Result<()> {
let state = AppState {
agent: Arc::new(agent),
sessions: Arc::new(RwLock::new(std::collections::HashMap::new())),
cancellations: Arc::new(RwLock::new(std::collections::HashMap::new())),
};
@@ -169,8 +169,15 @@ pub async fn handle_web(port: u16, host: String, open: bool) -> Result<()> {
Ok(())
}
async fn serve_index() -> Html<&'static str> {
Html(include_str!("../../static/index.html"))
async fn serve_index() -> Result<Redirect, (http::StatusCode, String)> {
let session = SessionManager::create_session(
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")),
"Web session".to_string(),
)
.await
.map_err(|err| (http::StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?;
Ok(Redirect::to(&format!("/session/{}", session.id)))
}
async fn serve_session(
@@ -222,23 +229,19 @@ async fn health_check() -> Json<serde_json::Value> {
}
async fn list_sessions() -> Json<serde_json::Value> {
match session::list_sessions() {
match SessionManager::list_sessions().await {
Ok(sessions) => {
let session_info: Vec<serde_json::Value> = sessions
.into_iter()
.filter_map(|(name, path)| {
session::read_metadata(&path).ok().map(|metadata| {
serde_json::json!({
"name": name,
"path": path,
"description": metadata.description,
"message_count": metadata.message_count,
"working_dir": metadata.working_dir
})
})
})
.collect();
let mut session_info = Vec::new();
for session in sessions {
session_info.push(serde_json::json!({
"name": session.id,
"path": session.id,
"description": session.description,
"message_count": session.message_count,
"working_dir": session.working_dir
}));
}
Json(serde_json::json!({
"sessions": session_info
}))
@@ -251,30 +254,14 @@ async fn list_sessions() -> Json<serde_json::Value> {
async fn get_session(
axum::extract::Path(session_id): axum::extract::Path<String>,
) -> Json<serde_json::Value> {
let session_file = match session::get_path(session::Identifier::Name(session_id)) {
Ok(path) => path,
Err(e) => {
return Json(serde_json::json!({
"error": format!("Invalid session ID: {}", e)
}));
}
};
let error_response = |e: Box<dyn std::error::Error>| {
Json(serde_json::json!({
match SessionManager::get_session(&session_id, true).await {
Ok(session) => Json(serde_json::json!({
"metadata": session,
"messages": session.conversation.unwrap_or_default().messages()
})),
Err(e) => Json(serde_json::json!({
"error": e.to_string()
}))
};
match session::read_messages(&session_file) {
Ok(messages) => match session::read_metadata(&session_file) {
Ok(metadata) => Json(serde_json::json!({
"metadata": metadata,
"messages": messages
})),
Err(e) => error_response(e.into()),
},
Err(e) => error_response(e.into()),
})),
}
}
@@ -299,46 +286,14 @@ async fn handle_socket(socket: WebSocket, state: AppState) {
session_id,
..
}) => {
// Get session file path from session_id
let session_file = match session::get_path(session::Identifier::Name(
session_id.clone(),
)) {
Ok(path) => path,
Err(e) => {
tracing::error!("Failed to get session path: {}", e);
continue;
}
};
// Get or create session in memory (for fast access during processing)
let session_messages = {
let sessions = state.sessions.read().await;
if let Some(session) = sessions.get(&session_id) {
session.clone()
} else {
drop(sessions);
let mut sessions = state.sessions.write().await;
// Load existing messages from JSONL file if it exists
let existing_messages =
session::read_messages(&session_file).unwrap_or_default();
let new_session = Arc::new(Mutex::new(existing_messages));
sessions.insert(session_id.clone(), new_session.clone());
new_session
}
};
// Clone sender for async processing
let sender_clone = sender.clone();
let agent = state.agent.clone();
let session_id_clone = session_id.clone();
// Process message in a separate task to allow streaming
let task_handle = tokio::spawn(async move {
let result = process_message_streaming(
&agent,
session_messages,
session_file,
session_id_clone,
content,
sender_clone,
)
@@ -349,25 +304,21 @@ async fn handle_socket(socket: WebSocket, state: AppState) {
}
});
// Store the abort handle
{
let mut cancellations = state.cancellations.write().await;
cancellations
.insert(session_id.clone(), task_handle.abort_handle());
}
// Wait for task completion and handle abort
// Handle task completion and cleanup
let sender_for_abort = sender.clone();
let session_id_for_cleanup = session_id.clone();
let cancellations_for_cleanup = state.cancellations.clone();
tokio::spawn(async move {
match task_handle.await {
Ok(_) => {
// Task completed normally
}
Ok(_) => {}
Err(e) if e.is_cancelled() => {
// Task was aborted
let mut sender = sender_for_abort.lock().await;
let _ = sender
.send(Message::Text(
@@ -387,11 +338,8 @@ async fn handle_socket(socket: WebSocket, state: AppState) {
}
}
// Clean up cancellation token
{
let mut cancellations = cancellations_for_cleanup.write().await;
cancellations.remove(&session_id_for_cleanup);
}
let mut cancellations = cancellations_for_cleanup.write().await;
cancellations.remove(&session_id_for_cleanup);
});
}
Ok(WebSocketMessage::Cancel { session_id }) => {
@@ -436,27 +384,16 @@ async fn handle_socket(socket: WebSocket, state: AppState) {
async fn process_message_streaming(
agent: &Agent,
session_messages: Arc<Mutex<Conversation>>,
session_file: std::path::PathBuf,
session_id: String,
content: String,
sender: Arc<Mutex<futures::stream::SplitSink<WebSocket, Message>>>,
) -> Result<()> {
use futures::StreamExt;
use goose::agents::SessionConfig;
use goose::conversation::message::MessageContent;
use goose::session;
// Create a user message
let user_message = GooseMessage::user().with_text(content.clone());
// Messages will be auto-compacted in agent.reply() if needed
let messages: Conversation = {
let mut session_msgs = session_messages.lock().await;
session_msgs.push(user_message.clone());
session_msgs.clone()
};
// Persist messages to JSONL file with provider for automatic description generation
let provider = agent.provider().await;
if provider.is_err() {
let error_msg = "I'm not properly configured yet. Please configure a provider through the CLI first using `goose configure`.".to_string();
@@ -475,19 +412,13 @@ async fn process_message_streaming(
return Ok(());
}
let provider = provider.unwrap();
let working_dir = Some(std::env::current_dir()?);
session::persist_messages(
&session_file,
&messages,
Some(provider.clone()),
working_dir.clone(),
)
.await?;
let session = SessionManager::get_session(&session_id, true).await?;
let mut messages = session.conversation.unwrap_or_default();
messages.push(user_message);
let session_config = SessionConfig {
id: session::Identifier::Path(session_file.clone()),
working_dir: std::env::current_dir()?,
id: session.id.clone(),
working_dir: session.working_dir,
schedule_id: None,
execution_mode: None,
max_turns: None,
@@ -502,29 +433,11 @@ async fn process_message_streaming(
while let Some(result) = stream.next().await {
match result {
Ok(AgentEvent::Message(message)) => {
// Add message to our session
{
let mut session_msgs = session_messages.lock().await;
session_msgs.push(message.clone());
}
SessionManager::add_message(&session_id, &message).await?;
// Persist messages to JSONL file (no provider needed for assistant messages)
let current_messages = {
let session_msgs = session_messages.lock().await;
session_msgs.clone()
};
session::persist_messages(
&session_file,
&current_messages,
None,
working_dir.clone(),
)
.await?;
// Handle different message content types
for content in &message.content {
match content {
MessageContent::Text(text) => {
// Send the text response
let mut sender = sender.lock().await;
let _ = sender
.send(Message::Text(
@@ -539,7 +452,6 @@ async fn process_message_streaming(
.await;
}
MessageContent::ToolRequest(req) => {
// Send tool request notification
let mut sender = sender.lock().await;
if let Ok(tool_call) = &req.tool_call {
let _ = sender
@@ -557,13 +469,8 @@ async fn process_message_streaming(
.await;
}
}
MessageContent::ToolResponse(_resp) => {
// Tool responses are already included in the complete message stream
// and will be persisted to session history. No need to send separate
// WebSocket messages as this would cause duplicates.
}
MessageContent::ToolResponse(_resp) => {}
MessageContent::ToolConfirmationRequest(confirmation) => {
// Send tool confirmation request
let mut sender = sender.lock().await;
let _ = sender
.send(Message::Text(
@@ -580,8 +487,6 @@ async fn process_message_streaming(
))
.await;
// For now, auto-approve in web mode
// TODO: Implement proper confirmation UI
agent.handle_confirmation(
confirmation.id.clone(),
goose::permission::PermissionConfirmation {
@@ -591,7 +496,6 @@ async fn process_message_streaming(
).await;
}
MessageContent::Thinking(thinking) => {
// Send thinking indicator
let mut sender = sender.lock().await;
let _ = sender
.send(Message::Text(
@@ -604,7 +508,6 @@ async fn process_message_streaming(
.await;
}
MessageContent::ContextLengthExceeded(msg) => {
// Send context exceeded notification
let mut sender = sender.lock().await;
let _ = sender
.send(Message::Text(
@@ -618,55 +521,27 @@ async fn process_message_streaming(
))
.await;
// For now, auto-summarize in web mode
// TODO: Implement proper UI for context handling
let (summarized_messages, _, _) =
agent.summarize_context(messages.messages()).await?;
{
let mut session_msgs = session_messages.lock().await;
*session_msgs = summarized_messages;
}
}
_ => {
// Handle other message types as needed
SessionManager::replace_conversation(
&session_id,
&summarized_messages,
)
.await?;
}
_ => {}
}
}
}
Ok(AgentEvent::HistoryReplaced(new_messages)) => {
// Replace the session's message history with the compacted messages
{
let mut session_msgs = session_messages.lock().await;
*session_msgs = Conversation::new_unvalidated(new_messages);
}
// Persist the updated messages to the JSONL file
let current_messages = {
let session_msgs = session_messages.lock().await;
session_msgs.clone()
};
if let Err(e) = session::persist_messages(
&session_file,
&current_messages,
None, // No provider needed for persisting
working_dir.clone(),
)
.await
{
error!("Failed to persist compacted messages: {}", e);
}
Ok(AgentEvent::HistoryReplaced(_new_messages)) => {
tracing::info!("History replaced, compacting happened in reply");
}
Ok(AgentEvent::McpNotification(_notification)) => {
// Handle MCP notifications if needed
// For now, we'll just log them
tracing::info!("Received MCP notification in web interface");
}
Ok(AgentEvent::ModelChange { model, mode }) => {
// Log model change
tracing::info!("Model changed to {} in {} mode", model, mode);
}
Err(e) => {
error!("Error in message stream: {}", e);
let mut sender = sender.lock().await;
@@ -699,7 +574,6 @@ async fn process_message_streaming(
}
}
// Send completion message
let mut sender = sender.lock().await;
let _ = sender
.send(Message::Text(
@@ -713,6 +587,3 @@ async fn process_message_streaming(
Ok(())
}
// Add webbrowser dependency for opening browser
use webbrowser;
+1 -1
View File
@@ -10,7 +10,7 @@ pub mod session;
pub mod signal;
// Re-export commonly used types
pub use session::Session;
pub use session::CliSession;
pub static APP_STRATEGY: Lazy<AppStrategyArgs> = Lazy::new(|| AppStrategyArgs {
top_level_domain: "Block".to_string(),
@@ -4,7 +4,7 @@ use goose::conversation::Conversation;
use crate::scenario_tests::message_generator::MessageGenerator;
use crate::scenario_tests::mock_client::weather_client;
use crate::scenario_tests::provider_configs::{get_provider_configs, ProviderConfig};
use crate::session::Session;
use crate::session::CliSession;
use anyhow::Result;
use goose::agents::Agent;
use goose::model::ModelConfig;
@@ -218,7 +218,7 @@ where
.update_provider(provider_arc as Arc<dyn goose::providers::base::Provider>)
.await?;
let mut session = Session::new(agent, None, false, None, None, None, None);
let mut session = CliSession::new(agent, None, false, None, None, None, None);
let mut error = None;
for message in &messages {
+49 -70
View File
@@ -1,28 +1,27 @@
use super::output;
use super::CliSession;
use console::style;
use goose::agents::types::RetryConfig;
use goose::agents::Agent;
use goose::config::{Config, ExtensionConfig, ExtensionConfigManager};
use goose::providers::create;
use goose::recipe::{Response, SubRecipe};
use goose::session;
use goose::session::Identifier;
use goose::session::SessionManager;
use rustyline::EditMode;
use std::collections::HashSet;
use std::process;
use std::sync::Arc;
use tokio::task::JoinSet;
use super::output;
use super::Session;
/// Configuration for building a new Goose session
///
/// This struct contains all the parameters needed to create a new session,
/// including session identification, extension configuration, and debug settings.
#[derive(Default, Clone, Debug)]
pub struct SessionBuilderConfig {
/// Optional identifier for the session (name or path)
pub identifier: Option<Identifier>,
/// Optional identifier for the session
pub session_id: Option<String>,
/// Whether to resume an existing session
pub resume: bool,
/// Whether to run without a session file
@@ -129,20 +128,8 @@ async fn offer_extension_debugging_help(
}
}
// Create a temporary session file for this debugging session
let temp_session_file =
std::env::temp_dir().join(format!("goose_debug_extension_{}.jsonl", extension_name));
// Create the debugging session
let mut debug_session = Session::new(
debug_agent,
Some(temp_session_file.clone()),
false,
None,
None,
None,
None,
);
let mut debug_session = CliSession::new(debug_agent, None, false, None, None, None, None);
// Process the debugging request
println!("{}", style("Analyzing the extension failure...").yellow());
@@ -160,10 +147,6 @@ async fn offer_extension_debugging_help(
);
}
}
// Clean up the temporary session file
let _ = std::fs::remove_file(temp_session_file);
Ok(())
}
@@ -174,7 +157,7 @@ pub struct SessionSettings {
pub temperature: Option<f32>,
}
pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
// Load config and get provider/model
let config = Config::global();
@@ -257,68 +240,64 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
});
// Handle session file resolution and resuming
let session_file: Option<std::path::PathBuf> = if session_config.no_session {
let session_id: Option<String> = if session_config.no_session {
None
} else if session_config.resume {
if let Some(identifier) = session_config.identifier {
let session_file = match session::get_path(identifier) {
Err(e) => {
output::render_error(&format!("Invalid session identifier: {}", e));
if let Some(session_id) = session_config.session_id {
match SessionManager::get_session(&session_id, false).await {
Ok(_) => Some(session_id),
Err(_) => {
output::render_error(&format!(
"Cannot resume session {} - no such session exists",
style(&session_id).cyan()
));
process::exit(1);
}
Ok(path) => path,
};
if !session_file.exists() {
output::render_error(&format!(
"Cannot resume session {} - no such session exists",
style(session_file.display()).cyan()
));
process::exit(1);
}
Some(session_file)
} else {
// Try to resume most recent session
match session::get_most_recent_session() {
Ok(file) => Some(file),
match SessionManager::list_sessions().await {
Ok(sessions) => {
if sessions.is_empty() {
output::render_error("Cannot resume - no previous sessions found");
process::exit(1);
}
Some(sessions[0].id.clone())
}
Err(_) => {
output::render_error("Cannot resume - no previous sessions found");
process::exit(1);
}
}
}
} else if let Some(session_id) = session_config.session_id {
Some(session_id)
} else {
// Create new session with provided name/path or generated name
let id = match session_config.identifier {
Some(identifier) => identifier,
None => Identifier::Name(session::generate_session_id()),
};
// Just get the path - file will be created when needed
match session::get_path(id) {
Ok(path) => Some(path),
Err(e) => {
output::render_error(&format!("Failed to create session path: {}", e));
process::exit(1);
}
}
let session = SessionManager::create_session(
std::env::current_dir().unwrap(),
"CLI Session".to_string(),
)
.await
.unwrap();
Some(session.id)
};
if session_config.resume {
if let Some(session_file) = session_file.as_ref() {
// Read the session metadata
let metadata = session::read_metadata(session_file).unwrap_or_else(|e| {
output::render_error(&format!("Failed to read session metadata: {}", e));
process::exit(1);
});
if let Some(session_id) = session_id.as_ref() {
// Read the session metadata from database
let metadata = SessionManager::get_session(session_id, false)
.await
.unwrap_or_else(|e| {
output::render_error(&format!("Failed to read session metadata: {}", e));
process::exit(1);
});
let current_workdir =
std::env::current_dir().expect("Failed to get current working directory");
if current_workdir != metadata.working_dir {
// Ask user if they want to change the working directory
let change_workdir = cliclack::confirm(format!("{} The original working directory of this session was set to {}. Your current directory is {}. Do you want to switch back to the original working directory?", style("WARNING:").yellow(), style(metadata.working_dir.display()).cyan(), style(current_workdir.display()).cyan()))
.initial_value(true)
.interact().expect("Failed to get user input");
.initial_value(true)
.interact().expect("Failed to get user input");
if change_workdir {
if !metadata.working_dir.exists() {
@@ -417,9 +396,9 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
});
// Create new session
let mut session = Session::new(
let mut session = CliSession::new(
Arc::try_unwrap(agent_ptr).unwrap_or_else(|_| panic!("There should be no more references")),
session_file.clone(),
session_id.clone(),
session_config.debug,
session_config.scheduled_job_id.clone(),
session_config.max_turns,
@@ -586,7 +565,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
session_config.resume,
&provider_name,
&model_name,
&session_file,
&session_id,
Some(&provider_for_display),
);
}
@@ -600,7 +579,7 @@ mod tests {
#[test]
fn test_session_builder_config_creation() {
let config = SessionBuilderConfig {
identifier: Some(Identifier::Name("test".to_string())),
session_id: Some("test".to_string()),
resume: false,
no_session: false,
extensions: vec!["echo test".to_string()],
@@ -639,7 +618,7 @@ mod tests {
fn test_session_builder_config_default() {
let config = SessionBuilderConfig::default();
assert!(config.identifier.is_none());
assert!(config.session_id.is_none());
assert!(!config.resume);
assert!(!config.no_session);
assert!(config.extensions.is_empty());
+83 -248
View File
@@ -21,7 +21,6 @@ use goose::permission::permission_confirmation::PrincipalType;
use goose::permission::Permission;
use goose::permission::PermissionConfirmation;
use goose::providers::base::Provider;
pub use goose::session::Identifier;
use goose::utils::safe_truncate;
use anyhow::{Context, Result};
@@ -39,6 +38,7 @@ use rmcp::model::ServerNotification;
use rmcp::model::{ErrorCode, ErrorData};
use goose::conversation::message::{Message, MessageContent};
use goose::session::SessionManager;
use rand::{distributions::Alphanumeric, Rng};
use rustyline::EditMode;
use serde_json::Value;
@@ -54,13 +54,12 @@ pub enum RunMode {
Plan,
}
pub struct Session {
pub struct CliSession {
agent: Agent,
messages: Conversation,
session_file: Option<PathBuf>,
// Cache for completion data - using std::sync for thread safety without async
session_id: Option<String>,
completion_cache: Arc<std::sync::RwLock<CompletionCache>>,
debug: bool, // New field for debug mode
debug: bool,
run_mode: RunMode,
scheduled_job_id: Option<String>, // ID of the scheduled job that triggered this session
max_turns: Option<u32>,
@@ -111,8 +110,6 @@ pub async fn classify_planner_response(
)
.await?;
// println!("classify_planner_response: {result:?}\n"); // TODO: remove
let predicted = result.as_concat_text();
if predicted.to_lowercase().contains("plan") {
Ok(PlannerResponseType::Plan)
@@ -121,30 +118,33 @@ pub async fn classify_planner_response(
}
}
impl Session {
impl CliSession {
pub fn new(
agent: Agent,
session_file: Option<PathBuf>,
session_id: Option<String>,
debug: bool,
scheduled_job_id: Option<String>,
max_turns: Option<u32>,
edit_mode: Option<EditMode>,
retry_config: Option<RetryConfig>,
) -> Self {
let messages = if let Some(session_file) = &session_file {
session::read_messages(session_file).unwrap_or_else(|e| {
eprintln!("Warning: Failed to load message history: {}", e);
Conversation::new_unvalidated(Vec::new())
let messages = if let Some(session_id) = &session_id {
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async {
SessionManager::get_session(session_id, true)
.await
.map(|session| session.conversation.unwrap_or_default())
.unwrap()
})
})
} else {
// Don't try to read messages if we're not saving sessions
Conversation::new_unvalidated(Vec::new())
};
Session {
CliSession {
agent,
messages,
session_file,
session_id,
completion_cache: Arc::new(std::sync::RwLock::new(CompletionCache::new())),
debug,
run_mode: RunMode::Normal,
@@ -155,13 +155,15 @@ impl Session {
}
}
/// Helper function to summarize context messages
pub fn session_id(&self) -> Option<&String> {
self.session_id.as_ref()
}
async fn summarize_context_messages(
messages: &mut Conversation,
agent: &Agent,
message_suffix: &str,
) -> Result<()> {
// Summarize messages to fit within context length
let (summarized_messages, _, _) = agent.summarize_context(messages.messages()).await?;
let msg = format!("Context maxed out\n{}\n{}", "-".repeat(50), message_suffix);
output::render_text(&msg, Some(Color::Yellow), true);
@@ -179,7 +181,6 @@ impl Session {
let mut parts: Vec<&str> = extension_command.split_whitespace().collect();
let mut envs = HashMap::new();
// Parse environment variables (format: KEY=value)
while let Some(part) = parts.first() {
if !part.contains('=') {
break;
@@ -194,7 +195,6 @@ impl Session {
}
let cmd = parts.remove(0).to_string();
// Generate a random name for the ephemeral extension
let name: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(8)
@@ -374,46 +374,10 @@ impl Session {
cancel_token: CancellationToken,
) -> Result<()> {
let cancel_token = cancel_token.clone();
let message_text = message.as_concat_text();
// TODO(Douwe): Make sure we generate the description here still:
self.push_message(message);
// Get the provider from the agent for description generation
let provider = self.agent.provider().await?;
// Persist messages with provider for automatic description generation
if let Some(session_file) = &self.session_file {
let working_dir = Some(
std::env::current_dir().expect("failed to get current session working directory"),
);
session::persist_messages_with_schedule_id(
session_file,
&self.messages,
Some(provider),
self.scheduled_job_id.clone(),
working_dir,
)
.await?;
}
// Track the current directory and last instruction in projects.json
let session_id = self
.session_file
.as_ref()
.and_then(|p| p.file_stem())
.and_then(|s| s.to_str())
.map(|s| s.to_string());
if let Err(e) = crate::project_tracker::update_project_tracker(
Some(&message_text),
session_id.as_deref(),
) {
eprintln!(
"Warning: Failed to update project tracker with instruction: {}",
e
);
}
self.process_agent_response(false, cancel_token).await?;
Ok(())
}
@@ -493,35 +457,14 @@ impl Session {
self.push_message(Message::user().with_text(&content));
// Track the current directory and last instruction in projects.json
let session_id = self
.session_file
.as_ref()
.and_then(|p| p.file_stem())
.and_then(|s| s.to_str())
.map(|s| s.to_string());
if let Err(e) = crate::project_tracker::update_project_tracker(
Some(&content),
session_id.as_deref(),
self.session_id.as_deref(),
) {
eprintln!("Warning: Failed to update project tracker with instruction: {}", e);
}
let provider = self.agent.provider().await?;
// Persist messages with provider for automatic description generation
if let Some(session_file) = &self.session_file {
let working_dir = Some(std::env::current_dir().unwrap_or_default());
session::persist_messages_with_schedule_id(
session_file,
&self.messages,
Some(provider),
self.scheduled_job_id.clone(),
working_dir,
)
.await?;
}
let _provider = self.agent.provider().await?;
output::show_thinking();
let start_time = Instant::now();
@@ -659,16 +602,25 @@ impl Session {
input::InputResult::Clear => {
save_history(&mut editor);
if let Some(session_id) = &self.session_id {
if let Err(e) = SessionManager::replace_conversation(
session_id,
&Conversation::default(),
)
.await
{
output::render_error(&format!("Failed to clear session: {}", e));
continue;
}
}
self.messages.clear();
tracing::info!("Chat context cleared by user.");
output::render_message(
&Message::assistant().with_text("Chat context cleared."),
self.debug,
);
if let Some(file) = self.session_file.as_ref().filter(|f| f.exists()) {
std::fs::remove_file(file)?;
std::fs::File::create(file)?;
}
continue;
}
input::InputResult::PromptCommand(opts) => {
@@ -729,44 +681,30 @@ impl Session {
output::show_thinking();
// Get the provider for summarization
let provider = self.agent.provider().await?;
let _provider = self.agent.provider().await?;
// Call the summarize_context method which uses the summarize_messages function
// Call the summarize_context method
let (summarized_messages, _token_counts, summarization_usage) = self
.agent
.summarize_context(self.messages.messages())
.await?;
// Update the session messages with the summarized ones
self.messages = summarized_messages;
self.messages = summarized_messages.clone();
// Persist the summarized messages and update session metadata with new token counts
if let Some(session_file) = &self.session_file {
let working_dir = std::env::current_dir().ok();
session::persist_messages_with_schedule_id(
session_file,
&self.messages,
Some(provider),
self.scheduled_job_id.clone(),
working_dir,
)
.await?;
// Persist the summarized messages and update session metadata
if let Some(session_id) = &self.session_id {
// Replace all messages with the summarized version
SessionManager::replace_conversation(session_id, &summarized_messages)
.await?;
// Update session metadata with the new token counts from summarization
if let Some(usage) = summarization_usage {
let session_file_path = session::storage::get_path(
session::storage::Identifier::Path(session_file.to_path_buf()),
)?;
let mut metadata =
session::storage::read_metadata(&session_file_path)?;
let session =
SessionManager::get_session(session_id, false).await?;
// Update token counts with the summarization usage
// Use output tokens as total since that's what's actually in the context going forward
let summary_tokens = usage.usage.output_tokens.unwrap_or(0);
metadata.total_tokens = Some(summary_tokens);
metadata.input_tokens = None; // Clear input tokens since we now have a summary
metadata.output_tokens = Some(summary_tokens);
metadata.message_count = self.messages.len();
// Update accumulated tokens (add the summarization cost)
let accumulate = |a: Option<i32>, b: Option<i32>| -> Option<i32> {
@@ -775,20 +713,28 @@ impl Session {
_ => a.or(b),
}
};
metadata.accumulated_total_tokens = accumulate(
metadata.accumulated_total_tokens,
let accumulated_total = accumulate(
session.accumulated_total_tokens,
usage.usage.total_tokens,
);
metadata.accumulated_input_tokens = accumulate(
metadata.accumulated_input_tokens,
let accumulated_input = accumulate(
session.accumulated_input_tokens,
usage.usage.input_tokens,
);
metadata.accumulated_output_tokens = accumulate(
metadata.accumulated_output_tokens,
let accumulated_output = accumulate(
session.accumulated_output_tokens,
usage.usage.output_tokens,
);
session::storage::update_metadata(&session_file_path, &metadata)
SessionManager::update_session(session_id)
.total_tokens(Some(summary_tokens))
.input_tokens(None)
.output_tokens(Some(summary_tokens))
.accumulated_total_tokens(accumulated_total)
.accumulated_input_tokens(accumulated_input)
.accumulated_output_tokens(accumulated_output)
.apply()
.await?;
}
}
@@ -808,19 +754,10 @@ impl Session {
} else {
println!("{}", console::style("Summarization cancelled.").yellow());
}
continue;
}
}
}
println!(
"\nClosing session.{}",
self.session_file
.as_ref()
.map(|p| format!(" Recorded to {}", p.display()))
.unwrap_or_default()
);
Ok(())
}
@@ -919,16 +856,13 @@ impl Session {
) -> Result<()> {
let cancel_token_clone = cancel_token.clone();
let session_config = self.session_file.as_ref().map(|s| {
let session_id = session::Identifier::Path(s.clone());
SessionConfig {
id: session_id.clone(),
working_dir: std::env::current_dir().unwrap_or_default(),
schedule_id: self.scheduled_job_id.clone(),
execution_mode: None,
max_turns: self.max_turns,
retry_config: self.retry_config.clone(),
}
let session_config = self.session_id.as_ref().map(|session_id| SessionConfig {
id: session_id.clone(),
working_dir: std::env::current_dir().unwrap_or_default(),
schedule_id: self.scheduled_job_id.clone(),
execution_mode: None,
max_turns: self.max_turns,
retry_config: self.retry_config.clone(),
});
let mut stream = self
.agent
@@ -998,17 +932,6 @@ impl Session {
Err(ErrorData { code: ErrorCode::INVALID_REQUEST, message: std::borrow::Cow::from("Tool call cancelled by user".to_string()), data: None })
));
self.messages.push(response_message);
if let Some(session_file) = &self.session_file {
let working_dir = std::env::current_dir().ok();
session::persist_messages_with_schedule_id(
session_file,
&self.messages,
None,
self.scheduled_job_id.clone(),
working_dir,
)
.await?;
}
cancel_token_clone.cancel();
drop(stream);
break;
@@ -1140,22 +1063,8 @@ impl Session {
);
}
}
self.messages.push(message.clone());
// No need to update description on assistant messages
if let Some(session_file) = &self.session_file {
let working_dir = std::env::current_dir().ok();
session::persist_messages_with_schedule_id(
session_file,
&self.messages,
None,
self.scheduled_job_id.clone(),
working_dir,
)
.await?;
}
if interactive {output::hide_thinking()};
let _ = progress_bars.hide();
output::render_message(&message, self.debug);
@@ -1267,26 +1176,10 @@ impl Session {
_ => (),
}
}
Some(Ok(AgentEvent::HistoryReplaced(new_messages))) => {
// Replace the session's message history with the compacted messages
self.messages = Conversation::new_unvalidated(new_messages);
// Persist the updated messages to the session file
if let Some(session_file) = &self.session_file {
let provider = self.agent.provider().await.ok();
let working_dir = std::env::current_dir().ok();
if let Err(e) = session::persist_messages_with_schedule_id(
session_file,
&self.messages,
provider,
self.scheduled_job_id.clone(),
working_dir,
).await {
eprintln!("Failed to persist compacted messages: {}", e);
}
}
}
Some(Ok(AgentEvent::ModelChange { model, mode })) => {
Some(Ok(AgentEvent::HistoryReplaced(new_messages))) => {
self.messages = Conversation::new_unvalidated(new_messages.clone());
}
Some(Ok(AgentEvent::ModelChange { model, mode })) => {
// Log model change if in debug mode
if self.debug {
eprintln!("Model changed to {} in {} mode", model, mode);
@@ -1308,20 +1201,8 @@ impl Session {
// Try auto-compaction first - keep the stream alive!
if let Ok(compact_result) = goose::context_mgmt::auto_compact::perform_compaction(&self.agent, self.messages.messages()).await {
self.messages = compact_result.messages;
// Persist the compacted messages
if let Some(session_file) = &self.session_file {
let provider = self.agent.provider().await.ok();
let working_dir = std::env::current_dir().ok();
if let Err(e) = session::persist_messages_with_schedule_id(
session_file,
&self.messages,
provider,
self.scheduled_job_id.clone(),
working_dir,
).await {
eprintln!("Failed to persist compacted messages: {}", e);
}
if let Some(session_id) = &self.session_id {
SessionManager::replace_conversation(session_id, &self.messages).await?;
}
output::render_text(
@@ -1465,40 +1346,13 @@ impl Session {
}),
));
}
// TODO(Douwe): update also db
self.push_message(response_message);
// No need for description update here
if let Some(session_file) = &self.session_file {
let working_dir = std::env::current_dir().ok();
session::persist_messages_with_schedule_id(
session_file,
&self.messages,
None,
self.scheduled_job_id.clone(),
working_dir,
)
.await?;
}
let prompt = format!(
"The existing call to {} was interrupted. How would you like to proceed?",
last_tool_name
);
self.push_message(Message::assistant().with_text(&prompt));
// No need for description update here
if let Some(session_file) = &self.session_file {
let working_dir = std::env::current_dir().ok();
session::persist_messages_with_schedule_id(
session_file,
&self.messages,
None,
self.scheduled_job_id.clone(),
working_dir,
)
.await?;
}
output::render_message(&Message::assistant().with_text(&prompt), self.debug);
} else {
// An interruption occurred outside of a tool request-response.
@@ -1509,20 +1363,6 @@ impl Session {
// Interruption occurred after a tool had completed but not assistant reply
let prompt = "The tool calling loop was interrupted. How would you like to proceed?";
self.push_message(Message::assistant().with_text(prompt));
// No need for description update here
if let Some(session_file) = &self.session_file {
let working_dir = std::env::current_dir().ok();
session::persist_messages_with_schedule_id(
session_file,
&self.messages,
None,
self.scheduled_job_id.clone(),
working_dir,
)
.await?;
}
output::render_message(
&Message::assistant().with_text(prompt),
self.debug,
@@ -1545,10 +1385,6 @@ impl Session {
Ok(())
}
pub fn session_file(&self) -> Option<PathBuf> {
self.session_file.clone()
}
/// Update the completion cache with fresh data
/// This should be called before the interactive session starts
pub async fn update_completion_cache(&mut self) -> Result<()> {
@@ -1619,17 +1455,16 @@ impl Session {
);
}
pub fn get_metadata(&self) -> Result<session::SessionMetadata> {
if !self.session_file.as_ref().is_some_and(|f| f.exists()) {
return Err(anyhow::anyhow!("Session file does not exist"));
pub async fn get_metadata(&self) -> Result<session::Session> {
match &self.session_id {
Some(id) => SessionManager::get_session(id, false).await,
None => Err(anyhow::anyhow!("No session available")),
}
session::read_metadata(self.session_file.as_ref().unwrap())
}
// Get the session's total token usage
pub fn get_total_token_usage(&self) -> Result<Option<i32>> {
let metadata = self.get_metadata()?;
pub async fn get_total_token_usage(&self) -> Result<Option<i32>> {
let metadata = self.get_metadata().await?;
Ok(metadata.total_tokens)
}
@@ -1661,7 +1496,7 @@ impl Session {
}
}
match self.get_metadata() {
match self.get_metadata().await {
Ok(metadata) => {
let total_tokens = metadata.total_tokens.unwrap_or(0) as usize;
+3 -11
View File
@@ -14,7 +14,7 @@ use serde_json::Value;
use std::cell::RefCell;
use std::collections::HashMap;
use std::io::{Error, IsTerminal, Write};
use std::path::{Path, PathBuf};
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
@@ -685,12 +685,12 @@ pub fn display_session_info(
resume: bool,
provider: &str,
model: &str,
session_file: &Option<PathBuf>,
session_id: &Option<String>,
provider_instance: Option<&Arc<dyn goose::providers::base::Provider>>,
) {
let start_session_msg = if resume {
"resuming session |"
} else if session_file.is_none() {
} else if session_id.is_none() {
"running without session |"
} else {
"starting session |"
@@ -732,14 +732,6 @@ pub fn display_session_info(
);
}
if let Some(session_file) = session_file {
println!(
" {} {}",
style("logging to").dim(),
style(session_file.display()).dim().cyan(),
);
}
println!(
" {} {}",
style("working directory:").dim(),
+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,
))
+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(
@@ -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");
+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"
);
+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) => {
+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
}
}
+13 -13
View File
@@ -31,18 +31,18 @@ thiserror = "1.0"
futures = "0.3"
dirs = "5.0"
reqwest = { version = "0.12.9", features = [
"rustls-tls-native-roots",
"json",
"cookies",
"gzip",
"brotli",
"deflate",
"zstd",
"charset",
"http2",
"stream",
"blocking"
], default-features = false }
"rustls-tls-native-roots",
"json",
"cookies",
"gzip",
"brotli",
"deflate",
"zstd",
"charset",
"http2",
"stream",
"blocking"
], default-features = false }
tokio = { version = "1.43", features = ["full"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
@@ -79,6 +79,7 @@ rand = "0.8.5"
utoipa = { version = "4.1", features = ["chrono"] }
tokio-cron-scheduler = "0.14.0"
urlencoding = "2.1"
sqlx = { version = "0.7", features = ["runtime-tokio-rustls", "sqlite", "chrono", "json"] }
# For Bedrock provider
aws-config = { version = "1.5.16", features = ["behavior-version-latest"] }
@@ -100,7 +101,6 @@ ahash = "0.8"
tokio-util = "0.7.15"
unicode-normalization = "0.1"
arrow = "52.2"
oauth2 = "5.0.0"
[target.'cfg(target_os = "windows")'.dependencies]
+138 -95
View File
@@ -42,8 +42,6 @@ use crate::providers::errors::ProviderError;
use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe};
use crate::scheduler_trait::SchedulerTrait;
use crate::security::security_inspector::SecurityInspector;
use crate::session;
use crate::session::extension_data::ExtensionState;
use crate::tool_inspection::ToolInspectionManager;
use crate::tool_monitor::RepetitionInspector;
use crate::utils::is_token_cancelled;
@@ -55,7 +53,7 @@ use rmcp::model::{
use serde_json::Value;
use tokio::sync::{mpsc, Mutex};
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info, instrument};
use tracing::{debug, error, info, instrument, warn};
use super::final_output_tool::FinalOutputTool;
use super::model_selector::autopilot::AutoPilot;
@@ -66,12 +64,14 @@ use crate::agents::todo_tools::{
todo_read_tool, todo_write_tool, TODO_READ_TOOL_NAME, TODO_WRITE_TOOL_NAME,
};
use crate::conversation::message::{Message, ToolRequest};
use crate::session::extension_data::ExtensionState;
use crate::session::{extension_data, SessionManager};
const DEFAULT_MAX_TURNS: u32 = 1000;
/// Context needed for the reply function
pub struct ReplyContext {
pub messages: Conversation,
pub conversation: Conversation,
pub tools: Vec<Tool>,
pub toolshim_tools: Vec<Tool>,
pub system_prompt: String,
@@ -268,7 +268,7 @@ impl Agent {
.await;
Ok(ReplyContext {
messages: conversation,
conversation,
tools,
toolshim_tools,
system_prompt,
@@ -506,17 +506,12 @@ impl Agent {
)))
} else if tool_call.name == TODO_READ_TOOL_NAME {
// Handle task planner read tool
let session_file_path = if let Some(session_config) = session {
session::storage::get_path(session_config.id.clone()).ok()
} else {
None
};
let todo_content = if let Some(path) = session_file_path {
session::storage::read_metadata(&path)
let todo_content = if let Some(session_config) = session {
SessionManager::get_session(&session_config.id, false)
.await
.ok()
.and_then(|m| {
session::TodoState::from_extension_data(&m.extension_data)
.and_then(|metadata| {
extension_data::TodoState::from_extension_data(&metadata.extension_data)
.map(|state| state.content)
})
.unwrap_or_default()
@@ -551,43 +546,39 @@ impl Agent {
None,
)))
} else if let Some(session_config) = session {
// Update session metadata with new TODO content
match session::storage::get_path(session_config.id.clone()) {
Ok(path) => match session::storage::read_metadata(&path) {
Ok(mut metadata) => {
let todo_state = session::TodoState::new(content);
todo_state
.to_extension_data(&mut metadata.extension_data)
.ok();
let path_clone = path.clone();
let metadata_clone = metadata.clone();
let update_result = tokio::task::spawn(async move {
session::storage::update_metadata(&path_clone, &metadata_clone)
.await
})
.await;
match update_result {
Ok(Ok(_)) => ToolCallResult::from(Ok(vec![Content::text(
format!("Updated ({} chars)", char_count),
)])),
_ => ToolCallResult::from(Err(ErrorData::new(
match SessionManager::get_session(&session_config.id, false).await {
Ok(mut session) => {
let todo_state = extension_data::TodoState::new(content);
if todo_state
.to_extension_data(&mut session.extension_data)
.is_ok()
{
match SessionManager::update_session(&session_config.id)
.extension_data(session.extension_data)
.apply()
.await
{
Ok(_) => ToolCallResult::from(Ok(vec![Content::text(format!(
"Updated ({} chars)",
char_count
))])),
Err(_) => ToolCallResult::from(Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
"Failed to update session metadata".to_string(),
None,
))),
}
} else {
ToolCallResult::from(Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
"Failed to serialize TODO state".to_string(),
None,
)))
}
Err(_) => ToolCallResult::from(Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
"Failed to read session metadata".to_string(),
None,
))),
},
}
Err(_) => ToolCallResult::from(Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
"Failed to get session path".to_string(),
"Failed to read session metadata".to_string(),
None,
))),
}
@@ -911,10 +902,9 @@ impl Agent {
> {
// Try to get session metadata for more accurate token counts
let session_metadata = if let Some(session_config) = session {
match session::storage::get_path(session_config.id.clone()) {
Ok(session_file_path) => session::storage::read_metadata(&session_file_path).ok(),
Err(_) => None,
}
SessionManager::get_session(&session_config.id, false)
.await
.ok()
} else {
None
};
@@ -960,7 +950,7 @@ impl Agent {
cancel_token: Option<CancellationToken>,
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
// Handle auto-compaction before processing
let (messages, compaction_msg, _summarization_usage) = match self
let (conversation, compaction_msg, _summarization_usage) = match self
.handle_auto_compaction(unfixed_conversation.messages(), &session)
.await?
{
@@ -969,7 +959,7 @@ impl Agent {
let context = self
.prepare_reply_context(unfixed_conversation, &session)
.await?;
(context.messages, None, None)
(context.conversation, None, None)
}
};
@@ -977,10 +967,13 @@ impl Agent {
if let Some(compaction_msg) = compaction_msg {
return Ok(Box::pin(async_stream::try_stream! {
yield AgentEvent::Message(Message::assistant().with_summarization_requested(compaction_msg));
yield AgentEvent::HistoryReplaced(messages.messages().clone());
yield AgentEvent::HistoryReplaced(conversation.messages().clone());
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &conversation).await?
}
// Continue with normal reply processing using compacted messages
let mut reply_stream = self.reply_internal(messages, session, cancel_token).await?;
let mut reply_stream = self.reply_internal(conversation, session, cancel_token).await?;
while let Some(event) = reply_stream.next().await {
yield event?;
}
@@ -988,19 +981,20 @@ impl Agent {
}
// No compaction needed, proceed with normal processing
self.reply_internal(messages, session, cancel_token).await
self.reply_internal(conversation, session, cancel_token)
.await
}
/// Main reply method that handles the actual agent processing
async fn reply_internal(
&self,
messages: Conversation,
conversation: Conversation,
session: Option<SessionConfig>,
cancel_token: Option<CancellationToken>,
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
let context = self.prepare_reply_context(messages, &session).await?;
let context = self.prepare_reply_context(conversation, &session).await?;
let ReplyContext {
mut messages,
mut conversation,
mut tools,
mut toolshim_tools,
mut system_prompt,
@@ -1011,12 +1005,52 @@ impl Agent {
let reply_span = tracing::Span::current();
self.reset_retry_attempts().await;
if let Some(content) = messages
.last()
.and_then(|msg| msg.content.first())
.and_then(|c| c.as_text())
{
debug!("user_message" = &content);
// This will need further refactoring. In the ideal world we pass the new message into
// reply and load the existing conversation. Until we get to that point, fetch the conversation
// so far and append the last (user) message that the caller already added.
if let Some(session_config) = &session {
let stored_conversation = SessionManager::get_session(&session_config.id, true)
.await?
.conversation
.ok_or_else(|| {
anyhow::anyhow!("Session {} has no conversation", session_config.id)
})?;
match conversation.len().cmp(&stored_conversation.len()) {
std::cmp::Ordering::Equal => {
if conversation != stored_conversation {
warn!("Session messages mismatch - replacing with incoming");
SessionManager::replace_conversation(&session_config.id, &conversation)
.await?;
}
}
std::cmp::Ordering::Greater
if conversation.len() == stored_conversation.len() + 1 =>
{
let last_message = conversation.last().unwrap();
if let Some(content) = last_message.content.first().and_then(|c| c.as_text()) {
debug!("user_message" = &content);
}
SessionManager::add_message(&session_config.id, last_message).await?;
}
_ => {
warn!(
"Unexpected session state: stored={}, incoming={}. Replacing.",
stored_conversation.len(),
conversation.len()
);
SessionManager::replace_conversation(&session_config.id, &conversation).await?;
}
}
let provider = self.provider().await?;
let session_id = session_config.id.clone();
tokio::spawn(async move {
if let Err(e) =
SessionManager::maybe_update_description(&session_id, provider).await
{
warn!("Failed to generate session description: {}", e);
}
});
}
Ok(Box::pin(async_stream::try_stream! {
@@ -1054,7 +1088,7 @@ impl Agent {
{
let mut autopilot = self.autopilot.lock().await;
if let Some((new_provider, role, model)) = autopilot.check_for_switch(&messages, self.provider().await?).await? {
if let Some((new_provider, role, model)) = autopilot.check_for_switch(&conversation, self.provider().await?).await? {
debug!("AutoPilot switching to {} role with model {}", role, model);
self.update_provider(new_provider).await?;
@@ -1065,17 +1099,16 @@ impl Agent {
}
}
let mut stream = Self::stream_response_from_provider(
self.provider().await?,
&system_prompt,
messages.messages(),
conversation.messages(),
&tools,
&toolshim_tools,
).await?;
let mut added_message = false;
let mut messages_to_add = Vec::new();
let mut no_tools_called = true;
let mut messages_to_add = Conversation::default();
let mut tools_updated = false;
while let Some(next) = stream.next().await {
@@ -1109,12 +1142,12 @@ impl Agent {
// Record usage for the session
if let Some(ref session_config) = &session {
if let Some(ref usage) = usage {
Self::update_session_metrics(session_config, usage, messages.len())
.await?;
Self::update_session_metrics(session_config, usage).await?;
}
}
if let Some(response) = response {
messages_to_add.push(response.clone());
let ToolCategorizeResult {
frontend_requests,
remaining_requests,
@@ -1161,7 +1194,7 @@ impl Agent {
let inspection_results = self.tool_inspection_manager
.inspect_tools(
&remaining_requests,
messages.messages(),
conversation.messages(),
)
.await?;
@@ -1260,25 +1293,26 @@ impl Agent {
let final_message_tool_resp = message_tool_response.lock().await.clone();
yield AgentEvent::Message(final_message_tool_resp.clone());
added_message = true;
messages_to_add.push(response);
no_tools_called = false;
messages_to_add.push(final_message_tool_resp);
}
}
Err(ProviderError::ContextLengthExceeded(error_msg)) => {
info!("Context length exceeded, attempting compaction");
match auto_compact::perform_compaction(self, messages.messages()).await {
match auto_compact::perform_compaction(self, conversation.messages()).await {
Ok(compact_result) => {
messages = compact_result.messages;
conversation = compact_result.messages;
yield AgentEvent::Message(
Message::assistant().with_summarization_requested(
"Context limit reached. Conversation has been automatically compacted to continue."
)
);
yield AgentEvent::HistoryReplaced(messages.messages().to_vec());
yield AgentEvent::HistoryReplaced(conversation.messages().to_vec());
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &conversation).await?
}
continue;
}
Err(_) => {
@@ -1301,40 +1335,49 @@ impl Agent {
if tools_updated {
(tools, toolshim_tools, system_prompt) = self.prepare_tools_and_prompt().await?;
}
if !added_message {
let mut exit_chat = false;
if no_tools_called {
if let Some(final_output_tool) = self.final_output_tool.lock().await.as_ref() {
if final_output_tool.final_output.is_none() {
tracing::warn!("Final output tool has not been called yet. Continuing agent loop.");
warn!("Final output tool has not been called yet. Continuing agent loop.");
let message = Message::user().with_text(FINAL_OUTPUT_CONTINUATION_MESSAGE);
messages_to_add.push(message.clone());
yield AgentEvent::Message(message);
messages.extend(messages_to_add);
continue
} else {
let message = Message::assistant().with_text(final_output_tool.final_output.clone().unwrap());
messages_to_add.push(message.clone());
yield AgentEvent::Message(message);
exit_chat = true;
}
}
match self.handle_retry_logic(&mut messages, &session, &initial_messages).await {
Ok(should_retry) => {
if should_retry {
info!("Retry logic triggered, restarting agent loop");
continue;
} else {
match self.handle_retry_logic(&mut conversation, &session, &initial_messages).await {
Ok(should_retry) => {
if should_retry {
info!("Retry logic triggered, restarting agent loop");
} else {
exit_chat = true;
}
}
Err(e) => {
error!("Retry logic failed: {}", e);
yield AgentEvent::Message(Message::assistant().with_text(
format!("Retry logic encountered an error: {}", e)
));
exit_chat = true;
}
}
Err(e) => {
error!("Retry logic failed: {}", e);
yield AgentEvent::Message(Message::assistant().with_text(
format!("Retry logic encountered an error: {}", e)
));
}
}
break;
}
messages.extend(messages_to_add);
if let Some(session_config) = &session {
for msg in &messages_to_add {
SessionManager::add_message(&session_config.id, msg).await?;
}
}
conversation.extend(messages_to_add);
if exit_chat {
break;
}
tokio::task::yield_now().await;
}
+20 -26
View File
@@ -15,7 +15,7 @@ use crate::providers::toolshim::{
modify_system_prompt_for_tool_json, OllamaInterpreter,
};
use crate::session;
use crate::session::SessionManager;
use rmcp::model::Tool;
async fn toolshim_postprocess(
@@ -276,23 +276,9 @@ impl Agent {
pub(crate) async fn update_session_metrics(
session_config: &crate::agents::types::SessionConfig,
usage: &ProviderUsage,
messages_length: usize,
) -> Result<()> {
let session_file_path = match session::storage::get_path(session_config.id.clone()) {
Ok(path) => path,
Err(e) => {
return Err(anyhow::anyhow!("Failed to get session file path: {}", e));
}
};
let mut metadata = session::storage::read_metadata(&session_file_path)?;
metadata.schedule_id = session_config.schedule_id.clone();
metadata.total_tokens = usage.usage.total_tokens;
metadata.input_tokens = usage.usage.input_tokens;
metadata.output_tokens = usage.usage.output_tokens;
metadata.message_count = messages_length + 1;
let session_id = session_config.id.as_str();
let session = SessionManager::get_session(session_id, false).await?;
let accumulate = |a: Option<i32>, b: Option<i32>| -> Option<i32> {
match (a, b) {
@@ -300,16 +286,24 @@ impl Agent {
_ => a.or(b),
}
};
metadata.accumulated_total_tokens =
accumulate(metadata.accumulated_total_tokens, usage.usage.total_tokens);
metadata.accumulated_input_tokens =
accumulate(metadata.accumulated_input_tokens, usage.usage.input_tokens);
metadata.accumulated_output_tokens = accumulate(
metadata.accumulated_output_tokens,
usage.usage.output_tokens,
);
session::storage::update_metadata(&session_file_path, &metadata).await?;
let accumulated_total =
accumulate(session.accumulated_total_tokens, usage.usage.total_tokens);
let accumulated_input =
accumulate(session.accumulated_input_tokens, usage.usage.input_tokens);
let accumulated_output =
accumulate(session.accumulated_output_tokens, usage.usage.output_tokens);
SessionManager::update_session(session_id)
.schedule_id(session_config.schedule_id.clone())
.total_tokens(usage.usage.total_tokens)
.input_tokens(usage.usage.input_tokens)
.output_tokens(usage.usage.output_tokens)
.accumulated_total_tokens(accumulated_total)
.accumulated_input_tokens(accumulated_input)
.accumulated_output_tokens(accumulated_output)
.apply()
.await?;
Ok(())
}
+8 -55
View File
@@ -421,12 +421,12 @@ impl Agent {
} else {
let sessions_info: Vec<String> = sessions
.into_iter()
.map(|(session_name, metadata)| {
.map(|(session_name, session)| {
format!(
"- Session: {} (Messages: {}, Working Dir: {})",
session_name,
metadata.message_count,
metadata.working_dir.display()
session.conversation.unwrap_or_default().len(),
session.working_dir.display()
)
})
.collect();
@@ -462,55 +462,19 @@ impl Agent {
)
})?;
// Get the session file path
let session_path = match crate::session::storage::get_path(
crate::session::storage::Identifier::Name(session_id.to_string()),
) {
Ok(path) => path,
Err(e) => {
return Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Invalid session ID '{}': {}", session_id, e),
None,
));
}
};
// Check if session file exists
if !session_path.exists() {
return Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Session '{}' not found", session_id),
None,
));
}
// Read session metadata
let metadata = match crate::session::storage::read_metadata(&session_path) {
let session = match crate::session::SessionManager::get_session(session_id, true).await {
Ok(metadata) => metadata,
Err(e) => {
return Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to read session metadata: {}", e),
None,
));
}
};
// Read session messages
let messages = match crate::session::storage::read_messages(&session_path) {
Ok(messages) => messages,
Err(e) => {
return Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to read session messages: {}", e),
format!("Failed to read session for '{}': {}", session_id, e),
None,
));
}
};
// Format the response with metadata and messages
let metadata_json = match serde_json::to_string_pretty(&metadata) {
let metadata_json = match serde_json::to_string_pretty(&session) {
Ok(json) => json,
Err(e) => {
return Err(ErrorData::new(
@@ -521,20 +485,9 @@ impl Agent {
}
};
let messages_json = match serde_json::to_string_pretty(&messages) {
Ok(json) => json,
Err(e) => {
return Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to serialize messages: {}", e),
None,
));
}
};
Ok(vec![Content::text(format!(
"Session '{}' Content:\n\nMetadata:\n{}\n\nMessages:\n{}",
session_id, metadata_json, messages_json
"Session '{}' Content:\n\nSession:\n{}",
session_id, metadata_json
))])
}
}
+1 -2
View File
@@ -1,4 +1,3 @@
use crate::session;
use mcp_core::ToolResult;
use rmcp::model::{Content, Tool};
use serde::{Deserialize, Serialize};
@@ -82,7 +81,7 @@ pub struct FrontendTool {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionConfig {
/// Unique identifier for the session
pub id: session::Identifier,
pub id: String,
/// Working directory for the session
pub working_dir: PathBuf,
/// ID of the schedule that triggered this session, if any
+43 -24
View File
@@ -55,7 +55,7 @@ pub async fn check_compaction_needed(
agent: &Agent,
messages: &[Message],
threshold_override: Option<f64>,
session_metadata: Option<&crate::session::storage::SessionMetadata>,
session_metadata: Option<&crate::session::Session>,
) -> Result<CompactionCheckResult> {
// Get threshold from config or use override
let config = Config::global();
@@ -182,7 +182,7 @@ pub async fn check_and_compact_messages(
agent: &Agent,
messages: &[Message],
threshold_override: Option<f64>,
session_metadata: Option<&crate::session::storage::SessionMetadata>,
session_metadata: Option<&crate::session::Session>,
) -> Result<AutoCompactResult> {
// First check if compaction is needed
let check_result =
@@ -242,6 +242,7 @@ pub async fn check_and_compact_messages(
mod tests {
use super::*;
use crate::conversation::message::{Message, MessageContent};
use crate::session::extension_data;
use crate::{
agents::Agent,
model::ModelConfig,
@@ -303,21 +304,32 @@ mod tests {
fn create_test_session_metadata(
message_count: usize,
working_dir: &str,
) -> crate::session::storage::SessionMetadata {
) -> crate::session::Session {
use crate::conversation::Conversation;
use std::path::PathBuf;
crate::session::storage::SessionMetadata {
message_count,
let mut conversation = Conversation::default();
for i in 0..message_count {
conversation.push(create_test_message(format!("message {}", i).as_str()));
}
crate::session::Session {
id: "test_session".to_string(),
working_dir: PathBuf::from(working_dir),
description: "Test session".to_string(),
created_at: "2024-01-01T00:00:00Z".to_string(),
updated_at: "2024-01-01T00:00:00Z".to_string(),
schedule_id: Some("test_job".to_string()),
recipe: None,
total_tokens: Some(100),
input_tokens: Some(50),
output_tokens: Some(50),
accumulated_total_tokens: Some(100),
accumulated_input_tokens: Some(50),
accumulated_output_tokens: Some(50),
extension_data: crate::session::ExtensionData::new(),
recipe: None,
extension_data: extension_data::ExtensionData::new(),
conversation: Some(conversation),
message_count,
}
}
@@ -540,7 +552,7 @@ mod tests {
#[tokio::test]
async fn test_auto_compact_uses_session_metadata() {
use crate::session::storage::SessionMetadata;
use crate::session::Session;
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
@@ -557,22 +569,22 @@ mod tests {
create_test_message("Second message"),
];
// Create session metadata with specific token counts
// Create session with specific token counts
#[allow(clippy::field_reassign_with_default)]
let mut session_metadata = SessionMetadata::default();
let mut session = Session::default();
{
session_metadata.total_tokens = Some(8000); // High token count to trigger compaction
session_metadata.accumulated_total_tokens = Some(15000); // Even higher accumulated count
session_metadata.input_tokens = Some(5000);
session_metadata.output_tokens = Some(3000);
session.total_tokens = Some(8000); // High token count to trigger compaction
session.accumulated_total_tokens = Some(15000); // Even higher accumulated count
session.input_tokens = Some(5000);
session.output_tokens = Some(3000);
}
// Test with session metadata - should use total_tokens for compaction (not accumulated)
// Test with session - should use total_tokens for compaction (not accumulated)
let result_with_metadata = check_compaction_needed(
&agent,
&messages,
Some(0.3), // 30% threshold
Some(&session_metadata),
Some(&session),
)
.await
.unwrap();
@@ -595,8 +607,8 @@ mod tests {
assert!(!result_without_metadata.needs_compaction);
assert!(result_without_metadata.current_tokens < 8000);
// Test with metadata that has only accumulated tokens (no total_tokens)
let mut session_metadata_no_total = SessionMetadata::default();
// Test with session that has only accumulated tokens (no total_tokens)
let mut session_metadata_no_total = Session::default();
#[allow(clippy::field_reassign_with_default)]
{
session_metadata_no_total.accumulated_total_tokens = Some(7500);
@@ -616,7 +628,7 @@ mod tests {
assert!(result_with_no_total.current_tokens < 7500);
// Test with metadata that has no token counts - should fall back to estimation
let empty_metadata = SessionMetadata::default();
let empty_metadata = Session::default();
let result_with_empty_metadata = check_compaction_needed(
&agent,
@@ -634,7 +646,7 @@ mod tests {
#[tokio::test]
async fn test_auto_compact_end_to_end_with_metadata() {
use crate::session::storage::SessionMetadata;
use crate::session::Session;
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
@@ -655,10 +667,10 @@ mod tests {
];
// Create session metadata with high token count to trigger compaction
let mut session_metadata = SessionMetadata::default();
let mut session = Session::default();
#[allow(clippy::field_reassign_with_default)]
{
session_metadata.total_tokens = Some(9000); // High enough to trigger compaction
session.total_tokens = Some(9000); // High enough to trigger compaction
}
// Test full compaction flow with session metadata
@@ -666,7 +678,7 @@ mod tests {
&agent,
&messages,
Some(0.3), // 30% threshold
Some(&session_metadata),
Some(&session),
)
.await
.unwrap();
@@ -704,7 +716,14 @@ mod tests {
let comprehensive_metadata = create_test_session_metadata(3, "/test/working/dir");
// Verify the helper created non-null metadata
assert_eq!(comprehensive_metadata.message_count, 3);
assert_eq!(
comprehensive_metadata
.clone()
.conversation
.unwrap_or_default()
.len(),
3
);
assert_eq!(
comprehensive_metadata.working_dir.to_str().unwrap(),
"/test/working/dir"
+19 -1
View File
@@ -3,11 +3,12 @@ use rmcp::model::Role;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use thiserror::Error;
use utoipa::ToSchema;
pub mod message;
mod tool_result_serde;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)]
pub struct Conversation(Vec<Message>);
#[derive(Error, Debug)]
@@ -122,6 +123,23 @@ impl Default for Conversation {
}
}
impl IntoIterator for Conversation {
type Item = Message;
type IntoIter = std::vec::IntoIter<Message>;
fn into_iter(self) -> Self::IntoIter {
self.0.into_iter()
}
}
impl<'a> IntoIterator for &'a Conversation {
type Item = &'a Message;
type IntoIter = std::slice::Iter<'a, Message>;
fn into_iter(self) -> Self::IntoIter {
self.0.iter()
}
}
/// Fix a conversation that we're about to send to an LLM. So the last and first
/// messages should always be from the user.
pub fn fix_conversation(conversation: Conversation) -> (Conversation, Vec<String>) {
+9 -2
View File
@@ -31,6 +31,8 @@ pub fn get_current_model() -> Option<String> {
CURRENT_MODEL.lock().ok().and_then(|model| model.clone())
}
pub static MSG_COUNT_FOR_SESSION_NAME_GENERATION: usize = 3;
/// Information about a model's capabilities
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)]
pub struct ModelInfo {
@@ -454,7 +456,7 @@ pub trait Provider: Send + Sync {
messages
.iter()
.filter(|m| m.role == rmcp::model::Role::User)
.take(3)
.take(MSG_COUNT_FOR_SESSION_NAME_GENERATION)
.map(|m| m.as_concat_text())
.collect()
}
@@ -476,7 +478,12 @@ pub trait Provider: Send + Sync {
)
.await?;
let description = result.0.as_concat_text();
let description = result
.0
.as_concat_text()
.split_whitespace()
.collect::<Vec<_>>()
.join(" ");
Ok(safe_truncate(&description, 100))
}
+79 -160
View File
@@ -21,8 +21,7 @@ use crate::providers::base::Provider as GooseProvider; // Alias to avoid conflic
use crate::providers::create;
use crate::recipe::Recipe;
use crate::scheduler_trait::SchedulerTrait;
use crate::session;
use crate::session::storage::SessionMetadata;
use crate::session::{Session, SessionManager};
// Track running tasks with their abort handles
type RunningTasksMap = HashMap<String, tokio::task::AbortHandle>;
@@ -649,38 +648,24 @@ impl Scheduler {
&self,
sched_id: &str,
limit: usize,
) -> Result<Vec<(String, SessionMetadata)>, SchedulerError> {
// Changed return type
let all_session_files = session::storage::list_sessions()
) -> Result<Vec<(String, Session)>, SchedulerError> {
let all_sessions = SessionManager::list_sessions()
.await
.map_err(|e| SchedulerError::StorageError(io::Error::other(e)))?;
let mut schedule_sessions: Vec<(String, SessionMetadata)> = Vec::new();
let mut schedule_sessions: Vec<(String, Session)> = Vec::new();
for (session_name, session_path) in all_session_files {
match session::storage::read_metadata(&session_path) {
Ok(metadata) => {
// metadata is not mutable here, and SessionMetadata is original
if metadata.schedule_id.as_deref() == Some(sched_id) {
schedule_sessions.push((session_name, metadata)); // Keep the tuple
}
}
Err(e) => {
tracing::warn!(
"Failed to read metadata for session file {}: {}. Skipping.",
session_path.display(),
e
);
}
for session in all_sessions {
if session.schedule_id.as_deref() == Some(sched_id) {
schedule_sessions.push((session.id.clone(), session));
}
}
schedule_sessions.sort_by(|a, b| b.0.cmp(&a.0));
schedule_sessions.sort_by(|a, b| b.0.cmp(&a.0)); // Sort by session_name (timestamp string)
// Keep the tuple, just take the limit
let result_sessions: Vec<(String, SessionMetadata)> =
let result_sessions: Vec<(String, Session)> =
schedule_sessions.into_iter().take(limit).collect();
Ok(result_sessions) // Return the Vec of tuples
Ok(result_sessions)
}
pub async fn run_now(&self, sched_id: &str) -> Result<String, SchedulerError> {
@@ -1066,7 +1051,7 @@ struct JobExecutionError {
async fn run_scheduled_job_internal(
job: ScheduledJob,
provider_override: Option<Arc<dyn GooseProvider>>, // New optional parameter
provider_override: Option<Arc<dyn GooseProvider>>,
jobs_arc: Option<Arc<Mutex<JobsMap>>>,
job_id: Option<String>,
) -> std::result::Result<String, JobExecutionError> {
@@ -1116,7 +1101,7 @@ async fn run_scheduled_job_internal(
let agent: Agent = Agent::new();
let agent_provider: Arc<dyn GooseProvider>; // Use the aliased GooseProvider
let agent_provider: Arc<dyn GooseProvider>;
if let Some(provider) = provider_override {
agent_provider = provider;
@@ -1155,7 +1140,8 @@ async fn run_scheduled_job_internal(
),
})?;
}
if let Some(recipe_extensions) = recipe.extensions {
if let Some(ref recipe_extensions) = recipe.extensions {
for extension in recipe_extensions {
agent
.add_extension(extension.clone())
@@ -1174,49 +1160,53 @@ async fn run_scheduled_job_internal(
});
}
tracing::info!("Agent configured with provider for job '{}'", job.id);
// Log the execution mode
let execution_mode = job.execution_mode.as_deref().unwrap_or("background");
tracing::info!("Job '{}' running in {} mode", job.id, execution_mode);
let session_id_for_return = session::generate_session_id();
let current_dir = match std::env::current_dir() {
Ok(cd) => cd,
Err(e) => {
return Err(JobExecutionError {
job_id: job.id.clone(),
error: format!("Failed to get current directory for job execution: {}", e),
});
}
};
// Create session upfront for both cases
let session = match SessionManager::create_session(
current_dir.clone(),
if recipe.prompt.is_some() {
format!("Scheduled job: {}", job.id)
} else {
"Empty job - no prompt".to_string()
},
)
.await
{
Ok(s) => s,
Err(e) => {
return Err(JobExecutionError {
job_id: job.id.clone(),
error: format!("Failed to create session: {}", e),
});
}
};
// Update the job with the session ID if we have access to the jobs arc
if let (Some(jobs_arc), Some(job_id_str)) = (jobs_arc.as_ref(), job_id.as_ref()) {
let mut jobs_guard = jobs_arc.lock().await;
if let Some((_, job_def)) = jobs_guard.get_mut(job_id_str) {
job_def.current_session_id = Some(session_id_for_return.clone());
job_def.current_session_id = Some(session.id.clone());
}
}
let session_file_path = match crate::session::storage::get_path(
crate::session::storage::Identifier::Name(session_id_for_return.clone()),
) {
Ok(path) => path,
Err(e) => {
return Err(JobExecutionError {
job_id: job.id.clone(),
error: format!("Failed to get session file path: {}", e),
});
}
};
if let Some(prompt_text) = recipe.prompt {
if let Some(ref prompt_text) = recipe.prompt {
let mut all_session_messages =
Conversation::new_unvalidated(vec![Message::user().with_text(prompt_text.clone())]);
let current_dir = match std::env::current_dir() {
Ok(cd) => cd,
Err(e) => {
return Err(JobExecutionError {
job_id: job.id.clone(),
error: format!("Failed to get current directory for job execution: {}", e),
});
}
};
let session_config = SessionConfig {
id: crate::session::storage::Identifier::Name(session_id_for_return.clone()),
id: session.id.clone(),
working_dir: current_dir.clone(),
schedule_id: Some(job.id.clone()),
execution_mode: job.execution_mode.clone(),
@@ -1236,7 +1226,6 @@ async fn run_scheduled_job_internal(
use futures::StreamExt;
while let Some(message_result) = stream.next().await {
// Check if the task has been cancelled
tokio::task::yield_now().await;
match message_result {
@@ -1246,15 +1235,9 @@ async fn run_scheduled_job_internal(
}
all_session_messages.push(msg);
}
Ok(AgentEvent::McpNotification(_)) => {
// Handle notifications if needed
}
Ok(AgentEvent::ModelChange { .. }) => {
// Model change events are informational, just continue
}
Ok(AgentEvent::HistoryReplaced(_)) => {
// Handle history replacement events if needed
}
Ok(AgentEvent::McpNotification(_)) => {}
Ok(AgentEvent::ModelChange { .. }) => {}
Ok(AgentEvent::HistoryReplaced(_)) => {}
Err(e) => {
tracing::error!(
"[Job {}] Error receiving message from agent: {}",
@@ -1265,51 +1248,6 @@ async fn run_scheduled_job_internal(
}
}
}
match crate::session::storage::read_metadata(&session_file_path) {
Ok(mut updated_metadata) => {
updated_metadata.message_count = all_session_messages.len();
if let Err(e) = crate::session::storage::save_messages_with_metadata(
&session_file_path,
&updated_metadata,
&all_session_messages,
) {
tracing::error!(
"[Job {}] Failed to persist final messages: {}",
job.id,
e
);
}
}
Err(e) => {
tracing::error!(
"[Job {}] Failed to read updated metadata before final save: {}",
job.id,
e
);
let fallback_metadata = crate::session::storage::SessionMetadata {
working_dir: current_dir.clone(),
description: String::new(),
schedule_id: Some(job.id.clone()),
message_count: all_session_messages.len(),
total_tokens: None,
input_tokens: None,
output_tokens: None,
accumulated_total_tokens: None,
accumulated_input_tokens: None,
accumulated_output_tokens: None,
extension_data: crate::session::ExtensionData::new(),
recipe: None,
};
if let Err(e_fb) = crate::session::storage::save_messages_with_metadata(
&session_file_path,
&fallback_metadata,
&all_session_messages,
) {
tracing::error!("[Job {}] Failed to persist final messages with fallback metadata: {}", job.id, e_fb);
}
}
}
}
Err(e) => {
return Err(JobExecutionError {
@@ -1324,28 +1262,19 @@ async fn run_scheduled_job_internal(
job.id,
job.source
);
let metadata = crate::session::storage::SessionMetadata {
working_dir: std::env::current_dir().unwrap_or_default(),
description: "Empty job - no prompt".to_string(),
schedule_id: Some(job.id.clone()),
message_count: 0,
..Default::default()
};
if let Err(e) = crate::session::storage::save_messages_with_metadata(
&session_file_path,
&metadata,
&Conversation::new_unvalidated(vec![]),
) {
tracing::error!(
"[Job {}] Failed to persist metadata for empty job: {}",
job.id,
e
);
}
}
if let Err(e) = SessionManager::update_session(&session.id)
.schedule_id(Some(job.id.clone()))
.recipe(Some(recipe))
.apply()
.await
{
tracing::error!("[Job {}] Failed to update session metadata: {}", job.id, e);
}
tracing::info!("Finished job: {}", job.id);
Ok(session_id_for_return)
Ok(session.id)
}
#[async_trait]
@@ -1378,7 +1307,7 @@ impl SchedulerTrait for Scheduler {
&self,
sched_id: &str,
limit: usize,
) -> Result<Vec<(String, SessionMetadata)>, SchedulerError> {
) -> Result<Vec<(String, Session)>, SchedulerError> {
self.sessions(sched_id, limit).await
}
@@ -1407,15 +1336,12 @@ mod tests {
use super::*;
use crate::recipe::Recipe;
use crate::{
model::ModelConfig, // Use the actual ModelConfig for the mock's field
model::ModelConfig,
providers::base::{ProviderMetadata, ProviderUsage, Usage},
providers::errors::ProviderError,
};
use rmcp::model::Tool;
use rmcp::model::{AnnotateAble, RawTextContent, Role};
// Removed: use crate::session::storage::{get_most_recent_session, read_metadata};
// `read_metadata` is still used by the test itself, so keep it or its module.
use crate::session::storage::read_metadata;
use crate::conversation::message::{Message, MessageContent};
use std::env;
@@ -1488,7 +1414,8 @@ mod tests {
let recipe_dir = temp_dir.path().join("recipes_for_test_scheduler");
fs::create_dir_all(&recipe_dir)?;
let _ = session::storage::ensure_session_dir().expect("Failed to ensure app session dir");
let _ = crate::session::session_manager::ensure_session_dir()
.expect("Failed to ensure app session dir");
let schedule_id_str = "test_schedule_001_scheduler_check".to_string();
let recipe_filename = recipe_dir.join(format!("{}.json", schedule_id_str));
@@ -1539,39 +1466,31 @@ mod tests {
.await
.expect("run_scheduled_job_internal failed");
let session_dir = session::storage::ensure_session_dir()?;
let expected_session_path = session_dir.join(format!("{}.jsonl", created_session_id));
assert!(
expected_session_path.exists(),
"Expected session file {} was not created",
expected_session_path.display()
);
let metadata = read_metadata(&expected_session_path)?;
let session = SessionManager::get_session(&created_session_id, true).await?;
let schedule_id = session.schedule_id.clone();
assert_eq!(
metadata.schedule_id,
schedule_id,
Some(schedule_id_str.clone()),
"Session metadata schedule_id ({:?}) does not match the job ID ({}). File: {}",
metadata.schedule_id,
"Session metadata schedule_id ({:?}) does not match the job ID ({}). Session: {}",
schedule_id,
schedule_id_str,
expected_session_path.display()
created_session_id
);
// Check if messages were written
let messages_in_file = crate::session::storage::read_messages(&expected_session_path)?;
// Check if messages were written using SessionManager
let messages_in_session = session.conversation.unwrap_or_default();
assert!(
!messages_in_file.is_empty(),
"No messages were written to the session file: {}",
expected_session_path.display()
!messages_in_session.is_empty(),
"No messages were written to the session: {}",
created_session_id
);
// We expect at least a user prompt and an assistant response
assert!(
messages_in_file.len() >= 2,
"Expected at least 2 messages (prompt + response), found {} in file: {}",
messages_in_file.len(),
expected_session_path.display()
messages_in_session.len() >= 2,
"Expected at least 2 messages (prompt + response), found {} in session: {}",
messages_in_session.len(),
created_session_id
);
// Clean up environment variables
+2 -2
View File
@@ -2,7 +2,7 @@ use async_trait::async_trait;
use chrono::{DateTime, Utc};
use crate::scheduler::{ScheduledJob, SchedulerError};
use crate::session::storage::SessionMetadata;
use crate::session::Session;
/// Common trait for all scheduler implementations
#[async_trait]
@@ -30,7 +30,7 @@ pub trait SchedulerTrait: Send + Sync {
&self,
sched_id: &str,
limit: usize,
) -> Result<Vec<(String, SessionMetadata)>, SchedulerError>;
) -> Result<Vec<(String, Session)>, SchedulerError>;
/// Update a schedule's cron expression
async fn update_schedule(&self, sched_id: &str, new_cron: String)
-99
View File
@@ -1,99 +0,0 @@
use crate::session::{self, SessionMetadata};
use anyhow::Result;
use serde::Serialize;
use std::cmp::Ordering;
use utoipa::ToSchema;
#[derive(Clone, Serialize, ToSchema)]
pub struct SessionInfo {
pub id: String,
pub path: String,
pub modified: String,
pub metadata: SessionMetadata,
}
/// Sort order for listing sessions
pub enum SortOrder {
Ascending,
Descending,
}
pub fn get_valid_sorted_sessions(sort_order: SortOrder) -> Result<Vec<SessionInfo>> {
let sessions = match session::list_sessions() {
Ok(sessions) => sessions,
Err(e) => {
tracing::error!("Failed to list sessions: {:?}", e);
return Err(anyhow::anyhow!("Failed to list sessions"));
}
};
let mut session_infos: Vec<SessionInfo> = Vec::new();
let mut corrupted_count = 0;
for (id, path) in sessions {
// Get file modification time with fallback
let modified = path
.metadata()
.and_then(|m| m.modified())
.map(|time| {
chrono::DateTime::<chrono::Utc>::from(time)
.format("%Y-%m-%d %H:%M:%S UTC")
.to_string()
})
.unwrap_or_else(|_| {
tracing::warn!("Failed to get modification time for session: {}", id);
"Unknown".to_string()
});
// Try to read metadata with error handling
match session::read_metadata(&path) {
Ok(metadata) => {
session_infos.push(SessionInfo {
id,
path: path.to_string_lossy().to_string(),
modified,
metadata,
});
}
Err(e) => {
corrupted_count += 1;
tracing::warn!(
"Failed to read metadata for session '{}': {}. Skipping corrupted session.",
id,
e
);
// Optionally, we could create a placeholder entry for corrupted sessions
// to show them in the UI with an error indicator, but for now we skip them
continue;
}
}
}
if corrupted_count > 0 {
tracing::warn!(
"Skipped {} corrupted sessions during listing",
corrupted_count
);
}
// Sort sessions by modified date
// Since all dates are in ISO format (YYYY-MM-DD HH:MM:SS UTC), we can just use string comparison
// This works because the ISO format ensures lexicographical ordering matches chronological ordering
session_infos.sort_by(|a, b| {
if a.modified == "Unknown" && b.modified == "Unknown" {
return Ordering::Equal;
} else if a.modified == "Unknown" {
return Ordering::Greater; // Unknown dates go last
} else if b.modified == "Unknown" {
return Ordering::Less;
}
match sort_order {
SortOrder::Ascending => a.modified.cmp(&b.modified),
SortOrder::Descending => b.modified.cmp(&a.modified),
}
});
Ok(session_infos)
}
+113
View File
@@ -0,0 +1,113 @@
use crate::conversation::Conversation;
use crate::session::Session;
use anyhow::Result;
use chrono::NaiveDateTime;
use std::fs;
use std::io::{self, BufRead};
use std::path::{Path, PathBuf};
use std::time::SystemTime;
const MAX_FILE_SIZE: u64 = 50 * 1024 * 1024;
pub fn list_sessions(session_dir: &PathBuf) -> Result<Vec<(String, PathBuf)>> {
let entries = fs::read_dir(session_dir)?
.filter_map(|entry| {
let entry = entry.ok()?;
let path = entry.path();
if path.extension().is_some_and(|ext| ext == "jsonl") {
let name = path.file_stem()?.to_string_lossy().to_string();
Some((name, path))
} else {
None
}
})
.collect::<Vec<_>>();
Ok(entries)
}
pub fn load_session(session_name: &str, session_path: &Path) -> Result<Session> {
let file = fs::File::open(session_path).map_err(|e| {
anyhow::anyhow!(
"Failed to open session file {}: {}",
session_path.display(),
e
)
})?;
let file_metadata = file.metadata()?;
if file_metadata.len() > MAX_FILE_SIZE {
return Err(anyhow::anyhow!("Session file too large"));
}
if file_metadata.len() == 0 {
return Err(anyhow::anyhow!("Empty session file"));
}
let modified_time = file_metadata.modified().unwrap_or(SystemTime::now());
let created_time = file_metadata
.created()
.unwrap_or_else(|_| parse_session_timestamp(session_name).unwrap_or(modified_time));
let reader = io::BufReader::new(file);
let mut lines = reader.lines();
let mut messages = Vec::new();
let mut session = Session {
id: session_name.to_string(),
..Default::default()
};
if let Some(Ok(line)) = lines.next() {
let mut metadata_json: serde_json::Value = serde_json::from_str(&line)
.map_err(|_| anyhow::anyhow!("Invalid session metadata JSON"))?;
if let Some(obj) = metadata_json.as_object_mut() {
obj.entry("id").or_insert(serde_json::json!(session_name));
obj.entry("created_at")
.or_insert(serde_json::json!(format_timestamp(created_time)?));
obj.entry("updated_at")
.or_insert(serde_json::json!(format_timestamp(modified_time)?));
obj.entry("extension_data").or_insert(serde_json::json!({}));
obj.entry("message_count").or_insert(serde_json::json!(0));
if let Some(desc) = obj.get_mut("description") {
if let Some(desc_str) = desc.as_str() {
*desc = serde_json::json!(desc_str
.split_whitespace()
.collect::<Vec<_>>()
.join(" "));
}
}
}
session = serde_json::from_value(metadata_json)?;
session.id = session_name.to_string();
}
for line in lines.map_while(Result::ok) {
if let Ok(message) = serde_json::from_str(&line) {
messages.push(message);
}
}
if !messages.is_empty() {
session.conversation = Some(Conversation::new_unvalidated(messages));
}
Ok(session)
}
fn format_timestamp(time: SystemTime) -> Result<String> {
let duration = time.duration_since(std::time::UNIX_EPOCH)?;
let timestamp = chrono::DateTime::from_timestamp(duration.as_secs() as i64, 0)
.unwrap_or_default()
.format("%Y-%m-%d %H:%M:%S")
.to_string();
Ok(timestamp)
}
fn parse_session_timestamp(session_name: &str) -> Option<SystemTime> {
NaiveDateTime::parse_from_str(session_name, "%Y%m%d_%H%M%S")
.ok()
.map(|dt| SystemTime::from(dt.and_utc()))
}
+3 -12
View File
@@ -1,14 +1,5 @@
pub mod extension_data;
pub mod info;
pub mod storage;
mod legacy;
pub mod session_manager;
// Re-export common session types and functions
pub use storage::{
ensure_session_dir, generate_description, generate_description_with_schedule_id,
generate_session_id, get_most_recent_session, get_path, list_sessions, persist_messages,
persist_messages_with_schedule_id, read_messages, read_metadata, update_metadata, Identifier,
SessionMetadata,
};
pub use extension_data::{ExtensionData, ExtensionState, TodoState};
pub use info::{get_valid_sorted_sessions, SessionInfo};
pub use session_manager::{Session, SessionInsights, SessionManager};
+858
View File
@@ -0,0 +1,858 @@
use crate::config::APP_STRATEGY;
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::providers::base::{Provider, MSG_COUNT_FOR_SESSION_NAME_GENERATION};
use crate::recipe::Recipe;
use crate::session::extension_data::ExtensionData;
use anyhow::Result;
use etcetera::{choose_app_strategy, AppStrategy};
use rmcp::model::Role;
use serde::{Deserialize, Serialize};
use sqlx::sqlite::SqliteConnectOptions;
use sqlx::{Pool, Sqlite};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::OnceCell;
use tracing::{info, warn};
use utoipa::ToSchema;
const CURRENT_SCHEMA_VERSION: i32 = 1;
static SESSION_STORAGE: OnceCell<Arc<SessionStorage>> = OnceCell::const_new();
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct Session {
pub id: String,
#[schema(value_type = String)]
pub working_dir: PathBuf,
pub description: String,
pub created_at: String,
pub updated_at: String,
pub extension_data: ExtensionData,
pub total_tokens: Option<i32>,
pub input_tokens: Option<i32>,
pub output_tokens: Option<i32>,
pub accumulated_total_tokens: Option<i32>,
pub accumulated_input_tokens: Option<i32>,
pub accumulated_output_tokens: Option<i32>,
pub schedule_id: Option<String>,
pub recipe: Option<Recipe>,
pub conversation: Option<Conversation>,
pub message_count: usize,
}
pub struct SessionUpdateBuilder {
session_id: String,
description: Option<String>,
working_dir: Option<PathBuf>,
extension_data: Option<ExtensionData>,
total_tokens: Option<Option<i32>>,
input_tokens: Option<Option<i32>>,
output_tokens: Option<Option<i32>>,
accumulated_total_tokens: Option<Option<i32>>,
accumulated_input_tokens: Option<Option<i32>>,
accumulated_output_tokens: Option<Option<i32>>,
schedule_id: Option<Option<String>>,
recipe: Option<Option<Recipe>>,
}
#[derive(Serialize, ToSchema, Debug)]
#[serde(rename_all = "camelCase")]
pub struct SessionInsights {
/// Total number of sessions
total_sessions: usize,
/// Total tokens used across all sessions
total_tokens: i64,
}
impl SessionUpdateBuilder {
fn new(session_id: String) -> Self {
Self {
session_id,
description: None,
working_dir: None,
extension_data: None,
total_tokens: None,
input_tokens: None,
output_tokens: None,
accumulated_total_tokens: None,
accumulated_input_tokens: None,
accumulated_output_tokens: None,
schedule_id: None,
recipe: None,
}
}
pub fn description(mut self, description: impl Into<String>) -> Self {
self.description = Some(description.into());
self
}
pub fn working_dir(mut self, working_dir: PathBuf) -> Self {
self.working_dir = Some(working_dir);
self
}
pub fn extension_data(mut self, data: ExtensionData) -> Self {
self.extension_data = Some(data);
self
}
pub fn total_tokens(mut self, tokens: Option<i32>) -> Self {
self.total_tokens = Some(tokens);
self
}
pub fn input_tokens(mut self, tokens: Option<i32>) -> Self {
self.input_tokens = Some(tokens);
self
}
pub fn output_tokens(mut self, tokens: Option<i32>) -> Self {
self.output_tokens = Some(tokens);
self
}
pub fn accumulated_total_tokens(mut self, tokens: Option<i32>) -> Self {
self.accumulated_total_tokens = Some(tokens);
self
}
pub fn accumulated_input_tokens(mut self, tokens: Option<i32>) -> Self {
self.accumulated_input_tokens = Some(tokens);
self
}
pub fn accumulated_output_tokens(mut self, tokens: Option<i32>) -> Self {
self.accumulated_output_tokens = Some(tokens);
self
}
pub fn schedule_id(mut self, schedule_id: Option<String>) -> Self {
self.schedule_id = Some(schedule_id);
self
}
pub fn recipe(mut self, recipe: Option<Recipe>) -> Self {
self.recipe = Some(recipe);
self
}
pub async fn apply(self) -> Result<()> {
SessionManager::apply_update(self).await
}
}
pub struct SessionManager;
impl SessionManager {
pub async fn instance() -> Result<Arc<SessionStorage>> {
SESSION_STORAGE
.get_or_try_init(|| async { SessionStorage::new().await.map(Arc::new) })
.await
.map(Arc::clone)
}
pub async fn create_session(working_dir: PathBuf, description: String) -> Result<Session> {
let today = chrono::Utc::now().format("%Y%m%d").to_string();
let storage = Self::instance().await?;
let mut tx = storage.pool.begin().await?;
let max_idx = sqlx::query_scalar::<_, Option<i32>>(
"SELECT MAX(CAST(SUBSTR(id, 10) AS INTEGER)) FROM sessions WHERE id LIKE ?",
)
.bind(format!("{}_%", today))
.fetch_one(&mut *tx)
.await?
.unwrap_or(0);
let session_id = format!("{}_{}", today, max_idx + 1);
sqlx::query(
r#"
INSERT INTO sessions (id, description, working_dir, extension_data)
VALUES (?, ?, ?, '{}')
"#,
)
.bind(&session_id)
.bind(&description)
.bind(working_dir.to_string_lossy().as_ref())
.execute(&mut *tx)
.await?;
tx.commit().await?;
Self::get_session(&session_id, false).await
}
pub async fn get_session(id: &str, include_messages: bool) -> Result<Session> {
Self::instance()
.await?
.get_session(id, include_messages)
.await
}
pub fn update_session(id: &str) -> SessionUpdateBuilder {
SessionUpdateBuilder::new(id.to_string())
}
async fn apply_update(builder: SessionUpdateBuilder) -> Result<()> {
Self::instance().await?.apply_update(builder).await
}
pub async fn add_message(id: &str, message: &Message) -> Result<()> {
Self::instance().await?.add_message(id, message).await
}
pub async fn replace_conversation(id: &str, conversation: &Conversation) -> Result<()> {
Self::instance()
.await?
.replace_conversation(id, conversation)
.await
}
pub async fn list_sessions() -> Result<Vec<Session>> {
Self::instance().await?.list_sessions().await
}
pub async fn delete_session(id: &str) -> Result<()> {
Self::instance().await?.delete_session(id).await
}
pub async fn get_insights() -> Result<SessionInsights> {
Self::instance().await?.get_insights().await
}
pub async fn maybe_update_description(id: &str, provider: Arc<dyn Provider>) -> Result<()> {
let session = Self::get_session(id, true).await?;
let conversation = session
.conversation
.ok_or_else(|| anyhow::anyhow!("No messages found"))?;
let user_message_count = conversation
.messages()
.iter()
.filter(|m| matches!(m.role, Role::User))
.count();
if user_message_count <= MSG_COUNT_FOR_SESSION_NAME_GENERATION {
let description = provider.generate_session_name(&conversation).await?;
Self::update_session(id)
.description(description)
.apply()
.await
} else {
Ok(())
}
}
}
pub struct SessionStorage {
pool: Pool<Sqlite>,
}
pub fn ensure_session_dir() -> Result<PathBuf> {
let data_dir = choose_app_strategy(APP_STRATEGY.clone())
.expect("goose requires a home dir")
.data_dir()
.join("sessions");
if !data_dir.exists() {
fs::create_dir_all(&data_dir)?;
}
Ok(data_dir)
}
fn role_to_string(role: &Role) -> &'static str {
match role {
Role::User => "user",
Role::Assistant => "assistant",
}
}
impl Default for Session {
fn default() -> Self {
Self {
id: String::new(),
working_dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
description: String::new(),
created_at: String::new(),
updated_at: String::new(),
extension_data: ExtensionData::default(),
total_tokens: None,
input_tokens: None,
output_tokens: None,
accumulated_total_tokens: None,
accumulated_input_tokens: None,
accumulated_output_tokens: None,
schedule_id: None,
recipe: None,
conversation: None,
message_count: 0,
}
}
}
impl Session {
pub fn without_messages(mut self) -> Self {
self.conversation = None;
self
}
}
impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session {
fn from_row(row: &sqlx::sqlite::SqliteRow) -> Result<Self, sqlx::Error> {
use sqlx::Row;
let recipe_json: Option<String> = row.try_get("recipe_json")?;
let recipe = recipe_json.and_then(|json| serde_json::from_str(&json).ok());
Ok(Session {
id: row.try_get("id")?,
working_dir: PathBuf::from(row.try_get::<String, _>("working_dir")?),
description: row.try_get("description")?,
created_at: row.try_get("created_at")?,
updated_at: row.try_get("updated_at")?,
extension_data: serde_json::from_str(&row.try_get::<String, _>("extension_data")?)
.unwrap_or_default(),
total_tokens: row.try_get("total_tokens")?,
input_tokens: row.try_get("input_tokens")?,
output_tokens: row.try_get("output_tokens")?,
accumulated_total_tokens: row.try_get("accumulated_total_tokens")?,
accumulated_input_tokens: row.try_get("accumulated_input_tokens")?,
accumulated_output_tokens: row.try_get("accumulated_output_tokens")?,
schedule_id: row.try_get("schedule_id")?,
recipe,
conversation: None,
message_count: row.try_get("message_count").unwrap_or(0) as usize,
})
}
}
impl SessionStorage {
async fn new() -> Result<Self> {
let session_dir = ensure_session_dir()?;
let db_path = session_dir.join("sessions.db");
let storage = if db_path.exists() {
Self::open(&db_path).await?
} else {
let storage = Self::create(&db_path).await?;
if let Err(e) = storage.import_legacy(&session_dir).await {
warn!("Failed to import some legacy sessions: {}", e);
}
storage
};
Ok(storage)
}
async fn get_pool(db_path: &Path, create_if_missing: bool) -> Result<Pool<Sqlite>> {
let options = SqliteConnectOptions::new()
.filename(db_path)
.create_if_missing(create_if_missing);
sqlx::SqlitePool::connect_with(options).await.map_err(|e| {
anyhow::anyhow!(
"Failed to open SQLite database at '{}': {}",
db_path.display(),
e
)
})
}
async fn open(db_path: &Path) -> Result<Self> {
let pool = Self::get_pool(db_path, false).await?;
let storage = Self { pool };
storage.run_migrations().await?;
Ok(storage)
}
async fn create(db_path: &Path) -> Result<Self> {
let pool = Self::get_pool(db_path, true).await?;
sqlx::query(
r#"
CREATE TABLE schema_version (
version INTEGER PRIMARY KEY,
applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
"#,
)
.execute(&pool)
.await?;
sqlx::query("INSERT INTO schema_version (version) VALUES (?)")
.bind(CURRENT_SCHEMA_VERSION)
.execute(&pool)
.await?;
sqlx::query(
r#"
CREATE TABLE sessions (
id TEXT PRIMARY KEY,
description TEXT NOT NULL DEFAULT '',
working_dir TEXT NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
extension_data TEXT DEFAULT '{}',
total_tokens INTEGER,
input_tokens INTEGER,
output_tokens INTEGER,
accumulated_total_tokens INTEGER,
accumulated_input_tokens INTEGER,
accumulated_output_tokens INTEGER,
schedule_id TEXT,
recipe_json TEXT
)
"#,
)
.execute(&pool)
.await?;
sqlx::query(
r#"
CREATE TABLE messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL REFERENCES sessions(id),
role TEXT NOT NULL,
content_json TEXT NOT NULL,
created_timestamp INTEGER NOT NULL,
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
tokens INTEGER
)
"#,
)
.execute(&pool)
.await?;
sqlx::query("CREATE INDEX idx_messages_session ON messages(session_id)")
.execute(&pool)
.await?;
sqlx::query("CREATE INDEX idx_messages_timestamp ON messages(timestamp)")
.execute(&pool)
.await?;
sqlx::query("CREATE INDEX idx_sessions_updated ON sessions(updated_at DESC)")
.execute(&pool)
.await?;
Ok(Self { pool })
}
async fn import_legacy(&self, session_dir: &PathBuf) -> Result<()> {
use crate::session::legacy;
let sessions = match legacy::list_sessions(session_dir) {
Ok(sessions) => sessions,
Err(_) => {
warn!("No legacy sessions found to import");
return Ok(());
}
};
if sessions.is_empty() {
return Ok(());
}
let mut imported_count = 0;
let mut failed_count = 0;
for (session_name, session_path) in sessions {
match legacy::load_session(&session_name, &session_path) {
Ok(session) => match self.import_legacy_session(&session).await {
Ok(_) => {
imported_count += 1;
info!(" ✓ Imported: {}", session_name);
}
Err(e) => {
failed_count += 1;
info!(" ✗ Failed to import {}: {}", session_name, e);
}
},
Err(e) => {
failed_count += 1;
info!(" ✗ Failed to load {}: {}", session_name, e);
}
}
}
info!(
"Import complete: {} successful, {} failed",
imported_count, failed_count
);
Ok(())
}
async fn import_legacy_session(&self, session: &Session) -> Result<()> {
let recipe_json = match &session.recipe {
Some(recipe) => Some(serde_json::to_string(recipe)?),
None => None,
};
sqlx::query(
r#"
INSERT INTO sessions (
id, description, working_dir, created_at, updated_at, extension_data,
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
schedule_id, recipe_json
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&session.id)
.bind(&session.description)
.bind(session.working_dir.to_string_lossy().as_ref())
.bind(&session.created_at)
.bind(&session.updated_at)
.bind(serde_json::to_string(&session.extension_data)?)
.bind(session.total_tokens)
.bind(session.input_tokens)
.bind(session.output_tokens)
.bind(session.accumulated_total_tokens)
.bind(session.accumulated_input_tokens)
.bind(session.accumulated_output_tokens)
.bind(&session.schedule_id)
.bind(recipe_json)
.execute(&self.pool)
.await?;
if let Some(conversation) = &session.conversation {
self.replace_conversation(&session.id, conversation).await?;
}
Ok(())
}
async fn run_migrations(&self) -> Result<()> {
let current_version = self.get_schema_version().await?;
if current_version < CURRENT_SCHEMA_VERSION {
info!(
"Running database migrations from v{} to v{}...",
current_version, CURRENT_SCHEMA_VERSION
);
for version in (current_version + 1)..=CURRENT_SCHEMA_VERSION {
info!(" Applying migration v{}...", version);
self.apply_migration(version).await?;
self.update_schema_version(version).await?;
info!(" ✓ Migration v{} complete", version);
}
info!("All migrations complete");
}
Ok(())
}
async fn get_schema_version(&self) -> Result<i32> {
let table_exists = sqlx::query_scalar::<_, bool>(
r#"
SELECT EXISTS (
SELECT name FROM sqlite_master
WHERE type='table' AND name='schema_version'
)
"#,
)
.fetch_one(&self.pool)
.await?;
if !table_exists {
return Ok(0);
}
let version = sqlx::query_scalar::<_, i32>("SELECT MAX(version) FROM schema_version")
.fetch_one(&self.pool)
.await?;
Ok(version)
}
async fn update_schema_version(&self, version: i32) -> Result<()> {
sqlx::query("INSERT INTO schema_version (version) VALUES (?)")
.bind(version)
.execute(&self.pool)
.await?;
Ok(())
}
async fn apply_migration(&self, version: i32) -> Result<()> {
match version {
1 => {
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS schema_version (
version INTEGER PRIMARY KEY,
applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
"#,
)
.execute(&self.pool)
.await?;
}
_ => {
anyhow::bail!("Unknown migration version: {}", version);
}
}
Ok(())
}
async fn get_session(&self, id: &str, include_messages: bool) -> Result<Session> {
let mut session = sqlx::query_as::<_, Session>(
r#"
SELECT id, working_dir, description, created_at, updated_at, extension_data,
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
schedule_id, recipe_json
FROM sessions
WHERE id = ?
"#,
)
.bind(id)
.fetch_optional(&self.pool)
.await?
.ok_or_else(|| anyhow::anyhow!("Session not found"))?;
if include_messages {
let conv = self.get_conversation(&session.id).await?;
session.message_count = conv.messages().len();
session.conversation = Some(conv);
} else {
let count =
sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM messages WHERE session_id = ?")
.bind(&session.id)
.fetch_one(&self.pool)
.await? as usize;
session.message_count = count;
}
Ok(session)
}
async fn apply_update(&self, builder: SessionUpdateBuilder) -> Result<()> {
let mut updates = Vec::new();
let mut query = String::from("UPDATE sessions SET ");
macro_rules! add_update {
($field:expr, $name:expr) => {
if $field.is_some() {
if !updates.is_empty() {
query.push_str(", ");
}
updates.push($name);
query.push_str($name);
query.push_str(" = ?");
}
};
}
add_update!(builder.description, "description");
add_update!(builder.working_dir, "working_dir");
add_update!(builder.extension_data, "extension_data");
add_update!(builder.total_tokens, "total_tokens");
add_update!(builder.input_tokens, "input_tokens");
add_update!(builder.output_tokens, "output_tokens");
add_update!(builder.accumulated_total_tokens, "accumulated_total_tokens");
add_update!(builder.accumulated_input_tokens, "accumulated_input_tokens");
add_update!(
builder.accumulated_output_tokens,
"accumulated_output_tokens"
);
add_update!(builder.schedule_id, "schedule_id");
add_update!(builder.recipe, "recipe_json");
if updates.is_empty() {
return Ok(());
}
if !updates.is_empty() {
query.push_str(", ");
}
query.push_str("updated_at = datetime('now') WHERE id = ?");
let mut q = sqlx::query(&query);
if let Some(desc) = builder.description {
q = q.bind(desc);
}
if let Some(wd) = builder.working_dir {
q = q.bind(wd.to_string_lossy().to_string());
}
if let Some(ed) = builder.extension_data {
q = q.bind(serde_json::to_string(&ed)?);
}
if let Some(tt) = builder.total_tokens {
q = q.bind(tt);
}
if let Some(it) = builder.input_tokens {
q = q.bind(it);
}
if let Some(ot) = builder.output_tokens {
q = q.bind(ot);
}
if let Some(att) = builder.accumulated_total_tokens {
q = q.bind(att);
}
if let Some(ait) = builder.accumulated_input_tokens {
q = q.bind(ait);
}
if let Some(aot) = builder.accumulated_output_tokens {
q = q.bind(aot);
}
if let Some(sid) = builder.schedule_id {
q = q.bind(sid);
}
if let Some(recipe) = builder.recipe {
let recipe_json = recipe.map(|r| serde_json::to_string(&r)).transpose()?;
q = q.bind(recipe_json);
}
q = q.bind(&builder.session_id);
q.execute(&self.pool).await?;
Ok(())
}
async fn get_conversation(&self, session_id: &str) -> Result<Conversation> {
let rows = sqlx::query_as::<_, (String, String, i64)>(
"SELECT role, content_json, created_timestamp FROM messages WHERE session_id = ? ORDER BY timestamp",
)
.bind(session_id)
.fetch_all(&self.pool)
.await?;
let mut messages = Vec::new();
for (role_str, content_json, created_timestamp) in rows {
let role = match role_str.as_str() {
"user" => Role::User,
"assistant" => Role::Assistant,
_ => continue,
};
let content = serde_json::from_str(&content_json)?;
let message = Message::new(role, created_timestamp, content);
messages.push(message);
}
Ok(Conversation::new_unvalidated(messages))
}
async fn add_message(&self, session_id: &str, message: &Message) -> Result<()> {
sqlx::query(
r#"
INSERT INTO messages (session_id, role, content_json, created_timestamp)
VALUES (?, ?, ?, ?)
"#,
)
.bind(session_id)
.bind(role_to_string(&message.role))
.bind(serde_json::to_string(&message.content)?)
.bind(message.created)
.execute(&self.pool)
.await?;
sqlx::query("UPDATE sessions SET updated_at = datetime('now') WHERE id = ?")
.bind(session_id)
.execute(&self.pool)
.await?;
Ok(())
}
async fn replace_conversation(
&self,
session_id: &str,
conversation: &Conversation,
) -> Result<()> {
let mut tx = self.pool.begin().await?;
sqlx::query("DELETE FROM messages WHERE session_id = ?")
.bind(session_id)
.execute(&mut *tx)
.await?;
for message in conversation.messages() {
sqlx::query(
r#"
INSERT INTO messages (session_id, role, content_json, created_timestamp)
VALUES (?, ?, ?, ?)
"#,
)
.bind(session_id)
.bind(role_to_string(&message.role))
.bind(serde_json::to_string(&message.content)?)
.bind(message.created)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
Ok(())
}
async fn list_sessions(&self) -> Result<Vec<Session>> {
sqlx::query_as::<_, Session>(
r#"
SELECT s.id, s.working_dir, s.description, s.created_at, s.updated_at, s.extension_data,
s.total_tokens, s.input_tokens, s.output_tokens,
s.accumulated_total_tokens, s.accumulated_input_tokens, s.accumulated_output_tokens,
s.schedule_id, s.recipe_json,
COUNT(m.id) as message_count
FROM sessions s
INNER JOIN messages m ON s.id = m.session_id
GROUP BY s.id
ORDER BY s.updated_at DESC
"#,
)
.fetch_all(&self.pool)
.await
.map_err(Into::into)
}
async fn delete_session(&self, session_id: &str) -> Result<()> {
let exists =
sqlx::query_scalar::<_, bool>("SELECT EXISTS(SELECT 1 FROM sessions WHERE id = ?)")
.bind(session_id)
.fetch_one(&self.pool)
.await?;
if !exists {
return Err(anyhow::anyhow!("Session not found"));
}
sqlx::query("DELETE FROM messages WHERE session_id = ?")
.bind(session_id)
.execute(&self.pool)
.await?;
sqlx::query("DELETE FROM sessions WHERE id = ?")
.bind(session_id)
.execute(&self.pool)
.await?;
Ok(())
}
async fn get_insights(&self) -> Result<SessionInsights> {
let row = sqlx::query_as::<_, (i64, Option<i64>)>(
r#"
SELECT COUNT(*) as total_sessions,
COALESCE(SUM(COALESCE(accumulated_total_tokens, total_tokens, 0)), 0) as total_tokens
FROM sessions
"#,
)
.fetch_one(&self.pool)
.await?;
Ok(SessionInsights {
total_sessions: row.0 as usize,
total_tokens: row.1.unwrap_or(0),
})
}
}
File diff suppressed because it is too large Load Diff
+56 -81
View File
@@ -11,7 +11,7 @@ use tracing::{info, warn};
use crate::scheduler::{normalize_cron_expression, ScheduledJob, SchedulerError};
use crate::scheduler_trait::SchedulerTrait;
use crate::session::storage::SessionMetadata;
use crate::session::{Session, SessionManager};
const TEMPORAL_SERVICE_STARTUP_TIMEOUT: Duration = Duration::from_secs(15);
const TEMPORAL_SERVICE_HEALTH_CHECK_INTERVAL: Duration = Duration::from_millis(500);
@@ -699,37 +699,23 @@ impl TemporalScheduler {
}
}
// Note: This method fetches sessions from the session storage directly
// since Temporal service doesn't track session metadata
pub async fn sessions(
&self,
sched_id: &str,
limit: usize,
) -> Result<Vec<(String, SessionMetadata)>, SchedulerError> {
use crate::session::storage;
) -> Result<Vec<(String, Session)>, SchedulerError> {
use crate::session::SessionManager;
// Get all session files
let all_session_files = storage::list_sessions().map_err(|e| {
// Get all sessions from the database
let all_sessions = SessionManager::list_sessions().await.map_err(|e| {
SchedulerError::SchedulerInternalError(format!("Failed to list sessions: {}", e))
})?;
let mut schedule_sessions: Vec<(String, SessionMetadata)> = Vec::new();
let mut schedule_sessions: Vec<(String, Session)> = Vec::new();
for (session_name, session_path) in all_session_files {
match storage::read_metadata(&session_path) {
Ok(metadata) => {
// Check if this session belongs to the requested schedule
if metadata.schedule_id.as_deref() == Some(sched_id) {
schedule_sessions.push((session_name, metadata));
}
}
Err(e) => {
tracing::warn!(
"Failed to read metadata for session file {}: {}. Skipping.",
session_path.display(),
e
);
}
for session in all_sessions {
if session.schedule_id.as_deref() == Some(sched_id) {
schedule_sessions.push((session.id.clone(), session));
}
}
@@ -737,10 +723,10 @@ impl TemporalScheduler {
schedule_sessions.sort_by(|a, b| b.0.cmp(&a.0));
// Take only the requested limit
let result_sessions: Vec<(String, SessionMetadata)> =
let result_sessions: Vec<(String, Session)> =
schedule_sessions.into_iter().take(limit).collect();
tracing::info!(
info!(
"Found {} sessions for schedule '{}'",
result_sessions.len(),
sched_id
@@ -858,38 +844,41 @@ impl TemporalScheduler {
let recent_sessions = self.sessions(&job.id, 3).await?;
let mut has_active_session = false;
for (session_name, _) in recent_sessions {
let session_path = match crate::session::storage::get_path(
crate::session::storage::Identifier::Name(session_name.clone()),
) {
Ok(path) => path,
for (session_id, _) in recent_sessions {
// Get session info from database to check last update time
match SessionManager::list_sessions().await {
Ok(all_sessions) => {
if let Some(session_info) =
all_sessions.iter().find(|s| s.id == session_id)
{
// Parse the updated_at timestamp from the database
if let Ok(modified_dt) = DateTime::parse_from_str(
&session_info.updated_at,
"%Y-%m-%d %H:%M:%S UTC",
) {
let modified_utc = modified_dt.with_timezone(&Utc);
let now = Utc::now();
let time_diff = now.signed_duration_since(modified_utc);
// Increased tolerance to 5 minutes to reduce false positives
if time_diff.num_minutes() < 5 {
has_active_session = true;
tracing::debug!(
"Found active session for job '{}' modified {} minutes ago",
job.id,
time_diff.num_minutes()
);
break;
}
}
}
}
Err(e) => {
tracing::warn!(
"Failed to get session path for '{}': {}",
session_name,
"Failed to list sessions to check activity for job '{}': {}",
job.id,
e
);
continue;
}
};
// Check if session file was modified recently (within last 5 minutes instead of 2)
if let Ok(metadata) = std::fs::metadata(&session_path) {
if let Ok(modified) = metadata.modified() {
let modified_dt: DateTime<Utc> = modified.into();
let now = Utc::now();
let time_diff = now.signed_duration_since(modified_dt);
// Increased tolerance to 5 minutes to reduce false positives
if time_diff.num_minutes() < 5 {
has_active_session = true;
tracing::debug!(
"Found active session for job '{}' modified {} minutes ago",
job.id,
time_diff.num_minutes()
);
break;
}
}
}
}
@@ -897,9 +886,9 @@ impl TemporalScheduler {
// Only mark as completed if both Temporal service check failed AND no recent session activity
if !has_active_session {
tracing::info!(
"No active sessions found for job '{}' in the last 5 minutes, marking as completed",
job.id
);
"No active sessions found for job '{}' in the last 5 minutes, marking as completed",
job.id
);
let request = JobRequest {
action: "mark_completed".to_string(),
@@ -966,14 +955,16 @@ impl TemporalScheduler {
let recent_sessions = self.sessions(sched_id, 1).await?;
if let Some((session_name, _session_metadata)) = recent_sessions.first() {
// Check if this session is still active by looking at the session file
let session_path = match crate::session::storage::get_path(
crate::session::storage::Identifier::Name(session_name.clone()),
) {
Ok(path) => path,
// Check if this session still exists in the new system
match SessionManager::get_session(session_name, false).await {
Ok(_) => {
// Session exists, consider it active
let start_time = Utc::now();
return Ok(Some((session_name.clone(), start_time)));
}
Err(e) => {
tracing::warn!(
"Failed to get session path for '{}': {}",
"Session '{}' no longer exists: {}",
session_name,
e
);
@@ -983,21 +974,6 @@ impl TemporalScheduler {
let start_time = Utc::now();
return Ok(Some((session_id, start_time)));
}
};
// If the session file was modified recently (within last 5 minutes),
// consider it as the current running session
if let Ok(metadata) = std::fs::metadata(&session_path) {
if let Ok(modified) = metadata.modified() {
let modified_dt: DateTime<Utc> = modified.into();
let now = Utc::now();
let time_diff = now.signed_duration_since(modified_dt);
if time_diff.num_minutes() < 5 {
// This looks like an active session
return Ok(Some((session_name.clone(), modified_dt)));
}
}
}
}
@@ -1023,10 +999,9 @@ impl TemporalScheduler {
async fn make_request(&self, request: JobRequest) -> Result<JobResponse, SchedulerError> {
let url = format!("{}/jobs", self.service_url);
tracing::info!(
info!(
"TemporalScheduler: Making HTTP request to {} with action '{}'",
url,
request.action
url, request.action
);
let response = self
@@ -1212,7 +1187,7 @@ impl SchedulerTrait for TemporalScheduler {
&self,
sched_id: &str,
limit: usize,
) -> Result<Vec<(String, SessionMetadata)>, SchedulerError> {
) -> Result<Vec<(String, Session)>, SchedulerError> {
self.sessions(sched_id, limit).await
}
+5 -31
View File
@@ -364,10 +364,9 @@ mod schedule_tool_tests {
use goose::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME;
use goose::scheduler::{ScheduledJob, SchedulerError};
use goose::scheduler_trait::SchedulerTrait;
use goose::session::storage::SessionMetadata;
use goose::session::Session;
use std::sync::Arc;
// Mock scheduler for testing
struct MockScheduler {
jobs: tokio::sync::Mutex<Vec<ScheduledJob>>,
}
@@ -419,7 +418,7 @@ mod schedule_tool_tests {
&self,
_sched_id: &str,
_limit: usize,
) -> Result<Vec<(String, SessionMetadata)>, SchedulerError> {
) -> Result<Vec<(String, Session)>, SchedulerError> {
Ok(vec![])
}
@@ -853,7 +852,7 @@ mod final_output_tool_tests {
mod retry_tests {
use super::*;
use async_trait::async_trait;
use goose::agents::types::{RetryConfig, SessionConfig, SuccessCheck};
use goose::agents::types::{RetryConfig, SuccessCheck};
use goose::conversation::message::Message;
use goose::conversation::Conversation;
use goose::model::ModelConfig;
@@ -939,21 +938,10 @@ mod retry_tests {
"Valid config should pass validation"
);
let session_config = SessionConfig {
id: goose::session::Identifier::Name("test-retry".to_string()),
working_dir: std::env::current_dir()?,
schedule_id: None,
execution_mode: None,
max_turns: None,
retry_config: Some(retry_config),
};
let conversation =
Conversation::new(vec![Message::user().with_text("Complete this task")]).unwrap();
let reply_stream = agent
.reply(conversation, Some(session_config), None)
.await?;
let reply_stream = agent.reply(conversation, None, None).await?;
tokio::pin!(reply_stream);
let mut responses = Vec::new();
@@ -1051,10 +1039,8 @@ mod max_turns_tests {
use goose::model::ModelConfig;
use goose::providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage};
use goose::providers::errors::ProviderError;
use goose::session::storage::Identifier;
use mcp_core::tool::ToolCall;
use rmcp::model::Tool;
use std::path::PathBuf;
struct MockToolProvider {}
@@ -1116,21 +1102,9 @@ mod max_turns_tests {
let provider = Arc::new(MockToolProvider::new());
agent.update_provider(provider).await?;
// The mock provider will call a non-existent tool, which will fail and allow the loop to continue
// Create session config with max_turns = 1
let session_config = goose::agents::SessionConfig {
id: Identifier::Name("test_session".to_string()),
working_dir: PathBuf::from("/tmp"),
schedule_id: None,
execution_mode: None,
max_turns: Some(1),
retry_config: None,
};
let conversation = Conversation::new(vec![Message::user().with_text("Hello")]).unwrap();
let reply_stream = agent
.reply(conversation, Some(session_config), None)
.await?;
let reply_stream = agent.reply(conversation, None, None).await?;
tokio::pin!(reply_stream);
let mut responses = Vec::new();
-79
View File
@@ -757,85 +757,6 @@ async fn test_schedule_tool_sessions_action_empty() {
assert!(calls.contains(&"sessions".to_string()));
}
#[tokio::test]
async fn test_schedule_tool_session_content_action() {
let (agent, _) = ScheduleToolTestBuilder::new().build().await;
// Test with a non-existent session
let arguments = json!({
"action": "session_content",
"session_id": "non_existent_session"
});
let result = agent
.handle_schedule_management(arguments, "test_req".to_string())
.await;
assert!(result.is_err());
if let Err(err) = result {
assert!(err
.message
.contains("Session 'non_existent_session' not found"));
} else {
panic!("Expected ExecutionError");
}
}
#[tokio::test]
async fn test_schedule_tool_session_content_action_with_real_session() {
let (agent, _) = ScheduleToolTestBuilder::new().build().await;
// Create a temporary session file in the proper session directory
let session_dir = goose::session::storage::ensure_session_dir().unwrap();
let session_id = "test_session_real";
let session_path = session_dir.join(format!("{}.jsonl", session_id));
// Create test metadata and messages
let metadata = create_test_session_metadata(2, "/tmp");
let messages = goose::conversation::Conversation::new_unvalidated(vec![
goose::conversation::message::Message::user().with_text("Hello"),
goose::conversation::message::Message::assistant().with_text("Hi there!"),
]);
// Save the session file
goose::session::storage::save_messages_with_metadata(&session_path, &metadata, &messages)
.unwrap();
// Test the session_content action
let arguments = json!({
"action": "session_content",
"session_id": session_id
});
let result = agent
.handle_schedule_management(arguments, "test_req".to_string())
.await;
// Clean up the test session file
let _ = std::fs::remove_file(&session_path);
// Verify the result
assert!(result.is_ok());
if let Ok(content) = result {
assert_eq!(content.len(), 1);
if let Some(text_content) = content[0].as_text() {
assert!(text_content
.text
.contains("Session 'test_session_real' Content:"));
assert!(text_content.text.contains("Metadata:"));
assert!(text_content.text.contains("Messages:"));
assert!(text_content.text.contains("Hello"));
assert!(text_content.text.contains("Hi there!"));
assert!(text_content.text.contains("Test session"));
} else {
panic!("Expected text content");
}
} else {
panic!("Expected successful result");
}
}
#[tokio::test]
async fn test_schedule_tool_session_content_action_missing_session_id() {
let (agent, _) = ScheduleToolTestBuilder::new().build().await;
+12 -13
View File
@@ -12,7 +12,7 @@ use tokio::sync::Mutex;
use goose::agents::Agent;
use goose::scheduler::{ScheduledJob, SchedulerError};
use goose::scheduler_trait::SchedulerTrait;
use goose::session::storage::SessionMetadata;
use goose::session::Session;
#[derive(Debug, Clone)]
pub enum MockBehavior {
@@ -30,7 +30,7 @@ pub struct ConfigurableMockScheduler {
call_log: Arc<Mutex<Vec<String>>>,
behaviors: Arc<Mutex<HashMap<String, MockBehavior>>>,
#[allow(clippy::type_complexity)]
sessions_data: Arc<Mutex<HashMap<String, Vec<(String, SessionMetadata)>>>>,
sessions_data: Arc<Mutex<HashMap<String, Vec<(String, Session)>>>>,
}
#[allow(dead_code)]
@@ -184,7 +184,7 @@ impl SchedulerTrait for ConfigurableMockScheduler {
&self,
sched_id: &str,
limit: usize,
) -> Result<Vec<(String, SessionMetadata)>, SchedulerError> {
) -> Result<Vec<(String, Session)>, SchedulerError> {
self.log_call("sessions").await;
match self.get_behavior("sessions").await {
@@ -362,11 +362,7 @@ impl ScheduleToolTestBuilder {
self
}
pub async fn with_sessions_data(
self,
job_id: &str,
sessions: Vec<(String, SessionMetadata)>,
) -> Self {
pub async fn with_sessions_data(self, job_id: &str, sessions: Vec<(String, Session)>) -> Self {
{
let mut sessions_data = self.scheduler.sessions_data.lock().await;
sessions_data.insert(job_id.to_string(), sessions);
@@ -381,13 +377,14 @@ impl ScheduleToolTestBuilder {
}
}
// Helper function to create test session metadata
pub fn create_test_session_metadata(message_count: usize, working_dir: &str) -> SessionMetadata {
SessionMetadata {
message_count,
pub fn create_test_session_metadata(message_count: usize, working_dir: &str) -> Session {
Session {
id: "".to_string(),
working_dir: PathBuf::from(working_dir),
description: "Test session".to_string(),
created_at: "".to_string(),
schedule_id: Some("test_job".to_string()),
recipe: None,
total_tokens: Some(100),
input_tokens: Some(50),
output_tokens: Some(50),
@@ -395,6 +392,8 @@ pub fn create_test_session_metadata(message_count: usize, working_dir: &str) ->
accumulated_input_tokens: Some(50),
accumulated_output_tokens: Some(50),
extension_data: Default::default(),
recipe: None,
updated_at: "".to_string(),
conversation: None,
message_count,
}
}
@@ -1,496 +0,0 @@
use futures::StreamExt;
use goose::agents::types::SessionConfig;
use goose::agents::{Agent, AgentEvent};
use goose::conversation::message::Message;
use goose::conversation::Conversation;
use goose::model::ModelConfig;
use goose::providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage};
use goose::providers::errors::ProviderError;
use goose::session;
use goose::session::storage::SessionMetadata;
use rmcp::model::Tool;
use std::sync::Arc;
use tempfile::TempDir;
use uuid::Uuid;
// Mock provider implementation for testing
struct MockProvider {
model_config: ModelConfig,
}
impl MockProvider {
fn new() -> Self {
Self {
model_config: ModelConfig::new_or_fail("mock-model"),
}
}
}
#[async_trait::async_trait]
impl Provider for MockProvider {
fn metadata() -> ProviderMetadata
where
Self: Sized,
{
ProviderMetadata::new(
"mock",
"Mock Provider",
"A mock provider for testing",
"mock-model",
vec!["mock-model"],
"https://example.com",
vec![],
)
}
async fn complete(
&self,
_system: &str,
_messages: &[Message],
_tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
// Return a simple mock response
Ok((
Message::assistant().with_text("Mock response"),
ProviderUsage::new(
"mock-model".to_string(),
Usage::new(Some(10), Some(20), Some(30)),
),
))
}
async fn complete_with_model(
&self,
_model_config: &ModelConfig,
_system: &str,
_messages: &[Message],
_tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
// Return a simple mock response
Ok((
Message::assistant().with_text("Mock response"),
ProviderUsage::new(
"mock-model".to_string(),
Usage::new(Some(10), Some(20), Some(30)),
),
))
}
fn get_model_config(&self) -> ModelConfig {
self.model_config.clone()
}
async fn stream(
&self,
_system: &str,
_messages: &[Message],
_tools: &[Tool],
) -> Result<goose::providers::base::MessageStream, ProviderError> {
// Return a simple mock stream
let message = Message::assistant().with_text("Mock stream response");
let usage = ProviderUsage::new(
"mock-model".to_string(),
Usage::new(Some(10), Some(20), Some(30)),
);
Ok(goose::providers::base::stream_from_single_message(
message, usage,
))
}
fn supports_streaming(&self) -> bool {
true
}
async fn generate_session_name(
&self,
_messages: &Conversation,
) -> Result<String, ProviderError> {
Ok("Mock session description".to_string())
}
}
async fn create_test_session_dir() -> TempDir {
TempDir::new().unwrap()
}
async fn create_test_agent_with_mock_provider() -> Agent {
let agent = Agent::new();
let mock_provider = Arc::new(MockProvider::new());
agent.update_provider(mock_provider).await.unwrap();
agent
}
#[tokio::test]
async fn test_todo_add_persists_to_session() {
let temp_dir = create_test_session_dir().await;
let session_id = session::Identifier::Name(format!("test_session_{}", uuid::Uuid::new_v4()));
let agent = create_test_agent_with_mock_provider().await;
// Create a conversation with a TODO add request
let messages =
vec![Message::user().with_text("Add these tasks to my todo list: Buy milk, Call dentist")];
let conversation = Conversation::new(messages).unwrap();
let session_config = SessionConfig {
id: session_id.clone(),
working_dir: temp_dir.path().to_path_buf(),
schedule_id: None,
max_turns: Some(10),
execution_mode: Some("auto".to_string()),
retry_config: None,
};
// Process the conversation
let mut stream = agent
.reply(conversation, Some(session_config.clone()), None)
.await
.unwrap();
// Collect all events
while let Some(event) = stream.next().await {
if let Ok(_event) = event {
// Process events
}
}
// Verify TODO was persisted to session
let session_path = goose::session::storage::get_path(session_id).unwrap();
let metadata = goose::session::storage::read_metadata(&session_path).unwrap();
// Since we're using a mock provider, we can't test the actual TODO content
// but we can verify the metadata structure is correct
assert!(
metadata.extension_data.extension_states.is_empty()
|| !metadata.extension_data.extension_states.is_empty()
);
}
#[tokio::test]
async fn test_todo_list_reads_from_session() {
let temp_dir = create_test_session_dir().await;
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let agent = create_test_agent_with_mock_provider().await;
// Pre-populate session with TODO content
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let mut metadata = SessionMetadata::default();
use goose::session::extension_data::{ExtensionState, TodoState};
let todo_state = TodoState::new("- Task 1\n- Task 2\n- Task 3".to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
// Create a conversation requesting TODO list
let messages = vec![Message::user().with_text("Show me my todo list")];
let conversation = Conversation::new(messages).unwrap();
let session_config = SessionConfig {
id: session_id.clone(),
working_dir: temp_dir.path().to_path_buf(),
schedule_id: None,
max_turns: Some(10),
execution_mode: Some("auto".to_string()),
retry_config: None,
};
// Process the conversation
let mut stream = agent
.reply(conversation, Some(session_config), None)
.await
.unwrap();
// Collect all events
while let Some(event) = stream.next().await {
if let Ok(AgentEvent::Message(msg)) = event {
let _text = msg.as_concat_text();
// With mock provider, we can't verify the actual content
}
}
// Verify the TODO content is still in session
let metadata_after = goose::session::storage::read_metadata(&session_path).unwrap();
let todo_state_after = TodoState::from_extension_data(&metadata_after.extension_data);
assert!(todo_state_after.is_some());
assert_eq!(
todo_state_after.unwrap().content,
"- Task 1\n- Task 2\n- Task 3".to_string()
);
}
#[tokio::test]
async fn test_todo_isolation_between_sessions() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session1_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let session2_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
// Add TODO to session1
let session1_path = goose::session::storage::get_path(session1_id.clone()).unwrap();
let mut metadata1 = SessionMetadata::default();
let todo_state1 = TodoState::new("Session 1 tasks".to_string());
todo_state1
.to_extension_data(&mut metadata1.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session1_path, &metadata1)
.await
.unwrap();
// Add different TODO to session2
let session2_path = goose::session::storage::get_path(session2_id.clone()).unwrap();
let mut metadata2 = SessionMetadata::default();
let todo_state2 = TodoState::new("Session 2 tasks".to_string());
todo_state2
.to_extension_data(&mut metadata2.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session2_path, &metadata2)
.await
.unwrap();
// Verify isolation
let metadata1_read = goose::session::storage::read_metadata(&session1_path).unwrap();
let metadata2_read = goose::session::storage::read_metadata(&session2_path).unwrap();
let todo1 = TodoState::from_extension_data(&metadata1_read.extension_data).unwrap();
let todo2 = TodoState::from_extension_data(&metadata2_read.extension_data).unwrap();
assert_eq!(todo1.content, "Session 1 tasks");
assert_eq!(todo2.content, "Session 2 tasks");
}
#[tokio::test]
async fn test_todo_clear_removes_from_session() {
use goose::session::extension_data::{ExtensionState, TodoState};
let temp_dir = create_test_session_dir().await;
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let agent = create_test_agent_with_mock_provider().await;
// Pre-populate session with TODO content
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let mut metadata = SessionMetadata::default();
let todo_state = TodoState::new("- Task to clear".to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
// Create a conversation to clear TODO
let messages = vec![Message::user().with_text("Clear my entire todo list")];
let conversation = Conversation::new(messages).unwrap();
let session_config = SessionConfig {
id: session_id.clone(),
working_dir: temp_dir.path().to_path_buf(),
schedule_id: None,
max_turns: Some(10),
execution_mode: Some("auto".to_string()),
retry_config: None,
};
// Process the conversation
let mut stream = agent
.reply(conversation, Some(session_config), None)
.await
.unwrap();
// Consume the stream
while (stream.next().await).is_some() {}
// With mock provider, the TODO won't actually be cleared via tool calls
// but we can verify the structure is correct
let metadata_after = goose::session::storage::read_metadata(&session_path).unwrap();
let todo_state_after = TodoState::from_extension_data(&metadata_after.extension_data);
assert!(todo_state_after.is_some()); // Will still have the original content with mock
}
#[tokio::test]
async fn test_todo_persistence_across_agent_instances() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
// First agent instance adds TODO
{
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let mut metadata = SessionMetadata::default();
let todo_state = TodoState::new("Persistent task".to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
}
// Second agent instance reads TODO
{
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let metadata = goose::session::storage::read_metadata(&session_path).unwrap();
let todo_state = TodoState::from_extension_data(&metadata.extension_data).unwrap();
assert_eq!(todo_state.content, "Persistent task");
}
}
#[tokio::test]
async fn test_todo_max_chars_limit() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
// Set a small limit for testing
std::env::set_var("GOOSE_TODO_MAX_CHARS", "50");
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let mut metadata = SessionMetadata::default();
// Try to set content that exceeds the limit
let long_content = "x".repeat(100);
let todo_state = TodoState::new(long_content.clone());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
// This should succeed at the storage level (storage doesn't enforce limits)
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
// But when the agent tries to write through the TODO tool, it should enforce the limit
// This would be tested through the agent's dispatch_todo_tool_with_session method
// Clean up
std::env::remove_var("GOOSE_TODO_MAX_CHARS");
}
#[tokio::test]
async fn test_todo_with_special_characters() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let mut metadata = SessionMetadata::default();
// Test with various special characters
let special_content = r#"
- Task with "quotes"
- Task with 'single quotes'
- Task with emoji 🎉
- Task with unicode:
- Task with newline
continuation
- Task with tab separation
"#;
let todo_state = TodoState::new(special_content.to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
// Read back and verify
let metadata_read = goose::session::storage::read_metadata(&session_path).unwrap();
let todo_state_read = TodoState::from_extension_data(&metadata_read.extension_data).unwrap();
assert_eq!(todo_state_read.content, special_content);
}
#[tokio::test]
async fn test_todo_concurrent_access() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
// Spawn multiple concurrent TODO operations
let mut handles = vec![];
for i in 0..5 {
let session_id_clone = session_id.clone();
let handle = tokio::spawn(async move {
let session_path = goose::session::storage::get_path(session_id_clone).unwrap();
let mut metadata = goose::session::storage::read_metadata(&session_path)
.unwrap_or_else(|_| SessionMetadata::default());
let current_content = TodoState::from_extension_data(&metadata.extension_data)
.map(|t| t.content)
.unwrap_or_default();
let new_todo = TodoState::new(format!("{}\n- Task {}", current_content, i));
new_todo
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata).await
});
handles.push(handle);
}
// Wait for all operations to complete
for handle in handles {
handle.await.unwrap().unwrap();
}
// Verify final state contains at least one task
let session_path = goose::session::storage::get_path(session_id).unwrap();
let metadata = goose::session::storage::read_metadata(&session_path).unwrap();
let todo_state = TodoState::from_extension_data(&metadata.extension_data).unwrap();
// Should contain at least one task (concurrent writes may overwrite)
assert!(todo_state.content.contains("Task"));
}
#[tokio::test]
async fn test_todo_empty_session_returns_empty() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
let metadata = goose::session::storage::read_metadata(&session_path)
.unwrap_or_else(|_| SessionMetadata::default());
let todo_state = TodoState::from_extension_data(&metadata.extension_data);
assert!(todo_state.is_none() || todo_state.unwrap().content.is_empty());
}
#[tokio::test]
async fn test_todo_update_preserves_other_metadata() {
use goose::session::extension_data::{ExtensionState, TodoState};
let session_id = session::Identifier::Name(format!("test_session_{}", Uuid::new_v4()));
let session_path = goose::session::storage::get_path(session_id.clone()).unwrap();
// Set initial metadata with various fields
let mut metadata = SessionMetadata::default();
#[allow(clippy::field_reassign_with_default)]
{
metadata.message_count = 5;
metadata.description = "Test session".to_string();
metadata.total_tokens = Some(1000);
}
let todo_state = TodoState::new("Initial TODO".to_string());
todo_state
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
// Update only TODO content
let todo_state_updated = TodoState::new("Updated TODO".to_string());
todo_state_updated
.to_extension_data(&mut metadata.extension_data)
.unwrap();
goose::session::storage::update_metadata(&session_path, &metadata)
.await
.unwrap();
// Verify other fields are preserved
let metadata_read = goose::session::storage::read_metadata(&session_path).unwrap();
assert_eq!(metadata_read.message_count, 5);
assert_eq!(metadata_read.description, "Test session");
assert_eq!(metadata_read.total_tokens, Some(1000));
let todo_state_read = TodoState::from_extension_data(&metadata_read.extension_data).unwrap();
assert_eq!(todo_state_read.content, "Updated TODO");
}
+238 -143
View File
@@ -107,7 +107,7 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/StartAgentResponse"
"$ref": "#/components/schemas/Session"
}
}
}
@@ -185,7 +185,7 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/StartAgentResponse"
"$ref": "#/components/schemas/Session"
}
}
}
@@ -1500,12 +1500,43 @@
]
}
},
"/sessions/insights": {
"get": {
"tags": [
"Session Management"
],
"operationId": "get_session_insights",
"responses": {
"200": {
"description": "Session insights retrieved successfully",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/SessionInsights"
}
}
}
},
"401": {
"description": "Unauthorized - Invalid or missing API key"
},
"500": {
"description": "Internal server error"
}
},
"security": [
{
"api_key": []
}
]
}
},
"/sessions/{session_id}": {
"get": {
"tags": [
"Session Management"
],
"operationId": "get_session_history",
"operationId": "get_session",
"parameters": [
{
"name": "session_id",
@@ -1523,7 +1554,7 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/SessionHistoryResponse"
"$ref": "#/components/schemas/Session"
}
}
}
@@ -1543,6 +1574,93 @@
"api_key": []
}
]
},
"delete": {
"tags": [
"Session Management"
],
"operationId": "delete_session",
"parameters": [
{
"name": "session_id",
"in": "path",
"description": "Unique identifier for the session",
"required": true,
"schema": {
"type": "string"
}
}
],
"responses": {
"200": {
"description": "Session deleted successfully"
},
"401": {
"description": "Unauthorized - Invalid or missing API key"
},
"404": {
"description": "Session not found"
},
"500": {
"description": "Internal server error"
}
},
"security": [
{
"api_key": []
}
]
}
},
"/sessions/{session_id}/description": {
"put": {
"tags": [
"Session Management"
],
"operationId": "update_session_description",
"parameters": [
{
"name": "session_id",
"in": "path",
"description": "Unique identifier for the session",
"required": true,
"schema": {
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/UpdateSessionDescriptionRequest"
}
}
},
"required": true
},
"responses": {
"200": {
"description": "Session description updated successfully"
},
"400": {
"description": "Bad request - Description too long (max 200 characters)"
},
"401": {
"description": "Unauthorized - Invalid or missing API key"
},
"404": {
"description": "Session not found"
},
"500": {
"description": "Internal server error"
}
},
"security": [
{
"api_key": []
}
]
}
},
"/status": {
@@ -1815,6 +1933,12 @@
}
}
},
"Conversation": {
"type": "array",
"items": {
"$ref": "#/components/schemas/Message"
}
},
"CreateCustomProviderRequest": {
"type": "object",
"required": [
@@ -3426,6 +3550,92 @@
}
}
},
"Session": {
"type": "object",
"required": [
"id",
"working_dir",
"description",
"created_at",
"updated_at",
"extension_data",
"message_count"
],
"properties": {
"accumulated_input_tokens": {
"type": "integer",
"format": "int32",
"nullable": true
},
"accumulated_output_tokens": {
"type": "integer",
"format": "int32",
"nullable": true
},
"accumulated_total_tokens": {
"type": "integer",
"format": "int32",
"nullable": true
},
"conversation": {
"allOf": [
{
"$ref": "#/components/schemas/Conversation"
}
],
"nullable": true
},
"created_at": {
"type": "string"
},
"description": {
"type": "string"
},
"extension_data": {
"$ref": "#/components/schemas/ExtensionData"
},
"id": {
"type": "string"
},
"input_tokens": {
"type": "integer",
"format": "int32",
"nullable": true
},
"message_count": {
"type": "integer",
"minimum": 0
},
"output_tokens": {
"type": "integer",
"format": "int32",
"nullable": true
},
"recipe": {
"allOf": [
{
"$ref": "#/components/schemas/Recipe"
}
],
"nullable": true
},
"schedule_id": {
"type": "string",
"nullable": true
},
"total_tokens": {
"type": "integer",
"format": "int32",
"nullable": true
},
"updated_at": {
"type": "string"
},
"working_dir": {
"type": "string"
}
}
},
"SessionConfigRequest": {
"type": "object",
"required": [
@@ -3507,50 +3717,22 @@
}
}
},
"SessionHistoryResponse": {
"SessionInsights": {
"type": "object",
"required": [
"sessionId",
"metadata",
"messages"
"totalSessions",
"totalTokens"
],
"properties": {
"messages": {
"type": "array",
"items": {
"$ref": "#/components/schemas/Message"
},
"description": "List of messages in the session conversation"
"totalSessions": {
"type": "integer",
"description": "Total number of sessions",
"minimum": 0
},
"metadata": {
"$ref": "#/components/schemas/SessionMetadata"
},
"sessionId": {
"type": "string",
"description": "Unique identifier for the session"
}
}
},
"SessionInfo": {
"type": "object",
"required": [
"id",
"path",
"modified",
"metadata"
],
"properties": {
"id": {
"type": "string"
},
"metadata": {
"$ref": "#/components/schemas/SessionMetadata"
},
"modified": {
"type": "string"
},
"path": {
"type": "string"
"totalTokens": {
"type": "integer",
"format": "int64",
"description": "Total tokens used across all sessions"
}
}
},
@@ -3563,89 +3745,12 @@
"sessions": {
"type": "array",
"items": {
"$ref": "#/components/schemas/SessionInfo"
"$ref": "#/components/schemas/Session"
},
"description": "List of available session information objects"
}
}
},
"SessionMetadata": {
"type": "object",
"description": "Metadata for a session, stored as the first line in the session file",
"required": [
"working_dir",
"description",
"message_count"
],
"properties": {
"accumulated_input_tokens": {
"type": "integer",
"format": "int32",
"description": "The number of input tokens used in the session. Accumulated across all messages.",
"nullable": true
},
"accumulated_output_tokens": {
"type": "integer",
"format": "int32",
"description": "The number of output tokens used in the session. Accumulated across all messages.",
"nullable": true
},
"accumulated_total_tokens": {
"type": "integer",
"format": "int32",
"description": "The total number of tokens used in the session. Accumulated across all messages (useful for tracking cost over an entire session).",
"nullable": true
},
"description": {
"type": "string",
"description": "A short description of the session, typically 3 words or less"
},
"extension_data": {
"$ref": "#/components/schemas/ExtensionData"
},
"input_tokens": {
"type": "integer",
"format": "int32",
"description": "The number of input tokens used in the session. Retrieved from the provider's last usage.",
"nullable": true
},
"message_count": {
"type": "integer",
"description": "Number of messages in the session",
"minimum": 0
},
"output_tokens": {
"type": "integer",
"format": "int32",
"description": "The number of output tokens used in the session. Retrieved from the provider's last usage.",
"nullable": true
},
"recipe": {
"allOf": [
{
"$ref": "#/components/schemas/Recipe"
}
],
"nullable": true
},
"schedule_id": {
"type": "string",
"description": "ID of the schedule that triggered this session, if any",
"nullable": true
},
"total_tokens": {
"type": "integer",
"format": "int32",
"description": "The total number of tokens used in the session. Retrieved from the provider's last usage.",
"nullable": true
},
"working_dir": {
"type": "string",
"description": "Working directory for the session",
"example": "/home/user/sessions/session1"
}
}
},
"SessionsQuery": {
"type": "object",
"properties": {
@@ -3708,28 +3813,6 @@
}
}
},
"StartAgentResponse": {
"type": "object",
"required": [
"session_id",
"metadata",
"messages"
],
"properties": {
"messages": {
"type": "array",
"items": {
"$ref": "#/components/schemas/Message"
}
},
"metadata": {
"$ref": "#/components/schemas/SessionMetadata"
},
"session_id": {
"type": "string"
}
}
},
"SubRecipe": {
"type": "object",
"required": [
@@ -4029,6 +4112,18 @@
}
}
},
"UpdateSessionDescriptionRequest": {
"type": "object",
"required": [
"description"
],
"properties": {
"description": {
"type": "string",
"description": "Updated description (name) for the session (max 200 characters)"
}
}
},
"UpsertConfigQuery": {
"type": "object",
"required": [
+1 -2
View File
@@ -15,7 +15,6 @@ import { ExtensionInstallModal } from './components/ExtensionInstallModal';
import { ToastContainer } from 'react-toastify';
import { GoosehintsModal } from './components/GoosehintsModal';
import AnnouncementModal from './components/AnnouncementModal';
import { generateSessionId } from './sessions';
import ProviderGuard from './components/ProviderGuard';
import { ChatType } from './types/chat';
@@ -321,7 +320,7 @@ export function AppInner() {
const [_searchParams, setSearchParams] = useSearchParams();
const [chat, setChat] = useState<ChatType>({
sessionId: generateSessionId(),
sessionId: '',
title: 'Pair Chat',
messages: [],
messageHistoryIndex: 0,
+28 -3
View File
@@ -1,7 +1,7 @@
// This file is auto-generated by @hey-api/openapi-ts
import type { Options as ClientOptions, TDataShape, Client } from './client';
import type { AddSubRecipesData, AddSubRecipesResponses, AddSubRecipesErrors, ExtendPromptData, ExtendPromptResponses, ExtendPromptErrors, ResumeAgentData, ResumeAgentResponses, ResumeAgentErrors, UpdateSessionConfigData, UpdateSessionConfigResponses, UpdateSessionConfigErrors, StartAgentData, StartAgentResponses, StartAgentErrors, GetToolsData, GetToolsResponses, GetToolsErrors, UpdateAgentProviderData, UpdateAgentProviderResponses, UpdateAgentProviderErrors, UpdateRouterToolSelectorData, UpdateRouterToolSelectorResponses, UpdateRouterToolSelectorErrors, ReadAllConfigData, ReadAllConfigResponses, BackupConfigData, BackupConfigResponses, BackupConfigErrors, CreateCustomProviderData, CreateCustomProviderResponses, CreateCustomProviderErrors, RemoveCustomProviderData, RemoveCustomProviderResponses, RemoveCustomProviderErrors, GetExtensionsData, GetExtensionsResponses, GetExtensionsErrors, AddExtensionData, AddExtensionResponses, AddExtensionErrors, RemoveExtensionData, RemoveExtensionResponses, RemoveExtensionErrors, InitConfigData, InitConfigResponses, InitConfigErrors, UpsertPermissionsData, UpsertPermissionsResponses, UpsertPermissionsErrors, ProvidersData, ProvidersResponses, GetProviderModelsData, GetProviderModelsResponses, GetProviderModelsErrors, ReadConfigData, ReadConfigResponses, ReadConfigErrors, RecoverConfigData, RecoverConfigResponses, RecoverConfigErrors, RemoveConfigData, RemoveConfigResponses, RemoveConfigErrors, UpsertConfigData, UpsertConfigResponses, UpsertConfigErrors, ValidateConfigData, ValidateConfigResponses, ValidateConfigErrors, ConfirmPermissionData, ConfirmPermissionResponses, ConfirmPermissionErrors, ManageContextData, ManageContextResponses, ManageContextErrors, StartOpenrouterSetupData, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponses, CreateRecipeData, CreateRecipeResponses, CreateRecipeErrors, DecodeRecipeData, DecodeRecipeResponses, DecodeRecipeErrors, DeleteRecipeData, DeleteRecipeResponses, DeleteRecipeErrors, EncodeRecipeData, EncodeRecipeResponses, EncodeRecipeErrors, ListRecipesData, ListRecipesResponses, ListRecipesErrors, ScanRecipeData, ScanRecipeResponses, CreateScheduleData, CreateScheduleResponses, CreateScheduleErrors, DeleteScheduleData, DeleteScheduleResponses, DeleteScheduleErrors, ListSchedulesData, ListSchedulesResponses, ListSchedulesErrors, UpdateScheduleData, UpdateScheduleResponses, UpdateScheduleErrors, InspectRunningJobData, InspectRunningJobResponses, InspectRunningJobErrors, KillRunningJobData, KillRunningJobResponses, PauseScheduleData, PauseScheduleResponses, PauseScheduleErrors, RunNowHandlerData, RunNowHandlerResponses, RunNowHandlerErrors, SessionsHandlerData, SessionsHandlerResponses, SessionsHandlerErrors, UnpauseScheduleData, UnpauseScheduleResponses, UnpauseScheduleErrors, ListSessionsData, ListSessionsResponses, ListSessionsErrors, GetSessionHistoryData, GetSessionHistoryResponses, GetSessionHistoryErrors, StatusData, StatusResponses } from './types.gen';
import type { AddSubRecipesData, AddSubRecipesResponses, AddSubRecipesErrors, ExtendPromptData, ExtendPromptResponses, ExtendPromptErrors, ResumeAgentData, ResumeAgentResponses, ResumeAgentErrors, UpdateSessionConfigData, UpdateSessionConfigResponses, UpdateSessionConfigErrors, StartAgentData, StartAgentResponses, StartAgentErrors, GetToolsData, GetToolsResponses, GetToolsErrors, UpdateAgentProviderData, UpdateAgentProviderResponses, UpdateAgentProviderErrors, UpdateRouterToolSelectorData, UpdateRouterToolSelectorResponses, UpdateRouterToolSelectorErrors, ReadAllConfigData, ReadAllConfigResponses, BackupConfigData, BackupConfigResponses, BackupConfigErrors, CreateCustomProviderData, CreateCustomProviderResponses, CreateCustomProviderErrors, RemoveCustomProviderData, RemoveCustomProviderResponses, RemoveCustomProviderErrors, GetExtensionsData, GetExtensionsResponses, GetExtensionsErrors, AddExtensionData, AddExtensionResponses, AddExtensionErrors, RemoveExtensionData, RemoveExtensionResponses, RemoveExtensionErrors, InitConfigData, InitConfigResponses, InitConfigErrors, UpsertPermissionsData, UpsertPermissionsResponses, UpsertPermissionsErrors, ProvidersData, ProvidersResponses, GetProviderModelsData, GetProviderModelsResponses, GetProviderModelsErrors, ReadConfigData, ReadConfigResponses, ReadConfigErrors, RecoverConfigData, RecoverConfigResponses, RecoverConfigErrors, RemoveConfigData, RemoveConfigResponses, RemoveConfigErrors, UpsertConfigData, UpsertConfigResponses, UpsertConfigErrors, ValidateConfigData, ValidateConfigResponses, ValidateConfigErrors, ConfirmPermissionData, ConfirmPermissionResponses, ConfirmPermissionErrors, ManageContextData, ManageContextResponses, ManageContextErrors, StartOpenrouterSetupData, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponses, CreateRecipeData, CreateRecipeResponses, CreateRecipeErrors, DecodeRecipeData, DecodeRecipeResponses, DecodeRecipeErrors, DeleteRecipeData, DeleteRecipeResponses, DeleteRecipeErrors, EncodeRecipeData, EncodeRecipeResponses, EncodeRecipeErrors, ListRecipesData, ListRecipesResponses, ListRecipesErrors, ScanRecipeData, ScanRecipeResponses, CreateScheduleData, CreateScheduleResponses, CreateScheduleErrors, DeleteScheduleData, DeleteScheduleResponses, DeleteScheduleErrors, ListSchedulesData, ListSchedulesResponses, ListSchedulesErrors, UpdateScheduleData, UpdateScheduleResponses, UpdateScheduleErrors, InspectRunningJobData, InspectRunningJobResponses, InspectRunningJobErrors, KillRunningJobData, KillRunningJobResponses, PauseScheduleData, PauseScheduleResponses, PauseScheduleErrors, RunNowHandlerData, RunNowHandlerResponses, RunNowHandlerErrors, SessionsHandlerData, SessionsHandlerResponses, SessionsHandlerErrors, UnpauseScheduleData, UnpauseScheduleResponses, UnpauseScheduleErrors, ListSessionsData, ListSessionsResponses, ListSessionsErrors, GetSessionInsightsData, GetSessionInsightsResponses, GetSessionInsightsErrors, DeleteSessionData, DeleteSessionResponses, DeleteSessionErrors, GetSessionData, GetSessionResponses, GetSessionErrors, UpdateSessionDescriptionData, UpdateSessionDescriptionResponses, UpdateSessionDescriptionErrors, StatusData, StatusResponses } from './types.gen';
import { client as _heyApiClient } from './client.gen';
export type Options<TData extends TDataShape = TDataShape, ThrowOnError extends boolean = boolean> = ClientOptions<TData, ThrowOnError> & {
@@ -424,13 +424,38 @@ export const listSessions = <ThrowOnError extends boolean = false>(options?: Opt
});
};
export const getSessionHistory = <ThrowOnError extends boolean = false>(options: Options<GetSessionHistoryData, ThrowOnError>) => {
return (options.client ?? _heyApiClient).get<GetSessionHistoryResponses, GetSessionHistoryErrors, ThrowOnError>({
export const getSessionInsights = <ThrowOnError extends boolean = false>(options?: Options<GetSessionInsightsData, ThrowOnError>) => {
return (options?.client ?? _heyApiClient).get<GetSessionInsightsResponses, GetSessionInsightsErrors, ThrowOnError>({
url: '/sessions/insights',
...options
});
};
export const deleteSession = <ThrowOnError extends boolean = false>(options: Options<DeleteSessionData, ThrowOnError>) => {
return (options.client ?? _heyApiClient).delete<DeleteSessionResponses, DeleteSessionErrors, ThrowOnError>({
url: '/sessions/{session_id}',
...options
});
};
export const getSession = <ThrowOnError extends boolean = false>(options: Options<GetSessionData, ThrowOnError>) => {
return (options.client ?? _heyApiClient).get<GetSessionResponses, GetSessionErrors, ThrowOnError>({
url: '/sessions/{session_id}',
...options
});
};
export const updateSessionDescription = <ThrowOnError extends boolean = false>(options: Options<UpdateSessionDescriptionData, ThrowOnError>) => {
return (options.client ?? _heyApiClient).put<UpdateSessionDescriptionResponses, UpdateSessionDescriptionErrors, ThrowOnError>({
url: '/sessions/{session_id}/description',
...options,
headers: {
'Content-Type': 'application/json',
...options.headers
}
});
};
export const status = <ThrowOnError extends boolean = false>(options?: Options<StatusData, ThrowOnError>) => {
return (options?.client ?? _heyApiClient).get<StatusResponses, unknown, ThrowOnError>({
url: '/status',
+141 -76
View File
@@ -103,6 +103,8 @@ export type ContextManageResponse = {
tokenCounts: Array<number>;
};
export type Conversation = Array<Message>;
export type CreateCustomProviderRequest = {
api_key: string;
api_url: string;
@@ -698,6 +700,25 @@ export type ScheduledJob = {
source: string;
};
export type Session = {
accumulated_input_tokens?: number | null;
accumulated_output_tokens?: number | null;
accumulated_total_tokens?: number | null;
conversation?: Conversation | null;
created_at: string;
description: string;
extension_data: ExtensionData;
id: string;
input_tokens?: number | null;
message_count: number;
output_tokens?: number | null;
recipe?: Recipe | null;
schedule_id?: string | null;
total_tokens?: number | null;
updated_at: string;
working_dir: string;
};
export type SessionConfigRequest = {
response?: Response | null;
session_id: string;
@@ -718,78 +739,22 @@ export type SessionDisplayInfo = {
workingDir: string;
};
export type SessionHistoryResponse = {
export type SessionInsights = {
/**
* List of messages in the session conversation
* Total number of sessions
*/
messages: Array<Message>;
metadata: SessionMetadata;
totalSessions: number;
/**
* Unique identifier for the session
* Total tokens used across all sessions
*/
sessionId: string;
};
export type SessionInfo = {
id: string;
metadata: SessionMetadata;
modified: string;
path: string;
totalTokens: number;
};
export type SessionListResponse = {
/**
* List of available session information objects
*/
sessions: Array<SessionInfo>;
};
/**
* Metadata for a session, stored as the first line in the session file
*/
export type SessionMetadata = {
/**
* The number of input tokens used in the session. Accumulated across all messages.
*/
accumulated_input_tokens?: number | null;
/**
* The number of output tokens used in the session. Accumulated across all messages.
*/
accumulated_output_tokens?: number | null;
/**
* The total number of tokens used in the session. Accumulated across all messages (useful for tracking cost over an entire session).
*/
accumulated_total_tokens?: number | null;
/**
* A short description of the session, typically 3 words or less
*/
description: string;
extension_data?: ExtensionData;
/**
* The number of input tokens used in the session. Retrieved from the provider's last usage.
*/
input_tokens?: number | null;
/**
* Number of messages in the session
*/
message_count: number;
/**
* The number of output tokens used in the session. Retrieved from the provider's last usage.
*/
output_tokens?: number | null;
recipe?: Recipe | null;
/**
* ID of the schedule that triggered this session, if any
*/
schedule_id?: string | null;
/**
* The total number of tokens used in the session. Retrieved from the provider's last usage.
*/
total_tokens?: number | null;
/**
* Working directory for the session
*/
working_dir: string;
sessions: Array<Session>;
};
export type SessionsQuery = {
@@ -812,12 +777,6 @@ export type StartAgentRequest = {
working_dir: string;
};
export type StartAgentResponse = {
messages: Array<Message>;
metadata: SessionMetadata;
session_id: string;
};
export type SubRecipe = {
description?: string | null;
name: string;
@@ -930,6 +889,13 @@ export type UpdateScheduleRequest = {
cron: string;
};
export type UpdateSessionDescriptionRequest = {
/**
* Updated description (name) for the session (max 200 characters)
*/
description: string;
};
export type UpsertConfigQuery = {
is_secret: boolean;
key: string;
@@ -1020,7 +986,7 @@ export type ResumeAgentResponses = {
/**
* Agent started successfully
*/
200: StartAgentResponse;
200: Session;
};
export type ResumeAgentResponse = ResumeAgentResponses[keyof ResumeAgentResponses];
@@ -1082,10 +1048,10 @@ export type StartAgentResponses = {
/**
* Agent started successfully
*/
200: StartAgentResponse;
200: Session;
};
export type StartAgentResponse2 = StartAgentResponses[keyof StartAgentResponses];
export type StartAgentResponse = StartAgentResponses[keyof StartAgentResponses];
export type GetToolsData = {
body?: never;
@@ -2144,7 +2110,34 @@ export type ListSessionsResponses = {
export type ListSessionsResponse = ListSessionsResponses[keyof ListSessionsResponses];
export type GetSessionHistoryData = {
export type GetSessionInsightsData = {
body?: never;
path?: never;
query?: never;
url: '/sessions/insights';
};
export type GetSessionInsightsErrors = {
/**
* Unauthorized - Invalid or missing API key
*/
401: unknown;
/**
* Internal server error
*/
500: unknown;
};
export type GetSessionInsightsResponses = {
/**
* Session insights retrieved successfully
*/
200: SessionInsights;
};
export type GetSessionInsightsResponse = GetSessionInsightsResponses[keyof GetSessionInsightsResponses];
export type DeleteSessionData = {
body?: never;
path: {
/**
@@ -2156,7 +2149,7 @@ export type GetSessionHistoryData = {
url: '/sessions/{session_id}';
};
export type GetSessionHistoryErrors = {
export type DeleteSessionErrors = {
/**
* Unauthorized - Invalid or missing API key
*/
@@ -2171,14 +2164,86 @@ export type GetSessionHistoryErrors = {
500: unknown;
};
export type GetSessionHistoryResponses = {
export type DeleteSessionResponses = {
/**
* Session deleted successfully
*/
200: unknown;
};
export type GetSessionData = {
body?: never;
path: {
/**
* Unique identifier for the session
*/
session_id: string;
};
query?: never;
url: '/sessions/{session_id}';
};
export type GetSessionErrors = {
/**
* Unauthorized - Invalid or missing API key
*/
401: unknown;
/**
* Session not found
*/
404: unknown;
/**
* Internal server error
*/
500: unknown;
};
export type GetSessionResponses = {
/**
* Session history retrieved successfully
*/
200: SessionHistoryResponse;
200: Session;
};
export type GetSessionHistoryResponse = GetSessionHistoryResponses[keyof GetSessionHistoryResponses];
export type GetSessionResponse = GetSessionResponses[keyof GetSessionResponses];
export type UpdateSessionDescriptionData = {
body: UpdateSessionDescriptionRequest;
path: {
/**
* Unique identifier for the session
*/
session_id: string;
};
query?: never;
url: '/sessions/{session_id}/description';
};
export type UpdateSessionDescriptionErrors = {
/**
* Bad request - Description too long (max 200 characters)
*/
400: unknown;
/**
* Unauthorized - Invalid or missing API key
*/
401: unknown;
/**
* Session not found
*/
404: unknown;
/**
* Internal server error
*/
500: unknown;
};
export type UpdateSessionDescriptionResponses = {
/**
* Session description updated successfully
*/
200: unknown;
};
export type StatusData = {
body?: never;
+1 -1
View File
@@ -265,7 +265,7 @@ function BaseChatContent({
sessionOutputTokens,
localInputTokens,
localOutputTokens,
sessionMetadata,
session: sessionMetadata,
});
useEffect(() => {
@@ -94,8 +94,7 @@ describe('ExtensionInstallModal', () => {
expect(screen.getAllByRole('button')).toHaveLength(3);
});
it("should handle i-ching-mcp-server as allowed command", async () => {
it('should handle i-ching-mcp-server as allowed command', async () => {
mockElectron.getAllowedExtensions.mockResolvedValue([]);
render(<ExtensionInstallModal addExtension={mockAddExtension} />);
@@ -103,13 +102,16 @@ describe('ExtensionInstallModal', () => {
const eventHandler = getAddExtensionEventHandler();
await act(async () => {
await eventHandler({}, "goose://extension?cmd=i-ching-mcp-server&id=i-ching&name=I%20Ching&description=I%20Ching%20divination");
await eventHandler(
{},
'goose://extension?cmd=i-ching-mcp-server&id=i-ching&name=I%20Ching&description=I%20Ching%20divination'
);
});
expect(screen.getByRole("dialog")).toBeInTheDocument();
expect(screen.getByText("Confirm Extension Installation")).toBeInTheDocument();
expect(screen.getByRole('dialog')).toBeInTheDocument();
expect(screen.getByText('Confirm Extension Installation')).toBeInTheDocument();
expect(screen.getByText(/I Ching extension/)).toBeInTheDocument();
expect(screen.getAllByRole("button")).toHaveLength(3);
expect(screen.getAllByRole('button')).toHaveLength(3);
});
it('should handle blocked extension', async () => {
mockElectron.getAllowedExtensions.mockResolvedValue(['uvx allowed-package']);
@@ -1,356 +0,0 @@
import React, { useEffect, useState, useCallback, useRef } from 'react';
import { Search, ChevronDown, Folder, Loader2 } from 'lucide-react';
import { fetchSessions, type Session } from '../../sessions';
import { Input } from '../ui/input';
import {
SidebarMenu,
SidebarMenuItem,
SidebarMenuButton,
SidebarGroup,
SidebarGroupLabel,
SidebarGroupContent,
} from '../ui/sidebar';
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from '../ui/collapsible';
import { useTextAnimator } from '../../hooks/use-text-animator';
interface SessionsSectionProps {
onSelectSession: (sessionId: string) => void;
refreshTrigger?: number;
}
interface GroupedSessions {
today: Session[];
yesterday: Session[];
older: { [key: string]: Session[] };
}
export const SessionsSection: React.FC<SessionsSectionProps> = ({
onSelectSession,
refreshTrigger,
}) => {
const [sessions, setSessions] = useState<Session[]>([]);
const [searchTerm, setSearchTerm] = useState('');
const [groupedSessions, setGroupedSessions] = useState<GroupedSessions>({
today: [],
yesterday: [],
older: {},
});
const [sessionsWithDescriptions, setSessionsWithDescriptions] = useState<Set<string>>(new Set());
const refreshTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
const groupSessions = useCallback((sessionsToGroup: Session[]) => {
const now = new Date();
const today = new Date(now.getFullYear(), now.getMonth(), now.getDate());
const yesterday = new Date(today);
yesterday.setDate(yesterday.getDate() - 1);
const grouped: GroupedSessions = {
today: [],
yesterday: [],
older: {},
};
sessionsToGroup.forEach((session) => {
const sessionDate = new Date(session.modified);
const sessionDateOnly = new Date(
sessionDate.getFullYear(),
sessionDate.getMonth(),
sessionDate.getDate()
);
if (sessionDateOnly.getTime() === today.getTime()) {
grouped.today.push(session);
} else if (sessionDateOnly.getTime() === yesterday.getTime()) {
grouped.yesterday.push(session);
} else {
const dateKey = sessionDateOnly.toISOString().split('T')[0];
if (!grouped.older[dateKey]) {
grouped.older[dateKey] = [];
}
grouped.older[dateKey].push(session);
}
});
// Sort older sessions by date (newest first)
const sortedOlder: { [key: string]: Session[] } = {};
Object.keys(grouped.older)
.sort()
.reverse()
.forEach((key) => {
sortedOlder[key] = grouped.older[key];
});
grouped.older = sortedOlder;
setGroupedSessions(grouped);
}, []);
const loadSessions = useCallback(async () => {
try {
const sessions = await fetchSessions();
setSessions(sessions);
groupSessions(sessions);
} catch (err) {
console.error('Failed to load sessions:', err);
setSessions([]);
setGroupedSessions({ today: [], yesterday: [], older: {} });
}
}, [groupSessions]);
// Debounced refresh function
const debouncedRefresh = useCallback(() => {
console.log('SessionsSection: Debounced refresh triggered');
// Clear any existing timeout
if (refreshTimeoutRef.current) {
window.clearTimeout(refreshTimeoutRef.current);
}
// Set new timeout - reduced to 200ms for faster response
refreshTimeoutRef.current = setTimeout(() => {
console.log('SessionsSection: Executing debounced refresh');
loadSessions();
refreshTimeoutRef.current = null;
}, 200);
}, [loadSessions]);
// Cleanup timeout on unmount
useEffect(() => {
return () => {
if (refreshTimeoutRef.current) {
window.clearTimeout(refreshTimeoutRef.current);
}
};
}, []);
useEffect(() => {
console.log('SessionsSection: Initial load');
loadSessions();
}, [loadSessions]);
// Add effect to refresh sessions when refreshTrigger changes
useEffect(() => {
if (refreshTrigger) {
console.log('SessionsSection: Refresh trigger changed, triggering refresh');
debouncedRefresh();
}
}, [refreshTrigger, debouncedRefresh]);
// Add effect to listen for session creation events
useEffect(() => {
const handleSessionCreated = () => {
console.log('SessionsSection: Session created event received');
debouncedRefresh();
};
const handleMessageStreamFinish = () => {
console.log('SessionsSection: Message stream finished event received');
// Always refresh when message stream finishes
debouncedRefresh();
};
// Listen for custom events that indicate a session was created
window.addEventListener('session-created', handleSessionCreated);
// Also listen for message stream finish events
window.addEventListener('message-stream-finished', handleMessageStreamFinish);
return () => {
window.removeEventListener('session-created', handleSessionCreated);
window.removeEventListener('message-stream-finished', handleMessageStreamFinish);
};
}, [debouncedRefresh]);
useEffect(() => {
if (searchTerm) {
const filtered = sessions.filter((session) =>
(session.metadata.description || session.id)
.toLowerCase()
.includes(searchTerm.toLowerCase())
);
groupSessions(filtered);
} else {
groupSessions(sessions);
}
}, [searchTerm, sessions, groupSessions]);
// Component for individual session items with loading and animation states
const SessionItem = ({ session }: { session: Session }) => {
const hasDescription =
session.metadata.description && session.metadata.description.trim() !== '';
const isNewSession = session.id.match(/^\d{8}_\d{6}$/);
const messageCount = session.metadata.message_count || 0;
// Show loading for new sessions with few messages and no description
// Only show loading for sessions created in the last 5 minutes
const sessionDate = new Date(session.modified);
const fiveMinutesAgo = new Date(Date.now() - 5 * 60 * 1000);
const isRecentSession = sessionDate > fiveMinutesAgo;
const shouldShowLoading =
!hasDescription && isNewSession && messageCount <= 2 && isRecentSession;
const [isAnimating, setIsAnimating] = useState(false);
// Use text animator only for sessions that need animation
const descriptionRef = useTextAnimator({
text: isAnimating ? session.metadata.description : '',
});
// Track when description becomes available and trigger animation
useEffect(() => {
if (hasDescription && !sessionsWithDescriptions.has(session.id)) {
setSessionsWithDescriptions((prev) => new Set(prev).add(session.id));
// Only animate for new sessions that were showing loading
if (shouldShowLoading) {
setIsAnimating(true);
}
}
}, [hasDescription, session.id, shouldShowLoading]);
const handleClick = () => {
console.log('SessionItem: Clicked on session:', session.id);
onSelectSession(session.id);
};
return (
<SidebarMenuItem key={session.id}>
<SidebarMenuButton
onClick={handleClick}
className="cursor-pointer w-56 transition-all duration-300 ease-in-out hover:bg-background-medium hover:shadow-sm rounded-xl text-text-muted hover:text-text-default h-fit flex items-start transform hover:scale-[1.02] active:scale-[0.98]"
>
<div className="flex flex-col w-full">
<div className="text-sm w-48 truncate mb-1 px-1 text-ellipsis text-text-default flex items-center gap-2">
{shouldShowLoading ? (
<div className="flex items-center gap-2 animate-in fade-in duration-300">
<Loader2 className="size-3 animate-spin text-text-default" />
<span className="text-text-default animate-pulse">Generating description...</span>
</div>
) : (
<span
ref={isAnimating ? descriptionRef : undefined}
className={`transition-all duration-300 ${isAnimating ? 'animate-in fade-in duration-300' : ''}`}
>
{hasDescription ? session.metadata.description : `Session ${session.id}`}
</span>
)}
</div>
<div className="text-xs w-48 truncate px-1 flex items-center gap-2 text-ellipsis transition-colors duration-300">
<Folder className="size-4 transition-transform duration-300 group-hover:scale-110" />
<span className="transition-all duration-300">{session.metadata.working_dir}</span>
</div>
</div>
</SidebarMenuButton>
</SidebarMenuItem>
);
};
const renderSessionGroup = (sessions: Session[], title: string, index: number) => {
if (sessions.length === 0) return null;
const isFirstTwoGroups = index < 2;
return (
<Collapsible defaultOpen={isFirstTwoGroups} className="group/collapsible">
<SidebarGroup>
<CollapsibleTrigger className="w-full">
<SidebarGroupLabel className="flex cursor-pointer items-center justify-between text-text-default hover:text-text-default h-12 pl-3 transition-all duration-200 rounded-lg">
<div className="flex min-w-0 items-center">
<span className="opacity-100 transition-all duration-300 text-xs font-medium">
{title}
</span>
</div>
<ChevronDown className="size-4 text-text-muted flex-shrink-0 opacity-100 transition-all duration-300 ease-in-out group-data-[state=open]/collapsible:rotate-180" />
</SidebarGroupLabel>
</CollapsibleTrigger>
<CollapsibleContent className="data-[state=open]:animate-collapsible-down data-[state=closed]:animate-collapsible-up overflow-hidden transition-all duration-300 ease-in-out">
<SidebarGroupContent>
<SidebarMenu className="mb-2 space-y-1">
{sessions.map((session, sessionIndex) => (
<div
key={session.id}
className="animate-in slide-in-from-left-2 fade-in duration-300"
style={{
animationDelay: `${sessionIndex * 50}ms`,
animationFillMode: 'both',
}}
>
<SessionItem session={session} />
</div>
))}
</SidebarMenu>
</SidebarGroupContent>
</CollapsibleContent>
</SidebarGroup>
</Collapsible>
);
};
return (
<Collapsible defaultOpen={false} className="group/collapsible rounded-xl">
<SidebarGroup className="px-1">
<CollapsibleTrigger className="w-full">
<SidebarGroupLabel className="flex cursor-pointer items-center py-6 justify-between text-text-default px-4 transition-all duration-200 hover:bg-background-default rounded-lg">
<div className="flex min-w-0 items-center">
<span className="text-sm">Sessions</span>
</div>
<ChevronDown className="size-4 text-text-muted flex-shrink-0 opacity-100 transition-all duration-300 ease-in-out group-data-[state=open]/collapsible:rotate-180" />
</SidebarGroupLabel>
</CollapsibleTrigger>
<CollapsibleContent className="data-[state=open]:animate-collapsible-down data-[state=closed]:animate-collapsible-up overflow-hidden transition-all duration-300 ease-in-out">
<SidebarGroupContent>
{/* Search Input */}
<div className="p-1 pb-2 animate-in slide-in-from-top-2 fade-in duration-300">
<div className="relative flex flex-row items-center gap-2">
<Search className="absolute top-2.5 left-2.5 size-4 text-muted-foreground" />
<Input
type="search"
placeholder="Search sessions..."
className="pl-8 transition-all duration-200 focus:ring-2 focus:ring-borderProminent"
value={searchTerm}
onChange={(e: React.ChangeEvent<HTMLInputElement>) =>
setSearchTerm(e.target.value)
}
/>
</div>
</div>
{/* Sessions Groups */}
<div className="space-y-2">
{(() => {
let groupIndex = 0;
const groups = [
{ sessions: groupedSessions.today, title: 'Today' },
{ sessions: groupedSessions.yesterday, title: 'Yesterday' },
...Object.entries(groupedSessions.older).map(([date, sessions]) => ({
sessions,
title: new Date(date).toLocaleDateString('en-US', {
weekday: 'long',
year: 'numeric',
month: 'long',
day: 'numeric',
}),
})),
];
return groups.map(({ sessions, title }) => {
if (sessions.length === 0) return null;
const currentIndex = groupIndex++;
return (
<div
key={title}
className="animate-in slide-in-from-left-2 fade-in duration-300"
style={{
animationDelay: `${currentIndex * 100}ms`,
animationFillMode: 'both',
}}
>
{renderSessionGroup(sessions, title, currentIndex)}
</div>
);
});
})()}
</div>
</SidebarGroupContent>
</CollapsibleContent>
</SidebarGroup>
</Collapsible>
);
};
@@ -3,7 +3,6 @@ import { Button } from '../ui/button';
import { ScrollArea } from '../ui/scroll-area';
import BackButton from '../ui/BackButton';
import { Card } from '../ui/card';
import { fetchSessionDetails, SessionDetails } from '../../sessions';
import {
getScheduleSessions,
runScheduleNow,
@@ -21,6 +20,7 @@ import { toastError, toastSuccess } from '../../toasts';
import { Loader2, Pause, Play, Edit, Square, Eye } from 'lucide-react';
import cronstrue from 'cronstrue';
import { formatToLocalDateWithTimezone } from '../../utils/date';
import { getSession, Session } from '../../api';
interface ScheduleSessionMeta {
id: string;
@@ -146,7 +146,7 @@ const ScheduleDetailView: React.FC<ScheduleDetailViewProps> = ({ scheduleId, onN
// Track if we explicitly killed a job to distinguish from natural completion
const [jobWasKilled, setJobWasKilled] = useState(false);
const [selectedSessionDetails, setSelectedSessionDetails] = useState<SessionDetails | null>(null);
const [selectedSessionDetails, setSelectedSessionDetails] = useState<Session | null>(null);
const [isLoadingSessionDetails, setIsLoadingSessionDetails] = useState(false);
const [sessionDetailsError, setSessionDetailsError] = useState<string | null>(null);
const [isEditModalOpen, setIsEditModalOpen] = useState(false);
@@ -430,8 +430,11 @@ const ScheduleDetailView: React.FC<ScheduleDetailViewProps> = ({ scheduleId, onN
setSessionDetailsError(null);
setSelectedSessionDetails(null);
try {
const details = await fetchSessionDetails(sessionId);
setSelectedSessionDetails(details);
const response = await getSession<true>({
path: { session_id: sessionId },
throwOnError: true,
});
setSelectedSessionDetails(response.data);
} catch (err) {
console.error(`Failed to load session details for ${sessionId}:`, err);
const errorMsg = err instanceof Error ? err.message : 'Failed to load session details.';
@@ -459,7 +462,7 @@ const ScheduleDetailView: React.FC<ScheduleDetailViewProps> = ({ scheduleId, onN
setSelectedSessionDetails(null);
setSessionDetailsError(null);
}}
onRetry={() => loadAndShowSessionDetails(selectedSessionDetails?.sessionId)}
onRetry={() => loadAndShowSessionDetails(selectedSessionDetails?.id)}
showActionButtons={true}
/>
);
@@ -11,7 +11,7 @@ import {
LoaderCircle,
AlertCircle,
} from 'lucide-react';
import { resumeSession, type SessionDetails } from '../../sessions';
import { resumeSession } from '../../sessions';
import { Button } from '../ui/button';
import { toast } from 'react-toastify';
import { MainPanelLayout } from '../Layout/MainPanelLayout';
@@ -32,6 +32,8 @@ import { ContextManagerProvider } from '../context_management/ContextManager';
import { Message } from '../../types/message';
import BackButton from '../ui/BackButton';
import { Tooltip, TooltipContent, TooltipTrigger } from '../ui/Tooltip';
import { Session } from '../../api';
import { convertApiMessageToFrontendMessage } from '../context_management';
// Helper function to determine if a message is a user message (same as useChatEngine)
const isUserMessage = (message: Message): boolean => {
@@ -46,7 +48,7 @@ const filterMessagesForDisplay = (messages: Message[]): Message[] => {
};
interface SessionHistoryViewProps {
session: SessionDetails;
session: Session;
isLoading: boolean;
error: string | null;
onBack: () => void;
@@ -73,14 +75,12 @@ const SessionHeader: React.FC<{
);
};
// Session messages component that uses the same rendering as BaseChat
const SessionMessages: React.FC<{
messages: Message[];
isLoading: boolean;
error: string | null;
onRetry: () => void;
}> = ({ messages, isLoading, error, onRetry }) => {
// Filter messages for display (same as BaseChat)
const filteredMessages = filterMessagesForDisplay(messages);
return (
@@ -153,6 +153,8 @@ const SessionHistoryView: React.FC<SessionHistoryViewProps> = ({
const [isCopied, setIsCopied] = useState(false);
const [canShare, setCanShare] = useState(false);
const messages = (session.conversation || []).map(convertApiMessageToFrontendMessage);
useEffect(() => {
const savedSessionConfig = localStorage.getItem('session_sharing_config');
if (savedSessionConfig) {
@@ -183,10 +185,10 @@ const SessionHistoryView: React.FC<SessionHistoryViewProps> = ({
const shareToken = await createSharedSession(
config.baseUrl,
session.metadata.working_dir,
session.messages,
session.metadata.description || 'Shared Session',
session.metadata.total_tokens || 0
session.working_dir,
messages,
session.description || 'Shared Session',
session.total_tokens || 0
);
const shareableLink = `goose://sessions/${shareToken}`;
@@ -270,32 +272,32 @@ const SessionHistoryView: React.FC<SessionHistoryViewProps> = ({
<div className="flex-1 flex flex-col min-h-0 px-8">
<SessionHeader
onBack={onBack}
title={session.metadata.description || 'Session Details'}
title={session.description || 'Session Details'}
actionButtons={!isLoading ? actionButtons : null}
>
<div className="flex flex-col">
{!isLoading && session.messages.length > 0 ? (
{!isLoading ? (
<>
<div className="flex items-center text-text-muted text-sm space-x-5 font-mono">
<span className="flex items-center">
<Calendar className="w-4 h-4 mr-1" />
{formatMessageTimestamp(session.messages[0]?.created)}
{formatMessageTimestamp(messages[0]?.created)}
</span>
<span className="flex items-center">
<MessageSquareText className="w-4 h-4 mr-1" />
{session.metadata.message_count}
{session.message_count}
</span>
{session.metadata.total_tokens !== null && (
{session.total_tokens !== null && (
<span className="flex items-center">
<Target className="w-4 h-4 mr-1" />
{(session.metadata.total_tokens || 0).toLocaleString()}
{(session.total_tokens || 0).toLocaleString()}
</span>
)}
</div>
<div className="flex items-center text-text-muted text-sm mt-1 font-mono">
<span className="flex items-center">
<Folder className="w-4 h-4 mr-1" />
{session.metadata.working_dir}
{session.working_dir}
</span>
</div>
</>
@@ -309,7 +311,7 @@ const SessionHistoryView: React.FC<SessionHistoryViewProps> = ({
</SessionHeader>
<SessionMessages
messages={session.messages}
messages={messages}
isLoading={isLoading}
error={error}
onRetry={onRetry}
@@ -1,7 +1,7 @@
import React from 'react';
import { Session } from '../../sessions';
import { Card } from '../ui/card';
import { formatDate } from '../../utils/date';
import { Session } from '../../api';
interface SessionItemProps {
session: Session;
@@ -12,11 +12,11 @@ const SessionItem: React.FC<SessionItemProps> = ({ session, extraActions }) => {
return (
<Card className="p-4 mb-2 hover:bg-accent/50 cursor-pointer flex justify-between items-center">
<div>
<div className="font-medium">{session.metadata.description || `Session ${session.id}`}</div>
<div className="font-medium">{session.description || `Session ${session.id}`}</div>
<div className="text-sm text-muted-foreground">
{formatDate(session.modified)} {session.metadata.message_count} messages
{formatDate(session.updated_at)} {session.message_count} messages
</div>
<div className="text-sm text-muted-foreground">{session.metadata.working_dir}</div>
<div className="text-sm text-muted-foreground">{session.working_dir}</div>
</div>
{extraActions && <div>{extraActions}</div>}
</Card>
@@ -8,7 +8,6 @@ import {
Edit2,
Trash2,
} from 'lucide-react';
import { fetchSessions, updateSessionMetadata, deleteSession, type Session } from '../../sessions';
import { Card } from '../ui/card';
import { Button } from '../ui/button';
import { ScrollArea } from '../ui/scroll-area';
@@ -21,6 +20,7 @@ import { groupSessionsByDate, type DateGroup } from '../../utils/dateUtils';
import { Skeleton } from '../ui/skeleton';
import { toast } from 'react-toastify';
import { ConfirmationModal } from '../ui/ConfirmationModal';
import { deleteSession, listSessions, Session, updateSessionDescription } from '../../api';
interface EditSessionModalProps {
session: Session | null;
@@ -37,7 +37,7 @@ const EditSessionModal = React.memo<EditSessionModalProps>(
useEffect(() => {
if (session && isOpen) {
setDescription(session.metadata.description || session.id);
setDescription(session.description || session.id);
} else if (!isOpen) {
// Reset state when modal closes
setDescription('');
@@ -49,14 +49,18 @@ const EditSessionModal = React.memo<EditSessionModalProps>(
if (!session || disabled) return;
const trimmedDescription = description.trim();
if (trimmedDescription === session.metadata.description) {
if (trimmedDescription === session.description) {
onClose();
return;
}
setIsUpdating(true);
try {
await updateSessionMetadata(session.id, trimmedDescription);
await updateSessionDescription({
path: { session_id: session.id },
body: { description: trimmedDescription },
throwOnError: true,
});
await onSave(session.id, trimmedDescription);
// Close modal, then show success toast on a timeout to let the UI update complete.
@@ -68,8 +72,7 @@ const EditSessionModal = React.memo<EditSessionModalProps>(
const errorMessage = error instanceof Error ? error.message : 'Unknown error occurred';
console.error('Failed to update session description:', errorMessage);
toast.error(`Failed to update session description: ${errorMessage}`);
// Reset to original description on error
setDescription(session.metadata.description || session.id);
setDescription(session.description || session.id);
} finally {
setIsUpdating(false);
}
@@ -213,7 +216,8 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
setShowContent(false);
setError(null);
try {
const sessions = await fetchSessions();
const resp = await listSessions<true>({ throwOnError: true });
const sessions = resp.data.sessions;
// Use startTransition to make state updates non-blocking
startTransition(() => {
setSessions(sessions);
@@ -291,20 +295,20 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
startTransition(() => {
const searchTerm = caseSensitive ? debouncedSearchTerm : debouncedSearchTerm.toLowerCase();
const filtered = sessions.filter((session) => {
const description = session.metadata.description || session.id;
const path = session.path;
const workingDir = session.metadata.working_dir;
const description = session.description || session.id;
const workingDir = session.working_dir;
const sessionId = session.id;
if (caseSensitive) {
return (
description.includes(searchTerm) ||
path.includes(searchTerm) ||
sessionId.includes(searchTerm) ||
workingDir.includes(searchTerm)
);
} else {
return (
description.toLowerCase().includes(searchTerm) ||
path.toLowerCase().includes(searchTerm) ||
sessionId.toLowerCase().includes(searchTerm) ||
workingDir.toLowerCase().includes(searchTerm)
);
}
@@ -355,11 +359,7 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
const handleModalSave = useCallback(async (sessionId: string, newDescription: string) => {
// Update state immediately for optimistic UI
setSessions((prevSessions) =>
prevSessions.map((s) =>
s.id === sessionId
? { ...s, metadata: { ...s.metadata, description: newDescription } }
: s
)
prevSessions.map((s) => (s.id === sessionId ? { ...s, description: newDescription } : s))
);
}, []);
@@ -378,18 +378,21 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
setShowDeleteConfirmation(false);
const sessionToDeleteId = sessionToDelete.id;
const sessionName = sessionToDelete.metadata.description || sessionToDelete.id;
const sessionName = sessionToDelete.description || sessionToDelete.id;
setSessionToDelete(null);
try {
await deleteSession(sessionToDeleteId);
await deleteSession({
path: { session_id: sessionToDeleteId },
throwOnError: true,
});
toast.success('Session deleted successfully');
loadSessions();
} catch (error) {
console.error('Error deleting session:', error);
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
toast.error(`Failed to delete session "${sessionName}": ${errorMessage}`);
}
await loadSessions();
}, [sessionToDelete, loadSessions]);
const handleCancelDelete = useCallback(() => {
@@ -451,16 +454,16 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
<div className="flex-1">
<h3 className="text-base mb-1 pr-16 break-words">
{session.metadata.description || session.id}
{session.description || session.id}
</h3>
<div className="flex items-center text-text-muted text-xs mb-1">
<Calendar className="w-3 h-3 mr-1 flex-shrink-0" />
<span>{formatMessageTimestamp(Date.parse(session.modified) / 1000)}</span>
<span>{formatMessageTimestamp(Date.parse(session.updated_at) / 1000)}</span>
</div>
<div className="flex items-center text-text-muted text-xs mb-1">
<Folder className="w-3 h-3 mr-1 flex-shrink-0" />
<span className="truncate">{session.metadata.working_dir}</span>
<span className="truncate">{session.working_dir}</span>
</div>
</div>
@@ -468,14 +471,12 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
<div className="flex items-center space-x-3 text-xs text-text-muted">
<div className="flex items-center">
<MessageSquareText className="w-3 h-3 mr-1" />
<span className="font-mono">{session.metadata.message_count}</span>
<span className="font-mono">{session.message_count}</span>
</div>
{session.metadata.total_tokens !== null && (
{session.total_tokens !== null && (
<div className="flex items-center">
<Target className="w-3 h-3 mr-1" />
<span className="font-mono">
{(session.metadata.total_tokens || 0).toLocaleString()}
</span>
<span className="font-mono">{(session.total_tokens || 0).toLocaleString()}</span>
</div>
)}
</div>
@@ -675,7 +676,7 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
<ConfirmationModal
isOpen={showDeleteConfirmation}
title="Delete Session"
message={`Are you sure you want to delete the session "${sessionToDelete?.metadata.description || sessionToDelete?.id}"? This action cannot be undone.`}
message={`Are you sure you want to delete the session "${sessionToDelete?.description || sessionToDelete?.id}"? This action cannot be undone.`}
confirmLabel="Delete Session"
cancelLabel="Cancel"
confirmVariant="destructive"
@@ -72,7 +72,7 @@ export const SessionHeaderCard: React.FC<SessionHeaderCardProps> = ({ onBack, ch
/**
* Props for the SessionMessages component
*/
export interface SessionMessagesProps {
interface SessionMessagesProps {
messages: Message[];
isLoading: boolean;
error: string | null;
@@ -1,23 +1,21 @@
import { useEffect, useState } from 'react';
import { Card, CardContent, CardDescription } from '../ui/card';
import { getApiUrl } from '../../config';
import { Greeting } from '../common/Greeting';
import { fetchSessions, type Session, resumeSession } from '../../sessions';
import { useNavigate } from 'react-router-dom';
import { Button } from '../ui/button';
import { ChatSmart } from '../icons/';
import { Goose } from '../icons/Goose';
import { Skeleton } from '../ui/skeleton';
interface SessionInsightsType {
totalSessions: number;
mostActiveDirs: [string, number][];
avgSessionDuration: number;
totalTokens: number;
}
import {
getSessionInsights,
listSessions,
Session,
SessionInsights as ApiSessionInsights,
} from '../../api';
import { resumeSession } from '../../sessions';
export function SessionInsights() {
const [insights, setInsights] = useState<SessionInsightsType | null>(null);
const [insights, setInsights] = useState<ApiSessionInsights | null>(null);
const [error, setError] = useState<string | null>(null);
const [recentSessions, setRecentSessions] = useState<Session[]>([]);
const [isLoading, setIsLoading] = useState(true);
@@ -29,31 +27,14 @@ export function SessionInsights() {
const loadInsights = async () => {
try {
const response = await fetch(getApiUrl('/sessions/insights'), {
headers: {
Accept: 'application/json',
'Content-Type': 'application/json',
'X-Secret-Key': await window.electron.getSecretKey(),
},
});
if (!response.ok) {
const errorText = await response.text();
throw new Error(`Failed to fetch insights: ${response.status} ${errorText}`);
}
const data = await response.json();
setInsights(data);
// Clear any previous error when insights load successfully
const response = await getSessionInsights({ throwOnError: true });
setInsights(response.data);
setError(null);
} catch (error) {
console.error('Failed to load insights:', error);
setError(error instanceof Error ? error.message : 'Failed to load insights');
// Set fallback insights data so the UI can still render
setInsights({
totalSessions: 0,
mostActiveDirs: [],
avgSessionDuration: 0,
totalTokens: 0,
});
} finally {
@@ -63,10 +44,8 @@ export function SessionInsights() {
const loadRecentSessions = async () => {
try {
const sessions = await fetchSessions();
setRecentSessions(sessions.slice(0, 3));
} catch (error) {
console.error('Failed to load recent sessions:', error);
const response = await listSessions<true>({ throwOnError: true });
setRecentSessions(response.data.sessions.slice(0, 3));
} finally {
setIsLoadingSessions(false);
}
@@ -85,6 +64,7 @@ export function SessionInsights() {
mostActiveDirs: [],
avgSessionDuration: 0,
totalTokens: 0,
recentActivity: [],
};
}
// If we already have insights, just make sure loading is false
@@ -155,16 +135,6 @@ export function SessionInsights() {
</CardContent>
</Card>
{/* Average Duration Card Skeleton */}
{/*<Card className="w-full py-6 px-6 border-none rounded-2xl bg-background-default">*/}
{/* <CardContent className="flex flex-col justify-end h-full p-0">*/}
{/* <div className="flex flex-col justify-end">*/}
{/* <Skeleton className="h-10 w-20 mb-1" />*/}
{/* <span className="text-xs text-text-muted">Avg. chat length</span>*/}
{/* </div>*/}
{/* </CardContent>*/}
{/*</Card>*/}
{/* Total Tokens Card Skeleton */}
<Card className="w-full py-6 px-6 border-none rounded-2xl bg-background-default">
<CardContent className="flex flex-col justify-end h-full p-0">
@@ -363,11 +333,11 @@ export function SessionInsights() {
<div className="flex items-center space-x-2">
<ChatSmart className="h-4 w-4 text-text-muted" />
<span className="truncate max-w-[300px]">
{session.metadata.description || session.id}
{session.description || session.id}
</span>
</div>
<span className="text-text-muted font-mono font-light">
{formatDateOnly(session.modified)}
{formatDateOnly(session.updated_at)}
</span>
</div>
))
@@ -1,17 +1,16 @@
import React, { useState, useEffect, useCallback } from 'react';
import { View, ViewOptions } from '../../utils/navigationUtils';
import { fetchSessionDetails, type SessionDetails } from '../../sessions';
import SessionListView from './SessionListView';
import SessionHistoryView from './SessionHistoryView';
import { toastError } from '../../toasts';
import { useLocation } from 'react-router-dom';
import { getSession, Session } from '../../api';
interface SessionsViewProps {
setView: (view: View, viewOptions?: ViewOptions) => void;
}
const SessionsView: React.FC<SessionsViewProps> = ({ setView }) => {
const [selectedSession, setSelectedSession] = useState<SessionDetails | null>(null);
const [selectedSession, setSelectedSession] = useState<Session | null>(null);
const [showSessionHistory, setShowSessionHistory] = useState(false);
const [isLoadingSession, setIsLoadingSession] = useState(false);
const [error, setError] = useState<string | null>(null);
@@ -23,21 +22,17 @@ const SessionsView: React.FC<SessionsViewProps> = ({ setView }) => {
setError(null);
setShowSessionHistory(true);
try {
const sessionDetails = await fetchSessionDetails(sessionId);
setSelectedSession(sessionDetails);
const response = await getSession<true>({
path: { session_id: sessionId },
throwOnError: true,
});
setSelectedSession(response.data);
} catch (err) {
console.error(`Failed to load session details for ${sessionId}:`, err);
setError('Failed to load session details. Please try again later.');
// Keep the selected session null if there's an error
setSelectedSession(null);
setShowSessionHistory(false);
const errorMessage = err instanceof Error ? err.message : String(err);
toastError({
title: 'Failed to load session. The file may be corrupted.',
msg: 'Please try again later.',
traceback: errorMessage,
});
} finally {
setIsLoadingSession(false);
setInitialSessionId(null);
@@ -68,7 +63,7 @@ const SessionsView: React.FC<SessionsViewProps> = ({ setView }) => {
const handleRetryLoadSession = () => {
if (selectedSession) {
loadSessionDetails(selectedSession.sessionId);
loadSessionDetails(selectedSession.id);
}
};
@@ -78,14 +73,15 @@ const SessionsView: React.FC<SessionsViewProps> = ({ setView }) => {
<SessionHistoryView
session={
selectedSession || {
sessionId: initialSessionId || '',
messages: [],
metadata: {
description: 'Loading...',
working_dir: '',
message_count: 0,
total_tokens: 0,
},
id: initialSessionId || '',
conversation: [],
description: 'Loading...',
working_dir: '',
message_count: 0,
total_tokens: 0,
created_at: '',
updated_at: '',
extension_data: {},
}
}
isLoading={isLoadingSession}
@@ -97,7 +93,7 @@ const SessionsView: React.FC<SessionsViewProps> = ({ setView }) => {
<SessionListView
setView={setView}
onSelectSession={handleSelectSession}
selectedSessionId={selectedSession?.sessionId ?? null}
selectedSessionId={selectedSession?.id ?? null}
/>
);
};
+3 -6
View File
@@ -1,6 +1,5 @@
import React, { createContext, useContext, ReactNode } from 'react';
import { ChatType } from '../types/chat';
import { generateSessionId } from '../sessions';
import { Recipe } from '../recipe';
import { useDraftContext } from './DraftContext';
@@ -54,16 +53,14 @@ export const ChatProvider: React.FC<ChatProviderProps> = ({
};
const resetChat = () => {
const newSessionId = generateSessionId();
setChat({
sessionId: newSessionId,
sessionId: '',
title: DEFAULT_CHAT_TITLE,
messages: [],
messageHistoryIndex: 0,
recipeConfig: null, // Clear recipe when resetting chat
recipeParameters: null, // Clear when resetting chat
recipeConfig: null,
recipeParameters: null,
});
// Clear draft when resetting chat
clearDraft();
};
+2 -1
View File
@@ -34,7 +34,8 @@ export const checkServerStatus = async (client: Client): Promise<boolean> => {
try {
await status({ client, throwOnError: true });
return true;
} catch {
} catch (error) {
log.error('failure to connect, will retry', error);
if (attempt === maxAttempts) {
log.error(`Server failed to respond after ${(interval * maxAttempts) / 1000} seconds`);
}
+18 -20
View File
@@ -1,4 +1,4 @@
import { useState, useCallback, useRef } from 'react';
import { useCallback, useRef, useState } from 'react';
import { useConfig } from '../components/ConfigContext';
import { ChatType } from '../types/chat';
import { initializeSystem } from '../utils/providerUtils';
@@ -8,11 +8,11 @@ import {
initConfig,
Message as ApiMessage,
readAllConfig,
Recipe,
recoverConfig,
resumeAgent,
startAgent,
validateConfig,
Recipe,
} from '../api';
import { COST_TRACKING_ENABLED } from '../updates';
import { convertApiMessageToFrontendMessage } from '../components/context_management';
@@ -72,19 +72,17 @@ export function useAgent(): UseAgentReturn {
throwOnError: true,
});
const agentSessionInfo = agentResponse.data;
const sessionMetadata = agentSessionInfo.metadata;
let chat: ChatType = {
sessionId: agentSessionInfo.session_id,
title: sessionMetadata.recipe?.title || sessionMetadata.description,
const agentSession = agentResponse.data;
const messages = agentSession.conversation || [];
return {
sessionId: agentSession.id,
title: agentSession.recipe?.title || agentSession.description,
messageHistoryIndex: 0,
messages: agentSessionInfo.messages.map((message: ApiMessage) =>
messages: messages?.map((message: ApiMessage) =>
convertApiMessageToFrontendMessage(message)
),
recipeConfig: sessionMetadata.recipe,
recipeConfig: agentSession.recipe,
};
return chat;
}
if (initPromiseRef.current) {
@@ -121,11 +119,11 @@ export function useAgent(): UseAgentReturn {
throwOnError: true,
});
const agentSessionInfo = agentResponse.data;
if (!agentSessionInfo) {
const agentSession = agentResponse.data;
if (!agentSession) {
throw Error('Failed to get session info');
}
setSessionId(agentSessionInfo.session_id);
setSessionId(agentSession.id);
agentWaitingMessage('Agent is loading config');
@@ -139,7 +137,7 @@ export function useAgent(): UseAgentReturn {
}
agentWaitingMessage('Extensions are loading');
await initializeSystem(agentSessionInfo.session_id, provider as string, model as string, {
await initializeSystem(agentSession.id, provider as string, model as string, {
getExtensions,
addExtension,
setIsExtensionsLoading: initContext.setIsExtensionsLoading,
@@ -153,15 +151,15 @@ export function useAgent(): UseAgentReturn {
}
}
const sessionMetadata = agentSessionInfo.metadata;
const messages = agentSession.conversation || [];
let initChat: ChatType = {
sessionId: agentSessionInfo.session_id,
title: sessionMetadata.recipe?.title || sessionMetadata.description,
sessionId: agentSession.id,
title: agentSession.recipe?.title || agentSession.description,
messageHistoryIndex: 0,
messages: agentSessionInfo.messages.map((message: ApiMessage) =>
messages: messages.map((message: ApiMessage) =>
convertApiMessageToFrontendMessage(message)
),
recipeConfig: sessionMetadata.recipe,
recipeConfig: agentSession.recipe,
};
setAgentState(AgentState.INITIALIZED);
+17 -14
View File
@@ -1,7 +1,6 @@
import { useCallback, useEffect, useMemo, useState } from 'react';
import { getApiUrl } from '../config';
import { useMessageStream } from './useMessageStream';
import { fetchSessionDetails } from '../sessions';
import { LocalMessageStorage } from '../utils/localMessageStorage';
import {
Message,
@@ -16,6 +15,7 @@ import {
} from '../types/message';
import { ChatType } from '../types/chat';
import { ChatState } from '../types/chatState';
import { getSession } from '../api';
// Helper function to determine if a message is a user message
const isUserMessage = (message: Message): boolean => {
@@ -85,7 +85,7 @@ export const useChatEngine = ({
handleInputChange: _handleInputChange,
updateMessageStreamBody,
notifications,
sessionMetadata,
session,
setError,
} = useMessageStream({
api: getApiUrl('/reply'),
@@ -199,14 +199,17 @@ export const useChatEngine = ({
setChat((prevChat: ChatType) => ({ ...prevChat, messages }));
}, [messages, setChat]);
// Fetch session metadata to get token count
useEffect(() => {
const fetchSessionTokens = async () => {
try {
const sessionDetails = await fetchSessionDetails(chat.sessionId);
setSessionTokenCount(sessionDetails.metadata.total_tokens || 0);
setSessionInputTokens(sessionDetails.metadata.accumulated_input_tokens || 0);
setSessionOutputTokens(sessionDetails.metadata.accumulated_output_tokens || 0);
const response = await getSession<true>({
path: { session_id: chat.sessionId },
throwOnError: true,
});
const sessionDetails = response.data;
setSessionTokenCount(sessionDetails.total_tokens || 0);
setSessionInputTokens(sessionDetails.accumulated_input_tokens || 0);
setSessionOutputTokens(sessionDetails.accumulated_output_tokens || 0);
} catch (err) {
console.error('Error fetching session token count:', err);
}
@@ -219,13 +222,13 @@ export const useChatEngine = ({
// Update token counts when sessionMetadata changes from the message stream
useEffect(() => {
console.log('Session metadata received:', sessionMetadata);
if (sessionMetadata) {
setSessionTokenCount(sessionMetadata.total_tokens || 0);
setSessionInputTokens(sessionMetadata.accumulated_input_tokens || 0);
setSessionOutputTokens(sessionMetadata.accumulated_output_tokens || 0);
console.log('Session metadata received:', session);
if (session) {
setSessionTokenCount(session.total_tokens || 0);
setSessionInputTokens(session.accumulated_input_tokens || 0);
setSessionOutputTokens(session.accumulated_output_tokens || 0);
}
}, [sessionMetadata]);
}, [session]);
useEffect(() => {
return () => {
@@ -473,7 +476,7 @@ export const useChatEngine = ({
// Stream utilities
updateMessageStreamBody,
sessionMetadata,
sessionMetadata: session,
// Utilities
isUserMessage,
+4 -4
View File
@@ -1,14 +1,14 @@
import { useEffect, useRef, useState } from 'react';
import { useModelAndProvider } from '../components/ModelAndProviderContext';
import { getCostForModel } from '../utils/costDatabase';
import { SessionMetadata } from '../api';
import { Session } from '../api';
interface UseCostTrackingProps {
sessionInputTokens: number;
sessionOutputTokens: number;
localInputTokens: number;
localOutputTokens: number;
sessionMetadata?: SessionMetadata | null;
session?: Session | null;
}
export const useCostTracking = ({
@@ -16,7 +16,7 @@ export const useCostTracking = ({
sessionOutputTokens,
localInputTokens,
localOutputTokens,
sessionMetadata,
session,
}: UseCostTrackingProps) => {
const [sessionCosts, setSessionCosts] = useState<{
[key: string]: {
@@ -79,7 +79,7 @@ export const useCostTracking = ({
sessionOutputTokens,
localInputTokens,
localOutputTokens,
sessionMetadata,
session,
]);
return {
+11 -16
View File
@@ -1,7 +1,7 @@
import { useCallback, useEffect, useId, useReducer, useRef, useState } from 'react';
import useSWR from 'swr';
import { createUserMessage, hasCompletedToolCalls, Message, Role } from '../types/message';
import { getSessionHistory, SessionMetadata } from '../api';
import { getSession, Session } from '../api';
import { ChatState } from '../types/chatState';
let messageIdCounter = 0;
@@ -152,8 +152,8 @@ export interface UseMessageStreamHelpers {
/** Current model info from the backend */
currentModelInfo: { model: string; mode: string } | null;
/** Session metadata including token counts */
sessionMetadata: SessionMetadata | null;
/** Session including token counts */
session: Session | null;
/** Clear error state */
setError: (error: Error | undefined) => void;
@@ -188,7 +188,7 @@ export function useMessageStream({
const [currentModelInfo, setCurrentModelInfo] = useState<{ model: string; mode: string } | null>(
null
);
const [sessionMetadata, setSessionMetadata] = useState<SessionMetadata | null>(null);
const [session, setSession] = useState<Session | null>(null);
// expose a way to update the body so we can update the session id when CLE occurs
const updateMessageStreamBody = useCallback((newBody: object) => {
@@ -333,26 +333,21 @@ export function useMessageStream({
}
case 'Finish': {
// Call onFinish with the last message if available
if (onFinish && currentMessages.length > 0) {
const lastMessage = currentMessages[currentMessages.length - 1];
onFinish(lastMessage, parsedEvent.reason);
}
// Fetch updated session metadata with token counts
const sessionId = (extraMetadataRef.current.body as Record<string, unknown>)
?.session_id as string;
if (sessionId) {
try {
const sessionResponse = await getSessionHistory({
path: { session_id: sessionId },
});
const sessionResponse = await getSession({
path: { session_id: sessionId },
throwOnError: true,
});
if (sessionResponse.data?.metadata) {
setSessionMetadata(sessionResponse.data?.metadata);
}
} catch (error) {
console.error('Failed to fetch session metadata:', error);
if (sessionResponse.data) {
setSession(sessionResponse.data);
}
}
break;
@@ -631,7 +626,7 @@ export function useMessageStream({
updateMessageStreamBody,
notifications,
currentModelInfo,
sessionMetadata,
session: session,
setError,
};
}
+6 -175
View File
@@ -1,185 +1,16 @@
import { Message } from './types/message';
import {
getSessionHistory,
listSessions,
SessionInfo,
Message as ApiMessage,
SessionMetadata,
} from './api';
import { convertApiMessageToFrontendMessage } from './components/context_management';
import { getApiUrl } from './config';
import { Session } from './api';
// Helper function to ensure working directory is set
export function ensureWorkingDir(metadata: Partial<SessionMetadata>): SessionMetadata {
return {
description: metadata.description || '',
message_count: metadata.message_count || 0,
total_tokens: metadata.total_tokens || null,
working_dir: metadata.working_dir || process.env.HOME || '',
accumulated_input_tokens: metadata.accumulated_input_tokens || null,
accumulated_output_tokens: metadata.accumulated_output_tokens || null,
accumulated_total_tokens: metadata.accumulated_total_tokens || null,
};
}
export interface Session {
id: string;
path: string;
modified: string;
metadata: SessionMetadata;
}
export interface SessionDetails {
sessionId: string;
metadata: SessionMetadata;
messages: Message[];
}
/**
* Generate a session ID in the format yyyymmdd_hhmmss
*/
export function generateSessionId(): string {
const now = new Date();
const year = now.getFullYear();
const month = String(now.getMonth() + 1).padStart(2, '0');
const day = String(now.getDate()).padStart(2, '0');
const hours = String(now.getHours()).padStart(2, '0');
const minutes = String(now.getMinutes()).padStart(2, '0');
const seconds = String(now.getSeconds()).padStart(2, '0');
return `${year}${month}${day}_${hours}${minutes}${seconds}`;
}
/**
* Fetches all available sessions from the API
* @returns Promise with sessions data
*/
/**
* Fetches all available sessions from the API
* @returns Promise with an array of Session objects
*/
export async function fetchSessions(): Promise<Session[]> {
const response = await listSessions<true>();
// Check if the response has the expected structure
if (response && response.data && response.data.sessions) {
// Since the API returns SessionInfo, we need to convert to Session
const sessions = response.data.sessions
.filter(
(sessionInfo: SessionInfo) => sessionInfo.metadata && sessionInfo.metadata.message_count > 0
)
.map(
(sessionInfo: SessionInfo): Session => ({
id: sessionInfo.id,
path: sessionInfo.path,
modified: sessionInfo.modified,
metadata: ensureWorkingDir(sessionInfo.metadata),
})
);
// order sessions by 'modified' date descending
sessions.sort(
(a: Session, b: Session) => new Date(b.modified).getTime() - new Date(a.modified).getTime()
);
return sessions;
} else {
throw new Error('Unexpected response format from listSessions');
}
}
/**
* Fetches details for a specific session
* @param sessionId The ID of the session to fetch
* @returns Promise with session details
*/
export async function fetchSessionDetails(sessionId: string): Promise<SessionDetails> {
const response = await getSessionHistory<true>({
path: { session_id: sessionId },
});
try {
// Convert the SessionHistoryResponse to a SessionDetails object
return {
sessionId: response.data.sessionId,
metadata: ensureWorkingDir(response.data.metadata),
messages: response.data.messages.map((message: ApiMessage) =>
convertApiMessageToFrontendMessage(message)
), // slight diffs between backend and frontend Message obj
};
} catch (error) {
console.error(`Error fetching session details for ${sessionId}:`, error);
throw error;
}
}
/**
* Updates the metadata for a specific session
* @param sessionId The ID of the session to update
* @param description The new description (name) for the session
* @returns Promise that resolves when the update is complete
*/
export async function updateSessionMetadata(sessionId: string, description: string): Promise<void> {
const url = getApiUrl(`/sessions/${sessionId}/metadata`);
const secretKey = await window.electron.getSecretKey();
const response = await fetch(url, {
method: 'PUT',
headers: {
'Content-Type': 'application/json',
'X-Secret-Key': secretKey,
},
body: JSON.stringify({ description }),
});
if (!response.ok) {
const errorText = await response.text();
throw new Error(`Failed to update session metadata: ${response.statusText} - ${errorText}`);
}
}
/**
* Resumes a session. Currently, this opens a new window with the session loaded.
*/
export function resumeSession(session: SessionDetails | Session) {
const resumedSessionId = 'sessionId' in session ? session.sessionId : session.id;
console.log('Launching session in new window:', resumedSessionId);
const workingDir = session.metadata?.working_dir;
export function resumeSession(session: Session) {
console.log('Launching session in new window:', session.description || session.id);
const workingDir = session.working_dir;
if (!workingDir) {
throw new Error('Cannot resume session: working directory is missing in session metadata');
throw new Error('Cannot resume session: working directory is missing in session');
}
window.electron.createChatWindow(
undefined, // query
workingDir,
undefined, // version
resumedSessionId
session.id
);
}
/**
* Deletes a specific session
* @param sessionId The ID of the session to delete
* @returns Promise that resolves when the deletion is complete
*/
export async function deleteSession(sessionId: string): Promise<void> {
try {
const url = getApiUrl(`/sessions/${sessionId}/delete`);
const secretKey = await window.electron.getSecretKey();
const response = await fetch(url, {
method: 'DELETE',
headers: {
'X-Secret-Key': secretKey,
},
});
if (!response.ok) {
const errorText = await response.text();
throw new Error(`Failed to delete session: ${response.statusText} - ${errorText}`);
}
} catch (error) {
console.error(`Error deleting session ${sessionId}:`, error);
throw error;
}
}
+2 -2
View File
@@ -1,4 +1,4 @@
import { Session } from '../sessions';
import { Session } from '../api';
export interface DateGroup {
label: string;
@@ -16,7 +16,7 @@ export function groupSessionsByDate(sessions: Session[]): DateGroup[] {
const groups: { [key: string]: DateGroup } = {};
sessions.forEach((session) => {
const sessionDate = new Date(session.modified);
const sessionDate = new Date(session.updated_at);
const sessionDateStart = new Date(sessionDate);
sessionDateStart.setHours(0, 0, 0, 0);