Merge pull request #12271 from Kilo-Org/agent-manager-stop-session

feat(agent-manager): stop managed sessions from tool
This commit is contained in:
Marius
2026-07-17 11:01:10 +02:00
committed by GitHub
17 changed files with 312 additions and 40 deletions
+6
View File
@@ -0,0 +1,6 @@
---
"@kilocode/cli": minor
"kilo-code": minor
---
Allow agents to stop and remove a targeted Agent Manager session.
@@ -194,6 +194,11 @@ export class AgentManagerProvider implements Disposable {
},
stats: (refresh) => this.statsPoller.snapshot(refresh),
prs: () => this.prBridge.snapshot(),
managed: (id) => this.panelSessions.has(id) || !!this.state?.getSession(id),
close: async (id) => {
await this.onCloseSession(id)
this.postToWebview({ type: "agentManager.sessionClosed", sessionId: id })
},
log: (...args) => this.log(...args),
})
this.unsubTool = this.connectionService.onEventFiltered(
@@ -25,10 +25,12 @@ interface RequestBase {
type Request =
| (RequestBase & { operation: "overview"; filter?: OverviewFilter })
| (RequestBase & { operation: "prompt"; targetSessionID: string; prompt: string })
| (RequestBase & { operation: "stop"; targetSessionID: string })
type Result =
| { operation: "overview"; overview: Overview }
| { operation: "prompt"; sessionID: string; delivered: true }
| { operation: "stop"; sessionID: string; stopped: true }
interface Failure {
code: FailureCode | "cancelled" | "disconnected" | "timeout"
@@ -41,6 +43,8 @@ interface Options {
state(): WorktreeStateManager | undefined
stats(refresh?: boolean): Promise<{ worktrees: WorktreeStats[]; local?: LocalStats }>
prs(): Map<string, PRStatus>
managed(sessionID: string): boolean
close(sessionID: string): Promise<void>
log(...args: unknown[]): void
}
@@ -262,17 +266,25 @@ export class AgentManagerOrchestrationBridge {
})
return { result: { operation: "overview", overview: result } }
}
await prompt({
client,
root,
state,
sessionID: request.targetSessionID,
text: request.prompt,
messageID: request.id,
signal: active.controller.signal,
})
if (request.operation === "prompt") {
await prompt({
client,
root,
state,
sessionID: request.targetSessionID,
text: request.prompt,
messageID: request.id,
signal: active.controller.signal,
})
if (this.disposed || active.cancelled) return
return { result: { operation: "prompt", sessionID: request.targetSessionID, delivered: true } }
}
if (!this.options.managed(request.targetSessionID)) {
throw new OrchestrationError("unknown_session", "The session is not managed by this Agent Manager workspace")
}
await this.options.close(request.targetSessionID)
if (this.disposed || active.cancelled) return
return { result: { operation: "prompt", sessionID: request.targetSessionID, delivered: true } }
return { result: { operation: "stop", sessionID: request.targetSessionID, stopped: true } }
} catch (error) {
if (this.disposed || active.cancelled) return
return { error: failure(error) }
@@ -190,6 +190,11 @@ interface SessionForkedMessage {
worktreeId?: string
}
interface SessionClosedMessage {
type: "agentManager.sessionClosed"
sessionId: string
}
interface MultiVersionProgressMessage {
type: "agentManager.multiVersionProgress"
status: "creating" | "done"
@@ -307,6 +312,7 @@ export type AgentManagerOutMessage =
| ErrorOutMessage
| SessionAddedMessage
| SessionForkedMessage
| SessionClosedMessage
| MultiVersionProgressMessage
| SetSessionModelMessage
| SendInitialMessage
@@ -5,6 +5,7 @@ const { AgentManagerProvider } = await import("../../src/agent-manager/AgentMana
type Manager = {
connectionService: { getClient: () => unknown }
panel: {
postMessage: (message: unknown) => void
sessions: {
getSessionDirectories: () => ReadonlyMap<string, string>
clearSessionDirectory: (id: string) => void
@@ -24,6 +25,7 @@ function createManager(options?: { dir?: string; panelDir?: string; state?: bool
const aborted: string[][] = []
const cleared: string[] = []
const removed: string[] = []
const messages: unknown[] = []
const events: string[] = []
const client = {
backgroundProcess: {
@@ -44,6 +46,7 @@ function createManager(options?: { dir?: string; panelDir?: string; state?: bool
const manager = Object.create(AgentManagerProvider.prototype) as Manager
manager.connectionService = { getClient: () => client }
manager.panel = {
postMessage: (message) => messages.push(message),
sessions: {
getSessionDirectories: () => new Map(options?.panelDir ? [["s1", options.panelDir]] : []),
clearSessionDirectory: (id) => cleared.push(id),
@@ -59,12 +62,12 @@ function createManager(options?: { dir?: string; panelDir?: string; state?: bool
manager.pushState = mock(() => undefined)
manager.log = mock(() => undefined)
return { manager, stopped, aborted, cleared, removed, events }
return { manager, stopped, aborted, cleared, removed, messages, events }
}
describe("AgentManagerProvider closeSession", () => {
it("aborts the agent before stopping processes and removing its tab", async () => {
const { manager, stopped, aborted, cleared, removed, events } = createManager({ dir: "/repo/worktree" })
const { manager, stopped, aborted, cleared, removed, messages, events } = createManager({ dir: "/repo/worktree" })
await manager.onCloseSession("s1")
@@ -73,6 +76,7 @@ describe("AgentManagerProvider closeSession", () => {
expect(events).toEqual(["abort", "processes", "remove"])
expect(removed).toEqual(["s1"])
expect(cleared).toEqual(["s1"])
expect(messages).toEqual([])
expect(manager.panelSessions.has("s1")).toBe(false)
})
@@ -42,7 +42,9 @@ describe("AgentManagerOrchestrationBridge", () => {
state?: (state: "connecting" | "connected" | "disconnected" | "error") => void
} = {}
const status = { failList: "", failReply: false }
const managed = new Set(["ses_target"])
const promptAsync = mock(async () => ({ data: undefined }))
const close = mock(async () => undefined)
const client = {
session: {
get: mock(async () => ({
@@ -95,6 +97,8 @@ describe("AgentManagerOrchestrationBridge", () => {
state: () => state,
stats: async () => ({ worktrees: [] }),
prs: () => new Map(),
managed: (id) => managed.has(id),
close,
log: () => undefined,
})
const request = (value: AgentManagerRequest, directory = root) =>
@@ -102,7 +106,7 @@ describe("AgentManagerOrchestrationBridge", () => {
{ id: `event-${value.id}`, type: "kilocode.agent_manager.requested", properties: value } as SSEPayload,
directory,
)
return { bridge, client, handlers, lists, promptAsync, rejections, replies, request, status }
return { bridge, client, close, handlers, lists, managed, promptAsync, rejections, replies, request, status }
}
const request: AgentManagerRequest = {
@@ -148,6 +152,85 @@ describe("AgentManagerOrchestrationBridge", () => {
test.bridge.dispose()
})
it("stops a managed session through the same close operation as the UI", async () => {
const test = harness()
test.status.failReply = true
const stop: AgentManagerRequest = {
id: "amr_stop",
sessionID: "ses_caller",
operation: "stop",
targetSessionID: "ses_target",
}
test.request(stop)
await waitFor(() => test.replies.length === 1)
test.status.failReply = false
test.request(stop)
await waitFor(() => test.replies.length === 2)
expect(test.close).toHaveBeenCalledTimes(1)
expect(test.close).toHaveBeenCalledWith("ses_target")
expect(test.replies).toEqual([
{
requestID: "amr_stop",
directory: root,
result: { operation: "stop", sessionID: "ses_target", stopped: true },
},
{
requestID: "amr_stop",
directory: root,
result: { operation: "stop", sessionID: "ses_target", stopped: true },
},
])
test.bridge.dispose()
})
it("stops a live panel session before it is persisted", async () => {
const test = harness()
test.managed.add("ses_live")
test.request({
id: "amr_stop_live",
sessionID: "ses_caller",
operation: "stop",
targetSessionID: "ses_live",
})
await waitFor(() => test.replies.length === 1)
expect(state.getSession("ses_live")).toBeUndefined()
expect(test.close).toHaveBeenCalledWith("ses_live")
expect(test.replies[0]).toEqual({
requestID: "amr_stop_live",
directory: root,
result: { operation: "stop", sessionID: "ses_live", stopped: true },
})
test.bridge.dispose()
})
it("rejects stopping a session not managed by the current workspace", async () => {
const test = harness()
test.request({
id: "amr_stop_unknown",
sessionID: "ses_caller",
operation: "stop",
targetSessionID: "ses_unknown",
})
await waitFor(() => test.rejections.length === 1)
expect(test.close).not.toHaveBeenCalled()
expect(test.rejections).toEqual([
{
requestID: "amr_stop_unknown",
directory: root,
error: {
code: "unknown_session",
message: "The session is not managed by this Agent Manager workspace",
},
},
])
test.bridge.dispose()
})
it("rejects request origins outside the current Agent Manager workspace", async () => {
const test = harness()
@@ -1221,8 +1221,8 @@ const AgentManagerContent: Component = () => {
// Mark sessions loaded as soon as the session context receives data (even if empty)
const unsubSessions = vscode.onMessage((msg) => {
if (msg.type === "sessionsLoaded" && !sessionsLoaded()) setSessionsLoaded(true)
if (msg.type === "agentManager.sessionClosed") handleCloseTab(msg.sessionId, false)
})
const unsubRun = vscode.onMessage((msg) => {
if (msg.type !== "agentManager.runStatus") return
const ev = msg as RunStatus
@@ -1979,7 +1979,7 @@ const AgentManagerContent: Component = () => {
if (!sel || sel === LOCAL) return vscode.postMessage(msg)
vscode.postMessage({ ...msg, worktreeId: sel })
}
const handleCloseTab = (sessionId: string) => {
const handleCloseTab = (sessionId: string, notify = true) => {
freezeTabs()
const pending = isPending(sessionId)
const isActive = pending ? sessionId === activePendingId() : session.currentSessionID() === sessionId
@@ -2009,7 +2009,7 @@ const AgentManagerContent: Component = () => {
if (session.isSubmitting(sessionId) || isPendingSend(sessionId)) discardPendingDraft(sessionId)
queueMicrotask(() => deletePendingDraft(sessionId))
}
vscode.postMessage({ type: "agentManager.closeSession", sessionId })
if (notify) vscode.postMessage({ type: "agentManager.closeSession", sessionId })
tabFocus.restore()
}
@@ -643,6 +643,11 @@ export interface AgentManagerSessionForkedMessage {
worktreeId?: string
}
export interface AgentManagerSessionClosedMessage {
type: "agentManager.sessionClosed"
sessionId: string
}
// Full state push from extension to webview
export interface AgentManagerStateMessage {
type: "agentManager.state"
@@ -1162,6 +1167,7 @@ export type ExtensionMessage =
| AgentManagerWorktreeSetupMessage
| AgentManagerSessionAddedMessage
| AgentManagerSessionForkedMessage
| AgentManagerSessionClosedMessage
| AgentManagerStateMessage
| AgentManagerRunStatusMessage
| AgentManagerKeybindingsMessage
@@ -99,7 +99,13 @@ export const PromptRequest = Schema.Struct({
prompt: Prompt,
}).annotate({ identifier: "AgentManagerPromptRequest" })
export const Request = Schema.Union([OverviewRequest, PromptRequest]).annotate({
export const StopRequest = Schema.Struct({
...Base,
operation: Schema.Literal("stop"),
targetSessionID: SessionID,
}).annotate({ identifier: "AgentManagerStopRequest" })
export const Request = Schema.Union([OverviewRequest, PromptRequest, StopRequest]).annotate({
identifier: "AgentManagerRequest",
})
export type Request = Schema.Schema.Type<typeof Request>
@@ -115,7 +121,13 @@ export const PromptResult = Schema.Struct({
delivered: Schema.Literal(true),
}).annotate({ identifier: "AgentManagerPromptResult" })
export const Result = Schema.Union([OverviewResult, PromptResult]).annotate({
export const StopResult = Schema.Struct({
operation: Schema.Literal("stop"),
sessionID: SessionID,
stopped: Schema.Literal(true),
}).annotate({ identifier: "AgentManagerStopResult" })
export const Result = Schema.Union([OverviewResult, PromptResult, StopResult]).annotate({
identifier: "AgentManagerResult",
})
export type Result = Schema.Schema.Type<typeof Result>
@@ -37,7 +37,7 @@ interface State {
function matches(request: Request, result: Result) {
if (request.operation === "overview") return result.operation === "overview"
return result.operation === "prompt" && result.sessionID === request.targetSessionID
return result.operation === request.operation && result.sessionID === request.targetSessionID
}
export interface Interface {
@@ -2,12 +2,12 @@ import { type Rule } from "./rule"
export namespace AgentManagerPermission {
/**
* Prompting an existing Agent Manager session has an external side effect.
* Prompting or stopping an existing Agent Manager session has an external side effect.
* Broad approvals for legacy session creation must not silently grant it.
*/
export function harden(permission: string, pattern: string, rule: Rule): Rule {
if (permission !== "agent_manager" || pattern !== "prompt" || rule.action !== "allow") return rule
if (rule.permission === "agent_manager" && rule.pattern === "prompt") return rule
if (permission !== "agent_manager" || !["prompt", "stop"].includes(pattern) || rule.action !== "allow") return rule
if (rule.permission === "agent_manager" && rule.pattern === pattern) return rule
return { permission, pattern, action: "ask" }
}
}
@@ -75,13 +75,18 @@ const PromptParams = Schema.Struct({
),
})
export const Params = Schema.Union([StartParams, ListParams, PromptParams])
const StopParams = Schema.Struct({
action: Schema.Literal("stop"),
sessionID: SessionID,
})
export const Params = Schema.Union([StartParams, ListParams, PromptParams, StopParams])
const WireParams = Schema.Struct({
mode: Schema.optional(StartParams.fields.mode),
versions: Schema.optional(StartParams.fields.versions),
tasks: Schema.optional(StartParams.fields.tasks),
action: Schema.optional(Schema.Literals(["list", "prompt"])),
action: Schema.optional(Schema.Literals(["list", "prompt", "stop"])),
filter: Schema.optional(ListParams.fields.filter),
sessionID: Schema.optional(PromptParams.fields.sessionID),
prompt: Schema.optional(PromptParams.fields.prompt),
@@ -236,7 +241,7 @@ function select(
export const AgentManagerTool = Tool.define<
typeof Params,
{ action: "start" | "list" | "prompt"; requestID?: string; count?: number; sessionID?: string },
{ action: "start" | "list" | "prompt" | "stop"; requestID?: string; count?: number; sessionID?: string },
AgentManager.Service | Bus.Service | Provider.Service,
"agent_manager"
>(
@@ -275,27 +280,50 @@ export const AgentManagerTool = Tool.define<
metadata: { action: "list", count },
}
}
if (params.action === "prompt") {
yield* ctx.ask({
permission: "agent_manager",
patterns: ["prompt"],
always: ["prompt"],
metadata: { action: "prompt", sessionID: params.sessionID },
})
const result = yield* run(
host.request({
operation: "prompt",
sessionID: ctx.sessionID,
targetSessionID: params.sessionID,
prompt: params.prompt.trim(),
}),
ctx.abort,
)
if (result.operation !== "prompt")
return yield* Effect.die(new Error("Agent Manager host returned the wrong result type"))
return {
title: "Prompt delivered",
output: `Delivered the prompt to Agent Manager session ${result.sessionID}. The session accepted it asynchronously.`,
metadata: { action: "prompt", sessionID: result.sessionID },
}
}
yield* ctx.ask({
permission: "agent_manager",
patterns: ["prompt"],
always: ["prompt"],
metadata: { action: "prompt", sessionID: params.sessionID },
patterns: ["stop"],
always: ["stop"],
metadata: { action: "stop", sessionID: params.sessionID },
})
const result = yield* run(
host.request({
operation: "prompt",
operation: "stop",
sessionID: ctx.sessionID,
targetSessionID: params.sessionID,
prompt: params.prompt.trim(),
}),
ctx.abort,
)
if (result.operation !== "prompt")
if (result.operation !== "stop")
return yield* Effect.die(new Error("Agent Manager host returned the wrong result type"))
return {
title: "Prompt delivered",
output: `Delivered the prompt to Agent Manager session ${result.sessionID}. The session accepted it asynchronously.`,
metadata: { action: "prompt", sessionID: result.sessionID },
title: "Session stopped",
output: `Stopped Agent Manager session ${result.sessionID} and removed it from Agent Manager.`,
metadata: { action: "stop", sessionID: result.sessionID },
}
}
@@ -1,6 +1,6 @@
Inspect and orchestrate Agent Manager sessions, or start new sessions, in the VS Code extension.
Use `action: "list"` to inspect the compact Agent Manager overview and `action: "prompt"` to send one instruction to one existing managed session. List results include user-defined sections, ungrouped worktrees, and managed local sessions. Optional filters can narrow by section ID or by `idle`, `busy`, `retry`, `offline`, or `waiting` state. Prompting is targeted only: it does not broadcast, create a session, or wait for the target to finish.
Use `action: "list"` to inspect the compact Agent Manager overview, `action: "prompt"` to send one instruction to one existing managed session, and `action: "stop"` to stop and remove one managed session. List results include user-defined sections, ungrouped worktrees, and managed local sessions. Optional filters can narrow by section ID or by `idle`, `busy`, `retry`, `offline`, or `waiting` state. Prompting and stopping are targeted only: they do not broadcast or create sessions, and prompting does not wait for the target to finish.
To start sessions, keep using the existing `mode` and `tasks` input without an action. Use start mode when the user explicitly asks you to fan out work into Agent Manager, create Agent Manager worktrees, or start multiple Agent Manager sessions for independent tasks.
@@ -158,6 +158,8 @@ describe("agent_manager tool", () => {
expect(schema.anyOf).toBeUndefined()
expect(schema.oneOf).toBeUndefined()
expect(schema.allOf).toBeUndefined()
const action = schema.properties?.action
expect(action && typeof action === "object" ? action.enum : undefined).toEqual(["list", "prompt", "stop"])
expect(Object.keys(schema.properties ?? {})).toEqual([
"mode",
"versions",
@@ -299,6 +301,50 @@ describe("agent_manager tool", () => {
await rt.dispose()
})
test("stops one existing session with a separate mutation permission pattern", async () => {
const requests: unknown[] = []
const rt = makeRuntime("test", {
request: (input) =>
Effect.sync(() => {
requests.push(input)
return { operation: "stop" as const, sessionID: SessionID.make("ses_target"), stopped: true as const }
}),
})
const tool = await rt.runPromise(
Effect.gen(function* () {
return yield* Tool.init(yield* AgentManagerTool)
}),
)
const permissions: unknown[] = []
const result = await rt.runPromise(
provideTmpdirInstance(() =>
tool.execute(
{ action: "stop", sessionID: SessionID.make("ses_target") },
{ ...ctx, ask: (input: unknown) => Effect.sync(() => permissions.push(input)) },
),
).pipe(Effect.scoped),
)
expect(permissions).toEqual([
{
permission: "agent_manager",
patterns: ["stop"],
always: ["stop"],
metadata: { action: "stop", sessionID: "ses_target" },
},
])
expect(requests).toEqual([
{
operation: "stop",
sessionID: ctx.sessionID,
targetSessionID: "ses_target",
},
])
expect(result.output).toContain("removed it from Agent Manager")
expect(result.metadata).toEqual(expect.objectContaining({ action: "stop", sessionID: "ses_target" }))
await rt.dispose()
})
test("inherits the latest invoking model and variant when omitted", async () => {
const task = await publish(runtime, { prompt: "Fix" }, [
message("msg_current", "kilo", "kilo/shared", "low", 2),
@@ -3,9 +3,10 @@ import { Permission } from "../../../src/permission"
const broad = Permission.fromConfig({ agent_manager: "allow" })
describe("Agent Manager prompt permissions", () => {
describe("Agent Manager side-effect permissions", () => {
test("requires consent despite a broad Agent Manager allow rule", () => {
expect(Permission.resolve("agent_manager", "prompt", broad).action).toBe("ask")
expect(Permission.resolve("agent_manager", "stop", broad).action).toBe("ask")
expect(Permission.resolve("agent_manager", "local", broad).action).toBe("allow")
expect(Permission.resolve("agent_manager", "worktree", broad).action).toBe("allow")
})
@@ -13,16 +14,19 @@ describe("Agent Manager prompt permissions", () => {
test("requires consent despite a global allow rule", () => {
const rules = [{ permission: "*", pattern: "*", action: "allow" as const }]
expect(Permission.resolve("agent_manager", "prompt", rules).action).toBe("ask")
expect(Permission.resolve("agent_manager", "stop", rules).action).toBe("ask")
})
test("requires consent despite a saved wildcard approval", () => {
const rules = Permission.fromConfig({ agent_manager: "ask" })
const saved = [{ permission: "agent_manager", pattern: "*", action: "allow" as const }]
expect(Permission.resolve("agent_manager", "prompt", rules, saved).action).toBe("ask")
expect(Permission.resolve("agent_manager", "stop", rules, saved).action).toBe("ask")
})
test("allows only an explicit prompt approval", () => {
const rules = Permission.fromConfig({ agent_manager: { prompt: "allow" } })
test("allows only explicit side-effect approvals", () => {
const rules = Permission.fromConfig({ agent_manager: { prompt: "allow", stop: "allow" } })
expect(Permission.resolve("agent_manager", "prompt", rules).action).toBe("allow")
expect(Permission.resolve("agent_manager", "stop", rules).action).toBe("allow")
})
})
+15 -2
View File
@@ -273,7 +273,14 @@ export type AgentManagerPromptRequest = {
prompt: string
}
export type AgentManagerRequest = AgentManagerOverviewRequest | AgentManagerPromptRequest
export type AgentManagerStopRequest = {
id: AgentManagerRequestId
sessionID: string
operation: "stop"
targetSessionID: string
}
export type AgentManagerRequest = AgentManagerOverviewRequest | AgentManagerPromptRequest | AgentManagerStopRequest
export type NotebookRequestId = string
@@ -3311,7 +3318,13 @@ export type AgentManagerPromptResult = {
delivered: true
}
export type AgentManagerResult = AgentManagerOverviewResult | AgentManagerPromptResult
export type AgentManagerStopResult = {
operation: "stop"
sessionID: string
stopped: true
}
export type AgentManagerResult = AgentManagerOverviewResult | AgentManagerPromptResult | AgentManagerStopResult
export type AgentManagerFailure = {
code:
+47
View File
@@ -22883,6 +22883,28 @@
"required": ["id", "sessionID", "operation", "targetSessionID", "prompt"],
"additionalProperties": false
},
"AgentManagerStopRequest": {
"type": "object",
"properties": {
"id": {
"$ref": "#/components/schemas/AgentManagerRequestID"
},
"sessionID": {
"type": "string",
"pattern": "^ses"
},
"operation": {
"type": "string",
"enum": ["stop"]
},
"targetSessionID": {
"type": "string",
"pattern": "^ses"
}
},
"required": ["id", "sessionID", "operation", "targetSessionID"],
"additionalProperties": false
},
"AgentManagerRequest": {
"anyOf": [
{
@@ -22890,6 +22912,9 @@
},
{
"$ref": "#/components/schemas/AgentManagerPromptRequest"
},
{
"$ref": "#/components/schemas/AgentManagerStopRequest"
}
]
},
@@ -31977,6 +32002,25 @@
"required": ["operation", "sessionID", "delivered"],
"additionalProperties": false
},
"AgentManagerStopResult": {
"type": "object",
"properties": {
"operation": {
"type": "string",
"enum": ["stop"]
},
"sessionID": {
"type": "string",
"pattern": "^ses"
},
"stopped": {
"type": "boolean",
"enum": [true]
}
},
"required": ["operation", "sessionID", "stopped"],
"additionalProperties": false
},
"AgentManagerResult": {
"anyOf": [
{
@@ -31984,6 +32028,9 @@
},
{
"$ref": "#/components/schemas/AgentManagerPromptResult"
},
{
"$ref": "#/components/schemas/AgentManagerStopResult"
}
]
},