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

fix(acp): separate acp sessions from user sessions (#7857)

Browse Source 
Signed-off-by: Rabi Mishra <mishra.rabi@gmail.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Rabi Mishra
2026-03-26 22:22:42 +05:30
committed by GitHub
parent 3d8d2aa2dc
commit f5ad3384b8
7 changed files with 251 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-acp/src/server.rs
+3 -3
View File
@@ -848,7 +848,7 @@ impl GooseAcpAgent {
.create_session(
args.cwd.clone(),
"ACP Session".to_string(),
SessionType::User,
SessionType::Acp,
self.goose_mode,
)
.await
@@ -1246,7 +1246,7 @@ impl GooseAcpAgent {
async fn on_list_sessions(&self) -> Result<ListSessionsResponse, sacp::Error> {
let sessions = self
.session_manager
.list_sessions()
.list_sessions_by_types(&[SessionType::Acp])
.await
.map_err(|e| sacp::Error::internal_error().data(e.to_string()))?;
let session_infos: Vec<SessionInfo> = sessions
@@ -1422,7 +1422,7 @@ impl GooseAcpAgent {
) -> Result<ImportSessionResponse, sacp::Error> {
let session = self
.session_manager
.import_session(&req.data)
.import_session(&req.data, Some(SessionType::Acp))
.await
.map_err(|e| sacp::Error::internal_error().data(e.to_string()))?;
let session_json = serde_json::to_value(&session)
crates/goose-server/src/routes/session.rs
+10 -3
View File
@@ -11,7 +11,7 @@ use axum::{
};
use goose::agents::ExtensionConfig;
use goose::recipe::Recipe;
use goose::session::session_manager::SessionInsights;
use goose::session::session_manager::{SessionInsights, SessionType};
use goose::session::{EnabledExtensionsState, Session};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
@@ -357,7 +357,7 @@ async fn import_session(
) -> Result<Json<Session>, StatusCode> {
let session = state
.session_manager()
.import_session(&request.json)
.import_session(&request.json, Some(SessionType::User))
.await
.map_err(|_| StatusCode::BAD_REQUEST)?;
@@ -583,7 +583,14 @@ async fn search_sessions(
let search_results = state
.session_manager()
.search_chat_history(query, Some(limit), after_date, before_date, None)
.search_chat_history(
query,
Some(limit),
after_date,
before_date,
None,
vec![SessionType::User, SessionType::Scheduled],
)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
crates/goose/src/agents/platform_extensions/chatrecall.rs
+9
View File
@@ -1,6 +1,7 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use crate::session::session_manager::SessionType;
use anyhow::Result;
use async_trait::async_trait;
use indoc::indoc;
@@ -58,6 +59,13 @@ impl ChatRecallClient {
Ok(Self { info, context })
}
fn search_session_types(&self) -> Vec<SessionType> {
match self.context.session.as_ref().map(|s| s.session_type) {
Some(SessionType::Acp) => vec![SessionType::Acp],
_ => vec![SessionType::User, SessionType::Scheduled],
}
}
#[allow(clippy::too_many_lines)]
async fn handle_chatrecall(
&self,
@@ -177,6 +185,7 @@ impl ChatRecallClient {
after_date,
before_date,
exclude_session_id,
self.search_session_types(),
)
.await
{
crates/goose/src/session/chat_history_search.rs
+18
View File
@@ -1,4 +1,5 @@
use crate::conversation::message::MessageContent;
use crate::session::session_manager::SessionType;
use anyhow::Result;
use chrono::{DateTime, Utc};
use serde::Serialize;
@@ -52,6 +53,7 @@ pub struct ChatHistorySearch<'a> {
after_date: Option<DateTime<Utc>>,
before_date: Option<DateTime<Utc>>,
exclude_session_id: Option<String>,
session_types: Vec<SessionType>,
}
impl<'a> ChatHistorySearch<'a> {
@@ -62,6 +64,7 @@ impl<'a> ChatHistorySearch<'a> {
after_date: Option<DateTime<Utc>>,
before_date: Option<DateTime<Utc>>,
exclude_session_id: Option<String>,
session_types: Vec<SessionType>,
) -> Self {
Self {
pool,
@@ -70,6 +73,7 @@ impl<'a> ChatHistorySearch<'a> {
after_date,
before_date,
exclude_session_id,
session_types,
}
}
@@ -102,6 +106,10 @@ impl<'a> ChatHistorySearch<'a> {
query_builder = query_builder.bind(exclude_id);
}
for t in &self.session_types {
query_builder = query_builder.bind(t.to_string());
}
if let Some(after) = self.after_date {
query_builder = query_builder.bind(after);
}
@@ -159,6 +167,16 @@ impl<'a> ChatHistorySearch<'a> {
sql.push_str(" AND s.id != ?");
}
if !self.session_types.is_empty() {
let placeholders: String = self
.session_types
.iter()
.map(|_| "?")
.collect::<Vec<_>>()
.join(", ");
sql.push_str(&format!(" AND s.session_type IN ({})", placeholders));
}
if self.after_date.is_some() {
sql.push_str(" AND m.timestamp >= ?");
}
crates/goose/src/session/session_manager.rs
+208 -44
View File
@@ -19,12 +19,25 @@ use std::sync::{Arc, LazyLock};
use tracing::{info, warn};
use utoipa::ToSchema;
pub const CURRENT_SCHEMA_VERSION: i32 = 8;
pub const CURRENT_SCHEMA_VERSION: i32 = 9;
pub const SESSIONS_FOLDER: &str = "sessions";
pub const DB_NAME: &str = "sessions.db";
#[derive(Debug, Clone, Copy, Serialize, Deserialize, ToSchema, PartialEq, Eq, Default)]
#[derive(
Debug,
Clone,
Copy,
Serialize,
Deserialize,
ToSchema,
PartialEq,
Eq,
Default,
strum::Display,
strum::EnumString,
)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum SessionType {
#[default]
User,
@@ -33,35 +46,7 @@ pub enum SessionType {
Hidden,
Terminal,
Gateway,
}
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"),
SessionType::Terminal => write!(f, "terminal"),
SessionType::Gateway => write!(f, "gateway"),
}
}
}
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),
"terminal" => Ok(SessionType::Terminal),
"gateway" => Ok(SessionType::Gateway),
_ => Err(anyhow::anyhow!("Invalid session type: {}", s)),
}
}
Acp,
}
static SESSION_STORAGE: LazyLock<Arc<SessionStorage>> =
@@ -323,15 +308,23 @@ impl SessionManager {
}
pub async fn get_insights(&self) -> Result<SessionInsights> {
self.storage.get_insights().await
self.storage
.get_insights(&[SessionType::User, SessionType::Scheduled])
.await
}
pub async fn export_session(&self, id: &str) -> Result<String> {
self.storage.export_session(id).await
}
pub async fn import_session(&self, json: &str) -> Result<Session> {
self.storage.import_session(self, json).await
pub async fn import_session(
&self,
json: &str,
session_type_override: Option<SessionType>,
) -> Result<Session> {
self.storage
.import_session(self, json, session_type_override)
.await
}
pub async fn copy_session(&self, session_id: &str, new_name: String) -> Result<Session> {
@@ -376,9 +369,17 @@ impl SessionManager {
after_date: Option<chrono::DateTime<chrono::Utc>>,
before_date: Option<chrono::DateTime<chrono::Utc>>,
exclude_session_id: Option<String>,
session_types: Vec<SessionType>,
) -> Result<crate::session::chat_history_search::ChatRecallResults> {
self.storage
.search_chat_history(query, limit, after_date, before_date, exclude_session_id)
.search_chat_history(
query,
limit,
after_date,
before_date,
exclude_session_id,
session_types,
)
.await
}
@@ -919,6 +920,19 @@ impl SessionStorage {
.execute(&mut **tx)
.await?;
}
9 => {
sqlx::query(
r#"
UPDATE sessions
SET session_type = 'acp'
WHERE session_type = 'user'
AND name = 'ACP Session'
AND user_set_name = FALSE
"#,
)
.execute(&mut **tx)
.await?;
}
_ => {
anyhow::bail!("Unknown migration version: {}", version);
}
@@ -1316,17 +1330,32 @@ impl SessionStorage {
Ok(())
}
async fn get_insights(&self) -> Result<SessionInsights> {
let pool = self.pool().await?;
let row = sqlx::query_as::<_, (i64, Option<i64>)>(
async fn get_insights(&self, types: &[SessionType]) -> Result<SessionInsights> {
if types.is_empty() {
return Ok(SessionInsights {
total_sessions: 0,
total_tokens: 0,
});
}
let placeholders: String = types.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
let query = format!(
r#"
SELECT COUNT(*) as total_sessions,
COALESCE(SUM(COALESCE(accumulated_total_tokens, total_tokens, 0)), 0) as total_tokens
FROM sessions
WHERE session_type IN ({})
"#,
)
.fetch_one(pool)
.await?;
placeholders
);
let pool = self.pool().await?;
let mut q = sqlx::query_as::<_, (i64, Option<i64>)>(&query);
for t in types {
q = q.bind(t.to_string());
}
let row = q.fetch_one(pool).await?;
Ok(SessionInsights {
total_sessions: row.0 as usize,
@@ -1343,6 +1372,7 @@ impl SessionStorage {
&self,
session_manager: &SessionManager,
json: &str,
session_type_override: Option<SessionType>,
) -> Result<Session> {
let import: Session = serde_json::from_str(json)?;
@@ -1350,7 +1380,7 @@ impl SessionStorage {
.create_session(
import.working_dir.clone(),
import.name.clone(),
import.session_type,
session_type_override.unwrap_or(import.session_type),
import.goose_mode,
)
.await?;
@@ -1445,6 +1475,7 @@ impl SessionStorage {
after_date: Option<chrono::DateTime<chrono::Utc>>,
before_date: Option<chrono::DateTime<chrono::Utc>>,
exclude_session_id: Option<String>,
session_types: Vec<SessionType>,
) -> Result<crate::session::chat_history_search::ChatRecallResults> {
use crate::session::chat_history_search::ChatHistorySearch;
@@ -1456,6 +1487,7 @@ impl SessionStorage {
after_date,
before_date,
exclude_session_id,
session_types,
)
.execute()
.await
@@ -1753,7 +1785,7 @@ mod tests {
.unwrap();
let exported = sm.export_session(&original.id).await.unwrap();
let imported = sm.import_session(&exported).await.unwrap();
let imported = sm.import_session(&exported, None).await.unwrap();
assert_ne!(imported.id, original.id);
assert_eq!(imported.name, DESCRIPTION);
@@ -1770,6 +1802,69 @@ mod tests {
assert_eq!(conversation.messages()[1].role, Role::Assistant);
}
#[tokio::test]
async fn test_list_sessions_filters_by_type() {
let temp_dir = TempDir::new().unwrap();
let sm = SessionManager::new(temp_dir.path().to_path_buf());
let user_session = sm
.create_session(
PathBuf::from("/tmp/test"),
"User session".to_string(),
SessionType::User,
GooseMode::default(),
)
.await
.unwrap();
sm.add_message(
&user_session.id,
&Message {
id: None,
role: Role::User,
created: chrono::Utc::now().timestamp_millis(),
content: vec![MessageContent::text("hello world")],
metadata: Default::default(),
},
)
.await
.unwrap();
let acp_session = sm
.create_session(
PathBuf::from("/tmp/test"),
"ACP session".to_string(),
SessionType::Acp,
GooseMode::default(),
)
.await
.unwrap();
sm.add_message(
&acp_session.id,
&Message {
id: None,
role: Role::User,
created: chrono::Utc::now().timestamp_millis(),
content: vec![MessageContent::text("hello acp")],
metadata: Default::default(),
},
)
.await
.unwrap();
let default_sessions = sm.list_sessions().await.unwrap();
assert_eq!(default_sessions.len(), 1);
assert_eq!(default_sessions[0].name, "User session");
let acp_sessions = sm
.list_sessions_by_types(&[SessionType::Acp])
.await
.unwrap();
assert_eq!(acp_sessions.len(), 1);
assert_eq!(acp_sessions[0].name, "ACP session");
}
#[tokio::test]
async fn test_import_session_with_description_field() {
const OLD_FORMAT_JSON: &str = r#"{
@@ -1786,7 +1881,7 @@ mod tests {
let temp_dir = TempDir::new().unwrap();
let sm = SessionManager::new(temp_dir.path().to_path_buf());
let imported = sm.import_session(OLD_FORMAT_JSON).await.unwrap();
let imported = sm.import_session(OLD_FORMAT_JSON, None).await.unwrap();
assert_eq!(imported.name, "Old format session");
assert!(imported.user_set_name);
@@ -1865,4 +1960,73 @@ mod tests {
let reloaded = sm.get_session(&session.id, false).await.unwrap();
assert_eq!(reloaded.goose_mode, GooseMode::default());
}
#[tokio::test]
async fn test_acp_session_migration() {
let temp_dir = TempDir::new().unwrap();
let db_path = temp_dir.path().join(SESSIONS_FOLDER).join(DB_NAME);
if let Some(parent) = db_path.parent() {
std::fs::create_dir_all(parent).unwrap();
}
let pool = SqlitePoolOptions::new()
.connect_with(
SqliteConnectOptions::new()
.filename(&db_path)
.create_if_missing(true),
)
.await
.unwrap();
SessionStorage::create_schema(&pool).await.unwrap();
// Demote the schema back to v8 to simulate a database
// that has never seen migration 9.
sqlx::query("UPDATE schema_version SET version = 8")
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO sessions (id, name, user_set_name, session_type, working_dir, extension_data, goose_mode)
VALUES (?, ?, ?, ?, ?, ?, ?)",
)
.bind("user_id")
.bind("User Session")
.bind(false)
.bind("user")
.bind("/tmp")
.bind("{}")
.bind("auto")
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO sessions (id, name, user_set_name, session_type, working_dir, extension_data, goose_mode)
VALUES (?, ?, ?, ?, ?, ?, ?)",
)
.bind("acp_id")
.bind("ACP Session")
.bind(false)
.bind("user")
.bind("/tmp")
.bind("{}")
.bind("auto")
.execute(&pool)
.await
.unwrap();
pool.close().await;
let sm = SessionManager::new(temp_dir.path().to_path_buf());
sm.storage().pool().await.unwrap(); // Triggers migration
let user_session = sm.storage().get_session("user_id", false).await.unwrap();
assert_eq!(user_session.session_type, SessionType::User);
let acp_session = sm.storage().get_session("acp_id", false).await.unwrap();
assert_eq!(acp_session.session_type, SessionType::Acp);
}
}
ui/desktop/openapi.json
+2 -1
View File
@@ -7904,7 +7904,8 @@
"sub_agent",
"hidden",
"terminal",
"gateway"
"gateway",
"acp"
]
},
"SessionsQuery": {
ui/desktop/src/api/types.gen.ts
+1 -1
View File
@@ -1276,7 +1276,7 @@ export type SessionReplyResponse = {
request_id: string;
};
export type SessionType = 'user' | 'scheduled' | 'sub_agent' | 'hidden' | 'terminal' | 'gateway';
export type SessionType = 'user' | 'scheduled' | 'sub_agent' | 'hidden' | 'terminal' | 'gateway' | 'acp';
export type SessionsQuery = {
limit: number;