feat(caretaker-triage): implement main worker execution loop and egress action publisher (#28306)

2026-07-10 10:00:53 -07:00 · 2026-07-09 15:51:26 -07:00
parent 9a023bbba0
commit a4c91ce191
7 changed files with 436 additions and 3 deletions
@@ -65,7 +65,7 @@ class IssuesStore:
        
        # Early exit for terminal states
        terminal_states = {
-            "TRIAGED", "LOW_QUALITY", "NEEDS_INFO", "NEEDS_HUMAN"
+            "TRIAGED", "AUTO_CLOSE", "NEEDS_INFO", "NEEDS_HUMAN"
        }
        if current_status in terminal_states:
            return ClaimAction.SKIP
@@ -209,7 +209,7 @@ class IssuesStore:
            success: Whether AI triage completed successfully.
            workable_spec: Parsed workable specification to persist when status
              is TRIAGED.
-            status: Target issue status (TRIAGED, NEEDS_INFO, LOW_QUALITY,
+            status: Target issue status (TRIAGED, NEEDS_INFO, AUTO_CLOSE,
              or NEEDS_HUMAN).

        Returns:
@@ -0,0 +1,152 @@
+import os
+import json
+import base64
+import sys
+
+from google.cloud import firestore
+from triage_orchestrator import process_issue_triage
+from utils.validator import validate_triage_result
+from utils.egress import send_label_action, send_comment_action
+from db.issues_store import IssuesStore, ClaimAction, ReleaseAction
+
+
+def main() -> None:
+    """
+    Orchestrates the Cloud Run Job execution loop for Caretaker Triage.
+
+    Assumptions:
+        - Assumes ISSUE_DETAILS env var contains base64-encoded JSON payload.
+        - Assumes WORKFLOW_EXECUTION_ID contains unique lock holder ID.
+    """
+    # Cloud Run Jobs inject data via environment variables
+    encoded_data = os.environ.get("ISSUE_DETAILS")
+    
+    if not encoded_data:
+        print("[PROD] Error: No data provided in ISSUE_DETAILS.")
+        sys.exit(1)
+
+    try:
+        payload = json.loads(base64.b64decode(encoded_data))
+    except Exception as e:
+        print(f"[PROD] Error decoding payload: {e}")
+        sys.exit(1)
+    
+    try:
+        issue_number = int(payload.get("issue_number"))
+    except (TypeError, ValueError):
+        print("[PROD] Error: issue_number is not a valid number. Exiting.")
+        sys.exit(1)
+
+    try:
+        owner, repo = payload.get("repository", "").split("/")
+        if not owner or not repo:
+            raise ValueError
+    except (TypeError, ValueError):
+        print("[PROD] Error: Malformed repository format (expected 'owner/repo'). Exiting.")
+        sys.exit(1)
+
+    lock_holder = os.environ.get("WORKFLOW_EXECUTION_ID", "local-exec")
+    
+    # Initialize Firestore Client & IssuesStore
+    project_id = os.environ.get("PROJECT_ID")
+    db_id = os.environ.get("FIRESTORE_DATABASE")
+    collection_name = os.environ.get("FIRESTORE_COLLECTION", "issues")
+    
+    db_client = firestore.Client(project=project_id, database=db_id)
+    store = IssuesStore(db_client, collection_name)
+
+    # Claim the lock
+    claim_action = store.acquire_lock(owner, repo, issue_number, lock_holder)
+    
+    if claim_action == ClaimAction.SKIP:
+        print(
+            f"[WORKER] Issue #{issue_number} already handled or active lock "
+            "present. Exiting."
+        )
+        sys.exit(0)
+    elif claim_action == ClaimAction.NEEDS_HUMAN:
+        print(f"[WORKER] Issue #{issue_number} requires human review. Exiting.")
+        sys.exit(0)
+        
+    print(f"[WORKER] Starting triage for issue #{issue_number}...")
+    try:
+        success, raw_output = process_issue_triage(payload)
+    except Exception as e:
+        print(f"[WORKER] Triage process failed with exception: {e}")
+        success, raw_output = False, ""
+    
+    if success:
+        try:
+            triage_result = json.loads(raw_output)
+            validate_triage_result(triage_result)
+
+            quality = triage_result.get("triage_metadata", {}).get("quality")
+            workable_spec = triage_result.get("workable_spec", {})
+            
+            if quality in ["SPAM", "EMPTY", "FEATURE"]:
+                print(f"[WORKER] Quality: {quality}. Applying auto-close label.")
+                send_label_action(owner, repo, issue_number, ["auto-close"])
+                store.release_lock(
+                    owner,
+                    repo,
+                    issue_number,
+                    lock_holder,
+                    success=True,
+                    status="AUTO_CLOSE",
+                )
+                sys.exit(0)
+            elif quality == "NEEDS_INFO":
+                print(f"[WORKER] Quality: NEEDS_INFO. Leaving comment.")
+                comment_body = (
+                    triage_result.get("triage_metadata", {})
+                    .get("comment", "")
+                    .strip()
+                )
+                send_comment_action(owner, repo, issue_number, comment_body)
+                store.release_lock(
+                    owner,
+                    repo,
+                    issue_number,
+                    lock_holder,
+                    success=True,
+                    status="NEEDS_INFO",
+                )
+                sys.exit(0)
+            else:
+                effort = triage_result.get("triage_metadata", {}).get(
+                    "effort_estimate"
+                )
+                print(
+                    f"[WORKER] Quality: OK. Effort: {effort}. Applying "
+                    "effort label."
+                )
+                send_label_action(
+                    owner, repo, issue_number, [f"effort/{effort.lower()}"]
+                )
+                store.release_lock(
+                    owner,
+                    repo,
+                    issue_number,
+                    lock_holder,
+                    success=True,
+                    status="TRIAGED",
+                    workable_spec=workable_spec,
+                )
+                print(f"[WORKER] Triage success.")
+                sys.exit(0)
+                
+        except Exception as e:
+            print(f"[WORKER] Validation failed: {e}")
+            success = False
+    
+    # If an exception happens in json.loads or validate_triage_result 
+    # If LLM inference itself fails inside process_issue_triage
+    if not success:
+        release_action = store.release_lock(
+            owner, repo, issue_number, lock_holder, success=False
+        )
+        sys.exit(1 if release_action == ReleaseAction.RETRY else 0)
+
+
+if __name__ == "__main__":
+    main()
@@ -1 +1,2 @@
 google-cloud-firestore>=2.15.0, <3.0.0
+google-cloud-pubsub
@@ -29,7 +29,7 @@ class TestIssuesStore(unittest.TestCase):

    def test_acquire_lock_terminal_states(self):
        """acquire lock on terminal status docs should skip triage"""
-        terminal_statuses = ["TRIAGED", "LOW_QUALITY", "NEEDS_INFO", "NEEDS_HUMAN"]
+        terminal_statuses = ["TRIAGED", "AUTO_CLOSE", "NEEDS_INFO", "NEEDS_HUMAN"]
        self.snapshot.exists = True
        for status in terminal_statuses:
            with self.subTest(status=status):
@@ -0,0 +1,171 @@
+"""
+Unit tests for main.py execution loop.
+
+Verifies input payload decoding, locking claim actions, LLM quality routing,
+and exit codes for Cloud Run Jobs.
+"""
+
+import unittest
+from unittest.mock import patch, MagicMock
+import os
+import json
+import base64
+
+from main import main
+from db.issues_store import ClaimAction, ReleaseAction
+
+VALID_SPEC = {
+    "issue_id": "owner/repo#42",
+    "summary": {"problem": "p", "root_cause": "r", "context": "c"},
+    "implementation_plan": {
+        "files_to_modify": ["src/app.ts"], "steps": ["Fix bug"]
+    },
+    "testing_strategy": {
+        "test_file": "tests/app.test.ts",
+        "expected_behavior": "Pass",
+        "verification_steps": ["Check"],
+        "framework": "Vitest"
+    }
+}
+
+
+class TestMainExecutionLoop(unittest.TestCase):
+
+    def setUp(self):
+        payload = {
+            "issue_number": 42,
+            "repository": "owner/repo",
+            "title": "Fix crash",
+            "body": "App crashes on startup"
+        }
+        encoded = base64.b64encode(json.dumps(payload).encode("utf-8")).decode()
+
+        self.env_patcher = patch.dict(os.environ, {
+            "ISSUE_DETAILS": encoded,
+            "WORKFLOW_EXECUTION_ID": "exec-123",
+            "PROJECT_ID": "test-project",
+            "EGRESS_TOPIC_ID": "test-topic"
+        })
+        self.env_patcher.start()
+
+        self.mock_store = MagicMock()
+        self.store_patcher = patch(
+            "main.IssuesStore", return_value=self.mock_store
+        )
+        self.store_patcher.start()
+
+        self.db_patcher = patch("main.firestore.Client")
+        self.db_patcher.start()
+
+    def tearDown(self):
+        self.db_patcher.stop()
+        self.store_patcher.stop()
+        self.env_patcher.stop()
+
+    @patch.dict(os.environ, {"ISSUE_DETAILS": ""})
+    def test_main_missing_issue_details_exits_one(self):
+        """Missing ISSUE_DETAILS env var exits 1 to signal container error."""
+        with self.assertRaises(SystemExit) as ctx:
+            main()
+        self.assertEqual(ctx.exception.code, 1)
+
+    def test_main_claim_action_early_exits_zero(self):
+        """SKIP and NEEDS_HUMAN claim actions exit 0 without retrying."""
+        for action in [ClaimAction.SKIP, ClaimAction.NEEDS_HUMAN]:
+            with self.subTest(action=action):
+                self.mock_store.acquire_lock.return_value = action
+                with self.assertRaises(SystemExit) as ctx:
+                    main()
+                self.assertEqual(ctx.exception.code, 0)
+
+    @patch("main.process_issue_triage")
+    @patch("main.send_label_action")
+    def test_main_auto_close_quality_flow(self, mock_send_label, mock_triage):
+        """SPAM/EMPTY/FEATURE issues dispatch auto-close label."""
+        self.mock_store.acquire_lock.return_value = ClaimAction.PROCEED
+        output = json.dumps({"triage_metadata": {"quality": "SPAM"}})
+        mock_triage.return_value = (True, output)
+
+        with self.assertRaises(SystemExit) as ctx:
+            main()
+
+        self.assertEqual(ctx.exception.code, 0)
+        mock_send_label.assert_called_once_with(
+            "owner", "repo", 42, ["auto-close"]
+        )
+        self.mock_store.release_lock.assert_called_once_with(
+            "owner", "repo", 42, "exec-123", success=True, status="AUTO_CLOSE"
+        )
+
+    @patch("main.process_issue_triage")
+    @patch("main.send_comment_action")
+    def test_main_needs_info_quality_flow(
+        self, mock_send_comment, mock_triage
+    ):
+        """NEEDS_INFO issues dispatch comment action and release NEEDS_INFO."""
+        self.mock_store.acquire_lock.return_value = ClaimAction.PROCEED
+        output = json.dumps({
+            "triage_metadata": {
+                "quality": "NEEDS_INFO",
+                "comment": "Please provide logs."
+            }
+        })
+        mock_triage.return_value = (True, output)
+
+        with self.assertRaises(SystemExit) as ctx:
+            main()
+
+        self.assertEqual(ctx.exception.code, 0)
+        mock_send_comment.assert_called_once_with(
+            "owner", "repo", 42, "Please provide logs."
+        )
+        self.mock_store.release_lock.assert_called_once_with(
+            "owner", "repo", 42, "exec-123", success=True, status="NEEDS_INFO"
+        )
+
+    @patch("main.process_issue_triage")
+    @patch("main.send_label_action")
+    def test_main_ok_quality_flow(self, mock_send_label, mock_triage):
+        """OK quality issues dispatch effort label and release TRIAGED spec."""
+        self.mock_store.acquire_lock.return_value = ClaimAction.PROCEED
+        output = json.dumps({
+            "triage_metadata": {"quality": "OK", "effort_estimate": "SMALL"},
+            "workable_spec": VALID_SPEC
+        })
+        mock_triage.return_value = (True, output)
+
+        with self.assertRaises(SystemExit) as ctx:
+            main()
+
+        self.assertEqual(ctx.exception.code, 0)
+        mock_send_label.assert_called_once_with(
+            "owner", "repo", 42, ["effort/small"]
+        )
+        self.mock_store.release_lock.assert_called_once_with(
+            "owner",
+            "repo",
+            42,
+            "exec-123",
+            success=True,
+            status="TRIAGED",
+            workable_spec=VALID_SPEC,
+        )
+
+    @patch("main.process_issue_triage")
+    def test_main_failure_triggers_retry_release(self, mock_triage):
+        """Process/egress failure releases lock success=False and exits 1."""
+        self.mock_store.acquire_lock.return_value = ClaimAction.PROCEED
+        mock_triage.return_value = (False, "LLM failed")
+        self.mock_store.release_lock.return_value = ReleaseAction.RETRY
+
+        with self.assertRaises(SystemExit) as ctx:
+            main()
+
+        self.assertEqual(ctx.exception.code, 1)
+        self.mock_store.release_lock.assert_called_once_with(
+            "owner", "repo", 42, "exec-123", success=False
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,16 @@
+"""
+Handles LLM inference for issue triage.
+(Stubbed implementation for execution loop integration).
+"""
+
+def process_issue_triage(payload: dict) -> tuple[bool, str]:
+    """
+    Stubbed entrypoint for issue triage processing.
+
+    Args:
+        payload: Dictionary containing issue details (issue_number, repository).
+
+    Returns:
+        tuple[bool, str]: (success, raw_output)
+    """
+    return False, "Not implemented: LLM triage orchestrator"
@@ -0,0 +1,93 @@
+import os
+import json
+from typing import TypedDict, Literal, Union
+from google.cloud import pubsub_v1
+
+
+class BasePayload(TypedDict):
+    owner: str
+    repo: str
+    issueNumber: int
+
+
+class LabelPayload(BasePayload):
+    labels: list[str]
+
+
+class CommentPayload(BasePayload):
+    commentBody: str
+
+
+class LabelEvent(TypedDict):
+    action: Literal["LABEL"]
+    payload: LabelPayload
+
+
+class CommentEvent(TypedDict):
+    action: Literal["COMMENT"]
+    payload: CommentPayload
+
+
+EgressEvent = Union[LabelEvent, CommentEvent]
+
+
+def _publish_egress_action(egress_event: EgressEvent) -> None:
+    """
+    [Internal] Publishes an EgressEvent JSON payload to Pub/Sub.
+    """
+    project_id = os.environ.get("PROJECT_ID")
+    egress_topic_id = os.environ.get("EGRESS_TOPIC_ID")
+
+    if not project_id or not egress_topic_id:
+        print(
+            f"[WORKER] Warning: Missing PROJECT_ID ({project_id}) or "
+            f"EGRESS_TOPIC_ID ({egress_topic_id}), skipping egress."
+        )
+        return
+    try:
+        publisher = pubsub_v1.PublisherClient()
+        topic_path = publisher.topic_path(project_id, egress_topic_id)
+        data = json.dumps(egress_event).encode("utf-8")
+        future = publisher.publish(topic_path, data)
+        message_id = future.result()
+        print(
+            f"[WORKER] Published egress action to Pub/Sub ({egress_topic_id}). "
+            f"Message ID: {message_id}"
+        )
+    except Exception as e:
+        print(f"[WORKER] Error publishing to Pub/Sub: {e}")
+        raise
+
+
+def send_label_action(
+    owner: str, repo: str, issue_number: int, labels: list[str]
+) -> None:
+    """
+    Helper to publish a LABEL action to egress.
+    """
+    _publish_egress_action({
+        "action": "LABEL",
+        "payload": {
+            "owner": owner,
+            "repo": repo,
+            "issueNumber": issue_number,
+            "labels": labels,
+        },
+    })
+
+
+def send_comment_action(
+    owner: str, repo: str, issue_number: int, comment_body: str
+) -> None:
+    """
+    Helper to publish a COMMENT action to egress.
+    """
+    _publish_egress_action({
+        "action": "COMMENT",
+        "payload": {
+            "owner": owner,
+            "repo": repo,
+            "issueNumber": issue_number,
+            "commentBody": comment_body,
+        },
+    })