Support agentic test session cancellation

packages/cli/src/commands/pw-test.ts

@@ -39,6 +39,7 @@ import { cased } from '../sourcegen/index.js'
 import { shellQuote } from '../services/shell.js'
 import { Runtime } from '../runtimes/index.js'
 import { Bundler } from '../services/check-parser/bundler.js'
+import { registerTestSessionCancelHandler } from '../services/test-session-cancel.js'
 
 export default class PwTestCommand extends AuthCommand {
   static coreCommand = true
@@ -348,11 +349,7 @@ export default class PwTestCommand extends AuthCommand {
       }, links))
     })
 
-    runner.on(Events.CANCEL, async testSessionId => {
-      reporters.forEach(r => r.onCancel())
-      if (!testSessionId) return
-      await api.cancel.cancelTestSession({ testSessionId })
-    })
+    registerTestSessionCancelHandler(runner, reporters)
 
     runner.on(Events.DETACH, () => reporters.forEach(r => r.onDetach()))
 

packages/cli/src/commands/test.ts

@@ -25,6 +25,7 @@ import { BrowserCheckBundle } from '../constructs/browser-check-bundle.js'
 import { prepareReportersTypes, prepareRunLocation, validateDetachReporterTypes } from '../helpers/test-helper.js'
 import { Runtime } from '../runtimes/index.js'
 import { Bundler } from '../services/check-parser/bundler.js'
+import { registerTestSessionCancelHandler } from '../services/test-session-cancel.js'
 
 const MAX_RETRIES = 3
 
@@ -379,11 +380,7 @@ export default class Test extends AuthCommand {
       detach,
     )
 
-    runner.on(Events.CANCEL, async testSessionId => {
-      reporters.forEach(r => r.onCancel())
-      if (!testSessionId) return
-      await api.cancel.cancelTestSession({ testSessionId })
-    })
+    registerTestSessionCancelHandler(runner, reporters)
 
     runner.on(Events.DETACH, () => reporters.forEach(r => r.onDetach()))
 

packages/cli/src/commands/trigger.ts

@@ -21,6 +21,7 @@ import { NoMatchingChecksError, TestResultsShortLinks } from '../rest/test-sessi
 import { Session, RetryStrategyBuilder } from '../constructs/index.js'
 import { DEFAULT_REGION } from '../helpers/constants.js'
 import { validateDetachReporterTypes } from '../helpers/test-helper.js'
+import { registerTestSessionCancelHandler } from '../services/test-session-cancel.js'
 
 const MAX_RETRIES = 3
 
@@ -214,11 +215,7 @@ export default class Trigger extends AuthCommand {
       reporters.forEach(r => r.onError(err))
       process.exitCode = 1
     })
-    runner.on(Events.CANCEL, async testSessionId => {
-      reporters.forEach(r => r.onCancel())
-      if (!testSessionId) return
-      await api.cancel.cancelTestSession({ testSessionId })
-    })
+    registerTestSessionCancelHandler(runner, reporters)
     runner.on(Events.DETACH, () => reporters.forEach(r => r.onDetach()))
     await runner.run()
   }

packages/cli/src/services/__tests__/test-session-cancel.spec.ts

@@ -0,0 +1,54 @@
+import { describe, expect, it, vi } from 'vitest'
+import { EventEmitter } from 'node:events'
+
+import { Events } from '../abstract-check-runner.js'
+import { registerTestSessionCancelHandler } from '../test-session-cancel.js'
+
+function makeReporter () {
+  return {
+    onBegin: vi.fn(),
+    onCheckInProgress: vi.fn(),
+    onCheckAttemptResult: vi.fn(),
+    onCheckEnd: vi.fn(),
+    onEnd: vi.fn(),
+    onError: vi.fn(),
+    onSchedulingDelayExceeded: vi.fn(),
+    onStreamLogs: vi.fn(),
+    onCancel: vi.fn(),
+    onDetach: vi.fn(),
+  }
+}
+
+describe('registerTestSessionCancelHandler', () => {
+  it('cancels the whole test session for an agentic-shaped run', async () => {
+    const runner = new EventEmitter()
+    const reporter = makeReporter()
+    const cancelClient = {
+      cancelTestSession: vi.fn().mockResolvedValue(undefined),
+    }
+
+    registerTestSessionCancelHandler(runner, [reporter], cancelClient)
+    runner.emit(Events.CANCEL, 'ts-agentic')
+    await vi.waitFor(() => expect(cancelClient.cancelTestSession).toHaveBeenCalled())
+
+    expect(reporter.onCancel).toHaveBeenCalledTimes(1)
+    expect(cancelClient.cancelTestSession).toHaveBeenCalledWith({ testSessionId: 'ts-agentic' })
+    expect(cancelClient.cancelTestSession).not.toHaveBeenCalledWith(
+      expect.objectContaining({ sequenceId: expect.anything() }),
+    )
+  })
+
+  it('still marks reporters as cancelling when no test session was created', () => {
+    const runner = new EventEmitter()
+    const reporter = makeReporter()
+    const cancelClient = {
+      cancelTestSession: vi.fn().mockResolvedValue(undefined),
+    }
+
+    registerTestSessionCancelHandler(runner, [reporter], cancelClient)
+    runner.emit(Events.CANCEL)
+
+    expect(reporter.onCancel).toHaveBeenCalledTimes(1)
+    expect(cancelClient.cancelTestSession).not.toHaveBeenCalled()
+  })
+})

packages/cli/src/services/__tests__/test-session-runners.spec.ts

@@ -0,0 +1,126 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+import type { RunLocation } from '../abstract-check-runner.js'
+import TestRunner from '../test-runner.js'
+import TriggerRunner from '../trigger-runner.js'
+import { testSessions } from '../../rest/api.js'
+
+vi.mock('../../rest/api.js', () => ({
+  testSessions: {
+    run: vi.fn(),
+    trigger: vi.fn(),
+  },
+}))
+
+const RUN_LOCATION: RunLocation = { type: 'PUBLIC', region: 'eu-west-1' }
+
+describe('test-session runners', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('schedules local Agentic checks as cancellable test-session jobs', async () => {
+    vi.mocked(testSessions.run).mockResolvedValue({
+      data: {
+        testSessionId: 'ts-agentic',
+        sequenceIds: {
+          'agentic-logical-id': 'seq-agentic',
+        },
+      },
+    } as any)
+
+    const agenticCheck = {
+      logicalId: 'agentic-logical-id',
+      groupId: undefined,
+      getSourceFile: () => 'agentic.check.ts',
+    }
+    const agenticBundle = {
+      synthesize: vi.fn(() => ({
+        checkType: 'AGENTIC',
+        name: 'Agentic Check',
+      })),
+    }
+    const projectBundle = {
+      project: { name: 'Agentic Project', logicalId: 'agentic-project' },
+      data: {
+        'check-group': {},
+      },
+    }
+
+    const runner = new TestRunner(
+      'account-id',
+      projectBundle as any,
+      [{ construct: agenticCheck, bundle: agenticBundle }] as any,
+      [],
+      RUN_LOCATION,
+      60,
+      false,
+      true,
+      null,
+      null,
+      false,
+      '.',
+      null,
+    )
+
+    const scheduled = await runner.scheduleChecks('suite-id')
+    const payload = vi.mocked(testSessions.run).mock.calls[0][0]
+
+    expect(payload.checkRunJobs[0]).toMatchObject({
+      checkType: 'AGENTIC',
+      logicalId: 'agentic-logical-id',
+      filePath: 'agentic.check.ts',
+      sourceInfo: {
+        checkRunSuiteId: 'suite-id',
+        updateSnapshots: false,
+      },
+    })
+    expect(scheduled).toEqual({
+      testSessionId: 'ts-agentic',
+      checks: [{ check: agenticCheck, sequenceId: 'seq-agentic' }],
+    })
+  })
+
+  it('maps triggered Agentic checks back to the test-session sequence IDs', async () => {
+    const agenticCheck = {
+      id: 'agentic-check-id',
+      name: 'Triggered Agentic Check',
+      checkType: 'AGENTIC',
+    }
+    vi.mocked(testSessions.trigger).mockResolvedValue({
+      data: {
+        checks: [agenticCheck],
+        testSessionId: 'ts-trigger-agentic',
+        sequenceIds: {
+          'agentic-check-id': 'seq-trigger-agentic',
+        },
+      },
+    } as any)
+
+    const runner = new TriggerRunner(
+      'account-id',
+      60,
+      false,
+      true,
+      RUN_LOCATION,
+      [],
+      ['agentic-check-id'],
+      [],
+      null,
+      null,
+      undefined,
+      null,
+    )
+
+    const scheduled = await runner.scheduleChecks('suite-id')
+
+    expect(testSessions.trigger).toHaveBeenCalledWith(expect.objectContaining({
+      checkRunSuiteId: 'suite-id',
+      checkId: ['agentic-check-id'],
+    }))
+    expect(scheduled).toEqual({
+      testSessionId: 'ts-trigger-agentic',
+      checks: [{ check: agenticCheck, sequenceId: 'seq-trigger-agentic' }],
+    })
+  })
+})

packages/cli/src/services/test-session-cancel.ts

@@ -0,0 +1,30 @@
+import type { Reporter } from '../reporters/reporter.js'
+import { cancel } from '../rest/api.js'
+import { Events } from './abstract-check-runner.js'
+
+type TestSessionCancelEmitter = {
+  on(event: Events.CANCEL, listener: (testSessionId?: string) => Promise<void>): unknown
+}
+
+type TestSessionCancelClient = {
+  cancelTestSession(input: { testSessionId: string }): Promise<unknown>
+}
+
+export function registerTestSessionCancelHandler (
+  runner: TestSessionCancelEmitter,
+  reporters: Reporter[],
+  cancelClient: TestSessionCancelClient = cancel,
+) {
+  runner.on(Events.CANCEL, async testSessionId => {
+    reporters.forEach(r => r.onCancel())
+
+    if (!testSessionId) {
+      return
+    }
+
+    // Cancellation is intentionally test-session scoped. The backend resolves
+    // which running results are cancellable (currently Playwright and Agentic),
+    // so the CLI must not narrow the request by check type or sequence ID here.
+    await cancelClient.cancelTestSession({ testSessionId })
+  })
+}