返回 DeepSeek-Reasonix
session_export_test.go
根目录 / internal / serve / session_export_test.go
1 package serve
2
3 import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "fmt"
8 "io"
9 "net/http"
10 "net/http/httptest"
11 "net/url"
12 "path/filepath"
13 "strings"
14 "testing"
15
16 "reasonix/internal/config"
17 "reasonix/internal/control"
18 "reasonix/internal/provider"
19 "reasonix/internal/session"
20 )
21
22 func TestSessionExportHTTPFixedCompleteSnapshot(t *testing.T) {
23 service, err := session.NewService("serve", session.NewFilesystemPersistence(filepath.Join(t.TempDir(), "sessions")))
24 if err != nil {
25 t.Fatal(err)
26 }
27 t.Cleanup(func() { _ = service.CloseAll(context.Background()) })
28 runtime, err := service.Create(t.Context(), session.CreateOptions{SessionID: "export"})
29 if err != nil {
30 t.Fatal(err)
31 }
32 appendMessage := func(id string) {
33 t.Helper()
34 body, _ := json.Marshal(map[string]any{"message": provider.Message{ID: id, Role: provider.RoleUser, Content: id, Origin: provider.MessageOrigin("user")}})
35 if _, err := runtime.Session().AppendBatch(t.Context(), id, []session.Event{{Kind: "message/complete", Payload: body}}); err != nil {
36 t.Fatal(err)
37 }
38 }
39 for i := range 110 {
40 appendMessage(fmt.Sprintf("question-%03d", i))
41 }
42 bc := NewBroadcaster()
43 ctrl := control.New(control.Options{SessionService: service, SessionRuntime: runtime, ExclusiveSession: true, Sink: bc})
44 defer ctrl.Close()
45 server := httptest.NewServer(New(ctrl, bc, config.ServeConfig{}).Handler())
46 defer server.Close()
47 response, err := http.Get(server.URL + "/session-export/snapshot")
48 if err != nil {
49 t.Fatal(err)
50 }
51 var snapshot session.ExportSnapshot
52 err = json.NewDecoder(response.Body).Decode(&snapshot)
53 response.Body.Close()
54 if err != nil || response.StatusCode != http.StatusOK {
55 t.Fatalf("snapshot: %d %v", response.StatusCode, err)
56 }
57 appendMessage("AFTER-SNAPSHOT")
58 request, _ := json.Marshal(map[string]any{"snapshot": snapshot, "format": "json"})
59 response, err = http.Post(server.URL+"/session-export/document", "application/json", bytes.NewReader(request))
60 if err != nil {
61 t.Fatal(err)
62 }
63 body, err := io.ReadAll(response.Body)
64 response.Body.Close()
65 if err != nil {
66 t.Fatal(err)
67 }
68 if response.StatusCode != http.StatusOK || response.Header.Get("X-Reasonix-Export-Records") != "110" || !json.Valid(body) {
69 t.Fatalf("response: %d %s", response.StatusCode, body)
70 }
71 if !bytes.Contains(body, []byte("question-000")) || !bytes.Contains(body, []byte("question-109")) || bytes.Contains(body, []byte("AFTER-SNAPSHOT")) {
72 t.Fatal("export was truncated or changed its snapshot")
73 }
74 snapshot.Ref.SessionID = "other"
75 request, _ = json.Marshal(map[string]any{"snapshot": snapshot, "format": "json"})
76 response, err = http.Post(server.URL+"/session-export/document", "application/json", bytes.NewReader(request))
77 if err != nil {
78 t.Fatal(err)
79 }
80 response.Body.Close()
81 if response.StatusCode != http.StatusConflict {
82 t.Fatalf("accepted mismatched identity: %d", response.StatusCode)
83 }
84 malicious, _ := json.Marshal(map[string]any{"snapshot": snapshot, "format": "../../manifest.json"})
85 response, err = http.Post(server.URL+"/session-export/document", "application/json", bytes.NewReader(malicious))
86 if err != nil {
87 t.Fatal(err)
88 }
89 response.Body.Close()
90 if response.StatusCode != http.StatusBadRequest {
91 t.Fatalf("accepted path-like export format: %d", response.StatusCode)
92 }
93 response, err = http.Post(server.URL+"/session-export/diagnostic", "application/json", strings.NewReader(`{}`))
94 if err != nil {
95 t.Fatal(err)
96 }
97 body, err = io.ReadAll(response.Body)
98 response.Body.Close()
99 if err != nil || response.StatusCode != http.StatusOK || !json.Valid(body) {
100 t.Fatalf("diagnostic: %d %v", response.StatusCode, err)
101 }
102 }
103
104 func appendSessionExportTestMessage(t *testing.T, runtime *session.Runtime, id string) {
105 t.Helper()
106 body, _ := json.Marshal(map[string]any{"message": provider.Message{ID: id, Role: provider.RoleUser, Content: id, Origin: provider.MessageOrigin("user")}})
107 if _, err := runtime.Session().AppendBatch(t.Context(), id, []session.Event{{Kind: "message/complete", Payload: body}}); err != nil {
108 t.Fatal(err)
109 }
110 }
111
112 func captureSessionExportTestSnapshot(t *testing.T, serverURL string, ref session.SessionRef, diagnostic bool) session.ExportSnapshot {
113 t.Helper()
114 endpoint := serverURL + "/session-export/snapshot?sessionId=" + ref.SessionID
115 if diagnostic {
116 endpoint += "&diagnostic=1"
117 }
118 req, err := http.NewRequest(http.MethodGet, endpoint, nil)
119 if err != nil {
120 t.Fatal(err)
121 }
122 req.Header.Set(expectedSessionIDHeader, ref.SessionID)
123 resp, err := http.DefaultClient.Do(req)
124 if err != nil {
125 t.Fatal(err)
126 }
127 defer resp.Body.Close()
128 if resp.StatusCode != http.StatusOK {
129 body, _ := io.ReadAll(resp.Body)
130 t.Fatalf("snapshot status = %d: %s", resp.StatusCode, body)
131 }
132 var snapshot session.ExportSnapshot
133 if err = json.NewDecoder(resp.Body).Decode(&snapshot); err != nil {
134 t.Fatal(err)
135 }
136 return snapshot
137 }
138
139 func postFixedSessionExport(t *testing.T, serverURL, path, queryID, headerID string, body any) (*http.Response, []byte) {
140 t.Helper()
141 encoded, err := json.Marshal(body)
142 if err != nil {
143 t.Fatal(err)
144 }
145 req, err := http.NewRequest(http.MethodPost, serverURL+path+"?sessionId="+url.QueryEscape(queryID), bytes.NewReader(encoded))
146 if err != nil {
147 t.Fatal(err)
148 }
149 req.Header.Set("Content-Type", "application/json")
150 if headerID != "" {
151 req.Header.Set(expectedSessionIDHeader, headerID)
152 }
153 resp, err := http.DefaultClient.Do(req)
154 if err != nil {
155 t.Fatal(err)
156 }
157 data, err := io.ReadAll(resp.Body)
158 resp.Body.Close()
159 if err != nil {
160 t.Fatal(err)
161 }
162 return resp, data
163 }
164
165 func TestSessionExportRemainsPinnedAfterForegroundSwitch(t *testing.T) {
166 srv, ctrl, service, source := newExclusiveSessionServe(t)
167 runtime, ok := service.Runtime(source)
168 if !ok {
169 t.Fatal("source runtime is unavailable")
170 }
171 appendSessionExportTestMessage(t, runtime, "SOURCE-A")
172 ts := httptest.NewServer(srv.Handler())
173 defer ts.Close()
174 snapshot := captureSessionExportTestSnapshot(t, ts.URL, source, true)
175
176 target, err := ctrl.BindFreshSession(t.Context(), "target-b")
177 if err != nil {
178 t.Fatal(err)
179 }
180 targetRuntime, ok := service.Runtime(target)
181 if !ok {
182 t.Fatal("target runtime is unavailable")
183 }
184 appendSessionExportTestMessage(t, targetRuntime, "TARGET-B")
185
186 resp, body := postFixedSessionExport(t, ts.URL, "/session-export/document", source.SessionID, source.SessionID, map[string]any{"snapshot": snapshot, "format": "json"})
187 if resp.StatusCode != http.StatusOK || !bytes.Contains(body, []byte("SOURCE-A")) || bytes.Contains(body, []byte("TARGET-B")) {
188 t.Fatalf("fixed document status=%d body=%s", resp.StatusCode, body)
189 }
190 resp, body = postFixedSessionExport(t, ts.URL, "/session-export/validate", source.SessionID, source.SessionID, snapshot)
191 if resp.StatusCode != http.StatusNoContent {
192 t.Fatalf("fixed validation status=%d body=%s", resp.StatusCode, body)
193 }
194 resp, body = postFixedSessionExport(t, ts.URL, "/session-export/diagnostic", source.SessionID, source.SessionID, map[string]any{"exportSnapshot": snapshot})
195 if resp.StatusCode != http.StatusOK || !json.Valid(body) || !bytes.Contains(body, []byte("SOURCE-A")) || bytes.Contains(body, []byte("TARGET-B")) || !bytes.Contains(body, []byte("cold session")) {
196 t.Fatalf("fixed diagnostic status=%d body=%s", resp.StatusCode, body)
197 }
198
199 resp, body = postFixedSessionExport(t, ts.URL, "/session-export/document", target.SessionID, source.SessionID, map[string]any{"snapshot": snapshot, "format": "json"})
200 if resp.StatusCode != http.StatusConflict {
201 t.Fatalf("conflicting query/header/body status=%d body=%s", resp.StatusCode, body)
202 }
203 if err = service.Close(t.Context(), source); err != nil {
204 t.Fatal(err)
205 }
206 if err = service.Delete(t.Context(), source); err != nil {
207 t.Fatal(err)
208 }
209 resp, body = postFixedSessionExport(t, ts.URL, "/session-export/document", source.SessionID, source.SessionID, map[string]any{"snapshot": snapshot, "format": "json"})
210 if resp.StatusCode != http.StatusConflict {
211 t.Fatalf("deleted source export status=%d body=%s", resp.StatusCode, body)
212 }
213 }
214
215 func TestSessionExportRejectsPathLikeTargetIdentities(t *testing.T) {
216 srv, _, _, source := newExclusiveSessionServe(t)
217 ts := httptest.NewServer(srv.Handler())
218 defer ts.Close()
219 snapshot := captureSessionExportTestSnapshot(t, ts.URL, source, false)
220
221 for _, target := range []string{
222 "../" + source.SessionID,
223 source.SessionID + "/child",
224 `C:\\outside\\` + source.SessionID,
225 ".query-cache",
226 } {
227 t.Run(target, func(t *testing.T) {
228 resp, body := postFixedSessionExport(t, ts.URL, "/session-export/document", target, "", map[string]any{
229 "snapshot": snapshot,
230 "format": "json",
231 })
232 if resp.StatusCode != http.StatusBadRequest {
233 t.Fatalf("path-like query target status=%d body=%s", resp.StatusCode, body)
234 }
235 })
236 }
237
238 invalidSnapshot := snapshot
239 invalidSnapshot.Ref.SessionID = "../" + source.SessionID
240 resp, body := postFixedSessionExport(t, ts.URL, "/session-export/validate", "", "", invalidSnapshot)
241 if resp.StatusCode != http.StatusBadRequest {
242 t.Fatalf("path-like snapshot target status=%d body=%s", resp.StatusCode, body)
243 }
244 }
245
246 func TestSessionExportRemainsPinnedAfterTakeover(t *testing.T) {
247 srv, ctrl, service, source := newExclusiveSessionServe(t)
248 runtime, ok := service.Runtime(source)
249 if !ok {
250 t.Fatal("source runtime is unavailable")
251 }
252 appendSessionExportTestMessage(t, runtime, "BEFORE-TAKEOVER")
253 ts := httptest.NewServer(srv.Handler())
254 defer ts.Close()
255 snapshot := captureSessionExportTestSnapshot(t, ts.URL, source, false)
256
257 resp, raw := serveBody(t, http.MethodPost, ts.URL+"/handoff", `{"sessionPath":"session-id:`+source.SessionID+`","targetWriterId":"test-taker","force":true,"mode":"wait","timeoutMs":2000}`)
258 if resp.StatusCode != http.StatusOK {
259 t.Fatalf("handoff status=%d body=%s", resp.StatusCode, raw)
260 }
261 bodyResp, body := postFixedSessionExport(t, ts.URL, "/session-export/document", source.SessionID, source.SessionID, map[string]any{"snapshot": snapshot, "format": "json"})
262 if bodyResp.StatusCode != http.StatusOK || !bytes.Contains(body, []byte("BEFORE-TAKEOVER")) {
263 t.Fatalf("takeover export status=%d body=%s", bodyResp.StatusCode, body)
264 }
265 retireExclusiveForeground(t, ctrl, service)
266 }
267
267 lines GO