diff --git a/cli/src/agent/runners/runAgentSession.test.ts b/cli/src/agent/runners/runAgentSession.test.ts index 4b3e302e..e2f34f1c 100644 --- a/cli/src/agent/runners/runAgentSession.test.ts +++ b/cli/src/agent/runners/runAgentSession.test.ts @@ -17,6 +17,7 @@ vi.mock('@/agent/sessionFactory', () => ({ onUserMessage: vi.fn((handler) => { harness.userMessageHandler = handler }), + onCancelQueuedMessage: vi.fn(), keepAlive: vi.fn(), sendSessionEvent: vi.fn(), sendAgentMessage: vi.fn(), diff --git a/cli/src/agent/runners/runAgentSession.ts b/cli/src/agent/runners/runAgentSession.ts index 654bcd85..4a0b0e9f 100644 --- a/cli/src/agent/runners/runAgentSession.ts +++ b/cli/src/agent/runners/runAgentSession.ts @@ -56,6 +56,12 @@ export async function runAgentSession(opts: { messageQueue.push(formattedText, {}, localId); }); + session.onCancelQueuedMessage((localId) => { + const removed = messageQueue.cancelByLocalId(localId); + logger.debug(`[agent] cancelByLocalId(${localId}): ${removed ? 'removed' : 'not found (best-effort)'}`); + return removed; + }); + let currentPermissionMode: SessionPermissionMode = opts.permissionMode ?? sessionInfo.permissionMode ?? 'default'; const backend: AgentBackend = AgentRegistry.create(opts.agentType); diff --git a/cli/src/api/apiSession.ts b/cli/src/api/apiSession.ts index 2ac631c4..c19253de 100644 --- a/cli/src/api/apiSession.ts +++ b/cli/src/api/apiSession.ts @@ -81,6 +81,7 @@ export class ApiSessionClient extends EventEmitter { private readonly socket: Socket private pendingMessages: { message: UserMessage; localId?: string }[] = [] private pendingMessageCallback: ((message: UserMessage, localId?: string) => void) | null = null + private cancelQueuedMessageCallback: ((localId: string) => boolean) | null = null private lastSeenMessageSeq: number | null = null private backfillInFlight: Promise | null = null private needsBackfill = false @@ -200,7 +201,7 @@ export class ApiSessionClient extends EventEmitter { this.terminalManager.close(payload.terminalId) })) - this.socket.on('update', (data: Update) => { + this.socket.on('update', (data: Update, ack?: (response: { removed: boolean }) => void) => { try { if (!data.body) return @@ -209,6 +210,14 @@ export class ApiSessionClient extends EventEmitter { return } + if (data.body.t === 'cancel-queued-message') { + const removed = (data.body.localId && this.cancelQueuedMessageCallback) + ? this.cancelQueuedMessageCallback(data.body.localId) + : false + ack?.({ removed }) + return + } + if (data.body.t === 'update-session') { if (data.body.metadata && data.body.metadata.version > this.metadataVersion) { const parsed = MetadataSchema.safeParse(data.body.metadata.value) @@ -253,6 +262,10 @@ export class ApiSessionClient extends EventEmitter { } } + onCancelQueuedMessage(callback: (localId: string) => boolean): void { + this.cancelQueuedMessageCallback = callback + } + private enqueueUserMessage(message: UserMessage, localId?: string): void { if (this.pendingMessageCallback) { this.pendingMessageCallback(message, localId) diff --git a/cli/src/claude/runClaude.ts b/cli/src/claude/runClaude.ts index 9874d833..32b3d80e 100644 --- a/cli/src/claude/runClaude.ts +++ b/cli/src/claude/runClaude.ts @@ -332,6 +332,12 @@ export async function runClaude(options: StartOptions = {}): Promise { logger.debugLargeJson('User message pushed to queue:', message) }); + session.onCancelQueuedMessage((localId) => { + const removed = messageQueue.cancelByLocalId(localId); + logger.debug(`[claude] cancelByLocalId(${localId}): ${removed ? 'removed' : 'not found (best-effort)'}`); + return removed; + }); + const resolvePermissionMode = (value: unknown): PermissionMode => { const parsed = PermissionModeSchema.safeParse(value); if (!parsed.success || !isPermissionModeAllowedForFlavor(parsed.data, 'claude')) { diff --git a/cli/src/codex/runCodex.ts b/cli/src/codex/runCodex.ts index 90423028..301ed5f3 100644 --- a/cli/src/codex/runCodex.ts +++ b/cli/src/codex/runCodex.ts @@ -205,6 +205,12 @@ export async function runCodex(opts: { }); }); + session.onCancelQueuedMessage((localId) => { + const removed = messageQueue.cancelByLocalId(localId); + logger.debug(`[codex] cancelByLocalId(${localId}): ${removed ? 'removed' : 'not found (best-effort)'}`); + return removed; + }); + const formatFailureReason = (message: string): string => { const maxLength = 200; if (message.length <= maxLength) { diff --git a/cli/src/cursor/runCursor.ts b/cli/src/cursor/runCursor.ts index 66ed2377..9de0dc4b 100644 --- a/cli/src/cursor/runCursor.ts +++ b/cli/src/cursor/runCursor.ts @@ -86,6 +86,12 @@ export async function runCursor(opts: { messageQueue.push(formattedText, enhancedMode, localId); }); + session.onCancelQueuedMessage((localId) => { + const removed = messageQueue.cancelByLocalId(localId); + logger.debug(`[cursor] cancelByLocalId(${localId}): ${removed ? 'removed' : 'not found (best-effort)'}`); + return removed; + }); + const resolvePermissionMode = (value: unknown): PermissionMode => { const parsed = PermissionModeSchema.safeParse(value); if (!parsed.success || !isPermissionModeAllowedForFlavor(parsed.data, 'cursor')) { diff --git a/cli/src/gemini/runGemini.test.ts b/cli/src/gemini/runGemini.test.ts index 21d0feec..b6ef8e58 100644 --- a/cli/src/gemini/runGemini.test.ts +++ b/cli/src/gemini/runGemini.test.ts @@ -14,6 +14,7 @@ const harness = vi.hoisted(() => ({ geminiLoopError: null as Error | null, session: { onUserMessage: vi.fn(), + onCancelQueuedMessage: vi.fn(), rpcHandlerManager: { registerHandler: vi.fn() } diff --git a/cli/src/gemini/runGemini.ts b/cli/src/gemini/runGemini.ts index 612859ff..b4b7360a 100644 --- a/cli/src/gemini/runGemini.ts +++ b/cli/src/gemini/runGemini.ts @@ -128,6 +128,12 @@ export async function runGemini(opts: { messageQueue.push(formattedText, mode, localId); }); + session.onCancelQueuedMessage((localId) => { + const removed = messageQueue.cancelByLocalId(localId); + logger.debug(`[gemini] cancelByLocalId(${localId}): ${removed ? 'removed' : 'not found (best-effort)'}`); + return removed; + }); + const resolvePermissionMode = (value: unknown): PermissionMode => { const parsed = PermissionModeSchema.safeParse(value); if (!parsed.success || !isPermissionModeAllowedForFlavor(parsed.data, 'gemini')) { diff --git a/cli/src/opencode/runOpencode.test.ts b/cli/src/opencode/runOpencode.test.ts index 35aba81d..6f94b960 100644 --- a/cli/src/opencode/runOpencode.test.ts +++ b/cli/src/opencode/runOpencode.test.ts @@ -14,6 +14,7 @@ const harness = vi.hoisted(() => ({ opencodeLoopError: null as Error | null, session: { onUserMessage: vi.fn(), + onCancelQueuedMessage: vi.fn(), rpcHandlerManager: { registerHandler: vi.fn() } diff --git a/cli/src/opencode/runOpencode.ts b/cli/src/opencode/runOpencode.ts index fd45b8b0..f8e0dc63 100644 --- a/cli/src/opencode/runOpencode.ts +++ b/cli/src/opencode/runOpencode.ts @@ -108,6 +108,12 @@ export async function runOpencode(opts: { messageQueue.push(formattedText, mode, localId); }); + session.onCancelQueuedMessage((localId) => { + const removed = messageQueue.cancelByLocalId(localId); + logger.debug(`[opencode] cancelByLocalId(${localId}): ${removed ? 'removed' : 'not found (best-effort)'}`); + return removed; + }); + const resolvePermissionMode = (value: unknown): PermissionMode => { const parsed = PermissionModeSchema.safeParse(value); if (!parsed.success || !isPermissionModeAllowedForFlavor(parsed.data, 'opencode')) { diff --git a/cli/src/utils/MessageQueue2.test.ts b/cli/src/utils/MessageQueue2.test.ts index a4b63472..8c03f7a6 100644 --- a/cli/src/utils/MessageQueue2.test.ts +++ b/cli/src/utils/MessageQueue2.test.ts @@ -489,6 +489,68 @@ describe('MessageQueue2', () => { expect(consumedCount).toBe(0); }); + describe('cancelByLocalId', () => { + it('should remove the message with matching localId and return true', () => { + const queue = new MessageQueue2(mode => mode); + queue.push('msg1', 'local', 'id-abc'); + queue.push('msg2', 'local', 'id-def'); + + const removed = queue.cancelByLocalId('id-abc'); + expect(removed).toBe(true); + expect(queue.size()).toBe(1); + expect(queue.queue[0].localId).toBe('id-def'); + }); + + it('should return false when localId is not found', () => { + const queue = new MessageQueue2(mode => mode); + queue.push('msg1', 'local', 'id-abc'); + + const removed = queue.cancelByLocalId('id-nonexistent'); + expect(removed).toBe(false); + expect(queue.size()).toBe(1); + }); + + it('should return false when queue is empty', () => { + const queue = new MessageQueue2(mode => mode); + const removed = queue.cancelByLocalId('id-abc'); + expect(removed).toBe(false); + }); + + it('should not remove a message without localId even if localId param matches empty string', () => { + const queue = new MessageQueue2(mode => mode); + queue.push('msg-no-localid', 'local'); // no localId + + const removed = queue.cancelByLocalId(''); + expect(removed).toBe(false); + expect(queue.size()).toBe(1); + }); + + it('should only remove the first matching localId when duplicates exist', () => { + const queue = new MessageQueue2(mode => mode); + queue.push('msg1', 'local', 'id-dup'); + queue.push('msg2', 'local', 'id-dup'); + + const removed = queue.cancelByLocalId('id-dup'); + expect(removed).toBe(true); + expect(queue.size()).toBe(1); + // msg2 still remains + expect(queue.queue[0].message).toBe('msg2'); + }); + + it('should not affect messages without localId when cancelling by id', () => { + const queue = new MessageQueue2(mode => mode); + queue.push('msg-no-id', 'local'); + queue.push('msg-with-id', 'local', 'target-id'); + queue.push('msg-no-id-2', 'local'); + + const removed = queue.cancelByLocalId('target-id'); + expect(removed).toBe(true); + expect(queue.size()).toBe(2); + expect(queue.queue[0].message).toBe('msg-no-id'); + expect(queue.queue[1].message).toBe('msg-no-id-2'); + }); + }); + it('should differentiate between pushImmediate and pushIsolateAndClear behavior', async () => { const queue = new MessageQueue2<{ type: string }>((mode) => mode.type); diff --git a/cli/src/utils/MessageQueue2.ts b/cli/src/utils/MessageQueue2.ts index ed4b5141..fb59ec16 100644 --- a/cli/src/utils/MessageQueue2.ts +++ b/cli/src/utils/MessageQueue2.ts @@ -182,6 +182,20 @@ export class MessageQueue2 { logger.debug(`[MessageQueue2] unshift() completed. Queue size: ${this.queue.length}`); } + /** + * Remove the first queued message that matches the given localId. + * Returns true if a message was removed, false if not found. + * Best-effort: if the CLI is offline when cancel is issued, the message + * may already have been collected for invocation and won't be found here. + */ + cancelByLocalId(localId: string): boolean { + if (!localId) return false; + const idx = this.queue.findIndex(item => item.localId === localId); + if (idx === -1) return false; + this.queue.splice(idx, 1); + return true; + } + /** * Reset the queue - clears all messages and resets to empty state */ diff --git a/hub/src/store/index.ts b/hub/src/store/index.ts index f8c08cde..ee97856c 100644 --- a/hub/src/store/index.ts +++ b/hub/src/store/index.ts @@ -16,6 +16,7 @@ export type { StoredUser, VersionedUpdateResult } from './types' +export type { CancelQueuedMessageResult, LookupQueuedMessageResult } from './messages' export { MachineStore } from './machineStore' export { MessageStore } from './messageStore' export { PushStore } from './pushStore' diff --git a/hub/src/store/messageStore.ts b/hub/src/store/messageStore.ts index 3dec8002..b7173caa 100644 --- a/hub/src/store/messageStore.ts +++ b/hub/src/store/messageStore.ts @@ -1,7 +1,7 @@ import type { Database } from 'bun:sqlite' import type { StoredMessage } from './types' -import { addMessage, getMessages, getMessagesAfter, getMessagesByPosition, getUninvokedLocalMessages, markMessagesInvoked, mergeSessionMessages } from './messages' +import { addMessage, cancelQueuedMessage, deleteQueuedMessageById, lookupQueuedMessage, getMessages, getMessagesAfter, getMessagesByPosition, getUninvokedLocalMessages, markMessagesInvoked, mergeSessionMessages, type CancelQueuedMessageResult, type LookupQueuedMessageResult } from './messages' export class MessageStore { private readonly db: Database @@ -30,6 +30,18 @@ export class MessageStore { return getUninvokedLocalMessages(this.db, sessionId) } + cancelQueuedMessage(sessionId: string, messageId: string): CancelQueuedMessageResult { + return cancelQueuedMessage(this.db, sessionId, messageId) + } + + lookupQueuedMessage(sessionId: string, messageId: string): LookupQueuedMessageResult { + return lookupQueuedMessage(this.db, sessionId, messageId) + } + + deleteQueuedMessageById(sessionId: string, messageId: string): void { + deleteQueuedMessageById(this.db, sessionId, messageId) + } + markMessagesInvoked(sessionId: string, localIds: string[], invokedAt: number): void { markMessagesInvoked(this.db, sessionId, localIds, invokedAt) } diff --git a/hub/src/store/messages.test.ts b/hub/src/store/messages.test.ts new file mode 100644 index 00000000..c7c53c37 --- /dev/null +++ b/hub/src/store/messages.test.ts @@ -0,0 +1,174 @@ +import { describe, expect, it } from 'bun:test' +import { Store } from './index' + +function makeStore(): Store { + return new Store(':memory:') +} + +function makeSession(store: Store, tag: string) { + return store.sessions.getOrCreateSession(tag, { path: `/tmp/${tag}` }, null, 'default') +} + +describe('cancelQueuedMessage', () => { + it('happy path: deletes queued message, returns status=cancelled with localId', () => { + const store = makeStore() + const session = makeSession(store, 'cancel-happy') + const msg = store.messages.addMessage(session.id, { role: 'user', content: { type: 'text', text: 'hello' } }, 'lid-1') + + const result = store.messages.cancelQueuedMessage(session.id, msg.id) + expect(result.status).toBe('cancelled') + if (result.status === 'cancelled') { + expect(result.localId).toBe('lid-1') + } + + // Row should be gone from uninvoked list + const remaining = store.messages.getUninvokedLocalMessages(session.id) + expect(remaining).toHaveLength(0) + }) + + it('already-invoked: returns status=invoked with full message row, row stays in DB', () => { + const store = makeStore() + const session = makeSession(store, 'cancel-already-invoked') + const content = { role: 'user', content: { type: 'text', text: 'hello' } } + const msg = store.messages.addMessage(session.id, content, 'lid-2') + + const invokedAt = Date.now() + // Simulate CLI invoke ack + store.messages.markMessagesInvoked(session.id, ['lid-2'], invokedAt) + + const result = store.messages.cancelQueuedMessage(session.id, msg.id) + expect(result.status).toBe('invoked') + + // Must include the invoked row so the web client can restore authoritative state + if (result.status === 'invoked') { + expect(result.message.id).toBe(msg.id) + expect(result.message.localId).toBe('lid-2') + expect(result.message.invokedAt).toBe(invokedAt) + } + + // Row still exists (with invoked_at set) + const messages = store.messages.getMessages(session.id) + expect(messages.some(m => m.id === msg.id)).toBe(true) + }) + + it('cancel × 2 idempotent: second call returns status=cancelled with localId=null (row gone)', () => { + const store = makeStore() + const session = makeSession(store, 'cancel-idempotent') + const msg = store.messages.addMessage(session.id, { role: 'user', content: { type: 'text', text: 'hello' } }, 'lid-3') + + const first = store.messages.cancelQueuedMessage(session.id, msg.id) + expect(first.status).toBe('cancelled') + if (first.status === 'cancelled') { + expect(first.localId).toBe('lid-3') + } + + const second = store.messages.cancelQueuedMessage(session.id, msg.id) + expect(second.status).toBe('cancelled') + if (second.status === 'cancelled') { + expect(second.localId).toBeNull() + } + }) + + it('non-existent messageId: returns status=cancelled with localId=null', () => { + const store = makeStore() + const session = makeSession(store, 'cancel-nonexistent') + + const result = store.messages.cancelQueuedMessage(session.id, 'nonexistent-id') + expect(result.status).toBe('cancelled') + if (result.status === 'cancelled') { + expect(result.localId).toBeNull() + } + }) + + it('wrong sessionId: returns status=cancelled with localId=null, message from other session untouched', () => { + const store = makeStore() + const sessionA = makeSession(store, 'cancel-session-a') + const sessionB = makeSession(store, 'cancel-session-b') + const msg = store.messages.addMessage(sessionA.id, { role: 'user', content: { type: 'text', text: 'hello' } }, 'lid-A') + + const result = store.messages.cancelQueuedMessage(sessionB.id, msg.id) + expect(result.status).toBe('cancelled') + if (result.status === 'cancelled') { + expect(result.localId).toBeNull() + } + + // Original message still exists + const remaining = store.messages.getUninvokedLocalMessages(sessionA.id) + expect(remaining).toHaveLength(1) + }) + + it('cancelled localId is propagated from the deleted row', () => { + const store = makeStore() + const session = makeSession(store, 'cancel-localid-propagate') + const msg = store.messages.addMessage(session.id, { role: 'user', content: { type: 'text', text: 'hello' } }, 'lid-propagate') + + const result = store.messages.cancelQueuedMessage(session.id, msg.id) + expect(result.status).toBe('cancelled') + if (result.status === 'cancelled') { + expect(result.localId).toBe('lid-propagate') + } + }) + + it('cancel by localId before server echo: localId match returns status=cancelled with localId', () => { + const store = makeStore() + const session = makeSession(store, 'cancel-by-localid') + // Simulate the optimistic row: server has stored it with local_id but web client + // still holds msg.id === localId (server echo not yet received). + const localId = 'local:pre-echo-id' + store.messages.addMessage(session.id, { role: 'user', content: { type: 'text', text: 'hello' } }, localId) + + // The web client passes localId as messageId (before server echo replaces it) + const result = store.messages.cancelQueuedMessage(session.id, localId) + expect(result.status).toBe('cancelled') + if (result.status === 'cancelled') { + expect(result.localId).toBe(localId) + } + + // Row should be gone + const remaining = store.messages.getUninvokedLocalMessages(session.id) + expect(remaining).toHaveLength(0) + }) + + it('cancel by localId × 2 idempotent: second call returns status=cancelled with localId=null', () => { + const store = makeStore() + const session = makeSession(store, 'cancel-by-localid-idempotent') + const localId = 'local:idem-id' + store.messages.addMessage(session.id, { role: 'user', content: { type: 'text', text: 'hello' } }, localId) + + const first = store.messages.cancelQueuedMessage(session.id, localId) + expect(first.status).toBe('cancelled') + if (first.status === 'cancelled') { + expect(first.localId).toBe(localId) + } + + // Second cancel by the same localId — row is already gone + const second = store.messages.cancelQueuedMessage(session.id, localId) + expect(second.status).toBe('cancelled') + if (second.status === 'cancelled') { + expect(second.localId).toBeNull() + } + }) + + it('cancel by localId when invoked: returns status=invoked with message row', () => { + const store = makeStore() + const session = makeSession(store, 'cancel-by-localid-invoked') + const localId = 'local:invoked-id' + const msg = store.messages.addMessage(session.id, { role: 'user', content: { type: 'text', text: 'hello' } }, localId) + + const invokedAt = Date.now() + store.messages.markMessagesInvoked(session.id, [localId], invokedAt) + + // Web client passes localId as messageId — should detect invoked_at IS NOT NULL + const result = store.messages.cancelQueuedMessage(session.id, localId) + expect(result.status).toBe('invoked') + if (result.status === 'invoked') { + expect(result.message.id).toBe(msg.id) + expect(result.message.localId).toBe(localId) + expect(result.message.invokedAt).toBe(invokedAt) + } + + // Row still exists + const messages = store.messages.getMessages(session.id) + expect(messages.some(m => m.id === msg.id)).toBe(true) + }) +}) diff --git a/hub/src/store/messages.ts b/hub/src/store/messages.ts index c315ea7e..aa00bcf0 100644 --- a/hub/src/store/messages.ts +++ b/hub/src/store/messages.ts @@ -164,6 +164,112 @@ export function getMaxSeq(db: Database, sessionId: string): number { return row?.maxSeq ?? 0 } +export type CancelQueuedMessageResult = + | { status: 'cancelled'; localId: string | null } + | { status: 'invoked'; message: StoredMessage } + +/** Delete a queued (invoked_at IS NULL) message by session + message id. + * + * Runs inside a transaction to eliminate the SELECT-then-DELETE race window. + * Returns a discriminated union so callers can distinguish two zero-delete cases: + * - 'cancelled': row was absent (already cancelled, or wrong id/session) — treat as success. + * - 'invoked': row exists but invoked_at IS NOT NULL (CLI consumed it first) — + * caller must revert any optimistic removal using the returned row, + * not a stale client-side snapshot, so invokedAt is authoritative. + * + * The invoked_at IS NULL guard ensures cancel and invoke are mutually exclusive at + * the DB level (first-write-wins, mirrors markMessagesInvoked). */ +export function cancelQueuedMessage( + db: Database, + sessionId: string, + messageId: string +): CancelQueuedMessageResult { + return db.transaction(() => { + // Accept either the server-assigned uuid (id) or the client localId. + // This handles the pre-echo window where the web client still holds + // msg.id === localId and passes that as the messageId parameter. + // Note: local_id = ? evaluates to NULL (no match) when local_id IS NULL, + // which is safe — messages without a localId are inserted with invoked_at set + // and are never queued, so they cannot reach this code path anyway. + const row = db.prepare(` + SELECT * FROM messages + WHERE session_id = ? AND (id = ? OR local_id = ?) + LIMIT 1 + `).get(sessionId, messageId, messageId) as DbMessageRow | undefined + + if (!row) { + // Row absent: already cancelled or wrong id — fold into 'cancelled' + return { status: 'cancelled' as const, localId: null } + } + + if (row.invoked_at !== null) { + // CLI already consumed this message before the cancel arrived. + // Return the full row so the web client can restore authoritative invoked state + // rather than reverting to a stale queued snapshot (invokedAt: null). + return { status: 'invoked' as const, message: toStoredMessage(row) } + } + + db.prepare(` + DELETE FROM messages + WHERE session_id = ? AND (id = ? OR local_id = ?) AND invoked_at IS NULL + `).run(sessionId, messageId, messageId) + + return { status: 'cancelled' as const, localId: row.local_id } + })() +} + +export type LookupQueuedMessageResult = + | { status: 'absent' } + | { status: 'invoked'; message: StoredMessage } + | { status: 'queued'; localId: string | null; resolvedId: string } + +/** Look up a queued message without deleting it. + * + * Returns one of three discriminated states: + * - 'absent': row not found (already cancelled or wrong id). + * - 'invoked': row exists but invoked_at IS NOT NULL (CLI consumed it first). + * - 'queued': row exists and is cancellable; resolvedId is the server-assigned uuid. + * + * Used by the service layer to inspect state before issuing a CLI ack round-trip. + * The actual DELETE (after CLI ack) is performed by deleteQueuedMessageById. */ +export function lookupQueuedMessage( + db: Database, + sessionId: string, + messageId: string +): LookupQueuedMessageResult { + const row = db.prepare(` + SELECT * FROM messages + WHERE session_id = ? AND (id = ? OR local_id = ?) + LIMIT 1 + `).get(sessionId, messageId, messageId) as DbMessageRow | undefined + + if (!row) { + return { status: 'absent' as const } + } + + if (row.invoked_at !== null) { + return { status: 'invoked' as const, message: toStoredMessage(row) } + } + + return { status: 'queued' as const, localId: row.local_id, resolvedId: row.id } +} + +/** Delete a queued (invoked_at IS NULL) message by id or local_id. + * + * This is the "confirmed DELETE" step after the service layer has received a + * CLI ack with removed:true. Uses the same first-write-wins guard as the + * original cancelQueuedMessage. */ +export function deleteQueuedMessageById( + db: Database, + sessionId: string, + messageId: string +): void { + db.prepare(` + DELETE FROM messages + WHERE session_id = ? AND (id = ? OR local_id = ?) AND invoked_at IS NULL + `).run(sessionId, messageId, messageId) +} + /** Mark messages as invoked at the given server timestamp. * Only updates rows whose local_id is in localIds. * First-write-wins: rows with a non-NULL invoked_at are not updated. A duplicate diff --git a/hub/src/sync/messageService.test.ts b/hub/src/sync/messageService.test.ts new file mode 100644 index 00000000..55ee0c85 --- /dev/null +++ b/hub/src/sync/messageService.test.ts @@ -0,0 +1,319 @@ +/** + * MessageService.cancelQueuedMessage race scenario tests + * + * Race-A: CLI ack returns { removed: true } → DB DELETE + status='cancelled' + * Race-B: CLI ack returns { removed: false } (already shift()-ed) → markMessagesInvoked + status='invoked' + * Race-C: CLI ack times out (500 ms) → markMessagesInvoked + status='invoked' + * Race-D (CLI offline): no CLI socket in room → immediate DELETE, message-cancelled emit, no ack call + * Race-E (partial ack): broadcast ack receives err + [{ removed: true }] → DELETE + status='cancelled' + */ +import { describe, expect, it } from 'bun:test' +import { MessageService } from './messageService' +import { Store } from '../store' +import type { Server } from 'socket.io' +import type { SyncEvent } from '@hapi/protocol/types' + +// --------------------------------------------------------------------------- +// Test helpers +// --------------------------------------------------------------------------- + +function makeStore(): Store { + return new Store(':memory:') +} + +function makeSession(store: Store, tag: string) { + return store.sessions.getOrCreateSession(tag, { path: `/tmp/${tag}` }, null, 'default') +} + +type AckCallback = (err: Error | null, responses: Array<{ removed: boolean }>) => void + +function makeIo(onEmit: (ack: AckCallback) => void, socketCount = 1): Server { + const broadcastRoom = { + timeout: (_ms: number) => ({ + emit: (_event: string, _data: unknown, callback: AckCallback) => { + onEmit(callback) + } + }), + emit: () => {} + } + + // Pre-built set reused on every rooms.get() call (socketCount=0 → undefined) + const socketSet = socketCount > 0 + ? new Set(Array.from({ length: socketCount }, (_, i) => `socket-${i}`)) + : undefined + + return { + of: (_ns: string) => ({ + to: (_room: string) => broadcastRoom, + adapter: { rooms: { get: (_roomName: string) => socketSet } } + }) + } as unknown as Server +} + +function makePublisher() { + const events: SyncEvent[] = [] + return { + emit: (event: SyncEvent) => { events.push(event) }, + events + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe('MessageService.cancelQueuedMessage race scenarios', () => { + describe('Race-A: CLI ack removed:true → DELETE + status=cancelled', () => { + it('returns cancelled and emits message-cancelled SSE after CLI confirms removal', async () => { + const store = makeStore() + const session = makeSession(store, 'race-a') + const msg = store.messages.addMessage( + session.id, + { role: 'user', content: { type: 'text', text: 'hello' } }, + 'local-a' + ) + + const publisher = makePublisher() + const io = makeIo((callback) => { + // CLI confirms it removed the item + callback(null, [{ removed: true }]) + }) + + const service = new MessageService(store, io, publisher as any) + const result = await service.cancelQueuedMessage(session.id, msg.id) + + expect(result.status).toBe('cancelled') + + // Row must be gone from the DB + const remaining = store.messages.getUninvokedLocalMessages(session.id) + expect(remaining).toHaveLength(0) + + // message-cancelled SSE must have been broadcast + const cancelled = publisher.events.find(e => e.type === 'message-cancelled') + expect(cancelled).toBeDefined() + + // No messages-consumed for cancelled path (row is deleted, not invoked) + const consumedCount = publisher.events.filter(e => e.type === 'messages-consumed').length + expect(consumedCount).toBe(0) + }) + }) + + describe('Race-B: CLI ack removed:false (already shift()-ed) → markMessagesInvoked + status=invoked', () => { + it('returns invoked with message row when CLI says item was already consumed', async () => { + const store = makeStore() + const session = makeSession(store, 'race-b') + const msg = store.messages.addMessage( + session.id, + { role: 'user', content: { type: 'text', text: 'hello' } }, + 'local-b' + ) + + const publisher = makePublisher() + const io = makeIo((callback) => { + // CLI already shifted the item before the cancel arrived + callback(null, [{ removed: false }]) + }) + + const service = new MessageService(store, io, publisher as any) + const result = await service.cancelQueuedMessage(session.id, msg.id) + + expect(result.status).toBe('invoked') + if (result.status === 'invoked') { + expect(result.message.id).toBe(msg.id) + expect(result.message.localId).toBe('local-b') + expect(result.message.invokedAt).not.toBeNull() + } + + // Row must still exist but now have invoked_at set + const rows = store.messages.getMessages(session.id) + const row = rows.find(r => r.id === msg.id) + expect(row).toBeDefined() + expect(row!.invokedAt).not.toBeNull() + + // No message-cancelled SSE should have been emitted + const cancelled = publisher.events.find(e => e.type === 'message-cancelled') + expect(cancelled).toBeUndefined() + + // messages-consumed SSE must be broadcast so other web clients clear the queued row + const consumed = publisher.events.find(e => e.type === 'messages-consumed') + expect(consumed).toBeDefined() + if (consumed?.type === 'messages-consumed') { + expect(consumed.sessionId).toBe(session.id) + expect(consumed.localIds).toEqual(['local-b']) + expect(typeof consumed.invokedAt).toBe('number') + } + + // messages-consumed must be emitted exactly once + const consumedCount = publisher.events.filter(e => e.type === 'messages-consumed').length + expect(consumedCount).toBe(1) + }) + }) + + describe('Race-C: CLI ack timeout → markMessagesInvoked + status=invoked', () => { + it('returns invoked with message row when CLI does not respond within timeout', async () => { + const store = makeStore() + const session = makeSession(store, 'race-c') + const msg = store.messages.addMessage( + session.id, + { role: 'user', content: { type: 'text', text: 'hello' } }, + 'local-c' + ) + + const publisher = makePublisher() + const io = makeIo((callback) => { + // Simulate timeout: socket.io passes an error as first arg + callback(new Error('operation has timed out'), []) + }) + + const service = new MessageService(store, io, publisher as any) + const result = await service.cancelQueuedMessage(session.id, msg.id) + + expect(result.status).toBe('invoked') + if (result.status === 'invoked') { + expect(result.message.id).toBe(msg.id) + expect(result.message.invokedAt).not.toBeNull() + } + + // Row must still exist with invoked_at stamped + const rows = store.messages.getMessages(session.id) + const row = rows.find(r => r.id === msg.id) + expect(row).toBeDefined() + expect(row!.invokedAt).not.toBeNull() + + // No message-cancelled SSE + const cancelled = publisher.events.find(e => e.type === 'message-cancelled') + expect(cancelled).toBeUndefined() + + // messages-consumed SSE must be broadcast so other web clients clear the queued row + const consumed = publisher.events.find(e => e.type === 'messages-consumed') + expect(consumed).toBeDefined() + if (consumed?.type === 'messages-consumed') { + expect(consumed.sessionId).toBe(session.id) + expect(consumed.localIds).toEqual(['local-c']) + expect(typeof consumed.invokedAt).toBe('number') + } + + // messages-consumed must be emitted exactly once + const consumedCount = publisher.events.filter(e => e.type === 'messages-consumed').length + expect(consumedCount).toBe(1) + }) + }) + + describe('Race-D: CLI offline (room socket count === 0) → immediate DELETE, no ack', () => { + it('returns cancelled and emits message-cancelled without calling ack when no CLI socket is connected', async () => { + const store = makeStore() + const session = makeSession(store, 'race-d-offline') + const msg = store.messages.addMessage( + session.id, + { role: 'user', content: { type: 'text', text: 'hello' } }, + 'local-offline' + ) + + let ackCalled = false + // socketCount=0 → adapter.rooms.get() returns undefined → cliCount = 0 + const io = makeIo(() => { ackCalled = true }, 0) + const publisher = makePublisher() + + const service = new MessageService(store, io, publisher as any) + const result = await service.cancelQueuedMessage(session.id, msg.id) + + // Hub must return cancelled immediately + expect(result.status).toBe('cancelled') + + // CLI ack must NOT have been called + expect(ackCalled).toBe(false) + + // Row must be gone from the DB (immediate DELETE) + const remaining = store.messages.getUninvokedLocalMessages(session.id) + expect(remaining).toHaveLength(0) + + // message-cancelled SSE must have been emitted with localId + const cancelled = publisher.events.find(e => e.type === 'message-cancelled') + expect(cancelled).toBeDefined() + if (cancelled?.type === 'message-cancelled') { + expect(cancelled.localId).toBe('local-offline') + } + + // No messages-consumed (row was deleted, not invoked) + const consumedCount = publisher.events.filter(e => e.type === 'messages-consumed').length + expect(consumedCount).toBe(0) + + // No invoked_at stamped (row deleted, not marked invoked) + const rows = store.messages.getMessages(session.id) + expect(rows.find(r => r.id === msg.id)).toBeUndefined() + }) + }) + + describe('existing store-level invoked guard (DB first-write-wins) still respected', () => { + it('returns invoked without contacting CLI when DB row already has invoked_at', async () => { + const store = makeStore() + const session = makeSession(store, 'race-d-already-invoked') + const msg = store.messages.addMessage( + session.id, + { role: 'user', content: { type: 'text', text: 'hello' } }, + 'local-d' + ) + + // DB row was already marked invoked (e.g. by a concurrent messages-consumed) + const invokedAt = Date.now() + store.messages.markMessagesInvoked(session.id, ['local-d'], invokedAt) + + let cliContacted = false + const io = makeIo(() => { cliContacted = true }) + const publisher = makePublisher() + + const service = new MessageService(store, io, publisher as any) + const result = await service.cancelQueuedMessage(session.id, msg.id) + + expect(result.status).toBe('invoked') + // CLI must NOT have been contacted — DB guard should short-circuit before ack + expect(cliContacted).toBe(false) + + if (result.status === 'invoked') { + expect(result.message.invokedAt).toBe(invokedAt) + } + + // DB guard path: messages-consumed was already published by the prior + // messages-consumed flow that set invoked_at. No additional emit here. + const consumedCount = publisher.events.filter(e => e.type === 'messages-consumed').length + expect(consumedCount).toBe(0) + }) + }) + + describe('Race-E: partial ack — broadcast callback receives err + [{ removed: true }]', () => { + it('returns cancelled and deletes row when at least one socket acked removal, even if err is set', async () => { + const store = makeStore() + const session = makeSession(store, 'race-e') + const msg = store.messages.addMessage( + session.id, + { role: 'user', content: { type: 'text', text: 'hello' } }, + 'local-e' + ) + + const publisher = makePublisher() + // Reconnect-overlap scenario: one socket timed out (err set by Socket.IO), + // but the live socket confirmed removal in responses. + const io = makeIo((callback) => { + callback(new Error('operation has timed out'), [{ removed: true }]) + }) + + const service = new MessageService(store, io, publisher as any) + const result = await service.cancelQueuedMessage(session.id, msg.id) + + // The live socket's ack must win — cancel is confirmed + expect(result.status).toBe('cancelled') + + // Row must be deleted + const remaining = store.messages.getUninvokedLocalMessages(session.id) + expect(remaining).toHaveLength(0) + + // message-cancelled SSE must have been emitted + const cancelled = publisher.events.find(e => e.type === 'message-cancelled') + expect(cancelled).toBeDefined() + + // No messages-consumed (row deleted, not invoked) + const consumedCount = publisher.events.filter(e => e.type === 'messages-consumed').length + expect(consumedCount).toBe(0) + }) + }) +}) diff --git a/hub/src/sync/messageService.ts b/hub/src/sync/messageService.ts index 0aba66a1..07425861 100644 --- a/hub/src/sync/messageService.ts +++ b/hub/src/sync/messageService.ts @@ -1,6 +1,7 @@ import type { AttachmentMetadata, DecryptedMessage } from '@hapi/protocol/types' import type { Server } from 'socket.io' -import type { Store } from '../store' +import { randomUUID } from 'node:crypto' +import type { Store, CancelQueuedMessageResult } from '../store' import { EventPublisher } from './eventPublisher' export class MessageService { @@ -135,6 +136,174 @@ export class MessageService { })) } + async cancelQueuedMessage( + sessionId: string, + messageId: string + ): Promise { + // Phase 1: look up the row WITHOUT deleting it. + // This lets us ask the CLI first and only DELETE if the CLI confirms removal. + const lookup = this.store.messages.lookupQueuedMessage(sessionId, messageId) + + if (lookup.status === 'absent') { + // Row not found — already cancelled or wrong id. + return { status: 'cancelled', localId: null } + } + + if (lookup.status === 'invoked') { + // DB row already has invoked_at — CLI consumed it before we arrived. + // Return the full invoked row so the web client can restore authoritative + // state (with correct invokedAt) instead of a stale queued snapshot. + return lookup + } + + // Phase 2: row is still queued. Ask the CLI whether it already shifted the item + // (race window between collectBatch() shift and messages-consumed ack). + const { localId, resolvedId } = lookup + + if (!localId) { + // No localId — row exists but has no cancel path; treat as cancelled. + this.store.messages.deleteQueuedMessageById(sessionId, resolvedId) + this.publisher.emit({ type: 'message-cancelled', sessionId, messageId }) + return { status: 'cancelled', localId: null } + } + + // Phase 2a: if no CLI socket is currently in the session room, the CLI is + // offline and there is nobody to ack with. Delete the row immediately so a + // later CLI reconnect cannot pick it up via seq-backfill and re-enqueue the + // cancelled message. + // + // TOCTOU note: deleteQueuedMessageById already has an invoked_at IS NULL guard, + // so if a CLI socket joins between the cliCount read and the DELETE and wins the + // race by calling markMessagesInvoked first, the DELETE becomes a no-op. + // We re-read the row after the delete to detect that case and handle it exactly + // like Race-B (ack returned removed:false). + const roomName = `session:${sessionId}` + const cliCount = this.io.of('/cli').adapter.rooms.get(roomName)?.size ?? 0 + if (cliCount === 0) { + this.store.messages.deleteQueuedMessageById(sessionId, resolvedId) + // Re-check: if CLI joined and invoked the message between our cliCount read + // and the DELETE, the delete was a no-op and the row now has invoked_at set. + const recheck = this.store.messages.lookupQueuedMessage(sessionId, resolvedId) + if (recheck.status === 'invoked') { + // CLI beat us — treat identically to Race-B (ack returned not-found). + this.publisher.emit({ + type: 'messages-consumed', + sessionId, + localIds: [localId], + invokedAt: recheck.message.invokedAt!, + }) + return recheck + } + // Row is gone (absent) — clean cancel. + this.publisher.emit({ + type: 'message-cancelled', + sessionId, + messageId, + localId, + }) + return { status: 'cancelled', localId } + } + + const ackResult = await this.requestCliCancelAck(sessionId, localId, messageId, 500) + + if (ackResult === 'not-found' || ackResult === 'timeout') { + // CLI could not remove the item — it was already shift()-ed or CLI is + // offline. Stamp invoked_at immediately so the message lands in the thread + // as 'sent' instead of disappearing. The agent's later assistant message + // (if it produced one) joins the same thread normally. + const invokedAt = Date.now() + try { + this.store.messages.markMessagesInvoked(sessionId, [localId], invokedAt) + } catch (err) { + console.error('cancelQueuedMessage: markMessagesInvoked failed', err) + // DB write failed — let the HTTP 500 surface to the caller. + throw err + } + // Notify all SSE subscribers (other open tabs) that this queued row is now + // invoked so they remove it from the floating bar. Without this emit, only + // the tab that sent the DELETE request learns about the status change via the + // HTTP response; every other subscriber keeps the row in the queued bar until + // a refresh or a later event. Mirrors the identical publish in the normal + // CLI-driven path (sessionHandlers.ts messages-consumed handler). + this.publisher.emit({ + type: 'messages-consumed', + sessionId, + localIds: [localId], + invokedAt, + }) + // Re-fetch the single row via lookupQueuedMessage to avoid the 200-row + // pagination cap of getMessages. After markMessagesInvoked the row will + // have invoked_at set, so lookupQueuedMessage returns status='invoked'. + const recheck = this.store.messages.lookupQueuedMessage(sessionId, localId) + if (recheck.status === 'invoked') { + return recheck + } + // Row absent from DB after markMessagesInvoked — edge case, treat as cancelled + return { status: 'cancelled', localId } + } + + // Phase 3: CLI confirmed removal. Now DELETE the DB row and broadcast SSE. + this.store.messages.deleteQueuedMessageById(sessionId, resolvedId) + this.publisher.emit({ + type: 'message-cancelled', + sessionId, + messageId + }) + + return { status: 'cancelled', localId } + } + + /** + * Ask the CLI (via socket.io ack) whether it removed the in-memory queue item. + * Returns 'removed', 'not-found', or 'timeout'. + * + * Re-uses the existing 'update' event channel with a cancel-queued-message body, + * matching the ack pattern already used by rpcGateway + * (socket.timeout(ms).emitWithAck / BroadcastOperator.timeout(ms).emit + ack cb). + */ + private requestCliCancelAck( + sessionId: string, + localId: string, + messageId: string, + timeoutMs: number + ): Promise<'removed' | 'not-found' | 'timeout'> { + return new Promise((resolve) => { + const room = this.io.of('/cli').to(`session:${sessionId}`) + // socket.io v4 BroadcastOperator: .timeout(ms).emit(event, data, ackCb) + // ack signature: (err: Error | null, responses: T[]) + room.timeout(timeoutMs).emit( + 'update', + { + id: randomUUID(), + seq: 0, + createdAt: Date.now(), + body: { + t: 'cancel-queued-message' as const, + sid: sessionId, + messageId, + localId + } + }, + (err: Error | null, responses: Array<{ removed: boolean }>) => { + // Check responses before err: in a reconnect overlap or any room with + // multiple CLI sockets, Socket.IO may set err (one socket timed out) + // while still delivering successful responses from the sockets that did + // ack. Any confirmed removal wins over the partial timeout. + const removed = responses?.some((r) => r.removed === true) ?? false + if (removed) { + resolve('removed') + return + } + if (err) { + resolve('timeout') + return + } + resolve('not-found') + } + ) + }) + } + async sendMessage( sessionId: string, payload: { diff --git a/hub/src/sync/syncEngine.ts b/hub/src/sync/syncEngine.ts index c3c59bd3..8fffc0b6 100644 --- a/hub/src/sync/syncEngine.ts +++ b/hub/src/sync/syncEngine.ts @@ -9,7 +9,7 @@ import type { CodexCollaborationMode, DecryptedMessage, PermissionMode, Session, SyncEvent } from '@hapi/protocol/types' import type { Server } from 'socket.io' -import type { Store } from '../store' +import type { Store, CancelQueuedMessageResult } from '../store' import type { RpcRegistry } from '../socket/rpcRegistry' import type { SSEManager } from '../sse/sseManager' import { EventPublisher, type SyncEventListener } from './eventPublisher' @@ -312,6 +312,13 @@ export class SyncEngine { this.sessionCache.markMessageQueued(sessionId) } + async cancelQueuedMessage( + sessionId: string, + messageId: string + ): Promise { + return this.messageService.cancelQueuedMessage(sessionId, messageId) + } + async approvePermission( sessionId: string, requestId: string, diff --git a/hub/src/web/routes/messages.ts b/hub/src/web/routes/messages.ts index 1ce8e026..fb1e298c 100644 --- a/hub/src/web/routes/messages.ts +++ b/hub/src/web/routes/messages.ts @@ -51,6 +51,23 @@ export function createMessagesRoutes(getSyncEngine: () => SyncEngine | null): Ho return c.json(engine.getMessagesPage(sessionId, { limit, beforeSeq })) }) + app.delete('/sessions/:id/messages/:messageId', async (c) => { + const engine = requireSyncEngine(c, getSyncEngine) + if (engine instanceof Response) { + return engine + } + + const sessionResult = requireSessionFromParam(c, engine) + if (sessionResult instanceof Response) { + return sessionResult + } + const sessionId = sessionResult.sessionId + const messageId = c.req.param('messageId') + + const result = await engine.cancelQueuedMessage(sessionId, messageId) + return c.json(result) + }) + app.post('/sessions/:id/messages', async (c) => { const engine = requireSyncEngine(c, getSyncEngine) if (engine instanceof Response) { diff --git a/shared/src/schemas.ts b/shared/src/schemas.ts index 9a4ca011..a13eae0a 100644 --- a/shared/src/schemas.ts +++ b/shared/src/schemas.ts @@ -240,6 +240,11 @@ export const SyncEventSchema = z.discriminatedUnion('type', [ localIds: z.array(z.string()), invokedAt: z.number().optional() }), + SessionChangedSchema.extend({ + type: z.literal('message-cancelled'), + messageId: z.string(), + localId: z.string().optional() + }), SessionEventBaseSchema.extend({ type: z.literal('heartbeat'), data: z.object({ @@ -256,3 +261,10 @@ export const SyncEventSchema = z.discriminatedUnion('type', [ ]) export type SyncEvent = z.infer + +export const CancelMessageResponseSchema = z.discriminatedUnion('status', [ + z.object({ status: z.literal('cancelled'), localId: z.string().nullable() }), + z.object({ status: z.literal('invoked'), message: DecryptedMessageSchema }), +]) + +export type CancelMessageResponse = z.infer diff --git a/shared/src/socket.ts b/shared/src/socket.ts index cdbe1992..159c1139 100644 --- a/shared/src/socket.ts +++ b/shared/src/socket.ts @@ -114,17 +114,32 @@ export const UpdateMachineBodySchema = z.object({ export type UpdateMachineBody = z.infer +export const UpdateCancelQueuedMessageBodySchema = z.object({ + t: z.literal('cancel-queued-message'), + sid: z.string(), + messageId: z.string(), + localId: z.string().optional() +}) + +export type UpdateCancelQueuedMessageBody = z.infer + +export const CancelQueuedMessageAckSchema = z.object({ + removed: z.boolean() +}) + +export type CancelQueuedMessageAck = z.infer + export const UpdateSchema = z.object({ id: z.string(), seq: z.number(), - body: z.union([UpdateNewMessageBodySchema, UpdateSessionBodySchema, UpdateMachineBodySchema]), + body: z.union([UpdateNewMessageBodySchema, UpdateSessionBodySchema, UpdateMachineBodySchema, UpdateCancelQueuedMessageBodySchema]), createdAt: z.number() }) export type Update = z.infer export interface ServerToClientEvents { - update: (data: Update) => void + update: (data: Update, ack?: (response: CancelQueuedMessageAck) => void) => void 'rpc-request': (data: { method: string; params: string }, callback: (response: string) => void) => void 'terminal:open': (data: TerminalOpenPayload) => void 'terminal:write': (data: TerminalWritePayload) => void diff --git a/web/src/api/client.ts b/web/src/api/client.ts index 2b12e81f..86d2c961 100644 --- a/web/src/api/client.ts +++ b/web/src/api/client.ts @@ -25,6 +25,7 @@ import type { SessionResponse, SessionsResponse } from '@/types/api' +import type { CancelMessageResponse } from '@hapi/protocol/schemas' type ApiClientOptions = { baseUrl?: string @@ -307,6 +308,14 @@ export class ApiClient { }) } + async cancelMessage(sessionId: string, messageId: string): Promise { + const response = await this.request( + `/api/sessions/${encodeURIComponent(sessionId)}/messages/${encodeURIComponent(messageId)}`, + { method: 'DELETE' } + ) + return response as CancelMessageResponse + } + async abortSession(sessionId: string): Promise { await this.request(`/api/sessions/${encodeURIComponent(sessionId)}/abort`, { method: 'POST', diff --git a/web/src/components/AssistantChat/QueuedMessagesBar.test.tsx b/web/src/components/AssistantChat/QueuedMessagesBar.test.tsx new file mode 100644 index 00000000..e60fc7b9 --- /dev/null +++ b/web/src/components/AssistantChat/QueuedMessagesBar.test.tsx @@ -0,0 +1,58 @@ +import { describe, expect, it } from 'vitest' +import { computeCanCancel } from './QueuedMessagesBar' + +/** + * Unit tests for computeCanCancel — the race guard that prevents sending + * DELETE before the hub has a row to delete (pre-server-echo scenario). + * + * Key invariant: useSendMessage.onMutate creates an optimistic message with + * { id: localId, localId } + * so id === localId until the server echo (message-received SSE) arrives and + * message-window-store replaces the row with the server-assigned UUID id. + * After that replace, id !== localId. + * + * canCancel = hasServerEcho && !isPending + */ +describe('computeCanCancel', () => { + describe('hasServerEcho detection', () => { + it('is false when id === localId (purely optimistic, no server echo)', () => { + // useSendMessage.onMutate sets id = localId before POST /messages completes. + const localId = 'local-abc-123' + expect(computeCanCancel({ id: localId, localId, isPending: false })).toBe(false) + }) + + it('is true when id !== localId (server echo replaced id with server UUID)', () => { + const localId = 'local-abc-123' + const serverId = 'server-uuid-456' + expect(computeCanCancel({ id: serverId, localId, isPending: false })).toBe(true) + }) + + it('is true when localId is undefined/null (server-only row, no local tracking)', () => { + // Rows from server-loaded history have no localId — treat as already echoed. + expect(computeCanCancel({ id: 'server-uuid-789', localId: undefined, isPending: false })).toBe(true) + expect(computeCanCancel({ id: 'server-uuid-789', localId: null, isPending: false })).toBe(true) + }) + }) + + describe('isPending guard', () => { + it('is false when a cancel mutation is already in-flight, even with server echo', () => { + const localId = 'local-abc-123' + const serverId = 'server-uuid-456' + expect(computeCanCancel({ id: serverId, localId, isPending: true })).toBe(false) + }) + + it('is false when purely optimistic AND isPending', () => { + const localId = 'local-abc-123' + expect(computeCanCancel({ id: localId, localId, isPending: true })).toBe(false) + }) + }) + + describe('combined conditions', () => { + it('is true only when server echo received AND no in-flight cancel', () => { + const localId = 'local-abc-123' + const serverId = 'server-uuid-456' + // The normal case: user can click ✕ or ✎ + expect(computeCanCancel({ id: serverId, localId, isPending: false })).toBe(true) + }) + }) +}) diff --git a/web/src/components/AssistantChat/QueuedMessagesBar.tsx b/web/src/components/AssistantChat/QueuedMessagesBar.tsx index d2fc9846..bc036e6c 100644 --- a/web/src/components/AssistantChat/QueuedMessagesBar.tsx +++ b/web/src/components/AssistantChat/QueuedMessagesBar.tsx @@ -1,9 +1,12 @@ +import { useAssistantApi } from '@assistant-ui/react' import { useCallback, useSyncExternalStore } from 'react' +import type { ApiClient } from '@/api/client' import { getMessageWindowState, subscribeMessageWindow } from '@/lib/message-window-store' import { isQueuedForInvocation } from '@/lib/messages' import { EMPTY_STATE } from '@/hooks/queries/useMessages' import { normalizeDecryptedMessage } from '@/chat/normalize' import type { DecryptedMessage } from '@/types/api' +import { useCancelQueuedMessage } from '@/hooks/mutations/useCancelQueuedMessage' function ClockIcon() { return ( @@ -63,13 +66,44 @@ function getTextFromMessage(msg: DecryptedMessage): string { } /** - * Floating bar above the composer showing queued (pending invocation) messages. - * Disappears automatically when all queued messages are invoked or consumed. + * Determines whether the user can cancel or edit a queued message. * - * TODO PR 2: add cancel/edit buttons per item. + * Two conditions must both be true: + * 1. hasServerEcho: the hub has persisted the row. + * useSendMessage.onMutate creates { id: localId, localId } before POST /messages + * completes. Only after the server echo (message-received SSE) does the store + * replace the row with a server-assigned UUID id, making id !== localId. + * Sending DELETE before that echo would find no row in the hub and return + * cancelled/localId:null; the original POST could then still insert and broadcast + * the message, letting a canceled message reappear and be invoked. + * 2. !isPending: no cancel mutation is already in-flight. + * + * @internal Exported for unit testing. */ -export function QueuedMessagesBar({ sessionId }: { sessionId: string }) { +export function computeCanCancel({ + id, + localId, + isPending, +}: { + id: string + localId: string | null | undefined + isPending: boolean +}): boolean { + const hasServerEcho = localId ? id !== localId : true + return hasServerEcho && !isPending +} + +/** + * Floating bar above the composer showing queued (pending invocation) messages. + * Each item has an edit button (✎) and a cancel button (✕). + * + * Edit = client-side cancel + prefill composer with message text (Codex dialect). + * Cancel = DELETE /sessions/:id/messages/:messageId with optimistic removal. + */ +export function QueuedMessagesBar({ sessionId, api }: { sessionId: string; api: ApiClient | null }) { const queued = useQueuedMessages(sessionId) + const assistantApi = useAssistantApi() + const cancelMutation = useCancelQueuedMessage(api) if (queued.length === 0) { return null @@ -92,13 +126,100 @@ export function QueuedMessagesBar({ sessionId }: { sessionId: string }) { > {queued.map((msg) => { const text = getTextFromMessage(msg) + const localId = msg.localId ?? msg.id + const isPending = cancelMutation.isPending && cancelMutation.variables?.localId === localId + const canCancel = computeCanCancel({ id: msg.id, localId: msg.localId, isPending }) + + const handleCancel = () => { + if (!canCancel) return + cancelMutation.mutate({ + sessionId, + messageId: msg.id, + localId, + snapshot: msg, + }) + } + + const handleEdit = () => { + if (!canCancel) return + // Edit = cancel + prefill composer (Codex dialect: no separate edit mode). + cancelMutation.mutate( + { + sessionId, + messageId: msg.id, + localId, + snapshot: msg, + }, + { + onSuccess: (result) => { + // Race guard: if the agent already consumed this message, skip prefill. + // The hook's own onSuccess already reverted the optimistic removal. + if (result.status === 'invoked') return + // Only prefill if text is available; attachment-only rows get empty string. + const prefillText = text + if (prefillText) { + assistantApi.composer().setText(prefillText) + } + }, + } + ) + } + return (
  • - {text} - {/* TODO PR 2: cancel/edit buttons */} + + {text} + +
    + + +
    • ) })} diff --git a/web/src/components/SessionChat.tsx b/web/src/components/SessionChat.tsx index a6fe892b..5561b05b 100644 --- a/web/src/components/SessionChat.tsx +++ b/web/src/components/SessionChat.tsx @@ -454,7 +454,7 @@ export function SessionChat(props: { ) : null}
    - +
    { + if (!api) { + throw new Error('API unavailable') + } + return api.cancelMessage(input.sessionId, input.messageId) + }, + onMutate: (input) => { + // Optimistic: remove from the floating bar immediately. + removeOptimisticMessage(input.sessionId, input.localId) + }, + onSuccess: (result, input) => { + if (result.status === 'invoked') { + // Race: CLI consumed this message before cancel arrived. + // Restore using the server-validated invoked row so invokedAt is correct. + // Without this, messages-consumed SSE was a no-op (web row was missing) + // so the chip would be stuck as queued forever. + appendOptimisticMessage(input.sessionId, { + id: result.message.id, + seq: result.message.seq, + localId: result.message.localId, + content: result.message.content, + createdAt: result.message.createdAt, + invokedAt: result.message.invokedAt, + status: 'sent', + }) + } + // status === 'cancelled': optimistic removal stands — nothing extra to do. + }, + onError: (_error, input) => { + // Revert: put the message back so it re-appears in the bar. + appendOptimisticMessage(input.sessionId, input.snapshot) + haptic.notification('error') + }, + }) + + return mutation +} diff --git a/web/src/hooks/useSSE.ts b/web/src/hooks/useSSE.ts index 1899edb3..63cf5671 100644 --- a/web/src/hooks/useSSE.ts +++ b/web/src/hooks/useSSE.ts @@ -11,7 +11,7 @@ import type { SyncEvent } from '@/types/api' import { queryKeys } from '@/lib/query-keys' -import { clearMessageWindow, getMessageWindowState, ingestIncomingMessages, markMessagesConsumed, updateMessageStatus } from '@/lib/message-window-store' +import { clearMessageWindow, getMessageWindowState, ingestIncomingMessages, markMessagesConsumed, removeOptimisticMessage, updateMessageStatus } from '@/lib/message-window-store' type SSESubscription = { all?: boolean @@ -501,6 +501,12 @@ export function useSSE(options: { markMessagesConsumed(event.sessionId, event.localIds, event.invokedAt) } + if (event.type === 'message-cancelled') { + // Remove the cancelled message from the store. If the local + // optimistic removal already cleared it, this is a no-op. + removeOptimisticMessage(event.sessionId, event.messageId) + } + if (event.type === 'message-received') { ingestIncomingMessages(event.sessionId, [event.message]) } diff --git a/web/src/lib/message-window-store.test.ts b/web/src/lib/message-window-store.test.ts index 96f16bf9..089076c4 100644 --- a/web/src/lib/message-window-store.test.ts +++ b/web/src/lib/message-window-store.test.ts @@ -6,10 +6,26 @@ import { getMessageWindowState, ingestIncomingMessages, markMessagesConsumed, + removeOptimisticMessage, updateMessageStatus, } from '@/lib/message-window-store' -const SESSION_ID = 'session-message-window-store-test' +function makeMsg(overrides: Partial = {}): DecryptedMessage { + const id = overrides.id ?? 'msg-1' + return { + id, + seq: null, + localId: overrides.localId ?? id, + content: { + role: 'user', + content: { type: 'text', text: 'hello' } + }, + createdAt: Date.now(), + invokedAt: null, + status: 'queued', + ...overrides, + } +} function makeUserMessage(props: { id: string @@ -35,7 +51,81 @@ function makeUserMessage(props: { } as DecryptedMessage } +describe('removeOptimisticMessage', () => { + const SESSION = 'test-session-remove' + + afterEach(() => { + clearMessageWindow(SESSION) + }) + + it('removes a message matched by localId from the messages list', () => { + const msg = makeMsg({ id: 'msg-a', localId: 'local-a' }) + appendOptimisticMessage(SESSION, msg) + + removeOptimisticMessage(SESSION, 'local-a') + + const state = getMessageWindowState(SESSION) + expect(state.messages.find((m) => m.id === 'msg-a')).toBeUndefined() + }) + + it('removes a message matched by id (when localId equals id)', () => { + const msg = makeMsg({ id: 'msg-b', localId: 'msg-b' }) + appendOptimisticMessage(SESSION, msg) + + removeOptimisticMessage(SESSION, 'msg-b') + + const state = getMessageWindowState(SESSION) + expect(state.messages).toHaveLength(0) + }) + + it('is a no-op when localId does not match any message', () => { + const msg = makeMsg({ id: 'msg-c', localId: 'local-c' }) + appendOptimisticMessage(SESSION, msg) + + removeOptimisticMessage(SESSION, 'nonexistent') + + const state = getMessageWindowState(SESSION) + expect(state.messages).toHaveLength(1) + }) + + it('is a no-op when called with an empty string', () => { + const msg = makeMsg({ id: 'msg-d', localId: 'local-d' }) + appendOptimisticMessage(SESSION, msg) + + removeOptimisticMessage(SESSION, '') + + const state = getMessageWindowState(SESSION) + expect(state.messages).toHaveLength(1) + }) + + it('does not remove other messages when removing one', () => { + const msgA = makeMsg({ id: 'msg-e1', localId: 'local-e1' }) + const msgB = makeMsg({ id: 'msg-e2', localId: 'local-e2' }) + appendOptimisticMessage(SESSION, msgA) + appendOptimisticMessage(SESSION, msgB) + + removeOptimisticMessage(SESSION, 'local-e1') + + const state = getMessageWindowState(SESSION) + expect(state.messages.find((m) => m.id === 'msg-e1')).toBeUndefined() + expect(state.messages.find((m) => m.id === 'msg-e2')).toBeDefined() + }) + + it('is idempotent: second call is a no-op', () => { + const msg = makeMsg({ id: 'msg-f', localId: 'local-f' }) + appendOptimisticMessage(SESSION, msg) + + removeOptimisticMessage(SESSION, 'local-f') + removeOptimisticMessage(SESSION, 'local-f') + + const state = getMessageWindowState(SESSION) + expect(state.messages).toHaveLength(0) + }) +}) + describe('message-window-store status updates', () => { + const SESSION_ID = 'session-message-window-store-test' + afterEach(() => { clearMessageWindow(SESSION_ID) }) diff --git a/web/src/lib/message-window-store.ts b/web/src/lib/message-window-store.ts index 8f5ced24..0f32be07 100644 --- a/web/src/lib/message-window-store.ts +++ b/web/src/lib/message-window-store.ts @@ -620,6 +620,37 @@ export function updateMessageStatus(sessionId: string, localId: string, status: }) } +/** Remove an optimistic (not-yet-confirmed) message by its localId or server id. + * Used by the cancel affordance: optimistically drop the row immediately so the + * floating bar clears before the DELETE /messages/:id round-trip completes. If + * the request fails, the caller is responsible for re-inserting the row (e.g. + * via ingestIncomingMessages). Matches against both `localId` and `id` so that + * rows loaded from the server (which may have a stable uuid `id` + a localId) are + * also handled. + */ +export function removeOptimisticMessage(sessionId: string, localId: string): void { + if (!localId) return + updateState(sessionId, (prev) => { + let changed = false + const filterList = (list: DecryptedMessage[]) => { + const next = list.filter((message) => { + const matchesLocalId = message.localId === localId + const matchesId = message.id === localId + if (matchesLocalId || matchesId) { + changed = true + return false + } + return true + }) + return next + } + const messages = filterList(prev.messages) + const pending = filterList(prev.pending) + if (!changed) return prev + return buildState(prev, { messages, pending }) + }, true) +} + /** Transition the queued messages whose localIds match to 'sent' and record invokedAt. * Driven by the CLI ack (messages-consumed). Unmatched messages remain queued. * Also handles server-loaded messages (status=undefined) that have a matching localId.