From 68159b7a4eeb52688dfb530d9eb58fbfbc009d64 Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Thu, 3 Sep 2026 15:22:47 +0800 Subject: [PATCH] fix(agent-core-v2): keep running sessions' models alive across concurrent provider refreshes --- .changeset/olive-camels-refresh.md | 5 ++ .../src/tui/controllers/auth-flow.ts | 61 +++++++++++-- .../test/tui/utils/refresh-providers.test.ts | 13 ++- .../src/app/kosongConfig/discoveryService.ts | 89 ++++++++++--------- .../src/kosong/model/catalogService.ts | 4 +- .../test/app/kosongConfig/discovery.test.ts | 86 ++++++++---------- .../test/kosong/model/catalog.test.ts | 6 +- packages/oauth/src/refreshProviderModels.ts | 51 +++++++++-- 8 files changed, 201 insertions(+), 114 deletions(-) create mode 100644 .changeset/olive-camels-refresh.md diff --git a/.changeset/olive-camels-refresh.md b/.changeset/olive-camels-refresh.md new file mode 100644 index 00000000000..5b2aa206616 --- /dev/null +++ b/.changeset/olive-camels-refresh.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix models in use by a running session failing with "not configured in config.toml" when another Kimi Code process refreshes the provider model list at the same time. diff --git a/apps/kimi-code/src/tui/controllers/auth-flow.ts b/apps/kimi-code/src/tui/controllers/auth-flow.ts index 0940e445a7c..b50167e7ee4 100644 --- a/apps/kimi-code/src/tui/controllers/auth-flow.ts +++ b/apps/kimi-code/src/tui/controllers/auth-flow.ts @@ -217,10 +217,11 @@ export class AuthFlowController { * `replaceSections`), the orchestrator's two-phase contract (removeProvider * then setConfig) is absorbed the same way the v2 engine's own refresh path * does it: the removal is staged in memory only, and the following - * setConfig persists the complete records in a single write — so a process - * exit mid-refresh can never leave config.toml in a "provider removed, not - * yet restored" state. The v1 harness keeps the legacy host (two - * whole-document writes, each atomic on its own). + * setConfig merges the sparse patch onto a fresh read and persists + * everything in a single write — so a process exit mid-refresh can never + * leave config.toml in a "provider removed, not yet restored" state, and + * concurrent writes from other processes survive. The v1 harness keeps the + * legacy host (two whole-document writes, each atomic on its own). */ private buildRefreshHost(): RefreshProviderHost { const { host } = this; @@ -239,6 +240,7 @@ export class AuthFlowController { }; } let staged: KimiConfig | undefined; + const pendingRemovals = new Set(); const requireStaged = (): KimiConfig => { if (staged === undefined) { throw new Error('refresh host: getConfig must be called before writes'); @@ -251,17 +253,60 @@ export class AuthFlowController { return staged; }, removeProvider: (id) => { + pendingRemovals.add(id); staged = removeProviderFromConfig(requireStaged(), id); return Promise.resolve(staged); }, setConfig: async (patch) => { - // The orchestrator always passes complete records (built from a full - // clone), so the Partial-shaped patch is a full KimiConfig overlay. - staged = { ...requireStaged(), ...patch } as KimiConfig; + // The patch is sparse — only the refreshed providers' entries and + // their owned aliases — so merge it onto a fresh read for unrelated + // keys to survive, and fold the staged removals into the same write. + const fresh = await host.harness.getConfig({ reload: true }); + const sections: Record = {}; + if (pendingRemovals.size > 0 || patch.providers !== undefined) { + const providers: Record = { ...fresh.providers }; + for (const id of pendingRemovals) delete providers[id]; + if (patch.providers !== undefined) Object.assign(providers, patch.providers); + sections['providers'] = providers; + } + if (pendingRemovals.size > 0 || patch.models !== undefined) { + const models: Record = { ...fresh.models }; + for (const [key, record] of Object.entries(models)) { + const owner = + typeof record === 'object' && record !== null + ? (record as { provider?: string }).provider + : undefined; + if (owner !== undefined && pendingRemovals.has(owner)) delete models[key]; + } + if (patch.models !== undefined) Object.assign(models, patch.models); + sections['models'] = models; + } // Object.entries keeps keys whose value is `undefined`, so a cleared // section (e.g. a dangling defaultModel) is expressed as a removal in // the atomic write; sections absent from the patch stay untouched. - await host.harness.replaceConfigSections(Object.fromEntries(Object.entries(patch))); + for (const [key, value] of Object.entries(patch)) { + if (key === 'providers' || key === 'models') continue; + sections[key] = value; + } + if (pendingRemovals.size > 0 && !('defaultModel' in patch)) { + const defaultModel = fresh.defaultModel; + const owner = + defaultModel === undefined + ? undefined + : (fresh.models?.[defaultModel] as { provider?: string } | undefined)?.provider; + if (owner !== undefined && pendingRemovals.has(owner)) { + sections['defaultModel'] = undefined; + } + } + if (pendingRemovals.size > 0 && !('defaultProvider' in patch)) { + const defaultProvider = fresh.defaultProvider; + if (defaultProvider !== undefined && pendingRemovals.has(defaultProvider)) { + sections['defaultProvider'] = undefined; + } + } + pendingRemovals.clear(); + await host.harness.replaceConfigSections(sections); + staged = { ...fresh, ...sections } as KimiConfig; return staged; }, resolveOAuthToken, diff --git a/apps/kimi-code/test/tui/utils/refresh-providers.test.ts b/apps/kimi-code/test/tui/utils/refresh-providers.test.ts index 18d43fec7a7..af3a45b02e9 100644 --- a/apps/kimi-code/test/tui/utils/refresh-providers.test.ts +++ b/apps/kimi-code/test/tui/utils/refresh-providers.test.ts @@ -40,7 +40,18 @@ function makeRefreshHost(initial: KimiConfig): { return structuredClone(persisted); }); const setConfig = vi.fn(async (patch: Partial) => { - persisted = { ...persisted, ...patch }; + const next = { ...persisted }; + if (patch.providers !== undefined) { + next.providers = { ...persisted.providers, ...patch.providers }; + } + if (patch.models !== undefined) { + next.models = { ...persisted.models, ...patch.models }; + } + for (const [key, value] of Object.entries(patch)) { + if (key === 'providers' || key === 'models') continue; + (next as Record)[key] = value; + } + persisted = next; return structuredClone(persisted); }); return { diff --git a/packages/agent-core-v2/src/app/kosongConfig/discoveryService.ts b/packages/agent-core-v2/src/app/kosongConfig/discoveryService.ts index e04413b9bba..322f6417922 100644 --- a/packages/agent-core-v2/src/app/kosongConfig/discoveryService.ts +++ b/packages/agent-core-v2/src/app/kosongConfig/discoveryService.ts @@ -25,6 +25,7 @@ import { getProviderDefinition } from '#/kosong/provider/providerDefinition'; import { DEFAULT_MODEL_SECTION, + DEFAULT_PROVIDER_SECTION, MODELS_SECTION, PROVIDERS_SECTION, THINKING_SECTION, @@ -142,10 +143,11 @@ export class ProviderDiscoveryService implements IProviderDiscoveryService { } private buildRefreshHost(exclusion: StaticExclusion, userAgent: string): RefreshProviderHost { + const pendingRemovals = new Set(); return { getConfig: async () => this.readUserConfigShape(exclusion), - removeProvider: (providerId) => this.shapeWithoutProvider(providerId), - setConfig: (patch) => this.applyRefreshPatch(patch, exclusion), + removeProvider: (providerId) => this.queueProviderRemoval(pendingRemovals, providerId), + setConfig: (patch) => this.applyRefreshPatch(patch, pendingRemovals), resolveOAuthToken: (providerName, oauthRef) => this.resolveOAuthToken(providerName, oauthRef), userAgent, }; @@ -173,7 +175,11 @@ export class ProviderDiscoveryService implements IProviderDiscoveryService { }; } - private shapeWithoutProvider(providerId: string): Promise { + private queueProviderRemoval( + pendingRemovals: Set, + providerId: string, + ): Promise { + pendingRemovals.add(providerId); const current = this.readUserConfigShape(); const providers = current.providers as Record; const restProviders = Object.fromEntries( @@ -192,57 +198,54 @@ export class ProviderDiscoveryService implements IProviderDiscoveryService { private async applyRefreshPatch( patch: ManagedKimiConfigShape, - exclusion: StaticExclusion, + pendingRemovals: Set, ): Promise { - const userProviders = + await this.config.reload(); + const removals = [...pendingRemovals]; + pendingRemovals.clear(); + const removed = new Set(removals); + const providers = this.config.inspect>(PROVIDERS_SECTION).userValue ?? {}; - const userModels = + const models = this.config.inspect>(MODELS_SECTION).userValue ?? {}; const sections: Record = {}; - if (patch.providers !== undefined) { - sections[PROVIDERS_SECTION] = { - ...exclusion.providers, - ...patch.providers, - }; + if (removals.length > 0 || patch.providers !== undefined) { + const nextProviders: Record = Object.fromEntries( + Object.entries(providers).filter(([id]) => !removed.has(id)), + ); + if (patch.providers !== undefined) Object.assign(nextProviders, patch.providers); + sections[PROVIDERS_SECTION] = nextProviders; } - if (patch.models !== undefined) { - sections[MODELS_SECTION] = { - ...exclusion.models, - ...(patch.models as Record), - }; + if (removals.length > 0 || patch.models !== undefined) { + const nextModels: Record = Object.fromEntries( + Object.entries(models).filter( + ([, record]) => record.provider === undefined || !removed.has(record.provider), + ), + ); + if (patch.models !== undefined) Object.assign(nextModels, patch.models); + sections[MODELS_SECTION] = nextModels; } - const restoreDefault = exclusion.defaultModel !== undefined; if ('defaultModel' in patch) { - sections[DEFAULT_MODEL_SECTION] = restoreDefault - ? exclusion.defaultModel - : patch.defaultModel; + sections[DEFAULT_MODEL_SECTION] = patch.defaultModel; + } else if (removals.length > 0) { + const defaultModel = this.config.inspect(DEFAULT_MODEL_SECTION).userValue; + if (defaultModel !== undefined && removed.has(models[defaultModel]?.provider ?? '')) { + sections[DEFAULT_MODEL_SECTION] = undefined; + } } if ('thinking' in patch) { - sections[THINKING_SECTION] = restoreDefault ? exclusion.thinking : patch.thinking; + sections[THINKING_SECTION] = patch.thinking; + } + if ('defaultProvider' in patch) { + sections[DEFAULT_PROVIDER_SECTION] = patch['defaultProvider']; + } else if (removals.length > 0) { + const defaultProvider = this.config.inspect(DEFAULT_PROVIDER_SECTION).userValue; + if (defaultProvider !== undefined && removed.has(defaultProvider)) { + sections[DEFAULT_PROVIDER_SECTION] = undefined; + } } await this.config.replaceSections(sections); - return { - providers: - patch.providers !== undefined - ? ({ ...exclusion.providers, ...patch.providers } as ManagedKimiConfigShape['providers']) - : (userProviders as ManagedKimiConfigShape['providers']), - models: - patch.models !== undefined - ? ({ ...exclusion.models, ...patch.models } as ManagedKimiConfigShape['models']) - : (userModels as ManagedKimiConfigShape['models']), - defaultModel: - 'defaultModel' in patch - ? restoreDefault - ? exclusion.defaultModel - : patch.defaultModel - : this.config.inspect(DEFAULT_MODEL_SECTION).userValue, - thinking: - 'thinking' in patch - ? restoreDefault - ? exclusion.thinking - : patch.thinking - : this.config.inspect(THINKING_SECTION).userValue, - }; + return this.readUserConfigShape(); } private async resolveOAuthToken( diff --git a/packages/agent-core-v2/src/kosong/model/catalogService.ts b/packages/agent-core-v2/src/kosong/model/catalogService.ts index 903ebceea3d..9949ab7cc69 100644 --- a/packages/agent-core-v2/src/kosong/model/catalogService.ts +++ b/packages/agent-core-v2/src/kosong/model/catalogService.ts @@ -96,7 +96,9 @@ export class ModelCatalog extends Disposable implements IModelCatalog { } notifyConfigChanged(): void { - this.cache.clear(); + for (const id of this.cache.keys()) { + if (this.models.get(id) !== undefined) this.cache.delete(id); + } } get(id: string): Model { diff --git a/packages/agent-core-v2/test/app/kosongConfig/discovery.test.ts b/packages/agent-core-v2/test/app/kosongConfig/discovery.test.ts index d33a36added..d6297e3db06 100644 --- a/packages/agent-core-v2/test/app/kosongConfig/discovery.test.ts +++ b/packages/agent-core-v2/test/app/kosongConfig/discovery.test.ts @@ -639,69 +639,53 @@ describe('refreshProviderModels defaultModel self-heal', () => { } }); - it('keeps a default model the user selected while the catalog fetch was in flight', async () => { - const twoModels = { - 'kimi-code/kimi-k2': { - provider: KIMI_CODE_PROVIDER_NAME, - model: 'kimi-k2', - maxContextSize: 131072, - capabilities: ['thinking', 'tool_use'], - displayName: 'Kimi K2', - }, - 'kimi-code/kimi-k3': { - provider: KIMI_CODE_PROVIDER_NAME, - model: 'kimi-k3', - maxContextSize: 131072, - capabilities: ['thinking', 'tool_use'], - displayName: 'Kimi K3', - }, - }; - const { host, config, discovery, events } = await createHost( + it('keeps user writes that landed while the catalog fetch was in flight', async () => { + const { host, config, discovery, models } = await createHost( { providers: managedProviders, - models: twoModels, + models: managedModels, }, stubOAuthService(stubTokenProvider(['access-token'])), ); try { vi.stubGlobal( 'fetch', - vi.fn( - async () => { - await config.set('defaultModel', 'kimi-code/kimi-k3'); - return new Response( - JSON.stringify({ - data: [ - { - id: 'kimi-k2', - context_length: 131072, - supports_reasoning: true, - display_name: 'Kimi K2', - }, - { - id: 'kimi-k3', - context_length: 131072, - supports_reasoning: true, - display_name: 'Kimi K3', - }, - ], - }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, - ); - }, - ), + vi.fn(async () => { + await config.set('defaultModel', 'kimi-code/kimi-k3'); + await config.set('models', { + 'other/keep': { provider: 'other', model: 'keep', maxContextSize: 1000 }, + }); + return new Response( + JSON.stringify({ + data: [ + { + id: 'kimi-k2', + context_length: 131072, + supports_reasoning: true, + display_name: 'Kimi K2', + }, + { + id: 'kimi-k3', + context_length: 131072, + supports_reasoning: true, + display_name: 'Kimi K3', + }, + ], + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }), ); - const replaceSections = vi.spyOn(config, 'replaceSections'); const result = await discovery.refreshProviderModels({ scope: 'all' }); - expect(result).toEqual({ - changed: [], - unchanged: [KIMI_CODE_PROVIDER_NAME], - failed: [], - }); - expect(replaceSections).not.toHaveBeenCalled(); - expect(events.published).toEqual([]); + expect(result.failed).toEqual([]); + expect(result.changed).toEqual([ + { provider_id: KIMI_CODE_PROVIDER_NAME, provider_name: 'Kimi Code', added: 1, removed: 0 }, + ]); expect(config.get('defaultModel')).toBe('kimi-code/kimi-k3'); + const modelRecords = models.list(); + expect(modelRecords['kimi-code/kimi-k3']).toBeDefined(); + expect(modelRecords['other/keep']).toBeDefined(); } finally { host.dispose(); } diff --git a/packages/agent-core-v2/test/kosong/model/catalog.test.ts b/packages/agent-core-v2/test/kosong/model/catalog.test.ts index 98a51a8797e..96b29a47a5d 100644 --- a/packages/agent-core-v2/test/kosong/model/catalog.test.ts +++ b/packages/agent-core-v2/test/kosong/model/catalog.test.ts @@ -517,7 +517,7 @@ describe('ModelCatalog caching and config-event invalidation', () => { } }); - it('drops the cache when a watched config section changes', async () => { + it('drops the cache on edits but keeps entries whose model was removed', async () => { const { host, catalog, models, providers } = createHost(kimiSections); try { const before = catalog.get('k1'); @@ -528,6 +528,10 @@ describe('ModelCatalog caching and config-event invalidation', () => { await providers.set('kimi', { type: 'kimi', apiKey: 'sk-2', baseUrl: 'https://other.example.test/v1' }); expect(catalog.get('k1').baseUrl).toBe('https://other.example.test/v1'); + + const updated = catalog.get('k1'); + await models.delete('k1'); + expect(catalog.get('k1')).toBe(updated); } finally { host.dispose(); } diff --git a/packages/oauth/src/refreshProviderModels.ts b/packages/oauth/src/refreshProviderModels.ts index 2bd78037ef4..9bcb3371d5c 100644 --- a/packages/oauth/src/refreshProviderModels.ts +++ b/packages/oauth/src/refreshProviderModels.ts @@ -33,7 +33,18 @@ import { isRecord } from './utils'; */ export interface RefreshProviderHost { getConfig(): Promise; + /** + * Remove the provider entry and every alias it owns + * (`record.provider === providerId`) from the persisted config. The host may + * defer the deletion to the following `setConfig` write so the pair lands as + * one atomic update. + */ removeProvider(providerId: string): Promise; + /** + * Persist a sparse patch: only the provider/model entries carried here are + * written (entry-level replace); every other key on disk is left untouched, + * so a concurrent writer's changes survive the refresh. + */ setConfig(patch: ManagedKimiConfigShape): Promise; resolveOAuthToken(providerName: string, oauthRef?: ManagedKimiOAuthRef): Promise; /** @@ -207,6 +218,19 @@ function computeChanges(oldIds: Set, newIds: Set): { added: numb return { added, removed }; } +function pickDefined( + record: Readonly> | undefined, + keys: Iterable, +): Record { + const out: Record = {}; + if (record === undefined) return out; + for (const key of keys) { + const value = record[key]; + if (value !== undefined) out[key] = value; + } + return out; +} + interface ProviderModelSnapshot { readonly alias: string; readonly model: ManagedKimiModelAlias; @@ -379,7 +403,10 @@ function pickDefaultModel( * 3. Custom registries (models.dev-style, keyed by `provider.source`). * * Each branch diffs old vs new and only writes when something actually changed - * (`removeProvider` then `setConfig`). Failures are collected per-provider and + * (`removeProvider` then `setConfig`). Writes are scoped to the refreshed + * providers: `setConfig` receives only their entries, so keys belonging to + * other providers or written concurrently by another process survive. + * Failures are collected per-provider and * never abort the whole refresh. Pass `providerId` to scope the refresh to a * single provider; pass `scope: 'oauth'` to refresh only the managed provider. */ @@ -449,8 +476,8 @@ export async function refreshProviderModels( ); await host.removeProvider(KIMI_CODE_PROVIDER_NAME); config = await host.setConfig({ - providers: next.providers, - models: next.models, + providers: pickDefined(next.providers, [KIMI_CODE_PROVIDER_NAME]), + models: pickDefined(next.models, providerAliasKeys(next, KIMI_CODE_PROVIDER_NAME)), defaultModel: next.defaultModel, thinking: next.thinking, }); @@ -529,8 +556,8 @@ export async function refreshProviderModels( ); await host.removeProvider(providerId); config = await host.setConfig({ - providers: next.providers, - models: next.models, + providers: pickDefined(next.providers, [providerId]), + models: pickDefined(next.models, providerAliasKeys(next, providerId)), defaultModel: next.defaultModel, thinking: next.thinking, }); @@ -603,8 +630,8 @@ export async function refreshProviderModels( ); await host.removeProvider(providerId); config = await host.setConfig({ - providers: next.providers, - models: next.models, + providers: pickDefined(next.providers, [providerId]), + models: pickDefined(next.models, providerAliasKeys(next, providerId)), defaultModel: next.defaultModel, thinking: next.thinking, // The v1 `removeProvider` RPC clears `defaultProvider` when it points @@ -682,6 +709,7 @@ export async function refreshProviderModels( readonly removed: number; }> = []; const providersToRemoveBeforeSet = new Set(); + const batchProviderIds = new Set(); let hasUnreportedConfigChange = false; const remoteEntries = Object.values(entries); const remoteEntriesByProviderId = new Map( @@ -713,6 +741,7 @@ export async function refreshProviderModels( const existed = config.providers[providerId] !== undefined; applyCustomRegistryProvider(next, entry, source); const refreshedAliasKeys = providerRefreshAliasKeys(config, next, providerId, `${providerId}/`); + batchProviderIds.add(providerId); if (existed) { restoreProviderAliases(next, preserveUserProviderAliases(config, providerId, refreshedAliasKeys)); } @@ -749,9 +778,13 @@ export async function refreshProviderModels( for (const providerId of providersToRemoveBeforeSet) { await host.removeProvider(providerId); } + const batchAliasKeys = new Set(); + for (const providerId of batchProviderIds) { + for (const key of providerAliasKeys(next, providerId)) batchAliasKeys.add(key); + } config = await host.setConfig({ - providers: next.providers, - models: next.models, + providers: pickDefined(next.providers, batchProviderIds), + models: pickDefined(next.models, batchAliasKeys), defaultModel: next.defaultModel, thinking: next.thinking, });