diff --git a/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs b/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs index 70c058d998..d82123993e 100644 --- a/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs +++ b/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs @@ -5041,7 +5041,7 @@ mod tests { .result_for_assistant .as_deref() .unwrap_or_default() - .contains("Call GetToolSpec first")); + .contains("Call GetToolSpec with {\"tool_name\":\"get_weather\"}")); } #[tokio::test] @@ -5289,7 +5289,7 @@ mod tests { assert_failed_task_contains( &pipeline, "tool_1", - "Call GetToolSpec first with {\"tool_name\":\"WebFetch\"}", + "Call GetToolSpec with {\"tool_name\":\"WebFetch\"}", ); } @@ -5714,7 +5714,7 @@ mod tests { assert!(err .to_string() - .contains("Call GetToolSpec first with {\"tool_name\":\"WebFetch\"}")); + .contains("Call GetToolSpec with {\"tool_name\":\"WebFetch\"}")); } #[test] diff --git a/src/web-ui/src/flow_chat/services/flow-chat-manager/SessionModule.test.ts b/src/web-ui/src/flow_chat/services/flow-chat-manager/SessionModule.test.ts index 4701a3e0b6..87a746a710 100644 --- a/src/web-ui/src/flow_chat/services/flow-chat-manager/SessionModule.test.ts +++ b/src/web-ui/src/flow_chat/services/flow-chat-manager/SessionModule.test.ts @@ -26,6 +26,7 @@ import { } from '../sessionOpenIntent'; import type { Session } from '../../types/flow-chat'; import type { ReviewTeamRunManifest } from '@/shared/services/reviewTeamService'; +import { workspaceManager } from '@/infrastructure/services/business/workspaceManager'; const agentApiMocks = vi.hoisted(() => ({ ensureCoordinatorSession: vi.fn(), @@ -98,6 +99,7 @@ vi.mock('@/infrastructure/services/business/workspaceManager', () => ({ getState: () => ({ currentWorkspace: { id: 'workspace-1', rootPath: '/home/wsp/projects/Test', workspaceKind: 'normal' }, openedWorkspaces: new Map([['workspace-1', { id: 'workspace-1', rootPath: '/home/wsp/projects/Test', workspaceKind: 'normal' }]]), + recentWorkspaces: [], }), }, })); @@ -659,6 +661,8 @@ describe('SessionModule historical session coordination', () => { beforeEach(() => { vi.useFakeTimers(); dispatchStoreMocks.jobs = {}; + agentApiMocks.ensureCoordinatorSession.mockReset(); + agentApiMocks.createSession.mockReset(); }); afterEach(async () => { @@ -1406,6 +1410,156 @@ describe('SessionModule historical session coordination', () => { expect(agentApiMocks.createSession).not.toHaveBeenCalled(); }); + it.each(['local', 'peer-worktree'])('prepares a worktree follow-up through its owning project on %s', async surfaceId => { + activateSurface(surfaceId); + const session = createSession({ + isHistorical: false, + historyState: 'ready', + workspaceId: 'worktree-not-in-catalog', + workspacePath: '/worktrees/task', + projectWorkspaceId: 'workspace-1', + projectWorkspacePath: '/home/wsp/projects/Test', + config: { + workspaceId: 'worktree-not-in-catalog', + executionTarget: { kind: 'managedWorktree', worktreeId: 'worktree-1', rootPath: '/worktrees/task' }, + }, + dialogTurns: [{ id: 'turn-1', status: 'completed' } as any], + }); + const { context } = createContext(session); + agentApiMocks.ensureCoordinatorSession.mockResolvedValueOnce(undefined); + + try { + await expect(ensureBackendSession(context, session.sessionId)).resolves.toBeUndefined(); + expect(agentApiMocks.ensureCoordinatorSession).toHaveBeenCalledWith({ + sessionId: session.sessionId, + workspaceId: 'workspace-1', + includeInternal: false, + }); + expect(agentApiMocks.createSession).not.toHaveBeenCalled(); + expect(context.flowChatStore.getState().sessions.get(session.sessionId)).toMatchObject({ + workspaceId: 'worktree-not-in-catalog', + workspacePath: '/worktrees/task', + }); + } finally { + activateSurface(LOCAL_SURFACE_ID); + } + }); + + it.each(['ensure', 'retry'])('recreates an empty worktree session through its project during %s', async operation => { + const { context } = createContext(createSession({ + isHistorical: false, + historyState: 'ready', + workspaceId: 'worktree-not-in-catalog', + workspacePath: '/worktrees/task', + projectWorkspaceId: 'workspace-1', + projectWorkspacePath: '/home/wsp/projects/Test', + config: { + executionTarget: { kind: 'managedWorktree', worktreeId: 'worktree-1', rootPath: '/worktrees/task' }, + }, + })); + agentApiMocks.ensureCoordinatorSession.mockRejectedValueOnce(new Error('Session metadata not found')); + agentApiMocks.createSession.mockResolvedValueOnce(undefined); + + await (operation === 'ensure' ? ensureBackendSession : retryCreateBackendSession)(context, 'history-1'); + + expect(agentApiMocks.createSession).toHaveBeenCalledWith(expect.objectContaining({ + sessionId: 'history-1', + workspaceId: 'workspace-1', + projectWorkspacePath: '/home/wsp/projects/Test', + executionTarget: { kind: 'existingWorktree', worktreeId: 'worktree-1' }, + })); + }); + + it('restores a worktree session whose project identity is only in legacy config', async () => { + const { context } = createContext(createSession({ + workspaceId: undefined, + isHistorical: false, + historyState: 'ready', + contextRestoreState: 'pending', + config: { + workspaceId: 'worktree-not-in-catalog', + projectWorkspaceId: 'workspace-1', + executionTarget: { kind: 'managedWorktree', worktreeId: 'worktree-1', rootPath: '/worktrees/task' }, + }, + dialogTurns: [{ id: 'turn-1', status: 'completed' } as any], + })); + agentApiMocks.ensureCoordinatorSession.mockResolvedValueOnce(undefined); + + await ensureBackendSession(context, 'history-1'); + + expect(agentApiMocks.ensureCoordinatorSession).toHaveBeenCalledWith({ + sessionId: 'history-1', workspaceId: 'workspace-1', includeInternal: false, + }); + expect(agentApiMocks.createSession).not.toHaveBeenCalled(); + expect(context.flowChatStore.getState().sessions.get('history-1')?.contextRestoreState).toBe('ready'); + }); + + it('keeps an ordinary linked-workspace session addressed to its own workspace', async () => { + const { context } = createContext(createSession({ + isHistorical: false, + historyState: 'ready', + projectWorkspaceId: 'another-project', + config: { executionTarget: { kind: 'local', rootPath: '/home/wsp/projects/Test' } }, + dialogTurns: [{ id: 'turn-1', status: 'completed' } as any], + })); + agentApiMocks.ensureCoordinatorSession.mockResolvedValueOnce(undefined); + + await ensureBackendSession(context, 'history-1'); + + expect(agentApiMocks.ensureCoordinatorSession).toHaveBeenCalledWith({ + sessionId: 'history-1', workspaceId: 'workspace-1', includeInternal: false, + }); + }); + + it('does not fall back to the active local workspace when the owning project is unavailable', async () => { + const { context } = createContext(createSession({ + isHistorical: false, + historyState: 'ready', + projectWorkspaceId: 'missing-project', + config: { executionTarget: { kind: 'managedWorktree', worktreeId: 'worktree-1', rootPath: '/worktrees/task' } }, + dialogTurns: [{ id: 'turn-1', status: 'completed' } as any], + })); + + await expect(ensureBackendSession(context, 'history-1')).rejects.toThrow('Workspace ID is unavailable: missing-project'); + expect(agentApiMocks.ensureCoordinatorSession).not.toHaveBeenCalled(); + expect(agentApiMocks.createSession).not.toHaveBeenCalled(); + }); + + it('preserves SSH workspace identity and connection when recreating an empty session', async () => { + const remoteWorkspace = { + id: 'ssh-workspace', rootPath: '/remote/project', workspaceKind: 'remote', + connectionId: 'ssh-connection', sshHost: 'remote-host', + }; + const catalog = workspaceManager.getState(); + const getState = vi.spyOn(workspaceManager, 'getState').mockReturnValue({ + ...catalog, + openedWorkspaces: new Map([[remoteWorkspace.id, remoteWorkspace]]), + } as ReturnType); + const { context } = createContext(createSession({ + workspaceId: remoteWorkspace.id, + workspacePath: remoteWorkspace.rootPath, + isHistorical: false, + historyState: 'ready', + })); + agentApiMocks.ensureCoordinatorSession.mockRejectedValueOnce(new Error('Session metadata not found')); + agentApiMocks.createSession.mockResolvedValueOnce(undefined); + + try { + await ensureBackendSession(context, 'history-1'); + expect(agentApiMocks.ensureCoordinatorSession).toHaveBeenCalledWith({ + sessionId: 'history-1', workspaceId: remoteWorkspace.id, includeInternal: false, + }); + expect(agentApiMocks.createSession).toHaveBeenCalledWith(expect.objectContaining({ + workspaceId: remoteWorkspace.id, + workspacePath: remoteWorkspace.rootPath, + remoteConnectionId: remoteWorkspace.connectionId, + remoteSshHost: remoteWorkspace.sshHost, + })); + } finally { + getState.mockRestore(); + } + }); + it('restores pending backend context for a view-restored session before send', async () => { const { context } = createContext(createSession({ isHistorical: false, diff --git a/src/web-ui/src/flow_chat/services/flow-chat-manager/SessionModule.ts b/src/web-ui/src/flow_chat/services/flow-chat-manager/SessionModule.ts index 9ec7dbcc88..543d551844 100644 --- a/src/web-ui/src/flow_chat/services/flow-chat-manager/SessionModule.ts +++ b/src/web-ui/src/flow_chat/services/flow-chat-manager/SessionModule.ts @@ -1,5 +1,4 @@ import { requireSessionOwningWorkspaceId } from '../../utils/sessionOrdering'; -import { requireSessionWorkspaceId } from '../../utils/sessionWorkspace'; /** * Session management module * Handles session creation, switching, deletion, and other operations @@ -850,7 +849,9 @@ export async function ensureBackendSession( } const latestSession = context.flowChatStore.getState().sessions.get(sessionId) ?? session; - const workspaceId = requireSessionWorkspaceId(latestSession); + // Coordinator state belongs to the project; a managed worktree's execution + // record need not be in the opened/recent workspace catalog after first send. + const workspaceId = requireSessionOwningWorkspaceId(latestSession); const workspace = resolveSessionWorkspace({ workspaceId }); const workspacePath = workspace.rootPath; const projectWorkspacePath = requireSessionProjectWorkspacePath(latestSession, sessionId); @@ -984,7 +985,7 @@ export async function ensureBackendSession( worktreeId: latestSession.config.executionTarget.worktreeId, } : { kind: 'local' }, - workspaceId: latestSession.workspaceId, + workspaceId, remoteConnectionId: effectiveConnectionId, remoteSshHost: effectiveSshHost, relationship: buildCreateSessionRelationship(latestSession), @@ -1036,7 +1037,7 @@ export async function retryCreateBackendSession( worktreeId: session.config.executionTarget.worktreeId, } : { kind: 'local' }, - workspaceId: session.workspaceId, + workspaceId: requireSessionOwningWorkspaceId(session), remoteConnectionId: session.remoteConnectionId, remoteSshHost: session.remoteSshHost, relationship: buildCreateSessionRelationship(session), diff --git a/src/web-ui/src/flow_chat/session-drivers/local/LocalSessionDriver.test.ts b/src/web-ui/src/flow_chat/session-drivers/local/LocalSessionDriver.test.ts index 6189b768da..be100d2bbd 100644 --- a/src/web-ui/src/flow_chat/session-drivers/local/LocalSessionDriver.test.ts +++ b/src/web-ui/src/flow_chat/session-drivers/local/LocalSessionDriver.test.ts @@ -5,9 +5,10 @@ import type { DialogTurn } from '../../types/flow-chat'; import { getActiveSurfaceScope } from '@/infrastructure/peer-device/deviceSurface'; import { consumeSubmittedMessageArrival } from '../../services/submittedMessagePresentation'; -const { mockStartAcpDialogTurn, mockStartAgenticDialogTurn, mockTransition, mockUpdateSessionMetadata, mockGetMode, mockUpdateMode } = vi.hoisted(() => ({ +const { mockStartAcpDialogTurn, mockStartAgenticDialogTurn, mockEnsureCoordinatorSession, mockTransition, mockUpdateSessionMetadata, mockGetMode, mockUpdateMode } = vi.hoisted(() => ({ mockStartAcpDialogTurn: vi.fn(), mockStartAgenticDialogTurn: vi.fn(), + mockEnsureCoordinatorSession: vi.fn(), mockTransition: vi.fn(), mockUpdateSessionMetadata: vi.fn(), mockGetMode: vi.fn(), @@ -21,6 +22,7 @@ vi.mock('@/infrastructure/api/service-api/ACPClientAPI', () => ({ vi.mock('@/infrastructure/api/service-api/AgentAPI', () => ({ agentAPI: { startDialogTurn: mockStartAgenticDialogTurn, + ensureCoordinatorSession: mockEnsureCoordinatorSession, getSessionPermissionMode: mockGetMode, updateSessionPermissionMode: mockUpdateMode, }, @@ -28,6 +30,16 @@ vi.mock('@/infrastructure/api/service-api/AgentAPI', () => ({ vi.mock('@/infrastructure/api/service-api/SessionAPI', () => ({ sessionAPI: {} })); vi.mock('@/infrastructure/api/service-api/WorktreeAPI', () => ({ worktreeAPI: {} })); +vi.mock('@/infrastructure/services/business/workspaceManager', () => ({ + workspaceManager: { + getState: () => ({ + openedWorkspaces: new Map([['project-workspace', { + id: 'project-workspace', rootPath: WORKSPACE_PATH, workspaceKind: 'normal', + }]]), + recentWorkspaces: [], + }), + }, +})); vi.mock('../../state-machine', () => ({ stateMachineManager: { @@ -240,3 +252,56 @@ describe('host queue submissions', () => { expect(tracker.hostAcceptedTurn).toBe(true); }); }); + +describe('worktree follow-up submissions', () => { + beforeEach(() => { + vi.clearAllMocks(); + mockEnsureCoordinatorSession.mockResolvedValue(undefined); + mockTransition.mockResolvedValue(true); + mockStartAgenticDialogTurn.mockResolvedValue(undefined); + queueMocks.submit.mockResolvedValue({ receipt: { status: 'queued' } }); + }); + + it.each([false, true])('prepares the owning project and sends the next turn with host queue support %s', async supportsQueue => { + queueMocks.supported.mockReturnValue(supportsQueue); + const { context, session } = createHarness([persistedTurn('turn-1', 0)]); + Object.assign(session, { + mode: 'Standard', + isHistorical: false, + historyState: 'ready', + workspaceId: 'worktree-not-in-catalog', + workspacePath: '/worktrees/task', + projectWorkspaceId: 'project-workspace', + projectWorkspacePath: WORKSPACE_PATH, + config: { + executionTarget: { kind: 'managedWorktree', worktreeId: 'worktree-1', rootPath: '/worktrees/task' }, + }, + }); + const tracker = { createdLocalTurnId: null, hostAcceptedTurn: false }; + + await localSessionDriver.ensureReady(context, SESSION_ID); + await localSessionDriver.startTurn(context, { + ...startTurnInput(session), acpClientId: undefined, currentAgentType: 'Standard', + }, tracker); + + expect(mockEnsureCoordinatorSession).toHaveBeenCalledWith({ + sessionId: SESSION_ID, workspaceId: 'project-workspace', includeInternal: false, + }); + if (supportsQueue) { + expect(queueMocks.submit).toHaveBeenCalledWith( + expect.objectContaining({ content: 'hello' }), expect.any(Object), undefined, + ); + expect(mockStartAgenticDialogTurn).not.toHaveBeenCalled(); + } else { + expect(mockStartAgenticDialogTurn).toHaveBeenCalledWith(expect.objectContaining({ + sessionId: SESSION_ID, + workspaceId: 'worktree-not-in-catalog', + workspacePath: '/worktrees/task', + projectWorkspacePath: WORKSPACE_PATH, + })); + } + expect(tracker.hostAcceptedTurn).toBe(true); + expect(session.workspacePath).toBe('/worktrees/task'); + expect(session.dialogTurns[0].id).toBe('turn-1'); + }); +});