From 4a9e363cd530797069aef365a092ae56a3b58f45 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Mon, 13 Jul 2026 14:39:09 +0200 Subject: [PATCH] feat(cloud-agent-next): discover and send remote slash commands Add end-to-end remote slash command support across the CLI live transport, web composer, and mobile chat composer. - CLI live transport discovers commands via `list_commands` and publishes `commands.available` events, with per-owner in-flight dedupe, owner-change staleness guards, and catalog size validation. - Web session manager ignores stale command catalogs from replaced sessions; `useSlashCommandSets` selects commands by session type instead of always falling back to the pinned default catalog. - Mobile chat composer surfaces slash command suggestions and routes structured command submissions through the session manager. - UserConnectionDO treats `list_commands` like `list_models` for pending-request dedupe and 512 KiB response size enforcement. --- .../chat-composer-slash-commands.test.ts | 64 ++++ .../agents/chat-composer-slash-commands.ts | 54 +++ .../src/components/agents/chat-composer.tsx | 53 ++- .../agents/session-detail-content.tsx | 18 + .../agents/slash-command-suggestions.tsx | 44 +++ apps/web/src/hooks/slash-command-selection.ts | 26 ++ .../web/src/hooks/useSlashCommandSets.test.ts | 42 +++ apps/web/src/hooks/useSlashCommandSets.ts | 26 +- .../cli-live-transport.test.ts | 352 ++++++++++++++++-- .../lib/cloud-agent-sdk/cli-live-transport.ts | 140 ++++++- .../remote-command-catalog.test.ts | 112 ++++++ apps/web/src/lib/cloud-agent-sdk/schemas.ts | 39 ++ .../cloud-agent-sdk/session-manager.test.ts | 17 + .../lib/cloud-agent-sdk/session-manager.ts | 1 + .../src/dos/UserConnectionDO.test.ts | 64 ++++ .../src/dos/UserConnectionDO.ts | 32 +- 16 files changed, 1028 insertions(+), 56 deletions(-) create mode 100644 apps/mobile/src/components/agents/chat-composer-slash-commands.test.ts create mode 100644 apps/mobile/src/components/agents/chat-composer-slash-commands.ts create mode 100644 apps/mobile/src/components/agents/slash-command-suggestions.tsx create mode 100644 apps/web/src/hooks/slash-command-selection.ts create mode 100644 apps/web/src/hooks/useSlashCommandSets.test.ts create mode 100644 apps/web/src/lib/cloud-agent-sdk/remote-command-catalog.test.ts diff --git a/apps/mobile/src/components/agents/chat-composer-slash-commands.test.ts b/apps/mobile/src/components/agents/chat-composer-slash-commands.test.ts new file mode 100644 index 0000000000..ffb5706eba --- /dev/null +++ b/apps/mobile/src/components/agents/chat-composer-slash-commands.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, it } from 'vitest'; + +import { + getSlashCommandCandidate, + getSlashCommandSuggestions, + prepareChatComposerSubmission, +} from '@/components/agents/chat-composer-slash-commands'; + +const commands = [ + { name: 'compact', description: 'Compact context', hints: [] }, + { name: 'review', description: 'Review changes', hints: ['$ARGUMENTS'] }, + { name: 'rename', description: 'Rename a symbol', hints: ['$ARGUMENTS'] }, +]; + +describe('getSlashCommandCandidate', () => { + it('keeps input that can still match a command name', () => { + expect(getSlashCommandCandidate('/')).toBe('/'); + expect(getSlashCommandCandidate('/re')).toBe('/re'); + }); + + it('collapses prose and argument text to null', () => { + expect(getSlashCommandCandidate('hello')).toBeNull(); + expect(getSlashCommandCandidate('/review main')).toBeNull(); + expect(getSlashCommandCandidate('')).toBeNull(); + }); +}); + +describe('getSlashCommandSuggestions', () => { + it('filters the current catalog by the command-name prefix', () => { + expect(getSlashCommandSuggestions('/re', commands)).toEqual([commands[1], commands[2]]); + }); + + it('closes after command arguments begin or when input is not slash-prefixed', () => { + expect(getSlashCommandSuggestions('/review main', commands)).toEqual([]); + expect(getSlashCommandSuggestions('review', commands)).toEqual([]); + }); +}); + +describe('prepareChatComposerSubmission', () => { + it('parses a recognized slash command and preserves its argument text', () => { + expect(prepareChatComposerSubmission(' /review main branch ', commands, false)).toEqual({ + type: 'command', + command: 'review', + arguments: 'main branch', + }); + }); + + it('keeps unknown slash-prefixed input as a prompt', () => { + expect(prepareChatComposerSubmission(' /unknown keep this ', commands, true)).toEqual({ + type: 'prompt', + prompt: '/unknown keep this', + }); + }); + + it('rejects attachments only for recognized commands', () => { + expect(prepareChatComposerSubmission('/compact', commands, true)).toEqual({ + type: 'attachment-error', + }); + expect(prepareChatComposerSubmission('/not-a-command', commands, true)).toEqual({ + type: 'prompt', + prompt: '/not-a-command', + }); + }); +}); diff --git a/apps/mobile/src/components/agents/chat-composer-slash-commands.ts b/apps/mobile/src/components/agents/chat-composer-slash-commands.ts new file mode 100644 index 0000000000..4a8ae05c5d --- /dev/null +++ b/apps/mobile/src/components/agents/chat-composer-slash-commands.ts @@ -0,0 +1,54 @@ +import { type SlashCommandInfo } from 'cloud-agent-sdk'; + +type ChatComposerSubmission = + | { type: 'prompt'; prompt: string } + | { type: 'command'; command: string; arguments: string } + | { type: 'attachment-error' }; + +const SLASH_COMMAND_PATTERN = /^\/([\w.-]+)(?:\s+([\s\S]*))?$/; +const SLASH_CANDIDATE_PATTERN = /^\/[\w.-]*$/; + +/** + * Returns the input when it can still match a command name, `null` otherwise. + * Keeping non-candidates collapsed to `null` lets the composer skip + * re-rendering on every keystroke of ordinary prose. + */ +export function getSlashCommandCandidate(input: string): string | null { + return SLASH_CANDIDATE_PATTERN.test(input) ? input : null; +} + +export function getSlashCommandSuggestions( + input: string, + commands: SlashCommandInfo[] +): SlashCommandInfo[] { + const match = /^\/([\w.-]*)$/.exec(input); + if (!match) { + return []; + } + + const prefix = match[1] ?? ''; + return commands.filter(command => command.name.startsWith(prefix)); +} + +export function prepareChatComposerSubmission( + input: string, + commands: SlashCommandInfo[], + hasAttachments: boolean +): ChatComposerSubmission { + const trimmed = input.trim(); + const match = SLASH_COMMAND_PATTERN.exec(trimmed); + const commandName = match?.[1]; + const recognized = commandName ? commands.some(command => command.name === commandName) : false; + + if (!recognized || !commandName) { + return { type: 'prompt', prompt: trimmed }; + } + if (hasAttachments) { + return { type: 'attachment-error' }; + } + return { + type: 'command', + command: commandName, + arguments: match[2]?.trim() ?? '', + }; +} diff --git a/apps/mobile/src/components/agents/chat-composer.tsx b/apps/mobile/src/components/agents/chat-composer.tsx index 56a9e7bee3..998068cabd 100644 --- a/apps/mobile/src/components/agents/chat-composer.tsx +++ b/apps/mobile/src/components/agents/chat-composer.tsx @@ -1,5 +1,6 @@ import * as Haptics from 'expo-haptics'; import { useActionSheet } from '@expo/react-native-action-sheet'; +import { type SlashCommandInfo } from 'cloud-agent-sdk'; import { ArrowUp, Paperclip, Square } from 'lucide-react-native'; import { useCallback, useRef, useState } from 'react'; import { @@ -14,9 +15,15 @@ import Animated, { FadeIn, FadeOut } from 'react-native-reanimated'; import { toast } from 'sonner-native'; import { AttachmentPreviewStrip } from '@/components/agents/attachment-preview-strip'; +import { + getSlashCommandCandidate, + getSlashCommandSuggestions, + prepareChatComposerSubmission, +} from '@/components/agents/chat-composer-slash-commands'; import { ChatToolbar } from '@/components/agents/chat-toolbar'; import { type AgentMode } from '@/components/agents/mode-selector'; import { pickAgentAttachments } from '@/components/agents/attachment-picker'; +import { SlashCommandSuggestions } from '@/components/agents/slash-command-suggestions'; import { useTextHeight } from '@/components/agents/use-text-height'; import { BlurBar } from '@/components/ui/blur-bar'; import { AGENT_ATTACHMENT_MAX_FILES } from '@/lib/agent-attachments/constants'; @@ -36,9 +43,11 @@ const TEXT_INPUT_MAX_HEIGHT = TEXT_INPUT_LINE_HEIGHT * TEXT_INPUT_MAX_LINES + TEXT_INPUT_VERTICAL_PADDING; const PAPERCLIP_HIT_SLOP = { top: 8, bottom: 8, left: 8, right: 8 } as const; +const EMPTY_COMMANDS: SlashCommandInfo[] = []; type ChatComposerProps = { onSend: (text: string, attachments?: AgentAttachmentWire) => void | Promise; + onSendCommand: (command: string, argumentsText: string) => void | Promise; onStop?: () => void | Promise; disabled?: boolean; isStreaming?: boolean; @@ -52,10 +61,12 @@ type ChatComposerProps = { organizationId?: string; /** Only Cloud Agent sessions can receive attachments. */ attachmentsEnabled?: boolean; + commands?: SlashCommandInfo[]; }; export function ChatComposer({ onSend, + onSendCommand, onStop, disabled = false, isStreaming = false, @@ -68,12 +79,14 @@ export function ChatComposer({ onModelSelect, organizationId, attachmentsEnabled = true, + commands = EMPTY_COMMANDS, }: Readonly) { const colors = useThemeColors(); const { showActionSheetWithOptions } = useActionSheet(); const textRef = useRef(''); const inputRef = useRef(null); const [hasText, setHasText] = useState(false); + const [slashCommandInput, setSlashCommandInput] = useState(null); const [inputWidth, setInputWidth] = useState(0); const [isFocused, setIsFocused] = useState(false); const upload = useAgentAttachmentUpload({ organizationId }); @@ -91,11 +104,14 @@ export function ChatComposer({ const canSend = hasText && !disabled && !isStreaming; const showToolbar = isFocused || hasText || upload.attachments.length > 0; const toolbarDisabled = disabled || isStreaming; + const slashCommandSuggestions = + slashCommandInput === null ? [] : getSlashCommandSuggestions(slashCommandInput, commands); function handleChangeText(value: string) { textRef.current = value; measure.setText(value); setHasText(value.trim().length > 0); + setSlashCommandInput(getSlashCommandCandidate(value)); } function handleSend() { @@ -107,22 +123,44 @@ export function ChatComposer({ toast.error('Wait for attachments to finish uploading.'); return; } - if (trimmed.startsWith('/') && upload.attachments.length > 0) { + const submission = prepareChatComposerSubmission( + trimmed, + commands, + upload.attachments.length > 0 + ); + if (submission.type === 'attachment-error') { toast.error('Attachments cannot be sent with slash commands.'); return; } void Haptics.impactAsync(Haptics.ImpactFeedbackStyle.Light); - const payload = upload.toWirePayload(); - void onSend(trimmed, payload); + if (submission.type === 'command') { + void onSendCommand(submission.command, submission.arguments); + } else { + void onSend(submission.prompt, upload.toWirePayload()); + } textRef.current = ''; setHasText(false); + setSlashCommandInput(null); measure.reset(); inputRef.current?.clear(); upload.reset(); Keyboard.dismiss(); } + function handleSelectSlashCommand(command: SlashCommandInfo) { + const value = `/${command.name} `; + textRef.current = value; + measure.setText(value); + setHasText(true); + setSlashCommandInput(null); + inputRef.current?.setNativeProps({ + text: value, + selection: { start: value.length, end: value.length }, + }); + inputRef.current?.focus(); + } + function handleStop() { void Haptics.impactAsync(Haptics.ImpactFeedbackStyle.Medium); void onStop?.(); @@ -173,6 +211,15 @@ export function ChatComposer({ ) : null} + {slashCommandSuggestions.length > 0 ? ( + + + + ) : null} + {attachmentsEnabled ? ( { + // send() reports failures via its return value instead of throwing. + const accepted = await manager.send({ + payload: { type: 'command', command, arguments: argumentsText }, + }); + if (!accepted) { + toast.error('Failed to send command. Please try again.'); + return; + } + captureEvent(MESSAGE_SENT_EVENT, { surface: analyticsSurface }); + }, + [analyticsSurface, manager] + ); + return ( diff --git a/apps/mobile/src/components/agents/slash-command-suggestions.tsx b/apps/mobile/src/components/agents/slash-command-suggestions.tsx new file mode 100644 index 0000000000..1db223e0bb --- /dev/null +++ b/apps/mobile/src/components/agents/slash-command-suggestions.tsx @@ -0,0 +1,44 @@ +import { type SlashCommandInfo } from 'cloud-agent-sdk'; +import { Pressable, ScrollView } from 'react-native'; + +import { Text } from '@/components/ui/text'; + +type SlashCommandSuggestionsProps = { + commands: SlashCommandInfo[]; + onSelect: (command: SlashCommandInfo) => void; +}; + +export function SlashCommandSuggestions({ + commands, + onSelect, +}: Readonly) { + if (commands.length === 0) { + return null; + } + + return ( + + {commands.map(command => ( + { + onSelect(command); + }} + className="rounded-md px-3 py-2 active:bg-neutral-200 active:opacity-70 dark:active:bg-neutral-700" + accessibilityRole="button" + accessibilityLabel={`Use /${command.name}`} + > + /{command.name} + {command.description ? ( + + {command.description} + + ) : null} + + ))} + + ); +} diff --git a/apps/web/src/hooks/slash-command-selection.ts b/apps/web/src/hooks/slash-command-selection.ts new file mode 100644 index 0000000000..87e9bb145a --- /dev/null +++ b/apps/web/src/hooks/slash-command-selection.ts @@ -0,0 +1,26 @@ +import type { SlashCommandInfo } from '@/lib/cloud-agent-sdk'; +import type { ActiveSessionType } from '@/lib/cloud-agent-sdk/session-manager'; +import type { SlashCommand } from '@/lib/cloud-agent/slash-commands'; +import { commandsOrDefault } from '@cloud-agent-shared'; + +export function selectSlashCommands( + sessionType: ActiveSessionType | null, + commands: SlashCommandInfo[] +): SlashCommand[] { + const selectedCommands = + sessionType === 'cloud-agent' + ? commandsOrDefault(commands) + : sessionType === 'remote' + ? commands + : []; + return selectedCommands.map(toSlashCommand); +} + +function toSlashCommand(info: SlashCommandInfo): SlashCommand { + return { + trigger: info.name, + label: info.name, + description: info.description ?? '', + expansion: '', + }; +} diff --git a/apps/web/src/hooks/useSlashCommandSets.test.ts b/apps/web/src/hooks/useSlashCommandSets.test.ts new file mode 100644 index 0000000000..8cd58c3e38 --- /dev/null +++ b/apps/web/src/hooks/useSlashCommandSets.test.ts @@ -0,0 +1,42 @@ +import { selectSlashCommands } from './slash-command-selection'; + +jest.mock( + '@cloud-agent-shared', + () => ({ + commandsOrDefault: (commands: typeof reportedCommands) => + commands.length > 0 ? commands : [{ name: 'compact', hints: [] }], + }), + { virtual: true } +); + +const reportedCommands = [{ name: 'review', description: 'Review changes', hints: ['$ARGUMENTS'] }]; + +describe('selectSlashCommands', () => { + it('uses only the reported catalog for remote sessions', () => { + expect(selectSlashCommands('remote', reportedCommands)).toEqual([ + { + trigger: 'review', + label: 'review', + description: 'Review changes', + expansion: '', + }, + ]); + expect(selectSlashCommands('remote', [])).toEqual([]); + }); + + it('keeps Cloud Agent defaults while its reported catalog is unavailable', () => { + expect(selectSlashCommands('cloud-agent', [])).toEqual([ + { + trigger: 'compact', + label: 'compact', + description: '', + expansion: '', + }, + ]); + }); + + it('exposes no commands before resolution or for read-only sessions', () => { + expect(selectSlashCommands(null, reportedCommands)).toEqual([]); + expect(selectSlashCommands('read-only', reportedCommands)).toEqual([]); + }); +}); diff --git a/apps/web/src/hooks/useSlashCommandSets.ts b/apps/web/src/hooks/useSlashCommandSets.ts index 4c80b29244..4ebe3f35a5 100644 --- a/apps/web/src/hooks/useSlashCommandSets.ts +++ b/apps/web/src/hooks/useSlashCommandSets.ts @@ -1,9 +1,8 @@ import { useMemo } from 'react'; import { useAtomValue } from 'jotai'; import { useManager } from '@/components/cloud-agent-next/CloudAgentProvider'; -import type { SlashCommandInfo } from '@/lib/cloud-agent-sdk'; -import { commandsOrDefault } from '@cloud-agent-shared'; import type { SlashCommand } from '@/lib/cloud-agent/slash-commands'; +import { selectSlashCommands } from './slash-command-selection'; /** * Source of slash commands for the chat composer. @@ -11,10 +10,9 @@ import type { SlashCommand } from '@/lib/cloud-agent/slash-commands'; * The list comes from the cloud-agent session manager's `availableCommands` * Jotai atom, which is hydrated by `commands.available` events sent by the * cloud-agent worker on every /stream connect (and any time the wrapper - * re-pushes the catalog). When the live list is empty (no wrapper connection - * yet, or wrapper reported no commands), this hook falls back to the pinned - * default catalog so the new-session screen and empty-wrapper cases still - * get autocomplete. + * re-pushes the catalog). Cloud Agent sessions fall back to the pinned default + * catalog while the wrapper list is unavailable. Remote sessions use only the + * exact CLI catalog, while read-only and unresolved sessions expose nothing. * * `expansion` is vestigial — kept for type compatibility with the existing * `SlashCommand` UI shape, but unused now that ChatInput invokes the @@ -23,10 +21,11 @@ import type { SlashCommand } from '@/lib/cloud-agent/slash-commands'; export function useSlashCommandSets() { const manager = useManager(); const commands = useAtomValue(manager.atoms.availableCommands); + const activeSessionType = useAtomValue(manager.atoms.activeSessionType); const availableCommands: SlashCommand[] = useMemo( - () => commandsOrDefault(commands).map(toSlashCommand), - [commands] + () => selectSlashCommands(activeSessionType, commands), + [activeSessionType, commands] ); return { @@ -37,7 +36,7 @@ export function useSlashCommandSets() { { id: 'kilo', name: 'Kilo', - description: 'Project, MCP, and skill commands available in this session', + description: 'Project and MCP commands available in this session', prefix: '', commands: availableCommands, }, @@ -46,12 +45,3 @@ export function useSlashCommandSets() { ), }; } - -function toSlashCommand(info: SlashCommandInfo): SlashCommand { - return { - trigger: info.name, - label: info.name, - description: info.description ?? '', - expansion: '', - }; -} diff --git a/apps/web/src/lib/cloud-agent-sdk/cli-live-transport.test.ts b/apps/web/src/lib/cloud-agent-sdk/cli-live-transport.test.ts index 93b87feee2..6d19cc6eaf 100644 --- a/apps/web/src/lib/cloud-agent-sdk/cli-live-transport.test.ts +++ b/apps/web/src/lib/cloud-agent-sdk/cli-live-transport.test.ts @@ -84,6 +84,22 @@ const REMOTE_CATALOG = { defaultModel: { providerID: 'anthropic', modelID: 'claude-sonnet-4' }, truncated: false, } satisfies RemoteModelCatalogV1; +const COMMAND_WIRE_CATALOG = { + protocolVersion: 1, + commands: [ + { + name: 'review', + description: 'Review changes', + source: 'command' as const, + hints: ['$ARGUMENTS'], + }, + { + name: 'compact', + description: 'compact the current session context', + hints: [], + }, + ], +}; type FakeUserWebConnection = UserWebConnection & { emitCli: (event: UserWebCliEvent) => void; @@ -178,6 +194,239 @@ function emitMessageUpdated(connection: FakeUserWebConnection, sessionId = KILO_ } describe('CliLiveTransport unified user web connection', () => { + it('discovers and publishes commands for the current owner', async () => { + const connection = createConnection(); + jest + .mocked(connection.sendCommand) + .mockImplementation((_sessionId, command) => + Promise.resolve(command === 'list_models' ? WIRE_CATALOG : COMMAND_WIRE_CATALOG) + ); + const { transport, serviceEvents } = createTransportWithSinks({ connection }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + + expect(connection.sendCommand).toHaveBeenCalledWith( + KILO_SESSION_ID, + 'list_commands', + { protocolVersion: 1 }, + 'owner' + ); + expect(serviceEvents).toContainEqual({ + type: 'commands.available', + commands: COMMAND_WIRE_CATALOG.commands, + }); + transport.destroy(); + }); + + it('publishes an empty command catalog for a legacy CLI', async () => { + const connection = createConnection(); + jest + .mocked(connection.sendCommand) + .mockImplementation((_sessionId, command) => + command === 'list_models' + ? Promise.resolve(WIRE_CATALOG) + : Promise.reject(new Error('unknown command: list_commands')) + ); + const { transport, serviceEvents } = createTransportWithSinks({ connection }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + + expect(serviceEvents.filter(event => event.type === 'commands.available')).toEqual([ + { type: 'commands.available', commands: [] }, + ]); + transport.destroy(); + }); + + it('clears a previous catalog when a same-owner refresh is malformed', async () => { + const connection = createConnection(); + const onError = jest.fn(); + let commandCatalogRequest = 0; + jest.mocked(connection.sendCommand).mockImplementation((_sessionId, command) => { + if (command === 'list_models') return Promise.resolve(WIRE_CATALOG); + commandCatalogRequest += 1; + return Promise.resolve( + commandCatalogRequest === 1 ? COMMAND_WIRE_CATALOG : { invalid: true } + ); + }); + const { transport, serviceEvents } = createTransportWithSinks({ connection, onError }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + connection.emitReconnect(); + await Promise.resolve(); + await Promise.resolve(); + + expect(serviceEvents.filter(event => event.type === 'commands.available').at(-1)).toEqual({ + type: 'commands.available', + commands: [], + }); + expect(onError).toHaveBeenCalledWith('Invalid remote command catalog'); + transport.destroy(); + }); + + it('clears a previous catalog when the relay rejects an oversized refresh', async () => { + const connection = createConnection(); + const onError = jest.fn(); + let commandCatalogRequest = 0; + jest.mocked(connection.sendCommand).mockImplementation((_sessionId, command) => { + if (command === 'list_models') return Promise.resolve(WIRE_CATALOG); + commandCatalogRequest += 1; + return commandCatalogRequest === 1 + ? Promise.resolve(COMMAND_WIRE_CATALOG) + : Promise.reject( + new UserWebCommandError({ + code: 'CATALOG_TOO_LARGE', + message: 'Command catalog response is too large', + }) + ); + }); + const { transport, serviceEvents } = createTransportWithSinks({ connection, onError }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + connection.emitReconnect(); + await Promise.resolve(); + await Promise.resolve(); + + expect(serviceEvents.filter(event => event.type === 'commands.available').at(-1)).toEqual({ + type: 'commands.available', + commands: [], + }); + expect(onError).toHaveBeenCalledWith('Command catalog response is too large'); + transport.destroy(); + }); + + it('retains a same-owner catalog during refresh and after a transient failure', async () => { + const connection = createConnection(); + const onError = jest.fn(); + let rejectRefresh: ((error: Error) => void) | undefined; + let commandCatalogRequest = 0; + jest.mocked(connection.sendCommand).mockImplementation((_sessionId, command) => { + if (command === 'list_models') return Promise.resolve(WIRE_CATALOG); + commandCatalogRequest += 1; + if (commandCatalogRequest === 1) return Promise.resolve(COMMAND_WIRE_CATALOG); + return new Promise((_resolve, reject) => { + rejectRefresh = reject; + }); + }); + const { transport, serviceEvents } = createTransportWithSinks({ connection, onError }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + const publishedBeforeRefresh = serviceEvents.filter( + event => event.type === 'commands.available' + ); + + connection.emitReconnect(); + expect(serviceEvents.filter(event => event.type === 'commands.available')).toEqual( + publishedBeforeRefresh + ); + + rejectRefresh?.(new Error('command catalog timed out')); + await Promise.resolve(); + await Promise.resolve(); + expect(serviceEvents.filter(event => event.type === 'commands.available')).toEqual( + publishedBeforeRefresh + ); + expect(onError).toHaveBeenCalledWith('command catalog timed out'); + transport.destroy(); + }); + + it('ignores a late command catalog from a replaced owner', async () => { + const connection = createConnection(); + let resolveFirstCatalog: ((catalog: typeof COMMAND_WIRE_CATALOG) => void) | undefined; + const firstCatalog = new Promise(resolve => { + resolveFirstCatalog = resolve; + }); + const replacementCatalog = { + protocolVersion: 1, + commands: [{ name: 'compact', hints: [] }], + }; + jest.mocked(connection.sendCommand).mockImplementation((_sessionId, command, _data, owner) => { + if (command === 'list_models') return Promise.resolve(WIRE_CATALOG); + return owner === 'owner-a' ? firstCatalog : Promise.resolve(replacementCatalog); + }); + const { transport, serviceEvents } = createTransportWithSinks({ connection }); + + transport.connect(); + emitOwner(connection, 'owner-a'); + emitOwner(connection, 'owner-b'); + await Promise.resolve(); + await Promise.resolve(); + + resolveFirstCatalog?.(COMMAND_WIRE_CATALOG); + await Promise.resolve(); + await Promise.resolve(); + + expect(serviceEvents.filter(event => event.type === 'commands.available').at(-1)).toEqual({ + type: 'commands.available', + commands: replacementCatalog.commands, + }); + transport.destroy(); + }); + + it('keeps one command catalog request in flight per owner', async () => { + const connection = createConnection(); + let resolveCatalog: ((catalog: typeof COMMAND_WIRE_CATALOG) => void) | undefined; + jest.mocked(connection.sendCommand).mockImplementation((_sessionId, command) => { + if (command === 'list_models') return Promise.resolve(WIRE_CATALOG); + return new Promise(resolve => { + resolveCatalog = resolve; + }); + }); + const { transport } = createTransportWithSinks({ connection }); + + transport.connect(); + emitOwner(connection); + connection.emitReconnect(); + connection.emitReconnect(); + + expect( + jest + .mocked(connection.sendCommand) + .mock.calls.filter(([, command]) => command === 'list_commands') + ).toHaveLength(1); + + resolveCatalog?.(COMMAND_WIRE_CATALOG); + await Promise.resolve(); + await Promise.resolve(); + transport.destroy(); + }); + + it('clears commands when the current owner disconnects', async () => { + const connection = createConnection(); + jest + .mocked(connection.sendCommand) + .mockImplementation((_sessionId, command) => + Promise.resolve(command === 'list_models' ? WIRE_CATALOG : COMMAND_WIRE_CATALOG) + ); + const { transport, serviceEvents } = createTransportWithSinks({ connection }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + connection.emitSystem({ event: 'cli.disconnected', data: { connectionId: 'owner' } }); + + expect(serviceEvents.filter(event => event.type === 'commands.available').at(-1)).toEqual({ + type: 'commands.available', + commands: [], + }); + transport.destroy(); + }); + it('discovers and publishes a v1 catalog for the current owner', async () => { const connection = createConnection(); jest @@ -344,10 +593,10 @@ describe('CliLiveTransport unified user web connection', () => { ...REMOTE_CATALOG, providers: [{ ...REMOTE_CATALOG.providers[0], id: 'replacement-provider' }], } satisfies RemoteModelCatalogV1; - jest - .mocked(connection.sendCommand) - .mockReturnValueOnce(firstCatalog) - .mockResolvedValueOnce(replacementWireCatalog); + jest.mocked(connection.sendCommand).mockImplementation((_sessionId, command, _data, owner) => { + if (command === 'list_commands') return Promise.resolve(COMMAND_WIRE_CATALOG); + return owner === 'owner-a' ? firstCatalog : Promise.resolve(replacementWireCatalog); + }); const states: RemoteModelState[] = []; const { transport } = createTransportWithSinks({ connection, @@ -379,10 +628,12 @@ describe('CliLiveTransport unified user web connection', () => { it('keeps one catalog request in flight for an owner', async () => { const connection = createConnection(); let resolveCatalog: ((catalog: RemoteModelCatalogWireV1) => void) | undefined; - jest.mocked(connection.sendCommand).mockReturnValue( - new Promise(resolve => { - resolveCatalog = resolve; - }) + jest.mocked(connection.sendCommand).mockImplementation((_sessionId, command) => + command === 'list_commands' + ? Promise.resolve(COMMAND_WIRE_CATALOG) + : new Promise(resolve => { + resolveCatalog = resolve; + }) ); const { transport } = createTransportWithSinks({ connection }); @@ -391,7 +642,11 @@ describe('CliLiveTransport unified user web connection', () => { transport.retryRemoteModels?.(); connection.emitReconnect(); - expect(connection.sendCommand).toHaveBeenCalledTimes(1); + expect( + jest + .mocked(connection.sendCommand) + .mock.calls.filter(([, command]) => command === 'list_models') + ).toHaveLength(1); resolveCatalog?.(WIRE_CATALOG); await Promise.resolve(); @@ -401,10 +656,14 @@ describe('CliLiveTransport unified user web connection', () => { it('retains a v1 catalog when a same-owner reconnect refresh fails', async () => { const connection = createConnection(); - jest - .mocked(connection.sendCommand) - .mockResolvedValueOnce(WIRE_CATALOG) - .mockRejectedValueOnce(new Error('catalog refresh timed out')); + let modelCatalogRequest = 0; + jest.mocked(connection.sendCommand).mockImplementation((_sessionId, command) => { + if (command === 'list_commands') return Promise.resolve(COMMAND_WIRE_CATALOG); + modelCatalogRequest += 1; + return modelCatalogRequest === 1 + ? Promise.resolve(WIRE_CATALOG) + : Promise.reject(new Error('catalog refresh timed out')); + }); const states: RemoteModelState[] = []; const { transport } = createTransportWithSinks({ connection, @@ -428,13 +687,21 @@ describe('CliLiveTransport unified user web connection', () => { refresh: 'error', error: 'catalog refresh timed out', }); - expect(connection.sendCommand).toHaveBeenCalledTimes(2); + expect( + jest + .mocked(connection.sendCommand) + .mock.calls.filter(([, command]) => command === 'list_models') + ).toHaveLength(2); transport.destroy(); }); it('clears owner-scoped catalog state and rediscovers after session reappearance', async () => { const connection = createConnection(); - jest.mocked(connection.sendCommand).mockResolvedValue(WIRE_CATALOG); + jest + .mocked(connection.sendCommand) + .mockImplementation((_sessionId, command) => + Promise.resolve(command === 'list_models' ? WIRE_CATALOG : COMMAND_WIRE_CATALOG) + ); const states: RemoteModelState[] = []; const { transport } = createTransportWithSinks({ connection, @@ -458,13 +725,13 @@ describe('CliLiveTransport unified user web connection', () => { await Promise.resolve(); await Promise.resolve(); - expect(connection.sendCommand).toHaveBeenLastCalledWith( + expect(connection.sendCommand).toHaveBeenCalledWith( KILO_SESSION_ID, 'list_models', { protocolVersion: 1 }, 'owner-b' ); - expect(connection.sendCommand).toHaveBeenCalledTimes(2); + expect(connection.sendCommand).toHaveBeenCalledTimes(4); transport.destroy(); }); @@ -563,7 +830,10 @@ describe('CliLiveTransport unified user web connection', () => { userWebConnection.emitSystem({ event: 'cli.disconnected', data: { connectionId: 'owner' } }); userWebConnection.emitSystem({ event: 'cli.disconnected', data: { connectionId: 'owner' } }); - expect(serviceEvents).toEqual([{ type: 'stopped', reason: 'disconnected' }]); + expect(serviceEvents).toEqual([ + { type: 'commands.available', commands: [] }, + { type: 'stopped', reason: 'disconnected' }, + ]); userWebConnection.emitSystem({ event: 'sessions.heartbeat', @@ -600,6 +870,7 @@ describe('CliLiveTransport unified user web connection', () => { // a session.status event clears the disconnected UI state. expect(serviceEvents).toEqual([ { type: 'session.status', sessionId: KILO_SESSION_ID, status: { type: 'idle' } }, + { type: 'commands.available', commands: [] }, { type: 'stopped', reason: 'disconnected' }, { type: 'session.status', sessionId: KILO_SESSION_ID, status: { type: 'idle' } }, ]); @@ -1208,13 +1479,46 @@ describe('CliLiveTransport unified user web connection', () => { } ); - it('rejects structured slash commands without sending a viewer command', async () => { - const { userWebConnection, transport } = createTransportWithSinks(); + it('sends a structured command with arguments, message ID, model, and variant', async () => { + const connection = createConnection(); + jest.mocked(connection.sendCommand).mockImplementation((_sessionId, command) => { + if (command === 'list_models') return Promise.resolve(WIRE_CATALOG); + if (command === 'list_commands') return Promise.resolve(COMMAND_WIRE_CATALOG); + return Promise.resolve({ ok: true }); + }); + const { userWebConnection, transport } = createTransportWithSinks({ connection }); - await expect( - transport.send!({ payload: { type: 'command', command: 'review', arguments: '' } }) - ).rejects.toThrow('Slash commands are not supported on the CLI live transport yet'); - expect(userWebConnection.sendCommand).not.toHaveBeenCalled(); + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + jest.mocked(userWebConnection.sendCommand).mockClear(); + + await transport.send?.({ + payload: { type: 'command', command: 'review', arguments: 'main --fix' }, + messageId: 'msg-command-1', + remoteModelOverride: { + source: 'cli-catalog', + selection: { + model: { providerID: 'anthropic', modelID: 'claude-sonnet-4' }, + variant: 'high', + }, + }, + }); + + expect(userWebConnection.sendCommand).toHaveBeenCalledWith( + KILO_SESSION_ID, + 'send_command', + { + protocolVersion: 1, + command: 'review', + arguments: 'main --fix', + messageID: 'msg-command-1', + model: { providerID: 'anthropic', modelID: 'claude-sonnet-4' }, + variant: 'high', + }, + 'owner' + ); transport.destroy(); }); }); diff --git a/apps/web/src/lib/cloud-agent-sdk/cli-live-transport.ts b/apps/web/src/lib/cloud-agent-sdk/cli-live-transport.ts index ce96139c10..6edde3476b 100644 --- a/apps/web/src/lib/cloud-agent-sdk/cli-live-transport.ts +++ b/apps/web/src/lib/cloud-agent-sdk/cli-live-transport.ts @@ -6,12 +6,13 @@ import { normalizeCliEvent, isChatEvent } from './normalizer'; import { cliConnectionDataSchema, heartbeatDataSchema, + remoteCommandCatalogV1Schema, remoteModelCatalogV1Schema, sessionsListDataSchema, } from './schemas'; import type { RemoteModelState } from './remote-model-catalog'; import type { TransportFactory, TransportSendInput, TransportSink } from './transport'; -import type { KiloSessionId, SessionSnapshot } from './types'; +import type { KiloSessionId, SessionSnapshot, SlashCommandInfo } from './types'; import { UserWebCommandError, type UserWebCliEvent, @@ -41,6 +42,11 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor let lastForwardedHeartbeatStatus: string | null = null; let catalogRequestGeneration = 0; let catalogRequestInFlight: { ownerConnectionId: string; generation: number } | null = null; + let commandCatalogRequestGeneration = 0; + let commandCatalogRequestInFlight: { + ownerConnectionId: string; + generation: number; + } | null = null; let remoteModelState: RemoteModelState = { ownerConnectionId: null, protocol: 'unknown', @@ -52,12 +58,20 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor config.onRemoteModelStateChange?.(next); } + function publishCommands(commands: SlashCommandInfo[]): void { + sink.onServiceEvent({ type: 'commands.available', commands }); + } + function setOwnerConnectionId(nextOwnerConnectionId: string | null): void { if (ownerConnectionId === nextOwnerConnectionId) return; + const previousOwnerConnectionId = ownerConnectionId; ownerConnectionId = nextOwnerConnectionId; catalogRequestGeneration += 1; catalogRequestInFlight = null; + commandCatalogRequestGeneration += 1; + commandCatalogRequestInFlight = null; + if (previousOwnerConnectionId) publishCommands([]); publishRemoteModelState({ ownerConnectionId: nextOwnerConnectionId, protocol: 'unknown', @@ -65,7 +79,10 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor }); config.onCapabilityChange?.(); - if (nextOwnerConnectionId) discoverModels(nextOwnerConnectionId); + if (nextOwnerConnectionId) { + discoverModels(nextOwnerConnectionId); + discoverCommands(nextOwnerConnectionId); + } } function handleCatalogFailure( @@ -172,6 +189,99 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor }); } + function handleCommandCatalogFailure( + error: unknown, + expectedOwnerConnectionId: string, + expectedGeneration: number, + expectedRequestGeneration: number, + clearCatalog = false + ): void { + if ( + expectedGeneration !== generation || + expectedRequestGeneration !== commandCatalogRequestGeneration || + ownerConnectionId !== expectedOwnerConnectionId + ) { + return; + } + + if (error instanceof UserWebCommandError && error.code === 'SESSION_OWNER_CHANGED') { + setOwnerConnectionId(null); + return; + } + + if (error instanceof Error && error.message.includes('unknown command')) { + publishCommands([]); + return; + } + + if ( + clearCatalog || + (error instanceof UserWebCommandError && error.code === 'CATALOG_TOO_LARGE') + ) { + publishCommands([]); + } + config.onError?.( + error instanceof Error ? error.message : 'Failed to discover remote commands' + ); + } + + function discoverCommands(expectedOwnerConnectionId: string): void { + if (commandCatalogRequestInFlight?.ownerConnectionId === expectedOwnerConnectionId) return; + + commandCatalogRequestGeneration += 1; + const expectedRequestGeneration = commandCatalogRequestGeneration; + const expectedGeneration = generation; + commandCatalogRequestInFlight = { + ownerConnectionId: expectedOwnerConnectionId, + generation: expectedRequestGeneration, + }; + + void config.userWebConnection + .sendCommand( + config.kiloSessionId, + 'list_commands', + { protocolVersion: 1 }, + expectedOwnerConnectionId + ) + .then( + result => { + if ( + expectedGeneration !== generation || + expectedRequestGeneration !== commandCatalogRequestGeneration || + ownerConnectionId !== expectedOwnerConnectionId + ) { + return; + } + + const parsed = remoteCommandCatalogV1Schema.safeParse(result); + if (!parsed.success) { + handleCommandCatalogFailure( + new Error('Invalid remote command catalog'), + expectedOwnerConnectionId, + expectedGeneration, + expectedRequestGeneration, + true + ); + return; + } + + publishCommands(parsed.data.commands); + }, + error => + handleCommandCatalogFailure( + error, + expectedOwnerConnectionId, + expectedGeneration, + expectedRequestGeneration + ) + ) + .finally(() => { + if (commandCatalogRequestInFlight?.generation === expectedRequestGeneration) { + commandCatalogRequestInFlight = null; + } + }); + } + function replaySnapshot(snapshot: SessionSnapshot): void { sink.onServiceEvent({ type: 'session.created', info: snapshot.info }); @@ -357,6 +467,8 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor lastForwardedHeartbeatStatus = null; catalogRequestGeneration += 1; catalogRequestInFlight = null; + commandCatalogRequestGeneration += 1; + commandCatalogRequestInFlight = null; publishRemoteModelState({ ownerConnectionId: null, protocol: 'unknown', @@ -457,7 +569,10 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor resyncTimer = null; replayCurrentSnapshot(false); }, RECONNECT_RESYNC_DELAY_MS); - if (ownerConnectionId) discoverModels(ownerConnectionId); + if (ownerConnectionId) { + discoverModels(ownerConnectionId); + discoverCommands(ownerConnectionId); + } }); const releaseSubscription = config.userWebConnection.subscribeToCliSession( config.kiloSessionId @@ -480,9 +595,22 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor }, send: async (input: TransportSendInput) => { if (input.payload.type === 'command') { - return Promise.reject( - new Error('Slash commands are not supported on the CLI live transport yet') - ); + const remoteModel = getRemoteModelFields(input); + return sendCommand('send_command', { + protocolVersion: 1, + command: input.payload.command, + arguments: input.payload.arguments, + ...(input.messageId ? { messageID: input.messageId } : {}), + ...(remoteModel.kind === 'none' + ? {} + : { + model: + remoteModel.kind === 'structured' + ? remoteModel.model + : { providerID: 'kilo', modelID: remoteModel.model }, + ...(remoteModel.variant ? { variant: remoteModel.variant } : {}), + }), + }); } const payload = input.payload; const remoteModel = getRemoteModelFields(input); diff --git a/apps/web/src/lib/cloud-agent-sdk/remote-command-catalog.test.ts b/apps/web/src/lib/cloud-agent-sdk/remote-command-catalog.test.ts new file mode 100644 index 0000000000..a44573b116 --- /dev/null +++ b/apps/web/src/lib/cloud-agent-sdk/remote-command-catalog.test.ts @@ -0,0 +1,112 @@ +import { remoteCommandCatalogV1Schema } from './schemas'; + +describe('remote command catalog schema', () => { + it('accepts a strict v1 catalog and excludes skill commands', () => { + expect( + remoteCommandCatalogV1Schema.parse({ + protocolVersion: 1, + commands: [ + { + name: 'review', + description: 'Review changes', + source: 'command', + hints: ['$ARGUMENTS'], + }, + { + name: 'hidden-skill', + description: 'Not part of the remote surface', + source: 'skill', + hints: [], + }, + { + name: 'compact', + description: 'compact the current session context', + hints: [], + }, + ], + }) + ).toEqual({ + protocolVersion: 1, + commands: [ + { + name: 'review', + description: 'Review changes', + source: 'command', + hints: ['$ARGUMENTS'], + }, + { + name: 'compact', + description: 'compact the current session context', + hints: [], + }, + ], + }); + }); + + const command = { + name: 'review', + description: 'Review changes', + agent: 'code', + model: 'anthropic/claude-sonnet-4', + source: 'command' as const, + hints: ['$ARGUMENTS'], + }; + const catalog = (commands: unknown[]) => ({ protocolVersion: 1, commands }); + + it('rejects unsupported protocols and unknown wire fields', () => { + expect( + remoteCommandCatalogV1Schema.safeParse({ ...catalog([command]), protocolVersion: 2 }).success + ).toBe(false); + expect( + remoteCommandCatalogV1Schema.safeParse({ ...catalog([command]), extra: true }).success + ).toBe(false); + expect( + remoteCommandCatalogV1Schema.safeParse( + catalog([{ ...command, template: 'private implementation detail' }]) + ).success + ).toBe(false); + }); + + it('rejects more than 256 commands', () => { + expect( + remoteCommandCatalogV1Schema.safeParse(catalog(Array.from({ length: 257 }, () => command))) + .success + ).toBe(false); + }); + + it('rejects strings longer than 2,000 characters', () => { + for (const field of ['name', 'description', 'agent', 'model'] as const) { + expect( + remoteCommandCatalogV1Schema.safeParse( + catalog([{ ...command, [field]: 'x'.repeat(2_001) }]) + ).success + ).toBe(false); + } + + expect( + remoteCommandCatalogV1Schema.safeParse(catalog([{ ...command, hints: ['x'.repeat(2_001)] }])) + .success + ).toBe(false); + }); + + it('rejects more than 32 hints per command', () => { + expect( + remoteCommandCatalogV1Schema.safeParse( + catalog([{ ...command, hints: Array.from({ length: 33 }, () => 'hint') }]) + ).success + ).toBe(false); + }); + + it('measures the 512 KiB serialized bound in UTF-8 bytes', () => { + const multibyteCatalog = catalog( + Array.from({ length: 256 }, () => ({ + ...command, + description: 'é'.repeat(1_000), + })) + ); + const serialized = JSON.stringify(multibyteCatalog); + expect(serialized.length).toBeLessThan(512 * 1024); + expect(new TextEncoder().encode(serialized).byteLength).toBeGreaterThan(512 * 1024); + expect(remoteCommandCatalogV1Schema.safeParse(multibyteCatalog).success).toBe(false); + }); +}); diff --git a/apps/web/src/lib/cloud-agent-sdk/schemas.ts b/apps/web/src/lib/cloud-agent-sdk/schemas.ts index db4fbf2630..82b330f72b 100644 --- a/apps/web/src/lib/cloud-agent-sdk/schemas.ts +++ b/apps/web/src/lib/cloud-agent-sdk/schemas.ts @@ -333,6 +333,45 @@ export const remoteModelCatalogV1Schema = remoteModelCatalogWireV1Schema.transfo }); export type RemoteModelCatalogV1 = z.output; +export const REMOTE_COMMAND_MAX_COMMANDS = 256; +export const REMOTE_COMMAND_MAX_STRING_LENGTH = 2_000; +export const REMOTE_COMMAND_MAX_HINTS = 32; +export const REMOTE_COMMAND_CATALOG_MAX_SERIALIZED_BYTES = 512 * 1024; + +const remoteCommandStringSchema = z.string().max(REMOTE_COMMAND_MAX_STRING_LENGTH); +const remoteSlashCommandInfoSchema = z + .object({ + name: remoteCommandStringSchema.min(1), + description: remoteCommandStringSchema.optional(), + agent: remoteCommandStringSchema.optional(), + model: remoteCommandStringSchema.optional(), + source: z.enum(['command', 'mcp', 'skill']).optional(), + hints: z.array(remoteCommandStringSchema).max(REMOTE_COMMAND_MAX_HINTS), + subtask: z.boolean().optional(), + }) + .strict(); + +export const remoteCommandCatalogV1Schema = z + .object({ + protocolVersion: z.literal(1), + commands: z.array(remoteSlashCommandInfoSchema).max(REMOTE_COMMAND_MAX_COMMANDS), + }) + .strict() + .superRefine((catalog, context) => { + const serializedBytes = new TextEncoder().encode(JSON.stringify(catalog)).byteLength; + if (serializedBytes > REMOTE_COMMAND_CATALOG_MAX_SERIALIZED_BYTES) { + context.addIssue({ + code: 'custom', + message: `Catalog cannot exceed ${REMOTE_COMMAND_CATALOG_MAX_SERIALIZED_BYTES} serialized bytes`, + }); + } + }) + .transform(catalog => ({ + protocolVersion: 1 as const, + commands: catalog.commands.filter(command => command.source !== 'skill'), + })); +export type RemoteCommandCatalogV1 = z.output; + export const userWebCommandErrorDataSchema = z .object({ source: z.literal('relay'), diff --git a/apps/web/src/lib/cloud-agent-sdk/session-manager.test.ts b/apps/web/src/lib/cloud-agent-sdk/session-manager.test.ts index b0d706f9f3..09ffadc15c 100644 --- a/apps/web/src/lib/cloud-agent-sdk/session-manager.test.ts +++ b/apps/web/src/lib/cloud-agent-sdk/session-manager.test.ts @@ -2688,6 +2688,23 @@ describe('createSessionManager', () => { await switchPromise; }); + + it('ignores a command catalog emitted by the previous session after switching', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + await mgr.switchSession(kiloId('ses-1')); + + const mockedCreate = jest.mocked(createCloudAgentSession); + const previousSessionEvent = mockedCreate.mock.calls[0][0].onEvent; + + await mgr.switchSession(kiloId('ses-2')); + previousSessionEvent?.({ + type: 'commands.available', + commands: [{ name: 'stale', description: 'Old workspace command', hints: [] }], + }); + + expect(atomValue(config.store, mgr.atoms.availableCommands)).toEqual([]); + }); }); // ------------------------------------------------------------------------- diff --git a/apps/web/src/lib/cloud-agent-sdk/session-manager.ts b/apps/web/src/lib/cloud-agent-sdk/session-manager.ts index 9b823f2326..0bc4dc811a 100644 --- a/apps/web/src/lib/cloud-agent-sdk/session-manager.ts +++ b/apps/web/src/lib/cloud-agent-sdk/session-manager.ts @@ -911,6 +911,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { }, onEvent: event => { if (event.type === 'commands.available') { + if (expectedGeneration !== switchGeneration) return; // Replace the catalog wholesale. The DO sends the full list on // every connect, so we never need to merge incrementally. store.set(availableCommandsAtom, event.commands); diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 83fddbf466..9d97fee875 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -1352,6 +1352,38 @@ describe('UserConnectionDO', () => { expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(1); }); + it('rejects a duplicate in-flight list_commands request for the same viewer session and owner', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', + }); + sendCommand(doInstance, webWs, { + id: 'cmd-2', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', + }); + + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-2', + error: { + source: 'relay', + code: 'CATALOG_REQUEST_PENDING', + message: 'Command catalog request already pending', + }, + }); + expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(1); + }); + it('expires pending commands before handling another command', () => { const now = 1_000_000; vi.spyOn(Date, 'now').mockReturnValue(now); @@ -1524,6 +1556,38 @@ describe('UserConnectionDO', () => { }); }); + it('rejects a list_commands result one byte over 512 KiB', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1), + }); + + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-1', + error: { + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Command catalog response is too large', + }, + }); + }); + it('rejects a multibyte list_models result over 512 KiB', () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index e79014dc7f..8317d4a71b 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -49,12 +49,24 @@ const CATALOG_TOO_LARGE_ERROR = { message: 'Model catalog response is too large', }; +const COMMAND_CATALOG_TOO_LARGE_ERROR = { + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Command catalog response is too large', +}; + const CATALOG_REQUEST_PENDING_ERROR = { source: 'relay', code: 'CATALOG_REQUEST_PENDING', message: 'Model catalog request already pending', }; +const COMMAND_CATALOG_REQUEST_PENDING_ERROR = { + source: 'relay', + code: 'CATALOG_REQUEST_PENDING', + message: 'Command catalog request already pending', +}; + const PENDING_COMMAND_LIMIT_ERROR = { source: 'relay', code: 'PENDING_COMMAND_LIMIT', @@ -72,6 +84,10 @@ const CLI_COMMAND_ERROR = { message: 'Command failed', }; +function isCatalogCommand(command: string): command is 'list_models' | 'list_commands' { + return command === 'list_models' || command === 'list_commands'; +} + export class UserConnectionDO extends DurableObject { private static readonly HEARTBEAT_TIMEOUT_MS = 30_000; private static readonly PENDING_COMMAND_TTL_MS = 35_000; @@ -468,14 +484,17 @@ export class UserConnectionDO extends DurableObject { if (!entry || entry.targetCliWs !== respondingWs) return; this.pendingCommands.delete(id); - if (entry.command === 'list_models' && result !== undefined) { + if (isCatalogCommand(entry.command) && result !== undefined) { const serializedResult = JSON.stringify(result); const resultBytes = new TextEncoder().encode(serializedResult).byteLength; if (resultBytes > MAX_CATALOG_RESULT_BYTES) { this.sendToWeb(entry.ws, { type: 'response', id: entry.originalId, - error: CATALOG_TOO_LARGE_ERROR, + error: + entry.command === 'list_models' + ? CATALOG_TOO_LARGE_ERROR + : COMMAND_CATALOG_TOO_LARGE_ERROR, }); return; } @@ -633,11 +652,11 @@ export class UserConnectionDO extends DurableObject { const targetConnectionId = targetAttachment.connectionId; if ( - msg.command === 'list_models' && + isCatalogCommand(msg.command) && [...this.pendingCommands.values()].some( entry => entry.ws === ws && - entry.command === 'list_models' && + entry.command === msg.command && entry.sessionId === msg.sessionId && entry.targetConnectionId === targetConnectionId ) @@ -645,7 +664,10 @@ export class UserConnectionDO extends DurableObject { this.sendToWeb(ws, { type: 'response', id: msg.id, - error: CATALOG_REQUEST_PENDING_ERROR, + error: + msg.command === 'list_models' + ? CATALOG_REQUEST_PENDING_ERROR + : COMMAND_CATALOG_REQUEST_PENDING_ERROR, }); return; }