From d69ed4b331c3fb86250fb55916b9d376fcaece7c Mon Sep 17 00:00:00 2001 From: _Kerman Date: Thu, 9 Jul 2026 14:34:24 +0800 Subject: [PATCH] fix(agent-core-v2): apply KIMI_CODE_CUSTOM_HEADERS and host identity headers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port v1's provider-manager outbound header logic to agent-core-v2 so `KIMI_CODE_CUSTOM_HEADERS` and host identity headers are applied to outbound LLM requests, closing the migration gap from #1186: - env `KIMI_CODE_CUSTOM_HEADERS` is the lowest-precedence header layer; - host identity headers (User-Agent + X-Msh-*) are sent for Kimi providers, only the User-Agent for every other provider — a Kimi provider routed through the Anthropic protocol still gets the full set, matching v1; - provider `customHeaders` always win on conflict. Host headers are seeded by the CLI via `createKimiDefaultHeaders` and a new `IHostRequestHeaders` App-scope token (defaulting to empty), so the model resolver can layer them without the host threading them through every call site. --- .../kimi-code/src/cli/v2/create-v2-harness.ts | 11 +++- .../src/app/model/hostRequestHeaders.ts | 39 ++++++++++++ .../src/app/model/modelResolverService.ts | 35 ++++++++++- packages/agent-core-v2/src/index.ts | 1 + packages/agent-core-v2/test/harness/agent.ts | 4 +- .../test/model/modelResolver.test.ts | 63 ++++++++++++++++++- 6 files changed, 149 insertions(+), 4 deletions(-) create mode 100644 packages/agent-core-v2/src/app/model/hostRequestHeaders.ts diff --git a/apps/kimi-code/src/cli/v2/create-v2-harness.ts b/apps/kimi-code/src/cli/v2/create-v2-harness.ts index b525931cf..cac3a6b23 100644 --- a/apps/kimi-code/src/cli/v2/create-v2-harness.ts +++ b/apps/kimi-code/src/cli/v2/create-v2-harness.ts @@ -15,10 +15,12 @@ import { ISessionLifecycleService, bootstrap, ensureMainAgent, + hostRequestHeadersSeed, logSeed, resolveLoggingConfig, type Scope, } from '@moonshot-ai/agent-core-v2'; +import { createKimiDefaultHeaders } from '@moonshot-ai/kimi-code-oauth'; import { KimiAuthFacade, resolveConfigPath, @@ -49,7 +51,14 @@ export async function createV2Harness(options: KimiHarnessOptions): Promise>; +} + +export const IHostRequestHeaders = createDecorator('hostRequestHeaders'); + +export class HostRequestHeaders implements IHostRequestHeaders { + constructor(readonly headers: Readonly> = {}) {} +} + +/** Seed the host-provided outbound identity headers into an App scope. */ +export function hostRequestHeadersSeed(headers: Readonly>): ScopeSeed { + return [[IHostRequestHeaders as ServiceIdentifier, new HostRequestHeaders(headers)]]; +} + +registerScopedService( + LifecycleScope.App, + IHostRequestHeaders, + HostRequestHeaders, + InstantiationType.Delayed, + 'model', +); diff --git a/packages/agent-core-v2/src/app/model/modelResolverService.ts b/packages/agent-core-v2/src/app/model/modelResolverService.ts index ae01404e9..9b11e91d0 100644 --- a/packages/agent-core-v2/src/app/model/modelResolverService.ts +++ b/packages/agent-core-v2/src/app/model/modelResolverService.ts @@ -15,6 +15,7 @@ * the Model itself; no Platform is required. */ +import { parseKimiCodeCustomHeaders } from '@moonshot-ai/kimi-code-oauth'; import { InstantiationType } from '#/_base/di/extensions'; import { Disposable } from '#/_base/di/lifecycle'; import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; @@ -31,6 +32,7 @@ import { IProviderService } from '#/app/provider/provider'; import { IProtocolAdapterRegistry, type Protocol, type ProtocolProviderOptions } from '#/app/protocol/protocol'; import { type ProtocolAdapterRegistry } from '#/app/protocol/protocolAdapterRegistry'; +import { IHostRequestHeaders } from './hostRequestHeaders'; import type { ModelConfig } from './model'; import { IModelService } from './model'; import { @@ -67,6 +69,7 @@ export class ModelResolverService extends Disposable implements IModelResolver { @IOAuthService private readonly oauth: IOAuthService, @IProtocolAdapterRegistry private readonly protocolRegistry: IProtocolAdapterRegistry, + @IHostRequestHeaders private readonly hostRequestHeaders: IHostRequestHeaders, ) { super(); } @@ -132,7 +135,11 @@ export class ModelResolverService extends Disposable implements IModelResolver { aliases: model.aliases ?? [], protocol, baseUrl: resolvedBaseUrl, - headers: providerConfig?.customHeaders ?? {}, + headers: resolveOutboundHeaders( + providerConfig?.type, + providerConfig?.customHeaders, + this.hostRequestHeaders.headers, + ), capabilities, maxContextSize: model.maxContextSize, maxOutputSize: model.maxOutputSize, @@ -294,6 +301,32 @@ export class ModelResolverService extends Disposable implements IModelResolver { } } +/** + * Resolve the outbound `defaultHeaders` for a Model, layering lowest to highest + * precedence (matches v1's `provider-manager`): + * + * 1. `KIMI_CODE_CUSTOM_HEADERS` env (re-read on every resolve so env changes + * take effect without restarting the session); + * 2. host identity headers — the full set (`User-Agent` + `X-Msh-*`) for a + * Kimi provider, only the `User-Agent` for every other provider so device + * identity never leaks to third-party endpoints (a Kimi provider routed + * through the Anthropic protocol still gets the full set, matching v1); + * 3. provider `customHeaders` (always win on conflict). + */ +export function resolveOutboundHeaders( + providerType: string | undefined, + customHeaders: Readonly> | undefined, + hostHeaders: Readonly>, +): Readonly> { + const hostLayer = providerType === 'kimi' ? hostHeaders : userAgentOnly(hostHeaders); + return { ...parseKimiCodeCustomHeaders(), ...hostLayer, ...(customHeaders ?? {}) }; +} + +function userAgentOnly(headers: Readonly>): Record { + const userAgent = headers['User-Agent']; + return userAgent === undefined ? {} : { 'User-Agent': userAgent }; +} + function resolveModelCapabilities( declaredCapabilities: readonly string[] | undefined, protocol: Protocol, diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index 457c68b5e..1db8dcc1d 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -91,6 +91,7 @@ export * from '#/app/protocol/protocolAdapterRegistry'; import '#/app/model/configSection'; import '#/app/model/envOverlay'; export * from '#/app/model/completionBudget'; +export * from '#/app/model/hostRequestHeaders'; export * from '#/app/model/model'; export type { AuthProvider, diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index 566176d97..19fdaf123 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -121,6 +121,7 @@ import type { PersistedRecord } from '#/wire/wireService'; import { WireService } from '#/wire/wireServiceImpl'; import { IModelService } from '#/app/model/model'; import { type Model } from '#/app/model/modelInstance'; +import { IHostRequestHeaders } from '#/app/model/hostRequestHeaders'; import { IModelResolver } from '#/app/model/modelResolver'; import { ModelResolverService } from '#/app/model/modelResolverService'; import { IPlatformService } from '#/app/platform/platform'; @@ -827,8 +828,9 @@ class ConfigBackedModelResolver extends ModelResolverService { @IModelService models: IModelService, @IOAuthService oauth: IOAuthService, @IProtocolAdapterRegistry protocolRegistry: IProtocolAdapterRegistry, + @IHostRequestHeaders hostRequestHeaders: IHostRequestHeaders, ) { - super(config, providers, platforms, models, oauth, protocolRegistry); + super(config, providers, platforms, models, oauth, protocolRegistry, hostRequestHeaders); } override resolve(id: string): Model { diff --git a/packages/agent-core-v2/test/model/modelResolver.test.ts b/packages/agent-core-v2/test/model/modelResolver.test.ts index 3d30322a3..2b00278e5 100644 --- a/packages/agent-core-v2/test/model/modelResolver.test.ts +++ b/packages/agent-core-v2/test/model/modelResolver.test.ts @@ -21,8 +21,9 @@ import { IOAuthService } from '#/app/auth/auth'; import { IConfigService } from '#/app/config/config'; import { APIStatusError } from '#/app/llmProtocol/errors'; import { type ModelConfig, IModelService } from '#/app/model/model'; +import { HostRequestHeaders, IHostRequestHeaders } from '#/app/model/hostRequestHeaders'; import { IModelResolver } from '#/app/model/modelResolver'; -import { ModelResolverService } from '#/app/model/modelResolverService'; +import { ModelResolverService, resolveOutboundHeaders } from '#/app/model/modelResolverService'; import { type PlatformConfig, IPlatformService } from '#/app/platform/platform'; import { type ProviderConfig, IProviderService } from '#/app/provider/provider'; import { type ChatProvider } from '#/app/llmProtocol/provider'; @@ -89,6 +90,7 @@ describe('ModelResolverService', () => { createChatProvider(input: ProtocolAdapterConfig): ChatProvider; }); reg.define(IModelResolver, ModelResolverService); + reg.defineInstance(IHostRequestHeaders, new HostRequestHeaders()); }, }); }); @@ -442,6 +444,65 @@ describe('ModelResolverService', () => { }); expect(createdProtocolConfigs[0]).not.toHaveProperty('customHeaders'); }); + + it('merges KIMI_CODE_CUSTOM_HEADERS env headers below provider customHeaders', async () => { + process.env['KIMI_CODE_CUSTOM_HEADERS'] = 'X-Env: env-val\nX-Shared: from-env'; + try { + providers['p'] = { + type: 'kimi', + baseUrl: 'https://example.test/v1', + apiKey: 'sk', + customHeaders: { 'X-Shared': 'from-provider', 'X-Provider': 'p' }, + }; + models['m'] = { provider: 'p', model: 'wire-name', maxContextSize: 1000 }; + + const model = ix.get(IModelResolver).resolve('m'); + for await (const _event of model.request({ systemPrompt: '', tools: [], messages: [] })) { + void _event; + } + + expect(createdProtocolConfigs).toHaveLength(1); + expect(createdProtocolConfigs[0]).toMatchObject({ + defaultHeaders: { + 'X-Env': 'env-val', + // provider customHeaders override the env header on conflict + 'X-Shared': 'from-provider', + 'X-Provider': 'p', + }, + }); + } finally { + delete process.env['KIMI_CODE_CUSTOM_HEADERS']; + } + }); + + it('resolveOutboundHeaders: kimi provider gets full host headers, others get only User-Agent', () => { + const saved = process.env['KIMI_CODE_CUSTOM_HEADERS']; + delete process.env['KIMI_CODE_CUSTOM_HEADERS']; + try { + const host = { 'User-Agent': 'kimi-code-cli/1.0', 'X-Msh-Device-Id': 'dev' }; + + // kimi provider → full identity (even when routed through anthropic) + expect(resolveOutboundHeaders('kimi', undefined, host)).toEqual({ + 'User-Agent': 'kimi-code-cli/1.0', + 'X-Msh-Device-Id': 'dev', + }); + // non-kimi providers → User-Agent only + expect(resolveOutboundHeaders('openai', undefined, host)).toEqual({ + 'User-Agent': 'kimi-code-cli/1.0', + }); + expect(resolveOutboundHeaders('anthropic', undefined, host)).toEqual({ + 'User-Agent': 'kimi-code-cli/1.0', + }); + // provider customHeaders win on conflict + expect(resolveOutboundHeaders('kimi', { 'User-Agent': 'custom' }, host)).toEqual({ + 'User-Agent': 'custom', + 'X-Msh-Device-Id': 'dev', + }); + } finally { + if (saved === undefined) delete process.env['KIMI_CODE_CUSTOM_HEADERS']; + else process.env['KIMI_CODE_CUSTOM_HEADERS'] = saved; + } + }); }); describe('provider options', () => {