perf: reduce chat latency — dedup, prewarm, parallelism, caching - #533
Merged
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
WalkthroughThe PR adds a per-worker prewarm mechanism that initializes NeMo Guardrails, the MCP preloader, and a singleton Agent on Celery worker startup. Input guardrail evaluation is moved from a synchronous pre-stream gate to a concurrent background ChangesWorker Performance, Concurrent Guardrails & Background Task Improvements
Sequence Diagram(s)sequenceDiagram
participant CeleryWorker
participant _prewarm_worker
participant _ensure_rails_in_thread
participant _get_worker_agent
participant _prewarm_ready
participant run_background_chat
participant RedisLock
participant _execute_background_chat
rect rgba(173, 216, 230, 0.5)
note over CeleryWorker, _prewarm_ready: Worker startup prewarm
CeleryWorker->>_prewarm_worker: worker_process_init
_prewarm_worker->>_ensure_rails_in_thread: build LLMRails under _rails_thread_lock
_prewarm_worker->>_get_worker_agent: create Agent singleton
_prewarm_worker->>_prewarm_ready: set()
end
rect rgba(144, 238, 144, 0.5)
note over run_background_chat, _execute_background_chat: Task execution with dedup guard
run_background_chat->>_prewarm_ready: wait(timeout=30)
run_background_chat->>RedisLock: SETNX(task_id)
alt already running
run_background_chat-->>run_background_chat: return deduplicated
else acquired
run_background_chat->>_execute_background_chat: run with singleton Agent
end
end
sequenceDiagram
participant Workflow
participant _guardrail_task_var
participant check_input
participant MessageStore
participant agentic_tool_flow
participant ThreadPoolExecutor
Workflow->>check_input: asyncio.create_task(check_input(msg_text))
Workflow->>_guardrail_task_var: store task
Workflow->>MessageStore: handle_immediate_save (optimistic)
Workflow->>agentic_tool_flow: begin streaming
par prompt preparation
ThreadPoolExecutor->>ThreadPoolExecutor: build_prompt_segments(...)
and tool preparation
agentic_tool_flow->>agentic_tool_flow: get_cloud_tools()
end
agentic_tool_flow->>_guardrail_task_var: await task result
alt blocked
agentic_tool_flow-->>Workflow: emit_block_event, guardrail_blocked=True, return early
else allowed
agentic_tool_flow-->>Workflow: stream tokens to client
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
OlivierTrudeau
force-pushed
the
feat/latency-instrumentation
branch
from
June 18, 2026 19:56
1446a29 to
4bb4c94
Compare
![aurora-test-app1[bot]](https://avatars.githubusercontent.com/in/3646679?s=60&v=4){kind=small}
There was a problem hiding this comment.
Aurora Risk Review
Verdict: RISKY
This PR introduces three deployment-day risks: a CI environment-variable validation gate that is actively failing on the head SHA (confirmed run ID 27785708973), a guardrail safety regression where optimistic message persistence combined with a silently-swallowed delete failure leaves blocked messages in LLM context, and a fragile token-count heuristic that mislabels LLM stream errors as guardrail blocks. The latency improvements themselves are sound, but these three issues need resolution before merge.
Findings
| # | Severity | File | Finding |
|---|---|---|---|
| 1 | HIGH | .github/workflows/validate-env-vars.yml:1 |
Validate Environment Variables CI gate is actively failing on the head SHA |
| 2 | HIGH | server/chat/backend/agent/workflow.py:1082 |
Optimistic message persist + silent delete failure leaves blocked messages in LLM context |
| 3 | MEDIUM | server/chat/backend/agent/workflow.py:1316 |
Zero-token heuristic conflates LLM stream errors with guardrail blocks, showing wrong error message to users |
Aurora reviews PRs for incident prevention. This is advisory only and does not block merge.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 13
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
server/chat/backend/agent/tools/cloud_tools.py (1)
1118-1162:⚠️ Potential issue | 🟠 MajorProtect reads/writes of
_tc.current_tool_callswith_tc.lockin matching paths.Multiple iterations over
_tc.current_tool_callsat lines 1118–1162 (and again at lines 1211–1221, 1250–1253) lack lock protection while mutations occur outside the lock. This contradicts the protected access pattern already used elsewhere in the wrapper (line 1095+ for placeholder updates). With concurrent tool callbacks and cleanup threads, unprotected dictionary iteration can throwRuntimeError: dictionary changed size during iteration, and mutations can race with deletions.Apply the snapshot pattern used in the codebase: acquire
_tc.lock, snapshotcurrent_tool_calls.items()to a list, then iterate; protect all mutations with the lock as well.Pattern example
- for call_id, call_info in _tc.current_tool_calls.items(): + with _tc.lock: + calls_snapshot = list(_tc.current_tool_calls.items()) + for call_id, call_info in calls_snapshot: ... - call_info = _tc.current_tool_calls[matching_tool_call_id] - call_info['input'] = signature_payload + with _tc.lock: + call_info = _tc.current_tool_calls.get(matching_tool_call_id) + if call_info: + call_info['input'] = signature_payload🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1118 - 1162, The multiple iterations over _tc.current_tool_calls dictionary (in the for loops checking signatures, tool names, and candidates from lines 1118-1162) lack lock protection, creating race conditions with concurrent mutations. Wrap all reads from _tc.current_tool_calls with _tc.lock by acquiring the lock before the iterations, converting _tc.current_tool_calls.items() to a list snapshot to prevent RuntimeError during concurrent modifications, and ensure all mutations to call_info dictionaries (such as setting call_info['input'] and call_info['signature']) also occur within the locked section. Follow the same protected access pattern already used elsewhere in the wrapper for consistency.server/routes/slack/slack_events_helpers.py (1)
691-696:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winCapture
dispatch_tsat the Celery enqueue boundary.The worker logs
time.time() - trigger_metadata["dispatch_ts"]as Celery queue wait, but the current caller captures this beforesend_message_to_aurora()may create a session and prepare metadata. That helper/database time will be misreported as queue wait; setdispatch_tsimmediately before.delay()instead.Proposed fix
- if dispatch_ts: - trigger_metadata["dispatch_ts"] = dispatch_ts - # Launch background chat task # For Slack `@mentions`, link to incident but don't send "investigation started" notifications + import time as _aurora_time + trigger_metadata["dispatch_ts"] = _aurora_time.time() run_background_chat.delay(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/routes/slack/slack_events_helpers.py` around lines 691 - 696, The dispatch_ts is currently being set too early in the flow, before send_message_to_aurora() executes, which means any time spent in that helper function gets incorrectly attributed to Celery queue wait time in worker logs. Move the assignment of trigger_metadata["dispatch_ts"] = dispatch_ts to immediately before the run_background_chat.delay() call so that the timestamp accurately captures only the queue wait time, not the pre-enqueueing processing time.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/celery_config.py`:
- Around line 256-298: The _prewarm_worker function connected to
worker_process_init is performing blocking operations that exceed Celery's
4-second timeout, causing worker processes to be terminated. Move the heavy
initialization work (NeMo Guardrails via _build_rails_sync, MCP preloader via
start_mcp_preloader, and Agent singleton via _get_worker_agent) out of the
worker_process_init signal handler and into a background warmup task that runs
asynchronously or implement lazy initialization that occurs on first task
execution instead of blocking the worker startup process.
In `@server/chat/backend/agent/agent.py`:
- Around line 347-360: The get_cloud_tools() call is being executed in a
ThreadPoolExecutor worker thread, but it depends on thread-local user context
that was set on the main thread via set_user_context() and set_tool_capture().
Since thread-local context doesn't propagate to worker threads,
get_cloud_tools() returns tools with default context instead of the intended
capture/cache behavior. Remove the _tools_future submission from the thread pool
and instead call get_cloud_tools() directly on the current thread after
obtaining the segments result from _prompt_future.result(). Keep only the
build_prompt_segments call in the ThreadPoolExecutor worker.
In `@server/chat/backend/agent/providers/bedrock_provider.py`:
- Around line 148-155: The cache key construction at line 148 in the
BedrockProvider class only includes native_model, region, and temperature, but
ignores other kwargs that remain after the streaming parameter is removed at
line 143. These remaining kwargs are later applied to the model config before
instantiation, meaning different runtime kwargs will incorrectly reuse a stale
cached model_instance. Include the effective kwargs (all kwargs except
streaming) in the cache_key construction at line 148, or alternatively bypass
the cache check when non-streaming kwargs are present.
- Around line 150-158: The BedrockProvider class accesses the shared
_native_client_cache dictionary without synchronization, creating race
conditions in multithreaded Flask environments. Add a class-level threading lock
to BedrockProvider (e.g., _cache_lock), then wrap all accesses to
_native_client_cache with this lock using a context manager. This includes the
read-check-evict block around lines 150-157 where the cache is checked and
entries are deleted, and the write operation around line 184 where new clients
are stored in the cache. Ensure every read, delete, and write operation on
_native_client_cache is protected by acquiring the lock to prevent concurrent
modifications that could cause KeyError exceptions or duplicate client
instantiations.
In `@server/chat/backend/agent/utils/immediate_save_handler.py`:
- Around line 57-72: The immediate_save_handler code performs an unprotected
read and blindly removes the last user message from the messages array, which
can delete the wrong message if concurrent operations modify the chat_sessions
row (another tab appends a message, or _append_new_turn_ui_messages writes
simultaneously). Fix this by adding a FOR UPDATE clause to the SELECT query to
lock the row before modification, and modify the function to accept a
current-turn message identifier (generated message id is preferred over text)
passed from the caller in agentic_tool_flow. Instead of popping the last user
message, iterate through the messages array to find and remove only the message
that matches the provided identifier, ensuring the correct message is rolled
back even under concurrent access.
In `@server/chat/backend/agent/workflow.py`:
- Around line 1072-1077: The guardrail check via check_input is only enforced in
the agentic_tool_flow path, but the workflow can route to alternative paths like
triage before direct_react (around lines 180-188) which bypass the guardrail
enforcement and allow unvetted input to reach LLM nodes. Add a graph-level
guardrail gate that runs before _route_start to ensure all input paths are
protected, or alternatively make every first-hop node (agentic_tool_flow,
triage, direct_react) await the guardrail task stored in _guardrail_task_var
before executing any LLM calls.
- Around line 1320-1343: The guardrail fallback path in the code starting around
line 1320 is not fail-closed because it treats pending or errored guardrail
tasks as None, potentially ignoring slow or failed safety checks. Replace the
conditional check for _guardrail_task.done() with an unconditional await of
_guardrail_task.result(), and modify the exception handling to fail closed
(raise or yield an error message) instead of silently setting rail_result to
None. Additionally, ensure that when guardrails block the message in this
fallback path, the code sets the guardrail_blocked flag and performs the same
audit/rollback/state cleanup that happens in the main agentic_tool_flow branch,
rather than just yielding a message and returning immediately.
In `@server/chat/background/task.py`:
- Around line 47-65: The cached PostgreSQLClient created in _get_worker_agent()
reuses the same database connection across multiple background tasks without
resetting the RLS context for each task, which can cause different users' tasks
to see incorrect or stale data. In the _execute_background_chat() function,
after calling _get_worker_agent() to retrieve the cached agent, you must call
set_rls_context() from utils.auth.stateless_auth to set the current user and
organization context on the PostgreSQLClient's connection before the workflow
executes any database operations. This ensures each task runs with the correct
RLS context for its user.
- Line 1344: Remove the f-string prefix from the logger.info call on line 1344
since the string is constant with no variable interpolation. Additionally,
locate the logging call in the except block around lines 1489-1490 and replace
the standard logger call (such as logger.error) with logger.exception() to
properly capture and log exception information within the exception handler.
- Around line 486-494: The current deduplication logic in the task method only
checks if the Redis key exists but doesn't verify whether the original task
execution is actually still running or has already completed. This can cause the
task to return a "deduplicated" status prematurely even if the worker crashed
after acquiring the lock, leaving the session stuck. Instead of storing just "1"
in the _dedup_key, store an owner identifier or heartbeat timestamp. When the
dedup key exists, check if the owner/heartbeat indicates the original task is
verifiably still running (within a reasonable heartbeat window) or has completed
successfully. Only skip execution and return the "deduplicated" status if you
can confirm one of these conditions; otherwise, allow the task to proceed as a
legitimate retry. This ensures that genuinely concurrent duplicates are still
prevented while crashed tasks can be properly retried.
In `@server/main_chatbot.py`:
- Around line 1672-1685: The REDIS_URL logging in the startup environment info
block is a security risk because slicing to 30 characters still exposes
authentication credentials in the format `redis://:password@...`. Remove the
slicing logic on the REDIS_URL environment variable and instead log only the
presence or absence of the URL (such as "set" or "unset"), or parse the URL and
log only safe components like the scheme, host, and port while completely
excluding any credentials.
In `@server/routes/slack/slack_events_helpers.py`:
- Around line 630-631: The function signature containing the parameters
channel_context, thinking_message_ts, and dispatch_ts has implicit None defaults
without explicit type union syntax, violating Ruff RUF013. Update the type
annotations for these parameters to explicitly include None in the union (e.g.,
change str to str | None and float to float | None). At minimum, fix the
parameters on the changed lines 630-631, but ideally update the entire function
signature starting at line 628 for consistency.
In `@server/routes/slack/slack_events.py`:
- Line 183: In the logger.info call that logs the `@Aurora` mention processing,
remove the raw text snippet `{text[:100]}` from the log message as it may expose
secrets or sensitive customer data. Replace the f-string logging with lazy
logging by using the old-style % formatting or passing parameters separately to
the logger.info method, so that string interpolation only occurs if the log
level is actually enabled. Keep only the metadata fields like channel and
final_thread_ts in the log output.
---
Outside diff comments:
In `@server/chat/backend/agent/tools/cloud_tools.py`:
- Around line 1118-1162: The multiple iterations over _tc.current_tool_calls
dictionary (in the for loops checking signatures, tool names, and candidates
from lines 1118-1162) lack lock protection, creating race conditions with
concurrent mutations. Wrap all reads from _tc.current_tool_calls with _tc.lock
by acquiring the lock before the iterations, converting
_tc.current_tool_calls.items() to a list snapshot to prevent RuntimeError during
concurrent modifications, and ensure all mutations to call_info dictionaries
(such as setting call_info['input'] and call_info['signature']) also occur
within the locked section. Follow the same protected access pattern already used
elsewhere in the wrapper for consistency.
In `@server/routes/slack/slack_events_helpers.py`:
- Around line 691-696: The dispatch_ts is currently being set too early in the
flow, before send_message_to_aurora() executes, which means any time spent in
that helper function gets incorrectly attributed to Celery queue wait time in
worker logs. Move the assignment of trigger_metadata["dispatch_ts"] =
dispatch_ts to immediately before the run_background_chat.delay() call so that
the timestamp accurately captures only the queue wait time, not the
pre-enqueueing processing time.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 842f112d-f6f3-4500-a8ec-30cfca07db99
📒 Files selected for processing (16)
deploy/helm/aurora/templates/celery-worker-deployment.yamldeploy/helm/aurora/templates/chatbot-deployment.yamlserver/celery_config.pyserver/chat/backend/agent/agent.pyserver/chat/backend/agent/prompt/composer.pyserver/chat/backend/agent/providers/bedrock_provider.pyserver/chat/backend/agent/tools/cloud_tools.pyserver/chat/backend/agent/utils/immediate_save_handler.pyserver/chat/backend/agent/workflow.pyserver/chat/background/task.pyserver/guardrails/input_rail.pyserver/main_chatbot.pyserver/routes/slack/slack_events.pyserver/routes/slack/slack_events_helpers.pyserver/utils/db/connection_pool.pyserver/utils/secrets/secret_ref_utils.py
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
server/chat/backend/agent/agent.py (1)
877-877:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winRemove unused imports from guardrails module.
_BLOCKED_REASON,_FAIL_CLOSED_AUTH, and_FAIL_CLOSED_CONNECTIVITYare imported but never used in this function.🧹 Proposed fix
- from guardrails.input_rail import InputRailResult, _BLOCKED_REASON, _FAIL_CLOSED_AUTH, _FAIL_CLOSED_CONNECTIVITY + from guardrails.input_rail import InputRailResult🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/agent.py` at line 877, Remove the three unused imports from the guardrails.input_rail import statement on line 877 of agent.py. The constants _BLOCKED_REASON, _FAIL_CLOSED_AUTH, and _FAIL_CLOSED_CONNECTIVITY are not used anywhere in the function and should be deleted from the import statement, keeping only InputRailResult which is actively used.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@server/chat/backend/agent/agent.py`:
- Line 877: Remove the three unused imports from the guardrails.input_rail
import statement on line 877 of agent.py. The constants _BLOCKED_REASON,
_FAIL_CLOSED_AUTH, and _FAIL_CLOSED_CONNECTIVITY are not used anywhere in the
function and should be deleted from the import statement, keeping only
InputRailResult which is actively used.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: f05f5a85-acbf-42f9-ac65-55aa62d13ed2
📒 Files selected for processing (4)
server/celery_config.pyserver/chat/backend/agent/agent.pyserver/chat/background/task.pyserver/routes/slack/slack_events.py
Adds [LATENCY]-prefixed timing logs across the full chat message pipeline to diagnose ~30s response times in the deployed K8s environment. Instruments: auth/provider resolution, guardrails input rail, prompt building, tool loading, LLM model creation, Bedrock instantiation, Vault secret fetches, DB pool acquisition, and end-to-end time-to-first-token. Also adds POD_NAME/NODE_NAME env vars to the chatbot Helm deployment for pod identification. Co-authored-by: Cursor <cursoragent@cursor.com>
The _pwa_first_token_logged variable needs nonlocal declaration to be accessible from the nested process_stream closure. Co-authored-by: Cursor <cursoragent@cursor.com>
…allelism) - Cache ChatBedrockConverse instances per (model, region, temp) to skip repeated boto3 session + TLS handshake (~300-570ms savings per message) - Per-worker singleton Agent in Celery to avoid recreating PostgreSQLClient, WeaviateClient, and LLMManager on every background task (~2s savings) - Parallelize build_prompt_segments + get_cloud_tools in agentic_tool_flow (overlaps CPU-bound prompt assembly with tool loading) - Parallelize Slack "Thinking..." message send with channel context fetch using ThreadPoolExecutor (~800ms savings on new top-level messages) Co-authored-by: Cursor <cursoragent@cursor.com>
Fire the NeMo Guardrails input rail as a concurrent asyncio.Task instead of awaiting it synchronously. The guardrails LLM call (~1.2s) now runs in parallel with provider resolution, tool loading, prompt assembly, and LLM client creation inside agentic_tool_flow. Messages are persisted optimistically on receipt; if guardrails blocks, the persisted message is rolled back via delete_last_saved_message. The LLM is never reached when blocked — safety guarantees are preserved. Uses a contextvars.ContextVar to pass the pending task from workflow.stream() to agentic_tool_flow without modifying the Pydantic State schema. Co-authored-by: Cursor <cursoragent@cursor.com>
… dedup - Make worker_process_init prewarm synchronous so children are fully warm before accepting tasks (was a daemon thread that raced with incoming tasks) - Remove verbose TOOL CAPTURE debug logging from cloud_tools.py - Simplify dedup guard Redis access (backend.client is always available) - Fix persist condition in workflow.py to match original behavior (save even when _rail_msg_text is empty, matching pre-refactor logic) - Clean up DNS check in main_chatbot.py startup Co-authored-by: Cursor <cursoragent@cursor.com>
- Remove secret ref from latency log (CodeQL clear-text sensitive data) - Use single import style for guardrails.input_rail (code quality) - Remove unused sent_msg assignment in non-background path (code quality) Co-authored-by: Cursor <cursoragent@cursor.com>
The Bedrock client cache saved ~300-570ms by reusing boto3 sessions, but production runs on OpenRouter where ChatOpenAI instantiation is already cheap (~10ms). Reverting to keep the diff minimal and avoid dead code for the active provider path. Co-authored-by: Cursor <cursoragent@cursor.com>
Debugging-only code that adds noise to the diff without lasting value. Co-authored-by: Cursor <cursoragent@cursor.com>
Blocked messages sitting in chat history are harmless — it's what the user typed. The optimistic persist doesn't need a rollback path. Co-authored-by: Cursor <cursoragent@cursor.com>
- Move prewarm to background thread with Event gate (Celery 4s timeout) - Keep get_cloud_tools on main thread (needs thread-local context) - Fix constant f-string lint warning - Remove user message text from Slack processing log (privacy) Co-authored-by: Cursor <cursoragent@cursor.com>
- Guardrail fallback now fails closed on exceptions (was silently ignoring) - Use logger.exception for early Slack send failure - Add Optional type annotations on dispatch_ts parameter Co-authored-by: Cursor <cursoragent@cursor.com>
OlivierTrudeau
force-pushed
the
feat/latency-instrumentation
branch
from
June 18, 2026 21:21
0bf1b32 to
d3aef04
Compare
|
🔍 Aurora is reviewing this PR for incident risk. This usually takes a minute or two — findings will appear as a review when it's done. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
server/chat/backend/agent/tools/cloud_tools.py (3)
1057-1058:⚠️ Potential issue | 🟠 Major | ⚡ Quick winNormalize
Nonevalues the same way as workflow signatures.
Workflow._build_tool_signature()dropsNonevalues, but this wrapper keeps them. Optional defaults likeoutput_file=Nonecan miss exact matching and force unsafe fallbacks.Proposed fix
def _sanitize_kwargs_for_signature(kwargs: Dict[str, Any]) -> Dict[str, Any]: - return {k: v for k, v in kwargs.items() if k not in INTERNAL_CONTEXT_KEYS} + return { + k: v + for k, v in kwargs.items() + if k not in INTERNAL_CONTEXT_KEYS and v is not None + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1057 - 1058, The `_sanitize_kwargs_for_signature` function currently filters out internal context keys but does not drop `None` values, while `Workflow._build_tool_signature()` drops them during normalization. This mismatch causes optional parameters with `None` defaults (like `output_file=None`) to fail exact matching. Modify the dictionary comprehension in `_sanitize_kwargs_for_signature` to exclude both keys present in `INTERNAL_CONTEXT_KEYS` and any entries where the value is `None`, ensuring consistency with the workflow signature building behavior.
1137-1164:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDo not attach output to an arbitrary oldest call when matching is ambiguous.
When multiple incomplete calls for the same tool remain, choosing the oldest can persist a successful output under the wrong
tool_call_id. Mirror the error path’s fail-safe behavior unless there is an explicit sequential-execution signal.Proposed fix
if len(candidate_ids) == 1: matching_tool_call_id = candidate_ids[0][0] call_info = _tc.current_tool_calls[matching_tool_call_id] 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, ) elif len(candidate_ids) > 1: - candidate_ids.sort(key=lambda x: x[1] if x[1] else datetime.min) - matching_tool_call_id = candidate_ids[0][0] - call_info = _tc.current_tool_calls[matching_tool_call_id] - 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}. " - f"This is expected for OpenAI sequential execution." - ) + logging.error( + "Ambiguous %s completion: %d incomplete calls remain and no signature matched; " + "leaving capture unmatched rather than misattributing output", + tool_name, + len(candidate_ids), + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1137 - 1164, In the branch where len(candidate_ids) > 1 in the tool call matching logic, the code currently sorts candidates by start_time and assigns the output to the oldest incomplete call without verification. This can incorrectly attach output to the wrong tool_call_id. Instead of assigning to the oldest candidate, implement a fail-safe approach that handles the ambiguous matching case (multiple incomplete calls of the same tool with no clear match) by avoiding arbitrary assignment. Look for how error cases are handled elsewhere in the matching logic to mirror similar fail-safe behavior when sequential execution cannot be definitively determined.
1292-1334:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftApply the PR-review read-only gate to all write-capable tools.
The new gate removes
iac_tool,cloud_exec,github_commit,tailscale_ssh, and kubectl, but PR review can still receive write-capable tools later in this function, including GitLab write actions, Bitbucket mutators, Jira mutators, SharePoint page creation, andcloudflare_action.Also applies to: 2031-2156, 2235-2311, 2438-2477
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1292 - 1334, The PR review read-only gate using the `not is_pr_review` check has been inconsistently applied across write-capable tools in this function. While `iac_tool`, `cloud_exec`, `github_commit`, and `tailscale_ssh` are properly gated, other write-capable tools later in the function (including GitLab write actions, Bitbucket mutators, Jira mutators, SharePoint page creation, and `cloudflare_action`) are missing this protection. For each of these unprotected write-capable tools, add `and not is_pr_review` to their conditional checks to ensure they cannot be used during PR review mode, maintaining consistency with the existing patterns used for other write tools in the function.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/chat/backend/agent/tools/cloud_tools.py`:
- Around line 1024-1034: The cache key construction in the code starting around
the cache_key variable assignment is missing the trigger_action_id value, which
can cause different trigger actions for the same user and mode to incorrectly
reuse the same cached entry. Extract the trigger_action_id from the
state_context (similar to how other attributes like trigger_rca_requested and
is_background are extracted using getattr with a default value), and then
include this value in the cache_key string template alongside the other
parameters like rca_flag, is_background, is_postmortem_action, is_pr_review, and
is_rca_context to ensure each unique action has a unique cache entry.
- Around line 1037-1044: The cache hit in the caching logic (returning early
from _langchain_tools_cache) bypasses the connector capability and connectivity
checks that occur later in the function, allowing stale tools to be exposed
after credentials are removed or org-scoped changes occur. To fix this, either
incorporate connector and organization capability version information into the
cache_key itself so that cache invalidation happens automatically when
capabilities change, or ensure that is_<connector>_connected(user_id) checks are
performed even on cache hits before returning the cached tools. This ensures
agent tools are always gated behind proper connectivity validation regardless of
whether the result comes from cache or not. Apply this same validation pattern
to all other cache hit locations mentioned in the same file.
- Around line 1115-1123: The code iterates over _tc.current_tool_calls in
multiple places without holding _tc.lock, creating a race condition where other
tool executions can modify or delete entries during iteration. To fix this,
acquire _tc.lock before creating a snapshot of _tc.current_tool_calls (e.g., via
.copy() or list()), then release the lock and iterate over the snapshot for
matching logic. Apply this pattern to all three locations where
current_tool_calls is iterated: around the first loop checking signature
matching, the fallback loop matching by tool_name and command, and any other
iteration loops in the specified line ranges (1141-1149, 1208-1218). When you
need to mutate state based on the matched result, re-acquire _tc.lock before
accessing or modifying _tc.current_tool_calls again.
In `@server/chat/background/task.py`:
- Around line 47-65: The _get_worker_agent() function lacks thread
synchronization, allowing multiple threads to simultaneously create separate
Agent instances when called from both the prewarm daemon thread and task
execution threads. Add a global threading.Lock variable (similar to how
_worker_agent is declared globally) and use it to wrap the entire singleton
creation logic in _get_worker_agent() so that only one thread can initialize the
Agent, PostgreSQLClient, and WeaviateClient instances at a time. Acquire the
lock before the null check on _worker_agent and release it after initialization
completes to prevent resource leaks from duplicate instances.
- Around line 1447-1448: The `_slack_early_sent` flag is being set to True
unconditionally after calling `_send_response_to_slack()`, even when the
function returns normally without actually sending anything (no channel, no
assistant message, or missing Slack client). This causes the parent code to
incorrectly skip the fallback send. Modify the `_send_response_to_slack()`
function to return a boolean value indicating whether it actually sent a message
successfully, and then only set `_slack_early_sent = True` when the function
returns True, not unconditionally.
- Around line 483-495: The Redis deduplication guard check using the SETNX
operation on `_dedup_key` is currently running after the
`_prewarm_ready.wait(timeout=30)` call, which causes duplicate task redeliveries
to unnecessarily block on prewarm for up to 30 seconds. Move the entire dedup
guard block (the lines that create `_dedup_key`, instantiate `_redis`, and
perform the `_redis.set(_dedup_key, "1", nx=True, ex=1800)` SETNX check) to
execute before the prewarm wait so that duplicate redeliveries fail the dedup
check immediately without occupying worker slots.
---
Outside diff comments:
In `@server/chat/backend/agent/tools/cloud_tools.py`:
- Around line 1057-1058: The `_sanitize_kwargs_for_signature` function currently
filters out internal context keys but does not drop `None` values, while
`Workflow._build_tool_signature()` drops them during normalization. This
mismatch causes optional parameters with `None` defaults (like
`output_file=None`) to fail exact matching. Modify the dictionary comprehension
in `_sanitize_kwargs_for_signature` to exclude both keys present in
`INTERNAL_CONTEXT_KEYS` and any entries where the value is `None`, ensuring
consistency with the workflow signature building behavior.
- Around line 1137-1164: In the branch where len(candidate_ids) > 1 in the tool
call matching logic, the code currently sorts candidates by start_time and
assigns the output to the oldest incomplete call without verification. This can
incorrectly attach output to the wrong tool_call_id. Instead of assigning to the
oldest candidate, implement a fail-safe approach that handles the ambiguous
matching case (multiple incomplete calls of the same tool with no clear match)
by avoiding arbitrary assignment. Look for how error cases are handled elsewhere
in the matching logic to mirror similar fail-safe behavior when sequential
execution cannot be definitively determined.
- Around line 1292-1334: The PR review read-only gate using the `not
is_pr_review` check has been inconsistently applied across write-capable tools
in this function. While `iac_tool`, `cloud_exec`, `github_commit`, and
`tailscale_ssh` are properly gated, other write-capable tools later in the
function (including GitLab write actions, Bitbucket mutators, Jira mutators,
SharePoint page creation, and `cloudflare_action`) are missing this protection.
For each of these unprotected write-capable tools, add `and not is_pr_review` to
their conditional checks to ensure they cannot be used during PR review mode,
maintaining consistency with the existing patterns used for other write tools in
the function.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: df06d68c-2186-4bd3-9ebf-921055e61e67
📒 Files selected for processing (8)
deploy/helm/aurora/templates/celery-worker-deployment.yamlserver/celery_config.pyserver/chat/backend/agent/agent.pyserver/chat/backend/agent/tools/cloud_tools.pyserver/chat/backend/agent/workflow.pyserver/chat/background/task.pyserver/routes/slack/slack_events.pyserver/routes/slack/slack_events_helpers.py
OlivierTrudeau
force-pushed
the
feat/latency-instrumentation
branch
3 times, most recently
from
June 18, 2026 22:11
d94b278 to
2fa5743
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
server/chat/backend/agent/tools/cloud_tools.py (1)
1190-1215:⚠️ Potential issue | 🟠 Major | ⚡ Quick winCap tool output before capturing it.
capture_tool_end()receivesoutput_strbeforecap_tool_output()runs, so collected ToolMessages—especially on cancellation—can persist the raw oversized output that the returned LangChain tool result avoids.Proposed fix
+ from chat.backend.agent.utils.tool_output_cap import cap_tool_output + result_str = json.dumps(result) if isinstance(result, dict) else str(result) + capped_result = cap_tool_output(result_str, tool_name) + if matching_tool_call_id and tool_capture: @@ if not already_captured: try: - output_str = json.dumps(result) if isinstance(result, dict) else str(result) - tool_capture.capture_tool_end(matching_tool_call_id, output_str, is_error=False) + tool_capture.capture_tool_end(matching_tool_call_id, capped_result, is_error=False) logging.info(f"Wrapper called capture_tool_end for {matching_tool_call_id}") except Exception as capture_error: logging.error(f"Failed to capture tool end in wrapper: {capture_error}") @@ - # Cap tool output before returning to LangChain so the ReAct - # loop never accumulates oversized ToolMessages. - from chat.backend.agent.utils.tool_output_cap import cap_tool_output - result_str = json.dumps(result) if isinstance(result, dict) else str(result) - result = cap_tool_output(result_str, tool_name) - - return result + return capped_result🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1190 - 1215, The cap_tool_output function is being called after capture_tool_end, which means the captured tool end message contains the raw uncapped output while the returned result is capped. Move the cap_tool_output call and result_str computation (currently at lines 1210-1212) to execute before the capture_tool_end call (around line 1194), so that capture_tool_end receives the already-capped output_str instead of the raw output, ensuring consistency between captured and returned tool results.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/chat/backend/agent/tools/cloud_tools.py`:
- Around line 1107-1123: Remove the sensitive payload logging statements in the
tool capture block that are exposing raw kwargs, current_tool_calls, call_info
details, and result data. Specifically, remove or replace the logging.info calls
that log the KWARGS, the tool_capture.current_tool_calls dictionary, the
call_info, and the result variable, as these can contain sensitive information
like secrets, commands, or cloud resource details. Replace these with logging
statements that only capture non-sensitive metadata such as tool_name and
call_id for debugging purposes.
- Around line 1165-1170: The sorting lambda in the candidate_ids.sort() call
uses datetime.min as a fallback when x[1] (start_time) is None, but this causes
a TypeError because datetime.min is timezone-naive while start_time is always
timezone-aware (created with datetime.now(timezone.utc)). Replace the
datetime.min fallback with datetime.min.replace(tzinfo=timezone.utc) to ensure
both the fallback value and the actual start_time values are timezone-aware and
can be compared without error.
In `@server/chat/backend/agent/workflow.py`:
- Around line 1075-1082: The synchronous handle_immediate_save() function blocks
the event loop immediately after scheduling the check_input() guardrail task,
preventing the guardrails from running concurrently. Move the
handle_immediate_save() call into a worker thread using asyncio utilities (such
as loop.run_in_executor()) so that the event loop remains available to execute
the guardrails task concurrently with the database persistence operation.
- Around line 1061-1077: The is_scaffold variable is computed to identify
system-authored background RCA messages but is not being used to prevent
guardrail checks. Add is_scaffold to the conditional guard on the line that
schedules the check_input task, so the condition reads if msg_text and not
is_pr_review and not is_scaffold: to skip input rails for scaffold messages just
as PR review messages are skipped.
---
Outside diff comments:
In `@server/chat/backend/agent/tools/cloud_tools.py`:
- Around line 1190-1215: The cap_tool_output function is being called after
capture_tool_end, which means the captured tool end message contains the raw
uncapped output while the returned result is capped. Move the cap_tool_output
call and result_str computation (currently at lines 1210-1212) to execute before
the capture_tool_end call (around line 1194), so that capture_tool_end receives
the already-capped output_str instead of the raw output, ensuring consistency
between captured and returned tool results.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: f4987ec8-2505-4859-88d7-52c22443f1fd
📒 Files selected for processing (2)
server/chat/backend/agent/tools/cloud_tools.pyserver/chat/backend/agent/workflow.py
Contributor
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
server/chat/backend/agent/tools/cloud_tools.py (1)
1190-1215:⚠️ Potential issue | 🟠 Major | ⚡ Quick winCap tool output before capturing it.
capture_tool_end()receivesoutput_strbeforecap_tool_output()runs, so collected ToolMessages—especially on cancellation—can persist the raw oversized output that the returned LangChain tool result avoids.Proposed fix
+ from chat.backend.agent.utils.tool_output_cap import cap_tool_output + result_str = json.dumps(result) if isinstance(result, dict) else str(result) + capped_result = cap_tool_output(result_str, tool_name) + if matching_tool_call_id and tool_capture: @@ if not already_captured: try: - output_str = json.dumps(result) if isinstance(result, dict) else str(result) - tool_capture.capture_tool_end(matching_tool_call_id, output_str, is_error=False) + tool_capture.capture_tool_end(matching_tool_call_id, capped_result, is_error=False) logging.info(f"Wrapper called capture_tool_end for {matching_tool_call_id}") except Exception as capture_error: logging.error(f"Failed to capture tool end in wrapper: {capture_error}") @@ - # Cap tool output before returning to LangChain so the ReAct - # loop never accumulates oversized ToolMessages. - from chat.backend.agent.utils.tool_output_cap import cap_tool_output - result_str = json.dumps(result) if isinstance(result, dict) else str(result) - result = cap_tool_output(result_str, tool_name) - - return result + return capped_result🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1190 - 1215, The cap_tool_output function is being called after capture_tool_end, which means the captured tool end message contains the raw uncapped output while the returned result is capped. Move the cap_tool_output call and result_str computation (currently at lines 1210-1212) to execute before the capture_tool_end call (around line 1194), so that capture_tool_end receives the already-capped output_str instead of the raw output, ensuring consistency between captured and returned tool results.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/chat/backend/agent/tools/cloud_tools.py`:
- Around line 1107-1123: Remove the sensitive payload logging statements in the
tool capture block that are exposing raw kwargs, current_tool_calls, call_info
details, and result data. Specifically, remove or replace the logging.info calls
that log the KWARGS, the tool_capture.current_tool_calls dictionary, the
call_info, and the result variable, as these can contain sensitive information
like secrets, commands, or cloud resource details. Replace these with logging
statements that only capture non-sensitive metadata such as tool_name and
call_id for debugging purposes.
- Around line 1165-1170: The sorting lambda in the candidate_ids.sort() call
uses datetime.min as a fallback when x[1] (start_time) is None, but this causes
a TypeError because datetime.min is timezone-naive while start_time is always
timezone-aware (created with datetime.now(timezone.utc)). Replace the
datetime.min fallback with datetime.min.replace(tzinfo=timezone.utc) to ensure
both the fallback value and the actual start_time values are timezone-aware and
can be compared without error.
In `@server/chat/backend/agent/workflow.py`:
- Around line 1075-1082: The synchronous handle_immediate_save() function blocks
the event loop immediately after scheduling the check_input() guardrail task,
preventing the guardrails from running concurrently. Move the
handle_immediate_save() call into a worker thread using asyncio utilities (such
as loop.run_in_executor()) so that the event loop remains available to execute
the guardrails task concurrently with the database persistence operation.
- Around line 1061-1077: The is_scaffold variable is computed to identify
system-authored background RCA messages but is not being used to prevent
guardrail checks. Add is_scaffold to the conditional guard on the line that
schedules the check_input task, so the condition reads if msg_text and not
is_pr_review and not is_scaffold: to skip input rails for scaffold messages just
as PR review messages are skipped.
---
Outside diff comments:
In `@server/chat/backend/agent/tools/cloud_tools.py`:
- Around line 1190-1215: The cap_tool_output function is being called after
capture_tool_end, which means the captured tool end message contains the raw
uncapped output while the returned result is capped. Move the cap_tool_output
call and result_str computation (currently at lines 1210-1212) to execute before
the capture_tool_end call (around line 1194), so that capture_tool_end receives
the already-capped output_str instead of the raw output, ensuring consistency
between captured and returned tool results.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: f4987ec8-2505-4859-88d7-52c22443f1fd
📒 Files selected for processing (2)
server/chat/backend/agent/tools/cloud_tools.pyserver/chat/backend/agent/workflow.py
🛑 Comments failed to post (4)
server/chat/backend/agent/tools/cloud_tools.py (2)
1107-1123:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winRemove sensitive tool payload logging.
These INFO logs dump raw
kwargs,current_tool_calls, and pre-cappedresult. Tool arguments and outputs can include commands, cloud resources, logs, or secrets; keep only metadata.Proposed fix
- logging.info(f"TOOL CAPTURE: KWARGS: {kwargs}") signature_payload = _sanitize_kwargs_for_signature(kwargs) @@ - 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}") + logging.debug( + "TOOL CAPTURE: matching %s call with %d sanitized kwargs", + tool_name, + len(signature_payload), + ) for call_id, call_info in tool_capture.current_tool_calls.items(): - 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}") if call_info.get('signature') == tool_signature: matching_tool_call_id = call_id🧰 Tools
🪛 ast-grep (0.43.0)
[info] 1109-1109: use jsonify instead of json.dumps for JSON output
Context: json.dumps(signature_payload, sort_keys=True)
Note: Security best practice.(use-jsonify)
🪛 Ruff (0.15.17)
[warning] 1107-1107: Logging statement uses f-string
(G004)
[warning] 1118-1118: Logging statement uses f-string
(G004)
[warning] 1119-1119: Logging statement uses f-string
(G004)
[warning] 1121-1121: Logging statement uses f-string
(G004)
[warning] 1122-1122: Logging statement uses f-string
(G004)
[warning] 1123-1123: Logging statement uses f-string
(G004)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1107 - 1123, Remove the sensitive payload logging statements in the tool capture block that are exposing raw kwargs, current_tool_calls, call_info details, and result data. Specifically, remove or replace the logging.info calls that log the KWARGS, the tool_capture.current_tool_calls dictionary, the call_info, and the result variable, as these can contain sensitive information like secrets, commands, or cloud resource details. Replace these with logging statements that only capture non-sensitive metadata such as tool_name and call_id for debugging purposes.
1165-1170:
⚠️ Potential issue | 🟡 Minor🧩 Analysis chain
🏁 Script executed:
#!/bin/bash # Description: Confirm the mixed aware/naive datetime sort failure and the timestamp-based fix. python - <<'PY' from datetime import datetime, timezone items = [("aware", datetime.now(timezone.utc)), ("missing", None)] try: sorted(items, key=lambda x: x[1] if x[1] else datetime.min) except TypeError as exc: print("expected failure:", exc) fixed = sorted(items, key=lambda x: x[1].timestamp() if x[1] else float("-inf")) print("fixed order:", [name for name, _ in fixed]) PYRepository: Arvo-AI/aurora
Length of output: 164
🏁 Script executed:
# Check the actual code at the specified lines if [ -f "server/chat/backend/agent/tools/cloud_tools.py" ]; then wc -l server/chat/backend/agent/tools/cloud_tools.py sed -n '1160,1175p' server/chat/backend/agent/tools/cloud_tools.py else echo "File not found" fiRepository: Arvo-AI/aurora
Length of output: 1275
🏁 Script executed:
# Search for where start_time is set and its type to confirm if it's timezone-aware rg -B 5 -A 2 'start_time' server/chat/backend/agent/tools/cloud_tools.py | head -100Repository: Arvo-AI/aurora
Length of output: 1731
🏁 Script executed:
# Check for datetime imports and usage patterns in the file rg '(from datetime|import datetime)' server/chat/backend/agent/tools/cloud_tools.pyRepository: Arvo-AI/aurora
Length of output: 88
🏁 Script executed:
# Search for where candidate_ids is populated to understand the structure of x[1] rg -B 10 'candidate_ids' server/chat/backend/agent/tools/cloud_tools.py | head -150Repository: Arvo-AI/aurora
Length of output: 2460
🏁 Script executed:
# Search for where start_time is set to confirm it's timezone-aware rg -B 3 -A 3 "start_time.*=" server/chat/backend/agent/tools/cloud_tools.py | head -80Repository: Arvo-AI/aurora
Length of output: 40
🏁 Script executed:
# Check if datetime.utcnow, timezone.utc, or other aware datetime creation is used rg "(utcnow|timezone\.|aware|astimezone)" server/chat/backend/agent/tools/cloud_tools.pyRepository: Arvo-AI/aurora
Length of output: 40
🏁 Script executed:
# Search more broadly for start_time assignments in the entire agent tools directory rg "start_time" server/chat/backend/agent/tools/ -B 2 -A 2Repository: Arvo-AI/aurora
Length of output: 12748
🏁 Script executed:
# Check what datetime methods are used in the file rg "datetime\." server/chat/backend/agent/tools/cloud_tools.pyRepository: Arvo-AI/aurora
Length of output: 507
🏁 Script executed:
# Look at the tool_capture module to understand the structure of current_tool_calls fd -t f "tool_capture" server/chat/backend/agent/tools/ 2>/dev/null || fd -t f "tool_capture" server/ 2>/dev/null | head -5Repository: Arvo-AI/aurora
Length of output: 40
🏁 Script executed:
# Find where current_tool_calls is defined and populated rg "current_tool_calls" server/chat/backend/agent/tools/ -B 3 -A 3 | grep -E "(start_time|datetime)" | head -50Repository: Arvo-AI/aurora
Length of output: 461
🏁 Script executed:
# Search for tool_capture module definition find server -name "*tool_capture*" -o -name "*capture*" 2>/dev/null | head -10Repository: Arvo-AI/aurora
Length of output: 114
🏁 Script executed:
# Check if there's a place where start_time is explicitly set rg "\['start_time'\]|\"start_time\"" server/chat/backend/agent/tools/cloud_tools.py -B 2 -A 2Repository: Arvo-AI/aurora
Length of output: 40
🏁 Script executed:
# Examine the tool_context_capture.py file cat -n server/chat/backend/agent/utils/tool_context_capture.py | head -200Repository: Arvo-AI/aurora
Length of output: 11166
🏁 Script executed:
# Search for start_time in tool_context_capture rg "start_time" server/chat/backend/agent/utils/tool_context_capture.py -B 3 -A 3Repository: Arvo-AI/aurora
Length of output: 1601
Replace naive
datetime.minfallback with timezone-aware comparison.When sorting
candidate_idsbystart_time, mixing timezone-aware datetimes with naivedatetime.minraisesTypeError: can't compare offset-naive and offset-aware datetimes. Sincestart_timeis always set asdatetime.now(timezone.utc)(timezone-aware), the fallback must also be timezone-aware.Proposed fix
- candidate_ids.sort(key=lambda x: x[1] if x[1] else datetime.min) + candidate_ids.sort( + key=lambda x: x[1].timestamp() if x[1] else float("-inf") + )🧰 Tools
🪛 Ruff (0.15.17)
[warning] 1168-1168: Use of
datetime.datetime.minwithout timezone information(DTZ901)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1165 - 1170, The sorting lambda in the candidate_ids.sort() call uses datetime.min as a fallback when x[1] (start_time) is None, but this causes a TypeError because datetime.min is timezone-naive while start_time is always timezone-aware (created with datetime.now(timezone.utc)). Replace the datetime.min fallback with datetime.min.replace(tzinfo=timezone.utc) to ensure both the fallback value and the actual start_time values are timezone-aware and can be compared without error.Source: Linters/SAST tools
server/chat/backend/agent/workflow.py (2)
1061-1077:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winSkip input rails for RCA scaffold messages.
is_scaffoldis computed on Line 1063 but Line 1075 still schedulescheck_input(msg_text)for scaffoldHumanMessages. These are system-authored background RCA prompts, so a rail block or fail-closed result can incorrectly abort an RCA run.Proposed fix
- if msg_text and not is_pr_review: + if msg_text and not is_pr_review and not is_scaffold: _guardrail_task = asyncio.create_task(check_input(msg_text)) _guardrail_task_var.set(_guardrail_task)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/workflow.py` around lines 1061 - 1077, The is_scaffold variable is computed to identify system-authored background RCA messages but is not being used to prevent guardrail checks. Add is_scaffold to the conditional guard on the line that schedules the check_input task, so the condition reads if msg_text and not is_pr_review and not is_scaffold: to skip input rails for scaffold messages just as PR review messages are skipped.
1075-1082: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick win
Do not block the event loop immediately after scheduling the rail task.
asyncio.create_task()will not runcheck_input()while the coroutine stays in the synchronoushandle_immediate_save()DB path. Move the save into a worker thread so guardrails can actually overlap with persistence.Proposed fix
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 or last_msg.content) + await asyncio.to_thread( + handle_immediate_save, + input_state.session_id, + input_state.user_id, + msg_text or last_msg.content, + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/workflow.py` around lines 1075 - 1082, The synchronous handle_immediate_save() function blocks the event loop immediately after scheduling the check_input() guardrail task, preventing the guardrails from running concurrently. Move the handle_immediate_save() call into a worker thread using asyncio utilities (such as loop.run_in_executor()) so that the event loop remains available to execute the guardrails task concurrently with the database persistence operation.
OlivierTrudeau
force-pushed
the
feat/latency-instrumentation
branch
3 times, most recently
from
June 18, 2026 22:34
cd9dad2 to
52998b5
Compare
Keep only the cache key fix (remove id(tool_capture)) and the dynamic resolution of tool_capture via get_tool_capture() at call time. Remove TOOL CAPTURE debug logs. No cosmetic renames. Co-authored-by: Cursor <cursoragent@cursor.com>
OlivierTrudeau
force-pushed
the
feat/latency-instrumentation
branch
from
June 18, 2026 22:35
52998b5 to
0ea1554
Compare
…ntation # Conflicts: # server/chat/background/task.py
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
beng360
reviewed
Jun 19, 2026
beng360
reviewed
Jun 19, 2026
beng360
reviewed
Jun 19, 2026
beng360
reviewed
Jun 19, 2026
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
server/chat/backend/agent/tools/cloud_tools.py (1)
1027-1033:⚠️ Potential issue | 🟠 MajorCache key missing
trigger_action_id– stale action IDs can execute.The cache key (line 1033) omits
trigger_action_id, but_action_idis extracted at line 1340 and pinned to tools via closure. When a different action later hits the cache with identical other parameters, the cached tool carries the original action ID, not the current one.Include
trigger_action_idin the cache key to differentiate cache entries by action:Proposed fix
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) + trigger_action_id = getattr(state_context, 'trigger_action_id', None) if state_context else None 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}" + cache_key = ( + f"{user_id}:{capture_tag}:{mode_suffix}:background={is_background}:" + f"rca={rca_flag}:postmortem={is_postmortem_action}:is_rca_ctx={is_rca_context}:" + f"pr_review={is_pr_review}:action_id={trigger_action_id or ''}" + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1027 - 1033, The cache_key construction in the line containing the f-string assignment is missing the trigger_action_id parameter, which causes different actions with identical other parameters to share the same cached tool and reuse stale action IDs. Extract trigger_action_id from state_context using getattr (similar to how trigger_rca_requested, is_background, is_postmortem_action, and is_pr_review are extracted) and append it to the cache_key f-string to ensure each action gets its own distinct cache entry.server/chat/background/task.py (1)
829-830:⚠️ Potential issue | 🟠 Major | ⚡ Quick winSend the current in-memory assistant response to Slack, not the last persisted one.
The early Slack path calls
_send_response_to_slack()and then suppresses the parent fallback when it returnsTrue, but the helper querieschat_sessions.messagesand picks the last persisted assistant message. If the current turn has not flushed yet—or a guardrail-blocked Slack turn produced no new assistant message—this can update Slack with a previous response and skip the real fallback.Suggested direction
- _slack_early_sent = _send_response_to_slack(user_id, session_id, trigger_metadata) + _slack_early_sent = _send_response_to_slack( + user_id, + session_id, + trigger_metadata, + response_text=_latest_assistant_text(state.messages), + ) @@ -def _send_response_to_slack(user_id: str, session_id: str, trigger_metadata: Dict[str, Any]) -> bool: +def _send_response_to_slack( + user_id: str, + session_id: str, + trigger_metadata: Dict[str, Any], + response_text: Optional[str] = None, +) -> bool: @@ - # Get the last assistant message from the chat session - with db_pool.get_admin_connection() as conn: + last_assistant_message = response_text + if not last_assistant_message: + # Keep the existing DB lookup as the parent-task fallback path. + with db_pool.get_admin_connection() as conn:Add
_latest_assistant_text(...)near the Slack helper and returnFalsefor early send when the current state has no assistant response.Also applies to: 1461-1465, 1970-2016
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/background/task.py` around lines 829 - 830, The issue is that the early Slack send path in the code around line 829 calls _send_response_to_slack() which retrieves the last persisted assistant message from the database, but this may be stale if the current turn hasn't flushed yet or was blocked by a guardrail. To fix this, create a new helper function _latest_assistant_text() that returns the current in-memory assistant response for the present turn rather than querying persisted messages, then use this helper in the early send check to ensure Slack receives the actual current response, and return False if no current assistant response exists to avoid sending outdated messages.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/chat/background/task.py`:
- Around line 692-698: The `notify_action_started` call and other dispatcher
notification calls (such as `notify_investigation_failed`,
`notify_action_completed`) are not gated by the `send_notifications` flag,
causing notifications to be sent even when callers opt out. Wrap the try-except
block containing the `notify_action_started` call with an `if
send_notifications:` guard at the beginning, and apply the same gating pattern
to all other `notify_*` calls throughout the file at the locations mentioned
(around lines 843-851, 872-888, 901-917, 1304, 1541-1567). Keep any database
status update operations outside these guards so they always execute regardless
of the notification flag.
In `@server/guardrails/input_rail.py`:
- Around line 177-189: The issue is that `_rails_thread_lock` is being acquired
on the event loop in the async code path (referenced in the range 207-214),
which blocks the workflow before the work can be moved to a worker thread,
preventing concurrent guardrail/prompt work during cold start. Extract the lock
acquisition and rails construction logic into a new synchronous helper function
(similar in structure to `prewarm_rails_sync()`), then call this helper via
`asyncio.to_thread()` from the async code paths instead of acquiring the lock
directly in the coroutine. This moves the lock contention off the event loop and
into the worker thread where it belongs.
In `@server/routes/slack/slack_events.py`:
- Line 363: In the logger.warning call, convert the f-string formatting to lazy
formatting using percent placeholders. Replace the f-string with a string
containing `%s` placeholders for clicker_user_id, clicker_org_id, and
incident_org_id, and pass the variables as a tuple argument to the
logger.warning method using the `%` operator. This ensures string formatting
only occurs when the log level is actually enabled, avoiding unnecessary
computation when warnings are disabled.
---
Outside diff comments:
In `@server/chat/backend/agent/tools/cloud_tools.py`:
- Around line 1027-1033: The cache_key construction in the line containing the
f-string assignment is missing the trigger_action_id parameter, which causes
different actions with identical other parameters to share the same cached tool
and reuse stale action IDs. Extract trigger_action_id from state_context using
getattr (similar to how trigger_rca_requested, is_background,
is_postmortem_action, and is_pr_review are extracted) and append it to the
cache_key f-string to ensure each action gets its own distinct cache entry.
In `@server/chat/background/task.py`:
- Around line 829-830: The issue is that the early Slack send path in the code
around line 829 calls _send_response_to_slack() which retrieves the last
persisted assistant message from the database, but this may be stale if the
current turn hasn't flushed yet or was blocked by a guardrail. To fix this,
create a new helper function _latest_assistant_text() that returns the current
in-memory assistant response for the present turn rather than querying persisted
messages, then use this helper in the early send check to ensure Slack receives
the actual current response, and return False if no current assistant response
exists to avoid sending outdated messages.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 63c9984d-3e80-4cd4-b99d-15386ad83809
📒 Files selected for processing (8)
server/celery_config.pyserver/chat/backend/agent/agent.pyserver/chat/backend/agent/tools/cloud_tools.pyserver/chat/backend/agent/workflow.pyserver/chat/background/task.pyserver/guardrails/input_rail.pyserver/routes/slack/slack_events.pyserver/routes/slack/slack_events_helpers.py
Contributor
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
server/chat/backend/agent/tools/cloud_tools.py (1)
1027-1033:⚠️ Potential issue | 🟠 MajorCache key missing
trigger_action_id– stale action IDs can execute.The cache key (line 1033) omits
trigger_action_id, but_action_idis extracted at line 1340 and pinned to tools via closure. When a different action later hits the cache with identical other parameters, the cached tool carries the original action ID, not the current one.Include
trigger_action_idin the cache key to differentiate cache entries by action:Proposed fix
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) + trigger_action_id = getattr(state_context, 'trigger_action_id', None) if state_context else None 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}" + cache_key = ( + f"{user_id}:{capture_tag}:{mode_suffix}:background={is_background}:" + f"rca={rca_flag}:postmortem={is_postmortem_action}:is_rca_ctx={is_rca_context}:" + f"pr_review={is_pr_review}:action_id={trigger_action_id or ''}" + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/backend/agent/tools/cloud_tools.py` around lines 1027 - 1033, The cache_key construction in the line containing the f-string assignment is missing the trigger_action_id parameter, which causes different actions with identical other parameters to share the same cached tool and reuse stale action IDs. Extract trigger_action_id from state_context using getattr (similar to how trigger_rca_requested, is_background, is_postmortem_action, and is_pr_review are extracted) and append it to the cache_key f-string to ensure each action gets its own distinct cache entry.server/chat/background/task.py (1)
829-830:⚠️ Potential issue | 🟠 Major | ⚡ Quick winSend the current in-memory assistant response to Slack, not the last persisted one.
The early Slack path calls
_send_response_to_slack()and then suppresses the parent fallback when it returnsTrue, but the helper querieschat_sessions.messagesand picks the last persisted assistant message. If the current turn has not flushed yet—or a guardrail-blocked Slack turn produced no new assistant message—this can update Slack with a previous response and skip the real fallback.Suggested direction
- _slack_early_sent = _send_response_to_slack(user_id, session_id, trigger_metadata) + _slack_early_sent = _send_response_to_slack( + user_id, + session_id, + trigger_metadata, + response_text=_latest_assistant_text(state.messages), + ) @@ -def _send_response_to_slack(user_id: str, session_id: str, trigger_metadata: Dict[str, Any]) -> bool: +def _send_response_to_slack( + user_id: str, + session_id: str, + trigger_metadata: Dict[str, Any], + response_text: Optional[str] = None, +) -> bool: @@ - # Get the last assistant message from the chat session - with db_pool.get_admin_connection() as conn: + last_assistant_message = response_text + if not last_assistant_message: + # Keep the existing DB lookup as the parent-task fallback path. + with db_pool.get_admin_connection() as conn:Add
_latest_assistant_text(...)near the Slack helper and returnFalsefor early send when the current state has no assistant response.Also applies to: 1461-1465, 1970-2016
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/background/task.py` around lines 829 - 830, The issue is that the early Slack send path in the code around line 829 calls _send_response_to_slack() which retrieves the last persisted assistant message from the database, but this may be stale if the current turn hasn't flushed yet or was blocked by a guardrail. To fix this, create a new helper function _latest_assistant_text() that returns the current in-memory assistant response for the present turn rather than querying persisted messages, then use this helper in the early send check to ensure Slack receives the actual current response, and return False if no current assistant response exists to avoid sending outdated messages.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/chat/background/task.py`:
- Around line 692-698: The `notify_action_started` call and other dispatcher
notification calls (such as `notify_investigation_failed`,
`notify_action_completed`) are not gated by the `send_notifications` flag,
causing notifications to be sent even when callers opt out. Wrap the try-except
block containing the `notify_action_started` call with an `if
send_notifications:` guard at the beginning, and apply the same gating pattern
to all other `notify_*` calls throughout the file at the locations mentioned
(around lines 843-851, 872-888, 901-917, 1304, 1541-1567). Keep any database
status update operations outside these guards so they always execute regardless
of the notification flag.
In `@server/guardrails/input_rail.py`:
- Around line 177-189: The issue is that `_rails_thread_lock` is being acquired
on the event loop in the async code path (referenced in the range 207-214),
which blocks the workflow before the work can be moved to a worker thread,
preventing concurrent guardrail/prompt work during cold start. Extract the lock
acquisition and rails construction logic into a new synchronous helper function
(similar in structure to `prewarm_rails_sync()`), then call this helper via
`asyncio.to_thread()` from the async code paths instead of acquiring the lock
directly in the coroutine. This moves the lock contention off the event loop and
into the worker thread where it belongs.
In `@server/routes/slack/slack_events.py`:
- Line 363: In the logger.warning call, convert the f-string formatting to lazy
formatting using percent placeholders. Replace the f-string with a string
containing `%s` placeholders for clicker_user_id, clicker_org_id, and
incident_org_id, and pass the variables as a tuple argument to the
logger.warning method using the `%` operator. This ensures string formatting
only occurs when the log level is actually enabled, avoiding unnecessary
computation when warnings are disabled.
---
Outside diff comments:
In `@server/chat/backend/agent/tools/cloud_tools.py`:
- Around line 1027-1033: The cache_key construction in the line containing the
f-string assignment is missing the trigger_action_id parameter, which causes
different actions with identical other parameters to share the same cached tool
and reuse stale action IDs. Extract trigger_action_id from state_context using
getattr (similar to how trigger_rca_requested, is_background,
is_postmortem_action, and is_pr_review are extracted) and append it to the
cache_key f-string to ensure each action gets its own distinct cache entry.
In `@server/chat/background/task.py`:
- Around line 829-830: The issue is that the early Slack send path in the code
around line 829 calls _send_response_to_slack() which retrieves the last
persisted assistant message from the database, but this may be stale if the
current turn hasn't flushed yet or was blocked by a guardrail. To fix this,
create a new helper function _latest_assistant_text() that returns the current
in-memory assistant response for the present turn rather than querying persisted
messages, then use this helper in the early send check to ensure Slack receives
the actual current response, and return False if no current assistant response
exists to avoid sending outdated messages.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 63c9984d-3e80-4cd4-b99d-15386ad83809
📒 Files selected for processing (8)
server/celery_config.pyserver/chat/backend/agent/agent.pyserver/chat/backend/agent/tools/cloud_tools.pyserver/chat/backend/agent/workflow.pyserver/chat/background/task.pyserver/guardrails/input_rail.pyserver/routes/slack/slack_events.pyserver/routes/slack/slack_events_helpers.py
🛑 Comments failed to post (3)
server/chat/background/task.py (1)
692-698:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winHonor
send_notificationsfor all per-task dispatcher notifications.
send_notificationsgates investigation-start notifications, but the new action-start/action-complete and failure dispatcher calls still run when a caller opts out;_execute_background_chat()also accepts the flag without using it. Gate only thenotify_*calls on this flag and keep DB status updates unchanged.Suggested gating pattern
- if is_action_source: + if is_action_source and send_notifications: try: from utils.notifications.dispatcher import notify_action_started notify_action_started(user_id, trigger_metadata, session_id) @@ - if trigger_metadata and trigger_metadata.get('source') == 'action': + if send_notifications and trigger_metadata and trigger_metadata.get('source') == 'action': if not result.get("action_notification_sent"): action_status = 'error' if result.get("guardrail_blocked") else 'success' @@ - if action_status: + if send_notifications and action_status: try: from utils.notifications.dispatcher import notify_action_completed notify_action_completed(user_id, trigger_metadata, session_id, status=action_status, error_message=action_error_msg)Apply the same
if send_notifications:guard around thenotify_investigation_failed(...)and timeout/errornotify_action_completed(...)blocks.Also applies to: 843-851, 872-888, 901-917, 1304-1304, 1541-1567
🧰 Tools
🪛 Ruff (0.15.17)
[warning] 697-697: Do not catch blind exception:
Exception(BLE001)
[warning] 698-698: Logging statement uses f-string
(G004)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/chat/background/task.py` around lines 692 - 698, The `notify_action_started` call and other dispatcher notification calls (such as `notify_investigation_failed`, `notify_action_completed`) are not gated by the `send_notifications` flag, causing notifications to be sent even when callers opt out. Wrap the try-except block containing the `notify_action_started` call with an `if send_notifications:` guard at the beginning, and apply the same gating pattern to all other `notify_*` calls throughout the file at the locations mentioned (around lines 843-851, 872-888, 901-917, 1304, 1541-1567). Keep any database status update operations outside these guards so they always execute regardless of the notification flag.Source: Linters/SAST tools
server/guardrails/input_rail.py (1)
177-189:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winMove the rails construction lock off the event loop.
_get_rails()can run while the prewarm thread owns_rails_thread_lock; acquiring thatthreading.Lockinside the coroutine blocks the workflow event loop beforeasyncio.to_thread()runs, defeating the concurrent guardrail/prompt work during cold start. Put the shared lock acquisition inside a sync helper that runs in the worker thread, and call that helper from both paths.Suggested lock handoff
+def _ensure_rails_sync(): + """Build or return the cached rails instance under the thread lock.""" + global _rails_instance, _last_init_failure_ts + if _rails_instance is not None: + return _rails_instance + with _rails_thread_lock: + if _rails_instance is not None: + return _rails_instance + try: + _rails_instance = _build_rails_sync() + except Exception: + _last_init_failure_ts = time.monotonic() + raise + return _rails_instance + + def prewarm_rails_sync() -> None: """Sync init for Celery worker prewarm — thread-safe with async lazy init.""" - global _rails_instance, _last_init_failure_ts - if _rails_instance is not None: - return - with _rails_thread_lock: - if _rails_instance is not None: - return - try: - _rails_instance = _build_rails_sync() - except Exception: - _last_init_failure_ts = time.monotonic() - raise + _ensure_rails_sync() @@ async with _get_lock(): - with _rails_thread_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 _rails_instance is not None: + return _rails_instance + return await asyncio.to_thread(_ensure_rails_sync)Also applies to: 207-214
🧰 Tools
🪛 Ruff (0.15.17)
[warning] 179-179: Using the global statement to update
_rails_instanceis discouraged(PLW0603)
[warning] 179-179: Using the global statement to update
_last_init_failure_tsis discouraged(PLW0603)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/guardrails/input_rail.py` around lines 177 - 189, The issue is that `_rails_thread_lock` is being acquired on the event loop in the async code path (referenced in the range 207-214), which blocks the workflow before the work can be moved to a worker thread, preventing concurrent guardrail/prompt work during cold start. Extract the lock acquisition and rails construction logic into a new synchronous helper function (similar in structure to `prewarm_rails_sync()`), then call this helper via `asyncio.to_thread()` from the async code paths instead of acquiring the lock directly in the coroutine. This moves the lock contention off the event loop and into the worker thread where it belongs.server/routes/slack/slack_events.py (1)
363-363:
⚠️ Potential issue | 🟡 Minor🧩 Analysis chain
🏁 Script executed:
#!/bin/bash # Verify remaining logger f-strings in this file rg -nP 'logger\.(debug|info|warning|error|exception|critical)\(\s*f"' server/routes/slack/slack_events.pyRepository: Arvo-AI/aurora
Length of output: 2599
Use lazy logger formatting at Line 363.
The logger call uses an f-string, which triggers Ruff G004 and eagerly formats even when the log level is disabled. Convert to lazy formatting using
%placeholders:Proposed fix
- logger.warning(f"User {clicker_user_id} (org {clicker_org_id}) tried to run suggestion for incident in org {incident_org_id}") + logger.warning( + "User %s (org %s) tried to run suggestion for incident in org %s", + clicker_user_id, + clicker_org_id, + incident_org_id, + )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.logger.warning( "User %s (org %s) tried to run suggestion for incident in org %s", clicker_user_id, clicker_org_id, incident_org_id, )🧰 Tools
🪛 Ruff (0.15.17)
[warning] 363-363: Logging statement uses f-string
(G004)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/routes/slack/slack_events.py` at line 363, In the logger.warning call, convert the f-string formatting to lazy formatting using percent placeholders. Replace the f-string with a string containing `%s` placeholders for clicker_user_id, clicker_org_id, and incident_org_id, and pass the variables as a tuple argument to the logger.warning method using the `%` operator. This ensures string formatting only occurs when the log level is actually enabled, avoiding unnecessary computation when warnings are disabled.Source: Linters/SAST tools
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/guardrails/input_rail.py`:
- Around line 183-191: In the prewarm_rails_sync function, the direct
acquisition of _rails_thread_lock blocks the event loop when the Celery prewarm
daemon thread holds it during _build_rails_sync execution. Move the critical
section (lock acquisition, backoff validation, and _build_rails_sync call) into
a worker thread using asyncio.to_thread() so the event loop remains unblocked
for concurrent requests. Additionally, re-check the backoff condition inside the
critical section after acquiring _rails_thread_lock to prevent the race
condition where a caller passes the backoff check before the lock but gets
delayed waiting for it, then bypasses the backoff window when finally acquiring
the lock.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: c20d39fa-ba8f-447e-b3cd-dc0ec0a7aa48
📒 Files selected for processing (1)
server/guardrails/input_rail.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/guardrails/input_rail.py`:
- Around line 183-202: The error message "input rail init recently failed;
backing off" is duplicated three times in the module. Extract this literal
string to a module-level constant near the top of the file with other constants,
naming it something like _INIT_BACKOFF_ERROR. Then replace all three occurrences
of the literal string (one in the _ensure_rails_in_thread function at line 196
and two other locations in the file) with references to this new constant to
improve maintainability.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 57c33ee1-c41f-45f6-86f5-65279525deb2
📒 Files selected for processing (2)
server/celery_config.pyserver/guardrails/input_rail.py
beng360
approved these changes
Jun 22, 2026
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this does
Fixes the slow chat replies we've been seeing on Slack and UI (10-30s down to ~5s consistent).
The main problems were:
acks_late=True+ Redis broker caused redelivery while the task was still executing. This made Slack replies "change" after appearing.Changes
Reliability fixes:
execto celery worker command so the process gets SIGTERM directly from K8s (prevents ghost consumers after deploys)Performance:
worker_process_initpre-warms guardrails, MCP, and Agent singleton on every child fork — no more cold starts even whenmax_tasks_per_childrecycles workerscheck_inputfires concurrently with prompt assembly and tool loading (only blocks right before the LLM call)build_prompt_segmentsandget_cloud_toolsrun in parallel via ThreadPoolExecutorid(tool_capture)— guaranteed a cache miss every single messageInstrumentation:
[LATENCY]logs at every stage: queue wait, agent init, guardrails, prompt build, LLM TTFT, Slack send, E2E totalsSummary by CodeRabbit
Performance Improvements
Reliability
Slack Integration
Safety