ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(acp): introduce threads (#8344)

Browse Source 
Signed-off-by: Bradley Axen <baxen@squareup.com>
This commit is contained in:
Bradley Axen
2026-04-08 22:14:28 -07:00
committed by GitHub
parent 331d1e2efb
commit 9d7de43eec
23 changed files with 2207 additions and 880 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-acp/acp-meta.json
+20 -15
View File
@@ -25,26 +25,11 @@
"requestType": "UpdateWorkingDirRequest",
"responseType": "EmptyResponse"
},
{
"method": "session/get",
"requestType": "GetSessionRequest",
"responseType": "GetSessionResponse"
},
{
"method": "session/delete",
"requestType": "DeleteSessionRequest",
"responseType": "EmptyResponse"
},
{
"method": "_goose/session/export",
"requestType": "ExportSessionRequest",
"responseType": "ExportSessionResponse"
},
{
"method": "_goose/session/import",
"requestType": "ImportSessionRequest",
"responseType": "ImportSessionResponse"
},
{
"method": "_goose/config/extensions",
"requestType": "GetExtensionsRequest",
@@ -89,6 +74,26 @@
"method": "_goose/secret/remove",
"requestType": "RemoveSecretRequest",
"responseType": "EmptyResponse"
},
{
"method": "_goose/session/export",
"requestType": "ExportSessionRequest",
"responseType": "ExportSessionResponse"
},
{
"method": "_goose/session/import",
"requestType": "ImportSessionRequest",
"responseType": "ImportSessionResponse"
},
{
"method": "_goose/session/archive",
"requestType": "ArchiveSessionRequest",
"responseType": "EmptyResponse"
},
{
"method": "_goose/session/unarchive",
"requestType": "UnarchiveSessionRequest",
"responseType": "EmptyResponse"
}
]
}
crates/goose-acp/acp-schema.json
+154 -135
View File
@@ -125,36 +125,6 @@
"x-side": "agent",
"x-method": "_goose/working_dir/update"
},
"GetSessionRequest": {
"type": "object",
"properties": {
"sessionId": {
"type": "string"
},
"includeMessages": {
"type": "boolean",
"default": false
}
},
"required": [
"sessionId"
],
"description": "Get a session by ID.",
"x-side": "agent",
"x-method": "session/get"
},
"GetSessionResponse": {
"type": "object",
"properties": {
"session": {
"description": "The session object with id, name, working_dir, timestamps, tokens, etc.",
"default": null
}
},
"description": "Get a session response.",
"x-side": "agent",
"x-method": "session/get"
},
"DeleteSessionRequest": {
"type": "object",
"properties": {
@@ -169,60 +139,6 @@
"x-side": "agent",
"x-method": "session/delete"
},
"ExportSessionRequest": {
"type": "object",
"properties": {
"sessionId": {
"type": "string"
}
},
"required": [
"sessionId"
],
"description": "Export a session as a JSON string.",
"x-side": "agent",
"x-method": "_goose/session/export"
},
"ExportSessionResponse": {
"type": "object",
"properties": {
"data": {
"type": "string"
}
},
"required": [
"data"
],
"description": "Export session response.",
"x-side": "agent",
"x-method": "_goose/session/export"
},
"ImportSessionRequest": {
"type": "object",
"properties": {
"data": {
"type": "string"
}
},
"required": [
"data"
],
"description": "Import a session from a JSON string.",
"x-side": "agent",
"x-method": "_goose/session/import"
},
"ImportSessionResponse": {
"type": "object",
"properties": {
"session": {
"description": "The imported session object.",
"default": null
}
},
"description": "Import session response.",
"x-side": "agent",
"x-method": "_goose/session/import"
},
"GetExtensionsRequest": {
"type": "object",
"description": "List configured extensions and any warnings.",
@@ -458,6 +374,108 @@
"x-side": "agent",
"x-method": "_goose/secret/remove"
},
"ExportSessionRequest": {
"type": "object",
"properties": {
"sessionId": {
"type": "string"
}
},
"required": [
"sessionId"
],
"description": "Export a session as a JSON string.",
"x-side": "agent",
"x-method": "_goose/session/export"
},
"ExportSessionResponse": {
"type": "object",
"properties": {
"data": {
"type": "string"
}
},
"required": [
"data"
],
"description": "Export session response — raw JSON of the goose session with `conversation`.",
"x-side": "agent",
"x-method": "_goose/session/export"
},
"ImportSessionRequest": {
"type": "object",
"properties": {
"data": {
"type": "string"
}
},
"required": [
"data"
],
"description": "Import a session from a JSON string.",
"x-side": "agent",
"x-method": "_goose/session/import"
},
"ImportSessionResponse": {
"type": "object",
"properties": {
"sessionId": {
"type": "string"
},
"title": {
"type": [
"string",
"null"
]
},
"updatedAt": {
"type": [
"string",
"null"
]
},
"messageCount": {
"type": "integer",
"format": "uint64",
"minimum": 0
}
},
"required": [
"sessionId",
"messageCount"
],
"description": "Import session response — metadata about the newly created session.",
"x-side": "agent",
"x-method": "_goose/session/import"
},
"ArchiveSessionRequest": {
"type": "object",
"properties": {
"sessionId": {
"type": "string"
}
},
"required": [
"sessionId"
],
"description": "Archive a session (soft delete).",
"x-side": "agent",
"x-method": "_goose/session/archive"
},
"UnarchiveSessionRequest": {
"type": "object",
"properties": {
"sessionId": {
"type": "string"
}
},
"required": [
"sessionId"
],
"description": "Unarchive a previously archived session.",
"x-side": "agent",
"x-method": "_goose/session/unarchive"
},
"ExtRequest": {
"properties": {
"id": {
@@ -515,15 +533,6 @@
"description": "Params for _goose/working_dir/update",
"title": "UpdateWorkingDirRequest"
},
{
"allOf": [
{
"$ref": "#/$defs/GetSessionRequest"
}
],
"description": "Params for session/get",
"title": "GetSessionRequest"
},
{
"allOf": [
{
@@ -533,24 +542,6 @@
"description": "Params for session/delete",
"title": "DeleteSessionRequest"
},
{
"allOf": [
{
"$ref": "#/$defs/ExportSessionRequest"
}
],
"description": "Params for _goose/session/export",
"title": "ExportSessionRequest"
},
{
"allOf": [
{
"$ref": "#/$defs/ImportSessionRequest"
}
],
"description": "Params for _goose/session/import",
"title": "ImportSessionRequest"
},
{
"allOf": [
{
@@ -631,6 +622,42 @@
],
"description": "Params for _goose/secret/remove",
"title": "RemoveSecretRequest"
},
{
"allOf": [
{
"$ref": "#/$defs/ExportSessionRequest"
}
],
"description": "Params for _goose/session/export",
"title": "ExportSessionRequest"
},
{
"allOf": [
{
"$ref": "#/$defs/ImportSessionRequest"
}
],
"description": "Params for _goose/session/import",
"title": "ImportSessionRequest"
},
{
"allOf": [
{
"$ref": "#/$defs/ArchiveSessionRequest"
}
],
"description": "Params for _goose/session/archive",
"title": "ArchiveSessionRequest"
},
{
"allOf": [
{
"$ref": "#/$defs/UnarchiveSessionRequest"
}
],
"description": "Params for _goose/session/unarchive",
"title": "UnarchiveSessionRequest"
}
]
},
@@ -686,30 +713,6 @@
],
"title": "ReadResourceResponse"
},
{
"allOf": [
{
"$ref": "#/$defs/GetSessionResponse"
}
],
"title": "GetSessionResponse"
},
{
"allOf": [
{
"$ref": "#/$defs/ExportSessionResponse"
}
],
"title": "ExportSessionResponse"
},
{
"allOf": [
{
"$ref": "#/$defs/ImportSessionResponse"
}
],
"title": "ImportSessionResponse"
},
{
"allOf": [
{
@@ -749,6 +752,22 @@
}
],
"title": "CheckSecretResponse"
},
{
"allOf": [
{
"$ref": "#/$defs/ExportSessionResponse"
}
],
"title": "ExportSessionResponse"
},
{
"allOf": [
{
"$ref": "#/$defs/ImportSessionResponse"
}
],
"title": "ImportSessionResponse"
}
]
},
crates/goose-acp/src/adapters.rs
+12 -23
View File
@@ -6,7 +6,6 @@ use std::{
task::{Context, Poll},
};
use tokio::sync::mpsc;
use tracing::error;
/// Converts an mpsc::Receiver<String> to AsyncRead
/// Each message is terminated with a newline for JSON-RPC framing
@@ -61,15 +60,18 @@ impl tokio::io::AsyncRead for ReceiverToAsyncRead {
}
}
/// Converts an mpsc::Sender<String> to AsyncWrite
/// Splits incoming data on newlines for JSON-RPC framing
/// Converts an unbounded mpsc::Sender<String> to AsyncWrite.
/// Splits incoming data on newlines for JSON-RPC framing.
///
/// Uses an unbounded sender so that bursts of outgoing messages (e.g. replaying
/// a long session history) are never silently dropped due to backpressure.
pub(crate) struct SenderToAsyncWrite {
tx: mpsc::Sender<String>,
tx: mpsc::UnboundedSender<String>,
buffer: Vec<u8>,
}
impl SenderToAsyncWrite {
pub(crate) fn new(tx: mpsc::Sender<String>) -> Self {
pub(crate) fn new(tx: mpsc::UnboundedSender<String>) -> Self {
Self {
tx,
buffer: Vec::new(),
@@ -89,24 +91,11 @@ impl tokio::io::AsyncWrite for SenderToAsyncWrite {
let line = String::from_utf8_lossy(&self.buffer[..pos]).to_string();
self.buffer.drain(..=pos);
if !line.is_empty() {
if let Err(e) = self.tx.try_send(line.clone()) {
match e {
mpsc::error::TrySendError::Full(_) => {
let truncated: String = line.chars().take(100).collect();
error!(
"Channel full, dropping message (backpressure): {}",
truncated
);
}
mpsc::error::TrySendError::Closed(_) => {
return Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::BrokenPipe,
"Channel closed",
)));
}
}
}
if !line.is_empty() && self.tx.send(line).is_err() {
return Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::BrokenPipe,
"Channel closed",
)));
}
}
crates/goose-acp/src/server.rs
+935 -410
View File
File diff suppressed because it is too large Load Diff
crates/goose-acp/src/transport.rs
+1 -1
View File
@@ -26,7 +26,7 @@ pub(crate) const JSON_MIME_TYPE: &str = "application/json";
pub(crate) struct TransportSession {
pub to_agent_tx: mpsc::Sender<String>,
pub from_agent_rx: Arc<Mutex<mpsc::Receiver<String>>>,
pub from_agent_rx: Arc<Mutex<mpsc::UnboundedReceiver<String>>>,
pub handle: tokio::task::JoinHandle<()>,
}
crates/goose-acp/src/transport/http.rs
+3 -3
View File
@@ -32,7 +32,7 @@ impl HttpState {
async fn create_session(&self) -> Result<String, StatusCode> {
let (to_agent_tx, to_agent_rx) = mpsc::channel::<String>(256);
let (from_agent_tx, from_agent_rx) = mpsc::channel::<String>(256);
let (from_agent_tx, from_agent_rx) = mpsc::unbounded_channel::<String>();
let agent = self.server.create_agent().await.map_err(|e| {
error!("Failed to create agent: {}", e);
@@ -87,7 +87,7 @@ impl HttpState {
async fn get_receiver(
&self,
acp_session_id: &str,
) -> Result<Arc<Mutex<mpsc::Receiver<String>>>, StatusCode> {
) -> Result<Arc<Mutex<mpsc::UnboundedReceiver<String>>>, StatusCode> {
let sessions = self.sessions.read().await;
let session = sessions.get(acp_session_id).ok_or(StatusCode::NOT_FOUND)?;
Ok(session.from_agent_rx.clone())
@@ -95,7 +95,7 @@ impl HttpState {
}
fn create_sse_stream(
receiver: Arc<Mutex<mpsc::Receiver<String>>>,
receiver: Arc<Mutex<mpsc::UnboundedReceiver<String>>>,
cleanup: Option<(Arc<HttpState>, String)>,
) -> Sse<impl futures::Stream<Item = Result<axum::response::sse::Event, Infallible>>> {
let stream = async_stream::stream! {
crates/goose-acp/src/transport/websocket.rs
+1 -1
View File
@@ -30,7 +30,7 @@ impl WsState {
async fn create_connection(&self) -> Result<String> {
let (to_agent_tx, to_agent_rx) = mpsc::channel::<String>(256);
let (from_agent_tx, from_agent_rx) = mpsc::channel::<String>(256);
let (from_agent_tx, from_agent_rx) = mpsc::unbounded_channel::<String>();
let agent = self.server.create_agent().await?;
crates/goose-acp/tests/common_tests/mod.rs
+7 -1
View File
@@ -58,13 +58,19 @@ pub async fn run_list_sessions<C: Connection>() {
for s in &mut response.sessions {
s.updated_at = None;
}
let mut expected_meta = serde_json::Map::new();
expected_meta.insert(
"messageCount".to_string(),
serde_json::Value::Number(2.into()),
);
assert_eq!(
response,
ListSessionsResponse::new(vec![SessionInfo::new(
session.session_id().clone(),
session.work_dir()
)
.title("ACP Session".to_string())])
.title("New Chat".to_string())
.meta(expected_meta)])
);
}
crates/goose-acp/tests/custom_requests_test.rs
+55 -92
View File
@@ -3,14 +3,15 @@ mod common_tests;
use common_tests::fixtures::server::AcpServerConnection;
use common_tests::fixtures::{
run_test, send_custom, Connection, Session, SessionData, TestConnectionConfig,
run_test, send_custom, Connection, PermissionDecision, Session, SessionData,
TestConnectionConfig,
};
use goose::model::ModelConfig;
use goose::providers::base::{MessageStream, Provider};
use goose::providers::errors::ProviderError;
use goose_acp::server::AcpProviderFactory;
use goose_test_support::EnforceSessionId;
use std::sync::Arc;
use goose_test_support::{EnforceSessionId, IgnoreSessionId};
use std::sync::{Arc, Mutex};
use common_tests::fixtures::OpenAiFixture;
@@ -65,34 +66,6 @@ fn mock_provider_factory() -> AcpProviderFactory {
})
}
#[test]
fn test_custom_session_get() {
run_test(async {
let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await;
let mut conn = AcpServerConnection::new(TestConnectionConfig::default(), openai).await;
let SessionData { session, .. } = conn.new_session().await.unwrap();
let session_id = session.session_id().0.clone();
let result = send_custom(
conn.cx(),
"session/get",
serde_json::json!({
"sessionId": session_id,
}),
)
.await;
assert!(result.is_ok(), "expected ok, got: {:?}", result);
let response = result.unwrap();
let returned_session = response.get("session").expect("missing 'session' field");
assert_eq!(
returned_session.get("id").and_then(|v| v.as_str()),
Some(session_id.as_ref())
);
});
}
#[test]
fn test_custom_get_tools() {
run_test(async {
@@ -237,27 +210,6 @@ fn test_provider_switching_updates_session_state() {
.await
.expect("provider config option should succeed");
let response = send_custom(
conn.cx(),
"session/get",
serde_json::json!({
"sessionId": session_id,
}),
)
.await
.expect("session/get should succeed");
let session_value = response.get("session").expect("missing session");
assert_eq!(
session_value.get("provider_name"),
Some(&serde_json::json!("anthropic"))
);
assert_eq!(
session_value
.get("model_config")
.and_then(|value| value.get("model_name")),
Some(&serde_json::json!("current"))
);
let response = send_custom(
conn.cx(),
"_goose/session/provider/update",
@@ -278,27 +230,6 @@ fn test_provider_switching_updates_session_state() {
"expected refreshed config options"
);
let response = send_custom(
conn.cx(),
"session/get",
serde_json::json!({
"sessionId": session_id,
}),
)
.await
.expect("session/get after provider update should succeed");
let session_value = response.get("session").expect("missing session");
assert_eq!(
session_value.get("provider_name"),
Some(&serde_json::json!("openai"))
);
assert_eq!(
session_value
.get("model_config")
.and_then(|value| value.get("model_name")),
Some(&serde_json::json!("o4-mini"))
);
let response = send_custom(
conn.cx(),
"_goose/session/provider/update",
@@ -319,25 +250,6 @@ fn test_provider_switching_updates_session_state() {
.any(|option| option.get("id") == Some(&serde_json::json!("provider"))),
"missing provider config option after reset"
);
let response = send_custom(
conn.cx(),
"session/get",
serde_json::json!({
"sessionId": session_id,
}),
)
.await
.expect("session/get after provider reset should succeed");
let session_value = response.get("session").expect("missing session");
assert_eq!(
session_value.get("provider_name"),
Some(&serde_json::json!("goose"))
);
assert_eq!(
session_value.get("model_config"),
Some(&serde_json::Value::Null)
);
});
}
@@ -351,3 +263,54 @@ fn test_custom_unknown_method() {
assert!(result.is_err(), "expected method_not_found error");
});
}
#[test]
fn test_developer_fs_requests_use_acp_session_id() {
run_test(async {
let seen_session_id = Arc::new(Mutex::new(None::<String>));
let seen_session_id_clone = Arc::clone(&seen_session_id);
let openai = OpenAiFixture::new(
vec![
(
"Use the read tool to read /tmp/test_acp_read.txt and output only its contents."
.to_string(),
include_str!("test_data/openai_fs_read_tool_call.txt"),
),
(
r#""content":"test-read-content-12345""#.into(),
include_str!("test_data/openai_fs_read_tool_result.txt"),
),
],
Arc::new(IgnoreSessionId),
)
.await;
let config = TestConnectionConfig {
read_text_file: Some(Arc::new(move |req| {
*seen_session_id_clone.lock().unwrap() = Some(req.session_id.0.to_string());
Ok(sacp::schema::ReadTextFileResponse::new(
"test-read-content-12345",
))
})),
..Default::default()
};
let mut conn = AcpServerConnection::new(config, openai).await;
let SessionData { mut session, .. } = conn.new_session().await.unwrap();
let acp_session_id = session.session_id().0.to_string();
let output = session
.prompt(
"Use the read tool to read /tmp/test_acp_read.txt and output only its contents.",
PermissionDecision::Cancel,
)
.await
.expect("prompt should succeed");
assert_eq!(output.text, "test-read-content-12345");
assert_eq!(
seen_session_id.lock().unwrap().as_deref(),
Some(acp_session_id.as_str()),
"ACP read request should use the ACP session/thread ID",
);
});
}
crates/goose-acp/tests/fixtures/server.rs
Vendored
+5 -2
View File
@@ -4,7 +4,7 @@ use super::{
};
use async_trait::async_trait;
use goose::config::PermissionManager;
use goose_test_support::{EnforceSessionId, ExpectedSessionId};
use goose_test_support::{ExpectedSessionId, IgnoreSessionId};
use sacp::schema::{
ClientCapabilities, CloseSessionRequest, ContentBlock, CreateTerminalRequest,
FileSystemCapabilities, ImageContent, InitializeRequest, KillTerminalRequest,
@@ -99,7 +99,10 @@ impl Connection for AcpServerConnection {
type Session = AcpServerSession;
fn expected_session_id() -> Arc<dyn ExpectedSessionId> {
Arc::new(EnforceSessionId::default())
// The ACP session ID returned to clients is now a thread ID, which is
// intentionally different from the internal session ID the agent sends
// to the LLM provider. Skip strict matching.
Arc::new(IgnoreSessionId)
}
async fn new(config: TestConnectionConfig, openai: super::OpenAiFixture) -> Self {
crates/goose-sdk/src/custom_requests.rs
+47 -47
View File
@@ -83,24 +83,6 @@ pub struct UpdateWorkingDirRequest {
pub working_dir: String,
}
/// Get a session by ID.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(method = "session/get", response = GetSessionResponse)]
#[serde(rename_all = "camelCase")]
pub struct GetSessionRequest {
pub session_id: String,
#[serde(default)]
pub include_messages: bool,
}
/// Get a session response.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcResponse)]
pub struct GetSessionResponse {
/// The session object with id, name, working_dir, timestamps, tokens, etc.
#[serde(default)]
pub session: serde_json::Value,
}
/// Delete a session.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(method = "session/delete", response = EmptyResponse)]
@@ -109,35 +91,6 @@ pub struct DeleteSessionRequest {
pub session_id: String,
}
/// Export a session as a JSON string.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(method = "_goose/session/export", response = ExportSessionResponse)]
#[serde(rename_all = "camelCase")]
pub struct ExportSessionRequest {
pub session_id: String,
}
/// Export session response.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcResponse)]
pub struct ExportSessionResponse {
pub data: String,
}
/// Import a session from a JSON string.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(method = "_goose/session/import", response = ImportSessionResponse)]
pub struct ImportSessionRequest {
pub data: String,
}
/// Import session response.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcResponse)]
pub struct ImportSessionResponse {
/// The imported session object.
#[serde(default)]
pub session: serde_json::Value,
}
/// List configured extensions and any warnings.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(method = "_goose/config/extensions", response = GetExtensionsResponse)]
@@ -254,6 +207,53 @@ pub struct ListProvidersResponse {
pub providers: Vec<ProviderListEntry>,
}
/// Archive a session (soft delete).
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(method = "_goose/session/archive", response = EmptyResponse)]
#[serde(rename_all = "camelCase")]
pub struct ArchiveSessionRequest {
pub session_id: String,
}
/// Unarchive a previously archived session.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(method = "_goose/session/unarchive", response = EmptyResponse)]
#[serde(rename_all = "camelCase")]
pub struct UnarchiveSessionRequest {
pub session_id: String,
}
/// Export a session as a JSON string.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(method = "_goose/session/export", response = ExportSessionResponse)]
#[serde(rename_all = "camelCase")]
pub struct ExportSessionRequest {
pub session_id: String,
}
/// Export session response — raw JSON of the goose session with `conversation`.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcResponse)]
pub struct ExportSessionResponse {
pub data: String,
}
/// Import a session from a JSON string.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(method = "_goose/session/import", response = ImportSessionResponse)]
pub struct ImportSessionRequest {
pub data: String,
}
/// Import session response — metadata about the newly created session.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcResponse)]
#[serde(rename_all = "camelCase")]
pub struct ImportSessionResponse {
pub session_id: String,
pub title: Option<String>,
pub updated_at: Option<String>,
pub message_count: u64,
}
/// Empty success response for operations that return no data.
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcResponse)]
pub struct EmptyResponse {}
crates/goose/src/providers/canonical/mod.rs
+44 -1
View File
@@ -3,7 +3,9 @@ mod name_builder;
mod registry;
pub use model::{CanonicalModel, Limit, Modalities, Modality, Pricing};
pub use name_builder::{canonical_name, map_to_canonical_model, strip_version_suffix};
pub use name_builder::{
canonical_name, map_provider_name, map_to_canonical_model, strip_version_suffix,
};
pub use registry::CanonicalModelRegistry;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
@@ -21,6 +23,47 @@ impl ModelMapping {
}
}
/// Return recommended model names for a provider using only the bundled canonical registry.
///
/// This avoids network calls by looking up all known models for the provider,
/// filtering to text-input + tool-calling models, and sorting by release date.
/// The returned names are the canonical short names (e.g. "claude-3.5-sonnet").
///
/// TODO: This trades speed for correctness — the canonical registry may not perfectly
/// match what the provider API returns (new models not yet in the registry, deprecated
/// models still listed, or locally-installed models for providers like Ollama). Consider
/// whether to reconcile with a live API call in the background.
pub fn recommended_models_from_registry(provider: &str) -> Vec<String> {
let registry = match CanonicalModelRegistry::bundled() {
Ok(r) => r,
Err(_) => return vec![],
};
let registry_provider = map_provider_name(provider);
let all = registry.get_all_models_for_provider(registry_provider);
let mut models_with_dates: Vec<(String, Option<String>)> = all
.iter()
.filter(|m| m.modalities.input.contains(&Modality::Text) && m.tool_call)
.filter_map(|m| {
let (_, name) = m.id.split_once('/')?;
Some((name.to_string(), m.release_date.clone()))
})
.collect();
models_with_dates.sort_by(|a, b| match (&a.1, &b.1) {
(Some(date_a), Some(date_b)) => date_b.cmp(date_a),
(Some(_), None) => std::cmp::Ordering::Less,
(None, Some(_)) => std::cmp::Ordering::Greater,
(None, None) => a.0.cmp(&b.0),
});
models_with_dates
.into_iter()
.map(|(name, _)| name)
.collect()
}
pub fn maybe_get_canonical_model(provider: &str, model: &str) -> Option<CanonicalModel> {
let registry = CanonicalModelRegistry::bundled().ok()?;
crates/goose/src/providers/canonical/name_builder.rs
+1 -1
View File
@@ -39,7 +39,7 @@ fn is_meta_provider(provider: &str) -> bool {
matches!(provider, "databricks" | "tetrate" | "bedrock" | "azure")
}
fn map_provider_name(provider: &str) -> &str {
pub fn map_provider_name(provider: &str) -> &str {
match provider {
// Goose provider names that differ from models.dev names
"xai" => "x-ai",
crates/goose/src/session/mod.rs
+2
View File
@@ -3,6 +3,7 @@ mod diagnostics;
pub mod extension_data;
mod legacy;
pub mod session_manager;
pub mod thread_manager;
pub use diagnostics::{
config_path, generate_diagnostics, get_system_info, latest_llm_log_path,
@@ -12,3 +13,4 @@ pub use extension_data::{EnabledExtensionsState, ExtensionData, ExtensionState,
pub use session_manager::{
Session, SessionInsights, SessionManager, SessionType, SessionUpdateBuilder,
};
pub use thread_manager::{Thread, ThreadManager, ThreadMetadata};
crates/goose/src/session/session_manager.rs
+134 -8
View File
@@ -19,7 +19,7 @@ use std::sync::{Arc, LazyLock};
use tracing::{info, warn};
use utoipa::ToSchema;
pub const CURRENT_SCHEMA_VERSION: i32 = 9;
pub const CURRENT_SCHEMA_VERSION: i32 = 10;
pub const SESSIONS_FOLDER: &str = "sessions";
pub const DB_NAME: &str = "sessions.db";
@@ -81,6 +81,8 @@ pub struct Session {
pub model_config: Option<ModelConfig>,
#[serde(default)]
pub goose_mode: GooseMode,
#[serde(default)]
pub thread_id: Option<String>,
}
pub struct SessionUpdateBuilder<'a> {
@@ -103,6 +105,7 @@ pub struct SessionUpdateBuilder<'a> {
provider_name: Option<Option<String>>,
model_config: Option<Option<ModelConfig>>,
goose_mode: Option<GooseMode>,
thread_id: Option<Option<String>>,
}
#[derive(Serialize, ToSchema, Debug)]
@@ -134,6 +137,7 @@ impl<'a> SessionUpdateBuilder<'a> {
provider_name: None,
model_config: None,
goose_mode: None,
thread_id: None,
}
}
@@ -241,6 +245,11 @@ impl<'a> SessionUpdateBuilder<'a> {
self.goose_mode = Some(mode);
self
}
pub fn thread_id(mut self, thread_id: Option<String>) -> Self {
self.thread_id = Some(thread_id);
self
}
}
pub struct SessionManager {
@@ -361,7 +370,22 @@ impl SessionManager {
if user_message_count <= MSG_COUNT_FOR_SESSION_NAME_GENERATION {
let name = provider.generate_session_name(id, &conversation).await?;
self.update(id).system_generated_name(name).apply().await
self.update(id)
.system_generated_name(name.clone())
.apply()
.await?;
// Also update the thread name so ACP clients see it via session/list.
if let Some(ref thread_id) = session.thread_id {
let thread_mgr = super::thread_manager::ThreadManager::new(self.storage.clone());
let thread = thread_mgr.get_thread(thread_id).await?;
if !thread.user_set_name {
thread_mgr
.update_thread(thread_id, Some(name), Some(false), None)
.await?;
}
}
Ok(())
} else {
Ok(())
}
@@ -407,7 +431,7 @@ pub struct SessionStorage {
session_dir: PathBuf,
}
fn role_to_string(role: &Role) -> &'static str {
pub(crate) fn role_to_string(role: &Role) -> &'static str {
match role {
Role::User => "user",
Role::Assistant => "assistant",
@@ -439,6 +463,7 @@ impl Default for Session {
provider_name: None,
model_config: None,
goose_mode: GooseMode::default(),
thread_id: None,
}
}
}
@@ -508,6 +533,7 @@ impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session {
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or_default(),
thread_id: row.try_get("thread_id").ok().flatten(),
})
}
}
@@ -537,7 +563,7 @@ impl SessionStorage {
}
}
async fn pool(&self) -> Result<&Pool<Sqlite>> {
pub(crate) async fn pool(&self) -> Result<&Pool<Sqlite>> {
self.initialized
.get_or_try_init(|| async {
let schema_exists = sqlx::query_scalar::<_, bool>(
@@ -607,7 +633,8 @@ impl SessionStorage {
user_recipe_values_json TEXT,
provider_name TEXT,
model_config_json TEXT,
goose_mode TEXT NOT NULL DEFAULT 'auto'
goose_mode TEXT NOT NULL DEFAULT 'auto',
thread_id TEXT
)
"#,
)
@@ -647,6 +674,48 @@ impl SessionStorage {
sqlx::query("CREATE INDEX idx_sessions_type ON sessions(session_type)")
.execute(pool)
.await?;
sqlx::query("CREATE INDEX IF NOT EXISTS idx_sessions_thread ON sessions(thread_id)")
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS threads (
id TEXT PRIMARY KEY,
name TEXT NOT NULL DEFAULT 'New Chat',
user_set_name BOOLEAN DEFAULT FALSE,
working_dir TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
archived_at TIMESTAMP,
metadata_json TEXT DEFAULT '{}'
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS thread_messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
thread_id TEXT NOT NULL REFERENCES threads(id),
session_id TEXT,
message_id TEXT,
role TEXT NOT NULL,
content_json TEXT NOT NULL,
created_timestamp INTEGER NOT NULL,
metadata_json TEXT DEFAULT '{}'
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_thread_messages_thread ON thread_messages(thread_id)",
)
.execute(pool)
.await?;
sqlx::query("CREATE INDEX IF NOT EXISTS idx_thread_messages_message_id ON thread_messages(message_id)")
.execute(pool)
.await?;
Ok(())
}
@@ -938,6 +1007,59 @@ impl SessionStorage {
.execute(&mut **tx)
.await?;
}
10 => {
// Check if thread_id column already exists (e.g. fresh schema)
let has_thread_id = sqlx::query_scalar::<_, i32>(
"SELECT COUNT(*) FROM pragma_table_info('sessions') WHERE name = 'thread_id'",
)
.fetch_one(&mut **tx)
.await?
> 0;
if !has_thread_id {
sqlx::query("ALTER TABLE sessions ADD COLUMN thread_id TEXT")
.execute(&mut **tx)
.await?;
}
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_sessions_thread ON sessions(thread_id)",
)
.execute(&mut **tx)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS threads (
id TEXT PRIMARY KEY,
name TEXT NOT NULL DEFAULT 'New Chat',
user_set_name BOOLEAN DEFAULT FALSE,
working_dir TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
archived_at TIMESTAMP,
metadata_json TEXT DEFAULT '{}'
)",
)
.execute(&mut **tx)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS thread_messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
thread_id TEXT NOT NULL REFERENCES threads(id),
session_id TEXT,
message_id TEXT,
role TEXT NOT NULL,
content_json TEXT NOT NULL,
created_timestamp INTEGER NOT NULL,
metadata_json TEXT DEFAULT '{}'
)",
)
.execute(&mut **tx)
.await?;
sqlx::query("CREATE INDEX IF NOT EXISTS idx_thread_messages_thread ON thread_messages(thread_id)")
.execute(&mut **tx)
.await?;
sqlx::query("CREATE INDEX IF NOT EXISTS idx_thread_messages_message_id ON thread_messages(message_id)")
.execute(&mut **tx)
.await?;
}
_ => {
anyhow::bail!("Unknown migration version: {}", version);
}
@@ -999,7 +1121,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, goose_mode
provider_name, model_config_json, goose_mode, thread_id
FROM sessions
WHERE id = ?
"#,
@@ -1063,6 +1185,7 @@ impl SessionStorage {
add_update!(builder.provider_name, "provider_name");
add_update!(builder.model_config, "model_config_json");
add_update!(builder.goose_mode, "goose_mode");
add_update!(builder.thread_id, "thread_id");
if updates.is_empty() {
return Ok(());
@@ -1131,6 +1254,9 @@ impl SessionStorage {
if let Some(goose_mode) = builder.goose_mode {
q = q.bind(goose_mode.to_string());
}
if let Some(thread_id) = builder.thread_id {
q = q.bind(thread_id);
}
let pool = self.pool().await?;
let mut tx = pool.begin_with("BEGIN IMMEDIATE").await?;
@@ -1282,10 +1408,10 @@ 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.goose_mode,
s.provider_name, s.model_config_json, s.goose_mode, s.thread_id,
COUNT(m.id) as message_count
FROM sessions s
INNER JOIN messages m ON s.id = m.session_id
LEFT JOIN messages m ON s.id = m.session_id
{}
GROUP BY s.id
ORDER BY s.updated_at DESC
crates/goose/src/session/thread_manager.rs
+433
View File
@@ -0,0 +1,433 @@
use super::session_manager::{role_to_string, SessionStorage};
use crate::conversation::message::Message;
use anyhow::Result;
use chrono::{DateTime, Utc};
use rmcp::model::Role;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Thread {
pub id: String,
pub name: String,
pub user_set_name: bool,
pub working_dir: Option<String>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
pub archived_at: Option<DateTime<Utc>>,
pub metadata: ThreadMetadata,
#[serde(default)]
pub current_session_id: Option<String>,
#[serde(default)]
pub message_count: i64,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ThreadMetadata {
#[serde(default)]
pub persona_id: Option<String>,
#[serde(default)]
pub project_id: Option<String>,
#[serde(default)]
pub provider_id: Option<String>,
#[serde(default)]
pub model_name: Option<String>,
#[serde(default)]
pub mode: Option<String>,
#[serde(flatten)]
pub extra: std::collections::HashMap<String, serde_json::Value>,
}
pub struct ThreadManager {
storage: Arc<SessionStorage>,
}
const THREAD_SELECT: &str = "\
SELECT t.id, t.name, t.user_set_name, t.working_dir, t.created_at, t.updated_at, \
t.archived_at, t.metadata_json, \
(SELECT s.id FROM sessions s WHERE s.thread_id = t.id ORDER BY s.created_at DESC LIMIT 1) as current_session_id, \
(SELECT COUNT(*) FROM thread_messages WHERE thread_id = t.id) as message_count \
FROM threads t";
type ThreadRow = (
String,
String,
bool,
Option<String>,
String,
String,
Option<String>,
String,
Option<String>,
i64,
);
fn thread_from_row(
(
id,
name,
user_set_name,
working_dir,
created_at,
updated_at,
archived_at_str,
metadata_json,
current_session_id,
message_count,
): ThreadRow,
) -> Result<Thread> {
let metadata: ThreadMetadata = serde_json::from_str(&metadata_json).unwrap_or_default();
let archived_at = archived_at_str.as_deref().and_then(|s| s.parse().ok());
Ok(Thread {
id,
name,
user_set_name,
working_dir,
created_at: created_at.parse().unwrap_or_else(|_| Utc::now()),
updated_at: updated_at.parse().unwrap_or_else(|_| Utc::now()),
archived_at,
metadata,
current_session_id,
message_count,
})
}
impl ThreadManager {
pub fn new(storage: Arc<SessionStorage>) -> Self {
Self { storage }
}
pub async fn create_thread(
&self,
name: Option<String>,
metadata: Option<ThreadMetadata>,
working_dir: Option<String>,
) -> Result<Thread> {
let pool = self.storage.pool().await?;
let id = uuid::Uuid::new_v4().to_string();
let name = name.unwrap_or_else(|| "New Chat".to_string());
let meta = metadata.unwrap_or_default();
let metadata_json = serde_json::to_string(&meta)?;
sqlx::query(
"INSERT INTO threads (id, name, user_set_name, working_dir, metadata_json) VALUES (?, ?, FALSE, ?, ?)",
)
.bind(&id)
.bind(&name)
.bind(&working_dir)
.bind(&metadata_json)
.execute(pool)
.await?;
self.get_thread(&id).await
}
pub async fn get_thread(&self, id: &str) -> Result<Thread> {
let pool = self.storage.pool().await?;
let sql = format!("{} WHERE t.id = ?", THREAD_SELECT);
let row = sqlx::query_as::<_, ThreadRow>(&sql)
.bind(id)
.fetch_one(pool)
.await?;
thread_from_row(row)
}
pub async fn update_thread(
&self,
id: &str,
name: Option<String>,
user_set_name: Option<bool>,
metadata: Option<ThreadMetadata>,
) -> Result<Thread> {
let pool = self.storage.pool().await?;
let mut sets = Vec::new();
if name.is_some() {
sets.push("name = ?");
sets.push("user_set_name = ?");
}
if metadata.is_some() {
sets.push("metadata_json = ?");
}
if !sets.is_empty() {
let sql = format!(
"UPDATE threads SET {}, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
sets.join(", ")
);
let mut q = sqlx::query(&sql);
if let Some(ref n) = name {
q = q.bind(n);
q = q.bind(user_set_name.unwrap_or(true));
}
if let Some(ref meta) = metadata {
q = q.bind(serde_json::to_string(meta)?);
}
q = q.bind(id);
q.execute(pool).await?;
}
self.get_thread(id).await
}
pub async fn list_threads(&self, include_archived: bool) -> Result<Vec<Thread>> {
let pool = self.storage.pool().await?;
let sql = if include_archived {
format!("{} ORDER BY t.updated_at DESC", THREAD_SELECT)
} else {
format!(
"{} WHERE t.archived_at IS NULL ORDER BY t.updated_at DESC",
THREAD_SELECT
)
};
let rows = sqlx::query_as::<_, ThreadRow>(&sql).fetch_all(pool).await?;
rows.into_iter().map(thread_from_row).collect()
}
pub async fn archive_thread(&self, id: &str) -> Result<Thread> {
let pool = self.storage.pool().await?;
sqlx::query("UPDATE threads SET archived_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP WHERE id = ?")
.bind(id)
.execute(pool)
.await?;
self.get_thread(id).await
}
pub async fn unarchive_thread(&self, id: &str) -> Result<Thread> {
let pool = self.storage.pool().await?;
sqlx::query(
"UPDATE threads SET archived_at = NULL, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
)
.bind(id)
.execute(pool)
.await?;
self.get_thread(id).await
}
pub async fn update_metadata(
&self,
id: &str,
f: impl FnOnce(&mut ThreadMetadata),
) -> Result<Thread> {
let thread = self.get_thread(id).await?;
let mut meta = thread.metadata;
f(&mut meta);
self.update_thread(id, None, None, Some(meta)).await
}
pub async fn update_working_dir(&self, id: &str, working_dir: &str) -> Result<()> {
let pool = self.storage.pool().await?;
sqlx::query(
"UPDATE threads SET working_dir = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
)
.bind(working_dir)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn delete_thread(&self, id: &str) -> Result<()> {
let pool = self.storage.pool().await?;
let mut tx = pool.begin().await?;
sqlx::query("DELETE FROM thread_messages WHERE thread_id = ?")
.bind(id)
.execute(&mut *tx)
.await?;
sqlx::query(
"DELETE FROM messages WHERE session_id IN (SELECT id FROM sessions WHERE thread_id = ?)",
)
.bind(id)
.execute(&mut *tx)
.await?;
sqlx::query("DELETE FROM sessions WHERE thread_id = ?")
.bind(id)
.execute(&mut *tx)
.await?;
sqlx::query("DELETE FROM threads WHERE id = ?")
.bind(id)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
pub async fn append_message(
&self,
thread_id: &str,
session_id: Option<&str>,
message: &Message,
) -> Result<Message> {
let pool = self.storage.pool().await?;
let role_str = role_to_string(&message.role);
let metadata_json = serde_json::to_string(&message.metadata)?;
// When the incoming message is text-only, try to coalesce it with the
// last stored row if that row has the same role and is also text-only.
// This avoids storing one row per streaming token while keeping the UI
// streaming path unchanged (callers still forward every chunk).
if message.has_only_text_content() && !message.content.is_empty() {
let new_text = message.as_concat_text();
let maybe_last = sqlx::query_as::<_, (i64, String, String, String, String)>(
"SELECT id, message_id, role, content_json, metadata_json \
FROM thread_messages \
WHERE thread_id = ? \
ORDER BY id DESC LIMIT 1",
)
.bind(thread_id)
.fetch_optional(pool)
.await?;
if let Some((
row_id,
existing_msg_id,
last_role,
last_content_json,
last_metadata_json,
)) = maybe_last
{
if last_role == role_str
&& last_metadata_json == metadata_json
&& is_text_only_json(&last_content_json)
{
// Append text into the existing row's single text element.
let updated_json = append_text_json(&last_content_json, &new_text)?;
sqlx::query("UPDATE thread_messages SET content_json = ? WHERE id = ?")
.bind(&updated_json)
.bind(row_id)
.execute(pool)
.await?;
sqlx::query("UPDATE threads SET updated_at = CURRENT_TIMESTAMP WHERE id = ?")
.bind(thread_id)
.execute(pool)
.await?;
let mut stored = message.clone();
stored.id = Some(existing_msg_id);
return Ok(stored);
}
}
}
// Default path: insert a new row.
let content_json = serde_json::to_string(&message.content)?;
let message_id = message
.id
.clone()
.unwrap_or_else(|| format!("tmsg_{}", uuid::Uuid::new_v4()));
sqlx::query(
"INSERT INTO thread_messages (thread_id, session_id, message_id, role, content_json, created_timestamp, metadata_json) VALUES (?, ?, ?, ?, ?, ?, ?)",
)
.bind(thread_id)
.bind(session_id)
.bind(&message_id)
.bind(role_str)
.bind(&content_json)
.bind(message.created)
.bind(&metadata_json)
.execute(pool)
.await?;
sqlx::query("UPDATE threads SET updated_at = CURRENT_TIMESTAMP WHERE id = ?")
.bind(thread_id)
.execute(pool)
.await?;
let mut stored = message.clone();
stored.id = Some(message_id);
Ok(stored)
}
pub async fn fork_thread(&self, source_thread_id: &str) -> Result<Thread> {
let source = self.get_thread(source_thread_id).await?;
let pool = self.storage.pool().await?;
let new_id = uuid::Uuid::new_v4().to_string();
let name = format!("Fork of {}", source.name);
let metadata_json = serde_json::to_string(&source.metadata)?;
sqlx::query(
"INSERT INTO threads (id, name, user_set_name, working_dir, metadata_json) VALUES (?, ?, FALSE, ?, ?)",
)
.bind(&new_id)
.bind(&name)
.bind(&source.working_dir)
.bind(&metadata_json)
.execute(pool)
.await?;
// Copy all thread messages
sqlx::query(
"INSERT INTO thread_messages (thread_id, session_id, message_id, role, content_json, created_timestamp, metadata_json) \
SELECT ?, session_id, 'tmsg_' || hex(randomblob(16)), role, content_json, created_timestamp, metadata_json \
FROM thread_messages WHERE thread_id = ? ORDER BY id ASC",
)
.bind(&new_id)
.bind(source_thread_id)
.execute(pool)
.await?;
self.get_thread(&new_id).await
}
pub async fn list_messages(&self, thread_id: &str) -> Result<Vec<Message>> {
let pool = self.storage.pool().await?;
let rows = sqlx::query_as::<_, (Option<String>, String, Option<String>, String, i64, String)>(
"SELECT message_id, role, session_id, content_json, created_timestamp, metadata_json FROM thread_messages WHERE thread_id = ? ORDER BY id ASC",
)
.bind(thread_id)
.fetch_all(pool)
.await?;
let mut messages = Vec::new();
for (message_id, role_str, _session_id, content_json, created_timestamp, metadata_json) in
rows
{
let role = match role_str.as_str() {
"user" => Role::User,
"assistant" => Role::Assistant,
_ => continue,
};
let content = serde_json::from_str(&content_json)?;
let metadata = serde_json::from_str(&metadata_json).unwrap_or_default();
let mut msg = Message::new(role, created_timestamp, content);
msg.metadata = metadata;
if let Some(id) = message_id {
msg = msg.with_id(id);
}
messages.push(msg);
}
Ok(messages)
}
}
/// Check whether a `content_json` string represents a single text-only element.
/// Avoids a full deserialize by inspecting the JSON structure directly.
fn is_text_only_json(content_json: &str) -> bool {
let Ok(items) = serde_json::from_str::<Vec<serde_json::Value>>(content_json) else {
return false;
};
items.len() == 1
&& items[0].get("type").and_then(|v| v.as_str()) == Some("text")
&& items[0].get("text").is_some()
}
/// Append `new_text` to the single text element in a text-only `content_json` array.
fn append_text_json(content_json: &str, new_text: &str) -> anyhow::Result<String> {
let mut items: Vec<serde_json::Value> = serde_json::from_str(content_json)?;
if let Some(text_val) = items.get_mut(0).and_then(|v| v.get_mut("text")) {
let existing = text_val.as_str().unwrap_or("");
*text_val = serde_json::Value::String(format!("{}{}", existing, new_text));
}
Ok(serde_json::to_string(&items)?)
}
crates/goose/tests/thread_message_coalescing_test.rs
+197
View File
@@ -0,0 +1,197 @@
use goose::conversation::message::Message;
use goose::session::session_manager::SessionStorage;
use goose::session::thread_manager::ThreadManager;
use rmcp::model::CallToolRequestParams;
use std::sync::Arc;
use tempfile::TempDir;
async fn setup() -> (ThreadManager, TempDir) {
let tmp = TempDir::new().unwrap();
let storage = SessionStorage::create(tmp.path()).await.unwrap();
let tm = ThreadManager::new(Arc::new(storage));
(tm, tmp)
}
#[tokio::test]
async fn consecutive_text_chunks_are_coalesced() {
let (tm, _tmp) = setup().await;
let thread = tm.create_thread(None, None, None).await.unwrap();
// Simulate streaming: three consecutive assistant text chunks.
tm.append_message(
&thread.id,
Some("s1"),
&Message::assistant().with_text("Hello"),
)
.await
.unwrap();
tm.append_message(
&thread.id,
Some("s1"),
&Message::assistant().with_text(" world"),
)
.await
.unwrap();
tm.append_message(&thread.id, Some("s1"), &Message::assistant().with_text("!"))
.await
.unwrap();
let messages = tm.list_messages(&thread.id).await.unwrap();
assert_eq!(messages.len(), 1, "should coalesce into a single row");
assert_eq!(messages[0].as_concat_text(), "Hello world!");
}
#[tokio::test]
async fn role_change_prevents_coalescing() {
let (tm, _tmp) = setup().await;
let thread = tm.create_thread(None, None, None).await.unwrap();
tm.append_message(&thread.id, Some("s1"), &Message::user().with_text("Hi"))
.await
.unwrap();
tm.append_message(
&thread.id,
Some("s1"),
&Message::assistant().with_text("Hey"),
)
.await
.unwrap();
let messages = tm.list_messages(&thread.id).await.unwrap();
assert_eq!(messages.len(), 2, "different roles should not coalesce");
assert_eq!(messages[0].as_concat_text(), "Hi");
assert_eq!(messages[1].as_concat_text(), "Hey");
}
#[tokio::test]
async fn non_text_content_breaks_coalescing() {
let (tm, _tmp) = setup().await;
let thread = tm.create_thread(None, None, None).await.unwrap();
// Text, then tool request, then more text — should be 3 rows.
tm.append_message(
&thread.id,
Some("s1"),
&Message::assistant().with_text("Let me check"),
)
.await
.unwrap();
let tool_msg = Message::assistant().with_tool_request(
"call_1",
Ok(CallToolRequestParams::new("shell").with_arguments(
serde_json::json!({"command": "ls"})
.as_object()
.unwrap()
.clone(),
)),
);
tm.append_message(&thread.id, Some("s1"), &tool_msg)
.await
.unwrap();
tm.append_message(
&thread.id,
Some("s1"),
&Message::assistant().with_text("Done"),
)
.await
.unwrap();
let messages = tm.list_messages(&thread.id).await.unwrap();
assert_eq!(messages.len(), 3, "tool request should break coalescing");
assert_eq!(messages[0].as_concat_text(), "Let me check");
assert_eq!(messages[2].as_concat_text(), "Done");
}
#[tokio::test]
async fn text_after_tool_response_not_coalesced_with_tool() {
let (tm, _tmp) = setup().await;
let thread = tm.create_thread(None, None, None).await.unwrap();
// A tool request message (non-text) followed by text — should not coalesce.
let tool_msg = Message::assistant().with_tool_request(
"call_1",
Ok(CallToolRequestParams::new("shell").with_arguments(
serde_json::json!({"command": "ls"})
.as_object()
.unwrap()
.clone(),
)),
);
tm.append_message(&thread.id, Some("s1"), &tool_msg)
.await
.unwrap();
tm.append_message(
&thread.id,
Some("s1"),
&Message::assistant().with_text("Result"),
)
.await
.unwrap();
let messages = tm.list_messages(&thread.id).await.unwrap();
assert_eq!(
messages.len(),
2,
"text should not coalesce with non-text predecessor"
);
}
#[tokio::test]
async fn empty_message_not_coalesced() {
let (tm, _tmp) = setup().await;
let thread = tm.create_thread(None, None, None).await.unwrap();
tm.append_message(
&thread.id,
Some("s1"),
&Message::assistant().with_text("Hello"),
)
.await
.unwrap();
// An empty assistant message (no content items).
let empty = Message::assistant();
tm.append_message(&thread.id, Some("s1"), &empty)
.await
.unwrap();
let messages = tm.list_messages(&thread.id).await.unwrap();
// Empty message should be inserted as a new row (not coalesced).
assert_eq!(messages.len(), 2);
assert_eq!(messages[0].as_concat_text(), "Hello");
}
#[tokio::test]
async fn metadata_change_prevents_coalescing() {
let (tm, _tmp) = setup().await;
let thread = tm.create_thread(None, None, None).await.unwrap();
tm.append_message(
&thread.id,
Some("s1"),
&Message::assistant().with_text("Visible"),
)
.await
.unwrap();
tm.append_message(
&thread.id,
Some("s1"),
&Message::assistant().with_text(" hidden").agent_only(),
)
.await
.unwrap();
let messages = tm.list_messages(&thread.id).await.unwrap();
assert_eq!(
messages.len(),
2,
"metadata boundary should break coalescing"
);
assert!(messages[0].metadata.user_visible);
assert!(!messages[1].metadata.user_visible);
assert_eq!(messages[0].as_concat_text(), "Visible");
assert_eq!(messages[1].as_concat_text(), " hidden");
}
ui/acp/src/generated/client.gen.ts
+24 -22
View File
@@ -9,6 +9,7 @@ export interface ExtMethodProvider {
import type {
AddExtensionRequest,
ArchiveSessionRequest,
CheckSecretRequest,
CheckSecretResponse,
DeleteSessionRequest,
@@ -16,8 +17,6 @@ import type {
ExportSessionResponse,
GetExtensionsRequest,
GetExtensionsResponse,
GetSessionRequest,
GetSessionResponse,
GetToolsRequest,
GetToolsResponse,
ImportSessionRequest,
@@ -31,6 +30,7 @@ import type {
RemoveConfigRequest,
RemoveExtensionRequest,
RemoveSecretRequest,
UnarchiveSessionRequest,
UpdateProviderRequest,
UpdateProviderResponse,
UpdateWorkingDirRequest,
@@ -41,7 +41,6 @@ import {
zCheckSecretResponse,
zExportSessionResponse,
zGetExtensionsResponse,
zGetSessionResponse,
zGetToolsResponse,
zImportSessionResponse,
zListProvidersResponse,
@@ -77,29 +76,10 @@ export class GooseExtClient {
await this.conn.extMethod("_goose/working_dir/update", params);
}
async sessionGet(params: GetSessionRequest): Promise<GetSessionResponse> {
const raw = await this.conn.extMethod("session/get", params);
return zGetSessionResponse.parse(raw) as GetSessionResponse;
}
async sessionDelete(params: DeleteSessionRequest): Promise<void> {
await this.conn.extMethod("session/delete", params);
}
async GooseSessionExport(
params: ExportSessionRequest,
): Promise<ExportSessionResponse> {
const raw = await this.conn.extMethod("_goose/session/export", params);
return zExportSessionResponse.parse(raw) as ExportSessionResponse;
}
async GooseSessionImport(
params: ImportSessionRequest,
): Promise<ImportSessionResponse> {
const raw = await this.conn.extMethod("_goose/session/import", params);
return zImportSessionResponse.parse(raw) as ImportSessionResponse;
}
async GooseConfigExtensions(
params: GetExtensionsRequest,
): Promise<GetExtensionsResponse> {
@@ -153,4 +133,26 @@ export class GooseExtClient {
async GooseSecretRemove(params: RemoveSecretRequest): Promise<void> {
await this.conn.extMethod("_goose/secret/remove", params);
}
async GooseSessionExport(
params: ExportSessionRequest,
): Promise<ExportSessionResponse> {
const raw = await this.conn.extMethod("_goose/session/export", params);
return zExportSessionResponse.parse(raw) as ExportSessionResponse;
}
async GooseSessionImport(
params: ImportSessionRequest,
): Promise<ImportSessionResponse> {
const raw = await this.conn.extMethod("_goose/session/import", params);
return zImportSessionResponse.parse(raw) as ImportSessionResponse;
}
async GooseSessionArchive(params: ArchiveSessionRequest): Promise<void> {
await this.conn.extMethod("_goose/session/archive", params);
}
async GooseSessionUnarchive(params: UnarchiveSessionRequest): Promise<void> {
await this.conn.extMethod("_goose/session/unarchive", params);
}
}
ui/acp/src/generated/index.ts
+21 -16
View File
@@ -1,6 +1,6 @@
// This file is auto-generated by @hey-api/openapi-ts
export type { AddExtensionRequest, CheckSecretRequest, CheckSecretResponse, DeleteSessionRequest, EmptyResponse, ExportSessionRequest, ExportSessionResponse, ExtRequest, ExtResponse, GetExtensionsRequest, GetExtensionsResponse, GetSessionRequest, GetSessionResponse, GetToolsRequest, GetToolsResponse, ImportSessionRequest, ImportSessionResponse, ListProvidersRequest, ListProvidersResponse, ProviderListEntry, ReadConfigRequest, ReadConfigResponse, ReadResourceRequest, ReadResourceResponse, RemoveConfigRequest, RemoveExtensionRequest, RemoveSecretRequest, UpdateProviderRequest, UpdateProviderResponse, UpdateWorkingDirRequest, UpsertConfigRequest, UpsertSecretRequest } from './types.gen.js';
export type { AddExtensionRequest, ArchiveSessionRequest, CheckSecretRequest, CheckSecretResponse, DeleteSessionRequest, EmptyResponse, ExportSessionRequest, ExportSessionResponse, ExtRequest, ExtResponse, GetExtensionsRequest, GetExtensionsResponse, GetToolsRequest, GetToolsResponse, ImportSessionRequest, ImportSessionResponse, ListProvidersRequest, ListProvidersResponse, ProviderListEntry, ReadConfigRequest, ReadConfigResponse, ReadResourceRequest, ReadResourceResponse, RemoveConfigRequest, RemoveExtensionRequest, RemoveSecretRequest, UnarchiveSessionRequest, UpdateProviderRequest, UpdateProviderResponse, UpdateWorkingDirRequest, UpsertConfigRequest, UpsertSecretRequest } from './types.gen.js';
export const GOOSE_EXT_METHODS = [
{
@@ -28,26 +28,11 @@ export const GOOSE_EXT_METHODS = [
requestType: "UpdateWorkingDirRequest",
responseType: "EmptyResponse",
},
{
method: "session/get",
requestType: "GetSessionRequest",
responseType: "GetSessionResponse",
},
{
method: "session/delete",
requestType: "DeleteSessionRequest",
responseType: "EmptyResponse",
},
{
method: "_goose/session/export",
requestType: "ExportSessionRequest",
responseType: "ExportSessionResponse",
},
{
method: "_goose/session/import",
requestType: "ImportSessionRequest",
responseType: "ImportSessionResponse",
},
{
method: "_goose/config/extensions",
requestType: "GetExtensionsRequest",
@@ -93,6 +78,26 @@ export const GOOSE_EXT_METHODS = [
requestType: "RemoveSecretRequest",
responseType: "EmptyResponse",
},
{
method: "_goose/session/export",
requestType: "ExportSessionRequest",
responseType: "ExportSessionResponse",
},
{
method: "_goose/session/import",
requestType: "ImportSessionRequest",
responseType: "ImportSessionResponse",
},
{
method: "_goose/session/archive",
requestType: "ArchiveSessionRequest",
responseType: "EmptyResponse",
},
{
method: "_goose/session/unarchive",
requestType: "UnarchiveSessionRequest",
responseType: "EmptyResponse",
},
] as const;
export type GooseExtMethod = (typeof GOOSE_EXT_METHODS)[number];
ui/acp/src/generated/types.gen.ts
+47 -51
View File
@@ -71,24 +71,6 @@ export type UpdateWorkingDirRequest = {
workingDir: string;
};
/**
* Get a session by ID.
*/
export type GetSessionRequest = {
sessionId: string;
includeMessages?: boolean;
};
/**
* Get a session response.
*/
export type GetSessionResponse = {
/**
* The session object with id, name, working_dir, timestamps, tokens, etc.
*/
session?: unknown;
};
/**
* Delete a session.
*/
@@ -96,37 +78,6 @@ export type DeleteSessionRequest = {
sessionId: string;
};
/**
* Export a session as a JSON string.
*/
export type ExportSessionRequest = {
sessionId: string;
};
/**
* Export session response.
*/
export type ExportSessionResponse = {
data: string;
};
/**
* Import a session from a JSON string.
*/
export type ImportSessionRequest = {
data: string;
};
/**
* Import session response.
*/
export type ImportSessionResponse = {
/**
* The imported session object.
*/
session?: unknown;
};
/**
* List configured extensions and any warnings.
*/
@@ -245,17 +196,62 @@ export type RemoveSecretRequest = {
key: string;
};
/**
* Export a session as a JSON string.
*/
export type ExportSessionRequest = {
sessionId: string;
};
/**
* Export session response — raw JSON of the goose session with `conversation`.
*/
export type ExportSessionResponse = {
data: string;
};
/**
* Import a session from a JSON string.
*/
export type ImportSessionRequest = {
data: string;
};
/**
* Import session response — metadata about the newly created session.
*/
export type ImportSessionResponse = {
sessionId: string;
title?: string | null;
updatedAt?: string | null;
messageCount: number;
};
/**
* Archive a session (soft delete).
*/
export type ArchiveSessionRequest = {
sessionId: string;
};
/**
* Unarchive a previously archived session.
*/
export type UnarchiveSessionRequest = {
sessionId: string;
};
export type ExtRequest = {
id: string;
method: string;
params?: AddExtensionRequest | RemoveExtensionRequest | GetToolsRequest | ReadResourceRequest | UpdateWorkingDirRequest | GetSessionRequest | DeleteSessionRequest | ExportSessionRequest | ImportSessionRequest | GetExtensionsRequest | UpdateProviderRequest | ListProvidersRequest | ReadConfigRequest | UpsertConfigRequest | RemoveConfigRequest | CheckSecretRequest | UpsertSecretRequest | RemoveSecretRequest | {
params?: AddExtensionRequest | RemoveExtensionRequest | GetToolsRequest | ReadResourceRequest | UpdateWorkingDirRequest | DeleteSessionRequest | GetExtensionsRequest | UpdateProviderRequest | ListProvidersRequest | ReadConfigRequest | UpsertConfigRequest | RemoveConfigRequest | CheckSecretRequest | UpsertSecretRequest | RemoveSecretRequest | ExportSessionRequest | ImportSessionRequest | ArchiveSessionRequest | UnarchiveSessionRequest | {
[key: string]: unknown;
} | null;
};
export type ExtResponse = {
id: string;
result?: EmptyResponse | GetToolsResponse | ReadResourceResponse | GetSessionResponse | ExportSessionResponse | ImportSessionResponse | GetExtensionsResponse | UpdateProviderResponse | ListProvidersResponse | ReadConfigResponse | CheckSecretResponse | unknown;
result?: EmptyResponse | GetToolsResponse | ReadResourceResponse | GetExtensionsResponse | UpdateProviderResponse | ListProvidersResponse | ReadConfigResponse | CheckSecretResponse | ExportSessionResponse | ImportSessionResponse | unknown;
} | {
error: {
code: number;
ui/acp/src/generated/zod.gen.ts
+59 -51
View File
@@ -61,21 +61,6 @@ export const zUpdateWorkingDirRequest = z.object({
workingDir: z.string()
});
/**
* Get a session by ID.
*/
export const zGetSessionRequest = z.object({
sessionId: z.string(),
includeMessages: z.boolean().optional().default(false)
});
/**
* Get a session response.
*/
export const zGetSessionResponse = z.object({
session: z.unknown().optional().default(null)
});
/**
* Delete a session.
*/
@@ -83,34 +68,6 @@ export const zDeleteSessionRequest = z.object({
sessionId: z.string()
});
/**
* Export a session as a JSON string.
*/
export const zExportSessionRequest = z.object({
sessionId: z.string()
});
/**
* Export session response.
*/
export const zExportSessionResponse = z.object({
data: z.string()
});
/**
* Import a session from a JSON string.
*/
export const zImportSessionRequest = z.object({
data: z.string()
});
/**
* Import session response.
*/
export const zImportSessionResponse = z.object({
session: z.unknown().optional().default(null)
});
/**
* List configured extensions and any warnings.
*/
@@ -226,6 +183,57 @@ export const zRemoveSecretRequest = z.object({
key: z.string()
});
/**
* Export a session as a JSON string.
*/
export const zExportSessionRequest = z.object({
sessionId: z.string()
});
/**
* Export session response — raw JSON of the goose session with `conversation`.
*/
export const zExportSessionResponse = z.object({
data: z.string()
});
/**
* Import a session from a JSON string.
*/
export const zImportSessionRequest = z.object({
data: z.string()
});
/**
* Import session response — metadata about the newly created session.
*/
export const zImportSessionResponse = z.object({
sessionId: z.string(),
title: z.union([
z.string(),
z.null()
]).optional(),
updatedAt: z.union([
z.string(),
z.null()
]).optional(),
messageCount: z.coerce.bigint().gte(BigInt(0)).max(BigInt('18446744073709551615'), { message: 'Invalid value: Expected uint64 to be <= 18446744073709551615' })
});
/**
* Archive a session (soft delete).
*/
export const zArchiveSessionRequest = z.object({
sessionId: z.string()
});
/**
* Unarchive a previously archived session.
*/
export const zUnarchiveSessionRequest = z.object({
sessionId: z.string()
});
export const zExtRequest = z.object({
id: z.string(),
method: z.string(),
@@ -236,10 +244,7 @@ export const zExtRequest = z.object({
zGetToolsRequest,
zReadResourceRequest,
zUpdateWorkingDirRequest,
zGetSessionRequest,
zDeleteSessionRequest,
zExportSessionRequest,
zImportSessionRequest,
zGetExtensionsRequest,
zUpdateProviderRequest,
zListProvidersRequest,
@@ -248,7 +253,11 @@ export const zExtRequest = z.object({
zRemoveConfigRequest,
zCheckSecretRequest,
zUpsertSecretRequest,
zRemoveSecretRequest
zRemoveSecretRequest,
zExportSessionRequest,
zImportSessionRequest,
zArchiveSessionRequest,
zUnarchiveSessionRequest
]),
z.union([
z.record(z.unknown()),
@@ -265,14 +274,13 @@ export const zExtResponse = z.union([
zEmptyResponse,
zGetToolsResponse,
zReadResourceResponse,
zGetSessionResponse,
zExportSessionResponse,
zImportSessionResponse,
zGetExtensionsResponse,
zUpdateProviderResponse,
zListProvidersResponse,
zReadConfigResponse,
zCheckSecretResponse
zCheckSecretResponse,
zExportSessionResponse,
zImportSessionResponse
]),
z.unknown()
]).optional()
ui/desktop/openapi.json
+4
View File
@@ -7748,6 +7748,10 @@
"session_type": {
"$ref": "#/components/schemas/SessionType"
},
"thread_id": {
"type": "string",
"nullable": true
},
"total_tokens": {
"type": "integer",
"format": "int32",
ui/desktop/src/api/types.gen.ts
+1
View File
@@ -1232,6 +1232,7 @@ export type Session = {
recipe?: Recipe | null;
schedule_id?: string | null;
session_type?: SessionType;
thread_id?: string | null;
total_tokens?: number | null;
updated_at: string;
user_recipe_values?: {