返回 DeepSeek-Reasonix
control_test.go
根目录 / internal / taskmonitor / control_test.go
1 package taskmonitor
2
3 import (
4 "context"
5 "os"
6 "path/filepath"
7 "strings"
8 "sync"
9 "testing"
10 "time"
11 )
12
13 func TestControlService_StopTask(t *testing.T) {
14 s := NewInMemoryStore()
15 cs := NewControlService(s)
16 ctx := context.Background()
17
18 mustUpsertControl(t, s, "/p", TaskSnapshot{
19 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
20 State: TaskStateRunning, Version: 1,
21 CreatedAt: time.Now(), UpdatedAt: time.Now(),
22 })
23
24 res, err := cs.StopTaskWithKiller(ctx, "/p", "t1", 1, "user request", "idem-1", &mockKiller{fn: func(string, string) bool { return true }})
25 if err != nil {
26 t.Fatalf("StopTask: %v", err)
27 }
28 if !res.Accepted {
29 t.Errorf("expected accepted, got %+v", res)
30 }
31 if res.State != TaskStateCancelled {
32 t.Errorf("expected cancelled, got %q", res.State)
33 }
34 if res.Version != 2 {
35 t.Errorf("expected version 2, got %d", res.Version)
36 }
37 }
38
39 func TestControlService_StopRoutesNamespacedTaskToRuntimeJobID(t *testing.T) {
40 s := NewInMemoryStore()
41 cs := NewControlService(s)
42 now := time.Now()
43 mustUpsertControl(t, s, "/p", TaskSnapshot{
44 SchemaVersion: 1, TaskID: "session-1--task-1", JobID: "task-1", SessionID: "session-1",
45 State: TaskStateRunning, RuntimeState: RuntimeStateAlive, Version: 1,
46 CreatedAt: now, UpdatedAt: now,
47 })
48
49 killer := &mockKiller{fn: func(sessionID, jobID string) bool {
50 return sessionID == "session-1" && jobID == "task-1"
51 }}
52 res, err := cs.StopTaskWithKiller(context.Background(), "/p", "session-1--task-1", 1, "", "", killer)
53 if err != nil || !res.Accepted {
54 t.Fatalf("namespaced stop: result=%+v err=%v", res, err)
55 }
56 }
57
58 func TestRuntimeJobIDSupportsSnapshotsBeforeJobIDField(t *testing.T) {
59 longSession := strings.Repeat("s", maxFieldLen)
60 for _, tc := range []struct {
61 name string
62 snap TaskSnapshot
63 want string
64 }{
65 {name: "legacy raw id", snap: TaskSnapshot{TaskID: "task-1", SessionID: "session-1"}, want: "task-1"},
66 {name: "namespaced id", snap: TaskSnapshot{TaskID: "session-1--task-1", SessionID: "session-1"}, want: "task-1"},
67 {name: "hashed namespace", snap: TaskSnapshot{TaskID: monitorTaskID(longSession, "task-1"), SessionID: longSession}, want: "task-1"},
68 {name: "explicit id", snap: TaskSnapshot{TaskID: "monitor-id", JobID: "bash-2", SessionID: "session-1"}, want: "bash-2"},
69 } {
70 t.Run(tc.name, func(t *testing.T) {
71 if got := runtimeJobID(&tc.snap); got != tc.want {
72 t.Fatalf("runtimeJobID() = %q, want %q", got, tc.want)
73 }
74 })
75 }
76 }
77
78 func TestControlService_StopRequiresRuntimeOwner(t *testing.T) {
79 s := NewInMemoryStore()
80 cs := NewControlService(s)
81 mustUpsertControl(t, s, "/p", TaskSnapshot{
82 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
83 State: TaskStateRunning, Version: 1,
84 CreatedAt: time.Now(), UpdatedAt: time.Now(),
85 })
86
87 res, err := cs.StopTask(context.Background(), "/p", "t1", 1, "", "")
88 if err != nil || res.Accepted || res.Error == nil || res.Error.Code != ErrTaskRuntimeUnavailable {
89 t.Fatalf("expected unavailable runtime, got result=%+v err=%v", res, err)
90 }
91 snap, _ := s.GetTask(context.Background(), "/p", "t1")
92 if snap.State != TaskStateRunning || snap.Version != 1 {
93 t.Fatalf("failed stop mutated task: %+v", snap)
94 }
95 }
96
97 func TestControlService_CancelRejectsUnreachableRuntime(t *testing.T) {
98 s := NewInMemoryStore()
99 cs := NewControlService(s)
100 mustUpsertControl(t, s, "/p", TaskSnapshot{
101 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
102 State: TaskStateRunning, Version: 1,
103 CreatedAt: time.Now(), UpdatedAt: time.Now(),
104 })
105
106 killer := &mockKiller{fn: func(string, string) bool { return false }}
107 res, err := cs.CancelTaskWithKiller(context.Background(), "/p", "t1", 1, "", "", killer)
108 if err != nil || res.Accepted || res.Error == nil || res.Error.Code != ErrTaskRuntimeUnavailable {
109 t.Fatalf("expected rejected runtime control, got result=%+v err=%v", res, err)
110 }
111 snap, _ := s.GetTask(context.Background(), "/p", "t1")
112 if snap.State != TaskStateRunning || snap.Version != 1 {
113 t.Fatalf("failed cancel mutated task: %+v", snap)
114 }
115 }
116
117 func TestControlService_VersionConflict(t *testing.T) {
118 s := NewInMemoryStore()
119 cs := NewControlService(s)
120
121 mustUpsertControl(t, s, "/p", TaskSnapshot{
122 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
123 State: TaskStateRunning, Version: 3,
124 CreatedAt: time.Now(), UpdatedAt: time.Now(),
125 })
126
127 res, _ := cs.StopTask(context.Background(), "/p", "t1", 1, "", "")
128 if res.Accepted || res.Error == nil || res.Error.Code != ErrTaskVersionConflict {
129 t.Errorf("expected version conflict, got %+v", res)
130 }
131 }
132
133 func TestControlService_NotFound(t *testing.T) {
134 cs := NewControlService(NewInMemoryStore())
135 res, _ := cs.StopTask(context.Background(), "/p", "ghost", 1, "", "")
136 if res.Error == nil || res.Error.Code != ErrTaskNotFound {
137 t.Errorf("expected not_found, got %+v", res.Error)
138 }
139 }
140
141 func TestControlService_TerminalGuard(t *testing.T) {
142 s := NewInMemoryStore()
143 cs := NewControlService(s)
144 mustUpsertControl(t, s, "/p", TaskSnapshot{
145 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
146 State: TaskStateSucceeded, Version: 1,
147 CreatedAt: time.Now(), UpdatedAt: time.Now(),
148 })
149 res, _ := cs.StopTask(context.Background(), "/p", "t1", 1, "", "")
150 if res.Error == nil || res.Error.Code != ErrTaskAlreadyTerminal {
151 t.Errorf("expected terminal guard, got %+v", res.Error)
152 }
153 }
154
155 func TestControlService_RequeueFailedTaskDoesNotClaimLiveRuntime(t *testing.T) {
156 s := NewInMemoryStore()
157 cs := NewControlService(s)
158 mustUpsertControl(t, s, "/p", TaskSnapshot{
159 SchemaVersion: 1, TaskID: "failed", SessionID: "s1",
160 State: TaskStateFailed, RuntimeState: RuntimeStateExited, Version: 3,
161 CreatedAt: time.Now(), UpdatedAt: time.Now(),
162 })
163 res, err := cs.RequeueTask(context.Background(), "/p", "failed", 3, "requeue-1")
164 if err != nil || !res.Accepted || res.State != TaskStateQueued || res.RuntimeState != RuntimeStateExited || res.Version != 4 {
165 t.Fatalf("expected failed task to be requeued without a live runtime, got result=%+v err=%v", res, err)
166 }
167 snap, _ := s.GetTask(context.Background(), "/p", "failed")
168 if snap.RuntimeState != RuntimeStateExited {
169 t.Fatalf("requeue changed runtime state to %q, want exited", snap.RuntimeState)
170 }
171 }
172
173 func TestControlService_RequeueRejectsLiveRuntime(t *testing.T) {
174 s := NewInMemoryStore()
175 cs := NewControlService(s)
176 mustUpsertControl(t, s, "/p", TaskSnapshot{
177 SchemaVersion: 1, TaskID: "failed", SessionID: "s1",
178 State: TaskStateFailed, RuntimeState: RuntimeStateAlive, Version: 3,
179 CreatedAt: time.Now(), UpdatedAt: time.Now(),
180 })
181 res, err := cs.RequeueTask(context.Background(), "/p", "failed", 3, "")
182 if err != nil || res.Error == nil || res.Error.Code != ErrTaskInProgress {
183 t.Fatalf("expected live-runtime guard, got result=%+v err=%v", res, err)
184 }
185 }
186
187 func TestControlService_RequeueAllowsExpiredRuntimeLease(t *testing.T) {
188 now := time.Now().UTC()
189 s := NewInMemoryStore()
190 cs := NewControlService(s)
191 mustUpsertControl(t, s, "/p", TaskSnapshot{
192 SchemaVersion: 1, TaskID: "failed", SessionID: "s1",
193 State: TaskStateFailed, RuntimeState: RuntimeStateAlive, RuntimeLeaseUntil: now.Add(-time.Minute), Version: 3,
194 CreatedAt: now.Add(-time.Hour), UpdatedAt: now.Add(-time.Minute),
195 })
196 res, err := cs.RequeueTask(context.Background(), "/p", "failed", 3, "")
197 if err != nil || !res.Accepted || res.State != TaskStateQueued || res.RuntimeState != RuntimeStateExited {
198 t.Fatalf("expected expired lease to requeue, got result=%+v err=%v", res, err)
199 }
200 }
201
202 func TestControlService_RequeueRejectsNonFailedState(t *testing.T) {
203 s := NewInMemoryStore()
204 cs := NewControlService(s)
205 mustUpsertControl(t, s, "/p", TaskSnapshot{
206 SchemaVersion: 1, TaskID: "done", SessionID: "s1",
207 State: TaskStateSucceeded, RuntimeState: RuntimeStateExited, Version: 3,
208 CreatedAt: time.Now(), UpdatedAt: time.Now(),
209 })
210 res, err := cs.RequeueTask(context.Background(), "/p", "done", 3, "")
211 if err != nil || res.Error == nil || res.Error.Code != ErrTaskNotRequeueable {
212 t.Fatalf("expected not-requeueable guard, got result=%+v err=%v", res, err)
213 }
214 }
215
216 func TestControlService_Idempotency(t *testing.T) {
217 s := NewInMemoryStore()
218 cs := NewControlService(s)
219
220 mustUpsertControl(t, s, "/p", TaskSnapshot{
221 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
222 State: TaskStateRunning, Version: 1,
223 CreatedAt: time.Now(), UpdatedAt: time.Now(),
224 })
225
226 // First call
227 killer := &mockKiller{fn: func(string, string) bool { return true }}
228 res1, err := cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "key-1", killer)
229 if err != nil || !res1.Accepted {
230 t.Fatalf("first call failed: %v, %+v", err, res1)
231 }
232
233 // Second call with same key, op, task, version — idempotent
234 res2, err := cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "key-1", killer)
235 if err != nil {
236 t.Fatalf("second call: %v", err)
237 }
238 if !res2.Idempotent || !res2.Accepted {
239 t.Errorf("expected idempotent accepted, got %+v", res2)
240 }
241 }
242
243 func TestControlService_IdempotencyConflict_DifferentOp(t *testing.T) {
244 s := NewInMemoryStore()
245 cs := NewControlService(s)
246 killer := &mockKiller{fn: func(string, string) bool { return true }}
247
248 mustUpsertControl(t, s, "/p", TaskSnapshot{
249 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
250 State: TaskStateRunning, Version: 1,
251 CreatedAt: time.Now(), UpdatedAt: time.Now(),
252 })
253
254 cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "key-1", killer)
255 // Same key but different command
256 res, _ := cs.CancelTask(context.Background(), "/p", "t1", 1, "", "key-1")
257 if !strings.Contains(res.Error.Code, "idempotency") {
258 t.Errorf("expected idempotency conflict, got %+v", res.Error)
259 }
260 }
261
262 func TestControlService_IdempotencyConflict_DifferentVersion(t *testing.T) {
263 s := NewInMemoryStore()
264 cs := NewControlService(s)
265 killer := &mockKiller{fn: func(string, string) bool { return true }}
266
267 mustUpsertControl(t, s, "/p", TaskSnapshot{
268 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
269 State: TaskStateRunning, Version: 1,
270 CreatedAt: time.Now(), UpdatedAt: time.Now(),
271 })
272
273 cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "key-1", killer)
274 res, _ := cs.StopTaskWithKiller(context.Background(), "/p", "t1", 2, "", "key-1", killer)
275 if !strings.Contains(res.Error.Code, "idempotency") {
276 t.Errorf("expected idempotency conflict for different version, got %+v", res.Error)
277 }
278 }
279
280 func TestControlService_AuditEvent(t *testing.T) {
281 s := NewInMemoryStore()
282 cs := NewControlService(s)
283 killer := &mockKiller{fn: func(string, string) bool { return true }}
284
285 mustUpsertControl(t, s, "/p", TaskSnapshot{
286 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
287 State: TaskStateRunning, Version: 1,
288 CreatedAt: time.Now(), UpdatedAt: time.Now(),
289 })
290
291 cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, `stop command "rm -rf ./private" in /Users/alice/project`, "", killer)
292
293 events, _ := s.ListEvents(context.Background(), "/p", "t1", 0)
294 found := false
295 for _, ev := range events {
296 if ev.EventType == "control_stop" {
297 found = true
298 if ev.Sequence < 1 {
299 t.Errorf("expected positive sequence, got %d", ev.Sequence)
300 }
301 if ev.ErrorSummary != "" {
302 t.Errorf("control reason leaked into event: %q", ev.ErrorSummary)
303 }
304 if ev.SessionID != "s1" {
305 t.Errorf("expected session s1, got %q", ev.SessionID)
306 }
307 if ev.TaskID != "t1" {
308 t.Errorf("expected task t1, got %q", ev.TaskID)
309 }
310 }
311 }
312 if !found {
313 t.Error("expected audit event for stop")
314 }
315 }
316
317 func TestControlService_StopPreservesRuntimeLeaseUntilExit(t *testing.T) {
318 s := NewInMemoryStore()
319 cs := NewControlService(s)
320 now := time.Now()
321 leaseUntil := now.Add(time.Minute)
322 mustUpsertControl(t, s, "/p", TaskSnapshot{
323 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
324 State: TaskStateRunning, RuntimeState: RuntimeStateAlive,
325 RuntimeLeaseUntil: leaseUntil, RuntimeOwnerID: "owner-1", Version: 1,
326 CreatedAt: now, UpdatedAt: now,
327 })
328
329 res, err := cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "", &mockKiller{fn: func(string, string) bool { return true }})
330 if err != nil || !res.Accepted {
331 t.Fatalf("stop: result=%+v err=%v", res, err)
332 }
333 snap, err := s.GetTask(context.Background(), "/p", "t1")
334 if err != nil || snap == nil {
335 t.Fatalf("snapshot: %+v err=%v", snap, err)
336 }
337 if snap.RuntimeState != RuntimeStateAlive || snap.RuntimeOwnerID != "owner-1" || !snap.RuntimeLeaseUntil.Equal(leaseUntil) {
338 t.Fatalf("stop discarded live runtime ownership: %+v", snap)
339 }
340 reconciled := *snap
341 reconcileRuntime(&reconciled, leaseUntil.Add(time.Second))
342 if reconciled.State != TaskStateCancelled || reconciled.RuntimeState != RuntimeStateExited {
343 t.Fatalf("expired cancelled runtime did not reconcile: %+v", reconciled)
344 }
345 }
346
347 func TestControlService_StopBoundsLegacyLeaseLessRuntime(t *testing.T) {
348 s := NewInMemoryStore()
349 cs := NewControlService(s)
350 now := time.Now()
351 mustUpsertControl(t, s, "/p", TaskSnapshot{
352 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
353 State: TaskStateRunning, RuntimeState: RuntimeStateAlive,
354 RuntimeOwnerID: "owner-1", Version: 1, CreatedAt: now, UpdatedAt: now,
355 })
356
357 res, err := cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "", &mockKiller{fn: func(string, string) bool { return true }})
358 if err != nil || !res.Accepted {
359 t.Fatalf("stop: result=%+v err=%v", res, err)
360 }
361 snap, err := s.GetTask(context.Background(), "/p", "t1")
362 if err != nil || snap == nil {
363 t.Fatalf("snapshot: %+v err=%v", snap, err)
364 }
365 if snap.RuntimeState != RuntimeStateAlive || snap.RuntimeLeaseUntil.IsZero() || snap.RuntimeOwnerID != "owner-1" {
366 t.Fatalf("legacy runtime did not receive bounded lease: %+v", snap)
367 }
368 if got := snap.RuntimeLeaseUntil.Sub(snap.UpdatedAt); got != runtimeLeaseTTL {
369 t.Fatalf("lease duration = %v, want %v", got, runtimeLeaseTTL)
370 }
371 reconciled := *snap
372 reconcileRuntime(&reconciled, snap.RuntimeLeaseUntil.Add(time.Second))
373 if reconciled.State != TaskStateCancelled || reconciled.RuntimeState != RuntimeStateExited {
374 t.Fatalf("expired legacy runtime did not reconcile: %+v", reconciled)
375 }
376 }
377
378 func TestControlService_FileStoreClaimsIdempotencyBeforeSideEffects(t *testing.T) {
379 project := t.TempDir()
380 store := NewFileStore(".reasonix/tasks")
381 now := time.Now()
382 if err := store.SaveTask(context.Background(), project, TaskSnapshot{
383 SchemaVersion: 1, TaskID: "t1", SessionID: "s1", State: TaskStateRunning,
384 RuntimeState: RuntimeStateAlive, Version: 1, CreatedAt: now, UpdatedAt: now,
385 }); err != nil {
386 t.Fatal(err)
387 }
388 claimed := make(chan struct{})
389 release := make(chan struct{})
390 killer := &mockKiller{fn: func(string, string) bool {
391 close(claimed)
392 <-release
393 return true
394 }}
395 firstDone := make(chan ControlResult, 1)
396 go func() {
397 res, _ := NewControlService(store).StopTaskWithKiller(context.Background(), project, "t1", 1, "", "same-key", killer)
398 firstDone <- res
399 }()
400 <-claimed
401 second, err := NewControlService(store).StopTaskWithKiller(context.Background(), project, "t1", 1, "", "same-key", &mockKiller{fn: func(string, string) bool { t.Fatal("second request reached runtime"); return true }})
402 if err != nil || second.Error == nil || second.Error.Code != ErrTaskInProgress {
403 t.Fatalf("expected pending idempotency claim, got result=%+v err=%v", second, err)
404 }
405 close(release)
406 first := <-firstDone
407 if !first.Accepted || first.State != TaskStateCancelled {
408 t.Fatalf("first operation not accepted: %+v", first)
409 }
410 }
411
412 func TestInMemoryStore_IdempotencyClaimIsPendingUntilFinalized(t *testing.T) {
413 store := NewInMemoryStore()
414 r := IdempotencyRecord{Key: "same-key", Op: "stop", TaskID: "t1", Version: 1}
415 first, err := store.ClaimIdempotency(context.Background(), "/p", r)
416 if err != nil || first != nil {
417 t.Fatalf("first claim = %+v, err=%v", first, err)
418 }
419 second, err := store.ClaimIdempotency(context.Background(), "/p", r)
420 if err != nil || second == nil || !second.Pending {
421 t.Fatalf("second claim = %+v, err=%v; want pending record", second, err)
422 }
423 if err := store.FinalizeIdempotency(context.Background(), "/p", r); err != nil {
424 t.Fatal(err)
425 }
426 final, err := store.ClaimIdempotency(context.Background(), "/p", r)
427 if err != nil || final == nil || final.Pending {
428 t.Fatalf("final claim = %+v, err=%v; want finalized record", final, err)
429 }
430 }
431
432 func TestFileStore_IdempotencyClaimQuarantinesCorruptRecord(t *testing.T) {
433 root := t.TempDir()
434 store := NewFileStore(filepath.Join(".reasonix", "tasks"))
435 key := "broken-key"
436 idemDir := filepath.Join(root, ".reasonix", "tasks", ".idempotency")
437 if err := os.MkdirAll(idemDir, 0o700); err != nil {
438 t.Fatal(err)
439 }
440 target := filepath.Join(idemDir, key+".json")
441 if err := os.WriteFile(target, []byte(`{"pending":`), 0o600); err != nil {
442 t.Fatal(err)
443 }
444 rec := IdempotencyRecord{Key: key, Op: "stop", TaskID: "t1", Version: 1}
445 claimed, err := store.ClaimIdempotency(context.Background(), root, rec)
446 if err != nil || claimed != nil {
447 t.Fatalf("claim = %+v, err=%v; want fresh claim", claimed, err)
448 }
449 if _, err := os.Stat(target); err != nil {
450 t.Fatalf("fresh claim was not published: %v", err)
451 }
452 backups, err := filepath.Glob(target + ".corrupt-*")
453 if err != nil || len(backups) != 1 {
454 t.Fatalf("corrupt record backups = %v, err=%v; want one quarantined record", backups, err)
455 }
456 }
457
458 func TestControlService_AuditSequenceMonotonic(t *testing.T) {
459 s := NewInMemoryStore()
460 cs := NewControlService(s)
461 killer := &mockKiller{fn: func(string, string) bool { return true }}
462
463 mustUpsertControl(t, s, "/p", TaskSnapshot{
464 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
465 State: TaskStateRunning, Version: 1,
466 CreatedAt: time.Now(), UpdatedAt: time.Now(),
467 })
468
469 // Stop creates audit event sequence 1
470 cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "", killer)
471
472 // Reset task to running (simulate a new execution lifecycle)
473 s.UpsertTask("/p", TaskSnapshot{
474 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
475 State: TaskStateRunning, Version: 2,
476 CreatedAt: time.Now(), UpdatedAt: time.Now(),
477 })
478
479 // Cancel should get sequence 2 from NextSequence
480 res, _ := cs.CancelTaskWithKiller(context.Background(), "/p", "t1", 2, "", "", killer)
481 if !res.Accepted {
482 t.Fatalf("cancel failed: %+v", res)
483 }
484
485 events, _ := s.ListEvents(context.Background(), "/p", "t1", 0)
486 if len(events) != 2 {
487 t.Fatalf("expected 2 events, got %d", len(events))
488 }
489 if events[1].Sequence != 2 {
490 t.Errorf("expected sequence 2, got %d", events[1].Sequence)
491 }
492 }
493
494 func TestControlService_KillJob(t *testing.T) {
495 s := NewInMemoryStore()
496 cs := NewControlService(s)
497
498 killed := false
499 mk := &mockKiller{fn: func(sessionID, id string) bool {
500 killed = true
501 return sessionID == "s1" && id == "t1"
502 }}
503
504 mustUpsertControl(t, s, "/p", TaskSnapshot{
505 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
506 State: TaskStateRunning, Version: 1,
507 CreatedAt: time.Now(), UpdatedAt: time.Now(),
508 })
509
510 res, _ := cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "", mk)
511 if !res.Accepted {
512 t.Fatalf("stop failed: %+v", res)
513 }
514 if !killed {
515 t.Error("expected Kill to be called for stop")
516 }
517 }
518
519 func TestControlService_KillNotCalledForTerminalTask(t *testing.T) {
520 s := NewInMemoryStore()
521 cs := NewControlService(s)
522
523 killed := false
524 mk := &mockKiller{fn: func(_, _ string) bool { killed = true; return true }}
525
526 mustUpsertControl(t, s, "/p", TaskSnapshot{
527 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
528 State: TaskStateSucceeded, Version: 1,
529 CreatedAt: time.Now(), UpdatedAt: time.Now(),
530 })
531
532 cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "", mk)
533 if killed {
534 t.Error("Kill should not be called for terminal tasks")
535 }
536 }
537
538 func TestControlService_ConcurrentKillersRemainCallScoped(t *testing.T) {
539 s := NewInMemoryStore()
540 cs := NewControlService(s)
541 now := time.Now()
542 for _, snap := range []TaskSnapshot{
543 {SchemaVersion: 1, TaskID: "task-a", SessionID: "session-a", State: TaskStateRunning, RuntimeState: RuntimeStateAlive, Version: 1, CreatedAt: now, UpdatedAt: now},
544 {SchemaVersion: 1, TaskID: "task-b", SessionID: "session-b", State: TaskStateRunning, RuntimeState: RuntimeStateAlive, Version: 1, CreatedAt: now, UpdatedAt: now},
545 } {
546 mustUpsertControl(t, s, "/p", snap)
547 }
548
549 started := make(chan struct{})
550 killed := make(chan string, 2)
551 var wg sync.WaitGroup
552 for _, tc := range []struct {
553 taskID, sessionID string
554 }{
555 {taskID: "task-a", sessionID: "session-a"},
556 {taskID: "task-b", sessionID: "session-b"},
557 } {
558 wg.Go(func() {
559 <-started
560 killer := &mockKiller{fn: func(sessionID, taskID string) bool {
561 killed <- sessionID + "/" + taskID
562 return sessionID == tc.sessionID && taskID == tc.taskID
563 }}
564 res, err := cs.StopTaskWithKiller(context.Background(), "/p", tc.taskID, 1, "", "", killer)
565 if err != nil || !res.Accepted {
566 t.Errorf("StopTaskWithKiller(%s): result=%+v err=%v", tc.taskID, res, err)
567 }
568 })
569 }
570 close(started)
571 wg.Wait()
572 close(killed)
573
574 got := map[string]bool{}
575 for target := range killed {
576 got[target] = true
577 }
578 for _, want := range []string{"session-a/task-a", "session-b/task-b"} {
579 if !got[want] {
580 t.Fatalf("missing call-scoped kill %q; got %v", want, got)
581 }
582 }
583 }
584
585 func TestControlService_ConcurrentAccess(t *testing.T) {
586 s := NewInMemoryStore()
587 cs := NewControlService(s)
588 killer := &mockKiller{fn: func(string, string) bool { return true }}
589
590 mustUpsertControl(t, s, "/p", TaskSnapshot{
591 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
592 State: TaskStateRunning, Version: 1,
593 CreatedAt: time.Now(), UpdatedAt: time.Now(),
594 })
595
596 var wg sync.WaitGroup
597 success := 0
598 var mu sync.Mutex
599
600 for range 10 {
601 wg.Go(func() {
602 res, _ := cs.StopTaskWithKiller(context.Background(), "/p", "t1", 1, "", "", killer)
603 if res.Accepted {
604 mu.Lock()
605 success++
606 mu.Unlock()
607 }
608 })
609 }
610 wg.Wait()
611 // Exactly one caller should succeed due to mutex + version CAS
612 if success != 1 {
613 t.Errorf("expected exactly 1 success, got %d", success)
614 }
615 }
616
617 func TestControlService_CancelTask(t *testing.T) {
618 s := NewInMemoryStore()
619 cs := NewControlService(s)
620 killer := &mockKiller{fn: func(string, string) bool { return true }}
621
622 mustUpsertControl(t, s, "/p", TaskSnapshot{
623 SchemaVersion: 1, TaskID: "t1", SessionID: "s1",
624 State: TaskStateWaiting, Version: 1,
625 CreatedAt: time.Now(), UpdatedAt: time.Now(),
626 })
627
628 res, _ := cs.CancelTaskWithKiller(context.Background(), "/p", "t1", 1, "timeout", "", killer)
629 if !res.Accepted || res.State != TaskStateCancelled {
630 t.Errorf("expected cancelled, got %+v", res)
631 }
632 }
633
634 func TestControlService_OpenSession(t *testing.T) {
635 s := NewInMemoryStore()
636 cs := NewControlService(s)
637
638 mustUpsertControl(t, s, "/p", TaskSnapshot{
639 SchemaVersion: 1, TaskID: "t1", SessionID: "sess-abc",
640 State: TaskStateRunning, Version: 1,
641 CreatedAt: time.Now(), UpdatedAt: time.Now(),
642 })
643
644 res, _ := cs.OpenTaskSession(context.Background(), "/p", "t1")
645 if res.SessionID != "sess-abc" || !res.Accepted {
646 t.Errorf("expected sess-abc, got %+v", res)
647 }
648 }
649
650 // mockKiller implements JobKiller for tests.
651 type mockKiller struct {
652 fn func(string, string) bool
653 }
654
655 func (m *mockKiller) Kill(sessionID, id string) bool {
656 if m.fn != nil {
657 return m.fn(sessionID, id)
658 }
659 return false
660 }
661
662 func mustUpsertControl(t *testing.T, s *InMemoryStore, proj string, snap TaskSnapshot) {
663 t.Helper()
664 if err := s.UpsertTask(proj, snap); err != nil {
665 t.Fatal(err)
666 }
667 }
668
668 lines GO