| 1 | package bot |
| 2 | |
| 3 | import ( |
| 4 | "encoding/json" |
| 5 | "io" |
| 6 | "net/http" |
| 7 | "net/http/httptest" |
| 8 | "strings" |
| 9 | "testing" |
| 10 | "time" |
| 11 | |
| 12 | "lchat/internal/e2ee" |
| 13 | ) |
| 14 | |
| 15 | // Tests for research sessions (research.go): one consent tap → sealed read |
| 16 | // grants → a bounded read-only tool loop → one posted answer. The security |
| 17 | // properties pinned here: write ops never enter the palette, expired grants |
| 18 | // never reach the provider, a grant sealed to another bot stays silent, and |
| 19 | // no token ever appears in what the model is sent. |
| 20 | |
| 21 | // waitNonTyping skips the typing beats the research loop emits and returns |
| 22 | // the next real payload a thread member would decrypt. |
| 23 | func (h *connectHarness) waitNonTyping(t *testing.T) map[string]any { |
| 24 | t.Helper() |
| 25 | deadline := time.After(10 * time.Second) |
| 26 | for { |
| 27 | select { |
| 28 | case p := <-h.replies: |
| 29 | if p["type"] == "chat.typing" { |
| 30 | continue |
| 31 | } |
| 32 | return p |
| 33 | case <-deadline: |
| 34 | t.Fatal("no payload arrived") |
| 35 | return nil |
| 36 | } |
| 37 | } |
| 38 | } |
| 39 | |
| 40 | // expectSilence asserts nothing (typing aside) lands within the window. |
| 41 | func (h *connectHarness) expectSilence(t *testing.T, d time.Duration) { |
| 42 | t.Helper() |
| 43 | deadline := time.After(d) |
| 44 | for { |
| 45 | select { |
| 46 | case p := <-h.replies: |
| 47 | if p["type"] == "chat.typing" { |
| 48 | continue |
| 49 | } |
| 50 | t.Fatalf("unexpected payload: %v", p) |
| 51 | case <-deadline: |
| 52 | return |
| 53 | } |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | // scriptedClaude answers each request in order with the given raw bodies and |
| 58 | // records every request body for assertions. |
| 59 | func scriptedClaude(t *testing.T, bodies ...string) *[][]byte { |
| 60 | t.Helper() |
| 61 | var reqs [][]byte |
| 62 | i := 0 |
| 63 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 64 | raw, _ := io.ReadAll(r.Body) |
| 65 | reqs = append(reqs, raw) |
| 66 | w.Header().Set("Content-Type", "application/json") |
| 67 | if i < len(bodies) { |
| 68 | io.WriteString(w, bodies[i]) |
| 69 | i++ |
| 70 | return |
| 71 | } |
| 72 | io.WriteString(w, `{"content":[{"type":"text","text":"(script exhausted)"}],"stop_reason":"end_turn"}`) |
| 73 | })) |
| 74 | t.Cleanup(srv.Close) |
| 75 | t.Setenv("ANTHROPIC_BASE_URL", srv.URL) |
| 76 | return &reqs |
| 77 | } |
| 78 | |
| 79 | func TestResearchEndToEnd(t *testing.T) { |
| 80 | gh := newFakeGitHub(t, `[{"full_name":"paul/lchat"}]`) |
| 81 | reqs := scriptedClaude(t, |
| 82 | `{"content":[{"type":"tool_use","id":"tu1","name":"github__repos_list","input":{}}],"stop_reason":"tool_use"}`, |
| 83 | `{"content":[{"type":"text","text":"You have one repo: paul/lchat."}],"stop_reason":"end_turn"}`, |
| 84 | ) |
| 85 | h := newConnectHarness(t) |
| 86 | h.rt.ai = newAnthropicClient("test-key", nil) |
| 87 | |
| 88 | grant := h.seal(t, map[string]any{"service": "github", "token": "ghp_research", "exp": time.Now().Add(time.Minute).UnixMilli(), "access": "read"}) |
| 89 | h.rt.handleResearch("t1", h.key, &researchPayload{ |
| 90 | Type: "connect.research", Goal: "what repos do I have?", |
| 91 | Grants: map[string]string{"github": grant}, RID: "rr1", UID: "u1", Name: "Paul", |
| 92 | }) |
| 93 | |
| 94 | ack := h.waitNonTyping(t) |
| 95 | if ack["type"] != "connect.result" || ack["rid"] != "rr1" || ack["data"] != "started" { |
| 96 | t.Fatalf("ack: %v", ack) |
| 97 | } |
| 98 | answer := h.waitNonTyping(t) |
| 99 | if txt, _ := answer["text"].(string); !strings.Contains(txt, "paul/lchat") { |
| 100 | t.Fatalf("answer: %v", answer) |
| 101 | } |
| 102 | if gh.hits != 1 || gh.auth != "Bearer ghp_research" { |
| 103 | t.Fatalf("provider saw hits=%d auth=%q", gh.hits, gh.auth) |
| 104 | } |
| 105 | |
| 106 | // The palette offered to the model: read ops only, and never a token. |
| 107 | if len(*reqs) != 2 { |
| 108 | t.Fatalf("model calls: %d", len(*reqs)) |
| 109 | } |
| 110 | first := string((*reqs)[0]) |
| 111 | if !strings.Contains(first, "github__repos_list") { |
| 112 | t.Error("read op missing from the palette") |
| 113 | } |
| 114 | if strings.Contains(first, "issue_create") { |
| 115 | t.Error("write op leaked into the research palette") |
| 116 | } |
| 117 | for i, raw := range *reqs { |
| 118 | if strings.Contains(string(raw), "ghp_research") { |
| 119 | t.Errorf("request %d carries the provider token — credentials must never reach the model", i) |
| 120 | } |
| 121 | } |
| 122 | // The op result is fenced as a tool_result on the second call. |
| 123 | var second struct { |
| 124 | Messages []struct { |
| 125 | Role string `json:"role"` |
| 126 | Content json.RawMessage `json:"content"` |
| 127 | } `json:"messages"` |
| 128 | } |
| 129 | if err := json.Unmarshal((*reqs)[1], &second); err != nil { |
| 130 | t.Fatal(err) |
| 131 | } |
| 132 | last := second.Messages[len(second.Messages)-1] |
| 133 | if last.Role != "user" || !strings.Contains(string(last.Content), "tool_result") || !strings.Contains(string(last.Content), "paul/lchat") { |
| 134 | t.Fatalf("tool_result turn: %s", last.Content) |
| 135 | } |
| 136 | } |
| 137 | |
| 138 | func TestResearchForeignGrantStaysSilent(t *testing.T) { |
| 139 | newFakeGitHub(t, `[]`) |
| 140 | h := newConnectHarness(t) |
| 141 | h.rt.ai = newAnthropicClient("test-key", nil) |
| 142 | other, err := e2ee.GenerateUserKeypair() |
| 143 | if err != nil { |
| 144 | t.Fatal(err) |
| 145 | } |
| 146 | sealed, err := e2ee.SealForPub(other.PublicKey, map[string]any{"service": "github", "token": "x", "access": "read"}) |
| 147 | if err != nil { |
| 148 | t.Fatal(err) |
| 149 | } |
| 150 | h.rt.handleResearch("t1", h.key, &researchPayload{ |
| 151 | Type: "connect.research", Goal: "anything", Grants: map[string]string{"github": sealed}, RID: "rr2", |
| 152 | }) |
| 153 | h.expectSilence(t, 400*time.Millisecond) |
| 154 | } |
| 155 | |
| 156 | func TestResearchExpiredGrantNeverReachesProvider(t *testing.T) { |
| 157 | gh := newFakeGitHub(t, `[]`) |
| 158 | reqs := scriptedClaude(t, |
| 159 | `{"content":[{"type":"tool_use","id":"tu1","name":"github__repos_list","input":{}}],"stop_reason":"tool_use"}`, |
| 160 | `{"content":[{"type":"text","text":"That look expired."}],"stop_reason":"end_turn"}`, |
| 161 | ) |
| 162 | h := newConnectHarness(t) |
| 163 | h.rt.ai = newAnthropicClient("test-key", nil) |
| 164 | grant := h.seal(t, map[string]any{"service": "github", "token": "ghp_x", |
| 165 | "exp": time.Now().Add(-time.Minute).UnixMilli(), "access": "read"}) |
| 166 | h.rt.handleResearch("t1", h.key, &researchPayload{ |
| 167 | Type: "connect.research", Goal: "repos?", Grants: map[string]string{"github": grant}, RID: "rr3", |
| 168 | }) |
| 169 | h.waitNonTyping(t) // ack |
| 170 | h.waitNonTyping(t) // answer |
| 171 | if gh.hits != 0 { |
| 172 | t.Fatalf("expired grant reached the provider (%d hits)", gh.hits) |
| 173 | } |
| 174 | if !strings.Contains(string((*reqs)[1]), "expired") { |
| 175 | t.Error("the model wasn't told the authorization expired") |
| 176 | } |
| 177 | } |
| 178 | |
| 179 | func TestResearchWithoutModelAnswersHonestly(t *testing.T) { |
| 180 | h := newConnectHarness(t) // ai == nil |
| 181 | grant := h.seal(t, map[string]any{"service": "github", "token": "ghp_x", "access": "read"}) |
| 182 | h.rt.handleResearch("t1", h.key, &researchPayload{ |
| 183 | Type: "connect.research", Goal: "repos?", Grants: map[string]string{"github": grant}, RID: "rr4", |
| 184 | }) |
| 185 | rep := h.waitNonTyping(t) |
| 186 | if rep["rid"] != "rr4" || rep["error"] == nil { |
| 187 | t.Fatalf("want an honest rid-matched error, got %v", rep) |
| 188 | } |
| 189 | } |
| 190 | |
| 191 | func TestActionsFromToolUseServiceResearch(t *testing.T) { |
| 192 | content := []anthBlock{ |
| 193 | {Type: "tool_use", Name: "service_research", Input: map[string]any{ |
| 194 | "services": []any{"github", "dnsimple", "github", "not-a-service"}, |
| 195 | "goal": "which domains point at GitHub Pages?", |
| 196 | }}, |
| 197 | } |
| 198 | got := actionsFromToolUse(content) |
| 199 | if len(got) != 1 { |
| 200 | t.Fatalf("actions: %+v", got) |
| 201 | } |
| 202 | a := got[0] |
| 203 | if a["type"] != "service_research" || a["goal"] != "which domains point at GitHub Pages?" { |
| 204 | t.Fatalf("chip: %+v", a) |
| 205 | } |
| 206 | services, _ := a["services"].([]string) |
| 207 | if len(services) != 2 || services[0] != "github" || services[1] != "dnsimple" { |
| 208 | t.Fatalf("services (deduped, registry-checked): %v", services) |
| 209 | } |
| 210 | if label, _ := a["label"].(string); !strings.Contains(label, "GitHub & DNSimple") || !strings.Contains(label, "read-only") { |
| 211 | t.Fatalf("label: %q", label) |
| 212 | } |
| 213 | |
| 214 | // A goal-less or unknown-services-only proposal never becomes a chip. |
| 215 | bad := actionsFromToolUse([]anthBlock{ |
| 216 | {Type: "tool_use", Name: "service_research", Input: map[string]any{"services": []any{"github"}, "goal": " "}}, |
| 217 | {Type: "tool_use", Name: "service_research", Input: map[string]any{"services": []any{"nope"}, "goal": "x"}}, |
| 218 | }) |
| 219 | if len(bad) != 0 { |
| 220 | t.Fatalf("bad proposals became chips: %+v", bad) |
| 221 | } |
| 222 | } |
| 223 | |
amadan