From 98be51b74cf7714a71513aea2743034a36da9548 Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Mon, 20 Jul 2026 15:24:57 -0400 Subject: [PATCH] fix(core): harden structured output invariants --- packages/core/src/session/runner/llm.ts | 18 +- packages/core/src/tool-output-store.ts | 64 ++++- packages/core/src/tool/edit.ts | 2 +- packages/core/src/tool/patch.ts | 2 +- packages/core/src/tool/read.ts | 1 + packages/core/src/tool/registry.ts | 76 ++--- packages/core/src/tool/shell.ts | 1 + packages/core/src/tool/structured.ts | 24 ++ .../test/session-runner-tool-registry.test.ts | 261 ++++++++++------- packages/core/test/session-runner.test.ts | 262 +++++++++++------- packages/core/test/tool-edit.test.ts | 9 +- packages/core/test/tool-output-store.test.ts | 47 ++++ packages/core/test/tool-patch.test.ts | 9 +- packages/core/test/tool-shell.test.ts | 29 +- 14 files changed, 540 insertions(+), 265 deletions(-) diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 959005754a..aecf919330 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -305,9 +305,17 @@ const layer = Layer.effect( // settles so the model may recover. A typed infrastructure failure (tool output // could not be persisted) also fails the assistant and then fails the drain. const settledFailure = settledCauses.find((cause) => !Cause.hasInterrupts(cause) && !isUserDeclined(cause)) - const typedError = - settledFailure === undefined ? undefined : Option.getOrUndefined(Cause.findErrorOption(settledFailure)) - const storageDefect = settledFailure?.reasons.find( + const infrastructureFailure = settledCauses.find( + (cause) => + Option.getOrUndefined(Cause.findErrorOption(cause)) instanceof ToolOutputStore.StorageError || + cause.reasons.some( + (reason) => Cause.isDieReason(reason) && reason.defect instanceof ToolOutputStore.StorageError, + ), + ) + const typedError = infrastructureFailure + ? Option.getOrUndefined(Cause.findErrorOption(infrastructureFailure)) + : undefined + const storageDefect = infrastructureFailure?.reasons.find( (reason) => Cause.isDieReason(reason) && reason.defect instanceof ToolOutputStore.StorageError, ) const infraError = @@ -365,8 +373,8 @@ const layer = Layer.effect( if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause) if (userDeclined) return yield* Effect.interrupt - if ((toolsInterrupted || infraError !== undefined) && settledFailure) - return yield* Effect.failCause(settledFailure) + if (infraError !== undefined && infrastructureFailure) return yield* Effect.failCause(infrastructureFailure) + if (toolsInterrupted && settledFailure) return yield* Effect.failCause(settledFailure) if (toolsInterrupted && settled._tag === "Failure") return yield* Effect.failCause(settled.cause) if (stepFailure) return yield* new StepFailedError({ error: stepFailure }) return { diff --git a/packages/core/src/tool-output-store.ts b/packages/core/src/tool-output-store.ts index c0cb912ce0..d1fe89362b 100644 --- a/packages/core/src/tool-output-store.ts +++ b/packages/core/src/tool-output-store.ts @@ -17,6 +17,11 @@ export const RETENTION = Duration.days(7) export const MANAGED_DIRECTORY = "tool-output" +export interface Limits { + readonly maxLines: number + readonly maxBytes: number +} + export interface BoundInput { readonly sessionID: SessionSchema.ID readonly callID: string @@ -42,7 +47,7 @@ export class StorageError extends Schema.TaggedErrorClass()("ToolO export type Error = StorageError export interface Interface { - readonly limits: () => Effect.Effect<{ readonly maxLines: number; readonly maxBytes: number }> + readonly limits: () => Effect.Effect readonly bound: (input: BoundInput) => Effect.Effect readonly cleanup: () => Effect.Effect } @@ -116,6 +121,27 @@ const lineCount = (text: string) => { return count } +export const contextualOverflow = (output: ToolOutput, limits: Limits) => { + const contextual = output.content + .filter((item) => item.type === "text") + .map((item) => item.text) + .join("") + return Buffer.byteLength(contextual, "utf-8") > limits.maxBytes || lineCount(contextual) > limits.maxLines +} + +export const propagateTruncation = (output: ToolOutput, limits: Limits): ToolOutput => + contextualOverflow(output, limits) && + Predicate.isObject(output.structured) && + "truncated" in output.structured && + typeof output.structured.truncated === "boolean" + ? { ...output, structured: { ...output.structured, truncated: true } } + : output + +const record = (value: unknown): Record => + typeof value === "object" && value !== null && !Array.isArray(value) + ? Object.fromEntries(Object.entries(value)) + : { value } + const layer = Layer.effect( Service, Effect.gen(function* () { @@ -146,41 +172,49 @@ const layer = Layer.effect( const media = input.output.content.filter((item) => item.type === "file") const text = input.output.content.filter((item) => item.type === "text") const encoded = yield* Effect.try({ - try: () => JSON.stringify(input.output.structured) ?? String(input.output.structured), + try: () => JSON.stringify(record(input.output.structured)), catch: (cause) => new StorageError({ operation: "encode", cause }), }) + const decoded: unknown = JSON.parse(encoded) + const structured = + typeof decoded === "object" && decoded !== null && !Array.isArray(decoded) + ? Object.fromEntries(Object.entries(decoded)) + : yield* Effect.die(new Error("Durable tool structured output must be an object")) const encodedBytes = Buffer.byteLength(encoded, "utf-8") const contextual = input.output.content.length === 0 ? encoded : text.map((item) => item.text).join("") const contextualBytes = input.output.content.length === 0 ? encodedBytes : Buffer.byteLength(contextual, "utf-8") const structuredOverflow = encodedBytes > MAX_STRUCTURED_BYTES - const contextualOverflow = - contextualBytes > outputLimits.maxBytes || lineCount(contextual) > outputLimits.maxLines - if (!structuredOverflow && !contextualOverflow) + const contentOverflow = contextualBytes > outputLimits.maxBytes || lineCount(contextual) > outputLimits.maxLines + if (!structuredOverflow && !contentOverflow) return { - output: input.output, + output: { ...input.output, structured }, outputPaths: [], } yield* fs.ensureDir(directory).pipe(Effect.mapError((cause) => new StorageError({ operation: "write", cause }))) const structuredPath = structuredOverflow ? yield* write(encoded) : undefined - const contextualPath = contextualOverflow + const contextualPath = contentOverflow ? input.output.content.length === 0 && structuredPath ? structuredPath - : yield* write(contextual) + : yield* write(contextual).pipe( + Effect.onError(() => + structuredPath ? fs.remove(structuredPath).pipe(Effect.catch(() => Effect.void)) : Effect.void, + ), + ) : undefined return { output: { structured: structuredPath ? { _truncated: true, _bytes: encodedBytes, _outputPath: structuredPath } - : contextualOverflow && + : contentOverflow && input.propagateTruncation === true && - Predicate.isObject(input.output.structured) && - "truncated" in input.output.structured && - typeof input.output.structured.truncated === "boolean" - ? { ...input.output.structured, truncated: true } - : input.output.structured, - content: contextualOverflow + Predicate.isObject(structured) && + "truncated" in structured && + typeof structured.truncated === "boolean" + ? { ...structured, truncated: true } + : structured, + content: contentOverflow ? [ { type: "text" as const, diff --git a/packages/core/src/tool/edit.ts b/packages/core/src/tool/edit.ts index 69b1dc592f..9856d3a797 100644 --- a/packages/core/src/tool/edit.ts +++ b/packages/core/src/tool/edit.ts @@ -110,7 +110,7 @@ export const Plugin = { ...output, files: output.files.map((file) => ({ ...file, - patch: ToolStructured.truncate(file.patch, maximumBytes), + patch: ToolStructured.patch(file.patch, maximumBytes), })), })), toModelOutput: ({ input, output }) => [ diff --git a/packages/core/src/tool/patch.ts b/packages/core/src/tool/patch.ts index 399a0029d5..7c939ff7aa 100644 --- a/packages/core/src/tool/patch.ts +++ b/packages/core/src/tool/patch.ts @@ -79,7 +79,7 @@ export const Plugin = { ...output, files: output.files.map((file) => ({ ...file, - patch: ToolStructured.truncate(file.patch, maximumBytes), + patch: ToolStructured.patch(file.patch, maximumBytes), })), })), toModelOutput: ({ output }) => [{ type: "text", text: toModelOutput(output) }], diff --git a/packages/core/src/tool/read.ts b/packages/core/src/tool/read.ts index eff11b206d..6715d9678f 100644 --- a/packages/core/src/tool/read.ts +++ b/packages/core/src/tool/read.ts @@ -75,6 +75,7 @@ export const Plugin = { output: Output, structured: StructuredOutput, codeModeOutput: "output", + contentTruncation: true, toStructuredOutput: ({ output }) => { if ("encoding" in output) return { diff --git a/packages/core/src/tool/registry.ts b/packages/core/src/tool/registry.ts index ec341fab31..06ce3cc0e1 100644 --- a/packages/core/src/tool/registry.ts +++ b/packages/core/src/tool/registry.ts @@ -83,18 +83,16 @@ const registryLayer = Layer.effect( const base64 = /^data:[^,]*;base64,(.*)$/s.exec(item.uri)?.[1] if (base64 === undefined) return Effect.succeed(item) const resource = item.name ?? `${item.mime} tool output` - return image - .normalize(resource, { uri: resource, content: base64, encoding: "base64", mime: item.mime }) - .pipe( - Effect.map((result) => ({ - ...item, - uri: `data:${result.mime};base64,${result.content}`, - mime: result.mime, - })), - Effect.catchTag("Image.ResizerUnavailableError", () => Effect.succeed(item)), - Effect.catchTag("Image.DecodeError", () => Effect.succeed("decode" as const)), - Effect.catchTag("Image.SizeError", () => Effect.succeed("size" as const)), - ) + return image.normalize(resource, { uri: resource, content: base64, encoding: "base64", mime: item.mime }).pipe( + Effect.map((result) => ({ + ...item, + uri: `data:${result.mime};base64,${result.content}`, + mime: result.mime, + })), + Effect.catchTag("Image.ResizerUnavailableError", () => Effect.succeed(item)), + Effect.catchTag("Image.DecodeError", () => Effect.succeed("decode" as const)), + Effect.catchTag("Image.SizeError", () => Effect.succeed("size" as const)), + ) }) const note = (reason: "decode" | "size", text: string) => { const count = normalized.filter((item) => item === reason).length @@ -150,7 +148,28 @@ const registryLayer = Layer.effect( name: part.name, }, ), - ).pipe(Effect.flatMap((content) => progress({ structured: update.structured, content }))) + ).pipe( + Effect.flatMap((content) => + propagateTruncation + ? resources + .limits() + .pipe( + Effect.map((limits) => + ToolOutputStore.propagateTruncation({ structured: update.structured, content }, limits), + ), + ) + : Effect.succeed({ structured: update.structured, content }), + ), + Effect.flatMap((output) => { + const structured = output.structured + const entries = + typeof structured === "object" && structured !== null && !Array.isArray(structured) + ? Object.entries(structured) + : undefined + if (!entries) return Effect.die(new Error("Tool progress structured output must be an object")) + return progress({ structured: Object.fromEntries(entries), content: output.content }) + }), + ) }, }, ).pipe( @@ -166,22 +185,13 @@ const registryLayer = Layer.effect( if ("result" in pending) { settlement = pending } else { - const bounded = yield* resources.bound({ - sessionID: input.sessionID, - callID: input.call.id, - output: { structured: pending.output.structured, content: yield* normalizeImages(pending.output.content) }, - propagateTruncation, - }) - const result = ToolOutput.toResultValue(bounded.output) - settlement = - result.type === "error" - ? bounded.outputPaths.length > 0 - ? { result, outputPaths: bounded.outputPaths } - : { result } - : bounded.outputPaths.length > 0 - ? { result, output: bounded.output, outputPaths: bounded.outputPaths } - : { result, output: bounded.output } + const output = { + structured: pending.output.structured, + content: yield* normalizeImages(pending.output.content), + } + settlement = { result: ToolOutput.toResultValue(output), output } } + const initialResult = settlement.result const afterEvent: ToolHooks.AfterEvent = { tool: input.call.name, sessionID: input.sessionID, @@ -199,13 +209,15 @@ const registryLayer = Layer.effect( sessionID: input.sessionID, callID: input.call.id, output: afterEvent.output, + propagateTruncation, }) : undefined - const outputPaths = [ - ...new Set([...(afterEvent.outputPaths ?? []), ...(finalOutput?.outputPaths ?? [])]), - ] + const outputPaths = [...new Set([...(afterEvent.outputPaths ?? []), ...(finalOutput?.outputPaths ?? [])])] return { - result: afterEvent.result, + result: + finalOutput !== undefined && afterEvent.result === initialResult && !("result" in pending) + ? ToolOutput.toResultValue(finalOutput.output) + : afterEvent.result, ...(finalOutput !== undefined ? { output: finalOutput.output } : {}), ...(outputPaths.length > 0 ? { outputPaths } : {}), ...(settlement.error !== undefined ? { error: settlement.error } : {}), diff --git a/packages/core/src/tool/shell.ts b/packages/core/src/tool/shell.ts index 89e9ea381c..2e30e06ddf 100644 --- a/packages/core/src/tool/shell.ts +++ b/packages/core/src/tool/shell.ts @@ -149,6 +149,7 @@ export const Plugin = { output: Output, structured: StructuredOutput, codeModeOutput: "output", + contentTruncation: true, toStructuredOutput: ({ output }) => ({ truncated: output.truncated, ...(output.exit === undefined ? {} : { exit: output.exit }), diff --git a/packages/core/src/tool/structured.ts b/packages/core/src/tool/structured.ts index 2d34610ecb..50a691fc1b 100644 --- a/packages/core/src/tool/structured.ts +++ b/packages/core/src/tool/structured.ts @@ -1,5 +1,6 @@ export * as ToolStructured from "./structured" +import { formatPatch, parsePatch } from "diff" import { ToolOutputStore } from "../tool-output-store" export function fit(project: (maxStringBytes: number) => A) { @@ -37,3 +38,26 @@ export function truncate(input: string, maximumBytes: number) { } return input.slice(0, end).replace(/\r?\n$/, "") + (maximumBytes >= markerBytes ? marker : "") } + +export function patch(input: string, maximumBytes: number) { + if (Buffer.byteLength(input, "utf-8") <= maximumBytes) return input + const parsed = parsePatch(input)[0] + const hunk = parsed?.hunks[0] + if (!parsed || !hunk) return truncate(input, maximumBytes) + const changed = [hunk.lines.find((line) => line.startsWith("-")), hunk.lines.find((line) => line.startsWith("+"))] + .filter((line) => line !== undefined) + .map((line) => line[0] + truncate(line.slice(1), Math.max(0, Math.floor(maximumBytes / 2) - 1))) + const lines = [...changed, " ... diff truncated ..."] + return formatPatch({ + ...parsed, + hunks: [ + { + oldStart: hunk.oldStart, + oldLines: lines.filter((line) => line.startsWith("-") || line.startsWith(" ")).length, + newStart: hunk.newStart, + newLines: lines.filter((line) => line.startsWith("+") || line.startsWith(" ")).length, + lines, + }, + ], + }) +} diff --git a/packages/core/test/session-runner-tool-registry.test.ts b/packages/core/test/session-runner-tool-registry.test.ts index 8e9060820e..b5010d5dcb 100644 --- a/packages/core/test/session-runner-tool-registry.test.ts +++ b/packages/core/test/session-runner-tool-registry.test.ts @@ -101,14 +101,17 @@ describe("ToolRegistry", () => { return Effect.gen(function* () { const service = yield* ToolRegistry.Service const text = "x".repeat(ToolOutputStore.MAX_STRUCTURED_BYTES) - yield* service.register({ - structured_only: Tool.make({ - description: "Return structured output", - input: Schema.Struct({}), - output: Schema.Struct({ text: Schema.String }), - execute: () => Effect.succeed({ text }), - }), - }, { codemode: false }) + yield* service.register( + { + structured_only: Tool.make({ + description: "Return structured output", + input: Schema.Struct({}), + output: Schema.Struct({ text: Schema.String }), + execute: () => Effect.succeed({ text }), + }), + }, + { codemode: false }, + ) const settled = yield* settleTool(service, { sessionID, ...identity, @@ -173,6 +176,42 @@ describe("ToolRegistry", () => { ), ) + live.live("allows execute-after hooks to redact output before retention", () => + Effect.acquireUseRelease( + Effect.promise(() => tmpdir()), + (tmp) => { + const layer = AppNodeBuilder.build(LayerNode.group([ToolRegistry.node, ToolHooks.node]), [ + [Global.node, Global.layerWith({ data: tmp.path })], + [Image.node, imageStore], + [ToolOutputStore.node, ToolOutputStore.nodeWithoutConfig], + ]) + return Effect.gen(function* () { + const service = yield* ToolRegistry.Service + const hooks = yield* ToolHooks.Service + const text = "sensitive".repeat(ToolOutputStore.MAX_STRUCTURED_BYTES) + let observed: unknown + yield* hooks.hook.after((event) => { + observed = event.output?.structured + event.output = { structured: { redacted: true }, content: [{ type: "text", text: "redacted" }] } + }) + yield* service.register({ secret: constant(text) }, { codemode: false }) + const settled = yield* settleTool(service, { + sessionID, + ...identity, + call: { type: "tool-call", id: "call-secret", name: "secret", input: { text: "ignored" } }, + }) + + expect(observed).toEqual({ text }) + expect(settled).toEqual({ + result: { type: "text", value: "redacted" }, + output: { structured: { redacted: true }, content: [{ type: "text", text: "redacted" }] }, + }) + }).pipe(Effect.provide(layer)) + }, + (tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]()), + ), + ) + it.effect("rejects invalid dotted namespaces", () => Effect.gen(function* () { const service = yield* ToolRegistry.Service @@ -202,12 +241,15 @@ describe("ToolRegistry", () => { it.effect("filters disabled tools with edit aliases and ordered wildcard precedence", () => Effect.gen(function* () { const service = yield* ToolRegistry.Service - yield* service.register({ - question: make(), - bash: make(), - edit: make("edit"), - write: make("edit"), - }, { codemode: false }) + yield* service.register( + { + question: make(), + bash: make(), + edit: make("edit"), + write: make("edit"), + }, + { codemode: false }, + ) const names = (permissions: PermissionV2.Ruleset) => toolDefinitions(service, permissions).pipe(Effect.map((definitions) => definitions.map((tool) => tool.name))) @@ -280,14 +322,17 @@ describe("ToolRegistry", () => { it.effect("returns model errors without swallowing interruption or defects", () => Effect.gen(function* () { const service = yield* ToolRegistry.Service - yield* service.register({ - failed: Tool.make({ - description: "Failed", - input: Schema.Struct({}), - output: Schema.Struct({ ok: Schema.Boolean }), - execute: () => Effect.fail(new Tool.Failure({ message: "Denied" })), - }), - }, { codemode: false }) + yield* service.register( + { + failed: Tool.make({ + description: "Failed", + input: Schema.Struct({}), + output: Schema.Struct({ ok: Schema.Boolean }), + execute: () => Effect.fail(new Tool.Failure({ message: "Denied" })), + }), + }, + { codemode: false }, + ) expect( yield* executeTool(service, { sessionID, @@ -303,14 +348,17 @@ describe("ToolRegistry", () => { }), ).toEqual({ type: "error", value: "Unknown tool: missing" }) - yield* service.register({ - defect: Tool.make({ - description: "Defect", - input: Schema.Struct({}), - output: Schema.Struct({}), - execute: () => Effect.die("unexpected executor defect"), - }), - }, { codemode: false }) + yield* service.register( + { + defect: Tool.make({ + description: "Defect", + input: Schema.Struct({}), + output: Schema.Struct({}), + execute: () => Effect.die("unexpected executor defect"), + }), + }, + { codemode: false }, + ) expect( yield* service.materialize().pipe( Effect.flatMap((materialized) => @@ -353,22 +401,23 @@ describe("ToolRegistry", () => { Effect.gen(function* () { const service = yield* ToolRegistry.Service const contexts: Tool.Context[] = [] - yield* service.register({ - context: Tool.make({ - description: "Context", - input: Schema.Struct({}), - output: Schema.Struct({ ok: Schema.Boolean }), - execute: (_, context) => Effect.sync(() => contexts.push(context)).pipe(Effect.as({ ok: true })), - }), - }, { codemode: false }) + yield* service.register( + { + context: Tool.make({ + description: "Context", + input: Schema.Struct({}), + output: Schema.Struct({ ok: Schema.Boolean }), + execute: (_, context) => Effect.sync(() => contexts.push(context)).pipe(Effect.as({ ok: true })), + }), + }, + { codemode: false }, + ) yield* executeTool(service, { sessionID, ...identity, call: { type: "tool-call", id: "call-context", name: "context", input: {} }, }) - expect(contexts).toEqual([ - { sessionID, ...identity, callID: "call-context", progress: expect.any(Function) }, - ]) + expect(contexts).toEqual([{ sessionID, ...identity, callID: "call-context", progress: expect.any(Function) }]) }), ) @@ -388,27 +437,30 @@ describe("ToolRegistry", () => { output: { structured: {}, content: [{ type: "text", text: "bounded reference" }] }, outputPaths: ["/managed/generic"], }) - expect(bounds).toHaveLength(2) + expect(bounds).toHaveLength(1) }), ) it.effect("normalizes image tool output at settlement and drops unresizable images", () => Effect.gen(function* () { const service = yield* ToolRegistry.Service - yield* service.register({ - snapshot: Tool.make({ - description: "Return images", - input: Schema.Struct({ text: Schema.String }), - output: Schema.Struct({ text: Schema.String }), - execute: ({ text }) => Effect.succeed({ text }), - toModelOutput: ({ output }) => [ - { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "frame.png" }, - { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "too-large.png" }, - { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "corrupt.png" }, - { type: "text", text: output.text }, - ], - }), - }, { codemode: false }) + yield* service.register( + { + snapshot: Tool.make({ + description: "Return images", + input: Schema.Struct({ text: Schema.String }), + output: Schema.Struct({ text: Schema.String }), + execute: ({ text }) => Effect.succeed({ text }), + toModelOutput: ({ output }) => [ + { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "frame.png" }, + { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "too-large.png" }, + { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "corrupt.png" }, + { type: "text", text: output.text }, + ], + }), + }, + { codemode: false }, + ) const settlement = yield* settleTool(service, call("snapshot")) expect(settlement.output?.content).toEqual([ @@ -423,23 +475,26 @@ describe("ToolRegistry", () => { it.effect("normalizes image progress content before it is published", () => Effect.gen(function* () { const service = yield* ToolRegistry.Service - yield* service.register({ - progressive: Tool.make({ - description: "Emit image progress", - input: Schema.Struct({ text: Schema.String }), - output: Schema.Struct({ text: Schema.String }), - execute: ({ text }, context) => - context - .progress({ - structured: { stage: "capture" }, - content: [ - { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "frame.png" }, - { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "too-large.png" }, - ], - }) - .pipe(Effect.as({ text })), - }), - }, { codemode: false }) + yield* service.register( + { + progressive: Tool.make({ + description: "Emit image progress", + input: Schema.Struct({ text: Schema.String }), + output: Schema.Struct({ text: Schema.String }), + execute: ({ text }, context) => + context + .progress({ + structured: { stage: "capture" }, + content: [ + { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "frame.png" }, + { type: "file", data: "aW1hZ2U=", mime: "image/png", name: "too-large.png" }, + ], + }) + .pipe(Effect.as({ text })), + }), + }, + { codemode: false }, + ) const updates: ToolRegistry.Progress[] = [] yield* settleTool(service, { @@ -471,15 +526,18 @@ describe("ToolRegistry", () => { encode: SchemaGetter.transform((value) => value === "yes"), }), ) - yield* service.register({ - transformed: Tool.make({ - description: "Transform values", - input: Schema.Struct({ value: Transformed }), - output: Schema.Struct({ value: Transformed }), - execute: ({ value }) => Effect.sync(() => executed.push(value)).pipe(Effect.as({ value })), - toModelOutput: ({ output }) => [{ type: "text", text: String(output.value) }], - }), - }, { codemode: false }) + yield* service.register( + { + transformed: Tool.make({ + description: "Transform values", + input: Schema.Struct({ value: Transformed }), + output: Schema.Struct({ value: Transformed }), + execute: ({ value }) => Effect.sync(() => executed.push(value)).pipe(Effect.as({ value })), + toModelOutput: ({ output }) => [{ type: "text", text: String(output.value) }], + }), + }, + { codemode: false }, + ) expect( yield* executeTool(service, { @@ -498,25 +556,28 @@ describe("ToolRegistry", () => { ).toMatchObject({ type: "error", value: expect.stringContaining("Invalid tool input") }) expect(executed).toEqual(["yes"]) - yield* service.register({ - invalid_output: Tool.make({ - description: "Return invalid output", - input: Schema.Struct({}), - output: Schema.Struct({ - value: Schema.Boolean.pipe( - Schema.decodeTo(Schema.String, { - decode: SchemaGetter.transform((value) => String(value)), - encode: SchemaGetter.transformOrFail((value) => - value === "valid" - ? Effect.succeed(true) - : Effect.fail(new SchemaIssue.InvalidValue(Option.some(value), { message: "invalid output" })), - ), - }), - ), + yield* service.register( + { + invalid_output: Tool.make({ + description: "Return invalid output", + input: Schema.Struct({}), + output: Schema.Struct({ + value: Schema.Boolean.pipe( + Schema.decodeTo(Schema.String, { + decode: SchemaGetter.transform((value) => String(value)), + encode: SchemaGetter.transformOrFail((value) => + value === "valid" + ? Effect.succeed(true) + : Effect.fail(new SchemaIssue.InvalidValue(Option.some(value), { message: "invalid output" })), + ), + }), + ), + }), + execute: () => Effect.succeed({ value: "invalid" }), }), - execute: () => Effect.succeed({ value: "invalid" }), - }), - }, { codemode: false }) + }, + { codemode: false }, + ) expect( yield* executeTool(service, { sessionID, diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index b42b5c2053..2499484dc9 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -239,43 +239,46 @@ const permission = Layer.succeed( ) const echo = Layer.effectDiscard( ToolRegistry.Service.use((registry) => - registry.register({ - echo: Tool.make({ - description: "Echo text", - input: Schema.Struct({ text: Schema.String }), - output: Schema.Struct({ text: Schema.String }), - toModelOutput: ({ output }) => [{ type: "text", text: output.text }], - execute: ({ text }, context) => - Effect.gen(function* () { - authorizations.push(context) - executions.push(text) - activeToolExecutions++ - maxActiveToolExecutions = Math.max(maxActiveToolExecutions, activeToolExecutions) - if (activeToolExecutions === toolExecutionsReady && toolExecutionsStarted) { - yield* Deferred.succeed(toolExecutionsStarted, undefined) - } - if (toolExecutionGate) yield* Deferred.await(toolExecutionGate) - return { text } - }).pipe(Effect.ensuring(Effect.sync(() => activeToolExecutions--))), - }), - defect: Tool.make({ - description: "Fail unexpectedly", - input: Schema.Struct({}), - output: Schema.Struct({}), - execute: () => - (toolExecutionGate ? Deferred.await(toolExecutionGate) : Effect.void).pipe( - Effect.andThen(Effect.die("unexpected tool defect")), - ), - }), - // BigInt output with no model content forces ToolOutputStore.bound onto its - // JSON.stringify encode path, which fails with a typed StorageError. - storefail: Tool.make({ - description: "Produce output that cannot be persisted", - input: Schema.Struct({}), - output: Schema.Any, - execute: () => Effect.succeed({ big: 1n }), - }), - }, { codemode: false }), + registry.register( + { + echo: Tool.make({ + description: "Echo text", + input: Schema.Struct({ text: Schema.String }), + output: Schema.Struct({ text: Schema.String }), + toModelOutput: ({ output }) => [{ type: "text", text: output.text }], + execute: ({ text }, context) => + Effect.gen(function* () { + authorizations.push(context) + executions.push(text) + activeToolExecutions++ + maxActiveToolExecutions = Math.max(maxActiveToolExecutions, activeToolExecutions) + if (activeToolExecutions === toolExecutionsReady && toolExecutionsStarted) { + yield* Deferred.succeed(toolExecutionsStarted, undefined) + } + if (toolExecutionGate) yield* Deferred.await(toolExecutionGate) + return { text } + }).pipe(Effect.ensuring(Effect.sync(() => activeToolExecutions--))), + }), + defect: Tool.make({ + description: "Fail unexpectedly", + input: Schema.Struct({}), + output: Schema.Struct({}), + execute: () => + (toolExecutionGate ? Deferred.await(toolExecutionGate) : Effect.void).pipe( + Effect.andThen(Effect.die("unexpected tool defect")), + ), + }), + // BigInt output with no model content forces ToolOutputStore.bound onto its + // JSON.stringify encode path, which fails with a typed StorageError. + storefail: Tool.make({ + description: "Produce output that cannot be persisted", + input: Schema.Struct({}), + output: Schema.Any, + execute: () => Effect.succeed({ big: 1n }), + }), + }, + { codemode: false }, + ), ), ) const echoNode = makeLocationNode({ name: "test/session-runner-tools", layer: echo, deps: [ToolRegistry.node] }) @@ -841,19 +844,22 @@ describe("SessionRunnerLLM", () => { const registry = yield* ToolRegistry.Service const contexts: Tool.Context[] = [] const progress = "x".repeat(ToolOutputStore.MAX_STRUCTURED_BYTES) - yield* registry.register({ - location_context: Tool.make({ - description: "Read application context", - input: Schema.Struct({ query: Schema.String }), - output: Schema.Struct({ answer: Schema.String }), - execute: ({ query }, context) => - Effect.gen(function* () { - contexts.push(context) - yield* context.progress({ structured: { phase: progress } }) - return { answer: query.toUpperCase() } - }), - }), - }, { codemode: false }) + yield* registry.register( + { + location_context: Tool.make({ + description: "Read application context", + input: Schema.Struct({ query: Schema.String }), + output: Schema.Struct({ answer: Schema.String }), + execute: ({ query }, context) => + Effect.gen(function* () { + contexts.push(context) + yield* context.progress({ structured: { phase: progress } }) + return { answer: query.toUpperCase() } + }), + }), + }, + { codemode: false }, + ) yield* admit(session, "Use application context") responses = [reply.tool("call-location", "location_context", { query: "hello" }), []] const events = yield* EventV2.Service @@ -906,14 +912,17 @@ describe("SessionRunnerLLM", () => { const scope = yield* Scope.make() const executions: string[] = [] yield* registry - .register({ - reloaded: Tool.make({ - description: "Record the advertised tool", - input: Schema.Struct({}), - output: Schema.Struct({ value: Schema.String }), - execute: () => Effect.sync(() => executions.push("advertised")).pipe(Effect.as({ value: "advertised" })), - }), - }, { codemode: false }) + .register( + { + reloaded: Tool.make({ + description: "Record the advertised tool", + input: Schema.Struct({}), + output: Schema.Struct({ value: Schema.String }), + execute: () => Effect.sync(() => executions.push("advertised")).pipe(Effect.as({ value: "advertised" })), + }), + }, + { codemode: false }, + ) .pipe(Scope.provide(scope)) yield* admit(session, "Use the reloaded tool") responses = [ @@ -931,14 +940,17 @@ describe("SessionRunnerLLM", () => { const run = yield* session.resume(sessionID).pipe(Effect.forkChild) yield* Deferred.await(streamStarted) yield* Scope.close(scope, Exit.void) - yield* registry.register({ - reloaded: Tool.make({ - description: "Record the replacement tool", - input: Schema.Struct({}), - output: Schema.Struct({ value: Schema.String }), - execute: () => Effect.sync(() => executions.push("replacement")).pipe(Effect.as({ value: "replacement" })), - }), - }, { codemode: false }) + yield* registry.register( + { + reloaded: Tool.make({ + description: "Record the replacement tool", + input: Schema.Struct({}), + output: Schema.Struct({ value: Schema.String }), + execute: () => Effect.sync(() => executions.push("replacement")).pipe(Effect.as({ value: "replacement" })), + }), + }, + { codemode: false }, + ) yield* Deferred.succeed(streamGate, undefined) yield* Fiber.join(run) @@ -3374,17 +3386,20 @@ describe("SessionRunnerLLM", () => { Effect.gen(function* () { const session = yield* setup const registry = yield* ToolRegistry.Service - yield* registry.register({ - blocked: Tool.make({ - description: "Fail because policy blocked execution", - input: Schema.Struct({}), - output: Schema.Struct({}), - execute: () => - Effect.fail(new PermissionV2.BlockedError({ rules: [], permission: "blocked", resources: ["*"] })).pipe( - Effect.mapError(() => new Tool.Failure({ message: "Permission blocked" })), - ), - }), - }, { codemode: false }) + yield* registry.register( + { + blocked: Tool.make({ + description: "Fail because policy blocked execution", + input: Schema.Struct({}), + output: Schema.Struct({}), + execute: () => + Effect.fail(new PermissionV2.BlockedError({ rules: [], permission: "blocked", resources: ["*"] })).pipe( + Effect.mapError(() => new Tool.Failure({ message: "Permission blocked" })), + ), + }), + }, + { codemode: false }, + ) yield* admit(session, "Call blocked") responses = [reply.tool("call-blocked", "blocked", {}), reply.stop()] @@ -3409,14 +3424,17 @@ describe("SessionRunnerLLM", () => { Effect.gen(function* () { const session = yield* setup const registry = yield* ToolRegistry.Service - yield* registry.register({ - declined: Tool.make({ - description: "Fail because the user declined approval", - input: Schema.Struct({}), - output: Schema.Struct({}), - execute: () => Effect.die(new PermissionV2.DeclinedError()), - }), - }, { codemode: false }) + yield* registry.register( + { + declined: Tool.make({ + description: "Fail because the user declined approval", + input: Schema.Struct({}), + output: Schema.Struct({}), + execute: () => Effect.die(new PermissionV2.DeclinedError()), + }), + }, + { codemode: false }, + ) yield* admit(session, "Call declined") response = reply.tool("call-declined", "declined", {}) @@ -3446,17 +3464,20 @@ describe("SessionRunnerLLM", () => { Effect.gen(function* () { const session = yield* setup const registry = yield* ToolRegistry.Service - yield* registry.register({ - corrected: Tool.make({ - description: "Fail with user correction feedback", - input: Schema.Struct({}), - output: Schema.Struct({}), - execute: () => - Effect.fail(new PermissionV2.CorrectedError({ feedback: "Use another tool" })).pipe( - Effect.mapError(() => new Tool.Failure({ message: "Use another tool" })), - ), - }), - }, { codemode: false }) + yield* registry.register( + { + corrected: Tool.make({ + description: "Fail with user correction feedback", + input: Schema.Struct({}), + output: Schema.Struct({}), + execute: () => + Effect.fail(new PermissionV2.CorrectedError({ feedback: "Use another tool" })).pipe( + Effect.mapError(() => new Tool.Failure({ message: "Use another tool" })), + ), + }), + }, + { codemode: false }, + ) yield* admit(session, "Call corrected") responses = [reply.tool("call-corrected", "corrected", {}), reply.stop()] @@ -3512,6 +3533,36 @@ describe("SessionRunnerLLM", () => { }), ) + it.effect("does not hide output persistence failure behind another concurrent tool defect", () => + Effect.gen(function* () { + const session = yield* setup + yield* admit(session, "Call defect and storefail") + responses = [ + [ + LLMEvent.stepStart({ index: 0 }), + LLMEvent.toolCall({ id: "call-defect", name: "defect", input: {} }), + LLMEvent.toolCall({ id: "call-storefail", name: "storefail", input: {} }), + LLMEvent.stepFinish({ index: 0, reason: "tool-calls" }), + LLMEvent.finish({ reason: "tool-calls" }), + ], + [], + ] + + const exit = yield* session.resume(sessionID).pipe(Effect.exit) + + expect(Exit.isFailure(exit)).toBe(true) + expect(requests).toHaveLength(1) + expect(yield* session.context(sessionID)).toMatchObject([ + { type: "user", text: "Call defect and storefail" }, + { + type: "assistant", + finish: "error", + error: { type: "unknown", message: expect.stringContaining("Failed to encode tool output") }, + }, + ]) + }), + ) + it.effect("returns configured permission denials to the model and continues", () => Effect.gen(function* () { const session = yield* setup @@ -3554,14 +3605,17 @@ describe("SessionRunnerLLM", () => { Effect.gen(function* () { const session = yield* setup const registry = yield* ToolRegistry.Service - yield* registry.register({ - question: Tool.make({ - description: "Ask the user", - input: Schema.Struct({}), - output: Schema.Struct({}), - execute: () => Effect.die(new QuestionTool.CancelledError()), - }), - }, { codemode: false }) + yield* registry.register( + { + question: Tool.make({ + description: "Ask the user", + input: Schema.Struct({}), + output: Schema.Struct({}), + execute: () => Effect.die(new QuestionTool.CancelledError()), + }), + }, + { codemode: false }, + ) yield* admit(session, "Ask then stop") responses = [reply.tool("call-question", "question", {}), []] diff --git a/packages/core/test/tool-edit.test.ts b/packages/core/test/tool-edit.test.ts index cd8b0c6d47..47182a59ce 100644 --- a/packages/core/test/tool-edit.test.ts +++ b/packages/core/test/tool-edit.test.ts @@ -181,8 +181,8 @@ describe("EditTool", () => { (tmp) => { reset() const target = path.join(tmp.path, "large.txt") - const before = "x".repeat(9_000) - const after = "y".repeat(9_000) + const before = "x".repeat(20_000) + const after = "y".repeat(20_000) return Effect.promise(() => fs.writeFile(target, before)).pipe( Effect.andThen( withTool(tmp.path, (registry) => @@ -195,7 +195,10 @@ describe("EditTool", () => { expect(Buffer.byteLength(JSON.stringify(structured))).toBeLessThanOrEqual( ToolOutputStore.MAX_STRUCTURED_BYTES, ) - expect(() => parsePatch(structured.files[0]?.patch ?? "")).not.toThrow() + const hunk = parsePatch(structured.files[0]?.patch ?? "")[0]?.hunks[0] + expect(hunk?.lines.some((line) => line.startsWith("-"))).toBe(true) + expect(hunk?.lines.some((line) => line.startsWith("+"))).toBe(true) + expect(hunk).toMatchObject({ oldLines: 2, newLines: 2 }) expect(structured).toMatchObject({ replacements: 1, files: [ diff --git a/packages/core/test/tool-output-store.test.ts b/packages/core/test/tool-output-store.test.ts index 2193ad1cfe..1b67c07c79 100644 --- a/packages/core/test/tool-output-store.test.ts +++ b/packages/core/test/tool-output-store.test.ts @@ -168,6 +168,53 @@ describe("ToolOutputStore", () => { ), ) + it.live("normalizes structured values to their durable JSON record", () => + withStore(({ store }) => + Effect.gen(function* () { + const primitive = yield* store.bound({ + sessionID, + callID: "call-primitive", + output: { structured: "value", content: [] }, + }) + const nonFinite = yield* store.bound({ + sessionID, + callID: "call-non-finite", + output: { structured: { value: Number.NaN }, content: [] }, + }) + const omitted = yield* store.bound({ + sessionID, + callID: "call-omitted", + output: { structured: undefined, content: [] }, + }) + + expect(primitive.output.structured).toEqual({ value: "value" }) + expect(nonFinite.output.structured).toEqual({ value: null }) + expect(omitted.output.structured).toEqual({}) + }), + ), + ) + + it.live("measures primitive overflow after durable record normalization", () => + withStore(({ store, fs }) => + Effect.gen(function* () { + const value = "x".repeat(ToolOutputStore.MAX_STRUCTURED_BYTES) + const encoded = JSON.stringify({ value }) + const result = yield* store.bound({ + sessionID, + callID: "call-primitive-overflow", + output: { structured: value, content: [] }, + }) + + expect(result.output.structured).toEqual({ + _truncated: true, + _bytes: Buffer.byteLength(encoded), + _outputPath: result.outputPaths[0], + }) + expect(yield* fs.readFileString(result.outputPaths[0])).toBe(encoded) + }), + ), + ) + it.live("preserves native media and structured metadata without applying a settlement media limit", () => withStore(({ store }) => Effect.gen(function* () { diff --git a/packages/core/test/tool-patch.test.ts b/packages/core/test/tool-patch.test.ts index 1f7f5b42b3..ca9e132bd5 100644 --- a/packages/core/test/tool-patch.test.ts +++ b/packages/core/test/tool-patch.test.ts @@ -224,8 +224,8 @@ describe("PatchTool", () => { (tmp) => { reset() const target = path.join(tmp.path, "large.txt") - const before = "x".repeat(9_000) - const after = "y".repeat(9_000) + const before = "x".repeat(20_000) + const after = "y".repeat(20_000) return Effect.promise(() => fs.writeFile(target, `${before}\n`)).pipe( Effect.andThen( withTool(tmp.path, (registry) => @@ -238,7 +238,10 @@ describe("PatchTool", () => { expect(Buffer.byteLength(JSON.stringify(structured))).toBeLessThanOrEqual( ToolOutputStore.MAX_STRUCTURED_BYTES, ) - expect(() => parsePatch(structured.files[0]?.patch ?? "")).not.toThrow() + const hunk = parsePatch(structured.files[0]?.patch ?? "")[0]?.hunks[0] + expect(hunk?.lines.some((line) => line.startsWith("-"))).toBe(true) + expect(hunk?.lines.some((line) => line.startsWith("+"))).toBe(true) + expect(hunk).toMatchObject({ oldLines: 2, newLines: 2 }) expect(structured).toMatchObject({ applied: [{ type: "update", resource: "large.txt" }], files: [{ file: "large.txt", patch: expect.stringContaining("... truncated ...") }], diff --git a/packages/core/test/tool-shell.test.ts b/packages/core/test/tool-shell.test.ts index 1e59e788e4..0c95878de2 100644 --- a/packages/core/test/tool-shell.test.ts +++ b/packages/core/test/tool-shell.test.ts @@ -223,6 +223,30 @@ describe("ShellTool", () => { ), ) + it.live("marks output truncated when generic settlement limits are lower than shell capture limits", () => + Effect.acquireUseRelease( + Effect.promise(() => tmpdir()), + (tmp) => { + reset() + return withSession(tmp.path, (registry) => + Effect.gen(function* () { + const settled = yield* settleTool( + registry, + call({ command: overflowCommand(ToolOutputStore.MAX_BYTES + 1_000) }), + ) + + expect(settled.output?.structured).toMatchObject({ exit: 0, truncated: true }) + expect(settled.output?.content[0]).toMatchObject({ + type: "text", + text: expect.stringContaining("output truncated; full content saved to"), + }) + }), + ) + }, + (tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]().then(() => undefined)), + ), + ) + it.live("resolves a relative workdir from the active Location", () => Effect.acquireUseRelease( Effect.promise(() => tmpdir()), @@ -439,7 +463,10 @@ describe("ShellTool", () => { if (update.structured.truncated !== true) return const content = update.content[0] if (content?.type !== "text") return - if (content.text.indexOf("\n\n[output truncated; full output saved to:") !== ShellTool.MAX_CAPTURE_BYTES) + if ( + content.text.indexOf("\n\n[output truncated; full output saved to:") !== + ShellTool.MAX_CAPTURE_BYTES + ) return yield* Deferred.succeed(observed, update) yield* Effect.promise(() => fs.writeFile(releasePath, ""))