diff --git a/apps/mobile/app.config.ts b/apps/mobile/app.config.ts index bfa9299c14..bfea015e26 100644 --- a/apps/mobile/app.config.ts +++ b/apps/mobile/app.config.ts @@ -42,6 +42,11 @@ const config: ExpoConfig = { AdAttributionKit: { PostbackCopyURL: 'https://appsflyer-skadnetwork.com/', }, + // Make the app's Documents directory user-visible in the iOS Files + // app so downloaded any-file attachments (uploaded via the cloud-agent + // composer) can be opened in place. + UIFileSharingEnabled: true, + LSSupportsOpeningDocumentsInPlace: true, }, }, splash: { diff --git a/apps/mobile/src/app/(app)/agent-chat/new.tsx b/apps/mobile/src/app/(app)/agent-chat/new.tsx index 15806fbc3c..a76272cb78 100644 --- a/apps/mobile/src/app/(app)/agent-chat/new.tsx +++ b/apps/mobile/src/app/(app)/agent-chat/new.tsx @@ -210,7 +210,10 @@ export default function NewSessionScreen() { const { addCandidates } = attachments; const handleAddAttachment = useCallback(async () => { - addCandidates(await pickAgentAttachments(showActionSheetWithOptions)); + // Fire-and-forget: the upload hook owns its own progress + error toasts, + // and `canCreate` (computed from `attachments.isUploading` / + // `attachments.hasFailedAttachments`) gates the start-session button. + void addCandidates(await pickAgentAttachments(showActionSheetWithOptions)); }, [addCandidates, showActionSheetWithOptions]); // Cloud-Agent vs. remote-instance submit safety. diff --git a/apps/mobile/src/components/agents/attachment-chip-description.ts b/apps/mobile/src/components/agents/attachment-chip-description.ts new file mode 100644 index 0000000000..a9fae3d488 --- /dev/null +++ b/apps/mobile/src/components/agents/attachment-chip-description.ts @@ -0,0 +1,83 @@ +import { formatFileSize } from '@kilocode/kilo-chat'; + +export type ChipStateInput = { + filename: string; + size: number; + status: 'pending' | 'uploading' | 'uploaded' | 'error'; + progress: number | null; + terminal?: boolean; +}; + +type ChipDescription = { + /** Always present — the original filename. */ + filename: string; + /** Human-readable file size (e.g. "1 KB"). Always present. */ + sizeText: string; + /** + * Progress label for in-flight and uploaded states (e.g. "42%" or + * "Uploaded"). Empty string when no progress is surfaced. + */ + progressText: string; + /** + * Error message for failed states. `null` for non-error states — the + * chip renders the message in place of the size/progress secondary + * text and never both. + */ + message: string | null; + /** True ONLY for retryable error chips — drives the tap-to-retry affordance. */ + showRetry: boolean; + /** True for every rendered chip — drives the X remove affordance. */ + showRemove: boolean; +}; + +/** + * Human-readable progress label. `progress` is `null` when the upload task + * could not determine a total (server did not advertise Content-Length); + * we surface that honestly instead of fabricating a percentage. + */ +export function progressLabel(progress: number | null): string { + if (progress === null) { + return 'Uploading…'; + } + if (progress >= 1) { + return 'Uploaded'; + } + return `${Math.round(progress * 100)}%`; +} + +/** + * Pure projection of the document-kind attachment chip. Kept separate + * from the React component so every feature state (happy, retryable, + * terminal, indeterminate progress) can be unit-tested without a + * renderer. The image-kind chip is a thumbnail overlay and shares no + * state-dependent text; the description below covers the document chip. + */ +export function describeAttachmentChip(state: ChipStateInput): ChipDescription { + const { filename, size, status, progress, terminal } = state; + const isErrored = status === 'error'; + // Explicit `terminal === false` check: a chip is only retryable when + // the server-side flag is set to `false` (transient failure). An + // undefined `terminal` flag defaults to the conservative non-retryable + // bucket so a stray missing field cannot silently enable retry on a + // server-rejected file. + const isRetryable = isErrored && terminal === false; + + let message: string | null = null; + if (isErrored) { + message = isRetryable ? 'Upload failed. Tap to retry.' : "This file can't be uploaded."; + } + + let progressText = ''; + if (status === 'pending' || status === 'uploading' || status === 'uploaded') { + progressText = progressLabel(progress); + } + + return { + filename, + sizeText: formatFileSize(size), + progressText, + message, + showRetry: isRetryable, + showRemove: true, + }; +} diff --git a/apps/mobile/src/components/agents/attachment-picker.test.ts b/apps/mobile/src/components/agents/attachment-picker.test.ts index 6c288aabdf..eb37f47d12 100644 --- a/apps/mobile/src/components/agents/attachment-picker.test.ts +++ b/apps/mobile/src/components/agents/attachment-picker.test.ts @@ -1,4 +1,5 @@ import { type ActionSheetProps } from '@expo/react-native-action-sheet'; +import * as DocumentPicker from 'expo-document-picker'; import { describe, expect, it, vi } from 'vitest'; import { pickAgentAttachments } from './attachment-picker'; @@ -23,11 +24,34 @@ vi.mock('expo-image-picker', () => ({ requestCameraPermissionsAsync: vi.fn(), })); +const getDocumentAsyncMock = vi.mocked(DocumentPicker.getDocumentAsync); + +type ShowActionSheet = ActionSheetProps['showActionSheetWithOptions']; +type SheetButtonHandler = Parameters[1]; + +/** + * Drive `pickAgentAttachments` by capturing the sheet handler the + * production code registers, then invoking it with a button index. + * Avoids inline callback-shaped mocks that trip prefer-await-to-callbacks. + */ +async function pickWithSheetSelection( + buttonIndex: number +): Promise>> { + const showActionSheet = vi.fn() as unknown as ShowActionSheet & { + mock: { calls: [unknown, SheetButtonHandler][] }; + }; + const resultPromise = pickAgentAttachments(showActionSheet); + const registered = showActionSheet.mock.calls[0]?.[1]; + expect(registered).toEqual(expect.any(Function)); + await Promise.resolve(registered?.(buttonIndex)); + return resultPromise; +} + describe('agent attachment picker', () => { it('opens a native action sheet that keeps all sources and the cancel action', () => { - const showActionSheet: ActionSheetProps['showActionSheetWithOptions'] = vi.fn( - (_options, _callback) => undefined - ); + const showActionSheet = vi.fn() as unknown as ShowActionSheet & { + mock: { calls: unknown[][] }; + }; void pickAgentAttachments(showActionSheet); @@ -41,3 +65,71 @@ describe('agent attachment picker', () => { expect(reactNativeMock.alert).not.toHaveBeenCalled(); }); }); + +describe('agent attachment picker (document MIME derivation)', () => { + it('derives MIME from the extension, never from the picker MIME', async () => { + getDocumentAsyncMock.mockResolvedValueOnce({ + canceled: false, + assets: [ + { + uri: 'file:///cache/notes.md', + name: 'notes.md', + mimeType: 'application/octet-stream', + size: 12, + }, + ], + } as unknown as Awaited>); + + const candidates = await pickWithSheetSelection(2); + expect(candidates).toHaveLength(1); + expect(candidates[0]?.name).toBe('notes.md'); + // The picker reported `application/octet-stream`; the picker must + // ignore it and return the extension-derived MIME. + expect(candidates[0]?.mimeType).toBe('text/plain'); + }); + + it('returns application/octet-stream for an extension outside the canonical table', async () => { + getDocumentAsyncMock.mockResolvedValueOnce({ + canceled: false, + assets: [ + { + uri: 'file:///cache/clip.mov', + name: 'clip.mov', + mimeType: 'video/quicktime', + size: 12, + }, + ], + } as unknown as Awaited>); + + const candidates = await pickWithSheetSelection(2); + expect(candidates[0]?.mimeType).toBe('application/octet-stream'); + }); + + it('falls back to application/octet-stream for a filename with no usable extension', async () => { + getDocumentAsyncMock.mockResolvedValueOnce({ + canceled: false, + assets: [ + { + uri: 'file:///cache/README', + name: 'README', + mimeType: 'text/plain', + size: 12, + }, + ], + } as unknown as Awaited>); + + const candidates = await pickWithSheetSelection(2); + expect(candidates[0]?.mimeType).toBe('application/octet-stream'); + }); + + it('returns an empty list on cancel', async () => { + getDocumentAsyncMock.mockResolvedValueOnce({ + canceled: true, + assets: [], + } as unknown as Awaited>); + + // Cancel is button index 3; Files (2) with a canceled document result + // also yields []. Use Files so the document path is exercised. + expect(await pickWithSheetSelection(2)).toEqual([]); + }); +}); diff --git a/apps/mobile/src/components/agents/attachment-picker.ts b/apps/mobile/src/components/agents/attachment-picker.ts index 9fedefc37c..d280c40fc3 100644 --- a/apps/mobile/src/components/agents/attachment-picker.ts +++ b/apps/mobile/src/components/agents/attachment-picker.ts @@ -4,10 +4,7 @@ import * as ImagePicker from 'expo-image-picker'; import { type ActionSheetProps } from '@expo/react-native-action-sheet'; import { Alert, Linking } from 'react-native'; -import { - AGENT_ATTACHMENT_MIME_BY_EXTENSION, - type AgentAttachmentExtension, -} from '@/lib/agent-attachments/constants'; +import { mimeForExtension, normalizeAttachmentExtension } from '@/lib/agent-attachments/validate'; import { type AgentAttachmentCandidate } from '@/lib/agent-attachments/use-agent-attachment-upload'; const IMAGE_PICKER_OPTIONS = { @@ -22,17 +19,18 @@ function showPermissionSettingsAlert({ message, title }: { message: string; titl ]); } -function mimeTypeForExtension(extension: AgentAttachmentExtension): string { - return AGENT_ATTACHMENT_MIME_BY_EXTENSION[extension]; -} - function normalizeImageAsset(asset: { uri: string; fileName?: string | null; mimeType?: string | null; fileSize?: number | null; }): AgentAttachmentCandidate { - const name = asset.fileName ?? `image.${(asset.mimeType ?? 'image/png').split('/')[1] ?? 'png'}`; + // Image picker cannot return a filename with an arbitrary extension; + // synthesize one from the picker's MIME so `normalizeAttachmentExtension` + // can resolve a known key. The actual byte size is re-measured by the + // upload hook via `getInfoAsync`; `size` here is informational. + const fallbackName = `image.${(asset.mimeType ?? 'image/png').split('/')[1] ?? 'png'}`; + const name = asset.fileName ?? fallbackName; return { name, uri: asset.uri, @@ -47,14 +45,21 @@ function normalizeDocumentAsset(asset: { mimeType?: string; size?: number; }): AgentAttachmentCandidate { - const dot = asset.name.lastIndexOf('.'); - const ext = - dot > 0 ? (asset.name.slice(dot + 1).toLowerCase() as AgentAttachmentExtension) : null; - const mimeType = asset.mimeType ?? (ext ? mimeTypeForExtension(ext) : undefined); + // The picker MIME is intentionally NOT consulted. The cloud-agent + // storage layer rejects anything outside the canonical extension + // table, and iOS pickers report `application/octet-stream` for any + // extension the platform doesn't ship a UTI for. Resolving MIME from + // the extension makes the picker → upload hook contract exact. + const extension = normalizeAttachmentExtension(asset.name); return { name: asset.name, uri: asset.uri, - mimeType, + // The candidate shape carries MIME for kilochat-picker parity, but + // the agent-attachments classifier ignores it and re-derives from + // the extension. No closed-union cast — the extension is whatever + // survives `normalizeAttachmentExtension`, including the `bin` + // fallback. + mimeType: mimeForExtension(extension), size: asset.size, }; } diff --git a/apps/mobile/src/components/agents/attachment-preview-strip.test.ts b/apps/mobile/src/components/agents/attachment-preview-strip.test.ts new file mode 100644 index 0000000000..ab17c0e4c2 --- /dev/null +++ b/apps/mobile/src/components/agents/attachment-preview-strip.test.ts @@ -0,0 +1,183 @@ +import { describe, expect, it } from 'vitest'; + +import { + type ChipStateInput, + describeAttachmentChip, + progressLabel, +} from './attachment-chip-description'; + +const baseState: ChipStateInput = { + filename: 'doc.pdf', + size: 1024, + status: 'uploaded', + progress: 1, +}; + +describe('progressLabel', () => { + it('renders a percentage for determinate progress', () => { + expect(progressLabel(0.42)).toBe('42%'); + expect(progressLabel(0)).toBe('0%'); + }); + + it('renders "Uploaded" once progress reaches 1', () => { + expect(progressLabel(1)).toBe('Uploaded'); + }); + + it('renders "Uploading…" when progress is null (indeterminate)', () => { + expect(progressLabel(null)).toBe('Uploading…'); + }); +}); + +describe('describeAttachmentChip — happy / in-flight', () => { + it('shows filename and human-readable size for a fully uploaded chip', () => { + const desc = describeAttachmentChip(baseState); + expect(desc.filename).toBe('doc.pdf'); + expect(desc.sizeText).toBe('1 KB'); + expect(desc.progressText).toBe('Uploaded'); + expect(desc.message).toBeNull(); + expect(desc.showRetry).toBe(false); + expect(desc.showRemove).toBe(true); + }); + + it('shows determinate progress while uploading', () => { + const desc = describeAttachmentChip({ + ...baseState, + status: 'uploading', + progress: 0.42, + }); + expect(desc.progressText).toBe('42%'); + expect(desc.message).toBeNull(); + expect(desc.showRetry).toBe(false); + expect(desc.showRemove).toBe(true); + // Size and filename are still surfaced so the user knows what is uploading. + expect(desc.sizeText).toBe('1 KB'); + expect(desc.filename).toBe('doc.pdf'); + }); + + it('shows indeterminate progress honestly when the server did not advertise a total', () => { + const desc = describeAttachmentChip({ + ...baseState, + status: 'uploading', + progress: null, + }); + expect(desc.progressText).toBe('Uploading…'); + expect(desc.message).toBeNull(); + expect(desc.showRetry).toBe(false); + }); + + it('surfaces filename + human size for a pending chip before any progress is known', () => { + const desc = describeAttachmentChip({ + ...baseState, + status: 'pending', + progress: 0, + }); + expect(desc.filename).toBe('doc.pdf'); + expect(desc.sizeText).toBe('1 KB'); + expect(desc.progressText).toBe('0%'); + expect(desc.message).toBeNull(); + }); +}); + +describe('describeAttachmentChip — unhappy, retryable (network/timeout/408/429/5xx/PUT)', () => { + it('shows the retry copy and enables the tap-to-retry affordance', () => { + const desc = describeAttachmentChip({ + ...baseState, + status: 'error', + progress: null, + terminal: false, + }); + expect(desc.message).toBe('Upload failed. Tap to retry.'); + expect(desc.showRetry).toBe(true); + expect(desc.showRemove).toBe(true); + // Size + filename are still surfaced so the user knows which file failed. + expect(desc.sizeText).toBe('1 KB'); + expect(desc.filename).toBe('doc.pdf'); + }); + + it('keeps showRetry=false when terminal is undefined (defensive default = non-retryable)', () => { + // A missing `terminal` flag must not silently turn into a retryable chip; + // the explicit contract is `terminal: true` for permanent failures. + const desc = describeAttachmentChip({ + ...baseState, + status: 'error', + progress: null, + }); + expect(desc.showRetry).toBe(false); + expect(desc.message).toBe("This file can't be uploaded."); + }); +}); + +describe('describeAttachmentChip — unhappy, non-retryable (presign BAD_REQUEST/FORBIDDEN/UNPROCESSABLE_CONTENT)', () => { + it('shows the terminal copy and disables the tap-to-retry affordance', () => { + const desc = describeAttachmentChip({ + ...baseState, + status: 'error', + progress: null, + terminal: true, + }); + expect(desc.message).toBe("This file can't be uploaded."); + expect(desc.showRetry).toBe(false); + // Remove (X) MUST stay available so the user can clear the chip and + // unblock send. + expect(desc.showRemove).toBe(true); + // Filename + size are still surfaced so the user knows which file is + // permanently rejected. + expect(desc.filename).toBe('doc.pdf'); + expect(desc.sizeText).toBe('1 KB'); + }); + + it('keeps progressText empty for terminal failures (no percent leaked into the message)', () => { + const desc = describeAttachmentChip({ + ...baseState, + status: 'error', + progress: 0.5, + terminal: true, + }); + expect(desc.progressText).toBe(''); + }); +}); + +describe('describeAttachmentChip — feature-state invariants', () => { + const states: ChipStateInput[] = [ + { ...baseState, status: 'pending', progress: 0 }, + { ...baseState, status: 'uploading', progress: 0.42 }, + { ...baseState, status: 'uploaded', progress: 1 }, + { ...baseState, status: 'error', terminal: false, progress: null }, + { ...baseState, status: 'error', terminal: true, progress: null }, + ]; + + for (const state of states) { + it(`always surfaces filename + human size for status=${state.status} terminal=${Boolean(state.terminal)}`, () => { + const desc = describeAttachmentChip(state); + expect(desc.filename).toBe(state.filename); + expect(desc.sizeText.length).toBeGreaterThan(0); + }); + } + + for (const state of states) { + it(`always keeps the remove (X) affordance for status=${state.status} terminal=${Boolean(state.terminal)}`, () => { + const desc = describeAttachmentChip(state); + expect(desc.showRemove).toBe(true); + }); + } + + it('never merges retryable and non-retryable error affordances', () => { + const retryable = describeAttachmentChip({ + ...baseState, + status: 'error', + terminal: false, + }); + const terminal = describeAttachmentChip({ + ...baseState, + status: 'error', + terminal: true, + }); + expect(retryable.showRetry).toBe(true); + expect(terminal.showRetry).toBe(false); + // Both keep the remove affordance — that is the only affordance in common. + expect(retryable.showRemove).toBe(true); + expect(terminal.showRemove).toBe(true); + // The error messages must not collapse to the same string. + expect(retryable.message).not.toBe(terminal.message); + }); +}); diff --git a/apps/mobile/src/components/agents/attachment-preview-strip.tsx b/apps/mobile/src/components/agents/attachment-preview-strip.tsx index c07fcff51a..ede100931f 100644 --- a/apps/mobile/src/components/agents/attachment-preview-strip.tsx +++ b/apps/mobile/src/components/agents/attachment-preview-strip.tsx @@ -1,4 +1,3 @@ -import { formatFileSize } from '@kilocode/kilo-chat'; import { ActivityIndicator, Pressable, ScrollView, View } from 'react-native'; import { AlertCircle, File as FileIcon, X } from 'lucide-react-native'; @@ -7,6 +6,7 @@ import { Text } from '@/components/ui/text'; import { useThemeColors } from '@/lib/hooks/use-theme-colors'; import { cn } from '@/lib/utils'; import { type AgentAttachment } from '@/lib/agent-attachments/use-agent-attachment-upload'; +import { describeAttachmentChip } from '@/components/agents/attachment-chip-description'; type Props = { attachments: AgentAttachment[]; @@ -29,19 +29,27 @@ function AttachmentChip({ const isImage = attachment.kind === 'image'; const isUploading = attachment.status === 'pending' || attachment.status === 'uploading'; const isErrored = attachment.status === 'error'; + const description = describeAttachmentChip({ + filename: attachment.filename, + size: attachment.size, + status: attachment.status, + progress: attachment.progress, + terminal: attachment.terminal, + }); return ( - {attachment.filename} + {description.filename} - {formatFileSize(attachment.size)} + {description.message ?? `${description.sizeText} · ${description.progressText}`} @@ -83,15 +91,17 @@ function AttachmentChip({ ) : null} - - - + {description.showRemove ? ( + + + + ) : null} ); } @@ -107,6 +117,7 @@ export function AttachmentPreviewStrip({ attachments, onRemove, onRetry }: Reado className="mb-2" contentContainerClassName="items-center" keyboardShouldPersistTaps="handled" + accessibilityRole="summary" > {attachments.map(attachment => ( void | Promise; + onSend: ( + text: string, + attachments?: AgentAttachmentWire, + submission?: AgentAttachmentSubmissionPayload + ) => void | Promise; onSendCommand: (command: string, argumentsText: string) => Promise; onCreateSession: () => Promise; onExitCli: (onAccepted: () => void) => Promise; @@ -249,7 +254,7 @@ export function ChatComposer({ onExitCli, confirmExitCli: showRemoteCliExitConfirmation, onSendPrompt: async prompt => { - await onSend(prompt, upload.toWirePayload()); + await onSend(prompt, upload.toWirePayload(), upload.toSubmissionPayload()); }, }, { @@ -315,7 +320,10 @@ export function ChatComposer({ const { addCandidates, removeAttachment, retryAttachment } = upload; const handleAddAttachment = useCallback(async () => { - addCandidates(await pickAgentAttachments(showActionSheetWithOptions)); + // Fire-and-forget: the upload hook owns its own progress + error toasts, + // and the composer's send flow consults `upload.isUploading` / + // `upload.hasFailedAttachments` to gate admission. + void addCandidates(await pickAgentAttachments(showActionSheetWithOptions)); }, [addCandidates, showActionSheetWithOptions]); const textInputStyle: TextStyle = { diff --git a/apps/mobile/src/components/agents/mobile-session-manager-helpers.ts b/apps/mobile/src/components/agents/mobile-session-manager-helpers.ts new file mode 100644 index 0000000000..8b93d69ae2 --- /dev/null +++ b/apps/mobile/src/components/agents/mobile-session-manager-helpers.ts @@ -0,0 +1,45 @@ +import { type RemoteAttachmentPart } from 'cloud-agent-sdk'; + +import { type AgentAttachmentSubmissionPayload } from '@/lib/agent-attachments/agent-attachment-types'; +import { mimeForExtension, normalizeAttachmentExtension } from '@/lib/agent-attachments/validate'; +import { trpcClient } from '@/lib/trpc'; + +/** + * Build the `RemoteAttachmentPart[]` payload for a CAPABLE remote CLI + * session (the active CLI advertised `capabilities.attachments: true`). + * For each `file` in the composer's submission payload this mints a + * presigned GET via `trpcClient.cloudAgentNext.getAttachmentDownloadUrl` + * and returns the wire part the SDK appends to `send_message.parts` + * (after the text part). + * + * Contract: + * - `filename` on the wire is the SERVER-ISSUED `remoteName` + * (`.`), NEVER the original picker filename. + * - `mime` is derived SOLELY from the `remoteName` extension via the + * canonical table in `agent-attachments/validate` — the picker MIME + * is not consulted. + * + * Sequencing: callers `await` the full array; one failed presign + * propagates so the composer can surface the failure through the + * existing send pipeline. + */ +export async function buildRemoteAttachmentParts( + submission: AgentAttachmentSubmissionPayload +): Promise { + const parts = await Promise.all( + submission.files.map(async file => { + const result = await trpcClient.cloudAgentNext.getAttachmentDownloadUrl.mutate({ + messageUuid: submission.messageUuid, + filename: file.remoteName, + }); + const mime = mimeForExtension(normalizeAttachmentExtension(file.remoteName)); + return { + type: 'file', + mime, + filename: file.remoteName, + url: result.signedUrl, + } satisfies RemoteAttachmentPart; + }) + ); + return parts; +} diff --git a/apps/mobile/src/components/agents/mobile-session-manager.test.ts b/apps/mobile/src/components/agents/mobile-session-manager.test.ts index a664cb53f2..ffb80238b2 100644 --- a/apps/mobile/src/components/agents/mobile-session-manager.test.ts +++ b/apps/mobile/src/components/agents/mobile-session-manager.test.ts @@ -1,275 +1,68 @@ -import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { createStore } from 'jotai'; -import { type UserWebConnection } from 'cloud-agent-sdk'; +import { describe, expect, it, vi } from 'vitest'; -const mocks = vi.hoisted(() => ({ - createSessionManager: vi.fn(config => ({ config })), - createNativeUserWebConnectionLifecycleHooks: vi.fn(() => ({ marker: 'native-lifecycle-hooks' })), - getWithRuntimeStateQuery: vi.fn(), - getSessionQuery: vi.fn(), - getSessionMessagesQuery: vi.fn(), - sendMessageMutate: vi.fn(), - prepareSessionMutate: vi.fn(), -})); +import { type AgentAttachmentSubmissionPayload } from '@/lib/agent-attachments/agent-attachment-types'; +import { trpcClient } from '@/lib/trpc'; +import { buildRemoteAttachmentParts } from '@/components/agents/mobile-session-manager-helpers'; -function noCleanup(): void { - return undefined; -} - -const userWebConnection: UserWebConnection = { - retain: vi.fn(() => noCleanup), - connect: vi.fn(() => undefined), - disconnect: vi.fn(() => undefined), - destroy: vi.fn(() => undefined), - subscribeToCliSession: vi.fn(() => noCleanup), - sendCommand: vi.fn(), - sendCommandToConnection: vi.fn(), - onCliEvent: vi.fn(() => noCleanup), - onSystemEvent: vi.fn(() => noCleanup), - onReconnect: vi.fn(() => noCleanup), - onSessionEvent: vi.fn(() => noCleanup), -}; - -vi.mock('cloud-agent-sdk', () => ({ - createSessionManager: mocks.createSessionManager, -})); - -vi.mock('expo-secure-store', () => ({ - getItemAsync: vi.fn(), -})); - -vi.mock('sonner-native', () => ({ - toast: { error: vi.fn() }, -})); - -vi.mock('@/components/agents/mode-options', () => ({ - normalizeAgentMode: vi.fn(mode => mode), -})); - -vi.mock('@/components/agents/mobile-session-diagnostics', () => ({ - formatSafeCloudAgentFailureDiagnostic: vi.fn(() => null), - withCloudAgentDiagnostics: vi.fn( - async (_operation: string, _organizationId: string | undefined, run: () => Promise) => { - const result = await run(); - return result; - } - ), -})); - -vi.mock('@/lib/user-web-connection-lifecycle', () => ({ - createNativeUserWebConnectionLifecycleHooks: mocks.createNativeUserWebConnectionLifecycleHooks, -})); - -vi.mock('@/lib/config', () => ({ - API_BASE_URL: 'https://api.example.com', - CLOUD_AGENT_WS_URL: 'wss://agent.example.com', - WEB_BASE_URL: 'https://web.example.com', -})); - -vi.mock('@/lib/trpc', () => ({ - trpcClient: { - cliSessionsV2: { - get: { query: mocks.getSessionQuery }, - getSessionMessages: { query: mocks.getSessionMessagesQuery }, - getWithRuntimeState: { query: mocks.getWithRuntimeStateQuery }, - }, - activeSessions: { - list: { query: vi.fn() }, - }, - cloudAgentNext: { - sendMessage: { mutate: mocks.sendMessageMutate }, - prepareSession: { mutate: mocks.prepareSessionMutate }, - }, - organizations: { +vi.mock('@/lib/trpc', () => { + const mutate = vi.fn(); + return { + trpcClient: { cloudAgentNext: { - sendMessage: { mutate: mocks.sendMessageMutate }, - prepareSession: { mutate: mocks.prepareSessionMutate }, + getAttachmentDownloadUrl: { mutate }, }, }, - }, -})); - -type CapturedSessionManagerConfig = { - userWebConnection: unknown; - cliWebsocketUrl?: string; - getAuthToken?: () => Promise; - lifecycleHooks?: unknown; - fetchSession: (kiloSessionId: string) => Promise<{ associatedPr: unknown }>; - fetchSnapshot: (kiloSessionId: string) => Promise<{ info: unknown; messages: unknown[] }>; - fetchSnapshotPage: (kiloSessionId: string, options: { cursor?: string }) => Promise; - prepare: (input: { - prompt: string; - mode: string; - model: string; - initialPayload?: unknown; - }) => Promise; -}; - -describe('createMobileAgentSessionManager', () => { - beforeEach(() => { - vi.clearAllMocks(); - mocks.prepareSessionMutate.mockResolvedValue({ - cloudAgentSessionId: 'agent_123', - kiloSessionId: 'ses_123', - }); - }); - - it('injects the app-scoped user web connection without raw viewer transport options', async () => { - const { createMobileAgentSessionManager } = - await import('@/components/agents/mobile-session-manager'); - - createMobileAgentSessionManager({ - store: createStore(), - userWebConnection, - }); - - const config = mocks.createSessionManager.mock.calls[0]?.[0] as CapturedSessionManagerConfig; - expect(config.userWebConnection).toBe(userWebConnection); - expect(config.cliWebsocketUrl).toBeUndefined(); - expect(config.getAuthToken).toBeUndefined(); - }); - - it('passes native lifecycle hooks to Cloud Agent stream connections', async () => { - const { createMobileAgentSessionManager } = - await import('@/components/agents/mobile-session-manager'); + }; +}); - createMobileAgentSessionManager({ - store: createStore(), - userWebConnection, +describe('buildRemoteAttachmentParts', () => { + it('mints a presigned URL per file using remoteName and derives MIME from extension', async () => { + const mutate = vi.mocked(trpcClient.cloudAgentNext.getAttachmentDownloadUrl.mutate); + mutate.mockResolvedValue({ + signedUrl: 'https://r2.example.com/signed', + key: 'user-id/cloud-agent/msg-uuid/file', + expiresAt: new Date().toISOString(), }); - const config = mocks.createSessionManager.mock.calls[0]?.[0] as CapturedSessionManagerConfig; - expect(mocks.createNativeUserWebConnectionLifecycleHooks).toHaveBeenCalledTimes(1); - expect(config.lifecycleHooks).toEqual({ marker: 'native-lifecycle-hooks' }); - }); - - it('converts an initial Kilo model ref to the Cloud Agent prepare payload', async () => { - const { createMobileAgentSessionManager } = - await import('@/components/agents/mobile-session-manager'); + const submission: AgentAttachmentSubmissionPayload = { + messageUuid: 'msg-uuid', + wire: { + path: 'msg-uuid', + files: ['msg-uuid.zip', 'msg-uuid.txt', 'msg-uuid.png'], + }, + files: [ + { remoteName: 'msg-uuid.zip', originalName: 'archive.zip', size: 100 }, + { remoteName: 'msg-uuid.txt', originalName: 'notes.txt', size: 50 }, + { remoteName: 'msg-uuid.png', originalName: 'image.png', size: 200 }, + ], + }; - createMobileAgentSessionManager({ - store: createStore(), - userWebConnection, - }); + const parts = await buildRemoteAttachmentParts(submission); - const config = mocks.createSessionManager.mock.calls[0]?.[0] as CapturedSessionManagerConfig; - await config.prepare({ - prompt: 'Initial prompt', - mode: 'code', - model: 'fallback-model', - initialPayload: { - type: 'prompt', - prompt: 'Initial prompt', - mode: 'code', - model: { providerID: 'kilo', modelID: 'anthropic/claude-sonnet-4' }, - variant: 'high', - }, - }); + expect(mutate).toHaveBeenCalledTimes(3); + expect(mutate).toHaveBeenCalledWith({ messageUuid: 'msg-uuid', filename: 'msg-uuid.zip' }); + expect(mutate).toHaveBeenCalledWith({ messageUuid: 'msg-uuid', filename: 'msg-uuid.txt' }); + expect(mutate).toHaveBeenCalledWith({ messageUuid: 'msg-uuid', filename: 'msg-uuid.png' }); - expect(mocks.prepareSessionMutate).toHaveBeenCalledWith( + expect(parts).toEqual([ { - prompt: 'Initial prompt', - mode: 'code', - model: 'fallback-model', - initialPayload: { - type: 'prompt', - prompt: 'Initial prompt', - mode: 'code', - model: 'anthropic/claude-sonnet-4', - variant: 'high', - }, + type: 'file', + mime: 'application/octet-stream', + filename: 'msg-uuid.zip', + url: 'https://r2.example.com/signed', }, - { context: { skipBatch: true } } - ); - }); - - it('rejects a non-Kilo initial model ref before Cloud Agent prepare', async () => { - const { createMobileAgentSessionManager } = - await import('@/components/agents/mobile-session-manager'); - - createMobileAgentSessionManager({ - store: createStore(), - userWebConnection, - }); - - const config = mocks.createSessionManager.mock.calls[0]?.[0] as CapturedSessionManagerConfig; - await expect( - config.prepare({ - prompt: 'Initial prompt', - mode: 'code', - model: 'fallback-model', - initialPayload: { - type: 'prompt', - prompt: 'Initial prompt', - mode: 'code', - model: { providerID: 'anthropic', modelID: 'claude-sonnet-4' }, - }, - }) - ).rejects.toThrow('Cloud Agent only supports Kilo models'); - expect(mocks.prepareSessionMutate).not.toHaveBeenCalled(); - }); - - it('preserves snapshot model metadata from the session export', async () => { - const { createMobileAgentSessionManager } = - await import('@/components/agents/mobile-session-manager'); - mocks.getSessionQuery.mockResolvedValue({ - session_id: 'ses_123', - parent_session_id: null, - }); - mocks.getSessionMessagesQuery.mockResolvedValue({ - info: { - id: 'ses_123', - model: { providerID: 'anthropic', id: 'claude-sonnet-4', variant: 'high' }, + { + type: 'file', + mime: 'text/plain', + filename: 'msg-uuid.txt', + url: 'https://r2.example.com/signed', }, - messages: [], - }); - - createMobileAgentSessionManager({ - store: createStore(), - userWebConnection, - }); - - const config = mocks.createSessionManager.mock.calls[0]?.[0] as CapturedSessionManagerConfig; - await expect(config.fetchSnapshot('ses_123')).resolves.toEqual({ - info: { - id: 'ses_123', - parentID: undefined, - model: { providerID: 'anthropic', id: 'claude-sonnet-4', variant: 'high' }, + { + type: 'file', + mime: 'image/png', + filename: 'msg-uuid.png', + url: 'https://r2.example.com/signed', }, - messages: [], - }); - }); - - it('propagates associatedPr from fetched session data', async () => { - const { createMobileAgentSessionManager } = - await import('@/components/agents/mobile-session-manager'); - const associatedPr = { - url: 'https://github.com/Kilo-Org/cloud/pull/3383', - number: 3383, - state: 'open', - title: 'Refactor cloud agent session management', - headSha: 'abc123', - lastSyncedAt: '2026-05-22T20:00:00.000Z', - }; - - mocks.getWithRuntimeStateQuery.mockResolvedValue({ - cloud_agent_session_id: 'agent_123', - title: 'Session title', - organization_id: null, - git_url: 'https://github.com/Kilo-Org/cloud.git', - git_branch: 'feature/pr', - associatedPr, - runtimeState: null, - }); - - createMobileAgentSessionManager({ - store: createStore(), - userWebConnection, - }); - - const config = mocks.createSessionManager.mock.calls[0]?.[0] as CapturedSessionManagerConfig; - const session = await config.fetchSession('ses_123'); - - expect(session.associatedPr).toBe(associatedPr); + ]); }); }); diff --git a/apps/mobile/src/components/agents/mobile-session-manager.ts b/apps/mobile/src/components/agents/mobile-session-manager.ts index 772c631c4c..4b2b380e2b 100644 --- a/apps/mobile/src/components/agents/mobile-session-manager.ts +++ b/apps/mobile/src/components/agents/mobile-session-manager.ts @@ -9,19 +9,17 @@ import { type ResolvedSession, type SessionManager, type SessionSnapshot, - type TransportSendPayload, type UserWebConnection, } from 'cloud-agent-sdk'; -import { normalizeAgentMode } from '@/components/agents/mode-options'; +import { normalizeTransportPayload } from '@/components/agents/mobile-session-transport-payload'; import { formatSafeCloudAgentFailureDiagnostic, withCloudAgentDiagnostics, } from '@/components/agents/mobile-session-diagnostics'; import { fetchMobileSessionSnapshotPage } from '@/components/agents/mobile-session-page-adapter'; -import { trpcClient } from '@/lib/trpc'; import { API_BASE_URL, CLOUD_AGENT_WS_URL, WEB_BASE_URL } from '@/lib/config'; +import { trpcClient } from '@/lib/trpc'; import { AUTH_TOKEN_KEY } from '@/lib/storage-keys'; -import { type SendMessagePayload } from '@/lib/cloud-agent-next/types'; import { createNativeUserWebConnectionLifecycleHooks } from '@/lib/user-web-connection-lifecycle'; type CreateMobileAgentSessionManagerOptions = { @@ -34,31 +32,6 @@ type AgentMode = 'code' | 'plan' | 'debug' | 'orchestrator' | 'ask'; const skipBatchOptions = { context: { skipBatch: true } }; -function normalizeTransportPayload(payload: TransportSendPayload): SendMessagePayload { - if (payload.type === 'prompt') { - if (!payload.model) { - throw new Error('Model is required'); - } - if (payload.model.providerID !== 'kilo') { - throw new Error('Cloud Agent only supports Kilo models'); - } - - return { - type: 'prompt', - prompt: payload.prompt, - mode: normalizeAgentMode(payload.mode), - model: payload.model.modelID, - variant: payload.variant, - }; - } - - return { - type: 'command', - command: payload.command, - arguments: payload.arguments, - }; -} - export function createMobileAgentSessionManager({ store, userWebConnection, @@ -84,8 +57,21 @@ export function createMobileAgentSessionManager({ }; } const active = await trpcClient.activeSessions.list.query(); - const isRemote = active.sessions.some(s => s.id === kiloSessionId); - return { type: isRemote ? 'remote' : 'read-only', kiloSessionId }; + const activeSession = active.sessions.find(s => s.id === kiloSessionId); + if (!activeSession) { + return { type: 'read-only', kiloSessionId }; + } + // Surface the owning CLI's per-session capabilities so the initial + // `supportsAttachments` gate reflects whatever the mobile adapter + // observed at resolution time. Heartbeat upgrades / downgrades + // arrive later via `onTransportCapabilitiesChange` from the + // cli-live-transport; the seed here just covers the window before + // the first heartbeat lands. + return { + type: 'remote', + kiloSessionId, + ...(activeSession.capabilities ? { capabilities: activeSession.capabilities } : {}), + }; }, getTicket: async (sessionId: CloudAgentSessionId): Promise => { const ticket = await withCloudAgentDiagnostics('getTicket', organizationId, async () => { diff --git a/apps/mobile/src/components/agents/mobile-session-transport-payload.ts b/apps/mobile/src/components/agents/mobile-session-transport-payload.ts new file mode 100644 index 0000000000..32580a7c68 --- /dev/null +++ b/apps/mobile/src/components/agents/mobile-session-transport-payload.ts @@ -0,0 +1,37 @@ +import { type TransportSendPayload } from 'cloud-agent-sdk'; + +import { normalizeAgentMode } from '@/components/agents/mode-options'; +import { type SendMessagePayload } from '@/lib/cloud-agent-next/types'; + +/** + * Normalize a transport send payload into the wire `SendMessagePayload`. + * + * Kept in its own module (rather than alongside `buildRemoteAttachmentParts`) + * because it depends on `mode-options`, which transitively pulls in React + * Native / Expo modules. Isolating it keeps the attachment helper importable + * from the Node-based unit test environment. + */ +export function normalizeTransportPayload(payload: TransportSendPayload): SendMessagePayload { + if (payload.type === 'prompt') { + if (!payload.model) { + throw new Error('Model is required'); + } + if (payload.model.providerID !== 'kilo') { + throw new Error('Cloud Agent only supports Kilo models'); + } + + return { + type: 'prompt', + prompt: payload.prompt, + mode: normalizeAgentMode(payload.mode), + model: payload.model.modelID, + variant: payload.variant, + }; + } + + return { + type: 'command', + command: payload.command, + arguments: payload.arguments, + }; +} diff --git a/apps/mobile/src/components/agents/session-detail-content.test.ts b/apps/mobile/src/components/agents/session-detail-content.test.ts new file mode 100644 index 0000000000..dc6981ea13 --- /dev/null +++ b/apps/mobile/src/components/agents/session-detail-content.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, it } from 'vitest'; + +import { resolveSendAttachmentKind } from '@/components/agents/session-detail-send-attachment'; + +describe('resolveSendAttachmentKind', () => { + it.each([ + { activeSessionType: 'cloud-agent' as const, supports: true, has: true, expected: 'cloud' }, + { activeSessionType: 'cloud-agent' as const, supports: false, has: true, expected: 'cloud' }, + { activeSessionType: 'remote' as const, supports: true, has: true, expected: 'remote-capable' }, + { activeSessionType: 'remote' as const, supports: false, has: true, expected: 'none' }, + { activeSessionType: 'read-only' as const, supports: true, has: true, expected: 'none' }, + { activeSessionType: null, supports: true, has: true, expected: 'none' }, + { activeSessionType: undefined, supports: true, has: true, expected: 'none' }, + { activeSessionType: 'cloud-agent' as const, supports: true, has: false, expected: 'none' }, + { activeSessionType: 'remote' as const, supports: true, has: false, expected: 'none' }, + ])( + 'returns $expected for sessionType=$activeSessionType, supports=$supports, has=$has', + ({ activeSessionType, supports, has, expected }) => { + expect(resolveSendAttachmentKind(activeSessionType, supports, has)).toBe(expected); + } + ); +}); diff --git a/apps/mobile/src/components/agents/session-detail-content.tsx b/apps/mobile/src/components/agents/session-detail-content.tsx index 10e1de9734..ce834a7f3b 100644 --- a/apps/mobile/src/components/agents/session-detail-content.tsx +++ b/apps/mobile/src/components/agents/session-detail-content.tsx @@ -27,6 +27,11 @@ import { SessionContextMetrics, } from '@/components/agents/session-context-metrics'; import { SessionContextSheet } from '@/components/agents/session-context-sheet'; +import { buildRemoteAttachmentParts } from '@/components/agents/mobile-session-manager-helpers'; +import { + buildRemoteAttachmentPartsWithRetryableFeedback, + resolveSendAttachmentKind, +} from '@/components/agents/session-detail-send-attachment'; import { useSessionManager } from '@/components/agents/session-provider'; import { SessionStatusIndicator } from '@/components/agents/session-status-indicator'; import { PreparationGroup } from '@/components/agents/preparation-group'; @@ -58,6 +63,7 @@ import { ScreenHeader } from '@/components/screen-header'; import { Button } from '@/components/ui/button'; import { Skeleton } from '@/components/ui/skeleton'; import { Text } from '@/components/ui/text'; +import { type AgentAttachmentSubmissionPayload } from '@/lib/agent-attachments/agent-attachment-types'; import { type AgentAttachmentWire } from '@/lib/agent-attachments/use-agent-attachment-upload'; import { type AnalyticsSurface, @@ -460,11 +466,42 @@ export function SessionDetailContent({ const keyboardContainerKind = getSessionKeyboardContainerKind(Platform.OS); const handleSend = useCallback( - async (text: string, attachments?: AgentAttachmentWire) => { + async ( + text: string, + attachments?: AgentAttachmentWire, + submission?: AgentAttachmentSubmissionPayload + ) => { if (requiresModel && !currentModel) { toast.error('Select a model before sending'); return; } + // Pick the wire shape via the same pure helper the unit test covers: + // - cloud-agent → unchanged `{path, files}` (S3a) + // - remote + supportsAttachments → materialize presigned GETs and + // forward as `attachmentParts` (S3b) + // - everything else → no attachment field on the wire + const kind = resolveSendAttachmentKind( + activeSessionType, + supportsAttachments, + attachments !== undefined + ); + let attachmentParts: Awaited> | undefined = + undefined; + if (kind === 'remote-capable' && submission) { + const result = await buildRemoteAttachmentPartsWithRetryableFeedback( + submission, + buildRemoteAttachmentParts + ); + if (!result.ok) { + // Retryable presign failure: the manager never reached send(), so + // its onSendFailed toast does not fire. Surface the retryable message + // through the same toast channel and throw so the composer keeps the + // draft/attachments for a retry. + toast.error(result.message); + throw new Error(result.message); + } + attachmentParts = result.parts; + } // manager.send() reports failures via its own return value (and toasts // through the manager's onSendFailed hook) rather than rejecting — it // is the single toast owner for send failures. Throw here, without a @@ -478,7 +515,8 @@ export function SessionDetailContent({ model: currentModel, variant: currentVariant || undefined, }, - ...(supportsAttachments && attachments ? { attachments } : {}), + ...(kind === 'cloud' && attachments ? { attachments } : {}), + ...(kind === 'remote-capable' && attachmentParts ? { attachmentParts } : {}), }); if (!sent) { throw new Error('Failed to send message'); @@ -491,6 +529,7 @@ export function SessionDetailContent({ currentModel, currentVariant, requiresModel, + activeSessionType, supportsAttachments, analyticsSurface, ] diff --git a/apps/mobile/src/components/agents/session-detail-send-attachment.test.ts b/apps/mobile/src/components/agents/session-detail-send-attachment.test.ts new file mode 100644 index 0000000000..a72f81ead1 --- /dev/null +++ b/apps/mobile/src/components/agents/session-detail-send-attachment.test.ts @@ -0,0 +1,46 @@ +import { type RemoteAttachmentPart } from 'cloud-agent-sdk'; +import { describe, expect, it } from 'vitest'; + +import { buildRemoteAttachmentPartsWithRetryableFeedback } from './session-detail-send-attachment'; +import { type AgentAttachmentSubmissionPayload } from '@/lib/agent-attachments/agent-attachment-types'; + +const submission: AgentAttachmentSubmissionPayload = { + wire: { path: 'msg-uuid', files: ['server-name.txt'] }, + messageUuid: 'msg-uuid', + files: [ + { + remoteName: 'server-name.txt', + originalName: 'original-name.txt', + size: 1024, + }, + ], +}; + +describe('buildRemoteAttachmentPartsWithRetryableFeedback', () => { + it('returns parts when the underlying builder succeeds', async () => { + const parts: RemoteAttachmentPart[] = [ + { + type: 'file', + mime: 'text/plain', + filename: 'server-name.txt', + url: 'https://r2.example.com/signed.txt', + }, + ]; + const result = await buildRemoteAttachmentPartsWithRetryableFeedback(submission, async () => { + await Promise.resolve(); + return parts; + }); + expect(result).toEqual({ ok: true, parts }); + }); + + it('returns a retryable message when the underlying builder rejects', async () => { + const result = await buildRemoteAttachmentPartsWithRetryableFeedback(submission, async () => { + await Promise.resolve(); + throw new Error('presign network failure'); + }); + expect(result).toEqual({ + ok: false, + message: "Couldn't attach files. Tap send to try again.", + }); + }); +}); diff --git a/apps/mobile/src/components/agents/session-detail-send-attachment.ts b/apps/mobile/src/components/agents/session-detail-send-attachment.ts new file mode 100644 index 0000000000..6a82d4b869 --- /dev/null +++ b/apps/mobile/src/components/agents/session-detail-send-attachment.ts @@ -0,0 +1,69 @@ +import { type RemoteAttachmentPart, type ResolvedSession } from 'cloud-agent-sdk'; + +import { type AgentAttachmentSubmissionPayload } from '@/lib/agent-attachments/agent-attachment-types'; + +type BuildRemoteAttachmentPartsResult = + | { ok: true; parts: RemoteAttachmentPart[] } + | { ok: false; message: string }; + +/** + * Pure decision for the send path: given the active session type and + * capability state, pick which wire shape `manager.send()` should + * receive. Extracted from the component so the cloud-vs-remote gate + * is unit-testable without React. + * + * - `cloud` — pass the unchanged `{path, files}` S3a wire as + * `attachments`. Returned for cloud-agent sessions regardless of + * capability (S3a is a no-op for non-cloud but the existing branch + * stays so the wire is stable). + * - `remote-capable` — call sites must `await buildRemoteAttachmentParts` + * and pass the result as `attachmentParts`. + * - `none` — no attachments path applies; the call site must omit + * both `attachments` and `attachmentParts`. This covers non-capable + * remote sessions (paperclip hidden) and read-only sessions. + * + * `hasAttachments` is the composer's pre-computed signal: false means + * the user did not add any files in the first place, so the decision + * short-circuits to `none` regardless of session type. + */ +export function resolveSendAttachmentKind( + activeSessionType: ResolvedSession['type'] | null | undefined, + supportsAttachments: boolean, + hasAttachments: boolean +): 'cloud' | 'remote-capable' | 'none' { + if (!hasAttachments) { + return 'none'; + } + if (activeSessionType === 'cloud-agent') { + return 'cloud'; + } + if (activeSessionType === 'remote' && supportsAttachments) { + return 'remote-capable'; + } + return 'none'; +} + +/** + * Build remote attachment parts for a capable remote session, mapping a + * transient presign failure to a retryable user-facing message. The caller + * is responsible for surfacing `message` through the same toast/error surface + * used for send failures and for preserving the composer draft so the user + * can retry by sending again. + * + * `buildParts` is injected so the test can stub it without pulling in the + * real tRPC client. + */ +export async function buildRemoteAttachmentPartsWithRetryableFeedback( + submission: AgentAttachmentSubmissionPayload, + buildParts: (payload: AgentAttachmentSubmissionPayload) => Promise +): Promise { + try { + const parts = await buildParts(submission); + return { ok: true, parts }; + } catch { + return { + ok: false, + message: "Couldn't attach files. Tap send to try again.", + }; + } +} diff --git a/apps/mobile/src/lib/agent-attachments/agent-attachment-types.ts b/apps/mobile/src/lib/agent-attachments/agent-attachment-types.ts new file mode 100644 index 0000000000..b5e056a14a --- /dev/null +++ b/apps/mobile/src/lib/agent-attachments/agent-attachment-types.ts @@ -0,0 +1,170 @@ +/** + * Pure helpers extracted from `use-agent-attachment-upload` so the + * upload contract is unit-testable without React or React Native + * (vitest's node environment cannot parse the Flow-typed + * `react-native/index.js` that `sonner-native` and friends pull in). + * The hook re-exports these symbols so existing call sites keep + * importing from one place. + */ + +import { type AgentAttachmentExtension, type AgentAttachmentMime } from './constants'; + +export type AgentAttachmentKind = 'image' | 'document'; +export type AgentAttachmentStatus = 'pending' | 'uploading' | 'uploaded' | 'error'; + +export type AgentAttachment = { + id: string; + /** Original filename as supplied by the picker. Display-only. */ + filename: string; + /** Basename of the server-side R2 key; set once the upload succeeds. */ + remoteFilename?: string; + kind: AgentAttachmentKind; + /** Normalized extension used for the R2 key suffix and MIME derivation. */ + extension: AgentAttachmentExtension; + mimeType: AgentAttachmentMime; + /** + * Measured local byte size (via `getInfoAsync`). The picker-reported + * `size` is unreliable on iOS so we always re-measure before adding + * the candidate to state. + */ + size: number; + localUri: string; + status: AgentAttachmentStatus; + /** Human-readable error for retryable failures; the chip is the affordance. */ + error?: string; + /** True when the server rejected this attachment permanently. */ + terminal?: boolean; + /** 0..1 determinate progress; `null` when progress is unavailable. */ + progress: number | null; +}; + +export type AgentAttachmentWire = { + path: string; + files: string[]; +}; + +/** + * Composer submission payload. The wire is the existing + * `{path, files}` shape consumed by the cloud-agent SDK; `messageUuid` + * and the per-file descriptor are new in S2 and are the contract S3b + * consumes to materialize the cloud `{path, files}` materialization. + * + * NOTE: there is NO `mime` field on the descriptor. Every consumer + * derives MIME from the validated `remoteName` extension. + */ +export type AgentAttachmentSubmissionFile = { + remoteName: string; + originalName: string; + size: number; +}; + +export type AgentAttachmentSubmissionPayload = { + wire: AgentAttachmentWire; + messageUuid: string; + files: AgentAttachmentSubmissionFile[]; +}; + +/** + * Status code → retryable / terminal classifier (pinned by the plan). + * Terminal set: presign BAD_REQUEST / FORBIDDEN / UNPROCESSABLE_CONTENT / + * UNAUTHORIZED / NOT_FOUND — these are policy rejections and must never + * be retried. Anything else the upload task throws (network, timeout, + * 408/429/5xx, generic PUT failure) is retryable. + */ +export function classifyUploadFailure(error: unknown): { retryable: boolean; reason: string } { + // The mutation throws `TRPCClientError` for + // BAD_REQUEST / FORBIDDEN / UNPROCESSABLE_CONTENT — those are TERMINAL. + // Any other thrown object (network, timeout, expiry, etc.) is RETRYABLE. + const code = (error as { data?: { code?: string; message?: string } } | null)?.data?.code; + const dataMessage = (error as { data?: { code?: string; message?: string } } | null)?.data + ?.message; + if (code === 'BAD_REQUEST' || code === 'FORBIDDEN' || code === 'UNPROCESSABLE_CONTENT') { + return { retryable: false, reason: dataMessage ?? "This file can't be uploaded." }; + } + if (code === 'UNAUTHORIZED' || code === 'NOT_FOUND') { + return { retryable: false, reason: dataMessage ?? "This file can't be uploaded." }; + } + if (error instanceof TypeError) { + return { retryable: true, reason: 'Network error' }; + } + if (error instanceof Error) { + if (/abort|cancel|expir/i.test(error.message)) { + return { retryable: true, reason: 'Upload failed' }; + } + if (/status (408|429|5\d\d)/.test(error.message)) { + return { retryable: true, reason: 'Upload failed' }; + } + if (/status \d{3}/.test(error.message)) { + // Any other HTTP error (4xx other than the terminal codes above, + // 5xx, etc.) is retryable — the plan pins PUT failures as retryable. + return { retryable: true, reason: 'Upload failed' }; + } + return { retryable: true, reason: 'Upload failed' }; + } + return { retryable: true, reason: 'Upload failed' }; +} + +/** + * Build the `{path, files}` wire payload from the current attachments. + * Pure: the hook is just a memoized wrapper. + */ +export function buildWirePayload( + attachments: AgentAttachment[], + path: string +): AgentAttachmentWire | undefined { + const files = attachments + .filter(item => item.status === 'uploaded') + .map(item => item.remoteFilename) + .filter((filename): filename is string => filename !== undefined); + if (files.length === 0) { + return undefined; + } + return { path, files }; +} + +/** + * Build the S2 submission payload from the current attachments. `wire` is + * the existing `{path, files}` shape; `messageUuid` and the per-file + * descriptor are new in S2 and are what S3b consumes to materialize the + * cloud `{path, files}` materialization. There is NO `mime` field on the + * descriptor — every consumer derives MIME from the validated `remoteName` + * extension. Pure: the hook is just a memoized wrapper. + */ +export function buildSubmissionPayload( + attachments: AgentAttachment[], + path: string, + messageUuid: string +): AgentAttachmentSubmissionPayload | undefined { + const ready: { remoteFilename: string; filename: string; size: number }[] = []; + for (const item of attachments) { + if (item.status === 'uploaded' && item.remoteFilename !== undefined) { + ready.push({ + remoteFilename: item.remoteFilename, + filename: item.filename, + size: item.size, + }); + } + } + if (ready.length === 0) { + return undefined; + } + return { + wire: { path, files: ready.map(item => item.remoteFilename) }, + messageUuid, + files: ready.map(item => ({ + remoteName: item.remoteFilename, + originalName: item.filename, + size: item.size, + })), + }; +} + +/** True when any chip is mid-flight (pending or uploading). */ +export function isAnyAttachmentUploading(attachments: AgentAttachment[]): boolean { + return attachments.some(item => item.status === 'pending' || item.status === 'uploading'); +} + +/** True when any chip is in a failed state (retryable or terminal). */ +export function hasAnyFailedAttachment(attachments: AgentAttachment[]): boolean { + return attachments.some(item => item.status === 'error'); +} diff --git a/apps/mobile/src/lib/agent-attachments/constants.ts b/apps/mobile/src/lib/agent-attachments/constants.ts index 852dc29e9d..13aab50f92 100644 --- a/apps/mobile/src/lib/agent-attachments/constants.ts +++ b/apps/mobile/src/lib/agent-attachments/constants.ts @@ -1,28 +1,92 @@ export const AGENT_ATTACHMENT_MAX_FILES = 5; export const AGENT_ATTACHMENT_MAX_BYTES = 5 * 1024 * 1024; -export const AGENT_ATTACHMENT_EXTENSIONS = [ - 'png', - 'jpg', - 'jpeg', - 'webp', - 'gif', - 'pdf', - 'txt', - 'md', - 'csv', -] as const; +/** + * Extensions that are NEVER allowed regardless of the picker's reported MIME + * type. The cloud-agent storage layer rejects these as executables, so we + * mirror the deny list client-side to give a precise error at the chip + * instead of a round-trip rejection. + */ +export const AGENT_ATTACHMENT_DENIED_EXTENSIONS = new Set([ + 'exe', + 'dll', + 'msi', + 'com', + 'scr', + 'apk', + 'ipa', + 'dmg', + 'pkg', +]); -export type AgentAttachmentExtension = (typeof AGENT_ATTACHMENT_EXTENSIONS)[number]; +/** + * Normalized extensions must match this regex. Anything that does not is + * coerced to `bin` and treated as opaque binary. Keeping the bound tight + * protects the server's storage key and the wire contract. + */ +export const AGENT_ATTACHMENT_EXTENSION_REGEX = /^[a-z0-9]{1,16}$/; +/** + * Fallback extension when the candidate's filename has no parseable + * extension or the extension does not match the allowed regex. Consumers + * derive MIME from the extension, so the fallback MUST be present in + * `AGENT_ATTACHMENT_MIME_BY_EXTENSION`. + */ +export const AGENT_ATTACHMENT_FALLBACK_EXTENSION = 'bin'; + +/** + * Canonical extension → MIME table. MUST stay in lock-step with the server's + * allowed content type set; the parity test in `validate.test.ts` asserts + * every entry resolves to a defined MIME. + */ export const AGENT_ATTACHMENT_MIME_BY_EXTENSION = { + // Images png: 'image/png', jpg: 'image/jpeg', jpeg: 'image/jpeg', webp: 'image/webp', gif: 'image/gif', + // Documents pdf: 'application/pdf', + // Text-ish source — server treats all of these as text/plain txt: 'text/plain', - md: 'text/markdown', - csv: 'text/csv', -} as const satisfies Record; + md: 'text/plain', + csv: 'text/plain', + log: 'text/plain', + json: 'text/plain', + xml: 'text/plain', + yaml: 'text/plain', + yml: 'text/plain', + toml: 'text/plain', + ini: 'text/plain', + html: 'text/plain', + css: 'text/plain', + js: 'text/plain', + jsx: 'text/plain', + ts: 'text/plain', + tsx: 'text/plain', + py: 'text/plain', + rb: 'text/plain', + go: 'text/plain', + rs: 'text/plain', + java: 'text/plain', + c: 'text/plain', + h: 'text/plain', + cpp: 'text/plain', + hpp: 'text/plain', + sh: 'text/plain', + sql: 'text/plain', + // Opaque binary fallback + bin: 'application/octet-stream', +} as const satisfies Record; + +export type AgentAttachmentMime = + (typeof AGENT_ATTACHMENT_MIME_BY_EXTENSION)[keyof typeof AGENT_ATTACHMENT_MIME_BY_EXTENSION]; + +/** + * Any extension that survives normalization. Callers MUST go through + * `AGENT_ATTACHMENT_MIME_BY_EXTENSION` to resolve MIME — there is no + * closed union of accepted extensions, and a picker can supply extensions + * outside this list. + */ +export type AgentAttachmentExtension = keyof typeof AGENT_ATTACHMENT_MIME_BY_EXTENSION; diff --git a/apps/mobile/src/lib/agent-attachments/upload-task.test.ts b/apps/mobile/src/lib/agent-attachments/upload-task.test.ts new file mode 100644 index 0000000000..68a2d4f742 --- /dev/null +++ b/apps/mobile/src/lib/agent-attachments/upload-task.test.ts @@ -0,0 +1,97 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { trpcClient } from '@/lib/trpc'; +import { classifyAttachment } from './validate'; +import { uploadOne } from './upload-task'; + +vi.mock('@/lib/trpc', () => ({ + trpcClient: { + cloudAgentNext: { + getAttachmentUploadUrl: { mutate: vi.fn() }, + }, + organizations: { + cloudAgentNext: { + getAttachmentUploadUrl: { mutate: vi.fn() }, + }, + }, + }, +})); + +vi.mock('expo-file-system/legacy', () => ({ + createUploadTask: vi.fn(() => ({ + uploadAsync: vi.fn().mockResolvedValue({ status: 200 }), + })), + FileSystemUploadType: { BINARY_CONTENT: 'binary-content' }, + getInfoAsync: vi.fn().mockResolvedValue({ + exists: true, + isDirectory: false, + size: 100, + }), +})); + +describe('uploadOne', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(trpcClient.cloudAgentNext.getAttachmentUploadUrl.mutate).mockResolvedValue({ + signedUrl: 'https://r2.example.com/signed', + key: 'key.bin', + expiresAt: new Date().toISOString(), + }); + }); + + it('sends the classified extension to the presign (invalid-suffix filename → bin)', async () => { + const classified = classifyAttachment({ name: 'report.v-2', size: 100 }); + expect(classified.ok).toBe(true); + if (!classified.ok) { + return; + } + expect(classified.extension).toBe('bin'); + + await uploadOne({ + attachmentId: 'att-1', + path: 'path-1', + extension: classified.extension, + contentType: 'application/octet-stream', + contentLength: 100, + localUri: 'file:///cache/report.v-2.bin', + onProgress: () => undefined, + }); + + const mutate = vi.mocked(trpcClient.cloudAgentNext.getAttachmentUploadUrl.mutate); + expect(mutate).toHaveBeenCalledTimes(1); + expect(mutate).toHaveBeenCalledWith( + expect.objectContaining({ + messageUuid: 'path-1', + attachmentId: 'att-1', + contentType: 'application/octet-stream', + contentLength: 100, + extension: 'bin', + }) + ); + }); + + it('sends extension: "bin" for a suffix longer than 16 characters', async () => { + const classified = classifyAttachment({ + name: `archive.${'a'.repeat(32)}`, + size: 100, + }); + expect(classified.ok).toBe(true); + if (!classified.ok) { + return; + } + expect(classified.extension).toBe('bin'); + + await uploadOne({ + attachmentId: 'att-2', + path: 'path-2', + extension: classified.extension, + contentType: 'application/octet-stream', + contentLength: 100, + localUri: 'file:///cache/archive.bin', + onProgress: () => undefined, + }); + + const mutate = vi.mocked(trpcClient.cloudAgentNext.getAttachmentUploadUrl.mutate); + expect(mutate).toHaveBeenCalledWith(expect.objectContaining({ extension: 'bin' })); + }); +}); diff --git a/apps/mobile/src/lib/agent-attachments/upload-task.ts b/apps/mobile/src/lib/agent-attachments/upload-task.ts new file mode 100644 index 0000000000..ef64d87f2a --- /dev/null +++ b/apps/mobile/src/lib/agent-attachments/upload-task.ts @@ -0,0 +1,95 @@ +import { createUploadTask, FileSystemUploadType, getInfoAsync } from 'expo-file-system/legacy'; + +import { trpcClient } from '@/lib/trpc'; +import { + type AgentAttachmentExtension, + type AgentAttachmentMime, +} from '@/lib/agent-attachments/constants'; + +export function normalizeFilename(name: string, extension: AgentAttachmentExtension): string { + // If the original filename had no usable extension we append the + // normalized one so the display value and the server's R2 key agree. + const dot = name.lastIndexOf('.'); + if (dot > 0 && dot < name.length - 1) { + return name; + } + return `${name}.${extension}`; +} + +export async function measureLocalSize(uri: string): Promise { + try { + const info = await getInfoAsync(uri); + if (info.exists && !info.isDirectory) { + return info.size; + } + } catch { + return null; + } + return null; +} + +type UploadOutcome = { key: string }; + +/** + * Presign + PUT a single local file. Progress is reported via `onProgress` + * (`null` when the server omits Content-Length). + */ +export async function uploadOne(args: { + organizationId?: string; + attachmentId: string; + path: string; + extension: AgentAttachmentExtension; + contentType: AgentAttachmentMime; + contentLength: number; + localUri: string; + onProgress: (progress: number | null) => void; +}): Promise { + const { organizationId, attachmentId, path, contentType, contentLength, localUri, onProgress } = + args; + const baseInput = { + messageUuid: path, + attachmentId, + contentType, + contentLength, + extension: args.extension, + }; + const result = organizationId + ? await trpcClient.organizations.cloudAgentNext.getAttachmentUploadUrl.mutate({ + ...baseInput, + organizationId, + }) + : await trpcClient.cloudAgentNext.getAttachmentUploadUrl.mutate(baseInput); + + // Per-chip determinate progress via `createUploadTask` (the + // main-module `createUploadTask` throws at runtime in SDK 55, so we + // import from `expo-file-system/legacy`). A signed-URL PUT only + // reports progress when the response advertises Content-Length; we + // fall back to `null` (indeterminate) when the server omits it. + const task = createUploadTask( + result.signedUrl, + localUri, + { + uploadType: FileSystemUploadType.BINARY_CONTENT, + httpMethod: 'PUT', + headers: { 'Content-Type': contentType }, + }, + progress => { + const total = progress.totalBytesExpectedToSend; + if (total > 0) { + onProgress(progress.totalBytesSent / total); + } else { + onProgress(null); + } + } + ); + const uploadResult = await task.uploadAsync(); + if (!uploadResult || uploadResult.status < 200 || uploadResult.status >= 300) { + throw new Error(`Upload failed with status ${uploadResult?.status ?? 'no response'}`); + } + return { key: result.key }; +} + +/** Chip/toast copy for terminal (non-retryable) upload failures. */ +export function describeTerminalReason(_reason: string): string { + return "This file can't be uploaded."; +} diff --git a/apps/mobile/src/lib/agent-attachments/use-agent-attachment-upload.test.ts b/apps/mobile/src/lib/agent-attachments/use-agent-attachment-upload.test.ts new file mode 100644 index 0000000000..a2a5c73924 --- /dev/null +++ b/apps/mobile/src/lib/agent-attachments/use-agent-attachment-upload.test.ts @@ -0,0 +1,285 @@ +import { describe, expect, it } from 'vitest'; + +import { AGENT_ATTACHMENT_MAX_BYTES } from './constants'; +import { + type AgentAttachment, + type AgentAttachmentSubmissionPayload, + type AgentAttachmentWire, + buildSubmissionPayload, + buildWirePayload, + classifyUploadFailure, + hasAnyFailedAttachment, + isAnyAttachmentUploading, +} from './agent-attachment-types'; +// Tests import pure helpers from their owning module (not the hook barrel). + +function makeAttachment(overrides: Partial): AgentAttachment { + return { + id: 'a1', + filename: 'doc.pdf', + kind: 'document', + extension: 'pdf', + mimeType: 'application/pdf', + size: 1024, + localUri: 'file:///cache/doc.pdf', + status: 'uploaded', + progress: 1, + remoteFilename: 'org/2026/07/uuid/doc.pdf', + ...overrides, + }; +} + +describe('classifyUploadFailure — terminal (presign policy rejections)', () => { + it('marks BAD_REQUEST as terminal', () => { + expect( + classifyUploadFailure({ data: { code: 'BAD_REQUEST', message: 'extension not allowed' } }) + ).toEqual({ retryable: false, reason: 'extension not allowed' }); + }); + + it('marks FORBIDDEN as terminal', () => { + expect(classifyUploadFailure({ data: { code: 'FORBIDDEN', message: 'org policy' } })).toEqual({ + retryable: false, + reason: 'org policy', + }); + }); + + it('marks UNPROCESSABLE_CONTENT as terminal', () => { + expect( + classifyUploadFailure({ data: { code: 'UNPROCESSABLE_CONTENT', message: 'bad extension' } }) + ).toEqual({ retryable: false, reason: 'bad extension' }); + }); + + it('marks UNAUTHORIZED as terminal', () => { + expect(classifyUploadFailure({ data: { code: 'UNAUTHORIZED', message: 'expired' } })).toEqual({ + retryable: false, + reason: 'expired', + }); + }); + + it('marks NOT_FOUND as terminal', () => { + expect(classifyUploadFailure({ data: { code: 'NOT_FOUND', message: 'gone' } })).toEqual({ + retryable: false, + reason: 'gone', + }); + }); +}); + +describe('classifyUploadFailure — retryable (network/timeout/408/429/5xx/PUT)', () => { + it('marks a TypeError as retryable (network failure)', () => { + expect(classifyUploadFailure(new TypeError('Network request failed'))).toEqual({ + retryable: true, + reason: 'Network error', + }); + }); + + it('marks an abort/cancel/expiry message as retryable', () => { + expect(classifyUploadFailure(new Error('request aborted'))).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + expect(classifyUploadFailure(new Error('Upload was canceled'))).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + expect(classifyUploadFailure(new Error('URL expired'))).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + }); + + it('marks HTTP 408 / 429 / 5xx PUT failures as retryable', () => { + expect(classifyUploadFailure(new Error('Upload failed with status 408'))).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + expect(classifyUploadFailure(new Error('Upload failed with status 429'))).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + expect(classifyUploadFailure(new Error('Upload failed with status 500'))).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + expect(classifyUploadFailure(new Error('Upload failed with status 503'))).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + }); + + it('marks any other HTTP error as retryable (the plan pins PUT failures as retryable)', () => { + expect(classifyUploadFailure(new Error('Upload failed with status 418'))).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + }); + + it('marks an unknown thrown value as retryable with a generic reason', () => { + expect(classifyUploadFailure('something weird')).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + expect(classifyUploadFailure(undefined)).toEqual({ + retryable: true, + reason: 'Upload failed', + }); + }); + + it('never collapses retryable and terminal into a single bucket', () => { + const retryable = classifyUploadFailure(new TypeError('Network request failed')); + const terminal = classifyUploadFailure({ data: { code: 'BAD_REQUEST', message: 'x' } }); + expect(retryable.retryable).toBe(true); + expect(terminal.retryable).toBe(false); + }); +}); + +describe('buildWirePayload', () => { + it('returns undefined when no chips are uploaded', () => { + expect(buildWirePayload([], 'path-1')).toBeUndefined(); + }); + + it('returns undefined when chips are still uploading', () => { + const list = [ + makeAttachment({ status: 'uploading', progress: 0.5 }), + makeAttachment({ status: 'pending', progress: 0 }), + ]; + expect(buildWirePayload(list, 'path-1')).toBeUndefined(); + }); + + it('returns {path, files} for uploaded chips only', () => { + const list = [ + makeAttachment({ id: 'a', remoteFilename: 'org/uuid/a.pdf' }), + makeAttachment({ id: 'b', status: 'uploading', progress: 0.2 }), + makeAttachment({ id: 'c', status: 'error', terminal: false }), + ]; + const payload: AgentAttachmentWire | undefined = buildWirePayload(list, 'path-1'); + expect(payload).toEqual({ path: 'path-1', files: ['org/uuid/a.pdf'] }); + }); +}); + +describe('buildSubmissionPayload', () => { + it('returns undefined when no chips are uploaded', () => { + expect(buildSubmissionPayload([], 'path-1', 'uuid-1')).toBeUndefined(); + }); + + it('builds the S2 contract: wire + messageUuid + per-file descriptor with NO mime field', () => { + const list = [ + makeAttachment({ + id: 'a', + filename: 'a.pdf', + size: 1024, + remoteFilename: 'org/uuid/a.pdf', + }), + ]; + const payload: AgentAttachmentSubmissionPayload | undefined = buildSubmissionPayload( + list, + 'path-1', + 'uuid-1' + ); + expect(payload).toEqual({ + wire: { path: 'path-1', files: ['org/uuid/a.pdf'] }, + messageUuid: 'uuid-1', + files: [ + { + remoteName: 'org/uuid/a.pdf', + originalName: 'a.pdf', + size: 1024, + }, + ], + }); + // No `mime` field on the descriptor — every consumer derives MIME from + // the validated `remoteName` extension. + expect(payload).toBeDefined(); + const firstFile = payload?.files[0]; + expect(firstFile).toBeDefined(); + expect(Object.keys(firstFile ?? {})).toEqual(['remoteName', 'originalName', 'size']); + }); + + it('omits in-flight and failed chips from the payload', () => { + const list = [ + makeAttachment({ id: 'a', remoteFilename: 'org/uuid/a.pdf' }), + makeAttachment({ id: 'b', status: 'uploading', progress: 0.5 }), + makeAttachment({ id: 'c', status: 'error', terminal: false }), + ]; + const payload = buildSubmissionPayload(list, 'path-1', 'uuid-1'); + expect(payload?.files).toHaveLength(1); + expect(payload?.files[0]?.remoteName).toBe('org/uuid/a.pdf'); + }); +}); + +describe('isAnyAttachmentUploading / hasAnyFailedAttachment (send-admission signals)', () => { + it('reports uploading=true while a chip is in flight', () => { + const list = [makeAttachment({ status: 'uploading', progress: 0.5 })]; + expect(isAnyAttachmentUploading(list)).toBe(true); + expect(hasAnyFailedAttachment(list)).toBe(false); + }); + + it('reports uploading=true for a pending chip (not yet started)', () => { + const list = [makeAttachment({ status: 'pending', progress: 0 })]; + expect(isAnyAttachmentUploading(list)).toBe(true); + }); + + it('reports hasFailed=true after a chip errors (retryable OR terminal)', () => { + const retryable = [makeAttachment({ status: 'error', terminal: false, progress: null })]; + const terminal = [makeAttachment({ status: 'error', terminal: true, progress: null })]; + expect(hasAnyFailedAttachment(retryable)).toBe(true); + expect(hasAnyFailedAttachment(terminal)).toBe(true); + }); + + it('reports both signals false once every chip is uploaded', () => { + const list = [makeAttachment({ status: 'uploaded', progress: 1 })]; + expect(isAnyAttachmentUploading(list)).toBe(false); + expect(hasAnyFailedAttachment(list)).toBe(false); + }); + + it('reports both signals false on an empty list (send can proceed)', () => { + expect(isAnyAttachmentUploading([])).toBe(false); + expect(hasAnyFailedAttachment([])).toBe(false); + }); +}); + +describe('feature-state matrix — send-admission behavior', () => { + it('blocks send while ANY chip is uploading OR failed', () => { + const cases: { list: AgentAttachment[]; shouldBlock: boolean }[] = [ + { list: [], shouldBlock: false }, + { list: [makeAttachment({ status: 'uploaded', progress: 1 })], shouldBlock: false }, + { list: [makeAttachment({ status: 'uploading', progress: 0.5 })], shouldBlock: true }, + { list: [makeAttachment({ status: 'pending', progress: 0 })], shouldBlock: true }, + { + list: [makeAttachment({ status: 'error', terminal: false, progress: null })], + shouldBlock: true, + }, + { + list: [makeAttachment({ status: 'error', terminal: true, progress: null })], + shouldBlock: true, + }, + // Mixed: one uploaded + one uploading still blocks. + { + list: [ + makeAttachment({ id: 'a', status: 'uploaded', progress: 1 }), + makeAttachment({ id: 'b', status: 'uploading', progress: 0.3 }), + ], + shouldBlock: true, + }, + // Mixed: one uploaded + one terminal still blocks until the terminal chip is removed. + { + list: [ + makeAttachment({ id: 'a', status: 'uploaded', progress: 1 }), + makeAttachment({ id: 'b', status: 'error', terminal: true, progress: null }), + ], + shouldBlock: true, + }, + ]; + for (const { list, shouldBlock } of cases) { + const blocked = isAnyAttachmentUploading(list) || hasAnyFailedAttachment(list); + expect(blocked, `unexpected admission for ${JSON.stringify(list.map(a => a.status))}`).toBe( + shouldBlock + ); + } + }); +}); + +describe('size limits (5 MB / 5 files) — constant parity', () => { + it('exposes the 5 MB constant for both web and mobile parity', () => { + expect(AGENT_ATTACHMENT_MAX_BYTES).toBe(5 * 1024 * 1024); + }); +}); diff --git a/apps/mobile/src/lib/agent-attachments/use-agent-attachment-upload.ts b/apps/mobile/src/lib/agent-attachments/use-agent-attachment-upload.ts index 955e59e8b5..84514085b3 100644 --- a/apps/mobile/src/lib/agent-attachments/use-agent-attachment-upload.ts +++ b/apps/mobile/src/lib/agent-attachments/use-agent-attachment-upload.ts @@ -2,32 +2,32 @@ import * as Crypto from 'expo-crypto'; import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { toast } from 'sonner-native'; -import { trpcClient } from '@/lib/trpc'; +import { AGENT_ATTACHMENT_MAX_FILES } from '@/lib/agent-attachments/constants'; import { - AGENT_ATTACHMENT_MAX_BYTES, - AGENT_ATTACHMENT_MAX_FILES, - AGENT_ATTACHMENT_MIME_BY_EXTENSION, - type AgentAttachmentExtension, -} from '@/lib/agent-attachments/constants'; -import { canAddAttachments, classifyAttachment } from '@/lib/agent-attachments/validate'; - -export type AgentAttachmentKind = 'image' | 'document'; -export type AgentAttachmentStatus = 'pending' | 'uploading' | 'uploaded' | 'error'; - -type AllowedContentType = (typeof AGENT_ATTACHMENT_MIME_BY_EXTENSION)[AgentAttachmentExtension]; + canAddAttachments, + classifyAttachment, + describeClassificationFailure, + mimeForExtension, +} from '@/lib/agent-attachments/validate'; +import { + type AgentAttachment, + type AgentAttachmentSubmissionPayload, + type AgentAttachmentWire, + buildSubmissionPayload, + buildWirePayload, + classifyUploadFailure, + hasAnyFailedAttachment, + isAnyAttachmentUploading, +} from '@/lib/agent-attachments/agent-attachment-types'; +import { + describeTerminalReason, + measureLocalSize, + normalizeFilename, + uploadOne, +} from '@/lib/agent-attachments/upload-task'; -export type AgentAttachment = { - id: string; - filename: string; - /** Basename of the server-side R2 key; set once the upload succeeds. */ - remoteFilename?: string; - kind: AgentAttachmentKind; - mimeType: AllowedContentType; - size: number; - localUri: string; - status: AgentAttachmentStatus; - error?: string; -}; +// Re-export only the types consumers import from this module. +export type { AgentAttachment, AgentAttachmentSubmissionPayload, AgentAttachmentWire }; export type AgentAttachmentCandidate = { name: string; @@ -36,78 +36,31 @@ export type AgentAttachmentCandidate = { size?: number; }; -export type AgentAttachmentWire = { - path: string; - files: string[]; -}; - type UseAgentAttachmentUploadOptions = { organizationId?: string; }; type UseAgentAttachmentUploadReturn = { attachments: AgentAttachment[]; - addCandidates: (candidates: AgentAttachmentCandidate[]) => void; + addCandidates: (candidates: AgentAttachmentCandidate[]) => Promise; removeAttachment: (id: string) => void; retryAttachment: (id: string) => void; reset: () => void; isUploading: boolean; hasFailedAttachments: boolean; + /** Wire payload for the existing `chat-composer` send path. */ toWirePayload: () => AgentAttachmentWire | undefined; + /** The S2 submission payload. `undefined` when there are no uploads. */ + toSubmissionPayload: () => AgentAttachmentSubmissionPayload | undefined; }; -function ensureExtension(name: string, fallback: string): string { - const dot = name.lastIndexOf('.'); - if (dot > 0 && dot < name.length - 1) { - return name; - } - return `${name}.${fallback}`; -} - -async function uploadOne(args: { - organizationId?: string; - attachmentId: string; - path: string; - contentType: AllowedContentType; - localUri: string; -}): Promise<{ key: string }> { - const { organizationId, attachmentId, path, contentType, localUri } = args; - // expo-file-system's `File` is not a `Blob`; materialize a real `Blob` from - // the file:// URI so the PUT body matches the signed Content-Length. - const localFileResponse = await fetch(localUri); - const blob = await localFileResponse.blob(); - if (blob.size > AGENT_ATTACHMENT_MAX_BYTES) { - throw new Error('File is larger than 5 MB'); - } - const baseInput = { - messageUuid: path, - attachmentId, - contentType, - contentLength: blob.size, - }; - const result = organizationId - ? await trpcClient.organizations.cloudAgentNext.getAttachmentUploadUrl.mutate({ - ...baseInput, - organizationId, - }) - : await trpcClient.cloudAgentNext.getAttachmentUploadUrl.mutate(baseInput); - const response = await fetch(result.signedUrl, { - method: 'PUT', - headers: { 'Content-Type': contentType }, - body: blob, - }); - if (!response.ok) { - throw new Error(`Upload failed with status ${response.status}`); - } - return { key: result.key }; -} - export function useAgentAttachmentUpload( options: UseAgentAttachmentUploadOptions = {} ): UseAgentAttachmentUploadReturn { const { organizationId } = options; const [attachments, setAttachments] = useState([]); const pathRef = useRef(Crypto.randomUUID()); + const messageUuidRef = useRef(Crypto.randomUUID()); const isMountedRef = useRef(true); useEffect(() => { @@ -117,44 +70,62 @@ export function useAgentAttachmentUpload( }; }, []); + const updateAttachment = useCallback((id: string, patch: Partial) => { + if (!isMountedRef.current) { + return; + } + setAttachments(current => current.map(item => (item.id === id ? { ...item, ...patch } : item))); + }, []); + const startUpload = useCallback( (attachment: AgentAttachment, path: string) => { - const update = (patch: Partial) => { - if (!isMountedRef.current) { - return; - } - setAttachments(current => - current.map(item => (item.id === attachment.id ? { ...item, ...patch } : item)) - ); - }; - const run = async () => { - update({ status: 'uploading', error: undefined }); + updateAttachment(attachment.id, { + status: 'uploading', + error: undefined, + terminal: undefined, + progress: 0, + }); try { const { key } = await uploadOne({ organizationId, attachmentId: attachment.id, path, + extension: attachment.extension, contentType: attachment.mimeType, + contentLength: attachment.size, localUri: attachment.localUri, + onProgress: progress => { + updateAttachment(attachment.id, { progress }); + }, + }); + updateAttachment(attachment.id, { + status: 'uploaded', + remoteFilename: key.split('/').at(-1), + progress: 1, }); - // The wire payload must reference the object the server actually - // stored, so take the filename from the returned R2 key. - update({ status: 'uploaded', remoteFilename: key.split('/').at(-1) }); } catch (error) { - const message = error instanceof Error ? error.message : 'Upload failed'; - toast.error(`Failed to upload file: ${message}`); - update({ status: 'error', error: message }); + const { retryable, reason } = classifyUploadFailure(error); + updateAttachment(attachment.id, { + status: 'error', + error: retryable ? reason : describeTerminalReason(reason), + terminal: !retryable, + progress: null, + }); + // Single toast per failed chip. Terminal surfaces its own chip + // copy so the toast only needs to echo the same intent. + toast.error( + retryable ? `Failed to upload file: ${reason}` : describeTerminalReason(reason) + ); } }; - void run(); }, - [organizationId] + [organizationId, updateAttachment] ); const addCandidates = useCallback( - (candidates: AgentAttachmentCandidate[]) => { + async (candidates: AgentAttachmentCandidate[]) => { if (candidates.length === 0) { return; } @@ -170,42 +141,45 @@ export function useAgentAttachmentUpload( ); } + // We pre-classify synchronously; the *measured* size comes from + // `getInfoAsync`. Candidates that the picker reports as zero-size + // (common on iOS) are re-measured before the size/empty rule fires. + const measured = await Promise.all( + accepted.map(async candidate => { + const measuredSize = await measureLocalSize(candidate.uri); + const size = measuredSize ?? candidate.size ?? 0; + return { candidate, size }; + }) + ); + const additions: AgentAttachment[] = []; - for (const candidate of accepted) { - const classified = classifyAttachment({ - name: candidate.name, - mimeType: candidate.mimeType, - size: candidate.size, - }); + for (const { candidate, size } of measured) { + const classified = classifyAttachment({ name: candidate.name, size }); if (!classified.ok) { - toast.error( - classified.reason === 'too-large' - ? `File too large: ${candidate.name}. Max size is 5 MB.` - : `File type not supported: ${candidate.name}. Attach PNG, JPEG, WebP, GIF, PDF, TXT, MD, or CSV files.` - ); + toast.error(describeClassificationFailure(classified.reason)); } else { const ext = classified.extension; + const filename = normalizeFilename(candidate.name, ext); additions.push({ id: Crypto.randomUUID(), - filename: ensureExtension(candidate.name, ext), + filename, kind: classified.kind, - // Always derive the content type from the extension: OS pickers - // report generic types (e.g. application/octet-stream) that the - // backend's allowed-type enum rejects. - mimeType: AGENT_ATTACHMENT_MIME_BY_EXTENSION[ext], - size: candidate.size ?? 0, + extension: ext, + mimeType: mimeForExtension(ext), + size: classified.size, localUri: candidate.uri, status: 'pending', + progress: 0, }); } } if (additions.length === 0) { return; } + setAttachments(current => [...current, ...additions]); for (const addition of additions) { startUpload(addition, pathRef.current); } - setAttachments(current => [...current, ...additions]); }, [attachments.length, startUpload] ); @@ -217,7 +191,9 @@ export function useAgentAttachmentUpload( const retryAttachment = useCallback( (id: string) => { const attachment = attachments.find(item => item.id === id); - if (!attachment) { + if (!attachment || attachment.terminal) { + // Terminal chips have no retry affordance; bail so a stray + // tap cannot re-upload a server-rejected file. return; } startUpload(attachment, pathRef.current); @@ -228,23 +204,22 @@ export function useAgentAttachmentUpload( const reset = useCallback(() => { setAttachments([]); pathRef.current = Crypto.randomUUID(); + messageUuidRef.current = Crypto.randomUUID(); }, []); - const toWirePayload = useCallback((): AgentAttachmentWire | undefined => { - const files = attachments - .filter(item => item.status === 'uploaded') - .map(item => item.remoteFilename) - .filter((filename): filename is string => filename !== undefined); - if (files.length === 0) { - return undefined; - } - return { path: pathRef.current, files }; - }, [attachments]); + const toWirePayload = useCallback( + (): AgentAttachmentWire | undefined => buildWirePayload(attachments, pathRef.current), + [attachments] + ); - const isUploading = attachments.some( - item => item.status === 'pending' || item.status === 'uploading' + const toSubmissionPayload = useCallback( + (): AgentAttachmentSubmissionPayload | undefined => + buildSubmissionPayload(attachments, pathRef.current, messageUuidRef.current), + [attachments] ); - const failedAttachments = attachments.some(item => item.status === 'error'); + + const isUploading = isAnyAttachmentUploading(attachments); + const hasFailedAttachments = hasAnyFailedAttachment(attachments); return useMemo( () => ({ @@ -254,8 +229,9 @@ export function useAgentAttachmentUpload( retryAttachment, reset, isUploading, - hasFailedAttachments: failedAttachments, + hasFailedAttachments, toWirePayload, + toSubmissionPayload, }), [ attachments, @@ -264,8 +240,9 @@ export function useAgentAttachmentUpload( retryAttachment, reset, isUploading, - failedAttachments, + hasFailedAttachments, toWirePayload, + toSubmissionPayload, ] ); } diff --git a/apps/mobile/src/lib/agent-attachments/validate.test.ts b/apps/mobile/src/lib/agent-attachments/validate.test.ts index 548d79f3f3..b3022cb4a6 100644 --- a/apps/mobile/src/lib/agent-attachments/validate.test.ts +++ b/apps/mobile/src/lib/agent-attachments/validate.test.ts @@ -1,34 +1,130 @@ import { describe, expect, it } from 'vitest'; -import { canAddAttachments, classifyAttachment } from './validate'; +import { + AGENT_ATTACHMENT_DENIED_EXTENSIONS, + AGENT_ATTACHMENT_MIME_BY_EXTENSION, + type AgentAttachmentExtension, +} from './constants'; +import { + canAddAttachments, + classifyAttachment, + describeClassificationFailure, + mimeForExtension, + normalizeAttachmentExtension, +} from './validate'; + +describe('normalizeAttachmentExtension', () => { + it('lowercases a known extension', () => { + expect(normalizeAttachmentExtension('NOTES.PDF')).toBe('pdf'); + }); + + it('returns the fallback when the filename has no extension', () => { + expect(normalizeAttachmentExtension('README')).toBe('bin'); + }); + + it('returns the fallback for an extension that violates the regex', () => { + expect(normalizeAttachmentExtension('evil.tar!@#')).toBe('bin'); + expect(normalizeAttachmentExtension(`archive.${'a'.repeat(32)}`)).toBe('bin'); + }); + + it('returns the fallback when the filename ends in a dot', () => { + expect(normalizeAttachmentExtension('weird.')).toBe('bin'); + }); + + it('accepts an extension that is not in the canonical table but matches the regex', () => { + // `mov` is not in AGENT_ATTACHMENT_MIME_BY_EXTENSION; the classifier + // must still produce a normalized extension and let the MIME fallback + // take over. We rely on the fallback only — the caller is expected to + // also fall through to `bin` if the MIME lookup would fail. + const ext = normalizeAttachmentExtension('clip.mov'); + expect(ext).toBe('mov'); + }); +}); describe('classifyAttachment', () => { - it('accepts a png by extension', () => { - expect(classifyAttachment({ name: 'a.PNG', mimeType: 'image/png', size: 10 })).toEqual({ + it('accepts a PNG and reports the image kind', () => { + expect(classifyAttachment({ name: 'a.PNG', size: 10 })).toEqual({ ok: true, kind: 'image', extension: 'png', + size: 10, }); }); - it('accepts a markdown file by extension', () => { - expect(classifyAttachment({ name: 'notes.md', mimeType: 'text/markdown', size: 10 })).toEqual({ + it('accepts a markdown file and reports the document kind', () => { + expect(classifyAttachment({ name: 'notes.md', size: 10 })).toEqual({ ok: true, kind: 'document', extension: 'md', + size: 10, + }); + }); + + it('accepts an extension outside the image/document allow-list as a generic binary', () => { + expect(classifyAttachment({ name: 'archive.zip', size: 10 })).toEqual({ + ok: true, + kind: 'document', + extension: 'zip', + size: 10, + }); + }); + + it('rejects a zero-byte file with reason=empty', () => { + expect(classifyAttachment({ name: 'empty.pdf', size: 0 })).toEqual({ + ok: false, + reason: 'empty', + }); + }); + + it('rejects a negative size (defensive) with reason=empty', () => { + expect(classifyAttachment({ name: 'neg.pdf', size: -1 })).toEqual({ + ok: false, + reason: 'empty', + }); + }); + + it('rejects a file whose extension is on the deny list', () => { + expect(classifyAttachment({ name: 'malware.exe', size: 10 })).toEqual({ + ok: false, + reason: 'denied', + }); + }); + + it('rejects every entry in AGENT_ATTACHMENT_DENIED_EXTENSIONS', () => { + for (const ext of AGENT_ATTACHMENT_DENIED_EXTENSIONS) { + const result = classifyAttachment({ name: `virus.${ext}`, size: 10 }); + expect(result, `expected ${ext} to be denied`).toEqual({ ok: false, reason: 'denied' }); + } + }); + + it('rejects a file over the 5 MB cap', () => { + expect(classifyAttachment({ name: 'big.pdf', size: 6 * 1024 * 1024 })).toEqual({ + ok: false, + reason: 'too-large', }); }); - it('rejects unsupported extension', () => { - expect( - classifyAttachment({ name: 'a.exe', mimeType: 'application/octet-stream', size: 10 }).ok - ).toBe(false); + it('accepts a file exactly at the 5 MB cap', () => { + expect(classifyAttachment({ name: 'edge.pdf', size: 5 * 1024 * 1024 })).toEqual({ + ok: true, + kind: 'document', + extension: 'pdf', + size: 5 * 1024 * 1024, + }); }); - it('rejects a file over the size cap', () => { - expect( - classifyAttachment({ name: 'a.pdf', mimeType: 'application/pdf', size: 6 * 1024 * 1024 }).ok - ).toBe(false); + it('checks the deny list before the size cap (a 5 MB .exe is denied, not too-large)', () => { + expect(classifyAttachment({ name: 'big.exe', size: 5 * 1024 * 1024 + 1 })).toEqual({ + ok: false, + reason: 'denied', + }); + }); + + it('checks empty before the deny list (a 0-byte .exe is empty, not denied)', () => { + expect(classifyAttachment({ name: 'zero.exe', size: 0 })).toEqual({ + ok: false, + reason: 'empty', + }); }); }); @@ -45,3 +141,55 @@ describe('canAddAttachments', () => { expect(canAddAttachments(5, 1)).toEqual({ ok: false, acceptedCount: 0 }); }); }); + +describe('describeClassificationFailure', () => { + it('returns the locked copy for each reason', () => { + expect(describeClassificationFailure('denied')).toMatch(/can't be attached/i); + expect(describeClassificationFailure('empty')).toMatch(/empty/i); + expect(describeClassificationFailure('too-large')).toMatch(/5 MB/); + }); +}); + +describe('mimeForExtension (cross-surface parity)', () => { + it('returns the canonical MIME for every documented extension', () => { + expect(mimeForExtension('png')).toBe('image/png'); + expect(mimeForExtension('jpg')).toBe('image/jpeg'); + expect(mimeForExtension('jpeg')).toBe('image/jpeg'); + expect(mimeForExtension('webp')).toBe('image/webp'); + expect(mimeForExtension('gif')).toBe('image/gif'); + expect(mimeForExtension('pdf')).toBe('application/pdf'); + expect(mimeForExtension('txt')).toBe('text/plain'); + expect(mimeForExtension('md')).toBe('text/plain'); + expect(mimeForExtension('ts')).toBe('text/plain'); + expect(mimeForExtension('bin')).toBe('application/octet-stream'); + }); + + it('falls back to application/octet-stream for an extension not in the canonical table', () => { + // The picker must NEVER trust the picker's reported MIME; when the + // extension is not in the table the fallback is `application/octet-stream`. + // The compose-time `normalizeAttachmentExtension` keeps the original + // extension so the server can still distinguish `mov` vs `mp4`; the + // caller's `mimeForExtension` lookup is the safety net. + expect(mimeForExtension('mov' as AgentAttachmentExtension)).toBe('application/octet-stream'); + }); +}); + +describe('AGENT_ATTACHMENT_MIME_BY_EXTENSION (server parity)', () => { + it('resolves every key to a defined MIME string', () => { + for (const [ext, mime] of Object.entries(AGENT_ATTACHMENT_MIME_BY_EXTENSION)) { + expect(typeof mime).toBe('string'); + expect(mime.length).toBeGreaterThan(0); + expect(ext.length).toBeGreaterThan(0); + } + }); + + it('contains the fallback `bin` key', () => { + expect(AGENT_ATTACHMENT_MIME_BY_EXTENSION.bin).toBe('application/octet-stream'); + }); + + it('does not contain a denied extension as a key', () => { + for (const ext of AGENT_ATTACHMENT_DENIED_EXTENSIONS) { + expect(Object.hasOwn(AGENT_ATTACHMENT_MIME_BY_EXTENSION, ext)).toBe(false); + } + }); +}); diff --git a/apps/mobile/src/lib/agent-attachments/validate.ts b/apps/mobile/src/lib/agent-attachments/validate.ts index ce09c9d77f..28191fc686 100644 --- a/apps/mobile/src/lib/agent-attachments/validate.ts +++ b/apps/mobile/src/lib/agent-attachments/validate.ts @@ -1,30 +1,90 @@ import { - AGENT_ATTACHMENT_EXTENSIONS, + AGENT_ATTACHMENT_DENIED_EXTENSIONS, + AGENT_ATTACHMENT_EXTENSION_REGEX, + AGENT_ATTACHMENT_FALLBACK_EXTENSION, AGENT_ATTACHMENT_MAX_BYTES, AGENT_ATTACHMENT_MAX_FILES, + AGENT_ATTACHMENT_MIME_BY_EXTENSION, type AgentAttachmentExtension, + type AgentAttachmentMime, } from './constants'; -const IMAGE_EXTENSIONS = new Set(['png', 'jpg', 'jpeg', 'webp', 'gif']); +const IMAGE_EXTENSIONS = new Set(['png', 'jpg', 'jpeg', 'webp', 'gif']); + +/** + * Normalize a candidate's filename extension. + * + * - Lower-cased, dot-stripped. + * - Falls back to {@link AGENT_ATTACHMENT_FALLBACK_EXTENSION} when the + * filename has no extension or the extension does not match + * {@link AGENT_ATTACHMENT_EXTENSION_REGEX}. + * - Never throws; the result is always a known key in + * {@link AGENT_ATTACHMENT_MIME_BY_EXTENSION}. + */ +export function normalizeAttachmentExtension(name: string): AgentAttachmentExtension { + const dot = name.lastIndexOf('.'); + if (dot === -1 || dot === name.length - 1) { + return AGENT_ATTACHMENT_FALLBACK_EXTENSION; + } + const raw = name.slice(dot + 1).toLowerCase(); + return AGENT_ATTACHMENT_EXTENSION_REGEX.test(raw) + ? (raw as AgentAttachmentExtension) + : AGENT_ATTACHMENT_FALLBACK_EXTENSION; +} + +/** + * Resolve the canonical MIME for a normalized extension. The picker MIME + * is intentionally NOT consulted: OS pickers report generic types like + * `application/octet-stream` for anything the platform doesn't recognize, + * and the cloud-agent storage layer rejects anything outside the + * allow-list. The extension is the single source of truth. + * + * When the extension is not in {@link AGENT_ATTACHMENT_MIME_BY_EXTENSION} + * the lookup falls back to `application/octet-stream` so the picker + * contract is total and never leaks `undefined` into the upload path. + */ +export function mimeForExtension(extension: string): AgentAttachmentMime { + // Index via a string map so missing keys are `undefined` at the type + // level (the const table alone would make `??` look unnecessary). + const mimeByExtension: Readonly> = + AGENT_ATTACHMENT_MIME_BY_EXTENSION; + return mimeByExtension[extension] ?? 'application/octet-stream'; +} type ClassifiedAttachment = - | { ok: true; kind: 'image' | 'document'; extension: AgentAttachmentExtension } - | { ok: false; reason: 'unsupported' | 'too-large' }; + | { ok: true; kind: 'image' | 'document'; extension: AgentAttachmentExtension; size: number } + | { + ok: false; + reason: 'denied' | 'empty' | 'too-large'; + }; -type Candidate = { name: string; mimeType?: string; size?: number }; +type AttachmentCandidate = { name: string; size: number }; -export function classifyAttachment(candidate: Candidate): ClassifiedAttachment { - const ext = candidate.name.split('.').pop()?.toLowerCase(); - if (!ext || !(AGENT_ATTACHMENT_EXTENSIONS as readonly string[]).includes(ext)) { - return { ok: false, reason: 'unsupported' }; +/** + * Classify a candidate attachment for the cloud-agent composer. + * + * The picker `mimeType` is intentionally ignored: the extension is the + * single source of truth for both the deny list and the MIME we send to + * the server. `size` is the **measured** local byte size (see + * `getInfoAsync` in `use-agent-attachment-upload`); the picker-reported + * size is unreliable on iOS and must not be trusted. + */ +export function classifyAttachment(candidate: AttachmentCandidate): ClassifiedAttachment { + if (candidate.size <= 0) { + return { ok: false, reason: 'empty' }; } - if (typeof candidate.size === 'number' && candidate.size > AGENT_ATTACHMENT_MAX_BYTES) { + const extension = normalizeAttachmentExtension(candidate.name); + if (AGENT_ATTACHMENT_DENIED_EXTENSIONS.has(extension)) { + return { ok: false, reason: 'denied' }; + } + if (candidate.size > AGENT_ATTACHMENT_MAX_BYTES) { return { ok: false, reason: 'too-large' }; } return { ok: true, - kind: IMAGE_EXTENSIONS.has(ext) ? 'image' : 'document', - extension: ext as AgentAttachmentExtension, + kind: IMAGE_EXTENSIONS.has(extension) ? 'image' : 'document', + extension, + size: candidate.size, }; } @@ -41,3 +101,17 @@ export function canAddAttachments( } return { ok: true, acceptedCount: remaining, truncated: true }; } + +const CLASSIFICATION_FAILURE_MESSAGES = { + denied: "Executable files can't be attached", + empty: 'File is empty', + 'too-large': 'Files must be 5 MB or smaller', +} as const satisfies Record<'denied' | 'empty' | 'too-large', string>; + +/** + * Human-readable copy for a single classification outcome. Centralized so + * the picker, the upload hook, and the chip surface use the same strings. + */ +export function describeClassificationFailure(reason: 'denied' | 'empty' | 'too-large'): string { + return CLASSIFICATION_FAILURE_MESSAGES[reason]; +} 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 8fc568bbb1..cc8554412d 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 @@ -15,6 +15,7 @@ import { } from './user-web-connection'; import type { KiloSessionId, SessionSnapshot, SessionSnapshotPageOutcome } from './types'; import { kiloId, makeSnapshot, stubTextPart, stubUserMessage } from './test-helpers'; +import type { RemoteAttachmentPart } from './transport'; const KILO_SESSION_ID = kiloId('kilo-ses-1'); const NEW_KILO_SESSION_ID = 'ses_12345678901234567890123456'; @@ -156,6 +157,7 @@ function createTransportWithSinks(opts?: { onRemoteModelStateChange?: (state: RemoteModelState) => void; onRemoteCommandStateChange?: (state: RemoteCommandState) => void; onCapabilityChange?: () => void; + onCapabilitiesChange?: (capabilities: { attachments?: boolean } | undefined) => void; }) { const userWebConnection = opts?.connection ?? createConnection(); const chatEvents: ChatEvent[] = []; @@ -169,6 +171,7 @@ function createTransportWithSinks(opts?: { onRemoteModelStateChange: opts?.onRemoteModelStateChange, onRemoteCommandStateChange: opts?.onRemoteCommandStateChange, onCapabilityChange: opts?.onCapabilityChange, + onCapabilitiesChange: opts?.onCapabilitiesChange, })({ onChatEvent: event => chatEvents.push(event), onServiceEvent: event => serviceEvents.push(event), @@ -204,6 +207,22 @@ function emitMessageUpdated(connection: FakeUserWebConnection, sessionId = KILO_ }); } +function emitHeartbeat( + connection: FakeUserWebConnection, + sessions: Array<{ + id: string; + status: string; + title: string; + capabilities?: { attachments?: boolean }; + }>, + connectionId = 'owner' +): void { + connection.emitSystem({ + event: 'sessions.heartbeat', + data: { connectionId, sessions }, + }); +} + describe('CliLiveTransport unified user web connection', () => { it('discovers and publishes a v1 catalog for the current owner', async () => { const connection = createConnection(); @@ -3123,3 +3142,212 @@ describe('CliLiveTransport page-seam + reconnect', () => { transport.destroy(); }); }); + +describe('CliLiveTransport remote attachment capabilities', () => { + it('publishes heartbeat attachments: true via onCapabilitiesChange', async () => { + const connection = createConnection(); + jest + .mocked(connection.sendCommand) + .mockImplementation((_sessionId, command) => + Promise.resolve(command === 'list_models' ? WIRE_CATALOG : { ok: true }) + ); + const capabilities: ({ attachments?: boolean } | undefined)[] = []; + const { transport, userWebConnection } = createTransportWithSinks({ + connection, + onCapabilitiesChange: capability => capabilities.push(capability), + }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + + emitHeartbeat(userWebConnection, [ + { + id: KILO_SESSION_ID, + status: 'active', + title: 'Tracked', + capabilities: { attachments: true }, + }, + ]); + + expect(capabilities.at(-1)).toEqual({ attachments: true }); + + emitHeartbeat(userWebConnection, [ + { + id: KILO_SESSION_ID, + status: 'active', + title: 'Tracked', + capabilities: { attachments: false }, + }, + ]); + + expect(capabilities.at(-1)).toEqual({ attachments: false }); + transport.destroy(); + }); + + it('publishes sessions.list capability changes via onCapabilitiesChange', async () => { + const connection = createConnection(); + jest + .mocked(connection.sendCommand) + .mockImplementation((_sessionId, command) => + Promise.resolve(command === 'list_models' ? WIRE_CATALOG : { ok: true }) + ); + const capabilities: ({ attachments?: boolean } | undefined)[] = []; + const { transport, userWebConnection } = createTransportWithSinks({ + connection, + onCapabilitiesChange: capability => capabilities.push(capability), + }); + + transport.connect(); + userWebConnection.emitSystem({ + event: 'sessions.list', + data: { + sessions: [ + { + id: KILO_SESSION_ID, + status: 'active', + title: 'Tracked', + connectionId: 'owner', + capabilities: { attachments: true }, + }, + ], + }, + }); + await Promise.resolve(); + await Promise.resolve(); + + expect(capabilities.at(-1)).toEqual({ attachments: true }); + + userWebConnection.emitSystem({ + event: 'sessions.list', + data: { + sessions: [ + { + id: KILO_SESSION_ID, + status: 'active', + title: 'Tracked', + connectionId: 'owner', + capabilities: { attachments: false }, + }, + ], + }, + }); + + expect(capabilities.at(-1)).toEqual({ attachments: false }); + transport.destroy(); + }); + + it('publishes undefined on reconnect then re-advertises via heartbeat', async () => { + const connection = createConnection(); + jest + .mocked(connection.sendCommand) + .mockImplementation((_sessionId, command) => + Promise.resolve(command === 'list_models' ? WIRE_CATALOG : { ok: true }) + ); + const capabilities: ({ attachments?: boolean } | undefined)[] = []; + const { transport, userWebConnection } = createTransportWithSinks({ + connection, + onCapabilitiesChange: capability => capabilities.push(capability), + }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + + emitHeartbeat(userWebConnection, [ + { + id: KILO_SESSION_ID, + status: 'active', + title: 'Tracked', + capabilities: { attachments: true }, + }, + ]); + expect(capabilities.at(-1)).toEqual({ attachments: true }); + + userWebConnection.emitReconnect(); + await Promise.resolve(); + await Promise.resolve(); + + // The gate must close immediately on reconnect, before the next + // heartbeat/sessions.list has a chance to re-advertise. + expect(capabilities.at(-1)).toBeUndefined(); + + emitHeartbeat(userWebConnection, [ + { + id: KILO_SESSION_ID, + status: 'active', + title: 'Tracked', + capabilities: { attachments: true }, + }, + ]); + expect(capabilities.at(-1)).toEqual({ attachments: true }); + + transport.destroy(); + }); +}); + +describe('CliLiveTransport send_message parts', () => { + it('appends file parts after the text part in send_message', async () => { + const connection = createConnection(); + jest + .mocked(connection.sendCommand) + .mockImplementation((_sessionId, command) => + Promise.resolve(command === 'list_models' ? WIRE_CATALOG : { ok: true }) + ); + const { userWebConnection, transport } = createTransportWithSinks({ connection }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + await Promise.resolve(); + jest.mocked(userWebConnection.sendCommand).mockClear(); + + const attachmentParts: RemoteAttachmentPart[] = [ + { + type: 'file', + mime: 'text/plain', + filename: 'msg-uuid.txt', + url: 'https://r2.example.com/msg-uuid.txt', + }, + { + type: 'file', + mime: 'image/png', + filename: 'msg-uuid.png', + url: 'https://r2.example.com/msg-uuid.png', + }, + ]; + + await transport.send?.({ + payload: { type: 'prompt', prompt: 'look at these files' }, + attachmentParts, + }); + + expect(userWebConnection.sendCommand).toHaveBeenCalledWith( + KILO_SESSION_ID, + 'send_message', + { + sessionID: KILO_SESSION_ID, + parts: [ + { type: 'text', text: 'look at these files' }, + { + type: 'file', + mime: 'text/plain', + filename: 'msg-uuid.txt', + url: 'https://r2.example.com/msg-uuid.txt', + }, + { + type: 'file', + mime: 'image/png', + filename: 'msg-uuid.png', + url: 'https://r2.example.com/msg-uuid.png', + }, + ], + }, + '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 09c0100982..c22cef94a6 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 @@ -52,6 +52,15 @@ type CliLiveTransportConfig = { onRemoteModelStateChange?: (state: RemoteModelState) => void; onRemoteCommandStateChange?: (state: RemoteCommandState) => void; onCapabilityChange?: () => void; + /** + * Fired whenever the per-session capabilities advertised by the owning + * CLI in `sessions.heartbeat` / `sessions.list` change (upgrade, downgrade, + * reconnect, or absent). The payload is the latest capabilities — `undefined` + * means the CLI has not reported any (older CLIs, mid-reconnect, or a CLI + * whose session list dropped this session). The session manager uses this + * to recompute the `supportsAttachments` gate. + */ + onCapabilitiesChange?: (capabilities: { attachments?: boolean } | undefined) => void; }; // How long after a reconnect to re-fetch the snapshot a second time. Covers @@ -83,6 +92,21 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor let sessionStopped = false; let ownerConnectionId: string | null = null; let lastForwardedHeartbeatStatus: string | null = null; + /** + * Latest per-session capabilities observed in a `sessions.heartbeat` or + * `sessions.list` payload for this session. `undefined` means the CLI + * has not reported any (older CLIs, mid-reconnect, or a heartbeat/list + * payload that omitted this session). Compared structurally on every + * new observation; only emitted on a real change. + */ + let currentCapabilities: { attachments?: boolean } | undefined = undefined; + function publishCapabilities(next: { attachments?: boolean } | undefined): void { + const previousAttachments = currentCapabilities?.attachments; + const nextAttachments = next?.attachments; + if (previousAttachments === nextAttachments) return; + currentCapabilities = next; + config.onCapabilitiesChange?.(next); + } let catalogRequestGeneration = 0; let catalogRequestInFlight: { ownerConnectionId: string; generation: number } | null = null; // Command catalog discovery runs on its own generation so it stays @@ -176,6 +200,11 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor commands: snapshotCommands(), }); } + // A CLI handoff or a permanent drop invalidates whatever the prior + // owner reported — the new owner has to re-advertise before any + // capability re-enables. Empty currentCapabilities also drives the + // existing 'idle' reset on the consumer side. + publishCapabilities(undefined); config.onCapabilityChange?.(); if (nextOwnerConnectionId) { @@ -502,6 +531,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor setOwnerConnectionId(session.connectionId); sessionStopped = false; forwardHeartbeatStatus(session.status); + publishCapabilities(session.capabilities); return; } @@ -519,6 +549,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor setOwnerConnectionId(parsed.data.connectionId); sessionStopped = false; forwardHeartbeatStatus(session.status); + publishCapabilities(session.capabilities); return; } @@ -615,6 +646,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor catalogRequestInFlight = null; commandCatalogRequestGeneration += 1; commandCatalogRequestInFlight = null; + publishCapabilities(undefined); publishRemoteModelState({ ownerConnectionId: null, protocol: 'unknown', @@ -748,6 +780,10 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor handleSystemMessage(msg.event, msg.data); }); const offReconnect = config.userWebConnection.onReconnect(() => { + // Recompute the capability gate fail-closed immediately. The prior + // owner may have been attachment-capable, but after a reconnect we + // must wait for the next heartbeat / sessions.list to re-advertise. + publishCapabilities(undefined); replayCurrentSnapshot(false); // The snapshot store lags the live stream, and the CLI only forwards // events "from now" after a resubscribe — parts finalized while the @@ -848,9 +884,28 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor } const payload = input.payload; const remoteModel = getRemoteModelFields(input); + const parts: Array< + | { type: 'text'; text: string } + | { + type: 'file'; + mime: string; + filename: string; + url: string; + } + > = [{ type: 'text', text: payload.prompt }]; + if (input.attachmentParts && input.attachmentParts.length > 0) { + for (const part of input.attachmentParts) { + parts.push({ + type: 'file', + mime: part.mime, + filename: part.filename, + url: part.url, + }); + } + } return sendCommand('send_message', { sessionID: config.kiloSessionId, - parts: [{ type: 'text', text: payload.prompt }], + parts, ...(payload.mode ? { agent: payload.mode } : {}), ...(remoteModel.kind === 'none' ? {} diff --git a/apps/web/src/lib/cloud-agent-sdk/index.ts b/apps/web/src/lib/cloud-agent-sdk/index.ts index 6bf5dbc203..f33e608955 100644 --- a/apps/web/src/lib/cloud-agent-sdk/index.ts +++ b/apps/web/src/lib/cloud-agent-sdk/index.ts @@ -110,6 +110,7 @@ export type { CloudAgentApi, CloudAgentStreamTicket, CloudAgentStreamTicketResult, + RemoteAttachmentPart, TransportFactory, TransportSink, Transport, diff --git a/apps/web/src/lib/cloud-agent-sdk/schemas.test.ts b/apps/web/src/lib/cloud-agent-sdk/schemas.test.ts new file mode 100644 index 0000000000..c9a662d06f --- /dev/null +++ b/apps/web/src/lib/cloud-agent-sdk/schemas.test.ts @@ -0,0 +1,57 @@ +import { activeSessionSchema } from './schemas'; + +describe('activeSessionSchema capabilities', () => { + it('parses a session whose `capabilities` is absent', () => { + const parsed = activeSessionSchema.parse({ + id: 'ses_remote_a', + status: 'idle', + title: 'Test', + connectionId: 'conn-1', + }); + expect(parsed.capabilities).toBeUndefined(); + }); + + it('parses a session whose `capabilities.attachments` is false', () => { + const parsed = activeSessionSchema.parse({ + id: 'ses_remote_b', + status: 'idle', + title: 'Test', + connectionId: 'conn-1', + capabilities: { attachments: false }, + }); + expect(parsed.capabilities).toEqual({ attachments: false }); + }); + + it('parses a session whose `capabilities.attachments` is true', () => { + const parsed = activeSessionSchema.parse({ + id: 'ses_remote_c', + status: 'idle', + title: 'Test', + connectionId: 'conn-1', + capabilities: { attachments: true }, + }); + expect(parsed.capabilities).toEqual({ attachments: true }); + }); + + it('parses a session whose `capabilities` is an empty object (no attachments key)', () => { + const parsed = activeSessionSchema.parse({ + id: 'ses_remote_d', + status: 'idle', + title: 'Test', + connectionId: 'conn-1', + capabilities: {}, + }); + expect(parsed.capabilities).toEqual({}); + }); + + it('rejects a session whose `capabilities.attachments` is not a boolean', () => { + const result = activeSessionSchema.safeParse({ + id: 'ses_remote_e', + status: 'idle', + title: 'Test', + connectionId: 'conn-1', + capabilities: { attachments: 'yes' }, + }); + expect(result.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 0727ad0ece..dcc7761fe4 100644 --- a/apps/web/src/lib/cloud-agent-sdk/schemas.ts +++ b/apps/web/src/lib/cloud-agent-sdk/schemas.ts @@ -435,6 +435,13 @@ export type WebInboundMessage = z.infer; // Active CLI sessions // --------------------------------------------------------------------------- +export const activeSessionCapabilitiesSchema = z + .object({ + attachments: z.boolean().optional(), + }) + .optional(); +export type ActiveSessionCapabilities = z.infer; + export const activeSessionSchema = z .object({ id: z.string(), @@ -443,6 +450,18 @@ export const activeSessionSchema = z gitUrl: z.string().optional(), gitBranch: z.string().optional(), parentSessionId: z.string().optional(), + /** + * Per-session capabilities advertised by the owning CLI in its + * `sessions.heartbeat` / `sessions.list` payload. `attachments: true` + * gates the remote-CLI attachment path; absent / false means the CLI + * either predates the capability or has not yet reported it. + * + * Declared here (rather than relying on `.passthrough()`) so typed + * consumers can reason about the shape without an extra `as` cast. The + * surrounding `activeSessionSchema` is `.passthrough()` so older + * forward-compatible fields continue to be preserved verbatim. + */ + capabilities: activeSessionCapabilitiesSchema, }) .passthrough(); export type ActiveSessionData = z.infer; 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 84a27bca21..ca1a1ecd29 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 @@ -35,6 +35,7 @@ import type { } from './types'; import type { RemoteModelState } from './remote-model-catalog'; import type { RemoteCommandState } from './remote-command-catalog'; +import type { RemoteAttachmentPart } from './transport'; import type { NormalizedEvent } from './normalizer'; // --------------------------------------------------------------------------- @@ -116,6 +117,7 @@ const mockSessionCallbacks: { onRemoteModelStateChange?: (state: RemoteModelState) => void; onRemoteCommandStateChange?: (state: RemoteCommandState) => void; onTransportCapabilityChange?: () => void; + onTransportCapabilitiesChange?: (capabilities: { attachments?: boolean } | undefined) => void; onEvent?: (event: NormalizedEvent) => void; onMessageQueued?: (messageId: string) => void; onMessageCompleted?: (messageId: string) => void; @@ -150,6 +152,7 @@ jest.mock('./session', () => ({ onRemoteModelStateChange?: (state: RemoteModelState) => void; onRemoteCommandStateChange?: (state: RemoteCommandState) => void; onTransportCapabilityChange?: () => void; + onTransportCapabilitiesChange?: (capabilities: { attachments?: boolean } | undefined) => void; onEvent?: (event: NormalizedEvent) => void; onMessageQueued?: (messageId: string) => void; onMessageCompleted?: (messageId: string) => void; @@ -220,6 +223,8 @@ jest.mock('./session', () => ({ mockSessionCallbacks.onRemoteModelStateChange = sessionConfig.onRemoteModelStateChange; mockSessionCallbacks.onRemoteCommandStateChange = sessionConfig.onRemoteCommandStateChange; mockSessionCallbacks.onTransportCapabilityChange = sessionConfig.onTransportCapabilityChange; + mockSessionCallbacks.onTransportCapabilitiesChange = + sessionConfig.onTransportCapabilitiesChange; mockSessionCallbacks.onEvent = sessionConfig.onEvent; mockSessionCallbacks.onMessageQueued = sessionConfig.onMessageQueued; mockSessionCallbacks.onMessageCompleted = sessionConfig.onMessageCompleted; @@ -415,6 +420,7 @@ describe('createSessionManager', () => { mockSessionCallbacks.onResolved = undefined; mockSessionCallbacks.onRemoteModelStateChange = undefined; mockSessionCallbacks.onTransportCapabilityChange = undefined; + mockSessionCallbacks.onTransportCapabilitiesChange = undefined; mockSessionCallbacks.onEvent = undefined; mockSessionCallbacks.onMessageQueued = undefined; mockSessionCallbacks.onMessageCompleted = undefined; @@ -4104,4 +4110,220 @@ describe('createSessionManager — paginated initial snapshot + loadOlderMessage atomValue(config.store, mgr.atoms.messagesList).map(m => m.info.id) ).toEqual(['msg-current']); }); + + // ------------------------------------------------------------------------- + // supportsAttachments gate + // ------------------------------------------------------------------------- + + describe('supportsAttachments gate', () => { + it('heartbeat absent -> true upgrade flips supportsAttachments true', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + + await mgr.switchSession(kiloId('ses-1')); + mockSessionCallbacks.onResolved?.({ type: 'remote', kiloSessionId: kiloId('ses-1') }); + + expect(atomValue(config.store, mgr.atoms.supportsAttachments)).toBe(false); + + mockSessionCallbacks.onTransportCapabilitiesChange?.({ attachments: true }); + expect(atomValue(config.store, mgr.atoms.supportsAttachments)).toBe(true); + }); + + it('true -> false downgrade flips supportsAttachments false', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + + await mgr.switchSession(kiloId('ses-1')); + mockSessionCallbacks.onResolved?.({ type: 'remote', kiloSessionId: kiloId('ses-1') }); + + mockSessionCallbacks.onTransportCapabilitiesChange?.({ attachments: true }); + expect(atomValue(config.store, mgr.atoms.supportsAttachments)).toBe(true); + + mockSessionCallbacks.onTransportCapabilitiesChange?.({ attachments: false }); + expect(atomValue(config.store, mgr.atoms.supportsAttachments)).toBe(false); + }); + + it('true -> absent flips supportsAttachments false', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + + await mgr.switchSession(kiloId('ses-1')); + mockSessionCallbacks.onResolved?.({ type: 'remote', kiloSessionId: kiloId('ses-1') }); + + mockSessionCallbacks.onTransportCapabilitiesChange?.({ attachments: true }); + expect(atomValue(config.store, mgr.atoms.supportsAttachments)).toBe(true); + + mockSessionCallbacks.onTransportCapabilitiesChange?.(undefined); + expect(atomValue(config.store, mgr.atoms.supportsAttachments)).toBe(false); + }); + }); + + // ------------------------------------------------------------------------- + // send attachment gating + // ------------------------------------------------------------------------- + + describe('send attachment gating', () => { + it('cloud-agent send with attachments forwards to session.send', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + + await mgr.switchSession(kiloId('ses-1')); + const attachments = { + path: '12345678-1234-4234-9234-123456789abc', + files: ['87654321-4321-4321-8321-cba987654321.md'], + }; + mockSession.send.mockResolvedValue(undefined); + + const accepted = await mgr.send({ + payload: { type: 'prompt', prompt: 'Hello', mode: 'code', model: 'claude-3-5-sonnet' }, + attachments, + }); + + expect(accepted).toBe(true); + expect(mockSession.send).toHaveBeenCalledWith({ + messageId: expect.stringMatching(/^msg_/), + payload: { + type: 'prompt', + prompt: 'Hello', + mode: 'code', + model: { providerID: 'kilo', modelID: 'claude-3-5-sonnet' }, + }, + attachments, + images: undefined, + }); + }); + + it('cloud-agent send with attachments: undefined drops the field', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + + await mgr.switchSession(kiloId('ses-1')); + mockSession.send.mockResolvedValue(undefined); + + await mgr.send({ + payload: { type: 'prompt', prompt: 'Hello', mode: 'code', model: 'claude-3-5-sonnet' }, + attachments: undefined, + }); + + expect(mockSession.send).toHaveBeenCalledWith({ + messageId: expect.stringMatching(/^msg_/), + payload: { + type: 'prompt', + prompt: 'Hello', + mode: 'code', + model: { providerID: 'kilo', modelID: 'claude-3-5-sonnet' }, + }, + images: undefined, + }); + expect( + ( + mockSession.send.mock.calls[0] as [ + { + attachments?: unknown; + }, + ] + )[0].attachments + ).toBeUndefined(); + }); + + it('non-capable remote + non-empty attachmentParts rejects before transport send', async () => { + const onSendFailed = jest.fn(); + const config = createMockConfig({ onSendFailed }); + const mgr = createSessionManager(config); + + await mgr.switchSession(kiloId('ses-1')); + mockSessionCallbacks.onResolved?.({ type: 'remote', kiloSessionId: kiloId('ses-1') }); + mockSession.send.mockResolvedValue(undefined); + + const attachmentParts: RemoteAttachmentPart[] = [ + { + type: 'file', + mime: 'text/plain', + filename: 'file.txt', + url: 'https://example.com/file.txt', + }, + ]; + + const accepted = await mgr.send({ + payload: { type: 'prompt', prompt: 'Hello', mode: 'code', model: 'claude-3-5-sonnet' }, + attachmentParts, + }); + + expect(accepted).toBe(false); + expect(mockSession.send).not.toHaveBeenCalled(); + expect(atomValue(config.store, mgr.atoms.failedPrompt)).toBe('Hello'); + expect(onSendFailed).toHaveBeenCalledWith( + 'Hello', + expect.any(String), + expect.objectContaining({ + message: 'Only capable remote CLI sessions support attachments', + }) + ); + }); + + it('non-cloud + attachments rejects before transport send', async () => { + const onSendFailed = jest.fn(); + const config = createMockConfig({ onSendFailed }); + const mgr = createSessionManager(config); + + await mgr.switchSession(kiloId('ses-1')); + mockSessionCallbacks.onResolved?.({ type: 'remote', kiloSessionId: kiloId('ses-1') }); + mockSession.send.mockResolvedValue(undefined); + + const accepted = await mgr.send({ + payload: { type: 'prompt', prompt: 'Hello', mode: 'code', model: 'claude-3-5-sonnet' }, + attachments: { + path: '12345678-1234-4234-9234-123456789abc', + files: ['87654321-4321-4321-8321-cba987654321.md'], + }, + }); + + expect(accepted).toBe(false); + expect(mockSession.send).not.toHaveBeenCalled(); + expect(atomValue(config.store, mgr.atoms.failedPrompt)).toBe('Hello'); + expect(onSendFailed).toHaveBeenCalledWith( + 'Hello', + expect.any(String), + expect.objectContaining({ + message: 'Only Cloud Agent sessions support attachments', + }) + ); + }); + + it('capable remote + attachmentParts forwards to session.send', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + + await mgr.switchSession(kiloId('ses-1')); + mockSessionCallbacks.onResolved?.({ type: 'remote', kiloSessionId: kiloId('ses-1') }); + mockSessionCallbacks.onTransportCapabilitiesChange?.({ attachments: true }); + mockSession.send.mockResolvedValue(undefined); + + const attachmentParts: RemoteAttachmentPart[] = [ + { + type: 'file', + mime: 'text/plain', + filename: 'file.txt', + url: 'https://example.com/file.txt', + }, + ]; + + const accepted = await mgr.send({ + payload: { type: 'prompt', prompt: 'Hello', mode: 'code', model: 'claude-3-5-sonnet' }, + attachmentParts, + }); + + expect(accepted).toBe(true); + expect(mockSession.send).toHaveBeenCalledWith({ + messageId: expect.stringMatching(/^msg_/), + payload: { + type: 'prompt', + prompt: 'Hello', + mode: 'code', + }, + images: undefined, + attachmentParts, + }); + }); + }); }); 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 8443e8b9b4..bfca62d6d8 100644 --- a/apps/web/src/lib/cloud-agent-sdk/session-manager.ts +++ b/apps/web/src/lib/cloud-agent-sdk/session-manager.ts @@ -1,7 +1,12 @@ import type { CloudAgentAttachments } from '@/lib/cloud-agent/constants'; import type { Images } from '@/lib/images-schema'; import { errorShapeSchema } from './schemas'; -import type { SendCommandPayload, SendPromptPayload, TransportSendPayload } from './transport'; +import type { + RemoteAttachmentPart, + SendCommandPayload, + SendPromptPayload, + TransportSendPayload, +} from './transport'; import { modelRefsEqual } from './remote-model-catalog'; import type { ModelRef, @@ -269,6 +274,16 @@ type SessionManager = { payload: SessionManagerSendPayload; attachments?: CloudAgentAttachments; images?: Images; + /** + * Ready file parts to forward to a CAPABLE remote CLI session (the CLI + * advertised `capabilities.attachments: true` in its most recent + * heartbeat). Distinct from the cloud-only `attachments` field: cloud + * sessions use `attachments`, remote sessions use `attachmentParts`. + * Session-manager enforces the gate — a non-null payload for a + * non-capable session is rejected with a typed error before it can + * reach the transport. + */ + attachmentParts?: RemoteAttachmentPart[]; }): Promise; setRemoteModelOverride(override: RemoteModelOverride | null): void; retryRemoteModels(): void; @@ -502,6 +517,16 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { let switchGeneration = 0; let currentSession: CloudAgentSession | null = null; let activeSessionType: ActiveSessionType | null = null; + /** + * Latest per-session capabilities reported by the live CLI transport's + * `onTransportCapabilitiesChange` callback. Captured here so the + * `supportsAttachments` gate can be recomputed on every heartbeat-driven + * capability change (upgrade, downgrade, reconnect, absent) — not only at + * the initial `onResolved` moment. `undefined` means the CLI has not + * reported any (older CLIs, mid-reconnect, or a session that the active + * CLI no longer claims). + */ + let currentCapabilities: { attachments?: boolean } | undefined = undefined; let observedModelSource: ObservedModelSource | null = null; // True while a connect/reconnect cycle is still replaying its message // history; false once live events are flowing. See clearOverrideIfDiverged. @@ -582,6 +607,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { loadOlderGeneration += 1; olderMessagesInFlight = null; olderMessagesTerminal = false; + currentCapabilities = undefined; } function setChildSessionHydrationState( @@ -662,6 +688,35 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { store.set(canInterruptAtom, session.canInterrupt); } + /** + * Recompute the `supportsAttachments` gate for the active session. Called + * on every `onResolved` (initial resolution) AND every + * `onTransportCapabilitiesChange` (heartbeat upgrade/downgrade/reconnect/ + * absent) so the UI gate tracks the CLI's most recent advertisement. + * + * Rules: + * - `cloud-agent`: always supports attachments (S3a is a no-op for + * cloud-agent sessions, but cloud-agent attachments flow through + * the existing `attachments` field, not the new `attachmentParts`). + * - `remote`: supports attachments only when the live CLI reported + * `capabilities.attachments === true` in its most recent heartbeat + * or `sessions.list`. Any other state (absent, false, mid-reconnect) + * → `false`, matching today's "no paperclip" parity for non-capable + * remote sessions. + * - `read-only`: never supports attachments. + */ + function recomputeSupportsAttachments(sessionType: ActiveSessionType | null): void { + let supports: boolean; + if (sessionType === 'cloud-agent') { + supports = true; + } else if (sessionType === 'remote') { + supports = currentCapabilities?.attachments === true; + } else { + supports = false; + } + store.set(supportsAttachmentsAtom, supports); + } + function updateObservedModel(model: ModelSelection, source: ObservedModelSource): void { observedModelSource = source; // Only churn the atom when the selection actually changes: the incoming @@ -1093,7 +1148,13 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { activeSessionType = resolved.type; store.set(sessionTypeAtom, resolved.type); store.set(activeSessionTypeAtom, resolved.type); - store.set(supportsAttachmentsAtom, resolved.type === 'cloud-agent'); + // Seed capabilities from the resolved session so the initial gate + // reflects whatever the mobile-side `resolveSession` adapter had + // to work with (e.g. the current `activeSessions.list` snapshot). + // Subsequent heartbeat changes arrive via `onTransportCapabilitiesChange` + // and overwrite this. + currentCapabilities = resolved.type === 'remote' ? resolved.capabilities : undefined; + recomputeSupportsAttachments(resolved.type); updateCapabilityAtoms(session); }, onRemoteModelStateChange: handleRemoteModelStateChange, @@ -1105,6 +1166,11 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { if (expectedGeneration !== switchGeneration) return; if (currentSession === session) updateCapabilityAtoms(session); }, + onTransportCapabilitiesChange: capabilities => { + if (expectedGeneration !== switchGeneration) return; + currentCapabilities = capabilities; + recomputeSupportsAttachments(activeSessionType); + }, onReplayComplete: () => { if (expectedGeneration !== switchGeneration) return; remoteHistoryReplaying = false; @@ -1212,6 +1278,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { payload: SessionManagerSendPayload; attachments?: CloudAgentAttachments; images?: Images; + attachmentParts?: RemoteAttachmentPart[]; }): Promise { store.set(errorAtom, null); if (store.get(agentStatusAtom).type !== 'disconnected') { @@ -1261,14 +1328,31 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { try { if (!currentSession) throw new Error('No active session'); if (input.attachments && sessionType !== 'cloud-agent') { + // The cloud-only `attachments` field is exclusive to cloud-agent + // sessions. Remote CLI sessions (capable or not) go through the + // new `attachmentParts` path. Reject loudly if a caller mixes them + // up — this is a programmer error, not a user-recoverable state. throw new Error('Only Cloud Agent sessions support attachments'); } + if (input.attachmentParts && input.attachmentParts.length > 0) { + if (sessionType !== 'remote' || currentCapabilities?.attachments !== true) { + // A non-null `attachmentParts` for a non-capable session is a + // UI-bug: the paperclip is supposed to be hidden whenever this + // gate fails, so we should never see payload here. Refuse to + // forward rather than silently drop — same policy as the + // cloud-only branch above. + throw new Error('Only capable remote CLI sessions support attachments'); + } + } await currentSession.send({ payload: transportPayload, messageId, ...(input.attachments ? { attachments: input.attachments } : {}), images: input.images, ...(sessionType === 'remote' && remoteModelOverride ? { remoteModelOverride } : {}), + ...(input.attachmentParts && input.attachmentParts.length > 0 + ? { attachmentParts: input.attachmentParts } + : {}), }); if (sessionType === 'remote' && kiloSessionId) { diff --git a/apps/web/src/lib/cloud-agent-sdk/session.ts b/apps/web/src/lib/cloud-agent-sdk/session.ts index 13f0b196e1..62c895d9e0 100644 --- a/apps/web/src/lib/cloud-agent-sdk/session.ts +++ b/apps/web/src/lib/cloud-agent-sdk/session.ts @@ -24,6 +24,7 @@ import type { UserWebConnection } from './user-web-connection'; import type { CloudAgentApi, CloudAgentStreamTicketResult, + RemoteAttachmentPart, TransportFactory, TransportSink, TransportSendPayload, @@ -75,6 +76,7 @@ type CloudAgentSessionConfig = { onRemoteModelStateChange?: (state: RemoteModelState) => void; onRemoteCommandStateChange?: (state: RemoteCommandState) => void; onTransportCapabilityChange?: () => void; + onTransportCapabilitiesChange?: (capabilities: { attachments?: boolean } | undefined) => void; onSessionCreated?: (info: SessionInfo) => void; onSessionUpdated?: (info: SessionInfo) => void; onReplayComplete?: () => void; @@ -93,6 +95,14 @@ type CloudAgentSessionSendInput = { attachments?: CloudAgentAttachments; images?: Images; remoteModelOverride?: RemoteModelOverride; + /** + * Ready file parts to forward to a CAPABLE remote CLI session. The + * session-manager gate already enforces that this is only non-empty for + * a `remote` session whose CLI has advertised `capabilities.attachments: + * true`; transports that don't support the path (cloud-agent, read-only, + * non-capable remote) simply ignore it. + */ + attachmentParts?: RemoteAttachmentPart[]; }; type CloudAgentSessionAnswerInput = { @@ -288,6 +298,7 @@ function createCloudAgentSession(config: CloudAgentSessionConfig): CloudAgentSes onRemoteModelStateChange: config.onRemoteModelStateChange, onRemoteCommandStateChange: config.onRemoteCommandStateChange, onCapabilityChange: config.onTransportCapabilityChange, + onCapabilitiesChange: config.onTransportCapabilitiesChange, }); } case 'cloud-agent': { diff --git a/apps/web/src/lib/cloud-agent-sdk/transport.ts b/apps/web/src/lib/cloud-agent-sdk/transport.ts index f4e99896ae..b3578b01f4 100644 --- a/apps/web/src/lib/cloud-agent-sdk/transport.ts +++ b/apps/web/src/lib/cloud-agent-sdk/transport.ts @@ -10,6 +10,24 @@ import type { Images } from '@/lib/images-schema'; import type { CloudAgentSessionId, KiloSessionId } from './types'; import type { ModelRef, RemoteModelOverride } from './remote-model-catalog'; +/** + * Ready-to-render file part for a remote-CLI `send_message` call. Distinct + * from the cloud-only `attachments` field on {@link TransportSendInput}: + * the remote CLI fetches these from the `url` and classifies them by + * their `filename` (which MUST be the server-issued `.`) and + * `mime` (which the caller derives from the validated extension). + * + * Deliberately the minimum the CLI needs: `id` / `sessionID` / `messageID` + * are not part of the send payload because the CLI assigns them when it + * materializes the user message. + */ +type RemoteAttachmentPart = { + type: 'file'; + mime: string; + filename: string; + url: string; +}; + type CloudAgentStreamTicket = { ticket: string; /** Unix timestamp in seconds when the ticket expires. */ @@ -67,6 +85,15 @@ type TransportSendInput = { attachments?: CloudAgentAttachments; images?: Images; remoteModelOverride?: RemoteModelOverride; + /** + * Ready file parts to append to the remote CLI's `send_message` `parts` + * array (after the text part). Distinct from the cloud-only `attachments` + * field: this path is for CAPABLE remote CLI sessions (the session + * advertised `capabilities.attachments: true` in its most recent + * heartbeat). Transports that don't support the path (cloud-agent, read- + * only, non-capable remote) ignore it. + */ + attachmentParts?: RemoteAttachmentPart[]; }; /** Lifecycle interface for a transport. */ @@ -141,6 +168,7 @@ export type { CloudAgentSendPayload, CloudAgentStreamTicket, CloudAgentStreamTicketResult, + RemoteAttachmentPart, TransportFactory, TransportSink, Transport, diff --git a/apps/web/src/lib/cloud-agent-sdk/types.ts b/apps/web/src/lib/cloud-agent-sdk/types.ts index 18781e3c93..27659f9244 100644 --- a/apps/web/src/lib/cloud-agent-sdk/types.ts +++ b/apps/web/src/lib/cloud-agent-sdk/types.ts @@ -198,7 +198,18 @@ export type ServiceStateSnapshot = { // --------------------------------------------------------------------------- export type ResolvedSession = - | { type: 'remote'; kiloSessionId: KiloSessionId } + | { + type: 'remote'; + kiloSessionId: KiloSessionId; + /** + * Per-session capabilities reported by the owning CLI's most recent + * heartbeat or `sessions.list`. A `remote` session supports attachments + * only when `capabilities?.attachments === true`; absent / false is a + * structural no-attachments state (older CLIs, CLIs that have not yet + * advertised the capability, or CLIs that have downgraded). + */ + capabilities?: { attachments?: boolean }; + } | { type: 'cloud-agent'; kiloSessionId: KiloSessionId; cloudAgentSessionId: CloudAgentSessionId } | { type: 'read-only'; kiloSessionId: KiloSessionId }; diff --git a/apps/web/src/lib/cloud-agent/constants.ts b/apps/web/src/lib/cloud-agent/constants.ts index a33cf2a78d..69896be5df 100644 --- a/apps/web/src/lib/cloud-agent/constants.ts +++ b/apps/web/src/lib/cloud-agent/constants.ts @@ -53,6 +53,128 @@ export const CLOUD_AGENT_ATTACHMENT_MAX_COUNT = 5; export const CLOUD_AGENT_ATTACHMENT_MAX_SIZE_BYTES = 5 * 1024 * 1024; // 5MB export const CLOUD_AGENT_ATTACHMENT_PRESIGNED_URL_EXPIRY_SECONDS = 900; // 15 min +/** + * Filename-derived extensions that are never stored or downloaded, regardless + * of the caller-supplied MIME. The allow-list is enforced at every layer that + * validates a stored filename: the web tRPC schemas, the cloud-agent-next + * persistence schema, the cloud-agent-next runtime download helper, and the + * R2 download-presign helper. + */ +export const CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS = [ + 'exe', + 'dll', + 'msi', + 'com', + 'scr', + 'apk', + 'ipa', + 'dmg', + 'pkg', +] as const; + +export type CloudAgentAttachmentDeniedExtension = + (typeof CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS)[number]; + +/** + * Wire-format pattern for a caller-supplied extension on the upload-presign + * input. The extension is lowercased, alpha-numeric, and at most 16 chars. + * The deny-list is checked separately so this regex stays a pure shape rule. + */ +export const CLOUD_AGENT_ATTACHMENT_EXTENSION_REGEX = /^[a-z0-9]{1,16}$/; + +/** + * Relaxed contentType shape used when the caller supplies an extension and we + * accept any reasonable MIME. The 128-char cap is a safety net; the regex + * matches `type/subtype` plus a small set of structural characters. + */ +export const CLOUD_AGENT_ATTACHMENT_RELAXED_CONTENT_TYPE_REGEX = /^[\w.+-]+\/[\w.+-]+$/; +export const CLOUD_AGENT_ATTACHMENT_MAX_CONTENT_TYPE_LENGTH = 128; + +/** + * Server-stored filename regex. The UUID prefix is preserved so the + * presign/download helpers can derive the R2 key from caller-supplied + * metadata alone. The suffix is the union of the legacy 9-extension allow-list + * and the relaxed `^[a-z0-9]{1,16}$` pattern; the deny-list is enforced by + * the shared validation helper. + */ +export const CLOUD_AGENT_ATTACHMENT_RELAXED_FILENAME_REGEX = + /^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.[a-z0-9]{1,16}$/; + +/** + * Canonical extension → MIME map. The STORED extension is the single source + * of MIME truth for the worker (which derives the prompt MIME from the + * filename suffix). A new extension without a mapping falls back to + * `application/octet-stream`; the `bin` no-extension fallback also resolves + * to `application/octet-stream`. + */ +export const CLOUD_AGENT_ATTACHMENT_EXTENSION_TO_MIME: Record = { + // images + png: 'image/png', + jpg: 'image/jpeg', + jpeg: 'image/jpeg', + webp: 'image/webp', + gif: 'image/gif', + // pdf + pdf: 'application/pdf', + // text-y formats — surfaced as `text/plain` for the prompt; the worker's + // prompt-shape rule still routes these through the file part path. + txt: 'text/plain', + md: 'text/plain', + csv: 'text/plain', + log: 'text/plain', + json: 'text/plain', + xml: 'text/plain', + yaml: 'text/plain', + yml: 'text/plain', + toml: 'text/plain', + ini: 'text/plain', + html: 'text/plain', + css: 'text/plain', + js: 'text/plain', + jsx: 'text/plain', + ts: 'text/plain', + tsx: 'text/plain', + py: 'text/plain', + rb: 'text/plain', + go: 'text/plain', + rs: 'text/plain', + java: 'text/plain', + c: 'text/plain', + h: 'text/plain', + cpp: 'text/plain', + hpp: 'text/plain', + sh: 'text/plain', + sql: 'text/plain', +}; + +export const CLOUD_AGENT_ATTACHMENT_FALLBACK_MIME = 'application/octet-stream'; +export const CLOUD_AGENT_ATTACHMENT_FALLBACK_EXTENSION = 'bin'; + +/** + * Resolve the canonical MIME for a stored extension. Returns the fallback + * for unknown / missing / invalid extensions; the worker relies on this as + * the single source of truth so `getPromptMime` cannot diverge. + */ +export function getPromptMimeForExtension(extension: string | undefined | null): string { + if (!extension) return CLOUD_AGENT_ATTACHMENT_FALLBACK_MIME; + const normalized = extension.toLowerCase(); + return ( + CLOUD_AGENT_ATTACHMENT_EXTENSION_TO_MIME[normalized] ?? CLOUD_AGENT_ATTACHMENT_FALLBACK_MIME + ); +} + +/** + * Normalize a caller-supplied extension: empty / non-conforming → `bin`. The + * R2 key, the stored filename, and the download-presign helper all use the + * normalized value so the wire-payload and the R2 basename can never disagree. + */ +export function normalizeAttachmentExtension(extension: string | undefined | null): string { + if (extension && CLOUD_AGENT_ATTACHMENT_EXTENSION_REGEX.test(extension)) { + return extension.toLowerCase(); + } + return CLOUD_AGENT_ATTACHMENT_FALLBACK_EXTENSION; +} + export type CloudAgentAttachments = { path: string; files: string[]; diff --git a/apps/web/src/lib/r2/cloud-agent-attachments.test.ts b/apps/web/src/lib/r2/cloud-agent-attachments.test.ts index 894dc716e4..b1a94d0388 100644 --- a/apps/web/src/lib/r2/cloud-agent-attachments.test.ts +++ b/apps/web/src/lib/r2/cloud-agent-attachments.test.ts @@ -1,5 +1,6 @@ import { beforeAll, beforeEach, describe, expect, it, jest } from '@jest/globals'; import type { + generateCloudAgentAttachmentDownloadUrl as GenerateCloudAgentAttachmentDownloadUrl, generateCloudAgentAttachmentUploadUrl as GenerateCloudAgentAttachmentUploadUrl, generateImageUploadUrl as GenerateImageUploadUrl, } from './cloud-agent-attachments'; @@ -17,6 +18,9 @@ jest.mock('@aws-sdk/client-s3', () => ({ PutObjectCommand: jest .fn() .mockImplementation((input: unknown) => ({ input, name: 'PutObjectCommand' })), + GetObjectCommand: jest + .fn() + .mockImplementation((input: unknown) => ({ input, name: 'GetObjectCommand' })), })); const MESSAGE_UUID = '12345678-1234-4234-9234-123456789abc'; @@ -25,6 +29,7 @@ const ATTACHMENT_ID = '87654321-4321-4321-8321-cba987654321'; let mockGetSignedUrl: jest.Mock< (client: unknown, command: unknown, options: unknown) => Promise >; +let generateCloudAgentAttachmentDownloadUrl: typeof GenerateCloudAgentAttachmentDownloadUrl; let generateCloudAgentAttachmentUploadUrl: typeof GenerateCloudAgentAttachmentUploadUrl; let generateImageUploadUrl: typeof GenerateImageUploadUrl; @@ -35,6 +40,7 @@ describe('cloud-agent attachment upload URL signing', () => { mockGetSignedUrl = signer.getSignedUrl as unknown as jest.Mock< (client: unknown, command: unknown, options: unknown) => Promise >; + generateCloudAgentAttachmentDownloadUrl = attachments.generateCloudAgentAttachmentDownloadUrl; generateCloudAgentAttachmentUploadUrl = attachments.generateCloudAgentAttachmentUploadUrl; generateImageUploadUrl = attachments.generateImageUploadUrl; }); @@ -48,7 +54,7 @@ describe('cloud-agent attachment upload URL signing', () => { ['md', 'text/markdown'], ['csv', 'text/csv'], ] as const)( - 'issues a server-derived .%s suffix and attachment metadata for %s', + 'issues a server-derived .%s suffix and attachment metadata for %s (legacy no-extension path)', async (extension, contentType) => { const result = await generateCloudAgentAttachmentUploadUrl({ userId: 'user-1', @@ -58,7 +64,10 @@ describe('cloud-agent attachment upload URL signing', () => { contentLength: 42, }); - const command = mockGetSignedUrl.mock.calls[0]?.[1] as { input: Record }; + const command = mockGetSignedUrl.mock.calls[0]?.[1] as { + input: Record; + name: string; + }; expect(command.input).toMatchObject({ Bucket: 'attachment-bucket', Key: `user-1/cloud-agent/${MESSAGE_UUID}/${ATTACHMENT_ID}.${extension}`, @@ -74,6 +83,45 @@ describe('cloud-agent attachment upload URL signing', () => { } ); + it('derives the R2 key suffix from a validated extension when extension is provided', async () => { + const result = await generateCloudAgentAttachmentUploadUrl({ + userId: 'user-1', + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + contentType: 'application/x-kilo-binary', + contentLength: 42, + extension: 'kilo', + }); + + const command = mockGetSignedUrl.mock.calls[0]?.[1] as { input: Record }; + expect(command.input).toMatchObject({ + Bucket: 'attachment-bucket', + Key: `user-1/cloud-agent/${MESSAGE_UUID}/${ATTACHMENT_ID}.kilo`, + ContentType: 'application/x-kilo-binary', + ContentLength: 42, + Metadata: { + userId: 'user-1', + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + }, + }); + expect(result.key).toBe(`user-1/cloud-agent/${MESSAGE_UUID}/${ATTACHMENT_ID}.kilo`); + }); + + it('rejects a deny-listed extension in the upload helper before signing', async () => { + await expect( + generateCloudAgentAttachmentUploadUrl({ + userId: 'user-1', + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + contentType: 'application/octet-stream', + contentLength: 42, + extension: 'exe', + }) + ).rejects.toThrow('exe'); + expect(mockGetSignedUrl).not.toHaveBeenCalled(); + }); + it('preserves image-only App Builder key generation and metadata', async () => { await generateImageUploadUrl({ service: 'app-builder', @@ -91,3 +139,40 @@ describe('cloud-agent attachment upload URL signing', () => { }); }); }); + +describe('cloud-agent attachment download URL signing', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetSignedUrl.mockResolvedValue('https://example.test/signed'); + }); + + it('derives the R2 key from the caller (author) and signs a GET with the validated filename', async () => { + const result = await generateCloudAgentAttachmentDownloadUrl({ + userId: 'user-1', + messageUuid: MESSAGE_UUID, + filename: `${ATTACHMENT_ID}.kilo`, + }); + + const command = mockGetSignedUrl.mock.calls[0]?.[1] as { + input: Record; + name: string; + }; + expect(command.input).toMatchObject({ + Bucket: 'attachment-bucket', + Key: `user-1/cloud-agent/${MESSAGE_UUID}/${ATTACHMENT_ID}.kilo`, + }); + expect(command.name).toBe('GetObjectCommand'); + expect(result.key).toBe(`user-1/cloud-agent/${MESSAGE_UUID}/${ATTACHMENT_ID}.kilo`); + }); + + it('rejects a deny-listed extension in the filename before signing', async () => { + await expect( + generateCloudAgentAttachmentDownloadUrl({ + userId: 'user-1', + messageUuid: MESSAGE_UUID, + filename: `${ATTACHMENT_ID}.exe`, + }) + ).rejects.toThrow(); + expect(mockGetSignedUrl).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/src/lib/r2/cloud-agent-attachments.ts b/apps/web/src/lib/r2/cloud-agent-attachments.ts index e1b7704b24..b4b97195d4 100644 --- a/apps/web/src/lib/r2/cloud-agent-attachments.ts +++ b/apps/web/src/lib/r2/cloud-agent-attachments.ts @@ -1,15 +1,16 @@ -import { PutObjectCommand } from '@aws-sdk/client-s3'; +import { GetObjectCommand, PutObjectCommand } from '@aws-sdk/client-s3'; import { getSignedUrl } from '@aws-sdk/s3-request-presigner'; import { + CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS, CLOUD_AGENT_ATTACHMENT_MIME_TO_EXTENSION, CLOUD_AGENT_ATTACHMENT_PRESIGNED_URL_EXPIRY_SECONDS, CLOUD_AGENT_IMAGE_MIME_TO_EXTENSION, CLOUD_AGENT_IMAGE_PRESIGNED_URL_EXPIRY_SECONDS, + normalizeAttachmentExtension, + type CloudAgentAttachmentAllowedType, + type CloudAgentImageAllowedType, } from '@/lib/cloud-agent/constants'; -import type { - CloudAgentAttachmentAllowedType, - CloudAgentImageAllowedType, -} from '@/lib/cloud-agent/constants'; +import { cloudAgentRelaxedAttachmentFilenameSchema } from '@/routers/cloud-agent-next-schemas'; import { r2Client, r2CloudAgentAttachmentsBucketName } from '@/lib/r2/client'; type Service = 'app-builder' | 'cloud-agent'; @@ -86,8 +87,15 @@ export type GenerateCloudAgentAttachmentUploadUrlParams = { userId: string; messageUuid: string; attachmentId: string; - contentType: CloudAgentAttachmentAllowedType; + contentType: CloudAgentAttachmentAllowedType | (string & {}); contentLength: number; + /** + * Optional caller-supplied extension. When provided, the R2 key suffix + * derives from the validated extension (after deny-list filtering and + * normalization) and `contentType` is treated as a relaxed MIME label. + * When omitted, the legacy MIME→extension map is used. + */ + extension?: string; }; export type GenerateCloudAgentAttachmentUploadUrlResult = { @@ -102,9 +110,18 @@ export async function generateCloudAgentAttachmentUploadUrl({ attachmentId, contentType, contentLength, + extension, }: GenerateCloudAgentAttachmentUploadUrlParams): Promise { - const extension = CLOUD_AGENT_ATTACHMENT_MIME_TO_EXTENSION[contentType]; - const key = `${userId}/cloud-agent/${messageUuid}/${attachmentId}.${extension}`; + const suffix = extension + ? (() => { + const normalized = normalizeAttachmentExtension(extension); + if ((CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS as readonly string[]).includes(normalized)) { + throw new Error(`Attachment extension "${normalized}" is not allowed`); + } + return normalized; + })() + : CLOUD_AGENT_ATTACHMENT_MIME_TO_EXTENSION[contentType as CloudAgentAttachmentAllowedType]; + const key = `${userId}/cloud-agent/${messageUuid}/${attachmentId}.${suffix}`; const command = new PutObjectCommand({ Bucket: r2CloudAgentAttachmentsBucketName, Key: key, @@ -130,3 +147,50 @@ export async function generateCloudAgentAttachmentUploadUrl({ ).toISOString(), }; } + +export type GenerateCloudAgentAttachmentDownloadUrlParams = { + userId: string; + messageUuid: string; + filename: string; +}; + +export type GenerateCloudAgentAttachmentDownloadUrlResult = { + signedUrl: string; + key: string; + expiresAt: string; +}; + +/** + * Presign a GET for a stored attachment. The key prefix is derived from the + * caller (author) — do NOT reproduce the session-owner divergence here. The + * filename is re-validated through the shared relaxed-regex + deny-list + * schema so the key can never contain an unsafe basename regardless of the + * caller. No org mirror: remote CLI sessions are personal. + */ +export async function generateCloudAgentAttachmentDownloadUrl({ + userId, + messageUuid, + filename, +}: GenerateCloudAgentAttachmentDownloadUrlParams): Promise { + // The schema throws on violation; the caller is responsible for passing a + // parsed value, but we re-validate here as a defense in depth. + const parsed = cloudAgentRelaxedAttachmentFilenameSchema.parse(filename); + const key = `${userId}/cloud-agent/${messageUuid}/${parsed}`; + + const command = new GetObjectCommand({ + Bucket: r2CloudAgentAttachmentsBucketName, + Key: key, + }); + + const signedUrl = await getSignedUrl(r2Client, command, { + expiresIn: CLOUD_AGENT_ATTACHMENT_PRESIGNED_URL_EXPIRY_SECONDS, + }); + + return { + signedUrl, + key, + expiresAt: new Date( + Date.now() + CLOUD_AGENT_ATTACHMENT_PRESIGNED_URL_EXPIRY_SECONDS * 1000 + ).toISOString(), + }; +} diff --git a/apps/web/src/routers/active-sessions-router.schema.test.ts b/apps/web/src/routers/active-sessions-router.schema.test.ts new file mode 100644 index 0000000000..c14abf8cda --- /dev/null +++ b/apps/web/src/routers/active-sessions-router.schema.test.ts @@ -0,0 +1,80 @@ +import { describe, it, expect } from '@jest/globals'; +import { activeSessionSchema } from './active-sessions-router'; + +describe('activeSessionSchema capabilities', () => { + it('accepts a session row with capabilities.attachments: true', () => { + const row = { + id: 's1', + status: 'busy', + title: 'Fix bug', + connectionId: 'cli-1', + capabilities: { attachments: true }, + }; + expect(activeSessionSchema.safeParse(row).success).toBe(true); + }); + + it('accepts a session row with capabilities.attachments: false', () => { + const row = { + id: 's1', + status: 'busy', + title: 'Fix bug', + connectionId: 'cli-1', + capabilities: { attachments: false }, + }; + expect(activeSessionSchema.safeParse(row).success).toBe(true); + }); + + it('accepts a session row with an empty capabilities object', () => { + const row = { + id: 's1', + status: 'busy', + title: 'Fix bug', + connectionId: 'cli-1', + capabilities: {}, + }; + expect(activeSessionSchema.safeParse(row).success).toBe(true); + }); + + it('accepts a session row with an absent capabilities field (legacy CLI)', () => { + const row = { + id: 's1', + status: 'busy', + title: 'Fix bug', + connectionId: 'cli-1', + }; + const result = activeSessionSchema.safeParse(row); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.capabilities).toBeUndefined(); + } + }); + + it('rejects a non-boolean capabilities.attachments value', () => { + const row = { + id: 's1', + status: 'busy', + title: 'Fix bug', + connectionId: 'cli-1', + capabilities: { attachments: 'yes' }, + }; + expect(activeSessionSchema.safeParse(row).success).toBe(false); + }); + + it('rejects unknown capability keys (strict object)', () => { + const row = { + id: 's1', + status: 'busy', + title: 'Fix bug', + connectionId: 'cli-1', + capabilities: { terminal: true }, + }; + // The default zod behavior strips unknown keys rather than rejecting — + // assert that the unknown key is dropped, not preserved, so consumers + // never see a flag the cloud did not advertise. + const result = activeSessionSchema.safeParse(row); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.capabilities).toEqual({}); + } + }); +}); diff --git a/apps/web/src/routers/active-sessions-router.test.ts b/apps/web/src/routers/active-sessions-router.test.ts index b43739bd29..5dd9a68080 100644 --- a/apps/web/src/routers/active-sessions-router.test.ts +++ b/apps/web/src/routers/active-sessions-router.test.ts @@ -6,6 +6,13 @@ import { organizations, organization_memberships } from '@kilocode/db/schema'; import type { User, Organization } from '@kilocode/db/schema'; import type { createCallerForUser as CreateCallerForUser } from '@/routers/test-utils'; +// NOTE: `activeSessionSchema` validation is covered in the sibling +// `active-sessions-router.schema.test.ts`. It is deliberately NOT imported +// here: a static import of the router module evaluates `config.server.ts` +// (and freezes `SESSION_INGEST_WORKER_URL` at the empty `.env.test` value) +// before the `beforeAll` below can set the env — which is exactly what the +// dynamic-import dance is designed to avoid. + // `.env.test` sets SESSION_INGEST_WORKER_URL to '' (shared fixture used by // other test files too — do not change it here). `createCallerForUser`'s // import chain (test-utils -> trpc/init -> ...) transitively loads diff --git a/apps/web/src/routers/active-sessions-router.ts b/apps/web/src/routers/active-sessions-router.ts index ac3471f88d..5266e59e4c 100644 --- a/apps/web/src/routers/active-sessions-router.ts +++ b/apps/web/src/routers/active-sessions-router.ts @@ -5,13 +5,19 @@ import { z } from 'zod'; import { SESSION_INGEST_WORKER_URL } from '@/lib/config.server'; import { generateInternalServiceToken } from '@/lib/tokens'; -const activeSessionSchema = z.object({ +export const activeSessionSchema = z.object({ id: z.string(), status: z.string(), title: z.string(), connectionId: z.string(), gitUrl: z.string().optional(), gitBranch: z.string().optional(), + /** + * Capabilities advertised by the CLI connection that owns this session. + * Omitted when the owning connection's latest heartbeat did not include a + * capabilities object (legacy CLI, or a CLI that predates the field). + */ + capabilities: z.object({ attachments: z.boolean().optional() }).optional(), // Optional: legacy CLIs (predating the `kilo remote` spawner) never // report a platform. Only present in the response when the CLI supplied it. platform: z.string().optional(), diff --git a/apps/web/src/routers/cloud-agent-next-router.test.ts b/apps/web/src/routers/cloud-agent-next-router.test.ts index 95e4a13c12..20ca8ed159 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -42,6 +42,14 @@ const mockGenerateCloudAgentAttachmentUploadUrl = jest.fn< }) => Promise<{ signedUrl: string; key: string; expiresAt: string }> >(() => Promise.resolve({ signedUrl: 'signed', key: 'key', expiresAt: 'expires' })); +const mockGenerateCloudAgentAttachmentDownloadUrl = jest.fn< + (input: { userId: string; messageUuid: string; filename: string }) => Promise<{ + signedUrl: string; + key: string; + expiresAt: string; + }> +>(() => Promise.resolve({ signedUrl: 'signed', key: 'key', expiresAt: 'expires' })); + const mockGetSession = jest.fn<(cloudAgentSessionId: string) => Promise<{ model?: string }>>(); const mockCreateCloudAgentNextClient = jest.fn(() => ({ @@ -125,6 +133,7 @@ jest.mock('@/lib/cloud-agent/gitlab-integration-helpers', () => ({ jest.mock('@/lib/r2/cloud-agent-attachments', () => ({ generateImageUploadUrl: jest.fn(), generateCloudAgentAttachmentUploadUrl: mockGenerateCloudAgentAttachmentUploadUrl, + generateCloudAgentAttachmentDownloadUrl: mockGenerateCloudAgentAttachmentDownloadUrl, })); jest.mock('@/lib/cloud-agent/session-ownership', () => ({ @@ -163,6 +172,7 @@ let createCaller: (ctx: { user: User }) => { contentType: 'application/pdf'; contentLength: number; }) => Promise; + getAttachmentDownloadUrl: (input: { messageUuid: string; filename: string }) => Promise; checkEligibility: () => Promise<{ balance: number; minBalance: number; @@ -359,6 +369,44 @@ describe('cloudAgentNextRouter attachment forwarding', () => { contentLength: 42, }); }); + + it('presigns Cloud Agent attachment downloads with the caller-scoped key prefix', async () => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + await caller.getAttachmentDownloadUrl({ + messageUuid: '12345678-1234-4234-9234-123456789abc', + filename: '87654321-4321-4321-8321-cba987654321.kilo', + }); + + expect(mockGenerateCloudAgentAttachmentDownloadUrl).toHaveBeenCalledWith({ + userId: 'user-1', + messageUuid: '12345678-1234-4234-9234-123456789abc', + filename: '87654321-4321-4321-8321-cba987654321.kilo', + }); + }); + + it('rejects a deny-listed extension before reaching the presign helper', async () => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + + await expect( + caller.getAttachmentDownloadUrl({ + messageUuid: '12345678-1234-4234-9234-123456789abc', + filename: '87654321-4321-4321-8321-cba987654321.exe', + }) + ).rejects.toThrow(); + expect(mockGenerateCloudAgentAttachmentDownloadUrl).not.toHaveBeenCalled(); + }); + + it('rejects a filename that violates the relaxed-regex shape', async () => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + + await expect( + caller.getAttachmentDownloadUrl({ + messageUuid: '12345678-1234-4234-9234-123456789abc', + filename: 'not-a-uuid.kilo', + }) + ).rejects.toThrow(); + expect(mockGenerateCloudAgentAttachmentDownloadUrl).not.toHaveBeenCalled(); + }); }); describe('cloudAgentNextRouter helper procedures', () => { diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index 1eb76cbc10..8c87dd9853 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -37,9 +37,11 @@ import { baseCloseTerminalNextOutputSchema, cloudAgentGetAttachmentUploadUrlSchema, cloudAgentGetImageUploadUrlSchema, + cloudAgentGetAttachmentDownloadUrlSchema, } from './cloud-agent-next-schemas'; import { generateCloudAgentAttachmentUploadUrl, + generateCloudAgentAttachmentDownloadUrl, generateImageUploadUrl, } from '@/lib/r2/cloud-agent-attachments'; import * as z from 'zod'; @@ -358,6 +360,22 @@ export const cloudAgentNextRouter = createTRPCRouter({ attachmentId: input.attachmentId, contentType: input.contentType, contentLength: input.contentLength, + ...(input.extension ? { extension: input.extension } : {}), + }); + }), + + /** + * Generate a presigned download URL for a stored Cloud Agent attachment. + * Personal scope only: the key prefix is derived from the caller (author). + * No org mirror — remote CLI sessions are personal. + */ + getAttachmentDownloadUrl: baseProcedure + .input(cloudAgentGetAttachmentDownloadUrlSchema) + .mutation(async ({ ctx, input }) => { + return generateCloudAgentAttachmentDownloadUrl({ + userId: ctx.user.id, + messageUuid: input.messageUuid, + filename: input.filename, }); }), diff --git a/apps/web/src/routers/cloud-agent-next-schemas.test.ts b/apps/web/src/routers/cloud-agent-next-schemas.test.ts index 27b5146317..62e1013447 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.test.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.test.ts @@ -1,148 +1,174 @@ import { describe, expect, it } from '@jest/globals'; import { - basePrepareSessionNextSchema, - baseSendMessageNextSchema, - personalPrepareSessionNextSchema, - cloudAgentAttachmentsSchema, + cloudAgentGetAttachmentDownloadUrlSchema, cloudAgentGetAttachmentUploadUrlSchema, + cloudAgentRelaxedAttachmentFilenameSchema, } from './cloud-agent-next-schemas'; const MESSAGE_UUID = '12345678-1234-4234-9234-123456789abc'; -const PDF_FILENAME = '87654321-4321-4321-8321-cba987654321.pdf'; -const TXT_FILENAME = '87654321-4321-4321-8321-cba987654321.txt'; -const MD_FILENAME = '87654321-4321-4321-8321-cba987654321.md'; -const CSV_FILENAME = '87654321-4321-4321-8321-cba987654321.csv'; +const ATTACHMENT_ID = '87654321-4321-4321-8321-cba987654321'; -describe('cloudAgentAttachmentsSchema', () => { - it('accepts document attachments and up to five ordered files', () => { - expect( - cloudAgentAttachmentsSchema.parse({ - path: MESSAGE_UUID, - files: [PDF_FILENAME, TXT_FILENAME, MD_FILENAME, CSV_FILENAME], - }) - ).toEqual({ - path: MESSAGE_UUID, - files: [PDF_FILENAME, TXT_FILENAME, MD_FILENAME, CSV_FILENAME], +describe('cloudAgentGetAttachmentUploadUrlSchema', () => { + it('preserves the legacy 9-MIME contract when extension is absent', () => { + const result = cloudAgentGetAttachmentUploadUrlSchema.safeParse({ + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + contentType: 'image/png', + contentLength: 1024, }); + expect(result.success).toBe(true); }); - it('rejects unsupported suffixes and more than five attachments', () => { - expect( - cloudAgentAttachmentsSchema.safeParse({ path: MESSAGE_UUID, files: ['file.docx'] }).success - ).toBe(false); - expect( - cloudAgentAttachmentsSchema.safeParse({ - path: MESSAGE_UUID, - files: Array.from({ length: 6 }, () => PDF_FILENAME), - }).success - ).toBe(false); + it('preserves the existing web-hook request shape (no extension field)', () => { + const result = cloudAgentGetAttachmentUploadUrlSchema.parse({ + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + contentType: 'text/markdown', + contentLength: 4096, + }); + expect(result.contentType).toBe('text/markdown'); + expect(result.extension).toBeUndefined(); }); -}); -describe('cloudAgentGetAttachmentUploadUrlSchema', () => { - it.each(['image/png', 'application/pdf', 'text/plain', 'text/markdown', 'text/csv'] as const)( - 'accepts %s document upload requests', - contentType => { - expect( - cloudAgentGetAttachmentUploadUrlSchema.parse({ - messageUuid: MESSAGE_UUID, - attachmentId: MESSAGE_UUID, - contentType, - contentLength: 5 * 1024 * 1024, - }).contentType - ).toBe(contentType); + it('accepts a relaxed contentType when extension is provided', () => { + const result = cloudAgentGetAttachmentUploadUrlSchema.safeParse({ + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + contentType: 'application/x-kilo-binary', + contentLength: 4096, + extension: 'kilo', + }); + expect(result.success).toBe(true); + }); + + it('rejects a malformed contentType even when extension is provided', () => { + const result = cloudAgentGetAttachmentUploadUrlSchema.safeParse({ + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + contentType: 'not a mime', + contentLength: 4096, + extension: 'kilo', + }); + expect(result.success).toBe(false); + }); + + it('rejects a contentType outside the legacy allow-list when extension is absent', () => { + const result = cloudAgentGetAttachmentUploadUrlSchema.safeParse({ + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + contentType: 'application/x-kilo-binary', + contentLength: 4096, + }); + expect(result.success).toBe(false); + }); + + it('rejects deny-listed extensions on the upload input', () => { + for (const extension of ['exe', 'dll', 'msi', 'com', 'scr', 'apk', 'ipa', 'dmg', 'pkg']) { + const result = cloudAgentGetAttachmentUploadUrlSchema.safeParse({ + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + contentType: 'application/octet-stream', + contentLength: 4096, + extension, + }); + expect(result.success).toBe(false); + if (!result.success) { + const extensionIssues = result.error.issues.filter(issue => issue.path[0] === 'extension'); + expect(extensionIssues[0]?.message).toContain(extension); + } } - ); + }); - it('rejects unsupported MIME and oversized uploads', () => { + it('rejects extensions that exceed the 16-character shape or include non-alphanumerics', () => { expect( cloudAgentGetAttachmentUploadUrlSchema.safeParse({ messageUuid: MESSAGE_UUID, - attachmentId: MESSAGE_UUID, - contentType: 'application/msword', - contentLength: 1, + attachmentId: ATTACHMENT_ID, + contentType: 'application/octet-stream', + contentLength: 4096, + extension: 'abcdefghijklmnopq', }).success ).toBe(false); expect( cloudAgentGetAttachmentUploadUrlSchema.safeParse({ messageUuid: MESSAGE_UUID, - attachmentId: MESSAGE_UUID, - contentType: 'application/pdf', - contentLength: 5 * 1024 * 1024 + 1, + attachmentId: ATTACHMENT_ID, + contentType: 'application/octet-stream', + contentLength: 4096, + extension: 'tar.gz', }).success ).toBe(false); }); -}); - -describe('basePrepareSessionNextSchema', () => { - it('preserves structured initial slash command payloads', () => { - const initialPayload = { - type: 'command' as const, - command: 'review', - arguments: 'main', - }; - const result = basePrepareSessionNextSchema.parse({ - githubRepo: 'kilocode/cloud', - prompt: '/review main', - mode: 'code', - model: 'anthropic/claude-sonnet', - initialPayload, + it('preserves the 5 MB positive contentLength cap even with an extension', () => { + const result = cloudAgentGetAttachmentUploadUrlSchema.safeParse({ + messageUuid: MESSAGE_UUID, + attachmentId: ATTACHMENT_ID, + contentType: 'application/octet-stream', + contentLength: 5 * 1024 * 1024 + 1, + extension: 'kilo', }); + expect(result.success).toBe(false); + }); +}); - expect(result.initialPayload).toEqual(initialPayload); +describe('cloudAgentRelaxedAttachmentFilenameSchema', () => { + it('accepts any 1-16 char alphanumeric extension after the UUID prefix', () => { + for (const filename of [ + `${ATTACHMENT_ID}.kilo`, + `${ATTACHMENT_ID}.docx`, + `${ATTACHMENT_ID}.tar`, + `${ATTACHMENT_ID}.a`, + `${ATTACHMENT_ID}.123`, + ]) { + expect(cloudAgentRelaxedAttachmentFilenameSchema.safeParse(filename).success).toBe(true); + } }); - it('accepts canonical attachment references and rejects ambiguous legacy input', () => { - const prompt = { - githubRepo: 'kilocode/cloud', - prompt: 'Review this file', - mode: 'code', - model: 'anthropic/claude-sonnet', - attachments: { path: MESSAGE_UUID, files: [MD_FILENAME] }, - }; + it('rejects filenames whose extension is in the deny-list', () => { + for (const extension of ['exe', 'dll', 'msi', 'com', 'scr', 'apk', 'ipa', 'dmg', 'pkg']) { + expect( + cloudAgentRelaxedAttachmentFilenameSchema.safeParse(`${ATTACHMENT_ID}.${extension}`).success + ).toBe(false); + } + }); - expect(basePrepareSessionNextSchema.parse(prompt).attachments).toEqual(prompt.attachments); + it('rejects filenames outside the UUID + 1-16 alphanumeric shape', () => { + expect(cloudAgentRelaxedAttachmentFilenameSchema.safeParse('not-a-uuid.kilo').success).toBe( + false + ); + expect(cloudAgentRelaxedAttachmentFilenameSchema.safeParse(`${ATTACHMENT_ID}`).success).toBe( + false + ); expect( - basePrepareSessionNextSchema.safeParse({ - ...prompt, - images: { path: MESSAGE_UUID, files: ['87654321-4321-4321-8321-cba987654321.png'] }, - }).success + cloudAgentRelaxedAttachmentFilenameSchema.safeParse(`${ATTACHMENT_ID}.abcdefghijklmnopq`) + .success ).toBe(false); }); }); -describe('personalPrepareSessionNextSchema', () => { - it('rejects Bitbucket repository identity in personal context', () => { - expect( - personalPrepareSessionNextSchema.safeParse({ - bitbucketRepo: { - fullName: 'acme/api', - workspaceUuid: '11111111-1111-4111-8111-111111111111', - repositoryUuid: '22222222-2222-4222-8222-222222222222', - }, - prompt: 'Inspect the repository', - mode: 'code', - model: 'anthropic/claude-sonnet', - }).success - ).toBe(false); +describe('cloudAgentGetAttachmentDownloadUrlSchema', () => { + it('accepts a relaxed UUID.filename pair', () => { + const result = cloudAgentGetAttachmentDownloadUrlSchema.safeParse({ + messageUuid: MESSAGE_UUID, + filename: `${ATTACHMENT_ID}.kilo`, + }); + expect(result.success).toBe(true); }); -}); -describe('baseSendMessageNextSchema', () => { - it('accepts canonical attachment references and rejects both attachment fields', () => { - const send = { - cloudAgentSessionId: 'agent_1', - payload: { type: 'prompt' as const, prompt: 'Summarize', mode: 'code', model: 'test' }, - attachments: { path: MESSAGE_UUID, files: [PDF_FILENAME] }, - }; + it('rejects a deny-listed extension on the download input', () => { + const result = cloudAgentGetAttachmentDownloadUrlSchema.safeParse({ + messageUuid: MESSAGE_UUID, + filename: `${ATTACHMENT_ID}.exe`, + }); + expect(result.success).toBe(false); + }); - expect(baseSendMessageNextSchema.parse(send).attachments).toEqual(send.attachments); - expect( - baseSendMessageNextSchema.safeParse({ - ...send, - images: { path: MESSAGE_UUID, files: ['87654321-4321-4321-8321-cba987654321.png'] }, - }).success - ).toBe(false); + it('rejects an unparseable filename', () => { + const result = cloudAgentGetAttachmentDownloadUrlSchema.safeParse({ + messageUuid: MESSAGE_UUID, + filename: 'not-a-uuid.exe', + }); + expect(result.success).toBe(false); }); }); diff --git a/apps/web/src/routers/cloud-agent-next-schemas.ts b/apps/web/src/routers/cloud-agent-next-schemas.ts index e7e4b24758..0f9b98b007 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.ts @@ -1,8 +1,13 @@ import * as z from 'zod'; import { CLOUD_AGENT_ATTACHMENT_ALLOWED_TYPES, + CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS, + CLOUD_AGENT_ATTACHMENT_EXTENSION_REGEX, + CLOUD_AGENT_ATTACHMENT_MAX_CONTENT_TYPE_LENGTH, CLOUD_AGENT_ATTACHMENT_MAX_COUNT, CLOUD_AGENT_ATTACHMENT_MAX_SIZE_BYTES, + CLOUD_AGENT_ATTACHMENT_RELAXED_CONTENT_TYPE_REGEX, + CLOUD_AGENT_ATTACHMENT_RELAXED_FILENAME_REGEX, CLOUD_AGENT_IMAGE_ALLOWED_TYPES, CLOUD_AGENT_IMAGE_MAX_COUNT, CLOUD_AGENT_IMAGE_MAX_SIZE_BYTES, @@ -40,16 +45,107 @@ const cloudAgentAttachmentFilenameSchema = z /^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.(?:png|jpg|jpeg|webp|gif|pdf|txt|md|csv)$/ ); +const DENIED_EXTENSION_SET = new Set(CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS); + +/** + * Shared relaxed-filename regex for the post-S2 world: a UUID with any + * `^[a-z0-9]{1,16}$` suffix. The deny-list is enforced separately by + * `cloudAgentRelaxedAttachmentFilenameSchema` so callers that need a + * deny-list-validated basename (download presign, attachment list) can use + * the same regex without re-implementing the check. + */ +const cloudAgentRelaxedAttachmentFilenameShapeSchema = z + .string() + .regex(CLOUD_AGENT_ATTACHMENT_RELAXED_FILENAME_REGEX); + +/** + * Filename schema that combines the relaxed regex with the deny-list check. + * Used by every layer that accepts an R2 object basename. + */ +export const cloudAgentRelaxedAttachmentFilenameSchema = z.string().superRefine((value, ctx) => { + const shape = cloudAgentRelaxedAttachmentFilenameShapeSchema.safeParse(value); + if (!shape.success) { + ctx.addIssue({ + code: 'custom', + message: + 'Attachment filename must be a UUID with a 1-16 character lowercase alphanumeric extension', + path: [], + }); + return; + } + const suffix = value.slice(value.lastIndexOf('.') + 1).toLowerCase(); + if (DENIED_EXTENSION_SET.has(suffix)) { + ctx.addIssue({ + code: 'custom', + message: `Attachment extension "${suffix}" is not allowed`, + path: [], + }); + } +}); + export const cloudAgentAttachmentsSchema = z.object({ path: z.uuid(), files: z.array(cloudAgentAttachmentFilenameSchema).min(1).max(CLOUD_AGENT_ATTACHMENT_MAX_COUNT), }); -export const cloudAgentGetAttachmentUploadUrlSchema = z.object({ +/** + * Upload-presign input. When `extension` is provided, the caller is opting + * into the relaxed surface: `contentType` only needs to match the + * `type/subtype` shape (≤128 chars) and the R2 key suffix is derived from the + * validated extension, not from `contentType`. When `extension` is absent, + * the legacy behavior is preserved: `contentType` must be one of the + * existing allowed MIMEs and the suffix derives from the legacy MIME map. + */ +export const cloudAgentGetAttachmentUploadUrlSchema = z + .object({ + messageUuid: z.uuid(), + attachmentId: z.uuid(), + contentType: z.string().min(1).max(CLOUD_AGENT_ATTACHMENT_MAX_CONTENT_TYPE_LENGTH), + contentLength: z.number().int().positive().max(CLOUD_AGENT_ATTACHMENT_MAX_SIZE_BYTES), + extension: z + .string() + .regex(CLOUD_AGENT_ATTACHMENT_EXTENSION_REGEX) + .superRefine((extension, ctx) => { + if (DENIED_EXTENSION_SET.has(extension.toLowerCase())) { + ctx.addIssue({ + code: 'custom', + message: `Attachment extension "${extension}" is not allowed`, + path: [], + }); + } + }) + .optional(), + }) + .superRefine((data, ctx) => { + if (data.extension) { + // Relaxed surface: contentType only needs the type/subtype shape. + if (!CLOUD_AGENT_ATTACHMENT_RELAXED_CONTENT_TYPE_REGEX.test(data.contentType)) { + ctx.addIssue({ + code: 'custom', + path: ['contentType'], + message: 'contentType must match type/subtype', + }); + } + } else if ( + !(CLOUD_AGENT_ATTACHMENT_ALLOWED_TYPES as readonly string[]).includes(data.contentType) + ) { + ctx.addIssue({ + code: 'custom', + path: ['contentType'], + message: `contentType must be one of: ${CLOUD_AGENT_ATTACHMENT_ALLOWED_TYPES.join(', ')}`, + }); + } + }); + +/** + * Download-presign input. The presign helper re-validates the filename with + * the same relaxed regex + deny-list combination so the key prefix can never + * accept an unsafe basename. There is intentionally no org mirror — remote + * CLI sessions are personal. + */ +export const cloudAgentGetAttachmentDownloadUrlSchema = z.object({ messageUuid: z.uuid(), - attachmentId: z.uuid(), - contentType: z.enum(CLOUD_AGENT_ATTACHMENT_ALLOWED_TYPES), - contentLength: z.number().int().positive().max(CLOUD_AGENT_ATTACHMENT_MAX_SIZE_BYTES), + filename: cloudAgentRelaxedAttachmentFilenameSchema, }); function hasOnlyOneAttachmentField(data: { images?: unknown; attachments?: unknown }): boolean { diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts index 54ebf0a13e..8d70a21dcf 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts @@ -505,6 +505,7 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ attachmentId: input.attachmentId, contentType: input.contentType, contentLength: input.contentLength, + ...(input.extension ? { extension: input.extension } : {}), }); }), diff --git a/services/cloud-agent-next/src/execution/attachment-prompt-parts.test.ts b/services/cloud-agent-next/src/execution/attachment-prompt-parts.test.ts index 9d3d780cd5..41063a5c19 100644 --- a/services/cloud-agent-next/src/execution/attachment-prompt-parts.test.ts +++ b/services/cloud-agent-next/src/execution/attachment-prompt-parts.test.ts @@ -1,5 +1,7 @@ import { describe, expect, it, vi } from 'vitest'; import { + PROMPT_MIME_BY_EXTENSION, + PROMPT_MIME_FALLBACK, assertR2AttachmentDownloadConfigured, buildSignedPromptAttachments, } from './attachment-prompt-parts.js'; @@ -24,18 +26,64 @@ const createEnv = (overrides: Partial = {}): Env => ...overrides, }) as Env; +/** + * Exact canonical extension → MIME table mirroring the worker + * `PROMPT_MIME_BY_EXTENSION`. Every entry must be covered by the test + * `buildSignedPromptAttachments` mapping below. + */ +const CANONICAL_TABLE: ReadonlyArray = [ + ['png', 'image/png'], + ['jpg', 'image/jpeg'], + ['jpeg', 'image/jpeg'], + ['webp', 'image/webp'], + ['gif', 'image/gif'], + ['pdf', 'application/pdf'], + ['txt', 'text/plain'], + ['md', 'text/plain'], + ['csv', 'text/plain'], + ['log', 'text/plain'], + ['json', 'text/plain'], + ['xml', 'text/plain'], + ['yaml', 'text/plain'], + ['yml', 'text/plain'], + ['toml', 'text/plain'], + ['ini', 'text/plain'], + ['html', 'text/plain'], + ['css', 'text/plain'], + ['js', 'text/plain'], + ['jsx', 'text/plain'], + ['ts', 'text/plain'], + ['tsx', 'text/plain'], + ['py', 'text/plain'], + ['rb', 'text/plain'], + ['go', 'text/plain'], + ['rs', 'text/plain'], + ['java', 'text/plain'], + ['c', 'text/plain'], + ['h', 'text/plain'], + ['cpp', 'text/plain'], + ['hpp', 'text/plain'], + ['sh', 'text/plain'], + ['sql', 'text/plain'], +]; + +describe('PROMPT_MIME_BY_EXTENSION canonical table', () => { + it('maps every documented extension to its declared MIME exactly once', () => { + for (const [extension, mime] of CANONICAL_TABLE) { + expect(PROMPT_MIME_BY_EXTENSION[extension]).toBe(mime); + } + }); + + it('resolves unknown / extensionless filenames to the octet-stream fallback', () => { + expect(PROMPT_MIME_BY_EXTENSION['zip']).toBeUndefined(); + expect(PROMPT_MIME_BY_EXTENSION['bin']).toBeUndefined(); + expect(PROMPT_MIME_BY_EXTENSION['']).toBeUndefined(); + expect(PROMPT_MIME_FALLBACK).toBe('application/octet-stream'); + }); +}); + describe('buildSignedPromptAttachments', () => { - it.each([ - ['png', 'image/png'], - ['jpg', 'image/jpeg'], - ['jpeg', 'image/jpeg'], - ['webp', 'image/webp'], - ['gif', 'image/gif'], - ['pdf', 'application/pdf'], - ['txt', 'text/plain'], - ['md', 'text/plain'], - ['csv', 'text/plain'], - ])('maps .%s wrapper attachments to %s for Kilo', async (suffix, mime) => { + it.each(CANONICAL_TABLE)('maps .%s wrapper attachments to %s for Kilo', async (suffix, mime) => { const attachments = { path: '00000000-0000-4000-8000-000000000000', files: [`11111111-1111-4111-8111-111111111111.${suffix}`], @@ -50,6 +98,27 @@ describe('buildSignedPromptAttachments', () => { expect(result[0]).toEqual(expect.objectContaining({ filename: attachments.files[0], mime })); }); + + it('resolves an extension not in the canonical table to application/octet-stream', async () => { + const attachments = { + path: '00000000-0000-4000-8000-000000000000', + files: ['22222222-2222-4222-8222-222222222222.zip'], + } satisfies Attachments; + + const result = await buildSignedPromptAttachments({ + env: createEnv(), + userId: 'user_test', + sessionId: 'agent_test', + attachments, + }); + + expect(result[0]).toEqual( + expect.objectContaining({ + filename: attachments.files[0], + mime: 'application/octet-stream', + }) + ); + }); }); describe('assertR2AttachmentDownloadConfigured', () => { diff --git a/services/cloud-agent-next/src/execution/attachment-prompt-parts.ts b/services/cloud-agent-next/src/execution/attachment-prompt-parts.ts index e8ef3c6251..df75afcfc7 100644 --- a/services/cloud-agent-next/src/execution/attachment-prompt-parts.ts +++ b/services/cloud-agent-next/src/execution/attachment-prompt-parts.ts @@ -15,18 +15,50 @@ type R2AttachmentDownloadEnv = { R2_ATTACHMENTS_BUCKET?: string; }; -const PROMPT_MIME_BY_SUFFIX: Record = { - '.png': 'image/png', - '.jpg': 'image/jpeg', - '.jpeg': 'image/jpeg', - '.webp': 'image/webp', - '.gif': 'image/gif', - '.pdf': 'application/pdf', - '.txt': 'text/plain', - '.md': 'text/plain', - '.csv': 'text/plain', +/** + * Canonical extension → MIME table (no leading dot). The stored extension is + * the single source of truth for the prompt MIME. Mirrors the web-side + * `CLOUD_AGENT_ATTACHMENT_EXTENSION_TO_MIME` constant; the wrapper + * `getPromptMime` resolves this exact table for the same filename suffix. + */ +export const PROMPT_MIME_BY_EXTENSION: Record = { + png: 'image/png', + jpg: 'image/jpeg', + jpeg: 'image/jpeg', + webp: 'image/webp', + gif: 'image/gif', + pdf: 'application/pdf', + txt: 'text/plain', + md: 'text/plain', + csv: 'text/plain', + log: 'text/plain', + json: 'text/plain', + xml: 'text/plain', + yaml: 'text/plain', + yml: 'text/plain', + toml: 'text/plain', + ini: 'text/plain', + html: 'text/plain', + css: 'text/plain', + js: 'text/plain', + jsx: 'text/plain', + ts: 'text/plain', + tsx: 'text/plain', + py: 'text/plain', + rb: 'text/plain', + go: 'text/plain', + rs: 'text/plain', + java: 'text/plain', + c: 'text/plain', + h: 'text/plain', + cpp: 'text/plain', + hpp: 'text/plain', + sh: 'text/plain', + sql: 'text/plain', }; +export const PROMPT_MIME_FALLBACK = 'application/octet-stream'; + export function assertR2AttachmentDownloadConfigured( env: T ): asserts env is T & { @@ -92,10 +124,9 @@ export async function buildSignedPromptAttachments({ function getPromptMime(filename: string): string { const dotIndex = filename.lastIndexOf('.'); - const suffix = dotIndex >= 0 ? filename.slice(dotIndex).toLowerCase() : ''; - const mime = PROMPT_MIME_BY_SUFFIX[suffix]; - if (!mime) { - throw new Error('Invalid attachment filename'); - } - return mime; + const extension = dotIndex >= 0 ? filename.slice(dotIndex + 1).toLowerCase() : ''; + // No MIME-based extension inference anywhere: an unknown / missing / + // invalid extension (including the legacy `bin` no-extension fallback) is + // always treated as `application/octet-stream`. + return PROMPT_MIME_BY_EXTENSION[extension] ?? PROMPT_MIME_FALLBACK; } diff --git a/services/cloud-agent-next/src/persistence/schemas.test.ts b/services/cloud-agent-next/src/persistence/schemas.test.ts index f33367c373..1576a5de7e 100644 --- a/services/cloud-agent-next/src/persistence/schemas.test.ts +++ b/services/cloud-agent-next/src/persistence/schemas.test.ts @@ -27,12 +27,15 @@ describe('AttachmentsSchema', () => { }); it('rejects unsupported document filename suffixes and more than five attachments', () => { + // `.docx` matches the relaxed `[a-z0-9]{1,16}` pattern. The deny-list + // covers the unsafe execution set; the storage layer now accepts any + // non-denied extension. expect( AttachmentsSchema.safeParse({ path: '123e4567-e89b-12d3-a456-426614174000', files: ['123e4567-e89b-12d3-a456-426614174001.docx'], }).success - ).toBe(false); + ).toBe(true); expect( AttachmentsSchema.safeParse({ path: '123e4567-e89b-12d3-a456-426614174000', @@ -47,6 +50,58 @@ describe('AttachmentsSchema', () => { }).success ).toBe(false); }); + + it('accepts any non-denied alphanumeric extension up to 16 characters', () => { + expect( + AttachmentsSchema.safeParse({ + path: '123e4567-e89b-12d3-a456-426614174000', + files: ['123e4567-e89b-12d3-a456-426614174001.docx'], + }).success + ).toBe(true); + expect( + AttachmentsSchema.safeParse({ + path: '123e4567-e89b-12d3-a456-426614174000', + files: ['123e4567-e89b-12d3-a456-426614174001.tar'], + }).success + ).toBe(true); + expect( + AttachmentsSchema.safeParse({ + path: '123e4567-e89b-12d3-a456-426614174000', + files: ['123e4567-e89b-12d3-a456-426614174001.kilocode'], + }).success + ).toBe(true); + }); + + it('rejects filenames whose extension is in the deny-list', () => { + for (const extension of ['exe', 'dll', 'msi', 'com', 'scr', 'apk', 'ipa', 'dmg', 'pkg']) { + const result = AttachmentsSchema.safeParse({ + path: '123e4567-e89b-12d3-a456-426614174000', + files: [`123e4567-e89b-12d3-a456-426614174001.${extension}`], + }); + expect(result.success).toBe(false); + } + }); + + it('rejects extensions longer than 16 characters and extensions with non-alphanumeric characters', () => { + expect( + AttachmentsSchema.safeParse({ + path: '123e4567-e89b-12d3-a456-426614174000', + files: ['123e4567-e89b-12d3-a456-426614174001.abcdefghijklmnopq'], + }).success + ).toBe(false); + expect( + AttachmentsSchema.safeParse({ + path: '123e4567-e89b-12d3-a456-426614174000', + files: ['123e4567-e89b-12d3-a456-426614174001.tar.gz'], + }).success + ).toBe(false); + expect( + AttachmentsSchema.safeParse({ + path: '123e4567-e89b-12d3-a456-426614174000', + files: ['123e4567-e89b-12d3-a456-426614174001.TAR'], + }).success + ).toBe(false); + }); }); describe('ImagesSchema', () => { diff --git a/services/cloud-agent-next/src/persistence/schemas.ts b/services/cloud-agent-next/src/persistence/schemas.ts index 8f96b0d579..1f489534ce 100644 --- a/services/cloud-agent-next/src/persistence/schemas.ts +++ b/services/cloud-agent-next/src/persistence/schemas.ts @@ -3,6 +3,26 @@ import { MESSAGE_ID_FORMAT_DESCRIPTION, MESSAGE_ID_PATTERN } from '../session/me import { BUILTIN_AGENT_MODES, Limits } from '../schema.js'; import type { SandboxId } from '../types.js'; +/** + * Attachment filename extension deny-list shared by every worker-layer + * validator (persistence schema, runtime download helper). Centralized here + * so the worker never ships a path that the storage layer would later + * reject. + */ +export const CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS = [ + 'exe', + 'dll', + 'msi', + 'com', + 'scr', + 'apk', + 'ipa', + 'dmg', + 'pkg', +] as const; + +const CLOUD_AGENT_DENIED_EXTENSION_SET = new Set(CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS); + /** * Schema for callback target configuration. * Defined here to avoid circular dependency with router/schemas.ts. @@ -21,14 +41,31 @@ const attachmentMessageUuidSchema = z .uuid() .describe('Bare message upload UUID; service prefix is derived by the worker'); +/** + * Filename regex shape for stored attachments. The relaxed post-S2 surface + * accepts any `^[a-z0-9]{1,16}$` suffix after the UUID prefix. The deny-list + * is enforced as a refinement. + */ +const ATTACHMENT_RELAXED_FILENAME_REGEX = + /^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.[a-z0-9]{1,16}$/; + const attachmentFilenameSchema = z .string() .min(1) .max(128) .regex( - /^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.(png|jpg|jpeg|webp|gif|pdf|txt|md|csv)$/, - 'Attachment filename must be a UUID with extension png, jpg, jpeg, webp, gif, pdf, txt, md, or csv' - ); + ATTACHMENT_RELAXED_FILENAME_REGEX, + 'Attachment filename must be a UUID with a 1-16 character lowercase alphanumeric extension' + ) + .superRefine((filename, ctx) => { + const suffix = filename.slice(filename.lastIndexOf('.') + 1).toLowerCase(); + if (CLOUD_AGENT_DENIED_EXTENSION_SET.has(suffix)) { + ctx.addIssue({ + code: 'custom', + message: `Attachment extension "${suffix}" is not allowed`, + }); + } + }); export const AttachmentsSchema = z.object({ path: attachmentMessageUuidSchema, diff --git a/services/cloud-agent-next/src/router/schemas.ts b/services/cloud-agent-next/src/router/schemas.ts index a6f5c14f31..00de4c12ea 100644 --- a/services/cloud-agent-next/src/router/schemas.ts +++ b/services/cloud-agent-next/src/router/schemas.ts @@ -15,6 +15,7 @@ import { EncryptedSecretsSchema, CallbackTargetSchema, AttachmentsSchema, + CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS, ImagesSchema, RuntimeSkillSchema, RuntimeSkillsSchema, @@ -36,6 +37,7 @@ export { EncryptedSecretsSchema, CallbackTargetSchema, AttachmentsSchema, + CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS, ImagesSchema, RuntimeSkillSchema, RuntimeSkillsSchema, diff --git a/services/cloud-agent-next/src/utils/attachment-download.test.ts b/services/cloud-agent-next/src/utils/attachment-download.test.ts index 07262b01fe..e30fa91a35 100644 --- a/services/cloud-agent-next/src/utils/attachment-download.test.ts +++ b/services/cloud-agent-next/src/utils/attachment-download.test.ts @@ -70,7 +70,7 @@ describe('buildPresignedAttachments', () => { ).rejects.toThrow('Invalid attachment message UUID'); }); - it('rejects filenames outside the server-permitted suffix allowlist', async () => { + it('rejects filenames with non-alphanumeric extensions outside the relaxed shape', async () => { await expect( buildPresignedAttachments( createR2Client(), @@ -80,9 +80,56 @@ describe('buildPresignedAttachments', () => { 'cloud-agent', { path: '123e4567-e89b-12d3-a456-426614174000', - files: ['123e4567-e89b-12d3-a456-426614174001.svg'], + files: ['123e4567-e89b-12d3-a456-426614174001.docx.tar'], } ) ).rejects.toThrow('Invalid attachment filename'); }); + + it('accepts any non-denied 1-16 char alphanumeric extension and signs the matching R2 key', async () => { + const r2Client = createR2Client(); + const messageUuid = '123e4567-e89b-12d3-a456-426614174000'; + const files = [ + '123e4567-e89b-12d3-a456-426614174001.docx', + '123e4567-e89b-12d3-a456-426614174002.kilo', + '123e4567-e89b-12d3-a456-426614174003.py', + ]; + + const result = await buildPresignedAttachments( + r2Client, + 'attachments', + 'session', + 'user-123', + 'cloud-agent', + { path: messageUuid, files } + ); + + expect(r2Client.getSignedURL).toHaveBeenCalledWith( + 'attachments', + `user-123/cloud-agent/${messageUuid}/${files[0]}` + ); + expect(result.map(entry => entry.filename)).toEqual(files); + }); + + it('rejects filenames whose extension is in the deny-list at download time', async () => { + for (const extension of ['exe', 'dll', 'msi', 'com', 'scr', 'apk', 'ipa', 'dmg', 'pkg']) { + // The Attachments type is already constrained by AttachmentsSchema at + // the boundary, so we cast through `unknown` to reach the runtime + // validator inside `buildPresignedAttachments`. + const malformed = { + path: '123e4567-e89b-12d3-a456-426614174000', + files: [`123e4567-e89b-12d3-a456-426614174001.${extension}`], + } as unknown as Parameters[5]; + await expect( + buildPresignedAttachments( + createR2Client(), + 'attachments', + 'session', + 'user-123', + 'cloud-agent', + malformed + ) + ).rejects.toThrow(`Attachment extension "${extension}" is not allowed`); + } + }); }); diff --git a/services/cloud-agent-next/src/utils/attachment-download.ts b/services/cloud-agent-next/src/utils/attachment-download.ts index 0a4717e393..242985d3bb 100644 --- a/services/cloud-agent-next/src/utils/attachment-download.ts +++ b/services/cloud-agent-next/src/utils/attachment-download.ts @@ -1,5 +1,5 @@ import type { R2Client } from '@kilocode/worker-utils'; -import type { Attachments } from '../router/schemas.js'; +import { CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS, type Attachments } from '../router/schemas.js'; export type AttachmentService = 'app-builder' | 'cloud-agent'; @@ -11,8 +11,16 @@ export type PresignedAttachment = { const MESSAGE_UUID_REGEX = /^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$/; -const ATTACHMENT_FILENAME_REGEX = - /^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.(png|jpg|jpeg|webp|gif|pdf|txt|md|csv)$/; + +/** + * Relaxed post-S2 filename regex: UUID + any 1-16 character lowercase + * alphanumeric extension. The deny-list is enforced as a second pass in + * `validateAttachments`. + */ +const ATTACHMENT_RELAXED_FILENAME_REGEX = + /^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.[a-z0-9]{1,16}$/; + +const DENIED_EXTENSION_SET = new Set(CLOUD_AGENT_ATTACHMENT_DENIED_EXTENSIONS); export function deriveAttachmentService(createdOnPlatform?: string): AttachmentService { return createdOnPlatform === 'app-builder' ? 'app-builder' : 'cloud-agent'; @@ -32,9 +40,13 @@ function validateAttachments(attachments: Attachments): void { } for (const filename of attachments.files) { - if (!ATTACHMENT_FILENAME_REGEX.test(filename)) { + if (!ATTACHMENT_RELAXED_FILENAME_REGEX.test(filename)) { throw new Error('Invalid attachment filename'); } + const suffix = filename.slice(filename.lastIndexOf('.') + 1).toLowerCase(); + if (DENIED_EXTENSION_SET.has(suffix)) { + throw new Error(`Attachment extension "${suffix}" is not allowed`); + } } } diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts index 8dd64e7cb9..1173dd8f6c 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts @@ -63,6 +63,26 @@ function asFetch( return Object.assign(fn, { preconnect: fetch.preconnect }); } +function makeByteStream(totalBytes: number): ReadableStream { + const chunk = new Uint8Array(64 * 1024); + return new ReadableStream({ + start(controller) { + let pushed = 0; + while (pushed < totalBytes) { + const remaining = totalBytes - pushed; + if (remaining >= chunk.byteLength) { + controller.enqueue(chunk); + pushed += chunk.byteLength; + } else { + controller.enqueue(chunk.subarray(0, remaining)); + pushed += remaining; + } + } + controller.close(); + }, + }); +} + async function createCompleteGitWorkspace(workspacePath: string): Promise { const gitPath = path.join(workspacePath, '.git'); await fsp.mkdir(gitPath, { recursive: true }); @@ -1324,10 +1344,6 @@ describe('prepareWrapperBootstrapWorkspace', () => { const result = await materializePromptAttachments(prompt, { fetch: asFetch(async () => new Response('image-bytes', { status: 200 })), - writeResponse: async (filePath, response) => { - await fsp.writeFile(filePath, await response.text()); - return 11; - }, }); expect(result.message.parts).toEqual([ @@ -1339,6 +1355,7 @@ describe('prepareWrapperBootstrapWorkspace', () => { filename: 'diagram.png', }, ]); + expect(await fsp.readFile(path.join(tmpDir, 'diagram.png'), 'utf8')).toBe('image-bytes'); expect(result.message.attachments).toBeUndefined(); }); @@ -1368,11 +1385,6 @@ describe('prepareWrapperBootstrapWorkspace', () => { const result = await materializePromptAttachments(prompt, { fetch: asFetch(async () => new Response('pdf-bytes', { status: 200 })), - writeResponse: async (filePath, response) => { - const bytes = await response.text(); - await fsp.writeFile(filePath, bytes); - return bytes.length; - }, }); expect(result.message.parts).toEqual([ @@ -1384,6 +1396,7 @@ describe('prepareWrapperBootstrapWorkspace', () => { filename: 'spec.pdf', }, ]); + expect(await fsp.readFile(pdfPath, 'utf8')).toBe('pdf-bytes'); }); it.each([ @@ -1415,11 +1428,6 @@ describe('prepareWrapperBootstrapWorkspace', () => { const result = await materializePromptAttachments(prompt, { fetch: asFetch(async () => new Response(content, { status: 200 })), - writeResponse: async (filePath, response) => { - const bytes = await response.text(); - await fsp.writeFile(filePath, bytes); - return bytes.length; - }, }); expect(result.message.parts).toContainEqual({ @@ -1428,20 +1436,20 @@ describe('prepareWrapperBootstrapWorkspace', () => { url: `file://${localPath}`, filename, }); + expect(await fsp.readFile(localPath, 'utf8')).toBe(content); }); - it('rejects attachments with an oversized content-length before writing', async () => { - const localPath = path.join(tmpDir, 'too-large.pdf'); - const writeResponse = mock(async () => 0); + it('replaces an oversized attachment with an explanatory text part and deletes the partial file', async () => { + const localPath = path.join(tmpDir, 'too-large.bin'); const prompt: WrapperPromptRequest = { message: { - id: 'msg_header_limit', - prompt: 'Read this PDF', + id: 'msg_overflow', + prompt: 'Process this binary', attachments: [ { - filename: 'too-large.pdf', - mime: 'application/pdf', - signedUrl: 'https://r2.example.com/too-large.pdf', + filename: 'too-large.bin', + mime: 'application/octet-stream', + signedUrl: 'https://r2.example.com/too-large.bin', localPath, }, ], @@ -1455,24 +1463,316 @@ describe('prepareWrapperBootstrapWorkspace', () => { }, }; - let error: Error | undefined; - try { - await materializePromptAttachments(prompt, { - fetch: asFetch( - async () => - new Response('not-written', { - status: 200, - headers: { 'content-length': '5242881' }, - }) - ), - writeResponse, - }); - } catch (caught) { - error = caught instanceof Error ? caught : new Error(String(caught)); - } + // Stream 6 MiB of bytes (one more than the 5 MiB + 1 cap) so the + // bounded reader triggers the overflow branch deterministically. + const total = 6 * 1024 * 1024; + const chunk = new Uint8Array(64 * 1024); + const body = new ReadableStream({ + start(controller) { + let pushed = 0; + while (pushed < total) { + const remaining = total - pushed; + if (remaining >= chunk.byteLength) { + controller.enqueue(chunk); + pushed += chunk.byteLength; + } else { + controller.enqueue(chunk.subarray(0, remaining)); + pushed += remaining; + } + } + controller.close(); + }, + }); + + const result = await materializePromptAttachments(prompt, { + fetch: asFetch(async () => new Response(body, { status: 200 })), + }); + + expect(result.message.parts).toEqual([ + { type: 'text', text: 'Process this binary' }, + { + type: 'text', + text: 'attachment too-large.bin could not be retrieved (Attachment too large: bytes exceeded the 5 MiB cap)', + }, + ]); + expect(fs.existsSync(localPath)).toBe(false); + }); + + it('materializes a file exactly at the 5 MiB cap as a binary text part', async () => { + const localPath = path.join(tmpDir, 'exactly-5-mib.bin'); + const prompt: WrapperPromptRequest = { + message: { + id: 'msg_exact', + prompt: 'Process this file', + attachments: [ + { + filename: 'exactly-5-mib.bin', + mime: 'application/octet-stream', + signedUrl: 'https://r2.example.com/exactly-5-mib.bin', + localPath, + }, + ], + }, + session: { + ingestUrl: 'wss://worker.example.com/sessions/user/agent/ingest', + workerAuthToken: 'token', + wrapperRunId: 'wr_test', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_test', + }, + }; + + const result = await materializePromptAttachments(prompt, { + fetch: asFetch(async () => new Response(makeByteStream(5 * 1024 * 1024), { status: 200 })), + }); + + expect(result.message.parts).toEqual([ + { type: 'text', text: 'Process this file' }, + { + type: 'text', + text: `binary attachment saved: filename=exactly-5-mib.bin mime=application/octet-stream size=${5 * 1024 * 1024} path=${localPath}`, + }, + ]); + expect(fs.existsSync(localPath)).toBe(true); + }); + + it('rejects a file one byte over the 5 MiB cap and deletes the partial file', async () => { + const localPath = path.join(tmpDir, 'one-byte-over.bin'); + const prompt: WrapperPromptRequest = { + message: { + id: 'msg_over', + prompt: 'Process this file', + attachments: [ + { + filename: 'one-byte-over.bin', + mime: 'application/octet-stream', + signedUrl: 'https://r2.example.com/one-byte-over.bin', + localPath, + }, + ], + }, + session: { + ingestUrl: 'wss://worker.example.com/sessions/user/agent/ingest', + workerAuthToken: 'token', + wrapperRunId: 'wr_test', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_test', + }, + }; + + const result = await materializePromptAttachments(prompt, { + fetch: asFetch( + async () => new Response(makeByteStream(5 * 1024 * 1024 + 1), { status: 200 }) + ), + }); + + expect(result.message.parts).toEqual([ + { type: 'text', text: 'Process this file' }, + { + type: 'text', + text: 'attachment one-byte-over.bin could not be retrieved (Attachment too large: bytes exceeded the 5 MiB cap)', + }, + ]); + expect(fs.existsSync(localPath)).toBe(false); + }); + + it('deletes the partial file when a mid-transfer read fails', async () => { + const localPath = path.join(tmpDir, 'interrupted.bin'); + const prompt: WrapperPromptRequest = { + message: { + id: 'msg_interrupted', + prompt: 'Process this file', + attachments: [ + { + filename: 'interrupted.bin', + mime: 'application/octet-stream', + signedUrl: 'https://r2.example.com/interrupted.bin', + localPath, + }, + ], + }, + session: { + ingestUrl: 'wss://worker.example.com/sessions/user/agent/ingest', + workerAuthToken: 'token', + wrapperRunId: 'wr_test', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_test', + }, + }; + + const chunk = new Uint8Array(64 * 1024); + const body = new ReadableStream({ + pull(controller) { + controller.enqueue(chunk); + controller.error(new Error('stream reset')); + }, + }); + + const result = await materializePromptAttachments(prompt, { + fetch: asFetch(async () => new Response(body, { status: 200 })), + }); - expect(error?.message).toBe('Attachment too large: too-large.pdf'); - expect(writeResponse).not.toHaveBeenCalled(); + expect(result.message.parts).toEqual([ + { type: 'text', text: 'Process this file' }, + { + type: 'text', + text: 'attachment interrupted.bin could not be retrieved (stream reset)', + }, + ]); expect(fs.existsSync(localPath)).toBe(false); }); + + it('continues with text-only parts when a non-2xx response aborts one attachment', async () => { + const okPath = path.join(tmpDir, 'ok.png'); + const failPath = path.join(tmpDir, 'missing.png'); + const prompt: WrapperPromptRequest = { + message: { + id: 'msg_403', + prompt: 'Show me both', + attachments: [ + { + filename: 'missing.png', + mime: 'image/png', + signedUrl: 'https://r2.example.com/missing.png', + localPath: failPath, + }, + { + filename: 'ok.png', + mime: 'image/png', + signedUrl: 'https://r2.example.com/ok.png', + localPath: okPath, + }, + ], + }, + session: { + ingestUrl: 'wss://worker.example.com/sessions/user/agent/ingest', + workerAuthToken: 'token', + wrapperRunId: 'wr_test', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_test', + }, + }; + + const result = await materializePromptAttachments(prompt, { + fetch: asFetch(async input => { + const url = typeof input === 'string' ? input : (input as Request).url; + if (url === 'https://r2.example.com/missing.png') { + return new Response('forbidden', { status: 403 }); + } + return new Response('png-bytes', { status: 200 }); + }), + }); + + expect(result.message.parts).toEqual([ + { type: 'text', text: 'Show me both' }, + { + type: 'text', + text: 'attachment missing.png could not be retrieved (HTTP 403)', + }, + { + type: 'file', + mime: 'image/png', + url: `file://${okPath}`, + filename: 'ok.png', + }, + ]); + expect(fs.existsSync(failPath)).toBe(false); + expect(await fsp.readFile(okPath, 'utf8')).toBe('png-bytes'); + }); + + it('converts a network/timeout failure into an explanatory text part and continues', async () => { + const okPath = path.join(tmpDir, 'ok.json'); + const failPath = path.join(tmpDir, 'bad.json'); + const prompt: WrapperPromptRequest = { + message: { + id: 'msg_network', + prompt: 'Read both', + attachments: [ + { + filename: 'bad.json', + mime: 'text/plain', + signedUrl: 'https://r2.example.com/bad.json', + localPath: failPath, + }, + { + filename: 'ok.json', + mime: 'text/plain', + signedUrl: 'https://r2.example.com/ok.json', + localPath: okPath, + }, + ], + }, + session: { + ingestUrl: 'wss://worker.example.com/sessions/user/agent/ingest', + workerAuthToken: 'token', + wrapperRunId: 'wr_test', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_test', + }, + }; + + const result = await materializePromptAttachments(prompt, { + fetch: asFetch(async input => { + const url = typeof input === 'string' ? input : (input as Request).url; + if (url === 'https://r2.example.com/bad.json') { + throw new Error('socket hang up'); + } + return new Response('{"ok":true}', { status: 200 }); + }), + }); + + expect(result.message.parts).toEqual([ + { type: 'text', text: 'Read both' }, + { + type: 'text', + text: 'attachment bad.json could not be retrieved (socket hang up)', + }, + { + type: 'file', + mime: 'text/plain', + url: `file://${okPath}`, + filename: 'ok.json', + }, + ]); + expect(fs.existsSync(failPath)).toBe(false); + expect(await fsp.readFile(okPath, 'utf8')).toBe('{"ok":true}'); + }); + + it('materializes a generic binary attachment as a text part describing the saved file', async () => { + const localPath = path.join(tmpDir, 'payload.zip'); + const prompt: WrapperPromptRequest = { + message: { + id: 'msg_zip', + prompt: 'Extract the archive', + attachments: [ + { + filename: 'payload.zip', + mime: 'application/octet-stream', + signedUrl: 'https://r2.example.com/payload.zip', + localPath, + }, + ], + }, + session: { + ingestUrl: 'wss://worker.example.com/sessions/user/agent/ingest', + workerAuthToken: 'token', + wrapperRunId: 'wr_test', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_test', + }, + }; + + const result = await materializePromptAttachments(prompt, { + fetch: asFetch(async () => new Response('zip-payload', { status: 200 })), + }); + + expect(result.message.parts).toEqual([ + { type: 'text', text: 'Extract the archive' }, + { + type: 'text', + text: `binary attachment saved: filename=payload.zip mime=application/octet-stream size=11 path=${localPath}`, + }, + ]); + expect(await fsp.readFile(localPath, 'utf8')).toBe('zip-payload'); + }); }); diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts index 4391ddb01c..815f197bcc 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts @@ -37,6 +37,7 @@ const SETUP_COMMAND_ERROR_OUTPUT_MAX_BYTES = 4_096; const SETUP_COMMAND_DIAGNOSTIC_MAX_BYTES = 1_024; const GIT_BOOTSTRAP_MARKER = 'kilo-bootstrap-complete'; const MAX_ATTACHMENT_BYTES = 5_242_880; +const MAX_ATTACHMENT_DOWNLOAD_BYTES = MAX_ATTACHMENT_BYTES + 1; function cleanTerminalOutput(text: string): string { return stripAnsi(text) @@ -66,6 +67,16 @@ function boundedUtf8Tail(text: string, maxBytes: number): string { .replace(/^\uFFFD/, ''); } +/** + * True for MIME classes that the prompt must surface as a `file://` part. Any + * other class (generic binary) is materialized to disk and exposed to the + * agent as an explanatory text part with the absolute path, filename, MIME, + * and size. + */ +function isPromptFileMime(mime: string): boolean { + return mime.startsWith('text/') || mime.startsWith('image/') || mime === 'application/pdf'; +} + function createSetupOutputReporter( progress: BootstrapProgress | undefined, stepId: string @@ -807,53 +818,188 @@ function createSetupCommandDiagnostic(command: string, result: ExecResult): stri return parts.join(', '); } -async function downloadAttachment( +async function safeUnlink(filePath: string): Promise { + try { + await fs.unlink(filePath); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'ENOENT') { + logToFile(`attachment: failed to remove partial file ${filePath}: ${String(error)}`); + } + } +} + +/** + * Bounded streaming read. We never trust the server's `content-length` header + * alone: the response body is pulled at most `MAX_ATTACHMENT_DOWNLOAD_BYTES` + * bytes. If the producer keeps producing after the cap, the read is aborted, + * the partial file is removed, and a typed overflow error is thrown. + */ +async function downloadBounded( + filePath: string, + response: Response +): Promise<{ bytesWritten: number }> { + const body = response.body; + if (!body) { + throw new Error('Attachment download failed: empty body'); + } + + const handle = await fs.open(filePath, 'w'); + let bytesWritten = 0; + let overflowed = false; + try { + const reader = body.getReader(); + while (true) { + const { value, done } = await reader.read(); + if (done) break; + if (!value) continue; + const remaining = MAX_ATTACHMENT_DOWNLOAD_BYTES - bytesWritten; + if (value.byteLength >= remaining) { + const writable = MAX_ATTACHMENT_BYTES - bytesWritten; + await handle.write(value.subarray(0, writable)); + bytesWritten += writable; + overflowed = true; + try { + await reader.cancel(); + } catch { + // Ignore: we're tearing the connection down anyway. + } + break; + } + await handle.write(value); + bytesWritten += value.byteLength; + } + } catch (error) { + await safeUnlink(filePath); + throw error; + } finally { + await handle.close(); + } + + if (overflowed) { + await safeUnlink(filePath); + throw new Error('Attachment too large: bytes exceeded the 5 MiB cap'); + } + + return { bytesWritten }; +} + +export type DownloadResult = + | { kind: 'ok'; part: WrapperPromptPart; bytesWritten: number } + | { kind: 'failed'; part: WrapperPromptPart }; + +/** + * Download a single attachment. Per-file failure (non-2xx response, + * read/timeout error, overflow) is converted to an explanatory text part so + * the rest of the prompt can still proceed; the whole-message abort path is + * reserved for non-attachment failures. + */ +async function downloadAndMaterializeAttachment( attachment: WrapperBootstrapAttachment, fetchImpl: typeof fetch, - writeResponse: (filePath: string, response: Response) => Promise -): Promise { + abortController: AbortController +): Promise { await fs.mkdir(path.dirname(attachment.localPath), { recursive: true }); - const response = await fetchImpl(attachment.signedUrl, { - signal: AbortSignal.timeout(120_000), - }); + + let response: Response; + try { + response = await fetchImpl(attachment.signedUrl, { + signal: abortController.signal, + }); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + return { + kind: 'failed', + part: { + type: 'text', + text: `attachment ${attachment.filename} could not be retrieved (${message})`, + }, + }; + } + if (!response.ok) { - throw new Error(`Attachment download failed: ${attachment.filename}`); + return { + kind: 'failed', + part: { + type: 'text', + text: `attachment ${attachment.filename} could not be retrieved (HTTP ${response.status})`, + }, + }; + } + + let result: { bytesWritten: number }; + try { + result = await downloadBounded(attachment.localPath, response); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + return { + kind: 'failed', + part: { + type: 'text', + text: `attachment ${attachment.filename} could not be retrieved (${message})`, + }, + }; } - const contentLength = response.headers.get('content-length'); - if (contentLength && Number(contentLength) > MAX_ATTACHMENT_BYTES) { - throw new Error(`Attachment too large: ${attachment.filename}`); + + if (isPromptFileMime(attachment.mime)) { + return { + kind: 'ok', + part: { + type: 'file', + mime: attachment.mime, + url: `file://${attachment.localPath}`, + filename: attachment.filename, + }, + bytesWritten: result.bytesWritten, + }; } - await writeResponse(attachment.localPath, response); + + // Generic binary: surface as a text part with the absolute path so the + // agent can use shell tools against the file instead of receiving a + // `file://` part the Kilo SDK may not handle. return { - type: 'file', - mime: attachment.mime, - url: `file://${attachment.localPath}`, - filename: attachment.filename, + kind: 'ok', + part: { + type: 'text', + text: `binary attachment saved: filename=${attachment.filename} mime=${attachment.mime} size=${result.bytesWritten} path=${attachment.localPath}`, + }, + bytesWritten: result.bytesWritten, }; } +export type MaterializeDeps = { + fetch?: typeof fetch; +}; + export async function materializePromptAttachments( prompt: WrapperPromptRequest, - deps: { - fetch?: typeof fetch; - writeResponse?: (filePath: string, response: Response) => Promise; - } = {} + deps: MaterializeDeps = {} ): Promise { if (!prompt.message.attachments?.length) return prompt; const fetchImpl = deps.fetch ?? fetch; - const writeResponse = - deps.writeResponse ?? ((filePath: string, response: Response) => Bun.write(filePath, response)); - const fileParts: WrapperPromptPart[] = []; + + const parts: WrapperPromptPart[] = []; for (const attachment of prompt.message.attachments) { - fileParts.push(await downloadAttachment(attachment, fetchImpl, writeResponse)); + const abortController = new AbortController(); + const timeout = setTimeout( + () => abortController.abort(new Error('attachment download timeout')), + 120_000 + ); + let result: DownloadResult; + try { + result = await downloadAndMaterializeAttachment(attachment, fetchImpl, abortController); + } finally { + clearTimeout(timeout); + } + parts.push(result.part); } + return { ...prompt, message: { ...prompt.message, parts: [ ...(prompt.message.parts ?? [{ type: 'text', text: prompt.message.prompt ?? '' }]), - ...fileParts, + ...parts, ], prompt: undefined, attachments: undefined, diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 06b365e0f5..0333a8d87b 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -231,14 +231,18 @@ function sendHeartbeat( title: string; platform?: string; }>, - protocolVersion?: string, - instance?: { name: string; projectName: string; version?: string } + options: { + protocolVersion?: string; + capabilities?: { attachments?: boolean }; + instance?: { name: string; projectName: string; version?: string }; + } = {} ) { const msg = JSON.stringify({ type: 'heartbeat', sessions, - ...(protocolVersion ? { protocolVersion } : {}), - ...(instance ? { instance } : {}), + ...(options.protocolVersion ? { protocolVersion: options.protocolVersion } : {}), + ...(options.capabilities ? { capabilities: options.capabilities } : {}), + ...(options.instance ? { instance: options.instance } : {}), }); doInstance.webSocketMessage(cliWs as never, msg); } @@ -528,11 +532,11 @@ describe('UserConnectionDO', () => { const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], '1'); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { protocolVersion: '1' }); sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], '1'); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { protocolVersion: '1' }); expect(parseSent(webWs)).toMatchObject({ type: 'system', @@ -605,6 +609,179 @@ describe('UserConnectionDO', () => { const msgs = allSent(cliWs); expect(msgs).toContainEqual({ type: 'heartbeat_ack' }); }); + + // ----------------------------------------------------------------------- + // capabilities transitions (decision 8): exercise true→absent, true→false, + // and absent/false→true through the actual DO event path and assert the + // projected value in BOTH aggregateSessions() and the sessions.heartbeat + // event rows, including omission when the latest heartbeat omits + // capabilities. + // ----------------------------------------------------------------------- + + it('projects capabilities.attachments=true on every aggregateSessions row when the owning CLI advertises it', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { + capabilities: { attachments: true }, + }); + + const rows = doInstance.getActiveSessions(); + expect(rows).toHaveLength(2); + for (const row of rows) { + expect(row.capabilities).toEqual({ attachments: true }); + } + }); + + it('projects capabilities.attachments=true on every session row of the sessions.heartbeat event', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + // Establish ownership and subscription first so the heartbeat broadcast + // is targeted at this web socket. + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); + sendSubscribe(doInstance, webWs, 's1'); + webWs.send.mockClear(); + + // Now advertise capabilities: the broadcast heartbeat must carry the + // capabilities on the message envelope, but each session row only + // carries its own owning-connection identity (capabilities is read from + // the connection, not the per-session row of the broadcast). The S3b + // SDK consumer re-attaches capabilities per-row in the consumer layer. + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { + capabilities: { attachments: true }, + }); + + const sent = parseSent(webWs) as { data: { capabilities?: unknown; sessions: unknown[] } }; + expect(sent.data.capabilities).toEqual({ attachments: true }); + expect(sent.data.sessions).toHaveLength(2); + + // aggregateSessions() must still surface the latest capabilities, which + // is the S3b consumer's source of truth for the per-row projection. + const rows = doInstance.getActiveSessions(); + for (const row of rows) { + expect(row.capabilities).toEqual({ attachments: true }); + } + }); + + it('omits capabilities from aggregateSessions rows when the latest heartbeat omits the field (legacy CLI)', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + + // First heartbeat with capabilities + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + capabilities: { attachments: true }, + }); + // Second heartbeat omits capabilities (CLI rollback / legacy) + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + + const rows = doInstance.getActiveSessions(); + expect(rows).toHaveLength(1); + expect(rows[0]).not.toHaveProperty('capabilities'); + }); + + it('omits capabilities from sessions.heartbeat event envelope when the latest heartbeat omits the field', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + capabilities: { attachments: true }, + }); + sendSubscribe(doInstance, webWs, 's1'); + webWs.send.mockClear(); + + // Latest heartbeat omits capabilities. + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + + const sent = parseSent(webWs) as { data: Record }; + expect(sent.data).not.toHaveProperty('capabilities'); + }); + + it('flips capabilities.attachments from true to false on the next heartbeat (CLI revocation)', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + capabilities: { attachments: true }, + }); + expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ attachments: true }); + + // CLI advertises attachments=false (e.g. feature gated, profile change) + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + capabilities: { attachments: false }, + }); + + const rows = doInstance.getActiveSessions(); + expect(rows).toHaveLength(1); + expect(rows[0].capabilities).toEqual({ attachments: false }); + }); + + it('flips capabilities.attachments from absent to true when a legacy CLI starts advertising it', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + + // Legacy heartbeat — no capabilities field + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + const legacy = doInstance.getActiveSessions(); + expect(legacy[0]).not.toHaveProperty('capabilities'); + + // Upgraded CLI starts advertising attachments=true + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + capabilities: { attachments: true }, + }); + + const upgraded = doInstance.getActiveSessions(); + expect(upgraded[0].capabilities).toEqual({ attachments: true }); + }); + + it('flips capabilities.attachments from false to true on the next heartbeat', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + capabilities: { attachments: false }, + }); + expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ attachments: false }); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + capabilities: { attachments: true }, + }); + + expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ attachments: true }); + }); + + it('projects the same owning-connection capabilities on every session row of a multi-session heartbeat', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { + capabilities: { attachments: false }, + }); + + const rows = doInstance.getActiveSessions(); + expect(rows).toHaveLength(2); + expect(rows[0].capabilities).toEqual({ attachments: false }); + expect(rows[1].capabilities).toEqual({ attachments: false }); + }); + + it('reconstructs capabilities from a hibernated CLI attachment', () => { + const { doInstance, mockCtx } = setup(); + // Pre-existing attachment with capabilities — simulates a socket that + // was accepted before the DO was evicted. + const cliWs = createMockWs(['cli'], { + role: 'cli', + connectionId: 'cli-hiber', + sessions: [makeSession('s1')], + capabilities: { attachments: true }, + }); + mockCtx.addSocket(cliWs); + + const rows = doInstance.getActiveSessions(); + expect(rows).toHaveLength(1); + expect(rows[0].capabilities).toEqual({ attachments: true }); + }); }); // ------------------------------------------------------------------------- @@ -2916,7 +3093,9 @@ describe('UserConnectionDO', () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'busy', 'Fix bug')], '1'); + sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'busy', 'Fix bug')], { + protocolVersion: '1', + }); const result = doInstance.getActiveSessions(); expect(result).toEqual([ @@ -3082,10 +3261,9 @@ describe('UserConnectionDO', () => { it('persists `instance` in the WS attachment across heartbeats', () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [], '1', { - name: 'laptop-1', - projectName: 'kilo', - version: '0.1.0', + sendHeartbeat(doInstance, cliWs, [], { + protocolVersion: '1', + instance: { name: 'laptop-1', projectName: 'kilo', version: '0.1.0' }, }); const att = cliWs.deserializeAttachment() as { instance?: { name: string } }; @@ -3096,9 +3274,8 @@ describe('UserConnectionDO', () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); // First heartbeat: with instance. - sendHeartbeat(doInstance, cliWs, [], undefined, { - name: 'laptop-1', - projectName: 'kilo', + sendHeartbeat(doInstance, cliWs, [], { + instance: { name: 'laptop-1', projectName: 'kilo' }, }); // Second heartbeat: instance removed (legacy fallback). The DO must not // keep a stale `instance` value in the attachment. diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 05132eacf5..3596642c0e 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -27,6 +27,8 @@ type HeartbeatSession = { platform?: string; }; +type ConnectionCapabilities = { attachments?: boolean }; + type WSAttachment = | { role: 'cli'; @@ -36,6 +38,11 @@ type WSAttachment = // CLI hasn't sent its first heartbeat, or it's a legacy build that // predates this field entirely. Both cases fall back to legacy behavior. protocolVersion?: string; + // Latest capabilities advertised by this connection. Undefined means + // either the CLI hasn't sent its first heartbeat, or it's a legacy + // build that predates the capabilities field — both surface as no + // opt-in features. + capabilities?: ConnectionCapabilities; // Set from the authenticated /user/cli route; undefined on sockets // accepted before this field existed. Needed for the session-ready push. kiloUserId?: string; @@ -162,6 +169,8 @@ export class UserConnectionDO extends DurableObject { private connectionSessions = new Map(); // Protocol version per CLI connection (from heartbeat); absent = legacy CLI private connectionProtocolVersion = new Map(); + // Capabilities per CLI connection (from heartbeat); absent = legacy CLI + private connectionCapabilities = new Map(); // Pending command responses: correlationId → originating web socket private pendingCommands = new Map< string, @@ -194,9 +203,10 @@ export class UserConnectionDO extends DurableObject { if (attachment.role === 'cli') { cliCount++; - const { connectionId, sessions, protocolVersion } = attachment; + const { connectionId, sessions, protocolVersion, capabilities } = attachment; this.connectionSessions.set(connectionId, sessions); this.connectionProtocolVersion.set(connectionId, protocolVersion); + this.connectionCapabilities.set(connectionId, capabilities); sessionCount += sessions.length; for (const session of sessions) { this.sessionOwners.set(session.id, connectionId); @@ -398,7 +408,14 @@ export class UserConnectionDO extends DurableObject { switch (msg.type) { case 'heartbeat': - this.handleHeartbeat(ws, attachment, msg.sessions, msg.protocolVersion, msg.instance); + this.handleHeartbeat( + ws, + attachment, + msg.sessions, + msg.protocolVersion, + msg.capabilities, + msg.instance + ); break; case 'event': this.handleCliEvent(msg.sessionId, msg.parentSessionId, msg.event, msg.data); @@ -414,12 +431,14 @@ export class UserConnectionDO extends DurableObject { attachment: WSAttachment & { role: 'cli' }, sessions: HeartbeatSession[], protocolVersion: string | undefined, + capabilities: ConnectionCapabilities | undefined, instance: Instance | undefined ): void { const { connectionId } = attachment; const now = Date.now(); this.lastHeartbeatAt.set(connectionId, now); this.connectionProtocolVersion.set(connectionId, protocolVersion); + this.connectionCapabilities.set(connectionId, capabilities); this.scheduleNextAlarm(now); // Remove sessions this connection previously owned but no longer reports @@ -462,6 +481,7 @@ export class UserConnectionDO extends DurableObject { connectionId, sessions, protocolVersion, + capabilities, kiloUserId: attachment.kiloUserId, ...(instance ? { instance } : {}), }; @@ -484,7 +504,7 @@ export class UserConnectionDO extends DurableObject { const msg: WebInboundMessage = { type: 'system', event: 'sessions.heartbeat', - data: { connectionId, protocolVersion, sessions }, + data: { connectionId, protocolVersion, capabilities, sessions }, }; for (const ws2 of subscribers) { this.sendToWeb(ws2, msg); @@ -851,6 +871,7 @@ export class UserConnectionDO extends DurableObject { } this.connectionSessions.delete(connectionId); this.connectionProtocolVersion.delete(connectionId); + this.connectionCapabilities.delete(connectionId); this.lastHeartbeatAt.delete(connectionId); console.log('CLI socket disconnected', { @@ -911,7 +932,11 @@ export class UserConnectionDO extends DurableObject { // --------------------------------------------------------------------------- getActiveSessions(): Array< - HeartbeatSession & { connectionId: string; protocolVersion?: string } + HeartbeatSession & { + connectionId: string; + protocolVersion?: string; + capabilities?: ConnectionCapabilities; + } > { this.ensureState(); return this.aggregateSessions(); @@ -1123,7 +1148,11 @@ export class UserConnectionDO extends DurableObject { } private aggregateSessions(): Array< - HeartbeatSession & { connectionId: string; protocolVersion?: string } + HeartbeatSession & { + connectionId: string; + protocolVersion?: string; + capabilities?: ConnectionCapabilities; + } > { // Build set of connectionIds that still have a live CLI WebSocket. // This guards against stale entries that persist if a close event is delayed. @@ -1133,16 +1162,24 @@ export class UserConnectionDO extends DurableObject { if (att?.role === 'cli') liveConnectionIds.add(att.connectionId); } - const result: Array = []; + const result: Array< + HeartbeatSession & { + connectionId: string; + protocolVersion?: string; + capabilities?: ConnectionCapabilities; + } + > = []; for (const [connectionId, sessions] of this.connectionSessions) { if (!liveConnectionIds.has(connectionId)) continue; const protocolVersion = this.connectionProtocolVersion.get(connectionId); + const capabilities = this.connectionCapabilities.get(connectionId); for (const session of sessions) { if (session.parentSessionId) continue; result.push({ ...session, connectionId, ...(protocolVersion ? { protocolVersion } : {}), + ...(capabilities ? { capabilities } : {}), // Preserve byte-identical responses for legacy senders that never // include a `platform`: only forward the field when present. ...(session.platform ? { platform: session.platform } : {}), diff --git a/services/session-ingest/src/types/user-connection-protocol.test.ts b/services/session-ingest/src/types/user-connection-protocol.test.ts index 5e92c6b548..eb2f02dfcf 100644 --- a/services/session-ingest/src/types/user-connection-protocol.test.ts +++ b/services/session-ingest/src/types/user-connection-protocol.test.ts @@ -197,6 +197,69 @@ describe('CLIOutboundMessageSchema', () => { }); }); +describe('CLIOutboundMessageSchema capabilities', () => { + const baseSession = { id: 'ses_cap_1', status: 'busy', title: 'cap' }; + + it('accepts capabilities.attachments: true on a heartbeat', () => { + const msg = { + type: 'heartbeat', + protocolVersion: '1', + capabilities: { attachments: true }, + sessions: [baseSession], + }; + const result = CLIOutboundMessageSchema.safeParse(msg); + expect(result.success).toBe(true); + if (result.success && result.data.type === 'heartbeat') { + expect(result.data.capabilities).toEqual({ attachments: true }); + } + }); + + it('accepts capabilities.attachments: false on a heartbeat', () => { + const msg = { + type: 'heartbeat', + capabilities: { attachments: false }, + sessions: [baseSession], + }; + const result = CLIOutboundMessageSchema.safeParse(msg); + expect(result.success).toBe(true); + if (result.success && result.data.type === 'heartbeat') { + expect(result.data.capabilities).toEqual({ attachments: false }); + } + }); + + it('accepts an absent capabilities field on a heartbeat (legacy CLI)', () => { + const msg = { type: 'heartbeat', sessions: [baseSession] }; + const result = CLIOutboundMessageSchema.safeParse(msg); + expect(result.success).toBe(true); + if (result.success && result.data.type === 'heartbeat') { + expect(result.data.capabilities).toBeUndefined(); + } + }); + + it('accepts an empty capabilities object on a heartbeat', () => { + const msg = { + type: 'heartbeat', + capabilities: {}, + sessions: [baseSession], + }; + const result = CLIOutboundMessageSchema.safeParse(msg); + expect(result.success).toBe(true); + if (result.success && result.data.type === 'heartbeat') { + expect(result.data.capabilities).toEqual({}); + } + }); + + it('rejects a non-boolean capabilities.attachments value', () => { + const msg = { + type: 'heartbeat', + capabilities: { attachments: 'yes' }, + sessions: [baseSession], + }; + const result = CLIOutboundMessageSchema.safeParse(msg); + expect(result.success).toBe(false); + }); +}); + describe('CLIInboundMessageSchema', () => { it('parses valid subscribe', () => { const msg = { type: 'subscribe', sessionId: validSessionId }; diff --git a/services/session-ingest/src/types/user-connection-protocol.ts b/services/session-ingest/src/types/user-connection-protocol.ts index 05d5058a13..a7156b2e13 100644 --- a/services/session-ingest/src/types/user-connection-protocol.ts +++ b/services/session-ingest/src/types/user-connection-protocol.ts @@ -24,6 +24,10 @@ export const CLIOutboundMessageSchema = z.discriminatedUnion('type', [ // Absent on CLI builds older than the protocolVersion field itself — treat // a missing value as a legacy CLI with no negotiated wire protocol. protocolVersion: z.string().optional(), + // Per-connection capabilities advertised by the CLI. Absent on CLIs that + // predate the field — treated as a legacy CLI with no opt-in features + // (e.g. attachment uploads from the mobile viewer). + capabilities: z.object({ attachments: z.boolean().optional() }).optional(), // Optional identity of the spawning CLI process. Absent on legacy CLIs // (which are not spawned by `kilo remote`). When present, the DO // persists it in the WebSocket attachment and exposes it via