Skip to content

Commit 1903287

Browse files
wangxiaolong100ellipsis-dev[bot]roomote[bot]
authored andcommitted
fix: Fix the issue of Moonshot's maximum return token count being limited to 1024 (RooCodeInc#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>
1 parent 4d3c958 commit 1903287

File tree

3 files changed

+76
-2
lines changed

3 files changed

+76
-2
lines changed

src/api/providers/__tests__/moonshot.spec.ts

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,4 +294,66 @@ describe("MoonshotHandler", () => {
294294
expect(result.cacheReadTokens).toBeUndefined()
295295
})
296296
})
297+
298+
describe("addMaxTokensIfNeeded", () => {
299+
it("should always add max_tokens regardless of includeMaxTokens option", () => {
300+
// Create a test subclass to access the protected method
301+
class TestMoonshotHandler extends MoonshotHandler {
302+
public testAddMaxTokensIfNeeded(requestOptions: any, modelInfo: any) {
303+
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
304+
}
305+
}
306+
307+
const testHandler = new TestMoonshotHandler(mockOptions)
308+
const requestOptions: any = {}
309+
const modelInfo = {
310+
maxTokens: 32_000,
311+
}
312+
313+
// Test with includeMaxTokens set to false - should still add max tokens
314+
testHandler.testAddMaxTokensIfNeeded(requestOptions, modelInfo)
315+
316+
expect(requestOptions.max_tokens).toBe(32_000)
317+
})
318+
319+
it("should use modelMaxTokens when provided", () => {
320+
class TestMoonshotHandler extends MoonshotHandler {
321+
public testAddMaxTokensIfNeeded(requestOptions: any, modelInfo: any) {
322+
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
323+
}
324+
}
325+
326+
const customMaxTokens = 5000
327+
const testHandler = new TestMoonshotHandler({
328+
...mockOptions,
329+
modelMaxTokens: customMaxTokens,
330+
})
331+
const requestOptions: any = {}
332+
const modelInfo = {
333+
maxTokens: 32_000,
334+
}
335+
336+
testHandler.testAddMaxTokensIfNeeded(requestOptions, modelInfo)
337+
338+
expect(requestOptions.max_tokens).toBe(customMaxTokens)
339+
})
340+
341+
it("should fall back to modelInfo.maxTokens when modelMaxTokens is not provided", () => {
342+
class TestMoonshotHandler extends MoonshotHandler {
343+
public testAddMaxTokensIfNeeded(requestOptions: any, modelInfo: any) {
344+
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
345+
}
346+
}
347+
348+
const testHandler = new TestMoonshotHandler(mockOptions)
349+
const requestOptions: any = {}
350+
const modelInfo = {
351+
maxTokens: 16_000,
352+
}
353+
354+
testHandler.testAddMaxTokensIfNeeded(requestOptions, modelInfo)
355+
356+
expect(requestOptions.max_tokens).toBe(16_000)
357+
})
358+
})
297359
})

src/api/providers/moonshot.ts

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
import { moonshotModels, moonshotDefaultModelId } from "@roo-code/types"
1+
import OpenAI from "openai"
2+
import { moonshotModels, moonshotDefaultModelId, type ModelInfo } from "@roo-code/types"
23

34
import type { ApiHandlerOptions } from "../../shared/api"
45

@@ -36,4 +37,15 @@ export class MoonshotHandler extends OpenAiHandler {
3637
cacheReadTokens: usage?.cached_tokens,
3738
}
3839
}
40+
41+
// Override to always include max_tokens for Moonshot (not max_completion_tokens)
42+
protected override addMaxTokensIfNeeded(
43+
requestOptions:
44+
| OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming
45+
| OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming,
46+
modelInfo: ModelInfo,
47+
): void {
48+
// Moonshot uses max_tokens instead of max_completion_tokens
49+
requestOptions.max_tokens = this.options.modelMaxTokens || modelInfo.maxTokens
50+
}
3951
}

src/api/providers/openai.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -435,7 +435,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
435435
* Note: max_tokens is deprecated in favor of max_completion_tokens as per OpenAI documentation
436436
* O3 family models handle max_tokens separately in handleO3FamilyMessage
437437
*/
438-
private addMaxTokensIfNeeded(
438+
protected addMaxTokensIfNeeded(
439439
requestOptions:
440440
| OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming
441441
| OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming,

0 commit comments

Comments
 (0)