| 💾 Saved. d722711 k33g 4h ago | 1 | package acp |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "context" |
| 6 | "encoding/json" |
| 7 | "errors" |
| 8 | "io" |
| 9 | "strings" |
| 10 | "testing" |
| 11 | "time" |
| 12 | |
| 13 | "mm/internal/config" |
| 14 | history "mm/internal/session" |
| 15 | |
| 16 | sdk "github.com/coder/acp-go-sdk" |
| 17 | "github.com/firebase/genkit/go/ai" |
| 18 | ) |
| 19 | |
| 20 | // A prompt mixes text and attached files (resource_link): the text is kept |
| 21 | // as-is and the attachment becomes a path the model can read with its tools — |
| 22 | // bob has no embeddedContext capability, so the content itself never travels. |
| 23 | func TestFlattenTextAndResourceLink(t *testing.T) { |
| 24 | got := flatten([]sdk.ContentBlock{ |
| 25 | sdk.TextBlock("fix the greeting"), |
| 26 | sdk.ResourceLinkBlock("main.go", "file:///work/project/hello/main.go"), |
| 27 | }) |
| 28 | want := "fix the greeting\n[attached file: /work/project/hello/main.go]" |
| 29 | if got != want { |
| 30 | t.Errorf("flatten:\n got %q\nwant %q", got, want) |
| 31 | } |
| 32 | } |
| 33 | |
| 34 | // Unknown kinds must render, not fail: the kind only drives the client's |
| 35 | // display, so "execute" is the safe default. |
| 36 | func TestToolKindDefaultsToExecute(t *testing.T) { |
| 37 | cases := map[string]sdk.ToolKind{ |
| 38 | "read": sdk.ToolKindRead, |
| 39 | "edit": sdk.ToolKindEdit, |
| 40 | "execute": sdk.ToolKindExecute, |
| 41 | "": sdk.ToolKindExecute, |
| 42 | "exotic": sdk.ToolKindExecute, |
| 43 | } |
| 44 | for in, want := range cases { |
| 45 | if got := toolKind(in); got != want { |
| 46 | t.Errorf("toolKind(%q) = %q, want %q", in, got, want) |
| 47 | } |
| 48 | } |
| 49 | } |
| 50 | |
| 51 | // The list the editor is shown must be exactly what session/prompt |
| 52 | // intercepts — no more (a /quit or /abort promised here would reach the model |
| 53 | // as a question; the editor's own stop key cancels a turn) and no less (a |
| 54 | // command that works but is not listed is one the user never finds). Names |
| 55 | // are bare: ACP clients add the slash. |
| 56 | func TestAvailableCommandsMatchWhatPromptIntercepts(t *testing.T) { |
| 57 | want := []string{history.NewCommandName, "compact"} |
| 58 | cmds := availableCommands() |
| 59 | if len(cmds) != len(want) { |
| 60 | t.Fatalf("availableCommands: %d command(s), want %d: %+v", len(cmds), len(want), cmds) |
| 61 | } |
| 62 | for i, c := range cmds { |
| 63 | if c.Name != want[i] { |
| 64 | t.Errorf("command %d is %q, want %q", i, c.Name, want[i]) |
| 65 | } |
| 66 | if c.Description == "" { |
| 67 | t.Errorf("%q has no description: the editor would show a bare name", c.Name) |
| 68 | } |
| 69 | if c.Input != nil { |
| 70 | t.Errorf("%q takes no input, but an input spec is advertised", c.Name) |
| 71 | } |
| 72 | // Each advertised name must be recognised when it comes back as "/name". |
| 73 | typed := "/" + c.Name |
| 74 | if !history.IsNewCommand(typed) && !isCommand(typed, compactCommand) { |
| 75 | t.Errorf("%q is advertised but session/prompt would hand it to the model", typed) |
| 76 | } |
| 77 | } |
| 78 | } |
| 79 | |
| 80 | // Same rule as /new: the command is the whole line, spaces aside. |
| 81 | func TestIsCommandTrimsAndRequiresExactMatch(t *testing.T) { |
| 82 | cases := map[string]bool{ |
| 83 | "/compact": true, |
| 84 | " /compact\n": true, |
| 85 | "/compact now": false, |
| 86 | "/compacts": false, |
| 87 | "please /compact": false, |
| 88 | } |
| 89 | for in, want := range cases { |
| 90 | if got := isCommand(in, compactCommand); got != want { |
| 91 | t.Errorf("isCommand(%q) = %v, want %v", in, got, want) |
| 92 | } |
| 93 | } |
| 94 | } |
| 95 | |
| 96 | // resetSession is what /new does to a session: the history goes back to the |
| 97 | // system prompt alone, the count says how much was dropped, and everything |
| 98 | // that is NOT the history — cwd, the "allow always" grants — survives, because |
| 99 | // the editor still sees the same session. |
| 100 | func TestResetSessionKeepsIdentityDropsHistory(t *testing.T) { |
| 101 | s := &session{ |
| 102 | cwd: "/work/project", |
| 103 | messages: history.Fresh("you are bob"), |
| 104 | allowAlways: map[string]bool{"bash": true}, |
| 105 | } |
| 106 | s.messages = append(s.messages, ai.NewUserTextMessage("ls"), ai.NewModelTextMessage("main.go")) |
| 107 | |
| 108 | if got := resetSession(s, "you are bob"); got != 2 { |
| 109 | t.Errorf("forgotten = %d, want 2", got) |
| 110 | } |
| 111 | if len(s.messages) != 1 || s.messages[0].Role != ai.RoleSystem || s.messages[0].Text() != "you are bob" { |
| 112 | t.Errorf("history after reset: %+v, want the system prompt alone", s.messages) |
| 113 | } |
| 114 | if s.cwd != "/work/project" { |
| 115 | t.Errorf("cwd changed to %q", s.cwd) |
| 116 | } |
| 117 | if !s.allowAlways["bash"] { |
| 118 | t.Error("the 'allow always' grant for bash was dropped") |
| 119 | } |
| 120 | |
| 121 | // A second /new on a fresh session forgets nothing — and says so. |
| 122 | if got := resetSession(s, "you are bob"); got != 0 { |
| 123 | t.Errorf("forgotten on a fresh session = %d, want 0", got) |
| 124 | } |
| 125 | } |
| 126 | |
| 127 | // The command is recognised from a flattened prompt — the editor sends it as |
| 128 | // a text block, exactly as typed. |
| 129 | func TestNewCommandIsRecognisedFromPromptBlocks(t *testing.T) { |
| 130 | if !history.IsNewCommand(flatten([]sdk.ContentBlock{sdk.TextBlock("/new")})) { |
| 131 | t.Error("a text block '/new' is not recognised as the command") |
| 132 | } |
| 133 | if history.IsNewCommand(flatten([]sdk.ContentBlock{sdk.TextBlock("/new"), sdk.ResourceLinkBlock("main.go", "file:///w/main.go")})) { |
| 134 | t.Error("'/new' with an attached file is a question about that file, not the command") |
| 135 | } |
| 136 | } |
| 137 | |
| 138 | // responseSessionID must recognise exactly the line the commands hook waits |
| 139 | // for: a response whose result carries a sessionId. A notification (no id) |
| 140 | // or a prompt response (no sessionId) must not fire the hook — the first |
| 141 | // would announce the commands before the editor knows the session, the |
| 142 | // second would never come from session/new at all. |
| 143 | func TestResponseSessionIDMatchesOnlySessionResponses(t *testing.T) { |
| 144 | cases := map[string]struct { |
| 145 | line string |
| 146 | want sdk.SessionId |
| 147 | }{ |
| 148 | "session/new response": {`{"jsonrpc":"2.0","id":1,"result":{"sessionId":"bob-1-1"}}`, "bob-1-1"}, |
| 149 | "string id": {`{"jsonrpc":"2.0","id":"a","result":{"sessionId":"bob-1-2"}}`, "bob-1-2"}, |
| 150 | "prompt response": {`{"jsonrpc":"2.0","id":2,"result":{"stopReason":"end_turn"}}`, ""}, |
| 151 | "session/update": {`{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"bob-1-1","update":{}}}`, ""}, |
| 152 | "request from the agent": {`{"jsonrpc":"2.0","id":3,"method":"session/request_permission","params":{"sessionId":"bob-1-1"}}`, ""}, |
| 153 | "garbage": {`not json`, ""}, |
| 154 | } |
| 155 | for name, c := range cases { |
| 156 | got, ok := responseSessionID([]byte(c.line)) |
| 157 | if ok != (c.want != "") || got != c.want { |
| 158 | t.Errorf("%s: responseSessionID = (%q, %v), want (%q, %v)", name, got, ok, c.want, c.want != "") |
| 159 | } |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | // The order on the wire is the whole point: the editor must read the |
| 164 | // session/new response — and so learn the session id — before the |
| 165 | // available_commands_update that names it. This drives the real SDK over |
| 166 | // pipes, playing the client by hand, and reads the raw lines back. |
| 167 | func TestAvailableCommandsAreSentAfterTheSessionNewResponse(t *testing.T) { |
| 168 | clientToAgentR, clientToAgentW := io.Pipe() |
| 169 | agentToClientR, agentToClientW := io.Pipe() |
| 170 | t.Cleanup(func() { |
| 171 | _ = clientToAgentW.Close() |
| 172 | _ = agentToClientR.Close() |
| 173 | }) |
| 174 | newFront(nil, "system prompt", nil, agentToClientW, clientToAgentR) |
| 175 | |
| 176 | go func() { |
| 177 | _, _ = io.WriteString(clientToAgentW, `{"jsonrpc":"2.0","id":7,"method":"session/new","params":{"cwd":"/work","mcpServers":[]}}`+"\n") |
| 178 | }() |
| 179 | |
| 180 | lines := make(chan string) |
| 181 | go func() { |
| 182 | sc := bufio.NewScanner(agentToClientR) |
| 183 | for sc.Scan() { |
| 184 | lines <- sc.Text() |
| 185 | } |
| 186 | close(lines) |
| 187 | }() |
| 188 | next := func() string { |
| 189 | select { |
| 190 | case l, ok := <-lines: |
| 191 | if !ok { |
| 192 | t.Fatal("agent closed its output before sending both lines") |
| 193 | } |
| 194 | return l |
| 195 | case <-time.After(5 * time.Second): |
| 196 | t.Fatal("timed out waiting for the agent's next line") |
| 197 | } |
| 198 | return "" |
| 199 | } |
| 200 | |
| 201 | first, second := next(), next() |
| 202 | |
| 203 | sid, ok := responseSessionID([]byte(first)) |
| 204 | if !ok { |
| 205 | t.Fatalf("first line is not the session/new response:\n%s", first) |
| 206 | } |
| 207 | var upd struct { |
| 208 | Method string `json:"method"` |
| 209 | Params struct { |
| 210 | SessionId sdk.SessionId `json:"sessionId"` |
| 211 | Update struct { |
| 212 | SessionUpdate string `json:"sessionUpdate"` |
| 213 | AvailableCommands []sdk.AvailableCommand `json:"availableCommands"` |
| 214 | } `json:"update"` |
| 215 | } `json:"params"` |
| 216 | } |
| 217 | if err := json.Unmarshal([]byte(second), &upd); err != nil { |
| 218 | t.Fatalf("second line is not JSON: %v\n%s", err, second) |
| 219 | } |
| 220 | if upd.Method != "session/update" || upd.Params.Update.SessionUpdate != "available_commands_update" { |
| 221 | t.Fatalf("second line is not an available_commands_update:\n%s", second) |
| 222 | } |
| 223 | if upd.Params.SessionId != sid { |
| 224 | t.Errorf("commands announced for session %q, response created %q", upd.Params.SessionId, sid) |
| 225 | } |
| 226 | want := availableCommands() |
| 227 | if len(upd.Params.Update.AvailableCommands) != len(want) { |
| 228 | t.Fatalf("announced %d command(s), want %d: %+v", len(upd.Params.Update.AvailableCommands), len(want), upd.Params.Update.AvailableCommands) |
| 229 | } |
| 230 | for i, c := range upd.Params.Update.AvailableCommands { |
| 231 | if c.Name != want[i].Name { |
| 232 | t.Errorf("announced command %d is %q, want %q", i, c.Name, want[i].Name) |
| 233 | } |
| 234 | } |
| 235 | } |
| 236 | |
| 237 | // --- /compact and /abort -------------------------------------------------------- |
| 238 | |
| 239 | // oneTurn is a question answered after one bash command: the smallest unit |
| 240 | // compaction cuts on. |
| 241 | func oneTurn(q string) []*ai.Message { |
| 242 | return []*ai.Message{ |
| 243 | ai.NewUserTextMessage(q), |
| 244 | ai.NewModelMessage(ai.NewToolRequestPart(&ai.ToolRequest{Ref: q, Name: "bash", Input: map[string]any{"command": "echo " + q}})), |
| 245 | ai.NewMessage(ai.RoleTool, nil, ai.NewToolResponsePart(&ai.ToolResponse{Ref: q, Name: "bash", Output: q})), |
| 246 | ai.NewModelMessage(ai.NewTextPart("done " + q)), |
| 247 | } |
| 248 | } |
| 249 | |
| 250 | func sessionWithTurns(qs ...string) *session { |
| 251 | s := &session{cwd: "/work", messages: history.Fresh("You are Bob."), allowAlways: map[string]bool{"bash": true}} |
| 252 | for _, q := range qs { |
| 253 | s.messages = append(s.messages, oneTurn(q)...) |
| 254 | } |
| 255 | return s |
| 256 | } |
| 257 | |
| 258 | // The three outcomes of /compact, in the REPL's words. Only the first one |
| 259 | // touches the history; the two others leave it — and the server's token |
| 260 | // count, which the caller resets on `compressed` alone — exactly as they were. |
| 261 | func TestCompactHistoryOutcomes(t *testing.T) { |
| 262 | cfg := config.ContextConfig{Enabled: true, KeepLastTurns: 1, Threshold: 75, SummaryMaxTokens: 100} |
| 263 | summary := func(context.Context, []*ai.Message) (string, error) { return "the user ran a and b", nil } |
| 264 | explain := func(err error) string { return "explained: " + err.Error() } |
| 265 | |
| 266 | t.Run("compresses the turns older than the kept ones", func(t *testing.T) { |
| 267 | s := sessionWithTurns("a", "b", "c") |
| 268 | before := len(s.messages) |
| 269 | line, compressed := compactHistory(context.Background(), s, cfg, summary, explain) |
| 270 | if !compressed { |
| 271 | t.Fatalf("compressed = false, line %q", line) |
| 272 | } |
| 273 | if !strings.HasPrefix(line, "🗜️ compressed 8 messages → 1 summary + 4 kept") { |
| 274 | t.Errorf("report %q does not open with the REPL's words and the right counts", line) |
| 275 | } |
| 276 | if len(s.messages) >= before { |
| 277 | t.Errorf("history has %d messages after compaction, %d before", len(s.messages), before) |
| 278 | } |
| 279 | if s.messages[0].Role != ai.RoleSystem || s.allowAlways["bash"] != true || s.cwd != "/work" { |
| 280 | t.Error("compaction touched something other than the turns") |
| 281 | } |
| 282 | }) |
| 283 | |
| 284 | t.Run("nothing to compact is said, not swallowed", func(t *testing.T) { |
| 285 | s := sessionWithTurns("a") |
| 286 | before := len(s.messages) |
| 287 | line, compressed := compactHistory(context.Background(), s, cfg, summary, explain) |
| 288 | if compressed || len(s.messages) != before { |
| 289 | t.Errorf("compressed = %v, %d messages (was %d): a single kept turn must not move", compressed, len(s.messages), before) |
| 290 | } |
| 291 | if line != "🗜️ nothing to compact: 5 message(s), no turn older than the last 1" { |
| 292 | t.Errorf("line %q", line) |
| 293 | } |
| 294 | }) |
| 295 | |
| 296 | t.Run("a failed summary keeps the history and names the cause", func(t *testing.T) { |
| 297 | s := sessionWithTurns("a", "b", "c") |
| 298 | before := len(s.messages) |
| 299 | failing := func(context.Context, []*ai.Message) (string, error) { return "", errors.New("server down") } |
| 300 | line, compressed := compactHistory(context.Background(), s, cfg, failing, explain) |
| 301 | if compressed || len(s.messages) != before { |
| 302 | t.Errorf("compressed = %v, %d messages (was %d): a failure must leave the history alone", compressed, len(s.messages), before) |
| 303 | } |
| 304 | if line != "[compact: failed, history kept: explained: server down]" { |
| 305 | t.Errorf("line %q", line) |
| 306 | } |
| 307 | }) |
| 308 | } |