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 langgraph.runtime import Runtime
|
||||||
|
|
||||||
from deerflow.config.app_config import get_app_config
|
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
|
from deerflow.skills.review_worker import maybe_start_review
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -47,6 +48,7 @@ class SkillReviewMiddleware(AgentMiddleware[SkillReviewMiddlewareState]):
|
|||||||
thread_id: str = runtime.context.get("thread_id", "") if runtime.context else ""
|
thread_id: str = runtime.context.get("thread_id", "") if runtime.context else ""
|
||||||
if not thread_id:
|
if not thread_id:
|
||||||
return
|
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:
|
except Exception:
|
||||||
logger.exception("SkillReviewMiddleware.after_agent: unexpected error (ignored)")
|
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.
|
"""Core review logic: create_agent with full middleware chain, no 'silent creation' constraint.
|
||||||
|
|
||||||
Uses create_agent (the deerflow agent factory) directly so we can:
|
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},
|
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:
|
try:
|
||||||
result = await asyncio.wait_for(
|
result = await asyncio.wait_for(
|
||||||
asyncio.get_event_loop().run_in_executor(
|
asyncio.get_event_loop().run_in_executor(None, _invoke_with_user_context),
|
||||||
None,
|
|
||||||
lambda: agent.invoke({"messages": [first_message]}, config=run_config),
|
|
||||||
),
|
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
)
|
)
|
||||||
# Log the final AI message so we can see what the review agent decided.
|
# 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,
|
thread_id: str,
|
||||||
messages: list[BaseMessage],
|
messages: list[BaseMessage],
|
||||||
app_config: "AppConfig",
|
app_config: "AppConfig",
|
||||||
|
user_id: str = "default",
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Run one review; if a newer snapshot arrived while we ran, drain it."""
|
"""Run one review; if a newer snapshot arrived while we ran, drain it."""
|
||||||
lock = _get_lock(thread_id)
|
lock = _get_lock(thread_id)
|
||||||
async with lock:
|
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)
|
pending = _pending.pop(thread_id, None)
|
||||||
if pending is not 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(
|
def maybe_start_review(
|
||||||
thread_id: str,
|
thread_id: str,
|
||||||
messages: list[BaseMessage],
|
messages: list[BaseMessage],
|
||||||
app_config: "AppConfig",
|
app_config: "AppConfig",
|
||||||
|
user_id: str = "default",
|
||||||
) -> bool:
|
) -> bool:
|
||||||
"""Schedule a background review for ``thread_id``.
|
"""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)
|
logger.debug("skill_review_worker: queued review for thread %s (review in progress)", thread_id)
|
||||||
return False
|
return False
|
||||||
|
|
||||||
asyncio.create_task(_run_review_with_pending(thread_id, messages, app_config))
|
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)", thread_id, tool_count)
|
logger.info("skill_review_worker: started review for thread %s (%d tool calls) user=%s", thread_id, tool_count, user_id)
|
||||||
return True
|
return True
|
||||||
Reference in new issue
Block a user