Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 35 additions & 1 deletion src/forge/orchestrator/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,12 +102,46 @@ def __init__(
"""
self.settings = get_settings()
self.consumer_name = consumer_name or f"worker-{uuid.uuid4().hex[:8]}"
self.consumer = QueueConsumer(self.consumer_name)
self.consumer = QueueConsumer(
self.consumer_name,
terminal_failure_handler=self._handle_terminal_failure,
)
self.router = router or create_default_router()
self._shutdown_event = asyncio.Event()
self._checkpointer = None
self._compiled_workflows: dict[str, Any] = {} # Cache compiled workflows by name

async def _handle_terminal_failure(self, message: QueueMessage, error: str) -> None:
"""Post one Jira comment after queue retries are exhausted."""
jira = JiraClient()
event_marker = f"Event/correlation ID: {message.event_id}"
try:
comments = await jira.get_comments(message.ticket_key)
if any(event_marker in comment.body for comment in comments):
logger.info(
f"Terminal failure notification already exists for event {message.event_id}"
)
return

safe_error = redact_secrets(error)
if len(safe_error) > 500:
safe_error = f"{safe_error[:500]}..."
details = (
f"{safe_error}\n\n"
f"Ticket: {message.ticket_key}\n"
f"{event_marker}\n"
"Recovery: inspect the dead-letter entry, resolve the root cause, "
"then requeue the event."
)
await jira.add_error_comment(
issue_key=message.ticket_key,
error_message=details,
node_name="queue execution (retries exhausted)",
)
logger.info(f"Posted terminal queue failure notification to {message.ticket_key}")
finally:
await jira.close()

async def _handle_jira_event(self, message: QueueMessage) -> None:
"""Handle a Jira webhook event.

Expand Down
77 changes: 74 additions & 3 deletions src/forge/queue/consumer.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import asyncio
import logging
import uuid
from collections import defaultdict
from collections.abc import Callable, Coroutine
from typing import Any
Expand All @@ -24,8 +25,15 @@
# How often (seconds) the retry-queue poller wakes up
POLL_INTERVAL_SECONDS = 10

# Terminal-failure callbacks are claimed across workers because every worker
# polls the shared retry sorted set. The short lease recovers if a worker dies
# while posting the notification; completed claims are retained for idempotency.
TERMINAL_FAILURE_CLAIM_PREFIX = "forge:queue:terminal-failure:"
TERMINAL_FAILURE_CLAIM_TTL_SECONDS = 300

# Handler type for message processing
MessageHandler = Callable[[QueueMessage], Coroutine[Any, Any, None]]
TerminalFailureHandler = Callable[[QueueMessage, str], Coroutine[Any, Any, None]]


class QueueConsumer:
Expand All @@ -41,6 +49,7 @@ def __init__(
redis_client: redis.Redis | None = None,
jira_client: JiraClient | None = None,
max_concurrent_tasks: int | None = None,
terminal_failure_handler: TerminalFailureHandler | None = None,
):
"""Initialize the queue consumer.

Expand All @@ -50,6 +59,8 @@ def __init__(
jira_client: Optional Jira client for freshness checks.
max_concurrent_tasks: Maximum concurrent in-flight tasks. Defaults
to ``settings.queue_max_concurrent_tasks`` when not provided.
terminal_failure_handler: Optional callback invoked after a message
is moved to the dead-letter queue.
"""
self.consumer_name = consumer_name
self._redis = redis_client
Expand All @@ -65,6 +76,7 @@ def __init__(
self._semaphore = asyncio.Semaphore(concurrency)
self._active_tasks: set[asyncio.Task[None]] = set()
self._retry_queue = RetryQueue()
self._terminal_failure_handler = terminal_failure_handler

async def _get_redis(self) -> redis.Redis:
"""Get or create Redis client."""
Expand Down Expand Up @@ -134,6 +146,59 @@ async def _ack(self, stream: str, message_id: str) -> None:
redis_client = await self._get_redis()
await redis_client.xack(stream, CONSUMER_GROUP, message_id)

async def _notify_terminal_failure(self, message: QueueMessage, error: str) -> None:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for adding the cross-worker notification claim here. One concern: this claims only the callback; the retry entry itself is still read non-atomically by every worker via get_due_messages(), so two pollers can process the same event concurrently. If one succeeds while another exhausts the counter, we could move a successfully processed event to the DLQ and post a misleading terminal-failure comment; concurrent failures can also advance the retry counter more quickly than intended. Could we atomically claim/lease the retry entry before dispatching it (for example, with a Redis transaction or Lua operation), and add a multi-worker test?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 3f3036b. Due retry entries are now atomically claimed with a Redis Lua operation that advances their score to a 15-minute visibility lease before dispatch. A crashed worker’s entry becomes claimable again after the lease. Added both a concurrent-poller unit test and a real-Redis integration test; the full unit suite and PR CI pass.

"""Invoke the terminal callback once per event without blocking DLQ handling."""
if self._terminal_failure_handler is None:
return

claim_key = f"{TERMINAL_FAILURE_CLAIM_PREFIX}{message.event_id}"
claim_token = uuid.uuid4().hex
try:
redis_client = await self._get_redis()
claimed = await redis_client.set(
claim_key,
claim_token,
ex=TERMINAL_FAILURE_CLAIM_TTL_SECONDS,
nx=True,
)
except Exception as claim_error:
logger.error(
f"Could not claim terminal failure callback for {message.event_id}: {claim_error}"
)
return

if not claimed:
logger.info(f"Terminal failure already handled for event {message.event_id}")
return

try:
await self._terminal_failure_handler(message, error)
except Exception as callback_error:
logger.error(
f"Terminal failure callback failed for {message.event_id}: {callback_error}"
)
return

try:
# A completed claim intentionally has no expiry: DLQ entries are
# retained too, and replaying one must not create another comment.
await redis_client.set(claim_key, "completed")
except Exception as completion_error:
logger.error(
f"Could not complete terminal failure claim for "
f"{message.event_id}: {completion_error}"
)

async def _ack_terminal_message(self, message: QueueMessage, stream: str) -> None:
"""Best-effort acknowledge a message after it has reached the DLQ."""
try:
await self._ack(stream, message.message_id)
except Exception as ack_error:
logger.warning(
f"xack failed for terminal event {message.event_id}; "
f"PEL entry may linger: {ack_error}"
)

async def _process_message(
self,
message: QueueMessage,
Expand Down Expand Up @@ -193,7 +258,8 @@ async def _process_message(
try:
moved_to_dlq = not await self._retry_queue.enqueue_for_retry(message, str(e))
if moved_to_dlq:
await self._ack(stream, message.message_id)
await self._notify_terminal_failure(message, str(e))
await self._ack_terminal_message(message, stream)
except Exception as retry_err:
logger.error(
f"Failed to enqueue {message.event_id} for retry: {retry_err}. "
Expand Down Expand Up @@ -249,7 +315,7 @@ async def _process_retry_queue(self) -> None:
"""
while self._running:
try:
entries = await self._retry_queue.get_due_messages()
entries = await self._retry_queue.claim_due_messages()
for entry in entries:
retry_stream = (
JIRA_STREAM if entry.message.source == EventSource.JIRA else GITHUB_STREAM
Expand All @@ -268,7 +334,12 @@ async def _process_retry_queue(self) -> None:
# the counter keeps accumulating and the message can eventually
# reach the dead-letter queue.
await self._retry_queue.remove_from_retry_without_counter_reset(entry)
await self._retry_queue.enqueue_for_retry(entry.message, str(e))
queued_for_retry = await self._retry_queue.enqueue_for_retry(
entry.message, str(e)
)
if not queued_for_retry:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice job keeping notification failures from interrupting DLQ handling. Could we also acknowledge the original stream entry when queued_for_retry is false? This branch currently notifies and then continues, so terminal failures reached through the retry poller remain indefinitely in the Redis pending-entry list. A best-effort xack, regardless of whether the callback succeeds, plus a test for this path would keep the PEL clean.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 3f3036b. The retry-poller terminal path now performs a best-effort xack after the callback attempt, including when the callback fails. Added focused tests for both callback success and failure; the full unit suite and PR CI pass.

await self._notify_terminal_failure(entry.message, str(e))
await self._ack_terminal_message(entry.message, retry_stream)
continue

# Message processing succeeded — clean up retry state and
Expand Down
32 changes: 24 additions & 8 deletions src/forge/queue/retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,20 @@
RETRY_BACKOFF_MULTIPLIER = 2
INITIAL_RETRY_DELAY_SECONDS = 30
MAX_RETRY_DELAY_SECONDS = 3600 # 1 hour
RETRY_CLAIM_LEASE_SECONDS = 900 # 15 minutes

# Atomically move due entries beyond the visibility window before returning
# them. Other pollers cannot claim the same entries while this lease is active;
# if a worker exits before cleanup, the entries become visible again.
_CLAIM_DUE_MESSAGES_SCRIPT = """
local entries = redis.call(
"ZRANGEBYSCORE", KEYS[1], "-inf", ARGV[1], "LIMIT", 0, ARGV[3]
)
for _, entry in ipairs(entries) do
redis.call("ZADD", KEYS[1], ARGV[2], entry)
end
return entries
"""


@dataclass
Expand Down Expand Up @@ -153,25 +167,27 @@ async def _move_to_dead_letter(
f"{attempt} attempts. Error: {error}"
)

async def get_due_messages(self, limit: int = 10) -> list[RetryEntry]:
"""Get messages that are due for retry.
async def claim_due_messages(self, limit: int = 10) -> list[RetryEntry]:
"""Atomically lease and return messages that are due for retry.
Args:
limit: Maximum number of messages to return.
Returns:
List of retry entries ready to be processed.
List of retry entries exclusively leased to this poller for the
visibility window.
"""
redis = await self._get_redis()
now = datetime.utcnow().timestamp()
lease_until = now + RETRY_CLAIM_LEASE_SECONDS

# Get messages with score <= now
entries = await redis.zrangebyscore(
entries = await redis.eval(
_CLAIM_DUE_MESSAGES_SCRIPT,
1,
RETRY_QUEUE_KEY,
"-inf",
now,
start=0,
num=limit,
lease_until,
limit,
)

results = []
Expand Down
45 changes: 45 additions & 0 deletions tests/integration/redis/test_queue_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,16 @@
"""

import asyncio
import json
from datetime import datetime, timedelta

import pytest

from forge.models.events import EventSource
from forge.queue.consumer import CONSUMER_GROUP, QueueConsumer
from forge.queue.models import QueueMessage
from forge.queue.producer import GITHUB_STREAM, JIRA_STREAM, QueueProducer
from forge.queue.retry import RETRY_QUEUE_KEY, RetryEntry, RetryQueue


@pytest.mark.integration
Expand Down Expand Up @@ -262,6 +265,48 @@ async def handler(message: QueueMessage):
assert len(messages) == 0


@pytest.mark.integration
class TestRetryQueue:
"""Test retry leasing against a real shared Redis instance."""

async def test_due_entry_is_claimed_by_only_one_worker(self, redis_client):
"""Concurrent pollers must not receive the same due retry entry."""
message = QueueMessage(
message_id="1-0",
event_id="lease-test-1",
source=EventSource.JIRA,
event_type="issue_updated",
ticket_key="TEST-LEASE",
)
entry = RetryEntry(
message=message,
attempt=1,
next_retry=datetime.utcnow() - timedelta(seconds=1),
last_error="temporary failure",
)
serialized_entry = json.dumps(entry.to_dict())
await redis_client.zadd(
RETRY_QUEUE_KEY,
{serialized_entry: entry.next_retry.timestamp()},
)

first = RetryQueue()
second = RetryQueue()
first._redis = redis_client
second._redis = redis_client

claims = await asyncio.gather(
first.claim_due_messages(),
second.claim_due_messages(),
)

assert sorted(len(claim) for claim in claims) == [0, 1]
assert sum(claim[0].message.event_id == message.event_id for claim in claims if claim) == 1
lease_score = await redis_client.zscore(RETRY_QUEUE_KEY, serialized_entry)
assert lease_score is not None
assert lease_score > datetime.utcnow().timestamp()


@pytest.mark.integration
class TestQueueMessageSerialization:
"""Test message serialization through the queue."""
Expand Down
51 changes: 51 additions & 0 deletions tests/unit/orchestrator/test_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,57 @@ async def test_report_new_workflow_error_skips_non_reportable_errors(
notify.assert_not_awaited()


@pytest.mark.asyncio
async def test_terminal_failure_posts_sanitized_recovery_comment():
worker = OrchestratorWorker(consumer_name="test-worker")
message = QueueMessage(
message_id="1-0",
event_id="evt-terminal-1",
source=EventSource.JIRA,
event_type="issue_updated",
ticket_key="TEST-123",
)
jira = AsyncMock()
jira.get_comments = AsyncMock(return_value=[])

with patch("forge.orchestrator.worker.JiraClient", return_value=jira):
await worker._handle_terminal_failure(
message,
"clone https://ghp_abcdefghijklmnopqrstuvwxyz123456@github.com/acme/repo failed",
)

jira.add_error_comment.assert_awaited_once()
kwargs = jira.add_error_comment.await_args.kwargs
assert kwargs["issue_key"] == "TEST-123"
assert "[REDACTED]" in kwargs["error_message"]
assert "ghp_" not in kwargs["error_message"]
assert "Event/correlation ID: evt-terminal-1" in kwargs["error_message"]
assert "Recovery:" in kwargs["error_message"]
jira.close.assert_awaited_once()


@pytest.mark.asyncio
async def test_terminal_failure_skips_existing_event_comment():
worker = OrchestratorWorker(consumer_name="test-worker")
message = QueueMessage(
message_id="1-0",
event_id="evt-terminal-1",
source=EventSource.JIRA,
event_type="issue_updated",
ticket_key="TEST-123",
)
jira = AsyncMock()
jira.get_comments = AsyncMock(
return_value=[MagicMock(body="Event/correlation ID: evt-terminal-1")]
)

with patch("forge.orchestrator.worker.JiraClient", return_value=jira):
await worker._handle_terminal_failure(message, "failed")

jira.add_error_comment.assert_not_awaited()
jira.close.assert_awaited_once()


class TestQuestionDetection:
"""Tests for Q&A mode question detection."""

Expand Down
2 changes: 1 addition & 1 deletion tests/unit/queue/test_consumer.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ def _make_consumer(redis_mock: MagicMock, max_tasks: int = 20) -> QueueConsumer:
# the message for retry.
retry_mock = MagicMock(spec=RetryQueue)
retry_mock.enqueue_for_retry = AsyncMock(return_value=True) # queued, not DLQ
retry_mock.get_due_messages = AsyncMock(return_value=[])
retry_mock.claim_due_messages = AsyncMock(return_value=[])
retry_mock.remove_from_retry = AsyncMock()
retry_mock.remove_from_retry_without_counter_reset = AsyncMock()
consumer._retry_queue = retry_mock
Expand Down
Loading
Loading