Merge pull request #1810 from code-yeongyu/fix/resolve-subagent-type-for-tui-display

fix(tool-execute-before): resolve subagent_type for TUI display

src/plugin/session-agent-resolver.test.ts

@@ -0,0 +1,96 @@
+import { describe, expect, test } from "bun:test"
+import { resolveSessionAgent } from "./session-agent-resolver"
+
+describe("resolveSessionAgent", () => {
+  test("returns agent from first message with agent field", async () => {
+    //#given
+    const client = {
+      session: {
+        messages: async () => ({
+          data: [
+            { info: { role: "user" } },
+            { info: { role: "assistant", agent: "explore" } },
+            { info: { role: "assistant", agent: "oracle" } },
+          ],
+        }),
+      },
+    }
+
+    //#when
+    const agent = await resolveSessionAgent(client, "ses_test")
+
+    //#then
+    expect(agent).toBe("explore")
+  })
+
+  test("skips messages without agent field", async () => {
+    //#given
+    const client = {
+      session: {
+        messages: async () => ({
+          data: [
+            { info: { role: "user" } },
+            { info: { role: "system" } },
+            { info: { role: "assistant", agent: "plan" } },
+          ],
+        }),
+      },
+    }
+
+    //#when
+    const agent = await resolveSessionAgent(client, "ses_test")
+
+    //#then
+    expect(agent).toBe("plan")
+  })
+
+  test("returns undefined when no messages have agent", async () => {
+    //#given
+    const client = {
+      session: {
+        messages: async () => ({
+          data: [
+            { info: { role: "user" } },
+            { info: { role: "assistant" } },
+          ],
+        }),
+      },
+    }
+
+    //#when
+    const agent = await resolveSessionAgent(client, "ses_test")
+
+    //#then
+    expect(agent).toBeUndefined()
+  })
+
+  test("returns undefined when session has no messages", async () => {
+    //#given
+    const client = {
+      session: {
+        messages: async () => ({ data: [] }),
+      },
+    }
+
+    //#when
+    const agent = await resolveSessionAgent(client, "ses_test")
+
+    //#then
+    expect(agent).toBeUndefined()
+  })
+
+  test("returns undefined when API call fails", async () => {
+    //#given
+    const client = {
+      session: {
+        messages: async () => { throw new Error("API error") },
+      },
+    }
+
+    //#when
+    const agent = await resolveSessionAgent(client, "ses_test")
+
+    //#then
+    expect(agent).toBeUndefined()
+  })
+})

src/plugin/session-agent-resolver.ts

@@ -0,0 +1,36 @@
+import { log } from "../shared"
+
+interface SessionMessage {
+  info?: {
+    agent?: string
+    role?: string
+  }
+}
+
+type SessionClient = {
+  session: {
+    messages: (opts: { path: { id: string } }) => Promise<{ data?: SessionMessage[] }>
+  }
+}
+
+export async function resolveSessionAgent(
+  client: SessionClient,
+  sessionId: string,
+): Promise<string | undefined> {
+  try {
+    const messagesResp = await client.session.messages({ path: { id: sessionId } })
+    const messages = (messagesResp.data ?? []) as SessionMessage[]
+
+    for (const msg of messages) {
+      if (msg.info?.agent) {
+        return msg.info.agent
+      }
+    }
+  } catch (error) {
+    log("[session-agent-resolver] Failed to resolve agent from session", {
+      sessionId,
+      error: String(error),
+    })
+  }
+  return undefined
+}

src/plugin/tool-execute-before.test.ts

@@ -0,0 +1,138 @@
+import { describe, expect, test } from "bun:test"
+import { createToolExecuteBeforeHandler } from "./tool-execute-before"
+import type { CreatedHooks } from "../create-hooks"
+
+describe("createToolExecuteBeforeHandler", () => {
+  describe("task tool subagent_type normalization", () => {
+    const emptyHooks = {} as CreatedHooks
+
+    function createCtxWithSessionMessages(messages: Array<{ info?: { agent?: string; role?: string } }> = []) {
+      return {
+        client: {
+          session: {
+            messages: async () => ({ data: messages }),
+          },
+        },
+      } as unknown as Parameters<typeof createToolExecuteBeforeHandler>[0]["ctx"]
+    }
+
+    test("sets subagent_type to sisyphus-junior when category is provided without subagent_type", async () => {
+      //#given
+      const ctx = createCtxWithSessionMessages()
+      const handler = createToolExecuteBeforeHandler({ ctx, hooks: emptyHooks })
+      const input = { tool: "task", sessionID: "ses_123", callID: "call_1" }
+      const output = { args: { category: "quick", description: "Test" } as Record<string, unknown> }
+
+      //#when
+      await handler(input, output)
+
+      //#then
+      expect(output.args.subagent_type).toBe("sisyphus-junior")
+    })
+
+    test("preserves existing subagent_type when explicitly provided", async () => {
+      //#given
+      const ctx = createCtxWithSessionMessages()
+      const handler = createToolExecuteBeforeHandler({ ctx, hooks: emptyHooks })
+      const input = { tool: "task", sessionID: "ses_123", callID: "call_1" }
+      const output = { args: { subagent_type: "plan", description: "Plan test" } as Record<string, unknown> }
+
+      //#when
+      await handler(input, output)
+
+      //#then
+      expect(output.args.subagent_type).toBe("plan")
+    })
+
+    test("sets subagent_type to sisyphus-junior when category provided with different subagent_type", async () => {
+      //#given
+      const ctx = createCtxWithSessionMessages()
+      const handler = createToolExecuteBeforeHandler({ ctx, hooks: emptyHooks })
+      const input = { tool: "task", sessionID: "ses_123", callID: "call_1" }
+      const output = { args: { category: "quick", subagent_type: "oracle", description: "Test" } as Record<string, unknown> }
+
+      //#when
+      await handler(input, output)
+
+      //#then
+      expect(output.args.subagent_type).toBe("sisyphus-junior")
+    })
+
+    test("resolves subagent_type from session first message when session_id provided without subagent_type", async () => {
+      //#given
+      const ctx = createCtxWithSessionMessages([
+        { info: { role: "user" } },
+        { info: { role: "assistant", agent: "explore" } },
+        { info: { role: "assistant", agent: "oracle" } },
+      ])
+      const handler = createToolExecuteBeforeHandler({ ctx, hooks: emptyHooks })
+      const input = { tool: "task", sessionID: "ses_123", callID: "call_1" }
+      const output = { args: { session_id: "ses_abc123", description: "Continue task", prompt: "fix it" } as Record<string, unknown> }
+
+      //#when
+      await handler(input, output)
+
+      //#then
+      expect(output.args.subagent_type).toBe("explore")
+    })
+
+    test("falls back to 'continue' when session has no agent info", async () => {
+      //#given
+      const ctx = createCtxWithSessionMessages([
+        { info: { role: "user" } },
+        { info: { role: "assistant" } },
+      ])
+      const handler = createToolExecuteBeforeHandler({ ctx, hooks: emptyHooks })
+      const input = { tool: "task", sessionID: "ses_123", callID: "call_1" }
+      const output = { args: { session_id: "ses_abc123", description: "Continue task", prompt: "fix it" } as Record<string, unknown> }
+
+      //#when
+      await handler(input, output)
+
+      //#then
+      expect(output.args.subagent_type).toBe("continue")
+    })
+
+    test("preserves subagent_type when session_id is provided with explicit subagent_type", async () => {
+      //#given
+      const ctx = createCtxWithSessionMessages()
+      const handler = createToolExecuteBeforeHandler({ ctx, hooks: emptyHooks })
+      const input = { tool: "task", sessionID: "ses_123", callID: "call_1" }
+      const output = { args: { session_id: "ses_abc123", subagent_type: "explore", description: "Continue explore" } as Record<string, unknown> }
+
+      //#when
+      await handler(input, output)
+
+      //#then
+      expect(output.args.subagent_type).toBe("explore")
+    })
+
+    test("does not modify args for non-task tools", async () => {
+      //#given
+      const ctx = createCtxWithSessionMessages()
+      const handler = createToolExecuteBeforeHandler({ ctx, hooks: emptyHooks })
+      const input = { tool: "bash", sessionID: "ses_123", callID: "call_1" }
+      const output = { args: { command: "ls" } as Record<string, unknown> }
+
+      //#when
+      await handler(input, output)
+
+      //#then
+      expect(output.args.subagent_type).toBeUndefined()
+    })
+
+    test("does not set subagent_type when neither category nor session_id is provided and subagent_type is present", async () => {
+      //#given
+      const ctx = createCtxWithSessionMessages()
+      const handler = createToolExecuteBeforeHandler({ ctx, hooks: emptyHooks })
+      const input = { tool: "task", sessionID: "ses_123", callID: "call_1" }
+      const output = { args: { subagent_type: "oracle", description: "Oracle task" } as Record<string, unknown> }
+
+      //#when
+      await handler(input, output)
+
+      //#then
+      expect(output.args.subagent_type).toBe("oracle")
+    })
+  })
+})

src/plugin/tool-execute-before.ts

@@ -3,6 +3,7 @@ import type { PluginContext } from "./types"
 import { getMainSessionID } from "../features/claude-code-session-state"
 import { clearBoulderState } from "../features/boulder-state"
 import { log } from "../shared"
+import { resolveSessionAgent } from "./session-agent-resolver"
 
 import type { CreatedHooks } from "../create-hooks"
 
@@ -34,8 +35,13 @@ export function createToolExecuteBeforeHandler(args: {
       const argsObject = output.args
       const category = typeof argsObject.category === "string" ? argsObject.category : undefined
       const subagentType = typeof argsObject.subagent_type === "string" ? argsObject.subagent_type : undefined
-      if (category && !subagentType) {
+      const sessionId = typeof argsObject.session_id === "string" ? argsObject.session_id : undefined
+
+      if (category) {
         argsObject.subagent_type = "sisyphus-junior"
+      } else if (!subagentType && sessionId) {
+        const resolvedAgent = await resolveSessionAgent(ctx.client, sessionId)
+        argsObject.subagent_type = resolvedAgent ?? "continue"
       }
     }