tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 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-mcp/src/developer/rmcp_developer.rs
+5 -1
View File
@@ -968,6 +968,10 @@ impl DeveloperServer {
.and_then(|s| s.to_str())
.unwrap_or("bash");
let working_dir = std::env::var("GOOSE_WORKING_DIR")
.ok()
.map(std::path::PathBuf::from);
if let Some(ref env_file) = self.bash_env_file {
if shell_name == "bash" {
shell_config.envs.push((
@@ -977,7 +981,7 @@ impl DeveloperServer {
}
}
let mut command = configure_shell_command(&shell_config, command);
let mut command = configure_shell_command(&shell_config, command, working_dir.as_deref());
if self.extend_path_with_shell {
if let Err(e) = get_shell_path_dirs()
crates/goose-mcp/src/developer/shell.rs
+6
View File
@@ -109,8 +109,14 @@ pub fn normalize_line_endings(text: &str) -> String {
pub fn configure_shell_command(
shell_config: &ShellConfig,
command: &str,
working_dir: Option<&std::path::Path>,
) -> tokio::process::Command {
let mut command_builder = tokio::process::Command::new(&shell_config.executable);
if let Some(dir) = working_dir {
command_builder.current_dir(dir);
}
command_builder
.stdout(Stdio::piped())
.stderr(Stdio::piped())
crates/goose-server/src/openapi.rs
+9
View File
@@ -354,6 +354,8 @@ derive_utoipa!(Icon as IconSchema);
super::routes::config_management::get_pricing,
super::routes::agent::start_agent,
super::routes::agent::resume_agent,
super::routes::agent::restart_agent,
super::routes::agent::update_working_dir,
super::routes::agent::get_tools,
super::routes::agent::read_resource,
super::routes::agent::call_tool,
@@ -372,6 +374,7 @@ derive_utoipa!(Icon as IconSchema);
super::routes::session::import_session,
super::routes::session::update_session_user_recipe_values,
super::routes::session::edit_message,
super::routes::session::get_session_extensions,
super::routes::schedule::create_schedule,
super::routes::schedule::list_schedules,
super::routes::schedule::delete_schedule,
@@ -431,6 +434,7 @@ derive_utoipa!(Icon as IconSchema);
super::routes::session::EditType,
super::routes::session::EditMessageRequest,
super::routes::session::EditMessageResponse,
super::routes::session::SessionExtensionsResponse,
Message,
MessageContent,
MessageMetadata,
@@ -529,9 +533,14 @@ derive_utoipa!(Icon as IconSchema);
super::routes::agent::CallToolResponse,
super::routes::agent::StartAgentRequest,
super::routes::agent::ResumeAgentRequest,
super::routes::agent::RestartAgentRequest,
super::routes::agent::UpdateWorkingDirRequest,
super::routes::agent::UpdateFromSessionRequest,
super::routes::agent::AddExtensionRequest,
super::routes::agent::RemoveExtensionRequest,
super::routes::agent::ResumeAgentResponse,
super::routes::agent::RestartAgentResponse,
goose::agents::ExtensionLoadResult,
super::routes::setup::SetupResponse,
super::tunnel::TunnelInfo,
super::tunnel::TunnelState,
crates/goose-server/src/routes/agent.rs
+310 -88
View File
@@ -10,6 +10,7 @@ use axum::{
routing::{get, post},
Json, Router,
};
use goose::agents::ExtensionLoadResult;
use goose::config::PermissionManager;
use base64::Engine;
@@ -20,8 +21,9 @@ use goose::prompt_template::render_global_file;
use goose::providers::create;
use goose::recipe::Recipe;
use goose::recipe_deeplink;
use goose::session::extension_data::ExtensionState;
use goose::session::session_manager::SessionType;
use goose::session::{Session, SessionManager};
use goose::session::{EnabledExtensionsState, Session, SessionManager};
use goose::{
agents::{extension::ToolInfo, extension_manager::get_parameter_names},
config::permission::PermissionLevel,
@@ -34,7 +36,7 @@ use std::path::PathBuf;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use tracing::{error, warn};
use tracing::error;
#[derive(Deserialize, utoipa::ToSchema)]
pub struct UpdateFromSessionRequest {
@@ -63,6 +65,8 @@ pub struct StartAgentRequest {
recipe_id: Option<String>,
#[serde(default)]
recipe_deeplink: Option<String>,
#[serde(default)]
extension_overrides: Option<Vec<ExtensionConfig>>,
}
#[derive(Deserialize, utoipa::ToSchema)]
@@ -70,6 +74,17 @@ pub struct StopAgentRequest {
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct RestartAgentRequest {
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct UpdateWorkingDirRequest {
session_id: String,
working_dir: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct ResumeAgentRequest {
session_id: String,
@@ -122,6 +137,18 @@ pub struct CallToolResponse {
_meta: Option<Value>,
}
#[derive(Serialize, utoipa::ToSchema)]
pub struct ResumeAgentResponse {
pub session: Session,
#[serde(skip_serializing_if = "Option::is_none")]
pub extension_results: Option<Vec<ExtensionLoadResult>>,
}
#[derive(Serialize, utoipa::ToSchema)]
pub struct RestartAgentResponse {
pub extension_results: Vec<ExtensionLoadResult>,
}
#[utoipa::path(
post,
path = "/agent/start",
@@ -133,6 +160,7 @@ pub struct CallToolResponse {
(status = 500, description = "Internal server error", body = ErrorResponse)
)
)]
#[allow(clippy::too_many_lines)]
async fn start_agent(
State(state): State<Arc<AppState>>,
Json(payload): Json<StartAgentRequest>,
@@ -144,6 +172,7 @@ async fn start_agent(
recipe,
recipe_id,
recipe_deeplink,
extension_overrides,
} = payload;
let original_recipe = if let Some(deeplink) = recipe_deeplink {
@@ -191,6 +220,27 @@ async fn start_agent(
}
})?;
// Initialize session with extensions (either overrides from hub or global defaults)
let extensions_to_use =
extension_overrides.unwrap_or_else(goose::config::get_enabled_extensions);
let mut extension_data = session.extension_data.clone();
let extensions_state = EnabledExtensionsState::new(extensions_to_use);
if let Err(e) = extensions_state.to_extension_data(&mut extension_data) {
tracing::warn!("Failed to initialize session with extensions: {}", e);
} else {
SessionManager::update_session(&session.id)
.extension_data(extension_data.clone())
.apply()
.await
.map_err(|err| {
error!("Failed to save initial extension state: {}", err);
ErrorResponse {
message: format!("Failed to save initial extension state: {}", err),
status: StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
}
if let Some(recipe) = original_recipe {
SessionManager::update_session(&session.id)
.recipe(Some(recipe))
@@ -203,18 +253,50 @@ async fn start_agent(
status: StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
session = SessionManager::get_session(&session.id, false)
.await
.map_err(|err| {
error!("Failed to get updated session: {}", err);
ErrorResponse {
message: format!("Failed to get updated session: {}", err),
status: StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
}
// Refetch session to get all updates
session = SessionManager::get_session(&session.id, false)
.await
.map_err(|err| {
error!("Failed to get updated session: {}", err);
ErrorResponse {
message: format!("Failed to get updated session: {}", err),
status: StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
// Eagerly start loading extensions in the background
let session_for_spawn = session.clone();
let state_for_spawn = state.clone();
let session_id_for_task = session.id.clone();
let task = tokio::spawn(async move {
match state_for_spawn
.get_agent(session_for_spawn.id.clone())
.await
{
Ok(agent) => {
let results = agent.load_extensions_from_session(&session_for_spawn).await;
tracing::debug!(
"Background extension loading completed for session {}",
session_for_spawn.id
);
results
}
Err(e) => {
tracing::warn!(
"Failed to create agent for background extension loading: {}",
e
);
vec![]
}
}
});
state
.set_extension_loading_task(session_id_for_task, task)
.await;
Ok(Json(session))
}
@@ -223,7 +305,7 @@ async fn start_agent(
path = "/agent/resume",
request_body = ResumeAgentRequest,
responses(
(status = 200, description = "Agent started successfully", body = Session),
(status = 200, description = "Agent started successfully", body = ResumeAgentResponse),
(status = 400, description = "Bad request - invalid working directory"),
(status = 401, description = "Unauthorized - invalid secret key"),
(status = 500, description = "Internal server error")
@@ -232,7 +314,7 @@ async fn start_agent(
async fn resume_agent(
State(state): State<Arc<AppState>>,
Json(payload): Json<ResumeAgentRequest>,
) -> Result<Json<Session>, ErrorResponse> {
) -> Result<Json<ResumeAgentResponse>, ErrorResponse> {
goose::posthog::set_session_context("desktop", true);
let session = SessionManager::get_session(&payload.session_id, true)
@@ -246,7 +328,7 @@ async fn resume_agent(
}
})?;
if payload.load_model_and_extensions {
let extension_results = if payload.load_model_and_extensions {
let agent = state
.get_agent_for_route(payload.session_id.clone())
.await
@@ -255,81 +337,41 @@ async fn resume_agent(
status: code,
})?;
let config = Config::global();
agent
.restore_provider_from_session(&session)
.await
.map_err(|e| ErrorResponse {
message: e.to_string(),
status: StatusCode::INTERNAL_SERVER_ERROR,
})?;
let provider_result = async {
let provider_name = session
.provider_name
.clone()
.or_else(|| config.get_goose_provider().ok())
.ok_or_else(|| ErrorResponse {
message: "Could not configure agent: missing provider".into(),
status: StatusCode::INTERNAL_SERVER_ERROR,
})?;
let model_config = match session.model_config.clone() {
Some(saved_config) => saved_config,
None => {
let model_name = config.get_goose_model().map_err(|_| ErrorResponse {
message: "Could not configure agent: missing model".into(),
status: StatusCode::INTERNAL_SERVER_ERROR,
})?;
ModelConfig::new(&model_name).map_err(|e| ErrorResponse {
message: format!("Could not configure agent: invalid model {}", e),
status: StatusCode::INTERNAL_SERVER_ERROR,
})?
}
let extension_results =
if let Some(results) = state.take_extension_loading_task(&payload.session_id).await {
tracing::debug!(
"Using background extension loading results for session {}",
payload.session_id
);
state
.remove_extension_loading_task(&payload.session_id)
.await;
results
} else {
tracing::debug!(
"No background task found, loading extensions for session {}",
payload.session_id
);
agent.load_extensions_from_session(&session).await
};
let provider =
create(&provider_name, model_config)
.await
.map_err(|e| ErrorResponse {
message: format!("Could not create provider: {}", e),
status: StatusCode::INTERNAL_SERVER_ERROR,
})?;
Some(extension_results)
} else {
None
};
agent
.update_provider(provider, &payload.session_id)
.await
.map_err(|e| ErrorResponse {
message: format!("Could not configure agent: {}", e),
status: StatusCode::INTERNAL_SERVER_ERROR,
})
};
let extensions_result = async {
let enabled_configs = goose::config::get_enabled_extensions();
let agent_clone = agent.clone();
let extension_futures = enabled_configs
.into_iter()
.map(|config| {
let config_clone = config.clone();
let agent_ref = agent_clone.clone();
async move {
if let Err(e) = agent_ref.add_extension(config_clone.clone()).await {
warn!("Failed to load extension {}: {}", config_clone.name(), e);
goose::posthog::emit_error(
"extension_load_failed",
&format!("{}: {}", config_clone.name(), e),
);
}
Ok::<_, ErrorResponse>(())
}
})
.collect::<Vec<_>>();
futures::future::join_all(extension_futures).await;
Ok::<(), ErrorResponse>(()) // Fixed type annotation
};
let (provider_result, _) = tokio::join!(provider_result, extensions_result);
provider_result?;
}
Ok(Json(session))
Ok(Json(ResumeAgentResponse {
session,
extension_results,
}))
}
#[utoipa::path(
@@ -519,7 +561,8 @@ async fn agent_add_extension(
Json(request): Json<AddExtensionRequest>,
) -> Result<StatusCode, ErrorResponse> {
let extension_name = request.config.name();
let agent = state.get_agent(request.session_id).await?;
let agent = state.get_agent(request.session_id.clone()).await?;
agent.add_extension(request.config).await.map_err(|e| {
goose::posthog::emit_error(
"extension_add_failed",
@@ -527,6 +570,18 @@ async fn agent_add_extension(
);
ErrorResponse::internal(format!("Failed to add extension: {}", e))
})?;
// Persist here rather than in add_extension to ensure we only save state
// after the extension successfully loads. This prevents failed extensions
// from being persisted as enabled in the session.
agent
.persist_extension_state(&request.session_id)
.await
.map_err(|e| {
error!("Failed to persist extension state: {}", e);
ErrorResponse::internal(format!("Failed to persist extension state: {}", e))
})?;
Ok(StatusCode::OK)
}
@@ -545,8 +600,20 @@ async fn agent_remove_extension(
State(state): State<Arc<AppState>>,
Json(request): Json<RemoveExtensionRequest>,
) -> Result<StatusCode, ErrorResponse> {
let agent = state.get_agent(request.session_id).await?;
let agent = state.get_agent(request.session_id.clone()).await?;
agent.remove_extension(&request.name).await?;
agent
.persist_extension_state(&request.session_id)
.await
.map_err(|e| {
error!("Failed to persist extension state: {}", e);
ErrorResponse {
message: format!("Failed to persist extension state: {}", e),
status: StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
Ok(StatusCode::OK)
}
@@ -578,6 +645,159 @@ async fn stop_agent(
Ok(StatusCode::OK)
}
async fn restart_agent_internal(
state: &Arc<AppState>,
session_id: &str,
session: &Session,
) -> Result<Vec<ExtensionLoadResult>, ErrorResponse> {
// Remove existing agent (ignore error if not found)
let _ = state.agent_manager.remove_session(session_id).await;
let agent = state
.get_agent_for_route(session_id.to_string())
.await
.map_err(|code| ErrorResponse {
message: "Failed to create new agent during restart".into(),
status: code,
})?;
let provider_future = agent.restore_provider_from_session(session);
let extensions_future = agent.load_extensions_from_session(session);
let (provider_result, extension_results) = tokio::join!(provider_future, extensions_future);
provider_result.map_err(|e| ErrorResponse {
message: e.to_string(),
status: StatusCode::INTERNAL_SERVER_ERROR,
})?;
let context: HashMap<&str, Value> = HashMap::new();
let desktop_prompt =
render_global_file("desktop_prompt.md", &context).expect("Prompt should render");
let mut update_prompt = desktop_prompt;
if let Some(ref recipe) = session.recipe {
match build_recipe_with_parameter_values(
recipe,
session.user_recipe_values.clone().unwrap_or_default(),
)
.await
{
Ok(Some(recipe)) => {
if let Some(prompt) = apply_recipe_to_agent(&agent, &recipe, true).await {
update_prompt = prompt;
}
}
Ok(None) => {
// Recipe has missing parameters - use default prompt
}
Err(e) => {
return Err(ErrorResponse {
message: e.to_string(),
status: StatusCode::INTERNAL_SERVER_ERROR,
});
}
}
}
agent.extend_system_prompt(update_prompt).await;
Ok(extension_results)
}
#[utoipa::path(
post,
path = "/agent/restart",
request_body = RestartAgentRequest,
responses(
(status = 200, description = "Agent restarted successfully", body = RestartAgentResponse),
(status = 401, description = "Unauthorized - invalid secret key"),
(status = 404, description = "Session not found"),
(status = 500, description = "Internal server error")
)
)]
async fn restart_agent(
State(state): State<Arc<AppState>>,
Json(payload): Json<RestartAgentRequest>,
) -> Result<Json<RestartAgentResponse>, ErrorResponse> {
let session_id = payload.session_id.clone();
let session = SessionManager::get_session(&session_id, false)
.await
.map_err(|err| {
error!("Failed to get session during restart: {}", err);
ErrorResponse {
message: format!("Failed to get session: {}", err),
status: StatusCode::NOT_FOUND,
}
})?;
let extension_results = restart_agent_internal(&state, &session_id, &session).await?;
Ok(Json(RestartAgentResponse { extension_results }))
}
#[utoipa::path(
post,
path = "/agent/update_working_dir",
request_body = UpdateWorkingDirRequest,
responses(
(status = 200, description = "Working directory updated and agent restarted successfully"),
(status = 400, description = "Bad request - invalid directory path"),
(status = 401, description = "Unauthorized - invalid secret key"),
(status = 404, description = "Session not found"),
(status = 500, description = "Internal server error")
)
)]
async fn update_working_dir(
State(state): State<Arc<AppState>>,
Json(payload): Json<UpdateWorkingDirRequest>,
) -> Result<StatusCode, ErrorResponse> {
let session_id = payload.session_id.clone();
let working_dir = payload.working_dir.trim();
if working_dir.is_empty() {
return Err(ErrorResponse {
message: "Working directory cannot be empty".into(),
status: StatusCode::BAD_REQUEST,
});
}
let path = PathBuf::from(working_dir);
if !path.exists() || !path.is_dir() {
return Err(ErrorResponse {
message: "Invalid directory path".into(),
status: StatusCode::BAD_REQUEST,
});
}
// Update the session's working directory
SessionManager::update_session(&session_id)
.working_dir(path)
.apply()
.await
.map_err(|e| {
error!("Failed to update session working directory: {}", e);
ErrorResponse {
message: format!("Failed to update working directory: {}", e),
status: StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
// Get the updated session and restart the agent
let session = SessionManager::get_session(&session_id, false)
.await
.map_err(|err| {
error!("Failed to get session after working dir update: {}", err);
ErrorResponse {
message: format!("Failed to get session: {}", err),
status: StatusCode::NOT_FOUND,
}
})?;
restart_agent_internal(&state, &session_id, &session).await?;
Ok(StatusCode::OK)
}
#[utoipa::path(
post,
path = "/agent/read_resource",
@@ -702,6 +922,8 @@ pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/agent/start", post(start_agent))
.route("/agent/resume", post(resume_agent))
.route("/agent/restart", post(restart_agent))
.route("/agent/update_working_dir", post(update_working_dir))
.route("/agent/tools", get(get_tools))
.route("/agent/read_resource", post(read_resource))
.route("/agent/call_tool", post(call_tool))
crates/goose-server/src/routes/session.rs
+45 -1
View File
@@ -9,9 +9,11 @@ use axum::{
routing::{delete, get, put},
Json, Router,
};
use goose::agents::ExtensionConfig;
use goose::recipe::Recipe;
use goose::session::extension_data::ExtensionState;
use goose::session::session_manager::SessionInsights;
use goose::session::{Session, SessionManager};
use goose::session::{EnabledExtensionsState, Session, SessionManager};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
@@ -393,6 +395,44 @@ async fn edit_message(
}
}
#[derive(Serialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct SessionExtensionsResponse {
extensions: Vec<ExtensionConfig>,
}
#[utoipa::path(
get,
path = "/sessions/{session_id}/extensions",
params(
("session_id" = String, Path, description = "Unique identifier for the session")
),
responses(
(status = 200, description = "Session extensions retrieved successfully", body = SessionExtensionsResponse),
(status = 401, description = "Unauthorized - Invalid or missing API key"),
(status = 404, description = "Session not found"),
(status = 500, description = "Internal server error")
),
security(
("api_key" = [])
),
tag = "Session Management"
)]
async fn get_session_extensions(
Path(session_id): Path<String>,
) -> Result<Json<SessionExtensionsResponse>, StatusCode> {
let session = SessionManager::get_session(&session_id, false)
.await
.map_err(|_| StatusCode::NOT_FOUND)?;
// Try to get session-specific extensions, fall back to global config
let extensions = EnabledExtensionsState::from_extension_data(&session.extension_data)
.map(|state| state.extensions)
.unwrap_or_else(goose::config::get_enabled_extensions);
Ok(Json(SessionExtensionsResponse { extensions }))
}
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/sessions", get(list_sessions))
@@ -407,5 +447,9 @@ pub fn routes(state: Arc<AppState>) -> Router {
put(update_session_user_recipe_values),
)
.route("/sessions/{session_id}/edit_message", post(edit_message))
.route(
"/sessions/{session_id}/extensions",
get(get_session_extensions),
)
.with_state(state)
}
crates/goose-server/src/state.rs
+44
View File
@@ -6,8 +6,13 @@ use std::path::PathBuf;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use crate::tunnel::TunnelManager;
use goose::agents::ExtensionLoadResult;
type ExtensionLoadingTasks =
Arc<Mutex<HashMap<String, Arc<Mutex<Option<JoinHandle<Vec<ExtensionLoadResult>>>>>>>>;
#[derive(Clone)]
pub struct AppState {
@@ -17,6 +22,7 @@ pub struct AppState {
/// Tracks sessions that have already emitted recipe telemetry to prevent double counting.
recipe_session_tracker: Arc<Mutex<HashSet<String>>>,
pub tunnel_manager: Arc<TunnelManager>,
pub extension_loading_tasks: ExtensionLoadingTasks,
}
impl AppState {
@@ -30,9 +36,47 @@ impl AppState {
session_counter: Arc::new(AtomicUsize::new(0)),
recipe_session_tracker: Arc::new(Mutex::new(HashSet::new())),
tunnel_manager,
extension_loading_tasks: Arc::new(Mutex::new(HashMap::new())),
}))
}
pub async fn set_extension_loading_task(
&self,
session_id: String,
task: JoinHandle<Vec<ExtensionLoadResult>>,
) {
let mut tasks = self.extension_loading_tasks.lock().await;
tasks.insert(session_id, Arc::new(Mutex::new(Some(task))));
}
pub async fn take_extension_loading_task(
&self,
session_id: &str,
) -> Option<Vec<ExtensionLoadResult>> {
let task_holder = {
let tasks = self.extension_loading_tasks.lock().await;
tasks.get(session_id).cloned()
};
if let Some(holder) = task_holder {
let task = holder.lock().await.take();
if let Some(handle) = task {
match handle.await {
Ok(results) => return Some(results),
Err(e) => {
tracing::warn!("Background extension loading task failed: {}", e);
}
}
}
}
None
}
pub async fn remove_extension_loading_task(&self, session_id: &str) {
let mut tasks = self.extension_loading_tasks.lock().await;
tasks.remove(session_id);
}
pub fn scheduler(&self) -> Arc<dyn SchedulerTrait> {
self.agent_manager.scheduler()
}
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);