2323from utils .prompt_template_utils import normalize_prompt_generate_template_content
2424from consts .const import MEMORY_SEARCH_START_MSG , MEMORY_SEARCH_DONE_MSG , MEMORY_SEARCH_FAIL_MSG , TOOL_TYPE_MAPPING , \
2525 LANGUAGE , MESSAGE_ROLE , MODEL_CONFIG_MAPPING , CAN_EDIT_ALL_USER_ROLES , PERMISSION_PRIVATE , STREAM_STATUS_EVENT , \
26- DEFAULT_EN_TITLE , DEFAULT_ZH_TITLE
26+ DEFAULT_EN_TITLE , DEFAULT_ZH_TITLE , RUNTIME_CANCEL_POLL_INTERVAL_SECONDS
2727from consts .exceptions import AppException , MemoryPreparationException , SkillDuplicateError
2828from consts .error_code import ErrorCode
2929from consts .agent_unavailable_reasons import AgentUnavailableReason
107107)
108108from services .memory_config_service import build_memory_context
109109from services .streaming_channel import streaming_channel_manager
110+ from services .runtime_state_service import runtime_state_service
110111from utils .auth_utils import get_current_user_info , get_user_language
111112from utils .config_utils import tenant_config_manager
112113from utils .memory_utils import build_memory_config
@@ -133,6 +134,34 @@ async def _cleanup_channel_later(conversation_id: int, user_id: str, delay: floa
133134 await streaming_channel_manager .remove_channel (conversation_id , user_id )
134135
135136
137+ async def _poll_runtime_cancel_signal (conversation_id : int , user_id : str , stop_event ) -> None :
138+ """Mirror Redis cancel signal into the local agent stop_event."""
139+ while not stop_event .is_set ():
140+ if await runtime_state_service .is_cancelled_async (user_id = user_id , conversation_id = conversation_id ):
141+ stop_event .set ()
142+ logger .info (
143+ "Runtime cancel signal received, user_id=%s, conversation_id=%s" ,
144+ user_id ,
145+ conversation_id ,
146+ )
147+ return
148+ await asyncio .sleep (RUNTIME_CANCEL_POLL_INTERVAL_SECONDS )
149+
150+
151+ async def _cancel_task_on_runtime_signal (conversation_id : int , user_id : str , task : asyncio .Task ) -> None :
152+ """Cancel a local asyncio task when another Pod writes the runtime cancel signal."""
153+ while not task .done ():
154+ if await runtime_state_service .is_cancelled_async (user_id = user_id , conversation_id = conversation_id ):
155+ task .cancel ()
156+ logger .info (
157+ "Runtime cancel signal cancelled task, user_id=%s, conversation_id=%s" ,
158+ user_id ,
159+ conversation_id ,
160+ )
161+ return
162+ await asyncio .sleep (RUNTIME_CANCEL_POLL_INTERVAL_SECONDS )
163+
164+
136165def _extract_json_objects_from_text (text : str ) -> list [dict ]:
137166 """Extract all JSON objects embedded in a text blob."""
138167 if not text :
@@ -934,6 +963,14 @@ async def _stream_agent_chunks(
934963 user_id = user_id
935964 )
936965
966+ cancel_poll_task = asyncio .create_task (
967+ _poll_runtime_cancel_signal (
968+ conversation_id = agent_request .conversation_id ,
969+ user_id = user_id ,
970+ stop_event = agent_run_info .stop_event ,
971+ )
972+ )
973+
937974 # In resume mode, emit a status event first
938975 if is_resume_mode :
939976 await channel .publish (STREAM_STATUS_EVENT )
@@ -1196,7 +1233,8 @@ async def _stream_agent_chunks(
11961233 except Exception :
11971234 logger .exception ("Failed to mark last unit as completed" )
11981235
1199- terminal_status = "completed" if stream_completed_normally else "failed"
1236+ was_stopped = getattr (agent_run_info , "stop_event" , None ) and agent_run_info .stop_event .is_set ()
1237+ terminal_status = "stopped" if was_stopped else "completed" if stream_completed_normally else "failed"
12001238 try :
12011239 update_message_status (
12021240 streaming_message_id ,
@@ -1206,12 +1244,17 @@ async def _stream_agent_chunks(
12061244 except Exception :
12071245 logger .exception ("Failed to mark assistant message as %s" , terminal_status )
12081246
1247+ if not cancel_poll_task .done ():
1248+ cancel_poll_task .cancel ()
1249+
1250+ was_stopped = getattr (agent_run_info , "stop_event" , None ) and agent_run_info .stop_event .is_set ()
1251+ terminal_status = 'stopped' if was_stopped else 'completed' if stream_completed_normally else 'failed'
1252+
12091253 agent_run_manager .unregister_agent_run (
1210- agent_request .conversation_id , user_id )
1254+ agent_request .conversation_id , user_id , status = terminal_status )
12111255
12121256 # Mark channel as completed and schedule cleanup
12131257 if channel is not None :
1214- terminal_status = 'completed' if stream_completed_normally else 'failed'
12151258 await streaming_channel_manager .complete_channel (
12161259 conversation_id = agent_request .conversation_id ,
12171260 user_id = user_id ,
@@ -2747,6 +2790,11 @@ async def generate_stream_with_memory(
27472790 preprocess_manager .register_preprocess_task (
27482791 task_id , conversation_id , current_task
27492792 )
2793+ cancel_poll_task = (
2794+ asyncio .create_task (_cancel_task_on_runtime_signal (conversation_id , user_id , current_task ))
2795+ if current_task
2796+ else None
2797+ )
27502798
27512799 # Helper to emit memory_search token
27522800 def _memory_token (message_text : str ) -> str :
@@ -2845,6 +2893,8 @@ def _memory_token(message_text: str) -> str:
28452893 yield _safe_agent_stream_error_chunk ()
28462894 return
28472895 finally :
2896+ if cancel_poll_task and not cancel_poll_task .done ():
2897+ cancel_poll_task .cancel ()
28482898 # Always unregister preprocess task
28492899 preprocess_manager .unregister_preprocess_task (task_id )
28502900
@@ -3035,8 +3085,13 @@ async def run_agent_stream(
30353085 user_id = resolved_user_id ,
30363086 conversation_id = agent_request .conversation_id
30373087 )
3088+ run_state = await runtime_state_service .get_run_state_async (
3089+ user_id = resolved_user_id ,
3090+ conversation_id = agent_request .conversation_id ,
3091+ )
3092+ is_remote_running = run_state .get ("status" ) == "running"
30383093
3039- if existing_run_info is None :
3094+ if existing_run_info is None and not is_remote_running :
30403095 # Agent has finished while frontend was disconnected
30413096 # Update message status to completed if it's still streaming
30423097 try :
@@ -3060,8 +3115,82 @@ async def run_agent_stream(
30603115 conversation_id = agent_request .conversation_id ,
30613116 user_id = resolved_user_id
30623117 )
3118+ last_unit_index = resume_info ["resume_from_unit_index" ] - 1
3119+
3120+ def _resume_status_chunk (replay_chunk_count : int ) -> str :
3121+ payload = {
3122+ 'status' : 'resumed' ,
3123+ 'last_unit_index' : last_unit_index ,
3124+ 'replay_chunk_count' : replay_chunk_count ,
3125+ }
3126+ return f"data: { json .dumps (payload )} \n \n "
3127+
3128+ def _resume_completed_chunk (status : str = "completed" ) -> str :
3129+ payload = {
3130+ 'status' : status ,
3131+ 'last_unit_index' : last_unit_index ,
3132+ }
3133+ return f"data: { json .dumps (payload )} \n \n "
30633134
30643135 if channel is None :
3136+ if runtime_state_service .enabled and is_remote_running :
3137+ async def redis_channel_stream ():
3138+ replay_events = await runtime_state_service .read_stream_events_async (
3139+ user_id = resolved_user_id ,
3140+ conversation_id = agent_request .conversation_id ,
3141+ )
3142+ replay_chunk_count = len (replay_events )
3143+
3144+ yield STREAM_STATUS_EVENT
3145+ yield _resume_status_chunk (replay_chunk_count )
3146+
3147+ last_event_id = "0-0"
3148+ for event_id , chunk in replay_events :
3149+ last_event_id = event_id
3150+ if chunk :
3151+ yield chunk
3152+
3153+ while True :
3154+ events = await runtime_state_service .wait_for_stream_events_async (
3155+ user_id = resolved_user_id ,
3156+ conversation_id = agent_request .conversation_id ,
3157+ last_id = last_event_id ,
3158+ )
3159+ for event_id , chunk in events :
3160+ last_event_id = event_id
3161+ if chunk :
3162+ yield chunk
3163+
3164+ stream_status = await runtime_state_service .get_stream_status_async (
3165+ user_id = resolved_user_id ,
3166+ conversation_id = agent_request .conversation_id ,
3167+ )
3168+ latest_run_state = await runtime_state_service .get_run_state_async (
3169+ user_id = resolved_user_id ,
3170+ conversation_id = agent_request .conversation_id ,
3171+ )
3172+ if stream_status .get ("status" ) or latest_run_state .get ("status" ) in {
3173+ "completed" ,
3174+ "failed" ,
3175+ "stopped" ,
3176+ }:
3177+ break
3178+
3179+ terminal_status = stream_status .get ("status" ) or latest_run_state .get ("status" ) or "completed"
3180+ yield STREAM_STATUS_EVENT
3181+ yield _resume_completed_chunk (terminal_status )
3182+
3183+ return StreamingResponse (
3184+ redis_channel_stream (),
3185+ media_type = "text/event-stream" ,
3186+ headers = {
3187+ "Cache-Control" : "no-cache" ,
3188+ "Connection" : "keep-alive" ,
3189+ "X-Stream-Status" : "resumed" ,
3190+ "X-Last-Unit-Index" : str (resume_info ['resume_from_unit_index' ]),
3191+ },
3192+ )
3193+
30653194 # No channel exists, agent might be in a different state
30663195 return JSONResponse (
30673196 status_code = HTTPStatus .OK ,
@@ -3078,7 +3207,7 @@ async def channel_stream():
30783207
30793208 # Emit status event first with chunk count for skip tracking
30803209 yield STREAM_STATUS_EVENT
3081- yield f'data: {{"status": "resumed", "last_unit_index": { resume_info [ "resume_from_unit_index" ] - 1 } , " replay_chunk_count": { replay_chunk_count } }} \n \n '
3210+ yield _resume_status_chunk ( replay_chunk_count )
30823211
30833212 # Use subscribe_with_history(0) to replay ALL chunks from the buffer
30843213 # This ensures no chunks are lost even if frontend disconnected during streaming
@@ -3088,7 +3217,7 @@ async def channel_stream():
30883217
30893218 # Mark as complete when channel ends
30903219 yield STREAM_STATUS_EVENT
3091- yield f'data: {{"status": "completed", "last_unit_index": { resume_info [ "resume_from_unit_index" ] - 1 } }} \n \n '
3220+ yield _resume_completed_chunk ()
30923221
30933222 return StreamingResponse (
30943223 channel_stream (),
@@ -3102,6 +3231,11 @@ async def channel_stream():
31023231 )
31033232
31043233 # Normal mode: start new stream
3234+ await runtime_state_service .reset_stream_async (
3235+ user_id = resolved_user_id ,
3236+ conversation_id = agent_request .conversation_id ,
3237+ )
3238+
31053239 if not agent_request .is_debug and not skip_user_save :
31063240 save_messages (
31073241 agent_request ,
0 commit comments