Download
108 lines 4.3 KiBGoModified 2026-10-31 raw
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)
	}
}