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

Sessions required (#5548)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2025-11-03 21:04:44 -05:00
committed by GitHub
parent 86c3e42e43
commit 5b93ee587f
31 changed files with 956 additions and 2678 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+104 -182
View File
@@ -61,7 +61,7 @@ use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DEC
use crate::agents::subagent_task_config::TaskConfig;
use crate::conversation::message::{Message, MessageContent, SystemNotificationType, ToolRequest};
use crate::session::extension_data::{EnabledExtensionsState, ExtensionState};
use crate::session::SessionManager;
use crate::session::{Session, SessionManager};
const DEFAULT_MAX_TURNS: u32 = 1000;
const COMPACTION_THINKING_TEXT: &str = "goose is compacting the conversation...";
@@ -75,7 +75,6 @@ pub struct ReplyContext {
pub system_prompt: String,
pub goose_mode: GooseMode,
pub initial_messages: Vec<Message>,
pub config: &'static Config,
}
pub struct ToolCategorizeResult {
@@ -219,16 +218,20 @@ impl Agent {
self.retry_manager.get_attempts().await
}
/// Handle retry logic for the agent reply loop
async fn handle_retry_logic(
&self,
messages: &mut Conversation,
session: &Option<SessionConfig>,
session_config: &SessionConfig,
initial_messages: &[Message],
) -> Result<bool> {
let result = self
.retry_manager
.handle_retry_logic(messages, session, initial_messages, &self.final_output_tool)
.handle_retry_logic(
messages,
session_config,
initial_messages,
&self.final_output_tool,
)
.await?;
match result {
@@ -242,7 +245,6 @@ impl Agent {
async fn prepare_reply_context(
&self,
unfixed_conversation: Conversation,
session: &Option<SessionConfig>,
) -> Result<ReplyContext> {
let unfixed_messages = unfixed_conversation.messages().clone();
let (conversation, issues) = fix_conversation(unfixed_conversation.clone());
@@ -260,9 +262,8 @@ impl Agent {
let config = Config::global();
let (tools, toolshim_tools, system_prompt) = self.prepare_tools_and_prompt().await?;
let goose_mode = Self::determine_goose_mode(session.as_ref(), config);
let goose_mode = config.get_goose_mode().unwrap_or(GooseMode::Auto);
// Update permission inspector mode to match the session mode
self.tool_inspection_manager
.update_permission_inspector_mode(goose_mode)
.await;
@@ -274,7 +275,6 @@ impl Agent {
system_prompt,
goose_mode,
initial_messages,
config,
})
}
@@ -299,7 +299,7 @@ impl Agent {
permission_check_result: &PermissionCheckResult,
message_tool_response: Arc<Mutex<Message>>,
cancel_token: Option<tokio_util::sync::CancellationToken>,
session: Option<SessionConfig>,
session: &Session,
) -> Result<Vec<(String, ToolStream)>> {
let mut tool_futures: Vec<(String, ToolStream)> = Vec::new();
@@ -311,7 +311,7 @@ impl Agent {
tool_call,
request.id.clone(),
cancel_token.clone(),
session.clone(),
session,
)
.await;
@@ -392,7 +392,7 @@ impl Agent {
tool_call: CallToolRequestParam,
request_id: String,
cancellation_token: Option<CancellationToken>,
session: Option<SessionConfig>,
session: &Session,
) -> (String, Result<ToolCallResult, ErrorData>) {
if tool_call.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME {
let arguments = tool_call
@@ -451,17 +451,13 @@ impl Agent {
);
}
};
let (parent_session_id, parent_working_dir) = match session.as_ref() {
Some(s) => (Some(s.id.clone()), s.working_dir.clone()),
None => (None, std::env::current_dir().unwrap_or_default()),
};
// Get extensions from the agent's runtime state rather than global config
// This ensures subagents inherit extensions that were dynamically enabled by the parent
let extensions = self.get_extension_configs().await;
let task_config =
TaskConfig::new(provider, parent_session_id, parent_working_dir, extensions);
TaskConfig::new(provider, &session.id, &session.working_dir, extensions);
let arguments = match tool_call.arguments.clone() {
Some(args) => Value::Object(args),
@@ -731,117 +727,110 @@ impl Agent {
}
}
#[instrument(skip(self, unfixed_conversation, session), fields(user_message))]
#[instrument(skip(self, user_message, session_config), fields(user_message))]
pub async fn reply(
&self,
unfixed_conversation: Conversation,
session: Option<SessionConfig>,
user_message: Message,
session_config: SessionConfig,
cancel_token: Option<CancellationToken>,
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
let is_manual_compact = unfixed_conversation.messages().last().is_some_and(|msg| {
msg.content.iter().any(|c| {
if let MessageContent::Text(text) = c {
text.text.trim() == MANUAL_COMPACT_TRIGGER
} else {
false
}
})
let is_manual_compact = user_message.content.iter().any(|c| {
if let MessageContent::Text(text) = c {
text.text.trim() == MANUAL_COMPACT_TRIGGER
} else {
false
}
});
if !is_manual_compact {
let session_metadata = if let Some(session_config) = &session {
SessionManager::get_session(&session_config.id, false)
.await
.ok()
} else {
None
};
SessionManager::add_message(&session_config.id, &user_message).await?;
let session = SessionManager::get_session(&session_config.id, true).await?;
let needs_auto_compact = crate::context_mgmt::check_if_compaction_needed(
self,
&unfixed_conversation,
None,
session_metadata.as_ref(),
)
.await?;
let conversation = session
.conversation
.clone()
.ok_or_else(|| anyhow::anyhow!("Session {} has no conversation", session_config.id))?;
if !needs_auto_compact {
return self
.reply_internal(unfixed_conversation, session, cancel_token)
.await;
}
}
let needs_auto_compact =
crate::context_mgmt::check_if_compaction_needed(self, &conversation, None, &session)
.await?;
let conversation_to_compact = unfixed_conversation.clone();
let conversation_to_compact = conversation.clone();
Ok(Box::pin(async_stream::try_stream! {
if !is_manual_compact {
let config = crate::config::Config::global();
let threshold = config
.get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
.unwrap_or(DEFAULT_COMPACTION_THRESHOLD);
let threshold_percentage = (threshold * 100.0) as u32;
let final_conversation = if !needs_auto_compact {
conversation
} else {
if !is_manual_compact {
let config = crate::config::Config::global();
let threshold = config
.get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
.unwrap_or(DEFAULT_COMPACTION_THRESHOLD);
let threshold_percentage = (threshold * 100.0) as u32;
let inline_msg = format!(
"Exceeded auto-compact threshold of {}%. Performing auto-compaction...",
threshold_percentage
);
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::InlineMessage,
inline_msg,
)
);
}
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::ThinkingMessage,
COMPACTION_THINKING_TEXT,
)
);
match crate::context_mgmt::compact_messages(self, &conversation_to_compact, false).await {
Ok((compacted_conversation, summarization_usage)) => {
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &compacted_conversation).await?;
Self::update_session_metrics(session_to_store, &summarization_usage, true).await?;
}
yield AgentEvent::HistoryReplaced(compacted_conversation.clone());
let inline_msg = format!(
"Exceeded auto-compact threshold of {}%. Performing auto-compaction...",
threshold_percentage
);
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::InlineMessage,
"Compaction complete",
inline_msg,
)
);
}
if !is_manual_compact {
let mut reply_stream = self.reply_internal(compacted_conversation, session, cancel_token).await?;
while let Some(event) = reply_stream.next().await {
yield event?;
}
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::ThinkingMessage,
COMPACTION_THINKING_TEXT,
)
);
match crate::context_mgmt::compact_messages(self, &conversation_to_compact, false).await {
Ok((compacted_conversation, summarization_usage)) => {
SessionManager::replace_conversation(&session_config.id, &compacted_conversation).await?;
Self::update_session_metrics(&session_config, &summarization_usage, true).await?;
yield AgentEvent::HistoryReplaced(compacted_conversation.clone());
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::InlineMessage,
"Compaction complete",
)
);
compacted_conversation
}
Err(e) => {
yield AgentEvent::Message(
Message::assistant().with_text(
format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session")
)
);
return;
}
}
Err(e) => {
yield AgentEvent::Message(Message::assistant().with_text(
format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session")
));
};
if !is_manual_compact {
let mut reply_stream = self.reply_internal(final_conversation, session_config, session, cancel_token).await?;
while let Some(event) = reply_stream.next().await {
yield event?;
}
}
}))
}
/// Main reply method that handles the actual agent processing
async fn reply_internal(
&self,
conversation: Conversation,
session: Option<SessionConfig>,
session_config: SessionConfig,
session: Session,
cancel_token: Option<CancellationToken>,
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
let context = self.prepare_reply_context(conversation, &session).await?;
let context = self.prepare_reply_context(conversation).await?;
let ReplyContext {
mut conversation,
mut tools,
@@ -849,66 +838,22 @@ impl Agent {
mut system_prompt,
goose_mode,
initial_messages,
config,
} = context;
let reply_span = tracing::Span::current();
self.reset_retry_attempts().await;
// This will need further refactoring. In the ideal world we pass the new message into
// reply and load the existing conversation. Until we get to that point, fetch the conversation
// so far and append the last (user) message that the caller already added.
if let Some(session_config) = &session {
let stored_conversation = SessionManager::get_session(&session_config.id, true)
.await?
.conversation
.ok_or_else(|| {
anyhow::anyhow!("Session {} has no conversation", session_config.id)
})?;
match conversation.len().cmp(&stored_conversation.len()) {
std::cmp::Ordering::Equal => {
if conversation != stored_conversation {
warn!("Session messages mismatch - replacing with incoming");
SessionManager::replace_conversation(&session_config.id, &conversation)
.await?;
}
}
std::cmp::Ordering::Greater
if conversation.len() == stored_conversation.len() + 1 =>
{
let last_message = conversation.last().unwrap();
if let Some(content) = last_message.content.first().and_then(|c| c.as_text()) {
debug!("user_message" = &content);
}
SessionManager::add_message(&session_config.id, last_message).await?;
}
_ => {
warn!(
"Unexpected session state: stored={}, incoming={}. Replacing.",
stored_conversation.len(),
conversation.len()
);
SessionManager::replace_conversation(&session_config.id, &conversation).await?;
}
let provider = self.provider().await?;
let session_id = session_config.id.clone();
tokio::spawn(async move {
if let Err(e) = SessionManager::maybe_update_name(&session_id, provider).await {
warn!("Failed to generate session description: {}", e);
}
let provider = self.provider().await?;
let session_id = session_config.id.clone();
tokio::spawn(async move {
if let Err(e) = SessionManager::maybe_update_name(&session_id, provider).await {
warn!("Failed to generate session description: {}", e);
}
});
}
});
Ok(Box::pin(async_stream::try_stream! {
let _ = reply_span.enter();
let mut turns_taken = 0u32;
let max_turns = session
.as_ref()
.and_then(|s| s.max_turns)
.unwrap_or_else(|| {
config.get_param("GOOSE_MAX_TURNS").unwrap_or(DEFAULT_MAX_TURNS)
});
let max_turns = session_config.max_turns.unwrap_or(DEFAULT_MAX_TURNS);
loop {
if is_token_cancelled(&cancel_token) {
@@ -989,11 +934,8 @@ impl Agent {
}
}
// Record usage for the session
if let Some(ref session_config) = &session {
if let Some(ref usage) = usage {
Self::update_session_metrics(session_config, usage, false).await?;
}
if let Some(ref usage) = usage {
Self::update_session_metrics(&session_config, usage, false).await?;
}
if let Some(response) = response {
@@ -1078,18 +1020,17 @@ impl Agent {
&permission_check_result,
message_tool_response.clone(),
cancel_token.clone(),
session.clone(),
&session,
).await?;
let tool_futures_arc = Arc::new(Mutex::new(tool_futures));
// Process tools requiring approval
let mut tool_approval_stream = self.handle_approval_tool_requests(
&permission_check_result.needs_approval,
tool_futures_arc.clone(),
message_tool_response.clone(),
cancel_token.clone(),
session.clone(),
&session,
&inspection_results,
);
@@ -1136,10 +1077,8 @@ impl Agent {
}
if all_install_successful && !enable_extension_request_ids.is_empty() {
if let Some(ref session_config) = session {
if let Err(e) = self.save_extension_state(session_config).await {
warn!("Failed to save extension state after runtime changes: {}", e);
}
if let Err(e) = self.save_extension_state(&session_config).await {
warn!("Failed to save extension state after runtime changes: {}", e);
}
tools_updated = true;
}
@@ -1168,14 +1107,10 @@ impl Agent {
match crate::context_mgmt::compact_messages(self, &conversation, true).await {
Ok((compacted_conversation, usage)) => {
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &compacted_conversation).await?;
Self::update_session_metrics(session_to_store, &usage, true).await?;
}
SessionManager::replace_conversation(&session_config.id, &compacted_conversation).await?;
Self::update_session_metrics(&session_config, &usage, true).await?;
conversation = compacted_conversation;
did_recovery_compact_this_iteration = true;
yield AgentEvent::HistoryReplaced(conversation.clone());
continue;
}
@@ -1221,7 +1156,7 @@ impl Agent {
} else if did_recovery_compact_this_iteration {
// Avoid setting exit_chat; continue from last user message in the conversation
} else {
match self.handle_retry_logic(&mut conversation, &session, &initial_messages).await {
match self.handle_retry_logic(&mut conversation, &session_config, &initial_messages).await {
Ok(should_retry) => {
if should_retry {
info!("Retry logic triggered, restarting agent loop");
@@ -1242,10 +1177,8 @@ impl Agent {
}
}
if let Some(session_config) = &session {
for msg in &messages_to_add {
SessionManager::add_message(&session_config.id, msg).await?;
}
for msg in &messages_to_add {
SessionManager::add_message(&session_config.id, msg).await?;
}
conversation.extend(messages_to_add);
if exit_chat {
@@ -1257,17 +1190,6 @@ impl Agent {
}))
}
fn determine_goose_mode(session: Option<&SessionConfig>, config: &Config) -> GooseMode {
let mode = session.and_then(|s| s.execution_mode.as_deref());
match mode {
Some("foreground") => GooseMode::Chat,
Some("background") => GooseMode::Auto,
_ => config.get_goose_mode().unwrap_or(GooseMode::Auto),
}
}
/// Extend the system prompt with one line of additional instruction
pub async fn extend_system_prompt(&self, instruction: String) {
let mut prompt_manager = self.prompt_manager.lock().await;
prompt_manager.add_system_prompt_extra(instruction);
crates/goose/src/agents/platform_tools.rs
-1
View File
@@ -33,7 +33,6 @@ pub fn manage_schedule_tool() -> Tool {
"job_id": {"type": "string", "description": "Job identifier for operations on existing jobs"},
"recipe_path": {"type": "string", "description": "Path to recipe file for create action"},
"cron_expression": {"type": "string", "description": "A cron expression for create action. Supports both 5-field (minute hour day month weekday) and 6-field (second minute hour day month weekday) formats. 5-field expressions are automatically converted to 6-field by prepending '0' for seconds."},
"execution_mode": {"type": "string", "description": "Execution mode for create action: 'foreground' or 'background'", "enum": ["foreground", "background"], "default": "background"},
"limit": {"type": "integer", "description": "Limit for sessions list", "default": 50},
"session_id": {"type": "string", "description": "Session identifier for session_content action"}
}
crates/goose/src/agents/retry.rs
+1 -6
View File
@@ -108,18 +108,13 @@ impl RetryManager {
}
}
/// Handle retry logic for the agent reply loop
pub async fn handle_retry_logic(
&self,
messages: &mut Conversation,
session: &Option<SessionConfig>,
session_config: &SessionConfig,
initial_messages: &[Message],
final_output_tool: &Arc<Mutex<Option<crate::agents::final_output_tool::FinalOutputTool>>>,
) -> Result<RetryResult> {
let Some(session_config) = session else {
return Ok(RetryResult::Skipped);
};
let Some(retry_config) = &session_config.retry_config else {
return Ok(RetryResult::Skipped);
};
crates/goose/src/agents/schedule_tool.rs
-1
View File
@@ -186,7 +186,6 @@ impl Agent {
paused: false,
current_session_id: None,
process_start_time: None,
execution_mode: Some(execution_mode.to_string()),
};
match scheduler.add_scheduled_job(job).await {
crates/goose/src/agents/subagent_handler.rs
+24 -52
View File
@@ -1,8 +1,6 @@
use crate::session::session_manager::SessionType;
use crate::{
agents::{
extension::PlatformExtensionContext, subagent_task_config::TaskConfig, Agent, AgentEvent,
SessionConfig,
},
agents::{subagent_task_config::TaskConfig, AgentEvent, SessionConfig},
conversation::{message::Message, Conversation},
execution::manager::AgentManager,
session::SessionManager,
@@ -10,8 +8,8 @@ use crate::{
use anyhow::{anyhow, Result};
use futures::StreamExt;
use rmcp::model::{ErrorCode, ErrorData};
use std::future::Future;
use std::pin::Pin;
use std::{future::Future, sync::Arc};
use tracing::debug;
/// Standalone function to run a complete subagent task with output options
@@ -104,34 +102,18 @@ fn get_agent_messages(
.map_err(|e| anyhow!("Failed to create AgentManager: {}", e))?;
let parent_session_id = task_config.parent_session_id;
let working_dir = task_config.parent_working_dir;
let (agent, session_id) = match parent_session_id {
Some(parent_session_id) => {
let session = SessionManager::create_session(
working_dir.clone(),
format!("Subagent task for: {}", parent_session_id),
)
.await
.map_err(|e| anyhow!("Failed to create a session for sub agent: {}", e))?;
let session = SessionManager::create_session(
working_dir.clone(),
format!("Subagent task for: {}", parent_session_id),
SessionType::SubAgent,
)
.await
.map_err(|e| anyhow!("Failed to create a session for sub agent: {}", e))?;
let agent = agent_manager
.get_or_create_agent(session.id.clone())
.await
.map_err(|e| anyhow!("Failed to get sub agent session file path: {}", e))?;
(agent, Some(session.id))
}
None => {
let agent = Arc::new(Agent::new());
agent
.extension_manager
.set_context(PlatformExtensionContext {
session_id: None,
extension_manager: Some(Arc::downgrade(&agent.extension_manager)),
tool_route_manager: Some(Arc::downgrade(&agent.tool_route_manager)),
})
.await;
(agent, None)
}
};
let agent = agent_manager
.get_or_create_agent(session.id.clone())
.await
.map_err(|e| anyhow!("Failed to get sub agent session file path: {}", e))?;
agent
.update_provider(task_config.provider)
@@ -148,28 +130,18 @@ fn get_agent_messages(
}
}
let mut conversation =
Conversation::new_unvalidated(
vec![Message::user().with_text(text_instruction.clone())],
);
let session_config = if let Some(session_id) = session_id {
Some(SessionConfig {
id: session_id,
working_dir,
schedule_id: None,
execution_mode: None,
max_turns: task_config.max_turns.map(|v| v as u32),
retry_config: None,
})
} else {
None
let user_message = Message::user().with_text(text_instruction);
let mut conversation = Conversation::new_unvalidated(vec![user_message.clone()]);
let session_config = SessionConfig {
id: session.id.clone(),
schedule_id: None,
max_turns: task_config.max_turns.map(|v| v as u32),
retry_config: None,
};
let session_id = session_config.as_ref().map(|s| s.id.clone());
let mut stream = crate::session_context::with_session_id(session_id, async {
agent
.reply(conversation.clone(), session_config, None)
.await
let mut stream = crate::session_context::with_session_id(Some(session.id.clone()), async {
agent.reply(user_message, session_config, None).await
})
.await
.map_err(|e| anyhow!("Failed to get reply from agent: {}", e))?;
crates/goose/src/agents/subagent_task_config.rs
+6 -7
View File
@@ -2,7 +2,7 @@ use crate::agents::ExtensionConfig;
use crate::providers::base::Provider;
use std::env;
use std::fmt;
use std::path::PathBuf;
use std::path::{Path, PathBuf};
use std::sync::Arc;
/// Default maximum number of turns for task execution
@@ -15,7 +15,7 @@ pub const GOOSE_SUBAGENT_MAX_TURNS_ENV_VAR: &str = "GOOSE_SUBAGENT_MAX_TURNS";
#[derive(Clone)]
pub struct TaskConfig {
pub provider: Arc<dyn Provider>,
pub parent_session_id: Option<String>,
pub parent_session_id: String,
pub parent_working_dir: PathBuf,
pub extensions: Vec<ExtensionConfig>,
pub max_turns: Option<usize>,
@@ -34,17 +34,16 @@ impl fmt::Debug for TaskConfig {
}
impl TaskConfig {
/// Create a new TaskConfig with all required dependencies
pub fn new(
provider: Arc<dyn Provider>,
parent_session_id: Option<String>,
parent_working_dir: PathBuf,
parent_session_id: &str,
parent_working_dir: &Path,
extensions: Vec<ExtensionConfig>,
) -> Self {
Self {
provider,
parent_session_id,
parent_working_dir,
parent_session_id: parent_session_id.to_owned(),
parent_working_dir: parent_working_dir.to_owned(),
extensions,
max_turns: Some(
env::var(GOOSE_SUBAGENT_MAX_TURNS_ENV_VAR)
crates/goose/src/agents/tool_execution.rs
+4 -3
View File
@@ -29,8 +29,9 @@ impl From<ToolResult<Vec<Content>>> for ToolCallResult {
}
use super::agent::{tool_stream, ToolStream};
use crate::agents::{Agent, SessionConfig};
use crate::agents::Agent;
use crate::conversation::message::{Message, ToolRequest};
use crate::session::Session;
use crate::tool_inspection::get_security_finding_id_from_results;
pub const DECLINED_RESPONSE: &str = "The user has declined to run this tool. \
@@ -53,7 +54,7 @@ impl Agent {
tool_futures: Arc<Mutex<Vec<(String, ToolStream)>>>,
message_tool_response: Arc<Mutex<Message>>,
cancellation_token: Option<CancellationToken>,
session: Option<SessionConfig>,
session: &'a Session,
inspection_results: &'a [crate::tool_inspection::InspectionResult],
) -> BoxStream<'a, anyhow::Result<Message>> {
try_stream! {
@@ -93,7 +94,7 @@ impl Agent {
}
if confirmation.permission == Permission::AllowOnce || confirmation.permission == Permission::AlwaysAllow {
let (req_id, tool_result) = self.dispatch_tool_call(tool_call.clone(), request.id.clone(), cancellation_token.clone(), session.clone()).await;
let (req_id, tool_result) = self.dispatch_tool_call(tool_call.clone(), request.id.clone(), cancellation_token.clone(), session).await;
let mut futures = tool_futures.lock().await;
futures.push((req_id, match tool_result {
crates/goose/src/agents/types.rs
+1 -6
View File
@@ -2,7 +2,6 @@ use crate::mcp_utils::ToolResult;
use crate::providers::base::Provider;
use rmcp::model::{Content, Tool};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex};
use utoipa::ToSchema;
@@ -84,14 +83,10 @@ pub struct FrontendTool {
/// Session configuration for an agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionConfig {
/// Unique identifier for the session
/// Identifier of the underlying Session
pub id: String,
/// Working directory for the session
pub working_dir: PathBuf,
/// ID of the schedule that triggered this session, if any
pub schedule_id: Option<String>,
/// Execution mode for scheduled jobs: "foreground" or "background"
pub execution_mode: Option<String>,
/// Maximum number of turns (iterations) allowed without user input
pub max_turns: Option<u32>,
/// Retry configuration for automated validation and recovery
crates/goose/src/context_mgmt/mod.rs
+2 -3
View File
@@ -145,11 +145,10 @@ pub async fn check_if_compaction_needed(
agent: &Agent,
conversation: &Conversation,
threshold_override: Option<f64>,
session_metadata: Option<&crate::session::Session>,
session: &crate::session::Session,
) -> Result<bool> {
let messages = conversation.messages();
let config = Config::global();
// TODO(Douwe): check the default here; it seems to reset to 0.3 sometimes
let threshold = threshold_override.unwrap_or_else(|| {
config
.get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
@@ -159,7 +158,7 @@ pub async fn check_if_compaction_needed(
let provider = agent.provider().await?;
let context_limit = provider.get_model_config().context_limit();
let (current_tokens, token_source) = match session_metadata.and_then(|m| m.total_tokens) {
let (current_tokens, token_source) = match session.total_tokens {
Some(tokens) => (tokens as usize, "session metadata"),
None => {
let token_counter = create_token_counter()
crates/goose/src/scheduler.rs
+5 -13
View File
@@ -21,6 +21,7 @@ use crate::providers::base::Provider as GooseProvider; // Alias to avoid conflic
use crate::providers::create;
use crate::recipe::Recipe;
use crate::scheduler_trait::SchedulerTrait;
use crate::session::session_manager::SessionType;
use crate::session::{Session, SessionManager};
// Track running tasks with their abort handles
@@ -152,8 +153,6 @@ pub struct ScheduledJob {
pub current_session_id: Option<String>,
#[serde(default)]
pub process_start_time: Option<DateTime<Utc>>,
#[serde(default)]
pub execution_mode: Option<String>, // "foreground" or "background"
}
async fn persist_jobs_from_arc(
@@ -1160,8 +1159,6 @@ async fn run_scheduled_job_internal(
});
}
tracing::info!("Agent configured with provider for job '{}'", job.id);
let execution_mode = job.execution_mode.as_deref().unwrap_or("background");
tracing::info!("Job '{}' running in {} mode", job.id, execution_mode);
let current_dir = match std::env::current_dir() {
Ok(cd) => cd,
@@ -1173,10 +1170,10 @@ async fn run_scheduled_job_internal(
}
};
// Create session upfront
let session = match SessionManager::create_session(
current_dir.clone(),
format!("Scheduled job: {}", job.id),
SessionType::Scheduled,
)
.await
{
@@ -1204,23 +1201,19 @@ async fn run_scheduled_job_internal(
.or(recipe.instructions.as_ref())
.unwrap();
let mut conversation =
Conversation::new_unvalidated(vec![Message::user().with_text(prompt_text.clone())]);
let user_message = Message::user().with_text(prompt_text);
let mut conversation = Conversation::new_unvalidated(vec![user_message.clone()]);
let session_config = SessionConfig {
id: session.id.clone(),
working_dir: current_dir.clone(),
schedule_id: Some(job.id.clone()),
execution_mode: job.execution_mode.clone(),
max_turns: None,
retry_config: None,
};
let session_id = Some(session_config.id.clone());
match crate::session_context::with_session_id(session_id, async {
agent
.reply(conversation.clone(), Some(session_config.clone()), None)
.await
agent.reply(user_message, session_config, None).await
})
.await
{
@@ -1455,7 +1448,6 @@ mod tests {
paused: false,
current_session_id: None,
process_start_time: None,
execution_mode: Some("background".to_string()), // Default for test
};
let mock_model_config = ModelConfig::new_or_fail("test_model");
crates/goose/src/session/mod.rs
+1 -1
View File
@@ -6,4 +6,4 @@ pub mod session_manager;
pub use diagnostics::generate_diagnostics;
pub use extension_data::{EnabledExtensionsState, ExtensionData, ExtensionState, TodoState};
pub use session_manager::{Session, SessionInsights, SessionManager};
pub use session_manager::{Session, SessionInsights, SessionManager, SessionType};
crates/goose/src/session/session_manager.rs
+141 -45
View File
@@ -18,7 +18,47 @@ use tokio::sync::OnceCell;
use tracing::{info, warn};
use utoipa::ToSchema;
const CURRENT_SCHEMA_VERSION: i32 = 4;
const CURRENT_SCHEMA_VERSION: i32 = 5;
#[derive(Debug, Clone, Copy, Serialize, Deserialize, ToSchema, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum SessionType {
User,
Scheduled,
SubAgent,
Hidden,
}
impl Default for SessionType {
fn default() -> Self {
Self::User
}
}
impl std::fmt::Display for SessionType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SessionType::User => write!(f, "user"),
SessionType::SubAgent => write!(f, "sub_agent"),
SessionType::Hidden => write!(f, "hidden"),
SessionType::Scheduled => write!(f, "scheduled"),
}
}
}
impl std::str::FromStr for SessionType {
type Err = anyhow::Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"user" => Ok(SessionType::User),
"sub_agent" => Ok(SessionType::SubAgent),
"hidden" => Ok(SessionType::Hidden),
"scheduled" => Ok(SessionType::Scheduled),
_ => Err(anyhow::anyhow!("Invalid session type: {}", s)),
}
}
}
static SESSION_STORAGE: OnceCell<Arc<SessionStorage>> = OnceCell::const_new();
@@ -27,11 +67,12 @@ pub struct Session {
pub id: String,
#[schema(value_type = String)]
pub working_dir: PathBuf,
// Allow importing session exports from before 'description' was renamed to 'name'
#[serde(alias = "description")]
pub name: String,
#[serde(default)]
pub user_set_name: bool,
#[serde(default)]
pub session_type: SessionType,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
pub extension_data: ExtensionData,
@@ -52,6 +93,7 @@ pub struct SessionUpdateBuilder {
session_id: String,
name: Option<String>,
user_set_name: Option<bool>,
session_type: Option<SessionType>,
working_dir: Option<PathBuf>,
extension_data: Option<ExtensionData>,
total_tokens: Option<Option<i32>>,
@@ -78,6 +120,7 @@ impl SessionUpdateBuilder {
session_id,
name: None,
user_set_name: None,
session_type: None,
working_dir: None,
extension_data: None,
total_tokens: None,
@@ -110,6 +153,11 @@ impl SessionUpdateBuilder {
self
}
pub fn session_type(mut self, session_type: SessionType) -> Self {
self.session_type = Some(session_type);
self
}
pub fn working_dir(mut self, working_dir: PathBuf) -> Self {
self.working_dir = Some(working_dir);
self
@@ -183,10 +231,14 @@ impl SessionManager {
.map(Arc::clone)
}
pub async fn create_session(working_dir: PathBuf, name: String) -> Result<Session> {
pub async fn create_session(
working_dir: PathBuf,
name: String,
session_type: SessionType,
) -> Result<Session> {
Self::instance()
.await?
.create_session(working_dir, name)
.create_session(working_dir, name, session_type)
.await
}
@@ -306,6 +358,7 @@ impl Default for Session {
working_dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
name: String::new(),
user_set_name: false,
session_type: SessionType::default(),
created_at: Default::default(),
updated_at: Default::default(),
extension_data: ExtensionData::default(),
@@ -353,11 +406,17 @@ impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session {
let user_set_name = row.try_get("user_set_name").unwrap_or(false);
let session_type_str: String = row
.try_get("session_type")
.unwrap_or_else(|_| "user".to_string());
let session_type = session_type_str.parse().unwrap_or_default();
Ok(Session {
id: row.try_get("id")?,
working_dir: PathBuf::from(row.try_get::<String, _>("working_dir")?),
name,
user_set_name,
session_type,
created_at: row.try_get("created_at")?,
updated_at: row.try_get("updated_at")?,
extension_data: serde_json::from_str(&row.try_get::<String, _>("extension_data")?)
@@ -446,6 +505,7 @@ impl SessionStorage {
name TEXT NOT NULL DEFAULT '',
description TEXT NOT NULL DEFAULT '',
user_set_name BOOLEAN DEFAULT FALSE,
session_type TEXT NOT NULL DEFAULT 'user',
working_dir TEXT NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
@@ -491,6 +551,9 @@ impl SessionStorage {
sqlx::query("CREATE INDEX idx_sessions_updated ON sessions(updated_at DESC)")
.execute(&pool)
.await?;
sqlx::query("CREATE INDEX idx_sessions_type ON sessions(session_type)")
.execute(&pool)
.await?;
Ok(Self { pool })
}
@@ -553,31 +616,32 @@ impl SessionStorage {
sqlx::query(
r#"
INSERT INTO sessions (
id, name, user_set_name, working_dir, created_at, updated_at, extension_data,
id, name, user_set_name, session_type, working_dir, created_at, updated_at, extension_data,
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
schedule_id, recipe_json, user_recipe_values_json
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&session.id)
.bind(&session.name)
.bind(session.user_set_name)
.bind(session.working_dir.to_string_lossy().as_ref())
.bind(session.created_at)
.bind(session.updated_at)
.bind(serde_json::to_string(&session.extension_data)?)
.bind(session.total_tokens)
.bind(session.input_tokens)
.bind(session.output_tokens)
.bind(session.accumulated_total_tokens)
.bind(session.accumulated_input_tokens)
.bind(session.accumulated_output_tokens)
.bind(&session.schedule_id)
.bind(recipe_json)
.bind(user_recipe_values_json)
.execute(&self.pool)
.await?;
.bind(&session.id)
.bind(&session.name)
.bind(session.user_set_name)
.bind(session.session_type.to_string())
.bind(session.working_dir.to_string_lossy().as_ref())
.bind(session.created_at)
.bind(session.updated_at)
.bind(serde_json::to_string(&session.extension_data)?)
.bind(session.total_tokens)
.bind(session.input_tokens)
.bind(session.output_tokens)
.bind(session.accumulated_total_tokens)
.bind(session.accumulated_input_tokens)
.bind(session.accumulated_output_tokens)
.bind(&session.schedule_id)
.bind(recipe_json)
.bind(user_recipe_values_json)
.execute(&self.pool)
.await?;
if let Some(conversation) = &session.conversation {
self.replace_conversation(&session.id, conversation).await?;
@@ -687,6 +751,19 @@ impl SessionStorage {
.execute(&self.pool)
.await?;
}
5 => {
sqlx::query(
r#"
ALTER TABLE sessions ADD COLUMN session_type TEXT NOT NULL DEFAULT 'user'
"#,
)
.execute(&self.pool)
.await?;
sqlx::query("CREATE INDEX idx_sessions_type ON sessions(session_type)")
.execute(&self.pool)
.await?;
}
_ => {
anyhow::bail!("Unknown migration version: {}", version);
}
@@ -695,11 +772,16 @@ impl SessionStorage {
Ok(())
}
async fn create_session(&self, working_dir: PathBuf, name: String) -> Result<Session> {
async fn create_session(
&self,
working_dir: PathBuf,
name: String,
session_type: SessionType,
) -> Result<Session> {
let today = chrono::Utc::now().format("%Y%m%d").to_string();
Ok(sqlx::query_as(
r#"
INSERT INTO sessions (id, name, user_set_name, working_dir, extension_data)
INSERT INTO sessions (id, name, user_set_name, session_type, working_dir, extension_data)
VALUES (
? || '_' || CAST(COALESCE((
SELECT MAX(CAST(SUBSTR(id, 10) AS INTEGER))
@@ -709,23 +791,25 @@ impl SessionStorage {
?,
FALSE,
?,
?,
'{}'
)
RETURNING *
"#,
)
.bind(&today)
.bind(&today)
.bind(&name)
.bind(working_dir.to_string_lossy().as_ref())
.fetch_one(&self.pool)
.await?)
.bind(&today)
.bind(&today)
.bind(&name)
.bind(session_type.to_string())
.bind(working_dir.to_string_lossy().as_ref())
.fetch_one(&self.pool)
.await?)
}
async fn get_session(&self, id: &str, include_messages: bool) -> Result<Session> {
let mut session = sqlx::query_as::<_, Session>(
r#"
SELECT id, working_dir, name, description, user_set_name, created_at, updated_at, extension_data,
SELECT id, working_dir, name, description, user_set_name, session_type, created_at, updated_at, extension_data,
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
schedule_id, recipe_json, user_recipe_values_json
@@ -733,10 +817,10 @@ impl SessionStorage {
WHERE id = ?
"#,
)
.bind(id)
.fetch_optional(&self.pool)
.await?
.ok_or_else(|| anyhow::anyhow!("Session not found"))?;
.bind(id)
.fetch_optional(&self.pool)
.await?
.ok_or_else(|| anyhow::anyhow!("Session not found"))?;
if include_messages {
let conv = self.get_conversation(&session.id).await?;
@@ -773,6 +857,7 @@ impl SessionStorage {
add_update!(builder.name, "name");
add_update!(builder.user_set_name, "user_set_name");
add_update!(builder.session_type, "session_type");
add_update!(builder.working_dir, "working_dir");
add_update!(builder.extension_data, "extension_data");
add_update!(builder.total_tokens, "total_tokens");
@@ -803,6 +888,9 @@ impl SessionStorage {
if let Some(user_set_name) = builder.user_set_name {
q = q.bind(user_set_name);
}
if let Some(session_type) = builder.session_type {
q = q.bind(session_type.to_string());
}
if let Some(wd) = builder.working_dir {
q = q.bind(wd.to_string_lossy().to_string());
}
@@ -872,7 +960,6 @@ impl SessionStorage {
let mut message = Message::new(role, created_timestamp, content);
message.metadata = metadata;
// TODO(Douwe): make id required
message = message.with_id(format!("msg_{}_{}", session_id, idx));
messages.push(message);
}
@@ -942,20 +1029,21 @@ impl SessionStorage {
async fn list_sessions(&self) -> Result<Vec<Session>> {
sqlx::query_as::<_, Session>(
r#"
SELECT s.id, s.working_dir, s.name, s.description, s.user_set_name, s.created_at, s.updated_at, s.extension_data,
SELECT s.id, s.working_dir, s.name, s.description, s.user_set_name, s.session_type, s.created_at, s.updated_at, s.extension_data,
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,
COUNT(m.id) as message_count
FROM sessions s
INNER JOIN messages m ON s.id = m.session_id
WHERE s.session_type = 'user' OR s.session_type = 'scheduled'
GROUP BY s.id
ORDER BY s.updated_at DESC
"#,
)
.fetch_all(&self.pool)
.await
.map_err(Into::into)
.fetch_all(&self.pool)
.await
.map_err(Into::into)
}
async fn delete_session(&self, session_id: &str) -> Result<()> {
@@ -1008,7 +1096,11 @@ impl SessionStorage {
let import: Session = serde_json::from_str(json)?;
let session = self
.create_session(import.working_dir.clone(), import.name.clone())
.create_session(
import.working_dir.clone(),
import.name.clone(),
import.session_type,
)
.await?;
let mut builder = SessionUpdateBuilder::new(session.id.clone())
@@ -1084,7 +1176,7 @@ mod tests {
let description = format!("Test session {}", i);
let session = session_storage
.create_session(working_dir.clone(), description)
.create_session(working_dir.clone(), description, SessionType::User)
.await
.unwrap();
@@ -1176,7 +1268,11 @@ mod tests {
let storage = Arc::new(SessionStorage::create(&db_path).await.unwrap());
let original = storage
.create_session(PathBuf::from("/tmp/test"), DESCRIPTION.to_string())
.create_session(
PathBuf::from("/tmp/test"),
DESCRIPTION.to_string(),
SessionType::User,
)
.await
.unwrap();