| 1 | package main |
| 2 | |
| 3 | // End-to-end test of Eleven's memory (internal/bot/memory.go + |
| 4 | // docs/eleven-memory.md) against the real router: a DM remember lands in the |
| 5 | // thread's "memory" state slot (encrypted under the DM key — decrypted here |
| 6 | // as the member's client would), the fact rides the NEXT call's system |
| 7 | // context with its id, the reply carries the sidecar-attached open_memory |
| 8 | // chip, and forget(all) wipes the slot. |
| 9 | |
| 10 | import ( |
| 11 | "bytes" |
| 12 | "context" |
| 13 | "encoding/json" |
| 14 | "io" |
| 15 | "net/http" |
| 16 | "net/http/httptest" |
| 17 | "strings" |
| 18 | "sync" |
| 19 | "testing" |
| 20 | "time" |
| 21 | |
| 22 | "lchat/internal/bot" |
| 23 | "lchat/internal/e2ee" |
| 24 | ) |
| 25 | |
| 26 | func TestMemoryRememberRecallForget(t *testing.T) { |
| 27 | var mu sync.Mutex |
| 28 | var bodies []string |
| 29 | calls := 0 |
| 30 | fakeAI := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 31 | b, _ := io.ReadAll(r.Body) |
| 32 | mu.Lock() |
| 33 | bodies = append(bodies, string(b)) |
| 34 | i := calls |
| 35 | calls++ |
| 36 | mu.Unlock() |
| 37 | responses := [][]map[string]any{ |
| 38 | // Turn 1: remember + a conversational ack. |
| 39 | {{"type": "text", "text": "Noted!"}, |
| 40 | {"type": "tool_use", "name": "remember", "input": map[string]any{"fact": "Runs the Eleven Messenger project"}}}, |
| 41 | // Turn 2: a plain reply (the test only wants this call's request). |
| 42 | {{"type": "text", "text": "You run the Eleven Messenger project."}}, |
| 43 | // Turn 3: the member asked for a wipe. |
| 44 | {{"type": "text", "text": "All forgotten."}, |
| 45 | {"type": "tool_use", "name": "forget", "input": map[string]any{"all": true}}}, |
| 46 | } |
| 47 | if i >= len(responses) { |
| 48 | i = len(responses) - 1 |
| 49 | } |
| 50 | json.NewEncoder(w).Encode(map[string]any{"content": responses[i], "stop_reason": "end_turn"}) |
| 51 | })) |
| 52 | defer fakeAI.Close() |
| 53 | t.Setenv("ANTHROPIC_BASE_URL", fakeAI.URL) |
| 54 | |
| 55 | ts, s := newTestServer(t) |
| 56 | owner := provisionNamed(t, ts.URL, "Owner", false) |
| 57 | eleven := provisionNamed(t, ts.URL, "Eleven", true) |
| 58 | ownerUID, _ := s.db.deviceTokenUser(owner.Token) |
| 59 | dmID, dmKey := createThreadKind(t, ts.URL, s, "dm", owner, eleven) |
| 60 | |
| 61 | dir := t.TempDir() |
| 62 | sessionPath := dir + "/eleven.session.json" |
| 63 | if err := eleven.Save(sessionPath); err != nil { |
| 64 | t.Fatal(err) |
| 65 | } |
| 66 | ready := make(chan struct{}) |
| 67 | ctx, cancel := context.WithCancel(context.Background()) |
| 68 | defer cancel() |
| 69 | eng := &bot.Engine{Base: ts.URL, Session: eleven, SessionPath: sessionPath, |
| 70 | Recipe: bot.DefaultRecipe(), AIKey: "test-key", |
| 71 | Ready: func() { close(ready) }, Logf: t.Logf} |
| 72 | go eng.Run(ctx) |
| 73 | awaitReady(t, ready) |
| 74 | |
| 75 | // memorySlot reads the slot the way the member's scratchpad would. |
| 76 | memorySlot := func() (map[string]any, bool) { |
| 77 | resp, err := http.Get(ts.URL + "/api/threads/" + dmID + "/state?token=" + owner.Token + "&slot=memory") |
| 78 | if err != nil { |
| 79 | t.Fatal(err) |
| 80 | } |
| 81 | defer resp.Body.Close() |
| 82 | var out struct { |
| 83 | Blob string `json:"blob"` |
| 84 | } |
| 85 | if json.NewDecoder(resp.Body).Decode(&out) != nil || out.Blob == "" { |
| 86 | return nil, false |
| 87 | } |
| 88 | var mem map[string]any |
| 89 | if err := e2ee.DecryptJSON(dmKey, out.Blob, &mem); err != nil { |
| 90 | t.Fatalf("memory slot undecryptable with the DM key: %v", err) |
| 91 | } |
| 92 | return mem, true |
| 93 | } |
| 94 | |
| 95 | // Turn 1 — remember: the fact lands in the slot, and the reply carries the |
| 96 | // sidecar-attached open_memory chip labeled as a note. |
| 97 | postAsMember(t, ts.URL, dmID, dmKey, owner, ownerUID, "Please remember that I run the Eleven Messenger project.") |
| 98 | waitFor(t, "memory slot written", func() bool { |
| 99 | mem, ok := memorySlot() |
| 100 | if !ok { |
| 101 | return false |
| 102 | } |
| 103 | facts, _ := mem["facts"].([]any) |
| 104 | return len(facts) == 1 |
| 105 | }) |
| 106 | mem, _ := memorySlot() |
| 107 | fact := mem["facts"].([]any)[0].(map[string]any) |
| 108 | if fact["text"] != "Runs the Eleven Messenger project" || fact["id"] == "" { |
| 109 | t.Fatalf("stored fact: %+v", fact) |
| 110 | } |
| 111 | waitFor(t, "the open_memory chip on the reply", func() bool { |
| 112 | for _, rp := range botReplies(t, s, dmID, dmKey, "Eleven") { |
| 113 | for _, a := range rp.Actions { |
| 114 | if a["type"] == "open_memory" { |
| 115 | return true |
| 116 | } |
| 117 | } |
| 118 | } |
| 119 | return false |
| 120 | }) |
| 121 | |
| 122 | // Turn 2 — recall: the next call's system context carries the fact + id. |
| 123 | postAsMember(t, ts.URL, dmID, dmKey, owner, ownerUID, "What do I do?") |
| 124 | waitFor(t, "the recall reply", func() bool { |
| 125 | return len(botReplies(t, s, dmID, dmKey, "Eleven")) >= 2 |
| 126 | }) |
| 127 | mu.Lock() |
| 128 | secondReq := bodies[1] |
| 129 | mu.Unlock() |
| 130 | var call struct { |
| 131 | System []struct { |
| 132 | Text string `json:"text"` |
| 133 | } `json:"system"` |
| 134 | } |
| 135 | if err := json.Unmarshal([]byte(secondReq), &call); err != nil { |
| 136 | t.Fatal(err) |
| 137 | } |
| 138 | var note string |
| 139 | for _, b := range call.System[1:] { // after the cached block |
| 140 | if strings.Contains(b.Text, "MEMORY") { |
| 141 | note = b.Text |
| 142 | } |
| 143 | } |
| 144 | if !strings.Contains(note, "Runs the Eleven Messenger project") || !strings.Contains(note, fact["id"].(string)) { |
| 145 | t.Fatalf("memory note missing the fact or its id:\n%s", note) |
| 146 | } |
| 147 | |
| 148 | // Turn 3 — forget everything: the slot empties (the blob stays, sealed, |
| 149 | // holding zero facts) and the wipe announces itself with the chip. |
| 150 | postAsMember(t, ts.URL, dmID, dmKey, owner, ownerUID, "Forget everything about me please.") |
| 151 | waitFor(t, "memory wiped", func() bool { |
| 152 | mem, ok := memorySlot() |
| 153 | if !ok { |
| 154 | return false |
| 155 | } |
| 156 | facts, _ := mem["facts"].([]any) |
| 157 | return len(facts) == 0 |
| 158 | }) |
| 159 | } |
| 160 | |
| 161 | // TestMemberScratchpadWriteWinsConflict pins the scratchpad's side of the |
| 162 | // shared slot: a member's own write (the panel saving at the version it read) |
| 163 | // bumps the version, and a stale second write 409s instead of clobbering. |
| 164 | func TestMemberScratchpadWriteWinsConflict(t *testing.T) { |
| 165 | ts, s := newTestServer(t) |
| 166 | owner := provisionNamed(t, ts.URL, "Owner", false) |
| 167 | eleven := provisionNamed(t, ts.URL, "Eleven", true) |
| 168 | _ = s |
| 169 | dmID, dmKey := createThreadKind(t, ts.URL, s, "dm", owner, eleven) |
| 170 | |
| 171 | put := func(version int64) int { |
| 172 | blob, err := e2ee.EncryptJSON(dmKey, map[string]any{"v": 1, "facts": []map[string]any{ |
| 173 | {"id": "m1", "text": "Prefers short answers", "ts": time.Now().UnixMilli()}, |
| 174 | }}) |
| 175 | if err != nil { |
| 176 | t.Fatal(err) |
| 177 | } |
| 178 | body, _ := json.Marshal(map[string]any{"token": owner.Token, "slot": "memory", "blob": blob, "version": version}) |
| 179 | req, _ := http.NewRequest(http.MethodPut, ts.URL+"/api/threads/"+dmID+"/state", bytes.NewReader(body)) |
| 180 | req.Header.Set("Content-Type", "application/json") |
| 181 | resp, err := http.DefaultClient.Do(req) |
| 182 | if err != nil { |
| 183 | t.Fatal(err) |
| 184 | } |
| 185 | resp.Body.Close() |
| 186 | return resp.StatusCode |
| 187 | } |
| 188 | if code := put(0); code != http.StatusOK { |
| 189 | t.Fatalf("first write: %d", code) |
| 190 | } |
| 191 | if code := put(0); code != http.StatusConflict { |
| 192 | t.Fatalf("stale write should 409, got %d", code) |
| 193 | } |
| 194 | } |
| 195 | |
amadan