diff --git a/packages/core/src/runtime/history.test.ts b/packages/core/src/runtime/history.test.ts index 820d3902..27996aa3 100644 --- a/packages/core/src/runtime/history.test.ts +++ b/packages/core/src/runtime/history.test.ts @@ -79,6 +79,43 @@ describe('JsonlHistorySink', () => { await expect(sink.flush()).resolves.toBeUndefined() }) + test('serializes concurrent append calls without losing lines', async () => { + const sink = new JsonlHistorySink(filePath) + const events = Array.from({ length: 20 }, (_, idx) => + createHistoryEvent({ + sessionId: 'concurrent', + type: 'action', + step: idx, + content: `event-${idx}`, + }), + ) + + await Promise.all(events.map((event) => sink.append(event))) + await sink.flush() + + const lines = (await readFile(filePath, 'utf8')).trim().split('\n').filter(Boolean) + expect(lines).toHaveLength(20) + const parsed = lines.map((line) => JSON.parse(line) as { content?: string }) + expect(parsed[0]?.content).toBe('event-0') + expect(parsed[19]?.content).toBe('event-19') + }) + + test('close flushes pending writes and rejects future appends', async () => { + const sink = new JsonlHistorySink(filePath) + const first = sink.append(createHistoryEvent({ sessionId: 's-close', type: 'turn_start' })) + const second = sink.append( + createHistoryEvent({ sessionId: 's-close', type: 'assistant', content: 'done' }), + ) + await sink.close() + await Promise.all([first, second]) + + const lines = (await readFile(filePath, 'utf8')).trim().split('\n').filter(Boolean) + expect(lines).toHaveLength(2) + await expect( + sink.append(createHistoryEvent({ sessionId: 's-close', type: 'final' })), + ).rejects.toThrow('History sink is closed') + }) + test('handles special characters in content', async () => { const sink = new JsonlHistorySink(filePath) const event = createHistoryEvent({ diff --git a/packages/core/src/runtime/history.ts b/packages/core/src/runtime/history.ts index 0657b5c6..d9b848ac 100644 --- a/packages/core/src/runtime/history.ts +++ b/packages/core/src/runtime/history.ts @@ -5,21 +5,40 @@ import type { HistoryEvent, HistorySink, Role } from '@memo/core/types' /** JSONL history writer: one event per line. */ export class JsonlHistorySink implements HistorySink { - private ready = false + private ensureDirPromise: Promise | null = null + private writeQueue: Promise = Promise.resolve() + private closed = false constructor(private filePath: string) {} + private ensureDirectory() { + if (!this.ensureDirPromise) { + this.ensureDirPromise = mkdir(dirname(this.filePath), { recursive: true }).then( + () => {}, + ) + } + return this.ensureDirPromise + } + async append(event: HistoryEvent) { - if (!this.ready) { - await mkdir(dirname(this.filePath), { recursive: true }) - this.ready = true + if (this.closed) { + throw new Error('History sink is closed') } - await appendFile(this.filePath, `${JSON.stringify(event)}\n`, 'utf8') + this.writeQueue = this.writeQueue.then(async () => { + await this.ensureDirectory() + await appendFile(this.filePath, `${JSON.stringify(event)}\n`, 'utf8') + }) + return this.writeQueue } async flush() { - // appendFile already ensures persistence, here only for interface compatibility - return Promise.resolve() + await this.writeQueue + } + + async close() { + if (this.closed) return + this.closed = true + await this.flush() } } diff --git a/packages/core/src/runtime/session_runtime.ts b/packages/core/src/runtime/session_runtime.ts index 3c3fd8d9..fdc6a754 100644 --- a/packages/core/src/runtime/session_runtime.ts +++ b/packages/core/src/runtime/session_runtime.ts @@ -925,12 +925,14 @@ export class AgentSessionImpl implements AgentSession { }, }) for (const sink of this.sinks) { - if (sink.flush) { - try { + try { + if (sink.close) { + await sink.close() + } else if (sink.flush) { await sink.flush() - } catch (err) { - console.error(`History flush failed: ${(err as Error).message}`) } + } catch (err) { + console.error(`History flush failed: ${(err as Error).message}`) } } } diff --git a/packages/core/src/runtime/session_runtime_helpers.test.ts b/packages/core/src/runtime/session_runtime_helpers.test.ts new file mode 100644 index 00000000..aaa56cc0 --- /dev/null +++ b/packages/core/src/runtime/session_runtime_helpers.test.ts @@ -0,0 +1,79 @@ +import { describe, expect, test, vi } from 'vitest' +import type { HistorySink } from '@memo/core/types' +import { + accumulateUsage, + emitEventToSinks, + emptyUsage, + stableStringify, +} from '@memo/core/runtime/session_runtime_helpers' + +describe('accumulateUsage', () => { + test('uses explicit total when provided', () => { + const usage = emptyUsage() + accumulateUsage(usage, { prompt: 2, completion: 3, total: 100 }) + expect(usage).toEqual({ prompt: 2, completion: 3, total: 100 }) + }) + + test('falls back to prompt + completion when total is absent', () => { + const usage = emptyUsage() + accumulateUsage(usage, { prompt: 2, completion: 3 }) + expect(usage).toEqual({ prompt: 2, completion: 3, total: 5 }) + }) +}) + +describe('emitEventToSinks', () => { + test('writes structured error payload to stderr when sink append fails', async () => { + const writes: string[] = [] + const writeSpy = vi.spyOn(process.stderr, 'write').mockImplementation(((chunk: unknown) => { + writes.push(String(chunk)) + return true + }) as typeof process.stderr.write) + + const failingSink: HistorySink = { + append: async () => { + throw new Error('disk full') + }, + } + + try { + await emitEventToSinks( + { + ts: '2026-01-01T00:00:00.000Z', + sessionId: 's-1', + type: 'assistant', + content: 'hello', + }, + [failingSink], + ) + } finally { + writeSpy.mockRestore() + } + + expect(writes.length).toBeGreaterThan(0) + const parsed = JSON.parse(writes.join('').trim()) as Record + expect(parsed.level).toBe('error') + expect(parsed.event).toBe('history_sink_append_failed') + expect(parsed.message).toBe('disk full') + expect(parsed.sink).toBe('Object') + }) +}) + +describe('stableStringify', () => { + test('serializes self-referencing object without throwing', () => { + const root: Record = {} + root.self = root + + const serialized = stableStringify(root) + expect(serialized).toBe('{"self":"[Circular]"}') + }) + + test('serializes indirect circular references with circular marker', () => { + const parent: Record = { name: 'parent' } + const child: Record = { name: 'child', parent } + parent.child = child + + const serialized = stableStringify(parent) + expect(serialized).toContain('"child":{"name":"child","parent":"[Circular]"}') + expect(serialized).toContain('"name":"parent"') + }) +}) diff --git a/packages/core/src/runtime/session_runtime_helpers.ts b/packages/core/src/runtime/session_runtime_helpers.ts index fa7167cb..396b0579 100644 --- a/packages/core/src/runtime/session_runtime_helpers.ts +++ b/packages/core/src/runtime/session_runtime_helpers.ts @@ -38,6 +38,10 @@ export type ResolvedToolPermission = { approvalMode: 'auto' | 'strict' } +function writeStructuredError(payload: Record) { + process.stderr.write(`${JSON.stringify(payload)}\n`) +} + export function resolveToolPermission(options: AgentSessionOptions): ResolvedToolPermission { if (options.toolPermissionMode === 'none') { return { @@ -122,7 +126,12 @@ export async function emitEventToSinks(event: HistoryEvent, sinks: HistorySink[] try { await sink.append(event) } catch (err) { - console.error(`Failed to write history event: ${(err as Error).message}`) + writeStructuredError({ + level: 'error', + event: 'history_sink_append_failed', + sink: sink.constructor?.name || 'anonymous_sink', + message: (err as Error).message, + }) } } } @@ -133,14 +142,37 @@ export function isAbortError(err: unknown): err is Error { // Stable serialization for duplicate action detection (ensures consistent key ordering) export function stableStringify(value: unknown): string { + return stableStringifyWithSeen(value, new WeakSet(), 0) +} + +const MAX_STABLE_STRINGIFY_DEPTH = 100 + +function stableStringifyWithSeen(value: unknown, seen: WeakSet, depth: number): string { + if (depth > MAX_STABLE_STRINGIFY_DEPTH) { + return JSON.stringify('[MaxDepthExceeded]') + } + if (typeof value === 'bigint') { + return JSON.stringify(value.toString()) + } if (value === null || typeof value !== 'object') return JSON.stringify(value) ?? 'null' + if (seen.has(value)) { + return JSON.stringify('[Circular]') + } + + seen.add(value) if (Array.isArray(value)) { - return `[${value.map((v) => stableStringify(v)).join(',')}]` + const result = `[${value.map((v) => stableStringifyWithSeen(v, seen, depth + 1)).join(',')}]` + seen.delete(value) + return result } const entries = Object.entries(value as Record).sort(([a], [b]) => a.localeCompare(b), ) - return `{${entries.map(([k, v]) => `${JSON.stringify(k)}:${stableStringify(v)}`).join(',')}}` + const result = `{${entries + .map(([k, v]) => `${JSON.stringify(k)}:${stableStringifyWithSeen(v, seen, depth + 1)}`) + .join(',')}}` + seen.delete(value) + return result } export function buildAssistantToolCalls( diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index 7cb0c1bc..7ac5942a 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -371,4 +371,6 @@ export type HistorySink = { append: (event: HistoryEvent) => Promise | void /** 可选:flush 持久化。 */ flush?: () => Promise | void + /** 可选:关闭资源并确保落盘。 */ + close?: () => Promise | void } diff --git a/packages/tools/src/orchestrator/index.test.ts b/packages/tools/src/orchestrator/index.test.ts index 1dffc4dc..080f1404 100644 --- a/packages/tools/src/orchestrator/index.test.ts +++ b/packages/tools/src/orchestrator/index.test.ts @@ -183,4 +183,70 @@ describe('tool orchestrator', () => { } } }) + + test('rejects oversized raw string input before json parsing', async () => { + let executed = false + const orchestrator = createToolOrchestrator({ + tools: { + read_file: { + name: 'read_file', + execute: async () => { + executed = true + return { content: [{ type: 'text', text: 'ok' }] } + }, + }, + }, + }) + + const result = await orchestrator.executeAction({ + name: 'read_file', + input: 'x'.repeat(100_001), + }) + + assert.strictEqual(result.success, false) + assert.strictEqual(result.status, 'input_invalid') + assert.ok(result.observation.includes('input string too large')) + assert.strictEqual(executed, false) + }) + + test('rejects non-object json input payloads', async () => { + const orchestrator = createToolOrchestrator({ + tools: { + read_file: { + name: 'read_file', + execute: async () => ({ content: [{ type: 'text', text: 'ok' }] }), + }, + }, + }) + + const result = await orchestrator.executeAction({ + name: 'read_file', + input: '[]', + }) + + assert.strictEqual(result.success, false) + assert.strictEqual(result.status, 'input_invalid') + assert.ok(result.observation.includes('expected object')) + }) + + test('rejects validateInput success payload when shape is not object', async () => { + const orchestrator = createToolOrchestrator({ + tools: { + read_file: { + name: 'read_file', + validateInput: () => ({ ok: true, data: 'not-object' }), + execute: async () => ({ content: [{ type: 'text', text: 'ok' }] }), + }, + }, + }) + + const result = await orchestrator.executeAction({ + name: 'read_file', + input: { file_path: '/tmp/a' }, + }) + + assert.strictEqual(result.success, false) + assert.strictEqual(result.status, 'input_invalid') + assert.ok(result.observation.includes('expected object')) + }) }) diff --git a/packages/tools/src/orchestrator/index.ts b/packages/tools/src/orchestrator/index.ts index dd229343..e3dcf6ff 100644 --- a/packages/tools/src/orchestrator/index.ts +++ b/packages/tools/src/orchestrator/index.ts @@ -13,6 +13,7 @@ import type { } from './types' const DEFAULT_MAX_TOOL_RESULT_CHARS = 12_000 +const MAX_TOOL_INPUT_STRING_CHARS = 100_000 function getMaxToolResultChars() { const raw = process.env.MEMO_TOOL_RESULT_MAX_CHARS?.trim() @@ -76,9 +77,23 @@ function flattenCallToolResult(result: CallToolResult): string { return texts.join('\n') } -function parseToolInput(tool: OrchestratorTool, rawInput: unknown) { +type ParseToolInputResult = + | { ok: true; data: Record } + | { ok: false; error: string } + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +function parseToolInput(tool: OrchestratorTool, rawInput: unknown): ParseToolInputResult { let candidate: unknown = rawInput if (typeof rawInput === 'string') { + if (rawInput.length > MAX_TOOL_INPUT_STRING_CHARS) { + return { + ok: false as const, + error: `${tool.name} invalid input: input string too large (max ${MAX_TOOL_INPUT_STRING_CHARS} chars)`, + } + } const trimmed = rawInput.trim() if (trimmed) { try { @@ -91,12 +106,17 @@ function parseToolInput(tool: OrchestratorTool, rawInput: unknown) { } } - if (typeof candidate !== 'object' || candidate === null) { + if (!isRecord(candidate)) { return { ok: false as const, error: `${tool.name} invalid input: expected object` } } if (typeof tool.validateInput === 'function') { - return tool.validateInput(candidate) + const validated = tool.validateInput(candidate) + if (!validated.ok) return validated + if (!isRecord(validated.data)) { + return { ok: false as const, error: `${tool.name} invalid input: expected object` } + } + return { ok: true as const, data: validated.data } } return { ok: true as const, data: candidate } diff --git a/packages/tools/src/tools/apply_patch.test.ts b/packages/tools/src/tools/apply_patch.test.ts index 0c9cd2ed..3699cda9 100644 --- a/packages/tools/src/tools/apply_patch.test.ts +++ b/packages/tools/src/tools/apply_patch.test.ts @@ -1,5 +1,5 @@ import assert from 'node:assert' -import { access, mkdir, readFile, rm, writeFile } from 'node:fs/promises' +import { access, mkdir, readFile, rm, symlink, writeFile } from 'node:fs/promises' import { join } from 'node:path' import { tmpdir } from 'node:os' import { afterAll, beforeAll, describe, test } from 'vitest' @@ -181,7 +181,25 @@ describe('apply_patch tool (direct replace)', () => { assertPatchError(result) const text = textPayload(result) - assert.ok(text.includes('sandbox') || text.includes('不在允许目录')) + assert.ok(text.includes('sandbox write denied')) + await rm(outside, { force: true }) + }) + + test('denies writes when symlink target is outside writable roots', async () => { + const outside = '/tmp/memo-tools-apply-patch-symlink-target.txt' + const linkPath = join(tempDir, 'symlink-outside.txt') + await writeFile(outside, 'secret', 'utf8') + await symlink(outside, linkPath) + + const result = await applyPatchTool.execute({ + file_path: linkPath, + old_string: 'secret', + new_string: 'patched', + }) + + assertPatchError(result, 'sandbox write denied') + assert.strictEqual(await readText(outside), 'secret') + await rm(linkPath, { force: true }) await rm(outside, { force: true }) }) }) diff --git a/packages/tools/src/tools/apply_patch.ts b/packages/tools/src/tools/apply_patch.ts index 8defa8f2..26ce7ba1 100644 --- a/packages/tools/src/tools/apply_patch.ts +++ b/packages/tools/src/tools/apply_patch.ts @@ -1,4 +1,4 @@ -import { access, readFile, writeFile } from 'node:fs/promises' +import { access, readFile, realpath, writeFile } from 'node:fs/promises' import { z } from 'zod' import { textResult } from '@memo/tools/tools/mcp' import { normalizePath, writePathDenyReason } from '@memo/tools/tools/helpers' @@ -102,13 +102,14 @@ export const applyPatchTool = defineMcpTool({ supportsParallelToolCalls: false, isMutating: true, execute: async (input) => { - const filePath = normalizePath(input.file_path) + const userPath = normalizePath(input.file_path) try { - ensureWritable(filePath) - await access(filePath) + await access(userPath) + const targetPath = normalizePath(await realpath(userPath)) + ensureWritable(targetPath) - const original = await readFile(filePath, 'utf8') + const original = await readFile(targetPath, 'utf8') const operations = toOperations(input) let working = original @@ -146,14 +147,14 @@ export const applyPatchTool = defineMcpTool({ return textResult('No changes made.') } - await writeFile(filePath, working, 'utf8') + await writeFile(targetPath, working, 'utf8') return textResult( - `Success. Updated file: ${filePath}\nEdits: ${operations.length}\nReplacements: ${replacementCount}`, + `Success. Updated file: ${targetPath}\nEdits: ${operations.length}\nReplacements: ${replacementCount}`, ) } catch (err) { const code = (err as NodeJS.ErrnoException).code if (code === 'ENOENT') { - return textResult(`apply_patch failed: file does not exist: ${filePath}`, true) + return textResult(`apply_patch failed: file does not exist: ${userPath}`, true) } return textResult(`apply_patch failed: ${(err as Error).message}`, true) } diff --git a/packages/tools/src/tools/codex_tools.test.ts b/packages/tools/src/tools/codex_tools.test.ts index 87644cad..6b8ddd5c 100644 --- a/packages/tools/src/tools/codex_tools.test.ts +++ b/packages/tools/src/tools/codex_tools.test.ts @@ -37,6 +37,13 @@ function textPayload(result: { content?: Array<{ type: string; text?: string }> return first?.text ?? '' } +function outputPayload(text: string) { + const marker = 'Output:\n' + const index = text.indexOf(marker) + if (index < 0) return '' + return text.slice(index + marker.length) +} + beforeAll(async () => { tempDir = await makeTempDir('memo-tools-codex') prevWritableRoots = process.env.MEMO_SANDBOX_WRITABLE_ROOTS @@ -123,6 +130,44 @@ describe('codex shell family', () => { const resumedText = textPayload(resumed) assert.ok(resumedText.includes('still-alive')) }) + + test('write_stdin can fetch unread output tail after truncation', async () => { + const started = await execCommandTool.execute({ + cmd: `node -e "process.stdout.write('X'.repeat(5000)); setTimeout(() => {}, 2000)"`, + yield_time_ms: 300, + max_output_tokens: 10, + }) + const startedText = textPayload(started) + const match = startedText.match(/session ID (\d+)/) + assert.ok(match, `expected running session id, got: ${startedText}`) + const firstChunk = outputPayload(startedText) + assert.strictEqual(firstChunk.length, 40) + + const sessionId = Number(match?.[1]) + const next = await writeStdinTool.execute({ + session_id: sessionId, + yield_time_ms: 100, + max_output_tokens: 2000, + }) + const nextText = textPayload(next) + const nextChunk = outputPayload(nextText) + assert.ok(nextChunk.length > 0, `expected unread tail, got: ${nextText}`) + assert.ok(nextChunk.includes('X')) + }) + + test('exec_command rejects when active session cap is exceeded', async () => { + const results = [] + for (let i = 0; i < 70; i += 1) { + results.push(await execCommandTool.execute({ cmd: 'sleep 2', yield_time_ms: 0 })) + } + const overflow = results.find( + (result) => result.isError && textPayload(result).includes('too many active sessions'), + ) + assert.ok(overflow, 'expected active-session cap error') + + // Allow spawned sessions to exit before subsequent tests. + await new Promise((resolve) => setTimeout(resolve, 2200)) + }) }) describe('codex file/search family', () => { diff --git a/packages/tools/src/tools/exec_runtime.ts b/packages/tools/src/tools/exec_runtime.ts index 1542109e..188b6fee 100644 --- a/packages/tools/src/tools/exec_runtime.ts +++ b/packages/tools/src/tools/exec_runtime.ts @@ -19,6 +19,7 @@ type StartExecRequest = { login?: boolean tty?: boolean yield_time_ms?: number + execution_timeout_ms?: number max_output_tokens?: number source_tool?: string } @@ -102,12 +103,14 @@ function truncateByTokens(text: string, maxOutputTokens?: number) { return { output: text, originalTokenCount, + deliveredChars: text.length, } } return { output: text.slice(0, maxChars), originalTokenCount, + deliveredChars: maxChars, } } @@ -153,6 +156,27 @@ async function waitForWindow(session: SessionState, yieldMs: number): Promise { + if (session.exited || timeoutMs <= 0) return + await Promise.race([ + new Promise((resolve) => { + const timer = setTimeout(() => { + cleanup() + resolve() + }, timeoutMs) + const onExit = () => { + clearTimeout(timer) + cleanup() + resolve() + } + const cleanup = () => { + session.eventBus.off('exit', onExit) + } + session.eventBus.on('exit', onExit) + }), + ]) +} + class UnifiedExecManager { private sessions = new Map() private nextId = 1 @@ -169,12 +193,35 @@ class UnifiedExecManager { } } + private activeSessionCount() { + let count = 0 + for (const session of this.sessions.values()) { + if (!session.exited) count += 1 + } + return count + } + + private async terminateForTimeout(session: SessionState) { + if (session.exited) return + session.proc.kill('SIGTERM') + await waitForExit(session, 200) + if (!session.exited) { + session.proc.kill('SIGKILL') + await waitForExit(session, 200) + } + } + async start(request: StartExecRequest): Promise { const cmd = request.cmd.trim() if (!cmd) { throw new Error('cmd must not be empty') } + this.cleanupSessions() + if (this.activeSessionCount() >= MAX_SESSIONS) { + throw new Error(`too many active sessions (max ${MAX_SESSIONS})`) + } + const blocked = guardDangerousCommand({ toolName: request.source_tool ?? 'exec_command', command: cmd, @@ -235,8 +282,25 @@ class UnifiedExecManager { this.sessions.set(id, session) this.cleanupSessions() + const executionTimeoutMs = + typeof request.execution_timeout_ms === 'number' && request.execution_timeout_ms > 0 + ? Math.floor(request.execution_timeout_ms) + : null const yieldMs = clampYield(request.yield_time_ms, DEFAULT_EXEC_YIELD_TIME_MS) - await waitForWindow(session, yieldMs) + const waitMs = + executionTimeoutMs !== null + ? Math.min(yieldMs, Math.max(0, executionTimeoutMs)) + : yieldMs + await waitForWindow(session, waitMs) + + if (executionTimeoutMs !== null && !session.exited) { + const elapsedMs = Date.now() - startedAtMs + if (elapsedMs >= executionTimeoutMs) { + await this.terminateForTimeout(session) + this.cleanupSessions() + throw new Error(`command timed out after ${executionTimeoutMs}ms`) + } + } return this.buildResponseText(session, request.max_output_tokens) } @@ -277,8 +341,8 @@ class UnifiedExecManager { private buildResponseText(session: SessionState, maxOutputTokens?: number): string { const delta = session.output.slice(session.readOffset) - session.readOffset = session.output.length const truncated = truncateByTokens(delta, maxOutputTokens) + session.readOffset += truncated.deliveredChars const payload: SessionResponse = { sessionId: session.id, diff --git a/packages/tools/src/tools/helpers.test.ts b/packages/tools/src/tools/helpers.test.ts index 6ce2e697..6649241e 100644 --- a/packages/tools/src/tools/helpers.test.ts +++ b/packages/tools/src/tools/helpers.test.ts @@ -2,7 +2,7 @@ import assert from 'node:assert' import { join } from 'node:path' import { tmpdir } from 'node:os' import { describe, test } from 'vitest' -import { normalizePath, isWritePathAllowed } from '@memo/tools/tools/helpers' +import { normalizePath, isWritePathAllowed, writePathDenyReason } from '@memo/tools/tools/helpers' describe('helpers.normalizePath', () => { test('normalizes relative paths to absolute', () => { @@ -33,6 +33,9 @@ describe('helpers.sandbox', () => { try { const outside = normalizePath('/etc/passwd') assert.strictEqual(isWritePathAllowed(outside), false) + const reason = writePathDenyReason(outside) ?? '' + assert.ok(reason.includes('sandbox write denied')) + assert.ok(reason.includes('is not within allowed directories')) } finally { if (prev === undefined) delete process.env.MEMO_SANDBOX_WRITABLE_ROOTS else process.env.MEMO_SANDBOX_WRITABLE_ROOTS = prev diff --git a/packages/tools/src/tools/helpers.ts b/packages/tools/src/tools/helpers.ts index 305cd58a..e0d76cd9 100644 --- a/packages/tools/src/tools/helpers.ts +++ b/packages/tools/src/tools/helpers.ts @@ -1,6 +1,6 @@ import { normalize, resolve, dirname, join, relative, isAbsolute } from 'node:path' import { homedir } from 'node:os' -import { existsSync, statSync } from 'node:fs' +import { existsSync, statSync, realpathSync } from 'node:fs' import { readFile } from 'node:fs/promises' import ignore, { type Ignore } from 'ignore' @@ -11,6 +11,15 @@ export function normalizePath(rawPath: string) { return normalize(resolve(rawPath)) } +function canonicalizePath(absPath: string) { + const normalized = normalizePath(absPath) + try { + return normalizePath(realpathSync(normalized)) + } catch { + return normalized + } +} + function isSubPath(targetPath: string, rootPath: string) { const rel = relative(rootPath, targetPath) return rel === '' || (!rel.startsWith('..') && !isAbsolute(rel)) @@ -23,14 +32,14 @@ function parseWritableRootsFromEnv() { .split(',') .map((item) => item.trim()) .filter(Boolean) - .map((item) => normalizePath(item)) + .map((item) => canonicalizePath(item)) } export function getWritableRoots() { const roots = new Set() - roots.add(normalizePath(process.cwd())) + roots.add(canonicalizePath(process.cwd())) const memoHome = process.env.MEMO_HOME?.trim() || join(homedir(), '.memo') - roots.add(normalizePath(memoHome)) + roots.add(canonicalizePath(memoHome)) for (const item of parseWritableRootsFromEnv()) { roots.add(item) } @@ -38,14 +47,15 @@ export function getWritableRoots() { } export function isWritePathAllowed(absPath: string) { + const canonicalTarget = canonicalizePath(absPath) const roots = getWritableRoots() - return roots.some((root) => isSubPath(absPath, root)) + return roots.some((root) => isSubPath(canonicalTarget, root)) } export function writePathDenyReason(absPath: string) { if (isWritePathAllowed(absPath)) return null const roots = getWritableRoots() - return `sandbox 拒绝写入: ${absPath} 不在允许目录内 (${roots.join(', ')})` + return `sandbox write denied: ${absPath} is not within allowed directories (${roots.join(', ')})` } export const DEFAULT_IGNORE_PATTERNS = [ diff --git a/packages/tools/src/tools/shell.ts b/packages/tools/src/tools/shell.ts index e29f58d7..77dfee88 100644 --- a/packages/tools/src/tools/shell.ts +++ b/packages/tools/src/tools/shell.ts @@ -16,13 +16,16 @@ const SHELL_INPUT_SCHEMA = z type ShellInput = z.infer +const SAFE_SHELL_ARG = /^[A-Za-z0-9_./:@%+-]+$/ + +function shellQuote(part: string) { + if (part.length === 0) return "''" + if (SAFE_SHELL_ARG.test(part)) return part + return `'${part.replace(/'/g, `'\"'\"'`)}'` +} + function shellJoin(argv: string[]) { - return argv - .map((part) => { - if (/^[A-Za-z0-9_./:@%+-]+$/.test(part)) return part - return JSON.stringify(part) - }) - .join(' ') + return argv.map((part) => shellQuote(part)).join(' ') } export const shellTool = defineMcpTool({ @@ -38,6 +41,7 @@ export const shellTool = defineMcpTool({ workdir, login: false, yield_time_ms: timeout_ms, + execution_timeout_ms: timeout_ms, source_tool: 'shell', }) return textResult(content) diff --git a/packages/tools/src/tools/shell_command.ts b/packages/tools/src/tools/shell_command.ts index 399a986c..ce174602 100644 --- a/packages/tools/src/tools/shell_command.ts +++ b/packages/tools/src/tools/shell_command.ts @@ -30,6 +30,7 @@ export const shellCommandTool = defineMcpTool({ workdir, login, yield_time_ms: timeout_ms, + execution_timeout_ms: timeout_ms, source_tool: 'shell_command', }) return textResult(content) diff --git a/packages/tools/src/tools/shell_update_plan.test.ts b/packages/tools/src/tools/shell_update_plan.test.ts index 38ebd2f9..50281046 100644 --- a/packages/tools/src/tools/shell_update_plan.test.ts +++ b/packages/tools/src/tools/shell_update_plan.test.ts @@ -21,6 +21,19 @@ describe('shell wrappers and update_plan', () => { assert.ok(text.includes('shell-wrapper-ok')) }) + test('shell tool quotes dangerous shell metacharacters in argv', async () => { + const literal = '$HOME $(echo hacked);`date`' + const result = await shellTool.execute({ + command: ['printf', '%s', literal], + }) + + const text = textPayload(result) + assert.ok(!result.isError) + assert.ok(text.includes(literal)) + assert.ok(text.includes('$(echo hacked)')) + assert.ok(text.includes('`date`')) + }) + test('shell tool blocks dangerous argv command', async () => { const result = await shellTool.execute({ command: ['mkfs.ext4', '/dev/sda'], @@ -32,6 +45,20 @@ describe('shell wrappers and update_plan', () => { assert.ok(text.includes('reason="dangerous_command"')) }) + test('shell tool enforces timeout_ms as execution deadline', async () => { + const startedAt = Date.now() + const result = await shellTool.execute({ + command: ['sleep', '2'], + timeout_ms: 100, + }) + + const elapsedMs = Date.now() - startedAt + const text = textPayload(result) + assert.strictEqual(result.isError, true) + assert.ok(text.includes('timed out')) + assert.ok(elapsedMs < 1_500) + }) + test('shell_command executes script command form', async () => { const result = await shellCommandTool.execute({ command: 'echo shell-command-ok', @@ -57,6 +84,21 @@ describe('shell wrappers and update_plan', () => { assert.ok(text.includes('reason="dangerous_command"')) }) + test('shell_command enforces timeout_ms as execution deadline', async () => { + const startedAt = Date.now() + const result = await shellCommandTool.execute({ + command: 'sleep 2; echo too-late', + login: false, + timeout_ms: 100, + }) + + const elapsedMs = Date.now() - startedAt + const text = textPayload(result) + assert.strictEqual(result.isError, true) + assert.ok(text.includes('timed out')) + assert.ok(elapsedMs < 1_500) + }) + test('write_stdin fails for unknown session id', async () => { const result = await writeStdinTool.execute({ session_id: 999999, chars: 'noop' }) assert.strictEqual(result.isError, true) diff --git a/packages/tools/src/tools/webfetch.test.ts b/packages/tools/src/tools/webfetch.test.ts index 79d261c1..ad90ad94 100644 --- a/packages/tools/src/tools/webfetch.test.ts +++ b/packages/tools/src/tools/webfetch.test.ts @@ -2,38 +2,72 @@ import assert from 'node:assert' import { describe, test } from 'vitest' import { webfetchTool } from '@memo/tools/tools/webfetch' +function textPayload(result: { content?: Array<{ type: string; text?: string }> }) { + return result.content?.find((item) => item.type === 'text')?.text ?? '' +} + +async function withMockFetch( + mock: typeof globalThis.fetch, + run: () => Promise, +): Promise { + const original = globalThis.fetch + Object.assign(globalThis, { fetch: mock }) + try { + await run() + } finally { + Object.assign(globalThis, { fetch: original }) + } +} + describe('webfetch tool', () => { test('requires url', async () => { const res = webfetchTool.validateInput?.({ url: '' }) assert.ok(res && !res.ok) }) - test('fetches data url content', async () => { + test('rejects data protocol', async () => { const res = await webfetchTool.execute({ url: 'data:text/plain,hello' }) - const text = res.content?.[0]?.type === 'text' ? res.content[0].text : '' - assert.ok(text.includes('status=200')) - assert.ok(text.includes('text="hello"')) + const text = textPayload(res) + assert.strictEqual(res.isError, true) + assert.ok(text.includes('Unsupported protocol')) }) test('returns plain text when body is html', async () => { - const html = encodeURIComponent('

Hello

World

') - const res = await webfetchTool.execute({ url: `data:text/html,${html}` }) - const text = res.content?.[0]?.type === 'text' ? res.content[0].text : '' - assert.ok(!text.includes('

'), 'should strip HTML tags') - assert.ok(text.includes('Hello'), 'should keep visible text') - assert.ok(text.includes('World'), 'should keep paragraph text') + await withMockFetch( + async () => + new Response('

Hello

World

', { + status: 200, + headers: { 'content-type': 'text/html; charset=utf-8' }, + }), + async () => { + const res = await webfetchTool.execute({ url: 'https://example.com' }) + const text = textPayload(res) + assert.ok(!text.includes('

'), 'should strip HTML tags') + assert.ok(text.includes('Hello'), 'should keep visible text') + assert.ok(text.includes('World'), 'should keep paragraph text') + }, + ) }) test('rejects unsupported protocol', async () => { const res = await webfetchTool.execute({ url: 'file:///etc/hosts' }) - const text = res.content?.[0]?.type === 'text' ? res.content[0].text : '' + const text = textPayload(res) assert.ok(text.includes('Unsupported protocol'), 'should block file:// scheme') }) test('aborts when body is too large', async () => { - const large = 'a'.repeat(600_000) - const res = await webfetchTool.execute({ url: `data:text/plain,${large}` }) - const text = res.content?.[0]?.type === 'text' ? res.content[0].text : '' - assert.ok(text.includes('response body exceeds'), 'should enforce size limit') + await withMockFetch( + async () => + new Response('ignored', { + status: 200, + headers: { 'content-length': '600000' }, + }), + async () => { + const res = await webfetchTool.execute({ url: 'https://example.com/large' }) + const text = textPayload(res) + assert.strictEqual(res.isError, true) + assert.ok(text.includes('response body too large'), 'should enforce size limit') + }, + ) }) }) diff --git a/packages/tools/src/tools/webfetch.ts b/packages/tools/src/tools/webfetch.ts index 74dc1604..033b5aa6 100644 --- a/packages/tools/src/tools/webfetch.ts +++ b/packages/tools/src/tools/webfetch.ts @@ -13,7 +13,7 @@ type WebFetchInput = z.infer const WEBFETCH_TIMEOUT_MS = 10_000 const MAX_BODY_BYTES = 512_000 const MAX_BODY_PREVIEW = 4_000 -const ALLOWED_PROTOCOLS = new Set(['http:', 'https:', 'data:']) +const ALLOWED_PROTOCOLS = new Set(['http:', 'https:']) const HTML_BREAK_TAG = /<\/\s*(p|div|section|article|header|footer|aside|main|h[1-6]|li|tr|table|blockquote)\s*>/gi const HTML_LINE_BREAK = /<\s*(br|hr)\s*\/?>/gi