mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-09 03:39:37 -10:00
Compare commits
29
Commits
No files matched your search
@@ -41,4 +41,4 @@ jobs:
|
||||
|
||||
- name: Run tests
|
||||
run: |
|
||||
cd src/praisonai && python -m pytest tests/unit/ -v --tb=short --disable-warnings --timeout=60 || echo "Some tests failed"
|
||||
cd src/praisonai && python -m pytest tests/unit/ -v --tb=short --disable-warnings --timeout=60
|
||||
@@ -155,7 +155,9 @@ jobs:
|
||||
tests/unit/plugins/ tests/unit/runtime/ tests/unit/tools/
|
||||
tests/unit/bus/ tests/unit/skills/ tests/unit/agents/
|
||||
tests/unit/rag/ tests/unit/streaming/
|
||||
tests/unit/test_adapter_registry.py
|
||||
tests/unit/test_adapter_registry.py tests/unit/test_token_budget.py
|
||||
tests/unit/test_streaming_events.py tests/unit/test_tool_decorator_approval.py
|
||||
tests/unit/test_task_execution_config.py
|
||||
# Both tests previously deselected here now pass in combination.
|
||||
# test_lazy_litellm_import was looking the loader module up by
|
||||
# name and could get a different object than the one get_litellm()
|
||||
|
||||
@@ -3571,8 +3571,26 @@ Your Goal: {self.goal}
|
||||
|
||||
def _replace_chat_history(self, new_history: List[Dict[str, Any]]):
|
||||
"""Thread-safe replacement of entire chat history using proper locking."""
|
||||
with self._history_lock.lock():
|
||||
self._history_lock.value[:] = new_history
|
||||
# Share _auto_save_session's lock so a save cannot interleave with the
|
||||
# rewrite and the cursor update below (#6217).
|
||||
with self._auto_save_lock, self._history_lock.lock():
|
||||
history = self._history_lock.value
|
||||
# Re-sync the positional auto-save cursor (#6217): resume at the
|
||||
# first unsaved survivor, or right after the last survivor so
|
||||
# appended messages (e.g. durable replays) are saved; summaries
|
||||
# placed before survivors are not re-persisted.
|
||||
last = getattr(self, '_auto_save_last_index', 0)
|
||||
old = {id(m): i >= last for i, m in enumerate(history)}
|
||||
cursor = len(new_history)
|
||||
survivors = [i for i, m in enumerate(new_history) if id(m) in old]
|
||||
if survivors:
|
||||
cursor = survivors[-1] + 1
|
||||
for i in survivors:
|
||||
if old[id(new_history[i])]:
|
||||
cursor = i
|
||||
break
|
||||
history[:] = new_history
|
||||
self._auto_save_last_index = cursor
|
||||
|
||||
@property
|
||||
def _cache_lock(self):
|
||||
|
||||
@@ -4485,7 +4485,9 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
except Exception:
|
||||
# Catch any exception from the inner try block and re-raise to outer handler
|
||||
raise
|
||||
except ToolExecutionError:
|
||||
except (ToolExecutionError, BudgetExceededError):
|
||||
# Parity with the custom-LLM and async paths: a budget stop must
|
||||
# reach the caller, not turn into a ``None`` reply.
|
||||
raise
|
||||
except Exception as e:
|
||||
self._last_stop_reason = "error"
|
||||
|
||||
@@ -55,6 +55,27 @@ def _live_turn_message_ids() -> set:
|
||||
ids.update(id(m) for m in owned)
|
||||
return ids
|
||||
|
||||
|
||||
def _exchange_spans(history: list) -> list:
|
||||
"""Split history into ``(start, end)`` units that must stay together.
|
||||
|
||||
An assistant turn carrying ``tool_calls`` and the ``role="tool"`` results
|
||||
that follow it form one unit; every other message is its own unit. History
|
||||
mutators cut/delete on unit boundaries so they never leave an orphaned tool
|
||||
result or a dangling ``tool_calls`` turn, which providers reject (#6217).
|
||||
"""
|
||||
spans = []
|
||||
i, n = 0, len(history)
|
||||
while i < n:
|
||||
j = i + 1
|
||||
msg = history[i]
|
||||
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
||||
while j < n and history[j].get("role") == "tool":
|
||||
j += 1
|
||||
spans.append((i, j))
|
||||
i = j
|
||||
return spans
|
||||
|
||||
# Shared lazy display helpers (cached, thread-safe; avoid circular imports)
|
||||
from ._lazy_display import _get_console, _get_live, _get_display_functions
|
||||
|
||||
@@ -178,6 +199,11 @@ class MemoryMixin:
|
||||
length: Target length for chat history
|
||||
"""
|
||||
with self._history_lock:
|
||||
# Never cut inside a tool exchange: back off to its start (#6217).
|
||||
for start, end in _exchange_spans(self.chat_history):
|
||||
if start < length < end:
|
||||
length = start
|
||||
break
|
||||
self.chat_history = self.chat_history[:length]
|
||||
|
||||
def _begin_turn_tracking(self):
|
||||
@@ -299,8 +325,14 @@ class MemoryMixin:
|
||||
if len(self.chat_history) <= keep_last:
|
||||
return 0
|
||||
|
||||
deleted_count = len(self.chat_history) - keep_last
|
||||
self.chat_history = self.chat_history[-keep_last:]
|
||||
cut = len(self.chat_history) - keep_last
|
||||
# Never start the kept tail inside a tool exchange (#6217).
|
||||
for start, end in _exchange_spans(self.chat_history):
|
||||
if start < cut < end:
|
||||
cut = end
|
||||
break
|
||||
deleted_count = cut
|
||||
self.chat_history = self.chat_history[cut:]
|
||||
# Reset auto-save index to match new history length
|
||||
self._auto_save_last_index = len(self.chat_history)
|
||||
return deleted_count
|
||||
@@ -318,8 +350,17 @@ class MemoryMixin:
|
||||
True if deleted, False if index out of range
|
||||
"""
|
||||
with self._history_lock:
|
||||
n = len(self.chat_history)
|
||||
if not -n <= index < n:
|
||||
return False
|
||||
index %= n
|
||||
# Deleting part of a tool exchange removes the whole exchange so no
|
||||
# orphaned tool result / dangling tool_calls is left (#6217).
|
||||
for start, end in _exchange_spans(self.chat_history):
|
||||
if start <= index < end:
|
||||
break
|
||||
try:
|
||||
del self.chat_history[index]
|
||||
del self.chat_history[start:end]
|
||||
# Adjust auto-save index if deletion affects saved range
|
||||
if hasattr(self, '_auto_save_last_index'):
|
||||
self._auto_save_last_index = min(
|
||||
@@ -344,10 +385,18 @@ class MemoryMixin:
|
||||
"""
|
||||
with self._history_lock:
|
||||
original_len = len(self.chat_history)
|
||||
self.chat_history = [
|
||||
msg for msg in self.chat_history
|
||||
if pattern.lower() not in msg.get("content", "").lower()
|
||||
]
|
||||
needle = pattern.lower()
|
||||
history = self.chat_history
|
||||
kept = []
|
||||
# Match per exchange so a hit inside a tool exchange drops the whole
|
||||
# exchange; content may be None on tool-calling turns (#6217).
|
||||
for start, end in _exchange_spans(history):
|
||||
unit = history[start:end]
|
||||
if not any(
|
||||
needle in str(m.get("content") or "").lower() for m in unit
|
||||
):
|
||||
kept.extend(unit)
|
||||
self.chat_history = kept
|
||||
deleted_count = original_len - len(self.chat_history)
|
||||
# Adjust auto-save index if deletion affects saved range
|
||||
if deleted_count > 0 and hasattr(self, '_auto_save_last_index'):
|
||||
@@ -569,11 +618,10 @@ class MemoryMixin:
|
||||
if history:
|
||||
# Only restore if chat_history is empty (avoid duplicates)
|
||||
if not self.chat_history:
|
||||
# The store returns LLM-shaped dicts (to_llm_message), so
|
||||
# keep tool_calls / tool_call_id intact (#6217).
|
||||
for msg in history:
|
||||
self.chat_history.append({
|
||||
"role": msg["role"],
|
||||
"content": msg["content"]
|
||||
})
|
||||
self.chat_history.append(dict(msg))
|
||||
logging.info(f"Restored session {self._session_id} with {len(history)} messages from JSON store")
|
||||
|
||||
# Set agent info
|
||||
@@ -862,7 +910,7 @@ class MemoryMixin:
|
||||
except Exception as e:
|
||||
logging.debug(f"Discard partial assistant turn skipped: {e}")
|
||||
|
||||
def _add_message_with_tool_fields(self, role, content, **tool_fields):
|
||||
def _add_message_with_tool_fields(self, role, content, session_id=None, **tool_fields):
|
||||
"""Add a message carrying tool fields, tolerating stores without them.
|
||||
|
||||
The default JSON store and the SQLite store accept ``tool_calls`` /
|
||||
@@ -874,11 +922,13 @@ class MemoryMixin:
|
||||
back to a plain-text turn so no exchange is silently lost.
|
||||
"""
|
||||
add_message = self._session_store.add_message
|
||||
if self._store_supports_tool_fields(add_message):
|
||||
add_message(self._session_id, role, content, **tool_fields)
|
||||
if session_id is None:
|
||||
session_id = self._session_id
|
||||
if tool_fields and self._store_supports_tool_fields(add_message):
|
||||
add_message(session_id, role, content, **tool_fields)
|
||||
return
|
||||
# Legacy store: preserve the turn as plain text so resume still sees it.
|
||||
add_message(self._session_id, role, content)
|
||||
add_message(session_id, role, content)
|
||||
|
||||
def _store_supports_tool_fields(self, add_message) -> bool:
|
||||
"""Return True if ``add_message`` accepts the tool keyword fields."""
|
||||
@@ -1195,10 +1245,17 @@ class MemoryMixin:
|
||||
if self._session_store is not None:
|
||||
# Persist only NEW messages to SessionStore
|
||||
for msg in new_messages:
|
||||
self._session_store.add_message(
|
||||
self.auto_save, # Use auto_save name as session_id
|
||||
role=msg.get("role", "user"),
|
||||
content=msg.get("content", ""),
|
||||
tool_fields = {
|
||||
k: msg[k]
|
||||
for k in ("tool_calls", "tool_call_id")
|
||||
if msg.get(k) is not None
|
||||
}
|
||||
# Persist tool fields so resumed tool turns stay valid (#6217)
|
||||
self._add_message_with_tool_fields(
|
||||
msg.get("role", "user"),
|
||||
msg.get("content") or "",
|
||||
session_id=self.auto_save,
|
||||
**tool_fields,
|
||||
)
|
||||
# Update last saved index
|
||||
self._auto_save_last_index = len(clean_history)
|
||||
|
||||
@@ -957,7 +957,8 @@ def _validate_feature_config(feature: str, config: Dict[str, Any]) -> List[str]:
|
||||
"guardrails": {"enabled", "max_retries", "validator"},
|
||||
"web": {"enabled", "search", "fetch", "provider"},
|
||||
"output": {"enabled", "verbose", "markdown", "stream", "metrics", "preset"},
|
||||
"execution": {"enabled", "max_iter", "max_rpm", "max_execution_time", "preset"},
|
||||
"execution": {"enabled", "max_iter", "max_rpm", "max_execution_time", "preset",
|
||||
"max_budget", "on_budget_exceeded", "durable", "journal_path", "resume_run_id"},
|
||||
"caching": {"enabled", "prompt_caching", "ttl"},
|
||||
"autonomy": {"enabled", "max_steps", "approval_required"},
|
||||
"context": {"enabled", "max_tokens", "strategy"},
|
||||
|
||||
@@ -7,6 +7,15 @@ from datetime import datetime
|
||||
from typing import Dict, List, Optional, Any
|
||||
from .decorator import tool
|
||||
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"praisonaiagents.tools.jira_tools is deprecated and will be removed in a future release; "
|
||||
"use praisonai_tools.tools.jira_watch_tool instead (pip install praisonai-tools)",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def _get_jira_connection(
|
||||
|
||||
@@ -26,6 +26,15 @@ import logging
|
||||
import os
|
||||
from importlib import util
|
||||
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"praisonaiagents.tools.youdotcom_tools is deprecated and will be removed in a future release; "
|
||||
"use praisonai_tools.tools.youdotcom_tool instead (pip install praisonai-tools)",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
|
||||
def _check_youdotcom_available() -> tuple[bool, Optional[str]]:
|
||||
"""Check if You.com SDK is available and API key is set.
|
||||
|
||||
@@ -473,7 +473,10 @@ class YAMLWorkflowParser:
|
||||
|
||||
if 'reasoning' in role_config:
|
||||
agent['reasoning'] = role_config['reasoning']
|
||||
|
||||
|
||||
if isinstance(role_config.get('execution'), dict):
|
||||
agent['execution'] = role_config['execution']
|
||||
|
||||
agents[role_id] = agent
|
||||
|
||||
return agents
|
||||
@@ -591,7 +594,11 @@ class YAMLWorkflowParser:
|
||||
# so we don't override the Agent's own caching default.
|
||||
if cache is False:
|
||||
agent_kwargs['caching'] = False
|
||||
_execution: Dict[str, Any] = {}
|
||||
# An agent-level `execution:` mapping passes straight through to
|
||||
# ExecutionConfig (e.g. max_budget, on_budget_exceeded, durable,
|
||||
# journal_path, resume_run_id); flat max_iter/max_execution_time win.
|
||||
_yaml_execution = config.get('execution')
|
||||
_execution: Dict[str, Any] = dict(_yaml_execution) if isinstance(_yaml_execution, dict) else {}
|
||||
if max_iter is not None:
|
||||
_execution['max_iter'] = max_iter
|
||||
if max_execution_time is not None:
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
"""Sync OpenAI-client chat must surface BudgetExceededError, not return None."""
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent, BudgetExceededError, ExecutionConfig
|
||||
|
||||
|
||||
def test_sync_chat_reraises_budget_error(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test") # client is built lazily; no request is sent
|
||||
agent = Agent(
|
||||
name="b", instructions="x", llm="gpt-4o-mini", output="silent",
|
||||
execution=ExecutionConfig(max_budget=0.01),
|
||||
)
|
||||
with patch.object(
|
||||
agent, "_chat_completion",
|
||||
side_effect=BudgetExceededError("b", 0.02, 0.01),
|
||||
):
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("hello")
|
||||
@@ -0,0 +1,137 @@
|
||||
"""Chat-history integrity for tool turns (Issue #6217)."""
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.config import MemoryConfig
|
||||
from praisonaiagents.session import DefaultSessionStore
|
||||
|
||||
|
||||
def _tool_history():
|
||||
return [
|
||||
{"role": "user", "content": "weather in Paris?"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [{
|
||||
"id": "call_1", "type": "function",
|
||||
"function": {"name": "weather", "arguments": "{\"city\":\"Paris\"}"},
|
||||
}]},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "18C sunny"},
|
||||
{"role": "assistant", "content": "It is 18C and sunny."},
|
||||
]
|
||||
|
||||
|
||||
def _agent(session_id=None, store=None):
|
||||
kwargs = {}
|
||||
if session_id:
|
||||
kwargs["memory"] = MemoryConfig(history=True, auto_save=session_id)
|
||||
a = Agent(name="a", instructions="x", model="gpt-4o-mini", **kwargs)
|
||||
if store is not None:
|
||||
a._session_store = store
|
||||
a._init_session_store()
|
||||
return a
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path):
|
||||
return DefaultSessionStore(session_dir=str(tmp_path))
|
||||
|
||||
|
||||
def _roles(a):
|
||||
return [m["role"] for m in a.chat_history]
|
||||
|
||||
|
||||
def test_auto_save_round_trips_tool_fields(store):
|
||||
a = _agent("sess_rt", store)
|
||||
a.chat_history.extend(_tool_history())
|
||||
a._auto_save_session()
|
||||
|
||||
b = _agent("sess_rt", store)
|
||||
restored = b.chat_history
|
||||
assert restored[1]["tool_calls"][0]["id"] == "call_1"
|
||||
assert restored[2]["tool_call_id"] == "call_1"
|
||||
assert restored[2]["content"] == "18C sunny"
|
||||
|
||||
|
||||
def test_replace_chat_history_resyncs_auto_save_cursor(store):
|
||||
a = _agent("sess_cmp", store)
|
||||
|
||||
def turn(i):
|
||||
a._append_to_chat_history({"role": "user", "content": f"q{i}"})
|
||||
a._append_to_chat_history({"role": "assistant", "content": f"a{i}"})
|
||||
a._auto_save_session()
|
||||
|
||||
for i in range(3):
|
||||
turn(i)
|
||||
a._replace_chat_history([
|
||||
{"role": "system", "content": "summary"},
|
||||
{"role": "assistant", "content": "a2"},
|
||||
])
|
||||
for i in range(3, 6):
|
||||
turn(i)
|
||||
saved = [m["content"] for m in store.get_chat_history("sess_cmp")]
|
||||
assert saved == ["q0", "a0", "q1", "a1", "q2", "a2",
|
||||
"q3", "a3", "q4", "a4", "q5", "a5"]
|
||||
|
||||
|
||||
def test_replace_chat_history_keeps_unsaved_survivors(store):
|
||||
a = _agent("sess_unsaved", store)
|
||||
a._append_to_chat_history({"role": "user", "content": "q0"})
|
||||
a._auto_save_session()
|
||||
pending = {"role": "user", "content": "q1"}
|
||||
a._append_to_chat_history(pending)
|
||||
# Compaction mid-turn keeps the not-yet-saved message in its tail.
|
||||
a._replace_chat_history([{"role": "system", "content": "summary"}, pending])
|
||||
a._auto_save_session()
|
||||
saved = [m["content"] for m in store.get_chat_history("sess_unsaved")]
|
||||
assert saved == ["q0", "q1"]
|
||||
|
||||
|
||||
def test_replace_chat_history_saves_appended_messages(store):
|
||||
# Durable resume appends replayed tool pairs to a fully-saved history.
|
||||
a = _agent("sess_replay", store)
|
||||
a._append_to_chat_history({"role": "user", "content": "q0"})
|
||||
a._auto_save_session()
|
||||
replay = _tool_history()[1:3]
|
||||
a._replace_chat_history(list(a.chat_history) + replay)
|
||||
a._auto_save_session()
|
||||
saved = store.get_chat_history("sess_replay")
|
||||
assert [m["role"] for m in saved] == ["user", "assistant", "tool"]
|
||||
assert saved[2]["tool_call_id"] == "call_1"
|
||||
|
||||
|
||||
def test_prune_history_does_not_orphan_tool_result():
|
||||
a = _agent()
|
||||
a.chat_history.extend(_tool_history())
|
||||
a.prune_history(keep_last=2)
|
||||
assert _roles(a) == ["assistant"]
|
||||
|
||||
|
||||
def test_delete_history_removes_whole_exchange():
|
||||
for idx in (1, 2):
|
||||
a = _agent()
|
||||
a.chat_history.extend(_tool_history())
|
||||
assert a.delete_history(idx) is True
|
||||
assert _roles(a) == ["user", "assistant"]
|
||||
a = _agent()
|
||||
a.chat_history.extend(_tool_history())
|
||||
assert a.delete_history(10) is False
|
||||
assert a.delete_history(-1) is True
|
||||
assert len(a.chat_history) == 3
|
||||
|
||||
|
||||
def test_truncate_does_not_leave_dangling_tool_calls():
|
||||
a = _agent()
|
||||
a.chat_history.extend(_tool_history())
|
||||
a._truncate_chat_history(2)
|
||||
assert _roles(a) == ["user"]
|
||||
a = _agent()
|
||||
a.chat_history.extend(_tool_history())
|
||||
a._truncate_chat_history(3)
|
||||
assert len(a.chat_history) == 3
|
||||
|
||||
|
||||
def test_delete_history_matching_tolerates_none_content():
|
||||
a = _agent()
|
||||
a.chat_history.extend(_tool_history())
|
||||
assert a.delete_history_matching("image") == 0
|
||||
# A hit on the tool result drops its whole exchange.
|
||||
assert a.delete_history_matching("18c sunny") == 2
|
||||
assert _roles(a) == ["user", "assistant"]
|
||||
@@ -0,0 +1,57 @@
|
||||
"""An agent's YAML `execution:` block reaches ExecutionConfig (Issue #6206)."""
|
||||
from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
|
||||
|
||||
|
||||
def test_execution_block_passes_budget_and_durable_keys():
|
||||
agent = YAMLWorkflowParser()._create_agent("worker", {
|
||||
"role": "Worker",
|
||||
"execution": {
|
||||
"max_budget": 1.5,
|
||||
"on_budget_exceeded": "warn",
|
||||
"durable": True,
|
||||
"resume_run_id": "run-123",
|
||||
},
|
||||
"max_iter": 7,
|
||||
})
|
||||
assert agent._max_budget == 1.5
|
||||
assert agent._on_budget_exceeded == "warn"
|
||||
assert agent.execution.durable is True
|
||||
assert agent.execution.resume_run_id == "run-123"
|
||||
assert agent.execution.max_iter == 7
|
||||
|
||||
|
||||
def test_flat_max_iter_overrides_execution_block():
|
||||
agent = YAMLWorkflowParser()._create_agent("worker", {
|
||||
"execution": {"max_iter": 3},
|
||||
"max_iter": 9,
|
||||
})
|
||||
assert agent.execution.max_iter == 9
|
||||
|
||||
|
||||
def test_no_execution_block_keeps_defaults():
|
||||
agent = YAMLWorkflowParser()._create_agent("worker", {"role": "Worker"})
|
||||
assert not agent._max_budget
|
||||
|
||||
|
||||
def test_roles_format_keeps_execution_block():
|
||||
"""The agents.yaml `roles:` format must not drop `execution:` in conversion."""
|
||||
parser = YAMLWorkflowParser()
|
||||
parser.parse_string(
|
||||
"framework: praisonai\n"
|
||||
"topic: t\n"
|
||||
"roles:\n"
|
||||
" worker:\n"
|
||||
" role: Worker\n"
|
||||
" goal: g\n"
|
||||
" backstory: b\n"
|
||||
" execution:\n"
|
||||
" max_budget: 2.0\n"
|
||||
" on_budget_exceeded: warn\n"
|
||||
" tasks:\n"
|
||||
" t1:\n"
|
||||
" description: do it\n"
|
||||
" expected_output: done\n"
|
||||
)
|
||||
agent = parser._agents["worker"]
|
||||
assert agent._max_budget == 2.0
|
||||
assert agent._on_budget_exceeded == "warn"
|
||||
@@ -28,7 +28,9 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import itertools
|
||||
import logging
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Dict, Optional
|
||||
|
||||
@@ -106,9 +108,14 @@ class AdmissionRejected(Exception):
|
||||
busy acknowledgement to the user instead of an opaque error.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str = "Busy right now — please resend in a moment."):
|
||||
def __init__(
|
||||
self,
|
||||
message: str = "Busy right now — please resend in a moment.",
|
||||
retry_after: Optional[float] = None,
|
||||
):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.retry_after = retry_after
|
||||
|
||||
|
||||
class _TurnSlot:
|
||||
@@ -249,8 +256,16 @@ class AdmissionGate:
|
||||
*,
|
||||
resource_policy: Optional[object] = None,
|
||||
resource_sampler: Optional[object] = None,
|
||||
rate_limit_policy: Optional[object] = None,
|
||||
):
|
||||
self._policy = policy
|
||||
# Issue #6215: optional per-identity request-rate ceiling (a core
|
||||
# ``RateLimitPolicyProtocol``, e.g. ``SlidingWindowRateLimitPolicy``)
|
||||
# checked before a turn is admitted, keyed by the caller's rate key
|
||||
# (defaults to the session id).
|
||||
self._rate_limit_policy = rate_limit_policy
|
||||
self.rate_limited = 0
|
||||
self._rate_prune_after = 0.0
|
||||
# Optional memory/resource-pressure policy (Issue #3445). When set, a
|
||||
# lightweight RSS sample is folded into each admit decision so the gate
|
||||
# queues under soft pressure and sheds under hard pressure before the
|
||||
@@ -303,6 +318,16 @@ class AdmissionGate:
|
||||
# otherwise assume active (it was explicitly supplied).
|
||||
return bool(getattr(pol, "enabled", True))
|
||||
|
||||
def set_rate_limit_policy(self, policy: Optional[object]) -> None:
|
||||
"""Swap the per-identity rate budget in place (config hot-reload)."""
|
||||
self._rate_limit_policy = policy
|
||||
self._rate_prune_after = 0.0
|
||||
|
||||
@property
|
||||
def _rate_limit_enabled(self) -> bool:
|
||||
pol = self._rate_limit_policy
|
||||
return pol is not None and bool(getattr(pol, "enabled", True))
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
"""Whether the gate is active.
|
||||
@@ -313,7 +338,7 @@ class AdmissionGate:
|
||||
"""
|
||||
if self._policy is not None and self._max > 0:
|
||||
return True
|
||||
return self._resource_enabled
|
||||
return self._resource_enabled or self._rate_limit_enabled
|
||||
|
||||
@property
|
||||
def in_flight(self) -> int:
|
||||
@@ -337,6 +362,10 @@ class AdmissionGate:
|
||||
"admitted": self.admitted,
|
||||
"rejected": self.rejected,
|
||||
"shed": self.shed,
|
||||
"per_user_per_minute": int(
|
||||
getattr(self._rate_limit_policy, "max_requests", 0) or 0
|
||||
),
|
||||
"rate_limited": self.rate_limited,
|
||||
"max_rss_mb": float(
|
||||
getattr(self._resource_policy, "hard_rss_mb", 0.0) or 0.0
|
||||
),
|
||||
@@ -362,9 +391,14 @@ class AdmissionGate:
|
||||
return self._sem
|
||||
|
||||
@asynccontextmanager
|
||||
async def admit(self, *, session_id: str = ""):
|
||||
async def admit(self, *, session_id: str = "", rate_key: Optional[str] = None):
|
||||
"""Admit (or shed) an inbound turn, holding a slot for its duration.
|
||||
|
||||
``rate_key`` is the identity charged against the per-user request
|
||||
budget. It defaults to ``session_id``; surfaces where the session id is
|
||||
client-chosen (API headers, WebSocket sessions) pass a stable caller
|
||||
identity instead so a new conversation cannot reset the budget.
|
||||
|
||||
Yields when capacity is available (immediately or after waiting in the
|
||||
bounded queue). Raises :class:`AdmissionRejected` when the policy sheds
|
||||
the turn. A no-op pass-through when the gate is disabled.
|
||||
@@ -374,6 +408,8 @@ class AdmissionGate:
|
||||
yield
|
||||
return
|
||||
|
||||
self._check_rate_limit(rate_key or session_id)
|
||||
|
||||
# Issue #5168: the scope key is the caller-supplied session_id (the
|
||||
# per-user/per-tenant storage key at the call sites). Only tracked when
|
||||
# a per-scope sub-limit is active, so the unscoped path is unchanged.
|
||||
@@ -473,6 +509,54 @@ class AdmissionGate:
|
||||
pass
|
||||
slot.close()
|
||||
|
||||
def _check_rate_limit(self, session_id: str) -> None:
|
||||
"""Shed the turn if *session_id* is over its request-rate budget (#6215)."""
|
||||
if not self._rate_limit_enabled:
|
||||
return
|
||||
pol = self._rate_limit_policy
|
||||
now = time.monotonic()
|
||||
try:
|
||||
decision = pol.check(identity=session_id or "-", scope="ingress", now=now)
|
||||
except Exception as e: # pragma: no cover — never block on policy error
|
||||
logger.warning("AdmissionGate: rate-limit policy failed (%s); admitting", e)
|
||||
return
|
||||
finally:
|
||||
# Bounded cadence: a burst of distinct callers must not trigger a
|
||||
# full scan on every request while their windows are still live.
|
||||
if now >= self._rate_prune_after:
|
||||
self._rate_prune_after = now + _RATE_PRUNE_INTERVAL
|
||||
self._prune_rate_limit_state(pol, now)
|
||||
if decision.allowed:
|
||||
return
|
||||
self.rate_limited += 1
|
||||
retry = decision.retry_after_seconds
|
||||
wait = max(1, math.ceil(retry)) if retry is not None else None
|
||||
logger.info("AdmissionGate: rate-limited session=%s retry_after=%s", session_id or "-", wait)
|
||||
if wait is None:
|
||||
raise AdmissionRejected("Rate limit reached — please try again later.")
|
||||
raise AdmissionRejected(
|
||||
f"Rate limit reached — please try again in {wait}s.", retry_after=float(wait)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _prune_rate_limit_state(pol: object, now: float, cap: int = 4096) -> None:
|
||||
"""Drop expired per-identity entries so an open bot cannot grow them forever.
|
||||
|
||||
``SlidingWindowRateLimitPolicy`` reclaims a key only when that key is
|
||||
checked again and leaves reclamation of an unbounded identity space to
|
||||
the wrapper; inbound users are such a space.
|
||||
"""
|
||||
windows = getattr(pol, "_windows", None)
|
||||
lockouts = getattr(pol, "_lockouts", None)
|
||||
if not isinstance(windows, dict) or len(windows) + len(lockouts or {}) <= cap:
|
||||
return
|
||||
span = float(getattr(pol, "window_seconds", 0) or 0)
|
||||
for key in [k for k, (start, _) in windows.items() if now - start >= span]:
|
||||
windows.pop(key, None)
|
||||
if isinstance(lockouts, dict):
|
||||
for key in [k for k, until in lockouts.items() if now >= until]:
|
||||
lockouts.pop(key, None)
|
||||
|
||||
def _take_slot(self, scope: str) -> None:
|
||||
self._in_flight += 1
|
||||
if scope:
|
||||
@@ -596,6 +680,46 @@ def build_memory_pressure_policy(
|
||||
return MemoryPressurePolicy(soft_rss_mb=soft, hard_rss_mb=hard)
|
||||
|
||||
|
||||
_RATE_PRUNE_INTERVAL = 5.0
|
||||
|
||||
|
||||
def build_rate_limit_policy(per_user_per_minute: object) -> Optional[object]:
|
||||
"""Return a per-user request budget policy, or ``None`` when unlimited.
|
||||
|
||||
Accepts a non-negative whole number (``0``/``None`` = unlimited). Fractional,
|
||||
boolean, negative or non-numeric values raise ``ValueError`` so a typo never
|
||||
silently disables the budget the operator asked for.
|
||||
"""
|
||||
value = per_user_per_minute
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
|
||||
raise ValueError(
|
||||
f"per_user_per_minute must be a non-negative integer, got {value!r}"
|
||||
)
|
||||
try:
|
||||
if isinstance(value, float):
|
||||
if not value.is_integer():
|
||||
raise ValueError
|
||||
value = int(value)
|
||||
else:
|
||||
value = int(value)
|
||||
except (TypeError, ValueError) as err:
|
||||
raise ValueError(
|
||||
f"per_user_per_minute must be a non-negative integer, "
|
||||
f"got {per_user_per_minute!r}"
|
||||
) from err
|
||||
if value < 0:
|
||||
raise ValueError(
|
||||
f"per_user_per_minute must be >= 0, got {per_user_per_minute!r}"
|
||||
)
|
||||
if value == 0:
|
||||
return None
|
||||
from praisonaiagents.gateway import SlidingWindowRateLimitPolicy
|
||||
|
||||
return SlidingWindowRateLimitPolicy(max_requests=value, window_seconds=60.0)
|
||||
|
||||
|
||||
def build_admission_gate(
|
||||
max_concurrent_runs: int = 0,
|
||||
queue_depth: int = 0,
|
||||
@@ -603,6 +727,7 @@ def build_admission_gate(
|
||||
policy: Optional[object] = None,
|
||||
resource_policy: Optional[object] = None,
|
||||
max_concurrent_runs_per_scope: int = 0,
|
||||
per_user_per_minute: int = 0,
|
||||
) -> Optional[AdmissionGate]:
|
||||
"""Construct an :class:`AdmissionGate` from config, or return ``None``.
|
||||
|
||||
@@ -620,7 +745,13 @@ def build_admission_gate(
|
||||
gate so it queues under soft RSS pressure and sheds under hard pressure. A
|
||||
gate is built when *either* a concurrency ceiling *or* a resource policy is
|
||||
configured, so a memory-only deployment still gets backpressure.
|
||||
|
||||
``per_user_per_minute`` (Issue #6215) caps how many turns one identity may
|
||||
start per minute; ``0`` (the default) leaves ingress unlimited. Over the cap
|
||||
the turn is refused with a ``retry_after`` hint before any work starts.
|
||||
"""
|
||||
# Raises ValueError on a negative / non-integer value (fail fast).
|
||||
rate_limit_policy = build_rate_limit_policy(per_user_per_minute)
|
||||
if policy is None:
|
||||
# Only an explicit ``0`` (or ``None``) disables admission control. A
|
||||
# *negative* ceiling is a misconfiguration, not "disabled": fall through
|
||||
@@ -634,8 +765,12 @@ def build_admission_gate(
|
||||
if no_ceiling:
|
||||
# No concurrency ceiling: build a resource-only gate when a
|
||||
# resource policy is supplied, else no gate (legacy behaviour).
|
||||
if resource_policy is not None:
|
||||
return AdmissionGate(None, resource_policy=resource_policy)
|
||||
if resource_policy is not None or rate_limit_policy is not None:
|
||||
return AdmissionGate(
|
||||
None,
|
||||
resource_policy=resource_policy,
|
||||
rate_limit_policy=rate_limit_policy,
|
||||
)
|
||||
return None
|
||||
try:
|
||||
from praisonaiagents.gateway import ConcurrencyLimitPolicy
|
||||
@@ -647,4 +782,6 @@ def build_admission_gate(
|
||||
overflow_policy=overflow_policy,
|
||||
max_concurrent_runs_per_scope=max_concurrent_runs_per_scope,
|
||||
)
|
||||
return AdmissionGate(policy, resource_policy=resource_policy)
|
||||
return AdmissionGate(
|
||||
policy, resource_policy=resource_policy, rate_limit_policy=rate_limit_policy
|
||||
)
|
||||
@@ -912,6 +912,8 @@ class GatewayServerSchema(BaseModel):
|
||||
max_concurrent_runs: Optional[int] = Field(None, ge=0)
|
||||
queue_depth: Optional[int] = Field(None, ge=0)
|
||||
overflow_policy: Optional[str] = None
|
||||
# Per-user inbound request budget (#6215). 0 = unlimited (default).
|
||||
per_user_per_minute: Optional[int] = Field(None, ge=0)
|
||||
preauth_max_connections_per_ip: Optional[int] = Field(None, ge=0)
|
||||
max_unauthorized_frames: Optional[int] = Field(None, ge=0)
|
||||
# Graceful-drain windows (#2375 / #2533)
|
||||
|
||||
@@ -106,6 +106,49 @@ def _coerce_drain_timeout(value: Optional[Any]) -> Optional[float]:
|
||||
return timeout
|
||||
|
||||
|
||||
# Issue #6210: keys ``BotOS.from_config`` actually reads. Anything else in a
|
||||
# BotOS YAML used to be dropped silently.
|
||||
_BOTOS_TOP_KEYS = frozenset({
|
||||
"name", "agent", "platforms", "health", "supervision", "gateway",
|
||||
"drain_timeout", "max_concurrent_runs", "queue_depth", "overflow_policy",
|
||||
"max_concurrent_runs_per_scope", "max_rss_mb", "reliability",
|
||||
"per_user_per_minute",
|
||||
})
|
||||
_BOTOS_GATEWAY_KEYS = frozenset({
|
||||
"drain_timeout", "max_concurrent_runs", "queue_depth", "overflow_policy",
|
||||
"max_concurrent_runs_per_scope", "max_rss_mb", "reliability",
|
||||
"per_user_per_minute",
|
||||
})
|
||||
# Sections that only the gateway runtime (``praisonai gateway start``) honours.
|
||||
_GATEWAY_ONLY_KEYS = frozenset({
|
||||
"channels", "agents", "routing", "schedules", "rooms", "hooks",
|
||||
"access_groups", "mcp_servers", "identity",
|
||||
})
|
||||
|
||||
|
||||
def _check_botos_config_shape(raw: Dict[str, Any], path: str) -> None:
|
||||
"""Fail loudly on gateway-only YAML and warn on keys BotOS will not read."""
|
||||
gateway_only = sorted(k for k in raw if k in _GATEWAY_ONLY_KEYS)
|
||||
if gateway_only:
|
||||
raise ValueError(
|
||||
f"{path} is a gateway config (sections: {', '.join(gateway_only)}) "
|
||||
"which BotOS.from_config() cannot honour. Run it the same way as "
|
||||
"`praisonai gateway start --config`: "
|
||||
"asyncio.run(WebSocketGateway().start_with_config(path)) "
|
||||
"(from praisonai_bot import WebSocketGateway)."
|
||||
)
|
||||
ignored = [str(k) for k in raw if k not in _BOTOS_TOP_KEYS]
|
||||
gateway_cfg = raw.get("gateway")
|
||||
if isinstance(gateway_cfg, dict):
|
||||
ignored += [f"gateway.{k}" for k in gateway_cfg if k not in _BOTOS_GATEWAY_KEYS]
|
||||
if ignored:
|
||||
logger.warning(
|
||||
"BotOS.from_config(%s): ignoring unsupported keys %s "
|
||||
"(use WebSocketGateway().start_with_config() for full gateway configs)",
|
||||
path, ", ".join(sorted(ignored)),
|
||||
)
|
||||
|
||||
|
||||
class BotOS:
|
||||
"""Multi-platform bot orchestrator.
|
||||
|
||||
@@ -150,6 +193,7 @@ class BotOS:
|
||||
admission_policy: Optional[Any] = None,
|
||||
max_rss_mb: float = 0.0,
|
||||
reliability: Optional[str] = None,
|
||||
per_user_per_minute: int = 0,
|
||||
turn_lock: Optional[Any] = None,
|
||||
redis_config: Optional[Any] = None,
|
||||
degraded_registry: Optional[Any] = None,
|
||||
@@ -224,6 +268,8 @@ class BotOS:
|
||||
max_concurrent_runs_per_scope=max_concurrent_runs_per_scope,
|
||||
policy=admission_policy,
|
||||
resource_policy=build_memory_pressure_policy(max_rss_mb),
|
||||
# Issue #6215: per-user inbound request budget (0 = unlimited).
|
||||
per_user_per_minute=per_user_per_minute,
|
||||
)
|
||||
self._on_quiesce = None # optional callable(): host-suspend driver
|
||||
# W1: shared identity resolver applied to every managed bot —
|
||||
@@ -1600,11 +1646,22 @@ class BotOS:
|
||||
startup_grace: 60
|
||||
max_restarts_per_hour: 10
|
||||
|
||||
This is the lightweight BotOS dialect. A full ``gateway.yaml`` (with
|
||||
``channels:``/``agents:``/``routing:``/``schedules:`` …) is *not*
|
||||
accepted here — it raises ``ValueError`` instead of being silently
|
||||
degraded. Run such a file exactly like ``praisonai gateway start``::
|
||||
|
||||
from praisonai_bot import WebSocketGateway
|
||||
asyncio.run(WebSocketGateway().start_with_config("gateway.yaml"))
|
||||
|
||||
Args:
|
||||
path: Path to YAML config file.
|
||||
|
||||
Returns:
|
||||
Configured BotOS instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If the file uses gateway-only sections.
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
@@ -1620,6 +1677,8 @@ class BotOS:
|
||||
if not raw or not isinstance(raw, dict):
|
||||
raise ValueError(f"Invalid BotOS config: {path}")
|
||||
|
||||
_check_botos_config_shape(raw, path)
|
||||
|
||||
def _resolve_env(val):
|
||||
if isinstance(val, str):
|
||||
return re.sub(
|
||||
@@ -1788,6 +1847,11 @@ class BotOS:
|
||||
# ``reliability`` or ``gateway.reliability``; the preset composes drain
|
||||
# + admission while any explicit field above still overrides it.
|
||||
reliability = raw.get("reliability", gateway_cfg.get("reliability"))
|
||||
# Issue #6215: per-user inbound request budget. Default 0 = unlimited.
|
||||
# Validated (not truncated) by ``build_rate_limit_policy``.
|
||||
per_user_per_minute = raw.get(
|
||||
"per_user_per_minute", gateway_cfg.get("per_user_per_minute", 0)
|
||||
)
|
||||
|
||||
return cls(
|
||||
bots=bots,
|
||||
@@ -1800,6 +1864,7 @@ class BotOS:
|
||||
max_concurrent_runs_per_scope=max_concurrent_runs_per_scope,
|
||||
max_rss_mb=max_rss_mb,
|
||||
reliability=reliability,
|
||||
per_user_per_minute=per_user_per_minute,
|
||||
)
|
||||
|
||||
async def probe_all(self, timeout: float = 15.0) -> Dict[str, Any]:
|
||||
|
||||
@@ -1739,6 +1739,19 @@
|
||||
"default": null,
|
||||
"title": "Overflow Policy"
|
||||
},
|
||||
"per_user_per_minute": {
|
||||
"anyOf": [
|
||||
{
|
||||
"minimum": 0,
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"title": "Per User Per Minute"
|
||||
},
|
||||
"preauth_max_connections_per_ip": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
||||
@@ -16,6 +16,8 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
@@ -373,7 +375,13 @@ class GatewayApiEndpoints:
|
||||
if gate is not None and getattr(gate, "enabled", False):
|
||||
from ..bots._admission import AdmissionRejected
|
||||
try:
|
||||
async with gate.admit(session_id=session.session_id):
|
||||
# Issue #6215: charge a stable caller identity, only when a
|
||||
# budget is set (keeps duck-typed gates' signature intact).
|
||||
rk = (
|
||||
{"rate_key": self._rate_key_var.get() or "ip:unknown"}
|
||||
if getattr(gate, "_rate_limit_enabled", False) else {}
|
||||
)
|
||||
async with gate.admit(session_id=session.session_id, **rk):
|
||||
result = await self._gw._dispatch_agent_turn(
|
||||
agent, content, on_complete=_capture, **extra
|
||||
)
|
||||
@@ -554,6 +562,41 @@ class GatewayApiEndpoints:
|
||||
session_id = f"api:{key}:{agent_id}"
|
||||
return self._gw.create_session(agent_id=agent_id, session_id=session_id)
|
||||
|
||||
# Per-request stable identity for the per-user request budget (#6215);
|
||||
# inherited by the turn's task/stream context like ``_stream_turn_var``.
|
||||
_rate_key_var: "contextvars.ContextVar[Optional[str]]" = (
|
||||
contextvars.ContextVar("praisonai_gw_rate_key", default=None)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _rate_key(request: Any, auth_token: Optional[str] = None) -> str:
|
||||
"""Stable caller identity that a client-chosen value cannot reset.
|
||||
|
||||
A bearer token is used only when it *is* the gateway's configured
|
||||
``auth_token`` (so it was validated); any other bearer value is
|
||||
attacker-chosen and would mint a fresh budget per request, so such
|
||||
callers are charged to their peer address. Never returns ``None``, so the
|
||||
budget never falls back to the client-chosen session id.
|
||||
"""
|
||||
auth = request.headers.get("authorization", "")
|
||||
bearer = auth[7:] if auth.startswith("Bearer ") else ""
|
||||
if bearer and auth_token and hmac.compare_digest(
|
||||
bearer.encode("utf-8"), str(auth_token).encode("utf-8")
|
||||
):
|
||||
return "token:" + hashlib.sha256(bearer.encode("utf-8")).hexdigest()[:16]
|
||||
client = getattr(request, "client", None)
|
||||
host = getattr(client, "host", None)
|
||||
return f"ip:{host}" if host else "ip:unknown"
|
||||
|
||||
def _bind_rate_key(self, request: Any) -> None:
|
||||
"""Bind this request's per-user budget key for its turn (#6215).
|
||||
|
||||
Called by the auth-guarded route wrapper after auth has passed.
|
||||
"""
|
||||
cfg = getattr(self._gw, "config", None)
|
||||
token = getattr(cfg, "auth_token", None)
|
||||
self._rate_key_var.set(self._rate_key(request, token))
|
||||
|
||||
@staticmethod
|
||||
def _caller_key(request: Any) -> str:
|
||||
"""Derive a stable per-caller session key from the request.
|
||||
|
||||
@@ -9,6 +9,7 @@ from __future__ import annotations
|
||||
import os
|
||||
import socket
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
|
||||
@@ -86,19 +87,21 @@ def _process_create_time(pid: int) -> Optional[float]:
|
||||
class GatewayPIDLock:
|
||||
"""Manages PID lock file for single-instance gateway enforcement."""
|
||||
|
||||
def __init__(self, lock_dir: Optional[Path] = None, host: str = "127.0.0.1", port: int = 8765):
|
||||
def __init__(self, lock_dir: Optional[Path] = None, host: str = "127.0.0.1", port: int = 8765,
|
||||
lock_name: Optional[str] = None):
|
||||
"""Initialize PID lock manager.
|
||||
|
||||
Args:
|
||||
lock_dir: Directory for lock file (default: ~/.praisonai/)
|
||||
host: Gateway host
|
||||
port: Gateway port
|
||||
lock_name: Override the lock file name (used by the state lock)
|
||||
"""
|
||||
self.lock_dir = lock_dir or Path.home() / ".praisonai"
|
||||
# Make lock file specific to host and port to allow multiple instances
|
||||
safe_host = host.replace(":", "_").replace(".", "_")
|
||||
self.lock_file = self.lock_dir / f"gateway-{safe_host}-{port}.pid"
|
||||
self.lock_dir.mkdir(exist_ok=True)
|
||||
self.lock_file = self.lock_dir / (lock_name or f"gateway-{safe_host}-{port}.pid")
|
||||
self.lock_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def acquire_lock(self, host: str, port: int) -> bool:
|
||||
"""Acquire the PID lock for the gateway.
|
||||
@@ -110,6 +113,27 @@ class GatewayPIDLock:
|
||||
Returns:
|
||||
True if lock acquired successfully, False if conflict
|
||||
"""
|
||||
# Serialise check-then-claim across processes so two simultaneous
|
||||
# starts can never both observe "no live holder" and both win.
|
||||
with self._claim_guard():
|
||||
return self._acquire_locked(host, port)
|
||||
|
||||
@contextmanager
|
||||
def _claim_guard(self):
|
||||
"""Hold an OS advisory lock on a sidecar file for the claim (POSIX)."""
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # Windows: exclusive link() below still applies
|
||||
yield
|
||||
return
|
||||
guard = open(self.lock_file.with_name(self.lock_file.name + ".guard"), "a")
|
||||
try:
|
||||
fcntl.flock(guard, fcntl.LOCK_EX)
|
||||
yield
|
||||
finally:
|
||||
guard.close() # closing drops the flock
|
||||
|
||||
def _acquire_locked(self, host: str, port: int) -> bool:
|
||||
current_pid = os.getpid()
|
||||
|
||||
# Check if lock file exists
|
||||
@@ -120,6 +144,9 @@ class GatewayPIDLock:
|
||||
lines = lock_content.split('\n')
|
||||
if lines:
|
||||
existing_pid = int(lines[0])
|
||||
if existing_pid == current_pid:
|
||||
# Re-acquire by the same process (in-process restart).
|
||||
return True
|
||||
existing_create_time = None
|
||||
if len(lines) >= 5 and lines[4]:
|
||||
try:
|
||||
@@ -145,14 +172,29 @@ class GatewayPIDLock:
|
||||
create_time = _process_create_time(current_pid)
|
||||
create_time_line = "" if create_time is None else repr(create_time)
|
||||
lock_content = f"{current_pid}\n{host}\n{port}\n{int(time.time())}\n{create_time_line}\n"
|
||||
# Exclusive claim: write a private temp file, then hard-link it into
|
||||
# place. ``link`` fails if the lock already exists, so a racing claimant
|
||||
# can never overwrite a live holder (unlike rename/replace).
|
||||
temp_lock_file = self.lock_file.with_name(f"{self.lock_file.name}.{current_pid}.tmp")
|
||||
try:
|
||||
# Use a temporary file and atomic rename to prevent race conditions
|
||||
temp_lock_file = self.lock_file.with_suffix(".tmp")
|
||||
temp_lock_file.write_text(lock_content)
|
||||
temp_lock_file.replace(self.lock_file)
|
||||
try:
|
||||
os.link(temp_lock_file, self.lock_file)
|
||||
except FileExistsError:
|
||||
return False
|
||||
except OSError:
|
||||
# Filesystem without hard links: fall back to O_EXCL create.
|
||||
fd = os.open(self.lock_file, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644)
|
||||
with os.fdopen(fd, "w") as f:
|
||||
f.write(lock_content)
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
finally:
|
||||
try:
|
||||
temp_lock_file.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def release_lock(self) -> None:
|
||||
"""Release the PID lock."""
|
||||
@@ -253,6 +295,43 @@ class GatewayPIDLock:
|
||||
pass
|
||||
|
||||
|
||||
def resolve_state_dir() -> Path:
|
||||
"""Return the gateway's shared mutable state root (``~/.praisonai``).
|
||||
|
||||
The durable outbox, DLQ, pairing and hook stores are hard-wired under
|
||||
``~/.praisonai`` (they do not honour ``PRAISONAI_HOME``), so the home root -
|
||||
not the TCP port or ``PRAISONAI_HOME`` - is what identifies a gateway.
|
||||
"""
|
||||
return Path.home() / ".praisonai"
|
||||
|
||||
|
||||
def state_lock(state_dir: Optional[Path] = None, host: str = "127.0.0.1",
|
||||
port: int = 8765) -> GatewayPIDLock:
|
||||
"""Single-instance lock keyed on the shared state directory, regardless of port.
|
||||
|
||||
Two gateways sharing one state root would poll the same bot tokens and write
|
||||
the same SQLite stores concurrently (split-brain), so only one may own it.
|
||||
"""
|
||||
return GatewayPIDLock(lock_dir=state_dir or resolve_state_dir(), host=host,
|
||||
port=port, lock_name="gateway.state.pid")
|
||||
|
||||
|
||||
def format_state_collision_error(state_dir: Path, lock_info: Optional[dict] = None) -> str:
|
||||
"""Error shown when another live gateway already owns the shared state."""
|
||||
holder = ""
|
||||
if lock_info:
|
||||
holder = f" (PID {lock_info['pid']} on {lock_info['host']}:{lock_info['port']})"
|
||||
return "\n".join([
|
||||
f"Error: another gateway already owns the state at {state_dir}{holder}.",
|
||||
"",
|
||||
" Two gateways sharing state poll the same bot tokens and write the same",
|
||||
" queues, causing 409 Conflict errors and duplicate delivery.",
|
||||
" Stop it: praisonai gateway stop",
|
||||
" Or run an isolated gateway with its own state and bot tokens:",
|
||||
" HOME=/srv/gateway-b praisonai gateway start",
|
||||
])
|
||||
|
||||
|
||||
def format_collision_error(host: str, port: int, lock_info: Optional[dict] = None) -> str:
|
||||
"""Format a user-friendly error message for port collisions.
|
||||
|
||||
@@ -275,10 +354,11 @@ def format_collision_error(host: str, port: int, lock_info: Optional[dict] = Non
|
||||
pid = lock_info['pid']
|
||||
lines.extend([
|
||||
f" Another gateway may be running (PID {pid}).",
|
||||
f" Stop it: praisonai gateway stop",
|
||||
f" Or use a different port: GATEWAY_PORT={port + 1} praisonai gateway start",
|
||||
" Stop it: praisonai gateway stop",
|
||||
"",
|
||||
" Only ONE gateway process should poll each Telegram bot token.",
|
||||
" Only ONE gateway process may poll each bot token. To run a second,",
|
||||
" isolated gateway give it its own state, bot tokens and port:",
|
||||
f" HOME=/srv/gateway-b GATEWAY_PORT={port + 1} praisonai gateway start",
|
||||
])
|
||||
else:
|
||||
lines.extend([
|
||||
|
||||
@@ -1412,6 +1412,7 @@ class WebSocketGateway:
|
||||
|
||||
# PID lock for single-instance enforcement
|
||||
self._pid_lock: Optional[Any] = None
|
||||
self._state_lock: Optional[Any] = None
|
||||
|
||||
# Channel supervisor for resilient bot management. Share the gateway's
|
||||
# degraded-capability registry so the fleet crash-loop breaker (Issue
|
||||
@@ -1544,9 +1545,46 @@ class WebSocketGateway:
|
||||
logger.error("Gateway startup failed - another instance is running")
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
# Store PID lock for cleanup
|
||||
self._pid_lock = pid_lock
|
||||
|
||||
# When ``start_with_config`` already claimed the state root, its channel
|
||||
# bots are still polling if we fail here; only it may release that lock
|
||||
# (after channels stop), or a second gateway could claim the state.
|
||||
owns_state_lock = self._state_lock is None
|
||||
try:
|
||||
self._acquire_state_lock()
|
||||
await self._start_locked()
|
||||
except BaseException:
|
||||
# Any failure after claiming (bad config, bind error, cancel) must
|
||||
# not leave a live process holding the locks and blocking restarts.
|
||||
self._release_locks(state=owns_state_lock)
|
||||
raise
|
||||
|
||||
def _acquire_state_lock(self) -> None:
|
||||
"""Claim the shared state root (idempotent).
|
||||
|
||||
The port lock alone admits a second gateway on another port sharing the
|
||||
same state + bot tokens (split-brain), so the state root is locked too.
|
||||
"""
|
||||
if self._state_lock is not None:
|
||||
return
|
||||
from .port_utils import state_lock, resolve_state_dir, format_state_collision_error
|
||||
state_dir = resolve_state_dir()
|
||||
st_lock = state_lock(state_dir, host=self._host, port=self._port)
|
||||
if not st_lock.acquire_lock(self._host, self._port):
|
||||
logger.error("Gateway startup failed - another instance owns the shared state")
|
||||
raise RuntimeError(format_state_collision_error(state_dir, st_lock.get_lock_info()))
|
||||
self._state_lock = st_lock
|
||||
|
||||
def _release_locks(self, state: bool = True) -> None:
|
||||
"""Release the port lock and (if ``state``) the state lock (idempotent)."""
|
||||
for attr in ("_pid_lock", "_state_lock") if state else ("_pid_lock",):
|
||||
lock = getattr(self, attr, None)
|
||||
if lock:
|
||||
lock.release_lock()
|
||||
setattr(self, attr, None)
|
||||
|
||||
async def _start_locked(self) -> None:
|
||||
"""Body of :meth:`start` once the single-instance locks are held."""
|
||||
# Validate bind-aware auth configuration before starting
|
||||
from .auth import assert_external_bind_safe
|
||||
self.config.host = self._host
|
||||
@@ -2631,6 +2669,8 @@ class WebSocketGateway:
|
||||
auth_err = _check_auth(request)
|
||||
if auth_err:
|
||||
return auth_err
|
||||
# Issue #6215: bind the per-user budget key post-auth.
|
||||
self._api_endpoints._bind_rate_key(request)
|
||||
return await handler(request)
|
||||
return _wrapped
|
||||
|
||||
@@ -2768,14 +2808,8 @@ class WebSocketGateway:
|
||||
self._ledger_heartbeat_loop(), name="gateway-lifecycle-ledger"
|
||||
)
|
||||
try:
|
||||
# Lock cleanup on failure is handled by ``start()``.
|
||||
await self._server.serve()
|
||||
except Exception as e:
|
||||
# Clean up PID lock on any startup failure
|
||||
if hasattr(self, '_pid_lock') and self._pid_lock:
|
||||
self._pid_lock.release_lock()
|
||||
self._pid_lock = None
|
||||
# Re-raise the original exception
|
||||
raise
|
||||
finally:
|
||||
self._disarm_watchdog()
|
||||
await self._cancel_liveness_task()
|
||||
@@ -3993,9 +4027,7 @@ class WebSocketGateway:
|
||||
logger.debug("Failed to shut down pooled MCP client")
|
||||
self._mcp_pool.clear()
|
||||
|
||||
# Release PID lock
|
||||
if hasattr(self, '_pid_lock') and self._pid_lock:
|
||||
self._pid_lock.release_lock()
|
||||
self._release_locks()
|
||||
|
||||
logger.info("Gateway stopped")
|
||||
|
||||
@@ -5089,7 +5121,14 @@ class WebSocketGateway:
|
||||
# so enforce the shared gate here too.
|
||||
from ..bots._admission import AdmissionRejected
|
||||
try:
|
||||
async with gate.admit(session_id=session.session_id):
|
||||
# Issue #6215: stable per-user budget key (only
|
||||
# when a budget is set; duck-typed gates unchanged).
|
||||
_rk = (
|
||||
{"rate_key": self._client_rate_key(client_id)}
|
||||
if getattr(gate, "_rate_limit_enabled", False)
|
||||
else {}
|
||||
)
|
||||
async with gate.admit(session_id=session.session_id, **_rk):
|
||||
if getattr(session, "_stop_epoch", 0) != dequeued_stop_epoch:
|
||||
response = _stop_sentinel
|
||||
else:
|
||||
@@ -5894,6 +5933,18 @@ class WebSocketGateway:
|
||||
|
||||
# ── Client Management (Public API) ────────────────────────────────
|
||||
|
||||
def _client_rate_key(self, client_id: Optional[str]) -> str:
|
||||
"""Stable per-user budget key for a WebSocket turn (#6215).
|
||||
|
||||
Session ids are client-chosen, so charge the peer address instead.
|
||||
With no known peer, callers share one ``ip:unknown`` budget rather
|
||||
than falling back to the client-chosen session id.
|
||||
"""
|
||||
ws = self._clients.get(client_id) if client_id else None
|
||||
peer = getattr(ws, "client", None)
|
||||
host = getattr(peer, "host", None)
|
||||
return f"ip:{host}" if host else "ip:unknown"
|
||||
|
||||
def add_client(self, client_id: str, websocket: Any) -> None:
|
||||
"""Register a client WebSocket connection.
|
||||
|
||||
@@ -11281,6 +11332,29 @@ class WebSocketGateway:
|
||||
self._shell_routed_agents.clear()
|
||||
self._shell_routed_agents.update(candidate.get("shell", {}))
|
||||
|
||||
@staticmethod
|
||||
def _rate_policy_from_config(cfg: Any) -> Any:
|
||||
"""Build the ``gateway.per_user_per_minute`` policy (#6215); may raise."""
|
||||
from ..bots._admission import build_rate_limit_policy
|
||||
|
||||
gw_cfg = (cfg or {}).get("gateway") or {}
|
||||
return build_rate_limit_policy(gw_cfg.get("per_user_per_minute", 0))
|
||||
|
||||
def _apply_rate_policy(self, policy: Any) -> None:
|
||||
"""Install a reloaded per-user budget on the live admission gate (#6215).
|
||||
|
||||
An existing gate is updated in place (bots/API/WS already share it); a
|
||||
budget added to a gateway that had no gate builds a rate-only gate,
|
||||
which channels started afterwards pick up via ``_stamp_admission_gate``.
|
||||
"""
|
||||
gate = getattr(self, "_admission_gate", None)
|
||||
if gate is not None:
|
||||
gate.set_rate_limit_policy(policy)
|
||||
elif policy is not None:
|
||||
from ..bots._admission import AdmissionGate
|
||||
|
||||
self._admission_gate = AdmissionGate(None, rate_limit_policy=policy)
|
||||
|
||||
async def _restore_previous_config(
|
||||
self, error: str, changed_paths: Set[str]
|
||||
) -> None:
|
||||
@@ -11302,6 +11376,10 @@ class WebSocketGateway:
|
||||
)
|
||||
return
|
||||
try:
|
||||
try:
|
||||
self._apply_rate_policy(self._rate_policy_from_config(prev_cfg))
|
||||
except ValueError:
|
||||
pass # previous config was validated when it went live
|
||||
prev_agents_cfg = prev_cfg.get("agents", {})
|
||||
if prev_agents_cfg:
|
||||
provider_cfg = prev_cfg.get("provider", {})
|
||||
@@ -11431,6 +11509,13 @@ class WebSocketGateway:
|
||||
# drain window keep resolving the previous agents until the cutover
|
||||
# commits.
|
||||
report = self._validate_candidate(new_cfg)
|
||||
# Issue #6215: validate the per-user budget with the candidate so a
|
||||
# bad value rejects the reload instead of failing mid-cutover.
|
||||
try:
|
||||
rate_policy = self._rate_policy_from_config(new_cfg)
|
||||
except ValueError as e:
|
||||
report.ok = False
|
||||
report.failures.append(f"gateway.per_user_per_minute: {e}")
|
||||
if not report.ok:
|
||||
error = "; ".join(report.failures) or "candidate validation failed"
|
||||
logger.error(f"Reload rejected — {error}; keeping previous config")
|
||||
@@ -11453,6 +11538,7 @@ class WebSocketGateway:
|
||||
# channel is accepting turns (greptile #1 — never expose the
|
||||
# candidate before the cutover commits).
|
||||
self._activate_candidate(report.candidate)
|
||||
self._apply_rate_policy(rate_policy)
|
||||
|
||||
channels_cfg = new_cfg.get("channels", {})
|
||||
try:
|
||||
@@ -12181,6 +12267,8 @@ class WebSocketGateway:
|
||||
queue_depth=_queue_depth,
|
||||
overflow_policy=_overflow,
|
||||
max_concurrent_runs_per_scope=_max_runs_per_scope,
|
||||
# Issue #6215: per-user inbound request budget (0 = unlimited).
|
||||
per_user_per_minute=gw_cfg.get("per_user_per_minute", 0),
|
||||
)
|
||||
if self._admission_gate is not None:
|
||||
logger.info(
|
||||
@@ -12436,6 +12524,9 @@ class WebSocketGateway:
|
||||
except (NotImplementedError, OSError, ValueError):
|
||||
pass
|
||||
|
||||
# Claim the state root BEFORE channels start: channel setup replays the
|
||||
# inbound journal, which a rejected second gateway must never touch.
|
||||
self._acquire_state_lock()
|
||||
try:
|
||||
await _run_all()
|
||||
finally:
|
||||
@@ -12486,3 +12577,4 @@ class WebSocketGateway:
|
||||
logger.warning("Error draining websocket sessions: %s", e)
|
||||
finally:
|
||||
self._cancel_shutdown_deadline()
|
||||
self._release_locks()
|
||||
@@ -0,0 +1,125 @@
|
||||
"""A gateway's single-instance identity is its shared state, not its port (#6208)."""
|
||||
|
||||
import asyncio
|
||||
import multiprocessing
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonai_bot.gateway.port_utils import (
|
||||
format_collision_error,
|
||||
format_state_collision_error,
|
||||
resolve_state_dir,
|
||||
state_lock,
|
||||
)
|
||||
|
||||
|
||||
def _claim(d, start, done, results):
|
||||
start.wait() # release every worker at once so the claims genuinely race
|
||||
results.put(state_lock(Path(d)).acquire_lock("127.0.0.1", 8765))
|
||||
done.wait() # stay alive until all claimed, so no lock is seen as stale
|
||||
|
||||
|
||||
def _seed_foreign_holder(lock, pid):
|
||||
lock.lock_file.write_text(f"{pid}\n127.0.0.1\n8765\n0\n\n")
|
||||
|
||||
|
||||
def test_second_gateway_on_another_port_is_refused(tmp_path, monkeypatch):
|
||||
first = state_lock(tmp_path, port=8765)
|
||||
_seed_foreign_holder(first, os.getppid()) # a live process other than us
|
||||
second = state_lock(tmp_path, port=8766)
|
||||
assert second.lock_file == first.lock_file
|
||||
assert second.acquire_lock("127.0.0.1", 8766) is False
|
||||
msg = format_state_collision_error(tmp_path, second.get_lock_info())
|
||||
assert str(tmp_path) in msg and "HOME=" in msg and "PRAISONAI_HOME" not in msg
|
||||
|
||||
|
||||
def test_distinct_state_dirs_are_independent(tmp_path):
|
||||
a = state_lock(tmp_path / "a")
|
||||
_seed_foreign_holder(a, os.getppid())
|
||||
assert state_lock(tmp_path / "b").acquire_lock("127.0.0.1", 8766) is True
|
||||
|
||||
|
||||
def test_stale_state_lock_is_reclaimed(tmp_path, monkeypatch):
|
||||
lock = state_lock(tmp_path)
|
||||
_seed_foreign_holder(lock, 999999)
|
||||
monkeypatch.setattr(lock, "_is_process_running", lambda pid: False)
|
||||
assert lock.acquire_lock("127.0.0.1", 8765) is True
|
||||
|
||||
|
||||
def test_state_dir_ignores_praisonai_home(tmp_path, monkeypatch):
|
||||
# The stores live under ~/.praisonai regardless of PRAISONAI_HOME, so a
|
||||
# different PRAISONAI_HOME must not escape the lock.
|
||||
monkeypatch.setenv("PRAISONAI_HOME", str(tmp_path))
|
||||
assert resolve_state_dir() == Path.home() / ".praisonai"
|
||||
|
||||
|
||||
def test_concurrent_claims_have_exactly_one_winner(tmp_path):
|
||||
ctx = multiprocessing.get_context("spawn")
|
||||
n = 8
|
||||
start, done, results = ctx.Barrier(n), ctx.Barrier(n), ctx.Queue()
|
||||
procs = [ctx.Process(target=_claim, args=(str(tmp_path), start, done, results)) for _ in range(n)]
|
||||
for p in procs:
|
||||
p.start()
|
||||
outcomes = [results.get(timeout=60) for _ in range(n)]
|
||||
for p in procs:
|
||||
p.join(timeout=60)
|
||||
assert sum(outcomes) == 1
|
||||
|
||||
|
||||
def test_claim_never_overwrites_existing_lock(tmp_path):
|
||||
lock = state_lock(tmp_path)
|
||||
_seed_foreign_holder(lock, os.getppid())
|
||||
assert lock.acquire_lock("127.0.0.1", 8766) is False
|
||||
assert lock.lock_file.read_text().startswith(str(os.getppid()))
|
||||
|
||||
|
||||
def test_failed_startup_releases_both_locks(tmp_path, monkeypatch):
|
||||
from praisonai_bot.gateway import server as srv
|
||||
|
||||
monkeypatch.setenv("HOME", str(tmp_path)) # isolate both lock files
|
||||
gw = srv.WebSocketGateway.__new__(srv.WebSocketGateway)
|
||||
gw._is_running, gw._host, gw._port = False, "127.0.0.1", 0
|
||||
gw._pid_lock = gw._state_lock = None
|
||||
|
||||
async def boom():
|
||||
raise RuntimeError("bad config")
|
||||
|
||||
gw._start_locked = boom
|
||||
monkeypatch.setattr("praisonai_bot.gateway.port_utils.check_port_available", lambda h, p: (True, None))
|
||||
with pytest.raises(RuntimeError, match="bad config"):
|
||||
asyncio.run(gw.start())
|
||||
assert gw._pid_lock is None and gw._state_lock is None
|
||||
assert not list((tmp_path / ".praisonai").glob("*.pid"))
|
||||
|
||||
|
||||
def test_start_failure_keeps_state_lock_owned_by_start_with_config(tmp_path, monkeypatch):
|
||||
# start_with_config claims the state lock before channels start; if start()
|
||||
# then fails, channels are still polling, so start() must not drop it.
|
||||
from praisonai_bot.gateway import server as srv
|
||||
|
||||
monkeypatch.setenv("HOME", str(tmp_path))
|
||||
gw = srv.WebSocketGateway.__new__(srv.WebSocketGateway)
|
||||
gw._is_running, gw._host, gw._port = False, "127.0.0.1", 0
|
||||
gw._pid_lock = gw._state_lock = None
|
||||
gw._acquire_state_lock()
|
||||
held = gw._state_lock
|
||||
|
||||
async def boom():
|
||||
raise RuntimeError("bind error")
|
||||
|
||||
gw._start_locked = boom
|
||||
monkeypatch.setattr("praisonai_bot.gateway.port_utils.check_port_available", lambda h, p: (True, None))
|
||||
with pytest.raises(RuntimeError, match="bind error"):
|
||||
asyncio.run(gw.start())
|
||||
assert gw._pid_lock is None
|
||||
assert gw._state_lock is held and held.lock_file.exists()
|
||||
gw._release_locks() # start_with_config's finally, after channels stop
|
||||
assert gw._state_lock is None and not held.lock_file.exists()
|
||||
|
||||
|
||||
def test_collision_error_no_longer_suggests_bare_port_change():
|
||||
msg = format_collision_error("127.0.0.1", 8765, {"pid": 1, "is_running": True})
|
||||
assert "Or use a different port" not in msg
|
||||
assert "HOME=" in msg and "PRAISONAI_HOME" not in msg
|
||||
@@ -0,0 +1,240 @@
|
||||
"""Per-user inbound request budget in the admission gate (Issue #6215)."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonai_bot.bots._admission import (
|
||||
AdmissionGate,
|
||||
AdmissionRejected,
|
||||
build_admission_gate,
|
||||
)
|
||||
|
||||
|
||||
async def _turn(gate, user):
|
||||
async with gate.admit(session_id=user):
|
||||
return "ran"
|
||||
|
||||
|
||||
def test_unconfigured_means_no_gate():
|
||||
assert build_admission_gate(per_user_per_minute=0) is None
|
||||
|
||||
|
||||
def test_negative_budget_fails_fast():
|
||||
with pytest.raises(ValueError):
|
||||
build_admission_gate(per_user_per_minute=-1)
|
||||
|
||||
|
||||
def test_over_budget_user_refused_with_retry_after_others_unaffected():
|
||||
gate = build_admission_gate(per_user_per_minute=2)
|
||||
assert gate is not None and gate.enabled
|
||||
|
||||
async def run():
|
||||
assert await _turn(gate, "alice") == "ran"
|
||||
assert await _turn(gate, "alice") == "ran"
|
||||
with pytest.raises(AdmissionRejected) as exc:
|
||||
await _turn(gate, "alice")
|
||||
# A different user still has their own budget.
|
||||
assert await _turn(gate, "bob") == "ran"
|
||||
return exc.value
|
||||
|
||||
rej = asyncio.run(run())
|
||||
assert rej.retry_after is not None and 1 <= rej.retry_after <= 60
|
||||
assert "try again in" in rej.message
|
||||
stats = gate.stats()
|
||||
assert stats["rate_limited"] == 1
|
||||
assert stats["per_user_per_minute"] == 2
|
||||
assert stats["in_flight"] == 0
|
||||
|
||||
|
||||
def test_rate_limit_composes_with_concurrency_gate():
|
||||
gate = build_admission_gate(max_concurrent_runs=4, per_user_per_minute=1)
|
||||
assert gate.stats()["max_concurrent_runs"] == 4
|
||||
|
||||
async def run():
|
||||
await _turn(gate, "u")
|
||||
with pytest.raises(AdmissionRejected):
|
||||
await _turn(gate, "u")
|
||||
|
||||
asyncio.run(run())
|
||||
assert gate.stats()["in_flight"] == 0
|
||||
|
||||
|
||||
def test_stale_identities_are_pruned():
|
||||
from praisonaiagents.gateway import SlidingWindowRateLimitPolicy
|
||||
|
||||
pol = SlidingWindowRateLimitPolicy(max_requests=1, window_seconds=60)
|
||||
for i in range(10):
|
||||
pol.check(identity=f"u{i}", scope="ingress", now=0.0)
|
||||
AdmissionGate._prune_rate_limit_state(pol, now=120.0, cap=5)
|
||||
assert pol._windows == {}
|
||||
|
||||
|
||||
def test_botos_kwarg_and_yaml(tmp_path):
|
||||
try:
|
||||
from praisonai_bot.bots.botos import BotOS
|
||||
except ModuleNotFoundError: # pragma: no cover - optional deps
|
||||
pytest.skip("BotOS optional dependency not installed")
|
||||
assert BotOS(per_user_per_minute=5).admission_stats["per_user_per_minute"] == 5
|
||||
pytest.importorskip("yaml")
|
||||
cfg = tmp_path / "botos.yaml"
|
||||
cfg.write_text("gateway:\n per_user_per_minute: 7\n")
|
||||
assert BotOS.from_config(str(cfg)).admission_stats["per_user_per_minute"] == 7
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [0.5, -0.5, 1.5, "0.5", "abc", True, -1])
|
||||
def test_fractional_or_invalid_budget_rejected(bad):
|
||||
with pytest.raises(ValueError):
|
||||
build_admission_gate(per_user_per_minute=bad)
|
||||
|
||||
|
||||
def test_botos_yaml_fractional_budget_rejected(tmp_path):
|
||||
try:
|
||||
from praisonai_bot.bots.botos import BotOS
|
||||
except ModuleNotFoundError: # pragma: no cover - optional deps
|
||||
pytest.skip("BotOS optional dependency not installed")
|
||||
pytest.importorskip("yaml")
|
||||
cfg = tmp_path / "botos.yaml"
|
||||
cfg.write_text("gateway:\n per_user_per_minute: 0.5\n")
|
||||
with pytest.raises(ValueError):
|
||||
BotOS.from_config(str(cfg))
|
||||
with pytest.raises(ValueError):
|
||||
BotOS(per_user_per_minute=0.5)
|
||||
|
||||
|
||||
def test_rate_key_shared_across_new_sessions():
|
||||
"""A caller cannot reset the budget by switching conversation ids."""
|
||||
gate = build_admission_gate(per_user_per_minute=1)
|
||||
|
||||
async def run():
|
||||
async with gate.admit(session_id="conv-1", rate_key="ip:1.2.3.4"):
|
||||
pass
|
||||
with pytest.raises(AdmissionRejected):
|
||||
async with gate.admit(session_id="conv-2", rate_key="ip:1.2.3.4"):
|
||||
pass
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_prune_runs_at_bounded_interval(monkeypatch):
|
||||
gate = build_admission_gate(per_user_per_minute=5)
|
||||
calls = []
|
||||
monkeypatch.setattr(
|
||||
AdmissionGate, "_prune_rate_limit_state",
|
||||
staticmethod(lambda pol, now, cap=4096: calls.append(now)),
|
||||
)
|
||||
for i in range(50):
|
||||
gate._check_rate_limit(f"user-{i}")
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_set_rate_limit_policy_swaps_budget_in_place():
|
||||
from praisonai_bot.bots._admission import build_rate_limit_policy
|
||||
|
||||
gate = build_admission_gate(max_concurrent_runs=4)
|
||||
gate.set_rate_limit_policy(build_rate_limit_policy(1))
|
||||
assert gate.stats()["per_user_per_minute"] == 1
|
||||
|
||||
async def run():
|
||||
await _turn(gate, "u")
|
||||
with pytest.raises(AdmissionRejected):
|
||||
await _turn(gate, "u")
|
||||
|
||||
asyncio.run(run())
|
||||
gate.set_rate_limit_policy(build_rate_limit_policy(0))
|
||||
assert gate.stats()["per_user_per_minute"] == 0
|
||||
assert asyncio.run(_turn(gate, "u")) == "ran"
|
||||
|
||||
|
||||
def test_gateway_reload_applies_budget():
|
||||
try:
|
||||
from praisonai_bot.gateway.server import WebSocketGateway
|
||||
except ModuleNotFoundError: # pragma: no cover - optional deps
|
||||
pytest.skip("gateway optional dependency not installed")
|
||||
gw = WebSocketGateway.__new__(WebSocketGateway)
|
||||
gw._admission_gate = None
|
||||
gw._apply_rate_policy(
|
||||
gw._rate_policy_from_config({"gateway": {"per_user_per_minute": 3}})
|
||||
)
|
||||
assert gw._admission_gate.stats()["per_user_per_minute"] == 3
|
||||
gw._apply_rate_policy(gw._rate_policy_from_config({"gateway": {}}))
|
||||
assert gw._admission_gate.stats()["per_user_per_minute"] == 0
|
||||
with pytest.raises(ValueError):
|
||||
gw._rate_policy_from_config({"gateway": {"per_user_per_minute": 0.5}})
|
||||
|
||||
|
||||
def test_api_rate_key_ignores_session_header():
|
||||
try:
|
||||
from praisonai_bot.gateway.api_endpoints import GatewayApiEndpoints
|
||||
except ModuleNotFoundError: # pragma: no cover - optional deps
|
||||
pytest.skip("gateway optional dependency not installed")
|
||||
|
||||
class _Req:
|
||||
def __init__(self, headers, host="10.0.0.1"):
|
||||
self.headers = headers
|
||||
self.client = type("C", (), {"host": host})()
|
||||
|
||||
a = GatewayApiEndpoints._rate_key(_Req({"x-session-id": "one"}))
|
||||
b = GatewayApiEndpoints._rate_key(_Req({"x-session-id": "two"}))
|
||||
assert a == b == "ip:10.0.0.1"
|
||||
tok = GatewayApiEndpoints._rate_key(
|
||||
_Req({"authorization": "Bearer secret-token"}), "secret-token"
|
||||
)
|
||||
assert tok.startswith("token:") and "secret" not in tok
|
||||
|
||||
|
||||
def test_api_rate_key_rotating_unvalidated_bearer_shares_ip_budget():
|
||||
"""An unvalidated bearer (auth open, or wrong token) cannot mint budgets."""
|
||||
try:
|
||||
from praisonai_bot.gateway.api_endpoints import GatewayApiEndpoints
|
||||
except ModuleNotFoundError: # pragma: no cover - optional deps
|
||||
pytest.skip("gateway optional dependency not installed")
|
||||
|
||||
class _Req:
|
||||
def __init__(self, headers, host="10.0.0.1"):
|
||||
self.headers = headers
|
||||
self.client = type("C", (), {"host": host})()
|
||||
|
||||
for configured in (None, "", "real-token"):
|
||||
keys = {
|
||||
GatewayApiEndpoints._rate_key(
|
||||
_Req({"authorization": f"Bearer rotate-{i}"}), configured
|
||||
)
|
||||
for i in range(5)
|
||||
}
|
||||
assert keys == {"ip:10.0.0.1"}
|
||||
|
||||
|
||||
def test_api_rate_key_missing_peer_never_falls_back_to_session():
|
||||
try:
|
||||
from praisonai_bot.gateway.api_endpoints import GatewayApiEndpoints
|
||||
except ModuleNotFoundError: # pragma: no cover - optional deps
|
||||
pytest.skip("gateway optional dependency not installed")
|
||||
|
||||
class _Req:
|
||||
headers = {"x-session-id": "chosen"}
|
||||
client = None
|
||||
|
||||
assert GatewayApiEndpoints._rate_key(_Req()) == "ip:unknown"
|
||||
|
||||
gate = build_admission_gate(per_user_per_minute=1)
|
||||
|
||||
async def run():
|
||||
async with gate.admit(session_id="conv-1", rate_key="ip:unknown"):
|
||||
pass
|
||||
with pytest.raises(AdmissionRejected):
|
||||
async with gate.admit(session_id="conv-2", rate_key="ip:unknown"):
|
||||
pass
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_ws_rate_key_missing_peer_is_shared():
|
||||
try:
|
||||
from praisonai_bot.gateway.server import WebSocketGateway
|
||||
except ModuleNotFoundError: # pragma: no cover - optional deps
|
||||
pytest.skip("gateway optional dependency not installed")
|
||||
gw = WebSocketGateway.__new__(WebSocketGateway)
|
||||
gw._clients = {}
|
||||
assert gw._client_rate_key(None) == "ip:unknown"
|
||||
assert gw._client_rate_key("gone") == "ip:unknown"
|
||||
@@ -239,3 +239,60 @@ class TestBotOSFromConfigEnhanced:
|
||||
assert botos.list_bots() == []
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
class TestBotOSFromConfigGatewayDrift:
|
||||
"""Issue #6210: gateway.yaml must not be silently degraded by BotOS."""
|
||||
|
||||
def _write_yaml(self, data):
|
||||
f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
|
||||
yaml.dump(data, f)
|
||||
f.close()
|
||||
return f.name
|
||||
|
||||
@pytest.mark.parametrize("section", ["channels", "agents", "routing", "schedules"])
|
||||
def test_gateway_sections_raise(self, section):
|
||||
from praisonai_bot.bots.botos import BotOS
|
||||
|
||||
path = self._write_yaml({
|
||||
"platforms": {"telegram": {"token": "fake"}},
|
||||
section: {"x": {}},
|
||||
})
|
||||
try:
|
||||
with pytest.raises(ValueError, match="start_with_config"):
|
||||
BotOS.from_config(path)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_unknown_keys_warn(self, caplog):
|
||||
from praisonai_bot.bots.botos import BotOS
|
||||
|
||||
path = self._write_yaml({
|
||||
"platforms": {"telegram": {"token": "fake"}},
|
||||
"helth": {"interval": 5},
|
||||
"gateway": {"port": 9999, "drain_timeout": 5},
|
||||
})
|
||||
try:
|
||||
with caplog.at_level("WARNING"):
|
||||
botos = BotOS.from_config(path)
|
||||
assert "telegram" in botos.list_bots()
|
||||
assert "helth" in caplog.text and "gateway.port" in caplog.text
|
||||
assert "gateway.drain_timeout" not in caplog.text
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_non_string_unknown_key_warns_not_crashes(self, caplog):
|
||||
from praisonai_bot.bots.botos import BotOS
|
||||
|
||||
path = self._write_yaml({
|
||||
"platforms": {"telegram": {"token": "fake"}},
|
||||
123: "ignored",
|
||||
"zzz": 1,
|
||||
})
|
||||
try:
|
||||
with caplog.at_level("WARNING"):
|
||||
botos = BotOS.from_config(path)
|
||||
assert "telegram" in botos.list_bots()
|
||||
assert "123" in caplog.text and "zzz" in caplog.text
|
||||
finally:
|
||||
os.unlink(path)
|
||||
@@ -7,6 +7,7 @@ Provides CLI interface for running agents in autonomous loops with:
|
||||
- Iteration limits and timeouts
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
@@ -31,6 +32,8 @@ def loop_main(
|
||||
observe: bool = typer.Option(False, "--observe", "-o", help="Enable observability event logging"),
|
||||
recovery: bool = typer.Option(True, "--recovery/--no-recovery", help="Enable graduated doom loop recovery (default: on)"),
|
||||
level: Optional[str] = typer.Option(None, "--level", "-l", help="Autonomy level: suggest, auto_edit, full_auto"),
|
||||
max_budget: Optional[float] = typer.Option(None, "--max-budget", help="Hard USD spend cap for the whole run"),
|
||||
on_budget_exceeded: str = typer.Option("stop", "--on-budget-exceeded", help="When --max-budget is hit: stop or warn"),
|
||||
):
|
||||
"""Run an agent in an autonomous loop.
|
||||
|
||||
@@ -66,6 +69,8 @@ def loop_main(
|
||||
observe=observe,
|
||||
recovery=recovery,
|
||||
level=level,
|
||||
max_budget=max_budget,
|
||||
on_budget_exceeded=on_budget_exceeded,
|
||||
)
|
||||
|
||||
|
||||
@@ -80,8 +85,17 @@ def _run_loop(
|
||||
observe: bool = False,
|
||||
recovery: bool = True,
|
||||
level: Optional[str] = None,
|
||||
max_budget: Optional[float] = None,
|
||||
on_budget_exceeded: str = "stop",
|
||||
):
|
||||
"""Execute the autonomous loop."""
|
||||
if on_budget_exceeded not in ("stop", "warn"):
|
||||
typer.echo("Error: --on-budget-exceeded must be 'stop' or 'warn'", err=True)
|
||||
raise typer.Exit(2)
|
||||
# 0 would mean "no cap" downstream; negatives and NaN/inf are never a real cap.
|
||||
if max_budget is not None and not (math.isfinite(max_budget) and max_budget > 0):
|
||||
typer.echo("Error: --max-budget must be a positive number", err=True)
|
||||
raise typer.Exit(2)
|
||||
try:
|
||||
from praisonaiagents import Agent
|
||||
except ImportError:
|
||||
@@ -117,6 +131,13 @@ def _run_loop(
|
||||
|
||||
if not verbose:
|
||||
agent_kwargs["output"] = "silent"
|
||||
|
||||
# Only forward a budget when the user set one (None = no cap, zero overhead).
|
||||
if max_budget is not None:
|
||||
agent_kwargs["execution"] = {
|
||||
"max_budget": max_budget,
|
||||
"on_budget_exceeded": on_budget_exceeded,
|
||||
}
|
||||
|
||||
# Add completion promise instructions if provided
|
||||
if completion_promise:
|
||||
@@ -138,6 +159,8 @@ def _run_loop(
|
||||
typer.echo(" Context clearing: enabled")
|
||||
if timeout:
|
||||
typer.echo(f" Timeout: {timeout}s")
|
||||
if max_budget is not None:
|
||||
typer.echo(f" Max budget: ${max_budget} ({on_budget_exceeded})")
|
||||
typer.echo("")
|
||||
|
||||
# Run autonomous loop using unified API
|
||||
@@ -200,6 +223,8 @@ OPTIONS:
|
||||
-p, --completion-promise STR Promise text to signal completion
|
||||
-c, --clear-context Clear chat history between iterations
|
||||
-t, --timeout FLOAT Timeout in seconds
|
||||
--max-budget FLOAT Hard USD spend cap for the whole run
|
||||
--on-budget-exceeded STR stop (default) or warn when the cap is hit
|
||||
-m, --model STR LLM model to use
|
||||
-v, --verbose Show verbose output
|
||||
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
"""`praisonai loop --max-budget` reaches the agent's ExecutionConfig (Issue #6206)."""
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("typer")
|
||||
|
||||
import typer
|
||||
|
||||
from praisonai_code.cli.commands import loop as loop_module
|
||||
|
||||
|
||||
class _Result:
|
||||
success = True
|
||||
completion_reason = "promise"
|
||||
iterations = 1
|
||||
duration_seconds = 0.1
|
||||
started_at = None
|
||||
output = "done"
|
||||
error = None
|
||||
|
||||
|
||||
def _capture(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
class FakeAgent:
|
||||
def __init__(self, **kwargs):
|
||||
captured["agent_kwargs"] = kwargs
|
||||
|
||||
def start(self, prompt, **kwargs):
|
||||
return _Result()
|
||||
|
||||
import praisonaiagents
|
||||
monkeypatch.setattr(praisonaiagents, "Agent", FakeAgent)
|
||||
return captured
|
||||
|
||||
|
||||
def _run(**extra):
|
||||
loop_module._run_loop(
|
||||
prompt="build it", max_iterations=2, completion_promise=None,
|
||||
clear_context=False, timeout=None, model=None, verbose=False, **extra,
|
||||
)
|
||||
|
||||
|
||||
def test_max_budget_forwarded_to_execution(monkeypatch):
|
||||
captured = _capture(monkeypatch)
|
||||
_run(max_budget=0.5, on_budget_exceeded="warn")
|
||||
assert captured["agent_kwargs"]["execution"] == {
|
||||
"max_budget": 0.5, "on_budget_exceeded": "warn",
|
||||
}
|
||||
|
||||
|
||||
def test_no_budget_leaves_execution_unset(monkeypatch):
|
||||
captured = _capture(monkeypatch)
|
||||
_run()
|
||||
assert "execution" not in captured["agent_kwargs"]
|
||||
|
||||
|
||||
def test_invalid_on_budget_exceeded_rejected(monkeypatch):
|
||||
_capture(monkeypatch)
|
||||
with pytest.raises(typer.Exit):
|
||||
_run(max_budget=1.0, on_budget_exceeded="explode")
|
||||
|
||||
|
||||
def test_budget_dict_builds_real_execution_config():
|
||||
pytest.importorskip("praisonaiagents")
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="a", execution={"max_budget": 0.25, "on_budget_exceeded": "stop"})
|
||||
assert agent._max_budget == 0.25
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [0.0, -1.0, float("nan"), float("inf")])
|
||||
def test_invalid_max_budget_rejected(monkeypatch, bad):
|
||||
captured = _capture(monkeypatch)
|
||||
with pytest.raises(typer.Exit):
|
||||
_run(max_budget=bad)
|
||||
assert "agent_kwargs" not in captured
|
||||
@@ -1,6 +1,8 @@
|
||||
# PraisonAI TypeScript Node.js AI Agents Framework
|
||||
|
||||
PraisonAI is a production-ready Multi AI Agents framework, designed to create AI Agents to automate and solve problems ranging from simple tasks to complex challenges. It provides a low-code solution to streamline the building and management of multi-agent LLM systems, emphasising simplicity, customisation, and effective human-agent collaboration.
|
||||
PraisonAI runs AI agents on your own computers, even fully offline, so private data never has to leave the building. It is built for highly regulated teams, and free and open source for everyone.
|
||||
|
||||
This package is the TypeScript and JavaScript SDK of the open-source PraisonAI agent framework. Build single agents, multi-agent workflows, tools and memory in Node.js, on hosted models or on local models such as Ollama. Data stays on your own machines only when the model runs locally too; with a hosted model such as the default OpenAI one, prompts are sent to that provider. The Python framework and the full docs are at [github.com/MervinPraison/PraisonAI](https://github.com/MervinPraison/PraisonAI) and [praison.ai/docs](https://praison.ai/docs).
|
||||
|
||||
## Installation
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "praisonai",
|
||||
"version": "1.7.4",
|
||||
"description": "PraisonAI TypeScript AI Agents Framework - Node.js, npm, and Javascript AI Agents Framework",
|
||||
"description": "Open-source TypeScript and JavaScript AI agent framework for Node.js, on hosted or local models. The JavaScript SDK of PraisonAI.",
|
||||
"main": "dist/index.js",
|
||||
"module": "dist/esm/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -10,6 +10,7 @@ from typing import Optional, Dict, Any
|
||||
|
||||
from ..diff.diff_strategy import apply_search_replace_diff
|
||||
from ..utils.file_utils import (
|
||||
atomic_write,
|
||||
file_exists,
|
||||
is_path_within_directory,
|
||||
)
|
||||
@@ -142,8 +143,7 @@ def apply_diff(
|
||||
f.write(original_content)
|
||||
|
||||
# Write the modified content
|
||||
with open(abs_path, 'w', encoding=encoding) as f:
|
||||
f.write(result.content)
|
||||
atomic_write(abs_path, result.content, encoding=encoding)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
|
||||
@@ -10,6 +10,7 @@ import re
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
from ..utils.file_utils import (
|
||||
atomic_write,
|
||||
file_exists,
|
||||
is_path_within_directory,
|
||||
)
|
||||
@@ -191,8 +192,7 @@ def search_replace(
|
||||
f.write(original_content)
|
||||
|
||||
# Write the modified content
|
||||
with open(abs_path, 'w', encoding=encoding) as f:
|
||||
f.write(content)
|
||||
atomic_write(abs_path, content, encoding=encoding)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
|
||||
@@ -8,6 +8,7 @@ import os
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
from ..utils.file_utils import (
|
||||
atomic_write,
|
||||
file_exists,
|
||||
create_directories_for_file,
|
||||
is_path_within_directory,
|
||||
@@ -131,8 +132,7 @@ def write_file(
|
||||
pass # If we can't read, just use the content as-is
|
||||
|
||||
# Write the file
|
||||
with open(abs_path, 'w', encoding=encoding) as f:
|
||||
f.write(processed_content)
|
||||
atomic_write(abs_path, processed_content, encoding=encoding)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
|
||||
@@ -8,8 +8,10 @@ Provides utilities for file operations including:
|
||||
- Line ending normalization
|
||||
"""
|
||||
|
||||
import errno
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from typing import Optional, Tuple
|
||||
|
||||
|
||||
@@ -305,3 +307,68 @@ def count_file_lines(file_path: str) -> int:
|
||||
for _ in f:
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def atomic_write(path: str, content: str, *, encoding: Optional[str] = "utf-8") -> None:
|
||||
"""Write ``content`` to ``path`` atomically (temp file + fsync + os.replace).
|
||||
|
||||
``open(path, 'w')`` truncates the target immediately, so any interrupt
|
||||
(cancellation, SIGTERM, disk-full) between open and write leaves the file
|
||||
zero-byte or half-written and the original content lost. The code editing
|
||||
tools run on LLM-driven tool calls that regularly hit timeouts/cancellations,
|
||||
so they need the same crash-safe write the wrapper already uses for configs.
|
||||
"""
|
||||
# Resolve symlinks so we replace the link target (like open(path, 'w')
|
||||
# writes through the link) instead of turning the symlink into a plain file.
|
||||
path = os.path.realpath(path)
|
||||
target_dir = os.path.dirname(path) or "."
|
||||
# Atomic replace swaps in a fresh inode, so carry over the metadata that
|
||||
# matters for a replaced file: permission bits and ownership (uid/gid).
|
||||
# Hard-link identity and xattrs/ACLs cannot be preserved through an inode
|
||||
# swap and are out of scope for this tool.
|
||||
try:
|
||||
existing_stat = os.stat(path)
|
||||
except OSError:
|
||||
existing_stat = None
|
||||
# An inode swap would bypass a read-only target's permissions, so keep the
|
||||
# old open(path, 'w') behaviour and refuse to replace a file we can't write.
|
||||
if existing_stat is not None and not os.access(path, os.W_OK):
|
||||
raise PermissionError(errno.EACCES, os.strerror(errno.EACCES), path)
|
||||
# Create the temp file with the target's permission bits (setuid/setgid
|
||||
# dropped) so private files (e.g. 0600 .env) are never readable by others
|
||||
# while the new content is written. New files use 0o666 so the kernel
|
||||
# applies the umask, like open(path, 'w'). Never call os.umask(): it is
|
||||
# process-wide and would race with other threads.
|
||||
mode = (existing_stat.st_mode & 0o777) if existing_stat is not None else 0o666
|
||||
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_BINARY", 0)
|
||||
while True:
|
||||
tmp_path = os.path.join(target_dir, f".tmp_{uuid.uuid4().hex}")
|
||||
try:
|
||||
fd = os.open(tmp_path, flags, mode)
|
||||
break
|
||||
except FileExistsError:
|
||||
continue
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding=encoding) as f:
|
||||
f.write(content)
|
||||
f.flush()
|
||||
os.fsync(f.fileno()) # durability before rename
|
||||
if existing_stat is not None:
|
||||
# Restore bits the umask may have removed; never widens beyond the
|
||||
# target. If this fails the temp file is at most as open as the target.
|
||||
try:
|
||||
os.chmod(tmp_path, mode)
|
||||
except OSError:
|
||||
pass # best-effort
|
||||
if hasattr(os, "chown"):
|
||||
try:
|
||||
os.chown(tmp_path, existing_stat.st_uid, existing_stat.st_gid)
|
||||
except (OSError, AttributeError):
|
||||
pass # best-effort (needs privilege; no-op on Windows)
|
||||
os.replace(tmp_path, path) # atomic on POSIX + Windows
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
@@ -8,55 +8,12 @@ without multiple file read/write cycles.
|
||||
import contextlib
|
||||
import difflib
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from praisonai.code.utils.file_utils import atomic_write as _atomic_write
|
||||
from praisonai.code.utils.file_utils import is_path_within_directory
|
||||
|
||||
|
||||
def _atomic_write(path: str, content: str) -> None:
|
||||
"""Write ``content`` to ``path`` atomically (temp file + fsync + os.replace).
|
||||
|
||||
``open(path, 'w')`` truncates the target immediately, so any interrupt
|
||||
(cancellation, SIGTERM, disk-full) between open and write leaves the file
|
||||
zero-byte or half-written and the original content lost. multiedit runs on
|
||||
LLM-driven tool calls that regularly hit timeouts/cancellations, so it needs
|
||||
the same crash-safe write the rest of the wrapper already uses for configs.
|
||||
"""
|
||||
target_dir = os.path.dirname(os.path.abspath(path)) or "."
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".tmp_", dir=target_dir)
|
||||
try:
|
||||
# Atomic replace swaps in a fresh inode, so carry over the metadata that
|
||||
# matters for a replaced file: mode (mkstemp creates 0600 by default) and
|
||||
# ownership (uid/gid) so editing a group-/shared-owned file does not
|
||||
# silently change its owner. Hard-link identity and xattrs/ACLs cannot be
|
||||
# preserved through an inode swap and are out of scope for this tool.
|
||||
try:
|
||||
existing_stat = os.stat(path)
|
||||
except OSError:
|
||||
existing_stat = None
|
||||
with os.fdopen(fd, "w") as f:
|
||||
f.write(content)
|
||||
f.flush()
|
||||
os.fsync(f.fileno()) # durability before rename
|
||||
if existing_stat is not None:
|
||||
try:
|
||||
os.chmod(tmp_path, existing_stat.st_mode)
|
||||
except OSError:
|
||||
pass # best-effort
|
||||
if hasattr(os, "chown"):
|
||||
try:
|
||||
os.chown(tmp_path, existing_stat.st_uid, existing_stat.st_gid)
|
||||
except (OSError, AttributeError):
|
||||
pass # best-effort (needs privilege; no-op on Windows)
|
||||
os.replace(tmp_path, path) # atomic on POSIX + Windows
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _file_lock(path: str):
|
||||
@@ -250,7 +207,7 @@ def _run_edits(
|
||||
|
||||
# Write file atomically if not dry run (crash-safe: no zero-byte truncation)
|
||||
if not dry_run and result["edits_applied"] > 0:
|
||||
_atomic_write(filepath, content)
|
||||
_atomic_write(filepath, content, encoding=None) # match locale-default read above
|
||||
|
||||
result["success"] = result["edits_failed"] == 0
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ Tests for praisonai.code.tools module.
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import pytest
|
||||
|
||||
|
||||
class TestReadFile:
|
||||
@@ -286,7 +287,129 @@ class TestSearchReplace:
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
assert f.read() == "Hello, new world!"
|
||||
|
||||
|
||||
def test_interrupted_write_keeps_original(self, monkeypatch):
|
||||
"""A failure mid-write must not truncate the original file."""
|
||||
from praisonai.code.tools.search_replace import search_replace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "test.txt")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("Hello, old world!")
|
||||
|
||||
def boom(*args, **kwargs):
|
||||
raise KeyboardInterrupt
|
||||
|
||||
monkeypatch.setattr(os, "fsync", boom)
|
||||
try:
|
||||
search_replace(file_path, [
|
||||
{'search': 'old', 'replace': 'new'}
|
||||
], workspace=temp_dir)
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
assert f.read() == "Hello, old world!"
|
||||
assert os.listdir(temp_dir) == ["test.txt"]
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="POSIX permissions")
|
||||
def test_atomic_write_new_file_respects_umask(self):
|
||||
"""New files get 0666 & ~umask like open(path, 'w'), not mkstemp's 0600."""
|
||||
from praisonai.code.utils.file_utils import atomic_write
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
old = os.umask(0o022)
|
||||
try:
|
||||
path = os.path.join(temp_dir, "new.txt")
|
||||
atomic_write(path, "x")
|
||||
finally:
|
||||
os.umask(old)
|
||||
assert (os.stat(path).st_mode & 0o777) == 0o644
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="POSIX permissions")
|
||||
def test_atomic_write_does_not_touch_process_umask(self, monkeypatch):
|
||||
"""os.umask is process-wide; changing it races with other threads."""
|
||||
from praisonai.code.utils.file_utils import atomic_write
|
||||
|
||||
def no_umask(*args, **kwargs):
|
||||
raise AssertionError("atomic_write must not call os.umask")
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
monkeypatch.setattr(os, "umask", no_umask)
|
||||
new_path = os.path.join(temp_dir, "new.txt")
|
||||
atomic_write(new_path, "a")
|
||||
atomic_write(new_path, "b")
|
||||
with open(new_path) as f:
|
||||
assert f.read() == "b"
|
||||
assert os.listdir(temp_dir) == ["new.txt"]
|
||||
|
||||
@pytest.mark.skipif(not hasattr(os, "symlink"), reason="symlinks unsupported")
|
||||
def test_atomic_write_writes_through_symlink(self):
|
||||
"""Like open(path, 'w'), editing a symlink must update its target and keep the link."""
|
||||
from praisonai.code.utils.file_utils import atomic_write
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
real = os.path.join(temp_dir, "real.txt")
|
||||
link = os.path.join(temp_dir, "link.txt")
|
||||
with open(real, "w") as f:
|
||||
f.write("orig")
|
||||
try:
|
||||
os.symlink(real, link)
|
||||
except OSError:
|
||||
pytest.skip("cannot create symlink")
|
||||
atomic_write(link, "new")
|
||||
assert os.path.islink(link)
|
||||
with open(real) as f:
|
||||
assert f.read() == "new"
|
||||
assert sorted(os.listdir(temp_dir)) == ["link.txt", "real.txt"]
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="POSIX permission bits")
|
||||
def test_atomic_write_temp_file_never_wider_than_private_target(self, monkeypatch):
|
||||
"""A 0600 target's new content must never sit in a more readable temp file."""
|
||||
import stat
|
||||
from praisonai.code.utils import file_utils
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
path = os.path.join(temp_dir, ".env")
|
||||
with open(path, "w") as f:
|
||||
f.write("OLD=1")
|
||||
os.chmod(path, 0o600)
|
||||
|
||||
seen = []
|
||||
real_fdopen = os.fdopen
|
||||
|
||||
def spy_fdopen(fd, *a, **kw):
|
||||
seen.append(stat.S_IMODE(os.fstat(fd).st_mode))
|
||||
return real_fdopen(fd, *a, **kw)
|
||||
|
||||
monkeypatch.setattr(file_utils.os, "fdopen", spy_fdopen)
|
||||
file_utils.atomic_write(path, "NEW=2")
|
||||
|
||||
assert seen and all(m & 0o077 == 0 for m in seen), seen
|
||||
assert stat.S_IMODE(os.stat(path).st_mode) == 0o600
|
||||
with open(path) as f:
|
||||
assert f.read() == "NEW=2"
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt" or (hasattr(os, "geteuid") and os.geteuid() == 0),
|
||||
reason="POSIX permissions; root bypasses them")
|
||||
def test_atomic_write_refuses_read_only_target(self):
|
||||
"""Like open(path, 'w'), a read-only file must not be replaced via inode swap."""
|
||||
from praisonai.code.utils.file_utils import atomic_write
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
path = os.path.join(temp_dir, "locked.txt")
|
||||
with open(path, "w") as f:
|
||||
f.write("orig")
|
||||
os.chmod(path, 0o444)
|
||||
try:
|
||||
with pytest.raises(PermissionError):
|
||||
atomic_write(path, "new")
|
||||
with open(path) as f:
|
||||
assert f.read() == "orig"
|
||||
assert os.listdir(temp_dir) == ["locked.txt"]
|
||||
finally:
|
||||
os.chmod(path, 0o644)
|
||||
|
||||
def test_multiple_operations(self):
|
||||
"""Test multiple search/replace operations."""
|
||||
from praisonai.code.tools.search_replace import search_replace
|
||||
|
||||
Reference in new issue
Block a user