@raidou/pi-pm-subagents 0.1.1

package/src/pm-mode.ts

@@ -0,0 +1,142 @@
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+  ThemeColor,
+} from '@earendil-works/pi-coding-agent'
+
+import { BASH_READONLY_TOOL_NAME } from './bash-readonly.js'
+import type { PmMode, PmSubagentState } from './types.js'
+import type { PromptDefinition } from './utils/markdown.js'
+import { modelRefOf, resolveModelRef } from './utils/model-ref.js'
+import { persist } from './utils/state.js'
+import { composeTools, type ToolConfig } from './utils/tools.js'
+
+export type { ToolConfig } from './utils/tools.js'
+
+/** Restore the model captured before entering a read-only mode. */
+export async function restoreModel(
+  pi: ExtensionAPI,
+  state: PmSubagentState,
+  ctx: ExtensionContext,
+): Promise<void> {
+  const previousModel = state.previousModel
+  state.previousModel = undefined
+  if (!previousModel) return
+
+  const currentModel = ctx.model
+  if (currentModel) {
+    const currentRef = `${currentModel.provider}/${currentModel.id}`
+    if (currentRef === previousModel) return
+  }
+
+  const model = resolveModelRef(ctx, [previousModel])
+  if (model) await pi.setModel(model)
+}
+
+export function baseToolsOf(
+  pi: ExtensionAPI,
+  state: PmSubagentState,
+): string[] {
+  const diff = state.modeDiffTools
+  if (!diff) return pi.getActiveTools()
+  const current = pi.getActiveTools()
+  const added = new Set(diff.added)
+  const merged = [...new Set([...current, ...diff.removed])]
+  return merged.filter((name) => !added.has(name))
+}
+
+export function restoreTools(pi: ExtensionAPI, state: PmSubagentState): void {
+  pi.setActiveTools(baseToolsOf(pi, state))
+  state.modeDiffTools = undefined
+}
+
+export function assertModeIdle(state: PmSubagentState, entering: PmMode): void {
+  if (state.mode !== undefined) {
+    throw new Error(
+      `Cannot enter ${entering} mode while in ${state.mode} mode — exit the current mode first.`,
+    )
+  }
+}
+
+export const WRITE_TOOLS = new Set(['edit', 'write'])
+
+const BASH_REPLACEMENT: ReadonlyMap<string, string> = new Map([
+  ['bash', BASH_READONLY_TOOL_NAME],
+])
+
+export function calculateModeTools(
+  pi: ExtensionAPI,
+  state: PmSubagentState,
+  config: ToolConfig = {},
+): string[] {
+  const base = pi.getActiveTools()
+  const applied = composeTools(base, config)
+    .filter((name) => !WRITE_TOOLS.has(name))
+    .map((name) => BASH_REPLACEMENT.get(name) ?? name)
+
+  if (!state.modeDiffTools) {
+    state.modeDiffTools = {
+      added: applied.filter((name) => !base.includes(name)),
+      removed: base.filter((name) => !applied.includes(name)),
+    }
+  }
+  return applied
+}
+
+export async function exitModeFor(
+  pi: ExtensionAPI,
+  state: PmSubagentState,
+  ctx: ExtensionContext,
+  name: PmMode,
+): Promise<void> {
+  restoreTools(pi, state)
+  await restoreModel(pi, state, ctx)
+  ctx.ui.setStatus(name, undefined)
+  persist(pi, state)
+}
+
+export async function applyModeModel(
+  pi: ExtensionAPI,
+  state: PmSubagentState,
+  ctx: ExtensionContext,
+  modelRef: string | undefined,
+): Promise<void> {
+  if (!modelRef) return
+
+  const model = resolveModelRef(ctx, [modelRef])
+  if (!model) {
+    ctx.ui.notify(`Invalid model ref: ${modelRef}`, 'warning')
+    return
+  }
+
+  if (!state.previousModel) {
+    const current = ctx.model
+    if (current) state.previousModel = modelRefOf(current)
+  }
+
+  const success = await pi.setModel(model)
+  if (!success) {
+    ctx.ui.notify('No API key for this model', 'error')
+    state.previousModel = undefined
+  }
+}
+
+export interface ModeSetupOptions {
+  /** The mode's prompt definition. Used to derive the active tool set. */
+  promptDefinition: PromptDefinition
+  /** Theme color role used for the mode's status indicator. */
+  color: ThemeColor
+}
+
+export async function applyModeFor(
+  pi: ExtensionAPI,
+  state: PmSubagentState,
+  pmMode: PmMode,
+  ctx: ExtensionContext,
+  options: ModeSetupOptions,
+): Promise<void> {
+  pi.setActiveTools(calculateModeTools(pi, state, options.promptDefinition.fm))
+  await applyModeModel(pi, state, ctx, options.promptDefinition.fm.model)
+  ctx.ui.setStatus(pmMode, ctx.ui.theme.fg(options.color, pmMode))
+  persist(pi, state)
+}

package/src/prompts/mode.test.ts

@@ -0,0 +1,289 @@
+import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'
+import { tmpdir } from 'node:os'
+import { dirname, join } from 'node:path'
+import { fileURLToPath } from 'node:url'
+
+import { parseFrontmatter } from '@earendil-works/pi-coding-agent'
+import { afterEach, beforeEach, describe, expect, it } from 'vitest'
+
+import { readModePrompt } from './mode.js'
+
+const here = dirname(fileURLToPath(import.meta.url))
+
+const MOCK_AGENT_DIR_VAR = 'PI_CODING_AGENT_DIR'
+
+describe('mode', () => {
+  let originalAgentDir: string | undefined
+  let tempDir: string
+  let modesDir: string
+
+  beforeEach(async () => {
+    originalAgentDir = process.env[MOCK_AGENT_DIR_VAR]
+    tempDir = await mkdtemp(join(tmpdir(), 'pi-pm-subagents-test-mode-'))
+    modesDir = join(tempDir, 'pm-subagents-prompts')
+    await mkdir(modesDir, { recursive: true })
+    process.env[MOCK_AGENT_DIR_VAR] = tempDir
+  })
+
+  afterEach(async () => {
+    await rm(tempDir, { recursive: true, force: true })
+    process.env[MOCK_AGENT_DIR_VAR] = originalAgentDir
+  })
+
+  describe('readModePrompt', () => {
+    it('reads mode prompt with override', async () => {
+      await writeFile(
+        join(modesDir, 'test-mode.md'),
+        `---
+description: A test mode.
+tools: [read, write]
+---
+This is a test mode prompt.`,
+      )
+
+      const result = await readModePrompt('test-mode')
+      expect(result.fm.description).toBe('A test mode.')
+      expect(result.fm.tools).toEqual(['read', 'write'])
+      expect(result.systemPrompt).toBe('This is a test mode prompt.')
+    })
+
+    it('handles mode without frontmatter', async () => {
+      await writeFile(
+        join(modesDir, 'no-fm.md'),
+        'Plain prompt without frontmatter.',
+      )
+
+      const result = await readModePrompt('no-fm')
+      expect(result.fm.description).toBeUndefined()
+      expect(result.fm.tools).toBeUndefined()
+      expect(result.systemPrompt).toBe('Plain prompt without frontmatter.')
+    })
+
+    it('supports tools as array', async () => {
+      await writeFile(
+        join(modesDir, 'array-tools.md'),
+        `---
+tools: [read, write, edit]
+---
+Test prompt.`,
+      )
+
+      const result = await readModePrompt('array-tools')
+      expect(result.fm.tools).toEqual(['read', 'write', 'edit'])
+    })
+
+    it('supports extraTools as array', async () => {
+      await writeFile(
+        join(modesDir, 'extra-array.md'),
+        `---
+extraTools: [bash, grep]
+---
+Test prompt.`,
+      )
+
+      const result = await readModePrompt('extra-array')
+      expect(result.fm.extraTools).toEqual(['bash', 'grep'])
+    })
+
+    it('supports removeTools as array', async () => {
+      await writeFile(
+        join(modesDir, 'remove-array.md'),
+        `---
+removeTools: [write, bash]
+---
+Test prompt.`,
+      )
+
+      const result = await readModePrompt('remove-array')
+      expect(result.fm.removeTools).toEqual(['write', 'bash'])
+    })
+
+    it('supports model field', async () => {
+      await writeFile(
+        join(modesDir, 'model-test.md'),
+        `---
+model: anthropic/claude-sonnet-4-5
+---
+Test prompt.`,
+      )
+
+      const result = await readModePrompt('model-test')
+      expect(result.fm.model).toBe('anthropic/claude-sonnet-4-5')
+    })
+
+    it('appends extra content from -append.md', async () => {
+      await writeFile(
+        join(modesDir, 'with-append.md'),
+        `---
+description: Base mode.
+---
+Base prompt content.`,
+      )
+      await writeFile(
+        join(modesDir, 'with-append-append.md'),
+        'Appended extra content.',
+      )
+
+      const result = await readModePrompt('with-append')
+      expect(result.fm.description).toBe('Base mode.')
+      expect(result.systemPrompt).toBe(
+        'Base prompt content.\n\nAppended extra content.',
+      )
+    })
+
+    it('handles append without override', async () => {
+      await writeFile(
+        join(modesDir, 'coordinator-append.md'),
+        'Appended extra content.',
+      )
+
+      const bundled = await readFile(
+        join(here, '../../pm-subagents-prompts/coordinator.md'),
+        'utf8',
+      )
+      const { body } = parseFrontmatter(bundled)
+
+      const result = await readModePrompt('coordinator')
+      // Intentionally structural assertions instead of hardcoding the full
+      // body: tolerate wording edits to the bundled prompt while guarding
+      // against sections being accidentally dropped from it.
+      expect(result.systemPrompt).toContain('You are a COORDINATOR agent')
+      expect(result.systemPrompt).toContain('## Responsibilities')
+      expect(result.systemPrompt).toContain(
+        '- Delegate tasks and wait for me to tell you subagent',
+      )
+      expect(result.systemPrompt).toContain(
+        `- Do not trust a subagent's self-reported result blindly.`,
+      )
+      expect(result.systemPrompt).toContain(body.trim())
+      expect(result.systemPrompt).toContain(
+        `${body.trim()}\n\nAppended extra content.`,
+      )
+    })
+
+    it('handles empty tools array', async () => {
+      await writeFile(
+        join(modesDir, 'empty-tools.md'),
+        `---
+tools: []
+---
+Test prompt.`,
+      )
+
+      const result = await readModePrompt('empty-tools')
+      expect(result.fm.tools).toEqual([])
+    })
+
+    it('handles empty extraTools array', async () => {
+      await writeFile(
+        join(modesDir, 'empty-extra.md'),
+        `---
+extraTools: []
+---
+Test prompt.`,
+      )
+
+      const result = await readModePrompt('empty-extra')
+      expect(result.fm.extraTools).toEqual([])
+    })
+
+    it('handles empty removeTools array', async () => {
+      await writeFile(
+        join(modesDir, 'empty-remove.md'),
+        `---
+removeTools: []
+---
+Test prompt.`,
+      )
+
+      const result = await readModePrompt('empty-remove')
+      expect(result.fm.removeTools).toEqual([])
+    })
+
+    it('frontmatter from base is preserved when append exists', async () => {
+      await writeFile(
+        join(modesDir, 'fm-base.md'),
+        `---
+description: Frontmatter from base.
+tools: [read]
+---
+Base content.`,
+      )
+      await writeFile(join(modesDir, 'fm-base-append.md'), 'Append content.')
+
+      const result = await readModePrompt('fm-base')
+      expect(result.fm.description).toBe('Frontmatter from base.')
+      expect(result.fm.tools).toEqual(['read'])
+      expect(result.systemPrompt).toBe('Base content.\n\nAppend content.')
+    })
+
+    it('append tools merge with base tools (union, deduped)', async () => {
+      await writeFile(
+        join(modesDir, 'tools-merge.md'),
+        `---
+tools: [read, write]
+---
+Base content.`,
+      )
+      await writeFile(
+        join(modesDir, 'tools-merge-append.md'),
+        `---
+tools: [write, grep]
+---
+`,
+      )
+
+      const result = await readModePrompt('tools-merge')
+      expect(result.fm.tools).toEqual(['read', 'write', 'grep'])
+    })
+
+    it('merges description, model, tools, and systemPrompt from append', async () => {
+      await writeFile(
+        join(modesDir, 'all-merge.md'),
+        `---
+description: Base description.
+model: base/model
+tools: [read, write]
+---
+Base content.`,
+      )
+      await writeFile(
+        join(modesDir, 'all-merge-append.md'),
+        `---
+description: Append description.
+model: append/model
+tools: [grep]
+---
+Append body.`,
+      )
+
+      const result = await readModePrompt('all-merge')
+      expect(result.fm.description).toBe('Append description.')
+      expect(result.fm.model).toBe('append/model')
+      expect(result.fm.tools).toEqual(['read', 'write', 'grep'])
+      expect(result.systemPrompt).toBe('Base content.\n\nAppend body.')
+    })
+
+    it('append with frontmatter only and empty body merges cleanly', async () => {
+      await writeFile(
+        join(modesDir, 'empty-body.md'),
+        `---
+description: Base description.
+---
+Base content.`,
+      )
+      await writeFile(
+        join(modesDir, 'empty-body-append.md'),
+        `---
+model: append/model
+---
+`,
+      )
+
+      const result = await readModePrompt('empty-body')
+      expect(result.fm.description).toBe('Base description.')
+      expect(result.fm.model).toBe('append/model')
+      expect(result.systemPrompt).toBe('Base content.')
+    })
+  })
+})

package/src/prompts/mode.ts

@@ -0,0 +1,31 @@
+import { dirname, join } from 'node:path'
+import { fileURLToPath } from 'node:url'
+
+import { getAgentDir } from '@earendil-works/pi-coding-agent'
+
+import {
+  loadMarkdown,
+  mergePromptDefinitions,
+  type PromptDefinition,
+} from '../utils/markdown.js'
+
+export async function readModePrompt(name: string): Promise<PromptDefinition> {
+  const here = dirname(fileURLToPath(import.meta.url))
+  const agentDir = getAgentDir()
+  const modesPromptsDir = join(agentDir, 'pm-subagents-prompts')
+
+  const overridePath = join(modesPromptsDir, `${name}.md`)
+  const bundledPath = join(here, '../..', 'pm-subagents-prompts', `${name}.md`)
+  const appendPath = join(modesPromptsDir, `${name}-append.md`)
+
+  const base =
+    (await loadMarkdown(overridePath)) ?? (await loadMarkdown(bundledPath))
+
+  const append = await loadMarkdown(appendPath)
+
+  if (!base) throw new Error(`Mode prompt "${name}" not found`)
+
+  if (!append) return base
+
+  return mergePromptDefinitions(base, append)
+}