tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 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/examples/agent.rs
+2 -1
View File
@@ -1,7 +1,7 @@
use dotenvy::dotenv;
use futures::StreamExt;
use goose::agents::{Agent, AgentEvent, ExtensionConfig, SessionConfig};
use goose::config::{DEFAULT_EXTENSION_DESCRIPTION, DEFAULT_EXTENSION_TIMEOUT};
use goose::config::{GooseMode, DEFAULT_EXTENSION_DESCRIPTION, DEFAULT_EXTENSION_TIMEOUT};
use goose::conversation::message::Message;
use goose::providers::create_with_named_model;
use goose::providers::databricks::DATABRICKS_DEFAULT_MODEL;
@@ -24,6 +24,7 @@ async fn main() -> anyhow::Result<()> {
PathBuf::default(),
"max-turn-test".to_string(),
SessionType::Hidden,
GooseMode::default(),
)
.await?;
crates/goose/src/acp/provider.rs
+23 -18
View File
@@ -286,28 +286,33 @@ impl Provider for AcpProvider {
}
async fn update_mode(&self, session_id: &str, mode: GooseMode) -> Result<(), ProviderError> {
let _acp_session_id = self
.goose_to_acp_id
.lock()
let map = self.goose_to_acp_id.lock().await;
if map.is_empty() {
// Pre-initialization: no ACP session yet, just store the mode.
// The shared Arc<Mutex<GooseMode>> is read at session creation time.
drop(map);
} else if let Some(acp_session_id) = map.get(session_id).map(|r| r.session_id.clone()) {
drop(map);
self.send_untyped(
"session/set_mode",
serde_json::json!({
"sessionId": acp_session_id,
"modeId": mode.to_string().to_lowercase()
}),
)
.await
.get(session_id)
.map(|r| r.session_id.clone())
.ok_or_else(|| {
ProviderError::RequestFailed(format!("Session not found: {session_id}"))
})?;
let current = self
.goose_mode
.lock()
.map_err(|_| ProviderError::RequestFailed("Failed to read mode".into()))?;
if mode != *current {
// TODO: "session/set_mode" when session-scoped mode lands (#7603)
.map_err(|e| ProviderError::RequestFailed(format!("Failed to set mode: {e}")))?;
} else {
return Err(ProviderError::RequestFailed(format!(
"Mode change not supported: session is {}, requested {}",
current, mode
"Session not found: {session_id}"
)));
}
let mut current = self
.goose_mode
.lock()
.map_err(|_| ProviderError::RequestFailed("Failed to update mode".into()))?;
*current = mode;
Ok(())
}
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?;
crates/goose/src/execution/manager.rs
+90 -7
View File
@@ -21,6 +21,7 @@ pub struct AgentManager {
scheduler: Arc<dyn SchedulerTrait>,
session_manager: Arc<SessionManager>,
default_provider: Arc<RwLock<Option<Arc<dyn crate::providers::base::Provider>>>>,
default_mode: GooseMode,
}
impl AgentManager {
@@ -28,6 +29,7 @@ impl AgentManager {
session_manager: Arc<SessionManager>,
schedule_file_path: std::path::PathBuf,
max_sessions: Option<usize>,
default_mode: GooseMode,
) -> Result<Self> {
let scheduler = Scheduler::new(schedule_file_path, session_manager.clone()).await?;
@@ -39,6 +41,7 @@ impl AgentManager {
scheduler,
session_manager,
default_provider: Arc::new(RwLock::new(None)),
default_mode,
};
Ok(manager)
@@ -47,13 +50,20 @@ impl AgentManager {
pub async fn instance() -> Result<Arc<Self>> {
AGENT_MANAGER
.get_or_try_init(|| async {
let max_sessions = Config::global()
let config = Config::global();
let max_sessions = config
.get_goose_max_active_agents()
.unwrap_or(DEFAULT_MAX_SESSION);
let default_mode = config.get_goose_mode().unwrap_or_default();
let schedule_file_path = Paths::data_dir().join("schedule.json");
let session_manager = Arc::new(SessionManager::instance());
let manager =
Self::new(session_manager, schedule_file_path, Some(max_sessions)).await?;
let manager = Self::new(
session_manager,
schedule_file_path,
Some(max_sessions),
default_mode,
)
.await?;
Ok(Arc::new(manager))
})
.await
@@ -82,8 +92,14 @@ impl AgentManager {
}
}
let mode = Config::global().get_goose_mode().unwrap_or(GooseMode::Auto);
let mut mode = self.default_mode;
let permission_manager = PermissionManager::instance();
if let Ok(session) = self.session_manager.get_session(&session_id, false).await {
mode = session.goose_mode;
info!(goose_mode = %mode, session_id = %session_id, "Session loaded");
}
let config = AgentConfig::new(
Arc::clone(&self.session_manager),
permission_manager,
@@ -118,6 +134,10 @@ impl AgentManager {
agent
.update_provider(Arc::clone(provider), &session_id)
.await?;
provider
.update_mode(&session_id, mode)
.await
.map_err(|e| anyhow::anyhow!("Failed to propagate mode to provider: {}", e))?;
}
}
@@ -153,6 +173,9 @@ mod tests {
use std::sync::Arc;
use tempfile::TempDir;
use test_case::test_case;
use crate::config::GooseMode;
use crate::execution::SessionExecutionMode;
use crate::session::SessionManager;
@@ -161,9 +184,14 @@ mod tests {
async fn create_test_manager(temp_dir: &TempDir) -> AgentManager {
let session_manager = Arc::new(SessionManager::new(temp_dir.path().to_path_buf()));
let schedule_path = temp_dir.path().join("schedule.json");
AgentManager::new(session_manager, schedule_path, Some(100))
.await
.unwrap()
AgentManager::new(
session_manager,
schedule_path,
Some(100),
GooseMode::default(),
)
.await
.unwrap()
}
#[test]
@@ -369,4 +397,59 @@ mod tests {
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("not found"));
}
#[test_case(GooseMode::Approve ; "approve")]
#[test_case(GooseMode::Chat ; "chat")]
#[test_case(GooseMode::SmartApprove ; "smart_approve")]
#[tokio::test]
async fn test_agent_inherits_session_mode(mode: GooseMode) {
let temp_dir = TempDir::new().unwrap();
let manager = create_test_manager(&temp_dir).await;
let session = manager
.session_manager()
.create_session(
temp_dir.path().to_path_buf(),
"test".into(),
crate::session::SessionType::User,
mode,
)
.await
.unwrap();
let agent = manager.get_or_create_agent(session.id).await.unwrap();
assert_eq!(agent.goose_mode().await, mode);
}
#[tokio::test]
async fn test_session_mode_isolation() {
let temp_dir = TempDir::new().unwrap();
let manager = create_test_manager(&temp_dir).await;
let sm = manager.session_manager();
let s1 = sm
.create_session(
temp_dir.path().to_path_buf(),
"s1".into(),
crate::session::SessionType::User,
GooseMode::Approve,
)
.await
.unwrap();
let s2 = sm
.create_session(
temp_dir.path().to_path_buf(),
"s2".into(),
crate::session::SessionType::User,
GooseMode::Auto,
)
.await
.unwrap();
let a1 = manager.get_or_create_agent(s1.id).await.unwrap();
let a2 = manager.get_or_create_agent(s2.id).await.unwrap();
assert_eq!(a1.goose_mode().await, GooseMode::Approve);
assert_eq!(a2.goose_mode().await, GooseMode::Auto);
}
}
crates/goose/src/gateway/handler.rs
+15 -3
View File
@@ -128,14 +128,19 @@ impl GatewayHandler {
user.display_name.as_deref().unwrap_or(&user.user_id)
);
let config = Config::global();
let session = self
.agent_manager
.session_manager()
.create_session(working_dir, session_name, SessionType::Gateway)
.create_session(
working_dir,
session_name,
SessionType::Gateway,
config.get_goose_mode().unwrap_or_default(),
)
.await?;
let manager = self.agent_manager.session_manager();
let config = Config::global();
// Store the current provider and model config on the session so the agent
// can be restored after LRU eviction, matching the start_agent flow.
@@ -197,6 +202,7 @@ impl GatewayHandler {
let current_provider = config.get_goose_provider().ok();
let current_model_name = config.get_goose_model().ok();
let current_extensions = get_enabled_extensions();
let current_mode = config.get_goose_mode().unwrap_or_default();
// --- what the session has ---
let session_extensions: Vec<ExtensionConfig> =
@@ -208,8 +214,9 @@ impl GatewayHandler {
let model_changed = current_model_name.as_deref()
!= session.model_config.as_ref().map(|m| m.model_name.as_str());
let extensions_changed = current_extensions != session_extensions;
let mode_changed = current_mode != session.goose_mode;
if !provider_changed && !model_changed && !extensions_changed {
if !provider_changed && !model_changed && !extensions_changed && !mode_changed {
return Ok(false);
}
@@ -218,6 +225,7 @@ impl GatewayHandler {
provider_changed,
model_changed,
extensions_changed,
mode_changed,
"syncing gateway session with current config"
);
@@ -242,6 +250,10 @@ impl GatewayHandler {
}
}
if mode_changed {
update = update.goose_mode(current_mode);
}
update.apply().await?;
Ok(extensions_changed)
}
crates/goose/src/model.rs
+16
View File
@@ -429,6 +429,10 @@ mod tests {
#[test]
fn sets_limits_from_canonical_model() {
let _guard = env_lock::lock_env([
("GOOSE_MAX_TOKENS", None::<&str>),
("GOOSE_CONTEXT_LIMIT", None::<&str>),
]);
let config = ModelConfig::new_or_fail("gpt-4o").with_canonical_limits("openai");
assert_eq!(config.context_limit, Some(128_000));
@@ -438,6 +442,10 @@ mod tests {
#[test]
fn does_not_override_existing_context_limit() {
let _guard = env_lock::lock_env([
("GOOSE_MAX_TOKENS", None::<&str>),
("GOOSE_CONTEXT_LIMIT", None::<&str>),
]);
let mut config = ModelConfig::new_or_fail("gpt-4o");
config.context_limit = Some(64_000);
let config = config.with_canonical_limits("openai");
@@ -447,6 +455,10 @@ mod tests {
#[test]
fn does_not_override_existing_max_tokens() {
let _guard = env_lock::lock_env([
("GOOSE_MAX_TOKENS", None::<&str>),
("GOOSE_CONTEXT_LIMIT", None::<&str>),
]);
let mut config = ModelConfig::new_or_fail("gpt-4o");
config.max_tokens = Some(1_000);
let config = config.with_canonical_limits("openai");
@@ -456,6 +468,10 @@ mod tests {
#[test]
fn unknown_model_leaves_fields_none() {
let _guard = env_lock::lock_env([
("GOOSE_MAX_TOKENS", None::<&str>),
("GOOSE_CONTEXT_LIMIT", None::<&str>),
]);
let config =
ModelConfig::new_or_fail("totally-unknown-model").with_canonical_limits("openai");
crates/goose/src/providers/base.rs
+5 -6
View File
@@ -8,8 +8,7 @@ use super::canonical::{map_to_canonical_model, CanonicalModelRegistry};
use super::errors::ProviderError;
use super::retry::RetryConfig;
use crate::config::base::ConfigValue;
use crate::config::goose_mode::GooseMode;
use crate::config::ExtensionConfig;
use crate::config::{ExtensionConfig, GooseMode};
use crate::conversation::message::{Message, MessageContent};
use crate::conversation::Conversation;
use crate::model::ModelConfig;
@@ -708,6 +707,10 @@ pub trait Provider: Send + Sync {
))
}
async fn update_mode(&self, _session_id: &str, _mode: GooseMode) -> Result<(), ProviderError> {
Ok(())
}
fn permission_routing(&self) -> PermissionRouting {
PermissionRouting::Noop
}
@@ -719,10 +722,6 @@ pub trait Provider: Send + Sync {
) -> bool {
false
}
async fn update_mode(&self, _session_id: &str, _mode: GooseMode) -> Result<(), ProviderError> {
Ok(())
}
}
/// A message stream yields partial text content but complete tool calls, all within the Message object
crates/goose/src/providers/claude_code.rs
+17
View File
@@ -270,6 +270,8 @@ pub struct ClaudeCodeProvider {
#[serde(skip)]
pending_confirmations:
Arc<tokio::sync::Mutex<HashMap<String, oneshot::Sender<PermissionConfirmation>>>>,
#[serde(skip)]
initial_mode: tokio::sync::Mutex<Option<GooseMode>>,
}
impl ClaudeCodeProvider {
@@ -623,6 +625,7 @@ impl ProviderDef for ClaudeCodeProvider {
mcp_config_file,
cli_process: tokio::sync::OnceCell::new(),
pending_confirmations: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
initial_mode: tokio::sync::Mutex::new(None),
})
})
}
@@ -669,6 +672,19 @@ impl Provider for ClaudeCodeProvider {
Ok(extract_model_aliases(response.ok().flatten().as_ref()))
}
async fn update_mode(&self, _session_id: &str, mode: GooseMode) -> Result<(), ProviderError> {
// Mode is baked into the subprocess at spawn; claude-acp replaces
// this provider (#7801).
let mut guard = self.initial_mode.lock().await;
let current = *guard.get_or_insert(mode);
if current != mode {
return Err(ProviderError::RequestFailed(format!(
"Mode change not supported: session is {current}, requested {mode}",
)));
}
Ok(())
}
fn permission_routing(&self) -> PermissionRouting {
PermissionRouting::ActionRequired
}
@@ -1192,6 +1208,7 @@ mod tests {
mcp_config_file: None,
cli_process: tokio::sync::OnceCell::new(),
pending_confirmations: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
initial_mode: tokio::sync::Mutex::new(None),
}
}
crates/goose/src/providers/codex.rs
+50 -9
View File
@@ -3,6 +3,7 @@ use async_trait::async_trait;
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
use futures::future::BoxFuture;
use serde_json::json;
use std::collections::HashMap;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::Stdio;
@@ -54,6 +55,8 @@ pub struct CodexProvider {
skip_git_check: bool,
/// CLI config overrides for MCP servers
mcp_config_overrides: Vec<String>,
#[serde(skip)]
mode_by_session: tokio::sync::RwLock<HashMap<String, GooseMode>>,
}
impl CodexProvider {
@@ -69,11 +72,11 @@ impl CodexProvider {
true
}
/// Apply permission flags based on GOOSE_MODE setting
fn apply_permission_flags(cmd: &mut Command) -> Result<(), ProviderError> {
let config = Config::global();
let goose_mode = config.get_goose_mode().unwrap_or(GooseMode::Auto);
/// Apply permission flags based on GooseMode
fn apply_permission_flags(
cmd: &mut Command,
goose_mode: GooseMode,
) -> Result<(), ProviderError> {
match goose_mode {
GooseMode::Auto => {
// --yolo is shorthand for --dangerously-bypass-approvals-and-sandbox
@@ -101,6 +104,7 @@ impl CodexProvider {
system: &str,
messages: &[Message],
_tools: &[Tool],
goose_mode: GooseMode,
) -> Result<Vec<String>, ProviderError> {
// Single pass: text → prompt (stdin), images → temp files (-i flags)
let image_dir = Paths::state_dir().join("codex/images");
@@ -151,8 +155,7 @@ impl CodexProvider {
// JSON output format for structured parsing
cmd.arg("--json");
// Apply permission mode based on GOOSE_MODE
Self::apply_permission_flags(&mut cmd)?;
Self::apply_permission_flags(&mut cmd, goose_mode)?;
// Skip git repo check if configured
if self.skip_git_check {
@@ -653,6 +656,7 @@ impl ProviderDef for CodexProvider {
reasoning_effort,
skip_git_check,
mcp_config_overrides: codex_mcp_config_overrides(&resolved),
mode_by_session: tokio::sync::RwLock::new(HashMap::new()),
})
})
}
@@ -671,7 +675,7 @@ impl Provider for CodexProvider {
async fn stream(
&self,
model_config: &ModelConfig,
_session_id: &str, // CLI has no external session-id flag to propagate.
session_id: &str,
system: &str,
messages: &[Message],
tools: &[Tool],
@@ -687,7 +691,13 @@ impl Provider for CodexProvider {
));
}
let lines = self.execute_command(system, messages, tools).await?;
let goose_mode = {
let map = self.mode_by_session.read().await;
map.get(session_id).copied().unwrap_or_default()
};
let lines = self
.execute_command(system, messages, tools, goose_mode)
.await?;
let (message, usage) = self.parse_response(&lines)?;
@@ -720,6 +730,14 @@ impl Provider for CodexProvider {
))
}
async fn update_mode(&self, session_id: &str, mode: GooseMode) -> Result<(), ProviderError> {
self.mode_by_session
.write()
.await
.insert(session_id.to_string(), mode);
Ok(())
}
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
Ok(CODEX_KNOWN_MODELS.iter().map(|s| s.to_string()).collect())
}
@@ -908,6 +926,7 @@ mod tests {
reasoning_effort: "high".to_string(),
skip_git_check: false,
mcp_config_overrides: Vec::new(),
mode_by_session: tokio::sync::RwLock::new(HashMap::new()),
};
let lines = vec!["Hello, world!".to_string()];
@@ -928,6 +947,7 @@ mod tests {
reasoning_effort: "high".to_string(),
skip_git_check: false,
mcp_config_overrides: Vec::new(),
mode_by_session: tokio::sync::RwLock::new(HashMap::new()),
};
// Test with actual Codex CLI output format
@@ -961,6 +981,7 @@ mod tests {
reasoning_effort: "high".to_string(),
skip_git_check: false,
mcp_config_overrides: Vec::new(),
mode_by_session: tokio::sync::RwLock::new(HashMap::new()),
};
let lines: Vec<String> = vec![];
@@ -1009,6 +1030,7 @@ mod tests {
reasoning_effort: "high".to_string(),
skip_git_check: false,
mcp_config_overrides: Vec::new(),
mode_by_session: tokio::sync::RwLock::new(HashMap::new()),
};
let lines = vec![
@@ -1034,6 +1056,7 @@ mod tests {
reasoning_effort: "high".to_string(),
skip_git_check: false,
mcp_config_overrides: Vec::new(),
mode_by_session: tokio::sync::RwLock::new(HashMap::new()),
};
let lines = vec![
@@ -1106,6 +1129,7 @@ mod tests {
reasoning_effort: "high".to_string(),
skip_git_check: false,
mcp_config_overrides: Vec::new(),
mode_by_session: tokio::sync::RwLock::new(HashMap::new()),
};
let lines: Vec<String> = lines.iter().map(|s| s.to_string()).collect();
@@ -1122,6 +1146,7 @@ mod tests {
reasoning_effort: "high".to_string(),
skip_git_check: false,
mcp_config_overrides: Vec::new(),
mode_by_session: tokio::sync::RwLock::new(HashMap::new()),
};
let lines = vec![
@@ -1212,6 +1237,7 @@ mod tests {
reasoning_effort: "high".to_string(),
skip_git_check: false,
mcp_config_overrides: Vec::new(),
mode_by_session: tokio::sync::RwLock::new(HashMap::new()),
};
let lines = vec![
@@ -1239,4 +1265,19 @@ mod tests {
fn test_default_model() {
assert_eq!(CODEX_DEFAULT_MODEL, "gpt-5.2-codex");
}
#[test_case(GooseMode::Auto, &["--yolo"] ; "auto_yolo")]
#[test_case(GooseMode::SmartApprove, &["--full-auto"] ; "smart_approve_full_auto")]
#[test_case(GooseMode::Approve, &[] as &[&str] ; "approve_no_flags")]
#[test_case(GooseMode::Chat, &["--sandbox", "read-only"] ; "chat_read_only")]
fn test_apply_permission_flags(mode: GooseMode, expected: &[&str]) {
let mut cmd = tokio::process::Command::new("codex");
CodexProvider::apply_permission_flags(&mut cmd, mode).unwrap();
let args: Vec<&str> = cmd
.as_std()
.get_args()
.map(|a| a.to_str().unwrap())
.collect();
assert_eq!(args, expected);
}
}
crates/goose/src/providers/ollama.rs
+1 -10
View File
@@ -5,7 +5,6 @@ use super::openai_compatible::handle_status_openai_compat;
use super::retry::ProviderRetry;
use super::utils::{ImageFormat, RequestLog};
use crate::config::declarative_providers::DeclarativeProviderConfig;
use crate::config::GooseMode;
use crate::conversation::message::Message;
use crate::model::ModelConfig;
use crate::providers::formats::ollama::{create_request, response_to_streaming_message_ollama};
@@ -220,19 +219,11 @@ impl Provider for OllamaProvider {
messages: &[Message],
tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
let config = crate::config::Config::global();
let goose_mode = config.get_goose_mode().unwrap_or(GooseMode::Auto);
let filtered_tools = if goose_mode == GooseMode::Chat {
&[]
} else {
tools
};
let mut payload = create_request(
model_config,
system,
messages,
filtered_tools,
tools,
&ImageFormat::OpenAi,
true,
)?;
crates/goose/src/scheduler.rs
+1
View File
@@ -819,6 +819,7 @@ async fn execute_job(
std::env::current_dir()?,
format!("Scheduled job: {}", job.id),
SessionType::Scheduled,
agent.config.goose_mode,
)
.await?;
crates/goose/src/session/session_manager.rs
+131 -11
View File
@@ -1,4 +1,5 @@
use crate::config::paths::Paths;
use crate::config::GooseMode;
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::model::ModelConfig;
@@ -18,7 +19,7 @@ use std::sync::{Arc, LazyLock};
use tracing::{info, warn};
use utoipa::ToSchema;
pub const CURRENT_SCHEMA_VERSION: i32 = 7;
pub const CURRENT_SCHEMA_VERSION: i32 = 8;
pub const SESSIONS_FOLDER: &str = "sessions";
pub const DB_NAME: &str = "sessions.db";
@@ -93,6 +94,8 @@ pub struct Session {
pub message_count: usize,
pub provider_name: Option<String>,
pub model_config: Option<ModelConfig>,
#[serde(default)]
pub goose_mode: GooseMode,
}
pub struct SessionUpdateBuilder<'a> {
@@ -114,6 +117,7 @@ pub struct SessionUpdateBuilder<'a> {
user_recipe_values: Option<Option<HashMap<String, String>>>,
provider_name: Option<Option<String>>,
model_config: Option<Option<ModelConfig>>,
goose_mode: Option<GooseMode>,
}
#[derive(Serialize, ToSchema, Debug)]
@@ -144,6 +148,7 @@ impl<'a> SessionUpdateBuilder<'a> {
user_recipe_values: None,
provider_name: None,
model_config: None,
goose_mode: None,
}
}
@@ -241,6 +246,11 @@ impl<'a> SessionUpdateBuilder<'a> {
self.model_config = Some(Some(model_config));
self
}
pub fn goose_mode(mut self, mode: GooseMode) -> Self {
self.goose_mode = Some(mode);
self
}
}
pub struct SessionManager {
@@ -269,9 +279,10 @@ impl SessionManager {
working_dir: PathBuf,
name: String,
session_type: SessionType,
goose_mode: GooseMode,
) -> Result<Session> {
self.storage
.create_session(working_dir, name, session_type)
.create_session(working_dir, name, session_type, goose_mode)
.await
}
@@ -417,6 +428,7 @@ impl Default for Session {
message_count: 0,
provider_name: None,
model_config: None,
goose_mode: GooseMode::default(),
}
}
}
@@ -481,6 +493,11 @@ impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session {
message_count: row.try_get("message_count").unwrap_or(0) as usize,
provider_name: row.try_get("provider_name").ok().flatten(),
model_config,
goose_mode: row
.try_get::<String, _>("goose_mode")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or_default(),
})
}
}
@@ -579,7 +596,8 @@ impl SessionStorage {
recipe_json TEXT,
user_recipe_values_json TEXT,
provider_name TEXT,
model_config_json TEXT
model_config_json TEXT,
goose_mode TEXT NOT NULL DEFAULT 'auto'
)
"#,
)
@@ -692,8 +710,8 @@ impl SessionStorage {
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
schedule_id, recipe_json, user_recipe_values_json,
provider_name, model_config_json
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
provider_name, model_config_json, goose_mode
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&session.id)
@@ -715,6 +733,7 @@ impl SessionStorage {
.bind(user_recipe_values_json)
.bind(&session.provider_name)
.bind(model_config_json)
.bind(session.goose_mode.to_string())
.execute(&mut *tx)
.await?;
@@ -887,6 +906,15 @@ impl SessionStorage {
.execute(&mut **tx)
.await?;
}
8 => {
sqlx::query(
r#"
ALTER TABLE sessions ADD COLUMN goose_mode TEXT NOT NULL DEFAULT 'auto'
"#,
)
.execute(&mut **tx)
.await?;
}
_ => {
anyhow::bail!("Unknown migration version: {}", version);
}
@@ -900,6 +928,7 @@ impl SessionStorage {
working_dir: PathBuf,
name: String,
session_type: SessionType,
goose_mode: GooseMode,
) -> Result<Session> {
let pool = self.pool().await?;
let mut tx = pool.begin_with("BEGIN IMMEDIATE").await?;
@@ -907,7 +936,7 @@ impl SessionStorage {
let today = chrono::Utc::now().format("%Y%m%d").to_string();
let session = sqlx::query_as(
r#"
INSERT INTO sessions (id, name, user_set_name, session_type, working_dir, extension_data)
INSERT INTO sessions (id, name, user_set_name, session_type, working_dir, extension_data, goose_mode)
VALUES (
? || '_' || CAST(COALESCE((
SELECT MAX(CAST(SUBSTR(id, 10) AS INTEGER))
@@ -918,7 +947,8 @@ impl SessionStorage {
FALSE,
?,
?,
'{}'
'{}',
?
)
RETURNING *
"#,
@@ -928,6 +958,7 @@ impl SessionStorage {
.bind(&name)
.bind(session_type.to_string())
.bind(&*working_dir.to_string_lossy())
.bind(goose_mode.to_string())
.fetch_one(&mut *tx)
.await?;
@@ -944,7 +975,7 @@ impl SessionStorage {
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
schedule_id, recipe_json, user_recipe_values_json,
provider_name, model_config_json
provider_name, model_config_json, goose_mode
FROM sessions
WHERE id = ?
"#,
@@ -1007,6 +1038,7 @@ impl SessionStorage {
add_update!(builder.user_recipe_values, "user_recipe_values_json");
add_update!(builder.provider_name, "provider_name");
add_update!(builder.model_config, "model_config_json");
add_update!(builder.goose_mode, "goose_mode");
if updates.is_empty() {
return Ok(());
@@ -1072,6 +1104,9 @@ impl SessionStorage {
.transpose()?;
q = q.bind(model_config_json);
}
if let Some(goose_mode) = builder.goose_mode {
q = q.bind(goose_mode.to_string());
}
let pool = self.pool().await?;
let mut tx = pool.begin_with("BEGIN IMMEDIATE").await?;
@@ -1216,7 +1251,7 @@ impl SessionStorage {
s.total_tokens, s.input_tokens, s.output_tokens,
s.accumulated_total_tokens, s.accumulated_input_tokens, s.accumulated_output_tokens,
s.schedule_id, s.recipe_json, s.user_recipe_values_json,
s.provider_name, s.model_config_json,
s.provider_name, s.model_config_json, s.goose_mode,
COUNT(m.id) as message_count
FROM sessions s
INNER JOIN messages m ON s.id = m.session_id
@@ -1304,6 +1339,7 @@ impl SessionStorage {
import.working_dir.clone(),
import.name.clone(),
import.session_type,
import.goose_mode,
)
.await?;
@@ -1347,6 +1383,7 @@ impl SessionStorage {
original_session.working_dir.clone(),
new_name,
original_session.session_type,
original_session.goose_mode,
)
.await?;
@@ -1357,7 +1394,7 @@ impl SessionStorage {
.recipe(original_session.recipe)
.user_recipe_values(original_session.user_recipe_values);
// Preserve provider and model config from original session
// Preserve provider, model config, and goose_mode from original session
if let Some(provider_name) = original_session.provider_name {
builder = builder.provider_name(provider_name);
}
@@ -1366,6 +1403,8 @@ impl SessionStorage {
builder = builder.model_config(model_config);
}
builder = builder.goose_mode(original_session.goose_mode);
builder.apply().await?;
if let Some(conversation) = original_session.conversation {
@@ -1458,6 +1497,7 @@ mod tests {
use super::*;
use crate::conversation::message::{Message, MessageContent};
use tempfile::TempDir;
use test_case::test_case;
const NUM_CONCURRENT_SESSIONS: i32 = 10;
@@ -1529,6 +1569,7 @@ mod tests {
PathBuf::from("/tmp/lock-upgrade-test"),
"Lock Upgrade Session".to_string(),
SessionType::User,
GooseMode::default(),
)
.await
.unwrap();
@@ -1566,7 +1607,12 @@ mod tests {
let description = format!("Test session {}", i);
let session = sm
.create_session(working_dir.clone(), description, SessionType::User)
.create_session(
working_dir.clone(),
description,
SessionType::User,
GooseMode::default(),
)
.await
.unwrap();
@@ -1654,6 +1700,7 @@ mod tests {
PathBuf::from("/tmp/test"),
DESCRIPTION.to_string(),
SessionType::User,
GooseMode::default(),
)
.await
.unwrap();
@@ -1733,4 +1780,77 @@ mod tests {
assert!(imported.user_set_name);
assert_eq!(imported.working_dir, PathBuf::from("/tmp/test"));
}
#[test_case(GooseMode::Approve)]
#[test_case(GooseMode::SmartApprove)]
#[test_case(GooseMode::Chat)]
#[tokio::test]
async fn test_goose_mode_persists(mode: GooseMode) {
let temp_dir = TempDir::new().unwrap();
let sm = SessionManager::new(temp_dir.path().to_path_buf());
let session = sm
.create_session(
temp_dir.path().to_path_buf(),
"test".into(),
SessionType::User,
mode,
)
.await
.unwrap();
let reloaded = sm.get_session(&session.id, false).await.unwrap();
assert_eq!(reloaded.goose_mode, mode);
}
#[tokio::test]
async fn test_goose_mode_update() {
let temp_dir = TempDir::new().unwrap();
let sm = SessionManager::new(temp_dir.path().to_path_buf());
let session = sm
.create_session(
temp_dir.path().to_path_buf(),
"test".into(),
SessionType::User,
GooseMode::default(),
)
.await
.unwrap();
sm.update(&session.id)
.goose_mode(GooseMode::Approve)
.apply()
.await
.unwrap();
let reloaded = sm.get_session(&session.id, false).await.unwrap();
assert_eq!(reloaded.goose_mode, GooseMode::Approve);
}
#[tokio::test]
async fn test_goose_mode_malformed_defaults_to_auto() {
let temp_dir = TempDir::new().unwrap();
let sm = SessionManager::new(temp_dir.path().to_path_buf());
let session = sm
.create_session(
temp_dir.path().to_path_buf(),
"test".into(),
SessionType::User,
GooseMode::Approve,
)
.await
.unwrap();
let pool = &sm.storage().pool;
sqlx::query("UPDATE sessions SET goose_mode = 'garbage' WHERE id = ?")
.bind(&session.id)
.execute(pool)
.await
.unwrap();
let reloaded = sm.get_session(&session.id, false).await.unwrap();
assert_eq!(reloaded.goose_mode, GooseMode::default());
}
}
crates/goose/tests/agent.rs
+3
View File
@@ -340,6 +340,7 @@ mod tests {
use super::*;
use async_trait::async_trait;
use goose::agents::SessionConfig;
use goose::config::GooseMode;
use goose::conversation::message::{Message, MessageContent};
use goose::model::ModelConfig;
use goose::providers::base::{
@@ -427,6 +428,7 @@ mod tests {
PathBuf::default(),
"max-turn-test".to_string(),
SessionType::Hidden,
GooseMode::default(),
)
.await?;
@@ -542,6 +544,7 @@ mod tests {
std::path::PathBuf::from("."),
"Test Session".to_string(),
SessionType::Hidden,
GooseMode::default(),
)
.await
.expect("Failed to create session");
crates/goose/tests/compaction.rs
+2
View File
@@ -2,6 +2,7 @@ use anyhow::Result;
use async_trait::async_trait;
use futures::StreamExt;
use goose::agents::{Agent, AgentEvent, SessionConfig};
use goose::config::GooseMode;
use goose::conversation::message::{Message, MessageContent};
use goose::conversation::Conversation;
use goose::model::ModelConfig;
@@ -215,6 +216,7 @@ async fn setup_test_session(
temp_dir.path().to_path_buf(),
session_name.to_string(),
SessionType::Hidden,
GooseMode::default(),
)
.await?;
crates/goose/tests/providers.rs
+24
View File
@@ -119,6 +119,7 @@ struct ProviderTestConfig {
expected_session_id: fn() -> Arc<dyn ExpectedSessionId>,
test_permissions: bool,
test_smart_approve: bool,
test_mode_update: bool,
test_context_length_exceeded: bool,
expect_context_length_exceeded: bool,
context_length_exceeded: usize,
@@ -141,6 +142,7 @@ impl ProviderTestConfig {
expected_session_id: || Arc::new(EnforceSessionId::default()),
test_permissions: true,
test_smart_approve: true,
test_mode_update: true,
test_context_length_exceeded: true,
expect_context_length_exceeded: true,
context_length_exceeded: 600_000,
@@ -188,6 +190,7 @@ impl ProviderTestConfig {
skip,
expected_session_id: || Arc::new(IgnoreSessionId),
test_smart_approve: false,
test_mode_update: false,
test_context_length_exceeded: false,
..Self::with_llm_provider(name, model_name, &[])
}
@@ -246,6 +249,7 @@ impl ProviderFixture {
std::env::current_dir()?,
"provider_test".to_string(),
SessionType::User,
GooseMode::default(),
)
.await?;
let session_id = session.id;
@@ -586,6 +590,23 @@ impl ProviderFixture {
)
.await
}
async fn test_mode_update(&self) -> Result<()> {
// Start in Auto mode (fixture default), tools auto-approved.
// Switch to Approve mode dynamically via agent.
self.agent
.update_goose_mode(GooseMode::Approve, &self.session_id)
.await?;
// Verify tool call now requires permission (ActionRequired).
// Cancel prevents the task from completing → tool fails.
self.run_permission_test(
Permission::Cancel,
true,
"Use the get_code tool and output only its result.",
"mode_update",
)
.await
}
}
fn load_env() {
@@ -669,6 +690,9 @@ async fn test_provider(config: ProviderTestConfig) -> Result<()> {
.await?;
}
}
if config.test_mode_update {
run_test(GooseMode::Auto).await?.test_mode_update().await?;
}
Ok(())
}
.await;