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

feat: store working directory for sessions (#1559)

Browse Source
This commit is contained in:
Salman Mohammed
2025-03-07 11:12:57 -05:00
committed by GitHub
parent 2d0cd8e245
commit 32f20cd690
20 changed files with 334 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/builder.rs
+52 -33
View File
@@ -35,7 +35,59 @@ pub async fn build_session(
let mut agent = AgentFactory::create(&AgentFactory::configured_version(), provider)
.expect("Failed to create agent");
// Handle session file resolution and resuming
let session_file = if resume {
if let Some(identifier) = identifier {
let session_file = session::get_path(identifier);
if !session_file.exists() {
output::render_error(&format!(
"Cannot resume session {} - no such session exists",
style(session_file.display()).cyan()
));
process::exit(1);
}
session_file
} else {
// Try to resume most recent session
match session::get_most_recent_session() {
Ok(file) => file,
Err(_) => {
output::render_error("Cannot resume - no previous sessions found");
process::exit(1);
}
}
}
} else {
// Create new session with provided name/path or generated name
let id = match identifier {
Some(identifier) => identifier,
None => Identifier::Name(session::generate_session_id()),
};
// Just get the path - file will be created when needed
session::get_path(id)
};
if resume {
// 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);
});
// Ask user if they want to change the working directory
let change_workdir = cliclack::confirm(format!("{} The working directory of this session was set to {}. It does not match the current working directory. Would you like to change it?", style("WARNING:").yellow(), style(metadata.working_dir.display()).cyan()))
.initial_value(true)
.interact().expect("Failed to get user input");
if change_workdir {
std::env::set_current_dir(metadata.working_dir).unwrap();
}
}
// Setup extensions for the agent
// Extensions need to be added after the session is created because we change directory when resuming a session
for extension in ExtensionManager::get_all().expect("should load extensions") {
if extension.enabled {
let config = extension.config.clone();
@@ -59,39 +111,6 @@ pub async fn build_session(
}
}
// Handle session file resolution and resuming
let session_file = if resume {
if let Some(identifier) = identifier {
let session_file = session::get_path(identifier);
if !session_file.exists() {
output::render_error(&format!(
"Cannot resume session {} - no such session exists",
style(session_file.display()).cyan()
));
process::exit(1);
}
session_file
} else {
// Try to resume most recent session
match session::get_most_recent_session() {
Ok(file) => file,
Err(_) => {
output::render_error("Cannot resume - no previous sessions found");
process::exit(1);
}
}
}
} else {
// Create new session with provided name/path or generated name
let id = match identifier {
Some(identifier) => identifier,
None => Identifier::Name(session::generate_session_id()),
};
// Just get the path - file will be created when needed
session::get_path(id)
};
// Create new session
let mut session = Session::new(agent, session_file.clone(), debug);
crates/goose-cli/src/session/mod.rs
+12 -3
View File
@@ -13,7 +13,7 @@ use completion::GooseCompleter;
use etcetera::choose_app_strategy;
use etcetera::AppStrategy;
use goose::agents::extension::{Envs, ExtensionConfig};
use goose::agents::Agent;
use goose::agents::{Agent, SessionConfig};
use goose::config::Config;
use goose::message::{Message, MessageContent};
use goose::session;
@@ -199,7 +199,6 @@ 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));
// Get the provider from the agent for description generation
let provider = self.agent.provider().await;
@@ -436,7 +435,17 @@ impl Session {
async fn process_agent_response(&mut self, interactive: bool) -> Result<()> {
let session_id = session::Identifier::Path(self.session_file.clone());
let mut stream = self.agent.reply(&self.messages, Some(session_id)).await?;
let mut stream = self
.agent
.reply(
&self.messages,
Some(SessionConfig {
id: session_id,
working_dir: std::env::current_dir()
.expect("failed to get current session working directory"),
}),
)
.await?;
use futures::StreamExt;
loop {
crates/goose-cli/src/session/output.rs
+7
View File
@@ -479,6 +479,13 @@ pub fn display_session_info(resume: bool, provider: &str, model: &str, session_f
style("logging to").dim(),
style(session_file.display()).dim().cyan(),
);
println!(
" {} {}",
style("working directory:").dim(),
style(std::env::current_dir().unwrap().display())
.cyan()
.dim()
);
}
pub fn display_greeting() {
crates/goose-server/src/routes/reply.rs
+19 -4
View File
@@ -8,14 +8,18 @@ use axum::{
};
use bytes::Bytes;
use futures::{stream::StreamExt, Stream};
use goose::message::{Message, MessageContent};
use goose::session;
use goose::{
agents::SessionConfig,
message::{Message, MessageContent},
};
use mcp_core::role::Role;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{
convert::Infallible,
path::PathBuf,
pin::Pin,
task::{Context, Poll},
time::Duration,
@@ -29,6 +33,7 @@ use tokio_stream::wrappers::ReceiverStream;
struct ChatRequest {
messages: Vec<Message>,
session_id: Option<String>,
session_working_dir: String,
}
// Custom SSE response type for streaming messages
@@ -108,8 +113,8 @@ async fn handler(
let (tx, rx) = mpsc::channel(100);
let stream = ReceiverStream::new(rx);
// Get messages directly from the request
let messages = request.messages;
let session_working_dir = request.session_working_dir;
// Generate a new session ID if not provided in the request
let session_id = request
@@ -149,7 +154,10 @@ async fn handler(
let mut stream = match agent
.reply(
&messages,
Some(session::Identifier::Name(session_id.clone())),
Some(SessionConfig {
id: session::Identifier::Name(session_id.clone()),
working_dir: PathBuf::from(session_working_dir),
}),
)
.await
{
@@ -246,6 +254,7 @@ async fn handler(
struct AskRequest {
prompt: String,
session_id: Option<String>,
session_working_dir: String,
}
#[derive(Debug, Serialize)]
@@ -269,6 +278,8 @@ async fn ask_handler(
return Err(StatusCode::UNAUTHORIZED);
}
let session_working_dir = request.session_working_dir;
// Generate a new session ID if not provided in the request
let session_id = request
.session_id
@@ -289,7 +300,10 @@ async fn ask_handler(
let mut stream = match agent
.reply(
&messages,
Some(session::Identifier::Name(session_id.clone())),
Some(SessionConfig {
id: session::Identifier::Name(session_id.clone()),
working_dir: PathBuf::from(session_working_dir),
}),
)
.await
{
@@ -464,6 +478,7 @@ mod tests {
serde_json::to_string(&AskRequest {
prompt: "test prompt".to_string(),
session_id: Some("test-session".to_string()),
session_working_dir: "test-working-dir".to_string(),
})
.unwrap(),
))
crates/goose/src/agents/agent.rs
+12 -1
View File
@@ -1,8 +1,10 @@
use std::collections::HashMap;
use std::path::PathBuf;
use anyhow::Result;
use async_trait::async_trait;
use futures::stream::BoxStream;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::Arc;
@@ -13,6 +15,15 @@ use crate::session;
use mcp_core::prompt::Prompt;
use mcp_core::protocol::GetPromptResult;
/// Session configuration for an agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionConfig {
/// Unique identifier for the session
pub id: session::Identifier,
/// Working directory for the session
pub working_dir: PathBuf,
}
/// Core trait defining the behavior of an Agent
#[async_trait]
pub trait Agent: Send + Sync {
@@ -20,7 +31,7 @@ pub trait Agent: Send + Sync {
async fn reply(
&self,
messages: &[Message],
session_id: Option<session::Identifier>,
session: Option<SessionConfig>,
) -> Result<BoxStream<'_, Result<Message>>>;
/// Add a new MCP client to the agent
crates/goose/src/agents/mod.rs
+1 -1
View File
@@ -7,7 +7,7 @@ mod reference;
mod summarize;
mod truncate;
pub use agent::Agent;
pub use agent::{Agent, SessionConfig};
pub use capabilities::Capabilities;
pub use extension::ExtensionConfig;
pub use factory::{register_agent, AgentFactory};
crates/goose/src/agents/reference.rs
+6 -4
View File
@@ -7,6 +7,7 @@ use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{debug, instrument};
use super::agent::SessionConfig;
use super::Agent;
use crate::agents::capabilities::Capabilities;
use crate::agents::extension::{ExtensionConfig, ExtensionResult};
@@ -70,11 +71,11 @@ impl Agent for ReferenceAgent {
// TODO implement
}
#[instrument(skip(self, messages), fields(user_message))]
#[instrument(skip(self, messages, session), fields(user_message))]
async fn reply(
&self,
messages: &[Message],
session_id: Option<session::Identifier>,
session: Option<SessionConfig>,
) -> anyhow::Result<BoxStream<'_, anyhow::Result<Message>>> {
let mut messages = messages.to_vec();
let reply_span = tracing::Span::current();
@@ -148,10 +149,11 @@ impl Agent for ReferenceAgent {
capabilities.record_usage(usage.clone()).await;
// record usage for the session in the session file
if let Some(session_id) = session_id.clone() {
if let Some(session) = session.clone() {
// TODO: track session_id in langfuse tracing
let session_file = session::get_path(session_id);
let session_file = session::get_path(session.id);
let mut metadata = session::read_metadata(&session_file)?;
metadata.working_dir = session.working_dir;
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
crates/goose/src/agents/summarize.rs
+6 -4
View File
@@ -9,6 +9,7 @@ use tokio::sync::mpsc;
use tokio::sync::Mutex;
use tracing::{debug, error, instrument, warn};
use super::agent::SessionConfig;
use super::detect_read_only_tools;
use super::Agent;
use crate::agents::capabilities::Capabilities;
@@ -162,11 +163,11 @@ impl Agent for SummarizeAgent {
}
}
#[instrument(skip(self, messages), fields(user_message))]
#[instrument(skip(self, messages, session), fields(user_message))]
async fn reply(
&self,
messages: &[Message],
session_id: Option<session::Identifier>,
session: Option<SessionConfig>,
) -> anyhow::Result<BoxStream<'_, anyhow::Result<Message>>> {
let mut messages = messages.to_vec();
let reply_span = tracing::Span::current();
@@ -246,10 +247,11 @@ impl Agent for SummarizeAgent {
capabilities.record_usage(usage.clone()).await;
// record usage for the session in the session file
if let Some(session_id) = session_id.clone() {
if let Some(session) = session.clone() {
// TODO: track session_id in langfuse tracing
let session_file = session::get_path(session_id);
let session_file = session::get_path(session.id);
let mut metadata = session::read_metadata(&session_file)?;
metadata.working_dir = session.working_dir;
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
crates/goose/src/agents/truncate.rs
+6 -4
View File
@@ -8,6 +8,7 @@ use tokio::sync::mpsc;
use tokio::sync::Mutex;
use tracing::{debug, error, instrument, warn};
use super::agent::SessionConfig;
use super::detect_read_only_tools;
use super::Agent;
use crate::agents::capabilities::Capabilities;
@@ -145,11 +146,11 @@ impl Agent for TruncateAgent {
}
}
#[instrument(skip(self, messages, session_id), fields(user_message))]
#[instrument(skip(self, messages, session), fields(user_message))]
async fn reply(
&self,
messages: &[Message],
session_id: Option<session::Identifier>,
session: Option<SessionConfig>,
) -> anyhow::Result<BoxStream<'_, anyhow::Result<Message>>> {
let mut messages = messages.to_vec();
let reply_span = tracing::Span::current();
@@ -229,10 +230,11 @@ impl Agent for TruncateAgent {
capabilities.record_usage(usage.clone()).await;
// record usage for the session in the session file
if let Some(session_id) = session_id.clone() {
if let Some(session) = session.clone() {
// TODO: track session_id in langfuse tracing
let session_file = session::get_path(session_id);
let session_file = session::get_path(session.id);
let mut metadata = session::read_metadata(&session_file)?;
metadata.working_dir = session.working_dir;
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
crates/goose/src/session/storage.rs
+41 -6
View File
@@ -9,9 +9,18 @@ use std::io::{self, BufRead, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
fn get_home_dir() -> PathBuf {
choose_app_strategy(crate::config::APP_STRATEGY.clone())
.expect("goose requires a home dir")
.home_dir()
.to_path_buf()
}
/// Metadata for a session, stored as the first line in the session file
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone, Serialize)]
pub struct SessionMetadata {
/// Working directory for the session
pub working_dir: PathBuf,
/// A short description of the session, typically 3 words or less
pub description: String,
/// Number of messages in the session
@@ -20,9 +29,35 @@ pub struct SessionMetadata {
pub total_tokens: Option<i32>,
}
// Custom deserializer to handle old sessions without working_dir
impl<'de> Deserialize<'de> for SessionMetadata {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(Deserialize)]
struct Helper {
description: String,
message_count: usize,
total_tokens: Option<i32>,
working_dir: Option<PathBuf>,
}
let helper = Helper::deserialize(deserializer)?;
Ok(SessionMetadata {
description: helper.description,
message_count: helper.message_count,
total_tokens: helper.total_tokens,
working_dir: helper.working_dir.unwrap_or_else(get_home_dir),
})
}
}
impl SessionMetadata {
pub fn new() -> Self {
pub fn new(working_dir: PathBuf) -> Self {
Self {
working_dir,
description: String::new(),
message_count: 0,
total_tokens: None,
@@ -32,7 +67,7 @@ impl SessionMetadata {
impl Default for SessionMetadata {
fn default() -> Self {
Self::new()
Self::new(get_home_dir())
}
}
@@ -168,7 +203,7 @@ pub fn read_messages(session_file: &Path) -> Result<Vec<Message>> {
/// 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());
return Ok(SessionMetadata::default());
}
let file = fs::File::open(session_file)?;
@@ -182,12 +217,12 @@ pub fn read_metadata(session_file: &Path) -> Result<SessionMetadata> {
Ok(metadata) => Ok(metadata),
Err(_) => {
// If the first line isn't metadata, return default
Ok(SessionMetadata::new())
Ok(SessionMetadata::default())
}
}
} else {
// Empty file, return default
Ok(SessionMetadata::new())
Ok(SessionMetadata::default())
}
}