Merge pull request #12203 from Kilo-Org/fix-model-cache-race

fix(cli): isolate model cache refresh from caller cancellation

packages/opencode/src/kilocode/effect/cause.ts

@@ -0,0 +1,5 @@
+import { Cause } from "effect"
+
+export const isInterrupted = Cause.hasInterruptsOnly
+
+export const shouldReportPromptFailure = (cause: Cause.Cause<unknown>) => !isInterrupted(cause)

packages/opencode/src/provider/model-cache.ts

@@ -1,6 +1,6 @@
 // kilocode_change - new file
 import { fetchKiloModels, type KiloModelsResult } from "@kilocode/kilo-gateway"
-import { Context, Duration, Effect, Layer, Schema } from "effect"
+import { Context, Deferred, Duration, Effect, Exit, Layer, Schema, Scope } from "effect"
 import { FetchHttpClient, HttpClient, HttpClientRequest, HttpClientResponse } from "effect/unstable/http"
 import { Config } from "../config/config"
 import { Auth } from "../auth"
@@ -13,6 +13,7 @@ type Options = { -readonly [K in keyof KiloOptions]?: KiloOptions[K] } & { apiKe
 type Failure = NonNullable<KiloModelsResult["error"]>
 type Result = { readonly models: Models; readonly error?: Failure }
 type View = { models?: Models; timestamp?: number }
+type Flight = { readonly done: Deferred.Deferred<Result, unknown>; version: number }
 
 export interface KiloModels {
   readonly fetch: (options: KiloOptions) => Effect.Effect<KiloModelsResult, unknown>
@@ -28,9 +29,10 @@ export const kiloModelsLayer = Layer.succeed(
 )
 type Cell = {
   readonly providerID: string
+  readonly options: Options
   readonly view: View
-  readonly cached: Effect.Effect<Result, unknown>
-  readonly invalidate: Effect.Effect<void>
+  cached?: { readonly result: Result; readonly expires: number }
+  flight?: Flight
 }
 
 export interface Interface {
@@ -62,6 +64,7 @@ export const layer: Layer.Layer<
     const cfg = yield* Config.Service
     const kilo = yield* KiloModelsService
     const http = yield* HttpClient.HttpClient
+    const scope = yield* Scope.Scope
     const cells = new Map<string, Cell>()
     const active = new Map<string, Cell>()
     const versions = new Map<string, number>()
@@ -189,14 +192,24 @@ export const layer: Layer.Layer<
       const existing = cells.get(id)
       if (existing) return existing
       const view: View = {}
-      const [cached, invalidate] = yield* Effect.cachedInvalidateWithTTL(load(providerID, options), ttl)
-      const next = { providerID, view, cached, invalidate }
+      const next: Cell = { providerID, options, view }
       cells.set(id, next)
       return next
     })
 
-    // Failed loads are not cached so a temporary outage can recover on the next read.
-    const evaluate = (entry: Cell) => entry.cached.pipe(Effect.tapCause(() => entry.invalidate))
+    const invalidate = (entry: Cell) =>
+      Effect.sync(() => {
+        entry.cached = undefined
+      })
+
+    const detach = (entry: Cell) =>
+      invalidate(entry).pipe(
+        Effect.tap(() =>
+          Effect.sync(() => {
+            entry.flight = undefined
+          }),
+        ),
+      )
 
     const commit = (providerID: string, version: number, entry: Cell, result: Result) =>
       Effect.sync(() => {
@@ -214,6 +227,42 @@ export const layer: Layer.Layer<
         return result.models
       })
 
+    // A refresh belongs to the cache service, not the caller that happened to start it.
+    const evaluate = (entry: Cell, version: number) =>
+      Effect.uninterruptibleMask((restore) =>
+        Effect.gen(function* () {
+          const cached = entry.cached
+          if (cached && cached.expires > Date.now()) {
+            yield* commit(entry.providerID, version, entry, cached.result)
+            return cached.result
+          }
+
+          const existing = entry.flight
+          if (existing) {
+            existing.version = version
+            return yield* restore(Deferred.await(existing.done))
+          }
+
+          const done = yield* Deferred.make<Result, unknown>()
+          const flight = { done, version } satisfies Flight
+          entry.flight = flight
+          yield* Effect.uninterruptibleMask((restore) =>
+            Effect.gen(function* () {
+              const exit = yield* restore(load(entry.providerID, entry.options)).pipe(Effect.exit)
+              if (entry.flight === flight) {
+                entry.flight = undefined
+                if (Exit.isSuccess(exit)) {
+                  entry.cached = { result: exit.value, expires: Date.now() + Duration.toMillis(ttl) }
+                  yield* commit(entry.providerID, flight.version, entry, exit.value)
+                }
+              }
+              yield* Deferred.done(done, exit)
+            }),
+          ).pipe(Effect.forkIn(scope, { startImmediately: true }))
+          return yield* restore(Deferred.await(done))
+        }),
+      )
+
     const get = Effect.fn("ModelCache.get")(function* (providerID: string) {
       const entry = active.get(providerID)
       if (!entry?.view.models || entry.view.timestamp === undefined) {
@@ -226,7 +275,7 @@ export const layer: Layer.Layer<
         log.debug("cache expired", { providerID, age })
         entry.view.models = undefined
         entry.view.timestamp = undefined
-        yield* entry.invalidate
+        yield* invalidate(entry)
         return
       }
 
@@ -241,8 +290,8 @@ export const layer: Layer.Layer<
       versions.set(providerID, version)
       const entry = yield* cell(providerID, options)
       log.info("fetching models", { providerID })
-      const result = yield* evaluate(entry)
-      return yield* commit(providerID, version, entry, result)
+      const result = yield* evaluate(entry, version)
+      return result.models
     })
 
     const refresh = Effect.fn("ModelCache.refresh")(function* (providerID: string, options?: Options) {
@@ -250,16 +299,16 @@ export const layer: Layer.Layer<
       versions.set(providerID, version)
       const entry = yield* cell(providerID, options)
       log.info("refreshing models", { providerID })
-      yield* entry.invalidate
-      const result = yield* evaluate(entry)
-      return yield* commit(providerID, version, entry, result)
+      yield* invalidate(entry)
+      const result = yield* evaluate(entry, version)
+      return result.models
     })
 
     const clear = Effect.fn("ModelCache.clear")(function* (providerID: string) {
       versions.set(providerID, (versions.get(providerID) ?? 0) + 1)
       const entries = [...cells.entries()].filter(([, entry]) => entry.providerID === providerID)
       yield* Effect.all(
-        entries.map(([id, entry]) => entry.invalidate.pipe(Effect.tap(() => Effect.sync(() => cells.delete(id))))),
+        entries.map(([id, entry]) => detach(entry).pipe(Effect.tap(() => Effect.sync(() => cells.delete(id))))),
         { discard: true },
       )
       active.delete(providerID)

packages/opencode/src/session/prompt.ts

@@ -11,6 +11,7 @@ import { KiloSessionProcessor } from "@/kilocode/session/processor" // kilocode_
 import { KiloSessionOverflow } from "@/kilocode/session/overflow" // kilocode_change
 import { KiloReference } from "@/kilocode/reference/contains" // kilocode_change
 import { KiloReadObject } from "@/kilocode/tool/read-object" // kilocode_change
+import { isInterrupted } from "@/kilocode/effect/cause" // kilocode_change
 import * as SandboxPolicy from "@/kilocode/sandbox/policy" // kilocode_change
 import { CommandTimeout } from "@/kilocode/command-timeout" // kilocode_change
 import { Suggestion } from "@/kilocode/suggestion" // kilocode_change
@@ -716,6 +717,7 @@ export const layer = Layer.effect(
     ) {
       const exit = yield* provider.getModel(providerID, modelID).pipe(Effect.exit)
       if (Exit.isSuccess(exit)) return exit.value
+      if (isInterrupted(exit.cause)) return yield* Effect.interrupt // kilocode_change
       const err = Cause.squash(exit.cause)
       if (Provider.ModelNotFoundError.isInstance(err)) {
         const hint = err.suggestions?.length ? ` Did you mean: ${err.suggestions.join(", ")}?` : ""

packages/opencode/test/kilocode/effect-cause.test.ts

@@ -0,0 +1,13 @@
+import { expect, test } from "bun:test"
+import { Cause, Effect, Exit } from "effect"
+import { isInterrupted, shouldReportPromptFailure } from "../../src/kilocode/effect/cause"
+
+test("recognizes a pure interruption", () => {
+  const exit = Effect.runSync(Effect.exit(Effect.interrupt))
+  if (Exit.isSuccess(exit)) throw new Error("expected interruption")
+  expect(isInterrupted(exit.cause)).toBe(true)
+  expect(shouldReportPromptFailure(exit.cause)).toBe(false)
+  expect(isInterrupted(Cause.die(new Error("failure")))).toBe(false)
+  expect(shouldReportPromptFailure(Cause.die(new Error("failure")))).toBe(true)
+  expect(shouldReportPromptFailure(Cause.combine(exit.cause, Cause.die(new Error("failure"))))).toBe(true)
+})

packages/opencode/test/kilocode/model-cache-effect.test.ts

@@ -1,11 +1,11 @@
 // kilocode_change - new file
 import { expect } from "bun:test"
-import { Deferred, Effect, Fiber, Layer, Ref } from "effect"
+import { Deferred, Effect, Exit, Fiber, Layer, Option, Ref } from "effect"
 import { HttpClient, HttpClientResponse } from "effect/unstable/http"
 import { Auth } from "../../src/auth"
 import { ModelCache } from "../../src/provider/model-cache"
 import { TestConfig } from "../fixture/config"
-import { testEffect } from "../lib/effect"
+import { pollWithTimeout, testEffect } from "../lib/effect"
 
 type Hit = { readonly url: string }
 
@@ -19,19 +19,20 @@ function layer(
   hits: Ref.Ref<Hit[]>,
   cfg = TestConfig.layer(),
   access = auth,
-  gates?: { readonly started: Deferred.Deferred<void>; readonly wait: Deferred.Deferred<void> },
+  gates?: { readonly started: Deferred.Deferred<void>; readonly wait: Deferred.Deferred<void>; readonly count?: number },
+  fail?: number,
 ) {
   const http = HttpClient.make((request) =>
     Effect.gen(function* () {
       yield* Ref.update(hits, (list) => [...list, { url: request.url }])
       const count = (yield* Ref.get(hits)).length
-      if (gates && count === 1) {
+      if (gates && count === (gates.count ?? 1)) {
         yield* Deferred.succeed(gates.started, undefined)
         yield* Deferred.await(gates.wait)
       }
       return HttpClientResponse.fromWeb(
         request,
-        Response.json({ data: [{ id: `apertis-${count}`, owned_by: "apertis" }] }),
+        Response.json(count === fail ? null : { data: [{ id: `apertis-${count}`, owned_by: "apertis" }] }),
       )
     }),
   )
@@ -76,6 +77,96 @@ it.live("reuses cached values and refresh invalidates the provider cell", () =>
   }),
 )
 
+it.live("retries after a failed refresh", () =>
+  Effect.gen(function* () {
+    const hits = yield* Ref.make<Hit[]>([])
+    const out = yield* ModelCache.Service.use((cache) =>
+      Effect.gen(function* () {
+        const failed = yield* cache.fetch("apertis", { apiKey: "test-key" }).pipe(Effect.exit)
+        const models = yield* cache.fetch("apertis", { apiKey: "test-key" })
+        return { failed, models }
+      }),
+    ).pipe(Effect.provide(layer(hits, TestConfig.layer(), auth, undefined, 1)))
+
+    expect(Exit.isFailure(out.failed)).toBe(true)
+    expect(Object.keys(out.models)).toEqual(["apertis-2"])
+    expect((yield* Ref.get(hits)).length).toBe(2)
+  }),
+)
+
+it.live("keeps a shared refresh alive when one waiter times out", () =>
+  Effect.gen(function* () {
+    const hits = yield* Ref.make<Hit[]>([])
+    const started = yield* Deferred.make<void>()
+    const wait = yield* Deferred.make<void>()
+    const out = yield* ModelCache.Service.use((cache) =>
+      Effect.gen(function* () {
+        const first = yield* cache
+          .fetch("apertis", { apiKey: "test-key" })
+          .pipe(Effect.timeoutOption("10 millis"), Effect.forkChild)
+        yield* Deferred.await(started)
+        const second = yield* cache.fetch("apertis", { apiKey: "test-key" }).pipe(Effect.forkChild)
+        expect(Option.isNone(yield* Fiber.join(first))).toBe(true)
+        yield* Deferred.succeed(wait, undefined)
+        const models = yield* Fiber.join(second)
+        return { models, cached: yield* cache.get("apertis") }
+      }),
+    ).pipe(Effect.provide(layer(hits, TestConfig.layer(), auth, { started, wait })))
+
+    expect(Object.keys(out.models)).toEqual(["apertis-1"])
+    expect(out.cached).toEqual(out.models)
+    expect((yield* Ref.get(hits)).length).toBe(1)
+  }),
+)
+
+it.live("commits a refresh after its only waiter times out", () =>
+  Effect.gen(function* () {
+    const hits = yield* Ref.make<Hit[]>([])
+    const started = yield* Deferred.make<void>()
+    const wait = yield* Deferred.make<void>()
+    const cached = yield* ModelCache.Service.use((cache) =>
+      Effect.gen(function* () {
+        const caller = yield* cache
+          .fetch("apertis", { apiKey: "test-key" })
+          .pipe(Effect.timeoutOption("10 millis"), Effect.forkChild)
+        yield* Deferred.await(started)
+        expect(Option.isNone(yield* Fiber.join(caller))).toBe(true)
+        yield* Deferred.succeed(wait, undefined)
+        return yield* pollWithTimeout(
+          cache.get("apertis"),
+          "service-owned refresh did not commit after its waiter timed out",
+        )
+      }),
+    ).pipe(Effect.provide(layer(hits, TestConfig.layer(), auth, { started, wait })))
+
+    expect(Object.keys(cached)).toEqual(["apertis-1"])
+    expect((yield* Ref.get(hits)).length).toBe(1)
+  }),
+)
+
+it.live("deduplicates overlapping refresh calls", () =>
+  Effect.gen(function* () {
+    const hits = yield* Ref.make<Hit[]>([])
+    const started = yield* Deferred.make<void>()
+    const wait = yield* Deferred.make<void>()
+    const out = yield* ModelCache.Service.use((cache) =>
+      Effect.gen(function* () {
+        yield* cache.fetch("apertis", { apiKey: "test-key" })
+        const first = yield* cache.refresh("apertis", { apiKey: "test-key" }).pipe(Effect.forkChild)
+        yield* Deferred.await(started)
+        const second = yield* cache.refresh("apertis", { apiKey: "test-key" }).pipe(Effect.forkChild)
+        yield* Effect.yieldNow
+        yield* Deferred.succeed(wait, undefined)
+        return { first: yield* Fiber.join(first), second: yield* Fiber.join(second) }
+      }),
+    ).pipe(Effect.provide(layer(hits, TestConfig.layer(), auth, { started, wait, count: 2 })))
+
+    expect(Object.keys(out.first)).toEqual(["apertis-2"])
+    expect(out.second).toEqual(out.first)
+    expect((yield* Ref.get(hits)).length).toBe(2)
+  }),
+)
+
 it.live("keeps concurrent request options isolated", () =>
   Effect.gen(function* () {
     const hits = yield* Ref.make<Hit[]>([])
@@ -131,6 +222,34 @@ it.live("does not let an older fetch override a newer refresh", () =>
   }),
 )
 
+it.live("promotes a cached result after a newer option load fails", () =>
+  Effect.gen(function* () {
+    const hits = yield* Ref.make<Hit[]>([])
+    const started = yield* Deferred.make<void>()
+    const wait = yield* Deferred.make<void>()
+    const out = yield* ModelCache.Service.use((cache) =>
+      Effect.gen(function* () {
+        const first = yield* cache
+          .fetch("apertis", { apiKey: "first", baseURL: "https://first.test/v1" })
+          .pipe(Effect.forkChild)
+        yield* Deferred.await(started)
+        const failed = yield* cache
+          .fetch("apertis", { apiKey: "second", baseURL: "https://second.test/v1" })
+          .pipe(Effect.exit)
+        yield* Deferred.succeed(wait, undefined)
+        yield* Fiber.join(first)
+        const models = yield* cache.fetch("apertis", { apiKey: "first", baseURL: "https://first.test/v1" })
+        return { failed, models, current: yield* cache.get("apertis") }
+      }),
+    ).pipe(Effect.provide(layer(hits, TestConfig.layer(), auth, { started, wait }, 2)))
+
+    expect(Exit.isFailure(out.failed)).toBe(true)
+    expect(Object.keys(out.models)).toEqual(["apertis-1"])
+    expect(out.current).toEqual(out.models)
+    expect((yield* Ref.get(hits)).length).toBe(2)
+  }),
+)
+
 it.live("does not restore a fetch that was cleared while pending", () =>
   Effect.gen(function* () {
     const hits = yield* Ref.make<Hit[]>([])