返回 DeepSeek-Reasonix
pinned_files_test.go
根目录 / desktop / pinned_files_test.go
1 package main
2
3 import (
4 "context"
5 "os"
6 "path/filepath"
7 "reflect"
8 "strings"
9 "testing"
10
11 "reasonix/internal/agent"
12 "reasonix/internal/control"
13 "reasonix/internal/event"
14 "reasonix/internal/provider"
15 "reasonix/internal/tool"
16 )
17
18 func TestTabPinFile(t *testing.T) {
19 dir := t.TempDir()
20 filePath := filepath.Join(dir, "api_spec.md")
21 content := "# API Specification\n\n- Endpoint: /v1/chat\n- Method: POST\n"
22 if err := os.WriteFile(filePath, []byte(content), 0o644); err != nil {
23 t.Fatal(err)
24 }
25
26 tab := &WorkspaceTab{
27 ID: "tab-1",
28 WorkspaceRoot: dir,
29 }
30
31 info, err := tab.PinFile("api_spec.md")
32 if err != nil {
33 t.Fatalf("unexpected error pinning file: %v", err)
34 }
35 if info.Path != "api_spec.md" {
36 t.Fatalf("expected path api_spec.md, got %q", info.Path)
37 }
38 if info.SizeBytes != int64(len(content)) {
39 t.Fatalf("expected size %d, got %d", len(content), info.SizeBytes)
40 }
41 if info.TokenEstimate <= 0 {
42 t.Fatalf("expected positive token estimate, got %d", info.TokenEstimate)
43 }
44
45 // Test idempotency
46 info2, err := tab.PinFile("api_spec.md")
47 if err != nil {
48 t.Fatalf("unexpected error on duplicate pin: %v", err)
49 }
50 if info2.Path != "api_spec.md" {
51 t.Fatalf("expected path api_spec.md, got %q", info2.Path)
52 }
53 if len(tab.GetPinnedFiles()) != 1 {
54 t.Fatalf("expected exactly 1 pinned file, got %d", len(tab.GetPinnedFiles()))
55 }
56
57 // Test non-existent file
58 if _, err := tab.PinFile("missing.txt"); err == nil {
59 t.Fatal("expected error for non-existent file, got nil")
60 }
61
62 // Test directory pin rejection
63 subDir := filepath.Join(dir, "subdir")
64 if err := os.Mkdir(subDir, 0o755); err != nil {
65 t.Fatal(err)
66 }
67 if _, err := tab.PinFile("subdir"); err == nil {
68 t.Fatal("expected error for directory pin, got nil")
69 }
70
71 // Test path traversal rejection
72 if _, err := tab.PinFile("../outside.txt"); err == nil {
73 t.Fatal("expected error for path traversal, got nil")
74 }
75 }
76
77 func TestTabPinSymlinkOutsideWorkspaceForbidden(t *testing.T) {
78 wsDir := t.TempDir()
79 outsideDir := t.TempDir()
80 outsideFile := filepath.Join(outsideDir, "secret.key")
81 if err := os.WriteFile(outsideFile, []byte("SUPER_SECRET"), 0o600); err != nil {
82 t.Fatal(err)
83 }
84
85 symlinkPath := filepath.Join(wsDir, "escape_link.txt")
86 if err := os.Symlink(outsideFile, symlinkPath); err != nil {
87 t.Skipf("symlink creation not supported or permitted: %v", err)
88 }
89
90 tab := &WorkspaceTab{
91 ID: "tab-symlink",
92 WorkspaceRoot: wsDir,
93 }
94
95 _, err := tab.PinFile("escape_link.txt")
96 if err == nil {
97 t.Fatal("expected error when pinning symlink pointing outside workspace, got nil")
98 }
99 if !strings.Contains(err.Error(), "outside workspace") {
100 t.Fatalf("expected workspace escape error message, got: %v", err)
101 }
102 }
103
104 func TestTabPinSymlinkInsideWorkspaceAllowed(t *testing.T) {
105 wsDir := t.TempDir()
106 realFile := filepath.Join(wsDir, "real_config.json")
107 content := `{"allowed": true}`
108 if err := os.WriteFile(realFile, []byte(content), 0o644); err != nil {
109 t.Fatal(err)
110 }
111
112 symlinkPath := filepath.Join(wsDir, "config_link.json")
113 if err := os.Symlink(realFile, symlinkPath); err != nil {
114 t.Skipf("symlink creation not supported or permitted: %v", err)
115 }
116
117 tab := &WorkspaceTab{
118 ID: "tab-symlink-ok",
119 WorkspaceRoot: wsDir,
120 }
121
122 info, err := tab.PinFile("config_link.json")
123 if err != nil {
124 t.Fatalf("expected pinning valid inside-workspace symlink to succeed, got: %v", err)
125 }
126 if info.Path != "config_link.json" {
127 t.Fatalf("expected path config_link.json, got %q", info.Path)
128 }
129 if info.SizeBytes != int64(len(content)) {
130 t.Fatalf("expected size %d, got %d", len(content), info.SizeBytes)
131 }
132 }
133
134 func TestTabPinFileSizeLimit(t *testing.T) {
135 dir := t.TempDir()
136 bigFile := filepath.Join(dir, "big.dat")
137 data := make([]byte, maxPinnedFileSize+1)
138 if err := os.WriteFile(bigFile, data, 0o644); err != nil {
139 t.Fatal(err)
140 }
141
142 tab := &WorkspaceTab{
143 ID: "tab-1",
144 WorkspaceRoot: dir,
145 }
146
147 if _, err := tab.PinFile("big.dat"); err == nil {
148 t.Fatal("expected error for file exceeding size limit, got nil")
149 }
150 }
151
152 func TestTabUnpinFile(t *testing.T) {
153 dir := t.TempDir()
154 f1 := filepath.Join(dir, "f1.txt")
155 f2 := filepath.Join(dir, "f2.txt")
156 if err := os.WriteFile(f1, []byte("f1"), 0o644); err != nil {
157 t.Fatal(err)
158 }
159 if err := os.WriteFile(f2, []byte("f2"), 0o644); err != nil {
160 t.Fatal(err)
161 }
162
163 tab := &WorkspaceTab{
164 ID: "tab-1",
165 WorkspaceRoot: dir,
166 }
167
168 if _, err := tab.PinFile("f1.txt"); err != nil {
169 t.Fatal(err)
170 }
171 if _, err := tab.PinFile("f2.txt"); err != nil {
172 t.Fatal(err)
173 }
174 if len(tab.GetPinnedFiles()) != 2 {
175 t.Fatalf("expected 2 pinned files, got %d", len(tab.GetPinnedFiles()))
176 }
177
178 if err := tab.UnpinFile("f1.txt"); err != nil {
179 t.Fatal(err)
180 }
181 pinned := tab.GetPinnedFiles()
182 if len(pinned) != 1 || pinned[0] != "f2.txt" {
183 t.Fatalf("expected [f2.txt], got %v", pinned)
184 }
185
186 // Unpin non-existent should succeed without error
187 if err := tab.UnpinFile("nonexistent.txt"); err != nil {
188 t.Fatalf("unexpected error unpinning non-existent file: %v", err)
189 }
190 }
191
192 func TestTabPinnedContextSnapshot(t *testing.T) {
193 dir := t.TempDir()
194 f1 := filepath.Join(dir, "schema.sql")
195 sqlContent := "CREATE TABLE users (id INT PRIMARY KEY, name TEXT);"
196 if err := os.WriteFile(f1, []byte(sqlContent), 0o644); err != nil {
197 t.Fatal(err)
198 }
199
200 tab := &WorkspaceTab{
201 ID: "tab-1",
202 WorkspaceRoot: dir,
203 }
204
205 if _, err := tab.PinFile("schema.sql"); err != nil {
206 t.Fatal(err)
207 }
208
209 build := buildPinnedContext(dir, tab.GetPinnedFiles())
210 if len(build.Snapshot.Files) != 1 || build.Snapshot.Files[0].Path != "schema.sql" || build.Snapshot.Files[0].Content != sqlContent {
211 t.Fatalf("snapshot = %+v", build.Snapshot)
212 }
213
214 infoList := tab.GetPinnedFilesInfo()
215 if len(infoList) != 1 {
216 t.Fatalf("expected 1 info item, got %d", len(infoList))
217 }
218 if infoList[0].Path != "schema.sql" || infoList[0].SizeBytes != int64(len(sqlContent)) {
219 t.Fatalf("unexpected info item: %+v", infoList[0])
220 }
221 }
222
223 func TestPinnedContextEndToEndProviderRequest(t *testing.T) {
224 dir := t.TempDir()
225 docPath := filepath.Join(dir, "architecture.md")
226 docContent := "# Architecture\nUse clean layered architecture without global singletons.\n"
227 if err := os.WriteFile(docPath, []byte(docContent), 0o644); err != nil {
228 t.Fatal(err)
229 }
230
231 tab := &WorkspaceTab{
232 ID: "tab-e2e",
233 WorkspaceRoot: dir,
234 }
235
236 if _, err := tab.PinFile("architecture.md"); err != nil {
237 t.Fatalf("pin file: %v", err)
238 }
239
240 baseSystem := "You are a helpful coding assistant."
241 sessionPath := filepath.Join(dir, "session.jsonl")
242 if err := savePinnedContextState(sessionPath, tab.GetPinnedFiles()); err != nil {
243 t.Fatal(err)
244 }
245
246 prov := &capturingProvider{}
247 exec := agent.New(prov, tool.NewRegistry(), agent.NewSession(baseSystem), agent.Options{}, event.Discard)
248 ctrl := newFixtureController(t, control.Options{
249 Runner: exec,
250 Executor: exec,
251 SystemPrompt: baseSystem,
252 PinnedContextLoader: pinnedContextLoader(dir),
253 SessionDir: dir,
254 SessionPath: sessionPath,
255 Label: "test-e2e",
256 Sink: event.Discard,
257 })
258 tab.Ctrl = ctrl
259
260 // 1. Execute first user turn and assert provider receives a host revision.
261 if err := ctrl.RunTurn(context.Background(), "How should we design the service?"); err != nil {
262 t.Fatalf("RunTurn: %v", err)
263 }
264
265 reqMsgs := prov.lastRequestMessages(t)
266 if len(reqMsgs) == 0 {
267 t.Fatal("expected provider to receive messages, got none")
268 }
269 sysMsg := reqMsgs[0]
270 if sysMsg.Role != provider.RoleSystem {
271 t.Fatalf("expected first message to be system role, got %s", sysMsg.Role)
272 }
273 if sysMsg.Content != baseSystem {
274 t.Fatalf("pinned context changed system prompt: %q", sysMsg.Content)
275 }
276 if len(reqMsgs) < 2 || !strings.HasPrefix(reqMsgs[1].Content, "<pinned_context_revision") ||
277 !strings.Contains(reqMsgs[1].Content, `path="architecture.md"`) || !strings.Contains(reqMsgs[1].Content, "Use clean layered architecture") {
278 t.Fatalf("missing pinned revision: %+v", reqMsgs)
279 }
280
281 // 2. Unpin updates the sidecar; the next admitted turn appends a tombstone.
282 if err := tab.UnpinFile("architecture.md"); err != nil {
283 t.Fatalf("UnpinFile: %v", err)
284 }
285 if err := savePinnedContextState(sessionPath, tab.GetPinnedFiles()); err != nil {
286 t.Fatal(err)
287 }
288
289 // 3. Execute second turn and assert prior request bytes remain its prefix.
290 if err := ctrl.RunTurn(context.Background(), "Next question"); err != nil {
291 t.Fatalf("RunTurn 2: %v", err)
292 }
293
294 reqMsgs2 := prov.lastRequestMessages(t)
295 if reqMsgs2[0].Content != baseSystem {
296 t.Fatalf("system prompt drifted after unpin: %q", reqMsgs2[0].Content)
297 }
298 if len(reqMsgs2) <= len(reqMsgs) {
299 t.Fatalf("second request did not append: %+v", reqMsgs2)
300 }
301 if !reflect.DeepEqual(reqMsgs2[:len(reqMsgs)], reqMsgs) {
302 t.Fatal("unpin changed bytes from the previous provider request")
303 }
304 revocation := reqMsgs2[len(reqMsgs2)-2].Content
305 if !strings.HasPrefix(revocation, "<pinned_context_revision") ||
306 (!strings.Contains(revocation, `<remove path="architecture.md"></remove>`) &&
307 !strings.Contains(revocation, `kind="checkpoint"`)) {
308 t.Fatalf("missing unpin tombstone: %+v", reqMsgs2)
309 }
310 rows := historyMessagesWithPlannerDisplays(ctrl.History(), func(value string) string { return value }, nil, nil)
311 users := 0
312 for _, row := range rows {
313 if strings.Contains(row.Content, "<pinned_context_revision") {
314 t.Fatalf("pinned revision leaked into desktop history: %+v", rows)
315 }
316 if row.Role == string(provider.RoleUser) {
317 users++
318 }
319 }
320 if users != 2 {
321 t.Fatalf("desktop history user rows = %d, want 2: %+v", users, rows)
322 }
323 }
324
325 func TestLegacyTabPinnedFilesMigrateToSessionSidecar(t *testing.T) {
326 sessionPath := filepath.Join(t.TempDir(), "session.jsonl")
327 legacy := []string{"README.md", "docs/api.md"}
328
329 state, err := loadOrMigratePinnedContextState(sessionPath, legacy)
330 if err != nil {
331 t.Fatalf("migrate legacy pins: %v", err)
332 }
333 if strings.Join(state.Files, ",") != strings.Join(legacy, ",") {
334 t.Fatalf("migrated files = %v, want %v", state.Files, legacy)
335 }
336
337 loaded, err := loadPinnedContextState(sessionPath)
338 if err != nil {
339 t.Fatalf("load migrated sidecar: %v", err)
340 }
341 if strings.Join(loaded.Files, ",") != strings.Join(legacy, ",") {
342 t.Fatalf("sidecar files = %v, want %v", loaded.Files, legacy)
343 }
344 }
345
345 lines GO