File size: 4,554 Bytes
4e23b01 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | import {
createKimiDefaultHeaders,
KIMI_CODE_FLOW_CONFIG,
KIMI_CODE_PROVIDER_NAME,
KimiOAuthToolkit,
kimiCodeBaseUrl,
parseKimiCodeCustomHeaders,
resolveKimiCodeOAuthRef,
type KimiHostIdentity,
type ManagedKimiOAuthRef,
} from '@moonshot-ai/kimi-code-oauth';
import type {
ProviderConfig as KosongProviderConfig,
ProviderRequestAuth,
} from '@moonshot-ai/kosong';
import { APIStatusError, UNKNOWN_CAPABILITY } from '@moonshot-ai/kosong';
import { resolveKimiHome } from '@moonshot-ai/agent-core-v2';
import { ErrorCodes, KimiError } from '#/errors';
import type { Logger } from '#/logging/index';
import type { ModelProvider, ResolvedRuntimeProvider } from '#/model-provider';
import { mapOAuthTokenError } from '#/oauth-error';
export interface KimiForCodingProviderOptions extends KimiHostIdentity {
readonly homeDir?: string;
readonly model?: string;
readonly baseUrl?: string;
readonly promptCacheKey?: string;
readonly defaultHeaders?: Record<string, string>;
}
export class KimiForCodingProvider implements ModelProvider {
private readonly model: string;
private readonly baseUrl: string;
private readonly promptCacheKey: string | undefined;
private readonly defaultHeaders: Record<string, string> | undefined;
private readonly toolkit: KimiOAuthToolkit;
private readonly homeDir: string;
private readonly identity: KimiHostIdentity;
private readonly oauthRef: ManagedKimiOAuthRef;
constructor(options: KimiForCodingProviderOptions) {
this.model = options.model ?? 'kimi-for-coding';
this.baseUrl = options.baseUrl ?? kimiCodeBaseUrl();
this.promptCacheKey = options.promptCacheKey;
this.defaultHeaders = options.defaultHeaders;
this.homeDir = resolveKimiHome(options.homeDir);
this.identity = {
productName: options.productName,
version: options.version,
platform: options.platform,
userAgentSuffix: options.userAgentSuffix,
};
this.oauthRef = resolveKimiCodeOAuthRef({
oauthHost: KIMI_CODE_FLOW_CONFIG.oauthHost,
baseUrl: this.baseUrl,
});
this.toolkit = new KimiOAuthToolkit({
homeDir: this.homeDir,
identity: this.identity,
});
}
get defaultModel(): string {
return this.model;
}
resolveProviderConfig(model: string): ResolvedRuntimeProvider {
if (model !== this.model) {
throw new KimiError(
ErrorCodes.CONFIG_INVALID,
`Model "${model}" is not supported by KimiForCodingProvider.`,
);
}
const provider: KosongProviderConfig = {
type: 'kimi',
model: this.model,
baseUrl: this.baseUrl,
generationKwargs: this.promptCacheKey
? { prompt_cache_key: this.promptCacheKey }
: undefined,
defaultHeaders: {
...parseKimiCodeCustomHeaders(),
...createKimiDefaultHeaders({
homeDir: this.homeDir,
...this.identity,
}),
...this.defaultHeaders,
},
};
return {
providerName: 'kimi-for-coding',
provider,
modelCapabilities: UNKNOWN_CAPABILITY,
type: 'kimi',
protocol: undefined,
};
}
resolveAuth(_model: string, _options?: { readonly log?: Logger }) {
return async <T>(request: (auth: ProviderRequestAuth) => Promise<T>): Promise<T> => {
let auth = await this.buildAuth(false);
for (let refreshed = false; ; refreshed = true) {
try {
return await request(auth);
} catch (error) {
const is401 = error instanceof APIStatusError && error.statusCode === 401;
if (!is401) throw error;
if (refreshed) {
throw new KimiError(
ErrorCodes.AUTH_LOGIN_REQUIRED,
'OAuth token was rejected after refresh. Run /login to re-authenticate.',
{ cause: error },
);
}
auth = await this.buildAuth(true);
}
}
};
}
private async buildAuth(force: boolean): Promise<ProviderRequestAuth> {
try {
const apiKey = await this.toolkit.ensureFresh(KIMI_CODE_PROVIDER_NAME, {
force,
oauthRef: this.oauthRef,
});
return { apiKey };
} catch (error) {
// Classify OAuth token failures into the public KimiError protocol so the
// turn surfaces `auth.login_required` / `provider.connection_error`
// instead of collapsing everything to `internal`. Unrecognized errors are
// rethrown raw (see mapOAuthTokenError).
throw mapOAuthTokenError(error, KIMI_CODE_PROVIDER_NAME) ?? error;
}
}
}
|