fix(ai): preserve provider tool input identity

This commit is contained in:
starptech 2026-07-19 00:54:08 +02:00
commit 0f26246bfd
21 changed files with 833 additions and 74 deletions

View file

@ -324,6 +324,8 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
type: "tool",
id: event.data.callID,
name: event.data.name,
executed: event.data.executed,
providerState: event.data.state,
time: { created: event.created },
state: SessionMessage.ToolStateStreaming.make({ status: "streaming", input: "" }),
}),
@ -396,7 +398,7 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
status: "error",
error: event.data.error,
input: typeof match.state.input === "string" ? {} : match.state.input,
raw: typeof match.state.input === "string" ? match.state.input : undefined,
raw: event.data.raw ?? (typeof match.state.input === "string" ? match.state.input : undefined),
structured: match.state.status === "running" ? match.state.structured : {},
content: match.state.status === "running" ? match.state.content : [],
result: event.data.result,

View file

@ -28,6 +28,8 @@ import { toSessionError } from "../to-session-error"
import { SessionRunnerRetry } from "./retry"
import { SessionUsage } from "../usage"
const MAX_MALFORMED_TOOL_RETRIES = 4
const layer = Layer.effect(
Service,
Effect.gen(function* () {
@ -80,6 +82,7 @@ const layer = Layer.effect(
sessionID: SessionSchema.ID,
promotion: SessionPending.Delivery | undefined,
step: number,
recoverMalformedToolInput: boolean,
recoverOverflow?: typeof compaction.compact,
assistantMessageID?: SessionMessage.ID,
) {
@ -143,7 +146,7 @@ const layer = Layer.effect(
}
}
yield* publish(event)
if (event.type !== "tool-call" || event.providerExecuted) return
if (event.type !== "tool-call" || publisher.isProviderExecuted(event.id)) return
const tool = prepared.resolveToolCall(event.name)
if (tool.type === "reject") {
yield* serialized(publisher.failUnsettledTools(tool.error))
@ -244,13 +247,35 @@ const layer = Layer.effect(
if (overflowFailure) yield* publish(overflowFailure)
const llmFailure = streamFailure instanceof LLMError ? streamFailure : undefined
const malformedToolInput =
llmFailure?.reason._tag === "InvalidProviderOutput" && llmFailure.reason.source === "tool-input"
const recoveredMalformedToolInput = malformedToolInput
llmFailure?.reason._tag === "InvalidProviderOutput" &&
llmFailure.reason.source === "tool-input" &&
llmFailure.reason.toolCallID !== undefined &&
llmFailure.reason.toolName !== undefined &&
llmFailure.reason.raw !== undefined
? {
id: llmFailure.reason.toolCallID,
name: llmFailure.reason.toolName,
raw: llmFailure.reason.raw,
providerExecuted: llmFailure.reason.providerExecuted,
providerMetadata: llmFailure.reason.providerMetadata,
}
: undefined
const failedMalformedToolInput = malformedToolInput
? yield* serialized(
publisher.failUnsettledTools({
type: "provider.invalid-output",
message: "Tool call arguments were malformed JSON and were not executed. Retry with valid JSON.",
}),
publisher.failMalformedToolInput(
{
id: malformedToolInput.id,
name: malformedToolInput.name,
raw: malformedToolInput.raw,
providerExecuted: malformedToolInput.providerExecuted,
providerMetadata: malformedToolInput.providerMetadata,
},
{
type: "provider.invalid-output",
message: "Tool call arguments were malformed JSON and were not executed. Retry with valid JSON.",
},
toSessionError(llmFailure),
),
)
: false
if (llmFailure && !publisher.hasProviderError()) {
@ -267,6 +292,11 @@ const layer = Layer.effect(
}
// Provider error events only arrive from the stream, so the flag is final here.
const providerFailed = publisher.hasProviderError()
const recoveredMalformedToolInput =
recoverMalformedToolInput &&
!providerFailed &&
!streamInterrupted &&
(failedMalformedToolInput || (stream._tag === "Success" && publisher.hasMalformedToolInput()))
// Settle every owned tool fiber. FiberSet.join returns on the first failure, so retain
// the individual fibers and await all exits before publishing the terminal step event.
@ -346,6 +376,7 @@ const layer = Layer.effect(
return {
_tag: "Completed",
needsContinuation: needsContinuation || recoveredMalformedToolInput,
malformedToolInput: recoveredMalformedToolInput,
step: currentStep,
} as const
}),
@ -356,6 +387,7 @@ const layer = Layer.effect(
sessionID: SessionSchema.ID,
promotion: SessionPending.Delivery | undefined,
step: number,
recoverMalformedToolInput: boolean,
) {
// Compaction restarts rebuild the request from compacted history without re-promoting.
// Overflow recovery is one-shot: a post-compaction attempt must not recover another
@ -366,7 +398,14 @@ const layer = Layer.effect(
let assistantMessageID: SessionMessage.ID | undefined
while (true) {
const attempt = yield* Effect.suspend(() =>
attemptStep(sessionID, currentPromotion, currentStep, recoverOverflow, assistantMessageID),
attemptStep(
sessionID,
currentPromotion,
currentStep,
recoverMalformedToolInput,
recoverOverflow,
assistantMessageID,
),
).pipe(
Effect.tapError((error) =>
error instanceof SessionRunnerRetry.RetryableFailure
@ -388,7 +427,12 @@ const layer = Layer.effect(
.pipe(Effect.andThen(Effect.fail(error.cause)))
}),
)
if (attempt._tag === "Completed") return { needsContinuation: attempt.needsContinuation, step: attempt.step }
if (attempt._tag === "Completed")
return {
needsContinuation: attempt.needsContinuation,
malformedToolInput: attempt.malformedToolInput,
step: attempt.step,
}
if (attempt._tag === "RestartAfterOverflowCompaction") recoverOverflow = undefined
yield* Effect.yieldNow
currentPromotion = undefined
@ -446,12 +490,19 @@ const layer = Layer.effect(
while (shouldRun) {
let needsContinuation = true
let step = 1
let malformedToolRetries = 0
// Repeat steps while continuation is needed. A step needs continuation only
// when it recorded local tool calls whose results the model has not yet seen;
// a provider error suppresses it. Pending steers also continue the loop so
// interjections are answered before the session goes idle.
// malformed tool input can also continue within its bounded recovery budget.
// A provider error suppresses continuation. Pending steers also continue the
// loop so interjections are answered before the session goes idle.
while (needsContinuation) {
const result = yield* runStep(input.sessionID, promotion, step)
const result = yield* runStep(
input.sessionID,
promotion,
step,
malformedToolRetries < MAX_MALFORMED_TOOL_RETRIES,
)
// Steer/queue promotion inside runStep has already made the pending input a visible
// user message by this point, so the first-user-message check below is reliable.
if (!titleAttempted.has(input.sessionID)) {
@ -459,6 +510,7 @@ const layer = Layer.effect(
forkTitle(title.generateForFirstPrompt(yield* getSession(input.sessionID)).pipe(Effect.ignore))
}
needsContinuation = result.needsContinuation
malformedToolRetries = result.malformedToolInput ? malformedToolRetries + 1 : 0
step = result.step + 1
if (needsContinuation) {
yield* runPendingCompaction(input.sessionID)

View file

@ -53,6 +53,7 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
called: boolean
settled: boolean
providerExecuted: boolean
rawInput: string
}
>()
let assistantMessageID = input.assistantMessageID
@ -60,6 +61,7 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
let stepFailed = false
let providerFailed = false
let retryEvidence = false
let malformedToolInput = false
let stepFailure: SessionError.Error | undefined
let stepSettlement:
| {
@ -112,12 +114,12 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
if (state !== undefined) current.state = { ...current.state, ...state }
return Effect.succeed(current.ordinal)
})
const end = Effect.fnUntraced(function* (id: string, state?: Record<string, unknown>) {
const end = Effect.fnUntraced(function* (id: string, state?: Record<string, unknown>, value?: string) {
const current = chunks.get(id)
if (!current) return yield* Effect.die(new Error(`${name} end before start: ${id}`))
yield* ended(
id,
current.values.join(""),
value ?? current.values.join(""),
current.ordinal,
state === undefined ? current.state : { ...current.state, ...state },
)
@ -175,7 +177,12 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
yield* toolInput.flush()
})
const startToolInput = Effect.fnUntraced(function* (event: { readonly id: string; readonly name: string }) {
const startToolInput = Effect.fnUntraced(function* (event: {
readonly id: string
readonly name: string
readonly providerExecuted?: boolean
readonly providerMetadata?: ProviderMetadata
}) {
if (tools.has(event.id)) return yield* Effect.die(new Error(`Duplicate tool input start: ${event.id}`))
const assistantMessageID = yield* startAssistant()
tools.set(event.id, {
@ -183,7 +190,8 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
name: event.name,
called: false,
settled: false,
providerExecuted: false,
providerExecuted: event.providerExecuted === true,
rawInput: "",
})
yield* toolInput.start(event.id)
yield* events.publish(SessionEvent.Tool.Input.Started, {
@ -191,16 +199,60 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
assistantMessageID,
callID: event.id,
name: event.name,
executed: event.providerExecuted,
state: providerState(event.providerMetadata),
})
})
const endToolInput = Effect.fnUntraced(function* (event: { readonly id: string; readonly name: string }) {
const endToolInput = Effect.fnUntraced(function* (event: {
readonly id: string
readonly name: string
readonly input?: string
}) {
const tool = tools.get(event.id)
if (!tool) return yield* Effect.die(new Error(`Tool input end before start: ${event.id}`))
if (tool.name !== event.name)
return yield* Effect.die(new Error(`Tool input name changed for ${event.id}: ${tool.name} -> ${event.name}`))
if (!toolInput.has(event.id)) return yield* Effect.die(new Error(`Duplicate tool input end: ${event.id}`))
yield* toolInput.end(event.id)
tool.rawInput = event.input ?? tool.rawInput
yield* toolInput.end(event.id, undefined, event.input)
})
const failMalformedToolInput = Effect.fnUntraced(function* (
event: {
readonly id: string
readonly name: string
readonly raw: string
readonly providerExecuted?: boolean
readonly providerMetadata?: ProviderMetadata
},
error: SessionError.Error,
assistantError: SessionError.Error,
) {
if (!tools.has(event.id)) yield* startToolInput(event)
const tool = tools.get(event.id)
if (!tool || tool.called || tool.settled) return false
if (tool.name !== event.name)
return yield* Effect.die(new Error(`Tool input name changed for ${event.id}: ${tool.name} -> ${event.name}`))
tool.providerExecuted = event.providerExecuted === true || tool.providerExecuted
const correctedRaw = !toolInput.has(event.id) && tool.rawInput !== event.raw ? event.raw : undefined
if (toolInput.has(event.id)) {
tool.rawInput = event.raw
yield* toolInput.end(event.id, undefined, event.raw)
}
tool.settled = true
malformedToolInput = true
yield* events.publish(SessionEvent.Tool.Failed, {
sessionID: input.sessionID,
assistantMessageID: tool.assistantMessageID,
callID: event.id,
error,
raw: correctedRaw,
executed: tool.providerExecuted,
})
yield* startAssistant()
if (stepFailure === undefined) stepFailure = assistantError
return true
})
const flush = Effect.fn("SessionRunner.flush")(function* () {
@ -210,11 +262,7 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
const failTools = Effect.fnUntraced(function* (error: SessionError.Error, mode: "all" | "hosted" | "uncalled") {
let failed = false
for (const [callID, tool] of tools) {
if (
tool.settled ||
(mode === "hosted" && !tool.providerExecuted) ||
(mode === "uncalled" && tool.called)
)
if (tool.settled || (mode === "hosted" && !tool.providerExecuted) || (mode === "uncalled" && tool.called))
continue
tool.settled = true
failed = true
@ -325,6 +373,7 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
if (tool.name !== event.name)
return yield* Effect.die(new Error(`Tool input name changed for ${event.id}: ${tool.name} -> ${event.name}`))
if (!toolInput.has(event.id)) return yield* Effect.die(new Error(`Tool input delta after end: ${event.id}`))
tool.rawInput += event.text
yield* toolInput.append(event.id, event.text)
yield* events.publish(SessionEvent.Tool.Input.Delta, {
sessionID: input.sessionID,
@ -337,16 +386,30 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
case "tool-input-end":
yield* endToolInput(event)
return
case "tool-input-error":
retryEvidence = true
if (
!(yield* failMalformedToolInput(
event,
{
type: "provider.invalid-output",
message: "Tool call arguments were malformed JSON and were not executed. Retry with valid JSON.",
},
{ type: "provider.invalid-output", message: event.message },
))
)
return yield* Effect.die(new Error(`Malformed tool input after call settlement: ${event.id}`))
return
case "tool-call": {
retryEvidence = true
if (!tools.has(event.id)) yield* startToolInput(event)
const tool = tools.get(event.id)!
if (toolInput.has(event.id)) yield* endToolInput(event)
if (toolInput.has(event.id)) yield* endToolInput({ id: event.id, name: event.name })
if (tool.name !== event.name)
return yield* Effect.die(new Error(`Tool call name changed for ${event.id}: ${tool.name} -> ${event.name}`))
if (tool.called) return yield* Effect.die(new Error(`Duplicate tool call: ${event.id}`))
tool.called = true
tool.providerExecuted = event.providerExecuted === true
tool.providerExecuted = event.providerExecuted === true || tool.providerExecuted
yield* events.publish(SessionEvent.Tool.Called, {
sessionID: input.sessionID,
assistantMessageID: tool.assistantMessageID,
@ -436,10 +499,13 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
publish,
flush,
failAssistant,
failMalformedToolInput,
publishStepFailure,
failUnsettledTools,
hasProviderError: () => providerFailed,
hasMalformedToolInput: () => malformedToolInput,
hasRetryEvidence: () => retryEvidence,
isProviderExecuted: (callID: string) => tools.get(callID)?.providerExecuted === true,
stepFailure: () => stepFailure,
stepSettlement: () => stepSettlement,
startAssistant,

View file

@ -110,7 +110,12 @@ const toolResult = (tool: SessionMessage.AssistantTool, providerMetadata: Provid
result:
tool.executed === true && tool.state.result !== undefined
? tool.state.result
: { error: tool.state.error, content: tool.state.content, structured: tool.state.structured },
: {
error: tool.state.error,
content: tool.state.content,
structured: tool.state.structured,
...(tool.state.raw === undefined ? {} : { raw: tool.state.raw }),
},
resultType: "error",
providerExecuted: tool.executed,
providerMetadata,