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