99 lines
2.8 KiB
Go
99 lines
2.8 KiB
Go
|
package server
|
||
|
|
|
||
|
|
import (
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestToolCallStreamParser_PlainTextPassThrough(t *testing.T) {
|
||
|
|
p := NewToolCallStreamParser()
|
||
|
|
emit, calls, err := p.Feed("hello world\n")
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("unexpected error: %v", err)
|
||
|
|
}
|
||
|
|
if len(calls) != 0 {
|
||
|
|
t.Fatalf("expected no calls, got %+v", calls)
|
||
|
|
}
|
||
|
|
if emit != "hello world\n" {
|
||
|
|
t.Fatalf("emit = %q, want passthrough", emit)
|
||
|
|
}
|
||
|
|
rest, err := p.Flush()
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("flush error: %v", err)
|
||
|
|
}
|
||
|
|
if rest != "" {
|
||
|
|
t.Fatalf("flush leftover = %q, want empty", rest)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestToolCallStreamParser_ExtractsCompleteCall(t *testing.T) {
|
||
|
|
p := NewToolCallStreamParser()
|
||
|
|
in := "before\n<tool_call>\n{\"name\":\"bash\",\"input\":{\"command\":\"ls\"}}\n</tool_call>\nafter"
|
||
|
|
emit, calls, err := p.Feed(in)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("error: %v", err)
|
||
|
|
}
|
||
|
|
if len(calls) != 1 {
|
||
|
|
t.Fatalf("expected 1 call, got %d", len(calls))
|
||
|
|
}
|
||
|
|
if calls[0].Name != "bash" {
|
||
|
|
t.Fatalf("name = %q", calls[0].Name)
|
||
|
|
}
|
||
|
|
if !strings.Contains(string(calls[0].Input), `"command":"ls"`) {
|
||
|
|
t.Fatalf("input = %s", calls[0].Input)
|
||
|
|
}
|
||
|
|
if !strings.Contains(emit, "before") || !strings.Contains(emit, "after") {
|
||
|
|
t.Fatalf("emit lost surrounding text: %q", emit)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestToolCallStreamParser_HoldsPartialOpenSentinel(t *testing.T) {
|
||
|
|
p := NewToolCallStreamParser()
|
||
|
|
// Feed a chunk ending with a partial "<tool_ca". Parser must not emit it.
|
||
|
|
emit, calls, err := p.Feed("text<tool_ca")
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("error: %v", err)
|
||
|
|
}
|
||
|
|
if len(calls) != 0 {
|
||
|
|
t.Fatalf("calls = %+v", calls)
|
||
|
|
}
|
||
|
|
if emit != "text" {
|
||
|
|
t.Fatalf("emit = %q, want %q", emit, "text")
|
||
|
|
}
|
||
|
|
emit2, calls2, err := p.Feed("ll>{\"name\":\"x\"}</tool_call>")
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("error 2: %v", err)
|
||
|
|
}
|
||
|
|
if emit2 != "" {
|
||
|
|
t.Fatalf("emit2 = %q, want empty (only call extracted)", emit2)
|
||
|
|
}
|
||
|
|
if len(calls2) != 1 || calls2[0].Name != "x" {
|
||
|
|
t.Fatalf("calls2 = %+v", calls2)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestToolCallStreamParser_RejectsInvalidJSON(t *testing.T) {
|
||
|
|
p := NewToolCallStreamParser()
|
||
|
|
_, _, err := p.Feed("<tool_call>not json</tool_call>")
|
||
|
|
if err == nil {
|
||
|
|
t.Fatal("expected parse error for invalid JSON inside sentinels")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractAllToolCalls_MultipleAndCleanText(t *testing.T) {
|
||
|
|
in := "preamble\n<tool_call>{\"name\":\"a\",\"input\":{}}</tool_call>\nmiddle\n<tool_call>{\"tool\":\"b\",\"arguments\":{\"x\":1}}</tool_call>\nend"
|
||
|
|
clean, calls := ExtractAllToolCalls(in)
|
||
|
|
if len(calls) != 2 {
|
||
|
|
t.Fatalf("calls = %d", len(calls))
|
||
|
|
}
|
||
|
|
if calls[0].Name != "a" || calls[1].Name != "b" {
|
||
|
|
t.Fatalf("names = %q, %q", calls[0].Name, calls[1].Name)
|
||
|
|
}
|
||
|
|
if !strings.Contains(clean, "preamble") || !strings.Contains(clean, "middle") || !strings.Contains(clean, "end") {
|
||
|
|
t.Fatalf("clean text wrong: %q", clean)
|
||
|
|
}
|
||
|
|
if strings.Contains(clean, "<tool_call>") {
|
||
|
|
t.Fatalf("clean text still contains sentinels: %q", clean)
|
||
|
|
}
|
||
|
|
}
|