diff --git a/.changeset/model-roles.md b/.changeset/model-roles.md new file mode 100644 index 00000000..f671b21e --- /dev/null +++ b/.changeset/model-roles.md @@ -0,0 +1,5 @@ +--- +"@pythoughts/pythinker-code": minor +--- + +Add model roles: lock a model alias to the small, implementer, or advisor slot with `/model `, list assignments with `/model roles`, and reference roles as `@small`, `@implementer`, or `@advisor` wherever a subagent model can be set; an assigned implementer role becomes the default model for subagents. diff --git a/apps/pythinker-code/src/tui/commands/config.ts b/apps/pythinker-code/src/tui/commands/config.ts index 3194a940..36078248 100644 --- a/apps/pythinker-code/src/tui/commands/config.ts +++ b/apps/pythinker-code/src/tui/commands/config.ts @@ -41,7 +41,11 @@ import { openFileInExternalEditor, resolveEditorCommand, } from '#/utils/process/external-editor'; -import { LLM_NOT_SET_MESSAGE, NO_ACTIVE_SESSION_MESSAGE } from '../constant/pythinker-tui'; +import { + BUILT_IN_MODEL_ROLES, + LLM_NOT_SET_MESSAGE, + NO_ACTIVE_SESSION_MESSAGE, +} from '#/tui/constant/pythinker-tui'; import { formatErrorMessage } from '../utils/event-payload'; import { showUsage } from './info'; import { setExperimentalFeatures } from './experimental-flags'; @@ -480,6 +484,41 @@ function resolveWorkspaceConfigPath(input: string, workDir: string): string { export async function handleModelCommand(host: SlashCommandHost, args: string): Promise { const requestedAlias = args.trim(); + const tokens = requestedAlias.split(/\s+/u).filter(Boolean); + const config = await host.harness.getConfig({ reload: true }); + const roles = [...new Set([...BUILT_IN_MODEL_ROLES, ...Object.keys(config.modelRoles ?? {})])] + .filter((role) => role.length > 0 && role !== 'default'); + + if (tokens.length === 1 && tokens[0] === 'roles') { + host.showNotice( + 'Model roles', + roles + .map((role) => `${role}: ${config.modelRoles?.[role]?.trim() || '(not set)'}`) + .join('\n'), + ); + return; + } + + const role = tokens[0]; + if (role !== undefined && roles.includes(role)) { + if (tokens.length === 2 && (tokens[1] === 'clear' || tokens[1] === 'none')) { + await host.harness.setConfig({ modelRoles: { [role]: '' } }); + host.showStatus(`Cleared the ${role} model role.`, 'success'); + return; + } + if (tokens.length === 1) { + const picker = showModelPicker(host, config.modelRoles?.[role], undefined, { + assignToRole: role, + }); + if (picker !== undefined) { + void refreshModelsForOpenPicker(host, picker, config.modelRoles?.[role], { + assignToRole: role, + }); + } + return; + } + } + const normalized = normalizeModelChoices(host.state.appState.availableModels); const selectedValue = requestedAlias.length === 0 @@ -524,6 +563,7 @@ async function refreshModelsForOpenPicker( host: SlashCommandHost, picker: TabbedModelSelectorComponent, selectedValue: string | undefined, + options?: { assignToRole?: string }, ): Promise { const availableModels = host.state.appState.availableModels; const normalized = normalizeModelChoices(availableModels); @@ -574,7 +614,7 @@ async function refreshModelsForOpenPicker( } } - showModelPicker(host, refreshedSelected, activeTabId); + showModelPicker(host, refreshedSelected, activeTabId, options); } async function applyEditorChoice(host: SlashCommandHost, value: string): Promise { @@ -615,6 +655,7 @@ export function showModelPicker( host: SlashCommandHost, selectedValue?: string, initialTabId?: string, + options?: { assignToRole?: string }, ): TabbedModelSelectorComponent | undefined { const normalized = normalizeModelChoices(host.state.appState.availableModels); const entries = Object.entries(normalized.models); @@ -646,6 +687,10 @@ export function showModelPicker( initialTabId, onSelect: ({ alias, effort }) => { host.restoreEditor(); + if (options?.assignToRole !== undefined) { + void assignModelRole(host, options.assignToRole, alias); + return; + } void performModelSwitch(host, alias, effort); }, onCancel: () => { @@ -656,6 +701,17 @@ export function showModelPicker( return picker; } +async function assignModelRole(host: SlashCommandHost, role: string, alias: string): Promise { + // Model roles store aliases only; thinking effort stays with the active model. + try { + await host.harness.setConfig({ modelRoles: { [role]: alias } }); + } catch (error) { + host.showError(`Failed to lock the ${role} model: ${formatErrorMessage(error)}`); + return; + } + host.showStatus(`Locked ${alias} as the ${role} model.`, 'success'); +} + async function performModelSwitch(host: SlashCommandHost, alias: string, effort: string): Promise { if (host.state.appState.streamingPhase !== 'idle') { host.showError('Cannot switch models while streaming — press Esc or Ctrl-C first.'); diff --git a/apps/pythinker-code/src/tui/commands/registry.ts b/apps/pythinker-code/src/tui/commands/registry.ts index 7cb77b50..cf0a4bc5 100644 --- a/apps/pythinker-code/src/tui/commands/registry.ts +++ b/apps/pythinker-code/src/tui/commands/registry.ts @@ -164,7 +164,7 @@ export const BUILTIN_SLASH_COMMANDS = [ { name: 'model', aliases: [], - description: 'Switch LLM model', + description: 'Switch model; assign with /model , clear it, or list /model roles', priority: 100, availability: 'always', }, diff --git a/apps/pythinker-code/src/tui/constant/pythinker-tui.ts b/apps/pythinker-code/src/tui/constant/pythinker-tui.ts index cc2f00be..ae1be886 100644 --- a/apps/pythinker-code/src/tui/constant/pythinker-tui.ts +++ b/apps/pythinker-code/src/tui/constant/pythinker-tui.ts @@ -1,5 +1,8 @@ export { OAUTH_LOGIN_REQUIRED_CODE, PRODUCT_NAME } from '#/constant/app'; +/** Canonical model roles offered by `/model `, mirroring agent-core's list across the SDK package boundary. */ +export const BUILT_IN_MODEL_ROLES = ['small', 'implementer', 'advisor'] as const; + export const LLM_NOT_SET_MESSAGE = 'LLM not set, send "/login" to login'; export const NO_ACTIVE_SESSION_MESSAGE = 'No active session. Send /login to login.'; export const CTRL_D_HINT = 'Press Ctrl+D again to exit'; @@ -8,4 +11,3 @@ export const MAIN_AGENT_ID = 'main'; export const OAUTH_LOGIN_REQUIRED_STARTUP_NOTICE = 'OAuth login expired. Send /login to login.'; export const EXIT_CONFIRM_WINDOW_MS = 1500; export const MCP_STATUS_TRANSIENT_DURATION_MS = 750; - diff --git a/apps/pythinker-code/test/tui/commands/model-roles.test.ts b/apps/pythinker-code/test/tui/commands/model-roles.test.ts new file mode 100644 index 00000000..5cffdacb --- /dev/null +++ b/apps/pythinker-code/test/tui/commands/model-roles.test.ts @@ -0,0 +1,161 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { handleModelCommand } from '#/tui/commands/index'; +import type { SlashCommandHost } from '#/tui/commands/dispatch'; + +const ENTER = '\r'; + +interface TestPicker { + handleInput(data: string): void; +} + +function model(name: string) { + return { + provider: 'test', + model: name, + maxContextSize: 200_000, + displayName: name, + capabilities: [], + }; +} + +function makeHost(options: { + currentModel?: string; + availableModels?: Record>; + modelRoles?: Record; + setConfig?: ReturnType; +} = {}) { + const session = { + setModel: vi.fn(async () => {}), + setThinking: vi.fn(async () => {}), + }; + const getConfig = vi.fn(async () => ({ + providers: {}, + modelRoles: options.modelRoles, + })); + const setConfig = options.setConfig ?? vi.fn(async () => {}); + const host = { + state: { + appState: { + model: options.currentModel ?? 'worker', + thinkingLevel: 'off', + streamingPhase: 'idle', + availableModels: options.availableModels ?? { worker: model('worker') }, + }, + editorContainer: { children: [] }, + }, + session, + harness: { getConfig, setConfig }, + authFlow: { + refreshProviderModels: vi.fn(async () => ({ failed: [] })), + }, + setAppState: vi.fn((patch: Record) => Object.assign(host.state.appState, patch)), + showError: vi.fn(), + showStatus: vi.fn(), + showNotice: vi.fn(), + mountEditorReplacement: vi.fn(), + restoreEditor: vi.fn(), + track: vi.fn(), + } as unknown as SlashCommandHost; + return { host, session, setConfig }; +} + +function mountedPicker(host: SlashCommandHost, index = 0): TestPicker { + const mount = host.mountEditorReplacement as ReturnType; + return mount.mock.calls[index]?.[0] as TestPicker; +} + +describe('/model roles', () => { + it('lists every built-in role as not set when no assignments exist', async () => { + const { host } = makeHost(); + + await handleModelCommand(host, 'roles'); + + expect(host.showNotice).toHaveBeenCalledWith( + 'Model roles', + 'small: (not set)\nimplementer: (not set)\nadvisor: (not set)', + ); + }); + + it('locks a selected alias to a role without switching the session model', async () => { + const { host, session, setConfig } = makeHost(); + + await handleModelCommand(host, 'small'); + expect(host.authFlow.refreshProviderModels).toHaveBeenCalledOnce(); + mountedPicker(host).handleInput(ENTER); + + await vi.waitFor(() => { + expect(setConfig).toHaveBeenCalledWith({ modelRoles: { small: 'worker' } }); + }); + expect(session.setModel).not.toHaveBeenCalled(); + }); + + it('keeps role assignment active after the picker refreshes', async () => { + const { host, session, setConfig } = makeHost({ + currentModel: 'parent', + availableModels: { + parent: model('parent'), + worker: model('worker'), + }, + modelRoles: { small: 'worker' }, + }); + vi.mocked(host.mountEditorReplacement).mockImplementation((picker) => { + host.state.editorContainer.children[0] = picker; + }); + vi.mocked(host.authFlow.refreshProviderModels).mockImplementation(async () => { + host.state.appState.availableModels['reviewer'] = model('reviewer'); + return { changed: [], unchanged: [], failed: [] }; + }); + + await handleModelCommand(host, 'small'); + await vi.waitFor(() => { + expect(host.mountEditorReplacement).toHaveBeenCalledTimes(2); + }); + mountedPicker(host, 1).handleInput(ENTER); + + await vi.waitFor(() => { + expect(setConfig).toHaveBeenCalledWith({ modelRoles: { small: 'worker' } }); + }); + expect(session.setModel).not.toHaveBeenCalled(); + }); + + it('reports a role persistence failure without showing success', async () => { + const setConfig = vi.fn(async () => { + throw new Error('disk full'); + }); + const { host } = makeHost({ setConfig }); + + await handleModelCommand(host, 'small'); + mountedPicker(host).handleInput(ENTER); + + await vi.waitFor(() => { + expect(host.showError).toHaveBeenCalledWith(expect.stringContaining('disk full')); + }); + expect(host.showStatus).not.toHaveBeenCalled(); + }); + + it('clears a role with an empty-string tombstone', async () => { + const { host, setConfig } = makeHost({ modelRoles: { small: 'worker' } }); + + await handleModelCommand(host, 'small clear'); + + expect(setConfig).toHaveBeenCalledWith({ modelRoles: { small: '' } }); + }); + + it('keeps an existing model alias on the default switch path', async () => { + const { host, session } = makeHost({ + currentModel: 'parent', + availableModels: { + parent: model('parent'), + worker: model('worker'), + }, + }); + + await handleModelCommand(host, 'worker'); + mountedPicker(host).handleInput(ENTER); + + await vi.waitFor(() => { + expect(session.setModel).toHaveBeenCalledWith('worker'); + }); + }); +}); diff --git a/docs/configuration/config-files.md b/docs/configuration/config-files.md index 7cb082d9..42e9e76f 100644 --- a/docs/configuration/config-files.md +++ b/docs/configuration/config-files.md @@ -76,6 +76,7 @@ Fields in the config file fall into two categories: **top-level scalars** that d | Field | Type | Default | Description | | --- | --- | --- | --- | | `default_model` | `string` | — | Default model alias; must be defined in `models` | +| `model_roles` | `table` | — | Model role assignments → [`model_roles`](#model_roles) | | `default_thinking` | `boolean` | `false` | Whether new sessions enable Thinking (deep reasoning) mode by default; can be toggled from the model menu inside a session. Even when set to `true`, `[thinking].mode = "off"` will still force Thinking off | | `default_permission_mode` | `string` | `manual` | Default permission mode for new sessions; one of `manual` (prompt each time), `yolo` (auto-approve tool actions, but the agent may still ask questions), or `auto` (fully autonomous — the agent decides everything without asking, except a `DynamicWorkflow` call, which still shows its plan for approval) | | `default_plan_mode` | `boolean` | `false` | Whether new sessions start in Plan mode (produce a plan before executing) by default | @@ -94,7 +95,7 @@ Fields in the config file fall into two categories: **top-level scalars** that d | `permission` | `table` | — | Initial permission rules → [`permission`](#permission) | | `hooks` | `array` | — | Lifecycle hooks; see [Hooks](../customization/hooks.md) | -The following sections cover each of the nested tables in turn: `providers`, `models`, `thinking`, `loop_control`, `background`, `experimental`, `services`, and `permission`. +The following sections cover each of the nested tables in turn: `providers`, `models`, `model_roles`, `thinking`, `loop_control`, `background`, `experimental`, `services`, and `permission`. ## `providers` @@ -155,6 +156,24 @@ max_context_size = 1047576 You can also switch models temporarily without touching the config file — by setting `PYTHINKER_MODEL_*` environment variables, the CLI synthesizes a temporary provider in memory that does not persist after restart. See [Define a model from environment variables](./env-vars.md#define-a-model-from-environment-variables-pythinker_model). +## `model_roles` + +Each entry in the `model_roles` table locks a model alias to a named role. The built-in roles are `small`, `implementer`, and `advisor`; any other key except the reserved `default` defines a custom role. Values must be aliases defined in `models`; an empty string clears the role. + +```toml +[model_roles] +small = "haiku" +implementer = "worker-model" +advisor = "reviewer-model" +``` + +Roles take effect in two places: + +- Wherever a subagent model alias is accepted (the `Agent` and `DynamicWorkflow` tool `model` arguments, and agent profile frontmatter), a `@` reference such as `@small` resolves to the locked alias. An unassigned or unresolvable role falls back to the parent agent's model. +- When `implementer` is assigned, it becomes the default model for subagents that do not set an explicit or profile model. Subagents of those subagents inherit the same default. + +Inside the TUI, `/model ` assigns a role from the model picker, `/model clear` (or `/model none`) removes it, and `/model roles` lists the current assignments. See [Slash commands](../reference/slash-commands.md). + ## `thinking` `thinking` sets the global default behavior for Thinking mode. `mode = "off"` forces Thinking off even when the top-level `default_thinking = true`. diff --git a/docs/reference/slash-commands.md b/docs/reference/slash-commands.md index e4e96b99..39d9d796 100644 --- a/docs/reference/slash-commands.md +++ b/docs/reference/slash-commands.md @@ -15,7 +15,7 @@ Some commands are only available in the idle state. Executing these commands whi | `/login` | — | Select an account or platform and log in: Pythinker Code uses OAuth device-code flow; Pythinker Platform uses API key login | No | | `/logout` | — | Clear credentials for the currently selected account | No | | `/provider` | — | Open the interactive provider manager to view, add, and remove configured providers. See [Platforms & Models — `/provider` and provider management](../configuration/providers.md#provider-interactive-provider-management) | Yes | -| `/model` | — | Switch the LLM model used in the current session | Yes | +| `/model` | — | Switch the LLM model used in the current session. `/model ` locks a model alias to a model role (`small`, `implementer`, or `advisor`), `/model clear` (or `/model none`) removes the lock, and `/model roles` lists the current assignments. See [Model roles](../configuration/config-files.md#model_roles) | Yes | | `/settings` | `/config` | Open the settings panel inside the TUI | Yes | | `/experiments` | `/experimental` | Open the experimental feature panel | Yes | | `/permission` | — | Select a permission mode | Yes | diff --git a/packages/agent-core/src/config/index.ts b/packages/agent-core/src/config/index.ts index 41e5ca17..1031bf66 100644 --- a/packages/agent-core/src/config/index.ts +++ b/packages/agent-core/src/config/index.ts @@ -1,4 +1,5 @@ export * from './merge'; +export * from './model-roles'; export * from './path'; export * from './resolve'; export * from './schema'; diff --git a/packages/agent-core/src/config/model-roles.ts b/packages/agent-core/src/config/model-roles.ts new file mode 100644 index 00000000..311dae26 --- /dev/null +++ b/packages/agent-core/src/config/model-roles.ts @@ -0,0 +1,26 @@ +/** Built-in model roles a user can lock a model alias to. */ +export const BUILT_IN_MODEL_ROLES = ['small', 'implementer', 'advisor'] as const; +export type BuiltInModelRole = (typeof BUILT_IN_MODEL_ROLES)[number]; + +interface ModelRoleSource { + modelRoles?: Record; + defaultModel?: string; +} + +/** Resolve a role name to its locked model alias. Empty string means cleared. */ +export function resolveModelRoleAlias( + config: ModelRoleSource | undefined, + role: string, +): string | undefined { + if (role === 'default') return config?.defaultModel; + const alias = config?.modelRoles?.[role]?.trim(); + return alias === '' ? undefined : alias; +} + +/** Expand a "@role" model reference; non-@ strings pass through unchanged. */ +export function expandModelRef( + config: ModelRoleSource | undefined, + ref: string, +): string | undefined { + return ref.startsWith('@') ? resolveModelRoleAlias(config, ref.slice(1)) : ref; +} diff --git a/packages/agent-core/src/config/schema.ts b/packages/agent-core/src/config/schema.ts index 15cc1a1b..a8f93a7e 100644 --- a/packages/agent-core/src/config/schema.ts +++ b/packages/agent-core/src/config/schema.ts @@ -270,10 +270,16 @@ export const McpServerConfigSchema = z.preprocess((raw) => { export type McpServerConfig = z.infer; +const ModelRolesSchema = z.record(z.string(), z.string()).refine( + (roles) => !Object.hasOwn(roles, 'default'), + { message: '"default" is a reserved model role name' }, +); + export const PythinkerConfigSchema = z.object({ providers: z.record(z.string(), ProviderConfigSchema).default({}), defaultProvider: z.string().optional(), defaultModel: z.string().optional(), + modelRoles: ModelRolesSchema.optional(), outputStyle: z.string().trim().min(1).optional(), models: z.record(z.string(), ModelAliasSchema).optional(), thinking: ThinkingConfigSchema.optional(), @@ -319,6 +325,7 @@ export const PythinkerConfigPatchSchema = z providers: z.record(z.string(), ProviderConfigPatchSchema).optional(), defaultProvider: z.string().optional(), defaultModel: z.string().optional(), + modelRoles: ModelRolesSchema.optional(), outputStyle: z.string().trim().min(1).optional(), models: z.record(z.string(), ModelAliasPatchSchema).optional(), thinking: ThinkingConfigPatchSchema.optional(), diff --git a/packages/agent-core/src/config/toml.ts b/packages/agent-core/src/config/toml.ts index 2b3593e1..4c5eed91 100644 --- a/packages/agent-core/src/config/toml.ts +++ b/packages/agent-core/src/config/toml.ts @@ -200,7 +200,7 @@ export function loadRuntimeConfigSafe( } /** Sections keyed by user-chosen names where single entries can be dropped. */ -const ENTRY_KEYED_SECTIONS = new Set(['providers', 'models']); +const ENTRY_KEYED_SECTIONS = new Set(['providers', 'models', 'modelRoles']); interface SalvageResult { readonly config: PythinkerConfig | undefined; @@ -317,6 +317,8 @@ export function transformTomlData(data: Record): Record, ): { modelAlias: string | undefined; thinkingLevel: string | undefined; fastMode: boolean } { - const requested = options.modelAlias ?? profile?.model; + const config = parent.pythinkerConfig; + const requestedRaw = options.modelAlias ?? profile?.model; + const requested = + requestedRaw !== undefined + ? expandModelRef(config, requestedRaw) + : resolveModelRoleAlias(config, 'implementer'); + const tool = options.workflowRunId === undefined ? 'Agent' : 'DynamicWorkflow'; + const requestedDenied = + requested !== undefined && + ((requestedRaw !== undefined && + requestedRaw !== requested && + parent.permission.deniesModelOverride(tool, requestedRaw)) || + parent.permission.deniesModelOverride(tool, requested)); const modelAlias = requested !== undefined && child.config.canResolveModel(requested) && - !parent.permission.deniesModelOverride( - options.workflowRunId === undefined ? 'Agent' : 'DynamicWorkflow', - requested, - ) + !requestedDenied ? requested : parent.config.modelAlias; return { diff --git a/packages/agent-core/src/tools/builtin/collaboration/agent.ts b/packages/agent-core/src/tools/builtin/collaboration/agent.ts index 008409b9..4030a7c5 100644 --- a/packages/agent-core/src/tools/builtin/collaboration/agent.ts +++ b/packages/agent-core/src/tools/builtin/collaboration/agent.ts @@ -62,7 +62,9 @@ function createAgentToolInputSchema(forkContextEnabled: boolean, teamsEnabled = .trim() .min(1) .optional() - .describe('Optional configured model alias for this new subagent'), + .describe( + 'Optional configured model alias for this new subagent. References such as @small, @implementer, @advisor, and @ resolve through configured model roles and fall back to normal precedence when unassigned.', + ), cwd: z .string() .trim() diff --git a/packages/agent-core/src/tools/builtin/collaboration/dynamic-workflow.ts b/packages/agent-core/src/tools/builtin/collaboration/dynamic-workflow.ts index 002cce0f..e29588ae 100644 --- a/packages/agent-core/src/tools/builtin/collaboration/dynamic-workflow.ts +++ b/packages/agent-core/src/tools/builtin/collaboration/dynamic-workflow.ts @@ -102,7 +102,7 @@ export const DynamicWorkflowToolInputSchema = z .min(1) .optional() .describe( - 'Model alias for every subagent in this workflow, so the orchestrator can run on one model while the workers run on a cheaper or faster one. Defaults to the subagent type profile model, then this agent model.', + 'Model alias for every subagent in this workflow, so the orchestrator can run on one model while the workers run on a cheaper or faster one. References such as @small, @implementer, @advisor, and @ resolve through configured model roles and fall back to normal precedence when unassigned. Defaults to the subagent type profile model, then the configured implementer model role, then this agent model.', ), effort: z .string() diff --git a/packages/agent-core/test/config/configs.test.ts b/packages/agent-core/test/config/configs.test.ts index ef34647d..d17a309b 100644 --- a/packages/agent-core/test/config/configs.test.ts +++ b/packages/agent-core/test/config/configs.test.ts @@ -231,6 +231,47 @@ source = { kind = "apiJson", url = "https://registry.example/api.json", apiKey = }); }); + it('round-trips model roles from a fresh config', async () => { + const configPath = join(makeTempDir(), 'model-roles.toml'); + + await writeConfigFile(configPath, { providers: {}, modelRoles: { small: 'haiku' } }); + + const text = await readFile(configPath, 'utf-8'); + expect(text).toContain('[model_roles]'); + expect(text).toContain('small = "haiku"'); + expect(readConfigFile(configPath).modelRoles).toEqual({ small: 'haiku' }); + }); + + it('round-trips reassigned model roles instead of stale raw values', async () => { + const configPath = join(makeTempDir(), 'model-roles-reassigned.toml'); + const config = parseConfigString('[model_roles]\nsmall = "old"\n', configPath); + + await writeConfigFile(configPath, { ...config, modelRoles: { small: 'new' } }); + + expect(readConfigFile(configPath).modelRoles).toEqual({ small: 'new' }); + }); + + it('round-trips cleared model roles instead of stale raw values', async () => { + const configPath = join(makeTempDir(), 'model-roles-cleared.toml'); + const config = parseConfigString('[model_roles]\nsmall = "old"\n', configPath); + + await writeConfigFile(configPath, { ...config, modelRoles: { small: '' } }); + + expect(readConfigFile(configPath).modelRoles).toEqual({ small: '' }); + }); + + it('removes model roles when the role map is cleared', async () => { + const configPath = join(makeTempDir(), 'model-roles-removed.toml'); + await writeFile(configPath, '[model_roles]\nsmall = "old"\n'); + const config = readConfigFile(configPath); + + await writeConfigFile(configPath, { ...config, modelRoles: undefined }); + + const text = await readFile(configPath, 'utf-8'); + expect(text).not.toContain('[model_roles]'); + expect(readConfigFile(configPath).modelRoles).toBeUndefined(); + }); + it('round-trips an API key environment reference without an API key', async () => { const configPath = join(makeTempDir(), 'api-key-env-var.toml'); const config = parseConfigString( @@ -590,6 +631,15 @@ describe('harness config schema and patch merge', () => { expect(merged.raw?.['theme']).toBe('dark'); }); + it('deep-merges model role patches', () => { + const merged = mergeConfigPatch( + { providers: {}, modelRoles: { small: 'x', advisor: 'z' } }, + { modelRoles: { small: 'y' } }, + ); + + expect(merged.modelRoles).toEqual({ small: 'y', advisor: 'z' }); + }); + it('deep-merges experimental config patches', () => { const base = parseConfigString(` [experimental] @@ -886,6 +936,18 @@ max_context_size = -5 expect(result.fileWarnings[0]).toContain('models.broken'); }); + it('drops only the broken model role entry', async () => { + const configPath = await writeTempConfig(`${VALID_TOML} +[model_roles] +small = 123 +implementer = "k2" +`); + const result = loadRuntimeConfigSafe(configPath, {}); + expect(result.config.modelRoles?.['small']).toBeUndefined(); + expect(result.config.modelRoles?.['implementer']).toBe('k2'); + expect(result.fileWarnings[0]).toContain('model_roles.small'); + }); + it('drops the whole hooks list when one hook is invalid', async () => { const configPath = await writeTempConfig(`${VALID_TOML} [[hooks]] diff --git a/packages/agent-core/test/config/model-roles.test.ts b/packages/agent-core/test/config/model-roles.test.ts new file mode 100644 index 00000000..40a18784 --- /dev/null +++ b/packages/agent-core/test/config/model-roles.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, it } from 'vitest'; + +import { + expandModelRef, + PythinkerConfigPatchSchema, + PythinkerConfigSchema, + resolveModelRoleAlias, +} from '../../src/config'; + +describe('model roles', () => { + it('rejects the reserved default role in full and patch configs', () => { + expect( + PythinkerConfigSchema.safeParse({ modelRoles: { default: 'x' } }).success, + ).toBe(false); + expect( + PythinkerConfigPatchSchema.safeParse({ modelRoles: { default: 'x' } }).success, + ).toBe(false); + expect(PythinkerConfigSchema.safeParse({ modelRoles: { custom: 'x' } }).success).toBe(true); + expect(PythinkerConfigPatchSchema.safeParse({ modelRoles: { custom: 'x' } }).success).toBe(true); + }); + + it('resolves assigned roles and treats empty assignments as cleared', () => { + const config = { + modelRoles: { + small: 'haiku-4-5', + empty: '', + whitespace: ' ', + }, + }; + + expect(resolveModelRoleAlias(config, 'small')).toBe('haiku-4-5'); + expect(resolveModelRoleAlias(config, 'missing')).toBeUndefined(); + expect(resolveModelRoleAlias(config, 'empty')).toBeUndefined(); + expect(resolveModelRoleAlias(config, 'whitespace')).toBeUndefined(); + }); + + it('resolves the default role through defaultModel', () => { + expect(resolveModelRoleAlias({ defaultModel: 'opus-5' }, 'default')).toBe('opus-5'); + expect(resolveModelRoleAlias(undefined, 'default')).toBeUndefined(); + }); + + it('expands role references without recursion', () => { + const config = { + modelRoles: { + small: 'haiku-4-5', + nested: '@small', + }, + defaultModel: 'opus-5', + }; + + expect(expandModelRef(config, '@small')).toBe('haiku-4-5'); + expect(expandModelRef(config, '@nested')).toBe('@small'); + expect(expandModelRef(config, '@default')).toBe('opus-5'); + expect(expandModelRef(config, '@unassigned')).toBeUndefined(); + }); + + it('passes non-role aliases through unchanged, including without config', () => { + expect(expandModelRef(undefined, 'gpt-5-codex')).toBe('gpt-5-codex'); + expect(expandModelRef(undefined, '@small')).toBeUndefined(); + }); +}); diff --git a/packages/agent-core/test/session/subagent-host.test.ts b/packages/agent-core/test/session/subagent-host.test.ts index 65fe52b1..6924f260 100644 --- a/packages/agent-core/test/session/subagent-host.test.ts +++ b/packages/agent-core/test/session/subagent-host.test.ts @@ -1754,6 +1754,155 @@ describe('SessionSubagentHost', () => { expect(toolNamesPerCall.at(-1)).toContain('StructuredOutput'); }); + it('uses the implementer role when no explicit or profile model is set', async () => { + const parent = testAgent({ + initialConfig: { providers: {}, modelRoles: { implementer: 'implementer-model' } }, + }); + parent.configure(); + parent.agent.permission.setMode('yolo'); + + const child = testAgent(); + child.configure(); + child.configureRuntimeModel({ type: 'pythinker', apiKey: 'test-key', model: 'implementer-model' }); + child.mockNextResponse({ + type: 'text', + text: 'Completed the delegated implementation with enough detail for the parent agent to continue without repeating the work. '.repeat(2), + }); + + const host = new SessionSubagentHost(fakeSession(parent.agent, child.agent), 'main'); + const handle = await host.spawn({ + profileName: 'coder', + parentToolCallId: 'call_agent', + prompt: 'Implement the feature', + description: 'Implement feature', + runInBackground: false, + signal, + }); + await handle.completion; + + expect(child.agent.config.modelAlias).toBe('implementer-model'); + }); + + it('expands an explicit model role reference', async () => { + const parent = testAgent({ + initialConfig: { providers: {}, modelRoles: { small: 'small-model' } }, + }); + parent.configure(); + parent.agent.permission.setMode('yolo'); + + const child = testAgent(); + child.configure(); + child.configureRuntimeModel({ type: 'pythinker', apiKey: 'test-key', model: 'small-model' }); + child.mockNextResponse({ + type: 'text', + text: 'Completed the delegated task on the selected small model and returned enough detail for the parent to continue. '.repeat(2), + }); + + const host = new SessionSubagentHost(fakeSession(parent.agent, child.agent), 'main'); + const handle = await host.spawn({ + profileName: 'coder', + modelAlias: '@small', + parentToolCallId: 'call_agent', + prompt: 'Investigate the issue', + description: 'Investigate issue', + runInBackground: false, + signal, + }); + await handle.completion; + + expect(child.agent.config.modelAlias).toBe('small-model'); + }); + + it('falls back to the parent model when an explicit role is unassigned', async () => { + const parent = testAgent(); + parent.configure(); + parent.agent.permission.setMode('yolo'); + + const child = testAgent(); + child.configure(); + child.mockNextResponse({ + type: 'text', + text: 'Completed the delegated task with the inherited model and returned enough detail for the parent agent to continue. '.repeat(2), + }); + + const host = new SessionSubagentHost(fakeSession(parent.agent, child.agent), 'main'); + const handle = await host.spawn({ + profileName: 'coder', + modelAlias: '@small', + parentToolCallId: 'call_agent', + prompt: 'Investigate the issue', + description: 'Investigate issue', + runInBackground: false, + signal, + }); + await handle.completion; + + expect(child.agent.config.modelAlias).toBe(parent.agent.config.modelAlias); + }); + + it('falls back to the parent model when the implementer role is unresolvable', async () => { + const parent = testAgent({ + initialConfig: { providers: {}, modelRoles: { implementer: 'no-such-alias' } }, + }); + parent.configure(); + parent.agent.permission.setMode('yolo'); + + const child = testAgent(); + child.configure(); + child.mockNextResponse({ + type: 'text', + text: 'Completed the delegated task with the inherited model and returned enough detail for the parent agent to continue. '.repeat(2), + }); + + const host = new SessionSubagentHost(fakeSession(parent.agent, child.agent), 'main'); + const handle = await host.spawn({ + profileName: 'coder', + parentToolCallId: 'call_agent', + prompt: 'Implement the feature', + description: 'Implement feature', + runInBackground: false, + signal, + }); + await handle.completion; + + expect(child.agent.config.modelAlias).toBe(parent.agent.config.modelAlias); + }); + + it('applies model deny rules to the alias resolved from a role reference', async () => { + const parent = testAgent({ + initialConfig: { providers: {}, modelRoles: { small: 'small-model' } }, + }); + parent.configure(); + parent.agent.permission.setMode('yolo'); + parent.agent.permission.rules.push({ + decision: 'deny', + scope: 'session-runtime', + pattern: 'Agent(model:small-model)', + }); + + const child = testAgent(); + child.configure(); + child.configureRuntimeModel({ type: 'pythinker', apiKey: 'test-key', model: 'small-model' }); + child.mockNextResponse({ + type: 'text', + text: 'Completed the contained delegated task with enough detail for the parent agent to continue without repeating the work. '.repeat(2), + }); + + const host = new SessionSubagentHost(fakeSession(parent.agent, child.agent), 'main'); + const handle = await host.spawn({ + profileName: 'coder', + modelAlias: '@small', + parentToolCallId: 'call_agent', + prompt: 'Investigate the issue', + description: 'Investigate issue', + runInBackground: false, + signal, + }); + await handle.completion; + + expect(child.agent.config.modelAlias).toBe(parent.agent.config.modelAlias); + }); + it('realigns a resumed subagent to the parent agent current model', async () => { const parent = testAgent(); parent.configure();