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

feat: persist GooseMode per-session via session DB (#7854)

Browse Source 
Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-03-16 19:37:21 +08:00
committed by GitHub
parent 2631095f20
commit 94fdcdd07a
42 changed files with 953 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+48 -8
View File
@@ -137,6 +137,7 @@ impl AgentConfig {
pub struct Agent {
pub(super) provider: SharedProvider,
pub config: AgentConfig,
pub(super) current_goose_mode: Mutex<GooseMode>,
pub extension_manager: Arc<ExtensionManager>,
pub(super) final_output_tool: Arc<Mutex<Option<FinalOutputTool>>>,
@@ -203,14 +204,13 @@ where
impl Agent {
pub fn new() -> Self {
let config = Config::global();
Self::with_config(AgentConfig::new(
Arc::new(SessionManager::instance()),
PermissionManager::instance(),
None,
Config::global().get_goose_mode().unwrap_or(GooseMode::Auto),
Config::global()
.get_goose_disable_session_naming()
.unwrap_or(false),
config.get_goose_mode().unwrap_or_default(),
config.get_goose_disable_session_naming().unwrap_or(false),
GoosePlatform::GooseCli,
))
}
@@ -220,6 +220,7 @@ impl Agent {
let provider = Arc::new(Mutex::new(None));
let goose_platform = config.goose_platform.clone();
let initial_mode = config.goose_mode;
let capabilities = match config.goose_platform {
GoosePlatform::GooseDesktop => ExtensionManagerCapabilities { mcpui: true },
GoosePlatform::GooseCli => ExtensionManagerCapabilities { mcpui: false },
@@ -229,6 +230,7 @@ impl Agent {
Self {
provider: provider.clone(),
config,
current_goose_mode: Mutex::new(initial_mode),
extension_manager: Arc::new(ExtensionManager::new(
provider.clone(),
session_manager,
@@ -351,7 +353,9 @@ impl Agent {
.prepare_tools_and_prompt(session_id, working_dir)
.await?;
if self.config.goose_mode == GooseMode::SmartApprove {
let goose_mode = *self.current_goose_mode.lock().await;
if goose_mode == GooseMode::SmartApprove {
self.tool_inspection_manager.apply_tool_annotations(&tools);
}
@@ -360,7 +364,7 @@ impl Agent {
tools,
toolshim_tools,
system_prompt,
goose_mode: self.config.goose_mode,
goose_mode,
tool_call_cut_off: Config::global()
.get_param::<usize>("GOOSE_TOOL_CALL_CUTOFF")
.unwrap_or(10),
@@ -1721,6 +1725,28 @@ impl Agent {
.context("Failed to persist provider config to session")
}
pub async fn update_goose_mode(&self, mode: GooseMode, session_id: &str) -> Result<()> {
if let Some(provider) = self.provider.lock().await.as_ref() {
provider
.update_mode(session_id, mode)
.await
.map_err(|e| anyhow::anyhow!("Provider rejected mode update: {e}"))?;
}
*self.current_goose_mode.lock().await = mode;
self.config
.session_manager
.clone()
.update(session_id)
.goose_mode(mode)
.apply()
.await
.context("Failed to persist goose_mode to session")
}
pub async fn goose_mode(&self) -> GooseMode {
*self.current_goose_mode.lock().await
}
/// 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<()> {
@@ -1752,7 +1778,16 @@ impl Agent {
.await
.map_err(|e| anyhow!("Could not create provider: {}", e))?;
self.update_provider(provider, &session.id).await
self.update_provider(provider, &session.id).await?;
// Propagate session mode to the new provider
if let Some(provider) = self.provider.lock().await.as_ref() {
provider
.update_mode(&session.id, session.goose_mode)
.await
.map_err(|e| anyhow!("Failed to propagate mode to provider: {}", e))?;
}
*self.current_goose_mode.lock().await = session.goose_mode;
Ok(())
}
/// Override the system prompt with a custom template
@@ -1864,12 +1899,14 @@ impl Agent {
let model_name = &model_config.model_name;
tracing::debug!("Using model: {}", model_name);
let goose_mode = *self.current_goose_mode.lock().await;
let prompt_manager = self.prompt_manager.lock().await;
let system_prompt = prompt_manager
.builder()
.with_extensions(extensions_info.into_iter())
.with_frontend_instructions(self.frontend_instructions.lock().await.clone())
.with_extension_and_tool_counts(extension_count, tool_count)
.with_goose_mode(goose_mode)
.build();
let recipe_prompt = prompt_manager.get_recipe_prompt().await;
@@ -2220,7 +2257,10 @@ mod tests {
);
let prompt_manager = agent.prompt_manager.lock().await;
let system_prompt = prompt_manager.builder().build();
let system_prompt = prompt_manager
.builder()
.with_goose_mode(GooseMode::default())
.build();
let final_output_tool_ref = agent.final_output_tool.lock().await;
let final_output_tool_system_prompt =
crates/goose/src/agents/platform_extensions/summon.rs
+16 -4
View File
@@ -12,7 +12,7 @@ use crate::agents::subagent_task_config::{TaskConfig, DEFAULT_SUBAGENT_MAX_TURNS
use crate::agents::tool_execution::ToolCallContext;
use crate::agents::AgentConfig;
use crate::config::paths::Paths;
use crate::config::Config;
use crate::config::{Config, GooseMode};
use crate::providers;
use crate::recipe::build_recipe::build_recipe_from_template;
use crate::recipe::local_recipes::load_local_recipe_file;
@@ -1243,11 +1243,14 @@ impl SummonClient {
.await
.map_err(|e| format!("Failed to build task config: {}", e))?;
// Subagents must use Auto until get_agent_messages forwards
// ActionRequired messages to the parent. Until then, any mode
// that requires approval will hang on the subagent's confirmation_rx.
let agent_config = AgentConfig::new(
self.context.session_manager.clone(),
crate::config::permission::PermissionManager::instance(),
None,
crate::config::GooseMode::Auto,
GooseMode::Auto,
true, // disable session naming for subagents
crate::agents::GoosePlatform::GooseCli,
);
@@ -1259,6 +1262,7 @@ impl SummonClient {
working_dir,
"Delegated task".to_string(),
SessionType::SubAgent,
GooseMode::Auto,
)
.await
.map_err(|e| format!("Failed to create subagent session: {}", e))?;
@@ -1704,11 +1708,14 @@ impl SummonClient {
let description = truncate(&Self::get_task_description(&params), 40);
// Subagents must use Auto until get_agent_messages forwards
// ActionRequired messages to the parent. Until then, any mode
// that requires approval will hang on the subagent's confirmation_rx.
let agent_config = AgentConfig::new(
self.context.session_manager.clone(),
crate::config::permission::PermissionManager::instance(),
None,
crate::config::GooseMode::Auto,
GooseMode::Auto,
true, // disable session naming for subagents
crate::agents::GoosePlatform::GooseCli,
);
@@ -1716,7 +1723,12 @@ impl SummonClient {
let subagent_session = self
.context
.session_manager
.create_session(working_dir, description.clone(), SessionType::SubAgent)
.create_session(
working_dir,
description.clone(),
SessionType::SubAgent,
GooseMode::Auto,
)
.await
.map_err(|e| format!("Failed to create subagent session: {}", e))?;
crates/goose/src/agents/prompt_manager.rs
+12 -2
View File
@@ -55,6 +55,7 @@ pub struct SystemPromptBuilder<'a, M> {
subagents_enabled: bool,
hints: Option<String>,
code_execution_mode: bool,
goose_mode: Option<GooseMode>,
}
impl<'a> SystemPromptBuilder<'a, PromptManager> {
@@ -106,6 +107,11 @@ impl<'a> SystemPromptBuilder<'a, PromptManager> {
self
}
pub fn with_goose_mode(mut self, mode: GooseMode) -> Self {
self.goose_mode = Some(mode);
self
}
pub fn build(self) -> String {
let mut extensions_info = self.extensions_info;
@@ -128,8 +134,9 @@ impl<'a> SystemPromptBuilder<'a, PromptManager> {
})
.collect();
let config = Config::global();
let goose_mode = config.get_goose_mode().unwrap_or(GooseMode::Auto);
let goose_mode = self
.goose_mode
.unwrap_or_else(|| Config::global().get_goose_mode().unwrap_or_default());
let extension_tool_limits = self
.extension_tool_count
@@ -250,6 +257,7 @@ impl PromptManager {
subagents_enabled: false,
hints: None,
code_execution_mode: false,
goose_mode: None,
}
}
@@ -403,6 +411,7 @@ mod tests {
#[tokio::test]
async fn test_all_platform_extensions() {
use crate::agents::platform_extensions::{PlatformExtensionContext, PLATFORM_EXTENSIONS};
use crate::config::GooseMode;
use crate::session::SessionManager;
use std::sync::Arc;
@@ -413,6 +422,7 @@ mod tests {
tmp_dir.path().to_path_buf(),
"test session".to_owned(),
crate::session::SessionType::Hidden,
GooseMode::default(),
)
.await
.unwrap();
crates/goose/src/agents/reply_parts.rs
+5
View File
@@ -181,6 +181,8 @@ impl Agent {
let provider = self.provider().await?;
let model_config = provider.get_model_config();
let goose_mode = *self.current_goose_mode.lock().await;
let prompt_manager = self.prompt_manager.lock().await;
let mut system_prompt = prompt_manager
.builder()
@@ -189,6 +191,7 @@ impl Agent {
.with_extension_and_tool_counts(extension_count, tool_count)
.with_code_execution_mode(code_execution_active)
.with_hints(working_dir)
.with_goose_mode(goose_mode)
.build();
// Handle toolshim if enabled
@@ -435,6 +438,7 @@ impl Agent {
#[cfg(test)]
mod tests {
use super::*;
use crate::config::GooseMode;
use crate::conversation::message::Message;
use crate::model::ModelConfig;
use crate::providers::base::{Provider, ProviderUsage, Usage};
@@ -483,6 +487,7 @@ mod tests {
std::env::current_dir().unwrap(),
"test-prepare-tools".to_string(),
SessionType::Hidden,
GooseMode::default(),
)
.await?;