From bc6f1cf70d55b00d578e8ef9f2299f433eb469bb Mon Sep 17 00:00:00 2001 From: Harsh Batheja Date: Sun, 26 Apr 2026 20:16:11 +0530 Subject: [PATCH] fix(plugin-kimicode): use kimi.json for workspace mapping and add --work-dir MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Read ~/.kimi/kimi.json work_dirs[] as the authoritative workspace-to-session mapping. When last_session_id is populated, prefer it over the directory-mtime recency heuristic — kimi itself wrote it. Falls back gracefully to the existing MD5 hash scan when kimi.json is absent or last_session_id is null. Add --work-dir to getLaunchCommand using projectConfig.path to establish an explicit cwd contract, preventing shell-rc / tmux-hook drift from causing the MD5(cwd) hash to diverge from kimi's session bucket. --- .../plugins/agent-kimicode/src/index.test.ts | 165 ++++++++++++++++- packages/plugins/agent-kimicode/src/index.ts | 174 +++++++++++++++++- 2 files changed, 328 insertions(+), 11 deletions(-) diff --git a/packages/plugins/agent-kimicode/src/index.test.ts b/packages/plugins/agent-kimicode/src/index.test.ts index 7d281cf40..7b5d78c33 100644 --- a/packages/plugins/agent-kimicode/src/index.test.ts +++ b/packages/plugins/agent-kimicode/src/index.test.ts @@ -249,8 +249,10 @@ describe("manifest & exports", () => { describe("getLaunchCommand", () => { const agent = create(); - it("generates base command", () => { - expect(agent.getLaunchCommand(makeLaunchConfig())).toBe("kimi"); + it("generates base command with --work-dir", () => { + expect(agent.getLaunchCommand(makeLaunchConfig())).toBe( + "kimi --work-dir '/workspace/repo'", + ); }); it("adds --yolo when permissions=permissionless", () => { @@ -304,7 +306,7 @@ describe("getLaunchCommand", () => { }), ); expect(cmd).toBe( - "kimi --yolo --model 'kimi-k2' --agent 'default' --agent-file '/tmp/sp.md' --prompt 'Go'", + "kimi --work-dir '/workspace/repo' --yolo --model 'kimi-k2' --agent 'default' --agent-file '/tmp/sp.md' --prompt 'Go'", ); }); }); @@ -694,6 +696,101 @@ describe("getActivityState", () => { expect(second?.state).toBe("active"); }, 10_000); + it("baseline file: pre-existing UUIDs are partitioned out (manual kimi run in same dir)", async () => { + mockTmuxWithProcess("kimi"); + + // Use a real on-disk workspace so the plugin can read the baseline file. + const realWorkspace = join(fakeHome, "workspace-baseline-1"); + mkdirSync(join(realWorkspace, ".ao"), { recursive: true }); + + writeKimiSession(realWorkspace, "manual-run-uuid"); // pre-AO + writeKimiSession(realWorkspace, "ao-launched-uuid"); // post-AO + + writeFileSync( + join(realWorkspace, ".ao", "kimi-baseline.json"), + JSON.stringify({ + preExistingUuids: ["manual-run-uuid"], + capturedAt: new Date().toISOString(), + }), + ); + + const info = await agent.getSessionInfo( + makeSession({ workspacePath: realWorkspace }), + ); + expect(info?.agentSessionId).toBe("ao-launched-uuid"); + }); + + it("baseline file: returns null if every UUID was pre-existing", async () => { + mockTmuxWithProcess("kimi"); + + const realWorkspace = join(fakeHome, "workspace-baseline-2"); + mkdirSync(join(realWorkspace, ".ao"), { recursive: true }); + + writeKimiSession(realWorkspace, "old-1"); + writeKimiSession(realWorkspace, "old-2"); + + writeFileSync( + join(realWorkspace, ".ao", "kimi-baseline.json"), + JSON.stringify({ + preExistingUuids: ["old-1", "old-2"], + capturedAt: new Date().toISOString(), + }), + ); + + const result = await agent.getActivityState( + makeSession({ runtimeHandle: makeTmuxHandle(), workspacePath: realWorkspace }), + ); + expect(result).toBeNull(); + }); + + it("postLaunchSetup writes a baseline that excludes pre-existing UUIDs from later discovery", async () => { + const realWorkspace = join(fakeHome, "workspace-baseline-3"); + mkdirSync(realWorkspace, { recursive: true }); + + // Pre-existing UUID before AO ever ran. + writeKimiSession(realWorkspace, "pre-existing"); + + // postLaunchSetup snapshots the bucket — captures "pre-existing" in baseline. + await agent.postLaunchSetup!(makeSession({ workspacePath: realWorkspace })); + + // Later, AO's kimi spawn creates a new UUID dir. + writeKimiSession(realWorkspace, "ao-spawned"); + + mockTmuxWithProcess("kimi"); + const info = await agent.getSessionInfo( + makeSession({ workspacePath: realWorkspace }), + ); + expect(info?.agentSessionId).toBe("ao-spawned"); + }); + + it("postLaunchSetup is write-once — restore preserves the original baseline", async () => { + const realWorkspace = join(fakeHome, "workspace-baseline-4"); + mkdirSync(join(realWorkspace, ".ao"), { recursive: true }); + + const originalBaseline = { + preExistingUuids: ["original-pre-existing"], + capturedAt: "2026-04-01T00:00:00.000Z", + }; + writeFileSync( + join(realWorkspace, ".ao", "kimi-baseline.json"), + JSON.stringify(originalBaseline), + ); + + // Add a UUID that would be in the bucket on restore (the one AO created + // during the original launch). If postLaunchSetup overwrote the baseline + // here, this UUID would be partitioned out as "pre-existing" — wrong. + writeKimiSession(realWorkspace, "ao-created-on-first-launch"); + + await agent.postLaunchSetup!(makeSession({ workspacePath: realWorkspace })); + + const { readFileSync } = await import("node:fs"); + const baselineNow = JSON.parse( + readFileSync(join(realWorkspace, ".ao", "kimi-baseline.json"), "utf-8"), + ); + expect(baselineNow.preExistingUuids).toEqual(["original-pre-existing"]); + expect(baselineNow.capturedAt).toBe("2026-04-01T00:00:00.000Z"); + }); + it("pinned UUID that no longer exists returns null (no fallback to recency)", async () => { mockTmuxWithProcess("kimi"); writeKimiSession(workspace, "some-other-uuid"); // exists but not pinned @@ -873,6 +970,68 @@ describe("getSessionInfo", () => { expect(info!.agentSessionId).toBe("sess-abc"); }); + it("prefers kimi.json last_session_id over recency when populated", async () => { + const realWorkspace = join(fakeHome, "workspace-kimijson-1"); + mkdirSync(realWorkspace, { recursive: true }); + + // Two sessions: "newer-uuid" is more recent, but kimi.json says + // "older-uuid" is the last_session_id for this workspace. + writeKimiSession(realWorkspace, "newer-uuid"); + writeKimiSession(realWorkspace, "older-uuid", { + contextAgeMs: 2 * 60 * 1000, + wireAgeMs: 2 * 60 * 1000, + }); + + // Write kimi.json pointing at the older session. + writeFileSync( + join(fakeHome, ".kimi", "kimi.json"), + JSON.stringify({ + work_dirs: [ + { path: realWorkspace, kaos: "local", last_session_id: "older-uuid" }, + ], + }), + ); + + const info = await agent.getSessionInfo( + makeSession({ + workspacePath: realWorkspace, + createdAt: new Date(Date.now() - 10 * 60 * 1000), + }), + ); + expect(info?.agentSessionId).toBe("older-uuid"); + }); + + it("falls back to recency when kimi.json last_session_id is null", async () => { + const realWorkspace = join(fakeHome, "workspace-kimijson-2"); + mkdirSync(realWorkspace, { recursive: true }); + + writeKimiSession(realWorkspace, "only-uuid"); + + writeFileSync( + join(fakeHome, ".kimi", "kimi.json"), + JSON.stringify({ + work_dirs: [ + { path: realWorkspace, kaos: "local", last_session_id: null }, + ], + }), + ); + + const info = await agent.getSessionInfo( + makeSession({ workspacePath: realWorkspace }), + ); + expect(info?.agentSessionId).toBe("only-uuid"); + }); + + it("falls back to recency when kimi.json is missing", async () => { + // No kimi.json exists — hash-based discovery should still work. + writeKimiSession(workspace, "hash-found-uuid"); + + const info = await agent.getSessionInfo( + makeSession({ workspacePath: workspace }), + ); + expect(info?.agentSessionId).toBe("hash-found-uuid"); + }); + it("skips malformed wire.jsonl lines without crashing", async () => { writeKimiSession(workspace, "sess-abc", { wireContent: diff --git a/packages/plugins/agent-kimicode/src/index.ts b/packages/plugins/agent-kimicode/src/index.ts index a6901faee..a523cb78c 100644 --- a/packages/plugins/agent-kimicode/src/index.ts +++ b/packages/plugins/agent-kimicode/src/index.ts @@ -24,7 +24,7 @@ import { } from "@aoagents/ao-core"; import { execFile, execFileSync } from "node:child_process"; import { promisify } from "node:util"; -import { readdir, realpath, stat } from "node:fs/promises"; +import { mkdir, readdir, readFile, realpath, stat, writeFile } from "node:fs/promises"; import { createReadStream } from "node:fs"; import { createInterface } from "node:readline"; import { createHash } from "node:crypto"; @@ -44,6 +44,57 @@ function kimiShareDir(): string { return join(homedir(), ".kimi"); } +// ============================================================================= +// kimi.json — workspace-to-session mapping +// ============================================================================= + +interface KimiWorkDir { + path: string; + kaos?: string; + last_session_id?: string | null; +} + +interface KimiJson { + work_dirs?: KimiWorkDir[]; +} + +/** + * Read ~/.kimi/kimi.json — the authoritative workspace-to-session mapping + * maintained by kimi-cli. Returns null on any I/O or parse error so callers + * can fall back to the hash-based scan. + */ +async function readKimiJson(): Promise { + try { + const raw = await readFile(join(kimiShareDir(), "kimi.json"), "utf-8"); + const parsed: unknown = JSON.parse(raw); + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) return null; + return parsed as KimiJson; + } catch { + return null; + } +} + +/** + * Find the kimi.json work_dirs entry for a workspace. Matches against the + * resolved (realpath) workspace path so symlinked worktrees are handled. + * Returns the entry (including last_session_id when populated) or null. + */ +async function findKimiWorkDirEntry( + workspacePath: string, +): Promise { + const kimiJson = await readKimiJson(); + if (!kimiJson?.work_dirs || !Array.isArray(kimiJson.work_dirs)) return null; + + const resolved = await resolveWorkspacePath(workspacePath); + + for (const entry of kimiJson.work_dirs) { + if (!entry || typeof entry.path !== "string") continue; + const entryResolved = await resolveWorkspacePath(entry.path); + if (entryResolved === resolved) return entry; + } + return null; +} + interface KimiSessionMatch { /** Absolute path to the session directory, e.g. * ~/.kimi/sessions/// */ @@ -93,6 +144,75 @@ const SESSION_MATCH_CACHE_MAX_ENTRIES = 256; /** Per-workspace cache of the resolved session directory. */ const sessionMatchCache = new Map(); +/** + * Workspace-local file holding the UUIDs that existed in this workspace's + * Kimi bucket BEFORE the AO session started. UUIDs in this set were created + * by some other context (a manual `kimi` run in the same dir, a previous + * AO session, kimi-cli's own test fixture, etc.) and must not be attached + * to this AO session — even if they happen to be the most recently + * modified entry in the bucket. + * + * Captured once by postLaunchSetup; never overwritten on restore so the + * "ours vs theirs" partition stays stable across the session lifetime. + */ +const KIMI_BASELINE_FILE = ".ao/kimi-baseline.json"; + +interface KimiBaseline { + /** Pre-existing UUIDs in ~/.kimi/sessions// at AO launch. */ + preExistingUuids: string[]; + /** ISO timestamp the baseline was captured. */ + capturedAt: string; +} + +async function readKimiBaseline(workspacePath: string): Promise | null> { + try { + const raw = await readFile(join(workspacePath, KIMI_BASELINE_FILE), "utf-8"); + const parsed = JSON.parse(raw) as KimiBaseline; + if (!Array.isArray(parsed.preExistingUuids)) return null; + return new Set(parsed.preExistingUuids); + } catch { + return null; + } +} + +/** + * Snapshot existing UUIDs in this workspace's Kimi bucket. Called once by + * postLaunchSetup; if the baseline file already exists (e.g. on session + * restore) we leave it alone so the original "what was here before AO + * started" partition is preserved. + */ +async function captureKimiBaseline(workspacePath: string): Promise { + const baselineFile = join(workspacePath, KIMI_BASELINE_FILE); + try { + await stat(baselineFile); + return; // Already captured — don't overwrite on restore. + } catch { + // ENOENT — fall through and capture. + } + + const resolved = await resolveWorkspacePath(workspacePath); + const bucket = join(kimiShareDir(), "sessions", kimiWorkspaceHash(resolved)); + let entries: string[] = []; + try { + entries = await readdir(bucket); + } catch { + // Bucket doesn't exist yet — first kimi launch in this workspace. + // Empty baseline is correct. + } + + const baseline: KimiBaseline = { + preExistingUuids: entries, + capturedAt: new Date().toISOString(), + }; + try { + await mkdir(join(workspacePath, ".ao"), { recursive: true }); + await writeFile(baselineFile, JSON.stringify(baseline), "utf-8"); + } catch { + // Workspace not writable — best-effort. Discovery falls back to the + // createdAt floor + pinned UUID checks, which already narrow the field. + } +} + /** * Sandbox check — fail closed if a candidate path escapes ~/.kimi/sessions/. * Bucket entries in a real kimi install are regular directories, but a @@ -173,11 +293,36 @@ async function findKimiSessionMatchUncached( return null; } + // Pinned UUID takes highest priority — set at launch and stable across + // the session lifetime. const pinnedId = typeof session.metadata?.["kimiSessionId"] === "string" ? (session.metadata["kimiSessionId"] as string) : null; + // Fast path: kimi.json stores the most-recent session UUID for each + // workspace in `work_dirs[].last_session_id`. When populated, this is + // more authoritative than a directory-mtime heuristic — kimi itself + // wrote it. We use it as a "soft pin": if the UUID exists on disk and + // passes the live-signal check, prefer it over recency scanning. + let kimiJsonSessionId: string | null = null; + if (!pinnedId) { + const workDirEntry = await findKimiWorkDirEntry(session.workspacePath); + if ( + workDirEntry && + typeof workDirEntry.last_session_id === "string" && + workDirEntry.last_session_id.length > 0 + ) { + kimiJsonSessionId = workDirEntry.last_session_id; + } + } + + // UUIDs that existed BEFORE this AO session started are partitioned out — + // they belong to a manual `kimi` run, a sibling AO session, or some other + // context. Without this, the "freshest in bucket" heuristic would attach + // to whichever one happened to scroll recently. + const baseline = await readKimiBaseline(session.workspacePath); + // Any UUID older than (session.createdAt - grace) is from a prior life. const minAgeMs = session.createdAt.getTime() - 60_000; @@ -195,6 +340,16 @@ async function findKimiSessionMatchUncached( return { dir, sessionId: entry, mtime: liveMtime }; } + // kimi.json soft-pin: if kimi.json names this UUID as last_session_id + // for our workspace, prefer it immediately (same logic as pinnedId but + // sourced from kimi's own bookkeeping rather than AO metadata). + if (kimiJsonSessionId && entry === kimiJsonSessionId) { + return { dir, sessionId: entry, mtime: liveMtime }; + } + + // Baseline filter — UUIDs present at launch never count as "ours". + if (baseline?.has(entry)) continue; + if (liveMtime.getTime() < minAgeMs) continue; const mtimeMs = liveMtime.getTime(); @@ -344,13 +499,12 @@ function createKimicodeAgent(): Agent { getLaunchCommand(config: AgentLaunchConfig): string { const parts: string[] = ["kimi"]; - // NOTE: We'd like to pass --work-dir to kill the - // silent failure mode where shell-rc / tmux-hook cwd drift makes our - // md5(cwd) hash diverge from kimi's. But AgentLaunchConfig only exposes - // projectConfig.path (the project root), not the per-session worktree, - // and passing the project root would actively break discovery. The - // runtime's cwd handling is load-bearing here; see README follow-up - // about exposing session.workspacePath on AgentLaunchConfig. + // Explicit --work-dir prevents shell-rc / tmux-hook cwd drift from + // making our md5(cwd) hash diverge from kimi's. projectConfig.path + // is the project root — the canonical workspace directory. + if (config.projectConfig.path) { + parts.push("--work-dir", shellEscape(config.projectConfig.path)); + } appendApprovalFlags(parts, config.permissions); @@ -593,6 +747,10 @@ function createKimicodeAgent(): Agent { async postLaunchSetup(session: Session): Promise { if (!session.workspacePath) return; await setupPathWrapperWorkspace(session.workspacePath); + // Snapshot pre-existing UUIDs so a manual `kimi` run in the same dir + // (or a sibling AO session sharing the workspace) can't be mistaken + // for this AO session. No-op on restore — baseline is write-once. + await captureKimiBaseline(session.workspacePath); }, }; }