fix(skill-review): propagate user_id through review chain for correct DB ownership
This commit is contained in:
1 parent
857798ff2f
commit
335bc3cf48
2 files changed
+24
-10
No files matched your search
+3
-1
@@ -18,6 +18,7 @@ from langchain.agents.middleware import AgentMiddleware
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from deerflow.config.app_config import get_app_config
|
||||
from deerflow.runtime.user_context import get_effective_user_id
|
||||
from deerflow.skills.review_worker import maybe_start_review
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -47,6 +48,7 @@ class SkillReviewMiddleware(AgentMiddleware[SkillReviewMiddlewareState]):
|
||||
thread_id: str = runtime.context.get("thread_id", "") if runtime.context else ""
|
||||
if not thread_id:
|
||||
return
|
||||
maybe_start_review(thread_id, messages, app_config)
|
||||
user_id = get_effective_user_id()
|
||||
maybe_start_review(thread_id, messages, app_config, user_id)
|
||||
except Exception:
|
||||
logger.exception("SkillReviewMiddleware.after_agent: unexpected error (ignored)")
|
||||
@@ -122,7 +122,7 @@ _REVIEW_SYSTEM_PROMPT = textwrap.dedent("""\
|
||||
""")
|
||||
|
||||
|
||||
async def _run_review(thread_id: str, messages: list[BaseMessage], app_config: "AppConfig") -> None:
|
||||
async def _run_review(thread_id: str, messages: list[BaseMessage], app_config: "AppConfig", user_id: str = "default") -> None:
|
||||
"""Core review logic: create_agent with full middleware chain, no 'silent creation' constraint.
|
||||
|
||||
Uses create_agent (the deerflow agent factory) directly so we can:
|
||||
@@ -196,12 +196,22 @@ async def _run_review(thread_id: str, messages: list[BaseMessage], app_config: "
|
||||
additional_kwargs={"prompt_prefix": _REVIEW_SYSTEM_PROMPT},
|
||||
)
|
||||
|
||||
def _invoke_with_user_context() -> dict:
|
||||
from deerflow.runtime.user_context import set_current_user, reset_current_user
|
||||
|
||||
class _SimpleUser:
|
||||
def __init__(self, uid: str) -> None:
|
||||
self.id = uid
|
||||
|
||||
token = set_current_user(_SimpleUser(user_id))
|
||||
try:
|
||||
return agent.invoke({"messages": [first_message]}, config=run_config)
|
||||
finally:
|
||||
reset_current_user(token)
|
||||
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.get_event_loop().run_in_executor(
|
||||
None,
|
||||
lambda: agent.invoke({"messages": [first_message]}, config=run_config),
|
||||
),
|
||||
asyncio.get_event_loop().run_in_executor(None, _invoke_with_user_context),
|
||||
timeout=timeout,
|
||||
)
|
||||
# Log the final AI message so we can see what the review agent decided.
|
||||
@@ -227,21 +237,23 @@ async def _run_review_with_pending(
|
||||
thread_id: str,
|
||||
messages: list[BaseMessage],
|
||||
app_config: "AppConfig",
|
||||
user_id: str = "default",
|
||||
) -> None:
|
||||
"""Run one review; if a newer snapshot arrived while we ran, drain it."""
|
||||
lock = _get_lock(thread_id)
|
||||
async with lock:
|
||||
await _run_review(thread_id, messages, app_config)
|
||||
await _run_review(thread_id, messages, app_config, user_id)
|
||||
|
||||
pending = _pending.pop(thread_id, None)
|
||||
if pending is not None:
|
||||
asyncio.create_task(_run_review_with_pending(thread_id, pending, app_config))
|
||||
asyncio.create_task(_run_review_with_pending(thread_id, pending, app_config, user_id))
|
||||
|
||||
|
||||
def maybe_start_review(
|
||||
thread_id: str,
|
||||
messages: list[BaseMessage],
|
||||
app_config: "AppConfig",
|
||||
user_id: str = "default",
|
||||
) -> bool:
|
||||
"""Schedule a background review for ``thread_id``.
|
||||
|
||||
@@ -267,6 +279,6 @@ def maybe_start_review(
|
||||
logger.debug("skill_review_worker: queued review for thread %s (review in progress)", thread_id)
|
||||
return False
|
||||
|
||||
asyncio.create_task(_run_review_with_pending(thread_id, messages, app_config))
|
||||
logger.info("skill_review_worker: started review for thread %s (%d tool calls)", thread_id, tool_count)
|
||||
asyncio.create_task(_run_review_with_pending(thread_id, messages, app_config, user_id))
|
||||
logger.info("skill_review_worker: started review for thread %s (%d tool calls) user=%s", thread_id, tool_count, user_id)
|
||||
return True
|
||||
Reference in new issue
Block a user