Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 36 additions & 3 deletions apps/desktop/e2e/correction-session-switch.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,17 @@ async function send(page: Page, text: string): Promise<void> {
await page.keyboard.press('Enter')
}

async function steer(page: Page, text: string): Promise<void> {
const composer = page.locator('[contenteditable="true"]').first()
const primary = page.locator('[data-slot="composer-root"] button[type="submit"]')

await composer.waitFor({ state: 'visible', timeout: 15_000 })
await composer.click()
await composer.type(text, { delay: 5 })
await expect(primary).toHaveAttribute('aria-label', /Steer/)
await primary.click()
}

async function waitForTranscriptText(page: Page, text: string): Promise<void> {
await page.waitForFunction(
(expected: string) => (document.querySelector('[data-slot="aui_thread-viewport"]')?.textContent ?? '').includes(expected),
Expand Down Expand Up @@ -61,6 +72,17 @@ async function transcriptTextOrder(page: Page): Promise<string[]> {
})
}

async function transcriptMessageOrder(page: Page): Promise<string[]> {
return page.evaluate(() => {
const viewport = document.querySelector('[data-slot="aui_thread-viewport"]')
if (!viewport) return []

return Array.from(viewport.querySelectorAll<HTMLElement>('[data-role="user"], [data-role="assistant"]'))
.map(message => message.textContent?.trim() ?? '')
.filter(Boolean)
})
}

async function openFreshDraft(page: Page, priorSessionText: string): Promise<void> {
await page.locator('[data-slot="sidebar"] button[aria-label="New session"]').first().click()
await page.waitForFunction(
Expand All @@ -87,6 +109,16 @@ function relevantOrder(messages: string[]): string[] {
return messages.filter(message => message.includes(ORIGINAL_PROMPT) || message.includes(CORRECTION))
}

function steerTurnOrder(messages: string[]): string[] {
return messages.flatMap(message => {
if (message.includes(ORIGINAL_PROMPT)) return [ORIGINAL_PROMPT]
if (message.includes(CORRECTION)) return [CORRECTION]
if (message.includes(CORRECTED_REPLY)) return [CORRECTED_REPLY]

return []
})
}

test.describe('correction session switch', () => {
let fixture: MockBackendFixture | null = null

Expand All @@ -113,9 +145,9 @@ test.describe('correction session switch', () => {
await waitForTranscriptText(page, TOOL_STARTED)
await waitForTranscriptText(page, ORIGINAL_PROMPT)

// The historical session redirected while a foreground terminal task was
// running. Enter records the accepted correction at the next tool boundary.
await send(page, CORRECTION)
// The historical session redirects while a foreground terminal task is
// running. Use the visible Steer action to cover the real composer path.
await steer(page, CORRECTION)
await waitForTranscriptText(page, CORRECTION)

const orderBeforeSwitch = relevantOrder(await transcriptTextOrder(page))
Expand All @@ -136,5 +168,6 @@ test.describe('correction session switch', () => {
expect(await textNodeOccurrences(page, CORRECTION)).toBe(1)

await waitForTranscriptText(page, CORRECTED_REPLY)
expect(steerTurnOrder(await transcriptMessageOrder(page))).toEqual([ORIGINAL_PROMPT, CORRECTION, CORRECTED_REPLY])
})
})
49 changes: 49 additions & 0 deletions apps/desktop/e2e/queue-turn-boundary.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import { MOCK_REPLY } from './mock-server'

const ACTIVE_PROMPT = 'E2E_QUEUE_TURN_BOUNDARY_ACTIVE'
const QUEUED_PROMPT = 'E2E_QUEUE_TURN_BOUNDARY_QUEUED'
const STEER_PROMPT = 'E2E_STEER_TURN_BOUNDARY_CORRECTION'

async function send(page: Page, text: string): Promise<void> {
const composer = page.locator('[contenteditable="true"]').first()
Expand All @@ -23,6 +24,37 @@ async function send(page: Page, text: string): Promise<void> {
await page.keyboard.press('Enter')
}

async function steer(page: Page, text: string): Promise<void> {
const composer = page.locator('[contenteditable="true"]').first()
const primary = page.locator('[data-slot="composer-root"] button[type="submit"]')

await composer.click()
await composer.type(text, { delay: 5 })
await expect(primary).toHaveAttribute('aria-label', /Steer/)
await primary.click()
}

async function transcriptMessageOrder(page: Page): Promise<string[]> {
return page.evaluate(() => {
const viewport = document.querySelector('[data-slot="aui_thread-viewport"]')
if (!viewport) return []

return Array.from(viewport.querySelectorAll<HTMLElement>('[data-role="user"], [data-role="assistant"]'))
.map(message => message.textContent?.trim() ?? '')
.filter(Boolean)
})
}

function steerTurnOrder(messages: string[]): string[] {
return messages.flatMap(message => {
if (message.includes(ACTIVE_PROMPT)) return [ACTIVE_PROMPT]
if (message.includes(STEER_PROMPT)) return [STEER_PROMPT]
if (message.includes(MOCK_REPLY)) return [MOCK_REPLY]

return []
})
}

test.describe('queued prompt turn boundary', () => {
let fixture: MockBackendFixture | null = null

Expand Down Expand Up @@ -66,4 +98,21 @@ test.describe('queued prompt turn boundary', () => {
)
await expect.poll(() => mock.receivedPrompts.filter(prompt => prompt === QUEUED_PROMPT)).toHaveLength(1)
})

test('places a steer prompt before the reply it redirects', async () => {
const { mock, page } = fixture!

await send(page, ACTIVE_PROMPT)
await mock.waitForHeldStream()
await steer(page, STEER_PROMPT)
mock.releaseHeldStream()

await page.waitForFunction(
expected => (document.querySelector('[data-slot="aui_thread-viewport"]')?.textContent ?? '').includes(expected),
MOCK_REPLY,
{ timeout: 60_000 }
)

expect(steerTurnOrder(await transcriptMessageOrder(page))).toEqual([ACTIVE_PROMPT, STEER_PROMPT, MOCK_REPLY])
})
})
88 changes: 70 additions & 18 deletions apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -220,7 +220,13 @@ export function usePromptActions({
const copy = t.desktop

const appendSessionTextMessage = useCallback(
(sessionId: string, role: ChatMessage['role'], text: string, storedSessionId?: string | null) => {
(
sessionId: string,
role: ChatMessage['role'],
text: string,
storedSessionId?: string | null,
options: { insertBeforeActiveReply?: boolean } = {}
) => {
// Strip ANSI: slash-command output from the backend worker carries SGR
// color codes (e.g. "Unknown command" in red). The ESC byte is invisible
// in the chat panel, so without this the `[1;31m…[0m` payload leaks as
Expand All @@ -231,21 +237,33 @@ export function usePromptActions({
return
}

const messageId = `${role}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`

updateSessionState(
sessionId,
state => ({
...state,
messages: [
...state.messages,
{
id: `${role}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`,
role,
parts: [textPart(body)]
}
]
}),
state => {
const message: ChatMessage = {
id: messageId,
role,
parts: [textPart(body)]
}
const streamIndex = options.insertBeforeActiveReply && state.streamId
? state.messages.findIndex(candidate => candidate.id === state.streamId)
: -1
const lastAssistantIndex = options.insertBeforeActiveReply
? state.messages.map(candidate => candidate.role).lastIndexOf('assistant')
: -1
const insertionIndex = streamIndex >= 0 ? streamIndex : lastAssistantIndex
const messages = insertionIndex >= 0
? [...state.messages.slice(0, insertionIndex), message, ...state.messages.slice(insertionIndex)]
: [...state.messages, message]

return { ...state, messages }
},
storedSessionId ?? selectedStoredSessionIdRef.current
)

return messageId
},
[selectedStoredSessionIdRef, updateSessionState]
)
Expand Down Expand Up @@ -620,15 +638,49 @@ export function usePromptActions({
// message after the interrupted checkpoint, matching the durable core
// transcript rather than a system note that changes role after reload.
const send = async (id: string): Promise<boolean> => {
const result = await requestGateway<SessionRedirectResponse>('session.redirect', { session_id: id, text })
// Redirect aborts the model request, so the completion event can race
// its RPC response. Insert before the live reply *before* awaiting the
// gateway; appending after the response leaves the correction below a
// reply that the redirect has already replaced.
const messageId = appendSessionTextMessage(id, 'user', text, undefined, { insertBeforeActiveReply: true })
const discardOptimisticMessage = () =>
updateSessionState(id, state => ({
...state,
messages: state.messages.filter(message => message.id !== messageId)
}))
const moveOptimisticMessageToEnd = () =>
updateSessionState(id, state => {
const message = state.messages.find(candidate => candidate.id === messageId)

return message
? { ...state, messages: [...state.messages.filter(candidate => candidate.id !== messageId), message] }
: state
})

if (result?.status === 'redirected' || result?.status === 'queued') {
triggerHaptic('submit')
appendSessionTextMessage(id, 'user', text)
try {
const result = await requestGateway<SessionRedirectResponse>('session.redirect', { session_id: id, text })

if (result?.status === 'redirected') {
triggerHaptic('submit')

return true
}

return true
if (result?.status === 'queued') {
// Build-window redirects become the next turn, not part of the
// active reply, so retain the optimistic row at the tail.
moveOptimisticMessageToEnd()
triggerHaptic('submit')

return true
}
} catch (err) {
discardOptimisticMessage()
throw err
}

discardOptimisticMessage()

return false
}

Expand Down Expand Up @@ -661,7 +713,7 @@ export function usePromptActions({

return false
},
[activeSessionId, activeSessionIdRef, appendSessionTextMessage, requestGateway, selectedStoredSessionIdRef]
[activeSessionId, activeSessionIdRef, appendSessionTextMessage, requestGateway, selectedStoredSessionIdRef, updateSessionState]
)

const reloadFromMessage = useCallback(
Expand Down
Loading