@raidou/pi-pm-subagents 0.1.1

package/src/subagent/manager.test.ts

@@ -0,0 +1,392 @@
+import type {
+  AgentSession,
+  ContextUsage,
+} from '@earendil-works/pi-coding-agent'
+import { describe, expect, it, vi } from 'vitest'
+
+import {
+  formatSubagentSummary,
+  type LiveSubagent,
+  MAX_REUSE_FOLLOWUPS,
+  SubagentManager,
+} from './manager.js'
+
+type StubSession = Pick<
+  AgentSession,
+  'dispose' | 'abort' | 'steer' | 'subscribe' | 'prompt' | 'getContextUsage'
+> & { messages: AgentSession['messages'] }
+
+function makeStubSession(contextUsage?: ContextUsage): {
+  session: StubSession
+  steerMock: ReturnType<typeof vi.fn>
+  promptMock: ReturnType<typeof vi.fn>
+  abortMock: ReturnType<typeof vi.fn>
+} {
+  const steerMock = vi.fn().mockResolvedValue(undefined)
+  const promptMock = vi.fn().mockResolvedValue(undefined)
+  const abortMock = vi.fn().mockResolvedValue(undefined)
+
+  const session: StubSession = {
+    messages: [],
+    dispose: () => {},
+    abort: abortMock,
+    steer: steerMock,
+    subscribe: () => () => {},
+    prompt: promptMock,
+    getContextUsage: vi.fn().mockReturnValue(contextUsage),
+  }
+
+  return { session, steerMock, promptMock, abortMock }
+}
+
+function registerSubagent(
+  manager: SubagentManager,
+  session: StubSession,
+  overrides: Partial<{
+    id: number
+    status: 'running' | 'done' | 'failed' | 'killed'
+    followUpCount: number
+  }> = {},
+): LiveSubagent {
+  const id = overrides.id ?? 1
+  const status = overrides.status ?? 'done'
+  const followUpCount = overrides.followUpCount ?? 0
+
+  const subagent: LiveSubagent = {
+    id,
+    title: `Task ${id}`,
+    previousEntries: [],
+    prompt: `Do task ${id}`,
+    status,
+    session: session as unknown as AgentSession,
+    startedAt: Date.now() - 1000,
+    completedAt: status !== 'running' ? Date.now() : undefined,
+    followUpCount,
+    activeTools: [],
+    role: 'worker',
+  }
+
+  ;(
+    manager as unknown as { subagents: Map<number, LiveSubagent> }
+  ).subagents.set(id, subagent)
+  return subagent
+}
+
+describe('formatSubagentSummary', () => {
+  it('renders context usage between title and follow symbol', () => {
+    const { session } = makeStubSession()
+    const manager = new SubagentManager()
+    const subagent = registerSubagent(manager, session, {
+      id: 1,
+      status: 'done',
+      followUpCount: 2,
+    })
+    subagent.contextUsage = {
+      tokens: 60000,
+      contextWindow: 200000,
+      percent: 30,
+    }
+
+    const summary = formatSubagentSummary(subagent, 80)
+
+    expect(summary).toBe('done #1 Task 1 60k/200k ⟳ 2 1s')
+  })
+
+  it('renders ? when tokens is null', () => {
+    const { session } = makeStubSession()
+    const manager = new SubagentManager()
+    const subagent = registerSubagent(manager, session, { id: 1 })
+    subagent.contextUsage = {
+      tokens: null,
+      contextWindow: 200000,
+      percent: null,
+    }
+
+    expect(formatSubagentSummary(subagent, 80)).toContain(' Task 1 ? ')
+  })
+
+  it('renders ? when contextUsage is undefined', () => {
+    const { session } = makeStubSession()
+    const manager = new SubagentManager()
+    const subagent = registerSubagent(manager, session, { id: 1 })
+
+    expect(formatSubagentSummary(subagent, 80)).toContain(' Task 1 ? ')
+  })
+})
+
+describe('SubagentManager.followup', () => {
+  describe('subagent not found', () => {
+    it('throws when subagent id does not exist', async () => {
+      const manager = new SubagentManager()
+      await expect(manager.followup(99, 'title', 'task')).rejects.toThrow(
+        'Subagent #99 not found',
+      )
+    })
+  })
+
+  describe('done subagent followup (regression)', () => {
+    it('calls run, onStart, resets fields, increments followUpCount', async () => {
+      const { session, promptMock } = makeStubSession()
+      const onStartMock = vi.fn()
+      const onStatusChangeMock = vi.fn()
+      const manager = new SubagentManager({
+        onEachStart: onStartMock,
+        onStatusChange: onStatusChangeMock,
+      })
+
+      registerSubagent(manager, session, {
+        id: 1,
+        status: 'done',
+        followUpCount: 0,
+      })
+      promptMock.mockReturnValue(new Promise(() => {}))
+
+      const result = await manager.followup(1, 'New Title', 'New Task')
+
+      expect(result.id).toBe(1)
+      expect(result.status).toBe('running')
+      expect(result.title).toBe('New Title')
+      expect(result.prompt).toBe('New Task')
+      expect(result.followUpCount).toBe(1)
+      expect(result.startedAt).toBeGreaterThan(0)
+      expect(result.completedAt).toBeUndefined()
+      expect(onStartMock).toHaveBeenCalledOnce()
+      expect(onStatusChangeMock).toHaveBeenCalled()
+      expect(promptMock).toHaveBeenCalledOnce()
+    })
+
+    it(`throws when followUpCount reaches MAX_REUSE_FOLLOWUPS (${MAX_REUSE_FOLLOWUPS})`, async () => {
+      const { session } = makeStubSession()
+      const manager = new SubagentManager()
+      registerSubagent(manager, session, {
+        id: 1,
+        status: 'done',
+        followUpCount: MAX_REUSE_FOLLOWUPS,
+      })
+
+      await expect(manager.followup(1, 'title', 'task')).rejects.toThrow(
+        `Subagent #1 follow-up budget exhausted (${MAX_REUSE_FOLLOWUPS}/${MAX_REUSE_FOLLOWUPS}). Start a fresh subagent instead.`,
+      )
+
+      expect(session.prompt).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('running subagent followup (steer)', () => {
+    it('calls session.steer with prefixed task text', async () => {
+      const { session, steerMock, promptMock } = makeStubSession()
+      const onStartMock = vi.fn()
+      const onStatusChangeMock = vi.fn()
+      const manager = new SubagentManager({
+        onEachStart: onStartMock,
+        onStatusChange: onStatusChangeMock,
+      })
+
+      registerSubagent(manager, session, {
+        id: 1,
+        status: 'running',
+        followUpCount: 0,
+      })
+
+      const result = await manager.followup(1, 'New Title', 'New Task')
+
+      expect(result.id).toBe(1)
+      expect(result.status).toBe('running')
+      expect(result.title).toBe('New Title')
+      expect(result.prompt).toBe('New Task')
+      expect(result.followUpCount).toBe(1)
+      expect(steerMock).toHaveBeenCalledOnce()
+      expect(steerMock).toHaveBeenCalledWith('New Task')
+      expect(promptMock).not.toHaveBeenCalled()
+      expect(onStartMock).not.toHaveBeenCalled()
+    })
+
+    it('does not reset startedAt / completedAt / message', async () => {
+      const { session, steerMock } = makeStubSession()
+      const manager = new SubagentManager()
+      const subagent = registerSubagent(manager, session, {
+        id: 1,
+        status: 'running',
+        followUpCount: 0,
+      })
+      const originalStartedAt = Date.now() - 5000
+      subagent.startedAt = originalStartedAt
+
+      await manager.followup(1, 'title', 'task')
+
+      expect(steerMock).toHaveBeenCalled()
+      expect(subagent.startedAt).toBe(originalStartedAt)
+      expect(subagent.completedAt).toBeUndefined()
+    })
+
+    it('increments followUpCount each time', async () => {
+      const { session, steerMock } = makeStubSession()
+      const manager = new SubagentManager()
+      registerSubagent(manager, session, {
+        id: 1,
+        status: 'running',
+        followUpCount: 3,
+      })
+
+      await manager.followup(1, 't1', 'task 1')
+      expect(steerMock).toHaveBeenCalledTimes(1)
+
+      await manager.followup(1, 't2', 'task 2')
+      expect(steerMock).toHaveBeenCalledTimes(2)
+
+      const subagent = (
+        manager as unknown as { subagents: Map<number, LiveSubagent> }
+      ).subagents.get(1)
+      expect(subagent?.followUpCount).toBe(5)
+    })
+
+    it(`throws when followUpCount reaches MAX_REUSE_FOLLOWUPS (${MAX_REUSE_FOLLOWUPS})`, async () => {
+      const { session, steerMock } = makeStubSession()
+      const manager = new SubagentManager()
+      registerSubagent(manager, session, {
+        id: 1,
+        status: 'running',
+        followUpCount: MAX_REUSE_FOLLOWUPS,
+      })
+
+      await expect(manager.followup(1, 'title', 'task')).rejects.toThrow(
+        `Subagent #1 follow-up budget exhausted (${MAX_REUSE_FOLLOWUPS}/${MAX_REUSE_FOLLOWUPS}). Start a fresh subagent instead.`,
+      )
+
+      expect(steerMock).not.toHaveBeenCalled()
+    })
+
+    it('saves previousEntry.status as done (not running)', async () => {
+      const { session, steerMock } = makeStubSession()
+      const manager = new SubagentManager()
+      const subagent = registerSubagent(manager, session, {
+        id: 1,
+        status: 'running',
+        followUpCount: 0,
+      })
+
+      await manager.followup(1, 'New Title', 'task')
+
+      expect(steerMock).toHaveBeenCalled()
+      expect(subagent.previousEntries).toHaveLength(1)
+      expect(subagent.previousEntries[0]).toMatchObject({
+        title: 'Task 1',
+        status: 'done',
+        followUpCount: 0,
+      })
+    })
+  })
+
+  describe('previousEntries accumulation', () => {
+    it('accumulates previous titles with full metadata on each followup', async () => {
+      const { session } = makeStubSession()
+      const manager = new SubagentManager()
+      const subagent = registerSubagent(manager, session, {
+        id: 1,
+        status: 'done',
+        followUpCount: 0,
+      })
+
+      await manager.followup(1, 'Title 1', 'task 1')
+      await manager.followup(1, 'Title 2', 'task 2')
+      await manager.followup(1, 'Title 3', 'task 3')
+
+      expect(subagent.previousEntries).toHaveLength(3)
+      expect(subagent.previousEntries[0]).toMatchObject({
+        title: 'Title 2',
+        status: 'done',
+        followUpCount: 2,
+      })
+      expect(subagent.previousEntries[1]).toMatchObject({
+        title: 'Title 1',
+        status: 'done',
+        followUpCount: 1,
+      })
+      expect(subagent.previousEntries[2]).toMatchObject({
+        title: 'Task 1',
+        status: 'done',
+        followUpCount: 0,
+      })
+      expect(subagent.title).toBe('Title 3')
+    })
+
+    it('does not modify previousEntries when followup budget exhausted', async () => {
+      const { session } = makeStubSession()
+      const manager = new SubagentManager()
+      const subagent = registerSubagent(manager, session, {
+        id: 1,
+        status: 'done',
+        followUpCount: MAX_REUSE_FOLLOWUPS,
+      })
+
+      await expect(manager.followup(1, 'New Title', 'task')).rejects.toThrow()
+
+      expect(subagent.previousEntries).toEqual([])
+      expect(subagent.title).toBe('Task 1')
+    })
+  })
+})
+
+describe('SubagentManager subscribe contextUsage', () => {
+  it('initializes contextUsage on subscribe', () => {
+    const contextUsage: ContextUsage = {
+      tokens: 50000,
+      contextWindow: 200000,
+      percent: 25.0,
+    }
+    const { session } = makeStubSession(contextUsage)
+    const manager = new SubagentManager()
+    const subagent = registerSubagent(manager, session, { id: 1 })
+
+    expect(subagent.contextUsage).toBeUndefined()
+    ;(manager as unknown as { subscribe: (s: LiveSubagent) => void }).subscribe(
+      subagent,
+    )
+
+    expect(subagent.contextUsage).toEqual(contextUsage)
+  })
+
+  it('updates contextUsage on message_end event', () => {
+    const initialUsage: ContextUsage = {
+      tokens: 50000,
+      contextWindow: 200000,
+      percent: 25.0,
+    }
+    const updatedUsage: ContextUsage = {
+      tokens: 100000,
+      contextWindow: 200000,
+      percent: 50.0,
+    }
+
+    const subscribeMock = vi.fn().mockReturnValue(() => {})
+    const getContextUsageMock = vi.fn().mockReturnValue(initialUsage)
+
+    const session = {
+      messages: [],
+      dispose: () => {},
+      abort: vi.fn().mockResolvedValue(undefined),
+      steer: vi.fn().mockResolvedValue(undefined),
+      subscribe: subscribeMock,
+      prompt: vi.fn().mockResolvedValue(undefined),
+      getContextUsage: getContextUsageMock,
+    } as unknown as StubSession
+
+    const manager = new SubagentManager()
+    const subagent = registerSubagent(manager, session, { id: 1 })
+    ;(manager as unknown as { subscribe: (s: LiveSubagent) => void }).subscribe(
+      subagent,
+    )
+
+    expect(subagent.contextUsage).toEqual(initialUsage)
+
+    getContextUsageMock.mockReturnValue(updatedUsage)
+    const calls = subscribeMock.mock.calls
+    expect(calls.length).toBeGreaterThan(0)
+    const subscribeCallback = calls[0]?.[0]
+    if (subscribeCallback) {
+      subscribeCallback({ type: 'message_end', turnIndex: 0 })
+    }
+
+    expect(subagent.contextUsage).toEqual(updatedUsage)
+  })
+})

package/src/subagent/manager.ts

@@ -0,0 +1,277 @@
+import { mkdirSync } from 'node:fs'
+import { join } from 'node:path'
+
+import type { ThinkingLevel } from '@earendil-works/pi-agent-core'
+import type { Api, Model } from '@earendil-works/pi-ai'
+import {
+  type AgentSession,
+  type ContextUsage,
+  createAgentSession,
+  DefaultResourceLoader,
+  getAgentDir,
+  SessionManager,
+} from '@earendil-works/pi-coding-agent'
+
+import { formatContextUsage, formatElapsed } from '../utils/format.js'
+import { lastMessageText } from '../utils/messages.js'
+import { FOLLOW_SYMBOL } from './consts.ts'
+import type { FleetEntryBase } from './fleet.js'
+import {
+  runInSubagentSpawnContext,
+  SUBAGENT_SESSION_ID_PREFIX,
+} from './identity.js'
+
+const subagentDirFor = (cwd: string, agentDir: string): string => {
+  const safeCwd = cwd.replace(/^[/\\]/, '').replace(/[/\\:]/g, '-')
+  return join(agentDir, 'sessions', 'pi-pm-subagents', `--${safeCwd}--`)
+}
+
+const MAX_SUBAGENT_OUTPUT_BYTES = 50 * 1024
+
+export const MAX_REUSE_FOLLOWUPS = 50
+export const MAX_CONCURRENCY_SUBAGENT = 5
+
+export type SubagentStatus = 'running' | 'done' | 'failed' | 'killed'
+
+export interface SpawnOptions {
+  cwd: string
+  model?: Model<Api>
+  thinkingLevel?: ThinkingLevel
+  tools?: readonly string[]
+  systemPrompt?: string
+  role?: string
+  onComplete?: (subagent: LiveSubagent, lastMessage: string) => Promise<void>
+}
+
+export interface LiveSubagent {
+  id: number
+  title: string
+  previousEntries: FleetEntryBase[]
+  prompt: string
+  status: SubagentStatus
+  session: AgentSession
+  startedAt: number
+  completedAt?: number
+  followUpCount: number
+  activeTools: string[]
+  role: string
+  contextUsage?: ContextUsage
+  onComplete?: (subagent: LiveSubagent, lastMessage: string) => Promise<void>
+}
+
+export interface SubagentManagerOptions {
+  onStatusChange?: () => void
+  onEachStart?: (subagent: LiveSubagent) => void
+  onEachEnd?: (subagent: LiveSubagent) => void
+}
+
+export function formatSubagentSummary(
+  subagent: LiveSubagent,
+  titleWidth = 80,
+): string {
+  return [
+    subagent.status,
+    `#${subagent.id}`,
+    subagent.title.slice(0, titleWidth),
+    formatContextUsage(subagent.contextUsage),
+    `${FOLLOW_SYMBOL} ${subagent.followUpCount}`,
+    formatElapsed(subagent),
+  ].join(' ')
+}
+
+export class SubagentManager {
+  private subagents = new Map<number, LiveSubagent>()
+  private seq = 0
+
+  constructor(private options: SubagentManagerOptions = {}) {}
+
+  private countRunning(): number {
+    return [...this.subagents.values()].filter(
+      (subagent) => subagent.status === 'running',
+    ).length
+  }
+
+  list(): LiveSubagent[] {
+    return [...this.subagents.values()]
+  }
+
+  get(id: number): LiveSubagent | undefined {
+    return this.subagents.get(id)
+  }
+
+  latest(): LiveSubagent | undefined {
+    let latest: LiveSubagent | undefined
+    for (const subagent of this.subagents.values()) {
+      if (!latest || subagent.id > latest.id) latest = subagent
+    }
+    return latest
+  }
+
+  private async createLoader(
+    options: SpawnOptions,
+  ): Promise<DefaultResourceLoader> {
+    const loader = new DefaultResourceLoader({
+      cwd: options.cwd,
+      agentDir: getAgentDir(),
+      systemPromptOverride: (base) => `${base}\n\n${options.systemPrompt}`,
+    })
+    await loader.reload()
+    return loader
+  }
+
+  async createNewSubagent(
+    title: string,
+    prompt: string,
+    options: SpawnOptions,
+  ): Promise<LiveSubagent> {
+    if (this.countRunning() >= MAX_CONCURRENCY_SUBAGENT) {
+      throw new Error(
+        `Subagent concurrency limit reached (${MAX_CONCURRENCY_SUBAGENT}). Wait for an existing subagent to finish, or abort one.`,
+      )
+    }
+    this.seq += 1
+    const id = this.seq
+    const subagentSessionDir = subagentDirFor(options.cwd, getAgentDir())
+    mkdirSync(subagentSessionDir, { recursive: true })
+    const created = await runInSubagentSpawnContext(id, async () => {
+      const loader = await this.createLoader(options)
+      return createAgentSession({
+        cwd: options.cwd,
+        model: options.model,
+        thinkingLevel: options.thinkingLevel,
+        tools: options.tools ? [...options.tools] : undefined,
+        resourceLoader: loader,
+        sessionManager: SessionManager.create(options.cwd, subagentSessionDir, {
+          id: `${SUBAGENT_SESSION_ID_PREFIX}${id}-${Date.now()}`,
+        }),
+      })
+    })
+
+    const activeTools: string[] = options.tools ? [...options.tools] : []
+    const subagent: LiveSubagent = {
+      id,
+      title,
+      previousEntries: [],
+      prompt,
+      status: 'running',
+      session: created.session,
+      startedAt: Date.now(),
+      followUpCount: 0,
+      activeTools,
+      role: options.role ?? 'worker',
+      onComplete: options.onComplete,
+    }
+
+    this.subagents.set(id, subagent)
+    this.subscribe(subagent)
+    this.options.onEachStart?.(subagent)
+    this.options.onStatusChange?.()
+    void this.run(subagent, prompt)
+    return subagent
+  }
+
+  async followup(
+    id: number,
+    title: string,
+    prompt: string,
+  ): Promise<LiveSubagent> {
+    const followupSubagent = this.subagents.get(id)
+    if (!followupSubagent) throw new Error(`Subagent #${id} not found`)
+    if (followupSubagent.followUpCount >= MAX_REUSE_FOLLOWUPS)
+      throw new Error(
+        `Subagent #${id} follow-up budget exhausted (${MAX_REUSE_FOLLOWUPS}/${MAX_REUSE_FOLLOWUPS}). Start a fresh subagent instead.`,
+      )
+
+    const wasRunning = followupSubagent.status === 'running'
+    const prevStartedAt = followupSubagent.startedAt
+    followupSubagent.previousEntries.unshift({
+      title: followupSubagent.title,
+      status: wasRunning ? 'done' : followupSubagent.status,
+      followUpCount: followupSubagent.followUpCount,
+      startedAt: prevStartedAt,
+      completedAt: wasRunning
+        ? Date.now()
+        : (followupSubagent.completedAt ?? Date.now()),
+    })
+    followupSubagent.title = title
+    followupSubagent.prompt = prompt
+    followupSubagent.followUpCount += 1
+
+    if (followupSubagent.status === 'running') {
+      await followupSubagent.session.steer(prompt)
+      this.options.onStatusChange?.()
+      return followupSubagent
+    }
+
+    followupSubagent.status = 'running'
+    followupSubagent.startedAt = Date.now()
+    followupSubagent.completedAt = undefined
+    this.options.onStatusChange?.()
+    this.options.onEachStart?.(followupSubagent)
+    void this.run(followupSubagent, prompt)
+    return followupSubagent
+  }
+
+  async steer(id: number, text: string): Promise<boolean> {
+    const subagent = this.subagents.get(id)
+    if (!subagent || subagent.status !== 'running') return false
+    await subagent.session.steer(text)
+    return true
+  }
+
+  async abort(id: number): Promise<boolean> {
+    const subagent = this.subagents.get(id)
+    if (!subagent || subagent.status !== 'running') return false
+    subagent.status = 'killed'
+    this.options.onStatusChange?.()
+    await subagent.session.abort()
+    return true
+  }
+
+  disposeAll(): void {
+    const disposedSessions = new Set<AgentSession>()
+    for (const subagent of this.subagents.values()) {
+      if (subagent.status === 'running') {
+        subagent.status = 'killed'
+        this.options.onEachEnd?.(subagent)
+      }
+      if (!disposedSessions.has(subagent.session)) {
+        subagent.session.dispose()
+        disposedSessions.add(subagent.session)
+      }
+    }
+    this.subagents.clear()
+  }
+
+  private subscribe(subagent: LiveSubagent): void {
+    subagent.contextUsage = subagent.session.getContextUsage()
+    subagent.session.subscribe((event) => {
+      if (event.type === 'message_end') {
+        subagent.contextUsage = subagent.session.getContextUsage()
+      }
+    })
+  }
+
+  private async run(subagent: LiveSubagent, prompt: string): Promise<void> {
+    let lastMessage: string | undefined
+    try {
+      await subagent.session.prompt(prompt)
+      const finalText = lastMessageText(
+        subagent.session.messages,
+        MAX_SUBAGENT_OUTPUT_BYTES,
+      )
+      lastMessage = finalText ?? '(Subagent finished without a final message.)'
+      if (subagent.status === 'running') subagent.status = 'done'
+    } catch (error) {
+      if (subagent.status === 'running') {
+        lastMessage = error instanceof Error ? error.message : String(error)
+        subagent.status = 'failed'
+      }
+    } finally {
+      subagent.completedAt = Date.now()
+      this.options.onStatusChange?.()
+      await subagent.onComplete?.(subagent, lastMessage ?? '')
+      this.options.onEachEnd?.(subagent)
+    }
+  }
+}