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

Add support for changing working dir and extensions in same window/session (#6057)

Browse Source
This commit is contained in:
Zane
2026-01-08 16:48:15 -08:00
committed by GitHub
parent 4a60f026d5
commit 9a01fcb740
49 changed files with 1865 additions and 1204 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+124 -1
View File
@@ -13,7 +13,7 @@ use super::platform_tools;
use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
use crate::action_required_manager::ActionRequiredManager;
use crate::agents::extension::{ExtensionConfig, ExtensionResult, ToolInfo};
use crate::agents::extension_manager::{get_parameter_names, ExtensionManager};
use crate::agents::extension_manager::{get_parameter_names, normalize, ExtensionManager};
use crate::agents::extension_manager_extension::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE;
use crate::agents::final_output_tool::{FINAL_OUTPUT_CONTINUATION_MESSAGE, FINAL_OUTPUT_TOOL_NAME};
use crate::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME;
@@ -78,6 +78,14 @@ pub struct ToolCategorizeResult {
pub filtered_response: Message,
}
#[derive(Debug, Clone, serde::Serialize, utoipa::ToSchema)]
pub struct ExtensionLoadResult {
pub name: String,
pub success: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
/// The main goose Agent
pub struct Agent {
pub(super) provider: SharedProvider,
@@ -566,6 +574,91 @@ impl Agent {
Ok(())
}
/// Save current extension state to session by session_id
pub async fn persist_extension_state(&self, session_id: &str) -> Result<()> {
let extension_configs = self.extension_manager.get_extension_configs().await;
let extensions_state = EnabledExtensionsState::new(extension_configs);
let session = SessionManager::get_session(session_id, false).await?;
let mut extension_data = session.extension_data.clone();
extensions_state
.to_extension_data(&mut extension_data)
.map_err(|e| anyhow!("Failed to serialize extension state: {}", e))?;
SessionManager::update_session(session_id)
.extension_data(extension_data)
.apply()
.await?;
Ok(())
}
/// Load extensions from session into the agent
/// Skips extensions that are already loaded
pub async fn load_extensions_from_session(
self: &Arc<Self>,
session: &Session,
) -> Vec<ExtensionLoadResult> {
let session_extensions =
EnabledExtensionsState::from_extension_data(&session.extension_data);
let enabled_configs = match session_extensions {
Some(state) => state.extensions,
None => {
tracing::warn!(
"No extensions found in session {}. This is unexpected.",
session.id
);
return vec![];
}
};
let extension_futures = enabled_configs
.into_iter()
.map(|config| {
let config_clone = config.clone();
let agent_ref = self.clone();
async move {
let name = config_clone.name().to_string();
let normalized_name = normalize(&name);
if agent_ref
.extension_manager
.is_extension_enabled(&normalized_name)
.await
{
tracing::debug!("Extension {} already loaded, skipping", name);
return ExtensionLoadResult {
name,
success: true,
error: None,
};
}
match agent_ref.add_extension(config_clone).await {
Ok(_) => ExtensionLoadResult {
name,
success: true,
error: None,
},
Err(e) => {
let error_msg = e.to_string();
warn!("Failed to load extension {}: {}", name, error_msg);
ExtensionLoadResult {
name,
success: false,
error: Some(error_msg),
}
}
}
}
})
.collect::<Vec<_>>();
futures::future::join_all(extension_futures).await
}
pub async fn add_extension(&self, extension: ExtensionConfig) -> ExtensionResult<()> {
match &extension {
ExtensionConfig::Frontend {
@@ -937,6 +1030,7 @@ impl Agent {
let conversation_with_moim = super::moim::inject_moim(
conversation.clone(),
&self.extension_manager,
&working_dir,
).await;
let mut stream = Self::stream_response_from_provider(
@@ -1324,6 +1418,35 @@ impl Agent {
.context("Failed to persist provider config to session")
}
/// Restore the provider from session data or fall back to global config
/// This is used when resuming a session to restore the provider state
pub async fn restore_provider_from_session(&self, session: &Session) -> Result<()> {
let config = Config::global();
let provider_name = session
.provider_name
.clone()
.or_else(|| config.get_goose_provider().ok())
.ok_or_else(|| anyhow!("Could not configure agent: missing provider"))?;
let model_config = match session.model_config.clone() {
Some(saved_config) => saved_config,
None => {
let model_name = config
.get_goose_model()
.map_err(|_| anyhow!("Could not configure agent: missing model"))?;
crate::model::ModelConfig::new(&model_name)
.map_err(|e| anyhow!("Could not configure agent: invalid model {}", e))?
}
};
let provider = crate::providers::create(&provider_name, model_config)
.await
.map_err(|e| anyhow!("Could not create provider: {}", e))?;
self.update_provider(provider, &session.id).await
}
/// Override the system prompt with a custom template
pub async fn override_system_prompt(&self, template: String) {
let mut prompt_manager = self.prompt_manager.lock().await;
crates/goose/src/agents/extension_manager.rs
+90 -31
View File
@@ -133,7 +133,7 @@ impl ResourceItem {
/// Sanitizes a string by replacing invalid characters with underscores.
/// Valid characters match [a-zA-Z0-9_-]
fn normalize(input: String) -> String {
pub fn normalize(input: &str) -> String {
let mut result = String::with_capacity(input.len());
for c in input.chars() {
result.push(match c {
@@ -153,7 +153,7 @@ fn generate_extension_name(
let base = server_info
.and_then(|info| {
let name = info.server_info.name.as_str();
(!name.is_empty()).then(|| normalize(name.to_string()))
(!name.is_empty()).then(|| normalize(name))
})
.unwrap_or_else(|| "unnamed".to_string());
@@ -219,6 +219,7 @@ async fn child_process_client(
mut command: Command,
timeout: &Option<u64>,
provider: SharedProvider,
working_dir: Option<&PathBuf>,
) -> ExtensionResult<McpClient> {
#[cfg(unix)]
command.process_group(0);
@@ -228,6 +229,27 @@ async fn child_process_client(
command.env("PATH", path);
}
// Use explicitly passed working_dir, falling back to GOOSE_WORKING_DIR env var
let effective_working_dir = working_dir
.map(|p| p.to_path_buf())
.or_else(|| std::env::var("GOOSE_WORKING_DIR").ok().map(PathBuf::from));
if let Some(ref dir) = effective_working_dir {
if dir.exists() && dir.is_dir() {
tracing::info!("Setting MCP process working directory: {:?}", dir);
command.current_dir(dir);
// Also set GOOSE_WORKING_DIR env var for the child process
command.env("GOOSE_WORKING_DIR", dir);
} else {
tracing::warn!(
"Working directory doesn't exist or isn't a directory: {:?}",
dir
);
}
} else {
tracing::info!("No working directory specified, using default");
}
let (transport, mut stderr) = TokioChildProcess::builder(command)
.stderr(Stdio::piped())
.spawn()?;
@@ -422,25 +444,6 @@ async fn create_streamable_http_client(
}
}
async fn create_stdio_client(
cmd: &str,
args: &[String],
all_envs: HashMap<String, String>,
timeout: &Option<u64>,
provider: SharedProvider,
) -> ExtensionResult<Box<dyn McpClientTrait>> {
extension_malware_check::deny_if_malicious_cmd_args(cmd, args).await?;
let resolved_cmd = resolve_command(cmd);
let command = Command::new(resolved_cmd).configure(|command| {
command.args(args).envs(all_envs);
});
Ok(Box::new(
child_process_client(command, timeout, provider).await?,
))
}
impl ExtensionManager {
pub fn new(provider: SharedProvider) -> Self {
Self {
@@ -466,6 +469,22 @@ impl ExtensionManager {
self.context.lock().await.clone()
}
/// Resolve the working directory for an extension.
/// Priority: session working_dir > current_dir
async fn resolve_working_dir(&self) -> PathBuf {
// Try to get working_dir from session via context
if let Some(ref session_id) = self.context.lock().await.session_id {
if let Ok(session) =
crate::session::SessionManager::get_session(session_id, false).await
{
return session.working_dir;
}
}
// Fall back to current_dir
std::env::current_dir().unwrap_or_default()
}
pub async fn supports_resources(&self) -> bool {
self.extensions
.lock()
@@ -476,12 +495,15 @@ impl ExtensionManager {
pub async fn add_extension(&self, config: ExtensionConfig) -> ExtensionResult<()> {
let config_name = config.key().to_string();
let sanitized_name = normalize(config_name.clone());
let sanitized_name = normalize(&config_name);
if self.extensions.lock().await.contains_key(&sanitized_name) {
return Ok(());
}
// Resolve working_dir: session > current_dir
let effective_working_dir = self.resolve_working_dir().await;
let mut temp_dir = None;
let client: Box<dyn McpClientTrait> = match &config {
@@ -519,7 +541,24 @@ impl ExtensionManager {
..
} => {
let all_envs = merge_environments(envs, env_keys, &sanitized_name).await?;
create_stdio_client(cmd, args, all_envs, timeout, self.provider.clone()).await?
// Check for malicious packages before launching the process
extension_malware_check::deny_if_malicious_cmd_args(cmd, args).await?;
let cmd = resolve_command(cmd);
let command = Command::new(cmd).configure(|command| {
command.args(args).envs(all_envs);
});
let client = child_process_client(
command,
timeout,
self.provider.clone(),
Some(&effective_working_dir),
)
.await?;
Box::new(client)
}
ExtensionConfig::Builtin { name, timeout, .. } => {
let cmd = std::env::current_exe()
@@ -540,10 +579,17 @@ impl ExtensionManager {
let command = Command::new(cmd).configure(|command| {
command.arg("mcp").arg(name);
});
Box::new(child_process_client(command, timeout, self.provider.clone()).await?)
let client = child_process_client(
command,
timeout,
self.provider.clone(),
Some(&effective_working_dir),
)
.await?;
Box::new(client)
}
ExtensionConfig::Platform { name, .. } => {
let normalized_key = normalize(name.clone());
let normalized_key = normalize(name);
let def = PLATFORM_EXTENSIONS
.get(normalized_key.as_str())
.ok_or_else(|| {
@@ -572,7 +618,15 @@ impl ExtensionManager {
command.arg("python").arg(file_path.to_str().unwrap());
});
Box::new(child_process_client(command, timeout, self.provider.clone()).await?)
let client = child_process_client(
command,
timeout,
self.provider.clone(),
Some(&effective_working_dir),
)
.await?;
Box::new(client)
}
ExtensionConfig::Frontend { .. } => {
return Err(ExtensionError::ConfigError(
@@ -630,7 +684,7 @@ impl ExtensionManager {
/// Get aggregated usage statistics
pub async fn remove_extension(&self, name: &str) -> ExtensionResult<()> {
let sanitized_name = normalize(name.to_string());
let sanitized_name = normalize(name);
self.extensions.lock().await.remove(&sanitized_name);
Ok(())
}
@@ -1247,10 +1301,14 @@ impl ExtensionManager {
.map(|ext| ext.get_client())
}
pub async fn collect_moim(&self) -> Option<String> {
pub async fn collect_moim(&self, working_dir: &std::path::Path) -> Option<String> {
// Use minute-level granularity to prevent conversation changes every second
let timestamp = chrono::Local::now().format("%Y-%m-%d %H:%M:00").to_string();
let mut content = format!("<info-msg>\nIt is currently {}\n", timestamp);
let mut content = format!(
"<info-msg>\nIt is currently {}\nWorking directory: {}\n",
timestamp,
working_dir.display()
);
let platform_clients: Vec<(String, McpClientBox)> = {
let extensions = self.extensions.lock().await;
@@ -1308,7 +1366,7 @@ mod tests {
client: McpClientBox,
available_tools: Vec<String>,
) {
let sanitized_name = normalize(name.clone());
let sanitized_name = normalize(&name);
let config = ExtensionConfig::Builtin {
name: name.clone(),
display_name: Some(name.clone()),
@@ -1760,8 +1818,9 @@ mod tests {
#[tokio::test]
async fn test_collect_moim_uses_minute_granularity() {
let em = ExtensionManager::new_without_provider();
let working_dir = std::path::Path::new("/tmp");
if let Some(moim) = em.collect_moim().await {
if let Some(moim) = em.collect_moim(working_dir).await {
// Timestamp should end with :00 (seconds fixed to 00)
assert!(
moim.contains(":00\n"),
crates/goose/src/agents/mod.rs
+2 -2
View File
@@ -24,10 +24,10 @@ pub(crate) mod todo_extension;
mod tool_execution;
pub mod types;
pub use agent::{Agent, AgentEvent};
pub use agent::{Agent, AgentEvent, ExtensionLoadResult};
pub use execute_commands::COMPACT_TRIGGERS;
pub use extension::ExtensionConfig;
pub use extension_manager::ExtensionManager;
pub use extension_manager::{normalize, ExtensionManager};
pub use prompt_manager::PromptManager;
pub use subagent_task_config::TaskConfig;
pub use types::{FrontendTool, RetryConfig, SessionConfig, SuccessCheck};
crates/goose/src/agents/moim.rs
+12 -4
View File
@@ -2,6 +2,7 @@ use crate::agents::extension_manager::ExtensionManager;
use crate::conversation::message::Message;
use crate::conversation::{fix_conversation, Conversation};
use rmcp::model::Role;
use std::path::Path;
// Test-only utility. Do not use in production code. No `test` directive due to call outside crate.
thread_local! {
@@ -11,12 +12,13 @@ thread_local! {
pub async fn inject_moim(
conversation: Conversation,
extension_manager: &ExtensionManager,
working_dir: &Path,
) -> Conversation {
if SKIP.with(|f| f.get()) {
return conversation;
}
if let Some(moim) = extension_manager.collect_moim().await {
if let Some(moim) = extension_manager.collect_moim(working_dir).await {
let mut messages = conversation.messages().clone();
let idx = messages
.iter()
@@ -45,17 +47,19 @@ pub async fn inject_moim(
mod tests {
use super::*;
use rmcp::model::CallToolRequestParam;
use std::path::PathBuf;
#[tokio::test]
async fn test_moim_injection_before_assistant() {
let em = ExtensionManager::new_without_provider();
let working_dir = PathBuf::from("/test/dir");
let conv = Conversation::new_unvalidated(vec![
Message::user().with_text("Hello"),
Message::assistant().with_text("Hi"),
Message::user().with_text("Bye"),
]);
let result = inject_moim(conv, &em).await;
let result = inject_moim(conv, &em, &working_dir).await;
let msgs = result.messages();
assert_eq!(msgs.len(), 3);
@@ -70,14 +74,16 @@ mod tests {
.join("");
assert!(merged_content.contains("Hello"));
assert!(merged_content.contains("<info-msg>"));
assert!(merged_content.contains("Working directory: /test/dir"));
}
#[tokio::test]
async fn test_moim_injection_no_assistant() {
let em = ExtensionManager::new_without_provider();
let working_dir = PathBuf::from("/test/dir");
let conv = Conversation::new_unvalidated(vec![Message::user().with_text("Hello")]);
let result = inject_moim(conv, &em).await;
let result = inject_moim(conv, &em, &working_dir).await;
assert_eq!(result.messages().len(), 1);
@@ -89,11 +95,13 @@ mod tests {
.join("");
assert!(merged_content.contains("Hello"));
assert!(merged_content.contains("<info-msg>"));
assert!(merged_content.contains("Working directory: /test/dir"));
}
#[tokio::test]
async fn test_moim_with_tool_calls() {
let em = ExtensionManager::new_without_provider();
let working_dir = PathBuf::from("/test/dir");
let conv = Conversation::new_unvalidated(vec![
Message::user().with_text("Search for something"),
@@ -135,7 +143,7 @@ mod tests {
),
]);
let result = inject_moim(conv, &em).await;
let result = inject_moim(conv, &em, &working_dir).await;
let msgs = result.messages();
assert_eq!(msgs.len(), 6);