refactor(model-pool): scope disabled-models by company; cascade on delete (v42)

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.
This commit is contained in:
ZacharyZhang-NY committed 2026-04-18 22:14:09 -04:00
1 parent 2ace9183b0
commit d75bfd5851
22 files changed
+270 -135

No files matched your search

+5 -3
View File
@@ -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) {
+3 -3
View File
@@ -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); }
+64
View File
@@ -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:<connectorId>`) 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`);
}
},
},
];
}
+2
View File
@@ -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) {
+43 -9
View File
@@ -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 {
+5 -6
View File
@@ -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);
});
+2 -2
View File
@@ -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);
}
+5 -4
View File
@@ -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([]);
}
});
+9 -3
View File
@@ -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,
});
+4 -4
View File
@@ -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<TaskTier, number> = {
* 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) {
+21 -4
View File
@@ -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();
+1 -1
View File
@@ -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;
+1 -1
View File
@@ -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";
+1 -1
View File
@@ -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;
+7 -5
View File
@@ -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);
}
@@ -101,6 +101,7 @@ export function CompanyOnboarding({
totalEnabledModels,
toggleConnectorModel,
setAllConnectorModels,
poolSelections,
} = useConnectorModelSelection(connectors, currentLocale, t, setError);
const pickDirectory = useCallback(async (): Promise<string | null> => {
@@ -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 () => {
@@ -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<void>;
setAllConnectorModels: (connectorId: ConnectorId, enabled: boolean) => Promise<void>;
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<ConnectorPoolSelection[]>(
() =>
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,
};
}
+17
View File
@@ -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({
@@ -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();
@@ -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();
+11 -13
View File
@@ -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,
});
+1 -2
View File
@@ -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",