From 7c845aa0ec8a0ff46f25cff26545260c0126c3f7 Mon Sep 17 00:00:00 2001 From: Josh Lambert Date: Mon, 31 Aug 2026 11:03:07 -0400 Subject: [PATCH 1/5] fix(cli): drain background work before headless exit --- .changeset/headless-session-drain.md | 6 + packages/opencode/src/cli/cmd/run.ts | 200 +++++++----- packages/opencode/src/effect/app-runtime.ts | 10 +- packages/opencode/src/kilocode/cli/cmd/run.ts | 21 +- .../opencode/src/kilocode/cli/run-auto.ts | 2 +- .../opencode/src/kilocode/cli/run-drain.ts | 109 +++++++ .../server/httpapi/groups/kilocode.ts | 18 ++ .../server/httpapi/handlers/kilocode.ts | 15 + .../opencode/src/kilocode/session/drain.ts | 119 +++++++ .../server/routes/instance/httpapi/server.ts | 2 + packages/opencode/src/session/prompt.ts | 38 ++- packages/opencode/src/session/run-state.ts | 30 +- packages/opencode/src/tool/registry.ts | 2 + packages/opencode/src/tool/task.ts | 116 ++++--- .../test/kilocode/cli-run-auto-helper.test.ts | 9 +- .../kilocode/headless-session-drain.test.ts | 297 ++++++++++++++++++ .../opencode/test/kilocode/run-auto.test.ts | 113 ++++++- .../opencode/test/kilocode/run-drain.test.ts | 112 +++++++ .../test/kilocode/session-drain.test.ts | 123 ++++++++ .../test/kilocode/task-nesting.test.ts | 4 + .../test/kilocode/tool-task-model.test.ts | 4 + packages/opencode/test/tool/task.test.ts | 55 ++++ packages/schema/src/event-manifest.ts | 2 + packages/schema/src/kilocode/session-drain.ts | 12 + packages/sdk/js/src/v2/gen/sdk.gen.ts | 45 +++ packages/sdk/js/src/v2/gen/types.gen.ts | 84 +++++ packages/sdk/openapi.json | 249 +++++++++++++++ 27 files changed, 1649 insertions(+), 148 deletions(-) create mode 100644 .changeset/headless-session-drain.md create mode 100644 packages/opencode/src/kilocode/cli/run-drain.ts create mode 100644 packages/opencode/src/kilocode/session/drain.ts create mode 100644 packages/opencode/test/kilocode/headless-session-drain.test.ts create mode 100644 packages/opencode/test/kilocode/run-drain.test.ts create mode 100644 packages/opencode/test/kilocode/session-drain.test.ts create mode 100644 packages/schema/src/kilocode/session-drain.ts diff --git a/.changeset/headless-session-drain.md b/.changeset/headless-session-drain.md new file mode 100644 index 000000000000..52a52d606c00 --- /dev/null +++ b/.changeset/headless-session-drain.md @@ -0,0 +1,6 @@ +--- +"@kilocode/cli": patch +"@kilocode/sdk": patch +--- + +Wait for background subagents and their parent continuations before completing headless runs. diff --git a/packages/opencode/src/cli/cmd/run.ts b/packages/opencode/src/cli/cmd/run.ts index d4cab544c531..26fc55086377 100644 --- a/packages/opencode/src/cli/cmd/run.ts +++ b/packages/opencode/src/cli/cmd/run.ts @@ -279,12 +279,12 @@ export const RunCommand = effectCmd({ const { InstanceRef } = yield* Effect.promise(() => import("@/effect/instance-ref")) const { ServerAuth } = yield* Effect.promise(() => import("@/server/auth")) // kilocode_change start - lazy Kilo implementations (see top-of-file note) - const { createKiloClient } = yield* Effect.promise(() => import("@kilocode/sdk/v2")) const { buildRunMessage } = yield* Effect.promise(() => import("@/kilocode/cli/cmd/run-message")) const { importCloudSession, validateCloudFork, reportCloudImportError } = yield* Effect.promise( () => import("@/kilocode/cloud-session"), ) const { KiloRunAuto } = yield* Effect.promise(() => import("@/kilocode/cli/run-auto")) + const { KiloRunDrain } = yield* Effect.promise(() => import("@/kilocode/cli/run-drain")) const { KiloHeadless } = yield* Effect.promise(() => import("@/kilocode/permission/headless")) const { KiloRun, KiloRunDaemon } = yield* Effect.promise(() => import("@/kilocode/cli/cmd/run")) // kilocode_change end @@ -368,7 +368,8 @@ export const RunCommand = effectCmd({ ? ServerAuth.headers({ password: args.password, username: args.username }) : undefined const attachSDK = (dir?: string) => { - return createKiloClient({ + return KiloRunDrain.client({ + // kilocode_change baseUrl: args.attach!, directory: dir, headers: attachHeaders, @@ -769,6 +770,7 @@ export const RunCommand = effectCmd({ const sessionID = sess.id // kilocode_change start - track Task children; plain headless runs deny subagent asks instead of hanging (#11903) const tracked = KiloRunAuto.create(sessionID) // kilocode_change - named to avoid shadowing the `auto` flag + const drain = KiloRunDrain.create(sessionID) if (!args.attach && !args.auto && !skipPermissions) KiloHeadless.mark(sessionID) // kilocode_change - --yolo skips too // kilocode_change end @@ -800,6 +802,7 @@ export const RunCommand = effectCmd({ // kilocode_change start - revert to upstream: consume native events without normalizing sync copies for await (const event of events.stream) { + if (drain.event(event)) break // kilocode_change // kilocode_change end if ( @@ -903,19 +906,11 @@ export const RunCommand = effectCmd({ } // kilocode_change end - if ( - event.type === "session.status" && - event.properties.sessionID === sessionID && - event.properties.status.type === "idle" - ) { - break - } - // kilocode_change start - non-interactive runs dismiss suggestions so they don't block if (event.type === "suggestion.shown") { const suggestion = event.properties if (suggestion.sessionID === sessionID || KiloRunAuto.allowed(tracked, suggestion.sessionID)) { - await client.suggestion.dismiss({ requestID: suggestion.id }).catch(() => {}) + await client.suggestion.dismiss({ requestID: suggestion.id }, { signal: drain.signal }).catch(() => {}) } continue } @@ -923,20 +918,24 @@ export const RunCommand = effectCmd({ if (event.type === "permission.asked") { const permission = event.properties + if (!KiloRunAuto.allowed(tracked, permission.sessionID)) continue // kilocode_change // kilocode_change start - skill shell batches need an interactive human decision. The server ignores // non-interactive approvals, so headless runs must reject explicitly rather than leave them pending. if (permission.metadata?.["skillShell"] === true || permission.metadata?.["sandboxEscalation"] === true) { - await client.permission.reply({ requestID: permission.id, reply: "reject" }) + await client.permission.reply({ requestID: permission.id, reply: "reject" }, { signal: drain.signal }) continue } // kilocode_change end // kilocode_change start - approve root and tracked Task child permissions in auto mode if (args.auto) { if (!KiloRunAuto.allowed(tracked, permission.sessionID)) continue - await client.permission.reply({ - requestID: permission.id, - reply: "once", - }) + await client.permission.reply( + { + requestID: permission.id, + reply: "once", + }, + { signal: drain.signal }, + ) continue } // kilocode_change end @@ -947,10 +946,13 @@ export const RunCommand = effectCmd({ if (permission.sessionID !== sessionID) { if (!KiloRunAuto.allowed(tracked, permission.sessionID)) continue if (skipPermissions) { - await client.permission.reply({ - requestID: permission.id, - reply: "once", - }) + await client.permission.reply( + { + requestID: permission.id, + reply: "once", + }, + { signal: drain.signal }, + ) continue } UI.println( @@ -959,10 +961,13 @@ export const RunCommand = effectCmd({ `subagent permission requested: ${permission.permission} (${permission.patterns.join(", ")}); auto-rejecting`, ) autoRejected = true // kilocode_change - await client.permission.reply({ - requestID: permission.id, - reply: "reject", - }) + await client.permission.reply( + { + requestID: permission.id, + reply: "reject", + }, + { signal: drain.signal }, + ) continue } // kilocode_change end @@ -970,10 +975,13 @@ export const RunCommand = effectCmd({ if (permission.sessionID !== sessionID) continue if (skipPermissions) { - await client.permission.reply({ - requestID: permission.id, - reply: "once", - }) + await client.permission.reply( + { + requestID: permission.id, + reply: "once", + }, + { signal: drain.signal }, + ) } else { UI.println( UI.Style.TEXT_WARNING_BOLD + "!", @@ -981,29 +989,32 @@ export const RunCommand = effectCmd({ `permission requested: ${permission.permission} (${permission.patterns.join(", ")}); auto-rejecting`, ) autoRejected = true // kilocode_change - await client.permission.reply({ - requestID: permission.id, - reply: "reject", - }) + await client.permission.reply( + { + requestID: permission.id, + reply: "reject", + }, + { signal: drain.signal }, + ) } } // kilocode_change start - bounded network retry handling if (event.type === "session.network.asked") { const request = event.properties - if (request.sessionID !== sessionID) continue + if (!KiloRunAuto.allowed(tracked, request.sessionID)) continue retries++ if (retries > MAX_RETRIES) { UI.println( UI.Style.TEXT_WARNING_BOLD + "!", UI.Style.TEXT_NORMAL + `network retry limit reached (${MAX_RETRIES}); rejecting`, ) - await client.network.reject({ requestID: request.id }) + await client.network.reject({ requestID: request.id }, { signal: drain.signal }) continue } const delay = Math.min(5000 * Math.pow(2, retries - 1), 60000) - await new Promise((resolve) => setTimeout(resolve, delay)) - await client.network.reply({ requestID: request.id }) + await drain.pause(delay) + await client.network.reply({ requestID: request.id }, { signal: drain.signal }) } // kilocode_change end } @@ -1032,63 +1043,103 @@ export const RunCommand = effectCmd({ await share(client, sessionID) + // kilocode_change start if (!interactive) { - const events = await client.event.subscribe() - const completed = loop(client, events).catch((e) => { - console.error(e) - process.exitCode = 1 + const events = await client.event.subscribe(undefined, { + signal: drain.signal, + sseMaxRetryAttempts: 1, + onSseError: (error) => drain.end(error), }) + const completed = loop(client, events).then( + (error) => { + drain.end() + return error + }, + (error) => { + drain.end(error) + return undefined + }, + ) async function finish() { - if (args.attach) return + await drain.wait(client, cwd) const error = await completed if (error) process.exitCode = 1 } - // kilocode_change start - handle built-in session commands - if (builtin) { - const result = await KiloRun.runBuiltin(client, sessionID, builtin, args.model, sess.model, cwd) - if (result.error) { - if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error)) - process.exitCode = 1 + try { + await drain.race(KiloRunDrain.check(client, drain.signal)) + await drain.ready() + if (builtin) { + const result = await drain.race( + KiloRun.runBuiltin(client, sessionID, builtin, args.model, sess.model, cwd, drain.signal), + ) + if (result.error) { + if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error)) + process.exitCode = 1 + await finish() + return + } + await finish() + return } - return - } - // kilocode_change end - if (args.command) { - const result = await client.session.command({ - sessionID, - agent, - model: args.model, - command: args.command, - arguments: message, - variant: args.variant, - }) + if (args.command) { + const result = await drain.race( + client.session.command( + { + sessionID, + agent, + model: args.model, + command: args.command, + arguments: message, + variant: args.variant, + }, + { signal: drain.signal }, + ), + ) + if (result.error) { + if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error)) + process.exitCode = 1 + await finish() + return + } + await finish() + return + } + + const model = pick(args.model) + const result = await drain.race( + client.session.prompt( + { + sessionID, + agent, + model, + variant: args.variant, + parts: [...files, { type: "text", text: message }], + }, + { signal: drain.signal }, + ), + ) if (result.error) { if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error)) process.exitCode = 1 + await finish() return } await finish() return - } - - const model = pick(args.model) - const result = await client.session.prompt({ - sessionID, - agent, - model, - variant: args.variant, - parts: [...files, { type: "text", text: message }], - }) - if (result.error) { - if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error)) + } catch (error) { + const text = error instanceof Error ? error.message : String(error) + if (!emit("error", { error: text })) UI.error(text) process.exitCode = 1 - return + } finally { + drain.close() + await completed + await KiloRunDrain.flush() } - await finish() return } + // kilocode_change end const model = pick(args.model) const { runInteractiveMode } = await import("./run/runtime") @@ -1169,7 +1220,8 @@ export const RunCommand = effectCmd({ if (auth) headers.set("Authorization", auth) return Server.Default().app.fetch(new Request(request, { headers })) }) as typeof globalThis.fetch - const sdk = createKiloClient({ + const sdk = KiloRunDrain.client({ + // kilocode_change baseUrl: "http://kilo.internal", fetch: fetchFn, directory, diff --git a/packages/opencode/src/effect/app-runtime.ts b/packages/opencode/src/effect/app-runtime.ts index eb61a88fea01..0a2acfcaf336 100644 --- a/packages/opencode/src/effect/app-runtime.ts +++ b/packages/opencode/src/effect/app-runtime.ts @@ -55,6 +55,7 @@ import { BackgroundJob } from "@/background/job" import { RuntimeFlags } from "@/effect/runtime-flags" // kilocode_change start import { Notebook } from "@/kilocode/notebook/service" +import { SessionDrain } from "@/kilocode/session/drain" import { AgentManager } from "@/kilocode/agent-manager/service" // kilocode_change end import { EventV2Bridge } from "@/event-v2-bridge" @@ -70,7 +71,14 @@ import { Pty } from "@opencode-ai/core/pty" // kilocode_change // kilocode_change start - retain Kilo runtime services in the upstream node graph const memory = LayerNode.make({ service: MemoryService.Service, layer: MemoryService.layer, deps: [] }) -const kilo = LayerNode.group([Credential.node, ModelCache.node, AgentManager.node, Notebook.node, memory]) +const kilo = LayerNode.group([ + Credential.node, + ModelCache.node, + AgentManager.node, + Notebook.node, + SessionDrain.node, + memory, +]) // kilocode_change end export const AppLayer = AppNodeBuilderV1.build( diff --git a/packages/opencode/src/kilocode/cli/cmd/run.ts b/packages/opencode/src/kilocode/cli/cmd/run.ts index 2580db2b5dbb..363bbf900f2a 100644 --- a/packages/opencode/src/kilocode/cli/cmd/run.ts +++ b/packages/opencode/src/kilocode/cli/cmd/run.ts @@ -1,4 +1,5 @@ -import { createKiloClient, type KiloClient } from "@kilocode/sdk/v2" +import type { KiloClient } from "@kilocode/sdk/v2" +import { KiloRunDrain } from "../run-drain" import { UI } from "@/cli/ui" import { DaemonClient } from "@/kilocode/daemon/client" import { isBuiltinCommand, type BuiltinCommand } from "@/kilocode/session/builtin-commands" @@ -28,6 +29,7 @@ export namespace KiloRun { model?: string, current?: { id: string; providerID: string }, directory?: string, + signal?: AbortSignal, ) { const selected = resolve(model, current) if (!selected) { @@ -38,12 +40,15 @@ export namespace KiloRun { switch (command) { case "compact": case "summarize": - return sdk.session.summarize({ - sessionID, - directory, - providerID: selected.providerID, - modelID: selected.modelID, - }) + return sdk.session.summarize( + { + sessionID, + directory, + providerID: selected.providerID, + modelID: selected.modelID, + }, + { signal }, + ) } } } @@ -58,7 +63,7 @@ export namespace KiloRunDaemon { const daemon = await DaemonClient.maybe() if (!daemon) return false const dir = input.directory ?? Filesystem.resolve(process.cwd()) - const client = createKiloClient({ baseUrl: daemon.url, directory: dir, headers: daemon.headers }) + const client = KiloRunDrain.client({ baseUrl: daemon.url, directory: dir, headers: daemon.headers }) await input.execute(client) return true } diff --git a/packages/opencode/src/kilocode/cli/run-auto.ts b/packages/opencode/src/kilocode/cli/run-auto.ts index 84014997ef35..7d2ff8d478bd 100644 --- a/packages/opencode/src/kilocode/cli/run-auto.ts +++ b/packages/opencode/src/kilocode/cli/run-auto.ts @@ -26,7 +26,7 @@ export namespace KiloRunAuto { export function track(state: State, part: Part) { if (part.type !== "tool") return if (part.tool !== "task") return - if (part.sessionID !== state.root) return + if (!part.sessionID || !state.sessions.has(part.sessionID)) return const id = child(meta(part.state)) if (!id) return state.sessions.add(id) diff --git a/packages/opencode/src/kilocode/cli/run-drain.ts b/packages/opencode/src/kilocode/cli/run-drain.ts new file mode 100644 index 000000000000..1ceb9b843866 --- /dev/null +++ b/packages/opencode/src/kilocode/cli/run-drain.ts @@ -0,0 +1,109 @@ +import { createKiloClient, type Event, type KiloClient } from "@kilocode/sdk/v2" +import z from "zod" +import { setTimeout } from "node:timers/promises" + +const connections = new WeakMap[0]>>() +const capability = z.object({ + paths: z.object({ + "/kilocode/session/{sessionID}/drain": z.object({ + post: z.object({ operationId: z.literal("kilocode.drainSession") }), + }), + }), +}) + +export namespace KiloRunDrain { + export function client(config: NonNullable[0]>) { + const sdk = createKiloClient(config) + connections.set(sdk, config) + return sdk + } + + export async function check(sdk: KiloClient, signal: AbortSignal) { + const config = connections.get(sdk) + if (!config?.baseUrl) throw new Error("Missing headless server transport") + const headers = new Headers() + if (config.headers instanceof Headers) config.headers.forEach((value, name) => headers.set(name, value)) + else if (Array.isArray(config.headers)) for (const [name, value] of config.headers) headers.set(name, value) + else + for (const [name, value] of Object.entries(config.headers ?? {})) { + if (typeof value === "string") headers.set(name, value) + else if (value != null) throw new Error("Unsupported server header value") + } + const request = new Request(`${config.baseUrl.replace(/\/$/, "")}/doc`, { headers, signal }) + const response = await (config.fetch ?? globalThis.fetch)(request) + if (!response.ok || !response.headers.get("content-type")?.toLowerCase().includes("application/json")) { + throw new Error("Server does not support session draining; upgrade or restart the server") + } + if (!capability.safeParse(await response.json()).success) { + throw new Error("Server does not support session draining; upgrade or restart the server") + } + } + + export function create(sessionID: string) { + const token = crypto.randomUUID() + const abort = new AbortController() + const connected = Promise.withResolvers() + const acknowledged = Promise.withResolvers() + const failure = Promise.withResolvers() + let closing = false + let drained = false + const race = (work: Promise) => + Promise.race([ + work, + failure.promise.then((error): never => { + throw error + }), + ]) + return { + token, + signal: abort.signal, + race, + ready: () => race(connected.promise), + pause: (ms: number) => setTimeout(ms, undefined, { signal: abort.signal }), + event(event: Event) { + if (event.type === "server.connected") connected.resolve() + if ( + event.type === "session.turn.close" && + event.properties.sessionID === sessionID && + event.properties.reason === "interrupted" + ) { + failure.resolve(new Error("Session interrupted before completion")) + } + if ( + event.type !== "session.drained" || + event.properties.sessionID !== sessionID || + event.properties.token !== token + ) + return false + drained = true + acknowledged.resolve() + return true + }, + end(error?: unknown) { + if (closing || (drained && !error)) return + failure.resolve(error instanceof Error ? error : new Error("Session event stream ended before completion")) + }, + async wait(sdk: KiloClient, directory?: string) { + const result = await race(sdk.kilocode.drainSession({ sessionID, directory, token }, { signal: abort.signal })) + if (result.error || !z.literal(true).safeParse(result.data).success) + throw new Error("Server did not acknowledge session completion") + await race(acknowledged.promise) + }, + close() { + closing = true + abort.abort() + }, + } + } + + export async function flush() { + await Promise.all( + [process.stdout, process.stderr].map( + (stream) => + new Promise((resolve, reject) => { + stream.write("", (error) => (error ? reject(error) : resolve())) + }), + ), + ) + } +} diff --git a/packages/opencode/src/kilocode/server/httpapi/groups/kilocode.ts b/packages/opencode/src/kilocode/server/httpapi/groups/kilocode.ts index 6c2f00e1d113..1ee461f6d8c2 100644 --- a/packages/opencode/src/kilocode/server/httpapi/groups/kilocode.ts +++ b/packages/opencode/src/kilocode/server/httpapi/groups/kilocode.ts @@ -26,6 +26,7 @@ import { ModelUsage } from "@/kilocode/session/model-usage" import { MessageID, SessionID } from "@/session/schema" import { ApiNotFoundError, InvalidRequestError } from "@/server/routes/instance/httpapi/errors" import { CommandFiles } from "@/kilocode/command-files" +import { Token } from "@opencode-ai/schema/kilocode/session-drain" const root = "/kilocode" const Scope = Schema.Literals(["global", "project"]) @@ -68,6 +69,8 @@ export const ResumeSessionPayload = Schema.Struct({ snapshotInitialization: Schema.optional(Schema.Literal("wait")), }) +export const DrainSessionPayload = Schema.Struct({ token: Token }) + export const NotebookReplyPayload = Schema.Struct({ result: NotebookResult }) export const NotebookRejectPayload = Schema.Struct({ error: NotebookFailure }) export const AgentManagerReplyPayload = Schema.Struct({ result: AgentManagerResult }) @@ -90,6 +93,7 @@ export const KilocodePaths = { agentManagerReject: `${root}/agent-manager/:requestID/reject`, sessionModelUsage: `/session/:sessionID/model-usage`, resumeSession: `${root}/session/:sessionID/resume`, + drainSession: `${root}/session/:sessionID/drain`, backgroundJobs: `${root}/background-jobs`, backgroundJobCancel: `${root}/background-jobs/:jobID/cancel`, backgroundJobPromote: `${root}/background-jobs/:jobID/promote`, @@ -113,6 +117,20 @@ export const KilocodeApi = HttpApi.make("kilocode") "Resume the specified unfinished assistant turn without adding a user message. Active, completed, reverted, and blocked sessions cannot be resumed.", }), ), + HttpApiEndpoint.post("drainSession", KilocodePaths.drainSession, { + params: { sessionID: SessionID }, + query: WorkspaceRoutingQuery, + payload: DrainSessionPayload, + success: described(Schema.Boolean, "Session work drained"), + error: ApiNotFoundError, + }).annotateMerge( + OpenApi.annotations({ + identifier: "kilocode.drainSession", + summary: "Wait for session completion", + description: + "Wait for active session work and background result delivery, then publish the matching drain acknowledgment.", + }), + ), HttpApiEndpoint.post("heapSnapshot", KilocodePaths.heapSnapshot, { query: WorkspaceRoutingQuery, success: described(Schema.String, "Heap snapshot file path"), diff --git a/packages/opencode/src/kilocode/server/httpapi/handlers/kilocode.ts b/packages/opencode/src/kilocode/server/httpapi/handlers/kilocode.ts index 704f5b0f0563..bd7517aaaa4a 100644 --- a/packages/opencode/src/kilocode/server/httpapi/handlers/kilocode.ts +++ b/packages/opencode/src/kilocode/server/httpapi/handlers/kilocode.ts @@ -34,6 +34,8 @@ import { InvalidRequestError } from "@/server/routes/instance/httpapi/errors" import { Skill } from "@/skill" import { BackgroundJob } from "@/background/job" import { SessionRunState } from "@/session/run-state" +import { SessionDrain } from "@/kilocode/session/drain" +import { Drained } from "@opencode-ai/schema/kilocode/session-drain" import { SessionID } from "@/session/schema" import { RuntimeFlags } from "@/effect/runtime-flags" import { KiloSnapshotCleanup } from "@/kilocode/snapshot/cleanup" @@ -51,6 +53,7 @@ import { RemoveSkillPayload, RemoveSnapshotPayload, ResumeSessionPayload, + DrainSessionPayload, BackgroundJobInfo, BackgroundJobsQuery, } from "../groups/kilocode" @@ -66,6 +69,7 @@ export const kilocodeHandlers = HttpApiBuilder.group(InstanceHttpApi, "kilocode" const notebook = yield* Notebook.Service const background = yield* BackgroundJob.Service const runState = yield* SessionRunState.Service + const drain = yield* SessionDrain.Service const flags = yield* RuntimeFlags.Service const locations = yield* LocationServiceMap.Service const fs = yield* FSUtil.Service @@ -78,6 +82,16 @@ export const kilocodeHandlers = HttpApiBuilder.group(InstanceHttpApi, "kilocode" const events = yield* EventV2Bridge.Service const scope = yield* Scope.Scope + const drainSession = Effect.fn("KilocodeHttpApi.drainSession")(function* (ctx: { + params: { sessionID: SessionID } + payload: typeof DrainSessionPayload.Type + }) { + yield* mapStorageNotFound(sessions.get(ctx.params.sessionID)) + yield* drain.wait(ctx.params.sessionID) + yield* events.publish(Drained, { sessionID: ctx.params.sessionID, token: ctx.payload.token }) + return true + }) + const resumeSession = Effect.fn("KilocodeHttpApi.resumeSession")(function* (ctx: { params: { sessionID: SessionID } payload: typeof ResumeSessionPayload.Type @@ -334,6 +348,7 @@ export const kilocodeHandlers = HttpApiBuilder.group(InstanceHttpApi, "kilocode" return handlers .handle("resumeSession", resumeSession) + .handle("drainSession", drainSession) .handle("heapSnapshot", heapSnapshot) .handle("commandFiles", commandFiles) .handle("removeCommand", removeCommand) diff --git a/packages/opencode/src/kilocode/session/drain.ts b/packages/opencode/src/kilocode/session/drain.ts new file mode 100644 index 000000000000..48cb5c729fc6 --- /dev/null +++ b/packages/opencode/src/kilocode/session/drain.ts @@ -0,0 +1,119 @@ +import { Context, Deferred, Effect, Layer } from "effect" +import { LayerNode } from "@opencode-ai/core/effect/layer-node" +import { InstanceState } from "@/effect/instance-state" +import type { SessionID } from "@/session/schema" + +type Entry = { + count: number + parent?: Entry + waiters: Set> +} + +export type Delivery = { release: () => void; retained: boolean } + +export interface Interface { + readonly hold: (id: SessionID) => Effect.Effect<() => void> + readonly link: (child: SessionID, parent: SessionID) => Effect.Effect + readonly wait: (id: SessionID) => Effect.Effect + readonly track: (id: SessionID, work: Effect.Effect) => Effect.Effect +} + +export class Service extends Context.Service()("@kilo/SessionDrain") {} + +export const layer = Layer.effect( + Service, + Effect.gen(function* () { + const state = yield* InstanceState.make(() => + Effect.gen(function* () { + const entries = new Map() + const data = { entries, closed: false } + yield* Effect.addFinalizer(() => + Effect.sync(() => { + data.closed = true + for (const entry of entries.values()) { + for (const waiter of entry.waiters) Deferred.doneUnsafe(waiter, Effect.interrupt) + entry.waiters.clear() + } + entries.clear() + }), + ) + return data + }), + ) + + function entry(entries: Map, id: SessionID) { + const found = entries.get(id) + if (found) return found + const value: Entry = { count: 0, waiters: new Set() } + entries.set(id, value) + return value + } + + function update(value: Entry, delta: number) { + for (let current: Entry | undefined = value; current; current = current.parent) { + current.count += delta + if (current.count !== 0) continue + for (const waiter of current.waiters) Deferred.doneUnsafe(waiter, Effect.void) + current.waiters.clear() + } + } + + const hold = Effect.fn("SessionDrain.hold")(function* (id: SessionID) { + const data = yield* InstanceState.get(state) + if (data.closed) return yield* Effect.interrupt + const value = entry(data.entries, id) + update(value, 1) + let active = true + return () => { + if (!active) return + active = false + if (!data.closed) update(value, -1) + } + }) + + const link = Effect.fn("SessionDrain.link")(function* (child: SessionID, parent: SessionID) { + const data = yield* InstanceState.get(state) + if (data.closed) return yield* Effect.interrupt + const value = entry(data.entries, child) + const ancestor = entry(data.entries, parent) + if (value.parent === ancestor) return yield* Effect.void + if (value.parent) return yield* Effect.die(new Error("Session drain parent changed")) + for (let current: Entry | undefined = ancestor; current; current = current.parent) { + if (current === value) return yield* Effect.die(new Error("Cyclic session drain ancestry")) + } + value.parent = ancestor + update(ancestor, value.count) + return yield* Effect.void + }) + + const wait = Effect.fn("SessionDrain.wait")(function* (id: SessionID) { + const data = yield* InstanceState.get(state) + const value = entry(data.entries, id) + while (!data.closed) { + if (value.count === 0) return yield* Effect.void + const waiter = Deferred.makeUnsafe() + value.waiters.add(waiter) + yield* Deferred.await(waiter).pipe( + Effect.ensuring( + Effect.sync(() => { + value.waiters.delete(waiter) + }), + ), + ) + } + return yield* Effect.interrupt + }) + + const track = (id: SessionID, work: Effect.Effect) => + Effect.acquireUseRelease( + hold(id), + () => work, + (release) => Effect.sync(release), + ) + + return Service.of({ hold, link, wait, track }) + }), +) + +export const node = LayerNode.make({ service: Service, layer, deps: [] }) +export * as SessionDrain from "./drain" diff --git a/packages/opencode/src/server/routes/instance/httpapi/server.ts b/packages/opencode/src/server/routes/instance/httpapi/server.ts index dc35b2ddbf79..0a84ed84e500 100644 --- a/packages/opencode/src/server/routes/instance/httpapi/server.ts +++ b/packages/opencode/src/server/routes/instance/httpapi/server.ts @@ -42,6 +42,7 @@ import { SessionProcessor } from "@/session/processor" import { SessionPrompt } from "@/session/prompt" import { SessionRevert } from "@/session/revert" import { SessionRunState } from "@/session/run-state" +import { SessionDrain } from "@/kilocode/session/drain" // kilocode_change import { Session } from "@/session/session" import { SessionStatus } from "@/session/status" import { SessionSummary } from "@/session/summary" @@ -264,6 +265,7 @@ const app = LayerNode.group([ RuntimeFlags.node, EventV2Bridge.node, SessionRunState.node, + SessionDrain.node, // kilocode_change SessionProcessor.node, SessionCompaction.node, SessionRevert.node, diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 47de0add498a..5a5d51b9b83c 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -76,6 +76,7 @@ import { EffectBridge } from "@/effect/bridge" import { TaskTool, type TaskPromptOps } from "@/tool/task" import { assertExternalDirectoryEffect } from "@/tool/external-directory" // kilocode_change import { SessionRunState } from "./run-state" +import { SessionDrain } from "@/kilocode/session/drain" // kilocode_change import { RuntimeFlags } from "@/effect/runtime-flags" import { EventV2Bridge } from "@/event-v2-bridge" import { Database } from "@opencode-ai/core/database/database" @@ -139,14 +140,10 @@ function isOrphanedInterruptedTool(part: SessionV1.ToolPart) { export interface Interface { readonly cancel: (sessionID: SessionID) => Effect.Effect - readonly prompt: ( - input: PromptInput, - ) => Effect.Effect + readonly prompt: (input: PromptInput) => Effect.Effect readonly loop: (input: LoopInput) => Effect.Effect readonly shell: (input: ShellInput) => Effect.Effect - readonly command: ( - input: CommandInput, - ) => Effect.Effect + readonly command: (input: CommandInput) => Effect.Effect readonly resolvePromptParts: (template: string) => Effect.Effect } @@ -176,6 +173,7 @@ export const layer = Layer.effect( const scope = yield* Scope.Scope const instruction = yield* Instruction.Service const state = yield* SessionRunState.Service + const drain = yield* SessionDrain.Service // kilocode_change const revert = yield* SessionRevert.Service const summary = yield* SessionSummary.Service const sys = yield* SystemPrompt.Service @@ -1432,16 +1430,19 @@ export const layer = Layer.effect( // Queue tails and runner fibers can resume outside the HTTP request's // ambient instance context; bridge both Effect refs and legacy ALS. const bridge = yield* EffectBridge.make() - return yield* KiloSessionPromptQueue.enqueue( + return yield* drain.track( input.sessionID, - message.info.id, - bridge.run( - loop({ sessionID: input.sessionID, snapshotInitialization: input.snapshotInitialization }).pipe( - Effect.orDie, + KiloSessionPromptQueue.enqueue( + input.sessionID, + message.info.id, + bridge.run( + loop({ sessionID: input.sessionID, snapshotInitialization: input.snapshotInitialization }).pipe( + Effect.orDie, + ), ), - ), // kilocode_change - bridge.run(lastAssistant(input.sessionID)), - dismiss, + bridge.run(lastAssistant(input.sessionID)), + dismiss, + ), ) // kilocode_change end }, @@ -1620,7 +1621,13 @@ export const layer = Layer.effect( } compactionAttempts++ // kilocode_change end - yield* compaction.create({ sessionID, agent: lastUser.agent, model: lastUser.model, auto: true, overflow: false }) // kilocode_change + yield* compaction.create({ + sessionID, + agent: lastUser.agent, + model: lastUser.model, + auto: true, + overflow: false, + }) // kilocode_change continue } @@ -2660,6 +2667,7 @@ export const node = LayerNode.make({ CrossSpawnSpawner.node, Instruction.node, SessionRunState.node, + SessionDrain.node, // kilocode_change SessionRevert.node, SessionSummary.node, SystemPrompt.node, diff --git a/packages/opencode/src/session/run-state.ts b/packages/opencode/src/session/run-state.ts index e2666882466a..880ed3dc911f 100644 --- a/packages/opencode/src/session/run-state.ts +++ b/packages/opencode/src/session/run-state.ts @@ -8,6 +8,7 @@ import { Effect, Latch, Layer, Scope, Context } from "effect" import { Session } from "./session" import { SessionID } from "./schema" import { SessionStatus } from "./status" +import { SessionDrain } from "@/kilocode/session/drain" // kilocode_change export interface Interface { readonly assertNotBusy: (sessionID: SessionID) => Effect.Effect @@ -32,6 +33,7 @@ export const layer = Layer.effect( Effect.gen(function* () { const background = yield* BackgroundJob.Service const status = yield* SessionStatus.Service + const drain = yield* SessionDrain.Service // kilocode_change const state = yield* InstanceState.make( Effect.fn("SessionRunState.state")(function* () { @@ -86,12 +88,18 @@ export const layer = Layer.effect( yield* existing.cancel }) + // kilocode_change start const ensureRunning = Effect.fn("SessionRunState.ensureRunning")(function* ( sessionID: SessionID, onInterrupt: Effect.Effect, work: Effect.Effect, ) { - return yield* (yield* runner(sessionID, onInterrupt)).ensureRunning(work) + return yield* drain.track( + sessionID, + Effect.gen(function* () { + return yield* (yield* runner(sessionID, onInterrupt)).ensureRunning(work) + }), + ) // kilocode_change }) const startShell = Effect.fn("SessionRunState.startShell")(function* ( @@ -100,10 +108,16 @@ export const layer = Layer.effect( work: Effect.Effect, ready?: Latch.Latch, ) { - return yield* (yield* runner(sessionID, onInterrupt)) - .startShell(work, ready) - .pipe(Effect.catchTag("RunnerBusy", () => Effect.fail(busyError(sessionID)))) + return yield* drain.track( + sessionID, + Effect.gen(function* () { + return yield* (yield* runner(sessionID, onInterrupt)) + .startShell(work, ready) + .pipe(Effect.catchTag("RunnerBusy", () => Effect.fail(busyError(sessionID)))) + }), + ) }) + // kilocode_change end return Service.of({ assertNotBusy, cancel, ensureRunning, startShell }) }), @@ -149,6 +163,12 @@ function busyError(sessionID: SessionID) { return new Session.BusyError({ sessionID }) } -export const node = LayerNode.make({ service: Service, layer: layer, deps: [BackgroundJob.node, SessionStatus.node] }) +// kilocode_change start +export const node = LayerNode.make({ + service: Service, + layer: layer, + deps: [BackgroundJob.node, SessionStatus.node, SessionDrain.node], +}) +// kilocode_change end export * as SessionRunState from "./run-state" diff --git a/packages/opencode/src/tool/registry.ts b/packages/opencode/src/tool/registry.ts index 85b4be204c45..414ebbb3e84a 100644 --- a/packages/opencode/src/tool/registry.ts +++ b/packages/opencode/src/tool/registry.ts @@ -33,6 +33,7 @@ import { WebSearchTool } from "./websearch" import { KiloToolRegistry } from "../kilocode/tool/registry" // kilocode_change import { Notebook } from "@/kilocode/notebook/service" // kilocode_change import { AgentManager } from "@/kilocode/agent-manager/service" // kilocode_change +import { SessionDrain } from "@/kilocode/session/drain" // kilocode_change import { RepoOverviewTool } from "@/kilocode/tool/repo-overview" // kilocode_change import { RepoCloneTool } from "./repo_clone" // kilocode_change import { Flag } from "@opencode-ai/core/flag/flag" // kilocode_change @@ -516,6 +517,7 @@ export const node = LayerNode.suspend(() => Skill.node, Session.node, BackgroundJob.node, + SessionDrain.node, Provider.node, LSP.node, Instruction.node, diff --git a/packages/opencode/src/tool/task.ts b/packages/opencode/src/tool/task.ts index 87b822499eca..aea3fca1e17b 100644 --- a/packages/opencode/src/tool/task.ts +++ b/packages/opencode/src/tool/task.ts @@ -11,6 +11,8 @@ import { deriveSubagentSessionPermission } from "../agent/subagent-permissions" import type { SessionPrompt } from "../session/prompt" import { Config } from "@/config/config" import { Provider } from "@/provider/provider" // kilocode_change +import { SessionDrain } from "@/kilocode/session/drain" // kilocode_change +import { EventV2Bridge } from "@/event-v2-bridge" // kilocode_change import { KiloTask } from "../kilocode/tool/task" // kilocode_change import { KiloTaskBackgroundProcess } from "../kilocode/tool/task-background-process" // kilocode_change import { KiloCostPropagation } from "../kilocode/session/cost-propagation" // kilocode_change @@ -18,7 +20,7 @@ import { KiloSessionProcessor } from "../kilocode/session/processor" // kilocode import { KiloSession } from "../kilocode/session" // kilocode_change import { resumeHint } from "../kilocode/task-resume" // kilocode_change import { errorMessage } from "@/util/error" // kilocode_change -import { Effect, Exit, Schema, Scope } from "effect" +import { Cause, Effect, Exit, Schema, Scope } from "effect" // kilocode_change import { EffectBridge } from "@/effect/bridge" import { RuntimeFlags } from "@/effect/runtime-flags" import * as SandboxPolicy from "@/kilocode/sandbox/policy" // kilocode_change @@ -99,6 +101,8 @@ export const TaskTool = Tool.define( const config = yield* Config.Service const sessions = yield* Session.Service const provider = yield* Provider.Service // kilocode_change + const drain = yield* SessionDrain.Service // kilocode_change + const events = yield* EventV2Bridge.Service // kilocode_change const scope = yield* Scope.Scope const flags = yield* RuntimeFlags.Service const database = yield* Database.Service @@ -106,6 +110,7 @@ export const TaskTool = Tool.define( const run = Effect.fn("TaskTool.execute")(function* ( params: Schema.Schema.Type, ctx: Tool.Context, + delivery: SessionDrain.Delivery, // kilocode_change ) { const cfg = yield* config.get() const selection = cfg.experimental?.task_model_selection === true // kilocode_change @@ -228,6 +233,7 @@ export const TaskTool = Tool.define( // kilocode_change end // kilocode_change start - rebuild in-memory ancestry and inherit confinement after creation/resume KiloSession.register({ id: nextSession.id, parentID: ctx.sessionID, platform }) + yield* drain.link(nextSession.id, ctx.sessionID) yield* SandboxPolicy.inherit(ctx.sessionID, nextSession.id, fallback).pipe( Effect.provideService(Config.Service, config), ) @@ -253,7 +259,8 @@ export const TaskTool = Tool.define( function* () { const parts = yield* ops.resolvePromptParts(params.prompt) KiloSessionProcessor.markReviewTelemetry(parts, params.command) // kilocode_change - carry review command into child session telemetry - const result = yield* ops.prompt({ + // kilocode_change start + const initial = yield* ops.prompt({ messageID: MessageID.ascending(), sessionID: nextSession.id, model: { @@ -271,6 +278,10 @@ export const TaskTool = Tool.define( }, parts, }) + yield* drain.wait(nextSession.id) + const latest = (yield* sessions.messages({ sessionID: nextSession.id, limit: 1 })).at(-1) + const result = latest?.info.role === "assistant" && latest.info.id > initial.info.id ? latest : initial + // kilocode_change end // kilocode_change start - expose terminal child assistant errors through the task tool boundary, // including the resumable task_id so the parent agent can continue the subagent (#11620) if (result.info.role === "assistant" && result.info.error) { @@ -294,47 +305,64 @@ export const TaskTool = Tool.define( text: string, ) { const currentParent = yield* sessions.get(ctx.sessionID) - yield* ops - .prompt({ - sessionID: ctx.sessionID, - agent: currentParent.agent ?? ctx.agent, - model: selection - ? currentParent.model - ? { providerID: currentParent.model.providerID, modelID: currentParent.model.id } - : source - : undefined, - variant: selection ? (currentParent.model?.variant ?? reasoning) : variant, - parts: [ - { - type: "text", - synthetic: true, - text: renderOutput({ - sessionID: nextSession.id, - state, - summary: - state === "completed" - ? `Background task completed: ${params.description}` - : `Background task failed: ${params.description}`, - text, - }), - }, - ], - }) - .pipe(Effect.ignore, Effect.forkIn(scope, { startImmediately: true })) + yield* ops.prompt({ + sessionID: ctx.sessionID, + agent: currentParent.agent ?? ctx.agent, + model: selection + ? currentParent.model + ? { providerID: currentParent.model.providerID, modelID: currentParent.model.id } + : source + : undefined, + variant: selection ? (currentParent.model?.variant ?? reasoning) : variant, + parts: [ + { + type: "text", + synthetic: true, + text: renderOutput({ + sessionID: nextSession.id, + state, + summary: + state === "completed" + ? `Background task completed: ${params.description}` + : `Background task failed: ${params.description}`, + text, + }), + }, + ], + }) }) // kilocode_change end // kilocode_change start - background tasks propagate only cost accrued by this invocation - const notify = Effect.fn("TaskTool.notifyBackgroundResult")(function* (jobID: string) { - yield* background.wait({ id: jobID }).pipe( - Effect.flatMap((result) => { - if (result.info?.status === "completed") return inject("completed", result.info.output ?? "") - if (result.info?.status === "error") return inject("error", result.info.error ?? "") - return Effect.void + const notify = Effect.fn("TaskTool.notifyBackgroundResult")((jobID: string) => + Effect.uninterruptible( + Effect.gen(function* () { + yield* background.wait({ id: jobID }).pipe( + Effect.flatMap((result) => { + if (result.info?.status === "completed") return inject("completed", result.info.output ?? "") + if (result.info?.status === "error") return inject("error", result.info.error ?? "") + if (result.info?.status === "cancelled") return Effect.void + return Effect.die(new Error("Background task result is unavailable")) + }), + Effect.interruptible, + Effect.catchCause((cause) => + Cause.hasInterruptsOnly(cause) + ? Effect.void + : events.publish(Session.Event.Error, { + sessionID: ctx.sessionID, + error: new MessageV2.APIError({ + message: "Failed to deliver background task result", + isRetryable: false, + }).toObject(), + }), + ), + Effect.ensuring(Effect.sync(delivery.release)), + Effect.forkIn(scope, { startImmediately: true }), + ) + delivery.retained = true }), - Effect.forkIn(scope, { startImmediately: true }), - ) - }) + ), + ) const withCostPropagation = (task: Effect.Effect) => Effect.acquireUseRelease( @@ -396,6 +424,7 @@ export const TaskTool = Tool.define( }) function backgroundResult() { + delivery.retained = true // kilocode_change return { title: params.description, metadata: { @@ -494,7 +523,16 @@ export const TaskTool = Tool.define( }), ), execute: (params: Schema.Schema.Type, ctx: Tool.Context) => - run(params, ctx).pipe(Effect.orDie), + Effect.acquireUseRelease( + drain + .hold(ctx.sessionID) + .pipe(Effect.map((release): SessionDrain.Delivery => ({ release, retained: false }))), + (delivery) => run(params, ctx, delivery), + (delivery) => + Effect.sync(() => { + if (!delivery.retained) delivery.release() + }), + ).pipe(Effect.orDie), } }) // kilocode_change end diff --git a/packages/opencode/test/kilocode/cli-run-auto-helper.test.ts b/packages/opencode/test/kilocode/cli-run-auto-helper.test.ts index 2b843848ac47..8f9237c7ef7d 100644 --- a/packages/opencode/test/kilocode/cli-run-auto-helper.test.ts +++ b/packages/opencode/test/kilocode/cli-run-auto-helper.test.ts @@ -21,10 +21,17 @@ describe("KiloRunAuto", () => { }) expect(KiloRunAuto.allowed(state, "ses_child")).toBe(true) + KiloRunAuto.track(state, { + type: "tool", + tool: "task", + sessionID: "ses_child", + state: { metadata: { sessionId: "ses_grandchild" } }, + }) + expect(KiloRunAuto.allowed(state, "ses_grandchild")).toBe(true) expect(KiloRunAuto.allowed(state, "ses_other")).toBe(false) }) - test("ignores malformed or non-root task metadata", () => { + test("ignores malformed or untrusted task metadata", () => { const state = KiloRunAuto.create("ses_root") KiloRunAuto.track(state, { diff --git a/packages/opencode/test/kilocode/headless-session-drain.test.ts b/packages/opencode/test/kilocode/headless-session-drain.test.ts new file mode 100644 index 000000000000..732b7dd7e47c --- /dev/null +++ b/packages/opencode/test/kilocode/headless-session-drain.test.ts @@ -0,0 +1,297 @@ +import { expect } from "bun:test" +import { Effect, Fiber } from "effect" +import { createKiloClient, type Event } from "@kilocode/sdk/v2" +import path from "node:path" +import { pathToFileURL } from "node:url" +import { cliIt, type CliFixture } from "../lib/cli-process" +import { awaitWithTimeout } from "../lib/effect" +import { reply } from "../lib/llm-server" +import { testProviderConfig } from "../lib/test-provider" + +function gate() { + const deferred = Promise.withResolvers() + return { + promise: deferred.promise, + open: () => deferred.resolve(), + wait: (label: string) => + awaitWithTimeout( + Effect.promise(() => deferred.promise), + label, + "30 seconds", + ), + } +} + +function observer(url: string, busy: boolean) { + const sdk = pathToFileURL(path.resolve(import.meta.dir, "../../../sdk/js/src/v2/index.ts")).href + return `import { createKiloClient } from ${JSON.stringify(sdk)} + export default async ({ client: legacy }) => { + const client = createKiloClient(legacy._client.getConfig()) + const abort = new AbortController() + const ready = Promise.withResolvers() + let root + let pump + return { + "chat.message": async (input) => { + if (root) return + root = input.sessionID + const events = await client.event.subscribe(undefined, { signal: abort.signal, sseMaxRetryAttempts: 1, onSseError: (error) => { console.error(error); ready.reject(error) } }) + pump = (async () => { + let observed + try { + for await (const event of events.stream) { + if (event.type === "server.connected") { + ready.resolve() + continue + } + if (event.type === "server.instance.disposed") break + if (event.properties?.sessionID !== root) continue + if (event.type === "session.error") break + const matched = ${busy} + ? event.type === "session.queue.changed" && event.properties.queued.length > 0 + : event.type === "session.status" && event.properties.status.type === "idle" + if (!matched) continue + observed = event + break + } + } finally { + abort.abort() + } + if (!observed) return + const response = await fetch(${JSON.stringify(`${url}/events`)}, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(observed), + }) + if (!response.ok) throw new Error("event relay failed") + })().catch((error) => { + ready.reject(error) + if (!abort.signal.aborted) console.error(error) + }).finally(() => ready.reject(new Error("event relay ended"))) + await ready.promise + }, + dispose: async () => { abort.abort(); await pump }, + } + }` +} + +function probe(home: string, busy: boolean) { + return Effect.gen(function* () { + const entered = gate() + const release = gate() + const queued = gate() + const idle = gate() + const server = yield* Effect.acquireRelease( + Effect.sync(() => + Bun.serve({ + hostname: "127.0.0.1", + port: 0, + idleTimeout: 0, + async fetch(request) { + const pathname = new URL(request.url).pathname + if (pathname === "/tool") { + entered.open() + await release.promise + return new Response("PARENT_TOOL_DONE") + } + if (pathname === "/events") { + const event = (await request.json()) as Event + if (event.type === "session.queue.changed" && event.properties.queued.length > 0) queued.open() + if (event.type === "session.status" && event.properties.status.type === "idle") idle.open() + return new Response(null, { status: 204 }) + } + return new Response("unexpected request", { status: 404 }) + }, + }), + ), + (server) => + Effect.promise(async () => { + release.open() + await server.stop(true) + }), + ) + const url = `http://127.0.0.1:${server.port}` + const file = path.join(home, "drain-observer.mjs") + yield* Effect.promise(() => Bun.write(file, observer(url, busy))) + return { url, plugin: pathToFileURL(file).href, entered, release, queued, idle } + }) +} + +function scenario({ home, llm, opencode }: CliFixture, mode: "foreground" | "busy" | "idle", attached = false) { + return Effect.gen(function* () { + const tap = yield* probe(home, mode === "busy") + const child = gate() + const requested = gate() + yield* Effect.addFinalizer(() => Effect.sync(child.open)) + const base = testProviderConfig(llm.url) + const model = base.provider.test.models["test-model"] + const cfg = { + ...base, + model: "test/parent", + small_model: "test/parent", + enabled_providers: ["test"], + share: "disabled", + snapshot: false, + sandbox: { enabled: false }, + permission: { "*": "allow" }, + plugin: mode === "foreground" ? [] : [tap.plugin], + agent: { + code: { model: "test/parent" }, + general: { model: "test/child" }, + title: { disable: true }, + }, + provider: { + test: { + ...base.provider.test, + models: { + parent: { ...model, id: "parent", name: "Parent" }, + child: { ...model, id: "child", name: "Child" }, + }, + }, + }, + } + const script = `const r = await fetch(${JSON.stringify(`${tap.url}/tool`)}); if (!r.ok) throw new Error("tool gate failed"); console.log(await r.text())` + yield* llm.pushMatch( + ({ body }) => body.model === "parent", + reply().tool("task", { + description: "Produce background result", + prompt: "Produce the requested result.", + subagent_type: "general", + background: mode !== "foreground", + }), + ...(mode === "foreground" + ? [] + : [ + mode === "busy" + ? reply().tool("bash", { + command: `bun -e ${JSON.stringify(script)}`, + description: "Wait for the test request gate", + }) + : reply().text("WAITING_FOR_CHILD").stop(), + ]), + reply().text("FINAL_AFTER_CHILD").stop(), + ) + yield* llm.pushMatch(({ body }) => { + if (body.model !== "child") return false + requested.open() + return true + }, reply().wait(child.promise).text("CHILD_RESULT_SENTINEL").stop()) + if (mode === "foreground") child.open() + const env = { + PWD: home, + KILO_CONFIG_CONTENT: JSON.stringify(cfg), + KILO_CONFIG: "", + KILO_CONFIG_DIR: "", + KILO_DB: ":memory:", + KILO_AUTH_CONTENT: "{}", + KILO_SERVER_PASSWORD: "", + KILO_SERVER_USERNAME: "kilo", + KILO_PURE: "false", + KILO_DISABLE_DEFAULT_PLUGINS: "true", + KILO_EXPERIMENTAL_BACKGROUND_SUBAGENTS: "true", + KILO_TELEMETRY_LEVEL: "off", + KILO_AUTO_SHARE: "false", + } + const server = attached ? yield* opencode.serve({ env, readyTimeoutMs: 30_000 }) : undefined + const run = yield* opencode.startRun("Exercise background completion.", { + model: "test/parent", + agent: "code", + format: "json", + printLogs: true, + extraArgs: ["--auto", "--dir", home, "--title", "drain regression", ...(server ? ["--attach", server.url] : [])], + env, + }) + yield* Effect.raceFirst( + requested.wait("child never reached the mock LLM"), + run.result.pipe( + Effect.flatMap((result) => + Effect.fail(new Error(`CLI exited before child request: ${JSON.stringify(result)}`)), + ), + ), + ) + if (mode === "busy") { + yield* tap.entered.wait("parent never entered the held tool") + child.open() + yield* tap.queued.wait("child callback was not queued") + tap.release.open() + } + if (mode === "idle") { + yield* tap.idle.wait("parent never became idle") + child.open() + } + const result = yield* awaitWithTimeout(run.result, "CLI did not drain and exit", "45 seconds") + opencode.expectExit(result, 0) + const events = opencode.parseJsonEvents(result.stdout) + const texts = events.filter((event) => event.type === "text").map((event) => (event.part as { text: string }).text) + expect(texts.filter((text) => text === "FINAL_AFTER_CHILD")).toHaveLength(1) + expect(events.some((event) => event.type === "error")).toBe(false) + const inputs = yield* llm.inputs + const parents = inputs.filter((input) => input.model === "parent") + expect(parents).toHaveLength(mode === "foreground" ? 2 : 3) + expect(inputs.filter((input) => input.model === "child")).toHaveLength(1) + expect(JSON.stringify(parents.at(-1))).toContain("CHILD_RESULT_SENTINEL") + if (mode !== "foreground") expect(JSON.stringify(parents.at(-1))).toContain("Background task completed") + if (mode === "busy") expect(JSON.stringify(parents.at(-1))).toContain("PARENT_TOOL_DONE") + expect(yield* llm.pending).toBe(0) + }) +} + +cliIt.live("foreground Task completes before headless exit", (fixture) => scenario(fixture, "foreground"), 90_000) +cliIt.live( + "headless run drains a child callback queued behind a busy parent", + (fixture) => scenario(fixture, "busy"), + 90_000, +) +cliIt.live( + "headless run waits for a child after the parent becomes idle", + (fixture) => scenario(fixture, "idle"), + 90_000, +) +cliIt.live( + "attached run drains a callback queued behind the parent", + (fixture) => scenario(fixture, "busy", true), + 90_000, +) +cliIt.live( + "attached run waits for a child after the parent becomes idle", + (fixture) => scenario(fixture, "idle", true), + 90_000, +) + +cliIt.live( + "SDK drain validates requests and publishes its matching acknowledgment", + ({ home, opencode }) => + Effect.gen(function* () { + const server = yield* opencode.serve({ + env: { KILO_SERVER_PASSWORD: "", KILO_SERVER_USERNAME: "kilo" }, + readyTimeoutMs: 30_000, + }) + const sdk = createKiloClient({ baseUrl: server.url, directory: home }) + const session = yield* Effect.promise(() => sdk.session.create({})) + if (!session.data) throw new Error("Session creation failed") + const id = session.data.id + const abort = new AbortController() + yield* Effect.addFinalizer(() => Effect.sync(() => abort.abort())) + const ready = gate() + const token = crypto.randomUUID() + const stream = yield* Effect.promise(() => + sdk.event.subscribe(undefined, { signal: abort.signal, sseMaxRetryAttempts: 1 }), + ) + const received = yield* Effect.promise(async () => { + for await (const event of stream.stream) { + if (event.type === "server.connected") ready.open() + if (event.type === "session.drained" && event.properties.token === token) return event.properties + } + throw new Error("Drain acknowledgment was not delivered") + }).pipe(Effect.forkChild) + yield* ready.wait("Event stream did not connect") + const result = yield* Effect.promise(() => sdk.kilocode.drainSession({ sessionID: id, token })) + expect(result.data).toBe(true) + expect(yield* Fiber.join(received)).toEqual({ sessionID: id, token }) + const invalid = yield* Effect.promise(() => sdk.kilocode.drainSession({ sessionID: id, token: "" })) + expect(invalid.response?.status).toBe(400) + const missing = yield* Effect.promise(() => sdk.kilocode.drainSession({ sessionID: "ses_missing_drain", token })) + expect(missing.response?.status).toBe(404) + }), + 90_000, +) diff --git a/packages/opencode/test/kilocode/run-auto.test.ts b/packages/opencode/test/kilocode/run-auto.test.ts index 4d67e1607e58..c50dfb2f95ab 100644 --- a/packages/opencode/test/kilocode/run-auto.test.ts +++ b/packages/opencode/test/kilocode/run-auto.test.ts @@ -112,8 +112,10 @@ function args() { } const tty = Object.getOwnPropertyDescriptor(process.stdin, "isTTY") +const exitCode = process.exitCode afterEach(() => { + process.exitCode = exitCode if (tty) { Object.defineProperty(process.stdin, "isTTY", tty) return @@ -123,7 +125,15 @@ afterEach(() => { async function run(sdk: Record) { mock.module("@kilocode/sdk/v2", () => ({ - createKiloClient: () => sdk, + createKiloClient: (config: { fetch?: () => Promise }) => { + config.fetch = async () => + Response.json({ + paths: { + "/kilocode/session/{sessionID}/drain": { post: { operationId: "kilocode.drainSession" } }, + }, + }) + return sdk + }, })) Object.defineProperty(process.stdin, "isTTY", { @@ -147,7 +157,17 @@ describe("cli run auto permissions", () => { get: async () => ({ data: { share: "manual" } }), }, event: { - subscribe: async () => ({ stream: q.stream() }), + subscribe: async () => { + q.push({ type: "server.connected", properties: {} }) + return { stream: q.stream() } + }, + }, + kilocode: { + drainSession: async (input: { sessionID: string; token: string }) => { + q.push({ type: "session.drained", properties: input }) + q.end() + return { data: true } + }, }, permission: { reply: async (input: { requestID: string; reply: string }) => { @@ -165,8 +185,7 @@ describe("cli run auto permissions", () => { q.push(permission("perm_other", "ses_other")) q.push(permission("perm_child", "ses_child")) q.push(idle()) - await Promise.race([done.promise, new Promise((resolve) => setTimeout(resolve, 25))]) - q.end() + await done.promise return { data: undefined } }, }, @@ -176,4 +195,90 @@ describe("cli run auto permissions", () => { expect(calls).toEqual([{ requestID: "perm_child", reply: "once" }]) }) + + test("a failed prompt aborts a blocked permission reply during cleanup", async () => { + const q = feed() + const started = Promise.withResolvers() + let aborted = false + const sdk = { + config: { get: async () => ({ data: { share: "manual" } }) }, + event: { + subscribe: async () => { + q.push({ type: "server.connected", properties: {} }) + return { stream: q.stream() } + }, + }, + permission: { + reply: async (_input: unknown, options: { signal: AbortSignal }) => { + expect(options.signal).toBeInstanceOf(AbortSignal) + started.resolve() + return new Promise((_resolve, reject) => { + options.signal.addEventListener( + "abort", + () => { + aborted = true + reject(options.signal.reason) + }, + { once: true }, + ) + }) + }, + }, + session: { + get: async (input: { sessionID: string }) => ({ data: { id: input.sessionID, directory: "/tmp/project" } }), + prompt: async () => { + q.push(task("ses_child")) + q.push(permission("perm_child", "ses_child")) + await started.promise + throw new Error("intentional prompt failure") + }, + }, + } + await run(sdk) + q.end() + expect(aborted).toBe(true) + expect(process.exitCode).toBe(1) + }) + + test("handles a tracked child's network retry without touching another session", async () => { + const q = feed() + const done = Promise.withResolvers() + const calls: string[] = [] + const sdk = { + config: { get: async () => ({ data: { share: "manual" } }) }, + event: { + subscribe: async () => { + q.push({ type: "server.connected", properties: {} }) + return { stream: q.stream() } + }, + }, + network: { + reply: async (input: { requestID: string }, options: { signal: AbortSignal }) => { + expect(options.signal).toBeInstanceOf(AbortSignal) + calls.push(input.requestID) + done.resolve() + return { data: true } + }, + }, + kilocode: { + drainSession: async (input: { sessionID: string; token: string }) => { + q.push({ type: "session.drained", properties: input }) + q.end() + return { data: true } + }, + }, + session: { + get: async (input: { sessionID: string }) => ({ data: { id: input.sessionID, directory: "/tmp/project" } }), + prompt: async () => { + q.push(task("ses_child")) + q.push({ type: "session.network.asked", properties: { id: "net_other", sessionID: "ses_other" } }) + q.push({ type: "session.network.asked", properties: { id: "net_child", sessionID: "ses_child" } }) + await done.promise + return { data: undefined } + }, + }, + } + await run(sdk) + expect(calls).toEqual(["net_child"]) + }, 15_000) }) diff --git a/packages/opencode/test/kilocode/run-drain.test.ts b/packages/opencode/test/kilocode/run-drain.test.ts new file mode 100644 index 000000000000..0f8507038fd9 --- /dev/null +++ b/packages/opencode/test/kilocode/run-drain.test.ts @@ -0,0 +1,112 @@ +import { expect, test } from "bun:test" +import { KiloRunDrain } from "@/kilocode/cli/run-drain" + +function client(handle: (request: Request) => Promise) { + return KiloRunDrain.client({ + baseUrl: "http://drain.test/prefix", + headers: { Authorization: "Basic test-only" }, + fetch: Object.assign( + (input: RequestInfo | URL, init?: RequestInit) => + handle(input instanceof Request ? input : new Request(input, init)), + { preconnect: () => {} }, + ), + }) +} + +function event(token: string, sessionID = "ses_parent") { + return { id: "evt_drain", type: "session.drained" as const, properties: { sessionID, token } } +} + +test("checks capabilities through the selected transport before accepting completion", async () => { + const sdk = client(async (request) => { + expect(new URL(request.url).pathname).toBe("/prefix/doc") + expect(request.headers.get("authorization")).toBe("Basic test-only") + return Response.json({ + paths: { "/kilocode/session/{sessionID}/drain": { post: { operationId: "kilocode.drainSession" } } }, + }) + }) + await KiloRunDrain.check(sdk, new AbortController().signal) +}) + +test.each([ + () => new Response("old server", { headers: { "content-type": "text/html" } }), + () => Response.json({ paths: {} }), + () => Response.json({ error: "Not Found" }, { status: 404 }), +])("rejects unsupported or HTML-returning servers", async (response) => { + const sdk = client(async () => response()) + await expect(KiloRunDrain.check(sdk, new AbortController().signal)).rejects.toThrow( + "does not support session draining", + ) +}) + +test("accepts an acknowledgment that arrives before the HTTP result", async () => { + const response = Promise.withResolvers() + const drain = KiloRunDrain.create("ses_parent") + const sdk = client(() => response.promise) + const waiting = drain.wait(sdk) + expect(drain.event(event("wrong"))).toBe(false) + expect(drain.event(event(drain.token, "ses_other"))).toBe(false) + expect(drain.event(event(drain.token))).toBe(true) + response.resolve(Response.json(true)) + await waiting + drain.close() +}) + +test("accepts the HTTP result before the matching acknowledgment", async () => { + const received = Promise.withResolvers() + const drain = KiloRunDrain.create("ses_parent") + const sdk = client(async () => Response.json(true)) + const original = sdk.kilocode.drainSession.bind(sdk.kilocode) + sdk.kilocode.drainSession = (params, options) => { + const result = original(params, options) + void result.then(() => received.resolve()) + return result + } + const waiting = drain.wait(sdk) + await received.promise + expect(drain.event(event(drain.token))).toBe(true) + await waiting + drain.close() +}) + +test("stream loss fails a pending wait instead of reporting completion", async () => { + const response = Promise.withResolvers() + const drain = KiloRunDrain.create("ses_parent") + const sdk = client(() => response.promise) + const waiting = drain.wait(sdk) + drain.end() + await expect(waiting).rejects.toThrow("event stream ended before completion") + response.resolve(Response.json(true)) + drain.close() +}) + +test("closing a run aborts pending retry waits", async () => { + const drain = KiloRunDrain.create("ses_parent") + const waiting = drain.pause(60_000) + drain.close() + await expect(waiting).rejects.toThrow() +}) + +test("a root interruption fails but a superseded handoff and child interruption do not", async () => { + const drain = KiloRunDrain.create("ses_parent") + drain.event({ + id: "evt_superseded", + type: "session.turn.close", + properties: { sessionID: "ses_parent", reason: "superseded" }, + }) + drain.event({ + id: "evt_child", + type: "session.turn.close", + properties: { sessionID: "ses_child", reason: "interrupted" }, + }) + drain.event({ id: "evt_ready", type: "server.connected", properties: {} }) + await drain.ready() + const waiting = drain.race(new Promise(() => undefined)) + drain.event({ + id: "evt_root", + type: "session.turn.close", + properties: { sessionID: "ses_parent", reason: "interrupted" }, + }) + await expect(waiting).rejects.toThrow("Session interrupted") + drain.close() +}) diff --git a/packages/opencode/test/kilocode/session-drain.test.ts b/packages/opencode/test/kilocode/session-drain.test.ts new file mode 100644 index 000000000000..b3f4fd33065a --- /dev/null +++ b/packages/opencode/test/kilocode/session-drain.test.ts @@ -0,0 +1,123 @@ +import { expect } from "bun:test" +import { Effect, Fiber } from "effect" +import { SessionDrain } from "@/kilocode/session/drain" +import { SessionID } from "@/session/schema" +import { testEffect } from "../lib/effect" + +const it = testEffect(SessionDrain.layer) + +it.instance( + "drains only after execution and delivery reservations end", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const id = SessionID.make("ses_drain_parent") + const parent = yield* drain.hold(id) + const delivery = yield* drain.hold(id) + let done = false + const waiter = yield* drain.wait(id).pipe( + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + expect(done).toBe(false) + parent() + expect(done).toBe(false) + const callback = yield* drain.hold(id) + delivery() + expect(done).toBe(false) + callback() + yield* Fiber.join(waiter) + expect(done).toBe(true) + }), +) + +it.instance( + "includes live descendants without blocking other sessions", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const parent = SessionID.make("ses_drain_parent") + const child = SessionID.make("ses_drain_child") + const grandchild = SessionID.make("ses_drain_grandchild") + yield* drain.link(child, parent) + yield* drain.link(grandchild, child) + const release = yield* drain.hold(grandchild) + let done = false + const waiter = yield* drain.wait(parent).pipe( + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + yield* drain.wait(SessionID.make("ses_drain_other")) + expect(done).toBe(false) + release() + yield* Fiber.join(waiter) + yield* drain.wait(child) + yield* drain.wait(grandchild) + }), +) + +it.instance( + "a stale release cannot finish a newer reservation", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const id = SessionID.make("ses_drain_parent") + const previous = yield* drain.hold(id) + const next = yield* drain.hold(id) + let done = false + const waiter = yield* drain.wait(id).pipe( + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + previous() + previous() + expect(done).toBe(false) + next() + yield* Fiber.join(waiter) + }), +) + +it.instance( + "interrupting a wait does not cancel tracked work", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const id = SessionID.make("ses_drain_parent") + const release = yield* drain.hold(id) + const cancelled = yield* drain.wait(id).pipe(Effect.forkChild({ startImmediately: true })) + yield* Fiber.interrupt(cancelled) + let done = false + const waiter = yield* drain.wait(id).pipe( + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + expect(done).toBe(false) + release() + yield* Fiber.join(waiter) + }), +) + +it.instance( + "tracks failures without retaining work or poisoning later waits", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const id = SessionID.make("ses_drain_parent") + const result = yield* drain.track(id, Effect.fail("handled child failure")).pipe(Effect.exit) + expect(result._tag).toBe("Failure") + yield* drain.wait(id) + yield* drain.track(id, Effect.void) + yield* drain.wait(id) + }), +) diff --git a/packages/opencode/test/kilocode/task-nesting.test.ts b/packages/opencode/test/kilocode/task-nesting.test.ts index 16af26481bfc..a25950aa23c5 100644 --- a/packages/opencode/test/kilocode/task-nesting.test.ts +++ b/packages/opencode/test/kilocode/task-nesting.test.ts @@ -8,6 +8,8 @@ import { BackgroundJob } from "../../src/background/job" import { Bus } from "../../src/bus" import { SessionRunState } from "../../src/session/run-state" import { SessionStatus } from "../../src/session/status" +import { SessionDrain } from "@/kilocode/session/drain" +import { EventV2Bridge } from "@/event-v2-bridge" import { Config } from "../../src/config/config" import { RuntimeFlags } from "../../src/effect/runtime-flags" import * as CrossSpawnSpawner from "@opencode-ai/core/cross-spawn-spawner" @@ -45,6 +47,8 @@ const it = testEffect( RuntimeFlags.node, SessionRunState.node, SessionStatus.node, + SessionDrain.node, + EventV2Bridge.node, CrossSpawnSpawner.node, Session.node, SessionProjector.node, diff --git a/packages/opencode/test/kilocode/tool-task-model.test.ts b/packages/opencode/test/kilocode/tool-task-model.test.ts index 17ca53022e10..7a5863ca7f61 100644 --- a/packages/opencode/test/kilocode/tool-task-model.test.ts +++ b/packages/opencode/test/kilocode/tool-task-model.test.ts @@ -10,6 +10,8 @@ import { BackgroundJob } from "../../src/background/job" import { Bus } from "../../src/bus" import { SessionRunState } from "../../src/session/run-state" import { SessionStatus } from "../../src/session/status" +import { SessionDrain } from "@/kilocode/session/drain" +import { EventV2Bridge } from "@/event-v2-bridge" import { Config } from "../../src/config/config" import { RuntimeFlags } from "../../src/effect/runtime-flags" import * as CrossSpawnSpawner from "@opencode-ai/core/cross-spawn-spawner" @@ -110,6 +112,8 @@ const it = testEffect( RuntimeFlags.node, SessionRunState.node, SessionStatus.node, + SessionDrain.node, + EventV2Bridge.node, CrossSpawnSpawner.node, Session.node, SessionProjector.node, diff --git a/packages/opencode/test/tool/task.test.ts b/packages/opencode/test/tool/task.test.ts index 905026238bd0..5134ec790aac 100644 --- a/packages/opencode/test/tool/task.test.ts +++ b/packages/opencode/test/tool/task.test.ts @@ -15,6 +15,7 @@ import { MessageV2 } from "@/session/message-v2" // kilocode_change import type { SessionPrompt } from "../../src/session/prompt" import { MessageID, PartID, SessionID } from "../../src/session/schema" // kilocode_change - SessionID used by cost propagation tests import { SessionRunState } from "@/session/run-state" +import { SessionDrain } from "@/kilocode/session/drain" // kilocode_change import { SessionStatus } from "@/session/status" import { Provider } from "../../src/provider/provider" // kilocode_change import { KiloSession } from "../../src/kilocode/session" // kilocode_change @@ -47,6 +48,7 @@ const layer = (flags: Partial = {}) => Session.node, SessionProjector.node, SessionRunState.node, + SessionDrain.node, // kilocode_change SessionStatus.node, Truncate.node, ToolRegistry.node, @@ -1402,6 +1404,59 @@ describe("tool.task", () => { }) // kilocode_change start - subagent cost propagation coverage (#6321) +it.instance("returns the child's final answer after nested background delivery drains", () => + Effect.gen(function* () { + const sessions = yield* Session.Service + const drain = yield* SessionDrain.Service + const { chat, assistant } = yield* seed() + const entered = yield* Deferred.make<{ input: SessionPrompt.PromptInput; release: () => void }>() + const factory = yield* TaskTool + const tool = yield* factory.init() + const ops: TaskPromptOps = { + ...stubOps(), + prompt: (input) => + Effect.gen(function* () { + const grandchild = yield* sessions.create({ parentID: input.sessionID, title: "nested background" }) + yield* drain.link(grandchild.id, input.sessionID) + const release = yield* drain.hold(grandchild.id) + const initial = reply(input, "WAITING_FOR_NESTED_CHILD") + yield* sessions.updateMessage(initial.info) + for (const part of initial.parts) yield* sessions.updatePart(part) + yield* Deferred.succeed(entered, { input, release }) + return initial + }), + } + const running = yield* tool + .execute( + { + description: "Nested result", + prompt: "Return the completed result", + subagent_type: "general", + background: false, + }, + { + sessionID: chat.id, + messageID: assistant.id, + agent: "build", + abort: new AbortController().signal, + extra: { promptOps: ops }, + messages: [], + metadata: () => Effect.void, + ask: () => Effect.void, + }, + ) + .pipe(Effect.forkChild) + const pending = yield* Deferred.await(entered) + const final = reply(pending.input, "FINAL_NESTED_RESULT") + yield* sessions.updateMessage(final.info) + for (const part of final.parts) yield* sessions.updatePart(part) + pending.release() + const result = yield* Fiber.join(running) + expect(result.output).toContain("FINAL_NESTED_RESULT") + expect(result.output).not.toContain("WAITING_FOR_NESTED_CHILD") + }), +) + const assistantCost = Effect.fn("TaskToolTest.assistantCost")(function* (sessionID: string) { const sessions = yield* Session.Service const msgs = yield* sessions.messages({ sessionID: SessionID.make(sessionID) }) diff --git a/packages/schema/src/event-manifest.ts b/packages/schema/src/event-manifest.ts index b681362e83be..d015e89915b5 100644 --- a/packages/schema/src/event-manifest.ts +++ b/packages/schema/src/event-manifest.ts @@ -30,6 +30,7 @@ import { TuiEvent } from "./tui-event" import { VcsEvent } from "./vcs-event" import { WorkspaceEvent } from "./workspace-event" import { WorktreeEvent } from "./worktree-event" +import * as SessionDrain from "./kilocode/session-drain" // kilocode_change const sessionV1DurableDefinitions = SessionV1.Event.Definitions.filter((definition) => definition.durable !== undefined) const sessionV1LiveDefinitions = SessionV1.Event.Definitions.filter((definition) => definition.durable === undefined) @@ -79,6 +80,7 @@ export const Definitions = Event.inventory( ...WorkspaceEvent.Definitions, ...WorktreeEvent.Definitions, ...ServerEvent.Definitions, + ...SessionDrain.Definitions, // kilocode_change ) export const Latest = Event.latest(Definitions) export { Durable } diff --git a/packages/schema/src/kilocode/session-drain.ts b/packages/schema/src/kilocode/session-drain.ts new file mode 100644 index 000000000000..ef481e073448 --- /dev/null +++ b/packages/schema/src/kilocode/session-drain.ts @@ -0,0 +1,12 @@ +import { Schema } from "effect" +import { Event } from "../event" +import { SessionID } from "../session-id" + +export const Token = Schema.String.check(Schema.isMinLength(1), Schema.isMaxLength(128)) + +export const Drained = Event.define({ + type: "session.drained", + schema: { sessionID: SessionID, token: Token }, +}) + +export const Definitions = Event.inventory(Drained) diff --git a/packages/sdk/js/src/v2/gen/sdk.gen.ts b/packages/sdk/js/src/v2/gen/sdk.gen.ts index db4a92420832..92f81cea4dcd 100644 --- a/packages/sdk/js/src/v2/gen/sdk.gen.ts +++ b/packages/sdk/js/src/v2/gen/sdk.gen.ts @@ -191,6 +191,8 @@ import type { KilocodeBackgroundJobsResponses, KilocodeCommandFilesErrors, KilocodeCommandFilesResponses, + KilocodeDrainSessionErrors, + KilocodeDrainSessionResponses, KilocodeHeapSnapshotErrors, KilocodeHeapSnapshotResponses, KilocodeNotebookListErrors, @@ -8369,6 +8371,49 @@ export class Kilocode extends HeyApiClient { }) } + /** + * Wait for session completion + * + * Wait for active session work and background result delivery, then publish the matching drain acknowledgment. + */ + public drainSession( + parameters: { + sessionID: string + directory?: string + workspace?: string + token: string + }, + options?: Options, + ) { + const params = buildClientParams( + [parameters], + [ + { + args: [ + { in: "path", key: "sessionID" }, + { in: "query", key: "directory" }, + { in: "query", key: "workspace" }, + { in: "body", key: "token" }, + ], + }, + ], + ) + return (options?.client ?? this.client).post< + KilocodeDrainSessionResponses, + KilocodeDrainSessionErrors, + ThrowOnError + >({ + url: "/kilocode/session/{sessionID}/drain", + ...options, + ...params, + headers: { + "Content-Type": "application/json", + ...options?.headers, + ...params.headers, + }, + }) + } + /** * List command files * diff --git a/packages/sdk/js/src/v2/gen/types.gen.ts b/packages/sdk/js/src/v2/gen/types.gen.ts index f5d69d48001a..b7cf3a51deed 100644 --- a/packages/sdk/js/src/v2/gen/types.gen.ts +++ b/packages/sdk/js/src/v2/gen/types.gen.ts @@ -95,6 +95,7 @@ export type Event = | EventServerConnected1 | EventGlobalDisposed1 | EventGlobalConfigUpdated1 + | EventSessionDrained1 | EventServerInstanceDisposed | EventSessionTurnOpen | EventSessionTurnClose @@ -214,6 +215,7 @@ export type Event = | EventServerConnected | EventGlobalDisposed | EventGlobalConfigUpdated + | EventSessionDrained export type QuestionReplied = { sessionID: string @@ -1272,6 +1274,7 @@ export type GlobalEvent = { | EventServerConnected | EventGlobalDisposed | EventGlobalConfigUpdated + | EventSessionDrained | { id: string type: "models-dev.refreshed" @@ -2157,6 +2160,14 @@ export type GlobalEvent = { [key: string]: unknown } } + | { + id: string + type: "session.drained" + properties: { + sessionID: string + token: string + } + } | SyncEventSessionCreated | SyncEventSessionUpdated | SyncEventSessionDeleted @@ -4740,6 +4751,7 @@ export type V2Event = | ServerConnected | GlobalDisposed | GlobalConfigUpdated + | SessionDrained export type V2EventStream = string @@ -6230,6 +6242,15 @@ export type EventGlobalConfigUpdated = { } } +export type EventSessionDrained = { + id: string + type: "session.drained" + properties: { + sessionID: string + token: string + } +} + export type SyncEventSessionCreated = { type: "sync" id: string @@ -9204,6 +9225,24 @@ export type GlobalConfigUpdated = { } } +export type SessionDrained = { + id: string + metadata?: { + [key: string]: unknown + } + type: "session.drained" + durable?: { + aggregateID: string + seq: number + version: number + } + location?: LocationRef + data: { + sessionID: string + token: string + } +} + export type QuestionV2Request = { id: string sessionID: string @@ -10225,6 +10264,15 @@ export type EventGlobalConfigUpdated1 = { } } +export type EventSessionDrained1 = { + id: string + type: "session.drained" + properties: { + sessionID: string + token: string + } +} + export type EventMemoryStatus1 = { id: string type: "memory.status" @@ -16629,6 +16677,42 @@ export type KilocodeResumeSessionResponses = { export type KilocodeResumeSessionResponse = KilocodeResumeSessionResponses[keyof KilocodeResumeSessionResponses] +export type KilocodeDrainSessionData = { + body?: { + token: string + } + path: { + sessionID: string + } + query?: { + directory?: string + workspace?: string + } + url: "/kilocode/session/{sessionID}/drain" +} + +export type KilocodeDrainSessionErrors = { + /** + * Bad request + */ + 400: BadRequestError + /** + * NotFoundError + */ + 404: NotFoundError +} + +export type KilocodeDrainSessionError = KilocodeDrainSessionErrors[keyof KilocodeDrainSessionErrors] + +export type KilocodeDrainSessionResponses = { + /** + * Session work drained + */ + 200: boolean +} + +export type KilocodeDrainSessionResponse = KilocodeDrainSessionResponses[keyof KilocodeDrainSessionResponses] + export type KilocodeHeapSnapshotData = { body?: never path?: never diff --git a/packages/sdk/openapi.json b/packages/sdk/openapi.json index 47ec386a2c9e..7aa56fcd3f53 100644 --- a/packages/sdk/openapi.json +++ b/packages/sdk/openapi.json @@ -15200,6 +15200,98 @@ ] } }, + "/kilocode/session/{sessionID}/drain": { + "post": { + "tags": ["kilocode"], + "operationId": "kilocode.drainSession", + "parameters": [ + { + "name": "sessionID", + "in": "path", + "schema": { + "type": "string", + "pattern": "^ses.*" + }, + "required": true + }, + { + "name": "directory", + "in": "query", + "schema": { + "type": "string" + }, + "required": false + }, + { + "name": "workspace", + "in": "query", + "schema": { + "type": "string" + }, + "required": false + } + ], + "responses": { + "200": { + "description": "Session work drained", + "content": { + "application/json": { + "schema": { + "type": "boolean", + "description": "Session work drained" + } + } + } + }, + "400": { + "description": "Bad request", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/BadRequestError" + } + } + } + }, + "404": { + "description": "NotFoundError", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/NotFoundError" + } + } + } + } + }, + "description": "Wait for active session work and background result delivery, then publish the matching drain acknowledgment.", + "summary": "Wait for session completion", + "requestBody": { + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "token": { + "type": "string", + "minLength": 1, + "maxLength": 128 + } + }, + "required": ["token"], + "additionalProperties": false + } + } + } + }, + "x-codeSamples": [ + { + "lang": "js", + "source": "import { createKiloClient } from \"@kilocode/sdk\"\n\nconst client = createKiloClient()\nawait client.kilocode.drainSession({\n ...\n})" + } + ] + } + }, "/kilocode/heap/snapshot": { "post": { "tags": ["kilocode"], @@ -26685,6 +26777,9 @@ { "$ref": "#/components/schemas/EventGlobalConfigUpdated1" }, + { + "$ref": "#/components/schemas/EventSessionDrained1" + }, { "$ref": "#/components/schemas/EventServerInstanceDisposed" }, @@ -27042,6 +27137,9 @@ { "$ref": "#/components/schemas/EventGlobalConfigUpdated" }, + { + "$ref": "#/components/schemas/EventSessionDrained" + }, { "$ref": "#/components/schemas/EventServerInstanceDisposed" } @@ -30239,6 +30337,9 @@ { "$ref": "#/components/schemas/EventGlobalConfigUpdated" }, + { + "$ref": "#/components/schemas/EventSessionDrained" + }, { "type": "object", "properties": { @@ -33294,6 +33395,37 @@ "required": ["id", "type", "properties"], "additionalProperties": false }, + { + "type": "object", + "properties": { + "id": { + "type": "string", + "pattern": "^evt_" + }, + "type": { + "type": "string", + "enum": ["session.drained"] + }, + "properties": { + "type": "object", + "properties": { + "sessionID": { + "type": "string", + "pattern": "^ses" + }, + "token": { + "type": "string", + "minLength": 1, + "maxLength": 128 + } + }, + "required": ["sessionID", "token"], + "additionalProperties": false + } + }, + "required": ["id", "type", "properties"], + "additionalProperties": false + }, { "$ref": "#/components/schemas/EventServerInstanceDisposed" }, @@ -41135,6 +41267,9 @@ }, { "$ref": "#/components/schemas/GlobalConfigUpdated" + }, + { + "$ref": "#/components/schemas/SessionDrained" } ] }, @@ -46327,6 +46462,36 @@ "required": ["id", "type", "properties"], "additionalProperties": false }, + "EventSessionDrained": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "type": { + "type": "string", + "enum": ["session.drained"] + }, + "properties": { + "type": "object", + "properties": { + "sessionID": { + "type": "string", + "pattern": "^ses" + }, + "token": { + "type": "string", + "minLength": 1, + "maxLength": 128 + } + }, + "required": ["sessionID", "token"], + "additionalProperties": false + } + }, + "required": ["id", "type", "properties"], + "additionalProperties": false + }, "SyncEventSessionCreated": { "type": "object", "properties": { @@ -55282,6 +55447,59 @@ "required": ["id", "type", "data"], "additionalProperties": false }, + "SessionDrained": { + "type": "object", + "properties": { + "id": { + "type": "string", + "pattern": "^evt_" + }, + "metadata": { + "type": "object" + }, + "type": { + "type": "string", + "enum": ["session.drained"] + }, + "durable": { + "type": "object", + "properties": { + "aggregateID": { + "type": "string" + }, + "seq": { + "type": "integer" + }, + "version": { + "type": "integer" + } + }, + "required": ["aggregateID", "seq", "version"], + "additionalProperties": false + }, + "location": { + "$ref": "#/components/schemas/LocationRef" + }, + "data": { + "type": "object", + "properties": { + "sessionID": { + "type": "string", + "pattern": "^ses" + }, + "token": { + "type": "string", + "minLength": 1, + "maxLength": 128 + } + }, + "required": ["sessionID", "token"], + "additionalProperties": false + } + }, + "required": ["id", "type", "data"], + "additionalProperties": false + }, "QuestionV2Request": { "type": "object", "properties": { @@ -58453,6 +58671,37 @@ "required": ["id", "type", "properties"], "additionalProperties": false }, + "EventSessionDrained1": { + "type": "object", + "properties": { + "id": { + "type": "string", + "pattern": "^evt_" + }, + "type": { + "type": "string", + "enum": ["session.drained"] + }, + "properties": { + "type": "object", + "properties": { + "sessionID": { + "type": "string", + "pattern": "^ses" + }, + "token": { + "type": "string", + "minLength": 1, + "maxLength": 128 + } + }, + "required": ["sessionID", "token"], + "additionalProperties": false + } + }, + "required": ["id", "type", "properties"], + "additionalProperties": false + }, "EventMemoryStatus1": { "type": "object", "properties": { From a3f4b4674ae35c1a8046ec89b3b822ec996abe5d Mon Sep 17 00:00:00 2001 From: Josh Lambert Date: Mon, 31 Aug 2026 14:21:13 -0400 Subject: [PATCH 2/5] fix(cli): harden drain lifecycle and update CI fixtures --- .../opencode/src/effect/instance-registry.ts | 33 +++- .../opencode/src/kilocode/cli/run-drain.ts | 34 +++- .../opencode/src/kilocode/session/drain.ts | 146 ++++++++++++------ packages/opencode/src/tool/task.ts | 22 +-- packages/opencode/test/event-manifest.test.ts | 4 +- .../kilocode/headless-session-drain.test.ts | 4 +- .../opencode/test/kilocode/run-auto.test.ts | 2 +- .../opencode/test/kilocode/run-drain.test.ts | 58 +++++++ .../test/kilocode/run-network.test.ts | 86 ++++++----- .../server/httpapi-exercise-scenarios.ts | 18 +++ .../test/kilocode/session-drain.test.ts | 137 +++++++++++++++- packages/opencode/test/tool/task.test.ts | 82 ++++++++++ 12 files changed, 520 insertions(+), 106 deletions(-) diff --git a/packages/opencode/src/effect/instance-registry.ts b/packages/opencode/src/effect/instance-registry.ts index 8c4c85c9280a..925555a3854b 100644 --- a/packages/opencode/src/effect/instance-registry.ts +++ b/packages/opencode/src/effect/instance-registry.ts @@ -3,6 +3,20 @@ import type { WorkspaceV2 } from "@opencode-ai/core/workspace" // kilocode_chang const disposers = new Set<(directory: string, workspaceID?: WorkspaceV2.ID) => Promise>() // kilocode_change // kilocode_change start +const beforeDisposers = new Set<(directory: string, workspaceID?: WorkspaceV2.ID) => undefined>() +const closing = new Map() + +export function registerBeforeDisposer(disposer: (directory: string, workspaceID?: WorkspaceV2.ID) => undefined) { + beforeDisposers.add(disposer) + return () => { + beforeDisposers.delete(disposer) + } +} + +export function isDisposing(directory: string) { + return closing.has(directory) +} + export function registerDisposer( disposer: (directory: string, workspaceID?: WorkspaceV2.ID) => Promise, // kilocode_change ) { @@ -13,6 +27,23 @@ export function registerDisposer( } export async function disposeInstance(directory: string, workspaceID?: WorkspaceV2.ID) { - await Promise.allSettled([...disposers].map((disposer) => disposer(directory, workspaceID))) + closing.set(directory, (closing.get(directory) ?? 0) + 1) + try { + const errors: unknown[] = [] + const guards = [...beforeDisposers] + for (const disposer of guards) { + try { + disposer(directory, workspaceID) + } catch (error) { + errors.push(error) + } + } + await Promise.allSettled([...disposers].map((disposer) => disposer(directory, workspaceID))) + if (errors.length) throw new AggregateError(errors, "Instance pre-disposal failed") + } finally { + const remaining = (closing.get(directory) ?? 1) - 1 + if (remaining) closing.set(directory, remaining) + else closing.delete(directory) + } } // kilocode_change end diff --git a/packages/opencode/src/kilocode/cli/run-drain.ts b/packages/opencode/src/kilocode/cli/run-drain.ts index 1ceb9b843866..ac900ff12f33 100644 --- a/packages/opencode/src/kilocode/cli/run-drain.ts +++ b/packages/opencode/src/kilocode/cli/run-drain.ts @@ -31,10 +31,19 @@ export namespace KiloRunDrain { } const request = new Request(`${config.baseUrl.replace(/\/$/, "")}/doc`, { headers, signal }) const response = await (config.fetch ?? globalThis.fetch)(request) + if (response.status === 401 || response.status === 403) { + throw new Error(`Server rejected access (HTTP ${response.status}); check server authentication and permissions`) + } + if (!response.ok && response.status !== 404) { + throw new Error(`Server capability check failed (HTTP ${response.status})`) + } if (!response.ok || !response.headers.get("content-type")?.toLowerCase().includes("application/json")) { throw new Error("Server does not support session draining; upgrade or restart the server") } - if (!capability.safeParse(await response.json()).success) { + const body: unknown = await response.json().catch(() => { + throw new Error("Server returned invalid capability JSON") + }) + if (!capability.safeParse(body).success) { throw new Error("Server does not support session draining; upgrade or restart the server") } } @@ -96,12 +105,29 @@ export namespace KiloRunDrain { } } - export async function flush() { + export async function flush( + streams: readonly { write(chunk: string, callback: (error?: Error | null) => void): unknown }[] = [ + process.stdout, + process.stderr, + ], + ) { await Promise.all( - [process.stdout, process.stderr].map( + streams.map( (stream) => new Promise((resolve, reject) => { - stream.write("", (error) => (error ? reject(error) : resolve())) + const done = (error?: unknown) => { + const closed = + error instanceof Error && + "code" in error && + (error.code === "EPIPE" || error.code === "ERR_STREAM_DESTROYED") + if (error == null || closed) resolve() + else reject(error) + } + try { + stream.write("", done) + } catch (error) { + done(error) + } }), ), ) diff --git a/packages/opencode/src/kilocode/session/drain.ts b/packages/opencode/src/kilocode/session/drain.ts index 48cb5c729fc6..fd4cbc355742 100644 --- a/packages/opencode/src/kilocode/session/drain.ts +++ b/packages/opencode/src/kilocode/session/drain.ts @@ -1,14 +1,21 @@ import { Context, Deferred, Effect, Layer } from "effect" import { LayerNode } from "@opencode-ai/core/effect/layer-node" import { InstanceState } from "@/effect/instance-state" +import { isDisposing, registerBeforeDisposer } from "@/effect/instance-registry" +import { capture, type InstanceContext } from "@/kilocode/instance" import type { SessionID } from "@/session/schema" type Entry = { + id: SessionID count: number + children: number + pins: number parent?: Entry waiters: Set> } +type State = { entries: Map; closed: boolean } + export type Delivery = { release: () => void; retained: boolean } export interface Interface { @@ -23,85 +30,128 @@ export class Service extends Context.Service()("@kilo/Sessio export const layer = Layer.effect( Service, Effect.gen(function* () { - const state = yield* InstanceState.make(() => - Effect.gen(function* () { - const entries = new Map() - const data = { entries, closed: false } - yield* Effect.addFinalizer(() => - Effect.sync(() => { - data.closed = true - for (const entry of entries.values()) { - for (const waiter of entry.waiters) Deferred.doneUnsafe(waiter, Effect.interrupt) - entry.waiters.clear() - } - entries.clear() - }), - ) - return data + const active = new Map() + const instances = new WeakMap() + let stopped = false + + function close(data: State) { + if (data.closed) return + data.closed = true + for (const entry of data.entries.values()) { + for (const waiter of entry.waiters) Deferred.doneUnsafe(waiter, Effect.interrupt) + entry.waiters.clear() + } + data.entries.clear() + } + + const off = registerBeforeDisposer((directory) => { + const data = active.get(directory) ?? { entries: new Map(), closed: false } + close(data) + active.delete(directory) + const ctx = capture() + if (ctx?.directory === directory) instances.set(ctx, data) + return undefined + }) + yield* Effect.addFinalizer(() => + Effect.sync(() => { + stopped = true + off() + for (const data of active.values()) close(data) + active.clear() }), ) - function entry(entries: Map, id: SessionID) { - const found = entries.get(id) + const state = Effect.gen(function* () { + const ctx = yield* InstanceState.context + if (stopped || isDisposing(ctx.directory)) return yield* Effect.interrupt + const known = instances.get(ctx) + if (known) return known.closed ? yield* Effect.interrupt : known + const data = active.get(ctx.directory) ?? { entries: new Map(), closed: false } + active.set(ctx.directory, data) + instances.set(ctx, data) + return data + }) + + function entry(data: State, id: SessionID) { + const found = data.entries.get(id) if (found) return found - const value: Entry = { count: 0, waiters: new Set() } - entries.set(id, value) + const value: Entry = { id, count: 0, children: 0, pins: 0, waiters: new Set() } + data.entries.set(id, value) return value } - function update(value: Entry, delta: number) { + function prune(data: State, value: Entry) { + if (value.count || value.pins || value.parent || value.children || value.waiters.size) return + if (data.entries.get(value.id) === value) data.entries.delete(value.id) + } + + function update(data: State, value: Entry, delta: number) { for (let current: Entry | undefined = value; current; current = current.parent) { current.count += delta if (current.count !== 0) continue for (const waiter of current.waiters) Deferred.doneUnsafe(waiter, Effect.void) current.waiters.clear() + prune(data, current) } } const hold = Effect.fn("SessionDrain.hold")(function* (id: SessionID) { - const data = yield* InstanceState.get(state) - if (data.closed) return yield* Effect.interrupt - const value = entry(data.entries, id) - update(value, 1) - let active = true + const data = yield* state + const value = entry(data, id) + update(data, value, 1) + let held = true return () => { - if (!active) return - active = false - if (!data.closed) update(value, -1) + if (!held) return + held = false + if (!data.closed) update(data, value, -1) } }) const link = Effect.fn("SessionDrain.link")(function* (child: SessionID, parent: SessionID) { - const data = yield* InstanceState.get(state) - if (data.closed) return yield* Effect.interrupt - const value = entry(data.entries, child) - const ancestor = entry(data.entries, parent) + const data = yield* state + const value = entry(data, child) + const ancestor = entry(data, parent) if (value.parent === ancestor) return yield* Effect.void if (value.parent) return yield* Effect.die(new Error("Session drain parent changed")) for (let current: Entry | undefined = ancestor; current; current = current.parent) { if (current === value) return yield* Effect.die(new Error("Cyclic session drain ancestry")) } value.parent = ancestor - update(ancestor, value.count) + ancestor.children++ + update(data, ancestor, value.count) return yield* Effect.void }) const wait = Effect.fn("SessionDrain.wait")(function* (id: SessionID) { - const data = yield* InstanceState.get(state) - const value = entry(data.entries, id) - while (!data.closed) { - if (value.count === 0) return yield* Effect.void - const waiter = Deferred.makeUnsafe() - value.waiters.add(waiter) - yield* Deferred.await(waiter).pipe( - Effect.ensuring( - Effect.sync(() => { - value.waiters.delete(waiter) - }), - ), - ) - } - return yield* Effect.interrupt + const data = yield* state + const value = data.entries.get(id) + if (!value) return yield* Effect.void + return yield* Effect.acquireUseRelease( + Effect.sync(() => { + value.pins++ + }), + () => + Effect.gen(function* () { + while (!data.closed) { + if (value.count === 0) return yield* Effect.void + const waiter = Deferred.makeUnsafe() + value.waiters.add(waiter) + yield* Deferred.await(waiter).pipe( + Effect.ensuring( + Effect.sync(() => { + value.waiters.delete(waiter) + }), + ), + ) + } + return yield* Effect.interrupt + }), + () => + Effect.sync(() => { + value.pins-- + if (!data.closed) prune(data, value) + }), + ) }) const track = (id: SessionID, work: Effect.Effect) => diff --git a/packages/opencode/src/tool/task.ts b/packages/opencode/src/tool/task.ts index aea3fca1e17b..491b43bdae54 100644 --- a/packages/opencode/src/tool/task.ts +++ b/packages/opencode/src/tool/task.ts @@ -337,6 +337,9 @@ export const TaskTool = Tool.define( const notify = Effect.fn("TaskTool.notifyBackgroundResult")((jobID: string) => Effect.uninterruptible( Effect.gen(function* () { + if (delivery.retained) return + const owner = yield* Scope.fork(scope, "parallel") + yield* Scope.addFinalizer(owner, Effect.sync(delivery.release)) yield* background.wait({ id: jobID }).pipe( Effect.flatMap((result) => { if (result.info?.status === "completed") return inject("completed", result.info.output ?? "") @@ -356,7 +359,7 @@ export const TaskTool = Tool.define( }).toObject(), }), ), - Effect.ensuring(Effect.sync(delivery.release)), + Effect.ensuring(Scope.close(owner, Exit.void)), Effect.forkIn(scope, { startImmediately: true }), ) delivery.retained = true @@ -411,20 +414,16 @@ export const TaskTool = Tool.define( type: id, title: params.description, metadata, - onPromote: Effect.all([ - ctx.metadata({ - title: params.description, - metadata: { ...metadata, background: true, jobId: nextSession.id }, - }), - notify(nextSession.id), - ]), + onPromote: ctx.metadata({ + title: params.description, + metadata: { ...metadata, background: true, jobId: nextSession.id }, + }), // kilocode_change - only the initial-background start needs its own cost bracket; the // foreground/promoted path below is already wrapped by the acquireUseRelease at the bottom of run() run: runInBackground ? backgroundRun : runTask().pipe(Effect.onInterrupt(() => ops.cancel(nextSession.id))), }) function backgroundResult() { - delivery.retained = true // kilocode_change return { title: params.description, metadata: { @@ -466,7 +465,10 @@ export const TaskTool = Tool.define( background.wait({ id: nextSession.id }).pipe(Effect.map((waited) => waited.info)), background.waitForPromotion(nextSession.id), ) - if (result?.metadata?.background === true) return backgroundResult() + if (result?.metadata?.background === true) { + yield* notify(info.id) + return backgroundResult() + } if (result?.status === "error") return yield* Effect.fail(new Error(result.error ?? "Task failed")) if (result?.status === "cancelled") return yield* Effect.fail(new Error("Task cancelled")) return { diff --git a/packages/opencode/test/event-manifest.test.ts b/packages/opencode/test/event-manifest.test.ts index 99403925aade..bc87fc4708d0 100644 --- a/packages/opencode/test/event-manifest.test.ts +++ b/packages/opencode/test/event-manifest.test.ts @@ -3,13 +3,15 @@ import { SessionEvent } from "@opencode-ai/core/session/event" import { EventManifest as SchemaEventManifest } from "@opencode-ai/schema/event-manifest" import { Todo } from "@/session/todo" import { EventManifest } from "@/event-manifest" +import { Drained } from "@opencode-ai/schema/kilocode/session-drain" // kilocode_change describe("public event manifest", () => { test("contains every latest public wire type once", () => { expect(EventManifest.Definitions).toBe(SchemaEventManifest.Definitions) expect(EventManifest.Latest).toBe(SchemaEventManifest.Latest) expect(EventManifest.Durable).toBe(SchemaEventManifest.Durable) - expect(EventManifest.Latest.size).toBe(90) // kilocode_change - include global.config.updated and worktree.setup.ready + expect(EventManifest.Latest.size).toBe(91) // kilocode_change - include session.drained + expect(EventManifest.Latest.get("session.drained")).toBe(Drained) // kilocode_change expect(EventManifest.Latest.get("session.next.step.ended")).toBe(SessionEvent.Step.Ended) expect(EventManifest.Latest.get("todo.updated")).toBe(Todo.Event.Updated) expect(EventManifest.Latest.has("ide.installed")).toBe(false) diff --git a/packages/opencode/test/kilocode/headless-session-drain.test.ts b/packages/opencode/test/kilocode/headless-session-drain.test.ts index 732b7dd7e47c..1d9a87fcaee3 100644 --- a/packages/opencode/test/kilocode/headless-session-drain.test.ts +++ b/packages/opencode/test/kilocode/headless-session-drain.test.ts @@ -151,6 +151,8 @@ function scenario({ home, llm, opencode }: CliFixture, mode: "foreground" | "bus }, } const script = `const r = await fetch(${JSON.stringify(`${tap.url}/tool`)}); if (!r.ok) throw new Error("tool gate failed"); console.log(await r.text())` + const scriptPath = path.join(home, "parent-gate.mjs") + yield* Effect.promise(() => Bun.write(scriptPath, script)) yield* llm.pushMatch( ({ body }) => body.model === "parent", reply().tool("task", { @@ -164,7 +166,7 @@ function scenario({ home, llm, opencode }: CliFixture, mode: "foreground" | "bus : [ mode === "busy" ? reply().tool("bash", { - command: `bun -e ${JSON.stringify(script)}`, + command: `bun ${JSON.stringify(scriptPath.replaceAll("\\", "/"))}`, description: "Wait for the test request gate", }) : reply().text("WAITING_FOR_CHILD").stop(), diff --git a/packages/opencode/test/kilocode/run-auto.test.ts b/packages/opencode/test/kilocode/run-auto.test.ts index c50dfb2f95ab..3c107d8787c9 100644 --- a/packages/opencode/test/kilocode/run-auto.test.ts +++ b/packages/opencode/test/kilocode/run-auto.test.ts @@ -115,7 +115,7 @@ const tty = Object.getOwnPropertyDescriptor(process.stdin, "isTTY") const exitCode = process.exitCode afterEach(() => { - process.exitCode = exitCode + process.exitCode = exitCode ?? 0 if (tty) { Object.defineProperty(process.stdin, "isTTY", tty) return diff --git a/packages/opencode/test/kilocode/run-drain.test.ts b/packages/opencode/test/kilocode/run-drain.test.ts index 0f8507038fd9..0f64571ef00e 100644 --- a/packages/opencode/test/kilocode/run-drain.test.ts +++ b/packages/opencode/test/kilocode/run-drain.test.ts @@ -39,6 +39,64 @@ test.each([ ) }) +test.each([401, 403])("reports HTTP %i as an authentication or authorization failure", async (status) => { + const sdk = client(async () => new Response("private error body", { status })) + await expect(KiloRunDrain.check(sdk, new AbortController().signal)).rejects.toThrow( + `Server rejected access (HTTP ${status})`, + ) +}) + +test.each([429, 500, 503])("preserves HTTP %i instead of recommending a server upgrade", async (status) => { + const sdk = client(async () => new Response("private error body", { status })) + await expect(KiloRunDrain.check(sdk, new AbortController().signal)).rejects.toThrow( + `Server capability check failed (HTTP ${status})`, + ) +}) + +test("reports invalid JSON without exposing the response body", async () => { + const sdk = client( + async () => new Response("private invalid payload", { headers: { "content-type": "application/json" } }), + ) + await expect(KiloRunDrain.check(sdk, new AbortController().signal)).rejects.toThrow( + "Server returned invalid capability JSON", + ) +}) + +test.each(["EPIPE", "ERR_STREAM_DESTROYED"])("flush tolerates closed output %s", async (code) => { + const error = Object.assign(new Error("closed output"), { code }) + await KiloRunDrain.flush([{ write: (_chunk, callback) => callback(error) }]) + await KiloRunDrain.flush([ + { + write: () => { + throw error + }, + }, + ]) +}) + +test("flush preserves other I/O errors", async () => { + const error = Object.assign(new Error("output failed"), { code: "EIO" }) + await expect(KiloRunDrain.flush([{ write: (_chunk, callback) => callback(error) }])).rejects.toBe(error) +}) + +test("flush waits for all output callbacks", async () => { + const first = Promise.withResolvers() + const second = Promise.withResolvers() + let done = false + const output = KiloRunDrain.flush([ + { write: (_chunk, callback) => first.promise.then(() => callback()) }, + { write: (_chunk, callback) => second.promise.then(() => callback()) }, + ]).then(() => { + done = true + }) + first.resolve() + await first.promise + expect(done).toBe(false) + second.resolve() + await output + expect(done).toBe(true) +}) + test("accepts an acknowledgment that arrives before the HTTP result", async () => { const response = Promise.withResolvers() const drain = KiloRunDrain.create("ses_parent") diff --git a/packages/opencode/test/kilocode/run-network.test.ts b/packages/opencode/test/kilocode/run-network.test.ts index a4f15d56d440..8da63bea09ab 100644 --- a/packages/opencode/test/kilocode/run-network.test.ts +++ b/packages/opencode/test/kilocode/run-network.test.ts @@ -1,4 +1,4 @@ -import { afterEach, describe, expect, mock, test } from "bun:test" +import { afterEach, describe, expect, mock, spyOn, test } from "bun:test" type Event = { type: string @@ -92,12 +92,12 @@ function args() { } } -const timer = globalThis.setTimeout +const exitCode = process.exitCode const tty = Object.getOwnPropertyDescriptor(process.stdin, "isTTY") const text = Bun.stdin.text afterEach(() => { - globalThis.setTimeout = timer + process.exitCode = exitCode ?? 0 Bun.stdin.text = text if (tty) { Object.defineProperty(process.stdin, "isTTY", tty) @@ -106,18 +106,35 @@ afterEach(() => { delete (process.stdin as { isTTY?: boolean }).isTTY }) -function instant() { - globalThis.setTimeout = ((cb: TimerHandler) => { - if (typeof cb === "function") { - queueMicrotask(() => cb()) - } - return 0 as unknown as ReturnType - }) as unknown as typeof setTimeout -} - -async function run(sdk: Record, overrides: Record = {}, terminal = true) { - mock.module("@kilocode/sdk/v2", () => ({ - createKiloClient: () => sdk, +async function run( + sdk: Record, + q: ReturnType>, + overrides: Record = {}, + terminal = true, +) { + sdk.event = { + subscribe: async () => { + q.push({ type: "server.connected", properties: {} }) + return { stream: q.stream() } + }, + } + sdk.kilocode = { + drainSession: async (input: { sessionID: string; token: string }) => { + q.push({ type: "session.drained", properties: input }) + q.end() + return { data: true } + }, + } + await mock.module("@kilocode/sdk/v2", () => ({ + createKiloClient: (config: { fetch?: () => Promise }) => { + config.fetch = async () => + Response.json({ + paths: { + "/kilocode/session/{sessionID}/drain": { post: { operationId: "kilocode.drainSession" } }, + }, + }) + return sdk + }, })) Object.defineProperty(process.stdin, "isTTY", { @@ -127,12 +144,21 @@ async function run(sdk: Record, overrides: Record ({ + ...create(id), + pause: async () => undefined, + })) + try { + return await RunCommand.handler({ ...args(), ...overrides } as never) + } finally { + clock.mockRestore() + } } describe("cli run network retries", () => { test("rejects after repeated offline resumes without busy", async () => { - instant() const q = feed() const calls: string[] = [] const gate = Promise.withResolvers() @@ -142,9 +168,6 @@ describe("cli run network retries", () => { config: { get: async () => ({ data: { share: "manual" } }), }, - event: { - subscribe: async () => ({ stream: q.stream() }), - }, network: { reply: async (input: { requestID: string }) => { calls.push(input.requestID) @@ -153,7 +176,6 @@ describe("cli run network retries", () => { reject: async (input: { requestID: string }) => { state.reject = input.requestID q.push(idle()) - q.end() gate.resolve() }, }, @@ -169,14 +191,13 @@ describe("cli run network retries", () => { }, } - await run(sdk) + await run(sdk, q) expect(calls).toStrictEqual(["req_1", "req_2", "req_3"]) expect(state.reject).toBe("req_4") }) test("resets retry budget only after the session is busy again", async () => { - instant() const q = feed() const calls: string[] = [] const gate = Promise.withResolvers() @@ -186,9 +207,6 @@ describe("cli run network retries", () => { config: { get: async () => ({ data: { share: "manual" } }), }, - event: { - subscribe: async () => ({ stream: q.stream() }), - }, network: { reply: async (input: { requestID: string }) => { calls.push(input.requestID) @@ -202,13 +220,11 @@ describe("cli run network retries", () => { return } q.push(idle()) - q.end() gate.resolve() }, reject: async (input: { requestID: string }) => { state.reject = input.requestID q.push(idle()) - q.end() gate.resolve() }, }, @@ -224,7 +240,7 @@ describe("cli run network retries", () => { }, } - await run(sdk) + await run(sdk, q) expect(calls).toStrictEqual(["req_1", "req_2", "req_3", "req_4"]) expect(state.reject).toBeUndefined() @@ -244,9 +260,6 @@ describe("cli run network retries", () => { config: { get: async () => ({ data: { share: "manual" } }), }, - event: { - subscribe: async () => ({ stream: q.stream() }), - }, session: { get: async (input: { sessionID: string }) => ({ data: { @@ -258,13 +271,12 @@ describe("cli run network retries", () => { summarize: async (input: unknown) => { calls.push(input) q.push(idle()) - q.end() return { data: true } }, }, } - await run(sdk, { command: "compact", message: [] }, false) + await run(sdk, q, { command: "compact", message: [] }, false) expect(calls).toEqual([ { @@ -288,9 +300,6 @@ describe("cli run network retries", () => { config: { get: async () => ({ data: { share: "manual" } }), }, - event: { - subscribe: async () => ({ stream: q.stream() }), - }, session: { create: async () => ({ data: { id: "ses_created", directory: "/tmp/project" }, @@ -301,7 +310,6 @@ describe("cli run network retries", () => { type: "session.status", properties: { sessionID: "ses_created", status: { type: "idle" } }, }) - q.end() return { data: undefined } }, summarize: async () => { @@ -310,7 +318,7 @@ describe("cli run network retries", () => { }, } - await run(sdk, { command: "compact", continue: false, session: undefined, message: ["argument"] }, false) + await run(sdk, q, { command: "compact", continue: false, session: undefined, message: ["argument"] }, false) expect(calls).toEqual([ { diff --git a/packages/opencode/test/kilocode/server/httpapi-exercise-scenarios.ts b/packages/opencode/test/kilocode/server/httpapi-exercise-scenarios.ts index 8e04027f611a..71a7bf0a710c 100644 --- a/packages/opencode/test/kilocode/server/httpapi-exercise-scenarios.ts +++ b/packages/opencode/test/kilocode/server/httpapi-exercise-scenarios.ts @@ -561,6 +561,24 @@ export const kiloScenarios: Scenario[] = [ body: { messageID: "msg_httpapi_missing" }, })) .status(400), + http.protected + .post("/kilocode/session/{sessionID}/drain", "kilocode.drainSession") + .seeded((ctx) => ctx.session({ title: "Empty drain" })) + .at((ctx) => ({ + path: route("/kilocode/session/{sessionID}/drain", { sessionID: ctx.state.id }), + headers: ctx.headers(), + body: { token: "httpapi-drain" }, + })) + .json(200, (body) => check(body === true, "an empty session should drain")), + http.protected + .post("/kilocode/session/{sessionID}/drain", "kilocode.drainSession.invalid") + .seeded((ctx) => ctx.session({ title: "Invalid drain token" })) + .at((ctx) => ({ + path: route("/kilocode/session/{sessionID}/drain", { sessionID: ctx.state.id }), + headers: ctx.headers(), + body: { token: "" }, + })) + .status(400), http.protected .get("/session/{sessionID}/model-usage", "kilocode.sessionModelUsage") .seeded((ctx) => ctx.session({ title: "Model usage" })) diff --git a/packages/opencode/test/kilocode/session-drain.test.ts b/packages/opencode/test/kilocode/session-drain.test.ts index b3f4fd33065a..41ec28d290fd 100644 --- a/packages/opencode/test/kilocode/session-drain.test.ts +++ b/packages/opencode/test/kilocode/session-drain.test.ts @@ -1,5 +1,8 @@ import { expect } from "bun:test" -import { Effect, Fiber } from "effect" +import { Deferred, Effect, Exit, Fiber, Scheduler } from "effect" +import { InstanceState } from "@/effect/instance-state" +import { InstanceRef } from "@/effect/instance-ref" +import { disposeInstance, registerDisposer } from "@/effect/instance-registry" import { SessionDrain } from "@/kilocode/session/drain" import { SessionID } from "@/session/schema" import { testEffect } from "../lib/effect" @@ -109,6 +112,138 @@ it.instance( }), ) +it.instance( + "disposal closes waiters before cancellation releases their last hold", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const ctx = yield* InstanceState.context + const id = SessionID.make("ses_dispose_drain") + const release = yield* drain.hold(id) + const cancelled = yield* Deferred.make() + const finish = yield* Deferred.make() + const off = registerDisposer(async (directory) => { + if (directory !== ctx.directory) return + release() + Deferred.doneUnsafe(cancelled, Effect.void) + await Effect.runPromise(Deferred.await(finish)) + }) + yield* Effect.addFinalizer(() => + Effect.sync(() => { + off() + Deferred.doneUnsafe(finish, Effect.void) + }), + ) + const waiting = yield* drain.wait(id).pipe(Effect.exit, Effect.forkChild({ startImmediately: true })) + const disposing = yield* Effect.promise(() => disposeInstance(ctx.directory)).pipe(Effect.forkChild) + yield* Deferred.await(cancelled) + expect(Exit.hasInterrupts(yield* Fiber.join(waiting))).toBe(true) + expect(Exit.hasInterrupts(yield* drain.wait(id).pipe(Effect.exit))).toBe(true) + yield* Deferred.succeed(finish, undefined) + yield* Fiber.join(disposing) + expect(Exit.hasInterrupts(yield* drain.wait(id).pipe(Effect.exit))).toBe(true) + yield* drain.wait(id).pipe(Effect.provideService(InstanceRef, { ...ctx })) + }), +) + +it.instance( + "disposing one directory leaves another directory's waiter intact", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const ctx = yield* InstanceState.context + const other = { ...ctx, directory: `${ctx.directory}/other` } + const id = SessionID.make("ses_isolated_drain") + const release = yield* drain.hold(id).pipe(Effect.provideService(InstanceRef, other)) + let done = false + const waiting = yield* drain.wait(id).pipe( + Effect.provideService(InstanceRef, other), + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + yield* Effect.promise(() => disposeInstance(ctx.directory)) + yield* Effect.yieldNow + expect(done).toBe(false) + release() + yield* Fiber.join(waiting) + }), +) + +it.instance( + "reacquiring a released entry cannot strand an awakened waiter on stale state", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const id = SessionID.make("ses_reacquire_drain") + const queued: Array<() => void> = [] + const scheduler = new Scheduler.MixedScheduler("async", (task) => { + queued.push(task) + return () => { + const index = queued.indexOf(task) + if (index >= 0) queued.splice(index, 1) + } + }) + let paused = false + scheduler.shouldYield = () => paused + const flush = () => { + while (queued.length) queued.shift()?.() + } + yield* Effect.addFinalizer(() => + Effect.sync(() => { + paused = false + flush() + }), + ) + const first = yield* drain.hold(id) + let done = false + const waiting = yield* drain.wait(id).pipe( + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + Effect.provideService(Scheduler.Scheduler, scheduler), + ) + flush() + paused = true + first() + const second = yield* drain.hold(id) + paused = false + flush() + expect(done).toBe(false) + second() + flush() + yield* Fiber.join(waiting) + }), +) + +it.instance( + "resuming an idle child still contributes to its parent", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const parent = SessionID.make("ses_resume_parent") + const child = SessionID.make("ses_resume_child") + yield* drain.link(child, parent) + yield* drain.track(child, Effect.void) + yield* drain.wait(parent) + const release = yield* drain.hold(child) + let done = false + const waiting = yield* drain.wait(parent).pipe( + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + expect(done).toBe(false) + release() + yield* Fiber.join(waiting) + }), +) + it.instance( "tracks failures without retaining work or poisoning later waits", Effect.gen(function* () { diff --git a/packages/opencode/test/tool/task.test.ts b/packages/opencode/test/tool/task.test.ts index 5134ec790aac..c02c1bfbdbd3 100644 --- a/packages/opencode/test/tool/task.test.ts +++ b/packages/opencode/test/tool/task.test.ts @@ -1457,6 +1457,88 @@ it.instance("returns the child's final answer after nested background delivery d }), ) +it.instance("promotion installs delivery cleanup even when its metadata hook fails", () => + Effect.gen(function* () { + const background = yield* BackgroundJob.Service + const drain = yield* SessionDrain.Service + const { chat, assistant } = yield* seed() + const child = yield* Deferred.make() + const releaseChild = yield* Deferred.make() + const metadata = yield* Deferred.make() + const releaseMetadata = yield* Deferred.make() + const notification = yield* Deferred.make() + const releaseNotification = yield* Deferred.make() + yield* Effect.addFinalizer(() => + Effect.sync(() => { + for (const gate of [releaseChild, releaseMetadata, releaseNotification]) Deferred.doneUnsafe(gate, Effect.void) + }), + ) + let notifications = 0 + const factory = yield* TaskTool + const tool = yield* factory.init() + const ops: TaskPromptOps = { + ...stubOps(), + prompt: (input) => + Effect.gen(function* () { + if (input.sessionID === chat.id) { + notifications++ + yield* Deferred.succeed(notification, undefined) + yield* Deferred.await(releaseNotification) + } else { + yield* Deferred.succeed(child, input.sessionID) + yield* Deferred.await(releaseChild) + } + return reply(input, "completed") + }), + } + const running = yield* tool + .execute( + { description: "Promote", prompt: "Do work", subagent_type: "general", background: false }, + { + sessionID: chat.id, + messageID: assistant.id, + agent: "build", + abort: new AbortController().signal, + extra: { promptOps: ops }, + messages: [], + ask: () => Effect.void, + metadata: (input) => + input.metadata?.background === true + ? Effect.gen(function* () { + yield* Deferred.succeed(metadata, undefined) + yield* Deferred.await(releaseMetadata) + return yield* Effect.die(new Error("metadata hook failed")) + }) + : Effect.void, + }, + ) + .pipe(Effect.forkChild) + const childID = yield* Deferred.await(child) + const promotion = yield* background.promote(childID).pipe(Effect.exit, Effect.forkChild) + yield* Deferred.await(metadata) + const result = yield* Fiber.join(running) + expect(result.metadata.background).toBe(true) + let drained = false + const waiting = yield* drain.wait(chat.id).pipe( + Effect.tap(() => + Effect.sync(() => { + drained = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + yield* Deferred.succeed(releaseChild, undefined) + yield* Deferred.await(notification) + expect(drained).toBe(false) + yield* Deferred.succeed(releaseNotification, undefined) + yield* Fiber.join(waiting) + yield* Deferred.succeed(releaseMetadata, undefined) + expect(Exit.isFailure(yield* Fiber.join(promotion))).toBe(true) + expect(notifications).toBe(1) + yield* drain.wait(chat.id) + }), +) + const assistantCost = Effect.fn("TaskToolTest.assistantCost")(function* (sessionID: string) { const sessions = yield* Session.Service const msgs = yield* sessions.messages({ sessionID: SessionID.make(sessionID) }) From 383d9e9613e7dced277376671f14c6f36a6da09c Mon Sep 17 00:00:00 2001 From: Josh Lambert Date: Mon, 31 Aug 2026 14:26:32 -0400 Subject: [PATCH 3/5] fix(cli): snapshot drain waiters before resuming them --- .../opencode/src/kilocode/session/drain.ts | 3 +- .../test/kilocode/session-drain.test.ts | 28 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/packages/opencode/src/kilocode/session/drain.ts b/packages/opencode/src/kilocode/session/drain.ts index fd4cbc355742..6b3ee9032edb 100644 --- a/packages/opencode/src/kilocode/session/drain.ts +++ b/packages/opencode/src/kilocode/session/drain.ts @@ -89,8 +89,9 @@ export const layer = Layer.effect( for (let current: Entry | undefined = value; current; current = current.parent) { current.count += delta if (current.count !== 0) continue - for (const waiter of current.waiters) Deferred.doneUnsafe(waiter, Effect.void) + const waiters = [...current.waiters] current.waiters.clear() + for (const waiter of waiters) Deferred.doneUnsafe(waiter, Effect.void) prune(data, current) } } diff --git a/packages/opencode/test/kilocode/session-drain.test.ts b/packages/opencode/test/kilocode/session-drain.test.ts index 41ec28d290fd..6b658d9f6131 100644 --- a/packages/opencode/test/kilocode/session-drain.test.ts +++ b/packages/opencode/test/kilocode/session-drain.test.ts @@ -219,6 +219,34 @@ it.instance( }), ) +it.instance( + "a resumed waiter can acquire work without waking newly registered waiters repeatedly", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const id = SessionID.make("ses_reentrant_drain") + const release = yield* drain.hold(id) + const first = yield* drain.wait(id).pipe( + Effect.flatMap(() => drain.hold(id)), + Effect.forkChild({ startImmediately: true }), + ) + let done = false + const second = yield* drain.wait(id).pipe( + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + release() + const releaseAgain = yield* Fiber.join(first) + expect(done).toBe(false) + releaseAgain() + yield* Fiber.join(second) + expect(done).toBe(true) + }), +) + it.instance( "resuming an idle child still contributes to its parent", Effect.gen(function* () { From ac54cca97b7ad292630dfa895fda181e0999af11 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Tue, 1 Sep 2026 10:36:49 +0200 Subject: [PATCH 4/5] fix(cli): scope and simplify headless session draining Route drain requests and CLI output through the session owner directory. Simplify notification holds and completion waiters while preserving disposal safety. Cover local, attached, and daemon cross-directory runs. --- .changeset/headless-session-drain.md | 2 +- packages/opencode/src/cli/cmd/run.ts | 89 ++++++------------ .../opencode/src/kilocode/cli/run-drain.ts | 11 ++- .../opencode/src/kilocode/session/drain.ts | 29 ++---- .../src/server/shared/workspace-routing.ts | 1 + packages/opencode/src/tool/task.ts | 21 ++--- .../kilocode/headless-session-drain.test.ts | 94 +++++++++++++------ .../opencode/test/kilocode/run-drain.test.ts | 15 +++ .../test/kilocode/session-drain.test.ts | 40 ++++---- packages/opencode/test/tool/task.test.ts | 23 ++++- 10 files changed, 177 insertions(+), 148 deletions(-) diff --git a/.changeset/headless-session-drain.md b/.changeset/headless-session-drain.md index 52a52d606c00..0fcfe763690d 100644 --- a/.changeset/headless-session-drain.md +++ b/.changeset/headless-session-drain.md @@ -3,4 +3,4 @@ "@kilocode/sdk": patch --- -Wait for background subagents and their parent continuations before completing headless runs. +Wait for background subagents and their parent continuations before completing headless runs, including sessions resumed from another directory. diff --git a/packages/opencode/src/cli/cmd/run.ts b/packages/opencode/src/cli/cmd/run.ts index 26fc55086377..f2fde37a6bf1 100644 --- a/packages/opencode/src/cli/cmd/run.ts +++ b/packages/opencode/src/cli/cmd/run.ts @@ -1028,8 +1028,8 @@ export const RunCommand = effectCmd({ // kilocode_change end return error } - const cwd = args.attach ? (directory ?? sess.directory ?? (await current(sdk))) : (directory ?? root) - const client = args.attach ? attachSDK(cwd) : sdk + const cwd = sess.directory ?? directory ?? (await current(sdk)) // kilocode_change + const client = KiloRunDrain.scope(sdk, cwd) // kilocode_change // kilocode_change start - classify deferred attach commands in the session directory const builtin = deferred ? await KiloRun.resolveBuiltin(client, args.command, cwd) : initial if (deferred) { @@ -1060,74 +1060,41 @@ export const RunCommand = effectCmd({ return undefined }, ) - async function finish() { - await drain.wait(client, cwd) - const error = await completed - if (error) process.exitCode = 1 - } - try { await drain.race(KiloRunDrain.check(client, drain.signal)) await drain.ready() - if (builtin) { - const result = await drain.race( - KiloRun.runBuiltin(client, sessionID, builtin, args.model, sess.model, cwd, drain.signal), - ) - if (result.error) { - if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error)) - process.exitCode = 1 - await finish() - return - } - await finish() - return - } - - if (args.command) { - const result = await drain.race( - client.session.command( - { - sessionID, - agent, - model: args.model, - command: args.command, - arguments: message, - variant: args.variant, - }, - { signal: drain.signal }, - ), - ) - if (result.error) { - if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error)) - process.exitCode = 1 - await finish() - return - } - await finish() - return - } - - const model = pick(args.model) const result = await drain.race( - client.session.prompt( - { - sessionID, - agent, - model, - variant: args.variant, - parts: [...files, { type: "text", text: message }], - }, - { signal: drain.signal }, - ), + builtin + ? KiloRun.runBuiltin(client, sessionID, builtin, args.model, sess.model, cwd, drain.signal) + : args.command + ? client.session.command( + { + sessionID, + agent, + model: args.model, + command: args.command, + arguments: message, + variant: args.variant, + }, + { signal: drain.signal }, + ) + : client.session.prompt( + { + sessionID, + agent, + model: pick(args.model), + variant: args.variant, + parts: [...files, { type: "text", text: message }], + }, + { signal: drain.signal }, + ), ) if (result.error) { if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error)) process.exitCode = 1 - await finish() - return } - await finish() - return + await drain.wait(client, cwd) + if (await completed) process.exitCode = 1 } catch (error) { const text = error instanceof Error ? error.message : String(error) if (!emit("error", { error: text })) UI.error(text) diff --git a/packages/opencode/src/kilocode/cli/run-drain.ts b/packages/opencode/src/kilocode/cli/run-drain.ts index ac900ff12f33..d1236d8ea134 100644 --- a/packages/opencode/src/kilocode/cli/run-drain.ts +++ b/packages/opencode/src/kilocode/cli/run-drain.ts @@ -18,6 +18,12 @@ export namespace KiloRunDrain { return sdk } + export function scope(sdk: KiloClient, directory: string) { + const config = connections.get(sdk) + if (!config) throw new Error("Missing headless server transport") + return client({ ...config, directory }) + } + export async function check(sdk: KiloClient, signal: AbortSignal) { const config = connections.get(sdk) if (!config?.baseUrl) throw new Error("Missing headless server transport") @@ -56,7 +62,7 @@ export namespace KiloRunDrain { const failure = Promise.withResolvers() let closing = false let drained = false - const race = (work: Promise) => + const race = >(work: A) => Promise.race([ work, failure.promise.then((error): never => { @@ -94,8 +100,7 @@ export namespace KiloRunDrain { }, async wait(sdk: KiloClient, directory?: string) { const result = await race(sdk.kilocode.drainSession({ sessionID, directory, token }, { signal: abort.signal })) - if (result.error || !z.literal(true).safeParse(result.data).success) - throw new Error("Server did not acknowledge session completion") + if (result.error || result.data !== true) throw new Error("Server did not acknowledge session completion") await race(acknowledged.promise) }, close() { diff --git a/packages/opencode/src/kilocode/session/drain.ts b/packages/opencode/src/kilocode/session/drain.ts index 6b3ee9032edb..3c42f338dac7 100644 --- a/packages/opencode/src/kilocode/session/drain.ts +++ b/packages/opencode/src/kilocode/session/drain.ts @@ -11,13 +11,11 @@ type Entry = { children: number pins: number parent?: Entry - waiters: Set> + done: Deferred.Deferred } type State = { entries: Map; closed: boolean } -export type Delivery = { release: () => void; retained: boolean } - export interface Interface { readonly hold: (id: SessionID) => Effect.Effect<() => void> readonly link: (child: SessionID, parent: SessionID) => Effect.Effect @@ -37,10 +35,7 @@ export const layer = Layer.effect( function close(data: State) { if (data.closed) return data.closed = true - for (const entry of data.entries.values()) { - for (const waiter of entry.waiters) Deferred.doneUnsafe(waiter, Effect.interrupt) - entry.waiters.clear() - } + for (const entry of data.entries.values()) Deferred.doneUnsafe(entry.done, Effect.void) data.entries.clear() } @@ -75,13 +70,13 @@ export const layer = Layer.effect( function entry(data: State, id: SessionID) { const found = data.entries.get(id) if (found) return found - const value: Entry = { id, count: 0, children: 0, pins: 0, waiters: new Set() } + const value: Entry = { id, count: 0, children: 0, pins: 0, done: Deferred.makeUnsafe() } data.entries.set(id, value) return value } function prune(data: State, value: Entry) { - if (value.count || value.pins || value.parent || value.children || value.waiters.size) return + if (value.count || value.pins || value.parent || value.children) return if (data.entries.get(value.id) === value) data.entries.delete(value.id) } @@ -89,9 +84,9 @@ export const layer = Layer.effect( for (let current: Entry | undefined = value; current; current = current.parent) { current.count += delta if (current.count !== 0) continue - const waiters = [...current.waiters] - current.waiters.clear() - for (const waiter of waiters) Deferred.doneUnsafe(waiter, Effect.void) + const done = current.done + current.done = Deferred.makeUnsafe() + Deferred.doneUnsafe(done, Effect.void) prune(data, current) } } @@ -135,15 +130,7 @@ export const layer = Layer.effect( Effect.gen(function* () { while (!data.closed) { if (value.count === 0) return yield* Effect.void - const waiter = Deferred.makeUnsafe() - value.waiters.add(waiter) - yield* Deferred.await(waiter).pipe( - Effect.ensuring( - Effect.sync(() => { - value.waiters.delete(waiter) - }), - ), - ) + yield* Deferred.await(value.done) } return yield* Effect.interrupt }), diff --git a/packages/opencode/src/server/shared/workspace-routing.ts b/packages/opencode/src/server/shared/workspace-routing.ts index a18b69079240..0c041ed6a4c1 100644 --- a/packages/opencode/src/server/shared/workspace-routing.ts +++ b/packages/opencode/src/server/shared/workspace-routing.ts @@ -22,6 +22,7 @@ export function getWorkspaceRouteSessionID(url: URL) { if (url.pathname === "/session/viewed") return null // kilocode_change - Kilo static route is not a session ID const id = + url.pathname.match(/^\/kilocode\/session\/([^/]+)\/drain$/)?.[1] ?? // kilocode_change url.pathname.match(/^\/session\/([^/]+)(?:\/|$)/)?.[1] ?? url.pathname.match(/^\/experimental\/session\/([^/]+)\/background$/)?.[1] if (!id) return null diff --git a/packages/opencode/src/tool/task.ts b/packages/opencode/src/tool/task.ts index 491b43bdae54..72135a44ec02 100644 --- a/packages/opencode/src/tool/task.ts +++ b/packages/opencode/src/tool/task.ts @@ -110,7 +110,6 @@ export const TaskTool = Tool.define( const run = Effect.fn("TaskTool.execute")(function* ( params: Schema.Schema.Type, ctx: Tool.Context, - delivery: SessionDrain.Delivery, // kilocode_change ) { const cfg = yield* config.get() const selection = cfg.experimental?.task_model_selection === true // kilocode_change @@ -337,9 +336,9 @@ export const TaskTool = Tool.define( const notify = Effect.fn("TaskTool.notifyBackgroundResult")((jobID: string) => Effect.uninterruptible( Effect.gen(function* () { - if (delivery.retained) return const owner = yield* Scope.fork(scope, "parallel") - yield* Scope.addFinalizer(owner, Effect.sync(delivery.release)) + const release = yield* drain.hold(ctx.sessionID) + yield* Scope.addFinalizer(owner, Effect.sync(release)) yield* background.wait({ id: jobID }).pipe( Effect.flatMap((result) => { if (result.info?.status === "completed") return inject("completed", result.info.output ?? "") @@ -362,7 +361,6 @@ export const TaskTool = Tool.define( Effect.ensuring(Scope.close(owner, Exit.void)), Effect.forkIn(scope, { startImmediately: true }), ) - delivery.retained = true }), ), ) @@ -414,10 +412,12 @@ export const TaskTool = Tool.define( type: id, title: params.description, metadata, + // kilocode_change start onPromote: ctx.metadata({ title: params.description, metadata: { ...metadata, background: true, jobId: nextSession.id }, }), + // kilocode_change end // kilocode_change - only the initial-background start needs its own cost bracket; the // foreground/promoted path below is already wrapped by the acquireUseRelease at the bottom of run() run: runInBackground ? backgroundRun : runTask().pipe(Effect.onInterrupt(() => ops.cancel(nextSession.id))), @@ -465,10 +465,12 @@ export const TaskTool = Tool.define( background.wait({ id: nextSession.id }).pipe(Effect.map((waited) => waited.info)), background.waitForPromotion(nextSession.id), ) + // kilocode_change start if (result?.metadata?.background === true) { yield* notify(info.id) return backgroundResult() } + // kilocode_change end if (result?.status === "error") return yield* Effect.fail(new Error(result.error ?? "Task failed")) if (result?.status === "cancelled") return yield* Effect.fail(new Error("Task cancelled")) return { @@ -525,16 +527,7 @@ export const TaskTool = Tool.define( }), ), execute: (params: Schema.Schema.Type, ctx: Tool.Context) => - Effect.acquireUseRelease( - drain - .hold(ctx.sessionID) - .pipe(Effect.map((release): SessionDrain.Delivery => ({ release, retained: false }))), - (delivery) => run(params, ctx, delivery), - (delivery) => - Effect.sync(() => { - if (!delivery.retained) delivery.release() - }), - ).pipe(Effect.orDie), + drain.track(ctx.sessionID, run(params, ctx)).pipe(Effect.orDie), } }) // kilocode_change end diff --git a/packages/opencode/test/kilocode/headless-session-drain.test.ts b/packages/opencode/test/kilocode/headless-session-drain.test.ts index 1d9a87fcaee3..db9355d059b8 100644 --- a/packages/opencode/test/kilocode/headless-session-drain.test.ts +++ b/packages/opencode/test/kilocode/headless-session-drain.test.ts @@ -2,6 +2,7 @@ import { expect } from "bun:test" import { Effect, Fiber } from "effect" import { createKiloClient, type Event } from "@kilocode/sdk/v2" import path from "node:path" +import { mkdir } from "node:fs/promises" import { pathToFileURL } from "node:url" import { cliIt, type CliFixture } from "../lib/cli-process" import { awaitWithTimeout } from "../lib/effect" @@ -26,7 +27,8 @@ function observer(url: string, busy: boolean) { const sdk = pathToFileURL(path.resolve(import.meta.dir, "../../../sdk/js/src/v2/index.ts")).href return `import { createKiloClient } from ${JSON.stringify(sdk)} export default async ({ client: legacy }) => { - const client = createKiloClient(legacy._client.getConfig()) + const cfg = legacy._client.getConfig() + const client = createKiloClient({ ...cfg, headers: Object.fromEntries(new Headers(cfg.headers)) }) const abort = new AbortController() const ready = Promise.withResolvers() let root @@ -117,8 +119,14 @@ function probe(home: string, busy: boolean) { }) } -function scenario({ home, llm, opencode }: CliFixture, mode: "foreground" | "busy" | "idle", attached = false) { +function scenario( + { home, llm, opencode }: CliFixture, + mode: "foreground" | "busy" | "idle", + opts: { attach?: boolean; resume?: boolean; daemon?: boolean } = {}, +) { return Effect.gen(function* () { + const directory = opts.resume ? path.join(home, "session") : home + if (opts.resume) yield* Effect.promise(() => mkdir(directory)) const tap = yield* probe(home, mode === "busy") const child = gate() const requested = gate() @@ -184,7 +192,7 @@ function scenario({ home, llm, opencode }: CliFixture, mode: "foreground" | "bus KILO_CONFIG_CONTENT: JSON.stringify(cfg), KILO_CONFIG: "", KILO_CONFIG_DIR: "", - KILO_DB: ":memory:", + KILO_DB: opts.resume ? path.join(home, "resume.db") : ":memory:", KILO_AUTH_CONTENT: "{}", KILO_SERVER_PASSWORD: "", KILO_SERVER_USERNAME: "kilo", @@ -193,15 +201,45 @@ function scenario({ home, llm, opencode }: CliFixture, mode: "foreground" | "bus KILO_EXPERIMENTAL_BACKGROUND_SUBAGENTS: "true", KILO_TELEMETRY_LEVEL: "off", KILO_AUTO_SHARE: "false", + KILO_NO_DAEMON: opts.daemon ? "" : "1", + KILO_PARENT_PID: "", + KILO_TEST_DAEMON_STATE_DIR: path.join(home, "daemon-state"), + KILO_TEST_DAEMON_LOG_DIR: path.join(home, "daemon-log"), + } + const server = opts.attach || opts.resume ? yield* opencode.serve({ env, readyTimeoutMs: 30_000 }) : undefined + const session = + opts.resume && server + ? yield* Effect.promise(() => + createKiloClient({ baseUrl: server.url, directory }).session.create({}, { throwOnError: true }), + ) + : undefined + if (opts.daemon) { + yield* Effect.addFinalizer(() => + opencode + .spawn(["daemon", "stop", "--json"], { env }) + .pipe(Effect.tap((result) => Effect.sync(() => opencode.expectExit(result, 0)))), + ) + const result = yield* opencode.spawn(["daemon", "start", "--hostname", "127.0.0.1", "--port", "0", "--json"], { + env, + }) + opencode.expectExit(result, 0) + expect(JSON.parse(result.stdout).running).toBe(true) } - const server = attached ? yield* opencode.serve({ env, readyTimeoutMs: 30_000 }) : undefined const run = yield* opencode.startRun("Exercise background completion.", { model: "test/parent", agent: "code", format: "json", printLogs: true, - extraArgs: ["--auto", "--dir", home, "--title", "drain regression", ...(server ? ["--attach", server.url] : [])], - env, + extraArgs: [ + "--auto", + "--dir", + home, + "--title", + "drain regression", + ...(opts.attach && server ? ["--attach", server.url] : []), + ...(session ? ["--session", session.data.id] : []), + ], + env: opts.daemon ? { ...env, KILO_DB: path.join(home, "unused.db") } : env, }) yield* Effect.raceFirst( requested.wait("child never reached the mock LLM"), @@ -238,32 +276,25 @@ function scenario({ home, llm, opencode }: CliFixture, mode: "foreground" | "bus }) } -cliIt.live("foreground Task completes before headless exit", (fixture) => scenario(fixture, "foreground"), 90_000) -cliIt.live( - "headless run drains a child callback queued behind a busy parent", - (fixture) => scenario(fixture, "busy"), - 90_000, -) -cliIt.live( - "headless run waits for a child after the parent becomes idle", - (fixture) => scenario(fixture, "idle"), - 90_000, -) -cliIt.live( - "attached run drains a callback queued behind the parent", - (fixture) => scenario(fixture, "busy", true), - 90_000, -) -cliIt.live( - "attached run waits for a child after the parent becomes idle", - (fixture) => scenario(fixture, "idle", true), - 90_000, -) +for (const [name, mode, opts] of [ + ["foreground Task completes before headless exit", "foreground", {}], + ["headless run drains a child callback queued behind a busy parent", "busy", {}], + ["headless run waits for a child after the parent becomes idle", "idle", {}], + ["attached run drains a callback queued behind the parent", "busy", { attach: true }], + ["attached run waits for a child after the parent becomes idle", "idle", { attach: true }], + ["local run resumes and drains a session from another directory", "busy", { resume: true }], + ["attached run resumes and drains a session from another directory", "idle", { attach: true, resume: true }], + ["daemon run resumes and drains a session from another directory", "idle", { daemon: true, resume: true }], +] as const) { + cliIt.live(name, (fixture) => scenario(fixture, mode, opts), 90_000) +} cliIt.live( - "SDK drain validates requests and publishes its matching acknowledgment", + "SDK drain follows the session directory and validates acknowledgments", ({ home, opencode }) => Effect.gen(function* () { + const directory = path.join(home, "caller") + yield* Effect.promise(() => mkdir(directory)) const server = yield* opencode.serve({ env: { KILO_SERVER_PASSWORD: "", KILO_SERVER_USERNAME: "kilo" }, readyTimeoutMs: 30_000, @@ -287,9 +318,12 @@ cliIt.live( throw new Error("Drain acknowledgment was not delivered") }).pipe(Effect.forkChild) yield* ready.wait("Event stream did not connect") - const result = yield* Effect.promise(() => sdk.kilocode.drainSession({ sessionID: id, token })) + const result = yield* Effect.promise(() => sdk.kilocode.drainSession({ sessionID: id, directory, token })) expect(result.data).toBe(true) - expect(yield* Fiber.join(received)).toEqual({ sessionID: id, token }) + expect(yield* awaitWithTimeout(Fiber.join(received), "Drain acknowledged the wrong directory")).toEqual({ + sessionID: id, + token, + }) const invalid = yield* Effect.promise(() => sdk.kilocode.drainSession({ sessionID: id, token: "" })) expect(invalid.response?.status).toBe(400) const missing = yield* Effect.promise(() => sdk.kilocode.drainSession({ sessionID: "ses_missing_drain", token })) diff --git a/packages/opencode/test/kilocode/run-drain.test.ts b/packages/opencode/test/kilocode/run-drain.test.ts index 0f64571ef00e..ad6570804f4e 100644 --- a/packages/opencode/test/kilocode/run-drain.test.ts +++ b/packages/opencode/test/kilocode/run-drain.test.ts @@ -28,6 +28,21 @@ test("checks capabilities through the selected transport before accepting comple await KiloRunDrain.check(sdk, new AbortController().signal) }) +test("scopes requests to the session directory without changing the transport", async () => { + const paths: Array = [] + const sdk = client(async (request) => { + const url = new URL(request.url) + expect(url.origin).toBe("http://drain.test") + expect(url.pathname).toBe("/prefix/config") + expect(request.headers.get("authorization")).toBe("Basic test-only") + paths.push(url.searchParams.get("directory")) + return Response.json({}) + }) + await KiloRunDrain.scope(sdk, "/session owner").config.get() + await sdk.config.get() + expect(paths).toEqual(["/session owner", null]) +}) + test.each([ () => new Response("old server", { headers: { "content-type": "text/html" } }), () => Response.json({ paths: {} }), diff --git a/packages/opencode/test/kilocode/session-drain.test.ts b/packages/opencode/test/kilocode/session-drain.test.ts index 6b658d9f6131..76df7acc4e81 100644 --- a/packages/opencode/test/kilocode/session-drain.test.ts +++ b/packages/opencode/test/kilocode/session-drain.test.ts @@ -90,25 +90,29 @@ it.instance( ) it.instance( - "interrupting a wait does not cancel tracked work", + "interrupting a wait does not cancel other waiters or later generations", Effect.gen(function* () { const drain = yield* SessionDrain.Service - const id = SessionID.make("ses_drain_parent") - const release = yield* drain.hold(id) - const cancelled = yield* drain.wait(id).pipe(Effect.forkChild({ startImmediately: true })) - yield* Fiber.interrupt(cancelled) - let done = false - const waiter = yield* drain.wait(id).pipe( - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) - expect(done).toBe(false) - release() - yield* Fiber.join(waiter) + const id = SessionID.make("ses_drain_child") + yield* drain.link(id, SessionID.make("ses_drain_parent")) + for (let round = 0; round < 2; round++) { + const release = yield* drain.hold(id) + const cancelled = yield* drain.wait(id).pipe(Effect.forkChild({ startImmediately: true })) + let done = false + const waiter = yield* drain.wait(id).pipe( + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + yield* Fiber.interrupt(cancelled) + expect(done).toBe(false) + release() + yield* Fiber.join(waiter) + expect(done).toBe(true) + } }), ) @@ -134,9 +138,11 @@ it.instance( }), ) const waiting = yield* drain.wait(id).pipe(Effect.exit, Effect.forkChild({ startImmediately: true })) + const other = yield* drain.wait(id).pipe(Effect.exit, Effect.forkChild({ startImmediately: true })) const disposing = yield* Effect.promise(() => disposeInstance(ctx.directory)).pipe(Effect.forkChild) yield* Deferred.await(cancelled) expect(Exit.hasInterrupts(yield* Fiber.join(waiting))).toBe(true) + expect(Exit.hasInterrupts(yield* Fiber.join(other).pipe(Effect.timeout("2 seconds")))).toBe(true) expect(Exit.hasInterrupts(yield* drain.wait(id).pipe(Effect.exit))).toBe(true) yield* Deferred.succeed(finish, undefined) yield* Fiber.join(disposing) diff --git a/packages/opencode/test/tool/task.test.ts b/packages/opencode/test/tool/task.test.ts index c02c1bfbdbd3..ca2cb88fa535 100644 --- a/packages/opencode/test/tool/task.test.ts +++ b/packages/opencode/test/tool/task.test.ts @@ -1200,6 +1200,8 @@ describe("tool.task", () => { background.instance("background task completion does not wait for the parent async prompt", () => Effect.gen(function* () { const jobs = yield* BackgroundJob.Service + const drain = yield* SessionDrain.Service // kilocode_change + const finish = yield* Deferred.make() // kilocode_change const { chat, assistant } = yield* seed() const tool = yield* TaskTool const def = yield* tool.init() @@ -1219,8 +1221,12 @@ describe("tool.task", () => { extra: { promptOps: { ...stubOps({ text: "background done" }), + // kilocode_change start prompt: (input) => - input.sessionID === chat.id ? Effect.never : Effect.succeed(reply(input, "background done")), + input.sessionID === chat.id + ? Deferred.await(finish).pipe(Effect.as(reply(input, "parent done"))) + : Effect.succeed(reply(input, "background done")), + // kilocode_change end } satisfies TaskPromptOps, }, messages: [], @@ -1232,6 +1238,21 @@ describe("tool.task", () => { const waited = yield* jobs.wait({ id: result.metadata.sessionId, timeout: 1_000 }) expect(waited.timedOut).toBe(false) expect(waited.info?.status).toBe("completed") + // kilocode_change start + let done = false + const waiting = yield* drain.wait(chat.id).pipe( + Effect.tap(() => + Effect.sync(() => { + done = true + }), + ), + Effect.forkChild({ startImmediately: true }), + ) + expect(done).toBe(false) + yield* Deferred.succeed(finish, undefined) + yield* Fiber.join(waiting) + expect(done).toBe(true) + // kilocode_change end }), ) From 0a9ce2f0ef5c890f2dd446d5616df6df4708c105 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Tue, 1 Sep 2026 14:17:16 +0200 Subject: [PATCH 5/5] fix(cli): close remaining headless drain races Keep reservations through prompt intake, runner completion, and shell cleanup. Report cancellation after the parent turn has ended without emitting another turn close. Keep disposal bookkeeping in Kilo-owned code and consolidate regression fixtures without dropping the required coverage. --- .changeset/headless-session-drain.md | 2 +- packages/opencode/src/cli/cmd/run.ts | 103 +++----- .../opencode/src/effect/instance-registry.ts | 36 +-- packages/opencode/src/effect/runner.ts | 87 +++--- packages/opencode/src/kilocode/cli/cmd/run.ts | 16 +- .../opencode/src/kilocode/cli/run-drain.ts | 38 ++- .../src/kilocode/effect/instance-registry.ts | 36 +++ .../opencode/src/kilocode/effect/runner.ts | 41 ++- .../opencode/src/kilocode/session/drain.ts | 2 +- .../opencode/src/kilocode/session/prompt.ts | 65 +++-- packages/opencode/src/session/prompt.ts | 38 ++- packages/opencode/src/session/run-state.ts | 1 + packages/opencode/src/tool/task.ts | 3 +- packages/opencode/test/event-manifest.test.ts | 5 +- .../kilocode/headless-session-drain.test.ts | 248 ++++++++++++++---- .../opencode/test/kilocode/run-auto.test.ts | 54 ++-- .../opencode/test/kilocode/run-drain.test.ts | 46 +++- .../test/kilocode/run-network.test.ts | 11 +- .../test/kilocode/runner-start-order.test.ts | 131 ++++++++- .../test/kilocode/session-drain.test.ts | 122 +++------ .../test/kilocode/task-nesting.test.ts | 125 ++++++++- packages/opencode/test/tool/task.test.ts | 158 +---------- packages/schema/src/kilocode/session-drain.ts | 7 +- packages/sdk/js/src/v2/gen/types.gen.ts | 44 ++++ packages/sdk/openapi.json | 137 ++++++++++ 25 files changed, 1001 insertions(+), 555 deletions(-) create mode 100644 packages/opencode/src/kilocode/effect/instance-registry.ts diff --git a/.changeset/headless-session-drain.md b/.changeset/headless-session-drain.md index 0fcfe763690d..6643f7495946 100644 --- a/.changeset/headless-session-drain.md +++ b/.changeset/headless-session-drain.md @@ -3,4 +3,4 @@ "@kilocode/sdk": patch --- -Wait for background subagents and their parent continuations before completing headless runs, including sessions resumed from another directory. +Wait for background subagents and their parent continuations before completing headless runs, including sessions resumed from another directory. Report cancelled headless runs as failures. diff --git a/packages/opencode/src/cli/cmd/run.ts b/packages/opencode/src/cli/cmd/run.ts index f2fde37a6bf1..0bdac60ad2b1 100644 --- a/packages/opencode/src/cli/cmd/run.ts +++ b/packages/opencode/src/cli/cmd/run.ts @@ -910,7 +910,7 @@ export const RunCommand = effectCmd({ if (event.type === "suggestion.shown") { const suggestion = event.properties if (suggestion.sessionID === sessionID || KiloRunAuto.allowed(tracked, suggestion.sessionID)) { - await client.suggestion.dismiss({ requestID: suggestion.id }, { signal: drain.signal }).catch(() => {}) + await client.suggestion.dismiss({ requestID: suggestion.id }).catch(() => {}) } continue } @@ -922,20 +922,17 @@ export const RunCommand = effectCmd({ // kilocode_change start - skill shell batches need an interactive human decision. The server ignores // non-interactive approvals, so headless runs must reject explicitly rather than leave them pending. if (permission.metadata?.["skillShell"] === true || permission.metadata?.["sandboxEscalation"] === true) { - await client.permission.reply({ requestID: permission.id, reply: "reject" }, { signal: drain.signal }) + await client.permission.reply({ requestID: permission.id, reply: "reject" }) continue } // kilocode_change end // kilocode_change start - approve root and tracked Task child permissions in auto mode if (args.auto) { if (!KiloRunAuto.allowed(tracked, permission.sessionID)) continue - await client.permission.reply( - { - requestID: permission.id, - reply: "once", - }, - { signal: drain.signal }, - ) + await client.permission.reply({ + requestID: permission.id, + reply: "once", + }) continue } // kilocode_change end @@ -946,13 +943,10 @@ export const RunCommand = effectCmd({ if (permission.sessionID !== sessionID) { if (!KiloRunAuto.allowed(tracked, permission.sessionID)) continue if (skipPermissions) { - await client.permission.reply( - { - requestID: permission.id, - reply: "once", - }, - { signal: drain.signal }, - ) + await client.permission.reply({ + requestID: permission.id, + reply: "once", + }) continue } UI.println( @@ -961,13 +955,10 @@ export const RunCommand = effectCmd({ `subagent permission requested: ${permission.permission} (${permission.patterns.join(", ")}); auto-rejecting`, ) autoRejected = true // kilocode_change - await client.permission.reply( - { - requestID: permission.id, - reply: "reject", - }, - { signal: drain.signal }, - ) + await client.permission.reply({ + requestID: permission.id, + reply: "reject", + }) continue } // kilocode_change end @@ -975,13 +966,10 @@ export const RunCommand = effectCmd({ if (permission.sessionID !== sessionID) continue if (skipPermissions) { - await client.permission.reply( - { - requestID: permission.id, - reply: "once", - }, - { signal: drain.signal }, - ) + await client.permission.reply({ + requestID: permission.id, + reply: "once", + }) } else { UI.println( UI.Style.TEXT_WARNING_BOLD + "!", @@ -989,13 +977,10 @@ export const RunCommand = effectCmd({ `permission requested: ${permission.permission} (${permission.patterns.join(", ")}); auto-rejecting`, ) autoRejected = true // kilocode_change - await client.permission.reply( - { - requestID: permission.id, - reply: "reject", - }, - { signal: drain.signal }, - ) + await client.permission.reply({ + requestID: permission.id, + reply: "reject", + }) } } @@ -1009,12 +994,12 @@ export const RunCommand = effectCmd({ UI.Style.TEXT_WARNING_BOLD + "!", UI.Style.TEXT_NORMAL + `network retry limit reached (${MAX_RETRIES}); rejecting`, ) - await client.network.reject({ requestID: request.id }, { signal: drain.signal }) + await client.network.reject({ requestID: request.id }) continue } const delay = Math.min(5000 * Math.pow(2, retries - 1), 60000) await drain.pause(delay) - await client.network.reply({ requestID: request.id }, { signal: drain.signal }) + await client.network.reply({ requestID: request.id }) } // kilocode_change end } @@ -1029,7 +1014,7 @@ export const RunCommand = effectCmd({ return error } const cwd = sess.directory ?? directory ?? (await current(sdk)) // kilocode_change - const client = KiloRunDrain.scope(sdk, cwd) // kilocode_change + const client = KiloRunDrain.scope(sdk, cwd, interactive ? undefined : drain.signal) // kilocode_change // kilocode_change start - classify deferred attach commands in the session directory const builtin = deferred ? await KiloRun.resolveBuiltin(client, args.command, cwd) : initial if (deferred) { @@ -1065,29 +1050,23 @@ export const RunCommand = effectCmd({ await drain.ready() const result = await drain.race( builtin - ? KiloRun.runBuiltin(client, sessionID, builtin, args.model, sess.model, cwd, drain.signal) + ? KiloRun.runBuiltin(client, sessionID, builtin, args.model, sess.model, cwd) : args.command - ? client.session.command( - { - sessionID, - agent, - model: args.model, - command: args.command, - arguments: message, - variant: args.variant, - }, - { signal: drain.signal }, - ) - : client.session.prompt( - { - sessionID, - agent, - model: pick(args.model), - variant: args.variant, - parts: [...files, { type: "text", text: message }], - }, - { signal: drain.signal }, - ), + ? client.session.command({ + sessionID, + agent, + model: args.model, + command: args.command, + arguments: message, + variant: args.variant, + }) + : client.session.prompt({ + sessionID, + agent, + model: pick(args.model), + variant: args.variant, + parts: [...files, { type: "text", text: message }], + }), ) if (result.error) { if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error)) diff --git a/packages/opencode/src/effect/instance-registry.ts b/packages/opencode/src/effect/instance-registry.ts index 925555a3854b..f05e2dc932cd 100644 --- a/packages/opencode/src/effect/instance-registry.ts +++ b/packages/opencode/src/effect/instance-registry.ts @@ -1,22 +1,9 @@ import type { WorkspaceV2 } from "@opencode-ai/core/workspace" // kilocode_change +import { dispose } from "@/kilocode/effect/instance-registry" // kilocode_change const disposers = new Set<(directory: string, workspaceID?: WorkspaceV2.ID) => Promise>() // kilocode_change // kilocode_change start -const beforeDisposers = new Set<(directory: string, workspaceID?: WorkspaceV2.ID) => undefined>() -const closing = new Map() - -export function registerBeforeDisposer(disposer: (directory: string, workspaceID?: WorkspaceV2.ID) => undefined) { - beforeDisposers.add(disposer) - return () => { - beforeDisposers.delete(disposer) - } -} - -export function isDisposing(directory: string) { - return closing.has(directory) -} - export function registerDisposer( disposer: (directory: string, workspaceID?: WorkspaceV2.ID) => Promise, // kilocode_change ) { @@ -27,23 +14,8 @@ export function registerDisposer( } export async function disposeInstance(directory: string, workspaceID?: WorkspaceV2.ID) { - closing.set(directory, (closing.get(directory) ?? 0) + 1) - try { - const errors: unknown[] = [] - const guards = [...beforeDisposers] - for (const disposer of guards) { - try { - disposer(directory, workspaceID) - } catch (error) { - errors.push(error) - } - } - await Promise.allSettled([...disposers].map((disposer) => disposer(directory, workspaceID))) - if (errors.length) throw new AggregateError(errors, "Instance pre-disposal failed") - } finally { - const remaining = (closing.get(directory) ?? 1) - 1 - if (remaining) closing.set(directory, remaining) - else closing.delete(directory) - } + await dispose(directory, workspaceID, () => + Promise.allSettled([...disposers].map((disposer) => disposer(directory, workspaceID))), + ) } // kilocode_change end diff --git a/packages/opencode/src/effect/runner.ts b/packages/opencode/src/effect/runner.ts index f99ed873f014..154d81629cdc 100644 --- a/packages/opencode/src/effect/runner.ts +++ b/packages/opencode/src/effect/runner.ts @@ -43,6 +43,7 @@ export const make = ( onIdle?: Effect.Effect onBusy?: Effect.Effect onInterrupt?: Effect.Effect + lease?: Effect.Effect<() => void> // kilocode_change }, ): Runner => { const ref = SynchronizedRef.makeUnsafe>({ _tag: "Idle" }) @@ -82,11 +83,22 @@ export const make = ( ).pipe(Effect.flatten) // kilocode_change start - do not let work publish busy before the Running state is committed - const startRun = (work: Effect.Effect, done: Deferred.Deferred) => { + const current = (fiber: Fiber.Fiber) => { + const st = state() + return st._tag === "Running" && st.run.fiber === fiber + } + + const startRun = ( + work: Effect.Effect, + done: Deferred.Deferred, + record: (started: KiloRunner.Started) => void, + ) => { const id = next() return KiloRunner.start({ work, scope, + lease: opts?.lease, + record, finish: (exit) => finishRun(id, done, exit), handle: (fiber) => ({ id, done, fiber }) satisfies RunHandle, }) @@ -95,17 +107,19 @@ export const make = ( // kilocode_change start - open work only after the Running state is committed const finishShell = (id: number) => - SynchronizedRef.modifyEffect( - ref, - Effect.fnUntraced(function* (st) { - if (st._tag === "Shell" && st.shell.id === id) { - return [idle, { _tag: "Idle" }] as const - } - if (st._tag === "ShellThenRun" && st.shell.id === id) { - return yield* KiloRunner.commit(startRun(st.run.work, st.run.done), Effect.void) - } - return [Effect.void, st] as const - }), + KiloRunner.guard(current, (record) => + SynchronizedRef.modifyEffect( + ref, + Effect.fnUntraced(function* (st) { + if (st._tag === "Shell" && st.shell.id === id) { + return [idle, { _tag: "Idle" }] as const + } + if (st._tag === "ShellThenRun" && st.shell.id === id) { + return yield* KiloRunner.commit(startRun(st.run.work, st.run.done, record), Effect.void) + } + return [Effect.void, st] as const + }), + ), ).pipe(Effect.flatten) // kilocode_change end @@ -118,27 +132,29 @@ export const make = ( // kilocode_change start - open work only after the Running state is committed const ensureRunning = (work: Effect.Effect) => - SynchronizedRef.modifyEffect( - ref, - Effect.fnUntraced(function* (st) { - switch (st._tag) { - case "Running": - case "ShellThenRun": - return [awaitDone(st.run.done), st] as const - case "Shell": { - const run = { - id: next(), - done: yield* Deferred.make(), - work, - } satisfies PendingHandle - return [awaitDone(run.done), { _tag: "ShellThenRun", shell: st.shell, run }] as const - } - case "Idle": { - const done = yield* Deferred.make() - return yield* KiloRunner.commit(startRun(work, done), awaitDone(done)) + KiloRunner.guard(current, (record) => + SynchronizedRef.modifyEffect( + ref, + Effect.fnUntraced(function* (st) { + switch (st._tag) { + case "Running": + case "ShellThenRun": + return [awaitDone(st.run.done), st] as const + case "Shell": { + const run = { + id: next(), + done: yield* Deferred.make(), + work, + } satisfies PendingHandle + return [awaitDone(run.done), { _tag: "ShellThenRun", shell: st.shell, run }] as const + } + case "Idle": { + const done = yield* Deferred.make() + return yield* KiloRunner.commit(startRun(work, done, record), awaitDone(done)) + } } - } - }), + }), + ), ).pipe(Effect.flatten) const startShell = (work: Effect.Effect, ready?: Latch.Latch): Effect.Effect => @@ -153,7 +169,12 @@ export const make = ( yield* onBusy const id = next() const cancelled = yield* Deferred.make() - const fiber = yield* work.pipe(Effect.ensuring(finishShell(id)), Effect.forkChild) + // kilocode_change start + const fiber = yield* KiloRunner.fork( + work.pipe(Effect.ensuring(finishShell(id)), Effect.forkChild({ uninterruptible: false })), + opts?.lease, + ) + // kilocode_change end const shell = { id, cancelled, ready, fiber } satisfies ShellHandle return [ Effect.gen(function* () { diff --git a/packages/opencode/src/kilocode/cli/cmd/run.ts b/packages/opencode/src/kilocode/cli/cmd/run.ts index 363bbf900f2a..4115a33e0e14 100644 --- a/packages/opencode/src/kilocode/cli/cmd/run.ts +++ b/packages/opencode/src/kilocode/cli/cmd/run.ts @@ -29,7 +29,6 @@ export namespace KiloRun { model?: string, current?: { id: string; providerID: string }, directory?: string, - signal?: AbortSignal, ) { const selected = resolve(model, current) if (!selected) { @@ -40,15 +39,12 @@ export namespace KiloRun { switch (command) { case "compact": case "summarize": - return sdk.session.summarize( - { - sessionID, - directory, - providerID: selected.providerID, - modelID: selected.modelID, - }, - { signal }, - ) + return sdk.session.summarize({ + sessionID, + directory, + providerID: selected.providerID, + modelID: selected.modelID, + }) } } } diff --git a/packages/opencode/src/kilocode/cli/run-drain.ts b/packages/opencode/src/kilocode/cli/run-drain.ts index d1236d8ea134..e4cb71902dd8 100644 --- a/packages/opencode/src/kilocode/cli/run-drain.ts +++ b/packages/opencode/src/kilocode/cli/run-drain.ts @@ -1,6 +1,7 @@ import { createKiloClient, type Event, type KiloClient } from "@kilocode/sdk/v2" import z from "zod" import { setTimeout } from "node:timers/promises" +import { promisify } from "node:util" const connections = new WeakMap[0]>>() const capability = z.object({ @@ -18,10 +19,10 @@ export namespace KiloRunDrain { return sdk } - export function scope(sdk: KiloClient, directory: string) { + export function scope(sdk: KiloClient, directory: string, signal?: AbortSignal) { const config = connections.get(sdk) if (!config) throw new Error("Missing headless server transport") - return client({ ...config, directory }) + return client({ ...config, directory, signal: signal ?? config.signal }) } export async function check(sdk: KiloClient, signal: AbortSignal) { @@ -78,9 +79,9 @@ export namespace KiloRunDrain { event(event: Event) { if (event.type === "server.connected") connected.resolve() if ( - event.type === "session.turn.close" && - event.properties.sessionID === sessionID && - event.properties.reason === "interrupted" + (event.type === "session.drain.interrupted" || + (event.type === "session.turn.close" && event.properties.reason === "interrupted")) && + event.properties.sessionID === sessionID ) { failure.resolve(new Error("Session interrupted before completion")) } @@ -117,23 +118,16 @@ export namespace KiloRunDrain { ], ) { await Promise.all( - streams.map( - (stream) => - new Promise((resolve, reject) => { - const done = (error?: unknown) => { - const closed = - error instanceof Error && - "code" in error && - (error.code === "EPIPE" || error.code === "ERR_STREAM_DESTROYED") - if (error == null || closed) resolve() - else reject(error) - } - try { - stream.write("", done) - } catch (error) { - done(error) - } - }), + streams.map((stream) => + promisify(stream.write.bind(stream))("").catch((error: unknown) => { + if ( + error instanceof Error && + "code" in error && + (error.code === "EPIPE" || error.code === "ERR_STREAM_DESTROYED") + ) + return + throw error + }), ), ) } diff --git a/packages/opencode/src/kilocode/effect/instance-registry.ts b/packages/opencode/src/kilocode/effect/instance-registry.ts new file mode 100644 index 000000000000..eab43a284e38 --- /dev/null +++ b/packages/opencode/src/kilocode/effect/instance-registry.ts @@ -0,0 +1,36 @@ +import type { WorkspaceV2 } from "@opencode-ai/core/workspace" + +const before = new Set<(directory: string, workspaceID?: WorkspaceV2.ID) => undefined>() +const closing = new Map() + +export function registerBeforeDisposer(disposer: (directory: string, workspaceID?: WorkspaceV2.ID) => undefined) { + before.add(disposer) + return () => { + before.delete(disposer) + } +} + +export function isDisposing(directory: string) { + return closing.has(directory) +} + +export async function dispose(directory: string, workspaceID: WorkspaceV2.ID | undefined, run: () => Promise) { + closing.set(directory, (closing.get(directory) ?? 0) + 1) + try { + const errors: unknown[] = [] + const guards = [...before] + for (const disposer of guards) { + try { + disposer(directory, workspaceID) + } catch (error) { + errors.push(error) + } + } + await run() + if (errors.length) throw new AggregateError(errors, "Instance pre-disposal failed") + } finally { + const remaining = (closing.get(directory) ?? 1) - 1 + if (remaining) closing.set(directory, remaining) + else closing.delete(directory) + } +} diff --git a/packages/opencode/src/kilocode/effect/runner.ts b/packages/opencode/src/kilocode/effect/runner.ts index f880fd23a506..8b9e79f70ed5 100644 --- a/packages/opencode/src/kilocode/effect/runner.ts +++ b/packages/opencode/src/kilocode/effect/runner.ts @@ -1,17 +1,54 @@ import { Effect, Exit, Fiber, Latch, Scope } from "effect" export namespace KiloRunner { + export type Started = { + fiber: Fiber.Fiber + ready: Latch.Latch + } + + export const guard = ( + current: (fiber: Fiber.Fiber) => boolean, + work: (record: (started: Started) => void) => Effect.Effect, + ) => + Effect.suspend(() => { + let started: Started | undefined + return work((value) => { + started = value + }).pipe( + Effect.onExit(() => { + if (!started) return Effect.void + return current(started.fiber) ? started.ready.open : Fiber.interrupt(started.fiber) + }), + ) + }) + + export const fork = (work: Effect.Effect>, lease?: Effect.Effect<() => void>) => + Effect.gen(function* () { + const release = yield* lease ?? Effect.succeed(() => {}) + const fiber = yield* work.pipe(Effect.onError(() => Effect.sync(release))) + fiber.addObserver(release) + return fiber + }).pipe(Effect.uninterruptible) + export const start = (input: { work: Effect.Effect scope: Scope.Scope + lease?: Effect.Effect<() => void> + record: (started: Started) => void finish: (exit: Exit.Exit) => Effect.Effect handle: (fiber: Fiber.Fiber) => R }) => Effect.gen(function* () { const ready = yield* Latch.make() - const fiber = yield* ready.whenOpen(input.work).pipe(Effect.onExit(input.finish), Effect.forkIn(input.scope)) + const fiber = yield* fork( + ready + .whenOpen(input.work) + .pipe(Effect.onExit(input.finish), Effect.forkIn(input.scope, { uninterruptible: false })), + input.lease, + ) + input.record({ fiber, ready }) return { run: input.handle(fiber), ready } - }) + }).pipe(Effect.uninterruptible) export const commit = (start: Effect.Effect<{ run: R; ready: Latch.Latch }>, after: Effect.Effect) => Effect.gen(function* () { diff --git a/packages/opencode/src/kilocode/session/drain.ts b/packages/opencode/src/kilocode/session/drain.ts index 3c42f338dac7..fbedf3734e5f 100644 --- a/packages/opencode/src/kilocode/session/drain.ts +++ b/packages/opencode/src/kilocode/session/drain.ts @@ -1,7 +1,7 @@ import { Context, Deferred, Effect, Layer } from "effect" import { LayerNode } from "@opencode-ai/core/effect/layer-node" import { InstanceState } from "@/effect/instance-state" -import { isDisposing, registerBeforeDisposer } from "@/effect/instance-registry" +import { isDisposing, registerBeforeDisposer } from "@/kilocode/effect/instance-registry" import { capture, type InstanceContext } from "@/kilocode/instance" import type { SessionID } from "@/session/schema" diff --git a/packages/opencode/src/kilocode/session/prompt.ts b/packages/opencode/src/kilocode/session/prompt.ts index 5ac604508328..5aa478b2301e 100644 --- a/packages/opencode/src/kilocode/session/prompt.ts +++ b/packages/opencode/src/kilocode/session/prompt.ts @@ -1,4 +1,3 @@ -// kilocode_change - new file import path from "path" import fs from "fs/promises" import { Cause, Effect, Exit, Fiber, Scope } from "effect" @@ -12,6 +11,9 @@ import { Flag } from "@opencode-ai/core/flag/flag" import { PlanFollowup } from "@/kilocode/plan-followup" import { PlanFile } from "@/kilocode/plan-file" import { KiloSession } from "@/kilocode/session" +import type { SessionDrain } from "@/kilocode/session/drain" +import type { EventV2 } from "@opencode-ai/core/event" +import { Interrupted } from "@opencode-ai/schema/kilocode/session-drain" import { KiloSessionMessageOrder } from "@/kilocode/session/message-order" import { KiloSessionPromptQueue } from "@/kilocode/session/prompt-queue" import { Permission } from "@/permission" @@ -155,32 +157,41 @@ export namespace KiloSessionPrompt { return action === "continue" ? "continue" : "break" } - export const cancelTree = Effect.fn("KiloSessionPrompt.cancelTree")(function* (input: { - sessionID: SessionID - sessions: Pick - cancel: (sessionID: SessionID) => Effect.Effect - }) { - function descendants(sessionID: SessionID): Effect.Effect { - return Effect.gen(function* () { - const children = yield* input.sessions.children(sessionID) - const nested = yield* Effect.forEach(children, (child) => descendants(child.id), { concurrency: "unbounded" }) - return [...children.map((child) => child.id), ...nested.flat()] - }) - } - - const children = yield* descendants(input.sessionID) - yield* Effect.forEach( - [input.sessionID, ...children], - (sessionID) => - Effect.gen(function* () { - yield* KiloSessionPromptQueue.cancel(sessionID) - PlanFollowup.abort(sessionID) - yield* abortIntakes(sessionID) - yield* input.cancel(sessionID) - }), - { concurrency: "unbounded", discard: true }, - ) - }) + export const cancelTree = Effect.fn("KiloSessionPrompt.cancelTree")( + function* (input: { + sessionID: SessionID + sessions: Pick + drain: Pick + events: Pick + cancel: (sessionID: SessionID) => Effect.Effect + }) { + function descendants(sessionID: SessionID): Effect.Effect { + return Effect.gen(function* () { + const children = yield* input.sessions.children(sessionID) + const nested = yield* Effect.forEach(children, (child) => descendants(child.id), { concurrency: "unbounded" }) + return [...children.map((child) => child.id), ...nested.flat()] + }) + } + + const children = yield* descendants(input.sessionID) + yield* Effect.forEach( + [input.sessionID, ...children], + (sessionID) => + Effect.gen(function* () { + yield* KiloSessionPromptQueue.cancel(sessionID) + PlanFollowup.abort(sessionID) + yield* abortIntakes(sessionID) + yield* input.cancel(sessionID) + }), + { concurrency: "unbounded", discard: true }, + ) + }, + (work, input) => + input.drain.track( + input.sessionID, + work.pipe(Effect.ensuring(input.events.publish(Interrupted, { sessionID: input.sessionID }))), + ), + ) export const recoverDanglingAssistant = Effect.fn("KiloSessionPrompt.recoverDanglingAssistant")(function* (input: { sessionID: SessionID diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 3ccc0aa6493a..6ca4c8a7fb7f 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -141,10 +141,14 @@ function isOrphanedInterruptedTool(part: SessionV1.ToolPart) { export interface Interface { readonly cancel: (sessionID: SessionID) => Effect.Effect - readonly prompt: (input: PromptInput) => Effect.Effect + readonly prompt: ( + input: PromptInput, + ) => Effect.Effect readonly loop: (input: LoopInput) => Effect.Effect readonly shell: (input: ShellInput) => Effect.Effect - readonly command: (input: CommandInput) => Effect.Effect + readonly command: ( + input: CommandInput, + ) => Effect.Effect readonly resolvePromptParts: (template: string) => Effect.Effect } @@ -194,7 +198,7 @@ export const layer = Layer.effect( const cancel = Effect.fn("SessionPrompt.cancel")(function* (sessionID: SessionID) { yield* Effect.logInfo("cancel", { "session.id": sessionID }) - yield* KiloSessionPrompt.cancelTree({ sessionID, sessions, cancel: state.cancel }) // kilocode_change - stop queued work and subagents + yield* KiloSessionPrompt.cancelTree({ sessionID, sessions, drain, events, cancel: state.cancel }) // kilocode_change - stop queued work and subagents }) // kilocode_change start - preserve configured reference mentions on the Core reference architecture @@ -1431,23 +1435,21 @@ export const layer = Layer.effect( // Queue tails and runner fibers can resume outside the HTTP request's // ambient instance context; bridge both Effect refs and legacy ALS. const bridge = yield* EffectBridge.make() - return yield* drain.track( + return yield* KiloSessionPromptQueue.enqueue( input.sessionID, - KiloSessionPromptQueue.enqueue( - input.sessionID, - message.info.id, - bridge.run( - loop({ sessionID: input.sessionID, snapshotInitialization: input.snapshotInitialization }).pipe( - Effect.orDie, - ), + message.info.id, + bridge.run( + loop({ sessionID: input.sessionID, snapshotInitialization: input.snapshotInitialization }).pipe( + Effect.orDie, ), - bridge.run(lastAssistant(input.sessionID)), - dismiss, - ), + ), // kilocode_change + bridge.run(lastAssistant(input.sessionID)), + dismiss, ) // kilocode_change end }, Effect.catchTag("NotFoundError", Effect.die), + (work, input) => drain.track(input.sessionID, work), // kilocode_change ) const lastAssistant = Effect.fnUntraced(function* (sessionID: SessionID) { @@ -1622,13 +1624,7 @@ export const layer = Layer.effect( } compactionAttempts++ // kilocode_change end - yield* compaction.create({ - sessionID, - agent: lastUser.agent, - model: lastUser.model, - auto: true, - overflow: false, - }) // kilocode_change + yield* compaction.create({ sessionID, agent: lastUser.agent, model: lastUser.model, auto: true, overflow: false }) // kilocode_change continue } diff --git a/packages/opencode/src/session/run-state.ts b/packages/opencode/src/session/run-state.ts index 880ed3dc911f..be04864e6a30 100644 --- a/packages/opencode/src/session/run-state.ts +++ b/packages/opencode/src/session/run-state.ts @@ -66,6 +66,7 @@ export const layer = Layer.effect( }), onBusy: status.set(sessionID, { type: "busy" }), onInterrupt, + lease: drain.hold(sessionID), // kilocode_change }) data.runners.set(sessionID, next) return next diff --git a/packages/opencode/src/tool/task.ts b/packages/opencode/src/tool/task.ts index 72135a44ec02..a5ee9ff24cdd 100644 --- a/packages/opencode/src/tool/task.ts +++ b/packages/opencode/src/tool/task.ts @@ -20,7 +20,8 @@ import { KiloSessionProcessor } from "../kilocode/session/processor" // kilocode import { KiloSession } from "../kilocode/session" // kilocode_change import { resumeHint } from "../kilocode/task-resume" // kilocode_change import { errorMessage } from "@/util/error" // kilocode_change -import { Cause, Effect, Exit, Schema, Scope } from "effect" // kilocode_change +import { Effect, Exit, Schema, Scope } from "effect" +import { Cause } from "effect" // kilocode_change import { EffectBridge } from "@/effect/bridge" import { RuntimeFlags } from "@/effect/runtime-flags" import * as SandboxPolicy from "@/kilocode/sandbox/policy" // kilocode_change diff --git a/packages/opencode/test/event-manifest.test.ts b/packages/opencode/test/event-manifest.test.ts index bc87fc4708d0..1ca54fea79b2 100644 --- a/packages/opencode/test/event-manifest.test.ts +++ b/packages/opencode/test/event-manifest.test.ts @@ -3,15 +3,16 @@ import { SessionEvent } from "@opencode-ai/core/session/event" import { EventManifest as SchemaEventManifest } from "@opencode-ai/schema/event-manifest" import { Todo } from "@/session/todo" import { EventManifest } from "@/event-manifest" -import { Drained } from "@opencode-ai/schema/kilocode/session-drain" // kilocode_change +import { Drained, Interrupted } from "@opencode-ai/schema/kilocode/session-drain" // kilocode_change describe("public event manifest", () => { test("contains every latest public wire type once", () => { expect(EventManifest.Definitions).toBe(SchemaEventManifest.Definitions) expect(EventManifest.Latest).toBe(SchemaEventManifest.Latest) expect(EventManifest.Durable).toBe(SchemaEventManifest.Durable) - expect(EventManifest.Latest.size).toBe(91) // kilocode_change - include session.drained + expect(EventManifest.Latest.size).toBe(92) // kilocode_change - include session drain events expect(EventManifest.Latest.get("session.drained")).toBe(Drained) // kilocode_change + expect(EventManifest.Latest.get("session.drain.interrupted")).toBe(Interrupted) // kilocode_change expect(EventManifest.Latest.get("session.next.step.ended")).toBe(SessionEvent.Step.Ended) expect(EventManifest.Latest.get("todo.updated")).toBe(Todo.Event.Updated) expect(EventManifest.Latest.has("ide.installed")).toBe(false) diff --git a/packages/opencode/test/kilocode/headless-session-drain.test.ts b/packages/opencode/test/kilocode/headless-session-drain.test.ts index db9355d059b8..228512820311 100644 --- a/packages/opencode/test/kilocode/headless-session-drain.test.ts +++ b/packages/opencode/test/kilocode/headless-session-drain.test.ts @@ -1,6 +1,6 @@ import { expect } from "bun:test" import { Effect, Fiber } from "effect" -import { createKiloClient, type Event } from "@kilocode/sdk/v2" +import { createKiloClient, type Event, type KiloClient } from "@kilocode/sdk/v2" import path from "node:path" import { mkdir } from "node:fs/promises" import { pathToFileURL } from "node:url" @@ -23,6 +23,66 @@ function gate() { } } +function environment(home: string, url: string, opts: { plugin?: string; resume?: boolean; daemon?: boolean } = {}) { + const base = testProviderConfig(url) + const model = base.provider.test.models["test-model"] + const cfg = { + ...base, + model: "test/parent", + small_model: "test/parent", + enabled_providers: ["test"], + share: "disabled", + snapshot: false, + sandbox: { enabled: false }, + permission: { "*": "allow" }, + plugin: opts.plugin ? [opts.plugin] : [], + agent: { + code: { model: "test/parent" }, + general: { model: "test/child" }, + title: { disable: true }, + }, + provider: { + test: { + ...base.provider.test, + models: { + parent: { ...model, id: "parent", name: "Parent" }, + child: { ...model, id: "child", name: "Child" }, + }, + }, + }, + } + return { + PWD: home, + KILO_CONFIG_CONTENT: JSON.stringify(cfg), + KILO_CONFIG: "", + KILO_CONFIG_DIR: "", + KILO_DB: opts.resume ? path.join(home, "resume.db") : ":memory:", + KILO_AUTH_CONTENT: "{}", + KILO_SERVER_PASSWORD: "", + KILO_SERVER_USERNAME: "kilo", + KILO_PURE: "false", + KILO_DISABLE_DEFAULT_PLUGINS: "true", + KILO_EXPERIMENTAL_BACKGROUND_SUBAGENTS: "true", + KILO_TELEMETRY_LEVEL: "off", + KILO_AUTO_SHARE: "false", + KILO_NO_DAEMON: opts.daemon ? "" : "1", + KILO_PARENT_PID: "", + KILO_TEST_DAEMON_STATE_DIR: path.join(home, "daemon-state"), + KILO_TEST_DAEMON_LOG_DIR: path.join(home, "daemon-log"), + } +} + +function drain(client: KiloClient, id: string, timeout = 30_000) { + return client.kilocode.drainSession( + { sessionID: id, token: crypto.randomUUID() }, + { throwOnError: true, signal: AbortSignal.timeout(timeout) }, + ) +} + +function pending(client: KiloClient, id: string) { + return Effect.promise(() => drain(client, id, 1_000).catch((err: unknown) => err)) +} + function observer(url: string, busy: boolean) { const sdk = pathToFileURL(path.resolve(import.meta.dir, "../../../sdk/js/src/v2/index.ts")).href return `import { createKiloClient } from ${JSON.stringify(sdk)} @@ -77,7 +137,7 @@ function observer(url: string, busy: boolean) { }` } -function probe(home: string, busy: boolean) { +function probe(home: string, mode: "foreground" | "busy" | "idle" | "intake") { return Effect.gen(function* () { const entered = gate() const release = gate() @@ -114,7 +174,11 @@ function probe(home: string, busy: boolean) { ) const url = `http://127.0.0.1:${server.port}` const file = path.join(home, "drain-observer.mjs") - yield* Effect.promise(() => Bun.write(file, observer(url, busy))) + const plugin = + mode === "intake" + ? `export default async () => ({ "chat.message": async () => { const response = await fetch(${JSON.stringify(`${url}/tool`)}); if (!response.ok) throw new Error("intake gate failed") } })` + : observer(url, mode === "busy") + yield* Effect.promise(() => Bun.write(file, plugin)) return { url, plugin: pathToFileURL(file).href, entered, release, queued, idle } }) } @@ -122,42 +186,15 @@ function probe(home: string, busy: boolean) { function scenario( { home, llm, opencode }: CliFixture, mode: "foreground" | "busy" | "idle", - opts: { attach?: boolean; resume?: boolean; daemon?: boolean } = {}, + opts: { attach?: boolean; resume?: boolean; daemon?: boolean; abort?: boolean } = {}, ) { return Effect.gen(function* () { const directory = opts.resume ? path.join(home, "session") : home if (opts.resume) yield* Effect.promise(() => mkdir(directory)) - const tap = yield* probe(home, mode === "busy") + const tap = yield* probe(home, mode) const child = gate() const requested = gate() yield* Effect.addFinalizer(() => Effect.sync(child.open)) - const base = testProviderConfig(llm.url) - const model = base.provider.test.models["test-model"] - const cfg = { - ...base, - model: "test/parent", - small_model: "test/parent", - enabled_providers: ["test"], - share: "disabled", - snapshot: false, - sandbox: { enabled: false }, - permission: { "*": "allow" }, - plugin: mode === "foreground" ? [] : [tap.plugin], - agent: { - code: { model: "test/parent" }, - general: { model: "test/child" }, - title: { disable: true }, - }, - provider: { - test: { - ...base.provider.test, - models: { - parent: { ...model, id: "parent", name: "Parent" }, - child: { ...model, id: "child", name: "Child" }, - }, - }, - }, - } const script = `const r = await fetch(${JSON.stringify(`${tap.url}/tool`)}); if (!r.ok) throw new Error("tool gate failed"); console.log(await r.text())` const scriptPath = path.join(home, "parent-gate.mjs") yield* Effect.promise(() => Bun.write(scriptPath, script)) @@ -187,28 +224,10 @@ function scenario( return true }, reply().wait(child.promise).text("CHILD_RESULT_SENTINEL").stop()) if (mode === "foreground") child.open() - const env = { - PWD: home, - KILO_CONFIG_CONTENT: JSON.stringify(cfg), - KILO_CONFIG: "", - KILO_CONFIG_DIR: "", - KILO_DB: opts.resume ? path.join(home, "resume.db") : ":memory:", - KILO_AUTH_CONTENT: "{}", - KILO_SERVER_PASSWORD: "", - KILO_SERVER_USERNAME: "kilo", - KILO_PURE: "false", - KILO_DISABLE_DEFAULT_PLUGINS: "true", - KILO_EXPERIMENTAL_BACKGROUND_SUBAGENTS: "true", - KILO_TELEMETRY_LEVEL: "off", - KILO_AUTO_SHARE: "false", - KILO_NO_DAEMON: opts.daemon ? "" : "1", - KILO_PARENT_PID: "", - KILO_TEST_DAEMON_STATE_DIR: path.join(home, "daemon-state"), - KILO_TEST_DAEMON_LOG_DIR: path.join(home, "daemon-log"), - } + const env = environment(home, llm.url, { ...opts, plugin: mode === "foreground" ? undefined : tap.plugin }) const server = opts.attach || opts.resume ? yield* opencode.serve({ env, readyTimeoutMs: 30_000 }) : undefined const session = - opts.resume && server + (opts.resume || opts.abort) && server ? yield* Effect.promise(() => createKiloClient({ baseUrl: server.url, directory }).session.create({}, { throwOnError: true }), ) @@ -257,9 +276,52 @@ function scenario( } if (mode === "idle") { yield* tap.idle.wait("parent never became idle") - child.open() + if (!opts.abort) child.open() + } + if (opts.abort) { + if (!server || !session) throw new Error("Abort scenario requires an attached session") + const client = createKiloClient({ baseUrl: server.url, directory }) + const controller = new AbortController() + yield* Effect.addFinalizer(() => Effect.sync(() => controller.abort())) + const ready = gate() + const token = crypto.randomUUID() + const stream = yield* Effect.promise(() => + client.event.subscribe(undefined, { signal: controller.signal, sseMaxRetryAttempts: 1 }), + ) + const received = yield* Effect.promise(async () => { + const signals: string[] = [] + for await (const event of stream.stream) { + if (event.type === "server.connected") ready.open() + if ( + (event.type === "session.drain.interrupted" || + (event.type === "session.turn.close" && event.properties.reason === "interrupted")) && + event.properties.sessionID === session.data.id + ) + signals.push(event.type) + if (event.type === "session.drained" && event.properties.token === token) return signals + } + throw new Error("Cancellation event stream ended before the barrier") + }).pipe(Effect.forkChild) + yield* ready.wait("Cancellation event stream did not connect") + const aborted = yield* Effect.promise(() => + client.session.abort({ sessionID: session.data.id }, { throwOnError: true }), + ) + expect(aborted.data).toBe(true) + yield* Effect.promise(() => + client.kilocode.drainSession({ sessionID: session.data.id, token }, { throwOnError: true }), + ) + expect(yield* awaitWithTimeout(Fiber.join(received), "Cancellation barrier was not delivered")).toEqual([ + "session.drain.interrupted", + ]) } const result = yield* awaitWithTimeout(run.result, "CLI did not drain and exit", "45 seconds") + if (opts.abort) { + expect(result.exitCode).not.toBe(0) + expect(result.stdout).toContain("WAITING_FOR_CHILD") + expect(result.stdout).not.toContain("FINAL_AFTER_CHILD") + expect(yield* llm.pending).toBe(1) + return + } opencode.expectExit(result, 0) const events = opencode.parseJsonEvents(result.stdout) const texts = events.filter((event) => event.type === "text").map((event) => (event.part as { text: string }).text) @@ -282,6 +344,11 @@ for (const [name, mode, opts] of [ ["headless run waits for a child after the parent becomes idle", "idle", {}], ["attached run drains a callback queued behind the parent", "busy", { attach: true }], ["attached run waits for a child after the parent becomes idle", "idle", { attach: true }], + [ + "attached run fails when the idle parent is aborted with a child still running", + "idle", + { attach: true, abort: true }, + ], ["local run resumes and drains a session from another directory", "busy", { resume: true }], ["attached run resumes and drains a session from another directory", "idle", { attach: true, resume: true }], ["daemon run resumes and drains a session from another directory", "idle", { daemon: true, resume: true }], @@ -289,6 +356,81 @@ for (const [name, mode, opts] of [ cliIt.live(name, (fixture) => scenario(fixture, mode, opts), 90_000) } +for (const [name, mode] of [ + ["SDK drain waits for accepted promptAsync intake and work", "intake"], + ["SDK drain waits for summarize work after HTTP cancellation", "summary"], +] as const) { + cliIt.live( + name, + ({ home, llm, opencode }) => + Effect.gen(function* () { + const tap = mode === "intake" ? yield* probe(home, "intake") : undefined + const release = gate() + yield* Effect.addFinalizer(() => Effect.sync(release.open)) + const server = yield* opencode.serve({ + env: environment(home, llm.url, { plugin: tap?.plugin }), + readyTimeoutMs: 30_000, + }) + const client = createKiloClient({ baseUrl: server.url, directory: home }) + const session = yield* Effect.promise(() => client.session.create({}, { throwOnError: true })) + const id = session.data.id + const text = mode === "intake" ? "INTAKE_WORK_FINISHED" : "SUMMARY_AFTER_DISCONNECT" + if (mode === "summary") { + yield* llm.text("Seed response") + yield* Effect.promise(() => + client.session.prompt( + { sessionID: id, parts: [{ type: "text", text: "Seed the summary session" }] }, + { throwOnError: true }, + ), + ) + yield* Effect.promise(() => drain(client, id)) + } + yield* llm.hold(text, release.promise) + if (tap) { + const accepted = yield* Effect.promise(() => + client.session.promptAsync( + { sessionID: id, parts: [{ type: "text", text: "Run admitted work" }] }, + { throwOnError: true }, + ), + ) + expect(accepted.response.status).toBe(204) + yield* tap.entered.wait("accepted prompt did not enter the intake hook") + expect(yield* llm.calls).toBe(0) + expect(yield* pending(client, id)).toMatchObject({ name: "TimeoutError" }) + tap.release.open() + yield* awaitWithTimeout(llm.wait(1), "accepted prompt never reached the provider", "30 seconds") + } + if (mode === "summary") { + const abort = new AbortController() + yield* Effect.addFinalizer(() => Effect.sync(() => abort.abort())) + const caller = client.session + .summarize( + { sessionID: id, providerID: "test", modelID: "parent", auto: false }, + { throwOnError: true, signal: abort.signal }, + ) + .then( + () => undefined, + (err: unknown) => err, + ) + yield* awaitWithTimeout(llm.wait(2), "summary never reached the provider", "30 seconds") + abort.abort() + expect(yield* Effect.promise(() => caller)).toMatchObject({ name: "AbortError" }) + const status = yield* Effect.promise(() => client.session.status({}, { throwOnError: true })) + expect(status.data[id]?.type).toBe("busy") + } + const waiting = yield* pending(client, id) + release.open() + const finished = yield* Effect.promise(() => drain(client, id)) + expect(waiting).toMatchObject({ name: "TimeoutError" }) + expect(finished.data).toBe(true) + const messages = yield* Effect.promise(() => client.session.messages({ sessionID: id }, { throwOnError: true })) + expect(JSON.stringify(messages.data)).toContain(text) + expect(yield* llm.pending).toBe(0) + }), + 90_000, + ) +} + cliIt.live( "SDK drain follows the session directory and validates acknowledgments", ({ home, opencode }) => diff --git a/packages/opencode/test/kilocode/run-auto.test.ts b/packages/opencode/test/kilocode/run-auto.test.ts index 3c107d8787c9..21ea28b18a80 100644 --- a/packages/opencode/test/kilocode/run-auto.test.ts +++ b/packages/opencode/test/kilocode/run-auto.test.ts @@ -1,4 +1,8 @@ import { afterEach, describe, expect, mock, test } from "bun:test" +import * as SDK from "@kilocode/sdk/v2" +import { RunCommand } from "@/cli/cmd/run" + +const actual = { ...SDK } type Event = { type: string @@ -114,7 +118,8 @@ function args() { const tty = Object.getOwnPropertyDescriptor(process.stdin, "isTTY") const exitCode = process.exitCode -afterEach(() => { +afterEach(async () => { + await mock.module("@kilocode/sdk/v2", () => actual) process.exitCode = exitCode ?? 0 if (tty) { Object.defineProperty(process.stdin, "isTTY", tty) @@ -123,9 +128,12 @@ afterEach(() => { delete (process.stdin as { isTTY?: boolean }).isTTY }) -async function run(sdk: Record) { - mock.module("@kilocode/sdk/v2", () => ({ - createKiloClient: (config: { fetch?: () => Promise }) => { +type Transport = { signal?: AbortSignal | null } + +async function run(sdk: Record, transport: Transport = {}) { + await mock.module("@kilocode/sdk/v2", () => ({ + createKiloClient: (config: Transport & { fetch?: () => Promise }) => { + transport.signal = config.signal config.fetch = async () => Response.json({ paths: { @@ -141,8 +149,6 @@ async function run(sdk: Record) { value: true, }) - const key = JSON.stringify({ time: Date.now(), rand: Math.random() }) - const { RunCommand } = await import(`../../src/cli/cmd/run?${key}`) return RunCommand.handler(args() as never) } @@ -198,6 +204,7 @@ describe("cli run auto permissions", () => { test("a failed prompt aborts a blocked permission reply during cleanup", async () => { const q = feed() + const transport: Transport = {} const started = Promise.withResolvers() let aborted = false const sdk = { @@ -209,19 +216,21 @@ describe("cli run auto permissions", () => { }, }, permission: { - reply: async (_input: unknown, options: { signal: AbortSignal }) => { - expect(options.signal).toBeInstanceOf(AbortSignal) + reply: async () => { + const signal = transport.signal + expect(signal).toBeInstanceOf(AbortSignal) + if (!signal) throw new Error("Missing client cancellation signal") + const result = Promise.withResolvers() + signal.addEventListener( + "abort", + () => { + aborted = true + result.reject(signal.reason) + }, + { once: true }, + ) started.resolve() - return new Promise((_resolve, reject) => { - options.signal.addEventListener( - "abort", - () => { - aborted = true - reject(options.signal.reason) - }, - { once: true }, - ) - }) + return result.promise }, }, session: { @@ -234,7 +243,7 @@ describe("cli run auto permissions", () => { }, }, } - await run(sdk) + await run(sdk, transport) q.end() expect(aborted).toBe(true) expect(process.exitCode).toBe(1) @@ -242,6 +251,7 @@ describe("cli run auto permissions", () => { test("handles a tracked child's network retry without touching another session", async () => { const q = feed() + const transport: Transport = {} const done = Promise.withResolvers() const calls: string[] = [] const sdk = { @@ -253,8 +263,8 @@ describe("cli run auto permissions", () => { }, }, network: { - reply: async (input: { requestID: string }, options: { signal: AbortSignal }) => { - expect(options.signal).toBeInstanceOf(AbortSignal) + reply: async (input: { requestID: string }) => { + expect(transport.signal).toBeInstanceOf(AbortSignal) calls.push(input.requestID) done.resolve() return { data: true } @@ -278,7 +288,7 @@ describe("cli run auto permissions", () => { }, }, } - await run(sdk) + await run(sdk, transport) expect(calls).toEqual(["net_child"]) }, 15_000) }) diff --git a/packages/opencode/test/kilocode/run-drain.test.ts b/packages/opencode/test/kilocode/run-drain.test.ts index ad6570804f4e..92bc2f3114ae 100644 --- a/packages/opencode/test/kilocode/run-drain.test.ts +++ b/packages/opencode/test/kilocode/run-drain.test.ts @@ -1,5 +1,6 @@ import { expect, test } from "bun:test" import { KiloRunDrain } from "@/kilocode/cli/run-drain" +import { KiloRun } from "@/kilocode/cli/cmd/run" function client(handle: (request: Request) => Promise) { return KiloRunDrain.client({ @@ -28,19 +29,34 @@ test("checks capabilities through the selected transport before accepting comple await KiloRunDrain.check(sdk, new AbortController().signal) }) -test("scopes requests to the session directory without changing the transport", async () => { +test("scopes requests without losing the transport or cancellation signal", async () => { + const abort = new AbortController() const paths: Array = [] + const signals: AbortSignal[] = [] const sdk = client(async (request) => { const url = new URL(request.url) expect(url.origin).toBe("http://drain.test") - expect(url.pathname).toBe("/prefix/config") + expect(url.pathname).toBe(request.method === "GET" ? "/prefix/config" : "/prefix/session/ses_test/summarize") expect(request.headers.get("authorization")).toBe("Basic test-only") paths.push(url.searchParams.get("directory")) - return Response.json({}) + signals.push(request.signal) + return Response.json(request.method === "GET" ? {} : true) }) - await KiloRunDrain.scope(sdk, "/session owner").config.get() + const scoped = KiloRunDrain.scope(sdk, "/session owner", abort.signal) + await scoped.config.get() + await KiloRunDrain.scope(scoped, "/resumed owner").config.get() + await KiloRun.runBuiltin( + scoped, + "ses_test", + "compact", + undefined, + { providerID: "test", id: "model" }, + "/session owner", + ) await sdk.config.get() - expect(paths).toEqual(["/session owner", null]) + abort.abort() + expect(paths).toEqual(["/session owner", "/resumed owner", "/session owner", null]) + expect(signals.map((signal) => signal.aborted)).toEqual([true, true, true, false]) }) test.each([ @@ -99,8 +115,8 @@ test("flush waits for all output callbacks", async () => { const second = Promise.withResolvers() let done = false const output = KiloRunDrain.flush([ - { write: (_chunk, callback) => first.promise.then(() => callback()) }, - { write: (_chunk, callback) => second.promise.then(() => callback()) }, + { write: (_chunk, callback) => void first.promise.then(() => callback()) }, + { write: (_chunk, callback) => void second.promise.then(() => callback()) }, ]).then(() => { done = true }) @@ -160,6 +176,22 @@ test("closing a run aborts pending retry waits", async () => { await expect(waiting).rejects.toThrow() }) +test("session cancellation fails a drain after the parent turn already closed", async () => { + const drain = KiloRunDrain.create("ses_parent") + drain.event({ + id: "evt_done", + type: "session.turn.close", + properties: { sessionID: "ses_parent", reason: "completed" }, + }) + drain.event({ id: "evt_other", type: "session.drain.interrupted", properties: { sessionID: "ses_other" } }) + drain.event({ id: "evt_ready", type: "server.connected", properties: {} }) + await drain.ready() + const waiting = drain.race(Promise.withResolvers().promise) + drain.event({ id: "evt_cancel", type: "session.drain.interrupted", properties: { sessionID: "ses_parent" } }) + await expect(waiting).rejects.toThrow("Session interrupted") + drain.close() +}) + test("a root interruption fails but a superseded handoff and child interruption do not", async () => { const drain = KiloRunDrain.create("ses_parent") drain.event({ diff --git a/packages/opencode/test/kilocode/run-network.test.ts b/packages/opencode/test/kilocode/run-network.test.ts index 8da63bea09ab..dd6eb10d0aa7 100644 --- a/packages/opencode/test/kilocode/run-network.test.ts +++ b/packages/opencode/test/kilocode/run-network.test.ts @@ -1,4 +1,9 @@ import { afterEach, describe, expect, mock, spyOn, test } from "bun:test" +import * as SDK from "@kilocode/sdk/v2" +import { RunCommand } from "@/cli/cmd/run" +import { KiloRunDrain } from "@/kilocode/cli/run-drain" + +const actual = { ...SDK } type Event = { type: string @@ -96,7 +101,8 @@ const exitCode = process.exitCode const tty = Object.getOwnPropertyDescriptor(process.stdin, "isTTY") const text = Bun.stdin.text -afterEach(() => { +afterEach(async () => { + await mock.module("@kilocode/sdk/v2", () => actual) process.exitCode = exitCode ?? 0 Bun.stdin.text = text if (tty) { @@ -142,9 +148,6 @@ async function run( value: terminal, }) - const key = JSON.stringify({ time: Date.now(), rand: Math.random() }) - const { RunCommand } = await import(`../../src/cli/cmd/run?${key}`) - const { KiloRunDrain } = await import("../../src/kilocode/cli/run-drain") const create = KiloRunDrain.create const clock = spyOn(KiloRunDrain, "create").mockImplementation((id) => ({ ...create(id), diff --git a/packages/opencode/test/kilocode/runner-start-order.test.ts b/packages/opencode/test/kilocode/runner-start-order.test.ts index 1faedf21629c..554f238355ac 100644 --- a/packages/opencode/test/kilocode/runner-start-order.test.ts +++ b/packages/opencode/test/kilocode/runner-start-order.test.ts @@ -1,7 +1,12 @@ import { describe, expect } from "bun:test" -import { Deferred, Effect, Fiber, Scope } from "effect" +import { Deferred, Effect, Exit, Fiber, Scope } from "effect" import { Runner } from "@/effect/runner" -import { awaitWithTimeout, it, pollWithTimeout } from "../lib/effect" +import { KiloRunner } from "@/kilocode/effect/runner" +import { SessionDrain } from "@/kilocode/session/drain" +import { SessionID } from "@/session/schema" +import { awaitWithTimeout, pollWithTimeout, testEffect } from "../lib/effect" + +const it = testEffect(SessionDrain.layer) describe("Runner start ordering", () => { it.live( @@ -92,4 +97,126 @@ describe("Runner start ordering", () => { ) }), ) + + it.instance( + "keeps drain blocked during cancelled callers' shell cleanup and handoff", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const id = SessionID.make("ses_runner_shell_cleanup") + const runner = Runner.make(yield* Scope.Scope, { lease: drain.hold(id) }) + const entered = yield* Deferred.make() + const stopping = yield* Deferred.make() + const cleanup = yield* Deferred.make() + const started = yield* Deferred.make() + const finish = yield* Deferred.make() + yield* Effect.gen(function* () { + const work = Deferred.succeed(entered, undefined).pipe( + Effect.andThen(Effect.never), + Effect.onInterrupt(() => Deferred.succeed(stopping, undefined).pipe(Effect.andThen(Deferred.await(cleanup)))), + ) + const shell = yield* drain.track(id, runner.startShell(work)).pipe(Effect.forkChild) + yield* Deferred.await(entered) + const run = Deferred.succeed(started, undefined).pipe(Effect.andThen(Deferred.await(finish)), Effect.as("done")) + const caller = yield* drain.track(id, runner.ensureRunning(run)).pipe(Effect.forkChild) + yield* pollWithTimeout( + Effect.sync(() => (runner.state._tag === "ShellThenRun" ? true : undefined)), + "runner did not queue work", + ) + yield* Fiber.interrupt(caller) + const waiting = yield* drain.wait(id).pipe(Effect.forkChild({ startImmediately: true })) + const stop = yield* Fiber.interrupt(shell).pipe(Effect.forkChild) + yield* awaitWithTimeout(Deferred.await(stopping), "shell cleanup did not start") + expect(runner.state._tag).toBe("ShellThenRun") + expect(waiting.pollUnsafe()).toBeUndefined() + yield* Deferred.succeed(cleanup, undefined) + yield* awaitWithTimeout(Fiber.join(stop), "shell cleanup did not finish") + yield* awaitWithTimeout(Deferred.await(started), "queued work did not start") + expect(runner.busy).toBe(true) + expect(waiting.pollUnsafe()).toBeUndefined() + yield* awaitWithTimeout(runner.cancel, "queued work inherited an uninterruptible shell finalizer") + yield* awaitWithTimeout(Fiber.join(waiting), "shell handoff reservation was not released") + expect(runner.busy).toBe(false) + }).pipe( + Effect.ensuring(Deferred.succeed(cleanup, undefined).pipe(Effect.andThen(Deferred.succeed(finish, undefined)))), + ) + }), + ) + + it.instance( + "releases provisional work when its caller is interrupted before commit", + Effect.gen(function* () { + const drain = yield* SessionDrain.Service + const id = SessionID.make("ses_runner_provisional") + const runner = Runner.make(yield* Scope.Scope, { + lease: drain + .hold(id) + .pipe(Effect.tap(() => Effect.withFiber((fiber) => Effect.sync(() => fiber.interruptUnsafe())))), + }) + const started = yield* Deferred.make() + const work = Deferred.succeed(started, undefined).pipe(Effect.as("done")) + const caller = yield* drain.track(id, runner.ensureRunning(work)).pipe(Effect.forkChild) + + const exit = yield* awaitWithTimeout(Fiber.await(caller), "provisional work rollback did not finish") + expect(Exit.hasInterrupts(exit)).toBe(true) + expect(yield* Deferred.isDone(started)).toBe(false) + expect(runner.busy).toBe(false) + yield* awaitWithTimeout(drain.wait(id), "provisional work reservation was not released") + }), + ) + + it.instance( + "releases a closed-scope fiber before its first evaluation", + Effect.gen(function* () { + const scope = yield* Scope.make() + yield* Scope.close(scope, Exit.void) + const drain = yield* SessionDrain.Service + const id = SessionID.make("ses_runner_closed") + const work = yield* Deferred.make() + const finish = yield* Deferred.make() + const started = yield* KiloRunner.guard( + () => false, + (record) => + KiloRunner.start({ + scope, + lease: drain.hold(id), + record, + work: Deferred.succeed(work, undefined), + finish: () => Deferred.succeed(finish, undefined).pipe(Effect.asVoid), + handle: (fiber) => fiber, + }), + ) + + expect(Exit.hasInterrupts(yield* Fiber.await(started.run))).toBe(true) + expect(yield* Deferred.isDone(work)).toBe(false) + expect(yield* Deferred.isDone(finish)).toBe(false) + yield* awaitWithTimeout(drain.wait(id), "closed-scope fiber reservation was not released") + }), + ) + + it.live( + "can interrupt callers waiting for the runner mutex", + Effect.gen(function* () { + const entered = yield* Deferred.make() + const unlock = yield* Deferred.make() + const finish = yield* Deferred.make() + yield* Effect.addFinalizer(() => + Deferred.succeed(unlock, undefined).pipe(Effect.andThen(Deferred.succeed(finish, undefined))), + ) + const runner = Runner.make(yield* Scope.Scope, { + onBusy: Deferred.succeed(entered, undefined).pipe(Effect.andThen(Deferred.await(unlock))), + }) + const shell = yield* runner.startShell(Deferred.await(finish).pipe(Effect.as("shell"))).pipe(Effect.forkChild) + yield* Deferred.await(entered) + const started = yield* Deferred.make() + const work = Deferred.succeed(started, undefined).pipe(Effect.as("done")) + const caller = yield* runner.ensureRunning(work).pipe(Effect.forkChild({ startImmediately: true })) + + yield* awaitWithTimeout(Fiber.interrupt(caller), "caller could not cancel its mutex wait") + yield* Deferred.succeed(unlock, undefined) + yield* Deferred.succeed(finish, undefined) + expect(yield* Fiber.join(shell)).toBe("shell") + expect(yield* Deferred.isDone(started)).toBe(false) + expect(runner.busy).toBe(false) + }), + ) }) diff --git a/packages/opencode/test/kilocode/session-drain.test.ts b/packages/opencode/test/kilocode/session-drain.test.ts index 76df7acc4e81..8a6376c2ab6d 100644 --- a/packages/opencode/test/kilocode/session-drain.test.ts +++ b/packages/opencode/test/kilocode/session-drain.test.ts @@ -16,24 +16,16 @@ it.instance( const id = SessionID.make("ses_drain_parent") const parent = yield* drain.hold(id) const delivery = yield* drain.hold(id) - let done = false - const waiter = yield* drain.wait(id).pipe( - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) - expect(done).toBe(false) + const waiting = yield* drain.wait(id).pipe(Effect.forkChild({ startImmediately: true })) + expect(waiting.pollUnsafe()).toBeUndefined() parent() - expect(done).toBe(false) + expect(waiting.pollUnsafe()).toBeUndefined() const callback = yield* drain.hold(id) delivery() - expect(done).toBe(false) + expect(waiting.pollUnsafe()).toBeUndefined() callback() - yield* Fiber.join(waiter) - expect(done).toBe(true) + yield* Fiber.join(waiting) + expect(waiting.pollUnsafe()?._tag).toBe("Success") }), ) @@ -47,19 +39,11 @@ it.instance( yield* drain.link(child, parent) yield* drain.link(grandchild, child) const release = yield* drain.hold(grandchild) - let done = false - const waiter = yield* drain.wait(parent).pipe( - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) + const waiting = yield* drain.wait(parent).pipe(Effect.forkChild({ startImmediately: true })) yield* drain.wait(SessionID.make("ses_drain_other")) - expect(done).toBe(false) + expect(waiting.pollUnsafe()).toBeUndefined() release() - yield* Fiber.join(waiter) + yield* Fiber.join(waiting) yield* drain.wait(child) yield* drain.wait(grandchild) }), @@ -72,20 +56,12 @@ it.instance( const id = SessionID.make("ses_drain_parent") const previous = yield* drain.hold(id) const next = yield* drain.hold(id) - let done = false - const waiter = yield* drain.wait(id).pipe( - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) + const waiting = yield* drain.wait(id).pipe(Effect.forkChild({ startImmediately: true })) previous() previous() - expect(done).toBe(false) + expect(waiting.pollUnsafe()).toBeUndefined() next() - yield* Fiber.join(waiter) + yield* Fiber.join(waiting) }), ) @@ -98,20 +74,12 @@ it.instance( for (let round = 0; round < 2; round++) { const release = yield* drain.hold(id) const cancelled = yield* drain.wait(id).pipe(Effect.forkChild({ startImmediately: true })) - let done = false - const waiter = yield* drain.wait(id).pipe( - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) + const waiting = yield* drain.wait(id).pipe(Effect.forkChild({ startImmediately: true })) yield* Fiber.interrupt(cancelled) - expect(done).toBe(false) + expect(waiting.pollUnsafe()).toBeUndefined() release() - yield* Fiber.join(waiter) - expect(done).toBe(true) + yield* Fiber.join(waiting) + expect(waiting.pollUnsafe()?._tag).toBe("Success") } }), ) @@ -159,19 +127,12 @@ it.instance( const other = { ...ctx, directory: `${ctx.directory}/other` } const id = SessionID.make("ses_isolated_drain") const release = yield* drain.hold(id).pipe(Effect.provideService(InstanceRef, other)) - let done = false - const waiting = yield* drain.wait(id).pipe( - Effect.provideService(InstanceRef, other), - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) + const waiting = yield* drain + .wait(id) + .pipe(Effect.provideService(InstanceRef, other), Effect.forkChild({ startImmediately: true })) yield* Effect.promise(() => disposeInstance(ctx.directory)) yield* Effect.yieldNow - expect(done).toBe(false) + expect(waiting.pollUnsafe()).toBeUndefined() release() yield* Fiber.join(waiting) }), @@ -202,23 +163,16 @@ it.instance( }), ) const first = yield* drain.hold(id) - let done = false - const waiting = yield* drain.wait(id).pipe( - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - Effect.provideService(Scheduler.Scheduler, scheduler), - ) + const waiting = yield* drain + .wait(id) + .pipe(Effect.forkChild({ startImmediately: true }), Effect.provideService(Scheduler.Scheduler, scheduler)) flush() paused = true first() const second = yield* drain.hold(id) paused = false flush() - expect(done).toBe(false) + expect(waiting.pollUnsafe()).toBeUndefined() second() flush() yield* Fiber.join(waiting) @@ -235,21 +189,13 @@ it.instance( Effect.flatMap(() => drain.hold(id)), Effect.forkChild({ startImmediately: true }), ) - let done = false - const second = yield* drain.wait(id).pipe( - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) + const waiting = yield* drain.wait(id).pipe(Effect.forkChild({ startImmediately: true })) release() const releaseAgain = yield* Fiber.join(first) - expect(done).toBe(false) + expect(waiting.pollUnsafe()).toBeUndefined() releaseAgain() - yield* Fiber.join(second) - expect(done).toBe(true) + yield* Fiber.join(waiting) + expect(waiting.pollUnsafe()?._tag).toBe("Success") }), ) @@ -263,16 +209,8 @@ it.instance( yield* drain.track(child, Effect.void) yield* drain.wait(parent) const release = yield* drain.hold(child) - let done = false - const waiting = yield* drain.wait(parent).pipe( - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) - expect(done).toBe(false) + const waiting = yield* drain.wait(parent).pipe(Effect.forkChild({ startImmediately: true })) + expect(waiting.pollUnsafe()).toBeUndefined() release() yield* Fiber.join(waiting) }), diff --git a/packages/opencode/test/kilocode/task-nesting.test.ts b/packages/opencode/test/kilocode/task-nesting.test.ts index a25950aa23c5..ce9087774cf1 100644 --- a/packages/opencode/test/kilocode/task-nesting.test.ts +++ b/packages/opencode/test/kilocode/task-nesting.test.ts @@ -1,7 +1,7 @@ import { LayerNode } from "@opencode-ai/core/effect/layer-node" import { SessionProjector } from "@opencode-ai/core/session/projector" import { afterEach, describe, expect, test } from "bun:test" -import { Effect, Exit } from "effect" +import { Deferred, Effect, Exit, Fiber } from "effect" import { Database } from "@opencode-ai/core/database/database" import { Agent } from "../../src/agent/agent" import { BackgroundJob } from "../../src/background/job" @@ -25,6 +25,7 @@ import { ModelV2 } from "@opencode-ai/core/model" import { Provider } from "../../src/provider/provider" import { Permission } from "../../src/permission" import { TaskTool, type TaskPromptOps } from "../../src/tool/task" +import type { Context } from "../../src/tool/tool" import { KiloSessionPrompt } from "../../src/kilocode/session/prompt" import * as SandboxPolicy from "../../src/kilocode/sandbox/policy" import { Truncate } from "../../src/tool/truncate" @@ -107,7 +108,7 @@ async function script(dir: string) { return command } -function stubOps(opts?: { onPrompt?: (input: SessionPrompt.PromptInput) => void }): TaskPromptOps { +function stubOps(opts?: { onPrompt?: (input: SessionPrompt.PromptInput) => void; text?: string }): TaskPromptOps { const prompt = (input: SessionPrompt.PromptInput) => Effect.sync(() => { opts?.onPrompt?.(input) @@ -134,7 +135,7 @@ function stubOps(opts?: { onPrompt?: (input: SessionPrompt.PromptInput) => void messageID: id, sessionID: input.sessionID, type: "text", - text: "done", + text: opts?.text ?? "done", }, ], } satisfies MessageV2.WithParts @@ -533,3 +534,121 @@ describe("Kilo task nesting", () => { ), ) }) + +describe("Kilo task drain", () => { + function task(message: MessageV2.Assistant, ops: TaskPromptOps, metadata: Context["metadata"] = () => Effect.void) { + return Effect.gen(function* () { + const tool = yield* (yield* TaskTool).init() + return yield* tool.execute( + { + description: "Drain child", + prompt: "Return the completed result", + subagent_type: "general", + background: false, + }, + { + sessionID: message.sessionID, + messageID: message.id, + agent: message.agent, + abort: new AbortController().signal, + extra: { promptOps: ops }, + messages: [], + metadata, + ask: () => Effect.void, + }, + ) + }) + } + + it.instance("returns the child's final answer after nested background delivery drains", () => + Effect.gen(function* () { + const sessions = yield* Session.Service + const drain = yield* SessionDrain.Service + const { assistant } = yield* seed() + const entered = yield* Deferred.make<{ input: SessionPrompt.PromptInput; release: () => void }>() + const ops: TaskPromptOps = { + ...stubOps(), + prompt: (input) => + Effect.gen(function* () { + const grandchild = yield* sessions.create({ parentID: input.sessionID, title: "nested background" }) + yield* drain.link(grandchild.id, input.sessionID) + const release = yield* drain.hold(grandchild.id) + const initial = yield* stubOps({ text: "WAITING_FOR_NESTED_CHILD" }).prompt(input) + yield* sessions.updateMessage(initial.info) + for (const part of initial.parts) yield* sessions.updatePart(part) + yield* Deferred.succeed(entered, { input, release }) + return initial + }), + } + const running = yield* task(assistant, ops).pipe(Effect.forkChild) + const pending = yield* Deferred.await(entered) + const final = yield* stubOps({ text: "FINAL_NESTED_RESULT" }).prompt(pending.input) + yield* sessions.updateMessage(final.info) + for (const part of final.parts) yield* sessions.updatePart(part) + pending.release() + const result = yield* Fiber.join(running) + expect(result.output).toContain("FINAL_NESTED_RESULT") + expect(result.output).not.toContain("WAITING_FOR_NESTED_CHILD") + }), + ) + + it.instance("promotion installs delivery cleanup even when its metadata hook fails", () => + Effect.gen(function* () { + const background = yield* BackgroundJob.Service + const drain = yield* SessionDrain.Service + const { chat, assistant } = yield* seed() + const child = yield* Deferred.make() + const releaseChild = yield* Deferred.make() + const metadata = yield* Deferred.make() + const releaseMetadata = yield* Deferred.make() + const notification = yield* Deferred.make() + const releaseNotification = yield* Deferred.make() + yield* Effect.addFinalizer(() => + Effect.sync(() => { + for (const gate of [releaseChild, releaseMetadata, releaseNotification]) + Deferred.doneUnsafe(gate, Effect.void) + }), + ) + let notifications = 0 + const ops: TaskPromptOps = { + ...stubOps(), + prompt: (input) => + Effect.gen(function* () { + if (input.sessionID === chat.id) { + notifications++ + yield* Deferred.succeed(notification, undefined) + yield* Deferred.await(releaseNotification) + } else { + yield* Deferred.succeed(child, input.sessionID) + yield* Deferred.await(releaseChild) + } + return yield* stubOps({ text: "completed" }).prompt(input) + }), + } + const running = yield* task(assistant, ops, (input) => + input.metadata?.background === true + ? Effect.gen(function* () { + yield* Deferred.succeed(metadata, undefined) + yield* Deferred.await(releaseMetadata) + return yield* Effect.die(new Error("metadata hook failed")) + }) + : Effect.void, + ).pipe(Effect.forkChild) + const childID = yield* Deferred.await(child) + const promotion = yield* background.promote(childID).pipe(Effect.exit, Effect.forkChild) + yield* Deferred.await(metadata) + const result = yield* Fiber.join(running) + expect(result.metadata.background).toBe(true) + const waiting = yield* drain.wait(chat.id).pipe(Effect.forkChild({ startImmediately: true })) + yield* Deferred.succeed(releaseChild, undefined) + yield* Deferred.await(notification) + expect(waiting.pollUnsafe()).toBeUndefined() + yield* Deferred.succeed(releaseNotification, undefined) + yield* Fiber.join(waiting) + yield* Deferred.succeed(releaseMetadata, undefined) + expect(Exit.isFailure(yield* Fiber.join(promotion))).toBe(true) + expect(notifications).toBe(1) + yield* drain.wait(chat.id) + }), + ) +}) diff --git a/packages/opencode/test/tool/task.test.ts b/packages/opencode/test/tool/task.test.ts index ca2cb88fa535..12cdeb1fa06b 100644 --- a/packages/opencode/test/tool/task.test.ts +++ b/packages/opencode/test/tool/task.test.ts @@ -1200,8 +1200,6 @@ describe("tool.task", () => { background.instance("background task completion does not wait for the parent async prompt", () => Effect.gen(function* () { const jobs = yield* BackgroundJob.Service - const drain = yield* SessionDrain.Service // kilocode_change - const finish = yield* Deferred.make() // kilocode_change const { chat, assistant } = yield* seed() const tool = yield* TaskTool const def = yield* tool.init() @@ -1221,12 +1219,8 @@ describe("tool.task", () => { extra: { promptOps: { ...stubOps({ text: "background done" }), - // kilocode_change start prompt: (input) => - input.sessionID === chat.id - ? Deferred.await(finish).pipe(Effect.as(reply(input, "parent done"))) - : Effect.succeed(reply(input, "background done")), - // kilocode_change end + input.sessionID === chat.id ? Effect.never : Effect.succeed(reply(input, "background done")), } satisfies TaskPromptOps, }, messages: [], @@ -1238,21 +1232,6 @@ describe("tool.task", () => { const waited = yield* jobs.wait({ id: result.metadata.sessionId, timeout: 1_000 }) expect(waited.timedOut).toBe(false) expect(waited.info?.status).toBe("completed") - // kilocode_change start - let done = false - const waiting = yield* drain.wait(chat.id).pipe( - Effect.tap(() => - Effect.sync(() => { - done = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) - expect(done).toBe(false) - yield* Deferred.succeed(finish, undefined) - yield* Fiber.join(waiting) - expect(done).toBe(true) - // kilocode_change end }), ) @@ -1425,141 +1404,6 @@ describe("tool.task", () => { }) // kilocode_change start - subagent cost propagation coverage (#6321) -it.instance("returns the child's final answer after nested background delivery drains", () => - Effect.gen(function* () { - const sessions = yield* Session.Service - const drain = yield* SessionDrain.Service - const { chat, assistant } = yield* seed() - const entered = yield* Deferred.make<{ input: SessionPrompt.PromptInput; release: () => void }>() - const factory = yield* TaskTool - const tool = yield* factory.init() - const ops: TaskPromptOps = { - ...stubOps(), - prompt: (input) => - Effect.gen(function* () { - const grandchild = yield* sessions.create({ parentID: input.sessionID, title: "nested background" }) - yield* drain.link(grandchild.id, input.sessionID) - const release = yield* drain.hold(grandchild.id) - const initial = reply(input, "WAITING_FOR_NESTED_CHILD") - yield* sessions.updateMessage(initial.info) - for (const part of initial.parts) yield* sessions.updatePart(part) - yield* Deferred.succeed(entered, { input, release }) - return initial - }), - } - const running = yield* tool - .execute( - { - description: "Nested result", - prompt: "Return the completed result", - subagent_type: "general", - background: false, - }, - { - sessionID: chat.id, - messageID: assistant.id, - agent: "build", - abort: new AbortController().signal, - extra: { promptOps: ops }, - messages: [], - metadata: () => Effect.void, - ask: () => Effect.void, - }, - ) - .pipe(Effect.forkChild) - const pending = yield* Deferred.await(entered) - const final = reply(pending.input, "FINAL_NESTED_RESULT") - yield* sessions.updateMessage(final.info) - for (const part of final.parts) yield* sessions.updatePart(part) - pending.release() - const result = yield* Fiber.join(running) - expect(result.output).toContain("FINAL_NESTED_RESULT") - expect(result.output).not.toContain("WAITING_FOR_NESTED_CHILD") - }), -) - -it.instance("promotion installs delivery cleanup even when its metadata hook fails", () => - Effect.gen(function* () { - const background = yield* BackgroundJob.Service - const drain = yield* SessionDrain.Service - const { chat, assistant } = yield* seed() - const child = yield* Deferred.make() - const releaseChild = yield* Deferred.make() - const metadata = yield* Deferred.make() - const releaseMetadata = yield* Deferred.make() - const notification = yield* Deferred.make() - const releaseNotification = yield* Deferred.make() - yield* Effect.addFinalizer(() => - Effect.sync(() => { - for (const gate of [releaseChild, releaseMetadata, releaseNotification]) Deferred.doneUnsafe(gate, Effect.void) - }), - ) - let notifications = 0 - const factory = yield* TaskTool - const tool = yield* factory.init() - const ops: TaskPromptOps = { - ...stubOps(), - prompt: (input) => - Effect.gen(function* () { - if (input.sessionID === chat.id) { - notifications++ - yield* Deferred.succeed(notification, undefined) - yield* Deferred.await(releaseNotification) - } else { - yield* Deferred.succeed(child, input.sessionID) - yield* Deferred.await(releaseChild) - } - return reply(input, "completed") - }), - } - const running = yield* tool - .execute( - { description: "Promote", prompt: "Do work", subagent_type: "general", background: false }, - { - sessionID: chat.id, - messageID: assistant.id, - agent: "build", - abort: new AbortController().signal, - extra: { promptOps: ops }, - messages: [], - ask: () => Effect.void, - metadata: (input) => - input.metadata?.background === true - ? Effect.gen(function* () { - yield* Deferred.succeed(metadata, undefined) - yield* Deferred.await(releaseMetadata) - return yield* Effect.die(new Error("metadata hook failed")) - }) - : Effect.void, - }, - ) - .pipe(Effect.forkChild) - const childID = yield* Deferred.await(child) - const promotion = yield* background.promote(childID).pipe(Effect.exit, Effect.forkChild) - yield* Deferred.await(metadata) - const result = yield* Fiber.join(running) - expect(result.metadata.background).toBe(true) - let drained = false - const waiting = yield* drain.wait(chat.id).pipe( - Effect.tap(() => - Effect.sync(() => { - drained = true - }), - ), - Effect.forkChild({ startImmediately: true }), - ) - yield* Deferred.succeed(releaseChild, undefined) - yield* Deferred.await(notification) - expect(drained).toBe(false) - yield* Deferred.succeed(releaseNotification, undefined) - yield* Fiber.join(waiting) - yield* Deferred.succeed(releaseMetadata, undefined) - expect(Exit.isFailure(yield* Fiber.join(promotion))).toBe(true) - expect(notifications).toBe(1) - yield* drain.wait(chat.id) - }), -) - const assistantCost = Effect.fn("TaskToolTest.assistantCost")(function* (sessionID: string) { const sessions = yield* Session.Service const msgs = yield* sessions.messages({ sessionID: SessionID.make(sessionID) }) diff --git a/packages/schema/src/kilocode/session-drain.ts b/packages/schema/src/kilocode/session-drain.ts index ef481e073448..73c6eff22502 100644 --- a/packages/schema/src/kilocode/session-drain.ts +++ b/packages/schema/src/kilocode/session-drain.ts @@ -9,4 +9,9 @@ export const Drained = Event.define({ schema: { sessionID: SessionID, token: Token }, }) -export const Definitions = Event.inventory(Drained) +export const Interrupted = Event.define({ + type: "session.drain.interrupted", + schema: { sessionID: SessionID }, +}) + +export const Definitions = Event.inventory(Drained, Interrupted) diff --git a/packages/sdk/js/src/v2/gen/types.gen.ts b/packages/sdk/js/src/v2/gen/types.gen.ts index 2cd1e7892aa8..212795e977a0 100644 --- a/packages/sdk/js/src/v2/gen/types.gen.ts +++ b/packages/sdk/js/src/v2/gen/types.gen.ts @@ -96,6 +96,7 @@ export type Event = | EventGlobalDisposed1 | EventGlobalConfigUpdated1 | EventSessionDrained1 + | EventSessionDrainInterrupted1 | EventServerInstanceDisposed | EventSessionTurnOpen | EventSessionTurnClose @@ -216,6 +217,7 @@ export type Event = | EventGlobalDisposed | EventGlobalConfigUpdated | EventSessionDrained + | EventSessionDrainInterrupted export type QuestionReplied = { sessionID: string @@ -1275,6 +1277,7 @@ export type GlobalEvent = { | EventGlobalDisposed | EventGlobalConfigUpdated | EventSessionDrained + | EventSessionDrainInterrupted | { id: string type: "models-dev.refreshed" @@ -2168,6 +2171,13 @@ export type GlobalEvent = { token: string } } + | { + id: string + type: "session.drain.interrupted" + properties: { + sessionID: string + } + } | SyncEventSessionCreated | SyncEventSessionUpdated | SyncEventSessionDeleted @@ -4844,6 +4854,7 @@ export type V2Event = | GlobalDisposed | GlobalConfigUpdated | SessionDrained + | SessionDrainInterrupted export type V2EventStream = string @@ -6343,6 +6354,14 @@ export type EventSessionDrained = { } } +export type EventSessionDrainInterrupted = { + id: string + type: "session.drain.interrupted" + properties: { + sessionID: string + } +} + export type SyncEventSessionCreated = { type: "sync" id: string @@ -9335,6 +9354,23 @@ export type SessionDrained = { } } +export type SessionDrainInterrupted = { + id: string + metadata?: { + [key: string]: unknown + } + type: "session.drain.interrupted" + durable?: { + aggregateID: string + seq: number + version: number + } + location?: LocationRef + data: { + sessionID: string + } +} + export type QuestionV2Request = { id: string sessionID: string @@ -10365,6 +10401,14 @@ export type EventSessionDrained1 = { } } +export type EventSessionDrainInterrupted1 = { + id: string + type: "session.drain.interrupted" + properties: { + sessionID: string + } +} + export type EventMemoryStatus1 = { id: string type: "memory.status" diff --git a/packages/sdk/openapi.json b/packages/sdk/openapi.json index 9e67fd6a7dfa..f0c216321a55 100644 --- a/packages/sdk/openapi.json +++ b/packages/sdk/openapi.json @@ -26981,6 +26981,9 @@ { "$ref": "#/components/schemas/EventSessionDrained1" }, + { + "$ref": "#/components/schemas/EventSessionDrainInterrupted1" + }, { "$ref": "#/components/schemas/EventServerInstanceDisposed" }, @@ -27341,6 +27344,9 @@ { "$ref": "#/components/schemas/EventSessionDrained" }, + { + "$ref": "#/components/schemas/EventSessionDrainInterrupted" + }, { "$ref": "#/components/schemas/EventServerInstanceDisposed" } @@ -30541,6 +30547,9 @@ { "$ref": "#/components/schemas/EventSessionDrained" }, + { + "$ref": "#/components/schemas/EventSessionDrainInterrupted" + }, { "type": "object", "properties": { @@ -33627,6 +33636,32 @@ "required": ["id", "type", "properties"], "additionalProperties": false }, + { + "type": "object", + "properties": { + "id": { + "type": "string", + "pattern": "^evt_" + }, + "type": { + "type": "string", + "enum": ["session.drain.interrupted"] + }, + "properties": { + "type": "object", + "properties": { + "sessionID": { + "type": "string", + "pattern": "^ses" + } + }, + "required": ["sessionID"], + "additionalProperties": false + } + }, + "required": ["id", "type", "properties"], + "additionalProperties": false + }, { "$ref": "#/components/schemas/EventServerInstanceDisposed" }, @@ -41624,6 +41659,9 @@ }, { "$ref": "#/components/schemas/SessionDrained" + }, + { + "$ref": "#/components/schemas/SessionDrainInterrupted" } ] }, @@ -46846,6 +46884,31 @@ "required": ["id", "type", "properties"], "additionalProperties": false }, + "EventSessionDrainInterrupted": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "type": { + "type": "string", + "enum": ["session.drain.interrupted"] + }, + "properties": { + "type": "object", + "properties": { + "sessionID": { + "type": "string", + "pattern": "^ses" + } + }, + "required": ["sessionID"], + "additionalProperties": false + } + }, + "required": ["id", "type", "properties"], + "additionalProperties": false + }, "SyncEventSessionCreated": { "type": "object", "properties": { @@ -55854,6 +55917,54 @@ "required": ["id", "type", "data"], "additionalProperties": false }, + "SessionDrainInterrupted": { + "type": "object", + "properties": { + "id": { + "type": "string", + "pattern": "^evt_" + }, + "metadata": { + "type": "object" + }, + "type": { + "type": "string", + "enum": ["session.drain.interrupted"] + }, + "durable": { + "type": "object", + "properties": { + "aggregateID": { + "type": "string" + }, + "seq": { + "type": "integer" + }, + "version": { + "type": "integer" + } + }, + "required": ["aggregateID", "seq", "version"], + "additionalProperties": false + }, + "location": { + "$ref": "#/components/schemas/LocationRef" + }, + "data": { + "type": "object", + "properties": { + "sessionID": { + "type": "string", + "pattern": "^ses" + } + }, + "required": ["sessionID"], + "additionalProperties": false + } + }, + "required": ["id", "type", "data"], + "additionalProperties": false + }, "QuestionV2Request": { "type": "object", "properties": { @@ -59056,6 +59167,32 @@ "required": ["id", "type", "properties"], "additionalProperties": false }, + "EventSessionDrainInterrupted1": { + "type": "object", + "properties": { + "id": { + "type": "string", + "pattern": "^evt_" + }, + "type": { + "type": "string", + "enum": ["session.drain.interrupted"] + }, + "properties": { + "type": "object", + "properties": { + "sessionID": { + "type": "string", + "pattern": "^ses" + } + }, + "required": ["sessionID"], + "additionalProperties": false + } + }, + "required": ["id", "type", "properties"], + "additionalProperties": false + }, "EventMemoryStatus1": { "type": "object", "properties": {