diff --git a/packages/core/src/session/runner/run-turn.ts b/packages/core/src/session/runner/run-turn.ts index 404fa4dfed..273e84830f 100644 --- a/packages/core/src/session/runner/run-turn.ts +++ b/packages/core/src/session/runner/run-turn.ts @@ -18,7 +18,7 @@ import { isContextOverflowFailure, type ProviderErrorEvent, } from "@opencode-ai/llm" -import { DateTime, Effect, Schema, Semaphore, Stream } from "effect" +import { Cause, DateTime, Effect, FiberSet, Option, Schema, Semaphore, Stream } from "effect" import { AgentV2 } from "../../agent" import { Config } from "../../config" import { Database } from "../../database/database" @@ -26,9 +26,11 @@ import { EventV2 } from "../../event" import { Location } from "../../location" import { ModelV2 } from "../../model" import { ProviderV2 } from "../../provider" +import { QuestionV2 } from "../../question" import { SkillGuidance } from "../../skill/guidance" import { SystemContext } from "../../system-context/index" import { SystemContextRegistry } from "../../system-context/registry" +import { ToolOutputStore } from "../../tool-output-store" import { ToolRegistry } from "../../tool/registry" import { SessionCompaction } from "../compaction" import { SessionContextEpoch } from "../context-epoch" @@ -40,7 +42,6 @@ import { SessionStore } from "../store" import type { RunError } from "./index" import { SessionRunnerModel } from "./model" import { createLLMEventPublisher } from "./publish-llm-event" -import { SettleProviderTurn } from "./settle-provider-turn" import { toLLMMessages } from "./to-llm-message" export interface Input { @@ -53,6 +54,12 @@ const AttemptResult = Schema.TaggedUnion({ CompactedOverflow: {}, }) +const isQuestionRejected = (cause: Cause.Cause) => + cause.reasons.some((reason) => Cause.isDieReason(reason) && reason.defect instanceof QuestionV2.RejectedError) + +const awaitTools = (fibers: FiberSet.FiberSet) => + Effect.raceFirst(FiberSet.join(fibers), FiberSet.awaitEmpty(fibers)) + export const make = Effect.gen(function* () { const events = yield* EventV2.Service const llm = yield* LLMClient.Service @@ -202,42 +209,48 @@ export const make = Effect.gen(function* () { ), ) }) - const result = yield* SettleProviderTurn.run({ - stream: (runTool) => - llm.stream(prepared.request).pipe( - Stream.runForEach((event) => - Effect.gen(function* () { - if (overflowFailure || publisher.hasProviderError()) return - if ( - LLMEvent.is.providerError(event) && - isContextOverflowFailure(event) && - !publisher.hasAssistantStarted() - ) { - overflowFailure = event - return - } - yield* publish(event) - if (event.type !== "tool-call" || event.providerExecuted) return - yield* runTool(toolEffect(event)) - }), - ), - Effect.ensuring(withPublication(publisher.flush())), - ), - recoverOverflow: (failure, restore) => + const settleProviderTurn = Effect.fnUntraced(function* () { + const tools = yield* FiberSet.make() + return yield* Effect.uninterruptibleMask((restore) => Effect.gen(function* () { - if (!canRecoverOverflow || publisher.hasAssistantStarted()) return false - if (!isContextOverflowFailure(overflowFailure ?? failure)) return false - return yield* restore( - compaction.compactAfterOverflow({ - sessionID: prepared.session.id, - entries: prepared.entries, - model: prepared.model, - request: prepared.request, - }), + const stream = yield* restore( + llm.stream(prepared.request).pipe( + Stream.runForEach((event) => + Effect.gen(function* () { + if (overflowFailure || publisher.hasProviderError()) return + if ( + LLMEvent.is.providerError(event) && + isContextOverflowFailure(event) && + !publisher.hasAssistantStarted() + ) { + overflowFailure = event + return + } + yield* publish(event) + if (event.type !== "tool-call" || event.providerExecuted) return + yield* toolEffect(event).pipe(FiberSet.run(tools)) + }), + ), + Effect.ensuring(withPublication(publisher.flush())), + ), + ).pipe(Effect.exit) + const failure = + stream._tag === "Failure" ? Option.getOrUndefined(Cause.findErrorOption(stream.cause)) : undefined + if ( + canRecoverOverflow && + !publisher.hasAssistantStarted() && + isContextOverflowFailure(overflowFailure ?? failure) && + (yield* restore( + compaction.compactAfterOverflow({ + sessionID: prepared.session.id, + entries: prepared.entries, + model: prepared.model, + request: prepared.request, + }), + )) ) - }), - projectProviderFailure: (failure) => - Effect.gen(function* () { + return AttemptResult.cases.CompactedOverflow.make({}) + if (overflowFailure) yield* publish(overflowFailure) if (failure instanceof LLMError && !publisher.hasProviderError()) { yield* failUnsettled("Provider did not return a tool result", true) @@ -250,17 +263,35 @@ export const make = Effect.gen(function* () { }), ) } + const streamInterrupted = stream._tag === "Failure" && Cause.hasInterrupts(stream.cause) + if (streamInterrupted) yield* FiberSet.clear(tools) + const settled = yield* restore(awaitTools(tools)).pipe(Effect.exit) + if (settled._tag === "Failure" && isQuestionRejected(settled.cause)) { + yield* FiberSet.clear(tools) + yield* failUnsettled("Tool execution interrupted") + return yield* Effect.interrupt + } + const toolInterrupted = settled._tag === "Failure" && Cause.hasInterrupts(settled.cause) + if (toolInterrupted) yield* FiberSet.clear(tools) + if (streamInterrupted || toolInterrupted || publisher.hasProviderError()) + yield* failUnsettled("Tool execution interrupted") + if (settled._tag === "Failure" && !toolInterrupted) { + const failure = Cause.squash(settled.cause) + yield* failUnsettled( + `Tool execution failed: ${failure instanceof Error ? failure.message : String(failure)}`, + ) + } + if (stream._tag === "Success" && !publisher.hasProviderError()) + yield* failUnsettled("Provider did not return a tool result", true) + if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause) + if (settled._tag === "Failure") return yield* Effect.failCause(settled.cause) + return AttemptResult.cases.Complete.make({ + needsContinuation: !publisher.hasProviderError() && needsContinuation, + }) }), - hasProviderError: publisher.hasProviderError, - failUnsettled, - }) - return SettleProviderTurn.Result.match(result, { - RecoveredOverflow: () => AttemptResult.cases.CompactedOverflow.make({}), - Complete: () => - AttemptResult.cases.Complete.make({ - needsContinuation: !publisher.hasProviderError() && needsContinuation, - }), + ) }) + return yield* settleProviderTurn() }, Effect.scoped) const run = Effect.fn("SessionRunner.runTurn")(function* (input: Input): Effect.fn.Return { diff --git a/packages/core/src/session/runner/settle-provider-turn.ts b/packages/core/src/session/runner/settle-provider-turn.ts deleted file mode 100644 index 2e6302bfb2..0000000000 --- a/packages/core/src/session/runner/settle-provider-turn.ts +++ /dev/null @@ -1,71 +0,0 @@ -export * as SettleProviderTurn from "./settle-provider-turn" - -import { Cause, Effect, FiberSet, Option, Schema } from "effect" -import { QuestionV2 } from "../../question" -import { ToolOutputStore } from "../../tool-output-store" - -export const Result = Schema.TaggedUnion({ - Complete: {}, - RecoveredOverflow: {}, -}) - -const isQuestionRejected = (cause: Cause.Cause) => - cause.reasons.some((reason) => Cause.isDieReason(reason) && reason.defect instanceof QuestionV2.RejectedError) - -const awaitTools = (fibers: FiberSet.FiberSet) => - Effect.raceFirst(FiberSet.join(fibers), FiberSet.awaitEmpty(fibers)) - -/** - * Runs one provider response together with every local tool it starts. - * - * The provider and tools remain interruptible. Once the provider stops, cleanup - * cannot be interrupted before every started tool is observed, cancelled, or - * settled and its original failure is propagated. - */ -export const run = Effect.fn("SessionRunner.settleProviderTurn")(function* (input: { - readonly stream: ( - runTool: (effect: Effect.Effect) => Effect.Effect, - ) => Effect.Effect - readonly recoverOverflow: ( - failure: unknown, - restore: (effect: Effect.Effect) => Effect.Effect, - ) => Effect.Effect - readonly projectProviderFailure: (failure: unknown) => Effect.Effect - readonly hasProviderError: () => boolean - readonly failUnsettled: (message: string, providerExecuted?: boolean) => Effect.Effect -}) { - const tools = yield* FiberSet.make() - return yield* Effect.uninterruptibleMask((restore) => - Effect.gen(function* () { - const stream = yield* restore(input.stream((effect) => effect.pipe(FiberSet.run(tools)))).pipe(Effect.exit) - const failure = stream._tag === "Failure" ? Option.getOrUndefined(Cause.findErrorOption(stream.cause)) : undefined - if (yield* input.recoverOverflow(failure, restore)) return Result.cases.RecoveredOverflow.make({}) - - yield* input.projectProviderFailure(failure) - const streamInterrupted = stream._tag === "Failure" && Cause.hasInterrupts(stream.cause) - if (streamInterrupted) yield* FiberSet.clear(tools) - const settled = yield* restore(awaitTools(tools)).pipe(Effect.exit) - if (settled._tag === "Failure" && isQuestionRejected(settled.cause)) { - yield* FiberSet.clear(tools) - yield* input.failUnsettled("Tool execution interrupted") - return yield* Effect.interrupt - } - - const toolInterrupted = settled._tag === "Failure" && Cause.hasInterrupts(settled.cause) - if (toolInterrupted) yield* FiberSet.clear(tools) - if (streamInterrupted || toolInterrupted || input.hasProviderError()) - yield* input.failUnsettled("Tool execution interrupted") - if (settled._tag === "Failure" && !toolInterrupted) { - const failure = Cause.squash(settled.cause) - yield* input.failUnsettled( - `Tool execution failed: ${failure instanceof Error ? failure.message : String(failure)}`, - ) - } - if (stream._tag === "Success" && !input.hasProviderError()) - yield* input.failUnsettled("Provider did not return a tool result", true) - if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause) - if (settled._tag === "Failure") return yield* Effect.failCause(settled.cause) - return Result.cases.Complete.make({}) - }), - ) -}, Effect.scoped)