feat(api): add Kimi Code provider with OAuth device flow #945
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,29 @@ | ||
| import { | ||
| SECRET_STATE_KEYS, | ||
| dynamicProviders, | ||
| kimiCodeDefaultModelId, | ||
| providerSettingsSchema, | ||
| providerSettingsSchemaDiscriminated, | ||
| } from "../index.js" | ||
|
|
||
| describe("Kimi Code provider types", () => { | ||
| it("registers Kimi Code as a dynamic provider with a distinct secret", () => { | ||
| expect(dynamicProviders).toContain("kimi-code") | ||
| expect(SECRET_STATE_KEYS).toContain("kimiCodeApiKey") | ||
| expect(SECRET_STATE_KEYS).toContain("moonshotApiKey") | ||
| }) | ||
|
|
||
| it("parses OAuth and API-key settings independently from Moonshot", () => { | ||
| expect( | ||
| providerSettingsSchemaDiscriminated.parse({ | ||
| apiProvider: "kimi-code", | ||
| kimiCodeAuthMethod: "api-key", | ||
| kimiCodeApiKey: "kimi-key", | ||
| apiModelId: kimiCodeDefaultModelId, | ||
| }), | ||
| ).toMatchObject({ kimiCodeApiKey: "kimi-key" }) | ||
| expect(providerSettingsSchema.parse({ apiProvider: "kimi-code", kimiCodeAuthMethod: "oauth" })).toMatchObject({ | ||
| kimiCodeAuthMethod: "oauth", | ||
| }) | ||
| }) | ||
| }) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| import type { ModelInfo } from "../model.js" | ||
|
|
||
| export const KIMI_CODE_BASE_URL = "https://api.kimi.com/coding/v1" | ||
| export const kimiCodeDefaultModelId = "kimi-for-coding" | ||
|
|
||
| export const kimiCodeDefaultModelInfo: ModelInfo = { | ||
| contextWindow: 262_144, | ||
| maxTokens: 32_768, | ||
| supportsImages: false, | ||
| supportsPromptCache: false, | ||
| description: "Kimi Code's coding model for subscription and API-key access.", | ||
| } | ||
|
|
||
| export const kimiCodeModels = { | ||
| [kimiCodeDefaultModelId]: kimiCodeDefaultModelInfo, | ||
| } as const satisfies Record<string, ModelInfo> | ||
|
|
||
| export type KimiCodeModelId = keyof typeof kimiCodeModels |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,198 @@ | ||
| import { buildApiHandler } from "../../index" | ||
| import { KimiCodeHandler } from "../kimi-code" | ||
|
|
||
| const { mockGetAccessToken, mockForceRefreshAccessToken, mockGetModels } = vi.hoisted(() => ({ | ||
| mockGetAccessToken: vi.fn(), | ||
| mockForceRefreshAccessToken: vi.fn(), | ||
| mockGetModels: vi.fn(), | ||
| })) | ||
|
|
||
| vi.mock("../../../integrations/kimi-code/oauth", () => ({ | ||
| kimiCodeOAuthManager: { | ||
| getAccessToken: mockGetAccessToken, | ||
| forceRefreshAccessToken: mockForceRefreshAccessToken, | ||
| }, | ||
| })) | ||
|
|
||
| vi.mock("../fetchers/modelCache", () => ({ getModels: mockGetModels })) | ||
|
|
||
| describe("KimiCodeHandler", () => { | ||
| beforeEach(() => { | ||
| vi.clearAllMocks() | ||
| mockGetAccessToken.mockResolvedValue("oauth-token") | ||
| mockForceRefreshAccessToken.mockResolvedValue("refreshed-token") | ||
| mockGetModels.mockRejectedValue(new Error("offline")) | ||
| }) | ||
|
|
||
| it("is dispatched separately from Moonshot and preserves an unknown selected model", () => { | ||
| const handler = buildApiHandler({ | ||
| apiProvider: "kimi-code", | ||
| kimiCodeAuthMethod: "api-key", | ||
| kimiCodeApiKey: "kimi-key", | ||
| apiModelId: "future-kimi-model", | ||
| }) | ||
| expect(handler).toBeInstanceOf(KimiCodeHandler) | ||
| expect(handler.getModel().id).toBe("future-kimi-model") | ||
| }) | ||
|
|
||
| it("uses kimi-for-coding only when no model is selected", () => { | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "kimi-key" }) | ||
| expect(handler.getModel().id).toBe("kimi-for-coding") | ||
| }) | ||
|
|
||
| it("uses API key when auth method is api-key", async () => { | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "my-api-key" }) | ||
| const gen = handler.createMessage("system", [{ role: "user", content: "test" }]) | ||
| const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValueOnce( | ||
| new Response(JSON.stringify({ choices: [{ message: { content: "response" }, finish_reason: "stop" }] }), { | ||
| status: 200, | ||
| }), | ||
| ) | ||
| try { | ||
| for await (const chunk of gen) { | ||
| // consume | ||
| } | ||
| } catch { | ||
| // expected - mock is incomplete | ||
| } | ||
| expect(mockGetAccessToken).not.toHaveBeenCalled() | ||
| }) | ||
|
|
||
| it("uses OAuth token when auth method is oauth or not specified", async () => { | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" }) | ||
| const gen = handler.createMessage("system", [{ role: "user", content: "test" }]) | ||
| try { | ||
| for await (const chunk of gen) { | ||
| // consume | ||
| } | ||
| } catch { | ||
| // expected - mock will fail | ||
| } | ||
| expect(mockGetAccessToken).toHaveBeenCalled() | ||
| }) | ||
|
|
||
| it("throws error when OAuth is required but no token available", async () => { | ||
| mockGetAccessToken.mockResolvedValueOnce(null) | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" }) | ||
| const gen = handler.createMessage("system", [{ role: "user", content: "test" }]) | ||
| await expect(async () => { | ||
| for await (const chunk of gen) { | ||
| // consume | ||
| } | ||
| }).rejects.toThrow("Not authenticated with Kimi Code") | ||
| }) | ||
|
|
||
| it("throws error when API key auth is missing the key", async () => { | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key" }) | ||
| const gen = handler.createMessage("system", [{ role: "user", content: "test" }]) | ||
| await expect(async () => { | ||
| for await (const chunk of gen) { | ||
| // consume | ||
| } | ||
| }).rejects.toThrow("Kimi Code API key is required") | ||
| }) | ||
|
|
||
| it("retries with forced refresh on 401 when using OAuth", async () => { | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" }) | ||
| const fetchSpy = vi.spyOn(globalThis, "fetch") | ||
| fetchSpy.mockResolvedValueOnce(new Response(null, { status: 401 })) | ||
| fetchSpy.mockResolvedValueOnce( | ||
| new Response(JSON.stringify({ choices: [{ message: { content: "ok" }, finish_reason: "stop" }] }), { | ||
| status: 200, | ||
| }), | ||
| ) | ||
| const gen = handler.createMessage("system", [{ role: "user", content: "test" }]) | ||
| try { | ||
| for await (const chunk of gen) { | ||
| // consume | ||
| } | ||
| } catch { | ||
| // expected - mock is incomplete | ||
| } | ||
| expect(mockForceRefreshAccessToken).toHaveBeenCalledOnce() | ||
| }) | ||
|
|
||
| it("force-refreshes and retries exactly once after a non-streaming OAuth 401", async () => { | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" }) | ||
| const unauthorized = Object.assign(new Error("Unauthorized"), { status: 401 }) | ||
| const createCompletion = vi | ||
| .spyOn((handler as any).client.chat.completions, "create") | ||
| .mockRejectedValueOnce(unauthorized) | ||
| .mockResolvedValueOnce({ choices: [{ message: { content: "retried" } }] }) | ||
|
|
||
| await expect(handler.completePrompt("test")).resolves.toBe("retried") | ||
| expect(mockForceRefreshAccessToken).toHaveBeenCalledOnce() | ||
| expect(createCompletion).toHaveBeenCalledTimes(2) | ||
| }) | ||
|
|
||
| it("does not retry on 401 when using API key auth", async () => { | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" }) | ||
| const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValueOnce(new Response(null, { status: 401 })) | ||
| const gen = handler.createMessage("system", [{ role: "user", content: "test" }]) | ||
| await expect(async () => { | ||
| for await (const chunk of gen) { | ||
| // consume | ||
| } | ||
| }).rejects.toThrow() | ||
| expect(mockForceRefreshAccessToken).not.toHaveBeenCalled() | ||
| }) | ||
|
|
||
|
Comment on lines
+95
to
+139
Contributor
There was a problem hiding this comment. 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win Cover the forced-refresh retry contract for While Please extend the test suite to ensure comprehensive coverage of the error handling and retry decisions as previously requested. 🤖 Proposed tests to addAdd these within the it("retries completePrompt with forced refresh on 401 when using OAuth", async () => {
const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" })
const fetchSpy = vi.spyOn(globalThis, "fetch")
fetchSpy.mockResolvedValueOnce(new Response(null, { status: 401 }))
fetchSpy.mockResolvedValueOnce(
new Response(JSON.stringify({ choices: [{ text: "ok" }] }), { status: 200 }),
)
try {
await handler.completePrompt("test")
} catch {
// expected - mock is incomplete
}
expect(mockForceRefreshAccessToken as any).toHaveBeenCalled()
})
it("does not retry on non-401 errors when using OAuth", async () => {
const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" })
const fetchSpy = vi.spyOn(globalThis, "fetch")
fetchSpy.mockResolvedValueOnce(new Response(null, { status: 500 }))
const gen = handler.createMessage("system", [{ role: "user", content: "test" }])
await expect(async () => {
for await (const chunk of gen) {}
}).rejects.toThrow()
expect(mockForceRefreshAccessToken as any).not.toHaveBeenCalled()
})🤖 Prompt for AI Agents |
||
| it("fetches models during prepareRequest", async () => { | ||
| mockGetModels.mockResolvedValueOnce({ "test-model": { maxTokens: 1000 } }) | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" }) | ||
| const gen = handler.createMessage("system", [{ role: "user", content: "test" }]) | ||
| try { | ||
| for await (const chunk of gen) { | ||
| // consume | ||
| } | ||
| } catch { | ||
| // expected | ||
| } | ||
| expect(mockGetModels).toHaveBeenCalled() | ||
| }) | ||
|
|
||
| it("continues when model discovery fails", async () => { | ||
| mockGetModels.mockRejectedValueOnce(new Error("discovery failed")) | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" }) | ||
| const gen = handler.createMessage("system", [{ role: "user", content: "test" }]) | ||
| try { | ||
| for await (const chunk of gen) { | ||
| // consume | ||
| } | ||
| } catch { | ||
| // expected - different error | ||
| } | ||
| expect(mockGetModels).toHaveBeenCalled() | ||
| }) | ||
|
|
||
| it.each([ | ||
| ["failure", () => Promise.reject(new Error("offline"))], | ||
| ["empty response", () => Promise.resolve({})], | ||
| ])("does not repeatedly block requests after model discovery %s", async (_case, discovery) => { | ||
| mockGetModels.mockImplementationOnce(discovery) | ||
| vi.spyOn(globalThis, "fetch").mockImplementation( | ||
| async () => new Response(JSON.stringify({ choices: [{ message: { content: "ok" } }] }), { status: 200 }), | ||
| ) | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" }) | ||
|
|
||
| await handler.completePrompt("first") | ||
| await handler.completePrompt("second") | ||
|
|
||
| expect(mockGetModels).toHaveBeenCalledOnce() | ||
| }) | ||
|
|
||
| it("uses discovered model info when available", async () => { | ||
| mockGetModels.mockResolvedValueOnce({ "kimi-for-coding": { maxTokens: 8000, contextWindow: 128000 } }) | ||
| const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" }) | ||
| const gen = handler.createMessage("system", [{ role: "user", content: "test" }]) | ||
| try { | ||
| for await (const chunk of gen) { | ||
| // consume | ||
| } | ||
| } catch { | ||
| // expected | ||
| } | ||
| const model = handler.getModel() | ||
| expect(model.info.maxTokens).toBe(8000) | ||
| }) | ||
| }) | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.