perf: reduce chat latency — dedup, prewarm, parallelism, caching

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:

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()

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

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]