Skip to content

Commit ee0d788

Browse files
author
unraid
committed
fix: harden bounded agent communication review fixes
CodeRabbit and Codecov surfaced real gaps in UDS framing, peer discovery, mailbox retention, and summary context coverage. This tightens those paths without suppressing review or coverage signals. Constraint: PR #369 must address CodeRabbit and Codecov findings without warning suppression or fake fallbacks Rejected: Suppress Codecov or CodeRabbit warnings | leaves real receive-path and test-isolation gaps Rejected: Add unreachable feature-gated tests | bun:bundle keeps those branches compile-time gated in local tests Confidence: high Scope-risk: moderate Directive: Keep UDS auth-token rejection outside feature flags; do not reintroduce inline token fallbacks Tested: bun test --coverage --coverage-reporter lcov --coverage-dir coverage; bun run test:all; bun run lint; bun run build; bun run build:vite; bun audit; git diff --cached --check Not-tested: Remote Codecov/CodeRabbit refreshed reports until pushed
1 parent f353eb0 commit ee0d788

15 files changed

Lines changed: 651 additions & 147 deletions

File tree

packages/builtin-tools/src/tools/ListPeersTool/ListPeersTool.ts

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,19 +84,27 @@ Use this tool to discover messaging targets before sending cross-session message
8484
// UDS socket directory. The implementation scans for live sockets
8585
// and optionally includes Remote Control bridge peers.
8686
const peers: PeerInfo[] = []
87+
const seen = new Set<string>()
88+
const addPeer = (peer: PeerInfo): void => {
89+
if (seen.has(peer.address)) return
90+
seen.add(peer.address)
91+
peers.push(peer)
92+
}
8793

8894
/* eslint-disable @typescript-eslint/no-require-imports */
8995
const udsMessaging =
9096
require('src/utils/udsMessaging.js') as typeof import('src/utils/udsMessaging.js')
9197
const udsClient =
9298
require('src/utils/udsClient.js') as typeof import('src/utils/udsClient.js')
99+
const bridgePeers =
100+
require('src/bridge/peerSessions.js') as typeof import('src/bridge/peerSessions.js')
93101
/* eslint-enable @typescript-eslint/no-require-imports */
94102

95103
const messagingSocketPath = udsMessaging.getUdsMessagingSocketPath()
96104
if (messagingSocketPath) {
97105
// Self entry for reference
98106
if (_input.include_self) {
99-
peers.push({
107+
addPeer({
100108
address: udsMessaging.formatUdsAddress(messagingSocketPath),
101109
name: 'self',
102110
pid: process.pid,
@@ -106,14 +114,18 @@ Use this tool to discover messaging targets before sending cross-session message
106114

107115
for (const peer of await udsClient.listPeers()) {
108116
if (!peer.messagingSocketPath) continue
109-
peers.push({
117+
addPeer({
110118
address: udsMessaging.formatUdsAddress(peer.messagingSocketPath),
111119
name: peer.name ?? peer.kind,
112120
cwd: peer.cwd,
113121
pid: peer.pid,
114122
})
115123
}
116124

125+
for (const peer of await bridgePeers.listBridgePeers()) {
126+
addPeer(peer)
127+
}
128+
117129
return {
118130
data: { peers },
119131
}

packages/builtin-tools/src/tools/SendMessageTool/SendMessageTool.ts

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -672,17 +672,15 @@ export const SendMessageTool: Tool<InputSchema, SendMessageToolOutput> =
672672
errorCode: 9,
673673
}
674674
}
675-
if (feature('UDS_INBOX')) {
676-
if (
677-
addr.scheme === 'uds' &&
678-
(hasInlineUdsToken(input.to) || wasInlineUdsTokenRejected(input))
679-
) {
680-
return {
681-
result: false,
682-
message:
683-
'uds addresses must not include inline auth tokens; use the ListPeers address',
684-
errorCode: 9,
685-
}
675+
if (
676+
addr.scheme === 'uds' &&
677+
(hasInlineUdsToken(input.to) || wasInlineUdsTokenRejected(input))
678+
) {
679+
return {
680+
result: false,
681+
message:
682+
'uds addresses must not include inline auth tokens; use the ListPeers address',
683+
errorCode: 9,
686684
}
687685
}
688686
if (input.to.includes('@')) {
@@ -808,6 +806,22 @@ export const SendMessageTool: Tool<InputSchema, SendMessageToolOutput> =
808806
},
809807

810808
async call(input, context, canUseTool, assistantMessage) {
809+
if (typeof input.message === 'string') {
810+
const addr = parseAddress(input.to)
811+
if (
812+
addr.scheme === 'uds' &&
813+
(hasInlineUdsToken(input.to) || wasInlineUdsTokenRejected(input))
814+
) {
815+
return {
816+
data: {
817+
success: false,
818+
message:
819+
'uds addresses must not include inline auth tokens; use the ListPeers address',
820+
},
821+
}
822+
}
823+
}
824+
811825
if (feature('UDS_INBOX') && typeof input.message === 'string') {
812826
const addr = parseAddress(input.to)
813827
if (addr.scheme === 'bridge') {
@@ -827,10 +841,10 @@ export const SendMessageTool: Tool<InputSchema, SendMessageToolOutput> =
827841
const { postInterClaudeMessage } =
828842
require('src/bridge/peerSessions.js') as typeof import('src/bridge/peerSessions.js')
829843
/* eslint-enable @typescript-eslint/no-require-imports */
830-
const result = await postInterClaudeMessage(
844+
const result = (await postInterClaudeMessage(
831845
addr.target,
832846
input.message,
833-
) as { ok: boolean; error?: string }
847+
)) as { ok: boolean; error?: string }
834848
const preview = input.summary || truncate(input.message, 50)
835849
return {
836850
data: {

packages/builtin-tools/src/tools/SendMessageTool/__tests__/udsRecipientSanitization.test.ts

Lines changed: 77 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,4 @@
1-
import { describe, expect, mock, test } from 'bun:test'
2-
3-
mock.module('bun:bundle', () => ({
4-
feature: (name: string) => name === 'UDS_INBOX',
5-
}))
1+
import { describe, expect, test } from 'bun:test'
62

73
describe('SendMessageTool UDS recipient handling', () => {
84
test('redacts inline UDS tokens before classifier and observable paths', async () => {
@@ -25,6 +21,62 @@ describe('SendMessageTool UDS recipient handling', () => {
2521
).toBe('to uds:/tmp/peer.sock: hello')
2622
})
2723

24+
test('keeps redacted UDS token rejection through observable backfill', async () => {
25+
const { SendMessageTool } = await import('../SendMessageTool.js')
26+
const observableInput = {
27+
to: 'uds:/tmp/peer.sock#token=secret-token',
28+
message: {
29+
type: 'plan_approval_response',
30+
request_id: 'req-1',
31+
approve: false,
32+
reason: 'needs tests',
33+
},
34+
} as Record<string, unknown>
35+
36+
SendMessageTool.backfillObservableInput!(observableInput)
37+
38+
expect(observableInput.to).toBe('uds:/tmp/peer.sock')
39+
expect(observableInput.recipient).toBe('uds:/tmp/peer.sock')
40+
expect(observableInput.type).toBe('plan_approval_response')
41+
expect(observableInput.request_id).toBe('req-1')
42+
expect(observableInput.approve).toBe(false)
43+
expect(observableInput.content).toBe('needs tests')
44+
expect(JSON.stringify(observableInput)).not.toContain('secret-token')
45+
46+
const result = await SendMessageTool.validateInput!(
47+
observableInput as never,
48+
{} as never,
49+
)
50+
51+
expect(result.result).toBe(false)
52+
if (result.result !== false) {
53+
throw new Error('expected validation to reject redacted inline UDS token')
54+
}
55+
expect(result.message).toContain('inline auth tokens')
56+
})
57+
58+
test('redacts UDS tokens in structured classifier text', async () => {
59+
const { SendMessageTool } = await import('../SendMessageTool.js')
60+
const to = 'uds:/tmp/peer.sock#token=secret-token'
61+
62+
expect(
63+
SendMessageTool.toAutoClassifierInput({
64+
to,
65+
message: { type: 'shutdown_request' },
66+
}),
67+
).toBe('shutdown_request to uds:/tmp/peer.sock')
68+
expect(
69+
SendMessageTool.toAutoClassifierInput({
70+
to,
71+
message: {
72+
type: 'plan_approval_response',
73+
request_id: 'req-1',
74+
approve: true,
75+
},
76+
}),
77+
).toBe('plan_approval approve to uds:/tmp/peer.sock')
78+
})
79+
2880
test('rejects inline UDS tokens during validation', async () => {
2981
const { SendMessageTool } = await import('../SendMessageTool.js')
3082
const result = await SendMessageTool.validateInput!(
@@ -36,6 +88,26 @@ describe('SendMessageTool UDS recipient handling', () => {
3688
)
3789

3890
expect(result.result).toBe(false)
91+
if (result.result !== false) {
92+
throw new Error('expected validation to reject inline UDS token')
93+
}
94+
expect(result.message).toContain('inline auth tokens')
95+
expect(JSON.stringify(result)).not.toContain('secret-token')
96+
})
97+
98+
test('rejects inline UDS tokens during execution without leaking them', async () => {
99+
const { SendMessageTool } = await import('../SendMessageTool.js')
100+
const result = await SendMessageTool.call(
101+
{
102+
to: 'uds:/tmp/peer.sock#token=secret-token',
103+
message: 'hello',
104+
},
105+
{} as never,
106+
undefined as never,
107+
undefined as never,
108+
)
109+
110+
expect(result.data.success).toBe(false)
39111
expect(JSON.stringify(result)).not.toContain('secret-token')
40112
})
41113
})

src/bridge/peerSessions.ts

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,38 @@ import { getBridgeAccessToken } from './bridgeConfig.js'
66
import { getReplBridgeHandle } from './replBridgeHandle.js'
77
import { toCompatSessionId } from './sessionIdCompat.js'
88

9+
export type BridgePeerSession = {
10+
address: string
11+
name?: string
12+
cwd?: string
13+
pid?: number
14+
}
15+
16+
/**
17+
* List locally registered sessions that have published a Remote Control
18+
* session ID. The PID registry is the local source of truth for bridge peers
19+
* already known to this machine; SendMessage can use these bridge:<id>
20+
* addresses when the current process has an active bridge handle.
21+
*/
22+
export async function listBridgePeers(): Promise<BridgePeerSession[]> {
23+
const { listAllLiveSessions } = await import('../utils/udsClient.js')
24+
const sessions = await listAllLiveSessions()
25+
const peers: BridgePeerSession[] = []
26+
27+
for (const session of sessions) {
28+
if (session.pid === process.pid || !session.bridgeSessionId) continue
29+
const compatId = toCompatSessionId(session.bridgeSessionId)
30+
peers.push({
31+
address: `bridge:${compatId}`,
32+
name: session.name ?? session.kind,
33+
cwd: session.cwd,
34+
pid: session.pid,
35+
})
36+
}
37+
38+
return peers
39+
}
40+
941
/**
1042
* Send a plain-text message to another Claude session via the bridge API.
1143
*

src/cli/print.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2773,7 +2773,7 @@ function runHeadlessStreaming(
27732773
const value =
27742774
typeof entry.message.data === 'string'
27752775
? entry.message.data
2776-
: jsonStringify(entry.message)
2776+
: jsonStringify(entry.message.data)
27772777
enqueue({
27782778
mode: 'prompt',
27792779
value,
Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
import {
2+
afterAll,
3+
afterEach,
4+
beforeEach,
5+
describe,
6+
expect,
7+
mock,
8+
test,
9+
} from 'bun:test'
10+
import { debugMock } from '../../../../tests/mocks/debug'
11+
import { logMock } from '../../../../tests/mocks/log'
12+
import { asAgentId } from '../../../types/ids.js'
13+
import type { CacheSafeParams } from '../../../utils/forkedAgent.js'
14+
15+
const transcriptMessages = [
16+
{ type: 'user', message: { content: 'start' }, uuid: 'u1' },
17+
{
18+
type: 'assistant',
19+
message: { content: [{ type: 'text', text: 'working' }] },
20+
uuid: 'a1',
21+
},
22+
{ type: 'user', message: { content: 'continue' }, uuid: 'u2' },
23+
]
24+
25+
let poorModeActive = false
26+
let forkCalls = 0
27+
let updateCalls: Array<{ taskId: string; summary: string }> = []
28+
let transcript = { messages: transcriptMessages }
29+
const sessionStorageSnapshot = {
30+
...(require('../../../utils/sessionStorage.ts') as Record<string, unknown>),
31+
}
32+
33+
mock.module('src/commands/poor/poorMode.js', () => ({
34+
isPoorModeActive: () => poorModeActive,
35+
}))
36+
37+
mock.module('src/tasks/LocalAgentTask/LocalAgentTask.js', () => ({
38+
updateAgentSummary: (taskId: string, summary: string) => {
39+
updateCalls.push({ taskId, summary })
40+
},
41+
}))
42+
43+
mock.module(
44+
'@claude-code-best/builtin-tools/tools/AgentTool/runAgent.js',
45+
() => ({
46+
filterIncompleteToolCalls: <T>(messages: T) => messages,
47+
}),
48+
)
49+
50+
mock.module('src/utils/debug.js', debugMock)
51+
mock.module('src/utils/log.js', logMock)
52+
53+
mock.module('src/utils/forkedAgent.js', () => ({
54+
runForkedAgent: async () => {
55+
forkCalls += 1
56+
return {
57+
messages: [
58+
{
59+
type: 'assistant',
60+
message: {
61+
content: [{ type: 'text', text: 'Reading udsClient.ts' }],
62+
},
63+
},
64+
],
65+
}
66+
},
67+
}))
68+
69+
mock.module('src/utils/sessionStorage.js', () => ({
70+
...sessionStorageSnapshot,
71+
getAgentTranscript: async () => transcript,
72+
}))
73+
74+
afterAll(() => {
75+
mock.module('src/utils/sessionStorage.js', () =>
76+
require('../../../utils/sessionStorage.ts'),
77+
)
78+
})
79+
80+
describe('startAgentSummarization', () => {
81+
const realSetTimeout = globalThis.setTimeout
82+
const realClearTimeout = globalThis.clearTimeout
83+
let scheduled:
84+
| ((...args: Parameters<TimerHandler & ((...args: unknown[]) => void)>) => void)
85+
| undefined
86+
87+
beforeEach(() => {
88+
poorModeActive = false
89+
forkCalls = 0
90+
updateCalls = []
91+
transcript = { messages: transcriptMessages }
92+
scheduled = undefined
93+
globalThis.setTimeout = ((callback: TimerHandler) => {
94+
scheduled = callback as (...args: unknown[]) => void
95+
return 1 as unknown as ReturnType<typeof setTimeout>
96+
}) as unknown as typeof setTimeout
97+
globalThis.clearTimeout = (() => undefined) as typeof clearTimeout
98+
})
99+
100+
afterEach(() => {
101+
globalThis.setTimeout = realSetTimeout
102+
globalThis.clearTimeout = realClearTimeout
103+
})
104+
105+
test('summarizes bounded transcript once and skips unchanged fingerprints', async () => {
106+
const { startAgentSummarization } = await import('../agentSummary.js')
107+
108+
const handle = startAgentSummarization(
109+
'task-1',
110+
asAgentId('a0000000000000000'),
111+
{
112+
forkContextMessages: [{ type: 'user', message: { content: 'old' } }],
113+
model: 'claude-test',
114+
} as unknown as CacheSafeParams,
115+
() => undefined,
116+
)
117+
118+
expect(typeof scheduled).toBe('function')
119+
await scheduled!()
120+
121+
expect(forkCalls).toBe(1)
122+
expect(updateCalls).toEqual([
123+
{ taskId: 'task-1', summary: 'Reading udsClient.ts' },
124+
])
125+
126+
await scheduled!()
127+
128+
expect(forkCalls).toBe(1)
129+
expect(updateCalls).toHaveLength(1)
130+
131+
handle.stop()
132+
})
133+
})

0 commit comments

Comments
 (0)