fix(ai): preserve provider tool input identity
This commit is contained in:
parent
7840a0d378
commit
0f26246bfd
21 changed files with 833 additions and 74 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue