stop functionality & partial db flush #489
Conversation
WalkthroughThis PR adds cancellation and task management improvements across the agent execution pipeline: buffer flushing before task cancellation in agent tasks, enhanced status computation in base task cleanup, dependency injection for session and Redis management in conversation service, and explicit task ID tracking for revocation and session cleanup. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Router as conversations_router
participant ConvService as conversation_service
participant AgentTask as agent_tasks
participant Redis as redis_streaming
participant Celery
Client->>Router: POST cancel/stop_generation
Router->>ConvService: stop_generation(run_id)
alt run_id missing
ConvService->>ConvService: Fetch active session via session_service
alt No active session
ConvService-->>Client: Success (no active session)
end
end
ConvService->>Redis: Set cancellation flag
ConvService->>Redis: get_task_id(conversation_id, run_id)
Redis-->>ConvService: task_id
ConvService->>Celery: revoke(task_id, terminate=True)
Celery-->>AgentTask: Cancellation signal
AgentTask->>AgentTask: Flush buffered message chunks
AgentTask->>Redis: Publish flushed chunks
AgentTask->>Redis: Return completion status (False)
ConvService->>Redis: clear_session(conversation_id, run_id)
Redis->>Redis: Publish end event (status=cancelled)
Redis-->>ConvService: Session cleared
ConvService-->>Client: Cancellation signal sent & task revoked
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~25 minutes Areas requiring extra attention:
Possibly Related PRs
Suggested Reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 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 |
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (6)
app/celery/tasks/base_task.py (1)
85-91: DB cleanup improvement looks good; clarify retval contractWrapping the status log and DB cleanup in a
try/finallyis a solid improvement and ensuresSessionLocal()is always closed and dereferenced.This now implicitly assumes that BaseTask tasks return a strict boolean, where
Falsemeans “cancelled” and anything else means “completed successfully”. If any existing tasks usingBaseTaskreturnFalsein non‑cancellation cases, their logs will now say “cancelled”.If that contract is intentional and consistent across all
base=BaseTasktasks, this change is fine; otherwise, consider either:
- documenting the convention (“BaseTask tasks must return True/False for success/cancel”), or
- narrowing the “cancelled” check to only the specific tasks that actually use this pattern.
app/modules/conversations/utils/redis_streaming.py (1)
200-234: Task‑ID tracking and clear_session behavior look consistent; consider richer error loggingThe new
set_task_id/get_task_idAPIs andclear_sessionflow are consistent with the existing Redis naming/TTL patterns and givestop_generationwhat it needs to revoke Celery tasks and mark sessions as cancelled.Given
clear_sessionis intentionally best‑effort and non‑throwing, the broadexcept Exceptionis reasonable. If you want more debuggability with minimal behavior change, consider:
- logging with
exc_info=True(orlogger.exception) so stack traces are preserved, and/or- narrowing the catch to Redis‑related exceptions.
Functionally this looks correct as‑is.
app/celery/tasks/agent_tasks.py (3)
103-116: Cancellation now flushes partial AI output before endingThe added flush in the agent cancellation path ensures any buffered AI‑generated chunks are persisted before you emit the
"end"/"cancelled"event and bail out, which is exactly what you want for partial‑response durability.Swallowing all exceptions from
flush_message_bufferwith a warning keeps cancellation robust; if you want more diagnosability without changing behavior, you could log withexc_info=Trueor catchChatHistoryServiceErrorexplicitly and let truly unexpected exceptions propagate.
284-297: Regenerate cancellation flush mirrors agent behavior correctlyThe regeneration path now mirrors the agent path by flushing
MessageType.AI_GENERATEDfromservice.history_manageron cancellation, logging the resulting message ID, and warning (but not failing) if the flush raises.This keeps the regeneration flow’s persistence semantics aligned with the main agent execution and still ensures cancellation proceeds even if DB writes fail.
156-180: Returning a boolean completion flag—update type hints or document the conventionBoth background tasks now return
completed: boolfrom the outer Celery task function so thatBaseTask.on_successcan distinguish cancellation (False) from normal completion (True). Given the inner coroutines only ever returnTrueorFalse(exceptions go to theexceptblock), this is a clear and reliable signal.Two minor nits to consider:
- The function signatures still declare
-> None, which no longer matches the implementation. If you care about static typing, consider changing the annotations to-> boolor dropping the explicit return type.- It would be worth briefly documenting (in a docstring or comment) that these tasks follow the “True = completed, False = cancelled” convention so future BaseTask subclasses can depend on it consistently.
Also applies to: 347-371
app/modules/conversations/conversation/conversation_service.py (1)
1080-1145: stop_generation pipeline is well‑structured; response message could be more preciseThe updated
stop_generationflow is coherent end‑to‑end:
- If
run_idis omitted, it usesSessionService.get_active_sessionto find the latest session and returns early with a success message when none exists.- It always sets the Redis cancellation flag so background tasks polling
check_cancellationcan exit gracefully.- When a stored task ID is present, it revokes the Celery task with
terminate=True, covering both queued and running tasks, and logs either success or failure of the revoke.- Regardless of task ID, it calls
clear_sessionso clients see an"end"event with a cancelled status and the session status is updated, which also handles stale session IDs.One small improvement: the final response message is always
"Cancellation signal sent and task revoked", even when:
- no task ID was found, or
- revocation raised and you logged a warning.
To avoid misleading callers, consider wording like
"Cancellation signal sent; any associated task will be revoked if possible"or tailoring the message based on whether a task_id was actually found and revoked.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
app/celery/tasks/agent_tasks.py(5 hunks)app/celery/tasks/base_task.py(1 hunks)app/modules/conversations/conversation/conversation_service.py(6 hunks)app/modules/conversations/conversations_router.py(3 hunks)app/modules/conversations/utils/redis_streaming.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (4)
app/celery/tasks/agent_tasks.py (2)
app/modules/intelligence/memory/chat_history_service.py (1)
flush_message_buffer(83-133)app/modules/conversations/message/message_model.py (1)
MessageType(17-20)
app/modules/conversations/conversation/conversation_service.py (4)
app/modules/conversations/session/session_service.py (2)
SessionService(15-163)get_active_session(23-98)app/modules/conversations/utils/redis_streaming.py (4)
RedisStreamManager(11-247)set_cancellation(182-186)get_task_id(206-210)clear_session(212-233)app/modules/conversations/conversation/conversation_schema.py (1)
ActiveSessionErrorResponse(78-80)app/modules/conversations/conversations_router.py (1)
get_active_session(568-594)
app/modules/conversations/conversations_router.py (2)
app/celery/tasks/agent_tasks.py (2)
execute_agent_background(16-201)execute_regenerate_background(209-391)app/modules/conversations/utils/redis_streaming.py (1)
set_task_id(200-204)
app/modules/conversations/utils/redis_streaming.py (1)
tests/conftest.py (1)
get(142-150)
🪛 Ruff (0.14.4)
app/celery/tasks/agent_tasks.py
112-112: Do not catch blind exception: Exception
(BLE001)
114-114: Use explicit conversion flag
Replace with conversion flag
(RUF010)
180-180: Consider moving this statement to an else block
(TRY300)
293-293: Do not catch blind exception: Exception
(BLE001)
295-295: Use explicit conversion flag
Replace with conversion flag
(RUF010)
371-371: Consider moving this statement to an else block
(TRY300)
app/modules/conversations/conversation/conversation_service.py
1122-1122: Do not catch blind exception: Exception
(BLE001)
1123-1123: Use explicit conversion flag
Replace with conversion flag
(RUF010)
1136-1136: Do not catch blind exception: Exception
(BLE001)
1138-1138: Use explicit conversion flag
Replace with conversion flag
(RUF010)
app/modules/conversations/utils/redis_streaming.py
230-230: Do not catch blind exception: Exception
(BLE001)
231-233: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
232-232: Use explicit conversion flag
Replace with conversion flag
(RUF010)
🔇 Additional comments (3)
app/modules/conversations/conversations_router.py (2)
358-372: Capturing and storing agent task ID is correct and aligns with stop flowUsing the
AsyncResultfromexecute_agent_background.delay(...)and persistingtask_result.idviaredis_manager.set_task_idcleanly connects request initiation to the later revocation path inConversationService.stop_generation.Given that you already set status to
"queued"and publish the queued event before starting the task, this fits the existing semantics with no extra failure modes.
493-505: Regenerate task ID tracking mirrors agent path appropriatelyThe regenerate endpoint now mirrors the agent endpoint by capturing the Celery
AsyncResultand storingtask_result.idin Redis. This keeps the stop/cancellation semantics consistent across both flows and reuses the same Redis keying scheme.No functional issues here.
app/modules/conversations/conversation/conversation_service.py (1)
88-107: SessionService/RedisStreamManager DI looks sound and backward‑compatibleInjecting
SessionServiceandRedisStreamManagervia__init__while still defaulting to concrete instances—and wiring them throughConversationService.create()—keeps the public construction path intact but makes testing and future customization much easier.Storing
self.celery_apphere is a reasonable compromise to allow stop‑related logic inside the service without forcing every caller to pass a Celery app explicitly. I don’t see any obvious behavioral regressions from these changes.Also applies to: 128-147
Summary by CodeRabbit
Bug Fixes
Refactor