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/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())
}
}