Jelajahi Sumber

fix(core): simplify compaction semantics (#36267)

Kit Langton 1 bulan lalu
induk
melakukan
cf5e0adf02

+ 137 - 173
packages/core/src/session/compaction.ts

@@ -17,7 +17,6 @@ import { Token } from "../util/token"
 const DEFAULT_BUFFER = 20_000
 const DEFAULT_KEEP_TOKENS = 8_000
 const TOOL_OUTPUT_MAX_CHARS = 2_000
-const SUMMARY_OUTPUT_TOKENS = 4_096
 const SUMMARY_TEMPLATE = `Output exactly the Markdown structure shown inside <template> and keep the section order unchanged. Do not include the <template> tags in your response.
 <template>
 ## Objective
@@ -61,20 +60,14 @@ type Dependencies = {
   readonly llm: {
     readonly stream: (request: LLMRequest) => Stream.Stream<LLMEvent, LLMError>
   }
+  readonly models: SessionRunnerModel.Interface
   readonly config: Settings
 }
 
 export type AutoInput = {
-  readonly sessionID: SessionSchema.ID
-  readonly messages: readonly SessionMessage.Info[]
-  readonly request: LLMRequest
-}
-
-type CompactInput = {
   readonly sessionID: SessionSchema.ID
   readonly messages: readonly SessionMessage.Info[]
   readonly model: Model
-  readonly inputID?: SessionMessage.ID
 }
 
 export type ManualInput = {
@@ -83,16 +76,27 @@ export type ManualInput = {
   readonly inputID: SessionMessage.ID
 }
 
+type Plan = {
+  readonly sessionID: SessionSchema.ID
+  readonly model: Model
+  readonly reason: SessionMessage.Compaction["reason"]
+  readonly prompt: string
+  readonly recent: string
+  readonly inputID?: SessionMessage.ID
+}
+
+export type Outcome =
+  | Pick<SessionMessage.CompactionCompleted, "status">
+  | Pick<SessionMessage.CompactionFailed, "status" | "error">
+
 export interface Interface {
-  readonly compactIfNeeded: (input: AutoInput) => Effect.Effect<boolean>
-  readonly compactAfterOverflow: (input: AutoInput) => Effect.Effect<boolean>
-  readonly compactManual: (input: ManualInput) => Effect.Effect<boolean>
+  readonly required: (input: AutoInput) => boolean
+  readonly compact: (input: AutoInput) => Effect.Effect<Outcome>
+  readonly compactManual: (input: ManualInput) => Effect.Effect<Outcome>
 }
 
 export class Service extends Context.Service<Service, Interface>()("@opencode/v2/SessionCompaction") {}
 
-const estimate = (value: unknown) => Token.estimate(JSON.stringify(value))
-
 const truncate = (value: string) =>
   value.length <= TOOL_OUTPUT_MAX_CHARS ? value : `${value.slice(0, TOOL_OUTPUT_MAX_CHARS)}\n[truncated]`
 
@@ -156,30 +160,33 @@ const select = (
 ): { readonly head: string; readonly recent: string } | undefined => {
   const conversation = messages
     .filter((message) => message.type !== "compaction")
-    .map(serialize)
-    .filter(Boolean)
+    .flatMap((message) => {
+      const text = serialize(message)
+      return text ? [{ message, text }] : []
+    })
   if (conversation.length === 0) return undefined
   let total = 0
   let split = conversation.length
-  let splitPrefix = ""
-  let splitSuffix = ""
   for (let index = conversation.length - 1; index >= 0; index--) {
-    const next = total + Token.estimate(conversation[index])
-    if (next > tokens) {
-      const remaining = Math.max(0, tokens - total) * 4
-      if (remaining > 0) {
-        splitPrefix = conversation[index].slice(0, -remaining)
-        splitSuffix = conversation[index].slice(-remaining)
-        split = index + 1
-      }
-      break
-    }
+    const next = total + Token.estimate(conversation[index].text)
+    if (split < conversation.length && next > tokens) break
     total = next
     split = index
   }
+  while (split > 0 && conversation[split].message.type !== "user") split--
+  if (split === 0) {
+    const latestUser = conversation.findLastIndex((item) => item.message.type === "user")
+    if (latestUser > 0) split = latestUser
+  }
   return {
-    head: [...conversation.slice(0, split), splitPrefix].filter(Boolean).join("\n\n"),
-    recent: [splitSuffix, ...conversation.slice(split)].filter(Boolean).join("\n\n"),
+    head: conversation
+      .slice(0, split)
+      .map((item) => item.text)
+      .join("\n\n"),
+    recent: conversation
+      .slice(split)
+      .map((item) => item.text)
+      .join("\n\n"),
   }
 }
 
@@ -192,37 +199,50 @@ export const buildPrompt = (input: { readonly previousSummary?: string; readonly
     ...input.context,
   ].join("\n\n")
 
+const planContent = (messages: readonly SessionMessage.Info[], tokens: number) => {
+  const selected = select(messages, tokens)
+  if (!selected) return
+  const previousSummary = messages.findLast(
+    (message) => message.type === "compaction" && message.status === "completed",
+  )
+  const previousRecent = previousSummary?.type === "compaction" ? previousSummary.recent : ""
+  const summarizeRecent = !previousRecent && !selected.head
+  return {
+    prompt: buildPrompt({
+      previousSummary: previousSummary?.type === "compaction" ? previousSummary.summary : undefined,
+      context: summarizeRecent ? [selected.recent] : [previousRecent, selected.head].filter(Boolean),
+    }),
+    recent: summarizeRecent ? "" : selected.recent,
+  }
+}
+
 const make = (dependencies: Dependencies) => {
   const config = dependencies.config
-  const compact = Effect.fn("SessionCompaction.compact")(function* (input: {
+  const failed = Effect.fnUntraced(function* (input: {
     readonly sessionID: SessionSchema.ID
-    readonly model: Model
     readonly reason: SessionMessage.Compaction["reason"]
-    readonly previousSummary?: string
-    readonly context: readonly string[]
-    readonly recent: string
-    readonly output?: number
+    readonly error: SessionError.Error
     readonly inputID?: SessionMessage.ID
   }) {
-    const output = input.output ?? input.model.route.defaults.limits?.output ?? 0
-    const summaryPrompt = buildPrompt({ previousSummary: input.previousSummary, context: input.context })
-    const summaryOutput = Math.min(output || SUMMARY_OUTPUT_TOKENS, SUMMARY_OUTPUT_TOKENS)
+    yield* dependencies.events.publish(SessionEvent.Compaction.Failed, input)
+    return { status: "failed" as const, error: input.error }
+  })
+  const execute = Effect.fn("SessionCompaction.execute")(function* (plan: Plan) {
     yield* dependencies.events.publish(SessionEvent.Compaction.Started, {
-      sessionID: input.sessionID,
-      reason: input.reason,
-      recent: input.recent,
-      inputID: input.inputID,
+      sessionID: plan.sessionID,
+      reason: plan.reason,
+      recent: plan.recent,
+      inputID: plan.inputID,
     })
 
     const chunks: string[] = []
     let failure: SessionError.Error | undefined
-    const summarized = yield* dependencies.llm
+    yield* dependencies.llm
       .stream(
         LLM.request({
-          model: input.model,
-          messages: [Message.user(summaryPrompt)],
+          model: plan.model,
+          messages: [Message.user(plan.prompt)],
           tools: [],
-          generation: { maxTokens: summaryOutput },
         }),
       )
       .pipe(
@@ -235,130 +255,110 @@ const make = (dependencies: Dependencies) => {
           if (LLMEvent.is.textDelta(event)) {
             chunks.push(event.text)
             return dependencies.events.publish(SessionEvent.Compaction.Delta, {
-              sessionID: input.sessionID,
+              sessionID: plan.sessionID,
               text: event.text,
             })
           }
           return Effect.void
         }),
-        Effect.as(true),
         Effect.catchTag("LLM.Error", (error) =>
           Effect.sync(() => {
             failure = toSessionError(error)
-            return false
           }),
         ),
         Effect.onInterrupt(() =>
-          input.reason === "auto"
-            ? dependencies.events.publish(SessionEvent.Compaction.Failed, {
-                sessionID: input.sessionID,
-                reason: input.reason,
+          plan.reason === "auto"
+            ? failed({
+                sessionID: plan.sessionID,
+                reason: plan.reason,
                 error: { type: "compaction.interrupted", message: "Compaction was interrupted" },
-                inputID: input.inputID,
-              })
+                inputID: plan.inputID,
+              }).pipe(Effect.asVoid)
             : Effect.void,
         ),
       )
     const summary = chunks.join("")
-    if (!summarized || failure || !summary.trim()) {
-      yield* dependencies.events.publish(SessionEvent.Compaction.Failed, {
-        sessionID: input.sessionID,
-        reason: input.reason,
-        error: failure ?? { type: "compaction.failed", message: "Compaction produced no summary" },
-        inputID: input.inputID,
+    if (failure || !summary.trim()) {
+      const error = failure ?? { type: "compaction.failed" as const, message: "Compaction produced no summary" }
+      return yield* failed({
+        sessionID: plan.sessionID,
+        reason: plan.reason,
+        error,
+        inputID: plan.inputID,
       })
-      return false
     }
     yield* dependencies.events.publish(SessionEvent.Compaction.Ended, {
-      sessionID: input.sessionID,
-      reason: input.reason,
+      sessionID: plan.sessionID,
+      reason: plan.reason,
       text: summary,
-      recent: input.recent,
+      recent: plan.recent,
     })
-    return true
+    return { status: "completed" as const }
   })
-  const compactAvailable = Effect.fn("SessionCompaction.compactAvailable")(function* (
-    input: CompactInput & {
-      readonly reason: SessionMessage.Compaction["reason"]
-      readonly output?: number
-    },
-  ) {
-    const selected = select(input.messages, config.tokens)
-    if (!selected) return false
-    const previousSummary = input.messages.find(
-      (message) => message.type === "compaction" && message.status === "completed",
-    )
-    const summarizeRecent = selected.head.length === 0
-    const previousRecent = previousSummary?.type === "compaction" ? previousSummary.recent : ""
-    return yield* compact({
-      sessionID: input.sessionID,
-      model: input.model,
-      reason: input.reason,
-      previousSummary: previousSummary?.type === "compaction" ? previousSummary.summary : undefined,
-      context: (summarizeRecent ? [previousRecent, selected.recent] : [previousRecent, selected.head]).filter(
-        Boolean,
-      ),
-      recent: summarizeRecent ? "" : selected.recent,
-      output: input.output,
-      inputID: input.inputID,
-    })
-  })
-  const compactAfterOverflow = Effect.fn("SessionCompaction.compactAfterOverflow")(function* (input: AutoInput) {
-    return yield* compactAvailable({
+  const compact = Effect.fn("SessionCompaction.compact")(function* (input: AutoInput) {
+    const content = planContent(input.messages, config.tokens)
+    if (content)
+      return yield* execute({
+        sessionID: input.sessionID,
+        model: input.model,
+        reason: "auto",
+        ...content,
+      })
+    const error = { type: "compaction.unavailable" as const, message: "Nothing to compact yet" }
+    return yield* failed({
       sessionID: input.sessionID,
-      messages: input.messages,
-      model: input.request.model,
       reason: "auto",
-      output: input.request.generation?.maxTokens ?? input.request.model.route.defaults.limits?.output ?? 0,
+      error,
     })
   })
-  const compactManual = Effect.fn("SessionCompaction.compactManual")(function* (input: CompactInput) {
-    return yield* compactAvailable({ ...input, reason: "manual" })
-  })
-  const compactIfNeeded = Effect.fn("SessionCompaction.compactIfNeeded")(function* (input: AutoInput) {
+  const required = (input: AutoInput) => {
     if (!config.auto) return false
-    const context = input.request.model.route.defaults.limits?.context
+    const context = input.model.route.defaults.limits?.context
     if (context === undefined || context <= 0) return false
-    const output = input.request.generation?.maxTokens ?? input.request.model.route.defaults.limits?.output ?? 0
-    if (
-      estimate({ system: input.request.system, messages: input.request.messages, tools: input.request.tools }) <=
-      context - Math.max(output, config.buffer)
-    )
-      return false
-    const selected = select(input.messages, config.tokens)
-    if (!selected) return false
-    const previousSummary = input.messages.find(
-      (message) => message.type === "compaction" && message.status === "completed",
+    const last = input.messages.findLast(
+      (message): message is SessionMessage.Assistant & { tokens: NonNullable<SessionMessage.Assistant["tokens"]> } =>
+        message.type === "assistant" && message.tokens !== undefined,
     )
-    if (!selected.head && previousSummary?.type !== "compaction") return false
-    const previousRecent = previousSummary?.type === "compaction" ? previousSummary.recent : ""
-    const summaryContext = [previousRecent, selected.head].filter(Boolean)
-    const summaryOutput = Math.min(output || SUMMARY_OUTPUT_TOKENS, SUMMARY_OUTPUT_TOKENS)
-    if (
-      Token.estimate(
-        buildPrompt({
-          previousSummary: previousSummary?.type === "compaction" ? previousSummary.summary : undefined,
-          context: summaryContext,
+    if (!last) return false
+    const output = input.model.route.defaults.limits?.output ?? 0
+    const used =
+      last.tokens.input + last.tokens.output + last.tokens.reasoning + last.tokens.cache.read + last.tokens.cache.write
+    if (used <= 0) return false
+    return used >= context - (output || config.buffer)
+  }
+  const compactManual = Effect.fn("SessionCompaction.compactManual")(function* (input: ManualInput) {
+    const content = planContent(input.messages, config.tokens)
+    if (!content)
+      return yield* failed({
+        sessionID: input.session.id,
+        reason: "manual",
+        error: { type: "compaction.unavailable", message: "Nothing to compact yet" },
+        inputID: input.inputID,
+      })
+    const resolved = yield* dependencies.models.resolve(input.session).pipe(
+      Effect.catch((cause) =>
+        failed({
+          sessionID: input.session.id,
+          reason: "manual",
+          error: toSessionError(cause),
+          inputID: input.inputID,
         }),
-      ) >
-      context - summaryOutput
+      ),
     )
-      return false
-    return yield* compact({
-      sessionID: input.sessionID,
-      model: input.request.model,
-      reason: "auto",
-      previousSummary: previousSummary?.type === "compaction" ? previousSummary.summary : undefined,
-      context: summaryContext,
-      recent: selected.recent,
-      output,
+    if ("status" in resolved) return resolved
+    return yield* execute({
+      sessionID: input.session.id,
+      model: resolved.model,
+      reason: "manual",
+      inputID: input.inputID,
+      ...content,
     })
   })
-  return {
-    compactIfNeeded,
-    compactAfterOverflow,
+  return Service.of({
+    required,
+    compact,
     compactManual,
-  }
+  })
 }
 
 export const layer = Layer.effect(
@@ -368,43 +368,7 @@ export const layer = Layer.effect(
     const llm = yield* LLMClient.Service
     const config = yield* Config.Service
     const models = yield* SessionRunnerModel.Service
-    const configured = settings(yield* config.entries())
-    const compaction = make({ events, llm, config: configured })
-
-    return Service.of({
-      compactIfNeeded: compaction.compactIfNeeded,
-      compactAfterOverflow: compaction.compactAfterOverflow,
-      compactManual: Effect.fn("SessionCompaction.compactManual")(function* (input) {
-        if (!select(input.messages, configured.tokens)) {
-          yield* events.publish(SessionEvent.Compaction.Failed, {
-            sessionID: input.session.id,
-            reason: "manual",
-            error: { type: "compaction.unavailable", message: "Nothing to compact yet" },
-            inputID: input.inputID,
-          })
-          return false
-        }
-        const resolved = yield* models.resolve(input.session).pipe(
-          Effect.catch((error) =>
-            events
-              .publish(SessionEvent.Compaction.Failed, {
-                sessionID: input.session.id,
-                reason: "manual",
-                error: toSessionError(error),
-                inputID: input.inputID,
-              })
-              .pipe(Effect.as(undefined)),
-          ),
-        )
-        if (!resolved) return false
-        return yield* compaction.compactManual({
-          sessionID: input.session.id,
-          messages: input.messages,
-          model: resolved.model,
-          inputID: input.inputID,
-        })
-      }),
-    })
+    return make({ events, llm, models, config: settings(yield* config.entries()) })
   }),
 )
 

+ 12 - 20
packages/core/src/session/runner/llm.ts

@@ -156,7 +156,7 @@ const layer = Layer.effect(
       sessionID: SessionSchema.ID,
       promotion: SessionPending.Delivery | undefined,
       step: number,
-      recoverOverflow?: typeof compaction.compactAfterOverflow,
+      recoverOverflow?: typeof compaction.compact,
       assistantMessageID?: SessionMessage.ID,
     ) {
       const session = yield* getSession(sessionID)
@@ -189,6 +189,12 @@ const layer = Layer.effect(
       const providerMetadataKey = model.route.providerMetadataKey ?? model.provider
       const entries = yield* SessionHistory.entriesForRunner(db, session.id, checkpoint.baselineSeq)
       const context = entries.map((entry) => entry.message)
+      const compactionInput = { sessionID: session.id, messages: context, model }
+      if (compaction.required(compactionInput) && !(yield* SessionPending.compaction(db, session.id))) {
+        const compacted = yield* compaction.compact(compactionInput)
+        if (compacted.status === "completed") return { _tag: "RestartAfterCompaction", step: currentStep } as const
+        return yield* new StepFailedError({ error: compacted.error })
+      }
       const isLastStep = agentInfo.steps !== undefined && currentStep >= agentInfo.steps
       const toolMaterialization = isLastStep ? undefined : yield* tools.materialize(agentInfo.permissions)
       const promptCacheKey = /^ses_[0-9a-f]{64}$/.test(session.id) ? session.id.slice(4) : session.id
@@ -208,12 +214,6 @@ const layer = Layer.effect(
       const toolFibers = yield* FiberSet.make<void, ToolOutputStore.Error>()
       const ownedToolFibers: Array<Fiber.Fiber<void, ToolOutputStore.Error>> = []
       let needsContinuation = false
-      // Automatic compaction completed; rebuild the request from compacted history.
-      if (
-        !(yield* SessionPending.compaction(db, session.id)) &&
-        (yield* compaction.compactIfNeeded({ sessionID: session.id, messages: context, request }))
-      )
-        return { _tag: "RestartAfterCompaction", step: currentStep } as const
       const startSnapshot = yield* snapshots.capture()
       const publisher = createLLMEventPublisher(events, {
         sessionID: session.id,
@@ -326,7 +326,8 @@ const layer = Layer.effect(
             recoverOverflow &&
             !publisher.hasRetryEvidence() &&
             isContextOverflowFailure(overflowFailure ?? streamFailure) &&
-            (yield* restore(recoverOverflow({ sessionID: session.id, messages: context, request })))
+            (yield* restore(recoverOverflow({ sessionID: session.id, messages: context, model }))).status ===
+              "completed"
           )
             return { _tag: "RestartAfterOverflowCompaction", step: currentStep } as const
 
@@ -446,7 +447,7 @@ const layer = Layer.effect(
       // Compaction restarts rebuild the request from compacted history without re-promoting.
       // Overflow recovery is one-shot: a post-compaction attempt must not recover another
       // overflow, so the recovery hook is dropped after it fires.
-      let recoverOverflow: typeof compaction.compactAfterOverflow | undefined = compaction.compactAfterOverflow
+      let recoverOverflow: typeof compaction.compact | undefined = compaction.compact
       let currentPromotion = promotion
       let currentStep = step
       let assistantMessageID: SessionMessage.ID | undefined
@@ -486,7 +487,7 @@ const layer = Layer.effect(
       sessionID: SessionSchema.ID,
     ) {
       const pending = yield* SessionPending.compaction(db, sessionID)
-      if (!pending) return false
+      if (!pending) return
       const session = yield* getSession(sessionID)
       return yield* Effect.uninterruptibleMask((restore) =>
         Effect.gen(function* () {
@@ -499,7 +500,7 @@ const layer = Layer.effect(
               })
             }),
           ).pipe(Effect.exit)
-          if (Exit.isSuccess(compacted) && compacted.value) return true
+          if (Exit.isSuccess(compacted)) return
           if (Exit.isFailure(compacted)) {
             const unsettled = yield* SessionPending.compaction(db, sessionID)
             if (unsettled)
@@ -511,15 +512,6 @@ const layer = Layer.effect(
               })
             return yield* Effect.failCause(compacted.cause)
           }
-          const unsettled = yield* SessionPending.compaction(db, sessionID)
-          if (unsettled)
-            yield* events.publish(SessionEvent.Compaction.Failed, {
-              sessionID,
-              reason: "manual",
-              error: { type: "compaction.failed", message: "Compaction could not start" },
-              inputID: unsettled.id,
-            })
-          return true
         }),
       )
     })

+ 2 - 1
packages/core/test/session-compaction.test.ts

@@ -147,10 +147,11 @@ it.effect("manual compaction summarizes short context instead of no-op", () =>
         messages: [userMessage],
         inputID: SessionMessage.ID.make("msg_manual_compaction"),
       }),
-    ).toBe(true)
+    ).toEqual({ status: "completed" })
     expect(Array.from(yield* Fiber.join(delta)).map((event) => event.data.text)).toEqual(["manual summary"])
 
     expect(requests).toHaveLength(1)
+    expect(requests[0]?.generation).toBeUndefined()
     expect(JSON.stringify(requests[0]?.messages)).toContain("Manual compaction should include this short conversation.")
     expect(yield* store.context(sessionID)).toMatchObject([
       { type: "compaction", reason: "manual", summary: "manual summary", recent: "" },

+ 41 - 2
packages/core/test/session-runner.test.ts

@@ -113,6 +113,16 @@ const reply = {
     LLMEvent.finish({ reason: "stop" }),
   ],
   text: (text: string, id: string) => fragmentFixture("text", id, [text]).completeEvents,
+  textWithUsage: (text: string, id: string, inputTokens: number) =>
+    fragmentFixture("text", id, [text]).completeEvents.map((event) =>
+      LLMEvent.is.stepFinish(event)
+        ? LLMEvent.stepFinish({
+            index: event.index,
+            reason: event.reason,
+            usage: { inputTokens, nonCachedInputTokens: inputTokens },
+          })
+        : event,
+    ),
   tool: (id: string, name: string, input: unknown) => [
     LLMEvent.stepStart({ index: 0 }),
     LLMEvent.toolCall({ id, name, input }),
@@ -1696,7 +1706,7 @@ describe("SessionRunnerLLM", () => {
   it.effect("automatically compacts into a completed summary and retained recent turn", () =>
     Effect.gen(function* () {
       const session = yield* setup
-      response = reply.text("Earlier answer", "text-first")
+      response = reply.textWithUsage("Earlier answer", "text-first", 3_950)
       yield* admit(session, "Earlier question ".repeat(180))
       yield* session.resume(sessionID)
 
@@ -1704,7 +1714,7 @@ describe("SessionRunnerLLM", () => {
       requests.length = 0
       responses = [
         reply.text("## Objective\n- Preserve the task", "text-summary"),
-        reply.text("Continued", "text-final"),
+        reply.textWithUsage("Continued", "text-final", 3_950),
       ]
       yield* admit(session, "Recent exact request ".repeat(180))
       yield* session.resume(sessionID)
@@ -1743,6 +1753,35 @@ describe("SessionRunnerLLM", () => {
     }),
   )
 
+  it.effect("stops after required automatic compaction fails", () =>
+    Effect.gen(function* () {
+      const session = yield* setup
+      response = reply.textWithUsage("Earlier answer", "text-before-failed-compaction", 3_950)
+      yield* admit(session, "Earlier question ".repeat(180))
+      yield* session.resume(sessionID)
+
+      currentModel = compactModel
+      requests.length = 0
+      responses = [
+        [LLMEvent.providerError({ message: "Unsupported parameter: max_output_tokens" })],
+        reply.text("Must not run", "text-after-failed-compaction"),
+      ]
+      yield* admit(session, "Recent exact request ".repeat(180))
+      expect(yield* Effect.exit(session.resume(sessionID))).toMatchObject({ _tag: "Failure" })
+
+      expect(requests).toHaveLength(1)
+      expect(requests[0]?.generation).toBeUndefined()
+      expect(yield* session.context(sessionID)).toContainEqual(
+        expect.objectContaining({
+          type: "compaction",
+          status: "failed",
+          reason: "auto",
+          error: expect.objectContaining({ message: "Unsupported parameter: max_output_tokens" }),
+        }),
+      )
+    }),
+  )
+
   it.effect("forces one compaction and retries after provider context overflow", () =>
     Effect.gen(function* () {
       const session = yield* setupOverflowRecovery