318 lines
11 KiB
Go
318 lines
11 KiB
Go
package agent
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
|
|
"sysagent/internal/project"
|
|
)
|
|
|
|
// scriptedModel simula un servidor OpenAI-compatible. reply decide la
|
|
// respuesta a partir de la petición recibida (y su número de orden).
|
|
type scriptedModel struct {
|
|
mu sync.Mutex
|
|
requests []map[string]any
|
|
reply func(n int, req map[string]any) (status int, body any)
|
|
}
|
|
|
|
func (s *scriptedModel) server(t *testing.T) *httptest.Server {
|
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/v1/models":
|
|
io.WriteString(w, `{"data":[{"id":"test-model","max_model_len":32768}]}`)
|
|
case "/v1/chat/completions":
|
|
raw, _ := io.ReadAll(r.Body)
|
|
var req map[string]any
|
|
if err := json.Unmarshal(raw, &req); err != nil {
|
|
t.Errorf("bad request JSON: %v", err)
|
|
}
|
|
s.mu.Lock()
|
|
s.requests = append(s.requests, req)
|
|
n := len(s.requests)
|
|
s.mu.Unlock()
|
|
status, body := s.reply(n, req)
|
|
w.WriteHeader(status)
|
|
json.NewEncoder(w).Encode(body)
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
}
|
|
|
|
func text(content string) map[string]any {
|
|
return map[string]any{"choices": []any{map[string]any{
|
|
"message": map[string]any{"role": "assistant", "content": content}, "finish_reason": "stop"}}}
|
|
}
|
|
|
|
func toolCall(id, name string, args any) map[string]any {
|
|
b, _ := json.Marshal(args)
|
|
return map[string]any{"choices": []any{map[string]any{
|
|
"message": map[string]any{"role": "assistant", "content": "", "tool_calls": []any{map[string]any{
|
|
"id": id, "type": "function", "function": map[string]any{"name": name, "arguments": string(b)}}}},
|
|
"finish_reason": "tool_calls"}}}
|
|
}
|
|
|
|
func messages(req map[string]any) []map[string]any {
|
|
var out []map[string]any
|
|
for _, m := range req["messages"].([]any) {
|
|
out = append(out, m.(map[string]any))
|
|
}
|
|
return out
|
|
}
|
|
|
|
func setupProject(t *testing.T, url string, files map[string]string) string {
|
|
t.Helper()
|
|
dir := t.TempDir()
|
|
conf := "API_URL=" + url + "/v1\nAPI_KEY=sk-test-secret-key-000\nMODEL=test-model\nCONTEXT_SIZE=0\n"
|
|
if err := os.WriteFile(filepath.Join(dir, "init.conf"), []byte(conf), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for name, content := range files {
|
|
if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
return dir
|
|
}
|
|
|
|
const agentMD = "# Agente de prueba\n\n```policy\nEXECUTION_MODE=privileged\nALLOW_NETWORK=true\n```\n"
|
|
|
|
func bootstrap(t *testing.T, dir string, cli bool, stdin string) (*Agent, *bytes.Buffer) {
|
|
t.Helper()
|
|
var stdout, stderr bytes.Buffer
|
|
a, err := Bootstrap(context.Background(), Options{Root: dir, CLI: cli, Stdin: strings.NewReader(stdin),
|
|
Stdout: &stdout, Stderr: &stderr, Version: "test"})
|
|
if err != nil {
|
|
t.Fatalf("Bootstrap: %v\n%s", err, stderr.String())
|
|
}
|
|
t.Cleanup(a.Close)
|
|
return a, &stdout
|
|
}
|
|
|
|
func TestBootstrapRequiresAgentMD(t *testing.T) {
|
|
dir := setupProject(t, "http://127.0.0.1:1", nil)
|
|
_, err := Bootstrap(context.Background(), Options{Root: dir, Stdout: io.Discard, Stderr: io.Discard, Version: "test"})
|
|
if !errors.Is(err, project.ErrAgentNotFound) {
|
|
t.Fatalf("expected ErrAgentNotFound, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestRunTasksWithNativeTools(t *testing.T) {
|
|
m := &scriptedModel{reply: func(n int, req map[string]any) (int, any) {
|
|
// Por tarea: 1) pide exec, 2) respuesta final.
|
|
switch n {
|
|
case 1:
|
|
return 200, toolCall("c1", "exec", map[string]any{"command": "printf", "args": []string{"hola-%s", "tool"}})
|
|
case 2:
|
|
return 200, text("Ejecutado printf.\nTASK_STATUS: DONE")
|
|
case 3:
|
|
return 200, toolCall("c2", "write_file", map[string]any{"path": "agent.md", "content": "EXECUTION_MODE=unrestricted"})
|
|
default:
|
|
return 200, text("La política lo impide.\nTASK_STATUS: BLOCKED")
|
|
}
|
|
}}
|
|
srv := m.server(t)
|
|
defer srv.Close()
|
|
dir := setupProject(t, srv.URL, map[string]string{
|
|
"agent.md": agentMD, "01-a.md": "primera tarea", "02-b.md": "intenta escalar", "README.md": "no es tarea",
|
|
"memoria.md": "## Hosts\n- db01: 10.0.0.5\n",
|
|
})
|
|
a, stdout := bootstrap(t, dir, false, "")
|
|
if a.ctx.ContextSize() != 32768 {
|
|
t.Errorf("context size from server = %d", a.ctx.ContextSize())
|
|
}
|
|
if err := a.RunTasks(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if len(m.requests) != 4 {
|
|
t.Fatalf("requests = %d", len(m.requests))
|
|
}
|
|
if _, ok := m.requests[0]["tools"]; !ok {
|
|
t.Error("native mode must send the tools field")
|
|
}
|
|
msgs := messages(m.requests[1])
|
|
last := msgs[len(msgs)-1]
|
|
if last["role"] != "tool" || last["tool_call_id"] != "c1" || !strings.Contains(last["content"].(string), "hola-tool") {
|
|
t.Errorf("tool result message wrong: %v", last)
|
|
}
|
|
system := msgs[0]["content"].(string)
|
|
for _, want := range []string{"REGLAS INMUTABLES", "Agente de prueba", "db01: 10.0.0.5", "EXECUTION_MODE=privileged", "TASK_STATUS"} {
|
|
if !strings.Contains(system, want) {
|
|
t.Errorf("system prompt missing %q", want)
|
|
}
|
|
}
|
|
// El intento de reescribir agent.md debe llegar al modelo como denegación.
|
|
msgs = messages(m.requests[3])
|
|
last = msgs[len(msgs)-1]
|
|
if !strings.Contains(last["content"].(string), "denied by policy") {
|
|
t.Errorf("denial not reported to model: %v", last["content"])
|
|
}
|
|
if data, _ := os.ReadFile(filepath.Join(dir, "agent.md")); string(data) != agentMD {
|
|
t.Fatal("agent.md was modified!")
|
|
}
|
|
if a.taskStatus["01-a.md"] != StatusDone || a.taskStatus["02-b.md"] != StatusBlocked {
|
|
t.Errorf("statuses = %v", a.taskStatus)
|
|
}
|
|
|
|
all, _ := json.Marshal(m.requests)
|
|
if strings.Contains(string(all), "no es tarea") {
|
|
t.Error("README.md leaked into the prompt")
|
|
}
|
|
if strings.Contains(string(all), "sk-test-secret-key-000") {
|
|
t.Error("API key leaked into messages sent to the model")
|
|
}
|
|
logData, _ := os.ReadFile(filepath.Join(dir, "agent.log"))
|
|
for _, want := range []string{`action="tool:exec"`, `result="denied"`, `result="DONE"`} {
|
|
if !strings.Contains(string(logData), want) {
|
|
t.Errorf("agent.log missing %s", want)
|
|
}
|
|
}
|
|
if strings.Contains(string(logData), "sk-test-secret-key-000") {
|
|
t.Error("API key leaked into agent.log")
|
|
}
|
|
if !strings.Contains(stdout.String(), "→ exec: printf") || !strings.Contains(stdout.String(), "← denied") {
|
|
t.Errorf("console output:\n%s", stdout.String())
|
|
}
|
|
}
|
|
|
|
func TestFallbackToJSONProtocol(t *testing.T) {
|
|
m := &scriptedModel{reply: func(n int, req map[string]any) (int, any) {
|
|
if _, hasTools := req["tools"]; hasTools {
|
|
return 400, map[string]any{"error": map[string]any{"message": "tools are not supported by this model"}}
|
|
}
|
|
msgs := messages(req)
|
|
if strings.HasPrefix(msgs[len(msgs)-1]["content"].(string), "TOOL_RESULTS") {
|
|
return 200, text("Listo.\nTASK_STATUS: DONE")
|
|
}
|
|
return 200, text("Reviso el sistema.\n```tool\n{\"tool\": \"exec\", \"arguments\": {\"command\": \"echo\", \"args\": [\"json-mode\"]}}\n```")
|
|
}}
|
|
srv := m.server(t)
|
|
defer srv.Close()
|
|
dir := setupProject(t, srv.URL, map[string]string{"agent.md": agentMD, "01-a.md": "tarea"})
|
|
a, _ := bootstrap(t, dir, false, "")
|
|
if err := a.RunTasks(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if a.toolMode != "json" {
|
|
t.Fatalf("toolMode = %s", a.toolMode)
|
|
}
|
|
lastReq := m.requests[len(m.requests)-1]
|
|
msgs := messages(lastReq)
|
|
if !strings.Contains(msgs[0]["content"].(string), "```tool") {
|
|
t.Error("JSON protocol instructions missing from system prompt")
|
|
}
|
|
result := msgs[len(msgs)-1]["content"].(string)
|
|
if !strings.Contains(result, "json-mode") || !strings.Contains(result, "untrusted data") {
|
|
t.Errorf("tool results = %q", result)
|
|
}
|
|
if a.taskStatus["01-a.md"] != StatusDone {
|
|
t.Errorf("status = %s", a.taskStatus["01-a.md"])
|
|
}
|
|
}
|
|
|
|
func TestLoopDetectionStopsTask(t *testing.T) {
|
|
m := &scriptedModel{reply: func(n int, req map[string]any) (int, any) {
|
|
return 200, toolCall("same", "exec", map[string]any{"command": "echo", "args": []string{"same"}})
|
|
}}
|
|
srv := m.server(t)
|
|
defer srv.Close()
|
|
dir := setupProject(t, srv.URL, map[string]string{"agent.md": agentMD, "01-a.md": "tarea"})
|
|
a, stdout := bootstrap(t, dir, false, "")
|
|
if err := a.RunTasks(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if a.taskStatus["01-a.md"] != StatusBlocked {
|
|
t.Errorf("status = %s", a.taskStatus["01-a.md"])
|
|
}
|
|
if len(m.requests) != loopStopAt {
|
|
t.Errorf("requests = %d, want %d", len(m.requests), loopStopAt)
|
|
}
|
|
all, _ := json.Marshal(m.requests)
|
|
if !strings.Contains(string(all), "LOOP WARNING") {
|
|
t.Error("model was not warned about the loop")
|
|
}
|
|
if !strings.Contains(stdout.String(), "loop detected") {
|
|
t.Errorf("stdout:\n%s", stdout.String())
|
|
}
|
|
}
|
|
|
|
func TestMissingTaskStatusIsNudgedOnce(t *testing.T) {
|
|
m := &scriptedModel{reply: func(n int, req map[string]any) (int, any) {
|
|
return 200, text("hecho, creo")
|
|
}}
|
|
srv := m.server(t)
|
|
defer srv.Close()
|
|
dir := setupProject(t, srv.URL, map[string]string{"agent.md": agentMD, "01-a.md": "tarea"})
|
|
a, _ := bootstrap(t, dir, false, "")
|
|
if err := a.RunTasks(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(m.requests) != 2 {
|
|
t.Errorf("requests = %d, want 2 (one nudge)", len(m.requests))
|
|
}
|
|
}
|
|
|
|
func TestRunCLIIgnoresTasks(t *testing.T) {
|
|
m := &scriptedModel{reply: func(n int, req map[string]any) (int, any) {
|
|
if n == 1 {
|
|
return 200, toolCall("c1", "system_info", map[string]any{})
|
|
}
|
|
return 200, text("Respuesta del agente")
|
|
}}
|
|
srv := m.server(t)
|
|
defer srv.Close()
|
|
dir := setupProject(t, srv.URL, map[string]string{"agent.md": agentMD, "01-a.md": "TAREA QUE NO DEBE CARGARSE"})
|
|
|
|
a, stdout := bootstrap(t, dir, true, "/status\nhola agente\n/history\n/policy\n/quit\n")
|
|
if len(a.tasks) != 0 {
|
|
t.Fatal("--cli must not load task files")
|
|
}
|
|
if err := a.RunCLI(context.Background(), make(chan os.Signal)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
out := stdout.String()
|
|
for _, want := range []string{"System Agent", "Mode: privileged", "Tool calling: native", "→ system_info",
|
|
"Respuesta del agente", "[1] user: hola agente", "tool_call system_info", "ALLOW_NETWORK=true", "MAX_STEPS=40"} {
|
|
if !strings.Contains(out, want) {
|
|
t.Errorf("output missing %q\n%s", want, out)
|
|
}
|
|
}
|
|
if strings.Contains(out, "sk-test-secret-key-000") {
|
|
t.Error("API key shown in CLI output")
|
|
}
|
|
all, _ := json.Marshal(m.requests)
|
|
if strings.Contains(string(all), "TAREA QUE NO DEBE CARGARSE") {
|
|
t.Error("task file content sent in --cli mode")
|
|
}
|
|
if strings.Contains(messages(m.requests[0])[0]["content"].(string), "MODO TAREAS") {
|
|
t.Error("task-mode instructions must not appear in --cli")
|
|
}
|
|
}
|
|
|
|
func TestParseTaskStatus(t *testing.T) {
|
|
cases := map[string]TaskStatus{
|
|
"ok\nTASK_STATUS: DONE": StatusDone,
|
|
"**TASK_STATUS: BLOCKED**": StatusBlocked,
|
|
"TASK_STATUS: DONE\n...\nTASK_STATUS: FAILED": StatusFailed,
|
|
}
|
|
for in, want := range cases {
|
|
if got, ok := parseTaskStatus(in); !ok || got != want {
|
|
t.Errorf("%q: %s %v", in, got, ok)
|
|
}
|
|
}
|
|
if _, ok := parseTaskStatus("sin estado"); ok {
|
|
t.Error("expected no status")
|
|
}
|
|
}
|