| 1 | package control |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "path/filepath" |
| 6 | "reflect" |
| 7 | "slices" |
| 8 | "testing" |
| 9 | |
| 10 | "reasonix/internal/agent" |
| 11 | "reasonix/internal/event" |
| 12 | "reasonix/internal/provider" |
| 13 | ) |
| 14 | |
| 15 | func TestPinnedContextNeverChangesBasePrompt(t *testing.T) { |
| 16 | dir := t.TempDir() |
| 17 | exec := agent.New(nil, nil, agent.NewSession("legacy composed system"), agent.Options{}, event.Discard) |
| 18 | ctrl := newOwnedTestController(t, Options{ |
| 19 | Runner: exec, |
| 20 | Executor: exec, |
| 21 | SystemPrompt: "BASE", |
| 22 | SessionDir: dir, |
| 23 | SessionPath: filepath.Join(dir, "session.jsonl"), |
| 24 | Sink: event.Discard, |
| 25 | }) |
| 26 | if got := controlSystemMessage(ctrl.History()); got != "BASE" { |
| 27 | t.Fatalf("migrated system prompt = %q", got) |
| 28 | } |
| 29 | if reasons := exec.Session().DrainContentRewriteReasons(); !slices.Contains(reasons, "legacy_pinned_system_migration") { |
| 30 | t.Fatalf("migration reasons = %v", reasons) |
| 31 | } |
| 32 | |
| 33 | ctrl.ApplyExtensionSystemPrompt("EXTENSION") |
| 34 | if got := ctrl.SystemPrompt(); got != "EXTENSION" { |
| 35 | t.Fatalf("SystemPrompt = %q", got) |
| 36 | } |
| 37 | if got := controlSystemMessage(ctrl.History()); got != "EXTENSION" { |
| 38 | t.Fatalf("extension system prompt = %q", got) |
| 39 | } |
| 40 | |
| 41 | if err := ctrl.NewSession(); err != nil { |
| 42 | t.Fatalf("NewSession: %v", err) |
| 43 | } |
| 44 | if got := controlSystemMessage(ctrl.History()); got != "EXTENSION" { |
| 45 | t.Fatalf("new session system prompt = %q", got) |
| 46 | } |
| 47 | } |
| 48 | |
| 49 | func TestPinnedContextLoaderAppendsAtAdmittedTurns(t *testing.T) { |
| 50 | prov := &recordingProvider{streams: [][]provider.Chunk{ |
| 51 | {{Type: provider.ChunkText, Text: "one"}, {Type: provider.ChunkDone}}, |
| 52 | {{Type: provider.ChunkText, Text: "two"}, {Type: provider.ChunkDone}}, |
| 53 | {{Type: provider.ChunkText, Text: "three"}, {Type: provider.ChunkDone}}, |
| 54 | }} |
| 55 | exec := agent.New(prov, nil, agent.NewSession("BASE"), agent.Options{}, event.Discard) |
| 56 | content := "A" |
| 57 | loads := 0 |
| 58 | sessionPath := filepath.Join(t.TempDir(), "session.jsonl") |
| 59 | ctrl := newOwnedTestController(t, Options{ |
| 60 | Runner: exec, |
| 61 | Executor: exec, |
| 62 | SystemPrompt: "BASE", |
| 63 | SessionPath: sessionPath, |
| 64 | PinnedContextLoader: func(_ context.Context, path string) (agent.PinnedContextSnapshot, error) { |
| 65 | loads++ |
| 66 | if path != sessionPath { |
| 67 | t.Fatalf("loader path = %q", path) |
| 68 | } |
| 69 | return agent.PinnedContextSnapshot{Files: []agent.PinnedContextFile{{Path: "a.md", Content: content}}}, nil |
| 70 | }, |
| 71 | Sink: event.Discard, |
| 72 | }) |
| 73 | if err := ctrl.Run(context.Background(), "first"); err != nil { |
| 74 | t.Fatal(err) |
| 75 | } |
| 76 | if err := ctrl.Run(context.Background(), "second"); err != nil { |
| 77 | t.Fatal(err) |
| 78 | } |
| 79 | content = "B" |
| 80 | if err := ctrl.Run(context.Background(), "third"); err != nil { |
| 81 | t.Fatal(err) |
| 82 | } |
| 83 | if loads != 3 { |
| 84 | t.Fatalf("loader calls = %d", loads) |
| 85 | } |
| 86 | if got := controlSystemMessage(ctrl.History()); got != "BASE" { |
| 87 | t.Fatalf("system prompt changed: %q", got) |
| 88 | } |
| 89 | revisions := 0 |
| 90 | for _, message := range ctrl.History() { |
| 91 | if agent.IsPinnedContextRevision(message) { |
| 92 | revisions++ |
| 93 | } |
| 94 | } |
| 95 | if revisions != 2 { |
| 96 | t.Fatalf("revision messages = %d, want 2", revisions) |
| 97 | } |
| 98 | if len(prov.requests) != 3 { |
| 99 | t.Fatalf("provider requests = %d", len(prov.requests)) |
| 100 | } |
| 101 | for i := 1; i < len(prov.requests); i++ { |
| 102 | previous := prov.requests[i-1].Messages |
| 103 | current := prov.requests[i].Messages |
| 104 | if len(current) < len(previous) || !reflect.DeepEqual(current[:len(previous)], previous) { |
| 105 | t.Fatalf("request %d is not prefixed by request %d", i, i-1) |
| 106 | } |
| 107 | } |
| 108 | } |
| 109 |