Compare commits

..
Author SHA1 Message Date
praisonai-triage-agent[bot]andMervin Praison 76a64958b1 fix: honor render page timeout and prevent token leak via exception chain
- fetch_rendered_html now extends the HTTP client deadline to cover the
  requested page_timeout (+30s buffer) so slow synchronous renders are not
  aborted early by the default 60s client timeout.
- _request raises the sanitized MrScraperError outside the except block so the
  original token-bearing httpx exception is not attached via __context__/__cause__.

Co-authored-by: Mervin Praison <[email protected]>
2026-09-29 07:03:28 +00:00
riandradiva 86b97985c5 Add MrScraper client and PraisonAI callable tools example 2026-09-29 13:55:31 +07:00
292 changed files with 2514 additions and 23821 deletions

No files matched your search

+1 -29
View File
@@ -16,9 +16,6 @@ const PR_MAX_AUTO_ADDITIONS = 800;
const PR_MAX_AUTO_FILES = 30;
const MANUAL_ONLY_LABELS = new Set(['security', 'breaking-change', 'needs-manual-review', 'release']);
const WORKFLOW_ONLY_LABEL = 'merge-gate-ci-only';
const EXTERNAL_EXAMPLE_PATH_PREFIX = 'examples/tools/external/';
const EXTERNAL_EXAMPLE_LABEL = 'pipeline/external-example';
const MAINTAINER_ACCEPT_EXAMPLE_LABEL = 'maintainer-accept-example';
const CI_ONLY_PATH_PREFIXES = ['.github/workflows/', '.github/actions/', '.github/scripts/merge-gate'];
const SDK_PATH_PREFIXES = ['src/praisonai-agents/', 'src/praisonai/', 'src/praisonai-ts/'];
const TS_SDK_SOURCE_PREFIX = 'src/praisonai-ts/src/';
@@ -331,7 +328,7 @@ function finalClaudeCompletedOnSha(comments, headPushedAt) {
}
const FINAL_CLAUDE_REVIEW_BODY =
'@claude You are the FINAL architecture reviewer. If the branch is under MervinPraison/PraisonAI (not a fork), you are able to make modifications to this branch and push directly. SCOPE: Review changes in this PR. Python SDK: praisonaiagents, praisonai. TypeScript SDK: src/praisonai-ts/. Do NOT modify src/praisonai-rust. Read ALL comments above from Gemini, Qodo, CodeRabbit, and Copilot carefully before responding.\n\n**MANDATORY READ (before reviewing):**\n- Always read src/praisonai-agents/AGENTS.md\n- If this PR touches src/praisonai-ts/, also read src/praisonai-ts/AGENTS.md §2.1.2 (TS triage + PR review checklist)\n\n**Phase 1: Review per AGENTS.md**\n1. Protocol-driven: check heavy implementations vs core SDK\n2. Backward compatible: ensure zero feature regressions\n3. Performance: no hot-path regressions\n4. SDK value: review in depth whether the change genuinely adds value to the SDK — never add features for the sake of adding them. It must strengthen the SDK (simpler, more user-friendly, robust, world-class, secure). If it does not clearly add value, request changes or recommend rejecting/closing rather than merging scope creep\n5. Do not bloat the Agent class with additional params — only if absolutely required; we already support many params.\n6. Repo routing: agent-callable tools → PraisonAI-Tools; lifecycle plugins → PraisonAI-Plugins; optional sandbox backends → PraisonAI-Plugins (`praisonai.sandbox` entry point) — request changes if wrongly added to praisonaiagents/\n6b. **External tool examples (`examples/tools/external/**`):** If the diff is confined to that path, mark **SDK value** ❌ for monorepo merge — redirect the contributor to PraisonAI-Tools; do **not** Approve for merge unless the PR has the `maintainer-accept-example` label.\n\n**MANDATORY COMMENT FORMAT — include this Phase 1 table in your review comment:**\n#### Phase 1 — AGENTS.md review\n| Check | Result |\n|---|---|\n| Protocol-driven / no heavy impl in core | ✅ or ❌ + one-line rationale |\n| Backward compatible | ✅ or ❌ + one-line rationale |\n| Performance (hot path) | ✅ or ❌ + one-line rationale |\n| **SDK value** | ✅ or ❌ + one-line rationale (explicitly judge whether the change strengthens the SDK) |\n| No Agent param bloat | ✅ or ❌ + one-line rationale |\n| Repo routing | ✅ or ❌ + one-line rationale |\n\nFor TypeScript PRs (src/praisonai-ts/), also add:\n| TS types / parity / tests | ✅ or ❌ + one-line rationale (npm run build && npm test) |\n\n**Phase 2: FIX Valid Issues**\n7. For any VALID bugs or architectural flaws found by Gemini, CodeRabbit, Qodo, Copilot, or any other reviewer: implement the fix\n8. Also independently identify and fix any gaps or issues you find in the changed code — do not rely only on prior reviewer feedback\n9. Push all code fixes directly to THIS branch (do NOT create a new PR)\n10. Comment a summary of exact files modified and what you skipped\n\n**Phase 3: Final Verdict**\n11. If all issues are resolved, approve the PR / close the Issue\n12. If blocking issues remain, request changes / leave clear action items';
'@claude You are the FINAL architecture reviewer. If the branch is under MervinPraison/PraisonAI (not a fork), you are able to make modifications to this branch and push directly. SCOPE: Review changes in this PR. Python SDK: praisonaiagents, praisonai. TypeScript SDK: src/praisonai-ts/. Do NOT modify src/praisonai-rust. Read ALL comments above from Gemini, Qodo, CodeRabbit, and Copilot carefully before responding.\n\n**MANDATORY READ (before reviewing):**\n- Always read src/praisonai-agents/AGENTS.md\n- If this PR touches src/praisonai-ts/, also read src/praisonai-ts/AGENTS.md §2.1.2 (TS triage + PR review checklist)\n\n**Phase 1: Review per AGENTS.md**\n1. Protocol-driven: check heavy implementations vs core SDK\n2. Backward compatible: ensure zero feature regressions\n3. Performance: no hot-path regressions\n4. SDK value: review in depth whether the change genuinely adds value to the SDK — never add features for the sake of adding them. It must strengthen the SDK (simpler, more user-friendly, robust, world-class, secure). If it does not clearly add value, request changes or recommend rejecting/closing rather than merging scope creep\n5. Do not bloat the Agent class with additional params — only if absolutely required; we already support many params.\n6. Repo routing: agent-callable tools → PraisonAI-Tools; lifecycle plugins → PraisonAI-Plugins; optional sandbox backends → PraisonAI-Plugins (`praisonai.sandbox` entry point) — request changes if wrongly added to praisonaiagents/\n\n**MANDATORY COMMENT FORMAT — include this Phase 1 table in your review comment:**\n#### Phase 1 — AGENTS.md review\n| Check | Result |\n|---|---|\n| Protocol-driven / no heavy impl in core | ✅ or ❌ + one-line rationale |\n| Backward compatible | ✅ or ❌ + one-line rationale |\n| Performance (hot path) | ✅ or ❌ + one-line rationale |\n| **SDK value** | ✅ or ❌ + one-line rationale (explicitly judge whether the change strengthens the SDK) |\n| No Agent param bloat | ✅ or ❌ + one-line rationale |\n| Repo routing | ✅ or ❌ + one-line rationale |\n\nFor TypeScript PRs (src/praisonai-ts/), also add:\n| TS types / parity / tests | ✅ or ❌ + one-line rationale (npm run build && npm test) |\n\n**Phase 2: FIX Valid Issues**\n7. For any VALID bugs or architectural flaws found by Gemini, CodeRabbit, Qodo, Copilot, or any other reviewer: implement the fix\n8. Also independently identify and fix any gaps or issues you find in the changed code — do not rely only on prior reviewer feedback\n9. Push all code fixes directly to THIS branch (do NOT create a new PR)\n10. Comment a summary of exact files modified and what you skipped\n\n**Phase 3: Final Verdict**\n11. If all issues are resolved, approve the PR / close the Issue\n12. If blocking issues remain, request changes / leave clear action items';
async function getMergeState(github, owner, repo, prNumber) {
const query = `
@@ -586,11 +583,6 @@ function isCiOnlyChange(files) {
return files.every((f) => CI_ONLY_PATH_PREFIXES.some((p) => f.filename.startsWith(p)));
}
function isExternalExampleOnlyChange(files) {
if (!files.length) return false;
return files.every((f) => f.filename.startsWith(EXTERNAL_EXAMPLE_PATH_PREFIX));
}
function isInternalPullRequestLink(link, owner, repo) {
if (link?.base?.repo?.full_name === `${owner}/${repo}`) return true;
const baseUrl = link?.base?.repo?.url || '';
@@ -981,15 +973,6 @@ async function evaluatePipelineQuiescent(github, owner, repo, prNumber, core, op
const pullFiles = await listPullFiles(github, owner, repo, prNumber);
if (
isExternalExampleOnlyChange(pullFiles) &&
!ctx.labels.includes(MAINTAINER_ACCEPT_EXAMPLE_LABEL)
) {
reasons.push(
'external example only — belongs in PraisonAI-Tools (add maintainer-accept-example to opt in)'
);
}
const agentChange = getAgentPyChangeFromFiles(pullFiles);
const agentManual = manualReviewReasonForAgentPy(agentChange);
if (agentManual) reasons.push(agentManual);
@@ -1036,13 +1019,6 @@ async function listPrNumbersForMergeGateScan(github, owner, repo, core) {
}
const mergeReady = mergeReadyIssues
.filter((issue) => issue.pull_request)
.filter((issue) => {
const labels = (issue.labels || []).map((l) => (typeof l === 'string' ? l : l.name));
if (labels.includes(EXTERNAL_EXAMPLE_LABEL) && !labels.includes(MAINTAINER_ACCEPT_EXAMPLE_LABEL)) {
return false;
}
return true;
})
.map((issue) => issue.number)
.sort((a, b) => a - b);
@@ -1170,10 +1146,6 @@ module.exports = {
claudeRunBlocksPr,
hasBlockingClaudeRunForPr,
isCiOnlyChange,
isExternalExampleOnlyChange,
EXTERNAL_EXAMPLE_PATH_PREFIX,
EXTERNAL_EXAMPLE_LABEL,
MAINTAINER_ACCEPT_EXAMPLE_LABEL,
isInternalPullRequestLink,
resolvePrNumberFromLinkedPullRequests,
resolvePrNumberFromHeadBranch,
@@ -73,33 +73,4 @@ assert(
)
);
assert(
'external example only path',
mg.isExternalExampleOnlyChange([{ filename: 'examples/tools/external/arxiv/tool.py' }])
);
assert(
'mixed paths not external-example only',
!mg.isExternalExampleOnlyChange([
{ filename: 'examples/tools/external/foo.py' },
{ filename: 'src/praisonai-agents/foo.py' },
])
);
assert(
'upstream head repo',
ps.isUpstreamHeadRepo(
{ head: { repo: { full_name: 'MervinPraison/PraisonAI' } } },
'MervinPraison',
'PraisonAI'
)
);
assert(
'fork head repo still synced (not upstream)',
!ps.isUpstreamHeadRepo(
{ head: { repo: { full_name: 'ai-mrscraper/PraisonAI' } } },
'MervinPraison',
'PraisonAI'
)
);
process.exit(failed ? 1 : 0);
+5 -25
View File
@@ -37,17 +37,8 @@ const LABEL_SPECS = [
{ name: 'pipeline/blocked:cooldown', color: 'fef2c0', description: 'Blocked: post-push or @claude cooldown' },
{ name: 'pipeline/blocked:stale-final', color: 'f9d0c4', description: 'Blocked: FINAL stale after new commits' },
{ name: 'pipeline/blocked:no-final', color: 'ededed', description: 'Blocked: no FINAL @claude trigger yet' },
{
name: 'pipeline/external-example',
color: 'bfd4f2',
description: 'Diff only under examples/tools/external/ — not merge-ready without maintainer-accept-example',
},
];
// PRs with pipeline/external-example cannot reach pipeline/merge-ready unless
// maintainers add maintainer-accept-example (enforced in merge-gate evaluatePipelineQuiescent,
// same fail-closed pattern as fork PRs).
function deriveStage(comments, evalResult) {
if (evalResult.ready) return 'pipeline/merge-ready';
if (!chain.chainKickPosted(comments)) return 'pipeline/reviews-pending';
@@ -72,8 +63,7 @@ function reasonToBlockerLabel(reason) {
r.includes('without test') ||
r.includes('files changed') ||
r.includes('possible secret') ||
r.includes('agent.py') ||
r.includes('external example only')
r.includes('agent.py')
) return 'pipeline/blocked:manual-review';
if (r.includes('recent @claude') || r.includes('post-push buffer')) return 'pipeline/blocked:cooldown';
if (r.includes('stale final')) return 'pipeline/blocked:stale-final';
@@ -143,16 +133,11 @@ async function syncPipelineLabels(github, owner, repo, prNumber, core) {
github, owner, repo, prNumber, core
);
const { stage, blockers, all } = computePipelineLabels(ctx.comments, evalResult);
const pullFiles = await mergeGate.listPullFiles(github, owner, repo, prNumber);
const auxiliary = mergeGate.isExternalExampleOnlyChange(pullFiles)
? [mergeGate.EXTERNAL_EXAMPLE_LABEL]
: [];
const labelSet = [...all, ...auxiliary];
const current = ctx.labels.filter((l) => l.startsWith(PIPELINE_PREFIX));
const desired = new Set(labelSet);
const desired = new Set(all);
const toRemove = current.filter((l) => !desired.has(l) && l !== 'pipeline/merged');
const toAdd = labelSet.filter((l) => !current.includes(l));
const toAdd = all.filter((l) => !current.includes(l));
for (const name of toRemove) {
try {
@@ -214,11 +199,6 @@ async function dispatchMergeGateForOldestReady(github, owner, repo, readyCandida
return 0;
}
function isUpstreamHeadRepo(pr, owner, repo) {
const headFull = pr.head?.repo?.full_name;
return !headFull || headFull === `${owner}/${repo}`;
}
async function syncOpenPullRequests(github, owner, repo, options, core) {
const { maxPrs = 20, dispatchMergeGate = true } = options || {};
await ensurePipelineLabels(github, owner, repo, core);
@@ -243,8 +223,9 @@ async function syncOpenPullRequests(github, owner, repo, options, core) {
for (const pr of prs) {
if (synced >= maxPrs) break;
if (pr.draft) continue;
if (pr.head?.repo?.full_name && pr.head.repo.full_name !== `${owner}/${repo}`) continue;
const result = await syncPipelineLabels(github, owner, repo, pr.number, core);
if (result.ready && isUpstreamHeadRepo(pr, owner, repo)) {
if (result.ready) {
readyCandidates.push({
prNumber: pr.number,
createdAt: result.createdAt || new Date(pr.created_at).getTime(),
@@ -267,7 +248,6 @@ module.exports = {
STAGE_LABELS,
BLOCKER_LABELS,
ALL_PIPELINE_LABELS,
isUpstreamHeadRepo,
deriveStage,
deriveBlockerLabels,
computePipelineLabels,
+1 -1
View File
@@ -1,6 +1,6 @@
name: Pipeline Status Sync
# Sync pipeline/* stage and blocker labels on open PRs (fork heads included; merge gate only for upstream branches).
# Sync pipeline/* stage and blocker labels on open internal PRs.
on:
schedule:
+7 -21
View File
@@ -34,24 +34,10 @@ jobs:
id: config
env:
HAS_APP_CREDENTIALS: ${{ secrets.PRAISONAI_APP_ID != '' && secrets.PRAISONAI_APP_PRIVATE_KEY != '' }}
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
run: |
echo "configured=$HAS_APP_CREDENTIALS" >> "$GITHUB_OUTPUT"
openai_ready=true
if [ -z "$OPENAI_API_KEY" ] || [ "$OPENAI_API_KEY" = "not-needed" ]; then
openai_ready=false
if [ "$HAS_APP_CREDENTIALS" = "true" ]; then
echo "::notice::Skipping PraisonAI PR review: OPENAI_API_KEY is missing or a placeholder."
fi
fi
if [ "$HAS_APP_CREDENTIALS" = "true" ] && [ "$openai_ready" = "true" ]; then
echo "ready=true" >> "$GITHUB_OUTPUT"
else
echo "ready=false" >> "$GITHUB_OUTPUT"
fi
run: echo "configured=$HAS_APP_CREDENTIALS" >> "$GITHUB_OUTPUT"
- name: Determine checkout ref
if: steps.config.outputs.ready == 'true'
if: steps.config.outputs.configured == 'true'
id: dest
env:
ISSUE_NUMBER: ${{ github.event.issue.number }}
@@ -67,7 +53,7 @@ jobs:
fi
- name: Checkout Repository
if: steps.config.outputs.ready == 'true'
if: steps.config.outputs.configured == 'true'
uses: actions/checkout@v4
with:
persist-credentials: false
@@ -75,7 +61,7 @@ jobs:
fetch-depth: 0
- name: Generate GitHub App Token
if: steps.config.outputs.ready == 'true'
if: steps.config.outputs.configured == 'true'
id: generate_token
uses: tibdex/github-app-token@v2
with:
@@ -83,17 +69,17 @@ jobs:
private_key: ${{ secrets.PRAISONAI_APP_PRIVATE_KEY }}
- name: Set up Python
if: steps.config.outputs.ready == 'true'
if: steps.config.outputs.configured == 'true'
uses: actions/setup-python@v5
with:
python-version: '3.11'
- name: Install PraisonAI
if: steps.config.outputs.ready == 'true'
if: steps.config.outputs.configured == 'true'
run: pip install "praisonai[all]"
- name: Run PraisonAI PR Review
if: steps.config.outputs.ready == 'true'
if: steps.config.outputs.configured == 'true'
env:
GITHUB_TOKEN: ${{ steps.generate_token.outputs.token }}
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+1 -3
View File
@@ -773,9 +773,7 @@ jobs:
env:
# Bump deliberately: this is the version everything in this job was
# verified against. Latest at time of writing is v0.33.3.
# No leading "v": install.sh passes ?version=$OLLAMA_VERSION and ollama.com
# redirects to GitHub tag v$version (vv0.33.2 → 404 if set to v0.33.2).
OLLAMA_VERSION: "0.33.2"
OLLAMA_VERSION: v0.33.2
# 522 MB, capabilities: completion, tools, thinking.
OLLAMA_TEST_MODEL: qwen3:0.6b
-1
View File
@@ -122,4 +122,3 @@ jobs/*
# Local Android SDK/NDK/JDK for building praisonai-mobile on a laptop without
# Android Studio (how: .github/workflows/mobile-platforms.yml). Several GB.
.android-sdk/
.worktrees/
+2 -3
View File
@@ -776,7 +776,7 @@ Methods:
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">wait_for_completions</a>(timeout: Optional[float] = None, agent_ids: Optional[List[str]] = None) -> List[SubAgentCompletionEvent]</code>
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">where_does_it_run</a>() -> str</code>
* <code title="class AutoApproveBackend">AutoApproveBackend.<a href="./src/praisonai-agents/praisonaiagents/approval/backends.py">request_approval</a>(request: ApprovalRequest) -> ApprovalDecision</code>
* <code title="class AutoApproveBackend">AutoApproveBackend.<a href="./src/praisonai-agents/praisonaiagents/approval/backends.py">request_approval_sync</a>(request: ApprovalRequest, timeout: Optional[float] = None) -> ApprovalDecision</code>
* <code title="class AutoApproveBackend">AutoApproveBackend.<a href="./src/praisonai-agents/praisonaiagents/approval/backends.py">request_approval_sync</a>(request: ApprovalRequest) -> ApprovalDecision</code>
* <code title="class RunOutcome">RunOutcome.<a href="./src/praisonai-agents/praisonaiagents/agent/run_outcome.py">completed</a>(output: Optional[str] = None) -> 'RunOutcome'</code>
* <code title="class RunOutcome">RunOutcome.<a href="./src/praisonai-agents/praisonaiagents/agent/run_outcome.py">from_exception</a>(exc: BaseException, output: Optional[str] = None) -> 'RunOutcome'</code>
* <code title="class RunOutcome">RunOutcome.<a href="./src/praisonai-agents/praisonaiagents/agent/run_outcome.py">succeeded</a>() -> bool</code>
@@ -1018,8 +1018,7 @@ export { // Bot types
FailoverManager, // Enums
SandboxStatus, // Gateway types
type GatewayConfig, // Other types
type ProviderStatus, // Session-projection reducer (Issue #5324)
SessionProjection, AutonomyLevel, GatewayEventType, RagRetrievalPolicy } from "./gateway";
type ProviderStatus, AutonomyLevel, GatewayEventType, RagRetrievalPolicy } from "./gateway";
export { LLMGuardrail, createLLMGuardrail } from "./guardrails/llm-guardrail";
export { DisplayTypes, HooksManager, WorkflowHooksExecutor, clearAllCallbacks, clearApprovalCallback, createHooksManager, createLoggingOperationHooks, createLoggingWorkflowHooks, createTimingWorkflowHooks, createValidationOperationHooks, createWorkflowHooks, executeCallback, executeSyncCallback, getRegisteredDisplayTypes, hasApprovalCallback, registerApprovalCallback, registerDisplayCallback, requestApproval, unregisterDisplayCallback } from "./hooks";
export { // Computer Use
+1 -1
View File
@@ -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.11" \
"praisonai>=4.7.10" \
"praisonai[chat]" \
"embedchain[github,youtube]"
+1 -1
View File
@@ -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.11" \
"praisonai>=4.7.10" \
"praisonai[ui]" \
"praisonai[chat]" \
"praisonai[realtime]" \
+1 -1
View File
@@ -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.11" \
"praisonai>=4.7.10" \
"praisonai[ui]" \
"praisonai[crewai]"
@@ -0,0 +1,6 @@
"""MrScraper example client and agent callable tools."""
from .client import MrScraperClient, MrScraperError
from .tools import make_mrscraper_tools
__all__ = ["MrScraperClient", "MrScraperError", "make_mrscraper_tools"]
@@ -0,0 +1,276 @@
"""Small synchronous client for the MrScraper n8n operation contract.
The platform's response envelopes are passed through because they vary by API.
"""
from __future__ import annotations
import json
import os
from importlib import resources
from typing import Any, Dict, List, Mapping, Optional, Union
from urllib.parse import quote
import httpx
Json = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
PLATFORM = "https://api.app.mrscraper.com"
SERP = "https://sync.scraper.mrscraper.com"
RENDER = "https://api.mrscraper.com"
class MrScraperError(RuntimeError):
"""Sanitized HTTP or transport failure; never includes a token-bearing URL."""
def __init__(self, message: str, status_code: Optional[int] = None) -> None:
super().__init__(message)
self.status_code = status_code
def _present(values: Mapping[str, Any]) -> Dict[str, Any]:
return {key: value for key, value in values.items() if value is not None}
def _patterns(values: Optional[List[str]]) -> Optional[str]:
return "|".join(values) if values is not None else None
class MrScraperClient:
"""Call MrScraper with a token from the environment or constructor.
``http_timeout`` controls the Python request deadline. API ``timeout``
options control page loading and are sent independently.
"""
def __init__(self, token: Optional[str] = None, *, http_timeout: float = 60.0,
transport: Optional[httpx.BaseTransport] = None) -> None:
self._token = token or os.environ.get("MRSCRAPER_API_TOKEN")
if not self._token:
raise ValueError("Set MRSCRAPER_API_TOKEN or pass token")
self._http = httpx.Client(timeout=http_timeout, transport=transport)
def close(self) -> None:
self._http.close()
def __enter__(self) -> "MrScraperClient":
return self
def __exit__(self, *args: object) -> None:
self.close()
def _request(self, method: str, host: str, path: str, *,
params: Optional[Mapping[str, Any]] = None,
body: Optional[Mapping[str, Any]] = None,
timeout: Optional[float] = None) -> Json:
headers = {"Accept": "application/json"}
if body is not None:
headers["Content-Type"] = "application/json"
if host == SERP:
headers["Authorization"] = f"Bearer {self._token}"
# The n8n credential adds this header to every route as well.
headers["x-api-token"] = self._token
else:
headers["x-api-token"] = self._token
request_kwargs: Dict[str, Any] = {"headers": headers, "params": params, "json": body}
if timeout is not None:
request_kwargs["timeout"] = timeout
error: Optional[MrScraperError] = None
try:
response = self._http.request(method, host + path, **request_kwargs)
response.raise_for_status()
if not response.content:
return None
if "json" in response.headers.get("Content-Type", ""):
return response.json()
return response.text
except httpx.HTTPStatusError as exc:
status = exc.response.status_code
error = MrScraperError(f"MrScraper HTTP {status} on {method} {host}{path}", status)
except (httpx.RequestError, ValueError) as exc:
# httpx errors can contain the rendered URL and its token query.
error = MrScraperError(f"MrScraper request failed on {method} {host}{path}: {type(exc).__name__}")
# Raise outside the except block so the original token-bearing exception
# is not attached via __context__/__cause__ and cannot leak the token.
raise error
def get_account_info(self) -> Json:
"""Get subscription account details from the platform host."""
return self._request("GET", PLATFORM, "/api/v1/subscription-accounts")
def create_prompt_scraper(self, url: str, message: str, *, mode: str = "Super",
proxy_country: Optional[str] = None,
output_schema: Optional[Mapping[str, Any]] = None) -> Json:
"""Create a general AI scraper; append an optional expected JSON shape to its prompt."""
if mode not in ("Super", "Cheap"):
raise ValueError("mode must be Super or Cheap")
if output_schema is not None:
message += "\n\nReturn the output as JSON matching this schema:\n" + json.dumps(output_schema)
return self._request("POST", PLATFORM, "/api/v1/scrapers-ai", body=_present({
"graph": "general", "url": url, "message": message, "mode": mode,
"proxyCountry": proxy_country}))
def extract_page_by_prompt(self, url: str, message: str, *, mode: str = "Super",
proxy_country: Optional[str] = None,
output_schema: Optional[Mapping[str, Any]] = None) -> Json:
"""Create and start the same general scraper used by prompt-based creation."""
return self.create_prompt_scraper(url, message, mode=mode,
proxy_country=proxy_country,
output_schema=output_schema)
def create_listing_scraper(self, url: str, message: str, *, max_pages: int = 1,
proxy_country: Optional[str] = None,
item_schema: Optional[Mapping[str, Any]] = None) -> Json:
"""Create a listing scraper with pagination and optional item shape."""
if item_schema is not None:
message += "\n\nReturn each item as JSON matching this schema:\n" + json.dumps(item_schema)
return self._request("POST", PLATFORM, "/api/v1/scrapers-ai", body=_present({
"graph": "listing", "url": url, "message": message,
"maxPages": max_pages, "proxyCountry": proxy_country}))
def extract_listings(self, url: str, message: str, *, max_pages: int = 1,
proxy_country: Optional[str] = None,
item_schema: Optional[Mapping[str, Any]] = None) -> Json:
"""Create and start the same listing scraper used by listing creation."""
return self.create_listing_scraper(url, message, max_pages=max_pages,
proxy_country=proxy_country,
item_schema=item_schema)
def create_crawl_scraper(self, url: str, *, max_depth: int = 2,
max_pages: int = 50, limit: int = 50,
include_patterns: Optional[List[str]] = None,
exclude_patterns: Optional[List[str]] = None) -> Json:
"""Create a map scraper that discovers URLs from a starting page."""
return self._request("POST", PLATFORM, "/api/v1/scrapers-ai", body=_present({
"graph": "map", "url": url, "maxDepth": max_depth,
"maxPages": max_pages, "limit": limit,
"includePatterns": _patterns(include_patterns),
"excludePatterns": _patterns(exclude_patterns)}))
def crawl_website_urls(self, url: str, *, max_depth: int = 2,
max_pages: int = 50, limit: int = 50,
include_patterns: Optional[List[str]] = None,
exclude_patterns: Optional[List[str]] = None) -> Json:
"""Create and start the same map scraper used by crawl creation."""
return self.create_crawl_scraper(url, max_depth=max_depth,
max_pages=max_pages, limit=limit,
include_patterns=include_patterns,
exclude_patterns=exclude_patterns)
def extract_structured_data(self, url: str, category: str, *, mode: str = "Super",
proxy_country: Optional[str] = None) -> Json:
"""Create a general scraper with a local n8n preset prompt for ``category``."""
prompt_path = resources.files("praisonai_mrscraper").joinpath("structured_data_prompts.json")
prompts = json.loads(prompt_path.read_text(encoding="utf-8"))
if category not in prompts:
raise ValueError(f"Unknown category; choose from {', '.join(sorted(prompts))}")
return self.create_prompt_scraper(url, prompts[category], mode=mode,
proxy_country=proxy_country)
def search_google_serp(self, query: str, *, region: str = "us", language: str = "en",
page: int = 1, format: str = "json",
render_js: bool = False) -> Json:
"""Search Google with MrScraper's synchronous SERP endpoint."""
if format not in ("json", "html"):
raise ValueError("format must be json or html")
return self._request("POST", SERP, "/api/google/serp/v2/sync", body={
"query": query, "region": region, "language": language,
"page": page, "format": format, "renderJs": render_js})
def fetch_rendered_html(self, url: str, *, max_retries: int = 3,
page_timeout: int = 300, geo_code: str = "us",
html: bool = True, markdown: bool = False,
screenshot: Optional[str] = None, proxy_country: str = "us",
block_resources: Optional[bool] = None,
return_cookie: Optional[bool] = None,
super_mode: Optional[bool] = None,
wait_for_selector: Optional[str] = None,
wait_until: Optional[str] = None,
home_page: Optional[str] = None,
token_cap: Optional[int] = None) -> Json:
"""Fetch a rendered page; page controls are distinct from HTTP timeout."""
if screenshot is not None and screenshot not in ("full", "top"):
raise ValueError("screenshot must be full or top")
if wait_until is not None and wait_until not in ("domcontentloaded", "load", "networkidle"):
raise ValueError("unsupported wait_until")
query = _present({"token": self._token, "browserRendering": True,
"timeout": page_timeout, "geoCode": geo_code,
"html": html, "markdown": markdown, "screenshot": screenshot,
"proxyCountry": proxy_country, "blockResources": block_resources,
"returnCookie": return_cookie, "super": super_mode,
"waitForSelector": wait_for_selector, "waitUntil": wait_until})
query = {key: str(value).lower() if isinstance(value, bool) else value
for key, value in query.items()}
# Give the synchronous render the page timeout it requested, plus a
# buffer for network transfer, so the HTTP client does not abort early.
http_timeout = float(page_timeout) + 30.0
return self._request("POST", RENDER, "/", params=query, timeout=http_timeout,
body=_present({
"url": url, "maxRetries": max_retries, "homePage": home_page,
"tokenCap": token_cap}))
def get_results(self, scraper_id: str, *, page: int = 1, page_size: int = 10,
sort: str = "createdAt", sort_order: str = "DESC") -> Json:
"""List results for a scraper with URL-encoded query parameters."""
return self._request("GET", PLATFORM, "/api/v1/results", params={
"filters[scraperId]": scraper_id, "page": page, "pageSize": page_size,
"sort": sort, "sortOrder": sort_order})
def get_latest_results(self, scraper_id: str, *, page_size: int = 10) -> Json:
"""List newest results for a scraper."""
return self.get_results(scraper_id, page=1, page_size=page_size,
sort="createdAt", sort_order="DESC")
def get_result_detail(self, result_id: str) -> Json:
"""Fetch one result by its URL-encoded ID."""
return self._request("GET", PLATFORM,
"/api/v1/results/" + quote(result_id, safe=""))
def run_existing_scraper(self, scraper_id: str, url: str, *,
scraper_type: str = "ai", agent_type: str = "general",
max_retry: int = 3, proxy_country: Optional[str] = None,
options: Optional[Mapping[str, Any]] = None) -> Json:
"""Rerun an AI or manual scraper; options are validated per scraper type."""
if scraper_type not in ("ai", "manual"):
raise ValueError("scraper_type must be ai or manual")
if agent_type not in ("general", "listing", "map"):
raise ValueError("agent_type must be general, listing, or map")
common = {"bypassProxy", "html", "markdown", "screenshot"}
general = common | {"renderJavascript", "returnCookies", "useHomePage", "waitForSelector"}
listing = general | {"maxPages", "timeout", "stream"}
mapping = general if agent_type == "general" else listing if agent_type == "listing" else {
"maxDepth", "maxPages", "limit", "includePatterns", "excludePatterns"}
manual = {"bypassProxy", "cookieJar", "cookies", "homePage", "homePageTimeout",
"html", "markdown", "paginator", "proxy", "record", "returnCookie",
"screenshot", "stream", "timeout", "tokenCap"}
provided = dict(options or {})
unknown = provided.keys() - (manual if scraper_type == "manual" else mapping)
if unknown:
raise ValueError(f"Unsupported rerun options: {', '.join(sorted(unknown))}")
if scraper_type == "ai" and agent_type == "listing":
provided = {"maxPages": 5, "timeout": 300, **provided}
if scraper_type == "ai" and agent_type == "map":
provided = {"maxDepth": 2, "maxPages": 50, "limit": 50, **provided}
for key in ("includePatterns", "excludePatterns"):
if isinstance(provided.get(key), list):
provided[key] = _patterns(provided[key])
if scraper_type == "manual" and "screenshot" in provided:
if not isinstance(provided["screenshot"], bool):
raise ValueError("manual screenshot must be a boolean")
provided["screenshot"] = str(provided["screenshot"]).lower()
route = "/api/v1/scrapers-ai-rerun" if scraper_type == "ai" else "/api/v1/scrapers-manual-rerun"
return self._request("POST", PLATFORM, route, body={**_present({
"scraperId": scraper_id, "url": url, "maxRetry": max_retry,
"proxyCountry": proxy_country}), **provided})
def run_existing_scraper_batch(self, scraper_id: str, urls: List[str], *,
scraper_type: str = "ai") -> Json:
"""Rerun an AI or manual scraper for a nonempty list of URLs."""
if scraper_type not in ("ai", "manual"):
raise ValueError("scraper_type must be ai or manual")
if not isinstance(urls, list) or not urls or any(
not isinstance(url, str) or not url for url in urls
):
raise ValueError("urls must be a nonempty list of URL strings")
route = "/api/v1/scrapers-ai-rerun/bulk" if scraper_type == "ai" else "/api/v1/scrapers-manual-rerun/bulk"
return self._request("POST", PLATFORM, route, body={"scraperId": scraper_id,
"urls": urls})
@@ -0,0 +1 @@
{"article":"\nPlease extract the data by following the schema:\n{\n \"fields\": [\n {\n \"name\": \"headline\",\n \"type\": \"string\",\n \"description\": \"Article headline or title\"\n },\n {\n \"name\": \"articleBody\",\n \"type\": \"string\",\n \"description\": \"Full text content of the article\"\n },\n {\n \"name\": \"articleBodyHtml\",\n \"type\": \"string\",\n \"description\": \"HTML markup of the article body\"\n },\n {\n \"name\": \"description\",\n \"type\": \"string\",\n \"description\": \"Short summary or description of the article\"\n },\n {\n \"name\": \"datePublished\",\n \"type\": \"string\",\n \"description\": \"Publication date in ISO 8601 format\"\n },\n {\n \"name\": \"datePublishedRaw\",\n \"type\": \"string\",\n \"description\": \"Publication date as displayed on the page\"\n },\n {\n \"name\": \"dateModified\",\n \"type\": \"string\",\n \"description\": \"Last modified date in ISO 8601 format\"\n },\n {\n \"name\": \"dateModifiedRaw\",\n \"type\": \"string\",\n \"description\": \"Last modified date as displayed on the page\"\n },\n {\n \"name\": \"authors\",\n \"type\": \"array\",\n \"item_type\": \"object\",\n \"description\": \"List of article authors\",\n \"item_structure\": {\n \"name\": {\n \"type\": \"string\",\n \"description\": \"Parsed individual author name\"\n },\n \"nameRaw\": {\n \"type\": \"string\",\n \"description\": \"Raw author name as displayed on the page\"\n }\n }\n },\n {\n \"name\": \"inLanguage\",\n \"type\": \"string\",\n \"description\": \"Language code of the article (e.g., en)\"\n },\n {\n \"name\": \"breadcrumbs\",\n \"type\": \"array\",\n \"item_type\": \"object\",\n \"description\": \"Navigation breadcrumb trail\",\n \"item_structure\": {\n \"url\": {\n \"type\": \"string\",\n \"description\": \"Breadcrumb link URL\"\n },\n \"name\": {\n \"type\": \"string\",\n \"description\": \"Breadcrumb label\"\n }\n }\n },\n {\n \"name\": \"mainImage\",\n \"type\": \"object\",\n \"structure\": {\n \"url\": {\n \"type\": \"string\",\n \"description\": \"URL of the main image\"\n }\n },\n \"description\": \"Primary image of the article\"\n },\n {\n \"name\": \"images\",\n \"type\": \"array\",\n \"item_type\": \"object\",\n \"description\": \"All images found in the article\",\n \"item_structure\": {\n \"url\": {\n \"type\": \"string\",\n \"description\": \"Image URL\"\n }\n }\n },\n {\n \"name\": \"videos\",\n \"type\": \"array\",\n \"item_type\": \"object\",\n \"description\": \"All videos found in the article\",\n \"item_structure\": {\n \"url\": {\n \"type\": \"string\",\n \"description\": \"Video URL\"\n }\n }\n },\n {\n \"name\": \"audios\",\n \"type\": \"array\",\n \"item_type\": \"object\",\n \"description\": \"All audio files found in the article\",\n \"item_structure\": {\n \"url\": {\n \"type\": \"string\",\n \"description\": \"Audio URL\"\n }\n }\n },\n {\n \"name\": \"url\",\n \"type\": \"string\",\n \"description\": \"URL of the article page\"\n },\n {\n \"name\": \"canonicalUrl\",\n \"type\": \"string\",\n \"description\": \"Canonical URL of the article\"\n }\n ]\n}\n","forumThread":"\nPlease extract the data by following the schema:\n{\n \"fields\": [\n {\n \"name\": \"topic\",\n \"type\": \"object\",\n \"structure\": {\n \"name\": {\n \"type\": \"string\",\n \"description\": \"Topic title or question\"\n }\n },\n \"description\": \"Forum thread topic\"\n },\n {\n \"name\": \"posts\",\n \"type\": \"array\",\n \"item_type\": \"object\",\n \"description\": \"List of posts in the thread\",\n \"item_structure\": {\n \"text\": {\n \"type\": \"string\",\n \"description\": \"Post text content\"\n },\n \"reactions\": {\n \"type\": \"object\",\n \"structure\": {\n \"likes\": {\n \"type\": \"number\",\n \"description\": \"Number of likes\"\n },\n \"replies\": {\n \"type\": \"number\",\n \"description\": \"Number of replies\"\n }\n },\n \"description\": \"Reaction counts for the post\"\n },\n \"datePublished\": {\n \"type\": \"string\",\n \"description\": \"Post date in ISO 8601 format\Line truncated
@@ -0,0 +1,47 @@
"""Bounded, agent-callable wrappers around MrScraperClient."""
from __future__ import annotations
import json
from typing import Callable, List
from .client import MrScraperClient
def make_mrscraper_tools(client: MrScraperClient, *, max_chars: int = 12000) -> List[Callable[..., str]]:
"""Return callable tools for rendered pages, extraction, crawl, and search.
Screenshot and cookie fields are deliberately not sent to an LLM context.
"""
if max_chars <= 0:
raise ValueError("max_chars must be positive")
def _without_sensitive(value: object) -> object:
if isinstance(value, dict):
return {key: _without_sensitive(item) for key, item in value.items()
if "screenshot" not in key.lower() and "cookie" not in key.lower()}
if isinstance(value, list):
return [_without_sensitive(item) for item in value]
return value
def _bounded(value: object) -> str:
value = _without_sensitive(value)
raw = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
return raw[:max_chars]
def read_web_page(url: str) -> str:
"""Read rendered Markdown from a public web page for a short summary."""
return _bounded(client.fetch_rendered_html(url, html=False, markdown=True))
def extract_page(url: str, prompt: str) -> str:
"""Ask MrScraper to extract information from one page using a prompt."""
return _bounded(client.extract_page_by_prompt(url, prompt))
def crawl_urls(url: str) -> str:
"""Discover URLs linked from a starting website."""
return _bounded(client.crawl_website_urls(url))
def search_google(query: str) -> str:
"""Search Google SERP with MrScraper and return bounded results."""
return _bounded(client.search_google_serp(query))
return [read_web_page, extract_page, crawl_urls, search_google]
+13
View File
@@ -0,0 +1,13 @@
[build-system]
requires = ["setuptools>=68"]
build-backend = "setuptools.build_meta"
[project]
name = "praisonai-mrscraper-example"
version = "0.1.0"
description = "MrScraper client and PraisonAI callable tools example"
requires-python = ">=3.9"
dependencies = ["httpx>=0.25,<1"]
[tool.setuptools.package-data]
praisonai_mrscraper = ["structured_data_prompts.json"]
+135
View File
@@ -0,0 +1,135 @@
"""Contract tests for the n8n routes, without paid API calls."""
import json
import httpx
import pytest
from praisonai_mrscraper import MrScraperClient, MrScraperError, make_mrscraper_tools
@pytest.fixture
def calls():
seen = []
def respond(request):
seen.append(request)
return httpx.Response(200, json={"data": {"sample": True}})
return seen, httpx.MockTransport(respond)
def test_hosts_auth_and_verbs(calls):
seen, transport = calls
with MrScraperClient("private-token", transport=transport) as client:
client.get_account_info()
client.search_google_serp("example")
client.fetch_rendered_html("https://example.com")
account, serp, render = seen
assert (account.method, account.url.host, account.url.path) == (
"GET", "api.app.mrscraper.com", "/api/v1/subscription-accounts")
assert account.headers["x-api-token"] == "private-token"
assert (serp.method, serp.url.host, serp.url.path) == (
"POST", "sync.scraper.mrscraper.com", "/api/google/serp/v2/sync")
assert serp.headers["Authorization"] == "Bearer private-token"
assert serp.headers["x-api-token"] == "private-token"
assert json.loads(serp.content)["renderJs"] is False
assert (render.method, render.url.host, render.url.path) == (
"POST", "api.mrscraper.com", "/")
assert render.url.params["token"] == "private-token"
assert render.url.params["browserRendering"] == "true"
assert json.loads(render.content) == {"url": "https://example.com", "maxRetries": 3}
def test_create_results_and_reruns(calls):
seen, transport = calls
with MrScraperClient("secret", transport=transport) as client:
result = client.create_prompt_scraper("https://example.com", "Extract title",
output_schema={"title": "string"})
client.create_listing_scraper("https://example.com", "Extract links")
client.crawl_website_urls("https://example.com", include_patterns=["/a", "/b"])
client.get_results("id/a+b", page=2)
client.get_result_detail("result/a+b")
client.run_existing_scraper("s", "https://example.com", agent_type="listing",
options={"stream": False})
client.run_existing_scraper("s", "https://example.com", scraper_type="manual",
options={"screenshot": True, "cookies": []})
client.run_existing_scraper_batch("s", ["https://example.com"], scraper_type="manual")
assert result == {"data": {"sample": True}}
assert json.loads(seen[0].content)["graph"] == "general"
assert "Return the output as JSON" in json.loads(seen[0].content)["message"]
assert json.loads(seen[1].content)["maxPages"] == 1
assert json.loads(seen[2].content)["includePatterns"] == "/a|/b"
assert seen[3].url.params["filters[scraperId]"] == "id/a+b"
assert seen[3].url.params["page"] == "2"
assert seen[4].url.raw_path.startswith(b"/api/v1/results/result%2Fa%2Bb")
assert json.loads(seen[5].content)["stream"] is False
assert json.loads(seen[5].content)["maxPages"] == 5
assert seen[6].url.path == "/api/v1/scrapers-manual-rerun"
assert json.loads(seen[6].content)["screenshot"] == "true"
assert json.loads(seen[7].content)["urls"] == ["https://example.com"]
def test_render_query_and_errors_are_sanitized():
def fail(request):
return httpx.Response(401, json={"error": "private-token"})
with MrScraperClient("private-token", transport=httpx.MockTransport(fail)) as client:
with pytest.raises(MrScraperError) as error:
client.fetch_rendered_html("https://example.com", html=False,
markdown=True, screenshot="top", token_cap=20)
assert error.value.status_code == 401
assert "private-token" not in str(error.value)
assert "token=" not in str(error.value)
# The token-bearing request must not leak through the exception chain.
assert error.value.__context__ is None
assert error.value.__cause__ is None
def test_render_extends_http_timeout_to_page_timeout():
seen = []
def respond(request):
seen.append(request)
return httpx.Response(200, json={"data": {"markdown": "ok"}})
with MrScraperClient("secret", http_timeout=60.0,
transport=httpx.MockTransport(respond)) as client:
client.fetch_rendered_html("https://example.com", page_timeout=300)
assert seen[0].extensions["timeout"]["read"] == pytest.approx(330.0)
def test_structured_preset_and_bounded_tools(calls):
seen, transport = calls
with MrScraperClient("secret", transport=transport) as client:
client.extract_structured_data("https://example.com", "article")
tools = make_mrscraper_tools(client, max_chars=25)
assert len(tools[0]("https://example.com")) <= 25
assert "headline" in json.loads(seen[0].content)["message"]
assert seen[1].url.params["markdown"] == "true"
assert seen[1].url.params["html"] == "false"
def test_invalid_options_do_not_make_requests(calls):
seen, transport = calls
with MrScraperClient("secret", transport=transport) as client:
with pytest.raises(ValueError):
client.run_existing_scraper_batch("s", [])
with pytest.raises(ValueError):
client.run_existing_scraper("s", "https://example.com", options={"bad": True})
with pytest.raises(ValueError):
client.extract_structured_data("https://example.com", "unknown")
assert seen == []
def test_agent_tool_removes_nested_image_and_cookie_fields():
def respond(request):
return httpx.Response(200, json={"data": {"markdown": "Hello",
"screenshot": "base64-image",
"cookies": ["session-secret"]}})
with MrScraperClient("secret", transport=httpx.MockTransport(respond)) as client:
result = make_mrscraper_tools(client)[0]("https://example.com")
assert "Hello" in result
assert "base64-image" not in result
assert "session-secret" not in result
+32 -32
View File
@@ -18,76 +18,76 @@ This report compares **Python SDK feature categories** against **Python document
| Category | Features | Docs | Lines |
|----------|----------|------|-------|
| ✅ AGUI | 1 | 2 | 511 |
| ✅ Agent | 22 | 59 | 19076 |
| ✅ Agent | 22 | 59 | 18970 |
| ✅ Agent-to-Agent (A2A) | 1 | 7 | 2996 |
| ✅ Agent-to-User (A2U) | 1 | 3 | 701 |
| ✅ Approval | 1 | 8 | 3969 |
| ✅ Approval | 1 | 8 | 3941 |
| ✅ Audio | 2 | 12 | 1000 |
| ✅ Auto Generation | 5 | 12 | 3941 |
| ✅ Autonomy | 3 | 7 | 2805 |
| ✅ Bots | 7 | 38 | 15809 |
| ✅ Auto Generation | 5 | 12 | 3879 |
| ✅ Autonomy | 3 | 7 | 2794 |
| ✅ Bots | 7 | 36 | 14882 |
| ✅ Budget | 1 | 1 | 287 |
| ✅ CLI | 5 | 124 | 39191 |
| ✅ CLI | 5 | 124 | 38936 |
| ✅ Chunking | 2 | 2 | 435 |
| ✅ Citations | 2 | 1 | 202 |
| ✅ Code Execution | 2 | 14 | 5022 |
| ✅ Conditions | 1 | 3 | 1386 |
| ✅ Configuration | 3 | 8 | 3341 |
| ✅ Context Management | 16 | 36 | 14384 |
| ✅ Context Management | 16 | 36 | 14289 |
| ✅ Deep Research | 8 | 2 | 591 |
| ✅ Display | 6 | 3 | 844 |
| ✅ Embeddings | 6 | 23 | 2051 |
| ✅ Evaluation | 1 | 9 | 3637 |
| ✅ Events | 1 | 2 | 782 |
| ✅ Execution | 3 | 4 | 1694 |
| ✅ Failover | 2 | 1 | 537 |
| ✅ Evaluation | 1 | 9 | 3537 |
| ✅ Events | 1 | 2 | 780 |
| ✅ Execution | 3 | 4 | 1693 |
| ✅ Failover | 2 | 1 | 480 |
| ✅ Files | 2 | 7 | 2376 |
| ✅ Flow | 1 | 3 | 811 |
| ✅ Gateway | 7 | 97 | 37597 |
| ✅ Guardrails | 5 | 4 | 2642 |
| ✅ Gateway | 7 | 89 | 34188 |
| ✅ Guardrails | 5 | 4 | 2529 |
| ✅ Handoffs | 11 | 6 | 2748 |
| ✅ Hooks | 2 | 9 | 5230 |
| ✅ Image | 1 | 11 | 1295 |
| ✅ Knowledge | 4 | 20 | 6260 |
| ✅ LLM | 3 | 15 | 5958 |
| ✅ Loops | 4 | 5 | 1478 |
| ✅ MCP | 1 | 64 | 15116 |
| ✅ Memory | 6 | 21 | 8198 |
| ✅ Hooks | 2 | 9 | 4990 |
| ✅ Image | 1 | 11 | 1255 |
| ✅ Knowledge | 4 | 20 | 6184 |
| ✅ LLM | 3 | 15 | 5946 |
| ✅ Loops | 4 | 5 | 1475 |
| ✅ MCP | 1 | 64 | 15068 |
| ✅ Memory | 6 | 21 | 8186 |
| ✅ OCR | 2 | 1 | 237 |
| ✅ Observability | 2 | 23 | 2859 |
| ✅ Optimizer | 1 | 2 | 792 |
| ✅ Output | 3 | 5 | 1289 |
| ✅ Output | 3 | 5 | 1273 |
| ✅ Parallel Execution | 3 | 2 | 528 |
| ✅ Planning | 6 | 7 | 1827 |
| ✅ Plugins | 8 | 5 | 4108 |
| ✅ Planning | 6 | 7 | 1825 |
| ✅ Plugins | 8 | 5 | 3404 |
| ✅ Prompts | 2 | 9 | 2043 |
| ✅ Providers | 1 | 58 | 8751 |
| ✅ Providers | 1 | 58 | 8732 |
| ✅ Query | 1 | 2 | 736 |
| ✅ RAG | 5 | 15 | 3152 |
| ✅ Realtime | 2 | 5 | 668 |
| ✅ Reflection | 3 | 3 | 732 |
| ✅ Retrieval | 2 | 5 | 1046 |
| ✅ Routing | 1 | 2 | 430 |
| ✅ Sandbox | 5 | 7 | 3720 |
| ✅ Security | 1 | 3 | 3642 |
| ✅ Sessions | 4 | 21 | 8388 |
| ✅ Skills | 6 | 16 | 6045 |
| ✅ Sandbox | 5 | 7 | 3712 |
| ✅ Security | 1 | 3 | 3607 |
| ✅ Sessions | 4 | 21 | 8285 |
| ✅ Skills | 6 | 16 | 5988 |
| ✅ Tasks | 2 | 8 | 3266 |
| ✅ Telemetry | 1 | 2 | 660 |
| ✅ Templates | 1 | 9 | 1859 |
| ✅ Tools | 12 | 147 | 45378 |
| ✅ Tools | 12 | 145 | 44594 |
| ✅ Tracing | 3 | 2 | 139 |
| ✅ Vector Store | 1 | 12 | 1219 |
| ✅ Video | 2 | 10 | 1303 |
| ✅ Vision | 2 | 1 | 329 |
| ✅ Web | 3 | 10 | 2822 |
| ✅ Workflows | 5 | 21 | 8880 |
| ✅ Web | 3 | 10 | 2815 |
| ✅ Workflows | 5 | 21 | 8862 |
## Documentation Without Features
These docs exist but don't match any implemented feature category:
- ℹ️ Database (44 docs, 8206 lines)
- ℹ️ Database (43 docs, 7657 lines)
- ℹ️ Documents (1 docs, 810 lines)
- ℹ️ Teams (1 docs, 214 lines)
@@ -183,9 +183,7 @@ banner. Key fragments:
timeout-minutes: 15
env:
# Bump deliberately; the cache key and what we test both follow this.
# No leading "v": install.sh passes ?version=$OLLAMA_VERSION and ollama.com
# redirects to GitHub tag v$version (vv0.33.2 → 404 if set to v0.33.2).
OLLAMA_VERSION: "0.33.2"
OLLAMA_VERSION: v0.33.2
OLLAMA_TEST_MODEL: qwen3:0.6b
steps:
- uses: actions/checkout@v4
@@ -206,12 +204,9 @@ banner. Key fragments:
shell: bash
run: |
set -euo pipefail
# The official installer, VERSION-PINNED via OLLAMA_VERSION (which it
# honours). Since v0.33 the linux bundle ships as a .tar.zst with GPU
# runtimes bundled in, so a direct release-asset download is not a
# workable alternative; the installer handles the zst/tgz split.
sudo apt-get update -qq && sudo apt-get install -y -qq zstd
curl -fsSL https://ollama.com/install.sh | sh
curl -fsSL -o /tmp/ollama.tgz \
"https://github.com/ollama/ollama/releases/download/${OLLAMA_VERSION}/ollama-linux-amd64.tgz"
sudo tar -C /usr -xzf /tmp/ollama.tgz
ollama --version
- name: Start Ollama and pull model
shell: bash
+3 -35
View File
@@ -37,22 +37,6 @@ _cache_lock = threading.RLock()
# Global module cache
_module_cache: Dict[str, Any] = {}
# Import roots of praisonaiagents' declared *core* runtime dependencies (see
# pyproject.toml ``[project].dependencies``). A lazy import that fails because
# one of these is missing/broken indicates a genuinely broken install, which
# must be surfaced rather than masked as a graceful ``None``. Optional
# integrations whose own optional SDKs are absent are NOT in this set, so they
# keep the documented ``None`` fallback.
_CORE_REQUIRED_DEPENDENCIES = frozenset({
"pydantic",
"rich",
"openai",
"posthog",
"aiohttp",
"yaml", # PyYAML — core dep; import root differs from the package name
"croniter",
})
def lazy_import(
module_path: str,
@@ -236,25 +220,9 @@ def create_lazy_getattr_with_fallback(
_cache[name] = None
return None
except ImportError as exc:
# Distinguish "optional integration not installed" from
# "a required *core* dependency is broken". A missing/broken core
# dependency (e.g. ``pydantic``) must be surfaced loudly instead
# of being masked as ``None`` (which resurfaces later as a
# confusing ``NoneType`` error). But an optional integration whose
# own optional SDK is absent must keep the documented graceful
# ``None`` fallback — the failing module name there belongs to
# that optional SDK, not to us. We therefore only escalate when
# the failing module is one of praisonaiagents' declared *core*
# runtime dependencies; everything else stays a soft None.
failing = getattr(exc, "name", None)
failing_root = failing.split(".")[0] if failing else None
if failing_root in _CORE_REQUIRED_DEPENDENCIES:
raise ImportError(
f"Failed to load {name!r} because importing {module_path!r} "
f"raised: {exc}. This usually means a required core dependency "
f"({failing_root!r}) is missing or broken, not that {name!r} is "
f"an optional feature you chose not to install."
) from exc
# Optional module not available: preserve the graceful None
# fallback, but record why so real import failures remain
# diagnosable instead of surfacing later as NoneType errors.
_logger.debug(
"Lazy import of %r (%s.%s) failed: %s",
name, module_path, attr_name, exc,
@@ -7,7 +7,10 @@ from praisonaiagents._logging import get_logger
import asyncio
import contextlib
import threading
from typing import List, Optional, Any, Dict, Union, Literal, TYPE_CHECKING, Callable
import concurrent.futures
import random
import re
from typing import List, Optional, Any, Dict, Union, Literal, TYPE_CHECKING, Callable, Generator
from collections import OrderedDict
import inspect
@@ -1370,17 +1373,12 @@ class Agent(GoalLoopMixin, SteeringMixin, SandboxMixin, SkillReviewMixin, Unifie
array_mode=ArrayMode.PASSTHROUGH,
default=None,
)
# Per-agent event-hook registry (isolates BEFORE_TOOL/AFTER_AGENT/...
# hooks so they never fire on unrelated agents sharing the default
# registry). A bare HookRegistry passed as ``hooks=`` still works too.
_hooks_registry = None
if _hooks_config is not None:
if isinstance(_hooks_config, list):
_hooks_list = _hooks_config
elif isinstance(_hooks_config, HooksConfig):
step_callback = _hooks_config.on_step
_hooks_list = list(_hooks_config.middleware or [])
_hooks_registry = getattr(_hooks_config, "registry", None)
# Route on_step / on_tool_call onto the SAME middleware chain
# that already powers ``hooks=[...]`` (MiddlewareManager). Both
# keys used to be stored-and-never-called, so the TypeError that
@@ -2103,11 +2101,7 @@ class Agent(GoalLoopMixin, SteeringMixin, SandboxMixin, SkillReviewMixin, Unifie
# Lazy init for HookRunner and StreamEventEmitter (zero overhead when not used)
self.__hook_runner = None # Will be initialized on first access
self.__stream_emitter = None # Will be initialized on first access
# Prefer a per-agent registry declared via HooksConfig(registry=...) so
# this agent's event hooks stay isolated; else fall back to the raw
# ``hooks`` value (a bare HookRegistry still isolates, anything else
# resolves to the shared default registry in the _hook_runner property).
self._hooks_registry_param = _hooks_registry if _hooks_registry is not None else hooks
self._hooks_registry_param = hooks # Store for lazy init
# Handle llm= deprecation: model= is the preferred parameter name
# llm= still works but shows deprecation warning
@@ -2601,11 +2595,6 @@ class Agent(GoalLoopMixin, SteeringMixin, SandboxMixin, SkillReviewMixin, Unifie
self._on_budget_exceeded = _on_budget_exceeded
# Thread-safe cost/token tracking (Gap 1a fix)
self._cost_lock = threading.Lock()
# Serializes the read-modify-write of _auto_save_last_index. Auto-save
# now runs via asyncio.to_thread, so two concurrent async runs sharing
# one Agent could otherwise read the same last-saved index in separate
# worker threads and each persist the same turns, duplicating history.
self._auto_save_lock = threading.Lock()
self._total_cost = 0.0
self._total_tokens_in = 0
self._total_tokens_out = 0
@@ -2892,10 +2881,7 @@ Your Goal: {self.goal}
# Reuse the caller's manager as-is (no disk I/O side-effect).
_mgr = self._approval_permissions
else:
_mgr = PermissionManager(
agent_name=self.name,
workspace_root=os.getcwd(),
)
_mgr = PermissionManager(agent_name=self.name)
if isinstance(self._approval_permissions, dict):
_mgr.load_rules_from_config(self._approval_permissions)
elif isinstance(self._approval_permissions, (list, tuple)):
@@ -3066,12 +3052,6 @@ Your Goal: {self.goal}
# Agent-centric feature instances (lazy loaded for zero performance impact)
self._auto_memory = auto_memory
self._policy = policy
# Route enabled GUARDRAIL/POLICY/SKILL plugins into their typed
# subsystems now that guardrail (_setup_guardrail), _skills and _policy
# are all set. A plugin declared PluginType.GUARDRAIL/POLICY/SKILL thus
# participates in the GuardrailChain / PolicyEngine / SkillManager, not
# just a generic lifecycle hook. No-op when no such plugin is enabled.
self._merge_plugin_subsystems()
self._output_style = output_style
# Backward-compatible: `thinking_budget` property mirrors the legacy int
# budget when supplied via the alias (Issue #4452); the unified effort is
@@ -3252,7 +3232,7 @@ Your Goal: {self.goal}
for k, v in self.__dict__.items():
if k in ("_Agent__cache_lock",):
object.__setattr__(result, k, threading.RLock())
elif k in ("_cost_lock", "_auto_save_lock"):
elif k == "_cost_lock":
object.__setattr__(result, k, threading.Lock())
elif k == "_approvals_lock":
object.__setattr__(result, k, asyncio.Lock())
@@ -4176,44 +4156,12 @@ Summary:"""
except Exception:
pass
# Gate every code-mode tool call through the SAME PolicyEngine,
# permission deny-rules and per-tool ``input_guardrails`` the direct
# tool-call loop runs. The approval framework the proxy already honours
# does not cover these, so without this hook a model could reach an
# allow-listed tool from code and bypass a guardrail that blocks the
# same call on the normal path.
#
# Contract (mirrors the direct path in ``_check_tool_approval_sync``):
# returns ``(denial_reason, arguments)`` where ``denial_reason`` is a
# string to block or ``None`` to allow, and ``arguments`` is the
# possibly-guardrail-rewritten kwargs dict the tool must run with. The
# proxy uses the returned args so a sanitising guardrail's rewrite is
# honoured (not discarded) exactly as on the direct path. BYPASS mode
# skips the permission-deny gate, matching the direct contract, so a
# call permitted directly is not rejected only in code mode. Only runs
# in the tool-capable modes.
policy_hook = None
if code_tools and code_execution_mode in _tool_capable:
def policy_hook(name, arguments):
if not self._is_bypass_mode():
deny = self._check_permission_manager_deny(name, arguments)
if isinstance(deny, dict):
return deny.get("error", "denied by permission policy"), arguments
verdict = self._check_tool_policy_and_guardrails(name, arguments)
if isinstance(verdict, dict):
return verdict.get("error", "denied by policy"), arguments
# ``verdict`` is ``(None, rewritten_args)`` when allowed; hand the
# possibly-rewritten args back so the proxy dispatches with them.
_, rewritten = verdict
return None, rewritten
# An unknown code_mode raises out of here rather than silently
# producing nothing, which is the failure this whole change is about.
new_tools = build_code_execution_tools(
code_mode=code_execution_mode,
allowed_tools=allowed if code_tools else [],
registry=scoped_registry,
policy_hook=policy_hook,
)
existing = {getattr(t, "__name__", None) for t in (self.tools or [])}
@@ -5391,11 +5339,8 @@ Summary:"""
f"response_len={len(response_str)}"
)
# Auto-save session after each async iteration (memory integration).
# Offloaded to a worker thread so the synchronous FileLock/JSON
# write never stalls the shared event loop (mirrors the goal-
# completion-judge offload below).
await asyncio.to_thread(self._auto_save_session)
# Auto-save session after each async iteration (memory integration)
self._auto_save_session()
# ─────────────────────────────────────────────────────────────
# GOAL COMPLETION JUDGE (opt-in): independent acceptance-criteria
@@ -5535,8 +5480,8 @@ Summary:"""
# Yield control to allow other async tasks to run
await asyncio.sleep(0)
# Final auto-save before returning (offloaded off the event loop).
await asyncio.to_thread(self._auto_save_session)
# Final auto-save before returning
self._auto_save_session()
# Max iterations reached
return AutonomyResult(
@@ -7123,25 +7068,6 @@ Answer:"""
"""The most recent rejection message a guardrail sent back to the model."""
return getattr(self, "_last_guardrail_error", None)
def _guardrail_blocked_message(self, error=None):
"""Fixed, user-facing string returned when an output guardrail blocks a run.
A blocked response must never surface the rejected (potentially unsafe)
model text, but it must also never be ``None`` - callers expect a
``str`` and would otherwise crash on ``.strip()`` (issue #5596, FR-001).
The returned string is intentionally **fixed** and carries no diagnostic
detail: a guardrail's own rejection reason (especially an LLM validator's
free-text) can echo the blocked output, so appending it would leak the
very content the guardrail rejected. The reason is logged instead, where
it stays auditable without reaching the caller. The rejection is also
exposed programmatically via :attr:`last_guardrail_error`.
"""
reason = str(error) if error else getattr(self, "_last_guardrail_error", None)
if reason:
logging.info(f"Agent {self.name}: guardrail block reason: {reason}")
return "I'm sorry, but I can't share that response because it did not pass the safety guardrail."
def _process_guardrail(self, task_output):
"""Process the guardrail validation for a task output.
@@ -7157,16 +7083,8 @@ Answer:"""
return GuardrailResult(success=True, result=task_output)
try:
# Call the guardrail function, awaiting it if the caller supplied an
# async def. Calling a coroutine function synchronously would only
# build a coroutine object (never running its body), which then
# crashes GuardrailResult.from_tuple and is misreported as a
# validation failure.
if asyncio.iscoroutinefunction(self._guardrail_fn):
from ..utils.async_bridge import run_coroutine_from_any_context
result = run_coroutine_from_any_context(self._guardrail_fn(task_output))
else:
result = self._guardrail_fn(task_output)
# Call the guardrail function
result = self._guardrail_fn(task_output)
# Convert the result to a GuardrailResult
return GuardrailResult.from_tuple(result)
@@ -27,21 +27,27 @@ def run_async_in_sync_context(coro):
tool loop, otherwise a bare un-awaited coroutine would be handed to the
model as the tool result and the tool body would never run (silent data
loss).
Delegates to the single canonical bridge
(:func:`utils.async_bridge.run_coroutine_from_any_context`) so all sync→async
call sites share one implementation: the caller's ``contextvars`` (trace /
session / approval context) are carried into the coroutine, and the running
loop is handled on a dedicated worker thread rather than raising.
``timeout=None`` is passed deliberately: this helper backs the sync
tool-calling path, where a tool's own timeout policy (or lack of one) already
governs how long a tool may run. Imposing the bridge's default 5-minute
deadline here would silently cancel long-running tools that previously had no
limit — a backward-incompatible regression — so the bridge adds none.
"""
from ..utils.async_bridge import run_coroutine_from_any_context
return run_coroutine_from_any_context(coro, timeout=None)
try:
asyncio.get_running_loop()
except RuntimeError:
# No event loop - safe to use asyncio.run()
return asyncio.run(coro)
# Event loop exists - avoid deadlock by running in dedicated thread
import concurrent.futures
def run_in_thread():
new_loop = asyncio.new_event_loop()
asyncio.set_event_loop(new_loop)
try:
return new_loop.run_until_complete(coro)
finally:
new_loop.close()
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(run_in_thread)
return future.result(timeout=300) # 5 minute timeout
class DualLock:
@@ -69,24 +75,9 @@ class DualLock:
"""
def __init__(self):
"""Initialize with a single underlying mutex shared by both call styles.
A single ``threading.Lock`` backs both ``sync()`` and ``async_lock()``
so that holding one genuinely blocks the other — mutual exclusion is
real across sync and async callers of the same instance. A plain
``Lock`` (not ``RLock``) is used because the async path acquires on a
worker thread but releases on the event-loop thread, which ``RLock``'s
thread-ownership semantics forbid. Same-thread ``sync()`` reentrancy is
preserved explicitly via an owner/depth guard below. The per-event-loop
``asyncio.Lock`` only serialises coroutines on the same loop so they
queue for the underlying mutex cooperatively instead of every coroutine
racing straight into the blocking acquire.
"""
self._thread_lock = threading.Lock() # Shared mutex (cross-thread release ok)
self._owner_lock = threading.Lock() # Guards the reentrancy bookkeeping
self._owner_thread: Any = None # Thread id currently holding for sync reentrancy
self._owner_depth = 0
self._async_locks = WeakKeyDictionary() # Per-event-loop async gate
"""Initialize with separate threading and asyncio locks."""
self._thread_lock = threading.RLock() # Re-entrant lock to handle nested acquisitions
self._async_locks = WeakKeyDictionary() # Per-event-loop async locks
def __deepcopy__(self, memo):
"""Return an unlocked primitive with no event-loop state sharing."""
@@ -96,31 +87,9 @@ class DualLock:
@contextmanager
def sync(self):
"""Acquire lock in synchronous context using the shared mutex.
Re-entrant for the same thread: a nested ``sync()`` on the thread that
already holds the mutex does not re-acquire (and cannot self-deadlock).
"""
me = threading.get_ident()
with self._owner_lock:
reentrant = self._owner_thread == me and self._owner_depth > 0
if reentrant:
self._owner_depth += 1
if not reentrant:
self._thread_lock.acquire()
with self._owner_lock:
self._owner_thread = me
self._owner_depth = 1
try:
"""Acquire lock in synchronous context using threading.Lock."""
with self._thread_lock:
yield
finally:
with self._owner_lock:
self._owner_depth -= 1
released = self._owner_depth == 0
if released:
self._owner_thread = None
if released:
self._thread_lock.release()
def _get_async_lock(self):
"""Get or create an asyncio.Lock for the current event loop."""
@@ -133,64 +102,16 @@ class DualLock:
self._async_locks[loop] = asyncio.Lock()
return self._async_locks[loop]
async def _acquire_shared_mutex(self, loop):
"""Acquire the shared mutex off-loop, releasing it if we are cancelled.
``run_in_executor`` can complete (mutex acquired) after the awaiting
task is cancelled. If we let the ``CancelledError`` propagate without
checking, the worker would strand the mutex forever. We therefore wait
on the acquisition, and on cancellation ensure the mutex is released
once the worker finishes so no caller waits indefinitely.
"""
future = loop.run_in_executor(None, self._thread_lock.acquire)
try:
await asyncio.shield(future)
except asyncio.CancelledError:
# The worker may still be mid-acquire; wait for it to settle, then
# release if it managed to take the mutex so it is never orphaned.
try:
acquired = await future
except asyncio.CancelledError:
acquired = False
if acquired:
self._thread_lock.release()
raise
@asynccontextmanager
async def async_lock(self):
"""Acquire lock in asynchronous context using the shared mutex.
Backs onto the *same* ``threading.Lock`` as ``sync()`` so a concurrent
sync caller is genuinely excluded (and vice-versa). The blocking
``acquire`` is offloaded to a worker thread so the event loop is never
blocked while waiting; the matching ``release`` (allowed cross-thread on
a plain ``Lock``) happens on the loop thread. A per-event-loop
``asyncio.Lock`` serialises coroutines on this loop first so they don't
all spin up executor hops at once.
While the mutex is held, the *event-loop thread* is registered as the
sync owner so a nested ``with lock.sync()`` on that thread (used by the
async handoff seeding path) is treated as re-entrant instead of
self-deadlocking against the mutex we already hold.
"""Acquire lock in asynchronous context using asyncio.Lock.
Uses a per-event-loop asyncio.Lock to ensure proper async coordination
without blocking the event loop or violating thread ownership semantics.
"""
loop = asyncio.get_running_loop()
async_gate = self._get_async_lock()
async with async_gate:
await self._acquire_shared_mutex(loop)
me = threading.get_ident()
with self._owner_lock:
self._owner_thread = me
self._owner_depth = 1
try:
yield
finally:
with self._owner_lock:
self._owner_depth -= 1
released = self._owner_depth == 0
if released:
self._owner_thread = None
if released:
self._thread_lock.release()
async_lock = self._get_async_lock()
async with async_lock:
yield
def is_async_context(self) -> bool:
"""Check if we're currently in an async context."""
@@ -264,11 +185,6 @@ class AsyncSafeState:
@asynccontextmanager
async def async_lock(self):
"""Acquire lock in async context."""
# Fail fast BEFORE blocking on the shared mutex: an in-progress
# synchronous deepcopy holds the mutex for the whole copy, so waiting on
# it here would block until the copy finished instead of surfacing the
# documented RuntimeError (and could stall the loop indefinitely).
self._raise_if_copying()
async with self._lock.async_lock():
self._begin_async_access()
try:
@@ -276,14 +192,6 @@ class AsyncSafeState:
finally:
self._end_async_access()
def _raise_if_copying(self) -> None:
"""Reject entry while a synchronous deepcopy of this state is active."""
with self._activity_lock:
if self._copying:
raise RuntimeError(
"Cannot access AsyncSafeState during synchronous deepcopy"
)
def _begin_async_access(self) -> None:
"""Register async access unless a synchronous copy is in progress."""
with self._activity_lock:
@@ -300,32 +208,20 @@ class AsyncSafeState:
def __enter__(self):
"""Support for synchronous context manager protocol (backward compatibility)."""
self._sync_guard = self._lock.sync()
self._sync_guard.__enter__()
self._lock._thread_lock.acquire()
return self.value
def __exit__(self, exc_type, exc_val, exc_tb):
"""Support for synchronous context manager protocol (backward compatibility)."""
return self._sync_guard.__exit__(exc_type, exc_val, exc_tb)
self._lock._thread_lock.release()
return None
async def __aenter__(self):
"""Support for asynchronous context manager protocol."""
# Fail fast before blocking on the mutex an in-progress deepcopy holds.
self._raise_if_copying()
loop = asyncio.get_running_loop()
async_lock = self._lock._get_async_lock()
await async_lock.acquire()
try:
# Hold the SAME underlying mutex used by the sync path so async and
# sync accessors are mutually exclusive (matches DualLock.async_lock).
# Cancellation-safe acquire so the mutex is never stranded if the
# awaiting task is cancelled mid-acquisition.
await self._lock._acquire_shared_mutex(loop)
try:
self._begin_async_access()
except Exception:
self._lock._thread_lock.release()
raise
self._begin_async_access()
except Exception:
async_lock.release()
raise
@@ -337,10 +233,7 @@ class AsyncSafeState:
try:
self._end_async_access()
finally:
try:
self._lock._thread_lock.release()
finally:
async_lock.release()
async_lock.release()
return None
def get(self) -> Any:
@@ -751,43 +751,6 @@ Your Goal: {self.goal}"""
return None
def _coerce_structured_output(self, response_text, schema_model):
"""Parse a string response into the requested Pydantic model.
When ``output_pydantic``/``output_json`` is a Pydantic model, callers
expect a validated instance rather than a raw string. This tolerates
markdown code fences around the JSON payload.
Returns the parsed model instance on success, or the original
``response_text`` unchanged if parsing is not possible (preserves
backward compatibility for non-Pydantic schemas or malformed output).
"""
if not schema_model or not isinstance(response_text, str):
return response_text
if not hasattr(schema_model, 'model_validate_json'):
return response_text
candidate = response_text.strip()
try:
return schema_model.model_validate_json(candidate)
except Exception:
pass
# Strip a surrounding markdown code fence (```json ... ```) if present
fenced = re.search(r"```(?:json)?\s*(.*?)\s*```", candidate, re.DOTALL)
if fenced:
try:
return schema_model.model_validate_json(fenced.group(1).strip())
except Exception:
pass
logging.warning(
f"Agent {self.name}: output_pydantic={getattr(schema_model, '__name__', schema_model)} "
f"was requested but the response could not be parsed into the model; "
f"returning raw string."
)
return response_text
def _supports_native_structured_output(self):
"""Check if current model supports native structured output via response_format.
@@ -3599,13 +3562,6 @@ Your Goal: {self.goal}"""
# Apply guardrail validation for custom LLM response
try:
validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools, task_name, task_description, task_id, cancel_token=cancel_token, messages=processed_history)
# Coerce into the requested Pydantic model so callers
# get a validated instance, not a raw JSON string.
# Only output_pydantic changes the return type; output_json
# must stay a string for backward-compatible downstream
# JSON parsing (task parsing, workflow results).
if output_pydantic:
validated_response = self._coerce_structured_output(validated_response, output_pydantic)
# Execute callback and display after validation
self._execute_callback_and_display(prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
return self._trigger_after_agent_hook(prompt, validated_response, start_time)
@@ -3613,7 +3569,7 @@ Your Goal: {self.goal}"""
logging.error(f"Agent {self.name}: Guardrail validation failed for custom LLM: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
except ToolExecutionError:
raise
except Exception as e:
@@ -3650,48 +3606,41 @@ Your Goal: {self.goal}"""
)
# Track messages THIS turn appends so a failure rolls back only our
# own messages, never a concurrent turn's (see memory_mixin). The
# default (OpenAI) path used to skip this, so its positional rollback
# ``del chat_history[n:]`` erased a concurrent turn's interleaved
# messages and corrupted the shared transcript.
_turn_token = self._begin_turn_tracking()
# Store chat history length for potential rollback
chat_history_length = len(self.chat_history)
# Normalize original_prompt for consistent chat history storage
normalized_content = original_prompt
if isinstance(original_prompt, list):
# Extract text from multimodal prompts
normalized_content = next((item["text"] for item in original_prompt if item.get("type") == "text"), "")
# Add user message to chat history BEFORE LLM call so handoffs can
# access it. Use atomic check-then-act to prevent TOCTOU races: an
# unatomic read of chat_history[-1] followed by a separate append
# lets a concurrent turn's user message be mistaken for this turn's
# own duplicate, silently dropping this user turn from the record.
if self._add_to_chat_history_if_not_duplicate("user", normalized_content):
# Persist user message to DB (OpenAI path)
self._persist_message("user", normalized_content)
# Wrap entire turn in try/finally so per-turn tracking is always
# ended (the finally reads the ownership list during rollback).
reflection_count = 0
start_time = time.time()
# Apply context management before LLM call (auto-compaction)
# Zero overhead when context=False
system_prompt_content = messages[0].get("content", "") if messages and messages[0].get("role") == "system" else ""
processed_messages, context_result = self._apply_context_management(
messages=messages,
system_prompt=system_prompt_content,
tools=tools,
)
# Use processed messages for the LLM call
messages = processed_messages
# Wrap entire while loop in try-except for rollback on any failure
try:
# Normalize original_prompt for consistent chat history storage
normalized_content = original_prompt
if isinstance(original_prompt, list):
# Extract text from multimodal prompts
normalized_content = next((item["text"] for item in original_prompt if item.get("type") == "text"), "")
# Add user message to chat history BEFORE LLM call so handoffs can
# access it. Use atomic check-then-act to prevent TOCTOU races: an
# unatomic read of chat_history[-1] followed by a separate append
# lets a concurrent turn's user message be mistaken for this turn's
# own duplicate, silently dropping this user turn from the record.
if self._add_to_chat_history_if_not_duplicate("user", normalized_content):
# Persist user message to DB (OpenAI path)
self._persist_message("user", normalized_content)
reflection_count = 0
start_time = time.time()
# Apply context management before LLM call (auto-compaction)
# Zero overhead when context=False
system_prompt_content = messages[0].get("content", "") if messages and messages[0].get("role") == "system" else ""
processed_messages, context_result = self._apply_context_management(
messages=messages,
system_prompt=system_prompt_content,
tools=tools,
)
# Use processed messages for the LLM call
messages = processed_messages
while True:
try:
if self.verbose:
@@ -3737,11 +3686,6 @@ Your Goal: {self.goal}"""
# Apply guardrail validation even for JSON output
try:
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id, cancel_token=cancel_token, messages=messages)
# Coerce into the requested Pydantic model so callers
# get a validated instance, not a raw JSON string.
# output_json intentionally stays a string (backward compat).
if output_pydantic:
validated_response = self._coerce_structured_output(validated_response, output_pydantic)
# Execute callback after validation
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
return self._trigger_after_agent_hook(original_prompt, validated_response, start_time)
@@ -3749,7 +3693,7 @@ Your Goal: {self.goal}"""
logging.error(f"Agent {self.name}: Guardrail validation failed for JSON output: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
if not self.self_reflect:
# User message already added before LLM call via _build_messages
@@ -3770,7 +3714,7 @@ Your Goal: {self.goal}"""
logging.error(f"Agent {self.name}: Guardrail validation failed for reasoning content: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
# Apply guardrail to regular response
try:
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id, cancel_token=cancel_token, messages=messages)
@@ -3781,7 +3725,7 @@ Your Goal: {self.goal}"""
logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
reflection_prompt = f"""
Reflect on your previous response: '{response_text}'.
@@ -3849,7 +3793,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
self._end_run(None, "error", {"error": str(e)})
return self._guardrail_blocked_message(e)
return None
# Check if we've hit max reflections
if reflection_count >= self.max_reflect - 1:
@@ -3867,7 +3811,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed after max reflections: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
# If not satisfactory and not at max reflections, continue with regeneration
logging.debug(f"{self.name} reflection count {reflection_count + 1}, continuing reflection process")
@@ -3914,10 +3858,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# Rollback chat history
self._rollback_chat_history_to(chat_history_length)
return None
finally:
# End per-turn tracking after any rollback has run (rollback
# reads the ownership list, so this must come last).
self._end_turn_tracking(_turn_token)
def clean_json_output(self, output: str) -> str:
"""Clean JSON output while preserving the legacy agent method."""
@@ -4277,11 +4217,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# Apply guardrail validation for custom LLM response
try:
validated_response = await self._aapply_guardrail_with_retry(response_text, prompt, temperature, tools, task_name, task_description, task_id, messages=effective_history)
# Coerce into the requested Pydantic model so callers
# get a validated instance, not a raw JSON string.
# output_json intentionally stays a string (backward compat).
if output_pydantic:
validated_response = self._coerce_structured_output(validated_response, output_pydantic)
# Execute callback after validation
self._execute_callback_and_display(normalized_content, validated_response, time.time() - start_time, task_name, task_description, task_id)
return await self._atrigger_after_agent_hook(prompt, validated_response, start_time)
@@ -4289,7 +4224,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed for custom LLM: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
except ToolExecutionError:
raise
except Exception as e:
@@ -4315,15 +4250,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
if durable_context is not None:
messages = await durable_context.arestore_messages(messages)
# Track messages THIS turn appends so a failure rolls back only our
# own messages, never a concurrent turn's (see memory_mixin). The
# default (OpenAI) path used to skip this, so its positional rollback
# ``del chat_history[n:]`` erased a concurrent turn's interleaved
# messages. Each achat() turn runs in its own task/context and
# _clear_turn_tracking() ran at the top of this turn, so this list is
# cleaned up when the next turn on this task starts.
self._begin_turn_tracking()
# Store chat history length for potential rollback
chat_history_length = len(self.chat_history)
@@ -4427,11 +4353,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# Apply guardrail validation even for JSON output
try:
validated_response = await self._aapply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id, messages=messages)
# Coerce into the requested Pydantic model so callers
# get a validated instance, not a raw JSON string.
# output_json intentionally stays a string (backward compat).
if output_pydantic:
validated_response = self._coerce_structured_output(validated_response, output_pydantic)
# Execute callback after validation
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
return await self._atrigger_after_agent_hook(original_prompt, validated_response, start_time)
@@ -4439,7 +4360,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed for JSON output: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
# For regular responses (no self-reflection)
if not self.self_reflect:
@@ -4462,7 +4383,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed for reasoning content: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
else:
# Apply guardrail to regular response content
try:
@@ -4474,7 +4395,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
# If self-reflection is enabled, implement reflection logic
if self.self_reflect:
@@ -4514,7 +4435,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
_areflection_kwargs = {}
if temperature is not None:
@@ -4548,7 +4469,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed after reflection: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
# Check if we've hit max reflections
if reflection_count >= self.max_reflect - 1:
@@ -4567,7 +4488,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed after max reflections: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
# Regenerate response based on reflection
regenerate_messages = reflection_messages + [
@@ -4617,7 +4538,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed after reflection error: {guard_e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(guard_e)
return None
continue
# This should never be reached due to the returns above
@@ -4667,11 +4588,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
if get_logger().getEffectiveLevel() == logging.DEBUG:
total_time = time.time() - start_time
logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
# Coerce into the requested Pydantic model so callers
# get a validated instance, not a raw JSON string.
# output_json intentionally stays a string (backward compat).
if output_pydantic:
response_text = self._coerce_structured_output(response_text, output_pydantic)
# Execute callback after JSON/Pydantic completion
self._execute_callback_and_display(original_prompt, response_text, time.time() - start_time, task_name, task_description, task_id)
return await self._atrigger_after_agent_hook(original_prompt, response_text, start_time)
@@ -4786,7 +4702,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.error(f"Agent {self.name}: Guardrail validation failed for OpenAI client: {e}")
# Rollback chat history on guardrail failure
self._rollback_chat_history_to(chat_history_length)
return self._guardrail_blocked_message(e)
return None
except ToolExecutionError:
raise
except Exception as e:
@@ -4810,16 +4726,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# dangling, unanswered user message (mirrors the sync path).
self._rollback_chat_history_to(chat_history_length)
return None
finally:
# Unregister this turn's ownership list from the process-wide
# live-turn registry once the turn finishes (success or failure).
# The sync path does this in _end_turn_tracking's finally; the async
# path previously relied only on the *next* turn's
# _clear_turn_tracking(), so a terminal achat/arun/astart left its
# completed ownership globally registered and a later ephemeral()
# cleanup treated its committed messages as still live. Rollback in
# the except blocks above already read ownership before this runs.
self._clear_turn_tracking()
async def _achat_completion(self, response, tools, reasoning_steps=False):
"""Async version of _chat_completion method"""
@@ -5792,28 +5698,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
return llm_summarize_async
def _resolve_compaction_max_tokens(self, execution_cfg):
"""
Resolve the compaction trigger budget.
An explicit ``ExecutionConfig.max_context_tokens`` always wins. When it
is unset, size the trigger to the active model's context window (reusing
the model-aware ``budgeter`` the in-loop path already uses) instead of a
flat default, so large-context models are not compacted prematurely.
Falls back to a static default only for the offline / unknown-model case.
"""
_explicit = getattr(execution_cfg, 'max_context_tokens', None)
if _explicit:
return _explicit
try:
from ..context.budgeter import get_model_limit, get_output_reserve
model_name = self.llm if isinstance(self.llm, str) else "gpt-4o-mini"
limit = get_model_limit(model_name)
reserve = get_output_reserve(model_name)
return max(1, int(limit * 0.8) - reserve)
except Exception:
return 8000
def _apply_context_compaction(self, messages, hook_event_class):
"""
Apply context compaction to messages if enabled (sync version).
@@ -5833,7 +5717,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
from ..compaction import ContextCompactor
from ..compaction.strategy import CompactionStrategy
_max_tok = self._resolve_compaction_max_tokens(_execution_cfg)
_max_tok = getattr(_execution_cfg, 'max_context_tokens', None) or 8000
_strategy = getattr(_execution_cfg, 'compaction_strategy', None) or CompactionStrategy.TRUNCATE
# Create LLM summarization function if strategy is LLM_SUMMARIZE
@@ -5865,20 +5749,17 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# Perform compaction
if _strategy == CompactionStrategy.LLM_SUMMARIZE and _llm_fn:
# Drive the async summariser through the shared bridge so it runs
# correctly whether or not a loop is already running (FastAPI /
# Jupyter / a bot handler). The previous code caught bare
# RuntimeError from ``asyncio.run`` ("event loop already
# running") and silently downgraded LLM summarisation to naive
# truncation whenever a loop was live; the bridge removes that
# spurious fallback so the real summariser always runs. Provider
# errors inside the summariser remain best-effort (see
# ``_llm_summarize_async``, which falls back to a naive summary on
# failure) — compaction never hard-fails the turn.
from ..utils.async_bridge import run_coroutine_from_any_context
compacted_msgs, _cr = run_coroutine_from_any_context(
_compactor.compact_async(messages)
)
import asyncio
try:
# Run async compaction in event loop
compacted_msgs, _cr = asyncio.run(_compactor.compact_async(messages))
except RuntimeError:
# If already in async context, fall back to sync (naive) compaction
logging.warning(
f"[compaction] {self.name}: LLM_SUMMARIZE fell back to naive summarization "
f"(asyncio.run not available in sync context)"
)
compacted_msgs, _cr = _compactor.compact(messages)
else:
compacted_msgs, _cr = _compactor.compact(messages)
@@ -5923,7 +5804,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
from ..compaction import ContextCompactor
from ..compaction.strategy import CompactionStrategy
_max_tok = self._resolve_compaction_max_tokens(_execution_cfg)
_max_tok = getattr(_execution_cfg, 'max_context_tokens', None) or 8000
_strategy = getattr(_execution_cfg, 'compaction_strategy', None) or CompactionStrategy.TRUNCATE
# Create LLM summarization function if strategy is LLM_SUMMARIZE
@@ -31,66 +31,23 @@ from praisonaiagents._logging import get_logger
logger = get_logger(__name__)
class _DynamicLimiter:
"""Resizable counting limiter: capacity can change with holders outstanding.
Unlike swapping a fresh threading.Semaphore on every retune (which lets
holders of the old primitive and the release() path operate on different
objects), this keeps one long-lived object per agent name and mutates its
capacity in place. Every acquirer and every release() therefore always
operate on the same shared state, so a retune takes effect immediately and
the configured cap can never be silently exceeded. Loop-neutral: uses a
threading.Condition so it is a true global cap across event loops/threads.
"""
def __init__(self, limit: int):
self._limit = limit
self._held = 0
self._cond = threading.Condition()
def set_limit(self, limit: int) -> None:
with self._cond:
self._limit = limit
self._cond.notify_all()
def acquire(self, blocking: bool = True) -> bool:
with self._cond:
while self._limit > 0 and self._held >= self._limit:
if not blocking:
return False
self._cond.wait(timeout=0.05)
self._held += 1
return True
def release(self) -> None:
with self._cond:
self._held = max(0, self._held - 1)
self._cond.notify_all()
def has_holders(self) -> bool:
with self._cond:
return self._held > 0
class ConcurrencyRegistry:
"""Registry for per-agent concurrency limits.
Thread-safe. Each agent name maps to a loop-neutral, resizable limiter, so
the limit is a true global cap across every event loop and thread and a
retune takes effect immediately for in-flight holders.
Thread-safe. Each agent name maps to a loop-neutral threading.Semaphore, so
the limit is a true global cap across every event loop and thread.
Limit of 0 means unlimited (no throttling).
"""
def __init__(self, default_limit: int = 0):
self._default_limit = default_limit
self._limits: Dict[str, int] = {}
# One long-lived, loop-neutral _DynamicLimiter per agent. Its capacity is
# mutated in place on retune (never swapped for a fresh object), so every
# acquirer and every release() always share the same state and the
# configured cap can never be silently exceeded. Loop-neutral means the
# per-agent limit is a true global cap under multi-loop / multi-thread
# deployments; async callers poll without blocking their loop.
self._limiters: Dict[str, _DynamicLimiter] = {}
# One loop-neutral threading.Semaphore per agent. Unlike asyncio.Semaphore
# (which binds to the loop it is first awaited on), a threading.Semaphore
# is shared safely across every event loop and thread, so the per-agent
# limit is a true global cap under multi-loop / multi-thread deployments.
# Async callers acquire it in an executor so they never block their loop.
self._semaphores: Dict[str, threading.Semaphore] = {}
self._lock = threading.Lock()
def set_limit(self, agent_name: str, max_concurrent: int) -> None:
@@ -102,11 +59,8 @@ class ConcurrencyRegistry:
"""
with self._lock:
self._limits[agent_name] = max_concurrent
# Mutate the existing limiter in place so the retune takes effect
# immediately for in-flight holders without ever swapping the object.
limiter = self._limiters.get(agent_name)
if limiter is not None:
limiter.set_limit(max_concurrent)
# Reset semaphore so next acquire creates a fresh one at the new limit
self._semaphores.pop(agent_name, None)
def get_limit(self, agent_name: str) -> int:
"""Get concurrency limit for an agent."""
@@ -114,57 +68,39 @@ class ConcurrencyRegistry:
return self._limits.get(agent_name, self._default_limit)
def remove_limit(self, agent_name: str) -> None:
"""Remove concurrency limit for an agent (reverts to default).
The limiter object is kept alive while any permit is still held, so
in-flight acquire()/release() pairs keep operating on the same shared
state; dropping it mid-flight would strand those permits and let a
later set_limit() start from a fresh, empty count and admit work beyond
the configured cap. Once no permits are outstanding it is safe to drop.
"""
"""Remove concurrency limit for an agent (reverts to default)."""
with self._lock:
self._limits.pop(agent_name, None)
# Retune the surviving limiter to unlimited so it stops blocking
# while remaining the single accounting object for its holders.
limiter = self._limiters.get(agent_name)
if limiter is None:
return
limiter.set_limit(0)
if not limiter.has_holders():
self._limiters.pop(agent_name, None)
self._semaphores.pop(agent_name, None)
def _get_limiter(self, agent_name: str) -> Optional[_DynamicLimiter]:
"""Get or create the loop-neutral limiter for an agent.
def _get_semaphore(self, agent_name: str) -> Optional[threading.Semaphore]:
"""Get or create the loop-neutral semaphore for an agent.
Returns None only when the agent has never been given a limit (pure
unlimited fast path). Once a limiter exists it is always returned — even
while unlimited (limit 0) — so acquire() and release() stay symmetric
across a mid-flight retune to/from unlimited and the count can never be
corrupted. The same _DynamicLimiter is shared across every event loop
and thread, so a configured limit is a true global cap.
Returns None if unlimited. The same threading.Semaphore is shared across
every event loop and thread, so the configured limit is a true global cap.
"""
with self._lock:
limit = self._limits.get(agent_name, self._default_limit)
limiter = self._limiters.get(agent_name)
if limiter is None:
if limit <= 0:
return None
limiter = _DynamicLimiter(limit)
self._limiters[agent_name] = limiter
return limiter
if limit <= 0:
return None
sem = self._semaphores.get(agent_name)
if sem is None:
sem = threading.Semaphore(limit)
self._semaphores[agent_name] = sem
return sem
async def acquire(self, agent_name: str) -> None:
"""Acquire concurrency slot for agent. No-op if unlimited.
Waits on the loop-neutral limiter in short, cancellable polls so the
Waits on the loop-neutral semaphore in short, cancellable polls so the
running event loop is never blocked while other tasks hold permits, and
a cancelled/timed-out await never leaves a thread blocked on acquire().
"""
limiter = self._get_limiter(agent_name)
if limiter is None:
sem = self._get_semaphore(agent_name)
if sem is None:
return
while True:
if limiter.acquire(blocking=False):
if sem.acquire(blocking=False):
return
# Yield to the loop; on cancellation this raises and no permit leaks.
await asyncio.sleep(0.005)
@@ -174,18 +110,21 @@ class ConcurrencyRegistry:
Prefer async acquire() when possible. Blocks the calling thread until a
permit is available. Safe to call whether or not a loop is running in the
current thread, since the limiter is loop-neutral.
current thread, since the semaphore is loop-neutral.
"""
limiter = self._get_limiter(agent_name)
if limiter is not None:
limiter.acquire()
sem = self._get_semaphore(agent_name)
if sem is not None:
sem.acquire()
def release(self, agent_name: str) -> None:
"""Release concurrency slot for agent. No-op if unlimited."""
with self._lock:
limiter = self._limiters.get(agent_name)
if limiter is not None:
limiter.release()
sem = self._semaphores.get(agent_name)
if sem is not None:
try:
sem.release()
except ValueError:
pass # Already fully released
@asynccontextmanager
async def throttle(self, agent_name: str):
@@ -58,8 +58,8 @@ def save_cost_report(
}
filepath = COST_DIR / f"{safe_name}.json"
from ..utils.atomic_io import atomic_write_json
atomic_write_json(str(filepath), report, indent=2)
with open(filepath, "w") as f:
json.dump(report, f, indent=2)
logger.info(f"[costs] Saved cost report to {filepath}")
return str(filepath)
@@ -898,18 +898,14 @@ You are a professional research analyst. When conducting research:
if grounding_file_prompt:
create_params["grounding_file_prompt"] = grounding_file_prompt
interaction = await asyncio.to_thread(
self.gemini_client.interactions.create, **create_params
)
interaction = self.gemini_client.interactions.create(**create_params)
if self.verbose:
self.logger.debug(f"Gemini research started (async): {interaction.id}")
start_time = time.time()
while True:
interaction = await asyncio.to_thread(
self.gemini_client.interactions.get, interaction.id
)
interaction = self.gemini_client.interactions.get(interaction.id)
if interaction.status == "completed":
if self.verbose:
@@ -200,12 +200,8 @@ class ExecutionMixin:
response = await self.achat(prompt, **kwargs)
response_str = str(response) if response else ""
# Auto-save session after chat. Offloaded to a worker thread so
# the synchronous FileLock/JSON read-modify-write never stalls
# the shared event loop (the session store is a process-wide
# singleton, so a blocking write here blocks every concurrent
# async agent task on the same loop).
await asyncio.to_thread(self._auto_save_session)
# Auto-save session after chat
self._auto_save_session()
# Wrap in AutonomyResult for consistent API
from .autonomy import AutonomyResult
@@ -476,8 +472,7 @@ class ExecutionMixin:
extra={"session_id": sid, "managed_session_id": msid},
exc_info=True,
)
# Offload the blocking session write off the event loop.
await asyncio.to_thread(self._auto_save_session)
self._auto_save_session()
return result
@@ -1738,17 +1733,12 @@ Write the complete compiled report:"""
# Policy/guardrail gate (protocol-driven). Mirrors the sync path in
# _execute_tool_impl so async callers cannot bypass a PolicyEngine
# deny or a tool-call guardrail. A user-supplied guardrail/policy may
# do blocking work (an LLM moderation call, a DB/network lookup), so
# offload to a worker thread rather than run it inline and stall the
# event loop for every other coroutine sharing it.
# deny or a tool-call guardrail. The check is pure/sync (no awaits).
check = getattr(self, "_check_tool_policy_and_guardrails", None)
if check is not None:
# ``tools_override`` is threaded through so a run-scoped tool
# list cannot bypass the per-tool guardrails its tools declare.
policy_result = await asyncio.to_thread(
check, function_name, arguments, tools_override
)
policy_result = check(function_name, arguments, tools_override)
if isinstance(policy_result, dict):
return policy_result # Error dict
_, arguments = policy_result
@@ -1775,12 +1765,9 @@ Write the complete compiled report:"""
# Tool-result guardrail gate — the native path applies this
# below, but the MCP branch returns early, so gate here too
# or successful MCP output would bypass validation entirely.
# Offloaded to a thread: a user guardrail may block on I/O.
apply_result_guardrails = getattr(self, "_apply_tool_result_guardrails", None)
if apply_result_guardrails is not None:
mcp_result = await asyncio.to_thread(
apply_result_guardrails, function_name, mcp_result
)
mcp_result = apply_result_guardrails(function_name, mcp_result)
return mcp_result
# Try to find the function in the override tools list first, then agent's tools list.
@@ -2141,14 +2128,9 @@ Write the complete compiled report:"""
# sync path in tool_execution.py. Runs on the raw result before it
# re-enters the LLM context so a guardrail can inspect or redact
# unsafe tool output. Fail-closed. Zero overhead when unset.
# Offloaded to a thread: a user guardrail may block on I/O (an LLM
# moderation call, a network lookup), which would otherwise stall
# the event loop for every other coroutine sharing it.
apply_result_guardrails = getattr(self, "_apply_tool_result_guardrails", None)
if apply_result_guardrails is not None:
result = await asyncio.to_thread(
apply_result_guardrails, function_name, result, tools_override
)
result = apply_result_guardrails(function_name, result, tools_override)
# Loop guard (post-execution) — record the outcome and surface a
# block/halt decision back to the model on this same turn, mirroring
@@ -569,13 +569,17 @@ class Handoff:
handoffs to the same agent serialize on the event loop instead of
blocking it, keeping each seeded ``achat`` turn atomic.
``DualLock`` backs ``sync()`` and ``async_lock()`` onto one shared mutex,
so ``async_lock()`` alone already excludes a concurrent *synchronous*
handoff on another thread. We still take the inner ``sync()`` so the
seed/restore reads as one explicit critical section; because
``async_lock()`` registers the event-loop thread as the mutex owner, that
nested ``sync()`` is treated as re-entrant on the same thread instead of
self-deadlocking against the mutex the async side already holds.
``DualLock`` keeps its sync (threading ``RLock``) and async (per-loop
``asyncio.Lock``) sides independent, so the async lock alone would NOT
exclude a concurrent *synchronous* handoff running on another thread —
both could seed/restore the same shared ``chat_history`` at once, and a
cross-thread sync seed landing *between* our apply and restore would
corrupt the awaited turn. To close that gap we ALSO hold the sync
thread-lock for the whole seeded turn. The async lock (acquired first)
has already serialized every same-loop async handoff to this target, so
the thread-lock here only ever contends with a genuine cross-thread sync
handoff to the same agent — exactly the case that must be mutually
excluded — which is a rare mixed sync+async fan-out to one shared target.
"""
lock = _get_handoff_seed_lock(self.agent)
async with lock.async_lock():
@@ -764,50 +764,6 @@ class MemoryMixin:
except Exception as e:
logging.debug(f"Auto-memory extraction failed: {e}")
def _persist_memory_turn(self, user_message: str, assistant_response: str):
"""Persist a raw conversation turn to the memory store.
Called after each response when ``_memory_instance`` is set but
``auto_memory`` extraction is off. Stores the turn to short-term memory
so a later Agent sharing the same store (e.g. ``memory={"user_id": uid}``)
recalls it via ``get_memory_context()``. Best-effort: never raises.
Async-safe: a memory store write is blocking file/DB I/O. When this runs
inside a live event loop (the async ``achat``/``astart`` after-agent path),
the write is offloaded to a worker thread so it never stalls the loop
(AGENTS.md §4.5). On the sync path (no running loop) it writes inline.
"""
memory = getattr(self, "_memory_instance", None)
if memory is None or not assistant_response:
return
store = getattr(memory, "store_short_term", None) or getattr(memory, "add_short_term", None)
if store is None:
return
text = f"User: {user_message}\nAssistant: {assistant_response}"
metadata = {"agent_id": getattr(self, "agent_id", getattr(self, "name", None))}
def _do_store():
try:
store(text, metadata=metadata)
except Exception as e:
logging.debug(f"Memory turn persistence failed: {e}")
try:
import asyncio
asyncio.get_running_loop()
except RuntimeError:
# No running loop (sync chat/run/start path): write inline.
_do_store()
return
# Inside a live loop: offload the blocking write so it cannot stall the
# event loop. Fire-and-forget — persistence is best-effort.
try:
import asyncio
asyncio.get_running_loop().run_in_executor(None, _do_store)
except Exception as e:
logging.debug(f"Memory turn persistence scheduling failed: {e}")
_do_store()
def _process_auto_learning(self):
"""Process auto-learning extraction after agent response.
@@ -928,56 +884,45 @@ class MemoryMixin:
if not self.auto_save:
return
# Serialize the whole read-modify-write of _auto_save_last_index. This
# method now runs via asyncio.to_thread, so two concurrent async runs
# sharing one Agent can execute it in separate worker threads. Without
# this lock both could read the same last-saved index, then each persist
# the same new turns, duplicating conversation history — the store's
# file lock serializes the writes but does not protect this index.
lock = getattr(self, '_auto_save_lock', None)
if lock is None:
lock = threading.Lock()
self._auto_save_lock = lock
with lock:
try:
# Filter out history markers before saving
clean_history = [
{k: v for k, v in msg.items() if k != "_from_history"}
for msg in self.chat_history
]
# Issue 1 FIX: Only save NEW messages since last save
# Track last saved index to avoid duplicates
last_saved = getattr(self, '_auto_save_last_index', 0)
new_messages = clean_history[last_saved:]
if not new_messages:
return # Nothing new to save
# G-1 FIX: Use SessionStore for conversation history persistence
# This maintains clean separation: Memory = facts, SessionStore = turns
if self._session_store is not None:
# Persist only NEW messages to SessionStore
for msg in new_messages:
self._session_store.add_message(
self.auto_save, # Use auto_save name as session_id
role=msg.get("role", "user"),
content=msg.get("content", ""),
)
# Update last saved index
self._auto_save_last_index = len(clean_history)
logging.debug(f"Auto-saved {len(new_messages)} new messages to SessionStore: {self.auto_save}")
elif self._memory_instance and hasattr(self._memory_instance, 'save_session'):
# Fallback to Memory.save_session() for backward compatibility
# Memory.save_session() replaces entire history, so no duplicate issue
self._memory_instance.save_session(
name=self.auto_save,
conversation_history=clean_history,
metadata={"agent_name": self.name, "user_id": self.user_id}
try:
# Filter out history markers before saving
clean_history = [
{k: v for k, v in msg.items() if k != "_from_history"}
for msg in self.chat_history
]
# Issue 1 FIX: Only save NEW messages since last save
# Track last saved index to avoid duplicates
last_saved = getattr(self, '_auto_save_last_index', 0)
new_messages = clean_history[last_saved:]
if not new_messages:
return # Nothing new to save
# G-1 FIX: Use SessionStore for conversation history persistence
# This maintains clean separation: Memory = facts, SessionStore = turns
if self._session_store is not None:
# Persist only NEW messages to SessionStore
for msg in new_messages:
self._session_store.add_message(
self.auto_save, # Use auto_save name as session_id
role=msg.get("role", "user"),
content=msg.get("content", ""),
)
logging.debug(f"Auto-saved session to Memory: {self.auto_save}")
except Exception as e:
logging.debug(f"Error auto-saving session: {e}")
# Update last saved index
self._auto_save_last_index = len(clean_history)
logging.debug(f"Auto-saved {len(new_messages)} new messages to SessionStore: {self.auto_save}")
elif self._memory_instance and hasattr(self._memory_instance, 'save_session'):
# Fallback to Memory.save_session() for backward compatibility
# Memory.save_session() replaces entire history, so no duplicate issue
self._memory_instance.save_session(
name=self.auto_save,
conversation_history=clean_history,
metadata={"agent_name": self.name, "user_id": self.user_id}
)
logging.debug(f"Auto-saved session to Memory: {self.auto_save}")
except Exception as e:
logging.debug(f"Error auto-saving session: {e}")
def _save_output_to_file(self, content: str) -> bool:
"""Save agent output to file if output_file is configured.
@@ -29,11 +29,6 @@ logger = logging.getLogger(__name__)
# timeouts so they cannot exhaust process resources.
_MAX_ORPHANED_TOOL_EXECUTORS = 4
# 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
# 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).
# It is distinct from {} ("no arguments"): {} would silently execute the tool
@@ -44,57 +39,6 @@ _TOOL_POOL_MAX_WORKERS = 2
_TOOL_ARGUMENTS_PARSE_FAILED = object()
class _CallableGuardrailAdapter:
"""Expose a plain callable guardrail as a chainable guardrail object.
A user ``guardrail`` may be a bare callable returning ``(bool, output)``
(the Agent's final-output contract). To let it run *inside* a
``GuardrailChain`` alongside a plugin guardrail — so composing the two never
silently drops either — this wraps the callable with a ``validate_output``
(and passthrough ``validate_input``) the chain understands. Any
``validate_tool_call`` / ``validate_tool_result`` the underlying object
already exposes is forwarded so tool-surface behaviour is preserved.
"""
__slots__ = ("_fn",)
def __init__(self, fn):
self._fn = fn
def validate_output(self, content, **kwargs):
result = self._fn(content)
if isinstance(result, tuple) and len(result) == 2:
is_valid, processed = result
if is_valid:
# Normalise the passthrough value to a string so the chain can
# keep threading ``content`` through subsequent guardrails.
if isinstance(processed, str):
return True, processed
return True, content
return False, processed if isinstance(processed, str) else str(processed)
# A callable that does not follow the (bool, output) contract is treated
# as a no-op passthrough rather than blocking.
return True, content
def validate_input(self, content, **kwargs):
inner = getattr(self._fn, "validate_input", None)
if callable(inner):
return inner(content, **kwargs)
return True, content
def validate_tool_call(self, tool_name, arguments, **kwargs):
inner = getattr(self._fn, "validate_tool_call", None)
if callable(inner):
return inner(tool_name, arguments, **kwargs)
return True, arguments
def validate_tool_result(self, tool_name, result, **kwargs):
inner = getattr(self._fn, "validate_tool_result", None)
if callable(inner):
return inner(tool_name, result, **kwargs)
return True, result
def tool_arguments_parse_failed(arguments) -> bool:
"""True when tool-call arguments could not be parsed (vs. legitimately empty)."""
return arguments is _TOOL_ARGUMENTS_PARSE_FAILED
@@ -653,175 +597,6 @@ class ToolExecutionMixin:
display_name = tool.get("type", "spec")
logging.debug("Added plugin tool to Agent: %s", display_name or tool)
def _merge_plugin_subsystems(self):
"""Route enabled GUARDRAIL/POLICY/SKILL plugins into their subsystems.
Complements ``_merge_plugin_tools``: instead of contributing tools, a
plugin declared ``PluginType.GUARDRAIL``/``POLICY``/``SKILL`` now
participates in the typed subsystem its category promises —
* GUARDRAIL: ``as_guardrail()`` objects are folded into ``self.guardrail``
so ``_setup_guardrail()`` (which runs after this) wires them into the
same ``GuardrailChain`` the Agent already runs — including
``validate_tool_call`` / ``validate_tool_result``.
* POLICY: ``get_policies()`` rules are added to ``self._policy``
(a ``PolicyEngine``) so they gate tool calls.
* SKILL: ``get_skills()`` entries are appended to ``self._skills`` so the
lazy ``SkillManager`` discovers them as real, loadable skills.
Best-effort and fail-open on plugin errors: a broken plugin must never
prevent the Agent from constructing. Zero overhead when no such plugin
is enabled (each collector returns an empty list).
"""
try:
from ..plugins import get_plugin_manager
manager = get_plugin_manager()
except Exception as exc: # pragma: no cover - defensive plugin boundary
logging.warning("Failed to access plugin manager for subsystems: %s", exc)
return
# --- GUARDRAIL plugins -> self.guardrail (GuardrailChain) -------------
get_guardrails = getattr(manager, "get_all_guardrails", None)
if callable(get_guardrails):
try:
plugin_guardrails = get_guardrails() or []
except Exception as exc: # pragma: no cover - defensive boundary
logging.warning("Failed to collect plugin guardrails: %s", exc)
plugin_guardrails = []
if plugin_guardrails:
try:
from ..guardrails import GuardrailChain
from ..guardrails.protocols import is_guardrail_object
chain_items = [
g for g in plugin_guardrails if is_guardrail_object(g)
]
if chain_items:
chain = GuardrailChain(chain_items)
# Register the plugin guardrails' tool-call/result
# surface on the runtime lists ``_check_tool_policy_and_
# guardrails`` / ``_apply_tool_result_guardrails`` read,
# so a guardrail plugin finally sees tool calls and raw
# tool results — the blind spot ``after_llm`` never
# covered. Runs after ``_setup_guardrail()`` so we append
# rather than replace any user-supplied guardrail surface.
calls = getattr(self, "_tool_call_guardrails", None)
if isinstance(calls, list):
calls.append(chain)
else:
self._tool_call_guardrails = [chain]
results = getattr(self, "_tool_result_guardrails", None)
if isinstance(results, list):
results.append(chain)
else:
self._tool_result_guardrails = [chain]
# Also make the chain participate in the Agent's
# input/output validation (not just the tool surface).
existing_fn = getattr(self, "_guardrail_fn", None)
if existing_fn is None:
# No user guardrail: the plugin chain becomes the
# Agent's output/input validator directly.
self.guardrail = chain
self._guardrail_fn = chain
elif existing_fn is not chain:
# User already supplied a guardrail. Compose rather
# than replace so BOTH run on the final response:
# wrap the user's callable in an adapter exposing
# validate_output/validate_input, then chain it with
# the plugin chain. The user's guardrail runs first;
# a plugin guardrail can no longer be silently
# skipped just because a user guardrail exists.
adapter = _CallableGuardrailAdapter(existing_fn)
combined = GuardrailChain([adapter, chain])
self.guardrail = combined
self._guardrail_fn = combined
except Exception as exc: # pragma: no cover - defensive boundary
logging.warning("Failed to wire plugin guardrails: %s", exc)
# --- POLICY plugins -> self._policy (PolicyEngine) -------------------
get_policies = getattr(manager, "get_all_policies", None)
if callable(get_policies):
try:
plugin_policies = get_policies() or []
except Exception as exc: # pragma: no cover - defensive boundary
logging.warning("Failed to collect plugin policies: %s", exc)
plugin_policies = []
if plugin_policies:
try:
policy = getattr(self, "_policy", None)
if policy is None:
from ..policy import PolicyEngine
policy = PolicyEngine()
self._policy = policy
add_policy = getattr(policy, "add_policy", None)
if callable(add_policy):
# A user-supplied policy always wins: never let a plugin
# rule silently replace an existing policy of the same
# name (PolicyEngine.add_policy overwrites by name), which
# could otherwise loosen a user-configured restriction.
existing = set()
get_policy = getattr(policy, "get_policy", None)
for rule in plugin_policies:
rule_name = getattr(rule, "name", None)
if rule_name is not None and callable(get_policy):
try:
if get_policy(rule_name) is not None:
logging.warning(
"Skipping plugin policy %r: a policy "
"with that name already exists.",
rule_name,
)
continue
except Exception: # pragma: no cover
pass
if rule_name is not None and rule_name in existing:
logging.warning(
"Skipping duplicate plugin policy %r.",
rule_name,
)
continue
try:
add_policy(rule)
if rule_name is not None:
existing.add(rule_name)
except Exception as exc: # pragma: no cover
logging.warning(
"Failed to add plugin policy %r: %s", rule, exc
)
except Exception as exc: # pragma: no cover - defensive boundary
logging.warning("Failed to wire plugin policies: %s", exc)
# --- SKILL plugins -> self._skills (SkillManager discovers them) -----
get_skills = getattr(manager, "get_all_skills", None)
if callable(get_skills):
try:
plugin_skills = get_skills() or []
except Exception as exc: # pragma: no cover - defensive boundary
logging.warning("Failed to collect plugin skills: %s", exc)
plugin_skills = []
if plugin_skills:
current = getattr(self, "_skills", None)
if not isinstance(current, list):
current = [current] if current else []
for skill in plugin_skills:
# ``self._skills`` feeds ``SkillManager.add_skill(path)``,
# which resolves each entry as a filesystem path. Only
# string path selectors are loadable; a non-string (e.g. a
# source object) would raise inside ``Path()``, so skip it
# with a clear warning rather than corrupt discovery.
if not isinstance(skill, str):
logging.warning(
"Skipping plugin skill %r: get_skills() must return "
"filesystem path strings to skill directories.",
skill,
)
continue
if skill not in current:
current.append(skill)
self._skills = current
def _cast_arguments(self, func, arguments):
"""Cast arguments to their expected types based on function signature."""
if not callable(func) or not arguments:
@@ -1306,72 +1081,13 @@ class ToolExecutionMixin:
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}"
max_workers=2, 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:
logging.warning(
f"Tool worker pool exhausted; refusing to queue '{function_name}' "
f"behind {orphaned} hung tool call(s)."
)
result = {
"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,
}
break
future = executor.submit(ctx.run, execute_with_context)
try:
result = future.result(timeout=tool_timeout)
except concurrent.futures.TimeoutError:
@@ -1904,11 +1620,6 @@ class ToolExecutionMixin:
``_atrigger_after_agent_hook`` without any async-safety change.
"""
from ..hooks import HookEvent, AfterAgentInput
# AfterAgentInput.response is typed/serialized as a string (its
# to_dict truncates via ``response[:500]``). Structured-output turns
# (output_pydantic) return a Pydantic instance, so coerce to str here
# to keep string-only hooks (e.g. command hooks) from crashing while
# the public chat() return value stays the model.
return AfterAgentInput(
session_id=getattr(self, '_session_id', 'default'),
cwd=os.getcwd(),
@@ -1916,7 +1627,7 @@ class ToolExecutionMixin:
timestamp=str(time.time()),
agent_name=self.name,
prompt=prompt if isinstance(prompt, str) else str(prompt),
response=response if isinstance(response, str) else (str(response) if response else ""),
response=response or "",
tools_used=tools_used or [],
total_tokens=0,
execution_time_ms=(time.time() - start_time) * 1000
@@ -1933,16 +1644,7 @@ class ToolExecutionMixin:
# Auto-memory extraction (opt-in via MemoryConfig(auto_memory=True))
if response:
prompt_str = prompt if isinstance(prompt, str) else str(prompt)
if getattr(self, "_auto_memory", False):
self._process_auto_memory(prompt_str, str(response))
else:
# When memory is enabled but auto_memory extraction is off, still
# persist the raw turn so a later Agent sharing the same memory
# (e.g. memory={"user_id": uid}) can recall it. Without this the
# sync chat/run/start path wrote nothing and cross-session recall
# silently returned empty (issue #5595). Mirrors the write already
# done on the async unified path.
self._persist_memory_turn(prompt_str, str(response))
self._process_auto_memory(prompt_str, str(response))
# Auto-learning extraction (opt-in via LearnConfig(mode=LearnMode.AGENTIC))
self._process_auto_learning()
@@ -2196,30 +1898,6 @@ class ToolExecutionMixin:
tail = text[-tail_limit:] if tail_limit > 0 else ""
return f"{head}\n...[{len(text):,} chars, showing first/last portions]...\n{tail}"
@staticmethod
def _approval_backend_accepts_timeout(backend) -> bool:
"""Return True if ``backend.request_approval_sync`` accepts a ``timeout`` kwarg.
Determined by signature introspection so the caller can route the
per-call timeout without a call-and-catch that would re-invoke a
backend whose *own* body raised ``TypeError``. Backends that accept
``**kwargs`` are treated as accepting ``timeout``. On any introspection
failure we conservatively return False (call without ``timeout``).
"""
method = getattr(backend, 'request_approval_sync', None)
if method is None:
return False
try:
sig = inspect.signature(method)
except (TypeError, ValueError):
return False
params = sig.parameters
if 'timeout' in params:
return True
return any(
p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()
)
def _resolve_approval_decision(self, tool_name: str, tool_args: dict, is_async: bool = False):
"""Shared approval logic for both sync and async paths.
@@ -2393,51 +2071,41 @@ class ToolExecutionMixin:
return _ask_backend_async()
else:
# Sync path - handle timeout and sync/async backend compatibility.
#
# The per-call timeout is passed explicitly rather than
# mutated onto the (often SHARED) backend object. Mutating
# ``backend._timeout`` here leaked between agents that pass
# one backend instance: an overlapping call captured the
# first agent's temporary value as its "original" and
# restored it, leaving the backend permanently wrong.
# Sync path - handle timeout and sync/async backend compatibility
cfg_timeout = getattr(self, '_approval_timeout', 0)
# Compute effective timeout from agent configuration.
# None -> indefinite wait
# > 0 -> explicit per-call timeout
# == 0 -> backend default (or 60s fallback)
orig_timeout = None
if cfg_timeout is None:
effective_timeout = None
orig_timeout = getattr(backend, '_timeout', None)
if orig_timeout is not None:
backend._timeout = 86400 * 365
elif cfg_timeout > 0:
effective_timeout = cfg_timeout
else:
effective_timeout = getattr(backend, '_timeout', 60)
if hasattr(backend, 'request_approval_sync'):
# Decide *up front* whether the backend accepts a
# ``timeout`` kwarg by inspecting its signature, rather
# than calling and catching ``TypeError``. A try/except
# would re-invoke the backend when the *user callback*
# inside it raised its own ``TypeError`` — duplicating
# side effects and masking the real error. Signature
# introspection is side-effect free.
if self._approval_backend_accepts_timeout(backend):
return _remember(backend.request_approval_sync(
request, timeout=effective_timeout
orig_timeout = getattr(backend, '_timeout', None)
if orig_timeout is not None:
backend._timeout = cfg_timeout
try:
if hasattr(backend, 'request_approval_sync'):
return _remember(backend.request_approval_sync(request))
else:
# Use the shared utility to avoid code duplication and handle timeout correctly
from ..approval.utils import run_coroutine_safely
# Compute effective timeout from agent configuration
if cfg_timeout is None:
effective_timeout = None # indefinite wait
elif cfg_timeout > 0:
effective_timeout = cfg_timeout
else:
# cfg_timeout == 0: use backend default or fallback
effective_timeout = getattr(backend, '_timeout', 60)
return _remember(run_coroutine_safely(
backend.request_approval(request),
timeout=effective_timeout
))
# Legacy backend whose request_approval_sync has no
# timeout kwarg: honour its own default, no mutation of
# (often shared) backend state.
return _remember(backend.request_approval_sync(request))
else:
# Use the shared utility to avoid code duplication and handle timeout correctly
from ..approval.utils import run_coroutine_safely
return _remember(run_coroutine_safely(
backend.request_approval(request),
timeout=effective_timeout
))
finally:
if orig_timeout is not None and hasattr(backend, '_timeout'):
backend._timeout = orig_timeout
else:
if is_async:
# For async, wrap the decision in a coroutine
@@ -2846,42 +2514,6 @@ class ToolExecutionMixin:
logging.debug("permission manager is_denied failed for %s: %s", target, e)
return None
def _check_permission_manager_doom_loop(self, function_name, arguments=None):
"""Return an error dict if the PermissionManager detects a doom loop.
Deliver the doom-loop detection the ``PermissionManager`` advertises:
record this call in its detector and deny on a genuine loop. Previously
``check_doom_loop()`` was unreachable dead code so a configured manager
gave zero doom-loop protection.
This is intentionally **separate** from :meth:`_check_permission_manager_deny`
and invoked exactly **once per approved tool call**, with the final
(post-approval, possibly rewritten) arguments. Folding it into the deny
gate would record the call on every gate pass — counting a rejected
attempt, or counting an approved call twice when its args are rewritten —
which could trip the threshold early and deny legitimate repeats.
"""
manager = getattr(self, "_permission_manager", None)
if manager is None:
return None
check_loop = getattr(manager, "check_doom_loop", None)
if not callable(check_loop):
return None
try:
loop_result = check_loop(function_name, arguments)
if loop_result is not None and getattr(loop_result, "is_loop", False):
return {
"error": (
f"Tool '{function_name}' blocked: "
f"{getattr(loop_result, 'reason', 'doom loop detected')}"
),
"permission_denied": True,
"loop_blocked": True,
}
except Exception as e: # noqa: BLE001
logging.debug("permission manager doom-loop check failed for %s: %s", function_name, e)
return None
def _is_bypass_mode(self) -> bool:
"""Return ``True`` when ``PermissionMode.BYPASS`` is active.
@@ -2943,13 +2575,6 @@ class ToolExecutionMixin:
if manager_denial is not None:
return manager_denial
# Record the (approved, final-args) call in the doom-loop detector once,
# after approval, so rejected attempts and arg rewrites are not counted.
if not self._is_bypass_mode():
loop_denial = self._check_permission_manager_doom_loop(function_name, arguments)
if loop_denial is not None:
return loop_denial
from ..approval import get_approval_registry
get_approval_registry().mark_approved(
function_name, arguments, agent_name=getattr(self, "name", None),
@@ -2996,13 +2621,6 @@ class ToolExecutionMixin:
if manager_denial is not None:
return manager_denial
# Record the (approved, final-args) call in the doom-loop detector once,
# after approval, so rejected attempts and arg rewrites are not counted.
if not self._is_bypass_mode():
loop_denial = self._check_permission_manager_doom_loop(function_name, arguments)
if loop_denial is not None:
return loop_denial
from ..approval import get_approval_registry
get_approval_registry().mark_approved(
function_name, arguments, agent_name=getattr(self, "name", None),
@@ -454,10 +454,7 @@ def _build_execution_context(agents_instance, task_id, skip_memory_init=False):
if llm is None:
llm = getattr(executor_agent, 'llm', None)
if llm and hasattr(llm, 'set_current_agent'):
# Pass the agent's stable per-instance id so token usage is aggregated
# by identity, not by (collidable) display name.
agent_scope_id = getattr(executor_agent, '_approval_scope_id', None)
llm.set_current_agent(executor_agent.display_name, agent_id=agent_scope_id)
llm.set_current_agent(executor_agent.display_name)
# Place the agent's tools BEFORE snapshotting them. This snapshot is passed
# to chat() as an explicit `tools=`, which wins over `agent.tools` -- so if
@@ -2382,20 +2379,6 @@ class AgentTeam(SpawnAnnounceProtocol):
def run_all_tasks(self):
"""Synchronous version of run_all_tasks method"""
# The sync entry point (.start()/.run()) drives tasks one at a time and
# cannot batch async_execution=True tasks via asyncio.gather the way
# arun_all_tasks does. Warn once so the parallel fan-out feature the
# constructor docstring/ValueError advertise doesn't silently degrade to
# sequential execution -- point the caller at the async entry point.
if self.process in ("workflow", "sequential") and any(
getattr(t, 'async_execution', False) for t in self.tasks.values()
):
logger.warning(
"One or more tasks have async_execution=True, but the synchronous "
"entry point (.start()/.run()) runs them sequentially -- use "
"`await team.astart()` for parallel fan-out."
)
process = Process(
tasks=self.tasks,
agents=self.agents,
@@ -3275,45 +3258,21 @@ class AgentTeam(SpawnAnnounceProtocol):
if getattr(agent, "name", None)
}
def _own_agent_ids(self) -> set:
"""Stable per-instance identities of this instance's own agents.
Used as the collision-free key for token-report scoping so two agents
that share a display name (e.g. the default ``"Agent"``) are never
merged. Falls back to names elsewhere when ids are unavailable.
"""
return {
getattr(agent, "_approval_scope_id", None)
for agent in (self.agents or [])
if getattr(agent, "_approval_scope_id", None)
}
def _scoped_recent_interactions(
self, own_names: set, own_ids: Optional[set] = None
) -> List[Dict[str, Any]]:
def _scoped_recent_interactions(self, own_names: set) -> List[Dict[str, Any]]:
"""Return this instance's own recorded interactions, most-recent last.
The process-wide collector keeps a bounded window of every agent's
interactions; filter it to this instance's agents so per-model totals,
interaction counts, and the detailed report never leak another
concurrent instance's records. Filtering prefers the collision-free
``agent_id`` when present, falling back to the display name.
interactions; filter it to this instance's agent names so per-model
totals, interaction counts, and the detailed report never leak another
concurrent instance's records.
"""
own_ids = own_ids or set()
try:
interactions = get_token_collector().get_recent_interactions(
limit=get_token_collector()._max_recent
)
except Exception:
interactions = get_token_collector().get_recent_interactions(limit=100)
def _matches(i: Dict[str, Any]) -> bool:
agent_id = i.get("agent_id")
if agent_id is not None and own_ids:
return agent_id in own_ids
return i.get("agent") in own_names
return [i for i in interactions if _matches(i)]
return [i for i in interactions if i.get("agent") in own_names]
def _scoped_token_summary(self) -> Dict[str, Any]:
"""Return the global token summary filtered to this instance's agents.
@@ -3321,54 +3280,25 @@ class AgentTeam(SpawnAnnounceProtocol):
The process-wide TokenCollector aggregates every agent in the process,
so two concurrent PraisonAIAgents instances would otherwise read each
other's totals. Scope every reported field — ``by_agent``, ``by_model``,
``total_interactions`` and the totals — to this instance's own agents so
per-instance/per-tenant cost accounting reports only its own spend.
Scoping prefers the collector's collision-free ``by_agent_id``
breakdown; ``by_model`` and ``total_interactions`` are re-derived from
this instance's own interaction records (the aggregate cannot be split
by model per agent). Falls back to the unfiltered summary when this
instance has no agents to scope by.
``total_interactions`` and the totals — to this instance's own agent
names so per-instance/per-tenant cost accounting reports only its own
spend. ``by_agent`` is taken from the collector's aggregate breakdown;
``by_model`` and ``total_interactions`` are re-derived from this
instance's own interaction records (the aggregate cannot be split by
model per agent). Falls back to the unfiltered summary when this
instance has no named agents to scope by.
"""
summary = get_token_collector().get_session_summary()
own_names = self._own_agent_names()
own_ids = self._own_agent_ids()
by_agent = summary.get("by_agent", {})
by_agent_id = summary.get("by_agent_id", {})
# No agents to scope by (or nothing tracked yet): fall back to the
# No named agents to scope by (or nothing tracked yet): fall back to the
# unfiltered summary rather than silently reporting zeros.
if (not own_names and not own_ids) or (not by_agent and not by_agent_id):
if not own_names or not by_agent:
return summary
# Prefer the collision-free per-identity breakdown; fall back to
# name-based scoping when ids are unavailable (backward compatible).
scoped_by_agent: Dict[str, Dict[str, int]] = {}
if own_ids and by_agent_id:
for key, entry in by_agent_id.items():
name = entry.get("name") or key
# Match by stable id (collision-free) OR by display name for
# id-less buckets — e.g. the native OpenAI-client path records
# ``agent`` without an ``agent_id``, so the collector keys that
# bucket by name (``key == name``). Dropping those would zero out
# real OpenAI usage even though the interaction and per-model
# rows still show it (Greptile P1). We treat a bucket as id-less
# when its key equals its own display name.
is_id_match = key in own_ids
is_idless_name_match = key == name and name in own_names
if not (is_id_match or is_idless_name_match):
continue
metrics = entry.get("metrics", {})
if name in scoped_by_agent:
scoped_by_agent[name] = {
k: scoped_by_agent[name].get(k, 0) + v
for k, v in metrics.items()
if isinstance(v, (int, float))
}
else:
scoped_by_agent[name] = dict(metrics)
else:
scoped_by_agent = {
name: metrics for name, metrics in by_agent.items() if name in own_names
}
scoped_by_agent = {
name: metrics for name, metrics in by_agent.items() if name in own_names
}
totals = {
"input_tokens": 0, "output_tokens": 0, "cached_tokens": 0,
"reasoning_tokens": 0, "audio_input_tokens": 0, "audio_output_tokens": 0,
@@ -3384,7 +3314,7 @@ class AgentTeam(SpawnAnnounceProtocol):
# process-global summary (which mixes in other instances' models and
# counts). The collector's ``by_agent`` aggregate cannot be split by
# model, so the per-interaction window is the correct source.
own_interactions = self._scoped_recent_interactions(own_names, own_ids)
own_interactions = self._scoped_recent_interactions(own_names)
by_model: Dict[str, Dict[str, int]] = {}
for interaction in own_interactions:
model = interaction.get("model")
@@ -3418,11 +3348,10 @@ class AgentTeam(SpawnAnnounceProtocol):
summary = self._scoped_token_summary()
own_names = self._own_agent_names()
own_ids = self._own_agent_ids()
if own_names or own_ids:
if own_names:
# Only surface this instance's own interactions so the detailed
# report never exposes another concurrent instance's records.
recent = self._scoped_recent_interactions(own_names, own_ids)[-20:]
recent = self._scoped_recent_interactions(own_names)[-20:]
else:
recent = get_token_collector().get_recent_interactions(limit=20)
@@ -13,7 +13,7 @@ from __future__ import annotations
import asyncio
import re
from praisonaiagents._logging import get_logger
from typing import Any, Optional
from typing import Any
from .protocols import ApprovalDecision, ApprovalRequest
@@ -175,7 +175,7 @@ class AutoApproveBackend:
async def request_approval(self, request: ApprovalRequest) -> ApprovalDecision:
return ApprovalDecision(approved=True, reason="auto-approved", approver="system")
def request_approval_sync(self, request: ApprovalRequest, timeout: Optional[float] = None) -> ApprovalDecision:
def request_approval_sync(self, request: ApprovalRequest) -> ApprovalDecision:
return ApprovalDecision(approved=True, reason="auto-approved", approver="system")
class ConsoleBackend:
@@ -293,7 +293,7 @@ class ConsoleBackend:
return (True, "always", suggested or None, None)
return (True, "once", None, None)
def request_approval_sync(self, request: ApprovalRequest, timeout: Optional[float] = None) -> ApprovalDecision:
def request_approval_sync(self, request: ApprovalRequest) -> ApprovalDecision:
"""Synchronous approval via Rich console prompt."""
try:
approved, scope, scope_pattern, feedback = self._prompt_user(request)
@@ -474,12 +474,10 @@ class AgentApproval:
reason=f"Agent approval error: {e}",
)
def request_approval_sync(self, request: ApprovalRequest, timeout: Optional[float] = None) -> ApprovalDecision:
def request_approval_sync(self, request: ApprovalRequest) -> ApprovalDecision:
"""Synchronous wrapper."""
from .utils import run_coroutine_safely
return run_coroutine_safely(
self.request_approval(request), timeout=60 if timeout is None else timeout
)
return run_coroutine_safely(self.request_approval(request), timeout=60)
class CallbackBackend:
"""Wraps a legacy ``(function_name, arguments, risk_level) -> ApprovalDecision`` callback
@@ -491,7 +489,7 @@ class CallbackBackend:
def __init__(self, callback):
self._callback = callback
def request_approval_sync(self, request: ApprovalRequest, timeout: Optional[float] = None) -> ApprovalDecision:
def request_approval_sync(self, request: ApprovalRequest) -> ApprovalDecision:
result = self._callback(request.tool_name, request.arguments, request.risk_level)
if isinstance(result, ApprovalDecision):
return result
@@ -189,18 +189,6 @@ class ApprovalProtocol(Protocol):
"""
...
# Backends MAY also implement an optional synchronous entry point:
#
# def request_approval_sync(self, request, timeout=None) -> ApprovalDecision
#
# It is deliberately NOT declared on this runtime-checkable Protocol so an
# async-only backend still satisfies ``isinstance(b, ApprovalProtocol)``.
# When present, the per-call ``timeout`` (seconds; ``None`` = wait forever)
# is passed explicitly by the caller so shared backend state is never
# mutated — mutating a shared backend's timeout leaks between agents.
# Legacy sync backends without the ``timeout`` parameter are still
# supported; the caller falls back to a no-argument call on ``TypeError``.
@runtime_checkable
class ApprovalStoreProtocol(Protocol):
@@ -69,13 +69,6 @@ class BackgroundRunner:
"""
self.config = config or BackgroundConfig()
self._tasks: Dict[str, BackgroundTask] = {}
# Guards every read/write of ``_tasks``. This runner is a process-wide
# singleton (``get_background_runner``) whose sync wrappers insert on a
# separate background-loop thread via ``run_coroutine_threadsafe`` while
# other surfaces read/mutate the dict on their own thread; without this
# a concurrent insert + iterate/delete can raise
# ``RuntimeError: dictionary changed size during iteration``.
self._tasks_lock: threading.Lock = threading.Lock()
# Lazily create asyncio.Semaphore() to avoid Python 3.9 event loop issues
self._semaphore: Optional[asyncio.Semaphore] = None
self._semaphore_lock: threading.Lock = threading.Lock() # Thread-safe initialization
@@ -93,20 +86,17 @@ class BackgroundRunner:
@property
def tasks(self) -> List[BackgroundTask]:
"""Get all tasks."""
with self._tasks_lock:
return list(self._tasks.values())
return list(self._tasks.values())
@property
def running_tasks(self) -> List[BackgroundTask]:
"""Get currently running tasks."""
with self._tasks_lock:
return [t for t in self._tasks.values() if t.is_running]
return [t for t in self._tasks.values() if t.is_running]
@property
def pending_tasks(self) -> List[BackgroundTask]:
"""Get pending tasks."""
with self._tasks_lock:
return [t for t in self._tasks.values() if t.status == TaskStatus.PENDING]
return [t for t in self._tasks.values() if t.status == TaskStatus.PENDING]
async def submit(
self,
@@ -143,8 +133,7 @@ class BackgroundRunner:
}
)
with self._tasks_lock:
self._tasks[task.id] = task
self._tasks[task.id] = task
# Create the execution coroutine
async def execute():
@@ -246,8 +235,7 @@ class BackgroundRunner:
def get_task(self, task_id: str) -> Optional[BackgroundTask]:
"""Get a task by ID."""
with self._tasks_lock:
return self._tasks.get(task_id)
return self._tasks.get(task_id)
async def cancel_task(self, task_id: str) -> bool:
"""
@@ -259,8 +247,7 @@ class BackgroundRunner:
Returns:
True if cancelled, False if not found or already completed
"""
with self._tasks_lock:
task = self._tasks.get(task_id)
task = self._tasks.get(task_id)
if task is None:
return False
@@ -289,12 +276,10 @@ class BackgroundRunner:
Returns:
List of completed tasks
"""
with self._tasks_lock:
futures = [t._future for t in self._tasks.values() if t._future]
futures = [t._future for t in self._tasks.values() if t._future]
if not futures:
with self._tasks_lock:
return list(self._tasks.values())
return list(self._tasks.values())
try:
await asyncio.wait_for(
@@ -304,18 +289,16 @@ class BackgroundRunner:
except asyncio.TimeoutError:
pass
with self._tasks_lock:
return list(self._tasks.values())
return list(self._tasks.values())
def clear_completed(self):
"""Remove completed tasks from tracking."""
with self._tasks_lock:
completed_ids = [
task_id for task_id, task in self._tasks.items()
if task.is_completed
]
for task_id in completed_ids:
del self._tasks[task_id]
completed_ids = [
task_id for task_id, task in self._tasks.items()
if task.is_completed
]
for task_id in completed_ids:
del self._tasks[task_id]
def list_tasks(self, status: Optional[TaskStatus] = None) -> List[Dict[str, Any]]:
"""
@@ -327,8 +310,7 @@ class BackgroundRunner:
Returns:
List of task dictionaries
"""
with self._tasks_lock:
tasks = list(self._tasks.values())
tasks = self._tasks.values()
if status:
tasks = [t for t in tasks if t.status == status]
@@ -130,17 +130,6 @@ class BotConfig:
allow_silence: bool = False
silence_token: Optional[str] = None
# Bot-to-bot inbound (#5062). Off by default: a bot-authored inbound
# message is dropped as before. When ``allow_bots=True`` the channel accepts
# bot-authored messages and the runaway-loop guard (``BotLoopGuard``) is
# auto-active, tracking each participant pair and suppressing further replies
# once ``bot_loop_protection``'s sliding-window budget is exceeded.
# ``bot_loop_protection`` is a plain dict (enabled / max_events_per_window /
# window_seconds / cooldown_seconds) coerced into a ``BotLoopPolicy``; ``None``
# uses the default budget.
allow_bots: bool = False
bot_loop_protection: Optional[Dict[str, Any]] = None
def __post_init__(self) -> None:
if self.unknown_user_policy not in {"deny", "pair", "allow"}:
raise ValueError(
@@ -187,8 +176,6 @@ class BotConfig:
"stream_edit_interval_ms": self.stream_edit_interval_ms,
"allow_silence": self.allow_silence,
"silence_token": self.silence_token,
"allow_bots": self.allow_bots,
"bot_loop_protection": self.bot_loop_protection,
"metadata": self.metadata,
}
@@ -138,31 +138,15 @@ class BotLoopPolicy:
"""Build a policy from a config mapping (YAML/kwargs), tolerating None.
Unknown keys are ignored so a forward-compatible config never breaks
construction. ``None`` yields the enabled default policy. Malformed
numeric values (e.g. ``max_events_per_window: "abc"``) fall back to the
field default rather than raising, so a config typo cannot crash the
inbound path on the first bot message — the guard stays fail-safe.
construction. ``None`` yields the enabled default policy.
"""
if not data:
return cls()
def _as_int(value: object, default: int) -> int:
try:
return int(value) # type: ignore[arg-type]
except (TypeError, ValueError):
return default
def _as_float(value: object, default: float) -> float:
try:
return float(value) # type: ignore[arg-type]
except (TypeError, ValueError):
return default
return cls(
enabled=bool(data.get("enabled", True)),
max_events_per_window=_as_int(data.get("max_events_per_window", 20), 20),
window_seconds=_as_float(data.get("window_seconds", 60.0), 60.0),
cooldown_seconds=_as_float(data.get("cooldown_seconds", 60.0), 60.0),
max_events_per_window=int(data.get("max_events_per_window", 20)),
window_seconds=float(data.get("window_seconds", 60.0)),
cooldown_seconds=float(data.get("cooldown_seconds", 60.0)),
)
+23 -15
View File
@@ -267,13 +267,17 @@ class EventBus:
The published Event object
"""
type_str = event_type.value if isinstance(event_type, EventType) else event_type
# No pre-history no-subscriber shortcut here, matching publish_event's
# already-fixed shape: an event that was published happened, whether or
# not anyone was listening, and skipping the history append below would
# make it vanish from get_history(). The append is O(1) under a lock and
# the list is capped, so the only work worth skipping is subscriber
# matching/dispatch -- which the guard after the append still skips.
# Fast path: if no subscribers and no durable sinks, return a minimal
# event without expensive operations.
if not self._subscribers and not self._sinks:
return Event(
type=type_str,
data=data or {},
source=source,
metadata=metadata or {},
)
event = Event(
type=type_str,
data=data or {},
@@ -332,20 +336,24 @@ class EventBus:
"""
Get recent event history.
Every event passed to ``publish``/``publish_event``/``publish_async`` is
recorded here (capped at ``_max_history``), regardless of whether any
subscriber was registered at publish time. History is a debugging record
of what the system did; it does not silently depend on who was watching.
Both the sync and async publish paths append before any no-subscriber
fast path runs, so publishing to a bus with nothing subscribed still
shows up here.
Only events published while at least one subscriber was registered are
recorded. #2066 added a fast path that skips history (and the lock, and
Event construction) when ``has_subscribers`` is False, because memory
and sub-agent lifecycle publishes were paying uuid4 + lock + append on
every call with nothing listening.
The consequence is worth stating plainly, because it is surprising in
the case you would most want history -- debugging after the fact, with
nothing subscribed: publishing three events to a bus with no
subscribers and then calling this returns an empty list, not three
events. Subscribe (even a no-op) before publishing if you need a record.
Args:
event_type: Optional filter by event type
limit: Maximum number of events to return
Returns:
List of recent events published (subscribed or not)
List of recent events published while subscribers existed
"""
with self._lock:
events = self._event_history.copy()
@@ -83,8 +83,7 @@ class CheckpointService:
auto_checkpoint: bool = True,
max_checkpoints: int = 100,
user_name: str = "PraisonAI Checkpoints",
user_email: str = "[email protected]",
max_file_size: int = 2 * 1024 * 1024,
user_email: str = "[email protected]"
):
"""
Initialize the checkpoint service.
@@ -97,8 +96,6 @@ class CheckpointService:
max_checkpoints: Maximum checkpoints to keep
user_name: Git user.name for commits (default: "PraisonAI Checkpoints")
user_email: Git user.email for commits (default: "[email protected]")
max_file_size: Skip staging files larger than this many bytes
(default 2 MB); 0 disables the cap.
"""
self.config = CheckpointConfig(
workspace_dir=workspace_dir,
@@ -107,14 +104,10 @@ class CheckpointService:
auto_checkpoint=auto_checkpoint,
max_checkpoints=max_checkpoints,
user_name=user_name,
user_email=user_email,
max_file_size=max_file_size,
user_email=user_email
)
self._initialized = False
# Save counter used to run `git gc --auto` on a bounded cadence so the
# shadow object store is periodically pruned during long sessions.
self._saves_since_gc = 0
self._checkpoints: List[Checkpoint] = []
self._event_handlers: Dict[CheckpointEvent, List[Callable]] = {
event: [] for event in CheckpointEvent
@@ -191,78 +184,6 @@ class CheckpointService:
return True
async def _get_project_objects_dir(self) -> Optional[str]:
"""
Return the object database directory of the project's real git repo if
the workspace is inside one, else None.
Uses ``--git-common-dir`` so worktrees resolve to the shared object
store. Best-effort: any failure (no git, not a repo) returns None so the
shadow repo falls back to the standalone path unchanged.
"""
try:
env = self._get_sanitized_env()
for flag in ("--git-common-dir", "--git-dir"):
process = await asyncio.create_subprocess_exec(
"git", "rev-parse", flag,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
cwd=self.config.workspace_dir,
)
stdout, _ = await asyncio.wait_for(process.communicate(), timeout=30)
if process.returncode != 0:
continue
git_dir = stdout.decode().strip()
if not git_dir:
continue
if not os.path.isabs(git_dir):
git_dir = os.path.join(self.config.workspace_dir, git_dir)
objects_dir = os.path.join(git_dir, "objects")
if os.path.isdir(objects_dir):
return os.path.realpath(objects_dir)
except Exception as e:
logger.debug(f"Could not resolve project git objects dir: {e}")
return None
async def _apply_large_repo_config(self) -> None:
"""
Apply large-repo git settings so index writes and untracked scans stay
fast. Each setting is best-effort: older git versions may not support
every key, and a failure must never break checkpointing.
"""
settings = [
("feature.manyFiles", "true"),
("index.version", "4"),
("core.untrackedCache", "true"),
("core.fsmonitor", "true"),
]
for key, value in settings:
try:
await self._run_git("config", key, value)
except Exception as e:
logger.debug(f"Skipping large-repo config {key}={value}: {e}")
async def _seed_from_project_objects(self, objects_dir: str) -> bool:
"""
Point the shadow object store at the project's objects via a file-based
``objects/info/alternates`` so new blobs resolve against existing
objects instead of being re-hashed and re-stored.
A file is used (not GIT_ALTERNATE_OBJECT_DIRECTORIES) because
``_get_sanitized_env`` strips that env var. Returns True on success.
"""
try:
info_dir = os.path.join(self.git_dir, "objects", "info")
os.makedirs(info_dir, exist_ok=True)
with open(os.path.join(info_dir, "alternates"), "w") as f:
f.write(f"{objects_dir}\n")
logger.debug(f"Seeded shadow git objects from {objects_dir}")
return True
except Exception as e:
logger.debug(f"Could not seed shadow git from project objects: {e}")
return False
async def _init_shadow_git(self) -> bool:
"""Initialize the shadow git repository."""
try:
@@ -275,17 +196,7 @@ class CheckpointService:
# Set worktree to workspace
await self._run_git("config", "core.worktree", self.config.workspace_dir)
# Seed objects from the project's real repo (if any) so the first
# checkpoint only stores changed files, and tune for large repos.
# Both are best-effort: no real repo → unchanged standalone path.
objects_dir = await self._get_project_objects_dir()
if objects_dir and objects_dir != os.path.realpath(
os.path.join(self.git_dir, "objects")
):
await self._seed_from_project_objects(objects_dir)
await self._apply_large_repo_config()
# Create exclude file
exclude_dir = os.path.join(self.git_dir, "info")
os.makedirs(exclude_dir, exist_ok=True)
@@ -348,64 +259,6 @@ class CheckpointService:
env.pop(var, None)
return env
async def _unstage_oversized_files(self) -> None:
"""
Remove files larger than ``config.max_file_size`` from the index so they
are not committed into the shadow store. Best-effort: any failure leaves
staging exactly as ``git add -A`` produced it and never breaks the turn.
"""
max_size = self.config.max_file_size
if not max_size or max_size <= 0:
return
try:
# NUL-delimited so paths with spaces/newlines are handled safely.
staged = await self._run_git("diff", "--cached", "--name-only", "-z")
for path in staged.split("\0"):
if not path:
continue
abs_path = os.path.join(self.config.workspace_dir, path)
try:
# Use lstat so symlinks are measured by the link itself, not
# their target: a small link to a large file must stay in the
# checkpoint. Skip anything that isn't a regular file.
if os.path.islink(abs_path) or not os.path.isfile(abs_path):
continue
if os.lstat(abs_path).st_size > max_size:
# Keep the file on disk; only drop it from the index.
await self._run_git("rm", "--cached", "--quiet", "--", path)
logger.debug(f"Skipped oversized file from checkpoint: {path}")
except Exception as e:
logger.debug(f"Could not evaluate size of {path}: {e}")
except Exception as e:
logger.debug(f"Oversized-file skip pass failed (continuing): {e}")
async def _maybe_gc(self, cadence: int = 20) -> None:
"""
Every ``cadence`` saves, pack the shadow repo and make it self-contained.
We seed the shadow store from the project's objects via ``alternates``
for a fast first checkpoint, but a borrowed object can disappear if the
project repo later prunes or repacks it — which would make checkpoint
history unreadable. To keep ``/undo``/``/revert`` durable we periodically
run ``git repack -a -d`` (without ``--local``) so every reachable object,
including borrowed ones, is copied into a local pack the shadow store
owns. ``git gc --auto`` still handles pruning/packing of loose objects.
Best-effort: failures are swallowed so a gc issue never breaks a turn.
"""
self._saves_since_gc += 1
if self._saves_since_gc < cadence:
return
self._saves_since_gc = 0
try:
# -a: pack all reachable objects; -d: drop redundant packs. Omitting
# --local pulls borrowed (alternate) objects into the local pack so
# the shadow store no longer depends on the project object database.
await self._run_git("repack", "-a", "-d", "--quiet")
await self._run_git("gc", "--auto", "--quiet")
except Exception as e:
logger.debug(f"Shadow-store gc skipped: {e}")
async def save(
self,
@@ -442,13 +295,7 @@ class CheckpointService:
try:
# Stage all changes
await self._run_git("add", "-A")
# Drop oversized files from the index so the shadow store never
# copies giant blobs (build outputs, model weights, media). This
# mirrors the existing "never fail the turn" contract: any error
# here is swallowed and staging proceeds as before.
await self._unstage_oversized_files()
# Check if there are changes
try:
await self._run_git("diff", "--cached", "--quiet")
@@ -482,12 +329,7 @@ class CheckpointService:
# Prune old checkpoints if needed
await self._prune_checkpoints()
# Periodically let git pack loose objects and drop unreachable ones
# so the shadow store doesn't leak disk over a long session. Bounded
# cadence keeps the per-turn cost near zero; best-effort only.
await self._maybe_gc()
self._emit(CheckpointEvent.CHECKPOINT_CREATED, checkpoint)
logger.info(f"Created checkpoint: {checkpoint.short_id} - {message}")
@@ -40,11 +40,6 @@ class CheckpointConfig:
max_checkpoints: int = 100 # Maximum checkpoints to keep
user_name: str = field(default_factory=lambda: os.getenv("PRAISONAI_GIT_USER_NAME", "PraisonAI Checkpoints"))
user_email: str = field(default_factory=lambda: os.getenv("PRAISONAI_GIT_USER_EMAIL", "[email protected]"))
# Skip staging files larger than this (bytes); 0 disables the cap. Kept
# after the identity fields so existing positional construction of
# CheckpointConfig(workspace, storage, enabled, auto, max_checkpoints,
# user_name, user_email) is unaffected.
max_file_size: int = 2 * 1024 * 1024
exclude_patterns: List[str] = field(default_factory=lambda: [
".git",
".praisonai",
@@ -1395,21 +1395,13 @@ class HooksConfig:
# Middleware list
middleware: List[Any] = field(default_factory=list)
# Optional per-agent HookRegistry for event-based hooks (BEFORE_TOOL,
# AFTER_AGENT, ...). When None the agent uses the process-wide default
# registry (so plugins bridged into it still fire). Set this to a private
# ``HookRegistry()`` to keep event hooks registered for THIS agent from
# firing on unrelated agents that share the default registry.
registry: Optional[Any] = None
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"on_step": str(self.on_step) if self.on_step else None,
"on_tool_call": str(self.on_tool_call) if self.on_tool_call else None,
"middleware": [str(m) for m in self.middleware],
"registry": type(self.registry).__name__ if self.registry is not None else None,
}
@@ -7,16 +7,9 @@ another tool's harness, and results can be diffed/aggregated across tools.
Mapping (1:1 with EvalPort's ``Suite`` / ``ResultSet`` concepts):
EvalPackage <-> EvalPort Suite (id/name/description + test_cases + graders)
EvalCase <-> EvalPort TestCase (id / input / expected_output / grader refs)
EvalReport -> EvalPort ResultSet (summary + per-case grader_results)
Per-criterion scores map to native EvalPort ``grader_results`` (one per criterion,
matching the criterion graders referenced by the Suite). PraisonAI-only fields
with no EvalPort slot (package version, thresholds, seed, trajectory records, a
nonpositive ``timeout_seconds`` that the spec's ``timeout_ms`` cannot represent,
...) are kept under the ``praisonai`` key of the relevant ``metadata`` object so
round trips stay lossless.
EvalPackage <-> EvalPort Suite (name/description/version + cases + thresholds)
EvalCase <-> EvalPort test case (input / expected_output / grader refs)
EvalReport -> EvalPort ResultSet (aggregate + per-case pass/score/latency)
Design: this module is intentionally **dependency-free**. It emits and consumes
plain spec-shaped dicts (JSON), so the core SDK does not take on ``evalport-sdk``
@@ -34,90 +27,53 @@ Example:
>>> pkg2.name == pkg.name
True
"""
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List
from .package import EvalCase, EvalPackage, EvalReport
# EvalPort spec version the emitted documents conform to (semver, per spec).
EVALPORT_SPEC_VERSION = "1.0.0-rc.5"
# Grader that carries PraisonAI's own per-case verdict (EvalResult.passed/score).
# EvalPort requires every test case to reference at least one grader.
VERDICT_GRADER_ID = "praisonai_verdict"
_HANDLER = "praisonaiagents.eval"
_META_KEY = "praisonai"
def _grader_def(grader_id: str, description: str) -> Dict[str, Any]:
return {
"id": grader_id,
"type": "custom",
"params": {"handler": _HANDLER},
"description": description,
}
EVALPORT_SPEC_VERSION = "1.0"
def _case_to_evalport(case: EvalCase) -> Dict[str, Any]:
"""Map an ``EvalCase`` to an EvalPort TestCase dict.
"""Map an ``EvalCase`` to an EvalPort test-case dict.
``criteria`` become grader references (defined at suite level); a case
without criteria references the verdict grader. ``timeout_seconds`` maps to
the spec's ``timeout_ms``; user ``metadata`` is passed through untouched.
The native ``timeout_seconds`` is emitted as a dedicated top-level field so
it never collides with a user-supplied ``metadata["timeout_seconds"]`` key,
keeping round trips lossless for both.
"""
# A case's own criteria are its graders; the verdict grader is always
# referenced too so PraisonAI's pass/score verdict has a home in the
# ResultSet (spec requires >= 1 grader per case).
graders = [VERDICT_GRADER_ID] + [c for c in case.criteria if c != VERDICT_GRADER_ID]
item: Dict[str, Any] = {
"id": case.name,
"input": case.input,
"graders": list(dict.fromkeys(graders)),
}
if case.expected is not None:
item["expected_output"] = case.expected
if case.metadata:
item["metadata"] = dict(case.metadata)
if case.timeout_seconds is not None and case.timeout_seconds > 0:
item["timeout_ms"] = max(1, int(round(case.timeout_seconds * 1000)))
elif case.timeout_seconds is not None:
# Nonpositive timeout is not spec-representable (timeout_ms >= 1); keep
# the raw value under the reserved namespace so import restores it
# instead of silently applying the 30s default.
meta = item.setdefault("metadata", {})
native = dict(meta.get(_META_KEY) or {})
native["timeout_seconds"] = case.timeout_seconds
meta[_META_KEY] = native
if case.criteria:
item["graders"] = list(case.criteria)
if case.timeout_seconds is not None:
item["timeout_seconds"] = case.timeout_seconds
metadata = dict(case.metadata or {})
if metadata:
item["metadata"] = metadata
return item
def _case_from_evalport(item: Dict[str, Any]) -> EvalCase:
"""Map an EvalPort TestCase dict to an ``EvalCase``.
"""Map an EvalPort test-case dict to an ``EvalCase``.
Reads ``timeout_ms`` (spec); falls back to a legacy top-level or
``metadata["timeout_seconds"]`` value, without mutating caller metadata.
Prefers the dedicated top-level ``timeout_seconds`` field; falls back to a
``metadata["timeout_seconds"]`` value for suites emitted by other tools,
without mutating the caller's metadata.
"""
metadata = dict(item.get("metadata") or {})
native = metadata.get(_META_KEY) or {}
if "timeout_ms" in item:
timeout = item["timeout_ms"] / 1000.0
elif "timeout_seconds" in item:
if "timeout_seconds" in item:
timeout = item["timeout_seconds"]
elif "timeout_seconds" in native:
timeout = native["timeout_seconds"] # nonpositive timeout preserved on export
else:
timeout = metadata.get("timeout_seconds", 30.0)
criteria = []
for g in item.get("graders") or item.get("criteria") or []:
gid = g.get("id") if isinstance(g, dict) else g
if gid and gid != VERDICT_GRADER_ID:
criteria.append(gid)
return EvalCase(
name=item.get("id") or item.get("name") or "case",
input=item.get("input", ""),
expected=item.get("expected_output", item.get("expected")),
criteria=criteria,
criteria=list(item.get("graders") or item.get("criteria") or []),
metadata=metadata,
timeout_seconds=timeout,
)
@@ -132,28 +88,16 @@ def to_evalport(package: EvalPackage) -> Dict[str, Any]:
Returns:
A spec-shaped dict suitable for ``openeval.validate.validate_suite()``.
"""
criteria = dict.fromkeys(c for case in package.cases for c in case.criteria)
criteria.pop(VERDICT_GRADER_ID, None)
graders = [_grader_def(VERDICT_GRADER_ID, "PraisonAI per-case verdict (EvalResult.passed/score)")]
graders += [_grader_def(c, f"PraisonAI criterion: {c}") for c in criteria]
suite: Dict[str, Any] = {
"version": EVALPORT_SPEC_VERSION,
"id": package.name,
return {
"evalport_version": EVALPORT_SPEC_VERSION,
"kind": "suite",
"name": package.name,
"graders": graders,
"test_cases": [_case_to_evalport(c) for c in package.cases],
"metadata": {
_META_KEY: {
"version": package.version,
"thresholds": dict(package.thresholds or {}),
"seed": package.seed,
}
},
"description": package.description,
"version": package.version,
"cases": [_case_to_evalport(c) for c in package.cases],
"thresholds": dict(package.thresholds or {}),
"seed": package.seed,
}
if package.description:
suite["description"] = package.description
return suite
def from_evalport(suite: Dict[str, Any]) -> EvalPackage:
@@ -166,68 +110,22 @@ def from_evalport(suite: Dict[str, Any]) -> EvalPackage:
An ``EvalPackage`` whose cases can be run via ``EvalSuite`` /
``HarnessEvaluator`` as native ``EvalCase`` objects.
"""
items = suite.get("test_cases", suite.get("cases", []))
cases: List[EvalCase] = [_case_from_evalport(c) for c in items]
meta = (suite.get("metadata") or {}).get(_META_KEY) or {}
legacy = "evalport_version" in suite # documents emitted before spec alignment
cases: List[EvalCase] = [_case_from_evalport(c) for c in suite.get("cases", [])]
return EvalPackage(
name=suite.get("name") or suite.get("id") or "suite",
name=suite.get("name", "suite"),
description=suite.get("description", ""),
version=meta.get("version") or (suite.get("version") if legacy else None) or "1.0.0",
version=suite.get("version", "1.0.0"),
cases=cases,
thresholds=dict(meta.get("thresholds") or suite.get("thresholds") or {}),
seed=meta.get("seed", suite.get("seed")),
thresholds=dict(suite.get("thresholds") or {}),
seed=suite.get("seed"),
)
def _unit_score(value: Any) -> Optional[float]:
if isinstance(value, bool) or not isinstance(value, (int, float)) or value != value:
return None
return min(1.0, max(0.0, float(value)))
def _grader_result(grader_id: str, raw_score: Any, passed: bool) -> Dict[str, Any]:
"""Build one spec-shaped ``grader_result`` entry.
Scores are clamped to the spec's [0, 1] range; the raw value is preserved
under ``metadata`` when clamping changes it. A null score means "not
verified", so ``passed`` is forced ``False`` in that case.
"""
score = _unit_score(raw_score)
grader: Dict[str, Any] = {
"grader_id": grader_id,
"type": "custom",
"score": score,
"passed": bool(passed) if score is not None else False,
}
if score is not None and score != raw_score:
grader["metadata"] = {"openeval.raw_score": raw_score}
return grader
def _runner_info() -> Dict[str, str]:
runner = {"name": "praisonaiagents"}
try:
from importlib.metadata import version
runner["version"] = version("praisonaiagents")
except Exception:
pass
return runner
def report_to_evalport(
report: EvalReport,
*,
run_id: Optional[str] = None,
started_at: Optional[str] = None,
) -> Dict[str, Any]:
def report_to_evalport(report: EvalReport) -> Dict[str, Any]:
"""Convert an ``EvalReport`` to an EvalPort ``ResultSet`` dict.
Args:
report: The aggregated report from running an eval package.
run_id: Unique run identifier (defaults to a random UUID).
started_at: RFC 3339 run start time (defaults to now, UTC).
Returns:
A spec-shaped dict suitable for
@@ -235,52 +133,33 @@ def report_to_evalport(
"""
results: List[Dict[str, Any]] = []
for r in report.results:
verdict = _grader_result(VERDICT_GRADER_ID, r.score, r.passed)
grader_results = [verdict]
# Surface each per-criterion score as its own grader_result so EvalPort
# consumers can diff criteria across tools (they match the criterion
# grader ids referenced by the exported Suite). ``passed`` is unknown at
# the criterion level, so mirror the case verdict.
for name, score in (r.criteria_scores or {}).items():
if name == VERDICT_GRADER_ID:
continue
grader_results.append(_grader_result(name, score, r.passed))
item: Dict[str, Any] = {
"test_case_id": r.case_name,
"passed": verdict["passed"],
"grader_results": grader_results,
"case_id": r.case_name,
"passed": r.passed,
"score": r.score,
"latency_ms": r.latency_ms,
}
if r.actual_output is not None:
item["actual_output"] = str(r.actual_output)
if r.latency_ms is not None and r.latency_ms >= 0:
item["duration_ms"] = int(round(r.latency_ms))
item["actual_output"] = r.actual_output
if r.error is not None:
item["error"] = {"type": "runner_error", "message": str(r.error)}
item["error"] = r.error
if r.criteria_scores:
item["grader_scores"] = dict(r.criteria_scores)
if r.record is not None:
item["metadata"] = {_META_KEY: {"record": r.record}}
item["record"] = r.record
results.append(item)
summary: Dict[str, Any] = {
"total": report.total_cases,
"passed": report.passed_cases,
"failed": report.failed_cases,
"pass_rate": report.pass_rate,
}
meta: Dict[str, Any] = {"thresholds_met": dict(report.thresholds_met or {})}
avg = _unit_score(report.average_score)
if avg is not None:
summary["avg_score"] = avg
if avg != report.average_score:
meta["average_score"] = report.average_score # clamped
return {
"version": EVALPORT_SPEC_VERSION,
"suite_id": report.package_name,
"run_id": run_id or uuid.uuid4().hex,
"started_at": started_at or datetime.now(timezone.utc).isoformat(),
"runner": _runner_info(),
"evalport_version": EVALPORT_SPEC_VERSION,
"kind": "result_set",
"suite_name": report.package_name,
"summary": {
"total": report.total_cases,
"passed": report.passed_cases,
"failed": report.failed_cases,
"pass_rate": report.pass_rate,
"average_score": report.average_score,
"thresholds_met": dict(report.thresholds_met or {}),
},
"results": results,
"summary": summary,
"metadata": {_META_KEY: meta},
}
@@ -31,18 +31,6 @@ _LAZY_IMPORTS = {
"SessionVisibility": ("praisonaiagents.gateway.protocols", "SessionVisibility"),
"SessionSharingRole": ("praisonaiagents.gateway.protocols", "SessionSharingRole"),
"SessionObserverProtocol": ("praisonaiagents.gateway.protocols", "SessionObserverProtocol"),
# Session-projection reducer (Issue #5324)
"SessionProjection": ("praisonaiagents.gateway.session_projection", "SessionProjection"),
"SessionProjectionState": ("praisonaiagents.gateway.session_projection", "SessionProjectionState"),
"RunView": ("praisonaiagents.gateway.session_projection", "RunView"),
# Canonical session-address grammar (Issue #5383)
"SessionRef": ("praisonaiagents.gateway.addressing", "SessionRef"),
"build_session_path": ("praisonaiagents.gateway.addressing", "build_session_path"),
"parse_session_path": ("praisonaiagents.gateway.addressing", "parse_session_path"),
"derive_short_id": ("praisonaiagents.gateway.addressing", "derive_short_id"),
"slugify": ("praisonaiagents.gateway.addressing", "slugify"),
"RESERVED_NAMES": ("praisonaiagents.gateway.addressing", "RESERVED_NAMES"),
"SHORT_ID_LEN": ("praisonaiagents.gateway.addressing", "SHORT_ID_LEN"),
"GatewayCloseCode": ("praisonaiagents.gateway.protocols", "GatewayCloseCode"),
# Declarative method -> required-scope registry (Issue #3206)
"GatewayMethodDescriptor": ("praisonaiagents.gateway.protocols", "GatewayMethodDescriptor"),
@@ -57,11 +45,6 @@ _LAZY_IMPORTS = {
# Saturation / back-pressure observability (Issue #4265)
"HealthPressure": ("praisonaiagents.gateway.protocols", "HealthPressure"),
"evaluate_pressure": ("praisonaiagents.gateway.protocols", "evaluate_pressure"),
# Trusted-proxy ingress attribution (Issue #5312)
"IngressTrust": ("praisonaiagents.gateway.protocols", "IngressTrust"),
"IngressAttribution": ("praisonaiagents.gateway.protocols", "IngressAttribution"),
"resolve_ingress_attribution": ("praisonaiagents.gateway.protocols", "resolve_ingress_attribution"),
"parse_forwarded_for": ("praisonaiagents.gateway.protocols", "parse_forwarded_for"),
# Push protocols and dataclasses
"PushChannelProtocol": ("praisonaiagents.gateway.protocols", "PushChannelProtocol"),
"PresenceProtocol": ("praisonaiagents.gateway.protocols", "PresenceProtocol"),
@@ -90,18 +73,10 @@ _LAZY_IMPORTS = {
# Agent/gateway-callable thread creation (Issue #3987)
"ThreadResult": ("praisonaiagents.gateway.protocols", "ThreadResult"),
"ThreadStatus": ("praisonaiagents.gateway.protocols", "ThreadStatus"),
# Agent-callable message mutation: edit/delete (Issue #5054)
"MessageActionResult": ("praisonaiagents.gateway.protocols", "MessageActionResult"),
"MessageActionStatus": ("praisonaiagents.gateway.protocols", "MessageActionStatus"),
# Agent-callable cross-conversation request/reply (Issue #3689)
"ConversationReply": ("praisonaiagents.gateway.protocols", "ConversationReply"),
"ConversationReplyStatus": ("praisonaiagents.gateway.protocols", "ConversationReplyStatus"),
"ConversationRequestProtocol": ("praisonaiagents.gateway.protocols", "ConversationRequestProtocol"),
# Server→client interactive request/reply over the transport (Issue #5351)
"GatewayRequestKind": ("praisonaiagents.gateway.protocols", "GatewayRequestKind"),
"GatewayServerRequest": ("praisonaiagents.gateway.protocols", "GatewayServerRequest"),
"GatewayServerReply": ("praisonaiagents.gateway.protocols", "GatewayServerReply"),
"GatewayRequestChannelProtocol": ("praisonaiagents.gateway.protocols", "GatewayRequestChannelProtocol"),
# Agent-facing live status/health (Issue #3688)
"GatewayStatusProtocol": ("praisonaiagents.gateway.protocols", "GatewayStatusProtocol"),
"GatewayStatus": ("praisonaiagents.gateway.protocols", "GatewayStatus"),
@@ -252,8 +227,6 @@ _LAZY_IMPORTS = {
# liveness.py — Event-loop liveness watchdog (Issue #3385)
"LoopWatchdogPolicy": ("praisonaiagents.gateway.liveness", "LoopWatchdogPolicy"),
"LoopWatchdog": ("praisonaiagents.gateway.liveness", "LoopWatchdog"),
# liveness.py — Startup-phase watchdog (Issue #5265)
"StartupWatchdog": ("praisonaiagents.gateway.liveness", "StartupWatchdog"),
# degraded_state.py — Unified degraded-capability registry (Issue #3518)
"DegradedOwner": ("praisonaiagents.gateway.degraded_state", "DegradedOwner"),
"DegradedCapabilityProtocol": ("praisonaiagents.gateway.degraded_state", "DegradedCapabilityProtocol"),
@@ -396,18 +369,6 @@ __all__ = [
"SessionVisibility",
"SessionSharingRole",
"SessionObserverProtocol",
# Session-projection reducer (Issue #5324)
"SessionProjection",
"SessionProjectionState",
"RunView",
# Canonical session-address grammar (Issue #5383)
"SessionRef",
"build_session_path",
"parse_session_path",
"derive_short_id",
"slugify",
"RESERVED_NAMES",
"SHORT_ID_LEN",
"GatewayCloseCode",
# Declarative method -> required-scope registry (Issue #3206)
"GatewayMethodDescriptor",
@@ -421,11 +382,6 @@ __all__ = [
"compute_config_revision",
"HealthPressure",
"evaluate_pressure",
# Trusted-proxy ingress attribution (Issue #5312)
"IngressTrust",
"IngressAttribution",
"resolve_ingress_attribution",
"parse_forwarded_for",
# Push protocols (always available)
"PushChannelProtocol",
"PresenceProtocol",
@@ -449,18 +405,11 @@ __all__ = [
"ReactionResult",
"ReactionStatus",
"ThreadResult",
"MessageActionResult",
"MessageActionStatus",
"ThreadStatus",
# Agent-callable cross-conversation request/reply (Issue #3689)
"ConversationReply",
"ConversationReplyStatus",
"ConversationRequestProtocol",
# Server→client interactive request/reply over the transport (Issue #5351)
"GatewayRequestKind",
"GatewayServerRequest",
"GatewayServerReply",
"GatewayRequestChannelProtocol",
# Agent-facing live status/health (Issue #3688)
"GatewayStatusProtocol",
"GatewayStatus",
@@ -608,8 +557,6 @@ __all__ = [
# Event-loop liveness watchdog (Issue #3385)
"LoopWatchdogPolicy",
"LoopWatchdog",
# Startup-phase watchdog (Issue #5265)
"StartupWatchdog",
# Unified degraded-capability registry (Issue #3518)
"DegradedOwner",
"DegradedCapabilityProtocol",
@@ -1,227 +0,0 @@
"""Canonical session-address grammar for the Gateway (Issue #5383).
A pure, dependency-free grammar that maps between an opaque runtime
``session_key`` (a UUID, or an ``agent:<id>:<uuid>``-style key) and a stable,
human-readable, reconnect-survivable path — e.g. ``chat/<agentId>/<slug>-<shortId>``.
It gives every surface (dashboard deep-links, bot channels, the CLI, the
``praisonai-ts`` mirror, reconnecting clients) a single addressing contract to
agree on, instead of each one carrying the raw UUID or re-inventing a resolver.
This module lives in **core** because the grammar is a set of pure functions
that define a protocol/contract every surface must share. It has **no**
third-party dependencies and does no I/O — the concrete store-backed lookup
(``session.resolve``) is implemented in the wrapper behind this contract.
The grammar is deliberately lossy in one direction: ``build_session_path``
derives a short id from the trailing UUID so the whole opaque instance id never
appears in a URL, and an optional slug from the session label. ``parse_session_path``
recovers a :class:`SessionRef` (agent id + short id + slug hint) that a resolver
maps back to the concrete ``session_id``; it never assumes it can reconstruct the
key byte-for-byte, except via the exact-key escape hatch (a ``!``-prefixed final
segment) for sessions that must keep their literal key.
"""
from __future__ import annotations
import re
from dataclasses import dataclass
from typing import Optional
from urllib.parse import quote, unquote
__all__ = [
"SessionRef",
"SHORT_ID_LEN",
"RESERVED_NAMES",
"derive_short_id",
"slugify",
"build_session_path",
"parse_session_path",
]
# Length of the short id derived from the trailing hex of a session key. Eight
# hex chars (32 bits) keeps URLs terse while staying collision-resistant enough
# for a per-agent namespace, where a resolver disambiguates the rare clash.
SHORT_ID_LEN = 8
# Reserved final-segment names that denote a well-known / sentinel session
# (e.g. the agent's ``main`` conversation) rather than a ``slug-shortId`` pair.
# A resolver treats these as a friendly sentinel, not a short-id lookup.
RESERVED_NAMES = frozenset({"main", "global", "default", "root"})
# Trailing run of hex digits that terminates the key — the opaque UUID tail we
# shorten from. Anchored at end-of-string so a *hex run followed by non-hex*
# (e.g. ``report-deadbeef.js``) is NOT shortened and instead round-trips via the
# literal-key escape hatch; only a genuine hex terminus (a UUID tail) is lossy.
_HEX_TAIL_RE = re.compile(r"([0-9a-fA-F]{%d,})$" % SHORT_ID_LEN)
# A ``slug-shortId`` final segment: optional slug, then a hex short id.
_SLUG_SHORT_RE = re.compile(r"^(?:(?P<slug>.+)-)?(?P<short>[0-9a-fA-F]{%d})$" % SHORT_ID_LEN)
# Characters allowed verbatim in a slug; everything else collapses to ``-``.
_SLUG_STRIP_RE = re.compile(r"[^a-z0-9]+")
# Exact-key escape hatch: a final segment prefixed with ``!`` carries the
# percent-encoded literal key verbatim (for sessions that must not be shortened).
_LITERAL_PREFIX = "!"
@dataclass(frozen=True)
class SessionRef:
"""A parsed friendly reference to a session, for a resolver to map to an id.
Exactly one of the resolution hints is authoritative depending on how the
path was written:
- ``literal_key`` set -> resolve by exact key (escape hatch).
- ``short_id`` set -> resolve by short id (optionally aided by ``slug_hint``).
- ``slug_hint`` only, matching a reserved name -> resolve the sentinel session.
``agent_id`` scopes the lookup so short ids only need to be unique per agent.
"""
agent_id: str
short_id: Optional[str] = None
slug_hint: Optional[str] = None
literal_key: Optional[str] = None
def derive_short_id(session_key: str) -> Optional[str]:
"""Derive a stable short id from the terminating hex run of ``session_key``.
Returns the last :data:`SHORT_ID_LEN` hex chars of the key's *terminating*
hex run (e.g. a UUID tail), lower-cased. The run must end the key: a key
whose hex is followed by non-hex characters (e.g. ``report-deadbeef.js``)
is deliberately not shortened, so it round-trips via the literal-key escape
hatch instead of silently losing its suffix. Returns ``None`` when the key
has no terminating hex run long enough to shorten — the caller then falls
back to the literal key or reserved-name form.
"""
if not session_key:
return None
match = _HEX_TAIL_RE.search(session_key)
if not match:
return None
hex_tail = match.group(1)
return hex_tail[-SHORT_ID_LEN:].lower()
def slugify(display_name: Optional[str], max_len: int = 48) -> Optional[str]:
"""Turn a human label into a lowercase, dash-separated URL slug.
Non-alphanumeric runs collapse to a single ``-``; leading/trailing dashes are
trimmed. Returns ``None`` for an empty/whitespace-only or unusable name so
the path omits the slug rather than emitting an empty segment.
"""
if not display_name:
return None
slug = _SLUG_STRIP_RE.sub("-", display_name.strip().lower()).strip("-")
if not slug:
return None
if len(slug) > max_len:
slug = slug[:max_len].rstrip("-")
return slug or None
def build_session_path(
agent_id: str,
session_key: str,
display_name: Optional[str] = None,
namespace: str = "chat",
) -> Optional[str]:
"""Build a stable, shareable path for a session.
Produces ``<namespace>/<agentId>/<slug>-<shortId>`` when a short id can be
derived (with the slug omitted when there is no usable ``display_name``), so
the whole opaque UUID never appears in the URL. When the key has no hex tail
to shorten, falls back to an exact-key segment (``!<percent-encoded-key>``)
so the address round-trips losslessly. Returns ``None`` if ``agent_id`` or
``session_key`` is empty.
Path segments are percent-escaped, so a key that looks like a file
(e.g. ``report.js``) is never mistaken for a static asset.
"""
if not agent_id or not session_key:
return None
ns = _seg(namespace) or "chat"
agent_seg = _seg(agent_id)
if not agent_seg:
return None
short_id = derive_short_id(session_key)
if short_id:
slug = slugify(display_name)
tail = f"{slug}-{short_id}" if slug else short_id
else:
# No hex tail to shorten: keep the literal key behind the escape hatch.
tail = _LITERAL_PREFIX + quote(session_key, safe="")
return f"{ns}/{agent_seg}/{tail}"
def parse_session_path(path: str) -> Optional[SessionRef]:
"""Parse a session path back into a :class:`SessionRef`.
Accepts the forms produced by :func:`build_session_path`, with or without a
leading namespace segment:
- ``chat/<agent>/<slug>-<shortId>`` / ``<agent>/<slug>-<shortId>``
- ``chat/<agent>/<shortId>`` / ``<agent>/<shortId>``
- ``chat/<agent>/main`` (reserved sentinel) / ``<agent>/main``
- ``chat/<agent>/!<percent-encoded-key>`` (exact-key escape hatch)
Returns ``None`` when the path has no agent segment or is otherwise
unparseable. A resolver combines the returned hints with its store to reach
the concrete ``session_id``.
"""
if not path:
return None
parts = [p for p in path.strip().strip("/").split("/") if p]
# ``build_session_path`` percent-escapes the namespace and agent into single
# segments, so a canonical address is exactly ``<ns>/<agent>/<ref>`` (3) or
# ``<agent>/<ref>`` (2). Reject anything longer so a malformed link like
# ``chat/x/y/z`` never silently acquires a different meaning by using only
# its last two segments.
if len(parts) == 3:
# Drop the leading namespace segment; agent + final ref remain.
agent_id = unquote(parts[-2])
final = parts[-1]
elif len(parts) == 2:
agent_id = unquote(parts[0])
final = parts[1]
else:
return None
agent_id = agent_id.strip()
if not agent_id or not final:
return None
# Exact-key escape hatch.
if final.startswith(_LITERAL_PREFIX):
literal = unquote(final[len(_LITERAL_PREFIX):])
if not literal:
return None
return SessionRef(agent_id=agent_id, literal_key=literal)
final = unquote(final)
# Reserved sentinel (main/global/...): a friendly named session.
if final.lower() in RESERVED_NAMES:
return SessionRef(agent_id=agent_id, slug_hint=final.lower())
match = _SLUG_SHORT_RE.match(final)
if match:
return SessionRef(
agent_id=agent_id,
short_id=match.group("short").lower(),
slug_hint=match.group("slug") or None,
)
# No short id present: treat the whole segment as a slug hint the resolver
# can match against session labels.
return SessionRef(agent_id=agent_id, slug_hint=final)
def _seg(value: Optional[str]) -> str:
"""Percent-escape a single path segment (``/`` and dots are encoded)."""
if not value:
return ""
return quote(value.strip(), safe="")
@@ -753,17 +753,10 @@ class ApiConfig:
backed by the gateway's live agents and sessions.
mcp: Serve an MCP JSON-RPC endpoint (``/mcp``) exposing the gateway's
registered agents as callable tools.
stream: Emit true token-level deltas on the OpenAI-compatible
``/v1/chat/completions`` SSE surface (``stream:true``) by wiring the
agent's existing ``stream_emitter`` through the gateway hot path.
Off by default: streaming responses stay byte-for-byte the buffered
single-chunk path, so enabling it never regresses non-streaming
latency or correctness.
"""
openai: bool = False
mcp: bool = False
stream: bool = False
@property
def enabled(self) -> bool:
@@ -772,7 +765,7 @@ class ApiConfig:
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {"openai": self.openai, "mcp": self.mcp, "stream": self.stream}
return {"openai": self.openai, "mcp": self.mcp}
@classmethod
def from_dict(cls, data: Optional[Dict[str, Any]]) -> "ApiConfig":
@@ -782,7 +775,6 @@ class ApiConfig:
return cls(
openai=bool(data.get("openai", False)),
mcp=bool(data.get("mcp", False)),
stream=bool(data.get("stream", False)),
)
@@ -941,73 +933,6 @@ class AttachmentConfig:
)
def _parse_drain_timeout(value: Any) -> Optional[float]:
"""Coerce a ``gateway.drain_timeout`` YAML value to a float (or None).
Tolerant of env-substituted strings (e.g. ``"20"``); an unparseable /
negative value falls back to None (no drain) rather than raising, matching
the wrapper's fail-open handling of a malformed drain window (#5362).
"""
if value is None:
return None
try:
parsed = float(value)
except (TypeError, ValueError):
return None
return parsed if parsed >= 0 else None
def _parse_watchdog_enabled(value: Any) -> bool:
"""Whether a ``gateway.watchdog`` YAML value enables the loop watchdog.
Accepts both a bare truthy scalar and the documented block form
(``watchdog: { enabled: true, ... }``), mirroring the wrapper's watchdog
wiring so the public config reflects the same on/off decision (#5362).
"""
if isinstance(value, dict):
flag = value.get("enabled", False)
else:
flag = value
if isinstance(flag, str):
return flag.strip().lower() in ("1", "true", "yes", "on")
return bool(flag)
def _parse_watchdog_timeout(value: Any) -> Optional[float]:
"""Derive the loop-stall budget (seconds) from a ``gateway.watchdog`` block.
Reads ``liveness_interval * liveness_strikes`` from the block form when
present so the public config exposes the effective wedge budget the CLI
``--watchdog-timeout`` flag sets; returns None when unset/unparseable.
"""
if not isinstance(value, dict):
return None
try:
interval = float(value.get("liveness_interval", 5.0))
strikes = int(value.get("liveness_strikes", 3)) or 3
except (TypeError, ValueError):
return None
budget = interval * strikes
return budget if budget > 0 else None
def _parse_watchdog_timeout_scalar(value: Any) -> Optional[float]:
"""Coerce an explicit ``gateway.watchdog_timeout`` scalar to a float.
Mirrors the ``watchdog_timeout > 0`` invariant (see ``GatewayConfig``):
tolerant of env-substituted strings (e.g. ``"15"``); an unparseable or
non-positive value falls back to None so a malformed timeout disables the
budget rather than raising (#5362).
"""
if value is None:
return None
try:
parsed = float(value)
except (TypeError, ValueError):
return None
return parsed if parsed > 0 else None
@dataclass
class GatewayConfig:
"""Configuration for the gateway server.
@@ -1040,13 +965,6 @@ class GatewayConfig:
host: str = "127.0.0.1"
port: int = 8765
bind_host: Optional[str] = None
# Issue #5312: CIDRs/IPs of upstream proxies/tunnels the operator declares
# trusted. Empty (default) means the gateway is directly exposed: forwarded
# headers are never believed and per-IP policy keys on the socket peer.
# When set, the real client IP is resolved by walking the forwarded chain
# across trusted hops (see ``resolve_ingress_attribution``); anything
# proxy-shaped but unattributable fails closed to the socket peer.
trusted_proxies: List[str] = field(default_factory=list)
cors_origins: List[str] = field(default_factory=lambda: [])
allowed_origins: List[str] = field(default_factory=lambda: [])
auth_token: Optional[str] = None
@@ -1070,14 +988,6 @@ class GatewayConfig:
max_concurrent_runs: int = 0 # Aggregate concurrency ceiling (0 = unlimited)
queue_depth: int = 0 # Bounded wait queue when at the ceiling
overflow_policy: str = "reject" # reject | queue | shed_oldest
# Issue #5362: complete the back-pressure/lifecycle posture on the public
# config so a Python embedder / gateway.yaml can pin it declaratively —
# parity with the CLI ``--drain-timeout`` / ``--watchdog`` /
# ``--watchdog-timeout`` flags. ``None``/False keep today's behaviour
# (no drain, watchdog off); a CLI flag still overrides these at start.
drain_timeout: Optional[float] = None # graceful-drain window on shutdown (s)
watchdog: bool = False # event-loop liveness watchdog (opt-in)
watchdog_timeout: Optional[float] = None # loop-stall budget before restart (s)
# Issue #5168: per-tenant/per-scope concurrency sub-limit within the global
# ceiling so one tenant cannot occupy every run slot. 0 disables the
# sub-limit (today's global-only behaviour).
@@ -1157,14 +1067,6 @@ class GatewayConfig:
raise ValueError(
"overflow_policy must be one of 'reject', 'queue', 'shed_oldest'"
)
if self.drain_timeout is not None and self.drain_timeout < 0:
raise ValueError(
"drain_timeout must be >= 0 (use 0/None to disable the drain window)"
)
if self.watchdog_timeout is not None and self.watchdog_timeout <= 0:
raise ValueError(
"watchdog_timeout must be > 0 (leave unset to use the default budget)"
)
if self.max_concurrent_runs_per_scope < 0:
raise ValueError(
"max_concurrent_runs_per_scope must be >= 0 "
@@ -1237,7 +1139,6 @@ class GatewayConfig:
return {
"host": self.host,
"port": self.port,
"trusted_proxies": list(self.trusted_proxies),
"cors_origins": self.cors_origins,
"allowed_origins": self.allowed_origins,
"auth_token": "***" if self.auth_token else None,
@@ -1253,9 +1154,6 @@ class GatewayConfig:
"max_concurrent_runs": self.max_concurrent_runs,
"queue_depth": self.queue_depth,
"overflow_policy": self.overflow_policy,
"drain_timeout": self.drain_timeout,
"watchdog": self.watchdog,
"watchdog_timeout": self.watchdog_timeout,
"max_concurrent_runs_per_scope": self.max_concurrent_runs_per_scope,
"preauth_max_connections_per_ip": self.preauth_max_connections_per_ip,
"max_unauthorized_frames": self.max_unauthorized_frames,
@@ -1435,17 +1333,9 @@ class MultiChannelGatewayConfig:
for tok, scopes in gw_data["auth_scopes"].items()
}
raw_trusted = gw_data.get("trusted_proxies", [])
trusted_proxies = (
[str(p).strip() for p in raw_trusted if str(p).strip()]
if isinstance(raw_trusted, (list, tuple))
else []
)
gateway_config = GatewayConfig(
host=gw_data.get("host", "127.0.0.1"),
port=gw_data.get("port", 8765),
trusted_proxies=trusted_proxies,
cors_origins=gw_data.get("cors_origins", []),
allowed_origins=gw_data.get("allowed_origins", []),
auth_token=gw_data.get("auth_token"),
@@ -1462,17 +1352,6 @@ class MultiChannelGatewayConfig:
max_concurrent_runs=int(gw_data.get("max_concurrent_runs", 0) or 0),
queue_depth=int(gw_data.get("queue_depth", 0) or 0),
overflow_policy=str(gw_data.get("overflow_policy", "reject") or "reject"),
drain_timeout=_parse_drain_timeout(gw_data.get("drain_timeout")),
watchdog=_parse_watchdog_enabled(gw_data.get("watchdog")),
# Prefer an explicit ``gateway.watchdog_timeout`` scalar (as emitted
# by ``to_dict`` and settable directly, matching the CLI
# ``--watchdog-timeout`` flag); otherwise derive the budget from the
# ``watchdog`` block so both round-trip and block forms work (#5362).
watchdog_timeout=(
_parse_watchdog_timeout_scalar(gw_data.get("watchdog_timeout"))
if gw_data.get("watchdog_timeout") is not None
else _parse_watchdog_timeout(gw_data.get("watchdog"))
),
max_concurrent_runs_per_scope=int(
gw_data.get("max_concurrent_runs_per_scope", 0) or 0
),
@@ -52,7 +52,6 @@ from .protocols import GATEWAY_RESTART_EXIT_CODE
__all__ = [
"LoopWatchdogPolicy",
"LoopWatchdog",
"StartupWatchdog",
]
@@ -438,271 +437,3 @@ class LoopWatchdog:
faulthandler.dump_traceback(file=fh, all_threads=True)
except Exception: # pragma: no cover - fail open
pass
class StartupWatchdog:
"""Front-rung liveness guard for the pre-loop startup window (Issue #5265).
:class:`LoopWatchdog` only becomes meaningful once the asyncio event loop
is running — it probes the loop via ``loop.call_soon_threadsafe``, so it can
do nothing before a loop exists. That leaves the window between process
entry and "loop is live" (heavy imports, config/secret loading, the initial
adapter ``connect()``\\ s) uncovered. A wedge there — a hanging SDK import, a
blocking synchronous DNS/network call while minting credentials, a deadlock
inside ``connect()`` — leaves the process *alive but making no progress*: it
holds no serving loop, answers no probes, and never exits, so an external
supervisor never restarts it. It becomes a silent zombie with no diagnostics.
This primitive closes that gap. It arms a wall-clock **deadline** on a
dedicated daemon OS thread *before* the heavy work begins, so it keeps
working precisely when the main thread is wedged. If the deadline elapses
without :meth:`confirm_loop_live`, it dumps all-thread stacks via
:mod:`faulthandler` (so the hang is diagnosable) and hands the process back
to the supervisor via ``os._exit`` with a restart-intent code. A legitimately
slow-but-alive boot extends the deadline via :meth:`report_startup_progress`
(bounded, so a truly wedged phase still trips). :meth:`confirm_loop_live`
disarms it the instant the serving loop begins — at which point
:class:`LoopWatchdog` takes over.
Design mirrors :class:`LoopWatchdog`: **opt-in** (nothing arms it unless the
embedder calls :meth:`arm`), **fail-open** (any error inside the watchdog
must never kill a healthy startup), and **stdlib-only** (it must run when the
rest of the process cannot, so its correctness cannot depend on the code it
guards). Typical use at the gateway entry point::
sw = StartupWatchdog()
sw.arm(deadline_s=config.startup_watchdog_timeout)
load_config(config); sw.report_startup_progress("config loaded")
connect_adapters(); sw.report_startup_progress("adapters connected")
loop = asyncio.new_event_loop()
LoopWatchdog().arm(loop)
sw.confirm_loop_live() # disarm; loop watchdog now owns liveness
loop.run_forever()
"""
def __init__(
self,
*,
on_expire: Literal["dump_and_exit", "dump_only"] = "dump_and_exit",
exit_code: int = GATEWAY_RESTART_EXIT_CODE,
dump_file: Optional[str] = None,
max_extensions: int = 10,
) -> None:
if on_expire not in ("dump_and_exit", "dump_only"):
raise ValueError(
"on_expire must be 'dump_and_exit' or 'dump_only'"
)
if max_extensions < 0:
raise ValueError("max_extensions must be >= 0")
self.on_expire = on_expire
self.exit_code = exit_code
self.dump_file = dump_file
self.max_extensions = max_extensions
self._thread: Optional[threading.Thread] = None
# Each armed generation gets its own stop event handed straight to the
# thread, so a confirm/re-arm race can never let a stale generation
# observe a fresh generation's cleared flag and self-exit (mirrors the
# per-generation discipline of LoopWatchdog's deadline mode). The stop
# event, once set, stays set: a healthy confirm/disarm is monotonic.
self._stop = threading.Event()
# A separate wake event carries progress reports: it is briefly set to
# nudge the waiting thread to re-read the (extended) deadline, then
# cleared. Keeping it distinct from _stop means a progress nudge can
# never be mistaken for a disarm (and vice versa).
self._wake = threading.Event()
self._deadline_s = 0.0
self._extensions_used = 0
# Progress reports enqueue here; the watchdog thread drains this under
# the lock so budget is spent exactly once per report even when several
# reports batch between two wakes.
self._pending_progress = 0
self._expired = False
self._confirmed = False
self._lock = threading.Lock()
@property
def armed(self) -> bool:
return self._thread is not None and self._thread.is_alive()
@property
def expired(self) -> bool:
"""True once the startup deadline expired (mainly for ``dump_only``)."""
return self._expired
@property
def confirmed(self) -> bool:
"""True once :meth:`confirm_loop_live` has disarmed the watchdog."""
return self._confirmed
def arm(self, *, deadline_s: float, dump_stacks: bool = True) -> None:
"""Arm a startup deadline on a dedicated daemon thread. Fail-open.
Idempotent: arming while already armed is a no-op. A non-finite or
non-positive ``deadline_s`` is ignored (nothing to guard). ``dump_stacks``
toggles the all-thread stack dump on expiry.
"""
if not math.isfinite(deadline_s) or deadline_s <= 0:
return
with self._lock:
if self.armed:
return
stop = threading.Event()
wake = threading.Event()
self._stop = stop
self._wake = wake
self._deadline_s = deadline_s
self._extensions_used = 0
self._pending_progress = 0
self._expired = False
self._confirmed = False
self._dump_stacks_on_expire = dump_stacks
try:
thread = threading.Thread(
target=self._run,
args=(stop, wake),
name="praisonai-startup-watchdog",
daemon=True,
)
thread.start()
except Exception: # pragma: no cover - fail open
self._thread = None
return
self._thread = thread
def report_startup_progress(self, phase: str) -> None:
"""Signal forward progress, extending the deadline by one interval.
A legitimately slow boot (e.g. many adapters, cold DNS) calls this at
each phase so it is not killed. Extensions are bounded by
``max_extensions`` so a phase that itself wedges still trips the
deadline. Fail-open and a no-op when not armed.
Budget accounting is done by the watchdog thread, not here: each report
merely enqueues a pending nudge. If several reports land before the
thread wakes, each still costs exactly one extension *and* each still
earns one fresh deadline window when the thread drains them — so a later
slow-but-progressing phase can never find the budget spent on resets it
never received (Issue #5265 review: "progress reports lose extensions").
"""
if not self.armed:
return
with self._lock:
self._pending_progress += 1
wake = self._wake
# Wake the waiting thread so it drains the pending nudges, spends budget,
# and resets its window once per still-affordable report. The thread
# keeps waiting against its own (still-unset) stop event.
wake.set()
def confirm_loop_live(self) -> None:
"""Disarm the startup watchdog; the event loop is now live.
Called immediately before :class:`LoopWatchdog` is armed. Safe to call
multiple times / when not armed. Sets the *current* generation's stop
event so that exact thread suppresses its self-exit even if it is
mid-dump, then drops the reference.
"""
with self._lock:
self._confirmed = True
self._stop.set()
# Nudge the thread out of its wait immediately so the join is prompt.
self._wake.set()
thread = self._thread
self._thread = None
if thread is not None and thread is not threading.current_thread():
try:
thread.join(timeout=1.0)
except Exception: # pragma: no cover - fail open
pass
def _run(self, stop: "threading.Event", wake: "threading.Event") -> None:
# Wait out the deadline against *this generation's* events. ``stop`` is
# set (and stays set) by a confirm/disarm — that must break us out and
# suppress the self-exit. ``wake`` is a transient progress nudge — on
# each wake we drain the queued progress reports, spending one extension
# of budget per report and resetting the window each time we can still
# afford it. A plain timeout with neither event set (and no affordable
# pending progress) means the startup genuinely wedged.
start = time.monotonic()
while not stop.is_set():
remaining = (start + self._deadline_s) - time.monotonic()
if remaining <= 0:
break
if wake.wait(remaining):
wake.clear()
if stop.is_set():
return
# Drain queued progress reports. Each report costs one extension
# and, while budget remains, grants one fresh deadline window —
# so batched reports never lose the resets they paid for.
extended = False
with self._lock:
pending = self._pending_progress
self._pending_progress = 0
for _ in range(pending):
if self._extensions_used >= self.max_extensions:
break
self._extensions_used += 1
extended = True
if extended:
start = time.monotonic()
# Whether or not budget remained, loop back and re-check the
# (possibly unchanged) deadline against this generation's stop.
continue
self._on_expire(stop)
def _on_expire(self, stop: "threading.Event") -> None:
# Only fire if this generation was not disarmed. A confirm_loop_live()
# that raced in has set this exact event, so a stale generation stays
# inert and never self-exits a healthy startup.
if stop.is_set():
return
self._expired = True
if getattr(self, "_dump_stacks_on_expire", True):
try:
self._dump_stacks()
except Exception: # pragma: no cover - fail open
pass
if self.on_expire == "dump_and_exit":
# If confirm_loop_live() raced in while we were dumping, honour it.
if stop.is_set():
return
try:
sys.stderr.flush()
except Exception: # pragma: no cover - fail open
pass
# Final re-check the instant before the irreversible os._exit: the
# flush above can block (e.g. a full pipe) long enough for a
# concurrent confirm_loop_live() to land. A startup that has just
# confirmed a live loop must never be terminated (Issue #5265
# review: "confirmation cannot prevent late exit").
if stop.is_set():
return
os._exit(self.exit_code)
def _dump_stacks(self) -> None:
reason = (
f"gateway startup did not confirm a live event loop within "
f"~{self._deadline_s:.0f}s"
)
header = (
f"praisonai-startup-watchdog: {reason}; "
f"dumping all-thread stacks\n"
)
try:
sys.stderr.write(header)
except Exception: # pragma: no cover - fail open
pass
try:
faulthandler.dump_traceback(file=sys.stderr, all_threads=True)
except Exception: # pragma: no cover - fail open
pass
dump_file = self.dump_file
if dump_file:
try:
with open(dump_file, "a", encoding="utf-8") as fh:
fh.write(header)
fh.flush()
faulthandler.dump_traceback(file=fh, all_threads=True)
except Exception: # pragma: no cover - fail open
pass
@@ -2555,225 +2555,6 @@ def resolve_auth_mode(bind_host: str, configured: Optional[AuthMode] = None) ->
return "local" if is_loopback(bind_host) else "token"
# ---------------------------------------------------------------------------
# Ingress attribution / trusted-proxy real-client-IP resolution (Issue #5312)
# ---------------------------------------------------------------------------
#
# The gateway keys rate-limiting, the pre-auth connection budget, and the
# operator identity on the *raw socket peer*. Behind a reverse proxy or tunnel
# (nginx, Caddy, Traefik, a cloud LB, Cloudflare, Tailscale Serve/Funnel) that
# peer is the *proxy* for every request, so every per-IP bucket collapses to
# one and the operator id becomes identical for all callers. Trusting
# ``X-Forwarded-For`` unconditionally is worse — a directly-reachable gateway
# would accept a spoofed client IP any attacker can set.
#
# This pure, import-free contract resolves the real client address *only* by
# walking the forwarded chain across hops the operator has explicitly declared
# trusted, and **fails closed** for anything proxy-shaped but unattributable
# (rate-limited as one bucket keyed on the socket peer; never trusts a header).
# It lives in core so every seam consults one unit-testable decision, exactly
# like the sibling ``resolve_route`` / ``evaluate_pressure`` contracts.
IngressTrust = Literal[
"direct-local",
"direct-remote",
"trusted-proxy",
"tunnel",
"unattributable-proxy",
]
"""Classification of where an inbound request actually came from.
- ``direct-local``: loopback peer, no proxy headers (local CLI / same host).
- ``direct-remote``: a remote peer talking to us directly, no proxy headers.
- ``trusted-proxy``: arrived via one or more hops all in ``trusted_proxies``;
the real client IP was resolved from the forwarded chain.
- ``tunnel``: a trusted loopback/private tunnel hop (e.g. Tailscale/Cloudflare
sidecar on loopback) presented a forwarded client — resolved like a trusted
proxy.
- ``unattributable-proxy``: proxy-shaped (forwarded headers present) but the
immediate hop is not declared trusted — fails closed, header not trusted.
"""
@dataclass(frozen=True)
class IngressAttribution:
"""Resolved answer to "who/where is this request from", consulted before
any per-IP policy.
Attributes:
trust: The :data:`IngressTrust` classification.
client_ip: The subject to key per-IP policy on. The *real* client only
when it was resolved across trusted hops; otherwise the raw socket
peer (never a spoofable header value).
via_proxy: Whether forwarded headers were present at all.
fail_closed: True when the request was proxy-shaped but unattributable —
the caller must rate-limit it as a single bucket keyed on the socket
peer and must never trust a client IP from a header.
"""
trust: IngressTrust
client_ip: str
via_proxy: bool
fail_closed: bool
def parse_forwarded_for(header: Optional[str]) -> "List[str]":
"""Parse an ``X-Forwarded-For`` header into an ordered list of hops.
The chain is left→right, client-first: ``client, proxy1, proxy2``. Empty
tokens and surrounding whitespace are dropped; ``None``/empty yields an
empty list. Pure string handling — no name resolution, no validation of
whether each token is a real IP (that is the caller's trust decision).
"""
if not header:
return []
return [tok.strip() for tok in header.split(",") if tok.strip()]
def _is_valid_ip(host: Optional[str]) -> bool:
"""Return whether ``host`` parses as a literal IPv4/IPv6 address.
Forwarded-header tokens are attacker-influenced strings; a value such as
``unknown`` or an arbitrary label must never become a per-IP policy subject
(it would group unrelated callers or split one caller across buckets). The
caller uses this to fail closed to the socket peer for a non-IP hop.
"""
if not host:
return False
import ipaddress
try:
ipaddress.ip_address(str(host).strip())
return True
except ValueError:
return False
def _is_trusted_hop(host: str, trusted_proxies: "Sequence[str]") -> bool:
"""Return whether ``host`` is a member of the declared trusted set.
A trusted entry may be a bare IP or a CIDR (``10.0.0.0/8``). Loopback is
*not* implicitly trusted here — the caller decides how to treat a loopback
tunnel — so this stays a pure membership test.
"""
import ipaddress
try:
ip = ipaddress.ip_address(host)
except ValueError:
return False
for entry in trusted_proxies:
entry = str(entry).strip()
if not entry:
continue
try:
if "/" in entry:
if ip in ipaddress.ip_network(entry, strict=False):
return True
elif ip == ipaddress.ip_address(entry):
return True
except ValueError:
continue
return False
def resolve_ingress_attribution(
*,
peer_ip: str,
forwarded_for: "Sequence[str]" = (),
real_ip: Optional[str] = None,
trusted_proxies: "Sequence[str]" = (),
) -> IngressAttribution:
"""Resolve the real client attribution for an inbound request.
Pure and side-effect-free so it is unit-testable in isolation and every
seam (rate limiters, connection budget, operator id) consults the same
decision.
Args:
peer_ip: Raw socket peer (``request.client.host``).
forwarded_for: Parsed ``X-Forwarded-For`` chain, left→right
(client-first). See :func:`parse_forwarded_for`.
real_ip: Optional ``X-Real-IP`` value (used only when there is no
forwarded chain and the immediate hop is trusted).
trusted_proxies: CIDRs/IPs the operator declares trusted.
Returns:
An :class:`IngressAttribution`. The real client IP is resolved *only*
across hops in ``trusted_proxies``; anything proxy-shaped but
unattributable fails closed to the socket peer.
Resolution rules:
* No proxy headers → ``direct-local`` (loopback peer) or
``direct-remote``; ``client_ip`` is the peer, never fail-closed.
* Proxy headers present but the socket peer is not a trusted hop →
``unattributable-proxy``, ``fail_closed=True``, ``client_ip`` = peer
(a spoofable header is never trusted on a directly-reachable
gateway).
* Socket peer is a trusted hop → walk the forwarded chain right→left,
peeling trusted hops, and stop at the first untrusted address: that
address is the real client. If the chain is empty, fall back to
``real_ip`` then the peer. The resolved client must be a literal IP —
a malformed/non-IP hop fails closed to the peer. A loopback trusted
hop is classified ``tunnel``; otherwise ``trusted-proxy``.
"""
has_headers = bool(forwarded_for) or bool(real_ip)
# No forwarded headers at all — attribute straight to the socket peer.
if not has_headers:
trust: IngressTrust = "direct-local" if is_loopback(peer_ip) else "direct-remote"
return IngressAttribution(
trust=trust,
client_ip=peer_ip,
via_proxy=False,
fail_closed=False,
)
# Proxy-shaped. The immediate hop (socket peer) must itself be trusted for
# any header to be believed; otherwise fail closed to the socket peer.
if not _is_trusted_hop(peer_ip, trusted_proxies):
return IngressAttribution(
trust="unattributable-proxy",
client_ip=peer_ip,
via_proxy=True,
fail_closed=True,
)
# Trusted immediate hop. Walk the forwarded chain right→left, peeling hops
# that are themselves trusted; the first untrusted address is the real
# client. Stop as soon as we hit it so a client-supplied spoofed prefix
# cannot be peeled past.
client_ip = None
for hop in reversed(list(forwarded_for)):
if _is_trusted_hop(hop, trusted_proxies):
continue
client_ip = hop
break
if client_ip is None:
# Entire chain was trusted (or empty). Prefer X-Real-IP, then peer.
client_ip = (real_ip.strip() if real_ip else None) or peer_ip
# The resolved client must be a literal IP before it keys per-IP policy. A
# malformed hop (e.g. ``unknown`` or an arbitrary label injected upstream)
# would otherwise group unrelated callers or split one caller across
# buckets — fail closed to the socket peer instead (Greptile P2).
if not _is_valid_ip(client_ip):
return IngressAttribution(
trust="unattributable-proxy",
client_ip=peer_ip,
via_proxy=True,
fail_closed=True,
)
trust = "tunnel" if is_loopback(peer_ip) else "trusted-proxy"
return IngressAttribution(
trust=trust,
client_ip=client_ip,
via_proxy=True,
fail_closed=False,
)
# ---------------------------------------------------------------------------
# Weak / placeholder secret guard (Issue #3259)
# ---------------------------------------------------------------------------
@@ -3293,57 +3074,6 @@ ReactionStatus = Literal["ok", "unsupported", "failed", "no_route"]
"""
MessageActionStatus = Literal["ok", "unsupported", "failed", "no_route"]
"""Closed set of outcomes for a message-mutation verb (``edit`` / ``delete``).
Mirrors :data:`ReactionStatus` so every agent-callable message action returns
the same typed shape:
* ``ok`` — the message was edited/deleted.
* ``unsupported`` — the channel has no matching capability
(``supports_edit`` / ``supports_delete``).
* ``failed`` — the transport rejected the action (e.g. no such message, or the
message is not editable/deletable by this bot).
* ``no_route`` — the target could not be resolved to a reachable channel.
"""
@dataclass
class MessageActionResult:
"""Outcome of an agent-initiated message mutation (Issue #5054).
Shared typed result for the ``edit`` and ``delete`` verbs of the unified
message-action surface. Every call resolves to exactly one
:data:`MessageActionStatus`, so a channel that cannot edit/delete returns a
typed ``unsupported`` outcome instead of raising — identical to how
:class:`ReactionResult` and :class:`ThreadResult` degrade.
Attributes:
status: The outcome (``ok`` / ``unsupported`` / ``failed`` /
``no_route``).
target: The resolved target the action was routed to.
detail: Optional model-readable explanation.
"""
status: MessageActionStatus
target: str = ""
detail: Optional[str] = None
@property
def ok(self) -> bool:
"""Whether the action succeeded (``status == "ok"``)."""
return self.status == "ok"
def as_dict(self) -> Dict[str, Any]:
"""Convert to a serializable dictionary for the tool return value."""
data: Dict[str, Any] = {"status": self.status}
if self.target:
data["target"] = self.target
if self.detail:
data["detail"] = self.detail
return data
@dataclass
class ReactionResult:
"""Outcome of an agent-initiated message reaction (Issue #3917).
@@ -3529,51 +3259,6 @@ class OutboundMessengerProtocol(Protocol):
"""
...
async def edit(
self,
target: str,
message_id: str,
text: str,
) -> "MessageActionResult":
"""Edit a prior message this agent/session sent (Issue #5054).
Lets an agent update its own live status message in place rather than
posting a follow-up — e.g. editing "Working…" to "Done ✅". Gated on the
channel's ``PlatformCapabilities.supports_edit``; a channel that cannot
edit returns a typed ``unsupported`` outcome instead of raising.
Args:
target: Symbolic target token ("origin", "<platform>",
"<platform>:<chat_id>[:<thread_id>]", or a friendly alias).
message_id: The id of the message to edit.
text: The new message text.
Returns:
A :class:`MessageActionResult` describing the outcome.
"""
...
async def delete(
self,
target: str,
message_id: str,
) -> "MessageActionResult":
"""Delete/unsend a prior message this agent/session sent (Issue #5054).
Lets an agent retract a message it should not have sent. Gated on the
channel's ``PlatformCapabilities.supports_delete``; a channel that
cannot delete returns a typed ``unsupported`` outcome instead of raising.
Args:
target: Symbolic target token ("origin", "<platform>",
"<platform>:<chat_id>[:<thread_id>]", or a friendly alias).
message_id: The id of the message to delete.
Returns:
A :class:`MessageActionResult` describing the outcome.
"""
...
# ---------------------------------------------------------------------------
# Agent-callable cross-conversation request/reply (Issue #3689)
@@ -3680,122 +3365,6 @@ class ConversationRequestProtocol(Protocol):
...
# ---------------------------------------------------------------------------
# Server→client interactive request/reply over the gateway transport
# (Issue #5351)
#
# A blocked HITL turn (approval / choice / free-text input) needs to reach a
# generic gateway client — a web dashboard, TUI, or custom app — and correlate
# the client's answer back into the waiting turn. Structured elicitation today
# only renders through per-channel button backends; a client on the gateway
# WebSocket cannot see or answer one, and an in-flight prompt is lost when the
# client drops (the ``since`` cursor replays *past events*, not *still-open
# requests*).
#
# Core owns only the *shape*: the typed request (:class:`GatewayServerRequest`),
# the correlated reply (:class:`GatewayServerReply`), and the protocol seam
# (:class:`GatewayRequestChannelProtocol`) — including the ``open_requests``
# contract that the resume path honours so a reconnecting client is re-sent
# every still-open (and *only* still-open) request for its sessions. The socket
# delivery, correlation, and durable open-request set are bound by the running
# gateway (praisonai-bot) exactly as the outbound messenger is — no heavy
# import lives in core.
# ---------------------------------------------------------------------------
GatewayRequestKind = Literal["approval", "choice", "input"]
"""Closed set of server→client interactive request kinds.
* ``approval`` — allow/deny; the reply ``value`` is ``"allow"`` or ``"deny"``.
* ``choice`` — select one of ``options``; the reply ``value`` is the option.
* ``input`` — free text; the reply ``value`` is the entered text.
"""
@dataclass
class GatewayServerRequest:
"""A typed server→client interactive request (Issue #5351).
Delivered to a connected gateway client as a ``server_request`` frame while
a turn is blocked awaiting a human answer, and re-issued verbatim on resume
for as long as it stays open.
Attributes:
request_id: Correlation id echoed back in the :class:`GatewayServerReply`.
kind: One of :data:`GatewayRequestKind`.
prompt: Human-readable prompt to render.
options: Selectable options, populated only for ``kind == "choice"``.
session_id: Session the request belongs to (used for resume replay).
"""
request_id: str
kind: GatewayRequestKind
prompt: str
options: Optional[List[str]] = None
session_id: str = ""
def as_dict(self) -> Dict[str, Any]:
"""Convert to a serializable dictionary for the ``server_request`` frame."""
data: Dict[str, Any] = {
"request_id": self.request_id,
"kind": self.kind,
"prompt": self.prompt,
}
if self.options is not None:
data["options"] = list(self.options)
if self.session_id:
data["session_id"] = self.session_id
return data
@dataclass
class GatewayServerReply:
"""A client's correlated answer to a :class:`GatewayServerRequest`.
Attributes:
request_id: The id of the request being answered.
value: ``"allow"``/``"deny"`` for approval, the chosen option for
choice, or the entered text for input.
"""
request_id: str
value: str
def as_dict(self) -> Dict[str, Any]:
"""Convert to a serializable dictionary."""
return {"request_id": self.request_id, "value": self.value}
@runtime_checkable
class GatewayRequestChannelProtocol(Protocol):
"""Protocol for server→client interactive request/reply over the transport.
A concrete implementation is bound by the running gateway (in praisonai-bot)
and registered into the per-turn context so a blocked HITL turn can deliver
an :class:`GatewayServerRequest` to the connected client and await the
correlated :class:`GatewayServerReply`. ``open_requests`` returns the
still-open requests for a session so the gateway's resume path can re-issue
them alongside the event replay — a dropped connection never wedges a turn.
"""
async def request(
self,
req: "GatewayServerRequest",
*,
timeout_s: float = 120.0,
) -> Optional["GatewayServerReply"]:
"""Deliver ``req`` to the session's client and await its reply.
Returns the correlated :class:`GatewayServerReply`, or ``None`` when no
reply arrives within ``timeout_s`` (the request is then cleared from the
open set so it is not replayed on a later reconnect).
"""
...
def open_requests(self, session_id: str) -> List["GatewayServerRequest"]:
"""Return the still-open (unanswered) requests for ``session_id``."""
...
# ---------------------------------------------------------------------------
# Agent-callable live status/health (Issue #3688)
#
@@ -7795,16 +7364,8 @@ def _register_core_gateway_methods() -> None:
# the same action, so both names carry the WRITE scope in lockstep.
"abort": OperatorScope.WRITE,
"message_abort": OperatorScope.WRITE,
# Issue #5351: answering a server→client interactive request (approval /
# choice / input) on behalf of the blocked turn mutates it, so it needs
# the same WRITE scope as sending a message.
"server_reply": OperatorScope.WRITE,
"session.status": OperatorScope.READ,
"session.transcript": OperatorScope.READ,
# Resolve a friendly reference (short id / slug / label / ``main``
# sentinel, optionally agent-scoped) to a concrete session id (Issue
# #5383). Read-only lookup: it reveals no more than session.status.
"session.resolve": OperatorScope.READ,
"approvals.resolve": OperatorScope.APPROVALS,
"pairing.approve": OperatorScope.PAIRING,
"pairing.revoke": OperatorScope.PAIRING,
@@ -1,365 +0,0 @@
"""Session-projection reducer for Gateway clients (Issue #5324).
A pure, dependency-free state machine that folds an initial snapshot plus a
stream of :class:`~praisonaiagents.gateway.protocols.GatewayEvent`\\ s into a
consistent, de-duplicated, memory-bounded view of a session. It turns *a
stream of frames* into *a correct, resumable view* so every gateway client
(the bundled Python client, the ``praisonai-ts`` mirror, the realtime
dashboard, third-party integrations) stops re-implementing the same fiddly
reconciliation logic:
- **Snapshot + fold**: apply a snapshot, then fold live events into ``entries``.
- **Identity de-duplication**: a final assistant message that arrives both as a
streamed delta sequence *and* as a persisted transcript row renders once.
- **Optimistic reconciliation**: a locally-echoed outbound message is replaced
(not duplicated) by its durable counterpart, keyed on ``request_id``.
- **Gap handling**: a sequence gap flips ``has_transport_gap``; after a resync
the reducer converges rather than appending duplicates.
- **Bounded retention**: tracked runs are capped (LRU) so long-lived sessions
stay memory-bounded; a run with an active (unfinished) stream is never
evicted.
This module lives in **core** because it is a pure reducer built on types core
already owns (``GatewayEvent``/``GatewayMessage``/``EventType``); it must not be
trapped in the transport client so all clients can share it. It has **no**
third-party dependencies and does no I/O.
"""
from __future__ import annotations
from collections import OrderedDict
from dataclasses import dataclass, field, replace
from typing import Any, Dict, List, Optional, Tuple, Union
from .protocols import EventType, GatewayEvent, GatewayMessage
__all__ = [
"RunView",
"SessionProjectionState",
"SessionProjection",
]
def _event_type_value(event: GatewayEvent) -> str:
"""Return the wire string for an event type (enum or raw str)."""
etype = event.type
return etype.value if isinstance(etype, EventType) else str(etype)
def _run_id_of(event: GatewayEvent) -> Optional[str]:
"""Best-effort run/turn identifier for a streaming event."""
data = event.data or {}
for key in ("run_id", "turn_id", "message_id", "response_id"):
value = data.get(key)
if value:
return str(value)
return None
@dataclass(frozen=True)
class RunView:
"""Immutable view of a single agent run/turn within a session.
Attributes:
run_id: The run/turn identifier the streamed deltas belong to.
text: The accumulated streamed text so far.
done: Whether the stream has ended (``STREAM_END``) for this run.
final_message_id: ``message_id`` of the persisted final row, once the
durable transcript row for this run has been reconciled in.
"""
run_id: str
text: str = ""
done: bool = False
final_message_id: Optional[str] = None
@dataclass(frozen=True)
class SessionProjectionState:
"""Immutable, de-duplicated, bounded view of a session.
Attributes:
entries: Ordered transcript messages, de-duplicated by identity.
runs: Per-run streaming views, keyed by ``run_id``.
has_transport_gap: ``True`` when a sequence gap was observed and a
resync has not yet reconciled the state.
"""
entries: Tuple[GatewayMessage, ...] = ()
runs: Dict[str, RunView] = field(default_factory=dict)
has_transport_gap: bool = False
class SessionProjection:
"""Fold a snapshot + :class:`GatewayEvent`\\ s into a consistent view.
The reducer is idempotent: applying the same persisted message twice (e.g.
once from the live stream's ``STREAM_END`` and once as a durable transcript
row on reconnect) yields a single entry. Feed it with
:meth:`apply_snapshot` then :meth:`apply`; read :attr:`state` at any time.
Args:
max_tracked_runs: Cap on how many finished runs are tracked before the
oldest are evicted (LRU). Runs with an active/unfinished stream are
never evicted. Must be positive.
"""
def __init__(self, *, max_tracked_runs: int = 200) -> None:
if max_tracked_runs <= 0:
raise ValueError("max_tracked_runs must be positive")
self._max_tracked_runs = max_tracked_runs
self._entries: List[GatewayMessage] = []
# message_id -> index into self._entries, for identity de-dup.
self._by_message_id: Dict[str, int] = {}
# request_id -> index into self._entries, for optimistic reconciliation.
self._by_request_id: Dict[str, int] = {}
# run_id -> RunView, ordered for LRU eviction.
self._runs: "OrderedDict[str, RunView]" = OrderedDict()
self._expected_sequence: Optional[int] = None
self._has_gap: bool = False
# -- public API -------------------------------------------------------
@property
def state(self) -> SessionProjectionState:
"""Return the current immutable projection state."""
return SessionProjectionState(
entries=tuple(self._entries),
runs=dict(self._runs),
has_transport_gap=self._has_gap,
)
def apply_snapshot(
self, snapshot: Union[Dict[str, Any], Any]
) -> SessionProjectionState:
"""Replace all tracked state from a full snapshot.
A snapshot resets gap tracking and rebuilds the transcript from the
persisted rows it carries, so a resync converges rather than appending
duplicates. Accepts a mapping (``{"messages"/"entries": [...],
"cursor"/"sequence": N}``) or any object exposing ``messages``/
``entries`` and an optional ``sequence``/``cursor`` attribute.
"""
self._entries = []
self._by_message_id = {}
self._by_request_id = {}
self._runs = OrderedDict()
self._has_gap = False
self._expected_sequence = None
rows = _snapshot_rows(snapshot)
for row in rows:
message = _coerce_message(row)
if message is not None:
self._upsert_message(message)
seq = _snapshot_sequence(snapshot)
if seq is not None:
self._expected_sequence = seq + 1
return self.state
def apply(self, event: GatewayEvent) -> SessionProjectionState:
"""Fold a single event into the state and return the new state.
Idempotent for persisted messages (de-duped by ``message_id``) and for
optimistic echoes (reconciled by ``request_id``). Streaming deltas are
accumulated per run; ``STREAM_END`` closes the run. Sequence gaps set
:attr:`SessionProjectionState.has_transport_gap`.
"""
self._track_sequence(event)
etype = _event_type_value(event)
if etype in _STREAM_DELTA_TYPES:
self._apply_delta(event)
elif etype == EventType.STREAM_END.value:
self._apply_stream_end(event)
elif etype in _MESSAGE_TYPES:
self._apply_message(event)
# Other event types (health, presence, typing, ...) carry no transcript
# state; they are intentionally ignored by the projection.
return self.state
# -- sequence / gap tracking -----------------------------------------
def _track_sequence(self, event: GatewayEvent) -> None:
seq = event.sequence
if seq is None:
return
if self._expected_sequence is not None and seq > self._expected_sequence:
# A hole in the monotonic sequence: mark the view as gapped until a
# snapshot/resync reconciles it.
self._has_gap = True
# Advance the watermark to the highest sequence seen.
if self._expected_sequence is None or seq >= self._expected_sequence:
self._expected_sequence = seq + 1
# -- message handling -------------------------------------------------
def _apply_message(self, event: GatewayEvent) -> None:
message = _coerce_message(event.data)
if message is None:
return
request_id = _request_id_of(event, message)
self._upsert_message(message, request_id=request_id)
def _upsert_message(
self, message: GatewayMessage, *, request_id: Optional[str] = None
) -> None:
# Identity de-dup: same persisted row seen twice renders once.
existing = self._by_message_id.get(message.message_id)
if existing is not None:
self._entries[existing] = message
self._reindex(existing, message, request_id)
return
# Optimistic reconciliation: a durable row replaces its local echo.
if request_id is not None:
echoed = self._by_request_id.get(request_id)
if echoed is not None:
stale_id = self._entries[echoed].message_id
self._entries[echoed] = message
# Drop the echo's provisional message_id so a re-delivered echo
# cannot later overwrite the durable row via identity lookup.
if stale_id != message.message_id:
self._by_message_id.pop(stale_id, None)
self._reindex(echoed, message, request_id)
return
self._entries.append(message)
idx = len(self._entries) - 1
self._by_message_id[message.message_id] = idx
if request_id is not None:
self._by_request_id[request_id] = idx
def _reindex(
self,
idx: int,
message: GatewayMessage,
request_id: Optional[str],
) -> None:
self._by_message_id[message.message_id] = idx
if request_id is not None:
self._by_request_id[request_id] = idx
# -- streaming handling ----------------------------------------------
def _apply_delta(self, event: GatewayEvent) -> None:
run_id = _run_id_of(event)
if run_id is None:
return
chunk = _delta_text_of(event)
current = self._runs.get(run_id)
if current is None:
current = RunView(run_id=run_id, text=chunk)
else:
current = replace(current, text=current.text + chunk)
self._runs[run_id] = current
self._runs.move_to_end(run_id)
self._evict_runs()
def _apply_stream_end(self, event: GatewayEvent) -> None:
run_id = _run_id_of(event)
if run_id is None:
return
current = self._runs.get(run_id) or RunView(run_id=run_id)
final_id = None
data = event.data or {}
final_id = data.get("message_id") or data.get("final_message_id")
current = replace(
current,
done=True,
final_message_id=str(final_id) if final_id else current.final_message_id,
)
self._runs[run_id] = current
self._runs.move_to_end(run_id)
self._evict_runs()
def _evict_runs(self) -> None:
if len(self._runs) <= self._max_tracked_runs:
return
# Evict oldest finished runs first; never evict an active stream.
for run_id in list(self._runs.keys()):
if len(self._runs) <= self._max_tracked_runs:
break
if self._runs[run_id].done:
del self._runs[run_id]
# Streaming delta event types whose text is accumulated per run.
_STREAM_DELTA_TYPES = frozenset(
{
EventType.TOKEN_STREAM.value,
"delta_text",
}
)
# Event types carrying a persisted transcript message.
_MESSAGE_TYPES = frozenset(
{
EventType.MESSAGE.value,
"message_persisted",
}
)
def _delta_text_of(event: GatewayEvent) -> str:
data = event.data or {}
for key in ("delta", "text", "content", "token", "chunk"):
value = data.get(key)
if isinstance(value, str):
return value
return ""
def _request_id_of(
event: GatewayEvent, message: GatewayMessage
) -> Optional[str]:
data = event.data or {}
rid = data.get("request_id") or message.metadata.get("request_id")
return str(rid) if rid else None
def _coerce_message(row: Any) -> Optional[GatewayMessage]:
if isinstance(row, GatewayMessage):
return row
if isinstance(row, dict):
payload = row
# A "message" event wraps the row under a "message" key in some frames.
if "message" in row and isinstance(row["message"], (dict, GatewayMessage)):
inner = row["message"]
return inner if isinstance(inner, GatewayMessage) else _coerce_message(inner)
if "content" not in payload and "message_id" not in payload:
return None
try:
return GatewayMessage.from_dict(payload)
except Exception:
return None
return None
def _snapshot_rows(snapshot: Any) -> List[Any]:
if isinstance(snapshot, dict):
for key in ("messages", "entries", "transcript"):
rows = snapshot.get(key)
if isinstance(rows, list):
return rows
return []
for attr in ("messages", "entries", "transcript"):
rows = getattr(snapshot, attr, None)
if isinstance(rows, list):
return rows
return []
def _snapshot_sequence(snapshot: Any) -> Optional[int]:
if isinstance(snapshot, dict):
for key in ("sequence", "cursor", "seq"):
value = snapshot.get(key)
if isinstance(value, int):
return value
return None
for attr in ("sequence", "cursor", "seq"):
value = getattr(snapshot, attr, None)
if isinstance(value, int):
return value
return None
@@ -53,9 +53,6 @@ class HookRegistry:
}
self._enabled = enabled
self._global_timeout = 60.0
# Monotonic counter giving each registered hook a stable secondary sort
# key so equal-priority hooks keep deterministic registration order.
self._seq_counter = 0
# Guard mutation/iteration of the per-event hook lists. The default
# registry is a cross-agent process-wide singleton, so a plugin
# (un)registering a hook on one thread can race a get_hooks() iteration
@@ -83,8 +80,6 @@ class HookRegistry:
The hook ID
"""
with self._lock:
hook._seq = self._seq_counter
self._seq_counter += 1
self._hooks[hook.event].append(hook)
logger.debug(f"Registered hook '{hook.name}' for event '{hook.event.value}'")
return hook.id
@@ -98,8 +93,7 @@ class HookRegistry:
description: Optional[str] = None,
sequential: bool = False,
timeout: float = 60.0,
is_async: bool = False,
priority: int = 100
is_async: bool = False
) -> str:
"""
Register a Python function as a hook.
@@ -113,7 +107,6 @@ class HookRegistry:
sequential: Whether to run sequentially
timeout: Timeout in seconds
is_async: Whether the function is async
priority: Execution order (lower runs earlier). Default 100
Returns:
The hook ID
@@ -126,8 +119,7 @@ class HookRegistry:
description=description,
sequential=sequential,
timeout=timeout,
is_async=is_async,
priority=priority
is_async=is_async
)
return self.register(hook)
@@ -141,8 +133,7 @@ class HookRegistry:
sequential: bool = False,
timeout: float = 60.0,
env: Optional[Dict[str, str]] = None,
shell: bool = True,
priority: int = 100
shell: bool = True
) -> str:
"""
Register a shell command as a hook.
@@ -157,7 +148,6 @@ class HookRegistry:
timeout: Timeout in seconds
env: Additional environment variables
shell: Whether to run in shell mode
priority: Execution order (lower runs earlier). Default 100
Returns:
The hook ID
@@ -171,8 +161,7 @@ class HookRegistry:
sequential=sequential,
timeout=timeout,
env=env or {},
shell=shell,
priority=priority
shell=shell
)
return self.register(hook)
@@ -181,8 +170,7 @@ class HookRegistry:
event: HookEvent,
matcher: Optional[str] = None,
sequential: bool = False,
timeout: float = 60.0,
priority: int = 100
timeout: float = 60.0
) -> Callable:
"""
Decorator to register a function as a hook.
@@ -192,7 +180,6 @@ class HookRegistry:
matcher: Optional regex pattern to match targets
sequential: Whether to run sequentially
timeout: Timeout in seconds
priority: Execution order (lower runs earlier). Default 100
Returns:
Decorator function
@@ -212,8 +199,7 @@ class HookRegistry:
matcher=matcher,
sequential=sequential,
timeout=timeout,
is_async=is_async,
priority=priority
is_async=is_async
)
@wraps(func)
@@ -267,10 +253,6 @@ class HookRegistry:
with self._lock:
hooks = list(self._hooks.get(event, []))
# Sort by priority (lower runs earlier); ties fall back to registration
# order via the monotonic _seq for deterministic, reproducible ordering.
hooks.sort(key=lambda h: (h.priority, h._seq))
if target is None:
return [h for h in hooks if h.enabled]
@@ -302,25 +284,20 @@ class HookRegistry:
Dictionary mapping event names to hook info
"""
result = {}
with self._lock:
for event in HookEvent:
hooks = list(self._hooks[event])
if hooks:
# Present hooks in effective execution order so an operator
# can see the sequence get_hooks() will run them in.
hooks.sort(key=lambda h: (h.priority, h._seq))
result[event.value] = [
{
"id": h.id,
"name": h.name,
"type": "command" if isinstance(h, CommandHook) else "function",
"matcher": h.matcher,
"enabled": h.enabled,
"sequential": h.sequential,
"priority": h.priority
}
for h in hooks
]
for event in HookEvent:
hooks = self._hooks[event]
if hooks:
result[event.value] = [
{
"id": h.id,
"name": h.name,
"type": "command" if isinstance(h, CommandHook) else "function",
"matcher": h.matcher,
"enabled": h.enabled,
"sequential": h.sequential
}
for h in hooks
]
return result
def enable_hook(self, hook_id: str) -> bool:
@@ -382,7 +359,7 @@ def set_default_registry(registry: HookRegistry):
def add_hook(
event: Union[str, HookEvent],
callback: Optional[Callable[[HookInput], HookResult]] = None,
priority: int = 100,
priority: int = 10,
matcher: Optional[str] = None
) -> Union[str, Callable]:
"""Register a hook callback. Simplified API.
@@ -399,9 +376,7 @@ def add_hook(
Args:
event: Hook event name ('before_tool', 'after_llm', etc.) or HookEvent enum
callback: Function to call when hook fires (optional when using as decorator)
priority: Execution order (lower = earlier). Default 100 — same default
bucket as ``register_function``/``on``, so hooks added through either
API keep registration order unless an explicit priority is given.
priority: Execution order (lower = earlier). Default 10. (Reserved for future use)
matcher: Optional regex pattern to match specific targets (e.g., tool names)
Returns:
@@ -427,8 +402,7 @@ def add_hook(
get_default_registry().register_function(
event=event,
func=func,
matcher=matcher,
priority=priority
matcher=matcher
)
return func
return decorator
@@ -436,8 +410,7 @@ def add_hook(
return get_default_registry().register_function(
event=event,
func=callback,
matcher=matcher,
priority=priority
matcher=matcher
)
def remove_hook(hook_id: str) -> bool:
@@ -102,41 +102,25 @@ class HookRunner:
if not hooks:
return []
# The registry hands us hooks already sorted by (priority, _seq), so the
# list is in the exact order the operator asked them to run. Preserve
# that global order across parallel/sequential boundaries: walk the list
# and run each *contiguous* run of same-mode hooks together. A redact
# (sequential, priority 10) that must precede a trace (parallel,
# priority 20) therefore still runs first — splitting all parallel hooks
# to the front would have leaked the unredacted input to the tracer.
# Adjacent parallel hooks are still gathered concurrently; only their
# position relative to a lower-priority sequential hook is respected.
results: List[HookExecutionResult] = []
current_input = input_data
i = 0
n = len(hooks)
while i < n:
j = i + 1
is_seq = hooks[i].sequential
while j < n and hooks[j].sequential == is_seq:
j += 1
run = hooks[i:j]
if is_seq:
run_results = await self._execute_sequential(run, event, current_input)
# A sequential run may have blocked; stop the whole chain so a
# later parallel group never runs after a denial (matches the
# sequential-hook break semantics).
results.extend(run_results)
if any(
r.output and r.output.is_denied() for r in run_results
):
break
else:
run_results = await self._execute_parallel(
run, event, current_input
)
results.extend(run_results)
i = j
# Separate sequential and parallel hooks
sequential_hooks = [h for h in hooks if h.sequential]
parallel_hooks = [h for h in hooks if not h.sequential]
results = []
# Execute parallel hooks first
if parallel_hooks:
parallel_results = await self._execute_parallel(
parallel_hooks, event, input_data
)
results.extend(parallel_results)
# Execute sequential hooks
if sequential_hooks:
sequential_results = await self._execute_sequential(
sequential_hooks, event, input_data
)
results.extend(sequential_results)
return results
@@ -316,8 +316,6 @@ class HookDefinition:
name: Optional[str] = None
description: Optional[str] = None
timeout: float = 60.0 # Timeout in seconds
priority: int = 100 # Execution order (lower runs earlier)
_seq: int = field(default=0, repr=False) # Registration order for stable ties
def matches(self, target: str) -> bool:
"""Check if this hook matches the target (tool name, etc.)."""
@@ -211,8 +211,6 @@ class ChromaKnowledgeAdapter:
**kwargs: Any):
"""Search for relevant content in ChromaDB."""
from ..models import SearchResult, SearchResultItem
from ..protocols import require_scope
require_scope(user_id, agent_id, run_id, "search", backend="chroma")
# Get embedding for query
embedding_model, embedding_kwargs = self._embedding_call()
@@ -238,15 +236,9 @@ class ChromaKnowledgeAdapter:
if run_id is not None:
where_filter["run_id"] = run_id
# Merge with additional filters, but never let a caller-supplied filter
# overwrite the tenant-scope constraint: filters={"user_id": "bob"} must
# not widen an already-scoped query back onto another tenant's data
# (issue #5319).
# Merge with additional filters
if filters:
for key, value in filters.items():
if key in where_filter:
continue
where_filter[key] = value
where_filter.update(filters)
# Search ChromaDB
try:
@@ -297,8 +289,6 @@ class ChromaKnowledgeAdapter:
**kwargs: Any):
"""Add content to ChromaDB."""
from ..models import AddResult
from ..protocols import require_scope
require_scope(user_id, agent_id, run_id, "add", backend="chroma")
# Convert content to string
content_str = str(content)
@@ -385,8 +375,6 @@ class ChromaKnowledgeAdapter:
run_id: Optional[str] = None, limit: int = 100, **kwargs: Any):
"""Get all items from ChromaDB."""
from ..models import SearchResult, SearchResultItem
from ..protocols import require_scope
require_scope(user_id, agent_id, run_id, "get_all", backend="chroma")
try:
# ChromaDB doesn't support get_all with filters easily, so we'll use peek
@@ -522,11 +510,8 @@ class SQLiteKnowledgeAdapter:
**kwargs: Any):
"""Search for relevant content in SQLite."""
from ..models import SearchResult, SearchResultItem
from ..protocols import require_scope
import json
require_scope(user_id, agent_id, run_id, "search", backend="sqlite")
conn = self._get_conn()
# Build query with filters
@@ -543,19 +528,6 @@ class SQLiteKnowledgeAdapter:
sql += " AND run_id = ?"
params.append(run_id)
# Apply metadata filters against the stored JSON metadata column so a
# filtered search returns the same subset chroma/mem0 already return
# instead of silently discarding the filters (issue #5319).
if filters:
for key, value in filters.items():
# Quote the key in the JSON path so a top-level metadata key that
# itself contains dots (e.g. "source.type") is matched literally
# instead of being interpreted as a nested path (issue #5319).
escaped_key = key.replace('"', '""')
sql += " AND json_extract(metadata, ?) = ?"
params.append(f'$."{escaped_key}"')
params.append(value)
sql += " LIMIT ?"
params.append(limit)
@@ -579,12 +551,9 @@ class SQLiteKnowledgeAdapter:
**kwargs: Any):
"""Add content to SQLite."""
from ..models import AddResult
from ..protocols import require_scope
import json
import time
require_scope(user_id, agent_id, run_id, "add", backend="sqlite")
content_str = str(content)
doc_id = str(time.time_ns())
metadata_json = json.dumps(metadata or {})
@@ -627,11 +596,8 @@ class SQLiteKnowledgeAdapter:
run_id: Optional[str] = None, limit: int = 100, **kwargs: Any):
"""Get all items from SQLite."""
from ..models import SearchResult, SearchResultItem
from ..protocols import require_scope
import json
require_scope(user_id, agent_id, run_id, "get_all", backend="sqlite")
conn = self._get_conn()
# Build query with filters
@@ -169,12 +169,7 @@ class MongoDBKnowledgeAdapter:
Returns:
Document ID of inserted content
Raises:
ScopeRequiredError: If no scope identifier is provided
"""
from ..protocols import require_scope
require_scope(user_id, agent_id, run_id, "add", backend="mongodb")
try:
# Generate embedding if model is available
embedding = None
@@ -215,7 +210,6 @@ class MongoDBKnowledgeAdapter:
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
filters: Optional[Dict[str, Any]] = None,
**kwargs
) -> List[Dict[str, Any]]:
"""
@@ -227,17 +221,11 @@ class MongoDBKnowledgeAdapter:
user_id: Optional user scope for tenant isolation
agent_id: Optional agent scope for tenant isolation
run_id: Optional run scope for tenant isolation
filters: Optional metadata filters applied to the query
**kwargs: Additional search parameters
Returns:
List of matching documents
Raises:
ScopeRequiredError: If no scope identifier is provided
"""
from ..protocols import require_scope
require_scope(user_id, agent_id, run_id, "search", backend="mongodb")
try:
results = []
@@ -249,13 +237,6 @@ class MongoDBKnowledgeAdapter:
"run_id": run_id,
}.items() if v is not None
}
# Fold metadata filters in alongside the scope so a filtered search
# returns the same subset chroma/mem0 already return instead of
# silently discarding the filters (issue #5319).
if filters:
for key, value in filters.items():
scope_filter[f"metadata.{key}"] = value
# Try vector search first if available
if self.use_vector_search and self.embedding_model and self._is_atlas_connection():
@@ -351,127 +332,6 @@ class MongoDBKnowledgeAdapter:
logger.error(f"Failed to delete document {document_id}: {e}")
return False
def get(self, item_id: str, **kwargs):
"""Get a specific document by ID from the knowledge store.
Returns a ``SearchResultItem`` (or ``None``) so the result shape matches
the shared KnowledgeStoreProtocol and the other adapters — a caller that
reads ``.text``/``.metadata`` or calls ``.to_dict()`` must not break when
it switches to the MongoDB backend (issue #5319).
"""
from ..models import SearchResultItem
try:
from bson import ObjectId
doc = self.collection.find_one({"_id": ObjectId(item_id)})
if not doc:
return None
return SearchResultItem(
id=str(doc["_id"]),
text=doc.get("content", ""),
metadata=doc.get("metadata", {}) or {},
score=1.0,
)
except Exception as e:
logger.error(f"Failed to get document {item_id}: {e}")
return None
def get_all(
self,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
limit: int = 100,
**kwargs
):
"""Get all documents matching the given scope.
Returns a ``SearchResult`` to match the shared KnowledgeStoreProtocol
and the other adapters.
Raises:
ScopeRequiredError: If no scope identifier is provided
"""
from ..models import SearchResult, SearchResultItem
from ..protocols import require_scope
require_scope(user_id, agent_id, run_id, "get_all", backend="mongodb")
try:
scope_filter = {
k: v for k, v in {
"user_id": user_id,
"agent_id": agent_id,
"run_id": run_id,
}.items() if v is not None
}
items = []
for doc in self.collection.find(scope_filter).limit(limit):
items.append(SearchResultItem(
id=str(doc["_id"]),
text=doc.get("content", ""),
metadata=doc.get("metadata", {}) or {},
score=1.0,
))
return SearchResult(results=items)
except Exception as e:
logger.error(f"Failed to get_all from MongoDB: {e}")
return SearchResult(results=[])
def update(self, item_id: str, content: Any, **kwargs):
"""Update an existing document's content/metadata by ID.
Returns an ``AddResult`` to match the shared KnowledgeStoreProtocol and
the other adapters (issue #5319).
"""
from ..models import AddResult
try:
from bson import ObjectId
update_fields: Dict[str, Any] = {"content": str(content)}
if "metadata" in kwargs and kwargs["metadata"] is not None:
update_fields["metadata"] = kwargs["metadata"]
result = self.collection.update_one(
{"_id": ObjectId(item_id)},
{"$set": update_fields}
)
if result.matched_count > 0:
return AddResult(success=True, id=item_id)
return AddResult(success=False, message=f"No document matched id {item_id}")
except Exception as e:
logger.error(f"Failed to update document {item_id}: {e}")
return AddResult(success=False, message=str(e))
def delete_all(
self,
*,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs
) -> bool:
"""Delete all documents matching the given scope.
Enforces the same tenant-scope contract as the mem0 backend: an
unscoped call raises ``ScopeRequiredError`` instead of silently wiping
every tenant's documents from the (typically shared) collection.
Raises:
ScopeRequiredError: If no scope identifier is provided
"""
from ..protocols import require_scope
require_scope(user_id, agent_id, run_id, "delete_all", backend="mongodb")
try:
scope_filter = {
k: v for k, v in {
"user_id": user_id,
"agent_id": agent_id,
"run_id": run_id,
}.items() if v is not None
}
self.collection.delete_many(scope_filter)
return True
except Exception as e:
logger.error(f"Failed to delete_all from MongoDB: {e}")
return False
def close(self):
"""Close MongoDB connection."""
if hasattr(self, 'client'):
@@ -207,10 +207,9 @@ def _recover_json_tool_calls(response_text: str, tools: List[Dict[str, Any]]) ->
if isinstance(response_json, list):
tool_calls: List[Dict[str, Any]] = []
for idx, tool_json in enumerate(response_json):
name = tool_json.get("name") if isinstance(tool_json, dict) else None
if name:
if isinstance(tool_json, dict) and tool_json.get("name") in allowed:
tool_calls.append(_make_tool_call(
name, tool_json.get("arguments", {}), response_text, idx))
tool_json["name"], tool_json.get("arguments", {}), response_text, idx))
if tool_calls:
return tool_calls
except (json.JSONDecodeError, TypeError, KeyError):
@@ -588,48 +588,12 @@ def get_retry_delay(category: ErrorCategory, attempt: int = 1, base_delay: float
return 0
def _retry_after_http_date_delay(value: str) -> Optional[float]:
"""Return seconds until an HTTP-date ``Retry-After`` value, else ``None``.
RFC 7231 allows ``Retry-After`` to be either delta-seconds or an absolute
HTTP-date. Delta-seconds is handled by the caller; this parses the date
form using the stdlib (no new dependency) and returns the non-negative
delay from now. A value in the past or unparseable yields ``None`` so the
caller falls through to its other signals.
"""
try:
from email.utils import parsedate_to_datetime
except Exception: # pragma: no cover - stdlib always present
return None
try:
when = parsedate_to_datetime(value.strip())
except (TypeError, ValueError):
return None
if when is None:
return None
import time as _time
from datetime import timezone
# An HTTP-date is UTC; a naive parse result is treated as UTC per the RFC.
if when.tzinfo is None:
when = when.replace(tzinfo=timezone.utc)
delay = when.timestamp() - _time.time()
return delay if delay > 0 else None
def extract_retry_after(
error: Exception, cap_seconds: float = 300.0,
) -> Optional[float]:
def extract_retry_after(error: Exception) -> Optional[float]:
"""Extract Retry-After header value from rate limit errors.
Args:
error: Exception potentially containing Retry-After info
cap_seconds: Upper bound applied to the parsed value. Defaults to 300s
(5 minutes) for in-tick backoff callers. A scheduler *hold* passes a
larger cap so a job can be parked past a long provider reset window
(e.g. an hourly quota) instead of re-firing every tick — capping at
5 minutes there would defeat the park.
Returns:
Delay in seconds if found, None otherwise
"""
@@ -644,21 +608,14 @@ def extract_retry_after(
retry_after = None
if retry_after is not None:
try:
return min(float(retry_after), cap_seconds)
return min(float(retry_after), 300.0) # Cap at 5 minutes
except (ValueError, TypeError):
# Not a plain delta-seconds value — RFC 7231 also permits an
# HTTP-date. Parse it so a provider that sends an absolute reset
# instant still yields a usable delay (else a long quota window
# would leave the job un-parked). A past/unparseable date yields
# no delay and falls through to the other signals.
delay = _retry_after_http_date_delay(str(retry_after))
if delay is not None:
return min(delay, cap_seconds)
pass # Not a plain number (could be an HTTP-date); fall through
# 2. Some SDKs expose a numeric ``retry_after`` attribute directly.
retry_after_attr = getattr(error, "retry_after", None)
if isinstance(retry_after_attr, (int, float)):
return min(float(retry_after_attr), cap_seconds)
return min(float(retry_after_attr), 300.0)
error_str = str(error)
@@ -675,7 +632,7 @@ def extract_retry_after(
if match:
try:
delay = float(match.group(1))
return min(delay, cap_seconds)
return min(delay, 300.0) # Cap at 5 minutes
except (ValueError, IndexError):
continue
@@ -359,17 +359,10 @@ class FailoverManager:
if profile.cooldown_until and current_time >= profile.cooldown_until:
profile.reset()
# Find available profiles (already priority-sorted by add_profile).
available = [p for p in self._profiles if p.is_available]
if available:
if self.config.rotate_on_success:
# Round-robin among the current highest-priority tier only,
# so equal-priority credentials configured to spread load
# actually share traffic instead of one absorbing all of it.
top_priority = available[0].priority
tier = [p for p in available if p.priority == top_priority]
return tier[self._current_index % len(tier)]
return available[0]
# Find first available profile
for profile in self._profiles:
if profile.is_available:
return profile
# If none available, return the one with shortest remaining cooldown
available_profiles = [p for p in self._profiles if p.status != ProviderStatus.DISABLED]
@@ -450,10 +443,6 @@ class FailoverManager:
profile: The profile that succeeded
"""
with self._lock:
# Advance the round-robin cursor so the next get_next_profile()
# hands work to a different equal-priority credential.
if self.config.rotate_on_success:
self._current_index += 1
if profile.status != ProviderStatus.AVAILABLE:
# Snapshot the cooldown this success is recovering from BEFORE
# resetting, so we can avoid clobbering a newer fleet-wide bench.
+133 -470
View File
@@ -2,12 +2,10 @@ import logging
from praisonaiagents._logging import get_logger
import os
import copy
import contextvars
import warnings
import re
import inspect
import asyncio
import threading
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union, Literal, Callable, TYPE_CHECKING, Protocol
@@ -39,12 +37,6 @@ from ..tools.schema import build_tool_definition
_TOOL_ARGUMENTS_PARSE_FAILED = object()
# ``litellm.success_callback``/``_async_success_callback``/``callbacks`` are
# process-global lists shared by every LLM instance. Guard the read-modify-write
# in ``_setup_event_tracking`` so concurrent instances don't corrupt them.
_EVENT_TRACKING_LOCK = threading.Lock()
def _ollama_tool_result_is_successful(tool_result: Any) -> bool:
"""Return whether a raw Ollama callback result is safe to chain.
@@ -221,24 +213,6 @@ class LLM:
# Class-level flag for one-time logging configuration
_logging_configured = False
# Current agent name for per-call token/cost attribution is stored in a
# PER-INSTANCE ContextVar (see __init__: self._current_agent_name_var).
#
# ContextVar, not a plain attribute: a single LLM instance can back two
# different Agents (Agent(llm=<shared LLM>) is supported), and
# PraisonAIAgents.arun_all_tasks dispatches tasks concurrently via
# asyncio.gather on one thread. With a plain attribute, task A sets
# "Researcher", awaits the network call, task B sets "Writer" on the SAME
# object, and when A's response returns it reads "Writer" -- A's tokens get
# attributed to B. A ContextVar is copied per task at gather/create_task
# time, so each concurrently-gathered coroutine keeps its own value across
# the await (same reasoning as _tools_scope_depths in agents/agents.py).
#
# Per-instance (not class-level): two distinct LLM instances -- including an
# original and its clone -- must keep independent attribution state within
# the same context, otherwise the last writer's agent name leaks across LLMs.
# Class-level cache for LiteLLM module (avoids repeated import overhead)
_litellm_module = None
@@ -598,17 +572,7 @@ Respond with ONLY a valid JSON tool call in this format:
# Token tracking
self.last_token_metrics: Optional[TokenMetrics] = None
self.session_token_metrics: Optional[TokenMetrics] = None
# Agent attribution is task-local: a single LLM instance is often shared
# by several concurrently-running agents (gateway channels, asyncio.gather),
# so a plain attribute would let one agent's set_current_agent() clobber
# another's mid-await and misattribute token spend (issues #5052/#3933).
# A ContextVar keeps the value isolated per asyncio task / thread.
self._current_agent_name_var: contextvars.ContextVar[Optional[str]] = (
contextvars.ContextVar("current_agent_name", default=None)
)
self._current_agent_id_var: contextvars.ContextVar[Optional[str]] = (
contextvars.ContextVar("current_agent_id", default=None)
)
self.current_agent_name: Optional[str] = None
# Rate limiting and retry settings
self._rate_limiter = extra_settings.get('rate_limiter', None)
@@ -944,8 +908,7 @@ Respond with ONLY a valid JSON tool call in this format:
# Billing/quota issues (must be checked before generic 429/rate-limit)
if any(indicator in error_str for indicator in [
"insufficient quota", "quota exceeded", "billing", "credit",
"payment required", "subscription required", "subscription expired",
"plan limit"
"payment required", "subscription required", "plan limit"
]):
return "billing"
@@ -971,27 +934,17 @@ Respond with ONLY a valid JSON tool call in this format:
]):
return "model_not_found"
# Format errors (checked before empty_response so "malformed response"
# and "*parse error*" are classified as format rather than empty output)
# Empty or malformed responses
if any(indicator in error_str for indicator in [
"validation error", "invalid format", "parse error", "parsing error",
"malformed", "invalid json", "schema error", "decode error"
]):
return "format_error"
# Empty or missing response content
if any(indicator in error_str for indicator in [
"empty response", "no response", "no content", "blank output",
"null response", "invalid response format",
"json decode error", "unexpected end of json"
"empty response", "no response", "invalid response format",
"json decode error", "unexpected end of json", "malformed response"
]):
return "empty_response"
# Service overloaded
if any(indicator in error_str for indicator in [
"overloaded", "service unavailable", "temporarily unavailable",
"server overloaded", "server busy", "try again later",
"503", "502", "500"
"server overloaded", "503", "502", "500"
]):
return "overloaded"
@@ -1002,6 +955,13 @@ Respond with ONLY a valid JSON tool call in this format:
]):
return "idle_timeout"
# Format errors
if any(indicator in error_str for indicator in [
"validation error", "invalid format", "parse error",
"malformed", "invalid json", "schema error"
]):
return "format_error"
# Default fallback
return "unknown"
@@ -1084,21 +1044,13 @@ Respond with ONLY a valid JSON tool call in this format:
is_retryable=True
)
# Auth errors - try profile rotation if available, else surface.
# Without a failover manager there is no alternate credential to try,
# so blind retry cannot succeed — surface the auth failure immediately.
if error_kind == "auth":
if self._failover_manager:
return FailoverDecision(
action="rotate_profile",
reason=error_kind,
backoff_ms=1000, # Brief delay before trying new profile
is_retryable=True
)
# Auth errors - try profile rotation if available
if error_kind == "auth" and self._failover_manager:
return FailoverDecision(
action="surface_error",
action="rotate_profile",
reason=error_kind,
is_retryable=False
backoff_ms=1000, # Brief delay before trying new profile
is_retryable=True
)
# Overloaded/timeout - retry with exponential backoff
@@ -1177,33 +1129,6 @@ Respond with ONLY a valid JSON tool call in this format:
model = getattr(response, "model", None)
return model or self.model
def _select_active_profile(self, kwargs: dict):
"""Pick the profile a fresh call should start on, honoring rotation.
Defaults to the construction-time ``self._current_profile`` so existing
behaviour and multi-agent safety are unchanged (the shared instance
state is never mutated). Only when the manager is configured with
``rotate_on_success`` do we re-select via ``get_next_profile()`` so a
long-lived, shared ``LLM`` actually spreads successful traffic across
the equal-priority tier instead of pinning its first profile forever.
The chosen profile travels in per-call kwargs, never on ``self``.
"""
profile = self._current_profile
manager = self._failover_manager
if manager is None:
return profile, kwargs
rotate = getattr(getattr(manager, "config", None), "rotate_on_success", False)
if not rotate:
return profile, kwargs
try:
rotated = manager.get_next_profile()
except Exception: # pragma: no cover - defensive
return profile, kwargs
if rotated and rotated != profile:
profile = rotated
kwargs = self._apply_profile_to_kwargs(rotated, kwargs)
return profile, kwargs
def _apply_profile_to_kwargs(self, profile: "AuthProfile", kwargs: dict) -> dict:
"""Return a new kwargs dict with profile overrides applied.
@@ -1442,10 +1367,8 @@ Respond with ONLY a valid JSON tool call in this format:
last_error = None
# Track the failover profile per call (never self._current_profile) so
# concurrent calls on one shared LLM cannot corrupt each other's
# failover bookkeeping (issue #3613 gap 2). When rotate_on_success is
# configured, pick the next profile in the tier so a shared instance
# actually rotates instead of pinning its construction-time profile.
active_profile, kwargs = self._select_active_profile(kwargs)
# failover bookkeeping (issue #3613 gap 2).
active_profile = self._current_profile
for attempt in range(self._max_retries + 1):
try:
@@ -1528,10 +1451,8 @@ Respond with ONLY a valid JSON tool call in this format:
last_error = None
# Track the failover profile per call (never self._current_profile) so
# concurrent coroutines on one shared LLM cannot corrupt each other's
# failover bookkeeping (issue #3613 gap 2). When rotate_on_success is
# configured, pick the next profile in the tier so a shared instance
# actually rotates instead of pinning its construction-time profile.
active_profile, kwargs = self._select_active_profile(kwargs)
# failover bookkeeping (issue #3613 gap 2).
active_profile = self._current_profile
for attempt in range(self._max_retries + 1):
try:
@@ -1888,22 +1809,12 @@ Respond with ONLY a valid JSON tool call in this format:
self,
arguments: Dict[str, Any],
tool_result_mapping: Dict[str, Any],
function_name: Optional[str] = None,
tools: Optional[List] = None,
) -> Dict[str, Any]:
"""Delegate Ollama function-name substitution to the provider adapter.
When ``function_name``/``tools`` are supplied, a parameter the target
tool declares as string-accepting is left untouched even if its literal
value happens to match an earlier tool's name (e.g. ``city="get_weather"``).
Otherwise a legitimate string argument would be silently overwritten with
an unrelated prior result. The substitution only ever targets values the
model meant as result references, never declared string inputs.
"""
"""Delegate Ollama function-name substitution to the provider adapter."""
adapter = getattr(self, "_provider_adapter", None)
if adapter is not None:
try:
resolved = adapter.resolve_chained_arguments(
return adapter.resolve_chained_arguments(
arguments,
tool_result_mapping,
correlation_id=self._ollama_correlation_id(),
@@ -1912,48 +1823,14 @@ Respond with ONLY a valid JSON tool call in this format:
# Keep compatibility with lightweight test/custom adapters.
resolver = getattr(adapter, "resolve_chained_arguments", None)
if resolver is not None:
resolved = resolver(arguments, tool_result_mapping)
else:
resolved = None
if resolved is not None:
return self._preserve_declared_string_args(
arguments, resolved, tool_result_mapping, function_name, tools
)
return resolver(arguments, tool_result_mapping)
from .adapters import resolve_ollama_chained_arguments
resolved = resolve_ollama_chained_arguments(
return resolve_ollama_chained_arguments(
arguments,
tool_result_mapping,
correlation_id=self._ollama_correlation_id(),
)
return self._preserve_declared_string_args(
arguments, resolved, tool_result_mapping, function_name, tools
)
def _preserve_declared_string_args(
self,
original: Dict[str, Any],
resolved: Dict[str, Any],
tool_result_mapping: Dict[str, Any],
function_name: Optional[str],
tools: Optional[List],
) -> Dict[str, Any]:
"""Undo result-reference substitution for declared string parameters.
Only acts when schema info is available and a value was actually changed
by the resolver. A no-op for callers that pass no tools (backward
compatible with every existing call site).
"""
if not function_name or not tools or resolved is original:
return resolved
for arg_name, orig_value in original.items():
if not (isinstance(orig_value, str) and orig_value in tool_result_mapping):
continue
if resolved.get(arg_name) is orig_value:
continue # resolver left it alone already
if self._param_accepts_string(function_name, arg_name, tools):
resolved[arg_name] = orig_value
return resolved
def _record_ollama_tool_result(
self,
@@ -2043,93 +1920,16 @@ Respond with ONLY a valid JSON tool call in this format:
return False
if tool_calls:
return False
force_tool_usage = getattr(self, 'force_tool_usage', 'never')
if force_tool_usage == 'never':
if self.force_tool_usage == 'never':
return False
elif force_tool_usage == 'always':
elif self.force_tool_usage == 'always':
return True
elif force_tool_usage == 'auto':
elif self.force_tool_usage == 'auto':
# Auto mode: only for Ollama on first iteration when model ignores tools
return self._is_ollama_provider() and iteration_count == 0
return False
def _force_tool_usage_message(self, response_text: str, tool_calls: Optional[List], formatted_tools: Optional[List], iteration_count: int) -> Optional[str]:
"""Return the force-tool-usage nudge, or None when no nudge is due.
Shared by both the sync and async tool loops so a `force_tool_usage`
setting an agent accepted is honoured identically whether it calls or
awaits. Returns the user-message content string (containing the tool
names) when the model ignored the tools, else None.
"""
if not self._should_force_tool_usage(response_text, tool_calls, formatted_tools, iteration_count):
return None
tool_names = self._get_tool_names_for_prompt(formatted_tools)
return self.FORCE_TOOL_USAGE_PROMPT.format(tool_names=tool_names)
def _tool_repair_message(self, tool_calls: Optional[List], formatted_tools: Optional[List]) -> Optional[str]:
"""Return a repair nudge for malformed tool calls, or None.
Charges the `max_tool_repairs` budget once per repair attempt (via
`_current_repair_count`). Returns None when the budget is spent or the
tool calls validate cleanly. Shared by the sync and async loops.
"""
max_tool_repairs = getattr(self, 'max_tool_repairs', 0)
if not tool_calls or max_tool_repairs <= 0:
return None
repair_attempt_count = getattr(self, '_current_repair_count', 0)
if repair_attempt_count >= max_tool_repairs:
return None
validation_errors = []
for tc in tool_calls:
error = self._validate_tool_call(tc, formatted_tools)
if error:
validation_errors.append(error)
if not validation_errors:
return None
error_msg = "; ".join(validation_errors)
tool_schemas = self._get_tool_schemas_for_prompt(formatted_tools)
self._current_repair_count = repair_attempt_count + 1
return self.TOOL_CALL_REPAIR_PROMPT.format(error=error_msg, tool_schemas=tool_schemas)
@staticmethod
def _param_accepts_string(tool_name: str, param: str, formatted_tools: Optional[List]) -> bool:
"""True if the named parameter's declared schema accepts a string.
Used to narrow `tool_result_mapping` substitution: a declared string
parameter must keep the value the model sent, even when that value
happens to match a prior tool's name. Recognises plain `type`,
list-typed `type` unions, and `anyOf` unions. Never raises on a
malformed schema - an unreadable schema simply returns False.
"""
try:
if not formatted_tools:
return False
for tool in formatted_tools:
if not isinstance(tool, dict):
continue
func = tool.get('function')
if not isinstance(func, dict) or func.get('name') != tool_name:
continue
props = func.get('parameters', {}).get('properties', {})
schema = props.get(param)
if not isinstance(schema, dict):
return False
ptype = schema.get('type')
if ptype == 'string':
return True
if isinstance(ptype, list) and 'string' in ptype:
return True
any_of = schema.get('anyOf')
if isinstance(any_of, list):
for member in any_of:
if isinstance(member, dict) and member.get('type') == 'string':
return True
return False
return False
except Exception:
return False
def _validate_tool_call(self, tool_call: Dict, formatted_tools: Optional[List]) -> Optional[str]:
"""
Validate a tool call against available tools.
@@ -3928,12 +3728,13 @@ Respond with ONLY a valid JSON tool call in this format:
logging.debug(f"Parsed {len(tool_calls)} tool call(s) from XML format")
# Force tool usage logic: if model ignores tools but should use them
_force = self._force_tool_usage_message(response_text, tool_calls, formatted_tools, iteration_count)
if _force is not None:
logging.debug(f"[OLLAMA_RELIABILITY] Force tool usage triggered. Adding prompt for tools: {self._get_tool_names_for_prompt(formatted_tools)}")
if self._should_force_tool_usage(response_text, tool_calls, formatted_tools, iteration_count):
tool_names = self._get_tool_names_for_prompt(formatted_tools)
force_prompt = self.FORCE_TOOL_USAGE_PROMPT.format(tool_names=tool_names)
logging.debug(f"[OLLAMA_RELIABILITY] Force tool usage triggered. Adding prompt for tools: {tool_names}")
messages.append({
"role": "user",
"content": _force
"content": force_prompt
})
iteration_count += 1
continue
@@ -3948,16 +3749,32 @@ Respond with ONLY a valid JSON tool call in this format:
continue
# Tool call repair logic: validate and repair malformed tool calls
_repair = self._tool_repair_message(tool_calls, formatted_tools)
if _repair is not None:
logging.debug(f"[OLLAMA_RELIABILITY] Tool call repair attempt {self._current_repair_count}/{self.max_tool_repairs}")
messages.append({
"role": "user",
"content": _repair
})
iteration_count += 1
tool_calls = None # Clear invalid tool calls
continue
repair_attempt_count = getattr(self, '_current_repair_count', 0)
if tool_calls and self.max_tool_repairs > 0 and repair_attempt_count < self.max_tool_repairs:
# Validate each tool call
validation_errors = []
for tc in tool_calls:
error = self._validate_tool_call(tc, formatted_tools)
if error:
validation_errors.append(error)
if validation_errors:
# Tool call is invalid, attempt repair
error_msg = "; ".join(validation_errors)
tool_schemas = self._get_tool_schemas_for_prompt(formatted_tools)
repair_prompt = self.TOOL_CALL_REPAIR_PROMPT.format(
error=error_msg,
tool_schemas=tool_schemas
)
logging.debug(f"[OLLAMA_RELIABILITY] Tool call repair attempt {repair_attempt_count + 1}/{self.max_tool_repairs}: {error_msg}")
messages.append({
"role": "user",
"content": repair_prompt
})
self._current_repair_count = repair_attempt_count + 1
iteration_count += 1
tool_calls = None # Clear invalid tool calls
continue
# Reset repair count on successful tool call
self._current_repair_count = 0
@@ -4033,7 +3850,7 @@ Respond with ONLY a valid JSON tool call in this format:
# result, then filter unknown argument keys.
if is_ollama:
arguments = self._resolve_ollama_chained_args(
arguments, tool_result_mapping, function_name, tools
arguments, tool_result_mapping
)
# Validate and filter arguments for Ollama provider
if is_ollama and tools:
@@ -4473,7 +4290,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
agent_name=agent_name, agent_role=agent_role, agent_tools=agent_tools,
task_name=task_name, task_description=task_description, task_id=task_id)
interaction_displayed = True
return _prepare_return_value(response_text)
return response_text
if reflection_count >= max_reflect - 1:
if verbose and not interaction_displayed:
@@ -4482,7 +4299,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
agent_name=agent_name, agent_role=agent_role, agent_tools=agent_tools,
task_name=task_name, task_description=task_description, task_id=task_id)
interaction_displayed = True
return _prepare_return_value(response_text)
return response_text
reflection_count += 1
messages.extend([
@@ -4777,7 +4594,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# Resolve any same-turn result reference first.
if is_ollama:
arguments = self._resolve_ollama_chained_args(
arguments, ollama_tool_result_mapping, function_name, tools
arguments, ollama_tool_result_mapping
)
# Validate and filter before dispatch. Nothing has
# been yielded yet, so streaming consumers remain
@@ -4811,8 +4628,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
tool_results = []
for _tool_call in tool_calls_batch:
_tool_call.arguments = self._resolve_ollama_chained_args(
_tool_call.arguments, ollama_tool_result_mapping,
_tool_call.function_name, tools
_tool_call.arguments, ollama_tool_result_mapping
)
# The batch-preparation filter above has
# already removed unknown keys. Re-validate
@@ -4919,12 +4735,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
yield f"Task interrupted: {reason}"
return
# Trim the conversation before the follow-up request so
# the stream loop bounds its context like the sync and
# async loops do (each tool turn appends an assistant
# message plus one tool reply per call).
messages = self._manage_context_in_loop(messages)
# Continue conversation after tool execution - get follow-up response
try:
follow_up_response = self._completion_with_retry(
@@ -4991,7 +4801,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
max_fallback_iterations = kwargs.pop("max_iterations", self.max_iter)
while fallback_iterations < max_fallback_iterations:
fallback_iterations += 1
messages = self._manage_context_in_loop(messages)
# Chaining references are scoped to one assistant turn;
# do not let a later turn reuse an earlier result.
ollama_tool_result_mapping = {}
@@ -5087,7 +4896,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# user's tool with parameters it never declared.
if is_ollama and tools:
arguments = self._resolve_ollama_chained_args(
arguments, ollama_tool_result_mapping, function_name, tools
arguments, ollama_tool_result_mapping
)
arguments = self._validate_and_filter_ollama_arguments(
function_name, arguments, tools)
@@ -5208,9 +5017,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
max_tool_calls_per_turn: int = 10, # Loop guardrails
stream: bool = True,
parallel_tool_calls: bool = False,
return_token_usage: bool = False,
**kwargs
) -> Union[str, tuple[str, TokenUsage]]:
) -> str:
"""Async version of get_response with identical functionality."""
# G2: Cooperative cancellation - honour InterruptController between tool iterations
# so /stop (and cancellation generally) halts mid-flight runs on every provider.
@@ -5343,20 +5151,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
iteration_count = 0
tool_call_count = 0 # Track total tool calls for guardrails
final_response_text = ""
# Track the final raw LLM response so return_token_usage can extract
# usage, mirroring the sync get_response path.
_final_llm_response = None
def _prepare_return_value(text: str) -> Union[str, tuple]:
if not return_token_usage:
return text
token_usage = (
self._extract_token_usage(_final_llm_response)
if _final_llm_response else None
)
if token_usage is None:
token_usage = TokenUsage()
return text, token_usage
stored_reasoning_content = None # Store reasoning content from tool execution
accumulated_tool_results = [] # Store all tool results across iterations
# Structured stop reason (unified with the OpenAI-native path).
@@ -5368,7 +5162,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
if _is_cancelled():
reason = _cancel_reason()
logging.debug(f"Async LLM tool loop cancelled: {reason}")
return _prepare_return_value(f"Task interrupted: {reason}")
return f"Task interrupted: {reason}"
# G2: Mid-run steering - drain any pending steering notes and inject
# them as user messages so the model sees them on its next step.
_inject_steering(messages)
@@ -5404,11 +5198,9 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# usage from the final response.completed event.
if stream_response is not None:
self._track_token_usage(stream_response, self.model)
_final_llm_response = stream_response
else:
resp = await self._call_responses_api_async(**responses_params)
response_text, tool_calls, _reasoning = self._extract_from_responses_output(resp)
_final_llm_response = resp
# Build Chat-Completions-compatible final_response
final_response = {
@@ -5449,7 +5241,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
if _is_cancelled():
reason = _cancel_reason()
logging.debug(f"Async LLM tool loop cancelled before tool dispatch: {reason}")
return _prepare_return_value(f"Task interrupted: {reason}")
return f"Task interrupted: {reason}"
serializable_tool_calls = self._serialize_tool_calls(tool_calls)
messages.append({
"role": "assistant",
@@ -5526,7 +5318,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
)
)
self._record_finish_reason(resp)
_final_llm_response = resp
reasoning_content = resp["choices"][0]["message"].get("provider_specific_fields", {}).get("reasoning_content")
response_text = resp["choices"][0]["message"]["content"]
@@ -5636,7 +5427,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
)
)
self._record_finish_reason(tool_response)
_final_llm_response = tool_response
# Handle None content from Gemini
response_content = tool_response.choices[0].message.get("content")
response_text = response_content if response_content is not None else ""
@@ -5677,35 +5467,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
)
interaction_displayed = True
# Force tool usage / repair: applied AFTER both the streaming
# and non-streaming branches converge so a streaming-with-tools
# provider still honours the policy the sync path applies. Both
# lived inline in get_response only, so `max_tool_repairs` and
# `force_tool_usage` were accepted from the caller -- and set by
# OllamaAdapter for every Ollama LLM -- then silently ignored the
# moment the agent was awaited.
_force = self._force_tool_usage_message(response_text, tool_calls, formatted_tools, iteration_count)
if _force is not None:
logging.debug(f"[OLLAMA_RELIABILITY] Force tool usage triggered (async). Adding prompt for tools: {self._get_tool_names_for_prompt(formatted_tools)}")
messages.append({
"role": "user",
"content": _force
})
iteration_count += 1
continue
_repair = self._tool_repair_message(tool_calls, formatted_tools)
if _repair is not None:
logging.debug(f"[OLLAMA_RELIABILITY] Tool call repair attempt (async) {self._current_repair_count}/{self.max_tool_repairs}")
messages.append({
"role": "user",
"content": _repair
})
iteration_count += 1
tool_calls = None # Clear invalid tool calls
continue
self._current_repair_count = 0
# For Ollama, if response is empty but we have tools, prompt for tool usage
if self._is_ollama_provider() and (not response_text or response_text.strip() == "") and formatted_tools and iteration_count == 0:
messages.append({
@@ -5722,7 +5483,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
if _is_cancelled():
reason = _cancel_reason()
logging.debug(f"Async LLM tool loop cancelled before tool dispatch: {reason}")
return _prepare_return_value(f"Task interrupted: {reason}")
return f"Task interrupted: {reason}"
# Convert tool_calls to a serializable format for all providers
serializable_tool_calls = self._serialize_tool_calls(tool_calls)
# Check if it's Ollama provider
@@ -5783,7 +5544,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
if is_ollama and tools:
if parallel_tool_calls:
arguments = self._resolve_ollama_chained_args(
arguments, ollama_tool_result_mapping, function_name, tools
arguments, ollama_tool_result_mapping
)
arguments = self._validate_and_filter_ollama_arguments(
function_name, arguments, tools
@@ -5791,7 +5552,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
if is_ollama and not parallel_tool_calls:
arguments = self._resolve_ollama_chained_args(
arguments, ollama_tool_result_mapping, function_name, tools
arguments, ollama_tool_result_mapping
)
if tools:
arguments = self._validate_and_filter_ollama_arguments(
@@ -5889,7 +5650,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
)
)
self._record_finish_reason(resp)
_final_llm_response = resp
reasoning_content = resp["choices"][0]["message"].get("provider_specific_fields", {}).get("reasoning_content")
response_text = resp["choices"][0]["message"]["content"]
@@ -5942,7 +5702,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
)
)
self._record_finish_reason(resp)
_final_llm_response = resp
response_text = resp["choices"][0]["message"].get("content") or ""
# If the response also contains new tool_calls, treat this as a
# tool-calling round rather than a final answer (Anthropic pattern)
@@ -6063,7 +5822,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
agent_name=agent_name, agent_role=agent_role, agent_tools=agent_tools,
task_name=task_name, task_description=task_description, task_id=task_id)
interaction_displayed = True
return _prepare_return_value(response_text)
return response_text
if not self_reflect:
# Use final_response_text if we went through tool iterations
@@ -6094,8 +5853,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# Return reasoning content if reasoning_steps is True and we have it
if reasoning_steps and stored_reasoning_content:
return _prepare_return_value(stored_reasoning_content)
return _prepare_return_value(display_text)
return stored_reasoning_content
return display_text
# Handle self-reflection
reflection_prompt = f"""
@@ -6210,7 +5969,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
agent_name=agent_name, agent_role=agent_role, agent_tools=agent_tools,
task_name=task_name, task_description=task_description, task_id=task_id)
interaction_displayed = True
return _prepare_return_value(response_text)
return response_text
if reflection_count >= max_reflect - 1:
if verbose and not interaction_displayed:
@@ -6219,7 +5978,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
agent_name=agent_name, agent_role=agent_role, agent_tools=agent_tools,
task_name=task_name, task_description=task_description, task_id=task_id)
interaction_displayed = True
return _prepare_return_value(response_text)
return response_text
reflection_count += 1
messages.extend([
@@ -6233,7 +5992,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
except json.JSONDecodeError:
reflection_count += 1
if reflection_count >= max_reflect:
return _prepare_return_value(response_text)
return response_text
continue # Now properly in a loop
except Exception as error:
@@ -6289,16 +6048,11 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
def _setup_event_tracking(self, events: List[Any]) -> None:
"""Setup callback functions for tracking model usage.
``litellm.success_callback``, ``litellm._async_success_callback`` and
``litellm.callbacks`` are process-global lists shared by every LLM
instance. Overwriting (or type-based stripping of) them wipes callbacks
registered by other concurrent agents — or by the application itself via
litellm's documented ``litellm.success_callback.append(...)`` extension
point. So we merge into each list instead: only the callbacks *this*
instance previously registered are removed, and unrelated instances'
callbacks are never touched. An empty ``events`` list is a no-op. The
whole read-modify-write is guarded by a module-level lock so concurrent
instances can't corrupt the shared lists.
``litellm.callbacks`` is a process-global list shared by every LLM
instance. Overwriting it wipes callbacks registered by other concurrent
agents, so we merge into it instead: only the callbacks *this* instance
previously registered are removed, and unrelated instances' callbacks
are never touched. An empty ``events`` list is a no-op.
"""
if not events:
return
@@ -6311,37 +6065,29 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
"Please install it with: pip install 'praisonaiagents[llm]'"
)
with _EVENT_TRACKING_LOCK:
# Only remove the success callbacks *this* instance registered on a
# prior call — never strip by type, which would delete another
# instance's or the application's own callbacks of the same class
# (e.g. an app-registered ``litellm.success_callback.append(...)``).
for cb in getattr(self, "_registered_success_callbacks", []):
if cb in litellm.success_callback:
litellm.success_callback.remove(cb)
for cb in getattr(self, "_registered_async_success_callbacks", []):
if cb in litellm._async_success_callback:
litellm._async_success_callback.remove(cb)
for event in events:
if event not in litellm.success_callback:
litellm.success_callback.append(event)
if event not in litellm._async_success_callback:
litellm._async_success_callback.append(event)
self._registered_success_callbacks = list(events)
self._registered_async_success_callbacks = list(events)
event_types = [type(event) for event in events]
# Remove old events of same type
for event in litellm.success_callback[:]:
if type(event) in event_types:
litellm.success_callback.remove(event)
for event in litellm._async_success_callback[:]:
if type(event) in event_types:
litellm._async_success_callback.remove(event)
# Merge into the global list rather than replacing it. Only remove the
# callbacks this instance registered on a prior call, then append the
# current ones, preserving other instances' callbacks.
if litellm.callbacks is None:
litellm.callbacks = []
for cb in getattr(self, "_registered_callbacks", []):
if cb in litellm.callbacks:
litellm.callbacks.remove(cb)
for event in events:
if event not in litellm.callbacks:
litellm.callbacks.append(event)
self._registered_callbacks = list(events)
# Merge into the global list rather than replacing it. Only remove the
# callbacks this instance registered on a prior call, then append the
# current ones, preserving other instances' callbacks.
if litellm.callbacks is None:
litellm.callbacks = []
for cb in getattr(self, "_registered_callbacks", []):
if cb in litellm.callbacks:
litellm.callbacks.remove(cb)
for event in events:
if event not in litellm.callbacks:
litellm.callbacks.append(event)
self._registered_callbacks = list(events)
def _track_token_usage(self, response: Any, model: str) -> Optional[TokenMetrics]:
"""Extract and track token usage from LLM response."""
@@ -6423,8 +6169,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
metadata={
"provider": provider,
"stream": False
},
agent_id=self.current_agent_id,
}
)
return metrics
@@ -6489,118 +6234,38 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
if not usage:
return None
# Read a value under any of the given names, supporting both dict and
# object usage. The Responses API reports input_tokens/output_tokens
# instead of prompt_tokens/completion_tokens, so accept both spellings
# (mirroring _track_token_usage) to avoid returning zero counts.
def _usage_value(*names: str) -> int:
for name in names:
value = (
usage.get(name)
if isinstance(usage, dict)
else getattr(usage, name, None)
)
if value is not None:
return int(value or 0)
return 0
return TokenUsage(
prompt_tokens=_usage_value("prompt_tokens", "input_tokens"),
completion_tokens=_usage_value("completion_tokens", "output_tokens"),
total_tokens=_usage_value("total_tokens"),
cached_tokens=_usage_value("cached_tokens"),
reasoning_tokens=_usage_value("reasoning_tokens"),
audio_input_tokens=_usage_value("audio_input_tokens"),
audio_output_tokens=_usage_value("audio_output_tokens"),
)
# Extract token counts with support for both dict and object access
if isinstance(usage, dict):
return TokenUsage(
prompt_tokens=usage.get("prompt_tokens", 0),
completion_tokens=usage.get("completion_tokens", 0),
total_tokens=usage.get("total_tokens", 0),
cached_tokens=usage.get("cached_tokens", 0),
reasoning_tokens=usage.get("reasoning_tokens", 0),
audio_input_tokens=usage.get("audio_input_tokens", 0),
audio_output_tokens=usage.get("audio_output_tokens", 0),
)
else:
# Object-style access
return TokenUsage(
prompt_tokens=getattr(usage, 'prompt_tokens', 0) or 0,
completion_tokens=getattr(usage, 'completion_tokens', 0) or 0,
total_tokens=getattr(usage, 'total_tokens', 0) or 0,
cached_tokens=getattr(usage, 'cached_tokens', 0) or 0,
reasoning_tokens=getattr(usage, 'reasoning_tokens', 0) or 0,
audio_input_tokens=getattr(usage, 'audio_input_tokens', 0) or 0,
audio_output_tokens=getattr(usage, 'audio_output_tokens', 0) or 0,
)
except Exception as e:
if self.verbose:
logging.warning(f"Failed to extract token usage: {e}")
return None
@property
def current_agent_name(self) -> Optional[str]:
"""Task-local name of the agent currently driving this LLM.
Backed by a per-instance ContextVar so concurrent agents sharing one
LLM instance each read/write their own value instead of racing on a
shared attribute (issues #5052/#3933).
"""
var = getattr(self, "_current_agent_name_var", None)
if var is None:
return None
return var.get()
@current_agent_name.setter
def current_agent_name(self, agent_name: Optional[str]) -> None:
var = getattr(self, "_current_agent_name_var", None)
if var is None:
var = contextvars.ContextVar(
"current_agent_name", default=None
)
self._current_agent_name_var = var
var.set(agent_name)
@property
def current_agent_id(self) -> Optional[str]:
"""Stable per-agent identity for token aggregation (issue #5052)."""
var = getattr(self, "_current_agent_id_var", None)
if var is None:
return None
return var.get()
@current_agent_id.setter
def current_agent_id(self, agent_id: Optional[str]) -> None:
var = getattr(self, "_current_agent_id_var", None)
if var is None:
var = contextvars.ContextVar("current_agent_id", default=None)
self._current_agent_id_var = var
var.set(agent_id)
def set_current_agent(
self,
agent_name: Optional[str],
agent_id: Optional[str] = None,
) -> None:
"""Set task-local agent attribution for token tracking."""
def set_current_agent(self, agent_name: Optional[str]):
"""Set the current agent name for token tracking."""
self.current_agent_name = agent_name
self.current_agent_id = agent_id
def __deepcopy__(self, memo):
"""Deep-copy the LLM while giving the clone fresh attribution ContextVars.
``contextvars.ContextVar`` has no ``__deepcopy__``/``__reduce__`` and is
not copyable/pickleable, so ``Agent.__deepcopy__`` (which recursively
copies ``_llm_instance`` per gateway channel, issue #5052) would raise
``TypeError: cannot pickle ContextVar``. Every other attribute is copied
normally, brand-new ContextVars start the clone's attribution clean, and
thread locks on LLM subclasses (RateLimiter/FailoverManager) get fresh
locks so configured LLMs clone cleanly (issues #1746 / #5052 / #3933).
"""
cls = self.__class__
clone = cls.__new__(cls)
memo[id(self)] = clone
for key, value in self.__dict__.items():
if key == "_current_agent_name_var":
clone.__dict__[key] = contextvars.ContextVar(
"current_agent_name", default=None
)
continue
if key == "_current_agent_id_var":
clone.__dict__[key] = contextvars.ContextVar(
"current_agent_id", default=None
)
continue
if type(value) is type(threading.RLock()):
clone.__dict__[key] = threading.RLock()
continue
if type(value) is type(threading.Lock()):
clone.__dict__[key] = threading.Lock()
continue
clone.__dict__[key] = copy.deepcopy(value, memo)
return clone
def _resolve_openai_compatible_model(self) -> str:
"""Route a bare model name through the OpenAI-compatible client.
@@ -7028,9 +6693,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
# (input_text / input_image); local image paths become
# data URLs. Plain-string content is passed through.
from .openai_client import OpenAIClient
# Do not forward persisted metadata/provider extensions as
# Responses message properties.
item = {"role": role}
item = dict(msg)
item["content"] = OpenAIClient._build_responses_content(
msg.get("content", "")
)
@@ -781,9 +781,7 @@ class OpenAIClient:
"output": msg.get("content", ""),
})
else:
# Persisted history may carry local metadata/provider fields.
# Responses message items accept role/content, not that record.
item = {"role": role}
item = dict(msg)
item["content"] = self._build_responses_content(
msg.get("content", "")
)
@@ -2656,23 +2654,13 @@ class OpenAIClient:
tool_result, messages, _tool_call_id, function_name,
history_sink=deferred_history_sink,
)
results_str = json.dumps(tool_result) if tool_result else "Function returned an empty output"
except ToolExecutionError:
raise
except Exception as e:
tool_result = {"error": str(e)}
results_str = f"Error executing function: {str(e)}"
if verbose:
yield f"\n[Function error: {str(e)}]"
# Serialize the result separately from execution so a tool
# that SUCCEEDS but returns a non-JSON-serializable value
# (datetime, set, bytes, DataFrame, custom object) is not
# misreported as an execution error -- parity with the
# sync/async non-streaming loops' (TypeError, ValueError)
# fallback.
try:
results_str = json.dumps(tool_result) if tool_result else "Function returned an empty output"
except (TypeError, ValueError):
tool_result = {"result": str(tool_result)}
results_str = json.dumps(tool_result)
# Trigger callback with result
display_tool_call_fn(f"Function {function_name} returned: {results_str[:200]}{'...' if len(results_str) > 200 else ''}", console=None)
@@ -210,27 +210,7 @@ class MCPToolRunner(threading.Thread):
return f"Error: MCP tool call timed out after {self.timeout} seconds"
if not success:
return f"Error: {result}"
# A tool that fails at the application level returns a normal
# CallToolResult with isError=True per the MCP spec. Surface it as a
# structured error so retry/circuit-breaker/hook logic treats it as a
# failure instead of a successful string result (issue #5319).
if getattr(result, 'isError', False):
if hasattr(result, 'content') and result.content:
if hasattr(result.content[0], 'text'):
error_text = result.content[0].text
else:
error_text = str(result.content[0])
else:
error_text = str(result)
# Guarantee a non-empty error string: an isError result with an
# empty text block would otherwise yield {"error": ""}, which
# retry/circuit-breaker/hook logic can misread as a success
# (issue #5319).
if not error_text:
error_text = f"MCP tool '{tool_name}' reported an error (isError=True)"
return {"error": error_text, "mcp_error": True}
# Process result
if hasattr(result, 'content') and result.content:
if hasattr(result.content[0], 'text'):
@@ -421,8 +401,8 @@ class MCP:
base = re.sub(r'\.(cmd|exe|bat|ps1)$', '', base, flags=re.IGNORECASE).lower()
return base in cls._COLD_START_LAUNCHERS
def __init__(self, command_or_string=None, args=None, *, command=None, timeout=None, debug=False,
allowed_tools: Optional[List[str]] = None, disabled_tools: Optional[List[str]] = None, max_response_bytes: Optional[int] = None, **kwargs):
def __init__(self, command_or_string=None, args=None, *, command=None, timeout=None, debug=False,
allowed_tools: Optional[List[str]] = None, disabled_tools: Optional[List[str]] = None, **kwargs):
"""
Initialize the MCP connection and get tools.
@@ -442,7 +422,6 @@ class MCP:
debug: Enable debug logging for MCP operations (default: False)
allowed_tools: Include whitelist - only these tools will be available (default: None = all tools)
disabled_tools: Exclude blacklist - these tools will be filtered out (default: None = no exclusions)
max_response_bytes: Maximum body size for each Streamable HTTP response.
**kwargs: Additional parameters for StdioServerParameters
Note:
@@ -460,12 +439,6 @@ class MCP:
if command_or_string is None and command is not None:
command_or_string = command
if max_response_bytes is not None:
if isinstance(max_response_bytes, bool) or not isinstance(max_response_bytes, int) or max_response_bytes <= 0:
raise ValueError("max_response_bytes must be a positive integer")
if not isinstance(command_or_string, str) or not re.match(r'^https?://', command_or_string):
raise ValueError("max_response_bytes is only supported for Streamable HTTP MCP servers")
# Set up logging - default to WARNING level to hide INFO messages
if debug:
get_logger("mcp-wrapper").setLevel(logging.DEBUG)
@@ -537,8 +510,6 @@ class MCP:
# single source of truth (see mcp_transport.get_transport_type).
from .mcp_transport import get_transport_type
if get_transport_type(command_or_string) == "sse" and 'transport_type' not in kwargs:
if max_response_bytes is not None:
raise ValueError("max_response_bytes is only supported for Streamable HTTP MCP servers")
# Legacy SSE URL - use SSE transport for backward compatibility
from .mcp_sse import SSEMCPClient
self.sse_client = SSEMCPClient(command_or_string, debug=debug, timeout=timeout)
@@ -562,8 +533,6 @@ class MCP:
transport_options['session'] = kwargs.pop('session')
if 'resumability' in kwargs:
transport_options['resumability'] = kwargs.pop('resumability')
if max_response_bytes is not None:
transport_options['max_response_bytes'] = max_response_bytes
self.http_stream_client = HTTPStreamMCPClient(
command_or_string,
@@ -71,61 +71,6 @@ def _register_cleanup():
atexit.register(_cleanup_all_clients)
_cleanup_registered = True
def _bounded_httpx_client_factory(max_response_bytes):
import httpx
class BoundedStream(httpx.AsyncByteStream):
def __init__(self, stream):
self.stream = stream
async def __aiter__(self):
total_bytes = 0
async for chunk in self.stream:
total_bytes += len(chunk)
if total_bytes > max_response_bytes:
raise ValueError(f"MCP response exceeded maximum of {max_response_bytes} bytes")
yield chunk
async def aclose(self):
await self.stream.aclose()
class BoundedClient(httpx.AsyncClient):
async def send(self, request, *, stream=False, **kwargs):
response = await super().send(request, stream=True, **kwargs)
encoding = response.headers.get("content-encoding", "identity").lower()
if encoding not in ("", "identity"):
await response.aclose()
raise ValueError("Compressed MCP responses are not supported with a response byte limit")
content_length = response.headers.get("content-length")
try:
content_length = int(content_length) if content_length is not None else None
except ValueError:
content_length = None
if content_length is not None and content_length > max_response_bytes:
await response.aclose()
raise ValueError(f"MCP response exceeded maximum of {max_response_bytes} bytes")
response.stream = BoundedStream(response.stream)
if not stream:
try:
await response.aread()
except BaseException:
await response.aclose()
raise
await response.aclose()
return response
def create_client(headers=None, timeout=None, auth=None):
request_headers = dict(headers or {})
request_headers["Accept-Encoding"] = "identity"
return BoundedClient(
headers=request_headers,
timeout=timeout,
auth=auth,
trust_env=True,
)
return create_client
class HTTPStreamMCPTool:
"""A wrapper for an MCP tool that can be used with praisonaiagents."""
@@ -192,8 +137,6 @@ class HTTPStreamMCPTool:
logger.debug(f"Async calling tool {self.name} with args: {kwargs}")
try:
result = await self.session.call_tool(self.name, kwargs)
if getattr(result, 'isError', False):
raise RuntimeError(f"MCP tool {self.name} failed: {result.content}")
# Extract text from result
if hasattr(result, 'content') and result.content:
@@ -547,13 +490,6 @@ class HTTPStreamMCPClient:
self.debug = debug
self.timeout = timeout
self.options = options or {}
self.max_response_bytes = self.options.get('max_response_bytes')
if self.max_response_bytes is not None and (
isinstance(self.max_response_bytes, bool)
or not isinstance(self.max_response_bytes, int)
or self.max_response_bytes <= 0
):
raise ValueError("max_response_bytes must be a positive integer")
self.session = None
self.tools = []
self.transport = None
@@ -574,16 +510,17 @@ class HTTPStreamMCPClient:
def _initialize(self):
"""Initialize the connection and tools."""
# Use the global event loop
loop = get_event_loop()
self.loop_thread = None
if not loop.is_running():
def run_event_loop():
asyncio.set_event_loop(loop)
loop.run_forever()
self.loop_thread = threading.Thread(target=run_event_loop, daemon=True)
self.loop_thread.start()
# Start a background thread to run the event loop
def run_event_loop():
asyncio.set_event_loop(loop)
loop.run_forever()
self.loop_thread = threading.Thread(target=run_event_loop, daemon=True)
self.loop_thread.start()
# Run the initialization in the event loop
future = asyncio.run_coroutine_threadsafe(self._async_initialize(), loop)
self.tools = future.result(timeout=self.timeout)
@@ -605,16 +542,11 @@ class HTTPStreamMCPClient:
headers = self.options.get('headers') or None
client_options = {
"url": self.base_url,
"headers": headers,
"timeout": self.timeout,
}
if self.max_response_bytes is not None:
client_options["httpx_client_factory"] = _bounded_httpx_client_factory(
self.max_response_bytes
)
self._streams_context = streamablehttp_client(**client_options)
self._streams_context = streamablehttp_client(
url=self.base_url,
headers=headers,
timeout=self.timeout,
)
# If any step of the handshake fails after a context has been entered,
# unwind the already-entered contexts so we don't leak the httpx/anyio
@@ -757,4 +689,4 @@ class HTTPStreamMCPClient:
self._force_cleanup()
except Exception:
# Never raise exceptions in __del__
pass
pass
@@ -73,7 +73,7 @@ class SqliteMemoryAdapter:
def _get_stm_conn(self):
"""Get thread-local short-term memory connection."""
if getattr(self._local, 'stm_conn', None) is None:
if not hasattr(self._local, 'stm_conn'):
self._local.stm_conn = sqlite3.connect(
self.short_db,
check_same_thread=False,
@@ -97,7 +97,7 @@ class SqliteMemoryAdapter:
def _get_ltm_conn(self):
"""Get thread-local long-term memory connection."""
if getattr(self._local, 'ltm_conn', None) is None:
if not hasattr(self._local, 'ltm_conn'):
self._local.ltm_conn = sqlite3.connect(
self.long_db,
check_same_thread=False,
@@ -272,30 +272,4 @@ class SqliteMemoryAdapter:
conn.close()
except Exception:
pass
self._all_connections.clear()
def close_thread_connections(self):
"""Close only the calling thread's connections.
Leaves the shared adapter and other threads' connections intact so the
adapter can be reused; the calling thread lazily reopens its own
connections on next access.
A connection is only removed from the registry and cleared from
thread-local storage after it closes successfully. If ``close()`` raises,
the connection stays tracked so a later ``close_connections()`` /
``close_thread_connections()`` can retry it, and the failure is logged
rather than silently swallowed.
"""
for attr in ("stm_conn", "ltm_conn"):
conn = getattr(self._local, attr, None)
if conn is None:
continue
try:
conn.close()
except Exception as e:
logger.warning(f"Failed to close {attr}; keeping it tracked for retry: {e}")
continue
with self._connection_lock:
self._all_connections.discard(conn)
setattr(self._local, attr, None)
self._all_connections.clear()
@@ -266,12 +266,11 @@ class FileMemory:
return True
except (IOError, OSError) as e:
self._log(f"Error writing {filepath}: {e}", logging.ERROR)
return False
finally:
try:
tmp_path.unlink()
except OSError:
pass
return False
def commit_memory_batch(
self,
@@ -500,16 +500,18 @@ class Memory(SearchMixin, MemoryCoreMixin):
def _get_stm_conn(self):
"""Get thread-local short-term memory SQLite connection."""
if not hasattr(self._local, 'stm_conn') or self._local.stm_conn is None:
# Hardened factory: WAL where the filesystem supports it, safe
# DELETE fallback on NFS/SMB/FUSE/virtiofs (Issue #5264).
from ..storage.sqlite import connect as _sqlite_connect
self._local.stm_conn = _sqlite_connect(
self._local.stm_conn = sqlite3.connect(
self.short_db,
check_same_thread=False, # Allow cross-thread cleanup
timeout=30.0, # 30 second timeout for lock contention
busy_timeout_ms=30000, # 30 seconds
timeout=30.0 # 30 second timeout for lock contention
)
# Configure busy timeout for better contention handling
self._local.stm_conn.execute("PRAGMA busy_timeout=30000") # 30 seconds
# Enable WAL mode for concurrent read/write without blocking
result = self._local.stm_conn.execute("PRAGMA journal_mode=WAL").fetchone()
if result and result[0].upper() != 'WAL':
logger.warning(f"WAL mode not enabled for STM, got: {result[0]}")
self._local.stm_conn.commit()
# Register connection for cleanup
@@ -520,16 +522,18 @@ class Memory(SearchMixin, MemoryCoreMixin):
def _get_ltm_conn(self):
"""Get thread-local long-term memory SQLite connection."""
if not hasattr(self._local, 'ltm_conn') or self._local.ltm_conn is None:
# Hardened factory: WAL where the filesystem supports it, safe
# DELETE fallback on NFS/SMB/FUSE/virtiofs (Issue #5264).
from ..storage.sqlite import connect as _sqlite_connect
self._local.ltm_conn = _sqlite_connect(
self._local.ltm_conn = sqlite3.connect(
self.long_db,
check_same_thread=False, # Allow cross-thread cleanup
timeout=30.0, # 30 second timeout for lock contention
busy_timeout_ms=30000, # 30 seconds
timeout=30.0 # 30 second timeout for lock contention
)
# Configure busy timeout for better contention handling
self._local.ltm_conn.execute("PRAGMA busy_timeout=30000") # 30 seconds
# Enable WAL mode for concurrent read/write without blocking
result = self._local.ltm_conn.execute("PRAGMA journal_mode=WAL").fetchone()
if result and result[0].upper() != 'WAL':
logger.warning(f"WAL mode not enabled for LTM, got: {result[0]}")
self._local.ltm_conn.commit()
# Register connection for cleanup
@@ -2532,31 +2536,16 @@ class Memory(SearchMixin, MemoryCoreMixin):
# Other active threads should manage their own connections
# Only clear full registry in __del__ or explicit shutdown
# Close memory adapter if it exists (protocol-driven path).
#
# SQLite is thread-scoped by contract (#1562/#1563): close only the
# calling thread's adapter connections and KEEP the shared adapter so
# other threads and subsequent same-thread searches keep using the
# adapter's tables. Tearing it down here would make searches fall back
# to legacy SQL against short_mem/long_mem (#5486).
# Close memory adapter if it exists (protocol-driven path)
if hasattr(self, 'memory_adapter') and self.memory_adapter:
if isinstance(self.memory_adapter, SqliteMemoryAdapter):
try:
if hasattr(self.memory_adapter, 'close_thread_connections'):
self.memory_adapter.close_thread_connections()
logger.debug("SQLite adapter thread connections closed")
except Exception as e:
logger.warning(f"Error closing SQLite adapter thread connections: {e}")
# Intentionally preserve self.memory_adapter for reuse.
else:
try:
if hasattr(self.memory_adapter, 'close') and callable(self.memory_adapter.close):
self.memory_adapter.close()
logger.debug("Memory adapter closed successfully")
except Exception as e:
logger.warning(f"Error closing memory adapter: {e}")
finally:
self.memory_adapter = None
try:
if hasattr(self.memory_adapter, 'close') and callable(self.memory_adapter.close):
self.memory_adapter.close()
logger.debug("Memory adapter closed successfully")
except Exception as e:
logger.warning(f"Error closing memory adapter: {e}")
finally:
self.memory_adapter = None
# Close MongoDB client if it exists (legacy path)
if hasattr(self, 'mongo_client') and self.mongo_client:
@@ -113,19 +113,12 @@ class PermissionManager:
logger.warning(f"Failed to load rules: {e}")
def _save_rules(self):
"""Save rules to disk atomically.
Serialises to a temp file and ``os.replace``\\ s it into place so a
non-serialisable value (``TypeError``), an OOM kill, or a crash
mid-write can never truncate the previous good rules file. A downgraded
``deny`` rule after restart is a security regression, so the old bytes
must survive any failed save.
"""
from ..utils.atomic_io import atomic_write_json
"""Save rules to disk."""
path = self._get_rules_path()
try:
atomic_write_json(path, [r.to_dict() for r in self._rules], indent=2)
except (OSError, TypeError, ValueError) as e:
with open(path, "w") as f:
json.dump([r.to_dict() for r in self._rules], f, indent=2)
except IOError as e:
logger.error(f"Failed to save rules: {e}")
def save_rules(self):
@@ -147,8 +140,7 @@ class PermissionManager:
logger.warning(f"Failed to load approvals: {e}")
def _save_approvals(self):
"""Save approvals to disk atomically (see :meth:`_save_rules`)."""
from ..utils.atomic_io import atomic_write_json
"""Save approvals to disk."""
path = self._get_approvals_path()
try:
# Filter out expired and one-time approvals
@@ -156,8 +148,9 @@ class PermissionManager:
a for a in self._approvals
if a.is_valid() and a.scope in ("session", "always")
]
atomic_write_json(path, [a.to_dict() for a in persistent], indent=2)
except (OSError, TypeError, ValueError) as e:
with open(path, "w") as f:
json.dump([a.to_dict() for a in persistent], f, indent=2)
except IOError as e:
logger.error(f"Failed to save approvals: {e}")
def add_rule(self, rule: PermissionRule) -> str:
@@ -108,11 +108,10 @@ class PlanStorage:
# Convert plan to markdown
markdown = plan.to_markdown()
# Atomic write so a crash mid-write never truncates the previous plan.
from ..utils.atomic_io import atomic_write_text
atomic_write_text(path, markdown)
with open(path, "w", encoding="utf-8") as f:
f.write(markdown)
plan.file_path = path
logger.debug(f"Saved plan to {path}")
@@ -228,11 +227,10 @@ class PlanStorage:
"""
filename = f"{name}.json"
path = os.path.join(self.todos_dir, filename)
# Atomic write so a crash mid-write never truncates the previous todos.
from ..utils.atomic_io import atomic_write_text
atomic_write_text(path, todo.to_json())
with open(path, "w", encoding="utf-8") as f:
f.write(todo.to_json())
logger.debug(f"Saved todo list to {path}")
return path
@@ -617,74 +617,6 @@ class PluginManager:
def get_all_tools(self) -> List[Any]:
"""Get all tools from all enabled plugins."""
return [tool for _plugin_name, tool in self.get_all_tools_with_sources()]
def _enabled_of_type(self, plugin_type: "PluginType") -> List[Tuple[str, Plugin]]:
"""Snapshot enabled ``(name, plugin)`` pairs of a given declared type.
Dispatch key for the typed subsystems: a plugin only contributes a
guardrail / skill / policy when it *declares* the matching
``PluginType`` in its ``PluginInfo``. Snapshotting under the lock keeps
this safe against concurrent register/unregister in multi-agent runs.
"""
from .plugin import PluginType
with self._lock:
result: List[Tuple[str, Plugin]] = []
for name, plugin in self._plugins.items():
if not self._enabled.get(name, False):
continue
try:
declared = getattr(plugin.info, "plugin_type", PluginType.HOOK)
except Exception:
declared = PluginType.HOOK
if declared == plugin_type:
result.append((name, plugin))
return result
def get_all_guardrails(self) -> List[Any]:
"""Guardrail objects contributed by enabled ``GUARDRAIL`` plugins.
Calls ``as_guardrail()`` on each enabled plugin declared
``PluginType.GUARDRAIL`` and returns the non-``None`` guardrail objects
the Agent folds into its ``GuardrailChain`` (so they see
``validate_tool_call`` / ``validate_tool_result``, not just final text).
"""
from .plugin import PluginType
guardrails: List[Any] = []
for name, plugin in self._enabled_of_type(PluginType.GUARDRAIL):
try:
g = plugin.as_guardrail()
except Exception as e:
logger.error(f"Error getting guardrail from plugin {name}: {e}")
continue
if g is not None:
guardrails.append(g)
return guardrails
def get_all_skills(self) -> List[Any]:
"""Skill sources/dirs contributed by enabled ``SKILL`` plugins."""
from .plugin import PluginType
skills: List[Any] = []
for name, plugin in self._enabled_of_type(PluginType.SKILL):
try:
skills.extend(plugin.get_skills() or [])
except Exception as e:
logger.error(f"Error getting skills from plugin {name}: {e}")
return skills
def get_all_policies(self) -> List[Any]:
"""Policy rules contributed by enabled ``POLICY`` plugins."""
from .plugin import PluginType
policies: List[Any] = []
for name, plugin in self._enabled_of_type(PluginType.POLICY):
try:
policies.extend(plugin.get_policies() or [])
except Exception as e:
logger.error(f"Error getting policies from plugin {name}: {e}")
return policies
def shutdown(self):
"""Shutdown all plugins."""
@@ -1306,51 +1238,6 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
return HookResult.allow()
yield HookEvent.CLI_BACKEND_EXECUTE, cli_backend_execute_hook
# Generic, declaration-driven tail. The explicit branches above give
# rewrite/deny semantics to the conversation/tool/channel events that need
# them. Everything else a plugin declares in ``PluginInfo.hooks`` — the
# gateway/schedule/kanban/compaction/subagent/model-fallback lifecycle — is
# observe-only, so a single loop bridges it by method name (``event.value``)
# without a hand-written branch per event. This keeps the coverage matrix
# from drifting: any future emitted event becomes first-class the moment a
# plugin declares it and defines the matching method.
yielded = {
HookEvent.BEFORE_AGENT, HookEvent.AFTER_AGENT,
HookEvent.BEFORE_LLM, HookEvent.AFTER_LLM,
HookEvent.BEFORE_TOOL, HookEvent.AFTER_TOOL,
HookEvent.BEFORE_TOOL_DEFINITIONS,
HookEvent.MESSAGE_RECEIVED, HookEvent.MESSAGE_SENDING,
HookEvent.MESSAGE_SENT, HookEvent.MESSAGE_UNDELIVERED,
HookEvent.ON_PERMISSION_ASK, HookEvent.ON_CONFIG, HookEvent.ON_AUTH,
HookEvent.SESSION_START, HookEvent.SESSION_END, HookEvent.ON_ERROR,
HookEvent.CLI_BACKEND_EXECUTE,
}
def _make_observe_hook(method):
def observe_hook(data, _m=method):
payload = data.to_dict() if hasattr(data, "to_dict") else {}
_m(payload)
return HookResult.allow()
return observe_hook
declared = getattr(plugin.info, "hooks", None) or []
seen: set = set()
for event in declared:
if not isinstance(event, HookEvent) or event in yielded or event in seen:
continue
seen.add(event)
method = getattr(plugin, event.value, None)
if not callable(method):
continue
parent = getattr(base, event.value, None)
own = getattr(type(plugin), event.value, None)
# A base no-op that the plugin did not override yields nothing (avoids
# registering dead hooks). Methods with no base counterpart (e.g.
# ``kanban_task_created``) are treated as genuine overrides.
if parent is not None and own is parent:
continue
yield event, _make_observe_hook(method)
# Global plugin manager instance
_default_manager: Optional[PluginManager] = None
@@ -93,7 +93,6 @@ class PluginInfo:
author: str = ""
hooks: List[PluginHook] = field(default_factory=list)
dependencies: List[str] = field(default_factory=list)
plugin_type: PluginType = PluginType.HOOK
def to_dict(self) -> Dict[str, Any]:
return {
@@ -103,7 +102,6 @@ class PluginInfo:
"author": self.author,
"hooks": [h.value for h in self.hooks],
"dependencies": self.dependencies,
"plugin_type": self.plugin_type.value,
}
class Plugin(ABC):
@@ -248,42 +246,6 @@ class Plugin(ABC):
def cli_backend_execute(self, context: Dict[str, Any]) -> None:
"""Called after a CLI backend delegates a turn. Observe-only."""
pass
def gateway_start(self, context: Dict[str, Any]) -> None:
"""Called when the gateway/BotOS starts. Observe-only (platforms, bot_count)."""
pass
def gateway_stop(self, context: Dict[str, Any]) -> None:
"""Called when the gateway/BotOS stops. Observe-only (platforms, reason)."""
pass
def schedule_add(self, context: Dict[str, Any]) -> None:
"""Called when a scheduled job is persisted. Observe-only."""
pass
def schedule_remove(self, context: Dict[str, Any]) -> None:
"""Called when a scheduled job is deleted. Observe-only."""
pass
def schedule_trigger(self, context: Dict[str, Any]) -> None:
"""Called when a scheduled job fires. Observe-only."""
pass
def before_compaction(self, context: Dict[str, Any]) -> None:
"""Called before conversation memory is compacted. Observe-only."""
pass
def after_compaction(self, context: Dict[str, Any]) -> None:
"""Called after conversation memory is compacted. Observe-only."""
pass
def subagent_stop(self, context: Dict[str, Any]) -> None:
"""Called when a spawned subagent reaches a terminal state. Observe-only."""
pass
def model_fallback(self, context: Dict[str, Any]) -> None:
"""Called when the runtime falls back to a secondary model. Observe-only."""
pass
def on_permission_ask(self, target: str, reason: str) -> Optional[bool]:
"""Called when permission is requested. Return True/False to auto-approve/deny."""
@@ -307,40 +269,6 @@ class Plugin(ABC):
"""
return []
def as_guardrail(self) -> Optional[Any]:
"""Return a ``GuardrailProtocol`` object for a ``GUARDRAIL`` plugin.
Overriding this lets a plugin declared ``PluginType.GUARDRAIL``
participate in the Agent's ``GuardrailChain`` — including
``validate_tool_call`` / ``validate_tool_result`` — instead of only
rewriting the model's final text through ``after_llm``. Return an
object exposing any of ``validate_input`` / ``validate_output`` /
``validate_tool_call`` / ``validate_tool_result``, or ``None`` (default)
to opt out.
"""
return None
def get_skills(self) -> List[Any]:
"""Return skill directory paths for a ``SKILL`` plugin.
Mirrors :meth:`get_tools`. Entries are filesystem path **strings** to
skill directories (each containing a ``SKILL.md``); they are handed to
the Agent's ``SkillManager`` so an installed plugin contributes real,
discoverable skills. Non-string entries are ignored (the manager loads
by path). Empty by default (opt-in).
"""
return []
def get_policies(self) -> List[Any]:
"""Return ``Policy`` rules for a ``POLICY`` plugin.
Entries are ``praisonaiagents.policy.Policy`` objects added to the
Agent's ``PolicyEngine`` so a policy plugin enforces execution rules
rather than only answering ``on_permission_ask``. Empty by default
(opt-in).
"""
return []
class FunctionPlugin(Plugin):
"""
A simple plugin that wraps functions for specific hooks.
@@ -49,7 +49,6 @@ __all__ = [
"create_deny_tools_policy",
"create_allow_tools_policy",
"create_read_only_policy",
"create_network_policy",
]
@@ -65,7 +64,6 @@ _LAZY_IMPORTS = {
"create_deny_tools_policy": (f"{__name__}.engine", "create_deny_tools_policy"),
"create_allow_tools_policy": (f"{__name__}.engine", "create_allow_tools_policy"),
"create_read_only_policy": (f"{__name__}.engine", "create_read_only_policy"),
"create_network_policy": (f"{__name__}.engine", "create_network_policy"),
}
__getattr__ = create_lazy_getattr(_LAZY_IMPORTS, __name__)
@@ -13,18 +13,6 @@ from .config import PolicyConfig
logger = get_logger(__name__)
def _normalize_host(host: str) -> str:
"""Normalize a hostname for consistent egress matching.
DNS hostnames are case-insensitive and a single trailing dot denotes the
root; treating ``API.Example.com.`` and ``api.example.com`` as the same
host prevents policy bypass via case/trailing-dot variants.
"""
if not isinstance(host, str):
return host
return host.strip().rstrip(".").lower()
class PolicyEngine:
"""
Policy Engine for execution control.
@@ -197,34 +185,6 @@ class PolicyEngine:
{"operation": operation, "file_path": file_path}
)
def check_network(
self,
host: str,
context: Optional[Dict[str, Any]] = None
) -> PolicyResult:
"""
Check if an outbound network request to a host is allowed.
This is the deterministic egress decision every outbound request
(tool HTTP call, MCP request, provider/LLM call) can be routed
through. It mirrors ``check_tool``/``check_file`` and evaluates a
``net:<host>`` resource against the configured policies.
Args:
host: Destination host (e.g., "api.example.com")
context: Additional context (e.g., purpose, agent_id)
Returns:
PolicyResult indicating if the egress is allowed
"""
normalized = _normalize_host(host)
# Destination host stays authoritative: a caller-supplied context can
# add metadata but must not override the host seen by rule conditions.
return self.check(
f"net:{normalized}",
{**(context or {}), "host": normalized}
)
def clear(self):
"""Remove all policies."""
self._policies.clear()
@@ -347,67 +307,6 @@ def create_allow_tools_policy(
return Policy(name=name, rules=rules, priority=10)
def create_network_policy(
allow_hosts: Optional[List[str]] = None,
deny_hosts: Optional[List[str]] = None,
name: str = "network_egress"
) -> Policy:
"""
Create a network-egress policy for outbound host access.
Semantics (evaluated by ``PolicyEngine.check_network``):
- ``deny_hosts`` wins: a match on the deny-list is always denied.
- A non-empty ``allow_hosts`` means default-deny: only hosts matching
the allow-list are permitted; everything else is denied.
- If ``allow_hosts`` is empty/None, hosts not on the deny-list fall
through to the engine's default action (allow unless strict_mode).
Host patterns support fnmatch wildcards (e.g. ``*.mycorp.com``).
Args:
allow_hosts: Hosts permitted egress (non-empty => default-deny)
deny_hosts: Hosts always denied (takes precedence over allow)
name: Policy name
Returns:
Policy
"""
# Normalize patterns so they match the normalized host that
# ``check_network`` evaluates (case-insensitive, no trailing dot).
allow_hosts = [_normalize_host(h) for h in (allow_hosts or [])]
deny_hosts = [_normalize_host(h) for h in (deny_hosts or [])]
rules: List[PolicyRule] = []
# Deny-list wins: highest priority.
for pattern in deny_hosts:
rules.append(PolicyRule(
action=PolicyAction.DENY,
resource=f"net:{pattern}",
reason=f"Host denied by egress policy: {pattern}",
name=f"deny_host_{pattern}",
priority=100
))
if allow_hosts:
# Allow the listed hosts...
for pattern in allow_hosts:
rules.append(PolicyRule(
action=PolicyAction.ALLOW,
resource=f"net:{pattern}",
name=f"allow_host_{pattern}",
priority=10
))
# ...and default-deny everything else.
rules.append(PolicyRule(
action=PolicyAction.DENY,
resource="net:*",
reason="Host not in egress allow-list",
name="deny_host_default",
priority=0
))
return Policy(name=name, rules=rules, priority=100)
def create_read_only_policy(name: str = "read_only") -> Policy:
"""
Create a read-only policy that denies write operations.
@@ -844,13 +844,6 @@ Subtask: {st.name}
break
else:
logging.debug(f"No subtasks created yet for {current_task.name}")
# Mirror the sync workflow(): default a mid-workflow loop
# task (reached via next_tasks/condition, not the start task)
# with no explicit input_file to the tasks.csv convention so
# its subtasks are still expanded instead of being silently
# skipped (issue #5319).
if current_task is not start_task and not current_task.input_file:
current_task.input_file = "tasks.csv"
# Create subtasks if needed
if current_task.input_file:
self._create_loop_subtasks(current_task)
@@ -7,14 +7,6 @@ re-exports them lazily so user code can keep importing from
import importlib
from .models import ChannelMessage
from .policy import (
DEFAULT_PREFERENCE,
DetailLevel,
NotificationCategory,
NotificationDecision,
NotificationPreference,
evaluate,
)
from .protocols import PushTransportProtocol
_lazy_cache: dict = {}
@@ -41,12 +33,6 @@ def __getattr__(name: str):
__all__ = [
"ChannelMessage",
"PushTransportProtocol",
"NotificationCategory",
"DetailLevel",
"NotificationPreference",
"NotificationDecision",
"evaluate",
"DEFAULT_PREFERENCE",
"PushClient",
"WebSocketTransport",
"PollingTransport",
@@ -1,212 +0,0 @@
"""Notification policy for gateway push — pure, dependency-free core.
The push stack elsewhere is transport + delivery-guarantee only: it can get
bytes to a device reliably, but it has no notion of *whether* or *what* to
notify. This module supplies that missing policy layer as pure protocol code so
every notifier (the gateway's delivery manager, scheduler completions,
background-job alerts) can make the same decision identically:
decision = evaluate(category, prefs, now=..., is_urgent=..., content=...)
if decision.send:
deliver(recipient, decision.redacted_content)
Three concerns compose into one :class:`NotificationDecision`:
* **Categories** — a recipient opts in per event class (``approval_requested``,
``agent_finished`` …); an un-opted category is suppressed.
* **Quiet hours** — a per-recipient local window suppresses non-urgent
categories; urgent categories may override (``urgent_overrides_quiet``).
* **Detail level** — ``private`` / ``identified`` / ``detailed`` redacts how
much message content leaves the gateway in a device notification.
The heavy enforcement + durable preference store + CLI/YAML surface live in the
``praisonai``/``praisonai-bot`` wrapper; this stays pure so it carries no heavy
imports and is trivially testable.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from datetime import datetime, time
from enum import Enum
from typing import Optional, Set
__all__ = [
"NotificationCategory",
"DetailLevel",
"NotificationPreference",
"NotificationDecision",
"evaluate",
"DEFAULT_PREFERENCE",
]
class NotificationCategory(str, Enum):
"""Event classes a recipient can subscribe to independently."""
APPROVAL_REQUESTED = "approval_requested"
AGENT_FINISHED = "agent_finished"
AGENT_QUESTION = "agent_question"
HUMAN_MENTIONED = "human_mentioned"
SCHEDULED_TASK_FAILED = "scheduled_task_failed"
BACKGROUND_TASK_FAILED = "background_task_failed"
class DetailLevel(str, Enum):
"""How much content leaves the gateway in a notification.
* ``PRIVATE`` — a bare "you have activity"; no agent/session, no content.
* ``IDENTIFIED`` — which agent/session, but no message content.
* ``DETAILED`` — content included verbatim.
"""
PRIVATE = "private"
IDENTIFIED = "identified"
DETAILED = "detailed"
def _in_quiet_window(now: time, start: time, end: time) -> bool:
"""Whether ``now`` falls inside the ``[start, end)`` local window.
Handles a window that wraps past midnight (e.g. ``22:00``–``07:30``) by
treating start > end as "outside the daytime gap".
"""
if start == end:
return False
if start < end:
return start <= now < end
return now >= start or now < end
@dataclass
class NotificationPreference:
"""A recipient's durable notification policy.
Attributes:
categories: Opted-in categories; a category not present is suppressed.
quiet_hours: Optional ``(start, end)`` local time window during which
non-urgent categories are suppressed. ``None`` means no quiet hours.
timezone: IANA name used to localise ``now`` for the quiet-hours test.
detail: Redaction level applied to the notification content.
urgent_overrides_quiet: When True, an urgent event is delivered even
inside quiet hours.
"""
categories: Set[NotificationCategory] = field(default_factory=set)
quiet_hours: Optional[tuple] = None
timezone: str = "UTC"
detail: DetailLevel = DetailLevel.IDENTIFIED
urgent_overrides_quiet: bool = True
def allows(self, category: NotificationCategory) -> bool:
"""Whether this recipient has opted into ``category``."""
return category in self.categories
# A sensible default applied when a recipient has set nothing: the two
# interaction-blocking categories are on, at the identified (no-content) detail
# level, with no quiet hours.
DEFAULT_PREFERENCE = NotificationPreference(
categories={
NotificationCategory.APPROVAL_REQUESTED,
NotificationCategory.AGENT_QUESTION,
},
detail=DetailLevel.IDENTIFIED,
)
@dataclass
class NotificationDecision:
"""Outcome of :func:`evaluate` for one (category, recipient, moment).
Attributes:
send: Whether the notification should be delivered at all.
reason: Machine/operator-readable reason (e.g. ``"quiet_hours"``,
``"category_not_subscribed"``, ``"ok"``).
redacted_content: The content to actually push, redacted per the
recipient's detail level (``None`` when ``send`` is False).
"""
send: bool
reason: str
redacted_content: Optional[str] = None
def _redact(detail: DetailLevel, content: Optional[str], identity: Optional[str]) -> str:
"""Apply the detail-level redaction to ``content``.
``identity`` is an agent/session label used at the ``identified`` level so a
recipient learns *what* has activity without the transcript leaving the
gateway.
"""
if detail == DetailLevel.DETAILED:
return content if content is not None else "You have new activity"
if detail == DetailLevel.IDENTIFIED:
return f"Activity in {identity}" if identity else "You have new activity"
return "You have new activity"
def evaluate(
category: NotificationCategory,
prefs: Optional[NotificationPreference] = None,
*,
now: Optional[datetime] = None,
is_urgent: bool = False,
content: Optional[str] = None,
identity: Optional[str] = None,
) -> NotificationDecision:
"""Decide whether/what to notify for one event — pure, no side effects.
Args:
category: The event class being considered.
prefs: The recipient's preference; ``DEFAULT_PREFERENCE`` when ``None``.
now: The moment of evaluation (defaults to ``datetime.now`` localised to
the recipient's timezone). A naive ``now`` is treated as already in
the recipient's local zone.
is_urgent: Whether this specific event is urgent (may override quiet
hours when the recipient allows it).
content: The full message content, redacted per ``prefs.detail``.
identity: An agent/session label used at the ``identified`` detail level.
Returns:
A :class:`NotificationDecision`. ``send`` is False (with a ``reason``)
when the category is not subscribed or quiet hours suppress it.
"""
if prefs is None:
prefs = DEFAULT_PREFERENCE
if not prefs.allows(category):
return NotificationDecision(send=False, reason="category_not_subscribed")
if prefs.quiet_hours is not None:
local_now = _localised_time(now, prefs.timezone)
start, end = prefs.quiet_hours
if _in_quiet_window(local_now, start, end):
if not (is_urgent and prefs.urgent_overrides_quiet):
return NotificationDecision(send=False, reason="quiet_hours")
return NotificationDecision(
send=True,
reason="urgent_override" if is_urgent and prefs.quiet_hours else "ok",
redacted_content=_redact(prefs.detail, content, identity),
)
def _localised_time(now: Optional[datetime], tz_name: str) -> time:
"""Return the wall-clock ``time`` at ``now`` in ``tz_name``.
A ``None`` ``now`` uses the current time. An aware ``now`` is converted to
``tz_name``; a naive ``now`` is assumed to already be local. ``zoneinfo`` is
stdlib (3.9+) so this stays dependency-free; an unknown zone falls back to
the naive time rather than raising.
"""
if now is None:
now = datetime.now()
if now.tzinfo is not None:
try:
from zoneinfo import ZoneInfo
now = now.astimezone(ZoneInfo(tz_name))
except Exception:
pass
return now.timetz().replace(tzinfo=None) if now.tzinfo else now.time()
@@ -30,13 +30,10 @@ from .protocols import (
RoundRobinRoomPlanner,
extract_mentions,
)
from .runtime import Room, RoomConfig
__all__ = [
"RoomEvent",
"RoomTurnPlannerProtocol",
"RoundRobinRoomPlanner",
"extract_mentions",
"RoomConfig",
"Room",
]
@@ -1,244 +0,0 @@
"""Room config + a minimal, storage-agnostic runtime that *enacts* the planner.
Issue #5067. The pure turn planner (:mod:`praisonaiagents.rooms.protocols`)
decides *who* speaks next; it was previously orphaned — nothing consumed it. A
room needs one more thing to be usable end-to-end: something that appends each
message to a transcript and drains :meth:`RoomTurnPlannerProtocol.plan_next`,
running the scheduled agent each turn.
This module supplies that piece while staying faithful to the package's
"no heavy imports, no I/O" contract: the transcript is an in-memory list and the
agents are injected as a mapping of ``id -> async callable``. The *durable*
transcript store, platform delivery and CLI/YAML surface remain a wrapper
(``praisonai-bot``) concern; the wrapper can subclass or wrap :class:`Room`,
persist :attr:`Room.transcript`, and replay it after a restart — determinism of
the planner guarantees the identical schedule resumes.
"""
from __future__ import annotations
import asyncio
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Awaitable, Callable, Dict, List, Optional, Sequence
from .protocols import RoomEvent, RoomTurnPlannerProtocol, RoundRobinRoomPlanner
__all__ = ["RoomConfig", "Room"]
# An agent turn: given the rendered transcript, produce this speaker's reply.
AgentReply = Callable[[str], Awaitable[str]]
@dataclass
class RoomConfig:
"""Declarative description of a conversation room (roster + planner caps).
A room is the live, message-driven counterpart of ``AgentTeam``: a named
roster of agent ids that take bounded, ``@``-addressed turns in one shared
transcript. This is the small, pure config the wrapper binds to a
channel/thread exactly where a single ``agent`` is bound today; the planner
itself already exists in :mod:`praisonaiagents.rooms.protocols`.
Attributes:
name: Room identifier (used when binding to a channel, e.g. ``room:standup``).
roster: Ordered agent ids that take turns.
planner: Planner selector; ``"round_robin"`` -> :class:`RoundRobinRoomPlanner`.
max_rounds: Hard cap on turn-taking rounds (passed to the planner).
max_messages: Hard cap on messages per human-initiated activity.
"""
name: str
roster: List[str] = field(default_factory=list)
planner: str = "round_robin"
max_rounds: int = 3
max_messages: int = 20
def __post_init__(self) -> None:
if not self.name:
raise ValueError("RoomConfig.name must be a non-empty string")
if self.planner != "round_robin":
raise ValueError(
f"Unknown room planner {self.planner!r}; expected 'round_robin'"
)
if self.max_rounds < 1:
raise ValueError("max_rounds must be >= 1")
if self.max_messages < 1:
raise ValueError("max_messages must be >= 1")
def build_planner(self) -> RoomTurnPlannerProtocol:
"""Construct the planner this config selects."""
return RoundRobinRoomPlanner(
max_rounds=self.max_rounds, max_messages=self.max_messages
)
def to_dict(self) -> Dict[str, object]:
"""Convert to a plain dict (for YAML round-tripping / diagnostics)."""
return {
"name": self.name,
"roster": list(self.roster),
"planner": self.planner,
"max_rounds": self.max_rounds,
"max_messages": self.max_messages,
}
@classmethod
def from_dict(cls, name: str, data: Dict[str, object]) -> "RoomConfig":
"""Create from a parsed ``rooms.<name>`` YAML mapping.
Accepts either ``roster:`` or the issue's ``agents:`` key for the list of
agent ids so the declarative surface reads naturally.
"""
data = data or {}
roster = data.get("roster")
key = "roster"
if roster is None:
roster = data.get("agents", [])
key = "agents"
if roster is None:
roster = []
# A roster must be a sequence of ids. Reject a bare string (which would
# silently split into character-sized entries, e.g. "researcher" ->
# ['r','e','s',...]) and mappings, with an actionable error instead of a
# room that is silently unusable.
if isinstance(roster, (str, bytes)) or isinstance(roster, Mapping):
raise ValueError(
f"RoomConfig {name!r}: {key!r} must be a list of agent ids, "
f"got {type(roster).__name__}; "
f"use a YAML list, e.g. {key}: [researcher, critic]"
)
if not isinstance(roster, Sequence):
raise ValueError(
f"RoomConfig {name!r}: {key!r} must be a list of agent ids, "
f"got {type(roster).__name__}"
)
return cls(
name=name,
roster=[str(a) for a in roster],
planner=str(data.get("planner", "round_robin")),
max_rounds=int(data.get("max_rounds", 3)),
max_messages=int(data.get("max_messages", 20)),
)
class Room:
"""Minimal storage-agnostic runtime that enacts the turn planner.
:class:`Room` owns the (in-memory) transcript and drains the planner: after
each incoming :class:`RoomEvent` it repeatedly asks the planner for the next
speaker and runs that agent, appending the reply, until the planner settles.
It performs no I/O and holds no durable state — the wrapper supplies durable
persistence/delivery by subclassing or wrapping this and persisting
:attr:`transcript`. Because the planner is pure, replaying a persisted
transcript resumes the identical schedule.
Args:
config: The room's roster + planner caps.
agents: Mapping of agent id -> async callable ``reply(prompt) -> str``.
Typically wraps ``Agent.achat``. Only ids in ``config.roster`` are
scheduled; unknown roster ids are skipped (treated as a pass).
planner: Optional planner override (defaults to ``config.build_planner()``).
render: Optional transcript renderer -> the prompt handed to each agent.
"""
def __init__(
self,
config: RoomConfig,
agents: Dict[str, AgentReply],
planner: Optional[RoomTurnPlannerProtocol] = None,
render: Optional[Callable[[List[RoomEvent]], str]] = None,
) -> None:
self.config = config
self.agents = agents
self.planner = planner or config.build_planner()
self._render = render or self._default_render
self.transcript: List[RoomEvent] = []
# Serialize the whole append-and-drain per room: concurrent
# ``on_message`` calls must not interleave against the shared transcript,
# or two drains could schedule and execute the same agent twice.
self._lock = asyncio.Lock()
@staticmethod
def _default_render(transcript: List[RoomEvent]) -> str:
"""Render the transcript as ``speaker: content`` lines."""
return "\n".join(
f"{ev.speaker}: {ev.content}" for ev in transcript if ev.content
)
def _current_round(self) -> int:
"""The round currently being filled in the active (post-human) slice."""
roster_set = set(self.config.roster)
last_human = -1
for idx, ev in enumerate(self.transcript):
if ev.speaker not in roster_set:
last_human = idx
activity = self.transcript[last_human:] if last_human >= 0 else []
return max((ev.round for ev in activity), default=0)
def _scheduled_round(self, speaker: str) -> int:
"""Round the planner scheduled ``speaker`` into for the current activity.
Only :class:`RoundRobinRoomPlanner` defines round semantics (round 0
fan-out, later rounds admitted by ``@``-mentions). For any other custom
planner the round taxonomy is unknown, so we simply tag turns with the
activity's current round rather than inventing round-robin admission —
imposing round-robin rules on a custom planner would mislabel its turns
and could feed it an event that makes it re-select and loop unbounded.
"""
if not isinstance(self.planner, RoundRobinRoomPlanner):
return self._current_round()
from .protocols import extract_mentions
roster_set = set(self.config.roster)
last_human = -1
for idx, ev in enumerate(self.transcript):
if ev.speaker not in roster_set:
last_human = idx
activity = self.transcript[last_human:] if last_human >= 0 else []
current_round = max((ev.round for ev in activity), default=0)
max_rounds = getattr(self.planner, "max_rounds", self.config.max_rounds)
for rnd in range(current_round, max_rounds):
spoken = {ev.speaker for ev in activity if ev.round == rnd}
if rnd == 0:
admitted = set(roster_set)
else:
admitted = set()
for ev in activity:
if ev.round != rnd - 1 or ev.speaker not in roster_set:
continue
for name in extract_mentions(ev.content):
if name in roster_set:
admitted.add(name)
if speaker in admitted and speaker not in spoken:
return rnd
return current_round
async def on_message(self, event: RoomEvent) -> List[RoomEvent]:
"""Append ``event`` then drain the plan, returning the agent replies.
Each new human (non-roster) message opens a fresh round-0 activity; the
planner scopes its accounting to messages since that last human turn, so
a settled room reliably restarts when the next human speaks.
"""
async with self._lock:
self.transcript.append(event)
produced: List[RoomEvent] = []
while True:
nxt = self.planner.plan_next(self.config.roster, self.transcript)
if nxt is None:
break
rnd = self._scheduled_round(nxt)
reply = self.agents.get(nxt)
if reply is None:
# Unknown roster id: record a pass so the planner advances
# instead of re-scheduling the same missing agent forever.
turn = RoomEvent(speaker=nxt, content="", round=rnd, passed=True)
else:
content = await reply(self._render(self.transcript))
turn = RoomEvent(speaker=nxt, content=content or "", round=rnd)
self.transcript.append(turn)
produced.append(turn)
return produced
@@ -25,7 +25,6 @@ if TYPE_CHECKING:
from .store import FileScheduleStore
from .config_store import ConfigYamlScheduleStore
from .parser import parse_schedule
from .due import quota_hold_from_failure
from .runner import ScheduleRunner
from .loop import ScheduleLoop, InProcessScheduleProvider
from .protocols import (
@@ -107,11 +106,6 @@ def __getattr__(name: str):
_module_cache[name] = parse_schedule
return parse_schedule
if name == "quota_hold_from_failure":
from .due import quota_hold_from_failure
_module_cache[name] = quota_hold_from_failure
return quota_hold_from_failure
if name == "ScheduleRunner":
from .runner import ScheduleRunner
_module_cache[name] = ScheduleRunner
@@ -182,7 +176,6 @@ __all__ = [
"FileScheduleStore",
"ConfigYamlScheduleStore",
"parse_schedule",
"quota_hold_from_failure",
"ScheduleRunner",
"ScheduleLoop",
"InProcessScheduleProvider",
@@ -15,7 +15,7 @@ import time
from typing import Dict, Iterator, List, Optional
from .models import ScheduleJob, RunRecord
from .due import is_due as _is_due, is_misfire as _is_misfire, resolve_schedule_timezone
from .due import is_due as _is_due, resolve_schedule_timezone
from .hook_emit import emit_schedule_add, emit_schedule_remove
logger = get_logger(__name__)
@@ -202,10 +202,6 @@ class ConfigYamlScheduleStore:
# a later ``remove(job.id)`` by the runner is then a no-op, so a spent
# one-shot yields exactly one SCHEDULE_REMOVE.
auto_removed: List[ScheduleJob] = []
# (job_id, job_name) for occurrences suppressed by the misfire policy.
# Logged as ``missed`` history after the lock so a dropped occurrence is
# observable instead of vanishing silently.
missed: List[tuple] = []
with self._lock, self._file_lock():
# Re-read from disk so we observe cross-process claims/leases.
self._reload_locked()
@@ -228,16 +224,6 @@ class ConfigYamlScheduleStore:
self._job_state.pop(job.id, None)
changed = True
continue
# Misfire policy: an occurrence that landed while the process
# was down and is now older than ``misfire_grace_seconds`` is
# recorded as ``missed`` (observable) rather than fired stale or
# dropped silently. Advance past the slot so it is not seen due
# again on the next poll.
if _is_misfire(job, now, self._default_timezone):
job.last_run_at = now
missed.append((job.id, job.name))
changed = True
continue
# Win the claim: pre-advance + lease atomically.
job.last_run_at = now
# Count the fire here (needed so ``should_retire()`` can drop a
@@ -268,14 +254,6 @@ class ConfigYamlScheduleStore:
return []
for job in auto_removed:
emit_schedule_remove(job)
# Record suppressed occurrences so a missed run is visible in history
# rather than silently dropped (log_run takes only the thread lock).
for job_id, job_name in missed:
self.log_run(
job_id=job_id,
status="missed",
job_name=job_name,
)
return claimed
def complete(self, job_id: str, owner_id: str) -> None:
@@ -76,158 +76,6 @@ def next_fire_time(
return croniter(cron_expr, base_datetime).get_next(datetime).timestamp()
def scheduled_instant(
job: Any,
now: float,
default_timezone: str | None = None,
) -> float | None:
"""Return the epoch of the occurrence this due job is firing *for*.
When a recurring job becomes due at ``now`` it is firing for a specific
scheduled slot that may lie in the past (the process was down when the
slot arrived). This returns that slot's canonical instant so a misfire
policy can decide whether recovery is still within an acceptable grace
window. ``None`` is returned when there is no meaningful past slot (a
never-run interval job, or a kind/expression that cannot be evaluated).
- ``every``: the *first* slot that came due after ``last_run_at``
(``last_run_at + every_seconds``). Its age measures how long the job has
been overdue — the meaningful lateness for a coalesced recovery.
- ``cron``: the last cron fire at or before ``now`` (the slot being
recovered for a daily/weekly schedule).
- ``at``: the target instant itself.
"""
sched = job.schedule
if sched.kind == "every":
if sched.every_seconds is None or job.last_run_at is None:
return None
if now - job.last_run_at < sched.every_seconds:
return None
return job.last_run_at + sched.every_seconds
if sched.kind == "at":
if sched.at is None:
return None
try:
return localize_wall_clock(
datetime.fromisoformat(sched.at),
sched.tz or default_timezone,
).timestamp()
except (ValueError, TypeError):
return None
if sched.kind == "cron":
if sched.cron_expr is None:
return None
try:
from croniter import croniter # type: ignore[import-untyped]
except ImportError:
return None
try:
base_datetime = datetime.fromtimestamp(
now, resolve_schedule_timezone(sched.tz or default_timezone)
)
return croniter(sched.cron_expr, base_datetime).get_prev(datetime).timestamp()
except (ValueError, KeyError, TypeError):
return None
return None
def is_misfire(
job: Any,
now: float,
default_timezone: str | None = None,
) -> bool:
"""Whether a *due* job's occurrence should be suppressed as a misfire.
Shared by every atomic-claim store so the misfire decision is defined once.
Returns ``True`` when the occurrence the job is firing for is older than
``misfire_grace_seconds`` — a stale run whose slot passed while the process
was down. Such an occurrence is recorded as ``missed`` (observable) instead
of firing a late, no-longer-useful run. An occurrence still within the
grace window (or a first run) coalesces into a single normal fire on
recovery, the pre-existing behaviour.
With no grace set this is always ``False`` — existing jobs keep firing
exactly as before. A first run (``last_run_at is None``) is never a misfire
for any kind: there is no missed *recurrence* to suppress, only the initial
fire, which must always run (matching ``is_due``'s never-run branches).
"""
grace = getattr(job, "misfire_grace_seconds", None)
if grace is None:
return False
# A first run is the initial fire, not a recovered recurrence — always let
# it run. ``every`` already encodes this (``scheduled_instant`` returns
# ``None`` with no ``last_run_at``); make it explicit for ``cron``/``at``
# too so a never-run job past a stale slot is not marked missed unfired.
if getattr(job, "last_run_at", None) is None:
return False
# Tolerate a mis-typed grace (e.g. a hand-edited quoted YAML value that
# slipped past coercion): a bad policy value must never raise and abort the
# whole claim pass — treat it as "no misfire" so other due jobs still fire.
try:
grace_val = float(grace)
except (TypeError, ValueError):
return False
instant = scheduled_instant(job, now, default_timezone)
if instant is None:
return False
age = now - instant
if age <= 0:
return False
return age > grace_val
# Upper bound on a single quota hold (24h) so a bogus/huge provider Retry-After
# cannot park a recurring job indefinitely.
_MAX_HOLD_SECONDS = 86_400.0
def quota_hold_from_failure(
error: Any,
now: float,
slack_seconds: float = 60.0,
) -> float | None:
"""Return an epoch to park a job past a provider's rate-limit window.
Pure decision helper for the "provider-quota-aware hold": given a run's
failure it decides whether the job should stop re-firing until the
provider's own reset window elapses. Returns ``now + Retry-After + slack``
when the failure is a rate-limit / quota signal carrying a usable reset
hint, else ``None`` (no hold — a non-quota failure keeps today's behaviour).
Reuses the core error classifier so a 429/quota error is recognised by
exception type, HTTP status code, or message, and the reset window is read
from the provider's ``Retry-After`` header / ``retry_after`` attribute /
message text — no new parsing surface. The small ``slack_seconds`` avoids
re-firing the instant the window reopens (clock skew / provider rounding).
Args:
error: The failure — an ``Exception`` (preferred, so the classifier can
read the response headers) or an error string.
now: Epoch the failure was observed at.
slack_seconds: Extra seconds added past the provider's reset window.
Returns:
The epoch before which the job should be parked, or ``None`` when the
failure is not a quota signal or carries no usable reset window.
"""
from praisonaiagents.llm.error_classifier import (
ErrorCategory,
classify_error,
extract_retry_after,
)
exc = error if isinstance(error, Exception) else Exception(str(error or ""))
if classify_error(exc) != ErrorCategory.RATE_LIMIT:
return None
# Honour the provider's full reset window (unlike in-tick backoff, which caps
# at 5 minutes) so a long quota window parks the job instead of re-firing —
# but bound it to 24h so a bogus/huge value can't park a job indefinitely.
retry_after = extract_retry_after(exc, cap_seconds=_MAX_HOLD_SECONDS)
if not retry_after or retry_after <= 0:
return None
return now + float(retry_after) + max(float(slack_seconds), 0.0)
def is_due(
job: Any,
now: float,
@@ -253,17 +101,6 @@ def is_due(
run_count = getattr(job, "run_count", 0)
if max_runs is not None and run_count >= max_runs:
return False
# Provider-quota hold short-circuits before the kind check: while a job is
# parked past a provider's known-closed window (a 429/quota ``Retry-After``
# captured on a prior failure) it is not due, so it stops re-firing and
# re-failing every tick against a benched provider. Any intervening fires
# coalesce — the job becomes due again at the first legal instant once the
# hold elapses. The executor clears ``hold_until`` on the first run that
# reaches the model (proof the window reopened).
hold_until = getattr(job, "hold_until", None)
if hold_until is not None and now < hold_until:
return False
until = getattr(job, "until", None)
if until is not None:
try:
@@ -10,23 +10,6 @@ from dataclasses import dataclass, field
from typing import Any, Dict, List, Literal, Optional
def _coerce_optional_float(value: Any) -> Optional[float]:
"""Coerce a persisted/hand-edited value to ``Optional[float]``.
A hand-written config.yaml may quote a numeric field (``'3600'``) so it
arrives as a string. Left as-is it would raise ``TypeError`` in a later
numeric comparison (e.g. the misfire age check), aborting the whole claim
pass and starving every other due job. Coerce leniently: ``None`` stays
``None``; a non-numeric value is treated as unset rather than exploding.
"""
if value is None:
return None
try:
return float(value)
except (TypeError, ValueError):
return None
@dataclass
class Schedule:
"""When to run a scheduled job.
@@ -215,16 +198,12 @@ class RunRecord:
job_id: ID of the job that was executed.
job_name: Human-readable name of the job.
status: Execution status. One of ``"succeeded"``, ``"failed"``,
``"skipped"``, ``"no_change"``, or ``"missed"``. ``"no_change"``
is the stateful "monitor mode" outcome — a watched source was
``"skipped"``, or ``"no_change"``. ``"no_change"`` is the
stateful "monitor mode" outcome — a watched source was
unchanged since the last tick, so the model turn was suppressed
silently (no tokens, no delivery). It is distinct from
``"skipped"`` (a generic gate go/no-go) so an operator can tell
"nothing changed" apart from "the gate said don't run".
``"missed"`` records a recurring occurrence that landed while
the process was down and fell outside the job's
``misfire_grace_seconds`` — it was not run, but is recorded
(not silently dropped) so an operator can see the gap.
result: Agent response text (truncated if very long).
error: Error message if status is ``"failed"``.
duration: Wall-clock seconds for execution.
@@ -234,7 +213,7 @@ class RunRecord:
job_id: str
job_name: str = ""
status: Literal["succeeded", "failed", "skipped", "no_change", "missed"] = "succeeded"
status: Literal["succeeded", "failed", "skipped", "no_change"] = "succeeded"
result: Optional[str] = None
error: Optional[str] = None
duration: float = 0.0
@@ -395,28 +374,6 @@ class ScheduleJob:
``skipped`` (no tokens, no delivery) so a downstream never runs
on empty inputs. Consistent with the existing ``skipped``
outcome.
misfire_grace_seconds: Misfire policy for a recurring occurrence missed
while the process was down — the maximum age (seconds) of the
occurrence being recovered for it to still run. When the slot
the job is firing for is older than this on recovery, the
occurrence is recorded as ``missed`` (an observable
:class:`RunRecord`) rather than fired — so a daily brief whose
gateway was down for hours does not surface a stale run long
after it was useful, yet the gap is never silently dropped. An
occurrence still within the window (or a first run) coalesces
into a single normal fire on recovery. ``None`` (default)
disables the check, so existing jobs are byte-for-byte
unchanged and keep firing exactly as before.
hold_until: Optional epoch instant before which the job is parked and
not due, even when its schedule would otherwise fire. Set when
a run fails with a provider rate-limit / quota signal (the
provider's own ``Retry-After`` / reset window plus a small
slack): the recurring job stops re-firing and re-failing every
tick against a benched provider. Intervening fires coalesce
into the first legal instant after the hold; the executor
clears it on the first run that reaches the model (proof the
window reopened). ``None`` (default) keeps today's behaviour,
so jobs that never hit a quota wall are unchanged.
"""
name: str = ""
@@ -446,10 +403,8 @@ class ScheduleJob:
context_from: Optional[List[str]] = None
context_max_chars: int = 4000
on_missing_context: Literal["run", "skip"] = "run"
misfire_grace_seconds: Optional[float] = None
backend: Optional[str] = None
backend_options: Dict[str, Any] = field(default_factory=dict)
hold_until: Optional[float] = None
# ── bounded-run retirement ───────────────────────────────────────
@@ -564,10 +519,6 @@ class ScheduleJob:
d["context_max_chars"] = self.context_max_chars
if self.on_missing_context != "run":
d["on_missing_context"] = self.on_missing_context
# Misfire policy. Only persist when configured so a job with no misfire
# grace stays byte-for-byte unchanged.
if self.misfire_grace_seconds is not None:
d["misfire_grace_seconds"] = self.misfire_grace_seconds
# External CLI backend action. Only persist when configured so agent
# and command jobs stay byte-for-byte unchanged; options are opaque to
# the core (the executor validates them against the backend registry).
@@ -575,11 +526,6 @@ class ScheduleJob:
d["backend"] = self.backend
if self.backend_options:
d["backend_options"] = dict(self.backend_options)
# Provider-quota hold. Only persist when parked so a job that never hit
# a quota wall stays byte-for-byte unchanged; the parked instant must
# survive a restart so a benched job stays parked across processes.
if self.hold_until is not None:
d["hold_until"] = self.hold_until
# Atomic-claim lease metadata (set dynamically by stores that support
# ``claim_due``). Persisted so a lease is visible across processes and
# survives a restart; omitted when no lease is held.
@@ -631,16 +577,12 @@ class ScheduleJob:
context_from=context_from,
context_max_chars=d.get("context_max_chars", 4000),
on_missing_context=d.get("on_missing_context", "run"),
misfire_grace_seconds=_coerce_optional_float(
d.get("misfire_grace_seconds")
),
backend=d.get("backend"),
backend_options=(
dict(d["backend_options"])
if isinstance(d.get("backend_options"), dict)
else {}
),
hold_until=d.get("hold_until"),
)
# Restore atomic-claim lease metadata if present (see ``to_dict``).
job._lease_until = d.get("lease_until", 0.0) or 0.0
@@ -14,7 +14,7 @@ import time
from typing import Dict, Iterator, List, Optional
from .models import ScheduleJob, RunRecord
from .due import is_due as _is_due, is_misfire as _is_misfire
from .due import is_due as _is_due
from .hook_emit import emit_schedule_add, emit_schedule_remove
logger = get_logger(__name__)
@@ -187,10 +187,6 @@ class FileScheduleStore:
# a later ``remove(job.id)`` by the runner is then a no-op, so a spent
# one-shot yields exactly one SCHEDULE_REMOVE.
auto_removed: List[ScheduleJob] = []
# (job_id, job_name) for occurrences suppressed by the misfire policy.
# Logged as ``missed`` history after the lock so a dropped occurrence is
# observable instead of vanishing silently.
missed: List[tuple] = []
with self._lock, self._file_lock():
# Re-read from disk so we observe cross-process claims/leases.
self._reload_locked()
@@ -212,18 +208,6 @@ class FileScheduleStore:
auto_removed.append(self._jobs.pop(job.id))
changed = True
continue
# Misfire policy: the job is due, but the occurrence it is
# firing for may have landed while the process was down. Decide
# whether to run it now or record it as ``missed`` (observable)
# rather than firing a stale run or dropping it silently.
if _is_misfire(job, now):
# Advance the schedule past the missed slot (so the next
# poll does not see it as due again) and record the miss,
# without claiming or firing.
job.last_run_at = now
missed.append((job.id, job.name))
changed = True
continue
# Win the claim: pre-advance + lease atomically.
job.last_run_at = now
# Count the fire here (needed so ``should_retire()`` can drop a
@@ -253,15 +237,6 @@ class FileScheduleStore:
return []
for job in auto_removed:
emit_schedule_remove(job)
# Record suppressed occurrences so a missed run is visible in history
# rather than silently dropped. Logged outside the file lock (log_run
# takes only the thread lock) to keep the critical section short.
for job_id, job_name in missed:
self.log_run(
job_id=job_id,
status="missed",
job_name=job_name,
)
return claimed
def complete(self, job_id: str, owner_id: str) -> None:
@@ -39,8 +39,6 @@ __all__ = [
"register_resolver",
"register_secret_for_redaction",
"redact_secrets",
"redact_outbound",
"OutboundRedactor",
"is_secret_ref",
]
@@ -289,85 +287,3 @@ def redact_secrets(text: str) -> str:
if secret in text:
text = text.replace(secret, _REDACTED)
return text
# ────────────────────────────────────────────────────────────────────────────
# Outbound redaction — the safe-by-default reply-path primitive (Issue #5055).
#
# ``redact_secrets`` masks only values the framework has *seen* (resolved
# credentials). An unregistered secret surfaced in a tool result or a document
# would still slip through, so the outbound primitive additionally masks tokens
# by *shape* (API keys, bearer tokens, ``AKIA…``, private-key blocks). This is
# the counterpart to log redaction: it is applied to text on its way to a chat
# user, before dispatch, by the wrapper's delivery seam.
# ────────────────────────────────────────────────────────────────────────────
# Credential-shaped patterns, matched by value (not key=value), so a naked token
# echoed by the model or embedded in a tool result is caught even when never
# registered. Kept deliberately narrow to avoid over-redacting ordinary text.
_CREDENTIAL_SHAPE_PATTERNS = (
# OpenAI-style keys (sk-..., sk-proj-...), Anthropic (sk-ant-...)
r"\bsk-(?:ant-|proj-)?[A-Za-z0-9_-]{16,}\b",
# AWS access key id
r"\bAKIA[0-9A-Z]{16}\b",
# GitHub tokens (ghp_, gho_, ghu_, ghs_, ghr_, github_pat_)
r"\bgh[pousr]_[A-Za-z0-9]{16,}\b",
r"\bgithub_pat_[A-Za-z0-9_]{22,}\b",
# Google API key
r"\bAIza[0-9A-Za-z_-]{35}\b",
# Slack token
r"\bxox[baprs]-[0-9A-Za-z-]{10,}\b",
# Bearer token in an Authorization-style value
r"\bBearer\s+[A-Za-z0-9._-]{16,}\b",
# PEM private-key block
r"-----BEGIN (?:RSA |EC |OPENSSH |DSA |PGP )?PRIVATE KEY-----"
r"[\s\S]*?-----END (?:RSA |EC |OPENSSH |DSA |PGP )?PRIVATE KEY-----",
)
_credential_shape_re = None
def _get_credential_shape_re():
"""Compile the credential-shape matcher lazily (stdlib ``re`` only)."""
global _credential_shape_re
if _credential_shape_re is None:
import re
_credential_shape_re = re.compile(
"|".join(f"(?:{p})" for p in _CREDENTIAL_SHAPE_PATTERNS)
)
return _credential_shape_re
def redact_outbound(text: str) -> str:
"""Scrub secrets from text on its way to a chat user (safe-by-default).
Combines the registered-value registry (:func:`redact_secrets`, holding
every resolved gateway credential) with a credential-shape regex so an
*unregistered* token surfaced in a tool result is still masked. Personal
data is intentionally out of scope here — an opt-in PII policy layers on top
via the wrapper/plugin hook — so ordinary text is never over-redacted.
Zero dependency, never raises: a scrubber error returns the registered-value
result rather than blocking delivery.
"""
if not text or not isinstance(text, str):
return text
text = redact_secrets(text)
try:
text = _get_credential_shape_re().sub(_REDACTED, text)
except Exception: # pragma: no cover — never block delivery on a regex bug
pass
return text
@runtime_checkable
class OutboundRedactor(Protocol):
"""Overridable contract for scrubbing outbound text before dispatch.
The default implementation is :func:`redact_outbound`. A deployment may
supply its own (e.g. one that also applies a PII policy) and inject it at
the wrapper's delivery seam without importing anything heavy into core.
"""
def redact(self, text: str) -> str: ...
@@ -13,7 +13,6 @@ import uuid
import json
import time
import logging
from copy import deepcopy
from typing import Any, Dict, List, Optional, TYPE_CHECKING
logger = logging.getLogger(__name__)
@@ -503,10 +502,6 @@ class Session:
) -> bool:
"""Persist one sub-agent transcript onto this session's own record.
Built-in stores merge the entry under their session file lock. Older
custom stores retain the bounded read/write verification fallback below;
that fallback cannot guarantee atomic preservation of concurrent entries.
The store replaces the whole ``AGENT_HISTORY_KEY`` map per write, so a
naive read-then-replace would let a concurrent writer (another agent
under the same parent) clobber this agent's entry, or vice versa. We
@@ -514,14 +509,6 @@ class Session:
*other* key changed underneath us, retrying a bounded number of times.
Legacy per-agent records are always merged in (they stay on disk).
"""
merge_map = getattr(session_store, "merge_session_metadata_map", None)
if callable(merge_map):
# Built-in stores merge inside their locked read/modify/write path.
# A post-write comparison cannot detect entries lost to a stale map.
return merge_map(
self.session_id, AGENT_HISTORY_KEY, {agent_key: messages},
defaults=self._legacy_agent_histories(session_store),
)
update_meta = getattr(session_store, "update_session_metadata", None)
if not callable(update_meta):
return False
@@ -577,7 +564,6 @@ class Session:
{
"role": msg.get("role", "user"),
"content": msg.get("content", ""),
**deepcopy(msg),
}
for msg in chat_history
if isinstance(msg, dict)
@@ -39,7 +39,6 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Any
if TYPE_CHECKING:
from ..gateway.protocols import (
ConversationRequestProtocol,
GatewayRequestChannelProtocol,
OutboundMessengerProtocol,
SendPolicyProtocol,
GatewayStatusProtocol,
@@ -274,41 +273,6 @@ def clear_conversation_requester(token: Token) -> None:
_CONVERSATION_REQUESTER.set(None)
# ---------------------------------------------------------------------------
# Server→client interactive request channel (Issue #5351)
#
# The running gateway registers a concrete GatewayRequestChannelProtocol impl
# into this task-local slot so a blocked HITL turn can deliver an approval /
# choice / input request to the connected gateway client and await the
# correlated reply. When unbound (CLI / one-shot runs) the elicitation simply
# falls back to its existing path — a clean gate, never a hang.
# ---------------------------------------------------------------------------
_REQUEST_CHANNEL: ContextVar[Optional["GatewayRequestChannelProtocol"]] = ContextVar(
"praisonai_request_channel", default=None
)
def register_request_channel(
channel: Optional["GatewayRequestChannelProtocol"],
) -> Token:
"""Register the active server→client request channel for this task."""
return _REQUEST_CHANNEL.set(channel)
def get_request_channel() -> Optional["GatewayRequestChannelProtocol"]:
"""Return the active request channel, or ``None`` if no gateway is running."""
return _REQUEST_CHANNEL.get()
def clear_request_channel(token: Token) -> None:
"""Restore the previous request channel using the token from register."""
try:
_REQUEST_CHANNEL.reset(token)
except (LookupError, ValueError):
_REQUEST_CHANNEL.set(None)
# ---------------------------------------------------------------------------
# Outbound send-policy guard (Issue #2226)
#
@@ -422,9 +386,6 @@ __all__ = [
"register_conversation_requester",
"get_conversation_requester",
"clear_conversation_requester",
"register_request_channel",
"get_request_channel",
"clear_request_channel",
"register_send_policy",
"get_send_policy",
"clear_send_policy",
@@ -34,7 +34,7 @@ from typing import Any, Dict, List, Optional
from praisonaiagents._logging import get_logger
from .store import DefaultSessionStore, FileLock, SessionData
from .store import DefaultSessionStore, SessionData
logger = get_logger(__name__)
@@ -98,14 +98,9 @@ class SqliteSessionStore(DefaultSessionStore):
try:
if self.db_path != ":memory:":
os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True)
# Hardened factory: WAL where the filesystem supports it, safe
# DELETE fallback on NFS/SMB/FUSE/virtiofs (Issue #5264).
from ..storage.sqlite import connect as _sqlite_connect
# guard=True: quick_check on open + forensic backup + bounded
# least-destructive repair so a malformed index file recovers
# instead of taking the durable path down (Issue #5387).
conn = _sqlite_connect(self.db_path, isolation_level=None, guard=True)
conn = sqlite3.connect(
self.db_path, check_same_thread=False, isolation_level=None
)
self._fts_available = self._init_schema(conn)
self._conn = conn
except Exception as exc:
@@ -370,27 +365,6 @@ class SqliteSessionStore(DefaultSessionStore):
self._index_session(session)
return ok
def _save_imported_session(self, session: SessionData, *, overwrite: bool = True) -> bool:
"""Refresh recall and routing after a verbatim portable restore."""
ok = super()._save_imported_session(session, overwrite=overwrite)
if ok:
try:
filepath = self._get_session_path(session.session_id)
with FileLock(filepath, self.lock_timeout):
# Read-only cache fallback is unsafe after an overwrite:
# it can restore routing fields removed by the import.
fresh = self._load_session_from_disk(session.session_id, filepath)
self._reingest_spill(session.session_id, fresh)
with self._lock:
self._cache[session.session_id] = fresh
self._index_session(fresh)
except Exception as exc:
# The JSON write succeeded, but the old index is no longer a
# trustworthy view. Fail closed until a later refresh/rebuild.
self._deindex_session(session.session_id)
logger.debug("Post-import index refresh failed for %s: %s", session.session_id, exc)
return ok
def _modify_session_locked(self, session_id, mutator, **kwargs) -> bool:
"""Refresh the index after any locked read-modify-write.
@@ -92,24 +92,20 @@ class SqliteTranscriptStore(DefaultSessionStore):
with self._db_lock:
if self._db_ready:
return self._conn
import sqlite3 as _sqlite3 # noqa: F401 - availability guard
import sqlite3 # lazy import — stdlib, no heavy dependency
if self.db_path != ":memory:":
os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True)
# Hardened factory: WAL where the filesystem supports it, safe
# DELETE fallback on NFS/SMB/FUSE/virtiofs (Issue #5264). WAL lets
# readers proceed concurrently with a writer where available.
from ..storage.sqlite import connect as _sqlite_connect
# guard=True: quick_check on open + forensic backup + bounded
# least-destructive repair so a malformed transcript DB recovers
# instead of taking the durable path down (Issue #5387).
conn = _sqlite_connect(
self.db_path,
isolation_level=None,
busy_timeout_ms=int(self.lock_timeout * 1000),
guard=True,
conn = sqlite3.connect(
self.db_path, check_same_thread=False, isolation_level=None
)
try:
# WAL lets readers proceed concurrently with a writer.
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA synchronous=NORMAL")
conn.execute("PRAGMA busy_timeout=%d" % int(self.lock_timeout * 1000))
except Exception: # pragma: no cover - PRAGMA best-effort
pass
conn.execute(
"CREATE TABLE IF NOT EXISTS sessions ("
" session_id TEXT PRIMARY KEY,"
@@ -158,9 +154,7 @@ class SqliteTranscriptStore(DefaultSessionStore):
if not session_dir or not os.path.isdir(session_dir):
return
try:
filenames = sorted(
f for f in os.listdir(session_dir) if f.endswith(".json")
)
filenames = [f for f in os.listdir(session_dir) if f.endswith(".json")]
except OSError:
return
if not filenames:
@@ -177,7 +171,7 @@ class SqliteTranscriptStore(DefaultSessionStore):
try:
with open(filepath, "r", encoding="utf-8") as f:
data = json.load(f)
except (OSError, UnicodeDecodeError, json.JSONDecodeError, TypeError):
except (OSError, json.JSONDecodeError, TypeError):
continue
if not isinstance(data, dict):
continue
@@ -216,7 +210,7 @@ class SqliteTranscriptStore(DefaultSessionStore):
except (json.JSONDecodeError, TypeError):
return None
def _write_row(self, session: SessionData, conn=None, *, overwrite: bool = True) -> bool:
def _write_row(self, session: SessionData, conn=None) -> bool:
own_lock = conn is None
if conn is None:
conn = self._connect()
@@ -235,21 +229,17 @@ class SqliteTranscriptStore(DefaultSessionStore):
getattr(session, "updated_at", None),
)
sql = (
("INSERT OR REPLACE" if overwrite else "INSERT OR IGNORE") + " INTO sessions "
"INSERT OR REPLACE INTO sessions "
"(session_id, data, agent_name, gateway_session_id, agent_id, updated_at) "
"VALUES (?, ?, ?, ?, ?, ?)"
)
try:
if own_lock:
with self._db_lock:
cursor = conn.execute(sql, params)
conn.execute(sql, params)
else:
cursor = conn.execute(sql, params)
if not overwrite and cursor.rowcount == 0:
raise FileExistsError(session.session_id)
conn.execute(sql, params)
return True
except FileExistsError:
raise
except Exception as exc:
logger.error("Failed to write session %s: %s", session.session_id, exc)
return False
@@ -292,52 +282,6 @@ class SqliteTranscriptStore(DefaultSessionStore):
self._cache[session.session_id] = session
return True
def _save_imported_session(self, session: SessionData, *, overwrite: bool = True) -> bool:
"""Restore to the database without applying the destination's window."""
session.updated_at = datetime.now(timezone.utc).isoformat()
with self._db_lock:
# INSERT OR IGNORE makes no-clobber atomic across store instances;
# checking session_exists before a replace would still race.
if not self._write_row(session, overwrite=overwrite):
return False
with self._lock:
self._cache[session.session_id] = session
return True
def export_all(self) -> Dict[str, Any]:
"""Export durable database rows, not legacy JSON sidecar files."""
return {
"version": self.PORTABLE_VERSION,
"sessions": [data for data in self._all_rows() if isinstance(data, dict)],
}
def _collect_lineage(
self, session: SessionData, *, exclude: str
) -> List[Dict[str, Any]]:
"""Collect persisted database records from the same conversation chain."""
lineage = self._lineage_key(session.to_dict())
if not lineage:
return []
def record_lineage(raw):
# Preserve the writer's Python JSON semantics (including NaN)
# and the shared lineage precedence, without retaining other rows.
try:
data = json.loads(raw)
return self._lineage_key(data) if isinstance(data, dict) else None
except (json.JSONDecodeError, TypeError, AttributeError):
return None
conn = self._connect()
with self._db_lock:
conn.create_function("portable_session_lineage", 1, record_lineage)
rows = conn.execute(
"SELECT data FROM sessions WHERE session_id != ? "
"AND portable_session_lineage(data) = ? ORDER BY updated_at DESC",
(exclude, lineage),
).fetchall()
return [json.loads(row[0]) for row in rows]
def _modify_session_locked(
self,
session_id: str,
@@ -1622,27 +1622,6 @@ class DefaultSessionStore:
session_id, _apply, error_label="update session metadata"
)
def merge_session_metadata_map(
self, session_id: str, key: str, updates: Dict[str, Any],
*, defaults: Optional[Dict[str, Any]] = None,
) -> bool:
"""Merge nested metadata entries using the freshest map under FileLock.
Defaults fill missing entries; persisted entries win over defaults,
and the supplied updates win over both.
"""
def _apply(session: SessionData) -> None:
merged = dict(defaults or {})
current = session.metadata.get(key)
if isinstance(current, dict):
merged.update(current)
merged.update(updates)
session.metadata[key] = merged
return self._modify_session_locked(
session_id, _apply, error_label="merge session metadata map"
)
def rename_session(self, session_id: str, title: str) -> bool:
"""Give a session a human-readable title (Issue #3737).
@@ -2327,7 +2306,7 @@ class DefaultSessionStore:
out.append(data)
return out
def _save_imported_session(self, session: SessionData, *, overwrite: bool = True) -> bool:
def _save_imported_session(self, session: SessionData) -> bool:
"""Persist a restored session verbatim (no retention/window applied).
Mirrors ``_save_session`` (timestamp + atomic, file-locked write) but
@@ -2338,10 +2317,6 @@ class DefaultSessionStore:
filepath = self._get_session_path(session.session_id)
session.updated_at = datetime.now(timezone.utc).isoformat()
with FileLock(filepath, self.lock_timeout):
# The initial import check can race with a peer creating this file.
# Check again inside the same lock that protects the replacement.
if not overwrite and os.path.exists(filepath):
raise FileExistsError(filepath)
if not self._atomic_write_json(filepath, session.to_dict()):
logger.error(f"Failed to save imported session {session.session_id}")
return False
@@ -2441,7 +2416,7 @@ class DefaultSessionStore:
# Persist the imported record verbatim: an import is a restore,
# so the destination's retention/active_window must not truncate
# or compact a valid larger export before it lands on disk.
if not self._save_imported_session(session, overwrite=overwrite):
if not self._save_imported_session(session):
report.skipped.append(
{"session_id": session_id, "reason": "write failed"}
)
@@ -2449,10 +2424,6 @@ class DefaultSessionStore:
with self._lock:
self._cache[session_id] = session
report.imported += 1
except FileExistsError:
report.skipped.append(
{"session_id": session_id, "reason": "already exists (use overwrite)"}
)
except Exception as e: # pragma: no cover - defensive; one bad record
report.skipped.append(
{"session_id": session_id, "reason": f"import error: {e}"}
@@ -66,9 +66,6 @@ def generate_title(
... )
>>> # Returns something like: "Python Code Debugging Help"
"""
if max_length <= 0:
return ""
# Fallback title from user message if generation fails
fallback_title = _create_fallback_title(user_msg, max_length)
@@ -122,9 +119,6 @@ async def generate_title_async(
Returns:
Generated title string, or fallback based on user message
"""
if max_length <= 0:
return ""
fallback_title = _create_fallback_title(user_msg, max_length)
try:
@@ -158,7 +152,8 @@ Return ONLY the title text, no quotes, no explanation."""
title = title.replace('\n', ' ').replace('\r', '')
# Truncate if too long
title = _truncate_title(title, max_length)
if len(title) > max_length:
title = title[:max_length-3] + "..."
# Return if non-empty
if title and len(title.strip()) > 0:
@@ -181,17 +176,6 @@ Return ONLY the title text, no quotes, no explanation."""
return fallback_title
def _truncate_title(title: str, max_length: int) -> str:
"""Apply the character cap without using negative ellipsis slices."""
if max_length <= 0:
return ""
if len(title) <= max_length:
return title
if max_length <= 3:
return title[:max_length]
return title[:max_length - 3] + "..."
def _create_fallback_title(user_msg: str, max_length: int) -> str:
"""Create a fallback title from the user message.
@@ -203,7 +187,7 @@ def _create_fallback_title(user_msg: str, max_length: int) -> str:
Simple title based on user message
"""
if not user_msg or not user_msg.strip():
return _truncate_title("Chat Session", max_length)
return "Chat Session"
# Clean up the message
clean_msg = user_msg.strip()
@@ -216,4 +200,8 @@ def _create_fallback_title(user_msg: str, max_length: int) -> str:
if '.' in clean_msg:
clean_msg = clean_msg.split('.')[0]
return _truncate_title(clean_msg or "Chat Session", max_length)
# Truncate if too long
if len(clean_msg) > max_length:
clean_msg = clean_msg[:max_length-3] + "..."
return clean_msg if clean_msg else "Chat Session"
@@ -29,10 +29,6 @@ __all__ = [
"FileBackend",
"SQLiteBackend",
"get_backend",
# Hardened stdlib SQLite connection factory (WAL-with-safe-fallback)
"sqlite_connect",
# Structural integrity probe for durable SQLite stores (Issue #5387)
"sqlite_quick_check",
]
_LAZY_IMPORTS = {
@@ -50,8 +46,6 @@ _LAZY_IMPORTS = {
"FileBackend": ("backends", "FileBackend"),
"SQLiteBackend": ("backends", "SQLiteBackend"),
"get_backend": ("backends", "get_backend"),
"sqlite_connect": ("sqlite", "connect"),
"sqlite_quick_check": ("sqlite", "quick_check"),
}
@@ -69,17 +69,6 @@ class FileBackend:
# Sanitize key for filesystem
safe_key = "".join(c if c.isalnum() or c in "-_" else "_" for c in key)
return self.storage_dir / f"{safe_key}{self.suffix}"
def _path_to_key(self, file_path: Path) -> Optional[str]:
"""Recognize records using the full configured suffix."""
if not file_path.is_file():
return None
if self.suffix:
if file_path.name.endswith(self.suffix):
return file_path.name[:-len(self.suffix)]
return None
# Preserve the existing no-extension selection for an empty suffix.
return file_path.name if not file_path.suffix else None
def save(self, key: str, data: Dict[str, Any]) -> None:
"""Save data with the given key."""
@@ -95,11 +84,11 @@ class FileBackend:
delete=False,
suffix=".tmp"
) as f:
temp_path = f.name
if self.pretty:
json.dump(data, f, indent=2, default=str, ensure_ascii=False)
else:
json.dump(data, f, default=str, ensure_ascii=False)
temp_path = f.name
os.replace(temp_path, file_path)
except Exception as e:
@@ -121,7 +110,7 @@ class FileBackend:
try:
with open(file_path, "r", encoding="utf-8") as f:
return json.load(f)
except (json.JSONDecodeError, UnicodeDecodeError, IOError) as e:
except (json.JSONDecodeError, IOError) as e:
logger.warning(f"Failed to load {key}: {e}")
return None
@@ -144,8 +133,8 @@ class FileBackend:
keys = []
for file_path in self.storage_dir.iterdir():
key = self._path_to_key(file_path)
if key is not None:
if file_path.is_file() and file_path.suffix == self.suffix:
key = file_path.stem
if not prefix or key.startswith(prefix):
keys.append(key)
@@ -160,7 +149,7 @@ class FileBackend:
count = 0
with self._lock:
for file_path in self.storage_dir.iterdir():
if self._path_to_key(file_path) is not None:
if file_path.is_file() and file_path.suffix == self.suffix:
try:
file_path.unlink()
count += 1
@@ -202,34 +191,28 @@ class SQLiteBackend:
if not isinstance(table_name, str) or not _re.match(r'^[a-zA-Z0-9_]+$', table_name):
raise ValueError("table_name must contain only alphanumeric characters and underscores")
self.table_name = table_name
# Quote identifiers so validated names that are SQL keywords or start
# with a digit (e.g. "select", "9table") are valid in all statements.
self._quoted_table = f'"{table_name}"'
self._quoted_index = f'"idx_{table_name}_key"'
self._local = threading.local()
# Ensure directory exists
Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
if auto_create:
try:
self._create_table()
except Exception:
self.close()
raise
self._create_table()
def _get_conn(self):
"""Get thread-local connection."""
# Lazy import sqlite3
import sqlite3
if not hasattr(self._local, "conn") or self._local.conn is None:
# Hardened factory: WAL where the filesystem supports it, safe
# DELETE fallback on NFS/SMB/FUSE/virtiofs (Issue #5264).
from .sqlite import connect as _sqlite_connect
self._local.conn = _sqlite_connect(self.db_path, isolation_level=None)
self._local.conn = sqlite3.connect(
self.db_path,
check_same_thread=False
)
self._local.conn.row_factory = sqlite3.Row
# Enable WAL mode for better concurrent write safety
self._local.conn.execute("PRAGMA journal_mode=WAL")
self._local.conn.execute("PRAGMA synchronous=NORMAL")
self._local.conn.commit()
return self._local.conn
@@ -239,7 +222,7 @@ class SQLiteBackend:
cur = conn.cursor()
cur.execute(f"""
CREATE TABLE IF NOT EXISTS {self._quoted_table} (
CREATE TABLE IF NOT EXISTS {self.table_name} (
key TEXT PRIMARY KEY,
data TEXT NOT NULL,
created_at REAL DEFAULT (strftime('%s', 'now')),
@@ -249,8 +232,8 @@ class SQLiteBackend:
# Index for prefix queries
cur.execute(f"""
CREATE INDEX IF NOT EXISTS {self._quoted_index}
ON {self._quoted_table}(key)
CREATE INDEX IF NOT EXISTS idx_{self.table_name}_key
ON {self.table_name}(key)
""")
conn.commit()
@@ -263,7 +246,7 @@ class SQLiteBackend:
json_data = json.dumps(data, default=str, ensure_ascii=False)
cur.execute(f"""
INSERT INTO {self._quoted_table} (key, data, updated_at)
INSERT INTO {self.table_name} (key, data, updated_at)
VALUES (?, ?, strftime('%s', 'now'))
ON CONFLICT(key) DO UPDATE SET
data = excluded.data,
@@ -278,7 +261,7 @@ class SQLiteBackend:
cur = conn.cursor()
cur.execute(f"""
SELECT data FROM {self._quoted_table} WHERE key = ?
SELECT data FROM {self.table_name} WHERE key = ?
""", (key,))
row = cur.fetchone()
@@ -295,7 +278,7 @@ class SQLiteBackend:
cur = conn.cursor()
cur.execute(f"""
DELETE FROM {self._quoted_table} WHERE key = ?
DELETE FROM {self.table_name} WHERE key = ?
""", (key,))
deleted = cur.rowcount > 0
@@ -309,13 +292,13 @@ class SQLiteBackend:
if prefix:
cur.execute(f"""
SELECT key FROM {self._quoted_table}
SELECT key FROM {self.table_name}
WHERE key LIKE ?
ORDER BY key
""", (f"{prefix}%",))
else:
cur.execute(f"""
SELECT key FROM {self._quoted_table}
SELECT key FROM {self.table_name}
ORDER BY key
""")
@@ -327,7 +310,7 @@ class SQLiteBackend:
cur = conn.cursor()
cur.execute(f"""
SELECT 1 FROM {self._quoted_table} WHERE key = ? LIMIT 1
SELECT 1 FROM {self.table_name} WHERE key = ? LIMIT 1
""", (key,))
return cur.fetchone() is not None
@@ -337,10 +320,10 @@ class SQLiteBackend:
conn = self._get_conn()
cur = conn.cursor()
cur.execute(f"SELECT COUNT(*) as count FROM {self._quoted_table}")
cur.execute(f"SELECT COUNT(*) as count FROM {self.table_name}")
count = cur.fetchone()["count"]
cur.execute(f"DELETE FROM {self._quoted_table}")
cur.execute(f"DELETE FROM {self.table_name}")
conn.commit()
return count
@@ -262,8 +262,8 @@ class BaseJSONStore:
delete=False,
suffix=".tmp"
) as f:
temp_path = f.name
json.dump(self._data, f, indent=2, default=str, ensure_ascii=False)
temp_path = f.name
# Atomic rename
if self.use_file_lock:
@@ -343,7 +343,7 @@ class BaseJSONStore:
def list_json_sessions(
storage_dir: Path,
suffix: str = ".json",
limit: Optional[int] = 50,
limit: int = 50,
) -> List[BaseSessionInfo]:
"""
List all sessions in a directory.
@@ -353,7 +353,7 @@ def list_json_sessions(
Args:
storage_dir: Directory to search
suffix: File suffix to filter by
limit: Maximum number of sessions to return, or None for all sessions
limit: Maximum number of sessions to return
Returns:
List of BaseSessionInfo, sorted by modification time (newest first)
@@ -370,7 +370,7 @@ def list_json_sessions(
# Try to get item count from file
if suffix == ".json":
try:
with open(file_path, "r", encoding="utf-8") as f:
with open(file_path, "r") as f:
data = json.load(f)
# Common patterns for item count
if "iterations" in data:
@@ -383,7 +383,7 @@ def list_json_sessions(
pass
elif suffix == ".jsonl":
try:
with open(file_path, "r", encoding="utf-8") as f:
with open(file_path, "r") as f:
info.item_count = sum(1 for _ in f)
except Exception:
pass
@@ -424,7 +424,7 @@ def cleanup_old_sessions(
now = datetime.now()
# Get all sessions sorted by age (oldest first)
sessions = list_json_sessions(storage_dir, suffix, limit=None)
sessions = list_json_sessions(storage_dir, suffix, limit=10000)
sessions.sort(key=lambda s: s.modified_at)
# Delete old sessions
@@ -541,30 +541,11 @@ class AsyncBaseJSONStore:
else:
self._backend.save(self._storage_key, self._data)
else:
# File-based async write. Serialise to a string FIRST, write it
# to a temp file, then os.replace it into place so a crash/kill
# mid-write cannot truncate the previous good file (writing
# directly to the target truncates it before any bytes land).
# File-based async write
import aiofiles
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
payload = json.dumps(self._data, indent=2, default=str, ensure_ascii=False)
dir_path = self.storage_path.parent
fd, temp_path = tempfile.mkstemp(
dir=str(dir_path),
prefix=self.storage_path.name + ".",
suffix=".tmp",
)
os.close(fd)
try:
async with aiofiles.open(temp_path, "w", encoding="utf-8") as f:
await f.write(payload)
os.replace(temp_path, self.storage_path)
except BaseException:
try:
os.remove(temp_path)
except OSError:
pass
raise
async with aiofiles.open(self.storage_path, "w", encoding="utf-8") as f:
await f.write(json.dumps(self._data, indent=2, default=str, ensure_ascii=False))
async def exists_async(self) -> bool:
"""Check if storage exists asynchronously."""
@@ -1,557 +0,0 @@
"""Hardened stdlib-only SQLite connection helper for core durable stores.
Every SQLite-backed store in the gateway/session/memory/run subsystems used to
open its connection with a bare ``PRAGMA journal_mode=WAL`` duplicated inline.
WAL relies on a shared-memory (``-shm``) sidecar and byte-range locking that
several common *deployment substrates* do not honour: NFS/SMB home
directories, FUSE / Docker-Desktop gRPC-FUSE mounts, overlayfs, virtiofs/9p on
cross-VM setups. On those filesystems WAL either errors (``disk I/O error``,
``database is locked``, ``unable to open database file``) or — worst case on a
cross-VM share — silently corrupts.
This module provides a single, lightweight, dependency-free factory that every
store can call so the hardening lives in one place:
from praisonaiagents.storage.sqlite import connect
conn = connect(db_path)
It attempts WAL and, on a filesystem that rejects it, transparently falls back
to ``journal_mode=DELETE`` — but **never live-downgrades a database whose
on-disk header already reports WAL** (a peer connection may hold uncheckpointed
commits). A read-only header probe runs first so initialisation never unlinks
sidecars another connection is using.
Only the standard library is used (``sqlite3`` is stdlib), so this adds zero
heavy dependencies and no import-time cost — it is lazy-imported by callers.
"""
from __future__ import annotations
import os
from typing import Optional
from praisonaiagents._logging import get_logger
logger = get_logger(__name__)
# SQLite database header: bytes 18 (write) / 19 (read) hold the file-format
# version numbers; both are 2 for a WAL database and 1 for rollback-journal
# (DELETE/TRUNCATE/PERSIST). Reading them lets us detect an existing WAL header
# without opening a write connection, so we never live-downgrade a DB a peer
# may hold uncheckpointed commits in.
_WAL_FILE_FORMAT = 2
def _resolve_probe_path(path: str) -> Optional[str]:
"""Resolve a connection target to the on-disk file to probe, or None.
``:memory:`` and anonymous/temporary databases have no header to read.
``file:`` URIs (used with ``uri=True``) embed the real filename plus a query
string; parsing them back to the filesystem path keeps the no-downgrade
header probe honest instead of stat-ing the URI text as a literal filename.
"""
if not path or path == ":memory:":
return None
if path.startswith("file:"):
from urllib.parse import unquote, urlparse
parsed = urlparse(path)
# In-memory / anonymous URIs (``file::memory:``, ``file:?...``) have no
# backing file to probe.
if not parsed.path or parsed.path == ":memory:":
return None
return unquote(parsed.path)
return path
def _header_is_wal(path: str) -> Optional[bool]:
"""Return True/False if the on-disk header reports WAL, or None if unknown.
Reads only the fixed 100-byte SQLite header, so it never touches sidecars
or takes a lock. Returns None when the file is absent/empty/too short to
have a header yet (a fresh DB), so the caller treats it as "free to set".
Accepts either a plain filesystem path or a ``file:`` URI.
"""
probe = _resolve_probe_path(path)
if probe is None:
return None
try:
with open(probe, "rb") as fh:
header = fh.read(20)
except (IOError, OSError):
return None
if len(header) < 20:
return None
return header[18] == _WAL_FILE_FORMAT and header[19] == _WAL_FILE_FORMAT
def apply_wal_with_fallback(
conn, path: str, busy_timeout_ms: int = 5000, synchronous: str = "FULL"
) -> str:
"""Set a durable journal mode, falling back off WAL on hostile filesystems.
Attempts ``journal_mode=WAL``. If the filesystem rejects it (a locking or
``disk I/O`` style error, common on NFS/SMB/FUSE/virtiofs), falls back to
``journal_mode=DELETE`` — unless the on-disk header already reports WAL, in
which case the DB is left in WAL so a peer connection's uncheckpointed
commits are never stranded. Returns the journal mode actually in effect.
``synchronous`` controls the ``PRAGMA synchronous`` durability level and
defaults to ``FULL`` so committed writes survive an OS/power crash — this
preserves the pre-hardening durability of the core stores (Issue #5264).
Callers that favour throughput over the last-transaction durability window
may pass ``synchronous="NORMAL"`` (safe under WAL against process crashes).
"""
sync = (synchronous or "FULL").upper()
try:
conn.execute(f"PRAGMA busy_timeout={int(busy_timeout_ms)}")
except Exception: # pragma: no cover - busy_timeout is universally supported
pass
try:
row = conn.execute("PRAGMA journal_mode=WAL").fetchone()
mode = (row[0] if row else "").lower()
if mode == "wal":
try:
conn.execute(f"PRAGMA synchronous={sync}")
except Exception:
pass
return "wal"
except Exception as exc:
logger.info(
"WAL journal mode unavailable for %s (%s); trying DELETE fallback.",
path,
exc,
)
# WAL did not take. Never live-downgrade a DB whose header already reports
# WAL — another connection may hold uncheckpointed commits in the -wal file.
# ``_header_is_wal`` resolves ``file:`` URIs and rejects ``:memory:`` itself.
if _header_is_wal(path):
logger.warning(
"WAL requested but not applied for %s and its header already "
"reports WAL; leaving as-is to avoid stranding peer commits.",
path,
)
return "wal"
try:
row = conn.execute("PRAGMA journal_mode=DELETE").fetchone()
mode = (row[0] if row else "delete").lower()
except Exception as exc:
logger.warning("Could not set DELETE journal mode for %s: %s", path, exc)
return "unknown"
try:
# FULL is the safe default for rollback-journal mode on unreliable
# filesystems; NORMAL is only durable under WAL.
conn.execute("PRAGMA synchronous=FULL")
except Exception:
pass
return mode
# Maximum forensic backups of a malformed file kept before the oldest is
# pruned, so a boot-loop on a permanently-bad file cannot fill the disk.
_MAX_FORENSIC_BACKUPS = 3
# A malformed file is repaired at most this many times before the guard gives
# up (records the exhaustion in a sidecar ledger) and stops re-repairing on
# every boot, so a permanently-bad file is not churned each start.
_MAX_REPAIR_ATTEMPTS = 3
def _is_transient_lock_error(exc: Exception) -> bool:
"""True if an error is a transient lock/busy condition, not corruption.
A healthy database that a peer momentarily locks raises ``database is
locked`` / ``database is busy``. Treating that as corruption would trigger a
needless forensic copy and repair on a file that was never corrupt, so the
guard must distinguish it from a genuine "malformed image" error.
"""
text = str(exc).lower()
return "locked" in text or "database is busy" in text
def quick_check(conn) -> bool:
"""Return True if ``PRAGMA quick_check`` reports the database is intact.
``quick_check`` is the cheap structural probe (it skips the exhaustive
per-row index cross-check that ``integrity_check`` does), so it is safe to
run once on open without touching the hot path. Any SQLite-level error
(a malformed header, "database disk image is malformed", an I/O error) is
treated as *not intact* so the caller can route to backup/repair. A
transient lock/busy error is re-raised so callers can tell "locked" (leave
alone) apart from "corrupt" (back up + repair).
"""
try:
row = conn.execute("PRAGMA quick_check(1)").fetchone()
except Exception as exc: # DatabaseError / OperationalError on a bad file
if _is_transient_lock_error(exc):
raise
logger.warning("quick_check raised for database (%s); treating as corrupt.", exc)
return False
return bool(row) and str(row[0]).lower() == "ok"
def _sidecar_paths(path: str):
"""Return the WAL/SHM sidecar paths for a database file."""
return [path + "-wal", path + "-shm"]
def _forensic_backup(path: str) -> Optional[str]:
"""Copy a malformed database (+WAL/SHM sidecars) aside for forensics.
Mirrors the JSON store's ``_quarantine_corrupt`` behaviour: the raw,
possibly-recoverable bytes are preserved as ``<file>.corrupt-<epoch_ms>``
instead of being clobbered by a repair. Fails **closed** — returns ``None``
without copying — when free disk space is below the file size plus a margin,
so a corruption event never itself fills the disk. Retention-capped so a
boot-loop cannot accumulate unbounded copies.
"""
import shutil
import time as _time
try:
size = os.path.getsize(path)
except OSError:
return None
try:
free = shutil.disk_usage(os.path.dirname(path) or ".").free
except OSError:
free = None
# Fail closed on low disk: need room for the file plus a small margin.
if free is not None and free < size * 2 + (1 << 20):
logger.error(
"Skipping forensic backup of %s: insufficient free disk space.", path
)
return None
base = f"{path}.corrupt-{int(_time.time() * 1000)}"
dest = base
attempt = 1
while os.path.exists(dest):
dest = f"{base}-{attempt}"
attempt += 1
try:
shutil.copy2(path, dest)
except OSError as exc:
logger.error("Forensic backup of %s failed: %s", path, exc)
return None
# Best-effort sidecar preservation; their absence is not fatal. Each sidecar
# is ``path + suffix`` (e.g. ``-wal``), copied next to the primary backup so
# a later analysis can replay the write-ahead log.
for sidecar in _sidecar_paths(path):
if os.path.exists(sidecar):
suffix = sidecar[len(path):]
try:
shutil.copy2(sidecar, dest + suffix)
except OSError:
pass
_prune_forensic_backups(path)
return dest
def _prune_forensic_backups(path: str) -> None:
"""Keep only the newest ``_MAX_FORENSIC_BACKUPS`` forensic copies."""
directory = os.path.dirname(path) or "."
prefix = os.path.basename(path) + ".corrupt-"
try:
entries = [
os.path.join(directory, name)
for name in os.listdir(directory)
if name.startswith(prefix)
]
except OSError:
return
# Only prune the primary DB copies (no sidecar suffix like ``-wal``).
copies = [p for p in entries if not (p.endswith("-wal") or p.endswith("-shm"))]
if len(copies) <= _MAX_FORENSIC_BACKUPS:
return
copies.sort(key=lambda p: os.path.getmtime(p))
for stale in copies[: len(copies) - _MAX_FORENSIC_BACKUPS]:
# Remove the primary copy together with its own ``-wal``/``-shm``
# sidecars so a WAL database's forensic sidecars cannot accumulate past
# the retention cap.
for target in (stale, stale + "-wal", stale + "-shm"):
try:
os.remove(target)
except OSError:
pass
def _repair_ledger_path(path: str) -> str:
return path + ".repair"
def _repair_budget_ok(path: str) -> bool:
"""Return True while the file still has repair attempts left in its budget.
A sidecar ledger (``<file>.repair``) records how many repairs have been
attempted so a permanently-malformed file is not re-repaired on every boot
(bounded persistent retry). Missing/unreadable ledger means a fresh budget.
"""
try:
with open(_repair_ledger_path(path), "r", encoding="utf-8") as fh:
attempts = int(fh.read().strip() or "0")
except (OSError, ValueError):
attempts = 0
return attempts < _MAX_REPAIR_ATTEMPTS
def _record_repair_attempt(path: str) -> None:
try:
with open(_repair_ledger_path(path), "r", encoding="utf-8") as fh:
attempts = int(fh.read().strip() or "0")
except (OSError, ValueError):
attempts = 0
try:
with open(_repair_ledger_path(path), "w", encoding="utf-8") as fh:
fh.write(str(attempts + 1))
except OSError:
pass
def _clear_repair_ledger(path: str) -> None:
for p in (_repair_ledger_path(path),):
try:
os.remove(p)
except OSError:
pass
def _repair_via_online_backup(path: str) -> bool:
"""Attempt a least-destructive repair of a malformed database in place.
Recovers what SQLite's online-backup API can read into a fresh snapshot and
promotes it back with an atomic ``os.replace`` so the canonical file is only
swapped for a *verified-intact* rebuild — a failed attempt never touches the
original bytes (which are already preserved by :func:`_forensic_backup`).
Returns True only when the rebuilt file itself passes ``quick_check``.
"""
import sqlite3
_record_repair_attempt(path)
# A per-attempt-unique snapshot name so two concurrent repairs of the same
# file never delete or promote each other's in-flight rebuild.
snapshot = f"{path}.repair-tmp-{os.getpid()}-{id(path)}"
try:
os.remove(snapshot)
except OSError:
pass
src = dst = None
try:
src = sqlite3.connect(path)
dst = sqlite3.connect(snapshot)
src.backup(dst) # copies all readable pages into the fresh file
dst.commit()
except Exception as exc:
logger.warning("Online-backup repair of %s failed to read pages: %s", path, exc)
for c in (src, dst):
try:
if c is not None:
c.close()
except Exception:
pass
try:
os.remove(snapshot)
except OSError:
pass
return False
finally:
for c in (src, dst):
try:
if c is not None:
c.close()
except Exception:
pass
# Verify the rebuilt snapshot before promoting it over the canonical file.
verify = None
try:
verify = sqlite3.connect(snapshot)
intact = quick_check(verify)
except Exception:
intact = False
finally:
try:
if verify is not None:
verify.close()
except Exception:
pass
if not intact:
try:
os.remove(snapshot)
except OSError:
pass
return False
try:
# Drop stale sidecars of the malformed file so they cannot be replayed
# over the rebuilt database. The malformed original + its sidecars are
# already forensically backed up before we get here.
for sidecar in _sidecar_paths(path):
try:
os.remove(sidecar)
except OSError:
pass
os.replace(snapshot, path)
except OSError as exc:
logger.error("Could not promote repaired snapshot for %s: %s", path, exc)
return False
_clear_repair_ledger(path)
logger.warning("Repaired malformed SQLite database %s via online backup.", path)
return True
def _quarantine_unrepairable(path: str) -> None:
"""Move an unrepairable malformed file aside so a fresh DB can be created.
When the online-backup API cannot rebuild an intact snapshot (e.g. a
corrupt header that reads as "file is not a database"), the malformed bytes
are already preserved by :func:`_forensic_backup`; renaming the canonical
file out of the way lets the store open a fresh, empty database instead of
crashing on the first ``CREATE TABLE`` — mirroring the JSON store's
quarantine-and-continue behaviour (Issue #5387). Sidecars go too so a stale
``-wal`` cannot be replayed over the new file.
"""
for sidecar in _sidecar_paths(path):
try:
os.remove(sidecar)
except OSError:
pass
try:
os.remove(path)
except OSError as exc:
logger.error("Could not clear unrepairable malformed file %s: %s", path, exc)
def _guard_integrity(path: str, *, repair: bool, backup: bool) -> None:
"""Detect and (optionally) repair a malformed database before first use.
Runs a cheap ``quick_check`` on open; on failure it forensically backs up
the malformed file (+sidecars) and attempts a bounded, least-destructive
online-backup repair. If the file cannot be repaired, its (already backed
up) bytes are quarantined so the store can open a fresh database rather than
crash on first use. A no-op for ``:memory:``/absent/empty files. A transient
lock/busy error is treated as "healthy but busy" and left untouched.
"""
import sqlite3
if not path or path == ":memory:" or not os.path.exists(path):
return
try:
if os.path.getsize(path) == 0:
return
except OSError:
return
probe = None
try:
probe = sqlite3.connect(path)
if quick_check(probe):
return
except Exception as exc:
# A transient lock is not corruption — never back up or repair a file
# that a peer merely holds open. Leave it exactly as-is.
if _is_transient_lock_error(exc):
logger.info(
"SQLite database %s is locked/busy on probe; skipping guard.", path
)
return
finally:
try:
if probe is not None:
probe.close()
except Exception:
pass
logger.error("SQLite database %s failed integrity check; corrupt.", path)
# Fail closed: if a forensic backup was requested but could not be taken
# (low disk, copy error), do NOT proceed to repair — repair's atomic
# promotion would otherwise destroy the only copy of the malformed bytes.
if backup:
if _forensic_backup(path) is None:
logger.error(
"Forensic backup of %s failed; skipping repair to preserve "
"the malformed file for recovery.",
path,
)
return
if repair and _repair_budget_ok(path):
if _repair_via_online_backup(path):
return
# Repair could not rebuild an intact snapshot. The malformed bytes are
# safely backed up, so quarantine the canonical file to let the durable
# store open a fresh database instead of crashing on first query.
_quarantine_unrepairable(path)
def connect(
path,
*,
check_same_thread: bool = False,
isolation_level: Optional[str] = "",
busy_timeout_ms: int = 5000,
synchronous: str = "FULL",
guard: bool = False,
repair: bool = True,
backup: bool = True,
**kwargs,
):
"""Open a hardened SQLite connection used by core durable stores.
Applies :func:`apply_wal_with_fallback` so the connection uses WAL where the
filesystem supports it and DELETE where it does not, with a sensible
``busy_timeout``. All keyword arguments other than the journal handling are
forwarded to :func:`sqlite3.connect`.
Args:
path: Database path (``":memory:"`` supported).
check_same_thread: Forwarded to ``sqlite3.connect`` (defaults False so
stores may share a connection across threads under their own lock).
isolation_level: Forwarded to ``sqlite3.connect``. Defaults to ``""``
(sqlite3's default autocommit-off); pass ``None`` for autocommit.
busy_timeout_ms: ``PRAGMA busy_timeout`` in milliseconds.
synchronous: ``PRAGMA synchronous`` level. Defaults to ``FULL`` so
committed writes survive an OS/power crash (matches pre-hardening
durability). Pass ``"NORMAL"`` for higher throughput under WAL.
guard: When True, run a one-time ``quick_check`` on open and, on a
malformed file, forensically back it up and attempt a bounded
least-destructive online-backup repair before connecting (Issue
#5387). Off by default so existing callers are unchanged.
repair: Attempt an online-backup repair on corruption (only when
``guard`` is set). Bounded by a sidecar attempt-ledger.
backup: Take a retention-capped, disk-space-aware forensic backup of the
malformed file before repair (only when ``guard`` is set).
**kwargs: Forwarded to ``sqlite3.connect``.
"""
import sqlite3 # lazy import — stdlib, no heavy dependency
db_path = path
if isinstance(path, (str, os.PathLike)) and str(path) != ":memory:":
db_path = str(path)
if guard:
try:
_guard_integrity(str(db_path), repair=repair, backup=backup)
except Exception as exc: # detection/repair must never block open
logger.debug("Integrity guard skipped for %s: %s", db_path, exc)
connect_kwargs = dict(kwargs)
connect_kwargs["check_same_thread"] = check_same_thread
if isolation_level != "":
connect_kwargs["isolation_level"] = isolation_level
conn = sqlite3.connect(db_path, **connect_kwargs)
try:
apply_wal_with_fallback(
conn,
str(db_path),
busy_timeout_ms=busy_timeout_ms,
synchronous=synchronous,
)
except Exception as exc: # never let hardening break connection creation
logger.debug("Journal-mode hardening skipped for %s: %s", db_path, exc)
return conn
@@ -889,16 +889,8 @@ Context:
return GuardrailResult(success=True, result=task_output)
try:
# Call the guardrail function, awaiting it if the caller supplied an
# async def. Calling a coroutine function synchronously would only
# build a coroutine object (never running its body), which then
# crashes GuardrailResult.from_tuple and is misreported as a
# validation failure. This mirrors how Task.callback already detects
# and awaits coroutine callbacks.
if asyncio.iscoroutinefunction(self._guardrail_fn):
result = self._run_async_guardrail(task_output)
else:
result = self._guardrail_fn(task_output)
# Call the guardrail function
result = self._guardrail_fn(task_output)
# Check if result is already a GuardrailResult
if isinstance(result, GuardrailResult):
@@ -922,18 +914,7 @@ Context:
result=None,
error=f"Guardrail validation error: {str(e)}"
)
def _run_async_guardrail(self, task_output):
"""Run an async def guardrail from the synchronous guardrail path.
The async orchestration path already offloads _process_guardrail to a
worker thread (no running loop there), so asyncio.run is safe. If called
while a loop is running in this thread, surface a clear error instead of
the misattributed 'coroutine was never awaited' failure.
"""
from ..utils.async_bridge import run_coroutine_from_any_context
return run_coroutine_from_any_context(self._guardrail_fn(task_output))
async def _execute_callback_with_metadata(self, task_output):
"""Execute callback with metadata support while maintaining backward compatibility.
@@ -63,28 +63,11 @@ class SessionTokenMetrics:
"""Aggregated token metrics for a session."""
total_interactions: int = 0
metrics_by_model: Dict[str, TokenMetrics] = field(default_factory=dict)
# Aggregated per agent *identity* (a stable per-instance key such as an
# Agent's ``_approval_scope_id``), not the bare display name. Two unrelated
# agents that happen to share a name (e.g. the default ``"Agent"``) must not
# merge into one cost bucket, which would make per-tenant accounting
# unrecoverable. Each entry carries the display name for reporting.
metrics_by_agent_id: Dict[str, Dict] = field(default_factory=dict)
metrics_by_agent: Dict[str, TokenMetrics] = field(default_factory=dict)
total_metrics: TokenMetrics = field(default_factory=TokenMetrics)
def add_interaction(
self,
model: str,
agent: Optional[str],
metrics: TokenMetrics,
agent_id: Optional[str] = None,
):
"""Add a new interaction's metrics.
``agent`` is the display name; ``agent_id`` is a stable per-instance
identity used as the aggregation key so name collisions between
unrelated agents never merge their usage. When ``agent_id`` is omitted
the display name is used as the key (backward-compatible fallback).
"""
def add_interaction(self, model: str, agent: Optional[str], metrics: TokenMetrics):
"""Add a new interaction's metrics."""
self.total_interactions += 1
# Update total metrics
@@ -95,33 +78,14 @@ class SessionTokenMetrics:
self.metrics_by_model[model] = TokenMetrics()
self.metrics_by_model[model] = self.metrics_by_model[model] + metrics
# Update agent-specific metrics, keyed by identity (not display name).
key = agent_id or agent
if key:
bucket = self.metrics_by_agent_id.get(key)
if bucket is None:
bucket = {"name": agent, "metrics": TokenMetrics()}
self.metrics_by_agent_id[key] = bucket
bucket["metrics"] = bucket["metrics"] + metrics
# Update agent-specific metrics
if agent:
if agent not in self.metrics_by_agent:
self.metrics_by_agent[agent] = TokenMetrics()
self.metrics_by_agent[agent] = self.metrics_by_agent[agent] + metrics
def get_summary(self) -> Dict:
"""Get a summary of session token usage.
``by_agent`` is keyed by display name for readability (metrics for
same-named agents are summed only for display); ``by_agent_id`` keeps
the collision-free per-identity breakdown that callers use to scope
usage to a specific instance/tenant.
"""
by_agent: Dict[str, TokenMetrics] = {}
by_agent_id: Dict[str, Dict] = {}
for key, bucket in self.metrics_by_agent_id.items():
name = bucket.get("name") or key
agent_metrics = bucket["metrics"]
if name in by_agent:
by_agent[name] = by_agent[name] + agent_metrics
else:
by_agent[name] = agent_metrics
by_agent_id[key] = {"name": name, "metrics": agent_metrics.to_dict()}
"""Get a summary of session token usage."""
return {
"total_interactions": self.total_interactions,
"total_tokens": self.total_metrics.total_tokens,
@@ -131,10 +95,9 @@ class SessionTokenMetrics:
for model, metrics in self.metrics_by_model.items()
},
"by_agent": {
name: metrics.to_dict()
for name, metrics in by_agent.items()
},
"by_agent_id": by_agent_id,
agent: metrics.to_dict()
for agent, metrics in self.metrics_by_agent.items()
}
}
@@ -165,25 +128,18 @@ class TokenCollector:
model: str,
agent: Optional[str],
metrics: TokenMetrics,
metadata: Optional[Dict] = None,
agent_id: Optional[str] = None,
metadata: Optional[Dict] = None
):
"""Track token usage for an interaction.
``agent`` is the display name; ``agent_id`` is an optional stable
per-instance identity used as the aggregation key so unrelated agents
that share a name never merge into one cost bucket.
"""
"""Track token usage for an interaction."""
with self._lock:
# Add to session metrics
self._session_metrics.add_interaction(model, agent, metrics, agent_id=agent_id)
self._session_metrics.add_interaction(model, agent, metrics)
# Track recent interaction
interaction = {
"timestamp": datetime.now().isoformat(),
"model": model,
"agent": agent,
"agent_id": agent_id,
"metrics": metrics.to_dict(),
"metadata": metadata or {}
}
@@ -114,27 +114,6 @@ from praisonaiagents.tools import internet_search
results = internet_search("AI news")
```
### Parallel Search MCP
Install `praisonaiagents[parallel-search]` to use Parallel's free Search MCP
without an API key. Select it for one call with `providers="parallel"`, or set
`WEB_SEARCH_PROVIDER=parallel` for the existing provider-selection path:
```python
from praisonaiagents.tools.web_search import search_web
results = search_web("recent agent framework releases", providers="parallel")
```
Parallel is opt-in and does not change the automatic provider order. When you
select it, your search query is sent to Parallel Search MCP for processing.
Requests include a `praisonaiagents/<installed-version>` User-Agent so Parallel
can measure aggregate project usage. This identifier is shared across
installations and contains no user or installation ID. Anonymous access is
intended for exploration and light use; see the
[Search MCP documentation](https://docs.parallel.ai/integrations/mcp/search-mcp)
for access limits.
### 2. Class-Based Approach (Complex Tools)
Best for tools that do multiple related things or need to remember information. Like a smart calculator that remembers your previous calculations and can do many different math operations.
@@ -261,6 +261,10 @@ class CircuitBreaker:
self._stats.state = CircuitState.HALF_OPEN
self._stats.success_count = 0
# Start health check if enabled and not already running
if self.config.enable_health_check and not self._health_check_task:
self._start_health_check()
# Execute function
try:
if asyncio.iscoroutinefunction(func):
@@ -319,91 +323,45 @@ class CircuitBreaker:
if (self._stats.state in (CircuitState.CLOSED, CircuitState.HALF_OPEN) and
len(self._failure_times) >= self.config.failure_threshold):
self._stats.state = CircuitState.OPEN
just_opened = True
else:
just_opened = False
# Kick off health-check-driven recovery the moment the circuit opens, so
# both call() and acall() get it (was previously only wired into acall).
# Guarded so it only runs with a live loop and no probe already running.
if just_opened and self.config.enable_health_check:
self._start_health_check()
def _start_health_check(self) -> None:
"""Start periodic health check task.
Only starts when the circuit is actually OPEN and no probe is already
running, and clears ``_health_check_task`` when the loop exits so a
future OPEN can restart probing (the previous version pinned a completed
Task forever, permanently disabling recovery). Requires a running loop;
pure-sync callers with no loop simply fall back to ``recovery_timeout``.
"""
"""Start periodic health check task."""
if not self._health_check:
return
async def health_check_loop():
try:
while self._stats.state == CircuitState.OPEN:
try:
await asyncio.sleep(self.config.health_check_interval)
while self._stats.state == CircuitState.OPEN:
try:
await asyncio.sleep(self.config.health_check_interval)
# Perform health check
is_healthy = False
if hasattr(self._health_check, 'ahealth_check'):
is_healthy = await self._health_check.ahealth_check()
elif asyncio.iscoroutinefunction(self._health_check):
is_healthy = await self._health_check()
else:
is_healthy = self._health_check()
# If healthy, transition to half-open for recovery probe
if is_healthy:
with self._lock:
if self._stats.state == CircuitState.OPEN:
self._stats.state = CircuitState.HALF_OPEN
self._stats.success_count = 0
# Perform health check
is_healthy = False
if hasattr(self._health_check, 'ahealth_check'):
is_healthy = await self._health_check.ahealth_check()
elif asyncio.iscoroutinefunction(self._health_check):
is_healthy = await self._health_check()
else:
# Run a synchronous callback off the event loop so a
# blocking probe (network/disk I/O) never stalls
# unrelated coroutines for the check's duration.
loop = asyncio.get_running_loop()
is_healthy = await loop.run_in_executor(
None, self._health_check
)
# If healthy, transition to half-open for recovery probe
if is_healthy:
with self._lock:
if self._stats.state == CircuitState.OPEN:
self._stats.state = CircuitState.HALF_OPEN
self._stats.success_count = 0
except asyncio.CancelledError:
break
except Exception:
# Ignore health check errors
pass
finally:
# Clear our own reference so a future OPEN can restart probing.
# Only clear if the stored task is still us (a concurrent
# _start_health_check may have already installed a replacement).
current = asyncio.current_task()
restart = False
with self._lock:
if self._health_check_task is current:
self._health_check_task = None
# A concurrent failure may have re-OPENed the circuit
# after our loop-condition check but before this clear.
# Without restarting, that OPEN episode would be left
# with no probe (the start guard had rejected a
# replacement while our stale task was still referenced).
restart = (
self._stats.state == CircuitState.OPEN
and self.config.enable_health_check
)
if restart:
self._start_health_check()
except asyncio.CancelledError:
break
except Exception:
# Ignore health check errors
pass
with self._lock:
if self._stats.state != CircuitState.OPEN or self._health_check_task is not None:
return
try:
loop = asyncio.get_running_loop()
self._health_check_task = loop.create_task(health_check_loop())
except RuntimeError:
# No event loop running, skip health check
pass
try:
loop = asyncio.get_event_loop()
self._health_check_task = loop.create_task(health_check_loop())
except RuntimeError:
# No event loop running, skip health check
pass
class CircuitBreakerRegistry:
@@ -50,39 +50,6 @@ _NO_GATEWAY_ASK_MSG = (
_MEDIA_RE = re.compile(r"MEDIA:(\S+)")
# Closed, core-owned message-action vocabulary with a per-action authorisation
# contract (Issue #5054). Every verb the ``send_message`` tool accepts MUST
# appear here declaring its *authority* — how much conversation it may touch:
#
# * ``read-only`` — inspects reachable targets, sends nothing.
# * ``current-conversation`` — writes to the conversation the turn is on
# (or an operator-permitted target).
# * ``own-prior-message`` — mutates a specific earlier message the agent sent
# (edit/delete); the most sensitive class, since a
# steered agent could otherwise rewrite history.
#
# The map is the single source of truth: a verb missing from it is rejected
# fail-closed (see ``_authority_for``), so a new action can never silently
# inherit ``send``'s guard. Adding a verb *requires* declaring its authority.
MESSAGE_ACTION_AUTHORITY = {
"list": "read-only",
"send": "current-conversation",
"react": "current-conversation",
"unreact": "current-conversation",
"thread": "current-conversation",
"edit": "own-prior-message",
"delete": "own-prior-message",
}
def _authority_for(action: str) -> Optional[str]:
"""Return the declared authority for ``action``, or ``None`` if unknown.
``None`` means the verb is not in the closed vocabulary and must be rejected
fail-closed rather than dispatched.
"""
return MESSAGE_ACTION_AUTHORITY.get(action)
# Default and hard upper bound for ``ask_conversation``'s bounded wait. The
# target is model-controlled, so a steered/prompt-injected agent could pass a
# negative, non-finite (NaN/inf), or absurdly large timeout. Any such value is
@@ -266,16 +233,12 @@ def send_message(
target: str = "origin",
message: str = "",
action: str = "send",
message_id: str = "",
) -> str:
"""Manage the conversation through the active gateway (Issue #5054).
"""Proactively message the user through the active gateway.
One coherent, authorised action surface: reach the user mid-task, react,
open a thread, or edit/retract a message this agent sent — dispatched to
whichever channel the current turn is on. Requires a running bot/gateway;
it is unavailable for plain CLI/one-shot runs. Every action is drawn from a
closed vocabulary (:data:`MESSAGE_ACTION_AUTHORITY`); an unknown verb is
rejected fail-closed rather than dispatched.
Use this to reach the user mid-task on the channel this conversation came
from, or on another channel they have configured. Requires a running
bot/gateway; it is unavailable for plain CLI/one-shot runs.
Args:
target: Symbolic destination. One of:
@@ -283,31 +246,21 @@ def send_message(
- "<platform>": that platform's home channel (e.g. "telegram")
- "<platform>:<chat_id>[:<thread_id>]": an explicit chat
- "<alias>": a friendly alias for a known target
message: The text to send (for "send"), the emoji (for
"react"/"unreact"), the thread title (for "thread"), or the new
text (for "edit"). Append " MEDIA:<path>" to a send to attach a
local file, e.g. "Report ready MEDIA:/tmp/report.pdf".
action: The verb to perform:
- "send" — deliver a message (default)
- "list" — return the targets currently reachable
- "react" — add an emoji reaction (pass the emoji as ``message``)
- "unreact" — remove an emoji reaction
- "thread" — open a thread/topic titled ``message`` under ``target``
- "edit" — replace ``message_id`` with ``message`` (your own message)
- "delete" — delete/unsend ``message_id`` (your own message)
Capability-gated verbs (react/thread/edit/delete) return a typed
``unsupported`` outcome on a channel that cannot render them, rather
message: The text to send. Append " MEDIA:<path>" to attach a local
file, e.g. "Report ready MEDIA:/tmp/report.pdf".
action: "send" to deliver a message (default), "list" to return the
targets currently reachable so you can pick a destination, "react"
to add an emoji reaction to the message being handled (pass the
emoji as ``message``), or "unreact" to remove one. Reactions are a
lightweight acknowledgement for busy group channels; on a channel
that cannot react you get a typed ``unsupported`` outcome rather
than an error.
message_id: The id of the message to react to / edit / delete. For
"react"/"unreact" it may be omitted when ``target`` is "origin"
(the message being handled is used).
Returns:
For action="send": a short human-readable summary of the delivery.
For action="list": a JSON array of reachable targets.
For action="react"/"unreact"/"thread"/"edit"/"delete": a JSON object
with a typed ``status`` (``ok`` / ``unsupported`` / ``failed`` /
``no_route``).
For action="react"/"unreact": a JSON object with a typed ``status``
(``ok`` / ``unsupported`` / ``failed`` / ``no_route``).
"""
try:
from ..session.context import get_outbound_messenger
@@ -316,64 +269,10 @@ def send_message(
if messenger is None:
return _NO_GATEWAY_MSG
# Fail-closed vocabulary gate (Issue #5054): reject any verb that has
# not declared its authority, before any dispatch or policy work.
authority = _authority_for(action)
if authority is None:
return (
f"Unknown action '{action}'. Use 'send', 'list', 'react', "
"'unreact', 'thread', 'edit', or 'delete'."
)
if action == "list":
targets = messenger.list_targets()
return json.dumps([t.as_dict() for t in targets])
if action == "thread":
denied = _check_send_policy(target)
if denied is not None:
return json.dumps({"status": "failed", "detail": denied})
name = message.strip()
if not name:
return json.dumps(
{"status": "failed", "detail": "no thread title provided"}
)
create_thread = getattr(messenger, "create_thread", None)
if create_thread is None:
return json.dumps(
{
"status": "unsupported",
"detail": "the active gateway does not support threads",
}
)
result = _run_async(create_thread(target, name))
return json.dumps(result.as_dict())
if action in ("edit", "delete"):
denied = _check_send_policy(target)
if denied is not None:
return json.dumps({"status": "failed", "detail": denied})
mid = message_id.strip()
if not mid:
return json.dumps(
{"status": "failed", "detail": "no message_id provided"}
)
fn = getattr(messenger, action, None)
if fn is None:
return json.dumps(
{
"status": "unsupported",
"detail": (
f"the active gateway does not support {action}"
),
}
)
if action == "edit":
result = _run_async(fn(target, mid, message))
else:
result = _run_async(fn(target, mid))
return json.dumps(result.as_dict())
if action in ("react", "unreact"):
# Reuse the same operator send-policy guard as ``send`` so a
# steered/prompt-injected agent cannot react on a channel the
@@ -398,15 +297,16 @@ def send_message(
)
result = _run_async(
react(
target,
emoji,
message_id=message_id.strip(),
remove=(action == "unreact"),
)
react(target, emoji, remove=(action == "unreact"))
)
return json.dumps(result.as_dict())
if action != "send":
return (
f"Unknown action '{action}'. Use 'send', 'list', 'react', "
"or 'unreact'."
)
# Outbound send-policy guard (Issue #2226): the target is
# model-controlled, so a steered/prompt-injected agent could misdeliver
# to any reachable channel. Enforce the operator's allow/deny policy
@@ -487,4 +387,4 @@ def ask_conversation(
return json.dumps({"status": "undelivered", "detail": str(e)})
__all__ = ["send_message", "ask_conversation", "MESSAGE_ACTION_AUTHORITY"]
__all__ = ["send_message", "ask_conversation"]
@@ -455,7 +455,6 @@ def execute_code_with_tools(
timeout: int = 30,
max_output_size: int = 10000,
bridge: Optional[Any] = None,
policy_hook: Optional[Any] = None,
) -> Dict[str, Any]:
"""Execute model-generated code that may call the agent's registered tools.
@@ -507,27 +506,18 @@ def execute_code_with_tools(
# (allow-list, registry) and limits are forwarded explicitly so the
# isolated path is gated by exactly what this caller authorised — never
# the transport's own defaults, and never a weaker path.
run_kwargs: Dict[str, Any] = dict(
return bridge.run_code(
code,
allowed_tools=allowed,
registry=registry,
timeout=timeout,
max_output_size=max_output_size,
)
# Only forward ``policy_hook`` when set so third-party CodeToolBridge
# transports whose ``run_code`` predates this keyword remain compatible
# (the shipped LocalProcessBridge accepts and honours it).
if policy_hook is not None:
run_kwargs["policy_hook"] = policy_hook
return bridge.run_code(code, **run_kwargs)
injected: Dict[str, Any] = {}
if allowed:
injected.update(
build_tool_namespace(allowed, registry=registry, policy_hook=policy_hook)
)
injected["tools"] = ToolProxy(
allowed, registry=registry, policy_hook=policy_hook
)
injected.update(build_tool_namespace(allowed, registry=registry))
injected["tools"] = ToolProxy(allowed, registry=registry)
return _execute_code_direct(
code,
@@ -989,8 +979,7 @@ class PythonTools:
# ──────────────────────────────────────────────────────────────────────
def _build_unsafe_execute_code(
allowed_tools: List[str], timeout: int = 30, registry: Optional[Any] = None,
policy_hook: Optional[Any] = None,
allowed_tools: List[str], timeout: int = 30, registry: Optional[Any] = None
):
"""Build the in-process ``execute_code`` tool for ``code_mode="unsafe"``.
@@ -998,10 +987,6 @@ def _build_unsafe_execute_code(
a registry built from its own granted tools so code-mode cannot reach a
globally-registered tool (e.g. a plugin entry-point tool) the agent was
never given; omitting it falls back to the global registry.
``policy_hook`` is the agent's PolicyEngine / permission-deny / per-tool
guardrail gate, applied to every tool call the code makes so code-mode is
gated identically to the direct tool-call loop (not by approval alone).
"""
@require_approval(risk_level="critical")
@@ -1018,16 +1003,14 @@ def _build_unsafe_execute_code(
Dict with ``result``, ``stdout``, ``stderr`` and ``success``.
"""
return execute_code_with_tools(
code, allowed_tools=allowed_tools, timeout=timeout, registry=registry,
policy_hook=policy_hook,
code, allowed_tools=allowed_tools, timeout=timeout, registry=registry
)
return execute_code
def _build_isolated_execute_code(
allowed_tools: List[str], timeout: int = 30, registry: Optional[Any] = None,
policy_hook: Optional[Any] = None,
allowed_tools: List[str], timeout: int = 30, registry: Optional[Any] = None
):
"""Build the ``execute_code`` tool for ``code_mode="isolated"``.
@@ -1062,7 +1045,6 @@ def _build_isolated_execute_code(
timeout=timeout,
registry=registry,
bridge=LocalProcessBridge(),
policy_hook=policy_hook,
)
return execute_code
@@ -1073,7 +1055,6 @@ def build_code_execution_tools(
allowed_tools: Optional[List[str]] = None,
timeout: int = 30,
registry: Optional[Any] = None,
policy_hook: Optional[Any] = None,
) -> List[Any]:
"""Return the code-execution tools for ``ExecutionConfig(code_execution=True)``.
@@ -1120,15 +1101,13 @@ def build_code_execution_tools(
if code_mode == "unsafe":
return [
_build_unsafe_execute_code(
list(allowed_tools or []), timeout=timeout, registry=registry,
policy_hook=policy_hook,
list(allowed_tools or []), timeout=timeout, registry=registry
)
]
if code_mode == "isolated":
return [
_build_isolated_execute_code(
list(allowed_tools or []), timeout=timeout, registry=registry,
policy_hook=policy_hook,
list(allowed_tools or []), timeout=timeout, registry=registry
)
]
@@ -57,7 +57,6 @@ class CodeToolBridge(Protocol):
registry: Optional[ToolRegistry] = None,
timeout: int = 30,
max_output_size: int = 10000,
policy_hook: Optional[Callable[..., Any]] = None,
) -> Dict[str, Any]:
"""Run *code* under isolation, servicing tool calls over the transport.
@@ -80,7 +79,6 @@ def serve_tool_call(
kwargs: Dict[str, Any],
allowed: Iterable[str],
registry: Optional[ToolRegistry] = None,
policy_hook: Optional[Callable[..., Any]] = None,
) -> Any:
"""Parent-side handler for a single bridged tool call.
@@ -88,13 +86,6 @@ def serve_tool_call(
receives from isolated code. It reuses the exact allow-list resolution and
``require_approval`` gate as the in-process proxy, so an isolated call is
subject to the same policy as an in-process one — never a weaker path.
``policy_hook`` (when supplied) is the owning agent's PolicyEngine /
permission-deny / per-tool-guardrail gate. It is consulted with the tool's
keyword arguments before the tool runs; it may deny the call and/or rewrite
the arguments (see :func:`_run_policy_hook` for the accepted return shapes),
so a code-mode call is gated identically to the normal tool-call loop (which
the approval gate alone does not cover).
"""
allowed_set = frozenset(allowed)
# ``registry is not None`` (not truthiness): an empty ToolRegistry is a
@@ -108,9 +99,7 @@ def serve_tool_call(
tool = resolved_registry.get(name)
if tool is None:
raise NameError(f"tool '{name}' is not registered")
return _invoke_with_approval(
name, tool, tuple(args), dict(kwargs), policy_hook=policy_hook
)
return _invoke_with_approval(name, tool, tuple(args), dict(kwargs))
def _resolve_callable(tool: Any) -> Callable[..., Any]:
@@ -139,70 +128,8 @@ def _call_and_resolve(callable_tool: Callable[..., Any], args: tuple, kwargs: Di
return result
def _bind_kwargs(callable_tool: Callable[..., Any], args: tuple, kwargs: Dict[str, Any]) -> Dict[str, Any]:
"""Best-effort mapping of ``(args, kwargs)`` to a name->value dict.
Positional arguments are bound to their parameter names so a policy/guardrail
hook sees the same keyword-argument view the normal tool-call loop passes to
``_check_tool_policy_and_guardrails``. Falls back to ``argN`` keys when the
signature is unavailable/unbindable.
"""
bound_args: Dict[str, Any] = dict(kwargs)
try:
import inspect
signature = inspect.signature(callable_tool)
bound = signature.bind_partial(*args, **kwargs)
bound_args = dict(bound.arguments)
except (TypeError, ValueError):
for index, value in enumerate(args):
bound_args.setdefault(f"arg{index}", value)
return bound_args
def _run_policy_hook(
policy_hook: Callable[..., Any],
name: str,
arguments: Dict[str, Any],
) -> Dict[str, Any]:
"""Invoke ``policy_hook`` and normalise its return, raising on denial.
The hook mirrors the direct tool-call path: it may deny the call and/or
rewrite the arguments (a sanitising input-guardrail). Two return shapes are
accepted for backward compatibility:
* ``(denial_reason, rewritten_args)`` — the current agent contract. A
non-``None`` ``denial_reason`` blocks; ``rewritten_args`` (when a dict) is
the argument view the tool must run with, so a guardrail rewrite is
honoured rather than discarded.
* a bare ``denial_reason`` string / ``None`` — older/third-party hooks. A
string blocks; ``None`` allows with the arguments unchanged.
Returns the (possibly rewritten) argument dict to dispatch with. Raises
``PermissionError`` when the hook denies the call.
"""
verdict = policy_hook(name, arguments)
denial_reason: Optional[str]
rewritten: Dict[str, Any] = arguments
if isinstance(verdict, tuple):
denial_reason = verdict[0]
if len(verdict) > 1 and isinstance(verdict[1], dict):
rewritten = verdict[1]
else:
denial_reason = verdict
if denial_reason is not None:
raise PermissionError(
f"Execution of {name} denied by policy: {denial_reason}"
)
return rewritten
def _invoke_with_approval(
name: str,
tool: Any,
args: tuple,
kwargs: Dict[str, Any],
policy_hook: Optional[Callable[..., Any]] = None,
name: str, tool: Any, args: tuple, kwargs: Dict[str, Any]
) -> Any:
"""Run a registered tool honouring the existing approval framework.
@@ -213,16 +140,6 @@ def _invoke_with_approval(
Positional arguments are bound to their parameter names before being sent to
the approval backend so the human approver sees every argument value and can
rewrite any of them via ``decision.modified_args``.
``policy_hook`` gates the call through the owning agent's PolicyEngine,
permission deny-rules, and per-tool ``input_guardrails`` — the exact checks
the normal tool-call loop runs and which the approval gate alone does not
cover. It runs before the tool executes and raises ``PermissionError`` on
denial so a code-mode call cannot silently bypass a guardrail the same tool
is subject to via the direct tool-call path. A guardrail's argument rewrite
is honoured (the returned args are dispatched), and — mirroring the direct
path — the hook is re-run after the approval backend modifies the arguments
so an approval rewrite cannot smuggle a call past an argument-scoped policy.
"""
from ..approval import (
is_approval_required,
@@ -233,25 +150,6 @@ def _invoke_with_approval(
callable_tool = _resolve_callable(tool)
# Bind positional args to parameter names once so the policy hook and the
# approval backend share the same keyword view the direct path uses. When a
# hook rewrites the args we dispatch by keyword from that view (positional
# ``args`` is then dropped to avoid a double-bind of the same parameter).
effective_kwargs = _bind_kwargs(callable_tool, args, kwargs)
use_kwargs_only = False
if policy_hook is not None:
rewritten = _run_policy_hook(policy_hook, name, dict(effective_kwargs))
# Only switch to keyword-only redispatch when the hook *actually*
# changed the arguments. Compare by value (not identity): the hook is
# handed a copy, so an allow-with-no-rewrite still returns a distinct
# object whose contents equal ``effective_kwargs``. Forcing keyword-only
# dispatch on an unchanged call would rebind a positional argument to a
# decorated tool's ``*args``/``**kwargs`` parameter names and break it.
if rewritten != effective_kwargs:
effective_kwargs = dict(rewritten)
use_kwargs_only = True
# Note: unlike the regular agent path we deliberately do NOT honour the
# per-session ``is_already_approved`` sticky flag, nor do we call
# ``mark_approved`` after approval. In code mode the model controls the whole
@@ -275,38 +173,41 @@ def _invoke_with_approval(
f"async context. Configure a non-console approval backend."
)
# Bind positional args to parameter names so the approval backend sees
# (and can modify) every argument, not just the keyword ones.
approval_args: Dict[str, Any] = dict(kwargs)
bound = None
try:
import inspect
signature = inspect.signature(callable_tool)
bound = signature.bind_partial(*args, **kwargs)
approval_args = dict(bound.arguments)
except (TypeError, ValueError):
# Signature unavailable/unbindable: fall back to positional preview.
for index, value in enumerate(args):
approval_args.setdefault(f"arg{index}", value)
decision = run_coroutine_from_any_context(
request_approval(name, dict(effective_kwargs))
request_approval(name, approval_args)
)
if not decision.approved:
raise PermissionError(
f"Execution of {name} denied: {decision.reason}"
)
if decision.modified_args:
effective_kwargs.update(decision.modified_args)
use_kwargs_only = True
# Re-authorise: an approval backend may have rewritten the args into
# a call an argument-scoped policy/guardrail prohibits. The earlier
# gate only saw the original args, so re-run the hook on the final
# ones before dispatch — exactly as the direct path does.
if policy_hook is not None:
effective_kwargs = dict(
_run_policy_hook(policy_hook, name, dict(effective_kwargs))
)
approval_args.update(decision.modified_args)
if bound is not None:
try:
rebound = signature.bind_partial(**approval_args)
return _call_and_resolve(callable_tool, rebound.args, rebound.kwargs)
except TypeError:
pass
# Fall back: apply only keyword modifications.
kwargs.update(
{k: v for k, v in decision.modified_args.items() if k in kwargs}
)
if use_kwargs_only:
# The args were rewritten (by a guardrail and/or the approval backend).
# Prefer re-binding through the signature so a rewritten VAR_POSITIONAL /
# VAR_KEYWORD parameter maps back to real positional/keyword args; fall
# back to a plain keyword call when the signature is unavailable.
try:
import inspect
signature = inspect.signature(callable_tool)
rebound = signature.bind_partial(**effective_kwargs)
return _call_and_resolve(callable_tool, rebound.args, rebound.kwargs)
except (TypeError, ValueError):
return _call_and_resolve(callable_tool, (), effective_kwargs)
return _call_and_resolve(callable_tool, args, kwargs)
@@ -314,7 +215,6 @@ def _make_proxy(
name: str,
allowed: frozenset,
registry: ToolRegistry,
policy_hook: Optional[Callable[..., Any]] = None,
) -> Callable[..., Any]:
"""Build a single tool proxy enforcing the allow-list and approval gate."""
if name not in allowed:
@@ -324,7 +224,7 @@ def _make_proxy(
raise NameError(f"tool '{name}' is not registered")
def _proxy(*args: Any, **kwargs: Any) -> Any:
return _invoke_with_approval(name, tool, args, kwargs, policy_hook=policy_hook)
return _invoke_with_approval(name, tool, args, kwargs)
_proxy.__name__ = name
return _proxy
@@ -353,7 +253,6 @@ class ToolProxy:
self,
allowed: Iterable[str],
registry: Optional[ToolRegistry] = None,
policy_hook: Optional[Callable[..., Any]] = None,
) -> None:
allowed_set = frozenset(allowed)
# ``is not None`` so an explicitly-passed empty (agent-scoped) registry
@@ -361,7 +260,7 @@ class ToolProxy:
resolved_registry = registry if registry is not None else get_registry()
def _getter(name: str) -> Callable[..., Any]:
return _make_proxy(name, allowed_set, resolved_registry, policy_hook=policy_hook)
return _make_proxy(name, allowed_set, resolved_registry)
# Stash all state inside a closure; never as an instance attribute.
object.__setattr__(self, "_ToolProxy__getter", _getter)
@@ -427,7 +326,6 @@ class LocalProcessBridge:
registry: Optional[ToolRegistry] = None,
timeout: int = 30,
max_output_size: int = 10000,
policy_hook: Optional[Callable[..., Any]] = None,
) -> Dict[str, Any]:
import os
import subprocess
@@ -516,7 +414,7 @@ class LocalProcessBridge:
# same way as the in-process one — do not wrap this in a broad
# try/except that would swallow them.
return _pump_bridge(
process, allowed, registry, timeout, max_output_size, policy_hook
process, allowed, registry, timeout, max_output_size
)
finally:
if temp_file:
@@ -526,7 +424,7 @@ class LocalProcessBridge:
pass
def _pump_bridge(process, allowed, registry, timeout, max_output_size, policy_hook=None):
def _pump_bridge(process, allowed, registry, timeout, max_output_size):
"""Drive the child: service tool-call requests, collect the result.
Runs in the parent. Tool-call request lines are serviced by
@@ -553,7 +451,7 @@ def _pump_bridge(process, allowed, registry, timeout, max_output_size, policy_ho
if line.startswith(_BRIDGE_CALL):
payload = json.loads(line[len(_BRIDGE_CALL):])
response = _service_bridge_call(
payload, allowed, registry, policy_error, policy_hook
payload, allowed, registry, policy_error
)
try:
encoded = json.dumps(response)
@@ -650,15 +548,14 @@ def _pump_bridge(process, allowed, registry, timeout, max_output_size, policy_ho
return outcome
def _service_bridge_call(payload, allowed, registry, policy_error, policy_hook=None):
def _service_bridge_call(payload, allowed, registry, policy_error):
"""Run one bridged tool call in the parent under the shared gate."""
name = payload.get("name")
args = payload.get("args", [])
kwargs = payload.get("kwargs", {})
try:
value = serve_tool_call(
name, args, kwargs, allowed=allowed, registry=registry,
policy_hook=policy_hook,
name, args, kwargs, allowed=allowed, registry=registry
)
return {"ok": True, "value": value}
except (PermissionError, NameError) as exc:
@@ -788,17 +685,12 @@ _main()
def build_tool_namespace(
allowed: Iterable[str],
registry: Optional[ToolRegistry] = None,
policy_hook: Optional[Callable[..., Any]] = None,
) -> Dict[str, Callable[..., Any]]:
"""Build a dict of ``{tool_name: proxy_callable}`` for the allow-list.
The returned mapping can be merged directly into the executor globals so
the model can call ``fetch(...)`` by bare name, in addition to the
``tools.fetch(...)`` form via :class:`ToolProxy`.
``policy_hook`` is forwarded to each proxy so bare-name calls are gated by
the owning agent's PolicyEngine / permission-deny / per-tool guardrails,
identically to the normal tool-call loop.
"""
allowed_set = frozenset(allowed)
# ``is not None`` so an explicitly-passed empty (agent-scoped) registry is
@@ -807,9 +699,7 @@ def build_tool_namespace(
namespace: Dict[str, Callable[..., Any]] = {}
for name in sorted(allowed_set):
try:
namespace[name] = _make_proxy(
name, allowed_set, resolved_registry, policy_hook=policy_hook
)
namespace[name] = _make_proxy(name, allowed_set, resolved_registry)
except NameError:
# Tool on the allow-list but not registered yet — skip silently;
# calling it would raise NameError at runtime anyway.
@@ -3,7 +3,7 @@
This module provides a single `search_web` function that automatically tries
multiple search providers in order, falling back to the next if one fails.
Default Search Provider Priority:
Search Provider Priority:
1. Tavily (requires TAVILY_API_KEY + tavily-python)
2. Brave (requires BRAVE_API_KEY + requests)
3. Exa (requires EXA_API_KEY + exa_py)
@@ -11,8 +11,6 @@ Default Search Provider Priority:
5. DuckDuckGo (requires ddgs package, no API key)
6. SearxNG (requires requests + running SearxNG instance)
Parallel Search MCP is available when explicitly selected.
Usage:
from praisonaiagents.tools import search_web
@@ -26,19 +24,10 @@ Usage:
from typing import List, Dict, Any, Optional
from praisonaiagents._logging import get_logger
import os
import json
from importlib import util
logger = get_logger(__name__)
PARALLEL_SEARCH_MCP_URL = "https://search.parallel.ai/mcp"
PARALLEL_SEARCH_TIMEOUT_SECONDS = 30
PARALLEL_MAX_RESPONSE_BYTES = 1_000_000
PARALLEL_MAX_RESULTS = 10
PARALLEL_MAX_TITLE_CHARS = 256
PARALLEL_MAX_URL_CHARS = 2048
PARALLEL_MAX_SNIPPET_CHARS = 1200
def _check_tavily() -> tuple[bool, Optional[str]]:
"""Check if Tavily is available."""
if util.find_spec("tavily") is None:
@@ -83,94 +72,6 @@ def _check_searxng() -> tuple[bool, Optional[str]]:
return False, "requests package not installed"
return True, None
def _check_parallel() -> tuple[bool, Optional[str]]:
"""Check if compatible Streamable HTTP MCP support is installed."""
try:
from praisonaiagents.mcp.mcp_http_stream import streamablehttp_client
except ImportError:
streamablehttp_client = None
if streamablehttp_client is None:
return False, "install praisonaiagents[parallel-search] for Streamable HTTP MCP support"
from inspect import signature
if "httpx_client_factory" not in signature(streamablehttp_client).parameters:
return False, "install praisonaiagents[parallel-search] for bounded Streamable HTTP support"
return True, None
def _truncate_parallel_text(value: Any, limit: int) -> str:
text = str(value or "").strip()
if len(text) > limit:
return text[:limit - 1] + "…"
return text
def _search_parallel(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
"""Search with the keyless Parallel Search MCP endpoint."""
result_limit = max(0, min(int(max_results), PARALLEL_MAX_RESULTS))
if result_limit == 0:
return []
from importlib.metadata import PackageNotFoundError, version
from praisonaiagents.mcp import MCP
# Keep this project-level so attribution aggregates across installations.
# Fall back gracefully when dist metadata is unavailable (source checkouts,
# vendored/frozen installs) instead of crashing the whole search path.
try:
package_version = version("praisonaiagents")
except PackageNotFoundError:
package_version = "unknown"
headers = {"User-Agent": f"praisonaiagents/{package_version}"}
with MCP(
PARALLEL_SEARCH_MCP_URL,
timeout=PARALLEL_SEARCH_TIMEOUT_SECONDS,
allowed_tools=["web_search"],
headers=headers,
max_response_bytes=PARALLEL_MAX_RESPONSE_BYTES,
) as client:
search_tool = next(
(tool for tool in client.get_tools() if getattr(tool, "__name__", None) == "web_search"),
None,
)
if search_tool is None:
raise RuntimeError("Parallel Search MCP did not expose web_search")
response = search_tool(objective=query, search_queries=[query])
if not isinstance(response, str):
raise ValueError("Parallel Search MCP returned a non-text response")
if response.startswith("Error:"):
raise RuntimeError(response)
payload = json.loads(response)
if not isinstance(payload, dict) or "results" not in payload:
raise ValueError("Parallel Search MCP returned an invalid response shape")
raw_results = payload["results"]
if not isinstance(raw_results, list):
raise ValueError("Parallel Search MCP returned an invalid results list")
results = []
for result in raw_results:
if not isinstance(result, dict):
continue
url = result.get("url")
if not isinstance(url, str) or not url or len(url) > PARALLEL_MAX_URL_CHARS:
continue
excerpts = result.get("excerpts", [])
if isinstance(excerpts, str):
excerpts = [excerpts]
if not isinstance(excerpts, list):
excerpts = []
snippet = " ".join(
excerpt.strip() for excerpt in excerpts
if isinstance(excerpt, str) and excerpt.strip()
)
results.append({
"title": _truncate_parallel_text(result.get("title", ""), PARALLEL_MAX_TITLE_CHARS),
"url": url,
"snippet": _truncate_parallel_text(snippet, PARALLEL_MAX_SNIPPET_CHARS),
"provider": "parallel",
})
if len(results) >= result_limit:
break
return results
def _search_tavily(query: str, max_results: int = 5, search_depth: str = "basic", raw_content: bool = False) -> List[Dict[str, Any]]:
"""Search using Tavily."""
from tavily import TavilyClient
@@ -371,10 +272,6 @@ SEARCH_PROVIDERS = [
("searxng", _check_searxng, _search_searxng),
]
SELECTABLE_SEARCH_PROVIDERS = SEARCH_PROVIDERS + [
("parallel", _check_parallel, _search_parallel),
]
def search_web(
query: str,
max_results: int = 5,
@@ -387,8 +284,7 @@ def search_web(
) -> List[Dict[str, Any]]:
"""Search the web using one or more providers with automatic fallback.
VALID PROVIDER NAMES: tavily, brave, exa, youdotcom, duckduckgo, searxng, parallel
Parallel is opt-in and does not change the existing automatic provider order.
VALID PROVIDER NAMES: tavily, brave, exa, youdotcom, duckduckgo, searxng
Args:
query: Search query string
@@ -397,8 +293,7 @@ def search_web(
- Single string: "tavily" (uses just that provider)
- Comma-separated: "tavily, brave" (tries in order, falls back on failure)
- List: ["tavily", "brave"] (same as comma-separated)
If not specified, tries the default providers in order until one succeeds.
Select "parallel" to use the keyless Parallel Search MCP provider.
If not specified, tries all providers in order until one succeeds.
Provider-specific options (only used by their respective provider):
- searxng_url: [SearxNG] Custom instance URL
@@ -413,7 +308,7 @@ def search_web(
errors = []
# Valid provider names for reference
valid_provider_names = [p[0] for p in SELECTABLE_SEARCH_PROVIDERS]
valid_provider_names = [p[0] for p in SEARCH_PROVIDERS]
# Check for WEB_SEARCH_PROVIDER env var (set by CLI --web-provider flag)
if providers is None:
@@ -433,7 +328,7 @@ def search_web(
for name in providers:
name_lower = name.lower().strip()
matched = False
for pname, check_fn, search_fn in SELECTABLE_SEARCH_PROVIDERS:
for pname, check_fn, search_fn in SEARCH_PROVIDERS:
if pname == name_lower:
provider_order.append((pname, check_fn, search_fn))
matched = True
@@ -503,7 +398,7 @@ def get_available_providers() -> List[Dict[str, Any]]:
print(f"{p['name']}: {'✓' if p['available'] else '✗'} {p.get('reason', '')}")
"""
result = []
for name, check_func, _ in SELECTABLE_SEARCH_PROVIDERS:
for name, check_func, _ in SEARCH_PROVIDERS:
is_available, error = check_func()
result.append({
"name": name,
@@ -7,78 +7,48 @@ ensuring no RuntimeError: "This event loop is already running" crashes.
import asyncio
import concurrent.futures
import contextvars
import logging
from typing import Any, Awaitable, Optional, TypeVar
from typing import Any, Awaitable, TypeVar
logger = logging.getLogger(__name__)
T = TypeVar('T')
def run_coroutine_from_any_context(coro: Awaitable[T], timeout: Optional[float] = 300) -> T:
def run_coroutine_from_any_context(coro: Awaitable[T], timeout: float = 300) -> T:
"""
Safely run a coroutine to completion from any context (sync OR inside a
running event loop).
- No running loop: run it directly with ``asyncio.run`` (wrapped in
``asyncio.wait_for`` when a timeout is set).
- Running loop (FastAPI, Jupyter, a bot handler calling ``agent.chat()``):
run it on a dedicated worker thread with its own loop so we never
deadlock on the caller's loop. The caller's ``contextvars`` are copied
into that thread so trace/session/approval context set by the caller is
visible inside the coroutine, and the timeout is enforced *inside* the
worker loop via ``asyncio.wait_for``. The caller also waits on the worker
with the same ``timeout`` as a hard backstop, so a coroutine that ignores
cancellation (e.g. a blocking C call that never yields) can still not pin
the caller past ``timeout``.
Safely run a coroutine from sync context only.
Args:
coro: The coroutine to execute
timeout: Maximum execution time in seconds (default: 5 minutes).
``None`` disables the timeout.
timeout: Maximum execution time in seconds (default: 5 minutes)
Returns:
The result of the coroutine
Raises:
TimeoutError: If the coroutine takes longer than ``timeout``
RuntimeError: If called from async context (event loop is running)
TimeoutError: If the coroutine takes longer than timeout
Examples:
>>> async def my_async_function():
... return "hello"
>>> result = run_coroutine_from_any_context(my_async_function())
>>> print(result) # "hello"
"""
def _wrapped() -> Awaitable[T]:
return asyncio.wait_for(coro, timeout) if timeout is not None else coro
try:
asyncio.get_running_loop()
except RuntimeError:
# No running event loop — safe to drive it directly on this thread.
return asyncio.run(_wrapped())
# A loop is already running on this thread. Offload to a worker thread with
# its own loop, carrying the caller's contextvars across so context set by
# the caller (trace/session/approval) is visible inside the coroutine.
ctx = contextvars.copy_context()
executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
future = executor.submit(ctx.run, asyncio.run, _wrapped())
try:
# The timeout is normally enforced *inside* the worker loop by
# asyncio.wait_for. But a coroutine that ignores cancellation (a blocking
# C call, a shielded/uncancellable await) can keep the worker alive past
# the deadline, in which case that inner wait_for never fires. We
# therefore also bound the caller's wait here so a stuck worker can never
# pin the caller indefinitely.
return future.result(timeout=timeout)
except concurrent.futures.TimeoutError as exc:
raise TimeoutError(
f"Coroutine did not complete within {timeout}s"
) from exc
finally:
# Never block on a stuck worker: don't wait for it to wind down.
executor.shutdown(wait=False)
loop = asyncio.get_running_loop()
# If we get here, event loop is running - this is not allowed
raise RuntimeError(
"run_coroutine_from_any_context() cannot be called from async context. "
"Use 'await' instead when calling from async functions."
)
except RuntimeError as e:
if "no running event loop" in str(e).lower():
# No event loop — safe to use asyncio.run()
return asyncio.run(coro)
else:
# Re-raise our own RuntimeError
raise
def is_async_context() -> bool:
"""
@@ -1,72 +0,0 @@
"""
Atomic JSON persistence helper.
Consolidates the temp-file-plus-``os.replace`` pattern already used across the
SDK (scheduler stores, session store, skills manager, gateway protocols) into a
single shared writer. Serialisation happens on the temp file *before* the target
is touched, so a crash, OOM kill, or a non-serialisable value can never truncate
the previous good file: the old bytes stay intact and ``os.replace`` swaps in the
new content atomically on both POSIX and Windows.
"""
import json
import os
import tempfile
from typing import Any
__all__ = ["atomic_write_json", "atomic_write_text"]
def atomic_write_text(path: str, text: str, encoding: str = "utf-8") -> None:
"""Atomically write ``text`` to ``path`` (temp file + ``os.replace``).
For already-serialised payloads (markdown, pre-dumped JSON strings) where
the caller holds a string rather than a JSON-serialisable object. The old
file is untouched until the fully-written temp file is swapped in.
"""
directory = os.path.dirname(os.path.abspath(path)) or "."
fd, tmp = tempfile.mkstemp(
dir=directory, prefix=os.path.basename(path) + ".", suffix=".tmp"
)
try:
with os.fdopen(fd, "w", encoding=encoding) as f:
f.write(text)
f.flush()
os.fsync(f.fileno())
os.replace(tmp, path)
except BaseException:
try:
os.unlink(tmp)
except OSError:
pass
raise
def atomic_write_json(path: str, data: Any, **dump_kwargs: Any) -> None:
"""Atomically serialise ``data`` as JSON to ``path``.
Writes to a temp file in the target's directory, fsyncs it, then
``os.replace``\\ s it into place. If serialisation or the write fails the
temp file is removed and the original ``path`` is left untouched.
Args:
path: Destination file path.
data: JSON-serialisable object.
**dump_kwargs: Passed through to :func:`json.dump` (e.g. ``indent=2``).
"""
directory = os.path.dirname(os.path.abspath(path)) or "."
fd, tmp = tempfile.mkstemp(
dir=directory, prefix=os.path.basename(path) + ".", suffix=".tmp"
)
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(data, f, **dump_kwargs)
f.flush()
os.fsync(f.fileno())
os.replace(tmp, path)
except BaseException:
try:
os.unlink(tmp)
except OSError:
pass
raise
@@ -6137,13 +6137,9 @@ class WorkflowManager:
"saved_at_iso": datetime.now().isoformat()
}
# Atomic write: serialise to a temp file and os.replace into place so a
# crash/OOM/kill mid-write cannot destroy the previous good checkpoint
# (which would force the workflow to restart from step 0 and re-run and
# re-bill completed steps).
from ..utils.atomic_io import atomic_write_json
atomic_write_json(str(checkpoint_file), checkpoint_data, indent=2, default=str)
with open(checkpoint_file, 'w') as f:
json.dump(checkpoint_data, f, indent=2, default=str)
self._log(f"Saved checkpoint '{name}' at step {completed_steps}")
return str(checkpoint_file)
Loaded 100 of 292 files, more files were not shown because too many files have changed in this diff. Show more