feat: add persistent workbar browser tabs and scoped MCP routing
This commit is contained in:
@@ -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()
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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)) {
|
||||
|
||||
@@ -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 }],
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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 },
|
||||
|
||||
@@ -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
|
||||
),
|
||||
|
||||
@@ -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([])
|
||||
})
|
||||
})
|
||||
@@ -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}-`))
|
||||
}
|
||||
}
|
||||
@@ -229,7 +229,7 @@ describe('runtime selection', () => {
|
||||
})
|
||||
).toEqual({
|
||||
provider: 'opencode',
|
||||
profileId: secondProfileId
|
||||
profileId: defaultProfileId
|
||||
})
|
||||
expect(
|
||||
resolveConfiguredAgentRuntimeSelection(configured, {
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user