mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-09 17:31:52 -10:00
Compare commits
78
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()
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.22" \
|
||||
"praisonai>=4.7.23" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.22" \
|
||||
"praisonai>=4.7.23" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.22" \
|
||||
"praisonai>=4.7.23" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -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"
|
||||
@@ -7099,8 +7101,12 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
model_name=str(getattr(self, 'llm', '') or ''),
|
||||
),
|
||||
)
|
||||
# Copy ContextVars into the worker so the re-scheduled model call
|
||||
# keeps the caller's context (session identity, AG-UI run owner).
|
||||
import contextvars
|
||||
_ctx = contextvars.copy_context()
|
||||
response = await loop.run_in_executor(
|
||||
None, lambda: manager.execute_model_call(request, _model_fn)
|
||||
None, lambda: _ctx.run(manager.execute_model_call, request, _model_fn)
|
||||
)
|
||||
return self._unwrap_model_response(response)
|
||||
|
||||
|
||||
@@ -22,12 +22,6 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
from typing import List, Optional, Any, Dict, Union, Generator, TYPE_CHECKING
|
||||
|
||||
# Shared HTTP launch state — the process-wide registry is defined in
|
||||
# agent/server_registry.py and shared with PraisonAIAgents.launch() so a
|
||||
# standalone Agent and a team launched on the same port cannot double-bind it or
|
||||
# hide each other's endpoints.
|
||||
from .server_registry import _server_registry
|
||||
|
||||
|
||||
class _TimeoutCancelToken:
|
||||
"""cancel_token for caller-mode ``start(timeout=...)``.
|
||||
@@ -76,6 +70,8 @@ def cleanup_launch_registration(agent_id: str) -> None:
|
||||
``_AgentServerRegistry.teardown_routes_for``) so cleanup cannot race a
|
||||
concurrent launch that reserves and registers a route on the same port.
|
||||
"""
|
||||
from .server_registry import _server_registry # lazy: launch path only
|
||||
|
||||
_server_registry.teardown_routes_for(agent_id)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -2394,7 +2390,11 @@ Write the complete compiled report:"""
|
||||
host = resolve_launch_host(host)
|
||||
|
||||
# Get or create the shared FastAPI app for this port (shared with
|
||||
# PraisonAIAgents.launch via the same _server_registry singleton).
|
||||
# PraisonAIAgents.launch via the same _server_registry singleton, so a
|
||||
# standalone Agent and a team on one port cannot double-bind it).
|
||||
# Imported lazily: only the HTTP launch path needs the registry.
|
||||
from .server_registry import _server_registry
|
||||
|
||||
app, is_new_app = _server_registry.get_or_create_app(
|
||||
port,
|
||||
title=f"PraisonAI Agents API (Port {port})",
|
||||
|
||||
@@ -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)
|
||||
@@ -1228,34 +1285,18 @@ class MemoryMixin:
|
||||
return False
|
||||
|
||||
try:
|
||||
import os
|
||||
from praisonaiagents.tools.path_safety import resolve_within_root
|
||||
from ..utils.output_file import write_output_file
|
||||
|
||||
project_root = os.environ.get("PRAISONAI_PROJECT_ROOT") or os.getcwd()
|
||||
file_path = resolve_within_root(self._output_file, project_root)
|
||||
# Shared hardened writer (root containment, O_NOFOLLOW, hardlink-safe).
|
||||
file_path = write_output_file(self._output_file, str(content))
|
||||
if file_path is None:
|
||||
logging.warning(
|
||||
"Output file %r is outside project root %r; skipping save",
|
||||
self._output_file,
|
||||
project_root,
|
||||
)
|
||||
print(f"⚠️ Output path outside project root: {self._output_file}")
|
||||
return False
|
||||
|
||||
# Create parent directories if they don't exist
|
||||
parent_dir = os.path.dirname(file_path)
|
||||
if parent_dir and not os.path.exists(parent_dir):
|
||||
os.makedirs(parent_dir, exist_ok=True)
|
||||
|
||||
# Write content to file
|
||||
with open(file_path, 'w', encoding='utf-8') as f:
|
||||
f.write(str(content))
|
||||
|
||||
# Print success message to terminal
|
||||
|
||||
print(f"✅ Output saved to {file_path}")
|
||||
logging.debug(f"Output saved to file: {file_path}")
|
||||
return True
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to save output to file '{self._output_file}': {e}")
|
||||
print(f"⚠️ Failed to save output to {self._output_file}: {e}")
|
||||
|
||||
@@ -121,6 +121,12 @@ class PlatformCapabilities:
|
||||
anchoring). When False, ``create_thread`` returns a typed
|
||||
``unsupported`` outcome and callers fall back to the parent channel
|
||||
rather than raising.
|
||||
supports_status_text: Whether the adapter can render a short,
|
||||
replaceable textual status line (e.g. "🔍 Searching the web...")
|
||||
without editing a message or adding reactions. When True the
|
||||
adapter SHOULD implement ``async set_status_text(channel_id, text)``
|
||||
(an empty ``text`` clears the line when the turn ends); streaming ``auto``/``progress`` on a non-edit channel then routes the
|
||||
curated activity line there instead of degrading to ``off``.
|
||||
"""
|
||||
|
||||
max_message_length: int = 4096
|
||||
@@ -142,6 +148,7 @@ class PlatformCapabilities:
|
||||
# Appended after all pre-existing fields to preserve the positional
|
||||
# constructor contract for external callers (do not reorder above).
|
||||
supports_delete: bool = False
|
||||
supports_status_text: bool = False
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
@@ -163,6 +170,7 @@ class PlatformCapabilities:
|
||||
"supports_idempotency_token": self.supports_idempotency_token,
|
||||
"supports_media": self.supports_media,
|
||||
"supports_threads": self.supports_threads,
|
||||
"supports_status_text": self.supports_status_text,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
@@ -186,6 +194,7 @@ class PlatformCapabilities:
|
||||
supports_idempotency_token=data.get("supports_idempotency_token", False),
|
||||
supports_media=data.get("supports_media", False),
|
||||
supports_threads=data.get("supports_threads", False),
|
||||
supports_status_text=data.get("supports_status_text", False),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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"},
|
||||
|
||||
@@ -5,6 +5,7 @@ Implements progressive escalation from direct response to full autonomous mode.
|
||||
"""
|
||||
|
||||
from praisonaiagents._logging import get_logger
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Callable
|
||||
|
||||
@@ -284,6 +285,14 @@ class EscalationPipeline:
|
||||
else:
|
||||
return {"response": "", "success": False}
|
||||
|
||||
async def _run_chat(self, prompt: str, **kwargs) -> Any:
|
||||
"""Run the blocking ``agent.chat`` in a worker thread.
|
||||
|
||||
``asyncio.to_thread`` copies contextvars, so request-scoped state
|
||||
(user, trace, budget) survives, and the event loop stays responsive.
|
||||
"""
|
||||
return await asyncio.to_thread(self.agent.chat, prompt, **kwargs)
|
||||
|
||||
async def _execute_direct(self, prompt: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Execute direct response (Stage 0).
|
||||
@@ -298,7 +307,7 @@ class EscalationPipeline:
|
||||
|
||||
try:
|
||||
# Direct chat without tools
|
||||
response = self.agent.chat(prompt)
|
||||
response = await self._run_chat(prompt)
|
||||
|
||||
self._current_context.add_step(
|
||||
action="direct_response",
|
||||
@@ -335,16 +344,12 @@ class EscalationPipeline:
|
||||
# Limit tools for heuristic stage
|
||||
selected_tools = selected_tools[:self.config.heuristic_max_tools]
|
||||
|
||||
# Execute with selected tools
|
||||
# Execute with selected tools passed per call, never mutating
|
||||
# the shared, caller-owned agent (safe under concurrent execute()).
|
||||
if hasattr(self.agent, 'tools'):
|
||||
original_tools = self.agent.tools
|
||||
self.agent.tools = selected_tools
|
||||
|
||||
try:
|
||||
response = self.agent.chat(prompt)
|
||||
finally:
|
||||
if hasattr(self.agent, 'tools'):
|
||||
self.agent.tools = original_tools
|
||||
response = await self._run_chat(prompt, tools=selected_tools)
|
||||
else:
|
||||
response = await self._run_chat(prompt)
|
||||
|
||||
self._current_context.add_step(
|
||||
action="heuristic_response",
|
||||
@@ -381,7 +386,7 @@ class EscalationPipeline:
|
||||
|
||||
Respond with numbered steps, then execute them."""
|
||||
|
||||
response = self.agent.chat(plan_prompt)
|
||||
response = await self._run_chat(plan_prompt)
|
||||
|
||||
self._current_context.add_step(
|
||||
action="planned_response",
|
||||
@@ -427,7 +432,7 @@ Respond with numbered steps, then execute them."""
|
||||
|
||||
Use all available tools as needed. Verify your work before completing."""
|
||||
|
||||
response = self.agent.chat(autonomous_prompt)
|
||||
response = await self._run_chat(autonomous_prompt)
|
||||
|
||||
self._current_context.add_step(
|
||||
action="autonomous_response",
|
||||
|
||||
@@ -14,7 +14,6 @@ from typing import Any, Optional, Tuple, TYPE_CHECKING
|
||||
|
||||
from praisonaiagents._logging import get_logger
|
||||
|
||||
from .judge import judge_goal, _tail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .models import GoalCriteria, GoalState
|
||||
@@ -168,6 +167,8 @@ class GoalLoopMixin:
|
||||
self._persist_goal_state()
|
||||
return "continue", reason
|
||||
|
||||
from .judge import judge_goal, _tail # lazy: goal loop is opt-in
|
||||
|
||||
state.turns_used += 1
|
||||
verdict, reason = judge_goal(
|
||||
state,
|
||||
|
||||
@@ -310,8 +310,9 @@ class SubQuestionEngine:
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.query(question, context, **kwargs)
|
||||
"""Async version; runs the sync query (which may call the LLM) in a worker thread."""
|
||||
import asyncio
|
||||
return await asyncio.to_thread(self.query, question, context, **kwargs)
|
||||
|
||||
def _decompose_with_llm(self, question: str) -> List[str]:
|
||||
"""Decompose question using LLM."""
|
||||
|
||||
@@ -3598,6 +3598,19 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
reason = _cancel_reason()
|
||||
logging.debug(f"LLM tool loop cancelled before tool dispatch: {reason}")
|
||||
return _prepare_return_value(f"Task interrupted: {reason}")
|
||||
# Same per-turn cap as the chat-completions loop, applied
|
||||
# before recording the assistant turn so every declared
|
||||
# tool_call_id gets exactly one tool reply.
|
||||
if tool_call_count >= max_tool_calls_per_turn:
|
||||
logging.warning(f"Tool call limit reached ({max_tool_calls_per_turn}). Stopping to prevent infinite loop.")
|
||||
self._last_stop_reason = "max_steps"
|
||||
if not (final_response_text and final_response_text.strip()):
|
||||
final_response_text = f"Tool call limit reached ({max_tool_calls_per_turn} calls). Task may be too complex or there may be a broken tool causing repeated calls."
|
||||
break
|
||||
if tool_call_count + len(tool_calls) > max_tool_calls_per_turn:
|
||||
tool_calls = tool_calls[:max_tool_calls_per_turn - tool_call_count]
|
||||
logging.warning(f"Limiting batch to {len(tool_calls)} tool calls to stay within limit of {max_tool_calls_per_turn}.")
|
||||
tool_call_count += len(tool_calls)
|
||||
serializable_tool_calls = self._serialize_tool_calls(tool_calls)
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
@@ -3680,7 +3693,7 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
|
||||
# Append tool result as standard tool message
|
||||
messages.append(self._create_tool_message(
|
||||
function_name, tool_result,
|
||||
tool_result_obj.function_name, tool_result,
|
||||
tool_result_obj.tool_call_id, is_ollama=False))
|
||||
|
||||
# Report any unparseable tool-call arguments so the
|
||||
@@ -4316,6 +4329,22 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
reason = _cancel_reason()
|
||||
logging.debug(f"LLM tool loop cancelled before tool dispatch: {reason}")
|
||||
return _prepare_return_value(f"Task interrupted: {reason}")
|
||||
# Apply the cap BEFORE recording the assistant turn so every
|
||||
# declared tool_call_id gets exactly one tool reply.
|
||||
# Guardrail: Check tool call limit to prevent infinite loops
|
||||
# Allow execution if we haven't reached the limit yet, but limit the batch size
|
||||
if tool_call_count >= max_tool_calls_per_turn:
|
||||
logging.warning(f"Tool call limit reached ({max_tool_calls_per_turn}). Stopping to prevent infinite loop.")
|
||||
self._last_stop_reason = "max_steps"
|
||||
if not (final_response_text and final_response_text.strip()):
|
||||
final_response_text = f"Tool call limit reached ({max_tool_calls_per_turn} calls). Task may be too complex or there may be a broken tool causing repeated calls."
|
||||
break
|
||||
elif tool_call_count + len(tool_calls) > max_tool_calls_per_turn:
|
||||
# Limit the batch to stay within the total limit
|
||||
remaining_calls = max_tool_calls_per_turn - tool_call_count
|
||||
tool_calls = tool_calls[:remaining_calls]
|
||||
logging.warning(f"Limiting batch to {remaining_calls} tool calls to stay within limit of {max_tool_calls_per_turn}.")
|
||||
|
||||
# Convert tool_calls to a serializable format for all providers
|
||||
serializable_tool_calls = self._serialize_tool_calls(tool_calls)
|
||||
# Check if this is Ollama provider
|
||||
@@ -4337,20 +4366,6 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
tool_results = [] # Store current iteration tool results
|
||||
tool_result_mapping = {} # Store function results by name for Ollama chaining
|
||||
|
||||
# Guardrail: Check tool call limit to prevent infinite loops
|
||||
# Allow execution if we haven't reached the limit yet, but limit the batch size
|
||||
if tool_call_count >= max_tool_calls_per_turn:
|
||||
logging.warning(f"Tool call limit reached ({max_tool_calls_per_turn}). Stopping to prevent infinite loop.")
|
||||
self._last_stop_reason = "max_steps"
|
||||
if not (final_response_text and final_response_text.strip()):
|
||||
final_response_text = f"Tool call limit reached ({max_tool_calls_per_turn} calls). Task may be too complex or there may be a broken tool causing repeated calls."
|
||||
break
|
||||
elif tool_call_count + len(tool_calls) > max_tool_calls_per_turn:
|
||||
# Limit the batch to stay within the total limit
|
||||
remaining_calls = max_tool_calls_per_turn - tool_call_count
|
||||
tool_calls = tool_calls[:remaining_calls]
|
||||
logging.warning(f"Limiting batch to {remaining_calls} tool calls to stay within limit of {max_tool_calls_per_turn}.")
|
||||
|
||||
# Collect media-bearing follow-up messages so they are
|
||||
# flushed only AFTER every tool reply for this assistant
|
||||
# turn, keeping all tool replies consecutive (provider contract).
|
||||
@@ -5542,24 +5557,15 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
break
|
||||
last_tool_call_fingerprint = fingerprint
|
||||
|
||||
# Record the assistant turn before the tool replies, so
|
||||
# the follow-up request is a valid conversation.
|
||||
if is_ollama:
|
||||
messages.append({"role": "assistant", "content": content or ""})
|
||||
else:
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": content or "",
|
||||
"tool_calls": self._serialize_tool_calls(tool_calls),
|
||||
})
|
||||
|
||||
# Any prose the model emitted alongside the calls is
|
||||
# still part of the answer.
|
||||
# still part of the answer -- yield it before the cap
|
||||
# check so it is not lost when the limit is reached.
|
||||
if content:
|
||||
yield content
|
||||
|
||||
# Enforce the same per-turn guardrail the streaming and
|
||||
# non-streaming loops apply, so the fallback cannot run
|
||||
# non-streaming loops apply (before recording the assistant turn, so no
|
||||
# declared tool_call_id is left without a reply), so the fallback cannot run
|
||||
# tools past the caller's configured limit.
|
||||
if tool_call_count >= max_tool_calls_per_turn:
|
||||
logging.warning(
|
||||
@@ -5574,6 +5580,17 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
f"Limiting batch to {remaining_calls} tool calls to stay "
|
||||
f"within limit of {max_tool_calls_per_turn}.")
|
||||
|
||||
# Record the assistant turn before the tool replies, so
|
||||
# the follow-up request is a valid conversation.
|
||||
if is_ollama:
|
||||
messages.append({"role": "assistant", "content": content or ""})
|
||||
else:
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": content or "",
|
||||
"tool_calls": self._serialize_tool_calls(tool_calls),
|
||||
})
|
||||
|
||||
for tool_call in tool_calls:
|
||||
tool_call_count += 1
|
||||
function_name, arguments, tool_call_id = (
|
||||
@@ -5965,13 +5982,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
reason = _cancel_reason()
|
||||
logging.debug(f"Async LLM tool loop cancelled before tool dispatch: {reason}")
|
||||
return _prepare_return_value(f"Task interrupted: {reason}")
|
||||
serializable_tool_calls = self._serialize_tool_calls(tool_calls)
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": response_text,
|
||||
"tool_calls": serializable_tool_calls,
|
||||
})
|
||||
|
||||
# Apply the cap BEFORE recording the assistant turn so every
|
||||
# declared tool_call_id gets exactly one tool reply.
|
||||
# Guardrail: Check tool call limit to prevent infinite loops
|
||||
# Allow execution if we haven't reached the limit yet, but limit the batch size
|
||||
if tool_call_count >= max_tool_calls_per_turn:
|
||||
@@ -5986,6 +5998,13 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
tool_calls = tool_calls[:remaining_calls]
|
||||
logging.warning(f"Limiting batch to {remaining_calls} tool calls to stay within limit of {max_tool_calls_per_turn}.")
|
||||
|
||||
serializable_tool_calls = self._serialize_tool_calls(tool_calls)
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": response_text,
|
||||
"tool_calls": serializable_tool_calls,
|
||||
})
|
||||
|
||||
# Pre-resolve each tool call, then dispatch the valid ones
|
||||
# as a batch (concurrently when parallel_tool_calls is set).
|
||||
_parsed_calls = []
|
||||
@@ -6247,6 +6266,22 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
reason = _cancel_reason()
|
||||
logging.debug(f"Async LLM tool loop cancelled before tool dispatch: {reason}")
|
||||
return _prepare_return_value(f"Task interrupted: {reason}")
|
||||
# Apply the cap BEFORE recording the assistant turn so every
|
||||
# declared tool_call_id gets exactly one tool reply.
|
||||
# Guardrail: Check tool call limit to prevent infinite loops
|
||||
# Allow execution if we haven't reached the limit yet, but limit the batch size
|
||||
if tool_call_count >= max_tool_calls_per_turn:
|
||||
logging.warning(f"Tool call limit reached ({max_tool_calls_per_turn}). Stopping to prevent infinite loop.")
|
||||
self._last_stop_reason = "max_steps"
|
||||
if not (final_response_text and final_response_text.strip()):
|
||||
final_response_text = f"Tool call limit reached ({max_tool_calls_per_turn} calls). Task may be too complex or there may be a broken tool causing repeated calls."
|
||||
break
|
||||
elif tool_call_count + len(tool_calls) > max_tool_calls_per_turn:
|
||||
# Limit the batch to stay within the total limit
|
||||
remaining_calls = max_tool_calls_per_turn - tool_call_count
|
||||
tool_calls = tool_calls[:remaining_calls]
|
||||
logging.warning(f"Limiting batch to {remaining_calls} tool calls to stay within limit of {max_tool_calls_per_turn}.")
|
||||
|
||||
# Convert tool_calls to a serializable format for all providers
|
||||
serializable_tool_calls = self._serialize_tool_calls(tool_calls)
|
||||
# Check if it's Ollama provider
|
||||
@@ -6264,20 +6299,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
"tool_calls": serializable_tool_calls
|
||||
})
|
||||
|
||||
# Guardrail: Check tool call limit to prevent infinite loops
|
||||
# Allow execution if we haven't reached the limit yet, but limit the batch size
|
||||
if tool_call_count >= max_tool_calls_per_turn:
|
||||
logging.warning(f"Tool call limit reached ({max_tool_calls_per_turn}). Stopping to prevent infinite loop.")
|
||||
self._last_stop_reason = "max_steps"
|
||||
if not (final_response_text and final_response_text.strip()):
|
||||
final_response_text = f"Tool call limit reached ({max_tool_calls_per_turn} calls). Task may be too complex or there may be a broken tool causing repeated calls."
|
||||
break
|
||||
elif tool_call_count + len(tool_calls) > max_tool_calls_per_turn:
|
||||
# Limit the batch to stay within the total limit
|
||||
remaining_calls = max_tool_calls_per_turn - tool_call_count
|
||||
tool_calls = tool_calls[:remaining_calls]
|
||||
logging.warning(f"Limiting batch to {remaining_calls} tool calls to stay within limit of {max_tool_calls_per_turn}.")
|
||||
|
||||
tool_results = [] # Store current iteration tool results
|
||||
# Collect media-bearing follow-up messages so they are
|
||||
# flushed only AFTER every tool reply for this assistant
|
||||
|
||||
@@ -14,10 +14,11 @@ and only runs when a server actually challenges for OAuth, so servers without
|
||||
auth — or with pre-supplied client info/tokens — behave exactly as before.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import time
|
||||
import webbrowser
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from threading import Thread
|
||||
from threading import Lock, Thread
|
||||
from typing import Any, Dict, Optional
|
||||
from urllib.parse import urlencode, urlparse, parse_qs
|
||||
|
||||
@@ -69,6 +70,38 @@ def _make_callback_server(handler: OAuthCallbackHandler, port: int) -> HTTPServe
|
||||
return HTTPServer(("127.0.0.1", port), _Handler)
|
||||
|
||||
|
||||
# One loopback callback server per process, shared by every in-flight flow.
|
||||
# The redirect URI is fixed (it is what dynamic registration whitelisted), so
|
||||
# concurrent flows must not each try to bind the port; the state-keyed
|
||||
# ``OAuthCallbackHandler`` already routes each code to the right flow.
|
||||
_callback_lock = Lock()
|
||||
_callback_server: Optional[HTTPServer] = None
|
||||
_callback_handler: Optional[OAuthCallbackHandler] = None
|
||||
_callback_refs = 0
|
||||
|
||||
|
||||
def _acquire_callback() -> OAuthCallbackHandler:
|
||||
global _callback_server, _callback_handler, _callback_refs
|
||||
with _callback_lock:
|
||||
if _callback_server is None:
|
||||
handler = OAuthCallbackHandler()
|
||||
server = _make_callback_server(handler, OAUTH_CALLBACK_PORT)
|
||||
Thread(target=server.serve_forever, daemon=True).start()
|
||||
_callback_server, _callback_handler = server, handler
|
||||
_callback_refs += 1
|
||||
return _callback_handler
|
||||
|
||||
|
||||
def _release_callback() -> None:
|
||||
global _callback_server, _callback_handler, _callback_refs
|
||||
with _callback_lock:
|
||||
_callback_refs -= 1
|
||||
if _callback_refs <= 0 and _callback_server is not None:
|
||||
_callback_server.shutdown()
|
||||
_callback_server.server_close()
|
||||
_callback_server, _callback_handler, _callback_refs = None, None, 0
|
||||
|
||||
|
||||
class MCPOAuthProvider:
|
||||
"""
|
||||
Orchestrates zero-config OAuth for a single remote MCP server.
|
||||
@@ -198,11 +231,11 @@ class MCPOAuthProvider:
|
||||
if not self.open_browser:
|
||||
raise InteractiveAuthRequired(
|
||||
f"Interactive OAuth is required to authenticate with "
|
||||
f"{self.server_url}. Run this from an interactive session, or "
|
||||
f"{self.server_url}. Run `praisonai mcp auth {self.mcp_name}` "
|
||||
f"from a terminal on this host (works headless/over SSH), or "
|
||||
f"pre-supply tokens, to complete authentication."
|
||||
)
|
||||
|
||||
callback = OAuthCallbackHandler()
|
||||
state = generate_state()
|
||||
code_verifier = generate_code_verifier()
|
||||
code_challenge = generate_code_challenge(code_verifier)
|
||||
@@ -222,15 +255,30 @@ class MCPOAuthProvider:
|
||||
|
||||
auth_url = f"{metadata['authorization_endpoint']}?{urlencode(params)}"
|
||||
|
||||
server = _make_callback_server(callback, OAUTH_CALLBACK_PORT)
|
||||
thread = Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
callback = _acquire_callback()
|
||||
try:
|
||||
logger.debug(f"Opening browser for OAuth authorization: {auth_url}")
|
||||
webbrowser.open(auth_url)
|
||||
try:
|
||||
opened = webbrowser.open(auth_url)
|
||||
except Exception: # noqa: BLE001 - no usable browser
|
||||
opened = False
|
||||
if not opened:
|
||||
# Headless/remote host: let the operator consent on any device.
|
||||
# No stdin reading: a library must not consume the host's input.
|
||||
print(
|
||||
f"Open this URL on any device to authorise "
|
||||
f"{self.mcp_name}:\n {auth_url}\n"
|
||||
f"After consenting, the browser is redirected to "
|
||||
f"{get_redirect_url()}. On another device either forward "
|
||||
f"that port (ssh -L {OAUTH_CALLBACK_PORT}:127.0.0.1:"
|
||||
f"{OAUTH_CALLBACK_PORT} <host>) or copy the redirect URL "
|
||||
f"and run `curl '<redirect URL>'` on this host.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
code = callback.wait_for_callback(state, timeout=300.0)
|
||||
finally:
|
||||
server.shutdown()
|
||||
callback.clear_state(state)
|
||||
_release_callback()
|
||||
|
||||
tokens = self._exchange_code(
|
||||
metadata, client_info, code, code_verifier, timeout=timeout
|
||||
|
||||
@@ -281,18 +281,25 @@ class BaseStore(ABC):
|
||||
if not entries:
|
||||
return True
|
||||
try:
|
||||
import json
|
||||
from ...storage.base import FileLock
|
||||
from ...utils.atomic_io import atomic_write_json, quarantine_corrupt_file
|
||||
|
||||
archive_path = self._archive_path()
|
||||
existing: List[Dict[str, Any]] = []
|
||||
if archive_path.exists():
|
||||
with open(archive_path, "r", encoding="utf-8") as f:
|
||||
existing = json.load(f)
|
||||
if not isinstance(existing, list):
|
||||
existing = []
|
||||
existing.extend(e.to_dict() for e in entries)
|
||||
archive_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(archive_path, "w", encoding="utf-8") as f:
|
||||
json.dump(existing, f, ensure_ascii=False, indent=2)
|
||||
with FileLock(archive_path, timeout=10):
|
||||
existing: Any = []
|
||||
if archive_path.exists():
|
||||
try:
|
||||
with open(archive_path, "r", encoding="utf-8") as f:
|
||||
existing = json.load(f)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
existing = None
|
||||
if not isinstance(existing, list):
|
||||
# Keep the bad bytes as a sidecar and keep evicting.
|
||||
quarantine_corrupt_file(str(archive_path))
|
||||
existing = []
|
||||
existing.extend(e.to_dict() for e in entries)
|
||||
atomic_write_json(str(archive_path), existing, ensure_ascii=False, indent=2)
|
||||
return True
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
@@ -430,20 +430,18 @@ class RAG:
|
||||
Currently wraps sync version. Can be extended for true async.
|
||||
"""
|
||||
import asyncio
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.retrieve(
|
||||
query=query,
|
||||
top_k=top_k,
|
||||
hybrid=hybrid,
|
||||
rerank=rerank,
|
||||
filters=filters,
|
||||
user_id=user_id,
|
||||
agent_id=agent_id,
|
||||
run_id=run_id,
|
||||
**kwargs,
|
||||
),
|
||||
# asyncio.to_thread copies contextvars into the worker thread
|
||||
return await asyncio.to_thread(
|
||||
self.retrieve,
|
||||
query=query,
|
||||
top_k=top_k,
|
||||
hybrid=hybrid,
|
||||
rerank=rerank,
|
||||
filters=filters,
|
||||
user_id=user_id,
|
||||
agent_id=agent_id,
|
||||
run_id=run_id,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def query(
|
||||
@@ -523,19 +521,15 @@ class RAG:
|
||||
|
||||
Currently wraps sync version. Can be extended for true async.
|
||||
"""
|
||||
# For now, wrap sync version
|
||||
# TODO: Implement true async when Knowledge supports it
|
||||
import asyncio
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.query(
|
||||
question=question,
|
||||
user_id=user_id,
|
||||
agent_id=agent_id,
|
||||
run_id=run_id,
|
||||
**kwargs,
|
||||
),
|
||||
# asyncio.to_thread copies contextvars into the worker thread
|
||||
return await asyncio.to_thread(
|
||||
self.query,
|
||||
question=question,
|
||||
user_id=user_id,
|
||||
agent_id=agent_id,
|
||||
run_id=run_id,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def stream(
|
||||
@@ -580,6 +574,10 @@ class RAG:
|
||||
prompt = self._build_prompt(question, context)
|
||||
|
||||
# 5. Stream generation
|
||||
yield from self._stream_generate(prompt, **kwargs)
|
||||
|
||||
def _stream_generate(self, prompt: str, **kwargs) -> Iterator[str]:
|
||||
"""Synchronously stream the LLM answer for an already-built prompt."""
|
||||
if self.llm is None:
|
||||
yield "LLM not available."
|
||||
return
|
||||
@@ -615,10 +613,15 @@ class RAG:
|
||||
"""
|
||||
Async streaming RAG response.
|
||||
|
||||
Yields tokens as they are generated.
|
||||
Yields tokens as they are generated. Blocking retrieval, reranking
|
||||
and sync LLM streaming run in worker threads so the event loop
|
||||
stays responsive.
|
||||
"""
|
||||
# 1. Retrieve
|
||||
results = self._retrieve(
|
||||
import asyncio
|
||||
|
||||
# 1. Retrieve (embedding + vector search are blocking)
|
||||
results = await asyncio.to_thread(
|
||||
self._retrieve,
|
||||
query=question,
|
||||
user_id=user_id,
|
||||
agent_id=agent_id,
|
||||
@@ -627,7 +630,7 @@ class RAG:
|
||||
)
|
||||
|
||||
# 2. Rerank
|
||||
results = self._rerank(question, results)
|
||||
results = await asyncio.to_thread(self._rerank, question, results)
|
||||
|
||||
# 3. Build context
|
||||
context, used_results = self._build_context(results)
|
||||
@@ -651,8 +654,14 @@ class RAG:
|
||||
async for chunk in self.llm.astream(messages=messages, **kwargs):
|
||||
yield chunk
|
||||
else:
|
||||
# Fallback to sync streaming wrapped
|
||||
for chunk in self.stream(question, user_id, agent_id, run_id, **kwargs):
|
||||
# Fallback: pump the sync generator from a worker thread
|
||||
# (reuses the context already retrieved above)
|
||||
it = self._stream_generate(prompt, **kwargs)
|
||||
sentinel = object()
|
||||
while True:
|
||||
chunk = await asyncio.to_thread(next, it, sentinel)
|
||||
if chunk is sentinel:
|
||||
break
|
||||
yield chunk
|
||||
except Exception as e:
|
||||
logger.error(f"Async streaming failed: {e}")
|
||||
|
||||
@@ -39,6 +39,7 @@ if TYPE_CHECKING:
|
||||
from .sqlite_transcript_store import SqliteTranscriptStore
|
||||
from .protocols import SessionStoreProtocol, SessionMirrorProtocol
|
||||
from .hierarchy import HierarchicalSessionStore, SessionSnapshot, ExtendedSessionData
|
||||
from .encrypted_store import EncryptedSessionStore, SessionEncryptionError, generate_session_key
|
||||
|
||||
# Lazy loading for zero import overhead
|
||||
_module_cache = {}
|
||||
@@ -207,6 +208,12 @@ def __getattr__(name: str):
|
||||
_module_cache[name] = wrap_inter_agent
|
||||
return wrap_inter_agent
|
||||
|
||||
if name in ("EncryptedSessionStore", "SessionEncryptionError", "generate_session_key"):
|
||||
from . import encrypted_store
|
||||
value = getattr(encrypted_store, name)
|
||||
_module_cache[name] = value
|
||||
return value
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
@@ -249,9 +256,3 @@ __all__ = [
|
||||
"SessionEncryptionError",
|
||||
"generate_session_key",
|
||||
]
|
||||
|
||||
from .encrypted_store import ( # noqa: E402
|
||||
EncryptedSessionStore,
|
||||
SessionEncryptionError,
|
||||
generate_session_key,
|
||||
)
|
||||
@@ -474,6 +474,7 @@ class SqliteTranscriptStore(DefaultSessionStore):
|
||||
"model": data.get("model") or data.get("llm") or meta.get("model"),
|
||||
"total_tokens": data.get("total_tokens") or data.get("token_count") or meta.get("total_tokens"),
|
||||
"cost": data.get("cost") or meta.get("cost"),
|
||||
"tool_stats": meta.get("tool_stats"),
|
||||
})
|
||||
if len(sessions) >= limit:
|
||||
break
|
||||
|
||||
@@ -2076,6 +2076,8 @@ class DefaultSessionStore:
|
||||
"model": data.get("model") or data.get("llm") or (data.get("metadata") or {}).get("model"),
|
||||
"total_tokens": data.get("total_tokens") or data.get("token_count") or (data.get("metadata") or {}).get("total_tokens"),
|
||||
"cost": data.get("cost") or (data.get("metadata") or {}).get("cost"),
|
||||
# Per-tool outcome tally for `usage --by tool` (Issue #6158).
|
||||
"tool_stats": (data.get("metadata") or {}).get("tool_stats"),
|
||||
})
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError) as exc:
|
||||
self._report_unreadable_session(filename[:-5], filepath, exc)
|
||||
|
||||
@@ -21,6 +21,9 @@ logger = get_logger(__name__)
|
||||
# Default snapshot directory (uses centralized paths - DRY)
|
||||
DEFAULT_SNAPSHOT_DIR = str(get_snapshots_dir())
|
||||
|
||||
# Bound every git call and never block on an interactive prompt.
|
||||
_GIT_TIMEOUT = 60
|
||||
|
||||
@dataclass
|
||||
class FileDiff:
|
||||
"""Represents a diff for a single file."""
|
||||
@@ -158,8 +161,13 @@ class FileSnapshot:
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=check,
|
||||
timeout=_GIT_TIMEOUT,
|
||||
env={**os.environ, "GIT_TERMINAL_PROMPT": "0"},
|
||||
)
|
||||
return result
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.error("Git command timed out after %ss: %s", _GIT_TIMEOUT, " ".join(cmd))
|
||||
raise
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.error(f"Git command failed: {' '.join(cmd)}")
|
||||
logger.error(f"stderr: {e.stderr}")
|
||||
|
||||
@@ -18,6 +18,13 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
# Minimum similarity ratio for an accepted block-anchor (fuzzy) match.
|
||||
_BLOCK_ANCHOR_THRESHOLD = 0.7
|
||||
# Minimum similarity for a "closest match" hint on the no-match error path.
|
||||
_CLOSEST_HINT_THRESHOLD = 0.6
|
||||
# Bounds for the hint search/output: skip long targets, cap total comparison
|
||||
# work (sum of len(window) * len(target)), and truncate long displayed lines.
|
||||
_CLOSEST_HINT_MAX_CHARS = 2000
|
||||
_CLOSEST_HINT_BUDGET = 2_000_000
|
||||
_CLOSEST_HINT_LINE_CHARS = 200
|
||||
|
||||
# Maximum characters of diagnostics output appended to a tool result.
|
||||
_DIAGNOSTICS_MAX_CHARS = 2000
|
||||
@@ -800,6 +807,69 @@ class EditTools:
|
||||
return spans
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _closest_region(content: str, old: str, max_lines: int = 20000):
|
||||
"""Find the block of ``content`` most similar to ``old``.
|
||||
|
||||
Used only on the no-match path. Returns ``(start_index, n_lines)`` or
|
||||
``None`` when nothing is similar enough or the search would be too
|
||||
expensive (bounded by input size and a total comparison budget).
|
||||
"""
|
||||
if len(old) > _CLOSEST_HINT_MAX_CHARS:
|
||||
return None
|
||||
old_lines = [line.strip() for line in old.splitlines()]
|
||||
lines = content.splitlines()
|
||||
n = len(old_lines)
|
||||
if not n or not lines or len(lines) > max_lines:
|
||||
return None
|
||||
n = min(n, len(lines))
|
||||
stripped = [line.strip() for line in lines]
|
||||
target = '\n'.join(old_lines)
|
||||
matcher = difflib.SequenceMatcher(autojunk=False)
|
||||
matcher.set_seq2(target)
|
||||
budget = _CLOSEST_HINT_BUDGET
|
||||
best_ratio, best_i = _CLOSEST_HINT_THRESHOLD, None
|
||||
for i in range(len(lines) - n + 1):
|
||||
window = '\n'.join(stripped[i:i + n])
|
||||
matcher.set_seq1(window)
|
||||
if matcher.real_quick_ratio() <= best_ratio:
|
||||
continue
|
||||
budget -= len(window) # quick_ratio() is a pure-Python pass over window
|
||||
if budget < 0:
|
||||
break
|
||||
if matcher.quick_ratio() <= best_ratio:
|
||||
continue
|
||||
budget -= len(window) * len(target)
|
||||
if budget < 0:
|
||||
break # optional hint: keep the best found so far
|
||||
ratio = matcher.ratio()
|
||||
if ratio > best_ratio:
|
||||
best_ratio, best_i = ratio, i
|
||||
return None if best_i is None else (best_i, n)
|
||||
|
||||
@staticmethod
|
||||
def _format_region(lines, start: int, n: int) -> str:
|
||||
"""Render up to 10 lines from ``start`` with line numbers, each truncated."""
|
||||
shown = lines[start:start + min(n, 10)]
|
||||
width = _CLOSEST_HINT_LINE_CHARS
|
||||
body = '\n'.join(
|
||||
f"{start + k + 1:>5}| "
|
||||
+ (line if len(line) <= width else line[:width] + "... [truncated]")
|
||||
for k, line in enumerate(shown))
|
||||
if n > len(shown):
|
||||
body += "\n ..."
|
||||
return f"lines {start + 1}-{start + n}:\n{body}"
|
||||
|
||||
def _not_found_error(self, prefix: str, content: str, old: str, region) -> str:
|
||||
"""Build the no-match error, adding ``region`` (from _closest_region) if any."""
|
||||
preview = old[:50] + ("..." if len(old) > 50 else "")
|
||||
if not region:
|
||||
return f"{prefix}: '{preview}'"
|
||||
hint = self._format_region(content.splitlines(), *region)
|
||||
return (f"{prefix}: '{preview}'. Closest match in the file, {hint}\n"
|
||||
f"Adjust old_string to match these lines exactly, "
|
||||
f"or re-read the file if it changed.")
|
||||
|
||||
@require_approval(risk_level="high")
|
||||
def edit_file(self, filepath: str, old_string: str, new_string: str,
|
||||
replace_all: bool = False, expected_hash: Optional[str] = None,
|
||||
@@ -885,8 +955,13 @@ class EditTools:
|
||||
spans = self._find_spans(content, old_string)
|
||||
|
||||
if not spans:
|
||||
preview = old_string[:50] + ("..." if len(old_string) > 50 else "")
|
||||
return f"Error: String not found in file: '{preview}'"
|
||||
# No "already applied" shortcut: similarity cannot prove
|
||||
# *where* a prior edit landed (ties pick the wrong site),
|
||||
# and a false success silently drops an edit. The hint
|
||||
# shows the closest lines, so a retried edit is visible.
|
||||
return self._not_found_error(
|
||||
"Error: String not found in file", content, old_string,
|
||||
self._closest_region(content, old_string))
|
||||
|
||||
occurrences = len(spans)
|
||||
|
||||
@@ -1180,8 +1255,9 @@ class EditTools:
|
||||
return f"Error: Empty hunk in move for '{op['src']}'"
|
||||
spans = self._find_spans(updated, old_block)
|
||||
if not spans:
|
||||
preview = old_block[:50] + ("..." if len(old_block) > 50 else "")
|
||||
return (f"Error: Hunk not found in '{op['src']}': '{preview}'")
|
||||
return self._not_found_error(
|
||||
f"Error: Hunk not found in '{op['src']}'", updated, old_block,
|
||||
self._closest_region(updated, old_block))
|
||||
if len(spans) > 1:
|
||||
preview = old_block[:30] + ("..." if len(old_block) > 30 else "")
|
||||
return (f"Error: Ambiguous hunk in '{op['src']}': '{preview}' "
|
||||
@@ -1264,8 +1340,9 @@ class EditTools:
|
||||
return f"Error: Empty hunk in update for '{op['path']}'"
|
||||
spans = self._find_spans(updated, old_block)
|
||||
if not spans:
|
||||
preview = old_block[:50] + ("..." if len(old_block) > 50 else "")
|
||||
return (f"Error: Hunk not found in '{op['path']}': '{preview}'")
|
||||
return self._not_found_error(
|
||||
f"Error: Hunk not found in '{op['path']}'", updated, old_block,
|
||||
self._closest_region(updated, old_block))
|
||||
if len(spans) > 1:
|
||||
preview = old_block[:30] + ("..." if len(old_block) > 30 else "")
|
||||
return (f"Error: Ambiguous hunk in '{op['path']}': '{preview}' "
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -387,6 +387,8 @@ class A2A:
|
||||
user_input=user_input,
|
||||
task_id=task.id,
|
||||
context_id=task.context_id,
|
||||
task_store=self.task_store,
|
||||
running=self._running,
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
|
||||
@@ -4,6 +4,7 @@ A2A Streaming Events
|
||||
SSE streaming support for A2A task updates.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import AsyncIterator, Optional
|
||||
|
||||
@@ -15,7 +16,11 @@ from praisonaiagents.ui.a2a.types import (
|
||||
TaskArtifactUpdateEvent,
|
||||
Artifact,
|
||||
TextPart,
|
||||
Message,
|
||||
)
|
||||
from praisonaiagents._logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class A2AEventEncoder:
|
||||
@@ -65,6 +70,7 @@ def create_status_event(
|
||||
state: TaskState,
|
||||
context_id: Optional[str] = None,
|
||||
final: bool = False,
|
||||
message: Optional[Message] = None,
|
||||
) -> TaskStatusUpdateEvent:
|
||||
"""
|
||||
Create a TaskStatusUpdateEvent.
|
||||
@@ -74,6 +80,7 @@ def create_status_event(
|
||||
state: New task state
|
||||
context_id: Optional context ID
|
||||
final: Whether this is the final event
|
||||
message: Optional status message (e.g. failure reason)
|
||||
|
||||
Returns:
|
||||
TaskStatusUpdateEvent
|
||||
@@ -81,7 +88,7 @@ def create_status_event(
|
||||
return TaskStatusUpdateEvent(
|
||||
task_id=task_id,
|
||||
context_id=context_id,
|
||||
status=TaskStatus(state=state),
|
||||
status=TaskStatus(state=state, message=message),
|
||||
final=final,
|
||||
)
|
||||
|
||||
@@ -130,60 +137,123 @@ async def stream_agent_response(
|
||||
user_input: str,
|
||||
task_id: str,
|
||||
context_id: Optional[str] = None,
|
||||
*,
|
||||
task_store=None,
|
||||
running: Optional[dict] = None,
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Stream agent response as A2A events.
|
||||
|
||||
The sync agent call runs in a worker thread so the event loop stays
|
||||
responsive. When ``task_store``/``running`` are given, the task's
|
||||
lifecycle is mirrored into the store and the call is registered so
|
||||
``tasks/get`` and ``tasks/cancel`` work for streamed tasks.
|
||||
|
||||
Args:
|
||||
agent: PraisonAI Agent instance
|
||||
agent: PraisonAI Agent (uses ``chat``) or AgentTeam (uses ``start``)
|
||||
user_input: User input text
|
||||
task_id: Task ID
|
||||
context_id: Optional context ID
|
||||
task_store: Optional TaskStore to keep in sync
|
||||
running: Optional dict of task_id -> future, used for cancellation
|
||||
|
||||
Yields:
|
||||
SSE-formatted event strings
|
||||
"""
|
||||
from praisonaiagents.ui.a2a.conversion import praisonai_to_a2a_message
|
||||
|
||||
encoder = A2AEventEncoder()
|
||||
|
||||
# Emit working status
|
||||
working_event = create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.WORKING,
|
||||
context_id=context_id,
|
||||
)
|
||||
yield encoder.encode_task_status(working_event)
|
||||
if task_store is not None:
|
||||
task_store.update_status(task_id, TaskState.WORKING)
|
||||
# Keep the task while the agent runs: a cancel can make it
|
||||
# terminal (and evictable) before we finish.
|
||||
task_store._pinned.add(task_id)
|
||||
|
||||
fut = None
|
||||
try:
|
||||
# Get response from agent
|
||||
response = agent.chat(user_input)
|
||||
# Emit working status
|
||||
yield encoder.encode_task_status(create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.WORKING,
|
||||
context_id=context_id,
|
||||
))
|
||||
|
||||
# tasks/cancel may have landed while we were suspended at the
|
||||
# yield above, before the call was registered: don't start it.
|
||||
if task_store is not None:
|
||||
task = task_store.get_task(task_id)
|
||||
if task is not None and task.status.state == TaskState.CANCELLED:
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
call = agent.chat if hasattr(agent, "chat") else agent.start
|
||||
fut = asyncio.ensure_future(asyncio.to_thread(call, user_input))
|
||||
if running is not None:
|
||||
running[task_id] = fut
|
||||
response = str(await fut)
|
||||
|
||||
# Emit artifact with response
|
||||
artifact_event = create_artifact_event(
|
||||
task_id=task_id,
|
||||
content=str(response),
|
||||
content=response,
|
||||
context_id=context_id,
|
||||
last_chunk=True,
|
||||
)
|
||||
if task_store is not None:
|
||||
task_store.add_to_history(
|
||||
task_id, praisonai_to_a2a_message(response, task_id=task_id)
|
||||
)
|
||||
task_store.add_artifact(task_id, artifact_event.artifact)
|
||||
task_store.update_status(task_id, TaskState.COMPLETED)
|
||||
yield encoder.encode_task_artifact(artifact_event)
|
||||
|
||||
# Emit completed status
|
||||
completed_event = create_status_event(
|
||||
yield encoder.encode_task_status(create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.COMPLETED,
|
||||
context_id=context_id,
|
||||
final=True,
|
||||
)
|
||||
yield encoder.encode_task_status(completed_event)
|
||||
|
||||
except Exception:
|
||||
# Emit failed status
|
||||
failed_event = create_status_event(
|
||||
))
|
||||
|
||||
except asyncio.CancelledError:
|
||||
# tasks/cancel marks the task CANCELLED before cancelling the call;
|
||||
# report that to the client. Any other cancellation (e.g. client
|
||||
# disconnect) is propagated.
|
||||
task = task_store.get_task(task_id) if task_store is not None else None
|
||||
if task is None or task.status.state != TaskState.CANCELLED:
|
||||
if task_store is not None:
|
||||
task_store.update_status(task_id, TaskState.CANCELLED)
|
||||
raise
|
||||
yield encoder.encode_task_status(create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.CANCELLED,
|
||||
context_id=context_id,
|
||||
final=True,
|
||||
))
|
||||
|
||||
except Exception as e:
|
||||
logger.error("A2A stream failed for task %s: %s", task_id, e, exc_info=True)
|
||||
reason = praisonai_to_a2a_message(str(e), task_id=task_id)
|
||||
if task_store is not None:
|
||||
task_store.update_status(task_id, TaskState.FAILED, message=reason)
|
||||
yield encoder.encode_task_status(create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.FAILED,
|
||||
context_id=context_id,
|
||||
final=True,
|
||||
)
|
||||
yield encoder.encode_task_status(failed_event)
|
||||
message=reason,
|
||||
))
|
||||
|
||||
except GeneratorExit:
|
||||
# Stream closed early (client went away): leave a terminal state.
|
||||
# No-op if the task already finished.
|
||||
if task_store is not None:
|
||||
task_store.update_status(task_id, TaskState.CANCELLED)
|
||||
raise
|
||||
|
||||
finally:
|
||||
if running is not None and running.get(task_id) is fut:
|
||||
running.pop(task_id, None)
|
||||
if task_store is not None:
|
||||
task_store._pinned.discard(task_id)
|
||||
|
||||
# Emit done
|
||||
yield encoder.encode_done()
|
||||
@@ -5,6 +5,7 @@ Provides utilities for streaming PraisonAI agent responses as AG-UI events.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
@@ -30,6 +31,12 @@ from praisonaiagents.ui.agui.types import (
|
||||
from praisonaiagents.ui.protocols import A2UI_MIME_TYPE
|
||||
from praisonaiagents.streaming.events import StreamEventType as _SE
|
||||
|
||||
# Identifies which AG-UI run owns the current task/thread. The agent's
|
||||
# stream_emitter is shared by every request served by one Agent, and its
|
||||
# callbacks fire synchronously in the emitting context, so each run's callback
|
||||
# only accepts events emitted while its own owner token is active.
|
||||
_RUN_OWNER: contextvars.ContextVar = contextvars.ContextVar("agui_run_owner", default=None)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EventBuffer:
|
||||
@@ -466,20 +473,29 @@ async def async_stream_agent_response(
|
||||
text_started = False
|
||||
text_ended = False
|
||||
original_emitter = getattr(agent, "stream_emitter", None)
|
||||
owner = object()
|
||||
|
||||
def on_stream_event(event: Any) -> None:
|
||||
# Drop events emitted by another concurrent run on the same agent.
|
||||
if _RUN_OWNER.get() is not owner:
|
||||
return
|
||||
loop.call_soon_threadsafe(queue.put_nowait, ("stream", event))
|
||||
|
||||
if original_emitter is not None:
|
||||
original_emitter.add_callback(on_stream_event)
|
||||
|
||||
async def run_chat() -> Any:
|
||||
# Tasks run in a copy of the creator's context, so this set is local
|
||||
# to the chat task and never leaks into the caller's context.
|
||||
_RUN_OWNER.set(owner)
|
||||
try:
|
||||
if hasattr(agent, "achat"):
|
||||
return await agent.achat(user_input, stream=True)
|
||||
if hasattr(agent, "chat"):
|
||||
# run_in_executor does not copy ContextVars; carry the owner.
|
||||
ctx = contextvars.copy_context()
|
||||
return await loop.run_in_executor(
|
||||
None, lambda: agent.chat(user_input, stream=True)
|
||||
None, lambda: ctx.run(agent.chat, user_input, stream=True)
|
||||
)
|
||||
raise ValueError("Agent must have 'chat' or 'achat' method")
|
||||
finally:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -9,9 +9,13 @@ repository or git is unavailable.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
# Bound every git call and never block on an interactive prompt.
|
||||
_GIT_TIMEOUT = 60
|
||||
|
||||
|
||||
def _slugify(name: str) -> str:
|
||||
"""Return a git-ref-safe, collision-resistant slug for ``name``.
|
||||
@@ -67,12 +71,16 @@ class GitWorktreeAdapter:
|
||||
cwd=str(cwd or self._root),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=_GIT_TIMEOUT,
|
||||
env={**os.environ, "GIT_TERMINAL_PROMPT": "0"},
|
||||
)
|
||||
|
||||
def _is_git_repo(self) -> bool:
|
||||
import subprocess
|
||||
|
||||
try:
|
||||
result = self._git("rev-parse", "--is-inside-work-tree")
|
||||
except (FileNotFoundError, OSError):
|
||||
except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
|
||||
return False
|
||||
return result.returncode == 0 and result.stdout.strip() == "true"
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.20"
|
||||
version = "1.7.21"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Concurrent AG-UI runs on one agent must not receive each other's tokens (#6150)."""
|
||||
|
||||
import asyncio
|
||||
|
||||
from praisonaiagents.streaming.events import StreamEvent, StreamEventEmitter, StreamEventType
|
||||
from praisonaiagents.ui.agui.streaming import async_stream_agent_response
|
||||
|
||||
|
||||
class _AsyncAgent:
|
||||
def __init__(self):
|
||||
self.stream_emitter = StreamEventEmitter()
|
||||
|
||||
async def achat(self, text, stream=True):
|
||||
for tok in (f"[{text}-1]", f"[{text}-2]"):
|
||||
self.stream_emitter.emit(StreamEvent(type=StreamEventType.DELTA_TEXT, content=tok))
|
||||
await asyncio.sleep(0.02)
|
||||
return text
|
||||
|
||||
|
||||
class _SyncAgent:
|
||||
def __init__(self):
|
||||
self.stream_emitter = StreamEventEmitter()
|
||||
|
||||
def chat(self, text, stream=True):
|
||||
import time
|
||||
for tok in (f"[{text}-1]", f"[{text}-2]"):
|
||||
self.stream_emitter.emit(StreamEvent(type=StreamEventType.DELTA_TEXT, content=tok))
|
||||
time.sleep(0.02)
|
||||
return text
|
||||
|
||||
|
||||
async def _collect(agent, who):
|
||||
out = []
|
||||
async for ev in async_stream_agent_response(agent, who, "th-" + who, "run-" + who):
|
||||
if getattr(ev, "delta", None):
|
||||
out.append(ev.delta)
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _run_pair(agent):
|
||||
async def main():
|
||||
return await asyncio.gather(_collect(agent, "ALICE"), _collect(agent, "BOB"))
|
||||
return asyncio.run(main())
|
||||
|
||||
|
||||
def test_concurrent_async_runs_do_not_leak_tokens():
|
||||
alice, bob = _run_pair(_AsyncAgent())
|
||||
assert alice == "[ALICE-1][ALICE-2]"
|
||||
assert bob == "[BOB-1][BOB-2]"
|
||||
|
||||
|
||||
def test_concurrent_sync_runs_do_not_leak_tokens():
|
||||
alice, bob = _run_pair(_SyncAgent())
|
||||
assert alice == "[ALICE-1][ALICE-2]"
|
||||
assert bob == "[BOB-1][BOB-2]"
|
||||
|
||||
|
||||
def test_model_hook_path_keeps_run_owner():
|
||||
"""With a wrap_model_call hook, achat re-schedules the model call through an
|
||||
executor; the run owner must survive so events still reach the right run."""
|
||||
from praisonaiagents.agent.chat_mixin import ChatMixin
|
||||
from praisonaiagents.hooks import wrap_model_call
|
||||
|
||||
@wrap_model_call
|
||||
def passthrough(request, call_next):
|
||||
return call_next(request)
|
||||
|
||||
class _HookedAgent(ChatMixin):
|
||||
name = "hooked"
|
||||
llm = "fake"
|
||||
|
||||
def __init__(self):
|
||||
self.stream_emitter = StreamEventEmitter()
|
||||
self._hooks = [passthrough]
|
||||
|
||||
async def _achat_completion_with_retry_core(self, messages, *a, **kw):
|
||||
text = messages
|
||||
for tok in (f"[{text}-1]", f"[{text}-2]"):
|
||||
self.stream_emitter.emit(StreamEvent(type=StreamEventType.DELTA_TEXT, content=tok))
|
||||
await asyncio.sleep(0.02)
|
||||
return text
|
||||
|
||||
async def achat(self, text, stream=True):
|
||||
return await self._achat_completion_with_retry_dispatch(text, stream=stream)
|
||||
|
||||
agent = _HookedAgent()
|
||||
assert agent._get_model_middleware_manager() is not None
|
||||
alice, bob = _run_pair(agent)
|
||||
assert alice == "[ALICE-1][ALICE-2]"
|
||||
assert bob == "[BOB-1][BOB-2]"
|
||||
@@ -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"]
|
||||
@@ -173,7 +173,7 @@ class TestGoalLoopGate:
|
||||
verdicts = iter([("continue", "not yet"), ("done", "goal met")])
|
||||
|
||||
with patch.object(agent, "chat", side_effect=_unique_chat()), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal",
|
||||
patch("praisonaiagents.goal.judge.judge_goal",
|
||||
side_effect=lambda *a, **k: next(verdicts)):
|
||||
result = agent.run_goal("do the task", goal="finish it", max_turns=10)
|
||||
|
||||
@@ -193,7 +193,7 @@ class TestGoalLoopGate:
|
||||
return f"still working {counter['n']}"
|
||||
|
||||
with patch.object(agent, "chat", side_effect=fake_chat), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal",
|
||||
patch("praisonaiagents.goal.judge.judge_goal",
|
||||
side_effect=lambda *a, **k: next(verdicts)):
|
||||
result = agent.run_goal("task", goal="g", max_turns=10)
|
||||
|
||||
@@ -205,7 +205,7 @@ class TestGoalLoopGate:
|
||||
def test_budget_exhaust_pauses_recoverably(self):
|
||||
agent = _make_agent()
|
||||
with patch.object(agent, "chat", side_effect=_unique_chat()), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal",
|
||||
patch("praisonaiagents.goal.judge.judge_goal",
|
||||
return_value=("continue", "never done")):
|
||||
result = agent.run_goal("task", goal="g", max_turns=3)
|
||||
|
||||
@@ -218,7 +218,7 @@ class TestGoalLoopGate:
|
||||
agent = _make_agent()
|
||||
# judge_goal itself fails open to ("continue", "judge unavailable")
|
||||
with patch.object(agent, "chat", side_effect=_unique_chat()), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal",
|
||||
patch("praisonaiagents.goal.judge.judge_goal",
|
||||
return_value=("continue", "judge unavailable")):
|
||||
result = agent.run_goal("task", goal="g", max_turns=2)
|
||||
assert result.completion_reason == "budget_paused"
|
||||
@@ -226,7 +226,7 @@ class TestGoalLoopGate:
|
||||
def test_consecutive_parse_failures_autopause(self):
|
||||
agent = _make_agent()
|
||||
with patch.object(agent, "chat", side_effect=_unique_chat()), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal",
|
||||
patch("praisonaiagents.goal.judge.judge_goal",
|
||||
return_value=("continue", "judge response unparseable")):
|
||||
result = agent.run_goal("task", goal="g", max_turns=20)
|
||||
# auto-pause after 3 consecutive parse failures, well before max_turns
|
||||
@@ -259,7 +259,7 @@ class TestConstraintViolation:
|
||||
)
|
||||
# Simulate a judge that respects constraints: constraint violated → continue
|
||||
with patch.object(agent, "chat", side_effect=_unique_chat()), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal",
|
||||
patch("praisonaiagents.goal.judge.judge_goal",
|
||||
return_value=("continue", "constraint violated")):
|
||||
result = agent.run_goal("task", goal="g", criteria=criteria, max_turns=2)
|
||||
assert result.completion_reason == "budget_paused"
|
||||
@@ -282,7 +282,7 @@ class TestAsyncGoalGate:
|
||||
state = GoalState(goal="g", max_turns=5)
|
||||
state.status = "active"
|
||||
agent._goal_state = state
|
||||
with patch("praisonaiagents.goal.loop.judge_goal",
|
||||
with patch("praisonaiagents.goal.judge.judge_goal",
|
||||
return_value=("done", "goal met")):
|
||||
outcome = asyncio.run(agent._goal_gate_async("final response"))
|
||||
assert outcome == ("done", "goal met")
|
||||
@@ -314,7 +314,7 @@ class TestAsyncGoalGate:
|
||||
return ("done", "eventually met")
|
||||
|
||||
async def scenario():
|
||||
with patch("praisonaiagents.goal.loop.judge_goal",
|
||||
with patch("praisonaiagents.goal.judge.judge_goal",
|
||||
side_effect=slow_judge):
|
||||
task = asyncio.ensure_future(agent._goal_gate_async("resp"))
|
||||
await asyncio.to_thread(started.wait, 2.0)
|
||||
@@ -330,7 +330,7 @@ class TestAsyncGoalGate:
|
||||
def test_run_until_with_goal_delegates(self):
|
||||
agent = _make_agent()
|
||||
with patch.object(agent, "chat", side_effect=_unique_chat()), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal",
|
||||
patch("praisonaiagents.goal.judge.judge_goal",
|
||||
return_value=("done", "met")):
|
||||
result = agent.run_until("task", goal="reach the goal", max_iterations=5)
|
||||
assert result.completion_reason == "goal_met"
|
||||
@@ -7,7 +7,7 @@ at ``turns_used=0``. A crashed run (still ``active``) was also never resumed.
|
||||
|
||||
import pytest
|
||||
|
||||
import praisonaiagents.goal.loop as goal_loop
|
||||
import praisonaiagents.goal.judge as goal_judge
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.goal.models import GoalState
|
||||
from praisonaiagents.model_harness import ScriptedModel
|
||||
@@ -20,7 +20,7 @@ GOAL = "all tests pass"
|
||||
def _isolated_home(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("PRAISONAI_HOME", str(tmp_path / ".praisonai"))
|
||||
monkeypatch.setattr(
|
||||
goal_loop, "judge_goal", lambda state, *a, **k: ("continue", "not yet")
|
||||
goal_judge, "judge_goal", lambda state, *a, **k: ("continue", "not yet")
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -228,7 +228,7 @@ class TestGoalLoopEvidence:
|
||||
|
||||
with patch.object(agent, "chat",
|
||||
side_effect=lambda p: f"step {counter.__setitem__('n', counter['n'] + 1) or counter['n']}"), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal", side_effect=fake_judge):
|
||||
patch("praisonaiagents.goal.judge.judge_goal", side_effect=fake_judge):
|
||||
result = agent.run_goal("task", goal="g", max_turns=2)
|
||||
|
||||
# Blocking hook fails every turn → judge never consulted, ends paused.
|
||||
@@ -247,7 +247,7 @@ class TestGoalLoopEvidence:
|
||||
return ("done", "met")
|
||||
|
||||
with patch.object(agent, "chat", side_effect=lambda pr: "did it"), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal", side_effect=fake_judge):
|
||||
patch("praisonaiagents.goal.judge.judge_goal", side_effect=fake_judge):
|
||||
result = agent.run_goal("task", goal="g", max_turns=5)
|
||||
|
||||
assert result.completion_reason == "goal_met"
|
||||
@@ -283,7 +283,7 @@ class TestGoalLoopEvidence:
|
||||
return "trying"
|
||||
|
||||
with patch.object(agent, "chat", side_effect=fake_chat), \
|
||||
patch("praisonaiagents.goal.loop.judge_goal",
|
||||
patch("praisonaiagents.goal.judge.judge_goal",
|
||||
side_effect=lambda *a, **k: ("done", "met")):
|
||||
result = agent.run_goal("task", goal="g", max_turns=2)
|
||||
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Headless and concurrent OAuth callback handling (issue #6179)."""
|
||||
|
||||
import threading
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def provider_env(tmp_path, monkeypatch):
|
||||
from praisonaiagents.mcp import oauth_provider
|
||||
from praisonaiagents.mcp.mcp_auth_storage import MCPAuthStorage
|
||||
|
||||
started = []
|
||||
|
||||
class FakeServer:
|
||||
def serve_forever(self):
|
||||
pass
|
||||
|
||||
def shutdown(self):
|
||||
pass
|
||||
|
||||
def server_close(self):
|
||||
pass
|
||||
|
||||
def fake_make(handler, port):
|
||||
started.append(handler)
|
||||
return FakeServer()
|
||||
|
||||
monkeypatch.setattr(oauth_provider, "_make_callback_server", fake_make)
|
||||
monkeypatch.setattr(
|
||||
oauth_provider.MCPOAuthProvider,
|
||||
"_exchange_code",
|
||||
lambda self, m, c, code, v, timeout: {"access_token": f"tok-{code}"},
|
||||
)
|
||||
storage = MCPAuthStorage(filepath=str(tmp_path / "auth.json"))
|
||||
return oauth_provider, storage, started
|
||||
|
||||
|
||||
METADATA = {
|
||||
"authorization_endpoint": "https://auth.example.com/authorize",
|
||||
"token_endpoint": "https://auth.example.com/token",
|
||||
}
|
||||
|
||||
|
||||
def _state_of(url):
|
||||
return parse_qs(urlparse(url).query)["state"][0]
|
||||
|
||||
|
||||
def test_concurrent_flows_share_one_callback_server(provider_env, monkeypatch):
|
||||
oauth_provider, storage, started = provider_env
|
||||
urls = []
|
||||
both_opened = threading.Barrier(2)
|
||||
|
||||
def fake_open(url):
|
||||
urls.append((threading.current_thread().name, url))
|
||||
both_opened.wait(timeout=5)
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(oauth_provider.webbrowser, "open", fake_open)
|
||||
|
||||
results = {}
|
||||
|
||||
def run(name):
|
||||
p = oauth_provider.MCPOAuthProvider(name, f"https://{name}.example/mcp", storage=storage)
|
||||
results[name] = p._run_pkce_flow(METADATA, {"client_id": "c"}, None, timeout=5)
|
||||
|
||||
threads = [threading.Thread(target=run, args=(n,), name=n) for n in ("a", "b")]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for _ in range(100):
|
||||
if len(urls) == 2:
|
||||
break
|
||||
threading.Event().wait(0.05)
|
||||
assert len(started) == 1 # single bind, no port collision
|
||||
handler = started[0]
|
||||
for name, url in urls:
|
||||
handler.receive_callback(_state_of(url), f"code-{name}")
|
||||
for t in threads:
|
||||
t.join(timeout=5)
|
||||
|
||||
# Each flow must get the code sent for *its* state, not just any code.
|
||||
assert results == {"a": "tok-code-a", "b": "tok-code-b"}
|
||||
assert oauth_provider._callback_server is None # released after last flow
|
||||
|
||||
|
||||
def test_headless_prints_url_and_never_reads_stdin(provider_env, monkeypatch, capsys):
|
||||
"""Headless consent completes via the loopback callback; host stdin is untouched."""
|
||||
oauth_provider, storage, started = provider_env
|
||||
monkeypatch.setattr(oauth_provider.webbrowser, "open", lambda url: False)
|
||||
|
||||
def no_input(prompt=""):
|
||||
raise AssertionError("library must not read the host's stdin")
|
||||
|
||||
monkeypatch.setattr("builtins.input", no_input)
|
||||
|
||||
def deliver():
|
||||
# Simulates an SSH-forwarded browser or `curl <redirect>` on the host.
|
||||
for _ in range(100):
|
||||
err = capsys.readouterr().err
|
||||
if "Open this URL" in err:
|
||||
auth_url = err.split("\n")[1].strip()
|
||||
started[0].receive_callback(_state_of(auth_url), "xyz")
|
||||
return
|
||||
threading.Event().wait(0.05)
|
||||
|
||||
t = threading.Thread(target=deliver)
|
||||
t.start()
|
||||
p = oauth_provider.MCPOAuthProvider("n", "https://n.example/mcp", storage=storage)
|
||||
assert p._run_pkce_flow(METADATA, {"client_id": "c"}, None, timeout=5) == "tok-xyz"
|
||||
t.join(timeout=5)
|
||||
assert oauth_provider._callback_server is None
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Tests for A2A message/stream lifecycle (issue #6151)."""
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.ui.a2a.streaming import stream_agent_response
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TaskState, TextPart
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _restore_event_loop():
|
||||
# asyncio.run() leaves no current loop; later tests using
|
||||
# get_event_loop() would break.
|
||||
yield
|
||||
asyncio.set_event_loop(asyncio.new_event_loop())
|
||||
|
||||
|
||||
def _new_task(store):
|
||||
msg = Message(message_id="m1", role=Role.USER, parts=[TextPart(text="hi")])
|
||||
return store.create_task(msg)
|
||||
|
||||
|
||||
async def _consume(agen):
|
||||
return [e async for e in agen]
|
||||
|
||||
|
||||
class _Slow:
|
||||
def chat(self, x):
|
||||
time.sleep(0.5)
|
||||
return "done"
|
||||
|
||||
|
||||
class _Team:
|
||||
def start(self, x):
|
||||
return f"team:{x}"
|
||||
|
||||
|
||||
class _Boom:
|
||||
def chat(self, x):
|
||||
raise ValueError("kaboom")
|
||||
|
||||
|
||||
def test_stream_does_not_block_event_loop():
|
||||
async def main():
|
||||
ticks = 0
|
||||
|
||||
async def heartbeat():
|
||||
nonlocal ticks
|
||||
while True:
|
||||
await asyncio.sleep(0.02)
|
||||
ticks += 1
|
||||
|
||||
hb = asyncio.create_task(heartbeat())
|
||||
t = time.time()
|
||||
await asyncio.gather(
|
||||
_consume(stream_agent_response(_Slow(), "hi", "a")),
|
||||
_consume(stream_agent_response(_Slow(), "hi", "b")),
|
||||
)
|
||||
elapsed = time.time() - t
|
||||
hb.cancel()
|
||||
return elapsed, ticks
|
||||
|
||||
elapsed, ticks = asyncio.run(main())
|
||||
assert elapsed < 0.9
|
||||
assert ticks > 5
|
||||
|
||||
|
||||
def test_stream_updates_task_store_and_supports_team():
|
||||
store = TaskStore()
|
||||
task = _new_task(store)
|
||||
running = {}
|
||||
events = asyncio.run(_consume(stream_agent_response(
|
||||
_Team(), "x", task.id, task_store=store, running=running)))
|
||||
assert any("team:x" in e for e in events)
|
||||
assert any('"completed"' in e for e in events)
|
||||
t = store.get_task(task.id)
|
||||
assert t.status.state == TaskState.COMPLETED
|
||||
assert t.artifacts and running == {}
|
||||
assert task.id not in store._pinned
|
||||
|
||||
|
||||
def test_stream_failure_reports_reason():
|
||||
store = TaskStore()
|
||||
task = _new_task(store)
|
||||
events = asyncio.run(_consume(stream_agent_response(
|
||||
_Boom(), "x", task.id, task_store=store)))
|
||||
final = [e for e in events if '"failed"' in e]
|
||||
assert final and "kaboom" in final[0]
|
||||
assert store.get_task(task.id).status.state == TaskState.FAILED
|
||||
|
||||
|
||||
def test_stream_cancel_via_running_handle():
|
||||
store = TaskStore()
|
||||
task = _new_task(store)
|
||||
running = {}
|
||||
|
||||
async def main():
|
||||
consumer = asyncio.create_task(_consume(stream_agent_response(
|
||||
_Slow(), "x", task.id, task_store=store, running=running)))
|
||||
while task.id not in running:
|
||||
await asyncio.sleep(0.01)
|
||||
# Mirrors tasks/cancel: mark cancelled, then cancel the handle
|
||||
store.cancel_task(task.id)
|
||||
running[task.id].cancel()
|
||||
return await consumer
|
||||
|
||||
events = asyncio.run(main())
|
||||
assert any('"canceled"' in e or '"cancelled"' in e for e in events)
|
||||
assert store.get_task(task.id).status.state == TaskState.CANCELLED
|
||||
assert running == {}
|
||||
|
||||
|
||||
class _Counting:
|
||||
calls = 0
|
||||
|
||||
def chat(self, x):
|
||||
_Counting.calls += 1
|
||||
return "ran"
|
||||
|
||||
|
||||
def test_cancel_before_call_registered_skips_agent():
|
||||
store = TaskStore()
|
||||
task = _new_task(store)
|
||||
running = {}
|
||||
|
||||
async def main():
|
||||
agen = stream_agent_response(
|
||||
_Counting(), "x", task.id, task_store=store, running=running)
|
||||
first = await agen.__anext__()
|
||||
assert '"working"' in first
|
||||
# tasks/cancel arrives before the call is registered
|
||||
store.cancel_task(task.id)
|
||||
return [first] + [e async for e in agen]
|
||||
|
||||
events = asyncio.run(main())
|
||||
assert _Counting.calls == 0
|
||||
assert not any('"completed"' in e for e in events)
|
||||
assert store.get_task(task.id).status.state == TaskState.CANCELLED
|
||||
assert running == {} and task.id not in store._pinned
|
||||
|
||||
|
||||
def test_early_close_leaves_terminal_state_and_unpins():
|
||||
store = TaskStore()
|
||||
task = _new_task(store)
|
||||
|
||||
async def main():
|
||||
agen = stream_agent_response(_Team(), "x", task.id, task_store=store)
|
||||
await agen.__anext__()
|
||||
await agen.aclose()
|
||||
|
||||
asyncio.run(main())
|
||||
assert store.get_task(task.id).status.state == TaskState.CANCELLED
|
||||
assert task.id not in store._pinned
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Regression tests for #6152: async APIs must not block the event loop."""
|
||||
import asyncio
|
||||
import contextvars
|
||||
import time
|
||||
|
||||
from praisonaiagents.escalation.pipeline import EscalationPipeline
|
||||
from praisonaiagents.escalation.types import EscalationStage, StageContext
|
||||
from praisonaiagents.knowledge.query_engine import SubQuestionEngine
|
||||
from praisonaiagents.rag.pipeline import RAG
|
||||
|
||||
BLOCK = 0.3
|
||||
_request_var = contextvars.ContextVar("request_var", default=None)
|
||||
|
||||
|
||||
async def _heartbeat_ticks(coro):
|
||||
ticks = 0
|
||||
done = False
|
||||
|
||||
async def beat():
|
||||
nonlocal ticks
|
||||
while not done:
|
||||
await asyncio.sleep(0.02)
|
||||
ticks += 1
|
||||
|
||||
task = asyncio.create_task(beat())
|
||||
try:
|
||||
result = await coro
|
||||
finally:
|
||||
done = True
|
||||
await task
|
||||
return result, ticks
|
||||
|
||||
|
||||
class _BlockingAgent:
|
||||
def __init__(self):
|
||||
self.tools = ["a", "b", "c"]
|
||||
self.seen_tools = []
|
||||
|
||||
def chat(self, prompt, tools=None):
|
||||
self.seen_tools.append(tools)
|
||||
time.sleep(BLOCK)
|
||||
return "ok"
|
||||
|
||||
|
||||
def test_escalation_does_not_block_loop_or_mutate_agent_tools():
|
||||
agent = _BlockingAgent()
|
||||
pipeline = EscalationPipeline(agent=agent)
|
||||
pipeline._current_context = StageContext(stage=EscalationStage.HEURISTIC, prompt="")
|
||||
|
||||
async def run():
|
||||
return await _heartbeat_ticks(
|
||||
asyncio.gather(
|
||||
pipeline._execute_heuristic("hi"),
|
||||
pipeline._execute_heuristic("there"),
|
||||
)
|
||||
)
|
||||
|
||||
(r1, r2), ticks = asyncio.run(run())
|
||||
assert r1["success"] and r2["success"]
|
||||
assert ticks >= 5
|
||||
assert agent.tools == ["a", "b", "c"]
|
||||
assert all(t is not None for t in agent.seen_tools)
|
||||
|
||||
|
||||
class _BlockingRAG(RAG):
|
||||
llm = None # shadow the lazy LLM property
|
||||
|
||||
def __init__(self): # bypass heavy init
|
||||
from praisonaiagents.rag.models import RAGConfig
|
||||
self.config = RAGConfig()
|
||||
self.seen_var = None
|
||||
|
||||
def _retrieve(self, **kwargs):
|
||||
self.seen_var = _request_var.get()
|
||||
time.sleep(BLOCK)
|
||||
return []
|
||||
|
||||
def _rerank(self, question, results, rerank=None):
|
||||
return results
|
||||
|
||||
def _build_context(self, results):
|
||||
return "", []
|
||||
|
||||
def _build_prompt(self, question, context):
|
||||
return question
|
||||
|
||||
|
||||
def test_rag_astream_offloads_retrieval_and_keeps_contextvars():
|
||||
rag = _BlockingRAG()
|
||||
|
||||
async def consume():
|
||||
_request_var.set("req-1")
|
||||
return [c async for c in rag.astream("q")]
|
||||
|
||||
chunks, ticks = asyncio.run(_heartbeat_ticks(consume()))
|
||||
assert chunks == ["LLM not available."]
|
||||
assert ticks >= 5
|
||||
assert rag.seen_var == "req-1"
|
||||
|
||||
|
||||
def test_subquestion_aquery_offloads():
|
||||
engine = SubQuestionEngine()
|
||||
|
||||
def slow_query(question, context=None, **kwargs):
|
||||
time.sleep(BLOCK)
|
||||
return "done"
|
||||
|
||||
engine.query = slow_query
|
||||
result, ticks = asyncio.run(_heartbeat_ticks(engine.aquery("q")))
|
||||
assert result == "done"
|
||||
assert ticks >= 5
|
||||
|
||||
|
||||
class _SyncStreamLLM:
|
||||
"""Sync-only LLM (no ``astream``) whose ``stream()`` blocks between chunks."""
|
||||
|
||||
def stream(self, messages, **kwargs):
|
||||
for tok in ("a", "b", "c"):
|
||||
time.sleep(BLOCK / 3)
|
||||
yield tok
|
||||
|
||||
|
||||
def test_rag_astream_sync_fallback_offloads_and_retrieves_once():
|
||||
rag = _BlockingRAG()
|
||||
rag.llm = _SyncStreamLLM()
|
||||
calls = []
|
||||
orig_retrieve = rag._retrieve
|
||||
|
||||
def counting_retrieve(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return orig_retrieve(**kwargs)
|
||||
|
||||
rag._retrieve = counting_retrieve
|
||||
|
||||
async def consume():
|
||||
return [c async for c in rag.astream("q")]
|
||||
|
||||
chunks, ticks = asyncio.run(_heartbeat_ticks(consume()))
|
||||
assert chunks == ["a", "b", "c"]
|
||||
assert len(calls) == 1
|
||||
assert ticks >= 10
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Regression tests for issue #6171.
|
||||
|
||||
1. Agent.output_file must not overwrite a hardlinked file outside the root.
|
||||
2. Shadow-git / worktree git calls are bounded by a timeout and non-interactive.
|
||||
3. A torn learn-store archive is quarantined and eviction keeps working.
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_agent_output_file_hardlink_does_not_touch_outside_inode(tmp_path, monkeypatch):
|
||||
from praisonaiagents import Agent, OutputConfig
|
||||
|
||||
root = tmp_path / "root"
|
||||
root.mkdir()
|
||||
outside = tmp_path / "outside.txt"
|
||||
outside.write_text("SECRET")
|
||||
link = root / "out.md"
|
||||
try:
|
||||
os.link(outside, link)
|
||||
except (OSError, NotImplementedError):
|
||||
pytest.skip("hardlinks unsupported")
|
||||
|
||||
monkeypatch.setenv("PRAISONAI_PROJECT_ROOT", str(root))
|
||||
monkeypatch.chdir(root)
|
||||
agent = Agent(name="writer", output=OutputConfig(output_file="out.md"))
|
||||
assert agent._save_output_to_file("AGENT-OVERWRITE") is True
|
||||
|
||||
assert outside.read_text() == "SECRET"
|
||||
assert link.read_text() == "AGENT-OVERWRITE"
|
||||
|
||||
|
||||
def test_agent_output_file_outside_root_rejected(tmp_path, monkeypatch):
|
||||
from praisonaiagents import Agent, OutputConfig
|
||||
|
||||
root = tmp_path / "root"
|
||||
root.mkdir()
|
||||
monkeypatch.setenv("PRAISONAI_PROJECT_ROOT", str(root))
|
||||
agent = Agent(name="writer", output=OutputConfig(output_file=str(tmp_path / "escape.md")))
|
||||
assert agent._save_output_to_file("x") is False
|
||||
assert not (tmp_path / "escape.md").exists()
|
||||
|
||||
|
||||
def test_snapshot_run_git_is_bounded_and_non_interactive(tmp_path):
|
||||
from praisonaiagents.snapshot import snapshot as snap_mod
|
||||
|
||||
fs = snap_mod.FileSnapshot(str(tmp_path), snapshot_dir=str(tmp_path / "snaps"))
|
||||
with mock.patch.object(
|
||||
snap_mod.subprocess, "run",
|
||||
side_effect=subprocess.TimeoutExpired(["git"], snap_mod._GIT_TIMEOUT),
|
||||
) as run:
|
||||
with pytest.raises(subprocess.TimeoutExpired):
|
||||
fs._run_git("status")
|
||||
kwargs = run.call_args.kwargs
|
||||
assert kwargs["timeout"] == snap_mod._GIT_TIMEOUT
|
||||
assert kwargs["env"]["GIT_TERMINAL_PROMPT"] == "0"
|
||||
|
||||
|
||||
def test_git_worktree_timeout_treated_as_unavailable(tmp_path):
|
||||
from praisonaiagents.workspace.adapters import git_worktree as gw
|
||||
|
||||
with mock.patch(
|
||||
"subprocess.run",
|
||||
side_effect=subprocess.TimeoutExpired(["git"], gw._GIT_TIMEOUT),
|
||||
) as run:
|
||||
adapter = gw.GitWorktreeAdapter(tmp_path)
|
||||
assert adapter.available is False
|
||||
kwargs = run.call_args.kwargs
|
||||
assert kwargs["timeout"] == gw._GIT_TIMEOUT
|
||||
assert kwargs["env"]["GIT_TERMINAL_PROMPT"] == "0"
|
||||
|
||||
|
||||
def test_corrupt_archive_quarantined_and_eviction_continues(tmp_path):
|
||||
from praisonaiagents.memory.learn.stores import InsightStore
|
||||
|
||||
st = InsightStore(store_path=str(tmp_path / "insights.json"), max_entries=3)
|
||||
for i in range(5):
|
||||
st.add(f"insight {i}")
|
||||
archive = st._archive_path()
|
||||
archive.write_text('[{"id": "x", "cont') # torn write
|
||||
|
||||
for i in range(5, 12):
|
||||
st.add(f"insight {i}")
|
||||
|
||||
assert len(st._entries) == 3
|
||||
sidecars = list(tmp_path.glob("insights.archive.json.*.corrupt"))
|
||||
assert sidecars and sidecars[0].read_text() == '[{"id": "x", "cont'
|
||||
import json
|
||||
assert isinstance(json.loads(archive.read_text()), list)
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Regression guard (#6167): opt-in modules must not load on ``import Agent``."""
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
_SCRIPT = r"""
|
||||
import sys
|
||||
from praisonaiagents import Agent # noqa: F401
|
||||
import praisonaiagents.session as session
|
||||
|
||||
lazy = (
|
||||
"praisonaiagents.agent.server_registry",
|
||||
"praisonaiagents.goal.judge",
|
||||
"praisonaiagents.session.encrypted_store",
|
||||
)
|
||||
loaded = [m for m in lazy if m in sys.modules]
|
||||
assert not loaded, f"eagerly imported: {loaded}"
|
||||
|
||||
# Public encryption exports still resolve lazily.
|
||||
from praisonaiagents.session import ( # noqa: E402
|
||||
EncryptedSessionStore, SessionEncryptionError, generate_session_key,
|
||||
)
|
||||
assert "praisonaiagents.session.encrypted_store" in sys.modules
|
||||
assert session.EncryptedSessionStore is EncryptedSessionStore
|
||||
assert callable(generate_session_key) and issubclass(SessionEncryptionError, Exception)
|
||||
"""
|
||||
|
||||
|
||||
def test_opt_in_modules_not_loaded_on_agent_import():
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", _SCRIPT], capture_output=True, text=True, timeout=60
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Regression test for #6202: capping a tool batch must not leave orphaned tool_calls."""
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("litellm")
|
||||
|
||||
|
||||
def _litellm():
|
||||
# Resolve at test time: other tests may reload/replace ``litellm`` in
|
||||
# sys.modules after collection, and llm.py imports it lazily per call.
|
||||
import litellm
|
||||
return litellm
|
||||
|
||||
|
||||
def ModelResponse(**kw): # noqa: N802
|
||||
return _litellm().ModelResponse(**kw)
|
||||
|
||||
|
||||
def _tc(i):
|
||||
return {"id": f"c{i}", "type": "function",
|
||||
"function": {"name": "ping", "arguments": json.dumps({"n": i})}}
|
||||
|
||||
|
||||
def _ping(n: int) -> str:
|
||||
"ping"
|
||||
return f"pong{n}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["sync", "async"])
|
||||
def test_truncated_batch_has_no_orphaned_tool_calls(monkeypatch, mode):
|
||||
sent = []
|
||||
|
||||
def _respond(kwargs):
|
||||
sent.append(json.loads(json.dumps(kwargs["messages"], default=str)))
|
||||
if len(sent) == 1:
|
||||
return ModelResponse(choices=[{"message": {"role": "assistant", "content": "",
|
||||
"tool_calls": [_tc(i) for i in range(12)]},
|
||||
"finish_reason": "tool_calls", "index": 0}])
|
||||
return ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"},
|
||||
"finish_reason": "stop", "index": 0}])
|
||||
|
||||
async def _afake(*a, **k):
|
||||
return _respond(k)
|
||||
|
||||
monkeypatch.setattr(_litellm(), "completion", lambda *a, **k: _respond(k))
|
||||
monkeypatch.setattr(_litellm(), "acompletion", _afake)
|
||||
|
||||
from praisonaiagents import Agent
|
||||
agent = Agent(instructions="x", model="anthropic/claude-3-haiku-20240307",
|
||||
tools=[_ping], output="silent")
|
||||
result = agent.start("go") if mode == "sync" else asyncio.run(agent.astart("go"))
|
||||
|
||||
assert result == "done"
|
||||
follow_up = sent[1]
|
||||
declared = [t["id"] for m in follow_up if m.get("tool_calls") for t in m["tool_calls"]]
|
||||
replied = [m["tool_call_id"] for m in follow_up if m.get("role") == "tool"]
|
||||
assert sorted(declared) == sorted(replied)
|
||||
assert len(declared) == 10 # default max_tool_calls_per_turn
|
||||
|
||||
|
||||
def test_stream_fallback_keeps_text_sent_with_capped_tool_calls(monkeypatch):
|
||||
"""Text returned alongside tool calls must still be yielded once the cap is hit."""
|
||||
sent = []
|
||||
|
||||
def _respond(*a, **k):
|
||||
sent.append(k)
|
||||
n = len(sent)
|
||||
if n <= 2:
|
||||
return ModelResponse(choices=[{"message": {"role": "assistant", "content": f"text{n}",
|
||||
"tool_calls": [_tc(n)]},
|
||||
"finish_reason": "tool_calls", "index": 0}])
|
||||
return ModelResponse(choices=[{"message": {"role": "assistant", "content": "final"},
|
||||
"finish_reason": "stop", "index": 0}])
|
||||
|
||||
monkeypatch.setattr(_litellm(), "completion", _respond)
|
||||
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
llm = LLM(model="anthropic/claude-3-haiku-20240307")
|
||||
monkeypatch.setattr(llm, "_supports_streaming_tools", lambda: False)
|
||||
chunks = list(llm.get_response_stream(
|
||||
"go", tools=[_ping], execute_tool_fn=lambda name, args: _ping(**args),
|
||||
max_tool_calls_per_turn=1))
|
||||
|
||||
assert "text1" in chunks and "text2" in chunks # text2 arrived at the cap
|
||||
@@ -81,6 +81,111 @@ class TestExactMatchBackwardCompat:
|
||||
assert _read(p) == "x = 2\nx = 2\n"
|
||||
|
||||
|
||||
class TestNoMatchSelfCorrection:
|
||||
def test_retried_edit_shows_applied_line_in_hint(self, tools, tmp_path):
|
||||
# A retried edit is never reported as success; the hint shows the
|
||||
# already-edited line so the agent can see the change landed.
|
||||
p = tmp_path / "a.py"
|
||||
_write(p, "def foo():\n return 1\n")
|
||||
assert "Success" in tools.edit_file(str(p), "return 1", "return 2")
|
||||
before = _read(p)
|
||||
result = tools.edit_file(str(p), "return 1", "return 2")
|
||||
assert "String not found" in result and "No change needed" not in result
|
||||
assert " 2| return 2" in result
|
||||
assert _read(p) == before
|
||||
|
||||
def test_tied_candidate_elsewhere_is_not_reported_applied(self, tools, tmp_path):
|
||||
# Merge-gate case: g() already has `return 2` and ties with f()'s
|
||||
# `return 3`; the edit meant for f() must not report success.
|
||||
p = tmp_path / "a.py"
|
||||
src = "def g():\n return 2\n\ndef f():\n return 3\n"
|
||||
_write(p, src)
|
||||
result = tools.edit_file(str(p), "return 1", "return 2")
|
||||
assert "No change needed" not in result
|
||||
assert "String not found" in result
|
||||
assert _read(p) == src
|
||||
|
||||
def test_short_common_new_string_is_not_noop(self, tools, tmp_path):
|
||||
# new_string occurring more than once is not evidence of a prior edit.
|
||||
p = tmp_path / "a.py"
|
||||
_write(p, "x = 1\ny = x\n")
|
||||
result = tools.edit_file(str(p), "z = 1", "x")
|
||||
assert "No change needed" not in result
|
||||
|
||||
def test_new_inside_old_is_not_noop(self, tools, tmp_path):
|
||||
# new_string is a substring of old_string, so its presence proves nothing.
|
||||
p = tmp_path / "a.py"
|
||||
_write(p, "y = compute(a)\n")
|
||||
result = tools.edit_file(str(p), "x = compute(a)", "compute(a)")
|
||||
assert "Error" in result and "No change needed" not in result
|
||||
|
||||
def test_unrelated_new_string_elsewhere_is_not_noop(self, tools, tmp_path):
|
||||
# new_string exists in another function, not where old_string was expected.
|
||||
p = tmp_path / "a.py"
|
||||
_write(p, "def f(x):\n return x * 2\n\ndef g():\n return 0\n")
|
||||
result = tools.edit_file(str(p), "return x * 3", "return 0")
|
||||
assert "No change needed" not in result
|
||||
assert "String not found" in result
|
||||
assert " 2|" in result
|
||||
|
||||
def test_hint_truncates_long_lines(self, tools, tmp_path):
|
||||
p = tmp_path / "a.json"
|
||||
_write(p, '{"key": "' + "a" * 1000 + '"}\n')
|
||||
result = tools.edit_file(str(p), '{"key": "' + "a" * 900 + 'b"}', "x")
|
||||
assert "Closest match" in result and "[truncated]" in result
|
||||
assert len(result) < 1000
|
||||
|
||||
def test_hint_search_is_bounded_on_long_repetitive_lines(self, tools, tmp_path):
|
||||
import time
|
||||
p = tmp_path / "a.json"
|
||||
_write(p, "\n".join('{"k": "' + "ab" * 600 + '"}' for _ in range(500)) + "\n")
|
||||
start = time.monotonic()
|
||||
result = tools.edit_file(str(p), '{"k": "' + "ab" * 550 + 'c"}', "x")
|
||||
assert "String not found" in result
|
||||
assert time.monotonic() - start < 5
|
||||
|
||||
def test_hint_prefilter_is_budgeted(self):
|
||||
# Every window passes real_quick_ratio but fails quick_ratio (a
|
||||
# pure-Python pass); that prefilter work must count against the budget
|
||||
# (unbudgeted this took ~4s on a 20k-line file).
|
||||
import random
|
||||
import string
|
||||
import time
|
||||
rnd = random.Random(0)
|
||||
line = lambda alpha: "".join(rnd.choice(alpha) for _ in range(20))
|
||||
content = "\n".join(line(string.ascii_lowercase) for _ in range(19999))
|
||||
old = "\n".join(line("abcdefghijkl") for _ in range(95))
|
||||
start = time.monotonic()
|
||||
EditTools._closest_region(content, old)
|
||||
assert time.monotonic() - start < 2
|
||||
|
||||
def test_closest_region_hint_with_line_numbers(self, tools, tmp_path):
|
||||
p = tmp_path / "a.py"
|
||||
_write(p, "import os\n\ndef compute(x):\n total = x * 2\n return total\n")
|
||||
result = tools.edit_file(str(p), "def compute(y):\n total = y * 3", "x")
|
||||
assert "String not found" in result
|
||||
assert "Closest match" in result
|
||||
assert "lines 3-4" in result
|
||||
assert " 3| def compute(x):" in result
|
||||
|
||||
def test_no_hint_when_nothing_similar(self, tools, tmp_path):
|
||||
p = tmp_path / "a.py"
|
||||
_write(p, "hello world\n")
|
||||
result = tools.edit_file(str(p), "zzzzqqqq", "x")
|
||||
assert "String not found" in result
|
||||
assert "Closest match" not in result
|
||||
|
||||
def test_apply_patch_hunk_hint(self, tools, tmp_path):
|
||||
p = tmp_path / "a.py"
|
||||
_write(p, "def compute(x):\n return x * 2\n")
|
||||
patch = (f"*** Update File: {p}\n@@\n"
|
||||
"def compute(y):\n return y * 2\n===\n"
|
||||
"def compute(y):\n return y * 3\n")
|
||||
result = tools.apply_patch(patch)
|
||||
assert "Hunk not found" in result
|
||||
assert " 1| def compute(x):" in result
|
||||
|
||||
|
||||
class TestFuzzyMatching:
|
||||
def test_line_trimmed_indentation(self, tools, tmp_path):
|
||||
p = tmp_path / "a.py"
|
||||
|
||||
@@ -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"
|
||||
Generated
+1
-1
@@ -3162,7 +3162,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.20"
|
||||
version = "1.7.21"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.19"
|
||||
__version__ = "0.1.20"
|
||||
@@ -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)
|
||||
|
||||
@@ -23,7 +23,7 @@ wrapper's delivery stack, neither of which belong in the core SDK.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
from praisonaiagents.gateway import (
|
||||
DeliveryResult,
|
||||
@@ -61,6 +61,13 @@ class BotOutboundMessenger:
|
||||
media_policy: Optional :class:`OutboundMediaPolicy` governing which
|
||||
agent-supplied file paths may be uploaded (exfiltration guard).
|
||||
Defaults to a permissive-but-denylisted policy.
|
||||
liveness: Optional zero-arg predicate returning ``True`` while the
|
||||
originating turn is still live (same contract as
|
||||
``ApprovalRequest.liveness``). When it reports stale, every
|
||||
outbound effect (send/react/create_thread/edit/delete) is dropped
|
||||
fail-closed, so a stopped or superseded turn — or a retained
|
||||
handle to its messenger — cannot reach a user channel. ``None``
|
||||
(default) keeps today's unbound behaviour.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -69,10 +76,35 @@ class BotOutboundMessenger:
|
||||
*,
|
||||
origin: Optional[SessionSource] = None,
|
||||
media_policy: Optional[OutboundMediaPolicy] = None,
|
||||
liveness: Optional[Callable[[], bool]] = None,
|
||||
) -> None:
|
||||
self._router = router
|
||||
self._origin = origin
|
||||
self._media_policy = media_policy or OutboundMediaPolicy()
|
||||
self._liveness = liveness
|
||||
|
||||
_STALE_DETAIL = "turn no longer live (stopped or superseded)"
|
||||
|
||||
def _is_stale(self) -> bool:
|
||||
"""Whether the originating turn was stopped/superseded (#6160).
|
||||
|
||||
Mirrors ``ApprovalRegistry._reject_if_stale``: ``None`` liveness means
|
||||
always live, and a probe that raises is treated as live (fail-open on
|
||||
the *probe*, never on the effect) so a buggy predicate cannot wedge
|
||||
delivery.
|
||||
"""
|
||||
live = self._liveness
|
||||
if live is None:
|
||||
return False
|
||||
try:
|
||||
stale = not live()
|
||||
except Exception: # noqa: BLE001 - a liveness probe must never crash
|
||||
return False
|
||||
if stale:
|
||||
logger.warning(
|
||||
"BotOutboundMessenger: dropping outbound effect from a stale turn"
|
||||
)
|
||||
return stale
|
||||
|
||||
async def send(
|
||||
self,
|
||||
@@ -90,6 +122,13 @@ class BotOutboundMessenger:
|
||||
platform adapter. The result truthfully reports how many attachments
|
||||
were delivered, skipped by policy, or unsupported by the transport.
|
||||
"""
|
||||
if self._is_stale():
|
||||
return DeliveryResult(
|
||||
ok=False,
|
||||
target=target,
|
||||
summary=f"Not sent to {target}: {self._STALE_DETAIL}.",
|
||||
detail=self._STALE_DETAIL,
|
||||
)
|
||||
try:
|
||||
platform, channel_id, _thread_id = self._router.resolve(target, self._origin)
|
||||
except ValueError as e:
|
||||
@@ -244,7 +283,18 @@ class BotOutboundMessenger:
|
||||
# Only the first attachment carries the caption to avoid repeating the
|
||||
# full text on every file.
|
||||
first = True
|
||||
for path in media:
|
||||
for index, path in enumerate(media):
|
||||
# Re-check liveness before every upload (#6160): /stop only sets
|
||||
# the controller flag, it does not cancel this coroutine, so a turn
|
||||
# stopped mid-send must not keep uploading the remaining files.
|
||||
if self._is_stale():
|
||||
remaining = len(media) - index
|
||||
skipped += remaining
|
||||
notes.append(
|
||||
f"Stopped: {remaining} attachment(s) not sent, "
|
||||
f"{self._STALE_DETAIL}."
|
||||
)
|
||||
break
|
||||
try:
|
||||
safe = validate_media_delivery_path(
|
||||
path, policy=self._media_policy
|
||||
@@ -294,6 +344,10 @@ class BotOutboundMessenger:
|
||||
message with a single emoji instead of a full text reply. A channel that
|
||||
cannot react returns a typed ``unsupported`` outcome rather than raising.
|
||||
"""
|
||||
if self._is_stale():
|
||||
return ReactionResult(
|
||||
status="failed", target=target, detail=self._STALE_DETAIL
|
||||
)
|
||||
resolved_message_id = message_id
|
||||
if not resolved_message_id and target == "origin" and self._origin is not None:
|
||||
resolved_message_id = self._origin.message_id or ""
|
||||
@@ -335,6 +389,10 @@ class BotOutboundMessenger:
|
||||
target = f"slack:C123:{res.thread_id}" if res.ok else "slack:C123"
|
||||
await messenger.send(target, subtask_output)
|
||||
"""
|
||||
if self._is_stale():
|
||||
return ThreadResult(
|
||||
status="failed", target=target, detail=self._STALE_DETAIL
|
||||
)
|
||||
status, resolved, thread_id = await self._router.create_thread(
|
||||
target,
|
||||
name,
|
||||
@@ -367,6 +425,10 @@ class BotOutboundMessenger:
|
||||
cannot edit returns a typed ``unsupported`` outcome rather than raising,
|
||||
so the agent can fall back to sending a follow-up.
|
||||
"""
|
||||
if self._is_stale():
|
||||
return MessageActionResult(
|
||||
status="failed", target=target, detail=self._STALE_DETAIL
|
||||
)
|
||||
status, resolved = await self._router.edit(
|
||||
target, message_id, text, self._origin
|
||||
)
|
||||
@@ -388,6 +450,10 @@ class BotOutboundMessenger:
|
||||
channel that cannot delete returns a typed ``unsupported`` outcome
|
||||
rather than raising.
|
||||
"""
|
||||
if self._is_stale():
|
||||
return MessageActionResult(
|
||||
status="failed", target=target, detail=self._STALE_DETAIL
|
||||
)
|
||||
status, resolved = await self._router.delete(
|
||||
target, message_id, self._origin
|
||||
)
|
||||
|
||||
@@ -1342,8 +1342,24 @@ class BotSessionManager:
|
||||
thread_id=thread_id,
|
||||
message_id=message_id or None,
|
||||
)
|
||||
# Issue #6160: bind outbound effects to this turn's liveness.
|
||||
# The run-control InterruptController is created fresh per run
|
||||
# and set on /stop or supersede, so a stopped/superseded turn
|
||||
# (or a retained handle to this messenger) is dropped
|
||||
# fail-closed. A normally completed turn never sets it, so
|
||||
# post-turn background deliveries keep working. No controller
|
||||
# (run control off) -> unbound, today's behaviour.
|
||||
from praisonaiagents.agent.interrupt import InterruptController
|
||||
|
||||
_ic = getattr(agent, "interrupt_controller", None)
|
||||
messenger = BotOutboundMessenger(
|
||||
self._delivery_router, origin=origin_source
|
||||
self._delivery_router,
|
||||
origin=origin_source,
|
||||
liveness=(
|
||||
(lambda: not _ic.is_set())
|
||||
if isinstance(_ic, InterruptController)
|
||||
else None
|
||||
),
|
||||
)
|
||||
messenger_token = _register_messenger(messenger)
|
||||
except Exception as e: # pragma: no cover — defensive
|
||||
|
||||
@@ -217,7 +217,7 @@ class StreamingMode(Enum):
|
||||
OFF = "off" # Current behavior (single final message, chunked)
|
||||
DRAFT = "draft" # Send placeholder, edit in place with growing content
|
||||
PROGRESS = "progress" # Show compact status, then replace with final answer
|
||||
AUTO = "auto" # Stream (DRAFT) where the channel can edit, else OFF
|
||||
AUTO = "auto" # DRAFT where the channel can edit, else status text, else OFF
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -401,6 +401,24 @@ class DraftStreamer:
|
||||
if self._edit_rate_limit > 0:
|
||||
self._current_min_interval = max(self._current_min_interval, self._edit_rate_limit)
|
||||
|
||||
# Edit-less channels that declare ``supports_status_text`` (and back it
|
||||
# with ``set_status_text``) get the curated activity line as a textual
|
||||
# status instead of silence: degrade edit -> status-text -> off.
|
||||
pcaps = getattr(adapter, 'platform_capabilities', None)
|
||||
self._status_text_only = bool(
|
||||
not self._can_edit
|
||||
and self._config.mode in (StreamingMode.AUTO, StreamingMode.PROGRESS)
|
||||
and getattr(pcaps, 'supports_status_text', False) is True
|
||||
and callable(getattr(adapter, 'set_status_text', None))
|
||||
)
|
||||
self._last_status_text: Optional[str] = None
|
||||
if self._status_text_only:
|
||||
logger.info(
|
||||
"Channel %s can't edit; routing streaming mode '%s' to status text",
|
||||
channel_id, self._config.mode.value,
|
||||
)
|
||||
self._config = replace(self._config, mode=StreamingMode.PROGRESS)
|
||||
|
||||
# Resolve AUTO to a concrete mode: stream where the channel can edit,
|
||||
# otherwise fall back to OFF. This makes "stream where supported" an
|
||||
# explicit, opt-in choice instead of every platform silently sitting on
|
||||
@@ -419,7 +437,11 @@ class DraftStreamer:
|
||||
|
||||
# Override config if channel doesn't support editing. Logged (not
|
||||
# silent) so operators can see why a channel isn't streaming.
|
||||
if not self._can_edit and self._config.mode != StreamingMode.OFF:
|
||||
if (
|
||||
not self._can_edit
|
||||
and not self._status_text_only
|
||||
and self._config.mode != StreamingMode.OFF
|
||||
):
|
||||
logger.warning(
|
||||
"Channel %s doesn't support live editing; degrading streaming "
|
||||
"mode '%s' to 'off'",
|
||||
@@ -441,6 +463,16 @@ class DraftStreamer:
|
||||
if self._config.mode == StreamingMode.OFF:
|
||||
raise ValueError("Cannot start streaming when mode is OFF")
|
||||
|
||||
if self._status_text_only:
|
||||
# No placeholder message: the status line is the only live surface.
|
||||
self._message_id = ""
|
||||
self._last_edit_time = time.monotonic()
|
||||
try:
|
||||
await self._push_status_text(self._config.placeholder_text)
|
||||
except Exception as e: # noqa: BLE001 - status is advisory; never break the turn
|
||||
self._handle_edit_failure(e)
|
||||
return self._message_id
|
||||
|
||||
if self._rate_limiter:
|
||||
await self._rate_limiter.acquire(self._channel_id)
|
||||
|
||||
@@ -463,11 +495,33 @@ class DraftStreamer:
|
||||
logger.debug("DraftStreamer started, placeholder sent with message_id=%s", self._message_id)
|
||||
return self._message_id
|
||||
|
||||
async def _push_status_text(self, text: str) -> None:
|
||||
"""Push a textual status line; skips unchanged lines (no API cost).
|
||||
|
||||
Rate-limited like edits. Raises on adapter failure so callers apply the
|
||||
shared flood-control backoff (``_handle_edit_failure``).
|
||||
"""
|
||||
text = redact_outbound_text(text, self._adapter)
|
||||
if text == self._last_status_text:
|
||||
return
|
||||
if self._rate_limiter:
|
||||
await self._rate_limiter.acquire(self._channel_id)
|
||||
await self._adapter.set_status_text(self._channel_id, text)
|
||||
self._last_status_text = text
|
||||
|
||||
@property
|
||||
def _active(self) -> bool:
|
||||
"""Whether a live surface (placeholder or status line) is in use."""
|
||||
return self._message_id is not None and (
|
||||
bool(self._message_id) or self._status_text_only
|
||||
)
|
||||
|
||||
@property
|
||||
def _feed_style(self) -> bool:
|
||||
"""Whether the multi-line rolling progress feed is enabled."""
|
||||
return (
|
||||
self._config.mode == StreamingMode.PROGRESS
|
||||
not self._status_text_only
|
||||
and self._config.mode == StreamingMode.PROGRESS
|
||||
and self._config.progress_style == "feed"
|
||||
)
|
||||
|
||||
@@ -485,7 +539,7 @@ class DraftStreamer:
|
||||
|
||||
async def on_event(self, event: "StreamEvent") -> None:
|
||||
"""Stream event callback to handle progressive updates."""
|
||||
if self._is_finalized or not self._message_id:
|
||||
if self._is_finalized or not self._active:
|
||||
return
|
||||
|
||||
# Import here to avoid circular imports
|
||||
@@ -624,6 +678,8 @@ class DraftStreamer:
|
||||
content = self._config.placeholder_text
|
||||
elif self._current_tool:
|
||||
content = self._activity_line(self._current_tool)
|
||||
elif self._status_text_only:
|
||||
content = self._config.placeholder_text
|
||||
elif buffer:
|
||||
content = buffer
|
||||
else:
|
||||
@@ -645,20 +701,24 @@ class DraftStreamer:
|
||||
|
||||
async def _perform_update(self) -> None:
|
||||
"""Perform the actual message edit with flood-control handling."""
|
||||
if not self._message_id or self._is_finalized or not self._progressive:
|
||||
if not self._active or self._is_finalized or not self._progressive:
|
||||
return
|
||||
|
||||
content = self._render_content()
|
||||
if content is None:
|
||||
return
|
||||
|
||||
|
||||
try:
|
||||
# Apply rate limiting
|
||||
if self._rate_limiter:
|
||||
await self._rate_limiter.acquire(self._channel_id)
|
||||
|
||||
# Perform edit
|
||||
await self._adapter.edit_message(self._channel_id, self._message_id, content)
|
||||
if self._status_text_only:
|
||||
# Same rate limiting and flood backoff as edits.
|
||||
await self._push_status_text(content)
|
||||
else:
|
||||
# Apply rate limiting
|
||||
if self._rate_limiter:
|
||||
await self._rate_limiter.acquire(self._channel_id)
|
||||
|
||||
# Perform edit
|
||||
await self._adapter.edit_message(self._channel_id, self._message_id, content)
|
||||
|
||||
# Success - reset failure streak and update state
|
||||
self._fail_streak = 0
|
||||
@@ -709,7 +769,7 @@ class DraftStreamer:
|
||||
|
||||
async def finalize(self, final_content: str) -> None:
|
||||
"""Finalize with the complete response content."""
|
||||
if self._is_finalized or not self._message_id:
|
||||
if self._is_finalized or not self._active:
|
||||
return
|
||||
|
||||
self._is_finalized = True
|
||||
@@ -752,11 +812,12 @@ class DraftStreamer:
|
||||
# (to avoid leaving a stale "🤔 Thinking..." draft) and fall back to a
|
||||
# new send below.
|
||||
delivered = False
|
||||
try:
|
||||
await self._adapter.edit_message(self._channel_id, self._message_id, chunks[0])
|
||||
delivered = True
|
||||
except Exception as e: # noqa: BLE001 - adapter boundary; preserve final delivery
|
||||
logger.warning("DraftStreamer final edit failed: %s", e)
|
||||
if not self._status_text_only: # status-text mode has no placeholder to edit
|
||||
try:
|
||||
await self._adapter.edit_message(self._channel_id, self._message_id, chunks[0])
|
||||
delivered = True
|
||||
except Exception as e: # noqa: BLE001 - adapter boundary; preserve final delivery
|
||||
logger.warning("DraftStreamer final edit failed: %s", e)
|
||||
|
||||
# Fallback: if the edit failed (or progressive editing was disabled by
|
||||
# flooding), deliver the full answer as a fresh send so the user ALWAYS
|
||||
@@ -782,6 +843,15 @@ class DraftStreamer:
|
||||
delivered = False
|
||||
logger.warning("DraftStreamer follow-on chunk send failed: %s", chunk_exc)
|
||||
|
||||
# Clear the status line so a finished turn no longer looks busy. Runs
|
||||
# after delivery so it can never delay or block the answer.
|
||||
if self._status_text_only and self._last_status_text:
|
||||
try:
|
||||
await self._adapter.set_status_text(self._channel_id, "")
|
||||
self._last_status_text = ""
|
||||
except Exception as e: # noqa: BLE001 - status is advisory
|
||||
logger.debug("DraftStreamer status clear failed: %s", e)
|
||||
|
||||
logger.debug(
|
||||
"DraftStreamer finalized message %s, final_length=%d, progressive=%s, delivered=%s",
|
||||
self._message_id, len(final_content), self._progressive, delivered
|
||||
@@ -790,7 +860,7 @@ class DraftStreamer:
|
||||
@property
|
||||
def is_streaming(self) -> bool:
|
||||
"""Check if streaming is active."""
|
||||
return bool(self._message_id and not self._is_finalized)
|
||||
return bool(self._active and not self._is_finalized)
|
||||
|
||||
@property
|
||||
def has_content(self) -> bool:
|
||||
|
||||
@@ -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()
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-bot"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
description = "Bots, gateway, and channel CLI for PraisonAI — messaging platforms and WebSocket control plane extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
||||
@@ -808,3 +808,140 @@ class TestBotOSWiring:
|
||||
|
||||
os._wire_outbound_messenger()
|
||||
assert session._delivery_router is os._delivery_router
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Issue #6160 — outbound effects are bound to the originating turn's liveness
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTurnLivenessBinding:
|
||||
"""A stopped/superseded turn cannot deliver outbound effects."""
|
||||
|
||||
def _origin(self):
|
||||
return SessionSource(platform="telegram", channel_id="123", message_id="m1")
|
||||
|
||||
def test_stale_turn_drops_send_fail_closed(self):
|
||||
router, sent = _make_router()
|
||||
messenger = BotOutboundMessenger(
|
||||
router, origin=self._origin(), liveness=lambda: False
|
||||
)
|
||||
result = asyncio.run(messenger.send("origin", "late"))
|
||||
assert result.ok is False
|
||||
assert "no longer live" in (result.detail or "")
|
||||
assert sent == []
|
||||
|
||||
def test_stale_turn_drops_react_thread_edit_delete(self):
|
||||
router, _ = _make_router()
|
||||
router.react = MagicMock(side_effect=AssertionError("must not dispatch"))
|
||||
router.create_thread = MagicMock(side_effect=AssertionError("must not dispatch"))
|
||||
router.edit = MagicMock(side_effect=AssertionError("must not dispatch"))
|
||||
router.delete = MagicMock(side_effect=AssertionError("must not dispatch"))
|
||||
messenger = BotOutboundMessenger(
|
||||
router, origin=self._origin(), liveness=lambda: False
|
||||
)
|
||||
|
||||
async def run():
|
||||
return (
|
||||
await messenger.react("origin", "👍"),
|
||||
await messenger.create_thread("origin", "t"),
|
||||
await messenger.edit("origin", "m1", "x"),
|
||||
await messenger.delete("origin", "m1"),
|
||||
)
|
||||
|
||||
for res in asyncio.run(run()):
|
||||
assert res.status == "failed"
|
||||
assert "no longer live" in (res.detail or "")
|
||||
|
||||
def test_live_turn_and_unbound_deliver(self):
|
||||
router, sent = _make_router()
|
||||
live = BotOutboundMessenger(router, origin=self._origin(), liveness=lambda: True)
|
||||
unbound = BotOutboundMessenger(router, origin=self._origin())
|
||||
assert asyncio.run(live.send("origin", "a")).ok is True
|
||||
assert asyncio.run(unbound.send("origin", "b")).ok is True
|
||||
assert sent == [("123", "a"), ("123", "b")]
|
||||
|
||||
def test_raising_probe_is_treated_as_live(self):
|
||||
router, sent = _make_router()
|
||||
|
||||
def boom():
|
||||
raise RuntimeError("probe bug")
|
||||
|
||||
messenger = BotOutboundMessenger(router, origin=self._origin(), liveness=boom)
|
||||
assert asyncio.run(messenger.send("origin", "x")).ok is True
|
||||
assert sent == [("123", "x")]
|
||||
|
||||
def test_session_binds_retained_messenger_to_turn_interrupt(self):
|
||||
from praisonaiagents.agent.interrupt import InterruptController
|
||||
|
||||
router, sent = _make_router()
|
||||
agent = MagicMock(name="agent")
|
||||
agent.name = "Test"
|
||||
agent.chat_history = []
|
||||
ctrl = InterruptController()
|
||||
agent.interrupt_controller = ctrl
|
||||
captured = {}
|
||||
|
||||
def fake_chat(prompt, *args, **kwargs):
|
||||
captured["messenger"] = get_outbound_messenger()
|
||||
return "ok"
|
||||
|
||||
agent.chat.side_effect = fake_chat
|
||||
mgr = BotSessionManager(platform="telegram", delivery_router=router)
|
||||
asyncio.run(mgr.chat(agent, user_id="u1", prompt="hi", chat_id="123"))
|
||||
|
||||
retained = captured["messenger"]
|
||||
# Normal completion: a retained handle (background worker) still delivers.
|
||||
assert asyncio.run(retained.send("origin", "bg")).ok is True
|
||||
# Turn stopped/superseded: the same retained handle is now fail-closed.
|
||||
ctrl.request("user_stop_command")
|
||||
assert asyncio.run(retained.send("origin", "late")).ok is False
|
||||
assert sent == [("123", "bg")]
|
||||
|
||||
|
||||
def test_stop_mid_send_halts_remaining_attachment_uploads(tmp_path):
|
||||
# Issue #6160 (review): /stop only sets the controller flag; it does not
|
||||
# cancel an in-flight send(). Remaining uploads must re-check liveness.
|
||||
uploaded = []
|
||||
state = {"live": True}
|
||||
|
||||
class FakeBot:
|
||||
platform = "telegram"
|
||||
|
||||
async def send_message(self, channel_id, text):
|
||||
pass
|
||||
|
||||
async def send_media(self, channel_id, path, caption=None):
|
||||
uploaded.append(path)
|
||||
state["live"] = False # /stop arrives during the first upload
|
||||
|
||||
fake_bot = FakeBot()
|
||||
|
||||
class FakeBotOS:
|
||||
def get_bot(self, platform):
|
||||
return fake_bot if platform == "telegram" else None
|
||||
|
||||
def list_bots(self):
|
||||
return ["telegram"]
|
||||
|
||||
router = DeliveryRouter(FakeBotOS())
|
||||
router.directory._home_channels = {}
|
||||
router.directory._aliases = {}
|
||||
router.directory._observed = {}
|
||||
|
||||
files = []
|
||||
for name in ("a.png", "b.png", "c.png"):
|
||||
f = tmp_path / name
|
||||
f.write_bytes(b"\x89PNG\r\n\x1a\n" + b"0" * 32)
|
||||
files.append(str(f))
|
||||
|
||||
messenger = BotOutboundMessenger(
|
||||
router,
|
||||
origin=SessionSource(platform="telegram", channel_id="123"),
|
||||
liveness=lambda: state["live"],
|
||||
)
|
||||
result = asyncio.run(messenger.send("origin", "charts", media=files))
|
||||
|
||||
assert uploaded == [files[0]]
|
||||
assert "2 skipped" in result.summary
|
||||
assert "no longer live" in (result.detail or "")
|
||||
@@ -0,0 +1,155 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Tests for the textual status-line surface on edit-less channels (issue #6162).
|
||||
|
||||
A channel that can neither edit messages nor react, but declares
|
||||
``PlatformCapabilities.supports_status_text`` and implements
|
||||
``set_status_text``, gets the curated activity line instead of silence.
|
||||
Default (flag off) behaviour is unchanged: non-edit channels degrade to OFF.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.bots import PlatformCapabilities
|
||||
from praisonaiagents.streaming.events import StreamEvent, StreamEventType
|
||||
from praisonai_bot.bots._streaming import DraftStreamer, StreamingConfig, StreamingMode
|
||||
|
||||
|
||||
class _StatusAdapter:
|
||||
capabilities = {"live_edit": False}
|
||||
|
||||
def __init__(self, supports_status_text=True):
|
||||
self.platform_capabilities = PlatformCapabilities(
|
||||
supports_edit=False, supports_status_text=supports_status_text
|
||||
)
|
||||
self.status = []
|
||||
self.sent = []
|
||||
self.edits = []
|
||||
|
||||
async def set_status_text(self, channel_id, text):
|
||||
self.status.append((channel_id, text))
|
||||
|
||||
async def send_message(self, channel_id, content):
|
||||
self.sent.append(content)
|
||||
return {"message_id": "x"}
|
||||
|
||||
async def edit_message(self, channel_id, message_id, content):
|
||||
self.edits.append(content)
|
||||
|
||||
|
||||
def test_capability_defaults_off_and_round_trips():
|
||||
assert PlatformCapabilities().supports_status_text is False
|
||||
caps = PlatformCapabilities(supports_status_text=True)
|
||||
assert PlatformCapabilities.from_dict(caps.to_dict()).supports_status_text is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", [StreamingMode.AUTO, StreamingMode.PROGRESS])
|
||||
def test_non_edit_channel_without_flag_still_degrades_to_off(mode):
|
||||
s = DraftStreamer(_StatusAdapter(False), "c", StreamingConfig(mode=mode))
|
||||
assert s._config.mode == StreamingMode.OFF
|
||||
assert not s._status_text_only
|
||||
|
||||
|
||||
def test_flag_without_method_degrades_to_off():
|
||||
class _NoMethod:
|
||||
capabilities = {"live_edit": False}
|
||||
platform_capabilities = PlatformCapabilities(supports_status_text=True)
|
||||
|
||||
s = DraftStreamer(_NoMethod(), "c", StreamingConfig(mode=StreamingMode.AUTO))
|
||||
assert s._config.mode == StreamingMode.OFF
|
||||
|
||||
|
||||
def test_draft_mode_is_not_rerouted():
|
||||
s = DraftStreamer(_StatusAdapter(), "c", StreamingConfig(mode=StreamingMode.DRAFT))
|
||||
assert s._config.mode == StreamingMode.OFF
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", [StreamingMode.AUTO, StreamingMode.PROGRESS])
|
||||
def test_status_text_route_end_to_end(mode):
|
||||
adapter = _StatusAdapter()
|
||||
cfg = StreamingConfig(
|
||||
mode=mode, min_interval=0, progress_prefix="",
|
||||
activity_phrases={"web": "Searching the web…"},
|
||||
)
|
||||
s = DraftStreamer(adapter, "c", cfg)
|
||||
assert s._config.mode == StreamingMode.PROGRESS
|
||||
assert s._status_text_only
|
||||
|
||||
async def run():
|
||||
await s.start()
|
||||
assert s.is_streaming
|
||||
await s.on_event(StreamEvent(
|
||||
type=StreamEventType.TOOL_CALL_START, tool_call={"name": "web_search"},
|
||||
))
|
||||
await s._update_task
|
||||
# Text tokens must not re-send an unchanged status line.
|
||||
await s.on_event(StreamEvent(type=StreamEventType.DELTA_TEXT, content="partial"))
|
||||
if s._update_task:
|
||||
await s._update_task
|
||||
await s.finalize("final answer")
|
||||
|
||||
asyncio.run(run())
|
||||
# Status line is cleared ("") once the answer is delivered.
|
||||
assert [t for _, t in adapter.status] == ["🤔 Thinking...", "Searching the web…", ""]
|
||||
assert adapter.sent == ["final answer"] # no placeholder message was sent
|
||||
assert adapter.edits == []
|
||||
assert not s.is_streaming
|
||||
|
||||
|
||||
def test_status_text_failure_never_breaks_turn():
|
||||
adapter = _StatusAdapter()
|
||||
|
||||
async def boom(channel_id, text):
|
||||
raise RuntimeError("nope")
|
||||
|
||||
adapter.set_status_text = boom
|
||||
s = DraftStreamer(adapter, "c", StreamingConfig(mode=StreamingMode.AUTO))
|
||||
|
||||
async def run():
|
||||
await s.start()
|
||||
await s.finalize("done")
|
||||
|
||||
asyncio.run(run())
|
||||
assert adapter.sent == ["done"]
|
||||
|
||||
|
||||
def test_status_text_uses_rate_limiter_and_stops_after_flood_failures():
|
||||
adapter = _StatusAdapter()
|
||||
calls = []
|
||||
|
||||
async def flood(channel_id, text):
|
||||
calls.append(text)
|
||||
raise RuntimeError("429 Too Many Requests")
|
||||
|
||||
adapter.set_status_text = flood
|
||||
|
||||
class _Limiter:
|
||||
acquired = 0
|
||||
|
||||
async def acquire(self, channel_id):
|
||||
_Limiter.acquired += 1
|
||||
|
||||
cfg = StreamingConfig(
|
||||
mode=StreamingMode.AUTO, min_interval=0.01, disable_progressive_edits_after=2,
|
||||
)
|
||||
s = DraftStreamer(adapter, "c", cfg, rate_limiter=_Limiter())
|
||||
|
||||
async def run():
|
||||
await s.start() # failure 1
|
||||
for name in ("a", "b", "c"):
|
||||
await s.on_event(StreamEvent(
|
||||
type=StreamEventType.TOOL_CALL_START, tool_call={"name": name},
|
||||
))
|
||||
if s._update_task:
|
||||
await s._update_task
|
||||
await s.finalize("done")
|
||||
|
||||
asyncio.run(run())
|
||||
# Bounded: status pushes stop once the flood streak disables progress.
|
||||
assert len(calls) == 2
|
||||
assert not s._progressive
|
||||
assert s._current_min_interval > 0.01 # backoff widened the interval
|
||||
assert _Limiter.acquired >= 2 # status pushes go through the shared limiter
|
||||
assert adapter.sent == ["done"] # final delivery unaffected
|
||||
@@ -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)
|
||||
Generated
+2
-2
@@ -1428,7 +1428,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-bot"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -1510,7 +1510,7 @@ provides-extras = ["gateway", "bot", "bot-whatsapp-web", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.20"
|
||||
version = "1.7.21"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.19"
|
||||
__version__ = "0.1.20"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-browser"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
description = "Browser automation for PraisonAI — Chrome extension bridge, CDP, and Playwright agents extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+2
-2
@@ -932,7 +932,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-browser"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -976,7 +976,7 @@ provides-extras = ["server", "playwright", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.20"
|
||||
version = "1.7.21"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -12,6 +12,6 @@ the remaining terminal-agent modules, with PEP 562 shims left behind at the
|
||||
old ``praisonai.*`` paths for backward compatibility.
|
||||
"""
|
||||
|
||||
__version__ = "0.1.19"
|
||||
__version__ = "0.1.20"
|
||||
|
||||
__all__ = ["__version__"]
|
||||
@@ -686,6 +686,13 @@ def _run_print_code(
|
||||
get_token_collector().reset()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
from ..state.project_sessions import install_tool_stats_hook, reset_tool_stats
|
||||
|
||||
reset_tool_stats()
|
||||
install_tool_stats_hook()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Resolve/continue a session id so --resume/--continue compose with -p for
|
||||
# scripted multi-turn. Best-effort: continuity is layered via the shared
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -1515,6 +1515,15 @@ def run_main(
|
||||
output = get_output_controller()
|
||||
_ = get_current_context() # Initialize context
|
||||
|
||||
# Per-tool usage tally for `praisonai usage --by tool` (Issue #6158).
|
||||
try:
|
||||
from ..state.project_sessions import install_tool_stats_hook, reset_tool_stats
|
||||
|
||||
reset_tool_stats()
|
||||
install_tool_stats_hook()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# --pure / --no-plugins: suppression is scoped by the @scopes_no_plugins
|
||||
# decorator, which sets PRAISONAI_NO_PLUGINS (read by the core PluginManager)
|
||||
# for the duration of this call and always restores the prior value on
|
||||
@@ -3101,6 +3110,11 @@ def _wire_subtree_context_hook(
|
||||
matcher=file_tool_matcher(),
|
||||
name="subtree_instruction_injection",
|
||||
)
|
||||
# This per-agent registry replaces the global one, so register the
|
||||
# per-tool usage tally here too (Issue #6158).
|
||||
from ..state.project_sessions import install_tool_stats_hook
|
||||
|
||||
install_tool_stats_hook(registry)
|
||||
agent_config["hooks"] = registry
|
||||
except Exception:
|
||||
return
|
||||
|
||||
@@ -9,6 +9,7 @@ platform, no network egress, no configuration required.
|
||||
praisonai usage # last 30 days, grouped by day
|
||||
praisonai usage --by model # grouped by model
|
||||
praisonai usage --by project # grouped by project
|
||||
praisonai usage --by tool # per-tool calls, success rate, latency
|
||||
praisonai usage --days 7 # only the last week
|
||||
praisonai usage --json # machine-readable output
|
||||
"""
|
||||
@@ -51,6 +52,7 @@ def _rows_from_sessions(
|
||||
"project": project_label,
|
||||
"total_tokens": int(s.get("total_tokens") or 0),
|
||||
"cost": float(s.get("cost") or 0.0),
|
||||
"tool_stats": s.get("tool_stats") or {},
|
||||
}
|
||||
)
|
||||
return rows
|
||||
@@ -170,13 +172,87 @@ def _group(rows: List[Dict[str, Any]], by: str) -> List[Tuple[str, int, float]]:
|
||||
return grouped
|
||||
|
||||
|
||||
def _group_tools(rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Roll per-session ``tool_stats`` up into one row per tool (Issue #6158).
|
||||
|
||||
``last_ms``/``last_error`` are chosen by when the tool itself last ran
|
||||
(``last_at``/``last_error_at``), not by session ``updated_at``.
|
||||
"""
|
||||
tools: Dict[str, Dict[str, Any]] = {}
|
||||
for row in rows:
|
||||
stats = row.get("tool_stats")
|
||||
if not isinstance(stats, dict):
|
||||
continue
|
||||
for name, stat in stats.items():
|
||||
if not isinstance(stat, dict) or not stat.get("calls"):
|
||||
continue
|
||||
cur = tools.get(name)
|
||||
if cur is None:
|
||||
cur = tools[name] = {
|
||||
"tool": name, "calls": 0, "errors": 0, "total_ms": 0.0,
|
||||
"last_ms": 0.0, "last_at": -1.0,
|
||||
"last_error": None, "last_error_at": -1.0,
|
||||
}
|
||||
cur["calls"] += int(stat.get("calls") or 0)
|
||||
cur["errors"] += int(stat.get("errors") or 0)
|
||||
cur["total_ms"] += float(stat.get("total_ms") or 0.0)
|
||||
last_at = float(stat.get("last_at") or 0.0)
|
||||
if last_at > cur["last_at"]:
|
||||
cur["last_at"] = last_at
|
||||
cur["last_ms"] = float(stat.get("last_ms") or 0.0)
|
||||
err_at = float(stat.get("last_error_at") or 0.0)
|
||||
if stat.get("last_error") and err_at > cur["last_error_at"]:
|
||||
cur["last_error_at"] = err_at
|
||||
cur["last_error"] = stat["last_error"]
|
||||
result = []
|
||||
for cur in tools.values():
|
||||
calls = cur["calls"]
|
||||
result.append({
|
||||
"tool": cur["tool"],
|
||||
"calls": calls,
|
||||
"successes": calls - cur["errors"],
|
||||
"errors": cur["errors"],
|
||||
"success_rate": round((calls - cur["errors"]) / calls, 4) if calls else 0.0,
|
||||
"avg_ms": round(cur["total_ms"] / calls, 1) if calls else 0.0,
|
||||
"last_ms": round(cur["last_ms"], 1),
|
||||
"last_error": cur["last_error"],
|
||||
})
|
||||
result.sort(key=lambda r: (-r["calls"], r["tool"]))
|
||||
return result
|
||||
|
||||
|
||||
def _print_tool_report(output, tools: List[Dict[str, Any]], errors: List[str], meta: Dict[str, Any], json_: bool) -> None:
|
||||
if json_ or output.is_json_mode:
|
||||
output.print_json({**meta, "rows": tools, "errors": errors})
|
||||
return
|
||||
for message in errors:
|
||||
output.print_warning(f"Usage may be incomplete: {message}")
|
||||
if not tools:
|
||||
if not errors:
|
||||
output.print_info("No tool usage recorded yet")
|
||||
return
|
||||
headers = ["Tool", "Calls", "Success", "Avg ms", "Last ms", "Last error"]
|
||||
table_rows = [
|
||||
[
|
||||
t["tool"],
|
||||
f"{t['calls']:,}",
|
||||
f"{t['success_rate'] * 100:.1f}%",
|
||||
f"{t['avg_ms']:.0f}",
|
||||
f"{t['last_ms']:.0f}",
|
||||
(t["last_error"] or "-")[:60],
|
||||
]
|
||||
for t in tools
|
||||
]
|
||||
output.print_table(headers, table_rows, title="Tool usage")
|
||||
|
||||
|
||||
@app.command()
|
||||
def usage(
|
||||
days: int = typer.Option(
|
||||
30, "--days", "-d", help="Only include sessions updated in the last N days (0 = all)"
|
||||
),
|
||||
by: str = typer.Option(
|
||||
"day", "--by", "-b", help="Group by: day, model, or project"
|
||||
"day", "--by", "-b", help="Group by: day, model, project, or tool"
|
||||
),
|
||||
project: Optional[str] = typer.Option(
|
||||
None, "--project", "-p", help="Restrict to a specific project ID"
|
||||
@@ -189,12 +265,18 @@ def usage(
|
||||
output = get_output_controller()
|
||||
|
||||
by = (by or "day").lower()
|
||||
if by not in ("day", "model", "project"):
|
||||
output.print_error("--by must be one of: day, model, project")
|
||||
if by not in ("day", "model", "project", "tool"):
|
||||
output.print_error("--by must be one of: day, model, project, tool")
|
||||
raise typer.Exit(1)
|
||||
|
||||
errors: List[str] = []
|
||||
rows = _within_days(_collect_rows(project, on_error=errors.append), days)
|
||||
if by == "tool":
|
||||
_print_tool_report(
|
||||
output, _group_tools(rows), errors,
|
||||
{"by": by, "days": days, "project": project}, json_,
|
||||
)
|
||||
return
|
||||
grouped = _group(rows, by)
|
||||
|
||||
total_tokens = sum(r[1] for r in grouped)
|
||||
|
||||
@@ -1977,6 +1977,10 @@ class PraisonAI:
|
||||
return {}
|
||||
|
||||
cli_config = {}
|
||||
|
||||
# Extract --model so it overrides a YAML top-level ``llm``
|
||||
if getattr(self.args, 'model', None):
|
||||
cli_config['model'] = self.args.model
|
||||
|
||||
# Extract --trust flag
|
||||
if getattr(self.args, 'trust', False):
|
||||
|
||||
@@ -5,6 +5,8 @@ Provides session continuity within project boundaries.
|
||||
"""
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
@@ -506,6 +508,7 @@ def accumulate_session_usage(
|
||||
Best-effort: any failure leaves the session untouched and returns the
|
||||
current (possibly unchanged) totals.
|
||||
"""
|
||||
_flush_tool_stats(session_id, project_path)
|
||||
current = read_session_usage(session_id, project_path)
|
||||
try:
|
||||
from praisonaiagents.telemetry.token_collector import get_token_collector
|
||||
@@ -575,6 +578,142 @@ def accumulate_session_usage(
|
||||
return updated
|
||||
|
||||
|
||||
# ── Per-tool outcome tally (Issue #6158) ───────────────────────────────────
|
||||
# Fed by the existing AFTER_TOOL hook payload (tool name, error, timing) and
|
||||
# merged into the session's ``tool_stats`` metadata alongside token/cost, so
|
||||
# ``praisonai usage --by tool`` has durable data to report. No core changes.
|
||||
_tool_tally: Dict[str, Dict[str, Any]] = {}
|
||||
_tool_tally_lock = threading.Lock()
|
||||
_default_tool_hook_id: Optional[str] = None
|
||||
|
||||
|
||||
def _empty_tool_stat() -> Dict[str, Any]:
|
||||
return {
|
||||
"calls": 0, "errors": 0, "total_ms": 0.0, "last_ms": 0.0,
|
||||
"last_error": None, "last_at": 0.0, "last_error_at": 0.0,
|
||||
}
|
||||
|
||||
|
||||
def _on_after_tool(data: Any) -> None:
|
||||
"""AFTER_TOOL hook: tally one tool outcome (never blocks or alters it)."""
|
||||
try:
|
||||
name = getattr(data, "tool_name", "") or "unknown"
|
||||
error = getattr(data, "tool_error", None)
|
||||
ms = float(getattr(data, "execution_time_ms", 0.0) or 0.0)
|
||||
now = time.time()
|
||||
with _tool_tally_lock:
|
||||
stat = _tool_tally.setdefault(name, _empty_tool_stat())
|
||||
stat["calls"] += 1
|
||||
stat["total_ms"] += ms
|
||||
stat["last_ms"] = ms
|
||||
stat["last_at"] = now
|
||||
if error:
|
||||
stat["errors"] += 1
|
||||
stat["last_error"] = str(error)[:500]
|
||||
stat["last_error_at"] = now
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def install_tool_stats_hook(registry: Any = None) -> None:
|
||||
"""Register the per-tool tally on ``registry`` (default: global registry).
|
||||
|
||||
Registration on the global registry is idempotent; a per-agent registry
|
||||
(which replaces the global one for that agent) gets its own registration.
|
||||
"""
|
||||
global _default_tool_hook_id
|
||||
try:
|
||||
from praisonaiagents.hooks import HookEvent
|
||||
|
||||
if registry is not None:
|
||||
registry.register_function(
|
||||
event=HookEvent.AFTER_TOOL, func=_on_after_tool, name="tool_stats"
|
||||
)
|
||||
return
|
||||
if _default_tool_hook_id is None:
|
||||
from praisonaiagents.hooks import add_hook
|
||||
|
||||
_default_tool_hook_id = add_hook(HookEvent.AFTER_TOOL, _on_after_tool)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def reset_tool_stats() -> None:
|
||||
"""Clear the in-process tally (call at the start of a run)."""
|
||||
with _tool_tally_lock:
|
||||
_tool_tally.clear()
|
||||
|
||||
|
||||
def merge_tool_stats(
|
||||
base: Optional[Dict[str, Any]], delta: Dict[str, Any]
|
||||
) -> Dict[str, Dict[str, Any]]:
|
||||
"""Merge per-tool ``delta`` counters into ``base`` (returns a new dict)."""
|
||||
merged: Dict[str, Dict[str, Any]] = {}
|
||||
for name, stat in (base or {}).items():
|
||||
if isinstance(stat, dict):
|
||||
merged[name] = {**_empty_tool_stat(), **stat}
|
||||
for name, stat in (delta or {}).items():
|
||||
if not isinstance(stat, dict):
|
||||
continue
|
||||
cur = merged.setdefault(name, _empty_tool_stat())
|
||||
cur["calls"] = int(cur.get("calls") or 0) + int(stat.get("calls") or 0)
|
||||
cur["errors"] = int(cur.get("errors") or 0) + int(stat.get("errors") or 0)
|
||||
cur["total_ms"] = round(
|
||||
float(cur.get("total_ms") or 0.0) + float(stat.get("total_ms") or 0.0), 3
|
||||
)
|
||||
# Keep whichever side ran most recently (timestamps, not merge order).
|
||||
if stat.get("calls") and float(stat.get("last_at") or 0.0) >= float(cur.get("last_at") or 0.0):
|
||||
cur["last_ms"] = round(float(stat.get("last_ms") or 0.0), 3)
|
||||
cur["last_at"] = float(stat.get("last_at") or 0.0)
|
||||
if stat.get("last_error") and float(stat.get("last_error_at") or 0.0) >= float(
|
||||
cur.get("last_error_at") or 0.0
|
||||
):
|
||||
cur["last_error"] = stat["last_error"]
|
||||
cur["last_error_at"] = float(stat.get("last_error_at") or 0.0)
|
||||
return merged
|
||||
|
||||
|
||||
def _flush_tool_stats(session_id: str, project_path: Optional[str] = None) -> None:
|
||||
"""Merge (and drain) this run's tool tally into the session record.
|
||||
|
||||
The merge runs inside the store's locked read-modify-write when available
|
||||
so concurrent CLI processes don't overwrite each other's counts, and the
|
||||
drained tally is restored if the write fails so a later flush retries it.
|
||||
"""
|
||||
with _tool_tally_lock:
|
||||
if not _tool_tally:
|
||||
return
|
||||
delta = {k: dict(v) for k, v in _tool_tally.items()}
|
||||
_tool_tally.clear()
|
||||
ok = False
|
||||
try:
|
||||
store = _resolve_usage_store(session_id, project_path) or get_project_session_store(project_path)
|
||||
modify = getattr(store, "_modify_session_locked", None)
|
||||
if callable(modify):
|
||||
def _apply(session: Any) -> None:
|
||||
session.metadata["tool_stats"] = merge_tool_stats(
|
||||
session.metadata.get("tool_stats"), delta
|
||||
)
|
||||
|
||||
ok = bool(modify(session_id, _apply, error_label="merge tool stats"))
|
||||
else:
|
||||
try:
|
||||
existing = (getattr(store.get_session(session_id), "metadata", {}) or {}).get("tool_stats")
|
||||
except Exception:
|
||||
existing = None
|
||||
ok = bool(store.update_session_metadata(
|
||||
session_id, tool_stats=merge_tool_stats(existing, delta)
|
||||
))
|
||||
except Exception:
|
||||
ok = False
|
||||
if not ok:
|
||||
with _tool_tally_lock:
|
||||
restored = merge_tool_stats(delta, _tool_tally)
|
||||
_tool_tally.clear()
|
||||
_tool_tally.update(restored)
|
||||
|
||||
|
||||
def _first_exchange(history: List[Dict[str, Any]]) -> tuple:
|
||||
"""Return ``(user_msg, assistant_msg)`` for the first user↔assistant turn.
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-code"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
description = "Agentic terminal CLI for PraisonAI — the terminal-native agent (run, chat, code, runtime, cli backends) extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
@@ -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
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Tests for per-tool usage reporting: ``praisonai usage --by tool`` (Issue #6158)."""
|
||||
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
from typer.testing import CliRunner
|
||||
|
||||
import praisonai_code.cli.state.project_sessions as ps
|
||||
from praisonaiagents.hooks import HookRegistry
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(monkeypatch, tmp_path):
|
||||
project = DefaultSessionStore(session_dir=str(tmp_path / "project"))
|
||||
empty_global = DefaultSessionStore(session_dir=str(tmp_path / "global"))
|
||||
monkeypatch.setattr(ps, "get_project_session_store", lambda *a, **k: project)
|
||||
monkeypatch.setattr(ps, "_get_default_store", lambda: empty_global)
|
||||
monkeypatch.setattr(ps, "_resolve_usage_store", lambda *a, **k: project)
|
||||
ps.reset_tool_stats()
|
||||
yield project
|
||||
ps.reset_tool_stats()
|
||||
|
||||
|
||||
class _Evt:
|
||||
def __init__(self, name, ms, error=None):
|
||||
self.tool_name, self.execution_time_ms, self.tool_error = name, ms, error
|
||||
|
||||
|
||||
def _usage_json(by="tool"):
|
||||
from praisonai_code.cli.commands.usage import app
|
||||
|
||||
result = CliRunner().invoke(app, ["--by", by, "--json", "--days", "0"])
|
||||
assert result.exit_code == 0, result.output
|
||||
return json.loads(result.output[result.output.index("{"):])
|
||||
|
||||
|
||||
def test_tally_persists_and_merges_across_runs(store):
|
||||
store.add_message("s1", "user", "hi")
|
||||
ps._on_after_tool(_Evt("read_file", 10))
|
||||
ps._on_after_tool(_Evt("read_file", 30, error="boom"))
|
||||
ps._flush_tool_stats("s1")
|
||||
# Tally is drained after a flush, so a second flush does not double count.
|
||||
ps._flush_tool_stats("s1")
|
||||
ps._on_after_tool(_Evt("read_file", 20))
|
||||
ps._flush_tool_stats("s1")
|
||||
|
||||
stats = store.get_session("s1").metadata["tool_stats"]["read_file"]
|
||||
assert stats["calls"] == 3
|
||||
assert stats["errors"] == 1
|
||||
assert stats["total_ms"] == 60
|
||||
assert stats["last_ms"] == 20
|
||||
assert stats["last_error"] == "boom"
|
||||
|
||||
|
||||
|
||||
def test_failed_save_keeps_tally_for_retry(store, monkeypatch):
|
||||
store.add_message("s1", "user", "hi")
|
||||
ps._on_after_tool(_Evt("read_file", 10))
|
||||
monkeypatch.setattr(store, "_modify_session_locked", lambda *a, **k: False)
|
||||
ps._flush_tool_stats("s1")
|
||||
assert "tool_stats" not in store.get_session("s1").metadata
|
||||
monkeypatch.undo()
|
||||
monkeypatch.setattr(ps, "_resolve_usage_store", lambda *a, **k: store)
|
||||
ps._on_after_tool(_Evt("read_file", 20))
|
||||
ps._flush_tool_stats("s1")
|
||||
stats = store.get_session("s1").metadata["tool_stats"]["read_file"]
|
||||
assert stats["calls"] == 2 and stats["total_ms"] == 30
|
||||
|
||||
|
||||
def test_latest_values_follow_tool_time_not_session_time(store):
|
||||
store.add_message("a", "user", "hi")
|
||||
store.add_message("b", "user", "hi")
|
||||
store.update_session_metadata("a", tool_stats={"search": {
|
||||
"calls": 1, "errors": 1, "total_ms": 5, "last_ms": 5,
|
||||
"last_error": "old", "last_at": 100.0, "last_error_at": 100.0}})
|
||||
store.update_session_metadata("b", tool_stats={"search": {
|
||||
"calls": 1, "errors": 1, "total_ms": 9, "last_ms": 9,
|
||||
"last_error": "new", "last_at": 200.0, "last_error_at": 200.0}})
|
||||
store.add_message("a", "user", "text-only turn") # bumps a's updated_at
|
||||
row = _usage_json()["rows"][0]
|
||||
assert row["last_ms"] == 9.0 and row["last_error"] == "new"
|
||||
|
||||
def test_usage_by_tool_json_shape(store):
|
||||
store.add_message("s1", "user", "hi")
|
||||
store.add_message("s2", "user", "hi")
|
||||
ps._on_after_tool(_Evt("search", 100))
|
||||
ps._on_after_tool(_Evt("search", 300, error="timeout"))
|
||||
ps._flush_tool_stats("s1")
|
||||
ps._on_after_tool(_Evt("search", 200))
|
||||
ps._on_after_tool(_Evt("write_file", 5))
|
||||
ps._flush_tool_stats("s2")
|
||||
|
||||
data = _usage_json()
|
||||
assert data["by"] == "tool"
|
||||
rows = {r["tool"]: r for r in data["rows"]}
|
||||
assert rows["search"]["calls"] == 3
|
||||
assert rows["search"]["successes"] == 2
|
||||
assert rows["search"]["errors"] == 1
|
||||
assert rows["search"]["success_rate"] == pytest.approx(2 / 3, abs=1e-3)
|
||||
assert rows["search"]["avg_ms"] == 200.0
|
||||
assert rows["search"]["last_error"] == "timeout"
|
||||
assert rows["write_file"]["calls"] == 1
|
||||
assert data["rows"][0]["tool"] == "search" # most-called first
|
||||
|
||||
|
||||
def test_usage_by_tool_empty(store):
|
||||
assert _usage_json()["rows"] == []
|
||||
|
||||
|
||||
def test_existing_groupings_unaffected(store):
|
||||
store.add_message("s1", "user", "hi")
|
||||
store.update_session_metadata("s1", total_tokens=10, cost=0.5)
|
||||
data = _usage_json(by="model")
|
||||
assert data["total_tokens"] == 10
|
||||
|
||||
|
||||
def test_real_agent_tool_call_is_reported(store):
|
||||
"""A real Agent tool execution flows through AFTER_TOOL into the report."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def add(a: int, b: int) -> int:
|
||||
"""Add two numbers."""
|
||||
return a + b
|
||||
|
||||
registry = HookRegistry()
|
||||
ps.install_tool_stats_hook(registry)
|
||||
agent = Agent(name="calc", tools=[add], hooks=registry)
|
||||
assert agent.execute_tool("add", {"a": 2, "b": 3}) == 5
|
||||
|
||||
store.add_message("s1", "user", "hi")
|
||||
ps._flush_tool_stats("s1")
|
||||
rows = {r["tool"]: r for r in _usage_json()["rows"]}
|
||||
assert rows["add"]["calls"] == 1
|
||||
assert rows["add"]["errors"] == 0
|
||||
|
||||
|
||||
@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="needs OPENAI_API_KEY")
|
||||
def test_real_agentic_llm_tool_run_is_reported(store):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def get_weather(city: str) -> str:
|
||||
"""Get the weather for a city."""
|
||||
return f"Sunny in {city}"
|
||||
|
||||
ps.install_tool_stats_hook()
|
||||
agent = Agent(name="weather", tools=[get_weather], llm="gpt-4o-mini")
|
||||
print(agent.start("Use the get_weather tool for Paris and report it."))
|
||||
|
||||
store.add_message("s1", "user", "hi")
|
||||
ps._flush_tool_stats("s1")
|
||||
rows = {r["tool"]: r for r in _usage_json()["rows"]}
|
||||
assert rows["get_weather"]["calls"] >= 1
|
||||
Generated
+2
-2
@@ -1370,7 +1370,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-code"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1403,7 +1403,7 @@ requires-dist = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.20"
|
||||
version = "1.7.21"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.19"
|
||||
__version__ = "0.1.20"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-deploy"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
description = "Deployment for PraisonAI — API, Docker, and cloud (AWS, Azure, GCP) deployments extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+2
-2
@@ -977,7 +977,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-deploy"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1014,7 +1014,7 @@ provides-extras = ["api", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.20"
|
||||
version = "1.7.21"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.21"
|
||||
__version__ = "0.1.22"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-mcp"
|
||||
version = "0.1.21"
|
||||
version = "0.1.22"
|
||||
description = "MCP server hosting for PraisonAI — expose agents and tools to Cursor, Claude Desktop, and other MCP clients."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+2
-2
@@ -1044,7 +1044,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-mcp"
|
||||
version = "0.1.21"
|
||||
version = "0.1.22"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1084,7 +1084,7 @@ provides-extras = ["server", "auth", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.20"
|
||||
version = "1.7.21"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.19"
|
||||
__version__ = "0.1.20"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-sandbox"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
description = "Where PraisonAI agents run code — Docker, E2B, Modal, Daytona, Fly.io, Tenki, Novita, Sandlock, SSH and subprocess, as both single sandboxes and multi-instance compute providers."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+2
-2
@@ -1679,7 +1679,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-sandbox"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1744,7 +1744,7 @@ provides-extras = ["docker", "e2b", "sandlock", "ssh", "tenki", "flyio", "modal"
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.20"
|
||||
version = "1.7.21"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.19"
|
||||
__version__ = "0.1.20"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-train"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
description = "Training for PraisonAI — LLM fine-tuning (Unsloth) and iterative agent training extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
||||
Generated
+2
-2
@@ -2606,7 +2606,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-train"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -2696,7 +2696,7 @@ provides-extras = ["llm", "vision", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.20"
|
||||
version = "1.7.21"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -207,7 +207,7 @@ TS-only members: `dependencies`?
|
||||
|
||||
### `Agent.start`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/execution_mixin.py:819`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/execution_mixin.py:815`
|
||||
- TypeScript: `src/praisonai-ts/src/agent/simple.ts:2497`
|
||||
- Counts: 1 python params: 1 exact, 0 camelCase, 0 alias, 0 flattened, 0 missing; 0 mismatches; 0 waived; 20 TS-only of 21
|
||||
|
||||
@@ -386,7 +386,7 @@ TS-only members: `maxIdenticalActions`?, `maxSimilarActions`?, `maxConsecutiveFa
|
||||
|
||||
### `EscalationPipeline.__init__`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/escalation/pipeline.py:44`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/escalation/pipeline.py:45`
|
||||
- TypeScript: `src/praisonai-ts/src/escalation/pipeline.ts:51` (ctor `src/praisonai-ts/src/escalation/pipeline.ts:103`)
|
||||
- Counts: 5 python params: 3 exact, 2 camelCase, 0 alias, 0 flattened, 0 missing; 0 mismatches; 0 waived; 4 TS-only of 9
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -1876,7 +1876,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/execution_mixin.py:819",
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/execution_mixin.py:815",
|
||||
"resolved_class": "ExecutionMixin"
|
||||
},
|
||||
"ts_only": [
|
||||
@@ -3233,7 +3233,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/escalation/pipeline.py:44",
|
||||
"location": "src/praisonai-agents/praisonaiagents/escalation/pipeline.py:45",
|
||||
"resolved_class": "EscalationPipeline"
|
||||
},
|
||||
"ts_only": [
|
||||
|
||||
@@ -32,6 +32,9 @@ export class FileMemory {
|
||||
constructor(config: FileMemoryConfig) {
|
||||
this.filePath = config.filePath;
|
||||
this.maxEntries = config.maxEntries ?? 10000;
|
||||
if (this.maxEntries < 0) {
|
||||
throw new RangeError('maxEntries must be non-negative');
|
||||
}
|
||||
this.embeddingProvider = config.embeddingProvider;
|
||||
this.compactionThreshold = config.compactionThreshold ?? 1000;
|
||||
this.autoCompact = config.autoCompact ?? true;
|
||||
|
||||
@@ -38,6 +38,9 @@ export class Memory {
|
||||
|
||||
constructor(config: MemoryConfig = {}) {
|
||||
this.maxEntries = config.maxEntries ?? 1000;
|
||||
if (this.maxEntries < 0) {
|
||||
throw new RangeError('maxEntries must be non-negative');
|
||||
}
|
||||
this.maxTokens = config.maxTokens ?? 100000;
|
||||
this.embeddingProvider = config.embeddingProvider;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
import { Memory } from '../../../src/memory/memory';
|
||||
import { FileMemory } from '../../../src/memory/file-memory';
|
||||
import * as os from 'os';
|
||||
import * as path from 'path';
|
||||
|
||||
describe('maxEntries validation (#6146)', () => {
|
||||
const filePath = path.join(os.tmpdir(), `praisonai-maxentries-${process.pid}.jsonl`);
|
||||
|
||||
it.each([-1, -0.5, -Infinity])('Memory rejects negative maxEntries %p', (n) => {
|
||||
expect(() => new Memory({ maxEntries: n })).toThrow(RangeError);
|
||||
});
|
||||
|
||||
it.each([-1, -0.5, -Infinity])('FileMemory rejects negative maxEntries %p', (n) => {
|
||||
expect(() => new FileMemory({ filePath, maxEntries: n })).toThrow(RangeError);
|
||||
});
|
||||
|
||||
it('Memory with maxEntries 0 retains no entries', async () => {
|
||||
const m = new Memory({ maxEntries: 0 });
|
||||
await m.add('hello', 'user');
|
||||
expect(m.getAll()).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('Memory with positive maxEntries keeps the newest entries', async () => {
|
||||
const m = new Memory({ maxEntries: 2 });
|
||||
await m.add('a', 'user');
|
||||
await m.add('b', 'user');
|
||||
await m.add('c', 'user');
|
||||
expect(m.getAll().map(e => e.content)).toEqual(['b', 'c']);
|
||||
});
|
||||
});
|
||||
@@ -3,8 +3,8 @@ class Praisonai < Formula
|
||||
|
||||
desc "AI tools for various AI applications"
|
||||
homepage "https://github.com/MervinPraison/PraisonAI"
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.7.22.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.7.22.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.7.23.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.7.23.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
"""Private root for wrapper exceptions.
|
||||
|
||||
Joins every ``praisonai.*`` error to the core SDK hierarchy so downstream code
|
||||
can catch them all with ``except praisonaiagents.PraisonAIError``. Concrete
|
||||
classes keep their stdlib parents (``ValueError``, ``RuntimeError``,
|
||||
``TimeoutError``) via multiple inheritance, so existing ``except`` sites are
|
||||
unaffected.
|
||||
"""
|
||||
|
||||
from praisonaiagents.errors import PraisonAIError
|
||||
|
||||
|
||||
class _WrapperError(PraisonAIError):
|
||||
"""Base for wrapper errors; keeps plain ``str(e)`` and original ``args``."""
|
||||
|
||||
def __init__(self, *args):
|
||||
PraisonAIError.__init__(self, str(args[0]) if args else "")
|
||||
self.args = args
|
||||
|
||||
# Core's __str__ decorates with category/agent/run id; wrapper errors are
|
||||
# shown verbatim to CLI users, so keep the stdlib formatting.
|
||||
__str__ = Exception.__str__
|
||||
@@ -156,6 +156,7 @@ def _get_default_adapter_registry():
|
||||
|
||||
# Framework availability detection (lazy via __getattr__)
|
||||
from ._framework_availability import is_available
|
||||
from praisonai._errors import _WrapperError
|
||||
|
||||
# Lazy constants mapping for backward compatibility
|
||||
_AVAIL = {
|
||||
@@ -192,7 +193,7 @@ def __dir__():
|
||||
_DEFAULT_TOOL_TIMEOUT_WORKERS = 32
|
||||
|
||||
|
||||
class ToolTimeoutError(TimeoutError):
|
||||
class ToolTimeoutError(_WrapperError, TimeoutError):
|
||||
"""Raised when a wrapped tool exceeds its per-call timeout.
|
||||
|
||||
Preserves the tool's declared return-type contract: instead of silently
|
||||
@@ -1120,6 +1121,26 @@ class AgentsGenerator:
|
||||
'run_ctx': run_ctx,
|
||||
}
|
||||
|
||||
def _run_config_list(self, config):
|
||||
"""Return the ``config_list`` for this run with YAML top-level ``llm`` applied.
|
||||
|
||||
Precedence for agents without their own ``llm``/``model``: explicit CLI
|
||||
``--model`` (``cli_config['model']``) > top-level YAML ``llm`` >
|
||||
credential/catalogue-resolved ``config_list``. Returns a copy so the
|
||||
caller-owned ``self.config_list`` is never mutated.
|
||||
"""
|
||||
config_list = self.config_list
|
||||
if (self.cli_config or {}).get('model'):
|
||||
return config_list
|
||||
yaml_llm = config.get('llm')
|
||||
if isinstance(yaml_llm, dict):
|
||||
yaml_llm = yaml_llm.get('model')
|
||||
if not isinstance(yaml_llm, str) or not yaml_llm.strip():
|
||||
return config_list
|
||||
if config_list:
|
||||
return [{**config_list[0], 'model': yaml_llm.strip()}, *config_list[1:]]
|
||||
return [{'model': yaml_llm.strip()}]
|
||||
|
||||
def _run_adapter_setup(self, adapter):
|
||||
"""Run an adapter's setup hooks.
|
||||
|
||||
@@ -1696,7 +1717,7 @@ class AgentsGenerator:
|
||||
f"Framework {adapter.name!r} does not support building agents "
|
||||
"without running them; use generate_crew_and_kickoff()."
|
||||
)
|
||||
model_name = adapter._pick_model(self.config_list)
|
||||
model_name = adapter._pick_model(self._run_config_list(prep['config']))
|
||||
# Thread the per-agent tool_timeout wrap resolver (carried in prep by
|
||||
# _prepare_for_run) exactly as the real run paths do, so agents built for
|
||||
# evaluation carry the same heterogeneous per-agent timeout guards a
|
||||
@@ -1775,7 +1796,7 @@ class AgentsGenerator:
|
||||
with scoped_bridge():
|
||||
return prep['adapter'].run(
|
||||
prep['config'],
|
||||
self.config_list,
|
||||
self._run_config_list(prep['config']),
|
||||
prep['topic'],
|
||||
tools_dict=prep['tools_dict'],
|
||||
agent_callback=getattr(self, 'agent_callback', None),
|
||||
@@ -1844,7 +1865,7 @@ class AgentsGenerator:
|
||||
await asyncio.to_thread(self._run_adapter_setup, prep['adapter'])
|
||||
return await prep['adapter'].arun(
|
||||
prep['config'],
|
||||
self.config_list,
|
||||
self._run_config_list(prep['config']),
|
||||
prep['topic'],
|
||||
tools_dict=prep['tools_dict'],
|
||||
agent_callback=getattr(self, 'agent_callback', None),
|
||||
|
||||
Loaded 100 of 130 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user