fix(cli): make scan shutdown verifiable

- preselect and persist workflow identity before worker launch
- cancel first, then verify bounded Temporal termination
- reconcile Docker workers with Temporal open workflows
- fail closed on stale images and unavailable lifecycle state
- mark cancellation only after confirmed shutdown
This commit is contained in:
ajmallesh
2026-08-31 16:16:14 -07:00
parent 1b440c853c
commit 7b67302a39
9 changed files with 1246 additions and 191 deletions
+121 -5
View File
@@ -1,5 +1,5 @@
/**
* Thin Temporal client for reading one scan's state.
* Thin Temporal client for reading scan state and controlling scan workflow lifecycle.
*
* A running scan is queried live (getProgress) and read via pendingActivities for
* the in-flight agents; a closed scan is read once from its result. Everything goes
@@ -13,10 +13,20 @@ import { ACTIVITY_TO_PROGRESS, type PipelineState } from './scan/pipeline.js';
const ADDRESS = '127.0.0.1:7233';
const NAMESPACE = 'default';
const LIFECYCLE_RPC_DEADLINE_MS = 3_000;
const OPEN_SCAN_WORKFLOW_QUERY =
"WorkflowType = 'pentestPipelineWorkflow' AND (ExecutionStatus = 'Running' OR ExecutionStatus = 'Paused')";
// WorkflowExecutionStatusName values that mean the scan has closed. RUNNING (and the unused
// CONTINUED_AS_NEW) are the only non-terminal states.
const TERMINAL_STATUSES: ReadonlySet<string> = new Set(['COMPLETED', 'FAILED', 'CANCELLED', 'TERMINATED', 'TIMED_OUT']);
// WorkflowExecutionStatusName values that positively prove this execution has closed.
// PAUSED is open; UNSPECIFIED and UNKNOWN are not safe closure evidence.
const TERMINAL_STATUSES: ReadonlySet<string> = new Set([
'COMPLETED',
'FAILED',
'CANCELLED',
'TERMINATED',
'CONTINUED_AS_NEW',
'TIMED_OUT',
]);
export interface RunningAgent {
readonly agent: string;
@@ -66,11 +76,28 @@ export type TerminalOutcome =
| { readonly kind: 'success'; readonly state: PipelineState }
| { readonly kind: 'failed'; readonly message: string };
/**
* The authoritative Temporal state used by lifecycle commands. Transport failures deliberately
* remain errors instead of being represented as a closed workflow: callers must not report a
* scan stopped unless Temporal has positively confirmed it.
*/
export type WorkflowLifecycleState =
| { readonly kind: 'open'; readonly status: 'RUNNING' | 'PAUSED' }
| { readonly kind: 'terminal'; readonly status: string }
| { readonly kind: 'unknown'; readonly status: string }
| { readonly kind: 'not-found' };
/** A scan workflow returned by Temporal's eventually consistent open-workflow visibility query. */
export interface RunningScanWorkflow {
readonly workflowId: string;
readonly taskQueue: string;
}
let clientPromise: Promise<Client> | null = null;
function getClient(): Promise<Client> {
if (!clientPromise) {
const pending = Connection.connect({ address: ADDRESS }).then(
const pending = Connection.connect({ address: ADDRESS, connectTimeout: LIFECYCLE_RPC_DEADLINE_MS }).then(
(connection) => new Client({ connection, namespace: NAMESPACE }),
);
// A rejected connect must not be cached forever: clear the memo so the next call rebuilds
@@ -95,6 +122,95 @@ function resetClient(only?: Promise<Client>): void {
previous?.then((client) => client.connection.close()).catch(() => {});
}
/** Close the current channel and establish another before a termination retry. */
export async function refreshWorkflowLifecycleConnection(): Promise<void> {
const previous = clientPromise;
if (previous !== null) {
if (clientPromise === previous) clientPromise = null;
try {
const client = await previous;
await client.connection.close();
} catch {
// A failed prior connection is already detached. The new connection below is authoritative.
}
}
await getClient();
}
/**
* Run a bounded lifecycle RPC and discard the connection when Temporal did not positively say
* that the workflow is absent. A fresh connection is important after a gRPC timeout or transport
* failure: reusing a wedged channel can turn a recoverable stop into an indefinitely ambiguous one.
*/
async function runLifecycleRpc<T>(operation: (client: Client) => Promise<T>): Promise<T> {
const pending = getClient();
try {
const client = await pending;
return await client.withDeadline(Date.now() + LIFECYCLE_RPC_DEADLINE_MS, () => operation(client));
} catch (err) {
if (!(err instanceof WorkflowNotFoundError)) resetClient(pending);
throw err;
}
}
/** Describe a workflow for lifecycle control without reading its progress or pending activities. */
export async function describeWorkflowLifecycle(workflowId: string): Promise<WorkflowLifecycleState> {
try {
const desc = await runLifecycleRpc((client) => client.workflow.getHandle(workflowId).describe());
if (desc.status.name === 'RUNNING' || desc.status.name === 'PAUSED') {
return { kind: 'open', status: desc.status.name };
}
if (TERMINAL_STATUSES.has(desc.status.name)) return { kind: 'terminal', status: desc.status.name };
return { kind: 'unknown', status: desc.status.name };
} catch (err) {
if (err instanceof WorkflowNotFoundError) return { kind: 'not-found' };
throw err;
}
}
/** Request cooperative cancellation. This confirms request acceptance, not workflow closure. */
export async function requestWorkflowCancellation(workflowId: string): Promise<'requested' | 'not-found'> {
try {
await runLifecycleRpc((client) => client.workflow.getHandle(workflowId).cancel());
return 'requested';
} catch (err) {
if (err instanceof WorkflowNotFoundError) return 'not-found';
throw err;
}
}
/** Request forced termination. This confirms request acceptance, not workflow closure. */
export async function requestWorkflowTermination(
workflowId: string,
reason: string,
): Promise<'requested' | 'not-found'> {
try {
await runLifecycleRpc((client) => client.workflow.getHandle(workflowId).terminate(reason));
return 'requested';
} catch (err) {
if (err instanceof WorkflowNotFoundError) return 'not-found';
throw err;
}
}
/** List currently open Shannon scan workflows through Temporal visibility. */
export async function listRunningScanWorkflows(): Promise<readonly RunningScanWorkflow[]> {
return runLifecycleRpc(async (client) => {
const workflows: RunningScanWorkflow[] = [];
for await (const execution of client.workflow.list({ query: OPEN_SCAN_WORKFLOW_QUERY })) {
// Visibility is eventually consistent. Keep only the open scan rows returned by this page;
// each discovered workflow is described directly before `stop` accepts its closure.
if (
(execution.status.name === 'RUNNING' || execution.status.name === 'PAUSED') &&
execution.type === 'pentestPipelineWorkflow'
) {
workflows.push({ workflowId: execution.workflowId, taskQueue: execution.taskQueue });
}
}
return workflows;
});
}
/** Describe a scan: status, timing, and the agents currently running (from pendingActivities). Null if not found. */
export async function describeScan(workflowId: string): Promise<ScanDescription | null> {
const client = await getClient();