[CRCR] Use {run_id}:{run_attempt} instead of {check_run_id} in state machine key composition (#8170)

Use
`f"{_STATE_PREFIX}{delivery_id}:{downstream_repo}:{run_id}:{run_attempt}"`
instead of
`f"{_STATE_PREFIX}{delivery_id}:{downstream_repo}:{check_run_id}"` to
extend support of enabling downstream setting multiple jobs within one
workflow. Because `run_id` is workflow-level and `check_run_id` is
job-level and we don't want to see too many entries (one for each job
within a single workflow) displayed on HUD and PR CI panel.

cc @fffrog @can-gaa-hou @subinz1 @jewelkm89

Author: KarhouTam

aws/lambda/cross_repo_ci_relay/README.md

@@ -44,13 +44,13 @@ The callback endpoint validates incoming callbacks and forwards them to HUD for
 - **Identity**: the `Authorization: Bearer <oidc-token>` header is verified against GitHub's JWKS.  The OIDC `repository` claim is a trusted identity for the caller and is used for the L2+ allowlist check. Relay forwards this trusted value to HUD as a top-level `verified_repo` field; HUD should prefer it over anything self-reported in `callback_payload`.
 - **Repo level**: Relay determines the downstream repository's allowlist level (L1–L4) and forwards it to HUD as `downstream_repo_level`. This authoritative level information is determined once by the relay, ensuring HUD doesn't need to recompute it and avoiding synchronization/timing issues if tiering information becomes dynamic.
 - **Schema validation**: Relay validates that required fields (`delivery_id` and `workflow.status`) are present in the callback body.  Missing fields result in a `400` error to signal contract violations to the caller.  HUD receives validated data and does not need to perform schema checks.
-- **State machine**: Relay maintains a **unified state machine** in Redis to validate callback lifecycles, compute timing metrics, and support per-job tracking:
-  - **Unified structure**: Single enum `CallbackState` with states `DISPATCHED` (webhook side, keyed by sentinel `check_run_id="dispatched"`), `IN_PROGRESS`, and `COMPLETED` (callback side, per-job). State records stored as JSON: `{"state": "...", "timestamp": 1234.56, "job_name": "...", "run_id": "..."}`.
+- **State machine**: Relay maintains a **unified state machine** in Redis to validate callback lifecycles, compute timing metrics, and support per-workflow tracking:
+  - **Unified structure**: Single enum `CallbackState` with states `DISPATCHED` (webhook side, keyed by sentinel `run_id=0, run_attempt=0`), `IN_PROGRESS`, and `COMPLETED` (callback side, per-workflow). State records stored as JSON: `{"state": "...", "timestamp": 1234.56}`.
   - **Dispatch validation**: `DISPATCHED` state proves valid webhook origin. Callbacks without this state are rejected (no prior dispatch).
-  - **Job-level tracking**: Each job has independent state and timestamps keyed by `check_run_id` (`oot:state:{delivery_id}:{repo}:{check_run_id}`). Supports multiple jobs per webhook.
+  - **Workflow-level tracking**: Each workflow has independent state and timestamps keyed by `{run_id}:{run_attempt}` (`oot:state:{delivery_id}:{repo}:{run_id}:{run_attempt}`). Supports multiple workflows per webhook.
   - **Timing metrics**: `queue_time = dispatch_timestamp → in_progress_timestamp`, `execution_time = in_progress_timestamp → completed_timestamp`. Timestamps extracted from state records.
-  - **State transitions**: Rejects invalid flows (`COMPLETED` without prior `IN_PROGRESS`, duplicate `IN_PROGRESS` for the same `check_run_id`, duplicate `COMPLETED`, callbacks without a prior `DISPATCHED` record).
-    Note that the direction graph below is for a single check run, reruns have different `check_run_id` and are treated as separate jobs, so they won't violate the state machine since they won't have a prior `IN_PROGRESS` or `COMPLETED` record.
+  - **State transitions**: Rejects invalid flows (`COMPLETED` without prior `IN_PROGRESS`, duplicate `IN_PROGRESS` for the same `{run_id}:{run_attempt}`, duplicate `COMPLETED`, callbacks without a prior `DISPATCHED` record).
+    Note that the direction graph below is for a single check run, reruns have different `run_attempt` and are treated as separate workflows, so they won't violate the state machine since they won't have a prior `IN_PROGRESS` or `COMPLETED` record.
     ```mermaid
     stateDiagram-v2
         direction LR

aws/lambda/cross_repo_ci_relay/callback/callback_handler.py

@@ -11,7 +11,8 @@
 from utils.misc import (
     CallbackState,
     CallbackStateRecord,
-    DISPATCH_CHECK_RUN_ID,
+    DISPATCH_RUN_ATTEMPT,
+    DISPATCH_RUN_ID,
     HTTPException,
 )
 from utils.redis_helper import check_rate_limit
@@ -59,41 +60,41 @@ def _verify_access(
     return allowlist, repo_level
 
 
-def _parse_callback_body(body: dict) -> tuple[str, str, str, str, str]:
-    """Return (delivery_id, status, check_run_id, job_name, run_id) from ``body``.
+def _parse_callback_body(body: dict) -> tuple[str, str, int, int, str]:
+    """Return (delivery_id, status, run_id, run_attempt, workflow_name) from ``body``.
 
-    check_run_id is set by GitHub Actions (job.check_run_id context) and
-    cannot be tampered with, ensuring replay-attack detection integrity.
+    run_id and run_attempt uniquely identify a workflow run execution.
+    run_attempt defaults to 1 when not present in the callback body.
 
     Raises HTTPException(400) on any missing or mis-typed field.
     """
     try:
         delivery_id = body["delivery_id"]
         workflow_dict = body["workflow"]
         status = workflow_dict["status"]
-        check_run_id = workflow_dict["check_run_id"]  # Required
-        job_name = workflow_dict["job_name"]  # Required for HUD grouping
         run_id = workflow_dict["run_id"]  # Required for HUD grouping
+        run_attempt = workflow_dict.get("run_attempt", 1)  # Default to 1 if not present
+        workflow_name = workflow_dict["name"]  # Required for HUD grouping
     except (KeyError, TypeError) as exc:
         logger.warning(f"missing required field in callback body: {exc}")
         raise HTTPException(
             400, f"callback body missing required field: {exc}"
         ) from exc
-    return delivery_id, status, check_run_id, job_name, run_id
+    return delivery_id, status, run_id, run_attempt, workflow_name
 
 
 def _update_state_and_compute_metrics(
     config: RelayConfig,
     delivery_id: str,
     verified_repo: str,
-    check_run_id: str,
-    job_name: str,
-    run_id: str,
+    run_id: int,
+    run_attempt: int,
+    workflow_name: str,
     status: str,
     dispatch_record: CallbackStateRecord,
-    job_record: CallbackStateRecord | None,
+    workflow_record: CallbackStateRecord | None,
 ) -> dict:
-    """Persist the new job state to Redis and return CI timing metrics.
+    """Persist the new workflow state to Redis and return CI timing metrics.
 
     Writes IN_PROGRESS or COMPLETED state (with the current timestamp), then
     reads back the stored record to compute:
@@ -119,11 +120,11 @@ def _update_state_and_compute_metrics(
             config,
             delivery_id,
             verified_repo,
-            check_run_id,
+            run_id,
+            run_attempt,
             state,
             current_timestamp,
-            job_name,
-            run_id,
+            workflow_name,
         )
     except RedisError:
         raise HTTPException(
@@ -138,22 +139,24 @@ def _update_state_and_compute_metrics(
     except Exception:
         raise
 
-    updated_job_record = redis_helper.get_callback_state(
-        config, delivery_id, verified_repo, check_run_id
+    updated_workflow_record = redis_helper.get_callback_state(
+        config, delivery_id, verified_repo, run_id, run_attempt
     )
-    if updated_job_record is None:
+    if updated_workflow_record is None:
         return ci_metrics
 
     if state == CallbackState.IN_PROGRESS:
         ci_metrics["queue_time"] = _safe_delta(
             dispatch_record.timestamp,
-            updated_job_record.timestamp,
+            updated_workflow_record.timestamp,
             "queue_time",
         )
     else:
-        if job_record is not None:
+        if workflow_record is not None:
             ci_metrics["execution_time"] = _safe_delta(
-                job_record.timestamp, updated_job_record.timestamp, "execution_time"
+                workflow_record.timestamp,
+                updated_workflow_record.timestamp,
+                "execution_time",
             )
 
     return ci_metrics
@@ -181,10 +184,10 @@ def handle(config: RelayConfig, body: dict, verified_repo: str) -> dict:
         return {"ok": True, "status": "ignored"}
     _, repo_level = result
 
-    delivery_id, status, check_run_id, job_name, run_id = _parse_callback_body(body)
+    delivery_id, status, run_id, run_attempt, workflow_name = _parse_callback_body(body)
 
     dispatch_record = redis_helper.get_callback_state(
-        config, delivery_id, verified_repo, DISPATCH_CHECK_RUN_ID
+        config, delivery_id, verified_repo, DISPATCH_RUN_ID, DISPATCH_RUN_ATTEMPT
     )
     if not dispatch_record:
         logger.warning(
@@ -194,20 +197,20 @@ def handle(config: RelayConfig, body: dict, verified_repo: str) -> dict:
         )
         raise HTTPException(400, "callback rejected: no matching dispatch record")
 
-    job_record = redis_helper.get_callback_state(
-        config, delivery_id, verified_repo, check_run_id
+    workflow_record = redis_helper.get_callback_state(
+        config, delivery_id, verified_repo, run_id, run_attempt
     )
 
     ci_metrics = _update_state_and_compute_metrics(
         config,
         delivery_id,
         verified_repo,
-        check_run_id,
-        job_name,
         run_id,
+        run_attempt,
+        workflow_name,
         status,
         dispatch_record,
-        job_record,
+        workflow_record,
     )
 
     trusted = {

aws/lambda/cross_repo_ci_relay/tests/test_callback_handler.py

@@ -4,7 +4,7 @@
 
 from callback.callback_handler import handle
 from utils.allowlist import AllowlistLevel
-from utils.misc import CallbackState, DISPATCH_CHECK_RUN_ID, HTTPException
+from utils.misc import CallbackState, DISPATCH_RUN_ID, HTTPException
 from utils.redis_helper import CallbackStateRecord
 
 
@@ -19,7 +19,7 @@ def _cfg():
     return cfg
 
 
-def _body(status="completed", job_name="default", check_run_id="12345", run_id="99999"):
+def _body(status="completed", workflow_name="default", run_id=99999, run_attempt=1):
     return {
         "event_type": "pull_request",
         "delivery_id": "del-123",
@@ -32,9 +32,9 @@ def _body(status="completed", job_name="default", check_run_id="12345", run_id="
             "conclusion": "success" if status == "completed" else None,
             "name": "CI",
             "url": "http://ci.example.com/run/1",
-            "job_name": job_name,
-            "check_run_id": check_run_id,
+            "workflow_name": workflow_name,
             "run_id": run_id,
+            "run_attempt": run_attempt,
         },
     }
 
@@ -52,15 +52,19 @@ def setUp(self):
         self.mock_redis = self.patcher_redis.start()
         self.mock_redis.create_client.return_value = MagicMock()
 
-        # Setup default: dispatch exists, job state is None (in_progress not yet reported)
-        def default_get_state(cfg, delivery_id, repo, check_run_id_arg, client=None):
-            if check_run_id_arg == DISPATCH_CHECK_RUN_ID:
+        # Setup default: dispatch exists, workflow state is None (in_progress not yet reported)
+        def default_get_state(
+            cfg, delivery_id, repo, run_id_arg, run_attempt_arg, client=None
+        ):
+            if run_id_arg == DISPATCH_RUN_ID:
                 return CallbackStateRecord(
-                    CallbackState.DISPATCHED, time.time() - 30, "dispatch-job", 11111
+                    CallbackState.DISPATCHED,
+                    time.time() - 30,
                 )
-            elif check_run_id_arg == "12345":  # default check_run_id in _body()
+            elif run_id_arg == 99999:  # default run_id in _body()
                 return CallbackStateRecord(
-                    CallbackState.IN_PROGRESS, time.time() - 20, "default", 99999
+                    CallbackState.IN_PROGRESS,
+                    time.time() - 20,
                 )
             return None
 
@@ -110,15 +114,17 @@ def test_body_is_passed_to_hud_unchanged(self):
     def test_queue_time_calculated_from_state_records(self):
         """queue_time is the dispatch-to-in_progress delta."""
         dispatch_record = CallbackStateRecord(
-            CallbackState.DISPATCHED, 1000.0, "dispatch-job", 11111
+            CallbackState.DISPATCHED,
+            1000.0,
         )
-        job_record = CallbackStateRecord(
-            CallbackState.IN_PROGRESS, 1030.0, "default", 99999
+        workflow_record = CallbackStateRecord(
+            CallbackState.IN_PROGRESS,
+            1030.0,
         )
         self.mock_redis.get_callback_state.side_effect = [
             dispatch_record,  # dispatch lookup
-            None,  # job state: not yet set
-            job_record,  # re-read after set_callback_state
+            None,  # workflow state: not yet set
+            workflow_record,  # re-read after set_callback_state
         ]
 
         handle(_cfg(), _body(status="in_progress"), verified_repo="org/repo")
@@ -131,17 +137,20 @@ def test_queue_time_calculated_from_state_records(self):
     def test_execution_time_calculated_from_state_records(self):
         """execution_time is the in_progress-to-completed delta."""
         dispatch_record = CallbackStateRecord(
-            CallbackState.DISPATCHED, 1000.0, "dispatch-job", 11111
+            CallbackState.DISPATCHED,
+            1000.0,
         )
-        job_record = CallbackStateRecord(
-            CallbackState.IN_PROGRESS, 1030.0, "default", 99999
+        workflow_record = CallbackStateRecord(
+            CallbackState.IN_PROGRESS,
+            1030.0,
         )
         completed_record = CallbackStateRecord(
-            CallbackState.COMPLETED, 1060.0, "default", 99999
+            CallbackState.COMPLETED,
+            1060.0,
         )
         self.mock_redis.get_callback_state.side_effect = [
             dispatch_record,  # dispatch lookup
-            job_record,  # job state: in_progress
+            workflow_record,  # workflow state: in_progress
             completed_record,  # re-read after set_callback_state
         ]
 
@@ -199,16 +208,17 @@ def test_redis_error_fetching_dispatch_record_rejected(self):
             handle(_cfg(), _body(status="completed"), verified_repo="org/repo")
         self.assertEqual(ctx.exception.status_code, 400)
 
-    def test_redis_error_fetching_job_record_proceeds(self):
-        """Redis error on job record lookup returns None; callback proceeds."""
+    def test_redis_error_fetching_workflow_record_proceeds(self):
+        """Redis error on workflow record lookup returns None; callback proceeds."""
         dispatch_record = CallbackStateRecord(
-            CallbackState.DISPATCHED, 1000.0, "dispatch-job", 11111
+            CallbackState.DISPATCHED,
+            1000.0,
         )
-        # Three calls: dispatch lookup, job record lookup, re-read after set.
+        # Three calls: dispatch lookup, workflow record lookup, re-read after set.
         self.mock_redis.get_callback_state.side_effect = [
             dispatch_record,
-            None,  # job record lookup returns None (get_callback_state catches RedisError)
-            None,  # updated_job_record re-read → early return with empty metrics
+            None,  # workflow record lookup returns None (get_callback_state catches RedisError)
+            None,  # updated_workflow_record re-read → early return with empty metrics
         ]
 
         handle(_cfg(), _body(status="completed"), verified_repo="org/repo")