Merge pull request #72 from Kilo-Org/add-support-for-missing-cli-session-metadata

packages/opencode/src/kilo-sessions/ingest-queue.ts

@@ -1,4 +1,3 @@
-// kilocode_change - new file
 import { ulid } from "ulid"
 import type * as SDK from "@kilocode/sdk/v2"
 
@@ -9,6 +8,13 @@ export namespace IngestQueue {
   }
 
   export type Data =
+    | {
+        type: "kilo_meta"
+        data: {
+          platform: string
+          orgId?: string
+        }
+      }
     | {
         type: "session"
         data: SDK.Session
@@ -98,6 +104,7 @@ export namespace IngestQueue {
     function key(item: Data) {
       // Stable keys are important so updates for the same entity collapse to a single queued item.
       // If we can't derive a stable key, we fall back to a random key (ulid) so the item is still sent.
+      if (item.type === "kilo_meta") return "kilo_meta"
       if (item.type === "session") return "session"
       if (item.type === "session_diff") return "session_diff"
 

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

@@ -1,4 +1,3 @@
-// kilocode_change pretty much completely refactored - @iscekic for conflicts
 import { Bus } from "@/bus"
 import { Provider } from "@/provider/provider"
 import { Session } from "@/session"
@@ -6,17 +5,20 @@ import { MessageV2 } from "@/session/message-v2"
 import { Storage } from "@/storage/storage"
 import { Log } from "@/util/log"
 import { Auth } from "@/auth"
-import { IngestQueue } from "@/share/ingest-queue" // kilocode_change
+import { IngestQueue } from "@/kilo-sessions/ingest-queue"
 import type * as SDK from "@kilocode/sdk/v2"
 
-/**
- * Even though this is called "share-next", this is where we handle session stuff.
- */
-export namespace ShareNext {
+export namespace KiloSessions {
   const log = Log.create({ service: "share-next" })
 
   const authCache = new Map<string, { valid: boolean }>()
 
+  const orgCache = {
+    at: 0,
+    value: undefined as string | undefined,
+    inflight: undefined as Promise<string | undefined> | undefined,
+  }
+
   async function authValid(token: string) {
     const cached = authCache.get(token)
     if (cached) return cached.valid
@@ -124,6 +126,10 @@ export namespace ShareNext {
 
     Bus.subscribe(Session.Event.Updated, async (evt) => {
       await ingest.sync(evt.properties.info.id, [
+        {
+          type: "kilo_meta",
+          data: await meta(),
+        },
         {
           type: "session",
           data: evt.properties.info,
@@ -331,6 +337,10 @@ export namespace ShareNext {
     )
 
     await ingest.sync(sessionId, [
+      {
+        type: "kilo_meta",
+        data: await meta(),
+      },
       {
         type: "session",
         data: session,
@@ -350,4 +360,38 @@ export namespace ShareNext {
       },
     ])
   }
+
+  async function meta() {
+    const platform = process.env["KILO_PLATFORM"] || "cli"
+    const orgId = await getOrgId()
+
+    return {
+      platform,
+      ...(orgId ? { orgId } : {}),
+    }
+  }
+
+  async function getOrgId(): Promise<string | undefined> {
+    const env = process.env["KILO_ORG_ID"]
+    if (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 () => {
+      const auth = await Auth.get("kilo")
+      if (auth?.type === "oauth" && auth.accountId) return auth.accountId
+
+      return undefined
+    })()
+
+    try {
+      orgCache.value = await orgCache.inflight
+      return orgCache.value
+    } finally {
+      orgCache.inflight = undefined
+    }
+  }
 }

packages/opencode/src/project/bootstrap.ts

@@ -9,14 +9,14 @@ import { Command } from "../command"
 import { Instance } from "./instance"
 import { Vcs } from "./vcs"
 import { Log } from "@/util/log"
-import { ShareNext } from "@/share/share-next" // kilocode_change
+import { KiloSessions } from "@/kilo-sessions/kilo-sessions" // kilocode_change
 import { Snapshot } from "../snapshot"
 import { Truncate } from "../tool/truncation"
 
 export async function InstanceBootstrap() {
   Log.Default.info("bootstrapping", { directory: Instance.directory })
   await Plugin.init()
-  ShareNext.init() // kilocode_change
+  KiloSessions.init() // kilocode_change
   Format.init()
   await LSP.init()
   FileWatcher.init()

packages/opencode/src/session/index.ts

@@ -253,8 +253,8 @@ export namespace Session {
     if (cfg.share === "disabled") {
       throw new Error("Sharing is disabled in configuration")
     }
-    const { ShareNext } = await import("@/share/share-next")
-    const share = await ShareNext.share(id) // kilocode_change
+    const { KiloSessions } = await import("@/kilo-sessions/kilo-sessions")
+    const share = await KiloSessions.share(id) // kilocode_change
     await update(
       id,
       (draft) => {
@@ -268,9 +268,8 @@ export namespace Session {
   })
 
   export const unshare = fn(Identifier.schema("session"), async (id) => {
-    // Use ShareNext to remove the share (same as share function uses ShareNext to create)
-    const { ShareNext } = await import("@/share/share-next")
-    await ShareNext.unshare(id) // kilocode_change
+    const { KiloSessions } = await import("@/kilo-sessions/kilo-sessions")
+    await KiloSessions.unshare(id) // kilocode_change
     await update(
       id,
       (draft) => {
@@ -340,8 +339,8 @@ export namespace Session {
       for (const child of await children(sessionID)) {
         await remove(child.id)
       }
-      const { ShareNext } = await import("@/share/share-next")
-      await ShareNext.remove(sessionID).catch(() => {}) // kilocode_change
+      const { KiloSessions } = await import("@/kilo-sessions/kilo-sessions")
+      await KiloSessions.remove(sessionID).catch(() => {}) // kilocode_change
       for (const msg of await Storage.list(["message", sessionID])) {
         for (const part of await Storage.list(["part", msg.at(-1)!])) {
           await Storage.remove(part)

packages/opencode/test/share/ingest-queue.test.ts

@@ -1,6 +1,6 @@
 // kilocode_change - new file
 import { describe, expect, test, beforeEach } from "bun:test"
-import { IngestQueue } from "../../src/share/ingest-queue"
+import { IngestQueue } from "../../src/kilo-sessions/ingest-queue"
 
 function scheduler(now: () => number) {
   const tasks = new Map<number, { at: number; fn: () => void }>()
@@ -119,6 +119,39 @@ describe("share ingest queue", () => {
     expect((sent[0] as any).data[0].data.v).toBe(2)
   })
 
+  test("kilo_meta uses stable key and coalesces", async () => {
+    const sent: unknown[] = []
+    const sched = scheduler(() => clock.now)
+
+    const q = IngestQueue.create({
+      now: () => clock.now,
+      setTimeout: sched.setTimeout,
+      clearTimeout: sched.clearTimeout,
+      log: { error: () => {} },
+      getShare: async () => ({ ingestPath: "/ingest" }),
+      getClient: async () => ({
+        url: "https://ingest.test",
+        fetch: async (_input, init) => {
+          sent.push(JSON.parse((init?.body as string) ?? "{}"))
+          return new Response("{}", { status: 200 })
+        },
+      }),
+    })
+
+    await q.sync("s7", [{ type: "kilo_meta", data: { platform: "cli" } }])
+    clock.now = 100
+    await q.sync("s7", [{ type: "kilo_meta", data: { platform: "vscode", orgId: "org-1" } }])
+
+    clock.now = 1000
+    sched.run()
+    await Bun.sleep(0)
+    expect(sent.length).toBe(1)
+    expect((sent[0] as any).data.length).toBe(1)
+    expect((sent[0] as any).data[0].type).toBe("kilo_meta")
+    expect((sent[0] as any).data[0].data.platform).toBe("vscode")
+    expect((sent[0] as any).data[0].data.orgId).toBe("org-1")
+  })
+
   test("network failure retries and fill preserves newer updates", async () => {
     const sent: unknown[] = []
     const sched = scheduler(() => clock.now)