From d75bfd585128828fda3b58944e2b4be5d307e666 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sat, 18 Apr 2026 22:14:09 -0400 Subject: [PATCH] refactor(model-pool): scope disabled-models by company; cascade on delete (v42) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-connector "enabled/disabled model" settings lived under a global key in the settings key-value table (`disabled_models:${connectorId}`). No company dimension. Deleting a company left the pool choices behind, and creating a new company silently inherited them — the visible symptom was "onboarding's team- model configuration survives company deletion and pollutes the next company's CEO model list." Root fix, not a cleanup hook: the storage becomes a real relational table with `(company_id, connector_id)` primary key and an `ON DELETE CASCADE` FK from `companies`. Company deletion wipes pool config by construction; no cleanup step in `deleteCompany()` can drift. Schema (migration v42): - CREATE TABLE connector_model_pool_settings (company_id, connector_id, disabled_models_json, updated_at), PK (company_id, connector_id), FK companies(id) ON DELETE CASCADE. - Drop every `settings.disabled_models:*` key — these were the bug's residue; there is no right company to attribute them to, so the migration clears them. Users re-select on next onboarding. API surgery — every reader/writer now takes companyId: - db.getDisabledModels(companyId, connectorId) - db.setDisabledModels(companyId, connectorId, models) - db.hasModelPoolSettings(companyId, connectorId) - connector-model-policy: getConnectorModelPolicy(deps, companyId, connectorId); validateConnectorModelSelection(deps, companyId, connectorId, model) - model-dispatch: listEnabledModels(companyId, deps); pickModelForAgent resolves companyId from agent.companyId - orchestration heartbeat/task-runner/turn-context/run-dispatch pass the agent's or run's companyId through to the policy deps - IPC schema: connectorDisabledModelsSchema gains companyId; new connectorModelPoolQuerySchema pairs companyId+connectorId for reads Onboarding becomes commit-on-launch: - useConnectorModelSelection no longer persists via IPC during onboarding. Toggles mutate local state; it exposes `poolSelections` to the wizard. - bootstrapOnboardingSchema gains connectorPoolSelections. The onboarding service saves the company first, then persists each per-connector disabled list scoped to the new company id — CEO model validation runs against the company's freshly-persisted pool, not some stale global. - IPC setConnectorDisabledModels now requires companyId; the only consumer was onboarding, which no longer calls it at all. - Main-side listConnectorModels returns the raw provider catalog (no pool filtering), since that IPC is used from onboarding where no company exists yet; pool filtering happens at dispatch time. saveConnector IPC drops the per-company pool validation — connector config is a global hint, company-scoped enforcement belongs at run dispatch where the company is known. Tests: 603/603 green. Typecheck clean. Integration tests rewritten to set the pool via connectorPoolSelections on bootstrap input, not a pre-bootstrap global setDisabledModels call. --- src/main/connector-model-policy.ts | 8 +- src/main/database-base.ts | 6 +- src/main/database/migrations/v42-v42.ts | 64 ++++++++++ src/main/database/schema.ts | 2 + src/main/database/settings.ts | 52 ++++++-- src/main/ipc-core-admin-handlers.ts | 11 +- src/main/ipc-core-agent-graph-handlers.ts | 4 +- src/main/ipc-core-read-handlers.ts | 9 +- src/main/main-runtime-support.ts | 12 +- src/main/model-dispatch.ts | 8 +- src/main/onboarding-service.ts | 25 +++- src/main/orchestration/heartbeat.ts | 2 +- src/main/orchestration/run-dispatch.ts | 2 +- src/main/orchestration/task-runner.ts | 2 +- src/main/orchestration/turn-context.ts | 12 +- .../company-onboarding/OnboardingWizard.tsx | 4 +- .../useConnectorModelSelection.ts | 112 ++++++++---------- src/shared/contracts/schemas-company.ts | 17 +++ .../connector-model-enforcement.spec.ts | 4 +- .../integration/onboarding-model-pool.spec.ts | 22 ++-- tests/unit/connector-model-policy.test.ts | 24 ++-- tests/unit/ipc-agent-graph-handlers.test.ts | 3 +- 22 files changed, 270 insertions(+), 135 deletions(-) create mode 100644 src/main/database/migrations/v42-v42.ts diff --git a/src/main/connector-model-policy.ts b/src/main/connector-model-policy.ts index a28c8f5e..06f55de1 100644 --- a/src/main/connector-model-policy.ts +++ b/src/main/connector-model-policy.ts @@ -4,7 +4,7 @@ import { resolveModels } from "./model-resolver"; interface ConnectorModelPolicyDeps { getConnector: (id: ConnectorId) => { id: ConnectorId; label?: string | null; model: string | null }; - getDisabledModels: (connectorId: string) => Array<{ id: string; label: string }>; + getDisabledModels: (companyId: string, connectorId: string) => Array<{ id: string; label: string }>; } export interface ConnectorModelPolicyResult { @@ -39,11 +39,12 @@ function listProviderModels(provider: string): Array<{ id: string; label: string export function getConnectorModelPolicy( deps: ConnectorModelPolicyDeps, + companyId: string, connectorId: ConnectorId, ): ConnectorModelPolicyResult { const connector = deps.getConnector(connectorId); const definition = getConnectorDefinition(connectorId); - const disabledModels = new Set(deps.getDisabledModels(connectorId).map((model) => model.id)); + const disabledModels = new Set(deps.getDisabledModels(companyId, connectorId).map((model) => model.id)); const catalog = listProviderModels(definition.provider); const enabledModels = uniqueModels(catalog).filter((model) => !disabledModels.has(model.id)); @@ -58,10 +59,11 @@ export function getConnectorModelPolicy( export function validateConnectorModelSelection( deps: ConnectorModelPolicyDeps, + companyId: string, connectorId: ConnectorId, model: string | null | undefined, ): { ok: true; effectiveModel: string | null } | { ok: false; error: string } { - const policy = getConnectorModelPolicy(deps, connectorId); + const policy = getConnectorModelPolicy(deps, companyId, connectorId); const selectedModel = model?.trim() || null; if (policy.enabledModels.length === 0) { diff --git a/src/main/database-base.ts b/src/main/database-base.ts index 253c3e0f..ffcc24c7 100644 --- a/src/main/database-base.ts +++ b/src/main/database-base.ts @@ -263,9 +263,9 @@ export class AppDatabaseBase { getLogRetentionDays() { return this.settings.getLogRetentionDays(); } purgeExpiredRunLogs() { return this.settings.purgeExpiredRunLogs(); } canOpenPath(targetPath: string) { return this.settings.canOpenPath(targetPath); } - getDisabledModels(connectorId: string) { return this.settings.getDisabledModels(connectorId); } - setDisabledModels(connectorId: string, models: Array<{ id: string; label: string }>) { this.settings.setDisabledModels(connectorId, models); } - hasModelPoolSettings(connectorId: string) { return this.settings.hasModelPoolSettings(connectorId); } + getDisabledModels(companyId: string, connectorId: string) { return this.settings.getDisabledModels(companyId, connectorId); } + setDisabledModels(companyId: string, connectorId: string, models: Array<{ id: string; label: string }>) { this.settings.setDisabledModels(companyId, connectorId, models); } + hasModelPoolSettings(companyId: string, connectorId: string) { return this.settings.hasModelPoolSettings(companyId, connectorId); } getLastBudgetResetMonth(companyId: string) { return this.settings.getLastBudgetResetMonth(companyId); } setLastBudgetResetMonth(companyId: string, month: string) { this.settings.setLastBudgetResetMonth(companyId, month); } diff --git a/src/main/database/migrations/v42-v42.ts b/src/main/database/migrations/v42-v42.ts new file mode 100644 index 00000000..dc1a6f0a --- /dev/null +++ b/src/main/database/migrations/v42-v42.ts @@ -0,0 +1,64 @@ +import log from "electron-log/main.js"; +import type { Migration, MigrationDeps } from "./types"; +import { tableExists } from "./types"; + +/** + * Company-scope the connector model-pool settings. + * + * "Model pool" is the per-connector list of models a user has + * enabled/disabled during onboarding (Settings → Model pool). The + * original storage put it in the key-value `settings` table under + * `disabled_models:${connectorId}` — global, not tied to any + * company. That meant deleting a company left the pool config behind, + * and creating a new company inherited the previous company's + * enabled/disabled choices. + * + * This migration moves the data to a real relational table with a + * cascade FK. Company deletion now naturally wipes the settings — + * no manual cleanup in `deleteCompany()`, no risk of drift. + * + * Existing global entries cannot be assigned to a single company + * (the constraint they were written under had none), so the + * migration drops them. Users re-select their model pool on the + * next onboarding wake. Acceptable: this data is a one-click UI + * preference, not business-critical state, and the bug report was + * literally "these carry over wrongly — clear them." + */ +export function createMigrationsV42V42(deps: MigrationDeps): Migration[] { + const { db } = deps; + return [ + { + version: 42, + name: "scope_connector_model_pool_by_company", + fn: () => { + if (!tableExists(db, "connector_model_pool_settings")) { + db.exec(` + CREATE TABLE connector_model_pool_settings ( + company_id TEXT NOT NULL, + connector_id TEXT NOT NULL, + disabled_models_json TEXT NOT NULL DEFAULT '[]', + updated_at TEXT NOT NULL, + PRIMARY KEY (company_id, connector_id), + FOREIGN KEY (company_id) REFERENCES companies(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_connector_model_pool_company + ON connector_model_pool_settings(company_id); + `); + log.info("[migration v42] created connector_model_pool_settings"); + } + + // Purge stale global keys. Prefix scan is enough — both the + // key format (`disabled_models:`) and the + // semantics (one per connector) mean dropping them all is the + // correct reset. + const deleted = db + .prepare("DELETE FROM settings WHERE key LIKE 'disabled_models:%'") + .run(); + const changes = Number(deleted.changes); + if (changes > 0) { + log.info(`[migration v42] dropped ${changes} legacy disabled_models:* key(s) from settings`); + } + }, + }, + ]; +} diff --git a/src/main/database/schema.ts b/src/main/database/schema.ts index 263cda12..ba9e1e4a 100644 --- a/src/main/database/schema.ts +++ b/src/main/database/schema.ts @@ -6,6 +6,7 @@ import { createMigrationsV35V37 } from "./migrations/v35-v37"; import { createMigrationsV38V39 } from "./migrations/v38-v39"; import { createMigrationsV40V40 } from "./migrations/v40-v40"; import { createMigrationsV41V41 } from "./migrations/v41-v41"; +import { createMigrationsV42V42 } from "./migrations/v42-v42"; import { columnExists, nowIso, @@ -178,6 +179,7 @@ export function runMigrations( ...createMigrationsV38V39(deps), ...createMigrationsV40V40(deps), ...createMigrationsV41V41(deps), + ...createMigrationsV42V42(deps), ]; for (const migration of migrations) { diff --git a/src/main/database/settings.ts b/src/main/database/settings.ts index a20e1ca6..7a8d2282 100644 --- a/src/main/database/settings.ts +++ b/src/main/database/settings.ts @@ -77,22 +77,56 @@ export function buildSettingsMethods( return Number.isFinite(parsed) && parsed >= 1 ? parsed : 90; }, - getDisabledModels(connectorId: string): Array<{ id: string; label: string }> { - const raw = getSetting(`disabled_models:${connectorId}`); - if (!raw) return []; + /** + * Per-(company, connector) disabled model list. + * + * Pre-v42 the key-value `settings` table stored one global entry + * per connector (`disabled_models:${connectorId}`). That let the + * setting leak across companies — deleting a company did not + * clear the pool config, and new companies inherited the old + * selections. v42 moved the data to `connector_model_pool_settings` + * with `company_id` in the primary key and `ON DELETE CASCADE` + * from companies, so company removal wipes pool config by + * construction. + */ + getDisabledModels(companyId: string, connectorId: string): Array<{ id: string; label: string }> { + const row = getDb() + .prepare( + "SELECT disabled_models_json AS json FROM connector_model_pool_settings WHERE company_id = ? AND connector_id = ?", + ) + .get(companyId, connectorId) as { json: string } | undefined; + if (!row) return []; try { - return JSON.parse(raw); + return JSON.parse(row.json); } catch { - return raw.split(",").filter(Boolean).map((id) => ({ id, label: id })); + return []; } }, - setDisabledModels(connectorId: string, models: Array<{ id: string; label: string }>) { - setSetting(`disabled_models:${connectorId}`, JSON.stringify(models)); + setDisabledModels( + companyId: string, + connectorId: string, + models: Array<{ id: string; label: string }>, + ) { + const ts = new Date().toISOString(); + getDb() + .prepare( + `INSERT INTO connector_model_pool_settings (company_id, connector_id, disabled_models_json, updated_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(company_id, connector_id) DO UPDATE SET + disabled_models_json = excluded.disabled_models_json, + updated_at = excluded.updated_at`, + ) + .run(companyId, connectorId, JSON.stringify(models), ts); }, - hasModelPoolSettings(connectorId: string): boolean { - return getSetting(`disabled_models:${connectorId}`) !== null; + hasModelPoolSettings(companyId: string, connectorId: string): boolean { + const row = getDb() + .prepare( + "SELECT 1 AS found FROM connector_model_pool_settings WHERE company_id = ? AND connector_id = ? LIMIT 1", + ) + .get(companyId, connectorId) as { found: number } | undefined; + return !!row; }, getLastBudgetResetMonth(companyId: string): string | null { diff --git a/src/main/ipc-core-admin-handlers.ts b/src/main/ipc-core-admin-handlers.ts index e0a17876..2fd40ce3 100644 --- a/src/main/ipc-core-admin-handlers.ts +++ b/src/main/ipc-core-admin-handlers.ts @@ -7,7 +7,7 @@ import { deleteEntitySchema, openTerminalCommandSchema, } from "@shared/contracts"; -import { validateConnectorConfigurationModel, validateConnectorModelSelection } from "./connector-model-policy"; +import { validateConnectorConfigurationModel } from "./connector-model-policy"; import { reconcileExecutionModeChange } from "./execution-mode-reconciler"; import { fail, ok, type RegisterCoreHandlersDependencies } from "./ipc-core-common"; @@ -94,10 +94,9 @@ export function registerCoreAdminHandlers({ if (!connectorConfigValidation.ok) { return fail("CONNECTOR_MODEL_NOT_ALLOWED", connectorConfigValidation.error); } - const modelSelection = validateConnectorModelSelection(db, parsed.id, parsed.model ?? null); - if (!modelSelection.ok) { - return fail("CONNECTOR_MODEL_NOT_ALLOWED", modelSelection.error); - } + // The per-company disabled-models pool is not consulted here — + // connector config is a global hint, and pool enforcement + // happens at run dispatch time where the company is known. db.saveConnector({ ...parsed, command: existingConnector.command, @@ -133,7 +132,7 @@ export function registerCoreAdminHandlers({ registerHandle(IPC_CHANNELS.setConnectorDisabledModels, async (_event, payload) => { const parsed = connectorDisabledModelsSchema.parse(payload); - db.setDisabledModels(parsed.connectorId, parsed.disabledModels); + db.setDisabledModels(parsed.companyId, parsed.connectorId, parsed.disabledModels); publishDomainChanged(); return ok(true); }); diff --git a/src/main/ipc-core-agent-graph-handlers.ts b/src/main/ipc-core-agent-graph-handlers.ts index 7f6dbf4f..245e9813 100644 --- a/src/main/ipc-core-agent-graph-handlers.ts +++ b/src/main/ipc-core-agent-graph-handlers.ts @@ -94,7 +94,7 @@ export function registerCoreAgentGraphHandlers({ if (connectorIssue) { return fail("CONNECTOR_NOT_READY", connectorIssue); } - const modelSelection = validateConnectorModelSelection(db, parsed.connectorId, parsed.model ?? null); + const modelSelection = validateConnectorModelSelection(db, parsed.companyId, parsed.connectorId, parsed.model ?? null); if (!modelSelection.ok) { return fail("CONNECTOR_MODEL_NOT_ALLOWED", modelSelection.error); } @@ -180,7 +180,7 @@ export function registerCoreAgentGraphHandlers({ if (routingError) { return fail("INVALID_HIRE_ROUTING", routingError); } - const modelSelection = validateConnectorModelSelection(db, parsed.connectorId, parsed.model ?? null); + const modelSelection = validateConnectorModelSelection(db, parsed.companyId, parsed.connectorId, parsed.model ?? null); if (!modelSelection.ok) { return fail("CONNECTOR_MODEL_NOT_ALLOWED", modelSelection.error); } diff --git a/src/main/ipc-core-read-handlers.ts b/src/main/ipc-core-read-handlers.ts index 625c8a8c..95119f9c 100644 --- a/src/main/ipc-core-read-handlers.ts +++ b/src/main/ipc-core-read-handlers.ts @@ -4,6 +4,7 @@ import { z } from "zod"; import { IPC_CHANNELS, connectorInputSchema, + connectorModelPoolQuerySchema, entityGetterSchema, metricsQuerySchema, settingsSchema, @@ -105,12 +106,12 @@ export function registerCoreReadHandlers({ } }); - registerHandle(IPC_CHANNELS.getConnectorDisabledModels, (_event, connectorId) => { + registerHandle(IPC_CHANNELS.getConnectorDisabledModels, (_event, payload) => { try { - const parsed = connectorInputSchema.shape.id.parse(connectorId); - return ok(db.getDisabledModels(parsed)); + const parsed = connectorModelPoolQuerySchema.parse(payload); + return ok(db.getDisabledModels(parsed.companyId, parsed.connectorId)); } catch (err) { - logger.warn(`[ipc] getConnectorDisabledModels failed for connector ${String(connectorId)}:`, err instanceof Error ? err.message : String(err)); + logger.warn(`[ipc] getConnectorDisabledModels failed:`, err instanceof Error ? err.message : String(err)); return ok([]); } }); diff --git a/src/main/main-runtime-support.ts b/src/main/main-runtime-support.ts index 08a8bd36..0ebe9519 100644 --- a/src/main/main-runtime-support.ts +++ b/src/main/main-runtime-support.ts @@ -3,7 +3,6 @@ import type { AgentApiServer } from "./api-server"; import { initializeAgentTokenSigningKey } from "./agent-auth"; import { bootstrapApplication } from "./bootstrap-application"; import { openGenericTerminal as openGenericTerminalImpl } from "./connector-auth-service"; -import { getConnectorModelPolicy } from "./connector-model-policy"; import type { ConnectorId } from "@shared/types"; import type { AppDatabase } from "./database"; import { seedDefaultAutomationRules as seedDefaultAutomationRulesImpl } from "./default-automation-rules"; @@ -154,8 +153,15 @@ export function createMainRuntimeSupport(state: MainRuntimeState) { }, notify: (options) => requireIpcSupport().notify(options), listConnectorModels: async (connectorId) => { - const policy = getConnectorModelPolicy(requireDb(), connectorId as ConnectorId); - return policy.enabledModels; + // This IPC returns the raw provider catalog (pi-ai + curated + // Fireworks additions). Pool / disabled-model filtering is a + // company-scoped concern and happens in the renderer once the + // user has selected a company; the onboarding path uses this + // catalog directly because no company exists yet. + const { getConnectorDefinition } = await import("./connectors"); + const { resolveModels } = await import("./model-resolver"); + const definition = getConnectorDefinition(connectorId as ConnectorId); + return resolveModels(definition.provider).map((m) => ({ id: m.id, label: m.name })); }, logger: log, }); diff --git a/src/main/model-dispatch.ts b/src/main/model-dispatch.ts index 23c8d023..43b6baaf 100644 --- a/src/main/model-dispatch.ts +++ b/src/main/model-dispatch.ts @@ -48,7 +48,7 @@ export interface AgentAssignment { export interface PickModelDeps { getConnector: (id: ConnectorId) => ConnectorRecord; - getDisabledModels: (connectorId: string) => Array<{ id: string; label: string }>; + getDisabledModels: (companyId: string, connectorId: string) => Array<{ id: string; label: string }>; /** * Thin projection of the CEO for fallback. Returns `null` when the * company has no CEO yet (pre-bootstrap). The picker only needs enough @@ -85,7 +85,7 @@ const TIER_RANK: Record = { * model list. Returns a flat list of (provider, connectorId, modelId) triples * — the concrete universe the picker chooses from. */ -export function listEnabledModels(deps: PickModelDeps): EnabledModel[] { +export function listEnabledModels(companyId: string, deps: PickModelDeps): EnabledModel[] { const result: EnabledModel[] = []; for (const connector of deps.readConnectors()) { if (getConnectorExecutionReadinessIssue(connector)) continue; @@ -96,7 +96,7 @@ export function listEnabledModels(deps: PickModelDeps): EnabledModel[] { log.warn(`[model-dispatch] Skipping unknown connector ${connector.id}:`, err instanceof Error ? err.message : String(err)); continue; } - const policy = getConnectorModelPolicy(deps, connector.id as ConnectorId); + const policy = getConnectorModelPolicy(deps, companyId, connector.id as ConnectorId); for (const model of policy.enabledModels) { result.push({ provider: definition.provider, @@ -167,7 +167,7 @@ export function pickModelForAgent( return { connectorId: agent.connectorId, modelId: agent.model ?? null }; } - const enabled = listEnabledModels(deps); + const enabled = listEnabledModels(agent.companyId, deps); // Single-model install: there's nothing to pick from. Use it. if (enabled.length === 1) { diff --git a/src/main/onboarding-service.ts b/src/main/onboarding-service.ts index f06db3f8..94e3bc06 100644 --- a/src/main/onboarding-service.ts +++ b/src/main/onboarding-service.ts @@ -24,6 +24,10 @@ interface BootstrapInput { goalDescription: string; connectorId?: ConnectorId; connectorModel?: string; + connectorPoolSelections?: Array<{ + connectorId: ConnectorId; + disabledModels: Array<{ id: string; label: string }>; + }>; } interface OnboardingDependencies { @@ -215,10 +219,6 @@ export function createOnboardingService({ function bootstrapOnboarding(input: BootstrapInput): OnboardingBootstrapSummary { const connectors = db.listSnapshot({ companyId: null }).connectors; const selectedConnectorId = resolveOnboardingConnector(connectors, input.connectorId); - const modelSelection = validateConnectorModelSelection(db, selectedConnectorId, input.connectorModel ?? null); - if (!modelSelection.ok) { - throw new Error(modelSelection.error); - } const companyId = db.saveCompany({ name: input.companyName.trim(), @@ -228,6 +228,23 @@ export function createOnboardingService({ reviewDeliverables: input.reviewDeliverables, executionMode: input.executionMode ?? "ai_ceo", }); + + // Persist the user's per-connector disabled-model selections + // against the freshly-created company id. Onboarding is entirely + // pre-commit; the pool settings arrive in the input payload and + // land in the DB only once a company exists to scope them. + for (const pool of input.connectorPoolSelections ?? []) { + db.setDisabledModels(companyId, pool.connectorId, pool.disabledModels); + } + + // Validate the CEO model selection against the now-persisted pool + // scoped to this new company. Must run AFTER the pool is stored, + // otherwise the validator sees an empty pool and accepts anything. + const modelSelection = validateConnectorModelSelection(db, companyId, selectedConnectorId, input.connectorModel ?? null); + if (!modelSelection.ok) { + throw new Error(modelSelection.error); + } + const workspaceId = createDefaultWorkspace(companyId, input.workspacePath); const goalTitle = input.goalTitle.trim(); diff --git a/src/main/orchestration/heartbeat.ts b/src/main/orchestration/heartbeat.ts index c5f421e0..4cd5e6b2 100644 --- a/src/main/orchestration/heartbeat.ts +++ b/src/main/orchestration/heartbeat.ts @@ -44,7 +44,7 @@ export function createHeartbeat( } const assignment = pickModelForAgent(agent, { getConnector: (id) => db.getConnector(id), - getDisabledModels: (id) => db.getDisabledModels(id), + getDisabledModels: (id) => db.getDisabledModels(agent.companyId, id), getCeoAssignment: (cid) => { const ceo = db.getCeoForCompany(cid); if (!ceo) return null; diff --git a/src/main/orchestration/run-dispatch.ts b/src/main/orchestration/run-dispatch.ts index c82222af..ac849108 100644 --- a/src/main/orchestration/run-dispatch.ts +++ b/src/main/orchestration/run-dispatch.ts @@ -47,7 +47,7 @@ export function createRunDispatch(deps: OrchestrationDependencies) { logger.info(`[run] Agent ${agent.name} connector ${connector.id} blocked: ${issue}`); return "CONNECTOR_NOT_READY"; } - const modelSelection = validateConnectorModelSelection(db, connector.id, agent.model ?? null); + const modelSelection = validateConnectorModelSelection(db, agent.companyId, connector.id, agent.model ?? null); if (!modelSelection.ok) { logger.info(`[run] Agent ${agent.name} connector model blocked: ${modelSelection.error}`); return "CONNECTOR_MODEL_NOT_ALLOWED"; diff --git a/src/main/orchestration/task-runner.ts b/src/main/orchestration/task-runner.ts index 7f700237..b3907cb7 100644 --- a/src/main/orchestration/task-runner.ts +++ b/src/main/orchestration/task-runner.ts @@ -49,7 +49,7 @@ export function createTaskRunner( } const assignment = pickModelForAgent(agent, { getConnector: (id) => db.getConnector(id), - getDisabledModels: (id) => db.getDisabledModels(id), + getDisabledModels: (id) => db.getDisabledModels(agent.companyId, id), getCeoAssignment: (cid) => { const ceo = db.getCeoForCompany(cid); if (!ceo) return null; diff --git a/src/main/orchestration/turn-context.ts b/src/main/orchestration/turn-context.ts index 74884067..6d015329 100644 --- a/src/main/orchestration/turn-context.ts +++ b/src/main/orchestration/turn-context.ts @@ -132,9 +132,10 @@ function resolveSkills( function resolveEnabledConnectorModels( db: AppDatabase, + companyId: string, connector: { id: string; model: string | null }, ): string[] { - return getConnectorModelPolicy(db, connector.id as ConnectorId).enabledModels.map((model) => model.id); + return getConnectorModelPolicy(db, companyId, connector.id as ConnectorId).enabledModels.map((model) => model.id); } function resolveSession( @@ -261,16 +262,17 @@ function filterPendingApprovals( function buildConnectorCatalog( db: AppDatabase, + companyId: string, connectors: RunDispatchParams["snapshot"]["connectors"], ) { return connectors .filter((c) => c.status === "ready") - .filter((c) => db.hasModelPoolSettings(c.id)) + .filter((c) => db.hasModelPoolSettings(companyId, c.id)) .map((c) => ({ id: c.id, label: c.label, status: c.status, - enabledModels: resolveEnabledConnectorModels(db, c), + enabledModels: resolveEnabledConnectorModels(db, companyId, c), })) .filter((c) => c.enabledModels.length > 0); } @@ -325,7 +327,7 @@ export function buildRunPayload(params: RunDispatchParams, deps: TurnContextDeps relevantKnowledge: filterRelevantKnowledge(snapshot.knowledgeBase, companyId), activeSprint: snapshot.sprints.find((sprint) => sprint.companyId === companyId && sprint.status === "active") ?? null, recentDocuments: filterRecentDocuments(snapshot.documents, companyId, agent.id), - allConnectors: buildConnectorCatalog(db, snapshot.connectors), + allConnectors: buildConnectorCatalog(db, companyId, snapshot.connectors), sessionHandoff, skills, companySkillCatalog: db.listCompanySkills(companyId).map((skill) => ({ @@ -341,7 +343,7 @@ export function buildRunPayload(params: RunDispatchParams, deps: TurnContextDeps const thinkingEffort = extractThinkingEffort(agent.metadataJson, deps.logger, agent.id); const modelForRun = params.resolvedModel !== undefined ? params.resolvedModel : agent.model; - const modelSelection = validateConnectorModelSelection(db, connector.id, modelForRun); + const modelSelection = validateConnectorModelSelection(db, companyId, connector.id, modelForRun); if (!modelSelection.ok) { throw new Error(modelSelection.error); } diff --git a/src/renderer/components/company-onboarding/OnboardingWizard.tsx b/src/renderer/components/company-onboarding/OnboardingWizard.tsx index b731c77f..105b8f93 100644 --- a/src/renderer/components/company-onboarding/OnboardingWizard.tsx +++ b/src/renderer/components/company-onboarding/OnboardingWizard.tsx @@ -101,6 +101,7 @@ export function CompanyOnboarding({ totalEnabledModels, toggleConnectorModel, setAllConnectorModels, + poolSelections, } = useConnectorModelSelection(connectors, currentLocale, t, setError); const pickDirectory = useCallback(async (): Promise => { @@ -156,6 +157,7 @@ export function CompanyOnboarding({ goalDescription: values.goalDescription.trim(), connectorId: selectedConnector?.id as ConnectorId | undefined, connectorModel: selectedModel.trim(), + connectorPoolSelections: poolSelections, }), ); setLaunchSummary(result); @@ -168,7 +170,7 @@ export function CompanyOnboarding({ setLoading(false); } }, - [companyValues, executionMode, loading, selectedConnector?.id, selectedModel, t], + [companyValues, executionMode, loading, poolSelections, selectedConnector?.id, selectedModel, t], ); const finishLaunch = useCallback(async () => { diff --git a/src/renderer/components/company-onboarding/useConnectorModelSelection.ts b/src/renderer/components/company-onboarding/useConnectorModelSelection.ts index 04dbbee3..5e883bc9 100644 --- a/src/renderer/components/company-onboarding/useConnectorModelSelection.ts +++ b/src/renderer/components/company-onboarding/useConnectorModelSelection.ts @@ -6,7 +6,6 @@ import { import { isConnectorExecutionReady } from "@shared/connector-policy"; import type { AppLocale } from "@shared/locale"; import type { ConnectorId, ProfileSnapshot } from "@shared/types"; -import { unwrap } from "../../lib/desktop"; import type { TranslationFn } from "../../i18n"; /** @@ -36,6 +35,15 @@ function connectorCatalogsEqual( return true; } +/** + * Per-connector disabled-model list, shaped to feed straight into + * `bootstrapOnboarding`'s input. Empty array means "every model enabled." + */ +export interface ConnectorPoolSelection { + connectorId: ConnectorId; + disabledModels: Array<{ id: string; label: string }>; +} + interface UseConnectorModelSelectionResult { executionReadyConnectors: ProfileSnapshot["connectors"]; hasAtLeastOneReady: boolean; @@ -53,8 +61,13 @@ interface UseConnectorModelSelectionResult { loadingModels: boolean; totalEnabledModels: number; selectedConnectorEnabledModels: number; - toggleConnectorModel: (connectorId: ConnectorId, modelId: string) => Promise; - setAllConnectorModels: (connectorId: ConnectorId, enabled: boolean) => Promise; + toggleConnectorModel: (connectorId: ConnectorId, modelId: string) => void; + setAllConnectorModels: (connectorId: ConnectorId, enabled: boolean) => void; + /** + * Snapshot of the user's per-connector disabled selections, ready + * to be committed with the company creation in `bootstrapOnboarding`. + */ + poolSelections: ConnectorPoolSelection[]; } export function useConnectorModelSelection( @@ -91,6 +104,12 @@ export function useConnectorModelSelection( executionReadyConnectors[0] ?? null; + // Onboarding is a pre-commit phase — the company does not exist in + // the DB yet, so there is no persisted disabled-models list to read. + // Initial state is "every model enabled"; user toggles mutate local + // state only, and `bootstrapOnboarding` persists the final selection + // when the company is created. This removes the global-leak bug + // where a deleted company's disabled list inherited into the next. const loadConnectorModels = useCallback(async (connector: ProfileSnapshot["connectors"][number]) => { if (connector.id.startsWith("claude_compat")) { const configuredModels = (connector.model ?? "") @@ -98,18 +117,12 @@ export function useConnectorModelSelection( .map((model) => model.trim()) .filter(Boolean) .map((model) => ({ id: model, label: model })); - const disabledResult = await window.agentCompany.getConnectorDisabledModels(connector.id); - const disabledModels = disabledResult.ok ? disabledResult.data : []; - return mergeConnectorModelEntries(configuredModels, disabledModels); + return mergeConnectorModelEntries(configuredModels, []); } - const [enabledResult, disabledResult] = await Promise.all([ - window.agentCompany.listConnectorModels(connector.id), - window.agentCompany.getConnectorDisabledModels(connector.id), - ]); + const enabledResult = await window.agentCompany.listConnectorModels(connector.id); const enabledModels = enabledResult.ok ? enabledResult.data : []; - const disabledModels = disabledResult.ok ? disabledResult.data : []; - return mergeConnectorModelEntries(enabledModels, disabledModels); + return mergeConnectorModelEntries(enabledModels, []); }, []); useEffect(() => { @@ -215,24 +228,17 @@ export function useConnectorModelSelection( [connectorModels], ); - const toggleConnectorModel = useCallback(async (connectorId: ConnectorId, modelId: string) => { + const toggleConnectorModel = useCallback((connectorId: ConnectorId, modelId: string) => { const connectorModels = connectorModelsById[connectorId]; - if (!connectorModels) { - return; - } + if (!connectorModels) return; const currentEntry = connectorModels.find((model) => model.id === modelId); - if (!currentEntry) { - return; - } + if (!currentEntry) return; const updatedModels = connectorModels.map((model) => model.id === modelId ? { ...model, enabled: !model.enabled } : model, ); - const updatedModelMap = { - ...connectorModelsById, - [connectorId]: updatedModels, - }; + const updatedModelMap = { ...connectorModelsById, [connectorId]: updatedModels }; const nextEnabledCount = Object.values(updatedModelMap).reduce( (count, models) => count + models.filter((model) => model.enabled).length, 0, @@ -247,59 +253,40 @@ export function useConnectorModelSelection( return; } - const previousModels = connectorModelsById; setConnectorModelsById(updatedModelMap); setError(null); - - try { - unwrap( - await window.agentCompany.setConnectorDisabledModels({ - connectorId, - disabledModels: updatedModels - .filter((model) => !model.enabled) - .map((model) => ({ id: model.id, label: model.label })), - }), - ); - } catch (error) { - setConnectorModelsById(previousModels); - setError(error instanceof Error ? error.message : t("onboarding.unableToBootstrap")); - } + // Unused dependency: t is left in the closure list below for + // consistency with the error-copy callbacks; avoids a hook + // re-identity when the translator ref swaps locales. + void t; }, [connectorModelsById, currentLocale, setError, t]); /** - * Bulk-flip every model of one connector on or off. Uses the same - * single-call `setConnectorDisabledModels` IPC as the individual - * toggle so the UI and the DB agree after one round-trip. - * + * Bulk-flip every model of one connector on or off. Pure local state + * during onboarding; persistence happens at `bootstrapOnboarding`. * Disabling-all is allowed even if it zeroes the global enabled - * count — the continue gate at the wizard level catches that and - * blocks the next step, matching the single-toggle behavior. + * count — the continue gate at the wizard level catches that. */ const setAllConnectorModels = useCallback( - async (connectorId: ConnectorId, enabled: boolean) => { + (connectorId: ConnectorId, enabled: boolean) => { const connectorModels = connectorModelsById[connectorId]; if (!connectorModels || connectorModels.length === 0) return; - const updatedModels = connectorModels.map((model) => ({ ...model, enabled })); - const previousModels = connectorModelsById; setConnectorModelsById({ ...connectorModelsById, [connectorId]: updatedModels }); setError(null); - - try { - unwrap( - await window.agentCompany.setConnectorDisabledModels({ - connectorId, - disabledModels: enabled - ? [] - : updatedModels.map((model) => ({ id: model.id, label: model.label })), - }), - ); - } catch (error) { - setConnectorModelsById(previousModels); - setError(error instanceof Error ? error.message : t("onboarding.unableToBootstrap")); - } }, - [connectorModelsById, setError, t], + [connectorModelsById, setError], + ); + + const poolSelections = useMemo( + () => + Object.entries(connectorModelsById).map(([connectorId, models]) => ({ + connectorId: connectorId as ConnectorId, + disabledModels: models + .filter((model) => !model.enabled) + .map((model) => ({ id: model.id, label: model.label })), + })), + [connectorModelsById], ); return { @@ -318,5 +305,6 @@ export function useConnectorModelSelection( selectedConnectorEnabledModels, toggleConnectorModel, setAllConnectorModels, + poolSelections, }; } diff --git a/src/shared/contracts/schemas-company.ts b/src/shared/contracts/schemas-company.ts index 5d8667a2..0e5a9a04 100644 --- a/src/shared/contracts/schemas-company.ts +++ b/src/shared/contracts/schemas-company.ts @@ -60,10 +60,16 @@ export const connectorInputSchema = z.object({ }); export const connectorDisabledModelsSchema = z.object({ + companyId: z.string().uuid(), connectorId: connectorInputSchema.shape.id, disabledModels: z.array(z.object({ id: z.string(), label: z.string() })), }); +export const connectorModelPoolQuerySchema = z.object({ + companyId: z.string().uuid(), + connectorId: connectorInputSchema.shape.id, +}); + export const workspaceInputSchema = z.object({ id: optionalId, companyId: z.string().uuid(), @@ -93,6 +99,17 @@ export const bootstrapOnboardingSchema = z.object({ goalDescription: z.string().default(""), connectorId: connectorInputSchema.shape.id.optional(), connectorModel: z.string().default(""), + // Per-connector disabled-model selections gathered during onboarding. + // Persisted against the new company's id once it is created, so they + // are scoped correctly from the first moment the company exists. + connectorPoolSelections: z + .array( + z.object({ + connectorId: connectorInputSchema.shape.id, + disabledModels: z.array(z.object({ id: z.string(), label: z.string() })), + }), + ) + .default([]), }); export const socialAccountInputSchema = z.object({ diff --git a/tests/integration/connector-model-enforcement.spec.ts b/tests/integration/connector-model-enforcement.spec.ts index 2f914b0d..3399f71f 100644 --- a/tests/integration/connector-model-enforcement.spec.ts +++ b/tests/integration/connector-model-enforcement.spec.ts @@ -113,7 +113,7 @@ describe("connector model enforcement", () => { status: "todo", }); - db.setDisabledModels("openai-codex", [{ id: "gpt-5.4", label: "gpt-5.4" }]); + db.setDisabledModels(companyId, "openai-codex", [{ id: "gpt-5.4", label: "gpt-5.4" }]); const { taskRunner } = createOrchestrationHarness(db); const result = taskRunner.queueTaskRun(taskId, "manual", "board"); @@ -187,7 +187,7 @@ describe("connector model enforcement", () => { const checkoutSuccess = db.checkoutTask(taskId, run.id); expect(checkoutSuccess).toBe(true); - db.setDisabledModels("openai-codex", [{ id: "gpt-5.4", label: "gpt-5.4" }]); + db.setDisabledModels(companyId, "openai-codex", [{ id: "gpt-5.4", label: "gpt-5.4" }]); const { taskRunner } = createOrchestrationHarness(db); taskRunner.resumeQueuedRuns(); diff --git a/tests/integration/onboarding-model-pool.spec.ts b/tests/integration/onboarding-model-pool.spec.ts index 9ea4edcb..73601396 100644 --- a/tests/integration/onboarding-model-pool.spec.ts +++ b/tests/integration/onboarding-model-pool.spec.ts @@ -105,8 +105,6 @@ describe("onboarding model pool integration", () => { version: "test", lastError: null, }); - db.setDisabledModels("openai-codex", [{ id: disabledModel, label: disabledModel }]); - const summary = createBootstrapHarness(db).bootstrapOnboarding({ companyName: "Model Governance Co", companyDescription: "Tests onboarding model governance", @@ -118,6 +116,9 @@ describe("onboarding model pool integration", () => { goalDescription: "Ensure onboarding connector models flow into runtime decisions.", connectorId: "openai-codex", connectorModel: enabledModel, + connectorPoolSelections: [ + { connectorId: "openai-codex", disabledModels: [{ id: disabledModel, label: disabledModel }] }, + ], }); if (!summary.leadAgentId) { @@ -127,7 +128,7 @@ describe("onboarding model pool integration", () => { const agent = db.getAgent(summary.leadAgentId); expect(agent.connectorId).toBe("openai-codex"); expect(agent.model).toBe(enabledModel); - expect(db.getDisabledModels("openai-codex")).toEqual([{ id: disabledModel, label: disabledModel }]); + expect(db.getDisabledModels(summary.companyId, "openai-codex")).toEqual([{ id: disabledModel, label: disabledModel }]); expect(db.getCurrentCompanyId()).toBe(summary.companyId); const payload = buildKickoffPayload(db, summary); @@ -167,8 +168,6 @@ describe("onboarding model pool integration", () => { version: "test", lastError: null, }); - db.setDisabledModels("anthropic", [{ id: disabledModel, label: disabledModel }]); - const summary = createBootstrapHarness(db).bootstrapOnboarding({ companyName: "Anthropic Governance Co", companyDescription: "Tests anthropic model governance", @@ -180,6 +179,9 @@ describe("onboarding model pool integration", () => { goalDescription: "Ensure anthropic model pools survive onboarding and runtime prompt assembly.", connectorId: "anthropic", connectorModel: enabledModel, + connectorPoolSelections: [ + { connectorId: "anthropic", disabledModels: [{ id: disabledModel, label: disabledModel }] }, + ], }); const payload = buildKickoffPayload(db, summary); @@ -259,8 +261,6 @@ describe("onboarding model pool integration", () => { version: "test", lastError: null, }); - db.setDisabledModels("openai-codex", [{ id: defaultModel, label: defaultModel }]); - const summary = createBootstrapHarness(db).bootstrapOnboarding({ companyName: "Default Reroute Co", companyDescription: "Tests runtime reroute when the connector default is disabled", @@ -272,6 +272,9 @@ describe("onboarding model pool integration", () => { goalDescription: "Ensure runtime picks the first enabled model when connector default is disabled.", connectorId: "openai-codex", connectorModel: "", + connectorPoolSelections: [ + { connectorId: "openai-codex", disabledModels: [{ id: defaultModel, label: defaultModel }] }, + ], }); const payload = buildKickoffPayload(db, summary); @@ -300,8 +303,6 @@ describe("onboarding model pool integration", () => { version: "test", lastError: null, }); - db.setDisabledModels("openai-codex", [{ id: disabledModel, label: disabledModel }]); - expect(() => createBootstrapHarness(db).bootstrapOnboarding({ companyName: "Rejected Model Co", companyDescription: "Tests invalid onboarding model rejection", @@ -313,6 +314,9 @@ describe("onboarding model pool integration", () => { goalDescription: "Invalid explicit models should fail before the CEO is created.", connectorId: "openai-codex", connectorModel: disabledModel, + connectorPoolSelections: [ + { connectorId: "openai-codex", disabledModels: [{ id: disabledModel, label: disabledModel }] }, + ], })).toThrow(`${codexDefinition.label} does not allow model "${disabledModel}".`); db.close(); diff --git a/tests/unit/connector-model-policy.test.ts b/tests/unit/connector-model-policy.test.ts index af7d5a4c..d5063af2 100644 --- a/tests/unit/connector-model-policy.test.ts +++ b/tests/unit/connector-model-policy.test.ts @@ -12,13 +12,13 @@ function createDeps(overrides: { label: "Anthropic", model: overrides.connectorModel ?? null, }), - getDisabledModels: () => overrides.disabledModels ?? [], + getDisabledModels: (_c: string, _k: string) => overrides.disabledModels ?? [], }; } describe("connector model policy", () => { it("excludes disabled models from the enabled set", () => { - const policy = getConnectorModelPolicy(createDeps(), "anthropic"); + const policy = getConnectorModelPolicy(createDeps(), "test-co", "anthropic"); const allIds = policy.enabledModels.map((model) => model.id); if (allIds.length === 0) { throw new Error("anthropic provider should expose at least one model from pi-ai."); @@ -26,33 +26,31 @@ describe("connector model policy", () => { const firstModel = allIds[0]; const filteredPolicy = getConnectorModelPolicy( - createDeps({ disabledModels: [{ id: firstModel, label: firstModel }] }), - "anthropic", + createDeps({ disabledModels: [{ id: firstModel, label: firstModel }] }), "test-co", "anthropic", ); expect(filteredPolicy.enabledModels.some((model) => model.id === firstModel)).toBe(false); }); it("accepts an explicit enabled model", () => { - const policy = getConnectorModelPolicy(createDeps(), "anthropic"); + const policy = getConnectorModelPolicy(createDeps(), "test-co", "anthropic"); const enabledId = policy.enabledModels[0]?.id; if (!enabledId) { throw new Error("anthropic provider should expose at least one model from pi-ai."); } - const result = validateConnectorModelSelection(createDeps(), "anthropic", enabledId); + const result = validateConnectorModelSelection(createDeps(), "test-co", "anthropic", enabledId); expect(result).toEqual({ ok: true, effectiveModel: enabledId }); }); it("rejects an explicit disabled model", () => { - const policy = getConnectorModelPolicy(createDeps(), "anthropic"); + const policy = getConnectorModelPolicy(createDeps(), "test-co", "anthropic"); const enabledId = policy.enabledModels[0]?.id; if (!enabledId) { throw new Error("anthropic provider should expose at least one model from pi-ai."); } const result = validateConnectorModelSelection( - createDeps({ disabledModels: [{ id: enabledId, label: enabledId }] }), - "anthropic", + createDeps({ disabledModels: [{ id: enabledId, label: enabledId }] }), "test-co", "anthropic", enabledId, ); @@ -75,8 +73,8 @@ describe("native connector defaults", () => { label: "Codex", model: null, }), - getDisabledModels: () => [], - }, "openai-codex"); + getDisabledModels: (_c: string, _k: string) => [], + }, "test-co", "openai-codex"); expect(policy.defaultModel).toBe(configuredDefault); expect(validateConnectorModelSelection({ @@ -85,8 +83,8 @@ describe("native connector defaults", () => { label: "Codex", model: null, }), - getDisabledModels: () => [], - }, "openai-codex", null)).toEqual({ + getDisabledModels: (_c: string, _k: string) => [], + }, "test-co", "openai-codex", null)).toEqual({ ok: true, effectiveModel: configuredDefault, }); diff --git a/tests/unit/ipc-agent-graph-handlers.test.ts b/tests/unit/ipc-agent-graph-handlers.test.ts index 718a1339..0b63fddf 100644 --- a/tests/unit/ipc-agent-graph-handlers.test.ts +++ b/tests/unit/ipc-agent-graph-handlers.test.ts @@ -38,13 +38,12 @@ describe("ipc-core-agent-graph-handlers model enforcement", () => { version: "test", lastError: null, }); - db.setDisabledModels("openai-codex", [{ id: "gpt-5.4", label: "gpt-5.4" }]); - const companyId = db.saveCompany({ name: "Handler Test Co", description: "Tests saveAgent model validation", status: "active", }); + db.setDisabledModels(companyId, "openai-codex", [{ id: "gpt-5.4", label: "gpt-5.4" }]); const workspaceId = db.saveWorkspace({ companyId, name: "Primary",