fix(antigravity): preserve finish_reason tool_calls across streaming chunks #874
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
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,128 @@ | ||
| package chat_completions | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
|
|
||
| "github.com/tidwall/gjson" | ||
| ) | ||
|
|
||
| func TestFinishReasonToolCallsNotOverwritten(t *testing.T) { | ||
| ctx := context.Background() | ||
| var param any | ||
|
|
||
| // Chunk 1: Contains functionCall - should set SawToolCall = true | ||
| chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"functionCall":{"name":"list_files","args":{"path":"."}}}]}}]}}`) | ||
| result1 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, ¶m) | ||
|
|
||
| // Verify chunk1 has no finish_reason (null) | ||
| if len(result1) != 1 { | ||
| t.Fatalf("Expected 1 result from chunk1, got %d", len(result1)) | ||
| } | ||
| fr1 := gjson.Get(result1[0], "choices.0.finish_reason") | ||
| if fr1.Exists() && fr1.String() != "" && fr1.Type.String() != "Null" { | ||
| t.Errorf("Expected finish_reason to be null in chunk1, got: %v", fr1.String()) | ||
| } | ||
|
|
||
| // Chunk 2: Contains finishReason STOP + usage (final chunk, no functionCall) | ||
| // This simulates what the upstream sends AFTER the tool call chunk | ||
| chunk2 := []byte(`{"response":{"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":20,"totalTokenCount":30}}}`) | ||
| result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, ¶m) | ||
|
|
||
| // Verify chunk2 has finish_reason: "tool_calls" (not "stop") | ||
| if len(result2) != 1 { | ||
| t.Fatalf("Expected 1 result from chunk2, got %d", len(result2)) | ||
| } | ||
| fr2 := gjson.Get(result2[0], "choices.0.finish_reason").String() | ||
| if fr2 != "tool_calls" { | ||
| t.Errorf("Expected finish_reason 'tool_calls', got: %s", fr2) | ||
| } | ||
|
|
||
| // Verify native_finish_reason is lowercase upstream value | ||
| nfr2 := gjson.Get(result2[0], "choices.0.native_finish_reason").String() | ||
| if nfr2 != "stop" { | ||
| t.Errorf("Expected native_finish_reason 'stop', got: %s", nfr2) | ||
| } | ||
| } | ||
|
|
||
| func TestFinishReasonStopForNormalText(t *testing.T) { | ||
| ctx := context.Background() | ||
| var param any | ||
|
|
||
| // Chunk 1: Text content only | ||
| chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"Hello world"}]}}]}}`) | ||
| ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, ¶m) | ||
|
|
||
| // Chunk 2: Final chunk with STOP | ||
| chunk2 := []byte(`{"response":{"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}}`) | ||
| result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, ¶m) | ||
|
|
||
| // Verify finish_reason is "stop" (no tool calls were made) | ||
| fr := gjson.Get(result2[0], "choices.0.finish_reason").String() | ||
| if fr != "stop" { | ||
| t.Errorf("Expected finish_reason 'stop', got: %s", fr) | ||
| } | ||
| } | ||
|
|
||
| func TestFinishReasonMaxTokens(t *testing.T) { | ||
| ctx := context.Background() | ||
| var param any | ||
|
|
||
| // Chunk 1: Text content | ||
| chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"Hello"}]}}]}}`) | ||
| ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, ¶m) | ||
|
|
||
| // Chunk 2: Final chunk with MAX_TOKENS | ||
| chunk2 := []byte(`{"response":{"candidates":[{"finishReason":"MAX_TOKENS"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":100,"totalTokenCount":110}}}`) | ||
| result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, ¶m) | ||
|
|
||
| // Verify finish_reason is "max_tokens" | ||
| fr := gjson.Get(result2[0], "choices.0.finish_reason").String() | ||
| if fr != "max_tokens" { | ||
| t.Errorf("Expected finish_reason 'max_tokens', got: %s", fr) | ||
| } | ||
| } | ||
|
|
||
| func TestToolCallTakesPriorityOverMaxTokens(t *testing.T) { | ||
| ctx := context.Background() | ||
| var param any | ||
|
|
||
| // Chunk 1: Contains functionCall | ||
| chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"functionCall":{"name":"test","args":{}}}]}}]}}`) | ||
| ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, ¶m) | ||
|
|
||
| // Chunk 2: Final chunk with MAX_TOKENS (but we had a tool call, so tool_calls should win) | ||
| chunk2 := []byte(`{"response":{"candidates":[{"finishReason":"MAX_TOKENS"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":100,"totalTokenCount":110}}}`) | ||
| result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, ¶m) | ||
|
|
||
| // Verify finish_reason is "tool_calls" (takes priority over max_tokens) | ||
| fr := gjson.Get(result2[0], "choices.0.finish_reason").String() | ||
| if fr != "tool_calls" { | ||
| t.Errorf("Expected finish_reason 'tool_calls', got: %s", fr) | ||
| } | ||
| } | ||
|
Comment on lines
+10
to
+103
Contributor
There was a problem hiding this comment. These tests are great and cover all the necessary scenarios. To improve maintainability and reduce code duplication, consider refactoring them into a single table-driven test. Each entry in the table could represent a test case with a name, a sequence of input chunks, and the expected final |
||
|
|
||
| func TestNoFinishReasonOnIntermediateChunks(t *testing.T) { | ||
| ctx := context.Background() | ||
| var param any | ||
|
|
||
| // Chunk 1: Text content (no finish reason, no usage) | ||
| chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"Hello"}]}}]}}`) | ||
| result1 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, ¶m) | ||
|
|
||
| // Verify no finish_reason on intermediate chunk | ||
| fr1 := gjson.Get(result1[0], "choices.0.finish_reason") | ||
| if fr1.Exists() && fr1.String() != "" && fr1.Type.String() != "Null" { | ||
| t.Errorf("Expected no finish_reason on intermediate chunk, got: %v", fr1) | ||
| } | ||
|
|
||
| // Chunk 2: More text (no finish reason, no usage) | ||
| chunk2 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":" world"}]}}]}}`) | ||
| result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, ¶m) | ||
|
|
||
| // Verify no finish_reason on intermediate chunk | ||
| fr2 := gjson.Get(result2[0], "choices.0.finish_reason") | ||
| if fr2.Exists() && fr2.String() != "" && fr2.Type.String() != "Null" { | ||
| t.Errorf("Expected no finish_reason on intermediate chunk, got: %v", fr2) | ||
| } | ||
| } | ||
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The repeated type assertion
(*param).(*convertCliResponseToOpenAIChatParams)is used in several places (here, line 141, etc.), which can make the code harder to read. You've already extracted this into aparamsvariable on line 194. It would be cleaner to perform the type assertion once at the beginning of the function, right after the*paramnil check, and reuse the resulting variable throughout the function. This would improve readability and reduce redundancy.