417 lines
13 KiB
Go
417 lines
13 KiB
Go
package openai
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestToolCallTruncated(t *testing.T) {
|
|
valid := `{"tool_calls": [{"id": "a", "type": "function", "function": {"name": "write", "arguments": "{\"filePath\": \"x\"}"}}]}`
|
|
truncated := valid[:40]
|
|
prose := "the model mentions tool_calls but writes no JSON at all"
|
|
fenced := "```json\n" + truncated
|
|
|
|
cases := []struct {
|
|
in string
|
|
want bool
|
|
}{
|
|
{valid, false},
|
|
{truncated, true},
|
|
{prose, false},
|
|
{fenced, true},
|
|
{"", false},
|
|
{`{"other": true}`, false},
|
|
// prose before the payload: mirrored from TryParseToolCalls
|
|
{"我先说明一下：\n\n" + valid, false},
|
|
{"我先说明一下：\n\n" + truncated, true},
|
|
{"我先说明一下：\n\n" + fenced, true},
|
|
}
|
|
for _, c := range cases {
|
|
if got := ToolCallTruncated(c.in); got != c.want {
|
|
t.Errorf("ToolCallTruncated(%q) = %v, want %v", c.in, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestEscapeControlChars(t *testing.T) {
|
|
cases := []struct{ in, want string }{
|
|
{`{"a": "line1` + "\n" + `line2"}`, `{"a": "line1\nline2"}`},
|
|
{`{"a": "tab` + "\t" + `x"}`, `{"a": "tab\tx"}`},
|
|
// Escaped sequences inside strings stay as-is.
|
|
{`{"a": "already\nescaped"}`, `{"a": "already\nescaped"}`},
|
|
// Pretty-printed structure newlines outside strings survive.
|
|
{"{\n \"a\": \"x\"\n}", "{\n \"a\": \"x\"\n}"},
|
|
// Real newline inside a second string while first had escapes.
|
|
{`{"a": "x\ny", "b": "p` + "\n" + `q"}`, `{"a": "x\ny", "b": "p\nq"}`},
|
|
// Other control characters get \u-escaped.
|
|
{`{"a": "` + string(rune(0x01)) + `"}`, `{"a": "\u0001"}`},
|
|
}
|
|
for _, c := range cases {
|
|
if got := EscapeControlChars(c.in); got != c.want {
|
|
t.Errorf("EscapeControlChars(%q) = %q, want %q", c.in, got, c.want)
|
|
}
|
|
}
|
|
// Round-trip: repaired text must be valid JSON with the same value.
|
|
raw := "{\"a\": \"line1\nline2\"}"
|
|
fixed := EscapeControlChars(raw)
|
|
var v struct {
|
|
A string `json:"a"`
|
|
}
|
|
if err := json.Unmarshal([]byte(fixed), &v); err != nil {
|
|
t.Fatalf("repaired JSON invalid: %v", err)
|
|
}
|
|
if v.A != "line1\nline2" {
|
|
t.Fatalf("repaired value = %q", v.A)
|
|
}
|
|
}
|
|
|
|
func TestEscapeControlCharsDoublesInvalidEscapes(t *testing.T) {
|
|
cases := []struct{ in, want string }{
|
|
// \q is not a JSON escape — the backslash must be doubled
|
|
{`{"a": "x\qy"}`, `{"a": "x\\qy"}`},
|
|
// valid pairs survive, invalid ones heal
|
|
{`{"a": "ok\nstill\q"}`, `{"a": "ok\nstill\\q"}`},
|
|
// dangling escape at the end is carried, not doubled
|
|
{`{"a": "dangling\`, `{"a": "dangling\`},
|
|
// no backslash at all: untouched
|
|
{`{"a": "plain"}`, `{"a": "plain"}`},
|
|
}
|
|
for _, c := range cases {
|
|
if got := EscapeControlChars(c.in); got != c.want {
|
|
t.Errorf("EscapeControlChars(%q) = %q, want %q", c.in, got, c.want)
|
|
}
|
|
}
|
|
if !json.Valid([]byte(EscapeControlChars(`{"a": "x\qy"}`))) {
|
|
t.Fatal("healed text is not valid JSON")
|
|
}
|
|
}
|
|
|
|
// TestInnerEscapeMatchesParseHeal pins the contract that keeps streamed tool
|
|
// arguments identical to the final repaired parse.
|
|
func TestInnerEscapeMatchesParseHeal(t *testing.T) {
|
|
inner := `{"content": "a\qb` + "\n" + `c\qd"}`
|
|
streamed := innerEscape(inner)
|
|
healed := EscapeControlChars(inner)
|
|
if streamed != healed {
|
|
t.Fatalf("stream/parse heal mismatch:\n stream=%q\n parse =%q", streamed, healed)
|
|
}
|
|
if !json.Valid([]byte(healed)) {
|
|
t.Fatalf("healed inner invalid: %q", healed)
|
|
}
|
|
}
|
|
|
|
func TestTrimOverlap(t *testing.T) {
|
|
cases := []struct{ acc, head, want string }{
|
|
{"...fillStyle='#F4", "F4E2B8';draw()", "E2B8';draw()"},
|
|
{"...fillStyle='#F4", "E2B8';draw()", "E2B8';draw()"},
|
|
{"abcdef", "xyz", "xyz"},
|
|
{"abcabcabc", "abcabcTAIL", "TAIL"},
|
|
{"", "anything", "anything"},
|
|
}
|
|
for _, c := range cases {
|
|
if got := TrimOverlap(c.acc, c.head); got != c.want {
|
|
t.Errorf("TrimOverlap(%q, %q) = %q, want %q", c.acc, c.head, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestEscapeControlsAfterCarriesPrefixState(t *testing.T) {
|
|
prefix := `{"tool_calls": [{"function": {"arguments": "{\"content\": \"abc`
|
|
head := "def\n" + `ghi\""}}]}`
|
|
got := EscapeControlsAfter(prefix, head)
|
|
want := `def\nghi\""}}]}`
|
|
if got != want {
|
|
t.Fatalf("EscapeControlsAfter = %q, want %q", got, want)
|
|
}
|
|
// Without prefix context the same head starts outside a string, so the
|
|
// newline is structure whitespace and must survive.
|
|
if plain := EscapeControlChars(head); plain != head {
|
|
t.Fatalf("EscapeControlChars = %q, want unchanged", plain)
|
|
}
|
|
full := prefix + got
|
|
if !json.Valid([]byte(full)) {
|
|
t.Fatalf("repaired document invalid: %s", full)
|
|
}
|
|
}
|
|
|
|
func TestContinueInstructionMentionsTail(t *testing.T) {
|
|
partial := strings.Repeat("x", 200) + "TAIL80CHARS"
|
|
got := ContinueInstruction(partial, 3, 12)
|
|
if !strings.Contains(got, "3/12") {
|
|
t.Errorf("instruction missing round progress: %s", got)
|
|
}
|
|
if !strings.Contains(got, "TAIL80CHARS") {
|
|
t.Errorf("instruction missing tail: %s", got)
|
|
}
|
|
if strings.Contains(got, strings.Repeat("x", 100)) {
|
|
t.Errorf("instruction embedded too much of the partial")
|
|
}
|
|
if !strings.Contains(got, `\n`) {
|
|
t.Errorf("instruction missing escape-sequence rule")
|
|
}
|
|
}
|
|
|
|
func TestBalanceJSON(t *testing.T) {
|
|
cases := []struct{ name, in, want string }{
|
|
{"unterminated content", `{"filePath": "x.html", "content": "<html></html>`, `{"filePath": "x.html", "content": "<html></html>"}`},
|
|
{"trailing comma", `{"a": "b",`, `{"a": "b"}`},
|
|
{"trailing colon", `{"a":`, `{"a": ""}`},
|
|
{"dangling escape", `{"a": "b\`, `{"a": "b"}`},
|
|
{"nested array", `{"a": [1, 2`, `{"a": [1, 2]}`},
|
|
{"already valid", `{"a": "b"}`, `{"a": "b"}`},
|
|
{"comma inside string survives", `{"a": "b,c`, `{"a": "b,c"` + `}`},
|
|
}
|
|
for _, c := range cases {
|
|
if got := balanceJSON(c.in); got != c.want {
|
|
t.Errorf("%s: balanceJSON(%q) = %q, want %q", c.name, c.in, got, c.want)
|
|
}
|
|
}
|
|
// Unfixable damage passes through untouched.
|
|
broken := `{"a": 12.}`
|
|
if got := balanceJSON(broken); got != broken {
|
|
t.Errorf("unfixable: got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestTryParseToolCallsHealsBrokenInnerJSON(t *testing.T) {
|
|
text := `{"tool_calls": [{"id": "a", "type": "function", "function": {"name": "write", "arguments": "{\"content\": \"<html></html>"}}]}`
|
|
_, calls, ok := TryParseToolCalls(text)
|
|
if !ok {
|
|
t.Fatal("expected healed tool call")
|
|
}
|
|
var inner struct {
|
|
Content string `json:"content"`
|
|
}
|
|
if err := json.Unmarshal([]byte(calls[0].Function.Arguments), &inner); err != nil {
|
|
t.Fatalf("healed arguments invalid: %v", err)
|
|
}
|
|
if inner.Content != "<html></html>" {
|
|
t.Fatalf("content = %q", inner.Content)
|
|
}
|
|
}
|
|
|
|
func TestTryParseToolCallsRejectsUnfixableInnerJSON(t *testing.T) {
|
|
text := `{"tool_calls": [{"id": "a", "type": "function", "function": {"name": "write", "arguments": "{\"content\": 12.}"}}]}`
|
|
_, calls, ok := TryParseToolCalls(text)
|
|
if ok {
|
|
t.Fatalf("expected rejection, got calls=%v", calls)
|
|
}
|
|
}
|
|
|
|
func TestRunContinuedStitches(t *testing.T) {
|
|
turns := []string{
|
|
`{"tool_calls": [{"id": "a", "function": {"name": "write", "arguments": "{\"content\": \"part1`,
|
|
// Continuation repeats one char, emits a raw newline mid-string, then
|
|
// closes the inner JSON, the arguments string and the outer document.
|
|
"1\n" + `part2\"}"}}]}`,
|
|
}
|
|
prompts := 0
|
|
var got strings.Builder
|
|
err := RunContinued(context.Background(),
|
|
"first-prompt",
|
|
func(partial string, _ int) string {
|
|
if !strings.Contains(partial, "part1") {
|
|
t.Errorf("cont prompt built from unexpected partial: %s", partial)
|
|
}
|
|
return "continue-prompt"
|
|
},
|
|
3,
|
|
func(_ context.Context, prompt string, emit func(string) error) error {
|
|
prompts++
|
|
idx := prompts - 1
|
|
if idx >= len(turns) {
|
|
t.Fatalf("unexpected extra send #%d", prompts)
|
|
}
|
|
if idx > 0 && prompt != "continue-prompt" {
|
|
t.Errorf("round %d prompt = %q", idx, prompt)
|
|
}
|
|
return emit(turns[idx])
|
|
},
|
|
func(text string) error { got.WriteString(text); return nil },
|
|
func(string, ...any) {},
|
|
func(string, ...any) {},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("RunContinued: %v", err)
|
|
}
|
|
if prompts != 2 {
|
|
t.Fatalf("prompts = %d, want 2", prompts)
|
|
}
|
|
want := `{"tool_calls": [{"id": "a", "function": {"name": "write", "arguments": "{\"content\": \"part1\npart2\"}"}}]}`
|
|
if got.String() != want {
|
|
t.Fatalf("stitched = %q\nwant %q", got.String(), want)
|
|
}
|
|
if ToolCallTruncated(got.String()) {
|
|
t.Fatal("stitched reply still truncated")
|
|
}
|
|
_, calls, ok := TryParseToolCalls(got.String())
|
|
if !ok {
|
|
t.Fatalf("TryParseToolCalls failed on stitched reply")
|
|
}
|
|
var inner struct {
|
|
Content string `json:"content"`
|
|
}
|
|
if err := json.Unmarshal([]byte(calls[0].Function.Arguments), &inner); err != nil {
|
|
t.Fatalf("arguments JSON invalid: %v", err)
|
|
}
|
|
if inner.Content != "part1\npart2" {
|
|
t.Fatalf("content = %q", inner.Content)
|
|
}
|
|
}
|
|
|
|
func TestRunContinuedPassthroughWithoutCont(t *testing.T) {
|
|
prompts := 0
|
|
var got strings.Builder
|
|
err := RunContinued(context.Background(), "p", nil, 6,
|
|
func(_ context.Context, prompt string, emit func(string) error) error {
|
|
prompts++
|
|
return emit(`{"tool_calls": [{"trunc`)
|
|
},
|
|
func(text string) error { got.WriteString(text); return nil },
|
|
func(string, ...any) {}, func(string, ...any) {},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("RunContinued: %v", err)
|
|
}
|
|
if prompts != 1 {
|
|
t.Fatalf("prompts = %d, want 1", prompts)
|
|
}
|
|
if !ToolCallTruncated(got.String()) {
|
|
t.Fatal("passthrough should keep the truncated reply")
|
|
}
|
|
}
|
|
|
|
func TestRunContinuedContinuationFailureKeepsPartial(t *testing.T) {
|
|
prompts := 0
|
|
warns := 0
|
|
var got strings.Builder
|
|
err := RunContinued(context.Background(), "first",
|
|
func(string, int) string { return "next" }, 3,
|
|
func(_ context.Context, _ string, emit func(string) error) error {
|
|
prompts++
|
|
if prompts == 1 {
|
|
return emit(`{"tool_calls": [{"trunc`)
|
|
}
|
|
return errors.New("boom")
|
|
},
|
|
func(text string) error { got.WriteString(text); return nil },
|
|
func(string, ...any) {},
|
|
func(string, ...any) { warns++ },
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("RunContinued returned error despite partial text: %v", err)
|
|
}
|
|
if warns != 5 {
|
|
t.Fatalf("warns = %d, want 5 (3 retry notices + turn failure + exhaustion)", warns)
|
|
}
|
|
if got.String() != `{"tool_calls": [{"trunc` {
|
|
t.Fatalf("got = %q", got.String())
|
|
}
|
|
}
|
|
|
|
func TestRunContinuedStopsWhenClosed(t *testing.T) {
|
|
prompts := 0
|
|
err := RunContinued(context.Background(), "first",
|
|
func(string, int) string { return "next" }, 6,
|
|
func(_ context.Context, _ string, emit func(string) error) error {
|
|
prompts++
|
|
return emit(`{"tool_calls": [{"id": "a", "type": "function", "function": {"name": "f", "arguments": "{}"}}]}`)
|
|
},
|
|
func(string) error { return nil },
|
|
func(string, ...any) {}, func(string, ...any) {},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("RunContinued: %v", err)
|
|
}
|
|
if prompts != 1 {
|
|
t.Fatalf("prompts = %d, want 1 (valid JSON must not trigger continuation)", prompts)
|
|
}
|
|
}
|
|
|
|
func TestRunContinuedEmitErrorPropagates(t *testing.T) {
|
|
sentinel := errors.New("emit failed")
|
|
err := RunContinued(context.Background(), "first", nil, 0,
|
|
func(_ context.Context, _ string, emit func(string) error) error {
|
|
return emit("text")
|
|
},
|
|
func(string) error { return sentinel },
|
|
func(string, ...any) {}, func(string, ...any) {},
|
|
)
|
|
if !errors.Is(err, sentinel) {
|
|
t.Fatalf("err = %v, want sentinel", err)
|
|
}
|
|
}
|
|
|
|
func TestTryParseToolCallsShortCloseSequence(t *testing.T) {
|
|
inner := `{\"filePath\": \"a.html\", \"content\": \"<html></html>\"}`
|
|
raw := `{"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "write", "arguments": "` + inner + `}}]}`
|
|
|
|
remaining, calls, ok := TryParseToolCalls(raw)
|
|
if !ok {
|
|
t.Fatal("TryParseToolCalls ok = false, want true (short close sequence must heal)")
|
|
}
|
|
if remaining != "" {
|
|
t.Fatalf("remaining = %q, want empty", remaining)
|
|
}
|
|
if len(calls) != 1 {
|
|
t.Fatalf("calls = %d, want 1", len(calls))
|
|
}
|
|
args := calls[0].Function.Arguments
|
|
if !json.Valid([]byte(args)) {
|
|
t.Fatalf("args invalid: %s", args)
|
|
}
|
|
var obj struct {
|
|
FilePath string `json:"filePath"`
|
|
Content string `json:"content"`
|
|
}
|
|
if err := json.Unmarshal([]byte(args), &obj); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if obj.Content != "<html></html>" {
|
|
t.Fatalf("content = %q, want html", obj.Content)
|
|
}
|
|
}
|
|
|
|
func TestTryParseToolCallsLeakedQuoteCutoff(t *testing.T) {
|
|
raw := `{"tool_calls": [{"id": "call_001", "type": "function", "function": {"name": "write", "arguments": "{\"filePath\": \"tihu.html\", \"content\": \"<html><body>canvas scene leak "and then cut`
|
|
_, calls, ok := TryParseToolCalls(raw)
|
|
if !ok {
|
|
t.Fatal("TryParseToolCalls ok = false, want force-closed tool call")
|
|
}
|
|
args := calls[0].Function.Arguments
|
|
if !json.Valid([]byte(args)) {
|
|
t.Fatalf("args invalid after force close: %s", args)
|
|
}
|
|
var inner struct {
|
|
FilePath string `json:"filePath"`
|
|
Content string `json:"content"`
|
|
}
|
|
if err := json.Unmarshal([]byte(args), &inner); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if inner.FilePath != "tihu.html" {
|
|
t.Fatalf("filePath = %q", inner.FilePath)
|
|
}
|
|
if !strings.Contains(inner.Content, "canvas scene leak") || !strings.Contains(inner.Content, `"and then cut`) {
|
|
t.Fatalf("content lost pre-cut text: %q", inner.Content)
|
|
}
|
|
}
|
|
|
|
func TestTryParseToolCallsPlainCutoffReturnsPartial(t *testing.T) {
|
|
raw := `{"tool_calls": [{"id": "c", "type": "function", "function": {"name": "write", "arguments": "{\"filePath\": \"a.html\", \"content\": \"<html><body><canvas id=\"scene\">stopped mid`
|
|
_, calls, ok := TryParseToolCalls(raw)
|
|
if !ok {
|
|
t.Fatal("TryParseToolCalls ok = false, want partial tool call")
|
|
}
|
|
args := calls[0].Function.Arguments
|
|
if !json.Valid([]byte(args)) {
|
|
t.Fatalf("args invalid: %s", args)
|
|
}
|
|
if !strings.Contains(args, "stopped mid") {
|
|
t.Fatalf("args lost tail: %s", args)
|
|
}
|
|
}
|