119 lines
3.7 KiB
Go
119 lines
3.7 KiB
Go
package agent
|
|
|
|
import (
|
|
"log/slog"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestBuildPiArgsNoToolAllowlist(t *testing.T) {
|
|
// Extension tools registered via Pi's registerTool() must not be
|
|
// filtered out by a hardcoded --tools allowlist. Omitting --tools
|
|
// lets Pi use its full tool registry. See #2379.
|
|
args := buildPiArgs("test prompt", "/tmp/session.jsonl", ExecOptions{}, slog.Default())
|
|
for i, arg := range args {
|
|
if arg == "--tools" {
|
|
t.Errorf("buildPiArgs emits --tools %q; should not restrict tool registry (see #2379)", args[i+1])
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestBuildPiArgsBasicFlags(t *testing.T) {
|
|
args := buildPiArgs("hello world", "/tmp/s.jsonl", ExecOptions{
|
|
Model: "anthropic/claude-sonnet-4-20250514",
|
|
SystemPrompt: "be helpful",
|
|
}, slog.Default())
|
|
|
|
joined := strings.Join(args, " ")
|
|
for _, want := range []string{"-p", "--mode json", "--session /tmp/s.jsonl", "--provider anthropic", "--model claude-sonnet-4-20250514", "--append-system-prompt"} {
|
|
if !strings.Contains(joined, want) {
|
|
t.Errorf("expected %q in args, got: %v", want, args)
|
|
}
|
|
}
|
|
|
|
// Prompt must be the last positional argument.
|
|
if args[len(args)-1] != "hello world" {
|
|
t.Errorf("prompt should be last arg, got %q", args[len(args)-1])
|
|
}
|
|
}
|
|
|
|
func TestBuildPiArgsCustomArgsAppended(t *testing.T) {
|
|
// Users can still restrict tools via custom_args if desired.
|
|
args := buildPiArgs("prompt", "/tmp/s.jsonl", ExecOptions{
|
|
CustomArgs: []string{"--tools", "read,bash"},
|
|
}, slog.Default())
|
|
|
|
found := false
|
|
for i, arg := range args {
|
|
if arg == "--tools" && i+1 < len(args) && args[i+1] == "read,bash" {
|
|
found = true
|
|
}
|
|
}
|
|
if !found {
|
|
t.Errorf("custom --tools should pass through via custom_args, got: %v", args)
|
|
}
|
|
}
|
|
|
|
func TestStripPiToolCallMarkup(t *testing.T) {
|
|
tests := map[string]string{
|
|
`before call:bash{command:<|"|>cd repo/path && ls -F<|"|>}<tool_call|> after`: "before after",
|
|
`before call:read{path:<|"|>repo/path/roles/example/verify.yml<|"|>} after`: "before after",
|
|
`before response:bash{command:<|"|>multica issue comment list issue-id --all --output json<|"|>} after`: "before after",
|
|
`before call:bash{command:<|"|>printf '{"key":"value"}'<|"|>} after`: "before after",
|
|
`before <|turn>model after`: "before after",
|
|
}
|
|
for in, want := range tests {
|
|
got := stripPiToolCallMarkup(in)
|
|
if got != want {
|
|
t.Fatalf("unexpected stripped text: %q, want %q", got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestDrainPiTextBufferSplitToolCall(t *testing.T) {
|
|
chunks := []string{
|
|
"before ca",
|
|
`ll:bash{command:<|"|>ls -R repo/path`,
|
|
`/roles/example<|"|>}`,
|
|
" after",
|
|
}
|
|
var buf strings.Builder
|
|
var got strings.Builder
|
|
for _, chunk := range chunks {
|
|
got.WriteString(drainPiTextBuffer(&buf, chunk))
|
|
}
|
|
got.WriteString(flushPiTextBuffer(&buf))
|
|
if got.String() != "before after" {
|
|
t.Fatalf("unexpected streamed text: %q", got.String())
|
|
}
|
|
}
|
|
|
|
func TestDrainPiTextBufferSplitControlToken(t *testing.T) {
|
|
chunks := []string{"before <|tu", "rn>model after"}
|
|
var buf strings.Builder
|
|
var got strings.Builder
|
|
for _, chunk := range chunks {
|
|
got.WriteString(drainPiTextBuffer(&buf, chunk))
|
|
}
|
|
got.WriteString(flushPiTextBuffer(&buf))
|
|
if got.String() != "before after" {
|
|
t.Fatalf("unexpected streamed text: %q", got.String())
|
|
}
|
|
}
|
|
|
|
func TestFlushPiTextBufferKeepsUnmatchedToolPrefixes(t *testing.T) {
|
|
tests := []string{
|
|
"plain response: see below",
|
|
"plain call: see below",
|
|
`plain call:bash{command:<|"|>unterminated`,
|
|
}
|
|
for _, want := range tests {
|
|
var buf strings.Builder
|
|
got := drainPiTextBuffer(&buf, want)
|
|
got += flushPiTextBuffer(&buf)
|
|
if got != want {
|
|
t.Fatalf("unexpected flushed text: %q, want %q", got, want)
|
|
}
|
|
}
|
|
}
|