fix: Fix the issue of Moonshot's maximum return token count being limited to 1024 (#7673)

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
Co-authored-by: roomote[bot] <219738659+roomote[bot]@users.noreply.github.com>
This commit is contained in:
WangXiaolong 2025-09-05 11:19:58 +08:00 committed by GitHub
parent de12ec169b
commit bca6e86d2b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 76 additions and 2 deletions

View file

@ -294,4 +294,66 @@ describe("MoonshotHandler", () => {
expect(result.cacheReadTokens).toBeUndefined()
})
})
describe("addMaxTokensIfNeeded", () => {
it("should always add max_tokens regardless of includeMaxTokens option", () => {
// Create a test subclass to access the protected method
class TestMoonshotHandler extends MoonshotHandler {
public testAddMaxTokensIfNeeded(requestOptions: any, modelInfo: any) {
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
}
}
const testHandler = new TestMoonshotHandler(mockOptions)
const requestOptions: any = {}
const modelInfo = {
maxTokens: 32_000,
}
// Test with includeMaxTokens set to false - should still add max tokens
testHandler.testAddMaxTokensIfNeeded(requestOptions, modelInfo)
expect(requestOptions.max_tokens).toBe(32_000)
})
it("should use modelMaxTokens when provided", () => {
class TestMoonshotHandler extends MoonshotHandler {
public testAddMaxTokensIfNeeded(requestOptions: any, modelInfo: any) {
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
}
}
const customMaxTokens = 5000
const testHandler = new TestMoonshotHandler({
...mockOptions,
modelMaxTokens: customMaxTokens,
})
const requestOptions: any = {}
const modelInfo = {
maxTokens: 32_000,
}
testHandler.testAddMaxTokensIfNeeded(requestOptions, modelInfo)
expect(requestOptions.max_tokens).toBe(customMaxTokens)
})
it("should fall back to modelInfo.maxTokens when modelMaxTokens is not provided", () => {
class TestMoonshotHandler extends MoonshotHandler {
public testAddMaxTokensIfNeeded(requestOptions: any, modelInfo: any) {
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
}
}
const testHandler = new TestMoonshotHandler(mockOptions)
const requestOptions: any = {}
const modelInfo = {
maxTokens: 16_000,
}
testHandler.testAddMaxTokensIfNeeded(requestOptions, modelInfo)
expect(requestOptions.max_tokens).toBe(16_000)
})
})
})

View file

@ -1,4 +1,5 @@
import { moonshotModels, moonshotDefaultModelId } from "@roo-code/types"
import OpenAI from "openai"
import { moonshotModels, moonshotDefaultModelId, type ModelInfo } from "@roo-code/types"
import type { ApiHandlerOptions } from "../../shared/api"
@ -36,4 +37,15 @@ export class MoonshotHandler extends OpenAiHandler {
cacheReadTokens: usage?.cached_tokens,
}
}
// Override to always include max_tokens for Moonshot (not max_completion_tokens)
protected override addMaxTokensIfNeeded(
requestOptions:
| OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming
| OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming,
modelInfo: ModelInfo,
): void {
// Moonshot uses max_tokens instead of max_completion_tokens
requestOptions.max_tokens = this.options.modelMaxTokens || modelInfo.maxTokens
}
}

View file

@ -435,7 +435,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
* Note: max_tokens is deprecated in favor of max_completion_tokens as per OpenAI documentation
* O3 family models handle max_tokens separately in handleO3FamilyMessage
*/
private addMaxTokensIfNeeded(
protected addMaxTokensIfNeeded(
requestOptions:
| OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming
| OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming,