1package agent
2
3import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "strings"
8 "testing"
9 "time"
10
11 "github.com/patriceckhart/zot/packages/core"
12 "github.com/patriceckhart/zot/packages/provider"
13)
14
15func TestRPCRejectsMutationsDuringPromptPreparation(t *testing.T) {
16 models := provider.ModelsForProvider("anthropic")
17 if len(models) < 2 {
18 t.Fatal("need two catalog models")
19 }
20 modelJSON, err := json.Marshal(map[string]string{"model": models[1].ID})
21 if err != nil {
22 t.Fatal(err)
23 }
24 for _, tc := range []struct {
25 command string
26 raw []byte
27 }{
28 {"clear", nil},
29 {"set_model", modelJSON},
30 {"set_reasoning", []byte(`{"reasoning":"max"}`)},
31 } {
32 t.Run(tc.command, func(t *testing.T) {
33 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
34 defer cancel()
35 client := &startCaptureClient{}
36 ag := core.NewAgent(client, models[0].ID, "base", nil)
37 entered := make(chan struct{})
38 ag.BeforeStart = func(ctx context.Context, system string) string {
39 close(entered)
40 <-ctx.Done()
41 return "stale replacement"
42 }
43 var out bytes.Buffer
44 s := &rpcServer{ctx: ctx, agent: ag, provider: "anthropic", model: ag.Model, out: &out}
45 t.Cleanup(func() { cancel(); s.inFlight.Wait() })
46 s.dispatch("prompt", "prompt-1", []byte(`{"message":"hello"}`))
47 select {
48 case <-entered:
49 case <-ctx.Done():
50 t.Fatal("preparation did not start")
51 }
52 dispatched := make(chan struct{})
53 go func() { defer close(dispatched); s.dispatch(tc.command, "mutation", tc.raw) }()
54 select {
55 case <-dispatched:
56 case <-ctx.Done():
57 t.Fatal("mutation blocked the RPC read loop")
58 }
59 s.writeMu.Lock()
60 responses := append([]byte(nil), out.Bytes()...)
61 s.writeMu.Unlock()
62 var rejected bool
63 for _, line := range bytes.Split(responses, []byte("\n")) {
64 if len(line) == 0 {
65 continue
66 }
67 var response struct {
68 ID string
69 Success bool
70 Error string
71 }
72 if err := json.Unmarshal(line, &response); err != nil {
73 t.Fatal(err)
74 }
75 if response.ID == "mutation" {
76 rejected = !response.Success && strings.Contains(response.Error, "busy")
77 }
78 }
79 // Abort must remain usable while the hook waits. Wait before inspecting
80 // agent fields so the test itself does not race with prompt preparation.
81 s.dispatch("abort", "abort-1", nil)
82 s.inFlight.Wait()
83 if !rejected {
84 t.Errorf("mutation should return a busy error: %s", responses)
85 }
86 if ag.Model != models[0].ID || ag.Reasoning != "" || len(ag.Messages()) != 1 {
87 t.Fatal("busy mutation changed active agent state")
88 }
89 if ag.System != "base" || len(client.requests) != 0 {
90 t.Fatal("canceled hook result reached model")
91 }
92 // The same command succeeds when idle; the next preparation uses current
93 // state rather than marking the canceled extension response as prepared.
94 s.dispatch(tc.command, "idle-mutation", tc.raw)
95 switch tc.command {
96 case "clear":
97 if len(ag.Messages()) != 0 {
98 t.Fatal("idle clear failed")
99 }
100 case "set_model":
101 if ag.Model != models[1].ID {
102 t.Fatal("idle model change failed")
103 }
104 case "set_reasoning":
105 if ag.Reasoning != "max" {
106 t.Fatal("idle reasoning change failed")
107 }
108 }
109 ag.BeforeStart = func(_ context.Context, base string) string {
110 if base != "base" {
111 t.Error("stale prompt became base")
112 }
113 return "fresh replacement"
114 }
115 s.runPrompt("prompt-2", "again", nil)
116 if len(client.requests) != 1 || client.requests[0].System != "fresh replacement" || client.requests[0].Model != ag.Model {
117 t.Fatal("next prompt did not use freshly prepared state")
118 }
119 })
120 }
121}
122
123func TestRPCRejectsMutationsWhileTurnLockHeld(t *testing.T) {
124 // The same lock covers active model/tool work and compaction, not just
125 // preparation. Rejection must not depend on activeCancel being populated.
126 var out bytes.Buffer
127 s := &rpcServer{agent: core.NewAgent(nil, "model", "base", nil), out: &out}
128 s.turnMu.Lock()
129 defer s.turnMu.Unlock()
130 s.dispatch("clear", "1", nil)
131 if !strings.Contains(out.String(), `"success":false`) || !strings.Contains(out.String(), "busy") {
132 t.Fatalf("missing busy response: %s", out.String())
133 }
134}