Skip to content
Merged
Show file tree
Hide file tree
Changes from all 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
51 changes: 49 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 @@ -496,6 +523,7 @@ async def on_task_join(self, sid: str, data: dict) -> dict:
"started_at": streaming_info.get("started_at"),
"last_activity_at": streaming_info.get("last_activity_at"),
},
"status_updated": cached_context_metrics,
"subtasks": subtasks_dict,
}

Expand Down Expand Up @@ -989,12 +1017,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: 3 additions & 0 deletions backend/app/api/ws/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -762,6 +762,9 @@ class TaskJoinAck(BaseModel):
subtasks: Optional[List[Dict[str, Any]]] = Field(
None, description="Subtasks data for immediate message sync"
)
status_updated: Optional[Dict[str, Any]] = Field(
None, description="Latest cached chat:status_updated snapshot"
)
error: Optional[str] = None


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
18 changes: 17 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,23 @@ 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
)
if termination_reason:
runtime_result["termination_reason"] = termination_reason
else:
runtime_result.pop("termination_reason", 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.

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
Loading
Loading