claude-smart 0.2.44 → 0.2.46

package/plugin/vendor/reflexio/reflexio/server/llm/tools.py

@@ -13,7 +13,10 @@ logger = logging.getLogger(__name__)
 
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
-from reflexio.server.llm.llm_utils import make_strict_json_schema
+from reflexio.server.llm.llm_utils import (
+    assert_provider_safe_schema,
+    make_strict_json_schema,
+)
 from reflexio.server.llm.model_defaults import ModelRole, resolve_model_name
 
 if TYPE_CHECKING:
@@ -72,6 +75,11 @@ class Tool(BaseModel):
 
     def openai_spec(self) -> dict:
         parameters = self.args_model.model_json_schema()
+        # Boundary guard: tool-arg schemas bypass the response_format path and the
+        # registry contract test, so enforce provider-safety here too (raises under
+        # tests, warns in prod). Checked on the native schema to catch a model that
+        # forgot StrictStructuredOutput even when ``strict`` would later fold it.
+        assert_provider_safe_schema(parameters, name=self.args_model.__name__)
         if self.strict:
             parameters = make_strict_json_schema(parameters)
         return {
@@ -170,10 +178,17 @@ class ToolLoopResult(BaseModel):
 
     ctx: Any
     trace: ToolLoopTrace
-    finished_reason: Literal["finish_tool", "no_tool_call", "max_steps", "error"]
+    finished_reason: Literal[
+        "finish_tool", "structured_output", "no_tool_call", "max_steps", "error"
+    ]
     messages: list[dict[str, Any]] = Field(default_factory=list)
     pending_tool_call_ids: list[str] = Field(default_factory=list)
     max_steps_remaining: int = 0
+    # Set when finished_reason == "structured_output": the model ended the turn
+    # with a plain response parsed into the caller's ``response_format`` schema
+    # (the structured-output terminus, used by the extraction agent instead of a
+    # finish-sentinel tool call).
+    structured_output: BaseModel | None = None
 
 
 # Models we know support function calling per vendor docs but that litellm's
@@ -268,6 +283,36 @@ def _serialize_tool_result_for_history(result: dict[str, Any]) -> str:
     return f"{payload[:_MULTI_STAGE_RESULT_CHAR_CAP]}... [truncated]"
 
 
+def _normalize_tool_call_for_history(tool_call: Any) -> dict[str, Any]:
+    """Return an OpenAI-compatible plain dict for assistant tool-call history."""
+
+    if isinstance(tool_call, dict):
+        function = tool_call.get("function") or {}
+        if not isinstance(function, dict):
+            function = {
+                "name": getattr(function, "name", None),
+                "arguments": getattr(function, "arguments", None),
+            }
+        return {
+            "id": tool_call.get("id"),
+            "type": tool_call.get("type") or "function",
+            "function": {
+                "name": function.get("name"),
+                "arguments": function.get("arguments") or "{}",
+            },
+        }
+
+    function = getattr(tool_call, "function", None)
+    return {
+        "id": getattr(tool_call, "id", None),
+        "type": getattr(tool_call, "type", None) or "function",
+        "function": {
+            "name": getattr(function, "name", None),
+            "arguments": getattr(function, "arguments", None) or "{}",
+        },
+    }
+
+
 def _run_multi_stage_fallback(
     *,
     client: LiteLLMClient,
@@ -431,6 +476,7 @@ def run_tool_loop(
     fallback_schema: type[BaseModel] | None = None,
     fallback_tool_name: str | None = None,
     multi_stage_schema: type[BaseModel] | None = None,
+    response_format: type[BaseModel] | None = None,
     tool_choice: str | dict[str, Any] = "auto",
     log_label: str | None = None,
 ) -> ToolLoopResult:
@@ -471,6 +517,14 @@ def run_tool_loop(
             ``tool`` discriminator literal — that literal names the tool
             to dispatch, all other fields become its args. Takes priority
             over ``fallback_schema``.
+        response_format (type[BaseModel] | None): When set, the native tool loop
+            requests this schema as the model's structured response and treats a
+            turn with no tool call as the SUCCESS terminus — the parsed result is
+            returned on ``ToolLoopResult.structured_output`` with
+            ``finished_reason="structured_output"``. This is how the extraction
+            agent finishes (a direct structured answer) while still letting the
+            model call intermediate tools such as ``ask_human``. Leave unset for
+            finish-sentinel-tool loops.
         tool_choice (str | dict): Forwarded to each native tool-calling turn.
             Defaults to ``"auto"``. Pass an OpenAI tool-choice dict (e.g.
             ``{"type": "function", "function": {"name": "finish"}}``) to force a
@@ -500,7 +554,7 @@ def run_tool_loop(
     pending_tool_call_ids: list[str] = []
 
     # Lazily import the llm_io helpers only when logging is requested —
-    # matches classic's per-call lazy-import pattern in profile_deduplicator.py.
+    # matches classic's per-call lazy-import pattern in profile/components/consolidator.py.
     if log_label:
         from reflexio.server.services.service_utils import (
             log_llm_messages,
@@ -583,14 +637,16 @@ def run_tool_loop(
 
     local_msgs = list(messages)
     try:
+        tool_specs = registry.openai_specs()
         for _step in range(max_steps):
             if log_label:
                 log_llm_messages(logger, f"{log_label} (turn {_step + 1})", local_msgs)
             resp = client.generate_chat_response(
                 messages=local_msgs,
-                tools=registry.openai_specs(),
-                tool_choice=tool_choice,
+                tools=tool_specs or None,
+                tool_choice=tool_choice if tool_specs else None,
                 model_role=model_role,
+                response_format=response_format,
             )
             if log_label:
                 log_model_response(logger, f"{log_label} (turn {_step + 1})", resp)
@@ -611,12 +667,36 @@ def run_tool_loop(
 
             tool_calls = getattr(resp, "tool_calls", None)
             if not tool_calls:
-                # The model returned a plain-text turn with no tool calls. For a
-                # text-output agent this means "done", but the finish handler did
-                # NOT run, so no structured output was committed. Report a
-                # distinct reason so callers (and logs) don't conflate this with
-                # an actual finish_extraction call. Callers that require output
-                # (extraction) already gate success on committed output, so this
+                # No tool call this turn. When response_format was requested, this
+                # is the SUCCESS terminus: the model returned the final answer
+                # directly and the client parsed it into the schema (on the
+                # no-tools path the client returns the parsed BaseModel itself;
+                # with tools present it sets ``parsed_output``). This is the
+                # structured-output equivalent of a finish-sentinel tool call and
+                # is how the extraction agent commits its result.
+                if response_format is not None:
+                    structured = (
+                        resp
+                        if isinstance(resp, BaseModel)
+                        else getattr(resp, "parsed_output", None)
+                    )
+                    if isinstance(structured, BaseModel):
+                        trace.finished = True
+                        return ToolLoopResult(
+                            ctx=ctx,
+                            trace=trace,
+                            finished_reason="structured_output",
+                            structured_output=structured,
+                            messages=local_msgs,
+                            # The structured answer is committed on this turn —
+                            # one LLM call consumed, mirroring the finish_tool path.
+                            max_steps_remaining=max_steps - _step - 1,
+                        )
+                # No response_format requested (or nothing parseable): the finish
+                # handler did NOT run, so no structured output was committed.
+                # Report a distinct reason so callers (and logs) don't conflate
+                # this with an actual finish tool call. Callers that require
+                # output already gate success on committed output, so this
                 # surfaces accurately as a non-finish termination.
                 trace.finished = True
                 return ToolLoopResult(
@@ -627,21 +707,28 @@ def run_tool_loop(
                     pending_tool_call_ids=pending_tool_call_ids,
                     max_steps_remaining=max_steps - _step,
                 )
+            normalized_tool_calls = [
+                _normalize_tool_call_for_history(tc) for tc in tool_calls
+            ]
             # Emit ONE assistant message carrying ALL tool_calls from this turn.
             # OpenAI/Anthropic strict mode requires this shape.
             local_msgs.append(
-                {"role": "assistant", "content": None, "tool_calls": list(tool_calls)}
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": normalized_tool_calls,
+                }
             )
             # Process every tool call and append per-call tool result messages.
             # A single response's usage is attached to every turn it produced —
             # the summary helpers dedup by (model, prompt_tokens, completion_tokens).
-            for tc in tool_calls:
+            for tc in normalized_tool_calls:
                 # Time each tool individually — using the turn-start clock
                 # would inflate later tools' latencies with model time and
                 # earlier tools' work, masking the actual per-tool cost.
                 tool_t0 = time.monotonic()
-                name = tc.function.name
-                args_json = tc.function.arguments
+                name = tc["function"]["name"]
+                args_json = tc["function"]["arguments"]
                 outcome = registry.handle_outcome(name, args_json, ctx)
                 result = _tool_result_from_outcome(outcome, pending_tool_call_ids)
                 try:
@@ -664,13 +751,16 @@ def run_tool_loop(
                 local_msgs.append(
                     {
                         "role": "tool",
-                        "tool_call_id": tc.id,
+                        "tool_call_id": tc["id"],
                         "content": json.dumps(result),
                     }
                 )
             # After processing ALL tool calls, check whether the finish sentinel
             # appeared in this turn (may be alongside sibling calls).
-            if any(tc.function.name == finish_tool_name for tc in tool_calls):
+            if any(
+                tc["function"]["name"] == finish_tool_name
+                for tc in normalized_tool_calls
+            ):
                 trace.finished = True
                 return ToolLoopResult(
                     ctx=ctx,

package/plugin/vendor/reflexio/reflexio/server/operation_limiter.py

@@ -2,20 +2,25 @@
 
 from __future__ import annotations
 
+import logging
 import os
 import threading
 import time
 from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Literal
 
 from fastapi import HTTPException, status
 
+from reflexio.server.tracing import profile_step
 from reflexio.server.usage_metrics import record_usage_event
 
 OperationName = Literal["search", "publish", "aggregation"]
 
+logger = logging.getLogger(__name__)
+
 _DEFAULT_LIMITS: dict[OperationName, int] = {
     "search": 8,
     "publish": 4,
@@ -26,6 +31,12 @@ _DEFAULT_TIMEOUT_SECONDS: dict[OperationName, float] = {
     "publish": 5.0,
     "aggregation": 1.0,
 }
+_ESTIMATED_THREADS_PER_PUBLISH = 3
+_PUBLISH_WAIT_LOG_THRESHOLD_MS = 250
+_PUBLISH_PRESSURE_LOG_INTERVAL_SECONDS = 60.0
+_CGROUP_ROOT = Path("/sys/fs/cgroup")
+_PROC_SELF_CGROUP = Path("/proc/self/cgroup")
+_last_publish_pressure_log_by_org: dict[str, float] = {}
 
 
 @dataclass
@@ -33,6 +44,7 @@ class _LimiterState:
     semaphore: threading.BoundedSemaphore
     limit: int
     waiting: int = 0
+    active: int = 0
 
 
 _limiters: dict[tuple[str, OperationName], _LimiterState] = {}
@@ -60,6 +72,277 @@ def _timeout_for(operation: OperationName) -> float:
     return _DEFAULT_TIMEOUT_SECONDS[operation]
 
 
+def _resource_limit_nproc() -> int | str | None:
+    try:
+        import resource
+    except ImportError:  # pragma: no cover - platform-specific
+        return None
+    try:
+        soft, _hard = resource.getrlimit(resource.RLIMIT_NPROC)
+    except (OSError, ValueError, AttributeError):
+        return None
+    if soft == resource.RLIM_INFINITY:
+        return "unlimited"
+    return int(soft)
+
+
+def _linux_threads_max() -> int | None:
+    try:
+        with Path("/proc/sys/kernel/threads-max").open() as threads_max_file:
+            return int(threads_max_file.read().strip())
+    except (OSError, ValueError):
+        return None
+
+
+def _process_os_thread_count() -> int | None:
+    try:
+        with Path("/proc/self/status").open() as status_file:
+            for line in status_file:
+                if line.startswith("Threads:"):
+                    return int(line.split(":", maxsplit=1)[1].strip())
+    except (OSError, ValueError):
+        return None
+    return None
+
+
+def _available_memory_bytes() -> int | None:
+    try:
+        page_size = os.sysconf("SC_PAGE_SIZE")
+        available_pages = os.sysconf("SC_AVPHYS_PAGES")
+    except (OSError, ValueError, AttributeError):
+        return None
+    if isinstance(page_size, int) and isinstance(available_pages, int):
+        return page_size * available_pages
+    return None
+
+
+def _read_path_text(path: Path) -> str | None:
+    try:
+        return path.read_text().strip()
+    except OSError:
+        return None
+
+
+def _positive_int_from_text(text: str | None) -> int | None:
+    if text is None or text == "max":
+        return None
+    try:
+        value = int(text)
+    except ValueError:
+        return None
+    return value if value > 0 else None
+
+
+def _cgroup_candidate_paths(controller: str, filename: str) -> list[Path]:
+    candidates = [
+        _CGROUP_ROOT / filename,
+        _CGROUP_ROOT / controller / filename,
+    ]
+    cgroup_text = _read_path_text(_PROC_SELF_CGROUP)
+    if not cgroup_text:
+        return candidates
+
+    for line in cgroup_text.splitlines():
+        parts = line.split(":", maxsplit=2)
+        if len(parts) != 3:
+            continue
+        controllers_text = parts[1]
+        relative = parts[2].strip().lstrip("/")
+        relative_path = Path(relative) if relative else Path()
+        controllers = [item for item in controllers_text.split(",") if item]
+        if controllers_text == "":
+            candidates.append(_CGROUP_ROOT / relative_path / filename)
+            continue
+        if controller not in controllers:
+            continue
+        candidates.append(_CGROUP_ROOT / controller / relative_path / filename)
+        candidates.append(
+            _CGROUP_ROOT / ",".join(controllers) / relative_path / filename
+        )
+        candidates.append(_CGROUP_ROOT / relative_path / filename)
+    return candidates
+
+
+def _read_first_cgroup_value(controller: str, filename: str) -> str | None:
+    seen: set[Path] = set()
+    for path in _cgroup_candidate_paths(controller, filename):
+        if path in seen:
+            continue
+        seen.add(path)
+        value = _read_path_text(path)
+        if value:
+            return value
+    return None
+
+
+def _cgroup_v2_cpu_count() -> float | None:
+    raw = _read_first_cgroup_value("unified", "cpu.max")
+    if raw is None:
+        return None
+    parts = raw.split()
+    if len(parts) < 2 or parts[0] == "max":
+        return None
+    try:
+        quota = int(parts[0])
+        period = int(parts[1])
+    except ValueError:
+        return None
+    if quota <= 0 or period <= 0:
+        return None
+    return quota / period
+
+
+def _cgroup_v1_cpu_count() -> float | None:
+    quota = _positive_int_from_text(_read_first_cgroup_value("cpu", "cpu.cfs_quota_us"))
+    period = _positive_int_from_text(
+        _read_first_cgroup_value("cpu", "cpu.cfs_period_us")
+    )
+    if quota is None or period is None:
+        return None
+    return quota / period
+
+
+def _cgroup_cpu_count() -> float | None:
+    return _cgroup_v2_cpu_count() or _cgroup_v1_cpu_count()
+
+
+def _cgroup_memory_limit_bytes() -> int | None:
+    v2_limit = _positive_int_from_text(
+        _read_first_cgroup_value("unified", "memory.max")
+    )
+    if v2_limit is not None:
+        return v2_limit
+    return _positive_int_from_text(
+        _read_first_cgroup_value("memory", "memory.limit_in_bytes")
+    )
+
+
+def _cgroup_memory_current_bytes() -> int | None:
+    v2_current = _positive_int_from_text(
+        _read_first_cgroup_value("unified", "memory.current")
+    )
+    if v2_current is not None:
+        return v2_current
+    return _positive_int_from_text(
+        _read_first_cgroup_value("memory", "memory.usage_in_bytes")
+    )
+
+
+def _cgroup_pids_limit() -> int | None:
+    v2_limit = _positive_int_from_text(_read_first_cgroup_value("unified", "pids.max"))
+    if v2_limit is not None:
+        return v2_limit
+    return _positive_int_from_text(_read_first_cgroup_value("pids", "pids.max"))
+
+
+def _cgroup_pids_current() -> int | None:
+    v2_current = _positive_int_from_text(
+        _read_first_cgroup_value("unified", "pids.current")
+    )
+    if v2_current is not None:
+        return v2_current
+    return _positive_int_from_text(_read_first_cgroup_value("pids", "pids.current"))
+
+
+def _anyio_threadpool_stats() -> dict[str, int | None]:
+    try:
+        from anyio.to_thread import current_default_thread_limiter
+
+        stats = current_default_thread_limiter().statistics()
+        return {
+            "anyio_threadpool_total": int(stats.total_tokens),
+            "anyio_threadpool_borrowed": int(stats.borrowed_tokens),
+            "anyio_threadpool_waiting": int(stats.tasks_waiting),
+        }
+    except Exception:  # noqa: BLE001 - best-effort diagnostics only
+        return {
+            "anyio_threadpool_total": None,
+            "anyio_threadpool_borrowed": None,
+            "anyio_threadpool_waiting": None,
+        }
+
+
+def publish_hardware_capacity_snapshot() -> dict[str, int | str | float | None]:
+    """Return hardware/threadpool diagnostics for publish concurrency tuning."""
+    host_cpu_count = os.cpu_count()
+    cgroup_cpu_count = _cgroup_cpu_count()
+    effective_cpu_count = cgroup_cpu_count or host_cpu_count
+    process_os_threads = _process_os_thread_count()
+    nproc_limit = _resource_limit_nproc()
+    threads_max = _linux_threads_max()
+    cgroup_pids_limit = _cgroup_pids_limit()
+    cgroup_pids_current = _cgroup_pids_current()
+    cgroup_memory_limit = _cgroup_memory_limit_bytes()
+    cgroup_memory_current = _cgroup_memory_current_bytes()
+    cgroup_available_memory = (
+        max(0, cgroup_memory_limit - cgroup_memory_current)
+        if cgroup_memory_limit is not None and cgroup_memory_current is not None
+        else None
+    )
+    anyio_stats = _anyio_threadpool_stats()
+    host_thread_budget_candidates = [
+        value - process_os_threads
+        for value in (nproc_limit, threads_max)
+        if isinstance(value, int) and process_os_threads is not None
+    ]
+    cgroup_pid_budget = (
+        cgroup_pids_limit - cgroup_pids_current
+        if cgroup_pids_limit is not None and cgroup_pids_current is not None
+        else None
+    )
+    thread_budget_candidates = [
+        value
+        for value in (cgroup_pid_budget, *host_thread_budget_candidates)
+        if isinstance(value, int)
+    ]
+    max_by_os_threads = (
+        max(1, min(thread_budget_candidates) // _ESTIMATED_THREADS_PER_PUBLISH)
+        if thread_budget_candidates
+        else None
+    )
+    guidance_candidates = [
+        value
+        for value in (
+            anyio_stats["anyio_threadpool_total"],
+            max_by_os_threads,
+            max(1, int(effective_cpu_count * 2))
+            if effective_cpu_count is not None
+            else None,
+        )
+        if isinstance(value, int) and value > 0
+    ]
+    hardware_guidance = min(guidance_candidates) if guidance_candidates else None
+    return {
+        "configured_publish_limit": _limit_for("publish"),
+        "publish_timeout_seconds": _timeout_for("publish"),
+        "hardware_guidance_publish_limit": hardware_guidance,
+        "estimated_threads_per_publish": _ESTIMATED_THREADS_PER_PUBLISH,
+        "effective_cpu_count": effective_cpu_count,
+        "logical_cpu_count": host_cpu_count,
+        "cgroup_cpu_count": cgroup_cpu_count,
+        "python_active_threads": threading.active_count(),
+        "process_os_threads": process_os_threads,
+        "rlimit_nproc_soft": nproc_limit,
+        "linux_threads_max": threads_max,
+        "cgroup_pids_limit": cgroup_pids_limit,
+        "cgroup_pids_current": cgroup_pids_current,
+        "cgroup_memory_limit_bytes": cgroup_memory_limit,
+        "cgroup_memory_current_bytes": cgroup_memory_current,
+        "available_memory_bytes": cgroup_available_memory
+        if cgroup_available_memory is not None
+        else _available_memory_bytes(),
+        **anyio_stats,
+    }
+
+
+def log_publish_hardware_capacity() -> None:
+    """Log one startup snapshot for publish concurrency hardware tuning."""
+    logger.info(
+        "publish_concurrency_hardware_capacity %s",
+        publish_hardware_capacity_snapshot(),
+    )
+
+
 def _state_for(org_id: str, operation: OperationName) -> _LimiterState:
     key = (str(org_id), operation)
     desired_limit = _limit_for(operation)
@@ -83,6 +366,7 @@ def _record_limiter_event(
     wait_ms: int,
     waiting: int,
     limit: int,
+    active: int,
 ) -> None:
     record_usage_event(
         org_id=org_id,
@@ -95,31 +379,64 @@ def _record_limiter_event(
             "operation": operation,
             "waiting": waiting,
             "limit": limit,
+            "active": active,
         },
     )
 
 
+def _should_log_publish_pressure(org_id: str, now: float) -> bool:
+    with _limiters_lock:
+        last_logged = _last_publish_pressure_log_by_org.get(org_id)
+        if (
+            last_logged is not None
+            and now - last_logged < _PUBLISH_PRESSURE_LOG_INTERVAL_SECONDS
+        ):
+            return False
+        _last_publish_pressure_log_by_org[org_id] = now
+        return True
+
+
 @contextmanager
 def operation_limit(
     org_id: str,
     operation: OperationName,
     *,
     timeout_seconds: float | None = None,
+    wait_forever: bool = False,
 ) -> Any:
     """Acquire the per-org limiter for an operation, recording wait metrics."""
     state = _state_for(org_id, operation)
-    timeout = _timeout_for(operation) if timeout_seconds is None else timeout_seconds
+    timeout = (
+        None
+        if wait_forever
+        else (_timeout_for(operation) if timeout_seconds is None else timeout_seconds)
+    )
     start = time.perf_counter()
     with _limiters_lock:
         state.waiting += 1
         waiting = state.waiting
     acquired = False
     try:
-        acquired = state.semaphore.acquire(timeout=timeout)
-        wait_ms = int((time.perf_counter() - start) * 1000)
+        with profile_step(
+            f"{operation}.operation_limit.acquire",
+            limit=state.limit,
+            waiting=waiting,
+            wait_forever=wait_forever,
+        ) as span:
+            acquired = (
+                state.semaphore.acquire()
+                if timeout is None
+                else state.semaphore.acquire(timeout=timeout)
+            )
+            wait_ms = int((time.perf_counter() - start) * 1000)
+            span.set_data("wait_ms", wait_ms)
+            span.set_data("acquired", acquired)
         with _limiters_lock:
             state.waiting -= 1
             waiting_after = state.waiting
+            if acquired:
+                state.active += 1
+            active = state.active
         if not acquired:
             _record_limiter_event(
                 org_id=org_id,
@@ -129,10 +446,39 @@ def operation_limit(
                 wait_ms=wait_ms,
                 waiting=waiting_after,
                 limit=state.limit,
+                active=active,
             )
+            if operation == "publish":
+                logger.warning(
+                    "publish_limiter_timeout org_id=%s limit=%s active=%s "
+                    "waiting=%s wait_ms=%s timeout_seconds=%s hardware=%s",
+                    org_id,
+                    state.limit,
+                    active,
+                    waiting_after,
+                    wait_ms,
+                    timeout,
+                    publish_hardware_capacity_snapshot(),
+                )
             raise TimeoutError(
                 f"{operation} concurrency limit reached for org {org_id}"
             )
+        if (
+            operation == "publish"
+            and wait_ms >= _PUBLISH_WAIT_LOG_THRESHOLD_MS
+            and _should_log_publish_pressure(str(org_id), time.monotonic())
+        ):
+            logger.info(
+                "publish_limiter_wait org_id=%s limit=%s active=%s waiting=%s "
+                "wait_ms=%s timeout_seconds=%s hardware=%s",
+                org_id,
+                state.limit,
+                active,
+                waiting_after,
+                wait_ms,
+                timeout,
+                publish_hardware_capacity_snapshot(),
+            )
         _record_limiter_event(
             org_id=org_id,
             operation=operation,
@@ -141,10 +487,13 @@ def operation_limit(
             wait_ms=wait_ms,
             waiting=waiting,
             limit=state.limit,
+            active=active,
         )
         yield
     finally:
         if acquired:
+            with _limiters_lock:
+                state.active = max(0, state.active - 1)
             state.semaphore.release()
 
 
@@ -154,8 +503,14 @@ def run_with_operation_limit[T](
     operation: OperationName,
     fn: Callable[[], T],
     timeout_seconds: float | None = None,
+    wait_forever: bool = False,
 ) -> T:
-    with operation_limit(org_id, operation, timeout_seconds=timeout_seconds):
+    with operation_limit(
+        org_id,
+        operation,
+        timeout_seconds=timeout_seconds,
+        wait_forever=wait_forever,
+    ):
         return fn()
 
 
@@ -177,3 +532,4 @@ def limiter_http_exception(operation: OperationName) -> HTTPException:
 def reset_operation_limiters_for_tests() -> None:
     with _limiters_lock:
         _limiters.clear()
+        _last_publish_pressure_log_by_org.clear()