caoxiaozhu/JARVIS
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: harden streaming chat persistence and access control

Browse Source 
Persist streaming chat state during generator cleanup, close the SSE inner stream safely, and reject cross-user conversation access while locking the behavior with focused regressions.
This commit is contained in:
WIN-JHFT4D3SIVT\caoxiaozhu
2026-04-02 21:49:53 +08:00
parent 4251a79062
commit b3f9b5e715
3 changed files with 1232 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
backend/app/routers/conversation.py
View File

@@ -130,34 +130,42 @@ async def chat_stream(
agent_svc = AgentService(db)
async def stream_generator():
stream = None
msg_id = None
should_emit_done = False
try:
conv_id, msg_id, stream = await agent_svc.chat(
user_id=current_user.id,
message=data.message,
conversation_id=data.conversation_id,
file_ids=data.file_ids,
model_name=data.model_name,
)
except ValueError as exc:
yield f"event: error\ndata: {json.dumps({'error': str(exc)}, ensure_ascii=False)}\n\n"
return
try:
conv_id, msg_id, stream = await agent_svc.chat(
user_id=current_user.id,
message=data.message,
conversation_id=data.conversation_id,
file_ids=data.file_ids,
model_name=data.model_name,
)
except ValueError as exc:
yield f"event: error\ndata: {json.dumps({'error': str(exc)}, ensure_ascii=False)}\n\n"
return
yield f"event: metadata\ndata: {json.dumps({'conversation_id': conv_id, 'message_id': msg_id})}\n\n"
yield f"event: metadata\ndata: {json.dumps({'conversation_id': conv_id, 'message_id': msg_id})}\n\n"
try:
async for event in stream:
event_type = event.get('type', 'progress')
if event_type == 'chunk':
yield f"event: chunk\ndata: {json.dumps({'content': event.get('content', '')}, ensure_ascii=False)}\n\n"
elif event_type == 'error':
yield f"event: error\ndata: {json.dumps({'error': event.get('error', '未知错误')}, ensure_ascii=False)}\n\n"
else:
payload = {k: v for k, v in event.items() if k != 'type'}
yield f"event: progress\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n"
except Exception as e:
yield f"event: error\ndata: {json.dumps({'error': str(e)}, ensure_ascii=False)}\n\n"
try:
async for event in stream:
event_type = event.get('type', 'progress')
if event_type == 'chunk':
yield f"event: chunk\ndata: {json.dumps({'content': event.get('content', '')}, ensure_ascii=False)}\n\n"
elif event_type == 'error':
yield f"event: error\ndata: {json.dumps({'error': event.get('error', '未知错误')}, ensure_ascii=False)}\n\n"
else:
payload = {k: v for k, v in event.items() if k != 'type'}
yield f"event: progress\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n"
except Exception as e:
yield f"event: error\ndata: {json.dumps({'error': str(e)}, ensure_ascii=False)}\n\n"
should_emit_done = msg_id is not None
if should_emit_done:
yield f"event: done\ndata: {json.dumps({'message_id': msg_id})}\n\n"
finally:
yield f"event: done\ndata: {json.dumps({'message_id': msg_id})}\n\n"
if stream is not None:
await stream.aclose()
return StreamingResponse(
stream_generator(),

335
backend/app/services/agent_service.py
View File

@@ -53,6 +53,73 @@ def _is_streaming_rejection_error(error: Exception, user_llm_config: dict | None
return any(marker in error_text for marker in markers)
def _coerce_event_text(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
parts: list[str] = []
for item in content:
if isinstance(item, str):
parts.append(item)
elif isinstance(item, dict):
text = item.get("text")
if isinstance(text, str):
parts.append(text)
return "".join(parts)
return str(content) if content else ""
_CONTINUITY_STATE_VERSION = 1
_CONTINUITY_SNAPSHOT_FIELDS = (
"turn_context",
"routing_decision",
"continuity_state",
"pending_action",
"last_completed_action",
"clarification_context",
"tool_outcomes",
"pending_tasks",
"completed_tasks",
"created_entities",
"current_agent",
"next_step",
"agent_trace",
)
def _build_continuity_snapshot(state: dict[str, Any]) -> dict[str, Any] | None:
snapshot = {
field: state.get(field)
for field in _CONTINUITY_SNAPSHOT_FIELDS
if state.get(field) is not None
}
if not snapshot:
return None
return {
"version": _CONTINUITY_STATE_VERSION,
"state": snapshot,
}
def _extract_continuity_snapshot(payload: Any) -> dict[str, Any] | None:
if isinstance(payload, list):
for item in payload:
snapshot = _extract_continuity_snapshot(item)
if snapshot:
return snapshot
return None
if not isinstance(payload, dict):
return None
if payload.get("kind") != "agent_continuity_state":
return None
if payload.get("version") != _CONTINUITY_STATE_VERSION:
return None
state = payload.get("state")
if isinstance(state, dict):
return state
return None
class AgentService:
"""对话 Agent 服务"""
@@ -83,10 +150,23 @@ class AgentService:
"steps": steps or [],
}
def _build_current_datetime_context(self) -> tuple[str, dict[str, str]]:
now_utc = datetime.now(UTC)
reference = {
"current_time_iso": now_utc.isoformat(),
"current_date_iso": now_utc.date().isoformat(),
}
context = (
"【当前时间】\n"
f"- current_time_utc: {reference['current_time_iso']}\n"
f"- current_date_utc: {reference['current_date_iso']}\n"
"说明：解析‘今天/明天/后天/本周/下周’等相对时间时，请以 current_time_utc 为准。"
)
return context, reference
async def _get_user_llm_config(self, user_id: str, model_name: str | None = None) -> dict | None:
"""获取用户的 LLM 模型配置"""
result = await self.db.execute(select(User).where(User.id == user_id))
user = result.scalar_one_or_none()
user = await self.db.get(User, user_id)
if not user or not user.llm_config:
return None
@@ -106,6 +186,47 @@ class AgentService:
return None
async def _load_continuity_snapshot(self, conversation: Conversation) -> dict[str, Any] | None:
snapshot = _extract_continuity_snapshot(conversation.agent_state)
if snapshot:
return snapshot
result = await self.db.execute(
select(Message)
.where(Message.conversation_id == conversation.id, Message.role == "assistant")
.order_by(Message.created_at.desc())
)
for message in result.scalars():
snapshot = _extract_continuity_snapshot(message.attachments)
if snapshot:
return snapshot
return None
async def _build_agent_state(
self,
*,
user_id: str,
conversation: Conversation,
full_message: str,
memory_context: str | None,
current_datetime_context: str,
current_datetime_reference: dict[str, str],
user_llm_config: dict | None,
) -> dict[str, Any]:
state = initial_state(user_id, conversation.id)
state.update({
"messages": [HumanMessage(content=full_message)],
"memory_context": memory_context,
"current_datetime_context": current_datetime_context,
"current_datetime_reference": current_datetime_reference,
"user_llm_config": user_llm_config,
})
previous_snapshot = await self._load_continuity_snapshot(conversation)
if previous_snapshot:
state.update(previous_snapshot)
state["messages"] = [HumanMessage(content=full_message)]
return state
async def chat(
self,
user_id: str,
@@ -138,9 +259,14 @@ class AgentService:
if conversation_id:
result = await self.db.execute(
select(Conversation).where(Conversation.id == conversation_id)
select(Conversation).where(
Conversation.id == conversation_id,
Conversation.user_id == user_id,
)
)
conv = result.scalar_one_or_none()
if conv is None:
raise ValueError("会话不存在或无权访问")
else:
conv = None
@@ -203,33 +329,38 @@ class AgentService:
await self.db.commit()
await self.db.refresh(assistant_msg)
def _build_current_datetime_context() -> str:
now_utc = datetime.now(UTC)
return (
"【当前时间】\n"
f"- current_time_utc: {now_utc.isoformat()}\n"
f"- current_date_utc: {now_utc.date().isoformat()}\n"
"说明：解析‘今天/明天/后天/本周/下周’等相对时间时，请以 current_time_utc 为准。"
)
def _build_assistant_event_payload(content: str) -> dict[str, Any]:
return {
"source_type": "conversation",
"source_id": conversation_id,
"event_type": "message_created",
"title": "Assistant message",
"content_summary": content[:500],
"raw_excerpt": content[:2000],
"metadata_": {"role": "assistant"},
"importance_signal": 0.8,
}
async def run_agent():
collected = ""
state: dict[str, Any] | None = None
set_current_user(user_id)
try:
graph = get_agent_graph()
current_datetime_context = _build_current_datetime_context()
# 使用 initial_state 构建状态
state = initial_state(user_id, conversation_id)
state.update({
"messages": [HumanMessage(content=full_message)],
"memory_context": memory_ctx,
"current_datetime_context": current_datetime_context,
"user_llm_config": user_llm_config,
})
current_datetime_context, current_datetime_reference = self._build_current_datetime_context()
state = await self._build_agent_state(
user_id=user_id,
conversation=conv,
full_message=full_message,
memory_context=memory_ctx,
current_datetime_context=current_datetime_context,
current_datetime_reference=current_datetime_reference,
user_llm_config=user_llm_config,
)
yield self._build_progress_event("thinking", "Jarvis 正在分析请求", agent="master", step="理解你的问题")
collected = ""
try:
async for event in graph.astream_events(state, version="v2"):
kind = event.get("event")
@@ -272,49 +403,77 @@ class AgentService:
elif kind == "on_chat_model_stream":
chunk = data.get("chunk")
content = getattr(chunk, "content", "") if chunk else ""
content = _coerce_event_text(getattr(chunk, "content", "") if chunk else "")
if content:
collected += content
yield {"type": "chunk", "content": content}
elif kind == "on_chain_end" and event_name == "create_agent_graph":
# 最终输出通常在这里
elif kind == "on_chain_end":
output = data.get("output")
if isinstance(output, dict) and "final_response" in output:
final_resp = output["final_response"]
# 如果还没流式输出完整，补全它
if final_resp and not collected:
collected = final_resp
yield {"type": "chunk", "content": collected}
final_resp = None
if isinstance(output, dict):
state.update(output)
final_resp = output.get("final_response")
if final_resp:
final_text = str(final_resp)
if final_text != collected:
collected = final_text
yield {"type": "chunk", "content": final_text}
elif kind == "on_chat_model_end":
output = data.get("output")
final_content = _coerce_event_text(getattr(output, "content", "") if output else "")
if final_content:
final_text = final_content
if final_text != collected:
collected = final_text
yield {"type": "chunk", "content": final_text}
except Exception as e:
if _is_streaming_rejection_error(e, user_llm_config) and not collected:
yield self._build_progress_event("responding", "Jarvis 正在生成回复", agent="master", step="fallback")
try:
result_state = await graph.ainvoke(state)
if isinstance(result_state, dict):
state.update(result_state)
fallback_content = result_state.get("final_response") or str(result_state.get("messages", [AIMessage(content="")])[-1].content)
collected = str(fallback_content)
yield {"type": "chunk", "content": collected}
except Exception as fallback_error:
except Exception:
logger.exception("llm_sync_fallback_failed")
yield {"type": "error", "error": "模型服务暂不可用。"}
safe_error = "模型服务暂不可用，请稍后再试。"
yield {"type": "error", "error": safe_error}
collected = f"抱歉，发生错误: {safe_error}"
yield {"type": "chunk", "content": collected}
else:
logger.exception("agent_streaming_failed")
yield {"type": "error", "error": str(e)}
if not collected:
safe_error = "模型服务暂不可用，请稍后再试。"
yield {"type": "error", "error": safe_error}
collected = f"抱歉，发生错误: {safe_error}"
yield {"type": "chunk", "content": collected}
else:
yield {"type": "error", "error": str(e)}
finally:
clear_current_user()
asyncio.create_task(self._try_auto_summarize_background(user_id, conversation_id))
if collected:
try:
async with async_session() as session:
result2 = await session.execute(select(Message).where(Message.id == assistant_msg.id))
msg = result2.scalar_one_or_none()
if msg:
msg.content = collected
await session.commit()
if collected:
assistant_msg.content = collected
continuity_snapshot = _build_continuity_snapshot(state or {})
assistant_msg.attachments = ([{
"kind": "agent_continuity_state",
**continuity_snapshot,
}] if continuity_snapshot else None)
conv.agent_state = continuity_snapshot
await BrainService(self.db).create_event(
user_id,
**_build_assistant_event_payload(collected),
)
await self.db.commit()
await self.db.refresh(assistant_msg)
except Exception:
logger.exception("save_assistant_message_failed")
asyncio.create_task(self._try_auto_summarize_background(user_id, conversation_id))
return conversation_id, assistant_msg.id, run_agent()
@@ -331,32 +490,74 @@ class AgentService:
"""
user_llm_config = await self._get_user_llm_config(user_id, model_name)
model_name_used = model_name
if model_name and not user_llm_config:
raise ValueError("所选模型不可用于聊天，请切换到聊天模型")
if user_llm_config:
model_name_used = user_llm_config.get("name", model_name)
if not conversation_id:
if conversation_id:
result = await self.db.execute(
select(Conversation).where(
Conversation.id == conversation_id,
Conversation.user_id == user_id,
)
)
conv = result.scalar_one_or_none()
if conv is None:
raise ValueError("会话不存在或无权访问")
else:
conv = None
if not conv:
conv = Conversation(user_id=user_id, title=message[:50])
self.db.add(conv)
await self.db.commit()
await self.db.refresh(conv)
conversation_id = conv.id
else:
conversation_id = conv.id
user_msg = Message(conversation_id=conversation_id, role="user", content=message)
self.db.add(user_msg)
assistant_msg = Message(
conversation_id=conversation_id,
role="assistant",
content="",
model=model_name_used or "jarvis",
attachments=None,
)
self.db.add(assistant_msg)
brain_service = BrainService(self.db)
await brain_service.create_event(
user_id,
source_type="conversation",
source_id=conversation_id,
event_type="message_created",
title="User message",
content_summary=message[:500],
raw_excerpt=message[:2000],
metadata_={"role": "user"},
importance_signal=1.0,
)
memory_ctx = await memory_service.build_memory_context(self.db, user_id, conversation_id, message)
set_current_user(user_id)
try:
graph = get_agent_graph()
state = initial_state(user_id, conversation_id)
state.update({
"messages": [HumanMessage(content=message)],
"memory_context": memory_ctx,
"current_datetime_context": datetime.now(UTC).isoformat(),
"user_llm_config": user_llm_config,
})
current_datetime_context, current_datetime_reference = self._build_current_datetime_context()
state = await self._build_agent_state(
user_id=user_id,
conversation=conv,
full_message=message,
memory_context=memory_ctx,
current_datetime_context=current_datetime_context,
current_datetime_reference=current_datetime_reference,
user_llm_config=user_llm_config,
)
result_state = await graph.ainvoke(state)
response_content = result_state.get("final_response") or str(result_state.get("messages", [AIMessage(content="")])[-1].content)
except Exception as e:
@@ -365,13 +566,27 @@ class AgentService:
finally:
clear_current_user()
assistant_msg = Message(
conversation_id=conversation_id,
role="assistant",
content=response_content,
model=model_name_used or "jarvis",
brain_service = BrainService(self.db)
await brain_service.create_event(
user_id,
source_type="conversation",
source_id=conversation_id,
event_type="message_created",
title="Assistant message",
content_summary=response_content[:500],
raw_excerpt=response_content[:2000],
metadata_={"role": "assistant"},
importance_signal=0.8,
)
self.db.add(assistant_msg)
assistant_msg.content = response_content
continuity_snapshot = _build_continuity_snapshot(result_state) if 'result_state' in locals() else None
assistant_msg.attachments = ([{
"kind": "agent_continuity_state",
**continuity_snapshot,
}] if continuity_snapshot else None)
conv.agent_state = continuity_snapshot
await self.db.commit()
await self.db.refresh(assistant_msg)
return conversation_id, assistant_msg.id, response_content, model_name_used

927
backend/tests/backend/app/services/test_brain_ingestion.py
View File

File diff suppressed because it is too large Load Diff