fix: bound agent communication memory growth

UDS messaging now uses private local capabilities instead of exposing auth tokens through SDK metadata, environment variables, session registry, peer listing, or tool output. The receive path bounds NDJSON frames, response buffers, active clients, and pending inbox bytes, and strips auth metadata before messages enter the prompt queue.

Teammate mailboxes now validate file and message sizes, fail closed on corrupt mutation inputs, compact by count and retained bytes, and use stable message identity for in-process acknowledgements. Agent summaries now fork only a bounded recent context using lazy size estimation and content fingerprints instead of retaining or serializing unbounded histories.

Constraint: PR #361 was already merged; this branch is based on upstream/main@c2ac9a74.
Rejected: Default-disabling COORDINATOR_MODE/TEAMMEM only | explicit feature enablement still hit unbounded paths.
Rejected: Persisting UDS auth in SDK/env/session registry | bridge/remote metadata can leak local capability secrets.
Rejected: Inline uds #token addresses | observable/tool/classifier paths can reflect raw addresses outside the UDS request frame.
Rejected: Positional mailbox marking after compaction | compaction can shift indices across the lock boundary.
Confidence: high
Scope-risk: moderate
Directive: Do not expose UDS capability tokens through SDK messages, environment variables, session registry, peer-list output, or SendMessage result/classifier surfaces.
Directive: Do not reintroduce positional mailbox acknowledgements unless compaction is removed or read+mark is atomic under one lock.
Tested: bun test src/utils/__tests__/ndjsonFramer.test.ts src/utils/__tests__/udsMessaging.test.ts packages/builtin-tools/src/tools/SendMessageTool/__tests__/udsRecipientSanitization.test.ts
Tested: bunx tsc --noEmit --pretty false
Tested: bun run lint
Tested: bunx biome lint modified src/package files
Tested: bun run test:all (3704 pass, 0 fail, 6734 expects)
Tested: bun audit (No vulnerabilities found)
Tested: bun run build
Tested: bun run build:vite
Tested: git diff --check
Not-tested: End-to-end external UDS client driving a full production headless model turn.
This commit is contained in:
unraid
2026-04-26 21:44:42 +08:00
parent c2ac9a74c1
commit f353eb056a
17 changed files with 2087 additions and 137 deletions

View File

@@ -0,0 +1,91 @@
import { EventEmitter } from 'node:events'
import type { Socket } from 'node:net'
import { describe, expect, test } from 'bun:test'
import { attachNdjsonFramer } from '../ndjsonFramer.js'
type TestSocket = Socket & {
destroyed: boolean
emitData: (chunk: Buffer) => void
}
function createTestSocket(): TestSocket {
const emitter = new EventEmitter() as TestSocket
emitter.destroyed = false
emitter.destroy = ((_error?: Error) => {
emitter.destroyed = true
emitter.emit('close')
return emitter
}) as TestSocket['destroy']
emitter.emitData = (chunk: Buffer) => {
emitter.emit('data', chunk)
}
return emitter
}
describe('attachNdjsonFramer', () => {
test('accepts a complete frame at the configured byte limit', () => {
const socket = createTestSocket()
const messages: unknown[] = []
const errors: Error[] = []
attachNdjsonFramer(
socket,
msg => messages.push(msg),
text => JSON.parse(text) as unknown,
{
maxFrameBytes: Buffer.byteLength('{"a":1}', 'utf8'),
onFrameError: error => errors.push(error),
},
)
socket.emitData(Buffer.from('{"a":1}\n'))
expect(messages).toEqual([{ a: 1 }])
expect(errors).toEqual([])
expect(socket.destroyed).toBe(false)
})
test('destroys a complete frame over the configured byte limit', () => {
const socket = createTestSocket()
const messages: unknown[] = []
const errors: Error[] = []
attachNdjsonFramer(
socket,
msg => messages.push(msg),
text => JSON.parse(text) as unknown,
{
maxFrameBytes: 8,
onFrameError: error => errors.push(error),
},
)
socket.emitData(Buffer.from('{"long":true}\n'))
expect(messages).toEqual([])
expect(errors[0]?.message).toContain('NDJSON frame exceeded')
expect(socket.destroyed).toBe(true)
})
test('destroys oversized no-newline input before a frame can form', () => {
const socket = createTestSocket()
const messages: unknown[] = []
const errors: Error[] = []
attachNdjsonFramer(
socket,
msg => messages.push(msg),
text => JSON.parse(text) as unknown,
{
maxFrameBytes: 8,
onFrameError: error => errors.push(error),
},
)
socket.emitData(Buffer.from('x'.repeat(9)))
expect(messages).toEqual([])
expect(errors[0]?.message).toContain('NDJSON frame exceeded')
expect(socket.destroyed).toBe(true)
})
})

View File

@@ -0,0 +1,310 @@
import { afterEach, beforeEach, describe, expect, test } from 'bun:test'
import { mkdir, readFile, rm, writeFile } from 'node:fs/promises'
import { mkdtempSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { dirname, join } from 'node:path'
import {
compactMailboxMessages,
getInboxPath,
markMessageAsReadByIndex,
markMessageAsReadByIdentity,
markMessagesAsRead,
markMessagesAsReadByPredicate,
MAX_MAILBOX_MESSAGE_TEXT_BYTES,
MAX_MAILBOX_MESSAGES,
MAX_READ_MAILBOX_MESSAGES,
MAX_UNREAD_PROTOCOL_MAILBOX_MESSAGES,
readMailbox,
type TeammateMessage,
writeToMailbox,
} from '../teammateMailbox.js'
let tempHome = ''
let previousConfigDir: string | undefined
function message(
text: string,
read: boolean,
timestamp = new Date(0).toISOString(),
): TeammateMessage {
return {
from: 'team-lead',
text,
timestamp,
read,
}
}
async function seedMailbox(
agentName: string,
teamName: string,
messages: TeammateMessage[],
): Promise<void> {
const inboxPath = getInboxPath(agentName, teamName)
await mkdir(dirname(inboxPath), { recursive: true })
await writeFile(inboxPath, JSON.stringify(messages, null, 2), 'utf-8')
}
async function readRawMailbox(
agentName: string,
teamName: string,
): Promise<TeammateMessage[]> {
const content = await readFile(getInboxPath(agentName, teamName), 'utf-8')
return JSON.parse(content) as TeammateMessage[]
}
beforeEach(() => {
previousConfigDir = process.env.CLAUDE_CONFIG_DIR
tempHome = mkdtempSync(join(tmpdir(), 'teammate-mailbox-'))
process.env.CLAUDE_CONFIG_DIR = tempHome
})
afterEach(async () => {
if (previousConfigDir === undefined) {
delete process.env.CLAUDE_CONFIG_DIR
} else {
process.env.CLAUDE_CONFIG_DIR = previousConfigDir
}
await rm(tempHome, { recursive: true, force: true })
})
describe('compactMailboxMessages', () => {
test('prioritizes unread messages and keeps only recent read history', () => {
const compacted = compactMailboxMessages(
[
message('read-1', true),
message('read-2', true),
message('unread-1', false),
message('read-3', true),
message('unread-2', false),
message('read-4', true),
message('read-5', true),
message('unread-3', false),
],
{ maxMessages: 5, maxReadMessages: 2 },
)
expect(compacted.map(m => m.text)).toEqual([
'unread-1',
'unread-2',
'read-4',
'read-5',
'unread-3',
])
})
test('retains unread protocol messages separately from regular cap', () => {
const protocol = message(
JSON.stringify({ type: 'permission_response', request_id: 'req-1' }),
false,
)
const compacted = compactMailboxMessages(
[
protocol,
...Array.from({ length: 5 }, (_value, index) =>
message(`regular-${index}`, false),
),
],
{
maxMessages: 2,
maxReadMessages: 0,
maxUnreadProtocolMessages: 1,
},
)
expect(compacted.map(m => m.text)).toEqual([
protocol.text,
'regular-3',
'regular-4',
])
})
test('caps unread protocol messages with an independent bound', () => {
const compacted = compactMailboxMessages(
Array.from(
{ length: MAX_UNREAD_PROTOCOL_MAILBOX_MESSAGES + 1 },
(_value, index) =>
message(
JSON.stringify({
type: 'permission_response',
request_id: `req-${index}`,
}),
false,
),
),
)
expect(compacted).toHaveLength(MAX_UNREAD_PROTOCOL_MAILBOX_MESSAGES)
expect(compacted[0]?.text).toContain('req-1')
})
test('keeps retained mailbox bytes under an explicit budget', () => {
const compacted = compactMailboxMessages(
Array.from({ length: 20 }, (_value, index) =>
message(`msg-${index}-${'x'.repeat(200)}`, false),
),
{
maxMessages: 20,
maxReadMessages: 0,
maxRetainedBytes: 1_000,
},
)
expect(
Buffer.byteLength(JSON.stringify(compacted), 'utf8'),
).toBeLessThanOrEqual(1_000)
expect(compacted.length).toBeLessThan(20)
expect(compacted.at(-1)?.text).toContain('msg-19')
})
})
describe('teammate mailbox retention', () => {
test('writeToMailbox compacts oversized unread inbox files', async () => {
const existing = Array.from(
{ length: MAX_MAILBOX_MESSAGES + 20 },
(_value, index) => message(`old-${index}`, false),
)
await seedMailbox('worker', 'alpha', existing)
await writeToMailbox(
'worker',
{
from: 'team-lead',
text: 'newest',
timestamp: new Date(1).toISOString(),
},
'alpha',
)
const after = await readMailbox('worker', 'alpha')
expect(after).toHaveLength(MAX_MAILBOX_MESSAGES)
expect(after[0]?.text).toBe('old-21')
expect(after.at(-1)?.text).toBe('newest')
})
test('markMessagesAsRead compacts read history after consumption', async () => {
const existing = Array.from(
{ length: MAX_MAILBOX_MESSAGES + 20 },
(_value, index) => message(`msg-${index}`, false),
)
await seedMailbox('worker', 'alpha', existing)
await markMessagesAsRead('worker', 'alpha')
const after = await readRawMailbox('worker', 'alpha')
expect(after).toHaveLength(MAX_READ_MAILBOX_MESSAGES)
expect(after.every(m => m.read)).toBe(true)
expect(after[0]?.text).toBe(
`msg-${MAX_MAILBOX_MESSAGES + 20 - MAX_READ_MAILBOX_MESSAGES}`,
)
})
test('markMessagesAsReadByPredicate leaves structured messages unread', async () => {
await seedMailbox('worker', 'alpha', [
message('plain', false),
message(JSON.stringify({ type: 'permission_request' }), false),
])
await markMessagesAsReadByPredicate(
'worker',
m => !m.text.includes('permission_request'),
'alpha',
)
const after = await readRawMailbox('worker', 'alpha')
expect(after.map(m => m.read)).toEqual([true, false])
})
test('markMessageAsReadByIdentity survives compaction shifting indexes', async () => {
const permissionResponse = message(
JSON.stringify({ type: 'permission_response', request_id: 'req-1' }),
false,
)
await seedMailbox('worker', 'alpha', [
permissionResponse,
...Array.from({ length: MAX_MAILBOX_MESSAGES + 20 }, (_value, index) =>
message(`regular-${index}`, false),
),
])
await writeToMailbox(
'worker',
{
from: 'team-lead',
text: 'newest',
timestamp: new Date(2).toISOString(),
},
'alpha',
)
const marked = await markMessageAsReadByIdentity(
'worker',
'alpha',
permissionResponse,
)
const after = await readRawMailbox('worker', 'alpha')
expect(marked).toBe(true)
expect(after.some(m => m.text === permissionResponse.text && !m.read)).toBe(
false,
)
})
test('markMessageAsReadByIndex also compacts through the compatibility path', async () => {
const existing = Array.from(
{ length: MAX_MAILBOX_MESSAGES + 10 },
(_value, index) => message(`msg-${index}`, false),
)
await seedMailbox('worker', 'alpha', existing)
await markMessageAsReadByIndex('worker', 'alpha', existing.length - 1)
const after = await readRawMailbox('worker', 'alpha')
expect(after).toHaveLength(MAX_MAILBOX_MESSAGES)
expect(after.some(m => m.text === `msg-${existing.length - 1}`)).toBe(false)
expect(after.at(-1)?.text).toBe(`msg-${existing.length - 2}`)
})
test('writeToMailbox rejects oversized message text instead of storing it', async () => {
await expect(
writeToMailbox(
'worker',
{
from: 'team-lead',
text: 'x'.repeat(MAX_MAILBOX_MESSAGE_TEXT_BYTES + 1),
timestamp: new Date(3).toISOString(),
},
'alpha',
),
).rejects.toThrow('Mailbox message text exceeds')
expect(await readRawMailbox('worker', 'alpha')).toEqual([])
})
test('writeToMailbox fails closed when an existing mailbox is corrupt', async () => {
const inboxPath = getInboxPath('worker', 'alpha')
await mkdir(dirname(inboxPath), { recursive: true })
await writeFile(inboxPath, '{not-json', 'utf-8')
await expect(
writeToMailbox(
'worker',
{
from: 'team-lead',
text: 'new',
timestamp: new Date(4).toISOString(),
},
'alpha',
),
).rejects.toThrow()
expect(await readFile(inboxPath, 'utf-8')).toBe('{not-json')
})
test('readMailbox fails closed on corrupt mailbox content', async () => {
const inboxPath = getInboxPath('worker', 'alpha')
await mkdir(dirname(inboxPath), { recursive: true })
await writeFile(inboxPath, '{not-json', 'utf-8')
await expect(readMailbox('worker', 'alpha')).rejects.toThrow()
})
})

View File

@@ -0,0 +1,305 @@
import { afterEach, describe, expect, test } from 'bun:test'
import { chmod, mkdir, rm, stat, symlink, unlink } from 'node:fs/promises'
import { createConnection, createServer } from 'node:net'
import { dirname, join } from 'node:path'
import { tmpdir } from 'node:os'
import {
drainInbox,
MAX_UDS_INBOX_ENTRIES,
MAX_UDS_INBOX_BYTES,
MAX_UDS_FRAME_BYTES,
parseUdsTarget,
sendUdsMessage,
setOnEnqueue,
startUdsMessaging,
stopUdsMessaging,
} from '../udsMessaging.js'
function socketPath(label: string): string {
const suffix = `${process.pid}-${Date.now()}-${Math.random().toString(16).slice(2)}-${label}`
if (process.platform === 'win32') {
return `\\\\.\\pipe\\claude-code-test-${suffix}`
}
return join(tmpdir(), 'claude-code-test', `${suffix}.sock`)
}
function sleep(ms: number): Promise<void> {
return new Promise(resolve => setTimeout(resolve, ms))
}
async function waitForEnqueues(
expected: number,
sendMessages: () => Promise<void>,
): Promise<void> {
let count = 0
let resolveDone: (() => void) | undefined
const done = new Promise<void>(resolve => {
resolveDone = resolve
})
setOnEnqueue(() => {
count++
if (count >= expected) resolveDone?.()
})
await sendMessages()
await Promise.race([
done,
sleep(5_000).then(() => {
throw new Error(`Timed out waiting for ${expected} UDS enqueues`)
}),
])
setOnEnqueue(null)
}
afterEach(async () => {
setOnEnqueue(null)
drainInbox()
await stopUdsMessaging()
})
async function closeServer(server: ReturnType<typeof createServer>): Promise<void> {
await new Promise<void>(resolve => {
server.close(() => resolve())
})
}
describe('UDS inbox retention', () => {
test('drainInbox returns each pending socket message once', async () => {
const path = socketPath('drain')
await startUdsMessaging(path, { isExplicit: true })
expect(process.env.CLAUDE_CODE_MESSAGING_TOKEN).toBeUndefined()
await waitForEnqueues(2, async () => {
await sendUdsMessage(path, { type: 'text', data: 'one' })
await sendUdsMessage(path, { type: 'text', data: 'two' })
})
const drained = drainInbox()
expect(drained.map(entry => entry.message.data)).toEqual(['one', 'two'])
expect(drained.every(entry => entry.status === 'processed')).toBe(true)
expect(drainInbox()).toEqual([])
})
test('inbox is capped when messages arrive faster than they are drained', async () => {
const path = socketPath('cap')
await startUdsMessaging(path, { isExplicit: true })
await waitForEnqueues(MAX_UDS_INBOX_ENTRIES, async () => {
for (let i = 0; i < MAX_UDS_INBOX_ENTRIES; i++) {
await sendUdsMessage(path, { type: 'text', data: String(i) })
}
})
await expect(
sendUdsMessage(path, { type: 'text', data: 'overflow' }),
).rejects.toThrow('inbox full')
const drained = drainInbox()
expect(drained).toHaveLength(MAX_UDS_INBOX_ENTRIES)
expect(drained[0]?.message.data).toBe('0')
expect(drained.at(-1)?.message.data).toBe(String(MAX_UDS_INBOX_ENTRIES - 1))
})
test('inbox is capped by retained bytes before entry count', async () => {
const path = socketPath('byte-cap')
await startUdsMessaging(path, { isExplicit: true })
const payload = 'x'.repeat(32 * 1024)
let accepted = 0
for (;;) {
try {
await sendUdsMessage(path, { type: 'text', data: payload })
accepted++
if (accepted > MAX_UDS_INBOX_BYTES / payload.length + 20) {
throw new Error('byte cap was not enforced')
}
} catch (error) {
expect(error).toBeInstanceOf(Error)
expect((error as Error).message).toContain('inbox full')
break
}
}
const drained = drainInbox()
expect(drained.length).toBe(accepted)
expect(drained.length).toBeLessThan(MAX_UDS_INBOX_ENTRIES)
})
test('ping replies with pong without enqueueing inbox work', async () => {
const path = socketPath('ping')
await startUdsMessaging(path, { isExplicit: true })
await sendUdsMessage(path, { type: 'ping' })
expect(drainInbox()).toEqual([])
})
test('drained entries never expose the UDS auth token', async () => {
const path = socketPath('strip-token')
await startUdsMessaging(path, { isExplicit: true })
await waitForEnqueues(1, async () => {
await sendUdsMessage(path, {
type: 'notification',
meta: { keep: 'visible' },
})
})
const drained = drainInbox()
expect(drained).toHaveLength(1)
expect(drained[0]?.message.meta).toEqual({ keep: 'visible' })
expect(drained[0]?.message.meta).not.toHaveProperty('authToken')
})
test('rejects unauthenticated socket messages', async () => {
const path = socketPath('auth')
await startUdsMessaging(path, { isExplicit: true })
const response = await new Promise<string>((resolve, reject) => {
const conn = createConnection(path, () => {
conn.write(`${JSON.stringify({ type: 'text', data: 'bad' })}\n`)
})
conn.setTimeout(5_000, () => {
conn.destroy()
reject(new Error('Timed out waiting for auth rejection'))
})
conn.on('data', chunk => {
const text = chunk.toString('utf-8')
if (text.includes('\n')) {
conn.end()
resolve(text)
}
})
conn.on('error', reject)
})
expect(JSON.parse(response).type).toBe('error')
expect(drainInbox()).toEqual([])
})
test('destroys oversized frames before enqueueing inbox work', async () => {
const path = socketPath('oversized')
await startUdsMessaging(path, { isExplicit: true })
await new Promise<void>((resolve, reject) => {
const conn = createConnection(path, () => {
conn.write('x'.repeat(MAX_UDS_FRAME_BYTES + 1))
})
conn.setTimeout(5_000, () => {
conn.destroy()
reject(new Error('Timed out waiting for oversized frame close'))
})
conn.on('close', () => resolve())
conn.on('error', () => resolve())
})
expect(drainInbox()).toEqual([])
})
test('rejects oversized receiver responses before retaining them', async () => {
const path = socketPath('oversized-response')
if (process.platform !== 'win32') {
await mkdir(dirname(path), { recursive: true })
}
const receiver = createServer(socket => {
socket.on('data', () => {
socket.write('x'.repeat(MAX_UDS_FRAME_BYTES + 1))
})
})
await new Promise<void>((resolve, reject) => {
receiver.on('error', reject)
receiver.listen(path, () => resolve())
})
try {
await expect(
sendUdsMessage(
path,
{ type: 'text', data: 'hello' },
{ authToken: 'test-token' },
),
).rejects.toThrow('UDS response frame exceeded size limit')
} finally {
await closeServer(receiver)
if (process.platform !== 'win32') {
await unlink(path).catch(() => undefined)
}
}
})
test('rejects inline auth token UDS targets instead of parsing them', async () => {
const path = socketPath('inline-token')
const targetWithToken = `${path}#token=secret`
expect(() => parseUdsTarget(targetWithToken)).toThrow('inline auth token')
try {
parseUdsTarget(targetWithToken)
} catch (error) {
expect((error as Error).message).not.toContain('secret')
}
const { sendToUdsSocket } = await import('../udsClient.js')
await expect(sendToUdsSocket(targetWithToken, 'hello')).rejects.toThrow(
'inline auth token',
)
})
if (process.platform !== 'win32') {
test('creates the listening socket with owner-only permissions', async () => {
const path = socketPath('socket-mode')
await startUdsMessaging(path, { isExplicit: true })
const mode = (await stat(path)).mode & 0o777
expect(mode).toBe(0o600)
})
test('fails closed when the capability directory is not private', async () => {
const previousConfigDir = process.env.CLAUDE_CONFIG_DIR
const tempHome = join(
tmpdir(),
`uds-capability-${process.pid}-${Date.now()}-${Math.random().toString(16).slice(2)}`,
)
process.env.CLAUDE_CONFIG_DIR = tempHome
const capabilityDir = join(tempHome, 'messaging-capabilities')
await mkdir(capabilityDir, { recursive: true, mode: 0o755 })
await chmod(capabilityDir, 0o755)
try {
await expect(
startUdsMessaging(socketPath('broad-capdir'), { isExplicit: true }),
).rejects.toThrow('permissions are too broad')
} finally {
if (previousConfigDir === undefined) {
delete process.env.CLAUDE_CONFIG_DIR
} else {
process.env.CLAUDE_CONFIG_DIR = previousConfigDir
}
await rm(tempHome, { recursive: true, force: true })
}
})
test('fails closed when the capability directory is a symlink', async () => {
const previousConfigDir = process.env.CLAUDE_CONFIG_DIR
const tempHome = join(
tmpdir(),
`uds-capability-link-${process.pid}-${Date.now()}-${Math.random().toString(16).slice(2)}`,
)
const target = join(tempHome, 'target')
process.env.CLAUDE_CONFIG_DIR = tempHome
await mkdir(target, { recursive: true, mode: 0o700 })
await symlink(target, join(tempHome, 'messaging-capabilities'), 'dir')
try {
await expect(
startUdsMessaging(socketPath('symlink-capdir'), { isExplicit: true }),
).rejects.toThrow('not a private directory')
} finally {
if (previousConfigDir === undefined) {
delete process.env.CLAUDE_CONFIG_DIR
} else {
process.env.CLAUDE_CONFIG_DIR = previousConfigDir
}
await rm(tempHome, { recursive: true, force: true })
}
})
}
})