feat: add persistent workbar browser tabs and scoped MCP routing

This commit is contained in:
mesalogo
2026-09-10 01:26:56 +08:00
parent 428b4c0d08
commit 80b0b18712
54 changed files with 5446 additions and 1468 deletions
@@ -5,7 +5,6 @@ import { join, relative } from 'node:path'
import { afterEach, describe, expect, it, vi } from 'vitest'
import { LocalWorkspaceAccess } from '../workspace'
import {
DIRECT_MODEL_PROCESS_TRUNCATION_MARKER,
LocalDirectModelProcessService,
buildDirectModelProcessEnvironment,
type DirectModelProcessChild,
@@ -277,7 +276,7 @@ describe('LocalDirectModelProcessService', () => {
await service.dispose()
})
it('returns nonzero exits normally and bounds stdout and stderr by retaining both edges', async () => {
it('returns nonzero exits with bounded previews and retains full output for paging', async () => {
const { access } = await workspace()
const child = new FakeChild()
let now = 1_000
@@ -310,12 +309,19 @@ describe('LocalDirectModelProcessService', () => {
expect(result.durationMs).toBe(125)
expect(result.stdoutTruncated).toBe(true)
expect(result.stdout).toBe(
`${'A'.repeat(48 * 1024)}` +
DIRECT_MODEL_PROCESS_TRUNCATION_MARKER +
'Z'.repeat(48 * 1024)
'A'.repeat(48 * 1024) + 'M'.repeat(10) + 'Z'.repeat(48 * 1024 - 10)
)
expect(result.stderr).toBe('ordinary error output')
expect(result.stderrTruncated).toBe(false)
expect(result.stderrReference).toBeUndefined()
const reference = result.stdoutReference!
expect(reference.totalBytes).toBe(96 * 1024 + 10)
await expect(service.readOutput('output', reference.handle, reference.nextCursor)).resolves.toMatchObject({
content: 'Z'.repeat(10), eof: true, nextCursor: reference.totalBytes
})
await expect(service.readOutput('other', reference.handle)).rejects.toThrow()
await service.releaseConversation('output')
await expect(service.readOutput('output', reference.handle)).rejects.toThrow()
await service.dispose()
})
@@ -439,14 +445,18 @@ describe('LocalDirectModelProcessService', () => {
}
)
const first = run('first')
await vi.waitFor(() => expect(spawnProcess).toHaveBeenCalledOnce())
const second = run('second')
await vi.waitFor(() => expect(spawnProcess).toHaveBeenCalledTimes(2))
firstChild.stdout.emit('data', 'x'.repeat(120000))
secondChild.stderr.emit('data', 'y'.repeat(120000))
await service.releaseConversation('first')
expect(terminated).toHaveLength(1)
await expect(first).resolves.toMatchObject({
terminationReason: 'cancelled'
})
await expect(service.readOutput('first', (await first).stdoutReference!.handle)).rejects.toThrow()
const replacementChild = new FakeChild()
children.push(replacementChild)
const replacement = run('first')
@@ -464,6 +474,7 @@ describe('LocalDirectModelProcessService', () => {
await expect(second).resolves.toMatchObject({
terminationReason: 'cancelled'
})
await expect(service.readOutput('second', (await second).stderrReference!.handle)).rejects.toThrow()
await service.dispose()
})
@@ -496,4 +507,37 @@ describe('LocalDirectModelProcessService', () => {
expect(result.stderr).toContain('warn')
await service.dispose()
})
it('pages complete stdout and stderr from a real shell', async () => {
const { access } = await workspace()
const service = new LocalDirectModelProcessService()
try {
expect((await service.getCapability()).available).toBe(true)
const command = process.platform === 'win32'
? "[Console]::Out.Write(('x' * 120000) + 'stdout-end'); [Console]::Error.Write(('y' * 120000) + 'stderr-end')"
: "printf '%120000sstdout-end' '' | tr ' ' x; printf '%120000sstderr-end' '' | tr ' ' y >&2"
const result = await service.execute({ command }, {
conversationId: 'real-pages', workspace: access,
signal: new AbortController().signal
})
for (const [preview, reference, expected] of [
[result.stdout, result.stdoutReference, 'x'.repeat(120000) + 'stdout-end'],
[result.stderr, result.stderrReference, 'y'.repeat(120000) + 'stderr-end']
] as const) {
expect(reference).toBeDefined()
let output = preview
let cursor = reference!.nextCursor
while (cursor < reference!.totalBytes) {
const page = await service.readOutput('real-pages', reference!.handle, cursor, 4096)
output += page.content
cursor = page.nextCursor
}
expect(output).toBe(expected)
}
await service.dispose()
await expect(service.readOutput('real-pages', result.stdoutReference!.handle)).rejects.toThrow()
} finally {
await service.dispose()
}
})
})
+64 -96
View File
@@ -15,16 +15,18 @@ import {
type WaitableProcessTreeChild
} from './child-process-termination'
import { buildCredentialFilteredUserEnvironment } from './process-environment'
import {
PagedOutputStore,
type PagedOutputPage,
type PagedOutputReference,
type PagedOutputWriter
} from './paged-output-store'
const DEFAULT_TIMEOUT_MS = 120_000
const MAXIMUM_TIMEOUT_MS = 10 * 60_000
const MAXIMUM_OUTPUT_BYTES = 96 * 1024
const RETAINED_OUTPUT_EDGE_BYTES = MAXIMUM_OUTPUT_BYTES / 2
const TERMINATION_WAIT_MS = 2_000
export const DIRECT_MODEL_PROCESS_TRUNCATION_MARKER =
'\n...[GoodBuddy output truncated]...\n'
export const processExecuteInputSchema = z
.object({
command: z.string().trim().min(1).max(100_000),
@@ -70,6 +72,8 @@ export type ProcessExecuteResult = {
stderr: string
stdoutTruncated: boolean
stderrTruncated: boolean
stdoutReference?: PagedOutputReference
stderrReference?: PagedOutputReference
}
export type DirectModelProcessCapability =
@@ -94,6 +98,7 @@ export interface DirectModelProcessService {
input: ProcessExecuteInput,
context: DirectModelProcessExecutionContext
): Promise<ProcessExecuteResult>
readOutput(ownerId: string, handle: string, cursor?: number, limitBytes?: number): Promise<PagedOutputPage>
releaseConversation(conversationId: string): Promise<void>
dispose(): Promise<void>
}
@@ -271,67 +276,6 @@ async function resolveProcessShell(
return undefined
}
class BoundedProcessOutput {
private prefix = Buffer.alloc(0)
private suffix = Buffer.alloc(0)
private totalBytes = 0
append(chunk: Buffer | string): void {
const value = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)
this.totalBytes += value.byteLength
if (this.prefix.byteLength < RETAINED_OUTPUT_EDGE_BYTES) {
const remaining =
RETAINED_OUTPUT_EDGE_BYTES - this.prefix.byteLength
this.prefix = Buffer.concat([
this.prefix,
value.subarray(0, remaining)
])
}
if (value.byteLength >= RETAINED_OUTPUT_EDGE_BYTES) {
this.suffix = Buffer.from(
value.subarray(value.byteLength - RETAINED_OUTPUT_EDGE_BYTES)
)
return
}
const combined = Buffer.concat([this.suffix, value])
this.suffix =
combined.byteLength <= RETAINED_OUTPUT_EDGE_BYTES
? combined
: Buffer.from(
combined.subarray(
combined.byteLength - RETAINED_OUTPUT_EDGE_BYTES
)
)
}
result(): { text: string; truncated: boolean } {
if (this.totalBytes > MAXIMUM_OUTPUT_BYTES) {
return {
text: Buffer.concat([
this.prefix,
Buffer.from(DIRECT_MODEL_PROCESS_TRUNCATION_MARKER),
this.suffix
]).toString('utf8'),
truncated: true
}
}
const overlap = Math.max(
0,
this.prefix.byteLength +
this.suffix.byteLength -
this.totalBytes
)
return {
text: Buffer.concat([
this.prefix,
this.suffix.subarray(overlap)
]).toString('utf8'),
truncated: false
}
}
}
function shellArguments(
shell: ResolvedProcessShell,
command: string
@@ -399,6 +343,7 @@ export class LocalDirectModelProcessService
implements DirectModelProcessService
{
private readonly activeCalls = new Map<string, Set<ActiveCall>>()
private readonly outputs = new PagedOutputStore('process')
private readonly shell: Promise<ResolvedProcessShell | undefined>
private disposed = false
private disposePromise?: Promise<void>
@@ -496,31 +441,46 @@ export class LocalDirectModelProcessService
platform
)
const startedAt = (this.options.now ?? Date.now)()
let child: DirectModelProcessChild
try {
child = (
this.options.spawnProcess ??
(spawn as unknown as DirectModelProcessSpawn)
)(shell.executable, shellArguments(shell, input.command), {
cwd: workingDirectory.canonicalPath,
detached: platform !== 'win32',
env: environment,
shell: false,
stdio: ['ignore', 'pipe', 'pipe'],
windowsHide: true
})
} catch (error) {
throw new Error('无法启动命令 Shell', { cause: error })
}
return await this.waitForProcess(
child,
shell,
workingDirectory.displayPath,
input.timeoutMs,
signal,
startedAt
const stdout = await this.outputs.create(
context.conversationId, MAXIMUM_OUTPUT_BYTES
)
let stderr: PagedOutputWriter | undefined
try {
stderr = await this.outputs.create(
context.conversationId, MAXIMUM_OUTPUT_BYTES
)
signal.throwIfAborted()
let child: DirectModelProcessChild
try {
child = (
this.options.spawnProcess ??
(spawn as unknown as DirectModelProcessSpawn)
)(shell.executable, shellArguments(shell, input.command), {
cwd: workingDirectory.canonicalPath,
detached: platform !== 'win32',
env: environment,
shell: false,
stdio: ['ignore', 'pipe', 'pipe'],
windowsHide: true
})
} catch (error) {
throw new Error('无法启动命令 Shell', { cause: error })
}
return await this.waitForProcess(
child,
shell,
workingDirectory.displayPath,
input.timeoutMs,
signal,
startedAt,
stdout,
stderr
)
} catch (error) {
await Promise.all([stdout.abort(), stderr?.abort()])
throw error
}
}
private waitForProcess(
@@ -529,10 +489,10 @@ export class LocalDirectModelProcessService
cwd: string,
timeoutMs: number,
cancellationSignal: AbortSignal,
startedAt: number
startedAt: number,
stdout: PagedOutputWriter,
stderr: PagedOutputWriter
): Promise<ProcessExecuteResult> {
const stdout = new BoundedProcessOutput()
const stderr = new BoundedProcessOutput()
child.stdout?.on('data', (chunk) => stdout.append(chunk))
child.stderr?.on('data', (chunk) => stderr.append(chunk))
@@ -550,9 +510,9 @@ export class LocalDirectModelProcessService
child.stdout?.removeListener?.('error', failFromOutput)
child.stderr?.removeListener?.('error', failFromOutput)
}
const result = (): ProcessExecuteResult => {
const stdoutResult = stdout.result()
const stderrResult = stderr.result()
const result = async (): Promise<ProcessExecuteResult> => {
const stdoutResult = await stdout.finish()
const stderrResult = await stderr.finish()
return {
shell: { kind: shell.kind, label: shell.label },
cwd,
@@ -566,7 +526,9 @@ export class LocalDirectModelProcessService
stdout: stdoutResult.text,
stderr: stderrResult.text,
stdoutTruncated: stdoutResult.truncated,
stderrTruncated: stderrResult.truncated
stderrTruncated: stderrResult.truncated,
...(stdoutResult.reference ? { stdoutReference: stdoutResult.reference } : {}),
...(stderrResult.reference ? { stderrReference: stderrResult.reference } : {})
}
}
const finishResult = (): void => {
@@ -682,6 +644,11 @@ export class LocalDirectModelProcessService
call.controller.abort(abortError('对话进程资源已释放'))
}
await Promise.all(calls.map((call) => call.done))
await this.outputs.releaseOwner(conversationId)
}
readOutput(ownerId: string, handle: string, cursor?: number, limitBytes?: number): Promise<PagedOutputPage> {
return this.outputs.read(ownerId, handle, cursor, limitBytes)
}
dispose(): Promise<void> {
@@ -698,5 +665,6 @@ export class LocalDirectModelProcessService
call.controller.abort(abortError('进程执行服务已关闭'))
}
await Promise.all(calls.map((call) => call.done))
await this.outputs.dispose()
}
}
+108 -3
View File
@@ -22,6 +22,7 @@ import {
KnowledgeMcpGateway,
type MagicNotesDatabase
} from './knowledge-mcp-gateway'
import { browserTabIdSchema } from '../../shared/contracts'
const firstLibraryId = '11111111-1111-4111-8111-111111111111'
const secondLibraryId = '22222222-2222-4222-8222-222222222222'
@@ -185,6 +186,28 @@ function testTool(name: string): Tool {
}
const gateways: KnowledgeMcpGateway[] = []
it('honors a requested knowledge result count above eight', async () => {
const { service, searchHybridMany } = createService()
const [sample] = await searchHybridMany([firstLibraryId])
searchHybridMany.mockClear()
searchHybridMany.mockResolvedValue(Array.from({ length: 12 }, (_, index) => ({
...sample!,
result: {
...sample!.result,
chunk: { ...sample!.result.chunk, id: crypto.randomUUID() },
rank: index + 1
}
})))
const gateway = new KnowledgeMcpGateway(service)
gateways.push(gateway)
const controller = new AbortController()
const token = gateway.grant('large-search', [firstLibraryId], controller.signal)!
const references = await gateway.search(token, { query: 'evidence', limit: 12 })
expect(references).toHaveLength(12)
expect(searchHybridMany).toHaveBeenCalledWith([firstLibraryId], 'evidence', 12, controller.signal)
await expect(gateway.search(token, { query: 'evidence', limit: 101 })).rejects.toThrow()
})
const databases: AssistantDatabase[] = []
const temporaryDirectories: string[] = []
const httpServers: Server[] = []
@@ -210,6 +233,22 @@ afterEach(async () => {
describe('KnowledgeMcpGateway', () => {
it('exposes the assigned browser through a request-scoped conversation', async () => {
const { service } = createService()
const browserTabId = browserTabIdSchema.parse(
'00000000-0000-4000-8000-000000000301'
)
const primaryBrowserTabId = browserTabIdSchema.parse(
'00000000-0000-4000-8000-000000000302'
)
const releaseBrowserUsage = vi.fn()
const acquireTabUsage = vi.fn(
(conversationId: string, tabId: typeof browserTabId, owner: string) => ({
conversationId,
tabId,
owner,
signal: new AbortController().signal,
release: releaseBrowserUsage
})
)
const browserService = {
getOrigin: vi.fn(() => undefined),
navigate: vi.fn(async () => ({
@@ -222,21 +261,63 @@ describe('KnowledgeMcpGateway', () => {
select: vi.fn(),
back: vi.fn(),
screenshot: vi.fn(),
releaseConversation: vi.fn(async () => undefined)
releaseConversation: vi.fn(async () => undefined),
acquireTabUsage,
listTabs: vi.fn(() => [
{
conversationId: 'browser-conversation',
tabId: primaryBrowserTabId,
primary: true,
status: 'ready' as const,
isLoading: false,
canGoBack: false,
createdAt: 1,
updatedAt: 1
},
{
conversationId: 'browser-conversation',
tabId: browserTabId,
primary: false,
status: 'ready' as const,
isLoading: false,
canGoBack: false,
createdAt: 1,
updatedAt: 1
}
])
}
const gateway = new KnowledgeMcpGateway(service, {
browserService
})
gateways.push(gateway)
await gateway.start()
expect(() =>
gateway.grant(
'wrong-browser-tab',
[],
new AbortController().signal,
'none',
undefined,
'browser-conversation',
browserTabIdSchema.parse(
'00000000-0000-4000-8000-000000000399'
)
)
).toThrow('浏览器标签页不存在或不属于当前对话')
const token = gateway.grant(
'browser-request',
[],
new AbortController().signal,
'none',
undefined,
'browser-conversation'
'browser-conversation',
browserTabId
)!
expect(acquireTabUsage).toHaveBeenCalledWith(
'browser-conversation',
browserTabId,
'browser-request'
)
expect(gateway.getAvailableToolNames(token)).toEqual(browserToolNames)
const client = new Client({
name: 'browser-loopback-test',
@@ -276,8 +357,32 @@ describe('KnowledgeMcpGateway', () => {
expect(browserService.navigate).toHaveBeenCalledWith(
'browser-conversation',
'https://example.com',
expect.any(AbortSignal)
expect.any(AbortSignal),
browserTabId
)
expect(browserService.listTabs).toHaveBeenCalledTimes(2)
vi.mocked(browserService.click).mockRejectedValueOnce(
new Error('浏览器标签页不存在或不属于当前对话')
)
await expect(
client.callTool({
name: 'browser_click',
arguments: { ref: 'b_boundTabReference' }
})
).resolves.toMatchObject({ isError: true })
expect(browserService.click).toHaveBeenCalledWith(
'browser-conversation',
'b_boundTabReference',
expect.any(AbortSignal),
browserTabId
)
expect(browserService.listTabs).toHaveBeenCalledTimes(2)
gateway.revoke(token)
gateway.revoke(token)
expect(releaseBrowserUsage).toHaveBeenCalledOnce()
await expect(
client.callTool({ name: 'browser_snapshot', arguments: {} })
).rejects.toThrow()
} finally {
await client.close()
}
+104 -25
View File
@@ -19,7 +19,11 @@ import {
type Tool
} from '@modelcontextprotocol/sdk/types.js'
import { z } from 'zod'
import type { KnowledgeSearchReference } from '../../shared/contracts'
import type {
BrowserTabId,
BrowserTabSummary,
KnowledgeSearchReference
} from '../../shared/contracts'
import { stripKnowledgeHighlightTags } from '../../shared/knowledge-text'
import {
knowledgeToolNames,
@@ -70,6 +74,7 @@ import {
type BrowserToolName,
type BrowserToolService
} from '../browser/browser-model-tools'
import type { BrowserTabUsageLease } from '../browser/browser-service'
import { BrowserStaleReferenceError } from '../browser/cdp-browser-driver'
import { safeToolErrorDetail } from './approval-summary'
@@ -185,6 +190,8 @@ type Capability = {
configWorkspacePath?: string
authorizeConfigApply?: GoodBuddyConfigApplyAuthorizer
browserConversationId?: string
browserTabId?: BrowserTabId
browserUsageLease?: BrowserTabUsageLease
expiresAt: number
signal: AbortSignal
brokerController: AbortController
@@ -230,7 +237,14 @@ export type KnowledgeMcpGatewayOptions = {
now?: () => number
magicNotesDatabase?: MagicNotesDatabase
configService?: GoodBuddyConfigService
browserService?: BrowserToolService
browserService?: BrowserToolService & {
listTabs(conversationId: string): BrowserTabSummary[]
acquireTabUsage(
conversationId: string,
tabId: BrowserTabId,
owner: string
): BrowserTabUsageLease
}
launchEnvironmentProvider?: LaunchEnvironmentProvider
}
@@ -356,7 +370,14 @@ export class KnowledgeMcpGateway {
private readonly maximumBodyBytes: number
private readonly magicNotesDatabase?: MagicNotesDatabase
private readonly configService?: GoodBuddyConfigService
private readonly browserService?: BrowserToolService
private readonly browserService?: BrowserToolService & {
listTabs(conversationId: string): BrowserTabSummary[]
acquireTabUsage(
conversationId: string,
tabId: BrowserTabId,
owner: string
): BrowserTabUsageLease
}
private readonly launchEnvironmentProvider?: LaunchEnvironmentProvider
private server?: Server
private endpoint?: string
@@ -432,7 +453,9 @@ export class KnowledgeMcpGateway {
workspacePath: string
authorizeApply?: GoodBuddyConfigApplyAuthorizer
},
browserConversationId?: string
browserConversationId?: string,
browserTabId?: BrowserTabId,
browserUsageLease?: BrowserTabUsageLease
): string | undefined {
const effectiveMagicNotesAccess = this.magicNotesDatabase
? magicNotesAccess
@@ -443,6 +466,39 @@ export class KnowledgeMcpGateway {
const effectiveBrowserConversationId = this.browserService
? browserConversationId
: undefined
const effectiveBrowserTabId = effectiveBrowserConversationId
? browserTabId
: undefined
if (
effectiveBrowserConversationId &&
(!effectiveBrowserTabId ||
!this.browserService!
.listTabs(effectiveBrowserConversationId)
.some((tab) => tab.tabId === effectiveBrowserTabId))
) {
throw new Error('浏览器标签页不存在或不属于当前对话')
}
let effectiveBrowserUsageLease: BrowserTabUsageLease | undefined
if (effectiveBrowserConversationId && effectiveBrowserTabId) {
effectiveBrowserUsageLease = browserUsageLease ??
this.browserService!.acquireTabUsage(
effectiveBrowserConversationId,
effectiveBrowserTabId,
requestId
)
if (
effectiveBrowserUsageLease.conversationId !==
effectiveBrowserConversationId ||
effectiveBrowserUsageLease.tabId !== effectiveBrowserTabId ||
effectiveBrowserUsageLease.owner !== requestId
) {
effectiveBrowserUsageLease.release()
throw new Error('浏览器标签页使用租约与当前请求不匹配')
}
} else if (browserUsageLease) {
browserUsageLease.release()
throw new Error('浏览器标签页使用租约与当前请求不匹配')
}
if (
authorizedLibraryIds.length === 0 &&
effectiveMagicNotesAccess === 'none' &&
@@ -452,21 +508,28 @@ export class KnowledgeMcpGateway {
return undefined
}
signal.throwIfAborted()
return this.storeCapability({
requestId,
libraryIds: Object.freeze([...new Set(authorizedLibraryIds)]),
magicNotesAccess: effectiveMagicNotesAccess,
configAccess: effectiveConfigAccess,
browserConversationId: effectiveBrowserConversationId,
...(effectiveConfigAccess !== 'none'
? {
configWorkspacePath: config?.workspacePath,
authorizeConfigApply: config?.authorizeApply
}
: {}),
signal,
customMcpServers: []
})
try {
return this.storeCapability({
requestId,
libraryIds: Object.freeze([...new Set(authorizedLibraryIds)]),
magicNotesAccess: effectiveMagicNotesAccess,
configAccess: effectiveConfigAccess,
browserConversationId: effectiveBrowserConversationId,
browserTabId: effectiveBrowserTabId,
browserUsageLease: effectiveBrowserUsageLease,
...(effectiveConfigAccess !== 'none'
? {
configWorkspacePath: config?.workspacePath,
authorizeConfigApply: config?.authorizeApply
}
: {}),
signal,
customMcpServers: []
})
} catch (error) {
effectiveBrowserUsageLease?.release()
throw error
}
}
grantCustomMcp(
@@ -517,14 +580,27 @@ export class KnowledgeMcpGateway {
const abort = (): void => {
this.revoke(token)
}
const abortFromBrowser = (): void => {
this.revoke(token)
}
value.signal.addEventListener('abort', abort, { once: true })
value.browserUsageLease?.signal.addEventListener(
'abort',
abortFromBrowser,
{ once: true }
)
this.capabilities.set(token, {
...value,
expiresAt: this.now() + this.capabilityTtlMs,
brokerController,
references: new Map(),
removeAbortListener: () =>
removeAbortListener: () => {
value.signal.removeEventListener('abort', abort)
value.browserUsageLease?.signal.removeEventListener(
'abort',
abortFromBrowser
)
}
})
return token
}
@@ -542,6 +618,7 @@ export class KnowledgeMcpGateway {
capability.brokerController.abort(
new Error('MCP capability was revoked')
)
capability.browserUsageLease?.release()
for (const session of this.downstreamMcpSessions.values()) {
if (session.token === token) {
void this.closeDownstreamMcpSession(session)
@@ -724,7 +801,7 @@ export class KnowledgeMcpGateway {
...(capability.configAccess === 'write'
? goodbuddyConfigWriteToolNames
: []),
...(capability.browserConversationId
...(capability.browserConversationId && capability.browserTabId
? browserToolNames
: [])
]
@@ -1518,10 +1595,11 @@ export class KnowledgeMcpGateway {
}
)
const capability = this.getCapability(token)
const browserTools = capability.browserConversationId
const browserTools = capability.browserConversationId && capability.browserTabId
? new BrowserModelTools({
service: this.browserService!,
conversationId: capability.browserConversationId
conversationId: capability.browserConversationId,
browserTabId: capability.browserTabId
})
: undefined
const browserDefinitions = browserTools
@@ -1579,10 +1657,11 @@ export class KnowledgeMcpGateway {
}
}
const capability = this.getCapability(token)
const browserTools = capability.browserConversationId
const browserTools = capability.browserConversationId && capability.browserTabId
? new BrowserModelTools({
service: this.browserService!,
conversationId: capability.browserConversationId
conversationId: capability.browserConversationId,
browserTabId: capability.browserTabId
})
: undefined
if (browserTools?.ownsTool(name)) {
+34
View File
@@ -18,6 +18,40 @@ const toolPng = Buffer.from([
0x0d, 0x0a, 0x1a, 0x0a
]).toString('base64')
it('preserves full supplied and cached history beyond message and byte caps', async () => {
const fetcher = vi.fn<typeof fetch>(async () => new Response(
createEventStream('reply'),
{ headers: { 'content-type': 'text/event-stream' } }
))
const runtime = new ModelAgentRuntime({
apiKey: 'test-key',
baseUrl: 'https://example.test',
model: 'test-model',
protocol: 'anthropic-messages',
authentication: 'api-key',
fetcher
})
const history = Array.from({ length: 502 }, (_, index) => ({
role: index % 2 ? 'assistant' as const : 'user' as const,
content: `${index}: ${'x'.repeat(4_200)}`
}))
for (const suppliedHistory of [history, undefined]) {
for await (const event of runtime.run({
requestId: crypto.randomUUID(),
conversationId: 'full-history',
prompt: 'next',
history: suppliedHistory
}, new AbortController().signal)) {
expect(event).toBeDefined()
}
}
for (const [, init] of fetcher.mock.calls) {
const body = JSON.parse(String(init?.body))
expect(body.messages.slice(0, history.length)).toEqual(history)
}
expect(fetcher).toHaveBeenCalledTimes(2)
})
function createTextToolResult(text: string): ModelToolResult {
return {
parts: [{ type: 'text', text }],
+3 -20
View File
@@ -2625,24 +2625,8 @@ export class ModelAgentRuntime implements AgentRuntime {
conversationId: string,
messages: ConversationMessage[]
): void {
const retained: ConversationMessage[] = []
let bytes = 0
const compressionEnabled =
this.options.contextCompression?.settings.enabled === true
const maximumMessages = compressionEnabled ? 500 : 20
const maximumBytes = compressionEnabled
? 2 * 1024 * 1024
: 512 * 1024
for (const message of messages.slice(-maximumMessages).reverse()) {
const messageBytes = Buffer.byteLength(message.content)
if (bytes + messageBytes > maximumBytes) {
break
}
retained.unshift(message)
bytes += messageBytes
}
this.conversations.delete(conversationId)
this.conversations.set(conversationId, retained)
this.conversations.set(conversationId, [...messages])
while (this.conversations.size > 50) {
const oldest = this.conversations.keys().next().value
if (oldest) {
@@ -2658,9 +2642,7 @@ export class ModelAgentRuntime implements AgentRuntime {
request.history && request.history.length > 0
? request.history
: this.conversations.get(request.conversationId) ?? []
return this.options.contextCompression?.settings.enabled
? history
: history.slice(-20)
return history
}
private async *runImageGeneration(
@@ -3381,6 +3363,7 @@ export class ModelAgentRuntime implements AgentRuntime {
}
const toolContext: ModelToolCallContext = {
conversationId: request.conversationId,
browserTabId: request.browserTabId,
workMode: request.workMode ?? 'ask',
requestId: request.requestId,
runtimeTarget: 'model',
+291 -82
View File
@@ -7,9 +7,11 @@ import {
} from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { rgPath } from '@vscode/ripgrep'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import type { ResolvedMcpServer } from '../capabilities/capability-service'
import type { BrowserToolService } from '../browser/browser-model-tools'
import { browserTabIdSchema } from '../../shared/contracts'
import { BrowserStaleReferenceError } from '../browser/cdp-browser-driver'
import {
LocalWorkspaceAccess,
@@ -64,8 +66,15 @@ const png = Buffer.from([
0x0d, 0x0a, 0x1a, 0x0a
]).toString('base64')
const jpeg = Buffer.from([0xff, 0xd8, 0xff, 0xd9]).toString('base64')
const firstBrowserTabId = browserTabIdSchema.parse(
'00000000-0000-4000-8000-000000000201'
)
const secondBrowserTabId = browserTabIdSchema.parse(
'00000000-0000-4000-8000-000000000202'
)
const toolContext = {
conversationId: 'provider-test-conversation',
browserTabId: firstBrowserTabId,
workMode: 'execute'
} satisfies ModelToolCallContext
@@ -177,18 +186,159 @@ describe('ModelToolProvider', () => {
await provider.dispose()
})
it('provides bounded workspace read, list, and atomic write tools', async () => {
it.each(['connect', 'listTools'] as const)(
'isolates MCP %s failures and retries without reconnecting healthy servers',
async (failureMethod) => {
const workspace = await createWorkspace()
const failure = new Error('server unavailable')
const warning = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
const failedClient = {
...mocks.client,
connect: vi.fn().mockResolvedValue(undefined),
listTools: vi.fn().mockResolvedValue({ tools: [] }),
close: vi.fn().mockResolvedValue(undefined)
}
failedClient[failureMethod].mockRejectedValueOnce(failure)
mocks.Client.mockImplementationOnce(function () { return failedClient })
mocks.client.listTools.mockResolvedValue({
tools: [{ name: 'search', inputSchema: { type: 'object' } }]
})
const provider = new ModelToolProvider(workspace, [
createMcpServer(),
{ ...createMcpServer(), id: 'healthy', name: 'Healthy' }
], createBrowserService())
const signal = new AbortController().signal
try {
const tools = await provider.listTools(toolContext, signal)
expect(tools).toEqual(expect.arrayContaining([
expect.objectContaining({ name: 'workspace_read_text' }),
expect.objectContaining({ name: 'browser_snapshot' }),
expect.objectContaining({ serverName: 'Healthy' })
]))
expect(tools.filter((tool) => tool.source === 'mcp')).toHaveLength(1)
expect(failedClient.close).toHaveBeenCalledOnce()
expect(mocks.client.close).not.toHaveBeenCalled()
expect(warning).toHaveBeenCalledWith(
expect.stringContaining('Search MCP'), expect.any(Error)
)
const recovered = await provider.listTools(toolContext, signal)
expect(recovered.filter((tool) => tool.source === 'mcp')).toHaveLength(2)
expect(mocks.Client).toHaveBeenCalledTimes(3)
expect(mocks.client.connect).toHaveBeenCalledTimes(2)
await provider.callTool(
tools.find((tool) => tool.serverName === 'Healthy')!.name,
{}, signal, toolContext
)
expect(mocks.client.callTool).toHaveBeenCalledOnce()
} finally {
await provider.dispose()
warning.mockRestore()
}
}
)
it('loads tools from more than sixteen MCP servers', async () => {
mocks.client.listTools.mockResolvedValue({
tools: [{ name: 'search', inputSchema: { type: 'object' } }]
})
const provider = new ModelToolProvider(await createWorkspace(),
Array.from({ length: 17 }, (_, index) => ({
...createMcpServer(), id: `server-${index}`, name: `Server ${index}`
}))
)
try {
const tools = await provider.listTools(toolContext, new AbortController().signal)
expect(tools.filter((tool) => tool.source === 'mcp')).toHaveLength(17)
expect(mocks.client.connect).toHaveBeenCalledTimes(17)
} finally {
await provider.dispose()
}
})
it('preserves cancellation during MCP discovery and permits a later retry', async () => {
const controller = new AbortController()
const reason = new Error('cancel discovery')
mocks.client.connect.mockImplementationOnce(async () => {
controller.abort(reason)
throw reason
})
const provider = new ModelToolProvider(await createWorkspace(), [createMcpServer()])
try {
await expect(provider.listTools(toolContext, controller.signal)).rejects.toBe(reason)
expect(mocks.client.close).toHaveBeenCalledOnce()
await expect(provider.listTools(toolContext, new AbortController().signal))
.resolves.toEqual(expect.arrayContaining([
expect.objectContaining({ name: 'workspace_read_text' })
]))
expect(mocks.client.connect).toHaveBeenCalledTimes(2)
} finally {
await provider.dispose()
}
})
it('omits failed dynamic catalogs and retries refresh while healthy tools remain usable', async () => {
const warning = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
mocks.client.getServerCapabilities.mockReturnValue({ tools: { listChanged: true } })
mocks.client.listTools.mockResolvedValue({
tools: [{ name: 'search', inputSchema: { type: 'object' } }]
})
const provider = new ModelToolProvider(await createWorkspace(), [
createMcpServer(true),
{ ...createMcpServer(), id: 'healthy', name: 'Healthy' }
])
const signal = new AbortController().signal
try {
await provider.listTools(toolContext, signal)
const options = mocks.Client.mock.calls[0]![1] as {
listChanged: { tools: { onChanged(error: Error | null): void } }
}
options.listChanged.tools.onChanged(null)
mocks.client.listTools.mockRejectedValueOnce(new Error('refresh unavailable'))
const tools = await provider.listTools(toolContext, signal)
expect(tools.filter((tool) => tool.source === 'mcp')).toEqual([
expect.objectContaining({ serverName: 'Healthy' })
])
expect(mocks.client.close).not.toHaveBeenCalled()
const recovered = await provider.listTools(toolContext, signal)
expect(recovered.filter((tool) => tool.source === 'mcp')).toHaveLength(2)
expect(mocks.client.connect).toHaveBeenCalledTimes(2)
options.listChanged.tools.onChanged(null)
const controller = new AbortController()
const reason = new Error('cancel refresh')
mocks.client.listTools.mockImplementationOnce(async () => {
controller.abort(reason)
throw reason
})
await expect(provider.listTools(toolContext, controller.signal)).rejects.toBe(reason)
} finally {
await provider.dispose()
warning.mockRestore()
}
})
it('provides ripgrep, paged reads, and patch tools', async () => {
const workspace = await createWorkspace()
await mkdir(join(workspace, 'docs'))
await writeFile(join(workspace, 'docs', 'note.txt'), 'hello', 'utf8')
const provider = new ModelToolProvider(workspace)
await writeFile(join(workspace, 'docs', 'search.txt'), '中target\n', 'utf8')
const provider = new ModelToolProvider(
workspace,
[],
undefined,
undefined,
false,
{ ripgrepExecutablePath: rgPath }
)
const signal = new AbortController().signal
await expect(provider.listTools(toolContext, signal)).resolves.toEqual(
expect.arrayContaining([
expect.objectContaining({ name: 'workspace_rg' }),
expect.objectContaining({ name: 'workspace_read_text' }),
expect.objectContaining({ name: 'workspace_list_directory' }),
expect.objectContaining({ name: 'workspace_write_text' })
expect.objectContaining({ name: 'workspace_apply_patch' })
])
)
await expect(
@@ -198,32 +348,50 @@ describe('ModelToolProvider', () => {
signal,
toolContext
)
).resolves.toEqual({
parts: [{ type: 'text', text: 'hello' }],
contextBytes: 5
).resolves.toMatchObject({
parts: [{ type: 'text', text: expect.stringContaining('1: hello') }]
})
const listing = await provider.callTool(
'workspace_list_directory',
{ path: 'docs' },
'workspace_rg',
{ path: 'docs', filesOnly: true },
signal,
toolContext
)
expect(listing.parts).toEqual([
expect.objectContaining({
type: 'text',
text: expect.stringContaining('"note.txt"')
text: expect.stringContaining('docs/note.txt')
})
])
const matches = await provider.callTool(
'workspace_rg',
{ path: 'docs', pattern: 'target' },
signal,
toolContext
)
expect(matches.parts).toEqual([
expect.objectContaining({
type: 'text',
text: expect.stringContaining('docs/search.txt:1:2:中target')
})
])
const written = await provider.callTool(
'workspace_write_text',
{ path: 'docs/output.txt', content: 'saved' },
'workspace_apply_patch',
{
patch: [
'*** Begin Patch',
'*** Add File: docs/output.txt',
'+saved',
'*** End Patch'
].join('\n')
},
signal,
toolContext
)
expect(written.parts).toEqual([
expect.objectContaining({
type: 'text',
text: expect.stringContaining('"bytesWritten":5')
text: expect.stringContaining('docs/output.txt')
})
])
await expect(
@@ -231,12 +399,15 @@ describe('ModelToolProvider', () => {
).resolves.toBe('saved')
})
it('delegates workspace tools and disposal to WorkspaceAccess', async () => {
it('delegates paged workspace reads and disposal to WorkspaceAccess', async () => {
const readText = vi.fn(async () => ({
path: 'remote.txt',
name: 'remote.txt',
content: 'remote content',
size: 14
size: 14,
offsetBytes: 0,
bytesRead: 14,
truncated: false
}))
const listDirectory = vi.fn(async () => ({
path: '',
@@ -270,43 +441,19 @@ describe('ModelToolProvider', () => {
signal,
toolContext
)
).resolves.toEqual({
parts: [{ type: 'text', text: 'remote content' }],
contextBytes: 14
).resolves.toMatchObject({
parts: [{ type: 'text', text: expect.stringContaining('1: remote content') }]
})
await provider.callTool(
'workspace_list_directory',
{},
signal,
toolContext
)
await provider.callTool(
'workspace_write_text',
{ path: 'output.txt', content: 'saved' },
signal,
toolContext
)
expect(readText).toHaveBeenCalledWith(
expect.objectContaining({
path: 'remote.txt',
maximumBytes: 256 * 1024,
maximumBytes: 128 * 1024,
allowTruncated: true,
signal
})
)
expect(listDirectory).toHaveBeenCalledWith(
expect.objectContaining({
path: '.',
maximumEntries: 200,
includeOther: true,
signal
})
)
expect(writeTextAtomic).toHaveBeenCalledWith({
path: 'output.txt',
content: 'saved',
maximumBytes: 512 * 1024,
signal
})
expect(listDirectory).not.toHaveBeenCalled()
expect(writeTextAtomic).not.toHaveBeenCalled()
await provider.dispose()
expect(dispose).toHaveBeenCalledOnce()
})
@@ -361,8 +508,8 @@ describe('ModelToolProvider', () => {
for (const [name, argumentsValue] of [
[
'workspace_write_text',
{ path: 'output.txt', content: 'blocked' }
'workspace_apply_patch',
{ patch: '*** Begin Patch\n*** End Patch' }
],
['browser_click', { ref: 'b_target' }],
['note_create', { title: 'blocked' }],
@@ -389,8 +536,8 @@ describe('ModelToolProvider', () => {
expect(mocks.Client).not.toHaveBeenCalled()
const writeTool = {
name: 'workspace_write_text',
displayName: '写入工作区文本',
name: 'workspace_apply_patch',
displayName: '应用工作区补丁',
description: 'write',
inputSchema: {},
source: 'builtin'
@@ -398,7 +545,7 @@ describe('ModelToolProvider', () => {
expect(() =>
provider.getApproval(
writeTool,
{ path: 'output.txt', content: 'blocked' },
{ patch: '*** Begin Patch\n*** End Patch' },
'output.txt',
askContext
)
@@ -425,20 +572,10 @@ describe('ModelToolProvider', () => {
askContext
)
).resolves.toMatchObject({
parts: [{ type: 'text', text: 'remote content' }]
})
await expect(
provider.callTool(
'workspace_list_directory',
{},
signal,
askContext
)
).resolves.toMatchObject({
parts: [{ type: 'text', text: expect.any(String) }]
parts: [{ type: 'text', text: expect.stringContaining('1: remote content') }]
})
expect(readText).toHaveBeenCalledOnce()
expect(listDirectory).toHaveBeenCalledOnce()
expect(listDirectory).not.toHaveBeenCalled()
})
it('exposes scoped reads in Ask and Magic Notes writes only in Execute', async () => {
@@ -499,6 +636,7 @@ describe('ModelToolProvider', () => {
const askTools = await provider.listTools(askContext, signal)
expect(askTools.map((tool) => tool.name)).toEqual([
'workspace_read_text',
'knowledge_list',
'knowledge_search',
'note_list',
@@ -578,7 +716,11 @@ describe('ModelToolProvider', () => {
},
signal
)
).resolves.toEqual([])
).resolves.toEqual(
expect.arrayContaining([
expect.objectContaining({ name: 'workspace_read_text' })
])
)
const executeTools = await provider.listTools(
{ ...askContext, workMode: 'execute' },
signal
@@ -586,8 +728,7 @@ describe('ModelToolProvider', () => {
expect(executeTools.map((tool) => tool.name)).toEqual(
expect.arrayContaining([
'workspace_read_text',
'workspace_list_directory',
'workspace_write_text',
'workspace_apply_patch',
'knowledge_list',
'knowledge_search',
'note_search',
@@ -705,7 +846,9 @@ describe('ModelToolProvider', () => {
workspace,
[createMcpServer()],
undefined,
gateway
gateway,
false,
{ ripgrepExecutablePath: rgPath }
)
await expect(
validProvider.listTools(context, new AbortController().signal)
@@ -719,14 +862,16 @@ describe('ModelToolProvider', () => {
workspace,
[createMcpServer()],
undefined,
gateway
gateway,
false,
{ ripgrepExecutablePath: rgPath }
)
await expect(
overflowingProvider.listTools(
context,
new AbortController().signal
)
).rejects.toThrow('无法加载 MCP Server')
).resolves.toHaveLength(18)
await overflowingProvider.dispose()
})
@@ -799,6 +944,7 @@ describe('ModelToolProvider', () => {
stdoutTruncated: false,
stderrTruncated: false
})),
readOutput: vi.fn(),
releaseConversation: vi.fn(async () => undefined),
dispose: vi.fn(async () => undefined)
} satisfies DirectModelProcessService
@@ -994,10 +1140,12 @@ describe('ModelToolProvider', () => {
const provider = new ModelToolProvider(workspace, [], browserService)
const firstContext = {
conversationId: 'browser-conversation-one',
browserTabId: firstBrowserTabId,
workMode: 'execute'
} satisfies ModelToolCallContext
const secondContext = {
conversationId: 'browser-conversation-two',
browserTabId: secondBrowserTabId,
workMode: 'execute'
} satisfies ModelToolCallContext
const signal = new AbortController().signal
@@ -1051,12 +1199,14 @@ describe('ModelToolProvider', () => {
expect(browserService.screenshot).toHaveBeenNthCalledWith(
1,
firstContext.conversationId,
signal
signal,
firstBrowserTabId
)
expect(browserService.screenshot).toHaveBeenNthCalledWith(
2,
secondContext.conversationId,
signal
signal,
secondBrowserTabId
)
await provider.releaseConversation(firstContext.conversationId)
@@ -1090,6 +1240,63 @@ describe('ModelToolProvider', () => {
})
})
it('keeps a request on its bound tab when another tab changes or closes', async () => {
const workspace = await createWorkspace()
const browserService = createBrowserService()
const provider = new ModelToolProvider(workspace, [], browserService)
const conversationId = 'two-tab-conversation'
const firstRequest = {
conversationId,
browserTabId: firstBrowserTabId,
workMode: 'execute',
requestId: 'first-tab-request'
} satisfies ModelToolCallContext
const secondRequest = {
conversationId,
browserTabId: secondBrowserTabId,
workMode: 'execute',
requestId: 'second-tab-request'
} satisfies ModelToolCallContext
const signal = new AbortController().signal
await provider.callTool('browser_snapshot', {}, signal, firstRequest)
await provider.callTool('browser_snapshot', {}, signal, secondRequest)
expect(browserService.snapshot).toHaveBeenNthCalledWith(
1,
conversationId,
signal,
firstBrowserTabId
)
expect(browserService.snapshot).toHaveBeenNthCalledWith(
2,
conversationId,
signal,
secondBrowserTabId
)
vi.mocked(browserService.click).mockImplementation(
async (_conversationId, _ref, _signal, tabId) => {
if (tabId === firstBrowserTabId) {
throw new Error('浏览器标签页不存在或不属于当前对话')
}
}
)
await expect(
provider.callTool(
'browser_click',
{ ref: 'b_closedTabReference' },
signal,
firstRequest
)
).rejects.toThrow('浏览器标签页不存在或不属于当前对话')
expect(browserService.click).toHaveBeenCalledWith(
conversationId,
'b_closedTabReference',
signal,
firstBrowserTabId
)
})
it('exposes only allowlisted read-only Exa tools in Ask and Execute', async () => {
const workspace = await createWorkspace()
mocks.client.listTools.mockResolvedValue({
@@ -1130,18 +1337,20 @@ describe('ModelToolProvider', () => {
workMode: 'ask'
} satisfies ModelToolCallContext
await expect(provider.listTools(askContext, signal)).resolves.toEqual([
expect.objectContaining({
name: 'web_search',
displayName: '联网搜索',
source: 'builtin'
}),
expect.objectContaining({
name: 'web_fetch',
displayName: '读取网页',
source: 'builtin'
})
])
await expect(provider.listTools(askContext, signal)).resolves.toEqual(
expect.arrayContaining([
expect.objectContaining({
name: 'web_search',
displayName: '联网搜索',
source: 'builtin'
}),
expect.objectContaining({
name: 'web_fetch',
displayName: '读取网页',
source: 'builtin'
})
])
)
await provider.callTool(
'web_search',
{ query: 'GoodBuddy current release', numResults: 3 },
+43 -32
View File
@@ -24,6 +24,7 @@ import type {
} from './runtime'
import {
BrowserModelTools,
browserToolNames,
type BrowserToolService
} from '../browser/browser-model-tools'
import { BrowserStaleReferenceError } from '../browser/cdp-browser-driver'
@@ -45,12 +46,12 @@ import {
directModelSubagentInputSchema
} from '../assistant/direct-model-subagent-service'
import type { LaunchEnvironmentProvider } from '../local-tool-environment'
import type { BrowserTabId } from '../../shared/contracts'
import { searchWorkspaceWithRipgrep } from './direct-model-ripgrep'
import { applyWorkspacePatch } from './workspace-apply-patch'
import { readWorkspaceLines } from './workspace-read-lines'
const MAX_MODEL_TOOLS = 100
const MAX_MCP_SERVERS = 16
const MAX_TOOL_RESULT_BYTES = 256 * 1024
const MCP_TIMEOUT_MS = 30_000
const MCP_CALL_MAX_TOTAL_TIMEOUT_MS = 5 * 60_000
@@ -269,6 +270,7 @@ export type ModelToolResult = {
export type ModelToolCallContext = {
conversationId: string
browserTabId?: BrowserTabId
workMode: 'ask' | 'execute'
requestId?: string
runtimeTarget?: 'model'
@@ -635,7 +637,10 @@ function normalizeMcpResult(result: unknown): ModelToolResult {
}
export class ModelToolProvider implements ModelToolProviderLike {
private mcpConnections?: Promise<ConnectedMcp[]>
private readonly mcpConnections = new Map<
ResolvedMcpServer,
Promise<ConnectedMcp | undefined>
>()
private webSearchBindings?: Promise<Map<string, McpToolBinding>>
private readonly clients = new Set<Client>()
private readonly customMcpClients = new Set<Client>()
@@ -708,10 +713,13 @@ export class ModelToolProvider implements ModelToolProviderLike {
private getBrowserTools(
context: ModelToolCallContext
): BrowserModelTools | undefined {
return this.browserService && context.workMode === 'execute'
return this.browserService &&
context.workMode === 'execute' &&
context.browserTabId
? new BrowserModelTools({
service: this.browserService,
conversationId: context.conversationId
conversationId: context.conversationId,
browserTabId: context.browserTabId
})
: undefined
}
@@ -719,7 +727,7 @@ export class ModelToolProvider implements ModelToolProviderLike {
private getReservedToolCount(): number {
return (
this.getBuiltinTools().length +
(this.browserService ? 7 : 0) +
(this.browserService ? browserToolNames.length : 0) +
(this.webSearchEnabled ? 2 : 0) +
(this.programming.processService ? 1 : 0) +
(this.programming.subagentService ? 1 : 0) +
@@ -1077,23 +1085,27 @@ export class ModelToolProvider implements ModelToolProviderLike {
signal: AbortSignal,
refreshDynamic = false
): Promise<Map<string, McpToolBinding>> {
if (this.mcpServers.length > MAX_MCP_SERVERS) {
throw new Error('直连模型最多可加载 16 个 MCP Server')
}
this.mcpConnections ??= Promise.all(
this.mcpServers.map((server) => this.connectMcpServer(server, signal))
)
.catch(async (error) => {
this.mcpConnections = undefined
const clients = [...this.customMcpClients]
this.customMcpClients.clear()
clients.forEach((client) => this.clients.delete(client))
await Promise.allSettled(
clients.map((client) => client.close())
)
throw error
signal.throwIfAborted()
const loaded = await Promise.all(
this.mcpServers.map((server) => {
let pending = this.mcpConnections.get(server)
if (!pending) {
pending = this.connectMcpServer(server, signal).catch((error) => {
this.mcpConnections.delete(server)
if (!signal.aborted) {
console.warn(`MCP discovery failed: ${server.name}`, error)
}
return undefined
})
this.mcpConnections.set(server, pending)
}
return pending
})
const connections = await this.mcpConnections
)
signal.throwIfAborted()
const connections = loaded.filter(
(connection): connection is ConnectedMcp => connection !== undefined
)
if (refreshDynamic) {
await Promise.all(
connections.map(async (connection) => {
@@ -1123,14 +1135,18 @@ export class ModelToolProvider implements ModelToolProviderLike {
)
} catch (error) {
connection.dynamicToolsChanged = true
throw new Error(
`无法刷新 MCP Server「${connection.server.name}」的工具`,
{ cause: error }
)
connection.tools = []
if (!signal.aborted) {
console.warn(
`MCP refresh failed: ${connection.server.name}`,
error
)
}
}
})
)
}
signal.throwIfAborted()
const bindings = new Map<string, McpToolBinding>()
const reservedToolCount = this.getReservedToolCount()
for (const connection of connections) {
@@ -1782,7 +1798,7 @@ export class ModelToolProvider implements ModelToolProviderLike {
this.clients.clear()
this.customMcpClients.clear()
this.webSearchClients.clear()
this.mcpConnections = undefined
this.mcpConnections.clear()
this.webSearchBindings = undefined
await Promise.allSettled([
...clients.map((client) => client.close()),
@@ -1794,12 +1810,7 @@ export class ModelToolProvider implements ModelToolProviderLike {
async releaseConversation(conversationId: string): Promise<void> {
await Promise.allSettled([
this.browserService
? new BrowserModelTools({
service: this.browserService,
conversationId
}).release()
: undefined,
this.browserService?.releaseConversation(conversationId),
this.programming.processService?.releaseConversation(
conversationId
),
+104
View File
@@ -0,0 +1,104 @@
import { mkdtemp, readdir, rm } from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { afterEach, describe, expect, it } from 'vitest'
import { PagedOutputStore, PAGED_OUTPUT_MAX_PAGE_BYTES } from './paged-output-store'
const stores: PagedOutputStore[] = []
const directories: string[] = []
afterEach(async () => {
await Promise.all(stores.splice(0).map((store) => store.dispose()))
await Promise.all(directories.splice(0).map((directory) => rm(directory, { recursive: true, force: true })))
})
async function createStore() {
const directory = await mkdtemp(join(tmpdir(), 'goodbuddy-output-test-'))
directories.push(directory)
const store = new PagedOutputStore('test', { temporaryParentDirectory: directory })
stores.push(store)
return { store, directory }
}
describe('PagedOutputStore', () => {
it.each([1, 2, 3, 4, 7])('round trips split UTF-8 chunks with a %i byte page', async (limit) => {
const { store } = await createStore()
const text = 'a\u4f60\ud83d\ude00\ufffd\u597dZ'
const writer = await store.create('owner', 2)
for (const byte of Buffer.from(text)) writer.append(Buffer.from([byte]))
const saved = await writer.finish()
expect(saved.text).toBe('a')
let output = saved.text
let cursor = saved.reference!.nextCursor
while (cursor < saved.reference!.totalBytes) {
const page = await store.read('owner', writer.handle, cursor, limit)
expect(page.nextCursor).toBeGreaterThan(cursor)
expect(Buffer.byteLength(page.content)).toBeLessThanOrEqual(Math.max(limit, 4))
output += page.content
cursor = page.nextCursor
}
expect(output).toBe(text)
await expect(store.read('owner', writer.handle, cursor, limit)).resolves.toMatchObject({ content: '', eof: true, nextCursor: cursor })
await expect(store.read('owner', writer.handle, 2)).rejects.toThrow('UTF-8')
await expect(store.read('other', writer.handle)).rejects.toThrow()
for (const invalid of [-1, 0.5, cursor + 1]) {
await expect(store.read('owner', writer.handle, invalid)).rejects.toThrow()
}
for (const invalid of [0, 0.5, PAGED_OUTPUT_MAX_PAGE_BYTES + 1]) {
await expect(store.read('owner', writer.handle, 0, invalid)).rejects.toThrow()
}
})
it('keeps literal replacement characters in previews', async () => {
const { store } = await createStore()
const writer = await store.create('owner', 3)
writer.append('\ufffdrest')
const saved = await writer.finish()
expect(saved.text).toBe('\ufffd')
expect(saved.reference?.nextCursor).toBe(3)
})
it('retains output omitted to keep JSON-escaped previews bounded', async () => {
const { store } = await createStore()
const writer = await store.create('owner', 12)
writer.append('\u0000'.repeat(12))
const saved = await writer.finish()
expect(saved.text).toBe('\u0000\u0000')
expect(saved.reference?.nextCursor).toBe(2)
const page = await store.read('owner', writer.handle, saved.reference!.nextCursor)
expect(saved.text + page.content).toBe('\u0000'.repeat(12))
})
it('removes small, aborted, released and disposed output files', async () => {
const { store, directory } = await createStore()
const small = await store.create('owner', 10)
small.append('small')
expect(await small.finish()).toEqual({ text: 'small', truncated: false })
const retained = await store.create('owner', 1)
retained.append('retained')
await retained.finish()
const active = await store.create('owner', 1)
active.append('active')
const other = await store.create('other', 1)
other.append('other')
await other.finish()
const [childDirectory] = await readdir(directory)
await store.releaseOwner('owner')
expect(await readdir(join(directory, childDirectory!))).toHaveLength(1)
await expect(active.finish()).rejects.toThrow()
await expect(store.read('owner', retained.handle)).rejects.toThrow()
await expect(store.read('other', other.handle)).resolves.toMatchObject({ content: 'other' })
await store.create('other', 1)
await store.dispose()
expect(await readdir(directory)).toEqual([])
await expect(store.create('owner', 1)).rejects.toThrow()
})
it('cleans up when disposed during directory creation', async () => {
const { store, directory } = await createStore()
const pending = store.create('owner', 1)
const rejected = expect(pending).rejects.toThrow()
await store.dispose()
await rejected
expect(await readdir(directory)).toEqual([])
})
})
+296
View File
@@ -0,0 +1,296 @@
import { randomUUID } from 'node:crypto'
import {
closeSync,
openSync,
readSync,
writeSync
} from 'node:fs'
import { mkdir, mkdtemp, rm } from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
export const PAGED_OUTPUT_DEFAULT_PAGE_BYTES = 32 * 1024
export const PAGED_OUTPUT_MAX_PAGE_BYTES = 32 * 1024
export type PagedOutputReference = {
handle: string
nextCursor: number
totalBytes: number
}
export type PagedOutputPage = {
handle: string
content: string
cursor: number
nextCursor: number
totalBytes: number
eof: boolean
}
type OutputEntry = {
ownerId: string
path: string
totalBytes: number
}
export type PagedOutputStoreOptions = {
temporaryParentDirectory?: string
}
export class PagedOutputWriter {
private preview = Buffer.alloc(0)
private totalBytes = 0
private failure?: Error
private closed = false
constructor(
private readonly store: PagedOutputStore,
readonly handle: string,
readonly ownerId: string,
readonly path: string,
private readonly descriptor: number,
private readonly previewBytes: number
) {}
append(chunk: Buffer | string): void {
if (this.closed || this.failure) {
return
}
const value = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)
try {
let offset = 0
while (offset < value.byteLength) {
const written = writeSync(this.descriptor, value, offset)
if (written === 0) throw new Error('Unable to write tool output')
offset += written
}
this.totalBytes += value.byteLength
if (this.preview.byteLength < this.previewBytes) {
const remaining = this.previewBytes - this.preview.byteLength
this.preview = Buffer.concat([
this.preview,
value.subarray(0, remaining)
])
}
} catch (error) {
this.failure =
error instanceof Error ? error : new Error('无法保存分页输出')
}
}
async finish(): Promise<{
text: string
truncated: boolean
reference?: PagedOutputReference
}> {
if (this.closed) {
throw new Error('分页输出已经关闭')
}
this.closed = true
closeSync(this.descriptor)
if (this.failure) {
await this.store.discard(this)
throw new Error('无法保存完整工具输出', { cause: this.failure })
}
let previewLength = utf8PrefixLength(this.preview)
let text = this.preview.subarray(0, previewLength).toString('utf8')
// Tool results are JSON: escaped control characters also consume context.
if (Buffer.byteLength(JSON.stringify(text)) - 2 > this.previewBytes) {
let low = 0
let high = previewLength
while (low < high) {
const middle = Math.ceil((low + high) / 2)
const candidate = utf8Prefix(this.preview.subarray(0, middle))
if (Buffer.byteLength(JSON.stringify(candidate)) - 2 <= this.previewBytes) {
low = middle
} else {
high = middle - 1
}
}
previewLength = utf8PrefixLength(this.preview, low)
text = this.preview.subarray(0, previewLength).toString('utf8')
}
if (this.totalBytes <= previewLength) {
await this.store.discard(this)
return { text, truncated: false }
}
this.store.retain(this, this.totalBytes)
return {
text,
truncated: true,
reference: {
handle: this.handle,
nextCursor: previewLength,
totalBytes: this.totalBytes
}
}
}
async abort(): Promise<void> {
if (!this.closed) {
this.closed = true
closeSync(this.descriptor)
}
await this.store.discard(this)
}
}
function utf8PrefixLength(value: Buffer, end = value.byteLength): number {
if (end === 0) return 0
let start = end - 1
while (start > 0 && (value[start]! & 0xc0) === 0x80) start -= 1
const lead = value[start]!
const width = lead >= 0xf0 ? 4 : lead >= 0xe0 ? 3 : lead >= 0xc0 ? 2 : 1
return start + width > end ? start : end
}
function utf8Prefix(value: Buffer): string {
return value.subarray(0, utf8PrefixLength(value)).toString('utf8')
}
export class PagedOutputStore {
private readonly entries = new Map<string, OutputEntry>()
private readonly writers = new Set<PagedOutputWriter>()
private directoryPromise?: Promise<string>
private disposed = false
private disposePromise?: Promise<void>
constructor(
private readonly handlePrefix: string,
private readonly options: PagedOutputStoreOptions = {}
) {}
async create(
ownerId: string,
previewBytes: number
): Promise<PagedOutputWriter> {
if (this.disposed) {
throw new Error('分页输出存储已关闭')
}
const directory = await this.getDirectory()
if (this.disposed) throw new Error('分页输出存储已关闭')
const id = randomUUID()
const path = join(directory, id)
const descriptor = openSync(path, 'wx')
const writer = new PagedOutputWriter(
this,
`${this.handlePrefix}:${id}`,
ownerId,
path,
descriptor,
previewBytes
)
this.writers.add(writer)
return writer
}
async read(
ownerId: string,
handle: string,
cursor = 0,
limitBytes = PAGED_OUTPUT_DEFAULT_PAGE_BYTES
): Promise<PagedOutputPage> {
const entry = this.entries.get(handle)
if (!entry || entry.ownerId !== ownerId) {
throw new Error('分页输出句柄不存在或不属于当前会话')
}
if (!Number.isSafeInteger(cursor) || cursor < 0 || cursor > entry.totalBytes) {
throw new Error('分页输出 cursor 无效')
}
if (
!Number.isSafeInteger(limitBytes) ||
limitBytes < 1 ||
limitBytes > PAGED_OUTPUT_MAX_PAGE_BYTES
) {
throw new Error(`分页输出每次最多读取 ${PAGED_OUTPUT_MAX_PAGE_BYTES} 字节`)
}
const requestedBytes = Math.min(limitBytes + 3, entry.totalBytes - cursor)
const buffer = Buffer.alloc(requestedBytes)
const descriptor = openSync(entry.path, 'r')
let bytesRead: number
try {
bytesRead = readSync(descriptor, buffer, 0, requestedBytes, cursor)
} finally {
closeSync(descriptor)
}
if (bytesRead && (buffer[0]! & 0xc0) === 0x80) {
throw new Error('分页输出 cursor 必须位于 UTF-8 字符边界')
}
let consumedBytes = utf8PrefixLength(buffer, Math.min(limitBytes, bytesRead))
// A tiny page still consumes one complete character (at most four bytes).
if (consumedBytes === 0 && bytesRead > 0) {
consumedBytes = utf8PrefixLength(buffer, Math.min(4, bytesRead))
}
// Preserve forward progress even for an incomplete final UTF-8 sequence.
if (consumedBytes === 0) consumedBytes = bytesRead
const content = buffer.subarray(0, consumedBytes).toString('utf8')
const nextCursor = cursor + consumedBytes
return {
handle,
content,
cursor,
nextCursor,
totalBytes: entry.totalBytes,
eof: nextCursor >= entry.totalBytes
}
}
async releaseOwner(ownerId: string): Promise<void> {
await Promise.all([...this.writers]
.filter((writer) => writer.ownerId === ownerId)
.map((writer) => writer.abort()))
const owned = [...this.entries].filter(
([, entry]) => entry.ownerId === ownerId
)
for (const [handle] of owned) {
this.entries.delete(handle)
}
await Promise.allSettled(
owned.map(([, entry]) => rm(entry.path, { force: true }))
)
}
dispose(): Promise<void> {
this.disposePromise ??= this.disposeOnce()
return this.disposePromise
}
private async disposeOnce(): Promise<void> {
this.disposed = true
await Promise.all([...this.writers].map((writer) => writer.abort()))
this.entries.clear()
if (this.directoryPromise) {
const directory = await this.directoryPromise.catch(() => undefined)
if (directory) {
await rm(directory, { recursive: true, force: true })
}
}
}
retain(writer: PagedOutputWriter, totalBytes: number): void {
this.writers.delete(writer)
this.entries.set(writer.handle, {
ownerId: writer.ownerId,
path: writer.path,
totalBytes
})
}
async discard(writer: PagedOutputWriter): Promise<void> {
this.writers.delete(writer)
this.entries.delete(writer.handle)
await rm(writer.path, { force: true })
}
private getDirectory(): Promise<string> {
this.directoryPromise ??= this.createDirectory()
return this.directoryPromise
}
private async createDirectory(): Promise<string> {
const parent = this.options.temporaryParentDirectory ?? tmpdir()
await mkdir(parent, { recursive: true })
return await mkdtemp(join(parent, `goodbuddy-${this.handlePrefix}-`))
}
}
+1 -1
View File
@@ -229,7 +229,7 @@ describe('runtime selection', () => {
})
).toEqual({
provider: 'opencode',
profileId: secondProfileId
profileId: defaultProfileId
})
expect(
resolveConfiguredAgentRuntimeSelection(configured, {
+7 -2
View File
@@ -49,6 +49,7 @@ export function resolveConfiguredAgentRuntimeSelection(
settings: ResolvedRuntimeSettings,
selection: AgentRuntimeSelection
): AgentRuntimeSelection {
if ('profileId' in selection && selection.profileId) return selection
if (
selection.provider !== 'opencode' &&
selection.provider !== 'continue' &&
@@ -87,7 +88,7 @@ export function applyRuntimeSelection(
}
if (selection.provider === 'model') {
const profile = requireProfile(settings, selection.profileId)
const profile = requireProfile(settings, selection.profileId ?? settings.defaultModelProfileId)
return {
target: 'model',
settings: {
@@ -108,7 +109,11 @@ export function applyRuntimeSelection(
const profile = selection.profileId
? requireProfile(settings, selection.profileId)
: undefined
: selection.provider === 'opencode'
? settings.opencodeModelProfile
: selection.provider === 'continue'
? settings.continueModelProfile
: settings.deepseekHarnessModelProfile
if (selection.provider === 'opencode') {
if (profile && !isAgentRuntimeModelProtocol(profile.protocol)) {
throw new Error(
+3
View File
@@ -4,6 +4,7 @@ import type {
AgentQuestionAnswer,
AgentRequest,
AgentRuntimeStatus,
BrowserTabId,
RuntimeConversationCompactInput,
RuntimeConversationCompactResult,
RuntimeNativeSnapshot
@@ -137,6 +138,8 @@ export type AgentExecutionRequest = Omit<AgentRequest, 'workMode'> & {
trustedInstructions?: string
/** Main-process-only request-scoped authorization for built-in data tools. */
knowledgeCapabilityToken?: string
/** Main-process-only browser tab fixed for this request. */
browserTabId?: BrowserTabId
/**
* Main-process recovery cursor for an Agent-owned semantic prompt.
* It is never sent as prompt content and is ignored by local runtimes.