Merge pull request #73 from Kilo-Org/add-some-defensive-code-around-getting-org-id

packages/opencode/src/kilo-sessions/inflight-cache.ts

@@ -0,0 +1,71 @@
+type Entry<T> = {
+  at: number
+  value: T | undefined
+  inflight: Promise<T | undefined> | undefined
+}
+
+const store = new Map<string, Entry<unknown>>()
+
+export function withInFlightCache<T>(
+  key: string,
+  ttlMs: number,
+  cb: () => Promise<T | undefined>,
+): Promise<T | undefined> {
+  const now = Date.now()
+  const existing = store.get(key) as Entry<T> | 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<T> = 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<unknown>)
+  return inflight
+}
+
+export function clearInFlightCache(key: string) {
+  store.delete(key)
+}

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<string, { valid: boolean }>()
+  const Uuid = z.uuid()
+  type Uuid = z.infer<typeof Uuid>
 
-  const orgCache = {
-    at: 0,
-    value: undefined as string | undefined,
-    inflight: undefined as Promise<string | undefined> | 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<Response>
   }
 
-  const cache = {
-    at: 0,
-    value: undefined as Client | undefined,
-    inflight: undefined as Promise<Client | undefined> | undefined,
-  }
-
   async function getClient(): Promise<Client | undefined> {
-    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<string | undefined> {
+  async function getOrgId(): Promise<Uuid | undefined> {
     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
   }
 }

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<T>() {
+  const state = {
+    resolve: undefined as ((value: T) => void) | undefined,
+    reject: undefined as ((reason?: unknown) => void) | undefined,
+  }
+
+  const promise = new Promise<T>((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<number>()
+    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<number>()
+    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)
+  })
+})

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"