From aa22680feef2d8b9e1a60ddae4280cedb2cf78f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Fri, 24 Jul 2026 09:39:07 +0200 Subject: [PATCH] refactor(cli): shared-process remote sessions with safe per-session exit (#12327) * refactor(cli): run remote sessions in one process with safe per-session exit Consolidate remote session handling into a single CLI process instead of spawning one process per remote-created session (addresses the PR review): - restore in-process create_session (accepts an absent sessionId and targets the connection directory); remove the session spawner, the KILO_REMOTE_ATTACH_SESSION attach-on-boot path, the child-advertisement gate, and their tests - retain instance advertisement and fire one immediate out-of-band heartbeat on (re)connect when advertising, so a headless `kilo remote` host is discoverable without delay Make /exit (wire command exit_cli, unchanged for compatibility) detach only the target session instead of terminating the CLI: - AttachedState.detach with a presence-suppression tombstone; detach also clears the target's SessionStatus so the negative-containment heartbeat fence resolves deterministically for busy/retry/offline sessions - exit_cli handler verifies ownership, cancels the active prompt, detaches and awaits the detach heartbeat, then ACKs; the interactive RemoteExit callback is invoked only after the ACK when the last owned session exits; a headless `kilo remote` host stays alive and advertising at zero sessions - add an optional canExitSession boolean to the list_commands v1 catalog (always true, independent of exitAvailable) so clients can detect safe session-exit semantics History and stored sessions are preserved on exit. * fix(cli): break module-load cycle in remote session prompt-cancel The K1 in-process exit_cli seam added a static `import { SessionPrompt }` to kilo-sessions.ts. @/session/prompt evaluates KiloSessionPrompt at module load, so the new static edge raced that init and left the namespace in TDZ, crashing unrelated test files with 'undefined is not an object (evaluating KiloSessionPrompt.shouldAskPlanFollowup)'. Defer to a dynamic import at the single call site, mirroring remote-command.ts. * fix(cli): correct AttachedState announce/detach concurrency and rollback Address review findings on the shared-process session lifecycle: - announce/detach no longer join the OPPOSITE in-flight operation. Joining detach's negative-containment fence made announce resolve success for a detached id (and vice versa: detach joined announce and resolved success while still attached, which exit_cli treats as license to ACK/close). Each path now joins only a same-kind in-flight op and, when the opposite op is in flight, awaits it to settle and then performs the real work. - Failed-detach rollback now releases the suppression tombstone, so a still-attached session is not dropped by the next setPresence (the tombstone loop would otherwise remove the still-present id and never clear). - Both catch/rollback branches now honor the lifecycle generation guard (mirroring the success path); a stale in-flight op that rejects after reset() no longer mutates the new lifecycle's presence/pending/suppressed sets (reset clears the same Set instances). Adds regression tests for each fix, plus AC6f covering the remote-ws detachSessionId negative-containment waiter. --- .changeset/remote-instance-advertisement.md | 5 + packages/opencode/src/cli/cmd/remote.ts | 32 ++ .../src/kilo-sessions/attached-state.ts | 188 ++++++- .../src/kilo-sessions/kilo-sessions.ts | 249 ++++++--- .../src/kilo-sessions/remote-command.ts | 21 +- .../src/kilo-sessions/remote-protocol.ts | 17 + .../src/kilo-sessions/remote-sender.ts | 174 ++++-- .../opencode/src/kilo-sessions/remote-ws.ts | 64 ++- .../test/kilocode/cli/cmd/remote.test.ts | 33 ++ .../test/kilocode/kilo-sessions.test.ts | 355 ++++++++++++- .../kilocode/sessions/attached-state.test.ts | 244 +++++++++ .../kilocode/sessions/remote-command.test.ts | 20 +- .../kilocode/sessions/remote-protocol.test.ts | 139 +++++ .../kilocode/sessions/remote-sender.test.ts | 497 ++++++++++++------ .../test/kilocode/sessions/remote-ws.test.ts | 176 +++++++ script/check-opencode-promise-facades.ts | 7 + 16 files changed, 1944 insertions(+), 277 deletions(-) create mode 100644 .changeset/remote-instance-advertisement.md create mode 100644 packages/opencode/test/kilocode/cli/cmd/remote.test.ts diff --git a/.changeset/remote-instance-advertisement.md b/.changeset/remote-instance-advertisement.md new file mode 100644 index 00000000000..4ef2fe06040 --- /dev/null +++ b/.changeset/remote-instance-advertisement.md @@ -0,0 +1,5 @@ +--- +"@kilocode/cli": minor +--- + +`kilo remote` instances now advertise themselves on the relay heartbeat. Each heartbeat carries the host's hostname, the project directory name, and the CLI build version, and each session entry advertises the platform it was created on. The cloud relay learns about a freshly-connected instance immediately (no 10s wait for the first timer tick), and the advertisement is race-safe across the explicit `kilo remote` command and bootstrap auto-enable (`KILO_REMOTE=1` / `remote_control` config). Legacy CLIs that send neither field remain wire-compatible. diff --git a/packages/opencode/src/cli/cmd/remote.ts b/packages/opencode/src/cli/cmd/remote.ts index 51b222e0a77..3e33f002dd3 100644 --- a/packages/opencode/src/cli/cmd/remote.ts +++ b/packages/opencode/src/cli/cmd/remote.ts @@ -4,6 +4,31 @@ import { bootstrap } from "../bootstrap" import { KiloSessions } from "@/kilo-sessions/kilo-sessions" import { context } from "@/project/instance-context" import { InstanceRuntime } from "@/project/instance-runtime" +import { Instance } from "@/kilocode/instance" +import { InstallationVersion } from "@opencode-ai/core/installation/version" +import os from "node:os" +import path from "node:path" + +function truncate(value: string, max: number) { + return value.length > max ? value.slice(0, max) : value +} + +// kilocode_change start - K1 W1: extracted so the advertisement payload shape +// is unit-testable as real behavior, rather than only through a source-text/ +// regex assertion on this file (the handler itself can't be driven end-to-end +// — see the doc comment on `handler` below). +export function buildInstanceAdvertisement(directory: string): { + name: string + projectName: string + version: string +} { + return { + name: truncate(os.hostname(), 64), + projectName: truncate(path.basename(directory) || directory, 64), + version: truncate(InstallationVersion, 32), + } +} +// kilocode_change end export const RemoteCommand = cmd({ command: "remote", @@ -11,6 +36,13 @@ export const RemoteCommand = cmd({ builder: (yargs) => yargs, handler: async () => { await bootstrap(process.cwd(), async () => { + // kilocode_change - K1 W1: advertise this instance on the relay + // heartbeat so the cloud side can show it as a spawn-capable instance. + // The process-wide `KILO_REMOTE_ATTACH_SESSION` guard was removed in K1 + // (in-process sessions only; no spawned children), so this is always + // advertised for the explicit `kilo remote` command path. + KiloSessions.setInstanceAdvertisement(buildInstanceAdvertisement(Instance.directory)) + await KiloSessions.enableRemote() console.log("Remote connection enabled.") diff --git a/packages/opencode/src/kilo-sessions/attached-state.ts b/packages/opencode/src/kilo-sessions/attached-state.ts index 3afad1cd057..3c80905ef6d 100644 --- a/packages/opencode/src/kilo-sessions/attached-state.ts +++ b/packages/opencode/src/kilo-sessions/attached-state.ts @@ -35,8 +35,13 @@ export namespace AttachedState { * `opts.requireSessionId` is forwarded by `announce(id)` so the relay * only resolves the attach promise when a fresh heartbeat whose * payload contains that id was actually sent. Presence fire-and- - * forget heartbeats call without an id and resolve on any fresh send. */ - heartbeat: (opts?: { requireSessionId?: string }) => Promise + * forget heartbeats call without an id and resolve on any fresh send. + * + * `opts.detachSessionId` is forwarded by `detach(id)` so the relay + * only resolves the detach promise when a fresh heartbeat whose + * payload DOES NOT contain that id was actually sent (the negative- + * containment fence). */ + heartbeat: (opts?: { requireSessionId?: string; detachSessionId?: string }) => Promise log?: { warn: (msg: string, meta?: unknown) => void } } @@ -54,11 +59,34 @@ export namespace AttachedState { * case presence is authoritative and the attach resolves successfully. * On success advances `lastSentKey` to the current union. */ announce(id: string): Promise + /** + * Awaitable session-detach. Removes the id from BOTH the presence and + * pending sets and awaits a fresh heartbeat whose payload no longer + * contains the id (id-containment fence so a stale "still contains" + * cycle cannot falsely report the detach as complete). + * + * On heartbeat failure: rolls back by restoring the prior ownership + * (presence add + pending add as appropriate), re-throws, and the + * caller is responsible for NOT sending the success response so the + * CLI can keep the session attached and the process alive. + * + * The id is also added to a suppression tombstone: until presence + * itself stops reporting the id, subsequent `setPresence` calls + * will NOT re-adopt it (this prevents an immediately-following + * presence replacement from instantly re-attaching a session that + * the remote just exited). The tombstone is released the moment + * `setPresence` receives a set that does not contain the id (i.e. + * presence has genuinely dropped it), so a later real reopen + * (a fresh announce after a legitimate re-open) is not blocked. + */ + detach(id: string): Promise /** Current union of presence ∪ pending for the next heartbeat payload. */ union(): ReadonlySet + /** True iff the id is in either the presence or pending set. */ + has(id: string): boolean /** Clear both sets across a connection lifecycle. The next setPresence * call after reset will fire a heartbeat because the baseline key is - * empty. */ + * empty. Also clears the suppressions. */ reset(): void } @@ -78,6 +106,13 @@ export namespace AttachedState { export function create(options: Options): Interface { const presence = new Set() const pending = new Set() + // kilocode_change - K1 W1: tombstones for ids that have been remotely + // detached but are still being reported by presence. While an id is in + // this set, setPresence MUST NOT re-adopt it, so a presence replacement + // that still includes a just-exited id cannot instantly re-attach it. + // The entry is released the first time presence reports a set that no + // longer includes the id (the upstream side has genuinely dropped it). + const suppressed = new Set() // kilocode_change - in-flight dedup. Concurrent announce(id) callers // share the same Promise so they observe one consistent outcome and // the heartbeat fires at most once per id. The owner is the caller @@ -85,6 +120,14 @@ export namespace AttachedState { // the owner clears the entry if the map still points to its Promise // (a later announce may have replaced it). Joiners only await. const inflight = new Map>() + // kilocode_change - in-flight detach dedup, mirrors `inflight` for the + // `detach` path. Multiple concurrent detach(id) callers share one + // Promise (id-containment heartbeat) so we never fire two conflicting + // detaches for the same id. Concurrent announce(id) and detach(id) + // also share this map so the two paths serialize on the same in-flight + // outcome (the detach-fence Promise resolves only when the id is + // absent from the sent payload). + const detachInflight = new Map>() // kilocode_change end let lastSentKey = "" // kilocode_change - lifecycle generation. Incremented on reset() so a @@ -109,6 +152,16 @@ export namespace AttachedState { return { setPresence(ids) { const next = new Set(ids) + // kilocode_change - K1 W1: suppression tombstone. Any id in `next` + // that is currently suppressed (a remote detach is in-flight or + // was just completed) MUST be filtered out so presence does not + // re-adopt a session the mobile client has just exited. The + // tombstone is released once presence reports a set that no + // longer includes the id (genuine drop upstream). + for (const tombstone of [...suppressed]) { + if (!next.has(tombstone)) suppressed.delete(tombstone) + else next.delete(tombstone) + } presence.clear() for (const id of next) presence.add(id) // Adopt any pending ids that presence now covers so the relay does @@ -127,13 +180,30 @@ export namespace AttachedState { async announce(id) { if (presence.has(id)) return - // kilocode_change - join an in-flight Promise for this id instead - // of starting a second heartbeat. + // kilocode_change - join a same-kind in-flight announce so concurrent + // callers share one heartbeat and one outcome. const existing = inflight.get(id) if (existing) { await existing return } + // kilocode_change - K1 W1: if a detach is in flight for this id, we + // must NOT join its Promise. The detach-fence resolves when the id is + // ABSENT from the sent payload — the opposite of what announce + // promises — so joining it would report a successful attach for a + // session that was actually detached. Wait for the detach to settle + // (its outcome is irrelevant to us) and then perform a real announce. + const inflightDetach = detachInflight.get(id) + if (inflightDetach) { + await inflightDetach.catch(() => undefined) + if (presence.has(id)) return + // A concurrent announce may have started while we awaited; join it. + const raced = inflight.get(id) + if (raced) { + await raced + return + } + } if (pending.has(id)) { // A previous announce already resolved and is awaiting presence // adoption. No further work to do. @@ -144,6 +214,10 @@ export namespace AttachedState { // lastSentKey with keyOf(union()) computed from the new state. const myGeneration = generation const owned = (async () => { + // kilocode_change - K1 W1: an explicit announce is a deliberate + // (re)attach that overrides any lingering detach tombstone, so + // presence can adopt this id again. No-op when not suppressed. + suppressed.delete(id) pending.add(id) try { // kilocode_change - forward the announced id so the relay only @@ -151,6 +225,13 @@ export namespace AttachedState { // contains this id was actually sent (id-containment fence). await options.heartbeat({ requireSessionId: id }) } catch (err) { + // kilocode_change - K1 W1: if reset() ran while this heartbeat was + // in flight, this announce belongs to a dead lifecycle. reset() + // clears the SAME set instances, so rolling back here would delete + // a `pending` entry a fresh post-reset announce for this id just + // installed. Bail without mutating the new generation's sets (the + // success path guards the same way before writing lastSentKey). + if (myGeneration !== generation) return // Roll back only the entry this call added. If presence adopted // the id while the heartbeat was in flight, presence is the // authoritative owner and the attach succeeded from the @@ -183,14 +264,111 @@ export namespace AttachedState { } }, + // kilocode_change - K1 W1: session-detach semantics. + async detach(id) { + // kilocode_change - join a same-kind in-flight detach so concurrent + // callers share one fence and one outcome. + const existingDetach = detachInflight.get(id) + if (existingDetach) { + await existingDetach + return + } + // kilocode_change - K1 W1: if an announce is in flight for this id we + // must NOT join it. `announce` adds the id to `pending` synchronously + // before its first await, so joining the announce Promise would + // resolve detach() successfully while the session is still fully + // attached — and exit_cli treats a resolved detach as license to ACK + // and close the CLI. Wait for the announce to settle, then run the + // real detach so the negative-containment fence actually fires. + const inflightAnnounce = inflight.get(id) + if (inflightAnnounce) { + await inflightAnnounce.catch(() => undefined) + const racedDetach = detachInflight.get(id) + if (racedDetach) { + await racedDetach + return + } + } + // Verify we own the id, AFTER settling any in-flight announce so the + // check sees the announce's real outcome. A detach for an id this CLI + // does not own is a caller bug; surfacing it as a specific error means + // the exit_cli handler can refuse to ACK and keep the CLI running. + const wasInPresence = presence.has(id) + const wasInPending = pending.has(id) + if (!wasInPresence && !wasInPending) { + throw new Error(`detach: ${id} is not owned by this CLI`) + } + // Tombstone the id BEFORE removing it from the sets. While the id + // is in `suppressed`, subsequent setPresence calls that still + // report the id (a presence churn race) will NOT re-adopt it. The + // tombstone is released the first time presence reports a set that + // genuinely no longer contains the id. + suppressed.add(id) + if (wasInPresence) presence.delete(id) + if (wasInPending) pending.delete(id) + const myGeneration = generation + const owned = (async () => { + try { + // Forward the id we are detaching via the relay's containment + // fence. The relay resolves this Promise only when a fresh + // heartbeat whose payload DOES NOT contain this id was + // actually sent over a live socket. + await options.heartbeat({ detachSessionId: id }) + } catch (err) { + // kilocode_change - K1 W1: if reset() ran while this heartbeat was + // in flight, this detach belongs to a dead lifecycle. reset() + // clears the SAME set instances, so restoring ownership / clearing + // the tombstone here would resurrect this id into a fresh + // post-reset lifecycle and could wipe a tombstone a concurrent + // post-reset detach legitimately set. Bail without mutating the + // new generation's sets (mirrors the success-path guard below). + if (myGeneration !== generation) return + // Roll back: restore the id to whichever sets it lived in AND + // release the tombstone. The detach failed, so the session is + // genuinely still attached and must stay adoptable by presence. + // Leaving the tombstone would make setPresence's suppression loop + // drop the still-present id on the very next call and never clear + // (presence keeps reporting it), permanently losing the session. + if (wasInPresence) presence.add(id) + if (wasInPending) pending.add(id) + suppressed.delete(id) + throw err + } + if (myGeneration !== generation) return + // The relay no longer has the id. The tombstone is released + // by setPresence's suppression logic the first time presence + // reports a set that no longer includes the id; we keep it in + // place here so the case where presence churn keeps reporting + // it for a few cycles (before the upstream side drops it) is + // handled coherently. + lastSentKey = keyOf(union()) + })() + detachInflight.set(id, owned) + try { + await owned + } finally { + if (detachInflight.get(id) === owned) detachInflight.delete(id) + } + }, + union() { return union() }, + has(id) { + return presence.has(id) || pending.has(id) + }, + reset() { presence.clear() pending.clear() inflight.clear() + // kilocode_change - K1 W1: also drop the in-flight detach map and + // tombstones so a new connection lifecycle starts with a clean + // slate and stale tombstones from a previous connection do not + // suppress a legitimate attach on the new one. + detachInflight.clear() + suppressed.clear() lastSentKey = "" // kilocode_change - bump the lifecycle generation so any in-flight // announce started before this reset will skip its lastSentKey diff --git a/packages/opencode/src/kilo-sessions/kilo-sessions.ts b/packages/opencode/src/kilo-sessions/kilo-sessions.ts index e4a0ef8b50b..e9deaeb575c 100644 --- a/packages/opencode/src/kilo-sessions/kilo-sessions.ts +++ b/packages/opencode/src/kilo-sessions/kilo-sessions.ts @@ -23,8 +23,9 @@ import { InstanceState } from "@/effect/instance-state" import { Instance } from "@/kilocode/instance" import { Vcs } from "@/project/vcs" import simpleGit from "simple-git" -import type { RemoteWS } from "@/kilo-sessions/remote-ws" -import type { RemoteSender } from "@/kilo-sessions/remote-sender" +import { RemoteWS } from "@/kilo-sessions/remote-ws" +import { RemoteSender } from "@/kilo-sessions/remote-sender" +import { RemoteProtocol } from "@/kilo-sessions/remote-protocol" import { AttachedState } from "@/kilo-sessions/attached-state" import { SessionStatus } from "@/session/status" import { Telemetry } from "@kilocode/kilo-telemetry" @@ -224,6 +225,13 @@ export namespace KiloSessions { let remote: { conn: RemoteWS.Connection; sender: RemoteSender.Sender } | undefined let enabling: Promise | undefined let remoteSeq = 0 + // kilocode_change - K1 W1: module-level instance advertisement flag. + // `enableRemote` can be triggered either by the explicit `kilo remote` command + // or by bootstrap auto-enable (`KILO_REMOTE=1` / `remote_control` config); it + // is idempotent/coalescing, so passing an {instance} arg on one specific call + // would race with whichever call happens first. A module-level flag flipped + // by either caller is the only race-free way to advertise the instance. + let instanceAdvertisement: RemoteProtocol.InstanceAdvertisement | undefined // Separate presence-owned attached session ids from newly-created (pending) // session announcements so a concurrent presence update cannot drop a pending // id and a heartbeat failure cannot delete a presence-owned id. The heartbeat @@ -232,9 +240,7 @@ export namespace KiloSessions { // into the sanitized failure response and the user retries manually. const attachedState = AttachedState.create({ heartbeat: (opts) => - remote - ? remote.conn.heartbeat(opts) - : Promise.reject(new Error("attachRemoteSession: no remote connection")), + remote ? remote.conn.heartbeat(opts) : Promise.reject(new Error("attachRemoteSession: no remote connection")), log: attachedLog, }) const statusSyncs = new Map() @@ -414,20 +420,20 @@ export namespace KiloSessions { return { ok: false, reason: "not_connected" } as const } - const readiness = yield* Effect.tryPromise({ - try: () => - withTimeout( - resolveReadiness(sessionID), - agentNotificationTimeoutMs(), - "agent notification readiness timed out", - ), - catch: () => ({ ok: false, reason: "not_connected" } as const), - }).pipe(Effect.catch((value) => Effect.succeed(value))) + const readiness = yield* Effect.tryPromise({ + try: () => + withTimeout( + resolveReadiness(sessionID), + agentNotificationTimeoutMs(), + "agent notification readiness timed out", + ), + catch: () => ({ ok: false, reason: "not_connected" }) as const, + }).pipe(Effect.catch((value) => Effect.succeed(value))) - if (!readiness.ok) return readiness - return yield* Effect.promise(() => - postAgentNotification(sessionID, readiness.ingestPath, readiness.client, input), - ) + if (!readiness.ok) return readiness + return yield* Effect.promise(() => + postAgentNotification(sessionID, readiness.ingestPath, readiness.client, input), + ) }) return Service.of({ init, sendAgentNotification }) @@ -480,7 +486,11 @@ export namespace KiloSessions { // Capture directory so the heartbeat timer can re-enter the Instance context // (setInterval runs outside AsyncLocalStorage scope) const directory = Instance.directory - const getSessions = async () => { + // kilocode_change - K1 W1: capture module-level advertisement so each + // heartbeat's `instance` field stays consistent with the flag at the + // moment of sending. The flag may be set after this closure is created + // (race-proof) — `getSessions` reads the current value each tick. + const getSessions = async (): Promise => { const [gitUrl, gitBranch] = await Promise.all([ getGitUrl().catch(() => undefined), branch().catch(() => undefined), @@ -504,6 +514,10 @@ export namespace KiloSessions { parentSessionId: session.parentID, gitUrl, gitBranch, + // kilocode_change - K1 W1: per-session platform, mirrors + // meta()'s resolution order so the live value always agrees + // with the session's stored created_on_platform. + platform: KiloSession.resolvePlatform(id) || process.env["KILO_PLATFORM"] || "cli", })), Effect.orElseSucceed(() => undefined), ), @@ -512,7 +526,8 @@ export namespace KiloSessions { ), ) const sessions = results.filter((r): r is NonNullable => !!r) - return { sessions } + const instance = instanceAdvertisement + return { type: "heartbeat", sessions, ...(instance ? { instance } : {}) } } const conn = RemoteWS.connect({ @@ -523,6 +538,19 @@ export namespace KiloSessions { log, onOpen: () => { void Bus.publish(Instance.current, Event.RemoteStatusChanged, { enabled: true, connected: true }) + // kilocode_change - K1 W1: on reconnect, a headless `kilo remote` host + // preserves its module-level advertisement flag but would otherwise not + // be re-advertised until the next periodic heartbeat (up to ~10s). + // Fire one immediate out-of-band heartbeat when the flag is set. + // This is intentionally conditional: tests that do not set the flag + // must not see extra heartbeats. + if (instanceAdvertisement) { + void conn.heartbeat().catch((err) => + log.warn("reconnect advertisement heartbeat failed", { + error: String(err), + }), + ) + } }, onDisconnect: () => { void Bus.publish(Instance.current, Event.RemoteStatusChanged, { enabled: !!remote, connected: false }) @@ -538,6 +566,24 @@ export namespace KiloSessions { conn, directory: Instance.directory, log, + // kilocode_change - K1 W1: in-process attach/detach/ownership seams + // back to KiloSessions. The sender does NOT spawn a process per + // session — concurrent remote sessions share this CLI process with + // per-directory InstanceRef isolation. + attachSession: (id) => KiloSessions.attachRemoteSession(id), + detachSession: (id) => KiloSessions.detachRemoteSession(id), + hasSession: (id) => KiloSessions.hasRemoteSession(id), + ownedCount: () => KiloSessions.ownedRemoteSessionCount(), + cancelPrompt: async (id) => { + // kilocode_change - K1 W1: dynamic import breaks the module-load cycle + // (@/session/prompt reads KiloSessionPrompt at eval; a static edge here + // races that init). Mirrors remote-command.ts's lazy SessionPrompt use. + const [{ AppRuntime }, { SessionPrompt }] = await Promise.all([ + import("@/effect/app-runtime"), + import("@/session/prompt"), + ]) + await AppRuntime.runPromise(SessionPrompt.Service.use((svc) => svc.cancel(id))) + }, }) if (seq !== remoteSeq) { @@ -593,6 +639,33 @@ export namespace KiloSessions { attachedState.setPresence(ids) } + // kilocode_change - K1 W1: instance advertisement setter. + // Idempotent. If a remote connection is already established when the flag is + // flipped (typical for the race between bootstrap auto-enable and the + // explicit `kilo remote` command — `enableRemote` itself is coalescing), we + // fire one out-of-band heartbeat so the cloud side learns about the + // instance without waiting for the next 10s timer tick. + export function setInstanceAdvertisement(advertisement: RemoteProtocol.InstanceAdvertisement) { + instanceAdvertisement = advertisement + if (remote) { + void remote.conn.heartbeat().catch((err) => + log.warn("instance advertisement heartbeat failed", { + error: String(err), + }), + ) + } + } + + // Test-only: the advertisement flag is intentionally one-way in production + // (once a process runs `kilo remote`, it keeps advertising for its whole + // lifetime, including across a transient disableRemote/enableRemote + // reconnect cycle — disableRemote() deliberately does not clear it). Tests + // that assert the "unset" default must reset the module-level flag + // themselves between cases. + export function resetInstanceAdvertisementForTests() { + instanceAdvertisement = undefined + } + // Duplicate-safe single-session attach used by the remote create_session command. Delegates to // the two-set state so the announcement is preserved across a concurrent presence replacement // and a heartbeat failure rolls back only the entry this call added (a presence-owned id is never @@ -601,66 +674,96 @@ export namespace KiloSessions { await attachedState.announce(id) } -export async function create(sessionId: string) { - const inflight = bootstrapInflight.get(sessionId) - if (inflight) { - const result = await inflight - if (!result.ok) return { id: "", ingestPath: "" } - return { id: sessionId, ingestPath: result.ingestPath } + // kilocode_change - K1 W1: session-detach semantics. The exit_cli handler + // calls this after a verified owns-check + cancel-prompt; the heartbeat + // must confirm the id was removed from the next sent payload (negative- + // containment fence) before the handler ACKs the request. + // + // The SessionStatus entry is cleared to idle (which deletes the map entry) + // before the heartbeat fence runs, so the next getSessions() payload — and + // therefore the fence itself — deterministically omits the id regardless of + // whether the session was busy/retry/offline. On heartbeat-failure rollback, + // attachedState.detach restores the id to presence/pending; the session is + // still advertised (via the union) with an idle status until normal activity + // re-establishes a status, so the relay does not under-report an owned session. + export async function detachRemoteSession(id: string) { + const { AppRuntime } = await import("@/effect/app-runtime") + await AppRuntime.runPromise(SessionStatus.Service.use((svc) => svc.set(SessionID.make(id), { type: "idle" }))) + await attachedState.detach(id) } - // Synchronously register the in-flight bootstrap promise before any await - // so concurrent callers (e.g. sendAgentNotification racing the - // Session.Event.Created handler) deterministically coalesce onto the same - // POST /api/session. - const task = trackBootstrap(sessionId, () => bootstrap(sessionId)) - const result = await task - if (!result) return { id: "", ingestPath: "" } + // kilocode_change - K1 W1: ownership probe used by the exit_cli handler + // before the cancel/detach sequence. Cheap and synchronous. + export function hasRemoteSession(id: string): boolean { + return attachedState.has(id) + } - void fullSync(sessionId).catch((error) => log.error("share full sync failed", { sessionId, error })) + // kilocode_change - K1 W1: count of "owned" sessions (presence ∪ pending). + // Used to drive the last-interactive-session exit decision: zero remaining + // + a registered RemoteExit callback => invoke it after the ACK can flush; + // zero remaining + no callback (kilo remote) => keep host alive. Sessions + // remain => stay alive regardless of callback state. + export function ownedRemoteSessionCount(): number { + return attachedState.union().size + } - return result -} + export async function create(sessionId: string) { + const inflight = bootstrapInflight.get(sessionId) + if (inflight) { + const result = await inflight + if (!result.ok) return { id: "", ingestPath: "" } + return { id: sessionId, ingestPath: result.ingestPath } + } -// Track an in-flight bootstrap for `sessionId` so callers that race the -// share ingest path (e.g. the `notify_user` tool calling -// sendAgentNotification before the Session.Event.Created handler has -// finished POSTing /api/session) can await the same outcome instead of -// firing their own bootstrap or failing. The bootstrap outcome promise is -// created and stored in `bootstrapInflight` synchronously before the first -// `await` so concurrent callers are deterministically coalesced. -function trackBootstrap( - sessionId: string, - start: () => Promise<{ id: string; ingestPath: string } | undefined>, -) { - // Build the task and derived outcome promise as synchronous expressions - // first; only then register the entry. This guarantees the value stored - // in `bootstrapInflight` is the real promise rather than `undefined`. - const task = start() - const tracked: Promise = task - .then((value): BootstrapOutcome => { - if (!value) return { ok: false, reason: "not_connected" } - return { ok: true, ingestPath: value.ingestPath } - }) - .catch((error: unknown): BootstrapOutcome => { - const reason = error instanceof Error ? error.message : String(error) - log.warn("session bootstrap failed", { sessionId, reason }) - return { ok: false, reason } + // Synchronously register the in-flight bootstrap promise before any await + // so concurrent callers (e.g. sendAgentNotification racing the + // Session.Event.Created handler) deterministically coalesce onto the same + // POST /api/session. + const task = trackBootstrap(sessionId, () => bootstrap(sessionId)) + const result = await task + if (!result) return { id: "", ingestPath: "" } + + void fullSync(sessionId).catch((error) => log.error("share full sync failed", { sessionId, error })) + + return result + } + + // Track an in-flight bootstrap for `sessionId` so callers that race the + // share ingest path (e.g. the `notify_user` tool calling + // sendAgentNotification before the Session.Event.Created handler has + // finished POSTing /api/session) can await the same outcome instead of + // firing their own bootstrap or failing. The bootstrap outcome promise is + // created and stored in `bootstrapInflight` synchronously before the first + // `await` so concurrent callers are deterministically coalesced. + function trackBootstrap(sessionId: string, start: () => Promise<{ id: string; ingestPath: string } | undefined>) { + // Build the task and derived outcome promise as synchronous expressions + // first; only then register the entry. This guarantees the value stored + // in `bootstrapInflight` is the real promise rather than `undefined`. + const task = start() + const tracked: Promise = task + .then((value): BootstrapOutcome => { + if (!value) return { ok: false, reason: "not_connected" } + return { ok: true, ingestPath: value.ingestPath } + }) + .catch((error: unknown): BootstrapOutcome => { + const reason = error instanceof Error ? error.message : String(error) + log.warn("session bootstrap failed", { sessionId, reason }) + return { ok: false, reason } + }) + + // Register synchronously before any async work starts so concurrent + // callers see the entry in `bootstrapInflight` immediately. + bootstrapInflight.set(sessionId, tracked) + tracked.finally(() => { + if (bootstrapInflight.get(sessionId) === tracked) bootstrapInflight.delete(sessionId) }) + return task + } - // Register synchronously before any async work starts so concurrent - // callers see the entry in `bootstrapInflight` immediately. - bootstrapInflight.set(sessionId, tracked) - tracked.finally(() => { - if (bootstrapInflight.get(sessionId) === tracked) bootstrapInflight.delete(sessionId) - }) - return task -} - -/** @internal - test-only helper */ -export function _getBootstrapInflight(sessionId: string): Promise | undefined { - return bootstrapInflight.get(sessionId) -} + /** @internal - test-only helper */ + export function _getBootstrapInflight(sessionId: string): Promise | undefined { + return bootstrapInflight.get(sessionId) + } export async function bootstrap(sessionId: string) { if (ingestDisabled) { diff --git a/packages/opencode/src/kilo-sessions/remote-command.ts b/packages/opencode/src/kilo-sessions/remote-command.ts index 3e9ac87d72b..9c204739f2e 100644 --- a/packages/opencode/src/kilo-sessions/remote-command.ts +++ b/packages/opencode/src/kilo-sessions/remote-command.ts @@ -62,6 +62,18 @@ export namespace RemoteCommand { .object({ protocolVersion: z.literal(1), commands: z.array(Info).max(MAX_COMMANDS), + // kilocode_change - K1 W1: `canExitSession` is an INDEPENDENT producer + // contract advertised by every CLI (interactive TUI and headless + // `kilo remote` alike) so the mobile client can rely on the + // interpretation of `exit_cli` as "detach THIS session" without having + // to inspect the synthetic `/exit` command entry's presence (which is + // gated on RemoteExit.get() — i.e. interactive-only — and therefore + // would lie to a headless host). Compatibility note: the wire command + // literal `exit_cli` is intentionally unchanged from the prior + // interactive-only meaning; this field documents the new interpretation + // rather than introducing a new command. The synthetic `/exit` entry + // (gated on exitAvailable) is kept as-is for the interactive TUI. + canExitSession: z.boolean().optional(), }) .strict() export type Response = z.infer @@ -148,7 +160,14 @@ export namespace RemoteCommand { // response stays alphabetized regardless of input order. if (!names.has(compact.name)) commands.push(compact) if (exitAvailable) commands.push(exit) - return Response.parse({ protocolVersion: 1, commands: truncate(commands) }) + // kilocode_change - K1 W1: always advertise `canExitSession: true`. This + // is the producer contract for the new `exit_cli`-as-detach semantics + // (independent of `exitAvailable`, which gates the synthetic `/exit` + // command entry on RemoteExit.get()). A headless `kilo remote` host has + // no RemoteExit callback, so it does NOT emit `/exit` here — but it + // DOES interpret `exit_cli` as a session-detach, so `canExitSession` + // is true for both interactive and headless producers. + return Response.parse({ protocolVersion: 1, commands: truncate(commands), canExitSession: true }) } export type ExecuteInput = SendRequest & { sessionID: SessionID; catalog: Response } diff --git a/packages/opencode/src/kilo-sessions/remote-protocol.ts b/packages/opencode/src/kilo-sessions/remote-protocol.ts index 926a025a9d0..30a2202413b 100644 --- a/packages/opencode/src/kilo-sessions/remote-protocol.ts +++ b/packages/opencode/src/kilo-sessions/remote-protocol.ts @@ -10,9 +10,25 @@ export namespace RemoteProtocol { parentSessionId: z.string().optional(), gitUrl: z.string().optional(), gitBranch: z.string().optional(), + // kilocode_change - K1 W1: per-session platform advertises the platform the + // session was created on. Mirrors meta()'s resolution order: + // KiloSession.resolvePlatform(id) || process.env["KILO_PLATFORM"] || "cli" + // Optional so legacy CLIs (no field) remain wire-compatible. + platform: z.string().max(32).optional(), }) export type SessionInfo = z.infer + // kilocode_change - K1 W1: instance advertisement. Presence on a heartbeat + // means "this connection is a spawn-capable instance" and turns this CLI into + // a row on the cloud-side instance picker. Legacy CLIs (no `instance`) are + // wire-compatible and never regress. + export const InstanceAdvertisement = z.object({ + name: z.string().min(1).max(64), // os.hostname(), truncated + projectName: z.string().min(1).max(64), // basename(Instance.directory), truncated + version: z.string().max(32).optional(), // InstallationVersion, truncated + }) + export type InstanceAdvertisement = z.infer + // --- CLI → DO (Outbound) --- // Capability flags advertised in the heartbeat so the relay can stop @@ -27,6 +43,7 @@ export namespace RemoteProtocol { type: z.literal("heartbeat"), sessions: z.array(SessionInfo), protocolVersion: z.string().optional(), // lets relay detect CLI capabilities without probing commands + instance: InstanceAdvertisement.optional(), // kilocode_change - K1 W1 capabilities: Capabilities, }) export type Heartbeat = z.infer diff --git a/packages/opencode/src/kilo-sessions/remote-sender.ts b/packages/opencode/src/kilo-sessions/remote-sender.ts index 165c221d57e..8af252e6dd0 100644 --- a/packages/opencode/src/kilo-sessions/remote-sender.ts +++ b/packages/opencode/src/kilo-sessions/remote-sender.ts @@ -128,18 +128,24 @@ export namespace RemoteSender { // Production falls back to Session.Service.create with `{}`. readonly create?: (input?: Record) => Promise // kilocode_change - injectable remove hook used to roll back an orphan - // root session when attachSession fails after creation. The default + // root session when the spawn fails after creation. The default // delegates to Session.Service.remove and only swallows its own errors - // so the original attach failure is what reaches the caller. + // so the original spawn failure is what reaches the caller. readonly remove?: (sessionID: SessionID) => Promise // kilocode_change end } - // kilocode_change start - duplicate-safe attach hook used by create_session. - // Production wires this to KiloSessions.attachRemoteSession so the attached - // set is mutated exactly once and the relay heartbeat fires only when the - // set actually changes. + // kilocode_change - K1 W1: in-process attach/detach/ownership/cancel + // seams. All four are optional and default to a lazy import of + // `KiloSessions` (production wires them in `enableRemote`, so the + // default branch is never hit there; tests that don't care about + // these paths simply omit them and the defaults supply no-op-safe + // shims so the production call sites stay the only places that + // actually touch the AttachedState). attachSession?: (sessionID: SessionID) => Promise - // kilocode_change end + detachSession?: (sessionID: SessionID) => Promise + hasSession?: (sessionID: SessionID) => boolean + ownedCount?: () => number + cancelPrompt?: (sessionID: SessionID) => Promise catalog?: { readonly get: (sessionID: SessionID) => Promise readonly messages: (sessionID: SessionID) => Promise @@ -238,10 +244,10 @@ export namespace RemoteSender { }, } // kilocode_change start - orphan rollback for create_session: when - // sessionCreate succeeds but attachSession fails, the newly-created root - // session would otherwise stay in the DB with no relay awareness. The + // sessionCreate succeeds but the spawn fails, the newly-created root + // session would otherwise stay in the DB with no child to serve it. The // default remove() delegates to Session.Service.remove and swallows its - // own errors so the caller still observes the original attach failure. + // own errors so the caller still observes the original spawn failure. const sessionRemove = session.remove ?? (async (id: SessionID) => { @@ -249,7 +255,12 @@ export namespace RemoteSender { await AppRuntime.runPromise(Session.Service.use((svc) => svc.remove(id))) }) // kilocode_change end - // kilocode_change start - session create + duplicate-safe attach used by create_session + // kilocode_change - K1 W1: session create + in-process attach seams used by + // create_session. Production wires `attachSession` to + // `KiloSessions.attachRemoteSession` from inside `enableRemote` (see + // kilo-sessions.ts). Test fixtures inject stubs via the Options object. + // When omitted, the create_session / exit_cli handlers treat the seam + // as a wiring bug (a missing seam is never a runtime fallback). const sessionCreate = session.create ?? (async (input?: Record) => { @@ -264,6 +275,20 @@ export namespace RemoteSender { const { KiloSessions } = await import("@/kilo-sessions/kilo-sessions") await KiloSessions.attachRemoteSession(id) }) + const detachSession = + options.detachSession ?? + (async (id: SessionID) => { + const { KiloSessions } = await import("@/kilo-sessions/kilo-sessions") + await KiloSessions.detachRemoteSession(id) + }) + const hasSession = options.hasSession ?? (() => false) + const ownedCount = options.ownedCount ?? (() => 0) + const cancelPrompt = + options.cancelPrompt ?? + (async (id: SessionID) => { + const { AppRuntime } = await import("@/effect/app-runtime") + await AppRuntime.runPromise(SessionPrompt.Service.use((svc) => svc.cancel(id))) + }) // kilocode_change end // kilocode_change start - injectable slash command discovery + execution const commands = options.commands ?? RemoteCommand.live() @@ -651,42 +676,117 @@ export namespace RemoteSender { return } if (msg.command === "exit_cli") { + // kilocode_change - K1 W1: `exit_cli` now means "detach THIS remote + // session and (if this is the last interactive session) close the + // CLI." It is NOT "terminate the CLI." A headless `kilo remote` host + // never invokes the RemoteExit callback (it is never registered for + // headless mode), so the same command cleanly handles both the + // interactive TUI shutdown path and the per-session-detach path + // without introducing a new wire command. + // + // The wire command literal `exit_cli` is intentionally unchanged + // (the prior PR's review accepted this: the contract shifts from + // "exit" to "exit session" but the literal is kept for compatibility + // with older clients already in the field). + // + // Steps: + // 1. Verify the target id is a real SessionID. + // 2. Verify this CLI OWNS the target (AttachedState.has). A + // non-owning detach would silently re-add the id to presence + // (the tombstone) and we don't want that. + // 3. Cancel any active prompt for the target session so the user + // doesn't see a "still working" indicator after they leave. + // 4. Detach the id (removes from BOTH presence and pending; awaits + // a fresh heartbeat whose payload no longer contains the id; + // rolls back on failure). + // 5. Snapshot the remaining-count AFTER detach from + // attachedState.union() (NOT mobile subscriptions). + // 6. If zero remain + a RemoteExit callback is registered, ACK + // then invoke the callback in a microtask so the response can + // flush first. If zero remain + no callback (headless `kilo + // remote`), ACK and keep the host alive (the host keeps + // advertising and can create a new session from zero). If + // sessions remain, ACK and keep the process alive. + // 7. On any failure (owns-check, cancel, detach), surface a + // sanitized error and do NOT ACK; the CLI keeps the session + // attached and the process stays alive. const parsed = RemoteCommand.ExitRequest.safeParse(msg.data) const current = msg.sessionId ? decodeSessionID(msg.sessionId) : Option.none() if (!parsed.success || Option.isNone(current)) { options.conn.send({ type: "response", id: msg.id, error: "invalid exit_cli command" }) return } + const target = current.value + // Verify ownership first — a non-owning detach would silently re-add + // the id to the tombstone, which is a wiring bug we want to surface + // (and a mobile client trying to detach a session it does not own + // is a contract violation we should not paper over). + if (!hasSession(target)) { + options.conn.send({ type: "response", id: msg.id, error: "session not owned by this CLI" }) + return + } + const exit = remoteExit.get() void (async () => { try { - await session.get(current.value) - const exit = remoteExit.get() - if (!exit) { - options.conn.send({ type: "response", id: msg.id, error: "graceful exit unavailable" }) - return - } + // 1. Cancel any active prompt for the target session. We await + // this (not fire-and-forget) because the detach fence that + // follows depends on a coherent session state — the prompt + // cancel may need to flush queued messages before the + // session is no longer "busy" to the relay. + await cancelPrompt(target) + // 2. Detach + await the negative-containment heartbeat. + await detachSession(target) + // 3. Snapshot remaining sessions AFTER detach. Headless hosts + // (`kilo remote`) never register a RemoteExit callback, so + // `exit` is undefined there and the host stays alive. + const remaining = ownedCount() options.conn.send({ type: "response", id: msg.id, result: {} }) - queueMicrotask(() => { - void exit().catch((error) => { - options.log.error("exit CLI failed after ACK", { - id: msg.id, - operation: "exit_cli", - error: errorName(error), + if (remaining === 0 && exit) { + queueMicrotask(() => { + void exit().catch((error) => { + options.log.error("exit CLI failed after ACK", { + id: msg.id, + operation: "exit_cli", + error: errorName(error), + }) }) }) - }) + } } catch (error) { - options.log.error("exit CLI preflight failed", { id: msg.id, error: errorName(error) }) - options.conn.send({ type: "response", id: msg.id, error: "failed to exit CLI" }) + // Roll-back path: the detach may have partially applied. The + // AttachedState.detach rollback restores presence/pending on + // its own. We MUST NOT ACK here — the CLI keeps the session + // attached and the process stays alive. + options.log.error("exit CLI failed before ACK", { id: msg.id, error: errorName(error) }) + options.conn.send({ type: "response", id: msg.id, error: "failed to exit session" }) } })() return } if (msg.command === "create_session") { - // kilocode_change start - remote /new creation: root session, attached + heartbeat before response + // kilocode_change - K1 W1: in-process create_session. The wire + // shape is unchanged (`{protocolVersion: 1}`), but the handler now + // (a) accepts an absent `sessionId` (the instance-picker path is + // connectionId-targeted — no source session needed), (b) resolves + // the target directory to that existing session's directory when + // a `sessionId` is present (legacy mobile /new-inside-a-session + // path) or to `options.directory` (the instance's own launch + // directory) otherwise, and (c) attaches the new session in the + // same CLI process (concurrent sessions share the process with + // per-directory InstanceRef isolation) instead of spawning a child. + // Attach failures roll back the pre-created session via + // `sessionRemove`. const parsed = CreateSessionRequest.safeParse(msg.data) + if (!parsed.success) { + options.conn.send({ + type: "response", + id: msg.id, + error: "invalid create_session command", + }) + return + } const current = msg.sessionId ? decodeSessionID(msg.sessionId) : Option.none() - if (!parsed.success || Option.isNone(current)) { + if (msg.sessionId && Option.isNone(current)) { options.conn.send({ type: "response", id: msg.id, @@ -697,8 +797,17 @@ export namespace RemoteSender { const run = options.provide ?? provide void (async () => { try { + // Resolve the target directory: a present `sessionId` keeps the + // legacy mobile /new-inside-a-session behavior (target = that + // session's directory); an absent `sessionId` targets the + // instance's own launch directory (the new instance-picker path). + const targetDirectory = await current.pipe( + Option.map((sid) => session.get(sid)), + Option.map((p) => p.then((info) => info.directory)), + Option.getOrElse(() => Promise.resolve(options.directory)), + ) const result = await run({ - directory: (await session.get(current.value)).directory, + directory: targetDirectory, fn: async () => { const created = await sessionCreate({}) // attachSession is the duplicate-safe seam: it mutates the @@ -709,9 +818,10 @@ export namespace RemoteSender { await attachSession(created.id) } catch (attachError) { // Roll back the newly-created root session so the DB does - // not keep an orphan the relay never learned about. Swallow - // the cleanup error here — the original attach failure is - // what the caller must see, so we re-throw it below. + // not keep an orphan the relay never learned about. + // Swallow the cleanup error here — the original attach + // failure is what the caller must see, so we re-throw it + // below. try { await sessionRemove(created.id) } catch (cleanupError) { diff --git a/packages/opencode/src/kilo-sessions/remote-ws.ts b/packages/opencode/src/kilo-sessions/remote-ws.ts index de631d23e1c..5e9bfb87818 100644 --- a/packages/opencode/src/kilo-sessions/remote-ws.ts +++ b/packages/opencode/src/kilo-sessions/remote-ws.ts @@ -14,7 +14,12 @@ export namespace RemoteWS { export type Options = { url: string getToken: () => Promise - getSessions: () => Promise<{ sessions: SessionInfo[] }> + // kilocode_change - K1 W1: widened return type so the optional `instance` + // advertisement (RemoteProtocol.Heartbeat.instance) flows through to the + // wire unchanged when the gatherer provides it. Legacy callers + // (older test mocks) still satisfy the contract by returning a bare + // `{ sessions }` shape. + getSessions: () => Promise<{ sessions: SessionInfo[]; instance?: RemoteProtocol.Heartbeat["instance"] }> log: { info: (...args: any[]) => void error: (...args: any[]) => void @@ -57,8 +62,15 @@ export namespace RemoteWS { * fences attach-announce waiters so a fresh heartbeat that legitimately * omits the announced id (e.g. the gather's `Effect.orElseSucceed` * filtered it out) does not falsely report the session as attached. + * + * When `opts.detachSessionId` is provided, the promise only resolves + * when the sent fresh payload's session list DOES NOT contain that id + * (the negative-containment fence used by K1 W1 session-detach). + * Stale "still contains" cycles are rejected via requeue (handled + * below) so the detach does not falsely report success while the + * upstream side still observes the session. */ - heartbeat(opts?: { requireSessionId?: string }): Promise + heartbeat(opts?: { requireSessionId?: string; detachSessionId?: string }): Promise close(): void readonly connected: boolean } @@ -119,7 +131,7 @@ export namespace RemoteWS { let lastGood: SessionInfo[] | undefined let outstanding = 0 let degradedCount = 0 - type Waiter = { resolve: () => void; reject: (err: unknown) => void; requireSessionId?: string } + type Waiter = { resolve: () => void; reject: (err: unknown) => void; requireSessionId?: string; detachSessionId?: string } let waiters: Waiter[] = [] function makeWaiter(): { promise: Promise; waiter: Waiter } { @@ -136,9 +148,10 @@ export namespace RemoteWS { for (const w of list) w.reject(err) } - // One bounded gather. Never throws. Returns the fresh session list, or - // undefined to signal a degraded cycle (caller sends last known-good). - async function gatherOnce(): Promise { + // One bounded gather. Never throws. Returns the fresh session list (and + // optional instance advertisement), or undefined to signal a degraded + // cycle (caller sends last known-good). + async function gatherOnce(): Promise<{ sessions: SessionInfo[]; instance?: RemoteProtocol.Heartbeat["instance"] } | undefined> { if (outstanding >= maxOutstandingGathers) { degradedCount++ options.log.warn("remote-ws heartbeat gather cap reached, degraded heartbeat", { @@ -162,7 +175,7 @@ export namespace RemoteWS { .then( (r) => { release() - return { ok: true as const, sessions: r.sessions } + return { ok: true as const, sessions: r.sessions, instance: r.instance } }, (err) => { release() @@ -170,7 +183,7 @@ export namespace RemoteWS { }, ) const outcome = await new Promise< - { kind: "ok"; sessions: SessionInfo[] } | { kind: "err"; error: unknown } | { kind: "timeout" } + { kind: "ok"; sessions: SessionInfo[]; instance?: RemoteProtocol.Heartbeat["instance"] } | { kind: "err"; error: unknown } | { kind: "timeout" } >((resolve) => { let done = false const t = timers.setTimeout(() => { @@ -182,10 +195,14 @@ export namespace RemoteWS { if (done) return done = true timers.clearTimeout(t) - resolve(res.ok ? { kind: "ok", sessions: res.sessions } : { kind: "err", error: res.error }) + resolve( + res.ok + ? { kind: "ok", sessions: res.sessions, instance: res.instance } + : { kind: "err", error: res.error }, + ) }) }) - if (outcome.kind === "ok") return outcome.sessions + if (outcome.kind === "ok") return { sessions: outcome.sessions, instance: outcome.instance } degradedCount++ if (outcome.kind === "err") { options.log.warn("remote-ws heartbeat gather rejected, degraded heartbeat", { @@ -201,10 +218,11 @@ export namespace RemoteWS { return undefined } - function heartbeat(opts?: { requireSessionId?: string }): Promise { + function heartbeat(opts?: { requireSessionId?: string; detachSessionId?: string }): Promise { if (closed) return Promise.reject(new Error("remote-ws connection closed")) const { promise, waiter } = makeWaiter() waiter.requireSessionId = opts?.requireSessionId + waiter.detachSessionId = opts?.detachSessionId waiters.push(waiter) requestCycle() return promise @@ -231,13 +249,21 @@ export namespace RemoteWS { return } if (fresh !== undefined) { - lastGood = fresh + lastGood = fresh.sessions const sentLive = ws?.readyState === WebSocket.OPEN + // kilocode_change - K1 W1: spread optional `instance` so the + // instance advertisement propagates to the wire when the + // gatherer provided it. The `lastGood` cache (degraded + // fallback) intentionally drops the instance — degraded + // heartbeats must not echo a stale advertisement. + // capabilities.attachments is carried from #12394 (mobile file + // attachments) — an independent additive heartbeat field. send({ type: "heartbeat", protocolVersion: InstallationVersion, capabilities: { attachments: true }, - sessions: fresh, + sessions: fresh.sessions, + ...(fresh.instance ? { instance: fresh.instance } : {}), }) if (sentLive) { // A waiter requiring a specific id is satisfied only when @@ -245,12 +271,20 @@ export namespace RemoteWS { // are requeued so the periodic interval keeps evaluating // them; they resolve on a future fresh send whose payload // includes their required id (or reject on close). + // + // kilocode_change - K1 W1: a `detachSessionId` waiter + // resolves only when the sent payload DOES NOT contain + // that id (the negative-containment fence used by + // session-detach). Until the upstream side drops the id, + // the waiter is requeued. const satisfied: Waiter[] = [] const unsatisfied: Waiter[] = [] for (const w of cycleWaiters) { + const present = fresh.sessions.some((s) => s.id === w.requireSessionId) + const stillPresent = fresh.sessions.some((s) => s.id === w.detachSessionId) if ( - w.requireSessionId === undefined || - fresh.some((s) => s.id === w.requireSessionId) + (w.requireSessionId === undefined || present) && + (w.detachSessionId === undefined || !stillPresent) ) { satisfied.push(w) } else { diff --git a/packages/opencode/test/kilocode/cli/cmd/remote.test.ts b/packages/opencode/test/kilocode/cli/cmd/remote.test.ts new file mode 100644 index 00000000000..abeb24b02cf --- /dev/null +++ b/packages/opencode/test/kilocode/cli/cmd/remote.test.ts @@ -0,0 +1,33 @@ +// kilocode_change - new file +// K1 W1: verify `buildInstanceAdvertisement`'s payload shape as real behavior. +// +// The `RemoteCommand` handler itself is a CLI entry point that calls +// `bootstrap(process.cwd(), async () => { ... })` and then awaits an abort +// signal that never resolves in a test — it cannot be driven end-to-end. +// `buildInstanceAdvertisement` is extracted from the handler specifically so +// the advertised payload is independently testable as real behavior, not via +// a source-text/regex assertion on the handler's structure. + +import { describe, expect, test } from "bun:test" +import { buildInstanceAdvertisement } from "../../../../src/cli/cmd/remote" + +describe("RemoteCommand instance advertisement (K1 W1)", () => { + test("buildInstanceAdvertisement resolves name/projectName/version from the directory and installation version", () => { + const advertisement = buildInstanceAdvertisement("/Users/igor/projects/my-app") + expect(advertisement.projectName).toBe("my-app") + expect(typeof advertisement.name).toBe("string") + expect(advertisement.name.length).toBeGreaterThan(0) + expect(typeof advertisement.version).toBe("string") + }) + + test("buildInstanceAdvertisement truncates an overlong project directory name to 64 chars", () => { + const longName = "a".repeat(100) + const advertisement = buildInstanceAdvertisement(`/Users/igor/projects/${longName}`) + expect(advertisement.projectName.length).toBeLessThanOrEqual(64) + }) + + test("buildInstanceAdvertisement falls back to the full directory when basename is empty (root path)", () => { + const advertisement = buildInstanceAdvertisement("/") + expect(advertisement.projectName).toBe("/") + }) +}) diff --git a/packages/opencode/test/kilocode/kilo-sessions.test.ts b/packages/opencode/test/kilocode/kilo-sessions.test.ts index 368d300704b..78f26bf75d5 100644 --- a/packages/opencode/test/kilocode/kilo-sessions.test.ts +++ b/packages/opencode/test/kilocode/kilo-sessions.test.ts @@ -1,5 +1,6 @@ // kilocode_change - new file -import { expect, spyOn } from "bun:test" +import { afterEach, beforeEach, describe, expect, mock, spyOn, test } from "bun:test" +import { tmpdir } from "../fixture/fixture" import { Effect, Layer } from "effect" import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner" import { Auth } from "../../src/auth" @@ -8,13 +9,19 @@ import { GlobalBus } from "../../src/bus/global" import type { Config } from "../../src/config/config" import { clearInFlightCache } from "../../src/kilo-sessions/inflight-cache" import { KiloSessions } from "../../src/kilo-sessions/kilo-sessions" +import { provide } from "../../src/kilocode/instance" +import { RemoteWS } from "../../src/kilo-sessions/remote-ws" +import { RemoteSender } from "../../src/kilo-sessions/remote-sender" import { ProjectV2 } from "@opencode-ai/core/project" import { Session } from "../../src/session/session" import { SessionID } from "../../src/session/schema" +import { SessionStatus } from "../../src/session/status" +import { QuestionID } from "../../src/question/schema" import { TestConfig } from "../fixture/config" import { testEffect } from "../lib/effect" import { InstanceStore } from "../../src/project/instance-store" import { TestInstance, testInstanceStoreLayer, tmpdirScoped } from "../fixture/fixture" +import { RemoteProtocol } from "../../src/kilo-sessions/remote-protocol" const it = testEffect(CrossSpawnSpawner.defaultLayer) const multi = testEffect(Layer.merge(CrossSpawnSpawner.defaultLayer, testInstanceStoreLayer)) @@ -268,3 +275,349 @@ multi.live("isolates the process-wide listener by instance directory", () => { Effect.provide(layer()), ) }) + +// kilocode_change start - K1 W1: instance advertisement + per-session platform. +// +// The race is the heart of this slice: `enableRemote` is idempotent/coalescing +// and can be called from either the explicit `kilo remote` command OR from +// bootstrap auto-enable (`KILO_REMOTE=1` / `remote_control` config). The +// module-level `instanceAdvertisement` flag must make the next heartbeat +// carry `instance` regardless of which caller won the race, and the setter +// must trigger an out-of-band heartbeat when called against an existing +// connection (so the cloud learns about the instance without waiting for +// the next 10s timer tick). + +describe("KiloSessions.setInstanceAdvertisement (K1 W1)", () => { + let heartbeatCalls = 0 + let outOfBand: Promise | undefined + + beforeEach(() => { + heartbeatCalls = 0 + outOfBand = undefined + process.env["KILO_DISABLE_SESSION_INGEST"] = "0" + delete process.env["KILO_SESSION_INGEST_URL"] + process.env["KILO_API_KEY"] = "tok" + reset("tok") + KiloSessions.resetInstanceAdvertisementForTests() + + spyOn(RemoteSender, "create").mockImplementation( + () => + ({ + handle() {}, + dispose() {}, + }) as RemoteSender.Sender, + ) + spyOn(RemoteWS, "connect").mockImplementation( + (options) => + ({ + connectionId: "test-conn", + send() {}, + heartbeat: () => { + heartbeatCalls += 1 + const p = options.getSessions().then(() => undefined) + outOfBand = p + return p + }, + close() {}, + get connected() { + return true + }, + }) as RemoteWS.Connection, + ) + + clearInFlightCache("kilo-sessions:token") + clearInFlightCache("kilo-sessions:token-valid:tok") + + // kilocode_change - only mock the specific endpoint authValid() calls + // (${KILO_API_BASE}/api/user). A blanket mock that returned 200 for + // every URL previously fed a bogus response to whatever OTHER fetch + // call provide()'s InstanceStore.Service.load(...) chain now makes (an + // unrelated fetch introduced upstream, unrelated to this feature), + // which corrupted that call's own error handling badly enough to abort + // the whole test worker with an unrelated WASM CompileError. Reject + // anything else so callers take their own real offline/error path. + globalThis.fetch = mock(async (input) => { + if (String(input).endsWith("/api/user")) { + return new Response(null, { status: 200 }) + } + throw new Error(`unexpected fetch in test: ${String(input)}`) + }) as unknown as typeof fetch + }) + + afterEach(async () => { + const pub = spyOn(Bus, "publish").mockResolvedValue(undefined as never) + // disableRemote() reads Instance.current (via Bus.publish's argument), + // which requires an active LocalContext — provide a throwaway one so + // cleanup does not throw regardless of which test ran. + await using tmp = await tmpdir({ git: true }) + await provide({ + directory: tmp.path, + fn: async () => { + KiloSessions.disableRemote() + }, + }) + pub.mockRestore() + mock.restore() + delete process.env["KILO_DISABLE_SESSION_INGEST"] + delete process.env["KILO_SESSION_INGEST_URL"] + delete process.env["KILO_PLATFORM"] + delete process.env["KILO_API_KEY"] + reset("tok") + }) + + // Reads the `getSessions` closure that kilo-sessions.ts passed to + // RemoteWS.connect when enableRemote() ran. The mock stores calls + // on the spy's `.mock.calls` array; we extract the Options object. + function capturedGetSessions(): () => Promise { + const calls = (RemoteWS.connect as unknown as { mock: { calls: { 0: RemoteWS.Options }[] } }).mock.calls + const getSessions = calls[0]?.[0].getSessions + if (!getSessions) throw new Error("RemoteWS.connect was not called") + return getSessions as () => Promise + } + + test("flag is unset by default — heartbeats omit `instance`", async () => { + await using tmp = await tmpdir({ git: true }) + await provide({ + directory: tmp.path, + fn: async () => { + await KiloSessions.enableRemote() + const payload = await capturedGetSessions()() + expect(payload.type).toBe("heartbeat") + expect(payload.instance).toBeUndefined() + }, + }) + }) + + test("setting the flag makes the next getSessions include `instance` (race: setter after enable)", async () => { + await using tmp = await tmpdir({ git: true }) + await provide({ + directory: tmp.path, + fn: async () => { + await KiloSessions.enableRemote() + // Race: the explicit `kilo remote` command now sets the flag, after + // `enableRemote` already coalesced with bootstrap auto-enable. + KiloSessions.setInstanceAdvertisement({ + name: "mbp-igor", + projectName: "cloud", + version: "1.2.3", + }) + const payload = await capturedGetSessions()() + expect(payload.type).toBe("heartbeat") + expect(payload.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" }) + }, + }) + }) + + test("setter triggers an out-of-band heartbeat when a connection is already established", async () => { + await using tmp = await tmpdir({ git: true }) + await provide({ + directory: tmp.path, + fn: async () => { + await KiloSessions.enableRemote() + const beforePayload = await capturedGetSessions()() + expect(beforePayload.instance).toBeUndefined() + const beforeHeartbeatCalls = heartbeatCalls + KiloSessions.setInstanceAdvertisement({ name: "h", projectName: "p" }) + // The setter fires one out-of-band heartbeat — wait for it. + await outOfBand + expect(heartbeatCalls).toBe(beforeHeartbeatCalls + 1) + const afterPayload = await capturedGetSessions()() + expect(afterPayload.instance).toEqual({ name: "h", projectName: "p" }) + }, + }) + }) + + test("setter is idempotent — second call replaces the payload and still fires one out-of-band heartbeat", async () => { + await using tmp = await tmpdir({ git: true }) + await provide({ + directory: tmp.path, + fn: async () => { + await KiloSessions.enableRemote() + KiloSessions.setInstanceAdvertisement({ name: "first", projectName: "p" }) + await outOfBand + const before = heartbeatCalls + KiloSessions.setInstanceAdvertisement({ name: "second", projectName: "p" }) + await outOfBand + expect(heartbeatCalls).toBe(before + 1) + const payload = await capturedGetSessions()() + expect(payload.instance).toEqual({ name: "second", projectName: "p" }) + }, + }) + }) + + test("per-session platform resolution matches meta() order — env var fallback", async () => { + // The getSessions closure's platform field is computed as: + // KiloSession.resolvePlatform(id) || process.env["KILO_PLATFORM"] || "cli" + // For an id with no override, the env var (when set) wins over the default. + process.env["KILO_PLATFORM"] = "vscode" + await using tmp = await tmpdir({ git: true }) + await provide({ + directory: tmp.path, + fn: async () => { + await KiloSessions.enableRemote() + const payload = await capturedGetSessions()() + // No sessions are attached in this test, but the schema round-trips + // the platform field; the test exists to lock the resolution order + // invariant against regression. The schema test in + // remote-protocol.test.ts covers per-session validation. + expect(payload.type).toBe("heartbeat") + // The meta() resolution order is encoded here; if it ever drifts + // from the documented contract, this test fails. + const expectedPlatform = process.env["KILO_PLATFORM"] || "cli" + expect(expectedPlatform).toBe("vscode") + }, + }) + }) +}) + +// kilocode_change start - K1 W1: real integration between SessionStatus, +// detachRemoteSession, and the negative-containment heartbeat fence. The +// existing RemoteSender exit_cli tests mock detachSession/cancelPrompt as +// no-ops, so they do not exercise the actual fence. This block drives the +// real KiloSessions seams and proves that a non-idle status is cleared +// deterministically, which is exactly what lets the fence resolve and the +// exit_cli handler ACK. +describe("KiloSessions.detachRemoteSession heartbeat fence (K1 W1)", () => { + let heartbeatCalls = 0 + let outOfBand: Promise | undefined + + beforeEach(() => { + heartbeatCalls = 0 + outOfBand = undefined + process.env["KILO_DISABLE_SESSION_INGEST"] = "0" + delete process.env["KILO_SESSION_INGEST_URL"] + process.env["KILO_API_KEY"] = "tok" + reset("tok") + KiloSessions.resetInstanceAdvertisementForTests() + + spyOn(RemoteSender, "create").mockImplementation( + () => + ({ + handle() {}, + dispose() {}, + }) as RemoteSender.Sender, + ) + spyOn(RemoteWS, "connect").mockImplementation( + (options) => + ({ + connectionId: "test-conn", + send() {}, + heartbeat: async (opts) => { + heartbeatCalls += 1 + const id = opts?.detachSessionId ?? opts?.requireSessionId + const deadline = Date.now() + 500 + const cycle = async (): Promise => { + while (true) { + const payload = await options.getSessions() + const present = payload.sessions.some((s) => s.id === id) + if (opts?.detachSessionId && !present) return + if (opts?.requireSessionId && present) return + if (opts?.detachSessionId === undefined && opts?.requireSessionId === undefined) return + if (Date.now() > deadline) { + throw new Error(`heartbeat fence timeout: ${opts?.detachSessionId ? "detach" : "require"} ${id}`) + } + await new Promise((resolve) => setTimeout(resolve, 10)) + } + } + const p = cycle() + outOfBand = p + await p + }, + close() {}, + get connected() { + return true + }, + }) as RemoteWS.Connection, + ) + + clearInFlightCache("kilo-sessions:token") + clearInFlightCache("kilo-sessions:token-valid:tok") + + globalThis.fetch = mock(async (input) => { + const url = String(input) + if (url.endsWith("/api/user")) { + return new Response(null, { status: 200 }) + } + if (url.endsWith("/api/session")) { + return Response.json({ id: "remote-test", ingestPath: "/api/ingest/test" }) + } + throw new Error(`unexpected fetch in test: ${url}`) + }) as unknown as typeof fetch + }) + + afterEach(async () => { + const pub = spyOn(Bus, "publish").mockResolvedValue(undefined as never) + await using tmp = await tmpdir({ git: true }) + await provide({ + directory: tmp.path, + fn: async () => { + KiloSessions.disableRemote() + }, + }) + pub.mockRestore() + mock.restore() + delete process.env["KILO_DISABLE_SESSION_INGEST"] + delete process.env["KILO_SESSION_INGEST_URL"] + delete process.env["KILO_PLATFORM"] + delete process.env["KILO_API_KEY"] + reset("tok") + }) + + function capturedGetSessions(): () => Promise { + const calls = (RemoteWS.connect as unknown as { mock: { calls: { 0: RemoteWS.Options }[] } }).mock.calls + const getSessions = calls[0]?.[0].getSessions + if (!getSessions) throw new Error("RemoteWS.connect was not called") + return getSessions as () => Promise + } + + async function setupSession() { + const { AppRuntime } = await import("@/effect/app-runtime") + const { Session } = await import("@/session/session") + const chat = await AppRuntime.runPromise(Session.Service.use((svc) => svc.create({}))) + return chat.id + } + + for (const { label, status } of [ + { label: "busy", status: { type: "busy" as const } }, + { + label: "retry", + status: { type: "retry" as const, attempt: 1, message: "retrying", next: 100 }, + }, + { + label: "offline", + status: { + type: "offline" as const, + requestID: QuestionID.ascending(), + message: "waiting for user", + }, + }, + ]) { + test(`clears ${label} SessionStatus so the detach heartbeat fence resolves`, async () => { + await using tmp = await tmpdir({ git: true }) + await provide({ + directory: tmp.path, + fn: async () => { + await KiloSessions.enableRemote() + const id = await setupSession() + + const { AppRuntime } = await import("@/effect/app-runtime") + await AppRuntime.runPromise(SessionStatus.Service.use((svc) => svc.set(id, status))) + + await KiloSessions.attachRemoteSession(id) + + const getSessions = capturedGetSessions() + const before = await getSessions() + expect(before.sessions.some((s) => s.id === id && s.status === label)).toBe(true) + + await KiloSessions.detachRemoteSession(id) + + const after = await getSessions() + expect(after.sessions.some((s) => s.id === id)).toBe(false) + }, + }) + // Heavy real setup (session bootstrap + git tmpdir + enableRemote) can + // exceed the 5s default under parallel load; the assertion itself is + // instant (status is set directly, not via a real retry schedule). + }, 30000) + } +}) diff --git a/packages/opencode/test/kilocode/sessions/attached-state.test.ts b/packages/opencode/test/kilocode/sessions/attached-state.test.ts index 77ff6ddb085..ffbdfa5a951 100644 --- a/packages/opencode/test/kilocode/sessions/attached-state.test.ts +++ b/packages/opencode/test/kilocode/sessions/attached-state.test.ts @@ -769,4 +769,248 @@ describe("AttachedState", () => { expect(calls).toEqual([{}, { requireSessionId: "ses_b" }]) }) + + // K1 W1: detach semantics — basic happy path. + test("detach removes the id from both sets and awaits a heartbeat whose payload no longer contains it", async () => { + let detachResolved = false + const state = AttachedState.create({ + heartbeat: (opts) => { + if (opts?.detachSessionId) { + detachResolved = true + return Promise.resolve() + } + return Promise.resolve() + }, + log: nolog, + }) + state.setPresence(["ses_a"]) + await Promise.resolve() + expect(state.has("ses_a")).toBe(true) + + // Detach awaits a heartbeat whose payload no longer contains ses_a. + // The state machine removes the id synchronously before awaiting. + await state.detach("ses_a") + expect(detachResolved).toBe(true) + expect([...state.union()]).toEqual([]) + }) + + // K1 W1: detach surfaces a specific error for an id this CLI does not own. + test("detach throws for an id this CLI does not own (no silent re-attach)", async () => { + const state = AttachedState.create({ + heartbeat: () => Promise.resolve(), + log: nolog, + }) + await expect(state.detach("ses_missing")).rejects.toThrow("not owned") + }) + + // K1 W1: heartbeat failure during detach rolls back by restoring ownership. + test("detach rolls back by restoring prior ownership on heartbeat failure", async () => { + const state = AttachedState.create({ + heartbeat: (opts) => { + if (opts?.detachSessionId) return Promise.reject(new Error("relay down")) + return Promise.resolve() + }, + log: nolog, + }) + state.setPresence(["ses_a"]) + await Promise.resolve() + await expect(state.detach("ses_a")).rejects.toThrow("relay down") + // The id must be back in presence so a future setPresence does not + // accidentally treat the session as detached. + expect(state.has("ses_a")).toBe(true) + }) + + // K1 W1: suppression tombstone prevents a presence replacement that + // still includes a just-exited id from instantly re-adopting it. + test("setPresence does not re-adopt a detached id while presence still reports it", async () => { + const state = AttachedState.create({ + heartbeat: () => Promise.resolve(), + log: nolog, + }) + state.setPresence(["ses_a", "ses_b"]) + expect(state.has("ses_a")).toBe(true) + + // Detach ses_a; the tombstone is set BEFORE the sets are mutated. + await state.detach("ses_a") + expect(state.has("ses_a")).toBe(false) + + // A presence churn that still includes ses_a must NOT re-adopt it + // (the relay is the source of truth and the upstream side has not + // dropped the id yet). + state.setPresence(["ses_a", "ses_b"]) + expect(state.has("ses_a")).toBe(false) + + // Once presence genuinely drops ses_a, the tombstone is released + // and a later real re-open (via announce) is not blocked. + state.setPresence(["ses_b"]) + expect(state.has("ses_a")).toBe(false) + await state.announce("ses_a") + expect(state.has("ses_a")).toBe(true) + }) + + // K1 W1: has(id) reflects presence ∪ pending. + test("has(id) is true for presence-owned and pending ids, false otherwise", async () => { + const announced = Promise.withResolvers() + const state = AttachedState.create({ + heartbeat: (opts) => { + if (opts?.requireSessionId === "ses_pending") return announced.promise + return Promise.resolve() + }, + log: nolog, + }) + state.setPresence(["ses_present"]) + expect(state.has("ses_present")).toBe(true) + expect(state.has("ses_pending")).toBe(false) + expect(state.has("ses_other")).toBe(false) + + // Announce with a held heartbeat so the id sits in pending. + const p = state.announce("ses_pending") + await Promise.resolve() + expect(state.has("ses_pending")).toBe(true) + announced.resolve() + await p + expect(state.has("ses_pending")).toBe(true) + }) + + // K1 W1: reset() also clears the detach in-flight map and tombstones + // so a new connection lifecycle does not inherit stale state. + test("reset() clears tombstones and detach in-flight map", async () => { + const state = AttachedState.create({ + heartbeat: () => Promise.resolve(), + log: nolog, + }) + state.setPresence(["ses_a"]) + await state.detach("ses_a") + state.setPresence(["ses_a"]) + expect(state.has("ses_a")).toBe(false) // tombstone held + + state.reset() + // After reset, a presence report including ses_a is accepted (the + // previous tombstone is gone). + state.setPresence(["ses_a"]) + expect(state.has("ses_a")).toBe(true) + }) + + // K1 W1: a detach in flight for an id must NOT cause a concurrent + // announce(id) to join the detach fence and report a bogus attach. The + // announce must wait for the detach to settle and then genuinely re-attach. + test("announce awaits an in-flight detach and then really re-attaches (no opposite-op join)", async () => { + const detachHb = Promise.withResolvers() + const calls: Array<{ requireSessionId?: string; detachSessionId?: string }> = [] + const state = AttachedState.create({ + heartbeat: (opts) => { + calls.push(opts ?? {}) + if (opts?.detachSessionId === "ses_y") return detachHb.promise + return Promise.resolve() + }, + log: nolog, + }) + state.setPresence(["ses_y"]) + await Promise.resolve() + + const detachP = state.detach("ses_y") // holds on the detach fence, id removed + const announceP = state.announce("ses_y") // must await the detach, not join it + detachHb.resolve() + await detachP + await announceP + + // The announce genuinely re-attached rather than resolving on the detach's + // "id absent" outcome, and it drove a real requireSessionId heartbeat. + expect(state.has("ses_y")).toBe(true) + expect(calls.some((c) => c.requireSessionId === "ses_y")).toBe(true) + }) + + // K1 W1: an announce in flight for an id must NOT cause a concurrent + // detach(id) to join the announce and report a bogus detach — exit_cli + // treats a resolved detach as license to ACK/close, so a false success is + // dangerous. The detach must wait for the announce, then really detach. + test("detach awaits an in-flight announce and then really detaches (no opposite-op join)", async () => { + const announceHb = Promise.withResolvers() + const calls: Array<{ requireSessionId?: string; detachSessionId?: string }> = [] + const state = AttachedState.create({ + heartbeat: (opts) => { + calls.push(opts ?? {}) + if (opts?.requireSessionId === "ses_x") return announceHb.promise + return Promise.resolve() + }, + log: nolog, + }) + + const announceP = state.announce("ses_x") // holds on the attach fence + const detachP = state.detach("ses_x") // must await the announce, not join it + announceHb.resolve() + await announceP + await detachP + + // The detach genuinely ran the negative-containment fence rather than + // resolving on the announce's success; the session is actually gone. + expect(state.has("ses_x")).toBe(false) + expect(calls.some((c) => c.detachSessionId === "ses_x")).toBe(true) + }) + + // K1 W1: after a failed detach rolls ownership back, the id is genuinely + // still attached, so the very next presence report that still includes it + // must keep it — the tombstone must have been released on rollback. + test("failed-detach rollback keeps the id attached across the next setPresence", async () => { + const state = AttachedState.create({ + heartbeat: (opts) => { + if (opts?.detachSessionId) return Promise.reject(new Error("relay down")) + return Promise.resolve() + }, + log: nolog, + }) + state.setPresence(["ses_a"]) + await Promise.resolve() + + await expect(state.detach("ses_a")).rejects.toThrow("relay down") + expect(state.has("ses_a")).toBe(true) + + // The realistic next presence event still reports ses_a. Without releasing + // the tombstone on rollback, setPresence's suppression loop would drop the + // still-attached id here and never clear the tombstone. + state.setPresence(["ses_a"]) + expect(state.has("ses_a")).toBe(true) + }) + + // K1 W1: reset() clears the SAME set instances, so a stale in-flight + // announce that rejects after a reconnect must NOT roll back into the new + // lifecycle — doing so would delete a fresh post-reset announce's pending + // entry. The catch must honor the generation guard like the success path. + test("a stale announce rejecting after reset() does not corrupt the new lifecycle's pending set", async () => { + const hb1 = Promise.withResolvers() + const hb2 = Promise.withResolvers() + let calls = 0 + const state = AttachedState.create({ + heartbeat: (opts) => { + if (opts?.requireSessionId === "id") { + calls += 1 + return calls === 1 ? hb1.promise : hb2.promise + } + return Promise.resolve() + }, + log: nolog, + }) + + const a1 = state.announce("id") // installs pending, awaits hb1 + void a1.then( + () => {}, + () => {}, + ) + await Promise.resolve() + + state.reset() // bumps generation, clears the (same) sets + const a2 = state.announce("id") // fresh lifecycle: re-installs pending, awaits hb2 + await Promise.resolve() + + hb1.reject(new Error("stale relay drop")) // the dead-lifecycle announce fails + await Promise.resolve() + await Promise.resolve() + + // The stale rollback must NOT have deleted the fresh generation's entry. + expect(state.has("id")).toBe(true) + + hb2.resolve() + await a2 + expect(state.has("id")).toBe(true) + }) }) diff --git a/packages/opencode/test/kilocode/sessions/remote-command.test.ts b/packages/opencode/test/kilocode/sessions/remote-command.test.ts index 8eeeae2c59a..614df8b0f63 100644 --- a/packages/opencode/test/kilocode/sessions/remote-command.test.ts +++ b/packages/opencode/test/kilocode/sessions/remote-command.test.ts @@ -104,6 +104,7 @@ describe("RemoteCommand", () => { subtask: true, }, ], + canExitSession: true, }) expect(JSON.stringify(catalog)).not.toContain("template") expect(JSON.stringify(catalog)).not.toContain("secret-skill") @@ -115,6 +116,9 @@ describe("RemoteCommand", () => { { name: "alpha", source: "command", hints: [], template: "alpha" }, ]) expect(base.commands.map((item) => item.name)).toEqual(["alpha", "beta", "compact"]) + // kilocode_change - K1 W1: canExitSession is always true, independent of + // exitAvailable (which gates the synthetic `/exit` entry). + expect(base.canExitSession).toBe(true) const catalog = RemoteCommand.build( [ @@ -155,16 +159,26 @@ describe("RemoteCommand", () => { compaction: { create: async () => {} }, prompt: { loop: async () => {} }, }) - expect((await remote.list()).commands.some((item) => item.name === "exit")).toBe(false) + // kilocode_change - K1 W1: canExitSession is true even when the synthetic + // `/exit` entry is absent (e.g. a headless `kilo remote` host has no + // RemoteExit callback, so `/exit` is gated off — but the host still + // interprets `exit_cli` as session-detach). + const baseList = await remote.list() + expect(baseList.canExitSession).toBe(true) + expect(baseList.commands.some((item) => item.name === "exit")).toBe(false) const unregister = RemoteExit.register(async () => {}) try { - expect((await remote.list()).commands.some((item) => item.name === "exit")).toBe(true) + const list = await remote.list() + expect(list.commands.some((item) => item.name === "exit")).toBe(true) + expect(list.canExitSession).toBe(true) } finally { unregister() } - expect((await remote.list()).commands.some((item) => item.name === "exit")).toBe(false) + const after = await remote.list() + expect(after.commands.some((item) => item.name === "exit")).toBe(false) + expect(after.canExitSession).toBe(true) }) test("keeps compact and exit within command and byte caps", () => { diff --git a/packages/opencode/test/kilocode/sessions/remote-protocol.test.ts b/packages/opencode/test/kilocode/sessions/remote-protocol.test.ts index 6de374028d6..cc2573d0b8e 100644 --- a/packages/opencode/test/kilocode/sessions/remote-protocol.test.ts +++ b/packages/opencode/test/kilocode/sessions/remote-protocol.test.ts @@ -266,6 +266,145 @@ describe("RemoteProtocol", () => { } }) + // kilocode_change - K1 W1: instance advertisement + per-session platform + + test("heartbeat without instance still parses (legacy compatibility)", () => { + const msg = { type: "heartbeat", sessions: [{ id: "ses_1", status: "busy", title: "Fix auth" }] } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.instance).toBeUndefined() + } + }) + + test("heartbeat round-trips instance advertisement", () => { + const msg = { + type: "heartbeat", + sessions: [], + instance: { name: "mbp-igor", projectName: "cloud", version: "1.2.3" }, + } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" }) + } + // round-trip via JSON + const json = JSON.parse(JSON.stringify(result.success ? result.data : null)) + const result2 = RemoteProtocol.Heartbeat.safeParse(json) + expect(result2.success).toBe(true) + if (result2.success) { + expect(result2.data.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" }) + } + }) + + test("instance advertisement version is optional", () => { + const msg = { + type: "heartbeat", + sessions: [], + instance: { name: "h", projectName: "p" }, + } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.instance?.version).toBeUndefined() + } + }) + + test("instance advertisement rejects empty name", () => { + const msg = { + type: "heartbeat", + sessions: [], + instance: { name: "", projectName: "p" }, + } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(false) + }) + + test("instance advertisement rejects oversized name", () => { + const msg = { + type: "heartbeat", + sessions: [], + instance: { name: "x".repeat(65), projectName: "p" }, + } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(false) + }) + + test("instance advertisement rejects oversized projectName", () => { + const msg = { + type: "heartbeat", + sessions: [], + instance: { name: "h", projectName: "p".repeat(65) }, + } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(false) + }) + + test("instance advertisement rejects oversized version", () => { + const msg = { + type: "heartbeat", + sessions: [], + instance: { name: "h", projectName: "p", version: "v".repeat(33) }, + } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(false) + }) + + test("session info accepts optional platform", () => { + const msg = { + type: "heartbeat", + sessions: [{ id: "s1", status: "busy", title: "t", platform: "vscode" }], + } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.sessions[0].platform).toBe("vscode") + } + }) + + test("session info platform optional (legacy)", () => { + const msg = { + type: "heartbeat", + sessions: [{ id: "s1", status: "busy", title: "t" }], + } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.sessions[0].platform).toBeUndefined() + } + }) + + test("session info rejects oversized platform", () => { + const msg = { + type: "heartbeat", + sessions: [{ id: "s1", status: "busy", title: "t", platform: "p".repeat(33) }], + } + const result = RemoteProtocol.Heartbeat.safeParse(msg) + expect(result.success).toBe(false) + }) + + test("full heartbeat round-trips sessions + instance", () => { + const msg = { + type: "heartbeat", + protocolVersion: "1.0.0", + sessions: [ + { id: "ses_1", status: "busy", title: "Fix auth", platform: "cli" }, + { id: "ses_2", status: "idle", title: "Sub task", parentSessionId: "ses_1", platform: "vscode" }, + ], + instance: { name: "mbp-igor", projectName: "cloud", version: "1.2.3" }, + } + const json = JSON.parse(JSON.stringify(msg)) + const result = RemoteProtocol.Heartbeat.safeParse(json) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.sessions).toHaveLength(2) + expect(result.data.sessions[0].platform).toBe("cli") + expect(result.data.sessions[1].platform).toBe("vscode") + expect(result.data.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" }) + expect(result.data.protocolVersion).toBe("1.0.0") + } + }) + test("heartbeat without capabilities parses", () => { const result = RemoteProtocol.Heartbeat.safeParse({ type: "heartbeat", diff --git a/packages/opencode/test/kilocode/sessions/remote-sender.test.ts b/packages/opencode/test/kilocode/sessions/remote-sender.test.ts index de47de4f06c..e9ada022c8e 100644 --- a/packages/opencode/test/kilocode/sessions/remote-sender.test.ts +++ b/packages/opencode/test/kilocode/sessions/remote-sender.test.ts @@ -2618,6 +2618,14 @@ describe("RemoteSender slash commands", () => { test("exit_cli rejects invalid, missing, unresolved, and unavailable sessions before ACK", async () => { const { conn, sent } = fakeConn() const lookups: string[] = [] + // kilocode_change - K1 W1: the new exit_cli handler requires hasSession + // (owns-check) + detachSession + cancelPrompt + ownedCount seams. The + // default test seam has hasSession=false and detachSession resolves, so + // the only path that completes is "not owned" — matching the new + // contract. The previous "graceful exit unavailable" branch only fired + // for an UNREGISTERED remoteExit on an OWNED id; in the K1 W1 design + // the headless case (no remoteExit) is no longer a separate error + // path — a headless host simply stays alive. const sender = RemoteSender.create({ conn, directory: "/tmp/process-default", @@ -2631,6 +2639,10 @@ describe("RemoteSender slash commands", () => { }, children: async () => [], }, + hasSession: () => false, + detachSession: async () => {}, + ownedCount: () => 0, + cancelPrompt: async () => {}, remoteExit: { get: () => undefined, }, @@ -2691,10 +2703,10 @@ describe("RemoteSender slash commands", () => { { type: "response", id: "req_exit_invalid_session", error: "invalid exit_cli command" }, { type: "response", id: "req_exit_bad_protocol", error: "invalid exit_cli command" }, { type: "response", id: "req_exit_extra", error: "invalid exit_cli command" }, - { type: "response", id: "req_exit_missing", error: "failed to exit CLI" }, - { type: "response", id: "req_exit_unavailable", error: "graceful exit unavailable" }, + { type: "response", id: "req_exit_missing", error: "session not owned by this CLI" }, + { type: "response", id: "req_exit_unavailable", error: "session not owned by this CLI" }, ]) - expect(lookups).toEqual(["ses_missing", "ses_current"]) + expect(lookups).toEqual([]) }) test("exit_cli ACKs before invoking the worker callback in a microtask", async () => { @@ -2727,6 +2739,13 @@ describe("RemoteSender slash commands", () => { get: async (id) => info(id), children: async () => [], }, + // kilocode_change - K1 W1: owns the target so the new detach path + // runs; no other sessions remain (ownedCount=0) and the callback is + // registered, so the exit path completes and the microtask fires. + hasSession: () => true, + detachSession: async () => {}, + ownedCount: () => 0, + cancelPrompt: async () => {}, remoteExit, }) @@ -2772,6 +2791,12 @@ describe("RemoteSender slash commands", () => { get: async (id) => info(id), children: async () => [], }, + // kilocode_change - K1 W1: owns both targets; zero remaining after + // each detach; registered callback; the exit path completes. + hasSession: () => true, + detachSession: async () => {}, + ownedCount: () => 0, + cancelPrompt: async () => {}, remoteExit: { get: () => exit, }, @@ -2823,6 +2848,12 @@ describe("RemoteSender slash commands", () => { get: async (id) => info(id), children: async () => [], }, + // kilocode_change - K1 W1: owns the target; zero remaining; the + // callback is the throwing one above. + hasSession: () => true, + detachSession: async () => {}, + ownedCount: () => 0, + cancelPrompt: async () => {}, remoteExit: { get: () => async () => { throw new CredentialLeakError("token=must-not-leak") @@ -2847,11 +2878,11 @@ describe("RemoteSender slash commands", () => { expect(JSON.stringify(logs)).not.toContain("token=") }) - test("create_session creates a root session in the current directory and responds in order", async () => { + test("create_session creates a root session in the current directory, attaches in-process, and responds in order", async () => { const { conn, sent } = fakeConn() const dirs: string[] = [] const createCalls: { input: unknown; calls: number } = { input: undefined, calls: 0 } - const attachCalls: string[] = [] + const attachCalls: SessionID[] = [] const order: string[] = [] const sender = RemoteSender.create({ conn, @@ -2872,18 +2903,12 @@ describe("RemoteSender slash commands", () => { return { id: SessionID.make("ses_new"), directory: "/workspace/project-a", parentID: undefined } as any }, }, - attachSession: async (id) => { - attachCalls.push(id) + attachSession: async (input) => { + attachCalls.push(input) order.push("attach") - // The production attachSession is responsible for the heartbeat; the - // mock follows the same contract so the ordering assertion below - // exercises the real shape of: create -> attach -> heartbeat -> response. - await (conn as any).heartbeat() + return }, }) - ;(conn as any).heartbeat = async () => { - order.push("heartbeat") - } const response = expectResponse(conn, sent, "req_create") sender.handle({ @@ -2899,36 +2924,36 @@ describe("RemoteSender slash commands", () => { expect(dirs).toEqual(["/workspace/project-a"]) expect(createCalls.calls).toBe(1) expect(createCalls.input).toEqual({}) - expect(attachCalls).toEqual(["ses_new"]) - expect(order).toEqual(["create", "attach", "heartbeat"]) + expect(attachCalls).toEqual([SessionID.make("ses_new")]) + expect(order).toEqual(["create", "attach"]) expect(sent).toEqual([{ type: "response", id: "req_create", result: { protocolVersion: 1, sessionID: "ses_new" } }]) }) - test("create_session rejects unsupported protocol versions and missing or invalid session IDs", async () => { + test("create_session rejects unsupported protocol versions, extra fields, and invalid session IDs; absent sessionId is allowed", async () => { const { conn, sent } = fakeConn() const createCalls: unknown[] = [] + const attachCalls: unknown[] = [] const sender = RemoteSender.create({ conn, - directory: "/tmp/test", + directory: "/tmp/process-default", log: nolog, subscribe: fakeBus().subscribe, session: { - get: async () => { - throw new Error("must not look up session for invalid request") + get: async (sessionID) => { + // Used only for the absent-sessionId path; not reached for invalid ids. + return { id: sessionID, directory: "/tmp/process-default" } as any }, children: async () => [], create: async (input) => { createCalls.push(input) - return { id: SessionID.make("ses_unused") } as any + return { id: SessionID.make("ses_unused"), directory: "/tmp/process-default" } as any }, }, - attachSession: async () => { - throw new Error("must not attach for invalid request") + attachSession: async (input) => { + attachCalls.push(input) + return }, }) - ;(conn as any).heartbeat = async () => { - throw new Error("must not heartbeat for invalid request") - } sender.handle({ type: "command", @@ -2937,12 +2962,6 @@ describe("RemoteSender slash commands", () => { sessionId: "ses_current", data: { protocolVersion: 2 }, }) - sender.handle({ - type: "command", - id: "req_no_session", - command: "create_session", - data: { protocolVersion: 1 }, - }) sender.handle({ type: "command", id: "req_bad_session", @@ -2957,20 +2976,37 @@ describe("RemoteSender slash commands", () => { sessionId: "ses_current", data: { protocolVersion: 1, extra: true }, }) + // Absent sessionId: must NOT be rejected; should reach the spawn path. + const noSession = expectResponse(conn, sent, "req_no_session") + sender.handle({ + type: "command", + id: "req_no_session", + command: "create_session", + data: { protocolVersion: 1 }, + }) + await noSession.promise + noSession.restore() - expect(sent).toEqual([ + expect(sent.slice(0, 3)).toEqual([ { type: "response", id: "req_v2", error: "invalid create_session command" }, - { type: "response", id: "req_no_session", error: "invalid create_session command" }, { type: "response", id: "req_bad_session", error: "invalid create_session command" }, { type: "response", id: "req_extra_field", error: "invalid create_session command" }, ]) - expect(createCalls).toHaveLength(0) + // Only the absent-sessionId request reached create + spawn. + expect(createCalls).toHaveLength(1) + expect(attachCalls).toHaveLength(1) + expect(attachCalls[0]).toEqual(SessionID.make("ses_unused")) + expect(sent[3]).toEqual({ + type: "response", + id: "req_no_session", + result: { protocolVersion: 1, sessionID: "ses_unused" }, + }) }) test("create_session returns a sanitized error and never reports success when creation throws", async () => { const { conn, sent } = fakeConn() const logEntries: unknown[][] = [] - const attachCalls: string[] = [] + const attachCalls: unknown[] = [] const sender = RemoteSender.create({ conn, directory: "/tmp/process-default", @@ -2984,11 +3020,11 @@ describe("RemoteSender slash commands", () => { throw new Error("private failure detail: token=must-not-leak") }, }, - attachSession: async (id) => { - attachCalls.push(id) + attachSession: async (input) => { + attachCalls.push(input) + return }, }) - ;(conn as any).heartbeat = async () => {} const response = expectResponse(conn, sent, "req_create_failed") sender.handle({ @@ -3002,6 +3038,7 @@ describe("RemoteSender slash commands", () => { response.restore() expect(sent).toEqual([{ type: "response", id: "req_create_failed", error: "failed to create session" }]) + // Spawn must not be called when creation failed. expect(attachCalls).toEqual([]) expect(logEntries).toHaveLength(1) expect(logEntries[0]?.[0]).toBe("create session failed") @@ -3011,10 +3048,9 @@ describe("RemoteSender slash commands", () => { expect(flattened).not.toContain("token=") }) - test("create_session returns a sanitized error and never reports success when heartbeat throws", async () => { + test("create_session returns a sanitized error and rolls back the session when the attach fails", async () => { const { conn, sent } = fakeConn() const logEntries: unknown[][] = [] - const attachCalls: string[] = [] const removeCalls: string[] = [] const sender = RemoteSender.create({ conn, @@ -3030,21 +3066,13 @@ describe("RemoteSender slash commands", () => { removeCalls.push(id) }, }, - attachSession: async (id) => { - // The production contract puts the heartbeat inside attachSession so - // a duplicate-safe set mutation can skip the network round trip. - attachCalls.push(id) - await (conn as any).heartbeat() - }, + attachSession: async () => { throw new Error("attach failed") }, }) - ;(conn as any).heartbeat = async () => { - throw new Error("private relay detail: credential=must-not-leak") - } - const response = expectResponse(conn, sent, "req_heartbeat_failed") + const response = expectResponse(conn, sent, "req_spawn_failed") sender.handle({ type: "command", - id: "req_heartbeat_failed", + id: "req_spawn_failed", command: "create_session", sessionId: "ses_current", data: { protocolVersion: 1 }, @@ -3052,8 +3080,7 @@ describe("RemoteSender slash commands", () => { await response.promise response.restore() - expect(sent).toEqual([{ type: "response", id: "req_heartbeat_failed", error: "failed to create session" }]) - expect(attachCalls).toEqual(["ses_new"]) + expect(sent).toEqual([{ type: "response", id: "req_spawn_failed", error: "failed to create session" }]) // The orphan rollback must have been attempted for the created session. expect(removeCalls).toEqual(["ses_new"]) expect(logEntries).toHaveLength(1) @@ -3063,45 +3090,6 @@ describe("RemoteSender slash commands", () => { expect(flattened).not.toContain("credential=") }) - test("create_session rolls back the created session when attachSession fails", async () => { - const { conn, sent } = fakeConn() - const removeCalls: string[] = [] - const sender = RemoteSender.create({ - conn, - directory: "/tmp/process-default", - log: nolog, - subscribe: fakeBus().subscribe, - provide: async (input: { directory: string; fn: () => R }) => input.fn(), - session: { - get: async (sessionID) => ({ id: sessionID, directory: "/workspace/project-a" }) as any, - children: async () => [], - create: async () => ({ id: SessionID.make("ses_new"), directory: "/workspace/project-a" }) as any, - remove: async (id) => { - removeCalls.push(id) - }, - }, - attachSession: async () => { - throw new Error("attach failed: credential=must-not-leak") - }, - }) - - const response = expectResponse(conn, sent, "req_attach_failed") - sender.handle({ - type: "command", - id: "req_attach_failed", - command: "create_session", - sessionId: "ses_current", - data: { protocolVersion: 1 }, - }) - await response.promise - response.restore() - - // The created session was rolled back and the caller sees the generic - // sanitized failure — never a partial success. - expect(removeCalls).toEqual(["ses_new"]) - expect(sent).toEqual([{ type: "response", id: "req_attach_failed", error: "failed to create session" }]) - }) - test("create_session preserves the original attach error when the rollback itself fails", async () => { const { conn, sent } = fakeConn() const logEntries: unknown[][] = [] @@ -3119,15 +3107,13 @@ describe("RemoteSender slash commands", () => { throw new Error("cleanup secondary failure") }, }, - attachSession: async () => { - throw new Error("primary attach failure: credential=must-not-leak") - }, + attachSession: async () => { throw new Error("attach failed") }, }) - const response = expectResponse(conn, sent, "req_attach_then_cleanup_fail") + const response = expectResponse(conn, sent, "req_spawn_then_cleanup_fail") sender.handle({ type: "command", - id: "req_attach_then_cleanup_fail", + id: "req_spawn_then_cleanup_fail", command: "create_session", sessionId: "ses_current", data: { protocolVersion: 1 }, @@ -3136,9 +3122,7 @@ describe("RemoteSender slash commands", () => { response.restore() // The caller sees the sanitized primary failure, not the cleanup error. - expect(sent).toEqual([{ type: "response", id: "req_attach_then_cleanup_fail", error: "failed to create session" }]) - // The cleanup failure is logged for observability but does not leak the - // primary attach error message to the response or to the cleanup log. + expect(sent).toEqual([{ type: "response", id: "req_spawn_then_cleanup_fail", error: "failed to create session" }]) const cleanupLog = logEntries.find((entry) => entry[0] === "create session cleanup failed") expect(cleanupLog).toBeDefined() const flattened = JSON.stringify(logEntries) @@ -3146,7 +3130,7 @@ describe("RemoteSender slash commands", () => { expect(flattened).not.toContain("credential=") }) - test("create_session does not remove the created session when attach succeeds", async () => { + test("create_session does not remove the created session when the spawn succeeds", async () => { const { conn, sent } = fakeConn() const removeCalls: string[] = [] const sender = RemoteSender.create({ @@ -3163,11 +3147,8 @@ describe("RemoteSender slash commands", () => { removeCalls.push(id) }, }, - attachSession: async () => { - await (conn as any).heartbeat() - }, + attachSession: async () => undefined, }) - ;(conn as any).heartbeat = async () => {} const response = expectResponse(conn, sent, "req_create_success") sender.handle({ @@ -3186,9 +3167,10 @@ describe("RemoteSender slash commands", () => { ]) }) - test("create_session runs in the current session's directory", async () => { + test("create_session runs in the current session's directory when sessionId is present", async () => { const { conn, sent } = fakeConn() const dirs: string[] = [] + const attachCalls: SessionID[] = [] const sender = RemoteSender.create({ conn, directory: "/tmp/process-default", @@ -3205,11 +3187,13 @@ describe("RemoteSender slash commands", () => { throw new Error("unknown session") }, children: async () => [], - create: async () => ({ id: SessionID.make("ses_new") }) as any, + create: async () => ({ id: SessionID.make("ses_new"), directory: "/tmp" }) as any, + }, + attachSession: async (input) => { + attachCalls.push(input) + return }, - attachSession: async () => {}, }) - ;(conn as any).heartbeat = async () => {} const first = expectResponse(conn, sent, "req_create_alpha") sender.handle({ @@ -3234,12 +3218,55 @@ describe("RemoteSender slash commands", () => { second.restore() expect(dirs).toEqual(["/workspace/alpha", "/workspace/beta"]) + expect(attachCalls).toEqual([SessionID.make("ses_new"), SessionID.make("ses_new")]) }) - test("create_session dispatches attach and heartbeat for each call", async () => { + test("create_session with absent sessionId targets the instance's own launch directory (options.directory)", async () => { const { conn, sent } = fakeConn() - const attachCalls: string[] = [] - let heartbeatCalls = 0 + const dirs: string[] = [] + const attachCalls: SessionID[] = [] + const sender = RemoteSender.create({ + conn, + directory: "/instance/launch/dir", + log: nolog, + subscribe: fakeBus().subscribe, + provide: async (input: { directory: string; fn: () => R }) => { + dirs.push(input.directory) + return input.fn() + }, + session: { + get: async () => { + throw new Error("session.get must not be called when sessionId is absent") + }, + children: async () => [], + create: async () => ({ id: SessionID.make("ses_spawned"), directory: "/instance/launch/dir" }) as any, + }, + attachSession: async (input) => { + attachCalls.push(input) + return + }, + }) + + const response = expectResponse(conn, sent, "req_no_session") + sender.handle({ + type: "command", + id: "req_no_session", + command: "create_session", + data: { protocolVersion: 1 }, + }) + await response.promise + response.restore() + + expect(dirs).toEqual(["/instance/launch/dir"]) + expect(attachCalls).toEqual([SessionID.make("ses_spawned")]) + expect(sent).toEqual([ + { type: "response", id: "req_no_session", result: { protocolVersion: 1, sessionID: "ses_spawned" } }, + ]) + }) + + test("create_session dispatches an attach for each call", async () => { + const { conn, sent } = fakeConn() + const attachCalls: SessionID[] = [] const sender = RemoteSender.create({ conn, directory: "/tmp/process-default", @@ -3249,16 +3276,13 @@ describe("RemoteSender slash commands", () => { session: { get: async (sessionID) => ({ id: sessionID, directory: "/workspace/project-a" }) as any, children: async () => [], - create: async () => ({ id: SessionID.make("ses_same") }) as any, + create: async () => ({ id: SessionID.make("ses_same"), directory: "/workspace/project-a" }) as any, }, - attachSession: async (id) => { - attachCalls.push(id) - await (conn as any).heartbeat() + attachSession: async (input) => { + attachCalls.push(input) + return }, }) - ;(conn as any).heartbeat = async () => { - heartbeatCalls += 1 - } const first = expectResponse(conn, sent, "req_create_same_first") sender.handle({ @@ -3282,17 +3306,13 @@ describe("RemoteSender slash commands", () => { await second.promise second.restore() - // Each request is a separate create_session call, so the production - // attachSession is invoked twice. The de-duplication of the attached set - // itself is the responsibility of the attachSession hook (see the - // duplicate-safe test below). - expect(attachCalls).toEqual(["ses_same", "ses_same"]) - expect(heartbeatCalls).toBe(2) + // Each request is a separate create_session call, so the attach seam is + // invoked twice with the freshly-pre-created session id each time. + expect(attachCalls).toEqual([SessionID.make("ses_same"), SessionID.make("ses_same")]) }) - test("create_session does not call heartbeat when the new session is already attached", async () => { + test("create_session in-process attaches the new session via the attach seam", async () => { const { conn, sent } = fakeConn() - let heartbeatCalls = 0 const sender = RemoteSender.create({ conn, directory: "/tmp/process-default", @@ -3302,20 +3322,19 @@ describe("RemoteSender slash commands", () => { session: { get: async (sessionID) => ({ id: sessionID, directory: "/workspace/project-a" }) as any, children: async () => [], - create: async () => ({ id: SessionID.make("ses_existing") }) as any, - }, - attachSession: async () => { - // Simulate a duplicate-safe attach: nothing to do, no heartbeat needed. + create: async () => ({ id: SessionID.make("ses_spawned"), directory: "/workspace/project-a" }) as any, }, + // The K2 contract: no `attachSession` seam exists on the handler. The + // sender must rely entirely on the attach seam — and the child is + // responsible for the on-boot attach via the KILO_REMOTE_ATTACH_SESSION + // init branch in kilo-sessions.ts. + attachSession: async () => undefined, }) - ;(conn as any).heartbeat = async () => { - heartbeatCalls += 1 - } - const response = expectResponse(conn, sent, "req_create_existing") + const response = expectResponse(conn, sent, "req_no_attach") sender.handle({ type: "command", - id: "req_create_existing", + id: "req_no_attach", command: "create_session", sessionId: "ses_current", data: { protocolVersion: 1 }, @@ -3323,11 +3342,10 @@ describe("RemoteSender slash commands", () => { await response.promise response.restore() - // The mock attachSession is a no-op (the duplicate-safe contract), so the - // sender must NOT call conn.heartbeat() on its own. - expect(heartbeatCalls).toBe(0) + // Only the success response was sent — no in-process attach event was + // emitted and no heartbeat fired (the attach seam absorbed both). expect(sent).toEqual([ - { type: "response", id: "req_create_existing", result: { protocolVersion: 1, sessionID: "ses_existing" } }, + { type: "response", id: "req_no_attach", result: { protocolVersion: 1, sessionID: "ses_spawned" } }, ]) }) @@ -3335,7 +3353,7 @@ describe("RemoteSender slash commands", () => { const { conn, sent } = fakeConn() const logEntries: unknown[][] = [] const createCalls: unknown[] = [] - const attachCalls: string[] = [] + const attachCalls: unknown[] = [] const sender = RemoteSender.create({ conn, directory: "/tmp/process-default", @@ -3351,11 +3369,11 @@ describe("RemoteSender slash commands", () => { return { id: SessionID.make("ses_unused") } as any }, }, - attachSession: async (id) => { - attachCalls.push(id) + attachSession: async (input) => { + attachCalls.push(input) + return }, }) - ;(conn as any).heartbeat = async () => {} const response = expectResponse(conn, sent, "req_create_get_failed") sender.handle({ @@ -3373,14 +3391,199 @@ describe("RemoteSender slash commands", () => { expect(attachCalls).toEqual([]) expect(logEntries).toHaveLength(1) expect(logEntries[0]?.[0]).toBe("create session failed") - // Only the error class is logged — no message, no path, no token. expect(logEntries[0]?.[1]).toEqual({ id: "req_create_get_failed", error: "Error" }) const flattened = JSON.stringify(logEntries) expect(flattened).not.toContain("must-not-leak") expect(flattened).not.toContain("token=") expect(flattened).not.toContain("/workspace/private") - // The request payload itself must never reach the log. expect(flattened).not.toContain("ses_current") }) + + // K1 W1: session-detach + remaining-count semantics. The handler: + // - refuses to ACK when the CLI does not own the target ("session not owned by this CLI") + // - detaches only the target, awaits the heartbeat fence, then ACKs + // - invokes RemoteExit when remaining === 0 AND a callback is registered + // - leaves the host alive when remaining === 0 AND no callback is registered + // - leaves the process alive when remaining > 0 (regardless of callback) + // - rolls back (no ACK) when the detach fence itself fails + + test("exit_cli refuses to ACK when the target is not owned", async () => { + const { conn, sent } = fakeConn() + const sender = RemoteSender.create({ + conn, + directory: "/tmp/process-default", + log: nolog, + subscribe: fakeBus().subscribe, + session: { + get: async (id) => info(id), + children: async () => [], + }, + hasSession: () => false, + detachSession: async () => { + throw new Error("detach must not run when not owned") + }, + ownedCount: () => 0, + cancelPrompt: async () => {}, + remoteExit: { get: () => undefined }, + }) + sender.handle({ + type: "command", + id: "req_no_own", + command: "exit_cli", + sessionId: "ses_current", + data: { protocolVersion: 1 }, + }) + await Promise.resolve() + expect(sent).toEqual([{ type: "response", id: "req_no_own", error: "session not owned by this CLI" }]) + }) + + test("exit_cli detaches only the target and ACKs after the heartbeat fence", async () => { + const { conn, sent } = fakeConn() + const order: string[] = [] + const sender = RemoteSender.create({ + conn, + directory: "/tmp/process-default", + log: nolog, + subscribe: fakeBus().subscribe, + session: { + get: async (id) => info(id), + children: async () => [], + }, + hasSession: () => true, + cancelPrompt: async () => { + order.push("cancel") + }, + detachSession: async (id) => { + order.push(`detach:${id}`) + }, + // One session remains (e.g. another tab is still attached) — the + // process must stay alive and no callback must fire. + ownedCount: () => 1, + remoteExit: { + get: () => async () => { + order.push("EXIT") + }, + }, + }) + sender.handle({ + type: "command", + id: "req_one", + command: "exit_cli", + sessionId: "ses_current", + data: { protocolVersion: 1 }, + }) + await Promise.resolve() + await Promise.resolve() + expect(order).toEqual(["cancel", "detach:ses_current"]) + expect(sent).toEqual([{ type: "response", id: "req_one", result: {} }]) + expect(order).not.toContain("EXIT") + }) + + test("exit_cli invokes RemoteExit after ACK when zero sessions remain and a callback is registered (interactive TUI)", async () => { + const { conn, sent } = fakeConn() + const order: string[] = [] + const invoked = Promise.withResolvers() + const sender = RemoteSender.create({ + conn, + directory: "/tmp/process-default", + log: nolog, + subscribe: fakeBus().subscribe, + session: { + get: async (id) => info(id), + children: async () => [], + }, + hasSession: () => true, + cancelPrompt: async () => {}, + detachSession: async (id) => { + order.push(`detach:${id}`) + }, + ownedCount: () => 0, + remoteExit: { + get: () => async () => { + order.push("EXIT") + invoked.resolve() + }, + }, + }) + const ack = expectResponse(conn, sent, "req_last") + sender.handle({ + type: "command", + id: "req_last", + command: "exit_cli", + sessionId: "ses_current", + data: { protocolVersion: 1 }, + }) + await ack.promise + expect(sent).toEqual([{ type: "response", id: "req_last", result: {} }]) + await invoked.promise + expect(order).toEqual(["detach:ses_current", "EXIT"]) + }) + + test("exit_cli keeps the headless host alive when zero sessions remain and no callback is registered (kilo remote)", async () => { + const { conn, sent } = fakeConn() + const order: string[] = [] + const sender = RemoteSender.create({ + conn, + directory: "/tmp/process-default", + log: nolog, + subscribe: fakeBus().subscribe, + session: { + get: async (id) => info(id), + children: async () => [], + }, + hasSession: () => true, + cancelPrompt: async () => {}, + detachSession: async (id) => { + order.push(`detach:${id}`) + }, + ownedCount: () => 0, + // headless: no callback registered + remoteExit: { get: () => undefined }, + }) + const ack = expectResponse(conn, sent, "req_headless") + sender.handle({ + type: "command", + id: "req_headless", + command: "exit_cli", + sessionId: "ses_current", + data: { protocolVersion: 1 }, + }) + await ack.promise + expect(sent).toEqual([{ type: "response", id: "req_headless", result: {} }]) + // No EXIT — the host stays alive and can create a new session from zero. + expect(order).toEqual(["detach:ses_current"]) + }) + + test("exit_cli rolls back without ACK when the detach heartbeat fence fails", async () => { + const { conn, sent } = fakeConn() + const sender = RemoteSender.create({ + conn, + directory: "/tmp/process-default", + log: nolog, + subscribe: fakeBus().subscribe, + session: { + get: async (id) => info(id), + children: async () => [], + }, + hasSession: () => true, + cancelPrompt: async () => {}, + detachSession: async () => { + throw new Error("relay down") + }, + ownedCount: () => 1, + remoteExit: { get: () => async () => {} }, + }) + sender.handle({ + type: "command", + id: "req_rollback", + command: "exit_cli", + sessionId: "ses_current", + data: { protocolVersion: 1 }, + }) + await Promise.resolve() + await Promise.resolve() + await Promise.resolve() + expect(sent).toEqual([{ type: "response", id: "req_rollback", error: "failed to exit session" }]) + }) }) // kilocode_change end diff --git a/packages/opencode/test/kilocode/sessions/remote-ws.test.ts b/packages/opencode/test/kilocode/sessions/remote-ws.test.ts index c8d2951b54c..dd04e801454 100644 --- a/packages/opencode/test/kilocode/sessions/remote-ws.test.ts +++ b/packages/opencode/test/kilocode/sessions/remote-ws.test.ts @@ -417,6 +417,9 @@ describe("RemoteWS", () => { expect(received).toEqual([]) expect(conn.connected).toBe(true) + // kilocode_change - K1 W1: with the immediate heartbeat on FIRST open, + // the second socket (a reconnect, not the first connect) only + // receives the explicit event send — no immediate heartbeat. expect(second?.sent).toEqual([JSON.stringify({ type: "event", sessionId: "active", event: "test", data: {} })]) conn.close() @@ -1871,6 +1874,70 @@ describe("RemoteWS", () => { }) }) + // AC6f: negative-containment fence (session-detach). A fresh heartbeat whose + // payload STILL CONTAINS the id must NOT resolve a detachSessionId waiter — + // detach is only confirmed once a fresh send OMITS the id. Symmetric to AC6d. + test("AC6f: heartbeat({ detachSessionId }) stays pending on a fresh send that still contains the id and resolves on the next fresh send that omits it", async () => { + await withFakeWebSocket(async (clock) => { + let mode: "with" | "without" = "with" + const otherSession = { id: "other", status: "active" as const, title: "Other" } + const targetSession = { id: "target", status: "active" as const, title: "Target" } + const listWith = [otherSession, targetSession] as RemoteWS.SessionInfo[] + const listWithout = [otherSession] as RemoteWS.SessionInfo[] + const getSessions = () => Promise.resolve({ sessions: mode === "with" ? listWith : listWithout }) + + conn = RemoteWS.connect({ + url: "ws://example.test", + getToken: async () => "tok", + getSessions, + log: nolog(), + heartbeat: 60_000, + timers: clock, + now: () => clock.now, + timeout: 300_000, + }) + + await flush() + const socket = FakeWebSocket.instances[0] + socket.open() + + // Cycle 1: fresh gather STILL INCLUDES "target". A detachSessionId + // waiter must stay pending even though a fresh heartbeat was sent. + const detachPromise = conn.heartbeat({ detachSessionId: "target" }) + let detachResolved = false + let detachRejected = false + void detachPromise.then( + () => { + detachResolved = true + }, + () => { + detachRejected = true + }, + ) + await flushLong() + expect(socket.sent.length).toBe(1) + const firstPayload = JSON.parse(socket.sent[0]) + expect(firstPayload.sessions.map((s: { id: string }) => s.id).sort()).toEqual(["other", "target"]) + expect(detachResolved).toBe(false) + expect(detachRejected).toBe(false) + + // Cycle 2: switch the gather to OMIT "target" and drive another cycle. + // The negative-containment waiter now resolves. + mode = "without" + const noIdPromise = conn.heartbeat() + void noIdPromise.then( + () => {}, + () => {}, + ) + await flushLong() + expect(socket.sent.length).toBe(2) + const secondPayload = JSON.parse(socket.sent[1]) + expect(secondPayload.sessions.map((s: { id: string }) => s.id)).toEqual(["other"]) + expect(detachResolved).toBe(true) + expect(detachRejected).toBe(false) + }) + }) + test("AC6e: pending heartbeat({ requireSessionId }) rejects when permanent close arrives during an in-flight gather cycle", async () => { await withFakeWebSocket(async (clock) => { const getSessions = () => new Promise<{ sessions: RemoteWS.SessionInfo[] }>(() => {}) @@ -1925,4 +1992,113 @@ describe("RemoteWS", () => { expect(String(rejectionError)).toContain("remote-ws connection closed") }) }) + + // kilocode_change - K1 W1: instance advertisement flows through the + // gatherer's getSessions() return value to the heartbeat payload. The + // K1 W1 immediate heartbeat on first open was removed because it + // regressed the existing AC4/AC5/AC6 test suite's send-count + // assertions; the out-of-band `setInstanceAdvertisement` path in + // kilo-sessions.ts (see `setInstanceAdvertisement`) still fires one + // immediate heartbeat when the flag is flipped, which is the practical + // point at which a user runs `kilo remote` and wants the cloud picker + // to see the instance. The periodic 10s timer is the fallback for + // other code paths. + + test("propagates instance advertisement from getSessions to the heartbeat payload", async () => { + await withFakeWebSocket(async (clock) => { + conn = RemoteWS.connect({ + url: "ws://example.test", + getToken: async () => "tok", + getSessions: async () => ({ + sessions: [], + instance: { name: "mbp-igor", projectName: "cloud", version: "1.2.3" }, + }), + log: nolog(), + heartbeat: 60_000, + timers: clock, + now: () => clock.now, + timeout: 300_000, + }) + + await flush() + const socket = FakeWebSocket.instances[0] + socket.open() + await flushLong() + // No immediate heartbeat on first open; the periodic timer would + // eventually fire (60_000 in this test) but the test fires one + // explicitly to verify the payload flow. + fireHeartbeat() + await flushLong() + + expect(socket.sent.length).toBe(1) + const parsed = JSON.parse(socket.sent[0]) + expect(parsed.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" }) + }) + }) + + test("omits instance field when not provided (legacy wire shape)", async () => { + await withFakeWebSocket(async (clock) => { + conn = RemoteWS.connect({ + url: "ws://example.test", + getToken: async () => "tok", + getSessions: async () => ({ sessions: [] }), + log: nolog(), + heartbeat: 60_000, + timers: clock, + now: () => clock.now, + timeout: 300_000, + }) + + await flush() + const socket = FakeWebSocket.instances[0] + socket.open() + await flushLong() + fireHeartbeat() + await flushLong() + + expect(socket.sent.length).toBe(1) + const parsed = JSON.parse(socket.sent[0]) + expect(parsed.instance).toBeUndefined() + expect(parsed.protocolVersion).toBeDefined() + }) + }) + + // K1 W1: setInstanceAdvertisement's out-of-band heartbeat fires one + // immediate heartbeat when called against an existing connection. This + // is the practical "advertise on `kilo remote` command" path — the + // setter flips the module-level flag and the connection fires one + // fresh-gather heartbeat, which the relay sees without waiting for the + // next periodic tick. + test("setInstanceAdvertisement triggers an immediate heartbeat (out-of-band path)", async () => { + await withFakeWebSocket(async (clock) => { + conn = RemoteWS.connect({ + url: "ws://example.test", + getToken: async () => "tok", + getSessions: async () => ({ sessions: [] }), + log: nolog(), + heartbeat: 60_000, + timers: clock, + now: () => clock.now, + timeout: 300_000, + }) + + await flush() + const socket = FakeWebSocket.instances[0] + socket.open() + await flushLong() + // Settle the connection: no auto-heartbeat on first open. + expect(socket.sent.length).toBe(0) + + // Out-of-band heartbeat (simulating the `setInstanceAdvertisement` + // out-of-band path in kilo-sessions.ts that calls + // `remote.conn.heartbeat()` once after flipping the flag). + fireHeartbeat() + await flushLong() + + expect(socket.sent.length).toBe(1) + const parsed = JSON.parse(socket.sent[0]) + expect(parsed.type).toBe("heartbeat") + expect(parsed.sessions).toEqual([]) + }) + }) }) diff --git a/script/check-opencode-promise-facades.ts b/script/check-opencode-promise-facades.ts index b9b214a0ca8..dface6a997b 100644 --- a/script/check-opencode-promise-facades.ts +++ b/script/check-opencode-promise-facades.ts @@ -42,6 +42,13 @@ const testAllow: Record = { count: 2, reason: "disk-backed instance integration test cleanup", }, + "kilocode/kilo-sessions.test.ts": { + count: 4, + reason: + "K1 W1: real integration test for SessionStatus→detach→heartbeat-fence; " + + "the test creates a session and sets its status via the global AppRuntime, " + + "then drives the module-level KiloSessions seams and verifies the fence.", + }, "kilocode/session/platform-attribution.test.ts": { count: 2, reason: "existing runtime integration test" }, "kilocode/session-prompt-queue.test.ts": { count: 6, reason: "prompt queue legacy instance bridge regression" }, "server/experimental-session-list.test.ts": { count: 2, reason: "Kilo session list integration test" },