diff --git a/src/main/database/migrations/076-run-pins.ts b/src/main/database/migrations/076-run-pins.ts new file mode 100644 index 00000000..774f816a --- /dev/null +++ b/src/main/database/migrations/076-run-pins.ts @@ -0,0 +1,7 @@ +export const runPinsSchema = ` + alter table task_runs add column run_model text; + alter table task_runs add column run_credential_id text; + + insert into schema_migrations(version, name, applied_at) + values (76, 'run_pins', datetime('now')); +`; diff --git a/src/main/database/schema.ts b/src/main/database/schema.ts index 4db3b6a8..30346d35 100644 --- a/src/main/database/schema.ts +++ b/src/main/database/schema.ts @@ -74,6 +74,7 @@ import { spendSchema } from "./migrations/072-spend"; import { usageWindowTokensSchema } from "./migrations/073-usage-window-tokens"; import { runFailureKindSchema } from "./migrations/074-run-failure-kind"; import { providerCooldownsSchema } from "./migrations/075-provider-cooldowns"; +import { runPinsSchema } from "./migrations/076-run-pins"; interface Migration { version: number; @@ -157,6 +158,7 @@ const MIGRATIONS: Migration[] = [ { version: 73, name: "usage_window_tokens", sql: usageWindowTokensSchema }, { version: 74, name: "run_failure_kind", sql: runFailureKindSchema }, { version: 75, name: "provider_cooldowns", sql: providerCooldownsSchema }, + { version: 76, name: "run_pins", sql: runPinsSchema }, ]; const SCHEMA_VERSION = MIGRATIONS[MIGRATIONS.length - 1].version; diff --git a/src/main/services/provider-response-runner.ts b/src/main/services/provider-response-runner.ts index 2510f9f0..150b1f28 100644 --- a/src/main/services/provider-response-runner.ts +++ b/src/main/services/provider-response-runner.ts @@ -17,7 +17,7 @@ import { selectToolTurn, supportsToolTurn } from "./provider-tool-turns/router"; import { resolveTurnCredential } from "./provider-tool-turns/credential-resolver"; import { currentRunId } from "./run-control"; import { budgetBlock, isPriced, promptChars, providerPrices } from "./spend"; -import { effectiveModel } from "./task-run-preparation"; +import { assertModelSelectable, effectiveModel } from "./task-run-preparation"; import { HEALTH_STATE_HEALTHY, WORKSPACE_SCOPES, @@ -122,11 +122,15 @@ export class ProviderResponseRunner { const agent = taskRunId ? this.deps.db.prepare("select a.id, a.department, a.model_id from task_runs r join agents a on a.id = r.agent_id where r.id = ?").get(taskRunId) as { id: string; department: string; model_id: string | null } : null; - // Azure calls the deployment stored with its credential, so an agent there cannot pick another model. - if (provider.provider_key === "azure-openai" && agent?.model_id && agent.model_id !== provider.model) { - throw new Error(`Azure OpenAI calls the deployment configured on the provider; this agent asks for "${agent.model_id}". Clear the agent's model or set it to "${provider.model}".`); + const pins = taskRunId + ? this.deps.db.prepare("select run_model, run_credential_id from task_runs where id = ?").get(taskRunId) as { run_model: string | null; run_credential_id: string | null } + : null; + if (pins?.run_credential_id && pins.run_credential_id !== provider.credential_id) { + throw new Error("The provider's account was reconnected during this session; the next session continues with the new account."); } - const model = effectiveModel(agent?.model_id ?? null, provider.model); + // Every turn of a session calls the model the session started with (checked then), whatever the configuration says now. + if (!pins?.run_model) assertModelSelectable(provider.provider_key, agent?.model_id ?? null, provider.model); + const model = pins?.run_model ?? effectiveModel(agent?.model_id ?? null, provider.model); if (!model) throw new Error("Provider model is required."); const owner: ProviderRow & { model: string } = { ...provider, model }; const turn = this.deps.toolTurn ?? selectToolTurn(provider.provider_key); diff --git a/src/main/services/task-execution-service.ts b/src/main/services/task-execution-service.ts index 1adb76d7..38366a0c 100644 --- a/src/main/services/task-execution-service.ts +++ b/src/main/services/task-execution-service.ts @@ -37,7 +37,7 @@ import { ConnectorGrantsService } from "./composio/connector-grants"; import { ProviderCooldowns } from "./provider-cooldowns"; import { HeartbeatScheduler } from "./heartbeat-scheduler"; import { - assertExecutionApproval, createIsolationContext, effectiveModel, executionApprovalProblem, loadGoal, loadProject, loadWorkspace, + assertExecutionApproval, assertModelSelectable, createIsolationContext, effectiveModel, executionApprovalProblem, loadGoal, loadProject, loadWorkspace, requireAgentProvider, requireAssignedAgent, requireTask, sameWorkspaceAccess, workspaceAccess, } from "./task-run-preparation"; @@ -175,7 +175,15 @@ export class TaskExecutionService { let providerRequestId: string | null = null; let skillRuntimeSnapshot: SkillRuntimeSnapshotRow | null = null; try { - skillRuntimeSnapshot = this.resolveSkillRuntimeSnapshot(orgId, agent.id, runId, skillBudgetTokens(provider.context_window_tokens), `${task.title}\n${task.description}`, { providerKey: provider.provider_key, model: effectiveModel(agent.model_id, provider.model) }, task.writing_mode === WRITING_MODE.CONVERSATIONAL); + // The session keeps the model and account it started with; a resumed run keeps its original pins. + const pinned = (this.db.prepare("select run_model from task_runs where id = ?").get(runId) as { run_model: string | null }).run_model; + if (pinned === null) { + assertModelSelectable(provider.provider_key, agent.model_id, provider.model); + this.db.prepare("update task_runs set run_model = ?, run_credential_id = ? where id = ?") + .run(effectiveModel(agent.model_id, provider.model), provider.credential_id, runId); + } + const runModel = pinned ?? effectiveModel(agent.model_id, provider.model); + skillRuntimeSnapshot = this.resolveSkillRuntimeSnapshot(orgId, agent.id, runId, skillBudgetTokens(provider.context_window_tokens), `${task.title}\n${task.description}`, { providerKey: provider.provider_key, model: runModel }, task.writing_mode === WRITING_MODE.CONVERSATIONAL); // Watermark for reply delivery and mid-run message polling. const runStartedAt = new Date().toISOString(); const chatContext = loadChatWakeContext(this.db, orgId, agent.id, isolation.secretScopeJson); @@ -227,6 +235,7 @@ export class TaskExecutionService { handoffLines: renderTaskHandoff(this.db, orgId, task.handoff_json), knowledgeLines: relevantKnowledgeLines(this.db, { orgId, agentId: agent.id, projectId: task.project_id, goalId: task.goal_id, taskId: task.id }, `${task.title}\n${task.description}`), priorCompletedRuns: this.completedRunCount(task.id, runId), + runModel, journalLines: runJournalLines(this.db, task.id), subtasks: (this.db.prepare("select title, status = 'done' as done from tasks where parent_task_id = ? order by created_at").all(task.id) as Array<{ title: string; done: number }>) .map((row) => ({ title: row.title, done: row.done === 1 })), diff --git a/src/main/services/task-prompts.ts b/src/main/services/task-prompts.ts index 4dbb9757..9ceac0f3 100644 --- a/src/main/services/task-prompts.ts +++ b/src/main/services/task-prompts.ts @@ -1,6 +1,5 @@ import { WRITING_MODE, type WritingMode } from "../../shared/domain"; import { CEO_FIRST_HIRES, firstHiresFor, leadsPeople } from "./org-building"; -import { effectiveModel } from "./task-run-preparation"; import { canonicalMentionHandle } from "../../shared/mention-handles"; const RECORD_LEARNINGS_INSTRUCTION = @@ -41,6 +40,7 @@ export interface TaskExecutionProviderRow { label: string; provider_key: string; model: string | null; + credential_id: string | null; context_window_tokens: number; } @@ -202,6 +202,8 @@ export function buildTaskPrompt(input: { knowledgeLines: string[]; /** Completed runs this task already spent without reaching done. */ priorCompletedRuns: number; + /** The model this session calls (pinned at its start). */ + runModel: string | null; /** What earlier sessions on this task already did, and where the last one failed. */ journalLines: string[]; /** Names of the tools the model can call this run. */ @@ -221,7 +223,7 @@ export function buildTaskPrompt(input: { `Requires filesystem: ${input.task.requires_filesystem === 1 ? "yes" : "no"}`, `Assignee: ${input.agent.name} (${input.agent.role})`, ...(input.managerName ? [`Reports to: ${input.managerName}`] : []), - `Provider: ${input.provider.label} ${effectiveModel(input.agent.model_id, input.provider.model) ?? ""}`.trim(), + `Provider: ${input.provider.label} ${input.runModel ?? ""}`.trim(), `Execution context: ${input.isolation.executionContextId}`, `Service identity: ${input.isolation.serviceIdentityId}`, ]; diff --git a/src/main/services/task-run-preparation.ts b/src/main/services/task-run-preparation.ts index dd1e8bd7..61be5993 100644 --- a/src/main/services/task-run-preparation.ts +++ b/src/main/services/task-run-preparation.ts @@ -67,6 +67,13 @@ export function sameWorkspaceAccess(mountJson: string, current: Array<{ workspac return JSON.stringify(recorded) === JSON.stringify(current); } +/** Azure calls the deployment stored with its credential, so an agent there cannot pick another model. */ +export function assertModelSelectable(providerKey: string, agentModelId: string | null, providerModel: string | null): void { + if (providerKey === "azure-openai" && agentModelId && agentModelId !== providerModel) { + throw new Error(`Azure OpenAI calls the deployment configured on the provider; this agent asks for "${agentModelId}". Clear the agent's model or set it to "${providerModel}".`); + } +} + /** The model a run calls: the agent's own choice, or else its provider's default. */ export function effectiveModel(agentModelId: string | null, providerModel: string | null): string | null { return agentModelId?.trim() || providerModel; @@ -75,7 +82,7 @@ export function effectiveModel(agentModelId: string | null, providerModel: strin export function requireAgentProvider(db: DatabaseSync, orgId: string, agent: TaskExecutionAgentRow): TaskExecutionProviderRow { if (!agent.model_provider_id) throw new Error("Task assignee requires a model provider."); const provider = db.prepare(` - select id, label, provider_key, model, context_window_tokens + select id, label, provider_key, model, credential_id, context_window_tokens from providers where id = ? and org_id = ? `).get(agent.model_provider_id, orgId) as TaskExecutionProviderRow | undefined; diff --git a/tests/integration/agent-model.test.ts b/tests/integration/agent-model.test.ts index f9e66462..318c72d8 100644 --- a/tests/integration/agent-model.test.ts +++ b/tests/integration/agent-model.test.ts @@ -67,4 +67,51 @@ describe("agent model", () => { await expect(execution.runTask({ orgId, taskId })).rejects.toThrow('Azure OpenAI calls the deployment configured on the provider; this agent asks for "gpt-5.4-mini".'); connection.close(); }); + + test("a session keeps the model and account it started with", async () => { + const models: string[] = []; + let onTurn = (_turn: number) => {}; + let lastPrompt = ""; + const turn = async (input: ProviderToolTurnInput): Promise => { + models.push(input.model); + const first = input.messages[0]; + if (first?.role === "user") lastPrompt = first.text; + onTurn(models.length); + return models.length < 3 + ? { text: "Working.", toolCalls: [{ id: `c${models.length}`, name: "read_messages", arguments: {} }], rawResponseJson: "{}" } + : { text: "done", toolCalls: [], rawResponseJson: "{}" }; + }; + const connection = new DatabaseConnection(join(tmpdir(), `ac-model-pin-${randomUUID()}.db`)); + const organization = new OrganizationService(connection.db); + const providers = new ProviderService(connection.db, providerAdapters, testVault, undefined, turn); + const execution = new TaskExecutionService(connection.db, providers); + let orgId = ""; + let taskId = ""; + let providerId = ""; + connection.transaction(() => { + orgId = organization.createOrganization({ organizationName: "Acme", employeeName: "Dana", mission: "m", goal: "g" }); + providerId = (connection.db.prepare("select id from providers where org_id = ? and provider_key = 'openai'").get(orgId) as { id: string }).id; + providers.saveCredential({ orgId, providerId, secret: "sk-test", model: "gpt-5.4" }); + const agentId = organization.createAgent({ orgId, name: "Ada", role: "Engineer", department: "Engineering", modelProviderId: providerId, modelId: null, budgetCents: 1000, managerAgentId: null }); + taskId = organization.createTask({ orgId, title: "Work", description: "d", status: TASK_STATUS.IN_PROGRESS, assigneeAgentId: agentId, risk: RISK_LEVEL.LOW, projectId: null, goalId: null, workspaceId: null, priority: TASK_PRIORITY.NORMAL, requiresFilesystem: false }); + }); + onTurn = (count) => { if (count === 1) connection.db.prepare("update providers set model = 'gpt-6' where id = ?").run(providerId); }; + await execution.runTask({ orgId, taskId }); + expect(models).toEqual(["gpt-5.4", "gpt-5.4", "gpt-5.4"]); + + // A resumed run builds its prompt and calls with the model it was pinned to, not today's configuration. + models.length = 0; + onTurn = () => {}; + const resumedId = (connection.db.prepare("select id from task_runs where task_id = ?").get(taskId) as { id: string }).id; + connection.db.prepare("update task_runs set state = 'running', completed_at = null where id = ?").run(resumedId); + await execution.executeRun(orgId, resumedId); + expect(models).toEqual(["gpt-5.4", "gpt-5.4", "gpt-5.4"]); + expect(lastPrompt).toContain("Provider: OpenAI gpt-5.4"); + + models.length = 0; + onTurn = (count) => { if (count === 1) connection.db.prepare("update providers set credential_id = 'cred_reconnected' where id = ?").run(providerId); }; + await expect(execution.runTask({ orgId, taskId })).rejects.toThrow("The provider's account was reconnected during this session"); + expect(models).toEqual(["gpt-6"]); + connection.close(); + }); }); diff --git a/tests/unit/domain-services.test.ts b/tests/unit/domain-services.test.ts index df2b53b8..5975159f 100644 --- a/tests/unit/domain-services.test.ts +++ b/tests/unit/domain-services.test.ts @@ -24,7 +24,7 @@ describe("domain services", () => { const connection = createConnection(); const rows = connection.db.prepare("select version from schema_migrations order by version asc").all() as Array<{ version: number }>; - expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75]); + expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76]); connection.close(); });