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,68 @@
 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: initialize expensive singletons once at worker startup
+# so background chat tasks don't pay cold-start penalties on every invocation.
+# ---------------------------------------------------------------------------
+import threading
+from celery.signals import worker_process_init
+
+
+@worker_process_init.connect
+def _prewarm_worker(**kwargs):
+    """Pre-initialize expensive singletons at worker child process boot.
+
+    Celery's worker_process_init has a ~4s timeout — handlers that block longer
+    cause the parent to assume the child failed. We run the heavy init in a
+    background thread and expose a threading.Event that task code can wait on
+    before doing real work (avoids cold-start on first task).
+    """
+    import threading
+
+    _logger = logging.getLogger("celery.prewarm")
+
+    def _do_prewarm():
+        # 1. Pre-warm NeMo Guardrails (the _rails_instance global)
+        try:
+            import time as _pw_time
+            _t0 = _pw_time.perf_counter()
+            import guardrails.input_rail as _rail_mod
+            if _rail_mod._rails_instance is None:
+                _rail_mod._rails_instance = _rail_mod._build_rails_sync()
+                _ms = (_pw_time.perf_counter() - _t0) * 1000
+                _logger.info(f"[PREWARM] NeMo Guardrails initialized in {_ms:.0f} ms")
+            else:
+                _logger.info("[PREWARM] NeMo Guardrails already initialized")
+        except Exception as e:
+            _logger.warning(f"[PREWARM] Failed to pre-warm guardrails: {e}")
+
+        # 2. Start MCP preloader (background thread that refreshes tool schemas)
+        try:
+            from chat.backend.agent.tools.mcp_preloader import start_mcp_preloader
+            start_mcp_preloader()
+            _logger.info("[PREWARM] MCP Preloader service started")
+        except Exception as e:
+            _logger.warning(f"[PREWARM] Failed to start MCP preloader: {e}")
+
+        # 3. Pre-warm the per-worker Agent singleton (avoids ~4.6s cold start)
+        try:
+            import time as _pw_time2
+            _t1 = _pw_time2.perf_counter()
+            from chat.background.task import _get_worker_agent
+            _get_worker_agent()
+            _ms2 = (_pw_time2.perf_counter() - _t1) * 1000
+            _logger.info(f"[PREWARM] Agent singleton created in {_ms2:.0f} ms")
+        except Exception as e:
+            _logger.warning(f"[PREWARM] Failed to pre-warm Agent singleton: {e}")
+
+        _prewarm_ready.set()
+
+    t = threading.Thread(target=_do_prewarm, name="celery-prewarm", daemon=True)
+    t.start()
+
+
+# Event that task code can wait on to ensure prewarm completed
+_prewarm_ready = threading.Event()

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,29 @@ 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(None)
+                            if _pending_rail is not None:
+                                from guardrails.input_rail import InputRailResult, _BLOCKED_REASON, _FAIL_CLOSED_AUTH, _FAIL_CLOSED_CONNECTIVITY
+                                rail_result: InputRailResult = await _pending_rail
+                                # 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