Compare commits

...
Author SHA1 Message Date
MervinPraisonandClaude Opus 5.5 eb6c355a89 docs(ts): 'never has to leave' to match README and docs
Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-11-04 12:43:53 +00:00
MervinPraisonandClaude Opus 5.5 698a44260c docs(ts): match the README's main line; link praison.ai/docs
Keeps the hosted-model privacy qualifier from the review fix.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-11-04 12:42:45 +00:00
MervinPraisonandClaude Opus 5.5 40a3bae425 fix(ts): qualify README privacy claim for hosted models
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-04 11:51:50 +00:00
MervinPraisonandClaude Opus 5.5 379a25d1b6 docs(ts): lead the npm README and description with the new focus
Text only: the README intro paragraph and the package.json description.
No code, exports or keywords change.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-11-04 11:30:38 +00:00
4e419c4f05 fix: keep tool turns intact across auto-save, compaction and history mutators (#6218)
* fix: keep tool turns intact across auto-save, compaction and history mutators (fixes #6217)

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix: save messages appended by history replacement and serialize it with auto-save (#6217)

Co-Authored-By: Claude Opus 5.5 <[email protected]>

---------

Co-authored-by: MervinPraison <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-11-04 11:20:43 +01:00
7ae123e68d fix: deprecate You.com and Jira connectors in core in favour of PraisonAI-Tools (fixes #6168) (#6169)
Co-authored-by: MervinPraison <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-11-04 11:20:36 +01:00
Wu Shuwen f8b1f3ae6c fix: surface weekly package test failures (#6189) 2026-11-04 11:17:10 +01:00
Wu Shuwen 8f1d50b5ff test: include four core unit regression files in shard (#6187)
* test: gate agent token budget unit tests

* test: include more core unit regression files
2026-11-04 11:16:59 +01:00
praisonai-triage-agent[bot] ac94b9f8b7 Merge pull request #6216 from MervinPraison/claude/issue-6215-20261103
fix: optional per-user inbound request budget at bot/gateway admission
2026-11-03 11:56:00 +00:00
MervinPraisonandClaude Opus 5.5 d01c37d434 fix: key API rate budget only on a validated bearer, never fall back to client-chosen session id (#6215)
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 11:42:45 +00:00
praisonai-triage-agent[bot] feae4f936d Merge pull request #6212 from MervinPraison/claude/issue-6208-20261103
fix: gateway single-instance guard is state-scoped, not only port-scoped
2026-11-03 11:32:32 +00:00
MervinPraisonandClaude Opus 5.5 831490316f fix: start() keeps the state lock owned by start_with_config until channels stop; make concurrent-claim test deterministic
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 11:18:32 +00:00
MervinPraisonandClaude Opus 5.5 d20762ea11 fix: stable per-caller rate key, reload-applied and validated per-user budget, bounded pruning (#6215)
- API/WebSocket turns charge the budget to a stable caller identity
  (bearer-token digest or peer IP) so a new session id cannot reset it.
- gateway.per_user_per_minute changes now take effect on config reload
  (validated before cutover; restored on rollback).
- Fractional/boolean/negative budgets are rejected instead of truncated.
- Stale-identity pruning runs at most every 5s instead of on every request.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 11:05:38 +00:00
MervinPraisonandClaude Opus 5.5 8528cf3da5 fix: enforce optional per-user inbound request budget at bot/gateway admission (fixes #6215)
Wires the existing core SlidingWindowRateLimitPolicy into the shared
AdmissionGate so a single user cannot start unbounded turns. Configured
via gateway.per_user_per_minute (YAML) or BotOS(per_user_per_minute=...);
default 0 keeps ingress unlimited.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 10:56:28 +00:00
praisonai-triage-agent[bot] f16a6971bf Merge pull request #6211 from MervinPraison/claude/issue-6210-20261103
fix: BotOS.from_config rejects gateway.yaml sections instead of silently dropping them
2026-11-03 10:45:52 +00:00
MervinPraisonandClaude Opus 5.5 de3a1df83b fix: make gateway state lock an exclusive claim taken before channel recovery and released on any startup failure
- port_utils: serialise claim with flock + exclusive link()/O_EXCL so racing starts cannot both win
- key state lock on ~/.praisonai (where the stores actually live), not PRAISONAI_HOME
- server: acquire state lock in start_with_config before start_channels (journal replay)
- server: release both locks on any failure after claim, not only serve() errors

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 10:38:46 +00:00
MervinPraisonandClaude Opus 5.5 0b4cb4db60 fix: BotOS.from_config warns instead of crashing on non-string unknown YAML keys
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 10:32:13 +00:00
MervinPraisonandClaude Opus 5.5 2c8dd09aa5 fix: lock gateway single-instance guard on shared state dir, not just host:port (fixes #6208)
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 10:29:53 +00:00
MervinPraisonandClaude Opus 5.5 25729bec08 fix: BotOS.from_config rejects gateway-only YAML instead of silently dropping it (fixes #6210)
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 10:28:34 +00:00
praisonai-triage-agent[bot] 688bb55d51 Merge pull request #6205 from MervinPraison/claude/issue-6204-20261103
fix: crash-safe writes in praisonai.code edit tools
2026-11-03 09:56:01 +00:00
MervinPraisonandClaude Opus 5.5 60495e8d26 fix: atomic_write refuses read-only targets like open(path, 'w') instead of replacing them via inode swap
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 09:41:23 +00:00
MervinPraisonandClaude Opus 5.5 dfa448d57b fix: create atomic_write temp file with target's permission bits so private files are never exposed mid-write
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 09:39:26 +00:00
praisonai-triage-agent[bot] b5acce53be Merge pull request #6207 from MervinPraison/claude/issue-6206-20261103
fix: expose --max-budget on praisonai loop and pass YAML execution block through
2026-11-03 09:22:23 +00:00
MervinPraisonandClaude Opus 5.5 e128c541d5 fix: surface budget stops in sync chat, keep execution in roles YAML, validate --max-budget
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 09:07:53 +00:00
MervinPraisonandClaude Opus 5.5 91f0cc7727 fix: expose max-budget on praisonai loop and pass YAML execution block through (fixes #6206)
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 08:58:57 +00:00
MervinPraisonandClaude Opus 5.5 41acf94def fix: atomic_write writes through symlinks instead of replacing the link
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 08:53:47 +00:00
MervinPraisonandClaude Opus 5.5 40e04fdf44 fix: create atomic_write temp file with 0o666 instead of reading process-wide umask
os.umask() is process-wide, so the read-and-restore window could make files
created by other threads world-writable. The temp file is now created via
os.open(O_CREAT|O_EXCL, 0o666) so the kernel applies the current umask.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 08:51:47 +00:00
MervinPraisonandClaude Opus 5.5 d21d60600f fix: keep umask-based mode for new files and locale encoding in multiedit atomic writes
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 08:17:07 +00:00
MervinPraisonandClaude Opus 5.5 0b04d2f017 fix: write praisonai.code edit tools atomically so interrupts keep the original file (fixes #6204)
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-03 08:11:56 +00:00
32 changed files with 1533 additions and 110 deletions

No files matched your search

+1 -1
View File
@@ -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
+3 -1
View File
@@ -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([
+105 -13
View File
@@ -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
+3 -1
View File
@@ -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 -1
View File
@@ -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
+2 -45
View File
@@ -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
+124 -1
View File
@@ -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