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

Session manager (#4648)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2025-09-26 14:56:45 -04:00
committed by GitHub
parent d0aa14a1ea
commit dc292883b7
61 changed files with 2895 additions and 5572 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/cli.rs
+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,
crates/goose-cli/src/commands/bench.rs
+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");
crates/goose-cli/src/commands/schedule.rs
+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")
);
}
crates/goose-cli/src/commands/session.rs
+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"))
}
crates/goose-cli/src/commands/web.rs
+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;
crates/goose-cli/src/lib.rs
+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(),
crates/goose-cli/src/scenario_tests/scenario_runner.rs
+2 -2
View File
@@ -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 {
crates/goose-cli/src/session/builder.rs
+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());
crates/goose-cli/src/session/mod.rs
+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;
crates/goose-cli/src/session/output.rs
+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(),