Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 8 additions & 33 deletions packages/core/src/session/runner/llm.ts
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,6 @@ const layer = Layer.effect(
sessionID: SessionSchema.ID,
promotion: SessionPending.Delivery | undefined,
step: number,
recoverMalformedToolInput: boolean,
recoverOverflow?: typeof compaction.compact,
assistantMessageID?: SessionMessage.ID,
) {
Expand Down Expand Up @@ -144,6 +143,10 @@ const layer = Layer.effect(
}
}
yield* publish(event)
if (LLMEvent.is.toolInputError(event)) {
if (prepared.resolveToolCall(event.name).type === "settle") needsContinuation = true
return
}
if (event.type !== "tool-call" || event.providerExecuted) return
const tool = prepared.resolveToolCall(event.name)
if (tool.type === "reject") {
Expand Down Expand Up @@ -340,27 +343,15 @@ const layer = Layer.effect(
)
}

const recoveredMalformedToolInput =
recoverMalformedToolInput &&
publisher.hasMalformedToolInput() &&
stream._tag === "Success" &&
stepSettlement !== undefined &&
!providerFailed &&
!streamInterrupted &&
!userDeclined &&
!toolsInterrupted &&
infraError === undefined

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 (toolsInterrupted && settled._tag === "Failure") return yield* Effect.failCause(settled.cause)
if (stepFailure && !recoveredMalformedToolInput) return yield* new StepFailedError({ error: stepFailure })
if (stepFailure) return yield* new StepFailedError({ error: stepFailure })
return {
_tag: "Completed",
needsContinuation: needsContinuation || recoveredMalformedToolInput,
malformedToolInput: recoveredMalformedToolInput,
needsContinuation,
step: currentStep,
} as const
}),
Expand All @@ -371,7 +362,6 @@ 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
Expand All @@ -382,14 +372,7 @@ const layer = Layer.effect(
let assistantMessageID: SessionMessage.ID | undefined
while (true) {
const attempt = yield* Effect.suspend(() =>
attemptStep(
sessionID,
currentPromotion,
currentStep,
recoverMalformedToolInput,
recoverOverflow,
assistantMessageID,
),
attemptStep(sessionID, currentPromotion, currentStep, recoverOverflow, assistantMessageID),
).pipe(
Effect.tapError((error) =>
error instanceof SessionRunnerRetry.RetryableFailure
Expand All @@ -414,7 +397,6 @@ const layer = Layer.effect(
if (attempt._tag === "Completed")
return {
needsContinuation: attempt.needsContinuation,
malformedToolInput: attempt.malformedToolInput,
step: attempt.step,
}
if (attempt._tag === "RestartAfterOverflowCompaction") recoverOverflow = undefined
Expand Down Expand Up @@ -474,26 +456,19 @@ const layer = Layer.effect(
while (shouldRun) {
let needsContinuation = true
let step = 1
let canRecoverMalformedToolInput = true
// 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.
while (needsContinuation) {
const result = yield* runStep(
input.sessionID,
promotion,
step,
canRecoverMalformedToolInput,
)
const result = yield* runStep(input.sessionID, promotion, step)
// 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)) {
titleAttempted.add(input.sessionID)
forkTitle(title.generateForFirstPrompt(yield* getSession(input.sessionID)).pipe(Effect.ignore))
}
needsContinuation = result.needsContinuation
if (result.malformedToolInput) canRecoverMalformedToolInput = false
step = result.step + 1
if (needsContinuation) {
yield* runPendingCompaction(input.sessionID)
Expand Down
5 changes: 0 additions & 5 deletions packages/core/src/session/runner/publish-llm-event.ts
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,6 @@ 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:
| {
Expand Down Expand Up @@ -216,7 +215,6 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
readonly id: string
readonly name: string
readonly raw: string
readonly message: string
}) {
if (!tools.has(event.id)) yield* startToolInput(event)
const tool = tools.get(event.id)
Expand All @@ -226,7 +224,6 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
return yield* Effect.die(new Error(`Tool input name changed for ${event.id}: ${tool.name} -> ${event.name}`))
if (toolInput.has(event.id)) yield* endToolInput(event, event.raw)
tool.settled = true
malformedToolInput = true
yield* events.publish(SessionEvent.Tool.Failed, {
sessionID: input.sessionID,
assistantMessageID: tool.assistantMessageID,
Expand All @@ -237,7 +234,6 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
},
executed: false,
})
if (stepFailure === undefined) stepFailure = { type: "provider.invalid-output", message: event.message }
})

const flush = Effect.fn("SessionRunner.flush")(function* () {
Expand Down Expand Up @@ -482,7 +478,6 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
publishStepFailure,
failUnsettledTools,
hasProviderError: () => providerFailed,
hasMalformedToolInput: () => malformedToolInput,
hasRetryEvidence: () => retryEvidence,
stepFailure: () => stepFailure,
stepSettlement: () => stepSettlement,
Expand Down
51 changes: 39 additions & 12 deletions packages/core/test/session-runner.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4164,7 +4164,7 @@ describe("SessionRunnerLLM", () => {
}),
)

it.effect("continues once after malformed local tool input without exposing raw arguments", () =>
it.effect("continues after malformed local tool input without exposing raw arguments", () =>
Effect.gen(function* () {
const session = yield* setup
yield* admit(session, "Recover malformed tool input")
Expand Down Expand Up @@ -4226,7 +4226,6 @@ describe("SessionRunnerLLM", () => {
message.type === "assistant" && message.content.some((item) => item.type === "tool"),
)
expect(failed).toMatchObject({
error: { type: "provider.invalid-output", message: "Invalid JSON input for test tool call echo" },
content: [
{
type: "tool",
Expand All @@ -4244,10 +4243,11 @@ describe("SessionRunnerLLM", () => {
],
})
if (!failed) throw new Error("Malformed tool assistant missing")
expect(failed.error).toBeUndefined()
expect((yield* recordedStepSettlementEvents(sessionID, failed.id)).map((event) => event.type)).toEqual([
"session.step.started.1",
"session.tool.failed.1",
"session.step.failed.1",
"session.step.ended.1",
])
const database = (yield* Database.Service).db
const durable = yield* database
Expand Down Expand Up @@ -4422,7 +4422,7 @@ describe("SessionRunnerLLM", () => {
}),
)

it.effect("does not reset the malformed recovery budget after a valid tool step", () =>
it.effect("continues after repeated malformed tool input", () =>
Effect.gen(function* () {
const session = yield* setup
yield* admit(session, "Keep producing malformed tools")
Expand All @@ -4444,16 +4444,43 @@ describe("SessionRunnerLLM", () => {
reply.stop(),
]

expect(yield* session.resume(sessionID).pipe(Effect.flip)).toMatchObject({
error: {
type: "provider.invalid-output",
message: "Invalid JSON input for test tool call echo",
},
})
yield* session.resume(sessionID)

expect(requests).toHaveLength(3)
expect(requests).toHaveLength(4)
expect(executions).toEqual(["valid"])
expect((yield* recordedEventTypes(sessionID)).filter((type) => type === "session.step.failed.1")).toHaveLength(2)
expect((yield* recordedEventTypes(sessionID)).filter((type) => type === "session.step.failed.1")).toHaveLength(0)
}),
)

it.effect("does not continue malformed tool input past the agent step limit", () =>
Effect.gen(function* () {
const session = yield* setup
const agents = yield* AgentV2.Service
yield* agents.transform((editor) =>
editor.update(AgentV2.ID.make("build"), (agent) => {
agent.steps = 2
}),
)
yield* admit(session, "Stop malformed tools at the step limit")
const malformed = (id: string) => [
LLMEvent.stepStart({ index: 0 }),
LLMEvent.toolInputError({
id,
name: "echo",
raw: '{"text":"partial',
message: "Invalid JSON input for test tool call echo",
}),
LLMEvent.stepFinish({ index: 0, reason: "tool-calls" }),
LLMEvent.finish({ reason: "tool-calls" }),
]
responses = [malformed("call-first"), malformed("call-at-limit")]

yield* session.resume(sessionID)

expect(requests).toHaveLength(2)
expect(requests[0]?.toolChoice).toBeUndefined()
expect(requests[1]?.toolChoice).toMatchObject({ type: "none" })
expect((yield* recordedEventTypes(sessionID)).filter((type) => type === "session.tool.failed.1")).toHaveLength(2)
}),
)

Expand Down
Loading