Skip to content
Merged
Show file tree
Hide file tree
Changes from 10 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 48 additions & 2 deletions backend/app/api/ws/chat_namespace.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,11 @@
from app.services.chat.rag import process_context_and_rag
from app.services.chat.storage import session_manager
from app.services.chat.storage.db import get_db_session, run_sync_in_executor
from app.services.chat.trigger import trigger_ai_response_unified
from app.services.chat.trigger import (
collect_completed_result,
persist_completed_result,
trigger_ai_response_unified,
)
from app.services.chat.wework_task_defaults import apply_wework_task_defaults
from app.services.task_fork_history import task_fork_history_resolver
from app.utils.prompt_utils import extract_display_prompt
Expand All @@ -84,6 +88,29 @@
logger = logging.getLogger(__name__)


async def _finalize_failed_ai_trigger(
*,
task_id: int,
assistant_subtask_id: int,
error_message: str,
error_code: str,
) -> None:
"""Persist a failed async AI trigger so follow-up messages are not blocked."""
final_result = await collect_completed_result(
assistant_subtask_id,
status="FAILED",
error_message=error_message,
error_code=error_code,
)
await persist_completed_result(
subtask_id=assistant_subtask_id,
task_id=task_id,
status="FAILED",
result=final_result,
error=error_message,
)


class ChatNamespace(socketio.AsyncNamespace):
"""
Socket.IO namespace for chat functionality.
Expand Down Expand Up @@ -989,12 +1016,31 @@ async def _trigger_ai():
)

error_code = classify_error(e)
error_message = format_error_message(e)

try:
await _finalize_failed_ai_trigger(
task_id=task.id,
assistant_subtask_id=assistant_subtask.id,
error_message=error_message,
error_code=error_code,
)
except Exception as finalize_error:
logger.error(
"[WS] chat:send failed to persist async trigger error: "
"task_id=%s, subtask_id=%s, error=%s",
task.id,
assistant_subtask.id,
finalize_error,
exc_info=True,
)

# Emit error to frontend so user sees the failure
await self.emit(
ServerEvents.CHAT_ERROR,
ChatErrorPayload(
subtask_id=assistant_subtask.id,
error=format_error_message(e),
error=error_message,
type=error_code,
message_id=assistant_subtask.message_id,
task_id=task.id,
Expand Down
3 changes: 0 additions & 3 deletions backend/app/core/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -178,9 +178,6 @@ class Settings(BaseSettings):
CHAT_TOOL_MAX_TIME_SECONDS: float = (
60.0 # Maximum time for tool calling flow (5 minutes)
)
# Source-level tool output compaction default for users without an explicit
# preference override.
TOOL_OUTPUT_GUARD_ENABLED: bool = False
# Group chat history configuration
# In group chat mode, AI-bot sees: first N messages + last M messages (no duplicates)
# If total messages < N + M, all messages are kept
Expand Down
64 changes: 63 additions & 1 deletion backend/app/services/chat/storage/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,11 @@
# Internal field stored in Redis block metadata. It is stripped before returning
# blocks to callers so API consumers still see the existing block shape.
BLOCK_CONTENT_KEY_FIELD = "_content_key"
UNRESOLVED_PREVIEW_TOOL_BLOCK_STATUSES = {"pending", "generating_arguments"}
UNRESOLVED_PREVIEW_TOOL_BLOCK_MESSAGE = "Tool call was not executed before the turn completed. The turn may have hit the tool-call limit."
UNRESOLVED_PREVIEW_TOOL_BLOCK_GENERIC_MESSAGE = (
"Tool call preview did not complete before the turn ended."
)


class StreamContentType(str, Enum):
Expand Down Expand Up @@ -1317,7 +1322,12 @@ async def get_blocks(self, subtask_id: int) -> List[Dict[str, Any]]:
)
return []

async def finalize_and_get_blocks(self, subtask_id: int) -> List[Dict[str, Any]]:
async def finalize_and_get_blocks(
self,
subtask_id: int,
*,
termination_reason: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""Finalize any pending text block and return all blocks.

This should be called when the subtask completes (DONE event).
Expand All @@ -1338,6 +1348,10 @@ async def finalize_and_get_blocks(self, subtask_id: int) -> List[Dict[str, Any]]
redis_client = await self._cache._get_client()
try:
blocks = await self._load_blocks_from_client(redis_client, blocks_key)
blocks = self._finalize_unresolved_preview_tool_blocks(
blocks,
termination_reason=termination_reason,
)
logger.info(
f"[SessionManager] Finalized blocks for subtask {subtask_id}: "
f"count={len(blocks)}"
Expand All @@ -1351,6 +1365,54 @@ async def finalize_and_get_blocks(self, subtask_id: int) -> List[Dict[str, Any]]
)
return []

def _finalize_unresolved_preview_tool_blocks(
self,
blocks: List[Dict[str, Any]],
*,
termination_reason: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""Convert unresolved preview-only tool blocks into terminal error blocks.

Preview tool blocks are created while tool arguments stream, before real tool
execution begins. If the turn ends without a matching TOOL_RESULT, those
blocks would otherwise stay stuck in a pending state in the final result.

Only preview-only blocks without tool_output are rewritten. Legitimate
pending blocks that already carry a semantic output payload (for example
deferred interactive forms waiting for user input) are preserved.
"""
inferred_tool_limit = (
termination_reason == "completed_with_unexecuted_tool_calls"
)
has_explicit_tool_error = any(
block.get("type") == "tool"
and block.get("status") == BlockStatus.ERROR.value
and block.get("tool_output")
for block in blocks
)
unresolved_message = (
UNRESOLVED_PREVIEW_TOOL_BLOCK_MESSAGE
if inferred_tool_limit and not has_explicit_tool_error
else UNRESOLVED_PREVIEW_TOOL_BLOCK_GENERIC_MESSAGE
)
finalized_blocks: List[Dict[str, Any]] = []
for block in blocks:
if (
block.get("type") == "tool"
and block.get("status") in UNRESOLVED_PREVIEW_TOOL_BLOCK_STATUSES
and not block.get("tool_output")
):
finalized_blocks.append(
{
**block,
"status": BlockStatus.ERROR.value,
"tool_output": unresolved_message,
}
)
continue
finalized_blocks.append(block)
return finalized_blocks

async def get_accumulated_content(self, subtask_id: int) -> str:
"""Get accumulated content for a subtask.

Expand Down
14 changes: 13 additions & 1 deletion backend/app/services/chat/trigger/lifecycle.py
Original file line number Diff line number Diff line change
Expand Up @@ -408,7 +408,6 @@ async def collect_completed_result(
accumulated_content = await chat_storage.session_manager.get_accumulated_content(
subtask_id
)
blocks = await chat_storage.session_manager.finalize_and_get_blocks(subtask_id)
existing_result = await _get_existing_subtask_result(subtask_id)

if result is not None and not isinstance(result, dict):
Expand All @@ -419,6 +418,19 @@ async def collect_completed_result(
)

runtime_result = dict(result) if isinstance(result, dict) else {}
termination_reason = runtime_result.get("termination_reason")
if not isinstance(termination_reason, str) or not termination_reason:
existing_termination_reason = existing_result.get("termination_reason")
termination_reason = (
existing_termination_reason
if isinstance(existing_termination_reason, str)
and existing_termination_reason
else None
)
blocks = await chat_storage.session_manager.finalize_and_get_blocks(
subtask_id,
termination_reason=termination_reason,
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Resolved

has_payload = bool(
runtime_result
Expand Down
16 changes: 10 additions & 6 deletions backend/app/services/chat/webpage_ws_chat_emitter.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,16 +248,20 @@ async def emit_chat_status_updated(
subtask_id: int,
phase: str,
context_metrics: Dict[str, Any],
context_compaction: Dict[str, Any] | None = None,
) -> None:
"""Emit chat:status_updated event to task room."""
payload: Dict[str, Any] = {
"task_id": task_id,
"subtask_id": subtask_id,
"phase": phase,
"context_metrics": context_metrics,
}
if context_compaction is not None:
payload["context_compaction"] = context_compaction
await self.sio.emit(
ServerEvents.CHAT_STATUS_UPDATED,
{
"task_id": task_id,
"subtask_id": subtask_id,
"phase": phase,
"context_metrics": context_metrics,
},
payload,
room=f"task:{task_id}",
namespace=self.namespace,
)
Expand Down
3 changes: 3 additions & 0 deletions backend/app/services/execution/dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,8 +171,10 @@ def extract_completed_result(response_data: dict) -> dict:
),
"loaded_skills": response_data.get("loaded_skills"),
"stop_reason": response_data.get("stop_reason"),
"termination_reason": response_data.get("termination_reason"),
"messages_chain": response_data.get("messages_chain"),
"context_metrics": response_data.get("context_metrics"),
"context_compactions": response_data.get("context_compactions"),
"standalone_chat_workspace_path": response_data.get(
"standalone_chat_workspace_path"
),
Expand Down Expand Up @@ -249,6 +251,7 @@ def parse(
data={
"phase": data.get("phase"),
"context_metrics": data.get("context_metrics") or {},
"context_compaction": data.get("context_compaction"),
},
message_id=message_id,
)
Expand Down
46 changes: 31 additions & 15 deletions backend/app/services/execution/emitters/status_updating.py
Original file line number Diff line number Diff line change
Expand Up @@ -331,11 +331,17 @@ async def emit(self, event: ExecutionEvent) -> None:
await self._cancel_pending_storage_flush_task()

if event.type == EventType.DONE.value:
await self._handle_done(event)
final_result = await self._handle_done(event)
if final_result is not None:
event.result = final_result
elif event.type == EventType.ERROR.value:
await self._handle_error(event)
final_result = await self._handle_error(event)
if final_result is not None:
event.result = final_result
elif event.type == EventType.CANCELLED.value:
await self._handle_cancelled(event)
final_result = await self._handle_cancelled(event)
if final_result is not None:
event.result = final_result

# Forward event to wrapped emitter
await self._wrapped.emit(event)
Expand Down Expand Up @@ -414,43 +420,47 @@ async def close(self) -> None:
await self._cancel_pending_storage_flush_task()
await self._wrapped.close()

async def _handle_done(self, event: ExecutionEvent) -> None:
async def _handle_done(self, event: ExecutionEvent) -> Optional[Dict[str, Any]]:
"""Handle DONE event - update status to COMPLETED.

Args:
event: The DONE event
"""
if self._status_updated:
return
return event.result

await self._update_status_completed(event.result)
return await self._update_status_completed(event.result)

async def _handle_error(self, event: ExecutionEvent) -> None:
async def _handle_error(self, event: ExecutionEvent) -> Optional[Dict[str, Any]]:
"""Handle ERROR event - update status to FAILED.

Args:
event: The ERROR event
"""
if self._status_updated:
return
return event.result

error_message = event.error or "Unknown error"
await self._update_status_failed(error_message, error_code=event.error_code)
return await self._update_status_failed(
error_message, error_code=event.error_code
)

async def _handle_cancelled(self, event: ExecutionEvent) -> None:
async def _handle_cancelled(
self, event: ExecutionEvent
) -> Optional[Dict[str, Any]]:
"""Handle CANCELLED event - update status to CANCELLED.

Args:
event: The CANCELLED event
"""
if self._status_updated:
return
return event.result

await self._update_status_cancelled()
return await self._update_status_cancelled()

async def _update_status_completed(
self, result: Optional[Dict[str, Any]] = None
) -> None:
) -> Optional[Dict[str, Any]]:
"""Update subtask and task status to COMPLETED.

The actual task status (COMPLETED or PENDING_CONFIRMATION for pipeline mode)
Expand Down Expand Up @@ -485,16 +495,18 @@ async def _update_status_completed(
# Publish TaskCompletedEvent for unified handling
# This ensures subscription execution status is updated regardless of execution mode
await self._publish_task_completed_event("COMPLETED", final_result, None)
return final_result

except Exception as e:
logger.error(
f"[StatusUpdatingEmitter] Failed to update status to COMPLETED: {e}",
exc_info=True,
)
return result

async def _update_status_failed(
self, error_message: str, error_code: Optional[str] = None
) -> None:
) -> Optional[Dict[str, Any]]:
"""Update subtask and task status to FAILED.

Also publishes TaskCompletedEvent for unified handling by subscription
Expand Down Expand Up @@ -528,14 +540,16 @@ async def _update_status_failed(

# Publish TaskCompletedEvent for unified handling
await self._publish_task_completed_event("FAILED", result, error_message)
return result

except Exception as e:
logger.error(
f"[StatusUpdatingEmitter] Failed to update status to FAILED: {e}",
exc_info=True,
)
return None

async def _update_status_cancelled(self) -> None:
async def _update_status_cancelled(self) -> Optional[Dict[str, Any]]:
"""Update subtask and task status to CANCELLED.

Also publishes TaskCompletedEvent for unified handling by subscription
Expand All @@ -562,12 +576,14 @@ async def _update_status_cancelled(self) -> None:

# Publish TaskCompletedEvent for unified handling
await self._publish_task_completed_event("CANCELLED", result, None)
return result

except Exception as e:
logger.error(
f"[StatusUpdatingEmitter] Failed to update status for cancelled: {e}",
exc_info=True,
)
return None

async def _publish_task_completed_event(
self,
Expand Down
Loading
Loading