fix: isolate user concurrency from internal agent runs
This commit is contained in:
1 parent
124eb04a6b
commit
c23350f17b
8 files changed
+162
-10
No files matched your search
@@ -32,6 +32,7 @@ from app.gateway.roundtable_model_fallback import (
|
||||
)
|
||||
from app.gateway.roundtable_run_policy import roundtable_run_policy
|
||||
from app.gateway.roundtable_seat_skills import append_seat_skill_directive
|
||||
from app.gateway.services import RUN_SCOPE_BACKGROUND_CHILD
|
||||
from app.gateway.routers._roundtable_seed import (
|
||||
COORDINATOR_AGENT_ID,
|
||||
DASHBOARD_AGENT_ID,
|
||||
@@ -346,6 +347,7 @@ async def _run_agent_turn(
|
||||
body = RunCreateRequest(
|
||||
assistant_id="lead_agent",
|
||||
input={"messages": [{"type": "human", "content": [{"type": "text", "text": message}], "additional_kwargs": {}}]},
|
||||
metadata={"run_scope": RUN_SCOPE_BACKGROUND_CHILD},
|
||||
config={"recursion_limit": 1000},
|
||||
context={
|
||||
"agent_name": agent_name,
|
||||
@@ -374,7 +376,12 @@ async def _run_agent_turn(
|
||||
# 落 checkpoint)。额度满时本轮在此 await 排队,等同用户某个在跑轮次释放额度后再开。
|
||||
async with _user_background_semaphore(user_id):
|
||||
try:
|
||||
record = await start_run(body, thread_id, request) # type: ignore[arg-type]
|
||||
record = await start_run( # type: ignore[arg-type]
|
||||
body,
|
||||
thread_id,
|
||||
request,
|
||||
run_scope=RUN_SCOPE_BACKGROUND_CHILD,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — start_run 报错(409 抢占 / DB 抖动等)也要落诊断
|
||||
if diag is not None:
|
||||
await diag.record(
|
||||
|
||||
@@ -65,6 +65,7 @@ from langgraph.checkpoint.base import empty_checkpoint
|
||||
from app.gateway.roundtable_diag import record_foreground_diag
|
||||
from app.gateway.roundtable_model_fallback import fallback_model_chain, is_llm_error_message, reason_text
|
||||
from app.gateway.roundtable_run_policy import merge_skill_stop_names, roundtable_run_policy
|
||||
from app.gateway.services import RUN_SCOPE_ORCHESTRATION_CHILD, internal_run_scope_headers
|
||||
from deerflow.agents.roundtable_orchestrator.extraction_spec import build_summary_directive
|
||||
|
||||
from app.gateway.roundtable_seat_skills import append_seat_skill_directive
|
||||
@@ -325,7 +326,11 @@ def _auth_headers(request: Request) -> dict[str, str]:
|
||||
token = request.headers.get("authorization")
|
||||
if not token:
|
||||
raise HTTPException(status_code=401, detail="Missing Authorization header")
|
||||
return {"Authorization": token, "Content-Type": "application/json"}
|
||||
return {
|
||||
"Authorization": token,
|
||||
"Content-Type": "application/json",
|
||||
**internal_run_scope_headers(),
|
||||
}
|
||||
|
||||
|
||||
async def _create_thread(
|
||||
@@ -694,6 +699,7 @@ def _run_payload(
|
||||
}
|
||||
]
|
||||
},
|
||||
"metadata": {"run_scope": RUN_SCOPE_ORCHESTRATION_CHILD},
|
||||
"config": {"recursion_limit": 1000},
|
||||
"context": {
|
||||
"agent_name": agent_name,
|
||||
|
||||
@@ -10,6 +10,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
@@ -41,12 +42,68 @@ from deerflow.runtime.user_context import get_effective_user_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RUN_SCOPE_EXTERNAL_ROOT = "external_root"
|
||||
RUN_SCOPE_ORCHESTRATION_CHILD = "orchestration_child"
|
||||
RUN_SCOPE_BACKGROUND_CHILD = "background_child"
|
||||
_INTERNAL_RUN_SCOPES = {RUN_SCOPE_ORCHESTRATION_CHILD, RUN_SCOPE_BACKGROUND_CHILD}
|
||||
_ALL_RUN_SCOPES = {RUN_SCOPE_EXTERNAL_ROOT, *_INTERNAL_RUN_SCOPES}
|
||||
_INTERNAL_RUN_SCOPE_HEADER = "x-deerflow-internal-run-scope-token"
|
||||
_INTERNAL_RUN_SCOPE_TOKEN = (
|
||||
os.getenv("DEERFLOW_INTERNAL_RUN_SCOPE_TOKEN", "").strip()
|
||||
or uuid.uuid4().hex
|
||||
)
|
||||
|
||||
|
||||
def internal_run_scope_headers() -> dict[str, str]:
|
||||
"""Headers used by same-process loopback calls to mark internal child runs."""
|
||||
|
||||
return {_INTERNAL_RUN_SCOPE_HEADER: _INTERNAL_RUN_SCOPE_TOKEN}
|
||||
|
||||
|
||||
def _metadata_run_scope(body: Any) -> str | None:
|
||||
metadata = getattr(body, "metadata", None) or {}
|
||||
if not isinstance(metadata, dict):
|
||||
return None
|
||||
raw = str(metadata.get("run_scope") or "").strip()
|
||||
return raw or None
|
||||
|
||||
|
||||
def _request_is_loopback(request: Request) -> bool:
|
||||
client = getattr(request, "client", None)
|
||||
host = str(getattr(client, "host", "") or "").strip()
|
||||
return host in {"127.0.0.1", "::1", "localhost"}
|
||||
|
||||
|
||||
def _has_internal_run_scope_header(request: Request) -> bool:
|
||||
headers = getattr(request, "headers", None)
|
||||
token = headers.get(_INTERNAL_RUN_SCOPE_HEADER) if headers is not None else None
|
||||
return str(token or "") == _INTERNAL_RUN_SCOPE_TOKEN
|
||||
|
||||
|
||||
def _resolve_run_scope(body: Any, request: Request, run_scope: str | None) -> str:
|
||||
"""Return the trusted run scope used for concurrency accounting."""
|
||||
|
||||
if run_scope:
|
||||
scope = str(run_scope).strip()
|
||||
return scope if scope in _ALL_RUN_SCOPES else RUN_SCOPE_EXTERNAL_ROOT
|
||||
|
||||
scope = _metadata_run_scope(body)
|
||||
if (
|
||||
scope in _INTERNAL_RUN_SCOPES
|
||||
and _request_is_loopback(request)
|
||||
and _has_internal_run_scope_header(request)
|
||||
):
|
||||
return scope
|
||||
if scope == RUN_SCOPE_EXTERNAL_ROOT:
|
||||
return RUN_SCOPE_EXTERNAL_ROOT
|
||||
return RUN_SCOPE_EXTERNAL_ROOT
|
||||
|
||||
|
||||
async def _run_agent_with_concurrency_lease(
|
||||
*,
|
||||
concurrency_gate: Any | None,
|
||||
lease_acquired: bool,
|
||||
lease_metadata: dict[str, str],
|
||||
lease_metadata: dict[str, Any],
|
||||
heartbeat_interval_seconds: int,
|
||||
run_agent_kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
@@ -750,6 +807,7 @@ async def start_run(
|
||||
request: Request,
|
||||
*,
|
||||
enforce_access: bool = True,
|
||||
run_scope: str | None = None,
|
||||
) -> RunRecord:
|
||||
"""Create a RunRecord and launch the background agent task.
|
||||
|
||||
@@ -793,21 +851,29 @@ async def start_run(
|
||||
|
||||
effective_model_name = _resolve_effective_run_model_name(body)
|
||||
effective_user_id = get_effective_user_id()
|
||||
effective_run_scope = _resolve_run_scope(body, request, run_scope)
|
||||
model_concurrency_limit = _get_model_concurrency_limit(effective_model_name)
|
||||
user_model_concurrency_limit = _get_user_model_concurrency_limit(effective_model_name)
|
||||
user_model_concurrency_limit = (
|
||||
_get_user_model_concurrency_limit(effective_model_name)
|
||||
if effective_run_scope == RUN_SCOPE_EXTERNAL_ROOT
|
||||
else None
|
||||
)
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
run_metadata = dict(body.metadata or {})
|
||||
run_metadata["model_name"] = effective_model_name
|
||||
run_metadata["user_id"] = effective_user_id
|
||||
run_metadata["run_scope"] = effective_run_scope
|
||||
|
||||
redis_lease_acquired = False
|
||||
redis_gate_degraded = True
|
||||
user_model_accounted = user_model_concurrency_limit is not None
|
||||
lease_metadata = {
|
||||
"run_id": run_id,
|
||||
"thread_id": thread_id,
|
||||
"model_name": effective_model_name,
|
||||
"user_id": effective_user_id,
|
||||
"user_model_accounted": user_model_accounted,
|
||||
}
|
||||
if concurrency_gate is not None:
|
||||
acquire = await concurrency_gate.try_acquire(
|
||||
|
||||
@@ -18,7 +18,10 @@ concurrency:
|
||||
enabled: false
|
||||
url: redis://127.0.0.1:6379/0
|
||||
password: null
|
||||
namespace: deerflow:concurrency
|
||||
# Prefix with the deployment/system id because Redis may be shared by
|
||||
# multiple projects. Keep this unique per system to avoid cross-project
|
||||
# leases affecting each other.
|
||||
namespace: cmzs_deerflow:concurrency
|
||||
socket_connect_timeout_ms: 80
|
||||
socket_timeout_ms: 80
|
||||
acquire_timeout_ms: 120
|
||||
|
||||
+4
-1
@@ -14,7 +14,10 @@ class RedisConcurrencyConfig(BaseModel):
|
||||
default=None,
|
||||
description="Optional Redis password. You may also include credentials directly in url.",
|
||||
)
|
||||
namespace: str = Field(default="deerflow:concurrency", description="Redis key prefix for concurrency leases.")
|
||||
namespace: str = Field(
|
||||
default="cmzs_deerflow:concurrency",
|
||||
description="Redis key prefix for concurrency leases. Use a deployment/system prefix when Redis is shared.",
|
||||
)
|
||||
socket_connect_timeout_ms: int = Field(default=80, ge=1, le=10_000)
|
||||
socket_timeout_ms: int = Field(default=80, ge=1, le=10_000)
|
||||
acquire_timeout_ms: int = Field(default=120, ge=1, le=10_000)
|
||||
|
||||
+14
-4
@@ -110,7 +110,9 @@ end
|
||||
|
||||
redis.call('ZADD', KEYS[1], expires_at, run_id)
|
||||
redis.call('ZADD', KEYS[2], expires_at, run_id)
|
||||
redis.call('ZADD', KEYS[3], expires_at, run_id)
|
||||
if user_model_limit > 0 then
|
||||
redis.call('ZADD', KEYS[3], expires_at, run_id)
|
||||
end
|
||||
redis.call('ZADD', KEYS[4], expires_at, run_id)
|
||||
redis.call('SET', KEYS[5], payload, 'PX', ttl_ms)
|
||||
return {'acquired', ''}
|
||||
@@ -146,7 +148,7 @@ return {'acquired', ''}
|
||||
logger.debug("Failed to close Redis concurrency client", exc_info=True)
|
||||
|
||||
def _ns(self) -> str:
|
||||
return self.config.namespace.strip().rstrip(":") or "deerflow:concurrency"
|
||||
return self.config.namespace.strip().rstrip(":") or "cmzs_deerflow:concurrency"
|
||||
|
||||
def _keys(self, *, run_id: str, model_name: str, user_id: str, thread_id: str) -> list[str]:
|
||||
ns = self._ns()
|
||||
@@ -263,6 +265,7 @@ return {'acquired', ''}
|
||||
thread_id: str,
|
||||
model_name: str,
|
||||
user_id: str,
|
||||
user_model_accounted: bool = True,
|
||||
) -> None:
|
||||
if not self.enabled or not self._breaker.allow():
|
||||
return
|
||||
@@ -272,7 +275,10 @@ return {'acquired', ''}
|
||||
try:
|
||||
client = await self._redis()
|
||||
pipe = client.pipeline(transaction=False)
|
||||
for key in keys[:4]:
|
||||
refresh_keys = [keys[0], keys[1], keys[3]]
|
||||
if user_model_accounted:
|
||||
refresh_keys.append(keys[2])
|
||||
for key in refresh_keys:
|
||||
pipe.zadd(key, {run_id: expires_at})
|
||||
pipe.expire(keys[4], self.config.lease_ttl_seconds)
|
||||
await self._with_timeout(pipe.execute(), self.config.operation_timeout_ms)
|
||||
@@ -287,6 +293,7 @@ return {'acquired', ''}
|
||||
thread_id: str,
|
||||
model_name: str,
|
||||
user_id: str,
|
||||
user_model_accounted: bool = True,
|
||||
) -> None:
|
||||
if not self.enabled or not self._breaker.allow():
|
||||
return
|
||||
@@ -294,7 +301,10 @@ return {'acquired', ''}
|
||||
try:
|
||||
client = await self._redis()
|
||||
pipe = client.pipeline(transaction=False)
|
||||
for key in keys[:4]:
|
||||
release_keys = [keys[0], keys[1], keys[3]]
|
||||
if user_model_accounted:
|
||||
release_keys.append(keys[2])
|
||||
for key in release_keys:
|
||||
pipe.zrem(key, run_id)
|
||||
pipe.delete(keys[4])
|
||||
await self._with_timeout(pipe.execute(), self.config.operation_timeout_ms)
|
||||
|
||||
@@ -273,6 +273,7 @@ class RunManager:
|
||||
if r.status in (RunStatus.pending, RunStatus.running)
|
||||
and str((r.metadata or {}).get("model_name") or "").strip() == model_name
|
||||
and str((r.metadata or {}).get("user_id") or "").strip() == user_id
|
||||
and str((r.metadata or {}).get("run_scope") or "external_root").strip() == "external_root"
|
||||
)
|
||||
if active_for_user_model >= max(1, int(limit)):
|
||||
raise ModelConcurrencyLimitError(model_name, limit, scope="user_model", user_id=user_id)
|
||||
|
||||
@@ -10,11 +10,17 @@ import pytest
|
||||
|
||||
from app.gateway.routers import system_settings as ss_router
|
||||
from app.gateway.services import (
|
||||
RUN_SCOPE_EXTERNAL_ROOT,
|
||||
RUN_SCOPE_ORCHESTRATION_CHILD,
|
||||
_get_model_concurrency_limit,
|
||||
_get_user_model_concurrency_limit,
|
||||
_resolve_run_scope,
|
||||
_resolve_effective_run_model_name,
|
||||
internal_run_scope_headers,
|
||||
)
|
||||
from deerflow.config import system_settings as ss
|
||||
from deerflow.config.concurrency_config import RedisConcurrencyConfig
|
||||
from deerflow.runtime.runs.concurrency_gate import RedisConcurrencyGate
|
||||
from deerflow.runtime.runs.manager import ModelConcurrencyLimitError, RunManager
|
||||
|
||||
|
||||
@@ -212,6 +218,15 @@ async def test_run_manager_rejects_user_model_limit_atomically(
|
||||
assert _get_user_model_concurrency_limit("other") == 3
|
||||
|
||||
run_mgr = RunManager()
|
||||
await run_mgr.create_or_reject(
|
||||
"t-child",
|
||||
metadata={
|
||||
"model_name": "glm5",
|
||||
"user_id": "u1",
|
||||
"run_scope": RUN_SCOPE_ORCHESTRATION_CHILD,
|
||||
},
|
||||
user_model_concurrency_limit=("glm5", "u1", 1),
|
||||
)
|
||||
await run_mgr.create_or_reject("t1", metadata={"model_name": "glm5", "user_id": "u1"})
|
||||
await run_mgr.create_or_reject(
|
||||
"t2",
|
||||
@@ -247,3 +262,44 @@ def test_get_model_concurrency_limit_skips_when_disabled(
|
||||
)
|
||||
assert _get_model_concurrency_limit("glm5") is None
|
||||
assert _get_user_model_concurrency_limit("glm5") is None
|
||||
|
||||
|
||||
def test_internal_run_scope_is_only_trusted_for_loopback_or_server_argument() -> None:
|
||||
body = SimpleNamespace(metadata={"run_scope": RUN_SCOPE_ORCHESTRATION_CHILD})
|
||||
|
||||
trusted_loopback_request = SimpleNamespace(
|
||||
client=SimpleNamespace(host="127.0.0.1"),
|
||||
headers=internal_run_scope_headers(),
|
||||
)
|
||||
untrusted_loopback_request = SimpleNamespace(
|
||||
client=SimpleNamespace(host="127.0.0.1"),
|
||||
headers={},
|
||||
)
|
||||
public_request = SimpleNamespace(
|
||||
client=SimpleNamespace(host="10.0.0.8"),
|
||||
headers=internal_run_scope_headers(),
|
||||
)
|
||||
|
||||
assert (
|
||||
_resolve_run_scope(body, trusted_loopback_request, None)
|
||||
== RUN_SCOPE_ORCHESTRATION_CHILD
|
||||
)
|
||||
assert (
|
||||
_resolve_run_scope(body, untrusted_loopback_request, None)
|
||||
== RUN_SCOPE_EXTERNAL_ROOT
|
||||
)
|
||||
assert _resolve_run_scope(body, public_request, None) == RUN_SCOPE_EXTERNAL_ROOT
|
||||
assert (
|
||||
_resolve_run_scope(
|
||||
SimpleNamespace(metadata={}),
|
||||
public_request,
|
||||
RUN_SCOPE_ORCHESTRATION_CHILD,
|
||||
)
|
||||
== RUN_SCOPE_ORCHESTRATION_CHILD
|
||||
)
|
||||
|
||||
|
||||
def test_redis_concurrency_namespace_has_system_prefix() -> None:
|
||||
assert RedisConcurrencyConfig().namespace == "cmzs_deerflow:concurrency"
|
||||
gate = RedisConcurrencyGate(RedisConcurrencyConfig(namespace=""))
|
||||
assert gate._ns() == "cmzs_deerflow:concurrency"
|
||||
Reference in new issue
Block a user