diff --git a/.changeset/atomic-file-writes.md b/.changeset/atomic-file-writes.md new file mode 100644 index 000000000..10a70074f --- /dev/null +++ b/.changeset/atomic-file-writes.md @@ -0,0 +1,5 @@ +--- +"@pymodel/pythinker-code": patch +--- + +Write files atomically so an interrupted write leaves the previous content intact. diff --git a/.changeset/cancel-prompt-while-starting.md b/.changeset/cancel-prompt-while-starting.md new file mode 100644 index 000000000..b7fe5ac72 --- /dev/null +++ b/.changeset/cancel-prompt-while-starting.md @@ -0,0 +1,5 @@ +--- +"@pymodel/pythinker-code": patch +--- + +Allow cancelling a prompt while it is still starting. diff --git a/.changeset/dsml-parser-chunk-invariance.md b/.changeset/dsml-parser-chunk-invariance.md new file mode 100644 index 000000000..7410fca85 --- /dev/null +++ b/.changeset/dsml-parser-chunk-invariance.md @@ -0,0 +1,5 @@ +--- +"@pymodel/pythinker-code": patch +--- + +Fix DSML and Hermes tool calls in streamed responses being dropped, split, or mistaken for quoted documentation depending on how the response was chunked. diff --git a/.changeset/sensitive-file-symlink-alias.md b/.changeset/sensitive-file-symlink-alias.md new file mode 100644 index 000000000..907dd2cdb --- /dev/null +++ b/.changeset/sensitive-file-symlink-alias.md @@ -0,0 +1,5 @@ +--- +"@pymodel/pythinker-code": patch +--- + +Block file reads and writes that reach a sensitive file through a symlink alias. diff --git a/.changeset/subagent-usage-per-run.md b/.changeset/subagent-usage-per-run.md new file mode 100644 index 000000000..d92781b74 --- /dev/null +++ b/.changeset/subagent-usage-per-run.md @@ -0,0 +1,5 @@ +--- +"@pymodel/pythinker-code": patch +--- + +Report each subagent run's own token usage instead of the agent's lifetime total. diff --git a/.changeset/tool-cancel-holds-file-lease.md b/.changeset/tool-cancel-holds-file-lease.md new file mode 100644 index 000000000..a0090ab88 --- /dev/null +++ b/.changeset/tool-cancel-holds-file-lease.md @@ -0,0 +1,5 @@ +--- +"@pymodel/pythinker-code": patch +--- + +Keep a cancelled tool that ignores the stop signal from overlapping with the next tool on the same file. diff --git a/.changeset/web-fetch-streaming-limit.md b/.changeset/web-fetch-streaming-limit.md new file mode 100644 index 000000000..da731f49f --- /dev/null +++ b/.changeset/web-fetch-streaming-limit.md @@ -0,0 +1,5 @@ +--- +"@pymodel/pythinker-code": patch +--- + +Stop downloading a web page as soon as it exceeds the size limit instead of buffering it first. diff --git a/packages/agent-core-v2/src/agent/prompt/promptService.ts b/packages/agent-core-v2/src/agent/prompt/promptService.ts index 8532cfbb4..4716f68c7 100644 --- a/packages/agent-core-v2/src/agent/prompt/promptService.ts +++ b/packages/agent-core-v2/src/agent/prompt/promptService.ts @@ -187,6 +187,7 @@ interface Record extends PromptSnapshot { readonly launchedDeferred: Deferred; readonly completionDeferred: Deferred; handle: PromptHandle; + cancelReason?: Error; } function bundledSkillBlockCount(message: ContextMessage): number { @@ -227,6 +228,7 @@ export const promptLaunchingKey = defineState('prompt.launching', () => export class AgentPromptService implements IAgentPromptService { declare readonly _serviceBrand: undefined; private active: (Record & { turn: Turn }) | undefined; + private launchingRecord: Record | undefined; private readonly pending: Record[] = []; private readonly steered = new Map(); private readonly reservedPromptIds = new Set(); @@ -455,6 +457,7 @@ export class AgentPromptService implements IAgentPromptService { abort(promptId: string, reason: Error = userCancellationReason()): boolean { if (this.active?.id === promptId) { this.loop.cancel(this.active.turn.id, reason); return true; } + if (this.launchingRecord?.id === promptId) { this.launchingRecord.cancelReason ??= reason; return true; } const index = this.pending.findIndex((item) => item.id === promptId); if (index < 0) throw new Error2(ErrorCodes.PROMPT_NOT_FOUND, `prompt ${promptId} not found`); const [item] = this.pending.splice(index, 1) as [Record]; @@ -466,6 +469,11 @@ export class AgentPromptService implements IAgentPromptService { async drain(reason: Error = userCancellationReason()): Promise { for (const item of this.pending.slice()) this.abort(item.id, reason); + const launching = this.launchingRecord; + if (launching !== undefined) { + this.abort(launching.id, reason); + await launching.launchedDeferred.promise; + } if (this.active !== undefined) this.abort(this.active.id, reason); } @@ -488,6 +496,7 @@ export class AgentPromptService implements IAgentPromptService { clear(): void { for (const item of this.pending.slice()) this.abort(item.id); + if (this.launchingRecord !== undefined) this.abort(this.launchingRecord.id); if (this.active !== undefined) this.abort(this.active.id); this.context.clear(); } @@ -495,17 +504,20 @@ export class AgentPromptService implements IAgentPromptService { private async startNext(): Promise { if (this.active !== undefined || this.launching || this.steering > 0) return; const item = this.pending.shift(); if (item === undefined) return; + if (this.fullCompaction.compacting !== null && this.loop.status().state !== 'running') { this.pending.unshift(item); return; } this.launching = true; + this.launchingRecord = item; try { - if (this.fullCompaction.compacting !== null && this.loop.status().state !== 'running') { this.pending.unshift(item); return; } const { message, captions } = this.extractCompressionCaptions(item.message); await this.materializeDaemonRefs(message); + if (this.settleCancelledLaunch(item)) return; if (await this.blockedByHook(message, false)) { this.appendPrompt(message, captions); item.state = 'blocked'; item.launchedDeferred.resolve(undefined); item.completionDeferred.resolve({ promptId: item.id, result: undefined, state: 'blocked' }); this.publishCompleted(item.id, 'blocked'); return; } - const turn = (await this.loop.enqueue( + if (this.settleCancelledLaunch(item)) return; + const receipt = this.loop.enqueue( new PromptStepRequest( message, captions, @@ -513,10 +525,12 @@ export class AgentPromptService implements IAgentPromptService { item.maxOutputSize, item.infiniteRetry, ), - ).assigned).turn; + ); + const turn = (await receipt.assigned).turn; if (turn === undefined) { this.pending.unshift(item); return; } item.state = 'running'; item.launchedDeferred.resolve(turn); this.active = Object.assign(item, { turn }); this.publishStarted(item); + if (item.cancelReason !== undefined) this.loop.cancel(turn.id, item.cancelReason); void turn.result.then((result) => this.settle(item, result)); } catch { item.state = 'failed'; @@ -524,11 +538,20 @@ export class AgentPromptService implements IAgentPromptService { item.completionDeferred.resolve({ promptId: item.id, result: undefined, state: 'failed' }); this.publishCompleted(item.id, 'failed'); } finally { + this.launchingRecord = undefined; this.launching = false; if (this.active === undefined) void this.startNext(); } } + private settleCancelledLaunch(item: Record): boolean { + if (item.cancelReason === undefined) return false; + item.state = 'cancelled'; item.launchedDeferred.resolve(undefined); + item.completionDeferred.resolve({ promptId: item.id, result: undefined, state: 'cancelled' }); + this.publishAborted(item.id); + return true; + } + private settle(item: Record, result: TurnResult): void { if (this.active?.id !== item.id) return; this.active = undefined; diff --git a/packages/agent-core-v2/src/agent/toolExecutor/toolExecutorService.ts b/packages/agent-core-v2/src/agent/toolExecutor/toolExecutorService.ts index 712e97e34..8e8da042d 100644 --- a/packages/agent-core-v2/src/agent/toolExecutor/toolExecutorService.ts +++ b/packages/agent-core-v2/src/agent/toolExecutor/toolExecutorService.ts @@ -52,7 +52,7 @@ import { type UnavailableToolDescriber, } from './toolExecutor'; import { ToolCallStarted, ToolProgress, ToolResultEvent } from './toolExecutorEvents'; -import { ToolScheduler } from './toolScheduler'; +import { ToolScheduler, type OutstandingEffect } from './toolScheduler'; const ABORT_GRACE_MS = 2_000; const TOOL_OUTPUT_EMPTY = 'Tool output is empty.'; @@ -71,12 +71,15 @@ export interface ToolExecutionTask { export interface ToolExecutionRunResult { readonly result: ToolResult; readonly outcome: ToolExecutionOutcome; + readonly cancelled?: boolean; + readonly effectsSettled?: Promise; } interface TimedToolResult { readonly index: number; readonly result: ToolResult; readonly outcome: ToolExecutionOutcome; + readonly cancelled: boolean; readonly durationMs: number; } @@ -111,6 +114,7 @@ export const toolExecutorDupTypeTurnIdKey = defineState( export class AgentToolExecutorService implements IAgentToolExecutorService { declare readonly _serviceBrand: undefined; + private readonly outstandingEffects = new Set(); private readonly beforeExecuteEmitter = new BeforeToolExecuteEmitter(); readonly onBeforeExecuteTool: Event = this.beforeExecuteEmitter.event; private readonly willExecuteEmitter = new AsyncEmitter(); @@ -302,7 +306,13 @@ export class AgentToolExecutorService implements IAgentToolExecutorService { ); this.dispatchToolResult(call, finalized, options); - this.trackToolCall(call, finalized, timedResult.durationMs, options); + this.trackToolCall( + call, + finalized, + timedResult.durationMs, + options, + timedResult.cancelled || timedResult.outcome === 'aborted', + ); return { toolCallId: call.toolCall.id, @@ -316,8 +326,9 @@ export class AgentToolExecutorService implements IAgentToolExecutorService { result: ToolResult, durationMs: number, options: ToolExecutorExecuteOptions, + cancelled: boolean, ): void { - const outcome = toolTelemetryOutcome(result); + const outcome = toolTelemetryOutcome(result, cancelled); const toolCallId = call.toolCall.id; const dupType = this.toolCallDupTypes.get(toolCallId) ?? 'normal'; this.toolCallDupTypes.delete(toolCallId); @@ -385,6 +396,10 @@ export class AgentToolExecutorService implements IAgentToolExecutorService { return settleError(call.args, call.output, 'preflight-rejected'); } + if (options.signal.aborted) { + return settleError(call.args, abortedToolOutput(call.toolName, options.signal), 'aborted'); + } + let execution: ToolExecution; try { execution = await call.tool.resolveExecution(call.args); @@ -458,7 +473,7 @@ export class AgentToolExecutorService implements IAgentToolExecutorService { tasks: ToolExecutionTask[], signal: AbortSignal, ): AsyncIterable { - const scheduler = new ToolScheduler(); + const scheduler = new ToolScheduler(this.outstandingEffects); const allResults: Array> = []; const pendingResults = new Map>(); @@ -468,13 +483,24 @@ export class AgentToolExecutorService implements IAgentToolExecutorService { accesses: task.accesses, start: async () => { const startedAt = Date.now(); + const run = task.execute(signal); return { - result: task.execute(signal).then(({ result, outcome }) => ({ - index, - result, - outcome, - durationMs: Math.max(0, Date.now() - startedAt), - })), + result: run.then(({ result, outcome, cancelled, effectsSettled }) => { + if (effectsSettled !== undefined) { + this.trackOutstandingEffect(task.accesses, effectsSettled); + } + return { + index, + result, + outcome, + cancelled: cancelled === true, + durationMs: Math.max(0, Date.now() - startedAt), + }; + }), + effectsSettled: run.then( + (value) => value.effectsSettled, + () => undefined, + ), }; }, }); @@ -501,6 +527,14 @@ export class AgentToolExecutorService implements IAgentToolExecutorService { } } + private trackOutstandingEffect(accesses: ToolAccesses, settled: Promise): void { + const effect: OutstandingEffect = { accesses, settled }; + this.outstandingEffects.add(effect); + void settled.finally(() => { + this.outstandingEffects.delete(effect); + }); + } + private async runSingleExecution( call: RunnableToolCall, execution: RunnableToolExecution, @@ -516,12 +550,14 @@ export class AgentToolExecutorService implements IAgentToolExecutorService { abortedToolOutput(call.toolName, signal), ).result, outcome: 'aborted', + cancelled: true, }; } let rawResult: ExecutableToolResult; + let executePromise: Promise; try { - const executePromise = execution.execute({ + executePromise = execution.execute({ turnId: options.turnId, toolCallId: call.toolCall.id, trace: options.trace, @@ -532,7 +568,20 @@ export class AgentToolExecutorService implements IAgentToolExecutorService { this.dispatchToolProgress(call, update, options); }, }); - rawResult = await raceWithAbortGrace(executePromise, signal, call.toolName); + const raced = await raceWithAbortGrace(executePromise, signal); + if (raced.graceExpired) { + return { + result: makeErrorToolResult(call, call.args, abortedToolOutput(call.toolName, signal)) + .result, + outcome: 'executed', + cancelled: true, + effectsSettled: executePromise.then( + () => undefined, + () => undefined, + ), + }; + } + rawResult = raced.value; } catch (error) { const aborted = isAbortError(error) || signal.aborted; const output = aborted @@ -541,6 +590,7 @@ export class AgentToolExecutorService implements IAgentToolExecutorService { return { result: makeErrorToolResult(call, call.args, output).result, outcome: 'executed', + cancelled: aborted, }; } @@ -904,14 +954,12 @@ function normalizeToolResult(result: ExecutableToolResult): ToolResult { return base; } -function toolTelemetryOutcome(result: ToolResult): 'success' | 'error' | 'cancelled' { - if (result.isError !== true) return 'success'; - const text = toolOutputText(result.output).toLowerCase(); - return text.includes('aborted') || - text.includes('cancelled') || - text.includes('manually interrupted') - ? 'cancelled' - : 'error'; +function toolTelemetryOutcome( + result: ToolResult, + cancelled: boolean, +): 'success' | 'error' | 'cancelled' { + if (cancelled) return 'cancelled'; + return result.isError === true ? 'error' : 'success'; } function toolTelemetryErrorType(outcome: 'success' | 'error' | 'cancelled'): 'cancelled' | 'error' { @@ -919,13 +967,6 @@ function toolTelemetryErrorType(outcome: 'success' | 'error' | 'cancelled'): 'ca return 'error'; } -function toolOutputText(output: ToolResult['output']): string { - if (typeof output === 'string') return output; - return output - .filter((part): part is Extract => part.type === 'text') - .map((part) => part.text) - .join(''); -} function isMediaContentPart(part: ContentPart): boolean { return part.type === 'image_url' || part.type === 'audio_url' || part.type === 'video_url'; @@ -938,21 +979,21 @@ function abortedToolOutput(toolName: string, signal: AbortSignal): string { return `Tool "${toolName}" was aborted`; } +type AbortGraceOutcome = + | { readonly graceExpired: false; readonly value: Result } + | { readonly graceExpired: true }; + async function raceWithAbortGrace( executePromise: Promise, signal: AbortSignal, - toolName: string, -): Promise { +): Promise> { let graceTimer: ReturnType | undefined; let onAbort: (() => void) | undefined; - const graceSentinel: Promise = new Promise((resolve) => { + const graceSentinel: Promise> = new Promise((resolve) => { const armTimer = (): void => { graceTimer = setTimeout(() => { - resolve({ - output: abortedToolOutput(toolName, signal), - isError: true, - } as unknown as Result); + resolve({ graceExpired: true }); }, ABORT_GRACE_MS); }; if (signal.aborted) { @@ -964,7 +1005,10 @@ async function raceWithAbortGrace( }); try { - return await Promise.race([executePromise, graceSentinel]); + return await Promise.race([ + executePromise.then((value) => ({ graceExpired: false, value }) as const), + graceSentinel, + ]); } finally { if (graceTimer !== undefined) clearTimeout(graceTimer); if (onAbort !== undefined) { diff --git a/packages/agent-core-v2/src/agent/toolExecutor/toolScheduler.ts b/packages/agent-core-v2/src/agent/toolExecutor/toolScheduler.ts index 64be16d2a..c09c383cb 100644 --- a/packages/agent-core-v2/src/agent/toolExecutor/toolScheduler.ts +++ b/packages/agent-core-v2/src/agent/toolExecutor/toolScheduler.ts @@ -2,28 +2,64 @@ import { ToolAccesses } from '#/tool/toolContract'; export interface ToolCallTask { readonly accesses: ToolAccesses; - readonly start: () => Promise<{ readonly result: Promise }>; + readonly start: () => Promise<{ + readonly result: Promise; + readonly effectsSettled?: Promise; + }>; } +export interface OutstandingEffect { + readonly accesses: ToolAccesses; + readonly settled: Promise; +} + +export const DEFAULT_TOOL_CONCURRENCY = 16; + interface ScheduledToolCallTask extends ToolCallTask { readonly result: ControlledPromise; } +interface ActiveEntry { + readonly accesses: ToolAccesses; +} + type ControlledPromise = Promise & { readonly resolve: (value: Result | PromiseLike) => void; readonly reject: (reason?: unknown) => void; }; export class ToolScheduler { - private readonly activeTasks: Array> = []; + private readonly activeTasks: ActiveEntry[] = []; private queuedTasks: Array> = []; + constructor( + outstanding: Iterable = [], + private readonly maxConcurrency: number = DEFAULT_TOOL_CONCURRENCY, + ) { + for (const effect of outstanding) { + const entry: ActiveEntry = { accesses: effect.accesses }; + this.activeTasks.push(entry); + void effect.settled.then( + () => this.finish(entry), + () => this.finish(entry), + ); + } + } + + get running(): number { + return this.activeTasks.length; + } + + private get atCapacity(): boolean { + return this.activeTasks.length >= this.maxConcurrency; + } + add(task: ToolCallTask): Promise { const result = createControlledPromise(); void result.catch(() => undefined); const scheduledTask: ScheduledToolCallTask = { ...task, result }; - if (this.isBlocked(task, this.queuedTasks)) { + if (this.atCapacity || this.isBlocked(task, this.queuedTasks)) { this.queuedTasks.push(scheduledTask); } else { this.start(scheduledTask); @@ -43,7 +79,7 @@ export class ToolScheduler { private conflictsWithAny( task: ToolCallTask, - candidates: readonly ToolCallTask[], + candidates: readonly ActiveEntry[], ): boolean { return candidates.some((candidate) => ToolAccesses.conflict(task.accesses, candidate.accesses), @@ -52,7 +88,7 @@ export class ToolScheduler { private start(task: ScheduledToolCallTask): void { this.activeTasks.push(task); - let started: Promise<{ readonly result: Promise }>; + let started: ReturnType['start']>; try { started = task.start(); } catch (error) { @@ -62,14 +98,21 @@ export class ToolScheduler { } void started - .then(({ result }) => result) - .then(task.result.resolve, task.result.reject) + .then( + ({ result, effectsSettled }) => { + result.then(task.result.resolve, task.result.reject); + return Promise.allSettled([result, effectsSettled ?? result]); + }, + (error) => { + task.result.reject(error); + }, + ) .finally(() => { this.finish(task); }); } - private finish(task: ScheduledToolCallTask): void { + private finish(task: ActiveEntry): void { const index = this.activeTasks.indexOf(task); if (index >= 0) this.activeTasks.splice(index, 1); this.startQueuedTasks(); @@ -78,7 +121,7 @@ export class ToolScheduler { private startQueuedTasks(): void { const stillQueued: Array> = []; for (const task of this.queuedTasks) { - if (this.isBlocked(task, stillQueued)) { + if (this.atCapacity || this.isBlocked(task, stillQueued)) { stillQueued.push(task); } else { this.start(task); diff --git a/packages/agent-core-v2/src/agent/tools/os/read/readTool.ts b/packages/agent-core-v2/src/agent/tools/os/read/readTool.ts index 85fb3f32a..95d8aa2e8 100644 --- a/packages/agent-core-v2/src/agent/tools/os/read/readTool.ts +++ b/packages/agent-core-v2/src/agent/tools/os/read/readTool.ts @@ -12,6 +12,7 @@ import { import { registerAgentToolService } from '#/agent/toolRegistry/toolContribution'; import { resolvePathAccessPath, + sensitiveTargetError, type WorkspaceConfig, } from '#/tool/path-access'; import { MEDIA_SNIFF_BYTES, detectFileType } from '#/agent/media/file-type'; @@ -245,6 +246,8 @@ export class ReadTool implements IReadTool { if (lease.runtime.identity.generation !== inspected.identity.generation) { return { isError: true, output: 'Runtime changed before execution. Retry the tool call.' }; } + const denied = await sensitiveTargetError(lease.runtime.fs!, args.path, path); + if (denied !== undefined) return { isError: true, output: denied }; const result = await this.execution(lease.runtime.fs!, args, path); return this.resultTruncation.isSpillFilePath(path) ? { ...result, spillExempt: true as const } @@ -302,7 +305,7 @@ export class ReadTool implements IReadTool { output: notReadableFileOutput(args.path), }; } else { - lines = fs.readLines(safePath, { errors: 'strict' }); + lines = fs.readLines(safePath, { errors: 'strict', maxLineBytes: MAX_LINE_LENGTH * 4 }); } const lineOffset = args.line_offset ?? 1; diff --git a/packages/agent-core-v2/src/agent/tools/os/write/writeTool.ts b/packages/agent-core-v2/src/agent/tools/os/write/writeTool.ts index e9cc37ddb..2423b37ce 100644 --- a/packages/agent-core-v2/src/agent/tools/os/write/writeTool.ts +++ b/packages/agent-core-v2/src/agent/tools/os/write/writeTool.ts @@ -14,6 +14,7 @@ import { import { registerAgentToolService } from '#/agent/toolRegistry/toolContribution'; import { resolvePathAccessPath, + sensitiveTargetError, type WorkspaceConfig, } from '#/tool/path-access'; import { toInputJsonSchema } from '#/tool/input-schema'; @@ -70,6 +71,8 @@ export class WriteTool implements IWriteTool { if (lease.runtime.identity.generation !== inspected.identity.generation) { return { isError: true, output: 'Runtime changed before execution. Retry the tool call.' }; } + const denied = await sensitiveTargetError(lease.runtime.fs!, args.path, path); + if (denied !== undefined) return { isError: true, output: denied }; return await this.execution(lease.runtime.fs!, args, path); } finally { lease.dispose(); diff --git a/packages/agent-core-v2/src/app/web/providers/local-fetch-url.ts b/packages/agent-core-v2/src/app/web/providers/local-fetch-url.ts index ff7afe754..723d89fb2 100644 --- a/packages/agent-core-v2/src/app/web/providers/local-fetch-url.ts +++ b/packages/agent-core-v2/src/app/web/providers/local-fetch-url.ts @@ -74,6 +74,32 @@ export class LocalFetchURLProvider implements UrlFetcher { } } + private async readBoundedBody(response: Response): Promise { + if (response.body === null) return ''; + const reader = response.body.getReader(); + const chunks: Uint8Array[] = []; + let received = 0; + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + received += value.byteLength; + if (received > this.maxBytes) { + await reader.cancel().catch(() => {}); + throw new Error2( + ErrorCodes.WEB_FETCH_FAILED, + `Response body too large: more than ${String(this.maxBytes)} bytes were delivered (maxBytes ${String(this.maxBytes)}).`, + { details: { bytes: received, maxBytes: this.maxBytes } }, + ); + } + chunks.push(value); + } + } finally { + reader.releaseLock(); + } + return new TextDecoder('utf-8').decode(Buffer.concat(chunks)); + } + private async readResponse(response: Response): Promise { if (response.status >= 400) { await response.body?.cancel().catch(() => { @@ -98,16 +124,7 @@ export class LocalFetchURLProvider implements UrlFetcher { } } - const body = await response.text(); - - const actualBytes = Buffer.byteLength(body, 'utf8'); - if (actualBytes > this.maxBytes) { - throw new Error2( - ErrorCodes.WEB_FETCH_FAILED, - `Response body too large: ${String(actualBytes)} bytes exceeds maxBytes (${String(this.maxBytes)}).`, - { details: { bytes: actualBytes, maxBytes: this.maxBytes } }, - ); - } + const body = await this.readBoundedBody(response); const contentType = (response.headers.get('content-type') ?? '').toLowerCase(); if (contentType.startsWith('text/plain') || contentType.startsWith('text/markdown')) { diff --git a/packages/agent-core-v2/src/kosong/contract/usage.ts b/packages/agent-core-v2/src/kosong/contract/usage.ts index ad59bd6bc..fee365174 100644 --- a/packages/agent-core-v2/src/kosong/contract/usage.ts +++ b/packages/agent-core-v2/src/kosong/contract/usage.ts @@ -5,6 +5,22 @@ export interface TokenUsage { inputCacheCreation: number; } +export function usageDelta( + after: TokenUsage | undefined, + before: TokenUsage | undefined, +): TokenUsage | undefined { + if (after === undefined) return undefined; + return { + inputOther: Math.max(0, after.inputOther - (before?.inputOther ?? 0)), + output: Math.max(0, after.output - (before?.output ?? 0)), + inputCacheRead: Math.max(0, after.inputCacheRead - (before?.inputCacheRead ?? 0)), + inputCacheCreation: Math.max( + 0, + after.inputCacheCreation - (before?.inputCacheCreation ?? 0), + ), + }; +} + export function inputTotal(usage: TokenUsage): number { return usage.inputOther + usage.inputCacheRead + usage.inputCacheCreation; } diff --git a/packages/agent-core-v2/src/kosong/provider/bases/openai/dsml-tool-parser.ts b/packages/agent-core-v2/src/kosong/provider/bases/openai/dsml-tool-parser.ts index 998e3d0f8..7e7bc42e4 100644 --- a/packages/agent-core-v2/src/kosong/provider/bases/openai/dsml-tool-parser.ts +++ b/packages/agent-core-v2/src/kosong/provider/bases/openai/dsml-tool-parser.ts @@ -1,11 +1,22 @@ import type { StreamedMessagePart, ToolCall } from '#/kosong/contract/message'; -const CONTAINER_OPEN_RE = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*tool_calls\s*>/i; -const CONTAINER_CLOSE_RE = /^<\/\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*tool_calls\s*>/i; -const INVOKE_OPEN_RE = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*invoke(?:\s+[^>]*)?>/i; -const INVOKE_CLOSE_RE = /<\/\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*invoke\s*>/i; -const HERMES_OPEN_RE = /^/i; -const HERMES_CLOSE_RE = /<\/tool_call>/i; +const MARK = String.raw`\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*`; +const CONTAINER_OPEN_RE = new RegExp(String.raw`<${MARK}tool_calls\s*>`, 'yi'); +const CONTAINER_CLOSE_RE = new RegExp(String.raw``, 'yi'); +const INVOKE_OPEN_RE = new RegExp(String.raw`<${MARK}invoke(?:\s+[^>]*)?>`, 'yi'); +const INVOKE_CLOSE_RE = new RegExp(String.raw``, 'gi'); +const PARAM_OPEN_RE = new RegExp(String.raw`<${MARK}parameter\s+([^>]*?)>`, 'yi'); +const PARAM_CLOSE_RE = new RegExp(String.raw``, 'gi'); +const HERMES_OPEN_RE = //yi; +const HERMES_CLOSE_RE = /<\/tool_call>/gi; +const MARKED_RE = /<\/?\s*(?:[｜|]|DSML)/yi; +const NAME_ATTR_RE = /\bname\s*=\s*(?:"([^"]+)"|'([^']+)'|([^\s>]+))/i; +const STRING_ATTR_RE = /\bstring\s*=\s*(?:"(true|false)"|'(true|false)'|(true|false))/i; + +export const DSML_MAX_TAG_CHARS = 4096; +export const DSML_MAX_ENVELOPE_CHARS = 2 * 1024 * 1024; +const TAIL_MAX = 1024; +const TAG_NAMES = ['tool_calls', 'tool_call', 'invoke', 'parameter']; function unescapeXml(value: string): string { return value @@ -46,115 +57,165 @@ function parseParameterValue(rawVal: string, isStringAttr: boolean | undefined): return unescaped; } -function parseInvokeBody(invokeContent: string): Record { - const paramRegex = - /<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*parameter\s+([^>]*?)>([\s\S]*?)<\/\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*parameter\s*>/gi; - const args: Record = {}; - let paramFound = false; - let match: RegExpExecArray | null = null; - - while ((match = paramRegex.exec(invokeContent)) !== null) { - const attrStr = match[1] ?? ''; - const rawVal = match[2] ?? ''; - const nameMatch = /\bname\s*=\s*(?:"([^"]+)"|'([^']+)'|([^\s>]+))/i.exec(attrStr); - const paramName = nameMatch ? (nameMatch[1] ?? nameMatch[2] ?? nameMatch[3]) : undefined; - if (paramName) { - paramFound = true; - const stringAttrMatch = - /\bstring\s*=\s*(?:"(true|false)"|'(true|false)'|(true|false))/i.exec(attrStr); - const stringAttrVal = stringAttrMatch - ? (stringAttrMatch[1] ?? stringAttrMatch[2] ?? stringAttrMatch[3]) - : undefined; - const isStringAttr = - stringAttrVal !== undefined ? stringAttrVal.toLowerCase() === 'true' : undefined; - args[paramName] = parseParameterValue(rawVal, isStringAttr); - } - } +function attrValue(match: RegExpExecArray | null): string | undefined { + if (!match) return undefined; + return match[1] ?? match[2] ?? match[3]; +} - if (paramFound) { - return args; - } +function skipWhitespace(text: string, from: number): number { + let pos = from; + while (pos < text.length && /\s/.test(text[pos] as string)) pos += 1; + return pos; +} - const trimmed = invokeContent.trim(); - if (trimmed.startsWith('{') && trimmed.endsWith('}')) { +function stickyExec(re: RegExp, text: string, at: number): RegExpExecArray | null { + re.lastIndex = at; + return re.exec(text); +} + +export function parseInvokeBody(invokeContent: string): Record | null { + const args: Record = Object.create(null) as Record; + let pos = skipWhitespace(invokeContent, 0); + if (pos === invokeContent.length) return args; + + if (invokeContent[pos] === '{') { + const trimmed = invokeContent.trim(); + if (!trimmed.endsWith('}')) return null; try { - const parsed = JSON.parse(trimmed); - if (typeof parsed === 'object' && parsed !== null && !Array.isArray(parsed)) { - return parsed as Record; + const parsed: unknown = JSON.parse(trimmed); + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) return null; + for (const key of Object.keys(parsed)) { + args[key] = (parsed as Record)[key]; } - } catch {} + return args; + } catch { + return null; + } + } + + while (pos < invokeContent.length) { + const open = stickyExec(PARAM_OPEN_RE, invokeContent, pos); + if (!open) return null; + const attrStr = open[1] ?? ''; + const paramName = attrValue(NAME_ATTR_RE.exec(attrStr)); + if (!paramName) return null; + const valueStart = open.index + open[0].length; + const close = stickyExec(PARAM_CLOSE_RE, invokeContent, valueStart); + if (!close) return null; + const stringAttrVal = attrValue(STRING_ATTR_RE.exec(attrStr)); + const isStringAttr = + stringAttrVal !== undefined ? stringAttrVal.toLowerCase() === 'true' : undefined; + args[paramName] = parseParameterValue( + invokeContent.slice(valueStart, close.index), + isStringAttr, + ); + pos = skipWhitespace(invokeContent, close.index + close[0].length); } + return args; +} - return {}; +function newCallId(): string { + return `call_${crypto.randomUUID().replaceAll('-', '').slice(0, 24)}`; } function parseInvokeTag(invokeBlock: string): ToolCall | null { - const openMatch = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*invoke\s+([^>]*?)>/i.exec(invokeBlock); + const openMatch = stickyExec(INVOKE_OPEN_RE, invokeBlock, 0); if (!openMatch) return null; - - const attrStr = openMatch[1] ?? ''; - const nameMatch = /\bname\s*=\s*(?:"([^"]+)"|'([^']+)'|([^\s>]+))/i.exec(attrStr); - const toolName = nameMatch ? (nameMatch[1] ?? nameMatch[2] ?? nameMatch[3]) : undefined; + const toolName = attrValue(NAME_ATTR_RE.exec(openMatch[0].slice(0, -1))); if (!toolName) return null; - const closeMatch = INVOKE_CLOSE_RE.exec(invokeBlock); - if (!closeMatch) return null; - const innerContent = invokeBlock.slice(openMatch[0].length, closeMatch.index); - - const args = parseInvokeBody(innerContent); - return { - type: 'function', - id: `call_${crypto.randomUUID().replaceAll('-', '').slice(0, 24)}`, - name: toolName, - arguments: JSON.stringify(args), - }; + const closeMatch = stickyExec(INVOKE_CLOSE_RE, invokeBlock, openMatch[0].length); + if (!closeMatch || closeMatch.index + closeMatch[0].length !== invokeBlock.length) return null; + const args = parseInvokeBody(invokeBlock.slice(openMatch[0].length, closeMatch.index)); + if (args === null) return null; + return { type: 'function', id: newCallId(), name: toolName, arguments: JSON.stringify(args) }; } function parseHermesToolCall(toolCallBlock: string): ToolCall | null { - if (!HERMES_CLOSE_RE.test(toolCallBlock)) return null; const inner = toolCallBlock .replace(/^/i, '') .replace(/<\/tool_call>$/i, '') .trim(); try { - const parsed = JSON.parse(inner); - if (parsed && typeof parsed.name === 'string') { - const args = - typeof parsed.arguments === 'string' - ? parsed.arguments - : JSON.stringify(parsed.arguments ?? {}); - return { - type: 'function', - id: `call_${crypto.randomUUID().replaceAll('-', '').slice(0, 24)}`, - name: parsed.name, - arguments: args, - }; - } - } catch {} - return null; + const parsed: unknown = JSON.parse(inner); + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) return null; + const record = parsed as Record; + if (typeof record['name'] !== 'string') return null; + const rawArgs = record['arguments']; + const args = typeof rawArgs === 'string' ? rawArgs : JSON.stringify(rawArgs ?? {}); + return { type: 'function', id: newCallId(), name: record['name'], arguments: args }; + } catch { + return null; + } } -function isPotentialTagPrefix(s: string): boolean { - if (!s.startsWith('<')) return false; - const lower = s.toLowerCase(); +type TagKind = + | 'container-open' + | 'container-close' + | 'invoke-open' + | 'hermes-open' + | 'partial' + | 'none'; + +function markerPrefixState(lower: string): { rest: string; partial: boolean } { let rest = lower.startsWith(' 0 && rest.length < 4 && 'dsml'.startsWith(rest)) { + return { rest: '', partial: true }; } if (rest.startsWith('dsml')) { rest = rest.slice(4).trimStart(); - if (rest.startsWith('｜') || rest.startsWith('|')) { - rest = rest.slice(1).trimStart(); + if (rest.startsWith('｜') || rest.startsWith('|')) rest = rest.slice(1).trimStart(); + } + return { rest, partial: rest.length === 0 }; +} + +function classifyTag(buffer: string, at: number): { kind: TagKind; length: number } { + const containerOpen = stickyExec(CONTAINER_OPEN_RE, buffer, at); + if (containerOpen) return { kind: 'container-open', length: containerOpen[0].length }; + const containerClose = stickyExec(CONTAINER_CLOSE_RE, buffer, at); + if (containerClose) return { kind: 'container-close', length: containerClose[0].length }; + const invokeOpen = stickyExec(INVOKE_OPEN_RE, buffer, at); + if (invokeOpen) return { kind: 'invoke-open', length: invokeOpen[0].length }; + const hermesOpen = stickyExec(HERMES_OPEN_RE, buffer, at); + if (hermesOpen) return { kind: 'hermes-open', length: hermesOpen[0].length }; + + const lower = buffer.slice(at, at + DSML_MAX_TAG_CHARS + 1).toLowerCase(); + const { rest, partial } = markerPrefixState(lower); + if (partial) return { kind: 'partial', length: 0 }; + if (rest.includes('>')) return { kind: 'none', length: 0 }; + for (const name of TAG_NAMES) { + if (name.startsWith(rest)) return { kind: 'partial', length: 0 }; + if (rest.startsWith(name)) { + const after = rest.slice(name.length); + if (after.length === 0 || /^\s/.test(after)) return { kind: 'partial', length: 0 }; } } - if (rest.length === 0) return true; - const targets = ['tool_calls', 'tool_call', 'invoke', 'parameter']; - return targets.some((t) => t.startsWith(rest) || rest.startsWith(t)); + return { kind: 'none', length: 0 }; +} + +type Mode = 'text' | 'invoke' | 'hermes'; + +interface ContainerHold { + raw: string; + parts: StreamedMessagePart[]; } export class DsmlStreamParser { private _buffer = ''; + private _pos = 0; + private _mode: Mode = 'text'; + private _envelope: string[] = []; + private _envelopeLength = 0; + private _tail = ''; + private _stripWhitespace = false; + private _whitespaceHold = ''; + private _inContainer = false; + private _hold: ContainerHold | null = null; + private _lineStart = true; + private _linePrefix = ''; + private _fence: string | null = null; private _hasExtractedToolCalls = false; get hasExtractedToolCalls(): boolean { @@ -162,139 +223,276 @@ export class DsmlStreamParser { } feed(chunk: string): StreamedMessagePart[] { - this._buffer += chunk; const parts: StreamedMessagePart[] = []; + if (this._mode === 'text') { + this._buffer += chunk; + } else { + this._scanEnvelope(parts, chunk); + } + if (this._mode === 'text') { + this._drainText(parts); + } + return parts; + } - while (this._buffer.length > 0) { - const ltIdx = this._buffer.indexOf('<'); - if (ltIdx === -1) { - parts.push({ type: 'text', text: this._buffer }); - this._buffer = ''; - break; - } - - if (ltIdx > 0) { - parts.push({ type: 'text', text: this._buffer.slice(0, ltIdx) }); - this._buffer = this._buffer.slice(ltIdx); - } + flush(): StreamedMessagePart[] { + const parts: StreamedMessagePart[] = []; + if (this._mode !== 'text') { + const block = this._envelope.join('') + this._tail; + this._leaveEnvelope(); + this._buffer = block; + this._pos = 0; + } + if (this._whitespaceHold.length > 0) { + this._emitText(parts, this._whitespaceHold.replace(/^\r?\n/, '')); + this._whitespaceHold = ''; + } + this._stripWhitespace = false; + if (this._pos < this._buffer.length) { + this._emitText(parts, this._take(this._buffer.length - this._pos)); + } + this._buffer = ''; + this._pos = 0; + if (this._inContainer) { + this._releaseHold(parts, true); + this._inContainer = false; + } + return parts; + } - const openContainer = CONTAINER_OPEN_RE.exec(this._buffer); - if (openContainer) { - this._buffer = this._buffer.slice(openContainer[0].length); - if (/^\s* this._pos) { + this._whitespaceHold += this._take(end - this._pos); + } + if (this._pos >= this._buffer.length) break; + this._stripWhitespace = false; + if (this._buffer[this._pos] === '<') { + this._noteText(this._whitespaceHold); } else { - this._buffer = this._buffer.replace(/^\r?\n/, ''); + this._emitText(parts, this._whitespaceHold.replace(/^\r?\n/, '')); } - continue; + this._whitespaceHold = ''; } - const closeContainer = CONTAINER_CLOSE_RE.exec(this._buffer); - if (closeContainer) { - this._buffer = this._buffer.slice(closeContainer[0].length); - if (/^\s* this._pos) { + this._emitText(parts, this._take(ltIdx - this._pos)); continue; } - const invokeOpen = INVOKE_OPEN_RE.exec(this._buffer); - if (invokeOpen) { - const closeMatch = INVOKE_CLOSE_RE.exec(this._buffer); - if (!closeMatch) { - break; - } - const invokeEnd = closeMatch.index + closeMatch[0].length; - const invokeBlock = this._buffer.slice(0, invokeEnd); - const toolCall = parseInvokeTag(invokeBlock); - if (toolCall) { - this._hasExtractedToolCalls = true; - parts.push(toolCall); - this._buffer = this._buffer.slice(invokeEnd); - if (/^\s* DSML_MAX_TAG_CHARS) { + this._emitPlainTag(parts); + continue; } - const toolCallEnd = closeMatch.index + closeMatch[0].length; - const block = this._buffer.slice(0, toolCallEnd); - const toolCall = parseHermesToolCall(block); - if (toolCall) { - this._hasExtractedToolCalls = true; - parts.push(toolCall); - this._buffer = this._buffer.slice(toolCallEnd); - if (/^\s* DSML_MAX_ENVELOPE_CHARS) { + this._abandonEnvelope(parts, block, rest); + return; + } + const call = mode === 'invoke' ? parseInvokeTag(block) : parseHermesToolCall(block); + this._record(block); + if (call) { + this._emitCall(parts, call); + this._stripWhitespace = true; + } else { + this._emitText(parts, block); } + this._buffer = rest; + this._pos = 0; + return; + } - parts.push({ type: 'text', text: this._buffer.slice(0, nextLt) }); - this._buffer = this._buffer.slice(nextLt); + if (this._envelopeLength + probe.length >= DSML_MAX_ENVELOPE_CHARS) { + const block = this._envelope.join('') + probe; + this._leaveEnvelope(); + this._abandonEnvelope(parts, block, ''); + return; } - return parts; + const lastLt = probe.lastIndexOf('<'); + if (lastLt === -1 || probe.length - lastLt > TAIL_MAX) { + this._envelope.push(probe); + this._envelopeLength += probe.length; + this._tail = ''; + } else { + const committed = probe.slice(0, lastLt); + this._envelope.push(committed); + this._envelopeLength += committed.length; + this._tail = probe.slice(lastLt); + } } - flush(): StreamedMessagePart[] { - const parts: StreamedMessagePart[] = []; - if (this._buffer.length > 0) { - const invokeOpen = INVOKE_OPEN_RE.exec(this._buffer); - if (invokeOpen && INVOKE_CLOSE_RE.test(this._buffer)) { - const toolCall = parseInvokeTag(this._buffer); - if (toolCall) { - this._hasExtractedToolCalls = true; - parts.push(toolCall); - this._buffer = ''; - return parts; - } + private _abandonEnvelope(parts: StreamedMessagePart[], block: string, rest: string): void { + const nextLt = block.indexOf('<', 1); + const cut = nextLt === -1 ? block.length : nextLt; + const text = block.slice(0, cut); + this._record(text); + this._emitText(parts, text); + this._buffer = block.slice(cut) + rest; + this._pos = 0; + } + + private _take(length: number): string { + const taken = this._buffer.slice(this._pos, this._pos + length); + this._pos += length; + this._record(taken); + return taken; + } + + private _record(text: string): void { + if (this._hold === null) return; + this._hold.raw += text; + } + + private _releaseHold(parts: StreamedMessagePart[], asText: boolean): void { + const hold = this._hold; + this._hold = null; + if (hold === null) return; + if (asText) { + if (hold.raw.length > 0) parts.push({ type: 'text', text: hold.raw }); + } else { + parts.push(...hold.parts); + } + } + + private _emitCall(parts: StreamedMessagePart[], call: ToolCall): void { + this._hasExtractedToolCalls = true; + this._releaseHold(parts, false); + parts.push(call); + } + + private _emitText(parts: StreamedMessagePart[], text: string): void { + if (text.length === 0) return; + this._noteText(text); + const part: StreamedMessagePart = { type: 'text', text }; + if (this._hold !== null) { + this._hold.parts.push(part); + if (this._hold.raw.length > DSML_MAX_ENVELOPE_CHARS) { + this._releaseHold(parts, true); } - const hermesOpen = HERMES_OPEN_RE.exec(this._buffer); - if (hermesOpen && HERMES_CLOSE_RE.test(this._buffer)) { - const toolCall = parseHermesToolCall(this._buffer); - if (toolCall) { - this._hasExtractedToolCalls = true; - parts.push(toolCall); - this._buffer = ''; - return parts; + return; + } + parts.push(part); + } + + private _noteText(text: string): void { + for (const ch of text) { + if (ch === '\n') { + this._lineStart = true; + this._linePrefix = ''; + continue; + } + if (!this._lineStart) continue; + if (ch === ' ' && this._linePrefix.length < 3 && /^ *$/.test(this._linePrefix)) { + this._linePrefix += ch; + continue; + } + if (ch !== '`' && ch !== '~') { + this._lineStart = false; + continue; + } + const run = this._linePrefix.trimStart(); + if (run.length > 0 && run[0] !== ch) { + this._lineStart = false; + continue; + } + this._linePrefix += ch; + if (run.length + 1 === 3) { + this._lineStart = false; + if (this._fence === null) { + this._fence = ch; + } else if (this._fence === ch) { + this._fence = null; } } - parts.push({ type: 'text', text: this._buffer }); - this._buffer = ''; } - return parts; } } @@ -315,6 +513,6 @@ export function extractDsmlToolCalls(text: string): { } } - const cleanText = toolCalls.length > 0 ? textParts.join('').trim() : text; - return { cleanText, toolCalls }; + const joined = textParts.join(''); + return { cleanText: toolCalls.length > 0 ? joined.trim() : joined, toolCalls }; } diff --git a/packages/agent-core-v2/src/kosong/provider/bases/openai/openai-legacy.ts b/packages/agent-core-v2/src/kosong/provider/bases/openai/openai-legacy.ts index 25ad3daf7..bdd8e3ed9 100644 --- a/packages/agent-core-v2/src/kosong/provider/bases/openai/openai-legacy.ts +++ b/packages/agent-core-v2/src/kosong/provider/bases/openai/openai-legacy.ts @@ -429,7 +429,8 @@ export class OpenAILegacyStreamedMessage implements StreamedMessage { let text = message.content ?? null; let extractedToolCalls: ToolCall[] = []; - if (text) { + const hasNativeToolCalls = (message.tool_calls ?? []).some(isFunctionToolCall); + if (text && !hasNativeToolCalls) { const parsed = extractDsmlToolCalls(text); if (parsed.toolCalls.length > 0) { text = parsed.cleanText; @@ -465,6 +466,8 @@ export class OpenAILegacyStreamedMessage implements StreamedMessage { ): AsyncGenerator { const bufferedToolCalls = new Map(); const dsmlParser = new DsmlStreamParser(); + const recoveredToolCalls: ToolCall[] = []; + let nativeToolCallsSeen = false; try { for await (const chunk of response) { @@ -494,11 +497,16 @@ export class OpenAILegacyStreamedMessage implements StreamedMessage { if (delta.content) { for (const part of dsmlParser.feed(delta.content)) { - yield part; + if (part.type === 'function') { + recoveredToolCalls.push(part); + } else { + yield part; + } } } for (const toolCall of delta.tool_calls ?? []) { + nativeToolCallsSeen = true; for (const part of convertChatCompletionStreamToolCall(toolCall, bufferedToolCalls)) { yield part; } @@ -506,10 +514,17 @@ export class OpenAILegacyStreamedMessage implements StreamedMessage { } for (const part of dsmlParser.flush()) { - yield part; + if (part.type === 'function') { + recoveredToolCalls.push(part); + } else { + yield part; + } } - if (dsmlParser.hasExtractedToolCalls) { + if (!nativeToolCallsSeen && recoveredToolCalls.length > 0) { this._hasExtractedToolCalls = true; + for (const toolCall of recoveredToolCalls) { + yield toolCall; + } } } catch (error: unknown) { throw convertOpenAIError(error, this._convertErrorHook); diff --git a/packages/agent-core-v2/src/os/backends/node-local/hostFsService.ts b/packages/agent-core-v2/src/os/backends/node-local/hostFsService.ts index 9591a2e5b..ca776a700 100644 --- a/packages/agent-core-v2/src/os/backends/node-local/hostFsService.ts +++ b/packages/agent-core-v2/src/os/backends/node-local/hostFsService.ts @@ -8,15 +8,16 @@ import { realpath as nodeRealpath, rm, stat as nodeStat, - writeFile, } from 'node:fs/promises'; import { LifecycleScope } from '#/app/scopes'; import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; import { decodeTextWithErrors, type TextDecodeErrors } from '#/_base/execEnv/decodeText'; import { type HostDirEntry, type HostFileStat, IHostFileSystem } from '#/os/interface/hostFileSystem'; -import { toHostFsError } from '#/os/interface/hostFsErrors'; +import { toHostFsError, HostFsError, OsFsErrors } from '#/os/interface/hostFsErrors'; +import { atomicWrite } from '#/_base/utils/fs'; +const NEWLINE = Buffer.from([0x0a]); const READ_CHUNK_SIZE = 64 * 1024; function isUtf8Encoding(encoding: BufferEncoding): boolean { @@ -37,6 +38,33 @@ function* splitLinesKeepingTerminator(text: string): Generator { } } +function trimToValidUtf8Boundary(buf: Buffer): Buffer { + let i = buf.length - 1; + let continuations = 0; + while (i >= 0 && continuations < 3) { + const byte = buf[i]; + if (byte === undefined || (byte & 0xc0) !== 0x80) break; + continuations += 1; + i -= 1; + } + if (i < 0) return buf; + const lead = buf[i]; + if (lead === undefined) return buf; + if ((lead & 0x80) === 0) { + return continuations === 0 ? buf : buf.subarray(0, i + 1); + } + let needed = 0; + if ((lead & 0xe0) === 0xc0) needed = 1; + else if ((lead & 0xf0) === 0xe0) needed = 2; + else if ((lead & 0xf8) === 0xf0) needed = 3; + else return buf.subarray(0, i); + + if (continuations < needed) { + return buf.subarray(0, i); + } + return buf; +} + export class HostFileSystem implements IHostFileSystem { declare readonly _serviceBrand: undefined; @@ -57,8 +85,23 @@ export class HostFileSystem implements IHostFileSystem { } async writeText(path: string, data: string): Promise { + await this._replaceAtomically(path, data); + } + + private async _replaceAtomically(path: string, data: string | Uint8Array): Promise { try { - await writeFile(path, data, 'utf8'); + let target = path; + let mode: number | undefined; + try { + const existing = await nodeStat(path); + if (existing.isFile()) { + target = await nodeRealpath(path); + mode = existing.mode & 0o7777; + } + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'ENOENT') throw error; + } + await atomicWrite(target, data, undefined, mode); } catch (error) { throw toHostFsError(error, { path, op: 'write' }); } @@ -93,28 +136,27 @@ export class HostFileSystem implements IHostFileSystem { } async writeBytes(path: string, data: Uint8Array): Promise { - try { - await writeFile(path, data); - } catch (error) { - throw toHostFsError(error, { path, op: 'write' }); - } + await this._replaceAtomically(path, data); } async *readLines( path: string, - options?: { encoding?: BufferEncoding; errors?: TextDecodeErrors }, + options?: { encoding?: BufferEncoding; errors?: TextDecodeErrors; maxLineBytes?: number }, ): AsyncGenerator { try { const encoding = options?.encoding ?? 'utf-8'; const errors = options?.errors ?? 'strict'; if (!isUtf8Encoding(encoding)) { + if (options?.maxLineBytes !== undefined) { + throw new HostFsError(OsFsErrors.codes.OS_FS_UNKNOWN, 'maxLineBytes is only supported for UTF-8 encoding'); + } const content = decodeTextWithErrors(await readFile(path), encoding, errors); yield* splitLinesKeepingTerminator(content); return; } - yield* this._readUtf8Lines(path, errors); + yield* this._readUtf8Lines(path, errors, options?.maxLineBytes ?? Number.POSITIVE_INFINITY); } catch (error) { throw toHostFsError(error, { path, op: 'read' }); } @@ -123,13 +165,38 @@ export class HostFileSystem implements IHostFileSystem { private async *_readUtf8Lines( path: string, errors: TextDecodeErrors, + maxLineBytes: number, ): AsyncGenerator { const fh = await open(path, 'r'); try { const buf = Buffer.alloc(READ_CHUNK_SIZE); let pending: Buffer[] = []; + let pendingBytes = 0; + let pendingTruncated = false; let pendingOffset = 0; let fileOffset = 0; + const retain = (piece: Buffer): void => { + const room = maxLineBytes - pendingBytes; + if (room <= 0) { + pendingTruncated = true; + return; + } + if (piece.length > room) pendingTruncated = true; + const kept = Buffer.from(piece.subarray(0, Math.min(piece.length, room))); + pending.push(kept); + pendingBytes += kept.length; + }; + const takeLine = (piece: Buffer): Buffer => { + if (pending.length === 0 && piece.length <= maxLineBytes) return piece; + retain(piece); + const kept = Buffer.concat(pending); + const safe = pendingTruncated ? trimToValidUtf8Boundary(kept) : kept; + const line = pendingTruncated && piece.at(-1) === 0x0a ? Buffer.concat([safe, NEWLINE]) : safe; + pending = []; + pendingBytes = 0; + pendingTruncated = false; + return line; + }; while (true) { const { bytesRead } = await fh.read(buf, 0, buf.length, null); @@ -140,24 +207,22 @@ export class HostFileSystem implements IHostFileSystem { for (let i = 0; i < chunk.length; i += 1) { const byte = chunk[i]; if (byte !== 0x0a) continue; - const piece = chunk.subarray(lineStart, i + 1); const lineOffset = pending.length === 0 ? fileOffset + lineStart : pendingOffset; - const line = pending.length === 0 ? piece : Buffer.concat([...pending, piece]); + const line = takeLine(chunk.subarray(lineStart, i + 1)); yield decodeTextWithErrors(line, 'utf-8', errors, lineOffset !== 0); - pending = []; lineStart = i + 1; } if (lineStart < chunk.length) { - const tail = Buffer.from(chunk.subarray(lineStart)); if (pending.length === 0) pendingOffset = fileOffset + lineStart; - pending.push(tail); + retain(chunk.subarray(lineStart)); } fileOffset += bytesRead; } if (pending.length > 0) { - const line = Buffer.concat(pending); + const kept = Buffer.concat(pending); + const line = pendingTruncated ? trimToValidUtf8Boundary(kept) : kept; yield decodeTextWithErrors(line, 'utf-8', errors, pendingOffset !== 0); } } finally { diff --git a/packages/agent-core-v2/src/os/interface/hostFileSystem.ts b/packages/agent-core-v2/src/os/interface/hostFileSystem.ts index cadd91538..04929660a 100644 --- a/packages/agent-core-v2/src/os/interface/hostFileSystem.ts +++ b/packages/agent-core-v2/src/os/interface/hostFileSystem.ts @@ -30,7 +30,7 @@ export interface IHostFileSystem { writeBytes(path: string, data: Uint8Array): Promise; readLines( path: string, - options?: { encoding?: BufferEncoding; errors?: TextDecodeErrors }, + options?: { encoding?: BufferEncoding; errors?: TextDecodeErrors; maxLineBytes?: number }, ): AsyncGenerator; createExclusive(path: string, data: Uint8Array): Promise; stat(path: string): Promise; diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index 105e66a9a..8a58f9686 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -1,7 +1,7 @@ import { join } from 'pathe'; import { IInstantiationService } from '#/_base/di/instantiation'; -import { Disposable, toDisposable } from '#/_base/di/lifecycle'; +import { Disposable, toDisposable, type IDisposable } from '#/_base/di/lifecycle'; import { type CollectionView } from '#/_base/di/collection'; import { Emitter } from '#/_base/event'; import { onUnexpectedError } from '#/_base/errors/unexpectedError'; @@ -159,7 +159,7 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle if (this.records.get(id) === record) { const fallback = [...this.contributions.values()] .filter((candidate) => candidate.active && getAgentRuntimeDefinitionId(candidate.definition) === id) - .sort((left, right) => (right.providerGeneration ?? right.generation) - (left.providerGeneration ?? left.generation))[0]; + .toSorted((left, right) => (right.providerGeneration ?? right.generation) - (left.providerGeneration ?? left.generation))[0]; if (fallback !== undefined) this.records.set(id, fallback); else this.records.delete(id); } @@ -294,6 +294,7 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle await managed.handle.dispose(); } catch { } } + if (didCreate) await this.sessionMetadata.unregisterAgent(agentId).catch(() => undefined); if (!finalizerArmed) eventBus?.deactivateAgent(agent); if (didCreate) this.onDidCloseEmitter.fire(agent); throw error; @@ -443,29 +444,58 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle async remove(agent: AgentContext): Promise { const managed = this.roster.get(agent.agentId); - if (managed === undefined || managed.context !== agent || managed.closing) return; + if (managed === undefined || managed.context !== agent) return; + if (managed.closePromise !== undefined) return managed.closePromise; + if (managed.closing) return; managed.closing = true; + managed.closePromise = this.close(agent, managed); + return managed.closePromise; + } + + private async close(agent: AgentContext, managed: ManagedAgent): Promise { + const failures: unknown[] = []; + const phase = async (work: () => Promise | void): Promise => { + try { + await work(); + } catch (error) { + failures.push(error); + } + }; this.onWillCloseEmitter.fire(agent); const handle = managed.handle; - await handle.accessor.get(IAgentTaskService).stopAllOnExit('Session closed'); const loop = handle.accessor.get(IAgentLoopService); - const compaction = handle.accessor.get(IAgentFullCompactionService).compacting; - const compactionSettled = compaction?.promise.catch(() => undefined) ?? Promise.resolve(); const reason = abortError('Agent removed'); - const prompt = handle.accessor.get(IAgentPromptService); - for (const turnId of loop.status().pendingTurnIds) { - loop.cancel(turnId, reason); - } - loop.cancel(undefined, reason); - if (compaction !== null && !compaction.abortController.signal.aborted) { - compaction.abortController.abort(reason); + let quiescence: IDisposable | undefined; + try { + await phase(() => handle.accessor.get(IAgentTaskService).stopAllOnExit('Session closed')); + const compaction = handle.accessor.get(IAgentFullCompactionService).compacting; + const compactionSettled = compaction?.promise.catch(() => undefined) ?? Promise.resolve(); + const prompt = handle.accessor.get(IAgentPromptService); + await phase(async () => { + await prompt.drain(reason); + for (const turnId of loop.status().pendingTurnIds) { + loop.cancel(turnId, reason); + } + loop.cancel(undefined, reason); + if (compaction !== null && !compaction.abortController.signal.aborted) { + compaction.abortController.abort(reason); + } + await Promise.all([loop.settled(), compactionSettled]); + }); + await phase(() => { + quiescence = loop.tryAcquireQuiescence(); + }); + await phase(() => handle.accessor.get(IEventDispatcher).flush()); + await phase(() => managed.runtimeSet.close()); + managed.killSpace(); + await phase(() => handle.dispose()); + } finally { + quiescence?.dispose(); + if (this.roster.get(agent.agentId) === managed) this.roster.delete(agent.agentId); + this.onDidCloseEmitter.fire(agent); } - await Promise.all([loop.settled(), compactionSettled, prompt.drain(reason)]); - await managed.runtimeSet.close(); - managed.killSpace(); - await handle.dispose(); - if (this.roster.get(agent.agentId) === managed) this.roster.delete(agent.agentId); - this.onDidCloseEmitter.fire(agent); + if (failures.length === 1) throw failures[0]; + if (failures.length > 1) throw new AggregateError(failures, `Agent ${agent.agentId} removal failed`); } private managedFor(agent: AgentContext): ManagedAgent | undefined { diff --git a/packages/agent-core-v2/src/session/agentLifecycle/managedAgent.ts b/packages/agent-core-v2/src/session/agentLifecycle/managedAgent.ts index 31735f1de..78e3bd39d 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/managedAgent.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/managedAgent.ts @@ -8,6 +8,7 @@ import { IEventDispatcher } from '#/state/eventDispatcher'; export class ManagedAgent { active = false; closing = false; + closePromise: Promise | undefined; readonly runtimeSet: AgentRuntimeSet; constructor( diff --git a/packages/agent-core-v2/src/session/expertTalk/expertTalkService.ts b/packages/agent-core-v2/src/session/expertTalk/expertTalkService.ts index c46b47822..963a9353d 100644 --- a/packages/agent-core-v2/src/session/expertTalk/expertTalkService.ts +++ b/packages/agent-core-v2/src/session/expertTalk/expertTalkService.ts @@ -52,7 +52,7 @@ import { } from '#/kosong/contract/errors'; import { extractText } from '#/kosong/contract/message'; import { estimateTokensForContentParts } from '#/kosong/contract/tokens'; -import type { TokenUsage } from '#/kosong/contract/usage'; +import { usageDelta } from '#/kosong/contract/usage'; import { IModelCatalog, type Model } from '#/kosong/model/catalog'; import type { ModelRequester } from '#/kosong/model/modelRequester'; import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; @@ -1905,22 +1905,6 @@ function hasMarkdownSections(text: string, sections: readonly string[]): boolean return sections.every((section) => headings.has(section)); } -function usageDelta( - after: TokenUsage | undefined, - before: TokenUsage | undefined, -): TokenUsage | undefined { - if (after === undefined) return undefined; - return { - inputOther: Math.max(0, after.inputOther - (before?.inputOther ?? 0)), - output: Math.max(0, after.output - (before?.output ?? 0)), - inputCacheRead: Math.max(0, after.inputCacheRead - (before?.inputCacheRead ?? 0)), - inputCacheCreation: Math.max( - 0, - after.inputCacheCreation - (before?.inputCacheCreation ?? 0), - ), - }; -} - function estimateToolResultTokens(output: string | readonly ContentPart[]): number { return typeof output === 'string' ? Math.ceil(Buffer.byteLength(output, 'utf8') / 4) diff --git a/packages/agent-core-v2/src/session/sessionMetadata/sessionMetadata.ts b/packages/agent-core-v2/src/session/sessionMetadata/sessionMetadata.ts index d7b11c489..8ec0c041b 100644 --- a/packages/agent-core-v2/src/session/sessionMetadata/sessionMetadata.ts +++ b/packages/agent-core-v2/src/session/sessionMetadata/sessionMetadata.ts @@ -51,6 +51,7 @@ export interface ISessionMetadata { ): Promise; setArchived(archived: boolean): Promise; registerAgent(agentId: string, meta: AgentMeta): Promise; + unregisterAgent(agentId: string): Promise; } export const ISessionMetadata: ServiceIdentifier = diff --git a/packages/agent-core-v2/src/session/sessionMetadata/sessionMetadataService.ts b/packages/agent-core-v2/src/session/sessionMetadata/sessionMetadataService.ts index b5115b279..221f973e8 100644 --- a/packages/agent-core-v2/src/session/sessionMetadata/sessionMetadataService.ts +++ b/packages/agent-core-v2/src/session/sessionMetadata/sessionMetadataService.ts @@ -95,9 +95,10 @@ export class SessionMetadata extends Service implements ISessionMetadata { if (this.disposed) return false; const updatedAt = patch.updatedAt ?? (opts?.touchUpdatedAt === false ? this.data.updatedAt : Date.now()); - this.data = { ...this.data, ...patch, updatedAt }; - await this.store.set(this.scope, META_KEY, encodeSessionMeta(this.data)); + const next = { ...this.data, ...patch, updatedAt }; + await this.store.set(this.scope, META_KEY, encodeSessionMeta(next)); if (this.disposed) return false; + this.data = next; this.mirrorToReadModel(); this._onDidChangeMetadata.fire({ changed: Object.keys(patch) as (keyof SessionMeta)[], @@ -137,6 +138,15 @@ export class SessionMetadata extends Service implements ISessionMetadata { }); } + async unregisterAgent(agentId: string): Promise { + return this.enqueueUpdate(async () => { + await this.ready; + if (this.data.agents?.[agentId] === undefined) return; + const { [agentId]: _removed, ...agents } = this.data.agents; + await this.applyUpdate({ agents }, { touchUpdatedAt: false }); + }); + } + private enqueueUpdate(work: () => Promise): Promise { const run = this.updateQueue.then(work, work); const tracked: Promise = run.then( diff --git a/packages/agent-core-v2/src/session/subagent/runAgentTurn.ts b/packages/agent-core-v2/src/session/subagent/runAgentTurn.ts index e492ec20c..fc01acfff 100644 --- a/packages/agent-core-v2/src/session/subagent/runAgentTurn.ts +++ b/packages/agent-core-v2/src/session/subagent/runAgentTurn.ts @@ -1,5 +1,5 @@ import { APIProviderRateLimitError, isProviderRateLimitError } from '#/kosong/contract/errors'; -import { type TokenUsage } from '#/kosong/contract/usage'; +import { type TokenUsage, usageDelta } from '#/kosong/contract/usage'; import { linkAbortSignal, userCancellationReason } from '#/_base/utils/abort'; import type { IAgentScopeHandle } from '#/_base/di/scope'; @@ -36,6 +36,7 @@ export async function runAgentTurn( options: RunAgentTurnOptions, ): Promise { options.signal.throwIfAborted(); + const usageBefore = target.accessor.get(ISessionUsageService)?.status(agentContextOf(target)).total; const promptService = target.accessor.get(IAgentPromptService); const turn = request.kind === 'prompt' @@ -58,7 +59,7 @@ export async function runAgentTurn( void turn.ready.then(() => options.onReady?.()).catch(() => {}); } - const completion = awaitRun(target, turn, options); + const completion = awaitRun(target, turn, options, usageBefore); return { agentId: target.id, turn, completion }; } @@ -66,7 +67,8 @@ async function awaitRun( target: IAgentScopeHandle, turn: Turn, options: RunAgentTurnOptions, -): Promise<{ summary: string; usage?: TokenUsage }> { + usageBefore: TokenUsage | undefined, +): Promise<{ summary: string; usage?: TokenUsage; cumulativeUsage?: TokenUsage }> { const controller = new AbortController(); const unlink = linkAbortSignal(options.signal, controller); const loop = target.accessor.get(IAgentLoopService); @@ -86,8 +88,10 @@ async function awaitRun( }, cancelTurn, ); - const usage = target.accessor.get(ISessionUsageService)?.status(agentContextOf(target)).total; - return { summary, usage }; + const cumulativeUsage = target.accessor + .get(ISessionUsageService) + ?.status(agentContextOf(target)).total; + return { summary, usage: usageDelta(cumulativeUsage, usageBefore), cumulativeUsage }; } finally { unlink(); if (controller.signal.aborted) { diff --git a/packages/agent-core-v2/src/session/subagent/subagent.ts b/packages/agent-core-v2/src/session/subagent/subagent.ts index 35de4cbec..40e30eea3 100644 --- a/packages/agent-core-v2/src/session/subagent/subagent.ts +++ b/packages/agent-core-v2/src/session/subagent/subagent.ts @@ -27,7 +27,11 @@ export interface RunAgentOptions { export interface AgentRunHandle { readonly agentId: string; readonly turn: Turn; - readonly completion: Promise<{ readonly summary: string; readonly usage?: TokenUsage }>; + readonly completion: Promise<{ + readonly summary: string; + readonly usage?: TokenUsage; + readonly cumulativeUsage?: TokenUsage; + }>; } export class SubagentRunStartError extends Error { diff --git a/packages/agent-core-v2/src/tool/path-access.ts b/packages/agent-core-v2/src/tool/path-access.ts index 7e11c4631..68e6672d7 100644 --- a/packages/agent-core-v2/src/tool/path-access.ts +++ b/packages/agent-core-v2/src/tool/path-access.ts @@ -243,6 +243,33 @@ function relativeOutsideMessage(path: string, operation: PathAccessOperation): s ); } +export interface ResolvedTargetFs { + realpath(path: string): Promise; +} + +export async function resolveRealTarget(fs: ResolvedTargetFs, path: string): Promise { + try { + return await fs.realpath(path); + } catch { + const parent = pathe.dirname(path); + if (parent === path) return path; + return pathe.join(await resolveRealTarget(fs, parent), pathe.basename(path)); + } +} + +export async function sensitiveTargetError( + fs: ResolvedTargetFs, + requestedPath: string, + safePath: string, +): Promise { + const target = await resolveRealTarget(fs, safePath); + if (target === safePath || !isSensitiveFile(target)) return undefined; + return ( + `"${requestedPath}" resolves to "${target}", which matches a sensitive-file pattern ` + + `(env / credential / SSH key). Access is blocked to protect secrets.` + ); +} + export function resolvePathAccess( path: string, cwd: string, diff --git a/packages/agent-core-v2/test/agent/prompt/promptService.test.ts b/packages/agent-core-v2/test/agent/prompt/promptService.test.ts index c928e3d40..587f891ff 100644 --- a/packages/agent-core-v2/test/agent/prompt/promptService.test.ts +++ b/packages/agent-core-v2/test/agent/prompt/promptService.test.ts @@ -4,12 +4,12 @@ import { Readable } from 'node:stream'; import { DisposableStore } from '#/_base/di/lifecycle'; import { createServices } from '#/_base/di/test'; -import { Event } from '#/_base/event'; +import { Emitter } from '#/_base/event'; import { IAgentBlobService } from '#/agent/blob/agentBlobService'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; import type { ContextMessage } from '#/agent/contextMemory/types'; import type { ContentPart } from '#/kosong/contract/message'; -import { IAgentFullCompactionService } from '#/agent/fullCompaction/fullCompaction'; +import { IAgentFullCompactionService, type FullCompactionTask } from '#/agent/fullCompaction/fullCompaction'; import { IAgentLoopService } from '#/agent/loop/loop'; import { TurnSteer } from '#/agent/loop/turnOps'; import { IAgentPromptService } from '#/agent/prompt/prompt'; @@ -74,13 +74,19 @@ function harness(loopOptions: StubLoopOptions = { pendingTurnResult: true }) { }, }); const loop = stubLoopWithHooks(loopOptions); - const fullCompaction = { + let activeTask: FullCompactionTask | null = null; + const finishCompactionEmitter = new Emitter(); + disposables.add(finishCompactionEmitter); + const fullCompaction: IAgentFullCompactionService = { _serviceBrand: undefined, - compacting: null, + get compacting() { + return activeTask; + }, begin: () => false, + cancel: () => {}, hooks: createHooks(['onWillCompact']), - onDidFinishCompaction: Event.None, - } as unknown as IAgentFullCompactionService; + onDidFinishCompaction: finishCompactionEmitter.event, + }; const intake = { get: vi.fn(async () => ({ meta: { @@ -124,7 +130,21 @@ function harness(loopOptions: StubLoopOptions = { pendingTurnResult: true }) { (ix.get(IEventBus) as ISessionEventBus).activateAgent( ix.get(IAgentScopeContext).agentContext, ); - return { prompt: ix.get(IAgentPromptService), loop, context, fullCompaction, eventBus: ix.get(IEventBus), intake }; + return { + prompt: ix.get(IAgentPromptService), + loop, + context, + fullCompaction, + eventBus: ix.get(IEventBus), + intake, + setCompacting: (task: FullCompactionTask | null) => { + activeTask = task; + }, + finishCompaction: (task: FullCompactionTask) => { + activeTask = null; + finishCompactionEmitter.fire(task); + }, + }; } describe('AgentPromptService', () => { @@ -230,6 +250,100 @@ describe('AgentPromptService', () => { expect(events[0]).not.toHaveProperty('promptIds'); }); + it('abort during launch cancels the prompt before it reaches the loop', async () => { + const { prompt, loop, eventBus } = harness(); + const aborted: PromptAborted[] = []; + eventBus.subscribe(PromptAborted, (event) => aborted.push(event)); + const enqueueSpy = vi.spyOn(loop, 'enqueue'); + let release!: () => void; + const gate = new Promise((resolve) => { release = resolve; }); + let promptId: string | undefined; + prompt.hooks.onBeforeSubmitPrompt.register('gate', async (_ctx, next) => { promptId = prompt.list().active?.id; await gate; await next(); }); + const handlePromise = prompt.enqueue({ id: 'launching', message: message('slow start') }); + await vi.waitFor(() => { expect(prompt.list().pending).toEqual([]); }); + expect(promptId).toBeUndefined(); + expect(prompt.abort('launching')).toBe(true); + release(); + const handle = await handlePromise; + await expect(handle.completion).resolves.toMatchObject({ state: 'cancelled' }); + await expect(handle.launched).resolves.toBeUndefined(); + expect(enqueueSpy).not.toHaveBeenCalled(); + expect(prompt.list()).toEqual({ active: undefined, pending: [] }); + expect(aborted.map((event) => event.promptId)).toEqual(['launching']); + }); + + it('drain waits for a launching prompt and cancels it', async () => { + const { prompt, loop } = harness(); + const enqueueSpy = vi.spyOn(loop, 'enqueue'); + let release!: () => void; + const gate = new Promise((resolve) => { release = resolve; }); + prompt.hooks.onBeforeSubmitPrompt.register('gate', async (_ctx, next) => { await gate; await next(); }); + const handlePromise = prompt.enqueue({ id: 'launching', message: message('slow start') }); + await vi.waitFor(() => { expect(prompt.list().pending).toEqual([]); }); + let drained = false; + const drain = prompt.drain().then(() => { drained = true; }); + await Promise.resolve(); + expect(drained).toBe(false); + release(); + await drain; + const handle = await handlePromise; + await expect(handle.completion).resolves.toMatchObject({ state: 'cancelled' }); + expect(enqueueSpy).not.toHaveBeenCalled(); + }); + + it('clear cancels a launching prompt', async () => { + const { prompt, loop } = harness(); + const enqueueSpy = vi.spyOn(loop, 'enqueue'); + let release!: () => void; + const gate = new Promise((resolve) => { release = resolve; }); + prompt.hooks.onBeforeSubmitPrompt.register('gate', async (_ctx, next) => { await gate; await next(); }); + const handlePromise = prompt.enqueue({ id: 'launching', message: message('slow start') }); + await vi.waitFor(() => { expect(prompt.list().pending).toEqual([]); }); + prompt.clear(); + release(); + await expect((await handlePromise).completion).resolves.toMatchObject({ state: 'cancelled' }); + expect(enqueueSpy).not.toHaveBeenCalled(); + }); + + it('cancels the turn when abort lands after the loop assigned it', async () => { + const { prompt, loop } = harness(); + const cancelSpy = vi.spyOn(loop, 'cancel'); + const originalEnqueue = loop.enqueue.bind(loop); + let releaseAssignment!: () => void; + const assignmentGate = new Promise((resolve) => { releaseAssignment = resolve; }); + const enqueueSpy = vi.spyOn(loop, 'enqueue').mockImplementationOnce((request, options) => { + const receipt = originalEnqueue(request, options); + return { ...receipt, assigned: assignmentGate.then(() => receipt.assigned) }; + }); + const handlePromise = prompt.enqueue({ id: 'launching', message: message('slow start') }); + await vi.waitFor(() => { expect(enqueueSpy).toHaveBeenCalledOnce(); }); + expect(prompt.abort('launching')).toBe(true); + releaseAssignment(); + const handle = await handlePromise; + await expect(handle.launched).resolves.toBeDefined(); + expect(cancelSpy).toHaveBeenCalledOnce(); + }); + + it('parks a queued prompt while compaction runs instead of recursing', async () => { + const { prompt, setCompacting, finishCompaction } = harness(); + const task: FullCompactionTask = { + promise: new Promise(() => {}), + abortController: new AbortController(), + trigger: 'manual', + tokenCount: 100, + }; + setCompacting(task); + const handle = await prompt.enqueue({ id: 'parked', message: message('later') }); + expect(handle.state).toBe('pending'); + expect(prompt.list().pending.map((item) => item.id)).toEqual(['parked']); + expect(prompt.list().active).toBeUndefined(); + + finishCompaction(task); + await expect(handle.launched).resolves.toBeDefined(); + expect(prompt.list().pending).toEqual([]); + expect(prompt.list().active?.id).toBe('parked'); + }); + it('aborts pending prompts and settles completion', async () => { const { prompt, eventBus } = harness(); const aborted: PromptAborted[] = []; diff --git a/packages/agent-core-v2/test/agent/toolExecutor/toolExecutor.test.ts b/packages/agent-core-v2/test/agent/toolExecutor/toolExecutor.test.ts index 3d25f0cc0..d65ef62aa 100644 --- a/packages/agent-core-v2/test/agent/toolExecutor/toolExecutor.test.ts +++ b/packages/agent-core-v2/test/agent/toolExecutor/toolExecutor.test.ts @@ -647,6 +647,94 @@ describe('AgentToolExecutorService', () => { ]); }); + it('classifies telemetry by execution state, not by output text', async () => { + const tool = new TestTool('noisy', { + result: { output: 'upstream said: request aborted by peer', isError: true }, + }); + registry.register(tool); + + await execute([toolCall('call_noisy', 'noisy', {})]); + + expect(telemetryEvents).toContainEqual({ + event: 'tool_call', + properties: expect.objectContaining({ + tool_call_id: 'call_noisy', + outcome: 'error', + error_type: 'error', + }), + }); + }); + + it('reports cancelled telemetry for a tool aborted by the signal', async () => { + const controller = new AbortController(); + const tool = new ControlledTool('slow', ToolAccesses.writeFile('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/repo/a.ts')); + registry.register(tool); + + const execution = execute([toolCall('call_slow', 'slow', {})], controller.signal); + await tool.started; + controller.abort(); + await execution; + + expect(telemetryEvents).toContainEqual({ + event: 'tool_call', + properties: expect.objectContaining({ + tool_call_id: 'call_slow', + outcome: 'cancelled', + error_type: 'cancelled', + }), + }); + }); + + it('holds the resource lease past the abort grace until the ignored execution settles', async () => { + vi.useFakeTimers(); + try { + const controller = new AbortController(); + let finishStubborn: () => void = () => {}; + const stubbornDone = new Promise((resolve) => { + finishStubborn = resolve; + }); + const stubborn = new TestTool('stubborn', { + accesses: ToolAccesses.writeFile('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/repo/a.ts'), + execute: async () => { + await stubbornDone; + return { output: 'late write' }; + }, + }); + const follower = new TestTool('follower', { + accesses: ToolAccesses.writeFile('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/repo/a.ts'), + result: { output: 'follower ran' }, + }); + registry.register(stubborn); + registry.register(follower); + + const firstBatch = execute([toolCall('call_stubborn', 'stubborn', {})], controller.signal); + await vi.waitFor(() => expect(stubborn.calls).toHaveLength(1)); + controller.abort(); + await vi.advanceTimersByTimeAsync(2_000); + const results = await firstBatch; + expect(results).toEqual([ + expect.objectContaining({ output: 'Tool "stubborn" was aborted', isError: true }), + ]); + + let secondSettled = false; + const secondBatch = execute([toolCall('call_follower', 'follower', {})]).then((value) => { + secondSettled = true; + return value; + }); + await vi.advanceTimersByTimeAsync(10); + expect(follower.calls).toHaveLength(0); + expect(secondSettled).toBe(false); + + finishStubborn(); + await vi.advanceTimersByTimeAsync(10); + const secondResults = await secondBatch; + expect(follower.calls).toHaveLength(1); + expect(secondResults).toEqual([expect.objectContaining({ output: 'follower ran' })]); + } finally { + vi.useRealTimers(); + } + }); + it('does not start a queued conflicting tool after abort', async () => { const controller = new AbortController(); const first = new ControlledTool('first', ToolAccesses.writeFile('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/repo/a.ts')); diff --git a/packages/agent-core-v2/test/agent/toolExecutor/toolScheduler.test.ts b/packages/agent-core-v2/test/agent/toolExecutor/toolScheduler.test.ts index 1abb2d22c..9bf5767eb 100644 --- a/packages/agent-core-v2/test/agent/toolExecutor/toolScheduler.test.ts +++ b/packages/agent-core-v2/test/agent/toolExecutor/toolScheduler.test.ts @@ -239,6 +239,78 @@ interface ControlledTask { readonly reject: (error: unknown) => void; } +describe('ToolScheduler leases and budgets', () => { + it('keeps a conflicting task queued until the abandoned effect settles', async () => { + const started: string[] = []; + const drained: string[] = []; + const scheduler = makeScheduler(drained); + let settleEffects: () => void = () => {}; + const effectsSettled = new Promise((resolve) => { + settleEffects = resolve; + }); + const abandoned = makeControlledTask('abandoned', writePath('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/repo/a.ts'), started, effectsSettled); + const follower = makeControlledTask('follower', writePath('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/repo/a.ts'), started); + + scheduler.add(abandoned.task); + scheduler.add(follower.task); + abandoned.resolve(); + await waitOneMacrotask(); + + expect(started).toEqual(['abandoned']); + settleEffects(); + await waitOneMacrotask(); + expect(started).toEqual(['abandoned', 'follower']); + follower.resolve(); + await scheduler.collectResults(); + expect(drained).toEqual(['abandoned', 'follower']); + }); + + it('blocks on outstanding effects from a previous scheduler', async () => { + const started: string[] = []; + let settleEffects: () => void = () => {}; + const settled = new Promise((resolve) => { + settleEffects = resolve; + }); + const scheduler = new ToolScheduler([{ accesses: writePath('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/repo/a.ts'), settled }]); + const conflicting = makeControlledTask('conflicting', readPath('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/repo/a.ts'), started); + const unrelated = makeControlledTask('unrelated', readPath('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/repo/b.ts'), started); + const results = [scheduler.add(conflicting.task), scheduler.add(unrelated.task)]; + await waitOneMacrotask(); + + expect(started).toEqual(['unrelated']); + settleEffects(); + await waitOneMacrotask(); + expect(started).toEqual(['unrelated', 'conflicting']); + conflicting.resolve(); + unrelated.resolve(); + await Promise.all(results); + }); + + it('caps unrelated concurrency at the configured budget', async () => { + const started: string[] = []; + const scheduler = new ToolScheduler([], 2); + const tasks = ['a', 'b', 'c', 'd'].map((name) => + makeControlledTask(name, readPath(`/repo/${name}.ts`), started), + ); + const results = tasks.map((task) => scheduler.add(task.task)); + await waitOneMacrotask(); + + expect(started).toEqual(['a', 'b']); + expect(scheduler.running).toBe(2); + tasks[0]!.resolve(); + await waitOneMacrotask(); + expect(started).toEqual(['a', 'b', 'c']); + tasks[1]!.resolve(); + tasks[2]!.resolve(); + await waitOneMacrotask(); + expect(started).toEqual(['a', 'b', 'c', 'd']); + tasks[3]!.resolve(); + await Promise.all(results); + await waitOneMacrotask(); + expect(scheduler.running).toBe(0); + }); +}); + function makeScheduler(drained: string[]): { readonly add: (task: ToolCallTask) => void; readonly collectResults: () => Promise; @@ -265,6 +337,7 @@ function makeControlledTask( name: string, accesses: ToolAccesses, startedNames: string[], + effectsSettled?: Promise, ): ControlledTask { let resolveResult: (value: string) => void = () => {}; let rejectResult: (error: unknown) => void = () => {}; @@ -278,7 +351,7 @@ function makeControlledTask( accesses, start: async () => { startedNames.push(name); - return { result }; + return { result, effectsSettled }; }, }, resolve: () => { diff --git a/packages/agent-core-v2/test/app/sessionExport/sessionExport.test.ts b/packages/agent-core-v2/test/app/sessionExport/sessionExport.test.ts index d217cb0dc..f7a151fac 100644 --- a/packages/agent-core-v2/test/app/sessionExport/sessionExport.test.ts +++ b/packages/agent-core-v2/test/app/sessionExport/sessionExport.test.ts @@ -985,6 +985,7 @@ function stubSessionMetadata(meta: SessionMeta): ISessionMetadata { setGeneratedTitleIfUncustomized: async () => false, setArchived: async () => {}, registerAgent: async () => {}, + unregisterAgent: async () => {}, }; } diff --git a/packages/agent-core-v2/test/app/web/providers/local-fetch-url.test.ts b/packages/agent-core-v2/test/app/web/providers/local-fetch-url.test.ts index 88d07a10f..ceb43ef27 100644 --- a/packages/agent-core-v2/test/app/web/providers/local-fetch-url.test.ts +++ b/packages/agent-core-v2/test/app/web/providers/local-fetch-url.test.ts @@ -313,6 +313,46 @@ describe('LocalFetchURLProvider connection pinning', () => { expect((fetchImpl.mock.calls[0]![1] as RequestInit).dispatcher).toBeInstanceOf(Agent); }); + it('stops reading an oversized chunked body before buffering it', async () => { + let pulls = 0; + let cancelled = false; + const chunk = new Uint8Array(1024 * 1024); + const body = new ReadableStream({ + pull(controller) { + pulls += 1; + controller.enqueue(chunk); + }, + cancel() { + cancelled = true; + }, + }); + const fetchImpl = vi.fn().mockResolvedValue( + new Response(body, { status: 200, headers: { 'content-type': 'text/plain' } }), + ); + const provider = new LocalFetchURLProvider({ fetchImpl, maxBytes: 4 * 1024 * 1024 }); + + await expect(provider.fetch('https://example.com/stream')).rejects.toThrow( + 'Response body too large', + ); + + expect(cancelled).toBe(true); + expect(pulls).toBeLessThanOrEqual(8); + }); + + it('rejects a body larger than an understated content-length', async () => { + const fetchImpl = vi.fn().mockResolvedValue( + new Response('x'.repeat(2048), { + status: 200, + headers: { 'content-type': 'text/plain', 'content-length': '10' }, + }), + ); + const provider = new LocalFetchURLProvider({ fetchImpl, maxBytes: 1024 }); + + await expect(provider.fetch('https://example.com/lying')).rejects.toThrow( + 'Response body too large', + ); + }); + it('rejects oversized responses by content-length and still closes the pinned Agent', async () => { const fetchImpl = vi.fn().mockResolvedValue( new Response('short', { diff --git a/packages/agent-core-v2/test/features/dynamic_workflow/sessionDynamicWorkflow.test.ts b/packages/agent-core-v2/test/features/dynamic_workflow/sessionDynamicWorkflow.test.ts index 624ead8ff..f72792ee2 100644 --- a/packages/agent-core-v2/test/features/dynamic_workflow/sessionDynamicWorkflow.test.ts +++ b/packages/agent-core-v2/test/features/dynamic_workflow/sessionDynamicWorkflow.test.ts @@ -959,6 +959,9 @@ describe('SessionDynamicWorkflowService metadata compatibility', () => { registerAgent: async (agentId, meta) => { agents[agentId] = meta; }, + unregisterAgent: async (agentId) => { + delete agents[agentId]; + }, }); ix.stub(ISubagentRoutingService, { _serviceBrand: undefined, diff --git a/packages/agent-core-v2/test/features/externalHooks/integration.test.ts b/packages/agent-core-v2/test/features/externalHooks/integration.test.ts index e18764055..2fb8131a7 100644 --- a/packages/agent-core-v2/test/features/externalHooks/integration.test.ts +++ b/packages/agent-core-v2/test/features/externalHooks/integration.test.ts @@ -184,6 +184,7 @@ function stubSessionMetadata(title?: string): ISessionMetadata { setTitle: async () => {}, setArchived: async () => {}, registerAgent: async () => {}, + unregisterAgent: async () => {}, } as unknown as ISessionMetadata; } diff --git a/packages/agent-core-v2/test/kosong/provider/dsml-tool-parser-conformance.test.ts b/packages/agent-core-v2/test/kosong/provider/dsml-tool-parser-conformance.test.ts new file mode 100644 index 000000000..265d395bd --- /dev/null +++ b/packages/agent-core-v2/test/kosong/provider/dsml-tool-parser-conformance.test.ts @@ -0,0 +1,333 @@ +import { readFileSync } from 'node:fs'; +import { resolve } from 'node:path'; + +import { describe, expect, it } from 'vitest'; + +import type { StreamedMessagePart } from '#/kosong/contract/message'; +import { + DSML_MAX_ENVELOPE_CHARS, + DSML_MAX_TAG_CHARS, + DsmlStreamParser, + extractDsmlToolCalls, + parseInvokeBody, +} from '#/kosong/provider/bases/openai/dsml-tool-parser'; + +interface Normalized { + text: string; + calls: Array<{ name: string; arguments: string }>; +} + +function normalize(parts: StreamedMessagePart[]): Normalized { + const calls: Normalized['calls'] = []; + let text = ''; + for (const part of parts) { + if (part.type === 'function') calls.push({ name: part.name, arguments: part.arguments ?? '' }); + else if (part.type === 'text') text += part.text; + } + return { text, calls }; +} + +function runChunks(chunks: string[]): Normalized { + const parser = new DsmlStreamParser(); + const parts: StreamedMessagePart[] = []; + for (const chunk of chunks) parts.push(...parser.feed(chunk)); + parts.push(...parser.flush()); + return normalize(parts); +} + +function seeded(seed: number): () => number { + let state = seed >>> 0; + return () => { + state = (state * 1664525 + 1013904223) >>> 0; + return state / 0x100000000; + }; +} + +function randomPartition(input: string, random: () => number): string[] { + const chunks: string[] = []; + let pos = 0; + while (pos < input.length) { + const size = 1 + Math.floor(random() * 12); + chunks.push(input.slice(pos, pos + size)); + pos += size; + } + return chunks; +} + +const DSML_CALL = + '<｜DSML｜tool_calls>\n<｜DSML｜invoke name="Read">\n<｜DSML｜parameter name="filePath" string="true">src/app.ts\n\n'; + +const CORPUS: Record = { + plainText: 'Check if 5 < 10 and 20 > 15, or use
Hello
and vector.', + singleCall: `Looking into the code...\n\n${DSML_CALL}\nDone reading.`, + asciiBars: + '<|DSML|tool_calls>\n<|DSML|invoke name="Glob">\n<|DSML|parameter name="pattern" string="true">**/*.ts\n\n', + spacedMarkers: + '< | DSML | invoke name="Read">\n< | DSML | parameter name="filePath">src/app.ts\n', + twoInvokes: + '<｜DSML｜tool_calls>\n<｜DSML｜invoke name="Search">\n<｜DSML｜parameter name="query" string="true">export function\n<｜DSML｜parameter name="limit" string="false">25\n\n<｜DSML｜invoke name="Read">\n<｜DSML｜parameter name="path">src/main.ts\n\n', + bareInvoke: + 'Checking directory:\n<｜DSML｜invoke name="ListDir">\n<｜DSML｜parameter name="dir" string="true">packages\n', + jsonBody: '<｜DSML｜invoke name="Eval">{"code":"a < b && c > d"}', + hermes: 'Sure.\n\n{"name": "Read", "arguments": {"filePath": "package.json"}}\n\n', + hermesUnclosed: '{"name": "Read"}', + malformedInvoke: 'before <｜DSML｜invoke>malformed content without name after', + malformedBody: 'x <｜DSML｜invoke name="noop">definitely not JSON or parameters y', + containerNoCalls: 'before not a call after', + containerOnlyText: '<｜DSML｜tool_calls>\nplain prose\n', + unclosedContainer: 'a <｜DSML｜tool_calls>\nb', + fencedDoc: + 'Example only. Do not execute:\n```xml\n<｜DSML｜invoke name="noop">{"value":"documentation"}\n```\nEnd.', + fencedThenReal: `Use this syntax:\n\`\`\`\n${DSML_CALL}\n\`\`\`\nNow for real:\n${DSML_CALL}`, + tildeFence: '~~~\n{"name":"Read","arguments":{}}\n~~~\n', + unmarkedInvokeOutsideContainer: + 'Do not execute:\n{"value":"documentation"}\n', + unmarkedInvokeInsideContainer: + '\n\na.ts\n\n', + entities: + '<｜DSML｜invoke name="Eval">\n<｜DSML｜parameter name="code" string="true">a && b < c\n', + codeComparison: 'if (x < 5 && y > 2) { return a ', + multibyte: `café ✓ ${DSML_CALL} \u65E5\u672C\u8A9E 😀 done`, + strayClose: 'no container here still text', + emptyInvoke: '<｜DSML｜invoke name="Noop">', + whitespaceInvoke: '<｜DSML｜invoke name="Noop">\n \n', +}; + +describe('DsmlStreamParser conformance', () => { + describe('chunk invariance', () => { + for (const [name, input] of Object.entries(CORPUS)) { + it(`${name}: every two-way split equals the unsplit result`, () => { + const whole = runChunks([input]); + for (let offset = 1; offset < input.length; offset += 1) { + const split = runChunks([input.slice(0, offset), input.slice(offset)]); + expect(split, `split at ${offset}`).toEqual(whole); + } + }); + + it(`${name}: character-at-a-time equals the unsplit result`, () => { + const whole = runChunks([input]); + expect(runChunks(Array.from(input))).toEqual(whole); + }); + + it(`${name}: seeded random partitions equal the unsplit result`, () => { + const whole = runChunks([input]); + const random = seeded(name.length * 7919 + input.length); + for (let round = 0; round < 25; round += 1) { + const chunks = randomPartition(input, random); + expect(runChunks(chunks), chunks.join('|')).toEqual(whole); + } + }); + } + + it('handles a UTF-16 surrogate pair split between chunks', () => { + const input = `😀 ${DSML_CALL} 😀`; + const whole = runChunks([input]); + expect(runChunks([input.slice(0, 1), input.slice(1)])).toEqual(whole); + expect(runChunks([input.slice(0, input.length - 1), input.slice(-1)])).toEqual(whole); + }); + }); + + describe('streaming and non-streaming parity', () => { + for (const [name, input] of Object.entries(CORPUS)) { + it(`${name}: streamed text matches extractDsmlToolCalls`, () => { + const streamed = runChunks(Array.from(input)); + const extracted = extractDsmlToolCalls(input); + expect(extracted.toolCalls.map((c) => ({ name: c.name, arguments: c.arguments }))).toEqual( + streamed.calls, + ); + if (streamed.calls.length === 0) { + expect(extracted.cleanText).toBe(input); + expect(streamed.text).toBe(input); + } else { + expect(extracted.cleanText).toBe(streamed.text.trim()); + } + }); + } + }); + + describe('literal documentation stays text', () => { + it('does not promote a fenced DSML example to a tool call', () => { + const result = extractDsmlToolCalls(CORPUS['fencedDoc'] as string); + expect(result.toolCalls).toHaveLength(0); + expect(result.cleanText).toBe(CORPUS['fencedDoc']); + }); + + it('does not promote a fenced Hermes example to a tool call', () => { + const result = extractDsmlToolCalls(CORPUS['tildeFence'] as string); + expect(result.toolCalls).toHaveLength(0); + expect(result.cleanText).toBe(CORPUS['tildeFence']); + }); + + it('does not promote an unmarked invoke outside a container', () => { + const result = extractDsmlToolCalls(CORPUS['unmarkedInvokeOutsideContainer'] as string); + expect(result.toolCalls).toHaveLength(0); + expect(result.cleanText).toBe(CORPUS['unmarkedInvokeOutsideContainer']); + }); + + it('accepts an unmarked invoke inside a container', () => { + const result = extractDsmlToolCalls(CORPUS['unmarkedInvokeInsideContainer'] as string); + expect(result.toolCalls).toHaveLength(1); + expect(result.toolCalls[0]?.name).toBe('Read'); + expect(result.cleanText).toBe(''); + }); + + it('parses the real call that follows a fenced example', () => { + const result = extractDsmlToolCalls(CORPUS['fencedThenReal'] as string); + expect(result.toolCalls).toHaveLength(1); + expect(result.cleanText).toBe(`Use this syntax:\n\`\`\`\n${DSML_CALL}\n\`\`\`\nNow for real:`); + }); + + it('resumes recognition after the fence closes', () => { + const input = '```\n{"name":"A"}\n```\n{"name":"B","arguments":{}}'; + const result = extractDsmlToolCalls(input); + expect(result.toolCalls.map((c) => c.name)).toEqual(['B']); + }); + }); + + describe('invoke body validation', () => { + it('rejects malformed nonempty bodies instead of producing empty arguments', () => { + const result = extractDsmlToolCalls(CORPUS['malformedBody'] as string); + expect(result.toolCalls).toHaveLength(0); + expect(result.cleanText).toBe(CORPUS['malformedBody']); + }); + + it('accepts an empty body as empty arguments', () => { + expect(extractDsmlToolCalls(CORPUS['emptyInvoke'] as string).toolCalls[0]?.arguments).toBe('{}'); + expect(extractDsmlToolCalls(CORPUS['whitespaceInvoke'] as string).toolCalls[0]?.arguments).toBe( + '{}', + ); + }); + + it('rejects prose mixed between parameter blocks', () => { + const input = + '<｜DSML｜invoke name="Read">\n<｜DSML｜parameter name="a">1\nstray prose\n'; + expect(extractDsmlToolCalls(input).toolCalls).toHaveLength(0); + }); + + it('rejects unnamed parameters, unclosed parameters, arrays and scalars', () => { + expect(parseInvokeBody('1')).toBeNull(); + expect(parseInvokeBody('1')).toBeNull(); + expect(parseInvokeBody('[1,2]')).toBeNull(); + expect(parseInvokeBody('42')).toBeNull(); + expect(parseInvokeBody('{"a":')).toBeNull(); + }); + + it('lets the last duplicate parameter win', () => { + const body = '12'; + expect(parseInvokeBody(body)).toEqual({ a: 2 }); + }); + + it('preserves reserved keys as own properties without touching the global prototype', () => { + const input = + '<｜DSML｜invoke name="noop"><｜DSML｜parameter name="__proto__" string="false">{"value":"kept"}<｜DSML｜parameter name="constructor">c<｜DSML｜parameter name="prototype">p<｜DSML｜parameter name="normal">ok'; + const result = extractDsmlToolCalls(input); + expect(result.toolCalls).toHaveLength(1); + expect(JSON.parse(result.toolCalls[0]?.arguments ?? '{}')).toEqual( + JSON.parse('{"__proto__":{"value":"kept"},"constructor":"c","prototype":"p","normal":"ok"}'), + ); + expect(result.toolCalls[0]?.arguments).toContain('"__proto__":{"value":"kept"}'); + expect(({} as Record)['value']).toBeUndefined(); + }); + + it('keeps reserved keys from JSON bodies', () => { + const args = parseInvokeBody('{"__proto__":{"x":1},"y":2}'); + expect(JSON.stringify(args)).toBe('{"__proto__":{"x":1},"y":2}'); + }); + }); + + describe('text preservation', () => { + it('keeps container tags when the container yields no call', () => { + const streamed = runChunks([CORPUS['containerNoCalls'] as string]); + expect(streamed).toEqual({ text: CORPUS['containerNoCalls'], calls: [] }); + }); + + it('drops container tags once a call is accepted and keeps later text', () => { + const streamed = runChunks([CORPUS['singleCall'] as string]); + expect(streamed.text).toBe('Looking into the code...\n\nDone reading.'); + expect(streamed.calls).toEqual([{ name: 'Read', arguments: '{"filePath":"src/app.ts"}' }]); + }); + + it('emits held text in order when a container has an invalid block then a valid call', () => { + const input = `<｜DSML｜tool_calls>\n<｜DSML｜invoke>bad\n<｜DSML｜invoke name="Ok">\n`; + const parser = new DsmlStreamParser(); + const parts = [...parser.feed(input), ...parser.flush()]; + expect(parts.map((p) => p.type)).toEqual(['text', 'text', 'function']); + expect(normalize(parts).text).toBe('<｜DSML｜invoke>bad\n'); + }); + + it('preserves an unclosed container at flush', () => { + expect(runChunks([CORPUS['unclosedContainer'] as string])).toEqual({ + text: CORPUS['unclosedContainer'], + calls: [], + }); + }); + }); + + describe('resource budgets', () => { + it('stops holding a tag prefix beyond the tag budget', () => { + const input = ` { + const body = 'x'.repeat(DSML_MAX_ENVELOPE_CHARS + 16); + const input = `<｜DSML｜invoke name="Big">${body}`; + const whole = runChunks([input]); + expect(whole.calls).toHaveLength(0); + expect(whole.text).toBe(input); + const random = seeded(7); + const chunks: string[] = []; + let pos = 0; + while (pos < input.length) { + const size = 1 + Math.floor(random() * 70000); + chunks.push(input.slice(pos, pos + size)); + pos += size; + } + expect(runChunks(chunks)).toEqual(whole); + }); + + it('accepts an envelope exactly at the budget', () => { + const open = '<｜DSML｜invoke name="Big">'; + const close = ''; + const body = 'x'.repeat(DSML_MAX_ENVELOPE_CHARS - open.length - close.length); + const input = `${open}${body}${close}`; + expect(input.length).toBe(DSML_MAX_ENVELOPE_CHARS); + expect(runChunks([input]).calls).toHaveLength(0); + const jsonBody = `{"v":"${'x'.repeat(DSML_MAX_ENVELOPE_CHARS - open.length - close.length - 8)}"}`; + const valid = `${open}${jsonBody}${close}`; + expect(valid.length).toBe(DSML_MAX_ENVELOPE_CHARS); + expect(runChunks([valid]).calls).toHaveLength(1); + expect(runChunks([valid.slice(0, 5000), valid.slice(5000)]).calls).toHaveLength(1); + }); + + it('scans an unclosed envelope in linear time', () => { + const parser = new DsmlStreamParser(); + parser.feed('<｜DSML｜invoke name="Big">'); + const chunk = 'y'.repeat(64); + const rounds = Math.floor((DSML_MAX_ENVELOPE_CHARS - 64) / 64); + const started = performance.now(); + for (let i = 0; i < rounds; i += 1) parser.feed(chunk); + const elapsed = performance.now() - started; + expect(elapsed).toBeLessThan(2000); + expect(normalize(parser.flush()).calls).toHaveLength(0); + }); + }); + + describe('implementation parity with @pymodel/kosong', () => { + it('keeps both parser copies byte-identical except for the message import', () => { + const dir = import.meta.dirname; + const here = resolve(dir, '../../../src/kosong/provider/bases/openai/dsml-tool-parser.ts'); + const legacy = resolve(dir, '../../../../kosong/src/providers/dsml-tool-parser.ts'); + const v2 = readFileSync(here, 'utf8').replace( + "from '#/kosong/contract/message';", + "from '#/message';", + ); + expect(readFileSync(legacy, 'utf8')).toBe(v2); + }); + }); +}); diff --git a/packages/agent-core-v2/test/kosong/provider/dsml-tool-parser.test.ts b/packages/agent-core-v2/test/kosong/provider/dsml-tool-parser.test.ts index fc66fe0b7..41f8977db 100644 --- a/packages/agent-core-v2/test/kosong/provider/dsml-tool-parser.test.ts +++ b/packages/agent-core-v2/test/kosong/provider/dsml-tool-parser.test.ts @@ -295,6 +295,115 @@ describe('agent-core-v2: DsmlStreamParser and extractDsmlToolCalls', () => { }); }); + function mockClientFor(stream: unknown): unknown { + return { + chat: { + completions: { + create: () => ({ + withResponse: async () => ({ data: stream, response: { headers: new Headers() } }), + }), + }, + }, + }; + } + + async function* chunksOf(chunks: unknown[]) { + for (const chunk of chunks) yield chunk; + } + + const DSML_ECHO = + '<｜DSML｜tool_calls>\n<｜DSML｜invoke name="Read">\n<｜DSML｜parameter name="filePath" string="true">src/server.ts\n\n'; + + it('lets native streamed tool calls win over a DSML echo in content', async () => { + const provider = new OpenAILegacyChatProvider({ model: 'deepseek-chat', apiKey: 'k', stream: true }); + const chunks = [ + { id: 'c1', choices: [{ index: 0, delta: { content: DSML_ECHO }, finish_reason: null }] }, + { + id: 'c1', + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { index: 0, id: 'call_native', type: 'function', function: { name: 'Read', arguments: '{"filePath":"src/server.ts"}' } }, + ], + }, + finish_reason: 'tool_calls', + }, + ], + }, + ]; + (provider as unknown as { _client: unknown })._client = mockClientFor(chunksOf(chunks)); + const stream = await provider.generate('', [], []); + const parts: Array> = []; + for await (const p of stream) parts.push(p as unknown as Record); + const calls = parts.filter((p) => p['type'] === 'function'); + expect(calls).toHaveLength(1); + expect(calls[0]).toMatchObject({ id: 'call_native', name: 'Read' }); + expect(stream.finishReason).toBe('tool_calls'); + }); + + it('keeps intentional repeated native streamed calls distinct', async () => { + const provider = new OpenAILegacyChatProvider({ model: 'deepseek-chat', apiKey: 'k', stream: true }); + const native = (id: string, index: number) => ({ + index, + id, + type: 'function', + function: { name: 'Read', arguments: '{"filePath":"a.ts"}' }, + }); + const chunks = [ + { id: 'c1', choices: [{ index: 0, delta: { tool_calls: [native('call_a', 0)] }, finish_reason: null }] }, + { id: 'c1', choices: [{ index: 0, delta: { tool_calls: [native('call_b', 1)] }, finish_reason: 'tool_calls' }] }, + ]; + (provider as unknown as { _client: unknown })._client = mockClientFor(chunksOf(chunks)); + const stream = await provider.generate('', [], []); + const parts: Array> = []; + for await (const p of stream) parts.push(p as unknown as Record); + expect(parts.filter((p) => p['type'] === 'function').map((p) => p['id'])).toEqual(['call_a', 'call_b']); + }); + + it('defers recovered streamed calls until the response ends so a late native call can win', async () => { + const provider = new OpenAILegacyChatProvider({ model: 'deepseek-chat', apiKey: 'k', stream: true }); + const chunks = [ + { id: 'c1', choices: [{ index: 0, delta: { content: 'Reading.\n' + DSML_ECHO }, finish_reason: null }] }, + { id: 'c1', choices: [{ index: 0, delta: { content: '\nDone.' }, finish_reason: 'stop' }] }, + ]; + (provider as unknown as { _client: unknown })._client = mockClientFor(chunksOf(chunks)); + const stream = await provider.generate('', [], []); + const parts: Array> = []; + for await (const p of stream) parts.push(p as unknown as Record); + expect(parts.map((p) => p['type'])).toEqual(['text', 'text', 'function']); + expect(stream.finishReason).toBe('tool_calls'); + }); + + it('ignores a DSML echo in non-stream content when native tool calls are present', async () => { + const provider = new OpenAILegacyChatProvider({ model: 'deepseek-chat', apiKey: 'k', stream: false }); + const responseData = { + id: 'n1', + choices: [ + { + index: 0, + message: { + role: 'assistant', + content: DSML_ECHO, + tool_calls: [ + { id: 'call_native', type: 'function', function: { name: 'Read', arguments: '{"filePath":"src/server.ts"}' } }, + ], + }, + finish_reason: 'tool_calls', + }, + ], + }; + (provider as unknown as { _client: unknown })._client = mockClientFor(responseData); + const stream = await provider.generate('', [], []); + const parts: Array> = []; + for await (const p of stream) parts.push(p as unknown as Record); + const calls = parts.filter((p) => p['type'] === 'function'); + expect(calls).toHaveLength(1); + expect(calls[0]).toMatchObject({ id: 'call_native' }); + expect(parts.filter((p) => p['type'] === 'text').map((p) => p['text'])).toEqual([DSML_ECHO]); + }); + it('preserves surrounding whitespace and markdown hard breaks in non-stream response', async () => { const provider = new OpenAILegacyChatProvider({ model: 'deepseek-chat', diff --git a/packages/agent-core-v2/test/os/backends/node-local/hostFsService.test.ts b/packages/agent-core-v2/test/os/backends/node-local/hostFsService.test.ts index 54f2d8028..62ea014fc 100644 --- a/packages/agent-core-v2/test/os/backends/node-local/hostFsService.test.ts +++ b/packages/agent-core-v2/test/os/backends/node-local/hostFsService.test.ts @@ -1,4 +1,4 @@ -import { mkdir, mkdtemp, rm, symlink, writeFile } from 'node:fs/promises'; +import { chmod, lstat, mkdir, mkdtemp, readFile, readdir, rm, stat, symlink, writeFile } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'pathe'; @@ -18,6 +18,95 @@ afterEach(async () => { await rm(dir, { recursive: true, force: true }); }); +describe('HostFileSystem atomic writes', () => { + it('replaces the symlink target in place and keeps the link', async () => { + const target = join(dir, 'target.txt'); + await writeFile(target, 'old', 'utf-8'); + await chmod(target, 0o600); + const link = join(dir, 'link.txt'); + await symlink(target, link); + + await fs.writeText(link, 'new'); + + expect(await readFile(target, 'utf-8')).toBe('new'); + expect((await lstat(link)).isSymbolicLink()).toBe(true); + expect((await stat(target)).mode & 0o777).toBe(0o600); + expect((await readdir(dir)).toSorted()).toEqual(['link.txt', 'target.txt']); + }); + + it('creates a missing file and leaves no staging files behind', async () => { + const path = join(dir, 'fresh.txt'); + await fs.writeBytes(path, new Uint8Array([104, 105])); + expect(await readFile(path, 'utf-8')).toBe('hi'); + expect(await readdir(dir)).toEqual(['fresh.txt']); + }); + + it('keeps the previous content when the replacement cannot be staged', async () => { + const path = join(dir, 'keep.txt'); + await writeFile(path, 'intact', 'utf-8'); + await chmod(dir, 0o500); + try { + await expect(fs.writeText(path, 'partial')).rejects.toThrow(); + } finally { + await chmod(dir, 0o700); + } + expect(await readFile(path, 'utf-8')).toBe('intact'); + }); +}); + +describe('HostFileSystem readLines budget', () => { + it('bounds a newline-free line to maxLineBytes', async () => { + const path = join(dir, 'long.txt'); + await writeFile(path, `${'x'.repeat(300_000)}\nshort\n`, 'utf-8'); + + const lines: string[] = []; + for await (const line of fs.readLines(path, { maxLineBytes: 1024 })) lines.push(line); + + expect(lines).toHaveLength(2); + expect(lines[0]).toBe(`${'x'.repeat(1024)}\n`); + expect(lines[1]).toBe('short\n'); + }); + + it('bounds a trailing line without a terminator', async () => { + const path = join(dir, 'tail.txt'); + await writeFile(path, `a\n${'y'.repeat(5000)}`, 'utf-8'); + + const lines: string[] = []; + for await (const line of fs.readLines(path, { maxLineBytes: 16 })) lines.push(line); + + expect(lines).toEqual(['a\n', 'y'.repeat(16)]); + }); + + it('reads every byte when no budget is given', async () => { + const path = join(dir, 'full.txt'); + const long = 'z'.repeat(200_000); + await writeFile(path, `${long}\n`, 'utf-8'); + + const lines: string[] = []; + for await (const line of fs.readLines(path)) lines.push(line); + + expect(lines).toEqual([`${long}\n`]); + }); + + it('preserves complete UTF-8 code points when line budget cuts into multibyte characters', async () => { + const path = join(dir, 'multibyte.txt'); + await writeFile(path, `${'€'.repeat(100)}\n`, 'utf-8'); + + const lines: string[] = []; + for await (const line of fs.readLines(path, { maxLineBytes: 16, errors: 'strict' })) lines.push(line); + + expect(lines).toEqual([`${'€'.repeat(5)}\n`]); + }); + + it('rejects maxLineBytes when encoding is not UTF-8', async () => { + const path = join(dir, 'other.txt'); + await writeFile(path, 'content', 'utf-8'); + + const iter = fs.readLines(path, { encoding: 'latin1', maxLineBytes: 4 }); + await expect(iter.next()).rejects.toThrow('maxLineBytes is only supported for UTF-8 encoding'); + }); +}); + describe('HostFileSystem stat / lstat', () => { it('stat follows a symlink to a regular file while lstat stats the link', async () => { const target = join(dir, 'target.txt'); diff --git a/packages/agent-core-v2/test/os/backends/node-local/tools/read.test.ts b/packages/agent-core-v2/test/os/backends/node-local/tools/read.test.ts index 38cd4e8be..39056a2b9 100644 --- a/packages/agent-core-v2/test/os/backends/node-local/tools/read.test.ts +++ b/packages/agent-core-v2/test/os/backends/node-local/tools/read.test.ts @@ -234,6 +234,20 @@ describe('ReadTool', () => { }); }); + it('denies a benign alias that resolves to a sensitive file', async () => { + const { fs, readBytes } = createSpiedFs('SECRET=1\n'); + (fs as { realpath?: (path: string) => Promise }).realpath = vi.fn(async (path: string) => + path === '/tmp/notes.txt' ? '/home/user/.env' : path, + ); + const tool = createReadTool(fs, createTestEnv(), PERMISSIVE_WORKSPACE); + + const result = await execute(tool, { path: '/tmp/notes.txt' }); + + expect(result.isError).toBe(true); + expect(result.output).toContain('resolves to "/home/user/.env"'); + expect(readBytes).not.toHaveBeenCalled(); + }); + it('stats the resolved target so symlinked files stay readable', async () => { const { fs, stat } = createSpiedFs('alpha\n'); const tool = createReadTool(fs, createTestEnv(), PERMISSIVE_WORKSPACE); @@ -359,7 +373,7 @@ describe('ReadTool', () => { ), ); expect(readBytes).toHaveBeenCalledWith('/tmp/external.txt', MEDIA_SNIFF_BYTES); - expect(readLines).toHaveBeenCalledWith('/tmp/external.txt', { errors: 'strict' }); + expect(readLines).toHaveBeenCalledWith('/tmp/external.txt', { errors: 'strict', maxLineBytes: MAX_LINE_LENGTH * 4 }); }); it('returns a friendly error for missing files before sniffing bytes', async () => { @@ -405,7 +419,7 @@ describe('ReadTool', () => { ), ); expect(readBytes).toHaveBeenCalledWith('/home/test/notes/today.txt', MEDIA_SNIFF_BYTES); - expect(readLines).toHaveBeenCalledWith('/home/test/notes/today.txt', { errors: 'strict' }); + expect(readLines).toHaveBeenCalledWith('/home/test/notes/today.txt', { errors: 'strict', maxLineBytes: MAX_LINE_LENGTH * 4 }); }); it('blocks sensitive files independently from workspace access', async () => { diff --git a/packages/agent-core-v2/test/os/backends/node-local/tools/write.test.ts b/packages/agent-core-v2/test/os/backends/node-local/tools/write.test.ts index 57dc3f1cf..a2bed15f5 100644 --- a/packages/agent-core-v2/test/os/backends/node-local/tools/write.test.ts +++ b/packages/agent-core-v2/test/os/backends/node-local/tools/write.test.ts @@ -41,6 +41,7 @@ interface WriteFsOptions { appendText?: (path: string, data: string) => Promise; stat?: (path: string) => Promise; mkdir?: (path: string) => Promise; + readonly realpath?: (path: string) => Promise; } function createWriteFs(options: WriteFsOptions = {}) { @@ -56,8 +57,9 @@ function createWriteFs(options: WriteFsOptions = {}) { options.stat ?? (async () => ({ isFile: false, isDirectory: true, size: 0 })), ); const mkdir = vi.fn(options.mkdir ?? (async () => {})); - const fs = { cwd: '/', readText, writeText, appendText, stat, mkdir } as unknown as IHostFileSystem; - return { fs, readText, writeText, appendText, stat, mkdir }; + const realpath = vi.fn(options.realpath ?? (async (path: string) => path)); + const fs = { cwd: '/', readText, writeText, appendText, stat, mkdir, realpath } as unknown as IHostFileSystem; + return { fs, readText, writeText, appendText, stat, mkdir, realpath }; } function makeTool(options: WriteFsOptions = {}, workspace = PERMISSIVE_WORKSPACE) { @@ -335,6 +337,19 @@ describe('WriteTool', () => { expect(writeText).not.toHaveBeenCalled(); }); + it('denies a benign alias whose real target is a sensitive file', async () => { + const { tool, writeText } = makeTool({ + realpath: async (path) => (path === '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/workspace/notes.txt' ? '/home/user/.env' : path), + stat: async () => ({ isFile: true, isDirectory: false, size: 1 }), + }); + + const result = await execute(tool, { path: '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/workspace/notes.txt', content: 'x' }); + + expect(result.isError).toBe(true); + expect(result.output).toContain('resolves to "/home/user/.env"'); + expect(writeText).not.toHaveBeenCalled(); + }); + it('blocks sensitive file writes', async () => { const { tool, writeText } = makeTool({}, stubWorkspaceContext('/workspace')); diff --git a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts index e5936379a..b017184dd 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -196,6 +196,7 @@ describe('AgentLifecycleService', () => { let disposables: DisposableStore; let ix: TestInstantiationService; let registerAgent: ReturnType>; + let unregisterAgent: ReturnType>; let atomicDocs: Map; let permissionModeSetMode: ReturnType; let stopAllOnExit: ReturnType; @@ -220,6 +221,7 @@ describe('AgentLifecycleService', () => { ix.stub(IFileSystemStorageService, new InMemoryStorageService()); stubBlobPassThrough(ix); registerAgent = vi.fn().mockResolvedValue(undefined); + unregisterAgent = vi.fn().mockResolvedValue(undefined); atomicDocs = new Map(); ix.stub(ISessionContext, { _serviceBrand: undefined, @@ -255,6 +257,7 @@ describe('AgentLifecycleService', () => { setTitle: () => Promise.resolve(), setArchived: () => Promise.resolve(), registerAgent, + unregisterAgent, }); ix.stub(IBootstrapService, { _serviceBrand: undefined, @@ -352,6 +355,7 @@ describe('AgentLifecycleService', () => { }), cancel: loopCancel, settled: loopSettled, + tryAcquireQuiescence: () => ({ dispose: () => {} }), } as unknown as IAgentLoopService); promptDrain = vi.fn(async () => {}); ix.stub(IAgentPromptService, { @@ -685,6 +689,50 @@ describe('AgentLifecycleService', () => { } }); + it('remove runs every cleanup phase, rethrows the failure, and stays retry-safe', async () => { + stopAllOnExit.mockRejectedValueOnce(new Error('stop failed')); + const svc = ix.get(IAgentLifecycleService); + const main = await svc.create({ agentId: 'main' }); + const closed: string[] = []; + disposables.add(svc.onDidClose((agent) => closed.push(agent.agentId))); + + await expect(svc.remove(main)).rejects.toThrow('stop failed'); + + expect(promptDrain).toHaveBeenCalledOnce(); + expect(closed).toEqual(['main']); + expect(svc.get('main')).toBeUndefined(); + await expect(svc.remove(main)).resolves.toBeUndefined(); + const recreated = await svc.create({ agentId: 'main' }); + expect(recreated).not.toBe(main); + }); + + it('concurrent remove calls join a single close', async () => { + const svc = ix.get(IAgentLifecycleService); + const main = await svc.create({ agentId: 'main' }); + const closed: string[] = []; + disposables.add(svc.onDidClose((agent) => closed.push(agent.agentId))); + + await Promise.all([svc.remove(main), svc.remove(main), svc.remove(main)]); + + expect(stopAllOnExit).toHaveBeenCalledTimes(1); + expect(promptDrain).toHaveBeenCalledOnce(); + expect(closed).toEqual(['main']); + }); + + it('create failure after registration compensates the persisted agent metadata', async () => { + const config = ix.get(IConfigService); + vi.spyOn(config, 'get').mockImplementationOnce(() => { + throw new Error('bootstrap boom'); + }); + const svc = ix.get(IAgentLifecycleService); + + await expect(svc.create({ agentId: 'main' })).rejects.toThrow('bootstrap boom'); + + expect(registerAgent).toHaveBeenCalledOnce(); + expect(unregisterAgent).toHaveBeenCalledWith('main'); + expect(svc.get('main')).toBeUndefined(); + }); + it('remove stops the agent background tasks before disposal', async () => { const svc = ix.get(IAgentLifecycleService); const main = await svc.create({ agentId: 'main' }); @@ -802,6 +850,7 @@ describe('AgentLifecycleService', () => { setTitle: () => Promise.resolve(), setArchived: () => Promise.resolve(), registerAgent, + unregisterAgent: () => Promise.resolve(), }); const svc = ix.get(IAgentLifecycleService); diff --git a/packages/agent-core-v2/test/session/sessionMetadata/sessionMetadata.test.ts b/packages/agent-core-v2/test/session/sessionMetadata/sessionMetadata.test.ts index 5bde4d5bc..b94abfe45 100644 --- a/packages/agent-core-v2/test/session/sessionMetadata/sessionMetadata.test.ts +++ b/packages/agent-core-v2/test/session/sessionMetadata/sessionMetadata.test.ts @@ -214,6 +214,39 @@ describe('SessionMetadata', () => { expect(atUpdate).toMatchObject({ title: 'durable-first' }); }); + it('a rejected persist leaves memory at the last committed state', async () => { + const store = ix.get(IAtomicDocumentStore); + const meta = ix.get(ISessionMetadata); + await meta.ready; + await meta.update({ title: 'committed' }); + vi.spyOn(store, 'set').mockRejectedValueOnce(new Error('disk full')); + + await expect(meta.update({ title: 'lost' })).rejects.toThrow('disk full'); + expect((await meta.read()).title).toBe('committed'); + + await meta.setArchived(true); + const fresh = createFreshMetadata(ix); + expect(await fresh.read()).toMatchObject({ title: 'committed', archived: true }); + }); + + it('unregisterAgent removes the agent entry without bumping updatedAt', async () => { + const meta = ix.get(ISessionMetadata); + await meta.ready; + await meta.registerAgent('sub-1', { homedir: '/tmp/sub-1', type: 'sub', parentAgentId: 'main' }); + const before = (await meta.read()).updatedAt; + const nowSpy = vi.spyOn(Date, 'now').mockReturnValue(before + 10_000); + try { + await meta.unregisterAgent('sub-1'); + await meta.unregisterAgent('missing'); + } finally { + nowSpy.mockRestore(); + } + const after = await meta.read(); + expect(after.agents).toEqual({}); + expect(after.updatedAt).toBe(before); + expect((await createFreshMetadata(ix).read()).agents).toEqual({}); + }); + it('a mirror failure degrades the read model but never fails the metadata mutation', async () => { mirror.record = () => { throw new Error('mirror down'); diff --git a/packages/agent-core-v2/test/session/subagent/runAgentTurn.test.ts b/packages/agent-core-v2/test/session/subagent/runAgentTurn.test.ts new file mode 100644 index 000000000..f9c428782 --- /dev/null +++ b/packages/agent-core-v2/test/session/subagent/runAgentTurn.test.ts @@ -0,0 +1,72 @@ +import { describe, expect, it } from 'vitest'; + +import type { IAgentScopeHandle } from '#/_base/di/scope'; +import { LifecycleScope } from '#/app/scopes'; +import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; +import { IAgentPromptService } from '#/agent/prompt/prompt'; +import { IAgentLoopService, type Turn } from '#/agent/loop/loop'; +import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext'; +import { ISessionUsageService } from '#/session/usage/sessionUsage'; +import { runAgentTurn } from '#/session/subagent/runAgentTurn'; +import type { TokenUsage } from '#/kosong/contract/usage'; + +function usage(inputOther: number, output: number): TokenUsage { + return { inputOther, output, inputCacheRead: 0, inputCacheCreation: 0 }; +} + +function fakeHandle(totals: TokenUsage[]): IAgentScopeHandle { + let cursor = 0; + const agentContext = { agentId: 'sub', space: {} }; + const turn: Turn = { + id: 1, + signal: new AbortController().signal, + ready: Promise.resolve(), + result: Promise.resolve({ type: 'completed', steps: 1, truncated: false }), + cancel: () => true, + }; + return { + id: 'sub', + kind: LifecycleScope.Agent, + accessor: { + get: (serviceId: unknown): T => { + if (serviceId === ISessionUsageService) { + return { + status: () => ({ total: totals[Math.min(cursor, totals.length - 1)] }), + } as T; + } + if (serviceId === IAgentPromptService) { + return { + enqueue: async () => { + cursor += 1; + return { launched: Promise.resolve(turn) }; + }, + } as T; + } + if (serviceId === IAgentLoopService) return { cancel: () => true } as T; + if (serviceId === IAgentContextMemoryService) { + return { get: () => [{ role: 'assistant', content: [{ type: 'text', text: 'done' }] }] } as T; + } + if (serviceId === IAgentScopeContext) return { agentContext } as T; + throw new Error(`unexpected service ${String(serviceId)}`); + }, + }, + dispose: () => {}, + }; +} + +describe('runAgentTurn usage attribution', () => { + it('reports per-run usage as a delta and exposes the cumulative total separately', async () => { + const handle = fakeHandle([usage(100, 10), usage(160, 25), usage(200, 30)]); + const options = { signal: new AbortController().signal }; + + const first = await runAgentTurn(handle, { kind: 'prompt', prompt: 'one' }, options); + const firstResult = await first.completion; + expect(firstResult.usage).toEqual(usage(60, 15)); + expect(firstResult.cumulativeUsage).toEqual(usage(160, 25)); + + const second = await runAgentTurn(handle, { kind: 'prompt', prompt: 'two' }, options); + const secondResult = await second.completion; + expect(secondResult.usage).toEqual(usage(40, 5)); + expect(secondResult.cumulativeUsage).toEqual(usage(200, 30)); + }); +}); diff --git a/packages/agent-core-v2/test/tool/tool.test.ts b/packages/agent-core-v2/test/tool/tool.test.ts index d63567108..1030f9f12 100644 --- a/packages/agent-core-v2/test/tool/tool.test.ts +++ b/packages/agent-core-v2/test/tool/tool.test.ts @@ -602,6 +602,7 @@ function sessionMetadataStub(agents: Readonly>): ISess setGeneratedTitleIfUncustomized: async () => false, setArchived: async () => {}, registerAgent: async () => {}, + unregisterAgent: async () => {}, }; } diff --git a/packages/agent-gateway/src/start.ts b/packages/agent-gateway/src/start.ts index 062544b09..9c12d9b7e 100644 --- a/packages/agent-gateway/src/start.ts +++ b/packages/agent-gateway/src/start.ts @@ -292,35 +292,65 @@ export async function startServer(opts: ServerStartOptions): Promise => { - configChangedPublisher.close(); - await app.close(); - configWarningSubscription.dispose(); - pluginChangeSubscription.dispose(); - capabilityInstallSubscription.dispose(); - authFailureLimiter?.dispose(); - modelCatalogRefreshScheduler.dispose(); - try { - await shutdownServerTelemetry(telemetry); - } catch (error) { - logger.warn( - { err: error instanceof Error ? error.message : String(error) }, - 'telemetry shutdown failed; continuing server cleanup', - ); - } + let closePromise: Promise | undefined; + const close = (): Promise => { + closePromise ??= runClose(); + return closePromise; + }; + const runClose = async (): Promise => { + const failures: unknown[] = []; + const phase = async ( + name: string, + work: () => unknown, + required = true, + ): Promise => { + try { + await work(); + } catch (error) { + if (required) failures.push(error); + logger.warn( + { err: error instanceof Error ? error.message : String(error), phase: name }, + 'server cleanup phase failed; continuing', + ); + } + }; try { - await drainSessionMetadataWrites(); - await core.accessor.get(ISessionIndexMirror).drain(); - await core.accessor.get(IMcpOAuthService).shutdown(); - fsWatchBridge.dispose(); - const appendLogStore = core.accessor.get(IAppendLogStore); - core.dispose(); - await appendLogStore.drainRetirements(); - await drainSessionIndexMirror(); - await drainGlobalSearchDisposals(); - await drainQueryStoreDisposals(); - await drainSessionMetadataWrites(); - await drainLogCloses(); + await phase('config-publisher', () => configChangedPublisher.close()); + await phase('http', () => app.close()); + await phase('subscriptions', () => { + for (const sub of [ + configWarningSubscription, + pluginChangeSubscription, + capabilityInstallSubscription, + authFailureLimiter, + modelCatalogRefreshScheduler, + ]) { + try { + sub?.dispose(); + } catch (error) { + logger.warn( + { err: error instanceof Error ? error.message : String(error) }, + 'subscription disposal failed; continuing', + ); + } + } + }); + await phase('telemetry', () => shutdownServerTelemetry(telemetry), false); + await phase('session-metadata', () => drainSessionMetadataWrites()); + await phase('session-index', () => core.accessor.get(ISessionIndexMirror).drain()); + await phase('mcp-oauth', () => core.accessor.get(IMcpOAuthService).shutdown()); + await phase('fs-watch', () => fsWatchBridge.dispose()); + let appendLogStore: IAppendLogStore | undefined; + await phase('append-log', () => { + appendLogStore = core.accessor.get(IAppendLogStore); + }); + await phase('core', () => core.dispose()); + await phase('append-log-retirements', () => appendLogStore?.drainRetirements()); + await phase('session-index-drain', () => drainSessionIndexMirror()); + await phase('global-search', () => drainGlobalSearchDisposals()); + await phase('query-store', () => drainQueryStoreDisposals()); + await phase('session-metadata-final', () => drainSessionMetadataWrites()); + await phase('log-closes', () => drainLogCloses()); } finally { try { await registration.release(); @@ -329,6 +359,8 @@ export async function startServer(opts: ServerStartOptions): Promise 1) throw new AggregateError(failures, 'server cleanup failed'); }; const connectionRegistry = new ConnectionRegistry(); diff --git a/packages/agent-gateway/src/transport/ws/v1/registerWsV1.ts b/packages/agent-gateway/src/transport/ws/v1/registerWsV1.ts index 56c0a6513..02d2fdb00 100644 --- a/packages/agent-gateway/src/transport/ws/v1/registerWsV1.ts +++ b/packages/agent-gateway/src/transport/ws/v1/registerWsV1.ts @@ -6,7 +6,7 @@ import { type IConnectionRegistry } from '../connectionRegistry'; import type { SessionEventBroadcaster } from './sessionEventBroadcaster'; import type { FsWatchBridge } from './fsWatchBridge'; import type { JournalLogger } from './sessionEventJournal'; -import { WsConnectionV1 } from './wsConnectionV1'; +import { WS_MAX_PAYLOAD_BYTES, WsConnectionV1 } from './wsConnectionV1'; import { selectWsBearerProtocol } from '../bearerProtocol'; export const WS_PATH = '/api/v1/ws'; @@ -25,7 +25,11 @@ export interface RegisterWsV1Options { } export function registerWsV1(core: Scope, opts: RegisterWsV1Options): WebSocketServer { - const wss = new WebSocketServer({ noServer: true, handleProtocols: selectWsBearerProtocol }); + const wss = new WebSocketServer({ + noServer: true, + handleProtocols: selectWsBearerProtocol, + maxPayload: WS_MAX_PAYLOAD_BYTES, + }); const { registry, broadcaster } = opts; wss.on('connection', (socket, req) => { diff --git a/packages/agent-gateway/src/transport/ws/v1/wsConnectionV1.ts b/packages/agent-gateway/src/transport/ws/v1/wsConnectionV1.ts index 05fa7e374..ac1732de8 100644 --- a/packages/agent-gateway/src/transport/ws/v1/wsConnectionV1.ts +++ b/packages/agent-gateway/src/transport/ws/v1/wsConnectionV1.ts @@ -39,6 +39,11 @@ import { import { FsWatchBridge } from './fsWatchBridge'; const DEFAULT_MAX_BUFFER_SIZE = 1000; +export const WS_MAX_PAYLOAD_BYTES = 4 << 20; +export const WS_MAX_PENDING_CONTROLS = 64; +export const WS_MAX_SUBSCRIPTIONS = 256; +export const WS_CLOSE_OVERLOADED = 1013; +export const WS_CLOSE_POLICY = 1008; const DEFAULT_HEARTBEAT_INTERVAL_MS = 10_000; const HEARTBEAT_MISS_LIMIT = 2; @@ -50,6 +55,7 @@ const DEFAULT_MAX_BATCH_SIZE = 64; const DEFAULT_HIGH_WATER_MARK_BYTES = 1 << 20; const DEFAULT_BACKPRESSURE_RETRY_MS = 5; const DEFAULT_BACKPRESSURE_MAX_DELAY_MS = 100; +const DEFAULT_MAX_BUFFERED_BYTES_FACTOR = 8; interface InboundFrame { type: string; @@ -98,6 +104,7 @@ export class WsConnectionV1 implements BroadcastTarget { private expertTalkEvents = false; readonly subscriptions = new Map(); private controlQueue: Promise = Promise.resolve(); + private pendingControls = 0; private outbound: unknown[] = []; private flushTimer?: ReturnType; @@ -158,7 +165,7 @@ export class WsConnectionV1 implements BroadcastTarget { } get subscriptionSessionIds(): readonly string[] { - return Array.from(this.subscriptions.keys()).sort(); + return Array.from(this.subscriptions.keys()).toSorted(); } send(envelope: EventEnvelope, delivery: BroadcastDelivery = 'subscription'): void { @@ -206,9 +213,23 @@ export class WsConnectionV1 implements BroadcastTarget { } } + get pendingControlCount(): number { + return this.pendingControls; + } + private enqueueControl(task: () => Promise): void { - this.controlQueue = this.controlQueue.then(task).catch(() => { - }); + if (this.pendingControls >= WS_MAX_PENDING_CONTROLS) { + this.closeOverloaded('control queue overflow'); + return; + } + this.pendingControls += 1; + this.controlQueue = this.controlQueue + .then(() => (this.closed ? undefined : task())) + .catch(() => { + }) + .finally(() => { + this.pendingControls -= 1; + }); } private onHeartbeat(): void { @@ -420,10 +441,19 @@ export class WsConnectionV1 implements BroadcastTarget { }, ): Promise { const { accepted, resyncRequired, serverCursors, notFound } = collectors; + if (this.closed) return; + if (!this.subscriptions.has(sid) && this.subscriptions.size >= WS_MAX_SUBSCRIPTIONS) { + this.close(WS_CLOSE_POLICY, 'subscription limit exceeded'); + return; + } const ok = await this.broadcaster.subscribe(sid, this, filter, transcriptGrades, { deferTranscriptReset: cursor !== undefined, transcriptSince, }); + if (this.closed) { + if (ok) this.broadcaster.unsubscribe(sid, this); + return; + } if (!ok) { if (notFound !== undefined) notFound.push(sid); else resyncRequired.push(sid); @@ -490,6 +520,10 @@ export class WsConnectionV1 implements BroadcastTarget { private sendSubscribedFrame(msg: unknown): void { if (this.closed) return; + if (this.outbound.length >= this.maxBufferSize) { + this.closeOverloaded('outbound buffer overflow'); + return; + } this.outbound.push(msg); if (this.outbound.length >= this.maxBatchSize) { this.flush(); @@ -545,6 +579,10 @@ export class WsConnectionV1 implements BroadcastTarget { const now = Date.now(); if (this.backpressureSince === undefined) this.backpressureSince = now; if (now - this.backpressureSince >= DEFAULT_BACKPRESSURE_MAX_DELAY_MS) { + if (this.socket.bufferedAmount > this.highWaterMarkBytes * DEFAULT_MAX_BUFFERED_BYTES_FACTOR) { + this.closeOverloaded('slow consumer'); + return; + } this.flush(true); return; } @@ -556,6 +594,12 @@ export class WsConnectionV1 implements BroadcastTarget { this.backpressureRetryTimer.unref?.(); } + private closeOverloaded(reason: string): void { + if (this.closed) return; + this.outbound = []; + this.close(WS_CLOSE_OVERLOADED, reason); + } + close(code = 1000, reason?: string): void { if (this.closed) return; this.flush(true); @@ -585,7 +629,7 @@ export class WsConnectionV1 implements BroadcastTarget { function asStringArray(value: unknown): string[] { if (!Array.isArray(value)) return []; - return value.filter((v): v is string => typeof v === 'string'); + return value.filter((v): v is string => typeof v === 'string').slice(0, WS_MAX_SUBSCRIPTIONS); } function parseAgentFilter(value: unknown): Record | undefined { diff --git a/packages/agent-gateway/test/boot.test.ts b/packages/agent-gateway/test/boot.test.ts index e44ac2a7e..14e07905e 100644 --- a/packages/agent-gateway/test/boot.test.ts +++ b/packages/agent-gateway/test/boot.test.ts @@ -11,6 +11,7 @@ import { IBootstrapService, IFileSystemStorageService, IHostRequestHeaders, + IMcpOAuthService, InMemoryStorageService, ITelemetryService, noopTelemetryService, @@ -274,6 +275,49 @@ describe('server-v2 boot', () => { expect(process.listenerCount('uncaughtException')).toBe(exceptionBefore.length); }); + it('joins concurrent close calls into one shutdown', async () => { + home = await mkdtemp(join(tmpdir(), 'pythinker-server-v2-close-once-')); + server = await startServer({ + hostIdentity: TEST_HOST_IDENTITY, + host: '127.0.0.1', + port: 0, + homeDir: home, + logLevel: 'silent', + }); + const dispose = vi.spyOn(server.core, 'dispose'); + + await Promise.all([server.close(), server.close()]); + await server.close(); + server = undefined; + + expect(dispose).toHaveBeenCalledTimes(1); + expect(await listLiveServerInstances(home)).toEqual([]); + }); + + it('runs every cleanup phase and rethrows when a required drain fails', async () => { + home = await mkdtemp(join(tmpdir(), 'pythinker-server-v2-close-fail-')); + const rejectionBefore = process.listenerCount('unhandledRejection'); + const exceptionBefore = process.listenerCount('uncaughtException'); + server = await startServer({ + hostIdentity: TEST_HOST_IDENTITY, + host: '127.0.0.1', + port: 0, + homeDir: home, + logLevel: 'silent', + }); + const core = server.core; + const oauth = core.accessor.get(IMcpOAuthService); + vi.spyOn(oauth, 'shutdown').mockRejectedValueOnce(new Error('oauth shutdown failed')); + + await expect(server.close()).rejects.toThrow('oauth shutdown failed'); + server = undefined; + + expect(() => core.accessor.get(IBootstrapService)).toThrow(); + expect(await listLiveServerInstances(home)).toEqual([]); + expect(process.listenerCount('unhandledRejection')).toBe(rejectionBefore); + expect(process.listenerCount('uncaughtException')).toBe(exceptionBefore); + }); + it('does not leave process handlers installed when startup fails', async () => { home = await mkdtemp(join(tmpdir(), 'pythinker-server-v2-')); const emptyAssets = await mkdtemp(join(tmpdir(), 'pythinker-server-v2-assets-')); diff --git a/packages/agent-gateway/test/wsConnectionV1.test.ts b/packages/agent-gateway/test/wsConnectionV1.test.ts index bbf068eff..6d04dd870 100644 --- a/packages/agent-gateway/test/wsConnectionV1.test.ts +++ b/packages/agent-gateway/test/wsConnectionV1.test.ts @@ -7,6 +7,10 @@ import type { IConnectionRegistry } from '../src/transport/ws/connectionRegistry import type { SessionEventBroadcaster } from '../src/transport/ws/v1/sessionEventBroadcaster'; import { type WsConnectionV1Options, + WS_CLOSE_OVERLOADED, + WS_CLOSE_POLICY, + WS_MAX_PENDING_CONTROLS, + WS_MAX_SUBSCRIPTIONS, WsConnectionV1, coalesceFrames, } from '../src/transport/ws/v1/wsConnectionV1'; @@ -67,6 +71,10 @@ function makeBroadcaster(): SessionEventBroadcaster { } as unknown as SessionEventBroadcaster; } +function withBroadcaster(overrides: Partial): SessionEventBroadcaster { + return Object.assign(makeBroadcaster(), overrides); +} + function makeRegistry(): IConnectionRegistry { return { add: () => {}, @@ -764,6 +772,137 @@ describe('WsConnectionV1 outbound buffer', () => { conn.close(); }); + it('force-flushes a moderately slow consumer after the backpressure delay', async () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { flushIntervalMs: 16, highWaterMarkBytes: 100 }); + socket.sent = []; + + socket.bufferedAmount = 200; + conn.send(delta('s1', 'main', 1, 'Hello', 0)); + await vi.advanceTimersByTimeAsync(120); + + expect(socket.frames()).toHaveLength(1); + expect(socket.closeCalls).toEqual([]); + conn.close(); + }); + + it('disconnects a consumer whose socket buffer stays past the hard byte bound', async () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { flushIntervalMs: 16, highWaterMarkBytes: 100 }); + socket.sent = []; + + socket.bufferedAmount = 100 * 8 + 1; + conn.send(delta('s1', 'main', 1, 'Hello', 0)); + await vi.advanceTimersByTimeAsync(120); + + expect(socket.sent).toEqual([]); + expect(socket.closeCalls).toEqual([{ code: WS_CLOSE_OVERLOADED, reason: 'slow consumer' }]); + expect(conn.subscriptionSessionIds).toEqual([]); + }); + + it('disconnects when the outbound queue exceeds the advertised buffer size', () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { + flushIntervalMs: 16, + highWaterMarkBytes: 100, + maxBufferSize: 3, + maxBatchSize: 100, + }); + socket.sent = []; + socket.bufferedAmount = 200; + + for (let i = 0; i < 4; i += 1) conn.send(delta('s1', 'main', 1, `chunk${String(i)}`, i)); + + expect(socket.sent).toEqual([]); + expect(socket.closeCalls).toEqual([ + { code: WS_CLOSE_OVERLOADED, reason: 'outbound buffer overflow' }, + ]); + }); + + it('disconnects a client that floods control frames faster than they drain', () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { + broadcaster: withBroadcaster({ subscribe: () => new Promise(() => {}) }), + }); + socket.sent = []; + + for (let i = 0; i <= WS_MAX_PENDING_CONTROLS; i += 1) { + socket.emit( + 'message', + JSON.stringify({ type: 'subscribe', id: `f${String(i)}`, payload: { session_ids: ['s1'] } }), + ); + } + + expect(conn.pendingControlCount).toBe(WS_MAX_PENDING_CONTROLS); + expect(socket.closeCalls).toEqual([ + { code: WS_CLOSE_OVERLOADED, reason: 'control queue overflow' }, + ]); + }); + + it('releases a subscription that completes after the connection closed', async () => { + const socket = new FakeSocket(); + let releaseSubscribe: (ok: boolean) => void = () => {}; + const unsubscribe = vi.fn(); + const conn = makeConn(socket, { + broadcaster: withBroadcaster({ + subscribe: () => + new Promise((resolve) => { + releaseSubscribe = resolve; + }), + unsubscribe, + }), + }); + socket.emit( + 'message', + JSON.stringify({ type: 'subscribe', id: 'f1', payload: { session_ids: ['s1'] } }), + ); + await vi.advanceTimersByTimeAsync(0); + + socket.emit('close'); + releaseSubscribe(true); + await vi.advanceTimersByTimeAsync(0); + + expect(unsubscribe).toHaveBeenCalledWith('s1', conn); + expect(conn.subscriptionSessionIds).toEqual([]); + expect(conn.pendingControlCount).toBe(0); + }); + + it('skips queued controls once the connection is closed', async () => { + const socket = new FakeSocket(); + const subscribe = vi.fn(async () => true); + makeConn(socket, { broadcaster: withBroadcaster({ subscribe }) }); + socket.emit( + 'message', + JSON.stringify({ type: 'subscribe', id: 'f1', payload: { session_ids: ['s1'] } }), + ); + socket.emit('close'); + await vi.advanceTimersByTimeAsync(0); + + expect(subscribe).not.toHaveBeenCalled(); + }); + + it('caps the number of session subscriptions per connection', async () => { + const socket = new FakeSocket(); + const conn = makeConn(socket); + const ids = Array.from({ length: WS_MAX_SUBSCRIPTIONS + 1 }, (_, i) => `s${String(i)}`); + socket.emit( + 'message', + JSON.stringify({ type: 'subscribe', id: 'f1', payload: { session_ids: ids } }), + ); + await vi.advanceTimersByTimeAsync(0); + expect(conn.subscriptionSessionIds).toHaveLength(WS_MAX_SUBSCRIPTIONS); + + socket.emit( + 'message', + JSON.stringify({ type: 'subscribe', id: 'f2', payload: { session_ids: ['extra'] } }), + ); + await vi.advanceTimersByTimeAsync(0); + + expect(socket.closeCalls).toEqual([ + { code: WS_CLOSE_POLICY, reason: 'subscription limit exceeded' }, + ]); + }); + it('force-flushes buffered subscription frames on close', () => { const socket = new FakeSocket(); const conn = makeConn(socket, { flushIntervalMs: 1000 }); diff --git a/packages/kosong/src/providers/dsml-tool-parser.ts b/packages/kosong/src/providers/dsml-tool-parser.ts index 63bd8f69c..2b89bdabd 100644 --- a/packages/kosong/src/providers/dsml-tool-parser.ts +++ b/packages/kosong/src/providers/dsml-tool-parser.ts @@ -1,11 +1,22 @@ import type { StreamedMessagePart, ToolCall } from '#/message'; -const CONTAINER_OPEN_RE = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*tool_calls\s*>/i; -const CONTAINER_CLOSE_RE = /^<\/\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*tool_calls\s*>/i; -const INVOKE_OPEN_RE = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*invoke(?:\s+[^>]*)?>/i; -const INVOKE_CLOSE_RE = /<\/\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*invoke\s*>/i; -const HERMES_OPEN_RE = /^/i; -const HERMES_CLOSE_RE = /<\/tool_call>/i; +const MARK = String.raw`\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*`; +const CONTAINER_OPEN_RE = new RegExp(String.raw`<${MARK}tool_calls\s*>`, 'yi'); +const CONTAINER_CLOSE_RE = new RegExp(String.raw``, 'yi'); +const INVOKE_OPEN_RE = new RegExp(String.raw`<${MARK}invoke(?:\s+[^>]*)?>`, 'yi'); +const INVOKE_CLOSE_RE = new RegExp(String.raw``, 'gi'); +const PARAM_OPEN_RE = new RegExp(String.raw`<${MARK}parameter\s+([^>]*?)>`, 'yi'); +const PARAM_CLOSE_RE = new RegExp(String.raw``, 'gi'); +const HERMES_OPEN_RE = //yi; +const HERMES_CLOSE_RE = /<\/tool_call>/gi; +const MARKED_RE = /<\/?\s*(?:[｜|]|DSML)/yi; +const NAME_ATTR_RE = /\bname\s*=\s*(?:"([^"]+)"|'([^']+)'|([^\s>]+))/i; +const STRING_ATTR_RE = /\bstring\s*=\s*(?:"(true|false)"|'(true|false)'|(true|false))/i; + +export const DSML_MAX_TAG_CHARS = 4096; +export const DSML_MAX_ENVELOPE_CHARS = 2 * 1024 * 1024; +const TAIL_MAX = 1024; +const TAG_NAMES = ['tool_calls', 'tool_call', 'invoke', 'parameter']; function unescapeXml(value: string): string { return value @@ -46,115 +57,165 @@ function parseParameterValue(rawVal: string, isStringAttr: boolean | undefined): return unescaped; } -function parseInvokeBody(invokeContent: string): Record { - const paramRegex = - /<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*parameter\s+([^>]*?)>([\s\S]*?)<\/\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*parameter\s*>/gi; - const args: Record = {}; - let paramFound = false; - let match: RegExpExecArray | null = null; - - while ((match = paramRegex.exec(invokeContent)) !== null) { - const attrStr = match[1] ?? ''; - const rawVal = match[2] ?? ''; - const nameMatch = /\bname\s*=\s*(?:"([^"]+)"|'([^']+)'|([^\s>]+))/i.exec(attrStr); - const paramName = nameMatch ? (nameMatch[1] ?? nameMatch[2] ?? nameMatch[3]) : undefined; - if (paramName) { - paramFound = true; - const stringAttrMatch = - /\bstring\s*=\s*(?:"(true|false)"|'(true|false)'|(true|false))/i.exec(attrStr); - const stringAttrVal = stringAttrMatch - ? (stringAttrMatch[1] ?? stringAttrMatch[2] ?? stringAttrMatch[3]) - : undefined; - const isStringAttr = - stringAttrVal !== undefined ? stringAttrVal.toLowerCase() === 'true' : undefined; - args[paramName] = parseParameterValue(rawVal, isStringAttr); - } - } +function attrValue(match: RegExpExecArray | null): string | undefined { + if (!match) return undefined; + return match[1] ?? match[2] ?? match[3]; +} - if (paramFound) { - return args; - } +function skipWhitespace(text: string, from: number): number { + let pos = from; + while (pos < text.length && /\s/.test(text[pos] as string)) pos += 1; + return pos; +} - const trimmed = invokeContent.trim(); - if (trimmed.startsWith('{') && trimmed.endsWith('}')) { +function stickyExec(re: RegExp, text: string, at: number): RegExpExecArray | null { + re.lastIndex = at; + return re.exec(text); +} + +export function parseInvokeBody(invokeContent: string): Record | null { + const args: Record = Object.create(null) as Record; + let pos = skipWhitespace(invokeContent, 0); + if (pos === invokeContent.length) return args; + + if (invokeContent[pos] === '{') { + const trimmed = invokeContent.trim(); + if (!trimmed.endsWith('}')) return null; try { - const parsed = JSON.parse(trimmed); - if (typeof parsed === 'object' && parsed !== null && !Array.isArray(parsed)) { - return parsed as Record; + const parsed: unknown = JSON.parse(trimmed); + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) return null; + for (const key of Object.keys(parsed)) { + args[key] = (parsed as Record)[key]; } - } catch {} + return args; + } catch { + return null; + } + } + + while (pos < invokeContent.length) { + const open = stickyExec(PARAM_OPEN_RE, invokeContent, pos); + if (!open) return null; + const attrStr = open[1] ?? ''; + const paramName = attrValue(NAME_ATTR_RE.exec(attrStr)); + if (!paramName) return null; + const valueStart = open.index + open[0].length; + const close = stickyExec(PARAM_CLOSE_RE, invokeContent, valueStart); + if (!close) return null; + const stringAttrVal = attrValue(STRING_ATTR_RE.exec(attrStr)); + const isStringAttr = + stringAttrVal !== undefined ? stringAttrVal.toLowerCase() === 'true' : undefined; + args[paramName] = parseParameterValue( + invokeContent.slice(valueStart, close.index), + isStringAttr, + ); + pos = skipWhitespace(invokeContent, close.index + close[0].length); } + return args; +} - return {}; +function newCallId(): string { + return `call_${crypto.randomUUID().replaceAll('-', '').slice(0, 24)}`; } function parseInvokeTag(invokeBlock: string): ToolCall | null { - const openMatch = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*invoke\s+([^>]*?)>/i.exec(invokeBlock); + const openMatch = stickyExec(INVOKE_OPEN_RE, invokeBlock, 0); if (!openMatch) return null; - - const attrStr = openMatch[1] ?? ''; - const nameMatch = /\bname\s*=\s*(?:"([^"]+)"|'([^']+)'|([^\s>]+))/i.exec(attrStr); - const toolName = nameMatch ? (nameMatch[1] ?? nameMatch[2] ?? nameMatch[3]) : undefined; + const toolName = attrValue(NAME_ATTR_RE.exec(openMatch[0].slice(0, -1))); if (!toolName) return null; - const closeMatch = INVOKE_CLOSE_RE.exec(invokeBlock); - if (!closeMatch) return null; - const innerContent = invokeBlock.slice(openMatch[0].length, closeMatch.index); - - const args = parseInvokeBody(innerContent); - return { - type: 'function', - id: `call_${crypto.randomUUID().replaceAll('-', '').slice(0, 24)}`, - name: toolName, - arguments: JSON.stringify(args), - }; + const closeMatch = stickyExec(INVOKE_CLOSE_RE, invokeBlock, openMatch[0].length); + if (!closeMatch || closeMatch.index + closeMatch[0].length !== invokeBlock.length) return null; + const args = parseInvokeBody(invokeBlock.slice(openMatch[0].length, closeMatch.index)); + if (args === null) return null; + return { type: 'function', id: newCallId(), name: toolName, arguments: JSON.stringify(args) }; } function parseHermesToolCall(toolCallBlock: string): ToolCall | null { - if (!HERMES_CLOSE_RE.test(toolCallBlock)) return null; const inner = toolCallBlock .replace(/^/i, '') .replace(/<\/tool_call>$/i, '') .trim(); try { - const parsed = JSON.parse(inner); - if (parsed && typeof parsed.name === 'string') { - const args = - typeof parsed.arguments === 'string' - ? parsed.arguments - : JSON.stringify(parsed.arguments ?? {}); - return { - type: 'function', - id: `call_${crypto.randomUUID().replaceAll('-', '').slice(0, 24)}`, - name: parsed.name, - arguments: args, - }; - } - } catch {} - return null; + const parsed: unknown = JSON.parse(inner); + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) return null; + const record = parsed as Record; + if (typeof record['name'] !== 'string') return null; + const rawArgs = record['arguments']; + const args = typeof rawArgs === 'string' ? rawArgs : JSON.stringify(rawArgs ?? {}); + return { type: 'function', id: newCallId(), name: record['name'], arguments: args }; + } catch { + return null; + } } -function isPotentialTagPrefix(s: string): boolean { - if (!s.startsWith('<')) return false; - const lower = s.toLowerCase(); +type TagKind = + | 'container-open' + | 'container-close' + | 'invoke-open' + | 'hermes-open' + | 'partial' + | 'none'; + +function markerPrefixState(lower: string): { rest: string; partial: boolean } { let rest = lower.startsWith(' 0 && rest.length < 4 && 'dsml'.startsWith(rest)) { + return { rest: '', partial: true }; } if (rest.startsWith('dsml')) { rest = rest.slice(4).trimStart(); - if (rest.startsWith('｜') || rest.startsWith('|')) { - rest = rest.slice(1).trimStart(); + if (rest.startsWith('｜') || rest.startsWith('|')) rest = rest.slice(1).trimStart(); + } + return { rest, partial: rest.length === 0 }; +} + +function classifyTag(buffer: string, at: number): { kind: TagKind; length: number } { + const containerOpen = stickyExec(CONTAINER_OPEN_RE, buffer, at); + if (containerOpen) return { kind: 'container-open', length: containerOpen[0].length }; + const containerClose = stickyExec(CONTAINER_CLOSE_RE, buffer, at); + if (containerClose) return { kind: 'container-close', length: containerClose[0].length }; + const invokeOpen = stickyExec(INVOKE_OPEN_RE, buffer, at); + if (invokeOpen) return { kind: 'invoke-open', length: invokeOpen[0].length }; + const hermesOpen = stickyExec(HERMES_OPEN_RE, buffer, at); + if (hermesOpen) return { kind: 'hermes-open', length: hermesOpen[0].length }; + + const lower = buffer.slice(at, at + DSML_MAX_TAG_CHARS + 1).toLowerCase(); + const { rest, partial } = markerPrefixState(lower); + if (partial) return { kind: 'partial', length: 0 }; + if (rest.includes('>')) return { kind: 'none', length: 0 }; + for (const name of TAG_NAMES) { + if (name.startsWith(rest)) return { kind: 'partial', length: 0 }; + if (rest.startsWith(name)) { + const after = rest.slice(name.length); + if (after.length === 0 || /^\s/.test(after)) return { kind: 'partial', length: 0 }; } } - if (rest.length === 0) return true; - const targets = ['tool_calls', 'tool_call', 'invoke', 'parameter']; - return targets.some((t) => t.startsWith(rest) || rest.startsWith(t)); + return { kind: 'none', length: 0 }; +} + +type Mode = 'text' | 'invoke' | 'hermes'; + +interface ContainerHold { + raw: string; + parts: StreamedMessagePart[]; } export class DsmlStreamParser { private _buffer = ''; + private _pos = 0; + private _mode: Mode = 'text'; + private _envelope: string[] = []; + private _envelopeLength = 0; + private _tail = ''; + private _stripWhitespace = false; + private _whitespaceHold = ''; + private _inContainer = false; + private _hold: ContainerHold | null = null; + private _lineStart = true; + private _linePrefix = ''; + private _fence: string | null = null; private _hasExtractedToolCalls = false; get hasExtractedToolCalls(): boolean { @@ -162,139 +223,276 @@ export class DsmlStreamParser { } feed(chunk: string): StreamedMessagePart[] { - this._buffer += chunk; const parts: StreamedMessagePart[] = []; + if (this._mode === 'text') { + this._buffer += chunk; + } else { + this._scanEnvelope(parts, chunk); + } + if (this._mode === 'text') { + this._drainText(parts); + } + return parts; + } - while (this._buffer.length > 0) { - const ltIdx = this._buffer.indexOf('<'); - if (ltIdx === -1) { - parts.push({ type: 'text', text: this._buffer }); - this._buffer = ''; - break; - } - - if (ltIdx > 0) { - parts.push({ type: 'text', text: this._buffer.slice(0, ltIdx) }); - this._buffer = this._buffer.slice(ltIdx); - } + flush(): StreamedMessagePart[] { + const parts: StreamedMessagePart[] = []; + if (this._mode !== 'text') { + const block = this._envelope.join('') + this._tail; + this._leaveEnvelope(); + this._buffer = block; + this._pos = 0; + } + if (this._whitespaceHold.length > 0) { + this._emitText(parts, this._whitespaceHold.replace(/^\r?\n/, '')); + this._whitespaceHold = ''; + } + this._stripWhitespace = false; + if (this._pos < this._buffer.length) { + this._emitText(parts, this._take(this._buffer.length - this._pos)); + } + this._buffer = ''; + this._pos = 0; + if (this._inContainer) { + this._releaseHold(parts, true); + this._inContainer = false; + } + return parts; + } - const openContainer = CONTAINER_OPEN_RE.exec(this._buffer); - if (openContainer) { - this._buffer = this._buffer.slice(openContainer[0].length); - if (/^\s* this._pos) { + this._whitespaceHold += this._take(end - this._pos); + } + if (this._pos >= this._buffer.length) break; + this._stripWhitespace = false; + if (this._buffer[this._pos] === '<') { + this._noteText(this._whitespaceHold); } else { - this._buffer = this._buffer.replace(/^\r?\n/, ''); + this._emitText(parts, this._whitespaceHold.replace(/^\r?\n/, '')); } - continue; + this._whitespaceHold = ''; } - const closeContainer = CONTAINER_CLOSE_RE.exec(this._buffer); - if (closeContainer) { - this._buffer = this._buffer.slice(closeContainer[0].length); - if (/^\s* this._pos) { + this._emitText(parts, this._take(ltIdx - this._pos)); continue; } - const invokeOpen = INVOKE_OPEN_RE.exec(this._buffer); - if (invokeOpen) { - const closeMatch = INVOKE_CLOSE_RE.exec(this._buffer); - if (!closeMatch) { - break; - } - const invokeEnd = closeMatch.index + closeMatch[0].length; - const invokeBlock = this._buffer.slice(0, invokeEnd); - const toolCall = parseInvokeTag(invokeBlock); - if (toolCall) { - this._hasExtractedToolCalls = true; - parts.push(toolCall); - this._buffer = this._buffer.slice(invokeEnd); - if (/^\s* DSML_MAX_TAG_CHARS) { + this._emitPlainTag(parts); + continue; } - const toolCallEnd = closeMatch.index + closeMatch[0].length; - const block = this._buffer.slice(0, toolCallEnd); - const toolCall = parseHermesToolCall(block); - if (toolCall) { - this._hasExtractedToolCalls = true; - parts.push(toolCall); - this._buffer = this._buffer.slice(toolCallEnd); - if (/^\s* DSML_MAX_ENVELOPE_CHARS) { + this._abandonEnvelope(parts, block, rest); + return; + } + const call = mode === 'invoke' ? parseInvokeTag(block) : parseHermesToolCall(block); + this._record(block); + if (call) { + this._emitCall(parts, call); + this._stripWhitespace = true; + } else { + this._emitText(parts, block); } + this._buffer = rest; + this._pos = 0; + return; + } - parts.push({ type: 'text', text: this._buffer.slice(0, nextLt) }); - this._buffer = this._buffer.slice(nextLt); + if (this._envelopeLength + probe.length >= DSML_MAX_ENVELOPE_CHARS) { + const block = this._envelope.join('') + probe; + this._leaveEnvelope(); + this._abandonEnvelope(parts, block, ''); + return; } - return parts; + const lastLt = probe.lastIndexOf('<'); + if (lastLt === -1 || probe.length - lastLt > TAIL_MAX) { + this._envelope.push(probe); + this._envelopeLength += probe.length; + this._tail = ''; + } else { + const committed = probe.slice(0, lastLt); + this._envelope.push(committed); + this._envelopeLength += committed.length; + this._tail = probe.slice(lastLt); + } } - flush(): StreamedMessagePart[] { - const parts: StreamedMessagePart[] = []; - if (this._buffer.length > 0) { - const invokeOpen = INVOKE_OPEN_RE.exec(this._buffer); - if (invokeOpen && INVOKE_CLOSE_RE.test(this._buffer)) { - const toolCall = parseInvokeTag(this._buffer); - if (toolCall) { - this._hasExtractedToolCalls = true; - parts.push(toolCall); - this._buffer = ''; - return parts; - } + private _abandonEnvelope(parts: StreamedMessagePart[], block: string, rest: string): void { + const nextLt = block.indexOf('<', 1); + const cut = nextLt === -1 ? block.length : nextLt; + const text = block.slice(0, cut); + this._record(text); + this._emitText(parts, text); + this._buffer = block.slice(cut) + rest; + this._pos = 0; + } + + private _take(length: number): string { + const taken = this._buffer.slice(this._pos, this._pos + length); + this._pos += length; + this._record(taken); + return taken; + } + + private _record(text: string): void { + if (this._hold === null) return; + this._hold.raw += text; + } + + private _releaseHold(parts: StreamedMessagePart[], asText: boolean): void { + const hold = this._hold; + this._hold = null; + if (hold === null) return; + if (asText) { + if (hold.raw.length > 0) parts.push({ type: 'text', text: hold.raw }); + } else { + parts.push(...hold.parts); + } + } + + private _emitCall(parts: StreamedMessagePart[], call: ToolCall): void { + this._hasExtractedToolCalls = true; + this._releaseHold(parts, false); + parts.push(call); + } + + private _emitText(parts: StreamedMessagePart[], text: string): void { + if (text.length === 0) return; + this._noteText(text); + const part: StreamedMessagePart = { type: 'text', text }; + if (this._hold !== null) { + this._hold.parts.push(part); + if (this._hold.raw.length > DSML_MAX_ENVELOPE_CHARS) { + this._releaseHold(parts, true); } - const hermesOpen = HERMES_OPEN_RE.exec(this._buffer); - if (hermesOpen && HERMES_CLOSE_RE.test(this._buffer)) { - const toolCall = parseHermesToolCall(this._buffer); - if (toolCall) { - this._hasExtractedToolCalls = true; - parts.push(toolCall); - this._buffer = ''; - return parts; + return; + } + parts.push(part); + } + + private _noteText(text: string): void { + for (const ch of text) { + if (ch === '\n') { + this._lineStart = true; + this._linePrefix = ''; + continue; + } + if (!this._lineStart) continue; + if (ch === ' ' && this._linePrefix.length < 3 && /^ *$/.test(this._linePrefix)) { + this._linePrefix += ch; + continue; + } + if (ch !== '`' && ch !== '~') { + this._lineStart = false; + continue; + } + const run = this._linePrefix.trimStart(); + if (run.length > 0 && run[0] !== ch) { + this._lineStart = false; + continue; + } + this._linePrefix += ch; + if (run.length + 1 === 3) { + this._lineStart = false; + if (this._fence === null) { + this._fence = ch; + } else if (this._fence === ch) { + this._fence = null; } } - parts.push({ type: 'text', text: this._buffer }); - this._buffer = ''; } - return parts; } } @@ -315,6 +513,6 @@ export function extractDsmlToolCalls(text: string): { } } - const cleanText = toolCalls.length > 0 ? textParts.join('').trim() : text; - return { cleanText, toolCalls }; + const joined = textParts.join(''); + return { cleanText: toolCalls.length > 0 ? joined.trim() : joined, toolCalls }; } diff --git a/packages/kosong/src/providers/openai-legacy.ts b/packages/kosong/src/providers/openai-legacy.ts index ffbc6d5ab..c9008a8c3 100644 --- a/packages/kosong/src/providers/openai-legacy.ts +++ b/packages/kosong/src/providers/openai-legacy.ts @@ -409,7 +409,8 @@ export class OpenAILegacyStreamedMessage implements StreamedMessage { let text = message.content ?? null; let extractedToolCalls: ToolCall[] = []; - if (text) { + const hasNativeToolCalls = (message.tool_calls ?? []).some(isFunctionToolCall); + if (text && !hasNativeToolCalls) { const parsed = extractDsmlToolCalls(text); if (parsed.toolCalls.length > 0) { text = parsed.cleanText; @@ -445,6 +446,8 @@ export class OpenAILegacyStreamedMessage implements StreamedMessage { ): AsyncGenerator { const bufferedToolCalls = new Map(); const dsmlParser = new DsmlStreamParser(); + const recoveredToolCalls: ToolCall[] = []; + let nativeToolCallsSeen = false; try { for await (const chunk of response) { @@ -481,13 +484,18 @@ export class OpenAILegacyStreamedMessage implements StreamedMessage { // text content if (delta.content) { for (const part of dsmlParser.feed(delta.content)) { - yield part; + if (part.type === 'function') { + recoveredToolCalls.push(part); + } else { + yield part; + } } } // tool calls — preserve `index` on every yielded part so the generate // loop can route interleaved argument deltas from parallel tool calls. for (const toolCall of delta.tool_calls ?? []) { + nativeToolCallsSeen = true; for (const part of convertChatCompletionStreamToolCall(toolCall, bufferedToolCalls)) { yield part; } @@ -495,10 +503,17 @@ export class OpenAILegacyStreamedMessage implements StreamedMessage { } for (const part of dsmlParser.flush()) { - yield part; + if (part.type === 'function') { + recoveredToolCalls.push(part); + } else { + yield part; + } } - if (dsmlParser.hasExtractedToolCalls) { + if (!nativeToolCallsSeen && recoveredToolCalls.length > 0) { this._hasExtractedToolCalls = true; + for (const toolCall of recoveredToolCalls) { + yield toolCall; + } } } catch (error: unknown) { throw convertOpenAIError(error); diff --git a/packages/kosong/src/providers/pythinker.ts b/packages/kosong/src/providers/pythinker.ts index 9ce464109..c40e34114 100644 --- a/packages/kosong/src/providers/pythinker.ts +++ b/packages/kosong/src/providers/pythinker.ts @@ -339,7 +339,8 @@ class PythinkerStreamedMessage implements StreamedMessage { let text = message.content ?? null; let extractedToolCalls: ToolCall[] = []; - if (text) { + const hasNativeToolCalls = (message.tool_calls ?? []).some(isFunctionToolCall); + if (text && !hasNativeToolCalls) { const parsed = extractDsmlToolCalls(text); if (parsed.toolCalls.length > 0) { text = parsed.cleanText; @@ -374,6 +375,8 @@ class PythinkerStreamedMessage implements StreamedMessage { ): AsyncGenerator { const bufferedToolCalls = new Map(); const dsmlParser = new DsmlStreamParser(); + const recoveredToolCalls: ToolCall[] = []; + let nativeToolCallsSeen = false; try { for await (const chunk of response) { @@ -414,13 +417,18 @@ class PythinkerStreamedMessage implements StreamedMessage { // text content if (delta.content) { for (const part of dsmlParser.feed(delta.content)) { - yield part; + if (part.type === 'function') { + recoveredToolCalls.push(part); + } else { + yield part; + } } } // tool calls — preserve `index` on every yielded part so the generate // loop can route interleaved argument deltas from parallel tool calls. for (const toolCall of delta.tool_calls ?? []) { + nativeToolCallsSeen = true; for (const part of convertChatCompletionStreamToolCall(toolCall, bufferedToolCalls)) { yield part; } @@ -428,10 +436,17 @@ class PythinkerStreamedMessage implements StreamedMessage { } for (const part of dsmlParser.flush()) { - yield part; + if (part.type === 'function') { + recoveredToolCalls.push(part); + } else { + yield part; + } } - if (dsmlParser.hasExtractedToolCalls) { + if (!nativeToolCallsSeen && recoveredToolCalls.length > 0) { this._hasExtractedToolCalls = true; + for (const toolCall of recoveredToolCalls) { + yield toolCall; + } } } catch (error: unknown) { throw convertOpenAIError(error, classifyPythinkerQuotaError); diff --git a/packages/kosong/test/openai-legacy.test.ts b/packages/kosong/test/openai-legacy.test.ts index 687e4020a..0af904b40 100644 --- a/packages/kosong/test/openai-legacy.test.ts +++ b/packages/kosong/test/openai-legacy.test.ts @@ -1507,6 +1507,57 @@ describe('OpenAILegacyChatProvider', () => { expect(messages[0]).not.toHaveProperty('reasoning_content'); }); + it('lets native streamed tool calls win over a DSML echo in content', async () => { + const provider = new OpenAILegacyChatProvider({ + model: 'deepseek-chat', + apiKey: 'test-key', + stream: true, + }); + const echo = + '<｜DSML｜tool_calls>\n<｜DSML｜invoke name="Read">\n<｜DSML｜parameter name="filePath" string="true">a.ts\n\n'; + + async function* mockedStream(): AsyncIterable> { + yield { id: 'c1', choices: [{ index: 0, delta: { content: echo } }] }; + yield { + id: 'c1', + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + id: 'call_native', + type: 'function', + function: { name: 'Read', arguments: '{"filePath":"a.ts"}' }, + }, + ], + }, + finish_reason: 'tool_calls', + }, + ], + }; + } + + ( + provider as unknown as { _client: { chat: { completions: { create: unknown } } } } + )._client.chat.completions.create = vi + .fn() + .mockResolvedValue(mockedStream()); + + const stream = await provider.generate( + '', + [], + [{ role: 'user', content: [{ type: 'text', text: 'q' }], toolCalls: [] }], + ); + const parts: StreamedMessagePart[] = []; + for await (const part of stream) parts.push(part); + + const calls = parts.filter((part) => part.type === 'function'); + expect(calls).toHaveLength(1); + expect(calls[0]).toMatchObject({ id: 'call_native', name: 'Read' }); + }); + it('yields ThinkPart from streaming response even without explicit reasoningKey', async () => { const provider = new OpenAILegacyChatProvider({ model: 'deepseek-reasoner',