mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 08:11:07 -10:00
Compare commits
53
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2ead7ca927 | ||
|
|
ec0efc2b69 | ||
|
|
5a19ab07b2 | ||
|
|
a5dca259b0 | ||
|
|
59fb20af8e | ||
|
|
f9e058f939 | ||
|
|
e0bdf9648b | ||
|
|
fbfe8fa13c | ||
|
|
5aa15f5d3b | ||
|
|
7166fc5de6 | ||
|
|
88b9d8c875 | ||
|
|
a028c62125 | ||
|
|
4a9dcdcd8f | ||
|
|
e0e099caef | ||
|
|
1232cd8db2 | ||
|
|
8bd67b83bc | ||
|
|
d14fe808cb | ||
|
|
b8fd43457d | ||
|
|
78f7331738 | ||
|
|
78a648113d | ||
|
|
274acbdca1 | ||
|
|
6668341ab3 | ||
|
|
9c054f10ad | ||
|
|
ea1bca27b0 | ||
|
|
838e7affc2 | ||
|
|
4d1deaa839 | ||
|
|
a177d6bef4 | ||
|
|
df529e839a | ||
|
|
992a1d5453 | ||
|
|
c10e9f306a | ||
|
|
e0374be6da | ||
|
|
0c73f6354d | ||
|
|
0e6861e910 | ||
|
|
78e7e59dbd | ||
|
|
2eea6c6dcc | ||
|
|
7ff7041e62 | ||
|
|
98096d82d8 | ||
|
|
a805ae7d83 | ||
|
|
da0202a017 | ||
|
|
ce6bf2cb54 | ||
|
|
e7076d0949 | ||
|
|
6aabf3a856 | ||
|
|
5fe1a61697 | ||
|
|
cd5f61d14c | ||
|
|
153e4c8078 | ||
|
|
4be3be0797 | ||
|
|
ebf0a1b275 | ||
|
|
c25c5f042d | ||
|
|
58fbee09e4 | ||
|
|
7c64856641 | ||
|
|
ee454013b0 | ||
|
|
f4ac213ae2 | ||
|
|
ec978dd399 |
No files matched your search
@@ -313,7 +313,11 @@ runs:
|
||||
# Add allowed tools (include GitHub tools by default)
|
||||
TOOLS="mcp__github__get_issue,mcp__github__get_issue_comments,mcp__github__update_issue,mcp__github__search_issues,mcp__github__list_issues,mcp__github__add_issue_comment,mcp__github__get_pull_request,mcp__github__get_pull_request_files,mcp__github__get_pull_request_reviews,mcp__github__get_pull_request_status,Read,Write,Edit,Bash"
|
||||
if [[ -n "$INPUT_ALLOWED_TOOLS" ]]; then
|
||||
TOOLS="$TOOLS,$INPUT_ALLOWED_TOOLS"
|
||||
# Callers pass a YAML block (one tool per line). `read` below stops at
|
||||
# the first newline, so join lines with commas or every tool after
|
||||
# the first is silently dropped.
|
||||
EXTRA_TOOLS=$(printf '%s' "$INPUT_ALLOWED_TOOLS" | tr '\n' ',' | sed 's/,,*/,/g; s/^,//; s/,$//')
|
||||
TOOLS="$TOOLS,$EXTRA_TOOLS"
|
||||
fi
|
||||
|
||||
# Remove any explicitly disallowed tools from the effective allow-list so
|
||||
|
||||
@@ -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.18" \
|
||||
"praisonai>=4.7.21" \
|
||||
"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.18" \
|
||||
"praisonai>=4.7.21" \
|
||||
"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.18" \
|
||||
"praisonai>=4.7.21" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@ from collections import OrderedDict
|
||||
import inspect
|
||||
|
||||
# Decomposed agent functionality - real implementations in mixin files
|
||||
from .chat_mixin import ChatMixin
|
||||
from .chat_mixin import ChatMixin, _deadline_scope
|
||||
from .execution_mixin import ExecutionMixin
|
||||
from .memory_mixin import MemoryMixin
|
||||
from .async_memory_mixin import AsyncMemoryMixin
|
||||
@@ -463,7 +463,11 @@ class Agent(GoalLoopMixin, SteeringMixin, SandboxMixin, SkillReviewMixin, Unifie
|
||||
concurrent turns on a shared agent do not share per-turn counters."""
|
||||
if self._loop_guard is None:
|
||||
from ..escalation.loop_guard import LoopGuard, LoopGuardConfig
|
||||
self._loop_guard = LoopGuard(LoopGuardConfig(enabled=True))
|
||||
config = LoopGuardConfig(enabled=True)
|
||||
limit = getattr(self, "max_execution_time", None)
|
||||
if limit is not None:
|
||||
config.max_time_per_turn = float(limit) if limit > 0 else float("inf")
|
||||
self._loop_guard = LoopGuard(config)
|
||||
from .chat_mixin import _current_turn_loop_guard
|
||||
turn_guard = _current_turn_loop_guard(self)
|
||||
return turn_guard if turn_guard is not None else self._loop_guard
|
||||
@@ -1913,6 +1917,17 @@ class Agent(GoalLoopMixin, SteeringMixin, SandboxMixin, SkillReviewMixin, Unifie
|
||||
"Use Agent(policy=PolicyEngine(...)) for tool policy enforcement, "
|
||||
"or pass a validator via GuardrailConfig(validator=...)."
|
||||
)
|
||||
# GuardrailConfig.on_fail has a single implemented policy (retry, then
|
||||
# return the blocked message); refuse "skip"/"raise" rather than
|
||||
# silently behaving as "retry".
|
||||
if isinstance(_guardrails_config, GuardrailConfig) and (
|
||||
getattr(_guardrails_config.on_fail, "value", _guardrails_config.on_fail) != "retry"
|
||||
):
|
||||
raise ValueError(
|
||||
f"GuardrailConfig.on_fail={_guardrails_config.on_fail!r} is not supported; "
|
||||
"only 'retry' is implemented. Remove on_fail, or handle a blocked "
|
||||
"response via Agent.last_guardrail_error."
|
||||
)
|
||||
|
||||
# Fail loud when the value could not be turned into an enforceable
|
||||
# guardrail, instead of storing something that never runs. Silently
|
||||
@@ -3283,6 +3298,12 @@ Your Goal: {self.goal}
|
||||
if k in ("_db_run_context", "_db_session_init_lock"):
|
||||
# A clone starts outside the source agent's active DB turn.
|
||||
continue
|
||||
elif k == "_tool_executor":
|
||||
# Thread pools can't be copied; the clone lazily builds its own.
|
||||
continue
|
||||
elif k == "_Agent__hook_runner":
|
||||
# Holds locks; rebuilt lazily on first use.
|
||||
object.__setattr__(result, k, None)
|
||||
elif k in ("_Agent__cache_lock",):
|
||||
object.__setattr__(result, k, threading.RLock())
|
||||
elif k in ("_cost_lock", "_auto_save_lock"):
|
||||
@@ -4109,7 +4130,9 @@ Summary:"""
|
||||
step budget from ``ExecutionConfig.max_steps`` was reached and the run was
|
||||
truncated), a provider block/refusal/truncation
|
||||
(``"content_filtered" | "refused" | "length_truncated"``) derived from the
|
||||
LLM ``finish_reason``/refusal signal, or ``"error"``. Lets CLI/CI callers
|
||||
LLM ``finish_reason``/refusal signal, ``"timeout"`` (the turn passed
|
||||
``ExecutionConfig.max_execution_time`` or the ``run_autonomous``
|
||||
deadline), or ``"error"``. Lets CLI/CI callers
|
||||
branch on the terminal reason instead of parsing a magic string. Reads
|
||||
from whichever backend (OpenAI-native or LiteLLM) executed the last turn.
|
||||
"""
|
||||
@@ -4727,7 +4750,7 @@ Summary:"""
|
||||
if result.success:
|
||||
print(result.output)
|
||||
"""
|
||||
from .autonomy import AutonomyResult
|
||||
from .autonomy import AutonomyResult, _timeout_result
|
||||
import time as time_module
|
||||
from datetime import datetime, timezone
|
||||
|
||||
@@ -4766,6 +4789,11 @@ Summary:"""
|
||||
effective_promise = completion_promise
|
||||
if effective_promise is None:
|
||||
effective_promise = self.autonomy_config.get("completion_promise")
|
||||
# The promise is the only success signal, so tell the model about it
|
||||
# unless the caller already did (e.g. the `loop` CLI instructions).
|
||||
_promise_tag = f"<promise>{effective_promise}</promise>"
|
||||
if effective_promise and _promise_tag not in prompt and _promise_tag not in (self.instructions or ""):
|
||||
prompt = f"{prompt}\n\nWhen the task is fully complete, include {_promise_tag} in your response."
|
||||
|
||||
# Get clear_context from config if not explicitly set
|
||||
effective_clear_context = clear_context
|
||||
@@ -4796,17 +4824,10 @@ Summary:"""
|
||||
)
|
||||
|
||||
# Check timeout
|
||||
if timeout_seconds and (time_module.time() - start_time) > timeout_seconds:
|
||||
return AutonomyResult(
|
||||
success=False,
|
||||
output="Task timed out",
|
||||
completion_reason="timeout",
|
||||
iterations=iterations,
|
||||
stage=stage,
|
||||
actions=actions_taken,
|
||||
duration_seconds=time_module.time() - start_time,
|
||||
started_at=started_at,
|
||||
)
|
||||
_timed_out = _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at)
|
||||
if _timed_out is not None:
|
||||
return _timed_out
|
||||
|
||||
# G2: Check for interrupt request (cooperative cancellation) - sync version
|
||||
if self.interrupt_controller and self.interrupt_controller.is_set():
|
||||
@@ -4829,9 +4850,17 @@ Summary:"""
|
||||
# Reset the per-turn warning latch and inject any queued nudge.
|
||||
turn_prompt = self._consume_self_correction(prompt)
|
||||
try:
|
||||
response = self.chat(turn_prompt)
|
||||
with _deadline_scope(timeout_seconds and timeout_seconds - (time_module.time() - start_time)):
|
||||
response = self.chat(turn_prompt)
|
||||
except Exception as e:
|
||||
return AutonomyResult(
|
||||
# The run deadline cancels an overrunning turn; report it as a timeout.
|
||||
# A budget stop keeps its typed budget_exhausted result when a cap applies.
|
||||
return _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at
|
||||
) or (isinstance(e, BudgetExceededError) and self._autonomy_budget_result(
|
||||
iterations, stage, actions_taken, start_time, started_at,
|
||||
"", spend_baseline
|
||||
)) or AutonomyResult(
|
||||
success=False,
|
||||
output=str(e),
|
||||
completion_reason="error",
|
||||
@@ -4843,6 +4872,14 @@ Summary:"""
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
# A turn that overran the run deadline or max_execution_time
|
||||
# must not count as progress.
|
||||
_timed_out = _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at,
|
||||
turn_timed_out=getattr(self, "_last_stop_reason", None) == "timeout")
|
||||
if _timed_out is not None:
|
||||
return _timed_out
|
||||
|
||||
response_str = str(response)
|
||||
|
||||
# Spend kill-switch: halt if cumulative cost/tokens exceed cap.
|
||||
@@ -5008,7 +5045,7 @@ Summary:"""
|
||||
# the task naturally (model stopped calling tools = done).
|
||||
# This is the same signal the CLI/wrapper path trusts.
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
if self._autonomy_turn_tool_count > 0 and len(response_str) > 100:
|
||||
if not effective_promise and self._autonomy_turn_tool_count > 0 and len(response_str) > 100:
|
||||
# P3/G2: Emit completion callback
|
||||
execute_sync_callback('autonomy_complete',
|
||||
completion_reason="tool_completion",
|
||||
@@ -5045,7 +5082,8 @@ Summary:"""
|
||||
if obs is not None:
|
||||
obs.emit(_EvtType.STAGE_ESCALATE, {"from": prev_stage, "to": stage})
|
||||
|
||||
# Check for completion promise FIRST (structured signal)
|
||||
# Completion promise: when configured, it is the only success signal
|
||||
# (heuristic exits above/below are skipped in promise mode)
|
||||
if effective_promise:
|
||||
promise_tag = f"<promise>{effective_promise}</promise>"
|
||||
if promise_tag in response_str:
|
||||
@@ -5067,7 +5105,7 @@ Summary:"""
|
||||
)
|
||||
|
||||
# Check for keyword-based completion signals (word-boundary, G-COMPLETION-1 fix)
|
||||
if self._is_completion_signal(response_str):
|
||||
if not effective_promise and self._is_completion_signal(response_str):
|
||||
# P3/G2: Emit completion callback
|
||||
execute_sync_callback('autonomy_complete',
|
||||
completion_reason="goal",
|
||||
@@ -5086,7 +5124,7 @@ Summary:"""
|
||||
)
|
||||
|
||||
# For DIRECT stage, complete after first response
|
||||
if stage == "direct":
|
||||
if not effective_promise and stage == "direct":
|
||||
# P3/G2: Emit completion callback
|
||||
execute_sync_callback('autonomy_complete',
|
||||
completion_reason="goal",
|
||||
@@ -5114,7 +5152,7 @@ Summary:"""
|
||||
# Suppressed when doom recovery is active — once a doom loop
|
||||
# is detected, no_tool_calls (a success exit) should not fire;
|
||||
# the doom loop system manages termination instead.
|
||||
if (self._autonomy_turn_tool_count == 0 and iterations > 1
|
||||
if (not effective_promise and self._autonomy_turn_tool_count == 0 and iterations > 1
|
||||
and not self._doom_recovery_active):
|
||||
self._consecutive_no_tool_turns += 1
|
||||
if self._consecutive_no_tool_turns >= 2:
|
||||
@@ -5226,7 +5264,7 @@ Summary:"""
|
||||
|
||||
asyncio.run(main())
|
||||
"""
|
||||
from .autonomy import AutonomyResult
|
||||
from .autonomy import AutonomyResult, _timeout_result
|
||||
import time as time_module
|
||||
from datetime import datetime, timezone
|
||||
import asyncio
|
||||
@@ -5253,6 +5291,11 @@ Summary:"""
|
||||
effective_promise = completion_promise
|
||||
if effective_promise is None:
|
||||
effective_promise = self.autonomy_config.get("completion_promise")
|
||||
# The promise is the only success signal, so tell the model about it
|
||||
# unless the caller already did (e.g. the `loop` CLI instructions).
|
||||
_promise_tag = f"<promise>{effective_promise}</promise>"
|
||||
if effective_promise and _promise_tag not in prompt and _promise_tag not in (self.instructions or ""):
|
||||
prompt = f"{prompt}\n\nWhen the task is fully complete, include {_promise_tag} in your response."
|
||||
|
||||
# Get clear_context from config if not explicitly set
|
||||
effective_clear_context = clear_context
|
||||
@@ -5273,17 +5316,10 @@ Summary:"""
|
||||
iterations += 1
|
||||
|
||||
# Check timeout
|
||||
if timeout_seconds and (time_module.time() - start_time) > timeout_seconds:
|
||||
return AutonomyResult(
|
||||
success=False,
|
||||
output="Task timed out",
|
||||
completion_reason="timeout",
|
||||
iterations=iterations,
|
||||
stage=stage,
|
||||
actions=actions_taken,
|
||||
duration_seconds=time_module.time() - start_time,
|
||||
started_at=started_at,
|
||||
)
|
||||
_timed_out = _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at)
|
||||
if _timed_out is not None:
|
||||
return _timed_out
|
||||
|
||||
# G2: Check for interrupt request (cooperative cancellation)
|
||||
if self.interrupt_controller and self.interrupt_controller.is_set():
|
||||
@@ -5307,9 +5343,17 @@ Summary:"""
|
||||
# Reset the per-turn warning latch and inject any queued nudge.
|
||||
turn_prompt = self._consume_self_correction(prompt)
|
||||
try:
|
||||
response = await self.achat(turn_prompt)
|
||||
with _deadline_scope(timeout_seconds and timeout_seconds - (time_module.time() - start_time)):
|
||||
response = await self.achat(turn_prompt)
|
||||
except Exception as e:
|
||||
return AutonomyResult(
|
||||
# The run deadline cancels an overrunning turn; report it as a timeout.
|
||||
# A budget stop keeps its typed budget_exhausted result when a cap applies.
|
||||
return _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at
|
||||
) or (isinstance(e, BudgetExceededError) and self._autonomy_budget_result(
|
||||
iterations, stage, actions_taken, start_time, started_at,
|
||||
"", spend_baseline
|
||||
)) or AutonomyResult(
|
||||
success=False,
|
||||
output=str(e),
|
||||
completion_reason="error",
|
||||
@@ -5321,6 +5365,14 @@ Summary:"""
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
# A turn that overran the run deadline or max_execution_time
|
||||
# must not count as progress.
|
||||
_timed_out = _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at,
|
||||
turn_timed_out=getattr(self, "_last_stop_reason", None) == "timeout")
|
||||
if _timed_out is not None:
|
||||
return _timed_out
|
||||
|
||||
response_str = str(response)
|
||||
|
||||
# Spend kill-switch: halt if cumulative cost/tokens exceed cap.
|
||||
@@ -5475,7 +5527,7 @@ Summary:"""
|
||||
# produced a substantive response, the inner loop completed
|
||||
# the task naturally (model stopped calling tools = done).
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
if self._autonomy_turn_tool_count > 0 and len(response_str) > 100:
|
||||
if not effective_promise and self._autonomy_turn_tool_count > 0 and len(response_str) > 100:
|
||||
return AutonomyResult(
|
||||
success=True,
|
||||
output=response_str,
|
||||
@@ -5498,7 +5550,8 @@ Summary:"""
|
||||
if obs is not None:
|
||||
obs.emit(_EvtType.STAGE_ESCALATE, {"from": stage_order[idx], "to": stage})
|
||||
|
||||
# Check for completion promise FIRST (structured signal)
|
||||
# Completion promise: when configured, it is the only success signal
|
||||
# (heuristic exits above/below are skipped in promise mode)
|
||||
if effective_promise:
|
||||
promise_tag = f"<promise>{effective_promise}</promise>"
|
||||
if promise_tag in response_str:
|
||||
@@ -5514,7 +5567,7 @@ Summary:"""
|
||||
)
|
||||
|
||||
# Check for keyword-based completion signals (word-boundary, G-COMPLETION-1 fix)
|
||||
if self._is_completion_signal(response_str):
|
||||
if not effective_promise and self._is_completion_signal(response_str):
|
||||
return AutonomyResult(
|
||||
success=True,
|
||||
output=response_str,
|
||||
@@ -5527,7 +5580,7 @@ Summary:"""
|
||||
)
|
||||
|
||||
# For DIRECT stage, complete after first response
|
||||
if stage == "direct":
|
||||
if not effective_promise and stage == "direct":
|
||||
return AutonomyResult(
|
||||
success=True,
|
||||
output=response_str,
|
||||
@@ -5548,7 +5601,7 @@ Summary:"""
|
||||
# Suppressed when doom recovery is active — once a doom loop
|
||||
# is detected, no_tool_calls (a success exit) should not fire;
|
||||
# the doom loop system manages termination instead.
|
||||
if (self._autonomy_turn_tool_count == 0 and iterations > 1
|
||||
if (not effective_promise and self._autonomy_turn_tool_count == 0 and iterations > 1
|
||||
and not self._doom_recovery_active):
|
||||
self._consecutive_no_tool_turns += 1
|
||||
if self._consecutive_no_tool_turns >= 2:
|
||||
@@ -6662,10 +6715,19 @@ Summary:"""
|
||||
if hasattr(self, 'llm_instance') and self.llm_instance:
|
||||
llm = self.llm_instance
|
||||
|
||||
# rerank=True without rerank_model has no backend reranker,
|
||||
# so use the zero-dependency keyword reranker; with
|
||||
# rerank_model the vector store already reranks.
|
||||
reranker = None
|
||||
if rag_config_obj.rerank and not getattr(self._retrieval_config, "rerank_model", None):
|
||||
from praisonaiagents.rag.retriever import SimpleReranker
|
||||
reranker = SimpleReranker()
|
||||
|
||||
self._rag_instance = RAG(
|
||||
knowledge=self.knowledge,
|
||||
config=rag_config_obj,
|
||||
llm=llm,
|
||||
reranker=reranker,
|
||||
)
|
||||
except ImportError:
|
||||
logging.warning("RAG module not available. Install with: pip install 'praisonaiagents[rag]'")
|
||||
|
||||
@@ -400,6 +400,26 @@ class AutonomyResult:
|
||||
"""
|
||||
return self.output or ""
|
||||
|
||||
|
||||
def _timeout_result(timeout_seconds, iterations, stage, actions, start_time, started_at,
|
||||
turn_timed_out=False):
|
||||
"""Return a ``timeout`` AutonomyResult once a run is past ``timeout_seconds``
|
||||
or the last turn was cut by ``max_execution_time`` (``turn_timed_out``), else None."""
|
||||
import time
|
||||
if not turn_timed_out and (not timeout_seconds or (time.time() - start_time) <= timeout_seconds):
|
||||
return None
|
||||
return AutonomyResult(
|
||||
success=False,
|
||||
output="Task timed out",
|
||||
completion_reason="timeout",
|
||||
iterations=iterations,
|
||||
stage=stage,
|
||||
actions=actions,
|
||||
duration_seconds=time.time() - start_time,
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
|
||||
class DoomLoopTracker:
|
||||
"""Tracks actions to detect doom loops with graduated recovery.
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ and response processing. Extracted from agent.py for maintainability.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import contextvars
|
||||
import os
|
||||
import re
|
||||
@@ -29,6 +30,14 @@ _turns_lock = threading.Lock()
|
||||
# into tool worker threads with the context, so each turn has its own counters and
|
||||
# overlapping turns on a shared agent never exhaust each other's limits.
|
||||
_turn_loop_guard = contextvars.ContextVar("praisonai_turn_loop_guard", default=None)
|
||||
# Agent whose streamed turn already records agent_start/agent_end, set while its
|
||||
# non-streaming chat() fallback runs so that turn is not recorded twice.
|
||||
_stream_traced_agent = contextvars.ContextVar("praisonai_stream_traced_agent", default=None)
|
||||
|
||||
|
||||
def _agent_end_status(db_status: str) -> str:
|
||||
"""``agent_end`` status for replay: completed, failed or cancelled."""
|
||||
return "failed" if db_status == "error" else db_status
|
||||
|
||||
|
||||
def _enter_turn(agent, cancel):
|
||||
@@ -97,20 +106,50 @@ def _sleep_unless_cancelled(seconds: float, cancel_token: Any = None, step: floa
|
||||
time.sleep(min(step, remaining))
|
||||
|
||||
|
||||
# ``time.monotonic()`` deadline for turns started in this context, set by
|
||||
# ``run_autonomous(timeout_seconds=...)`` so a single long turn cannot overrun it.
|
||||
_run_deadline = contextvars.ContextVar("praisonai_run_deadline", default=None)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _deadline_scope(seconds):
|
||||
"""Bound every turn started in this block to ``seconds`` from now (no-op when falsy)."""
|
||||
if not seconds:
|
||||
yield
|
||||
return
|
||||
deadline = time.monotonic() + seconds
|
||||
outer = _run_deadline.get()
|
||||
token = _run_deadline.set(deadline if outer is None else min(deadline, outer))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_run_deadline.reset(token)
|
||||
|
||||
|
||||
class _TurnCancelToken:
|
||||
"""Track cancellation observed by one call without shared backend state."""
|
||||
"""Track cancellation observed by one call without shared backend state.
|
||||
|
||||
An optional ``deadline`` (``time.monotonic()``) makes the token report set
|
||||
once it passes, so the existing cancel checks also stop an overrunning turn.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source: Any,
|
||||
*,
|
||||
turn_id: Optional[int] = None,
|
||||
deadline: Optional[float] = None,
|
||||
):
|
||||
self._source = source
|
||||
self._observed = False
|
||||
self._turn_id = turn_id
|
||||
self._deadline = deadline
|
||||
self.timed_out = False
|
||||
|
||||
def is_set(self) -> bool:
|
||||
if self._deadline is not None and time.monotonic() >= self._deadline:
|
||||
self.timed_out = self._observed = True
|
||||
return True
|
||||
turn_checker = getattr(self._source, "_turn_is_cancelled", None)
|
||||
if self._turn_id is not None and callable(turn_checker):
|
||||
is_set = bool(turn_checker(self._turn_id))
|
||||
@@ -121,6 +160,8 @@ class _TurnCancelToken:
|
||||
|
||||
@property
|
||||
def reason(self):
|
||||
if self.timed_out:
|
||||
return "timeout"
|
||||
return getattr(self._source, "reason", None)
|
||||
|
||||
def check(self) -> None:
|
||||
@@ -157,13 +198,22 @@ class ChatMixin:
|
||||
return get_durable_run()
|
||||
|
||||
def _turn_cancel_token(self, source: Any, *, explicit: bool):
|
||||
"""Create per-turn cancellation state while consuming default requests once."""
|
||||
"""Create per-turn cancellation state while consuming default requests once.
|
||||
|
||||
``ExecutionConfig.max_execution_time`` and any enclosing
|
||||
``run_autonomous`` timeout become the token's deadline.
|
||||
"""
|
||||
limit = getattr(self, "max_execution_time", None)
|
||||
deadline = time.monotonic() + limit if limit else None
|
||||
outer = _run_deadline.get()
|
||||
if outer is not None:
|
||||
deadline = outer if deadline is None else min(deadline, outer)
|
||||
if source is None:
|
||||
return None
|
||||
return None if deadline is None else _TurnCancelToken(None, deadline=deadline)
|
||||
begin_turn = getattr(source, "_begin_turn", None)
|
||||
if explicit or not callable(begin_turn):
|
||||
return _TurnCancelToken(source)
|
||||
return _TurnCancelToken(source, turn_id=begin_turn())
|
||||
return _TurnCancelToken(source, deadline=deadline)
|
||||
return _TurnCancelToken(source, turn_id=begin_turn(), deadline=deadline)
|
||||
|
||||
@staticmethod
|
||||
def _notify_tool_call(tool_name, tool_input, tool_result, elapsed_time=None, success=True):
|
||||
@@ -1912,6 +1962,107 @@ Your Goal: {self.goal}"""
|
||||
self._openai_client = get_openai_client(api_key=getattr(self, 'api_key', None))
|
||||
return create_llm_dispatcher(openai_client=self._openai_client, model=model)
|
||||
|
||||
def _fallback_llm_for(self, model):
|
||||
"""Build a throwaway LLM bound to fallback ``model`` for the custom-LLM path.
|
||||
|
||||
Never assigned to ``self.llm_instance``, so concurrent turns keep the
|
||||
primary model. A fallback entry that is already a model object is used
|
||||
as-is. Option params (metrics, max_iter, retries, ...) are inherited;
|
||||
connection settings (base_url, api_key, auth) only when the fallback
|
||||
targets the same provider prefix as the primary.
|
||||
"""
|
||||
if not isinstance(model, str):
|
||||
return model
|
||||
from ..llm.llm import LLM
|
||||
params = dict(getattr(self, '_llm_init_params', None) or {})
|
||||
primary = str(params.get('model') or self.llm or '')
|
||||
# Bare names (no "provider/" prefix) share the configured server.
|
||||
prefix = lambda name: name.split('/', 1)[0] if '/' in name else ''
|
||||
if prefix(primary) != prefix(model):
|
||||
for key in ('base_url', 'api_key', 'auth'):
|
||||
params.pop(key, None)
|
||||
params['model'] = model
|
||||
return LLM(**params)
|
||||
|
||||
def _classify_custom_llm_fallback(self, error, failed_model, fallback_index):
|
||||
"""Return ``(next_model, category)`` if ``error`` warrants a model fallback."""
|
||||
next_model = self._get_next_fallback_model(fallback_index)
|
||||
if next_model is None:
|
||||
return None, None
|
||||
from ..llm.error_classifier import classify_llm_error
|
||||
model_name = str(failed_model)
|
||||
provider = model_name.split('/', 1)[0] if '/' in model_name else "openai"
|
||||
classification = classify_llm_error(error, provider=provider, model=model_name)
|
||||
if not classification.should_fallback_model:
|
||||
return None, None
|
||||
logging.info(f"[{self.name}] {model_name} unavailable — falling back to {next_model}")
|
||||
return next_model, classification.error_category
|
||||
|
||||
@staticmethod
|
||||
def _track_tool_runs(llm_kwargs):
|
||||
"""Return ``(llm_kwargs, ran)``; ``ran`` becomes non-empty once a tool executes.
|
||||
|
||||
A fallback replays the whole turn from the original messages, so it is
|
||||
only safe while no tool has run (otherwise side effects would repeat).
|
||||
"""
|
||||
import functools
|
||||
import inspect
|
||||
ran = []
|
||||
fn = llm_kwargs.get('execute_tool_fn')
|
||||
if fn is None:
|
||||
return llm_kwargs, ran
|
||||
if inspect.iscoroutinefunction(fn):
|
||||
@functools.wraps(fn)
|
||||
async def tracked(*args, **kwargs):
|
||||
ran.append(True)
|
||||
return await fn(*args, **kwargs)
|
||||
else:
|
||||
@functools.wraps(fn)
|
||||
def tracked(*args, **kwargs):
|
||||
ran.append(True)
|
||||
return fn(*args, **kwargs)
|
||||
return {**llm_kwargs, 'execute_tool_fn': tracked}, ran
|
||||
|
||||
def _custom_llm_get_response(self, llm_kwargs):
|
||||
"""Call ``llm_instance.get_response`` walking ``fallback_models`` on outage."""
|
||||
llm, index, current = self.llm_instance, 0, self.llm
|
||||
stream_cb = self.stream_emitter.emit if hasattr(self, 'stream_emitter') else None
|
||||
llm_kwargs, tools_ran = self._track_tool_runs(llm_kwargs)
|
||||
while True:
|
||||
try:
|
||||
with self._llm_usage_sink(llm):
|
||||
return llm.get_response(**llm_kwargs)
|
||||
except (ToolExecutionError, BudgetExceededError, InterruptedError):
|
||||
raise
|
||||
except Exception as e:
|
||||
if tools_ran:
|
||||
raise
|
||||
next_model, category = self._classify_custom_llm_fallback(e, current, index)
|
||||
if next_model is None:
|
||||
raise
|
||||
self._emit_model_fallback(current, next_model, category, index, stream_cb)
|
||||
llm, index, current = self._fallback_llm_for(next_model), index + 1, next_model
|
||||
|
||||
async def _custom_llm_get_response_async(self, llm_kwargs):
|
||||
"""Async twin of :meth:`_custom_llm_get_response`."""
|
||||
llm, index, current = self.llm_instance, 0, self.llm
|
||||
stream_cb = self.stream_emitter.emit if hasattr(self, 'stream_emitter') else None
|
||||
llm_kwargs, tools_ran = self._track_tool_runs(llm_kwargs)
|
||||
while True:
|
||||
try:
|
||||
with self._llm_usage_sink(llm):
|
||||
return await llm.get_response_async(**llm_kwargs)
|
||||
except (ToolExecutionError, BudgetExceededError, InterruptedError):
|
||||
raise
|
||||
except Exception as e:
|
||||
if tools_ran:
|
||||
raise
|
||||
next_model, category = self._classify_custom_llm_fallback(e, current, index)
|
||||
if next_model is None:
|
||||
raise
|
||||
await self._aemit_model_fallback(current, next_model, category, index, stream_cb)
|
||||
llm, index, current = self._fallback_llm_for(next_model), index + 1, next_model
|
||||
|
||||
def _get_model_fallback_lock(self):
|
||||
"""Lazily create the per-instance lock guarding model-fallback swaps.
|
||||
|
||||
@@ -2075,6 +2226,82 @@ Your Goal: {self.goal}"""
|
||||
return max_retries
|
||||
return 2
|
||||
|
||||
def _extract_llm_usage(self, response, model=None):
|
||||
"""Return ``(prompt_tokens, completion_tokens, cost_usd)`` for one LLM response priced as ``model``."""
|
||||
usage = getattr(response, 'usage', None) if response else None
|
||||
if not usage:
|
||||
return 0, 0, 0.0
|
||||
# Chat Completions names, then the Responses API names.
|
||||
prompt_tokens = getattr(usage, 'prompt_tokens', 0) or getattr(usage, 'input_tokens', 0) or 0
|
||||
completion_tokens = getattr(usage, 'completion_tokens', 0) or getattr(usage, 'output_tokens', 0) or 0
|
||||
# Calculate cost using litellm (if available) or fallback pricing
|
||||
cost_usd = self._calculate_llm_cost(prompt_tokens, completion_tokens, response=response, model=model)
|
||||
return prompt_tokens, completion_tokens, cost_usd
|
||||
|
||||
def _record_llm_usage(self, prompt_tokens, completion_tokens, cost_usd):
|
||||
"""Add one LLM call to the agent's cost/token totals and enforce ``max_budget``.
|
||||
|
||||
The single accounting point for every path: the OpenAI-native
|
||||
completions call it directly and the custom-LLM (``provider/model``)
|
||||
tool loop reaches it through the LLM usage sink, once per call.
|
||||
"""
|
||||
# Thread-safe cost tracking (Gap 1a fix)
|
||||
with self._cost_lock:
|
||||
self._total_cost += cost_usd
|
||||
self._total_tokens_in += prompt_tokens
|
||||
self._total_tokens_out += completion_tokens
|
||||
self._llm_call_count += 1
|
||||
budget_exceeded = self._max_budget and self._total_cost >= self._max_budget
|
||||
current_cost = self._total_cost
|
||||
|
||||
if budget_exceeded:
|
||||
if self._on_budget_exceeded == "stop":
|
||||
raise BudgetExceededError(
|
||||
f"Agent '{self.name}' exceeded budget: ${current_cost:.4f} >= ${self._max_budget:.4f}",
|
||||
budget_type="cost",
|
||||
limit=self._max_budget,
|
||||
used=current_cost,
|
||||
agent_id=self.name
|
||||
)
|
||||
elif self._on_budget_exceeded == "warn":
|
||||
logging.warning(
|
||||
f"[budget] {self.name}: ${current_cost:.4f} exceeded "
|
||||
f"${self._max_budget:.4f} budget"
|
||||
)
|
||||
elif callable(self._on_budget_exceeded):
|
||||
self._on_budget_exceeded(current_cost, self._max_budget)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _llm_usage_sink(self, llm=None):
|
||||
"""Route every call ``llm`` (default: the custom LLM) makes in this block into ``_record_llm_usage``.
|
||||
|
||||
Refuses to start when ``max_budget`` is already spent (``"stop"`` mode),
|
||||
so retrying after ``BudgetExceededError`` sends no further paid request.
|
||||
"""
|
||||
if self._max_budget and self._on_budget_exceeded == "stop":
|
||||
with self._cost_lock:
|
||||
current_cost = self._total_cost
|
||||
if current_cost >= self._max_budget:
|
||||
raise BudgetExceededError(
|
||||
f"Agent '{self.name}' exceeded budget: ${current_cost:.4f} >= ${self._max_budget:.4f}",
|
||||
budget_type="cost",
|
||||
limit=self._max_budget,
|
||||
used=current_cost,
|
||||
agent_id=self.name
|
||||
)
|
||||
from ..llm.llm import _usage_sink
|
||||
llm = llm if llm is not None else self.llm_instance
|
||||
# A fallback model is priced as itself, not as the primary ``self.llm``.
|
||||
model = None if llm is self.llm_instance else getattr(llm, 'model', None)
|
||||
token = _usage_sink.set((
|
||||
llm,
|
||||
lambda response: self._record_llm_usage(*self._extract_llm_usage(response, model)),
|
||||
))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_usage_sink.reset(token)
|
||||
|
||||
def _chat_completion(self, messages, temperature=None, tools=None, stream=None, reasoning_steps=False, task_name=None, task_description=None, task_id=None, response_format=None, _retry_depth=0, _fallback_index=0, cancel_token=None):
|
||||
start_time = time.time()
|
||||
|
||||
@@ -2274,19 +2501,8 @@ Your Goal: {self.goal}"""
|
||||
|
||||
# Emit LLM response trace event with token usage
|
||||
_duration_ms = (time.time() - start_time) * 1000
|
||||
_prompt_tokens = 0
|
||||
_completion_tokens = 0
|
||||
_cost_usd = 0.0
|
||||
|
||||
# Extract token usage from response if available
|
||||
if final_response:
|
||||
_usage = getattr(final_response, 'usage', None)
|
||||
if _usage:
|
||||
_prompt_tokens = getattr(_usage, 'prompt_tokens', 0) or 0
|
||||
_completion_tokens = getattr(_usage, 'completion_tokens', 0) or 0
|
||||
# Calculate cost using litellm (if available) or fallback pricing
|
||||
_cost_usd = self._calculate_llm_cost(_prompt_tokens, _completion_tokens, response=final_response)
|
||||
|
||||
_prompt_tokens, _completion_tokens, _cost_usd = self._extract_llm_usage(final_response)
|
||||
|
||||
_trace_emitter.llm_response(
|
||||
self.name,
|
||||
duration_ms=_duration_ms,
|
||||
@@ -2296,33 +2512,8 @@ Your Goal: {self.goal}"""
|
||||
cost_usd=_cost_usd,
|
||||
)
|
||||
|
||||
# Budget tracking & enforcement (zero overhead when _max_budget is None)
|
||||
# Thread-safe cost tracking (Gap 1a fix)
|
||||
with self._cost_lock:
|
||||
self._total_cost += _cost_usd
|
||||
self._total_tokens_in += _prompt_tokens
|
||||
self._total_tokens_out += _completion_tokens
|
||||
self._llm_call_count += 1
|
||||
budget_exceeded = self._max_budget and self._total_cost >= self._max_budget
|
||||
current_cost = self._total_cost
|
||||
|
||||
if budget_exceeded:
|
||||
if self._on_budget_exceeded == "stop":
|
||||
raise BudgetExceededError(
|
||||
f"Agent '{self.name}' exceeded budget: ${current_cost:.4f} >= ${self._max_budget:.4f}",
|
||||
budget_type="cost",
|
||||
limit=self._max_budget,
|
||||
used=current_cost,
|
||||
agent_id=self.name
|
||||
)
|
||||
elif self._on_budget_exceeded == "warn":
|
||||
logging.warning(
|
||||
f"[budget] {self.name}: ${current_cost:.4f} exceeded "
|
||||
f"${self._max_budget:.4f} budget"
|
||||
)
|
||||
elif callable(self._on_budget_exceeded):
|
||||
self._on_budget_exceeded(current_cost, self._max_budget)
|
||||
|
||||
self._record_llm_usage(_prompt_tokens, _completion_tokens, _cost_usd)
|
||||
|
||||
# Trigger AFTER_LLM hook
|
||||
# Only build the input if an AFTER_LLM hook is actually registered
|
||||
from ..hooks import HookEvent, AfterLLMInput
|
||||
@@ -2991,40 +3182,8 @@ Your Goal: {self.goal}"""
|
||||
)
|
||||
|
||||
# Post-call budget accounting (parity with sync _chat_completion).
|
||||
_prompt_tokens = 0
|
||||
_completion_tokens = 0
|
||||
_cost_usd = 0.0
|
||||
if final_response:
|
||||
_usage = getattr(final_response, 'usage', None)
|
||||
if _usage:
|
||||
_prompt_tokens = getattr(_usage, 'prompt_tokens', 0) or 0
|
||||
_completion_tokens = getattr(_usage, 'completion_tokens', 0) or 0
|
||||
_cost_usd = self._calculate_llm_cost(
|
||||
_prompt_tokens, _completion_tokens, response=final_response
|
||||
)
|
||||
with self._cost_lock:
|
||||
self._total_cost += _cost_usd
|
||||
self._total_tokens_in += _prompt_tokens
|
||||
self._total_tokens_out += _completion_tokens
|
||||
self._llm_call_count += 1
|
||||
budget_exceeded = self._max_budget and self._total_cost >= self._max_budget
|
||||
current_cost = self._total_cost
|
||||
if budget_exceeded:
|
||||
if self._on_budget_exceeded == "stop":
|
||||
raise BudgetExceededError(
|
||||
f"Agent '{self.name}' exceeded budget: ${current_cost:.4f} >= ${self._max_budget:.4f}",
|
||||
budget_type="cost",
|
||||
limit=self._max_budget,
|
||||
used=current_cost,
|
||||
agent_id=self.name
|
||||
)
|
||||
elif self._on_budget_exceeded == "warn":
|
||||
logging.warning(
|
||||
f"[budget] {self.name}: ${current_cost:.4f} exceeded "
|
||||
f"${self._max_budget:.4f} budget"
|
||||
)
|
||||
elif callable(self._on_budget_exceeded):
|
||||
self._on_budget_exceeded(current_cost, self._max_budget)
|
||||
_prompt_tokens, _completion_tokens, _cost_usd = self._extract_llm_usage(final_response)
|
||||
self._record_llm_usage(_prompt_tokens, _completion_tokens, _cost_usd)
|
||||
|
||||
# Trigger AFTER_LLM hook (parity with sync _chat_completion).
|
||||
from ..hooks import AfterLLMInput
|
||||
@@ -3450,7 +3609,9 @@ Your Goal: {self.goal}"""
|
||||
# Emit context trace event (zero overhead when not set)
|
||||
from ..trace.context_events import get_context_emitter
|
||||
_trace_emitter = get_context_emitter()
|
||||
_trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
|
||||
_trace_turn = _stream_traced_agent.get() is not self
|
||||
if _trace_turn:
|
||||
_trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
|
||||
durable_context = None
|
||||
durable_token = None
|
||||
previous_run_state = self._get_db_run_state()
|
||||
@@ -3519,12 +3680,15 @@ Your Goal: {self.goal}"""
|
||||
end_durable_run(durable_token)
|
||||
durable_context.close()
|
||||
try:
|
||||
_trace_emitter.agent_end(self.name)
|
||||
if _trace_turn:
|
||||
_trace_emitter.agent_end(self.name, {"status": _agent_end_status(db_status)})
|
||||
finally:
|
||||
if getattr(self, "_active_turn_token", None) is _cancel:
|
||||
self._active_turn_token = None
|
||||
_exit_turn(self, _cancel_ctx)
|
||||
if _cancel is not None:
|
||||
if getattr(_cancel, "timed_out", False):
|
||||
self._last_stop_reason = "timeout"
|
||||
_cancel.close()
|
||||
|
||||
def _chat_impl(self, prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice=None, seed=None, cancel_token=None):
|
||||
@@ -3851,7 +4015,7 @@ Your Goal: {self.goal}"""
|
||||
if hasattr(self, '_drain_steering_messages') and getattr(self, '_message_steering', None) is not None:
|
||||
llm_kwargs['steering_drain'] = self._drain_steering_messages
|
||||
|
||||
response_text = self.llm_instance.get_response(**llm_kwargs)
|
||||
response_text = self._custom_llm_get_response(llm_kwargs)
|
||||
|
||||
self._add_to_chat_history("assistant", response_text)
|
||||
# Persist assistant message to DB
|
||||
@@ -3882,7 +4046,7 @@ Your Goal: {self.goal}"""
|
||||
# Rollback chat history on guardrail failure
|
||||
self._rollback_chat_history_to(chat_history_length)
|
||||
return self._guardrail_blocked_message(e)
|
||||
except ToolExecutionError:
|
||||
except (ToolExecutionError, BudgetExceededError):
|
||||
raise
|
||||
except Exception as e:
|
||||
self._last_stop_reason = "error"
|
||||
@@ -3890,7 +4054,7 @@ Your Goal: {self.goal}"""
|
||||
self._rollback_chat_history_to(chat_history_length)
|
||||
_get_display_functions()['display_error'](f"Error in LLM chat: {e}")
|
||||
return None
|
||||
except ToolExecutionError:
|
||||
except (ToolExecutionError, BudgetExceededError):
|
||||
raise
|
||||
except Exception as e:
|
||||
self._last_stop_reason = "error"
|
||||
@@ -4368,12 +4532,14 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
await durable_context.aclose()
|
||||
finally:
|
||||
try:
|
||||
_trace_emitter.agent_end(self.name)
|
||||
_trace_emitter.agent_end(self.name, {"status": _agent_end_status(db_status)})
|
||||
finally:
|
||||
if getattr(self, "_active_turn_token", None) is _cancel:
|
||||
self._active_turn_token = None
|
||||
_exit_turn(self, _cancel_ctx)
|
||||
if _cancel is not None:
|
||||
if getattr(_cancel, "timed_out", False):
|
||||
self._last_stop_reason = "timeout"
|
||||
_cancel.close()
|
||||
|
||||
async def _achat_impl(self, prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice=None, seed=None, cancel_token=None):
|
||||
@@ -4584,7 +4750,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if hasattr(self, '_drain_steering_messages') and getattr(self, '_message_steering', None) is not None:
|
||||
llm_kwargs['steering_drain'] = self._drain_steering_messages
|
||||
|
||||
response_text = await self.llm_instance.get_response_async(**llm_kwargs)
|
||||
response_text = await self._custom_llm_get_response_async(llm_kwargs)
|
||||
|
||||
# LLM call succeeded - now it's safe to commit any compacted history
|
||||
if effective_history is not self.chat_history:
|
||||
@@ -4612,7 +4778,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# Rollback chat history on guardrail failure
|
||||
self._rollback_chat_history_to(chat_history_length)
|
||||
return self._guardrail_blocked_message(e)
|
||||
except ToolExecutionError:
|
||||
except (ToolExecutionError, BudgetExceededError):
|
||||
raise
|
||||
except Exception as e:
|
||||
self._last_stop_reason = "error"
|
||||
@@ -5126,6 +5292,9 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
except ToolExecutionError:
|
||||
raise
|
||||
except Exception as e:
|
||||
# The custom-LLM path raises on an exhausted budget (sync parity).
|
||||
if isinstance(e, BudgetExceededError) and self._using_custom_llm:
|
||||
raise
|
||||
self._last_stop_reason = "error"
|
||||
_get_display_functions()['display_error'](f"Error in achat: {e}")
|
||||
if get_logger().getEffectiveLevel() == logging.DEBUG:
|
||||
@@ -5327,9 +5496,26 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# via _enter_turn) so overlapping streams on a shared agent neither
|
||||
# exhaust nor inherit each other's per-turn limits.
|
||||
guard_token = _turn_loop_guard.set([self, None])
|
||||
# Record the streamed turn like chat()/achat() do, so a live replay can
|
||||
# show the agent running while its words stream (no-op when tracing is off).
|
||||
from ..trace.context_events import get_context_emitter
|
||||
_trace_emitter = get_context_emitter()
|
||||
status = "completed"
|
||||
if _trace_emitter._enabled:
|
||||
_trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
|
||||
try:
|
||||
yield from self._start_stream_turn(prompt, **kwargs)
|
||||
# _start_stream_turn returns "failed" when a guardrail or hook blocks
|
||||
# the turn, matching chat(), which records a blocked turn as failed.
|
||||
status = (yield from self._start_stream_turn(prompt, **kwargs)) or status
|
||||
except (GeneratorExit, InterruptedError):
|
||||
status = "cancelled"
|
||||
raise
|
||||
except BaseException:
|
||||
status = "failed"
|
||||
raise
|
||||
finally:
|
||||
if _trace_emitter._enabled:
|
||||
_trace_emitter.agent_end(self.name, {"status": status})
|
||||
try:
|
||||
_turn_loop_guard.reset(guard_token)
|
||||
except (ValueError, LookupError):
|
||||
@@ -5354,7 +5540,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if not _in_ok:
|
||||
logging.warning(f"Agent {getattr(self, 'name', '')}: input blocked by guardrail: {_in_err}")
|
||||
yield f"[Input blocked by guardrail: {_in_err}]"
|
||||
return
|
||||
return "failed"
|
||||
prompt = _in_prompt
|
||||
|
||||
# Same pre-turn checks as chat()/achat(): a streamed turn must not
|
||||
@@ -5375,7 +5561,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
))
|
||||
if self._hook_runner.is_blocked(hook_results):
|
||||
logging.warning(f"Agent {self.name} execution blocked by BEFORE_AGENT hook")
|
||||
return
|
||||
return "failed"
|
||||
for res in hook_results:
|
||||
if res.output and res.output.modified_input and "prompt" in res.output.modified_input:
|
||||
prompt = res.output.modified_input["prompt"]
|
||||
@@ -6195,6 +6381,9 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
else:
|
||||
fallback_kwargs.pop('stream', None)
|
||||
|
||||
# The streamed turn already records agent_start/agent_end; keep chat()
|
||||
# from recording the same turn a second time.
|
||||
_traced = _stream_traced_agent.set(self)
|
||||
try:
|
||||
return self.chat(prompt, **fallback_kwargs)
|
||||
except Exception as fallback_error:
|
||||
@@ -6207,6 +6396,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
except Exception:
|
||||
pass
|
||||
raise fallback_error from streaming_error
|
||||
finally:
|
||||
_stream_traced_agent.reset(_traced)
|
||||
|
||||
def _create_llm_summarize_function(self):
|
||||
"""
|
||||
@@ -6676,6 +6867,57 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
raise RuntimeError("Retry loop completed without returning or raising an exception")
|
||||
|
||||
async def _achat_completion_with_retry(self, messages, temperature=None, tools=None, stream=None, reasoning_steps=False, task_name=None, task_description=None, task_id=None, response_format=None, stream_callback=None, emit_events=True):
|
||||
"""Async model call with retry, recorded as ``llm_request``/``llm_response`` like the sync path.
|
||||
|
||||
``Agent.astart``/``achat`` recorded no model calls before this, so a replay
|
||||
of an async run showed the agent but never what it asked the model.
|
||||
"""
|
||||
from ..trace.context_events import get_context_emitter
|
||||
|
||||
_trace_emitter = get_context_emitter()
|
||||
call = self._achat_completion_with_retry_dispatch(
|
||||
messages, temperature, tools, stream, reasoning_steps, task_name,
|
||||
task_description, task_id, response_format,
|
||||
stream_callback=stream_callback, emit_events=emit_events,
|
||||
)
|
||||
if not _trace_emitter._enabled: # zero overhead when tracing is off
|
||||
return await call
|
||||
_trace_emitter.llm_request(
|
||||
self.name,
|
||||
messages_count=len(messages),
|
||||
tokens_used=0,
|
||||
model=self.llm if isinstance(self.llm, str) else None,
|
||||
messages=messages,
|
||||
)
|
||||
started = time.time()
|
||||
try:
|
||||
response = await call
|
||||
except Exception as e:
|
||||
_trace_emitter.llm_response(
|
||||
self.name,
|
||||
duration_ms=(time.time() - started) * 1000,
|
||||
finish_reason="error",
|
||||
response_content=str(e),
|
||||
)
|
||||
raise
|
||||
prompt_tokens = completion_tokens = 0
|
||||
cost_usd = 0.0
|
||||
usage = getattr(response, "usage", None)
|
||||
if usage:
|
||||
prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
|
||||
completion_tokens = getattr(usage, "completion_tokens", 0) or 0
|
||||
cost_usd = self._calculate_llm_cost(prompt_tokens, completion_tokens, response=response)
|
||||
_trace_emitter.llm_response(
|
||||
self.name,
|
||||
duration_ms=(time.time() - started) * 1000,
|
||||
response_content=self._extract_llm_response_content(response),
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
cost_usd=cost_usd,
|
||||
)
|
||||
return response
|
||||
|
||||
async def _achat_completion_with_retry_dispatch(self, messages, temperature=None, tools=None, stream=None, reasoning_steps=False, task_name=None, task_description=None, task_id=None, response_format=None, stream_callback=None, emit_events=True):
|
||||
"""
|
||||
Async wrapper for _execute_unified_achat_completion that adds jittered exponential backoff retry logic.
|
||||
|
||||
|
||||
@@ -497,157 +497,111 @@ class ExecutionMixin:
|
||||
# The backend's stream method is async, so we need to handle that
|
||||
import asyncio
|
||||
|
||||
async def stream_wrapper():
|
||||
async for chunk in self.backend.stream(prompt, **kwargs):
|
||||
yield chunk
|
||||
|
||||
# Convert async generator to sync generator
|
||||
import threading
|
||||
import queue
|
||||
|
||||
def _to_text(item):
|
||||
# Managed backends may yield event dicts; convert them to the
|
||||
# string chunks expected by Agent and skip other event types
|
||||
# (session.status_idle, etc.).
|
||||
if isinstance(item, str):
|
||||
return item
|
||||
if isinstance(item, dict) and item.get('type') == 'agent.message':
|
||||
content = item.get('content', [])
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
text_parts = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get('type') == 'text':
|
||||
text_parts.append(block.get('text', ''))
|
||||
elif isinstance(block, str):
|
||||
text_parts.append(block)
|
||||
if text_parts:
|
||||
return ''.join(text_parts)
|
||||
return None
|
||||
|
||||
# Convert async generator to sync generator. The backend's async
|
||||
# generator is always driven on a private loop in a background
|
||||
# thread, whether or not the caller already has a running loop, so
|
||||
# the caller receives chunks incrementally and a backend error
|
||||
# (including RuntimeError) can never re-run the stream.
|
||||
def sync_stream():
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
# Already in async context - need to handle differently
|
||||
import concurrent.futures
|
||||
import threading
|
||||
import queue
|
||||
|
||||
result_queue = queue.Queue()
|
||||
exception_holder = [None]
|
||||
|
||||
def run_in_thread():
|
||||
new_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(new_loop)
|
||||
# Bounded queue so an abandoned consumer applies back-pressure
|
||||
# instead of letting the producer fill memory unbounded.
|
||||
# ``cancel_event`` lets the producer stop early when the caller
|
||||
# closes the generator.
|
||||
bridge_queue = queue.Queue(maxsize=64)
|
||||
cancel_event = threading.Event()
|
||||
|
||||
def run_in_thread():
|
||||
new_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(new_loop)
|
||||
try:
|
||||
async def pump():
|
||||
agen = self.backend.stream(prompt, **kwargs)
|
||||
try:
|
||||
async for item in agen:
|
||||
if cancel_event.is_set():
|
||||
break
|
||||
# Block until space is available or the
|
||||
# consumer signals cancellation, so we
|
||||
# never spin nor grow the queue forever.
|
||||
while not cancel_event.is_set():
|
||||
try:
|
||||
bridge_queue.put(('item', item), timeout=0.1)
|
||||
break
|
||||
except queue.Full:
|
||||
continue
|
||||
if not cancel_event.is_set():
|
||||
bridge_queue.put(('done', None))
|
||||
except Exception as e:
|
||||
if not cancel_event.is_set():
|
||||
bridge_queue.put(('error', e))
|
||||
finally:
|
||||
# Run the backend's own cleanup (e.g. stopping
|
||||
# its producer) when the consumer closes early.
|
||||
aclose = getattr(agen, 'aclose', None)
|
||||
if aclose is not None:
|
||||
try:
|
||||
await aclose()
|
||||
except Exception as e:
|
||||
logger.debug(f"Backend stream cleanup failed: {e}")
|
||||
|
||||
new_loop.run_until_complete(pump())
|
||||
finally:
|
||||
try:
|
||||
async def collect():
|
||||
try:
|
||||
async for item in self.backend.stream(prompt, **kwargs):
|
||||
result_queue.put(('item', item))
|
||||
result_queue.put(('done', None))
|
||||
except Exception as e:
|
||||
exception_holder[0] = e
|
||||
result_queue.put(('error', e))
|
||||
|
||||
new_loop.run_until_complete(collect())
|
||||
new_loop.run_until_complete(new_loop.shutdown_asyncgens())
|
||||
finally:
|
||||
new_loop.close()
|
||||
|
||||
thread = threading.Thread(target=run_in_thread)
|
||||
thread.start()
|
||||
|
||||
|
||||
# Daemon thread so a stalled/abandoned producer can never
|
||||
# block interpreter shutdown.
|
||||
thread = threading.Thread(target=run_in_thread, daemon=True)
|
||||
thread.start()
|
||||
|
||||
try:
|
||||
while True:
|
||||
msg_type, data = result_queue.get()
|
||||
msg_type, data = bridge_queue.get()
|
||||
if msg_type == 'item':
|
||||
# For managed backends, we might get event objects
|
||||
# Convert to string format expected by Agent
|
||||
if isinstance(data, dict):
|
||||
if data.get('type') == 'agent.message':
|
||||
content = data.get('content', [])
|
||||
if isinstance(content, list):
|
||||
text_parts = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get('type') == 'text':
|
||||
text_parts.append(block.get('text', ''))
|
||||
elif isinstance(block, str):
|
||||
text_parts.append(block)
|
||||
if text_parts:
|
||||
yield ''.join(text_parts)
|
||||
elif isinstance(content, str):
|
||||
yield content
|
||||
# Skip other event types (session.status_idle, etc.)
|
||||
elif isinstance(data, str):
|
||||
yield data
|
||||
text = _to_text(data)
|
||||
if text is not None:
|
||||
yield text
|
||||
elif msg_type == 'done':
|
||||
break
|
||||
elif msg_type == 'error':
|
||||
raise data
|
||||
|
||||
thread.join()
|
||||
|
||||
except RuntimeError:
|
||||
# No event loop (the common script/CLI case). Drive the
|
||||
# backend's async generator from a background thread and
|
||||
# push each chunk through a queue as it arrives, so the
|
||||
# caller receives chunks incrementally instead of only
|
||||
# after the whole response has been buffered. Mirrors the
|
||||
# running-loop branch above.
|
||||
import threading
|
||||
import queue
|
||||
# Bounded queue so an abandoned consumer applies
|
||||
# back-pressure instead of letting the producer fill memory
|
||||
# unbounded. ``cancel_event`` lets the producer stop early
|
||||
# when the caller closes the generator.
|
||||
bridge_queue = queue.Queue(maxsize=64)
|
||||
cancel_event = threading.Event()
|
||||
|
||||
def run_in_thread_noloop():
|
||||
new_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(new_loop)
|
||||
try:
|
||||
async def pump():
|
||||
try:
|
||||
async for item in self.backend.stream(prompt, **kwargs):
|
||||
if cancel_event.is_set():
|
||||
break
|
||||
# Block until space is available or the
|
||||
# consumer signals cancellation, so we
|
||||
# never spin nor grow the queue forever.
|
||||
while not cancel_event.is_set():
|
||||
try:
|
||||
bridge_queue.put(('item', item), timeout=0.1)
|
||||
break
|
||||
except queue.Full:
|
||||
continue
|
||||
if not cancel_event.is_set():
|
||||
bridge_queue.put(('done', None))
|
||||
except Exception as e:
|
||||
if not cancel_event.is_set():
|
||||
bridge_queue.put(('error', e))
|
||||
|
||||
new_loop.run_until_complete(pump())
|
||||
finally:
|
||||
new_loop.close()
|
||||
|
||||
# Daemon thread so a stalled/abandoned producer can never
|
||||
# block interpreter shutdown.
|
||||
thread = threading.Thread(target=run_in_thread_noloop, daemon=True)
|
||||
thread.start()
|
||||
|
||||
finally:
|
||||
# Runs on normal completion AND on early generator
|
||||
# close/GC: signal the producer to stop, drain one slot
|
||||
# so a blocked put() can observe the cancel, and join.
|
||||
cancel_event.set()
|
||||
try:
|
||||
while True:
|
||||
msg_type, data = bridge_queue.get()
|
||||
if msg_type == 'item':
|
||||
item = data
|
||||
# Same conversion logic as the running-loop branch.
|
||||
if isinstance(item, dict):
|
||||
if item.get('type') == 'agent.message':
|
||||
content = item.get('content', [])
|
||||
if isinstance(content, list):
|
||||
text_parts = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get('type') == 'text':
|
||||
text_parts.append(block.get('text', ''))
|
||||
elif isinstance(block, str):
|
||||
text_parts.append(block)
|
||||
if text_parts:
|
||||
yield ''.join(text_parts)
|
||||
elif isinstance(content, str):
|
||||
yield content
|
||||
elif isinstance(item, str):
|
||||
yield item
|
||||
elif msg_type == 'done':
|
||||
break
|
||||
elif msg_type == 'error':
|
||||
raise data
|
||||
finally:
|
||||
# Runs on normal completion AND on early generator
|
||||
# close/GC: signal the producer to stop, drain one slot
|
||||
# so a blocked put() can observe the cancel, and join.
|
||||
cancel_event.set()
|
||||
try:
|
||||
bridge_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
pass
|
||||
thread.join(timeout=5)
|
||||
|
||||
bridge_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
pass
|
||||
thread.join(timeout=5)
|
||||
|
||||
return sync_stream()
|
||||
|
||||
except Exception as e:
|
||||
@@ -1427,6 +1381,10 @@ Write the complete compiled report:"""
|
||||
_loop_detection.record_tool_outcome(
|
||||
_ld_history, function_name, arguments, _result, _ld_config
|
||||
)
|
||||
# Per-turn tool count for autonomy completion heuristics (sync
|
||||
# parity with tool_execution.py). Counted only once the tool has
|
||||
# actually run, so a BEFORE_TOOL-blocked call cannot count.
|
||||
self._autonomy_turn_tool_count = getattr(self, '_autonomy_turn_tool_count', 0) + 1
|
||||
return _result
|
||||
|
||||
# Enforce BEFORE_TOOL/AFTER_TOOL security hooks for every async caller,
|
||||
@@ -1709,6 +1667,7 @@ Write the complete compiled report:"""
|
||||
async def _execute_tool_async_impl(self, function_name: str, arguments: Dict[str, Any], tool_call_id: Optional[str] = None, tools_override: Optional[List] = None) -> Any:
|
||||
"""Internal async tool execution implementation"""
|
||||
from ..errors import ToolExecutionError
|
||||
from .tool_execution import _CALLER_ERROR_TYPES
|
||||
try:
|
||||
logging.info(f"Executing async tool: {function_name} with arguments: {arguments}")
|
||||
|
||||
@@ -1940,20 +1899,50 @@ Write the complete compiled report:"""
|
||||
with tool_progress_channel(_progress_sink):
|
||||
return await call_target(**call_arguments)
|
||||
logging.debug(f"Executing sync function in executor: {function_name}")
|
||||
loop = asyncio.get_running_loop()
|
||||
from ..trace.context_events import copy_context_to_callable
|
||||
# copy_context_to_callable snapshots the CURRENT context (incl.
|
||||
# the injection contextvar) into the executor thread, so setting
|
||||
# it here propagates cancel_event to the sync tool body too.
|
||||
if _inject_state is not None:
|
||||
with with_injection_context(_inject_state), tool_progress_channel(_progress_sink):
|
||||
return await loop.run_in_executor(
|
||||
None, copy_context_to_callable(lambda: call_target(**call_arguments))
|
||||
)
|
||||
return await _run_sync(copy_context_to_callable(lambda: call_target(**call_arguments)))
|
||||
with tool_progress_channel(_progress_sink):
|
||||
return await loop.run_in_executor(
|
||||
None, copy_context_to_callable(lambda: call_target(**call_arguments))
|
||||
return await _run_sync(copy_context_to_callable(lambda: call_target(**call_arguments)))
|
||||
|
||||
async def _run_sync(fn):
|
||||
# With a tool timeout, a timed-out sync tool keeps its thread
|
||||
# (wait_for cannot stop it). Run it on the bounded per-agent
|
||||
# pool (sync-path parity) instead of the shared default
|
||||
# executor, so hung tools cannot starve asyncio.to_thread
|
||||
# callers elsewhere in the process.
|
||||
if not (tool_timeout and tool_timeout > 0) or not hasattr(self, "_submit_to_tool_pool"):
|
||||
return await asyncio.get_running_loop().run_in_executor(None, fn)
|
||||
# spill=True: when every worker is busy the call runs at once
|
||||
# on a counted overflow thread (executor is None) instead of
|
||||
# queueing into its timeout; refused once that cap is hit.
|
||||
future, executor, orphaned = self._submit_to_tool_pool(fn, spill=True)
|
||||
if future is None:
|
||||
logging.warning(
|
||||
f"Tool worker pool exhausted; refusing to queue '{function_name}' "
|
||||
f"behind {orphaned} hung tool call(s)."
|
||||
)
|
||||
return {
|
||||
"error": (
|
||||
f"Tool worker pool exhausted: {orphaned} earlier tool call(s) "
|
||||
f"are still running past their timeout. Refusing to queue "
|
||||
f"'{function_name}'."
|
||||
),
|
||||
"tool_pool_exhausted": True,
|
||||
"_praison_retryable": False,
|
||||
}
|
||||
try:
|
||||
return await asyncio.wrap_future(future)
|
||||
except asyncio.CancelledError:
|
||||
# Timed out (or caller cancelled) while still running:
|
||||
# retire the stuck pool so the next call gets a fresh worker.
|
||||
if executor is not None and not future.done():
|
||||
self._retire_tool_executor(executor)
|
||||
raise
|
||||
|
||||
# Circuit-breaker parity with the sync path (tool_execution.py).
|
||||
# The async tool path previously had no breaker, so a repeatedly
|
||||
@@ -1973,23 +1962,12 @@ Write the complete compiled report:"""
|
||||
_NeutralOutcome = _CircuitBreakerException
|
||||
try:
|
||||
from ..tools.circuit_breaker import (
|
||||
get_circuit_breaker,
|
||||
CircuitBreakerConfig,
|
||||
CircuitBreakerException as _CircuitBreakerException,
|
||||
_NeutralOutcome,
|
||||
)
|
||||
breaker_name = f"tool_{id(self)}_{function_name}"
|
||||
breaker = get_circuit_breaker(
|
||||
breaker_name,
|
||||
CircuitBreakerConfig(
|
||||
failure_threshold=5,
|
||||
recovery_timeout=60.0,
|
||||
timeout=30.0,
|
||||
graceful_degradation=True,
|
||||
),
|
||||
)
|
||||
if hasattr(self, "_register_breaker_finalizer"):
|
||||
self._register_breaker_finalizer(breaker_name)
|
||||
# Same per-agent breaker and policy as the sync path.
|
||||
if hasattr(self, "_get_tool_circuit_breaker"):
|
||||
breaker = self._get_tool_circuit_breaker(function_name)
|
||||
except ImportError:
|
||||
# Circuit breaker not available - fall back to direct invocation.
|
||||
logging.debug("Circuit breaker not available on async path, executing directly")
|
||||
@@ -2010,7 +1988,7 @@ Write the complete compiled report:"""
|
||||
# half-open one. Unwrapped in _invoke_guarded.
|
||||
try:
|
||||
r = await _invoke()
|
||||
except (ValueError, TypeError, AttributeError) as e:
|
||||
except _CALLER_ERROR_TYPES as e:
|
||||
raise _NeutralOutcome(e)
|
||||
if isinstance(r, dict) and r.get("error") and \
|
||||
r.get("_praison_retryable") is False:
|
||||
@@ -2051,13 +2029,7 @@ Write the complete compiled report:"""
|
||||
# _execute_tool_async_with_retry treats it as terminal.
|
||||
if loop_guard is not None:
|
||||
loop_guard.record(function_name, arguments, False, result=None)
|
||||
return {
|
||||
"error": f"Tool '{function_name}' circuit breaker open - too many recent failures",
|
||||
"circuit_open": True,
|
||||
"agent_name": getattr(self, "name", None),
|
||||
"session_id": getattr(self, "_session_id", None),
|
||||
"remediation": "Wait for recovery_timeout (60s) or investigate recent tool failures.",
|
||||
}
|
||||
return self._circuit_breaker_open_result(function_name)
|
||||
|
||||
# Apply the per-agent tool timeout (ToolConfig.timeout) so the async
|
||||
# path matches the sync path in tool_execution.py. asyncio.wait_for
|
||||
@@ -2217,7 +2189,7 @@ Write the complete compiled report:"""
|
||||
if isinstance(e, ToolExecutionError):
|
||||
retryable = e.is_retryable
|
||||
else:
|
||||
retryable = not isinstance(e, (ValueError, TypeError, AttributeError))
|
||||
retryable = not isinstance(e, _CALLER_ERROR_TYPES)
|
||||
# Private control-plane key: a tool's own result may legitimately
|
||||
# contain "retryable" (common for HTTP-API wrappers) and must not
|
||||
# collide with the framework's verdict nor leak to the model.
|
||||
|
||||
@@ -163,10 +163,14 @@ class ImageAgent(Agent):
|
||||
# Gemini provider doesn't support response_format parameter
|
||||
# Apply this filter if provider is explicitly 'gemini' or as fallback for gemini models
|
||||
config.pop('response_format', None)
|
||||
elif custom_llm_provider == "openai" and config.get('response_format') == "url":
|
||||
# "url" is OpenAI's default for DALL-E and is rejected by models that
|
||||
# don't accept response_format (e.g. gpt-image), so don't send it
|
||||
config.pop('response_format', None)
|
||||
elif custom_llm_provider == "openai":
|
||||
if config.get('response_format') == "url":
|
||||
# "url" is OpenAI's default for DALL-E and is rejected by models that
|
||||
# don't accept response_format (e.g. gpt-image), so don't send it
|
||||
config.pop('response_format', None)
|
||||
if 'dall-e-3' not in model_name.lower():
|
||||
# Only dall-e-3 accepts 'style'; other OpenAI image models reject it
|
||||
config.pop('style', None)
|
||||
from rich.progress import Progress, SpinnerColumn, TextColumn
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
|
||||
@@ -38,10 +38,26 @@ _active_turn_cancel: "contextvars.ContextVar[Optional[tuple]]" = contextvars.Con
|
||||
"praisonai_active_turn_cancel", default=None
|
||||
)
|
||||
|
||||
# Exceptions a tool body raises about its *input* (bad value, missing key or
|
||||
# path). They are terminal for the same arguments and say nothing about the
|
||||
# tool's health, so they are never retried and never count toward the per-tool
|
||||
# circuit breaker. Timeouts, connection errors and other OSErrors still do.
|
||||
_CALLER_ERROR_TYPES = (
|
||||
ValueError, TypeError, AttributeError, LookupError,
|
||||
FileNotFoundError, FileExistsError, IsADirectoryError, NotADirectoryError, PermissionError,
|
||||
)
|
||||
|
||||
# Worker count for the per-agent tool ThreadPoolExecutor. Used both when the
|
||||
# pool is created and to decide (with the per-executor in-flight counter)
|
||||
# whether a fresh call would have to queue behind hung workers.
|
||||
_TOOL_POOL_MAX_WORKERS = 2
|
||||
# whether a fresh call would have to queue behind hung workers. Matches the
|
||||
# default ParallelToolCallExecutor width so a parallel batch never queues.
|
||||
_TOOL_POOL_MAX_WORKERS = 5
|
||||
|
||||
# Upper bound on overflow ("spill") threads a single agent may run at once for
|
||||
# timed async tool calls that find every pool worker busy. A hung spill thread
|
||||
# cannot be reclaimed, so without this cap leaked threads would grow with
|
||||
# concurrency instead of staying bounded (#5356).
|
||||
_MAX_SPILL_TOOL_THREADS = 8
|
||||
|
||||
# Sentinel returned as the "arguments" value when a tool call's argument string
|
||||
# cannot be parsed (e.g. truncated by max_tokens or a dropped connection).
|
||||
@@ -1077,6 +1093,110 @@ class ToolExecutionMixin:
|
||||
from .async_safety import DualLock
|
||||
return self.__dict__.setdefault("_tool_executor_lock", DualLock())
|
||||
|
||||
def _submit_to_tool_pool(self, fn, *args, spill=False):
|
||||
"""Submit ``fn(*args)`` to the bounded per-agent tool pool.
|
||||
|
||||
Shared by the sync and async tool paths so a timed-out tool never
|
||||
strands a thread in the process-wide default executor. Returns
|
||||
``(future, executor, orphaned)``; ``future`` is ``None`` when the pool
|
||||
is exhausted by hung calls and the caller must refuse fast. With
|
||||
``spill=True`` and every worker busy (but not orphan-capped), the call
|
||||
runs immediately on a counted daemon thread instead of waiting for a
|
||||
worker, and ``(future, None, orphaned)`` is returned; once
|
||||
``_MAX_SPILL_TOOL_THREADS`` spill threads are still running the call is
|
||||
refused like an exhausted pool.
|
||||
"""
|
||||
with self._get_tool_executor_lock().sync():
|
||||
if getattr(self, '_tool_executor', None) is None:
|
||||
self._tool_executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=_TOOL_POOL_MAX_WORKERS,
|
||||
thread_name_prefix=f"tool-{self.name}"
|
||||
)
|
||||
# Per-executor in-flight counter so a retired pool's
|
||||
# late-finishing future can never corrupt the *new* pool's count.
|
||||
self._tool_executor._inflight = 0
|
||||
executor = self._tool_executor
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
inflight = getattr(executor, '_inflight', 0)
|
||||
# Once the orphan cap is hit we can no longer retire a stuck pool,
|
||||
# so we would otherwise queue behind hung workers and time out a
|
||||
# healthy tool that never got to start. Refuse fast and visibly
|
||||
# instead when no worker is free, and let the pool self-heal (the
|
||||
# orphan count drops as hung threads finish).
|
||||
if inflight >= _TOOL_POOL_MAX_WORKERS:
|
||||
if orphaned >= _MAX_ORPHANED_TOOL_EXECUTORS:
|
||||
return None, executor, orphaned
|
||||
if spill:
|
||||
spilled = getattr(self, '_tool_spill_inflight', 0)
|
||||
if spilled >= _MAX_SPILL_TOOL_THREADS:
|
||||
return None, executor, orphaned + spilled
|
||||
self._tool_spill_inflight = spilled + 1
|
||||
return self._start_spill_thread(fn, *args), None, orphaned
|
||||
executor._inflight = inflight + 1
|
||||
future = executor.submit(fn, *args)
|
||||
|
||||
def _on_tool_done(_f, _agent=self, _executor=executor): # noqa: ANN001
|
||||
with _agent._get_tool_executor_lock().sync():
|
||||
_executor._inflight = max(0, getattr(_executor, '_inflight', 1) - 1)
|
||||
# Orphan accounting is per-POOL, not per-future: retiring a pool
|
||||
# adds exactly one to the orphan count, so a pool may return
|
||||
# exactly one slot. Recover it only when this retired pool's
|
||||
# LAST in-flight worker exits (``_inflight`` == 0), and guard
|
||||
# with a one-shot flag so a multi-worker pool can never
|
||||
# decrement more than the single unit it contributed and steal
|
||||
# another orphaned pool's slot.
|
||||
if (
|
||||
getattr(_agent, '_tool_executor', None) is not _executor
|
||||
and _executor._inflight == 0
|
||||
and not getattr(_executor, '_orphan_recovered', False)
|
||||
):
|
||||
_executor._orphan_recovered = True
|
||||
_agent._tool_executor_orphaned = max(
|
||||
0, getattr(_agent, '_tool_executor_orphaned', 0) - 1
|
||||
)
|
||||
|
||||
future.add_done_callback(_on_tool_done)
|
||||
return future, executor, orphaned
|
||||
|
||||
def _start_spill_thread(self, fn, *args):
|
||||
"""Run ``fn(*args)`` on a daemon thread; caller already counted it."""
|
||||
future = concurrent.futures.Future()
|
||||
|
||||
def _target():
|
||||
try:
|
||||
if future.set_running_or_notify_cancel():
|
||||
try:
|
||||
future.set_result(fn(*args))
|
||||
except BaseException as e: # noqa: BLE001 - re-raised by awaiter
|
||||
future.set_exception(e)
|
||||
finally:
|
||||
with self._get_tool_executor_lock().sync():
|
||||
self._tool_spill_inflight = max(0, getattr(self, '_tool_spill_inflight', 1) - 1)
|
||||
|
||||
try:
|
||||
threading.Thread(target=_target, daemon=True, name=f"tool-{self.name}-spill").start()
|
||||
except BaseException:
|
||||
self._tool_spill_inflight = max(0, getattr(self, '_tool_spill_inflight', 1) - 1)
|
||||
raise
|
||||
return future
|
||||
|
||||
def _retire_tool_executor(self, executor):
|
||||
"""Retire a pool whose worker is stuck past its timeout.
|
||||
|
||||
The timed-out thread cannot be reclaimed. Retire this executor so the
|
||||
next call gets a fresh worker instead of queuing behind a permanently
|
||||
stuck one — but bound how many orphaned (still-blocked) pools we allow
|
||||
to accumulate: once the cap is reached we stop recycling and keep
|
||||
reusing the existing pool so repeated timeouts can't exhaust process
|
||||
resources with an unbounded number of leaked threads.
|
||||
"""
|
||||
with self._get_tool_executor_lock().sync():
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
if getattr(self, '_tool_executor', None) is executor and orphaned < _MAX_ORPHANED_TOOL_EXECUTORS:
|
||||
executor.shutdown(wait=False)
|
||||
self._tool_executor = None
|
||||
self._tool_executor_orphaned = orphaned + 1
|
||||
|
||||
def _get_turn_tools_lock(self):
|
||||
"""Return the DualLock guarding the per-turn tool buffer.
|
||||
|
||||
@@ -1342,65 +1462,16 @@ class ToolExecutionMixin:
|
||||
with tool_progress_channel(_attempt_sink):
|
||||
# Use copy_context to preserve injection context in executor thread
|
||||
ctx = contextvars.copy_context()
|
||||
started = threading.Event()
|
||||
|
||||
def execute_with_context():
|
||||
started.set()
|
||||
with with_injection_context(state):
|
||||
return self._execute_tool_with_circuit_breaker(function_name, arguments)
|
||||
|
||||
# Use reusable executor to prevent resource leaks
|
||||
pool_exhausted = False
|
||||
with self._get_tool_executor_lock().sync():
|
||||
if not hasattr(self, '_tool_executor') or self._tool_executor is None:
|
||||
self._tool_executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=_TOOL_POOL_MAX_WORKERS,
|
||||
thread_name_prefix=f"tool-{self.name}"
|
||||
)
|
||||
# Per-executor in-flight counter so a
|
||||
# retired pool's late-finishing future can
|
||||
# never corrupt the *new* pool's count.
|
||||
self._tool_executor._inflight = 0
|
||||
executor = self._tool_executor
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
inflight = getattr(executor, '_inflight', 0)
|
||||
# Once the orphan cap is hit we can no longer
|
||||
# retire a stuck pool, so we would otherwise
|
||||
# queue behind hung workers and time out a
|
||||
# healthy tool that never got to start. Refuse
|
||||
# fast and visibly instead when no worker is
|
||||
# free, and let the pool self-heal (the orphan
|
||||
# count drops as hung threads finish).
|
||||
if orphaned >= _MAX_ORPHANED_TOOL_EXECUTORS and inflight >= _TOOL_POOL_MAX_WORKERS:
|
||||
pool_exhausted = True
|
||||
else:
|
||||
executor._inflight = inflight + 1
|
||||
future = executor.submit(ctx.run, execute_with_context)
|
||||
|
||||
def _on_tool_done(_f, _agent=self, _executor=executor): # noqa: ANN001
|
||||
with _agent._get_tool_executor_lock().sync():
|
||||
_executor._inflight = max(0, getattr(_executor, '_inflight', 1) - 1)
|
||||
# Orphan accounting is per-POOL, not
|
||||
# per-future: retiring a pool adds
|
||||
# exactly one to the orphan count, so
|
||||
# a pool may return exactly one slot.
|
||||
# Recover it only when this retired
|
||||
# pool's LAST in-flight worker exits
|
||||
# (``_inflight`` == 0), and guard with
|
||||
# a one-shot flag so a multi-worker
|
||||
# pool can never decrement more than
|
||||
# the single unit it contributed and
|
||||
# steal another orphaned pool's slot.
|
||||
if (
|
||||
getattr(_agent, '_tool_executor', None) is not _executor
|
||||
and _executor._inflight == 0
|
||||
and not getattr(_executor, '_orphan_recovered', False)
|
||||
):
|
||||
_executor._orphan_recovered = True
|
||||
_agent._tool_executor_orphaned = max(
|
||||
0, getattr(_agent, '_tool_executor_orphaned', 0) - 1
|
||||
)
|
||||
|
||||
future.add_done_callback(_on_tool_done)
|
||||
if pool_exhausted:
|
||||
future, executor, orphaned = self._submit_to_tool_pool(ctx.run, execute_with_context)
|
||||
if future is None:
|
||||
logging.warning(
|
||||
f"Tool worker pool exhausted; refusing to queue '{function_name}' "
|
||||
f"behind {orphaned} hung tool call(s)."
|
||||
@@ -1415,6 +1486,13 @@ class ToolExecutionMixin:
|
||||
}
|
||||
break
|
||||
try:
|
||||
# Charge tool_timeout from when the body starts,
|
||||
# not from submit, so time spent queued for a
|
||||
# worker never times out a healthy tool. The
|
||||
# queue wait is itself bounded by tool_timeout
|
||||
# so a pool stuck behind hung workers still fails.
|
||||
if not started.wait(timeout=tool_timeout):
|
||||
raise concurrent.futures.TimeoutError()
|
||||
result = future.result(timeout=tool_timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
if _progress_active is not None:
|
||||
@@ -1427,19 +1505,7 @@ class ToolExecutionMixin:
|
||||
# inner budget never saw it, so this one still
|
||||
# earns an outer retry pass under the single budget.
|
||||
result = {"error": f"Tool timed out after {tool_timeout}s", "timeout": True, "_outer_timeout": True}
|
||||
# The timed-out thread cannot be reclaimed. Retire this
|
||||
# executor so the next call gets a fresh worker instead of
|
||||
# queuing behind a permanently stuck one — but bound how many
|
||||
# orphaned (still-blocked) threads we allow to accumulate:
|
||||
# once the cap is reached we stop recycling and keep reusing
|
||||
# the existing pool so repeated timeouts can't exhaust process
|
||||
# resources with an unbounded number of leaked threads.
|
||||
with self._get_tool_executor_lock().sync():
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
if getattr(self, '_tool_executor', None) is executor and orphaned < _MAX_ORPHANED_TOOL_EXECUTORS:
|
||||
executor.shutdown(wait=False)
|
||||
self._tool_executor = None
|
||||
self._tool_executor_orphaned = orphaned + 1
|
||||
self._retire_tool_executor(executor)
|
||||
else:
|
||||
with tool_progress_channel(_progress_sink), with_injection_context(state):
|
||||
result = self._execute_tool_with_circuit_breaker(function_name, arguments)
|
||||
@@ -1678,8 +1744,28 @@ class ToolExecutionMixin:
|
||||
# `if artifact_ref:` guards remain valid on short-result paths.
|
||||
artifact_ref = None
|
||||
|
||||
# ToolConfig.output_max_lines caps lines first, then the
|
||||
# char limit below applies; output_direction picks the end(s) kept.
|
||||
_tool_cfg = getattr(self, '_tool_config', None)
|
||||
direction = getattr(_tool_cfg, 'output_direction', 'both')
|
||||
max_lines = getattr(_tool_cfg, 'output_max_lines', None)
|
||||
capped_str = result_str
|
||||
if max_lines:
|
||||
lines = result_str.split("\n")
|
||||
if len(lines) > max_lines:
|
||||
if direction == "head":
|
||||
kept_head, kept_tail = lines[:max_lines], []
|
||||
elif direction == "tail":
|
||||
kept_head, kept_tail = [], lines[-max_lines:]
|
||||
else:
|
||||
n_tail = max_lines // 2
|
||||
kept_head = lines[:max_lines - n_tail]
|
||||
kept_tail = lines[-n_tail:] if n_tail else []
|
||||
marker = f"...[{len(lines):,} lines total, showing {max_lines:,}]..."
|
||||
capped_str = "\n".join(kept_head + [marker] + kept_tail)
|
||||
|
||||
# Check if we need to spill to artifact store
|
||||
if len(result_str) > limit:
|
||||
if len(capped_str) > limit or capped_str is not result_str:
|
||||
# Try to use artifact store if available
|
||||
if hasattr(self, '_artifact_store') and self._artifact_store is not None:
|
||||
try:
|
||||
@@ -1704,22 +1790,25 @@ class ToolExecutionMixin:
|
||||
except Exception as e:
|
||||
logging.debug(f"Failed to store artifact: {e}")
|
||||
|
||||
# Generate truncated preview
|
||||
tail_size = min(limit // 5, 2000)
|
||||
head = result_str[:limit - tail_size]
|
||||
tail = result_str[-tail_size:] if tail_size > 0 else ""
|
||||
|
||||
if len(capped_str) <= limit:
|
||||
# Only the line cap applied
|
||||
truncated = capped_str
|
||||
else:
|
||||
# Generate truncated preview
|
||||
if direction == "head":
|
||||
head, tail = capped_str[:limit], ""
|
||||
elif direction == "tail":
|
||||
head, tail = "", capped_str[-limit:]
|
||||
else:
|
||||
tail_size = min(limit // 5, 2000)
|
||||
head = capped_str[:limit - tail_size]
|
||||
tail = capped_str[-tail_size:] if tail_size > 0 else ""
|
||||
total = "chars total" if artifact_ref else "chars"
|
||||
truncated = f"{head}\n...[{len(result_str):,} {total}, showing first/last portions]...\n{tail}"
|
||||
|
||||
# If we stored an artifact, include reference in the output
|
||||
if artifact_ref:
|
||||
truncated = (
|
||||
f"{head}\n"
|
||||
f"...[{len(result_str):,} chars total, showing first/last portions]...\n"
|
||||
f"{tail}\n\n"
|
||||
f"{artifact_ref.to_inline()}"
|
||||
)
|
||||
else:
|
||||
# Fallback to simple truncation
|
||||
truncated = f"{head}\n...[{len(result_str):,} chars, showing first/last portions]...\n{tail}"
|
||||
truncated = f"{truncated}\n\n{artifact_ref.to_inline()}"
|
||||
else:
|
||||
truncated = result_str
|
||||
|
||||
@@ -1737,10 +1826,15 @@ class ToolExecutionMixin:
|
||||
else:
|
||||
truncated = self._truncate_tool_output(function_name, truncated, tool_call_id)
|
||||
|
||||
if len(truncated) < len(result_str):
|
||||
# A triggered line cap applies even when its marker makes the
|
||||
# preview no shorter than the original (many short lines).
|
||||
line_capped = capped_str is not result_str
|
||||
if line_capped or len(truncated) < len(result_str):
|
||||
logging.debug(f"Truncated {function_name} output from {len(result_str)} to {len(truncated)} chars")
|
||||
# For dicts, truncate large string fields (e.g., raw_content from search)
|
||||
if isinstance(result, dict):
|
||||
# For dicts, truncate large string fields (e.g., raw_content from search).
|
||||
# Field truncation keeps both ends, so a line cap or a head/tail
|
||||
# direction uses the string preview instead.
|
||||
if isinstance(result, dict) and direction == "both" and not line_capped:
|
||||
max_field_chars = getattr(self, 'tool_output_limit', DEFAULT_TOOL_OUTPUT_LIMIT) if not self.context_manager else None
|
||||
result = self._truncate_dict_fields(result, function_name, max_field_chars, tool_call_id)
|
||||
# Per-field truncation only shortens individual fields
|
||||
@@ -2126,7 +2220,7 @@ class ToolExecutionMixin:
|
||||
|
||||
return response
|
||||
|
||||
def _calculate_llm_cost(self, prompt_tokens: int, completion_tokens: int, response: any = None) -> float:
|
||||
def _calculate_llm_cost(self, prompt_tokens: int, completion_tokens: int, response: any = None, model: Optional[str] = None) -> float:
|
||||
"""Calculate estimated cost for LLM call.
|
||||
|
||||
Uses litellm for accurate pricing (1000+ models) when available,
|
||||
@@ -2136,6 +2230,7 @@ class ToolExecutionMixin:
|
||||
prompt_tokens: Number of tokens in the prompt
|
||||
completion_tokens: Number of tokens in the completion
|
||||
response: Optional LLM response object for more accurate cost calculation
|
||||
model: Model that answered (e.g. a fallback model); defaults to ``self.llm``
|
||||
|
||||
Returns:
|
||||
Estimated cost in USD
|
||||
@@ -2144,7 +2239,7 @@ class ToolExecutionMixin:
|
||||
return calculate_llm_cost(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
model=self.llm,
|
||||
model=model or self.llm,
|
||||
response=response,
|
||||
)
|
||||
|
||||
@@ -3227,23 +3322,28 @@ class ToolExecutionMixin:
|
||||
return None
|
||||
breaker_name = f"tool_{id(self)}_{function_name}"
|
||||
config = CircuitBreakerConfig(
|
||||
failure_threshold=5,
|
||||
recovery_timeout=60.0,
|
||||
timeout=30.0,
|
||||
graceful_degradation=True,
|
||||
failure_threshold=5, # Open after 5 failures
|
||||
recovery_timeout=60.0, # Wait 60s before trying half-open
|
||||
graceful_degradation=True, # Return error instead of raising exception
|
||||
)
|
||||
breaker = get_circuit_breaker(breaker_name, config)
|
||||
# Ensure the registry entry is removed the moment this Agent is
|
||||
# actually garbage-collected, regardless of whether close()/aclose()
|
||||
# was ever called. This closes the CPython id-reuse window where a
|
||||
# new Agent at the same address could inherit a stale OPEN breaker.
|
||||
self._register_breaker_finalizer(breaker_name)
|
||||
return breaker
|
||||
|
||||
def _circuit_breaker_open_result(self, function_name):
|
||||
"""Return the standard ``circuit_open`` error dict for *function_name*."""
|
||||
breaker = self._get_tool_circuit_breaker(function_name)
|
||||
recovery = breaker.config.recovery_timeout if breaker is not None else 60.0
|
||||
return {
|
||||
"error": f"Tool '{function_name}' circuit breaker open - too many recent failures",
|
||||
"circuit_open": True,
|
||||
"agent_name": getattr(self, "name", None),
|
||||
"session_id": getattr(self, "_session_id", None),
|
||||
"remediation": "Wait for recovery_timeout (60s) or investigate recent tool failures.",
|
||||
"remediation": f"Wait for recovery_timeout ({recovery:g}s) or investigate recent tool failures.",
|
||||
}
|
||||
|
||||
def _circuit_breaker_precheck(self, function_name):
|
||||
@@ -3301,7 +3401,7 @@ class ToolExecutionMixin:
|
||||
"""
|
||||
# Import circuit breaker components first (lazy import for performance)
|
||||
try:
|
||||
from ..tools.circuit_breaker import get_circuit_breaker, CircuitBreakerConfig, CircuitBreakerException, _NeutralOutcome
|
||||
from ..tools.circuit_breaker import CircuitBreakerException, _NeutralOutcome
|
||||
except ImportError:
|
||||
# Circuit breaker not available - fallback to direct execution
|
||||
logging.debug("Circuit breaker not available, falling back to direct tool execution")
|
||||
@@ -3309,29 +3409,15 @@ class ToolExecutionMixin:
|
||||
|
||||
try:
|
||||
|
||||
# Get or create circuit breaker for this tool.
|
||||
# Scope the key to this Agent instance so one agent's failing tool
|
||||
# cannot trip the breaker for another agent's same-named tool.
|
||||
breaker_name = f"tool_{id(self)}_{function_name}"
|
||||
config = CircuitBreakerConfig(
|
||||
failure_threshold=5, # Open after 5 failures
|
||||
recovery_timeout=60.0, # Wait 60s before trying half-open
|
||||
timeout=30.0, # Tool call timeout
|
||||
graceful_degradation=True # Return error instead of raising exception
|
||||
)
|
||||
breaker = get_circuit_breaker(breaker_name, config)
|
||||
|
||||
# Ensure the registry entry is removed the moment this Agent is
|
||||
# actually garbage-collected, regardless of whether close()/aclose()
|
||||
# was ever called. This closes the CPython id-reuse window where a
|
||||
# new Agent at the same address could inherit a stale OPEN breaker.
|
||||
self._register_breaker_finalizer(breaker_name)
|
||||
# Get or create this agent's circuit breaker for this tool (shared
|
||||
# with the async path so both use the same key and policy).
|
||||
breaker = self._get_tool_circuit_breaker(function_name)
|
||||
|
||||
# Execute tool through circuit breaker with failure detection wrapper
|
||||
def _tool_wrapper():
|
||||
result = self._execute_tool_impl(function_name, arguments)
|
||||
# Caller errors (bad input: ValueError/TypeError/AttributeError,
|
||||
# tagged non-retryable) say nothing about the tool's health: record
|
||||
# Caller errors (bad input: _CALLER_ERROR_TYPES, tagged
|
||||
# non-retryable) say nothing about the tool's health: record
|
||||
# them as neither failure nor success, so they can neither open the
|
||||
# breaker for other callers nor close a half-open one.
|
||||
if isinstance(result, dict) and result.get("error") and \
|
||||
@@ -3367,13 +3453,7 @@ class ToolExecutionMixin:
|
||||
except CircuitBreakerException as e:
|
||||
# Circuit breaker is open - return error dict instead of raising
|
||||
logging.warning(f"Tool '{function_name}' circuit breaker open: {e}")
|
||||
return {
|
||||
"error": f"Tool '{function_name}' circuit breaker open - too many recent failures",
|
||||
"circuit_open": True,
|
||||
"agent_name": getattr(self, "name", None),
|
||||
"session_id": getattr(self, "_session_id", None),
|
||||
"remediation": "Wait for recovery_timeout (60s) or investigate recent tool failures.",
|
||||
}
|
||||
return self._circuit_breaker_open_result(function_name)
|
||||
|
||||
def _find_declared_tool(self, function_name, tools=None):
|
||||
"""Return the tool OBJECT named ``function_name``, or ``None``.
|
||||
@@ -3823,7 +3903,7 @@ class ToolExecutionMixin:
|
||||
# model and cannot collide with a tool's own "retryable" field.
|
||||
return {
|
||||
"error": error_msg,
|
||||
"_praison_retryable": not isinstance(e, (ValueError, TypeError, AttributeError)),
|
||||
"_praison_retryable": not isinstance(e, _CALLER_ERROR_TYPES),
|
||||
}
|
||||
|
||||
# Unresolved: return a corrective, model-readable message so the model can
|
||||
|
||||
@@ -44,6 +44,9 @@ DEFAULT_DANGEROUS_TOOLS: Dict[str, str] = {
|
||||
"acp_edit_file": "high",
|
||||
"acp_delete_file": "high",
|
||||
"execute_query": "high",
|
||||
"github_create_branch": "high",
|
||||
"github_commit_and_push": "high",
|
||||
"github_create_pull_request": "high",
|
||||
"evaluate": "medium",
|
||||
"crawl": "medium",
|
||||
"scrape_page": "medium",
|
||||
@@ -76,10 +79,11 @@ DOOM_LOOP_TARGET = "__doom_loop__"
|
||||
# or set ``PRAISONAI_TOOL_SAFETY=off``; users who want stricter
|
||||
# controls can opt into ``approval="safe"`` or ``"read_only"``.
|
||||
PERMISSION_PRESETS = {
|
||||
# "default" — blocks delete + shell/code exec. Allows read/create/edit.
|
||||
# "default" — blocks delete + shell/code exec + git push/PR. Allows read/create/edit.
|
||||
"default": frozenset({
|
||||
"execute_command", "kill_process", "execute_code", "acp_execute_command",
|
||||
"delete_file", "move_file", "copy_file", "acp_delete_file",
|
||||
"github_create_branch", "github_commit_and_push", "github_create_pull_request",
|
||||
}),
|
||||
# "safe" — blocks all dangerous tools (file writes, shell exec, etc.)
|
||||
"safe": frozenset(DEFAULT_DANGEROUS_TOOLS.keys()),
|
||||
|
||||
@@ -590,8 +590,11 @@ class GuardrailConfig:
|
||||
# With validator function
|
||||
Agent(guardrails=my_validator_fn)
|
||||
|
||||
# With string preset
|
||||
Agent(guardrails="strict") # Uses strict preset
|
||||
# With a string: an LLM-validation prompt (not a preset name)
|
||||
Agent(guardrails="Ensure the response is polite")
|
||||
|
||||
# With a preset (sets max_retries only) plus a validator
|
||||
Agent(guardrails=["strict", {"validator": my_validator_fn}])
|
||||
|
||||
# For tool policy enforcement, use the dedicated `policy` param instead:
|
||||
Agent(policy=PolicyEngine(...)) # policy strings on guardrails are NOT enforced
|
||||
@@ -600,7 +603,7 @@ class GuardrailConfig:
|
||||
Agent(guardrails=GuardrailConfig(
|
||||
validator=my_validator_fn,
|
||||
max_retries=3,
|
||||
on_fail="retry",
|
||||
on_fail="retry", # only "retry" is implemented
|
||||
))
|
||||
|
||||
# With LLM-based validation
|
||||
@@ -909,7 +912,8 @@ class ExecutionConfig:
|
||||
# ``rate_limiter`` explicitly for burst/TPM control (it takes precedence).
|
||||
max_rpm: Optional[int] = None
|
||||
|
||||
# Time limits
|
||||
# Time limits: per-turn wall-clock seconds for the loop guard.
|
||||
# None = loop-guard default (600s); 0 = no time limit.
|
||||
max_execution_time: Optional[int] = None
|
||||
|
||||
# Retry settings
|
||||
|
||||
@@ -275,9 +275,11 @@ REFLECTION_PRESETS: Dict[str, Dict[str, Any]] = {
|
||||
# =============================================================================
|
||||
|
||||
GUARDRAIL_PRESETS: Dict[str, Dict[str, Any]] = {
|
||||
"strict": {"max_retries": 5, "on_fail": "raise"},
|
||||
"permissive": {"max_retries": 1, "on_fail": "skip"},
|
||||
"safety": {"max_retries": 3, "on_fail": "retry"},
|
||||
# Presets only set retry counts; supply a validator alongside them,
|
||||
# e.g. guardrails=["strict", {"validator": fn}].
|
||||
"strict": {"max_retries": 5},
|
||||
"permissive": {"max_retries": 1},
|
||||
"safety": {"max_retries": 3},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -107,7 +107,8 @@ class LoopGuardConfig:
|
||||
mutating_halt_threshold: int = 7
|
||||
|
||||
# Time-based limits
|
||||
max_time_per_turn: float = 120.0 # Max seconds per chat turn
|
||||
# Max seconds per chat turn; ExecutionConfig.max_execution_time overrides (0 = no limit there).
|
||||
max_time_per_turn: float = 600.0
|
||||
|
||||
# No-progress detection
|
||||
no_progress_warn: int = 4
|
||||
|
||||
@@ -11,6 +11,7 @@ import time
|
||||
import asyncio
|
||||
from praisonaiagents._logging import get_logger
|
||||
import subprocess
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
@@ -78,25 +79,14 @@ class HookRunner:
|
||||
# misbehaving hook can only exhaust hook capacity, never the loop's
|
||||
# shared default executor that guardrails/tools/locks also rely on.
|
||||
self._hook_pool: Optional[ThreadPoolExecutor] = None
|
||||
# hook.id -> in-flight Future for a sync hook. The slot is reserved
|
||||
# before submission and only cleared when the worker finishes, so a
|
||||
# still-running (timed-out or cancelled) hook caps new launches at one
|
||||
# thread per hook. Guarded by ``_stuck_lock`` so concurrent invocations
|
||||
# of the same hook cannot both reserve the slot and race past the cap.
|
||||
# hook.id -> worker Future of a sync hook that ran past its timeout.
|
||||
# It is cleared only when that worker finishes, so a stuck hook caps
|
||||
# new launches at one leaked thread per hook. Healthy concurrent calls
|
||||
# (e.g. parallel tool calls) are not tracked and never denied. A
|
||||
# threading lock, because execute_sync() drives each call on its own
|
||||
# event loop in the caller's thread.
|
||||
self._stuck: Dict[str, Any] = {}
|
||||
self._stuck_lock_obj: Optional[asyncio.Lock] = None
|
||||
|
||||
@property
|
||||
def _stuck_lock(self) -> asyncio.Lock:
|
||||
"""Lazily create the in-flight-slot lock on first use.
|
||||
|
||||
Created lazily (not in ``__init__``) because a ``HookRunner`` may be
|
||||
constructed before any event loop exists; the lock is only needed once
|
||||
we are already inside ``async def execute`` / ``_execute_function_hook``.
|
||||
"""
|
||||
if self._stuck_lock_obj is None:
|
||||
self._stuck_lock_obj = asyncio.Lock()
|
||||
return self._stuck_lock_obj
|
||||
self._stuck_lock = threading.Lock()
|
||||
|
||||
def _release_stuck(self, hook_id: str, fut: Any) -> None:
|
||||
"""Clear the in-flight slot once the worker Future completes.
|
||||
@@ -104,8 +94,9 @@ class HookRunner:
|
||||
Only removes the entry if it still points at *this* Future, so a newer
|
||||
reserved invocation is never clobbered by an older worker draining.
|
||||
"""
|
||||
if self._stuck.get(hook_id) is fut:
|
||||
self._stuck.pop(hook_id, None)
|
||||
with self._stuck_lock:
|
||||
if self._stuck.get(hook_id) is fut:
|
||||
self._stuck.pop(hook_id, None)
|
||||
|
||||
def _get_hook_pool(self) -> ThreadPoolExecutor:
|
||||
"""Lazily create the dedicated sync-hook executor."""
|
||||
@@ -439,14 +430,9 @@ class HookRunner:
|
||||
timeout=timeout
|
||||
)
|
||||
else:
|
||||
# Reserve the single in-flight slot for this hook *under a lock*
|
||||
# and *before* submitting, so two concurrent invocations can't
|
||||
# both see an empty ``_stuck`` and both launch a worker (which
|
||||
# would exceed the intended one-worker-per-hook cap and double
|
||||
# the side effects). The slot is only cleared once the worker
|
||||
# finishes; a still-running previous worker fails this one
|
||||
# closed immediately instead of leaking another thread.
|
||||
async with self._stuck_lock:
|
||||
# A previous invocation that timed out and is still running
|
||||
# fails this one closed instead of leaking another thread.
|
||||
with self._stuck_lock:
|
||||
prev = self._stuck.get(hook.id)
|
||||
if prev is not None and not prev.done():
|
||||
duration = (time.time() - start_time) * 1000
|
||||
@@ -463,29 +449,31 @@ class HookRunner:
|
||||
),
|
||||
)
|
||||
|
||||
# Execute sync function in the dedicated hook pool with
|
||||
# timeout. ``copy_context_to_callable`` preserves
|
||||
# trace/session contextvars across the thread, matching
|
||||
# other executor sites.
|
||||
from ..trace.context_events import copy_context_to_callable
|
||||
loop = asyncio.get_event_loop()
|
||||
func = hook.func
|
||||
data = input_data
|
||||
fut = loop.run_in_executor(
|
||||
self._get_hook_pool(),
|
||||
copy_context_to_callable(lambda: func(data)),
|
||||
)
|
||||
# Reserve the slot now; a completion callback clears it only
|
||||
# when the worker actually finishes, so a timeout *or* a
|
||||
# cancellation leaves the slot held until the thread drains.
|
||||
self._stuck[hook.id] = fut
|
||||
fut.add_done_callback(
|
||||
lambda f, hid=hook.id: self._release_stuck(hid, f)
|
||||
)
|
||||
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.shield(fut), timeout=timeout
|
||||
# Execute sync function in the dedicated hook pool with
|
||||
# timeout. ``copy_context_to_callable`` preserves
|
||||
# trace/session contextvars across the thread, matching
|
||||
# other executor sites. A concurrent Future (not a loop-bound
|
||||
# one) so the stuck-slot release works after the loop closes.
|
||||
from ..trace.context_events import copy_context_to_callable
|
||||
func = hook.func
|
||||
data = input_data
|
||||
fut = self._get_hook_pool().submit(
|
||||
copy_context_to_callable(lambda: func(data))
|
||||
)
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.shield(asyncio.wrap_future(fut)), timeout=timeout
|
||||
)
|
||||
except (asyncio.TimeoutError, asyncio.CancelledError):
|
||||
# The worker cannot be cancelled; hold the slot until it
|
||||
# actually finishes.
|
||||
if not fut.done():
|
||||
with self._stuck_lock:
|
||||
self._stuck[hook.id] = fut
|
||||
fut.add_done_callback(
|
||||
lambda f, hid=hook.id: self._release_stuck(hid, f)
|
||||
)
|
||||
raise
|
||||
|
||||
duration = (time.time() - start_time) * 1000
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@ from pydantic import BaseModel
|
||||
import time
|
||||
import json
|
||||
import xml.etree.ElementTree as ET
|
||||
from ..errors import AgentErrorKind, FailoverDecision, IdleTimeoutBreaker, ToolExecutionError
|
||||
from ..errors import AgentErrorKind, BudgetExceededError, FailoverDecision, IdleTimeoutBreaker, ToolExecutionError
|
||||
from ..model_harness.guard import check_model_request
|
||||
# Gap 2: Tool call execution imports
|
||||
from ..tools.call_executor import ToolCall, create_tool_call_executor
|
||||
@@ -344,6 +344,11 @@ except ImportError:
|
||||
# API calls per request. This is a DRY/maintenance concern, not a billing issue.
|
||||
from ..tools.hosted import is_hosted_tool # noqa: E402
|
||||
|
||||
# ``(llm, callback)`` bound by the owning Agent for one custom-LLM turn. Every
|
||||
# tracked call made by that ``llm`` reports its response to ``callback`` so the
|
||||
# Agent's cost/token counters and ``max_budget`` cover the whole tool loop.
|
||||
_usage_sink = contextvars.ContextVar("praisonai_llm_usage_sink", default=None)
|
||||
|
||||
|
||||
class LLMResponseError(Exception):
|
||||
"""Raised when the LLM tool-calling loop fails and cannot produce a response.
|
||||
@@ -3988,7 +3993,9 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
# Mark that fallback completed successfully
|
||||
fallback_completed = True
|
||||
streaming_success = False
|
||||
|
||||
|
||||
except BudgetExceededError:
|
||||
raise
|
||||
except Exception as fallback_error:
|
||||
# If non-streaming also fails, create a graceful fallback with partial streaming data
|
||||
logging.warning(f"Non-streaming fallback also failed: {fallback_error}")
|
||||
@@ -4514,7 +4521,7 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
final_response_text = response_text.strip() if response_text else ""
|
||||
break
|
||||
|
||||
except (LLMResponseError, ToolExecutionError):
|
||||
except (LLMResponseError, ToolExecutionError, BudgetExceededError):
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(f"Error in LLM iteration {iteration_count}: {e}")
|
||||
@@ -4861,6 +4868,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
interaction_displayed = True
|
||||
return _prepare_return_value(response_text)
|
||||
continue
|
||||
except BudgetExceededError:
|
||||
raise
|
||||
except Exception as e:
|
||||
_get_display_functions()['display_error'](f"Error in LLM response: {str(e)}")
|
||||
return None
|
||||
@@ -6856,6 +6865,16 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
def _track_token_usage(self, response: Any, model: str) -> Optional[TokenMetrics]:
|
||||
"""Extract and track token usage from LLM response."""
|
||||
try:
|
||||
return self._collect_token_metrics(response, model)
|
||||
finally:
|
||||
# Outside the collector's error swallowing so the Agent's budget
|
||||
# check can raise BudgetExceededError and stop the tool loop.
|
||||
sink = _usage_sink.get()
|
||||
if sink is not None and sink[0] is self:
|
||||
sink[1](response)
|
||||
|
||||
def _collect_token_metrics(self, response: Any, model: str) -> Optional[TokenMetrics]:
|
||||
if not TokenMetrics or not get_token_collector:
|
||||
return None
|
||||
|
||||
|
||||
@@ -127,43 +127,42 @@ class MCPToolRunner(threading.Thread):
|
||||
# Signal that initialization is complete
|
||||
self.initialized.set()
|
||||
|
||||
# Process requests
|
||||
while True:
|
||||
# Process requests. Each request runs as its own task so
|
||||
# one slow call never serializes the others behind it (the
|
||||
# session multiplexes requests by id), and the queue is read
|
||||
# with a short blocking get instead of a busy 10ms poll. The
|
||||
# bounded get keeps the reader thread from outliving the loop.
|
||||
loop = asyncio.get_running_loop()
|
||||
tasks = set()
|
||||
|
||||
def _next_item():
|
||||
try:
|
||||
# Check for new requests
|
||||
try:
|
||||
item = self.queue.get(block=False)
|
||||
if item is None: # Shutdown signal
|
||||
break
|
||||
|
||||
response_queue, kind, name, arguments, cancelled = item
|
||||
if cancelled.is_set():
|
||||
# Caller already gave up (timed out); never
|
||||
# execute a call it was told did not happen.
|
||||
continue
|
||||
try:
|
||||
if kind == "resource":
|
||||
result = await asyncio.wait_for(session.read_resource(name), timeout=self.timeout)
|
||||
elif kind == "prompt":
|
||||
result = await asyncio.wait_for(session.get_prompt(name, arguments or None), timeout=self.timeout)
|
||||
else:
|
||||
result = await asyncio.wait_for(session.call_tool(name, arguments), timeout=self.timeout)
|
||||
response_queue.put((True, result))
|
||||
except asyncio.TimeoutError:
|
||||
response_queue.put((False, f"MCP {kind} call timed out after {self.timeout} seconds (server side)"))
|
||||
except Exception as e:
|
||||
if _is_transport_closed(e):
|
||||
self._dead = f"MCP server connection lost: {str(e) or type(e).__name__}"
|
||||
response_queue.put((False, self._dead))
|
||||
break # leave stdio_client() so the child is reaped
|
||||
response_queue.put((False, str(e) or type(e).__name__))
|
||||
except queue.Empty:
|
||||
pass
|
||||
|
||||
# Give other tasks a chance to run
|
||||
await asyncio.sleep(0.01)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
return self.queue.get(timeout=0.25)
|
||||
except queue.Empty:
|
||||
return queue.Empty
|
||||
|
||||
try:
|
||||
while True:
|
||||
item = await loop.run_in_executor(None, _next_item)
|
||||
if item is queue.Empty:
|
||||
if self._dead:
|
||||
break # leave stdio_client() so the child is reaped
|
||||
continue
|
||||
if item is None: # Shutdown signal
|
||||
break
|
||||
if self._dead:
|
||||
self.queue.put(item) # answered by _fail_pending()
|
||||
break
|
||||
task = asyncio.create_task(self._handle_request(session, item))
|
||||
tasks.add(task)
|
||||
task.add_done_callback(tasks.discard)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
finally:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
except asyncio.CancelledError:
|
||||
# stop() cancelled a startup that never completed; leaving the
|
||||
# contexts above has already terminated the stdio child.
|
||||
@@ -197,6 +196,35 @@ class MCPToolRunner(threading.Thread):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _handle_request(self, session, item):
|
||||
"""Run one queued request on the session and post its result."""
|
||||
response_queue, kind, name, arguments, cancelled = item
|
||||
if cancelled.is_set():
|
||||
# Caller already gave up (timed out); never execute a call it was
|
||||
# told did not happen.
|
||||
return
|
||||
try:
|
||||
if kind == "resource":
|
||||
coro = session.read_resource(name)
|
||||
elif kind == "prompt":
|
||||
coro = session.get_prompt(name, arguments or None)
|
||||
else:
|
||||
coro = session.call_tool(name, arguments)
|
||||
result = await asyncio.wait_for(coro, timeout=self.timeout)
|
||||
response_queue.put((True, result))
|
||||
except asyncio.TimeoutError:
|
||||
response_queue.put((False, f"MCP {kind} call timed out after {self.timeout} seconds (server side)"))
|
||||
except asyncio.CancelledError:
|
||||
# Report the real cause when a sibling call saw the transport die.
|
||||
response_queue.put((False, self._dead or f"MCP {kind} call cancelled: runner shut down"))
|
||||
raise
|
||||
except Exception as e:
|
||||
if _is_transport_closed(e):
|
||||
self._dead = f"MCP server connection lost: {str(e) or type(e).__name__}"
|
||||
response_queue.put((False, self._dead))
|
||||
return
|
||||
response_queue.put((False, str(e) or type(e).__name__))
|
||||
|
||||
async def _discover_resources_and_prompts(self, session):
|
||||
"""Enumerate server resources, resource templates and prompts.
|
||||
|
||||
@@ -391,8 +419,8 @@ class MCPToolRunner(threading.Thread):
|
||||
Queues the shutdown sentinel and joins the thread so the request loop
|
||||
exits and its stdio subprocess is terminated instead of leaking a
|
||||
background thread (and its child) across the process lifetime. The
|
||||
request loop polls the queue between async operations, so once any
|
||||
in-flight call unwinds the sentinel is consumed promptly.
|
||||
request loop reads the sentinel promptly and cancels any in-flight
|
||||
calls, which are answered with an error instead of left waiting.
|
||||
|
||||
Args:
|
||||
timeout: Seconds to wait for the thread to exit. Defaults to the
|
||||
|
||||
@@ -199,6 +199,17 @@ class RAG:
|
||||
|
||||
return results
|
||||
|
||||
def _reranked(self, rerank: bool) -> bool:
|
||||
"""Whether a reranker actually ran: the local one, or a backend
|
||||
reranker configured on the knowledge (e.g. mem0 with rerank_model)."""
|
||||
if not rerank:
|
||||
return False
|
||||
if self.reranker:
|
||||
return True
|
||||
cfg = getattr(self.knowledge, "config", None)
|
||||
backend = cfg.get("reranker") if isinstance(cfg, dict) else None
|
||||
return bool(isinstance(backend, dict) and backend.get("default_rerank") and backend.get("provider"))
|
||||
|
||||
def _rerank(
|
||||
self,
|
||||
query: str,
|
||||
@@ -395,7 +406,7 @@ class RAG:
|
||||
"num_results": len(results),
|
||||
"num_citations": len(citations),
|
||||
"retrieval_strategy": self.config.retrieval_strategy.value,
|
||||
"reranked": eff_rerank,
|
||||
"reranked": self._reranked(eff_rerank),
|
||||
"top_k": eff_top_k,
|
||||
},
|
||||
)
|
||||
@@ -495,7 +506,7 @@ class RAG:
|
||||
"num_results": len(results),
|
||||
"num_citations": len(citations),
|
||||
"retrieval_strategy": self.config.retrieval_strategy.value,
|
||||
"reranked": self.config.rerank,
|
||||
"reranked": self._reranked(self.config.rerank),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@@ -332,7 +332,10 @@ class SimpleReranker:
|
||||
|
||||
scored = []
|
||||
for result in results:
|
||||
text = result.get('text', '').lower()
|
||||
# Knowledge adapters may return SearchResultItem objects
|
||||
if not isinstance(result, dict) and hasattr(result, 'to_dict'):
|
||||
result = result.to_dict()
|
||||
text = (result.get('text') or result.get('memory') or '').lower()
|
||||
text_words = set(text.split())
|
||||
|
||||
# Calculate overlap score
|
||||
|
||||
@@ -3,7 +3,7 @@ import re
|
||||
import subprocess
|
||||
import logging
|
||||
import secrets
|
||||
from typing import Optional
|
||||
from typing import List, Optional
|
||||
from .decorator import tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -11,22 +11,22 @@ logger = logging.getLogger(__name__)
|
||||
AGENT_BRANCH_PREFIX = "praisonai/"
|
||||
|
||||
|
||||
def _git(*args: str) -> subprocess.CompletedProcess:
|
||||
"""Run a git command capturing output as text."""
|
||||
return subprocess.run(["git", *args], check=True, capture_output=True, text=True)
|
||||
def _git(*args: str, cwd: Optional[str] = None) -> subprocess.CompletedProcess:
|
||||
"""Run a git command (in ``cwd`` when given) capturing output as text."""
|
||||
return subprocess.run(["git", *args], check=True, capture_output=True, text=True, cwd=cwd)
|
||||
|
||||
|
||||
def _default_branch() -> str:
|
||||
def _default_branch(cwd: Optional[str] = None) -> str:
|
||||
"""Detect the remote default branch, falling back to main/master."""
|
||||
try:
|
||||
ref = _git("symbolic-ref", "refs/remotes/origin/HEAD").stdout.strip()
|
||||
ref = _git("symbolic-ref", "refs/remotes/origin/HEAD", cwd=cwd).stdout.strip()
|
||||
# e.g. refs/remotes/origin/main -> main
|
||||
return ref.rsplit("/", 1)[-1]
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
for candidate in ("main", "master"):
|
||||
try:
|
||||
_git("show-ref", "--verify", "--quiet", f"refs/remotes/origin/{candidate}")
|
||||
_git("show-ref", "--verify", "--quiet", f"refs/remotes/origin/{candidate}", cwd=cwd)
|
||||
return candidate
|
||||
except subprocess.CalledProcessError:
|
||||
continue
|
||||
@@ -45,7 +45,7 @@ def _auto_branch_name(commit_message: str) -> str:
|
||||
return f"{AGENT_BRANCH_PREFIX}{_slugify(commit_message)}-{secrets.token_hex(3)}"
|
||||
|
||||
|
||||
def _divergence_refusal(branch: str) -> Optional[str]:
|
||||
def _divergence_refusal(branch: str, cwd: Optional[str] = None) -> Optional[str]:
|
||||
"""Return a refusal reason if pushing ``branch`` would require a force.
|
||||
|
||||
A force push is required when the remote branch exists but is not an
|
||||
@@ -54,13 +54,13 @@ def _divergence_refusal(branch: str) -> Optional[str]:
|
||||
agent-prefixed ones.
|
||||
"""
|
||||
try:
|
||||
_git("show-ref", "--verify", "--quiet", f"refs/remotes/origin/{branch}")
|
||||
_git("show-ref", "--verify", "--quiet", f"refs/remotes/origin/{branch}", cwd=cwd)
|
||||
remote_exists = True
|
||||
except subprocess.CalledProcessError:
|
||||
remote_exists = False
|
||||
if remote_exists:
|
||||
try:
|
||||
_git("merge-base", "--is-ancestor", f"origin/{branch}", "HEAD")
|
||||
_git("merge-base", "--is-ancestor", f"origin/{branch}", "HEAD", cwd=cwd)
|
||||
except subprocess.CalledProcessError:
|
||||
return (
|
||||
f"refusing to push to '{branch}': the remote branch has "
|
||||
@@ -69,7 +69,7 @@ def _divergence_refusal(branch: str) -> Optional[str]:
|
||||
return None
|
||||
|
||||
|
||||
def _branch_push_safety(branch: str) -> Optional[str]:
|
||||
def _branch_push_safety(branch: str, cwd: Optional[str] = None) -> Optional[str]:
|
||||
"""Return a refusal reason if pushing ``branch`` is unsafe, else None.
|
||||
|
||||
Rules (see github_commit_and_push docstring):
|
||||
@@ -82,14 +82,14 @@ def _branch_push_safety(branch: str) -> Optional[str]:
|
||||
(i.e. would require a force push) are refused.
|
||||
"""
|
||||
# Divergence (force-push) is refused for every branch, agent-prefixed or not.
|
||||
diverged = _divergence_refusal(branch)
|
||||
diverged = _divergence_refusal(branch, cwd=cwd)
|
||||
if diverged:
|
||||
return diverged
|
||||
|
||||
if branch.startswith(AGENT_BRANCH_PREFIX):
|
||||
return None
|
||||
|
||||
default = _default_branch()
|
||||
default = _default_branch(cwd=cwd)
|
||||
if branch == default:
|
||||
return (
|
||||
f"refusing to push to '{branch}': it is the repository "
|
||||
@@ -99,16 +99,16 @@ def _branch_push_safety(branch: str) -> Optional[str]:
|
||||
# Foreign-author check: commits on this branch since the merge-base with
|
||||
# the default branch must all be authored by the configured committer.
|
||||
try:
|
||||
committer = _git("config", "user.email").stdout.strip()
|
||||
committer = _git("config", "user.email", cwd=cwd).stdout.strip()
|
||||
except subprocess.CalledProcessError:
|
||||
committer = ""
|
||||
if committer:
|
||||
try:
|
||||
merge_base = _git(
|
||||
"merge-base", f"origin/{default}", "HEAD"
|
||||
"merge-base", f"origin/{default}", "HEAD", cwd=cwd
|
||||
).stdout.strip()
|
||||
authors = _git(
|
||||
"log", "--format=%ae", f"{merge_base}..HEAD"
|
||||
"log", "--format=%ae", f"{merge_base}..HEAD", cwd=cwd
|
||||
).stdout.split()
|
||||
foreign = {a for a in authors if a and a != committer}
|
||||
if foreign:
|
||||
@@ -123,25 +123,27 @@ def _branch_push_safety(branch: str) -> Optional[str]:
|
||||
return None
|
||||
|
||||
@tool
|
||||
def github_create_branch(branch_name: str) -> str:
|
||||
def github_create_branch(branch_name: str, repo_path: Optional[str] = None) -> str:
|
||||
"""Create and checkout a new git branch.
|
||||
|
||||
|
||||
Args:
|
||||
branch_name: The name of the branch to create and checkout.
|
||||
repo_path: Repository directory to run git in. Defaults to the
|
||||
current working directory.
|
||||
"""
|
||||
try:
|
||||
# Check if we are in a git repository
|
||||
subprocess.run(["git", "rev-parse", "--is-inside-work-tree"], check=True, capture_output=True)
|
||||
subprocess.run(["git", "rev-parse", "--is-inside-work-tree"], check=True, capture_output=True, cwd=repo_path)
|
||||
# Validate branch name to prevent misinterpretation as git options or invalid refs
|
||||
try:
|
||||
subprocess.run(
|
||||
["git", "check-ref-format", f"refs/heads/{branch_name}"],
|
||||
check=True, capture_output=True, text=True
|
||||
check=True, capture_output=True, text=True, cwd=repo_path
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
return f"Error: invalid branch name '{branch_name}': {e.stderr.strip() if e.stderr else 'branch name is not a valid git ref'}"
|
||||
subprocess.run(["git", "checkout", "-B", branch_name], check=True, capture_output=True, text=True)
|
||||
subprocess.run(["git", "checkout", "-B", branch_name], check=True, capture_output=True, text=True, cwd=repo_path)
|
||||
logger.debug(f"Branch '{branch_name}' checked out successfully.")
|
||||
logger.debug(f"Successfully checked out branch '{branch_name}'")
|
||||
return f"Successfully created and checked out branch '{branch_name}'"
|
||||
@@ -154,8 +156,15 @@ def github_commit_and_push(
|
||||
commit_message: str,
|
||||
branch: Optional[str] = None,
|
||||
allow_unsafe_branch: bool = False,
|
||||
repo_path: Optional[str] = None,
|
||||
paths: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
"""Stage all changes, commit them, and push to the remote repository safely.
|
||||
"""Stage tracked changes, commit them, and push to the remote repository safely.
|
||||
|
||||
Only changes to already-tracked files under the working directory are
|
||||
staged (``git add -u -- .``); new files are committed only when listed in
|
||||
``paths``. Staging happens after the safety checks below, so a refusal
|
||||
leaves the index untouched.
|
||||
|
||||
Pushes are safe by construction. Before pushing, these rules are enforced:
|
||||
|
||||
@@ -176,25 +185,27 @@ def github_commit_and_push(
|
||||
branch: Target branch to push to. Defaults to the current branch.
|
||||
allow_unsafe_branch: Opt-in override for the default/foreign-author
|
||||
refusals. Divergence (force-push) is never allowed.
|
||||
repo_path: Repository directory to run git in. Defaults to the
|
||||
current working directory.
|
||||
paths: Extra (e.g. untracked) files to stage alongside tracked changes.
|
||||
"""
|
||||
try:
|
||||
allow_unsafe = allow_unsafe_branch or os.environ.get(
|
||||
"PRAISONAI_GIT_ALLOW_UNSAFE_BRANCH", ""
|
||||
).lower() in ("1", "true", "yes")
|
||||
|
||||
# Add all changes
|
||||
subprocess.run(["git", "add", "."], check=True, capture_output=True)
|
||||
|
||||
# Check if there are changes to commit
|
||||
status = subprocess.run(["git", "status", "--porcelain"], check=True, capture_output=True, text=True)
|
||||
if not status.stdout.strip():
|
||||
# Check if there are changes to commit (tracked changes or explicit paths)
|
||||
status = _git(
|
||||
"status", "--porcelain", "--untracked-files=no", "--", ".", cwd=repo_path
|
||||
)
|
||||
if not status.stdout.strip() and not paths:
|
||||
return "No changes to commit."
|
||||
|
||||
# Resolve target branch (defaults to current)
|
||||
current_branch = _git("rev-parse", "--abbrev-ref", "HEAD").stdout.strip()
|
||||
current_branch = _git("rev-parse", "--abbrev-ref", "HEAD", cwd=repo_path).stdout.strip()
|
||||
target_branch = branch or current_branch
|
||||
|
||||
default = _default_branch()
|
||||
default = _default_branch(cwd=repo_path)
|
||||
auto_created = False
|
||||
|
||||
# Rule 2: default branch is never pushed directly (unless overridden).
|
||||
@@ -207,7 +218,7 @@ def github_commit_and_push(
|
||||
# retry can then succeed). Adding the new commit does not change the
|
||||
# divergence/foreign-author verdicts, which are about existing history.
|
||||
if not allow_unsafe:
|
||||
reason = _branch_push_safety(target_branch)
|
||||
reason = _branch_push_safety(target_branch, cwd=repo_path)
|
||||
if reason:
|
||||
return f"Error: {reason}"
|
||||
else:
|
||||
@@ -217,14 +228,14 @@ def github_commit_and_push(
|
||||
"checks for '%s'", target_branch
|
||||
)
|
||||
# Even when overriding, divergence (force) is never permitted.
|
||||
diverged = _divergence_refusal(target_branch)
|
||||
diverged = _divergence_refusal(target_branch, cwd=repo_path)
|
||||
if diverged:
|
||||
return f"Error: {diverged}"
|
||||
|
||||
# Move to the resolved branch only after passing safety checks.
|
||||
if auto_created:
|
||||
# Fresh agent branch created from current HEAD.
|
||||
_git("checkout", "-B", target_branch)
|
||||
_git("checkout", "-B", target_branch, cwd=repo_path)
|
||||
elif target_branch != current_branch:
|
||||
# Switch to the requested branch without discarding its commits:
|
||||
# ``-B`` would reset an existing branch to HEAD, so use a plain
|
||||
@@ -233,17 +244,24 @@ def github_commit_and_push(
|
||||
try:
|
||||
_git(
|
||||
"show-ref", "--verify", "--quiet",
|
||||
f"refs/heads/{target_branch}",
|
||||
f"refs/heads/{target_branch}", cwd=repo_path,
|
||||
)
|
||||
_git("checkout", target_branch)
|
||||
_git("checkout", target_branch, cwd=repo_path)
|
||||
except subprocess.CalledProcessError:
|
||||
_git("checkout", "-b", target_branch)
|
||||
_git("checkout", "-b", target_branch, cwd=repo_path)
|
||||
|
||||
# Stage tracked changes under cwd plus explicitly listed paths (never
|
||||
# ``git add .``). The ``.`` pathspec keeps ``-u`` from reaching sibling
|
||||
# directories when run from a repository subdirectory.
|
||||
_git("add", "-u", "--", ".", cwd=repo_path)
|
||||
if paths:
|
||||
_git("add", "--", *paths, cwd=repo_path)
|
||||
|
||||
# Commit changes on the resolved branch
|
||||
subprocess.run(["git", "commit", "-m", commit_message], check=True, capture_output=True, text=True)
|
||||
_git("commit", "-m", commit_message, cwd=repo_path)
|
||||
|
||||
# Push to remote (never with --force)
|
||||
subprocess.run(["git", "push", "-u", "origin", target_branch], check=True, capture_output=True, text=True)
|
||||
_git("push", "-u", "origin", target_branch, cwd=repo_path)
|
||||
if auto_created:
|
||||
return (
|
||||
f"Successfully committed and pushed changes to new agent branch "
|
||||
@@ -255,7 +273,13 @@ def github_commit_and_push(
|
||||
return f"Error committing and pushing: {e.stderr}"
|
||||
|
||||
@tool
|
||||
def github_create_pull_request(title: str, body: str, head_branch: str, base_branch: str = "main") -> str:
|
||||
def github_create_pull_request(
|
||||
title: str,
|
||||
body: str,
|
||||
head_branch: str,
|
||||
base_branch: str = "main",
|
||||
repo_path: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Create a Pull Request on GitHub using the gh CLI.
|
||||
|
||||
Args:
|
||||
@@ -267,6 +291,8 @@ def github_create_pull_request(title: str, body: str, head_branch: str, base_bra
|
||||
This is typically the feature branch you've been working on.
|
||||
base_branch: The target branch to merge your changes into. Usually the main
|
||||
development branch like "main", "master", or "develop". Defaults to "main".
|
||||
repo_path: Repository directory to run gh in. Defaults to the current
|
||||
working directory.
|
||||
|
||||
Returns:
|
||||
str: Success message with PR URL if created successfully, or error message if failed.
|
||||
@@ -282,7 +308,7 @@ def github_create_pull_request(title: str, body: str, head_branch: str, base_bra
|
||||
"""
|
||||
try:
|
||||
# Verify gh CLI is installed and authenticated
|
||||
subprocess.run(["gh", "auth", "status"], check=True, capture_output=True)
|
||||
subprocess.run(["gh", "auth", "status"], check=True, capture_output=True, cwd=repo_path)
|
||||
|
||||
# Create PR
|
||||
cmd = [
|
||||
@@ -292,7 +318,7 @@ def github_create_pull_request(title: str, body: str, head_branch: str, base_bra
|
||||
"--head", head_branch,
|
||||
"--base", base_branch
|
||||
]
|
||||
result = subprocess.run(cmd, check=True, capture_output=True, text=True)
|
||||
result = subprocess.run(cmd, check=True, capture_output=True, text=True, cwd=repo_path)
|
||||
return f"Successfully created Pull Request:\n{result.stdout.strip()}"
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.error(f"Failed to create PR: {e.stderr}")
|
||||
|
||||
@@ -85,6 +85,16 @@ def get_context_emitter() -> "ContextTraceEmitter":
|
||||
return _default_disabled_emitter
|
||||
|
||||
|
||||
def _agent_key(name: Optional[str]) -> str:
|
||||
"""The id an agent has in traces and graphs: its name, or "Agent" when unnamed.
|
||||
|
||||
``Agent(instructions=...)`` has ``name=None``; without this its events had no
|
||||
node id and its tools were recorded as ``tool:None/<tool>``, so nothing
|
||||
matched the ``agent:Agent`` node in its graph.
|
||||
"""
|
||||
return name or "Agent"
|
||||
|
||||
|
||||
def set_context_emitter(emitter: "ContextTraceEmitter") -> contextvars.Token:
|
||||
"""
|
||||
Set the context trace emitter for the current async context.
|
||||
@@ -786,7 +796,7 @@ class ContextTraceEmitter:
|
||||
session_id=self._session_id,
|
||||
agent_name=agent_name,
|
||||
data=metadata or {},
|
||||
node_id=f"agent:{agent_name}" if agent_name else None,
|
||||
node_id=f"agent:{_agent_key(agent_name)}",
|
||||
))
|
||||
|
||||
def agent_end(
|
||||
@@ -803,7 +813,7 @@ class ContextTraceEmitter:
|
||||
session_id=self._session_id,
|
||||
agent_name=agent_name,
|
||||
data=metadata or {},
|
||||
node_id=f"agent:{agent_name}" if agent_name else None,
|
||||
node_id=f"agent:{_agent_key(agent_name)}",
|
||||
))
|
||||
|
||||
def message_added(
|
||||
@@ -853,7 +863,7 @@ class ContextTraceEmitter:
|
||||
"tool_name": tool_name,
|
||||
"tool_args": self._redact_dict(tool_args) if tool_args else {},
|
||||
},
|
||||
node_id=f"tool:{agent_name}/{tool_name}",
|
||||
node_id=f"tool:{_agent_key(agent_name)}/{tool_name}",
|
||||
))
|
||||
|
||||
def tool_call_end(
|
||||
@@ -901,7 +911,7 @@ class ContextTraceEmitter:
|
||||
"duration_ms": duration_ms,
|
||||
"error": error,
|
||||
},
|
||||
node_id=f"tool:{agent_name}/{tool_name}",
|
||||
node_id=f"tool:{_agent_key(agent_name)}/{tool_name}",
|
||||
))
|
||||
|
||||
def _smart_truncate_result(self, result: str, tool_name: str) -> str:
|
||||
@@ -1030,7 +1040,7 @@ class ContextTraceEmitter:
|
||||
tokens_used=tokens_used,
|
||||
tokens_budget=tokens_budget,
|
||||
data=data,
|
||||
node_id=f"agent:{agent_name}" if agent_name else None,
|
||||
node_id=f"agent:{_agent_key(agent_name)}",
|
||||
)
|
||||
|
||||
def llm_response(
|
||||
@@ -1114,7 +1124,7 @@ class ContextTraceEmitter:
|
||||
completion_tokens=actual_completion_tokens,
|
||||
cost_usd=cost_usd,
|
||||
data=data,
|
||||
node_id=f"agent:{agent_name}" if agent_name else None,
|
||||
node_id=f"agent:{_agent_key(agent_name)}",
|
||||
)
|
||||
|
||||
def context_snapshot(
|
||||
|
||||
@@ -87,7 +87,8 @@ def _tool_name(tool: Any) -> str:
|
||||
|
||||
|
||||
def _is_agent(obj: Any) -> bool:
|
||||
return hasattr(obj, "chat") and isinstance(getattr(obj, "name", None), str)
|
||||
# An unnamed Agent (``Agent(instructions=...)``) has ``name=None``; it is still an agent.
|
||||
return hasattr(obj, "chat") and hasattr(obj, "name")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -204,8 +205,11 @@ class _Graph:
|
||||
|
||||
def _add_agent(g: _Graph, agent: Any, parent: Optional[str] = None, _seen: Optional[set] = None) -> str:
|
||||
"""Add an agent node with its tools and hand-off targets; return its id."""
|
||||
name = getattr(agent, "name", None) or type(agent).__name__
|
||||
agent_id = f"agent:{name}"
|
||||
from ..trace.context_events import _agent_key
|
||||
|
||||
name = getattr(agent, "name", None) or type(agent).__name__ # label
|
||||
key = _agent_key(getattr(agent, "name", None)) # the id the trace uses
|
||||
agent_id = f"agent:{key}"
|
||||
seen = _seen if _seen is not None else set()
|
||||
if agent_id in seen:
|
||||
return agent_id
|
||||
@@ -227,7 +231,7 @@ def _add_agent(g: _Graph, agent: Any, parent: Optional[str] = None, _seen: Optio
|
||||
tool_name = _tool_name(tool)
|
||||
if tool_name in handoff_tools:
|
||||
continue # drawn as the hand-off edge, not as a tool
|
||||
tool_id = f"tool:{name}/{tool_name}"
|
||||
tool_id = f"tool:{key}/{tool_name}"
|
||||
g.node(tool_id, "tool", tool_name, parent=parent)
|
||||
g.edge(agent_id, tool_id, "uses")
|
||||
for target in targets:
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
"""A configured completion promise must be the only success signal (#6027)."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.model_harness import ScriptedModel
|
||||
|
||||
NOT_DONE = (
|
||||
"I listed the files and started on the parser, but work remains: the tokenizer "
|
||||
"is still failing two tests and the CLI flag is not wired yet. Continuing next turn."
|
||||
)
|
||||
DONE_REPLY = "Fixed the tokenizer and wired the flag. <promise>DONE</promise>"
|
||||
PROMPT = "Refactor auth module" # stage=autonomous
|
||||
DIRECT_PROMPT = "Implement feature X" # stage=direct
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _offline(monkeypatch, tmp_path):
|
||||
for key in ("OPENAI_BASE_URL", "OPENAI_API_BASE", "ANTHROPIC_BASE_URL"):
|
||||
monkeypatch.setenv(key, "http://127.0.0.1:9")
|
||||
monkeypatch.setenv("HOME", str(tmp_path))
|
||||
|
||||
|
||||
def list_files(path: str) -> str:
|
||||
"""List files in a directory."""
|
||||
return "parser.py tokenizer.py cli.py"
|
||||
|
||||
|
||||
def _script(follow_ups=1):
|
||||
return (
|
||||
[ScriptedModel.tool_call("list_files", {"path": "."})]
|
||||
+ [NOT_DONE] * follow_ups
|
||||
+ [DONE_REPLY] * 3
|
||||
)
|
||||
|
||||
|
||||
def _make(model, **autonomy):
|
||||
return Agent(
|
||||
name="worker", instructions="Do the task.", llm=model, tools=[list_files],
|
||||
output="silent", autonomy={"level": "full_auto", **autonomy},
|
||||
)
|
||||
|
||||
|
||||
def _assert_promise(result):
|
||||
assert result.success is True
|
||||
assert result.completion_reason == "promise"
|
||||
assert "<promise>DONE</promise>" in result.output
|
||||
assert result.iterations >= 2
|
||||
|
||||
|
||||
def test_sync_tool_completion_does_not_preempt_promise():
|
||||
r = _make(ScriptedModel(_script())).run_autonomous(
|
||||
PROMPT, completion_promise="DONE", max_iterations=5)
|
||||
_assert_promise(r)
|
||||
|
||||
|
||||
def test_async_waits_for_promise():
|
||||
r = asyncio.run(_make(ScriptedModel(_script(2))).run_autonomous_async(
|
||||
PROMPT, completion_promise="DONE", max_iterations=5))
|
||||
_assert_promise(r)
|
||||
|
||||
|
||||
def test_sync_direct_stage_waits_for_promise():
|
||||
r = _make(ScriptedModel([NOT_DONE, DONE_REPLY])).run_autonomous(
|
||||
DIRECT_PROMPT, completion_promise="DONE", max_iterations=5)
|
||||
_assert_promise(r)
|
||||
|
||||
|
||||
def test_async_direct_stage_waits_for_promise():
|
||||
r = asyncio.run(_make(ScriptedModel([NOT_DONE, DONE_REPLY])).run_autonomous_async(
|
||||
DIRECT_PROMPT, completion_promise="DONE", max_iterations=5))
|
||||
_assert_promise(r)
|
||||
|
||||
|
||||
def test_promise_never_emitted_ends_unsuccessfully():
|
||||
r = _make(ScriptedModel([NOT_DONE] * 5)).run_autonomous(
|
||||
PROMPT, completion_promise="DONE", max_iterations=3)
|
||||
assert r.success is False
|
||||
assert r.completion_reason == "max_iterations"
|
||||
|
||||
|
||||
def test_iterative_mode_start_honours_promise():
|
||||
agent = Agent(
|
||||
name="loop_agent", instructions="Complete the task.", llm=ScriptedModel(_script()),
|
||||
tools=[list_files], output="silent",
|
||||
autonomy={"mode": "iterative", "max_iterations": 10,
|
||||
"completion_promise": "DONE", "observe": False},
|
||||
)
|
||||
_assert_promise(agent.start(PROMPT))
|
||||
|
||||
|
||||
def test_without_promise_tool_completion_unchanged():
|
||||
r = _make(ScriptedModel(_script())).run_autonomous(PROMPT, max_iterations=5)
|
||||
assert r.success is True
|
||||
assert r.completion_reason == "tool_completion"
|
||||
assert r.iterations == 1
|
||||
|
||||
|
||||
def test_async_tool_path_counts_tools():
|
||||
r = asyncio.run(_make(ScriptedModel(_script(2))).run_autonomous_async(
|
||||
PROMPT, max_iterations=5))
|
||||
assert r.success is True
|
||||
assert r.completion_reason == "tool_completion"
|
||||
assert r.iterations == 1
|
||||
|
||||
|
||||
def test_promise_tag_is_sent_to_model():
|
||||
model = ScriptedModel([DONE_REPLY])
|
||||
_make(model).run_autonomous(PROMPT, completion_promise="DONE", max_iterations=2)
|
||||
assert "<promise>DONE</promise>" in model.requests[0].last_user_message
|
||||
|
||||
|
||||
def test_async_blocked_tool_does_not_count_as_tool_completion():
|
||||
from praisonaiagents.hooks import HookEvent, HookRegistry, HookResult
|
||||
|
||||
registry = HookRegistry()
|
||||
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def deny(data):
|
||||
return HookResult.deny("blocked")
|
||||
|
||||
agent = Agent(
|
||||
name="worker", instructions="Do the task.", llm=ScriptedModel(_script(2)),
|
||||
tools=[list_files], output="silent", hooks=registry,
|
||||
autonomy={"level": "full_auto"},
|
||||
)
|
||||
r = asyncio.run(agent.run_autonomous_async(PROMPT, max_iterations=3))
|
||||
assert r.completion_reason != "tool_completion"
|
||||
@@ -154,9 +154,14 @@ class TestPromptReinjection:
|
||||
completion_promise="DONE"
|
||||
)
|
||||
|
||||
# All prompts should be the original, not "Continue with the task"
|
||||
# All prompts should be the original (plus the promise instruction),
|
||||
# not "Continue with the task"
|
||||
expected = (
|
||||
"Build a REST API\n\nWhen the task is fully complete, include "
|
||||
"<promise>DONE</promise> in your response."
|
||||
)
|
||||
for prompt in prompts_received:
|
||||
assert prompt == "Build a REST API", f"Expected original prompt, got: {prompt}"
|
||||
assert prompt == expected, f"Expected original prompt, got: {prompt}"
|
||||
|
||||
def test_prompt_not_continue_with_task(self):
|
||||
"""Prompt should NOT be 'Continue with the task' after first iteration."""
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
"""Regression tests for #6072: streaming bridge, async tool pool, breaker neutrality."""
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
|
||||
def _agent(**kw):
|
||||
return Agent(name="t", instructions="x", llm="gpt-4o-mini", **kw)
|
||||
|
||||
|
||||
class _Backend:
|
||||
def __init__(self, fail_first=False):
|
||||
self.produced = 0
|
||||
self.calls = 0
|
||||
self.fail_first = fail_first
|
||||
|
||||
async def stream(self, prompt, **kw):
|
||||
self.calls += 1
|
||||
if self.fail_first and self.calls == 1:
|
||||
yield "a"
|
||||
raise RuntimeError("boom")
|
||||
for i in range(100000):
|
||||
self.produced += 1
|
||||
yield f"c{i}"
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
def test_stream_close_in_running_loop_stops_producer():
|
||||
a = _agent()
|
||||
a.backend = _Backend()
|
||||
|
||||
async def main():
|
||||
g = a._delegate_streaming_to_backend("p")
|
||||
next(g)
|
||||
g.close()
|
||||
after_close = a.backend.produced
|
||||
time.sleep(0.3)
|
||||
return a.backend.produced - after_close
|
||||
|
||||
produced = asyncio.run(main())
|
||||
assert produced == 0
|
||||
|
||||
|
||||
def test_backend_runtime_error_does_not_rerun_stream():
|
||||
a = _agent()
|
||||
a.backend = _Backend(fail_first=True)
|
||||
|
||||
async def main():
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
for _ in a._delegate_streaming_to_backend("p"):
|
||||
pass
|
||||
|
||||
asyncio.run(main())
|
||||
assert a.backend.calls == 1
|
||||
|
||||
|
||||
def test_async_sync_tool_with_timeout_uses_agent_pool():
|
||||
seen = {}
|
||||
|
||||
def probe() -> str:
|
||||
seen["thread"] = threading.current_thread().name
|
||||
return "ok"
|
||||
|
||||
a = _agent(tools=[probe])
|
||||
a._tool_timeout = 5
|
||||
assert asyncio.run(a.execute_tool_async("probe", {})) == "ok"
|
||||
assert seen["thread"].startswith("tool-t")
|
||||
|
||||
|
||||
def test_async_timed_out_tool_does_not_hold_default_executor():
|
||||
def hang() -> str:
|
||||
time.sleep(2)
|
||||
return "ok"
|
||||
|
||||
a = _agent(tools=[hang])
|
||||
a._tool_timeout = 0.1
|
||||
|
||||
async def main():
|
||||
tasks = [asyncio.create_task(a.execute_tool_async("hang", {})) for _ in range(40)]
|
||||
await asyncio.gather(*tasks)
|
||||
start = time.monotonic()
|
||||
await asyncio.to_thread(lambda: 1)
|
||||
return time.monotonic() - start
|
||||
|
||||
assert asyncio.run(main()) < 1.0
|
||||
|
||||
|
||||
def test_caller_errors_do_not_open_breaker():
|
||||
def read_file(path: str) -> str:
|
||||
raise FileNotFoundError(path)
|
||||
|
||||
a = _agent(tools=[read_file])
|
||||
for i in range(7):
|
||||
result = a._execute_tool_with_circuit_breaker_impl("read_file", {"path": f"/nope{i}"})
|
||||
assert not result.get("circuit_open")
|
||||
|
||||
|
||||
def test_open_breaker_remediation_uses_live_config():
|
||||
a = _agent()
|
||||
breaker = a._get_tool_circuit_breaker("x")
|
||||
breaker.config.recovery_timeout = 7
|
||||
assert "(7s)" in a._circuit_breaker_open_result("x")["remediation"]
|
||||
|
||||
|
||||
def test_stream_close_runs_backend_cleanup():
|
||||
cleaned = threading.Event()
|
||||
|
||||
class _CleanupBackend:
|
||||
async def stream(self, prompt, **kw):
|
||||
try:
|
||||
while True:
|
||||
yield "c"
|
||||
await asyncio.sleep(0)
|
||||
finally:
|
||||
cleaned.set()
|
||||
|
||||
a = _agent()
|
||||
a.backend = _CleanupBackend()
|
||||
g = a._delegate_streaming_to_backend("p")
|
||||
next(g)
|
||||
g.close()
|
||||
assert cleaned.wait(2)
|
||||
|
||||
|
||||
def test_parallel_healthy_timed_tools_do_not_queue_into_timeout():
|
||||
def work(i: int) -> str:
|
||||
time.sleep(0.6)
|
||||
return "ok"
|
||||
|
||||
a = _agent(tools=[work])
|
||||
a._tool_timeout = 1
|
||||
|
||||
async def main():
|
||||
return await asyncio.gather(*(a.execute_tool_async("work", {"i": i}) for i in range(5)))
|
||||
|
||||
# Loop-guard may append advisory notes; what matters is no timeout.
|
||||
assert all(isinstance(r, str) and r.startswith("ok") for r in asyncio.run(main()))
|
||||
|
||||
|
||||
def test_hung_spill_threads_are_capped_across_bursts():
|
||||
from praisonaiagents.agent import tool_execution as te
|
||||
|
||||
release = threading.Event()
|
||||
a = _agent()
|
||||
started = 0
|
||||
try:
|
||||
for _ in range(5):
|
||||
# A burst of concurrent calls that all hang past their timeout;
|
||||
# each timed-out burst retires the stuck pool as the async path does.
|
||||
for _ in range(20):
|
||||
future, executor, _ = a._submit_to_tool_pool(lambda: release.wait(10), spill=True)
|
||||
started += future is not None
|
||||
a._retire_tool_executor(a._tool_executor)
|
||||
assert a._tool_spill_inflight <= te._MAX_SPILL_TOOL_THREADS
|
||||
cap = (
|
||||
te._MAX_SPILL_TOOL_THREADS
|
||||
+ (te._MAX_ORPHANED_TOOL_EXECUTORS + 1) * te._TOOL_POOL_MAX_WORKERS
|
||||
)
|
||||
assert started <= cap
|
||||
finally:
|
||||
release.set()
|
||||
deadline = time.monotonic() + 5
|
||||
while a._tool_spill_inflight and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert a._tool_spill_inflight == 0
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Budget and cost accounting on the custom-LLM (``provider/model``) path (#6029).
|
||||
|
||||
That path hands the whole turn, tool loop included, to ``LLM.get_response*``.
|
||||
Every call it makes must reach the Agent's cost/token counters so
|
||||
``ExecutionConfig.max_budget`` and the autonomy spend caps can fire.
|
||||
``litellm.completion`` is stubbed; nothing leaves the machine.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("litellm")
|
||||
|
||||
from praisonaiagents import Agent # noqa: E402
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig # noqa: E402
|
||||
from praisonaiagents.errors import BudgetExceededError # noqa: E402
|
||||
from praisonaiagents.utils.cost_utils import calculate_llm_cost # noqa: E402
|
||||
|
||||
MODEL = "anthropic/claude-3-5-sonnet-20240620"
|
||||
USAGE = {"prompt_tokens": 1500, "completion_tokens": 1500, "total_tokens": 3000}
|
||||
|
||||
|
||||
def _litellm():
|
||||
"""The litellm module LLM will import now.
|
||||
|
||||
Other tests evict ``litellm`` from sys.modules, so a module-level binding
|
||||
can go stale and stubbing it would let the real client run.
|
||||
"""
|
||||
import litellm
|
||||
return litellm
|
||||
|
||||
|
||||
def _text(model):
|
||||
return _litellm().ModelResponse(model=model, usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": "ok"}}])
|
||||
|
||||
|
||||
def _tool_call(model, n):
|
||||
return _litellm().ModelResponse(model=model, usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "tool_calls",
|
||||
"message": {"role": "assistant", "content": None, "tool_calls": [{
|
||||
"id": f"call_{n}", "type": "function",
|
||||
"function": {"name": "step", "arguments": "{}"}}]}}])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def calls(monkeypatch):
|
||||
"""Stub litellm; answer with a tool call while ``tool_rounds`` remain."""
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
# Agent caches its env-derived default model per process; keep the fake
|
||||
# key's Anthropic default from leaking into later tests.
|
||||
monkeypatch.setattr(Agent, "_default_model", None)
|
||||
monkeypatch.setattr(Agent, "_default_model_checked", False)
|
||||
state = {"n": 0, "tool_rounds": 0}
|
||||
|
||||
def fake(*a, **k):
|
||||
state["n"] += 1
|
||||
if k.get("tools") and state["tool_rounds"] > 0:
|
||||
state["tool_rounds"] -= 1
|
||||
return _tool_call(k.get("model"), state["n"])
|
||||
return _text(k.get("model"))
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
litellm = _litellm()
|
||||
monkeypatch.setattr(litellm, "completion", fake)
|
||||
monkeypatch.setattr(litellm, "acompletion", afake)
|
||||
return state
|
||||
|
||||
|
||||
def _agent(**kw):
|
||||
return Agent(name="p", instructions="x", llm=MODEL, output="silent", **kw)
|
||||
|
||||
|
||||
def step() -> str:
|
||||
"""One step of work."""
|
||||
return "done"
|
||||
|
||||
|
||||
def test_sync_chat_counts_cost_and_tokens(calls):
|
||||
agent = _agent()
|
||||
assert agent.chat("hi") == "ok"
|
||||
summary = agent.cost_summary
|
||||
assert summary["llm_calls"] == 1
|
||||
assert summary["tokens_in"] == 1500 and summary["tokens_out"] == 1500
|
||||
assert agent.total_cost > 0
|
||||
|
||||
|
||||
def test_sync_budget_raises_between_turns(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.05))
|
||||
with pytest.raises(BudgetExceededError):
|
||||
for i in range(5):
|
||||
agent.chat(f"hello {i}")
|
||||
assert calls["n"] < 5
|
||||
assert agent.total_cost >= 0.05
|
||||
|
||||
|
||||
def test_async_budget_raises_between_turns(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.05))
|
||||
|
||||
async def run():
|
||||
for i in range(5):
|
||||
await agent.achat(f"hello {i}")
|
||||
|
||||
with pytest.raises(BudgetExceededError):
|
||||
asyncio.run(run())
|
||||
assert calls["n"] < 5
|
||||
assert agent.total_cost >= 0.05
|
||||
|
||||
|
||||
def test_budget_trips_mid_tool_loop(calls):
|
||||
calls["tool_rounds"] = 10 # the model would keep calling tools for 11 calls
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=1.0))
|
||||
agent.chat("price one call")
|
||||
one_call = agent.total_cost
|
||||
calls["n"] = 0
|
||||
# A cap between one and two calls must stop the loop at the second call.
|
||||
agent = _agent(tools=[step], execution=ExecutionConfig(max_budget=one_call * 1.5))
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("go")
|
||||
assert calls["n"] == 2
|
||||
|
||||
|
||||
def test_warn_mode_does_not_raise(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.001, on_budget_exceeded="warn"))
|
||||
assert agent.chat("hi") == "ok"
|
||||
assert agent.total_cost > 0.001
|
||||
|
||||
|
||||
def test_autonomy_budget_exhausted_on_provider_path(calls):
|
||||
agent = _agent(autonomy={"max_budget_usd": 0.001})
|
||||
result = agent.run_autonomous("go", max_iterations=3)
|
||||
assert result.completion_reason == "budget_exhausted"
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_sink_ignores_other_llm_instances(calls):
|
||||
"""A call made by a different LLM object is not billed to this agent."""
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
|
||||
agent = _agent()
|
||||
other = LLM(model=MODEL)
|
||||
with agent._llm_usage_sink():
|
||||
other._track_token_usage(_text(MODEL), MODEL)
|
||||
assert agent.cost_summary["llm_calls"] == 0
|
||||
|
||||
|
||||
class _Unavailable(Exception):
|
||||
status_code = 503
|
||||
|
||||
|
||||
def test_fallback_model_spend_is_counted(calls, monkeypatch):
|
||||
"""A turn answered by a ``fallback_models`` entry is still billed to the agent."""
|
||||
backup = "anthropic/claude-3-haiku-20240307"
|
||||
|
||||
def fake(*a, **k):
|
||||
calls["n"] += 1
|
||||
if k.get("model") != backup:
|
||||
raise _Unavailable("503 overloaded")
|
||||
return _text(backup)
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
monkeypatch.setattr(_litellm(), "completion", fake)
|
||||
monkeypatch.setattr(_litellm(), "acompletion", afake)
|
||||
for run in (lambda a: a.chat("hi"), lambda a: asyncio.run(a.achat("hi"))):
|
||||
agent = _agent(retry=False) # fail over at once, no backoff sleeps
|
||||
agent.fallback_models = [backup]
|
||||
assert run(agent) == "ok"
|
||||
assert agent.cost_summary["llm_calls"] == 1
|
||||
# Priced as the model that answered, not the (pricier) primary.
|
||||
assert agent.total_cost == pytest.approx(calculate_llm_cost(1500, 1500, backup))
|
||||
assert agent.total_cost != pytest.approx(calculate_llm_cost(1500, 1500, MODEL))
|
||||
|
||||
|
||||
def test_retry_after_exhaustion_sends_no_request(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.001))
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("first")
|
||||
sent = calls["n"]
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("retry")
|
||||
with pytest.raises(BudgetExceededError):
|
||||
asyncio.run(agent.achat("retry async"))
|
||||
assert calls["n"] == sent
|
||||
|
||||
|
||||
def test_autonomy_budget_error_keeps_budget_exhausted(calls):
|
||||
"""A max_budget stop inside the turn still reports the autonomy cap result."""
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.001),
|
||||
autonomy={"max_budget_usd": 0.001})
|
||||
result = agent.run_autonomous("go", max_iterations=3)
|
||||
assert result.completion_reason == "budget_exhausted"
|
||||
result = asyncio.run(_agent(execution=ExecutionConfig(max_budget=0.001),
|
||||
autonomy={"max_budget_usd": 0.001}).run_autonomous_async("go", max_iterations=3))
|
||||
assert result.completion_reason == "budget_exhausted"
|
||||
@@ -388,7 +388,8 @@ def test_cancellation_during_final_save_finishes_record_and_cleanup(recorded_age
|
||||
assert db.starts[0]["run_id"] == db.ends[0]["run_id"]
|
||||
assert db.ends[0]["status"] == "completed"
|
||||
assert agent._active_turn_token is None
|
||||
emitter.agent_end.assert_called_once_with(agent.name)
|
||||
# The answer came back; only the final save was cancelled.
|
||||
emitter.agent_end.assert_called_once_with(agent.name, {"status": "completed"})
|
||||
token.close.assert_called_once()
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Regression tests for #5795: tool/MCP calls must not be timed out by queue wait."""
|
||||
import asyncio
|
||||
import concurrent.futures as cf
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.hooks import HookEvent, HookRegistry
|
||||
|
||||
|
||||
def _allow_all_hooks():
|
||||
"""Own registry (no leaked global hooks) with a sync BEFORE_TOOL hook,
|
||||
so parallel tool calls also run the hook concurrently."""
|
||||
registry = HookRegistry()
|
||||
registry.register_function(HookEvent.BEFORE_TOOL, lambda data: None)
|
||||
return registry
|
||||
|
||||
|
||||
def test_parallel_sync_tools_not_timed_out_by_queue_wait():
|
||||
def slow(n: int) -> str:
|
||||
"""slow tool"""
|
||||
time.sleep(0.5)
|
||||
return f"ok{n}"
|
||||
|
||||
agent = Agent(name="t", instructions="x", tools=[slow], llm="gpt-4o-mini", hooks=_allow_all_hooks())
|
||||
agent._tool_timeout = 0.8
|
||||
with cf.ThreadPoolExecutor(5) as ex:
|
||||
results = list(ex.map(lambda i: agent.execute_tool("slow", {"n": i}), range(5)))
|
||||
for i, result in enumerate(results):
|
||||
assert isinstance(result, str) and result.startswith(f"ok{i}"), result
|
||||
|
||||
|
||||
def test_parallel_async_sync_tools_not_timed_out_by_queue_wait():
|
||||
def slow(n: int) -> str:
|
||||
"""slow tool"""
|
||||
time.sleep(0.5)
|
||||
return f"ok{n}"
|
||||
|
||||
agent = Agent(name="a", instructions="x", tools=[slow], llm="gpt-4o-mini", hooks=_allow_all_hooks())
|
||||
agent._tool_timeout = 0.8
|
||||
|
||||
async def main():
|
||||
# Twice the pool size: the second half queues for a worker first.
|
||||
return await asyncio.gather(
|
||||
*(agent.execute_tool_async("slow", {"n": i}) for i in range(10))
|
||||
)
|
||||
|
||||
for i, result in enumerate(asyncio.run(main())):
|
||||
assert isinstance(result, str) and result.startswith(f"ok{i}"), result
|
||||
|
||||
|
||||
def test_hung_async_sync_tools_do_not_starve_default_executor():
|
||||
release = threading.Event()
|
||||
tools = []
|
||||
for i in range(10):
|
||||
def hang(x: int = 0) -> str:
|
||||
"""hang"""
|
||||
release.wait(10)
|
||||
return "x"
|
||||
hang.__name__ = f"hang{i}"
|
||||
tools.append(hang)
|
||||
|
||||
agent = Agent(name="u", instructions="x", tools=tools, llm="gpt-4o-mini", hooks=HookRegistry())
|
||||
agent._tool_timeout = 0.2
|
||||
|
||||
async def main():
|
||||
for i in range(10):
|
||||
result = await asyncio.wait_for(agent.execute_tool_async(f"hang{i}", {"x": 1}), 5)
|
||||
# Times out, or is refused fast once hung calls hit the pool cap.
|
||||
assert isinstance(result, dict) and (result.get("timeout") or result.get("tool_pool_exhausted")), result
|
||||
# The loop's default executor is still free for other to_thread users.
|
||||
assert await asyncio.wait_for(asyncio.to_thread(lambda: 42), 2) == 42
|
||||
|
||||
try:
|
||||
asyncio.run(main())
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
|
||||
def test_transport_wrapper_cancels_call_on_timeout():
|
||||
from praisonaiagents.mcp.mcp_sse import SSEMCPTool
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
threading.Thread(target=loop.run_forever, daemon=True).start()
|
||||
finished = threading.Event()
|
||||
|
||||
class _Session:
|
||||
async def call_tool(self, name, kwargs):
|
||||
await asyncio.sleep(1.0)
|
||||
finished.set()
|
||||
|
||||
try:
|
||||
tool = SSEMCPTool("t", "d", _Session(), {}, timeout=0.1)
|
||||
tool._loop = loop
|
||||
result = tool()
|
||||
assert isinstance(result, dict) and result.get("timeout"), result
|
||||
time.sleep(1.3)
|
||||
assert not finished.is_set()
|
||||
finally:
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
|
||||
|
||||
def test_stdio_runner_runs_requests_concurrently_and_cancels_on_shutdown(monkeypatch):
|
||||
import praisonaiagents.mcp.mcp as mcp_mod
|
||||
|
||||
class _Session:
|
||||
async def initialize(self):
|
||||
pass
|
||||
|
||||
async def list_tools(self):
|
||||
return SimpleNamespace(tools=[])
|
||||
|
||||
async def call_tool(self, name, arguments):
|
||||
await asyncio.sleep(10 if name == "slow" else 0)
|
||||
return name
|
||||
|
||||
@asynccontextmanager
|
||||
async def _stdio_client(params):
|
||||
yield None, None
|
||||
|
||||
@asynccontextmanager
|
||||
async def _client_session(read, write):
|
||||
yield _Session()
|
||||
|
||||
async def _no_discovery(self, session):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "stdio_client", _stdio_client)
|
||||
monkeypatch.setattr(mcp_mod, "ClientSession", _client_session)
|
||||
monkeypatch.setattr(mcp_mod.MCPToolRunner, "_discover_resources_and_prompts", _no_discovery)
|
||||
|
||||
runner = mcp_mod.MCPToolRunner(None, timeout=30)
|
||||
assert runner.initialized.wait(5) and runner._init_error is None
|
||||
|
||||
slow_q, fast_q = queue.Queue(), queue.Queue()
|
||||
runner.queue.put((slow_q, "tool", "slow", {}, threading.Event()))
|
||||
runner.queue.put((fast_q, "tool", "fast", {}, threading.Event()))
|
||||
# The fast call is not serialized behind the in-flight slow one.
|
||||
assert fast_q.get(timeout=2) == (True, "fast")
|
||||
|
||||
runner.queue.put(None)
|
||||
ok, message = slow_q.get(timeout=2)
|
||||
assert not ok and "shut down" in message
|
||||
runner.join(2)
|
||||
assert not runner.is_alive()
|
||||
|
||||
|
||||
def test_agent_deepcopy_after_async_tool_use():
|
||||
import copy
|
||||
|
||||
def echo(x: int) -> int:
|
||||
"""echo"""
|
||||
return x
|
||||
|
||||
agent = Agent(name="d", instructions="x", tools=[echo], llm="gpt-4o-mini")
|
||||
assert asyncio.run(agent.execute_tool_async("echo", {"x": 1})) == 1
|
||||
clone = copy.deepcopy(agent)
|
||||
assert "_tool_executor" not in clone.__dict__
|
||||
assert asyncio.run(clone.execute_tool_async("echo", {"x": 2})) == 2
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Loop guard per-turn time limit is configurable via ExecutionConfig (#6026)."""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
import praisonaiagents.escalation.loop_guard as lg
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig
|
||||
from praisonaiagents.escalation.loop_guard import GuardAction, LoopGuardConfig
|
||||
from praisonaiagents.model_harness import ScriptedModel
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clock(monkeypatch):
|
||||
state = {"t": 1_000_000.0}
|
||||
|
||||
class FakeTime:
|
||||
@staticmethod
|
||||
def time():
|
||||
return state["t"]
|
||||
|
||||
monkeypatch.setattr(lg, "time", FakeTime)
|
||||
return state
|
||||
|
||||
|
||||
def _agent(clock, tool_seconds, execution=None, calls=None):
|
||||
calls = [] if calls is None else calls
|
||||
|
||||
def run_tests(suite: str) -> str:
|
||||
"""Run a (slow) test suite."""
|
||||
calls.append(suite)
|
||||
clock["t"] += tool_seconds
|
||||
return f"{suite}: 412 passed"
|
||||
|
||||
model = ScriptedModel([
|
||||
ScriptedModel.tool_call("run_tests", {"suite": "unit"}),
|
||||
ScriptedModel.tool_call("run_tests", {"suite": "integration"}),
|
||||
"All suites passed.",
|
||||
])
|
||||
kwargs = {"execution": execution} if execution is not None else {}
|
||||
return Agent(name="coder", instructions="Run the tests.", llm=model,
|
||||
tools=[run_tests], output="silent", **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
def test_default_turn_with_130s_tool_completes(clock, use_async):
|
||||
agent = _agent(clock, 130.0)
|
||||
if use_async:
|
||||
out = asyncio.run(agent.achat("Run both suites"))
|
||||
else:
|
||||
out = agent.chat("Run both suites")
|
||||
assert "All suites passed" in str(out)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
def test_configured_limit_halts(clock, use_async):
|
||||
calls = []
|
||||
agent = _agent(clock, 130.0, ExecutionConfig(max_execution_time=60), calls)
|
||||
if use_async:
|
||||
asyncio.run(agent.achat("Run both suites"))
|
||||
else:
|
||||
with pytest.raises(Exception, match="Turn time limit exceeded"):
|
||||
agent.chat("Run both suites")
|
||||
# The second tool call is refused by the loop guard on both paths.
|
||||
assert calls == ["unit"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
def test_zero_disables_time_halt(clock, use_async):
|
||||
agent = _agent(clock, 10_000.0, ExecutionConfig(max_execution_time=0))
|
||||
if use_async:
|
||||
out = asyncio.run(agent.achat("Run both suites"))
|
||||
else:
|
||||
out = agent.chat("Run both suites")
|
||||
assert "All suites passed" in str(out)
|
||||
|
||||
|
||||
def test_ensure_loop_guard_reads_max_execution_time():
|
||||
agent = Agent(name="a", instructions="x", llm=ScriptedModel(["ok"]),
|
||||
output="silent", execution=ExecutionConfig(max_execution_time=900))
|
||||
assert agent._ensure_loop_guard().config.max_time_per_turn == 900.0
|
||||
|
||||
|
||||
def test_real_loop_still_halts_without_time_limit():
|
||||
guard = lg.LoopGuard(LoopGuardConfig(max_time_per_turn=float("inf")))
|
||||
actions = []
|
||||
for _ in range(20):
|
||||
decision = guard.check("read_file", {"path": "a.txt"})
|
||||
actions.append(decision.action)
|
||||
if decision.action == GuardAction.HALT:
|
||||
break
|
||||
guard.record("read_file", {"path": "a.txt"}, success=True, result="same")
|
||||
assert actions[-1] == GuardAction.HALT
|
||||
@@ -0,0 +1,126 @@
|
||||
"""``ExecutionConfig.max_execution_time`` and ``run_autonomous(timeout_seconds=)`` (#6029).
|
||||
|
||||
Both caps become a deadline on the turn's cancel token, so the LLM tool loop
|
||||
stops at its next boundary instead of running to completion.
|
||||
``litellm.completion`` is stubbed; nothing leaves the machine.
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("litellm")
|
||||
|
||||
from praisonaiagents import Agent # noqa: E402
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig # noqa: E402
|
||||
from praisonaiagents.escalation.loop_guard import LoopGuardConfig # noqa: E402
|
||||
|
||||
MODEL = "anthropic/claude-3-5-sonnet-20240620"
|
||||
USAGE = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def calls(monkeypatch):
|
||||
"""Stub litellm so the model asks for ``slow_step`` on every call."""
|
||||
# Import here, not at module level: other tests evict ``litellm`` from
|
||||
# sys.modules, and LLM then imports a fresh module object at call time.
|
||||
import litellm
|
||||
from litellm import ModelResponse
|
||||
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
# Agent caches its env-derived default model per process; keep the fake
|
||||
# key's Anthropic default from leaking into later tests.
|
||||
monkeypatch.setattr(Agent, "_default_model", None)
|
||||
monkeypatch.setattr(Agent, "_default_model_checked", False)
|
||||
state = {"n": 0}
|
||||
|
||||
def fake(*a, **k):
|
||||
state["n"] += 1
|
||||
if not k.get("tools"):
|
||||
return ModelResponse(model=k.get("model"), usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": "done"}}])
|
||||
return ModelResponse(model=k.get("model"), usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "tool_calls",
|
||||
"message": {"role": "assistant", "content": None, "tool_calls": [{
|
||||
"id": f"call_{state['n']}", "type": "function",
|
||||
"function": {"name": "slow_step", "arguments": "{}"}}]}}])
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
monkeypatch.setattr(litellm, "completion", fake)
|
||||
monkeypatch.setattr(litellm, "acompletion", afake)
|
||||
return state
|
||||
|
||||
|
||||
def slow_step() -> str:
|
||||
"""A slow step."""
|
||||
time.sleep(0.4)
|
||||
return "ok"
|
||||
|
||||
|
||||
def _agent(**kw):
|
||||
return Agent(name="t", instructions="x", llm=MODEL, tools=[slow_step],
|
||||
output="silent", **kw)
|
||||
|
||||
|
||||
def test_max_execution_time_stops_tool_loop(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_execution_time=1, max_iter=50))
|
||||
t0 = time.time()
|
||||
out = agent.chat("go")
|
||||
elapsed = time.time() - t0
|
||||
assert elapsed < 3
|
||||
assert "timeout" in (out or "")
|
||||
assert agent.last_stop_reason == "timeout"
|
||||
assert calls["n"] < 10
|
||||
|
||||
|
||||
def test_max_execution_time_stops_tool_loop_async(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_execution_time=1, max_iter=50))
|
||||
t0 = time.time()
|
||||
out = asyncio.run(agent.achat("go"))
|
||||
assert time.time() - t0 < 3
|
||||
assert "timeout" in (out or "")
|
||||
assert agent.last_stop_reason == "timeout"
|
||||
|
||||
|
||||
def test_unset_max_execution_time_creates_no_deadline():
|
||||
agent = Agent(name="t", instructions="x", output="silent")
|
||||
assert agent._turn_cancel_token(None, explicit=False) is None
|
||||
|
||||
|
||||
def test_loop_guard_turn_limit_follows_max_execution_time():
|
||||
agent = Agent(name="t", instructions="x", output="silent",
|
||||
execution=ExecutionConfig(max_execution_time=600))
|
||||
assert agent._ensure_loop_guard().config.max_time_per_turn == 600
|
||||
default = Agent(name="d", instructions="x", output="silent")
|
||||
assert default._ensure_loop_guard().config.max_time_per_turn == LoopGuardConfig.max_time_per_turn
|
||||
|
||||
|
||||
def test_run_autonomous_timeout_bounds_the_running_turn(calls):
|
||||
agent = _agent(autonomy=True, execution=ExecutionConfig(max_iter=50))
|
||||
t0 = time.time()
|
||||
result = agent.run_autonomous("go", max_iterations=1, timeout_seconds=1)
|
||||
assert time.time() - t0 < 3
|
||||
assert result.completion_reason == "timeout"
|
||||
assert not result.success
|
||||
|
||||
|
||||
def test_run_autonomous_async_timeout_bounds_the_running_turn(calls):
|
||||
agent = _agent(autonomy=True, execution=ExecutionConfig(max_iter=50))
|
||||
t0 = time.time()
|
||||
result = asyncio.run(agent.run_autonomous_async("go", max_iterations=1, timeout_seconds=1))
|
||||
assert time.time() - t0 < 3
|
||||
assert result.completion_reason == "timeout"
|
||||
|
||||
|
||||
def test_run_autonomous_reports_turn_max_execution_time(calls):
|
||||
"""A turn cut by max_execution_time is a failed timeout, not a goal."""
|
||||
agent = _agent(autonomy=True, execution=ExecutionConfig(max_execution_time=1, max_iter=50))
|
||||
result = agent.run_autonomous("go", max_iterations=3)
|
||||
assert result.completion_reason == "timeout"
|
||||
assert not result.success
|
||||
result = asyncio.run(agent.run_autonomous_async("go", max_iterations=3))
|
||||
assert result.completion_reason == "timeout"
|
||||
assert not result.success
|
||||
@@ -0,0 +1,44 @@
|
||||
"""GitHub write tools must go through the approval system (issue #6023)."""
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.approval import get_approval_registry
|
||||
from praisonaiagents.approval.registry import PERMISSION_PRESETS
|
||||
|
||||
GITHUB_WRITE_TOOLS = (
|
||||
"github_create_branch",
|
||||
"github_commit_and_push",
|
||||
"github_create_pull_request",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", GITHUB_WRITE_TOOLS)
|
||||
def test_registered_as_high_risk(name):
|
||||
reg = get_approval_registry()
|
||||
assert reg.is_required(name)
|
||||
# high (not critical) so YAML/declarative approve lists still apply
|
||||
assert reg.get_risk_level(name) == "high"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", GITHUB_WRITE_TOOLS)
|
||||
def test_in_deny_presets(name):
|
||||
for preset in ("default", "safe", "read_only"):
|
||||
assert name in PERMISSION_PRESETS[preset]
|
||||
assert name not in PERMISSION_PRESETS["full"]
|
||||
|
||||
|
||||
def test_unattended_agent_denied_by_default(monkeypatch):
|
||||
monkeypatch.delenv("PRAISONAI_TOOL_SAFETY", raising=False)
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(instructions="t", llm="gpt-4o-mini", approval=False)
|
||||
for name in ("github_commit_and_push", "github_create_pull_request"):
|
||||
result = agent._check_tool_approval_sync(name, {})
|
||||
assert isinstance(result, dict) and result.get("permission_denied")
|
||||
|
||||
|
||||
def test_explicit_full_preset_allows(monkeypatch):
|
||||
monkeypatch.setenv("PRAISONAI_TOOL_SAFETY", "off")
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(instructions="t", llm="gpt-4o-mini", approval=False)
|
||||
assert "github_commit_and_push" not in (agent._perm_deny or ())
|
||||
@@ -0,0 +1,176 @@
|
||||
"""Config options that were accepted but never read must either work or fail loud.
|
||||
|
||||
Covers GuardrailConfig.on_fail, ToolConfig.output_max_lines/output_direction and
|
||||
rerank=True on the Agent's RAG pipeline (issue #6042). No network calls.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.config import GuardrailConfig, ToolConfig
|
||||
from praisonaiagents.config.presets import GUARDRAIL_PRESETS
|
||||
|
||||
|
||||
def _ok(task_output):
|
||||
return True, task_output
|
||||
|
||||
|
||||
# --- GuardrailConfig.on_fail ------------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize("on_fail", ["skip", "raise"])
|
||||
def test_unsupported_on_fail_fails_loud(on_fail):
|
||||
with pytest.raises(ValueError, match="on_fail"):
|
||||
Agent(name="g", instructions="x", llm="gpt-4o-mini",
|
||||
guardrails=GuardrailConfig(validator=_ok, on_fail=on_fail))
|
||||
|
||||
|
||||
def test_retry_on_fail_still_accepted():
|
||||
a = Agent(name="g", instructions="x", llm="gpt-4o-mini",
|
||||
guardrails=GuardrailConfig(validator=_ok, max_retries=2, on_fail="retry"))
|
||||
assert a.max_guardrail_retries == 2
|
||||
|
||||
|
||||
def test_presets_carry_no_unimplemented_on_fail():
|
||||
assert all("on_fail" not in p for p in GUARDRAIL_PRESETS.values())
|
||||
a = Agent(name="g", instructions="x", llm="gpt-4o-mini",
|
||||
guardrails=["strict", {"validator": _ok}])
|
||||
assert a.max_guardrail_retries == 5
|
||||
assert a.guardrail is _ok
|
||||
|
||||
|
||||
# --- ToolConfig.output_max_lines / output_direction --------------------------
|
||||
|
||||
def many_lines() -> str:
|
||||
"""Return 200 short lines."""
|
||||
return "\n".join(f"line {i}" for i in range(200))
|
||||
|
||||
|
||||
def _run_tool(cfg):
|
||||
a = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[many_lines],
|
||||
tool_config=cfg, output="silent")
|
||||
return str(a.execute_tool("many_lines", {}))
|
||||
|
||||
|
||||
def test_default_tool_config_does_not_cap_lines():
|
||||
out = _run_tool(ToolConfig())
|
||||
assert out.splitlines()[-1] == "line 199"
|
||||
assert out.count("\n") == 199
|
||||
|
||||
|
||||
def test_output_max_lines_head():
|
||||
lines = _run_tool(ToolConfig(output_max_lines=5, output_direction="head")).splitlines()
|
||||
assert lines[:5] == [f"line {i}" for i in range(5)]
|
||||
assert "line 5" not in lines and "line 199" not in lines
|
||||
assert "200 lines total" in lines[5]
|
||||
|
||||
|
||||
def test_output_max_lines_tail():
|
||||
lines = _run_tool(ToolConfig(output_max_lines=5, output_direction="tail")).splitlines()
|
||||
assert lines[-5:] == [f"line {i}" for i in range(195, 200)]
|
||||
assert "line 0" not in lines
|
||||
|
||||
|
||||
def test_output_max_lines_both():
|
||||
lines = _run_tool(ToolConfig(output_max_lines=4)).splitlines()
|
||||
assert lines[:2] == ["line 0", "line 1"]
|
||||
assert lines[-2:] == ["line 198", "line 199"]
|
||||
assert len(lines) == 5 # 4 kept lines + marker
|
||||
|
||||
|
||||
def test_output_direction_applies_to_char_cap():
|
||||
out = _run_tool(ToolConfig(output_limit=50, output_direction="head"))
|
||||
assert out.startswith("line 0\nline 1")
|
||||
assert "line 199" not in out
|
||||
|
||||
|
||||
def short_lines() -> str:
|
||||
"""Return ten one-character lines."""
|
||||
return "\n".join("abcdefghij")
|
||||
|
||||
|
||||
def big_dict() -> dict:
|
||||
"""Return a dict with one large field."""
|
||||
return {"body": "START" + "x" * 45000 + "END"}
|
||||
|
||||
|
||||
def test_line_cap_applies_to_short_lines():
|
||||
a = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[short_lines],
|
||||
tool_config=ToolConfig(output_max_lines=5, output_direction="head"), output="silent")
|
||||
lines = str(a.execute_tool("short_lines", {})).splitlines()
|
||||
assert lines[:5] == list("abcde")
|
||||
assert "j" not in lines
|
||||
|
||||
|
||||
def test_output_direction_applies_to_dict_results():
|
||||
a = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[big_dict],
|
||||
tool_config=ToolConfig(output_limit=40000, output_direction="head"), output="silent")
|
||||
out = str(a.execute_tool("big_dict", {}))
|
||||
assert "START" in out and "END" not in out
|
||||
|
||||
|
||||
# --- rerank=True on Agent.rag -----------------------------------------------
|
||||
|
||||
class StubKnowledge:
|
||||
"""Minimal Knowledge stand-in: search() returns hits worst-first."""
|
||||
|
||||
def add(self, *a, **k):
|
||||
pass
|
||||
|
||||
def search(self, query, **kw):
|
||||
return {"results": [
|
||||
{"memory": "bananas are yellow", "score": 0.9, "metadata": {"source": "a"}},
|
||||
{"memory": "paris is the capital of france", "score": 0.8, "metadata": {"source": "b"}},
|
||||
]}
|
||||
|
||||
|
||||
def _rag(retrieval_config=None):
|
||||
a = Agent(name="r", instructions="x", llm="gpt-4o-mini", knowledge=StubKnowledge(),
|
||||
retrieval_config=retrieval_config)
|
||||
return a.rag
|
||||
|
||||
|
||||
def test_rerank_true_reranks_and_reports_it():
|
||||
pack = _rag({"rerank": True}).retrieve("capital of france")
|
||||
assert pack.metadata["reranked"] is True
|
||||
assert pack.context.index("paris") < pack.context.index("bananas")
|
||||
|
||||
|
||||
def test_rerank_false_reports_not_reranked():
|
||||
rag = _rag()
|
||||
assert rag.reranker is None
|
||||
assert rag.retrieve("capital of france").metadata["reranked"] is False
|
||||
|
||||
|
||||
def test_rerank_override_without_reranker_reports_false():
|
||||
assert _rag().retrieve("q", rerank=True).metadata["reranked"] is False
|
||||
|
||||
|
||||
def test_rerank_model_leaves_reranking_to_backend():
|
||||
rag = _rag({"rerank": True, "rerank_model": "cohere/rerank-v3"})
|
||||
assert rag.reranker is None
|
||||
|
||||
|
||||
def test_rerank_model_backend_reranking_reported():
|
||||
class BackendRerankKnowledge(StubKnowledge):
|
||||
config = {"reranker": {"default_rerank": True, "provider": "cohere"}}
|
||||
|
||||
a = Agent(name="r", instructions="x", llm="gpt-4o-mini", knowledge=BackendRerankKnowledge(),
|
||||
retrieval_config={"rerank": True, "rerank_model": "cohere/rerank-v3"})
|
||||
assert a.rag.reranker is None
|
||||
assert a.rag.retrieve("q").metadata["reranked"] is True
|
||||
assert a.rag.retrieve("q", rerank=False).metadata["reranked"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_true_async():
|
||||
pack = await _rag({"rerank": True}).aretrieve("capital of france")
|
||||
assert pack.metadata["reranked"] is True
|
||||
assert pack.context.index("paris") < pack.context.index("bananas")
|
||||
|
||||
|
||||
def test_simple_reranker_accepts_search_result_items():
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
from praisonaiagents.rag.retriever import SimpleReranker
|
||||
hits = [SearchResultItem(text="bananas are yellow", score=0.9),
|
||||
SearchResultItem(text="paris is the capital of france", score=0.8)]
|
||||
ranked = SimpleReranker().rerank("capital of france", hits)
|
||||
assert ranked[0]["text"].startswith("paris")
|
||||
@@ -275,3 +275,113 @@ def test_async_fallback_passes_override_and_leaves_shared_state(monkeypatch):
|
||||
assert used[0] is not original_dispatcher
|
||||
assert agent.llm == "gpt-4o-mini" # shared model untouched
|
||||
assert agent._unified_dispatcher is original_dispatcher # shared cache untouched
|
||||
|
||||
|
||||
# --- Custom-LLM path (Issue #6030) ---------------------------------------------
|
||||
|
||||
class _Unavailable(Exception):
|
||||
status_code = 503
|
||||
|
||||
|
||||
class _StubLLM:
|
||||
def __init__(self, model, fail, calls):
|
||||
self.model, self._fail, self._calls = model, fail, calls
|
||||
|
||||
def _respond(self):
|
||||
self._calls.append(self.model)
|
||||
if self._fail:
|
||||
raise self._fail
|
||||
return f"answer from {self.model}"
|
||||
|
||||
def get_response(self, **kwargs):
|
||||
return self._respond()
|
||||
|
||||
async def get_response_async(self, **kwargs):
|
||||
return self._respond()
|
||||
|
||||
|
||||
class _ToolThenFailLLM(_StubLLM):
|
||||
"""Runs a tool, then fails, like a 503 on the follow-up model request."""
|
||||
|
||||
def get_response(self, execute_tool_fn=None, **kwargs):
|
||||
execute_tool_fn("send_email", {})
|
||||
return self._respond()
|
||||
|
||||
|
||||
def _custom_agent(primary_error, fallbacks):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
calls = []
|
||||
primary = _StubLLM("primary-model", primary_error, calls)
|
||||
agent = Agent(name="c", instructions="x", llm=primary, output="silent")
|
||||
agent.fallback_models = [_StubLLM(m, None, calls) for m in fallbacks]
|
||||
seen = []
|
||||
agent._hook_runner.registry.register_function(
|
||||
HookEvent.MODEL_FALLBACK, lambda inp: seen.append(inp) or HookResult.allow()
|
||||
)
|
||||
return agent, primary, calls, seen
|
||||
|
||||
|
||||
def test_custom_llm_sync_falls_back_on_outage():
|
||||
agent, primary, calls, seen = _custom_agent(_Unavailable("503 overloaded"), ["backup"])
|
||||
assert agent.chat("hello") == "answer from backup"
|
||||
assert calls == ["primary-model", "backup"]
|
||||
assert len(seen) == 1 and seen[0].fallback_index == 0
|
||||
assert agent.llm_instance is primary # shared state untouched
|
||||
|
||||
|
||||
def test_custom_llm_async_falls_back_on_outage():
|
||||
agent, primary, calls, seen = _custom_agent(_Unavailable("503 overloaded"), ["backup"])
|
||||
assert asyncio.run(agent.achat("hello")) == "answer from backup"
|
||||
assert calls == ["primary-model", "backup"]
|
||||
assert len(seen) == 1
|
||||
assert agent.llm_instance is primary
|
||||
|
||||
|
||||
def test_custom_llm_non_fallback_error_does_not_walk_chain():
|
||||
err = Exception("401 Unauthorized: invalid api key")
|
||||
err.status_code = 401
|
||||
agent, _, calls, seen = _custom_agent(err, ["backup"])
|
||||
assert agent.chat("hello") is None
|
||||
assert calls == ["primary-model"]
|
||||
assert seen == []
|
||||
|
||||
|
||||
def test_fallback_llm_for_drops_connection_settings_across_providers(monkeypatch):
|
||||
import praisonaiagents.llm.llm as llm_mod
|
||||
|
||||
built = []
|
||||
monkeypatch.setattr(llm_mod, "LLM", lambda **kw: built.append(kw) or kw)
|
||||
agent = _DummyAgent()
|
||||
agent.llm = "ollama/llama3"
|
||||
agent._llm_init_params = {"model": "ollama/llama3", "base_url": "http://h", "api_key": "k", "max_iter": 3}
|
||||
agent._fallback_llm_for("ollama/qwen")
|
||||
agent._fallback_llm_for("openai/gpt-4o-mini")
|
||||
assert built[0] == {"model": "ollama/qwen", "base_url": "http://h", "api_key": "k", "max_iter": 3}
|
||||
assert built[1] == {"model": "openai/gpt-4o-mini", "max_iter": 3}
|
||||
assert agent._llm_init_params["model"] == "ollama/llama3"
|
||||
|
||||
|
||||
def test_fallback_llm_for_keeps_connection_settings_for_bare_names(monkeypatch):
|
||||
import praisonaiagents.llm.llm as llm_mod
|
||||
|
||||
built = []
|
||||
monkeypatch.setattr(llm_mod, "LLM", lambda **kw: built.append(kw) or kw)
|
||||
agent = _DummyAgent()
|
||||
agent.llm = "model-a"
|
||||
agent._llm_init_params = {"model": "model-a", "base_url": "http://custom/v1", "api_key": "k"}
|
||||
agent._fallback_llm_for("model-b")
|
||||
assert built[0] == {"model": "model-b", "base_url": "http://custom/v1", "api_key": "k"}
|
||||
|
||||
|
||||
def test_custom_llm_does_not_fall_back_after_a_tool_ran():
|
||||
from praisonaiagents import Agent
|
||||
|
||||
calls, sent = [], []
|
||||
primary = _ToolThenFailLLM("primary-model", _Unavailable("503 overloaded"), calls)
|
||||
agent = Agent(name="c", instructions="x", llm=primary, output="silent")
|
||||
agent.fallback_models = [_StubLLM("backup", None, calls)]
|
||||
agent.execute_tool = lambda name, args, **kw: sent.append(name)
|
||||
assert agent.chat("hello") is None # existing error handling, no replay
|
||||
assert calls == ["primary-model"]
|
||||
assert sent == ["send_email"]
|
||||
@@ -0,0 +1,31 @@
|
||||
"""OpenAI image calls must only forward style to models that accept it (dall-e-3)."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import ImageAgent
|
||||
|
||||
|
||||
@pytest.mark.parametrize("llm,kwargs,expected", [
|
||||
("openai/dall-e-2", {}, None),
|
||||
("openai/gpt-image-1", {}, None),
|
||||
("openai/gpt-image-1", {"style": "vivid"}, None),
|
||||
("openai/dall-e-3", {}, "natural"),
|
||||
("openai/dall-e-3", {"style": "vivid"}, "vivid"),
|
||||
])
|
||||
def test_style_forwarding(llm, kwargs, expected):
|
||||
agent = ImageAgent(llm=llm, verbose=False)
|
||||
received = {}
|
||||
|
||||
def fake_image_generation(**call_kwargs):
|
||||
received.update(call_kwargs)
|
||||
return {"data": [{"url": "https://example.com/image.png"}]}
|
||||
|
||||
agent._litellm = fake_image_generation
|
||||
try:
|
||||
agent.generate("A small green square icon", **kwargs)
|
||||
if expected is None:
|
||||
assert "style" not in received
|
||||
else:
|
||||
assert received["style"] == expected
|
||||
finally:
|
||||
agent.close()
|
||||
@@ -5,11 +5,13 @@ import pytest
|
||||
from praisonaiagents.tools import github_tools
|
||||
|
||||
|
||||
def _fake_git_factory(responses, calls):
|
||||
def _fake_git_factory(responses, calls, cwds=None):
|
||||
"""Build a fake subprocess.run that records calls and serves canned output."""
|
||||
cwds = cwds if cwds is not None else []
|
||||
|
||||
def fake_run(cmd, check=False, capture_output=False, text=False):
|
||||
def fake_run(cmd, check=False, capture_output=False, text=False, cwd=None):
|
||||
calls.append(cmd)
|
||||
cwds.append(cwd)
|
||||
assert cmd[0] == "git"
|
||||
key = tuple(cmd[1:])
|
||||
# Match on a prefix of the git args
|
||||
@@ -256,3 +258,50 @@ def test_create_branch_accepts_slash_names(monkeypatch):
|
||||
checkouts = [c for c in calls if c[:2] == ["git", "checkout"]]
|
||||
assert checkouts, "expected a checkout for a valid branch name"
|
||||
assert ["git", "checkout", "-B", "feature/new-thing"] in checkouts
|
||||
|
||||
|
||||
def test_repo_path_used_as_cwd_for_every_git_call(monkeypatch):
|
||||
responses = {
|
||||
("status", "--porcelain"): ("M file\n", 0),
|
||||
("rev-parse", "--abbrev-ref", "HEAD"): ("praisonai/work\n", 0),
|
||||
("symbolic-ref", "refs/remotes/origin/HEAD"): ("refs/remotes/origin/main\n", 0),
|
||||
}
|
||||
calls, cwds = [], []
|
||||
monkeypatch.setattr(subprocess, "run", _fake_git_factory(responses, calls, cwds))
|
||||
result = _run_tool(github_tools.github_commit_and_push, "fix", repo_path="/repo")
|
||||
assert "Successfully" in result
|
||||
assert calls and all(c == "/repo" for c in cwds)
|
||||
|
||||
|
||||
def test_never_stages_untracked_with_add_dot(monkeypatch):
|
||||
responses = {
|
||||
("status", "--porcelain"): ("M file\n", 0),
|
||||
("rev-parse", "--abbrev-ref", "HEAD"): ("praisonai/work\n", 0),
|
||||
("symbolic-ref", "refs/remotes/origin/HEAD"): ("refs/remotes/origin/main\n", 0),
|
||||
}
|
||||
calls = _call(monkeypatch, responses)
|
||||
_run_tool(github_tools.github_commit_and_push, "fix", paths=["new.py"])
|
||||
adds = [c for c in calls if c[:2] == ["git", "add"]]
|
||||
# scoped to cwd so a subdirectory call cannot stage sibling edits
|
||||
assert ["git", "add", "-u", "--", "."] in adds
|
||||
assert ["git", "add", "-u"] not in adds
|
||||
assert ["git", "add", "--", "new.py"] in adds
|
||||
assert ["git", "add", "."] not in adds
|
||||
# untracked files are not considered when checking for changes
|
||||
status = [c for c in calls if c[:2] == ["git", "status"]]
|
||||
assert all("--untracked-files=no" in c for c in status)
|
||||
|
||||
|
||||
def test_refusal_leaves_index_untouched(monkeypatch):
|
||||
responses = {
|
||||
("status", "--porcelain"): ("M file\n", 0),
|
||||
("rev-parse", "--abbrev-ref", "HEAD"): ("release/v2\n", 0),
|
||||
("symbolic-ref", "refs/remotes/origin/HEAD"): ("refs/remotes/origin/main\n", 0),
|
||||
("config", "user.email"): ("[email protected]\n", 0),
|
||||
("merge-base", "origin/main", "HEAD"): ("abc\n", 0),
|
||||
("log", "--format=%ae", "abc..HEAD"): ("[email protected]\n", 0),
|
||||
}
|
||||
calls = _call(monkeypatch, responses)
|
||||
result = _run_tool(github_tools.github_commit_and_push, "fix", branch="release/v2")
|
||||
assert "Error" in result
|
||||
assert not [c for c in calls if c[:2] == ["git", "add"]]
|
||||
@@ -0,0 +1,260 @@
|
||||
"""Unnamed agents in traces and graphs, agent_end status, and async model-call events.
|
||||
|
||||
Deterministic: the model call is replaced, so nothing reaches a provider.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.trace.context_events import (
|
||||
ContextListSink,
|
||||
ContextTraceEmitter,
|
||||
reset_context_emitter,
|
||||
set_context_emitter,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sink(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
s = ContextListSink()
|
||||
token = set_context_emitter(ContextTraceEmitter(sink=s, session_id="t", enabled=True))
|
||||
yield s
|
||||
reset_context_emitter(token)
|
||||
|
||||
|
||||
def _of(sink, kind):
|
||||
return [e for e in sink.get_events() if e.event_type.value == kind]
|
||||
|
||||
|
||||
def _fake_completion(text="Solar sails ride sunlight."):
|
||||
message = SimpleNamespace(content=text, tool_calls=None, role="assistant", reasoning_content=None)
|
||||
return SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=message, finish_reason="stop", index=0)],
|
||||
usage=SimpleNamespace(prompt_tokens=11, completion_tokens=7, total_tokens=18),
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
|
||||
|
||||
def search(query: str) -> str:
|
||||
"""Search the web."""
|
||||
return "result"
|
||||
|
||||
|
||||
# --- to_graph and ids ---------------------------------------------------------
|
||||
|
||||
def test_to_graph_accepts_an_unnamed_agent(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.workflows.diagram import to_graph
|
||||
|
||||
graph = to_graph(Agent(instructions="Answer briefly.", tools=[search]))
|
||||
ids = [n["id"] for n in graph["nodes"]]
|
||||
assert ids[:3] == ["start:begin", "agent:Agent", "tool:Agent/search"] or set(ids) >= {
|
||||
"start:begin", "agent:Agent", "tool:Agent/search", "end:done"}
|
||||
assert next(n for n in graph["nodes"] if n["id"] == "agent:Agent")["label"] == "Agent"
|
||||
|
||||
|
||||
def test_named_agent_graph_is_unchanged(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.workflows.diagram import to_graph
|
||||
|
||||
ids = [n["id"] for n in to_graph(Agent(name="Writer", instructions="x", tools=[search]))["nodes"]]
|
||||
assert "agent:Writer" in ids and "tool:Writer/search" in ids
|
||||
|
||||
|
||||
def test_emitter_ids_for_an_unnamed_agent():
|
||||
s = ContextListSink()
|
||||
em = ContextTraceEmitter(sink=s, session_id="t", enabled=True)
|
||||
em.agent_start(None)
|
||||
em.tool_call_start(None, "search", {"query": "x"})
|
||||
em.tool_call_end(None, "search", "ok", 1.0)
|
||||
em.agent_end(None)
|
||||
assert [e.node_id for e in s.get_events()] == [
|
||||
"agent:Agent", "tool:Agent/search", "tool:Agent/search", "agent:Agent"]
|
||||
|
||||
|
||||
def test_unnamed_agent_run_matches_its_graph(sink):
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.workflows.diagram import to_graph
|
||||
|
||||
agent = Agent(instructions="Answer briefly.")
|
||||
with patch.object(agent, "_chat_completion", return_value=_fake_completion()):
|
||||
agent.chat("Why do solar sails work?")
|
||||
graph_ids = {n["id"] for n in to_graph(agent)["nodes"]}
|
||||
for kind in ("agent_start", "agent_end"):
|
||||
(event,) = _of(sink, kind)
|
||||
assert event.agent_name is None # the name is not invented on the event
|
||||
assert event.node_id == "agent:Agent" and event.node_id in graph_ids
|
||||
|
||||
|
||||
# --- agent_end status ---------------------------------------------------------
|
||||
|
||||
def test_agent_end_records_completed(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(agent, "_chat_completion", return_value=_fake_completion()):
|
||||
assert agent.chat("hi")
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "completed"}
|
||||
|
||||
|
||||
def test_agent_end_records_failed(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(agent, "_chat_completion", side_effect=RuntimeError("provider down")):
|
||||
try:
|
||||
result = agent.chat("hi")
|
||||
except Exception:
|
||||
result = None
|
||||
assert result is None
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "failed"}
|
||||
|
||||
|
||||
# --- async model calls --------------------------------------------------------
|
||||
|
||||
def test_astart_records_the_model_call(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x", llm="gpt-4o-mini")
|
||||
|
||||
async def fake_dispatch(*args, **kwargs):
|
||||
return _fake_completion()
|
||||
|
||||
with patch.object(agent, "_achat_completion_with_retry_dispatch", side_effect=fake_dispatch):
|
||||
answer = asyncio.run(agent.astart("Why do solar sails work?"))
|
||||
assert "sunlight" in str(answer)
|
||||
(req,) = _of(sink, "llm_request")
|
||||
(resp,) = _of(sink, "llm_response")
|
||||
assert req.node_id == resp.node_id == "agent:Writer"
|
||||
assert resp.prompt_tokens == 11 and resp.completion_tokens == 7 and resp.cost_usd > 0
|
||||
assert "sunlight" in (resp.data.get("response_content") or "")
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "completed"}
|
||||
|
||||
|
||||
def test_async_model_error_is_recorded(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x", llm="gpt-4o-mini")
|
||||
|
||||
async def boom(*args, **kwargs):
|
||||
raise RuntimeError("rate limited")
|
||||
|
||||
with patch.object(agent, "_achat_completion_with_retry_dispatch", side_effect=boom):
|
||||
try:
|
||||
asyncio.run(agent.astart("hi"))
|
||||
except Exception:
|
||||
pass
|
||||
responses = _of(sink, "llm_response")
|
||||
assert responses and responses[0].data.get("finish_reason") == "error"
|
||||
assert "rate limited" in (responses[0].data.get("response_content") or "")
|
||||
|
||||
|
||||
def test_async_path_records_nothing_when_tracing_is_off(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
from praisonaiagents import Agent
|
||||
|
||||
s = ContextListSink()
|
||||
token = set_context_emitter(ContextTraceEmitter(sink=s, session_id="t", enabled=False))
|
||||
try:
|
||||
agent = Agent(name="Writer", instructions="x", llm="gpt-4o-mini")
|
||||
|
||||
async def fake_dispatch(*args, **kwargs):
|
||||
return _fake_completion()
|
||||
|
||||
with patch.object(agent, "_achat_completion_with_retry_dispatch", side_effect=fake_dispatch) as d:
|
||||
asyncio.run(agent.astart("hi"))
|
||||
assert d.called and s.get_events() == []
|
||||
finally:
|
||||
reset_context_emitter(token)
|
||||
|
||||
|
||||
# --- streamed turns (start(stream=True)) ---------------------------------------
|
||||
|
||||
def _fake_stream(self, prompt, **kwargs):
|
||||
yield "Solar "
|
||||
yield "sails."
|
||||
|
||||
|
||||
def test_streamed_turn_records_agent_start_and_end(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_turn", _fake_stream):
|
||||
text = "".join(agent._start_stream("hi"))
|
||||
assert text == "Solar sails."
|
||||
kinds = [(e.event_type.value, e.node_id, e.data.get("status")) for e in sink.get_events()]
|
||||
assert kinds == [("agent_start", "agent:Agent", None), ("agent_end", "agent:Agent", "completed")]
|
||||
|
||||
|
||||
def test_streamed_turn_stopped_early_is_cancelled(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_turn", _fake_stream):
|
||||
gen = agent._start_stream("hi")
|
||||
next(gen)
|
||||
gen.close() # the page stopped reading
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "cancelled"}
|
||||
|
||||
|
||||
def test_streamed_turn_that_raises_is_failed(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def boom(self, prompt, **kwargs):
|
||||
yield "x"
|
||||
raise RuntimeError("provider down")
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_turn", boom):
|
||||
with pytest.raises(RuntimeError):
|
||||
list(agent._start_stream("hi"))
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "failed"}
|
||||
|
||||
|
||||
def test_streamed_turn_records_nothing_when_tracing_is_off(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
from praisonaiagents import Agent
|
||||
|
||||
s = ContextListSink()
|
||||
token = set_context_emitter(ContextTraceEmitter(sink=s, session_id="t", enabled=False))
|
||||
try:
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_turn", _fake_stream):
|
||||
assert "".join(agent._start_stream("hi")) == "Solar sails."
|
||||
assert s.get_events() == []
|
||||
finally:
|
||||
reset_context_emitter(token)
|
||||
|
||||
|
||||
def test_streamed_turn_blocked_by_guardrail_is_failed(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(agent, "_validate_input_with_guardrail", return_value=(False, None, "no")):
|
||||
assert "blocked" in "".join(agent._start_stream("hi"))
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "failed"}
|
||||
|
||||
|
||||
def test_streamed_turn_fallback_records_the_turn_once(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def falls_back(self, prompt, **kwargs):
|
||||
yield self._stream_fallback_chat(prompt, kwargs, RuntimeError("stream broke"))
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_impl", falls_back), \
|
||||
patch.object(agent, "_chat_completion", return_value=_fake_completion()):
|
||||
assert "".join(agent._start_stream("hi")) == "Solar sails ride sunlight."
|
||||
kinds = [e.event_type.value for e in sink.get_events() if e.event_type.value.startswith("agent_")]
|
||||
assert kinds == ["agent_start", "agent_end"]
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "completed"}
|
||||
# chat() outside a stream still records its own turn.
|
||||
with patch.object(agent, "_chat_completion", return_value=_fake_completion()):
|
||||
agent.chat("again")
|
||||
assert len(_of(sink, "agent_start")) == 2
|
||||
Generated
+3
-3
@@ -3162,7 +3162,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -3292,8 +3292,8 @@ requires-dist = [
|
||||
{ name = "fastapi", marker = "extra == 'api'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'mcp'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'os'", specifier = ">=0.115.0" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "markitdown", extras = ["all"], marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
|
||||
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.20.0,<2" },
|
||||
{ name = "mcp", marker = "extra == 'parallel-search'", specifier = ">=1.30.0,<2" },
|
||||
|
||||
@@ -14,6 +14,7 @@ _LAZY_EXPORTS = {
|
||||
"praisonai_bot.gateway.admin",
|
||||
"provision_gateway_config",
|
||||
),
|
||||
"edit_gateway_channel": ("praisonai_bot.gateway.admin", "edit_gateway_channel"),
|
||||
}
|
||||
|
||||
|
||||
@@ -34,5 +35,6 @@ __all__ = [
|
||||
"GatewayRepairResult",
|
||||
"repair_gateway_config",
|
||||
"provision_gateway_config",
|
||||
"edit_gateway_channel",
|
||||
"__version__",
|
||||
]
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.15"
|
||||
__version__ = "0.1.18"
|
||||
@@ -468,10 +468,15 @@ async def deliver_media_to_adapter(
|
||||
)
|
||||
return False
|
||||
|
||||
target_id = thread_id or channel_id
|
||||
channel = client.get_channel(int(target_id))
|
||||
if channel is None and thread_id is not None:
|
||||
channel = client.get_channel(int(channel_id))
|
||||
# Resolve the thread *under* the parent channel (as ``send_message``
|
||||
# does) rather than by global id, so a target's thread segment can
|
||||
# never redirect the upload to an unrelated channel (#6084).
|
||||
channel = client.get_channel(int(channel_id))
|
||||
if channel is not None and thread_id is not None:
|
||||
get_thread = getattr(channel, "get_thread", None)
|
||||
thread = get_thread(int(thread_id)) if callable(get_thread) else None
|
||||
if thread is not None:
|
||||
channel = thread
|
||||
if channel is not None:
|
||||
await channel.send(
|
||||
content=caption or None, file=discord.File(path)
|
||||
|
||||
@@ -102,6 +102,27 @@ class BotOutboundMessenger:
|
||||
|
||||
resolved = f"{platform}:{channel_id}"
|
||||
|
||||
# Destination authorization (#6084): the target is model-controlled, so
|
||||
# an explicit "<platform>:<chat_id>" must name a chat this gateway
|
||||
# already knows (the turn's origin, or a home/alias/observed channel)
|
||||
# before anything — text or media — is delivered. Symbolic tokens
|
||||
# ("origin", a platform home, an alias) resolve to known chats by
|
||||
# construction. Operators widen the set via the alias overlay.
|
||||
if not self._is_attested(platform, channel_id):
|
||||
logger.warning(
|
||||
"BotOutboundMessenger: refusing send to unattested destination %s",
|
||||
resolved,
|
||||
)
|
||||
return DeliveryResult(
|
||||
ok=False,
|
||||
target=resolved,
|
||||
summary=f"Refused: {resolved} is not a known destination for this gateway.",
|
||||
detail=(
|
||||
"destination not attested (not the origin chat, a home "
|
||||
"channel, an alias, or an observed channel)"
|
||||
),
|
||||
)
|
||||
|
||||
# Idempotency covers the *whole* text+media envelope (issue #2578). A
|
||||
# duplicate proactive send must skip both the text and any media upload,
|
||||
# otherwise a re-fired job with attachments would double-upload media
|
||||
@@ -183,6 +204,21 @@ class BotOutboundMessenger:
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
def _is_attested(self, platform: str, channel_id: str) -> bool:
|
||||
"""Whether ``platform:channel_id`` is a chat this gateway knows."""
|
||||
origin = self._origin
|
||||
if (
|
||||
origin is not None
|
||||
and (origin.platform or "").lower() == platform.lower()
|
||||
and str(origin.channel_id) == str(channel_id)
|
||||
):
|
||||
return True
|
||||
try:
|
||||
return bool(self._router.directory.has_channel(platform, str(channel_id)))
|
||||
except Exception as e: # pragma: no cover — defensive, fail closed
|
||||
logger.warning("BotOutboundMessenger: attestation check failed: %s", e)
|
||||
return False
|
||||
|
||||
async def _deliver_media(
|
||||
self,
|
||||
target: str,
|
||||
|
||||
@@ -15,12 +15,26 @@ from dataclasses import dataclass, replace
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional, Protocol, TYPE_CHECKING
|
||||
|
||||
from ._chunk import _calculate_length, chunk_message, enforce_hard_cap
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents.streaming.events import StreamEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _fit_prefix(text: str, limit: int, unit: str) -> str:
|
||||
"""Longest prefix of *text* within *limit* in *unit*; single O(limit) pass."""
|
||||
if unit != "utf16":
|
||||
return text[:limit]
|
||||
used = 0
|
||||
for i, ch in enumerate(text):
|
||||
used += 2 if ord(ch) > 0xFFFF else 1
|
||||
if used > limit:
|
||||
return text[:i]
|
||||
return text
|
||||
|
||||
|
||||
# Matches <think>...</think> and <reasoning>...</reasoning> spans (case-insensitive,
|
||||
# spanning newlines). Also strips a trailing unclosed opening tag so partial
|
||||
# reasoning never leaks while a block is still streaming.
|
||||
@@ -375,7 +389,10 @@ class DraftStreamer:
|
||||
self._can_edit = caps.get('live_edit', True) # Assume true for backward compat
|
||||
self._text_limit = caps.get('text_limit', 0) or 0 # 0 = unlimited
|
||||
self._edit_rate_limit = caps.get('edit_rate_limit', 0) or 0
|
||||
|
||||
# Telegram measures its cap in UTF-16 units; default to codepoints.
|
||||
unit = getattr(getattr(adapter, 'platform_capabilities', None), 'length_unit', None)
|
||||
self._length_unit = unit if unit in ("codepoints", "utf16") else "codepoints"
|
||||
|
||||
# Per-stream runtime interval. Adaptive backoff mutates this, NOT the
|
||||
# shared config, so a flood in one stream never leaks into others.
|
||||
self._current_min_interval = self._config.min_interval
|
||||
@@ -621,8 +638,8 @@ class DraftStreamer:
|
||||
content = redact_outbound_text(content, self._adapter)
|
||||
|
||||
# Apply text limit if configured
|
||||
if self._text_limit > 0 and len(content) > self._text_limit:
|
||||
content = content[:self._text_limit - 3] + "..."
|
||||
if self._text_limit > 0 and _calculate_length(content, self._length_unit) > self._text_limit:
|
||||
content = _fit_prefix(content, max(1, self._text_limit - 3), self._length_unit) + "..."
|
||||
|
||||
return content
|
||||
|
||||
@@ -715,10 +732,20 @@ class DraftStreamer:
|
||||
# still gets the safe-by-default guarantee. The adapter carries any
|
||||
# injected redactor / opt-out so one policy governs every outbound seam.
|
||||
final_content = redact_outbound_text(final_content, self._adapter)
|
||||
|
||||
|
||||
# Over-cap answers cannot be delivered as one edit/send (both fail with
|
||||
# "message too long" and the reply is lost). Edit the placeholder with
|
||||
# the first chunk and send the rest as follow-ons (Issue #6078).
|
||||
chunks = [final_content]
|
||||
if self._text_limit > 0 and _calculate_length(final_content, self._length_unit) > self._text_limit:
|
||||
chunks = enforce_hard_cap(
|
||||
chunk_message(final_content, self._text_limit, length_unit=self._length_unit),
|
||||
self._text_limit, length_unit=self._length_unit,
|
||||
)
|
||||
|
||||
if self._rate_limiter:
|
||||
await self._rate_limiter.acquire(self._channel_id)
|
||||
|
||||
|
||||
# Primary delivery: when progressive edits are still healthy, edit the
|
||||
# existing placeholder in place. When they were disabled by flooding, a
|
||||
# fresh edit is likely to fail too, so try the placeholder edit first
|
||||
@@ -726,23 +753,35 @@ class DraftStreamer:
|
||||
# new send below.
|
||||
delivered = False
|
||||
try:
|
||||
await self._adapter.edit_message(self._channel_id, self._message_id, final_content)
|
||||
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
|
||||
# receives the completed answer regardless of the progressive path taken.
|
||||
if not delivered:
|
||||
try:
|
||||
await self._adapter.send_message(self._channel_id, final_content)
|
||||
await self._adapter.send_message(self._channel_id, chunks[0])
|
||||
delivered = True
|
||||
except Exception as send_exc: # noqa: BLE001 - last-resort delivery fallback
|
||||
logger.warning(
|
||||
"DraftStreamer final send fallback also failed: %s", send_exc
|
||||
)
|
||||
|
||||
|
||||
# Follow-on chunks: adapters' send_message already retries and parks
|
||||
# exhausted sends in the DLQ, so do NOT wrap it again (a second retry
|
||||
# layer would re-send a chunk that is already parked -> duplicates).
|
||||
for chunk in chunks[1:]:
|
||||
if self._rate_limiter:
|
||||
await self._rate_limiter.acquire(self._channel_id)
|
||||
try:
|
||||
await self._adapter.send_message(self._channel_id, chunk)
|
||||
except Exception as chunk_exc: # noqa: BLE001 - keep delivering remaining chunks
|
||||
delivered = False
|
||||
logger.warning("DraftStreamer follow-on chunk send failed: %s", chunk_exc)
|
||||
|
||||
logger.debug(
|
||||
"DraftStreamer finalized message %s, final_length=%d, progressive=%s, delivered=%s",
|
||||
self._message_id, len(final_content), self._progressive, delivered
|
||||
|
||||
@@ -551,27 +551,9 @@ class DiscordBot(OutboundResilienceMixin, ChatCommandMixin, MessageHookMixin):
|
||||
pass
|
||||
return
|
||||
final_content = send_result["content"]
|
||||
# A streamed answer longer than Discord's hard
|
||||
# cap cannot be delivered by editing the single
|
||||
# placeholder (both the edit and the send
|
||||
# fallback 400 on >2000 chars), which would drop
|
||||
# the whole reply. Fall back to the chunked
|
||||
# reply path so long answers are split and keep
|
||||
# the reply reference (Greptile #3 / #5).
|
||||
_cap = min(self.config.max_message_length, 2000)
|
||||
if len(final_content) > _cap:
|
||||
try:
|
||||
await self.delete_message(
|
||||
channel_id, placeholder_message_id
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
await self._send_long_message(
|
||||
message.channel, final_content,
|
||||
reference=message,
|
||||
)
|
||||
else:
|
||||
await streamer.finalize(final_content)
|
||||
# Over-cap answers are chunked by the shared
|
||||
# streamer (Issue #6078).
|
||||
await streamer.finalize(final_content)
|
||||
self.fire_message_sent(channel_id, final_content)
|
||||
except Exception:
|
||||
try:
|
||||
|
||||
@@ -1553,8 +1553,16 @@ def gateway_test(
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
@app.command("channels")
|
||||
channels_app = typer.Typer(
|
||||
help="List, probe and edit channels in gateway.yaml",
|
||||
invoke_without_command=True,
|
||||
)
|
||||
app.add_typer(channels_app, name="channels")
|
||||
|
||||
|
||||
@channels_app.callback()
|
||||
def gateway_channels(
|
||||
ctx: typer.Context,
|
||||
config: str = typer.Option("gateway.yaml", "--config", "-c", help="Path to gateway.yaml"),
|
||||
json_output: bool = typer.Option(False, "--json", help="Output JSON"),
|
||||
probe: bool = typer.Option(False, "--probe", help="Probe each channel's credentials"),
|
||||
@@ -1574,6 +1582,9 @@ def gateway_channels(
|
||||
import os
|
||||
import yaml
|
||||
|
||||
if ctx.invoked_subcommand is not None:
|
||||
return
|
||||
|
||||
if available:
|
||||
try:
|
||||
from praisonai_bot.bots._registry import list_platforms
|
||||
@@ -1662,6 +1673,125 @@ def gateway_channels(
|
||||
print(f"{name:<20} {platform:<12} {has_token:<12}")
|
||||
|
||||
|
||||
def _parse_channel_settings(pairs: Optional[List[str]]) -> dict:
|
||||
"""Parse repeated ``--set key=value`` flags; values are YAML-typed."""
|
||||
import yaml
|
||||
|
||||
settings = {}
|
||||
for pair in pairs or []:
|
||||
key, sep, value = pair.partition("=")
|
||||
if not sep or not key.strip():
|
||||
print(f"Error: --set expects key=value, got {pair!r}")
|
||||
raise typer.Exit(1)
|
||||
settings[key.strip()] = yaml.safe_load(value) if value else ""
|
||||
return settings
|
||||
|
||||
|
||||
def _edit_channel(
|
||||
ctx: typer.Context, config: Optional[str], name: str, settings=None, *, remove=False, must_exist=True
|
||||
) -> None:
|
||||
"""Run one validated channel edit, mapping failures to CLI exits (#6079)."""
|
||||
import os
|
||||
import yaml
|
||||
|
||||
from praisonai_bot.gateway.admin import _editable_channels, edit_gateway_channel
|
||||
|
||||
# ``gateway channels --config X <cmd>`` applies to the subcommand too.
|
||||
config = _resolve_doctor_config(config or ctx.parent.params.get("config") or "gateway.yaml")
|
||||
if not os.path.exists(config):
|
||||
print(f"Error: Config file not found: {config}")
|
||||
raise typer.Exit(1)
|
||||
with open(config) as fh:
|
||||
exists = name in _editable_channels(yaml.safe_load(fh) or {})
|
||||
if must_exist and not exists:
|
||||
print(f"Error: channel {name!r} is not configured in {config}")
|
||||
raise typer.Exit(1)
|
||||
if not must_exist and exists:
|
||||
print(f"Error: channel {name!r} already exists in {config}; use 'update'")
|
||||
raise typer.Exit(1)
|
||||
try:
|
||||
edit_gateway_channel(config, name, settings, remove=remove)
|
||||
except ValueError as exc:
|
||||
print(f"Error: {exc}\nNo changes written to {config}.")
|
||||
raise typer.Exit(1) from exc
|
||||
|
||||
|
||||
@channels_app.command("add")
|
||||
def gateway_channels_add(
|
||||
ctx: typer.Context,
|
||||
name: str = typer.Argument(..., help="Channel name, e.g. 'telegram'"),
|
||||
platform: Optional[str] = typer.Option(None, "--platform", help="Platform (default: the channel name)"),
|
||||
token: Optional[str] = typer.Option(
|
||||
None, "--token",
|
||||
help="Channel token (default: ${VAR} reference to the platform's credential env var)",
|
||||
),
|
||||
set_: Optional[List[str]] = typer.Option(None, "--set", help="Extra setting key=value (repeatable)"),
|
||||
config: Optional[str] = typer.Option(None, "--config", "-c", help="Path to gateway.yaml"),
|
||||
):
|
||||
"""Add one channel to gateway.yaml without touching the others.
|
||||
|
||||
Examples:
|
||||
praisonai gateway channels add telegram
|
||||
praisonai gateway channels add support --platform discord --set group_policy=mention_only
|
||||
"""
|
||||
from praisonai_bot.bots._config_schema import GatewayConfigSchema
|
||||
|
||||
platform = platform or name
|
||||
settings = {"platform": platform}
|
||||
settings["token"] = token if token is not None else GatewayConfigSchema._credential_env_ref(platform)
|
||||
settings.update(_parse_channel_settings(set_))
|
||||
_edit_channel(ctx, config, name, settings, must_exist=False)
|
||||
print(f"Added channel {name!r}.")
|
||||
|
||||
|
||||
@channels_app.command("update")
|
||||
def gateway_channels_update(
|
||||
ctx: typer.Context,
|
||||
name: str = typer.Argument(..., help="Channel name to update"),
|
||||
token: Optional[str] = typer.Option(None, "--token", help="New channel token"),
|
||||
set_: Optional[List[str]] = typer.Option(None, "--set", help="Setting key=value (repeatable)"),
|
||||
config: Optional[str] = typer.Option(None, "--config", "-c", help="Path to gateway.yaml"),
|
||||
):
|
||||
"""Update settings of one configured channel.
|
||||
|
||||
Examples:
|
||||
praisonai gateway channels update telegram --set group_policy=mention_only
|
||||
"""
|
||||
settings = _parse_channel_settings(set_)
|
||||
if token is not None:
|
||||
settings["token"] = token
|
||||
if not settings:
|
||||
print("Error: nothing to update; pass --set key=value, --token, or both")
|
||||
raise typer.Exit(1)
|
||||
_edit_channel(ctx, config, name, settings)
|
||||
print(f"Updated channel {name!r}.")
|
||||
|
||||
|
||||
@channels_app.command("disable")
|
||||
def gateway_channels_disable(
|
||||
ctx: typer.Context,
|
||||
name: str = typer.Argument(..., help="Channel name to disable"),
|
||||
config: Optional[str] = typer.Option(None, "--config", "-c", help="Path to gateway.yaml"),
|
||||
):
|
||||
"""Stop serving a channel but keep its config (sets ``enabled: false``).
|
||||
|
||||
Re-enable with ``praisonai gateway channels update <name> --set enabled=true``.
|
||||
"""
|
||||
_edit_channel(ctx, config, name, {"enabled": False})
|
||||
print(f"Disabled channel {name!r}.")
|
||||
|
||||
|
||||
@channels_app.command("remove")
|
||||
def gateway_channels_remove(
|
||||
ctx: typer.Context,
|
||||
name: str = typer.Argument(..., help="Channel name to remove"),
|
||||
config: Optional[str] = typer.Option(None, "--config", "-c", help="Path to gateway.yaml"),
|
||||
):
|
||||
"""Delete a channel block from gateway.yaml."""
|
||||
_edit_channel(ctx, config, name, remove=True)
|
||||
print(f"Removed channel {name!r}.")
|
||||
|
||||
|
||||
def _resolve_gateway_rest_url(
|
||||
url: Optional[str],
|
||||
host: Optional[str] = None,
|
||||
|
||||
@@ -15,6 +15,7 @@ one canonical flow (DRY).
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass, field
|
||||
@@ -201,7 +202,6 @@ def _repair_config_version(config_path: str):
|
||||
an interruption, full disk, or I/O error during ``doctor --fix`` can never
|
||||
leave the config truncated or half-written.
|
||||
"""
|
||||
import tempfile
|
||||
import yaml
|
||||
|
||||
with open(config_path) as fh:
|
||||
@@ -210,6 +210,14 @@ def _repair_config_version(config_path: str):
|
||||
from praisonaiagents.gateway.config import migrate_config_with_doctor
|
||||
|
||||
migrated, reasons = migrate_config_with_doctor(cfg)
|
||||
_atomic_write_yaml(config_path, migrated)
|
||||
return reasons
|
||||
|
||||
|
||||
def _atomic_write_yaml(config_path: str, data: dict) -> None:
|
||||
"""Serialise ``data`` to a temp file and ``os.replace`` it over the config."""
|
||||
import tempfile
|
||||
import yaml
|
||||
|
||||
directory = os.path.dirname(os.path.abspath(config_path))
|
||||
fd, tmp_path = tempfile.mkstemp(
|
||||
@@ -217,9 +225,7 @@ def _repair_config_version(config_path: str):
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "w") as fh:
|
||||
yaml.safe_dump(
|
||||
migrated, fh, default_flow_style=False, sort_keys=False
|
||||
)
|
||||
yaml.safe_dump(data, fh, default_flow_style=False, sort_keys=False)
|
||||
fh.flush()
|
||||
os.fsync(fh.fileno())
|
||||
os.replace(tmp_path, config_path)
|
||||
@@ -230,7 +236,127 @@ def _repair_config_version(config_path: str):
|
||||
pass
|
||||
raise
|
||||
|
||||
return reasons
|
||||
|
||||
def edit_gateway_channel(
|
||||
config_path: PathLike,
|
||||
name: str,
|
||||
settings: Optional[Dict] = None,
|
||||
*,
|
||||
remove: bool = False,
|
||||
) -> dict:
|
||||
"""Add, update or remove one ``channels:`` entry in place (Issue #6079).
|
||||
|
||||
A surgical read-modify-write: ``settings`` is merged into
|
||||
``channels[name]`` (created if absent, ``platform`` defaulting to ``name``)
|
||||
or, with ``remove=True``, the block is deleted. Other channels and
|
||||
top-level keys are untouched. The edited document is validated against
|
||||
``GatewayConfigSchema`` before the atomic write, so an invalid mutation is
|
||||
rejected rather than persisted. Disable a channel with
|
||||
``settings={"enabled": False}``.
|
||||
|
||||
Returns:
|
||||
The resulting channel block (``{}`` when removed).
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If ``config_path`` does not exist.
|
||||
KeyError: If ``remove=True`` and the channel is not configured.
|
||||
ValueError: If the edited config fails validation.
|
||||
"""
|
||||
import yaml
|
||||
|
||||
config_path = os.fspath(config_path)
|
||||
if not os.path.isfile(config_path):
|
||||
raise FileNotFoundError(config_path)
|
||||
# Hold a sidecar lock across read-modify-write so concurrent edits of
|
||||
# different channels cannot overwrite each other.
|
||||
with _config_edit_lock(config_path):
|
||||
with open(config_path) as fh:
|
||||
cfg = yaml.safe_load(fh) or {}
|
||||
return _apply_channel_edit(config_path, cfg, name, settings, remove)
|
||||
|
||||
|
||||
def _apply_channel_edit(config_path, cfg, name, settings, remove) -> dict:
|
||||
from praisonai_bot.bots._config_schema import validate_gateway_config
|
||||
|
||||
channels = _editable_channels(cfg)
|
||||
if remove:
|
||||
if name not in channels:
|
||||
raise KeyError(f"channel {name!r} is not configured in {config_path}")
|
||||
channels.pop(name)
|
||||
block: dict = {}
|
||||
else:
|
||||
block = dict(channels.get(name) or {})
|
||||
block.setdefault("platform", name)
|
||||
block.update(settings or {})
|
||||
channels[name] = block
|
||||
cfg["channels"] = channels
|
||||
|
||||
# Check the document itself: the validator may otherwise satisfy its
|
||||
# "at least one channel" rule from env credentials, which are not saved.
|
||||
if not any(
|
||||
not isinstance(b, dict) or b.get("enabled", True) is not False
|
||||
for b in channels.values()
|
||||
):
|
||||
raise ValueError("at least one enabled channel must remain configured")
|
||||
# Validate an env-resolved copy (as the runtime does; unset ``${VAR}``
|
||||
# refs are left as-is) while the saved file keeps the references.
|
||||
validate_gateway_config(cfg)
|
||||
_atomic_write_yaml(config_path, cfg)
|
||||
return block
|
||||
|
||||
|
||||
def _editable_channels(cfg: dict) -> dict:
|
||||
"""Return ``cfg``'s ``channels`` map, moving legacy declarations into it.
|
||||
|
||||
Single-bot (top-level ``platform``) and BotOS (``platforms``) files are
|
||||
only honoured by the loader while ``channels`` is empty, so they are
|
||||
converted in place before an edit adds a ``channels:`` block. ``${VAR}``
|
||||
references are copied verbatim.
|
||||
"""
|
||||
channels = cfg.get("channels") or {}
|
||||
if channels:
|
||||
return channels
|
||||
platform = cfg.get("platform")
|
||||
if platform:
|
||||
# Mirrors the loader's single-bot migration, which wins over BotOS.
|
||||
from praisonai_bot.bots._config_schema import GatewayConfigSchema
|
||||
|
||||
# Resolve a ``${VAR}`` platform first, as the loader does.
|
||||
raw = str(cfg.pop("platform"))
|
||||
platform = str(_resolve_env_token(raw) or raw).lower().strip()
|
||||
block = {"platform": platform}
|
||||
token = cfg.pop("token", None)
|
||||
block["token"] = token or GatewayConfigSchema._credential_env_ref(platform)
|
||||
if platform == "slack" and cfg.get("app_token") is not None:
|
||||
block["app_token"] = cfg.pop("app_token")
|
||||
for key in ("unknown_user_policy", "owner_user_id"):
|
||||
if cfg.get(key) is not None:
|
||||
block[key] = cfg[key]
|
||||
channels[platform] = block
|
||||
elif cfg.get("platforms"):
|
||||
for pname, pcfg in cfg.pop("platforms").items():
|
||||
channels[pname] = {"platform": pname, **(pcfg or {})}
|
||||
return channels
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _config_edit_lock(config_path: str):
|
||||
"""Exclusive OS advisory lock on ``<config>.lock`` (no-op if unsupported)."""
|
||||
with open(config_path + ".lock", "a+") as fh:
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # pragma: no cover - Windows
|
||||
fcntl = None
|
||||
if fcntl is not None:
|
||||
fcntl.flock(fh.fileno(), fcntl.LOCK_EX)
|
||||
else: # pragma: no cover - Windows
|
||||
try:
|
||||
import msvcrt
|
||||
|
||||
msvcrt.locking(fh.fileno(), msvcrt.LK_LOCK, 1)
|
||||
except Exception:
|
||||
pass
|
||||
yield # released when ``fh`` closes
|
||||
|
||||
|
||||
def _repair_gateway_secret(dry_run: bool = False, config_path: str = ""):
|
||||
|
||||
@@ -382,12 +382,16 @@ class PermissionAllowlist:
|
||||
path: Optional[Union[str, Path]] = None,
|
||||
) -> None:
|
||||
self._lock = threading.Lock()
|
||||
# In-memory mirror of grant keys (agent_id, tool_name, arg_signature)
|
||||
self._allowed: Set[tuple] = set()
|
||||
# In-memory mirror: (agent_id, tool_name, arg_signature) -> created_at.
|
||||
# Stamped so the grant TTL holds on the hot path for the whole process
|
||||
# lifetime, not only when the durable store is re-read on restart.
|
||||
self._allowed: Dict[tuple, float] = {}
|
||||
self._store: Optional[ScopedAllowlistStore] = None
|
||||
self._ttl_seconds = _DEFAULT_GRANT_TTL_SECONDS
|
||||
if durable:
|
||||
try:
|
||||
self._store = store or ScopedAllowlistStore(path=path)
|
||||
self._ttl_seconds = self._store.ttl_seconds
|
||||
self._rehydrate()
|
||||
except Exception:
|
||||
# Fail-closed on the persistence path: if the store cannot be
|
||||
@@ -409,13 +413,21 @@ class PermissionAllowlist:
|
||||
grant["tool_name"],
|
||||
grant.get("arg_signature") or "",
|
||||
)
|
||||
self._allowed.add(key)
|
||||
self._allowed[key] = float(grant.get("created_at") or time.time())
|
||||
if self._allowed:
|
||||
logger.info(
|
||||
"Rehydrated %d durable allow-list grant(s) from %s",
|
||||
len(self._allowed), self._store.path,
|
||||
)
|
||||
|
||||
def _evict_expired_locked(self) -> None:
|
||||
"""Drop grants older than the TTL from the in-memory mirror."""
|
||||
if self._ttl_seconds <= 0:
|
||||
return
|
||||
cutoff = time.time() - self._ttl_seconds
|
||||
for key in [k for k, ts in self._allowed.items() if ts <= cutoff]:
|
||||
del self._allowed[key]
|
||||
|
||||
# ── Scoped API ───────────────────────────────────────────────────
|
||||
def add_scoped(
|
||||
self,
|
||||
@@ -428,7 +440,7 @@ class PermissionAllowlist:
|
||||
"""Add a scoped grant, persisting it when durable."""
|
||||
key = (agent_id, tool_name, arg_signature or "")
|
||||
with self._lock:
|
||||
self._allowed.add(key)
|
||||
self._allowed[key] = time.time()
|
||||
if self._store is not None:
|
||||
try:
|
||||
self._store.add(
|
||||
@@ -457,11 +469,17 @@ class PermissionAllowlist:
|
||||
(empty signature) covers any arguments.
|
||||
"""
|
||||
sig = arg_signature or ""
|
||||
cutoff = time.time() - self._ttl_seconds if self._ttl_seconds > 0 else None
|
||||
with self._lock:
|
||||
# Only the candidate keys are checked (O(1)); expired ones are
|
||||
# evicted lazily. Full sweeps happen in list()/list_scoped().
|
||||
for a in (agent_id, ANY_AGENT):
|
||||
if (a, tool_name, "") in self._allowed:
|
||||
return True
|
||||
if sig and (a, tool_name, sig) in self._allowed:
|
||||
for key in ((a, tool_name, ""), (a, tool_name, sig) if sig else None):
|
||||
if key is None or key not in self._allowed:
|
||||
continue
|
||||
if cutoff is not None and self._allowed[key] <= cutoff:
|
||||
del self._allowed[key]
|
||||
continue
|
||||
return True
|
||||
return False
|
||||
|
||||
@@ -478,12 +496,11 @@ class PermissionAllowlist:
|
||||
if arg_signature is None:
|
||||
for key in list(self._allowed):
|
||||
if key[0] == agent_id and key[1] == tool_name:
|
||||
self._allowed.discard(key)
|
||||
del self._allowed[key]
|
||||
removed = True
|
||||
else:
|
||||
key = (agent_id, tool_name, arg_signature)
|
||||
if key in self._allowed:
|
||||
self._allowed.discard(key)
|
||||
if self._allowed.pop(key, None) is not None:
|
||||
removed = True
|
||||
if self._store is not None:
|
||||
try:
|
||||
@@ -508,6 +525,7 @@ class PermissionAllowlist:
|
||||
except Exception:
|
||||
logger.exception("Failed to list persisted grants; using memory")
|
||||
with self._lock:
|
||||
self._evict_expired_locked()
|
||||
items = [
|
||||
{
|
||||
"agent_id": a,
|
||||
@@ -530,7 +548,7 @@ class PermissionAllowlist:
|
||||
with self._lock:
|
||||
for key in list(self._allowed):
|
||||
if key[1] == tool_name:
|
||||
self._allowed.discard(key)
|
||||
del self._allowed[key]
|
||||
removed = True
|
||||
if self._store is not None:
|
||||
try:
|
||||
@@ -551,6 +569,7 @@ class PermissionAllowlist:
|
||||
grants are exposed via :meth:`list_scoped` / the ``grants`` view.
|
||||
"""
|
||||
with self._lock:
|
||||
self._evict_expired_locked()
|
||||
return sorted(
|
||||
{key[1] for key in self._allowed if key[0] == ANY_AGENT}
|
||||
)
|
||||
|
||||
@@ -11,11 +11,13 @@ import logging
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
from .._guarded_subprocess import run_guarded
|
||||
from .._guarded_subprocess import _kill_process_group, run_guarded
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -86,6 +88,17 @@ class KanbanDispatcher:
|
||||
self._task_runs: Dict[str, int] = {} # task_id -> open run_id
|
||||
self.worker_id = f"gateway_{os.getpid()}"
|
||||
|
||||
# Serialises blocking git/SQLite work handed to worker threads, so two
|
||||
# integrations never run git concurrently in the same base checkout.
|
||||
_blocking_lock = threading.Lock()
|
||||
|
||||
async def _run_blocking(self, fn, *args):
|
||||
"""Run a blocking (git / SQLite) helper off the event loop."""
|
||||
def _locked():
|
||||
with self._blocking_lock:
|
||||
return fn(*args)
|
||||
return await asyncio.to_thread(_locked)
|
||||
|
||||
def _close_run_safe(self, store, run_id, outcome, *, summary=None, metadata=None, error=None) -> bool:
|
||||
"""Close a run, swallowing errors so dispatch stays resilient.
|
||||
|
||||
@@ -196,8 +209,13 @@ class KanbanDispatcher:
|
||||
if not store:
|
||||
return 0
|
||||
|
||||
# Clean up completed processes first
|
||||
self._cleanup_completed_tasks(store)
|
||||
# Clean up completed processes first. Integration runs several git
|
||||
# commands, so keep it off the event loop.
|
||||
await self._run_blocking(self._cleanup_completed_tasks, store)
|
||||
|
||||
# Prove liveness for our own children before the reclaim sweep runs,
|
||||
# so a healthy long-running worker is never mistaken for a stale one.
|
||||
self._heartbeat_children(store)
|
||||
|
||||
# Reclaim stale/crashed claims before dispatching new work so that
|
||||
# tasks stranded by dead workers return to 'ready' and get re-dispatched.
|
||||
@@ -221,6 +239,14 @@ class KanbanDispatcher:
|
||||
|
||||
spawned = 0
|
||||
for task in ready_tasks[:available_slots]:
|
||||
live = self.running_tasks.get(task.id)
|
||||
if live is not None and live.poll() is None:
|
||||
# Never overwrite the handle of a still-running worker.
|
||||
logger.error(
|
||||
"Task %s is 'ready' but its worker (pid %s) is still running; "
|
||||
"skipping re-dispatch", task.id, live.pid,
|
||||
)
|
||||
continue
|
||||
try:
|
||||
# Try to claim the task atomically with a lease + owner PID.
|
||||
claimed = self._claim_task(store, task)
|
||||
@@ -245,8 +271,24 @@ class KanbanDispatcher:
|
||||
)
|
||||
continue
|
||||
|
||||
# Spawn worker process
|
||||
success = await self._spawn_worker(task, store)
|
||||
# Spawn worker process. If the dispatcher is stopped while the
|
||||
# worktree is being prepared, no worker exists yet: close the
|
||||
# run and hand the claim back before propagating.
|
||||
try:
|
||||
success = await self._spawn_worker(task, store)
|
||||
except asyncio.CancelledError:
|
||||
self._close_run_safe(
|
||||
store, run_id, 'failed',
|
||||
error='Dispatcher stopped before worker spawn',
|
||||
)
|
||||
try:
|
||||
store.release_claim(task.id, self.worker_id)
|
||||
except Exception as release_err:
|
||||
logger.error(
|
||||
"Failed to release claim for %s on cancel: %s",
|
||||
task.id, release_err,
|
||||
)
|
||||
raise
|
||||
if not success:
|
||||
# Failed to spawn: close run as failed, count it, and only
|
||||
# release the claim for retry when the failure was recorded
|
||||
@@ -394,6 +436,69 @@ class KanbanDispatcher:
|
||||
(pid, datetime.now(timezone.utc).isoformat(), task_id, self.worker_id),
|
||||
)
|
||||
|
||||
def _heartbeat_children(self, store: Any):
|
||||
"""Renew lease + heartbeat for every child worker that is still alive.
|
||||
|
||||
The spawned worker cannot heartbeat for itself, so the supervisor that
|
||||
holds its ``Popen`` handle does it. If the claim is gone (an operator
|
||||
released or moved the card), the child is terminated rather than left
|
||||
running unowned.
|
||||
"""
|
||||
beat = getattr(store, 'heartbeat', None)
|
||||
if beat is None:
|
||||
return
|
||||
for task_id, process in list(self.running_tasks.items()):
|
||||
if process.poll() is not None:
|
||||
continue
|
||||
try:
|
||||
owned = beat(task_id, self.worker_id, ttl_seconds=self.claim_ttl_seconds)
|
||||
except Exception as e:
|
||||
logger.debug(f"Heartbeat for task {task_id} failed: {e}")
|
||||
continue
|
||||
if not owned:
|
||||
logger.warning(
|
||||
f"Lost claim on task {task_id} while worker {process.pid} "
|
||||
f"was running; terminating it"
|
||||
)
|
||||
self._terminate_unowned(store, task_id, process)
|
||||
|
||||
def _terminate_unowned(self, store: Any, task_id: str, process: subprocess.Popen):
|
||||
"""Stop a child whose claim is no longer ours and forget its tracking."""
|
||||
try:
|
||||
process.terminate()
|
||||
process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
except OSError as e:
|
||||
logger.error(f"Failed to terminate worker for task {task_id}: {e}")
|
||||
self.running_tasks.pop(task_id, None)
|
||||
entry = getattr(self, '_worktrees', {}).get(task_id)
|
||||
if entry:
|
||||
# Lossless teardown: keep tracking (and say why) if work remains.
|
||||
preserved = self._remove_worktree(*entry)
|
||||
if preserved is None:
|
||||
self._worktrees.pop(task_id, None)
|
||||
else:
|
||||
try:
|
||||
store.add_comment(
|
||||
task_id, self.worker_id,
|
||||
f"worktree_preserved at {entry[0]}: {preserved}",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
log_path = getattr(self, '_temp_logs', {}).pop(task_id, None)
|
||||
if log_path:
|
||||
try:
|
||||
os.unlink(log_path)
|
||||
except OSError:
|
||||
pass
|
||||
run_id = self._task_runs.get(task_id)
|
||||
if run_id is None or self._close_run_safe(
|
||||
store, run_id, 'failed',
|
||||
error='Claim lost while worker was running; worker terminated',
|
||||
):
|
||||
self._task_runs.pop(task_id, None)
|
||||
|
||||
def _reclaim_stale_claims(self, store: Any):
|
||||
"""Sweep running tasks and reclaim those stranded by dead/stale workers."""
|
||||
reclaim = getattr(store, 'reclaim_stale_claims', None)
|
||||
@@ -409,6 +514,11 @@ class KanbanDispatcher:
|
||||
|
||||
for task_id in reclaimed or []:
|
||||
logger.warning(f"Reclaimed stale kanban task {task_id} back to 'ready'")
|
||||
# The store closed the abandoned run as 'crashed' atomically with
|
||||
# the reclaim; count the attempt so the circuit breaker sees it.
|
||||
self._record_failure_safe(
|
||||
store, task_id, error='Worker died or went stale; claim reclaimed'
|
||||
)
|
||||
# The running -> ready transition is announced as
|
||||
# KANBAN_TASK_MOVED (reclaimed=True) by the store itself;
|
||||
# 'KANBAN_TASK_RECLAIMED' was never a HookEvent member and so
|
||||
@@ -584,6 +694,11 @@ class KanbanDispatcher:
|
||||
return f"worktree removal error: {e}"
|
||||
return None
|
||||
|
||||
def _discard_worktree(self, path: str, branch: str) -> None:
|
||||
"""Remove a worktree and branch that no worker ever used."""
|
||||
if self._remove_worktree(path, force=True) is None:
|
||||
self._run_git("branch", "-D", branch)
|
||||
|
||||
def _commit_worktree_changes(self, path: str, branch: str) -> None:
|
||||
"""Commit any uncommitted edits the worker left in its worktree.
|
||||
|
||||
@@ -724,7 +839,22 @@ class KanbanDispatcher:
|
||||
# cwd so nothing regresses.
|
||||
worktree_path = None
|
||||
if getattr(task, 'workspace_kind', 'default') == 'worktree':
|
||||
worktree_path = self._prepare_worktree(task, store)
|
||||
prep = asyncio.ensure_future(
|
||||
self._run_blocking(self._prepare_worktree, task, store)
|
||||
)
|
||||
try:
|
||||
worktree_path = await asyncio.shield(prep)
|
||||
except asyncio.CancelledError:
|
||||
# The thread keeps running after cancellation. Let it
|
||||
# finish, then drop the fresh (unused) worktree and branch
|
||||
# so a retry can recreate them with `worktree add -b`.
|
||||
path = await asyncio.shield(prep)
|
||||
if path:
|
||||
await self._run_blocking(
|
||||
self._discard_worktree, path,
|
||||
f"kanban/{self._safe_task_id(task.id)}",
|
||||
)
|
||||
raise
|
||||
if worktree_path:
|
||||
if not hasattr(self, '_worktrees'):
|
||||
self._worktrees = {}
|
||||
@@ -753,6 +883,9 @@ class KanbanDispatcher:
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
cwd=worktree_path,
|
||||
# Own process group so shutdown can stop the worker's
|
||||
# descendants too (POSIX; ignored on Windows).
|
||||
start_new_session=(os.name != 'nt'),
|
||||
)
|
||||
# parent FD closed here; child still has its duped copy
|
||||
|
||||
@@ -1005,30 +1138,27 @@ class KanbanDispatcher:
|
||||
while self.running_tasks and (time.time() - start_time) < timeout:
|
||||
store = self._get_kanban_store()
|
||||
if store:
|
||||
self._cleanup_completed_tasks(store)
|
||||
await self._run_blocking(self._cleanup_completed_tasks, store)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Force terminate remaining tasks
|
||||
# Force terminate remaining tasks: signal every worker group
|
||||
# first, then wait for them together off the event loop.
|
||||
store = self._get_kanban_store()
|
||||
for task_id, process in self.running_tasks.items():
|
||||
logger.warning(f"Force terminating task {task_id}")
|
||||
try:
|
||||
process.terminate()
|
||||
process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
_terminate_worker(process)
|
||||
await asyncio.gather(*(
|
||||
asyncio.to_thread(_wait_quiet, process, 5)
|
||||
for process in self.running_tasks.values()
|
||||
))
|
||||
for task_id, process in self.running_tasks.items():
|
||||
if process.poll() is None:
|
||||
logger.warning(f"Task {task_id} did not terminate in 5s; sending SIGKILL")
|
||||
try:
|
||||
process.kill()
|
||||
except OSError as kill_err:
|
||||
logger.error(f"Failed to kill task {task_id}: {kill_err}")
|
||||
except OSError as os_err:
|
||||
logger.error(f"OS error while terminating task {task_id}: {os_err}")
|
||||
try:
|
||||
process.kill()
|
||||
except OSError as kill_err:
|
||||
logger.error(f"Failed to kill task {task_id}: {kill_err}")
|
||||
# KeyboardInterrupt, SystemExit, CancelledError now propagate as intended
|
||||
# Even if the worker exited, a descendant may have ignored
|
||||
# SIGTERM, so always finish off the whole group.
|
||||
_kill_worker_group(process)
|
||||
|
||||
for task_id in list(self.running_tasks):
|
||||
# Close any open run for the interrupted attempt so we do not
|
||||
# leave dangling task_runs rows / stale current_run_id pointers.
|
||||
run_id = self._task_runs.pop(task_id, None)
|
||||
@@ -1044,6 +1174,49 @@ class KanbanDispatcher:
|
||||
logger.info("Kanban dispatcher shutdown complete")
|
||||
|
||||
|
||||
def _terminate_worker(process: subprocess.Popen) -> None:
|
||||
"""Ask a worker and its descendants to exit (SIGTERM to its group)."""
|
||||
try:
|
||||
if os.name != 'nt':
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
else:
|
||||
process.terminate()
|
||||
except OSError:
|
||||
try:
|
||||
process.terminate()
|
||||
except OSError as err:
|
||||
logger.error(f"Failed to terminate worker {process.pid}: {err}")
|
||||
|
||||
|
||||
def _kill_worker_group(process: subprocess.Popen) -> None:
|
||||
"""SIGKILL a worker's group, even after the worker itself was reaped.
|
||||
|
||||
The worker leads its own session, so its group id equals its pid and the
|
||||
group stays addressable while any descendant is alive.
|
||||
"""
|
||||
if os.name == 'nt':
|
||||
if process.poll() is None:
|
||||
_kill_process_group(process)
|
||||
return
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
pass # group already empty
|
||||
if process.poll() is None:
|
||||
try:
|
||||
process.kill()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _wait_quiet(process: subprocess.Popen, timeout: float) -> None:
|
||||
"""Wait for a process, ignoring a timeout (the caller escalates)."""
|
||||
try:
|
||||
process.wait(timeout=timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
pass
|
||||
|
||||
|
||||
# Global dispatcher management functions
|
||||
|
||||
async def start_kanban_dispatcher(max_concurrent: int = 3, poll_interval: float = 5.0):
|
||||
|
||||
@@ -7686,6 +7686,14 @@ class WebSocketGateway:
|
||||
return None
|
||||
|
||||
cfg = (self._loaded_config or {}).get("scheduler", {}) or {}
|
||||
if not isinstance(cfg, dict):
|
||||
# A malformed (e.g. hot-reloaded) ``scheduler:`` block must not
|
||||
# kill the tick loop — fall back to the fail-closed defaults.
|
||||
logger.warning(
|
||||
"Ignoring non-mapping 'scheduler:' config (%s); using defaults",
|
||||
type(cfg).__name__,
|
||||
)
|
||||
cfg = {}
|
||||
denied = cfg.get("denied_toolsets")
|
||||
denied_toolsets = (
|
||||
set(denied) if isinstance(denied, (list, tuple, set))
|
||||
@@ -7762,16 +7770,33 @@ class WebSocketGateway:
|
||||
# keeps a long-lived gateway from pinning stale unattended-run
|
||||
# settings. Falls back to ``run_loop`` when the executor predates
|
||||
# per-tick policy refresh.
|
||||
# A tick that raises (store I/O such as EMFILE on the lock file, a
|
||||
# bad config) is logged and retried with capped backoff instead of
|
||||
# ending the task — otherwise scheduling would silently stop for
|
||||
# good while chat channels keep the gateway looking healthy.
|
||||
failures = 0
|
||||
try:
|
||||
while True:
|
||||
executor._run_policy = self._build_run_policy()
|
||||
await executor.tick_all()
|
||||
await asyncio.sleep(interval)
|
||||
try:
|
||||
executor._run_policy = self._build_run_policy()
|
||||
await executor.tick_all()
|
||||
failures = 0
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
failures += 1
|
||||
logger.exception(
|
||||
"Scheduler tick failed (%d in a row)", failures,
|
||||
)
|
||||
await asyncio.sleep(
|
||||
min(interval * 2 ** min(failures, 10), 300.0) if failures
|
||||
else interval
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
self._scheduler_task = asyncio.create_task(_run())
|
||||
logger.info("Scheduler tick started (interval=15s)")
|
||||
logger.info("Scheduler tick started (interval=%ss)", interval)
|
||||
|
||||
async def _start_session_cleanup(self) -> None:
|
||||
"""Start periodic cleanup of expired sessions."""
|
||||
|
||||
@@ -1094,7 +1094,8 @@ class SQLiteKanbanStore:
|
||||
|
||||
with self._get_connection() as conn:
|
||||
cursor = conn.execute("""
|
||||
SELECT id, claim_lock, claim_expires, worker_pid, last_heartbeat_at, version
|
||||
SELECT id, claim_lock, claim_expires, worker_pid, last_heartbeat_at, version,
|
||||
current_run_id
|
||||
FROM tasks
|
||||
WHERE status = 'running' AND claim_lock IS NOT NULL AND claim_lock != ''
|
||||
""")
|
||||
@@ -1127,7 +1128,7 @@ class SQLiteKanbanStore:
|
||||
result = conn.execute("""
|
||||
UPDATE tasks
|
||||
SET claim_lock = NULL, claim_expires = NULL, worker_pid = NULL,
|
||||
last_heartbeat_at = NULL,
|
||||
last_heartbeat_at = NULL, current_run_id = NULL,
|
||||
updated_at = ?, status = 'ready', version = version + 1
|
||||
WHERE id = ? AND status = 'running' AND claim_lock = ?
|
||||
AND version = ?
|
||||
@@ -1136,6 +1137,20 @@ class SQLiteKanbanStore:
|
||||
if result.rowcount > 0:
|
||||
reclaimed.append(row['id'])
|
||||
reason = 'dead' if not worker_alive else 'stale_heartbeat'
|
||||
# Close the abandoned attempt in the same transaction, so a
|
||||
# dispatcher that claims the task next can never have its
|
||||
# own fresh run mistaken for this one.
|
||||
if row.get('current_run_id') is not None and conn.execute(
|
||||
"UPDATE task_runs SET outcome = ?, error = ?, ended_at = ? "
|
||||
"WHERE id = ? AND ended_at IS NULL",
|
||||
(RunOutcome.CRASHED.value,
|
||||
f'Worker {reason}; claim reclaimed',
|
||||
now.isoformat(), row['current_run_id']),
|
||||
).rowcount:
|
||||
self._log_event(conn, row['id'], 'run_closed', {
|
||||
'run_id': row['current_run_id'],
|
||||
'outcome': RunOutcome.CRASHED.value,
|
||||
})
|
||||
reclaim_details[row['id']] = {
|
||||
'worker_id': row['claim_lock'],
|
||||
'worker_pid': pid,
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-bot"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
description = "Bots, gateway, and channel CLI for PraisonAI — messaging platforms and WebSocket control plane extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
||||
@@ -417,3 +417,50 @@ def test_discord_media_transient_error_propagates_for_retry(tmp_path):
|
||||
ok = asyncio.run(_run())
|
||||
assert ok is True
|
||||
assert channel.calls == 2
|
||||
|
||||
|
||||
def test_discord_media_thread_is_scoped_to_parent_channel(tmp_path):
|
||||
# #6084: the thread segment of "discord:A:B" must not redirect the upload
|
||||
# to an arbitrary channel B by global id — only a thread under the
|
||||
# (attested) parent A may receive it, otherwise the file lands in A.
|
||||
pytest.importorskip("discord")
|
||||
|
||||
f = tmp_path / "report.pdf"
|
||||
f.write_bytes(b"%PDF-1.4 data")
|
||||
|
||||
class _Channel:
|
||||
def __init__(self, threads=None):
|
||||
self.sent = 0
|
||||
self._threads = threads or {}
|
||||
|
||||
async def send(self, *args, **kwargs):
|
||||
self.sent += 1
|
||||
|
||||
def get_thread(self, tid):
|
||||
return self._threads.get(tid)
|
||||
|
||||
real_thread = _Channel()
|
||||
parent = _Channel(threads={555: real_thread})
|
||||
unrelated = _Channel()
|
||||
channels = {42: parent, 999: unrelated}
|
||||
|
||||
class Adapter:
|
||||
platform = "discord"
|
||||
|
||||
class _client:
|
||||
@staticmethod
|
||||
def get_channel(cid):
|
||||
return channels.get(cid)
|
||||
|
||||
def _run(thread_id):
|
||||
return asyncio.run(
|
||||
deliver_media_to_adapter(Adapter(), "42", str(f), thread_id=thread_id)
|
||||
)
|
||||
|
||||
assert _run("999") is True
|
||||
assert unrelated.sent == 0
|
||||
assert parent.sent == 1
|
||||
|
||||
assert _run("555") is True
|
||||
assert real_thread.sent == 1
|
||||
assert parent.sent == 1
|
||||
@@ -88,6 +88,34 @@ def test_send_to_platform_home_channel():
|
||||
assert sent == [("999", "Nightly summary")]
|
||||
|
||||
|
||||
def test_send_to_unattested_explicit_chat_is_refused():
|
||||
# #6084: a model-supplied "<platform>:<chat_id>" naming a chat the gateway
|
||||
# has no relationship with must fail closed instead of delivering.
|
||||
router, sent = _make_router(home="999")
|
||||
origin = SessionSource(platform="telegram", channel_id="123")
|
||||
messenger = BotOutboundMessenger(router, origin=origin)
|
||||
|
||||
result = asyncio.run(messenger.send("telegram:666", "leak"))
|
||||
|
||||
assert result.ok is False
|
||||
assert "Refused" in result.summary
|
||||
assert sent == []
|
||||
|
||||
|
||||
def test_send_to_known_explicit_chat_is_allowed():
|
||||
router, sent = _make_router(home="999", aliases={"ops": "555"})
|
||||
origin = SessionSource(platform="telegram", channel_id="123")
|
||||
messenger = BotOutboundMessenger(router, origin=origin)
|
||||
|
||||
for target, chat in (
|
||||
("telegram:123", "123"), # origin chat
|
||||
("telegram:999", "999"), # home channel
|
||||
("telegram:555:42", "555"), # aliased chat, thread under it
|
||||
):
|
||||
assert asyncio.run(messenger.send(target, "hi")).ok is True, target
|
||||
assert [c for c, _ in sent] == ["123", "999", "555"]
|
||||
|
||||
|
||||
def test_send_unresolvable_target_fails_cleanly():
|
||||
router, sent = _make_router()
|
||||
messenger = BotOutboundMessenger(router) # no origin
|
||||
|
||||
@@ -141,53 +141,100 @@ async def test_slack_stream_uploads_media(monkeypatch, tmp_path):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discord_stream_chunks_long_reply(monkeypatch):
|
||||
"""A streamed answer over Discord's cap is delivered via the chunked path.
|
||||
async def test_discord_stream_chunks_long_reply():
|
||||
"""A streamed answer over Discord's cap is chunked by the shared streamer.
|
||||
|
||||
Guards Greptile #3/#5: editing the single placeholder with >2000 chars
|
||||
would 400 and drop the whole reply, so the finalize path must fall back to
|
||||
the reference-preserving ``_send_long_message`` chunker.
|
||||
Editing the single placeholder with >2000 chars would 400 and drop the
|
||||
whole reply, so ``finalize`` edits the placeholder with the first chunk and
|
||||
sends the remainder as follow-ons (Issue #6078).
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
bot = DiscordBot(token="t", config=BotConfig(token="t", streaming=True))
|
||||
|
||||
long_reply = "x" * 5000
|
||||
bot._session.chat = AsyncMock(return_value=long_reply)
|
||||
bot.send_message = AsyncMock(return_value={"message_id": "m1"})
|
||||
bot.edit_message = AsyncMock(return_value=None)
|
||||
bot.delete_message = AsyncMock(return_value=True)
|
||||
|
||||
sent_chunks = []
|
||||
|
||||
long_sender = AsyncMock(
|
||||
side_effect=lambda channel, text, reference=None: sent_chunks.append(text)
|
||||
)
|
||||
monkeypatch.setattr(bot, "_send_long_message", long_sender)
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from praisonai_bot.bots._streaming import DraftStreamer
|
||||
|
||||
bot = DiscordBot(token="t", config=BotConfig(token="t", streaming=True))
|
||||
bot.send_message = AsyncMock(return_value={"message_id": "m1"})
|
||||
bot.edit_message = AsyncMock(return_value=None)
|
||||
|
||||
streamer = DraftStreamer(
|
||||
adapter=bot,
|
||||
channel_id="C1",
|
||||
config=bot._streaming_config,
|
||||
platform="discord",
|
||||
adapter=bot, channel_id="C1", config=bot._streaming_config, platform="discord",
|
||||
)
|
||||
await streamer.start()
|
||||
|
||||
# Simulate the finalize decision the adapter makes for an over-cap reply.
|
||||
send_result = bot.fire_message_sending("C1", long_reply)
|
||||
final_content = send_result["content"]
|
||||
cap = min(bot.config.max_message_length, 2000)
|
||||
assert len(final_content) > cap
|
||||
long_reply = "\n\n".join(f"Paragraph {i}. " + "x" * 900 for i in range(6))
|
||||
await streamer.finalize(long_reply)
|
||||
|
||||
message = MagicMock()
|
||||
message.channel = MagicMock()
|
||||
if len(final_content) > cap:
|
||||
await bot.delete_message("C1", "m1")
|
||||
await bot._send_long_message(message.channel, final_content, reference=message)
|
||||
else: # pragma: no cover - defensive
|
||||
await streamer.finalize(final_content)
|
||||
edited = bot.edit_message.await_args.args[2]
|
||||
follow_ons = [c.args[1] for c in bot.send_message.await_args_list[1:]]
|
||||
assert edited.startswith("Paragraph 0.")
|
||||
assert follow_ons, "remainder should be sent as follow-on messages"
|
||||
assert all(len(c) <= 2000 for c in [edited, *follow_ons])
|
||||
assert "".join([edited, *follow_ons]).replace("\n", "") == long_reply.replace("\n", "")
|
||||
|
||||
assert sent_chunks, "long reply should have been routed through the chunker"
|
||||
bot.delete_message.assert_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamer_chunks_by_utf16_and_falls_back_per_chunk():
|
||||
"""Telegram-style UTF-16 cap: over-cap reply is chunked, never dropped.
|
||||
|
||||
When the placeholder edit fails, the first chunk falls back to a fresh
|
||||
send, and every chunk stays within the UTF-16 limit (Issue #6078).
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from praisonai_bot.bots._chunk import _calculate_length
|
||||
from praisonai_bot.bots._streaming import DraftStreamer, StreamingConfig
|
||||
|
||||
adapter = SimpleNamespace(
|
||||
capabilities={"live_edit": True, "text_limit": 100},
|
||||
platform_capabilities=SimpleNamespace(length_unit="utf16"),
|
||||
send_message=AsyncMock(return_value={"message_id": "m1"}),
|
||||
edit_message=AsyncMock(side_effect=RuntimeError("message is too long")),
|
||||
)
|
||||
streamer = DraftStreamer(
|
||||
adapter=adapter, channel_id="C1",
|
||||
config=StreamingConfig(mode=StreamingMode.DRAFT), platform="telegram",
|
||||
)
|
||||
await streamer.start()
|
||||
|
||||
reply = "\U0001F600" * 120 # 240 UTF-16 units, 120 codepoints
|
||||
await streamer.finalize(reply)
|
||||
|
||||
sent = [c.args[1] for c in adapter.send_message.await_args_list[1:]]
|
||||
assert len(sent) >= 3
|
||||
assert all(_calculate_length(c, "utf16") <= 100 for c in sent)
|
||||
assert "".join(sent) == reply
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamer_follow_on_not_double_retried():
|
||||
"""Adapters already retry + park exhausted sends; the streamer must not
|
||||
retry again, or a parked chunk would be re-sent and duplicated (#6078)."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from praisonai_bot.bots._streaming import DraftStreamer, StreamingConfig
|
||||
|
||||
send = AsyncMock(side_effect=[{"message_id": "m1"}, ConnectionError("parked"), None])
|
||||
adapter = SimpleNamespace(
|
||||
capabilities={"live_edit": True, "text_limit": 10},
|
||||
send_message=send,
|
||||
edit_message=AsyncMock(return_value=None),
|
||||
)
|
||||
streamer = DraftStreamer(
|
||||
adapter=adapter, channel_id="C1",
|
||||
config=StreamingConfig(mode=StreamingMode.DRAFT), platform="telegram",
|
||||
)
|
||||
await streamer.start()
|
||||
await streamer.finalize("a" * 30)
|
||||
|
||||
# placeholder + exactly one attempt per follow-on chunk (2 chunks)
|
||||
assert send.await_count == 3
|
||||
|
||||
|
||||
def test_fit_prefix_utf16_is_bounded():
|
||||
from praisonai_bot.bots._streaming import _fit_prefix
|
||||
|
||||
assert _fit_prefix("\U0001F600" * 5, 5, "utf16") == "\U0001F600" * 2
|
||||
assert _fit_prefix("abc", 10, "utf16") == "abc"
|
||||
assert _fit_prefix("abcdef", 3, "codepoints") == "abc"
|
||||
@@ -0,0 +1,167 @@
|
||||
"""Per-channel config lifecycle: add / update / disable / remove (Issue #6079)."""
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from praisonai_bot.cli.commands.gateway import app
|
||||
from praisonai_bot.gateway.admin import edit_gateway_channel
|
||||
|
||||
BASE = {
|
||||
"agents": {"assistant": {"instructions": "hi"}},
|
||||
"channels": {
|
||||
"telegram": {"platform": "telegram", "token": "${TELEGRAM_BOT_TOKEN}"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config(tmp_path):
|
||||
path = tmp_path / "gateway.yaml"
|
||||
path.write_text(yaml.safe_dump(BASE, sort_keys=False))
|
||||
return path
|
||||
|
||||
|
||||
def _load(path):
|
||||
return yaml.safe_load(path.read_text())
|
||||
|
||||
|
||||
def _run(*args):
|
||||
return CliRunner().invoke(app, ["channels", *args])
|
||||
|
||||
|
||||
def test_add_channel_keeps_others_and_defaults_env_token(config):
|
||||
result = _run("add", "discord", "--config", str(config), "--set", "group_policy=mention_only")
|
||||
assert result.exit_code == 0, result.output
|
||||
channels = _load(config)["channels"]
|
||||
assert channels["telegram"] == BASE["channels"]["telegram"]
|
||||
assert channels["discord"] == {
|
||||
"platform": "discord",
|
||||
"token": "${DISCORD_BOT_TOKEN}",
|
||||
"group_policy": "mention_only",
|
||||
}
|
||||
assert _load(config)["agents"] == BASE["agents"]
|
||||
|
||||
|
||||
def test_add_existing_channel_is_rejected(config):
|
||||
result = _run("add", "telegram", "--config", str(config))
|
||||
assert result.exit_code == 1
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_update_and_disable_channel(config):
|
||||
assert _run("update", "telegram", "--config", str(config), "--token", "${NEW_TOKEN}").exit_code == 0
|
||||
assert _run("add", "discord", "--config", str(config)).exit_code == 0
|
||||
assert _run("disable", "telegram", "--config", str(config)).exit_code == 0
|
||||
telegram = _load(config)["channels"]["telegram"]
|
||||
assert telegram["token"] == "${NEW_TOKEN}"
|
||||
assert telegram["enabled"] is False
|
||||
|
||||
|
||||
def test_invalid_update_is_not_persisted(config):
|
||||
result = _run("update", "telegram", "--config", str(config), "--set", "allowed_users=[1, 2]")
|
||||
assert result.exit_code == 1
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_unknown_platform_is_rejected(config):
|
||||
assert _run("add", "bogus", "--config", str(config)).exit_code == 1
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_remove_channel(config):
|
||||
assert _run("add", "discord", "--config", str(config)).exit_code == 0
|
||||
assert _run("remove", "discord", "--config", str(config)).exit_code == 0
|
||||
assert "discord" not in _load(config)["channels"]
|
||||
assert _run("remove", "discord", "--config", str(config)).exit_code == 1
|
||||
|
||||
|
||||
def test_python_helper_matches_cli(config):
|
||||
block = edit_gateway_channel(config, "discord", {"token": "${D}"})
|
||||
assert block == {"platform": "discord", "token": "${D}"}
|
||||
with pytest.raises(KeyError):
|
||||
edit_gateway_channel(config, "nope", remove=True)
|
||||
with pytest.raises(ValueError):
|
||||
edit_gateway_channel(config, "bogus", {})
|
||||
assert "bogus" not in _load(config)["channels"]
|
||||
|
||||
|
||||
def test_list_still_works_without_subcommand(config):
|
||||
result = _run("--config", str(config), "--json")
|
||||
assert result.exit_code == 0
|
||||
assert '"telegram"' in result.output
|
||||
|
||||
|
||||
def test_parent_config_option_applies_to_subcommand(config, tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
other = tmp_path / "other.yaml"
|
||||
other.write_text(yaml.safe_dump(BASE, sort_keys=False))
|
||||
result = _run("--config", str(other), "update", "telegram", "--set", "group_policy=mention_only")
|
||||
assert result.exit_code == 0, result.output
|
||||
assert _load(other)["channels"]["telegram"]["group_policy"] == "mention_only"
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_single_bot_platform_is_converted_not_hidden(tmp_path):
|
||||
path = tmp_path / "bot.yaml"
|
||||
path.write_text(yaml.safe_dump({"agent": {"instructions": "hi"}, "platform": "telegram", "token": "${T}"}))
|
||||
edit_gateway_channel(path, "discord", {"token": "${D}"})
|
||||
cfg = _load(path)
|
||||
assert "platform" not in cfg and "token" not in cfg
|
||||
assert cfg["channels"]["telegram"] == {"platform": "telegram", "token": "${T}"}
|
||||
assert cfg["channels"]["discord"]["token"] == "${D}"
|
||||
|
||||
|
||||
def test_botos_platforms_are_converted(tmp_path):
|
||||
path = tmp_path / "bot.yaml"
|
||||
path.write_text(yaml.safe_dump({"agent": {"instructions": "hi"}, "platforms": {"telegram": {"token": "${T}"}}}))
|
||||
result = CliRunner().invoke(app, ["channels", "update", "telegram", "--config", str(path), "--set", "group_policy=mention_only"])
|
||||
assert result.exit_code == 0, result.output
|
||||
cfg = _load(path)
|
||||
assert "platforms" not in cfg
|
||||
assert cfg["channels"]["telegram"] == {"platform": "telegram", "token": "${T}", "group_policy": "mention_only"}
|
||||
|
||||
|
||||
def test_env_placeholder_in_typed_field_does_not_block_edit(config, monkeypatch):
|
||||
cfg = _load(config)
|
||||
cfg["channels"]["telegram"]["max_inbound_media_bytes"] = "${MEDIA_LIMIT}"
|
||||
config.write_text(yaml.safe_dump(cfg, sort_keys=False))
|
||||
monkeypatch.setenv("MEDIA_LIMIT", "1024")
|
||||
edit_gateway_channel(config, "telegram", {"token": "${NEW}"})
|
||||
saved = _load(config)["channels"]["telegram"]
|
||||
assert saved["max_inbound_media_bytes"] == "${MEDIA_LIMIT}"
|
||||
assert saved["token"] == "${NEW}"
|
||||
|
||||
|
||||
def test_removing_last_channel_is_rejected_even_with_env_token(config, monkeypatch):
|
||||
monkeypatch.setenv("TELEGRAM_BOT_TOKEN", "x")
|
||||
with pytest.raises(ValueError):
|
||||
edit_gateway_channel(config, "telegram", remove=True)
|
||||
with pytest.raises(ValueError):
|
||||
edit_gateway_channel(config, "telegram", {"enabled": False})
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_env_platform_reference_is_resolved_before_conversion(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("BOT_PLATFORM", "telegram")
|
||||
path = tmp_path / "bot.yaml"
|
||||
path.write_text(yaml.safe_dump({"agent": {"instructions": "hi"}, "platform": "${BOT_PLATFORM}", "token": "${T}"}))
|
||||
edit_gateway_channel(path, "discord", {"token": "${D}"})
|
||||
cfg = _load(path)
|
||||
assert cfg["channels"]["telegram"] == {"platform": "telegram", "token": "${T}"}
|
||||
assert cfg["channels"]["discord"]["token"] == "${D}"
|
||||
|
||||
|
||||
def test_concurrent_edits_are_not_lost(config):
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
names = [f"ch{i}" for i in range(8)]
|
||||
with ThreadPoolExecutor(len(names)) as pool:
|
||||
list(pool.map(lambda n: edit_gateway_channel(config, n, {"platform": "telegram", "token": "${X}"}), names))
|
||||
assert set(_load(config)["channels"]) == {"telegram", *names}
|
||||
|
||||
|
||||
def test_missing_config_raises_without_lock_file(tmp_path):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
edit_gateway_channel(tmp_path / "nope.yaml", "discord", {})
|
||||
assert not list(tmp_path.iterdir())
|
||||
@@ -378,3 +378,51 @@ def test_gateway_backend_forwards_authorized_reviewers(tmp_path):
|
||||
assert decision.approved is True
|
||||
|
||||
asyncio.run(go())
|
||||
|
||||
|
||||
def test_allowlist_hot_path_enforces_grant_ttl(tmp_path, monkeypatch):
|
||||
"""An allow-always grant expires in memory without a process restart."""
|
||||
import time as _time
|
||||
|
||||
from praisonai_bot.gateway.exec_approval import ScopedAllowlistStore
|
||||
|
||||
store = ScopedAllowlistStore(path=tmp_path / "allow.sqlite", ttl_seconds=60)
|
||||
al = PermissionAllowlist(store=store)
|
||||
al.add_scoped(agent_id="a1", tool_name="shell_exec", arg_signature="sig")
|
||||
assert al.allows(agent_id="a1", tool_name="shell_exec", arg_signature="sig")
|
||||
|
||||
now = _time.time()
|
||||
monkeypatch.setattr(_time, "time", lambda: now + 61)
|
||||
assert not al.allows(agent_id="a1", tool_name="shell_exec", arg_signature="sig")
|
||||
assert al.list_scoped() == []
|
||||
|
||||
|
||||
def test_allowlist_in_memory_grant_ttl(monkeypatch):
|
||||
"""Non-durable grants honour the default 90-day cap too."""
|
||||
import time as _time
|
||||
|
||||
al = PermissionAllowlist(durable=False)
|
||||
al.add("shell_exec")
|
||||
assert "shell_exec" in al
|
||||
now = _time.time()
|
||||
monkeypatch.setattr(_time, "time", lambda: now + 91 * 86400)
|
||||
assert "shell_exec" not in al
|
||||
assert al.list() == []
|
||||
|
||||
|
||||
def test_allowlist_rehydrate_keeps_original_grant_age(tmp_path, monkeypatch):
|
||||
"""Rehydrated grants keep their persisted created_at, not restart time."""
|
||||
import time as _time
|
||||
|
||||
from praisonai_bot.gateway.exec_approval import ScopedAllowlistStore
|
||||
|
||||
path = tmp_path / "allow.sqlite"
|
||||
PermissionAllowlist(store=ScopedAllowlistStore(path=path, ttl_seconds=60)).add(
|
||||
"shell_exec"
|
||||
)
|
||||
now = _time.time()
|
||||
monkeypatch.setattr(_time, "time", lambda: now + 30)
|
||||
al = PermissionAllowlist(store=ScopedAllowlistStore(path=path, ttl_seconds=60))
|
||||
assert "shell_exec" in al
|
||||
monkeypatch.setattr(_time, "time", lambda: now + 61)
|
||||
assert "shell_exec" not in al
|
||||
@@ -0,0 +1,167 @@
|
||||
"""Regression tests for #6025: the dispatcher heartbeats its own children.
|
||||
|
||||
A healthy worker that outlives the claim lease + stale timeout must not be
|
||||
reclaimed and re-dispatched while it is still running.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonai_bot.gateway.kanban_dispatcher import KanbanDispatcher
|
||||
from praisonai_bot.kanban.sqlite_store import SQLiteKanbanStore
|
||||
|
||||
pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="uses sleep stub")
|
||||
|
||||
|
||||
class _StubDispatcher(KanbanDispatcher):
|
||||
def __init__(self, store, seconds, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._store = store
|
||||
self._seconds = seconds
|
||||
|
||||
def _get_kanban_store(self):
|
||||
return self._store
|
||||
|
||||
def _build_execution_command(self, task):
|
||||
return [sys.executable, "-c", f"import time; time.sleep({self._seconds})"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("PRAISONAI_KANBAN_DB", str(tmp_path / "kanban.db"))
|
||||
return SQLiteKanbanStore()
|
||||
|
||||
|
||||
def _run(coro):
|
||||
return asyncio.run(coro)
|
||||
|
||||
|
||||
def test_long_running_worker_is_not_reclaimed(store):
|
||||
store.create_task({"id": "t1", "title": "long job", "status": "ready",
|
||||
"workspace_kind": "default"})
|
||||
d = _StubDispatcher(store, 5, max_concurrent=3, poll_interval=0.5,
|
||||
claim_ttl_seconds=1, stale_timeout_seconds=2)
|
||||
|
||||
async def drive():
|
||||
handles = []
|
||||
loop = asyncio.get_running_loop()
|
||||
t0 = loop.time()
|
||||
while loop.time() - t0 < 4.5:
|
||||
await d.dispatch_once()
|
||||
proc = d.running_tasks.get("t1")
|
||||
if proc is not None and proc not in handles:
|
||||
handles.append(proc)
|
||||
await asyncio.sleep(0.25)
|
||||
return handles
|
||||
|
||||
handles = _run(drive())
|
||||
try:
|
||||
assert len(handles) == 1
|
||||
runs = store.get_runs("t1")
|
||||
assert len(runs) == 1
|
||||
assert runs[0].ended_at is None
|
||||
assert not [e for e in store.list_events()
|
||||
if e.task_id == "t1" and e.event_type == "reclaimed"]
|
||||
|
||||
handles[0].wait(timeout=10)
|
||||
_run(d.dispatch_once())
|
||||
runs = store.get_runs("t1")
|
||||
assert len(runs) == 1
|
||||
assert runs[0].outcome == "completed"
|
||||
assert store.get_task("t1").status == "done"
|
||||
finally:
|
||||
for p in handles:
|
||||
if p.poll() is None:
|
||||
p.kill()
|
||||
|
||||
|
||||
def test_lost_claim_terminates_child(store):
|
||||
store.create_task({"id": "t2", "title": "job", "status": "ready",
|
||||
"workspace_kind": "default"})
|
||||
d = _StubDispatcher(store, 30, max_concurrent=1, poll_interval=0.5)
|
||||
_run(d.dispatch_once())
|
||||
proc = d.running_tasks["t2"]
|
||||
try:
|
||||
store.release_claim("t2", d.worker_id) # operator takes the card back
|
||||
d._heartbeat_children(store)
|
||||
assert "t2" not in d.running_tasks
|
||||
assert proc.wait(timeout=10) is not None
|
||||
runs = store.get_runs("t2")
|
||||
assert len(runs) == 1 and runs[0].outcome == "failed"
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
|
||||
|
||||
def test_reclaim_of_dead_worker_closes_run_and_counts_failure(store):
|
||||
store.create_task({"id": "t3", "title": "job", "status": "ready",
|
||||
"workspace_kind": "default"})
|
||||
# Simulate another dispatcher whose worker has died.
|
||||
assert store.claim_task("t3", "other", ttl_seconds=0, worker_pid=0)
|
||||
run_id = store.start_run("t3", profile="other")
|
||||
d = KanbanDispatcher(claim_ttl_seconds=1, stale_timeout_seconds=1)
|
||||
|
||||
d._reclaim_stale_claims(store)
|
||||
|
||||
run = [r for r in store.get_runs("t3") if r.id == run_id][0]
|
||||
assert run.outcome == "crashed"
|
||||
assert run.ended_at
|
||||
assert store.get_task("t3").consecutive_failures == 1
|
||||
|
||||
|
||||
def test_live_handle_is_never_overwritten(store):
|
||||
store.create_task({"id": "t4", "title": "job", "status": "ready",
|
||||
"workspace_kind": "default"})
|
||||
d = _StubDispatcher(store, 30, max_concurrent=2)
|
||||
proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
|
||||
d.running_tasks["t4"] = proc
|
||||
d._heartbeat_children = lambda _store: None # isolate the spawn guard
|
||||
try:
|
||||
assert _run(d.dispatch_once()) == 0
|
||||
assert d.running_tasks["t4"] is proc
|
||||
finally:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
|
||||
|
||||
def test_reclaim_closes_run_atomically_so_replacement_run_survives(store):
|
||||
store.create_task({"id": "t5", "title": "job", "status": "ready",
|
||||
"workspace_kind": "default"})
|
||||
assert store.claim_task("t5", "other", ttl_seconds=0, worker_pid=0)
|
||||
old_run = store.start_run("t5", profile="other")
|
||||
|
||||
assert store.reclaim_stale_claims(stale_timeout_seconds=1) == ["t5"]
|
||||
assert store.get_task("t5").current_run_id is None
|
||||
assert any(e.event_type == "run_closed" and e.task_id == "t5"
|
||||
for e in store.list_events())
|
||||
|
||||
# Another dispatcher wins the reclaimed task before ours reacts.
|
||||
assert store.claim_task("t5", "rival", ttl_seconds=60)
|
||||
new_run = store.start_run("t5", profile="rival")
|
||||
KanbanDispatcher()._reclaim_stale_claims(store)
|
||||
|
||||
runs = {r.id: r for r in store.get_runs("t5")}
|
||||
assert runs[old_run].outcome == "crashed"
|
||||
assert runs[new_run].ended_at is None
|
||||
|
||||
|
||||
def test_lost_claim_keeps_worktree_tracking_when_work_is_preserved(store):
|
||||
store.create_task({"id": "t6", "title": "job", "status": "ready",
|
||||
"workspace_kind": "default"})
|
||||
d = _StubDispatcher(store, 30, max_concurrent=1)
|
||||
_run(d.dispatch_once())
|
||||
proc = d.running_tasks["t6"]
|
||||
d._worktrees = {"t6": ("/tmp/wt-t6", "kanban/t6")}
|
||||
d._remove_worktree = lambda path, branch: "worktree has uncommitted changes"
|
||||
try:
|
||||
store.release_claim("t6", d.worker_id)
|
||||
d._heartbeat_children(store)
|
||||
assert "t6" not in d.running_tasks
|
||||
assert d._worktrees == {"t6": ("/tmp/wt-t6", "kanban/t6")}
|
||||
assert any("worktree_preserved" in c.text for c in store.get_comments("t6"))
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
@@ -0,0 +1,136 @@
|
||||
"""The kanban dispatcher must not run blocking git work on the event loop."""
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
import praisonai_bot.gateway.kanban_dispatcher as kd
|
||||
from praisonai_bot.kanban.sqlite_store import SQLiteKanbanStore
|
||||
|
||||
GIT_DELAY = 0.5
|
||||
|
||||
|
||||
class _Stub(kd.KanbanDispatcher):
|
||||
def _build_execution_command(self, task):
|
||||
return ["true"]
|
||||
|
||||
|
||||
async def _heartbeat(gaps, stop):
|
||||
last = time.monotonic()
|
||||
while not stop.is_set():
|
||||
await asyncio.sleep(0.05)
|
||||
now = time.monotonic()
|
||||
gaps.append(now - last)
|
||||
last = now
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def slow_git(tmp_path, monkeypatch):
|
||||
"""Patch ``run_guarded`` with a synchronous sleep, recording overlap."""
|
||||
monkeypatch.setenv("PRAISONAI_KANBAN_DB", str(tmp_path / "kanban.db"))
|
||||
monkeypatch.setenv("HOME", str(tmp_path))
|
||||
monkeypatch.chdir(tmp_path)
|
||||
state = {"active": 0, "max_active": 0, "calls": []}
|
||||
lock = threading.Lock()
|
||||
|
||||
def fake(cmd, *, timeout=None, cwd=None, **kw):
|
||||
with lock:
|
||||
state["active"] += 1
|
||||
state["max_active"] = max(state["max_active"], state["active"])
|
||||
state["calls"].append(" ".join(cmd[:3]))
|
||||
try:
|
||||
time.sleep(GIT_DELAY)
|
||||
if list(cmd[1:3]) == ["worktree", "add"]:
|
||||
os.makedirs(cmd[3], exist_ok=True)
|
||||
out = {"rev-parse": "main\n", "rev-list": "0\n"}.get(cmd[1], "")
|
||||
return subprocess.CompletedProcess(cmd, 0, stdout=out, stderr="")
|
||||
finally:
|
||||
with lock:
|
||||
state["active"] -= 1
|
||||
|
||||
monkeypatch.setattr(kd, "run_guarded", fake)
|
||||
return state
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform == "win32", reason="uses the POSIX `true` worker")
|
||||
async def test_dispatch_does_not_block_event_loop(slow_git):
|
||||
store = SQLiteKanbanStore()
|
||||
store.create_task({"id": "w1", "title": "isolated", "status": "ready",
|
||||
"workspace_kind": "worktree"})
|
||||
d = _Stub(max_concurrent=1)
|
||||
gaps, stop = [], asyncio.Event()
|
||||
hb = asyncio.create_task(_heartbeat(gaps, stop))
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
assert await d.dispatch_once() == 1 # spawn: rev-parse + worktree add
|
||||
await asyncio.to_thread(d.running_tasks["w1"].wait)
|
||||
await d.dispatch_once() # completion: commit/merge/cleanup
|
||||
await asyncio.sleep(0.1)
|
||||
stop.set()
|
||||
await hb
|
||||
|
||||
assert "git worktree add" in slow_git["calls"]
|
||||
assert "git merge --no-ff" in slow_git["calls"]
|
||||
assert max(gaps) < 0.3, f"event loop stalled for {max(gaps):.2f}s"
|
||||
|
||||
|
||||
async def test_blocking_work_is_serialised(slow_git):
|
||||
a, b = _Stub(), _Stub()
|
||||
await asyncio.gather(
|
||||
a._run_blocking(a._run_git, "status"),
|
||||
b._run_blocking(b._run_git, "status"),
|
||||
)
|
||||
assert slow_git["max_active"] == 1
|
||||
|
||||
|
||||
async def test_cancel_during_worktree_prep_cleans_up(slow_git):
|
||||
store = SQLiteKanbanStore()
|
||||
store.create_task({"id": "w1", "title": "isolated", "status": "ready",
|
||||
"workspace_kind": "worktree"})
|
||||
d = _Stub(max_concurrent=1)
|
||||
t = asyncio.create_task(d.dispatch_once())
|
||||
await asyncio.sleep(GIT_DELAY + 0.1) # inside `git worktree add`
|
||||
t.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await t
|
||||
|
||||
assert "git worktree remove" in slow_git["calls"]
|
||||
assert "git branch -D" in slow_git["calls"]
|
||||
assert not d.running_tasks
|
||||
assert store.get_task("w1").status == "ready"
|
||||
assert all(r.ended_at for r in store.get_runs("w1"))
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform == "win32", reason="POSIX process groups")
|
||||
async def test_shutdown_terminates_worker_descendants(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("PRAISONAI_KANBAN_DB", str(tmp_path / "kanban.db"))
|
||||
marker = tmp_path / "survived"
|
||||
|
||||
class _Spawner(kd.KanbanDispatcher):
|
||||
def _build_execution_command(self, task):
|
||||
# Worker exits on SIGTERM, but its child ignores SIGTERM and
|
||||
# would leave a marker if it outlived shutdown.
|
||||
return ["sh", "-c",
|
||||
f"(trap '' TERM; sleep 1; touch {marker}) & wait"]
|
||||
|
||||
class _Task:
|
||||
id, board, title, body = "t1", "default", "t1", ""
|
||||
|
||||
d = _Spawner()
|
||||
assert await d._spawn_worker(_Task(), store=None)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
# Skip the 30 s graceful window; go straight to force-terminate.
|
||||
ticks = itertools.count(step=1000)
|
||||
monkeypatch.setattr(kd, "time", types.SimpleNamespace(time=lambda: next(ticks)))
|
||||
await d._shutdown()
|
||||
|
||||
await asyncio.sleep(1.5)
|
||||
assert not marker.exists(), "worker's descendant survived shutdown"
|
||||
@@ -0,0 +1,61 @@
|
||||
"""The gateway scheduler tick loop survives a failing tick (issue #6035)."""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("praisonaiagents.scheduler")
|
||||
|
||||
import praisonaiagents.scheduler as sched_pkg
|
||||
from praisonaiagents.scheduler.store import FileScheduleStore
|
||||
|
||||
from praisonai_bot.gateway.server import WebSocketGateway
|
||||
from praisonai_bot.scheduler.executor import ScheduledAgentExecutor
|
||||
|
||||
|
||||
def _bare_gateway(config=None):
|
||||
gw = WebSocketGateway.__new__(WebSocketGateway) # skip heavy __init__
|
||||
gw._agents, gw._loaded_config, gw._scheduler_task = {}, config or {}, None
|
||||
|
||||
async def _deliver(*a, **k):
|
||||
return None
|
||||
|
||||
gw._deliver_scheduled_result = _deliver
|
||||
return gw
|
||||
|
||||
|
||||
async def test_tick_loop_survives_exception(monkeypatch):
|
||||
store = FileScheduleStore(store_dir=tempfile.mkdtemp())
|
||||
monkeypatch.setattr(sched_pkg, "get_default_store", lambda: store)
|
||||
calls = {"n": 0}
|
||||
orig = ScheduledAgentExecutor.tick_all
|
||||
|
||||
async def flaky_tick_all(self):
|
||||
calls["n"] += 1
|
||||
if calls["n"] == 1:
|
||||
raise OSError(24, "Too many open files")
|
||||
return await orig(self)
|
||||
|
||||
monkeypatch.setattr(ScheduledAgentExecutor, "tick_all", flaky_tick_all)
|
||||
|
||||
gw = _bare_gateway()
|
||||
gw._start_scheduler_tick(interval=0.01)
|
||||
try:
|
||||
await asyncio.sleep(0.3)
|
||||
assert calls["n"] >= 2
|
||||
assert not gw._scheduler_task.done()
|
||||
finally:
|
||||
gw._scheduler_task.cancel()
|
||||
await asyncio.gather(gw._scheduler_task, return_exceptions=True)
|
||||
|
||||
|
||||
def test_non_mapping_scheduler_block_uses_defaults():
|
||||
pytest.importorskip("praisonai.scheduler.run_policy")
|
||||
gw = _bare_gateway({"scheduler": ["x"]})
|
||||
policy = gw._build_run_policy()
|
||||
assert policy is not None
|
||||
assert policy.deliver_on_failure is True
|
||||
assert policy.scan_assembled_prompt is True
|
||||
from praisonai.scheduler.run_policy import DEFAULT_DENIED_TOOLSETS
|
||||
assert policy.denied_toolsets == set(DEFAULT_DENIED_TOOLSETS)
|
||||
Generated
+4
-4
@@ -1428,7 +1428,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-bot"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -1510,7 +1510,7 @@ provides-extras = ["gateway", "bot", "bot-whatsapp-web", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -1539,8 +1539,8 @@ requires-dist = [
|
||||
{ name = "fastapi", marker = "extra == 'api'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'mcp'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'os'", specifier = ">=0.115.0" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "markitdown", extras = ["all"], marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
|
||||
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.20.0,<2" },
|
||||
{ name = "mcp", marker = "extra == 'parallel-search'", specifier = ">=1.30.0,<2" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.15"
|
||||
__version__ = "0.1.18"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-browser"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
description = "Browser automation for PraisonAI — Chrome extension bridge, CDP, and Playwright agents extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+4
-4
@@ -932,7 +932,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-browser"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -976,7 +976,7 @@ provides-extras = ["server", "playwright", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -1005,8 +1005,8 @@ requires-dist = [
|
||||
{ name = "fastapi", marker = "extra == 'api'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'mcp'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'os'", specifier = ">=0.115.0" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "markitdown", extras = ["all"], marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
|
||||
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.20.0,<2" },
|
||||
{ name = "mcp", marker = "extra == 'parallel-search'", specifier = ">=1.30.0,<2" },
|
||||
|
||||
@@ -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.15"
|
||||
__version__ = "0.1.18"
|
||||
|
||||
__all__ = ["__version__"]
|
||||
@@ -88,8 +88,11 @@ def _run_loop(
|
||||
typer.echo("Error: praisonaiagents not installed. Run: pip install praisonaiagents", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Build autonomy config
|
||||
# Build autonomy config. `loop` always iterates: without an explicit mode
|
||||
# the default level (suggest) maps to single-call "caller" mode, which
|
||||
# would ignore --completion-promise and --max-iterations.
|
||||
autonomy_config = {
|
||||
"mode": "iterative",
|
||||
"max_iterations": max_iterations,
|
||||
"completion_promise": completion_promise,
|
||||
"clear_context": clear_context,
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-code"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
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,19 @@
|
||||
"""`praisonai loop` must run the iterative autonomy loop by default (#6027)."""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from praisonai_code.cli.commands.loop import _run_loop
|
||||
|
||||
|
||||
def test_loop_builds_iterative_autonomy_config():
|
||||
result = MagicMock(success=True, completion_reason="promise", iterations=2,
|
||||
duration_seconds=0.1, started_at=None)
|
||||
with patch("praisonaiagents.Agent") as agent_cls:
|
||||
agent_cls.return_value.start.return_value = result
|
||||
_run_loop("Refactor auth module", max_iterations=10, completion_promise="DONE",
|
||||
clear_context=False, timeout=None, model=None, verbose=False)
|
||||
|
||||
autonomy = agent_cls.call_args.kwargs["autonomy"]
|
||||
assert autonomy["mode"] == "iterative"
|
||||
assert autonomy["completion_promise"] == "DONE"
|
||||
assert autonomy["max_iterations"] == 10
|
||||
Generated
+2
-2
@@ -1370,7 +1370,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-code"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1403,7 +1403,7 @@ requires-dist = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.15"
|
||||
__version__ = "0.1.18"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-deploy"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
description = "Deployment for PraisonAI — API, Docker, and cloud (AWS, Azure, GCP) deployments extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
|
||||
Generated
+4
-4
@@ -977,7 +977,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-deploy"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1014,7 +1014,7 @@ provides-extras = ["api", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -1043,8 +1043,8 @@ requires-dist = [
|
||||
{ name = "fastapi", marker = "extra == 'api'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'mcp'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'os'", specifier = ">=0.115.0" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "markitdown", extras = ["all"], marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
|
||||
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.20.0,<2" },
|
||||
{ name = "mcp", marker = "extra == 'parallel-search'", specifier = ">=1.30.0,<2" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.17"
|
||||
__version__ = "0.1.20"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-mcp"
|
||||
version = "0.1.17"
|
||||
version = "0.1.20"
|
||||
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
+4
-4
@@ -1044,7 +1044,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-mcp"
|
||||
version = "0.1.17"
|
||||
version = "0.1.20"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1084,7 +1084,7 @@ provides-extras = ["server", "auth", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -1113,8 +1113,8 @@ requires-dist = [
|
||||
{ name = "fastapi", marker = "extra == 'api'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'mcp'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'os'", specifier = ">=0.115.0" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "markitdown", extras = ["all"], marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
|
||||
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.20.0,<2" },
|
||||
{ name = "mcp", marker = "extra == 'parallel-search'", specifier = ">=1.30.0,<2" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.15"
|
||||
__version__ = "0.1.18"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-sandbox"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
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
+4
-4
@@ -1679,7 +1679,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-sandbox"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -1744,7 +1744,7 @@ provides-extras = ["docker", "e2b", "sandlock", "ssh", "tenki", "flyio", "modal"
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -1773,8 +1773,8 @@ requires-dist = [
|
||||
{ name = "fastapi", marker = "extra == 'api'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'mcp'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'os'", specifier = ">=0.115.0" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "markitdown", extras = ["all"], marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
|
||||
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.20.0,<2" },
|
||||
{ name = "mcp", marker = "extra == 'parallel-search'", specifier = ">=1.30.0,<2" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.15"
|
||||
__version__ = "0.1.18"
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "praisonai-train"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
description = "Training for PraisonAI — LLM fine-tuning (Unsloth) and iterative agent training extracted from the praisonai wrapper."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
||||
Generated
+4
-4
@@ -2606,7 +2606,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-train"
|
||||
version = "0.1.15"
|
||||
version = "0.1.18"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
@@ -2696,7 +2696,7 @@ provides-extras = ["llm", "vision", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
source = { directory = "../praisonai-agents" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -2725,8 +2725,8 @@ requires-dist = [
|
||||
{ name = "fastapi", marker = "extra == 'api'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'mcp'", specifier = ">=0.115.0" },
|
||||
{ name = "fastapi", marker = "extra == 'os'", specifier = ">=0.115.0" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14" },
|
||||
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14,<1.106" },
|
||||
{ name = "markitdown", extras = ["all"], marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
|
||||
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.20.0,<2" },
|
||||
{ name = "mcp", marker = "extra == 'parallel-search'", specifier = ">=1.30.0,<2" },
|
||||
|
||||
@@ -46,7 +46,7 @@ This complements `PARITY.md`, which only tracks whether an export exists.
|
||||
|
||||
### `Agent.__init__`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/agent.py:624`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/agent.py:628`
|
||||
- TypeScript: `src/praisonai-ts/src/agent/simple.ts:197` (ctor `src/praisonai-ts/src/agent/simple.ts:1019`)
|
||||
- Counts: 42 python params: 29 exact, 8 camelCase, 2 alias, 3 flattened, 0 missing; 8 mismatches; 8 waived; 14 TS-only of 60
|
||||
|
||||
@@ -207,7 +207,7 @@ TS-only members: `dependencies`?
|
||||
|
||||
### `Agent.start`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/execution_mixin.py:799`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/execution_mixin.py:753`
|
||||
- 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
|
||||
|
||||
@@ -219,7 +219,7 @@ TS-only members: `previousResult`?, `onToken`?, `signal`?, `onEvent`?, `options`
|
||||
|
||||
### `Agent.chat`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3371`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3530`
|
||||
- TypeScript: `src/praisonai-ts/src/agent/simple.ts:3592`
|
||||
- Counts: 17 python params: 7 exact, 9 camelCase, 1 alias, 0 flattened, 0 missing; 0 mismatches; 0 waived; 2 TS-only of 19
|
||||
|
||||
@@ -279,7 +279,7 @@ TS-only members: `condition`?, `contextPolicy`?, `maxContextTokens`?, `maxContex
|
||||
|
||||
### `LLM.__init__`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/llm/llm.py:673`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/llm/llm.py:678`
|
||||
- TypeScript: `src/praisonai-ts/src/llm/index.ts:236` (ctor `src/praisonai-ts/src/llm/index.ts:416`)
|
||||
- Counts: 25 python params: 8 exact, 16 camelCase, 1 alias, 0 flattened, 0 missing; 1 mismatches; 1 waived; 1 TS-only of 26
|
||||
|
||||
|
||||
@@ -1185,7 +1185,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/agent.py:624",
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/agent.py:628",
|
||||
"resolved_class": "Agent"
|
||||
},
|
||||
"ts_only": [
|
||||
@@ -1797,7 +1797,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3371",
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3530",
|
||||
"resolved_class": "ChatMixin"
|
||||
},
|
||||
"ts_only": [
|
||||
@@ -1876,7 +1876,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/execution_mixin.py:799",
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/execution_mixin.py:753",
|
||||
"resolved_class": "ExecutionMixin"
|
||||
},
|
||||
"ts_only": [
|
||||
@@ -4659,7 +4659,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/llm/llm.py:673",
|
||||
"location": "src/praisonai-agents/praisonaiagents/llm/llm.py:678",
|
||||
"resolved_class": "LLM"
|
||||
},
|
||||
"ts_only": [
|
||||
|
||||
@@ -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.18.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.7.18.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.7.21.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v4.7.21.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -8,13 +8,20 @@ Install: pip install pymongo
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .base import StateStore
|
||||
from ...storage import MongoDBStorageAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _utcnow() -> datetime:
|
||||
"""Naive UTC now, matching the naive UTC datetimes pymongo returns."""
|
||||
return datetime.now(timezone.utc).replace(tzinfo=None)
|
||||
|
||||
|
||||
class MongoDBStateStore(StateStore):
|
||||
"""
|
||||
MongoDB-based state store.
|
||||
@@ -32,21 +39,17 @@ class MongoDBStateStore(StateStore):
|
||||
database: str = "praisonai",
|
||||
collection: str = "state",
|
||||
):
|
||||
try:
|
||||
from pymongo import MongoClient
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"pymongo is required for MongoDB support. "
|
||||
"Install with: pip install pymongo"
|
||||
)
|
||||
|
||||
self._client = MongoClient(url)
|
||||
self._db = self._client[database]
|
||||
self._collection = self._db[collection]
|
||||
|
||||
# Use the canonical storage adapter for the pymongo client (single
|
||||
# driver init path: pooling, timeouts, retries, import error message).
|
||||
self._adapter = MongoDBStorageAdapter(
|
||||
url=url, database=database, collection=collection,
|
||||
)
|
||||
self._collection = self._adapter._get_collection()
|
||||
self._client = self._adapter._client
|
||||
|
||||
# Create TTL index for automatic expiration
|
||||
self._collection.create_index("expires_at", expireAfterSeconds=0)
|
||||
|
||||
|
||||
logger.info(f"Connected to MongoDB: {database}.{collection}")
|
||||
|
||||
def get(self, key: str) -> Optional[Any]:
|
||||
@@ -56,7 +59,7 @@ class MongoDBStateStore(StateStore):
|
||||
return None
|
||||
|
||||
# Check TTL
|
||||
if doc.get("expires_at") and doc["expires_at"] <= time.time():
|
||||
if doc.get("expires_at") and doc["expires_at"] <= _utcnow():
|
||||
self._collection.delete_one({"_id": key})
|
||||
return None
|
||||
|
||||
@@ -72,8 +75,7 @@ class MongoDBStateStore(StateStore):
|
||||
doc = {"_id": key, "value": value, "updated_at": time.time()}
|
||||
|
||||
if ttl:
|
||||
from datetime import datetime, timedelta
|
||||
doc["expires_at"] = datetime.utcnow() + timedelta(seconds=ttl)
|
||||
doc["expires_at"] = _utcnow() + timedelta(seconds=ttl)
|
||||
|
||||
self._collection.replace_one({"_id": key}, doc, upsert=True)
|
||||
|
||||
@@ -87,7 +89,7 @@ class MongoDBStateStore(StateStore):
|
||||
doc = self._collection.find_one({"_id": key}, {"_id": 1, "expires_at": 1})
|
||||
if not doc:
|
||||
return False
|
||||
if doc.get("expires_at") and doc["expires_at"] <= time.time():
|
||||
if doc.get("expires_at") and doc["expires_at"] <= _utcnow():
|
||||
return False
|
||||
return True
|
||||
|
||||
@@ -115,18 +117,16 @@ class MongoDBStateStore(StateStore):
|
||||
if not doc or "expires_at" not in doc:
|
||||
return None
|
||||
|
||||
remaining = doc["expires_at"].timestamp() - time.time()
|
||||
remaining = (doc["expires_at"] - _utcnow()).total_seconds()
|
||||
if remaining <= 0:
|
||||
return None
|
||||
return int(remaining)
|
||||
|
||||
def expire(self, key: str, ttl: int) -> bool:
|
||||
"""Set TTL on existing key."""
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
result = self._collection.update_one(
|
||||
{"_id": key},
|
||||
{"$set": {"expires_at": datetime.utcnow() + timedelta(seconds=ttl)}}
|
||||
{"$set": {"expires_at": _utcnow() + timedelta(seconds=ttl)}}
|
||||
)
|
||||
return result.modified_count > 0
|
||||
|
||||
@@ -160,5 +160,5 @@ class MongoDBStateStore(StateStore):
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the store."""
|
||||
if self._client:
|
||||
self._client.close()
|
||||
self._adapter.close()
|
||||
self._client = None
|
||||
@@ -11,18 +11,16 @@ Heavy implementations that follow StorageBackendProtocol for:
|
||||
NOTE: ``praisonai.storage`` holds the canonical, low-level storage adapters.
|
||||
It is NOT dead code — ``praisonai.persistence`` builds on top of it. For
|
||||
example ``praisonai.persistence.state.redis.RedisStateStore`` is a thin wrapper
|
||||
around ``RedisStorageAdapter`` defined here, so this module is load-bearing and
|
||||
canonical for the backends that ``persistence`` wraps.
|
||||
around ``RedisStorageAdapter`` defined here, and ``MongoDBStateStore`` obtains
|
||||
its pymongo client from ``MongoDBStorageAdapter``, so this module is
|
||||
load-bearing and canonical for the backends that ``persistence`` wraps.
|
||||
Valkey shares driver init via ``persistence._valkey_client``.
|
||||
|
||||
Note on remaining duplication: ``persistence.state.mongodb`` / ``dynamodb`` /
|
||||
``valkey`` currently keep their own driver logic instead of wrapping the
|
||||
adapters here. This is because ``StateStore`` and these adapters implement
|
||||
different contracts — ``StateStore`` adds per-key TTL, hash operations
|
||||
(``hget``/``hset``) and counters (``incr``/``decr``), and the DynamoDB store
|
||||
uses a different partition-key schema (``pk``) than the adapter (``key``).
|
||||
Collapsing them therefore requires care to avoid changing stored-data layout;
|
||||
until that is done, treat these adapters as the canonical KV/blob layer and the
|
||||
``persistence.state`` stores as the richer, TTL/hash-aware state layer.
|
||||
Note on remaining duplication: ``persistence.state.dynamodb`` keeps its own
|
||||
boto3 setup because it uses a different partition-key schema (``pk``) than the
|
||||
adapter (``key``); collapsing it would change stored-data layout. The
|
||||
``persistence.state`` stores layer per-key TTL, hash operations
|
||||
(``hget``/``hset``) and counters on top of the KV/blob adapters here.
|
||||
|
||||
Application code that just needs a persistence layer should use
|
||||
``create_state_store``, ``create_conversation_store``, or
|
||||
|
||||
@@ -91,6 +91,18 @@ class MongoDBStorageAdapter:
|
||||
"retryReads": True,
|
||||
}
|
||||
|
||||
# Options set explicitly in the URI (e.g. ?retryWrites=false) win
|
||||
# over these defaults; URI option names are case-insensitive.
|
||||
query = self.url.split("?", 1)[1] if "?" in self.url else ""
|
||||
uri_opts = {
|
||||
part.split("=", 1)[0].lower()
|
||||
for part in re.split(r"[&;]", query) if part
|
||||
}
|
||||
client_kwargs = {
|
||||
k: v for k, v in client_kwargs.items()
|
||||
if k.lower() not in uri_opts
|
||||
}
|
||||
|
||||
if self.username and self.password:
|
||||
client_kwargs["username"] = self.username
|
||||
client_kwargs["password"] = self.password
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "4.7.18"
|
||||
__version__ = "4.7.21"
|
||||
@@ -12,14 +12,14 @@ dependencies = [
|
||||
"rich>=13.7",
|
||||
"markdown>=3.5",
|
||||
"pyparsing>=3.0.0",
|
||||
"praisonaiagents>=1.7.16",
|
||||
"praisonai-code>=0.1.15",
|
||||
"praisonai-bot>=0.1.15",
|
||||
"praisonai-train>=0.1.15",
|
||||
"praisonai-browser>=0.1.15",
|
||||
"praisonai-mcp>=0.1.17",
|
||||
"praisonai-sandbox>=0.1.15",
|
||||
"praisonai-deploy>=0.1.15",
|
||||
"praisonaiagents>=1.7.19",
|
||||
"praisonai-code>=0.1.18",
|
||||
"praisonai-bot>=0.1.18",
|
||||
"praisonai-train>=0.1.18",
|
||||
"praisonai-browser>=0.1.18",
|
||||
"praisonai-mcp>=0.1.20",
|
||||
"praisonai-sandbox>=0.1.18",
|
||||
"praisonai-deploy>=0.1.18",
|
||||
"python-dotenv>=0.19.0",
|
||||
"litellm>=1.83.14,<1.106", # 1.106 imports mcp.types.CacheableResult (mcp>=2.2); we pin mcp<2
|
||||
"PyYAML>=6.0",
|
||||
@@ -56,7 +56,7 @@ openai-compat = "praisonai.endpoints.providers.openai_compat:OpenAICompatProvide
|
||||
|
||||
[project.optional-dependencies]
|
||||
ui = [
|
||||
"aiui>=0.3.121,<0.4",
|
||||
"aiui>=0.3.148,<0.4",
|
||||
]
|
||||
gradio = ["gradio>=4.26.0"]
|
||||
api = [
|
||||
@@ -78,7 +78,7 @@ agentops = ["agentops>=0.3.12"]
|
||||
langfuse = ["langfuse>=3.0.0"]
|
||||
langextract = ["langextract>=1.0.0"]
|
||||
chat = [
|
||||
"aiui>=0.3.121,<0.4",
|
||||
"aiui>=0.3.148,<0.4",
|
||||
"tavily-python==0.5.0",
|
||||
"crawl4ai>=0.7.0",
|
||||
# Playwright also ships via praisonai-browser[all]; pin here for crawl4ai without [browser]
|
||||
@@ -86,7 +86,7 @@ chat = [
|
||||
"rich"
|
||||
]
|
||||
code = [
|
||||
"aiui>=0.3.121,<0.4",
|
||||
"aiui>=0.3.148,<0.4",
|
||||
"tavily-python==0.5.0",
|
||||
"crawl4ai>=0.7.0",
|
||||
# Playwright also ships via praisonai-browser[all]; pin here for crawl4ai without [browser]
|
||||
@@ -121,7 +121,7 @@ tenki = [
|
||||
"praisonai-sandbox[tenki]",
|
||||
]
|
||||
realtime = [
|
||||
"aiui>=0.3.121,<0.4",
|
||||
"aiui>=0.3.148,<0.4",
|
||||
"tavily-python==0.5.0",
|
||||
"crawl4ai>=0.7.0",
|
||||
"playwright>=1.47.0", # required by crawl4ai — browsers still need: playwright install chromium
|
||||
@@ -229,7 +229,7 @@ all = [
|
||||
"slack-bolt>=1.18.0",
|
||||
"brotlicffi>=1.0.0",
|
||||
"aiohttp>=3.8.0",
|
||||
"aiui>=0.3.121,<0.4", # PraisonAI UI — dashboard + chat for onboarded bots
|
||||
"aiui>=0.3.148,<0.4", # PraisonAI UI — dashboard + chat for onboarded bots
|
||||
# Storage backends
|
||||
"redis>=5.0.0",
|
||||
"pymongo>=4.6.0",
|
||||
@@ -246,7 +246,7 @@ observability = [
|
||||
]
|
||||
bot = [
|
||||
"praisonai-bot[bot]",
|
||||
"aiui>=0.3.121,<0.4",
|
||||
"aiui>=0.3.148,<0.4",
|
||||
]
|
||||
bot-whatsapp-web = [
|
||||
"praisonai-bot[bot-whatsapp-web]",
|
||||
@@ -283,7 +283,7 @@ xata = [
|
||||
claw = [
|
||||
# PraisonAI 🦞 — Full dashboard + bots + memory + knowledge + gateway
|
||||
# Install: pip install "praisonai[claw]"
|
||||
"aiui[all]>=0.3.121,<0.4",
|
||||
"aiui[all]>=0.3.148,<0.4",
|
||||
"praisonaiagents[all]>=1.5.40",
|
||||
"praisonai-bot[gateway,bot]>=0.0.6",
|
||||
# UI + persistence
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""MongoDBStateStore: shared driver init via MongoDBStorageAdapter + TTL handling."""
|
||||
|
||||
import sys
|
||||
import types
|
||||
from datetime import timedelta
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class _FakeCollection:
|
||||
"""Minimal in-memory stand-in for a pymongo collection."""
|
||||
|
||||
def __init__(self):
|
||||
self.docs = {}
|
||||
|
||||
def create_index(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
def find_one(self, query, projection=None):
|
||||
doc = self.docs.get(query["_id"])
|
||||
return dict(doc) if doc else None
|
||||
|
||||
def replace_one(self, query, doc, upsert=False):
|
||||
self.docs[query["_id"]] = dict(doc)
|
||||
|
||||
def update_one(self, query, update, upsert=False):
|
||||
doc = self.docs.get(query["_id"])
|
||||
if doc is None:
|
||||
return MagicMock(modified_count=0)
|
||||
doc.update(update.get("$set", {}))
|
||||
return MagicMock(modified_count=1)
|
||||
|
||||
def delete_one(self, query):
|
||||
return MagicMock(deleted_count=1 if self.docs.pop(query["_id"], None) else 0)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(monkeypatch):
|
||||
collection = _FakeCollection()
|
||||
client = MagicMock()
|
||||
client.__getitem__.return_value.__getitem__.return_value = collection
|
||||
mongo_client_cls = MagicMock(return_value=client)
|
||||
fake_pymongo = types.ModuleType("pymongo")
|
||||
fake_pymongo.MongoClient = mongo_client_cls
|
||||
monkeypatch.setitem(sys.modules, "pymongo", fake_pymongo)
|
||||
|
||||
from praisonai.persistence.state.mongodb import MongoDBStateStore
|
||||
|
||||
s = MongoDBStateStore(url="mongodb://example:27017", database="db", collection="state")
|
||||
s._mongo_client_cls = mongo_client_cls
|
||||
return s
|
||||
|
||||
|
||||
def test_uses_storage_adapter_for_driver(store):
|
||||
from praisonai.storage import MongoDBStorageAdapter
|
||||
|
||||
assert isinstance(store._adapter, MongoDBStorageAdapter)
|
||||
store._mongo_client_cls.assert_called_once()
|
||||
assert store._mongo_client_cls.call_args.args[0] == "mongodb://example:27017"
|
||||
|
||||
|
||||
def test_ttl_roundtrip_does_not_raise(store):
|
||||
store.set("k", {"a": 1}, ttl=60)
|
||||
assert store.get("k") == {"a": 1}
|
||||
assert store.exists("k") is True
|
||||
assert 0 < store.ttl("k") <= 60
|
||||
|
||||
|
||||
def test_expired_key_is_hidden(store):
|
||||
store.set("k", "v", ttl=60)
|
||||
store._collection.docs["k"]["expires_at"] -= timedelta(seconds=120)
|
||||
assert store.exists("k") is False
|
||||
assert store.ttl("k") is None
|
||||
assert store.get("k") is None
|
||||
|
||||
|
||||
def test_uri_options_override_adapter_defaults(store):
|
||||
from praisonai.persistence.state.mongodb import MongoDBStateStore
|
||||
|
||||
MongoDBStateStore(url="mongodb://example:27017/?RetryWrites=false&maxPoolSize=7")
|
||||
kwargs = store._mongo_client_cls.call_args.kwargs
|
||||
assert "retryWrites" not in kwargs
|
||||
assert "maxPoolSize" not in kwargs
|
||||
assert kwargs["retryReads"] is True
|
||||
@@ -52,10 +52,19 @@ class _MockOpenAI(BaseHTTPRequestHandler):
|
||||
self.wfile.write(data)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _project_root(tmp_path, monkeypatch):
|
||||
# Task output_file writes are contained to the project root (#5770); these
|
||||
# runs write under tmp_path, so make it the root.
|
||||
monkeypatch.setenv("PRAISONAI_PROJECT_ROOT", str(tmp_path))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_endpoint(monkeypatch):
|
||||
def mock_endpoint(monkeypatch, tmp_path):
|
||||
server = HTTPServer(("127.0.0.1", 0), _MockOpenAI)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
# output_file writes are contained to the project root (#5770).
|
||||
monkeypatch.setenv("PRAISONAI_PROJECT_ROOT", str(tmp_path))
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "placeholder")
|
||||
monkeypatch.setenv("OPENAI_BASE_URL", f"http://127.0.0.1:{server.server_port}/v1")
|
||||
monkeypatch.setenv("PRAISONAI_TELEMETRY_DISABLED", "true")
|
||||
@@ -103,7 +112,9 @@ def test_yaml_cli_config_forwards_metrics_flags():
|
||||
assert "metrics" not in config and "metrics_json" not in config
|
||||
|
||||
|
||||
def test_output_file_and_metrics_json(mock_endpoint, tmp_path, capsys):
|
||||
def test_output_file_and_metrics_json(mock_endpoint, tmp_path, capsys, monkeypatch):
|
||||
# output_file writes are contained in the project root (#5770).
|
||||
monkeypatch.setenv("PRAISONAI_PROJECT_ROOT", str(tmp_path))
|
||||
result = _run(_config(tmp_path), {"metrics_json": True, "trust": True})
|
||||
|
||||
assert "FINAL TEXT" in result
|
||||
|
||||
Reference in new issue
Block a user