diff --git a/packages/opencode/src/kilo-sessions/inflight-cache.ts b/packages/opencode/src/kilo-sessions/inflight-cache.ts new file mode 100644 index 0000000000..90ac6d8e66 --- /dev/null +++ b/packages/opencode/src/kilo-sessions/inflight-cache.ts @@ -0,0 +1,71 @@ +type Entry = { + at: number + value: T | undefined + inflight: Promise | undefined +} + +const store = new Map>() + +export function withInFlightCache( + key: string, + ttlMs: number, + cb: () => Promise, +): Promise { + const now = Date.now() + const existing = store.get(key) as Entry | undefined + + if (existing) { + // If a refresh is in-flight, always await it. + // This avoids returning a stale cached value while a newer one is being computed. + if (existing.inflight) return existing.inflight + + // `undefined` means "no value" (and is never cached). + if (existing.value !== undefined && now - existing.at < ttlMs) return Promise.resolve(existing.value) + } + + const next: Entry = existing + ? { + // Keep the original timestamp until the refresh succeeds. + // Otherwise, a failed refresh could make an old value look "fresh" and suppress retries. + at: existing.at, + value: existing.value, + inflight: undefined, + } + : { + at: now, + value: undefined, + inflight: undefined, + } + + // Guard against synchronous throws in `cb()` by forcing it into the promise chain. + // This ensures errors flow through the `.catch()` below and the cache entry is cleaned up. + const task = Promise.resolve().then(() => cb()) + .then((value) => { + if (value === undefined) { + // `undefined` is treated as a non-cacheable sentinel. + // Drop the entry entirely so future calls retry instead of serving stale data. + store.delete(key) + return undefined + } + + next.value = value + next.at = Date.now() + return value + }) + .catch((error) => { + store.delete(key) + throw error + }) + + const inflight = task.finally(() => { + next.inflight = undefined + }) + + next.inflight = inflight + store.set(key, next as Entry) + return inflight +} + +export function clearInFlightCache(key: string) { + store.delete(key) +} diff --git a/packages/opencode/src/kilo-sessions/kilo-sessions.ts b/packages/opencode/src/kilo-sessions/kilo-sessions.ts index 5de42889ab..ded76e59a6 100644 --- a/packages/opencode/src/kilo-sessions/kilo-sessions.ts +++ b/packages/opencode/src/kilo-sessions/kilo-sessions.ts @@ -6,43 +6,64 @@ import { Storage } from "@/storage/storage" import { Log } from "@/util/log" import { Auth } from "@/auth" import { IngestQueue } from "@/kilo-sessions/ingest-queue" +import { clearInFlightCache, withInFlightCache } from "@/kilo-sessions/inflight-cache" import type * as SDK from "@kilocode/sdk/v2" +import z from "zod" export namespace KiloSessions { - const log = Log.create({ service: "share-next" }) + const log = Log.create({ service: "kilo-sessions" }) - const authCache = new Map() + const Uuid = z.uuid() + type Uuid = z.infer - const orgCache = { - at: 0, - value: undefined as string | undefined, - inflight: undefined as Promise | undefined, + const tokenValidKeyTemplate = "kilo-sessions:token-valid:" + let tokenValidKey = tokenValidKeyTemplate + "unknown" + + const tokenKey = "kilo-sessions:token" + const orgKey = "kilo-sessions:org" + const clientKey = "kilo-sessions:client" + + const ttlMs = 10_000 + + function clearCache() { + clearInFlightCache(tokenKey) + clearInFlightCache(tokenValidKey) + clearInFlightCache(clientKey) + clearInFlightCache(orgKey) } async function authValid(token: string) { - const cached = authCache.get(token) - if (cached) return cached.valid + const newTokenValidKey = tokenValidKeyTemplate + token - const response = await fetch("https://app.kilo.ai/api/user", { - headers: { - Authorization: `Bearer ${token}`, - }, - }).catch(() => undefined) + if (newTokenValidKey !== tokenValidKey) { + clearInFlightCache(tokenValidKey) - // Don't cache transient network failures; allow future calls to retry. - if (!response) return false + tokenValidKey = newTokenValidKey + } - const valid = response.ok - authCache.set(token, { valid }) - return valid + return withInFlightCache(tokenValidKey, 15 * 60_000, async () => { + const response = await fetch("https://app.kilo.ai/api/user", { + headers: { + Authorization: `Bearer ${token}`, + }, + }).catch(() => undefined) + + // Don't cache transient network failures; allow future calls to retry. + if (!response) return undefined + + const valid = response.ok + return valid + }) } - export async function kilocodeToken() { - const auth = await Auth.get("kilo") - if (auth?.type === "api" && auth.key.length > 0) return auth.key - if (auth?.type === "oauth" && auth.access.length > 0) return auth.access - if (auth?.type === "wellknown" && auth.token.length > 0) return auth.token - return undefined + async function kilocodeToken() { + return withInFlightCache(tokenKey, ttlMs, async () => { + const auth = await Auth.get("kilo") + if (auth?.type === "api" && auth.key.length > 0) return auth.key + if (auth?.type === "oauth" && auth.access.length > 0) return auth.access + if (auth?.type === "wellknown" && auth.token.length > 0) return auth.token + return undefined + }) } type Client = { @@ -50,19 +71,8 @@ export namespace KiloSessions { fetch: (input: RequestInfo | URL, init?: RequestInit) => Promise } - const cache = { - at: 0, - value: undefined as Client | undefined, - inflight: undefined as Promise | undefined, - } - async function getClient(): Promise { - const now = Date.now() - if (cache.value && now - cache.at < 5_000) return cache.value - if (cache.inflight && now - cache.at < 5_000) return cache.inflight - - cache.at = now - cache.inflight = (async () => { + return withInFlightCache(clientKey, ttlMs, async () => { const token = await kilocodeToken() if (!token) return undefined @@ -88,14 +98,7 @@ export namespace KiloSessions { url: base, fetch: (input, init) => fetch(input, withHeaders(init)), } - })() - - try { - cache.value = await cache.inflight - return cache.value - } finally { - cache.inflight = undefined - } + }) } const ingest = IngestQueue.create({ @@ -105,10 +108,7 @@ export namespace KiloSessions { onAuthError: () => { // Non-retryable until credentials are fixed. // Clearing caches prevents repeated use of a now-invalid token/client. - authCache.clear() - cache.value = undefined - cache.inflight = undefined - cache.at = 0 + clearCache() }, }) @@ -371,27 +371,19 @@ export namespace KiloSessions { } } - async function getOrgId(): Promise { + async function getOrgId(): Promise { const env = process.env["KILO_ORG_ID"] - if (env) return env + if (isUuid(env)) return env - const now = Date.now() - if (orgCache.value && now - orgCache.at < 5_000) return orgCache.value - if (orgCache.inflight && now - orgCache.at < 5_000) return orgCache.inflight - - orgCache.at = now - orgCache.inflight = (async () => { + return withInFlightCache(orgKey, ttlMs, async () => { const auth = await Auth.get("kilo") - if (auth?.type === "oauth" && auth.accountId) return auth.accountId - + if (auth?.type === "oauth" && isUuid(auth.accountId)) return auth.accountId return undefined - })() + }) + } - try { - orgCache.value = await orgCache.inflight - return orgCache.value - } finally { - orgCache.inflight = undefined - } + function isUuid(value: string | undefined): value is Uuid { + if (!value) return false + return Uuid.safeParse(value).success } } diff --git a/packages/opencode/test/kilo-sessions/inflight-cache.test.ts b/packages/opencode/test/kilo-sessions/inflight-cache.test.ts new file mode 100644 index 0000000000..5927b5657a --- /dev/null +++ b/packages/opencode/test/kilo-sessions/inflight-cache.test.ts @@ -0,0 +1,157 @@ +import { afterEach, describe, expect, test } from "bun:test" +import { clearInFlightCache, withInFlightCache } from "../../src/kilo-sessions/inflight-cache" + +function deferred() { + const state = { + resolve: undefined as ((value: T) => void) | undefined, + reject: undefined as ((reason?: unknown) => void) | undefined, + } + + const promise = new Promise((resolve, reject) => { + state.resolve = resolve + state.reject = reject + }) + + return { promise, resolve: state.resolve!, reject: state.reject! } +} + +describe("withInFlightCache", () => { + const realNow = Date.now + const clock = { now: 0 } + + afterEach(() => { + Date.now = realNow + }) + + test("dedupes concurrent calls (inflight)", async () => { + Date.now = () => clock.now + const key = "inflight-dedupe" + clearInFlightCache(key) + + const job = deferred() + const calls = { count: 0 } + const cb = () => { + calls.count += 1 + return job.promise + } + + const a = withInFlightCache(key, 10_000, cb) + const b = withInFlightCache(key, 10_000, cb) + + expect(a).toBe(b) + // cb() is invoked via a microtask in withInFlightCache(). + // Allow it to run before asserting call count. + await Promise.resolve() + expect(calls.count).toBe(1) + + job.resolve(123) + expect(await a).toBe(123) + expect(await b).toBe(123) + + clearInFlightCache(key) + }) + + test("returns cached value within ttl without calling cb again", async () => { + Date.now = () => clock.now + const key = "cache-hit" + clearInFlightCache(key) + + const calls = { count: 0 } + const cb = async () => { + calls.count += 1 + return "ok" + } + + expect(await withInFlightCache(key, 10_000, cb)).toBe("ok") + expect(calls.count).toBe(1) + + clock.now += 5 + expect(await withInFlightCache(key, 10_000, cb)).toBe("ok") + expect(calls.count).toBe(1) + + clearInFlightCache(key) + }) + + test("does not cache undefined (treats as no value)", async () => { + Date.now = () => clock.now + const key = "cache-undefined" + clearInFlightCache(key) + + const calls = { count: 0 } + const cb = async () => { + calls.count += 1 + return undefined + } + + expect(await withInFlightCache(key, 10_000, cb)).toBeUndefined() + expect(calls.count).toBe(1) + + clock.now += 5 + expect(await withInFlightCache(key, 10_000, cb)).toBeUndefined() + expect(calls.count).toBe(2) + + clearInFlightCache(key) + }) + + test("expires cached entries after ttl and recomputes", async () => { + Date.now = () => clock.now + const key = "ttl-expire" + clearInFlightCache(key) + + const calls = { count: 0 } + const cb = async () => { + calls.count += 1 + return calls.count + } + + expect(await withInFlightCache(key, 10, cb)).toBe(1) + clock.now += 11 + expect(await withInFlightCache(key, 10, cb)).toBe(2) + + clearInFlightCache(key) + }) + + test("does not return stale cached value while refresh is in-flight", async () => { + Date.now = () => clock.now + const key = "refresh-inflight" + clearInFlightCache(key) + + // Seed a cached value. + expect(await withInFlightCache(key, 10, async () => 1)).toBe(1) + + // Expire it. + clock.now += 11 + + // Start a refresh that we can hold. + const job = deferred() + const refresh = withInFlightCache(key, 10, () => job.promise) + + // Subsequent callers must await refresh, not get the old cached value. + const follower = withInFlightCache(key, 10, async () => 999) + expect(follower).toBe(refresh) + + job.resolve(2) + expect(await refresh).toBe(2) + expect(await follower).toBe(2) + + clearInFlightCache(key) + }) + + test("clearInFlightCache forces recompute", async () => { + Date.now = () => clock.now + const key = "clear" + clearInFlightCache(key) + + const calls = { count: 0 } + const cb = async () => { + calls.count += 1 + return calls.count + } + + expect(await withInFlightCache(key, 10_000, cb)).toBe(1) + clearInFlightCache(key) + expect(await withInFlightCache(key, 10_000, cb)).toBe(2) + + clearInFlightCache(key) + }) +}) diff --git a/packages/opencode/test/share/ingest-queue.test.ts b/packages/opencode/test/kilo-sessions/ingest-queue.test.ts similarity index 99% rename from packages/opencode/test/share/ingest-queue.test.ts rename to packages/opencode/test/kilo-sessions/ingest-queue.test.ts index 939a891d8a..7819c87d56 100644 --- a/packages/opencode/test/share/ingest-queue.test.ts +++ b/packages/opencode/test/kilo-sessions/ingest-queue.test.ts @@ -1,4 +1,3 @@ -// kilocode_change - new file import { describe, expect, test, beforeEach } from "bun:test" import { IngestQueue } from "../../src/kilo-sessions/ingest-queue"