root/goodbuddy
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

fix: bound model streaming

Browse Source
This commit is contained in:
lofyer
2026-08-13 03:18:55 +08:00
parent 980f3a0c8f
commit 5b579ae100
8 changed files with 1081 additions and 475 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main/agent/bounded-response.test.ts
+48
View File
@@ -0,0 +1,48 @@
import { describe, expect, it } from 'vitest'
import { readBoundedResponseText } from './bounded-response'
describe('readBoundedResponseText', () => {
it('cancels an oversized response as soon as it crosses the byte limit', async () => {
const chunk = new Uint8Array(1024 * 1024)
let pulls = 0
const response = new Response(
new ReadableStream<Uint8Array>({
pull(controller) {
pulls += 1
controller.enqueue(chunk)
}
})
)
await expect(
readBoundedResponseText(response, {
maxBytes: 8 * 1024 * 1024,
tooLargeMessage: 'response too large'
})
).rejects.toThrow('response too large')
expect(pulls).toBeLessThan(20)
})
it('rejects an invalid declared response length without reading the body', async () => {
let pulls = 0
const response = new Response(
new ReadableStream<Uint8Array>({
pull(controller) {
pulls += 1
controller.enqueue(new Uint8Array([1]))
}
}),
{
headers: { 'content-length': 'invalid' }
}
)
await expect(
readBoundedResponseText(response, {
maxBytes: 1024,
tooLargeMessage: 'response too large'
})
).rejects.toThrow('response too large')
expect(pulls).toBe(0)
})
})
src/main/agent/bounded-response.ts
+54
View File
@@ -0,0 +1,54 @@
export type BoundedResponseTextOptions = {
maxBytes: number
missingBodyMessage?: string
tooLargeMessage: string
}
export async function readBoundedResponseText(
response: Response,
options: BoundedResponseTextOptions
): Promise<string> {
const declaredLength = response.headers.get('content-length')
if (declaredLength !== null) {
const parsedLength = Number(declaredLength)
if (
!Number.isSafeInteger(parsedLength) ||
parsedLength < 0 ||
parsedLength > options.maxBytes
) {
await response.body?.cancel().catch(() => undefined)
throw new Error(options.tooLargeMessage)
}
}
if (!response.body) {
if (options.missingBodyMessage) {
throw new Error(options.missingBodyMessage)
}
return ''
}
const reader = response.body.getReader()
const chunks: Uint8Array[] = []
let completed = false
let total = 0
try {
while (true) {
const { done, value } = await reader.read()
if (done) {
completed = true
break
}
total += value.byteLength
if (total > options.maxBytes) {
throw new Error(options.tooLargeMessage)
}
chunks.push(value)
}
} finally {
if (!completed) {
await reader.cancel().catch(() => undefined)
}
reader.releaseLock()
}
return Buffer.concat(chunks, total).toString('utf8')
}
src/main/agent/continue-host-adapter.test.ts
+59 -2
View File
@@ -7,7 +7,7 @@ import {
writeFile
} from 'node:fs/promises'
import { existsSync, readFileSync } from 'node:fs'
import { createHash } from 'node:crypto'
import { createHash, randomUUID } from 'node:crypto'
import { createServer } from 'node:http'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
@@ -165,8 +165,14 @@ describe('ContinueHostAdapter', () => {
'let r=[eS.join(hu.continueHome,AKt)],o='
)
expect(bundle).toContain('goodbuddyEvents:[]')
expect(bundle).toContain('goodbuddyEventsBytes:0')
expect(bundle).toContain('goodbuddyEventsBytes+=Buffer.byteLength')
expect(bundle).toContain('goodbuddyEventsBytes<=2097152')
expect(bundle).toContain('l.length<=1e5')
expect(bundle).toContain('goodbuddyEventsOverflow:!1')
expect(bundle).toContain('goodbuddyEventsOverflow=!0')
expect(bundle).toContain('goodbuddyEvents:ce')
expect(bundle).toContain('type:"text",delta:u')
expect(bundle).toContain('type:"text",delta:l')
expect(bundle).toContain('onToolStart?.(c.name,c.arguments,c.id)')
expect(bundle).toContain(
'function ZZo(e){let t=[];if(e.allow)'
@@ -994,6 +1000,57 @@ describe('ContinueHostAdapter', () => {
expect(killed).toBe(true)
})
it('fails when the patched host reports dropped stream events', async () => {
const distribution = await createDistribution()
let stateRequests = 0
vi.stubGlobal(
'fetch',
vi.fn(async (input: string | URL | Request) => {
if (String(input).endsWith('/state')) {
stateRequests += 1
return Response.json({
session: { history: [] },
isProcessing: stateRequests > 1,
messageQueueLength: 0,
pendingPermission: null,
goodbuddyEventsOverflow: stateRequests > 1
})
}
return Response.json({})
})
)
const adapter = new ContinueHostAdapter({
binaryPath: distribution.entryPath,
configPath: '',
workspace: process.cwd(),
cacheRoot: distribution.cacheRoot,
trustedBundleHashes: [distribution.sourceHash],
launchHost: () => ({
exitCode: null,
killed: false,
stderr: null,
once: () => undefined,
kill: () => true
}),
modelProfile: {
id: randomUUID(),
name: 'Local model',
baseUrl: 'http://127.0.0.1:11434/v1',
modelName: 'qwen3',
protocol: 'openai-chat-completions',
authentication: 'none'
}
})
await expect(
adapter.run(
'hello',
new AbortController().signal,
async () => 'deny'
)
).rejects.toThrow('流式事件超过安全限制')
})
it('uses auto mode and returns audit metadata for agent tools', async () => {
const distribution = await createDistribution()
let launchArgs: string[] = []
src/main/agent/continue-host-adapter.ts
+59 -23
View File
@@ -38,6 +38,7 @@ import {
safeToolErrorDetail
} from './approval-summary'
import { stageRuntimeSkillPackages } from './runtime-skill-packages'
import { readBoundedResponseText } from './bounded-response'
const supportedVersion = '1.5.47'
const supportedBundleHashes = new Set([
@@ -49,6 +50,8 @@ const maximumMessageBytes = 20 * 1024 * 1024
const maximumConfigBytes = 1024 * 1024
const maximumConfiguredMcpServers = 100
const maximumStreamEvents = 5_000
const maximumStreamEventBytes = 2 * 1024 * 1024
const maximumExecutionMilliseconds = 10 * 60_000
const knowledgeMcpName = 'goodbuddy-knowledge'
export const continueConfigurationRequiredMessage =
'Continue 尚未配置模型连接，请在设置中选择 GoodBuddy 模型连接或指定 Continue 配置文件'
@@ -116,7 +119,8 @@ const stateSchema = z.object({
goodbuddyEvents: z
.array(continueHostStreamEventSchema)
.max(maximumStreamEvents)
.optional()
.optional(),
goodbuddyEventsOverflow: z.boolean().optional()
})
type ContinueHostState = z.infer<typeof stateSchema>
@@ -704,17 +708,17 @@ export class ContinueHostAdapter {
patched = replaceExactly(
patched,
streamCallbacksMarker,
'a={onContent:u=>{u&&e.goodbuddyEvents.length<5e3&&e.goodbuddyEvents.push({type:"text",delta:u})},onContentComplete:u=>{},onToolStart:(u,l,c)=>{c&&e.goodbuddyEvents.length<5e3&&e.goodbuddyEvents.push({type:"tool",callId:c,name:u,state:"running",input:(()=>{try{return JSON.stringify(l).slice(0,4e3)}catch{return"[无法序列化]"}})()})},onToolResult:(u,l,c,d)=>{d&&e.goodbuddyEvents.length<5e3&&e.goodbuddyEvents.push({type:"tool",callId:d,name:l,state:c==="done"?"completed":"failed",output:String(u).slice(0,16e3)})},onToolError:(u,l,c)=>{c&&e.goodbuddyEvents.length<5e3&&e.goodbuddyEvents.push({type:"tool",callId:c,name:l??"unknown",state:"failed",error:String(u).slice(0,1e3)})},onToolPermissionRequest:'
'a={onContent:u=>{if(!u)return;let l=String(u);e.goodbuddyEventsBytes+=Buffer.byteLength(l);l.length<=1e5&&e.goodbuddyEvents.length<5e3&&e.goodbuddyEventsBytes<=2097152?e.goodbuddyEvents.push({type:"text",delta:l}):e.goodbuddyEventsOverflow=!0},onContentComplete:u=>{},onToolStart:(u,l,c)=>{if(!c)return;let d={type:"tool",callId:c,name:u,state:"running",input:(()=>{try{return JSON.stringify(l).slice(0,4e3)}catch{return"[无法序列化]"}})()};e.goodbuddyEventsBytes+=Buffer.byteLength(JSON.stringify(d));e.goodbuddyEvents.length<5e3&&e.goodbuddyEventsBytes<=2097152?e.goodbuddyEvents.push(d):e.goodbuddyEventsOverflow=!0},onToolResult:(u,l,c,d)=>{if(!d)return;let p={type:"tool",callId:d,name:l,state:c==="done"?"completed":"failed",output:String(u).slice(0,16e3)};e.goodbuddyEventsBytes+=Buffer.byteLength(JSON.stringify(p));e.goodbuddyEvents.length<5e3&&e.goodbuddyEventsBytes<=2097152?e.goodbuddyEvents.push(p):e.goodbuddyEventsOverflow=!0},onToolError:(u,l,c)=>{if(!c)return;let d={type:"tool",callId:c,name:l??"unknown",state:"failed",error:String(u).slice(0,1e3)};e.goodbuddyEventsBytes+=Buffer.byteLength(JSON.stringify(d));e.goodbuddyEvents.length<5e3&&e.goodbuddyEventsBytes<=2097152?e.goodbuddyEvents.push(d):e.goodbuddyEventsOverflow=!0},onToolPermissionRequest:'
)
patched = replaceExactly(
patched,
serverStateMarker,
'pendingPermission:null,goodbuddyEvents:[]},B='
'pendingPermission:null,goodbuddyEvents:[],goodbuddyEventsBytes:0,goodbuddyEventsOverflow:!1},B='
)
patched = replaceExactly(
patched,
serverStateEndpointMarker,
'j.get("/state",(we,Te)=>{M.lastActivity=Date.now(),B();let ue=e7e(M.session,M.isProcessing,rS.getQueueLength(),M.pendingPermission),ce=M.goodbuddyEvents.splice(0);Te.json({...ue,goodbuddyEvents:ce})})'
'j.get("/state",(we,Te)=>{M.lastActivity=Date.now(),B();let ue=e7e(M.session,M.isProcessing,rS.getQueueLength(),M.pendingPermission),ce=M.goodbuddyEvents.splice(0),de=M.goodbuddyEventsOverflow;M.goodbuddyEventsBytes=0,M.goodbuddyEventsOverflow=!1;Te.json({...ue,goodbuddyEvents:ce,goodbuddyEventsOverflow:de})})'
)
patched = replaceExactly(
patched,
@@ -839,14 +843,10 @@ export class ContinueHostAdapter {
redirect: 'error',
signal: init.signal
})
const contentLength = Number(response.headers.get('content-length') ?? 0)
if (contentLength > maximumStateBytes) {
throw new Error('Continue 宿主响应超过安全大小限制')
}
const body = await response.text()
if (Buffer.byteLength(body) > maximumStateBytes) {
throw new Error('Continue 宿主响应超过安全大小限制')
}
const body = await readBoundedResponseText(response, {
maxBytes: maximumStateBytes,
tooLargeMessage: 'Continue 宿主响应超过安全大小限制'
})
if (!response.ok) {
throw new Error(`Continue 宿主请求失败（HTTP ${response.status}）`)
}
@@ -861,8 +861,11 @@ export class ContinueHostAdapter {
signal: AbortSignal
): Promise<ContinueHostState> {
const expiresAt = Date.now() + 30_000
const timeoutSignal = AbortSignal.timeout(30_000)
const startupSignal = AbortSignal.any([signal, timeoutSignal])
try {
while (Date.now() < expiresAt) {
signal.throwIfAborted()
startupSignal.throwIfAborted()
const childFailure = getChildFailure()
if (childFailure) {
throw childFailure
@@ -872,12 +875,20 @@ export class ContinueHostAdapter {
}
try {
return stateSchema.parse(
await this.request(origin, token, '/state', { signal })
await this.request(origin, token, '/state', {
signal: startupSignal
})
)
} catch {
await delay(150, signal)
await delay(150, startupSignal)
}
}
} catch (error) {
if (timeoutSignal.aborted && !signal.aborted) {
throw new Error('Continue 宿主启动超时', { cause: error })
}
throw error
}
throw new Error('Continue 宿主启动超时')
}
@@ -1150,6 +1161,7 @@ export class ContinueHostAdapter {
let observedTools: ContinueHostTool[] = []
let streamedText = false
let executionTimeoutSignal: AbortSignal | undefined
try {
const initialState = await this.waitForStartup(
child,
@@ -1159,6 +1171,13 @@ export class ContinueHostAdapter {
signal
)
const startIndex = initialState.session.history.length
executionTimeoutSignal = AbortSignal.timeout(
maximumExecutionMilliseconds
)
const executionSignal = AbortSignal.any([
signal,
executionTimeoutSignal
])
const message =
runOptions.images && runOptions.images.length > 0
? [
@@ -1178,13 +1197,13 @@ export class ContinueHostAdapter {
await this.request(origin, token, '/message', {
method: 'POST',
body: messageBody,
signal
signal: executionSignal
})
const expiresAt = Date.now() + 10 * 60_000
const expiresAt = Date.now() + maximumExecutionMilliseconds
const handledPermissionIds = new Set<string>()
while (Date.now() < expiresAt) {
signal.throwIfAborted()
executionSignal.throwIfAborted()
if (childFailure) {
throw childFailure
}
@@ -1194,8 +1213,19 @@ export class ContinueHostAdapter {
)
}
const state = stateSchema.parse(
await this.request(origin, token, '/state', { signal })
await this.request(origin, token, '/state', {
signal: executionSignal
})
)
if (state.goodbuddyEventsOverflow) {
throw new Error('Continue 宿主流式事件超过安全限制')
}
const streamEventBytes = Buffer.byteLength(
JSON.stringify(state.goodbuddyEvents ?? [])
)
if (streamEventBytes > maximumStreamEventBytes) {
throw new Error('Continue 宿主流式事件超过安全限制')
}
observedTools = mergeContinueTools(
observedTools,
extractContinueTools(state.session.history, startIndex)
@@ -1265,7 +1295,7 @@ export class ContinueHostAdapter {
requestId: pending.requestId,
approved: decision !== 'deny'
}),
signal
signal: executionSignal
})
}
if (
@@ -1308,16 +1338,22 @@ export class ContinueHostAdapter {
: {})
}
}
await delay(150, signal)
await delay(150, executionSignal)
}
throw new Error('Continue 宿主执行超时')
} catch (error) {
if (error instanceof ContinueHostRunError) {
throw error
}
const normalizedError =
executionTimeoutSignal?.aborted && !signal.aborted
? new Error('Continue 宿主执行超时', { cause: error })
: error
throw new ContinueHostRunError(
error instanceof Error ? error.message : 'Continue 宿主执行失败',
{ cause: error, tools: observedTools }
normalizedError instanceof Error
? normalizedError.message
: 'Continue 宿主执行失败',
{ cause: normalizedError, tools: observedTools }
)
} finally {
signal.removeEventListener('abort', abort)
src/main/agent/continue-runtime.test.ts
+85
View File
@@ -1,5 +1,6 @@
import { beforeEach, describe, expect, it, vi } from 'vitest'
import type { RuntimeEvent } from './runtime'
import { randomUUID } from 'node:crypto'
import {
ContinueHostRunError,
type ContinueHostAdapterOptions
@@ -632,6 +633,90 @@ describe('ContinueAgentRuntime', () => {
])
})
it('fails instead of silently dropping an overflowing stream queue', async () => {
mocks.runHost.mockImplementation(
async (
_prompt,
_signal,
_authorize,
options
) => {
for (let index = 0; index < 1_001; index += 1) {
options?.onEvent?.({
type: 'text',
delta: String(index)
})
}
return { text: 'done', streamedText: true }
}
)
const stream = createRuntime().run(
{
requestId: randomUUID(),
conversationId: 'overflow-conversation',
prompt: 'test'
},
new AbortController().signal
)
await expect(async () => {
for await (const _event of stream) {
void _event
}
}).rejects.toThrow('流式事件积压超过安全限制')
})
it('aborts the host run when stream consumption ends early', async () => {
let resolveHost: (() => void) | undefined
const hostFinished = new Promise<void>((resolve) => {
resolveHost = resolve
})
let hostSignal: AbortSignal | undefined
mocks.runHost.mockImplementation(
async (
_prompt,
signal,
_authorize,
options
) => {
hostSignal = signal
await options?.onEvent?.({
type: 'text',
delta: 'partial'
})
await new Promise<void>((resolve) => {
signal.addEventListener(
'abort',
() => {
resolve()
resolveHost?.()
},
{ once: true }
)
})
throw signal.reason
}
)
const stream = createRuntime().run(
{
requestId: randomUUID(),
conversationId: 'early-close-conversation',
prompt: 'test'
},
new AbortController().signal
)
await expect(stream.next()).resolves.toMatchObject({
value: { type: 'status' }
})
await expect(stream.next()).resolves.toMatchObject({
value: { type: 'text', delta: 'partial' }
})
await stream.return()
await hostFinished
expect(hostSignal?.aborted).toBe(true)
})
it('emits terminal tool audits before a failed Continue run', async () => {
mocks.runHost.mockRejectedValue(
new ContinueHostRunError('Continue failed', {
src/main/agent/continue-runtime.ts
+26 -3
View File
@@ -51,6 +51,7 @@ export type ContinueRuntimeOptions = {
// The prompt reaches the Continue host through a local HTTP POST body, so no
// platform command-line limit applies to it.
const MAX_CONTINUE_PROMPT_CHARACTERS = 128_000
const MAX_QUEUED_STREAM_EVENTS = 1_000
const scopedReadToolNameSet = new Set<string>(scopedReadToolNames)
function continueToolFailureMessage(tool: ContinueHostTool): string {
@@ -331,7 +332,12 @@ export class ContinueAgentRuntime implements AgentRuntime {
let streamFinished = false
let streamResult: ContinueHostRunResult | undefined
let streamError: unknown
const hostController = new AbortController()
const hostSignal = AbortSignal.any([signal, hostController.signal])
const onEvent = (event: ContinueHostStreamEvent): void => {
if (queuedEvents.length >= MAX_QUEUED_STREAM_EVENTS) {
throw new Error('Continue 流式事件积压超过安全限制')
}
queuedEvents.push(event)
wakeStream?.()
wakeStream = undefined
@@ -339,7 +345,7 @@ export class ContinueAgentRuntime implements AgentRuntime {
const hostRun = host
.run(
conversationContext,
signal,
hostSignal,
authorize,
{
workMode: request.workMode,
@@ -361,7 +367,7 @@ export class ContinueAgentRuntime implements AgentRuntime {
wakeStream?.()
wakeStream = undefined
})
try {
while (!streamFinished || queuedEvents.length > 0) {
if (queuedEvents.length === 0) {
await new Promise<void>((resolve) => {
@@ -385,7 +391,12 @@ export class ContinueAgentRuntime implements AgentRuntime {
false
)
}
} finally {
hostController.abort(new Error('Continue 流式消费已结束'))
wakeStream?.()
wakeStream = undefined
await hostRun
}
if (streamError) {
throw streamError
}
@@ -396,7 +407,19 @@ export class ContinueAgentRuntime implements AgentRuntime {
} catch (error) {
if (error instanceof ContinueHostRunError) {
for (const tool of error.tools) {
yield toContinueToolEvent(request.requestId, tool, true)
const terminalEvent = toContinueToolEvent(
request.requestId,
tool,
true
)
const previous = emittedTools.get(tool.callId)
if (
!previous ||
previous.state !== tool.state ||
previous.error !== tool.error
) {
yield terminalEvent
}
}
}
throw error
src/main/agent/model-runtime.test.ts
+222
View File
@@ -321,6 +321,228 @@ describe('ModelAgentRuntime', () => {
await expect(consume()).rejects.toThrow('意外中断')
})
it('rejects malformed SSE JSON instead of silently skipping it', async () => {
const runtime = new ModelAgentRuntime({
apiKey: 'test-key',
baseUrl: 'https://bigtoken.ai',
model: 'sonnet-5',
protocol: 'anthropic-messages',
authentication: 'api-key',
fetcher: vi.fn<typeof fetch>(async () =>
new Response('data: {invalid}\n\n', {
status: 200,
headers: { 'content-type': 'text/event-stream' }
})
)
})
const consume = async (): Promise<void> => {
for await (const _event of runtime.run(
{
requestId: crypto.randomUUID(),
conversationId: crypto.randomUUID(),
prompt: 'test'
},
new AbortController().signal
)) {
void _event
}
}
await expect(consume()).rejects.toThrow('无效的流式 JSON')
})
it('parses CRLF event separators split across response chunks', async () => {
const payload = createEventStream('split CRLF').replaceAll('\n', '\r\n')
const splitAt = payload.indexOf('\r\n\r\n') + 3
const body = new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(
new TextEncoder().encode(payload.slice(0, splitAt))
)
controller.enqueue(
new TextEncoder().encode(payload.slice(splitAt))
)
controller.close()
}
})
const runtime = new ModelAgentRuntime({
apiKey: 'test-key',
baseUrl: 'https://bigtoken.ai',
model: 'sonnet-5',
protocol: 'anthropic-messages',
authentication: 'api-key',
fetcher: vi.fn<typeof fetch>(async () =>
new Response(body, {
status: 200,
headers: { 'content-type': 'text/event-stream' }
})
)
})
const events = []
for await (const event of runtime.run(
{
requestId: crypto.randomUUID(),
conversationId: crypto.randomUUID(),
prompt: 'test'
},
new AbortController().signal
)) {
events.push(event)
}
expect(events).toContainEqual(
expect.objectContaining({
type: 'text',
delta: 'split CRLF'
})
)
})
it('aborts a model request that exceeds the runtime timeout', async () => {
vi.useFakeTimers()
try {
const runtime = new ModelAgentRuntime({
apiKey: 'test-key',
baseUrl: 'https://bigtoken.ai',
model: 'sonnet-5',
protocol: 'anthropic-messages',
authentication: 'api-key',
requestTimeoutMs: 50,
fetcher: vi.fn<typeof fetch>(
async (_input, init) =>
new Promise<Response>((_resolve, reject) => {
init?.signal?.addEventListener(
'abort',
() => reject(init.signal?.reason),
{ once: true }
)
})
)
})
const stream = runtime.run(
{
requestId: crypto.randomUUID(),
conversationId: crypto.randomUUID(),
prompt: 'test'
},
new AbortController().signal
)
await expect(stream.next()).resolves.toMatchObject({
value: { type: 'status' }
})
const result = stream.next()
const assertion = expect(result).rejects.toThrow(
'模型接口请求超时'
)
await vi.advanceTimersByTimeAsync(50)
await assertion
} finally {
vi.useRealTimers()
}
})
it('aborts a stalled response body after headers arrive', async () => {
vi.useFakeTimers()
try {
let responseSignal: AbortSignal | null | undefined
const runtime = new ModelAgentRuntime({
apiKey: 'test-key',
baseUrl: 'https://bigtoken.ai',
model: 'sonnet-5',
protocol: 'anthropic-messages',
authentication: 'api-key',
requestTimeoutMs: 50,
fetcher: vi.fn<typeof fetch>(async (_input, init) => {
responseSignal = init?.signal
return new Response(
new ReadableStream<Uint8Array>({
start(controller) {
init?.signal?.addEventListener(
'abort',
() => controller.error(init.signal?.reason),
{ once: true }
)
}
}),
{
status: 200,
headers: { 'content-type': 'text/event-stream' }
}
)
})
})
const stream = runtime.run(
{
requestId: crypto.randomUUID(),
conversationId: crypto.randomUUID(),
prompt: 'test'
},
new AbortController().signal
)
await expect(stream.next()).resolves.toMatchObject({
value: { type: 'status' }
})
const result = stream.next()
const assertion = expect(result).rejects.toThrow(
'模型接口请求超时'
)
await vi.advanceTimersByTimeAsync(50)
await assertion
expect(responseSignal?.aborted).toBe(true)
} finally {
vi.useRealTimers()
}
})
it('bounds the total ordinary streaming response size', async () => {
const chunk = new TextEncoder().encode(
`data: ${JSON.stringify({
type: 'content_block_delta',
delta: {
type: 'text_delta',
text: 'x'.repeat(65_000)
}
})}\n\n`
)
const body = new ReadableStream<Uint8Array>({
pull(controller) {
controller.enqueue(chunk)
}
})
const runtime = new ModelAgentRuntime({
apiKey: 'test-key',
baseUrl: 'https://bigtoken.ai',
model: 'sonnet-5',
protocol: 'anthropic-messages',
authentication: 'api-key',
fetcher: vi.fn<typeof fetch>(async () =>
new Response(body, {
status: 200,
headers: { 'content-type': 'text/event-stream' }
})
)
})
const consume = async (): Promise<void> => {
for await (const _event of runtime.run(
{
requestId: crypto.randomUUID(),
conversationId: crypto.randomUUID(),
prompt: 'test'
},
new AbortController().signal
)) {
void _event
}
}
await expect(consume()).rejects.toThrow(
'流式响应超过安全限制'
)
})
it('preserves bounded provider error messages', async () => {
const runtime = new ModelAgentRuntime({
apiKey: 'test-key',
src/main/agent/model-runtime.ts
+229 -148
View File
@@ -37,6 +37,7 @@ import {
boundedToolDetail,
safeToolErrorDetail
} from './approval-summary'
import { readBoundedResponseText } from './bounded-response'
type ConversationMessage = {
role: 'user' | 'assistant'
@@ -98,12 +99,14 @@ type ModelToolResponse = {
const maxGeneratedImageBytes = 3_900_000
const maxImageResponseBytes = 5_300_000
const maxChatResponseBytes = 2 * 1024 * 1024
const maxStreamBlockBytes = 1024 * 1024
const maxToolArgumentBytes = 128 * 1024
const maxToolContextBytes = 1024 * 1024
const maxToolCallsPerRun = 40
const maxToolRounds = 24
const maxRepeatedIdenticalCalls = 3
const maxIdenticalRoundsWithoutProgress = 2
const defaultModelRequestTimeoutMs = 10 * 60_000
function getCurrentTimeInstruction(now = new Date()): string {
const systemTime = [
@@ -138,6 +141,7 @@ export type ModelRuntimeOptions = {
webSearchEnabled?: boolean
toolProvider?: ModelToolProviderLike
fetcher?: typeof fetch
requestTimeoutMs?: number
}
function getErrorMessage(value: unknown): string | undefined {
@@ -412,33 +416,48 @@ function createUsageEvent(
}
}
async function readBoundedText(
response: Response,
maxBytes: number
): Promise<string> {
if (!response.body) {
throw new Error('模型接口未返回响应内容')
function createRequestSignal(
signal: AbortSignal,
timeoutMs: number
): {
signal: AbortSignal
clear: () => void
timedOut: () => boolean
} {
const controller = new AbortController()
let timedOut = false
const abortFromCaller = (): void => {
controller.abort(signal.reason)
}
const reader = response.body.getReader()
const chunks: Uint8Array[] = []
let total = 0
try {
while (true) {
const { done, value } = await reader.read()
if (done) {
break
signal.addEventListener('abort', abortFromCaller, { once: true })
const timeout = setTimeout(() => {
if (controller.signal.aborted) {
return
}
total += value.byteLength
if (total > maxBytes) {
await reader.cancel().catch(() => undefined)
throw new Error('模型接口响应超过安全限制')
timedOut = true
controller.abort(new Error('模型接口请求超时'))
}, timeoutMs)
if (signal.aborted) {
abortFromCaller()
}
chunks.push(value)
return {
signal: controller.signal,
clear: () => {
clearTimeout(timeout)
signal.removeEventListener('abort', abortFromCaller)
},
timedOut: () => timedOut
}
} finally {
reader.releaseLock()
}
function normalizeRequestError(
error: unknown,
timedOut: boolean
): never {
if (timedOut) {
throw new Error('模型接口请求超时', { cause: error })
}
return Buffer.concat(chunks, total).toString('utf8')
throw error
}
function parseGeneratedImage(value: unknown): {
@@ -898,6 +917,14 @@ function parseModelToolResponse(
}
}
function getSseData(block: string): string {
return block
.split('\n')
.filter((line) => line.startsWith('data:'))
.map((line) => line.slice(5).trimStart())
.join('\n')
}
function parseStreamBlock(
block: string,
protocol: ModelProtocol
@@ -907,11 +934,7 @@ function parseStreamBlock(
stopped: boolean
usage?: ModelUsageUpdate
} {
const data = block
.split('\n')
.filter((line) => line.startsWith('data:'))
.map((line) => line.slice(5).trimStart())
.join('\n')
const data = getSseData(block)
if (!data) {
return { stopped: false }
}
@@ -925,8 +948,10 @@ function parseStreamBlock(
let event: unknown
try {
event = JSON.parse(data)
} catch {
return { stopped: false }
} catch (error) {
throw new Error('模型接口返回了无效的流式 JSON', {
cause: error
})
}
const error = getErrorMessage(event)
if (error) {
@@ -985,11 +1010,7 @@ function parseStreamBlock(
function parseSseData(
block: string
): { event?: unknown; stopped: boolean } {
const data = block
.split('\n')
.filter((line) => line.startsWith('data:'))
.map((line) => line.slice(5).trimStart())
.join('\n')
const data = getSseData(block)
if (!data) {
return { stopped: false }
}
@@ -998,8 +1019,56 @@ function parseSseData(
}
try {
return { event: JSON.parse(data), stopped: false }
} catch {
return { stopped: false }
} catch (error) {
throw new Error('模型接口返回了无效的流式 JSON', {
cause: error
})
}
}
async function* readBoundedSseBlocks(
response: Response
): AsyncGenerator<string, void, void> {
if (!response.body) {
throw new Error('模型接口未返回流式响应')
}
const reader = response.body.getReader()
const decoder = new TextDecoder()
let buffer = ''
let completed = false
let receivedBytes = 0
try {
while (true) {
const { done, value } = await reader.read()
receivedBytes += value?.byteLength ?? 0
if (receivedBytes > maxChatResponseBytes) {
throw new Error('模型接口流式响应超过安全限制')
}
buffer += decoder.decode(value, { stream: !done })
buffer = buffer.replaceAll('\r\n', '\n')
if (Buffer.byteLength(buffer) > maxStreamBlockBytes) {
throw new Error('模型接口流式响应块超过安全限制')
}
const blocks = buffer.split('\n\n')
buffer = blocks.pop() ?? ''
if (done && buffer.trim()) {
blocks.push(buffer)
buffer = ''
}
for (const block of blocks) {
yield block
}
if (done) {
completed = true
break
}
}
} finally {
if (!completed) {
await reader.cancel().catch(() => undefined)
}
reader.releaseLock()
}
}
@@ -1010,9 +1079,18 @@ export class ModelAgentRuntime implements AgentRuntime {
private readonly knownConversationIds = new Set<string>()
private readonly fetcher: typeof fetch
private readonly toolProvider: ModelToolProviderLike
private readonly requestTimeoutMs: number
constructor(private readonly options: ModelRuntimeOptions) {
this.fetcher = options.fetcher ?? fetch
this.requestTimeoutMs =
options.requestTimeoutMs ?? defaultModelRequestTimeoutMs
if (
!Number.isSafeInteger(this.requestTimeoutMs) ||
this.requestTimeoutMs < 1
) {
throw new Error('模型接口请求超时设置无效')
}
this.toolProvider =
options.toolProvider ??
new ModelToolProvider(
@@ -1073,6 +1151,31 @@ export class ModelAgentRuntime implements AgentRuntime {
return headers
}
private async fetchWithTimeout(
input: URL,
init: RequestInit,
signal: AbortSignal
): Promise<{
response: Response
clear: () => void
timedOut: () => boolean
}> {
const request = createRequestSignal(signal, this.requestTimeoutMs)
try {
return {
response: await this.fetcher(input, {
...init,
signal: request.signal
}),
clear: request.clear,
timedOut: request.timedOut
}
} catch (error) {
request.clear()
return normalizeRequestError(error, request.timedOut())
}
}
async getStatus(): Promise<AgentRuntimeStatus> {
const imageGeneration = this.capability === 'image-generation'
return {
@@ -1123,9 +1226,16 @@ export class ModelAgentRuntime implements AgentRuntime {
)
})
if (!response.ok) {
const responseText = await readBoundedResponseText(response, {
maxBytes: 128 * 1024,
missingBodyMessage: '模型接口未返回响应内容',
tooLargeMessage: '模型接口响应超过安全限制'
})
let detail: string | undefined
try {
detail = getErrorMessage(await response.json())
detail = getErrorMessage(
responseText.trim() ? JSON.parse(responseText) : undefined
)
} catch {
detail = undefined
}
@@ -1276,21 +1386,33 @@ export class ModelAgentRuntime implements AgentRuntime {
model: this.options.model,
prompt: request.prompt.slice(0, 100_000),
n: 1,
quality:
this.options.imageGenerationQuality ??
'auto',
quality: this.options.imageGenerationQuality ?? 'auto',
response_format: 'b64_json'
}
const response = await this.fetcher(this.getEndpoint(), {
const modelRequest = await this.fetchWithTimeout(
this.getEndpoint(),
{
method: 'POST',
headers: this.getHeaders(),
body: JSON.stringify(imageRequest),
body: JSON.stringify(imageRequest)
},
signal
})
const responseText = await readBoundedText(
response,
response.ok ? maxImageResponseBytes : 128 * 1024
)
const response = modelRequest.response
let responseText: string
try {
responseText = await readBoundedResponseText(response, {
maxBytes: response.ok
? maxImageResponseBytes
: 128 * 1024,
missingBodyMessage: '模型接口未返回响应内容',
tooLargeMessage: '模型接口响应超过安全限制'
})
} catch (error) {
return normalizeRequestError(error, modelRequest.timedOut())
} finally {
modelRequest.clear()
}
if (!response.ok) {
let errorPayload: unknown
try {
@@ -1421,17 +1543,23 @@ export class ModelAgentRuntime implements AgentRuntime {
if (Buffer.byteLength(body) > 2 * 1024 * 1024) {
throw new Error('模型工具请求上下文超过 2MB 安全限制')
}
const response = await this.fetcher(this.getEndpoint(), {
const request = await this.fetchWithTimeout(
this.getEndpoint(),
{
method: 'POST',
headers: this.getHeaders(),
body,
body
},
signal
})
if (!response.ok) {
const responseText = await readBoundedText(
response,
128 * 1024
)
const response = request.response
try {
if (!response.ok) {
const responseText = await readBoundedResponseText(response, {
maxBytes: 128 * 1024,
missingBodyMessage: '模型接口未返回响应内容',
tooLargeMessage: '模型接口响应超过安全限制'
})
let detail: string | undefined
try {
detail = getErrorMessage(
@@ -1443,8 +1571,7 @@ export class ModelAgentRuntime implements AgentRuntime {
detail = undefined
}
throw new Error(
detail ??
`模型接口请求失败（HTTP ${response.status}）`
detail ?? `模型接口请求失败（HTTP ${response.status}）`
)
}
if (
@@ -1454,11 +1581,6 @@ export class ModelAgentRuntime implements AgentRuntime {
?.toLocaleLowerCase()
.includes('text/event-stream')
) {
if (!response.body) {
throw new Error('模型接口未返回流式响应')
}
const reader = response.body.getReader()
const decoder = new TextDecoder()
const streamedToolCalls = new Map<
number,
{ arguments: string; id: string; name: string }
@@ -1468,33 +1590,9 @@ export class ModelAgentRuntime implements AgentRuntime {
}
let answer = ''
let reasoning = ''
let buffer = ''
let receivedStop = false
let receivedBytes = 0
let streamEnded = false
try {
while (!receivedStop) {
const { done, value } = await reader.read()
streamEnded = done
receivedBytes += value?.byteLength ?? 0
if (receivedBytes > maxChatResponseBytes) {
throw new Error('模型接口流式响应超过安全限制')
}
buffer += decoder.decode(value, { stream: !done }).replaceAll(
'\r\n',
'\n'
)
if (Buffer.byteLength(buffer) > maxChatResponseBytes) {
throw new Error('模型接口流式响应块超过安全限制')
}
const blocks = buffer.split('\n\n')
buffer = blocks.pop() ?? ''
if (done && buffer.trim()) {
blocks.push(buffer)
buffer = ''
}
for (const block of blocks) {
for await (const block of readBoundedSseBlocks(response)) {
const parsed = parseSseData(block)
if (parsed.stopped) {
receivedStop = true
@@ -1511,9 +1609,7 @@ export class ModelAgentRuntime implements AgentRuntime {
usage,
getUsageUpdate(parsed.event, 'openai')
)
const reasoningDelta = getOpenAIReasoningDelta(
parsed.event
)
const reasoningDelta = getOpenAIReasoningDelta(parsed.event)
if (reasoningDelta) {
reasoning += reasoningDelta
yield {
@@ -1590,16 +1686,6 @@ export class ModelAgentRuntime implements AgentRuntime {
streamedToolCalls.set(index as number, next)
}
}
if (done) {
break
}
}
} finally {
if (!streamEnded) {
await reader.cancel().catch(() => undefined)
}
reader.releaseLock()
}
if (!receivedStop) {
throw new Error('模型接口流式响应意外中断')
}
@@ -1640,17 +1726,20 @@ export class ModelAgentRuntime implements AgentRuntime {
streamed: true
}
}
const responseText = await readBoundedText(
response,
maxChatResponseBytes
)
const responseText = await readBoundedResponseText(response, {
maxBytes: maxChatResponseBytes,
missingBodyMessage: '模型接口未返回响应内容',
tooLargeMessage: '模型接口响应超过安全限制'
})
let payload: unknown
try {
payload = responseText.trim()
? JSON.parse(responseText)
: undefined
} catch (error) {
throw new Error('模型接口返回了无效 JSON', { cause: error })
throw new Error('模型接口返回了无效 JSON', {
cause: error
})
}
const providerError = getErrorMessage(payload)
if (providerError) {
@@ -1664,6 +1753,11 @@ export class ModelAgentRuntime implements AgentRuntime {
? 'anthropic'
: 'openai'
)
} catch (error) {
return normalizeRequestError(error, request.timedOut())
} finally {
request.clear()
}
}
private async *runToolExecution(
@@ -1742,10 +1836,18 @@ export class ModelAgentRuntime implements AgentRuntime {
request.requestId
)
let responseStep = await responseStream.next()
try {
while (!responseStep.done) {
yield responseStep.value
responseStep = await responseStream.next()
}
} finally {
if (!responseStep.done) {
await responseStream
.throw(new Error('模型流式消费已结束'))
.catch(() => undefined)
}
}
const response = responseStep.value
const usage = {
reported: false
@@ -2079,7 +2181,9 @@ export class ModelAgentRuntime implements AgentRuntime {
: responses
? this.getResponsesInput(request)
: this.getOpenAIMessages(request, system)
const response = await this.fetcher(this.getEndpoint(), {
const modelRequest = await this.fetchWithTimeout(
this.getEndpoint(),
{
method: 'POST',
headers: this.getHeaders(),
body: JSON.stringify(
@@ -2108,14 +2212,25 @@ export class ModelAgentRuntime implements AgentRuntime {
},
messages
}
),
)
},
signal
})
)
const response = modelRequest.response
try {
if (!response.ok) {
const responseText = await readBoundedResponseText(response, {
maxBytes: 128 * 1024,
missingBodyMessage: '模型接口未返回响应内容',
tooLargeMessage: '模型接口响应超过安全限制'
})
let detail: string | undefined
try {
detail = getErrorMessage(await response.json())
detail = getErrorMessage(
responseText.trim()
? JSON.parse(responseText)
: undefined
)
} catch {
detail = undefined
}
@@ -2124,41 +2239,13 @@ export class ModelAgentRuntime implements AgentRuntime {
)
}
if (!response.body) {
throw new Error('模型接口未返回流式响应')
}
const reader = response.body.getReader()
const decoder = new TextDecoder()
let buffer = ''
let answer = ''
let receivedStop = false
let streamEnded = false
const usage = {
reported: false
} satisfies ModelUsageAccumulator
try {
while (!receivedStop) {
const { done, value } = await reader.read()
streamEnded = done
buffer += decoder.decode(value, { stream: !done }).replaceAll(
'\r\n',
'\n'
)
if (Buffer.byteLength(buffer) > 1024 * 1024) {
throw new Error('模型接口流式响应块超过安全限制')
}
const blocks = buffer.split('\n\n')
buffer = blocks.pop() ?? ''
if (done && buffer.trim()) {
blocks.push(buffer)
buffer = ''
}
for (const block of blocks) {
for await (const block of readBoundedSseBlocks(response)) {
const parsed = parseStreamBlock(block, this.options.protocol)
if (parsed.usage) {
applyUsageUpdate(usage, parsed.usage)
@@ -2186,17 +2273,6 @@ export class ModelAgentRuntime implements AgentRuntime {
}
}
if (done) {
break
}
}
} finally {
if (!streamEnded) {
await reader.cancel().catch(() => undefined)
}
reader.releaseLock()
}
if (!receivedStop) {
throw new Error('模型接口流式响应意外中断')
}
@@ -2225,6 +2301,11 @@ export class ModelAgentRuntime implements AgentRuntime {
requestId: request.requestId,
type: 'done'
}
} catch (error) {
return normalizeRequestError(error, modelRequest.timedOut())
} finally {
modelRequest.clear()
}
}
async dispose(): Promise<void> {