From fccc043bd48874d47c2c19375a9ab923953119a9 Mon Sep 17 00:00:00 2001 From: joneba-google Date: Tue, 28 Jul 2026 19:21:50 +0000 Subject: [PATCH] feat(pr-generator-db): implement Firestore concurrency dual-locking and test ingestion utilities (#28432) --- .../pr-generator/workflow/db/__init__.py | 47 ++ .../pr-generator/workflow/db/db_interface.py | 428 ++++++++++++++++++ 2 files changed, 475 insertions(+) create mode 100644 tools/caretaker-agent/cloudrun/pr-generator/workflow/db/__init__.py create mode 100644 tools/caretaker-agent/cloudrun/pr-generator/workflow/db/db_interface.py diff --git a/tools/caretaker-agent/cloudrun/pr-generator/workflow/db/__init__.py b/tools/caretaker-agent/cloudrun/pr-generator/workflow/db/__init__.py new file mode 100644 index 0000000000..15ffd5c87b --- /dev/null +++ b/tools/caretaker-agent/cloudrun/pr-generator/workflow/db/__init__.py @@ -0,0 +1,47 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Firestore db package for code generation orchestrator.""" + +from .db_interface import ( + ClaimAction, + IssueStatus, + ReleaseAction, + acquire_lock, + create_issue, + get_firestore_client, + get_firestore_id, + get_issue, + get_issue_ref, + mark_needs_human, + mark_pr_created, + release_lock, + update_status, +) + +__all__ = [ + "ClaimAction", + "IssueStatus", + "ReleaseAction", + "acquire_lock", + "create_issue", + "get_firestore_client", + "get_firestore_id", + "get_issue", + "get_issue_ref", + "mark_needs_human", + "mark_pr_created", + "release_lock", + "update_status", +] diff --git a/tools/caretaker-agent/cloudrun/pr-generator/workflow/db/db_interface.py b/tools/caretaker-agent/cloudrun/pr-generator/workflow/db/db_interface.py new file mode 100644 index 0000000000..34b1f8df63 --- /dev/null +++ b/tools/caretaker-agent/cloudrun/pr-generator/workflow/db/db_interface.py @@ -0,0 +1,428 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Firestore database interface for code generation worker and orchestrator. + +Provides helper functions for worker.py and orchestrator.py to interface with +Firestore using the technical writeup specifications: +- Concurrency dual-lock validation (lock.holder, lock.expires_at 15 mins). +- Direct document ID resolution from the FIRESTORE_ID environment variable. +- State transitions (COMMIT_GENERATION, PR_EVALUATION_PENDING, NEEDS_HUMAN, etc.). +""" + +import os +from datetime import datetime, timedelta, timezone +from enum import Enum +from typing import Any +from google.cloud import firestore + + +class IssueStatus(str, Enum): + UNTRIAGED = "UNTRIAGED" + TRIAGING = "TRIAGING" + NEEDS_INFO = "NEEDS_INFO" + TRIAGED = "TRIAGED" + COMMIT_GENERATION = "COMMIT_GENERATION" + PR_EVALUATION_PENDING = "PR_EVALUATION_PENDING" + PR_REVISION = "PR_REVISION" + NEEDS_HUMAN = "NEEDS_HUMAN" + AUTO_CLOSE = "AUTO_CLOSE" + + +class ClaimAction(Enum): + PROCEED = "PROCEED" + SKIP = "SKIP" + NEEDS_HUMAN = "NEEDS_HUMAN" + + +class ReleaseAction(Enum): + COMPLETE = "COMPLETE" # Complete / no retry needed (Exit code 0) + RETRY = "RETRY" # Failed / trigger retry (Exit code 1) + + +PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", os.environ.get("PROJECT_ID")) +DATABASE_NAME = os.environ.get("FIRESTORE_DATABASE") +COLLECTION_NAME = os.environ.get("FIRESTORE_COLLECTION", "issues") + +_db_client: firestore.Client | None = None + + +def get_firestore_client() -> firestore.Client: + """Lazily initializes and returns the Firestore client.""" + global _db_client + if _db_client is None: + if DATABASE_NAME: + _db_client = firestore.Client(project=PROJECT_ID, database=DATABASE_NAME) + else: + _db_client = firestore.Client(project=PROJECT_ID) + return _db_client + + +def get_firestore_id( + doc_id: str | None = None, + owner: str | None = None, + repo: str | None = None, + issue_number: int | str | None = None, +) -> str: + """Resolves the Firestore document ID. + + Prioritizes the FIRESTORE_ID / firestore_id environment variable or explicit doc_id + over reconstructing the document ID from owner/repo/issue_number. + """ + resolved_id = ( + doc_id + or os.environ.get("FIRESTORE_ID") + or os.environ.get("firestore_id") + ) + if resolved_id: + return resolved_id + + if owner and repo and issue_number is not None: + return f"github_{owner}_{repo}_{issue_number}" + + raise ValueError( + "Firestore document ID could not be resolved. Please set the 'FIRESTORE_ID' " + "environment variable or provide explicit doc_id or owner, repo, and issue_number." + ) + + +def get_issue_ref( + owner: str | None = None, + repo: str | None = None, + issue_number: int | str | None = None, + doc_id: str | None = None, +): + """Generates the Firestore DocumentReference for an issue using the resolved document ID.""" + resolved_id = get_firestore_id(doc_id=doc_id, owner=owner, repo=repo, issue_number=issue_number) + return get_firestore_client().collection(COLLECTION_NAME).document(resolved_id) + + +@firestore.transactional +def _create_issue_tx( + transaction, + doc_ref, + owner: str, + repo: str, + issue_number: int, + title: str, + pr_number: str = "", + error: str = "", + doc_id: str | None = None, +) -> bool: + snapshot = doc_ref.get(transaction=transaction) + if not snapshot.exists: + resolved_id = doc_id or doc_ref.id + new_issue = { + "status": IssueStatus.UNTRIAGED.value, + "triage_attempts": 0, + "generation_attempts": 0, + "workable_spec": {}, + "lock": { + "holder": None, + "expires_at": None, + }, + "created_at": firestore.SERVER_TIMESTAMP, + "updated_at": firestore.SERVER_TIMESTAMP, + "github_metadata": { + "owner": owner, + "repo": repo, + "issue_number": issue_number, + "title": title, + "pr_number": pr_number, + }, + "error": error, + } + transaction.set(doc_ref, new_issue) + return True + return False + + +def create_issue( + owner: str, + repo: str, + issue_number: int, + title: str, + pr_number: str = "", + error: str = "", + doc_id: str | None = None, +) -> bool: + """Initializes a new issue document in a transaction.""" + doc_ref = get_issue_ref(owner=owner, repo=repo, issue_number=issue_number, doc_id=doc_id) + transaction = get_firestore_client().transaction() + return _create_issue_tx( + transaction, + doc_ref, + owner, + repo, + issue_number, + title, + pr_number, + error, + doc_id, + ) + + +@firestore.transactional +def _acquire_lock_tx( + transaction, + doc_ref, + lock_holder: str, + lock_duration_sec: int, + target_status: str, +) -> ClaimAction: + """Transactional logic to validate and claim concurrency locks. + + Step 1 & 2 (Lock Validation): + - If lock.expires_at is Null (or expired): no worker claimed, PROCEED. + - If lock.expires_at not elapsed, but lock.holder == current execution_id: crashed instance re-issue, PROCEED. + - Else (active lock held by another workflow): commit transaction with no changes and SKIP. + """ + snapshot = doc_ref.get(transaction=transaction) + if not snapshot.exists: + return ClaimAction.SKIP + + data = snapshot.to_dict() or {} + current_status = data.get("status") + attempts = data.get("generation_attempts", 0) + + # Only allow PR generation to start for TRIAGED issues, recovering COMMIT_GENERATION jobs, or PR_REVISION + allowed_start_states = { + IssueStatus.TRIAGED.value, + IssueStatus.COMMIT_GENERATION.value, + IssueStatus.PR_REVISION.value, # TODO: defensive programming for when PR revision is implemented + } + if current_status not in allowed_start_states: + return ClaimAction.SKIP + + if attempts >= 2: + transaction.update( + doc_ref, + { + "status": IssueStatus.NEEDS_HUMAN.value, + "lock.holder": None, + "lock.expires_at": None, + "updated_at": firestore.SERVER_TIMESTAMP, + }, + ) + return ClaimAction.NEEDS_HUMAN + + lock = data.get("lock") or {} + now = datetime.now(timezone.utc) + holder = lock.get("holder") + expires_at = lock.get("expires_at") + + # Check active lock condition + lock_is_active = (expires_at is not None) and (now <= expires_at) + + # If lock is active and held by another execution_id, exit cleanly + if lock_is_active and holder != lock_holder: + return ClaimAction.SKIP + + # Acquire lock and set status (Step 3: COMMIT_GENERATION) + new_expires_at = now + timedelta(seconds=lock_duration_sec) + new_attempts = attempts + 1 + + transaction.update( + doc_ref, + { + "status": target_status, + "generation_attempts": new_attempts, + "lock.holder": lock_holder, + "lock.expires_at": new_expires_at, + "updated_at": firestore.SERVER_TIMESTAMP, + }, + ) + return ClaimAction.PROCEED + + +def acquire_lock( + lock_holder: str, + doc_id: str | None = None, + owner: str | None = None, + repo: str | None = None, + issue_number: int | str | None = None, + lock_duration_sec: int = 900, # 15 minutes + target_status: str = IssueStatus.COMMIT_GENERATION.value, +) -> ClaimAction: + """Attempts to acquire the processing lock for a Cloud Run workflow execution.""" + doc_ref = get_issue_ref(owner=owner, repo=repo, issue_number=issue_number, doc_id=doc_id) + transaction = get_firestore_client().transaction() + return _acquire_lock_tx( + transaction, + doc_ref, + lock_holder, + lock_duration_sec, + target_status, + ) + + +@firestore.transactional +def _release_lock_tx( + transaction, + doc_ref, + lock_holder: str, + success: bool, + status: str | None = None, + pr_number: str | None = None, + error: str | None = None, + workable_spec: dict[str, Any] | None = None, +) -> ReleaseAction: + """Transactional logic to release the lock and update status.""" + snapshot = doc_ref.get(transaction=transaction) + if not snapshot.exists: + return ReleaseAction.COMPLETE + + data = snapshot.to_dict() or {} + lock = data.get("lock") or {} + + if lock.get("holder") != lock_holder: + return ReleaseAction.COMPLETE + + updates: dict[str, Any] = { + "lock.holder": None, + "lock.expires_at": None, + "updated_at": firestore.SERVER_TIMESTAMP, + } + + if pr_number is not None: + updates["github_metadata.pr_number"] = pr_number + + if error is not None: + updates["error"] = error + + if success: + updates["generation_attempts"] = 0 # Defensive reset for multi-stage runs + if status: + updates["status"] = status + if workable_spec is not None: + updates["workable_spec"] = workable_spec + transaction.update(doc_ref, updates) + return ReleaseAction.COMPLETE + else: + target_status = status if status else IssueStatus.TRIAGED.value + attempts = data.get("generation_attempts", 0) + if attempts < 2: + updates["status"] = target_status + transaction.update(doc_ref, updates) + return ReleaseAction.RETRY + else: + updates["status"] = IssueStatus.NEEDS_HUMAN.value + transaction.update(doc_ref, updates) + return ReleaseAction.COMPLETE + + +def release_lock( + lock_holder: str, + success: bool, + doc_id: str | None = None, + owner: str | None = None, + repo: str | None = None, + issue_number: int | str | None = None, + status: str | None = None, + pr_number: str | None = None, + error: str | None = None, + workable_spec: dict[str, Any] | None = None, +) -> ReleaseAction: + """Releases the processing lock for an issue and updates status.""" + doc_ref = get_issue_ref(owner=owner, repo=repo, issue_number=issue_number, doc_id=doc_id) + transaction = get_firestore_client().transaction() + return _release_lock_tx( + transaction, + doc_ref, + lock_holder, + success, + status, + pr_number, + error, + workable_spec, + ) + + +def mark_pr_created( + lock_holder: str, + pr_number: str, + doc_id: str | None = None, + owner: str | None = None, + repo: str | None = None, + issue_number: int | str | None = None, + status: str = IssueStatus.PR_EVALUATION_PENDING.value, +) -> ReleaseAction: + """Moves issue to PR_EVALUATION_PENDING, records pr_number, and releases lock.""" + return release_lock( + lock_holder=lock_holder, + success=True, + doc_id=doc_id, + owner=owner, + repo=repo, + issue_number=issue_number, + status=status, + pr_number=pr_number, + ) + + +def mark_needs_human( + lock_holder: str, + reason: str, + doc_id: str | None = None, + owner: str | None = None, + repo: str | None = None, + issue_number: int | str | None = None, +) -> ReleaseAction: + """Moves issue to NEEDS_HUMAN, records error, and releases lock.""" + return release_lock( + lock_holder=lock_holder, + success=False, + doc_id=doc_id, + owner=owner, + repo=repo, + issue_number=issue_number, + status=IssueStatus.NEEDS_HUMAN.value, + error=reason, + ) + + +def get_issue( + doc_id: str | None = None, + owner: str | None = None, + repo: str | None = None, + issue_number: int | str | None = None, +) -> dict[str, Any] | None: + """Retrieves an issue document snapshot as a dictionary using FIRESTORE_ID.""" + doc_ref = get_issue_ref(owner=owner, repo=repo, issue_number=issue_number, doc_id=doc_id) + snapshot = doc_ref.get() + if not snapshot.exists: + return None + return snapshot.to_dict() + + +def update_status( + status: str, + doc_id: str | None = None, + owner: str | None = None, + repo: str | None = None, + issue_number: int | str | None = None, + pr_number: str | None = None, + error: str | None = None, +) -> None: + """Updates issue status, PR number, and error message using FIRESTORE_ID.""" + doc_ref = get_issue_ref(owner=owner, repo=repo, issue_number=issue_number, doc_id=doc_id) + updates: dict[str, Any] = { + "status": status, + "updated_at": firestore.SERVER_TIMESTAMP, + } + if pr_number is not None: + updates["github_metadata.pr_number"] = pr_number + if error is not None: + updates["error"] = error + doc_ref.update(updates)