diff --git a/deploy/helm/aurora/templates/celery-worker-deployment.yaml b/deploy/helm/aurora/templates/celery-worker-deployment.yaml index 566a854d5..9d1a71ef6 100644 --- a/deploy/helm/aurora/templates/celery-worker-deployment.yaml +++ b/deploy/helm/aurora/templates/celery-worker-deployment.yaml @@ -52,7 +52,7 @@ spec: - name: celery-worker image: "{{ .Values.image.registry }}/aurora-server:{{ .Values.image.tag }}" imagePullPolicy: IfNotPresent - command: ["sh", "-c", "celery -A celery_config worker --loglevel=info --concurrency=${CELERY_CONCURRENCY:-4} --pidfile=/tmp/celery-worker.pid"] + command: ["sh", "-c", "exec celery -A celery_config worker --loglevel=info --concurrency=${CELERY_CONCURRENCY:-4} --pidfile=/tmp/celery-worker.pid"] securityContext: {{- include "aurora.containerSecurityContext" (dict "service" "celeryWorker" "global" $ "defaults" (dict)) | nindent 12 }} envFrom: diff --git a/server/celery_config.py b/server/celery_config.py index 31385c5b1..30480674a 100644 --- a/server/celery_config.py +++ b/server/celery_config.py @@ -245,3 +245,61 @@ if hasattr(celery_app, 'tasks'): non_celery_tasks = [t for t in celery_app.tasks.keys() if not t.startswith('celery.')] logging.info("Registered %d custom tasks: %s", len(non_celery_tasks), non_celery_tasks) + + +# --------------------------------------------------------------------------- +# Worker pre-warming: heavy singletons are initialized once per child process +# so the first task doesn't pay a cold-start penalty. +# --------------------------------------------------------------------------- +import threading + +try: + from celery.signals import worker_process_init +except (ImportError, ModuleNotFoundError): + # Tests stub celery as MagicMock when the package isn't installed (see + # tests/conftest.py), so celery.signals isn't importable outside workers. + worker_process_init = None + +_prewarm_ready = threading.Event() + +_prewarm_logger = logging.getLogger("celery.prewarm") + + +if worker_process_init is not None: + + @worker_process_init.connect + def _prewarm_worker(**kwargs): + """Kick off singleton init in a background thread. + + worker_process_init has a ~4s timeout before the parent kills the child, + so we can't block here. Task code calls _prewarm_ready.wait() instead. + """ + + def _do_prewarm(): + try: + from guardrails.input_rail import _ensure_rails_in_thread + _ensure_rails_in_thread() + _prewarm_logger.info("[PREWARM] NeMo Guardrails ready") + except Exception as e: + _prewarm_logger.warning("[PREWARM] Guardrails init failed: %s", e) + + try: + from chat.backend.agent.tools.mcp_preloader import start_mcp_preloader + start_mcp_preloader() + _prewarm_logger.info("[PREWARM] MCP Preloader started") + except Exception as e: + _prewarm_logger.warning("[PREWARM] MCP Preloader failed: %s", e) + + try: + from chat.background.task import _get_worker_agent + _get_worker_agent() + _prewarm_logger.info("[PREWARM] Agent singleton ready") + except Exception as e: + _prewarm_logger.warning("[PREWARM] Agent singleton failed: %s", e) + + _prewarm_ready.set() + + threading.Thread(target=_do_prewarm, name="celery-prewarm", daemon=True).start() +else: + # No worker signal hook (e.g. pytest stubs celery) — don't block task code. + _prewarm_ready.set() diff --git a/server/chat/backend/agent/agent.py b/server/chat/backend/agent/agent.py index bf273dbe1..9d71eef19 100644 --- a/server/chat/backend/agent/agent.py +++ b/server/chat/backend/agent/agent.py @@ -338,21 +338,23 @@ async def agentic_tool_flow( getattr(state, 'attachments', []), ) - # Build modular segments - segments = build_prompt_segments( - provider_preference=provider_preference, - mode=getattr(state, "mode", None), - has_zip_reference=has_zip_ref, - state=state, - ) + # Build prompt segments in background while getting tools on main thread + # (get_cloud_tools needs thread-local context for tool_capture/user resolution) + from concurrent.futures import ThreadPoolExecutor as _TPE + with _TPE(max_workers=1) as _pool: + _prompt_future = _pool.submit( + build_prompt_segments, + provider_preference=provider_preference, + mode=getattr(state, "mode", None), + has_zip_reference=has_zip_ref, + state=state, + ) + tools = get_cloud_tools() + segments = _prompt_future.result() - # Assemble final system prompt from segments system_prompt_text = assemble_system_prompt(segments) if system_prompt_override is not None: system_prompt_text = system_prompt_override - - # Get cloud tools - tools = get_cloud_tools() if tool_subset is not None: tools = tool_subset @@ -854,6 +856,43 @@ def on_llm_end(self, response, run_id=None, **kwargs): # Retry logic for network errors for attempt in range(3): try: + # --- Await the guardrails task that was fired concurrently --- + from chat.backend.agent.workflow import _guardrail_task_var + _pending_rail = _guardrail_task_var.get() + if _pending_rail is not None: + from guardrails.input_rail import InputRailResult, _FAIL_CLOSED_REASON + try: + rail_result: InputRailResult = await _pending_rail + except asyncio.CancelledError: + raise + except Exception as rail_exc: + logging.warning( + "Input rail task failed for session %s: %s", + state.session_id, + rail_exc, + ) + rail_result = InputRailResult( + blocked=True, + reason=_FAIL_CLOSED_REASON, + ) + finally: + # Clear the context var so it doesn't leak to the next turn + _guardrail_task_var.set(None) + + if rail_result.blocked: + from utils.security.audit_events import emit_block_event + emit_block_event( + user_id=state.user_id or "", + session_id=state.session_id or "", + layer="input_rail", + tool="agentic_tool_flow", + subject=getattr(state, "question", ""), + reason=rail_result.reason, + latency_ms=rail_result.latency_ms, + ) + state.guardrail_blocked = True + return state + # Use astream_events for token-by-token streaming logging.info(f"Starting agent token streaming for session {state.session_id}") result = None diff --git a/server/chat/backend/agent/tools/cloud_tools.py b/server/chat/backend/agent/tools/cloud_tools.py index e20365cc8..f44e15cdf 100644 --- a/server/chat/backend/agent/tools/cloud_tools.py +++ b/server/chat/backend/agent/tools/cloud_tools.py @@ -1021,19 +1021,16 @@ def get_cloud_tools(): mode = get_mode_from_context() mode_suffix = (mode or 'agent').lower() - # Create a cache key that accurately reflects the *specific* tool_capture instance (or lack thereof) - # - When no tool_capture is active we can safely cache per-user - # - When a tool_capture **is** active we additionally key on the `id()` of the object so each - # session gets its own wrapped functions that close over the *right* capture instance. + # Cache key uses stable identifiers — the tool_capture wrapper resolves the active + # capture instance dynamically at call time via get_tool_capture(), so we don't need + # id(tool_capture) which caused a cache miss every invocation. rca_flag = getattr(state_context, 'trigger_rca_requested', False) if state_context else False is_background = getattr(state_context, 'is_background', False) if state_context else False is_postmortem_action = getattr(state_context, 'is_postmortem_action', False) if state_context else False is_pr_review = getattr(state_context, 'is_pr_review', False) if state_context else False is_rca_context = _is_background_rca(state_context, is_background) - if tool_capture is None: - cache_key = f"{user_id}:nocapture:{mode_suffix}:background={is_background}:rca={rca_flag}:postmortem={is_postmortem_action}:is_rca_ctx={is_rca_context}:pr_review={is_pr_review}" - else: - cache_key = f"{user_id}:capture:{id(tool_capture)}:{mode_suffix}:background={is_background}:rca={rca_flag}:postmortem={is_postmortem_action}:is_rca_ctx={is_rca_context}:pr_review={is_pr_review}" + capture_tag = "capture" if tool_capture else "nocapture" + cache_key = f"{user_id}:{capture_tag}:{mode_suffix}:background={is_background}:rca={rca_flag}:postmortem={is_postmortem_action}:is_rca_ctx={is_rca_context}:pr_review={is_pr_review}" current_time = time.time() if ( @@ -1068,7 +1065,9 @@ def wrap_func_with_capture(func, tool_name): @wraps(func) def wrapped_func(**kwargs): - exec_lock = getattr(tool_capture, 'execution_lock', None) + # Resolve capture instance dynamically so cached tools work across sessions + tool_capture = get_tool_capture() + exec_lock = getattr(tool_capture, 'execution_lock', None) if tool_capture else None acquired = False try: if exec_lock: @@ -1118,7 +1117,9 @@ def wrapped_func(**kwargs): # The 'completed' flag just prevents cleanup race conditions logging.info(f"TOOL CAPTURE: Tool capture instance found: {tool_capture}") logging.info(f"TOOL CAPTURE: Tool capture instance found: {tool_capture.current_tool_calls}") - for call_id, call_info in tool_capture.current_tool_calls.items(): + with tool_capture.lock: + calls_snapshot = list(tool_capture.current_tool_calls.items()) + for call_id, call_info in calls_snapshot: logging.info(f"TOOL CAPTURE: Call info: {call_info}") logging.info(f"TOOL CAPTURE: Call signature right side: {tool_signature}") logging.info(f"TOOL CAPTURE: Call result: {result}") @@ -1129,7 +1130,9 @@ def wrapped_func(**kwargs): # Fallback – match by tool_name + command (provider may be missing) if not matching_tool_call_id: - for call_id, call_info in tool_capture.current_tool_calls.items(): + with tool_capture.lock: + calls_snapshot = list(tool_capture.current_tool_calls.items()) + for call_id, call_info in calls_snapshot: if call_info.get('tool_name') != tool_name: continue ci_input = call_info.get('input', {}) or {} @@ -1147,18 +1150,21 @@ def wrapped_func(**kwargs): # As a last resort, match by oldest incomplete call for sequential execution (OpenAI) # For parallel execution (Gemini), signature matching should have already succeeded if not matching_tool_call_id: - candidate_ids = [ - (call_id, call_info.get('start_time')) - for call_id, call_info in tool_capture.current_tool_calls.items() - if call_info.get('tool_name') == tool_name and not call_info.get('completed') - ] + with tool_capture.lock: + candidate_ids = [ + (call_id, call_info.get('start_time')) + for call_id, call_info in tool_capture.current_tool_calls.items() + if call_info.get('tool_name') == tool_name and not call_info.get('completed') + ] if len(candidate_ids) == 1: # Only one candidate - safe to use matching_tool_call_id = candidate_ids[0][0] - call_info = tool_capture.current_tool_calls[matching_tool_call_id] - call_info['input'] = signature_payload - call_info['signature'] = tool_signature + with tool_capture.lock: + call_info = tool_capture.current_tool_calls.get(matching_tool_call_id) + if call_info: + call_info['input'] = signature_payload + call_info['signature'] = tool_signature logging.info( "Matched tool call by single incomplete candidate: %s (updated signature)", matching_tool_call_id, @@ -1168,9 +1174,11 @@ def wrapped_func(**kwargs): # This handles OpenAI's sequential execution pattern candidate_ids.sort(key=lambda x: x[1] if x[1] else datetime.min) matching_tool_call_id = candidate_ids[0][0] - call_info = tool_capture.current_tool_calls[matching_tool_call_id] - call_info['input'] = signature_payload - call_info['signature'] = tool_signature + with tool_capture.lock: + call_info = tool_capture.current_tool_calls.get(matching_tool_call_id) + if call_info: + call_info['input'] = signature_payload + call_info['signature'] = tool_signature logging.warning( f"SEQUENTIAL FALLBACK: Found {len(candidate_ids)} incomplete {tool_name} calls, " f"matched to oldest: {matching_tool_call_id}. " @@ -1228,7 +1236,9 @@ def wrapped_func(**kwargs): serialized_payload = str(signature_payload) tool_signature = f"{tool_name}_{serialized_payload}" - for call_id, call_info in tool_capture.current_tool_calls.items(): + with tool_capture.lock: + calls_snapshot = list(tool_capture.current_tool_calls.items()) + for call_id, call_info in calls_snapshot: if call_info.get('signature') == tool_signature and not call_info.get('completed'): matching_tool_call_id = call_id logging.info(f"Matched error to tool call by signature: {matching_tool_call_id}") @@ -1237,10 +1247,11 @@ def wrapped_func(**kwargs): # Fallback: Only match if there's exactly ONE incomplete call for this tool # This prevents parallel tool calls from sharing the same error tracking if not matching_tool_call_id: - incomplete_calls = [ - call_id for call_id, call_info in tool_capture.current_tool_calls.items() - if call_info.get('tool_name') == tool_name and not call_info.get('completed', False) - ] + with tool_capture.lock: + incomplete_calls = [ + call_id for call_id, call_info in tool_capture.current_tool_calls.items() + if call_info.get('tool_name') == tool_name and not call_info.get('completed', False) + ] if len(incomplete_calls) == 1: matching_tool_call_id = incomplete_calls[0] diff --git a/server/chat/backend/agent/workflow.py b/server/chat/backend/agent/workflow.py index 8ea31f8c8..c2f0ba173 100644 --- a/server/chat/backend/agent/workflow.py +++ b/server/chat/backend/agent/workflow.py @@ -21,6 +21,13 @@ logger = logging.getLogger(__name__) +# Async-safe context variable: holds the pending guardrails check task so +# agentic_tool_flow can await it concurrently with its own setup work. +import contextvars +_guardrail_task_var: contextvars.ContextVar[Optional[asyncio.Task]] = contextvars.ContextVar( + '_guardrail_task_var', default=None +) + RCA_SUMMARY_PREFIX = "[RCA Investigation Summary" _USER_MESSAGE_RE = re.compile(r'\s*([\s\S]*?)\s*') @@ -1044,9 +1051,12 @@ async def stream(self, input_state: State): history_prefix_len = len(input_state.messages) - new_turn_input_count self._history_prefix_len = history_prefix_len - # --- Input rail: check user message for prompt injection --- - from guardrails.input_rail import check_input, InputRailResult + # --- Input rail: fire check concurrently with persistence + LangGraph setup --- + from guardrails.input_rail import check_input last_msg = input_state.messages[-1] if input_state.messages else None + msg_text: Optional[str] = None + is_scaffold = False + if last_msg and hasattr(last_msg, "type") and last_msg.type == "human": # Skip persistence for scaffold messages (background prompts, not user input) is_scaffold = getattr(last_msg, 'additional_kwargs', {}).get('is_rca_scaffold', False) @@ -1058,52 +1068,16 @@ async def stream(self, input_state: State): getattr(input_state, "question", None), last_msg.content, ) - # Skip rail when there is no untrusted text to evaluate (e.g. - # prediscovery prompts are entirely system-authored), or for PR - # change-gating reviews (PR title/body is GitHub API data, not - # user input to Aurora; injection defense is handled by - # _escape_prompt_data in the prompt builder). + # Skip guardrails for PR change-gating reviews (PR title/body is GitHub API + # data, not user input to Aurora). Otherwise fire concurrently — the result + # is awaited in agent.py (agentic_tool_flow) right before the LLM call. is_pr_review = getattr(input_state, "is_pr_review", False) - if not msg_text or is_pr_review: - rail_result = InputRailResult(blocked=False) - else: - rail_result = await check_input(msg_text) - if rail_result.blocked: - emit_block_event( - user_id=getattr(input_state, "user_id", "") or "", - session_id=getattr(input_state, "session_id", "") or "", - layer="input_rail", - tool="workflow", - subject=msg_text, - reason=rail_result.reason, - latency_ms=rail_result.latency_ms, - ) - from guardrails.input_rail import _BLOCKED_REASON, _FAIL_CLOSED_AUTH, _FAIL_CLOSED_CONNECTIVITY - _RAIL_USER_MESSAGES = { - _BLOCKED_REASON: "Your message was blocked by our safety system. Please rephrase your request.", - _FAIL_CLOSED_AUTH: "There is an issue with the AI service configuration. Please try again later.", - _FAIL_CLOSED_CONNECTIVITY: "The AI service is temporarily unavailable. Please try again in a moment.", - } - # Background chats have no interactive user: hard block stays. - # Foreground chats that were genuinely blocked: taint the session - # so every subsequent tool call goes through the command gate. - if getattr(input_state, "is_background", False) or rail_result.reason != _BLOCKED_REASON: - input_state.guardrail_blocked = True - yield ("token", _RAIL_USER_MESSAGES.get(rail_result.reason, "Something went wrong. Please try again.")) - return - from utils.auth.command_gate import mark_session_tainted - mark_session_tainted( - getattr(input_state, "session_id", None), - getattr(input_state, "user_id", None), - ) + if msg_text and not is_pr_review: + _guardrail_task_var.set(asyncio.create_task(check_input(msg_text))) - # Rail passed: NOW it's safe to persist the user message. - # Kept inside the rail gate so blocked messages never touch - # chat_sessions.messages (which legacy migration rehydrates into - # llm_context_history on the next turn). if input_state.session_id and input_state.user_id and not is_scaffold: from chat.backend.agent.utils.immediate_save_handler import handle_immediate_save - handle_immediate_save(input_state.session_id, input_state.user_id, msg_text) + handle_immediate_save(input_state.session_id, input_state.user_id, msg_text or last_msg.content) # Log initial state logger.info(f"Starting workflow with session_id={input_state.session_id}, user_id={input_state.user_id}") @@ -1339,6 +1313,7 @@ async def stream(self, input_state: State): self._last_state.update(stream_data) logger.info(f"[WORKFLOW STREAM] Completed: {_event_count} events, {_token_count} tokens streamed") + except Exception as stream_exception: logger.exception(f"[WORKFLOW STREAM ERROR] Exception in workflow stream for session {input_state.session_id}: {stream_exception}") raise diff --git a/server/chat/background/task.py b/server/chat/background/task.py index 5258ea839..215f57db3 100644 --- a/server/chat/background/task.py +++ b/server/chat/background/task.py @@ -9,6 +9,7 @@ import json import logging import os +import threading import uuid from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Optional @@ -28,6 +29,38 @@ logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Per-worker Agent singleton: avoids recreating PostgreSQLClient, WeaviateClient, +# and LLMManager on every background chat task (~2s cold-start savings). +# --------------------------------------------------------------------------- +_worker_agent = None +_worker_agent_lock = threading.Lock() + + +def _get_worker_agent(): + """Return a cached Agent instance for the current worker process.""" + global _worker_agent + if _worker_agent is None: + with _worker_agent_lock: + # Double-checked locking: prewarm thread and task thread can race here + if _worker_agent is None: + from chat.backend.agent.agent import Agent + from chat.backend.agent.db import PostgreSQLClient + from chat.backend.agent.weaviate_client import WeaviateClient + + pg = PostgreSQLClient() + wv = WeaviateClient(pg) + _worker_agent = Agent( + weaviate_client=wv, + postgres_client=pg, + websocket_sender=None, + event_loop=None, + ctx_len=15, + ) + logger.info("[BackgroundChat] Created per-worker singleton Agent") + return _worker_agent + + def _resolve_permitted_tools(user_id: str) -> Optional[set]: """Resolve permitted tools for background chats. Always fetches fresh from DB.""" try: @@ -445,10 +478,28 @@ def run_background_chat( """ from celery.exceptions import SoftTimeLimitExceeded + # Deduplication guard: acks_late=True can cause Redis to redeliver the same task + # to the worker while it's still running. Use a Redis SETNX lock on the task ID. + _dedup_key = f"celery:dedup:{self.request.id}" + _dedup_redis = None + _dedup_acquired = False + try: + _dedup_redis = celery_app.backend.client + _acquired = _dedup_redis.set(_dedup_key, "1", nx=True, ex=1800) + if not _acquired: + logger.warning(f"[BackgroundChat] DEDUP: Task {self.request.id} already running, skipping duplicate execution") + return {"session_id": session_id, "status": "deduplicated", "error": None} + _dedup_acquired = True + except Exception as _dedup_err: + logger.warning(f"[BackgroundChat] DEDUP check failed (proceeding anyway): {_dedup_err}") + + # Block until prewarm completes (avoids cold-start on first task after child fork) + from celery_config import _prewarm_ready + _prewarm_ready.wait(timeout=30) + logger.info(f"[BackgroundChat] Starting for user {user_id}, session {session_id}") logger.info(f"[BackgroundChat] Trigger: {trigger_metadata}") - # Eagerly persist the initial user message so it's visible in the UI immediately (opt-in) if os.environ.get("DISPLAY__RCA_USER_MSG", "").lower() in ("1", "true", "yes"): try: @@ -775,7 +826,8 @@ def run_background_chat( logger.error(f"[BackgroundChat] Failed to generate final visualization: {e}") # Send response back to Slack if this was triggered from Slack - if trigger_metadata and trigger_metadata.get('source') in ['slack', 'slack_button']: + # Skip if already sent inside _execute_background_chat (early send for lower latency) + if trigger_metadata and trigger_metadata.get('source') in ['slack', 'slack_button'] and not result.get('slack_sent_early'): try: _send_response_to_slack(user_id, session_id, trigger_metadata) except Exception as e: @@ -808,6 +860,7 @@ def run_background_chat( logger.debug("[BackgroundChat] Failed to dispatch after_rca actions") logger.info(f"[BackgroundChat] Completed for session {session_id}") + return result except SoftTimeLimitExceeded: @@ -869,6 +922,16 @@ def run_background_chat( } finally: + # Release dedup lock so Celery retries can re-acquire after failure + if _dedup_acquired and _dedup_redis is not None: + try: + _dedup_redis.delete(_dedup_key) + except Exception as dedup_cleanup_err: + logger.debug( + "[BackgroundChat] Failed to release dedup key %s: %s", + _dedup_key, + dedup_cleanup_err, + ) # Safety net: ensure session is never left in in_progress state if not completed_successfully: try: @@ -1257,32 +1320,18 @@ async def _execute_background_chat( from chat.background.background_websocket import BackgroundWebSocket from main_chatbot import process_workflow_async - weaviate_client = None - try: - # Initialize clients (same as handle_connection in main_chatbot.py) - postgres_client = PostgreSQLClient() - weaviate_client = WeaviateClient(postgres_client) + + # Reuse a per-worker Agent to avoid recreating PostgreSQLClient, WeaviateClient, + # and LLMManager on every task (~2s cold-start savings). + agent = _get_worker_agent() # Create background websocket (no-op, just discards messages) background_ws = BackgroundWebSocket() - # Create agent WITHOUT websocket_sender - tools will skip WebSocket messages - # Use reasonable ctx_len for RCAs - need enough history to build on previous tool calls - # But not too high to avoid context length errors (Azure has 128K limit) - # 15 is a good balance - allows agent to see its investigation progress while staying within limits - agent = Agent( - weaviate_client=weaviate_client, - postgres_client=postgres_client, - websocket_sender=None, - event_loop=None, - ctx_len=15, # Reasonable history for RCAs - allows agent to see investigation progress - ) - logger.info(f"[BackgroundChat] Created agent with ctx_len=15 (no WebSocket)") - # Create workflow for this session wf = Workflow(agent, session_id) - logger.info(f"[BackgroundChat] Created workflow for session {session_id}") + logger.info("[BackgroundChat] Created agent + workflow (singleton reuse)") # Build RCA context for system prompt (NOT added to user message) rca_context = _build_rca_context( @@ -1409,6 +1458,14 @@ async def _execute_background_chat( if hasattr(wf, '_wait_for_ongoing_tool_calls'): await wf._wait_for_ongoing_tool_calls() + # Send Slack response immediately — don't wait for post-processing + _slack_early_sent = False + if trigger_metadata and trigger_metadata.get('source') in ['slack', 'slack_button']: + try: + _slack_early_sent = _send_response_to_slack(user_id, session_id, trigger_metadata) + except Exception: + logger.exception("[BackgroundChat] Failed early Slack send") + # --- Phase 2: Jira action --- # Investigation is done. Now deterministically file in Jira. if rca_context and rca_context.get('integrations', {}).get('jira') \ @@ -1524,6 +1581,7 @@ async def _execute_background_chat( "tool_calls": tool_calls, "guardrail_blocked": guardrail_blocked, "after_rca_dispatched": True, + "slack_sent_early": _slack_early_sent, "action_notification_sent": action_notification_sent, } @@ -1539,13 +1597,6 @@ async def _execute_background_chat( await ContextManager._instance.async_queue.stop() except Exception as e: logger.error(f"[BackgroundChat] Failed to stop async save queue - potential resource leak: {e}") - - # Clean up weaviate client - if weaviate_client: - try: - weaviate_client.close() - except Exception as e: - logger.error(f"[BackgroundChat] Failed to close weaviate client - potential connection leak: {e}") TERMINAL_SESSION_STATUSES = frozenset({"completed", "failed", "cancelled"}) @@ -1916,8 +1967,11 @@ def _update_incident_status(incident_id: str, status: str, user_id: str) -> None logger.error(f"[BackgroundChat] Failed to update incident {incident_id} status to '{status}': {e}") -def _send_response_to_slack(user_id: str, session_id: str, trigger_metadata: Dict[str, Any]) -> None: - """Send Aurora's response back to the Slack channel after background chat completes.""" +def _send_response_to_slack(user_id: str, session_id: str, trigger_metadata: Dict[str, Any]) -> bool: + """Send Aurora's response back to the Slack channel after background chat completes. + + Returns True if a message was actually posted, False otherwise. + """ try: from connectors.slack_connector.client import get_slack_client_for_user from routes.slack.slack_events_helpers import format_response_for_slack @@ -1929,7 +1983,7 @@ def _send_response_to_slack(user_id: str, session_id: str, trigger_metadata: Dic if not channel: logger.warning(f"[BackgroundChat] No Slack channel in trigger_metadata for session {session_id}") - return + return False # Get the last assistant message from the chat session with db_pool.get_admin_connection() as conn: @@ -1943,7 +1997,7 @@ def _send_response_to_slack(user_id: str, session_id: str, trigger_metadata: Dic if not row or not row[0]: logger.warning(f"[BackgroundChat] No messages found in session {session_id}") - return + return False messages = row[0] if isinstance(messages, str): @@ -1959,7 +2013,7 @@ def _send_response_to_slack(user_id: str, session_id: str, trigger_metadata: Dic if not last_assistant_message: logger.warning(f"[BackgroundChat] No assistant message found in session {session_id}") - return + return False # Format the response for Slack (markdown conversion, length limits, etc.) formatted_message = format_response_for_slack(last_assistant_message) @@ -2019,7 +2073,7 @@ def _send_response_to_slack(user_id: str, session_id: str, trigger_metadata: Dic client = get_slack_client_for_user(user_id) if not client: logger.error(f"[BackgroundChat] Could not get Slack client for user {user_id}") - return + return False # Slack messages have a ~3000 char limit for update_message SLACK_MSG_LIMIT = 2900 @@ -2041,6 +2095,8 @@ def _send_response_to_slack(user_id: str, session_id: str, trigger_metadata: Dic thread_ts=thread_ts ) + return True + except Exception as e: logger.error(f"[BackgroundChat] Error sending response to Slack: {e}", exc_info=True) raise diff --git a/server/guardrails/input_rail.py b/server/guardrails/input_rail.py index 227a388d3..62b549c47 100644 --- a/server/guardrails/input_rail.py +++ b/server/guardrails/input_rail.py @@ -22,6 +22,7 @@ import asyncio import logging +import threading import time from dataclasses import dataclass @@ -33,7 +34,8 @@ _rails_instance = None _rails_lock: asyncio.Lock | None = None -_last_init_failure_ts: float = 0.0 +_rails_thread_lock = threading.Lock() +_last_init_failure_ts: float = float("-inf") _INIT_FAILURE_BACKOFF_S = 30.0 _FAIL_CLOSED_REASON = "input rail unavailable" @@ -50,6 +52,12 @@ def _get_lock() -> asyncio.Lock: return _rails_lock +def _record_init_failure() -> None: + """Stamp failure time so concurrent callers back off briefly.""" + global _last_init_failure_ts + _last_init_failure_ts = time.monotonic() + + @dataclass(frozen=True) class InputRailResult: blocked: bool @@ -172,6 +180,28 @@ def _build_rails_sync(): return LLMRails(config=rails_config, llm=_build_llm()) +def _ensure_rails_in_thread(): + """Build or return cached rails under ``_rails_thread_lock``. + + Runs in a worker thread (``asyncio.to_thread`` or Celery prewarm) so the + event loop never blocks on ``threading.Lock``. Backoff is checked here, not + only before lock acquisition, so a caller cannot pass an outer check, wait + on the lock, and rebuild after a concurrent failure stamped backoff. + """ + global _rails_instance + with _rails_thread_lock: + if _rails_instance is not None: + return _rails_instance + if time.monotonic() - _last_init_failure_ts < _INIT_FAILURE_BACKOFF_S: + raise RuntimeError("input rail init recently failed; backing off") + try: + _rails_instance = _build_rails_sync() + except Exception: + _record_init_failure() + raise + return _rails_instance + + async def _get_rails(): """Lazily build and cache the NeMo LLMRails instance. @@ -179,22 +209,20 @@ async def _get_rails(): the event loop. Failures are negative-cached for a short window so a flapping provider does not block the loop on every request. """ - global _rails_instance, _last_init_failure_ts + global _rails_instance if _rails_instance is not None: return _rails_instance + # Fast-fail without spawning a worker thread when backoff is active. if time.monotonic() - _last_init_failure_ts < _INIT_FAILURE_BACKOFF_S: raise RuntimeError("input rail init recently failed; backing off") async with _get_lock(): if _rails_instance is not None: return _rails_instance - try: - _rails_instance = await asyncio.to_thread(_build_rails_sync) - except Exception: - _last_init_failure_ts = time.monotonic() - raise - return _rails_instance + if time.monotonic() - _last_init_failure_ts < _INIT_FAILURE_BACKOFF_S: + raise RuntimeError("input rail init recently failed; backing off") + return await asyncio.to_thread(_ensure_rails_in_thread) def _triggered_rail_name(result) -> str: diff --git a/server/routes/slack/slack_events.py b/server/routes/slack/slack_events.py index 7a9c6870c..6ba2bdb47 100644 --- a/server/routes/slack/slack_events.py +++ b/server/routes/slack/slack_events.py @@ -120,18 +120,10 @@ def slack_events(): if client: if response_text: try: - sent_msg = client.send_message( - channel=channel, - text=response_text, - thread_ts=thread_ts - ) - - if trigger_background and sent_msg: - # Proceed with background task - # Determine session logic - msg_thread_ts = event.get('thread_ts') # Use original thread_ts for logic + if trigger_background: + # Determine session logic early to know if we need channel context + msg_thread_ts = event.get('thread_ts') - # Logic for new session vs thread incident_id = None session_id = None context_messages = [] @@ -139,18 +131,43 @@ def slack_events(): final_thread_ts = None if msg_thread_ts and msg_thread_ts != ts: - # Reply in thread - session_id, incident_id = get_session_from_thread(user_id, channel, msg_thread_ts) - context_messages = get_thread_messages(client, channel, msg_thread_ts) - final_thread_ts = msg_thread_ts + # Reply in thread — send "Thinking..." then fetch thread context + sent_msg = client.send_message( + channel=channel, + text=response_text, + thread_ts=thread_ts + ) + session_id, incident_id = get_session_from_thread(user_id, channel, msg_thread_ts) + context_messages = get_thread_messages(client, channel, msg_thread_ts) + final_thread_ts = msg_thread_ts else: - # New top level - fetch channel context with thread summaries - channel_context = get_channel_context_with_threads(client, channel, limit=5) - final_thread_ts = ts - + # New top-level message — parallelize "Thinking..." send with channel context fetch + from concurrent.futures import ThreadPoolExecutor + with ThreadPoolExecutor(max_workers=2) as pool: + thinking_future = pool.submit( + client.send_message, + channel=channel, + text=response_text, + thread_ts=thread_ts + ) + context_future = pool.submit( + get_channel_context_with_threads, client, channel, 5 + ) + sent_msg = thinking_future.result() + channel_context = context_future.result() + final_thread_ts = ts + else: + # Non-background response (e.g. auth error) — just send normally + client.send_message( + channel=channel, + text=response_text, + thread_ts=thread_ts + ) + + if trigger_background and sent_msg: thinking_ts = sent_msg.get('ts') - logger.info(f"Processing @Aurora mention in channel {channel}, thread {final_thread_ts}: {text[:100]}") + logger.info("Processing @Aurora mention in channel %s, thread %s", channel, final_thread_ts) send_message_to_aurora( user_id=user_id, @@ -161,7 +178,7 @@ def slack_events(): session_id=session_id, context_messages=context_messages, channel_context=channel_context, - thinking_message_ts=thinking_ts + thinking_message_ts=thinking_ts, ) except Exception as e: logger.error(f"Failed to send message to Slack: {e}") diff --git a/server/routes/slack/slack_events_helpers.py b/server/routes/slack/slack_events_helpers.py index f9bbf8e77..b8448eb4b 100644 --- a/server/routes/slack/slack_events_helpers.py +++ b/server/routes/slack/slack_events_helpers.py @@ -627,7 +627,7 @@ def get_session_from_thread(user_id: str, channel_id: str, thread_ts: str): def send_message_to_aurora(user_id: str, message_text: str, channel: str, thread_ts: str = None, incident_id: str = None, session_id: str = None, context_messages: list = None, - channel_context: str = None, thinking_message_ts: str = None): + channel_context: str | None = None, thinking_message_ts: str | None = None): """ Route a Slack message to Aurora's chat system. Uses background chat task to process the message.