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

feat: sessions api, view & resume prev sessions (#1453)

Browse Source 
* Centralize session files to goose::session module
* Write session metadata and messages in jsonl
* Refactor CLI build_session to use goose::session functions
* Track session's token usage by adding optional session_id in agent.reply(...)
* NOTE: Only sessions saved through the updates goose::session functions will show up in GUI

Co-authored-by: Bradley Axen <baxen@squareup.com>
This commit is contained in:
Salman Mohammed
2025-03-03 11:49:15 -05:00
committed by GitHub
parent 68b8c5d19d
commit 9ae9045584
25 changed files with 1413 additions and 257 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/builder.rs
+9 -25
View File
@@ -2,16 +2,16 @@ use console::style;
use goose::agents::extension::ExtensionError;
use goose::agents::AgentFactory;
use goose::config::{Config, ExtensionManager};
use goose::session;
use goose::session::Identifier;
use mcp_client::transport::Error as McpClientError;
use std::path::PathBuf;
use std::process;
use super::output;
use super::storage;
use super::Session;
pub async fn build_session(
identifier: Option<storage::Identifier>,
identifier: Option<Identifier>,
resume: bool,
extensions: Vec<String>,
builtins: Vec<String>,
@@ -65,7 +65,7 @@ pub async fn build_session(
// Handle session file resolution and resuming
let session_file = if resume {
if let Some(identifier) = identifier {
let session_file = storage::get_path(identifier);
let session_file = session::get_path(identifier);
if !session_file.exists() {
output::render_error(&format!(
"Cannot resume session {} - no such session exists",
@@ -76,7 +76,7 @@ pub async fn build_session(
session_file
} else {
// Try to resume most recent session
match storage::get_most_recent_session() {
match session::get_most_recent_session() {
Ok(file) => file,
Err(_) => {
output::render_error("Cannot resume - no previous sessions found");
@@ -88,10 +88,11 @@ pub async fn build_session(
// Create new session with provided name/path or generated name
let id = match identifier {
Some(identifier) => identifier,
None => storage::Identifier::Name(generate_session_name()),
None => Identifier::Name(session::generate_session_id()),
};
let session_file = storage::get_path(id);
create_new_session_file(session_file)
// Just get the path - file will be created when needed
session::get_path(id)
};
// Create new session
@@ -130,20 +131,3 @@ pub async fn build_session(
output::display_session_info(resume, &provider_name, &model, &session_file);
session
}
fn generate_session_name() -> String {
use rand::{distributions::Alphanumeric, Rng};
rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(8)
.map(char::from)
.collect()
}
fn create_new_session_file(session_file: PathBuf) -> PathBuf {
if session_file.exists() {
eprintln!("Session '{:?}' already exists", session_file);
process::exit(1);
}
session_file
}
crates/goose-cli/src/session/mod.rs
+43 -10
View File
@@ -3,11 +3,10 @@ mod completion;
mod input;
mod output;
mod prompt;
mod storage;
mod thinking;
pub use builder::build_session;
pub use storage::Identifier;
pub use goose::session::Identifier;
use anyhow::Result;
use completion::GooseCompleter;
@@ -16,6 +15,7 @@ use etcetera::AppStrategy;
use goose::agents::extension::{Envs, ExtensionConfig};
use goose::agents::Agent;
use goose::message::{Message, MessageContent};
use goose::session;
use mcp_core::handler::ToolError;
use mcp_core::prompt::PromptMessage;
@@ -56,7 +56,7 @@ impl CompletionCache {
impl Session {
pub fn new(agent: Box<dyn Agent>, session_file: PathBuf) -> Self {
let messages = match storage::read_messages(&session_file) {
let messages = match session::read_messages(&session_file) {
Ok(msgs) => msgs,
Err(e) => {
eprintln!("Warning: Failed to load message history: {}", e);
@@ -196,7 +196,13 @@ impl Session {
/// Process a single message and get the response
async fn process_message(&mut self, message: String) -> Result<()> {
self.messages.push(Message::user().with_text(&message));
storage::persist_messages(&self.session_file, &self.messages)?;
// Get the provider from the agent for description generation
let provider = self.agent.provider().await;
// Persist messages with provider for automatic description generation
session::persist_messages(&self.session_file, &self.messages, Some(provider)).await?;
self.process_agent_response(false).await?;
Ok(())
}
@@ -260,7 +266,13 @@ impl Session {
save_history(&mut editor);
self.messages.push(Message::user().with_text(&content));
storage::persist_messages(&self.session_file, &self.messages)?;
// Get the provider from the agent for description generation
let provider = self.agent.provider().await;
// Persist messages with provider for automatic description generation
session::persist_messages(&self.session_file, &self.messages, Some(provider))
.await?;
output::show_thinking();
self.process_agent_response(true).await?;
@@ -399,7 +411,8 @@ impl Session {
}
async fn process_agent_response(&mut self, interactive: bool) -> Result<()> {
let mut stream = self.agent.reply(&self.messages).await?;
let session_id = session::Identifier::Path(self.session_file.clone());
let mut stream = self.agent.reply(&self.messages, Some(session_id)).await?;
use futures::StreamExt;
loop {
@@ -421,7 +434,10 @@ impl Session {
// otherwise we have a model/tool to render
else {
self.messages.push(message.clone());
storage::persist_messages(&self.session_file, &self.messages)?;
// No need to update description on assistant messages
session::persist_messages(&self.session_file, &self.messages, None).await?;
if interactive {output::hide_thinking()};
output::render_message(&message);
if interactive {output::show_thinking()};
@@ -430,7 +446,9 @@ impl Session {
Some(Err(e)) => {
eprintln!("Error: {}", e);
drop(stream);
self.handle_interrupted_messages(false);
if let Err(e) = self.handle_interrupted_messages(false).await {
eprintln!("Error handling interruption: {}", e);
}
output::render_error(
"The error above was an exception we were not able to handle.\n\
These errors are often related to connection or authentication\n\
@@ -444,7 +462,9 @@ impl Session {
}
_ = tokio::signal::ctrl_c() => {
drop(stream);
self.handle_interrupted_messages(true);
if let Err(e) = self.handle_interrupted_messages(true).await {
eprintln!("Error handling interruption: {}", e);
}
break;
}
}
@@ -452,7 +472,7 @@ impl Session {
Ok(())
}
fn handle_interrupted_messages(&mut self, interrupt: bool) {
async fn handle_interrupted_messages(&mut self, interrupt: bool) -> Result<()> {
// First, get any tool requests from the last message if it exists
let tool_requests = self
.messages
@@ -493,11 +513,18 @@ impl Session {
}
self.messages.push(response_message);
// No need for description update here
session::persist_messages(&self.session_file, &self.messages, None).await?;
let prompt = format!(
"The existing call to {} was interrupted. How would you like to proceed?",
last_tool_name
);
self.messages.push(Message::assistant().with_text(&prompt));
// No need for description update here
session::persist_messages(&self.session_file, &self.messages, None).await?;
output::render_message(&Message::assistant().with_text(&prompt));
} else {
// An interruption occurred outside of a tool request-response.
@@ -508,6 +535,11 @@ 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.messages.push(Message::assistant().with_text(prompt));
// No need for description update here
session::persist_messages(&self.session_file, &self.messages, None)
.await?;
output::render_message(&Message::assistant().with_text(prompt));
}
Some(_) => {
@@ -521,6 +553,7 @@ impl Session {
}
}
}
Ok(())
}
pub fn session_file(&self) -> PathBuf {
crates/goose-cli/src/session/storage.rs
-180
View File
@@ -1,180 +0,0 @@
use anyhow::Result;
use etcetera::{choose_app_strategy, AppStrategy};
use goose::message::Message;
use std::fs::{self, File};
use std::io::{self, BufRead, Write};
use std::path::{Path, PathBuf};
pub enum Identifier {
Name(String),
Path(PathBuf),
}
pub fn get_path(id: Identifier) -> PathBuf {
match id {
Identifier::Name(name) => {
let session_dir = ensure_session_dir().expect("Failed to create session directory");
session_dir.join(format!("{}.jsonl", name))
}
Identifier::Path(path) => path,
}
}
/// Ensure the session directory exists and return its path
pub fn ensure_session_dir() -> Result<PathBuf> {
let data_dir = choose_app_strategy(crate::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)
}
/// Get the path to the most recently modified session file
pub fn get_most_recent_session() -> Result<PathBuf> {
let session_dir = ensure_session_dir()?;
let mut entries = fs::read_dir(&session_dir)?
.filter_map(|entry| entry.ok())
.filter(|entry| entry.path().extension().is_some_and(|ext| ext == "jsonl"))
.collect::<Vec<_>>();
if entries.is_empty() {
return Err(anyhow::anyhow!("No session files found"));
}
// Sort by modification time, most recent first
entries.sort_by(|a, b| {
b.metadata()
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH)
.cmp(
&a.metadata()
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH),
)
});
Ok(entries[0].path())
}
/// Read messages from a session file
///
/// Creates the file if it doesn't exist, reads and deserializes all messages if it does.
pub fn read_messages(session_file: &Path) -> Result<Vec<Message>> {
let file = fs::OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(session_file)?;
let reader = io::BufReader::new(file);
let mut messages = Vec::new();
for line in reader.lines() {
messages.push(serde_json::from_str::<Message>(&line?)?);
}
Ok(messages)
}
/// Write messages to a session file
///
/// Overwrites the file with all messages in JSONL format.
pub fn persist_messages(session_file: &Path, messages: &[Message]) -> Result<()> {
let file = File::create(session_file).expect("The path specified does not exist");
let mut writer = io::BufWriter::new(file);
for message in messages {
serde_json::to_writer(&mut writer, &message)?;
writeln!(writer)?;
}
writer.flush()?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use goose::message::MessageContent;
use tempfile::tempdir;
#[test]
fn test_read_write_messages() -> Result<()> {
let dir = tempdir()?;
let file_path = dir.path().join("test.jsonl");
// Create some test messages
let messages = vec![
Message::user().with_text("Hello"),
Message::assistant().with_text("Hi there"),
];
// Write messages
persist_messages(&file_path, &messages)?;
// Read them back
let read_messages = read_messages(&file_path)?;
// Compare
assert_eq!(messages.len(), read_messages.len());
for (orig, read) in messages.iter().zip(read_messages.iter()) {
assert_eq!(orig.role, read.role);
assert_eq!(orig.content.len(), read.content.len());
// Compare first text content
if let (Some(MessageContent::Text(orig_text)), Some(MessageContent::Text(read_text))) =
(orig.content.first(), read.content.first())
{
assert_eq!(orig_text.text, read_text.text);
} else {
panic!("Messages don't match expected structure");
}
}
Ok(())
}
#[test]
fn test_empty_file() -> Result<()> {
let dir = tempdir()?;
let file_path = dir.path().join("empty.jsonl");
// Reading an empty file should return empty vec
let messages = read_messages(&file_path)?;
assert!(messages.is_empty());
Ok(())
}
#[test]
fn test_get_most_recent() -> Result<()> {
let dir = tempdir()?;
let base_path = dir.path().join("sessions");
fs::create_dir_all(&base_path)?;
// Create a few session files with different timestamps
let old_file = base_path.join("old.jsonl");
let new_file = base_path.join("new.jsonl");
// Create files with some delay to ensure different timestamps
fs::write(&old_file, "dummy content")?;
std::thread::sleep(std::time::Duration::from_secs(1));
fs::write(&new_file, "dummy content")?;
// Override the home directory for testing
// This is a bit hacky but works for testing
std::env::set_var("HOME", dir.path());
if let Ok(most_recent) = get_most_recent_session() {
assert_eq!(most_recent.file_name().unwrap(), "new.jsonl");
}
Ok(())
}
}