fix: honor custom provider context windows (#830)

Persist runtime provider metadata for desktop sessions so context estimates stay tied to the provider that launched the session. Infer a unique saved provider window for older sessions without runtime metadata, and keep Xiaomi MiMo custom providers on thinking-only capability mapping instead of max-effort passthrough. Tested: bun test src/server/__tests__/conversations.test.ts -t "unique saved provider context window|pass the active provider id|provider model context windows" Tested: bun test src/server/__tests__/providers.test.ts -t "Xiaomi MiMo custom providers|custom providers declare|model context windows" Tested: bun test src/server/__tests__/sessions.test.ts src/server/__tests__/searchService.sessions.test.ts src/server/__tests__/ws-memory-events.test.ts Tested: cd desktop && bun run test -- src/stores/chatStore.test.ts src/stores/providerStore.test.ts src/components/controls/ModelSelector.test.tsx Not-tested: full bun run verify was not run for this local issue handoff. Confidence: high Scope-risk: moderate
2026-07-17 13:13:35 +08:00 · 2026-06-15 23:21:33 +08:00 · 2026-06-15 23:21:33 +08:00 · 35aa36c319
commit 35aa36c319
parent 572eba46f4
5 changed files with 194 additions and 4 deletions
--- a/src/server/tests/conversations.test.ts
+++ b/src/server/tests/conversations.test.ts
@ -699,6 +699,101 @@ describe('ConversationService', () => {
    }
  })

+  it('should infer a unique saved provider context window for sessions missing runtime metadata', async () => {
+    const previousConfigDir = process.env.CLAUDE_CONFIG_DIR
+    const previousNodeEnv = process.env.NODE_ENV
+    const previousModelContextWindows = process.env.CLAUDE_CODE_MODEL_CONTEXT_WINDOWS
+    const tmpConfigDir = await fs.mkdtemp(path.join(os.tmpdir(), 'claude-transcript-provider-infer-'))
+    const workDir = await fs.mkdtemp(path.join(os.tmpdir(), 'claude-workdir-provider-infer-'))
+    process.env.CLAUDE_CONFIG_DIR = tmpConfigDir
+    process.env.NODE_ENV = 'development'
+    delete process.env.CLAUDE_CODE_MODEL_CONTEXT_WINDOWS
+
+    try {
+      const providerService = new ProviderService()
+      await providerService.addProvider({
+        presetId: 'custom',
+        name: 'Xiaomi MiMo',
+        apiKey: 'provider-key',
+        authStrategy: 'auth_token',
+        baseUrl: 'https://token-plan-sgp.xiaomimimo.com/anthropic',
+        apiFormat: 'anthropic',
+        models: {
+          main: 'mimo-v2.5-pro[1m]',
+          haiku: 'mimo-v2.5-pro[1m]',
+          sonnet: 'mimo-v2.5-pro[1m]',
+          opus: 'mimo-v2.5-pro[1m]',
+        },
+        modelContextWindows: {
+          'mimo-v2.5-pro[1m]': 1_000_000,
+        },
+      })
+      const activeProvider = await providerService.addProvider({
+        presetId: 'custom',
+        name: 'Active DeepSeek',
+        apiKey: 'provider-key',
+        authStrategy: 'auth_token',
+        baseUrl: 'https://api.deepseek.com/anthropic',
+        apiFormat: 'anthropic',
+        models: {
+          main: 'deepseek-v4-pro',
+          haiku: 'deepseek-v4-flash',
+          sonnet: 'deepseek-v4-pro',
+          opus: 'deepseek-v4-pro',
+        },
+        modelContextWindows: {
+          'deepseek-v4-pro': 1_000_000,
+        },
+      })
+      await providerService.activateProvider(activeProvider.id)
+
+      const svc = new SessionService()
+      const { sessionId } = await svc.createSession(workDir)
+      const found = await svc.findSessionFile(sessionId)
+      expect(found).not.toBeNull()
+
+      await fs.appendFile(found!.filePath, JSON.stringify({
+        type: 'assistant',
+        uuid: crypto.randomUUID(),
+        timestamp: '2026-06-15T12:00:00.000Z',
+        cwd: workDir,
+        version: '999.0.0-test',
+        message: {
+          role: 'assistant',
+          model: 'mimo-v2.5-pro',
+          content: [{ type: 'text', text: 'hello' }],
+          usage: {
+            input_tokens: 100,
+            output_tokens: 20,
+          },
+        },
+      }) + '\n')
+
+      const contextEstimate = await svc.getTranscriptContextEstimate(sessionId)
+
+      expect(contextEstimate?.model).toBe('mimo-v2.5-pro')
+      expect(contextEstimate?.rawMaxTokens).toBe(1_000_000)
+    } finally {
+      if (previousConfigDir === undefined) {
+        delete process.env.CLAUDE_CONFIG_DIR
+      } else {
+        process.env.CLAUDE_CONFIG_DIR = previousConfigDir
+      }
+      if (previousNodeEnv === undefined) {
+        delete process.env.NODE_ENV
+      } else {
+        process.env.NODE_ENV = previousNodeEnv
+      }
+      if (previousModelContextWindows === undefined) {
+        delete process.env.CLAUDE_CODE_MODEL_CONTEXT_WINDOWS
+      } else {
+        process.env.CLAUDE_CODE_MODEL_CONTEXT_WINDOWS = previousModelContextWindows
+      }
+      await fs.rm(tmpConfigDir, { recursive: true, force: true })
+      await fs.rm(workDir, { recursive: true, force: true })
+    }
+  })
+
  it('should not report transcript context as full for low-trust media usage spikes', async () => {
    const previousConfigDir = process.env.CLAUDE_CONFIG_DIR
    const previousNodeEnv = process.env.NODE_ENV
@ -2126,6 +2221,10 @@ describe('WebSocket Chat Integration', () => {
          providerId: provider.id,
        },
      })
+      const launchInfo = await sessionService.getSessionLaunchInfo(sessionId)
+      expect(launchInfo).toMatchObject({
+        runtimeProviderId: provider.id,
+      })
    } finally {
      conversationService.startSession = originalStartSession
      conversationService.stopSession(sessionId)
--- a/src/server/tests/providers.test.ts
+++ b/src/server/tests/providers.test.ts
@ -230,6 +230,28 @@ describe('ProviderService', () => {
      )
    })

+    test('Xiaomi MiMo custom providers declare thinking without effort passthrough', async () => {
+      const svc = new ProviderService()
+      const provider = await svc.addProvider(sampleInput({
+        name: 'Xiaomi MiMo Custom',
+        baseUrl: 'https://token-plan-sgp.xiaomimimo.com/anthropic',
+        models: {
+          main: 'mimo-v2.5-pro[1m]',
+          haiku: 'mimo-v2.5-pro[1m]',
+          sonnet: 'mimo-v2.5-pro[1m]',
+          opus: 'mimo-v2.5-pro[1m]',
+        },
+      }))
+
+      await svc.activateProvider(provider.id)
+
+      const settings = await readSettings()
+      const env = settings.env as Record<string, string>
+      expect(env.ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES).toBe('thinking')
+      expect(env.ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES).toBe('thinking')
+      expect(env.ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES).toBe('thinking')
+    })
+
    test('DeepSeek preset follows the global thinking toggle instead of forcing disabled thinking', async () => {
      const svc = new ProviderService()
      const provider = await svc.addProvider(sampleInput({
--- a/src/server/services/conversationService.ts
+++ b/src/server/services/conversationService.ts
@ -398,12 +398,21 @@ export class ConversationService {

    session.startupPending = false

-    if (shouldReplacePlaceholder || !launchInfo) {
+    const shouldPersistRuntimeMetadata =
+      options?.providerId !== undefined ||
+      !!options?.model ||
+      !!options?.effort
+    if (shouldReplacePlaceholder || !launchInfo || shouldPersistRuntimeMetadata) {
      await sessionService.appendSessionMetadata(sessionId, {
        workDir: launchWorkDir,
        customTitle: launchInfo?.customTitle ?? null,
        repository: launchRepository,
        permissionMode: options?.permissionMode || launchInfo?.permissionMode,
+        ...(options?.providerId !== undefined
+          ? { runtimeProviderId: options.providerId }
+          : {}),
+        ...(options?.model ? { runtimeModelId: options.model } : {}),
+        ...(options?.effort ? { effortLevel: options.effort } : {}),
      })
    }

--- a/src/server/services/providerRuntimeEnv.ts
+++ b/src/server/services/providerRuntimeEnv.ts
@ -42,6 +42,7 @@ export const MANAGED_PROVIDER_ENV_KEYS = [
 ] as const

 const CUSTOM_PROVIDER_MODEL_CAPABILITIES = 'thinking,effort,adaptive_thinking,max_effort'
+const XIAOMI_MIMO_MODEL_CAPABILITIES = 'thinking'
 const AUTH_ENV_KEYS = new Set(['ANTHROPIC_API_KEY', 'ANTHROPIC_AUTH_TOKEN'])

 function isRecord(value: unknown): value is Record<string, unknown> {
@ -182,6 +183,25 @@ function getPresetModelContextWindows(presetId: string): Record<string, number>
  return PROVIDER_PRESETS.find((preset) => preset.id === presetId)?.modelContextWindows ?? {}
 }

+function isXiaomiMimoProvider(provider: SavedProvider, models: SavedProvider['models']): boolean {
+  const baseUrl = provider.baseUrl.toLowerCase()
+  const modelIds = Object.values(models).map((model) => model.toLowerCase())
+  return (
+    baseUrl.includes('xiaomimimo.com') ||
+    modelIds.some((model) => /^mimo-v\d/i.test(model))
+  )
+}
+
+function getCustomProviderModelCapabilities(
+  provider: SavedProvider,
+  models: SavedProvider['models'],
+): string {
+  if (isXiaomiMimoProvider(provider, models)) {
+    return XIAOMI_MIMO_MODEL_CAPABILITIES
+  }
+  return CUSTOM_PROVIDER_MODEL_CAPABILITIES
+}
+
 export function buildProviderAuthEnv(
  provider: SavedProvider,
  presetDefaultEnv: Record<string, string>,
@ -243,12 +263,13 @@ export function buildProviderManagedEnv(
  }

  const presetDefaultEnv = getPresetDefaultEnv(provider.presetId)
+  const customProviderCapabilities = getCustomProviderModelCapabilities(provider, models)
  const customProviderCapabilityEnv =
    provider.presetId === 'custom'
      ? {
-          ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES: CUSTOM_PROVIDER_MODEL_CAPABILITIES,
-          ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES: CUSTOM_PROVIDER_MODEL_CAPABILITIES,
-          ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES: CUSTOM_PROVIDER_MODEL_CAPABILITIES,
+          ANTHROPIC_DEFAULT_HAIKU_MODEL_SUPPORTED_CAPABILITIES: customProviderCapabilities,
+          ANTHROPIC_DEFAULT_SONNET_MODEL_SUPPORTED_CAPABILITIES: customProviderCapabilities,
+          ANTHROPIC_DEFAULT_OPUS_MODEL_SUPPORTED_CAPABILITIES: customProviderCapabilities,
        }
      : {}

--- a/src/server/services/sessionService.ts
+++ b/src/server/services/sessionService.ts
@ -1367,6 +1367,7 @@ export class SessionService {
  ): Promise<number | undefined> {
    const launchInfo = await this.getSessionLaunchInfo(sessionId).catch(() => null)
    const providerIds: string[] = []
+    const allowSavedProviderInference = launchInfo?.runtimeProviderId === undefined

    if (typeof launchInfo?.runtimeProviderId === 'string') {
      providerIds.push(launchInfo.runtimeProviderId)
@ -1389,9 +1390,47 @@ export class SessionService {
      }
    }

+    if (allowSavedProviderInference) {
+      return this.getUniqueSavedProviderContextWindow(model)
+    }
+
    return undefined
  }

+  private async getUniqueSavedProviderContextWindow(model: string): Promise<number | undefined> {
+    const { providers } = await this.providerService.listProviders().catch(() => ({ providers: [] }))
+    const matches: number[] = []
+
+    for (const provider of providers) {
+      const env = await this.providerService.getProviderRuntimeEnv(provider.id).catch(() => null)
+      const contextWindow = getModelContextWindowFromEnvValue(
+        model,
+        env?.[MODEL_CONTEXT_WINDOWS_ENV_KEY],
+      )
+      if (contextWindow !== undefined) {
+        matches.push(contextWindow)
+      }
+    }
+
+    if (matches.length === 0) {
+      return undefined
+    }
+
+    const uniqueWindows = new Set(matches)
+    if (uniqueWindows.size !== 1) {
+      return undefined
+    }
+
+    const [contextWindow] = uniqueWindows
+    if (
+      contextWindow > MODEL_CONTEXT_WINDOW_DEFAULT &&
+      is1mContextDisabled()
+    ) {
+      return MODEL_CONTEXT_WINDOW_DEFAULT
+    }
+    return contextWindow
+  }
+
  private async getTranscriptContextWindow(sessionId: string, model: string): Promise<number> {
    const providerContextWindow = await this.getProviderContextWindowForSession(sessionId, model)
    if (providerContextWindow !== undefined) {