diff --git a/tools/caretaker-agent/cloudrun/triage-worker/db/issues_store.py b/tools/caretaker-agent/cloudrun/triage-worker/db/issues_store.py index 7fdc25a416..b9bfe2643e 100644 --- a/tools/caretaker-agent/cloudrun/triage-worker/db/issues_store.py +++ b/tools/caretaker-agent/cloudrun/triage-worker/db/issues_store.py @@ -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: diff --git a/tools/caretaker-agent/cloudrun/triage-worker/main.py b/tools/caretaker-agent/cloudrun/triage-worker/main.py new file mode 100644 index 0000000000..ce5bd1e391 --- /dev/null +++ b/tools/caretaker-agent/cloudrun/triage-worker/main.py @@ -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() diff --git a/tools/caretaker-agent/cloudrun/triage-worker/requirements.txt b/tools/caretaker-agent/cloudrun/triage-worker/requirements.txt index 4a38e337d5..09d97bd1f1 100644 --- a/tools/caretaker-agent/cloudrun/triage-worker/requirements.txt +++ b/tools/caretaker-agent/cloudrun/triage-worker/requirements.txt @@ -1 +1,2 @@ google-cloud-firestore>=2.15.0, <3.0.0 +google-cloud-pubsub \ No newline at end of file diff --git a/tools/caretaker-agent/cloudrun/triage-worker/tests/test_issues_store.py b/tools/caretaker-agent/cloudrun/triage-worker/tests/test_issues_store.py index 6bbc9403d3..6bfed65d49 100644 --- a/tools/caretaker-agent/cloudrun/triage-worker/tests/test_issues_store.py +++ b/tools/caretaker-agent/cloudrun/triage-worker/tests/test_issues_store.py @@ -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): diff --git a/tools/caretaker-agent/cloudrun/triage-worker/tests/test_main.py b/tools/caretaker-agent/cloudrun/triage-worker/tests/test_main.py new file mode 100644 index 0000000000..34c3dd0488 --- /dev/null +++ b/tools/caretaker-agent/cloudrun/triage-worker/tests/test_main.py @@ -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() diff --git a/tools/caretaker-agent/cloudrun/triage-worker/triage_orchestrator.py b/tools/caretaker-agent/cloudrun/triage-worker/triage_orchestrator.py new file mode 100644 index 0000000000..53d53f7cd6 --- /dev/null +++ b/tools/caretaker-agent/cloudrun/triage-worker/triage_orchestrator.py @@ -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" diff --git a/tools/caretaker-agent/cloudrun/triage-worker/utils/egress.py b/tools/caretaker-agent/cloudrun/triage-worker/utils/egress.py new file mode 100644 index 0000000000..67665bf2fe --- /dev/null +++ b/tools/caretaker-agent/cloudrun/triage-worker/utils/egress.py @@ -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, + }, + })