mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 01:08:39 -10:00
Compare commits
358
Commits
No files matched your search
@@ -16,6 +16,9 @@ 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/';
|
||||
@@ -328,7 +331,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/\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/\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';
|
||||
|
||||
async function getMergeState(github, owner, repo, prNumber) {
|
||||
const query = `
|
||||
@@ -583,6 +586,11 @@ 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 || '';
|
||||
@@ -973,6 +981,15 @@ 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);
|
||||
@@ -1019,6 +1036,13 @@ 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);
|
||||
|
||||
@@ -1146,6 +1170,10 @@ module.exports = {
|
||||
claudeRunBlocksPr,
|
||||
hasBlockingClaudeRunForPr,
|
||||
isCiOnlyChange,
|
||||
isExternalExampleOnlyChange,
|
||||
EXTERNAL_EXAMPLE_PATH_PREFIX,
|
||||
EXTERNAL_EXAMPLE_LABEL,
|
||||
MAINTAINER_ACCEPT_EXAMPLE_LABEL,
|
||||
isInternalPullRequestLink,
|
||||
resolvePrNumberFromLinkedPullRequests,
|
||||
resolvePrNumberFromHeadBranch,
|
||||
|
||||
@@ -73,4 +73,33 @@ 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);
|
||||
@@ -37,8 +37,17 @@ 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';
|
||||
@@ -63,7 +72,8 @@ function reasonToBlockerLabel(reason) {
|
||||
r.includes('without test') ||
|
||||
r.includes('files changed') ||
|
||||
r.includes('possible secret') ||
|
||||
r.includes('agent.py')
|
||||
r.includes('agent.py') ||
|
||||
r.includes('external example only')
|
||||
) 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';
|
||||
@@ -133,11 +143,16 @@ 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(all);
|
||||
const desired = new Set(labelSet);
|
||||
const toRemove = current.filter((l) => !desired.has(l) && l !== 'pipeline/merged');
|
||||
const toAdd = all.filter((l) => !current.includes(l));
|
||||
const toAdd = labelSet.filter((l) => !current.includes(l));
|
||||
|
||||
for (const name of toRemove) {
|
||||
try {
|
||||
@@ -199,6 +214,11 @@ 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);
|
||||
@@ -223,9 +243,8 @@ 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) {
|
||||
if (result.ready && isUpstreamHeadRepo(pr, owner, repo)) {
|
||||
readyCandidates.push({
|
||||
prNumber: pr.number,
|
||||
createdAt: result.createdAt || new Date(pr.created_at).getTime(),
|
||||
@@ -248,6 +267,7 @@ module.exports = {
|
||||
STAGE_LABELS,
|
||||
BLOCKER_LABELS,
|
||||
ALL_PIPELINE_LABELS,
|
||||
isUpstreamHeadRepo,
|
||||
deriveStage,
|
||||
deriveBlockerLabels,
|
||||
computePipelineLabels,
|
||||
|
||||
@@ -21,6 +21,9 @@ on:
|
||||
- '.github/actions/install-monorepo-packages/**'
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
cross-repo-integration:
|
||||
runs-on: ubuntu-latest
|
||||
@@ -41,13 +44,13 @@ jobs:
|
||||
ref: main
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v4
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install dependencies
|
||||
run: python -m pip install --upgrade pip
|
||||
- uses: PraisonAI/.github/actions/install-monorepo-packages
|
||||
- uses: ./PraisonAI/.github/actions/install-monorepo-packages
|
||||
with:
|
||||
repo-root: PraisonAI/
|
||||
wrapper-extras: ".[ui,dev]"
|
||||
@@ -138,13 +141,13 @@ jobs:
|
||||
ref: main
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v4
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install dependencies
|
||||
run: python -m pip install --upgrade pip
|
||||
- uses: PraisonAI/.github/actions/install-monorepo-packages
|
||||
- uses: ./PraisonAI/.github/actions/install-monorepo-packages
|
||||
with:
|
||||
repo-root: PraisonAI/
|
||||
wrapper-extras: ".[ui,dev]"
|
||||
@@ -162,4 +165,4 @@ jobs:
|
||||
python -m pytest tests/integration/test_aiui_host_agentic.py -v
|
||||
else
|
||||
echo "OPENAI_API_KEY not available, skipping agentic tests"
|
||||
fi
|
||||
fi
|
||||
@@ -1,6 +1,6 @@
|
||||
name: Pipeline Status Sync
|
||||
|
||||
# Sync pipeline/* stage and blocker labels on open internal PRs.
|
||||
# Sync pipeline/* stage and blocker labels on open PRs (fork heads included; merge gate only for upstream branches).
|
||||
|
||||
on:
|
||||
schedule:
|
||||
|
||||
@@ -16,8 +16,6 @@ jobs:
|
||||
review:
|
||||
runs-on: ubuntu-latest
|
||||
if: >
|
||||
secrets.PRAISONAI_APP_ID != '' &&
|
||||
secrets.PRAISONAI_APP_PRIVATE_KEY != '' &&
|
||||
(
|
||||
(github.event_name == 'pull_request' && github.event.pull_request.draft == false) ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
@@ -32,7 +30,28 @@ jobs:
|
||||
)
|
||||
|
||||
steps:
|
||||
- name: Check reviewer configuration
|
||||
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
|
||||
|
||||
- name: Determine checkout ref
|
||||
if: steps.config.outputs.ready == 'true'
|
||||
id: dest
|
||||
env:
|
||||
ISSUE_NUMBER: ${{ github.event.issue.number }}
|
||||
@@ -48,6 +67,7 @@ jobs:
|
||||
fi
|
||||
|
||||
- name: Checkout Repository
|
||||
if: steps.config.outputs.ready == 'true'
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
@@ -55,6 +75,7 @@ jobs:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Generate GitHub App Token
|
||||
if: steps.config.outputs.ready == 'true'
|
||||
id: generate_token
|
||||
uses: tibdex/github-app-token@v2
|
||||
with:
|
||||
@@ -62,14 +83,17 @@ jobs:
|
||||
private_key: ${{ secrets.PRAISONAI_APP_PRIVATE_KEY }}
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.config.outputs.ready == 'true'
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install PraisonAI
|
||||
if: steps.config.outputs.ready == 'true'
|
||||
run: pip install "praisonai[all]"
|
||||
|
||||
- name: Run PraisonAI PR Review
|
||||
if: steps.config.outputs.ready == 'true'
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ steps.generate_token.outputs.token }}
|
||||
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
|
||||
@@ -773,7 +773,9 @@ jobs:
|
||||
env:
|
||||
# Bump deliberately: this is the version everything in this job was
|
||||
# verified against. Latest at time of writing is v0.33.3.
|
||||
OLLAMA_VERSION: v0.33.2
|
||||
# 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"
|
||||
# 522 MB, capabilities: completion, tools, thinking.
|
||||
OLLAMA_TEST_MODEL: qwen3:0.6b
|
||||
|
||||
|
||||
@@ -122,3 +122,4 @@ 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/
|
||||
@@ -12,7 +12,6 @@
|
||||
<a href="https://github.com/MervinPraison/PraisonAI"><img src="https://static.pepy.tech/badge/PraisonAI" alt="Total Downloads" /></a>
|
||||
<a href="https://github.com/MervinPraison/PraisonAI"><img src="https://img.shields.io/github/v/release/MervinPraison/PraisonAI" alt="Latest Stable Version" /></a>
|
||||
<a href="https://github.com/MervinPraison/PraisonAI"><img src="https://img.shields.io/badge/License-MIT-yellow.svg" alt="License" /></a>
|
||||
<a href="https://registry.modelcontextprotocol.io/servers/io.github.MervinPraison/praisonai"><img src="https://img.shields.io/badge/MCP-Registry-blue" alt="MCP Registry" /></a>
|
||||
</p>
|
||||
|
||||
<div align="center">
|
||||
|
||||
@@ -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) -> 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 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,7 +1018,8 @@ export { // Bot types
|
||||
FailoverManager, // Enums
|
||||
SandboxStatus, // Gateway types
|
||||
type GatewayConfig, // Other types
|
||||
type ProviderStatus, AutonomyLevel, GatewayEventType, RagRetrievalPolicy } from "./gateway";
|
||||
type ProviderStatus, // Session-projection reducer (Issue #5324)
|
||||
SessionProjection, 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
|
||||
|
||||
@@ -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.10" \
|
||||
"praisonai>=4.7.11" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.10" \
|
||||
"praisonai>=4.7.11" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.10" \
|
||||
"praisonai>=4.7.11" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
"""OpenRouter video generation via VideoAgent.
|
||||
|
||||
Requires OPENROUTER_API_KEY. Optionally set OPENROUTER_REFERER and
|
||||
OPENROUTER_APP_TITLE for OpenRouter dashboard attribution.
|
||||
|
||||
The openrouter/ prefix is required so the request routes to OpenRouter.
|
||||
Discover exact model ids at https://openrouter.ai/api/v1/videos/models.
|
||||
"""
|
||||
|
||||
from praisonaiagents import VideoAgent
|
||||
|
||||
agent = VideoAgent(llm="openrouter/google/veo-3.1")
|
||||
|
||||
video = agent.generate(
|
||||
"A serene mountain landscape at sunset, locked camera",
|
||||
duration=6,
|
||||
resolution="720p",
|
||||
aspect_ratio="16:9",
|
||||
generate_audio=False,
|
||||
)
|
||||
print(f"Video ID: {video.id}")
|
||||
|
||||
completed = agent.wait_for_completion(video.id)
|
||||
if completed.status == "completed":
|
||||
agent.download(video.id, "openrouter_veo.mp4")
|
||||
print("Saved openrouter_veo.mp4")
|
||||
else:
|
||||
print(f"Generation ended with status: {completed.status}")
|
||||
@@ -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 | 18870 |
|
||||
| ✅ Agent | 22 | 59 | 19076 |
|
||||
| ✅ Agent-to-Agent (A2A) | 1 | 7 | 2996 |
|
||||
| ✅ Agent-to-User (A2U) | 1 | 3 | 701 |
|
||||
| ✅ Approval | 1 | 8 | 3941 |
|
||||
| ✅ Approval | 1 | 8 | 3969 |
|
||||
| ✅ Audio | 2 | 12 | 1000 |
|
||||
| ✅ Auto Generation | 5 | 12 | 3879 |
|
||||
| ✅ Autonomy | 3 | 7 | 2794 |
|
||||
| ✅ Bots | 7 | 36 | 14750 |
|
||||
| ✅ Auto Generation | 5 | 12 | 3941 |
|
||||
| ✅ Autonomy | 3 | 7 | 2805 |
|
||||
| ✅ Bots | 7 | 38 | 15809 |
|
||||
| ✅ Budget | 1 | 1 | 287 |
|
||||
| ✅ CLI | 5 | 124 | 38915 |
|
||||
| ✅ CLI | 5 | 124 | 39191 |
|
||||
| ✅ 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 | 14275 |
|
||||
| ✅ Context Management | 16 | 36 | 14384 |
|
||||
| ✅ Deep Research | 8 | 2 | 591 |
|
||||
| ✅ Display | 6 | 3 | 844 |
|
||||
| ✅ Embeddings | 6 | 23 | 2051 |
|
||||
| ✅ Evaluation | 1 | 9 | 3537 |
|
||||
| ✅ Events | 1 | 2 | 780 |
|
||||
| ✅ Execution | 3 | 4 | 1693 |
|
||||
| ✅ Failover | 2 | 1 | 480 |
|
||||
| ✅ Evaluation | 1 | 9 | 3637 |
|
||||
| ✅ Events | 1 | 2 | 782 |
|
||||
| ✅ Execution | 3 | 4 | 1694 |
|
||||
| ✅ Failover | 2 | 1 | 537 |
|
||||
| ✅ Files | 2 | 7 | 2376 |
|
||||
| ✅ Flow | 1 | 3 | 811 |
|
||||
| ✅ Gateway | 7 | 87 | 33047 |
|
||||
| ✅ Guardrails | 5 | 4 | 2529 |
|
||||
| ✅ Gateway | 7 | 97 | 37597 |
|
||||
| ✅ Guardrails | 5 | 4 | 2642 |
|
||||
| ✅ Handoffs | 11 | 6 | 2748 |
|
||||
| ✅ Hooks | 2 | 9 | 4990 |
|
||||
| ✅ Image | 1 | 11 | 1255 |
|
||||
| ✅ Knowledge | 4 | 19 | 5974 |
|
||||
| ✅ LLM | 3 | 15 | 5946 |
|
||||
| ✅ Loops | 4 | 5 | 1475 |
|
||||
| ✅ MCP | 1 | 63 | 14859 |
|
||||
| ✅ Memory | 6 | 21 | 8186 |
|
||||
| ✅ 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 |
|
||||
| ✅ OCR | 2 | 1 | 237 |
|
||||
| ✅ Observability | 2 | 23 | 2842 |
|
||||
| ✅ Observability | 2 | 23 | 2859 |
|
||||
| ✅ Optimizer | 1 | 2 | 792 |
|
||||
| ✅ Output | 3 | 5 | 1273 |
|
||||
| ✅ Output | 3 | 5 | 1289 |
|
||||
| ✅ Parallel Execution | 3 | 2 | 528 |
|
||||
| ✅ Planning | 6 | 7 | 1825 |
|
||||
| ✅ Plugins | 8 | 5 | 3316 |
|
||||
| ✅ Planning | 6 | 7 | 1827 |
|
||||
| ✅ Plugins | 8 | 5 | 4108 |
|
||||
| ✅ Prompts | 2 | 9 | 2043 |
|
||||
| ✅ Providers | 1 | 57 | 8640 |
|
||||
| ✅ Providers | 1 | 58 | 8751 |
|
||||
| ✅ Query | 1 | 2 | 736 |
|
||||
| ✅ RAG | 5 | 15 | 3138 |
|
||||
| ✅ RAG | 5 | 15 | 3152 |
|
||||
| ✅ Realtime | 2 | 5 | 668 |
|
||||
| ✅ Reflection | 3 | 3 | 732 |
|
||||
| ✅ Retrieval | 2 | 5 | 1046 |
|
||||
| ✅ Routing | 1 | 2 | 430 |
|
||||
| ✅ Sandbox | 5 | 7 | 3712 |
|
||||
| ✅ Security | 1 | 3 | 3598 |
|
||||
| ✅ Sessions | 4 | 20 | 8058 |
|
||||
| ✅ Skills | 6 | 16 | 5984 |
|
||||
| ✅ Sandbox | 5 | 7 | 3720 |
|
||||
| ✅ Security | 1 | 3 | 3642 |
|
||||
| ✅ Sessions | 4 | 21 | 8388 |
|
||||
| ✅ Skills | 6 | 16 | 6045 |
|
||||
| ✅ Tasks | 2 | 8 | 3266 |
|
||||
| ✅ Telemetry | 1 | 2 | 660 |
|
||||
| ✅ Templates | 1 | 8 | 1667 |
|
||||
| ✅ Tools | 12 | 145 | 44511 |
|
||||
| ✅ Templates | 1 | 9 | 1859 |
|
||||
| ✅ Tools | 12 | 147 | 45378 |
|
||||
| ✅ Tracing | 3 | 2 | 139 |
|
||||
| ✅ Vector Store | 1 | 12 | 1219 |
|
||||
| ✅ Video | 2 | 6 | 728 |
|
||||
| ✅ Video | 2 | 10 | 1303 |
|
||||
| ✅ Vision | 2 | 1 | 329 |
|
||||
| ✅ Web | 3 | 10 | 2815 |
|
||||
| ✅ Workflows | 5 | 21 | 8862 |
|
||||
| ✅ Web | 3 | 10 | 2822 |
|
||||
| ✅ Workflows | 5 | 21 | 8880 |
|
||||
|
||||
## Documentation Without Features
|
||||
|
||||
These docs exist but don't match any implemented feature category:
|
||||
|
||||
- ℹ️ Database (43 docs, 7657 lines)
|
||||
- ℹ️ Database (44 docs, 8206 lines)
|
||||
- ℹ️ Documents (1 docs, 810 lines)
|
||||
- ℹ️ Teams (1 docs, 214 lines)
|
||||
|
||||
|
||||
@@ -183,7 +183,9 @@ banner. Key fragments:
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
# Bump deliberately; the cache key and what we test both follow this.
|
||||
OLLAMA_VERSION: v0.33.2
|
||||
# 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_TEST_MODEL: qwen3:0.6b
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
@@ -204,9 +206,12 @@ banner. Key fragments:
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
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
|
||||
# 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
|
||||
ollama --version
|
||||
- name: Start Ollama and pull model
|
||||
shell: bash
|
||||
|
||||
@@ -37,6 +37,22 @@ _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,
|
||||
@@ -220,9 +236,25 @@ def create_lazy_getattr_with_fallback(
|
||||
_cache[name] = None
|
||||
return None
|
||||
except ImportError as 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.
|
||||
# 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
|
||||
_logger.debug(
|
||||
"Lazy import of %r (%s.%s) failed: %s",
|
||||
name, module_path, attr_name, exc,
|
||||
|
||||
@@ -28,8 +28,10 @@ def _configure_environment():
|
||||
"PYDANTIC_WARNINGS_ENABLED": "False",
|
||||
}
|
||||
|
||||
# Use setdefault so a value the host exported before launch is preserved
|
||||
# (and not leaked/overwritten into subprocesses spawned afterwards).
|
||||
for key, value in env_vars.items():
|
||||
os.environ[key] = value
|
||||
os.environ.setdefault(key, value)
|
||||
|
||||
|
||||
# ========================================================================
|
||||
@@ -48,7 +50,11 @@ def _get_all_noisy_loggers() -> List[str]:
|
||||
# Markdown
|
||||
"markdown_it", "rich.markdown",
|
||||
# System
|
||||
"asyncio", "selector_events", "pydantic",
|
||||
# NOTE: "asyncio" is intentionally NOT suppressed here — it is the host
|
||||
# application's own logger, not a PraisonAI dependency's. Forcing it to
|
||||
# CRITICAL hides unrelated host errors (e.g. "Task exception was never
|
||||
# retrieved").
|
||||
"selector_events", "pydantic",
|
||||
"praisonaiagents.telemetry.telemetry",
|
||||
]
|
||||
|
||||
|
||||
@@ -7,10 +7,7 @@ from praisonaiagents._logging import get_logger
|
||||
import asyncio
|
||||
import contextlib
|
||||
import threading
|
||||
import concurrent.futures
|
||||
import random
|
||||
import re
|
||||
from typing import List, Optional, Any, Dict, Union, Literal, TYPE_CHECKING, Callable, Generator
|
||||
from typing import List, Optional, Any, Dict, Union, Literal, TYPE_CHECKING, Callable
|
||||
from collections import OrderedDict
|
||||
import inspect
|
||||
|
||||
@@ -1373,12 +1370,17 @@ 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
|
||||
@@ -2101,7 +2103,11 @@ 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
|
||||
self._hooks_registry_param = hooks # Store for lazy init
|
||||
# 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
|
||||
|
||||
# Handle llm= deprecation: model= is the preferred parameter name
|
||||
# llm= still works but shows deprecation warning
|
||||
@@ -2595,6 +2601,11 @@ 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
|
||||
@@ -2881,7 +2892,10 @@ 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)
|
||||
_mgr = PermissionManager(
|
||||
agent_name=self.name,
|
||||
workspace_root=os.getcwd(),
|
||||
)
|
||||
if isinstance(self._approval_permissions, dict):
|
||||
_mgr.load_rules_from_config(self._approval_permissions)
|
||||
elif isinstance(self._approval_permissions, (list, tuple)):
|
||||
@@ -3024,6 +3038,7 @@ Your Goal: {self.goal}
|
||||
# turns on the same Agent instance don't corrupt each other's buffer.
|
||||
self._turn_tools_used = []
|
||||
self._turn_tools_lock = DualLock()
|
||||
self._tool_executor_lock = DualLock()
|
||||
|
||||
# Database persistence (lazy - no imports until used)
|
||||
self._db = db
|
||||
@@ -3051,6 +3066,12 @@ 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
|
||||
@@ -3231,7 +3252,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 == "_cost_lock":
|
||||
elif k in ("_cost_lock", "_auto_save_lock"):
|
||||
object.__setattr__(result, k, threading.Lock())
|
||||
elif k == "_approvals_lock":
|
||||
object.__setattr__(result, k, asyncio.Lock())
|
||||
@@ -4155,12 +4176,44 @@ 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 [])}
|
||||
@@ -5338,14 +5391,19 @@ Summary:"""
|
||||
f"response_len={len(response_str)}"
|
||||
)
|
||||
|
||||
# Auto-save session after each async iteration (memory integration)
|
||||
self._auto_save_session()
|
||||
# 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)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
# GOAL COMPLETION JUDGE (opt-in): independent acceptance-criteria
|
||||
# gate. No-op when no goal loop is active (behaviour unchanged).
|
||||
# Offloaded to a worker thread (blocking judge LLM call / hooks)
|
||||
# so it never stalls the shared event loop.
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
_goal_gate = self._goal_gate(response_str)
|
||||
_goal_gate = await self._goal_gate_async(response_str)
|
||||
if _goal_gate is not None:
|
||||
_outcome, _reason = _goal_gate
|
||||
if _outcome == "done":
|
||||
@@ -5477,8 +5535,8 @@ Summary:"""
|
||||
# Yield control to allow other async tasks to run
|
||||
await asyncio.sleep(0)
|
||||
|
||||
# Final auto-save before returning
|
||||
self._auto_save_session()
|
||||
# Final auto-save before returning (offloaded off the event loop).
|
||||
await asyncio.to_thread(self._auto_save_session)
|
||||
|
||||
# Max iterations reached
|
||||
return AutonomyResult(
|
||||
@@ -7065,6 +7123,25 @@ 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.
|
||||
|
||||
@@ -7080,8 +7157,16 @@ Answer:"""
|
||||
return GuardrailResult(success=True, result=task_output)
|
||||
|
||||
try:
|
||||
# Call the guardrail function
|
||||
result = self._guardrail_fn(task_output)
|
||||
# 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)
|
||||
|
||||
# Convert the result to a GuardrailResult
|
||||
return GuardrailResult.from_tuple(result)
|
||||
@@ -7993,16 +8078,28 @@ Answer:"""
|
||||
except Exception as e:
|
||||
logger.warning(f"Task cleanup failed: {e}")
|
||||
|
||||
# ThreadPoolExecutor cleanup
|
||||
# ThreadPoolExecutor cleanup. Detach the executor under the same lock
|
||||
# that guards timeout retirement/replacement (see tool_execution.py) so
|
||||
# cleanup and a concurrent timeout can never act on different executor
|
||||
# generations. Shut down outside the critical section to avoid holding
|
||||
# the lock across a blocking call.
|
||||
try:
|
||||
if hasattr(self, '_tool_executor') and self._tool_executor:
|
||||
executor = None
|
||||
lock = getattr(self, '_tool_executor_lock', None)
|
||||
if lock is not None:
|
||||
with lock.sync():
|
||||
executor = getattr(self, '_tool_executor', None)
|
||||
self._tool_executor = None
|
||||
elif hasattr(self, '_tool_executor'):
|
||||
executor = self._tool_executor
|
||||
self._tool_executor = None
|
||||
if executor is not None:
|
||||
# Use cancel_futures only if supported (Python 3.9+)
|
||||
import sys
|
||||
if sys.version_info >= (3, 9):
|
||||
self._tool_executor.shutdown(wait=False, cancel_futures=True)
|
||||
executor.shutdown(wait=False, cancel_futures=True)
|
||||
else:
|
||||
self._tool_executor.shutdown(wait=False)
|
||||
delattr(self, '_tool_executor')
|
||||
executor.shutdown(wait=False)
|
||||
except Exception as e:
|
||||
logger.warning(f"ThreadPoolExecutor cleanup failed: {e}")
|
||||
|
||||
@@ -8105,22 +8202,35 @@ Answer:"""
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# ThreadPoolExecutor cleanup (async-safe)
|
||||
if hasattr(self, '_tool_executor') and self._tool_executor:
|
||||
# ThreadPoolExecutor cleanup (async-safe). Detach the executor under
|
||||
# the lock that guards timeout retirement/replacement so cleanup and
|
||||
# a concurrent timeout can never act on different executor
|
||||
# generations. Capture the reference before scheduling the shutdown
|
||||
# so the deferred call cannot resolve a replacement pool created by
|
||||
# another call.
|
||||
executor = None
|
||||
lock = getattr(self, '_tool_executor_lock', None)
|
||||
if lock is not None:
|
||||
async with lock.async_lock():
|
||||
executor = getattr(self, '_tool_executor', None)
|
||||
self._tool_executor = None
|
||||
elif hasattr(self, '_tool_executor'):
|
||||
executor = self._tool_executor
|
||||
self._tool_executor = None
|
||||
if executor is not None:
|
||||
import sys
|
||||
loop = asyncio.get_running_loop()
|
||||
# Use run_in_executor to avoid blocking the event loop
|
||||
if sys.version_info >= (3, 9):
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._tool_executor.shutdown(wait=False, cancel_futures=True)
|
||||
None,
|
||||
lambda: executor.shutdown(wait=False, cancel_futures=True)
|
||||
)
|
||||
else:
|
||||
await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._tool_executor.shutdown(wait=False)
|
||||
None,
|
||||
lambda: executor.shutdown(wait=False)
|
||||
)
|
||||
delattr(self, '_tool_executor')
|
||||
|
||||
# Approval scope cleanup (see close()).
|
||||
self._release_approval_scope()
|
||||
|
||||
@@ -27,27 +27,21 @@ 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.
|
||||
"""
|
||||
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
|
||||
from ..utils.async_bridge import run_coroutine_from_any_context
|
||||
return run_coroutine_from_any_context(coro, timeout=None)
|
||||
|
||||
|
||||
class DualLock:
|
||||
@@ -75,9 +69,24 @@ class DualLock:
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""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
|
||||
"""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
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
"""Return an unlocked primitive with no event-loop state sharing."""
|
||||
@@ -87,9 +96,31 @@ class DualLock:
|
||||
|
||||
@contextmanager
|
||||
def sync(self):
|
||||
"""Acquire lock in synchronous context using threading.Lock."""
|
||||
with self._thread_lock:
|
||||
"""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:
|
||||
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."""
|
||||
@@ -102,16 +133,64 @@ 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 asyncio.Lock.
|
||||
|
||||
Uses a per-event-loop asyncio.Lock to ensure proper async coordination
|
||||
without blocking the event loop or violating thread ownership semantics.
|
||||
"""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.
|
||||
"""
|
||||
async_lock = self._get_async_lock()
|
||||
async with async_lock:
|
||||
yield
|
||||
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()
|
||||
|
||||
def is_async_context(self) -> bool:
|
||||
"""Check if we're currently in an async context."""
|
||||
@@ -185,6 +264,11 @@ 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:
|
||||
@@ -192,6 +276,14 @@ 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:
|
||||
@@ -208,20 +300,32 @@ class AsyncSafeState:
|
||||
|
||||
def __enter__(self):
|
||||
"""Support for synchronous context manager protocol (backward compatibility)."""
|
||||
self._lock._thread_lock.acquire()
|
||||
self._sync_guard = self._lock.sync()
|
||||
self._sync_guard.__enter__()
|
||||
return self.value
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Support for synchronous context manager protocol (backward compatibility)."""
|
||||
self._lock._thread_lock.release()
|
||||
return None
|
||||
return self._sync_guard.__exit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
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:
|
||||
self._begin_async_access()
|
||||
# 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
|
||||
except Exception:
|
||||
async_lock.release()
|
||||
raise
|
||||
@@ -233,7 +337,10 @@ class AsyncSafeState:
|
||||
try:
|
||||
self._end_async_access()
|
||||
finally:
|
||||
async_lock.release()
|
||||
try:
|
||||
self._lock._thread_lock.release()
|
||||
finally:
|
||||
async_lock.release()
|
||||
return None
|
||||
|
||||
def get(self) -> Any:
|
||||
|
||||
@@ -751,6 +751,43 @@ 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.
|
||||
|
||||
@@ -1431,15 +1468,8 @@ Your Goal: {self.goal}"""
|
||||
from ..hooks import HookEvent as _HookEvent
|
||||
import logging
|
||||
|
||||
# Map policy strategy to compactor strategy
|
||||
strategy_map = {
|
||||
"truncate": LegacyStrategy.TRUNCATE,
|
||||
"summarise": LegacyStrategy.SUMMARIZE,
|
||||
"drop_oldest_tools": LegacyStrategy.PRUNE,
|
||||
"sliding_window": LegacyStrategy.SLIDING,
|
||||
}
|
||||
|
||||
compactor.strategy = strategy_map.get(policy.strategy.value, LegacyStrategy.PRUNE)
|
||||
# Map policy strategy to compactor strategy (single source of truth)
|
||||
compactor.strategy = LegacyStrategy.from_policy_value(policy.strategy.value)
|
||||
|
||||
# Execute hooks
|
||||
try:
|
||||
@@ -1624,15 +1654,8 @@ Your Goal: {self.goal}"""
|
||||
from ..hooks import HookEvent as _HookEvent
|
||||
import logging
|
||||
|
||||
# Map policy strategy to compactor strategy
|
||||
strategy_map = {
|
||||
"truncate": LegacyStrategy.TRUNCATE,
|
||||
"summarise": LegacyStrategy.SUMMARIZE,
|
||||
"drop_oldest_tools": LegacyStrategy.PRUNE,
|
||||
"sliding_window": LegacyStrategy.SLIDING,
|
||||
}
|
||||
|
||||
compactor.strategy = strategy_map.get(policy.strategy.value, LegacyStrategy.PRUNE)
|
||||
# Map policy strategy to compactor strategy (single source of truth)
|
||||
compactor.strategy = LegacyStrategy.from_policy_value(policy.strategy.value)
|
||||
|
||||
# Execute hooks
|
||||
try:
|
||||
@@ -3576,6 +3599,13 @@ 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)
|
||||
@@ -3583,7 +3613,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
except ToolExecutionError:
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -3620,41 +3650,48 @@ 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)
|
||||
|
||||
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
|
||||
# Wrap entire turn in try/finally so per-turn tracking is always
|
||||
# ended (the finally reads the ownership list during rollback).
|
||||
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:
|
||||
@@ -3700,6 +3737,11 @@ 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)
|
||||
@@ -3707,7 +3749,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
|
||||
if not self.self_reflect:
|
||||
# User message already added before LLM call via _build_messages
|
||||
@@ -3728,7 +3770,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
# 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)
|
||||
@@ -3739,7 +3781,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
|
||||
reflection_prompt = f"""
|
||||
Reflect on your previous response: '{response_text}'.
|
||||
@@ -3807,7 +3849,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
|
||||
# Check if we've hit max reflections
|
||||
if reflection_count >= self.max_reflect - 1:
|
||||
@@ -3825,7 +3867,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
|
||||
# If not satisfactory and not at max reflections, continue with regeneration
|
||||
logging.debug(f"{self.name} reflection count {reflection_count + 1}, continuing reflection process")
|
||||
@@ -3872,6 +3914,10 @@ 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."""
|
||||
@@ -4231,6 +4277,11 @@ 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)
|
||||
@@ -4238,7 +4289,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
except ToolExecutionError:
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -4264,6 +4315,15 @@ 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)
|
||||
|
||||
@@ -4367,6 +4427,11 @@ 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)
|
||||
@@ -4374,7 +4439,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
|
||||
# For regular responses (no self-reflection)
|
||||
if not self.self_reflect:
|
||||
@@ -4397,7 +4462,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
else:
|
||||
# Apply guardrail to regular response content
|
||||
try:
|
||||
@@ -4409,7 +4474,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
|
||||
# If self-reflection is enabled, implement reflection logic
|
||||
if self.self_reflect:
|
||||
@@ -4449,7 +4514,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
|
||||
_areflection_kwargs = {}
|
||||
if temperature is not None:
|
||||
@@ -4483,7 +4548,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
|
||||
# Check if we've hit max reflections
|
||||
if reflection_count >= self.max_reflect - 1:
|
||||
@@ -4502,7 +4567,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
|
||||
# Regenerate response based on reflection
|
||||
regenerate_messages = reflection_messages + [
|
||||
@@ -4552,7 +4617,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 None
|
||||
return self._guardrail_blocked_message(guard_e)
|
||||
continue
|
||||
|
||||
# This should never be reached due to the returns above
|
||||
@@ -4602,6 +4667,11 @@ 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)
|
||||
@@ -4716,7 +4786,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 None
|
||||
return self._guardrail_blocked_message(e)
|
||||
except ToolExecutionError:
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -4740,6 +4810,16 @@ 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"""
|
||||
@@ -5712,6 +5792,28 @@ 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).
|
||||
@@ -5731,7 +5833,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
from ..compaction import ContextCompactor
|
||||
from ..compaction.strategy import CompactionStrategy
|
||||
|
||||
_max_tok = getattr(_execution_cfg, 'max_context_tokens', None) or 8000
|
||||
_max_tok = self._resolve_compaction_max_tokens(_execution_cfg)
|
||||
_strategy = getattr(_execution_cfg, 'compaction_strategy', None) or CompactionStrategy.TRUNCATE
|
||||
|
||||
# Create LLM summarization function if strategy is LLM_SUMMARIZE
|
||||
@@ -5763,17 +5865,20 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
# Perform compaction
|
||||
if _strategy == CompactionStrategy.LLM_SUMMARIZE and _llm_fn:
|
||||
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)
|
||||
# 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)
|
||||
)
|
||||
else:
|
||||
compacted_msgs, _cr = _compactor.compact(messages)
|
||||
|
||||
@@ -5818,7 +5923,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
from ..compaction import ContextCompactor
|
||||
from ..compaction.strategy import CompactionStrategy
|
||||
|
||||
_max_tok = getattr(_execution_cfg, 'max_context_tokens', None) or 8000
|
||||
_max_tok = self._resolve_compaction_max_tokens(_execution_cfg)
|
||||
_strategy = getattr(_execution_cfg, 'compaction_strategy', None) or CompactionStrategy.TRUNCATE
|
||||
|
||||
# Create LLM summarization function if strategy is LLM_SUMMARIZE
|
||||
|
||||
@@ -31,23 +31,66 @@ 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 threading.Semaphore, so
|
||||
the limit is a true global cap across every event loop and thread.
|
||||
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.
|
||||
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 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] = {}
|
||||
# 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] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def set_limit(self, agent_name: str, max_concurrent: int) -> None:
|
||||
@@ -59,8 +102,11 @@ class ConcurrencyRegistry:
|
||||
"""
|
||||
with self._lock:
|
||||
self._limits[agent_name] = max_concurrent
|
||||
# Reset semaphore so next acquire creates a fresh one at the new limit
|
||||
self._semaphores.pop(agent_name, None)
|
||||
# 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)
|
||||
|
||||
def get_limit(self, agent_name: str) -> int:
|
||||
"""Get concurrency limit for an agent."""
|
||||
@@ -68,39 +114,57 @@ 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)."""
|
||||
"""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.
|
||||
"""
|
||||
with self._lock:
|
||||
self._limits.pop(agent_name, None)
|
||||
self._semaphores.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)
|
||||
|
||||
def _get_semaphore(self, agent_name: str) -> Optional[threading.Semaphore]:
|
||||
"""Get or create the loop-neutral semaphore for an agent.
|
||||
def _get_limiter(self, agent_name: str) -> Optional[_DynamicLimiter]:
|
||||
"""Get or create the loop-neutral limiter for an agent.
|
||||
|
||||
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.
|
||||
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.
|
||||
"""
|
||||
with self._lock:
|
||||
limit = self._limits.get(agent_name, self._default_limit)
|
||||
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
|
||||
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
|
||||
|
||||
async def acquire(self, agent_name: str) -> None:
|
||||
"""Acquire concurrency slot for agent. No-op if unlimited.
|
||||
|
||||
Waits on the loop-neutral semaphore in short, cancellable polls so the
|
||||
Waits on the loop-neutral limiter 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().
|
||||
"""
|
||||
sem = self._get_semaphore(agent_name)
|
||||
if sem is None:
|
||||
limiter = self._get_limiter(agent_name)
|
||||
if limiter is None:
|
||||
return
|
||||
while True:
|
||||
if sem.acquire(blocking=False):
|
||||
if limiter.acquire(blocking=False):
|
||||
return
|
||||
# Yield to the loop; on cancellation this raises and no permit leaks.
|
||||
await asyncio.sleep(0.005)
|
||||
@@ -110,21 +174,18 @@ 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 semaphore is loop-neutral.
|
||||
current thread, since the limiter is loop-neutral.
|
||||
"""
|
||||
sem = self._get_semaphore(agent_name)
|
||||
if sem is not None:
|
||||
sem.acquire()
|
||||
limiter = self._get_limiter(agent_name)
|
||||
if limiter is not None:
|
||||
limiter.acquire()
|
||||
|
||||
def release(self, agent_name: str) -> None:
|
||||
"""Release concurrency slot for agent. No-op if unlimited."""
|
||||
with self._lock:
|
||||
sem = self._semaphores.get(agent_name)
|
||||
if sem is not None:
|
||||
try:
|
||||
sem.release()
|
||||
except ValueError:
|
||||
pass # Already fully released
|
||||
limiter = self._limiters.get(agent_name)
|
||||
if limiter is not None:
|
||||
limiter.release()
|
||||
|
||||
@asynccontextmanager
|
||||
async def throttle(self, agent_name: str):
|
||||
|
||||
@@ -58,8 +58,8 @@ def save_cost_report(
|
||||
}
|
||||
|
||||
filepath = COST_DIR / f"{safe_name}.json"
|
||||
with open(filepath, "w") as f:
|
||||
json.dump(report, f, indent=2)
|
||||
from ..utils.atomic_io import atomic_write_json
|
||||
atomic_write_json(str(filepath), report, indent=2)
|
||||
|
||||
logger.info(f"[costs] Saved cost report to {filepath}")
|
||||
return str(filepath)
|
||||
|
||||
@@ -898,14 +898,18 @@ You are a professional research analyst. When conducting research:
|
||||
if grounding_file_prompt:
|
||||
create_params["grounding_file_prompt"] = grounding_file_prompt
|
||||
|
||||
interaction = self.gemini_client.interactions.create(**create_params)
|
||||
interaction = await asyncio.to_thread(
|
||||
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 = self.gemini_client.interactions.get(interaction.id)
|
||||
interaction = await asyncio.to_thread(
|
||||
self.gemini_client.interactions.get, interaction.id
|
||||
)
|
||||
|
||||
if interaction.status == "completed":
|
||||
if self.verbose:
|
||||
|
||||
@@ -22,49 +22,28 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
from typing import List, Optional, Any, Dict, Union, Generator, TYPE_CHECKING
|
||||
|
||||
# Shared HTTP launch state (Agent.launch() multi-endpoint registration)
|
||||
_server_lock = threading.Lock()
|
||||
_server_started: Dict[int, bool] = {}
|
||||
_registered_agents: Dict[int, Dict[str, str]] = {}
|
||||
_shared_apps: Dict[int, Any] = {}
|
||||
# Shared HTTP launch state — the process-wide registry is defined in
|
||||
# agent/server_registry.py and shared with PraisonAIAgents.launch() so a
|
||||
# standalone Agent and a team launched on the same port cannot double-bind it or
|
||||
# hide each other's endpoints.
|
||||
from .server_registry import _server_registry
|
||||
|
||||
|
||||
def cleanup_launch_registration(agent_id: str) -> None:
|
||||
"""Remove launch() endpoints registered for ``agent_id`` from the shared
|
||||
HTTP state that ``Agent.launch()`` actually populates.
|
||||
|
||||
``Agent.launch()`` writes routes into the module-level ``_registered_agents``
|
||||
and ``_shared_apps`` dicts above. This tears those routes back down so a
|
||||
closed agent's endpoint stops accepting requests and the request handler
|
||||
closure no longer keeps the Agent object graph alive.
|
||||
``Agent.launch()`` reserves routes on the shared ``_server_registry``. This
|
||||
tears those routes back down so a closed agent's endpoint stops accepting
|
||||
requests and the request handler closure no longer keeps the Agent object
|
||||
graph alive.
|
||||
|
||||
The endpoint-registry removal and the matching FastAPI route teardown are
|
||||
performed together under the registry lock (see
|
||||
``_AgentServerRegistry.teardown_routes_for``) so cleanup cannot race a
|
||||
concurrent launch that reserves and registers a route on the same port.
|
||||
"""
|
||||
with _server_lock:
|
||||
for port, paths in list(_registered_agents.items()):
|
||||
stale_paths = [p for p, aid in paths.items() if aid == agent_id]
|
||||
for p in stale_paths:
|
||||
del paths[p]
|
||||
app = _shared_apps.get(port)
|
||||
if app is not None:
|
||||
try:
|
||||
# Only drop the agent-owned POST route. launch() always
|
||||
# registers the handler via ``.post(path)``, so filtering
|
||||
# on both path AND method preserves unrelated routes that
|
||||
# share the path with a different verb (e.g. the built-in
|
||||
# GET /health and GET /).
|
||||
app.router.routes = [
|
||||
r for r in app.router.routes
|
||||
if not (
|
||||
getattr(r, "path", None) == p
|
||||
and "POST" in (getattr(r, "methods", None) or set())
|
||||
)
|
||||
]
|
||||
# Invalidate cached OpenAPI schema so removed routes
|
||||
# disappear from /openapi.json and /docs.
|
||||
app.openapi_schema = None
|
||||
except Exception:
|
||||
pass
|
||||
if not paths:
|
||||
_registered_agents.pop(port, None)
|
||||
_server_registry.teardown_routes_for(agent_id)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
@@ -221,8 +200,12 @@ class ExecutionMixin:
|
||||
response = await self.achat(prompt, **kwargs)
|
||||
response_str = str(response) if response else ""
|
||||
|
||||
# Auto-save session after chat
|
||||
self._auto_save_session()
|
||||
# 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)
|
||||
|
||||
# Wrap in AutonomyResult for consistent API
|
||||
from .autonomy import AutonomyResult
|
||||
@@ -493,7 +476,8 @@ class ExecutionMixin:
|
||||
extra={"session_id": sid, "managed_session_id": msid},
|
||||
exc_info=True,
|
||||
)
|
||||
self._auto_save_session()
|
||||
# Offload the blocking session write off the event loop.
|
||||
await asyncio.to_thread(self._auto_save_session)
|
||||
|
||||
return result
|
||||
|
||||
@@ -1754,12 +1738,17 @@ 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. The check is pure/sync (no awaits).
|
||||
# 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.
|
||||
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 = check(function_name, arguments, tools_override)
|
||||
policy_result = await asyncio.to_thread(
|
||||
check, function_name, arguments, tools_override
|
||||
)
|
||||
if isinstance(policy_result, dict):
|
||||
return policy_result # Error dict
|
||||
_, arguments = policy_result
|
||||
@@ -1786,9 +1775,12 @@ 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 = apply_result_guardrails(function_name, mcp_result)
|
||||
mcp_result = await asyncio.to_thread(
|
||||
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.
|
||||
@@ -2149,9 +2141,14 @@ 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 = apply_result_guardrails(function_name, result, tools_override)
|
||||
result = await asyncio.to_thread(
|
||||
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
|
||||
@@ -2274,18 +2271,13 @@ Write the complete compiled report:"""
|
||||
NOTE: This implementation will be moved to wrapper layer in future version.
|
||||
For now, it maintains backward compatibility while following lazy import patterns.
|
||||
"""
|
||||
global _server_started, _registered_agents, _shared_apps, _server_lock
|
||||
|
||||
from .launch_security import authorise_launch_request, resolve_launch_host
|
||||
|
||||
# Try to import FastAPI dependencies - lazy loading
|
||||
try:
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException, Request
|
||||
from fastapi import HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel
|
||||
import threading
|
||||
import time
|
||||
import asyncio
|
||||
|
||||
# Define the request model here since we need pydantic
|
||||
@@ -2311,130 +2303,102 @@ Write the complete compiled report:"""
|
||||
# process-wide launch token and retroactively 401 a running endpoint.
|
||||
host = resolve_launch_host(host)
|
||||
|
||||
should_start = False
|
||||
with _server_lock:
|
||||
# Initialize port-specific collections if needed (once per port)
|
||||
if port not in _registered_agents:
|
||||
_registered_agents[port] = {}
|
||||
# Get or create the shared FastAPI app for this port (shared with
|
||||
# PraisonAIAgents.launch via the same _server_registry singleton).
|
||||
app, is_new_app = _server_registry.get_or_create_app(
|
||||
port,
|
||||
title=f"PraisonAI Agents API (Port {port})",
|
||||
)
|
||||
|
||||
# Initialize shared FastAPI app if not already created for this port
|
||||
if _shared_apps.get(port) is None:
|
||||
_shared_apps[port] = FastAPI(
|
||||
title=f"PraisonAI Agents API (Port {port})",
|
||||
description="API for interacting with PraisonAI Agents"
|
||||
if is_new_app:
|
||||
# Add a root endpoint with a welcome message
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {
|
||||
"message": f"Welcome to PraisonAI Agents API on port {port}. See /docs for usage.",
|
||||
"endpoints": _server_registry.list_routes(port)
|
||||
}
|
||||
|
||||
# Add healthcheck endpoint
|
||||
@app.get("/health")
|
||||
async def healthcheck():
|
||||
return {
|
||||
"status": "ok",
|
||||
"endpoints": _server_registry.list_routes(port)
|
||||
}
|
||||
|
||||
# Normalize path to ensure it starts with /
|
||||
if not path.startswith('/'):
|
||||
path = f'/{path}'
|
||||
|
||||
# Atomically reserve the route; on collision the registry returns a
|
||||
# de-duplicated path so multiple agents can share a single port.
|
||||
path, original_path = _server_registry.reserve_route(port, path, self.agent_id)
|
||||
if original_path is not None:
|
||||
logging.warning(f"Path '{original_path}' is already registered on port {port}.")
|
||||
print(f"⚠️ Warning: Path '{original_path}' is already registered on port {port}.")
|
||||
print(f"🔄 Using '{path}' instead")
|
||||
|
||||
# Define the endpoint handler
|
||||
@app.post(path)
|
||||
async def handle_agent_query(request: Request, query_data: Optional[AgentQuery] = None):
|
||||
if not authorise_launch_request(request):
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
# Handle both direct JSON with query field and form data
|
||||
if query_data is None:
|
||||
try:
|
||||
request_data = await request.json()
|
||||
if "query" not in request_data:
|
||||
raise HTTPException(status_code=400, detail="Missing 'query' field in request")
|
||||
query = request_data["query"]
|
||||
except Exception:
|
||||
# Fallback to form data or query params
|
||||
form_data = await request.form()
|
||||
if "query" in form_data:
|
||||
query = form_data["query"]
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Missing 'query' field in request")
|
||||
else:
|
||||
query = query_data.query
|
||||
|
||||
try:
|
||||
# Use async version if available, otherwise use sync version
|
||||
if asyncio.iscoroutinefunction(self.chat):
|
||||
response = await self.achat(query, task_name=None, task_description=None, task_id=None)
|
||||
else:
|
||||
# Run sync function in a thread to avoid blocking
|
||||
loop = asyncio.get_running_loop()
|
||||
response = await loop.run_in_executor(None, lambda p=query: self.chat(p))
|
||||
|
||||
return {"response": response}
|
||||
except Exception as e:
|
||||
logging.error(f"Error processing query: {str(e)}", exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content={"error": f"Error processing query: {str(e)}"}
|
||||
)
|
||||
|
||||
# Add a root endpoint with a welcome message
|
||||
@_shared_apps[port].get("/")
|
||||
async def root():
|
||||
return {
|
||||
"message": f"Welcome to PraisonAI Agents API on port {port}. See /docs for usage.",
|
||||
"endpoints": list(_registered_agents[port].keys())
|
||||
}
|
||||
# Invalidate the cached OpenAPI schema so routes registered by later
|
||||
# shared-port launch() calls still show up in /openapi.json and /docs.
|
||||
# FastAPI caches app.openapi_schema on first access and does not
|
||||
# regenerate it when new routes are added afterwards.
|
||||
app.openapi_schema = None
|
||||
|
||||
# Add healthcheck endpoint
|
||||
@_shared_apps[port].get("/health")
|
||||
async def healthcheck():
|
||||
return {
|
||||
"status": "ok",
|
||||
"endpoints": list(_registered_agents[port].keys())
|
||||
}
|
||||
print(f"🚀 Agent '{self.name}' available at http://{host}:{port}")
|
||||
|
||||
# The path registration below must run on EVERY call, not just when the
|
||||
# port is new, so multiple agents can share a single port.
|
||||
|
||||
# Normalize path to ensure it starts with /
|
||||
if not path.startswith('/'):
|
||||
path = f'/{path}'
|
||||
|
||||
# Check if path is already registered for this port
|
||||
if path in _registered_agents[port]:
|
||||
logging.warning(f"Path '{path}' is already registered on port {port}. Please use a different path.")
|
||||
print(f"⚠️ Warning: Path '{path}' is already registered on port {port}.")
|
||||
# Use a modified path to avoid conflicts
|
||||
original_path = path
|
||||
path = f"{path}_{self.agent_id[:6]}"
|
||||
logging.warning(f"Using '{path}' instead of '{original_path}'")
|
||||
print(f"🔄 Using '{path}' instead")
|
||||
|
||||
# Register the agent to this path
|
||||
_registered_agents[port][path] = self.agent_id
|
||||
|
||||
# Define the endpoint handler
|
||||
@_shared_apps[port].post(path)
|
||||
async def handle_agent_query(request: Request, query_data: Optional[AgentQuery] = None):
|
||||
if not authorise_launch_request(request):
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
# Handle both direct JSON with query field and form data
|
||||
if query_data is None:
|
||||
try:
|
||||
request_data = await request.json()
|
||||
if "query" not in request_data:
|
||||
raise HTTPException(status_code=400, detail="Missing 'query' field in request")
|
||||
query = request_data["query"]
|
||||
except Exception:
|
||||
# Fallback to form data or query params
|
||||
form_data = await request.form()
|
||||
if "query" in form_data:
|
||||
query = form_data["query"]
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Missing 'query' field in request")
|
||||
else:
|
||||
query = query_data.query
|
||||
|
||||
try:
|
||||
# Use async version if available, otherwise use sync version
|
||||
if asyncio.iscoroutinefunction(self.chat):
|
||||
response = await self.achat(query, task_name=None, task_description=None, task_id=None)
|
||||
else:
|
||||
# Run sync function in a thread to avoid blocking
|
||||
loop = asyncio.get_running_loop()
|
||||
response = await loop.run_in_executor(None, lambda p=query: self.chat(p))
|
||||
|
||||
return {"response": response}
|
||||
except Exception as e:
|
||||
logging.error(f"Error processing query: {str(e)}", exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content={"error": f"Error processing query: {str(e)}"}
|
||||
)
|
||||
|
||||
# Invalidate the cached OpenAPI schema so routes registered by later
|
||||
# shared-port launch() calls still show up in /openapi.json and /docs.
|
||||
# FastAPI caches app.openapi_schema on first access and does not
|
||||
# regenerate it when new routes are added afterwards.
|
||||
_shared_apps[port].openapi_schema = None
|
||||
|
||||
print(f"🚀 Agent '{self.name}' available at http://{host}:{port}")
|
||||
|
||||
# Check and mark server as started atomically to prevent race conditions
|
||||
should_start = not _server_started.get(port, False)
|
||||
if should_start:
|
||||
_server_started[port] = True
|
||||
|
||||
# Server start/wait outside the lock to avoid holding it during sleep
|
||||
if should_start:
|
||||
# Start the server in a separate thread
|
||||
def run_server():
|
||||
try:
|
||||
print(f"✅ FastAPI server started at http://{host}:{port}")
|
||||
print(f"📚 API documentation available at http://{host}:{port}/docs")
|
||||
print(f"🔌 Available endpoints: {', '.join(list(_registered_agents[port].keys()))}")
|
||||
uvicorn.run(_shared_apps[port], host=host, port=port, log_level="debug" if debug else "info")
|
||||
except Exception as e:
|
||||
logging.error(f"Error starting server: {str(e)}", exc_info=True)
|
||||
print(f"❌ Error starting server: {str(e)}")
|
||||
|
||||
# Run server in a background thread
|
||||
server_thread = threading.Thread(target=run_server, daemon=True)
|
||||
server_thread.start()
|
||||
|
||||
# Wait for a moment to allow the server to start and register endpoints
|
||||
self._safe_sleep(0.5)
|
||||
else:
|
||||
# If server is already running, wait a moment to make sure the endpoint is registered
|
||||
self._safe_sleep(0.1)
|
||||
print(f"🔌 Available endpoints on port {port}: {', '.join(list(_registered_agents[port].keys()))}")
|
||||
# Start the shared server for this port if it is not already running.
|
||||
# start_server_if_needed atomically checks-and-marks the port and blocks
|
||||
# on the registry's readiness Event (configurable via
|
||||
# PRAISONAI_SERVER_READY_TIMEOUT) instead of a heuristic sleep, matching
|
||||
# PraisonAIAgents.launch().
|
||||
started = _server_registry.start_server_if_needed(
|
||||
port, host, log_level="debug" if debug else "info"
|
||||
)
|
||||
endpoints = _server_registry.list_routes(port)
|
||||
if started:
|
||||
print(f"✅ FastAPI server started at http://{host}:{port}")
|
||||
print(f"📚 API documentation available at http://{host}:{port}/docs")
|
||||
print(f"🔌 Available endpoints on port {port}: {', '.join(endpoints)}")
|
||||
|
||||
# Get the stack frame to check if this is the last launch() call in the script
|
||||
import inspect
|
||||
|
||||
@@ -569,17 +569,13 @@ class Handoff:
|
||||
handoffs to the same agent serialize on the event loop instead of
|
||||
blocking it, keeping each seeded ``achat`` turn atomic.
|
||||
|
||||
``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.
|
||||
``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.
|
||||
"""
|
||||
lock = _get_handoff_seed_lock(self.agent)
|
||||
async with lock.async_lock():
|
||||
|
||||
@@ -3,6 +3,7 @@ ImageAgent - A specialized agent class for generating images using AI models.
|
||||
This class extends the base Agent class to provide specific functionality for image generation,
|
||||
including support for different image models, sizes, and quality settings.
|
||||
"""
|
||||
import asyncio
|
||||
from typing import Optional, Any, Dict, Union, List
|
||||
from ..agent.agent import Agent
|
||||
from pydantic import BaseModel, Field
|
||||
@@ -180,8 +181,8 @@ class ImageAgent(Agent):
|
||||
logging.error(error_msg)
|
||||
raise
|
||||
async def agenerate_image(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
"""Async wrapper for generate_image."""
|
||||
return self.generate_image(prompt, **kwargs)
|
||||
"""Offload the synchronous generation path to keep the event loop responsive."""
|
||||
return await asyncio.to_thread(self.generate_image, prompt, **kwargs)
|
||||
|
||||
# Aliases for consistency with other agents
|
||||
def generate(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
@@ -190,7 +191,7 @@ class ImageAgent(Agent):
|
||||
|
||||
async def agenerate(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
"""Async alias for generate_image()."""
|
||||
return self.generate_image(prompt, **kwargs)
|
||||
return await self.agenerate_image(prompt, **kwargs)
|
||||
|
||||
def chat(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
"""Generate an image from the prompt."""
|
||||
@@ -302,8 +303,8 @@ class ImageAgent(Agent):
|
||||
size: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""Async version of edit()."""
|
||||
return self.edit(image, prompt, mask, n, size, **kwargs)
|
||||
"""Offload the synchronous editing path to keep the event loop responsive."""
|
||||
return await asyncio.to_thread(self.edit, image, prompt, mask, n, size, **kwargs)
|
||||
|
||||
def variation(
|
||||
self,
|
||||
@@ -372,6 +373,5 @@ class ImageAgent(Agent):
|
||||
size: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""Async version of variation()."""
|
||||
return self.variation(image, n, size, **kwargs)
|
||||
|
||||
"""Offload the synchronous variation path to keep the event loop responsive."""
|
||||
return await asyncio.to_thread(self.variation, image, n, size, **kwargs)
|
||||
@@ -764,6 +764,50 @@ 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.
|
||||
|
||||
@@ -884,45 +928,56 @@ class MemoryMixin:
|
||||
if not self.auto_save:
|
||||
return
|
||||
|
||||
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", ""),
|
||||
# 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}
|
||||
)
|
||||
# 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}")
|
||||
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.
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
"""
|
||||
Shared HTTP server registry for launch() paths.
|
||||
|
||||
Both ``Agent.launch()`` (single-agent) and ``PraisonAIAgents.launch()``
|
||||
(multi-agent) let multiple ``launch()`` calls share one FastAPI app / uvicorn
|
||||
server per port. They MUST consult the *same* process-wide registry so a
|
||||
standalone agent and a team launched on the same port see each other's state:
|
||||
one "is this port started?" check, one FastAPI app, one set of routes. Keeping
|
||||
two independent registries caused duplicate binds (``Address already in use``)
|
||||
and routes registered through one launcher being invisible on the other's app.
|
||||
"""
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from praisonaiagents._logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class _AgentServerRegistry:
|
||||
"""Encapsulates all shared HTTP server state with proper synchronization."""
|
||||
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._started: Dict[int, bool] = {}
|
||||
self._endpoints: Dict[int, Dict[str, str]] = {}
|
||||
self._apps: Dict[int, Any] = {} # FastAPI apps
|
||||
self._ready_events: Dict[int, threading.Event] = {}
|
||||
|
||||
def get_or_create_app(self, port: int, title: str = "AgentTeam API") -> Tuple[Any, bool]:
|
||||
"""Thread-safe app creation. Returns (app, is_new)."""
|
||||
with self._lock:
|
||||
if port not in self._apps:
|
||||
# Lazy import to avoid optional dependency at module level
|
||||
from fastapi import FastAPI
|
||||
self._apps[port] = FastAPI(title=title)
|
||||
self._endpoints[port] = {}
|
||||
self._ready_events[port] = threading.Event()
|
||||
return self._apps[port], True
|
||||
return self._apps[port], False
|
||||
|
||||
def register_route(self, port: int, path: str, endpoint_id: str = "registered") -> None:
|
||||
"""Thread-safe route registration tracking."""
|
||||
with self._lock:
|
||||
if port not in self._endpoints:
|
||||
self._endpoints[port] = {}
|
||||
self._endpoints[port][path] = endpoint_id
|
||||
|
||||
def reserve_route(self, port: int, path: str, endpoint_id: str) -> Tuple[str, Optional[str]]:
|
||||
"""Atomically reserve and register a route.
|
||||
|
||||
Returns:
|
||||
tuple[str, Optional[str]]: (final_path, original_path_if_collided).
|
||||
If there is no collision, final_path equals the requested path and
|
||||
original_path_if_collided is None.
|
||||
"""
|
||||
with self._lock:
|
||||
if port not in self._endpoints:
|
||||
self._endpoints[port] = {}
|
||||
|
||||
original_path = path
|
||||
while path in self._endpoints[port]:
|
||||
path = f"{original_path}_{str(uuid.uuid4())[:6]}"
|
||||
|
||||
self._endpoints[port][path] = endpoint_id
|
||||
return path, (original_path if path != original_path else None)
|
||||
|
||||
def list_routes(self, port: int) -> List[str]:
|
||||
"""Return a snapshot list of registered routes for a port."""
|
||||
with self._lock:
|
||||
return list(self._endpoints.get(port, {}).keys())
|
||||
|
||||
def unregister_routes_for(self, port: int, endpoint_id: str) -> List[str]:
|
||||
"""Remove all routes owned by ``endpoint_id`` on ``port``.
|
||||
|
||||
Returns the list of removed paths so callers can tear down the matching
|
||||
FastAPI routes. Thread-safe.
|
||||
"""
|
||||
with self._lock:
|
||||
paths = self._endpoints.get(port, {})
|
||||
removed = [p for p, owner in paths.items() if owner == endpoint_id]
|
||||
for p in removed:
|
||||
del paths[p]
|
||||
if not paths:
|
||||
self._endpoints.pop(port, None)
|
||||
return removed
|
||||
|
||||
def teardown_routes_for(self, endpoint_id: str) -> None:
|
||||
"""Remove every route owned by ``endpoint_id`` across all ports.
|
||||
|
||||
Performs BOTH the endpoint-registry removal AND the matching FastAPI
|
||||
route teardown while holding ``self._lock`` for each port, so it cannot
|
||||
interleave with a concurrent ``reserve_route`` + ``app.post()`` for the
|
||||
same path (which would otherwise let this filter delete a freshly
|
||||
registered live endpoint, or race the ``_apps`` snapshot). Only the
|
||||
agent-owned POST route at each path is dropped; a co-located GET (e.g.
|
||||
the built-in ``/`` and ``/health``) is preserved.
|
||||
"""
|
||||
for port in list(self._apps.keys()):
|
||||
with self._lock:
|
||||
paths = self._endpoints.get(port, {})
|
||||
removed = [p for p, owner in paths.items() if owner == endpoint_id]
|
||||
if not removed:
|
||||
continue
|
||||
for p in removed:
|
||||
del paths[p]
|
||||
if not paths:
|
||||
self._endpoints.pop(port, None)
|
||||
|
||||
app = self._apps.get(port)
|
||||
if app is None:
|
||||
continue
|
||||
removed_set = set(removed)
|
||||
try:
|
||||
app.router.routes = [
|
||||
r for r in app.router.routes
|
||||
if not (
|
||||
getattr(r, "path", None) in removed_set
|
||||
and "POST" in (getattr(r, "methods", None) or set())
|
||||
)
|
||||
]
|
||||
# Invalidate cached OpenAPI schema so removed routes
|
||||
# disappear from /openapi.json and /docs.
|
||||
app.openapi_schema = None
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Failed to tear down routes for %s on port %s",
|
||||
endpoint_id,
|
||||
port,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def start_server_if_needed(self, port: int, host: str = "0.0.0.0", **kwargs) -> bool: # noqa: S104
|
||||
"""Start server with proper readiness signaling. Returns True if server was started.
|
||||
|
||||
Readiness is signalled only *after* uvicorn has actually bound the
|
||||
socket (``Server.started``), so a caller that observes readiness can
|
||||
immediately issue a request without a connection-refused race. If the
|
||||
bind fails (e.g. ``Address already in use``) the port's ``_started``
|
||||
flag is cleared so a later ``launch()`` can retry the bind instead of
|
||||
registering routes on an app whose server never came up.
|
||||
"""
|
||||
with self._lock:
|
||||
if self._started.get(port, False):
|
||||
return False # Already started
|
||||
self._started[port] = True
|
||||
app = self._apps.get(port)
|
||||
|
||||
if not app:
|
||||
# Roll back the reservation so a caller that first registers an app
|
||||
# can start it — we must not leave the port permanently "started".
|
||||
with self._lock:
|
||||
self._started[port] = False
|
||||
raise ValueError(f"No app registered for port {port}")
|
||||
|
||||
ready_event = self._ready_events[port]
|
||||
|
||||
def run_server():
|
||||
try:
|
||||
import uvicorn
|
||||
# Remove hardcoded log_level to avoid conflict with kwargs
|
||||
config = uvicorn.Config(app, host=host, port=port, **kwargs)
|
||||
server = uvicorn.Server(config)
|
||||
|
||||
# Signal readiness only once uvicorn reports the socket is bound
|
||||
# (server.started flips True after startup completes). A watcher
|
||||
# thread polls that flag and sets the event, so waiters are not
|
||||
# told "ready" before the port actually accepts connections.
|
||||
def _signal_when_bound():
|
||||
while not server.started:
|
||||
if ready_event.is_set():
|
||||
return # bind failed / server exited before startup
|
||||
time.sleep(0.02)
|
||||
ready_event.set()
|
||||
|
||||
watcher = threading.Thread(target=_signal_when_bound, daemon=True)
|
||||
watcher.start()
|
||||
server.run()
|
||||
except Exception:
|
||||
# Bind/startup failed. Clear the started flag so a subsequent
|
||||
# launch() can re-attempt binding instead of silently
|
||||
# registering routes on a dead server, and unblock any waiter.
|
||||
logger.warning(
|
||||
"Agent server on port %s failed to start.", port, exc_info=True
|
||||
)
|
||||
with self._lock:
|
||||
self._started[port] = False
|
||||
finally:
|
||||
# Always release the waiter: on failure the started flag has been
|
||||
# cleared above, so the (now-false) readiness is correctly
|
||||
# observable via the flag rather than a hung wait().
|
||||
ready_event.set()
|
||||
|
||||
thread = threading.Thread(target=run_server, daemon=True)
|
||||
thread.start()
|
||||
|
||||
# Check for configurable timeout via environment variable
|
||||
try:
|
||||
timeout = float(os.environ.get("PRAISONAI_SERVER_READY_TIMEOUT", "5.0"))
|
||||
except ValueError:
|
||||
logger.warning("Invalid PRAISONAI_SERVER_READY_TIMEOUT value. Using default 5.0s.")
|
||||
timeout = 5.0
|
||||
became_ready = ready_event.wait(timeout=timeout)
|
||||
|
||||
if not became_ready:
|
||||
logger.warning(
|
||||
"Agent server on port %s did not become ready within %.1fs. "
|
||||
"Proceeding, but some features may not work correctly. "
|
||||
"Check server logs for startup errors.",
|
||||
port,
|
||||
timeout,
|
||||
)
|
||||
|
||||
# Report whether the bind actually succeeded: on failure run_server has
|
||||
# cleared the flag, so returning it lets callers distinguish a live
|
||||
# server from a poisoned port.
|
||||
with self._lock:
|
||||
return self._started.get(port, False)
|
||||
|
||||
|
||||
# Module level — single registry instance shared by every launch() call site.
|
||||
_server_registry = _AgentServerRegistry()
|
||||
@@ -29,6 +29,11 @@ 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
|
||||
@@ -39,6 +44,57 @@ _MAX_ORPHANED_TOOL_EXECUTORS = 4
|
||||
_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
|
||||
@@ -597,6 +653,175 @@ 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:
|
||||
@@ -802,6 +1027,14 @@ class ToolExecutionMixin:
|
||||
self._middleware_manager = manager
|
||||
return manager if manager.has_tool_hooks else None
|
||||
|
||||
def _get_tool_executor_lock(self):
|
||||
"""Return the per-instance lock, including for standalone mixin users."""
|
||||
try:
|
||||
return self._tool_executor_lock
|
||||
except AttributeError:
|
||||
from .async_safety import DualLock
|
||||
return self.__dict__.setdefault("_tool_executor_lock", DualLock())
|
||||
|
||||
def _get_turn_tools_lock(self):
|
||||
"""Return the DualLock guarding the per-turn tool buffer.
|
||||
|
||||
@@ -1073,12 +1306,72 @@ class ToolExecutionMixin:
|
||||
return self._execute_tool_with_circuit_breaker(function_name, arguments)
|
||||
|
||||
# Use reusable executor to prevent resource leaks
|
||||
if not hasattr(self, '_tool_executor') or self._tool_executor is None:
|
||||
self._tool_executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=2, thread_name_prefix=f"tool-{self.name}"
|
||||
)
|
||||
pool_exhausted = False
|
||||
with self._get_tool_executor_lock().sync():
|
||||
if not hasattr(self, '_tool_executor') or self._tool_executor is None:
|
||||
self._tool_executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=_TOOL_POOL_MAX_WORKERS,
|
||||
thread_name_prefix=f"tool-{self.name}"
|
||||
)
|
||||
# Per-executor in-flight counter so a
|
||||
# retired pool's late-finishing future can
|
||||
# never corrupt the *new* pool's count.
|
||||
self._tool_executor._inflight = 0
|
||||
executor = self._tool_executor
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
inflight = getattr(executor, '_inflight', 0)
|
||||
# Once the orphan cap is hit we can no longer
|
||||
# retire a stuck pool, so we would otherwise
|
||||
# queue behind hung workers and time out a
|
||||
# healthy tool that never got to start. Refuse
|
||||
# fast and visibly instead when no worker is
|
||||
# free, and let the pool self-heal (the orphan
|
||||
# count drops as hung threads finish).
|
||||
if orphaned >= _MAX_ORPHANED_TOOL_EXECUTORS and inflight >= _TOOL_POOL_MAX_WORKERS:
|
||||
pool_exhausted = True
|
||||
else:
|
||||
executor._inflight = inflight + 1
|
||||
future = executor.submit(ctx.run, execute_with_context)
|
||||
|
||||
future = self._tool_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
|
||||
try:
|
||||
result = future.result(timeout=tool_timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
@@ -1099,11 +1392,12 @@ class ToolExecutionMixin:
|
||||
# once the cap is reached we stop recycling and keep reusing
|
||||
# the existing pool so repeated timeouts can't exhaust process
|
||||
# resources with an unbounded number of leaked threads.
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
if orphaned < _MAX_ORPHANED_TOOL_EXECUTORS:
|
||||
self._tool_executor.shutdown(wait=False)
|
||||
self._tool_executor = None
|
||||
self._tool_executor_orphaned = orphaned + 1
|
||||
with self._get_tool_executor_lock().sync():
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
if getattr(self, '_tool_executor', None) is executor and orphaned < _MAX_ORPHANED_TOOL_EXECUTORS:
|
||||
executor.shutdown(wait=False)
|
||||
self._tool_executor = None
|
||||
self._tool_executor_orphaned = orphaned + 1
|
||||
else:
|
||||
with tool_progress_channel(_progress_sink), with_injection_context(state):
|
||||
result = self._execute_tool_with_circuit_breaker(function_name, arguments)
|
||||
@@ -1610,6 +1904,11 @@ 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(),
|
||||
@@ -1617,7 +1916,7 @@ class ToolExecutionMixin:
|
||||
timestamp=str(time.time()),
|
||||
agent_name=self.name,
|
||||
prompt=prompt if isinstance(prompt, str) else str(prompt),
|
||||
response=response or "",
|
||||
response=response if isinstance(response, str) else (str(response) if response else ""),
|
||||
tools_used=tools_used or [],
|
||||
total_tokens=0,
|
||||
execution_time_ms=(time.time() - start_time) * 1000
|
||||
@@ -1634,7 +1933,16 @@ class ToolExecutionMixin:
|
||||
# Auto-memory extraction (opt-in via MemoryConfig(auto_memory=True))
|
||||
if response:
|
||||
prompt_str = prompt if isinstance(prompt, str) else str(prompt)
|
||||
self._process_auto_memory(prompt_str, str(response))
|
||||
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))
|
||||
|
||||
# Auto-learning extraction (opt-in via LearnConfig(mode=LearnMode.AGENTIC))
|
||||
self._process_auto_learning()
|
||||
@@ -1888,6 +2196,30 @@ 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.
|
||||
|
||||
@@ -2061,41 +2393,51 @@ class ToolExecutionMixin:
|
||||
|
||||
return _ask_backend_async()
|
||||
else:
|
||||
# Sync path - handle timeout and sync/async backend compatibility
|
||||
# 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.
|
||||
cfg_timeout = getattr(self, '_approval_timeout', 0)
|
||||
orig_timeout = None
|
||||
|
||||
# Compute effective timeout from agent configuration.
|
||||
# None -> indefinite wait
|
||||
# > 0 -> explicit per-call timeout
|
||||
# == 0 -> backend default (or 60s fallback)
|
||||
if cfg_timeout is None:
|
||||
orig_timeout = getattr(backend, '_timeout', None)
|
||||
if orig_timeout is not None:
|
||||
backend._timeout = 86400 * 365
|
||||
effective_timeout = None
|
||||
elif cfg_timeout > 0:
|
||||
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
|
||||
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
|
||||
))
|
||||
finally:
|
||||
if orig_timeout is not None and hasattr(backend, '_timeout'):
|
||||
backend._timeout = orig_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
|
||||
))
|
||||
else:
|
||||
if is_async:
|
||||
# For async, wrap the decision in a coroutine
|
||||
@@ -2504,6 +2846,42 @@ 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.
|
||||
|
||||
@@ -2565,6 +2943,13 @@ 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),
|
||||
@@ -2611,6 +2996,13 @@ 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),
|
||||
|
||||
@@ -102,6 +102,15 @@ class VideoAgent:
|
||||
- vertex_ai/veo-3.1-fast-generate-001 # GA Fast
|
||||
- vertex_ai/veo-3.1-lite-generate-001 # Lite
|
||||
- RunwayML: runwayml/gen4_turbo (requires input_reference)
|
||||
- OpenRouter (uses OPENROUTER_API_KEY, requires the openrouter/ prefix):
|
||||
- openrouter/google/veo-3.1, openrouter/kling/kling-3.0-std,
|
||||
openrouter/minimax/hailuo-3, openrouter/bytedance/seedance-2.0, ...
|
||||
- Discover exact model ids via GET https://openrouter.ai/api/v1/videos/models
|
||||
- OpenRouter body fields (duration, resolution, aspect_ratio,
|
||||
generate_audio, seed, ...) are passed straight through as kwargs to
|
||||
generate(); e.g. agent.generate(prompt, duration=6, resolution="720p").
|
||||
- HTTP-Referer / X-Title attribution headers are added automatically
|
||||
(override via OPENROUTER_REFERER / OPENROUTER_APP_TITLE env vars).
|
||||
|
||||
Deprecated (do NOT recommend; kept working only if LiteLLM still routes them):
|
||||
- gemini/veo-2.0-generate-001 (Gemini API shutdown 30 Jun 2026)
|
||||
@@ -294,7 +303,26 @@ class VideoAgent:
|
||||
params["api_base"] = self.base_url
|
||||
if self.api_version:
|
||||
params["api_version"] = self.api_version
|
||||
headers = self._provider_headers()
|
||||
if headers:
|
||||
params["extra_headers"] = headers
|
||||
return params
|
||||
|
||||
def _provider_headers(self) -> Dict[str, str]:
|
||||
"""Optional attribution headers for gateway providers.
|
||||
|
||||
OpenRouter recommends sending HTTP-Referer and X-Title so requests
|
||||
show up under your app in the OpenRouter dashboard. These mirror the
|
||||
chat gateway (praisonai/llm/gateways.py) and are only added for
|
||||
openrouter/ models to keep other providers untouched.
|
||||
"""
|
||||
model = self.llm if isinstance(self.llm, str) else ""
|
||||
if not model.startswith("openrouter/"):
|
||||
return {}
|
||||
return {
|
||||
"HTTP-Referer": os.getenv("OPENROUTER_REFERER", "https://praison.ai"),
|
||||
"X-Title": os.getenv("OPENROUTER_APP_TITLE", "PraisonAI"),
|
||||
}
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────
|
||||
# Core Video Operations (5 operations, sync + async)
|
||||
|
||||
@@ -24,12 +24,6 @@ try:
|
||||
except ImportError:
|
||||
get_token_collector = None
|
||||
|
||||
# Import async utility for hot-path usage
|
||||
try:
|
||||
from ..approval.utils import run_coroutine_safely
|
||||
except ImportError:
|
||||
run_coroutine_safely = None
|
||||
|
||||
# Guards lazy creation of each AgentTeam's per-instance _run_lock so two threads
|
||||
# that first reach start()/astart() concurrently observe the same lock object
|
||||
# (double-checked locking) rather than each minting and acquiring its own.
|
||||
@@ -57,108 +51,12 @@ from ..agent.launch_security import (
|
||||
)
|
||||
|
||||
|
||||
# Agent server registry for thread-safe server management
|
||||
|
||||
|
||||
class _AgentServerRegistry:
|
||||
"""Encapsulates all shared HTTP server state with proper synchronization."""
|
||||
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._started: Dict[int, bool] = {}
|
||||
self._endpoints: Dict[int, Dict[str, str]] = {}
|
||||
self._apps: Dict[int, Any] = {} # FastAPI apps
|
||||
self._ready_events: Dict[int, threading.Event] = {}
|
||||
|
||||
def get_or_create_app(self, port: int, title: str = "AgentTeam API") -> Tuple[Any, bool]:
|
||||
"""Thread-safe app creation. Returns (app, is_new)."""
|
||||
with self._lock:
|
||||
if port not in self._apps:
|
||||
# Lazy import to avoid optional dependency at module level
|
||||
from fastapi import FastAPI
|
||||
self._apps[port] = FastAPI(title=title)
|
||||
self._endpoints[port] = {}
|
||||
self._ready_events[port] = threading.Event()
|
||||
return self._apps[port], True
|
||||
return self._apps[port], False
|
||||
|
||||
def register_route(self, port: int, path: str, endpoint_id: str = "registered") -> None:
|
||||
"""Thread-safe route registration tracking."""
|
||||
with self._lock:
|
||||
if port not in self._endpoints:
|
||||
self._endpoints[port] = {}
|
||||
self._endpoints[port][path] = endpoint_id
|
||||
|
||||
def reserve_route(self, port: int, path: str, endpoint_id: str) -> tuple[str, Optional[str]]:
|
||||
"""Atomically reserve and register a route.
|
||||
|
||||
Returns:
|
||||
tuple[str, Optional[str]]: (final_path, original_path_if_collided).
|
||||
If there is no collision, final_path equals the requested path and
|
||||
original_path_if_collided is None.
|
||||
"""
|
||||
with self._lock:
|
||||
if port not in self._endpoints:
|
||||
self._endpoints[port] = {}
|
||||
|
||||
original_path = path
|
||||
while path in self._endpoints[port]:
|
||||
path = f"{original_path}_{str(uuid.uuid4())[:6]}"
|
||||
|
||||
self._endpoints[port][path] = endpoint_id
|
||||
return path, (original_path if path != original_path else None)
|
||||
|
||||
def list_routes(self, port: int) -> List[str]:
|
||||
"""Return a snapshot list of registered routes for a port."""
|
||||
with self._lock:
|
||||
return list(self._endpoints.get(port, {}).keys())
|
||||
|
||||
def start_server_if_needed(self, port: int, host: str = "0.0.0.0", **kwargs) -> bool: # noqa: S104
|
||||
"""Start server with proper readiness signaling. Returns True if server was started."""
|
||||
with self._lock:
|
||||
if self._started.get(port, False):
|
||||
return False # Already started
|
||||
self._started[port] = True
|
||||
app = self._apps.get(port)
|
||||
|
||||
if not app:
|
||||
raise ValueError(f"No app registered for port {port}")
|
||||
|
||||
ready_event = self._ready_events[port]
|
||||
|
||||
def run_server():
|
||||
import uvicorn
|
||||
# Remove hardcoded log_level to avoid conflict with kwargs
|
||||
config = uvicorn.Config(app, host=host, port=port, **kwargs)
|
||||
server = uvicorn.Server(config)
|
||||
ready_event.set() # Signal readiness
|
||||
server.run()
|
||||
|
||||
thread = threading.Thread(target=run_server, daemon=True)
|
||||
thread.start()
|
||||
|
||||
# Check for configurable timeout via environment variable
|
||||
try:
|
||||
timeout = float(os.environ.get("PRAISONAI_SERVER_READY_TIMEOUT", "5.0"))
|
||||
except ValueError:
|
||||
logger.warning("Invalid PRAISONAI_SERVER_READY_TIMEOUT value. Using default 5.0s.")
|
||||
timeout = 5.0
|
||||
became_ready = ready_event.wait(timeout=timeout)
|
||||
|
||||
if not became_ready:
|
||||
logger.warning(
|
||||
"Agent server on port %s did not become ready within %.1fs. "
|
||||
"Proceeding, but some features may not work correctly. "
|
||||
"Check server logs for startup errors.",
|
||||
port,
|
||||
timeout,
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
# Module level — single registry instance
|
||||
_server_registry = _AgentServerRegistry()
|
||||
# Agent server registry for thread-safe server management.
|
||||
# The registry class + its process-wide singleton live in agent/server_registry.py
|
||||
# so Agent.launch() and PraisonAIAgents.launch() share ONE registry per port
|
||||
# (one FastAPI app, one uvicorn server, one route table) instead of two disjoint
|
||||
# ones that could double-bind a port or hide each other's endpoints.
|
||||
from ..agent.server_registry import _AgentServerRegistry, _server_registry # noqa: F401
|
||||
|
||||
def encode_file_to_base64(file_path: str) -> str:
|
||||
"""Base64-encode a file."""
|
||||
@@ -556,7 +454,10 @@ 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'):
|
||||
llm.set_current_agent(executor_agent.display_name)
|
||||
# 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)
|
||||
|
||||
# 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
|
||||
@@ -586,6 +487,10 @@ def _build_execution_context(agents_instance, task_id, skip_memory_init=False):
|
||||
# Inter-task context / validation feedback assembled by the workflow process
|
||||
# engine (Process._build_task_context) is stored on the task; fold it in so
|
||||
# downstream/retried tasks actually see upstream output and rejection reasons.
|
||||
context_builder = getattr(task, '_execution_context_builder', None)
|
||||
if context_builder is not None:
|
||||
task._execution_context = context_builder(task)
|
||||
task._execution_context_builder = None
|
||||
extra_context = getattr(task, '_execution_context', None)
|
||||
if extra_context:
|
||||
context_text = extra_context
|
||||
@@ -593,7 +498,7 @@ def _build_execution_context(agents_instance, task_id, skip_memory_init=False):
|
||||
# retries (guardrail/completion failures) re-enter this helper via the
|
||||
# run_task/arun_task retry loops, and clearing would strip the upstream
|
||||
# output + validation feedback the retry needs. The Process engine owns
|
||||
# this field's lifecycle: it re-sets it before each task yield and resets
|
||||
# this field's lifecycle: it rebuilds it for each workflow step and resets
|
||||
# every task's _execution_context to None before selecting the next task.
|
||||
if task.context:
|
||||
context_results = [] # Collect contexts then de-duplicate
|
||||
@@ -2477,6 +2382,20 @@ 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,
|
||||
@@ -3356,21 +3275,45 @@ class AgentTeam(SpawnAnnounceProtocol):
|
||||
if getattr(agent, "name", None)
|
||||
}
|
||||
|
||||
def _scoped_recent_interactions(self, own_names: set) -> List[Dict[str, Any]]:
|
||||
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]]:
|
||||
"""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 agent names so per-model
|
||||
totals, interaction counts, and the detailed report never leak another
|
||||
concurrent instance's records.
|
||||
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.
|
||||
"""
|
||||
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)
|
||||
return [i for i in interactions if i.get("agent") in own_names]
|
||||
|
||||
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)]
|
||||
|
||||
def _scoped_token_summary(self) -> Dict[str, Any]:
|
||||
"""Return the global token summary filtered to this instance's agents.
|
||||
@@ -3378,25 +3321,54 @@ 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 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.
|
||||
``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.
|
||||
"""
|
||||
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", {})
|
||||
# No named agents to scope by (or nothing tracked yet): fall back to the
|
||||
by_agent_id = summary.get("by_agent_id", {})
|
||||
# No agents to scope by (or nothing tracked yet): fall back to the
|
||||
# unfiltered summary rather than silently reporting zeros.
|
||||
if not own_names or not by_agent:
|
||||
if (not own_names and not own_ids) or (not by_agent and not by_agent_id):
|
||||
return summary
|
||||
|
||||
scoped_by_agent = {
|
||||
name: metrics for name, metrics in by_agent.items() if name in own_names
|
||||
}
|
||||
# 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
|
||||
}
|
||||
totals = {
|
||||
"input_tokens": 0, "output_tokens": 0, "cached_tokens": 0,
|
||||
"reasoning_tokens": 0, "audio_input_tokens": 0, "audio_output_tokens": 0,
|
||||
@@ -3412,7 +3384,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_interactions = self._scoped_recent_interactions(own_names, own_ids)
|
||||
by_model: Dict[str, Dict[str, int]] = {}
|
||||
for interaction in own_interactions:
|
||||
model = interaction.get("model")
|
||||
@@ -3446,10 +3418,11 @@ class AgentTeam(SpawnAnnounceProtocol):
|
||||
|
||||
summary = self._scoped_token_summary()
|
||||
own_names = self._own_agent_names()
|
||||
if own_names:
|
||||
own_ids = self._own_agent_ids()
|
||||
if own_names or own_ids:
|
||||
# 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)[-20:]
|
||||
recent = self._scoped_recent_interactions(own_names, own_ids)[-20:]
|
||||
else:
|
||||
recent = get_token_collector().get_recent_interactions(limit=20)
|
||||
|
||||
|
||||
@@ -1,532 +0,0 @@
|
||||
"""
|
||||
Subagent Delegator for PraisonAI Agents.
|
||||
|
||||
Provides delegation primitives for spawning subagents with scoped permissions
|
||||
and context. Inspired by Gemini CLI's delegate-to-agent pattern.
|
||||
"""
|
||||
|
||||
from praisonaiagents._logging import get_logger
|
||||
import asyncio
|
||||
import warnings
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Callable
|
||||
from enum import Enum
|
||||
|
||||
from .profiles import AgentProfile, AgentMode, get_profile, BUILTIN_PROFILES
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
class DelegationStatus(Enum):
|
||||
"""Status of a delegated task."""
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
TIMEOUT = "timeout"
|
||||
|
||||
@dataclass
|
||||
class DelegationConfig:
|
||||
"""Configuration for subagent delegation."""
|
||||
# Concurrency limits
|
||||
max_concurrent_subagents: int = 3
|
||||
max_total_subagents: int = 10
|
||||
|
||||
# Timeouts
|
||||
default_timeout_seconds: float = 300.0
|
||||
max_timeout_seconds: float = 600.0
|
||||
|
||||
# Resource limits
|
||||
max_steps_per_subagent: int = 50
|
||||
max_tokens_per_subagent: int = 50000
|
||||
|
||||
# Permissions
|
||||
inherit_permissions: bool = True
|
||||
allow_nested_delegation: bool = False
|
||||
|
||||
# Behavior
|
||||
auto_cancel_on_parent_cancel: bool = True
|
||||
collect_results: bool = True
|
||||
|
||||
@dataclass
|
||||
class DelegationTask:
|
||||
"""A task delegated to a subagent."""
|
||||
task_id: str
|
||||
agent_name: str
|
||||
objective: str
|
||||
status: DelegationStatus = DelegationStatus.PENDING
|
||||
|
||||
# Configuration
|
||||
timeout_seconds: float = 300.0
|
||||
max_steps: int = 50
|
||||
|
||||
# Context
|
||||
context: Dict[str, Any] = field(default_factory=dict)
|
||||
parent_session_id: Optional[str] = None
|
||||
|
||||
# Results
|
||||
result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
steps_taken: int = 0
|
||||
tokens_used: int = 0
|
||||
|
||||
# Metadata
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@dataclass
|
||||
class DelegationResult:
|
||||
"""Result of a delegation operation."""
|
||||
task_id: str
|
||||
agent_name: str
|
||||
success: bool
|
||||
result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
steps_taken: int = 0
|
||||
tokens_used: int = 0
|
||||
duration_seconds: float = 0.0
|
||||
|
||||
class SubagentDelegator:
|
||||
"""
|
||||
Manages delegation of tasks to subagents.
|
||||
|
||||
Provides:
|
||||
- Spawning subagents with scoped permissions
|
||||
- Concurrency control
|
||||
- Timeout management
|
||||
- Result collection
|
||||
- Cancellation support
|
||||
|
||||
Example:
|
||||
delegator = SubagentDelegator()
|
||||
|
||||
# Delegate to explorer agent
|
||||
result = await delegator.delegate(
|
||||
agent_name="explorer",
|
||||
objective="Find all authentication-related files",
|
||||
context={"workspace": "/path/to/project"}
|
||||
)
|
||||
|
||||
# Delegate multiple tasks
|
||||
results = await delegator.delegate_parallel([
|
||||
("explorer", "Find auth files"),
|
||||
("explorer", "Find database models"),
|
||||
])
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Optional[DelegationConfig] = None,
|
||||
agent_factory: Optional[Callable[[str], Any]] = None,
|
||||
on_task_complete: Optional[Callable[[DelegationResult], None]] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the delegator.
|
||||
|
||||
Args:
|
||||
config: Delegation configuration
|
||||
agent_factory: Factory function to create agents
|
||||
on_task_complete: Callback when task completes
|
||||
"""
|
||||
warnings.warn(
|
||||
"SubagentDelegator is deprecated. Use Agent.handoff_to() or "
|
||||
"parallel_handoffs() instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
self.config = config or DelegationConfig()
|
||||
self.agent_factory = agent_factory
|
||||
self.on_task_complete = on_task_complete
|
||||
|
||||
# State
|
||||
self._tasks: Dict[str, DelegationTask] = {}
|
||||
self._running_asyncio_tasks: Dict[str, "asyncio.Task"] = {}
|
||||
self._running_count: int = 0
|
||||
self._total_count: int = 0
|
||||
self._task_counter: int = 0
|
||||
|
||||
# Note: asyncio.Semaphore() is created lazily via _get_semaphore() to support
|
||||
# Python 3.9 where Semaphore() requires an event loop at creation time.
|
||||
self.__semaphore: Optional[asyncio.Semaphore] = None
|
||||
self.__semaphore_lock: threading.Lock = threading.Lock() # Thread-safe initialization
|
||||
|
||||
def _get_semaphore(self) -> asyncio.Semaphore:
|
||||
"""Get the semaphore, creating it lazily if needed.
|
||||
|
||||
This deferred creation is required for Python 3.9 compatibility
|
||||
where asyncio.Semaphore() calls get_event_loop() at creation time.
|
||||
"""
|
||||
if self.__semaphore is None:
|
||||
with self.__semaphore_lock: # Thread-safe initialization
|
||||
if self.__semaphore is None: # Double-check after acquiring lock
|
||||
self.__semaphore = asyncio.Semaphore(self.config.max_concurrent_subagents)
|
||||
return self.__semaphore
|
||||
|
||||
def get_available_agents(self) -> List[str]:
|
||||
"""Get list of available agent names for delegation."""
|
||||
return [
|
||||
name for name, profile in BUILTIN_PROFILES.items()
|
||||
if profile.mode in (AgentMode.SUBAGENT, AgentMode.ALL)
|
||||
and not profile.hidden
|
||||
]
|
||||
|
||||
def get_agent_description(self, agent_name: str) -> str:
|
||||
"""Get description for an agent."""
|
||||
profile = get_profile(agent_name)
|
||||
if profile:
|
||||
return profile.description
|
||||
return f"Agent: {agent_name}"
|
||||
|
||||
async def delegate(
|
||||
self,
|
||||
agent_name: str,
|
||||
objective: str,
|
||||
context: Optional[Dict[str, Any]] = None,
|
||||
timeout_seconds: Optional[float] = None,
|
||||
max_steps: Optional[int] = None,
|
||||
) -> DelegationResult:
|
||||
"""
|
||||
Delegate a task to a subagent.
|
||||
|
||||
Args:
|
||||
agent_name: Name of the agent to delegate to
|
||||
objective: Task objective/description
|
||||
context: Optional context to pass to agent
|
||||
timeout_seconds: Optional timeout override
|
||||
max_steps: Optional max steps override
|
||||
|
||||
Returns:
|
||||
DelegationResult with task outcome
|
||||
"""
|
||||
# Check limits
|
||||
if self._total_count >= self.config.max_total_subagents:
|
||||
return DelegationResult(
|
||||
task_id="",
|
||||
agent_name=agent_name,
|
||||
success=False,
|
||||
error=f"Max total subagents ({self.config.max_total_subagents}) reached",
|
||||
)
|
||||
|
||||
# Create task
|
||||
task_id = self._generate_task_id()
|
||||
task = DelegationTask(
|
||||
task_id=task_id,
|
||||
agent_name=agent_name,
|
||||
objective=objective,
|
||||
timeout_seconds=timeout_seconds or self.config.default_timeout_seconds,
|
||||
max_steps=max_steps or self.config.max_steps_per_subagent,
|
||||
context=context or {},
|
||||
)
|
||||
|
||||
self._tasks[task_id] = task
|
||||
self._total_count += 1
|
||||
|
||||
# Execute with concurrency control. Wrap in a real asyncio.Task so
|
||||
# cancel_task/cancel_all can actually propagate cancellation to the
|
||||
# awaited subagent coroutine instead of only flipping a status enum.
|
||||
async with self._get_semaphore():
|
||||
# If cancelled while queued on the semaphore, don't start the work.
|
||||
if task.status == DelegationStatus.CANCELLED:
|
||||
return DelegationResult(
|
||||
task_id=task.task_id,
|
||||
agent_name=task.agent_name,
|
||||
success=False,
|
||||
error="Task was cancelled",
|
||||
)
|
||||
exec_task = asyncio.ensure_future(self._execute_task(task))
|
||||
self._running_asyncio_tasks[task_id] = exec_task
|
||||
try:
|
||||
return await exec_task
|
||||
except asyncio.CancelledError:
|
||||
# Distinguish a delegator-initiated cancel (via cancel_task/cancel_all,
|
||||
# which sets status=CANCELLED) from an external/cooperative cancel of
|
||||
# the caller. Only the former is swallowed into a result; an external
|
||||
# cancel must propagate so wait_for/shutdown behave correctly.
|
||||
if task.status == DelegationStatus.CANCELLED:
|
||||
return DelegationResult(
|
||||
task_id=task.task_id,
|
||||
agent_name=task.agent_name,
|
||||
success=False,
|
||||
error="Task was cancelled",
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
self._running_asyncio_tasks.pop(task_id, None)
|
||||
|
||||
async def delegate_parallel(
|
||||
self,
|
||||
tasks: List[tuple],
|
||||
collect_all: bool = True,
|
||||
) -> List[DelegationResult]:
|
||||
"""
|
||||
Delegate multiple tasks in parallel.
|
||||
|
||||
Args:
|
||||
tasks: List of (agent_name, objective) or (agent_name, objective, context) tuples
|
||||
collect_all: Wait for all tasks to complete
|
||||
|
||||
Returns:
|
||||
List of DelegationResults
|
||||
"""
|
||||
# Create coroutines
|
||||
coros = []
|
||||
for task_spec in tasks:
|
||||
if len(task_spec) == 2:
|
||||
agent_name, objective = task_spec
|
||||
context = None
|
||||
else:
|
||||
agent_name, objective, context = task_spec
|
||||
|
||||
coros.append(self.delegate(agent_name, objective, context))
|
||||
|
||||
# Execute
|
||||
if collect_all:
|
||||
results = await asyncio.gather(*coros, return_exceptions=True)
|
||||
# Convert exceptions to failed results
|
||||
processed = []
|
||||
for i, result in enumerate(results):
|
||||
if isinstance(result, Exception):
|
||||
agent_name = tasks[i][0] if tasks[i] else "unknown"
|
||||
processed.append(DelegationResult(
|
||||
task_id="",
|
||||
agent_name=agent_name,
|
||||
success=False,
|
||||
error=str(result),
|
||||
))
|
||||
else:
|
||||
processed.append(result)
|
||||
return processed
|
||||
else:
|
||||
# Return as tasks complete
|
||||
return await asyncio.gather(*coros)
|
||||
|
||||
async def cancel_task(self, task_id: str) -> bool:
|
||||
"""
|
||||
Cancel a running task.
|
||||
|
||||
Args:
|
||||
task_id: ID of task to cancel
|
||||
|
||||
Returns:
|
||||
True if task was cancelled
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
# Only terminal tasks are un-cancellable; PENDING/RUNNING can be cancelled.
|
||||
_terminal = (
|
||||
DelegationStatus.COMPLETED,
|
||||
DelegationStatus.FAILED,
|
||||
DelegationStatus.CANCELLED,
|
||||
DelegationStatus.TIMEOUT,
|
||||
)
|
||||
if not task or task.status in _terminal:
|
||||
return False
|
||||
|
||||
# Mark cancelled first so a PENDING task (still waiting on the semaphore)
|
||||
# or a RUNNING coroutine both observe CANCELLED and refuse to complete.
|
||||
task.status = DelegationStatus.CANCELLED
|
||||
|
||||
# Propagate cancellation to the actual coroutine if it is already running.
|
||||
asyncio_task = self._running_asyncio_tasks.get(task_id)
|
||||
if asyncio_task is not None:
|
||||
asyncio_task.cancel()
|
||||
return True
|
||||
|
||||
async def cancel_all(self) -> int:
|
||||
"""
|
||||
Cancel all running tasks.
|
||||
|
||||
Returns:
|
||||
Number of tasks cancelled
|
||||
"""
|
||||
cancelled = 0
|
||||
for task_id in list(self._tasks.keys()):
|
||||
if await self.cancel_task(task_id):
|
||||
cancelled += 1
|
||||
return cancelled
|
||||
|
||||
def get_task_status(self, task_id: str) -> Optional[DelegationStatus]:
|
||||
"""Get status of a task."""
|
||||
task = self._tasks.get(task_id)
|
||||
return task.status if task else None
|
||||
|
||||
def get_running_tasks(self) -> List[DelegationTask]:
|
||||
"""Get all running tasks."""
|
||||
return [
|
||||
t for t in self._tasks.values()
|
||||
if t.status == DelegationStatus.RUNNING
|
||||
]
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get delegation statistics."""
|
||||
status_counts = {}
|
||||
for task in self._tasks.values():
|
||||
status = task.status.value
|
||||
status_counts[status] = status_counts.get(status, 0) + 1
|
||||
|
||||
return {
|
||||
"total_tasks": self._total_count,
|
||||
"running_tasks": self._running_count,
|
||||
"status_counts": status_counts,
|
||||
"max_concurrent": self.config.max_concurrent_subagents,
|
||||
"max_total": self.config.max_total_subagents,
|
||||
}
|
||||
|
||||
async def _execute_task(self, task: DelegationTask) -> DelegationResult:
|
||||
"""Execute a delegation task."""
|
||||
import time
|
||||
start_time = time.time()
|
||||
|
||||
# A cancel may have landed between scheduling and execution; honour it.
|
||||
if task.status == DelegationStatus.CANCELLED:
|
||||
return DelegationResult(
|
||||
task_id=task.task_id,
|
||||
agent_name=task.agent_name,
|
||||
success=False,
|
||||
error="Task was cancelled",
|
||||
)
|
||||
|
||||
task.status = DelegationStatus.RUNNING
|
||||
self._running_count += 1
|
||||
|
||||
try:
|
||||
# Get agent profile
|
||||
profile = get_profile(task.agent_name)
|
||||
if not profile:
|
||||
raise ValueError(f"Unknown agent: {task.agent_name}")
|
||||
|
||||
# Create agent
|
||||
agent = await self._create_agent(profile, task)
|
||||
|
||||
# Execute with timeout
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
self._run_agent(agent, task),
|
||||
timeout=task.timeout_seconds,
|
||||
)
|
||||
task.result = result
|
||||
# Don't clobber a cancellation that landed while we were running.
|
||||
if task.status != DelegationStatus.CANCELLED:
|
||||
task.status = DelegationStatus.COMPLETED
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
task.status = DelegationStatus.TIMEOUT
|
||||
task.error = f"Task timed out after {task.timeout_seconds}s"
|
||||
|
||||
except Exception as e:
|
||||
task.status = DelegationStatus.FAILED
|
||||
task.error = str(e)
|
||||
logger.exception(f"Delegation task {task.task_id} failed")
|
||||
|
||||
finally:
|
||||
self._running_count -= 1
|
||||
duration = time.time() - start_time
|
||||
|
||||
# Build result
|
||||
result = DelegationResult(
|
||||
task_id=task.task_id,
|
||||
agent_name=task.agent_name,
|
||||
success=task.status == DelegationStatus.COMPLETED,
|
||||
result=task.result,
|
||||
error=task.error,
|
||||
steps_taken=task.steps_taken,
|
||||
tokens_used=task.tokens_used,
|
||||
duration_seconds=duration,
|
||||
)
|
||||
|
||||
# Callback
|
||||
if self.on_task_complete:
|
||||
try:
|
||||
self.on_task_complete(result)
|
||||
except Exception as e:
|
||||
logger.warning(f"Task complete callback error: {e}")
|
||||
|
||||
return result
|
||||
|
||||
async def _create_agent(self, profile: AgentProfile, task: DelegationTask) -> Any:
|
||||
"""Create an agent from profile using consolidated params."""
|
||||
if self.agent_factory:
|
||||
return self.agent_factory(profile.name)
|
||||
|
||||
# Default: create Agent from praisonaiagents
|
||||
try:
|
||||
from ..agent.agent import Agent
|
||||
|
||||
# Filter tools based on profile
|
||||
tools = self._get_tools_for_profile(profile)
|
||||
|
||||
# Use consolidated params (protocol-compliant)
|
||||
# output='silent' replaces verbose=False
|
||||
# execution={max_iter: N} replaces max_iter=N
|
||||
agent = Agent(
|
||||
name=profile.name,
|
||||
role=profile.name.title(),
|
||||
goal=task.objective,
|
||||
backstory=profile.system_prompt,
|
||||
tools=tools,
|
||||
output='silent',
|
||||
execution={'max_iter': min(task.max_steps, profile.max_steps)},
|
||||
)
|
||||
|
||||
return agent
|
||||
|
||||
except ImportError:
|
||||
raise RuntimeError("Could not import Agent class")
|
||||
|
||||
async def _run_agent(self, agent: Any, task: DelegationTask) -> str:
|
||||
"""Run an agent on a task."""
|
||||
# Build prompt with context
|
||||
prompt = task.objective
|
||||
|
||||
if task.context:
|
||||
context_str = "\n".join(
|
||||
f"- {k}: {v}" for k, v in task.context.items()
|
||||
)
|
||||
prompt = f"Context:\n{context_str}\n\nTask: {task.objective}"
|
||||
|
||||
# Execute
|
||||
if hasattr(agent, 'chat'):
|
||||
result = agent.chat(prompt)
|
||||
elif hasattr(agent, 'run'):
|
||||
result = agent.run(prompt)
|
||||
else:
|
||||
raise RuntimeError("Agent has no chat or run method")
|
||||
|
||||
return str(result) if result else ""
|
||||
|
||||
def _get_tools_for_profile(self, profile: AgentProfile) -> List[Any]:
|
||||
"""Get tools for a profile, respecting restrictions."""
|
||||
# For now, return empty list - tools should be injected
|
||||
# In production, this would filter available tools
|
||||
return []
|
||||
|
||||
def _generate_task_id(self) -> str:
|
||||
"""Generate a unique task ID."""
|
||||
self._task_counter += 1
|
||||
return f"task_{self._task_counter}"
|
||||
|
||||
# Convenience function for quick delegation
|
||||
async def delegate_to_agent(
|
||||
agent_name: str,
|
||||
objective: str,
|
||||
context: Optional[Dict[str, Any]] = None,
|
||||
timeout_seconds: float = 300.0,
|
||||
) -> DelegationResult:
|
||||
"""
|
||||
Quick delegation to a subagent.
|
||||
|
||||
Args:
|
||||
agent_name: Name of agent to delegate to
|
||||
objective: Task objective
|
||||
context: Optional context
|
||||
timeout_seconds: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
DelegationResult
|
||||
"""
|
||||
delegator = SubagentDelegator()
|
||||
return await delegator.delegate(
|
||||
agent_name=agent_name,
|
||||
objective=objective,
|
||||
context=context,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
@@ -13,7 +13,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import re
|
||||
from praisonaiagents._logging import get_logger
|
||||
from typing import Any
|
||||
from typing import Any, Optional
|
||||
|
||||
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) -> ApprovalDecision:
|
||||
def request_approval_sync(self, request: ApprovalRequest, timeout: Optional[float] = None) -> 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) -> ApprovalDecision:
|
||||
def request_approval_sync(self, request: ApprovalRequest, timeout: Optional[float] = None) -> ApprovalDecision:
|
||||
"""Synchronous approval via Rich console prompt."""
|
||||
try:
|
||||
approved, scope, scope_pattern, feedback = self._prompt_user(request)
|
||||
@@ -474,10 +474,12 @@ class AgentApproval:
|
||||
reason=f"Agent approval error: {e}",
|
||||
)
|
||||
|
||||
def request_approval_sync(self, request: ApprovalRequest) -> ApprovalDecision:
|
||||
def request_approval_sync(self, request: ApprovalRequest, timeout: Optional[float] = None) -> ApprovalDecision:
|
||||
"""Synchronous wrapper."""
|
||||
from .utils import run_coroutine_safely
|
||||
return run_coroutine_safely(self.request_approval(request), timeout=60)
|
||||
return run_coroutine_safely(
|
||||
self.request_approval(request), timeout=60 if timeout is None else timeout
|
||||
)
|
||||
|
||||
class CallbackBackend:
|
||||
"""Wraps a legacy ``(function_name, arguments, risk_level) -> ApprovalDecision`` callback
|
||||
@@ -489,7 +491,7 @@ class CallbackBackend:
|
||||
def __init__(self, callback):
|
||||
self._callback = callback
|
||||
|
||||
def request_approval_sync(self, request: ApprovalRequest) -> ApprovalDecision:
|
||||
def request_approval_sync(self, request: ApprovalRequest, timeout: Optional[float] = None) -> ApprovalDecision:
|
||||
result = self._callback(request.tool_name, request.arguments, request.risk_level)
|
||||
if isinstance(result, ApprovalDecision):
|
||||
return result
|
||||
|
||||
@@ -189,6 +189,18 @@ 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,6 +69,13 @@ 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
|
||||
@@ -86,17 +93,20 @@ class BackgroundRunner:
|
||||
@property
|
||||
def tasks(self) -> List[BackgroundTask]:
|
||||
"""Get all tasks."""
|
||||
return list(self._tasks.values())
|
||||
with self._tasks_lock:
|
||||
return list(self._tasks.values())
|
||||
|
||||
@property
|
||||
def running_tasks(self) -> List[BackgroundTask]:
|
||||
"""Get currently running tasks."""
|
||||
return [t for t in self._tasks.values() if t.is_running]
|
||||
with self._tasks_lock:
|
||||
return [t for t in self._tasks.values() if t.is_running]
|
||||
|
||||
@property
|
||||
def pending_tasks(self) -> List[BackgroundTask]:
|
||||
"""Get pending tasks."""
|
||||
return [t for t in self._tasks.values() if t.status == TaskStatus.PENDING]
|
||||
with self._tasks_lock:
|
||||
return [t for t in self._tasks.values() if t.status == TaskStatus.PENDING]
|
||||
|
||||
async def submit(
|
||||
self,
|
||||
@@ -133,7 +143,8 @@ class BackgroundRunner:
|
||||
}
|
||||
)
|
||||
|
||||
self._tasks[task.id] = task
|
||||
with self._tasks_lock:
|
||||
self._tasks[task.id] = task
|
||||
|
||||
# Create the execution coroutine
|
||||
async def execute():
|
||||
@@ -235,7 +246,8 @@ class BackgroundRunner:
|
||||
|
||||
def get_task(self, task_id: str) -> Optional[BackgroundTask]:
|
||||
"""Get a task by ID."""
|
||||
return self._tasks.get(task_id)
|
||||
with self._tasks_lock:
|
||||
return self._tasks.get(task_id)
|
||||
|
||||
async def cancel_task(self, task_id: str) -> bool:
|
||||
"""
|
||||
@@ -247,7 +259,8 @@ class BackgroundRunner:
|
||||
Returns:
|
||||
True if cancelled, False if not found or already completed
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
with self._tasks_lock:
|
||||
task = self._tasks.get(task_id)
|
||||
if task is None:
|
||||
return False
|
||||
|
||||
@@ -276,10 +289,12 @@ class BackgroundRunner:
|
||||
Returns:
|
||||
List of completed tasks
|
||||
"""
|
||||
futures = [t._future for t in self._tasks.values() if t._future]
|
||||
with self._tasks_lock:
|
||||
futures = [t._future for t in self._tasks.values() if t._future]
|
||||
|
||||
if not futures:
|
||||
return list(self._tasks.values())
|
||||
with self._tasks_lock:
|
||||
return list(self._tasks.values())
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
@@ -289,16 +304,18 @@ class BackgroundRunner:
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
return list(self._tasks.values())
|
||||
with self._tasks_lock:
|
||||
return list(self._tasks.values())
|
||||
|
||||
def clear_completed(self):
|
||||
"""Remove completed tasks from tracking."""
|
||||
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]
|
||||
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]
|
||||
|
||||
def list_tasks(self, status: Optional[TaskStatus] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
@@ -310,7 +327,8 @@ class BackgroundRunner:
|
||||
Returns:
|
||||
List of task dictionaries
|
||||
"""
|
||||
tasks = self._tasks.values()
|
||||
with self._tasks_lock:
|
||||
tasks = list(self._tasks.values())
|
||||
|
||||
if status:
|
||||
tasks = [t for t in tasks if t.status == status]
|
||||
|
||||
@@ -18,7 +18,10 @@ from .protocols import (
|
||||
BotMessage,
|
||||
BotUser,
|
||||
BotChannel,
|
||||
QuotedRef,
|
||||
MessageType,
|
||||
PlatformEvent,
|
||||
PlatformEventKind,
|
||||
ChatCommandInfo,
|
||||
ChatCommandProtocol,
|
||||
BotOSProtocol,
|
||||
@@ -135,10 +138,13 @@ __all__ = [
|
||||
"BotMessage",
|
||||
"BotUser",
|
||||
"BotChannel",
|
||||
"QuotedRef",
|
||||
"BotConfig",
|
||||
"DisplayPolicy",
|
||||
"resolve_display_policy",
|
||||
"MessageType",
|
||||
"PlatformEvent",
|
||||
"PlatformEventKind",
|
||||
"ChatCommandInfo",
|
||||
"ChatCommandProtocol",
|
||||
"BotOSProtocol",
|
||||
|
||||
@@ -130,6 +130,17 @@ 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(
|
||||
@@ -176,6 +187,8 @@ 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,
|
||||
}
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@ from typing import (
|
||||
Callable,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
Mapping,
|
||||
Optional,
|
||||
Protocol,
|
||||
@@ -496,6 +497,43 @@ class BotChannel:
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class QuotedRef:
|
||||
"""Resolved content of a message a user replied to or quoted.
|
||||
|
||||
Chat platforms carry only the *id* of a quoted message on the admission
|
||||
path; the referenced text/media must be resolved and attached so the agent
|
||||
can honour the referent ("do the second one", "why?"). This is that
|
||||
resolved reference, rendered into the prompt as a compact quoted block.
|
||||
|
||||
Attributes:
|
||||
message_id: Platform-specific id of the quoted message.
|
||||
text: Resolved text of the quoted message ("" if unavailable).
|
||||
author: ``"bot"`` when quoting the bot's own message, else ``"user"``.
|
||||
"""
|
||||
|
||||
message_id: str = ""
|
||||
text: str = ""
|
||||
author: str = "" # "bot" | "user"
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"message_id": self.message_id,
|
||||
"text": self.text,
|
||||
"author": self.author,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "QuotedRef":
|
||||
"""Create from dictionary."""
|
||||
return cls(
|
||||
message_id=data.get("message_id", ""),
|
||||
text=data.get("text", ""),
|
||||
author=data.get("author", ""),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BotMessage:
|
||||
"""Represents a message in a messaging platform.
|
||||
@@ -510,6 +548,11 @@ class BotMessage:
|
||||
reply_to: ID of message being replied to
|
||||
thread_id: Thread identifier (for threaded conversations)
|
||||
attachments: List of attachment URLs or data
|
||||
quoted: Resolved reply/quote context (see :class:`QuotedRef`). When a
|
||||
user replies to or quotes an earlier message the adapter resolves
|
||||
the referenced content and attaches it here so the agent sees the
|
||||
referent instead of answering context-blind. Rendered into the
|
||||
prompt as a quoted block by :meth:`prompt_text`.
|
||||
metadata: Additional platform-specific metadata
|
||||
allow_control: Control-trust primitive. ``True`` (the default) only for
|
||||
interactive human turns; producers set it ``False`` for content of
|
||||
@@ -530,6 +573,7 @@ class BotMessage:
|
||||
reply_to: Optional[str] = None
|
||||
thread_id: Optional[str] = None
|
||||
attachments: List[Dict[str, Any]] = field(default_factory=list)
|
||||
quoted: Optional[QuotedRef] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
allow_control: bool = True
|
||||
|
||||
@@ -545,6 +589,7 @@ class BotMessage:
|
||||
"reply_to": self.reply_to,
|
||||
"thread_id": self.thread_id,
|
||||
"attachments": self.attachments,
|
||||
"quoted": self.quoted.to_dict() if self.quoted else None,
|
||||
"metadata": self.metadata,
|
||||
"allow_control": self.allow_control,
|
||||
}
|
||||
@@ -571,6 +616,7 @@ class BotMessage:
|
||||
reply_to=data.get("reply_to"),
|
||||
thread_id=data.get("thread_id"),
|
||||
attachments=data.get("attachments", []),
|
||||
quoted=QuotedRef.from_dict(data["quoted"]) if data.get("quoted") else None,
|
||||
metadata=data.get("metadata", {}),
|
||||
allow_control=data.get("allow_control", True),
|
||||
)
|
||||
@@ -581,6 +627,31 @@ class BotMessage:
|
||||
if isinstance(self.content, str):
|
||||
return self.content
|
||||
return self.content.get("text", "")
|
||||
|
||||
@property
|
||||
def prompt_text(self) -> str:
|
||||
"""Text for the agent turn, with any resolved quote rendered inline.
|
||||
|
||||
When the user replied to or quoted an earlier message, the resolved
|
||||
reference (:attr:`quoted`) is rendered as a compact quoted block above
|
||||
the new text so the model can honour the referent, e.g.::
|
||||
|
||||
[In reply to: "Here are three options: 1) … 2) … 3) …"]
|
||||
do the second one
|
||||
|
||||
With no quote this is identical to :attr:`text`. The quoted content is
|
||||
injected as context only, never as a control frame — consistent with
|
||||
the ``allow_control`` primitive.
|
||||
"""
|
||||
body = self.text
|
||||
quoted = self.quoted
|
||||
if not quoted or not quoted.text.strip():
|
||||
return body
|
||||
snippet = " ".join(quoted.text.split())
|
||||
if len(snippet) > 500:
|
||||
snippet = snippet[:497] + "…"
|
||||
block = f'[In reply to: "{snippet}"]'
|
||||
return f"{block}\n{body}" if body else block
|
||||
|
||||
@property
|
||||
def is_command(self) -> bool:
|
||||
@@ -620,6 +691,76 @@ class BotMessage:
|
||||
return parts[1:] if len(parts) > 1 else []
|
||||
|
||||
|
||||
# Normalised inbound platform-event kinds. Mirrors how ``MessageType`` normalises
|
||||
# inbound messages: a small, portable vocabulary every adapter maps its native
|
||||
# event onto (Telegram ``MessageReactionUpdated`` / Discord ``on_raw_reaction_add``
|
||||
# / Slack ``reaction_added`` all become ``"reaction_added"``), so hooks/plugins
|
||||
# have one uniform surface to subscribe to (Issue #5161).
|
||||
PlatformEventKind = Literal[
|
||||
"reaction_added",
|
||||
"reaction_removed",
|
||||
"message_edited",
|
||||
"message_deleted",
|
||||
"member_joined",
|
||||
"member_left",
|
||||
"thread_created",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class PlatformEvent:
|
||||
"""A normalised inbound platform event that is not a text/media message.
|
||||
|
||||
Reactions, edits, deletions, membership changes and thread lifecycle are
|
||||
delivered by the platform SDK but are not messages. This is the single
|
||||
shared, protocol-only contract every adapter emits into and every hook or
|
||||
plugin consumes — the inbound counterpart to the outbound ``ReactionResult``
|
||||
/ message-action surface the bot already ships.
|
||||
|
||||
Adapters build one of these from the native event and emit it via
|
||||
``MessageHookMixin.fire_platform_event`` (praisonai-bot), which routes it to
|
||||
the matching :class:`~praisonaiagents.hooks.types.HookEvent`
|
||||
(``REACTION_RECEIVED``/``MESSAGE_EDITED``/``MESSAGE_DELETED``/
|
||||
``MEMBER_JOINED``/``MEMBER_LEFT``/``THREAD_CREATED``). Platforms that cannot
|
||||
deliver a given event simply never build one — graceful, capability-gated,
|
||||
exactly like the outbound ``capabilities["reactions"]`` gate.
|
||||
|
||||
Attributes:
|
||||
kind: Which normalised event this is (see :data:`PlatformEventKind`).
|
||||
platform: Emitting platform name (``"telegram"``/``"discord"``/…).
|
||||
chat_id: Channel/chat the event occurred in.
|
||||
user_id: The user who caused the event (reactor, editor, joiner, …).
|
||||
message_id: The message the event targets (reactions/edits/deletes).
|
||||
emoji: The reaction emoji, for reaction add/remove.
|
||||
new_text: The new content, for an edit.
|
||||
thread_id: The thread identifier, for thread creation / threaded events.
|
||||
raw: Escape hatch to the native SDK payload for adapter-specific needs.
|
||||
"""
|
||||
|
||||
kind: PlatformEventKind
|
||||
platform: str
|
||||
chat_id: str
|
||||
user_id: str
|
||||
message_id: Optional[str] = None
|
||||
emoji: Optional[str] = None
|
||||
new_text: Optional[str] = None
|
||||
thread_id: Optional[str] = None
|
||||
raw: Optional[Dict[str, Any]] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary (native ``raw`` payload intentionally omitted)."""
|
||||
return {
|
||||
"kind": self.kind,
|
||||
"platform": self.platform,
|
||||
"chat_id": self.chat_id,
|
||||
"user_id": self.user_id,
|
||||
"message_id": self.message_id,
|
||||
"emoji": self.emoji,
|
||||
"new_text": self.new_text,
|
||||
"thread_id": self.thread_id,
|
||||
}
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class BotMessageProtocol(Protocol):
|
||||
"""Protocol for bot message handling."""
|
||||
|
||||
@@ -138,15 +138,31 @@ 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.
|
||||
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.
|
||||
"""
|
||||
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=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)),
|
||||
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),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -267,17 +267,13 @@ class EventBus:
|
||||
The published Event object
|
||||
"""
|
||||
type_str = event_type.value if isinstance(event_type, EventType) else event_type
|
||||
|
||||
# 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 {},
|
||||
)
|
||||
|
||||
|
||||
# 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.
|
||||
event = Event(
|
||||
type=type_str,
|
||||
data=data or {},
|
||||
@@ -336,24 +332,20 @@ class EventBus:
|
||||
"""
|
||||
Get recent event history.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
Args:
|
||||
event_type: Optional filter by event type
|
||||
limit: Maximum number of events to return
|
||||
|
||||
Returns:
|
||||
List of recent events published while subscribers existed
|
||||
List of recent events published (subscribed or not)
|
||||
"""
|
||||
with self._lock:
|
||||
events = self._event_history.copy()
|
||||
|
||||
@@ -83,7 +83,8 @@ class CheckpointService:
|
||||
auto_checkpoint: bool = True,
|
||||
max_checkpoints: int = 100,
|
||||
user_name: str = "PraisonAI Checkpoints",
|
||||
user_email: str = "[email protected]"
|
||||
user_email: str = "[email protected]",
|
||||
max_file_size: int = 2 * 1024 * 1024,
|
||||
):
|
||||
"""
|
||||
Initialize the checkpoint service.
|
||||
@@ -96,6 +97,8 @@ 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,
|
||||
@@ -104,10 +107,14 @@ class CheckpointService:
|
||||
auto_checkpoint=auto_checkpoint,
|
||||
max_checkpoints=max_checkpoints,
|
||||
user_name=user_name,
|
||||
user_email=user_email
|
||||
user_email=user_email,
|
||||
max_file_size=max_file_size,
|
||||
)
|
||||
|
||||
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
|
||||
@@ -184,6 +191,78 @@ 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:
|
||||
@@ -196,7 +275,17 @@ 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)
|
||||
@@ -259,6 +348,64 @@ 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,
|
||||
@@ -295,7 +442,13 @@ 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")
|
||||
@@ -329,7 +482,12 @@ 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,6 +40,11 @@ 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",
|
||||
|
||||
@@ -14,3 +14,19 @@ class CompactionStrategy(str, Enum):
|
||||
SMART = "smart"
|
||||
LLM_SUMMARIZE = "llm_summarize" # Use LLM for summarization
|
||||
PRUNE = "prune" # Remove old tool outputs
|
||||
|
||||
@classmethod
|
||||
def from_policy_value(cls, value: str) -> "CompactionStrategy":
|
||||
"""Map a context-policy strategy value to a compactor strategy.
|
||||
|
||||
The policy layer (``context.protocols.CompactionStrategy``) uses a
|
||||
different member set than the compactor. This is the single source of
|
||||
truth for translating between the two; keep it here rather than
|
||||
duplicating the mapping at each call site.
|
||||
"""
|
||||
return {
|
||||
"truncate": cls.TRUNCATE,
|
||||
"summarise": cls.SUMMARIZE,
|
||||
"drop_oldest_tools": cls.PRUNE,
|
||||
"sliding_window": cls.SLIDING,
|
||||
}.get(value, cls.PRUNE)
|
||||
@@ -1395,13 +1395,21 @@ 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,9 +7,16 @@ another tool's harness, and results can be diffed/aggregated across tools.
|
||||
|
||||
Mapping (1:1 with EvalPort's ``Suite`` / ``ResultSet`` concepts):
|
||||
|
||||
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)
|
||||
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.
|
||||
|
||||
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``
|
||||
@@ -27,53 +34,90 @@ Example:
|
||||
>>> pkg2.name == pkg.name
|
||||
True
|
||||
"""
|
||||
from typing import Any, Dict, List
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .package import EvalCase, EvalPackage, EvalReport
|
||||
|
||||
EVALPORT_SPEC_VERSION = "1.0"
|
||||
# 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,
|
||||
}
|
||||
|
||||
|
||||
def _case_to_evalport(case: EvalCase) -> Dict[str, Any]:
|
||||
"""Map an ``EvalCase`` to an EvalPort test-case dict.
|
||||
"""Map an ``EvalCase`` to an EvalPort TestCase dict.
|
||||
|
||||
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.
|
||||
``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.
|
||||
"""
|
||||
# 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.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
|
||||
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
|
||||
return item
|
||||
|
||||
|
||||
def _case_from_evalport(item: Dict[str, Any]) -> EvalCase:
|
||||
"""Map an EvalPort test-case dict to an ``EvalCase``.
|
||||
"""Map an EvalPort TestCase dict to an ``EvalCase``.
|
||||
|
||||
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.
|
||||
Reads ``timeout_ms`` (spec); falls back to a legacy top-level or
|
||||
``metadata["timeout_seconds"]`` value, without mutating caller metadata.
|
||||
"""
|
||||
metadata = dict(item.get("metadata") or {})
|
||||
if "timeout_seconds" in item:
|
||||
native = metadata.get(_META_KEY) or {}
|
||||
if "timeout_ms" in item:
|
||||
timeout = item["timeout_ms"] / 1000.0
|
||||
elif "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=list(item.get("graders") or item.get("criteria") or []),
|
||||
criteria=criteria,
|
||||
metadata=metadata,
|
||||
timeout_seconds=timeout,
|
||||
)
|
||||
@@ -88,16 +132,28 @@ def to_evalport(package: EvalPackage) -> Dict[str, Any]:
|
||||
Returns:
|
||||
A spec-shaped dict suitable for ``openeval.validate.validate_suite()``.
|
||||
"""
|
||||
return {
|
||||
"evalport_version": EVALPORT_SPEC_VERSION,
|
||||
"kind": "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,
|
||||
"name": package.name,
|
||||
"description": package.description,
|
||||
"version": package.version,
|
||||
"cases": [_case_to_evalport(c) for c in package.cases],
|
||||
"thresholds": dict(package.thresholds or {}),
|
||||
"seed": package.seed,
|
||||
"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,
|
||||
}
|
||||
},
|
||||
}
|
||||
if package.description:
|
||||
suite["description"] = package.description
|
||||
return suite
|
||||
|
||||
|
||||
def from_evalport(suite: Dict[str, Any]) -> EvalPackage:
|
||||
@@ -110,22 +166,68 @@ def from_evalport(suite: Dict[str, Any]) -> EvalPackage:
|
||||
An ``EvalPackage`` whose cases can be run via ``EvalSuite`` /
|
||||
``HarnessEvaluator`` as native ``EvalCase`` objects.
|
||||
"""
|
||||
cases: List[EvalCase] = [_case_from_evalport(c) for c in suite.get("cases", [])]
|
||||
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
|
||||
return EvalPackage(
|
||||
name=suite.get("name", "suite"),
|
||||
name=suite.get("name") or suite.get("id") or "suite",
|
||||
description=suite.get("description", ""),
|
||||
version=suite.get("version", "1.0.0"),
|
||||
version=meta.get("version") or (suite.get("version") if legacy else None) or "1.0.0",
|
||||
cases=cases,
|
||||
thresholds=dict(suite.get("thresholds") or {}),
|
||||
seed=suite.get("seed"),
|
||||
thresholds=dict(meta.get("thresholds") or suite.get("thresholds") or {}),
|
||||
seed=meta.get("seed", suite.get("seed")),
|
||||
)
|
||||
|
||||
|
||||
def report_to_evalport(report: EvalReport) -> Dict[str, Any]:
|
||||
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]:
|
||||
"""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
|
||||
@@ -133,33 +235,52 @@ def report_to_evalport(report: EvalReport) -> Dict[str, Any]:
|
||||
"""
|
||||
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] = {
|
||||
"case_id": r.case_name,
|
||||
"passed": r.passed,
|
||||
"score": r.score,
|
||||
"latency_ms": r.latency_ms,
|
||||
"test_case_id": r.case_name,
|
||||
"passed": verdict["passed"],
|
||||
"grader_results": grader_results,
|
||||
}
|
||||
if r.actual_output is not None:
|
||||
item["actual_output"] = r.actual_output
|
||||
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))
|
||||
if r.error is not None:
|
||||
item["error"] = r.error
|
||||
if r.criteria_scores:
|
||||
item["grader_scores"] = dict(r.criteria_scores)
|
||||
item["error"] = {"type": "runner_error", "message": str(r.error)}
|
||||
if r.record is not None:
|
||||
item["record"] = r.record
|
||||
item["metadata"] = {_META_KEY: {"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 {
|
||||
"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 {}),
|
||||
},
|
||||
"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(),
|
||||
"results": results,
|
||||
"summary": summary,
|
||||
"metadata": {_META_KEY: meta},
|
||||
}
|
||||
@@ -27,18 +27,41 @@ _LAZY_IMPORTS = {
|
||||
"GatewayMessage": ("praisonaiagents.gateway.protocols", "GatewayMessage"),
|
||||
"EventType": ("praisonaiagents.gateway.protocols", "EventType"),
|
||||
"OperatorScope": ("praisonaiagents.gateway.protocols", "OperatorScope"),
|
||||
# Multi-observer session sharing (Issue #5192)
|
||||
"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"),
|
||||
"GATEWAY_METHODS": ("praisonaiagents.gateway.protocols", "GATEWAY_METHODS"),
|
||||
"register_gateway_method": ("praisonaiagents.gateway.protocols", "register_gateway_method"),
|
||||
"resolve_required_scope": ("praisonaiagents.gateway.protocols", "resolve_required_scope"),
|
||||
"authorize_method": ("praisonaiagents.gateway.protocols", "authorize_method"),
|
||||
"GatewayUnauthorized": ("praisonaiagents.gateway.protocols", "GatewayUnauthorized"),
|
||||
# Config hot-reload observability (Issue #3049)
|
||||
"ReloadStatus": ("praisonaiagents.gateway.protocols", "ReloadStatus"),
|
||||
"compute_config_revision": ("praisonaiagents.gateway.protocols", "compute_config_revision"),
|
||||
# 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"),
|
||||
@@ -67,10 +90,18 @@ _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"),
|
||||
@@ -203,6 +234,9 @@ _LAZY_IMPORTS = {
|
||||
"ConnectRecoveryStep": ("praisonaiagents.gateway.protocols", "ConnectRecoveryStep"),
|
||||
"is_recoverable": ("praisonaiagents.gateway.protocols", "is_recoverable"),
|
||||
"MessageParams": ("praisonaiagents.gateway.protocols", "MessageParams"),
|
||||
# First-class attachment contract (Issue #5207)
|
||||
"AttachmentRef": ("praisonaiagents.gateway.protocols", "AttachmentRef"),
|
||||
"AttachmentStoreProtocol": ("praisonaiagents.gateway.protocols", "AttachmentStoreProtocol"),
|
||||
"LeaveParams": ("praisonaiagents.gateway.protocols", "LeaveParams"),
|
||||
"JoinParams": ("praisonaiagents.gateway.protocols", "JoinParams"),
|
||||
"FrameDecodeError": ("praisonaiagents.gateway.protocols", "FrameDecodeError"),
|
||||
@@ -218,6 +252,8 @@ _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"),
|
||||
@@ -240,6 +276,8 @@ _LAZY_IMPORTS = {
|
||||
"SessionConfig": ("praisonaiagents.gateway.config", "SessionConfig"),
|
||||
"ApiConfig": ("praisonaiagents.gateway.config", "ApiConfig"),
|
||||
"EmergencyStopConfig": ("praisonaiagents.gateway.config", "EmergencyStopConfig"),
|
||||
# Attachment ceilings config (Issue #5207)
|
||||
"AttachmentConfig": ("praisonaiagents.gateway.config", "AttachmentConfig"),
|
||||
"ChannelRouteConfig": ("praisonaiagents.gateway.config", "ChannelRouteConfig"),
|
||||
"MultiChannelGatewayConfig": ("praisonaiagents.gateway.config", "MultiChannelGatewayConfig"),
|
||||
# Config version stamp + doctor-driven migration (Issue #3841)
|
||||
@@ -264,6 +302,12 @@ _LAZY_IMPORTS = {
|
||||
# Reload scope classification (Issue #3440)
|
||||
"ReloadScope": ("praisonaiagents.gateway.config", "ReloadScope"),
|
||||
"classify_reload": ("praisonaiagents.gateway.config", "classify_reload"),
|
||||
# Candidate validation + rollback contract (Issue #5144)
|
||||
"CandidateReport": ("praisonaiagents.gateway.config", "CandidateReport"),
|
||||
"ReloadValidationProtocol": (
|
||||
"praisonaiagents.gateway.config",
|
||||
"ReloadValidationProtocol",
|
||||
),
|
||||
}
|
||||
|
||||
# Lazy loading cache (shared with wrapper-backed implementations below)
|
||||
@@ -348,17 +392,40 @@ __all__ = [
|
||||
"GatewayMessage",
|
||||
"EventType",
|
||||
"OperatorScope",
|
||||
# Multi-observer session sharing (Issue #5192)
|
||||
"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",
|
||||
"GATEWAY_METHODS",
|
||||
"register_gateway_method",
|
||||
"resolve_required_scope",
|
||||
"authorize_method",
|
||||
"GatewayUnauthorized",
|
||||
# Config hot-reload observability (Issue #3049)
|
||||
"ReloadStatus",
|
||||
"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",
|
||||
@@ -382,11 +449,18 @@ __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",
|
||||
@@ -516,6 +590,9 @@ __all__ = [
|
||||
"ConnectRecoveryStep",
|
||||
"is_recoverable",
|
||||
"MessageParams",
|
||||
# First-class attachment contract (Issue #5207)
|
||||
"AttachmentRef",
|
||||
"AttachmentStoreProtocol",
|
||||
"LeaveParams",
|
||||
"JoinParams",
|
||||
"FrameDecodeError",
|
||||
@@ -531,6 +608,8 @@ __all__ = [
|
||||
# Event-loop liveness watchdog (Issue #3385)
|
||||
"LoopWatchdogPolicy",
|
||||
"LoopWatchdog",
|
||||
# Startup-phase watchdog (Issue #5265)
|
||||
"StartupWatchdog",
|
||||
# Unified degraded-capability registry (Issue #3518)
|
||||
"DegradedOwner",
|
||||
"DegradedCapabilityProtocol",
|
||||
@@ -553,6 +632,8 @@ __all__ = [
|
||||
"SessionConfig",
|
||||
"ApiConfig",
|
||||
"EmergencyStopConfig",
|
||||
# Attachment ceilings config (Issue #5207)
|
||||
"AttachmentConfig",
|
||||
"ChannelRouteConfig",
|
||||
"MultiChannelGatewayConfig",
|
||||
# Config version stamp + doctor-driven migration (Issue #3841)
|
||||
@@ -575,6 +656,8 @@ __all__ = [
|
||||
"is_hot_appliable",
|
||||
"ReloadScope",
|
||||
"classify_reload",
|
||||
"CandidateReport",
|
||||
"ReloadValidationProtocol",
|
||||
# Implementations (lazy loaded from praisonai wrapper)
|
||||
"WebSocketGateway",
|
||||
"GatewaySession",
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
"""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="")
|
||||
@@ -297,6 +297,74 @@ class SupportsHotReload(Protocol):
|
||||
...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Candidate reload validation + rollback contract (Issue #5144)
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# ``classify_reload``/``ReloadScope`` decide *what* a config change touches; a
|
||||
# structural (``FULL``) change makes the runtime tear the live channels down
|
||||
# and bring the new config up. A config that passes the schema but fails at
|
||||
# *runtime* (unreachable model, an adapter that throws on load, bad wiring) is
|
||||
# not caught before that cutover, and there is no rollback — a single bad edit
|
||||
# can take a healthy always-on gateway offline.
|
||||
#
|
||||
# This is the small, canonical contract that makes a reload an *atomic,
|
||||
# validated cutover*: build/boot the candidate runtime in isolation and run the
|
||||
# existing readiness/turn pre-flight *before* the live one is withdrawn; only
|
||||
# swap when the candidate is proven healthy, otherwise keep the old config
|
||||
# serving. Core owns only the *shape* (a report + a protocol, no heavy
|
||||
# imports); the wrapper/bot gateway performs the concrete build+pre-flight.
|
||||
|
||||
|
||||
@dataclass
|
||||
class CandidateReport:
|
||||
"""Outcome of validating a candidate config *before* cutover.
|
||||
|
||||
Produced by :meth:`ReloadValidationProtocol.validate_candidate`: the
|
||||
gateway builds/boots the candidate runtime in isolation (or a pre-flight)
|
||||
and reports whether it is healthy enough to swap in. When ``ok`` is False
|
||||
the live config keeps serving and the failures feed
|
||||
:class:`~praisonaiagents.gateway.protocols.ReloadStatus.error` / ``health()``.
|
||||
|
||||
Attributes:
|
||||
ok: Whether the candidate built and passed pre-flight — only then may
|
||||
the live runtime be drained and the candidate swapped in.
|
||||
failures: Human-readable reasons the candidate was rejected (empty when
|
||||
``ok``); joined into the reload ``error`` an operator sees.
|
||||
candidate: Optional opaque handle to the already-built candidate runtime
|
||||
the wrapper activates on success. Left untyped so core carries no
|
||||
heavy runtime import; ``None`` when the runtime validates in place.
|
||||
"""
|
||||
|
||||
ok: bool
|
||||
failures: "List[str]" = field(default_factory=list)
|
||||
candidate: Optional[Any] = None
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ReloadValidationProtocol(Protocol):
|
||||
"""Protocol a gateway implements for atomic, validated config cutover.
|
||||
|
||||
The invariant: **never drain the live runtime until**
|
||||
:meth:`validate_candidate` **returns a report with** ``ok=True``. Owning
|
||||
this contract in core keeps the "never cut over to an unvalidated
|
||||
candidate" guarantee canonical so every gateway build honours it, rather
|
||||
than being re-implemented per entry point (SIGHUP, ``gateway reload``,
|
||||
``gateway restart``).
|
||||
"""
|
||||
|
||||
async def validate_candidate(
|
||||
self, new_config: Mapping[str, Any]
|
||||
) -> CandidateReport:
|
||||
"""Build/boot the candidate from ``new_config`` and pre-flight it.
|
||||
|
||||
Must not stop or mutate the live runtime; on any failure return a
|
||||
report with ``ok=False`` and the reason(s) so the caller keeps the
|
||||
previous config serving.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionConfig:
|
||||
"""Configuration for gateway sessions.
|
||||
@@ -685,10 +753,17 @@ 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:
|
||||
@@ -697,7 +772,7 @@ class ApiConfig:
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {"openai": self.openai, "mcp": self.mcp}
|
||||
return {"openai": self.openai, "mcp": self.mcp, "stream": self.stream}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Optional[Dict[str, Any]]) -> "ApiConfig":
|
||||
@@ -707,6 +782,7 @@ class ApiConfig:
|
||||
return cls(
|
||||
openai=bool(data.get("openai", False)),
|
||||
mcp=bool(data.get("mcp", False)),
|
||||
stream=bool(data.get("stream", False)),
|
||||
)
|
||||
|
||||
|
||||
@@ -775,6 +851,163 @@ class EmergencyStopConfig:
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttachmentConfig:
|
||||
"""Attachment/media ceilings for the gateway wire protocol (Issue #5207).
|
||||
|
||||
The gateway is the control plane custom ``/ws`` clients build on, but the
|
||||
``message`` frame had no attachment contract and no per-attachment size/type
|
||||
policy — only the whole-frame ``max_payload``. These ceilings are advertised
|
||||
to clients in the ``hello`` handshake (folded into ``HelloResult.policy``) so
|
||||
a well-behaved client self-limits — inlining small files, chunking larger
|
||||
ones at ``chunk_bytes`` — instead of discovering the limit by being
|
||||
disconnected. All fields are advisory limits; enforcement lives wherever the
|
||||
gateway materialises attachments (a wrapper/bot concern).
|
||||
|
||||
Backward-compatible: attachments are opt-in per turn, and a client that
|
||||
never sends one is unaffected by these values.
|
||||
|
||||
Attributes:
|
||||
max_attachment_bytes: Largest single attachment (inline or stored) the
|
||||
gateway accepts, in bytes. Default 10 MiB.
|
||||
max_attachments: Maximum number of attachments on one ``message`` turn.
|
||||
chunk_bytes: Advertised chunk size a client should use when streaming a
|
||||
large file into the attachment store (reserve → put chunk → close).
|
||||
allowed_types: Optional allow-list of MIME types / prefixes (e.g.
|
||||
``image/`` or ``application/pdf``). Empty (default) means no type
|
||||
restriction is advertised.
|
||||
"""
|
||||
|
||||
max_attachment_bytes: int = 10 * 1024 * 1024 # 10 MiB
|
||||
max_attachments: int = 10
|
||||
chunk_bytes: int = 256 * 1024 # 256 KiB
|
||||
allowed_types: List[str] = field(default_factory=list)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.max_attachment_bytes < 0:
|
||||
raise ValueError(
|
||||
"max_attachment_bytes must be >= 0 (use 0 to disable attachments)"
|
||||
)
|
||||
if self.max_attachments < 0:
|
||||
raise ValueError("max_attachments must be >= 0")
|
||||
if self.chunk_bytes <= 0:
|
||||
raise ValueError("chunk_bytes must be > 0")
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
"""Whether the gateway advertises any attachment capacity."""
|
||||
return self.max_attachment_bytes > 0 and self.max_attachments > 0
|
||||
|
||||
def to_policy(self) -> Dict[str, Any]:
|
||||
"""Build the policy fragment advertised in ``HelloResult.policy``.
|
||||
|
||||
Only advertises the attachment ceilings when enabled, so a gateway with
|
||||
attachments disabled advertises exactly today's policy shape. The type
|
||||
allow-list is included only when configured.
|
||||
"""
|
||||
if not self.enabled:
|
||||
return {}
|
||||
policy: Dict[str, Any] = {
|
||||
"max_attachment_bytes": self.max_attachment_bytes,
|
||||
"max_attachments": self.max_attachments,
|
||||
"chunk_bytes": self.chunk_bytes,
|
||||
}
|
||||
if self.allowed_types:
|
||||
policy["allowed_attachment_types"] = list(self.allowed_types)
|
||||
return policy
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"max_attachment_bytes": self.max_attachment_bytes,
|
||||
"max_attachments": self.max_attachments,
|
||||
"chunk_bytes": self.chunk_bytes,
|
||||
"allowed_types": list(self.allowed_types),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Optional[Dict[str, Any]]) -> "AttachmentConfig":
|
||||
"""Create from a parsed ``gateway.attachments`` mapping (tolerant of None)."""
|
||||
if not isinstance(data, dict):
|
||||
return cls()
|
||||
allowed = data.get("allowed_types")
|
||||
return cls(
|
||||
max_attachment_bytes=int(
|
||||
data.get("max_attachment_bytes", 10 * 1024 * 1024)
|
||||
),
|
||||
max_attachments=int(data.get("max_attachments", 10)),
|
||||
chunk_bytes=int(data.get("chunk_bytes", 256 * 1024)),
|
||||
allowed_types=[str(t) for t in allowed] if isinstance(allowed, (list, tuple)) else [],
|
||||
)
|
||||
|
||||
|
||||
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.
|
||||
@@ -807,6 +1040,13 @@ 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
|
||||
@@ -830,6 +1070,18 @@ 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).
|
||||
max_concurrent_runs_per_scope: int = 0
|
||||
# Issue #2620: pre-auth edge protections for internet-exposed deployments.
|
||||
# Cap concurrent *unauthenticated* WebSocket connections per source IP so a
|
||||
# hostile client cannot park many half-open sockets up to max_connections
|
||||
@@ -858,6 +1110,12 @@ class GatewayConfig:
|
||||
control: "EmergencyStopConfig" = field(
|
||||
default_factory=lambda: EmergencyStopConfig()
|
||||
)
|
||||
# Issue #5207: first-class attachment ceilings for the gateway wire
|
||||
# protocol, advertised to clients via HelloResult.policy so a custom /ws
|
||||
# client can self-limit/chunk before sending a file. Defaults keep a
|
||||
# sensible 10 MiB / 10-file / 256 KiB-chunk policy; attachments remain
|
||||
# opt-in per turn so a client that never sends one is unaffected.
|
||||
attachments: "AttachmentConfig" = field(default_factory=lambda: AttachmentConfig())
|
||||
# Issue #4766: per-session turn-execution seam. Selects *where* a session's
|
||||
# agent turn runs (see ``TurnExecutorProtocol``). ``None`` (the default)
|
||||
# resolves to ``InProcessTurnExecutor`` at runtime — today's on-loop
|
||||
@@ -899,6 +1157,19 @@ 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 "
|
||||
"(use 0 to disable the per-scope sub-limit)"
|
||||
)
|
||||
if self.preauth_max_connections_per_ip < 0:
|
||||
raise ValueError(
|
||||
"preauth_max_connections_per_ip must be >= 0 (use 0 to disable "
|
||||
@@ -966,6 +1237,7 @@ 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,
|
||||
@@ -981,6 +1253,10 @@ 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,
|
||||
"push": self.push.to_dict(),
|
||||
@@ -989,6 +1265,7 @@ class GatewayConfig:
|
||||
"liveness": self.liveness.to_dict(),
|
||||
"turn_lock": self.turn_lock.to_dict(),
|
||||
"control": self.control.to_dict(),
|
||||
"attachments": self.attachments.to_dict(),
|
||||
# Report the active executor type so ``gateway doctor`` can surface
|
||||
# it; ``None`` means the in-process default (today's behaviour).
|
||||
"executor": (
|
||||
@@ -1158,9 +1435,17 @@ 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"),
|
||||
@@ -1177,6 +1462,20 @@ 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
|
||||
),
|
||||
preauth_max_connections_per_ip=int(
|
||||
gw_data.get("preauth_max_connections_per_ip", 32)
|
||||
),
|
||||
@@ -1188,6 +1487,7 @@ class MultiChannelGatewayConfig:
|
||||
liveness=LivenessConfig.from_dict(gw_data.get("liveness")),
|
||||
turn_lock=TurnLockConfig.from_dict(gw_data.get("turn_lock")),
|
||||
control=EmergencyStopConfig.from_dict(gw_data.get("control")),
|
||||
attachments=AttachmentConfig.from_dict(gw_data.get("attachments")),
|
||||
)
|
||||
|
||||
# Parse agents section (pass through as dicts)
|
||||
|
||||
@@ -52,6 +52,7 @@ from .protocols import GATEWAY_RESTART_EXIT_CODE
|
||||
__all__ = [
|
||||
"LoopWatchdogPolicy",
|
||||
"LoopWatchdog",
|
||||
"StartupWatchdog",
|
||||
]
|
||||
|
||||
|
||||
@@ -437,3 +438,271 @@ 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
|
||||
@@ -48,6 +48,7 @@ MIN_CLIENT_PROTOCOL_VERSION = 1
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import asyncio
|
||||
from typing import BinaryIO
|
||||
from praisonai.gateway.pairing import PairedChannel
|
||||
from ..agent import Agent
|
||||
from ..bots.presentation import MessagePresentation
|
||||
@@ -268,6 +269,84 @@ class OperatorScope(str, Enum):
|
||||
return list(cls)
|
||||
|
||||
|
||||
class SessionVisibility(str, Enum):
|
||||
"""How widely a Gateway session may be observed by other clients.
|
||||
|
||||
A session's visibility gates *whether* a second authorised client may
|
||||
attach to watch (or co-drive) the same live turn stream, on top of the
|
||||
orthogonal :class:`OperatorScope` authorisation. The vocabulary lives in
|
||||
core so every client/impl agrees; enforcement happens in the wrapper.
|
||||
|
||||
Values:
|
||||
PRIVATE: Owner only. Today's behaviour and the default — a session
|
||||
stays bound 1:1 to its owning client.
|
||||
SHARED: Members may co-drive; additional authorised clients may
|
||||
attach with WRITE-capable roles.
|
||||
READ_ONLY: Additional clients may attach as viewers to observe the
|
||||
transcript, but not send.
|
||||
"""
|
||||
|
||||
PRIVATE = "private"
|
||||
SHARED = "shared"
|
||||
READ_ONLY = "read_only"
|
||||
|
||||
|
||||
class SessionSharingRole(str, Enum):
|
||||
"""The role an attached client holds within a shared Gateway session.
|
||||
|
||||
Roles describe *what an attached client may do* within the session and
|
||||
layer on top of :class:`OperatorScope`. Fan-out remains scope-aware: an
|
||||
observer without ``READ`` receives nothing regardless of role.
|
||||
|
||||
Values:
|
||||
OWNER: Controls sharing/visibility; the session's original driver.
|
||||
MEMBER: May co-drive (send) when the session is ``SHARED``.
|
||||
VIEWER: Observes the live turn stream but cannot send.
|
||||
"""
|
||||
|
||||
OWNER = "owner"
|
||||
MEMBER = "member"
|
||||
VIEWER = "viewer"
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SessionObserverProtocol(Protocol):
|
||||
"""Protocol for attaching multiple observers to one live session.
|
||||
|
||||
Extends the single-owner :class:`GatewaySessionProtocol` model so that a
|
||||
session can fan its live turn stream (tokens, tool calls, results, final)
|
||||
out to more than one authorised client. Implementations live in the
|
||||
wrapper Gateway; the contract lives here so every client agrees on the
|
||||
attach/detach/enumerate surface.
|
||||
|
||||
Default behaviour is unchanged: a ``PRIVATE`` session keeps a single
|
||||
``OWNER`` observer (today's 1:1 model).
|
||||
"""
|
||||
|
||||
def attach(
|
||||
self,
|
||||
session_id: str,
|
||||
client_id: str,
|
||||
role: SessionSharingRole = SessionSharingRole.VIEWER,
|
||||
) -> None:
|
||||
"""Attach ``client_id`` to ``session_id`` with the given role.
|
||||
|
||||
Adds an observer rather than re-pointing ownership, so a second
|
||||
client can watch a session without hijacking it.
|
||||
"""
|
||||
...
|
||||
|
||||
def detach(self, session_id: str, client_id: str) -> None:
|
||||
"""Detach ``client_id`` from ``session_id`` (stop observing)."""
|
||||
...
|
||||
|
||||
def observers(
|
||||
self, session_id: str
|
||||
) -> "List[Tuple[str, SessionSharingRole]]":
|
||||
"""Return ``(client_id, role)`` pairs currently attached to the session."""
|
||||
...
|
||||
|
||||
|
||||
@dataclass
|
||||
class HelloParams:
|
||||
"""Parameters for initiating a versioned handshake.
|
||||
@@ -297,14 +376,18 @@ class HelloResult:
|
||||
Attributes:
|
||||
protocol: The negotiated protocol version
|
||||
features: Supported methods and events
|
||||
policy: Gateway policy limits (max_payload, heartbeat_ms, etc.)
|
||||
policy: Gateway policy limits (max_payload, heartbeat_ms, and — Issue
|
||||
#5207 — attachment ceilings max_attachment_bytes / max_attachments /
|
||||
chunk_bytes so a client self-limits and chunks before sending). The
|
||||
optional ``allowed_attachment_types`` list, when present, is carried
|
||||
here too.
|
||||
session_id: The session ID for this connection
|
||||
resumed: Whether an existing session was resumed
|
||||
cursor: Current event cursor position
|
||||
"""
|
||||
protocol: int
|
||||
features: Dict[str, List[str]] # {"methods": [...], "events": [...]}
|
||||
policy: Dict[str, int] # {"max_payload": ..., "heartbeat_ms": ...}
|
||||
policy: Dict[str, Any] # {"max_payload": ..., "heartbeat_ms": ..., "max_attachment_bytes": ...}
|
||||
session_id: str
|
||||
resumed: bool
|
||||
cursor: int
|
||||
@@ -465,6 +548,149 @@ def _require_str(value: Any, *, field_name: str) -> str:
|
||||
return value
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttachmentRef:
|
||||
"""A first-class attachment reference on the gateway wire protocol (Issue #5207).
|
||||
|
||||
The gateway is the control plane every custom ``/ws`` client builds on, yet
|
||||
the ``message`` frame carried only text or a free-form dict — a client had
|
||||
no supported, size-bounded, validated way to hand the agent a file or to
|
||||
receive an agent-generated artefact back. ``AttachmentRef`` is that missing
|
||||
typed contract: one shape every client and gateway implementation agrees on,
|
||||
additive and backward-compatible (no existing frame changes shape).
|
||||
|
||||
An attachment is carried one of two ways:
|
||||
|
||||
* **Inline** — ``data`` holds base64 for a small file, bounded by the
|
||||
advertised ``max_attachment_bytes`` policy so a well-behaved client
|
||||
self-limits before sending rather than discovering the limit by being
|
||||
disconnected.
|
||||
* **By reference** — ``ref_id`` names an entry the gateway materialised in
|
||||
its attachment store (see :class:`AttachmentStoreProtocol`), used for
|
||||
larger files streamed via chunked upload instead of a single frame.
|
||||
|
||||
The same shape is reused outbound: agent-generated files surface as
|
||||
``AttachmentRef`` entries a generic client can fetch over the same
|
||||
connection, instead of relying on platform (Telegram/Slack/…) delivery.
|
||||
|
||||
Attributes:
|
||||
filename: Client-facing file name (display / download target).
|
||||
mime: MIME type of the payload (e.g. ``image/png``, ``application/pdf``).
|
||||
size: Declared size in bytes (used to validate against policy ceilings).
|
||||
data: Inline base64 payload for a small file, or ``None`` when the
|
||||
attachment is carried by ``ref_id``.
|
||||
ref_id: Identifier in the gateway attachment store for a chunked / large
|
||||
file, or ``None`` for a purely inline attachment.
|
||||
"""
|
||||
|
||||
filename: str
|
||||
mime: str
|
||||
size: int
|
||||
data: Optional[str] = None # inline base64 for small files
|
||||
ref_id: Optional[str] = None # id in the gateway attachment store (chunked/large)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Serialize to a wire dict, omitting the unused carrier field."""
|
||||
frame: Dict[str, Any] = {
|
||||
"filename": self.filename,
|
||||
"mime": self.mime,
|
||||
"size": self.size,
|
||||
}
|
||||
if self.data is not None:
|
||||
frame["data"] = self.data
|
||||
if self.ref_id is not None:
|
||||
frame["ref_id"] = self.ref_id
|
||||
return frame
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "AttachmentRef":
|
||||
"""Validate a raw attachment dict into a typed ref.
|
||||
|
||||
Rejects a malformed attachment deterministically with the same
|
||||
structured ``HelloError`` envelope the rest of the inbound codec uses,
|
||||
so handlers keep receiving already-validated objects.
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
raise FrameDecodeError(
|
||||
HelloError(
|
||||
code=ConnectErrorCode.CONFIGURATION_ERROR,
|
||||
message="Each attachment must be an object",
|
||||
next_step=ConnectRecoveryStep.DO_NOT_RETRY,
|
||||
)
|
||||
)
|
||||
filename = _require_str(data.get("filename"), field_name="attachment.filename")
|
||||
mime = _require_str(data.get("mime"), field_name="attachment.mime")
|
||||
size = _coerce_int(data.get("size"), field_name="attachment.size", default=0)
|
||||
if size < 0:
|
||||
raise FrameDecodeError(
|
||||
HelloError(
|
||||
code=ConnectErrorCode.CONFIGURATION_ERROR,
|
||||
message="Field 'attachment.size' must be >= 0",
|
||||
next_step=ConnectRecoveryStep.DO_NOT_RETRY,
|
||||
)
|
||||
)
|
||||
inline = _as_opt_str(data.get("data"))
|
||||
ref_id = _as_opt_str(data.get("ref_id"))
|
||||
if inline is None and ref_id is None:
|
||||
raise FrameDecodeError(
|
||||
HelloError(
|
||||
code=ConnectErrorCode.CONFIGURATION_ERROR,
|
||||
message=(
|
||||
"An attachment must carry either inline 'data' (base64) "
|
||||
"or a store 'ref_id'"
|
||||
),
|
||||
next_step=ConnectRecoveryStep.DO_NOT_RETRY,
|
||||
)
|
||||
)
|
||||
if inline is not None and ref_id is not None:
|
||||
raise FrameDecodeError(
|
||||
HelloError(
|
||||
code=ConnectErrorCode.CONFIGURATION_ERROR,
|
||||
message=(
|
||||
"An attachment must carry exactly one of inline 'data' "
|
||||
"(base64) or a store 'ref_id', not both"
|
||||
),
|
||||
next_step=ConnectRecoveryStep.DO_NOT_RETRY,
|
||||
)
|
||||
)
|
||||
return cls(
|
||||
filename=filename,
|
||||
mime=mime,
|
||||
size=size,
|
||||
data=inline,
|
||||
ref_id=ref_id,
|
||||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class AttachmentStoreProtocol(Protocol):
|
||||
"""Seam for a gateway-managed attachment store (Issue #5207).
|
||||
|
||||
Large files cannot ride inline inside a single ``message`` frame (bounded by
|
||||
``max_payload``), so they are streamed as chunks into a store keyed by a
|
||||
reserved ``ref_id`` (reserve → put chunk at offset → close), then referenced
|
||||
from :class:`AttachmentRef`. The contract lives in core so every client and
|
||||
implementation agree on one shape; the concrete durable store (filesystem /
|
||||
object store) is a wrapper/bot concern, kept out of the dependency-free core.
|
||||
"""
|
||||
|
||||
def reserve(self, filename: str, mime: str, size: int) -> str:
|
||||
"""Reserve an entry for an incoming file and return its ``ref_id``."""
|
||||
...
|
||||
|
||||
def put_chunk(self, ref_id: str, offset: int, data: bytes) -> None:
|
||||
"""Write ``data`` at ``offset`` for the reserved ``ref_id``."""
|
||||
...
|
||||
|
||||
def close(self, ref_id: str) -> AttachmentRef:
|
||||
"""Finalize the upload and return the resulting :class:`AttachmentRef`."""
|
||||
...
|
||||
|
||||
def open(self, ref_id: str) -> "BinaryIO":
|
||||
"""Open the stored attachment for reading (outbound download)."""
|
||||
...
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageParams:
|
||||
"""Validated ``message`` frame — a client turn sent to the agent.
|
||||
@@ -476,12 +702,23 @@ class MessageParams:
|
||||
content: The message body (text, or a structured payload).
|
||||
session_id: Optional session the message belongs to.
|
||||
message_id: Optional client-supplied idempotency/correlation id.
|
||||
request_id: Optional client-supplied request-idempotency key (Issue
|
||||
#5193). When present the gateway dedups a resend of the same
|
||||
``request_id`` — a request queued while reconnecting and flushed on
|
||||
reconnect runs the turn exactly once instead of being lost or
|
||||
double-run. Absent keeps the legacy fire-and-forget behaviour.
|
||||
attachments: Optional list of :class:`AttachmentRef` (Issue #5207) — a
|
||||
first-class, size-bounded way for any ``/ws`` client (not just
|
||||
platform bots) to hand the agent files. Defaults to empty, so a
|
||||
frame without attachments decodes exactly as before.
|
||||
metadata: Optional additional message metadata.
|
||||
"""
|
||||
|
||||
content: Union[str, Dict[str, Any]]
|
||||
session_id: Optional[str] = None
|
||||
message_id: Optional[str] = None
|
||||
request_id: Optional[str] = None
|
||||
attachments: List[AttachmentRef] = field(default_factory=list)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
type: str = field(default="message", init=False)
|
||||
@@ -509,10 +746,24 @@ class MessageParams:
|
||||
metadata = data.get("metadata")
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
raw_attachments = data.get("attachments")
|
||||
attachments: List[AttachmentRef] = []
|
||||
if raw_attachments is not None:
|
||||
if not isinstance(raw_attachments, (list, tuple)):
|
||||
raise FrameDecodeError(
|
||||
HelloError(
|
||||
code=ConnectErrorCode.CONFIGURATION_ERROR,
|
||||
message="Field 'attachments' must be an array",
|
||||
next_step=ConnectRecoveryStep.DO_NOT_RETRY,
|
||||
)
|
||||
)
|
||||
attachments = [AttachmentRef.from_dict(a) for a in raw_attachments]
|
||||
return cls(
|
||||
content=content,
|
||||
session_id=_as_opt_str(data.get("session_id")),
|
||||
message_id=_as_opt_str(data.get("message_id")),
|
||||
request_id=_as_opt_str(data.get("request_id")),
|
||||
attachments=attachments,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
@@ -2304,6 +2555,225 @@ 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)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -2823,6 +3293,57 @@ 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).
|
||||
@@ -3008,6 +3529,51 @@ 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)
|
||||
@@ -3114,6 +3680,122 @@ 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)
|
||||
#
|
||||
@@ -3760,6 +4442,14 @@ class GatewayConcurrencyPolicyProtocol(Protocol):
|
||||
A config-driven default (:class:`ConcurrencyLimitPolicy`) is provided for
|
||||
the common "N concurrent runs, bounded wait queue, declared overflow"
|
||||
case.
|
||||
|
||||
On a shared/multi-tenant gateway the decision may additionally be scoped
|
||||
by a tenant/profile token (Issue #5168), mirroring the ``scope`` axis of
|
||||
:class:`RateLimitPolicyProtocol` / :class:`SpendBudgetPolicyProtocol` — the
|
||||
wrapper passes the caller's ``scope`` plus the live per-scope in-flight and
|
||||
queued counts so a single tenant cannot occupy every global slot. When no
|
||||
per-scope sub-limit is configured the scope arguments are ignored and the
|
||||
global-only decision is returned unchanged (backward compatible).
|
||||
"""
|
||||
|
||||
max_concurrent_runs: int
|
||||
@@ -3771,6 +4461,9 @@ class GatewayConcurrencyPolicyProtocol(Protocol):
|
||||
in_flight: int,
|
||||
queued: int,
|
||||
session_id: str = "",
|
||||
scope: str = "",
|
||||
scope_in_flight: int = 0,
|
||||
scope_queued: int = 0,
|
||||
) -> AdmissionDecision:
|
||||
"""Return an :class:`AdmissionDecision` for the supplied facts."""
|
||||
...
|
||||
@@ -3803,10 +4496,18 @@ class ConcurrencyLimitPolicy:
|
||||
A ``max_concurrent_runs`` of ``0`` disables admission control entirely
|
||||
(today's behaviour: every inbound turn is admitted immediately).
|
||||
|
||||
On a shared/multi-tenant gateway an optional ``max_concurrent_runs_per_scope``
|
||||
caps concurrency *per tenant/scope* within the global ceiling (Issue #5168):
|
||||
a tenant that already holds its sub-limit is queued/shed against its own
|
||||
slice even when a global slot is notionally free, so a noisy neighbour
|
||||
cannot starve quiet tenants. A value of ``0`` (the default) disables the
|
||||
per-scope sub-limit, preserving the byte-for-byte global-only behaviour.
|
||||
|
||||
Example::
|
||||
|
||||
ConcurrencyLimitPolicy(max_concurrent_runs=32, queue_depth=128,
|
||||
overflow_policy="reject")
|
||||
overflow_policy="reject",
|
||||
max_concurrent_runs_per_scope=4)
|
||||
"""
|
||||
|
||||
_OVERFLOW = ("reject", "queue", "shed_oldest")
|
||||
@@ -3816,6 +4517,7 @@ class ConcurrencyLimitPolicy:
|
||||
max_concurrent_runs: int = 0,
|
||||
queue_depth: int = 0,
|
||||
overflow_policy: str = "reject",
|
||||
max_concurrent_runs_per_scope: int = 0,
|
||||
):
|
||||
try:
|
||||
ceiling = int(max_concurrent_runs)
|
||||
@@ -3842,9 +4544,22 @@ class ConcurrencyLimitPolicy:
|
||||
f"overflow_policy must be one of {self._OVERFLOW}, "
|
||||
f"got {overflow_policy!r}"
|
||||
)
|
||||
try:
|
||||
per_scope = int(max_concurrent_runs_per_scope)
|
||||
except (TypeError, ValueError):
|
||||
raise ValueError(
|
||||
f"max_concurrent_runs_per_scope must be an integer, "
|
||||
f"got {max_concurrent_runs_per_scope!r}"
|
||||
)
|
||||
if per_scope < 0:
|
||||
raise ValueError(
|
||||
f"max_concurrent_runs_per_scope must be >= 0, "
|
||||
f"got {max_concurrent_runs_per_scope!r}"
|
||||
)
|
||||
self.max_concurrent_runs = ceiling
|
||||
self.queue_depth = depth
|
||||
self.overflow_policy = overflow
|
||||
self.max_concurrent_runs_per_scope = per_scope
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
@@ -3857,13 +4572,30 @@ class ConcurrencyLimitPolicy:
|
||||
in_flight: int,
|
||||
queued: int,
|
||||
session_id: str = "",
|
||||
scope: str = "",
|
||||
scope_in_flight: int = 0,
|
||||
scope_queued: int = 0,
|
||||
) -> AdmissionDecision:
|
||||
# Disabled: preserve legacy always-admit behaviour.
|
||||
if self.max_concurrent_runs <= 0:
|
||||
return AdmissionDecision.ADMIT
|
||||
# Per-scope fairness (Issue #5168): a tenant already at its sub-limit
|
||||
# is queued/shed against its own slice even when a global slot is free,
|
||||
# so a noisy tenant cannot starve quiet ones. Applied first because it
|
||||
# is the tighter, per-tenant ceiling; the global gate below still caps
|
||||
# aggregate concurrency. ``scope`` must be non-empty for the sub-limit
|
||||
# to apply (an unscoped caller falls through to the global decision).
|
||||
if self.max_concurrent_runs_per_scope > 0 and scope:
|
||||
if scope_in_flight >= self.max_concurrent_runs_per_scope:
|
||||
return self._overflow(scope_queued)
|
||||
if in_flight < self.max_concurrent_runs:
|
||||
return AdmissionDecision.ADMIT
|
||||
# At the ceiling: consult the bounded wait queue.
|
||||
return self._overflow(queued)
|
||||
|
||||
def _overflow(self, queued: int) -> AdmissionDecision:
|
||||
# At a ceiling: queue while the bounded wait queue has room, else apply
|
||||
# the declared overflow behaviour.
|
||||
if queued < self.queue_depth:
|
||||
return AdmissionDecision.QUEUE
|
||||
# Queue is full: declared overflow behaviour.
|
||||
@@ -6968,14 +7700,82 @@ def resolve_required_scope(
|
||||
|
||||
Default-deny: an unclassified/unknown method requires ``ADMIN`` so new
|
||||
control surface is closed until explicitly classified — the omission fails
|
||||
**closed** rather than open.
|
||||
**closed** rather than open. A non-string (or unhashable) ``method`` is
|
||||
likewise treated as unclassified and requires ``ADMIN``, so a malformed
|
||||
frame fails closed deterministically instead of raising ``TypeError``.
|
||||
"""
|
||||
if not isinstance(method, str):
|
||||
return OperatorScope.ADMIN
|
||||
desc = GATEWAY_METHODS.get(method)
|
||||
if desc is None:
|
||||
return OperatorScope.ADMIN
|
||||
return desc.resolve(params)
|
||||
|
||||
|
||||
class GatewayUnauthorized(PermissionError):
|
||||
"""Raised when a caller lacks the scope a gateway method requires.
|
||||
|
||||
Carries the ``method`` and the ``required`` scope so the dispatcher can
|
||||
render a structured, machine-readable denial (matching today's
|
||||
``insufficient scope`` envelope) without re-deriving either.
|
||||
"""
|
||||
|
||||
def __init__(self, method: str, required: OperatorScope) -> None:
|
||||
self.method = method
|
||||
self.required = required
|
||||
super().__init__(
|
||||
f"method {method!r} requires scope {required.value!r}"
|
||||
)
|
||||
|
||||
|
||||
def _scope_satisfies(
|
||||
held: "Set[OperatorScope]", required: OperatorScope
|
||||
) -> bool:
|
||||
"""Whether the ``held`` scopes satisfy ``required``.
|
||||
|
||||
Two implication rules, matching the wrapper's long-standing behaviour so
|
||||
wiring the registry into dispatch does not regress already-classified
|
||||
methods:
|
||||
|
||||
- ``ADMIN`` implies every scope (top of the lattice).
|
||||
- ``READ`` is the baseline lifecycle/observation scope that every
|
||||
authorised operator implicitly holds. An operator provisioned with any
|
||||
actionable scope (``WRITE``/``APPROVALS``/``PAIRING``/``ADMIN``) can
|
||||
therefore still complete the READ-classified session lifecycle
|
||||
(``hello``/``join``/``leave``/status) — exactly as it could before this
|
||||
guard existed, when those frames carried no per-endpoint scope check.
|
||||
|
||||
Otherwise a caller is authorised only when it holds the exact required
|
||||
scope.
|
||||
"""
|
||||
if OperatorScope.ADMIN in held:
|
||||
return True
|
||||
if required == OperatorScope.READ:
|
||||
# Any authorised operator (holding at least one scope) may perform the
|
||||
# read-only lifecycle/observation surface.
|
||||
return bool(held)
|
||||
return required in held
|
||||
|
||||
|
||||
def authorize_method(
|
||||
method: str,
|
||||
client_scopes: "Set[OperatorScope]",
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Central, default-deny authorisation gate for a gateway method.
|
||||
|
||||
Resolves the required scope from the single source of truth
|
||||
(:func:`resolve_required_scope`) — so an unclassified/plugin-registered
|
||||
method is ``ADMIN``-only by omission rather than reachable ungated — and
|
||||
raises :class:`GatewayUnauthorized` when ``client_scopes`` do not satisfy
|
||||
it. Dispatch calls this once, centrally, instead of scattering per-endpoint
|
||||
checks that can drift from the registry.
|
||||
"""
|
||||
required = resolve_required_scope(method, params)
|
||||
if not _scope_satisfies(set(client_scopes), required):
|
||||
raise GatewayUnauthorized(method=method, required=required)
|
||||
|
||||
|
||||
# Core method classification. Registered once at import so the dispatcher can
|
||||
# consult the registry instead of scattered per-endpoint checks. Structural
|
||||
# mutations (channel control) demand ADMIN; sending as the agent needs WRITE;
|
||||
@@ -6990,10 +7790,21 @@ def _register_core_gateway_methods() -> None:
|
||||
"leave": OperatorScope.READ,
|
||||
"agent.message": OperatorScope.WRITE,
|
||||
"message": OperatorScope.WRITE,
|
||||
# Aborting a turn mutates it, so it carries the same scope as sending one.
|
||||
# Aborting a turn mutates it, so it carries the same scope as sending
|
||||
# one. The wire also accepts the ``message_abort`` event-type alias for
|
||||
# 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,
|
||||
|
||||
@@ -0,0 +1,365 @@
|
||||
"""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
|
||||
@@ -193,6 +193,41 @@ class GoalLoopMixin:
|
||||
self._persist_goal_state()
|
||||
return "continue", reason
|
||||
|
||||
async def _goal_gate_async(self, response: str) -> Optional[Tuple[str, str]]:
|
||||
"""Async counterpart of :meth:`_goal_gate`.
|
||||
|
||||
Both the configured verification hooks (which may shell out) and the
|
||||
completion judge's ``litellm.completion(...)`` call are blocking I/O, so
|
||||
the whole gate is offloaded to a worker thread. Mirrors
|
||||
``_verification_gate_async``'s existing pattern for the same reason: a
|
||||
slow judge/hook must never stall the event loop shared by other
|
||||
concurrent agents.
|
||||
|
||||
Cancellation safety: the gate mutates shared goal state
|
||||
(``_goal_state``, verdict counters, journal, persisted metadata). A bare
|
||||
``asyncio.to_thread`` worker keeps running after the awaiting coroutine
|
||||
is cancelled, so its late writes could race with ``run_goal_async``'s
|
||||
teardown (which nulls ``_goal_state``) or a subsequent run reusing the
|
||||
same agent. We therefore join the worker before propagating the
|
||||
cancellation, guaranteeing the gate finishes against the goal state it
|
||||
started on and no orphaned thread mutates a newer run.
|
||||
"""
|
||||
state = getattr(self, "_goal_state", None)
|
||||
if state is None or state.status != "active":
|
||||
return None
|
||||
import asyncio
|
||||
task = asyncio.ensure_future(asyncio.to_thread(self._goal_gate, response))
|
||||
try:
|
||||
return await task
|
||||
except asyncio.CancelledError:
|
||||
# Let the in-flight gate finish so its state mutations complete
|
||||
# before run_goal_async's finally clears _goal_state, then re-raise.
|
||||
try:
|
||||
await asyncio.shield(task)
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
# -- public API ------------------------------------------------------------
|
||||
|
||||
def _resume_or_new_goal_state(
|
||||
|
||||
@@ -80,6 +80,7 @@ __all__ = [
|
||||
"MessageSendingInput",
|
||||
"MessageSentInput",
|
||||
"MessageUndeliveredInput",
|
||||
"PlatformEventInput",
|
||||
# Gateway & schedule lifecycle event inputs
|
||||
"GatewayStartInput",
|
||||
"GatewayStopInput",
|
||||
@@ -161,6 +162,7 @@ _LAZY_GROUPS = {
|
||||
'MessageSendingInput': ('praisonaiagents.hooks.events', 'MessageSendingInput'),
|
||||
'MessageSentInput': ('praisonaiagents.hooks.events', 'MessageSentInput'),
|
||||
'MessageUndeliveredInput': ('praisonaiagents.hooks.events', 'MessageUndeliveredInput'),
|
||||
'PlatformEventInput': ('praisonaiagents.hooks.events', 'PlatformEventInput'),
|
||||
'GatewayStartInput': ('praisonaiagents.hooks.events', 'GatewayStartInput'),
|
||||
'GatewayStopInput': ('praisonaiagents.hooks.events', 'GatewayStopInput'),
|
||||
'ScheduleTriggerInput': ('praisonaiagents.hooks.events', 'ScheduleTriggerInput'),
|
||||
|
||||
@@ -388,6 +388,42 @@ class MessageUndeliveredInput(HookInput):
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class PlatformEventInput(HookInput):
|
||||
"""Input for inbound platform-event hooks (Issue #5161).
|
||||
|
||||
Carries a normalised, non-message platform event — a reaction, edit,
|
||||
deletion, membership change or thread creation — to
|
||||
``REACTION_RECEIVED`` / ``MESSAGE_EDITED`` / ``MESSAGE_DELETED`` /
|
||||
``MEMBER_JOINED`` / ``MEMBER_LEFT`` / ``THREAD_CREATED`` hooks. Fields mirror
|
||||
:class:`~praisonaiagents.bots.protocols.PlatformEvent`; the ``kind`` field
|
||||
disambiguates reaction add vs remove within ``REACTION_RECEIVED``. Emitted
|
||||
by ``MessageHookMixin.fire_platform_event`` (praisonai-bot).
|
||||
"""
|
||||
kind: str = ""
|
||||
platform: str = ""
|
||||
chat_id: str = ""
|
||||
user_id: str = ""
|
||||
message_id: Optional[str] = None
|
||||
emoji: Optional[str] = None
|
||||
new_text: Optional[str] = None
|
||||
thread_id: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"kind": self.kind,
|
||||
"platform": self.platform,
|
||||
"chat_id": self.chat_id,
|
||||
"user_id": self.user_id,
|
||||
"message_id": self.message_id,
|
||||
"emoji": self.emoji,
|
||||
"new_text": self.new_text[:500] if self.new_text else None,
|
||||
"thread_id": self.thread_id,
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class GatewayStartInput(HookInput):
|
||||
"""Input for GATEWAY_START hooks (gateway/BotOS started)."""
|
||||
|
||||
@@ -53,6 +53,9 @@ 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
|
||||
@@ -80,6 +83,8 @@ 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
|
||||
@@ -93,7 +98,8 @@ class HookRegistry:
|
||||
description: Optional[str] = None,
|
||||
sequential: bool = False,
|
||||
timeout: float = 60.0,
|
||||
is_async: bool = False
|
||||
is_async: bool = False,
|
||||
priority: int = 100
|
||||
) -> str:
|
||||
"""
|
||||
Register a Python function as a hook.
|
||||
@@ -107,6 +113,7 @@ 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
|
||||
@@ -119,7 +126,8 @@ class HookRegistry:
|
||||
description=description,
|
||||
sequential=sequential,
|
||||
timeout=timeout,
|
||||
is_async=is_async
|
||||
is_async=is_async,
|
||||
priority=priority
|
||||
)
|
||||
return self.register(hook)
|
||||
|
||||
@@ -133,7 +141,8 @@ class HookRegistry:
|
||||
sequential: bool = False,
|
||||
timeout: float = 60.0,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
shell: bool = True
|
||||
shell: bool = True,
|
||||
priority: int = 100
|
||||
) -> str:
|
||||
"""
|
||||
Register a shell command as a hook.
|
||||
@@ -148,6 +157,7 @@ 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
|
||||
@@ -161,7 +171,8 @@ class HookRegistry:
|
||||
sequential=sequential,
|
||||
timeout=timeout,
|
||||
env=env or {},
|
||||
shell=shell
|
||||
shell=shell,
|
||||
priority=priority
|
||||
)
|
||||
return self.register(hook)
|
||||
|
||||
@@ -170,7 +181,8 @@ class HookRegistry:
|
||||
event: HookEvent,
|
||||
matcher: Optional[str] = None,
|
||||
sequential: bool = False,
|
||||
timeout: float = 60.0
|
||||
timeout: float = 60.0,
|
||||
priority: int = 100
|
||||
) -> Callable:
|
||||
"""
|
||||
Decorator to register a function as a hook.
|
||||
@@ -180,6 +192,7 @@ 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
|
||||
@@ -199,7 +212,8 @@ class HookRegistry:
|
||||
matcher=matcher,
|
||||
sequential=sequential,
|
||||
timeout=timeout,
|
||||
is_async=is_async
|
||||
is_async=is_async,
|
||||
priority=priority
|
||||
)
|
||||
|
||||
@wraps(func)
|
||||
@@ -253,6 +267,10 @@ 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]
|
||||
|
||||
@@ -284,20 +302,25 @@ class HookRegistry:
|
||||
Dictionary mapping event names to hook info
|
||||
"""
|
||||
result = {}
|
||||
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
|
||||
]
|
||||
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
|
||||
]
|
||||
return result
|
||||
|
||||
def enable_hook(self, hook_id: str) -> bool:
|
||||
@@ -359,7 +382,7 @@ def set_default_registry(registry: HookRegistry):
|
||||
def add_hook(
|
||||
event: Union[str, HookEvent],
|
||||
callback: Optional[Callable[[HookInput], HookResult]] = None,
|
||||
priority: int = 10,
|
||||
priority: int = 100,
|
||||
matcher: Optional[str] = None
|
||||
) -> Union[str, Callable]:
|
||||
"""Register a hook callback. Simplified API.
|
||||
@@ -376,7 +399,9 @@ 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 10. (Reserved for future use)
|
||||
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.
|
||||
matcher: Optional regex pattern to match specific targets (e.g., tool names)
|
||||
|
||||
Returns:
|
||||
@@ -402,7 +427,8 @@ def add_hook(
|
||||
get_default_registry().register_function(
|
||||
event=event,
|
||||
func=func,
|
||||
matcher=matcher
|
||||
matcher=matcher,
|
||||
priority=priority
|
||||
)
|
||||
return func
|
||||
return decorator
|
||||
@@ -410,7 +436,8 @@ def add_hook(
|
||||
return get_default_registry().register_function(
|
||||
event=event,
|
||||
func=callback,
|
||||
matcher=matcher
|
||||
matcher=matcher,
|
||||
priority=priority
|
||||
)
|
||||
|
||||
def remove_hook(hook_id: str) -> bool:
|
||||
|
||||
@@ -102,25 +102,41 @@ class HookRunner:
|
||||
if not hooks:
|
||||
return []
|
||||
|
||||
# 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)
|
||||
# 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
|
||||
|
||||
return results
|
||||
|
||||
|
||||
@@ -81,7 +81,22 @@ class HookEvent(str, Enum):
|
||||
# the event that is actually emitted.
|
||||
BEFORE_MESSAGE = "message_received"
|
||||
AFTER_MESSAGE = "message_sending"
|
||||
|
||||
|
||||
# Inbound platform events (non-message): the gateway hearing more than text.
|
||||
# A user reacting/editing/deleting a message, joining/leaving a channel, or a
|
||||
# thread being created are delivered by the platform SDK but previously
|
||||
# dropped before reaching the agent or a hook. Adapters normalise the native
|
||||
# event into a ``PlatformEvent`` and emit it via
|
||||
# ``MessageHookMixin.fire_platform_event`` (praisonai-bot); platforms that
|
||||
# cannot deliver a given event simply never emit it (capability-gated,
|
||||
# opt-in). Payload: ``PlatformEventInput`` (Issue #5161).
|
||||
REACTION_RECEIVED = "reaction_received"
|
||||
MESSAGE_EDITED = "message_edited"
|
||||
MESSAGE_DELETED = "message_deleted"
|
||||
MEMBER_JOINED = "member_joined"
|
||||
MEMBER_LEFT = "member_left"
|
||||
THREAD_CREATED = "thread_created"
|
||||
|
||||
# Gateway lifecycle
|
||||
GATEWAY_START = "gateway_start"
|
||||
GATEWAY_STOP = "gateway_stop"
|
||||
@@ -301,6 +316,8 @@ 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,6 +211,8 @@ 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()
|
||||
@@ -236,9 +238,15 @@ class ChromaKnowledgeAdapter:
|
||||
if run_id is not None:
|
||||
where_filter["run_id"] = run_id
|
||||
|
||||
# Merge with additional filters
|
||||
# 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).
|
||||
if filters:
|
||||
where_filter.update(filters)
|
||||
for key, value in filters.items():
|
||||
if key in where_filter:
|
||||
continue
|
||||
where_filter[key] = value
|
||||
|
||||
# Search ChromaDB
|
||||
try:
|
||||
@@ -289,6 +297,8 @@ 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)
|
||||
@@ -375,6 +385,8 @@ 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
|
||||
@@ -510,8 +522,11 @@ 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
|
||||
@@ -528,6 +543,19 @@ 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)
|
||||
|
||||
@@ -551,9 +579,12 @@ 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 {})
|
||||
@@ -596,8 +627,11 @@ 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
|
||||
|
||||
@@ -137,26 +137,15 @@ class Mem0Adapter:
|
||||
def _safe_mem0_search(self, **kwargs):
|
||||
"""
|
||||
Defensive wrapper for mem0.search() to handle MongoDB vector store compatibility.
|
||||
|
||||
Catches TypeError about unexpected 'vectors' kwarg and returns safe fallback.
|
||||
This addresses the upstream mem0 bug: https://github.com/mem0ai/mem0/issues/3185
|
||||
|
||||
Delegates to the shared ``safe_mem0_search`` helper so the workaround for
|
||||
the upstream mem0 bug (https://github.com/mem0ai/mem0/issues/3185) lives in
|
||||
a single place and cannot drift from the memory-layer copy. The shared
|
||||
helper returns an empty ``list`` on the fallback branch, which
|
||||
``_normalize_mem0_results`` already handles alongside the dict shape.
|
||||
"""
|
||||
try:
|
||||
return self.memory.search(**kwargs)
|
||||
except TypeError as e:
|
||||
error_msg = str(e).lower()
|
||||
if "unexpected keyword argument" in error_msg and "vectors" in error_msg:
|
||||
logger.warning(
|
||||
"Detected mem0 MongoDB vector store compatibility issue. "
|
||||
"This is a known upstream bug: https://github.com/mem0ai/mem0/issues/3185. "
|
||||
"The MongoDB vector store requires Atlas and has signature mismatches. "
|
||||
"Returning empty results. Consider using Qdrant or Chroma as vector store backends."
|
||||
)
|
||||
# Return empty results structure that matches expected mem0 format
|
||||
return {"results": []}
|
||||
else:
|
||||
# Re-raise if it's a different TypeError
|
||||
raise
|
||||
from ...memory.adapters.factories import safe_mem0_search
|
||||
return safe_mem0_search(self.memory, **kwargs)
|
||||
|
||||
def _normalize_mem0_results(self, raw: Any) -> SearchResult:
|
||||
"""
|
||||
|
||||
@@ -169,7 +169,12 @@ 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
|
||||
@@ -210,6 +215,7 @@ 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]]:
|
||||
"""
|
||||
@@ -221,11 +227,17 @@ 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 = []
|
||||
|
||||
@@ -237,6 +249,13 @@ 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():
|
||||
@@ -332,6 +351,127 @@ 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'):
|
||||
|
||||
@@ -7,6 +7,7 @@ NO heavy imports - only stdlib and typing.
|
||||
Implementations are provided by praisonai-tools or wrapper layer.
|
||||
"""
|
||||
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from praisonaiagents._logging import get_logger
|
||||
@@ -81,11 +82,16 @@ class RerankerRegistry:
|
||||
"""
|
||||
|
||||
_instance: Optional["RerankerRegistry"] = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __new__(cls) -> "RerankerRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._rerankers: Dict[str, Callable[..., RerankerProtocol]] = {}
|
||||
with cls._lock:
|
||||
# Double-checked locking for multi-agent safety
|
||||
if cls._instance is None:
|
||||
instance = super().__new__(cls)
|
||||
instance._rerankers: Dict[str, Callable[..., RerankerProtocol]] = {}
|
||||
cls._instance = instance
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
@@ -93,30 +99,35 @@ class RerankerRegistry:
|
||||
name: str,
|
||||
factory: Callable[..., RerankerProtocol]
|
||||
) -> None:
|
||||
"""Register a reranker factory."""
|
||||
self._rerankers[name] = factory
|
||||
"""Register a reranker factory. Thread-safe."""
|
||||
with self._lock:
|
||||
self._rerankers[name] = factory
|
||||
|
||||
def get(
|
||||
self,
|
||||
name: str,
|
||||
**kwargs
|
||||
) -> Optional[RerankerProtocol]:
|
||||
"""Get a reranker by name."""
|
||||
if name in self._rerankers:
|
||||
try:
|
||||
return self._rerankers[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize reranker '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
"""Get a reranker by name. Thread-safe."""
|
||||
with self._lock:
|
||||
factory = self._rerankers.get(name)
|
||||
if factory is None:
|
||||
return None
|
||||
try:
|
||||
return factory(**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize reranker '{name}': {e}")
|
||||
return None
|
||||
|
||||
def list_rerankers(self) -> List[str]:
|
||||
"""List all registered reranker names."""
|
||||
return list(self._rerankers.keys())
|
||||
"""List all registered reranker names. Thread-safe."""
|
||||
with self._lock:
|
||||
return list(self._rerankers.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered rerankers."""
|
||||
self._rerankers.clear()
|
||||
"""Clear all registered rerankers. Thread-safe."""
|
||||
with self._lock:
|
||||
self._rerankers.clear()
|
||||
|
||||
def get_reranker_registry() -> RerankerRegistry:
|
||||
"""Get the global reranker registry instance."""
|
||||
|
||||
@@ -7,6 +7,7 @@ NO heavy imports - only stdlib and typing.
|
||||
Implementations are provided by the wrapper layer.
|
||||
"""
|
||||
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from enum import Enum
|
||||
@@ -99,11 +100,16 @@ class RetrieverRegistry:
|
||||
"""
|
||||
|
||||
_instance: Optional["RetrieverRegistry"] = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __new__(cls) -> "RetrieverRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._retrievers: Dict[str, Callable[..., RetrieverProtocol]] = {}
|
||||
with cls._lock:
|
||||
# Double-checked locking for multi-agent safety
|
||||
if cls._instance is None:
|
||||
instance = super().__new__(cls)
|
||||
instance._retrievers: Dict[str, Callable[..., RetrieverProtocol]] = {}
|
||||
cls._instance = instance
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
@@ -111,30 +117,35 @@ class RetrieverRegistry:
|
||||
name: str,
|
||||
factory: Callable[..., RetrieverProtocol]
|
||||
) -> None:
|
||||
"""Register a retriever factory."""
|
||||
self._retrievers[name] = factory
|
||||
"""Register a retriever factory. Thread-safe."""
|
||||
with self._lock:
|
||||
self._retrievers[name] = factory
|
||||
|
||||
def get(
|
||||
self,
|
||||
name: str,
|
||||
**kwargs
|
||||
) -> Optional[RetrieverProtocol]:
|
||||
"""Get a retriever by name."""
|
||||
if name in self._retrievers:
|
||||
try:
|
||||
return self._retrievers[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize retriever '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
"""Get a retriever by name. Thread-safe."""
|
||||
with self._lock:
|
||||
factory = self._retrievers.get(name)
|
||||
if factory is None:
|
||||
return None
|
||||
try:
|
||||
return factory(**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize retriever '{name}': {e}")
|
||||
return None
|
||||
|
||||
def list_retrievers(self) -> List[str]:
|
||||
"""List all registered retriever names."""
|
||||
return list(self._retrievers.keys())
|
||||
"""List all registered retriever names. Thread-safe."""
|
||||
with self._lock:
|
||||
return list(self._retrievers.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered retrievers."""
|
||||
self._retrievers.clear()
|
||||
"""Clear all registered retrievers. Thread-safe."""
|
||||
with self._lock:
|
||||
self._retrievers.clear()
|
||||
|
||||
def get_retriever_registry() -> RetrieverRegistry:
|
||||
"""Get the global retriever registry instance."""
|
||||
|
||||
@@ -87,6 +87,10 @@ def __getattr__(name):
|
||||
from .model_capabilities import supports_streaming_with_tools
|
||||
_lazy_cache[name] = supports_streaming_with_tools
|
||||
return supports_streaming_with_tools
|
||||
elif name == "max_output_tokens":
|
||||
from .model_capabilities import max_output_tokens
|
||||
_lazy_cache[name] = max_output_tokens
|
||||
return max_output_tokens
|
||||
elif name == "ModelRouter":
|
||||
from .model_router import ModelRouter
|
||||
_lazy_cache[name] = ModelRouter
|
||||
@@ -243,6 +247,7 @@ __all__ = [
|
||||
"process_stream_chunks",
|
||||
"supports_structured_outputs",
|
||||
"supports_streaming_with_tools",
|
||||
"max_output_tokens",
|
||||
"ModelRouter",
|
||||
"ModelProfile",
|
||||
"TaskComplexity",
|
||||
|
||||
@@ -10,8 +10,70 @@ and integrates with Gap 2 (parallel tool execution).
|
||||
|
||||
from ..protocols import LLMProviderAdapterProtocol
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def resolve_ollama_chained_arguments(
|
||||
arguments: Dict[str, Any],
|
||||
tool_result_mapping: Dict[str, Any],
|
||||
*,
|
||||
correlation_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Resolve exact Ollama function-name references without coercing values.
|
||||
|
||||
Weak local models sometimes emit the name of an earlier function as a
|
||||
later argument. Keep the complete result (including mappings, lists,
|
||||
negative numbers, and decimals) so the destination tool receives exactly
|
||||
what the first tool returned.
|
||||
"""
|
||||
if not tool_result_mapping:
|
||||
return arguments
|
||||
resolved = dict(arguments)
|
||||
for arg_name, arg_value in resolved.items():
|
||||
if isinstance(arg_value, str) and arg_value in tool_result_mapping:
|
||||
replacement = tool_result_mapping[arg_value]
|
||||
resolved[arg_name] = replacement
|
||||
extra = {"correlation_id": correlation_id} if correlation_id else {}
|
||||
_logger.debug(
|
||||
"[OLLAMA_FIX] Replaced %s with a prior result in %s arguments",
|
||||
arg_value,
|
||||
arg_name,
|
||||
extra=extra,
|
||||
)
|
||||
return resolved
|
||||
|
||||
|
||||
def record_ollama_tool_result(
|
||||
tool_result_mapping: Dict[str, Any],
|
||||
function_name: str,
|
||||
tool_result: Any,
|
||||
) -> None:
|
||||
"""Record an Ollama tool result exactly as returned by the callback."""
|
||||
if not function_name:
|
||||
return
|
||||
# Some execution paths return a ToolResult-like object rather than the
|
||||
# callback's raw payload. Never expose its failure diagnostic as a
|
||||
# successful same-turn value.
|
||||
if getattr(tool_result, "error", None) is not None:
|
||||
return
|
||||
# The agent executor represents failures as an ``error`` mapping (or a
|
||||
# one-item list containing one). Never feed that diagnostic back as a
|
||||
# successful value to a dependent call in the same turn.
|
||||
if isinstance(tool_result, dict) and tool_result.get("error"):
|
||||
return
|
||||
if (
|
||||
isinstance(tool_result, list)
|
||||
and tool_result
|
||||
and isinstance(tool_result[0], dict)
|
||||
and tool_result[0].get("error")
|
||||
):
|
||||
return
|
||||
tool_result_mapping[function_name] = tool_result
|
||||
|
||||
|
||||
def collapse_union_param_types(tools):
|
||||
"""Rewrite tool schemas a local server cannot parse.
|
||||
@@ -55,42 +117,147 @@ def collapse_union_param_types(tools):
|
||||
return [fix(tool) for tool in tools]
|
||||
|
||||
|
||||
def _advertised_tool_names(tools: List[Dict[str, Any]]) -> set:
|
||||
"""The set of tool names the model was actually offered.
|
||||
|
||||
Salvage must never invent a tool: a text-format call is only promoted when
|
||||
its name matches something advertised, so a model hallucinating a plausible
|
||||
name in prose cannot trigger an unexpected execution.
|
||||
"""
|
||||
names = set()
|
||||
for tool in tools or []:
|
||||
if isinstance(tool, dict):
|
||||
name = tool.get("function", {}).get("name") or tool.get("name")
|
||||
if name:
|
||||
names.add(name)
|
||||
return names
|
||||
|
||||
|
||||
_FENCE_RE = re.compile(
|
||||
r"(?P<fence>`{3,}|~{3,})" # opening fence: >=3 backticks or tildes
|
||||
r"[\s\S]*?" # body (lazy)
|
||||
r"(?:(?P=fence)|\Z)", # matching closing fence OR end-of-string
|
||||
re.MULTILINE,
|
||||
)
|
||||
|
||||
|
||||
def _mask_protected_ranges(text: str) -> str:
|
||||
"""Blank out fenced code blocks so a call *shown* as an example is not run.
|
||||
|
||||
Only the span content is replaced (with spaces of equal length) -- offsets
|
||||
are preserved, so nothing outside a fence shifts. Legitimate prose that
|
||||
merely contains angle brackets is untouched because it is not a tool-call
|
||||
dialect; the fence guard specifically protects documentation of a call.
|
||||
|
||||
Both Markdown fence styles are recognised (``` and ~~~), and an *unclosed*
|
||||
fence is treated as running to end-of-text: a truncated example that opens
|
||||
a fence and names a real tool must not slip through as an executable call.
|
||||
"""
|
||||
def blank(match):
|
||||
return " " * len(match.group(0))
|
||||
return _FENCE_RE.sub(blank, text)
|
||||
|
||||
|
||||
def _make_tool_call(name: str, arguments: Any, seed: str, idx: int) -> Dict[str, Any]:
|
||||
return {
|
||||
"id": f"call_{name}_{idx}_{hash(seed) % 10000}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"arguments": json.dumps(arguments) if isinstance(arguments, (dict, list)) else str(arguments),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# <tool_call>{"name": ..., "arguments": ...}</tool_call>
|
||||
_TOOL_CALL_TAG_RE = re.compile(r"<tool_call>\s*(\{[\s\S]*?\})\s*</tool_call>", re.IGNORECASE)
|
||||
# <function=name>{...}</function>
|
||||
_FUNCTION_TAG_RE = re.compile(r"<function=([a-zA-Z0-9_\-]+)>\s*(\{[\s\S]*?\})\s*</function>", re.IGNORECASE)
|
||||
|
||||
|
||||
def _recover_json_tool_calls(response_text: str, tools: List[Dict[str, Any]]) -> Optional[List[Dict[str, Any]]]:
|
||||
"""Recover a tool call a local model emitted as JSON text.
|
||||
"""Salvage a tool call a local model emitted as text instead of natively.
|
||||
|
||||
Shared by every locally-served engine: small models routinely answer with
|
||||
the call as content instead of using the tool_calls field. Deliberately NOT
|
||||
on DefaultAdapter -- a hosted model returning JSON prose must never have it
|
||||
parsed as a tool call.
|
||||
|
||||
Three dialects are recovered, in-process, with no extra LLM round trip:
|
||||
1. the whole content is a JSON object/array of ``{"name", "arguments"}``;
|
||||
2. one or more ``<tool_call>{...}</tool_call>`` blocks embedded in prose;
|
||||
3. one or more ``<function=name>{...}</function>`` blocks embedded in prose.
|
||||
|
||||
Every recovered name is checked against the advertised tool allowlist so a
|
||||
call is never invented, and fenced code blocks are masked first so a call
|
||||
shown as a documentation example is never executed.
|
||||
"""
|
||||
if not response_text or not tools:
|
||||
return None
|
||||
|
||||
|
||||
allowed = _advertised_tool_names(tools)
|
||||
if not allowed:
|
||||
return None
|
||||
|
||||
# Fast path: a clean full-content JSON payload (the common case, unchanged).
|
||||
try:
|
||||
import json
|
||||
response_json = json.loads(response_text.strip())
|
||||
|
||||
# Normalize to list so both single and multi-tool payloads are supported
|
||||
if isinstance(response_json, dict):
|
||||
response_json = [response_json]
|
||||
|
||||
if isinstance(response_json, list):
|
||||
tool_calls: List[Dict[str, Any]] = []
|
||||
for idx, tool_json in enumerate(response_json):
|
||||
if isinstance(tool_json, dict) and "name" in tool_json:
|
||||
tool_calls.append({
|
||||
"id": f"call_{tool_json['name']}_{idx}_{hash(response_text) % 10000}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool_json["name"],
|
||||
"arguments": json.dumps(tool_json.get("arguments", {}))
|
||||
}
|
||||
})
|
||||
return tool_calls if tool_calls else None
|
||||
name = tool_json.get("name") if isinstance(tool_json, dict) else None
|
||||
if name:
|
||||
tool_calls.append(_make_tool_call(
|
||||
name, tool_json.get("arguments", {}), response_text, idx))
|
||||
if tool_calls:
|
||||
return tool_calls
|
||||
except (json.JSONDecodeError, TypeError, KeyError):
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
# Embedded dialects: only scan when a marker is present, and only outside
|
||||
# fenced code so an example call in documentation is never run. The marker
|
||||
# probe is case-insensitive to match the tag regexes -- a model that emits
|
||||
# <TOOL_CALL> or <FUNCTION=...> must recover just like the lowercase form.
|
||||
lowered = response_text.lower()
|
||||
if not any(marker in lowered for marker in ("<tool_call>", "<function=")):
|
||||
return None
|
||||
|
||||
scan_text = _mask_protected_ranges(response_text)
|
||||
|
||||
# Scan both dialects in a single pass ordered by position in the text, so a
|
||||
# model that interleaves <tool_call> and <function=> blocks yields calls in
|
||||
# the order it actually wrote them -- a dependent or side-effecting call is
|
||||
# never reordered ahead of one it relies on.
|
||||
matches = []
|
||||
for match in _TOOL_CALL_TAG_RE.finditer(scan_text):
|
||||
try:
|
||||
data = json.loads(match.group(1))
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
continue
|
||||
if isinstance(data, dict):
|
||||
name = data.get("name") or data.get("tool")
|
||||
if name in allowed:
|
||||
args = data.get("arguments", data.get("parameters", {}))
|
||||
matches.append((match.start(), name, args))
|
||||
|
||||
for match in _FUNCTION_TAG_RE.finditer(scan_text):
|
||||
name = match.group(1)
|
||||
if name in allowed:
|
||||
try:
|
||||
args = json.loads(match.group(2))
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
continue
|
||||
matches.append((match.start(), name, args))
|
||||
|
||||
matches.sort(key=lambda m: m[0])
|
||||
salvaged = [
|
||||
_make_tool_call(name, args, response_text, idx)
|
||||
for idx, (_, name, args) in enumerate(matches)
|
||||
]
|
||||
|
||||
return salvaged if salvaged else None
|
||||
|
||||
|
||||
class DefaultAdapter:
|
||||
@@ -100,6 +267,25 @@ class DefaultAdapter:
|
||||
"""Hosted providers get the schema untouched -- they support all of it."""
|
||||
return tools
|
||||
|
||||
def resolve_chained_arguments(
|
||||
self,
|
||||
arguments: Dict[str, Any],
|
||||
tool_result_mapping: Dict[str, Any],
|
||||
*,
|
||||
correlation_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Provider hook for same-turn tool-result references."""
|
||||
return arguments
|
||||
|
||||
def record_tool_result(
|
||||
self,
|
||||
tool_result_mapping: Dict[str, Any],
|
||||
function_name: str,
|
||||
tool_result: Any,
|
||||
) -> None:
|
||||
"""Provider hook for recording results used by later tool calls."""
|
||||
return None
|
||||
|
||||
def supports_prompt_caching(self) -> bool:
|
||||
return False
|
||||
|
||||
@@ -178,6 +364,25 @@ class OllamaAdapter(DefaultAdapter):
|
||||
def supports_streaming_with_tools(self) -> bool:
|
||||
# Ollama doesn't reliably support streaming with tools
|
||||
return False
|
||||
|
||||
def resolve_chained_arguments(
|
||||
self,
|
||||
arguments: Dict[str, Any],
|
||||
tool_result_mapping: Dict[str, Any],
|
||||
*,
|
||||
correlation_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
return resolve_ollama_chained_arguments(
|
||||
arguments, tool_result_mapping, correlation_id=correlation_id
|
||||
)
|
||||
|
||||
def record_tool_result(
|
||||
self,
|
||||
tool_result_mapping: Dict[str, Any],
|
||||
function_name: str,
|
||||
tool_result: Any,
|
||||
) -> None:
|
||||
record_ollama_tool_result(tool_result_mapping, function_name, tool_result)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -588,12 +588,48 @@ def get_retry_delay(category: ErrorCategory, attempt: int = 1, base_delay: float
|
||||
return 0
|
||||
|
||||
|
||||
def extract_retry_after(error: Exception) -> Optional[float]:
|
||||
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]:
|
||||
"""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
|
||||
"""
|
||||
@@ -608,14 +644,21 @@ def extract_retry_after(error: Exception) -> Optional[float]:
|
||||
retry_after = None
|
||||
if retry_after is not None:
|
||||
try:
|
||||
return min(float(retry_after), 300.0) # Cap at 5 minutes
|
||||
return min(float(retry_after), cap_seconds)
|
||||
except (ValueError, TypeError):
|
||||
pass # Not a plain number (could be an HTTP-date); fall through
|
||||
# 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)
|
||||
|
||||
# 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), 300.0)
|
||||
return min(float(retry_after_attr), cap_seconds)
|
||||
|
||||
error_str = str(error)
|
||||
|
||||
@@ -632,7 +675,7 @@ def extract_retry_after(error: Exception) -> Optional[float]:
|
||||
if match:
|
||||
try:
|
||||
delay = float(match.group(1))
|
||||
return min(delay, 300.0) # Cap at 5 minutes
|
||||
return min(delay, cap_seconds)
|
||||
except (ValueError, IndexError):
|
||||
continue
|
||||
|
||||
|
||||
@@ -359,10 +359,17 @@ class FailoverManager:
|
||||
if profile.cooldown_until and current_time >= profile.cooldown_until:
|
||||
profile.reset()
|
||||
|
||||
# Find first available profile
|
||||
for profile in self._profiles:
|
||||
if profile.is_available:
|
||||
return profile
|
||||
# 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]
|
||||
|
||||
# If none available, return the one with shortest remaining cooldown
|
||||
available_profiles = [p for p in self._profiles if p.status != ProviderStatus.DISABLED]
|
||||
@@ -443,6 +450,10 @@ 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.
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -38,6 +38,27 @@ def _base_model_name(model_name: str) -> str:
|
||||
return name
|
||||
|
||||
|
||||
def _coerce_output_token_limit(value) -> int | None:
|
||||
"""Return a positive integer output ceiling, or ``None`` when invalid."""
|
||||
if isinstance(value, bool) or value is None:
|
||||
return None
|
||||
try:
|
||||
coerced = int(value)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return None
|
||||
# ``int`` truncates fractional metadata (for example, ``8192.5``), which
|
||||
# would turn an invalid ceiling into an apparently authoritative limit.
|
||||
# Keep numeric strings such as ``"8192"`` compatible with LiteLLM's
|
||||
# metadata while rejecting non-integral numeric values before conversion.
|
||||
if not isinstance(value, str):
|
||||
try:
|
||||
if value != coerced:
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
return coerced if coerced > 0 else None
|
||||
|
||||
|
||||
def _fallback_supports_structured_outputs(model_name: str) -> bool:
|
||||
"""Static heuristic used only when litellm is unavailable.
|
||||
|
||||
@@ -577,8 +598,8 @@ def max_output_tokens(model_name: str):
|
||||
try:
|
||||
info = litellm.get_model_info(model=model_name)
|
||||
if info:
|
||||
out = info.get("max_output_tokens")
|
||||
if out:
|
||||
out = _coerce_output_token_limit(info.get("max_output_tokens"))
|
||||
if out is not None:
|
||||
return out
|
||||
except Exception:
|
||||
pass
|
||||
@@ -594,7 +615,7 @@ def max_output_tokens(model_name: str):
|
||||
if info is None and "/" in model_name:
|
||||
info = model_cost.get(model_name.split("/")[-1].lower())
|
||||
if info:
|
||||
return info.get("max_output_tokens")
|
||||
return _coerce_output_token_limit(info.get("max_output_tokens"))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -617,4 +638,4 @@ def is_gemini_internal_tool(tool) -> bool:
|
||||
if isinstance(tool, dict) and len(tool) == 1:
|
||||
tool_name = next(iter(tool.keys()))
|
||||
return tool_name in GEMINI_INTERNAL_TOOLS
|
||||
return False
|
||||
return False
|
||||
@@ -781,7 +781,9 @@ class OpenAIClient:
|
||||
"output": msg.get("content", ""),
|
||||
})
|
||||
else:
|
||||
item = dict(msg)
|
||||
# Persisted history may carry local metadata/provider fields.
|
||||
# Responses message items accept role/content, not that record.
|
||||
item = {"role": role}
|
||||
item["content"] = self._build_responses_content(
|
||||
msg.get("content", "")
|
||||
)
|
||||
@@ -2654,13 +2656,23 @@ 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:
|
||||
results_str = f"Error executing function: {str(e)}"
|
||||
tool_result = {"error": 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,7 +210,27 @@ 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'):
|
||||
@@ -316,11 +336,12 @@ class MCP:
|
||||
```
|
||||
"""
|
||||
|
||||
# Process-level registry of sanitized MCP server names that have been
|
||||
# namespaced via with_tool_prefix(), mirroring how tools/registry.py tracks
|
||||
# tool names. Lets skills' CapabilityValidator discover connected servers
|
||||
# instead of always failing closed (issue #3307).
|
||||
_active_server_names: set = set()
|
||||
# Process-level registry backing skills' CapabilityValidator discovery
|
||||
# (issue #3307): name -> number of live MCP clients holding that name.
|
||||
# Counted rather than collected, because two live clients may legitimately
|
||||
# declare the same server name. A name is released, and stops satisfying a
|
||||
# skill requirement, once the last client holding it lets go (issue #5135).
|
||||
_active_server_names: dict = {}
|
||||
_active_server_names_lock = threading.Lock()
|
||||
|
||||
# Default init/tool-call timeout (seconds). Used when the caller does not
|
||||
@@ -338,9 +359,54 @@ class MCP:
|
||||
|
||||
@classmethod
|
||||
def list_active_server_names(cls) -> set:
|
||||
"""Return the set of sanitized names of MCP servers namespaced this run."""
|
||||
"""Return the set of MCP server names currently registered this run."""
|
||||
with cls._active_server_names_lock:
|
||||
return set(cls._active_server_names)
|
||||
return {name for name, count in cls._active_server_names.items() if count > 0}
|
||||
|
||||
@classmethod
|
||||
def _adjust_server_name_counts(cls, names: tuple, delta: int) -> None:
|
||||
"""Add ``delta`` to each name's live count, dropping names reaching zero.
|
||||
|
||||
Must be called while holding :attr:`_active_server_names_lock`.
|
||||
"""
|
||||
for name in names:
|
||||
remaining = cls._active_server_names.get(name, 0) + delta
|
||||
if remaining > 0:
|
||||
cls._active_server_names[name] = remaining
|
||||
else:
|
||||
cls._active_server_names.pop(name, None)
|
||||
|
||||
def _adopt_registered_names(self, names: tuple) -> None:
|
||||
"""Make ``names`` this instance's registration, releasing any previous ones.
|
||||
|
||||
The release, the registration and the ownership update share one
|
||||
critical section, so a concurrent :meth:`shutdown` cannot observe the
|
||||
previous names after they have been handed over and release them twice.
|
||||
"""
|
||||
with type(self)._active_server_names_lock:
|
||||
previous = getattr(self, "_registered_server_names", None)
|
||||
if previous == names:
|
||||
return
|
||||
if previous:
|
||||
type(self)._adjust_server_name_counts(previous, -1)
|
||||
if names:
|
||||
type(self)._adjust_server_name_counts(names, 1)
|
||||
self._registered_server_names = names or None
|
||||
|
||||
def _release_registered_names(self) -> None:
|
||||
"""Release this instance's registration, decrementing it exactly once.
|
||||
|
||||
Claiming the names and decrementing them share one critical section, so
|
||||
concurrent lifecycle calls on the same instance (explicit
|
||||
``shutdown()``, context-manager exit, agent cleanup, ``__del__``) cannot
|
||||
each consume a shared client's count.
|
||||
"""
|
||||
with type(self)._active_server_names_lock:
|
||||
names = getattr(self, "_registered_server_names", None)
|
||||
if not names:
|
||||
return
|
||||
self._registered_server_names = None
|
||||
type(self)._adjust_server_name_counts(names, -1)
|
||||
|
||||
@classmethod
|
||||
def _is_cold_start_launcher(cls, cmd) -> bool:
|
||||
@@ -355,8 +421,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, **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, max_response_bytes: Optional[int] = None, **kwargs):
|
||||
"""
|
||||
Initialize the MCP connection and get tools.
|
||||
|
||||
@@ -376,6 +442,7 @@ 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:
|
||||
@@ -393,6 +460,12 @@ 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)
|
||||
@@ -431,6 +504,11 @@ class MCP:
|
||||
# Optional prefix applied to tool names to avoid cross-server collisions
|
||||
# when multiple MCP servers are loaded together (see load_mcp_tools).
|
||||
self._tool_prefix: Optional[str] = None
|
||||
|
||||
# Server names this instance added to the process-level registry via
|
||||
# with_tool_prefix(). Released by shutdown() so the registry shrinks
|
||||
# when a server goes away (issue #5135).
|
||||
self._registered_server_names: Optional[tuple] = None
|
||||
|
||||
# Check if this is a WebSocket URL (ws:// or wss://)
|
||||
if isinstance(command_or_string, str) and re.match(r'^wss?://', command_or_string):
|
||||
@@ -459,6 +537,8 @@ 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)
|
||||
@@ -482,6 +562,8 @@ 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,
|
||||
@@ -990,16 +1072,24 @@ class MCP:
|
||||
"underscore character."
|
||||
)
|
||||
|
||||
self._tool_prefix = sanitized
|
||||
# Both spellings are registered so a skill requirement matches whichever
|
||||
# it declares. dict.fromkeys collapses them when they are identical, so
|
||||
# one shutdown releases both.
|
||||
names = tuple(dict.fromkeys(n for n in (prefix, sanitized) if n))
|
||||
|
||||
# Record this server in the process-level registry so skills'
|
||||
# CapabilityValidator can discover it (issue #3307). Store both the
|
||||
# original name and its sanitized form so a skill requirement matches
|
||||
# regardless of which spelling it declares.
|
||||
with type(self)._active_server_names_lock:
|
||||
if prefix:
|
||||
type(self)._active_server_names.add(prefix)
|
||||
type(self)._active_server_names.add(sanitized)
|
||||
# Re-applying the same prefix *and* the same spellings is a no-op: the
|
||||
# tools already carry the prefix and the names are already registered.
|
||||
# The spellings are compared too, because a rename that sanitizes to the
|
||||
# same prefix (``"my-server"`` -> ``"my server"``) still changes the
|
||||
# original spelling a skill requirement may name.
|
||||
if (getattr(self, "_tool_prefix", None) == sanitized
|
||||
and getattr(self, "_registered_server_names", None) == names):
|
||||
return self
|
||||
|
||||
# Adopting the new names releases the previous ones in the same critical
|
||||
# section, so a concurrent shutdown cannot release them twice.
|
||||
self._adopt_registered_names(names)
|
||||
self._tool_prefix = sanitized
|
||||
|
||||
# Rename already-generated callable tools. Dispatch inside each
|
||||
# wrapper closes over the original tool name, so only the public
|
||||
@@ -1296,6 +1386,12 @@ class MCP:
|
||||
self.websocket_client.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Drop this server from the process-level registry now that its
|
||||
# connections are gone, so skills' CapabilityValidator stops reporting
|
||||
# it as available (issue #5135). Runs last so the registry only shrinks
|
||||
# once the transports above have actually been torn down.
|
||||
self._release_registered_names()
|
||||
|
||||
def __del__(self):
|
||||
"""Clean up resources when the object is garbage collected.
|
||||
|
||||
@@ -71,6 +71,61 @@ 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."""
|
||||
|
||||
@@ -137,6 +192,8 @@ 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:
|
||||
@@ -490,6 +547,13 @@ 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
|
||||
@@ -510,17 +574,16 @@ class HTTPStreamMCPClient:
|
||||
|
||||
def _initialize(self):
|
||||
"""Initialize the connection and tools."""
|
||||
# Use the global event loop
|
||||
loop = get_event_loop()
|
||||
|
||||
# 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()
|
||||
|
||||
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()
|
||||
|
||||
# Run the initialization in the event loop
|
||||
future = asyncio.run_coroutine_threadsafe(self._async_initialize(), loop)
|
||||
self.tools = future.result(timeout=self.timeout)
|
||||
@@ -542,11 +605,16 @@ class HTTPStreamMCPClient:
|
||||
|
||||
headers = self.options.get('headers') or None
|
||||
|
||||
self._streams_context = streamablehttp_client(
|
||||
url=self.base_url,
|
||||
headers=headers,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
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)
|
||||
|
||||
# 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
|
||||
@@ -689,4 +757,4 @@ class HTTPStreamMCPClient:
|
||||
self._force_cleanup()
|
||||
except Exception:
|
||||
# Never raise exceptions in __del__
|
||||
pass
|
||||
pass
|
||||
@@ -17,6 +17,7 @@ out of the core Memory class while preserving backward compatibility.
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from asyncio import CancelledError
|
||||
from typing import Any, Dict, List, Optional
|
||||
from ..protocols import MemoryProtocol
|
||||
|
||||
@@ -250,6 +251,24 @@ class Mem0MemoryAdapter:
|
||||
return self.mem0_client.get_all(**kwargs)
|
||||
|
||||
|
||||
def _create_chroma_client(chromadb, chroma_settings, path):
|
||||
"""Translate native Chroma panics into catchable memory backend failures."""
|
||||
try:
|
||||
return chromadb.PersistentClient(
|
||||
path=path,
|
||||
settings=chroma_settings(anonymized_telemetry=False, allow_reset=True),
|
||||
)
|
||||
except (Exception, KeyboardInterrupt, SystemExit, GeneratorExit, CancelledError):
|
||||
# Preserve ordinary backend errors and interpreter control-flow signals.
|
||||
raise
|
||||
except BaseException as exc:
|
||||
# Chroma's pyo3 PanicException inherits BaseException, not Exception.
|
||||
raise RuntimeError(
|
||||
f"Chroma persist failed at {path!r} ({type(exc).__name__}: {exc}). "
|
||||
"Pass a fresh directory via memory config rag_db_path."
|
||||
) from exc
|
||||
|
||||
|
||||
class ChromaMemoryAdapter:
|
||||
"""
|
||||
Memory adapter that uses ChromaDB to implement MemoryProtocol.
|
||||
@@ -265,13 +284,7 @@ class ChromaMemoryAdapter:
|
||||
os.makedirs(rag_path, exist_ok=True)
|
||||
|
||||
# Initialize ChromaDB client
|
||||
self.client = chromadb.PersistentClient(
|
||||
path=rag_path,
|
||||
settings=chroma_settings(
|
||||
anonymized_telemetry=False,
|
||||
allow_reset=True
|
||||
)
|
||||
)
|
||||
self.client = _create_chroma_client(chromadb, chroma_settings, rag_path)
|
||||
|
||||
# Initialize collection
|
||||
collection_name = kwargs.get("collection_name", "memory_store")
|
||||
@@ -802,4 +815,4 @@ class DakeraMemoryAdapter:
|
||||
|
||||
def reset_long_term(self) -> None:
|
||||
"""Clear all long-term (``episodic``) memory for the agent."""
|
||||
self._reset_type(self.long_term_type)
|
||||
self._reset_type(self.long_term_type)
|
||||
@@ -73,7 +73,7 @@ class SqliteMemoryAdapter:
|
||||
|
||||
def _get_stm_conn(self):
|
||||
"""Get thread-local short-term memory connection."""
|
||||
if not hasattr(self._local, 'stm_conn'):
|
||||
if getattr(self._local, 'stm_conn', None) is None:
|
||||
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 not hasattr(self._local, 'ltm_conn'):
|
||||
if getattr(self._local, 'ltm_conn', None) is None:
|
||||
self._local.ltm_conn = sqlite3.connect(
|
||||
self.long_db,
|
||||
check_same_thread=False,
|
||||
@@ -272,4 +272,30 @@ class SqliteMemoryAdapter:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._all_connections.clear()
|
||||
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)
|
||||
@@ -266,11 +266,12 @@ 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,18 +500,16 @@ 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:
|
||||
self._local.stm_conn = sqlite3.connect(
|
||||
# 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.short_db,
|
||||
check_same_thread=False, # Allow cross-thread cleanup
|
||||
timeout=30.0 # 30 second timeout for lock contention
|
||||
timeout=30.0, # 30 second timeout for lock contention
|
||||
busy_timeout_ms=30000, # 30 seconds
|
||||
)
|
||||
# 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
|
||||
@@ -522,18 +520,16 @@ 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:
|
||||
self._local.ltm_conn = sqlite3.connect(
|
||||
# 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.long_db,
|
||||
check_same_thread=False, # Allow cross-thread cleanup
|
||||
timeout=30.0 # 30 second timeout for lock contention
|
||||
timeout=30.0, # 30 second timeout for lock contention
|
||||
busy_timeout_ms=30000, # 30 seconds
|
||||
)
|
||||
# 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
|
||||
@@ -588,13 +584,8 @@ class Memory(SearchMixin, MemoryCoreMixin):
|
||||
) from exc
|
||||
|
||||
# Initialize ChromaDB with persistent storage
|
||||
self.chroma_client = chromadb.PersistentClient(
|
||||
path=rag_path,
|
||||
settings=ChromaSettings(
|
||||
anonymized_telemetry=False,
|
||||
allow_reset=True
|
||||
)
|
||||
)
|
||||
from .adapters.factories import _create_chroma_client
|
||||
self.chroma_client = _create_chroma_client(chromadb, ChromaSettings, rag_path)
|
||||
|
||||
collection_name = self.cfg.get("collection_name", "memory_store")
|
||||
self._collection_name = collection_name
|
||||
@@ -2541,16 +2532,31 @@ 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)
|
||||
# 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).
|
||||
if hasattr(self, 'memory_adapter') and self.memory_adapter:
|
||||
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
|
||||
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
|
||||
|
||||
# Close MongoDB client if it exists (legacy path)
|
||||
if hasattr(self, 'mongo_client') and self.mongo_client:
|
||||
|
||||
@@ -132,7 +132,8 @@ class StatusOutput:
|
||||
def agent_start(self, agent_name: str) -> None:
|
||||
"""Record agent start."""
|
||||
ts = time.time()
|
||||
_agent_start_times[agent_name] = ts
|
||||
with _output_lock:
|
||||
_agent_start_times[agent_name] = ts
|
||||
|
||||
if self._format == "jsonl":
|
||||
self._emit_jsonl("agent_start", agent_name=agent_name, timestamp=ts)
|
||||
@@ -144,8 +145,9 @@ class StatusOutput:
|
||||
ts = time.time()
|
||||
|
||||
# Calculate duration if not provided
|
||||
if duration_ms is None and agent_name in _agent_start_times:
|
||||
start_ts = _agent_start_times.pop(agent_name, None)
|
||||
if duration_ms is None:
|
||||
with _output_lock:
|
||||
start_ts = _agent_start_times.pop(agent_name, None)
|
||||
if start_ts:
|
||||
duration_ms = (ts - start_ts) * 1000
|
||||
|
||||
@@ -160,11 +162,13 @@ class StatusOutput:
|
||||
"""Record LLM call start."""
|
||||
global _ai_call_count
|
||||
ts = time.time()
|
||||
self._llm_start_time = ts
|
||||
_ai_call_count += 1
|
||||
with _output_lock:
|
||||
self._llm_start_time = ts
|
||||
_ai_call_count += 1
|
||||
call_count = _ai_call_count
|
||||
|
||||
# Context based on call sequence
|
||||
if _ai_call_count == 1:
|
||||
if call_count == 1:
|
||||
context = "thinking"
|
||||
else:
|
||||
context = "responding"
|
||||
@@ -191,22 +195,26 @@ class StatusOutput:
|
||||
if latency_ms is not None and latency_ms > 0:
|
||||
duration_ms = latency_ms
|
||||
elif duration_ms is None and hasattr(self, '_llm_start_time'):
|
||||
start_ts = self._llm_start_time
|
||||
with _output_lock:
|
||||
start_ts = self._llm_start_time
|
||||
if start_ts:
|
||||
duration_ms = (ts - start_ts) * 1000
|
||||
|
||||
# Track session totals for summary
|
||||
if not hasattr(self, '_session_tokens_in'):
|
||||
self._session_tokens_in = 0
|
||||
self._session_tokens_out = 0
|
||||
self._session_cost = 0.0
|
||||
self._session_llm_calls = 0
|
||||
|
||||
self._session_tokens_in += tokens_in
|
||||
self._session_tokens_out += tokens_out
|
||||
if cost:
|
||||
self._session_cost += cost
|
||||
self._session_llm_calls += 1
|
||||
# Track session totals for summary (shared instance across all agents,
|
||||
# so guard the read-modify-write with the module lock to avoid
|
||||
# under-counting when two LLM calls land concurrently).
|
||||
with _output_lock:
|
||||
if not hasattr(self, '_session_tokens_in'):
|
||||
self._session_tokens_in = 0
|
||||
self._session_tokens_out = 0
|
||||
self._session_cost = 0.0
|
||||
self._session_llm_calls = 0
|
||||
|
||||
self._session_tokens_in += tokens_in
|
||||
self._session_tokens_out += tokens_out
|
||||
if cost:
|
||||
self._session_cost += cost
|
||||
self._session_llm_calls += 1
|
||||
|
||||
# Only show metrics line in debug mode (when show_metrics is enabled)
|
||||
show_metrics = getattr(self, '_show_metrics', False)
|
||||
@@ -227,7 +235,8 @@ class StatusOutput:
|
||||
def tool_start(self, tool_name: str, tool_args: Optional[Dict[str, Any]] = None, agent_name: Optional[str] = None) -> None:
|
||||
"""Record tool start - stores info for inline display with result."""
|
||||
ts = time.time()
|
||||
self._tool_start_times[tool_name] = ts
|
||||
with _output_lock:
|
||||
self._tool_start_times[tool_name] = ts
|
||||
self._pending_tool_args = tool_args # Store for display with result
|
||||
self._pending_tool_name = tool_name
|
||||
|
||||
@@ -246,8 +255,9 @@ class StatusOutput:
|
||||
ts = time.time()
|
||||
|
||||
# Calculate duration if not provided
|
||||
if duration_ms is None and tool_name in self._tool_start_times:
|
||||
start_ts = self._tool_start_times.pop(tool_name, None)
|
||||
if duration_ms is None:
|
||||
with _output_lock:
|
||||
start_ts = self._tool_start_times.pop(tool_name, None)
|
||||
if start_ts:
|
||||
duration_ms = (ts - start_ts) * 1000
|
||||
|
||||
|
||||
@@ -95,7 +95,8 @@ class TraceOutput:
|
||||
|
||||
def llm_start(self, model: str = None) -> None:
|
||||
"""Record LLM call start."""
|
||||
self._llm_start_time = time.time()
|
||||
with _output_lock:
|
||||
self._llm_start_time = time.time()
|
||||
model_str = f" ({model})" if model else ""
|
||||
self._emit(f"Calling LLM{model_str}...", "cyan")
|
||||
|
||||
@@ -112,8 +113,12 @@ class TraceOutput:
|
||||
# Otherwise calculate from internal tracking
|
||||
if duration_ms is not None and duration_ms > 0:
|
||||
pass # Use the passed value
|
||||
elif self._llm_start_time:
|
||||
duration_ms = (time.time() - self._llm_start_time) * 1000
|
||||
else:
|
||||
with _output_lock:
|
||||
start_ts = self._llm_start_time
|
||||
self._llm_start_time = None
|
||||
if start_ts:
|
||||
duration_ms = (time.time() - start_ts) * 1000
|
||||
|
||||
duration_str = f" [{duration_ms/1000:.1f}s]" if duration_ms and duration_ms > 0 else ""
|
||||
|
||||
@@ -125,11 +130,11 @@ class TraceOutput:
|
||||
self._emit(metrics_line, "dim")
|
||||
else:
|
||||
self._emit(f"LLM responded{duration_str}", "green")
|
||||
self._llm_start_time = None
|
||||
|
||||
def tool_start(self, tool_name: str, tool_args: Dict[str, Any] = None) -> None:
|
||||
"""Record tool execution start."""
|
||||
self._tool_start_times[tool_name] = time.time()
|
||||
with _output_lock:
|
||||
self._tool_start_times[tool_name] = time.time()
|
||||
|
||||
# Format args compactly
|
||||
args_str = ""
|
||||
@@ -153,8 +158,9 @@ class TraceOutput:
|
||||
) -> None:
|
||||
"""Record tool execution end."""
|
||||
# Calculate duration if not provided
|
||||
if duration_ms is None and tool_name in self._tool_start_times:
|
||||
start_ts = self._tool_start_times.pop(tool_name, None)
|
||||
if duration_ms is None:
|
||||
with _output_lock:
|
||||
start_ts = self._tool_start_times.pop(tool_name, None)
|
||||
if start_ts:
|
||||
duration_ms = (time.time() - start_ts) * 1000
|
||||
|
||||
|
||||
@@ -62,6 +62,13 @@ ARITY: Dict[str, int] = {
|
||||
"apt-get": 2,
|
||||
"brew": 2,
|
||||
"systemctl": 2,
|
||||
# Windows / PowerShell wrappers: keep the wrapper + inner verb so a
|
||||
# ``powershell -Command Get-ChildItem`` generalises to ``powershell *``
|
||||
# rather than staying fully literal (the inner cmdlet is gated separately by
|
||||
# the dialect-aware command parser).
|
||||
"powershell": 1,
|
||||
"pwsh": 1,
|
||||
"cmd": 1,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -13,12 +13,22 @@ Design goals:
|
||||
- Best-effort and conservative: on any parse failure we fall back to treating
|
||||
the whole command as a single operation (today's behaviour) so a rule is
|
||||
never silently weakened.
|
||||
|
||||
Windows dialects (PowerShell / ``cmd``): a ``powershell -Command "…"`` /
|
||||
``pwsh -c "…"`` / ``cmd /c "…"`` wrapper otherwise collapses to a single
|
||||
opaque op, hiding the inner file-mutating cmdlet from the workspace-boundary
|
||||
gate. Such wrappers are unwrapped and their inner command string is parsed so
|
||||
the same boundary check that fires for POSIX mutations also fires on Windows.
|
||||
POSIX parsing is byte-for-byte unchanged (the wrapper unwrap only triggers for
|
||||
the ``powershell``/``pwsh``/``cmd`` executables).
|
||||
"""
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import re
|
||||
import shlex
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
# Operators that separate simple-commands within a compound command.
|
||||
@@ -59,6 +69,90 @@ def has_unresolvable_expansion(cmd: str) -> bool:
|
||||
_WRITE_REDIRECTS = (">", ">>", ">|", "&>", "&>>")
|
||||
|
||||
|
||||
# Supported shell dialects for parsing. ``posix`` is the default and unchanged
|
||||
# path; ``powershell`` and ``cmd`` recognise Windows command shape.
|
||||
DIALECT_POSIX = "posix"
|
||||
DIALECT_POWERSHELL = "powershell"
|
||||
DIALECT_CMD = "cmd"
|
||||
|
||||
|
||||
# Wrapper executables whose quoted ``-Command``/``-c``/``/c`` argument carries an
|
||||
# inner command string that must be parsed in the matching dialect so the inner
|
||||
# cmdlet/built-in is gated rather than the opaque wrapper.
|
||||
_POWERSHELL_WRAPPERS = ("powershell", "powershell.exe", "pwsh", "pwsh.exe")
|
||||
_CMD_WRAPPERS = ("cmd", "cmd.exe")
|
||||
|
||||
|
||||
# File-mutating PowerShell cmdlets, their common aliases, and ``cmd`` built-ins.
|
||||
# Used by the permission engine's file-mutation classifier so the external-
|
||||
# directory boundary check recognises Windows mutations the same way it does
|
||||
# POSIX ones (``rm``/``cp``/``mv``/…). Matched case-insensitively.
|
||||
_POWERSHELL_MUTATING = frozenset(
|
||||
{
|
||||
"remove-item",
|
||||
"new-item",
|
||||
"move-item",
|
||||
"copy-item",
|
||||
"rename-item",
|
||||
"set-content",
|
||||
"add-content",
|
||||
"clear-content",
|
||||
"out-file",
|
||||
# PowerShell aliases for the cmdlets above.
|
||||
"ri",
|
||||
"rd",
|
||||
"del",
|
||||
"erase",
|
||||
"ni",
|
||||
"mi",
|
||||
"move",
|
||||
"cpi",
|
||||
"copy",
|
||||
"cp",
|
||||
"rni",
|
||||
"ren",
|
||||
"sc",
|
||||
"ac",
|
||||
"rm",
|
||||
"mv",
|
||||
}
|
||||
)
|
||||
|
||||
_CMD_MUTATING = frozenset(
|
||||
{
|
||||
"del",
|
||||
"erase",
|
||||
"rd",
|
||||
"rmdir",
|
||||
"md",
|
||||
"mkdir",
|
||||
"move",
|
||||
"copy",
|
||||
"xcopy",
|
||||
"ren",
|
||||
"rename",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def is_mutating_executable(executable: str, dialect: str = DIALECT_POSIX) -> bool:
|
||||
"""Return ``True`` if *executable* is a file-mutating command for *dialect*.
|
||||
|
||||
POSIX callers keep their existing (executable-agnostic, path-based) boundary
|
||||
behaviour; this helper is provided so Windows dialects can recognise
|
||||
PowerShell cmdlets/aliases and ``cmd`` built-ins by name. Matching is
|
||||
case-insensitive and ignores any directory prefix on the executable.
|
||||
"""
|
||||
if not executable:
|
||||
return False
|
||||
name = executable.rsplit("/", 1)[-1].rsplit("\\", 1)[-1].lower()
|
||||
if dialect == DIALECT_POWERSHELL:
|
||||
return name in _POWERSHELL_MUTATING
|
||||
if dialect == DIALECT_CMD:
|
||||
return name in _CMD_MUTATING
|
||||
return False
|
||||
|
||||
|
||||
def _looks_like_path(tok: str) -> bool:
|
||||
"""Return ``True`` if *tok* is path-like enough to warrant a boundary check.
|
||||
|
||||
@@ -81,6 +175,30 @@ def _looks_like_path(tok: str) -> bool:
|
||||
)
|
||||
|
||||
|
||||
# Drive-letter absolute path (``C:\x`` / ``C:/x``) or UNC share (``\\host\x``).
|
||||
_WINDOWS_ABS_RE = re.compile(r"^(?:[A-Za-z]:[\\/]|\\\\)")
|
||||
|
||||
|
||||
def _looks_like_windows_path(tok: str) -> bool:
|
||||
"""Return ``True`` if *tok* is a Windows/PowerShell path shape.
|
||||
|
||||
Recognises drive-letter absolutes (``C:\\x``), UNC shares (``\\\\host\\x``),
|
||||
explicit relatives (``.\\x``, ``..\\x``) and any bare token embedding a
|
||||
backslash separator, in addition to the POSIX shapes. Used only for
|
||||
``powershell``/``cmd`` ops so POSIX boundary behaviour is unchanged.
|
||||
"""
|
||||
if not tok:
|
||||
return False
|
||||
if _looks_like_path(tok):
|
||||
return True
|
||||
return (
|
||||
bool(_WINDOWS_ABS_RE.match(tok))
|
||||
or tok.startswith(".\\")
|
||||
or tok.startswith("..\\")
|
||||
or "\\" in tok
|
||||
)
|
||||
|
||||
|
||||
def _extract_flag_path(tok: str) -> str:
|
||||
"""Extract a path operand attached to a short flag (``-o/tmp/x``).
|
||||
|
||||
@@ -110,11 +228,14 @@ class ShellOp:
|
||||
executable: The command name (e.g. ``rm``), or ``""`` if unknown.
|
||||
args: Positional/flag arguments following the executable.
|
||||
write_targets: Paths that are written/truncated via redirection.
|
||||
dialect: The shell dialect this op was parsed under (``posix`` by
|
||||
default; ``powershell``/``cmd`` for unwrapped Windows commands).
|
||||
"""
|
||||
|
||||
executable: str = ""
|
||||
args: List[str] = field(default_factory=list)
|
||||
write_targets: List[str] = field(default_factory=list)
|
||||
dialect: str = DIALECT_POSIX
|
||||
|
||||
@property
|
||||
def command_string(self) -> str:
|
||||
@@ -134,6 +255,13 @@ class ShellOp:
|
||||
(``-o/tmp/x``) are also unwrapped so the value is boundary-checked.
|
||||
Plain flags and non-path values are ignored.
|
||||
"""
|
||||
# POSIX behaviour is unchanged; Windows dialects additionally recognise
|
||||
# drive-letter/UNC/backslash path shapes so cmdlet operands are gated.
|
||||
is_path = (
|
||||
_looks_like_path
|
||||
if self.dialect == DIALECT_POSIX
|
||||
else _looks_like_windows_path
|
||||
)
|
||||
paths: List[str] = []
|
||||
for tok in self.args:
|
||||
if not tok:
|
||||
@@ -143,14 +271,14 @@ class ShellOp:
|
||||
if tok.startswith("-"):
|
||||
if "=" in tok:
|
||||
value = tok.split("=", 1)[1]
|
||||
if value and _looks_like_path(value):
|
||||
if value and is_path(value):
|
||||
paths.append(value)
|
||||
continue
|
||||
attached = _extract_flag_path(tok)
|
||||
if attached:
|
||||
paths.append(attached)
|
||||
continue
|
||||
if _looks_like_path(tok):
|
||||
if is_path(tok):
|
||||
paths.append(tok)
|
||||
return paths
|
||||
|
||||
@@ -275,25 +403,197 @@ def _split_simple_commands(cmd: str) -> List[str]:
|
||||
return segments
|
||||
|
||||
|
||||
def _parse_segment(segment: str) -> List[ShellOp]:
|
||||
def _strip_quotes(tok: str) -> str:
|
||||
"""Strip a single pair of surrounding single/double quotes from *tok*."""
|
||||
if len(tok) >= 2 and tok[0] == tok[-1] and tok[0] in ("'", '"'):
|
||||
return tok[1:-1]
|
||||
return tok
|
||||
|
||||
|
||||
# Leading ``powershell``/``pwsh``/``cmd`` executable of a wrapper segment.
|
||||
# Matched case-insensitively; an optional ``.exe`` suffix is tolerated. The
|
||||
# switches and the command flag that follow are parsed separately (below) so
|
||||
# option-prefixed forms such as ``powershell -NoProfile -Command "…"`` or
|
||||
# ``cmd /d /c "…"`` are recognised, not just the flag-immediately-after form.
|
||||
_WRAPPER_EXE_RE = re.compile(
|
||||
r"^\s*(?P<exe>powershell|pwsh|cmd)(?:\.exe)?\s+(?P<rest>.+)$",
|
||||
re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
|
||||
# PowerShell command flags whose operand is the inner command string.
|
||||
_PS_COMMAND_FLAGS = ("-command", "-c", "-encodedcommand", "-e", "-ec")
|
||||
# ``cmd`` command flags whose operand is the inner command string.
|
||||
_CMD_COMMAND_FLAGS = ("/c", "/k")
|
||||
# PowerShell base64 flags whose operand is a UTF-16LE base64 payload.
|
||||
_PS_ENCODED_FLAGS = ("-encodedcommand", "-e", "-ec")
|
||||
|
||||
|
||||
def _decode_powershell_encoded(payload: str) -> Optional[str]:
|
||||
"""Decode a PowerShell ``-EncodedCommand`` base64 UTF-16LE *payload*.
|
||||
|
||||
Returns the decoded inner command string, or ``None`` if the payload is not
|
||||
valid base64 / UTF-16LE so the caller can fail closed (treat the wrapper as
|
||||
an opaque, un-inspectable op) rather than exposing the base64 token as a
|
||||
bogus executable that a deny rule can never match.
|
||||
"""
|
||||
token = payload.strip().strip("'\"")
|
||||
if not token:
|
||||
return None
|
||||
try:
|
||||
raw = base64.b64decode(token, validate=True)
|
||||
except (binascii.Error, ValueError):
|
||||
return None
|
||||
try:
|
||||
decoded = raw.decode("utf-16-le")
|
||||
except UnicodeDecodeError:
|
||||
return None
|
||||
return decoded.strip() or None
|
||||
|
||||
|
||||
# Leading PowerShell call-operator/scriptblock scaffolding (``& { … }`` /
|
||||
# ``. { … }``) that hides the real cmdlet behind a ``{`` token when tokenized.
|
||||
_PS_SCRIPTBLOCK_RE = re.compile(r"^\s*[&.]\s*\{(?P<body>.*)\}\s*$", re.DOTALL)
|
||||
|
||||
|
||||
def _strip_powershell_scaffolding(inner: str) -> str:
|
||||
"""Strip a leading call-operator + scriptblock wrapper from *inner*.
|
||||
|
||||
PowerShell's ``& { <cmd> }`` / ``. { <cmd> }`` invocation would otherwise
|
||||
tokenize with ``{`` as the executable, demoting the real cmdlet to an
|
||||
argument and letting a command-specific deny slip through. Unwrap the
|
||||
scriptblock body so the inner cmdlet surfaces as the executable. Returns
|
||||
*inner* unchanged when no such wrapper is present.
|
||||
"""
|
||||
stripped = inner.strip()
|
||||
match = _PS_SCRIPTBLOCK_RE.match(stripped)
|
||||
if match is not None:
|
||||
return match.group("body").strip()
|
||||
# Bare leading call operator (``& Remove-Item …``) without a scriptblock.
|
||||
if stripped[:1] in ("&", ".") and stripped[1:2].isspace():
|
||||
return stripped[1:].strip()
|
||||
return stripped
|
||||
|
||||
|
||||
def _split_wrapper_flag(rest: str, dialect: str) -> Optional[tuple]:
|
||||
"""Split a wrapper's post-executable *rest* into (flag, operand).
|
||||
|
||||
Skips leading switches (``-NoProfile``, ``/d``, …) until the command flag
|
||||
(``-Command``/``-c``/``-EncodedCommand`` for PowerShell, ``/c``/``/k`` for
|
||||
``cmd``) is found, then returns that flag plus the verbatim remainder as the
|
||||
inner command operand. Returns ``None`` if no command flag is present.
|
||||
"""
|
||||
command_flags = (
|
||||
_CMD_COMMAND_FLAGS if dialect == DIALECT_CMD else _PS_COMMAND_FLAGS
|
||||
)
|
||||
remaining = rest
|
||||
# Walk whitespace-separated leading tokens; the command flag's operand is
|
||||
# everything after it (captured verbatim to preserve backslash paths).
|
||||
while remaining:
|
||||
stripped = remaining.lstrip()
|
||||
parts = stripped.split(None, 1)
|
||||
token = parts[0]
|
||||
if token.lower() in command_flags:
|
||||
operand = parts[1] if len(parts) > 1 else ""
|
||||
return token.lower(), operand
|
||||
# Not the command flag: skip this switch and continue. A bare
|
||||
# ``powershell script.ps1`` (no command flag) yields ``None`` so the
|
||||
# wrapper stays opaque (fail-closed) instead of misparsing a filename.
|
||||
if not token.startswith("-") and not token.startswith("/"):
|
||||
return None
|
||||
if len(parts) == 1:
|
||||
return None
|
||||
remaining = parts[1]
|
||||
return None
|
||||
|
||||
|
||||
def _opaque_wrapper_op(segment: str, dialect: str) -> List[ShellOp]:
|
||||
"""Fail-closed op for an un-inspectable wrapper (e.g. undecodable base64).
|
||||
|
||||
Preserves the original ``powershell``/``pwsh``/``cmd`` executable so a
|
||||
command-specific deny/allow still applies to the wrapper itself, and no
|
||||
broad wildcard silently authorises an opaque payload.
|
||||
"""
|
||||
return [_fallback_op(segment, dialect)]
|
||||
|
||||
|
||||
def _unwrap_shell_wrapper(segment: str) -> Optional[List[ShellOp]]:
|
||||
"""Return ops for the inner command of a ``powershell``/``cmd`` wrapper.
|
||||
|
||||
Detects ``powershell -Command "…"`` / ``pwsh -c "…"`` / ``cmd /c "…"`` —
|
||||
including option-prefixed forms (``powershell -NoProfile -Command "…"`` /
|
||||
``cmd /d /c "…"``), ``-EncodedCommand`` base64 payloads, and ``& { … }``
|
||||
scriptblock invocations — in the *raw* segment and parses the inner command
|
||||
string in the matching dialect so the inner cmdlet/built-in and its paths
|
||||
are gated rather than the opaque wrapper. Returns ``None`` when *segment* is
|
||||
not such a wrapper, so callers fall back to normal per-segment parsing.
|
||||
Detection is executable-name based, so a POSIX command never triggers this
|
||||
path.
|
||||
"""
|
||||
match = _WRAPPER_EXE_RE.match(segment)
|
||||
if match is None:
|
||||
return None
|
||||
exe = match.group("exe").lower()
|
||||
dialect = DIALECT_CMD if exe == "cmd" else DIALECT_POWERSHELL
|
||||
|
||||
split = _split_wrapper_flag(match.group("rest"), dialect)
|
||||
if split is None:
|
||||
return None
|
||||
flag, operand = split
|
||||
|
||||
inner = _strip_quotes(operand.strip()).strip()
|
||||
if flag in _PS_ENCODED_FLAGS:
|
||||
# Base64 UTF-16LE payload: decode so the real cmdlet is inspected;
|
||||
# fail closed to an opaque wrapper op if it cannot be decoded.
|
||||
decoded = _decode_powershell_encoded(inner)
|
||||
if decoded is None:
|
||||
return _opaque_wrapper_op(segment, dialect)
|
||||
inner = decoded
|
||||
|
||||
if not inner:
|
||||
return None
|
||||
|
||||
if dialect == DIALECT_POWERSHELL:
|
||||
inner = _strip_powershell_scaffolding(inner)
|
||||
if not inner:
|
||||
return None
|
||||
|
||||
return parse_command(inner, dialect=dialect)
|
||||
|
||||
|
||||
def _parse_segment(segment: str, dialect: str = DIALECT_POSIX) -> List[ShellOp]:
|
||||
"""Parse a single simple-command segment into ShellOp(s).
|
||||
|
||||
May return multiple ops when the segment embeds command substitutions.
|
||||
"""
|
||||
ops: List[ShellOp] = []
|
||||
|
||||
# ``powershell -Command "…"`` / ``cmd /c "…"`` carries an inner command that
|
||||
# would otherwise be an opaque op. Unwrap from the *raw* segment (before
|
||||
# shlex) so the inner command — and its backslash paths — are preserved and
|
||||
# parsed in the matching dialect.
|
||||
wrapped = _unwrap_shell_wrapper(segment)
|
||||
if wrapped is not None:
|
||||
return wrapped
|
||||
|
||||
# First, recurse into any command substitutions so e.g. ``$(rm -rf x)``
|
||||
# is evaluated as an ``rm`` operation.
|
||||
for inner in _extract_substitutions(segment):
|
||||
ops.extend(parse_command(inner))
|
||||
ops.extend(parse_command(inner, dialect=dialect))
|
||||
|
||||
# POSIX tokenizes with ``posix=True`` (unchanged). Windows dialects use
|
||||
# ``posix=False`` so backslash path separators (``C:\tmp\x``) are preserved
|
||||
# rather than consumed as escape characters; surrounding quotes are then
|
||||
# stripped from each token.
|
||||
posix_mode = dialect == DIALECT_POSIX
|
||||
try:
|
||||
tokens = shlex.split(segment, comments=False, posix=True)
|
||||
tokens = shlex.split(segment, comments=False, posix=posix_mode)
|
||||
except ValueError:
|
||||
# Unbalanced quotes etc. — conservative fallback: whole segment.
|
||||
tokens = segment.split()
|
||||
if not posix_mode:
|
||||
tokens = [_strip_quotes(t) for t in tokens]
|
||||
|
||||
op = ShellOp()
|
||||
op = ShellOp(dialect=dialect)
|
||||
args: List[str] = []
|
||||
skip_next = False
|
||||
skip_input_target = False
|
||||
@@ -368,11 +668,19 @@ def _parse_segment(segment: str) -> List[ShellOp]:
|
||||
return ops
|
||||
|
||||
|
||||
def parse_command(cmd: str) -> List[ShellOp]:
|
||||
def parse_command(cmd: str, *, dialect: str = DIALECT_POSIX) -> List[ShellOp]:
|
||||
"""Parse a shell command string into a list of ShellOp operations.
|
||||
|
||||
Args:
|
||||
cmd: The raw shell command (without the ``bash:``/``shell:`` prefix).
|
||||
dialect: The shell dialect to parse under — ``posix`` (default,
|
||||
byte-for-byte unchanged), ``powershell`` or ``cmd``. Statement
|
||||
separators (``;``, ``|``, ``&&``/``||``, ``&``) are shared across
|
||||
dialects; the dialect is recorded on each op and drives the
|
||||
file-mutation classifier for Windows cmdlets/built-ins. A
|
||||
``powershell``/``pwsh``/``cmd`` wrapper is unwrapped and its inner
|
||||
command parsed in the matching dialect regardless of the requested
|
||||
*dialect*, so a POSIX gate still sees into a Windows wrapper.
|
||||
|
||||
Returns:
|
||||
A list of ShellOp. On parse failure, returns a single best-effort
|
||||
@@ -385,20 +693,20 @@ def parse_command(cmd: str) -> List[ShellOp]:
|
||||
segments = _split_simple_commands(cmd)
|
||||
ops: List[ShellOp] = []
|
||||
for seg in segments:
|
||||
ops.extend(_parse_segment(seg))
|
||||
ops.extend(_parse_segment(seg, dialect))
|
||||
|
||||
if not ops:
|
||||
# Nothing extracted — fall back to whole command.
|
||||
return [_fallback_op(cmd)]
|
||||
return [_fallback_op(cmd, dialect)]
|
||||
return ops
|
||||
except Exception: # noqa: BLE001 - conservative: never weaken a rule on parse failure
|
||||
# Any unexpected failure must not weaken the rule: fall back.
|
||||
return [_fallback_op(cmd)]
|
||||
return [_fallback_op(cmd, dialect)]
|
||||
|
||||
|
||||
def _fallback_op(cmd: str) -> ShellOp:
|
||||
def _fallback_op(cmd: str, dialect: str = DIALECT_POSIX) -> ShellOp:
|
||||
"""Build a single ShellOp representing the whole command (legacy path)."""
|
||||
tokens = cmd.split()
|
||||
if not tokens:
|
||||
return ShellOp()
|
||||
return ShellOp(executable=tokens[0], args=tokens[1:])
|
||||
return ShellOp(dialect=dialect)
|
||||
return ShellOp(executable=tokens[0], args=tokens[1:], dialect=dialect)
|
||||
@@ -113,12 +113,19 @@ class PermissionManager:
|
||||
logger.warning(f"Failed to load rules: {e}")
|
||||
|
||||
def _save_rules(self):
|
||||
"""Save rules to disk."""
|
||||
"""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
|
||||
path = self._get_rules_path()
|
||||
try:
|
||||
with open(path, "w") as f:
|
||||
json.dump([r.to_dict() for r in self._rules], f, indent=2)
|
||||
except IOError as e:
|
||||
atomic_write_json(path, [r.to_dict() for r in self._rules], indent=2)
|
||||
except (OSError, TypeError, ValueError) as e:
|
||||
logger.error(f"Failed to save rules: {e}")
|
||||
|
||||
def save_rules(self):
|
||||
@@ -140,7 +147,8 @@ class PermissionManager:
|
||||
logger.warning(f"Failed to load approvals: {e}")
|
||||
|
||||
def _save_approvals(self):
|
||||
"""Save approvals to disk."""
|
||||
"""Save approvals to disk atomically (see :meth:`_save_rules`)."""
|
||||
from ..utils.atomic_io import atomic_write_json
|
||||
path = self._get_approvals_path()
|
||||
try:
|
||||
# Filter out expired and one-time approvals
|
||||
@@ -148,9 +156,8 @@ class PermissionManager:
|
||||
a for a in self._approvals
|
||||
if a.is_valid() and a.scope in ("session", "always")
|
||||
]
|
||||
with open(path, "w") as f:
|
||||
json.dump([a.to_dict() for a in persistent], f, indent=2)
|
||||
except IOError as e:
|
||||
atomic_write_json(path, [a.to_dict() for a in persistent], indent=2)
|
||||
except (OSError, TypeError, ValueError) as e:
|
||||
logger.error(f"Failed to save approvals: {e}")
|
||||
|
||||
def add_rule(self, rule: PermissionRule) -> str:
|
||||
|
||||
@@ -108,10 +108,11 @@ class PlanStorage:
|
||||
|
||||
# Convert plan to markdown
|
||||
markdown = plan.to_markdown()
|
||||
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write(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)
|
||||
|
||||
plan.file_path = path
|
||||
logger.debug(f"Saved plan to {path}")
|
||||
|
||||
@@ -227,10 +228,11 @@ class PlanStorage:
|
||||
"""
|
||||
filename = f"{name}.json"
|
||||
path = os.path.join(self.todos_dir, filename)
|
||||
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write(todo.to_json())
|
||||
|
||||
|
||||
# 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())
|
||||
|
||||
logger.debug(f"Saved todo list to {path}")
|
||||
return path
|
||||
|
||||
|
||||
@@ -378,6 +378,21 @@ def get_plugin_registry() -> list:
|
||||
return manager.is_enabled(name)
|
||||
return name in config_enabled
|
||||
|
||||
# Surface the least-privilege conversation boundary so an operator can see,
|
||||
# per plugin, whether it *requested* prompt/message access (declared a
|
||||
# conversation hook) and whether it is *granted*. Imported once and reused
|
||||
# across all three registry sources so the capability columns are never
|
||||
# silently omitted for a given plugin-listing path.
|
||||
from .manager import CONVERSATION_HOOKS
|
||||
|
||||
def _requests_conversation(declared) -> bool:
|
||||
for h in declared or []:
|
||||
value = h.value if hasattr(h, "value") else h
|
||||
for ch in CONVERSATION_HOOKS:
|
||||
if value == ch or value == getattr(ch, "value", ch):
|
||||
return True
|
||||
return False
|
||||
|
||||
entries: list = []
|
||||
seen: set = set()
|
||||
|
||||
@@ -388,9 +403,10 @@ def get_plugin_registry() -> list:
|
||||
if not name or name in seen:
|
||||
continue
|
||||
seen.add(name)
|
||||
declared = getattr(info, "hooks", None) or []
|
||||
hooks = [
|
||||
h.value if hasattr(h, "value") else str(h)
|
||||
for h in (getattr(info, "hooks", None) or [])
|
||||
for h in declared
|
||||
]
|
||||
entries.append({
|
||||
"name": name,
|
||||
@@ -399,6 +415,8 @@ def get_plugin_registry() -> list:
|
||||
"source": "registered",
|
||||
"enabled": manager.is_enabled(name) or _config_says_enabled(name),
|
||||
"hooks": hooks,
|
||||
"requests_conversation": _requests_conversation(declared),
|
||||
"conversation_granted": manager._granted_conversation_access(name),
|
||||
})
|
||||
|
||||
# 2. Entry-point plugins present on the system but not yet loaded, so
|
||||
@@ -415,6 +433,10 @@ def get_plugin_registry() -> list:
|
||||
continue
|
||||
seen.add(ep.name)
|
||||
dist = getattr(getattr(ep, "dist", None), "name", None)
|
||||
# Installed-but-not-loaded entry points: hooks are unknown without
|
||||
# executing the plugin (which we deliberately avoid), so we report a
|
||||
# conservative view — it has not been granted conversation access and
|
||||
# its requested capability is unknown until loaded.
|
||||
entries.append({
|
||||
"name": ep.name,
|
||||
"version": "-",
|
||||
@@ -422,6 +444,8 @@ def get_plugin_registry() -> list:
|
||||
"source": f"entry_point:{dist}" if dist else "entry_point",
|
||||
"enabled": _config_says_enabled(ep.name),
|
||||
"hooks": [],
|
||||
"requests_conversation": None,
|
||||
"conversation_granted": manager._granted_conversation_access(ep.name),
|
||||
})
|
||||
except Exception as e:
|
||||
import logging
|
||||
@@ -436,14 +460,17 @@ def get_plugin_registry() -> list:
|
||||
if not name or name in seen:
|
||||
continue
|
||||
seen.add(name)
|
||||
single_hooks = list(meta.get("hooks", []) or [])
|
||||
entries.append({
|
||||
"name": name,
|
||||
"version": meta.get("version", "1.0.0"),
|
||||
"description": meta.get("description", ""),
|
||||
"source": "single_file",
|
||||
"enabled": manager.is_enabled(name) or _config_says_enabled(name),
|
||||
"hooks": list(meta.get("hooks", []) or []),
|
||||
"hooks": single_hooks,
|
||||
"path": meta.get("path"),
|
||||
"requests_conversation": _requests_conversation(single_hooks),
|
||||
"conversation_granted": manager._granted_conversation_access(name),
|
||||
})
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
@@ -22,6 +22,87 @@ if TYPE_CHECKING:
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# Lifecycle hooks that carry the user's prompt / conversation content and can
|
||||
# rewrite it (transform-style). Dispatching these to an arbitrary pip-installed
|
||||
# entry-point plugin lets a transitively-installed dependency read or tamper
|
||||
# with every user's messages, so they are gated behind an explicit grant.
|
||||
#
|
||||
# This deliberately covers BOTH directions of the conversation: the inbound
|
||||
# prompt/messages/model-response events AND the outbound message-delivery
|
||||
# events (``message_sending`` can rewrite the assistant reply sent to the
|
||||
# user; ``message_sent`` / ``message_undelivered`` expose that reply). Gating
|
||||
# only the inbound half would let an ungranted plugin still read or tamper with
|
||||
# every outbound reply.
|
||||
CONVERSATION_HOOKS = frozenset({
|
||||
PluginHook.MESSAGE_RECEIVED,
|
||||
PluginHook.BEFORE_LLM,
|
||||
PluginHook.AFTER_LLM,
|
||||
PluginHook.BEFORE_AGENT,
|
||||
PluginHook.AFTER_AGENT,
|
||||
PluginHook.MESSAGE_SENDING,
|
||||
PluginHook.MESSAGE_SENT,
|
||||
PluginHook.MESSAGE_UNDELIVERED,
|
||||
})
|
||||
|
||||
|
||||
def _package_root() -> Optional[Path]:
|
||||
"""Filesystem directory of the installed ``praisonaiagents`` package.
|
||||
|
||||
Used to verify first-party provenance by real file location rather than a
|
||||
self-asserted ``__module__`` string. Returns ``None`` if it cannot be
|
||||
resolved (in which case provenance falls back to failing closed).
|
||||
"""
|
||||
try:
|
||||
import praisonaiagents
|
||||
|
||||
file = getattr(praisonaiagents, "__file__", None)
|
||||
if not file:
|
||||
return None
|
||||
return Path(file).resolve().parent
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _is_first_party(plugin: Plugin) -> bool:
|
||||
"""True only for plugins genuinely bundled inside the framework.
|
||||
|
||||
First-party plugins are shipped and reviewed with the framework, so they
|
||||
are trusted with conversation content. A plugin's ``__module__`` string is
|
||||
self-asserted and therefore spoofable — a third-party pip package could set
|
||||
``__module__ = "praisonaiagents.evil"`` to falsely claim first-party trust.
|
||||
|
||||
To make the trust decision on provenance the plugin cannot forge, we
|
||||
require BOTH:
|
||||
|
||||
1. the class's ``__module__`` sits under the ``praisonaiagents`` namespace, and
|
||||
2. the module's real ``__file__`` resolves to a path *inside* the installed
|
||||
``praisonaiagents`` package directory.
|
||||
|
||||
Any plugin failing either check — notably an auto-discovered
|
||||
``praisonai.plugins`` entry point from an arbitrary distribution — is
|
||||
treated as third-party and denied conversation access by default.
|
||||
"""
|
||||
cls = type(plugin)
|
||||
module_name = getattr(cls, "__module__", "") or ""
|
||||
if module_name != "praisonaiagents" and not module_name.startswith("praisonaiagents."):
|
||||
return False
|
||||
|
||||
root = _package_root()
|
||||
if root is None:
|
||||
# Cannot verify provenance — fail closed rather than trust a string.
|
||||
return False
|
||||
|
||||
module = sys.modules.get(module_name)
|
||||
module_file = getattr(module, "__file__", None) if module is not None else None
|
||||
if not module_file:
|
||||
return False
|
||||
try:
|
||||
resolved = Path(module_file).resolve()
|
||||
except Exception:
|
||||
return False
|
||||
return root == resolved.parent or root in resolved.parents
|
||||
|
||||
|
||||
def _env_plugins_suppressed() -> bool:
|
||||
"""Return True when external plugins are suppressed via env for this run.
|
||||
|
||||
@@ -237,10 +318,13 @@ class PluginManager:
|
||||
Returns:
|
||||
True if unregistered successfully
|
||||
"""
|
||||
if name not in self._plugins:
|
||||
return False
|
||||
|
||||
plugin = self._plugins[name]
|
||||
with self._lock:
|
||||
plugin = self._plugins.get(name)
|
||||
if plugin is None:
|
||||
return False
|
||||
del self._plugins[name]
|
||||
self._enabled.pop(name, None)
|
||||
|
||||
try:
|
||||
plugin.on_shutdown()
|
||||
except Exception as e:
|
||||
@@ -249,18 +333,16 @@ class PluginManager:
|
||||
# Counterpart to the ON_INIT emit in register(): the plugin is going
|
||||
# away, and a listener that saw it arrive must be able to see it leave.
|
||||
_emit_plugin_shutdown(plugin)
|
||||
|
||||
del self._plugins[name]
|
||||
del self._enabled[name]
|
||||
|
||||
|
||||
logger.info(f"Unregistered plugin: {name}")
|
||||
return True
|
||||
|
||||
def enable(self, name: str) -> bool:
|
||||
"""Enable a plugin."""
|
||||
if name in self._enabled:
|
||||
self._enabled[name] = True
|
||||
return True
|
||||
with self._lock:
|
||||
if name in self._enabled:
|
||||
self._enabled[name] = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def disable(self, name: str) -> bool:
|
||||
@@ -270,18 +352,40 @@ class PluginManager:
|
||||
hook registry, so disabling actually stops its lifecycle methods from
|
||||
firing during subsequent agent execution.
|
||||
"""
|
||||
if name in self._enabled:
|
||||
with self._lock:
|
||||
if name not in self._enabled:
|
||||
return False
|
||||
self._enabled[name] = False
|
||||
try:
|
||||
self.unwire_from_hook_registry(name)
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to unwire plugin '{name}': {e}")
|
||||
return True
|
||||
return False
|
||||
try:
|
||||
self.unwire_from_hook_registry(name)
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to unwire plugin '{name}': {e}")
|
||||
return True
|
||||
|
||||
def is_enabled(self, name: str) -> bool:
|
||||
"""Check if a plugin is enabled."""
|
||||
return self._enabled.get(name, False)
|
||||
|
||||
def _granted_conversation_access(self, name: str) -> bool:
|
||||
"""Whether a plugin may receive conversation-content hooks.
|
||||
|
||||
Least-privilege, safe-by-default: a plugin sees ``before_llm`` /
|
||||
``after_llm`` / ``before_agent`` / ``after_agent`` / ``message_received``
|
||||
(which carry — and can rewrite — the user's prompt/messages) only when
|
||||
|
||||
- it is **first-party** (bundled in the framework), or
|
||||
- the operator **granted** it in the per-plugin config options map via
|
||||
``allow_conversation: true`` (e.g. ``plugins.<name>.allow_conversation``).
|
||||
|
||||
Any other plugin — notably an auto-discovered ``praisonai.plugins``
|
||||
entry point from an arbitrary pip package — is denied these hooks by
|
||||
default while still receiving its tool/channel/utility hooks.
|
||||
"""
|
||||
plugin = self._plugins.get(name)
|
||||
if plugin is not None and _is_first_party(plugin):
|
||||
return True
|
||||
options = self._plugin_options.get(name) or {}
|
||||
return options.get("allow_conversation") is True
|
||||
|
||||
def get_plugin(self, name: str) -> Optional[Plugin]:
|
||||
"""Get a plugin by name."""
|
||||
@@ -410,13 +514,27 @@ class PluginManager:
|
||||
The result (may be modified by plugins)
|
||||
"""
|
||||
result = args[0] if args else None
|
||||
|
||||
for name, plugin in self._plugins.items():
|
||||
if not self._enabled.get(name, False):
|
||||
continue
|
||||
|
||||
|
||||
# Snapshot the enabled (name, plugin) pairs under the lock so a
|
||||
# concurrent register/unregister cannot mutate the dict mid-iteration
|
||||
# (which would raise RuntimeError and crash the firing hook point for
|
||||
# every agent sharing this manager).
|
||||
with self._lock:
|
||||
enabled_plugins = [
|
||||
(name, plugin)
|
||||
for name, plugin in self._plugins.items()
|
||||
if self._enabled.get(name, False)
|
||||
]
|
||||
|
||||
for name, plugin in enabled_plugins:
|
||||
if hook not in plugin.info.hooks:
|
||||
continue
|
||||
|
||||
# Least-privilege gate: conversation-content hooks are default-deny
|
||||
# for third-party/ungranted plugins so a pip-installed dependency
|
||||
# cannot silently read or rewrite every user's prompt/messages.
|
||||
if hook in CONVERSATION_HOOKS and not self._granted_conversation_access(name):
|
||||
continue
|
||||
|
||||
try:
|
||||
method = getattr(plugin, hook.value, None)
|
||||
@@ -499,6 +617,74 @@ 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."""
|
||||
@@ -704,8 +890,16 @@ class PluginManager:
|
||||
if getattr(plugin, "_wired_into_registry", None) is registry:
|
||||
# Avoid double-registration on repeated enable() calls
|
||||
continue
|
||||
allow_conversation = self._granted_conversation_access(name)
|
||||
hook_ids: List[str] = []
|
||||
for event, func in _adapt_plugin_hooks(plugin):
|
||||
# Same least-privilege gate as execute_hook(): never bridge
|
||||
# a conversation-content hook of an ungranted plugin into the
|
||||
# runtime the Agent actually consults, so a third-party
|
||||
# entry-point plugin cannot read/rewrite prompts even though
|
||||
# its tool/channel hooks still wire in.
|
||||
if event in CONVERSATION_HOOKS and not allow_conversation:
|
||||
continue
|
||||
hook_id = registry.register_function(
|
||||
event=event,
|
||||
func=func,
|
||||
@@ -1112,6 +1306,51 @@ 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,6 +93,7 @@ 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 {
|
||||
@@ -102,6 +103,7 @@ 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):
|
||||
@@ -246,6 +248,42 @@ 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."""
|
||||
@@ -269,6 +307,40 @@ 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,6 +49,7 @@ __all__ = [
|
||||
"create_deny_tools_policy",
|
||||
"create_allow_tools_policy",
|
||||
"create_read_only_policy",
|
||||
"create_network_policy",
|
||||
]
|
||||
|
||||
|
||||
@@ -64,6 +65,7 @@ _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,6 +13,18 @@ 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.
|
||||
@@ -185,6 +197,34 @@ 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()
|
||||
@@ -307,6 +347,67 @@ 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.
|
||||
|
||||
@@ -757,10 +757,9 @@ Context tasks: {[getattr(t, 'name', str(t)) for t in current_task.context] if cu
|
||||
Description length: {len(current_task.description)}
|
||||
""")
|
||||
|
||||
# Build context and store separately instead of mutating description
|
||||
context = self._build_task_context(current_task)
|
||||
# Store context in dedicated field instead of concatenating to description
|
||||
current_task._execution_context = context if context else ""
|
||||
# Pending async predecessors may finish only after this task is
|
||||
# yielded. Resolve their outputs when the task actually executes.
|
||||
current_task._execution_context_builder = self._build_task_context
|
||||
|
||||
# Skip execution for loop tasks, only process their subtasks
|
||||
if current_task.task_type == "loop":
|
||||
@@ -845,6 +844,13 @@ 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)
|
||||
@@ -1446,6 +1452,7 @@ Description length: {len(current_task.description)}
|
||||
""")
|
||||
|
||||
# Build context and store separately instead of mutating description
|
||||
current_task._execution_context_builder = None
|
||||
context = self._build_task_context(current_task)
|
||||
# Store context in dedicated field instead of concatenating to description
|
||||
current_task._execution_context = context if context else ""
|
||||
|
||||
@@ -7,6 +7,14 @@ 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 = {}
|
||||
@@ -33,6 +41,12 @@ def __getattr__(name: str):
|
||||
__all__ = [
|
||||
"ChannelMessage",
|
||||
"PushTransportProtocol",
|
||||
"NotificationCategory",
|
||||
"DetailLevel",
|
||||
"NotificationPreference",
|
||||
"NotificationDecision",
|
||||
"evaluate",
|
||||
"DEFAULT_PREFERENCE",
|
||||
"PushClient",
|
||||
"WebSocketTransport",
|
||||
"PollingTransport",
|
||||
|
||||
@@ -0,0 +1,212 @@
|
||||
"""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,10 +30,13 @@ from .protocols import (
|
||||
RoundRobinRoomPlanner,
|
||||
extract_mentions,
|
||||
)
|
||||
from .runtime import Room, RoomConfig
|
||||
|
||||
__all__ = [
|
||||
"RoomEvent",
|
||||
"RoomTurnPlannerProtocol",
|
||||
"RoundRobinRoomPlanner",
|
||||
"extract_mentions",
|
||||
"RoomConfig",
|
||||
"Room",
|
||||
]
|
||||
@@ -0,0 +1,244 @@
|
||||
"""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
|
||||
@@ -291,11 +291,18 @@ def resolve_runtime(
|
||||
try:
|
||||
runtime = resolver.resolve(agent_id, model_ref, session_ctx, **kwargs)
|
||||
|
||||
# Cache the resolved runtime
|
||||
# Cache the resolved runtime and lazily start the background cleanup
|
||||
# thread — both under the same lock so the check-and-start in
|
||||
# _start_cleanup_thread() is atomic and two concurrent first-writers
|
||||
# can't each spawn a redundant daemon thread.
|
||||
with _runtime_cache_lock:
|
||||
if session_ctx.session_id not in _runtime_cache:
|
||||
_runtime_cache[session_ctx.session_id] = {}
|
||||
_runtime_cache[session_ctx.session_id][cache_key] = (runtime, current_time)
|
||||
|
||||
# Start the background cleanup thread lazily, only once something is
|
||||
# actually cached — not merely because the module was imported.
|
||||
_start_cleanup_thread()
|
||||
|
||||
logger.info(f"Successfully resolved runtime for {cache_key}")
|
||||
return runtime
|
||||
@@ -374,5 +381,7 @@ def _start_cleanup_thread():
|
||||
_cleanup_thread.start()
|
||||
logger.debug("Started runtime cache cleanup thread")
|
||||
|
||||
# Initialize cleanup thread when module is imported
|
||||
_start_cleanup_thread()
|
||||
# NOTE: the cleanup thread is started lazily on the first cache write
|
||||
# (see resolve_runtime), NOT at import time, so a bare
|
||||
# `import praisonaiagents` / touching `runtime.SessionContext` never spawns a
|
||||
# thread the caller didn't ask for (important for fork-based servers).
|
||||
@@ -25,6 +25,7 @@ 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 (
|
||||
@@ -106,6 +107,11 @@ 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
|
||||
@@ -176,6 +182,7 @@ __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, resolve_schedule_timezone
|
||||
from .due import is_due as _is_due, is_misfire as _is_misfire, resolve_schedule_timezone
|
||||
from .hook_emit import emit_schedule_add, emit_schedule_remove
|
||||
|
||||
logger = get_logger(__name__)
|
||||
@@ -202,6 +202,10 @@ 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()
|
||||
@@ -215,16 +219,40 @@ class ConfigYamlScheduleStore:
|
||||
if lease_until > now and lease_owner != owner_id:
|
||||
continue
|
||||
if not _is_due(job, now, self._default_timezone):
|
||||
# A job past its ``until`` end instant is not due, but must
|
||||
# still be retired (not left lingering in listings) and
|
||||
# never fired. It is removed here without being claimed so
|
||||
# there is no final stale send.
|
||||
if job.should_retire(now):
|
||||
auto_removed.append(self._jobs.pop(job.id))
|
||||
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
|
||||
# spent job before a competitor reclaims it) and flag it so the
|
||||
# later ``mark_run`` does not double-count the same fire.
|
||||
job.run_count += 1
|
||||
job._run_counted = True
|
||||
job._lease_until = now + lease_seconds
|
||||
job._lease_owner = owner_id
|
||||
self._held_leases[job.id] = owner_id
|
||||
claimed.append(job)
|
||||
changed = True
|
||||
if job.delete_after_run:
|
||||
# One-shot: remove now so no competitor re-claims it.
|
||||
if job.should_retire(now):
|
||||
# Spent (one-shot, run-count met, or past ``until``): remove
|
||||
# now so no competitor re-claims it.
|
||||
auto_removed.append(self._jobs.pop(job.id))
|
||||
self._job_state.pop(job.id, None)
|
||||
if changed:
|
||||
@@ -240,6 +268,14 @@ 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,6 +76,158 @@ 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,
|
||||
@@ -87,9 +239,48 @@ def is_due(
|
||||
ISO timestamp) and ``cron`` (5-field expression, requires optional
|
||||
``croniter``). Disabled jobs are the caller's concern and are not filtered
|
||||
here.
|
||||
|
||||
A bounded recurring job carries an optional stop condition — a ``max_runs``
|
||||
budget and/or an ``until`` end instant. Once either is met the job is no
|
||||
longer due, so it never fires or delivers again before the claim path
|
||||
retires it (see :meth:`ScheduleJob.should_retire`).
|
||||
"""
|
||||
sched = job.schedule
|
||||
|
||||
# Bounded-run stop conditions short-circuit before any kind check so a
|
||||
# spent job is neither due nor delivered (no final stale send).
|
||||
max_runs = getattr(job, "max_runs", None)
|
||||
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:
|
||||
until_ts = localize_wall_clock(
|
||||
datetime.fromisoformat(until),
|
||||
sched.tz or default_timezone,
|
||||
).timestamp()
|
||||
except (ValueError, TypeError):
|
||||
# Fail safe: a stop bound that cannot be parsed must NOT let the
|
||||
# job run past it forever. Treat the job as not-due; the claim path
|
||||
# then retires it via ``ScheduleJob.should_retire`` (which returns
|
||||
# True for an unparseable ``until``) rather than firing it.
|
||||
logger.warning("Invalid 'until' timestamp for job %s: %s", job.id, until)
|
||||
return False
|
||||
if now >= until_ts:
|
||||
return False
|
||||
|
||||
if sched.kind == "every":
|
||||
if sched.every_seconds is None:
|
||||
return False
|
||||
|
||||
@@ -10,6 +10,23 @@ 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.
|
||||
@@ -198,12 +215,16 @@ 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"``, or ``"no_change"``. ``"no_change"`` is the
|
||||
stateful "monitor mode" outcome — a watched source was
|
||||
``"skipped"``, ``"no_change"``, or ``"missed"``. ``"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.
|
||||
@@ -213,7 +234,7 @@ class RunRecord:
|
||||
|
||||
job_id: str
|
||||
job_name: str = ""
|
||||
status: Literal["succeeded", "failed", "skipped", "no_change"] = "succeeded"
|
||||
status: Literal["succeeded", "failed", "skipped", "no_change", "missed"] = "succeeded"
|
||||
result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
duration: float = 0.0
|
||||
@@ -260,6 +281,15 @@ class ScheduleJob:
|
||||
``"isolated"`` creates a fresh agent turn.
|
||||
enabled: Toggle without deleting.
|
||||
delete_after_run: Auto-remove after first execution (one-shot).
|
||||
Sugar for ``max_runs=1``; either retires the job.
|
||||
max_runs: Optional run-count bound for a recurring schedule — fire at
|
||||
most this many times, then auto-remove (generalises
|
||||
``delete_after_run``). ``None`` (default) means unbounded.
|
||||
run_count: Number of times the job has fired; incremented on each
|
||||
claim. Compared against ``max_runs`` to retire a spent job.
|
||||
until: Optional ISO 8601 wall-clock instant after which a recurring
|
||||
schedule stops firing and is retired (timezone-resolved exactly
|
||||
like ``at`` via ``localize_wall_clock``). ``None`` ⇒ no end date.
|
||||
created_at: Epoch timestamp.
|
||||
last_run_at: Epoch timestamp of most recent execution.
|
||||
delivery: Optional delivery target for routing results back
|
||||
@@ -365,6 +395,28 @@ 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 = ""
|
||||
@@ -375,6 +427,9 @@ class ScheduleJob:
|
||||
session_target: Literal["main", "isolated"] = "isolated"
|
||||
enabled: bool = True
|
||||
delete_after_run: bool = False
|
||||
max_runs: Optional[int] = None
|
||||
run_count: int = 0
|
||||
until: Optional[str] = None
|
||||
created_at: float = field(default_factory=time.time)
|
||||
last_run_at: Optional[float] = None
|
||||
delivery: Optional[DeliveryTarget] = None
|
||||
@@ -391,8 +446,59 @@ 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 ───────────────────────────────────────
|
||||
|
||||
def should_retire(self, now: Optional[float] = None) -> bool:
|
||||
"""Return ``True`` when the job has met a stop condition and is spent.
|
||||
|
||||
The single source of truth for "auto-remove this job now?", shared by
|
||||
every claim path so the retire decision is defined once. A job retires
|
||||
when it is:
|
||||
|
||||
* a one-shot (``delete_after_run``),
|
||||
* at or past its ``max_runs`` budget, or
|
||||
* past its ``until`` end instant (when ``now`` is supplied).
|
||||
|
||||
The ``until`` bound is *also* short-circuited earlier in
|
||||
:func:`due.is_due` (a spent-by-date job stops being due immediately),
|
||||
but it is repeated here so a job that has already elapsed is actively
|
||||
retired by the claim path rather than lingering forever in listings.
|
||||
|
||||
Assumes ``run_count`` has already been incremented for the current
|
||||
fire, so ``max_runs=1`` retires after the first run (matching the
|
||||
one-shot semantics ``delete_after_run`` provides).
|
||||
|
||||
Args:
|
||||
now: Epoch to evaluate ``until`` against. When ``None`` the
|
||||
``until`` check is skipped (callers that only ran a job and
|
||||
want the run-count/one-shot decision).
|
||||
"""
|
||||
if self.delete_after_run:
|
||||
return True
|
||||
if self.max_runs is not None and self.run_count >= self.max_runs:
|
||||
return True
|
||||
if now is not None and self.until is not None:
|
||||
# Fail safe: a stop bound that cannot be evaluated must retire the
|
||||
# job (never run forever). Import lazily to keep models.py free of
|
||||
# the timezone helpers at import time.
|
||||
from datetime import datetime
|
||||
from .due import localize_wall_clock
|
||||
|
||||
try:
|
||||
until_ts = localize_wall_clock(
|
||||
datetime.fromisoformat(self.until),
|
||||
self.schedule.tz,
|
||||
).timestamp()
|
||||
except (ValueError, TypeError):
|
||||
return True
|
||||
if now >= until_ts:
|
||||
return True
|
||||
return False
|
||||
|
||||
# ── serialisation ────────────────────────────────────────────────
|
||||
|
||||
@@ -409,6 +515,15 @@ class ScheduleJob:
|
||||
"created_at": self.created_at,
|
||||
"last_run_at": self.last_run_at,
|
||||
}
|
||||
# Bounded-run stop conditions. Only persist when configured / non-zero
|
||||
# so unbounded jobs stay byte-for-byte unchanged and fully backward
|
||||
# compatible; ``run_count`` is persisted only once it starts counting.
|
||||
if self.max_runs is not None:
|
||||
d["max_runs"] = self.max_runs
|
||||
if self.run_count:
|
||||
d["run_count"] = self.run_count
|
||||
if self.until is not None:
|
||||
d["until"] = self.until
|
||||
if self.delivery is not None:
|
||||
d["delivery"] = self.delivery.to_dict()
|
||||
if self.origin is not None:
|
||||
@@ -449,6 +564,10 @@ 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).
|
||||
@@ -456,6 +575,11 @@ 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.
|
||||
@@ -488,6 +612,9 @@ class ScheduleJob:
|
||||
session_target=d.get("session_target", "isolated"),
|
||||
enabled=d.get("enabled", True),
|
||||
delete_after_run=d.get("delete_after_run", False),
|
||||
max_runs=d.get("max_runs"),
|
||||
run_count=d.get("run_count", 0),
|
||||
until=d.get("until"),
|
||||
created_at=d.get("created_at", time.time()),
|
||||
last_run_at=d.get("last_run_at"),
|
||||
delivery=DeliveryTarget.from_dict(delivery_data) if isinstance(delivery_data, dict) else None,
|
||||
@@ -504,12 +631,16 @@ 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
|
||||
|
||||
@@ -96,6 +96,16 @@ class ScheduleRunner:
|
||||
delivered: Whether result was delivered to a channel bot.
|
||||
"""
|
||||
job.last_run_at = time.time()
|
||||
# Count this fire exactly once. An atomic ``claim_due`` already
|
||||
# incremented ``run_count`` when it reserved the job (and needs to, so
|
||||
# ``should_retire()`` can drop a spent job before a competitor reclaims
|
||||
# it); it flags the job with ``_run_counted`` so we do not double-count
|
||||
# here. The non-atomic ``get_due_jobs`` path never sets that flag, so
|
||||
# ``mark_run`` remains the single incrementer there.
|
||||
if getattr(job, "_run_counted", False):
|
||||
job._run_counted = False
|
||||
else:
|
||||
job.run_count += 1
|
||||
|
||||
# Log execution history if the store supports it
|
||||
if hasattr(self._store, "log_run"):
|
||||
@@ -109,7 +119,7 @@ class ScheduleRunner:
|
||||
job_name=job.name,
|
||||
)
|
||||
|
||||
if job.delete_after_run:
|
||||
if job.should_retire(job.last_run_at):
|
||||
self._store.remove(job.id)
|
||||
else:
|
||||
self._store.update(job)
|
||||
|
||||
@@ -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
|
||||
from .due import is_due as _is_due, is_misfire as _is_misfire
|
||||
from .hook_emit import emit_schedule_add, emit_schedule_remove
|
||||
|
||||
logger = get_logger(__name__)
|
||||
@@ -187,6 +187,10 @@ 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()
|
||||
@@ -200,16 +204,41 @@ class FileScheduleStore:
|
||||
if lease_until > now and lease_owner != owner_id:
|
||||
continue
|
||||
if not _is_due(job, now):
|
||||
# A job past its ``until`` end instant is not due, but must
|
||||
# still be retired (not left lingering in listings) and
|
||||
# never fired. It is removed here without being claimed so
|
||||
# there is no final stale send.
|
||||
if job.should_retire(now):
|
||||
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
|
||||
# spent job before a competitor reclaims it) and flag it so the
|
||||
# later ``mark_run`` does not double-count the same fire.
|
||||
job.run_count += 1
|
||||
job._run_counted = True
|
||||
job._lease_until = now + lease_seconds
|
||||
job._lease_owner = owner_id
|
||||
self._held_leases[job.id] = owner_id
|
||||
claimed.append(job)
|
||||
changed = True
|
||||
if job.delete_after_run:
|
||||
# One-shot: remove now so no competitor re-claims it.
|
||||
if job.should_retire(now):
|
||||
# Spent (one-shot, run-count met, or past ``until``): remove
|
||||
# now so no competitor re-claims it.
|
||||
auto_removed.append(self._jobs.pop(job.id))
|
||||
if changed:
|
||||
if not self._save():
|
||||
@@ -224,6 +253,15 @@ 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,6 +39,8 @@ __all__ = [
|
||||
"register_resolver",
|
||||
"register_secret_for_redaction",
|
||||
"redact_secrets",
|
||||
"redact_outbound",
|
||||
"OutboundRedactor",
|
||||
"is_secret_ref",
|
||||
]
|
||||
|
||||
@@ -287,3 +289,85 @@ 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: ...
|
||||
Loaded 100 of 397 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user