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