123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108 | package agent
import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"blinker/internal/llm"
"blinker/internal/policy"
)
// proposing answers each request with the next reply: a Propose_Policy
// call with those arguments, or, for one not starting with "{", words. It
// records the messages of each request.
func proposing(t *testing.T, replies ...string) (*llm.Client, *[][]llm.Message) {
var seen [][]llm.Message
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var req struct{ Messages []llm.Message }
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Errorf("request: %v", err)
}
seen = append(seen, req.Messages)
if len(seen) > len(replies) {
t.Errorf("unexpected request %d", len(seen))
w.WriteHeader(http.StatusBadRequest)
return
}
msg := map[string]any{"role": "assistant", "content": replies[len(seen)-1]}
if strings.HasPrefix(replies[len(seen)-1], "{") {
msg = map[string]any{"role": "assistant", "content": "", "tool_calls": []map[string]any{{
"id": fmt.Sprintf("c%d", len(seen)), "type": "function",
"function": map[string]any{"name": "Propose_Policy", "arguments": replies[len(seen)-1]},
}}}
}
if err := json.NewEncoder(w).Encode(map[string]any{"choices": []map[string]any{{"message": msg}}}); err != nil {
t.Errorf("response: %v", err)
}
}))
t.Cleanup(srv.Close)
return llm.New(srv.URL, "", "m"), &seen
}
// loadPolicies is a new project's policies, from the named policies
// blinker is built with.
func loadPolicies(t *testing.T, project string) *policy.Loaded {
t.Helper()
dir := t.TempDir()
s := policy.Store{ConfigDir: dir, Dirs: policy.NamedDirs(dir, nil), Builtin: os.DirFS("../../policies")}
l, err := s.Load(project)
if err != nil {
t.Fatal(err)
}
return l
}
func TestPolicyDraft(t *testing.T) {
env := policy.Env{Home: "/home/u", Project: "/home/u/proj"}
l := loadPolicies(t, env.Project)
client, seen := proposing(t,
`{"file": "unsafe", "policy": "project /home/u/proj\n", "extra": 1}`, // unknown field
`{"file": "unsafe", "policy": "file write docs\n"}`, // no project line
`{"file": "unsafe", "policy": "project /home/u/proj\ninclude agent-unsafe\n# Docs.\nfile write docs\n"}`,
`I need to know which directory you mean.`,
)
d := NewPolicyDraft(client, l, env, nil, "let it write docs")
checked := 0
p, reply, err := d.Propose(context.Background(), func(PolicyProposal) error { checked++; return nil })
if err != nil || reply != "" || p.Source.Path != l.Own[policy.UnsafeMode].Path || p.File.Directory != env.Project || len(p.File.Files) != 1 || checked != 1 {
t.Fatalf("proposal %+v, reply %q, err %v, checked %d", p, reply, err, checked)
}
// The invalid proposals went back to the model with the reasons.
last := (*seen)[2]
if got := last[len(last)-1]; got.Role != "tool" || !strings.Contains(got.Content, "Not valid") || !strings.Contains(got.Content, `needs the line "project /home/u/proj"`) {
t.Errorf("feedback = %+v", got)
}
// Asking for changes answers the proposal; the model may answer in words.
d.Revise("only docs/api")
p, reply, err = d.Propose(context.Background(), func(PolicyProposal) error { return nil })
if err != nil || p.File != nil || !strings.Contains(reply, "which directory") {
t.Fatalf("proposal %+v, reply %q, err %v", p, reply, err)
}
msgs := (*seen)[3]
if got := msgs[len(msgs)-1]; got.Role != "tool" || got.ToolCallID != "c3" || !strings.Contains(got.Content, "only docs/api") {
t.Errorf("revision = %+v", got)
}
// Only the request and the policies reach the model, each by the name
// it proposes it by.
if !strings.Contains(msgs[0].Content, `"agent-unsafe", policy agent-unsafe`) || msgs[0].Role != "system" || msgs[1].Content != "let it write docs" {
t.Errorf("conversation starts %+v", msgs[:2])
}
}
func TestPolicyDraftGivesUp(t *testing.T) {
env := policy.Env{Home: "/home/u", Project: "/home/u/proj"}
l := loadPolicies(t, env.Project)
bad := `{"file": "agent-safe", "policy": ""}`
client, _ := proposing(t, bad, bad, bad)
d := NewPolicyDraft(client, l, env, nil, "x")
_, _, err := d.Propose(context.Background(), func(PolicyProposal) error { return fmt.Errorf("does not compile") })
if err == nil || !strings.Contains(err.Error(), "does not compile") {
t.Errorf("err = %v", err)
}
}
|