@raidou/pi-pm-subagents 0.1.1

package/src/models-config/subagent-model-utils.test.ts

@@ -0,0 +1,250 @@
+import type { Api, Model } from '@earendil-works/pi-ai'
+import type { ExtensionContext } from '@earendil-works/pi-coding-agent'
+import { describe, expect, it } from 'vitest'
+
+import { MODEL_DEFAULT } from './subagent-model-constants.js'
+import {
+  cycleSubagentModel,
+  formatSubagentModelLabel,
+  resolveSubagentModelForSpawn,
+} from './subagent-model-utils.js'
+
+describe('cycleSubagentModel', () => {
+  describe('empty scope', () => {
+    it('throws', () => {
+      expect(() => cycleSubagentModel([], undefined, 1)).toThrow(
+        'subagent model scope must be non-empty',
+      )
+    })
+  })
+
+  describe('single element scope', () => {
+    it('always returns index 0 regardless of direction', () => {
+      const scope = ['anthropic/claude-haiku-4-5']
+      expect(cycleSubagentModel(scope, undefined, 1)).toEqual({
+        ref: 'anthropic/claude-haiku-4-5',
+        index: 0,
+        poolSize: 1,
+      })
+      expect(
+        cycleSubagentModel(scope, 'anthropic/claude-haiku-4-5', 1),
+      ).toEqual({
+        ref: 'anthropic/claude-haiku-4-5',
+        index: 0,
+        poolSize: 1,
+      })
+      expect(
+        cycleSubagentModel(scope, 'anthropic/claude-haiku-4-5', -1),
+      ).toEqual({
+        ref: 'anthropic/claude-haiku-4-5',
+        index: 0,
+        poolSize: 1,
+      })
+    })
+  })
+
+  describe('multiple elements scope', () => {
+    const scope = ['model-a', 'model-b', 'model-c']
+
+    it('cycles forward wrapping around', () => {
+      expect(cycleSubagentModel(scope, 'model-a', 1)).toEqual({
+        ref: 'model-b',
+        index: 1,
+        poolSize: 3,
+      })
+      expect(cycleSubagentModel(scope, 'model-b', 1)).toEqual({
+        ref: 'model-c',
+        index: 2,
+        poolSize: 3,
+      })
+      expect(cycleSubagentModel(scope, 'model-c', 1)).toEqual({
+        ref: 'model-a',
+        index: 0,
+        poolSize: 3,
+      })
+    })
+
+    it('cycles backward wrapping around', () => {
+      expect(cycleSubagentModel(scope, 'model-a', -1)).toEqual({
+        ref: 'model-c',
+        index: 2,
+        poolSize: 3,
+      })
+      expect(cycleSubagentModel(scope, 'model-c', -1)).toEqual({
+        ref: 'model-b',
+        index: 1,
+        poolSize: 3,
+      })
+      expect(cycleSubagentModel(scope, 'model-b', -1)).toEqual({
+        ref: 'model-a',
+        index: 0,
+        poolSize: 3,
+      })
+    })
+
+    it('handles undefined currentRef as first element', () => {
+      expect(cycleSubagentModel(scope, undefined, 1)).toEqual({
+        ref: 'model-b',
+        index: 1,
+        poolSize: 3,
+      })
+      expect(cycleSubagentModel(scope, undefined, -1)).toEqual({
+        ref: 'model-c',
+        index: 2,
+        poolSize: 3,
+      })
+    })
+
+    it('handles currentRef not in scope as first element', () => {
+      expect(cycleSubagentModel(scope, 'unknown-model', 1)).toEqual({
+        ref: 'model-b',
+        index: 1,
+        poolSize: 3,
+      })
+    })
+  })
+})
+
+describe('formatSubagentModelLabel', () => {
+  it('returns ref for empty scope', () => {
+    expect(formatSubagentModelLabel([], 'model-a')).toBe('model-a')
+    expect(formatSubagentModelLabel([], MODEL_DEFAULT)).toBe(MODEL_DEFAULT)
+    expect(formatSubagentModelLabel([], undefined)).toBe(MODEL_DEFAULT)
+  })
+
+  it('shows ref with index and pool size when scope is non-empty', () => {
+    const scope = ['model-a', 'model-b', 'model-c']
+    expect(formatSubagentModelLabel(scope, 'model-a')).toBe('model-a (1/3)')
+    expect(formatSubagentModelLabel(scope, 'model-b')).toBe('model-b (2/3)')
+    expect(formatSubagentModelLabel(scope, 'model-c')).toBe('model-c (3/3)')
+  })
+
+  it('shows ref without index when scope is empty', () => {
+    expect(formatSubagentModelLabel([], 'model-a')).toBe('model-a')
+    expect(formatSubagentModelLabel([], MODEL_DEFAULT)).toBe(MODEL_DEFAULT)
+    expect(formatSubagentModelLabel([], undefined)).toBe(MODEL_DEFAULT)
+  })
+
+  it('shows ref without index when ref not found in scope', () => {
+    const scope = ['model-a', 'model-b']
+    expect(formatSubagentModelLabel(scope, 'unknown')).toBe('unknown')
+  })
+
+  it('shows MODEL_DEFAULT without index when ref is undefined and not in scope', () => {
+    const scope = ['model-a', 'model-b']
+    expect(formatSubagentModelLabel(scope, undefined)).toBe(MODEL_DEFAULT)
+  })
+
+  it('shows MODEL_DEFAULT with index when ref is undefined and DEFAULT is in scope', () => {
+    const scope = [MODEL_DEFAULT, 'model-a', 'model-b']
+    expect(formatSubagentModelLabel(scope, undefined)).toBe('DEFAULT (1/3)')
+  })
+})
+
+describe('resolveSubagentModelForSpawn', () => {
+  function makeFakeModel(provider: string, id: string): Model<Api> {
+    return {
+      provider,
+      id,
+      name: `${provider}/${id}`,
+      api: 'openai-responses',
+      baseUrl: `https://${provider}.example.com`,
+      reasoning: false,
+      input: ['text'],
+      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+      contextWindow: 128000,
+      maxTokens: 8192,
+    }
+  }
+
+  function makeFakeCtx(
+    registry: Record<string, Model<Api>>,
+    currentModel: Model<Api>,
+  ): ExtensionContext {
+    return {
+      modelRegistry: {
+        find: (provider: string, id: string) =>
+          registry[`${provider}/${id}`] ?? null,
+        getAvailable: () => Object.values(registry),
+        getModel: (provider: string, id: string) =>
+          registry[`${provider}/${id}`] ?? null,
+        getProviders: () => [],
+        getModels: () => [],
+        checkAuth: async () => undefined,
+        refresh: async () => ({ aborted: false, errors: new Map() }),
+        getAuth: async () => undefined,
+        login: async () => {
+          throw new Error('not implemented')
+        },
+        logout: async () => {},
+        stream: () => {
+          throw new Error('not implemented')
+        },
+        complete: async () => {
+          throw new Error('not implemented')
+        },
+        streamSimple: () => {
+          throw new Error('not implemented')
+        },
+        completeSimple: async () => {
+          throw new Error('not implemented')
+        },
+      },
+      model: currentModel,
+    } as unknown as ExtensionContext
+  }
+
+  const modelA = makeFakeModel('provider-a', 'model-a')
+  const modelB = makeFakeModel('provider-b', 'model-b')
+  const ctxModel = makeFakeModel('default', 'ctx-model')
+
+  it('returns sessionModel when it is a valid model ref', () => {
+    const ctx = makeFakeCtx(
+      { 'provider-a/model-a': modelA, 'provider-b/model-b': modelB },
+      ctxModel,
+    )
+    const result = resolveSubagentModelForSpawn(
+      ctx,
+      undefined,
+      'provider-a/model-a',
+    )
+    expect(result).toBe(modelA)
+  })
+
+  it('falls back to ctx.model when both roleModel and sessionModel are undefined', () => {
+    const ctx = makeFakeCtx({ 'default/ctx-model': ctxModel }, ctxModel)
+    const result = resolveSubagentModelForSpawn(ctx, undefined, undefined)
+    expect(result).toBe(ctxModel)
+  })
+
+  it('roleModel takes priority over sessionModel', () => {
+    const ctx = makeFakeCtx(
+      { 'provider-a/model-a': modelA, 'provider-b/model-b': modelB },
+      ctxModel,
+    )
+    const result = resolveSubagentModelForSpawn(
+      ctx,
+      'provider-b/model-b',
+      'provider-a/model-a',
+    )
+    expect(result).toBe(modelB)
+  })
+
+  it('falls back to ctx.model when sessionModel is not in registry', () => {
+    const ctx = makeFakeCtx({}, ctxModel)
+    const result = resolveSubagentModelForSpawn(ctx, undefined, 'unknown/model')
+    expect(result).toBe(ctxModel)
+  })
+
+  it('falls back to ctx.model when sessionModel is MODEL_DEFAULT', () => {
+    const ctx = makeFakeCtx({ 'default/ctx-model': ctxModel }, ctxModel)
+    const result = resolveSubagentModelForSpawn(ctx, undefined, MODEL_DEFAULT)
+    expect(result).toBe(ctxModel)
+  })
+
+  it('sessionModel undefined falls back to MODEL_DEFAULT then ctx.model', () => {
+    const ctx = makeFakeCtx({ 'default/ctx-model': ctxModel }, ctxModel)
+    const result = resolveSubagentModelForSpawn(ctx, undefined, undefined)
+    expect(result).toBe(ctxModel)
+  })
+})

package/src/models-config/subagent-model-utils.ts

@@ -0,0 +1,52 @@
+import type { Api, Model } from '@earendil-works/pi-ai'
+import type { ExtensionContext } from '@earendil-works/pi-coding-agent'
+
+import { resolveModelRef } from '../utils/model-ref.js'
+import { MODEL_DEFAULT } from './subagent-model-constants.js'
+
+export interface CycleResult {
+  readonly ref: string
+  readonly index: number
+  readonly poolSize: number
+}
+
+export function cycleSubagentModel(
+  scope: readonly string[],
+  currentRef: string | undefined,
+  direction: 1 | -1,
+): CycleResult {
+  if (scope.length === 0) {
+    throw new Error('subagent model scope must be non-empty')
+  }
+  const currentIndex = currentRef ? scope.indexOf(currentRef) : -1
+  const safeIndex = currentIndex < 0 ? 0 : currentIndex
+  const clampedNextIndex =
+    (((safeIndex + direction) % scope.length) + scope.length) % scope.length
+  return {
+    ref: scope[clampedNextIndex] ?? MODEL_DEFAULT,
+    index: clampedNextIndex,
+    poolSize: scope.length,
+  }
+}
+
+export function resolveSubagentModelForSpawn(
+  ctx: ExtensionContext,
+  roleModel: string | undefined,
+  sessionModel: string | undefined,
+): Model<Api> | undefined {
+  const ref = roleModel ?? sessionModel ?? MODEL_DEFAULT
+  return resolveModelRef(ctx, [ref]) ?? ctx.model
+}
+
+export function formatSubagentModelLabel(
+  scope: readonly string[],
+  currentRef: string | undefined,
+): string {
+  if (scope.length === 0) {
+    return currentRef ?? MODEL_DEFAULT
+  }
+  const ref = currentRef ?? MODEL_DEFAULT
+  const index = scope.indexOf(ref)
+  if (index < 0) return ref
+  return `${ref} (${index + 1}/${scope.length})`
+}

package/src/pm-mode.test.ts

@@ -0,0 +1,324 @@
+import type { Api, Model } from '@earendil-works/pi-ai'
+import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { BASH_READONLY_TOOL_NAME } from './bash-readonly.js'
+import {
+  applyModeModel,
+  calculateModeTools,
+  type ToolConfig,
+  WRITE_TOOLS,
+} from './pm-mode.js'
+import { baseToolsOf, restoreTools } from './pm-mode.js'
+import type { PmSubagentState } from './types.js'
+import { createState } from './utils/state.js'
+
+function fakePi(initialActive: string[]): ExtensionAPI & {
+  activeTools: string[]
+} {
+  const pi = {
+    activeTools: initialActive,
+    getActiveTools() {
+      return [...pi.activeTools]
+    },
+    setActiveTools(names: string[]) {
+      pi.activeTools = [...names]
+    },
+  }
+  return pi as ExtensionAPI & { activeTools: string[] }
+}
+
+describe('calculateModeTools / restoreTools diff mechanism', () => {
+  it('preserves tools registered during the mode and restores bash on exit', () => {
+    const pi = fakePi(['read', 'write', 'bash'])
+    const state: PmSubagentState = createState()
+    pi.setActiveTools(calculateModeTools(pi, state))
+    expect(pi.activeTools).toEqual(['read', BASH_READONLY_TOOL_NAME])
+
+    pi.activeTools.push('newtool')
+
+    restoreTools(pi, state)
+    expect(pi.activeTools).toContain('newtool')
+    expect(pi.activeTools).toContain('write')
+    expect(pi.activeTools).toContain('bash')
+    expect(pi.activeTools).not.toContain(BASH_READONLY_TOOL_NAME)
+    expect(state.modeDiffTools).toBeUndefined()
+  })
+
+  it('replays persisted modeTools idempotently and restores on exit', () => {
+    const pi = fakePi(['read', BASH_READONLY_TOOL_NAME, 'grep'])
+    const state: PmSubagentState = {
+      ...createState(),
+      modeDiffTools: { added: [BASH_READONLY_TOOL_NAME], removed: ['bash'] },
+    }
+
+    pi.setActiveTools(calculateModeTools(pi, state))
+    expect(pi.activeTools).toEqual(['read', BASH_READONLY_TOOL_NAME, 'grep'])
+    expect(state.modeDiffTools).toEqual({
+      added: [BASH_READONLY_TOOL_NAME],
+      removed: ['bash'],
+    })
+
+    restoreTools(pi, state)
+    expect(pi.activeTools).toEqual(['read', 'grep', 'bash'])
+    expect(state.modeDiffTools).toBeUndefined()
+  })
+
+  it('baseToolsOf includes tools registered during the mode', () => {
+    const pi = fakePi(['read', BASH_READONLY_TOOL_NAME, 'newtool'])
+    const state: PmSubagentState = {
+      ...createState(),
+      modeDiffTools: { added: [BASH_READONLY_TOOL_NAME], removed: ['bash'] },
+    }
+
+    const base = baseToolsOf(pi, state)
+    expect(base).toContain('bash')
+    expect(base).toContain('newtool')
+    expect(base).toContain('read')
+    expect(base).not.toContain(BASH_READONLY_TOOL_NAME)
+  })
+
+  it('baseToolsOf falls back to current active tools without a diff', () => {
+    const pi = fakePi(['read', 'bash'])
+    const state = createState()
+    expect(baseToolsOf(pi, state)).toEqual(['read', 'bash'])
+  })
+})
+
+describe('calculateModeTools', () => {
+  it('defaults to read-only behavior when no config is provided', () => {
+    const base = ['read', 'write', 'edit', 'bash', 'grep', 'find']
+    const result = calculateModeTools(fakePi(base), createState(), {})
+    expect(result).toContain('read')
+    expect(result).toContain(BASH_READONLY_TOOL_NAME)
+    expect(result).toContain('grep')
+    expect(result).toContain('find')
+    expect(result).not.toContain('write')
+    expect(result).not.toContain('edit')
+    expect(result).not.toContain('bash')
+  })
+
+  it('adds extra tools to the result', () => {
+    const base = ['read', 'write', 'bash']
+    const config: ToolConfig = {
+      extraTools: ['subagent_delegate', 'subagent_kill'],
+    }
+    const result = calculateModeTools(fakePi(base), createState(), config)
+    expect(result).toContain('read')
+    expect(result).toContain(BASH_READONLY_TOOL_NAME)
+    expect(result).toContain('subagent_delegate')
+    expect(result).toContain('subagent_kill')
+  })
+
+  it('removes tools specified in removeTools', () => {
+    const base = ['read', 'write', 'bash', 'grep', 'find']
+    const config: ToolConfig = {
+      removeTools: ['grep', 'find'],
+    }
+    const result = calculateModeTools(fakePi(base), createState(), config)
+    expect(result).toContain('read')
+    expect(result).toContain(BASH_READONLY_TOOL_NAME)
+    expect(result).not.toContain('grep')
+    expect(result).not.toContain('find')
+  })
+
+  it('applies tools whitelist when tools is specified', () => {
+    const base = ['read', 'write', 'bash', 'grep', 'find']
+    const config: ToolConfig = {
+      tools: ['read', BASH_READONLY_TOOL_NAME],
+    }
+    const result = calculateModeTools(fakePi(base), createState(), config)
+    expect(result).toEqual(['read', BASH_READONLY_TOOL_NAME])
+  })
+
+  it('combines extraTools and removeTools correctly', () => {
+    const base = ['read', 'write', 'bash', 'grep', 'find']
+    const config: ToolConfig = {
+      extraTools: ['subagent_delegate'],
+      removeTools: ['grep'],
+    }
+    const result = calculateModeTools(fakePi(base), createState(), config)
+    expect(result).toContain('read')
+    expect(result).toContain(BASH_READONLY_TOOL_NAME)
+    expect(result).toContain('find')
+    expect(result).toContain('subagent_delegate')
+    expect(result).not.toContain('grep')
+  })
+
+  it('combines tools, extraTools, and removeTools correctly', () => {
+    const base = ['read', 'write', 'bash', 'grep', 'find', 'ls']
+    const config: ToolConfig = {
+      tools: ['read', 'grep', 'subagent_delegate'],
+      extraTools: ['subagent_delegate'],
+      removeTools: ['find', 'ls'],
+    }
+    const result = calculateModeTools(fakePi(base), createState(), config)
+    expect(result).toContain('read')
+    expect(result).toContain('grep')
+    expect(result).toContain('subagent_delegate')
+    expect(result).not.toContain(BASH_READONLY_TOOL_NAME)
+    expect(result).not.toContain('find')
+    expect(result).not.toContain('ls')
+  })
+
+  it('handles empty arrays correctly', () => {
+    const base = ['read', 'write', 'bash']
+    const config: ToolConfig = {
+      tools: [],
+      extraTools: [],
+      removeTools: [],
+    }
+    const result = calculateModeTools(fakePi(base), createState(), config)
+    expect(result).toEqual([])
+  })
+
+  it('handles undefined fields correctly', () => {
+    const base = ['read', 'write', 'bash', 'grep']
+    const config: ToolConfig = {
+      tools: undefined,
+      extraTools: undefined,
+      removeTools: undefined,
+    }
+    const result = calculateModeTools(fakePi(base), createState(), config)
+    expect(result).toContain('read')
+    expect(result).toContain(BASH_READONLY_TOOL_NAME)
+    expect(result).toContain('grep')
+  })
+
+  it('removes all write tools from base', () => {
+    const base = ['read', 'write', 'edit', 'bash']
+    const result = calculateModeTools(fakePi(base), createState(), {})
+    for (const tool of WRITE_TOOLS) {
+      expect(result).not.toContain(tool)
+    }
+  })
+
+  it('keeps extraTools even when tools whitelist does not include them', () => {
+    const config: ToolConfig = {
+      tools: ['read', 'grep'],
+      extraTools: ['subagent_delegate'],
+    }
+    const result = calculateModeTools(fakePi([]), createState(), config)
+    expect(result).toContain('read')
+    expect(result).toContain('grep')
+    expect(result).toContain('subagent_delegate')
+  })
+})
+
+describe('applyModeModel', () => {
+  const mockModel: Model<Api> = {
+    provider: 'anthropic',
+    id: 'claude-sonnet-4-5',
+    name: 'Claude Sonnet 4.5',
+  } as Model<Api>
+
+  const mockState: PmSubagentState = {
+    mode: undefined,
+    modeDiffTools: undefined,
+    previousModel: undefined,
+  }
+
+  const setModelMock = vi.fn().mockResolvedValue(true)
+  const findModelMock = vi.fn().mockReturnValue(mockModel)
+  const notifyMock = vi.fn()
+
+  const mockPi = {
+    setModel: setModelMock,
+  } as unknown as Parameters<typeof applyModeModel>[0]
+
+  const mockCtx = {
+    model: mockModel,
+    modelRegistry: {
+      find: findModelMock,
+    },
+    ui: {
+      notify: notifyMock,
+      theme: {},
+    },
+  } as unknown as Parameters<typeof applyModeModel>[2]
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockState.previousModel = undefined
+    setModelMock.mockResolvedValue(true)
+    findModelMock.mockReturnValue(mockModel)
+    mockCtx.model = mockModel
+  })
+
+  it('does nothing when modelRef is undefined', async () => {
+    await applyModeModel(mockPi, mockState, mockCtx, undefined)
+    expect(mockState.previousModel).toBeUndefined()
+    expect(setModelMock).not.toHaveBeenCalled()
+  })
+
+  it('does nothing when modelRef is empty string', async () => {
+    await applyModeModel(mockPi, mockState, mockCtx, '')
+    expect(mockState.previousModel).toBeUndefined()
+    expect(setModelMock).not.toHaveBeenCalled()
+  })
+
+  it('successfully switches to valid model', async () => {
+    const modelRef = 'anthropic/claude-sonnet-4-5'
+    await applyModeModel(mockPi, mockState, mockCtx, modelRef)
+
+    expect(mockState.previousModel).toBe('anthropic/claude-sonnet-4-5')
+    expect(findModelMock).toHaveBeenCalledWith('anthropic', 'claude-sonnet-4-5')
+    expect(setModelMock).toHaveBeenCalledWith(mockModel)
+    expect(notifyMock).not.toHaveBeenCalled()
+  })
+
+  it('shows warning for invalid model ref', async () => {
+    const modelRef = 'invalid-model-ref'
+    await applyModeModel(mockPi, mockState, mockCtx, modelRef)
+
+    expect(mockState.previousModel).toBeUndefined()
+    expect(notifyMock).toHaveBeenCalledWith(
+      'Invalid model ref: invalid-model-ref',
+      'warning',
+    )
+    expect(setModelMock).not.toHaveBeenCalled()
+  })
+
+  it('shows warning when model is not found', async () => {
+    findModelMock.mockReturnValue(undefined)
+    const modelRef = 'anthropic/claude-sonnet-4-5'
+    await applyModeModel(mockPi, mockState, mockCtx, modelRef)
+
+    expect(mockState.previousModel).toBeUndefined()
+    expect(notifyMock).toHaveBeenCalledWith(
+      'Invalid model ref: anthropic/claude-sonnet-4-5',
+      'warning',
+    )
+    expect(setModelMock).not.toHaveBeenCalled()
+  })
+
+  it('shows error when API key is not available', async () => {
+    setModelMock.mockResolvedValue(false)
+    const modelRef = 'anthropic/claude-sonnet-4-5'
+    await applyModeModel(mockPi, mockState, mockCtx, modelRef)
+
+    expect(mockState.previousModel).toBeUndefined()
+    expect(notifyMock).toHaveBeenCalledWith(
+      'No API key for this model',
+      'error',
+    )
+    expect(setModelMock).toHaveBeenCalledWith(mockModel)
+  })
+
+  it('skips recording previousModel if it already exists', async () => {
+    mockState.previousModel = 'openai/gpt-4'
+    const modelRef = 'anthropic/claude-sonnet-4-5'
+    await applyModeModel(mockPi, mockState, mockCtx, modelRef)
+
+    expect(mockState.previousModel).toBe('openai/gpt-4')
+    expect(setModelMock).toHaveBeenCalledWith(mockModel)
+  })
+
+  it('does not record previousModel when current model is undefined', async () => {
+    mockCtx.model = undefined
+    const modelRef = 'anthropic/claude-sonnet-4-5'
+    await applyModeModel(mockPi, mockState, mockCtx, modelRef)
+    expect(mockState.previousModel).toBeUndefined()
+    expect(setModelMock).toHaveBeenCalledWith(mockModel)
+  })
+})