diff --git a/frontend/src/components/ContextPicker.tsx b/frontend/src/components/ContextPicker.tsx index f26bb36..2ab11f7 100644 --- a/frontend/src/components/ContextPicker.tsx +++ b/frontend/src/components/ContextPicker.tsx @@ -19,8 +19,8 @@ function formatSize(bytes: number): string { if (bytes === 0) return '0 B'; if (bytes < 1024) return `${bytes} B`; const kb = bytes / 1024; - if (kb < 1024) return `${kb.toFixed(1)} KB`; - return `${(kb / 1024).toFixed(1)} MB`; + if (kb < 1024) return `${parseFloat(kb.toFixed(1))} KB`; + return `${parseFloat((kb / 1024).toFixed(1))} MB`; } export function ContextPicker({ selected, onToggle, onClose }: Props) { @@ -31,10 +31,13 @@ export function ContextPicker({ selected, onToggle, onClose }: Props) { // Fetch available context blocks from /api/config useEffect(() => { fetch('/api/config', { credentials: 'include' }) - .then((r) => r.json()) + .then((r) => { + if (!r.ok) throw new Error(`HTTP ${r.status}`); + return r.json(); + }) .then((data: { contextBlocks?: Record }) => { const entries: ContextBlockEntry[] = []; - if (data.contextBlocks) { + if (data.contextBlocks && typeof data.contextBlocks === 'object') { for (const [name, info] of Object.entries(data.contextBlocks)) { entries.push({ name, path: info.path, sizeBytes: info.sizeBytes }); } diff --git a/frontend/src/pages/ChatView.tsx b/frontend/src/pages/ChatView.tsx index 959dc54..76a2570 100644 --- a/frontend/src/pages/ChatView.tsx +++ b/frontend/src/pages/ChatView.tsx @@ -167,21 +167,18 @@ export function ChatView() { // Stop TTS playback when user sends a new message voice.stopSpeaking(); - const payload = buildSendPayload(text, images); - if (contextBlocks?.length) payload.contextBlocks = contextBlocks; + const payload = { + ...buildSendPayload(text, images), + ...(contextBlocks?.length ? { contextBlocks } : {}), + }; const previews = images?.map((img) => img.preview); - if (msgState.running) { - // Server queues it natively — no client-side stop+re-send needed. - wsSend(poolKey, payload); - dispatch({ type: 'USER_SEND', text, images: previews, contextBlocks }); - forceScrollToBottom(); - } else { + if (!msgState.running) { wsSetRunning(poolKey, true); - wsSend(poolKey, payload); - dispatch({ type: 'USER_SEND', text, images: previews, contextBlocks }); - forceScrollToBottom(); } + wsSend(poolKey, payload); + dispatch({ type: 'USER_SEND', text, images: previews, contextBlocks }); + forceScrollToBottom(); return true; } @@ -193,15 +190,12 @@ export function ChatView() { ): void { if (!wsIsOpen(poolKey) || !msgState.running) return; const imagePayload = images?.map((img) => ({ data: img.data, mediaType: img.mediaType })); - const previews = images?.map((img) => img.preview); wsSend(poolKey, { type: 'interrupt', prompt: text, images: imagePayload, ...(contextBlocks?.length ? { contextBlocks } : {}), }); - dispatch({ type: 'USER_SEND', text, images: previews, contextBlocks }); - forceScrollToBottom(); } const handleStop = useCallback(() => { diff --git a/frontend/src/types/chat.ts b/frontend/src/types/chat.ts index a219e23..63e8ea9 100644 --- a/frontend/src/types/chat.ts +++ b/frontend/src/types/chat.ts @@ -43,12 +43,15 @@ export interface FinishedBlock { toolError?: boolean; } +/** Serialized context block content (e.g. XML strings injected by the backend). */ +export type ContextBlockContent = string; + export interface FinishedMessage { messageId: string; role: 'user' | 'assistant'; blocks: FinishedBlock[]; images?: string[]; - contextBlocks?: string[]; + contextBlocks?: ContextBlockContent[]; } // --- Legacy flat Message type (used for restore/session history only) --- diff --git a/server/__tests__/assemble-prompt.test.ts b/server/__tests__/assemble-prompt.test.ts index d34de18..b35e6b1 100644 --- a/server/__tests__/assemble-prompt.test.ts +++ b/server/__tests__/assemble-prompt.test.ts @@ -1,8 +1,9 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { join } from 'path'; -import { mkdirSync, writeFileSync, rmSync } from 'fs'; +import { mkdirSync, mkdtempSync, writeFileSync, rmSync } from 'fs'; +import { tmpdir } from 'os'; -const TMP_DIR = join(import.meta.dirname, '..', '..', '.test-assemble-prompt'); +let TMP_DIR: string; // Mock repo-config to provide contextBlocks const mockContextBlocks: Record = {}; @@ -26,7 +27,7 @@ vi.mock('../repo-config.js', () => ({ const { assemblePrompt } = await import('../chat.js'); beforeEach(() => { - mkdirSync(TMP_DIR, { recursive: true }); + TMP_DIR = mkdtempSync(join(tmpdir(), 'assemble-prompt-')); // Clear mock context blocks for (const key of Object.keys(mockContextBlocks)) delete mockContextBlocks[key]; }); @@ -125,6 +126,45 @@ describe('assemblePrompt — context blocks', () => { expect(result).toContain('safe content'); }); + it('does not escape XML-unsafe characters in file content', () => { + const filePath = join(TMP_DIR, 'html-content.md'); + const unsafeContent = ' & more
'; + writeFileSync(filePath, unsafeContent); + mockContextBlocks['HtmlContent'] = filePath; + + const result = assemblePrompt('Q', TMP_DIR, undefined, ['HtmlContent']); + + // Content should be preserved verbatim (not escaped) + expect(result).toContain(unsafeContent); + expect(result).toContain(' { + const filePath = join(TMP_DIR, 'empty.md'); + writeFileSync(filePath, ''); + mockContextBlocks['Empty'] = filePath; + + const result = assemblePrompt('Q', TMP_DIR, undefined, ['Empty']); + + // Empty file should still produce a context tag (or be skipped) + // Verify the prompt still contains the user message + expect(result).toContain('Q'); + }); + + it('handles duplicate names in contextBlocks array', () => { + const filePath = join(TMP_DIR, 'workflow.md'); + writeFileSync(filePath, 'Workflow content'); + mockContextBlocks['Workflow'] = filePath; + + const result = assemblePrompt('Q', TMP_DIR, undefined, ['Workflow', 'Workflow']); + + // Count occurrences of the context block + const matches = result.match(/ { const filePath = join(TMP_DIR, 'large.md'); // Create a file just over 100 KB @@ -159,10 +199,12 @@ describe('assemblePrompt — context blocks', () => { writeFileSync(filePath, 'Workflow content'); mockContextBlocks['Workflow'] = filePath; - const images = [{ data: 'dGVzdA==', mediaType: 'image/png' }]; + const images: Array<{ data: string; mediaType: string }> = [ + { data: 'dGVzdA==', mediaType: 'image/png' }, + ]; const result = assemblePrompt('Describe', TMP_DIR, images, ['Workflow']); - // Should have both context blocks and image references + // Shouldhave both context blocks and image references expect(result).toContain('; return typeof obj.label === 'string' && typeof obj.desc === 'string'; } +function isStringRecord(val: unknown): val is Record { + return !!val && typeof val === 'object' && !Array.isArray(val); +} + export function loadRepoConfig(repoPath: string): RepoConfig { if (!repoPath) return { ...EMPTY_CONFIG }; @@ -99,11 +105,7 @@ export function loadRepoConfig(repoPath: string): RepoConfig { const validTiers = new Set(['safe', 'standard', 'elevated']); const toolTierOverrides: Record = {}; - if ( - obj.toolTierOverrides && - typeof obj.toolTierOverrides === 'object' && - !Array.isArray(obj.toolTierOverrides) - ) { + if (isStringRecord(obj.toolTierOverrides)) { for (const [tool, tier] of Object.entries(obj.toolTierOverrides as Record)) { if (typeof tier === 'string' && validTiers.has(tier)) { toolTierOverrides[tool] = tier as ToolTierOverride; @@ -115,7 +117,7 @@ export function loadRepoConfig(repoPath: string): RepoConfig { const resolvedInboxPath = inboxPath ? join(repoPath, inboxPath) : ''; const repos: Record = {}; - if (obj.repos && typeof obj.repos === 'object' && !Array.isArray(obj.repos)) { + if (isStringRecord(obj.repos)) { for (const [name, path] of Object.entries(obj.repos as Record)) { if (typeof path !== 'string') continue; if (!existsSync(path)) { @@ -130,15 +132,21 @@ export function loadRepoConfig(repoPath: string): RepoConfig { } } + const resolvedRepoPath = resolve(repoPath); const contextBlocks: Record = {}; - if ( - obj.contextBlocks && - typeof obj.contextBlocks === 'object' && - !Array.isArray(obj.contextBlocks) - ) { + if (isStringRecord(obj.contextBlocks)) { for (const [name, path] of Object.entries(obj.contextBlocks as Record)) { if (typeof path !== 'string') continue; - contextBlocks[name] = path.startsWith('/') ? path : join(repoPath, path); + if (!SAFE_NAME_RE.test(name)) { + log.warn(`contextBlocks: skipping invalid name: ${name}`); + continue; + } + const resolved = path.startsWith('/') ? path : resolve(repoPath, path); + if (!resolved.startsWith(resolvedRepoPath + '/') && resolved !== resolvedRepoPath) { + log.warn(`contextBlocks.${name}: path escapes repo root: ${path}`); + continue; + } + contextBlocks[name] = resolved; } } diff --git a/server/ws-schemas.ts b/server/ws-schemas.ts index 93099e8..5e00cf1 100644 --- a/server/ws-schemas.ts +++ b/server/ws-schemas.ts @@ -5,6 +5,8 @@ const ImageSchema = z.object({ mediaType: z.string(), }); +const ContextBlocksSchema = z.array(z.string().max(100_000)).max(20).optional(); + export const ReattachMessage = z.object({ type: z.literal('reattach'), clientId: z.string(), @@ -21,14 +23,14 @@ export const SendMessage = z.object({ extraTools: z.string().optional(), worktree: z.boolean().optional(), images: z.array(ImageSchema).optional(), - contextBlocks: z.array(z.string()).optional(), + contextBlocks: ContextBlocksSchema, }); export const InterruptMessage = z.object({ type: z.literal('interrupt'), prompt: z.string().min(1), images: z.array(ImageSchema).optional(), - contextBlocks: z.array(z.string()).optional(), + contextBlocks: ContextBlocksSchema, }); export const StopMessage = z.object({