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(())
}
}
crates/goose-server/src/routes/mod.rs
+3 -1
View File
@@ -5,6 +5,7 @@ pub mod configs;
pub mod extension;
pub mod health;
pub mod reply;
pub mod session;
use axum::Router;
@@ -16,5 +17,6 @@ pub fn configure(state: crate::state::AppState) -> Router {
.merge(agent::routes(state.clone()))
.merge(extension::routes(state.clone()))
.merge(configs::routes(state.clone()))
.merge(config_management::routes(state))
.merge(config_management::routes(state.clone()))
.merge(session::routes(state))
}
crates/goose-server/src/routes/reply.rs
+73 -3
View File
@@ -9,6 +9,7 @@ use axum::{
use bytes::Bytes;
use futures::{stream::StreamExt, Stream};
use goose::message::{Message, MessageContent};
use goose::session;
use mcp_core::role::Role;
use serde::{Deserialize, Serialize};
@@ -27,6 +28,7 @@ use tokio_stream::wrappers::ReceiverStream;
#[derive(Debug, Deserialize)]
struct ChatRequest {
messages: Vec<Message>,
session_id: Option<String>,
}
// Custom SSE response type for streaming messages
@@ -109,6 +111,11 @@ async fn handler(
// Get messages directly from the request
let messages = request.messages;
// Generate a new session ID if not provided in the request
let session_id = request
.session_id
.unwrap_or_else(session::generate_session_id);
// Get a lock on the shared agent
let agent = state.agent.clone();
@@ -136,7 +143,16 @@ async fn handler(
}
};
let mut stream = match agent.reply(&messages).await {
// Get the provider first, before starting the reply stream
let provider = agent.provider().await;
let mut stream = match agent
.reply(
&messages,
Some(session::Identifier::Name(session_id.clone())),
)
.await
{
Ok(stream) => stream,
Err(e) => {
tracing::error!("Failed to start reply stream: {:?}", e);
@@ -158,11 +174,16 @@ async fn handler(
}
};
// Collect all messages for storage
let mut all_messages = messages.clone();
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
loop {
tokio::select! {
response = timeout(Duration::from_millis(500), stream.next()) => {
match response {
Ok(Some(Ok(message))) => {
all_messages.push(message.clone());
if let Err(e) = stream_event(MessageEvent::Message { message }, &tx).await {
tracing::error!("Error sending message through channel: {}", e);
let _ = stream_event(
@@ -173,6 +194,16 @@ async fn handler(
).await;
break;
}
// Store messages and generate description in background
let session_path = session_path.clone();
let messages = all_messages.clone();
let provider = provider.clone();
tokio::spawn(async move {
if let Err(e) = session::persist_messages(&session_path, &messages, Some(provider)).await {
tracing::error!("Failed to store session history: {:?}", e);
}
});
}
Ok(Some(Err(e))) => {
tracing::error!("Error processing message: {}", e);
@@ -214,6 +245,7 @@ async fn handler(
#[derive(Debug, Deserialize, Serialize)]
struct AskRequest {
prompt: String,
session_id: Option<String>,
}
#[derive(Debug, Serialize)]
@@ -237,16 +269,30 @@ async fn ask_handler(
return Err(StatusCode::UNAUTHORIZED);
}
// Generate a new session ID if not provided in the request
let session_id = request
.session_id
.unwrap_or_else(session::generate_session_id);
let agent = state.agent.clone();
let agent = agent.write().await;
let agent = agent.as_ref().ok_or(StatusCode::NOT_FOUND)?;
// Get the provider first, before starting the reply stream
let provider = agent.provider().await;
// Create a single message for the prompt
let messages = vec![Message::user().with_text(request.prompt)];
// Get response from agent
let mut response_text = String::new();
let mut stream = match agent.reply(&messages).await {
let mut stream = match agent
.reply(
&messages,
Some(session::Identifier::Name(session_id.clone())),
)
.await
{
Ok(stream) => stream,
Err(e) => {
tracing::error!("Failed to start reply stream: {:?}", e);
@@ -254,15 +300,20 @@ async fn ask_handler(
}
};
// Collect all messages for storage
let mut all_messages = messages.clone();
let mut response_message = Message::assistant();
while let Some(response) = stream.next().await {
match response {
Ok(message) => {
if message.role == Role::Assistant {
for content in message.content {
for content in &message.content {
if let MessageContent::Text(text) = content {
response_text.push_str(&text.text);
response_text.push('\n');
}
response_message.content.push(content.clone());
}
}
}
@@ -273,6 +324,24 @@ async fn ask_handler(
}
}
// Add the complete response message to the conversation history
if !response_message.content.is_empty() {
all_messages.push(response_message);
}
// Get the session path - file will be created when needed
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
// Store messages and generate description in background
let session_path = session_path.clone();
let messages = all_messages.clone();
let provider = provider.clone();
tokio::spawn(async move {
if let Err(e) = session::persist_messages(&session_path, &messages, Some(provider)).await {
tracing::error!("Failed to store session history: {:?}", e);
}
});
Ok(Json(AskResponse {
response: response_text.trim().to_string(),
}))
@@ -394,6 +463,7 @@ mod tests {
.body(Body::from(
serde_json::to_string(&AskRequest {
prompt: "test prompt".to_string(),
session_id: Some("test-session".to_string()),
})
.unwrap(),
))
crates/goose-server/src/routes/session.rs
+128
View File
@@ -0,0 +1,128 @@
use crate::state::AppState;
use axum::{
extract::{Path, State},
http::{HeaderMap, StatusCode},
routing::get,
Json, Router,
};
use goose::message::Message;
use goose::session;
use serde::Serialize;
#[derive(Serialize)]
struct SessionInfo {
id: String,
path: String,
modified: String,
metadata: session::SessionMetadata,
}
#[derive(Serialize)]
struct SessionListResponse {
sessions: Vec<SessionInfo>,
}
#[derive(Serialize)]
struct SessionHistoryResponse {
session_id: String,
metadata: session::SessionMetadata,
messages: Vec<Message>,
}
// List all available sessions
async fn list_sessions(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<SessionListResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
let sessions = match session::list_sessions() {
Ok(sessions) => sessions,
Err(e) => {
tracing::error!("Failed to list sessions: {:?}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
let session_infos = sessions
.into_iter()
.map(|(id, path)| {
// Get last modified time as string
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(|_| "Unknown".to_string());
// Get session description
let metadata = session::read_metadata(&path).expect("Failed to read session metadata");
SessionInfo {
id,
path: path.to_string_lossy().to_string(),
modified,
metadata,
}
})
.collect();
Ok(Json(SessionListResponse {
sessions: session_infos,
}))
}
// Get a specific session's history
async fn get_session_history(
State(state): State<AppState>,
headers: HeaderMap,
Path(session_id): Path<String>,
) -> Result<Json<SessionHistoryResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
// Read metadata
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) => {
tracing::error!("Failed to read session messages: {:?}", e);
return Err(StatusCode::NOT_FOUND);
}
};
Ok(Json(SessionHistoryResponse {
session_id,
metadata,
messages,
}))
}
// Configure routes for this module
pub fn routes(state: AppState) -> Router {
Router::new()
.route("/sessions", get(list_sessions))
.route("/sessions/:session_id", get(get_session_history))
.with_state(state)
}
crates/goose/examples/agent.rs
+1 -1
View File
@@ -30,7 +30,7 @@ async fn main() {
let messages = vec![Message::user()
.with_text("can you summarize the readme.md in this dir using just a haiku?")];
let mut stream = agent.reply(&messages).await.unwrap();
let mut stream = agent.reply(&messages, None).await.unwrap();
while let Some(message) = stream.next().await {
println!(
"{}",
crates/goose/src/agents/agent.rs
+11 -2
View File
@@ -4,10 +4,12 @@ use anyhow::Result;
use async_trait::async_trait;
use futures::stream::BoxStream;
use serde_json::Value;
use std::sync::Arc;
use super::extension::{ExtensionConfig, ExtensionResult};
use crate::message::Message;
use crate::providers::base::ProviderUsage;
use crate::providers::base::{Provider, ProviderUsage};
use crate::session;
use mcp_core::prompt::Prompt;
use mcp_core::protocol::GetPromptResult;
@@ -15,7 +17,11 @@ use mcp_core::protocol::GetPromptResult;
#[async_trait]
pub trait Agent: Send + Sync {
/// Create a stream that yields each message as it's generated by the agent
async fn reply(&self, messages: &[Message]) -> Result<BoxStream<'_, Result<Message>>>;
async fn reply(
&self,
messages: &[Message],
session_id: Option<session::Identifier>,
) -> Result<BoxStream<'_, Result<Message>>>;
/// Add a new MCP client to the agent
async fn add_extension(&mut self, config: ExtensionConfig) -> ExtensionResult<()>;
@@ -48,4 +54,7 @@ pub trait Agent: Send + Sync {
/// Get a prompt result with the given name and arguments
/// Returns the prompt text that would be used as user input
async fn get_prompt(&self, name: &str, arguments: Value) -> Result<GetPromptResult>;
/// Get a reference to the provider used by this agent
async fn provider(&self) -> Arc<Box<dyn Provider>>;
}
crates/goose/src/agents/capabilities.rs
+4 -4
View File
@@ -30,7 +30,7 @@ pub struct Capabilities {
clients: HashMap<String, McpClientBox>,
instructions: HashMap<String, String>,
resource_capable_extensions: HashSet<String>,
provider: Box<dyn Provider>,
provider: Arc<Box<dyn Provider>>,
provider_usage: Mutex<Vec<ProviderUsage>>,
system_prompt_override: Option<String>,
system_prompt_extensions: Vec<String>,
@@ -90,7 +90,7 @@ impl Capabilities {
clients: HashMap::new(),
instructions: HashMap::new(),
resource_capable_extensions: HashSet::new(),
provider,
provider: Arc::new(provider),
provider_usage: Mutex::new(Vec::new()),
system_prompt_override: None,
system_prompt_extensions: Vec::new(),
@@ -202,8 +202,8 @@ impl Capabilities {
}
/// Get a reference to the provider
pub fn provider(&self) -> &dyn Provider {
&*self.provider
pub fn provider(&self) -> Arc<Box<dyn Provider>> {
Arc::clone(&self.provider)
}
/// Record provider usage
crates/goose/src/agents/reference.rs
+21 -2
View File
@@ -3,6 +3,7 @@
use async_trait::async_trait;
use futures::stream::BoxStream;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{debug, instrument};
@@ -12,8 +13,8 @@ use crate::agents::extension::{ExtensionConfig, ExtensionResult};
use crate::message::{Message, ToolRequest};
use crate::providers::base::Provider;
use crate::providers::base::ProviderUsage;
use crate::register_agent;
use crate::token_counter::TokenCounter;
use crate::{register_agent, session};
use anyhow::{anyhow, Result};
use indoc::indoc;
use mcp_core::prompt::Prompt;
@@ -73,6 +74,7 @@ impl Agent for ReferenceAgent {
async fn reply(
&self,
messages: &[Message],
session_id: Option<session::Identifier>,
) -> anyhow::Result<BoxStream<'_, anyhow::Result<Message>>> {
let mut messages = messages.to_vec();
let reply_span = tracing::Span::current();
@@ -143,7 +145,19 @@ impl Agent for ReferenceAgent {
&messages,
&tools,
).await?;
capabilities.record_usage(usage).await;
capabilities.record_usage(usage.clone()).await;
// record usage for the session in the session file
if let Some(session_id) = session_id.clone() {
// TODO: track session_id in langfuse tracing
let session_file = session::get_path(session_id);
let mut metadata = session::read_metadata(&session_file)?;
metadata.total_tokens = usage.usage.total_tokens;
// The message count is the number of messages in the session + 1 for the response
// The message count does not include the tool response till next iteration
metadata.message_count = messages.len() + 1;
session::update_metadata(&session_file, &metadata).await?;
}
// Yield the assistant's response
yield response.clone();
@@ -233,6 +247,11 @@ impl Agent for ReferenceAgent {
Err(anyhow!("Prompt '{}' not found", name))
}
async fn provider(&self) -> Arc<Box<dyn Provider>> {
let capabilities = self.capabilities.lock().await;
capabilities.provider()
}
}
register_agent!("reference", ReferenceAgent);
crates/goose/src/agents/truncate.rs
+22 -2
View File
@@ -3,6 +3,7 @@
use async_trait::async_trait;
use futures::stream::BoxStream;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::sync::Mutex;
use tracing::{debug, error, instrument, warn};
@@ -18,6 +19,7 @@ use crate::providers::base::Provider;
use crate::providers::base::ProviderUsage;
use crate::providers::errors::ProviderError;
use crate::register_agent;
use crate::session;
use crate::token_counter::TokenCounter;
use crate::truncate::{truncate_messages, OldestFirstTruncation};
use anyhow::{anyhow, Result};
@@ -143,10 +145,11 @@ impl Agent for TruncateAgent {
}
}
#[instrument(skip(self, messages), fields(user_message))]
#[instrument(skip(self, messages, session_id), fields(user_message))]
async fn reply(
&self,
messages: &[Message],
session_id: Option<session::Identifier>,
) -> anyhow::Result<BoxStream<'_, anyhow::Result<Message>>> {
let mut messages = messages.to_vec();
let reply_span = tracing::Span::current();
@@ -223,7 +226,19 @@ impl Agent for TruncateAgent {
&tools,
).await {
Ok((response, usage)) => {
capabilities.record_usage(usage).await;
capabilities.record_usage(usage.clone()).await;
// record usage for the session in the session file
if let Some(session_id) = session_id.clone() {
// TODO: track session_id in langfuse tracing
let session_file = session::get_path(session_id);
let mut metadata = session::read_metadata(&session_file)?;
metadata.total_tokens = usage.usage.total_tokens;
// The message count is the number of messages in the session + 1 for the response
// The message count does not include the tool response till next iteration
metadata.message_count = messages.len() + 1;
session::update_metadata(&session_file, &metadata).await?;
}
// Reset truncation attempt
truncation_attempt = 0;
@@ -435,6 +450,11 @@ impl Agent for TruncateAgent {
Err(anyhow!("Prompt '{}' not found", name))
}
async fn provider(&self) -> Arc<Box<dyn Provider>> {
let capabilities = self.capabilities.lock().await;
capabilities.provider()
}
}
register_agent!("truncate", TruncateAgent);
crates/goose/src/lib.rs
+1
View File
@@ -4,6 +4,7 @@ pub mod message;
pub mod model;
pub mod prompt_template;
pub mod providers;
pub mod session;
pub mod token_counter;
pub mod tracing;
pub mod truncate;
crates/goose/src/session/mod.rs
+8
View File
@@ -0,0 +1,8 @@
pub mod storage;
// Re-export common session types and functions
pub use storage::{
ensure_session_dir, generate_description, generate_session_id, get_most_recent_session,
get_path, list_sessions, persist_messages, read_messages, read_metadata, update_metadata,
Identifier, SessionMetadata,
};
crates/goose/src/session/storage.rs
+411
View File
@@ -0,0 +1,411 @@
use crate::message::Message;
use crate::providers::base::Provider;
use anyhow::Result;
use chrono::Local;
use etcetera::{choose_app_strategy, AppStrategy, AppStrategyArgs};
use serde::{Deserialize, Serialize};
use std::fs::{self, File};
use std::io::{self, BufRead, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
/// Metadata for a session, stored as the first line in the session file
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionMetadata {
/// A short description of the session, typically 3 words or less
pub description: String,
/// Number of messages in the session
pub message_count: usize,
/// The total number of tokens used in the session. Retrieved from the provider's last usage.
pub total_tokens: Option<i32>,
}
impl SessionMetadata {
pub fn new() -> Self {
Self {
description: String::new(),
message_count: 0,
total_tokens: None,
}
}
}
impl Default for SessionMetadata {
fn default() -> Self {
Self::new()
}
}
// The single app name used for all Goose applications
const APP_NAME: &str = "goose";
#[derive(Debug, Clone, Serialize, Deserialize)]
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 app_strategy = AppStrategyArgs {
top_level_domain: "Block".to_string(),
author: "Block".to_string(),
app_name: APP_NAME.to_string(),
};
let data_dir = choose_app_strategy(app_strategy)
.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())
}
/// List all available session files
pub fn list_sessions() -> Result<Vec<(String, PathBuf)>> {
let session_dir = ensure_session_dir()?;
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)
}
/// Generate a session ID using timestamp format (yyyymmdd_hhmmss)
pub fn generate_session_id() -> String {
Local::now().format("%Y%m%d_%H%M%S").to_string()
}
/// Read messages from a session file
///
/// Creates the file if it doesn't exist, reads and deserializes all messages if it does.
/// The first line of the file is expected to be metadata, and the rest are messages.
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 lines = reader.lines();
let mut messages = Vec::new();
// Read the first line as metadata or create default if empty/missing
if let Some(line) = lines.next() {
let line = line?;
// Try to parse as metadata, but if it fails, treat it as a message
if let Ok(_metadata) = serde_json::from_str::<SessionMetadata>(&line) {
// Metadata successfully parsed, continue with the rest of the lines as messages
} else {
// This is not metadata, it's a message
messages.push(serde_json::from_str::<Message>(&line)?);
}
}
// Read the rest of the lines as messages
for line in lines {
messages.push(serde_json::from_str::<Message>(&line?)?);
}
Ok(messages)
}
/// Read session metadata from a session file
///
/// Returns default empty metadata if the file doesn't exist or has no metadata.
pub fn read_metadata(session_file: &Path) -> Result<SessionMetadata> {
if !session_file.exists() {
return Ok(SessionMetadata::new());
}
let file = fs::File::open(session_file)?;
let mut reader = io::BufReader::new(file);
let mut first_line = String::new();
// Read just the first line
if reader.read_line(&mut first_line)? > 0 {
// Try to parse as metadata
match serde_json::from_str::<SessionMetadata>(&first_line) {
Ok(metadata) => Ok(metadata),
Err(_) => {
// If the first line isn't metadata, return default
Ok(SessionMetadata::new())
}
}
} else {
// Empty file, return default
Ok(SessionMetadata::new())
}
}
/// Write messages to a session file with metadata
///
/// Overwrites the file with metadata as the first line, followed by all messages in JSONL format.
/// If a provider is supplied, it will automatically generate a description when appropriate.
pub async fn persist_messages(
session_file: &Path,
messages: &[Message],
provider: Option<Arc<Box<dyn Provider>>>,
) -> Result<()> {
// Read existing metadata
let mut metadata = read_metadata(session_file)?;
// Count user messages
let user_message_count = messages
.iter()
.filter(|m| m.role == mcp_core::role::Role::User)
.filter(|m| !m.as_concat_text().trim().is_empty())
.count();
// Check if we need to update the description (after 1st or 3rd user message)
if let Some(provider) = provider {
if user_message_count < 4 {
// Generate description
let mut description_prompt = "Based on the conversation so far, provide a concise header for this session in 4 words or less. This will be used for finding the session later in a UI with limited space - reply *ONLY* with the header. Avoid filler words such as help, summary, exchange, request etc that do not help distinguish different conversations.".to_string();
// get context from messages so far
let context: Vec<String> = messages.iter().map(|m| m.as_concat_text()).collect();
if !context.is_empty() {
description_prompt = format!(
"Here are the first few user messages:\n{}\n\n{}",
context.join("\n"),
description_prompt
);
}
// Generate the description
let message = Message::user().with_text(&description_prompt);
match provider
.complete(
"Reply with only a description in four words or less.",
&[message],
&[],
)
.await
{
Ok((response, _)) => {
metadata.description = response.as_concat_text();
}
Err(e) => {
tracing::error!("Failed to generate session description: {:?}", e);
}
}
}
}
// Write the file with metadata and messages
save_messages_with_metadata(session_file, &metadata, messages)
}
/// Write messages to a session file with the provided metadata
///
/// Overwrites the file with metadata as the first line, followed by all messages in JSONL format.
pub fn save_messages_with_metadata(
session_file: &Path,
metadata: &SessionMetadata,
messages: &[Message],
) -> Result<()> {
let file = File::create(session_file).expect("The path specified does not exist");
let mut writer = io::BufWriter::new(file);
// Write metadata as the first line
serde_json::to_writer(&mut writer, &metadata)?;
writeln!(writer)?;
// Write all messages
for message in messages {
serde_json::to_writer(&mut writer, &message)?;
writeln!(writer)?;
}
writer.flush()?;
Ok(())
}
/// Generate a description for the session using the provider
///
/// This function is called when appropriate to generate a short description
/// of the session based on the conversation history.
pub async fn generate_description(
session_file: &Path,
messages: &[Message],
provider: &dyn Provider,
) -> Result<()> {
// Create a special message asking for a 3-word description
let mut description_prompt = "Based on the conversation so far, provide a concise description of this session in 4 words or less. This will be used for finding the session later in a UI with limited space - reply *ONLY* with the description".to_string();
// get context from messages so far
let context: Vec<String> = messages
.iter()
.filter(|m| m.role == mcp_core::role::Role::User)
.take(3) // Use up to first 3 user messages for context
.map(|m| m.as_concat_text())
.collect();
if !context.is_empty() {
description_prompt = format!(
"Here are the first few user messages:\n{}\n\n{}",
context.join("\n"),
description_prompt
);
}
// Generate the description
let message = Message::user().with_text(&description_prompt);
let result = provider
.complete(
"Reply with only a description in four words or less",
&[message],
&[],
)
.await?;
let description = result.0.as_concat_text();
// Read current metadata
let mut metadata = read_metadata(session_file)?;
// Update description
metadata.description = description;
// Update the file with the new metadata and existing messages
update_metadata(session_file, &metadata).await?;
Ok(())
}
/// Update only the metadata in a session file, preserving all messages
pub async fn update_metadata(session_file: &Path, metadata: &SessionMetadata) -> Result<()> {
// Read all messages from the file
let messages = read_messages(session_file)?;
// Rewrite the file with the new metadata and existing messages
save_messages_with_metadata(session_file, metadata, &messages)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::message::MessageContent;
use tempfile::tempdir;
#[tokio::test]
async 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, None).await?;
// 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_generate_session_id() {
let id = generate_session_id();
// Check that it follows the timestamp format (yyyymmdd_hhmmss)
assert_eq!(id.len(), 15); // 8 chars for date + 1 for underscore + 6 for time
assert!(id.contains('_'));
// Split by underscore and check parts
let parts: Vec<&str> = id.split('_').collect();
assert_eq!(parts.len(), 2);
// Date part should be 8 digits
assert_eq!(parts[0].len(), 8);
// Time part should be 6 digits
assert_eq!(parts[1].len(), 6);
}
}
crates/goose/tests/truncate_agent.rs
+1 -1
View File
@@ -123,7 +123,7 @@ async fn run_truncate_test(
),
];
let reply_stream = agent.reply(&messages).await?;
let reply_stream = agent.reply(&messages, None).await?;
tokio::pin!(reply_stream);
let mut responses = Vec::new();