Merge pull request #13197 from Kilo-Org/fix-agent-manager-tool-schema

fix(agent-manager): make tool requests strict

.changeset/strict-agent-manager-tool-requests.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/cli": patch
+---
+
+Make Agent Manager tool requests use a strict operation union so starting sessions and managing existing sessions cannot be confused.

packages/opencode/src/kilocode/tool/agent-manager.ts

@@ -14,6 +14,18 @@ import { Effect, Schema } from "effect"
 import { matchesQuery } from "./model-search"
 import DESCRIPTION from "./agent-manager.txt"
 
+function strict<const Fields extends Schema.Struct.Fields>(fields: Fields) {
+  const target = Schema.Struct(fields)
+  // Preserve unknown keys long enough for the branch check to reject mixed operations.
+  const source = Schema.StructWithRest(target, [Schema.Record(Schema.String, Schema.Unknown)]).check(
+    Schema.makeFilter((value) => {
+      const extra = Object.keys(value).find((key) => !Object.hasOwn(fields, key))
+      return extra === undefined ? undefined : `Unexpected Agent Manager parameter: ${extra}`
+    }),
+  )
+  return source.pipe(Schema.decodeTo(target))
+}
+
 const Task = Schema.Struct({
   prompt: Schema.optional(Schema.NullOr(Schema.String)).annotate({
     description: "Initial prompt to send to the new session",
@@ -46,7 +58,35 @@ const Task = Schema.Struct({
   ),
 )
 
-const StartParams = Schema.Struct({
+function wireSchema() {
+  const schema = structuredClone(ToolJsonSchema.fromSchema(Params))
+
+  // llama.cpp rejects the prefix-only SessionID pattern. Keep the runtime brand
+  // check, but omit that provider-incompatible hint from the advertised schema.
+  function strip(value: unknown): void {
+    if (Array.isArray(value)) {
+      value.forEach(strip)
+      return
+    }
+    if (!value || typeof value !== "object") return
+    const item = value as Record<string, unknown>
+    if (item.type === "object" && item.additionalProperties === undefined) {
+      item.additionalProperties = false
+    }
+    if (item.properties && typeof item.properties === "object") {
+      const properties = item.properties as Record<string, unknown>
+      if (properties.sessionID && typeof properties.sessionID === "object") {
+        delete (properties.sessionID as Record<string, unknown>).pattern
+      }
+    }
+    Object.values(item).forEach(strip)
+  }
+
+  strip(schema)
+  return schema
+}
+
+const StartParams = strict({
   mode: Schema.Literals(["worktree", "local"]).annotate({
     description: "Use worktree for isolated git worktrees, or local for same-directory Agent Manager sessions",
   }),
@@ -59,14 +99,14 @@ const StartParams = Schema.Struct({
     .annotate({ description: "Agent Manager sessions to start" }),
 })
 
-const ListParams = Schema.Struct({
+const ListParams = strict({
   action: Schema.Literal("list").annotate({
     description:
       "Read the current Agent Manager sections, worktrees, and sessions before any assignment. This is the source of truth for section and session IDs.",
   }),
   filter: Schema.optional(
     Schema.NullOr(
-      Schema.Struct({
+      strict({
         sectionIDs: Schema.optional(Schema.Array(Schema.String).check(Schema.isMaxLength(100))),
         states: Schema.optional(
           Schema.Array(Schema.Literals(["idle", "busy", "retry", "offline", "waiting"])).check(Schema.isMaxLength(5)),
@@ -78,20 +118,24 @@ const ListParams = Schema.Struct({
   }),
 })
 
-const PromptParams = Schema.Struct({
+const PromptParams = strict({
   action: Schema.Literal("prompt"),
-  sessionID: SessionID,
+  sessionID: SessionID.annotate({
+    description: "Session ID returned by action=list. Do not use a worktree name, branch, or section name.",
+  }),
   prompt: Schema.String.check(Schema.isMinLength(1), Schema.isMaxLength(100_000)).check(
     Schema.makeFilter((value) => (value.trim() ? undefined : "Prompt must not be empty")),
   ),
 })
 
-const StopParams = Schema.Struct({
+const StopParams = strict({
   action: Schema.Literal("stop"),
-  sessionID: SessionID,
+  sessionID: SessionID.annotate({
+    description: "Session ID returned by action=list. Do not use a worktree name, branch, or section name.",
+  }),
 })
 
-const MoveParams = Schema.Struct({
+const MoveParams = strict({
   action: Schema.Literal("move").annotate({
     description: "Move exactly one managed worktree by targeting one of its session IDs returned by action=list.",
   }),
@@ -103,25 +147,7 @@ const MoveParams = Schema.Struct({
   }),
 })
 
-export const Params = Schema.Union([StartParams, ListParams, PromptParams, StopParams, MoveParams])
-
-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", "stop", "move"]).annotate({
-      description:
-        "Use list first to discover IDs and assignments. Use move only after list, once per worktree. Never edit .kilo/agent-manager.json for these operations.",
-    }),
-  ),
-  filter: Schema.optional(ListParams.fields.filter),
-  sessionID: Schema.optional(
-    Schema.String.annotate({ description: "For move, use a session ID returned by action=list (IDs start with ses_)." }),
-  ),
-  prompt: Schema.optional(PromptParams.fields.prompt),
-  sectionID: Schema.optional(MoveParams.fields.sectionID),
-})
+export const Params = Schema.Union([StartParams, ListParams, PromptParams, MoveParams, StopParams])
 
 type Input = Schema.Schema.Type<typeof Task>
 type Selected = { task?: AgentManagerTask; error?: string }
@@ -281,18 +307,10 @@ export const AgentManagerTool = Tool.define<
     const bus = yield* Bus.Service
     const host = yield* AgentManager.Service
     const provider = yield* Provider.Service
-    const wire = ToolJsonSchema.fromSchema(WireParams)
-    const section = wire.properties?.sectionID
-    if (section && typeof section === "object" && wire.properties) {
-      wire.properties.sectionID = {
-        anyOf: [{ type: "string", minLength: 1 }, { type: "null" }],
-        description: "Section ID returned by action=list. Use null to unassign the worktree from its current section.",
-      }
-    }
     return {
       description: DESCRIPTION,
       parameters: Params,
-      jsonSchema: wire,
+      jsonSchema: wireSchema(),
       execute: (params, ctx) =>
         Effect.gen(function* () {
           if ("action" in params) {

packages/opencode/test/kilocode/agent-manager-tool.test.ts

@@ -1,6 +1,6 @@
 import { AppNodeBuilder } from "@opencode-ai/core/effect/app-node-builder"
 import { describe, expect, test } from "bun:test"
-import { Effect, Layer, ManagedRuntime, Queue, Schema } from "effect"
+import { Effect, Layer, ManagedRuntime, Queue, Result, Schema } from "effect"
 import { MessageID, SessionID } from "../../src/session/schema"
 import { provideTmpdirInstance } from "../fixture/fixture"
 import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner"
@@ -151,45 +151,87 @@ function publish(
 }
 
 describe("agent_manager tool", () => {
-  test("uses an object-root input schema without combinators", async () => {
+  test("advertises each operation as a strict union branch", async () => {
     const tool = await init()
     const schema = ToolJsonSchema.fromTool(tool)
 
-    expect(schema.type).toBe("object")
-    expect(schema.anyOf).toBeUndefined()
+    expect(schema.type).toBeUndefined()
+    expect(schema.anyOf).toHaveLength(5)
     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", "move"])
-    expect(action && typeof action === "object" ? action.description : undefined).toContain("Use list first")
-    expect(action && typeof action === "object" ? action.description : undefined).toContain("Never edit")
-    expect(schema.properties?.sessionID).toEqual(
-      expect.objectContaining({ description: expect.stringContaining("IDs start with ses_") }),
+    const branches = schema.anyOf as Array<Record<string, unknown>>
+    const properties = (branch: Record<string, unknown>) => branch.properties as Record<string, unknown>
+    expect(branches.map((branch) => branch.required)).toEqual([
+      ["mode", "tasks"],
+      ["action"],
+      ["action", "sessionID", "prompt"],
+      ["action", "sessionID", "sectionID"],
+      ["action", "sessionID"],
+    ])
+    expect(branches.every((branch) => branch.additionalProperties === false)).toBe(true)
+    expect(properties(branches[2]!).sessionID).not.toHaveProperty("pattern")
+    expect(properties(branches[3]!).sessionID).not.toHaveProperty("pattern")
+    expect(properties(branches[4]!).sessionID).not.toHaveProperty("pattern")
+    expect(properties(branches[0]!)).toEqual(
+      expect.objectContaining({ mode: expect.anything(), tasks: expect.anything() }),
     )
-    expect(schema.properties?.sessionID).not.toHaveProperty("pattern")
-    expect(schema.properties?.sectionID).toEqual(
-      expect.objectContaining({ description: expect.stringContaining("Use null to unassign") }),
+    expect(properties(branches[1]!)).toEqual(
+      expect.objectContaining({ action: expect.objectContaining({ enum: ["list"] }) }),
     )
-    expect(schema.properties?.sectionID).toEqual(
+    expect(properties(branches[2]!)).toEqual(
       expect.objectContaining({
-        anyOf: expect.arrayContaining([expect.objectContaining({ type: "string" }), { type: "null" }]),
+        action: expect.objectContaining({ enum: ["prompt"] }),
+        sessionID: expect.objectContaining({
+          description: expect.stringContaining("Session ID returned by action=list"),
+        }),
+      }),
+    )
+    expect(properties(branches[3]!)).toEqual(
+      expect.objectContaining({ action: expect.objectContaining({ enum: ["move"] }) }),
+    )
+    expect(properties(branches[4]!)).toEqual(
+      expect.objectContaining({
+        action: expect.objectContaining({ enum: ["stop"] }),
+        sessionID: expect.objectContaining({
+          description: expect.stringContaining("Session ID returned by action=list"),
+        }),
       }),
     )
-    expect(Object.keys(schema.properties ?? {})).toEqual([
-      "mode",
-      "versions",
-      "tasks",
-      "action",
-      "filter",
-      "sessionID",
-      "prompt",
-      "sectionID",
-    ])
   })
 
-  test("keeps session ID validation local", () => {
-    expect(Schema.is(Params)({ action: "stop", sessionID: "ses_target" })).toBe(true)
-    expect(Schema.is(Params)({ action: "stop", sessionID: "invalid" })).toBe(false)
+  test("accepts each operation branch and rejects ambiguous payloads", () => {
+    const task = { prompt: "Fix the issue" }
+    const accepts = (input: unknown) => Result.isSuccess(Schema.decodeUnknownResult(Params)(input))
+    expect(accepts({ mode: "local", tasks: [task] })).toBe(true)
+    expect(accepts({ action: "list" })).toBe(true)
+    expect(accepts({ action: "list", filter: null })).toBe(true)
+    expect(accepts({ action: "prompt", sessionID: "ses_target", prompt: "Continue" })).toBe(true)
+    expect(accepts({ action: "stop", sessionID: "ses_target" })).toBe(true)
+    expect(accepts({ action: "move", sessionID: "ses_target", sectionID: null })).toBe(true)
+    expect(accepts({ action: "stop", sessionID: "invalid" })).toBe(false)
+
+    expect(accepts({ mode: "local", tasks: [task], action: "list" })).toBe(false)
+    expect(accepts({ action: "list", mode: "local", tasks: [task] })).toBe(false)
+    expect(accepts({ action: "prompt", sessionID: "ses_target", prompt: "Continue", mode: "local" })).toBe(false)
+    expect(accepts({ action: "stop", sessionID: "ses_target", prompt: "Continue" })).toBe(false)
+    expect(accepts({ action: "move", sessionID: "ses_target", sectionID: null, filter: null })).toBe(false)
+  })
+
+  test("rejects mixed payloads before dispatch", async () => {
+    const tool = await init()
+    const calls: unknown[] = []
+
+    await expect(
+      runtime.runPromise(
+        provideTmpdirInstance(() =>
+          tool.execute(
+            { mode: "local", tasks: [{ prompt: "Fix issue" }], action: "list" },
+            { ...ctx, ask: (input: unknown) => Effect.sync(() => calls.push(input)) },
+          ),
+        ).pipe(Effect.scoped),
+      ),
+    ).rejects.toThrow("Unexpected Agent Manager parameter")
+    expect(calls).toEqual([])
   })
 
   test("asks for agent_manager permission", async () => {