From 659257316d266415bb8df1e8aa21f6c5e859bd69 Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Mon, 22 Jun 2026 11:33:19 +1000 Subject: [PATCH] feat(ui): use acp new session on desktop (only for non recipe session) (#9914) --- crates/goose/src/acp/server.rs | 12 +++- crates/goose/src/acp/server/extensions.rs | 16 +++-- crates/goose/src/acp/server/new_session.rs | 43 +++++++----- .../acp/__tests__/chatSessionStore.test.ts | 8 --- ui/desktop/src/acp/chatSessionController.ts | 23 ++++++- ui/desktop/src/acp/chatSessionStore.ts | 17 ----- ui/desktop/src/acp/extensions.ts | 12 +++- ui/desktop/src/acp/sessions.ts | 37 +++++++++- ui/desktop/src/sessions.ts | 67 +++++++++++++------ 9 files changed, 163 insertions(+), 72 deletions(-) diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 12bd5d4e4..2459b8132 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -985,13 +985,18 @@ impl GooseAcpAgent { &self, config: &Config, mcp_servers: Vec, + goose_extensions: Option>, ) -> Result, agent_client_protocol::Error> { let mut extensions = Vec::new(); for builtin in &self.builtins { push_or_replace_extension(&mut extensions, builtin_to_extension_config(builtin)); } - if mcp_servers.is_empty() { + if let Some(goose_extensions) = goose_extensions { + for extension in extensions::goose_extensions_to_configs(goose_extensions)? { + push_or_replace_extension(&mut extensions, extension); + } + } else if mcp_servers.is_empty() { for extension in get_enabled_extensions_with_config(config) { push_or_replace_extension(&mut extensions, extension); } @@ -1113,7 +1118,7 @@ impl GooseAcpAgent { || EnabledExtensionsState::from_extension_data(&session.extension_data).is_none() { let extension_data = - self.build_enabled_extensions_data(config, &session, mcp_servers)?; + self.build_enabled_extensions_data(config, &session, mcp_servers, None)?; builder = builder.extension_data(extension_data); session_needs_update = true; } @@ -1142,8 +1147,9 @@ impl GooseAcpAgent { config: &Config, session: &Session, mcp_servers: Vec, + goose_extensions: Option>, ) -> Result { - let extensions = self.initial_session_extensions(config, mcp_servers)?; + let extensions = self.initial_session_extensions(config, mcp_servers, goose_extensions)?; let mut extension_data = session.extension_data.clone(); EnabledExtensionsState::new(extensions) .to_extension_data(&mut extension_data) diff --git a/crates/goose/src/acp/server/extensions.rs b/crates/goose/src/acp/server/extensions.rs index eb43ef9b5..347a4fba7 100644 --- a/crates/goose/src/acp/server/extensions.rs +++ b/crates/goose/src/acp/server/extensions.rs @@ -9,7 +9,7 @@ impl GooseAcpAgent { req: AddSessionExtensionRequest, ) -> Result { let session_id = &req.session_id; - let config = goose_extension_to_session_config(req.extension)?; + let config = goose_extension_to_config_without_secrets(req.extension)?; let agent = self.get_session_agent(&req.session_id).await?; agent .add_extension(config, session_id) @@ -315,19 +315,27 @@ fn goose_extension_to_config( }) } -fn goose_extension_to_session_config( +fn goose_extension_to_config_without_secrets( extension: GooseExtension, ) -> Result { let conversion = goose_extension_to_config(extension)?; if !conversion.secret_updates.is_empty() { return Err(agent_client_protocol::Error::invalid_params().data( - "literal environment values are not supported for session extensions; use envKeys", + "extension env values must be passed via envKeys referencing stored secrets, not inline env", )); } - Ok(conversion.config) } +pub(super) fn goose_extensions_to_configs( + extensions: Vec, +) -> Result, agent_client_protocol::Error> { + extensions + .into_iter() + .map(goose_extension_to_config_without_secrets) + .collect() +} + fn config_entry_to_goose_entry( entry: ExtensionEntry, ) -> Result, agent_client_protocol::Error> { diff --git a/crates/goose/src/acp/server/new_session.rs b/crates/goose/src/acp/server/new_session.rs index a17d568de..0cb556351 100644 --- a/crates/goose/src/acp/server/new_session.rs +++ b/crates/goose/src/acp/server/new_session.rs @@ -1,22 +1,19 @@ -use crate::acp::server::{meta_string, sid_short, validate_absolute_cwd, ResultExt}; +use crate::acp::custom_requests::GooseExtension; +use crate::acp::server::{meta_string, validate_absolute_cwd, ResultExt}; use crate::config::{Config, GooseMode}; use crate::session::SessionType; use super::GooseAcpAgent; -use agent_client_protocol::schema::{NewSessionRequest, NewSessionResponse, SessionId}; +use agent_client_protocol::schema::{Meta, NewSessionRequest, NewSessionResponse, SessionId}; use agent_client_protocol::{Client, ConnectionTo}; use std::collections::HashMap; -use tracing::debug; impl GooseAcpAgent { - #[allow(dead_code)] pub(super) async fn handle_new_session( &self, cx: &ConnectionTo, args: NewSessionRequest, ) -> Result { - debug!(?args, "new session request"); - let t_start = std::time::Instant::now(); validate_absolute_cwd(&args.cwd)?; let project_id = meta_string(args.meta.as_ref(), "projectId")?; let session_type = match meta_string(args.meta.as_ref(), "client")? { @@ -33,8 +30,8 @@ impl GooseAcpAgent { } None => super::resolve_default_provider_model_config(config)?, }; + let goose_extensions = meta_goose_extensions(args.meta.as_ref())?; let current_mode: GooseMode = config.get_goose_mode().unwrap_or_default(); - let t0 = std::time::Instant::now(); let mut goose_session = self .session_manager .create_session( @@ -46,8 +43,12 @@ impl GooseAcpAgent { .await .internal_err_ctx("Failed to create session")?; let mut builder = self.session_manager.update(&goose_session.id); - let extension_data = - self.build_enabled_extensions_data(config, &goose_session, args.mcp_servers)?; + let extension_data = self.build_enabled_extensions_data( + config, + &goose_session, + args.mcp_servers, + goose_extensions, + )?; builder = builder .provider_name(resolved_provider) .model_config(resolved_model_config) @@ -66,8 +67,6 @@ impl GooseAcpAgent { .await .internal_err_ctx("Failed to reload session")?; let session_id_str = goose_session.id.clone(); - let sid = sid_short(&session_id_str); - debug!(target: "perf", sid = %sid, ms = t0.elapsed().as_millis() as u64, "perf: new_session create_session"); let (_agent, extension_results) = self .activate_acp_session(cx, &goose_session, HashMap::new()) @@ -101,12 +100,22 @@ impl GooseAcpAgent { &goose_session, self.supports_goose_custom_notifications(), )?; - debug!( - target: "perf", - sid = %sid, - ms = t_start.elapsed().as_millis() as u64, - "perf: new_session done" - ); Ok(response) } } + +fn meta_goose_extensions( + meta: Option<&Meta>, +) -> Result>, agent_client_protocol::Error> { + let Some(value) = meta.and_then(|m| m.get("enabledExtensions")) else { + return Ok(None); + }; + if value.is_null() { + return Ok(None); + } + serde_json::from_value(value.clone()) + .map(Some) + .map_err(|e| { + agent_client_protocol::Error::invalid_params().data(format!("enabledExtensions: {e}")) + }) +} diff --git a/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts b/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts index e9bca9eb6..feb750fdc 100644 --- a/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts +++ b/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts @@ -156,14 +156,6 @@ describe('acpChatSessionStore', () => { expect(snapshot.session?.id).toBe(currentSessionId); expect(snapshot.messages).toEqual([initialMessage]); - expect(snapshot.tokenState).toMatchObject({ - inputTokens: 1, - outputTokens: 2, - totalTokens: 3, - accumulatedInputTokens: 4, - accumulatedOutputTokens: 5, - accumulatedTotalTokens: 9, - }); expect(snapshot.chatState).toBe(ChatState.Idle); expect(snapshot.sessionLoadError).toBeUndefined(); }); diff --git a/ui/desktop/src/acp/chatSessionController.ts b/ui/desktop/src/acp/chatSessionController.ts index 2bf0c182a..cf05dc200 100644 --- a/ui/desktop/src/acp/chatSessionController.ts +++ b/ui/desktop/src/acp/chatSessionController.ts @@ -1,5 +1,6 @@ import { v7 as uuidv7 } from 'uuid'; -import { updateSessionUserRecipeValues, type Message } from '../api'; +import { updateSessionUserRecipeValues, type Message, type Session } from '../api'; +import type { GooseExtension } from '@aaif/goose-sdk'; import { AppEvents } from '../constants/events'; import { ChatState } from '../types/chatState'; import { errorMessage } from '../utils/conversionUtils'; @@ -17,6 +18,7 @@ import { acpCancelPrompt, acpPromptSession } from './prompt'; import { acpForkSession, acpLoadSession, + acpNewSession, acpTruncateSessionConversation, isAcpSessionLoadInFlight, sessionInfoToSession, @@ -35,6 +37,7 @@ export interface AcpSubmitMessageOptions extends AcpSnapshotOptions { } export interface AcpChatSessionController { + createSession(cwd: string, gooseExtensions: GooseExtension[]): Promise; loadSession(sessionId: string, options?: AcpLoadSessionOptions): Promise; submitMessage( sessionId: string, @@ -73,6 +76,19 @@ function createAcpCreditsExhaustedMessage(error: AcpCreditsExhaustedError): Mess }; } +async function createSession(cwd: string, gooseExtensions: GooseExtension[]): Promise { + const { sessionId, sessionInfo, meta } = await acpNewSession(cwd, gooseExtensions); + const session = sessionInfoToSession(sessionInfo, meta); + + showExtensionLoadResults(meta.extensionResults); + window.dispatchEvent( + new CustomEvent(AppEvents.SESSION_EXTENSIONS_LOADED, { detail: { sessionId } }) + ); + acpChatSessionActions.finishSessionLoad(sessionId, session); + + return session; +} + async function loadSession(sessionId: string, options: AcpLoadSessionOptions = {}): Promise { const cached = acpChatSessionStore.getSnapshot(sessionId); if (cached?.session) { @@ -106,6 +122,10 @@ async function submitMessage( userMessage: Message, options: AcpSubmitMessageOptions ): Promise { + if (acpChatSessionStore.getSnapshot(sessionId)?.activePromptAttemptId) { + return; + } + const promptAttemptId = uuidv7(); acpChatSessionActions.startPromptAttempt(sessionId, promptAttemptId); @@ -247,6 +267,7 @@ async function setRecipeUserParams( } export const acpChatSessionController: AcpChatSessionController = { + createSession, loadSession, submitMessage, stop, diff --git a/ui/desktop/src/acp/chatSessionStore.ts b/ui/desktop/src/acp/chatSessionStore.ts index 70cee7f13..7dfc62608 100644 --- a/ui/desktop/src/acp/chatSessionStore.ts +++ b/ui/desktop/src/acp/chatSessionStore.ts @@ -38,22 +38,6 @@ const initialTokenState: TokenState = { accumulatedTotalTokens: 0, }; -function tokenStateFromSession(session: Session | undefined): Partial { - return { - inputTokens: session?.usage?.input_tokens ?? 0, - outputTokens: session?.usage?.output_tokens ?? 0, - totalTokens: session?.usage?.total_tokens ?? 0, - cacheReadTokens: session?.usage?.cache_read_input_tokens ?? 0, - cacheWriteTokens: session?.usage?.cache_write_input_tokens ?? 0, - accumulatedInputTokens: session?.accumulated_usage?.input_tokens ?? 0, - accumulatedOutputTokens: session?.accumulated_usage?.output_tokens ?? 0, - accumulatedTotalTokens: session?.accumulated_usage?.total_tokens ?? 0, - accumulatedCacheReadTokens: session?.accumulated_usage?.cache_read_input_tokens ?? 0, - accumulatedCacheWriteTokens: session?.accumulated_usage?.cache_write_input_tokens ?? 0, - ...(session?.accumulated_cost != null ? { accumulatedCost: session.accumulated_cost } : {}), - }; -} - export interface AcpChatSessionStore { getSnapshot(sessionId: string): AcpChatSessionSnapshot | undefined; } @@ -180,7 +164,6 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { const finishSessionLoad: AcpChatSessionActions['finishSessionLoad'] = (sessionId, session) => { const entry = getOrCreateEntry(sessionId); entry.session = session; - entry.tokenState = { ...entry.tokenState, ...tokenStateFromSession(session) }; entry.sessionLoadError = undefined; entry.chatState = entry.activePromptAttemptId ? ChatState.Streaming : ChatState.Idle; return notify(sessionId, entry); diff --git a/ui/desktop/src/acp/extensions.ts b/ui/desktop/src/acp/extensions.ts index 7964210b8..2fad4536a 100644 --- a/ui/desktop/src/acp/extensions.ts +++ b/ui/desktop/src/acp/extensions.ts @@ -1,7 +1,11 @@ import type { ExtensionResponse, ExtensionEntry } from '../api'; -import type { GooseExtensionEntry, McpServer } from '@aaif/goose-sdk'; +import type { GooseExtension, GooseExtensionEntry, McpServer } from '@aaif/goose-sdk'; import { getAcpClient } from './acpConnection'; +export function gooseExtensionName(extension: GooseExtension): string { + return extension.type === 'mcp' ? extension.server.name : extension.name; +} + function headersToRecord(headers: { name: string; value: string }[] = []) { return Object.fromEntries(headers.map(({ name, value }) => [name, value])); } @@ -65,6 +69,12 @@ function gooseExtensionEntryToExtensionEntry(entry: GooseExtensionEntry): Extens return null; } +export async function getConfiguredGooseExtensions(): Promise { + const client = await getAcpClient(); + const response = await client.goose.configExtensionsList_unstable({}); + return response.extensions; +} + export async function getConfiguredExtensions(): Promise { const client = await getAcpClient(); const response = await client.goose.configExtensionsList_unstable({}); diff --git a/ui/desktop/src/acp/sessions.ts b/ui/desktop/src/acp/sessions.ts index d6db5a370..fb99b060f 100644 --- a/ui/desktop/src/acp/sessions.ts +++ b/ui/desktop/src/acp/sessions.ts @@ -2,8 +2,10 @@ import type { ForkSessionRequest, ListSessionsRequest, LoadSessionResponse, + NewSessionRequest, SessionInfo, } from '@agentclientprotocol/sdk'; +import type { GooseExtension } from '@aaif/goose-sdk'; import { getAcpClient } from './acpConnection'; import { DEFAULT_CHAT_TITLE } from '../contexts/ChatContext'; import type { ExtensionLoadResult, Recipe, Session } from '../api'; @@ -56,8 +58,8 @@ export interface AcpLoadSessionResult { const inFlightSessionLoads = new Map>(); -export function parseLoadMeta(response: LoadSessionResponse): LoadSessionMeta { - const meta = (response._meta ?? {}) as LoadSessionMeta; +function parseSessionResponseMeta(rawMeta: unknown): LoadSessionMeta { + const meta = (rawMeta ?? {}) as LoadSessionMeta; return { recipe: meta.recipe, userRecipeValues: meta.userRecipeValues, @@ -66,6 +68,10 @@ export function parseLoadMeta(response: LoadSessionResponse): LoadSessionMeta { }; } +export function parseLoadMeta(response: LoadSessionResponse): LoadSessionMeta { + return parseSessionResponseMeta(response._meta); +} + function sessionInfoMeta(s: SessionInfo): GooseSessionInfoMeta { return (s._meta ?? {}) as GooseSessionInfoMeta; } @@ -197,6 +203,33 @@ async function loadAcpSession(sessionId: string): Promise }; } +export interface AcpNewSessionResult { + sessionId: string; + sessionInfo: SessionInfo; + meta: LoadSessionMeta; +} + +export async function acpNewSession( + cwd: string, + gooseExtensions: GooseExtension[] +): Promise { + const client = await getAcpClient(); + const meta: Record = { client: 'goose-desktop' }; + if (gooseExtensions.length > 0) { + meta.enabledExtensions = gooseExtensions; + } + const request: NewSessionRequest = { cwd, mcpServers: [], _meta: meta }; + const response = await client.newSession(request); + const sessionId = String(response.sessionId); + const sessionInfoResponse = await client.goose.sessionInfo_unstable({ sessionId }); + + return { + sessionId, + sessionInfo: sessionInfoResponse.session, + meta: parseSessionResponseMeta(response._meta), + }; +} + export async function acpDeleteSession(sessionId: string): Promise { const client = await getAcpClient(); await client.goose.sessionDelete({ sessionId }); diff --git a/ui/desktop/src/sessions.ts b/ui/desktop/src/sessions.ts index 2fa808294..41e51aec2 100644 --- a/ui/desktop/src/sessions.ts +++ b/ui/desktop/src/sessions.ts @@ -4,6 +4,9 @@ import type { setViewType } from './hooks/useNavigation'; import type { FixedExtensionEntry } from './components/ConfigContext'; import { AppEvents } from './constants/events'; import { decodeRecipe, Recipe } from './recipe'; +import { USE_ACP_CHAT } from './acpChatFeatureFlag'; +import { acpChatSessionController } from './acp/chatSessionController'; +import { getConfiguredGooseExtensions, gooseExtensionName } from './acp/extensions'; export function getSessionDisplayName(session: Session): string { if (session.user_set_name) { @@ -40,15 +43,51 @@ export function resumeSession(session: Session, setView: setViewType) { }); } +interface CreateSessionOptions { + recipeDeeplink?: string; + recipeId?: string; + extensionConfigs?: ExtensionConfig[]; + allExtensions?: FixedExtensionEntry[]; +} + +function selectedExtensionConfigs(options?: CreateSessionOptions): ExtensionConfig[] { + if (options?.extensionConfigs && options.extensionConfigs.length > 0) { + return options.extensionConfigs; + } + if (options?.allExtensions) { + return options.allExtensions + .filter((extension) => extension.enabled) + .map((extension) => { + const { enabled: _enabled, ...config } = extension; + return config as ExtensionConfig; + }); + } + return []; +} + +async function createAcpSession( + workingDir: string, + options?: CreateSessionOptions +): Promise { + const selectedNames = new Set(selectedExtensionConfigs(options).map((config) => config.name)); + const gooseExtensions = + selectedNames.size > 0 + ? (await getConfiguredGooseExtensions()) + .filter((entry) => selectedNames.has(gooseExtensionName(entry.extension))) + .map((entry) => entry.extension) + : []; + return acpChatSessionController.createSession(workingDir, gooseExtensions); +} + export async function createSession( workingDir: string, - options?: { - recipeDeeplink?: string; - recipeId?: string; - extensionConfigs?: ExtensionConfig[]; - allExtensions?: FixedExtensionEntry[]; - } + options?: CreateSessionOptions ): Promise { + const hasRecipe = Boolean(options?.recipeId || options?.recipeDeeplink); + if (USE_ACP_CHAT && !hasRecipe) { + return createAcpSession(workingDir, options); + } + const body: { working_dir: string; recipe?: Recipe; @@ -64,19 +103,9 @@ export async function createSession( body.recipe = await decodeRecipe(options.recipeDeeplink); } - if (options?.extensionConfigs && options.extensionConfigs.length > 0) { - body.extension_overrides = options.extensionConfigs; - } else if (options?.allExtensions) { - const extensionConfigs = options.allExtensions - .filter((extension) => extension.enabled) - .map((extension) => { - const { enabled: _enabled, ...config } = extension; - return config as ExtensionConfig; - }); - - if (extensionConfigs.length > 0) { - body.extension_overrides = extensionConfigs; - } + const extensionConfigs = selectedExtensionConfigs(options); + if (extensionConfigs.length > 0) { + body.extension_overrides = extensionConfigs; } const newAgent = await startAgent({