mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 08:11:07 -10:00
Compare commits
No files matched your search
@@ -93,7 +93,7 @@ jobs:
|
||||
# a CI-shaped environment rather than trusting a developer machine
|
||||
# where every sibling is already installed.
|
||||
pip install -e src/praisonai
|
||||
pip install pytest pytest-asyncio pytest-timeout
|
||||
pip install pytest pytest-asyncio pytest-timeout aiofiles
|
||||
- name: Run the unit suite
|
||||
working-directory: src/praisonai-agents
|
||||
env:
|
||||
|
||||
@@ -117,13 +117,25 @@ jobs:
|
||||
const labels = issue.labels.map((l) => l.name);
|
||||
|
||||
if (status !== 'DIRTY' && labels.includes(LABEL)) {
|
||||
await github.rest.issues.removeLabel({
|
||||
owner,
|
||||
repo,
|
||||
issue_number: prNumber,
|
||||
name: LABEL,
|
||||
});
|
||||
core.info(`Removed ${LABEL} from PR #${prNumber}`);
|
||||
try {
|
||||
await github.rest.issues.removeLabel({
|
||||
owner,
|
||||
repo,
|
||||
issue_number: prNumber,
|
||||
name: LABEL,
|
||||
});
|
||||
core.info(`Removed ${LABEL} from PR #${prNumber}`);
|
||||
} catch (err) {
|
||||
// Fork PRs may return 403/404 when the token cannot mutate labels on the head repo's issue.
|
||||
// For upstream PRs a 403/404 is a genuine failure that must surface, otherwise the
|
||||
// merge-blocking label stays in place while the check reports success.
|
||||
const isFork = headRepo && headRepo !== baseRepo;
|
||||
if (isFork && (err.status === 403 || err.status === 404)) {
|
||||
core.info(`Skip remove ${LABEL} on fork PR #${prNumber}: ${err.status}`);
|
||||
} else {
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (status !== 'DIRTY') return { skipped: true, reason: 'not dirty' };
|
||||
|
||||
@@ -42,7 +42,7 @@ jobs:
|
||||
# ============================================
|
||||
smoke:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 12
|
||||
timeout-minutes: 16
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
|
||||
@@ -1001,6 +1001,8 @@ export type { DbAdapter, DbConfig, DbMessage, DbRun, DbTrace } from "./db";
|
||||
export { MemoryPostgresAdapter, NeonPostgresAdapter, PostgresSessionStorage, createMemoryPostgres, createNeonPostgres, createPostgresSessionStorage } from "./db/postgres";
|
||||
export { MemoryRedisAdapter, UpstashRedisAdapter, createMemoryRedis, createUpstashRedis } from "./db/redis";
|
||||
export { SQLiteAdapter, createSQLiteAdapter } from "./db/sqlite";
|
||||
export { KNOWN_DECISION_MODELS, choiceQuestion, defaultTicketTriageQuestions, getChoice, getNoul, getScore, isDecisionModel, noulQuestion, resolveSystemOneBaseUrl, scoreQuestion, systemOne, triagedStart } from "./decisions";
|
||||
export type { DecisionAgentLike, DecisionRoutePlan, QuestionSpec, RouteTarget, SystemOneAnswer, SystemOneOptions, SystemOneResult } from "./decisions";
|
||||
export { // Functions
|
||||
// NB: the unaliased registerDisplayCallback name belongs to the per-type
|
||||
// callback registry (callbacks module, // Types
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.11" \
|
||||
"praisonai>=4.7.12" \
|
||||
"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.11" \
|
||||
"praisonai>=4.7.12" \
|
||||
"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.11" \
|
||||
"praisonai>=4.7.12" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -18,22 +18,22 @@ This report compares **Python SDK feature categories** against **Python document
|
||||
| Category | Features | Docs | Lines |
|
||||
|----------|----------|------|-------|
|
||||
| ✅ AGUI | 1 | 2 | 511 |
|
||||
| ✅ Agent | 22 | 59 | 19076 |
|
||||
| ✅ Agent | 22 | 59 | 19085 |
|
||||
| ✅ Agent-to-Agent (A2A) | 1 | 7 | 2996 |
|
||||
| ✅ Agent-to-User (A2U) | 1 | 3 | 701 |
|
||||
| ✅ Approval | 1 | 8 | 3969 |
|
||||
| ✅ Approval | 1 | 8 | 4028 |
|
||||
| ✅ Audio | 2 | 12 | 1000 |
|
||||
| ✅ Auto Generation | 5 | 12 | 3941 |
|
||||
| ✅ Autonomy | 3 | 7 | 2805 |
|
||||
| ✅ Bots | 7 | 38 | 15809 |
|
||||
| ✅ Autonomy | 3 | 7 | 2808 |
|
||||
| ✅ Bots | 7 | 39 | 16300 |
|
||||
| ✅ Budget | 1 | 1 | 287 |
|
||||
| ✅ CLI | 5 | 124 | 39191 |
|
||||
| ✅ CLI | 5 | 124 | 39199 |
|
||||
| ✅ Chunking | 2 | 2 | 435 |
|
||||
| ✅ Citations | 2 | 1 | 202 |
|
||||
| ✅ Code Execution | 2 | 14 | 5022 |
|
||||
| ✅ Conditions | 1 | 3 | 1386 |
|
||||
| ✅ Configuration | 3 | 8 | 3341 |
|
||||
| ✅ Context Management | 16 | 36 | 14384 |
|
||||
| ✅ Configuration | 3 | 8 | 3350 |
|
||||
| ✅ Context Management | 16 | 37 | 14687 |
|
||||
| ✅ Deep Research | 8 | 2 | 591 |
|
||||
| ✅ Display | 6 | 3 | 844 |
|
||||
| ✅ Embeddings | 6 | 23 | 2051 |
|
||||
@@ -43,7 +43,7 @@ This report compares **Python SDK feature categories** against **Python document
|
||||
| ✅ Failover | 2 | 1 | 537 |
|
||||
| ✅ Files | 2 | 7 | 2376 |
|
||||
| ✅ Flow | 1 | 3 | 811 |
|
||||
| ✅ Gateway | 7 | 97 | 37597 |
|
||||
| ✅ Gateway | 7 | 99 | 38257 |
|
||||
| ✅ Guardrails | 5 | 4 | 2642 |
|
||||
| ✅ Handoffs | 11 | 6 | 2748 |
|
||||
| ✅ Hooks | 2 | 9 | 5230 |
|
||||
@@ -59,7 +59,7 @@ This report compares **Python SDK feature categories** against **Python document
|
||||
| ✅ Output | 3 | 5 | 1289 |
|
||||
| ✅ Parallel Execution | 3 | 2 | 528 |
|
||||
| ✅ Planning | 6 | 7 | 1827 |
|
||||
| ✅ Plugins | 8 | 5 | 4108 |
|
||||
| ✅ Plugins | 8 | 5 | 4112 |
|
||||
| ✅ Prompts | 2 | 9 | 2043 |
|
||||
| ✅ Providers | 1 | 58 | 8751 |
|
||||
| ✅ Query | 1 | 2 | 736 |
|
||||
@@ -70,24 +70,24 @@ This report compares **Python SDK feature categories** against **Python document
|
||||
| ✅ Routing | 1 | 2 | 430 |
|
||||
| ✅ Sandbox | 5 | 7 | 3720 |
|
||||
| ✅ Security | 1 | 3 | 3642 |
|
||||
| ✅ Sessions | 4 | 21 | 8388 |
|
||||
| ✅ Sessions | 4 | 21 | 8389 |
|
||||
| ✅ Skills | 6 | 16 | 6045 |
|
||||
| ✅ Tasks | 2 | 8 | 3266 |
|
||||
| ✅ Telemetry | 1 | 2 | 660 |
|
||||
| ✅ Templates | 1 | 9 | 1859 |
|
||||
| ✅ Tools | 12 | 147 | 45378 |
|
||||
| ✅ Tools | 12 | 147 | 45414 |
|
||||
| ✅ Tracing | 3 | 2 | 139 |
|
||||
| ✅ Vector Store | 1 | 12 | 1219 |
|
||||
| ✅ Video | 2 | 10 | 1303 |
|
||||
| ✅ Vision | 2 | 1 | 329 |
|
||||
| ✅ Web | 3 | 10 | 2822 |
|
||||
| ✅ Workflows | 5 | 21 | 8880 |
|
||||
| ✅ Workflows | 5 | 21 | 8933 |
|
||||
|
||||
## Documentation Without Features
|
||||
|
||||
These docs exist but don't match any implemented feature category:
|
||||
|
||||
- ℹ️ Database (44 docs, 8206 lines)
|
||||
- ℹ️ Database (44 docs, 8359 lines)
|
||||
- ℹ️ Documents (1 docs, 810 lines)
|
||||
- ℹ️ Teams (1 docs, 214 lines)
|
||||
|
||||
|
||||
@@ -6,8 +6,13 @@ Compares agent instantiation times across popular AI agent frameworks.
|
||||
|
||||
Usage:
|
||||
python benchmarks/simple_benchmark.py
|
||||
python benchmarks/simple_benchmark.py --skip-crewai --json
|
||||
python benchmarks/simple_benchmark.py --json --fail-over-threshold
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from typing import Literal
|
||||
import argparse
|
||||
@@ -16,6 +21,14 @@ import argparse
|
||||
ITERATIONS = 100
|
||||
WARMUP = 10
|
||||
|
||||
THRESHOLDS_PATH = os.path.join(os.path.dirname(__file__), 'thresholds.json')
|
||||
|
||||
|
||||
def load_thresholds(path: str = THRESHOLDS_PATH) -> dict:
|
||||
"""Load benchmark thresholds from thresholds.json."""
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def sample_tool(city: Literal['nyc', 'sf']):
|
||||
"""Sample tool for benchmark testing."""
|
||||
@@ -36,35 +49,44 @@ def measure_instantiation(create_fn, iterations=ITERATIONS, warmup=WARMUP):
|
||||
return sum(times) / len(times)
|
||||
|
||||
|
||||
def run_benchmark():
|
||||
def run_benchmark(skip_crewai: bool = False, quiet: bool = False,
|
||||
iterations: int = ITERATIONS, warmup: int = WARMUP):
|
||||
"""Run the benchmark across all available frameworks (without tools)."""
|
||||
results = {}
|
||||
|
||||
print('=' * 60)
|
||||
print('PraisonAI Agents - Performance Benchmark')
|
||||
print('=' * 60)
|
||||
print(f'\nIterations: {ITERATIONS}')
|
||||
print('Metric: Agent instantiation time (microseconds)\n')
|
||||
|
||||
def log(msg: str = ''):
|
||||
if not quiet:
|
||||
print(msg)
|
||||
|
||||
def measure(create_fn):
|
||||
return measure_instantiation(create_fn, iterations=iterations, warmup=warmup)
|
||||
|
||||
log('=' * 60)
|
||||
log('PraisonAI Agents - Performance Benchmark')
|
||||
log('=' * 60)
|
||||
log(f'\nIterations: {iterations}')
|
||||
log(f'Warmup: {warmup}')
|
||||
log('Metric: Agent instantiation time (microseconds)\n')
|
||||
|
||||
# PraisonAI (without tools)
|
||||
print("Testing PraisonAI...")
|
||||
log("Testing PraisonAI...")
|
||||
from praisonaiagents import Agent as PraisonAgent
|
||||
|
||||
results['PraisonAI'] = measure_instantiation(
|
||||
results['PraisonAI'] = measure(
|
||||
lambda: PraisonAgent(name='Test', model='gpt-4o-mini', output="silent")
|
||||
)
|
||||
|
||||
results['PraisonAI (LiteLLM)'] = measure_instantiation(
|
||||
results['PraisonAI (LiteLLM)'] = measure(
|
||||
lambda: PraisonAgent(name='Test', model='openai/gpt-4o-mini', output="silent")
|
||||
)
|
||||
|
||||
# Other frameworks for comparison
|
||||
print("Testing other frameworks...")
|
||||
log("Testing other frameworks...")
|
||||
|
||||
try:
|
||||
from agno.agent import Agent as AgnoAgent
|
||||
from agno.models.openai import OpenAIChat
|
||||
results['Agno'] = measure_instantiation(
|
||||
results['Agno'] = measure(
|
||||
lambda: AgnoAgent(model=OpenAIChat(id='gpt-4o-mini'))
|
||||
)
|
||||
except ImportError:
|
||||
@@ -72,7 +94,7 @@ def run_benchmark():
|
||||
|
||||
try:
|
||||
from pydantic_ai import Agent as PydanticAgent
|
||||
results['PydanticAI'] = measure_instantiation(
|
||||
results['PydanticAI'] = measure(
|
||||
lambda: PydanticAgent('openai:gpt-4o-mini')
|
||||
)
|
||||
except ImportError:
|
||||
@@ -80,7 +102,7 @@ def run_benchmark():
|
||||
|
||||
try:
|
||||
from agents import Agent as OpenAIAgent
|
||||
results['OpenAI Agents SDK'] = measure_instantiation(
|
||||
results['OpenAI Agents SDK'] = measure(
|
||||
lambda: OpenAIAgent(name='Test', model='gpt-4o-mini')
|
||||
)
|
||||
except ImportError:
|
||||
@@ -96,47 +118,50 @@ def run_benchmark():
|
||||
"""Get weather info."""
|
||||
return 'sunny' if city == 'sf' else 'cloudy'
|
||||
|
||||
results['LangGraph'] = measure_instantiation(
|
||||
results['LangGraph'] = measure(
|
||||
lambda: create_react_agent(model=ChatOpenAI(model='gpt-4o-mini'), tools=[get_weather_lg])
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
from crewai.agent import Agent as CrewAgent
|
||||
from crewai.tools import tool as crewai_tool
|
||||
|
||||
@crewai_tool("Weather Tool")
|
||||
def get_weather_crew(city: Literal['nyc', 'sf']):
|
||||
"""Get weather info."""
|
||||
return 'sunny' if city == 'sf' else 'cloudy'
|
||||
|
||||
results['CrewAI'] = measure_instantiation(
|
||||
lambda: CrewAgent(
|
||||
role='Weather Agent',
|
||||
goal='Provide weather info',
|
||||
backstory='A weather expert',
|
||||
tools=[get_weather_crew],
|
||||
verbose=False
|
||||
if skip_crewai:
|
||||
log("Skipping CrewAI (--skip-crewai): optional framework import can hang on some platforms.")
|
||||
else:
|
||||
try:
|
||||
from crewai.agent import Agent as CrewAgent
|
||||
from crewai.tools import tool as crewai_tool
|
||||
|
||||
@crewai_tool("Weather Tool")
|
||||
def get_weather_crew(city: Literal['nyc', 'sf']):
|
||||
"""Get weather info."""
|
||||
return 'sunny' if city == 'sf' else 'cloudy'
|
||||
|
||||
results['CrewAI'] = measure(
|
||||
lambda: CrewAgent(
|
||||
role='Weather Agent',
|
||||
goal='Provide weather info',
|
||||
backstory='A weather expert',
|
||||
tools=[get_weather_crew],
|
||||
verbose=False
|
||||
)
|
||||
)
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Print results
|
||||
print('\n' + '=' * 60)
|
||||
print('RESULTS')
|
||||
print('=' * 60)
|
||||
log('\n' + '=' * 60)
|
||||
log('RESULTS')
|
||||
log('=' * 60)
|
||||
|
||||
baseline = results.get('PraisonAI', 1)
|
||||
print(f"\n{'Framework':<25} {'Avg Time (μs)':<15} {'Relative':<10}")
|
||||
print('-' * 50)
|
||||
log(f"\n{'Framework':<25} {'Avg Time (μs)':<15} {'Relative':<10}")
|
||||
log('-' * 50)
|
||||
|
||||
for name, avg in sorted(results.items(), key=lambda x: x[1]):
|
||||
ratio = avg / baseline
|
||||
print(f'{name:<25} {avg:<15.2f} {ratio:.2f}x')
|
||||
log(f'{name:<25} {avg:<15.2f} {ratio:.2f}x')
|
||||
|
||||
print('\n' + '=' * 60)
|
||||
log('\n' + '=' * 60)
|
||||
return results
|
||||
|
||||
|
||||
@@ -163,8 +188,13 @@ def get_package_versions():
|
||||
return versions
|
||||
|
||||
|
||||
def save_results(results: dict, filename: str = 'BENCHMARK_RESULTS.md'):
|
||||
"""Save benchmark results to a markdown file and update README."""
|
||||
def save_results(results: dict, filename: str = 'BENCHMARK_RESULTS.md', *,
|
||||
iterations: int = ITERATIONS, log=print):
|
||||
"""Save benchmark results to a markdown file and update README.
|
||||
|
||||
Status messages are emitted via ``log`` (defaults to ``print``). Pass a
|
||||
stderr logger when stdout must stay machine-readable (e.g. ``--json``).
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
from datetime import datetime
|
||||
@@ -191,7 +221,7 @@ def save_results(results: dict, filename: str = 'BENCHMARK_RESULTS.md'):
|
||||
with open(filepath, 'w') as f:
|
||||
f.write('# PraisonAI Agents - Benchmark Results\n\n')
|
||||
f.write(f'**Generated:** {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n')
|
||||
f.write(f'**Iterations:** {ITERATIONS}\n')
|
||||
f.write(f'**Iterations:** {iterations}\n')
|
||||
f.write('**Test:** Agent instantiation (without tools)\n\n')
|
||||
f.write('## Results\n\n')
|
||||
f.write('| Framework | Avg Time (μs) | Relative |\n')
|
||||
@@ -208,17 +238,17 @@ def save_results(results: dict, filename: str = 'BENCHMARK_RESULTS.md'):
|
||||
f.write('python benchmarks/simple_benchmark.py\n')
|
||||
f.write('```\n')
|
||||
|
||||
print(f'\nResults saved to: {filepath}')
|
||||
log(f'\nResults saved to: {filepath}')
|
||||
|
||||
# Also update the main README.md
|
||||
readme_path = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'README.md')
|
||||
if os.path.exists(readme_path):
|
||||
update_readme(readme_path, table_rows)
|
||||
update_readme(readme_path, table_rows, log=log)
|
||||
|
||||
return filepath
|
||||
|
||||
|
||||
def update_readme(readme_path: str, table_rows: list):
|
||||
def update_readme(readme_path: str, table_rows: list, *, log=print):
|
||||
"""Update the performance section in README.md with latest results."""
|
||||
import re
|
||||
|
||||
@@ -239,17 +269,54 @@ def update_readme(readme_path: str, table_rows: list):
|
||||
with open(readme_path, 'w') as f:
|
||||
f.write(content)
|
||||
|
||||
print(f'README.md updated: {readme_path}')
|
||||
log(f'README.md updated: {readme_path}')
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='PraisonAI Agents - Performance Benchmark')
|
||||
parser.add_argument('--save', action='store_true', help='Save results to file')
|
||||
parser.add_argument('--skip-crewai', action='store_true',
|
||||
help='Skip the optional CrewAI benchmark (its import can hang on some platforms)')
|
||||
parser.add_argument('--json', action='store_true',
|
||||
help='Emit results as JSON to stdout (for CI baselines/artifacts)')
|
||||
parser.add_argument('--fail-over-threshold', action='store_true',
|
||||
help='Exit non-zero if PraisonAI instantiation exceeds the budget in thresholds.json')
|
||||
args = parser.parse_args()
|
||||
|
||||
thresholds = load_thresholds()
|
||||
iterations = int(thresholds.get('iterations', ITERATIONS))
|
||||
warmup = int(thresholds.get('warmup', WARMUP))
|
||||
|
||||
results = run_benchmark(
|
||||
skip_crewai=args.skip_crewai,
|
||||
quiet=args.json,
|
||||
iterations=iterations,
|
||||
warmup=warmup,
|
||||
)
|
||||
|
||||
if args.save:
|
||||
# Keep stdout machine-readable under --json by sending save status to stderr.
|
||||
save_log = (lambda msg='': print(msg, file=sys.stderr)) if args.json else print
|
||||
save_results(results, iterations=iterations, log=save_log)
|
||||
elif not args.json:
|
||||
print('\nResults not saved (use --save flag to save results to file)')
|
||||
|
||||
exit_code = 0
|
||||
if args.fail_over_threshold:
|
||||
max_us = thresholds['praisonai_instantiation_us_max']
|
||||
actual_us = results.get('PraisonAI')
|
||||
if actual_us is not None and actual_us > max_us:
|
||||
print(
|
||||
f"ERROR: PraisonAI instantiation {actual_us:.2f} us > threshold {max_us} us",
|
||||
file=sys.stderr,
|
||||
)
|
||||
exit_code = 1
|
||||
|
||||
if args.json:
|
||||
print(json.dumps({'results': results, 'iterations': iterations, 'warmup': warmup}))
|
||||
|
||||
return exit_code
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='PraisonAI Agents - Performance Benchmark')
|
||||
parser.add_argument('--save', action='store_true', help='Save results to file')
|
||||
args = parser.parse_args()
|
||||
|
||||
results = run_benchmark()
|
||||
|
||||
if args.save:
|
||||
save_results(results)
|
||||
else:
|
||||
print('\nResults not saved (use --save flag to save results to file)')
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"praisonai_instantiation_us_max": 450,
|
||||
"iterations": 100,
|
||||
"warmup": 10
|
||||
}
|
||||
@@ -227,6 +227,20 @@ _LAZY_IMPORTS = {
|
||||
'aembed': ('praisonaiagents.embedding.embed', 'aembed'),
|
||||
'EmbeddingResult': ('praisonaiagents.embedding.result', 'EmbeddingResult'),
|
||||
'get_dimensions': ('praisonaiagents.embedding.dimensions', 'get_dimensions'),
|
||||
|
||||
# System One / Jev decision models (Ollama /v1/systemone)
|
||||
'system_one': ('praisonaiagents.decisions', 'system_one'),
|
||||
'asystem_one': ('praisonaiagents.decisions', 'asystem_one'),
|
||||
'SystemOneResult': ('praisonaiagents.decisions', 'SystemOneResult'),
|
||||
'choice_question': ('praisonaiagents.decisions', 'choice_question'),
|
||||
'noul_question': ('praisonaiagents.decisions', 'noul_question'),
|
||||
'score_question': ('praisonaiagents.decisions', 'score_question'),
|
||||
'is_decision_model': ('praisonaiagents.decisions', 'is_decision_model'),
|
||||
'DecisionRoutePlan': ('praisonaiagents.decisions', 'DecisionRoutePlan'),
|
||||
'DecisionTriageRouter': ('praisonaiagents.decisions', 'DecisionTriageRouter'),
|
||||
'triaged_start': ('praisonaiagents.decisions', 'triaged_start'),
|
||||
'triage_decision': ('praisonaiagents.decisions', 'triage_decision'),
|
||||
'default_ticket_triage_questions': ('praisonaiagents.decisions', 'default_ticket_triage_questions'),
|
||||
|
||||
# Guardrails
|
||||
'GuardrailResult': ('praisonaiagents.guardrails', 'GuardrailResult'),
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Shared lazy run-lock helper for core orchestration modules.
|
||||
|
||||
AgentTeam (``agents/agents.py``) and Workflow/AgentFlow (``workflows/workflows.py``)
|
||||
both lazily mint a per-instance re-entrancy lock the first time a run begins.
|
||||
Creation is serialized through a single module-level guard (double-checked
|
||||
locking) so two threads racing into the first ``start()``/``astart()``/``run()``
|
||||
observe the *same* lock object rather than each minting and acquiring its own.
|
||||
"""
|
||||
import threading
|
||||
|
||||
# Guards lazy creation of each instance's per-instance ``_run_lock`` so two
|
||||
# threads that first reach start()/astart()/run() concurrently observe the same
|
||||
# lock object (double-checked locking) rather than each minting its own.
|
||||
_RUN_LOCK_INIT_GUARD = threading.Lock()
|
||||
|
||||
|
||||
def ensure_run_lock(obj):
|
||||
"""Lazily mint ``obj._run_lock`` under a shared module guard.
|
||||
|
||||
``getattr(obj, "_run_lock", None)`` works for both the dataclass-field case
|
||||
(Workflow) and the plain-attribute case (AgentTeam), so the behaviour is
|
||||
identical to the per-module implementations it replaces.
|
||||
"""
|
||||
lock = getattr(obj, "_run_lock", None)
|
||||
if lock is None:
|
||||
with _RUN_LOCK_INIT_GUARD:
|
||||
lock = getattr(obj, "_run_lock", None)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
obj._run_lock = lock
|
||||
return lock
|
||||
@@ -547,6 +547,15 @@ class ToolExecutionMixin:
|
||||
so disabling or unregistering a plugin revokes its tools from this
|
||||
agent's discovery and execution paths.
|
||||
"""
|
||||
# Hot-path fast exit: when no plugins are registered (the common case)
|
||||
# skip constructing/scanning the manager entirely on every Agent init.
|
||||
try:
|
||||
from ..plugins.manager import has_registered_plugins
|
||||
|
||||
if not has_registered_plugins():
|
||||
return
|
||||
except Exception: # pragma: no cover - defensive plugin boundary
|
||||
pass
|
||||
try:
|
||||
from ..plugins import get_plugin_manager
|
||||
|
||||
@@ -673,6 +682,15 @@ class ToolExecutionMixin:
|
||||
prevent the Agent from constructing. Zero overhead when no such plugin
|
||||
is enabled (each collector returns an empty list).
|
||||
"""
|
||||
# Hot-path fast exit: when no plugins are registered (the common case)
|
||||
# skip constructing/scanning the manager entirely on every Agent init.
|
||||
try:
|
||||
from ..plugins.manager import has_registered_plugins
|
||||
|
||||
if not has_registered_plugins():
|
||||
return
|
||||
except Exception: # pragma: no cover - defensive plugin boundary
|
||||
pass
|
||||
try:
|
||||
from ..plugins import get_plugin_manager
|
||||
|
||||
|
||||
@@ -24,10 +24,7 @@ try:
|
||||
except ImportError:
|
||||
get_token_collector = 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.
|
||||
_RUN_LOCK_INIT_GUARD = threading.Lock()
|
||||
from .._run_lock import ensure_run_lock
|
||||
|
||||
# Task status constants
|
||||
class TaskStatus(Enum):
|
||||
@@ -2088,14 +2085,7 @@ class AgentTeam(SpawnAnnounceProtocol):
|
||||
guard (double-checked locking) so two threads racing into start()/astart()
|
||||
observe the same lock object.
|
||||
"""
|
||||
lock = getattr(self, '_run_lock', None)
|
||||
if lock is None:
|
||||
with _RUN_LOCK_INIT_GUARD:
|
||||
lock = getattr(self, '_run_lock', None)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
self._run_lock = lock
|
||||
return lock
|
||||
return ensure_run_lock(self)
|
||||
|
||||
def _execution_identity(self):
|
||||
"""Identity of the current execution for re-entrancy checks.
|
||||
|
||||
@@ -125,6 +125,13 @@ class BotConfig:
|
||||
|
||||
# Edit interval for streaming responses in milliseconds (default: 700ms)
|
||||
stream_edit_interval_ms: int = 700
|
||||
|
||||
# Opt-in privacy-safe runtime footer on the final reply (default: False).
|
||||
# When True, a short trailing line (e.g. ``— gpt-4o · 4.2s``) reporting the
|
||||
# model and turn latency is appended to the final reply only. Privacy-safe:
|
||||
# never surfaces tool args, URLs, paths or message content. Unavailable
|
||||
# fields are omitted rather than shown as placeholders.
|
||||
footer: bool = False
|
||||
|
||||
# Intentional silence support
|
||||
allow_silence: bool = False
|
||||
@@ -185,6 +192,7 @@ class BotConfig:
|
||||
"owner_user_id": "***" if self.owner_user_id else None,
|
||||
"streaming": self.streaming,
|
||||
"stream_edit_interval_ms": self.stream_edit_interval_ms,
|
||||
"footer": self.footer,
|
||||
"allow_silence": self.allow_silence,
|
||||
"silence_token": self.silence_token,
|
||||
"allow_bots": self.allow_bots,
|
||||
|
||||
@@ -6,6 +6,7 @@ Manages context window by compacting messages when needed.
|
||||
|
||||
import re
|
||||
import threading
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional, Callable, Awaitable, Tuple
|
||||
import asyncio
|
||||
|
||||
@@ -14,6 +15,8 @@ from .strategy import CompactionStrategy
|
||||
from .result import CompactionResult
|
||||
from .protocols import ToolResultPrunerProtocol, MessageFormatterProtocol, SummaryBuilderProtocol
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Cache for the one-time offline-safe tiktoken probe (see estimate_tokens).
|
||||
# None = not yet probed; True/False = probe result. Guarded by a lock so a
|
||||
@@ -337,18 +340,28 @@ class ContextCompactor:
|
||||
compacted = self._prune(processed_messages)
|
||||
elif self.strategy == CompactionStrategy.LLM_SUMMARIZE:
|
||||
if self.llm_summarize_fn:
|
||||
# For sync calls with LLM function, we need to run async
|
||||
# Route through the canonical sync-to-async bridge so the
|
||||
# user-configured LLM summariser runs identically whether or not
|
||||
# a loop is already running (e.g. an agent hosted in FastAPI,
|
||||
# aiohttp, Jupyter or a bot handler). The naive fallback is kept
|
||||
# only for genuine failures, and is now logged so a silent
|
||||
# degradation in context quality is visible.
|
||||
#
|
||||
# Keep the bridge's bounded default timeout rather than
|
||||
# ``timeout=None`` so a summariser that never completes (a hung
|
||||
# network call that ignores cancellation) falls back to the
|
||||
# naive summary below instead of pinning the caller's thread
|
||||
# indefinitely. ``TimeoutError`` is caught by the ``except`` here.
|
||||
from ..utils.async_bridge import run_coroutine_from_any_context
|
||||
try:
|
||||
# Check if we're already in an async context
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
# If in async context, fallback to naive summarization
|
||||
compacted = self._summarize(processed_messages)
|
||||
except RuntimeError:
|
||||
# No running loop, safe to create one
|
||||
compacted = asyncio.run(self._llm_summarize_async(processed_messages, focus_topic))
|
||||
compacted = run_coroutine_from_any_context(
|
||||
self._llm_summarize_async(processed_messages, focus_topic),
|
||||
)
|
||||
except Exception:
|
||||
# Fallback to naive summarization if async fails
|
||||
logger.warning(
|
||||
"LLM summarisation failed; using naive summary",
|
||||
exc_info=True,
|
||||
)
|
||||
compacted = self._summarize(processed_messages)
|
||||
else:
|
||||
compacted = self._llm_summarize(processed_messages, focus_topic)
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
"""Jev / System One typed decision models (Ollama ``/v1/systemone``, TypeSafe API)."""
|
||||
|
||||
from .routing import (
|
||||
DecisionRoutePlan,
|
||||
DecisionTriageRouter,
|
||||
default_ticket_triage_questions,
|
||||
triage_decision,
|
||||
triaged_start,
|
||||
)
|
||||
from .system_one import (
|
||||
KNOWN_DECISION_MODELS,
|
||||
SystemOneResult,
|
||||
asystem_one,
|
||||
choice_question,
|
||||
is_decision_model,
|
||||
noul_question,
|
||||
resolve_system_one_base_url,
|
||||
score_question,
|
||||
system_one,
|
||||
system_one_via_typesafe_sdk,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DecisionRoutePlan",
|
||||
"DecisionTriageRouter",
|
||||
"default_ticket_triage_questions",
|
||||
"triage_decision",
|
||||
"triaged_start",
|
||||
"KNOWN_DECISION_MODELS",
|
||||
"SystemOneResult",
|
||||
"asystem_one",
|
||||
"choice_question",
|
||||
"is_decision_model",
|
||||
"noul_question",
|
||||
"resolve_system_one_base_url",
|
||||
"score_question",
|
||||
"system_one",
|
||||
"system_one_via_typesafe_sdk",
|
||||
]
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Route user input through System One decision models before chat agents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Mapping, Optional, Protocol, Union
|
||||
|
||||
from .system_one import (
|
||||
SystemOneResult,
|
||||
choice_question,
|
||||
system_one,
|
||||
)
|
||||
|
||||
|
||||
class AgentLike(Protocol):
|
||||
def start(self, prompt: str, **kwargs: Any) -> Any: ...
|
||||
|
||||
|
||||
RouteTarget = Union[AgentLike, str, Callable[[SystemOneResult], Any]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class DecisionRoutePlan:
|
||||
"""Configuration for decision-first routing."""
|
||||
|
||||
route_question: str = "team"
|
||||
decision_model: str = field(
|
||||
default_factory=lambda: os.environ.get("OLLAMA_DECISION_MODEL", "nimble")
|
||||
)
|
||||
api_base: Optional[str] = None
|
||||
state_key: str = "message"
|
||||
min_confidence: Optional[float] = None
|
||||
questions: Mapping[str, Mapping[str, Any]] = field(default_factory=dict)
|
||||
routes: Mapping[str, RouteTarget] = field(default_factory=dict)
|
||||
fallback_route: Optional[str] = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not self.questions:
|
||||
raise ValueError("DecisionRoutePlan requires a non-empty questions mapping")
|
||||
if self.route_question not in self.questions:
|
||||
raise ValueError(
|
||||
f"route_question {self.route_question!r} missing from questions keys "
|
||||
f"{list(self.questions.keys())}"
|
||||
)
|
||||
route_q = self.questions[self.route_question]
|
||||
if route_q.get("type") != "choice":
|
||||
raise ValueError(
|
||||
f"route_question {self.route_question!r} must be a choice question, "
|
||||
f"got type {route_q.get('type')!r}"
|
||||
)
|
||||
|
||||
|
||||
def default_ticket_triage_questions() -> dict[str, dict[str, Any]]:
|
||||
"""Preset for support ticket routing (Ollama blog example)."""
|
||||
return {
|
||||
"team": choice_question(
|
||||
"Which team should handle this ticket?",
|
||||
{
|
||||
"billing": "Payments and refunds",
|
||||
"technical": "Bugs and integrations",
|
||||
"other": "None of the above",
|
||||
},
|
||||
),
|
||||
"refund": {
|
||||
"type": "noul",
|
||||
"instructions": "Does the customer explicitly ask for a refund?",
|
||||
},
|
||||
"urgency": {
|
||||
"type": "score",
|
||||
"instructions": "How urgent is this ticket?",
|
||||
"criteria": ["Routine", "Soon", "Urgent"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def triage_decision(
|
||||
message: str,
|
||||
*,
|
||||
questions: Mapping[str, Mapping[str, Any]],
|
||||
decision_model: str = "nimble",
|
||||
api_base: Optional[str] = None,
|
||||
state_key: str = "message",
|
||||
extra_state: Optional[Mapping[str, Any]] = None,
|
||||
) -> SystemOneResult:
|
||||
"""Run System One on *message* without invoking a chat LLM."""
|
||||
state = dict(extra_state or {})
|
||||
state[state_key] = message
|
||||
return system_one(
|
||||
state=state,
|
||||
questions=questions,
|
||||
model=decision_model,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
|
||||
def resolve_route_key(result: SystemOneResult, route_question: str) -> Optional[str]:
|
||||
return result.choice(route_question)
|
||||
|
||||
|
||||
def _invoke_target(target: RouteTarget, prompt: str, decision: SystemOneResult) -> Any:
|
||||
if isinstance(target, str):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
return Agent(instructions="You are a helpful assistant.", llm=target).start(prompt)
|
||||
if callable(target):
|
||||
return target(decision)
|
||||
return target.start(prompt)
|
||||
|
||||
|
||||
def triaged_start(
|
||||
prompt: str,
|
||||
plan: DecisionRoutePlan,
|
||||
*,
|
||||
extra_state: Optional[Mapping[str, Any]] = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Triage with a decision model, then delegate to the mapped chat agent/model.
|
||||
|
||||
Returns a dict with ``decision``, ``route``, ``confidence``, and ``response``.
|
||||
"""
|
||||
decision = triage_decision(
|
||||
prompt,
|
||||
questions=plan.questions,
|
||||
decision_model=plan.decision_model,
|
||||
api_base=plan.api_base,
|
||||
state_key=plan.state_key,
|
||||
extra_state=extra_state,
|
||||
)
|
||||
route_key = resolve_route_key(decision, plan.route_question)
|
||||
conf = decision.confidence(plan.route_question)
|
||||
if plan.min_confidence is not None and (conf is None or conf < plan.min_confidence):
|
||||
route_key = plan.fallback_route or route_key
|
||||
if not route_key or route_key not in plan.routes:
|
||||
route_key = plan.fallback_route or next(iter(plan.routes), None)
|
||||
if not route_key or route_key not in plan.routes:
|
||||
raise RuntimeError(
|
||||
f"No route for decision choice {route_key!r}; "
|
||||
f"available routes: {list(plan.routes.keys())}"
|
||||
)
|
||||
target = plan.routes[route_key]
|
||||
response = _invoke_target(target, prompt, decision)
|
||||
return {
|
||||
"decision": decision,
|
||||
"route": route_key,
|
||||
"confidence": conf,
|
||||
"response": response,
|
||||
}
|
||||
|
||||
|
||||
class DecisionTriageRouter:
|
||||
"""Object-oriented wrapper for repeated triage + chat routing."""
|
||||
|
||||
def __init__(self, plan: DecisionRoutePlan):
|
||||
self.plan = plan
|
||||
|
||||
def decide(self, message: str, **extra_state: Any) -> SystemOneResult:
|
||||
return triage_decision(
|
||||
message,
|
||||
questions=self.plan.questions,
|
||||
decision_model=self.plan.decision_model,
|
||||
api_base=self.plan.api_base,
|
||||
state_key=self.plan.state_key,
|
||||
extra_state=extra_state or None,
|
||||
)
|
||||
|
||||
def start(self, prompt: str, **extra_state: Any) -> dict[str, Any]:
|
||||
return triaged_start(prompt, self.plan, extra_state=extra_state or None)
|
||||
@@ -0,0 +1,289 @@
|
||||
"""Call Ollama / TypeSafe System One decision endpoint (``POST /v1/systemone``).
|
||||
|
||||
See: https://ollama.com/blog/ollama-now-supports-jev-style-decision-models
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Mapping, MutableMapping, Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
KNOWN_DECISION_MODELS = frozenset({"nimble", "tev1", "tev1:0.8b"})
|
||||
|
||||
|
||||
def is_decision_model(model: str) -> bool:
|
||||
"""Return True if *model* looks like a known Ollama decision model id."""
|
||||
name = (model or "").strip().split("/")[-1]
|
||||
if name in KNOWN_DECISION_MODELS:
|
||||
return True
|
||||
return name.startswith("tev1")
|
||||
|
||||
|
||||
def _base_source(api_base: Optional[str] = None) -> tuple[str, Optional[str]]:
|
||||
"""Return ``(base_url, source)`` where *source* names the env var selected.
|
||||
|
||||
*source* is ``None`` when ``api_base`` is explicit, ``"TYPESAFE_BASE_URL"`` or
|
||||
``"OLLAMA_HOST"`` when resolved from the environment, or ``"default"`` for the
|
||||
loopback fallback. The source lets callers scope credentials to the host that
|
||||
was actually chosen (so a TypeSafe key is never sent to an Ollama host).
|
||||
"""
|
||||
if api_base:
|
||||
return api_base.rstrip("/"), None
|
||||
for key in ("TYPESAFE_BASE_URL", "OLLAMA_HOST"):
|
||||
raw = (os.environ.get(key) or "").strip()
|
||||
if raw:
|
||||
from praisonaiagents.local.target import parse_ollama_host
|
||||
|
||||
parsed = parse_ollama_host(raw) if not raw.startswith("http") else raw.rstrip("/")
|
||||
return (parsed or raw).rstrip("/"), key
|
||||
return "http://127.0.0.1:11434", "default"
|
||||
|
||||
|
||||
def resolve_system_one_base_url(api_base: Optional[str] = None) -> str:
|
||||
"""Resolve base URL for ``/v1/systemone`` (Ollama 0.35+ or TypeSafe cloud)."""
|
||||
return _base_source(api_base)[0]
|
||||
|
||||
|
||||
def _is_loopback_host(host: Optional[str]) -> bool:
|
||||
if not host:
|
||||
return False
|
||||
if host.lower() == "localhost":
|
||||
return True
|
||||
try:
|
||||
return ipaddress.ip_address(host).is_loopback
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _validate_base_url(base: str) -> None:
|
||||
"""Reject non-HTTP(S) schemes and plain HTTP to non-loopback hosts."""
|
||||
parsed = urlparse(base)
|
||||
if parsed.scheme not in {"http", "https"}:
|
||||
raise ValueError(f"system_one requires an http(s) base URL, got {base!r}")
|
||||
if parsed.scheme == "http" and not _is_loopback_host(parsed.hostname):
|
||||
raise ValueError(
|
||||
f"system_one refuses plain HTTP for non-loopback host {parsed.hostname!r}; use https"
|
||||
)
|
||||
|
||||
|
||||
def choice_question(instructions: str, criteria: Mapping[str, str]) -> dict[str, Any]:
|
||||
return {"type": "choice", "instructions": instructions, "criteria": dict(criteria)}
|
||||
|
||||
|
||||
def noul_question(instructions: str) -> dict[str, Any]:
|
||||
return {"type": "noul", "instructions": instructions}
|
||||
|
||||
|
||||
def score_question(instructions: str, criteria: list[str]) -> dict[str, Any]:
|
||||
return {"type": "score", "instructions": instructions, "criteria": list(criteria)}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SystemOneResult:
|
||||
"""Parsed ``/v1/systemone`` response."""
|
||||
|
||||
model: str
|
||||
answers: dict[str, Any]
|
||||
usage: dict[str, Any] = field(default_factory=dict)
|
||||
raw: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def choice(self, name: str) -> Optional[str]:
|
||||
entry = self.answers.get(name) or {}
|
||||
if entry.get("type") == "choice":
|
||||
return entry.get("choice")
|
||||
return None
|
||||
|
||||
def noul(self, name: str) -> Optional[float]:
|
||||
entry = self.answers.get(name) or {}
|
||||
if entry.get("type") == "noul":
|
||||
val = entry.get("noul")
|
||||
return float(val) if val is not None else None
|
||||
return None
|
||||
|
||||
def score(self, name: str) -> Optional[float]:
|
||||
entry = self.answers.get(name) or {}
|
||||
if entry.get("type") == "score":
|
||||
val = entry.get("score")
|
||||
return float(val) if val is not None else None
|
||||
return None
|
||||
|
||||
def confidence(self, name: str) -> Optional[float]:
|
||||
entry = self.answers.get(name) or {}
|
||||
val = entry.get("confidence")
|
||||
return float(val) if val is not None else None
|
||||
|
||||
|
||||
def _default_timeout(explicit: Optional[float]) -> float:
|
||||
if explicit is not None:
|
||||
return explicit
|
||||
raw = (os.environ.get("SYSTEM_ONE_TIMEOUT") or os.environ.get("OLLAMA_SYSTEM_ONE_TIMEOUT") or "").strip()
|
||||
if raw:
|
||||
try:
|
||||
return float(raw)
|
||||
except ValueError:
|
||||
pass
|
||||
return 120.0
|
||||
|
||||
|
||||
def _default_api_key(api_key: Optional[str], *, allow_typesafe_key: bool) -> Optional[str]:
|
||||
"""Resolve the bearer key, scoping ``TYPESAFE_API_KEY`` to TypeSafe hosts only."""
|
||||
if api_key is not None:
|
||||
return api_key
|
||||
if allow_typesafe_key and os.environ.get("TYPESAFE_API_KEY"):
|
||||
return os.environ["TYPESAFE_API_KEY"]
|
||||
return os.environ.get("OLLAMA_API_KEY") or "ollama"
|
||||
|
||||
|
||||
def _build_payload(
|
||||
*,
|
||||
model: str,
|
||||
state: Mapping[str, Any],
|
||||
questions: Mapping[str, Mapping[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"model": model,
|
||||
"state": dict(state),
|
||||
"questions": {k: dict(v) for k, v in questions.items()},
|
||||
}
|
||||
|
||||
|
||||
def _parse_response(data: dict[str, Any]) -> SystemOneResult:
|
||||
return SystemOneResult(
|
||||
model=str(data.get("model") or ""),
|
||||
answers=dict(data.get("answers") or {}),
|
||||
usage=dict(data.get("usage") or {}),
|
||||
raw=data,
|
||||
)
|
||||
|
||||
|
||||
def system_one(
|
||||
*,
|
||||
state: Mapping[str, Any],
|
||||
questions: Mapping[str, Mapping[str, Any]],
|
||||
model: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
timeout: Optional[float] = None,
|
||||
) -> SystemOneResult:
|
||||
"""
|
||||
Run a System One decision request against Ollama or TypeSafe.
|
||||
|
||||
Example::
|
||||
|
||||
from praisonaiagents.decisions import system_one, choice_question, noul_question
|
||||
|
||||
result = system_one(
|
||||
state={"ticket": "Please refund my duplicate charge."},
|
||||
questions={
|
||||
"team": choice_question(
|
||||
"Which team?",
|
||||
{"billing": "Payments", "technical": "Bugs", "other": "Other"},
|
||||
),
|
||||
"refund": noul_question("Does the user ask for a refund?"),
|
||||
},
|
||||
model="nimble",
|
||||
)
|
||||
assert result.choice("team") == "billing"
|
||||
"""
|
||||
resolved_model = (
|
||||
model
|
||||
or os.environ.get("TYPESAFE_DEFAULT_MODEL")
|
||||
or os.environ.get("OLLAMA_DECISION_MODEL")
|
||||
or "nimble"
|
||||
)
|
||||
base, source = _base_source(api_base)
|
||||
_validate_base_url(base)
|
||||
url = f"{base}/v1/systemone"
|
||||
body = json.dumps(_build_payload(model=resolved_model, state=state, questions=questions)).encode("utf-8")
|
||||
headers = {"Content-Type": "application/json", "Accept": "application/json"}
|
||||
key = _default_api_key(api_key, allow_typesafe_key=source == "TYPESAFE_BASE_URL")
|
||||
if key:
|
||||
headers["Authorization"] = f"Bearer {key}"
|
||||
timeout = _default_timeout(timeout)
|
||||
req = urllib.request.Request(url, data=body, headers=headers, method="POST")
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
data = json.loads(resp.read().decode("utf-8"))
|
||||
except urllib.error.HTTPError as exc:
|
||||
detail = exc.read().decode("utf-8", errors="replace")[:2000]
|
||||
raise RuntimeError(f"system_one HTTP {exc.code}: {detail}") from exc
|
||||
except urllib.error.URLError as exc:
|
||||
raise RuntimeError(
|
||||
f"system_one failed to reach {url!r}. Is Ollama 0.35+ running with a decision model pulled?"
|
||||
) from exc
|
||||
if not isinstance(data, dict):
|
||||
raise RuntimeError(f"system_one expected JSON object, got {type(data).__name__}")
|
||||
return _parse_response(data)
|
||||
|
||||
|
||||
async def asystem_one(
|
||||
*,
|
||||
state: Mapping[str, Any],
|
||||
questions: Mapping[str, Mapping[str, Any]],
|
||||
model: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
timeout: Optional[float] = None,
|
||||
) -> SystemOneResult:
|
||||
"""Async ``system_one`` using aiohttp (already a core dependency)."""
|
||||
import aiohttp
|
||||
|
||||
resolved_model = (
|
||||
model
|
||||
or os.environ.get("TYPESAFE_DEFAULT_MODEL")
|
||||
or os.environ.get("OLLAMA_DECISION_MODEL")
|
||||
or "nimble"
|
||||
)
|
||||
base, source = _base_source(api_base)
|
||||
_validate_base_url(base)
|
||||
url = f"{base}/v1/systemone"
|
||||
payload = _build_payload(model=resolved_model, state=state, questions=questions)
|
||||
headers: MutableMapping[str, str] = {"Content-Type": "application/json", "Accept": "application/json"}
|
||||
key = _default_api_key(api_key, allow_typesafe_key=source == "TYPESAFE_BASE_URL")
|
||||
if key:
|
||||
headers["Authorization"] = f"Bearer {key}"
|
||||
timeout_obj = aiohttp.ClientTimeout(total=_default_timeout(timeout))
|
||||
async with aiohttp.ClientSession(timeout=timeout_obj) as session:
|
||||
async with session.post(url, json=payload, headers=headers) as resp:
|
||||
text = await resp.text()
|
||||
if resp.status >= 400:
|
||||
raise RuntimeError(f"system_one HTTP {resp.status}: {text[:2000]}")
|
||||
data = json.loads(text)
|
||||
if not isinstance(data, dict):
|
||||
raise RuntimeError(f"system_one expected JSON object, got {type(data).__name__}")
|
||||
return _parse_response(data)
|
||||
|
||||
|
||||
def system_one_via_typesafe_sdk(
|
||||
*,
|
||||
state: Mapping[str, Any],
|
||||
questions: Mapping[str, Any],
|
||||
model: Optional[str] = None,
|
||||
timeout: float = 120.0,
|
||||
) -> SystemOneResult:
|
||||
"""Optional path through the official ``typesafe-sdk`` (``pip install typesafe-sdk``)."""
|
||||
try:
|
||||
from typesafe_sdk import TypeSafeClient
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"typesafe-sdk is not installed. Use: pip install typesafe-sdk "
|
||||
"or call system_one() which uses HTTP directly."
|
||||
) from exc
|
||||
|
||||
with TypeSafeClient(timeout=timeout) as client:
|
||||
result = client.system_one(state=dict(state), questions=dict(questions), model=model)
|
||||
answers: dict[str, Any] = {}
|
||||
for key, val in getattr(result, "answers", {}).items():
|
||||
answers[key] = val if isinstance(val, dict) else getattr(val, "__dict__", {"value": val})
|
||||
return SystemOneResult(
|
||||
model=str(getattr(result, "model", model or "")),
|
||||
answers=answers,
|
||||
usage=dict(getattr(result, "usage", {}) or {}),
|
||||
raw={"via": "typesafe-sdk"},
|
||||
)
|
||||
@@ -11,6 +11,7 @@ import time
|
||||
import asyncio
|
||||
from praisonaiagents._logging import get_logger
|
||||
import subprocess
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
from .types import (
|
||||
@@ -72,6 +73,47 @@ class HookRunner:
|
||||
self._registry = registry if registry is not None else HookRegistry()
|
||||
self._default_timeout = default_timeout
|
||||
self._cwd = cwd or os.getcwd()
|
||||
# Dedicated executor for sync hooks. A hook that times out cannot be
|
||||
# cancelled and keeps its worker slot; isolating it here means a
|
||||
# misbehaving hook can only exhaust hook capacity, never the loop's
|
||||
# shared default executor that guardrails/tools/locks also rely on.
|
||||
self._hook_pool: Optional[ThreadPoolExecutor] = None
|
||||
# hook.id -> in-flight Future for a sync hook. The slot is reserved
|
||||
# before submission and only cleared when the worker finishes, so a
|
||||
# still-running (timed-out or cancelled) hook caps new launches at one
|
||||
# thread per hook. Guarded by ``_stuck_lock`` so concurrent invocations
|
||||
# of the same hook cannot both reserve the slot and race past the cap.
|
||||
self._stuck: Dict[str, Any] = {}
|
||||
self._stuck_lock_obj: Optional[asyncio.Lock] = None
|
||||
|
||||
@property
|
||||
def _stuck_lock(self) -> asyncio.Lock:
|
||||
"""Lazily create the in-flight-slot lock on first use.
|
||||
|
||||
Created lazily (not in ``__init__``) because a ``HookRunner`` may be
|
||||
constructed before any event loop exists; the lock is only needed once
|
||||
we are already inside ``async def execute`` / ``_execute_function_hook``.
|
||||
"""
|
||||
if self._stuck_lock_obj is None:
|
||||
self._stuck_lock_obj = asyncio.Lock()
|
||||
return self._stuck_lock_obj
|
||||
|
||||
def _release_stuck(self, hook_id: str, fut: Any) -> None:
|
||||
"""Clear the in-flight slot once the worker Future completes.
|
||||
|
||||
Only removes the entry if it still points at *this* Future, so a newer
|
||||
reserved invocation is never clobbered by an older worker draining.
|
||||
"""
|
||||
if self._stuck.get(hook_id) is fut:
|
||||
self._stuck.pop(hook_id, None)
|
||||
|
||||
def _get_hook_pool(self) -> ThreadPoolExecutor:
|
||||
"""Lazily create the dedicated sync-hook executor."""
|
||||
if self._hook_pool is None:
|
||||
self._hook_pool = ThreadPoolExecutor(
|
||||
max_workers=8, thread_name_prefix="praison-hook"
|
||||
)
|
||||
return self._hook_pool
|
||||
|
||||
@property
|
||||
def registry(self) -> HookRegistry:
|
||||
@@ -167,27 +209,30 @@ class HookRunner:
|
||||
Returns:
|
||||
List of execution results
|
||||
"""
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
loop = None
|
||||
|
||||
if loop is not None and loop.is_running():
|
||||
# Cannot execute sync in running event loop - would block the loop
|
||||
# Guide users to use the async API instead
|
||||
raise RuntimeError(
|
||||
"execute_sync() cannot be called from within a running event loop. "
|
||||
"Use 'await runner.execute(event, input_data, target)' instead in async contexts."
|
||||
)
|
||||
|
||||
# Fast path: no hooks → no loop creation needed
|
||||
# This mirrors the same check in execute() but avoids event loop creation overhead
|
||||
hooks = self._registry.get_hooks(event, target)
|
||||
if not hooks:
|
||||
return []
|
||||
|
||||
# No running loop and hooks exist — safe to create one
|
||||
return asyncio.run(self.execute(event, input_data, target, _hooks=hooks))
|
||||
# Use the canonical sync-to-async bridge so blocking policy/guardrail
|
||||
# hooks still run (and can deny) from a running event loop — e.g. a
|
||||
# FastAPI/aiohttp handler, Jupyter, or a bot callback invoking
|
||||
# agent.start()/agent.chat(). Raising here would turn documented sync
|
||||
# entry points into hard failures whenever any hook is registered.
|
||||
#
|
||||
# Bound the bridge wait so a hook that ignores its own per-hook timeout
|
||||
# (e.g. a blocking C call on the worker loop) can never pin the caller's
|
||||
# thread forever. Each hook already enforces ``default_timeout``
|
||||
# internally; we give the whole chain headroom of one timeout per hook
|
||||
# so legitimately slow sequential chains are not cut short, while still
|
||||
# providing a hard backstop instead of ``timeout=None``.
|
||||
from ..utils.async_bridge import run_coroutine_from_any_context
|
||||
bridge_timeout = max(self._default_timeout, 1.0) * (len(hooks) + 1)
|
||||
return run_coroutine_from_any_context(
|
||||
self.execute(event, input_data, target, _hooks=hooks),
|
||||
timeout=bridge_timeout,
|
||||
)
|
||||
|
||||
async def _execute_parallel(
|
||||
self,
|
||||
@@ -394,11 +439,52 @@ class HookRunner:
|
||||
timeout=timeout
|
||||
)
|
||||
else:
|
||||
# Execute sync function in thread pool with timeout
|
||||
loop = asyncio.get_event_loop()
|
||||
# Reserve the single in-flight slot for this hook *under a lock*
|
||||
# and *before* submitting, so two concurrent invocations can't
|
||||
# both see an empty ``_stuck`` and both launch a worker (which
|
||||
# would exceed the intended one-worker-per-hook cap and double
|
||||
# the side effects). The slot is only cleared once the worker
|
||||
# finishes; a still-running previous worker fails this one
|
||||
# closed immediately instead of leaking another thread.
|
||||
async with self._stuck_lock:
|
||||
prev = self._stuck.get(hook.id)
|
||||
if prev is not None and not prev.done():
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error="previous invocation still running past its timeout",
|
||||
duration_ms=duration,
|
||||
output=HookResult(
|
||||
decision="deny",
|
||||
reason=f"Hook '{hook.name}' is still running past its timeout",
|
||||
),
|
||||
)
|
||||
|
||||
# Execute sync function in the dedicated hook pool with
|
||||
# timeout. ``copy_context_to_callable`` preserves
|
||||
# trace/session contextvars across the thread, matching
|
||||
# other executor sites.
|
||||
from ..trace.context_events import copy_context_to_callable
|
||||
loop = asyncio.get_event_loop()
|
||||
func = hook.func
|
||||
data = input_data
|
||||
fut = loop.run_in_executor(
|
||||
self._get_hook_pool(),
|
||||
copy_context_to_callable(lambda: func(data)),
|
||||
)
|
||||
# Reserve the slot now; a completion callback clears it only
|
||||
# when the worker actually finishes, so a timeout *or* a
|
||||
# cancellation leaves the slot held until the thread drains.
|
||||
self._stuck[hook.id] = fut
|
||||
fut.add_done_callback(
|
||||
lambda f, hid=hook.id: self._release_stuck(hid, f)
|
||||
)
|
||||
|
||||
result = await asyncio.wait_for(
|
||||
loop.run_in_executor(None, hook.func, input_data),
|
||||
timeout=timeout
|
||||
asyncio.shield(fut), timeout=timeout
|
||||
)
|
||||
|
||||
duration = (time.time() - start_time) * 1000
|
||||
|
||||
@@ -9,6 +9,7 @@ Provides both legacy API (classify_error) and new structured classification
|
||||
"""
|
||||
|
||||
import re
|
||||
import math
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
@@ -644,7 +645,9 @@ def extract_retry_after(
|
||||
retry_after = None
|
||||
if retry_after is not None:
|
||||
try:
|
||||
return min(float(retry_after), cap_seconds)
|
||||
delay = float(retry_after)
|
||||
if math.isfinite(delay) and delay >= 0:
|
||||
return min(delay, cap_seconds)
|
||||
except (ValueError, TypeError):
|
||||
# Not a plain delta-seconds value — RFC 7231 also permits an
|
||||
# HTTP-date. Parse it so a provider that sends an absolute reset
|
||||
@@ -658,24 +661,44 @@ def extract_retry_after(
|
||||
# 2. Some SDKs expose a numeric ``retry_after`` attribute directly.
|
||||
retry_after_attr = getattr(error, "retry_after", None)
|
||||
if isinstance(retry_after_attr, (int, float)):
|
||||
return min(float(retry_after_attr), cap_seconds)
|
||||
try:
|
||||
delay = float(retry_after_attr)
|
||||
except OverflowError:
|
||||
delay = float("inf")
|
||||
if math.isfinite(delay) and delay >= 0:
|
||||
return min(delay, cap_seconds)
|
||||
|
||||
error_str = str(error)
|
||||
|
||||
# 3. Fall back to parsing common Retry-After patterns from the message.
|
||||
# Capture the full signed token so a negative hint (e.g. "-3600 seconds")
|
||||
# is rejected by the finite/non-negative check below instead of being
|
||||
# silently read as a positive delay. Match a whole numeric token so
|
||||
# leading-dot fractions, exponents and malformed suffixes cannot be
|
||||
# mistaken for an unrelated positive integer.
|
||||
# A period may terminate a sentence, but may not split an embedded token.
|
||||
number = (
|
||||
r"(?<![\w.+-])([+-]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][+-]?\d+)?)"
|
||||
r"(?![\deE+-]|\.(?!\s|$|[\"')\]}])|\.\s*(?:seconds?|s)\b)"
|
||||
)
|
||||
patterns = [
|
||||
r"retry.?after[:\s]+(\d+)",
|
||||
r"retry[:\s]+(\d+)",
|
||||
r"wait[:\s]+(\d+)",
|
||||
r"(\d+).*second",
|
||||
# Quoted duration values must close immediately after the unit.
|
||||
r'"?retryDelay"?\s*:\s*"' + number + r's"',
|
||||
r'"?retryDelay"?\s*:\s*' + number + r's(?=$|[\s,;}\]])',
|
||||
r"try again in\s+" + number,
|
||||
r"retry.?after[:\s]+" + number,
|
||||
r"retry[:\s]+" + number,
|
||||
r"wait[:\s]+" + number,
|
||||
number + r"\s*second",
|
||||
]
|
||||
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, error_str, re.IGNORECASE)
|
||||
if match:
|
||||
# Invalid earlier hints must not hide a later valid hint of this form.
|
||||
for match in re.finditer(pattern, error_str, re.IGNORECASE):
|
||||
try:
|
||||
delay = float(match.group(1))
|
||||
return min(delay, cap_seconds)
|
||||
if math.isfinite(delay) and delay >= 0:
|
||||
return min(delay, cap_seconds)
|
||||
except (ValueError, IndexError):
|
||||
continue
|
||||
|
||||
|
||||
@@ -8,6 +8,8 @@ import re
|
||||
import inspect
|
||||
import asyncio
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from functools import wraps
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Union, Literal, Callable, TYPE_CHECKING, Protocol
|
||||
|
||||
@@ -45,6 +47,22 @@ _TOOL_ARGUMENTS_PARSE_FAILED = object()
|
||||
_EVENT_TRACKING_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def _request_stop_reason(method):
|
||||
"""Keep nested response calls from replacing their caller's outcome."""
|
||||
if inspect.iscoroutinefunction(method):
|
||||
@wraps(method)
|
||||
async def async_call(self, *args, **kwargs):
|
||||
with self._stop_reason_scope():
|
||||
return await method(self, *args, **kwargs)
|
||||
return async_call
|
||||
|
||||
@wraps(method)
|
||||
def sync_call(self, *args, **kwargs):
|
||||
with self._stop_reason_scope():
|
||||
return method(self, *args, **kwargs)
|
||||
return sync_call
|
||||
|
||||
|
||||
def _ollama_tool_result_is_successful(tool_result: Any) -> bool:
|
||||
"""Return whether a raw Ollama callback result is safe to chain.
|
||||
|
||||
@@ -542,7 +560,9 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
# Structured stop reason for the last tool-execution loop:
|
||||
# "completed" | "max_steps" | "error". Lets callers distinguish a finished
|
||||
# task from one truncated by the step budget (see ExecutionConfig.max_steps).
|
||||
self._last_stop_reason = "completed"
|
||||
self._last_stop_reason_var = contextvars.ContextVar("last_stop_reason", default=None)
|
||||
self._stop_reason_active_var = contextvars.ContextVar("stop_reason_active", default=None)
|
||||
self._last_stop_reason_fallback = (None, "completed")
|
||||
self._idle_timeout_breaker = IdleTimeoutBreaker() # Circuit breaker for idle timeouts
|
||||
self.chat_history = []
|
||||
# Optional Agent-supplied thread-safe append for deferred re-injection.
|
||||
@@ -855,29 +875,15 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
Returns:
|
||||
Retry delay in seconds, or default if not found
|
||||
"""
|
||||
# Try to find retry delay patterns like "retryDelay: 58s" or "retry after 58 seconds"
|
||||
patterns = [
|
||||
r'"retryDelay":\s*"(\d+)s"', # JSON format: "retryDelay": "58s"
|
||||
r'retryDelay:\s*"?(\d+)s"?', # retryDelay: 58s or retryDelay: "58s"
|
||||
r'retry.{0,10}?(\d+)\s*second', # retry after 58 seconds (non-greedy)
|
||||
r'wait\s+(\d+)\s*second', # wait 58 seconds
|
||||
r'try again in (\d+)', # try again in 58
|
||||
r'Retry-After:\s*(\d+)', # HTTP Retry-After header
|
||||
]
|
||||
|
||||
# Max delay cap to prevent unbounded sleep (5 minutes default)
|
||||
max_delay = 300
|
||||
if self._rate_limiter is not None and hasattr(self._rate_limiter, 'max_retry_delay'):
|
||||
max_delay = self._rate_limiter.max_retry_delay
|
||||
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, error_message, re.IGNORECASE)
|
||||
if match:
|
||||
delay = float(match.group(1))
|
||||
# Clamp to safe bounds [0, max_delay] to prevent unbounded sleep
|
||||
return max(0, min(delay, max_delay))
|
||||
from .error_classifier import extract_retry_after
|
||||
|
||||
return self._retry_delay
|
||||
delay = extract_retry_after(Exception(error_message), cap_seconds=max_delay)
|
||||
return self._retry_delay if delay is None else delay
|
||||
|
||||
def _is_rate_limit_error(self, error: Exception) -> bool:
|
||||
"""Check if an exception is a rate limit error.
|
||||
@@ -2946,6 +2952,7 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
logging.debug(f"In-loop context management skipped: {e}")
|
||||
return messages
|
||||
|
||||
@_request_stop_reason
|
||||
def get_response(
|
||||
self,
|
||||
prompt: Union[str, List[Dict]],
|
||||
@@ -3153,6 +3160,11 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
# OpenAI models use the Responses API which returns text
|
||||
# and tool calls as separate output items — no content:null.
|
||||
if self._supports_responses_api():
|
||||
# A refusal belongs to one provider response. If that
|
||||
# response also requested tools, classify the follow-up
|
||||
# independently without clearing stronger loop outcomes.
|
||||
if iteration_count > 0 and self._last_stop_reason == "refused":
|
||||
self._last_stop_reason = "completed"
|
||||
responses_params = self._build_responses_params(
|
||||
messages=messages,
|
||||
tools=formatted_tools,
|
||||
@@ -5183,6 +5195,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# Other errors are generally not recoverable
|
||||
return False
|
||||
|
||||
@_request_stop_reason
|
||||
async def get_response_async(
|
||||
self,
|
||||
prompt: Union[str, List[Dict]],
|
||||
@@ -5384,6 +5397,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
# ── Responses API path (async) ──────────────────────────
|
||||
if self._supports_responses_api():
|
||||
if iteration_count > 0 and self._last_stop_reason == "refused":
|
||||
self._last_stop_reason = "completed"
|
||||
responses_params = self._build_responses_params(
|
||||
messages=messages,
|
||||
tools=formatted_tools,
|
||||
@@ -6441,8 +6456,9 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
truncation reason (``content_filtered | refused | length_truncated``)
|
||||
when the last completion was blocked, so a blocked/refused/truncated turn
|
||||
is surfaced as an explicit terminal reason instead of a silent empty
|
||||
``completed``. Only updates when the reason is still ``"completed"`` so a
|
||||
prior ``max_steps`` (sticky truncation) is never downgraded. Absent or
|
||||
``completed``. A known Responses incomplete reason also replaces an
|
||||
earlier refusal delta so tool continuation cannot clear that outcome.
|
||||
Other prior terminal reasons, including ``max_steps``, remain sticky. Absent or
|
||||
unrecognised finish reasons are a no-op — zero overhead on success.
|
||||
"""
|
||||
try:
|
||||
@@ -6465,11 +6481,19 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
finish_reason = getattr(choice, "finish_reason", None)
|
||||
msg = getattr(choice, "message", None)
|
||||
refusal = getattr(msg, "refusal", None) if msg is not None else None
|
||||
if not choices:
|
||||
from .openai_client import OpenAIClient
|
||||
finish_reason = OpenAIClient._responses_incomplete_finish_reason(response)
|
||||
refusal = OpenAIClient._extract_responses_refusal(response)
|
||||
if finish_reason is None and not refusal:
|
||||
return
|
||||
from ..agent.run_outcome import classify_finish_reason
|
||||
reason = classify_finish_reason(finish_reason, refusal)
|
||||
if reason is not None and self._last_stop_reason == "completed":
|
||||
responses_incomplete = not choices and finish_reason is not None
|
||||
reason = classify_finish_reason(finish_reason, None if responses_incomplete else refusal)
|
||||
if reason is not None and (
|
||||
self._last_stop_reason == "completed"
|
||||
or (responses_incomplete and self._last_stop_reason == "refused")
|
||||
):
|
||||
self._last_stop_reason = reason
|
||||
except Exception:
|
||||
# Never let outcome classification break the response path.
|
||||
@@ -6520,6 +6544,79 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
logging.warning(f"Failed to extract token usage: {e}")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _stop_reason_context():
|
||||
"""Identify the calling thread and, when present, asyncio task."""
|
||||
try:
|
||||
task = asyncio.current_task()
|
||||
except RuntimeError:
|
||||
task = None
|
||||
return threading.get_ident(), id(task) if task is not None else None
|
||||
|
||||
@contextmanager
|
||||
def _stop_reason_scope(self):
|
||||
"""Restore an active parent's binding when a nested call returns."""
|
||||
owner = self._stop_reason_context()
|
||||
active = getattr(self, "_stop_reason_active_var", None)
|
||||
if active is None:
|
||||
active = contextvars.ContextVar("stop_reason_active", default=None)
|
||||
self._stop_reason_active_var = active
|
||||
nested = active.get() == owner
|
||||
active_token = active.set(owner)
|
||||
reason = getattr(self, "_last_stop_reason_var", None)
|
||||
if reason is None:
|
||||
reason = contextvars.ContextVar("last_stop_reason", default=None)
|
||||
self._last_stop_reason_var = reason
|
||||
reason_token = reason.set((owner, "completed"))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
result = reason.get()
|
||||
reason.reset(reason_token)
|
||||
active.reset(active_token)
|
||||
if nested:
|
||||
self._last_stop_reason_fallback = reason.get()
|
||||
else:
|
||||
reason.set(result)
|
||||
self._last_stop_reason_fallback = result
|
||||
|
||||
@property
|
||||
def _last_stop_reason(self) -> str:
|
||||
"""Terminal outcome in the calling task or thread.
|
||||
|
||||
Concurrent request callers retain their own reason. A synchronous
|
||||
observer after asyncio.run on the same thread reads the latest async
|
||||
result, rather than a stale reason left in its outer context.
|
||||
"""
|
||||
var = getattr(self, "_last_stop_reason_var", None)
|
||||
binding = var.get() if var is not None else None
|
||||
owner = self._stop_reason_context()
|
||||
latest_owner, latest_reason = getattr(
|
||||
self, "_last_stop_reason_fallback", (None, "completed")
|
||||
)
|
||||
if binding is None or binding[0] != owner:
|
||||
return latest_reason
|
||||
active = getattr(self, "_stop_reason_active_var", None)
|
||||
if active is not None and active.get() == owner:
|
||||
return binding[1]
|
||||
if (
|
||||
owner[1] is None and latest_owner is not None
|
||||
and latest_owner[0] == owner[0] and latest_owner[1] is not None
|
||||
):
|
||||
return latest_reason
|
||||
return binding[1]
|
||||
|
||||
@_last_stop_reason.setter
|
||||
def _last_stop_reason(self, reason: str) -> None:
|
||||
"""Record an outcome without changing concurrent callers' bindings."""
|
||||
var = getattr(self, "_last_stop_reason_var", None)
|
||||
if var is None:
|
||||
var = contextvars.ContextVar("last_stop_reason", default=None)
|
||||
self._last_stop_reason_var = var
|
||||
binding = (self._stop_reason_context(), reason)
|
||||
var.set(binding)
|
||||
self._last_stop_reason_fallback = binding
|
||||
|
||||
@property
|
||||
def current_agent_name(self) -> Optional[str]:
|
||||
"""Task-local name of the agent currently driving this LLM.
|
||||
@@ -6569,7 +6666,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
self.current_agent_id = agent_id
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
"""Deep-copy the LLM while giving the clone fresh attribution ContextVars.
|
||||
"""Deep-copy the LLM with fresh attribution and outcome ContextVars.
|
||||
|
||||
``contextvars.ContextVar`` has no ``__deepcopy__``/``__reduce__`` and is
|
||||
not copyable/pickleable, so ``Agent.__deepcopy__`` (which recursively
|
||||
@@ -6588,6 +6685,19 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
"current_agent_name", default=None
|
||||
)
|
||||
continue
|
||||
if key == "_last_stop_reason_var":
|
||||
clone.__dict__[key] = contextvars.ContextVar(
|
||||
"last_stop_reason", default=None
|
||||
)
|
||||
continue
|
||||
if key == "_stop_reason_active_var":
|
||||
clone.__dict__[key] = contextvars.ContextVar(
|
||||
"stop_reason_active", default=None
|
||||
)
|
||||
continue
|
||||
if key == "_last_stop_reason_fallback":
|
||||
clone.__dict__[key] = (None, "completed")
|
||||
continue
|
||||
if key == "_current_agent_id_var":
|
||||
clone.__dict__[key] = contextvars.ContextVar(
|
||||
"current_agent_id", default=None
|
||||
@@ -6884,10 +6994,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
web_fetch_tool['max_uses'] = 5
|
||||
|
||||
# Add web_fetch tool to existing tools or create tools list
|
||||
if 'tools' in params and params['tools']:
|
||||
params['tools'].append(web_fetch_tool)
|
||||
else:
|
||||
params['tools'] = [web_fetch_tool]
|
||||
params['tools'] = list(params.get('tools') or []) + [web_fetch_tool]
|
||||
|
||||
logging.debug(f"Web fetch enabled with tool: {web_fetch_tool}")
|
||||
|
||||
@@ -6899,17 +7006,13 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
memory_tool_def = memory_tool.get_tool_definition()
|
||||
|
||||
# Add memory tool to existing tools or create tools list
|
||||
if 'tools' in params and params['tools']:
|
||||
params['tools'].append(memory_tool_def)
|
||||
else:
|
||||
params['tools'] = [memory_tool_def]
|
||||
params['tools'] = list(params.get('tools') or []) + [memory_tool_def]
|
||||
|
||||
# Add the beta header for Anthropic
|
||||
beta_header = memory_tool.get_beta_header()
|
||||
if 'extra_headers' in params:
|
||||
params['extra_headers']['anthropic-beta'] = beta_header
|
||||
else:
|
||||
params['extra_headers'] = {'anthropic-beta': beta_header}
|
||||
extra_headers = dict(params.get('extra_headers') or {})
|
||||
extra_headers['anthropic-beta'] = beta_header
|
||||
params['extra_headers'] = extra_headers
|
||||
|
||||
logging.debug(f"Claude memory tool enabled with beta header: {beta_header}")
|
||||
|
||||
@@ -7158,6 +7261,15 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
self._track_token_usage(response, self._response_model_for_tracking(response))
|
||||
return response
|
||||
|
||||
def _raise_responses_failure(self, error) -> None:
|
||||
"""Surface a failed Responses request without accepting partial output."""
|
||||
code = error.get("code") if isinstance(error, dict) else getattr(error, "code", None)
|
||||
message = error.get("message") if isinstance(error, dict) else getattr(error, "message", None)
|
||||
if self._last_stop_reason == "completed":
|
||||
self._last_stop_reason = "error"
|
||||
label = f"Responses API failed ({code})" if code else "Responses API failed"
|
||||
raise LLMResponseError(f"{label}: {message or 'Provider returned a failed response.'}")
|
||||
|
||||
def _extract_from_responses_output(self, response) -> tuple:
|
||||
"""
|
||||
Parse a ``ResponsesAPIResponse`` into the same (text, tool_calls,
|
||||
@@ -7169,11 +7281,16 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
Unlike Chat Completions, text and tool calls are *always* separate
|
||||
items, so ``content`` is never null when text is present.
|
||||
"""
|
||||
status = response.get("status") if isinstance(response, dict) else getattr(response, "status", None)
|
||||
if status == "failed":
|
||||
error = response.get("error") if isinstance(response, dict) else getattr(response, "error", None)
|
||||
self._raise_responses_failure(error)
|
||||
self._record_finish_reason(response)
|
||||
response_text = ""
|
||||
tool_calls: List[Dict[str, Any]] = []
|
||||
reasoning_content = None
|
||||
|
||||
output_items = getattr(response, 'output', None) or []
|
||||
output_items = (response.get('output') if isinstance(response, dict) else getattr(response, 'output', None)) or []
|
||||
for item in output_items:
|
||||
# Handle both dict and object access
|
||||
if isinstance(item, dict):
|
||||
@@ -7231,6 +7348,34 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
return response_text, tool_calls if tool_calls else None, reasoning_content
|
||||
|
||||
@staticmethod
|
||||
def _close_responses_stream(stream) -> None:
|
||||
"""Close a stream or its LiteLLM transport without masking its error."""
|
||||
try:
|
||||
close = getattr(stream, "close", None)
|
||||
if not callable(close):
|
||||
close = getattr(getattr(stream, "response", None), "close", None)
|
||||
if callable(close):
|
||||
close()
|
||||
except Exception:
|
||||
logging.debug("Responses stream cleanup failed", exc_info=True)
|
||||
|
||||
@staticmethod
|
||||
async def _aclose_responses_stream(stream) -> None:
|
||||
"""Await transport cleanup on completion, failure or cancellation."""
|
||||
try:
|
||||
target = stream
|
||||
close = getattr(target, "aclose", None) or getattr(target, "close", None)
|
||||
if not callable(close):
|
||||
target = getattr(stream, "response", None)
|
||||
close = getattr(target, "aclose", None) or getattr(target, "close", None)
|
||||
if callable(close):
|
||||
result = close()
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
except Exception:
|
||||
logging.debug("Responses stream cleanup failed", exc_info=True)
|
||||
|
||||
def _stream_responses_api(
|
||||
self,
|
||||
params: Dict[str, Any],
|
||||
@@ -7262,95 +7407,112 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if _emit:
|
||||
from ..streaming.events import StreamEvent, StreamEventType
|
||||
|
||||
for event in litellm.responses(**params):
|
||||
# Access event type — handle both object and dict
|
||||
if isinstance(event, dict):
|
||||
evt_type = event.get("type", "")
|
||||
else:
|
||||
evt_type = getattr(event, "type", "")
|
||||
|
||||
# ── Completed — capture the final response for usage accounting ──
|
||||
if evt_type == "response.completed":
|
||||
_final_response = (event.get("response") if isinstance(event, dict)
|
||||
else getattr(event, "response", None))
|
||||
|
||||
# ── Text delta ──────────────────────────────────────────
|
||||
if evt_type == "response.output_text.delta":
|
||||
delta_text = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
response_text += delta_text
|
||||
if _emit and delta_text:
|
||||
stream_callback(StreamEvent(
|
||||
type=StreamEventType.DELTA_TEXT,
|
||||
timestamp=time.perf_counter(),
|
||||
content=delta_text,
|
||||
is_reasoning=False,
|
||||
))
|
||||
|
||||
# ── Reasoning delta ─────────────────────────────────────
|
||||
elif evt_type == "response.reasoning_summary_text.delta":
|
||||
delta_text = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
if _emit and delta_text:
|
||||
stream_callback(StreamEvent(
|
||||
type=StreamEventType.DELTA_TEXT,
|
||||
timestamp=time.perf_counter(),
|
||||
content=delta_text,
|
||||
is_reasoning=True,
|
||||
))
|
||||
|
||||
# ── Function call arguments delta ───────────────────────
|
||||
elif evt_type == "response.function_call_arguments.delta":
|
||||
idx = (event.get("output_index", 0) if isinstance(event, dict)
|
||||
else getattr(event, "output_index", 0))
|
||||
delta_args = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
if idx not in _pending_tools:
|
||||
item_id = (event.get("item_id", f"tool_{idx}") if isinstance(event, dict)
|
||||
else getattr(event, "item_id", f"tool_{idx}"))
|
||||
_pending_tools[idx] = {
|
||||
"id": item_id,
|
||||
"name": "",
|
||||
"arguments": "",
|
||||
}
|
||||
_pending_tools[idx]["arguments"] += delta_args
|
||||
|
||||
# ── Output item done — finalise tool call metadata ──────
|
||||
elif evt_type == "response.output_item.done":
|
||||
response_stream = litellm.responses(**params)
|
||||
try:
|
||||
for event in response_stream:
|
||||
# Access event type — handle both object and dict
|
||||
if isinstance(event, dict):
|
||||
item = event.get("item", {})
|
||||
idx = event.get("output_index", 0)
|
||||
evt_type = event.get("type", "")
|
||||
else:
|
||||
item = getattr(event, "item", None) or {}
|
||||
idx = getattr(event, "output_index", 0)
|
||||
evt_type = getattr(event, "type", "")
|
||||
|
||||
if isinstance(item, dict):
|
||||
i_type = item.get("type", "")
|
||||
else:
|
||||
i_type = getattr(item, "type", "")
|
||||
if evt_type == "response.failed":
|
||||
failed_response = event.get("response") if isinstance(event, dict) else getattr(event, "response", None)
|
||||
error = failed_response.get("error") if isinstance(failed_response, dict) else getattr(failed_response, "error", None)
|
||||
self._raise_responses_failure(error)
|
||||
if evt_type == "error":
|
||||
self._raise_responses_failure(event)
|
||||
|
||||
if i_type == "function_call":
|
||||
if isinstance(item, dict):
|
||||
call_id = item.get("call_id", item.get("id", f"tool_{idx}"))
|
||||
fn_name = item.get("name", "")
|
||||
fn_args = item.get("arguments", "")
|
||||
else:
|
||||
call_id = getattr(item, "call_id", None) or getattr(item, "id", f"tool_{idx}")
|
||||
fn_name = getattr(item, "name", "")
|
||||
fn_args = getattr(item, "arguments", "")
|
||||
# ── Completed — capture the final response for usage accounting ──
|
||||
if evt_type in ("response.completed", "response.incomplete"):
|
||||
_final_response = (event.get("response") if isinstance(event, dict)
|
||||
else getattr(event, "response", None))
|
||||
self._record_finish_reason(_final_response)
|
||||
|
||||
if idx in _pending_tools:
|
||||
_pending_tools[idx]["id"] = call_id
|
||||
_pending_tools[idx]["name"] = fn_name
|
||||
# Prefer the accumulated delta args if available, else use final
|
||||
if not _pending_tools[idx]["arguments"]:
|
||||
_pending_tools[idx]["arguments"] = fn_args
|
||||
else:
|
||||
if evt_type == "response.refusal.delta":
|
||||
refusal = event.get("delta", "") if isinstance(event, dict) else getattr(event, "delta", "")
|
||||
self._record_finish_reason({"choices": [{"message": {"refusal": refusal}}]})
|
||||
|
||||
# ── Text delta ──────────────────────────────────────────
|
||||
if evt_type == "response.output_text.delta":
|
||||
delta_text = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
response_text += delta_text
|
||||
if _emit and delta_text:
|
||||
stream_callback(StreamEvent(
|
||||
type=StreamEventType.DELTA_TEXT,
|
||||
timestamp=time.perf_counter(),
|
||||
content=delta_text,
|
||||
is_reasoning=False,
|
||||
))
|
||||
|
||||
# ── Reasoning delta ─────────────────────────────────────
|
||||
elif evt_type == "response.reasoning_summary_text.delta":
|
||||
delta_text = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
if _emit and delta_text:
|
||||
stream_callback(StreamEvent(
|
||||
type=StreamEventType.DELTA_TEXT,
|
||||
timestamp=time.perf_counter(),
|
||||
content=delta_text,
|
||||
is_reasoning=True,
|
||||
))
|
||||
|
||||
# ── Function call arguments delta ───────────────────────
|
||||
elif evt_type == "response.function_call_arguments.delta":
|
||||
idx = (event.get("output_index", 0) if isinstance(event, dict)
|
||||
else getattr(event, "output_index", 0))
|
||||
delta_args = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
if idx not in _pending_tools:
|
||||
item_id = (event.get("item_id", f"tool_{idx}") if isinstance(event, dict)
|
||||
else getattr(event, "item_id", f"tool_{idx}"))
|
||||
_pending_tools[idx] = {
|
||||
"id": call_id,
|
||||
"name": fn_name,
|
||||
"arguments": fn_args,
|
||||
"id": item_id,
|
||||
"name": "",
|
||||
"arguments": "",
|
||||
}
|
||||
_pending_tools[idx]["arguments"] += delta_args
|
||||
|
||||
# ── Output item done — finalise tool call metadata ──────
|
||||
elif evt_type == "response.output_item.done":
|
||||
if isinstance(event, dict):
|
||||
item = event.get("item", {})
|
||||
idx = event.get("output_index", 0)
|
||||
else:
|
||||
item = getattr(event, "item", None) or {}
|
||||
idx = getattr(event, "output_index", 0)
|
||||
|
||||
if isinstance(item, dict):
|
||||
i_type = item.get("type", "")
|
||||
else:
|
||||
i_type = getattr(item, "type", "")
|
||||
|
||||
if i_type == "function_call":
|
||||
if isinstance(item, dict):
|
||||
call_id = item.get("call_id", item.get("id", f"tool_{idx}"))
|
||||
fn_name = item.get("name", "")
|
||||
fn_args = item.get("arguments", "")
|
||||
else:
|
||||
call_id = getattr(item, "call_id", None) or getattr(item, "id", f"tool_{idx}")
|
||||
fn_name = getattr(item, "name", "")
|
||||
fn_args = getattr(item, "arguments", "")
|
||||
|
||||
if idx in _pending_tools:
|
||||
_pending_tools[idx]["id"] = call_id
|
||||
_pending_tools[idx]["name"] = fn_name
|
||||
# Prefer the accumulated delta args if available, else use final
|
||||
if not _pending_tools[idx]["arguments"]:
|
||||
_pending_tools[idx]["arguments"] = fn_args
|
||||
else:
|
||||
_pending_tools[idx] = {
|
||||
"id": call_id,
|
||||
"name": fn_name,
|
||||
"arguments": fn_args,
|
||||
}
|
||||
|
||||
finally:
|
||||
self._close_responses_stream(response_stream)
|
||||
|
||||
# ── Build tool_calls list in output_index order ─────────────
|
||||
tool_calls: List[Dict[str, Any]] = []
|
||||
@@ -7393,80 +7555,97 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if _emit:
|
||||
from ..streaming.events import StreamEvent, StreamEventType
|
||||
|
||||
async for event in await litellm.aresponses(**params):
|
||||
if isinstance(event, dict):
|
||||
evt_type = event.get("type", "")
|
||||
else:
|
||||
evt_type = getattr(event, "type", "")
|
||||
|
||||
if evt_type == "response.completed":
|
||||
_final_response = (event.get("response") if isinstance(event, dict)
|
||||
else getattr(event, "response", None))
|
||||
|
||||
if evt_type == "response.output_text.delta":
|
||||
delta_text = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
response_text += delta_text
|
||||
if _emit and delta_text:
|
||||
stream_callback(StreamEvent(
|
||||
type=StreamEventType.DELTA_TEXT,
|
||||
timestamp=time.perf_counter(),
|
||||
content=delta_text,
|
||||
is_reasoning=False,
|
||||
))
|
||||
|
||||
elif evt_type == "response.reasoning_summary_text.delta":
|
||||
delta_text = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
if _emit and delta_text:
|
||||
stream_callback(StreamEvent(
|
||||
type=StreamEventType.DELTA_TEXT,
|
||||
timestamp=time.perf_counter(),
|
||||
content=delta_text,
|
||||
is_reasoning=True,
|
||||
))
|
||||
|
||||
elif evt_type == "response.function_call_arguments.delta":
|
||||
idx = (event.get("output_index", 0) if isinstance(event, dict)
|
||||
else getattr(event, "output_index", 0))
|
||||
delta_args = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
if idx not in _pending_tools:
|
||||
item_id = (event.get("item_id", f"tool_{idx}") if isinstance(event, dict)
|
||||
else getattr(event, "item_id", f"tool_{idx}"))
|
||||
_pending_tools[idx] = {"id": item_id, "name": "", "arguments": ""}
|
||||
_pending_tools[idx]["arguments"] += delta_args
|
||||
|
||||
elif evt_type == "response.output_item.done":
|
||||
response_stream = await litellm.aresponses(**params)
|
||||
try:
|
||||
async for event in response_stream:
|
||||
if isinstance(event, dict):
|
||||
item = event.get("item", {})
|
||||
idx = event.get("output_index", 0)
|
||||
evt_type = event.get("type", "")
|
||||
else:
|
||||
item = getattr(event, "item", None) or {}
|
||||
idx = getattr(event, "output_index", 0)
|
||||
evt_type = getattr(event, "type", "")
|
||||
|
||||
if isinstance(item, dict):
|
||||
i_type = item.get("type", "")
|
||||
else:
|
||||
i_type = getattr(item, "type", "")
|
||||
if evt_type == "response.failed":
|
||||
failed_response = event.get("response") if isinstance(event, dict) else getattr(event, "response", None)
|
||||
error = failed_response.get("error") if isinstance(failed_response, dict) else getattr(failed_response, "error", None)
|
||||
self._raise_responses_failure(error)
|
||||
if evt_type == "error":
|
||||
self._raise_responses_failure(event)
|
||||
|
||||
if evt_type in ("response.completed", "response.incomplete"):
|
||||
_final_response = (event.get("response") if isinstance(event, dict)
|
||||
else getattr(event, "response", None))
|
||||
self._record_finish_reason(_final_response)
|
||||
|
||||
if evt_type == "response.refusal.delta":
|
||||
refusal = event.get("delta", "") if isinstance(event, dict) else getattr(event, "delta", "")
|
||||
self._record_finish_reason({"choices": [{"message": {"refusal": refusal}}]})
|
||||
|
||||
if evt_type == "response.output_text.delta":
|
||||
delta_text = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
response_text += delta_text
|
||||
if _emit and delta_text:
|
||||
stream_callback(StreamEvent(
|
||||
type=StreamEventType.DELTA_TEXT,
|
||||
timestamp=time.perf_counter(),
|
||||
content=delta_text,
|
||||
is_reasoning=False,
|
||||
))
|
||||
|
||||
elif evt_type == "response.reasoning_summary_text.delta":
|
||||
delta_text = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
if _emit and delta_text:
|
||||
stream_callback(StreamEvent(
|
||||
type=StreamEventType.DELTA_TEXT,
|
||||
timestamp=time.perf_counter(),
|
||||
content=delta_text,
|
||||
is_reasoning=True,
|
||||
))
|
||||
|
||||
elif evt_type == "response.function_call_arguments.delta":
|
||||
idx = (event.get("output_index", 0) if isinstance(event, dict)
|
||||
else getattr(event, "output_index", 0))
|
||||
delta_args = (event.get("delta", "") if isinstance(event, dict)
|
||||
else getattr(event, "delta", ""))
|
||||
if idx not in _pending_tools:
|
||||
item_id = (event.get("item_id", f"tool_{idx}") if isinstance(event, dict)
|
||||
else getattr(event, "item_id", f"tool_{idx}"))
|
||||
_pending_tools[idx] = {"id": item_id, "name": "", "arguments": ""}
|
||||
_pending_tools[idx]["arguments"] += delta_args
|
||||
|
||||
elif evt_type == "response.output_item.done":
|
||||
if isinstance(event, dict):
|
||||
item = event.get("item", {})
|
||||
idx = event.get("output_index", 0)
|
||||
else:
|
||||
item = getattr(event, "item", None) or {}
|
||||
idx = getattr(event, "output_index", 0)
|
||||
|
||||
if i_type == "function_call":
|
||||
if isinstance(item, dict):
|
||||
call_id = item.get("call_id", item.get("id", f"tool_{idx}"))
|
||||
fn_name = item.get("name", "")
|
||||
fn_args = item.get("arguments", "")
|
||||
i_type = item.get("type", "")
|
||||
else:
|
||||
call_id = getattr(item, "call_id", None) or getattr(item, "id", f"tool_{idx}")
|
||||
fn_name = getattr(item, "name", "")
|
||||
fn_args = getattr(item, "arguments", "")
|
||||
i_type = getattr(item, "type", "")
|
||||
|
||||
if idx in _pending_tools:
|
||||
_pending_tools[idx]["id"] = call_id
|
||||
_pending_tools[idx]["name"] = fn_name
|
||||
if not _pending_tools[idx]["arguments"]:
|
||||
_pending_tools[idx]["arguments"] = fn_args
|
||||
else:
|
||||
_pending_tools[idx] = {"id": call_id, "name": fn_name, "arguments": fn_args}
|
||||
if i_type == "function_call":
|
||||
if isinstance(item, dict):
|
||||
call_id = item.get("call_id", item.get("id", f"tool_{idx}"))
|
||||
fn_name = item.get("name", "")
|
||||
fn_args = item.get("arguments", "")
|
||||
else:
|
||||
call_id = getattr(item, "call_id", None) or getattr(item, "id", f"tool_{idx}")
|
||||
fn_name = getattr(item, "name", "")
|
||||
fn_args = getattr(item, "arguments", "")
|
||||
|
||||
if idx in _pending_tools:
|
||||
_pending_tools[idx]["id"] = call_id
|
||||
_pending_tools[idx]["name"] = fn_name
|
||||
if not _pending_tools[idx]["arguments"]:
|
||||
_pending_tools[idx]["arguments"] = fn_args
|
||||
else:
|
||||
_pending_tools[idx] = {"id": call_id, "name": fn_name, "arguments": fn_args}
|
||||
|
||||
finally:
|
||||
await self._aclose_responses_stream(response_stream)
|
||||
|
||||
tool_calls: List[Dict[str, Any]] = []
|
||||
for idx in sorted(_pending_tools.keys()):
|
||||
|
||||
@@ -898,16 +898,52 @@ class OpenAIClient:
|
||||
responses_content.append(part)
|
||||
return responses_content
|
||||
|
||||
@staticmethod
|
||||
def _responses_incomplete_finish_reason(response) -> Optional[str]:
|
||||
"""Map known incomplete reasons to existing Chat Completions outcomes."""
|
||||
status = response.get("status") if isinstance(response, dict) else getattr(response, "status", None)
|
||||
if status != "incomplete":
|
||||
return None
|
||||
details = response.get("incomplete_details") if isinstance(response, dict) else getattr(response, "incomplete_details", None)
|
||||
reason = details.get("reason") if isinstance(details, dict) else getattr(details, "reason", None)
|
||||
return {"max_output_tokens": "length", "content_filter": "content_filter"}.get(reason)
|
||||
|
||||
@staticmethod
|
||||
def _extract_responses_refusal(response) -> Optional[str]:
|
||||
"""Preserve refusal explanations separately from ordinary output text."""
|
||||
output = response.get("output", []) if isinstance(response, dict) else getattr(response, "output", [])
|
||||
refusals = []
|
||||
for item in output or []:
|
||||
item_type = item.get("type") if isinstance(item, dict) else getattr(item, "type", None)
|
||||
if item_type != "message":
|
||||
continue
|
||||
content = item.get("content", []) if isinstance(item, dict) else getattr(item, "content", [])
|
||||
for block in content or []:
|
||||
block_type = block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
|
||||
if block_type == "refusal":
|
||||
refusal = block.get("refusal") if isinstance(block, dict) else getattr(block, "refusal", None)
|
||||
if isinstance(refusal, str) and refusal:
|
||||
refusals.append(refusal)
|
||||
return "".join(refusals) or None
|
||||
|
||||
def _responses_to_chat_completion(self, response) -> ChatCompletion:
|
||||
"""
|
||||
Wrap a Responses API response into a ChatCompletion dataclass
|
||||
so that all downstream code (tool loop, display, callbacks)
|
||||
works without modification.
|
||||
"""
|
||||
status = response.get("status") if isinstance(response, dict) else getattr(response, "status", None)
|
||||
if status == "failed":
|
||||
from .llm import LLMResponseError
|
||||
error = response.get("error") if isinstance(response, dict) else getattr(response, "error", None)
|
||||
code = error.get("code") if isinstance(error, dict) else getattr(error, "code", None)
|
||||
message = error.get("message") if isinstance(error, dict) else getattr(error, "message", None)
|
||||
label = f"Responses API failed ({code})" if code else "Responses API failed"
|
||||
raise LLMResponseError(f"{label}: {message or 'Provider returned a failed response.'}")
|
||||
response_text = ""
|
||||
tool_calls_list: List[ToolCall] = []
|
||||
|
||||
output_items = getattr(response, 'output', None) or []
|
||||
output_items = (response.get('output') if isinstance(response, dict) else getattr(response, 'output', None)) or []
|
||||
for item in output_items:
|
||||
item_type = getattr(item, "type", "") if not isinstance(item, dict) else item.get("type", "")
|
||||
|
||||
@@ -942,22 +978,26 @@ class OpenAIClient:
|
||||
))
|
||||
|
||||
# Build usage
|
||||
raw_usage = getattr(response, 'usage', None)
|
||||
raw_usage = response.get('usage') if isinstance(response, dict) else getattr(response, 'usage', None)
|
||||
usage = None
|
||||
if raw_usage:
|
||||
def usage_value(name):
|
||||
return raw_usage.get(name, 0) if isinstance(raw_usage, dict) else getattr(raw_usage, name, 0)
|
||||
|
||||
usage = CompletionUsage(
|
||||
prompt_tokens=getattr(raw_usage, 'input_tokens', 0),
|
||||
completion_tokens=getattr(raw_usage, 'output_tokens', 0),
|
||||
total_tokens=getattr(raw_usage, 'total_tokens', 0),
|
||||
prompt_tokens=usage_value('input_tokens'),
|
||||
completion_tokens=usage_value('output_tokens'),
|
||||
total_tokens=usage_value('total_tokens'),
|
||||
)
|
||||
|
||||
message = ChatCompletionMessage(
|
||||
content=response_text if response_text else None,
|
||||
role="assistant",
|
||||
refusal=self._extract_responses_refusal(response),
|
||||
tool_calls=tool_calls_list if tool_calls_list else None,
|
||||
)
|
||||
choice = Choice(
|
||||
finish_reason="tool_calls" if tool_calls_list else "stop",
|
||||
finish_reason=self._responses_incomplete_finish_reason(response) or ("tool_calls" if tool_calls_list else "stop"),
|
||||
index=0,
|
||||
message=message,
|
||||
)
|
||||
@@ -1055,6 +1095,7 @@ class OpenAIClient:
|
||||
Returns:
|
||||
ChatCompletion object or None if error
|
||||
"""
|
||||
raw = None
|
||||
# Lazy import StreamEvent types only when needed
|
||||
_emit = emit_events and stream_callback is not None
|
||||
if _emit:
|
||||
@@ -1114,6 +1155,9 @@ class OpenAIClient:
|
||||
except (FileNotFoundError, ValueError):
|
||||
raise
|
||||
except Exception as e:
|
||||
status = raw.get("status") if isinstance(raw, dict) else getattr(raw, "status", None)
|
||||
if status == "failed":
|
||||
raise
|
||||
self.logger.warning(f"Responses API streaming failed, falling back: {e}")
|
||||
# Fall through to Chat Completions streaming
|
||||
|
||||
@@ -1284,6 +1328,9 @@ class OpenAIClient:
|
||||
return final_response
|
||||
|
||||
except Exception as e:
|
||||
status = raw.get("status") if isinstance(raw, dict) else getattr(raw, "status", None)
|
||||
if status == "failed":
|
||||
raise
|
||||
self.logger.error(f"Error in stream processing: {e}")
|
||||
return None
|
||||
|
||||
@@ -1320,6 +1367,7 @@ class OpenAIClient:
|
||||
Returns:
|
||||
ChatCompletion object or None if error
|
||||
"""
|
||||
raw = None
|
||||
# Lazy import StreamEvent types only when needed
|
||||
_emit = emit_events and stream_callback is not None
|
||||
if _emit:
|
||||
@@ -1388,6 +1436,9 @@ class OpenAIClient:
|
||||
except (FileNotFoundError, ValueError):
|
||||
raise
|
||||
except Exception as e:
|
||||
status = raw.get("status") if isinstance(raw, dict) else getattr(raw, "status", None)
|
||||
if status == "failed":
|
||||
raise
|
||||
self.logger.warning(f"Responses API async streaming failed, falling back: {e}")
|
||||
# Fall through to Chat Completions streaming
|
||||
|
||||
@@ -1550,6 +1601,9 @@ class OpenAIClient:
|
||||
return final_response
|
||||
|
||||
except Exception as e:
|
||||
status = raw.get("status") if isinstance(raw, dict) else getattr(raw, "status", None)
|
||||
if status == "failed":
|
||||
raise
|
||||
self.logger.error(f"Error in async stream processing: {e}")
|
||||
return None
|
||||
|
||||
@@ -1654,6 +1708,7 @@ class OpenAIClient:
|
||||
"""
|
||||
# ── Responses API path (non-streaming only) ────────────────────
|
||||
if not stream and self._use_responses_api(model):
|
||||
raw = None
|
||||
try:
|
||||
resp_params = self._build_responses_input(
|
||||
messages, model, temperature, tools,
|
||||
@@ -1664,6 +1719,9 @@ class OpenAIClient:
|
||||
except (FileNotFoundError, ValueError):
|
||||
raise
|
||||
except Exception as e:
|
||||
status = raw.get("status") if isinstance(raw, dict) else getattr(raw, "status", None)
|
||||
if status == "failed":
|
||||
raise
|
||||
self.logger.warning(f"Responses API failed, falling back to Chat Completions: {e}")
|
||||
# Fall through to Chat Completions
|
||||
|
||||
@@ -1718,6 +1776,7 @@ class OpenAIClient:
|
||||
"""
|
||||
# ── Responses API path (non-streaming only) ────────────────────
|
||||
if not stream and self._use_responses_api(model):
|
||||
raw = None
|
||||
try:
|
||||
resp_params = self._build_responses_input(
|
||||
messages, model, temperature, tools,
|
||||
@@ -1728,6 +1787,9 @@ class OpenAIClient:
|
||||
except (FileNotFoundError, ValueError):
|
||||
raise
|
||||
except Exception as e:
|
||||
status = raw.get("status") if isinstance(raw, dict) else getattr(raw, "status", None)
|
||||
if status == "failed":
|
||||
raise
|
||||
self.logger.warning(f"Responses API failed, falling back to Chat Completions: {e}")
|
||||
# Fall through to Chat Completions
|
||||
|
||||
|
||||
@@ -84,6 +84,7 @@ class RateLimiter:
|
||||
|
||||
def __post_init__(self):
|
||||
"""Initialize internal state."""
|
||||
self.max_retry_delay = max(0, self.max_retry_delay)
|
||||
# Request-based limiting
|
||||
if self.requests_per_minute is not None:
|
||||
self._tokens = float(self.burst)
|
||||
|
||||
@@ -53,6 +53,7 @@ class Cap(str, Enum):
|
||||
STREAMING_WITH_TOOLS = "streaming_with_tools"
|
||||
PARALLEL_TOOL_CALLS = "parallel_tool_calls"
|
||||
EMBEDDINGS_ENDPOINT = "embeddings_endpoint"
|
||||
SYSTEM_ONE = "system_one"
|
||||
|
||||
|
||||
class Evidence(str, Enum):
|
||||
@@ -105,6 +106,9 @@ _OLLAMA_CAP_MAP = {
|
||||
"thinking": (Cap.THINKING,),
|
||||
"embedding": (Cap.EMBEDDINGS_ENDPOINT,),
|
||||
"audio": (Cap.AUDIO_IN,),
|
||||
"decision": (Cap.SYSTEM_ONE,),
|
||||
"systemone": (Cap.SYSTEM_ONE,),
|
||||
"system_one": (Cap.SYSTEM_ONE,),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -17,6 +17,11 @@ import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from .._logging import get_logger
|
||||
from ..utils.atomic_io import update_json
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def get_default_auth_filepath() -> str:
|
||||
"""
|
||||
@@ -74,22 +79,56 @@ class MCPAuthStorage:
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _read(self) -> Dict[str, Any]:
|
||||
"""Read all entries from storage file."""
|
||||
"""Read all entries from storage file.
|
||||
|
||||
A corrupt file is moved aside to a unique ``.corrupt`` sidecar and
|
||||
logged rather than silently treated as "no credentials", so a torn
|
||||
write can never make every server's tokens vanish without a trace.
|
||||
|
||||
The quarantine happens *under the same FileLock the writers use*, and
|
||||
the file is re-read inside the lock first: if a concurrent writer
|
||||
repaired the store between the first decode error and acquiring the
|
||||
lock, we return the repaired data instead of quarantining valid
|
||||
credentials.
|
||||
"""
|
||||
try:
|
||||
with open(self.filepath, 'r') as f:
|
||||
with open(self.filepath, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except (FileNotFoundError, json.JSONDecodeError):
|
||||
except FileNotFoundError:
|
||||
return {}
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
|
||||
# Corrupt on the optimistic (lock-free) read. Take the writer lock and
|
||||
# re-read before quarantining so we never move a file another process
|
||||
# just repaired.
|
||||
from ..session.store import FileLock
|
||||
from ..utils.atomic_io import quarantine_corrupt_file
|
||||
|
||||
try:
|
||||
with FileLock(self.filepath, timeout=10):
|
||||
try:
|
||||
with open(self.filepath, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except FileNotFoundError:
|
||||
return {}
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
try:
|
||||
dest = quarantine_corrupt_file(self.filepath)
|
||||
logger.error(
|
||||
"MCP auth store was corrupt; moved aside to %s",
|
||||
dest,
|
||||
)
|
||||
except OSError:
|
||||
logger.error(
|
||||
"MCP auth store at %s is corrupt and could not be "
|
||||
"quarantined", self.filepath,
|
||||
)
|
||||
return {}
|
||||
except (OSError, IOError):
|
||||
# Could not acquire the lock; fall back to the previous behaviour
|
||||
# of reporting empty rather than blocking a read path.
|
||||
return {}
|
||||
|
||||
def _write(self, data: Dict[str, Any]) -> None:
|
||||
"""Write all entries to storage file with secure permissions."""
|
||||
# Write to file
|
||||
with open(self.filepath, 'w') as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
# Set file permissions to 0600 (owner read/write only)
|
||||
if os.name != 'nt': # Unix-like systems
|
||||
os.chmod(self.filepath, 0o600)
|
||||
|
||||
def get(self, mcp_name: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
@@ -143,20 +182,34 @@ class MCPAuthStorage:
|
||||
def _set(self, mcp_name: str, entry: Dict[str, Any], server_url: Optional[str] = None) -> None:
|
||||
"""
|
||||
Set an auth entry for an MCP server.
|
||||
|
||||
|
||||
Uses a locked atomic read-modify-write so concurrent OAuth flows (across
|
||||
threads or processes) no longer lose each other's entries, and a crash
|
||||
mid-write leaves the previous file intact. The temp file created by
|
||||
``atomic_write_json`` is mode 0600, so secrets are never world-readable,
|
||||
even for the moment between create and chmod.
|
||||
|
||||
Args:
|
||||
mcp_name: Name of the MCP server
|
||||
entry: Auth entry dict
|
||||
server_url: Optional server URL to track
|
||||
"""
|
||||
data = self._read()
|
||||
|
||||
# Update server_url if provided
|
||||
if server_url:
|
||||
entry["server_url"] = server_url
|
||||
|
||||
data[mcp_name] = entry
|
||||
self._write(data)
|
||||
|
||||
def mutate(data: Dict[str, Any]) -> None:
|
||||
data[mcp_name] = entry
|
||||
|
||||
update_json(self.filepath, mutate, indent=2)
|
||||
|
||||
# Tighten to 0600 after the atomic swap. mkstemp already created the temp
|
||||
# file as 0600 so there is no world-readable window; this keeps the mode
|
||||
# correct across platforms and pre-existing files.
|
||||
if os.name != 'nt':
|
||||
try:
|
||||
os.chmod(self.filepath, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def set_tokens(
|
||||
self,
|
||||
@@ -201,10 +254,15 @@ class MCPAuthStorage:
|
||||
Args:
|
||||
mcp_name: Name of the MCP server
|
||||
"""
|
||||
data = self._read()
|
||||
if mcp_name in data:
|
||||
del data[mcp_name]
|
||||
self._write(data)
|
||||
def mutate(data: Dict[str, Any]) -> None:
|
||||
data.pop(mcp_name, None)
|
||||
|
||||
update_json(self.filepath, mutate, indent=2)
|
||||
if os.name != 'nt':
|
||||
try:
|
||||
os.chmod(self.filepath, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def set_code_verifier(self, mcp_name: str, code_verifier: str) -> None:
|
||||
"""
|
||||
|
||||
@@ -6,6 +6,7 @@ No dependencies required - stores data as plain Python lists.
|
||||
"""
|
||||
|
||||
import threading
|
||||
from copy import deepcopy
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
@@ -15,6 +16,8 @@ class InMemoryAdapter:
|
||||
|
||||
Stores short-term and long-term memories as plain Python lists.
|
||||
Data is not persisted across process restarts.
|
||||
Metadata must support deepcopy; writes and reads own independent copies so
|
||||
caller edits cannot change stored records outside the adapter's lock.
|
||||
|
||||
Usage:
|
||||
adapter = InMemoryAdapter()
|
||||
@@ -48,7 +51,7 @@ class InMemoryAdapter:
|
||||
"id": str(self._next_id),
|
||||
"text": text,
|
||||
"type": "short",
|
||||
"metadata": metadata,
|
||||
"metadata": deepcopy(metadata),
|
||||
}
|
||||
self._next_id += 1
|
||||
self._data.append(entry)
|
||||
@@ -64,7 +67,7 @@ class InMemoryAdapter:
|
||||
for e in self._data
|
||||
if e["type"] == "short" and query.lower() in e["text"].lower()
|
||||
]
|
||||
return results[:limit]
|
||||
return deepcopy(results[:limit])
|
||||
|
||||
def store_long_term(
|
||||
self, text: str, metadata: Optional[Dict[str, Any]] = None, **kwargs
|
||||
@@ -74,7 +77,7 @@ class InMemoryAdapter:
|
||||
"id": str(self._next_id),
|
||||
"text": text,
|
||||
"type": "long",
|
||||
"metadata": metadata,
|
||||
"metadata": deepcopy(metadata),
|
||||
}
|
||||
self._next_id += 1
|
||||
self._data.append(entry)
|
||||
@@ -90,7 +93,7 @@ class InMemoryAdapter:
|
||||
for e in self._data
|
||||
if e["type"] == "long" and query.lower() in e["text"].lower()
|
||||
]
|
||||
return results[:limit]
|
||||
return deepcopy(results[:limit])
|
||||
|
||||
def delete_memory(self, memory_id: str, tier: Optional[str] = None, **kwargs) -> bool:
|
||||
"""Delete a memory by ID. Returns True if an entry was removed.
|
||||
@@ -123,5 +126,8 @@ class InMemoryAdapter:
|
||||
|
||||
def get_all_memories(self, **kwargs) -> List[Dict[str, Any]]:
|
||||
with self._lock:
|
||||
# Return defensive copy to prevent external mutation of internal state
|
||||
return [dict(entry) for entry in self._data]
|
||||
# Entries are owned on insertion and never mutated afterward.
|
||||
# Keep their membership snapshot under lock, then copy nested data
|
||||
# without blocking stores, deletes, or eviction.
|
||||
snapshot = list(self._data)
|
||||
return deepcopy(snapshot)
|
||||
@@ -19,6 +19,16 @@ from ..protocols import MemoryProtocol
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _metadata_user_id(raw):
|
||||
"""Match the existing Python JSON reader, including non-finite values."""
|
||||
try:
|
||||
metadata = json.loads(raw) if raw else {}
|
||||
return str(metadata["user_id"]) if isinstance(metadata, dict) and metadata.get("user_id") is not None else None
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
class SqliteMemoryAdapter:
|
||||
"""
|
||||
SQLite-based memory adapter implementing MemoryProtocol.
|
||||
@@ -79,6 +89,7 @@ class SqliteMemoryAdapter:
|
||||
check_same_thread=False,
|
||||
timeout=10.0
|
||||
)
|
||||
self._local.stm_conn.create_function("memory_user_id", 1, _metadata_user_id, deterministic=True)
|
||||
self._local.stm_conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS short_term_memory (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
@@ -103,6 +114,7 @@ class SqliteMemoryAdapter:
|
||||
check_same_thread=False,
|
||||
timeout=10.0
|
||||
)
|
||||
self._local.ltm_conn.create_function("memory_user_id", 1, _metadata_user_id, deterministic=True)
|
||||
self._local.ltm_conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS long_term_memory (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
@@ -143,10 +155,13 @@ class SqliteMemoryAdapter:
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Search short-term memory."""
|
||||
conn = self._get_stm_conn()
|
||||
user_id = kwargs.get("user_id")
|
||||
user_filter = " AND memory_user_id(metadata) = ?" if user_id is not None else ""
|
||||
params = (f"%{query}%", str(user_id), limit) if user_id is not None else (f"%{query}%", limit)
|
||||
cursor = conn.execute(
|
||||
"SELECT id, content, metadata, timestamp FROM short_term_memory "
|
||||
"WHERE content LIKE ? ORDER BY timestamp DESC LIMIT ?",
|
||||
(f"%{query}%", limit)
|
||||
"WHERE content LIKE ?" + user_filter + " ORDER BY timestamp DESC LIMIT ?",
|
||||
params
|
||||
)
|
||||
|
||||
results = []
|
||||
@@ -184,10 +199,13 @@ class SqliteMemoryAdapter:
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Search long-term memory."""
|
||||
conn = self._get_ltm_conn()
|
||||
user_id = kwargs.get("user_id")
|
||||
user_filter = " AND memory_user_id(metadata) = ?" if user_id is not None else ""
|
||||
params = (f"%{query}%", str(user_id), limit) if user_id is not None else (f"%{query}%", limit)
|
||||
cursor = conn.execute(
|
||||
"SELECT id, content, metadata, timestamp FROM long_term_memory "
|
||||
"WHERE content LIKE ? ORDER BY timestamp DESC LIMIT ?",
|
||||
(f"%{query}%", limit)
|
||||
"WHERE content LIKE ?" + user_filter + " ORDER BY timestamp DESC LIMIT ?",
|
||||
params
|
||||
)
|
||||
|
||||
results = []
|
||||
@@ -252,8 +270,10 @@ class SqliteMemoryAdapter:
|
||||
|
||||
def get_all_memories(self, **kwargs) -> List[Dict[str, Any]]:
|
||||
"""Get all memories from both short-term and long-term."""
|
||||
short_memories = self.search_short_term("", limit=1000)
|
||||
long_memories = self.search_long_term("", limit=1000)
|
||||
# SQLite's negative LIMIT means no upper bound. Enumeration must return
|
||||
# every row; ordinary searches retain their default/explicit limits.
|
||||
short_memories = self.search_short_term("", limit=-1)
|
||||
long_memories = self.search_long_term("", limit=-1)
|
||||
|
||||
# Mark memory types
|
||||
for memory in short_memories:
|
||||
@@ -298,4 +318,4 @@ class SqliteMemoryAdapter:
|
||||
continue
|
||||
with self._connection_lock:
|
||||
self._all_connections.discard(conn)
|
||||
setattr(self._local, attr, None)
|
||||
setattr(self._local, attr, None)
|
||||
@@ -237,8 +237,8 @@ class BaseStore(ABC):
|
||||
def clear(self) -> int:
|
||||
"""Clear all entries."""
|
||||
count = len(self._entries)
|
||||
self._store._reset({})
|
||||
self._entries = {}
|
||||
self._save()
|
||||
self._was_updated = True
|
||||
return count
|
||||
|
||||
|
||||
@@ -1361,3 +1361,15 @@ def get_plugin_manager() -> PluginManager:
|
||||
if _default_manager is None:
|
||||
_default_manager = PluginManager()
|
||||
return _default_manager
|
||||
|
||||
|
||||
def has_registered_plugins() -> bool:
|
||||
"""Return whether the global manager has any plugins registered.
|
||||
|
||||
Cheap, allocation-free fast path for the hot Agent-init merge calls: when
|
||||
the singleton has not been created or holds no plugins (the common case),
|
||||
callers can skip constructing/scanning the manager entirely. Does NOT
|
||||
create the singleton as a side effect.
|
||||
"""
|
||||
manager = _default_manager
|
||||
return manager is not None and bool(manager._plugins)
|
||||
@@ -333,7 +333,9 @@ class Session:
|
||||
histories = (getattr(data, "metadata", None) or {}).get(AGENT_HISTORY_KEY)
|
||||
if not isinstance(histories, dict):
|
||||
return {}
|
||||
return {k: v for k, v in histories.items() if isinstance(v, list) and v}
|
||||
# An empty list is an explicit clear, not a missing transcript. Keep
|
||||
# it so current records override legacy and Memory fallback history.
|
||||
return {k: v for k, v in histories.items() if isinstance(v, list)}
|
||||
|
||||
def _merged_agent_histories(self, session_store) -> Dict[str, List[Dict[str, Any]]]:
|
||||
"""All known sub-agent transcripts: legacy records under current ones.
|
||||
@@ -537,7 +539,7 @@ class Session:
|
||||
# read and write; if the persisted map matches what we wrote, we win.
|
||||
persisted = self._stored_agent_histories(session_store)
|
||||
if persisted.get(agent_key) == messages and all(
|
||||
persisted.get(k) == v for k, v in current.items()
|
||||
persisted.get(k) == v for k, v in current.items() if k != agent_key
|
||||
):
|
||||
return True
|
||||
logger.warning(
|
||||
@@ -582,12 +584,7 @@ class Session:
|
||||
for msg in chat_history
|
||||
if isinstance(msg, dict)
|
||||
]
|
||||
if not messages:
|
||||
logging.debug(
|
||||
f"No chat history to persist for agent {agent_key} "
|
||||
f"in session {self.session_id}"
|
||||
)
|
||||
elif not self._store_agent_history(
|
||||
if not self._store_agent_history(
|
||||
session_store, agent_key, messages
|
||||
):
|
||||
logging.warning(
|
||||
|
||||
@@ -59,17 +59,22 @@ def _progress_section(checkpoint: Optional[Dict[str, Any]]) -> List[str]:
|
||||
if not isinstance(checkpoint, dict):
|
||||
return []
|
||||
completed = checkpoint.get("completed_steps")
|
||||
total = checkpoint.get("total_steps") or checkpoint.get("step_count")
|
||||
last = checkpoint.get("last_step") or checkpoint.get("current_step")
|
||||
if completed is None and last is None:
|
||||
return []
|
||||
if total:
|
||||
head = f"PROGRESS: {completed or 0} of {total} workflow steps done"
|
||||
else:
|
||||
head = f"PROGRESS: {completed or 0} workflow steps done"
|
||||
if last:
|
||||
head += f"; last completed: {last}"
|
||||
return [head]
|
||||
total = checkpoint.get("total_steps")
|
||||
if total is None:
|
||||
total = checkpoint.get("step_count")
|
||||
last = checkpoint.get("last_step")
|
||||
current = checkpoint.get("current_step")
|
||||
parts: List[str] = []
|
||||
if completed is not None:
|
||||
if total is not None:
|
||||
parts.append(f"{completed} of {total} workflow steps done")
|
||||
else:
|
||||
parts.append(f"{completed} workflow steps done")
|
||||
if last is not None and last != "":
|
||||
parts.append(f"last completed: {last}")
|
||||
if current is not None and current != "":
|
||||
parts.append(f"current step: {current}")
|
||||
return ["PROGRESS: " + "; ".join(parts)] if parts else []
|
||||
|
||||
|
||||
def _recent_actions(recap_text: str) -> List[str]:
|
||||
|
||||
@@ -6,6 +6,7 @@ forking, and revert capabilities.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from praisonaiagents._logging import get_logger
|
||||
@@ -13,17 +14,27 @@ import os
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import dataclass, field, fields
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .store import SessionData, SessionMessage, DefaultSessionStore, FileLock
|
||||
from .store import (
|
||||
SessionData,
|
||||
SessionMessage,
|
||||
DefaultSessionStore,
|
||||
FileLock,
|
||||
_LRUSessionCache,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@dataclass
|
||||
class SessionSnapshot:
|
||||
"""A snapshot of session state at a point in time."""
|
||||
"""A transcript snapshot; old records may contain only a message index.
|
||||
|
||||
Durable records pool repeated transcript entries; portable exports inline them.
|
||||
Legacy records cannot recover history already discarded before upgrading.
|
||||
"""
|
||||
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
session_id: str = ""
|
||||
@@ -31,9 +42,11 @@ class SessionSnapshot:
|
||||
created_at: float = field(default_factory=time.time)
|
||||
label: Optional[str] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
transcript: Optional[Dict[str, Any]] = None
|
||||
invalidated: bool = False
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
def _header_dict(self) -> Dict[str, Any]:
|
||||
data = {
|
||||
"id": self.id,
|
||||
"session_id": self.session_id,
|
||||
"message_index": self.message_index,
|
||||
@@ -41,6 +54,15 @@ class SessionSnapshot:
|
||||
"label": self.label,
|
||||
"metadata": self.metadata,
|
||||
}
|
||||
if self.invalidated:
|
||||
data["invalidated"] = True
|
||||
return data
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
data = self._header_dict()
|
||||
if self.transcript is not None:
|
||||
data["transcript"] = copy.deepcopy(self.transcript)
|
||||
return data
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "SessionSnapshot":
|
||||
@@ -51,6 +73,8 @@ class SessionSnapshot:
|
||||
created_at=data.get("created_at", time.time()),
|
||||
label=data.get("label"),
|
||||
metadata=data.get("metadata", {}),
|
||||
transcript=copy.deepcopy(data.get("transcript")),
|
||||
invalidated=data.get("invalidated", False),
|
||||
)
|
||||
|
||||
@dataclass
|
||||
@@ -75,30 +99,130 @@ class ExtendedSessionData(SessionData):
|
||||
"title": self.title,
|
||||
})
|
||||
return d
|
||||
|
||||
def _to_storage_dict(self) -> Dict[str, Any]:
|
||||
"""Pool snapshot records in the atomic session file, not portable exports."""
|
||||
data = super().to_dict()
|
||||
data.update({
|
||||
"parent_id": self.parent_id,
|
||||
"forked_from_message_id": self.forked_from_message_id,
|
||||
"children_ids": self.children_ids,
|
||||
"is_shared": self.is_shared,
|
||||
"title": self.title,
|
||||
})
|
||||
records = {}
|
||||
record_ids = {}
|
||||
|
||||
def intern(value):
|
||||
# Most transcript records contain text plus scalar fields and empty
|
||||
# metadata. Reuse hashes of their immutable strings rather than
|
||||
# re-encoding every archived text for every snapshot on each write.
|
||||
if isinstance(value, dict) and all(
|
||||
item is None or isinstance(item, (str, bool, int, float))
|
||||
or isinstance(item, (dict, list)) and not item
|
||||
for item in value.values()
|
||||
):
|
||||
identity = ("flat", tuple(
|
||||
(name, type(item), None if isinstance(item, (dict, list)) else item)
|
||||
for name, item in value.items()
|
||||
))
|
||||
else:
|
||||
identity = ("json", json.dumps(value, ensure_ascii=False,
|
||||
separators=(",", ":")))
|
||||
key = record_ids.get(identity)
|
||||
if key is None:
|
||||
key = str(len(records))
|
||||
record_ids[identity] = key
|
||||
records[key] = copy.deepcopy(value)
|
||||
return key
|
||||
|
||||
snapshots = []
|
||||
for snapshot in self.snapshots:
|
||||
saved = snapshot._header_dict()
|
||||
required = {"messages", "archived_messages", "last_compaction"}
|
||||
if (snapshot.transcript is not None and required <= snapshot.transcript.keys()
|
||||
and all(isinstance(snapshot.transcript[name], list)
|
||||
and all(isinstance(item, dict) for item in snapshot.transcript[name])
|
||||
for name in ("messages", "archived_messages"))
|
||||
and (snapshot.transcript["last_compaction"] is None
|
||||
or isinstance(snapshot.transcript["last_compaction"], dict))):
|
||||
refs = {}
|
||||
for name, value in snapshot.transcript.items():
|
||||
if name in ("messages", "archived_messages") and isinstance(value, list):
|
||||
refs[name] = {"items": [intern(item) for item in value]}
|
||||
else:
|
||||
refs[name] = {"value": intern(value)}
|
||||
saved["transcript_refs"] = refs
|
||||
# Readers that only understand positional/inline snapshots
|
||||
# must reject this record instead of restoring a stale index.
|
||||
saved["invalidated"] = True
|
||||
saved["transcript_invalidated"] = snapshot.invalidated
|
||||
elif snapshot.transcript is not None:
|
||||
# Preserve legacy inline shapes rather than encoding an
|
||||
# incomplete transcript as a valid pooled snapshot.
|
||||
saved["transcript"] = copy.deepcopy(snapshot.transcript)
|
||||
snapshots.append(saved)
|
||||
data["snapshots"] = snapshots
|
||||
if any("transcript_refs" in snapshot for snapshot in snapshots):
|
||||
data["snapshot_storage"] = {"version": 1, "records": records}
|
||||
return data
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "ExtendedSessionData":
|
||||
messages = [
|
||||
SessionMessage.from_dict(m)
|
||||
for m in data.get("messages", [])
|
||||
]
|
||||
archived = [
|
||||
SessionMessage.from_dict(m)
|
||||
for m in (data.get("archived_messages") or [])
|
||||
]
|
||||
snapshots = [
|
||||
SessionSnapshot.from_dict(s)
|
||||
for s in data.get("snapshots", [])
|
||||
]
|
||||
base = SessionData.from_dict(data)
|
||||
snapshots = []
|
||||
for raw in data.get("snapshots", []):
|
||||
if "transcript_refs" in raw:
|
||||
storage = data.get("snapshot_storage")
|
||||
if (not isinstance(storage, dict) or storage.get("version") != 1
|
||||
or not isinstance(storage.get("records"), dict)
|
||||
or not isinstance(raw["transcript_refs"], dict)
|
||||
or "transcript" in raw):
|
||||
raise ValueError("invalid snapshot storage")
|
||||
records = storage["records"]
|
||||
if not {"messages", "archived_messages", "last_compaction"} <= raw["transcript_refs"].keys():
|
||||
raise ValueError("incomplete snapshot transcript references")
|
||||
transcript = {}
|
||||
|
||||
def clone_record(value):
|
||||
# JSON object keys/scalars are immutable. Copy mutable
|
||||
# fields without traversing every archived text again.
|
||||
if isinstance(value, dict):
|
||||
return {key: copy.deepcopy(item) if isinstance(item, (dict, list)) else item
|
||||
for key, item in value.items()}
|
||||
return copy.deepcopy(value)
|
||||
|
||||
try:
|
||||
for name, ref in raw["transcript_refs"].items():
|
||||
if not isinstance(ref, dict):
|
||||
raise ValueError("invalid snapshot reference")
|
||||
if name in ("messages", "archived_messages") and (
|
||||
set(ref) != {"items"} or not isinstance(ref["items"], list)):
|
||||
raise ValueError("invalid snapshot message references")
|
||||
if name == "last_compaction" and set(ref) != {"value"}:
|
||||
raise ValueError("invalid snapshot checkpoint reference")
|
||||
if set(ref) == {"items"} and isinstance(ref["items"], list):
|
||||
transcript[name] = [clone_record(records[key]) for key in ref["items"]]
|
||||
elif set(ref) == {"value"}:
|
||||
transcript[name] = clone_record(records[ref["value"]])
|
||||
else:
|
||||
raise ValueError("invalid snapshot reference")
|
||||
if name in ("messages", "archived_messages") and not all(
|
||||
isinstance(item, dict) for item in transcript[name]):
|
||||
raise ValueError("invalid snapshot message record")
|
||||
if name == "last_compaction" and transcript[name] is not None and not isinstance(transcript[name], dict):
|
||||
raise ValueError("invalid snapshot checkpoint record")
|
||||
except (KeyError, TypeError) as exc:
|
||||
raise ValueError("missing snapshot record") from exc
|
||||
snapshot = SessionSnapshot.from_dict({
|
||||
**raw, "invalidated": raw.get("transcript_invalidated", False),
|
||||
})
|
||||
snapshot.transcript = transcript
|
||||
else:
|
||||
snapshot = SessionSnapshot.from_dict(raw)
|
||||
snapshots.append(snapshot)
|
||||
return cls(
|
||||
session_id=data.get("session_id", ""),
|
||||
messages=messages,
|
||||
created_at=data.get("created_at", datetime.now(timezone.utc).isoformat()),
|
||||
updated_at=data.get("updated_at", datetime.now(timezone.utc).isoformat()),
|
||||
agent_name=data.get("agent_name"),
|
||||
user_id=data.get("user_id"),
|
||||
metadata=data.get("metadata", {}),
|
||||
archived_messages=archived,
|
||||
**{descriptor.name: getattr(base, descriptor.name) for descriptor in fields(SessionData)},
|
||||
parent_id=data.get("parent_id"),
|
||||
forked_from_message_id=data.get("forked_from_message_id"),
|
||||
children_ids=data.get("children_ids", []),
|
||||
@@ -111,13 +235,7 @@ class ExtendedSessionData(SessionData):
|
||||
def from_session_data(cls, session: SessionData) -> "ExtendedSessionData":
|
||||
"""Convert a basic SessionData to ExtendedSessionData."""
|
||||
return cls(
|
||||
session_id=session.session_id,
|
||||
messages=session.messages,
|
||||
created_at=session.created_at,
|
||||
updated_at=session.updated_at,
|
||||
agent_name=session.agent_name,
|
||||
user_id=session.user_id,
|
||||
metadata=session.metadata,
|
||||
**{descriptor.name: getattr(session, descriptor.name) for descriptor in fields(SessionData)},
|
||||
)
|
||||
|
||||
class HierarchicalSessionStore(DefaultSessionStore):
|
||||
@@ -148,8 +266,21 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._extended_cache: Dict[str, ExtendedSessionData] = {}
|
||||
self._cache_fingerprints: Dict[str, Tuple[ExtendedSessionData, bytes]] = {}
|
||||
# Bound the extended-session cache with the same LRU policy (and
|
||||
# maxsize) as the base ``_cache`` so a long-lived process using this
|
||||
# subclass does not retain every session it has ever loaded in RAM.
|
||||
# Disk stays the source of truth, so an eviction just reloads on next
|
||||
# access. When an entry is evicted its companion mtime is dropped too so
|
||||
# ``_cache_mtimes`` stays bounded alongside the cache.
|
||||
self._cache_mtimes: Dict[str, float] = {} # Track file modification times
|
||||
self._extended_cache: Dict[str, ExtendedSessionData] = _LRUSessionCache(
|
||||
getattr(self._cache, "maxsize", 0),
|
||||
on_evict=lambda key: (
|
||||
self._cache_mtimes.pop(key, None),
|
||||
self._cache_fingerprints.pop(key, None),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _load_session_from_disk(self, session_id: str, filepath: str) -> ExtendedSessionData:
|
||||
@@ -171,9 +302,15 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
if not os.path.exists(filepath):
|
||||
return ExtendedSessionData(session_id=session_id)
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return ExtendedSessionData.from_dict(data)
|
||||
with open(filepath, "rb") as f:
|
||||
raw = f.read()
|
||||
data = json.loads(raw.decode("utf-8"))
|
||||
session = ExtendedSessionData.from_dict(data)
|
||||
# Bind the fingerprint to this parsed object and these exact bytes.
|
||||
# A later writer must not certify an older payload via its stats.
|
||||
with self._lock:
|
||||
self._cache_fingerprints[session_id] = (session, hashlib.sha256(raw).digest())
|
||||
return session
|
||||
except (json.JSONDecodeError, UnicodeDecodeError) as e:
|
||||
# Invalid UTF-8 (``UnicodeDecodeError``) is handled alongside
|
||||
# malformed JSON so a corrupt binary file is quarantined here too
|
||||
@@ -194,6 +331,21 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
f"refusing to overwrite existing data: {e}"
|
||||
)
|
||||
raise
|
||||
except ValueError as e:
|
||||
# Unsupported/damaged references must not turn into a fresh record
|
||||
# or a legacy positional snapshot that restores unrelated history.
|
||||
raise OSError("invalid snapshot storage; preserving session") from e
|
||||
|
||||
def _session_to_storage(self, session: SessionData) -> Dict[str, Any]:
|
||||
if isinstance(session, ExtendedSessionData):
|
||||
return session._to_storage_dict()
|
||||
return super()._session_to_storage(session)
|
||||
|
||||
def _dump_session_json(self, data: Any, stream) -> None:
|
||||
if isinstance(data, dict) and "snapshot_storage" in data:
|
||||
json.dump(data, stream, ensure_ascii=False, separators=(",", ":"))
|
||||
else:
|
||||
super()._dump_session_json(data, stream)
|
||||
|
||||
def _modify_session_locked(
|
||||
self,
|
||||
@@ -201,33 +353,81 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
mutator,
|
||||
*,
|
||||
error_label: str = "modify session",
|
||||
apply_retention: bool = True,
|
||||
) -> bool:
|
||||
"""Locked read-modify-write preserving extended session fields."""
|
||||
result = super()._modify_session_locked(
|
||||
session_id, mutator, error_label=error_label
|
||||
)
|
||||
if result:
|
||||
with self._lock:
|
||||
cached = self._cache.get(session_id)
|
||||
if isinstance(cached, ExtendedSessionData):
|
||||
self._extended_cache[session_id] = cached
|
||||
return result
|
||||
|
||||
def _is_cache_valid(self, session_id: str) -> bool:
|
||||
"""Check if cached session is still valid based on file mtime."""
|
||||
if session_id not in self._extended_cache:
|
||||
return False
|
||||
|
||||
filepath = self._get_session_path(session_id)
|
||||
if not os.path.exists(filepath):
|
||||
return False
|
||||
|
||||
with FileLock(filepath, self.lock_timeout):
|
||||
try:
|
||||
session = self._load_session_from_disk(session_id, filepath)
|
||||
except OSError:
|
||||
logger.error("Failed to %s %s: could not read existing session", error_label, session_id)
|
||||
return False
|
||||
before = session.messages[:]
|
||||
mutator(session)
|
||||
self._invalidate_legacy_snapshots(session, before)
|
||||
session.updated_at = datetime.now(timezone.utc).isoformat()
|
||||
if apply_retention:
|
||||
self._enforce_window(session)
|
||||
if not self._atomic_write_json(filepath, self._session_to_storage(session)):
|
||||
logger.error("Failed to %s %s", error_label, session_id)
|
||||
return False
|
||||
self._cache_written_session(session, filepath)
|
||||
return True
|
||||
|
||||
def _enforce_window(self, session: SessionData) -> None:
|
||||
before = session.messages[:]
|
||||
super()._enforce_window(session)
|
||||
self._invalidate_legacy_snapshots(session, before)
|
||||
|
||||
@staticmethod
|
||||
def _invalidate_legacy_snapshots(session: SessionData, before: List[SessionMessage]) -> None:
|
||||
if isinstance(session, ExtendedSessionData):
|
||||
for snapshot in session.snapshots:
|
||||
index = snapshot.message_index
|
||||
if snapshot.transcript is None and index >= 0 and (
|
||||
index >= len(session.messages)
|
||||
or before[:index + 1] != session.messages[:index + 1]
|
||||
):
|
||||
snapshot.invalidated = True
|
||||
|
||||
def _cache_written_session(self, session: ExtendedSessionData, filepath: str) -> None:
|
||||
"""Bind the written object to exact bytes while the caller holds FileLock."""
|
||||
try:
|
||||
current_mtime = os.path.getmtime(filepath)
|
||||
cached_mtime = self._cache_mtimes.get(session_id, 0)
|
||||
return current_mtime <= cached_mtime
|
||||
except (OSError, IOError):
|
||||
with open(filepath, "rb") as f:
|
||||
fingerprint = hashlib.sha256(f.read()).digest()
|
||||
except OSError:
|
||||
# The write succeeded; unavailable validation bytes only disable reuse.
|
||||
fingerprint = None
|
||||
with self._lock:
|
||||
self._cache[session.session_id] = session
|
||||
self._extended_cache[session.session_id] = session
|
||||
if fingerprint is None:
|
||||
self._cache_fingerprints.pop(session.session_id, None)
|
||||
else:
|
||||
self._cache_fingerprints[session.session_id] = (session, fingerprint)
|
||||
|
||||
def _is_cache_valid(self, session_id: str) -> Optional[bool]:
|
||||
"""Validate cached bytes; None means a transient validation failure."""
|
||||
with self._lock:
|
||||
if session_id not in self._extended_cache:
|
||||
return False
|
||||
filepath = self._get_session_path(session_id)
|
||||
try:
|
||||
with FileLock(filepath, self.lock_timeout):
|
||||
with open(filepath, "rb") as f:
|
||||
fingerprint = hashlib.sha256(f.read()).digest()
|
||||
with self._lock:
|
||||
cached = self._cache_fingerprints.get(session_id)
|
||||
return (
|
||||
cached is not None
|
||||
and cached[0] is self._extended_cache.get(session_id)
|
||||
and cached[1] == fingerprint
|
||||
)
|
||||
except FileNotFoundError:
|
||||
return False
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
def _read_session_fresh(self, session_id: str) -> ExtendedSessionData:
|
||||
"""Reload from disk and keep _cache and _extended_cache in sync."""
|
||||
@@ -237,16 +437,8 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
with self._lock:
|
||||
self._cache[session_id] = session
|
||||
|
||||
# Update cache with fresh file mtime
|
||||
filepath = self._get_session_path(session_id)
|
||||
try:
|
||||
mtime = os.path.getmtime(filepath) if os.path.exists(filepath) else time.time()
|
||||
except (OSError, IOError):
|
||||
mtime = time.time()
|
||||
|
||||
with self._lock:
|
||||
self._extended_cache[session_id] = session
|
||||
self._cache_mtimes[session_id] = mtime
|
||||
|
||||
return session
|
||||
|
||||
@@ -279,31 +471,38 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
|
||||
def _apply(session: SessionData) -> None:
|
||||
session.messages.append(message)
|
||||
if len(session.messages) > self.max_messages:
|
||||
session.messages = session.messages[-self.max_messages :]
|
||||
|
||||
return self._modify_session_locked(
|
||||
session_id, _apply, error_label="add message to session"
|
||||
)
|
||||
|
||||
def _load_extended_session(self, session_id: str, force_reload: bool = False) -> ExtendedSessionData:
|
||||
"""Load extended session with smart caching based on file modification time."""
|
||||
"""Reuse parsed extended data only when its source bytes are unchanged."""
|
||||
# Force reload bypasses cache validation
|
||||
if force_reload or not self._is_cache_valid(session_id):
|
||||
if force_reload:
|
||||
return self._read_session_fresh(session_id)
|
||||
valid = self._is_cache_valid(session_id)
|
||||
if valid is False:
|
||||
return self._read_session_fresh(session_id)
|
||||
|
||||
# Cache is valid, return cached version
|
||||
# Invalidation may have raced with the content check.
|
||||
with self._lock:
|
||||
return self._extended_cache[session_id]
|
||||
cached = self._extended_cache.get(session_id)
|
||||
if cached is not None:
|
||||
return cached
|
||||
return self._read_session_fresh(session_id)
|
||||
|
||||
def _save_extended_session(self, session: ExtendedSessionData) -> bool:
|
||||
def _save_extended_session(
|
||||
self, session: ExtendedSessionData, *, apply_retention: bool = True
|
||||
) -> bool:
|
||||
"""Save extended session to disk."""
|
||||
filepath = self._get_session_path(session.session_id)
|
||||
session.updated_at = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
# Trim messages if over limit
|
||||
if len(session.messages) > self.max_messages:
|
||||
session.messages = session.messages[-self.max_messages:]
|
||||
# Match the inherited retention policy, including non-destructive
|
||||
# compaction and tool-exchange-aware truncation.
|
||||
if apply_retention:
|
||||
self._enforce_window(session)
|
||||
|
||||
with FileLock(filepath, self.lock_timeout):
|
||||
try:
|
||||
@@ -317,20 +516,12 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
delete=False,
|
||||
suffix=".tmp"
|
||||
) as f:
|
||||
json.dump(session.to_dict(), f, indent=2, ensure_ascii=False)
|
||||
self._dump_session_json(self._session_to_storage(session), f)
|
||||
temp_path = f.name
|
||||
|
||||
os.replace(temp_path, filepath)
|
||||
|
||||
# Update cache with current file mtime after successful write
|
||||
try:
|
||||
mtime = os.path.getmtime(filepath)
|
||||
except (OSError, IOError):
|
||||
mtime = time.time()
|
||||
|
||||
with self._lock:
|
||||
self._extended_cache[session.session_id] = session
|
||||
self._cache_mtimes[session.session_id] = mtime
|
||||
self._cache_written_session(session, filepath)
|
||||
|
||||
return True
|
||||
except (IOError, OSError) as e:
|
||||
@@ -362,8 +553,15 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
|
||||
Returns:
|
||||
The session ID
|
||||
|
||||
Raises:
|
||||
ValueError: The session would be its own parent.
|
||||
OSError: Saving the session or registering its parent failed.
|
||||
A saved child is retained if parent registration fails.
|
||||
"""
|
||||
sid = session_id or str(uuid.uuid4())
|
||||
if parent_id == sid:
|
||||
raise ValueError(f"Session {sid} cannot be its own parent")
|
||||
|
||||
session = ExtendedSessionData(
|
||||
session_id=sid,
|
||||
@@ -373,15 +571,18 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
metadata=metadata or {},
|
||||
)
|
||||
|
||||
# Update parent's children list without clobbering concurrent message writes
|
||||
if not self._save_extended_session(session):
|
||||
raise OSError(f"Failed to save session {sid}")
|
||||
|
||||
# Register only after the child exists on disk.
|
||||
if parent_id:
|
||||
def _apply(parent_session: SessionData) -> None:
|
||||
assert isinstance(parent_session, ExtendedSessionData)
|
||||
if sid not in parent_session.children_ids:
|
||||
parent_session.children_ids.append(sid)
|
||||
self._modify_session_locked(parent_id, _apply, error_label="update parent children")
|
||||
if not self._modify_session_locked(parent_id, _apply, error_label="update parent children"):
|
||||
raise OSError(f"Session {sid} was saved but registration with parent {parent_id} failed")
|
||||
|
||||
self._save_extended_session(session)
|
||||
return sid
|
||||
|
||||
def fork_session(
|
||||
@@ -395,14 +596,25 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
|
||||
Args:
|
||||
session_id: The session to fork from
|
||||
from_message_index: Message index to fork from (None = all messages)
|
||||
from_message_index: Zero-based, inclusive message index to fork from
|
||||
(None = all messages, including an empty session).
|
||||
title: Optional title for the forked session
|
||||
|
||||
Returns:
|
||||
The new forked session ID
|
||||
|
||||
Raises:
|
||||
ValueError: The explicit message index is outside the parent history.
|
||||
OSError: Saving the fork or registering it with the parent failed.
|
||||
A saved fork is retained if parent registration fails.
|
||||
"""
|
||||
# Force reload to get latest messages from disk
|
||||
parent = self._load_extended_session(session_id, force_reload=True)
|
||||
|
||||
if from_message_index is not None and (
|
||||
from_message_index < 0 or from_message_index >= len(parent.messages)
|
||||
):
|
||||
raise ValueError(f"Invalid message index {from_message_index}")
|
||||
|
||||
# Determine which messages to copy
|
||||
if from_message_index is None:
|
||||
@@ -426,15 +638,17 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
metadata=copy.deepcopy(parent.metadata),
|
||||
)
|
||||
|
||||
self._save_extended_session(forked)
|
||||
if not self._save_extended_session(forked):
|
||||
raise OSError(f"Failed to save forked session {new_id}")
|
||||
|
||||
def _register_fork(parent: SessionData) -> None:
|
||||
if new_id not in parent.children_ids:
|
||||
parent.children_ids.append(new_id)
|
||||
|
||||
self._modify_session_locked(
|
||||
if not self._modify_session_locked(
|
||||
session_id, _register_fork, error_label="register forked session"
|
||||
)
|
||||
):
|
||||
raise OSError(f"Forked session {new_id} was saved but registration with parent {session_id} failed")
|
||||
|
||||
return new_id
|
||||
|
||||
@@ -476,6 +690,9 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
) -> str:
|
||||
"""
|
||||
Create a snapshot of the current session state.
|
||||
|
||||
Persist the active messages, archive and compaction checkpoint so
|
||||
later retention cannot move the snapshot's restore point.
|
||||
|
||||
Args:
|
||||
session_id: The session to snapshot
|
||||
@@ -484,6 +701,9 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
|
||||
Returns:
|
||||
The snapshot ID
|
||||
|
||||
Raises:
|
||||
OSError: The snapshot could not be persisted.
|
||||
"""
|
||||
snapshot = SessionSnapshot(
|
||||
session_id=session_id,
|
||||
@@ -493,14 +713,30 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
)
|
||||
|
||||
def _record_snapshot(session: SessionData) -> None:
|
||||
snapshot.message_index = (
|
||||
len(session.messages) - 1 if session.messages else -1
|
||||
# The caller already holds the session write lock. Read the base
|
||||
# transcript schema so existing durable checkpoint fields are not
|
||||
# lost through the hierarchy's partial legacy field conversion.
|
||||
captured = DefaultSessionStore._load_session_from_disk(
|
||||
self, session_id, self._get_session_path(session_id)
|
||||
)
|
||||
snapshot.message_index = (
|
||||
len(captured.messages) - 1 if captured.messages else -1
|
||||
)
|
||||
snapshot.transcript = copy.deepcopy({
|
||||
"messages": [message.to_dict() for message in captured.messages],
|
||||
"archived_messages": [message.to_dict() for message in captured.archived_messages],
|
||||
"last_compaction": captured.last_compaction.to_dict() if captured.last_compaction else None,
|
||||
})
|
||||
session.last_compaction = captured.last_compaction
|
||||
session.snapshots.append(snapshot)
|
||||
|
||||
self._modify_session_locked(
|
||||
session_id, _record_snapshot, error_label="create snapshot"
|
||||
)
|
||||
# Recording a snapshot must not compact or truncate an imported live
|
||||
# transcript, so skip retention on this write.
|
||||
if not self._modify_session_locked(
|
||||
session_id, _record_snapshot, error_label="create snapshot",
|
||||
apply_retention=False,
|
||||
):
|
||||
raise OSError(f"Failed to save snapshot {snapshot.id} for session {session_id}")
|
||||
|
||||
return snapshot.id
|
||||
|
||||
@@ -518,7 +754,7 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
snapshot_id: The snapshot to revert to
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
True if successful; False for missing or invalidated legacy snapshots.
|
||||
"""
|
||||
def _apply(session: SessionData) -> None:
|
||||
assert isinstance(session, ExtendedSessionData)
|
||||
@@ -534,14 +770,25 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
logger.warning(f"Snapshot {snapshot_id} not found")
|
||||
raise ValueError(f"Snapshot {snapshot_id} not found")
|
||||
|
||||
# Revert messages
|
||||
if snapshot.message_index >= 0:
|
||||
if snapshot.invalidated:
|
||||
raise ValueError("Legacy snapshot position was invalidated by retention")
|
||||
if snapshot.transcript is not None:
|
||||
restored = SessionData.from_dict({"session_id": session_id, **snapshot.transcript})
|
||||
session.messages = restored.messages
|
||||
session.archived_messages = restored.archived_messages
|
||||
session.last_compaction = restored.last_compaction
|
||||
elif snapshot.message_index >= len(session.messages):
|
||||
raise ValueError("Legacy snapshot position is outside the retained history")
|
||||
elif snapshot.message_index >= 0:
|
||||
session.messages = session.messages[:snapshot.message_index + 1]
|
||||
else:
|
||||
session.messages = []
|
||||
|
||||
try:
|
||||
return self._modify_session_locked(session_id, _apply, error_label="revert to snapshot")
|
||||
return self._modify_session_locked(
|
||||
session_id, _apply, error_label="revert to snapshot",
|
||||
apply_retention=False,
|
||||
)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
@@ -673,12 +920,9 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
|
||||
def get_extended_session(self, session_id: str,
|
||||
force_reload: bool = False) -> ExtendedSessionData:
|
||||
"""Get extended session data, using the mtime cache when it is valid.
|
||||
"""Get extended data, reusing parsed data when the JSON bytes match.
|
||||
|
||||
This called _read_session_fresh unconditionally, so every read went to
|
||||
disk and re-parsed the file: _is_cache_valid, _cache_mtimes and
|
||||
_load_extended_session were all built and never reached from the public
|
||||
API. Pass force_reload=True to re-read regardless.
|
||||
Pass force_reload=True to reload regardless of content validation.
|
||||
"""
|
||||
return self._load_extended_session(session_id, force_reload=force_reload)
|
||||
|
||||
@@ -688,9 +932,13 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
if session_id:
|
||||
self._cache.pop(session_id, None)
|
||||
self._extended_cache.pop(session_id, None)
|
||||
self._cache_fingerprints.pop(session_id, None)
|
||||
self._cache_mtimes.pop(session_id, None)
|
||||
else:
|
||||
self._cache.clear()
|
||||
self._extended_cache.clear()
|
||||
self._cache_fingerprints.clear()
|
||||
self._cache_mtimes.clear()
|
||||
|
||||
def export_session(self, session_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
@@ -716,6 +964,9 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
|
||||
Returns:
|
||||
The imported session ID
|
||||
|
||||
Raises:
|
||||
OSError: The imported session could not be persisted.
|
||||
"""
|
||||
session = ExtendedSessionData.from_dict(data)
|
||||
|
||||
@@ -729,7 +980,10 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
session.children_ids = []
|
||||
session.forked_from_message_id = None
|
||||
|
||||
self._save_extended_session(session)
|
||||
# Restore the exported record in full, independent of the destination
|
||||
# window. Retention applies to subsequent ordinary writes.
|
||||
if not self._save_extended_session(session, apply_retention=False):
|
||||
raise OSError(f"Failed to save imported session {session.session_id}")
|
||||
return session.session_id
|
||||
|
||||
# Global hierarchical store instance
|
||||
|
||||
@@ -30,6 +30,7 @@ Usage::
|
||||
|
||||
import os
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from praisonaiagents._logging import get_logger
|
||||
@@ -74,10 +75,13 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
db_path = os.path.expanduser(db_path)
|
||||
self.db_path = db_path
|
||||
self._db_lock = threading.RLock()
|
||||
self._backfill_lock = threading.RLock()
|
||||
self._backfill_running = False
|
||||
self._conn = None
|
||||
self._fts_available = False
|
||||
self._db_ready = False
|
||||
self._backfilled = False
|
||||
self._backfill_upgraded_files = {}
|
||||
|
||||
# ── index lifecycle ───────────────────────────────────────────────
|
||||
|
||||
@@ -230,42 +234,60 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
parts.append(str(content))
|
||||
return "\n".join(parts)
|
||||
|
||||
def _index_session(self, session: SessionData) -> None:
|
||||
"""Insert/replace a session's content in the index (best-effort)."""
|
||||
def _index_session(self, session: SessionData) -> bool:
|
||||
"""Refresh the index best-effort, reporting whether it succeeded."""
|
||||
conn = self._connect()
|
||||
if conn is None:
|
||||
return
|
||||
return False
|
||||
content = self._flatten(session)
|
||||
sid = session.session_id
|
||||
gateway_session_id = getattr(session, "gateway_session_id", None)
|
||||
agent_id = getattr(session, "agent_id", None)
|
||||
try:
|
||||
with self._db_lock:
|
||||
conn.execute("DELETE FROM session_fts WHERE session_id = ?", (sid,))
|
||||
conn.execute(
|
||||
"INSERT INTO session_fts (session_id, content) VALUES (?, ?)",
|
||||
(sid, content),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO session_meta (session_id, updated_at) "
|
||||
"VALUES (?, ?)",
|
||||
(sid, session.updated_at),
|
||||
)
|
||||
# Keep the gateway/agent routing index in sync so inbound
|
||||
# routing is an indexed lookup, not a full-directory scan.
|
||||
if gateway_session_id or agent_id:
|
||||
# A separate WAL reader must see the old or new index record,
|
||||
# never the autocommitted gap between DELETE and INSERT.
|
||||
owns_transaction = not conn.in_transaction
|
||||
conn.execute("SAVEPOINT praisonai_index_refresh")
|
||||
try:
|
||||
conn.execute("DELETE FROM session_fts WHERE session_id = ?", (sid,))
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO session_route "
|
||||
"(session_id, gateway_session_id, agent_id) "
|
||||
"VALUES (?, ?, ?)",
|
||||
(sid, gateway_session_id, agent_id),
|
||||
"INSERT INTO session_fts (session_id, content) VALUES (?, ?)",
|
||||
(sid, content),
|
||||
)
|
||||
else:
|
||||
conn.execute(
|
||||
"DELETE FROM session_route WHERE session_id = ?", (sid,)
|
||||
"INSERT OR REPLACE INTO session_meta (session_id, updated_at) "
|
||||
"VALUES (?, ?)",
|
||||
(sid, session.updated_at),
|
||||
)
|
||||
if gateway_session_id or agent_id:
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO session_route "
|
||||
"(session_id, gateway_session_id, agent_id) "
|
||||
"VALUES (?, ?, ?)",
|
||||
(sid, gateway_session_id, agent_id),
|
||||
)
|
||||
else:
|
||||
conn.execute("DELETE FROM session_route WHERE session_id = ?", (sid,))
|
||||
conn.execute(
|
||||
"DELETE FROM session_index_meta WHERE key = ?",
|
||||
("import_pending:" + sid,),
|
||||
)
|
||||
conn.execute("RELEASE praisonai_index_refresh")
|
||||
except BaseException:
|
||||
if conn.in_transaction:
|
||||
if owns_transaction:
|
||||
# A failed outermost RELEASE leaves the transaction
|
||||
# active, including its uncommitted writes and locks.
|
||||
conn.execute("ROLLBACK")
|
||||
else:
|
||||
conn.execute("ROLLBACK TO praisonai_index_refresh")
|
||||
conn.execute("RELEASE praisonai_index_refresh")
|
||||
raise
|
||||
except Exception as exc: # never let indexing break a write
|
||||
logger.debug("Session index update failed for %s: %s", sid, exc)
|
||||
return False
|
||||
return True
|
||||
|
||||
def _deindex_session(self, session_id: str) -> None:
|
||||
conn = self._connect()
|
||||
@@ -280,22 +302,60 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
logger.debug("Session de-index failed for %s: %s", session_id, exc)
|
||||
|
||||
def _ensure_backfilled(self) -> None:
|
||||
"""Backfill the index from existing JSON transcripts exactly once.
|
||||
"""Backfill existing transcripts, retrying an incomplete pass.
|
||||
|
||||
Guarded by a one-time flag rather than by an empty-index check: a
|
||||
Guarded by a successful-pass flag rather than an empty-index check: a
|
||||
single ``add_message`` on a *new* session could otherwise make the
|
||||
index non-empty and permanently skip backfilling pre-existing JSON
|
||||
transcripts, silently omitting legacy sessions from search results.
|
||||
Only sessions not already present in the index are (re)indexed, so the
|
||||
pass is cheap on a warm index.
|
||||
"""
|
||||
if self._backfilled:
|
||||
return
|
||||
with self._db_lock:
|
||||
if self._backfilled:
|
||||
# Do not hold the database lock while waiting for transcript locks:
|
||||
# imports acquire their file lock before refreshing SQLite.
|
||||
with self._backfill_lock:
|
||||
if self._backfill_running:
|
||||
return
|
||||
self._backfilled = True
|
||||
self._reindex_all()
|
||||
# A corruption callback can query this store during the pass.
|
||||
# It sees the index built so far; other threads wait for completion.
|
||||
self._backfill_running = True
|
||||
try:
|
||||
self._retry_import_indexes()
|
||||
if not self._backfilled:
|
||||
self._backfilled = self._reindex_all()
|
||||
finally:
|
||||
self._backfill_running = False
|
||||
|
||||
def _retry_import_indexes(self) -> None:
|
||||
"""Retry only failed restores, including after reopening the store."""
|
||||
conn = self._connect()
|
||||
if conn is None:
|
||||
return
|
||||
try:
|
||||
with self._db_lock:
|
||||
keys = conn.execute(
|
||||
"SELECT key FROM session_index_meta WHERE key GLOB 'import_pending:*'"
|
||||
).fetchall()
|
||||
for (key,) in keys:
|
||||
sid = key[len("import_pending:"):]
|
||||
filepath = self._get_session_path(sid)
|
||||
try:
|
||||
# Match the writer's file -> database order. File identity
|
||||
# may be unavailable; the lock still protects this read.
|
||||
with FileLock(filepath, self.lock_timeout):
|
||||
if not os.path.isfile(filepath):
|
||||
continue
|
||||
session = self._load_session_from_disk(sid, filepath)
|
||||
# The loader may quarantine a corrupt transcript and
|
||||
# return an empty placeholder. That is not recovery.
|
||||
if not os.path.isfile(filepath):
|
||||
continue
|
||||
self._reingest_spill(sid, session)
|
||||
self._index_session(session)
|
||||
except Exception as exc:
|
||||
logger.debug("Import index retry failed for %s: %s", sid, exc)
|
||||
except Exception as exc:
|
||||
logger.debug("Cannot enumerate pending import indexes: %s", exc)
|
||||
|
||||
def _indexed_ids(self) -> set:
|
||||
"""Return the set of session_ids already fully indexed.
|
||||
@@ -307,8 +367,7 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
otherwise have its already-content-indexed sessions skipped by backfill,
|
||||
leaving ``session_route`` empty and breaking gateway routing. Requiring a
|
||||
route row forces those sessions to be re-indexed so their route rows are
|
||||
populated. Backfill runs at most once (guarded by ``_backfilled``), so
|
||||
this stays a cheap one-time pass.
|
||||
populated. Successful backfill runs once; incomplete passes can retry.
|
||||
"""
|
||||
conn = self._connect()
|
||||
if conn is None:
|
||||
@@ -327,7 +386,7 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
except Exception:
|
||||
return set()
|
||||
|
||||
def _reindex_all(self) -> None:
|
||||
def _reindex_all(self) -> bool:
|
||||
"""Backfill the index from existing JSON transcripts.
|
||||
|
||||
Normally cheap: sessions already fully indexed are skipped. When the
|
||||
@@ -336,31 +395,74 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
indexed, Issue #5031), every session's FTS row is rebuilt once so
|
||||
archived-only queries work immediately after upgrade rather than only
|
||||
after each session's next write. The new version is then persisted so
|
||||
subsequent startups fall back to the cheap skip behaviour.
|
||||
subsequent startups fall back to the cheap skip behaviour. Completed
|
||||
upgrades are retained in memory across failed passes; retries skip
|
||||
those sessions only while their complete index projection and locked
|
||||
transcript generation still match.
|
||||
"""
|
||||
try:
|
||||
filenames = os.listdir(self.session_dir)
|
||||
except (IOError, OSError):
|
||||
return
|
||||
return False
|
||||
conn = self._connect()
|
||||
stale = (
|
||||
conn is not None
|
||||
and self._index_content_version(conn) != self.INDEX_CONTENT_VERSION
|
||||
)
|
||||
already = set() if stale else self._indexed_ids()
|
||||
if conn is None:
|
||||
# No index is available; readers already use the parent JSON scan.
|
||||
return True
|
||||
with self._db_lock:
|
||||
stale = self._index_content_version(conn) != self.INDEX_CONTENT_VERSION
|
||||
try:
|
||||
pending = {
|
||||
row[0][len("import_pending:"):]
|
||||
for row in conn.execute(
|
||||
"SELECT key FROM session_index_meta WHERE key GLOB 'import_pending:*'"
|
||||
).fetchall()
|
||||
}
|
||||
except Exception as exc:
|
||||
logger.debug("Cannot backfill pending import indexes: %s", exc)
|
||||
return False
|
||||
indexed = self._indexed_ids()
|
||||
# A completed row alone cannot prove that a JSON-only peer has not
|
||||
# replaced its transcript between upgrade attempts.
|
||||
already = indexed if not stale else set()
|
||||
complete = True
|
||||
for filename in filenames:
|
||||
if not filename.endswith(".json"):
|
||||
continue
|
||||
sid = filename[:-5]
|
||||
if sid in pending:
|
||||
# Failed restore retries must not be bypassed by a synthetic
|
||||
# empty transcript from the initial/upgrade backfill path.
|
||||
complete = False
|
||||
continue
|
||||
if sid in already:
|
||||
continue
|
||||
try:
|
||||
session = self._read_session_fresh(sid)
|
||||
filepath = self._get_session_path(sid)
|
||||
with FileLock(filepath, self.lock_timeout):
|
||||
identity = self._session_file_identity(filepath)
|
||||
if (
|
||||
stale and sid in indexed and identity is not None
|
||||
and self._backfill_upgraded_files.get(sid) == identity
|
||||
):
|
||||
continue
|
||||
# Match import's file -> database order and keep a peer
|
||||
# replacement from making the just-read projection stale.
|
||||
session = self._load_session_from_disk(sid, filepath)
|
||||
self._reingest_spill(sid, session)
|
||||
with self._lock:
|
||||
self._cache[sid] = session
|
||||
if self._index_session(session):
|
||||
if stale:
|
||||
self._backfill_upgraded_files[sid] = identity
|
||||
else:
|
||||
complete = False
|
||||
except Exception:
|
||||
continue
|
||||
self._index_session(session)
|
||||
if conn is not None:
|
||||
self._set_index_content_version(conn, self.INDEX_CONTENT_VERSION)
|
||||
complete = False
|
||||
if complete:
|
||||
with self._db_lock:
|
||||
self._set_index_content_version(conn, self.INDEX_CONTENT_VERSION)
|
||||
self._backfill_upgraded_files.clear()
|
||||
return complete
|
||||
|
||||
# ── write path: keep the index in sync ────────────────────────────
|
||||
|
||||
@@ -381,16 +483,95 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
# it can restore routing fields removed by the import.
|
||||
fresh = self._load_session_from_disk(session.session_id, filepath)
|
||||
self._reingest_spill(session.session_id, fresh)
|
||||
with self._lock:
|
||||
self._cache[session.session_id] = fresh
|
||||
self._index_session(fresh)
|
||||
# Cleanup must describe the generation actually refreshed,
|
||||
# including a peer write that followed the durable import.
|
||||
session._import_file_identity = self._session_file_identity(filepath)
|
||||
session._import_index_session = fresh
|
||||
with self._lock:
|
||||
self._cache[session.session_id] = fresh
|
||||
if not self._index_session(fresh):
|
||||
raise RuntimeError("post-import index update failed")
|
||||
except Exception as exc:
|
||||
# The JSON write succeeded, but the old index is no longer a
|
||||
# trustworthy view. Fail closed until a later refresh/rebuild.
|
||||
self._deindex_session(session.session_id)
|
||||
self._invalidate_import_index(session)
|
||||
logger.debug("Post-import index refresh failed for %s: %s", session.session_id, exc)
|
||||
return ok
|
||||
|
||||
def _invalidate_import_index(self, session: SessionData) -> None:
|
||||
"""Clear a failed import index unless it matches the current transcript."""
|
||||
identity = getattr(session, "_import_file_identity", None)
|
||||
conn = self._connect()
|
||||
if conn is None:
|
||||
return
|
||||
try:
|
||||
with self._db_lock:
|
||||
# Persist recovery separately from cleanup: unreadable or
|
||||
# changing generations must remain retryable after restart.
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO session_index_meta (key, value) VALUES (?, '1')",
|
||||
("import_pending:" + session.session_id,),
|
||||
)
|
||||
# Serialize with other SQLite writers before checking the file.
|
||||
# A peer saving afterward must index after this transaction;
|
||||
# a peer already indexed has a different atomic file identity.
|
||||
owns_transaction = not conn.in_transaction
|
||||
conn.execute("SAVEPOINT praisonai_import_cleanup")
|
||||
try:
|
||||
# SAVEPOINT alone is deferred. Acquire the SQLite writer
|
||||
# reservation before observing the transcript generation.
|
||||
conn.execute(
|
||||
"UPDATE session_index_meta SET value = value WHERE key = ?",
|
||||
("import_pending:" + session.session_id,),
|
||||
)
|
||||
filepath = self._get_session_path(session.session_id)
|
||||
current_identity = self._session_file_identity(filepath)
|
||||
expected = None
|
||||
if identity is not None and current_identity == identity:
|
||||
expected = getattr(session, "_import_index_session", session)
|
||||
elif current_identity is not None:
|
||||
# A JSON-only peer may replace the file without updating
|
||||
# SQLite. Read without FileLock while holding the DB
|
||||
# transaction; never quarantine or mutate that file.
|
||||
import json
|
||||
with open(filepath, "r", encoding="utf-8") as handle:
|
||||
data = json.load(handle)
|
||||
if self._session_file_identity(filepath) == current_identity:
|
||||
expected = SessionData.from_dict(data)
|
||||
if expected is not None:
|
||||
sid = session.session_id
|
||||
content = conn.execute("SELECT content FROM session_fts WHERE session_id = ?", (sid,)).fetchall()
|
||||
metadata = conn.execute("SELECT updated_at FROM session_meta WHERE session_id = ?", (sid,)).fetchall()
|
||||
routes = conn.execute("SELECT gateway_session_id, agent_id FROM session_route WHERE session_id = ?", (sid,)).fetchall()
|
||||
expected_routes = (
|
||||
[(expected.gateway_session_id, expected.agent_id)]
|
||||
if expected.gateway_session_id or expected.agent_id else []
|
||||
)
|
||||
# A peer may have successfully indexed this same file
|
||||
# generation after our failed update. Preserve a complete
|
||||
# matching projection; file identity alone cannot tell
|
||||
# whether those index rows are already current.
|
||||
current = (
|
||||
content == [(self._flatten(expected),)]
|
||||
and metadata == [(expected.updated_at,)]
|
||||
and routes == expected_routes
|
||||
)
|
||||
if not current:
|
||||
conn.execute("DELETE FROM session_fts WHERE session_id = ?", (sid,))
|
||||
conn.execute("DELETE FROM session_meta WHERE session_id = ?", (sid,))
|
||||
conn.execute("DELETE FROM session_route WHERE session_id = ?", (sid,))
|
||||
conn.execute("RELEASE praisonai_import_cleanup")
|
||||
except BaseException:
|
||||
if conn.in_transaction:
|
||||
if owns_transaction:
|
||||
conn.execute("ROLLBACK")
|
||||
else:
|
||||
conn.execute("ROLLBACK TO praisonai_import_cleanup")
|
||||
conn.execute("RELEASE praisonai_import_cleanup")
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.debug("Post-import index invalidation failed for %s: %s", session.session_id, exc)
|
||||
|
||||
def _modify_session_locked(self, session_id, mutator, **kwargs) -> bool:
|
||||
"""Refresh the index after any locked read-modify-write.
|
||||
|
||||
@@ -432,12 +613,6 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
logger.debug("Post-add index refresh failed for %s: %s", session_id, exc)
|
||||
return ok
|
||||
|
||||
def clear_session(self, session_id: str) -> bool:
|
||||
ok = super().clear_session(session_id)
|
||||
if ok:
|
||||
self._deindex_session(session_id)
|
||||
return ok
|
||||
|
||||
def delete_session(self, session_id: str) -> bool:
|
||||
ok = super().delete_session(session_id)
|
||||
if ok:
|
||||
@@ -464,14 +639,14 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
match = self._to_fts_query(query)
|
||||
rows = conn.execute(
|
||||
"SELECT session_id FROM session_fts WHERE session_fts "
|
||||
"MATCH ? ORDER BY bm25(session_fts) LIMIT ?",
|
||||
"MATCH ? ORDER BY bm25(session_fts), session_id LIMIT ?",
|
||||
(match, fetch),
|
||||
).fetchall()
|
||||
else:
|
||||
like = "%" + query.replace("%", "").replace("_", "") + "%"
|
||||
rows = conn.execute(
|
||||
"SELECT session_id FROM session_fts WHERE content LIKE ? "
|
||||
"LIMIT ?",
|
||||
"ORDER BY session_id LIMIT ?",
|
||||
(like, fetch),
|
||||
).fetchall()
|
||||
except Exception as exc:
|
||||
@@ -489,14 +664,127 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
return '""'
|
||||
return " OR ".join('"%s"' % t for t in terms)
|
||||
|
||||
@contextmanager
|
||||
def _search_id_connection(self, conn):
|
||||
"""Keep WAL candidate enumeration off the writer connection."""
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
|
||||
with self._db_lock:
|
||||
journal = conn.execute("PRAGMA journal_mode").fetchone()[0]
|
||||
if self.db_path != ":memory:" and journal.lower() == "wal":
|
||||
reader = sqlite3.connect(
|
||||
Path(self.db_path).resolve().as_uri() + "?mode=ro",
|
||||
uri=True, isolation_level=None,
|
||||
)
|
||||
try:
|
||||
yield reader
|
||||
finally:
|
||||
reader.close()
|
||||
else:
|
||||
# Non-WAL stores retain the shared lock only for ID enumeration,
|
||||
# never for transcript reads or corruption callbacks.
|
||||
with self._db_lock:
|
||||
yield conn
|
||||
|
||||
@contextmanager
|
||||
def _search_candidate_ids(self, conn, query: str):
|
||||
"""Spool matching IDs in fixed batches; close SQL before file I/O.
|
||||
|
||||
A stable candidate set requires enumerating the matches. Spooling only
|
||||
IDs bounds Python ID memory without copying the index or keeping a
|
||||
SQLite read snapshot across transcript reads and corruption callbacks.
|
||||
"""
|
||||
import json
|
||||
import tempfile
|
||||
|
||||
with tempfile.TemporaryFile(mode="w+t", encoding="utf-8") as spool:
|
||||
with self._search_id_connection(conn) as reader:
|
||||
if self._fts_available:
|
||||
cursor = reader.execute(
|
||||
"SELECT session_id FROM session_fts WHERE session_fts "
|
||||
"MATCH ? ORDER BY bm25(session_fts), session_id",
|
||||
(self._to_fts_query(query),),
|
||||
)
|
||||
else:
|
||||
like = "%" + query.replace("%", "").replace("_", "") + "%"
|
||||
cursor = reader.execute(
|
||||
"SELECT session_id FROM session_fts WHERE content LIKE ? "
|
||||
"ORDER BY session_id", (like,),
|
||||
)
|
||||
try:
|
||||
while True:
|
||||
rows = cursor.fetchmany(128)
|
||||
if not rows:
|
||||
break
|
||||
for row in rows:
|
||||
spool.write(json.dumps(row[0], ensure_ascii=True) + "\n")
|
||||
finally:
|
||||
cursor.close()
|
||||
spool.seek(0)
|
||||
yield (json.loads(line) for line in spool)
|
||||
|
||||
def _read_search_candidates(self, query: str, limit: int, consume=None):
|
||||
"""Read usable lineages outside the SQL lock and read snapshot.
|
||||
|
||||
A consumer scores each payload immediately instead of retaining every
|
||||
continuation. Without a consumer, preserve the private list interface.
|
||||
"""
|
||||
import json
|
||||
|
||||
conn = self._connect()
|
||||
if conn is None:
|
||||
return None
|
||||
self._ensure_backfilled()
|
||||
allowance = max(limit * 5, limit, 1)
|
||||
candidates = []
|
||||
lineages = set()
|
||||
try:
|
||||
with self._search_candidate_ids(conn, query) as ids:
|
||||
for sid in ids:
|
||||
filepath = self._get_session_path(sid)
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as handle:
|
||||
data = json.load(handle)
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, OSError) as exc:
|
||||
self._report_unreadable_session(sid, filepath, exc)
|
||||
continue
|
||||
if consume is None:
|
||||
candidates.append((sid, data))
|
||||
else:
|
||||
consume(sid, data)
|
||||
lineage = self._lineage_key(data)
|
||||
lineages.add(("lineage", lineage) if lineage is not None else ("session", sid))
|
||||
if len(lineages) >= allowance:
|
||||
break
|
||||
except Exception as exc:
|
||||
logger.debug("Index query failed (%s); falling back to scan.", exc)
|
||||
return None
|
||||
return candidates
|
||||
|
||||
# ── gateway/agent routing: indexed key → session lookup ───────────
|
||||
|
||||
def _read_indexed_route(self, session_id: str) -> Optional[SessionData]:
|
||||
"""Verify a route from durable JSON without a stale cache fallback."""
|
||||
filepath = self._get_session_path(session_id)
|
||||
try:
|
||||
with FileLock(filepath, self.lock_timeout):
|
||||
session = self._load_session_from_disk(session_id, filepath)
|
||||
self._reingest_spill(session_id, session)
|
||||
with self._lock:
|
||||
self._cache[session_id] = session
|
||||
return session
|
||||
except (OSError, ValueError, TypeError) as exc:
|
||||
logger.debug("Cannot verify indexed route for %s: %s", session_id, exc)
|
||||
return None
|
||||
|
||||
def get_by_gateway_session(self, gateway_session_id: str) -> Optional[SessionData]:
|
||||
"""Resolve a session by ``gateway_session_id`` via the indexed route.
|
||||
|
||||
Overrides the parent's O(N) ``os.listdir`` + JSON-parse-every-file scan
|
||||
with a single indexed ``SELECT`` on ``session_route`` so inbound routing
|
||||
latency is independent of the number of stored sessions (Issue #2956).
|
||||
The loaded binding must still match; transient reads do not use cache.
|
||||
Falls back to the parent scan if the index is unavailable.
|
||||
"""
|
||||
conn = self._connect()
|
||||
@@ -515,33 +803,52 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
return super().get_by_gateway_session(gateway_session_id)
|
||||
if row is None:
|
||||
return None
|
||||
try:
|
||||
return self._read_session_fresh(row[0])
|
||||
except Exception:
|
||||
session = self._read_indexed_route(row[0])
|
||||
if session is None or session.gateway_session_id != gateway_session_id:
|
||||
return None
|
||||
return session
|
||||
|
||||
def list_sessions_by_gateway_agent(self, agent_id: str, limit: int = 50) -> List[str]:
|
||||
"""List session IDs for a gateway agent via the indexed route.
|
||||
|
||||
Overrides the parent's full-directory scan with an indexed ``SELECT``
|
||||
on ``session_route`` (Issue #2956). Falls back to the parent scan if the
|
||||
index is unavailable.
|
||||
on ``session_route`` (Issue #2956). Matching candidates are verified
|
||||
against JSON outside the database lock, and only verified IDs count
|
||||
toward the limit. Falls back to the parent scan if the index is unavailable.
|
||||
"""
|
||||
conn = self._connect()
|
||||
if conn is None:
|
||||
return super().list_sessions_by_gateway_agent(agent_id, limit)
|
||||
self._ensure_backfilled()
|
||||
if limit == 0:
|
||||
return []
|
||||
verified = []
|
||||
offset = 0
|
||||
try:
|
||||
with self._db_lock:
|
||||
rows = conn.execute(
|
||||
"SELECT session_id FROM session_route WHERE agent_id = ? "
|
||||
"LIMIT ?",
|
||||
(agent_id, limit),
|
||||
).fetchall()
|
||||
while True:
|
||||
with self._db_lock:
|
||||
rows = conn.execute(
|
||||
"SELECT session_id FROM session_route WHERE agent_id = ? "
|
||||
"LIMIT ? OFFSET ?",
|
||||
(agent_id, limit, offset),
|
||||
).fetchall()
|
||||
if not rows:
|
||||
break
|
||||
# File I/O stays outside the database lock. Unknown generations
|
||||
# retain their index rows, but cannot supply obsolete bindings.
|
||||
for (sid,) in rows:
|
||||
session = self._read_indexed_route(sid)
|
||||
if session is not None and session.agent_id == agent_id:
|
||||
verified.append(sid)
|
||||
if limit > 0 and len(verified) >= limit:
|
||||
return verified
|
||||
if limit < 0:
|
||||
break
|
||||
offset += len(rows)
|
||||
except Exception as exc:
|
||||
logger.debug("Route lookup failed (%s); falling back to scan.", exc)
|
||||
return super().list_sessions_by_gateway_agent(agent_id, limit)
|
||||
return [r[0] for r in rows]
|
||||
return verified
|
||||
|
||||
def search(
|
||||
self,
|
||||
@@ -564,29 +871,19 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
if not query:
|
||||
return []
|
||||
|
||||
candidate_ids = self._candidate_ids(query, limit)
|
||||
if candidate_ids is None:
|
||||
return super().search(query, limit=limit, window=window)
|
||||
if not candidate_ids:
|
||||
return []
|
||||
|
||||
needle = query.lower()
|
||||
terms = [t for t in needle.split() if t]
|
||||
hits: List[tuple] = []
|
||||
|
||||
for sid in candidate_ids:
|
||||
filepath = self._get_session_path(sid)
|
||||
try:
|
||||
import json
|
||||
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except (json.JSONDecodeError, IOError, OSError):
|
||||
continue
|
||||
hits: Dict[Any, tuple] = {}
|
||||
position = 0
|
||||
|
||||
def consume(sid, data):
|
||||
"""Retain each lineage's best hit and its original tie position."""
|
||||
nonlocal position
|
||||
ordinal = position
|
||||
position += 1
|
||||
messages = self._searchable_messages(data)
|
||||
if not messages:
|
||||
continue
|
||||
return
|
||||
|
||||
best_index = -1
|
||||
best_score = 0.0
|
||||
@@ -608,7 +905,7 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
best_index = idx
|
||||
|
||||
if best_index < 0:
|
||||
continue
|
||||
return
|
||||
|
||||
start = max(0, best_index - window)
|
||||
end = min(len(messages), best_index + window + 1)
|
||||
@@ -639,13 +936,26 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
messages=context,
|
||||
bookends=self._bookends(messages, self.BOOKEND_SIZE),
|
||||
)
|
||||
hits.append((self._lineage_key(data), hit))
|
||||
lineage = self._lineage_key(data)
|
||||
key = ("lineage", lineage) if lineage is not None else ("session", sid)
|
||||
previous = hits.get(key)
|
||||
if previous is None or (hit.score, hit.when or "") > (
|
||||
previous[1].score, previous[1].when or ""
|
||||
):
|
||||
hits[key] = (lineage, hit, ordinal)
|
||||
|
||||
hits.sort(key=lambda item: (item[1].score, item[1].when or ""), reverse=True)
|
||||
candidates = self._read_search_candidates(query, limit, consume=consume)
|
||||
if candidates is None:
|
||||
return super().search(query, limit=limit, window=window)
|
||||
ranked = sorted(
|
||||
hits.values(),
|
||||
key=lambda item: (item[1].score, item[1].when or "", -item[2]),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
deduped: List[Any] = []
|
||||
seen_lineage: set = set()
|
||||
for lineage, hit in hits:
|
||||
for lineage, hit, _ in ranked:
|
||||
if lineage is not None:
|
||||
if lineage in seen_lineage:
|
||||
continue
|
||||
|
||||
@@ -21,6 +21,7 @@ try:
|
||||
_HAS_FCNTL = True
|
||||
except ImportError:
|
||||
_HAS_FCNTL = False
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
@@ -87,6 +88,48 @@ DEFAULT_RETENTION = RETENTION_COMPACT
|
||||
# past this many entries, so operators can spot runaway sessions on disk.
|
||||
ARCHIVE_WARN_THRESHOLD = 10_000
|
||||
|
||||
# Bound the in-memory session cache so a long-lived gateway or bot process that
|
||||
# serves many users does not retain every session ever touched in RAM. The
|
||||
# on-disk file is already the source of truth, so an eviction just means the
|
||||
# next access reloads from disk. ``0`` or ``None`` disables the bound.
|
||||
DEFAULT_CACHE_MAXSIZE = 256
|
||||
|
||||
|
||||
class _LRUSessionCache(OrderedDict):
|
||||
"""An ``OrderedDict`` that evicts the least-recently-used entry once it
|
||||
grows past ``maxsize``. Reads and writes refresh recency so the hot set
|
||||
survives. ``maxsize <= 0`` disables eviction (unbounded, legacy behaviour).
|
||||
|
||||
All mutating/reading operations used by :class:`DefaultSessionStore`
|
||||
(``in``, ``[]`` get/set, ``del``, ``pop``, ``clear``) work unchanged; only
|
||||
insertion order is managed, so callers need no other changes.
|
||||
"""
|
||||
|
||||
def __init__(self, maxsize: int = DEFAULT_CACHE_MAXSIZE, on_evict=None):
|
||||
super().__init__()
|
||||
self.maxsize = maxsize
|
||||
# Optional callback invoked with the key of each LRU-evicted entry, so a
|
||||
# companion structure (e.g. the hierarchical store's _cache_mtimes) can
|
||||
# drop the same key and stay bounded alongside this cache.
|
||||
self._on_evict = on_evict
|
||||
|
||||
def __getitem__(self, key):
|
||||
value = super().__getitem__(key)
|
||||
self.move_to_end(key)
|
||||
return value
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
if key in self:
|
||||
super().__setitem__(key, value)
|
||||
self.move_to_end(key)
|
||||
else:
|
||||
super().__setitem__(key, value)
|
||||
if self.maxsize and self.maxsize > 0:
|
||||
while len(self) > self.maxsize:
|
||||
evicted_key, _ = self.popitem(last=False)
|
||||
if self._on_evict is not None:
|
||||
self._on_evict(evicted_key)
|
||||
|
||||
@dataclass
|
||||
class SessionMessage:
|
||||
"""A single message in a session.
|
||||
@@ -599,6 +642,7 @@ class DefaultSessionStore:
|
||||
retention: Optional[str] = None,
|
||||
active_window: Optional[int] = None,
|
||||
mirror: Optional[Any] = None,
|
||||
cache_maxsize: int = DEFAULT_CACHE_MAXSIZE,
|
||||
):
|
||||
"""
|
||||
Initialize session store.
|
||||
@@ -625,6 +669,11 @@ class DefaultSessionStore:
|
||||
appended to the mirror on a background thread (local-first: the
|
||||
local write always happens first and a mirror outage never
|
||||
blocks or fails the turn). Left ``None`` there is zero overhead.
|
||||
cache_maxsize: Maximum number of sessions kept in the in-memory LRU
|
||||
cache (default 256). A long-lived multi-tenant process otherwise
|
||||
retains every session ever touched in RAM. Disk is the source of
|
||||
truth, so an evicted session simply reloads on next access. Pass
|
||||
``0`` to disable the bound (legacy unbounded cache).
|
||||
"""
|
||||
self.session_dir = session_dir or _resolve_default_session_dir()
|
||||
self.max_messages = max_messages
|
||||
@@ -649,7 +698,11 @@ class DefaultSessionStore:
|
||||
self.active_window = active_window if active_window is not None else max_messages
|
||||
|
||||
self._lock = threading.RLock()
|
||||
self._cache: Dict[str, SessionData] = {}
|
||||
# LRU-bounded so a long-lived multi-tenant process does not retain every
|
||||
# session ever touched in RAM. Pass cache_maxsize=0 for the legacy
|
||||
# unbounded cache. Disk remains the source of truth, so an eviction just
|
||||
# triggers a reload on next access.
|
||||
self._cache: Dict[str, SessionData] = _LRUSessionCache(cache_maxsize)
|
||||
|
||||
self._mirror = mirror
|
||||
|
||||
@@ -979,6 +1032,16 @@ class DefaultSessionStore:
|
||||
except Exception: # pragma: no cover - observability must never break load
|
||||
logger.debug("SESSION corruption hook failed", exc_info=True)
|
||||
|
||||
def _report_unreadable_session(self, session_id: str, filepath: str, error: Exception) -> None:
|
||||
"""Report decode failures during scans without modifying the source file."""
|
||||
if isinstance(error, (UnicodeDecodeError, json.JSONDecodeError)):
|
||||
logger.warning("Skipping unreadable session file %s: %s", filepath, error)
|
||||
self._fire_corruption_hook(session_id, str(error), None)
|
||||
|
||||
def _dump_session_json(self, data: Any, stream) -> None:
|
||||
"""Write durable JSON; subclasses can format private storage records."""
|
||||
json.dump(data, stream, indent=2, ensure_ascii=False)
|
||||
|
||||
def _atomic_write_json(self, filepath: str, data: Any) -> bool:
|
||||
"""Atomically write JSON data to disk (temp file + os.replace)."""
|
||||
temp_path = None
|
||||
@@ -993,7 +1056,7 @@ class DefaultSessionStore:
|
||||
suffix=".tmp",
|
||||
) as f:
|
||||
temp_path = f.name
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
self._dump_session_json(data, f)
|
||||
|
||||
os.replace(temp_path, filepath)
|
||||
return True
|
||||
@@ -1153,16 +1216,55 @@ class DefaultSessionStore:
|
||||
except (IOError, OSError):
|
||||
return
|
||||
|
||||
seen = {
|
||||
(m.role, m.content, m.timestamp) for m in session.messages
|
||||
}
|
||||
def _freeze(value: Any) -> Any:
|
||||
# A flat, iterative key avoids recursion both while converting and
|
||||
# hashing deeply nested JSON. Sort object keys to preserve equality
|
||||
# independently of insertion order; keep scalar types distinct.
|
||||
tokens = []
|
||||
stack = [("visit", value)]
|
||||
ancestors = set()
|
||||
while stack:
|
||||
operation, item = stack.pop()
|
||||
if operation == "leave":
|
||||
ancestors.remove(item)
|
||||
elif operation == "key":
|
||||
tokens.append(("key", type(item), item))
|
||||
elif isinstance(item, (list, dict)):
|
||||
identity = id(item)
|
||||
if identity in ancestors:
|
||||
raise ValueError("cyclic spill content")
|
||||
ancestors.add(identity)
|
||||
stack.append(("leave", identity))
|
||||
tokens.append((type(item), len(item)))
|
||||
if isinstance(item, list):
|
||||
stack.extend(("visit", child) for child in reversed(item))
|
||||
else:
|
||||
for key in sorted(item, reverse=True):
|
||||
stack.append(("visit", item[key]))
|
||||
stack.append(("key", key))
|
||||
else:
|
||||
tokens.append((type(item), item))
|
||||
return tuple(tokens)
|
||||
|
||||
def _key(message: SessionMessage) -> tuple:
|
||||
return tuple(_freeze(value) for value in (
|
||||
message.role, message.content, message.timestamp,
|
||||
))
|
||||
|
||||
seen = set()
|
||||
for message in session.messages:
|
||||
try:
|
||||
seen.add(_key(message))
|
||||
except (RecursionError, TypeError, ValueError):
|
||||
# An unsupported existing value must not stop other salvage.
|
||||
continue
|
||||
recovered: List[tuple] = [] # (filepath, [SessionMessage])
|
||||
for filename in candidates:
|
||||
filepath = os.path.join(spill_dir, filename)
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except (json.JSONDecodeError, IOError, OSError):
|
||||
except (json.JSONDecodeError, UnicodeDecodeError, RecursionError, IOError, OSError):
|
||||
continue
|
||||
# A syntactically valid spill can still carry an unexpected shape
|
||||
# (non-object root, non-list messages, non-object message). Guard
|
||||
@@ -1176,15 +1278,21 @@ class DefaultSessionStore:
|
||||
if not isinstance(raw_messages, list):
|
||||
continue
|
||||
msgs = []
|
||||
for raw in raw_messages:
|
||||
if not isinstance(raw, dict):
|
||||
continue
|
||||
msg = SessionMessage.from_dict(raw)
|
||||
key = (msg.role, msg.content, msg.timestamp)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
msgs.append(msg)
|
||||
pending_keys = set()
|
||||
try:
|
||||
for raw in raw_messages:
|
||||
if not isinstance(raw, dict):
|
||||
continue
|
||||
msg = SessionMessage.from_dict(raw)
|
||||
key = _key(msg)
|
||||
if key in seen or key in pending_keys:
|
||||
continue
|
||||
pending_keys.add(key)
|
||||
msgs.append(msg)
|
||||
except (RecursionError, TypeError, ValueError):
|
||||
# Retain the entire spill, without poisoning neighbor dedup.
|
||||
continue
|
||||
seen.update(pending_keys)
|
||||
recovered.append((filepath, msgs))
|
||||
|
||||
if not recovered:
|
||||
@@ -1201,7 +1309,7 @@ class DefaultSessionStore:
|
||||
# expose an oversized transcript that stays inconsistent until a
|
||||
# later mutation happens to compact it.
|
||||
self._enforce_window(session)
|
||||
if not self._atomic_write_json(filepath, session.to_dict()):
|
||||
if not self._atomic_write_json(filepath, self._session_to_storage(session)):
|
||||
# Could not fold the salvage back in durably — leave the spill
|
||||
# files in place so a later load can retry.
|
||||
return
|
||||
@@ -1219,6 +1327,7 @@ class DefaultSessionStore:
|
||||
mutator: Callable[[SessionData], None],
|
||||
*,
|
||||
error_label: str = "modify session",
|
||||
apply_retention: bool = True,
|
||||
) -> bool:
|
||||
"""Apply mutator after reloading from disk under FileLock."""
|
||||
filepath = self._get_session_path(session_id)
|
||||
@@ -1234,9 +1343,10 @@ class DefaultSessionStore:
|
||||
mutator(session)
|
||||
session.updated_at = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
self._enforce_window(session)
|
||||
if apply_retention:
|
||||
self._enforce_window(session)
|
||||
|
||||
if not self._atomic_write_json(filepath, session.to_dict()):
|
||||
if not self._atomic_write_json(filepath, self._session_to_storage(session)):
|
||||
logger.error(f"Failed to {error_label} {session_id}")
|
||||
return False
|
||||
|
||||
@@ -1245,6 +1355,10 @@ class DefaultSessionStore:
|
||||
|
||||
return True
|
||||
|
||||
def _session_to_storage(self, session: SessionData) -> Dict[str, Any]:
|
||||
"""Serialize a durable record; subclasses may retain portable exports."""
|
||||
return session.to_dict()
|
||||
|
||||
def _save_session(self, session: SessionData) -> bool:
|
||||
"""Save session to disk with atomic write."""
|
||||
filepath = self._get_session_path(session.session_id)
|
||||
@@ -1254,7 +1368,7 @@ class DefaultSessionStore:
|
||||
self._enforce_window(session)
|
||||
|
||||
with FileLock(filepath, self.lock_timeout):
|
||||
if not self._atomic_write_json(filepath, session.to_dict()):
|
||||
if not self._atomic_write_json(filepath, self._session_to_storage(session)):
|
||||
logger.error(f"Failed to save session {session.session_id}")
|
||||
return False
|
||||
return True
|
||||
@@ -1321,7 +1435,7 @@ class DefaultSessionStore:
|
||||
self._enforce_window(session)
|
||||
|
||||
# Write atomically
|
||||
if not self._atomic_write_json(filepath, session.to_dict()):
|
||||
if not self._atomic_write_json(filepath, self._session_to_storage(session)):
|
||||
logger.error(f"Failed to save session {session_id}")
|
||||
# Issue #3597: durable write failed (disk-full / corruption).
|
||||
# Spill just this turn to a fallback file and fire the
|
||||
@@ -1714,7 +1828,8 @@ class DefaultSessionStore:
|
||||
"total_tokens": data.get("total_tokens") or data.get("token_count") or (data.get("metadata") or {}).get("total_tokens"),
|
||||
"cost": data.get("cost") or (data.get("metadata") or {}).get("cost"),
|
||||
})
|
||||
except (json.JSONDecodeError, IOError):
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError) as exc:
|
||||
self._report_unreadable_session(filename[:-5], filepath, exc)
|
||||
continue
|
||||
except (IOError, OSError):
|
||||
pass
|
||||
@@ -1748,7 +1863,8 @@ class DefaultSessionStore:
|
||||
data = json.load(f)
|
||||
if data.get("agent_name") == agent_name:
|
||||
session_ids.append(data.get("session_id", filename[:-5]))
|
||||
except (json.JSONDecodeError, IOError):
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError) as exc:
|
||||
self._report_unreadable_session(filename[:-5], filepath, exc)
|
||||
continue
|
||||
except (IOError, OSError):
|
||||
pass
|
||||
@@ -1856,7 +1972,8 @@ class DefaultSessionStore:
|
||||
data = json.load(f)
|
||||
if data.get("gateway_session_id") == gateway_session_id:
|
||||
return SessionData.from_dict(data)
|
||||
except (json.JSONDecodeError, IOError):
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError) as exc:
|
||||
self._report_unreadable_session(filename[:-5], filepath, exc)
|
||||
continue
|
||||
except (IOError, OSError):
|
||||
pass
|
||||
@@ -1886,7 +2003,8 @@ class DefaultSessionStore:
|
||||
data = json.load(f)
|
||||
if data.get("agent_id") == agent_id:
|
||||
session_ids.append(data.get("session_id", filename[:-5]))
|
||||
except (json.JSONDecodeError, IOError):
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError) as exc:
|
||||
self._report_unreadable_session(filename[:-5], filepath, exc)
|
||||
continue
|
||||
except (IOError, OSError):
|
||||
pass
|
||||
@@ -2096,7 +2214,8 @@ class DefaultSessionStore:
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except (json.JSONDecodeError, IOError):
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError) as exc:
|
||||
self._report_unreadable_session(filename[:-5], filepath, exc)
|
||||
continue
|
||||
|
||||
messages = self._searchable_messages(data)
|
||||
@@ -2277,12 +2396,9 @@ class DefaultSessionStore:
|
||||
return {"version": self.PORTABLE_VERSION, "sessions": sessions}
|
||||
|
||||
def export_all(self) -> Dict[str, Any]:
|
||||
"""Export every stored session to a portable, versioned payload."""
|
||||
"""Export every stored session; raise OSError for an incomplete backup."""
|
||||
sessions: List[Dict[str, Any]] = []
|
||||
try:
|
||||
filenames = os.listdir(self.session_dir)
|
||||
except (IOError, OSError):
|
||||
filenames = []
|
||||
filenames = os.listdir(self.session_dir)
|
||||
for filename in filenames:
|
||||
if not filename.endswith(".json"):
|
||||
continue
|
||||
@@ -2290,8 +2406,9 @@ class DefaultSessionStore:
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
sessions.append(json.load(f))
|
||||
except (json.JSONDecodeError, IOError, OSError):
|
||||
continue
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError, OSError) as exc:
|
||||
self._report_unreadable_session(filename[:-5], filepath, exc)
|
||||
raise OSError(f"Incomplete session export: cannot read {filename}") from exc
|
||||
return {"version": self.PORTABLE_VERSION, "sessions": sessions}
|
||||
|
||||
def _collect_lineage(
|
||||
@@ -2303,15 +2420,13 @@ class DefaultSessionStore:
|
||||
(``lineage_id`` / ``root_session_id`` / ``thread_id``) so a compacted /
|
||||
rotated continuation exports alongside its logical session. Returns the
|
||||
raw session dicts (already portable). Empty when no lineage is known.
|
||||
Raises OSError when an unreadable record makes lineage coverage unknown.
|
||||
"""
|
||||
lineage = self._lineage_key(session.to_dict())
|
||||
if not lineage:
|
||||
return []
|
||||
out: List[Dict[str, Any]] = []
|
||||
try:
|
||||
filenames = os.listdir(self.session_dir)
|
||||
except (IOError, OSError):
|
||||
return []
|
||||
filenames = os.listdir(self.session_dir)
|
||||
for filename in filenames:
|
||||
if not filename.endswith(".json"):
|
||||
continue
|
||||
@@ -2319,8 +2434,9 @@ class DefaultSessionStore:
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except (json.JSONDecodeError, IOError, OSError):
|
||||
continue
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError, OSError) as exc:
|
||||
self._report_unreadable_session(filename[:-5], filepath, exc)
|
||||
raise OSError(f"Incomplete session export: cannot read {filename}") from exc
|
||||
if data.get("session_id") == exclude:
|
||||
continue
|
||||
if self._lineage_key(data) == lineage:
|
||||
@@ -2342,11 +2458,31 @@ class DefaultSessionStore:
|
||||
# Check again inside the same lock that protects the replacement.
|
||||
if not overwrite and os.path.exists(filepath):
|
||||
raise FileExistsError(filepath)
|
||||
if not self._atomic_write_json(filepath, session.to_dict()):
|
||||
if not self._atomic_write_json(filepath, self._session_to_storage(session)):
|
||||
logger.error(f"Failed to save imported session {session.session_id}")
|
||||
return False
|
||||
session._import_file_identity = self._session_file_identity(filepath)
|
||||
# Persistence hooks own their cache value; a subclass may refresh
|
||||
# it from a newer durable generation before returning to import.
|
||||
with self._lock:
|
||||
self._cache[session.session_id] = session
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _session_file_identity(filepath):
|
||||
"""Identify an atomic file generation without depending on wall-clock ordering."""
|
||||
try:
|
||||
stat = os.stat(filepath)
|
||||
except OSError:
|
||||
# A path metadata failure need not make the atomic generation
|
||||
# unknowable: an opened descriptor can still identify that file.
|
||||
try:
|
||||
with open(filepath, "rb") as current:
|
||||
stat = os.fstat(current.fileno())
|
||||
except OSError:
|
||||
return None
|
||||
return stat.st_dev, stat.st_ino, stat.st_mtime_ns, stat.st_size
|
||||
|
||||
def import_sessions(
|
||||
self,
|
||||
payload: Dict[str, Any],
|
||||
@@ -2446,8 +2582,6 @@ class DefaultSessionStore:
|
||||
{"session_id": session_id, "reason": "write failed"}
|
||||
)
|
||||
continue
|
||||
with self._lock:
|
||||
self._cache[session_id] = session
|
||||
report.imported += 1
|
||||
except FileExistsError:
|
||||
report.skipped.append(
|
||||
|
||||
@@ -60,6 +60,7 @@ class FileBackend:
|
||||
self.suffix = suffix
|
||||
self.pretty = pretty
|
||||
self._lock = threading.Lock()
|
||||
self._unreadable_paths = set()
|
||||
|
||||
# Create directory if needed
|
||||
self.storage_dir.mkdir(parents=True, exist_ok=True)
|
||||
@@ -85,6 +86,10 @@ class FileBackend:
|
||||
"""Save data with the given key."""
|
||||
with self._lock:
|
||||
file_path = self._key_to_path(key)
|
||||
if file_path in self._unreadable_paths:
|
||||
if file_path.exists():
|
||||
raise OSError(f"Refusing to overwrite unreadable record: {file_path}")
|
||||
self._unreadable_paths.discard(file_path)
|
||||
|
||||
# Atomic write via temp file
|
||||
try:
|
||||
@@ -112,18 +117,25 @@ class FileBackend:
|
||||
raise
|
||||
|
||||
def load(self, key: str) -> Optional[Dict[str, Any]]:
|
||||
"""Load data by key."""
|
||||
"""Load data by key; failed reads protect the record from later saves.
|
||||
|
||||
A successful reload or deletion permits writes again. Returning None
|
||||
for a failed read must not let a caller persist defaults over its source.
|
||||
"""
|
||||
file_path = self._key_to_path(key)
|
||||
|
||||
if not file_path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError, IOError) as e:
|
||||
logger.warning(f"Failed to load {key}: {e}")
|
||||
return None
|
||||
with self._lock:
|
||||
if not file_path.exists():
|
||||
self._unreadable_paths.discard(file_path)
|
||||
return None
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
self._unreadable_paths.discard(file_path)
|
||||
return data
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError) as e:
|
||||
self._unreadable_paths.add(file_path)
|
||||
logger.warning(f"Failed to load {key}: {e}")
|
||||
return None
|
||||
|
||||
def delete(self, key: str) -> bool:
|
||||
"""Delete data by key."""
|
||||
@@ -133,6 +145,7 @@ class FileBackend:
|
||||
if file_path.exists():
|
||||
try:
|
||||
file_path.unlink()
|
||||
self._unreadable_paths.discard(file_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete {key}: {e}")
|
||||
@@ -163,6 +176,7 @@ class FileBackend:
|
||||
if self._path_to_key(file_path) is not None:
|
||||
try:
|
||||
file_path.unlink()
|
||||
self._unreadable_paths.discard(file_path)
|
||||
count += 1
|
||||
except Exception:
|
||||
pass
|
||||
@@ -310,9 +324,9 @@ class SQLiteBackend:
|
||||
if prefix:
|
||||
cur.execute(f"""
|
||||
SELECT key FROM {self._quoted_table}
|
||||
WHERE key LIKE ?
|
||||
WHERE instr(key, ?) = 1
|
||||
ORDER BY key
|
||||
""", (f"{prefix}%",))
|
||||
""", (prefix,))
|
||||
else:
|
||||
cur.execute(f"""
|
||||
SELECT key FROM {self._quoted_table}
|
||||
@@ -337,10 +351,8 @@ class SQLiteBackend:
|
||||
conn = self._get_conn()
|
||||
cur = conn.cursor()
|
||||
|
||||
cur.execute(f"SELECT COUNT(*) as count FROM {self._quoted_table}")
|
||||
count = cur.fetchone()["count"]
|
||||
|
||||
cur.execute(f"DELETE FROM {self._quoted_table}")
|
||||
count = cur.rowcount
|
||||
conn.commit()
|
||||
|
||||
return count
|
||||
|
||||
@@ -191,6 +191,7 @@ class BaseJSONStore:
|
||||
self._lock = threading.Lock()
|
||||
self._data: Dict[str, Any] = {}
|
||||
self._backend = backend
|
||||
self._read_failed = False
|
||||
|
||||
if backend is None:
|
||||
# File-based storage (default)
|
||||
@@ -236,7 +237,9 @@ class BaseJSONStore:
|
||||
else:
|
||||
with open(self.storage_path, "r", encoding="utf-8") as f:
|
||||
self._data = json.load(f)
|
||||
except (json.JSONDecodeError, IOError) as e:
|
||||
self._read_failed = False
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, IOError) as e:
|
||||
self._read_failed = not isinstance(e, FileNotFoundError)
|
||||
logger.warning(f"Failed to load {self.storage_path}: {e}")
|
||||
self._data = self._default_data()
|
||||
|
||||
@@ -250,6 +253,13 @@ class BaseJSONStore:
|
||||
raise
|
||||
return
|
||||
|
||||
# Defaults from a failed read must not replace their unreadable source.
|
||||
# A successful reload or explicit reset/deletion permits writes again.
|
||||
if self._read_failed:
|
||||
if self.storage_path.exists():
|
||||
raise OSError(f"Refusing to overwrite unreadable record: {self.storage_path}")
|
||||
self._read_failed = False
|
||||
|
||||
try:
|
||||
dir_path = self.storage_path.parent
|
||||
dir_path.mkdir(parents=True, exist_ok=True)
|
||||
@@ -312,9 +322,18 @@ class BaseJSONStore:
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all stored data."""
|
||||
self._reset(self._default_data())
|
||||
|
||||
def _reset(self, data: Dict[str, Any]) -> None:
|
||||
"""Atomically replace data for an explicit clear, restoring guards on failure."""
|
||||
with self._lock:
|
||||
self._data = self._default_data()
|
||||
self._save()
|
||||
previous_data, previous_failed = self._data, self._read_failed
|
||||
self._data, self._read_failed = data, False
|
||||
try:
|
||||
self._save()
|
||||
except BaseException:
|
||||
self._data, self._read_failed = previous_data, previous_failed
|
||||
raise
|
||||
|
||||
def delete(self) -> bool:
|
||||
"""
|
||||
@@ -334,6 +353,7 @@ class BaseJSONStore:
|
||||
if self.storage_path.exists():
|
||||
self.storage_path.unlink()
|
||||
self._data = self._default_data()
|
||||
self._read_failed = False
|
||||
return True
|
||||
return False
|
||||
except Exception as e:
|
||||
@@ -484,6 +504,7 @@ class AsyncBaseJSONStore:
|
||||
self.storage_path = Path(storage_path)
|
||||
self._backend = backend
|
||||
self._data: Dict[str, Any] = {}
|
||||
self._read_failed = False
|
||||
self._async_lock: Optional[Any] = None # Lazy init asyncio.Lock
|
||||
self._async_lock_init: threading.Lock = threading.Lock() # Thread-safe initialization
|
||||
|
||||
@@ -524,7 +545,12 @@ class AsyncBaseJSONStore:
|
||||
async with aiofiles.open(self.storage_path, "r", encoding="utf-8") as f:
|
||||
content = await f.read()
|
||||
self._data = json.loads(content)
|
||||
except (FileNotFoundError, json.JSONDecodeError):
|
||||
self._read_failed = False
|
||||
except FileNotFoundError:
|
||||
self._read_failed = False
|
||||
self._data = self._default_data()
|
||||
except (UnicodeDecodeError, json.JSONDecodeError):
|
||||
self._read_failed = True
|
||||
self._data = self._default_data()
|
||||
|
||||
return self._data.copy()
|
||||
@@ -541,6 +567,10 @@ class AsyncBaseJSONStore:
|
||||
else:
|
||||
self._backend.save(self._storage_key, self._data)
|
||||
else:
|
||||
if self._read_failed:
|
||||
if self.storage_path.exists():
|
||||
raise OSError(f"Refusing to overwrite unreadable record: {self.storage_path}")
|
||||
self._read_failed = False
|
||||
# File-based async write. Serialise to a string FIRST, write it
|
||||
# to a temp file, then os.replace it into place so a crash/kill
|
||||
# mid-write cannot truncate the previous good file (writing
|
||||
@@ -591,6 +621,7 @@ class AsyncBaseJSONStore:
|
||||
if self.storage_path.exists():
|
||||
self.storage_path.unlink()
|
||||
self._data = self._default_data()
|
||||
self._read_failed = False
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ import os
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
from ..approval import require_approval
|
||||
from ..utils.atomic_io import update_json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -40,7 +41,7 @@ class TodoTools:
|
||||
return self._todo_file
|
||||
|
||||
def _load_todos(self) -> List[Dict]:
|
||||
"""Load todos from file."""
|
||||
"""Load todos from file (read-only helper used by ``todo_list``)."""
|
||||
todo_file = self._get_todo_file()
|
||||
if not os.path.exists(todo_file):
|
||||
return []
|
||||
@@ -51,16 +52,40 @@ class TodoTools:
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
logger.warning(f"Failed to load todos: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def _save_todos(self, todos: List[Dict]) -> None:
|
||||
"""Save todos to file."""
|
||||
todo_file = self._get_todo_file()
|
||||
try:
|
||||
with open(todo_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(todos, f, indent=2, ensure_ascii=False)
|
||||
except OSError as e:
|
||||
logger.error(f"Failed to save todos: {e}")
|
||||
raise
|
||||
"""Atomically persist a full todo list.
|
||||
|
||||
Retained for the MCP CLI adapter
|
||||
(``praisonai_mcp.mcp_server.adapters.cli_tools``), which does its own
|
||||
read-modify-write via ``_load_todos`` + ``_save_todos`` and writes the
|
||||
owner's record shape. The write goes through ``atomic_write_json`` so a
|
||||
crash mid-write never truncates the previous good file. In-process
|
||||
mutations (``todo_add``/``todo_update``) use the locked ``_update_todos``
|
||||
below instead; this plain save keeps the cross-adapter contract intact.
|
||||
"""
|
||||
from ..utils.atomic_io import atomic_write_json
|
||||
atomic_write_json(
|
||||
self._get_todo_file(), todos, indent=2, ensure_ascii=False
|
||||
)
|
||||
|
||||
def _update_todos(self, mutate) -> List[Dict]:
|
||||
"""Locked atomic read-modify-write over the todo file.
|
||||
|
||||
Parallel tool calls and multiple agents sharing the default
|
||||
``~/.praisonai/todos.json`` previously raced through an unlocked
|
||||
read-modify-write plus a truncating ``open(..., 'w')``, so most
|
||||
concurrent adds were lost and a torn write reset the whole plan. Routing
|
||||
every mutation through ``update_json`` serialises them and writes
|
||||
atomically; a corrupt file is moved aside instead of silently emptied.
|
||||
"""
|
||||
return update_json(
|
||||
self._get_todo_file(),
|
||||
lambda todos: (mutate(todos), todos)[1],
|
||||
default=list,
|
||||
indent=2,
|
||||
ensure_ascii=False,
|
||||
)
|
||||
|
||||
def _emit_update(self, todos: List[Dict]) -> None:
|
||||
"""Publish the full ordered list so subscribed frontends can render live.
|
||||
@@ -88,19 +113,22 @@ class TodoTools:
|
||||
JSON string with result
|
||||
"""
|
||||
try:
|
||||
todos = self._load_todos()
|
||||
|
||||
new_todo = {
|
||||
"id": len(todos) + 1,
|
||||
"task": task,
|
||||
"priority": priority,
|
||||
"category": category,
|
||||
"status": "pending",
|
||||
"created_at": self._get_timestamp()
|
||||
}
|
||||
|
||||
todos.append(new_todo)
|
||||
self._save_todos(todos)
|
||||
new_todo: Dict = {}
|
||||
|
||||
def mutate(todos: List[Dict]) -> None:
|
||||
new_todo.update({
|
||||
# Derive the id from the current max inside the lock so
|
||||
# concurrent adds never collide on ``len(todos) + 1``.
|
||||
"id": max((t.get("id", 0) for t in todos), default=0) + 1,
|
||||
"task": task,
|
||||
"priority": priority,
|
||||
"category": category,
|
||||
"status": "pending",
|
||||
"created_at": self._get_timestamp(),
|
||||
})
|
||||
todos.append(dict(new_todo))
|
||||
|
||||
todos = self._update_todos(mutate)
|
||||
self._emit_update(todos)
|
||||
|
||||
return json.dumps({
|
||||
@@ -157,30 +185,31 @@ class TodoTools:
|
||||
JSON string with result
|
||||
"""
|
||||
try:
|
||||
todos = self._load_todos()
|
||||
|
||||
todo_found = False
|
||||
for todo in todos:
|
||||
if todo.get("id") == todo_id:
|
||||
if status:
|
||||
# Convention: exactly one item is in_progress at a time.
|
||||
if status == "in_progress":
|
||||
for other in todos:
|
||||
if other is not todo and other.get("status") == "in_progress":
|
||||
other["status"] = "pending"
|
||||
todo["status"] = status
|
||||
if task:
|
||||
todo["task"] = task
|
||||
if priority:
|
||||
todo["priority"] = priority
|
||||
todo["updated_at"] = self._get_timestamp()
|
||||
todo_found = True
|
||||
break
|
||||
|
||||
if not todo_found:
|
||||
found = {"value": False}
|
||||
|
||||
def mutate(todos: List[Dict]) -> None:
|
||||
for todo in todos:
|
||||
if todo.get("id") == todo_id:
|
||||
if status:
|
||||
# Convention: exactly one item is in_progress at a time.
|
||||
if status == "in_progress":
|
||||
for other in todos:
|
||||
if other is not todo and other.get("status") == "in_progress":
|
||||
other["status"] = "pending"
|
||||
todo["status"] = status
|
||||
if task:
|
||||
todo["task"] = task
|
||||
if priority:
|
||||
todo["priority"] = priority
|
||||
todo["updated_at"] = self._get_timestamp()
|
||||
found["value"] = True
|
||||
break
|
||||
|
||||
todos = self._update_todos(mutate)
|
||||
|
||||
if not found["value"]:
|
||||
return json.dumps({"success": False, "error": f"Todo {todo_id} not found"})
|
||||
|
||||
self._save_todos(todos)
|
||||
self._emit_update(todos)
|
||||
|
||||
return json.dumps({
|
||||
|
||||
@@ -12,9 +12,21 @@ new content atomically on both POSIX and Windows.
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Any
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Any, Callable
|
||||
|
||||
__all__ = ["atomic_write_json", "atomic_write_text"]
|
||||
__all__ = [
|
||||
"atomic_write_json",
|
||||
"atomic_write_text",
|
||||
"quarantine_corrupt_file",
|
||||
"update_json",
|
||||
]
|
||||
|
||||
# Per-path in-process mutexes so threads in one process serialise before they
|
||||
# even reach the cross-process ``FileLock``. Keyed by absolute path.
|
||||
_path_mutexes: "defaultdict[str, threading.RLock]" = defaultdict(threading.RLock)
|
||||
|
||||
|
||||
def atomic_write_text(path: str, text: str, encoding: str = "utf-8") -> None:
|
||||
@@ -70,3 +82,70 @@ def atomic_write_json(path: str, data: Any, **dump_kwargs: Any) -> None:
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
def quarantine_corrupt_file(path: str) -> str:
|
||||
"""Move a corrupt file aside to a unique ``.corrupt`` sidecar.
|
||||
|
||||
Each quarantine uses a distinct ``<path>.<pid>.<ns>.corrupt`` name so a
|
||||
file that goes corrupt more than once never overwrites an earlier recovery
|
||||
copy. The move uses :func:`os.replace` (atomic rename) and the chosen
|
||||
destination is reported back. Raises ``OSError`` on failure so the caller
|
||||
can stop instead of overwriting potentially recoverable bytes.
|
||||
"""
|
||||
dest = "{path}.{pid}.{ns}.corrupt".format(
|
||||
path=path, pid=os.getpid(), ns=time.time_ns()
|
||||
)
|
||||
os.replace(path, dest)
|
||||
return dest
|
||||
|
||||
|
||||
def update_json(
|
||||
path: str,
|
||||
mutate: Callable[[Any], Any],
|
||||
default: Callable[[], Any] = dict,
|
||||
**dump_kwargs: Any,
|
||||
) -> Any:
|
||||
"""Locked read-modify-write for a JSON file.
|
||||
|
||||
Holds a per-path in-process mutex *and* the cross-process
|
||||
:class:`~praisonaiagents.session.store.FileLock` for the whole
|
||||
read-modify-write, so concurrent threads and processes no longer lose each
|
||||
other's updates (last-writer-wins). The write goes through
|
||||
:func:`atomic_write_json`, so a crash mid-write never truncates the previous
|
||||
good file.
|
||||
|
||||
A corrupt JSON file is *not* treated as empty: it is moved aside to a
|
||||
unique ``<path>.<pid>.<ns>.corrupt`` sidecar and the update proceeds from
|
||||
``default()`` so a single torn file can never silently wipe real data
|
||||
without leaving evidence. If the quarantine move itself fails, the error
|
||||
propagates instead of overwriting the original file.
|
||||
|
||||
Args:
|
||||
path: JSON file path.
|
||||
mutate: Callable invoked with the loaded data; mutate it in place.
|
||||
default: Factory for the initial value when the file is missing/corrupt.
|
||||
**dump_kwargs: Passed through to :func:`json.dump` (e.g. ``indent=2``).
|
||||
|
||||
Returns:
|
||||
Whatever ``mutate`` returns (useful for surfacing the new item/length).
|
||||
"""
|
||||
# Imported lazily to avoid a package-level import cycle (session -> utils).
|
||||
from ..session.store import FileLock
|
||||
|
||||
abspath = os.path.abspath(path)
|
||||
with _path_mutexes[abspath], FileLock(path, timeout=10):
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except FileNotFoundError:
|
||||
data = default()
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
# Quarantine under a unique name; if the move itself fails we must
|
||||
# NOT fall through to the atomic write, which would replace the
|
||||
# original and destroy its potentially recoverable bytes.
|
||||
quarantine_corrupt_file(path)
|
||||
data = default()
|
||||
result = mutate(data)
|
||||
atomic_write_json(path, data, **dump_kwargs)
|
||||
return result
|
||||
@@ -30,12 +30,9 @@ __all__ = ["StepCacheProtocol", "InMemoryStepCache", "make_step_key", "resolve_s
|
||||
def _snapshot(value: Any) -> Any:
|
||||
# Deep copy so a cached value cannot be mutated in place by later execution
|
||||
# (or a caller) and corrupt a future hit. Some step outputs hold objects a
|
||||
# deepcopy cannot handle (locks, live clients); a cache must never crash the
|
||||
# workflow, so fall back to the original reference in that rare case.
|
||||
try:
|
||||
return copy.deepcopy(value)
|
||||
except Exception:
|
||||
return value
|
||||
# deepcopy cannot handle (locks, live clients). The caller must treat a
|
||||
# failed snapshot as uncacheable, never retain a mutable original reference.
|
||||
return copy.deepcopy(value)
|
||||
|
||||
|
||||
def _identify(step: Any) -> str:
|
||||
@@ -109,20 +106,33 @@ class InMemoryStepCache:
|
||||
with self._lock:
|
||||
entry = self._entries.get(key)
|
||||
if entry is not None:
|
||||
try:
|
||||
snapshot = _snapshot(entry)
|
||||
except Exception:
|
||||
del self._entries[key]
|
||||
self.misses += 1
|
||||
return None
|
||||
self._entries.move_to_end(key)
|
||||
self.hits += 1
|
||||
# Snapshot out: a cached value contains live mutable objects
|
||||
# (e.g. the workflow's variables dict) that later execution
|
||||
# mutates. Handing back the stored object would let one run's
|
||||
# later mutations corrupt a future hit.
|
||||
return _snapshot(entry)
|
||||
return snapshot
|
||||
self.misses += 1
|
||||
return None
|
||||
|
||||
def set(self, key: str, value: Dict[str, Any]) -> None:
|
||||
# Snapshot in for the same reason: freeze the value at store time so
|
||||
# the caller mutating it afterwards cannot alter what is cached.
|
||||
snapshot = _snapshot(value)
|
||||
try:
|
||||
snapshot = _snapshot(value)
|
||||
except Exception:
|
||||
# An optional cache must not fail the workflow or serve stale data
|
||||
# if an existing key is replaced with an uncacheable value.
|
||||
with self._lock:
|
||||
self._entries.pop(key, None)
|
||||
return
|
||||
with self._lock:
|
||||
self._entries[key] = snapshot
|
||||
self._entries.move_to_end(key)
|
||||
|
||||
@@ -25,7 +25,6 @@ import json
|
||||
import copy
|
||||
import time
|
||||
import logging
|
||||
import threading
|
||||
from praisonaiagents._logging import get_logger
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Callable, Tuple, Union
|
||||
@@ -48,10 +47,7 @@ DEFAULT_MAX_PARALLEL_WORKERS = 3
|
||||
# summarisation request consumes roughly as many tokens as it saves.
|
||||
MIN_BRANCHES_FOR_LLM_SUMMARY = 3
|
||||
|
||||
# Guards lazy creation of each Workflow's per-instance _run_lock so two threads
|
||||
# entering run()/astart() concurrently on a fresh instance cannot each create
|
||||
# and acquire a *different* lock object (which would defeat the run guard).
|
||||
_RUN_LOCK_INIT_GUARD = threading.Lock()
|
||||
from .._run_lock import ensure_run_lock
|
||||
|
||||
|
||||
class _WriteTrackingDict(dict):
|
||||
@@ -796,14 +792,7 @@ class AgentFlow:
|
||||
first reach run()/astart() concurrently observe the *same* lock object
|
||||
(double-checked locking) rather than each minting and acquiring its own.
|
||||
"""
|
||||
lock = self._run_lock
|
||||
if lock is None:
|
||||
with _RUN_LOCK_INIT_GUARD:
|
||||
lock = self._run_lock
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
self._run_lock = lock
|
||||
return lock
|
||||
return ensure_run_lock(self)
|
||||
|
||||
def __post_init__(self):
|
||||
"""Resolve consolidated params to internal values."""
|
||||
@@ -1624,14 +1613,37 @@ class AgentFlow:
|
||||
if _step_cache is not None:
|
||||
from .step_cache import make_step_key
|
||||
_cache_key = make_step_key(step, previous_output, input, all_variables)
|
||||
_cached = _step_cache.get(_cache_key)
|
||||
try:
|
||||
_cached = copy.deepcopy(_step_cache.get(_cache_key))
|
||||
except Exception:
|
||||
_cached = None
|
||||
logger.warning("Could not restore cached workflow step '%s'", step.name)
|
||||
if _cached is not None:
|
||||
if verbose:
|
||||
print(f"↩︎ cache hit: {step.name}")
|
||||
previous_output = _cached.get("output")
|
||||
results.append({"step": step.name, "output": previous_output})
|
||||
# Replay the full step delta recorded on the cold run so a
|
||||
# hit is indistinguishable from re-executing: the step's
|
||||
# status/retries, any handler-supplied variables (and the
|
||||
# generated output variable, present only when the cold run
|
||||
# did not stop first), and -- crucially -- the stop signal.
|
||||
# A hit that only appended {step, output} and continued lost
|
||||
# the stop flag and the handler's variables, so a workflow
|
||||
# that stopped on the cold run ran on through the cached one.
|
||||
cached_record = _cached.get("step_record") or {
|
||||
"step": step.name, "output": previous_output,
|
||||
"status": "completed", "retries": 0,
|
||||
}
|
||||
results.append(dict(cached_record))
|
||||
self.step_statuses[step.name] = cached_record["status"]
|
||||
if hasattr(step, "status"):
|
||||
step.status = cached_record["status"]
|
||||
if _cached.get("variables"):
|
||||
all_variables.update(_cached["variables"])
|
||||
if _cached.get("stop"):
|
||||
if verbose:
|
||||
print(f"🛑 Workflow stopped at: {step.name}")
|
||||
break
|
||||
i += 1
|
||||
continue
|
||||
|
||||
@@ -1672,6 +1684,7 @@ class AgentFlow:
|
||||
retry_count = 0
|
||||
validation_feedback = None
|
||||
guardrail_failed = False
|
||||
cached_variable_updates = {} if _cache_key is not None else None
|
||||
|
||||
while retry_count <= max_retries:
|
||||
step_error = None
|
||||
@@ -1688,7 +1701,14 @@ class AgentFlow:
|
||||
output = result.output
|
||||
stop = result.stop_workflow
|
||||
if result.variables:
|
||||
# Accumulate across retries so the cached
|
||||
# snapshot matches the cold run: a rejected
|
||||
# attempt that wrote {a, b} followed by an
|
||||
# accepted attempt that wrote {a} leaves both
|
||||
# keys in the live run, so both must be cached.
|
||||
all_variables.update(result.variables)
|
||||
if cached_variable_updates is not None:
|
||||
cached_variable_updates.update(result.variables)
|
||||
else:
|
||||
output = str(result)
|
||||
|
||||
@@ -1985,17 +2005,26 @@ class AgentFlow:
|
||||
# Only a SUCCESSFUL step is cached. Caching a failure would serve
|
||||
# the failure again on every re-run, turning a transient error into
|
||||
# a permanent one that no retry could clear. The step's output
|
||||
# variable is stored too: a fresh run starts with empty working
|
||||
# variables, so a cache HIT that only restored `output` would leave
|
||||
# `<step>_output` (or step.output_variable) missing and break the
|
||||
# next step's substitutions. We snapshot exactly the delta this step
|
||||
# writes below (var_name = output_variable or f"{name}_output").
|
||||
# variables and stop signal must be stored too: a cache hit must
|
||||
# preserve handler state and control flow, not just output text.
|
||||
# Store only this step's variable updates, including its output
|
||||
# variable unless it stopped before that variable was written.
|
||||
if _cache_key is not None and not step_failed:
|
||||
_cached_var_name = step.output_variable or f"{step.name}_output"
|
||||
_step_cache.set(
|
||||
_cache_key,
|
||||
{"output": output, "variables": {_cached_var_name: output}},
|
||||
)
|
||||
if not stop:
|
||||
_cached_var_name = step.output_variable or f"{step.name}_output"
|
||||
cached_variable_updates[_cached_var_name] = output
|
||||
try:
|
||||
_step_cache.set(
|
||||
_cache_key,
|
||||
copy.deepcopy({
|
||||
"output": output, "variables": cached_variable_updates,
|
||||
"stop": stop, "step_record": step_record,
|
||||
}),
|
||||
)
|
||||
except Exception:
|
||||
# Custom caches may only accept serializable state. Keep
|
||||
# the successful workflow result when storage rejects it.
|
||||
logger.warning("Could not cache workflow step '%s'", step.name)
|
||||
|
||||
if verbose:
|
||||
print(f"✅ {step.name}: {str(output)}")
|
||||
@@ -2708,7 +2737,11 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
|
||||
)
|
||||
from .step_cache import make_step_key
|
||||
key = make_step_key(step, previous_output, input, all_variables)
|
||||
hit = cache.get(key)
|
||||
try:
|
||||
hit = copy.deepcopy(cache.get(key))
|
||||
except Exception:
|
||||
hit = None
|
||||
logger.warning("Could not restore cached workflow step")
|
||||
if hit is not None:
|
||||
if verbose:
|
||||
print(f"↩︎ cache hit: {getattr(step, 'name', getattr(step, '__name__', step))}")
|
||||
@@ -2724,7 +2757,10 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
|
||||
# that would replay the failure on every identical re-run and never let
|
||||
# the step retry -- turning a transient error into a permanent one.
|
||||
if isinstance(result, dict) and not result.get("error"):
|
||||
cache.set(key, result)
|
||||
try:
|
||||
cache.set(key, copy.deepcopy(result))
|
||||
except Exception:
|
||||
logger.warning("Could not cache workflow step")
|
||||
return result
|
||||
|
||||
def _execute_single_step_uncached(
|
||||
@@ -5575,7 +5611,11 @@ class WorkflowManager:
|
||||
)
|
||||
from .step_cache import make_step_key
|
||||
key = make_step_key(step, None, original_input, all_variables)
|
||||
hit = cache.get(key)
|
||||
try:
|
||||
hit = copy.deepcopy(cache.get(key))
|
||||
except Exception:
|
||||
hit = None
|
||||
logger.warning("Could not restore cached workflow step")
|
||||
if hit is not None:
|
||||
if verbose:
|
||||
print(f"↩︎ cache hit: {getattr(step, 'name', getattr(step, '__name__', step))}")
|
||||
@@ -5587,7 +5627,10 @@ class WorkflowManager:
|
||||
# Only cache a successful result -- a failed step must be free to retry
|
||||
# on the next run rather than replay a cached failure forever.
|
||||
if isinstance(result, dict) and not result.get("error"):
|
||||
cache.set(key, result)
|
||||
try:
|
||||
cache.set(key, copy.deepcopy(result))
|
||||
except Exception:
|
||||
logger.warning("Could not cache workflow step")
|
||||
return result
|
||||
|
||||
def _execute_single_step_nocache(
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.9"
|
||||
version = "1.7.10"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
@@ -73,6 +73,11 @@ llm = [
|
||||
"pydantic>=2.10.0"
|
||||
]
|
||||
|
||||
# System One / Jev decision models (optional; HTTP client works without this)
|
||||
decisions = [
|
||||
"typesafe-sdk>=0.7.0",
|
||||
]
|
||||
|
||||
# Add API dependencies
|
||||
api = [
|
||||
"fastapi>=0.115.0",
|
||||
@@ -140,6 +145,7 @@ all = [
|
||||
"praisonaiagents[knowledge]",
|
||||
"praisonaiagents[graph]",
|
||||
"praisonaiagents[llm]",
|
||||
"praisonaiagents[decisions]",
|
||||
"praisonaiagents[mcp]",
|
||||
"praisonaiagents[api]",
|
||||
"praisonaiagents[telemetry]",
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Real agent output retains StepResult state and stop behavior on cache hits."""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
os.environ.get("RUN_REAL_KEY_TESTS") != "1"
|
||||
or not os.environ.get("OPENAI_API_KEY"),
|
||||
reason="Set RUN_REAL_KEY_TESTS=1 and OPENAI_API_KEY for the real agentic test",
|
||||
)
|
||||
|
||||
|
||||
def test_cached_real_agent_step_retains_variables_and_stop():
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
from praisonaiagents.workflows.workflows import AgentFlow, StepResult
|
||||
|
||||
agent = Agent(
|
||||
name="cache-smoke",
|
||||
instructions="Answer briefly.",
|
||||
llm=LLM(
|
||||
model=os.environ.get("PRAISONAI_TEST_MODEL", "gpt-4o-mini"),
|
||||
api_key=os.environ["OPENAI_API_KEY"],
|
||||
base_url=os.environ.get("OPENAI_BASE_URL"),
|
||||
),
|
||||
output="silent",
|
||||
)
|
||||
calls = []
|
||||
|
||||
def finish(ctx):
|
||||
calls.append("finish")
|
||||
response = agent.start("Say hello in one sentence.")
|
||||
print(response)
|
||||
assert isinstance(response, str) and response.strip()
|
||||
return StepResult(
|
||||
output=response, stop_workflow=True, variables={"answer": response}
|
||||
)
|
||||
|
||||
def downstream(ctx):
|
||||
calls.append("downstream")
|
||||
return "unexpected"
|
||||
|
||||
flow = AgentFlow(steps=[finish, downstream], cache=True)
|
||||
first = flow.run("same", verbose=False)
|
||||
second = flow.run("same", verbose=False)
|
||||
|
||||
assert calls == ["finish"]
|
||||
assert first == second
|
||||
assert second["variables"]["answer"] == second["output"]
|
||||
@@ -1191,6 +1191,74 @@ class TestPluginTypedSubsystems:
|
||||
mgr.disable(name)
|
||||
mgr.unregister(name)
|
||||
|
||||
def test_has_registered_plugins_reflects_registration(self):
|
||||
"""The hot-path gate is False with no plugins and True once registered."""
|
||||
from praisonaiagents.plugins.plugin import PluginType
|
||||
from praisonaiagents.plugins import get_plugin_manager
|
||||
from praisonaiagents.plugins.manager import has_registered_plugins
|
||||
|
||||
mgr = get_plugin_manager()
|
||||
# Clean baseline: no plugins registered on the singleton.
|
||||
if not mgr._plugins:
|
||||
assert has_registered_plugins() is False
|
||||
|
||||
plugin, name = self._make_plugin(
|
||||
PluginType.GUARDRAIL, name="gate_probe_test",
|
||||
as_guardrail=lambda self: object(),
|
||||
)
|
||||
mgr.register(plugin)
|
||||
try:
|
||||
assert has_registered_plugins() is True
|
||||
finally:
|
||||
mgr.disable(name)
|
||||
mgr.unregister(name)
|
||||
|
||||
def test_has_registered_plugins_does_not_create_singleton(self):
|
||||
"""The gate must stay False and NOT build the singleton when uncreated.
|
||||
|
||||
The hot-path optimisation relies on the common case (no plugins ever
|
||||
registered) being allocation-free. This isolates the
|
||||
``_default_manager is None`` branch, which the other tests never reach
|
||||
because they call ``get_plugin_manager()`` first.
|
||||
"""
|
||||
from praisonaiagents.plugins import manager as manager_mod
|
||||
|
||||
saved = manager_mod._default_manager
|
||||
manager_mod._default_manager = None
|
||||
try:
|
||||
assert manager_mod.has_registered_plugins() is False
|
||||
# The cheap check must not have constructed the singleton.
|
||||
assert manager_mod._default_manager is None
|
||||
finally:
|
||||
manager_mod._default_manager = saved
|
||||
|
||||
def test_registered_plugin_still_merges_despite_fast_path(self):
|
||||
"""Agent init must still route a registered plugin even with the gate."""
|
||||
from praisonaiagents.plugins.plugin import PluginType
|
||||
from praisonaiagents.plugins import get_plugin_manager
|
||||
from praisonaiagents.plugins.manager import has_registered_plugins
|
||||
from praisonaiagents import Agent
|
||||
|
||||
class _G:
|
||||
def validate_output(self, content, **kw):
|
||||
return True, content
|
||||
|
||||
plugin, name = self._make_plugin(
|
||||
PluginType.GUARDRAIL, name="fastpath_merge_test",
|
||||
as_guardrail=lambda self: _G(),
|
||||
)
|
||||
mgr = get_plugin_manager()
|
||||
mgr.register(plugin)
|
||||
mgr.enable(name)
|
||||
try:
|
||||
# Gate is True, so the merge methods do NOT early-return.
|
||||
assert has_registered_plugins() is True
|
||||
agent = Agent(name="t", instructions="x", llm="gpt-4o-mini")
|
||||
assert agent._tool_result_guardrails
|
||||
finally:
|
||||
mgr.disable(name)
|
||||
mgr.unregister(name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,108 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.decisions.routing import (
|
||||
DecisionRoutePlan,
|
||||
DecisionTriageRouter,
|
||||
triaged_start,
|
||||
)
|
||||
|
||||
|
||||
def test_triaged_start_delegates_by_choice():
|
||||
fake_decision = mock.Mock()
|
||||
fake_decision.choice.return_value = "billing"
|
||||
fake_decision.confidence.return_value = 0.95
|
||||
fake_decision.answers = {}
|
||||
|
||||
class FakeAgent:
|
||||
def __init__(self):
|
||||
self.prompt = None
|
||||
|
||||
def start(self, prompt, **kwargs):
|
||||
self.prompt = prompt
|
||||
return "BILLING-OK"
|
||||
|
||||
billing = FakeAgent()
|
||||
plan = DecisionRoutePlan(
|
||||
route_question="team",
|
||||
questions={"team": {"type": "choice", "instructions": "x", "criteria": {"billing": "b"}}},
|
||||
routes={"billing": billing, "technical": FakeAgent()},
|
||||
)
|
||||
|
||||
with mock.patch("praisonaiagents.decisions.routing.triage_decision", return_value=fake_decision):
|
||||
out = triaged_start("refund please", plan)
|
||||
|
||||
assert out["route"] == "billing"
|
||||
assert out["response"] == "BILLING-OK"
|
||||
assert billing.prompt == "refund please"
|
||||
|
||||
|
||||
def test_decision_triage_router():
|
||||
router = DecisionTriageRouter(
|
||||
DecisionRoutePlan(
|
||||
questions={"team": {"type": "choice", "instructions": "x", "criteria": {"other": "o"}}},
|
||||
routes={"other": lambda d: "OK"},
|
||||
)
|
||||
)
|
||||
fake = mock.Mock(choice=lambda k: "other", confidence=lambda k: 1.0, answers={}, model="nimble", raw={})
|
||||
with mock.patch("praisonaiagents.decisions.routing.triage_decision", return_value=fake):
|
||||
out = router.start("hello")
|
||||
assert out["response"] == "OK"
|
||||
|
||||
|
||||
def _choice_q(*criteria):
|
||||
return {"type": "choice", "instructions": "x", "criteria": {c: c for c in criteria}}
|
||||
|
||||
|
||||
def test_plan_rejects_empty_questions():
|
||||
with pytest.raises(ValueError):
|
||||
DecisionRoutePlan(questions={}, routes={"a": lambda d: "x"})
|
||||
|
||||
|
||||
def test_plan_rejects_non_choice_route_question():
|
||||
with pytest.raises(ValueError):
|
||||
DecisionRoutePlan(
|
||||
route_question="refund",
|
||||
questions={"refund": {"type": "noul", "instructions": "x"}},
|
||||
routes={"a": lambda d: "x"},
|
||||
)
|
||||
|
||||
|
||||
def test_missing_confidence_triggers_fallback():
|
||||
fake = mock.Mock()
|
||||
fake.choice.return_value = "billing"
|
||||
fake.confidence.return_value = None
|
||||
fake.answers = {}
|
||||
plan = DecisionRoutePlan(
|
||||
route_question="team",
|
||||
min_confidence=0.5,
|
||||
questions={"team": _choice_q("billing", "other")},
|
||||
routes={"billing": lambda d: "B", "other": lambda d: "FALLBACK"},
|
||||
fallback_route="other",
|
||||
)
|
||||
with mock.patch("praisonaiagents.decisions.routing.triage_decision", return_value=fake):
|
||||
out = triaged_start("hi", plan)
|
||||
assert out["route"] == "other"
|
||||
|
||||
|
||||
def test_extra_state_does_not_clobber_prompt():
|
||||
from praisonaiagents.decisions.routing import triage_decision
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_system_one(*, state, questions, model, api_base):
|
||||
captured["state"] = dict(state)
|
||||
return mock.Mock()
|
||||
|
||||
with mock.patch("praisonaiagents.decisions.routing.system_one", fake_system_one):
|
||||
triage_decision(
|
||||
"classify-me",
|
||||
questions={"team": _choice_q("a")},
|
||||
state_key="message",
|
||||
extra_state={"message": "SHOULD-NOT-WIN", "locale": "en"},
|
||||
)
|
||||
assert captured["state"]["message"] == "classify-me"
|
||||
assert captured["state"]["locale"] == "en"
|
||||
@@ -0,0 +1,173 @@
|
||||
"""Unit tests for System One / Jev decision API client."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.decisions import (
|
||||
SystemOneResult,
|
||||
choice_question,
|
||||
is_decision_model,
|
||||
noul_question,
|
||||
resolve_system_one_base_url,
|
||||
score_question,
|
||||
system_one,
|
||||
)
|
||||
from praisonaiagents.local.capabilities import Cap, parse_ollama_capabilities
|
||||
|
||||
|
||||
def test_is_decision_model():
|
||||
assert is_decision_model("nimble")
|
||||
assert is_decision_model("ollama/nimble")
|
||||
assert is_decision_model("tev1:0.8b")
|
||||
assert not is_decision_model("llama3.2")
|
||||
|
||||
|
||||
def test_question_helpers():
|
||||
c = choice_question("Pick team", {"a": "A", "b": "B"})
|
||||
assert c["type"] == "choice"
|
||||
n = noul_question("Yes or no?")
|
||||
assert n["type"] == "noul"
|
||||
s = score_question("Rate", ["Low", "High"])
|
||||
assert s["type"] == "score" and len(s["criteria"]) == 2
|
||||
|
||||
|
||||
def test_parse_ollama_decision_capability():
|
||||
caps = parse_ollama_capabilities({"capabilities": ["decision", "completion"]})
|
||||
assert Cap.SYSTEM_ONE in caps
|
||||
|
||||
|
||||
def test_system_one_result_accessors():
|
||||
raw = {
|
||||
"model": "nimble",
|
||||
"answers": {
|
||||
"team": {"type": "choice", "choice": "billing", "confidence": 0.9},
|
||||
"refund": {"type": "noul", "noul": 0.99},
|
||||
"urgency": {"type": "score", "score": 0.5, "confidence": 0.1},
|
||||
},
|
||||
"usage": {"input_tokens": 1, "output_tokens": 2},
|
||||
}
|
||||
r = SystemOneResult(model="nimble", answers=raw["answers"], usage=raw["usage"], raw=raw)
|
||||
assert r.choice("team") == "billing"
|
||||
assert r.noul("refund") == pytest.approx(0.99)
|
||||
assert r.score("urgency") == pytest.approx(0.5)
|
||||
assert r.confidence("team") == pytest.approx(0.9)
|
||||
|
||||
|
||||
def test_system_one_http_mock():
|
||||
response_body = {
|
||||
"model": "nimble",
|
||||
"answers": {
|
||||
"team": {"type": "choice", "choice": "billing", "confidence": 0.92},
|
||||
},
|
||||
"usage": {"input_tokens": 10, "output_tokens": 1},
|
||||
}
|
||||
payload_holder = {}
|
||||
|
||||
def fake_urlopen(req, timeout=0):
|
||||
payload_holder["url"] = req.full_url
|
||||
payload_holder["body"] = json.loads(req.data.decode())
|
||||
return io.BytesIO(json.dumps(response_body).encode())
|
||||
|
||||
with mock.patch("urllib.request.urlopen", fake_urlopen):
|
||||
result = system_one(
|
||||
state={"ticket": "refund please"},
|
||||
questions={"team": choice_question("Team?", {"billing": "pay"})},
|
||||
model="nimble",
|
||||
api_base="http://127.0.0.1:11434",
|
||||
)
|
||||
assert result.choice("team") == "billing"
|
||||
assert payload_holder["url"].endswith("/v1/systemone")
|
||||
assert payload_holder["body"]["model"] == "nimble"
|
||||
|
||||
|
||||
def test_resolve_base_url_env(monkeypatch):
|
||||
monkeypatch.setenv("TYPESAFE_BASE_URL", "http://127.0.0.1:11434")
|
||||
assert resolve_system_one_base_url() == "http://127.0.0.1:11434"
|
||||
|
||||
|
||||
def test_default_timeout_honours_explicit_120(monkeypatch):
|
||||
from praisonaiagents.decisions.system_one import _default_timeout
|
||||
|
||||
monkeypatch.delenv("SYSTEM_ONE_TIMEOUT", raising=False)
|
||||
monkeypatch.delenv("OLLAMA_SYSTEM_ONE_TIMEOUT", raising=False)
|
||||
assert _default_timeout(120.0) == 120.0
|
||||
assert _default_timeout(5.0) == 5.0
|
||||
assert _default_timeout(None) == 120.0
|
||||
monkeypatch.setenv("SYSTEM_ONE_TIMEOUT", "30")
|
||||
assert _default_timeout(None) == 30.0
|
||||
assert _default_timeout(7.0) == 7.0
|
||||
|
||||
|
||||
def test_validate_base_url_rejects_plain_http_remote():
|
||||
from praisonaiagents.decisions.system_one import _validate_base_url
|
||||
|
||||
_validate_base_url("http://127.0.0.1:11434")
|
||||
_validate_base_url("http://localhost:11434")
|
||||
_validate_base_url("https://api.typesafe.ai")
|
||||
with pytest.raises(ValueError):
|
||||
_validate_base_url("http://evil.example.com")
|
||||
with pytest.raises(ValueError):
|
||||
_validate_base_url("ftp://127.0.0.1")
|
||||
|
||||
|
||||
def test_typesafe_key_not_sent_to_ollama_host(monkeypatch):
|
||||
monkeypatch.delenv("TYPESAFE_BASE_URL", raising=False)
|
||||
monkeypatch.setenv("TYPESAFE_API_KEY", "secret-typesafe")
|
||||
monkeypatch.setenv("OLLAMA_HOST", "http://127.0.0.1:11434")
|
||||
captured = {}
|
||||
|
||||
def fake_urlopen(req, timeout=0):
|
||||
captured["auth"] = req.headers.get("Authorization")
|
||||
return io.BytesIO(json.dumps({"model": "nimble", "answers": {}}).encode())
|
||||
|
||||
with mock.patch("urllib.request.urlopen", fake_urlopen):
|
||||
system_one(
|
||||
state={"ticket": "x"},
|
||||
questions={"team": choice_question("Team?", {"billing": "pay"})},
|
||||
model="nimble",
|
||||
)
|
||||
assert captured["auth"] == "Bearer ollama"
|
||||
|
||||
|
||||
def test_typesafe_key_sent_to_typesafe_base(monkeypatch):
|
||||
monkeypatch.setenv("TYPESAFE_BASE_URL", "https://api.typesafe.ai")
|
||||
monkeypatch.setenv("TYPESAFE_API_KEY", "secret-typesafe")
|
||||
captured = {}
|
||||
|
||||
def fake_urlopen(req, timeout=0):
|
||||
captured["auth"] = req.headers.get("Authorization")
|
||||
return io.BytesIO(json.dumps({"model": "nimble", "answers": {}}).encode())
|
||||
|
||||
with mock.patch("urllib.request.urlopen", fake_urlopen):
|
||||
system_one(
|
||||
state={"ticket": "x"},
|
||||
questions={"team": choice_question("Team?", {"billing": "pay"})},
|
||||
model="nimble",
|
||||
)
|
||||
assert captured["auth"] == "Bearer secret-typesafe"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not __import__("os").environ.get("RUN_OLLAMA_DECISION_TESTS"),
|
||||
reason="Set RUN_OLLAMA_DECISION_TESTS=1 with Ollama 0.35+ and `ollama pull nimble`",
|
||||
)
|
||||
def test_system_one_live_ollama():
|
||||
result = system_one(
|
||||
state={"ticket": "I was charged twice. Please refund."},
|
||||
questions={
|
||||
"team": choice_question(
|
||||
"Which team?",
|
||||
{"billing": "Payments", "technical": "Bugs", "other": "Other"},
|
||||
),
|
||||
"refund": noul_question("Does the customer ask for a refund?"),
|
||||
},
|
||||
model="nimble",
|
||||
timeout=900,
|
||||
)
|
||||
assert result.choice("team") in {"billing", "technical", "other"}
|
||||
assert result.noul("refund") is not None
|
||||
@@ -256,32 +256,63 @@ KNOWN_DEAD_OUT_OF_SCOPE = {
|
||||
}
|
||||
|
||||
|
||||
def _product_event_references(root):
|
||||
"""Index references once across the same non-test Python source scope."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
|
||||
references = set()
|
||||
pattern = re.compile(r"(?:HookEvent|PluginHook)\.([A-Za-z_][A-Za-z_0-9]*)")
|
||||
for directory, folders, files in os.walk(root):
|
||||
folders[:] = [name for name in folders if name != "tests"]
|
||||
for name in files:
|
||||
path = Path(directory) / name
|
||||
if path.suffix != ".py" or "hooks/types.py" in path.as_posix():
|
||||
continue
|
||||
references.update(pattern.findall(path.read_text(encoding="utf-8", errors="replace")))
|
||||
return references
|
||||
|
||||
|
||||
def test_reference_scan_reads_product_files_once(tmp_path, monkeypatch):
|
||||
from pathlib import Path
|
||||
|
||||
product = tmp_path / "src" / "module.py"
|
||||
product.parent.mkdir()
|
||||
product.write_text("HookEvent.BEFORE_TOOL\nPluginHook.ON_INIT\n", encoding="utf-8")
|
||||
excluded = tmp_path / "tests" / "test_fake.py"
|
||||
excluded.parent.mkdir()
|
||||
excluded.write_text("HookEvent.TEST_ONLY\n", encoding="utf-8")
|
||||
definitions = tmp_path / "src" / "hooks" / "types.py"
|
||||
definitions.parent.mkdir()
|
||||
definitions.write_text("HookEvent.DEFINITION_ONLY\n", encoding="utf-8")
|
||||
reads = []
|
||||
original = Path.read_text
|
||||
|
||||
def read(path, *args, **kwargs):
|
||||
reads.append(path)
|
||||
return original(path, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(Path, "read_text", read)
|
||||
assert _product_event_references(tmp_path) == {"BEFORE_TOOL", "ON_INIT"}
|
||||
assert reads == [product]
|
||||
|
||||
|
||||
def test_no_new_hook_event_is_declared_without_an_emission_site():
|
||||
"""Every canonical HookEvent must be referenced by non-test product code."""
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
# tests/unit/hooks/<file> -> repo root
|
||||
root = Path(__file__).resolve().parents[5]
|
||||
|
||||
def _refs(name):
|
||||
hits = []
|
||||
for prefix in ("HookEvent.", "PluginHook."):
|
||||
hits += subprocess.run(
|
||||
["grep", "-rn", "--include=*.py", f"{prefix}{name}", str(root)],
|
||||
capture_output=True, text=True,
|
||||
).stdout.splitlines()
|
||||
return [
|
||||
h for h in hits
|
||||
if "/tests/" not in h and "hooks/types.py" not in h
|
||||
]
|
||||
references = _product_event_references(root)
|
||||
|
||||
# Control probe (must be non-zero): a known-live event.
|
||||
assert _refs("BEFORE_TOOL"), "control probe failed - the grep found nothing for a live event"
|
||||
assert "BEFORE_TOOL" in references, "control probe failed - no reference to a live event"
|
||||
# Control probe (must be zero): a name that does not exist.
|
||||
assert not _refs("DEFINITELY_NOT_AN_EVENT"), "control probe failed - grep matched a bogus name"
|
||||
assert "DEFINITELY_NOT_AN_EVENT" not in references, "control probe matched a bogus name"
|
||||
|
||||
dead = sorted(m.name for m in HookEvent if not _refs(m.name))
|
||||
dead = sorted(m.name for m in HookEvent if m.name not in references)
|
||||
unexpected = sorted(set(dead) - KNOWN_DEAD_OUT_OF_SCOPE)
|
||||
assert unexpected == [], (
|
||||
"HookEvent members declared with no emission site (registering on one "
|
||||
|
||||
@@ -238,6 +238,102 @@ class TestStructuredRetryAfter:
|
||||
# No numeric header, no message pattern -> None
|
||||
assert extract_retry_after(_ProviderError("please slow down")) is None
|
||||
|
||||
def test_nan_header_is_ignored(self):
|
||||
"""A ``nan`` Retry-After header must not poison the delay (issue #5424)."""
|
||||
class _ProviderError(Exception):
|
||||
headers = {"retry-after": "nan"}
|
||||
# No usable header, no message pattern -> None (not NaN)
|
||||
assert extract_retry_after(_ProviderError("rate limited")) is None
|
||||
|
||||
def test_infinite_header_is_ignored(self):
|
||||
class _ProviderError(Exception):
|
||||
headers = {"retry-after": "inf"}
|
||||
assert extract_retry_after(_ProviderError("rate limited")) is None
|
||||
|
||||
def test_negative_header_is_ignored(self):
|
||||
class _ProviderError(Exception):
|
||||
headers = {"retry-after": "-5"}
|
||||
assert extract_retry_after(_ProviderError("rate limited")) is None
|
||||
|
||||
def test_zero_header_is_preserved(self):
|
||||
class _ProviderError(Exception):
|
||||
headers = {"retry-after": "0"}
|
||||
assert extract_retry_after(_ProviderError("rate limited")) == 0.0
|
||||
|
||||
def test_fractional_header_is_preserved(self):
|
||||
class _ProviderError(Exception):
|
||||
headers = {"retry-after": "1.5"}
|
||||
assert extract_retry_after(_ProviderError("rate limited")) == 1.5
|
||||
|
||||
def test_invalid_header_falls_through_to_message(self):
|
||||
"""An invalid high-priority header must not suppress a valid message delay."""
|
||||
class _ProviderError(Exception):
|
||||
headers = {"retry-after": "nan"}
|
||||
assert extract_retry_after(_ProviderError("retry after 30 seconds")) == 30.0
|
||||
|
||||
def test_nan_attribute_is_ignored(self):
|
||||
import math as _math
|
||||
class _ProviderError(Exception):
|
||||
retry_after = _math.nan
|
||||
assert extract_retry_after(_ProviderError("rate limited")) is None
|
||||
|
||||
def test_infinite_attribute_is_ignored(self):
|
||||
import math as _math
|
||||
class _ProviderError(Exception):
|
||||
retry_after = _math.inf
|
||||
assert extract_retry_after(_ProviderError("rate limited")) is None
|
||||
|
||||
def test_negative_attribute_is_ignored(self):
|
||||
class _ProviderError(Exception):
|
||||
retry_after = -3.0
|
||||
assert extract_retry_after(_ProviderError("rate limited")) is None
|
||||
|
||||
def test_invalid_attribute_falls_through_to_message(self):
|
||||
import math as _math
|
||||
class _ProviderError(Exception):
|
||||
retry_after = _math.nan
|
||||
assert extract_retry_after(_ProviderError("retry after 12 seconds")) == 12.0
|
||||
|
||||
def test_negative_message_hint_is_ignored(self):
|
||||
"""An echoed negative Retry-After in the message must not become a
|
||||
positive delay that would park a scheduler hold (issue #5424)."""
|
||||
assert extract_retry_after(Exception("retry after -3600 seconds")) is None
|
||||
assert extract_retry_after(Exception("Retry-After: -60")) is None
|
||||
assert extract_retry_after(Exception("wait -120 seconds")) is None
|
||||
assert extract_retry_after(Exception("-90 second cooldown")) is None
|
||||
|
||||
def test_invalid_hints_fall_through_to_valid_message_delay(self):
|
||||
"""A negative header + echoed negative seconds yields no delay; a later
|
||||
positive number is not spuriously fabricated."""
|
||||
class _ProviderError(Exception):
|
||||
headers = {"retry-after": "-5"}
|
||||
assert extract_retry_after(_ProviderError("retry after -3600 seconds")) is None
|
||||
|
||||
def test_loose_intervening_word_phrasing_is_not_misread(self):
|
||||
"""Phrasings with words between the number and "second" are intentionally
|
||||
NOT honoured: a greedy match would otherwise read a status code (e.g. the
|
||||
``429`` in "429 Too Many Requests retry in 30 more seconds") as the delay
|
||||
and park the job far past the real reset window. Such messages fall
|
||||
through to None rather than producing a bogus positive delay."""
|
||||
assert extract_retry_after(Exception("retry in 30 more seconds")) is None
|
||||
assert (
|
||||
extract_retry_after(
|
||||
Exception("429 Too Many Requests retry in 30 more seconds")
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
def test_fractional_message_second_delay_preserved(self):
|
||||
"""A fractional "N second" message delay is preserved, not truncated."""
|
||||
assert extract_retry_after(Exception("cool down 2.5 seconds")) == 2.5
|
||||
|
||||
def test_later_valid_hint_not_shadowed_by_earlier_negative(self):
|
||||
"""An earlier echoed negative hint must not shadow a later valid one in
|
||||
the same message (regression: re.search only saw the first match)."""
|
||||
assert extract_retry_after(
|
||||
Exception("retry after -5 seconds; retry after 30 seconds")
|
||||
) == 30.0
|
||||
|
||||
|
||||
class TestRetryLogic:
|
||||
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Native failed Responses must surface without a second provider request."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from openai.types.responses import ResponseError
|
||||
from openai.types.chat import ChatCompletionChunk
|
||||
|
||||
from praisonaiagents.llm.llm import LLMResponseError
|
||||
from praisonaiagents.llm.openai_client import OpenAIClient
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('mode', ['sync', 'async', 'sync_stream', 'async_stream'])
|
||||
@pytest.mark.parametrize('shape', ['sdk', 'dict'])
|
||||
@pytest.mark.parametrize('profile', ['failed', 'missing_error', 'completed', 'unsupported'])
|
||||
async def test_native_responses_failure_and_compatibility_fallback(mode, shape, profile, monkeypatch):
|
||||
monkeypatch.delenv('OPENAI_API_BASE', raising=False)
|
||||
monkeypatch.delenv('OPENAI_BASE_URL', raising=False)
|
||||
calls = []
|
||||
error = ResponseError(code='server_error', message='Provider failed.') if profile == 'failed' else None
|
||||
fields = {
|
||||
'status': 'failed' if profile in ('failed', 'missing_error') else 'completed',
|
||||
'error': error if shape == 'sdk' or error is None else error.model_dump(),
|
||||
'output': [{'type': 'message', 'content': [{'type': 'output_text', 'text': 'answer'}]}],
|
||||
'usage': {'input_tokens': 3, 'output_tokens': 2, 'total_tokens': 5},
|
||||
}
|
||||
if shape == 'sdk':
|
||||
fields['usage'] = SimpleNamespace(**fields['usage'])
|
||||
raw = SimpleNamespace(**fields) if shape == 'sdk' else fields
|
||||
|
||||
def respond(**kwargs):
|
||||
calls.append('responses')
|
||||
if profile == 'unsupported':
|
||||
raise NotImplementedError('Responses unsupported')
|
||||
return raw
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return respond(**kwargs)
|
||||
|
||||
fallback = SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content='fallback'))])
|
||||
|
||||
def chat(**kwargs):
|
||||
calls.append('chat')
|
||||
if kwargs.get('stream'):
|
||||
return iter([ChatCompletionChunk(
|
||||
id='chat-test', object='chat.completion.chunk', created=0, model='gpt-4o-mini',
|
||||
choices=[{'index': 0, 'delta': {'content': 'fallback'}, 'finish_reason': 'stop'}],
|
||||
)])
|
||||
return fallback
|
||||
|
||||
async def achat(**kwargs):
|
||||
result = chat(**kwargs)
|
||||
if kwargs.get('stream'):
|
||||
async def chunks():
|
||||
for chunk in result:
|
||||
yield chunk
|
||||
return chunks()
|
||||
return result
|
||||
|
||||
client = OpenAIClient(api_key='sk-test-not-real')
|
||||
client._sync_client = SimpleNamespace(
|
||||
responses=SimpleNamespace(create=respond), chat=SimpleNamespace(completions=SimpleNamespace(create=chat)),
|
||||
)
|
||||
client._async_client = SimpleNamespace(
|
||||
responses=SimpleNamespace(create=arespond), chat=SimpleNamespace(completions=SimpleNamespace(create=achat)),
|
||||
)
|
||||
|
||||
async def invoke():
|
||||
messages = [{'role': 'user', 'content': 'question'}]
|
||||
if mode == 'sync_stream':
|
||||
return client.process_stream_response(messages, model='gpt-4o-mini')
|
||||
if mode == 'async_stream':
|
||||
return await client.process_stream_response_async(messages, model='gpt-4o-mini')
|
||||
return client.create_completion(messages) if mode == 'sync' else await client.acreate_completion(messages)
|
||||
|
||||
if profile in ('failed', 'missing_error'):
|
||||
with pytest.raises(LLMResponseError, match='Responses API failed') as caught:
|
||||
await invoke()
|
||||
if error:
|
||||
assert 'server_error' in str(caught.value)
|
||||
assert 'Provider failed.' in str(caught.value)
|
||||
assert calls == ['responses']
|
||||
else:
|
||||
result = await invoke()
|
||||
assert result.choices[0].message.content == ('fallback' if profile == 'unsupported' else 'answer')
|
||||
assert calls == (['responses', 'chat'] if profile == 'unsupported' else ['responses'])
|
||||
if profile == 'completed':
|
||||
assert result.usage is not None
|
||||
assert result.usage.prompt_tokens == 3
|
||||
assert result.usage.completion_tokens == 2
|
||||
assert result.usage.total_tokens == 5
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Provider additions must belong to the request rather than its caller."""
|
||||
|
||||
from copy import deepcopy
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
|
||||
|
||||
def memory_tool():
|
||||
return SimpleNamespace(
|
||||
get_tool_definition=lambda: {'type': 'memory_20250818', 'name': 'memory'},
|
||||
get_beta_header=lambda: 'context-management-2025-06-27',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('source', ['override', 'settings'])
|
||||
@pytest.mark.parametrize('web_fetch', [False, True])
|
||||
@pytest.mark.parametrize('memory', [False, True])
|
||||
def test_provider_tool_additions_do_not_mutate_inputs(source, web_fetch, memory):
|
||||
tools = [{'type': 'function', 'function': {'name': 'local_tool'}}]
|
||||
original = deepcopy(tools)
|
||||
llm = LLM(model='anthropic/claude-sonnet-4-5', web_fetch=web_fetch,
|
||||
**({'tools': tools} if source == 'settings' else {}))
|
||||
llm._supports_claude_memory = lambda: memory
|
||||
llm._get_claude_memory_tool = memory_tool
|
||||
overrides = {'tools': tools} if source == 'override' else {}
|
||||
first = llm._build_completion_params(**overrides)
|
||||
first_snapshot = deepcopy(first['tools'])
|
||||
second = llm._build_completion_params(**overrides)
|
||||
assert tools == original
|
||||
assert first['tools'] == first_snapshot
|
||||
assert second['tools'] == first_snapshot
|
||||
assert len(second['tools']) == 1 + int(web_fetch) + int(memory)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('source', ['override', 'settings'])
|
||||
def test_memory_beta_header_does_not_mutate_input_headers(source):
|
||||
headers = {'x-custom': 'keep'}
|
||||
llm = LLM(model='anthropic/claude-sonnet-4-5',
|
||||
**({'extra_headers': headers} if source == 'settings' else {}))
|
||||
llm._supports_claude_memory = lambda: True
|
||||
llm._get_claude_memory_tool = memory_tool
|
||||
result = llm._build_completion_params(
|
||||
**({'extra_headers': headers} if source == 'override' else {}),
|
||||
)
|
||||
assert headers == {'x-custom': 'keep'}
|
||||
assert result['extra_headers'] == {
|
||||
'x-custom': 'keep', 'anthropic-beta': 'context-management-2025-06-27',
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
"""Provider failures cannot become successful Responses answers."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from openai.types.responses import ResponseError
|
||||
|
||||
from praisonaiagents.llm.llm import LLM, LLMResponseError
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('entry', ['sync', 'async', 'stream', 'astream'])
|
||||
@pytest.mark.parametrize('shape', ['sdk', 'dict'])
|
||||
@pytest.mark.parametrize('profile', ['failed', 'missing_error', 'completed'])
|
||||
async def test_public_responses_failure(entry, shape, profile, monkeypatch):
|
||||
error = ResponseError(code='server_error', message='Provider failed.') if profile == 'failed' else None
|
||||
status = 'completed' if profile == 'completed' else 'failed'
|
||||
fields = {
|
||||
'status': status,
|
||||
'error': error if shape == 'sdk' or error is None else error.model_dump(),
|
||||
'output': [{'type': 'message', 'content': [{'type': 'output_text', 'text': 'partial'}]}],
|
||||
}
|
||||
response = SimpleNamespace(**fields) if shape == 'sdk' else fields
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
requests = []
|
||||
|
||||
def respond(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return response
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return respond(**kwargs)
|
||||
|
||||
llm._call_responses_api = respond
|
||||
llm._call_responses_api_async = arespond
|
||||
if entry in ('stream', 'astream'):
|
||||
import litellm
|
||||
monkeypatch.setattr('praisonaiagents.llm.llm.check_model_request', lambda *args: None)
|
||||
events = [
|
||||
{'type': 'response.output_text.delta', 'delta': 'partial'},
|
||||
{'type': 'response.' + status, 'response': response},
|
||||
]
|
||||
if shape == 'sdk':
|
||||
events = [SimpleNamespace(**event) for event in events]
|
||||
|
||||
def stream(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return iter(events)
|
||||
|
||||
async def astream(**kwargs):
|
||||
requests.append(kwargs)
|
||||
|
||||
async def iterate():
|
||||
for event in events:
|
||||
yield event
|
||||
|
||||
return iterate()
|
||||
|
||||
monkeypatch.setattr(litellm, 'responses', stream)
|
||||
monkeypatch.setattr(litellm, 'aresponses', astream)
|
||||
|
||||
async def invoke():
|
||||
kwargs = dict(stream=entry in ('stream', 'astream'), verbose=False)
|
||||
return await llm.get_response_async('question', **kwargs) if entry in ('async', 'astream') else llm.get_response('question', **kwargs)
|
||||
|
||||
if profile == 'completed':
|
||||
assert await invoke() == 'partial'
|
||||
assert llm._last_stop_reason == 'completed'
|
||||
else:
|
||||
with pytest.raises(LLMResponseError, match='Responses API failed') as caught:
|
||||
await invoke()
|
||||
assert llm._last_stop_reason == 'error'
|
||||
if error:
|
||||
assert 'server_error' in str(caught.value)
|
||||
assert 'Provider failed.' in str(caught.value)
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('mode', ['sync', 'async'])
|
||||
@pytest.mark.parametrize('shape', ['sdk', 'dict'])
|
||||
async def test_stream_error_event_prevents_pending_tool_execution(mode, shape, monkeypatch):
|
||||
import litellm
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
effects = []
|
||||
requests = []
|
||||
|
||||
def lookup() -> str:
|
||||
"""Return a lookup result."""
|
||||
effects.append('executed')
|
||||
return 'found'
|
||||
|
||||
events = [
|
||||
{'type': 'response.output_item.done', 'output_index': 0, 'item': {
|
||||
'type': 'function_call', 'call_id': 'call_1', 'name': 'lookup', 'arguments': '{}',
|
||||
}},
|
||||
{'type': 'error', 'code': 'server_error', 'message': 'Stream failed.'},
|
||||
]
|
||||
if shape == 'sdk':
|
||||
events = [SimpleNamespace(**event) for event in events]
|
||||
monkeypatch.setattr('praisonaiagents.llm.llm.check_model_request', lambda *args: None)
|
||||
|
||||
def stream(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return iter(events)
|
||||
|
||||
async def astream(**kwargs):
|
||||
requests.append(kwargs)
|
||||
|
||||
async def iterate():
|
||||
for event in events:
|
||||
yield event
|
||||
|
||||
return iterate()
|
||||
|
||||
monkeypatch.setattr(litellm, 'responses', stream)
|
||||
monkeypatch.setattr(litellm, 'aresponses', astream)
|
||||
with pytest.raises(LLMResponseError, match='Stream failed.'):
|
||||
kwargs = dict(stream=True, verbose=False, tools=[lookup], max_iterations=1)
|
||||
if mode == 'sync':
|
||||
llm.get_response('question', **kwargs)
|
||||
else:
|
||||
await llm.get_response_async('question', **kwargs)
|
||||
assert effects == []
|
||||
assert len(requests) == 1
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Incomplete Responses retain their terminal reason and streamed usage."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from openai.types.responses.response import IncompleteDetails
|
||||
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
from praisonaiagents.llm.openai_client import OpenAIClient
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolate_responses_endpoint(monkeypatch):
|
||||
monkeypatch.delenv('OPENAI_API_BASE', raising=False)
|
||||
monkeypatch.delenv('OPENAI_BASE_URL', raising=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('entry', ['llm_sync', 'llm_async', 'llm_stream', 'llm_astream', 'client_sync', 'client_async'])
|
||||
@pytest.mark.parametrize('shape', ['sdk', 'dict'])
|
||||
@pytest.mark.parametrize('status,reason,expected_finish,expected_outcome', [
|
||||
('incomplete', 'max_output_tokens', 'length', 'length_truncated'),
|
||||
('incomplete', 'content_filter', 'content_filter', 'content_filtered'),
|
||||
('completed', 'max_output_tokens', 'stop', 'completed'),
|
||||
])
|
||||
async def test_public_incomplete_responses(entry, shape, status, reason, expected_finish, expected_outcome, monkeypatch):
|
||||
details = IncompleteDetails(reason=reason)
|
||||
response = SimpleNamespace(
|
||||
status=status, incomplete_details=details if shape == 'sdk' else details.model_dump(),
|
||||
output=[{'type': 'message', 'content': [{'type': 'output_text', 'text': 'partial'}]}],
|
||||
usage=(
|
||||
SimpleNamespace(input_tokens=3, output_tokens=2, total_tokens=5)
|
||||
if shape == 'sdk' else {'input_tokens': 3, 'output_tokens': 2, 'total_tokens': 5}
|
||||
),
|
||||
)
|
||||
requests = []
|
||||
|
||||
def respond(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return response
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return respond(**kwargs)
|
||||
|
||||
if entry.startswith('llm'):
|
||||
import litellm
|
||||
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
monkeypatch.setattr('praisonaiagents.llm.llm.check_model_request', lambda *args: None)
|
||||
monkeypatch.setattr(litellm, 'responses', respond)
|
||||
monkeypatch.setattr(litellm, 'aresponses', arespond)
|
||||
if entry in ('llm_stream', 'llm_astream'):
|
||||
events = [
|
||||
{'type': 'response.output_text.delta', 'delta': 'partial'},
|
||||
{'type': 'response.' + status, 'response': response},
|
||||
]
|
||||
if shape == 'sdk':
|
||||
events = [SimpleNamespace(**event) for event in events]
|
||||
|
||||
def stream(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return iter(events)
|
||||
|
||||
async def astream(**kwargs):
|
||||
requests.append(kwargs)
|
||||
|
||||
async def iterate():
|
||||
for event in events:
|
||||
yield event
|
||||
|
||||
return iterate()
|
||||
|
||||
monkeypatch.setattr(litellm, 'responses', stream)
|
||||
monkeypatch.setattr(litellm, 'aresponses', astream)
|
||||
kwargs = dict(stream=entry in ('llm_stream', 'llm_astream'), verbose=False)
|
||||
answer = await llm.get_response_async('question', **kwargs) if entry in ('llm_async', 'llm_astream') else llm.get_response('question', **kwargs)
|
||||
assert answer == 'partial'
|
||||
assert llm._last_stop_reason == expected_outcome
|
||||
for metrics in (llm.last_token_metrics, llm.session_token_metrics):
|
||||
assert metrics is not None
|
||||
assert metrics.input_tokens == 3
|
||||
assert metrics.output_tokens == 2
|
||||
assert metrics.total_tokens == 5
|
||||
else:
|
||||
client = OpenAIClient(api_key='sk-test-not-real')
|
||||
client._sync_client = SimpleNamespace(responses=SimpleNamespace(create=respond))
|
||||
client._async_client = SimpleNamespace(responses=SimpleNamespace(create=arespond))
|
||||
messages = [{'role': 'user', 'content': 'question'}]
|
||||
answer = client.create_completion(messages) if entry == 'client_sync' else await client.acreate_completion(messages)
|
||||
assert answer.choices[0].message.content == 'partial'
|
||||
assert answer.choices[0].finish_reason == expected_finish
|
||||
assert answer.usage.prompt_tokens == 3
|
||||
assert answer.usage.completion_tokens == 2
|
||||
assert answer.usage.total_tokens == 5
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize('reason', ['max_output_tokens', 'content_filter'])
|
||||
@pytest.mark.parametrize('existing', ['max_steps', 'cancelled'])
|
||||
def test_incomplete_response_preserves_stronger_terminal_reason(reason, existing):
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
llm._last_stop_reason = existing
|
||||
llm._extract_from_responses_output(SimpleNamespace(
|
||||
status='incomplete', incomplete_details=IncompleteDetails(reason=reason), output=[],
|
||||
))
|
||||
assert llm._last_stop_reason == existing
|
||||
|
||||
|
||||
@pytest.mark.parametrize('reason', [None, 'unknown_future_reason'])
|
||||
def test_unrecognised_incomplete_reason_preserves_existing_behavior(reason):
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
llm._extract_from_responses_output(SimpleNamespace(
|
||||
status='incomplete', incomplete_details={'reason': reason}, output=[],
|
||||
))
|
||||
assert llm._last_stop_reason == 'completed'
|
||||
@@ -0,0 +1,481 @@
|
||||
"""Responses refusals retain their provider signal across public call paths."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal, ResponseOutputText
|
||||
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
from praisonaiagents.llm.openai_client import OpenAIClient
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolate_responses_endpoint(monkeypatch):
|
||||
"""Fake Responses clients must not inherit a configured chat endpoint."""
|
||||
monkeypatch.delenv('OPENAI_API_BASE', raising=False)
|
||||
monkeypatch.delenv('OPENAI_BASE_URL', raising=False)
|
||||
|
||||
|
||||
def output_response(shape, profile):
|
||||
blocks = []
|
||||
if profile != 'refusal':
|
||||
blocks.append(ResponseOutputText(type='output_text', text='answer', annotations=[]))
|
||||
if profile != 'normal':
|
||||
blocks.append(ResponseOutputRefusal(type='refusal', refusal='Cannot comply.'))
|
||||
message = ResponseOutputMessage(id='msg_test', type='message', role='assistant', status='completed', content=blocks)
|
||||
return SimpleNamespace(output=[message if shape == 'sdk' else message.model_dump()])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('entry', ['llm_sync', 'llm_async', 'llm_stream', 'llm_astream', 'client_sync', 'client_async'])
|
||||
@pytest.mark.parametrize('shape', ['sdk', 'dict'])
|
||||
@pytest.mark.parametrize('profile', ['refusal', 'mixed', 'normal'])
|
||||
async def test_public_responses_refusal_signal(entry, shape, profile, monkeypatch):
|
||||
response = output_response(shape, profile)
|
||||
requests = []
|
||||
|
||||
def respond(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return response
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return respond(**kwargs)
|
||||
|
||||
expected_text = '' if profile == 'refusal' else 'answer'
|
||||
if entry.startswith('llm'):
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
llm._call_responses_api = respond
|
||||
llm._call_responses_api_async = arespond
|
||||
if entry in ('llm_stream', 'llm_astream'):
|
||||
import litellm
|
||||
monkeypatch.setattr('praisonaiagents.llm.llm.check_model_request', lambda *args: None)
|
||||
events = []
|
||||
if expected_text:
|
||||
events.append({'type': 'response.output_text.delta', 'delta': expected_text})
|
||||
events.append({'type': 'response.completed', 'response': response})
|
||||
if shape == 'sdk':
|
||||
events = [SimpleNamespace(**event) for event in events]
|
||||
|
||||
def stream(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return iter(events)
|
||||
|
||||
async def astream(**kwargs):
|
||||
requests.append(kwargs)
|
||||
|
||||
async def iterate():
|
||||
for event in events:
|
||||
yield event
|
||||
|
||||
return iterate()
|
||||
|
||||
monkeypatch.setattr(litellm, 'responses', stream)
|
||||
monkeypatch.setattr(litellm, 'aresponses', astream)
|
||||
kwargs = dict(stream=entry in ('llm_stream', 'llm_astream'), verbose=False)
|
||||
answer = await llm.get_response_async('question', **kwargs) if entry in ('llm_async', 'llm_astream') else llm.get_response('question', **kwargs)
|
||||
assert answer == expected_text
|
||||
assert llm._last_stop_reason == ('completed' if profile == 'normal' else 'refused')
|
||||
else:
|
||||
client = OpenAIClient(api_key='sk-test-not-real')
|
||||
client._sync_client = SimpleNamespace(responses=SimpleNamespace(create=respond))
|
||||
client._async_client = SimpleNamespace(responses=SimpleNamespace(create=arespond))
|
||||
messages = [{'role': 'user', 'content': 'question'}]
|
||||
answer = client.create_completion(messages) if entry == 'client_sync' else await client.acreate_completion(messages)
|
||||
assert answer.choices[0].message.content == (expected_text or None)
|
||||
assert answer.choices[0].message.refusal == (None if profile == 'normal' else 'Cannot comply.')
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize('reason', ['max_steps', 'cancelled'])
|
||||
@pytest.mark.parametrize('incomplete_reason', [None, 'max_output_tokens', 'content_filter'])
|
||||
def test_refusal_does_not_downgrade_existing_terminal_reason(reason, incomplete_reason):
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
llm._last_stop_reason = reason
|
||||
response = output_response('sdk', 'refusal')
|
||||
if incomplete_reason:
|
||||
response.status = 'incomplete'
|
||||
response.incomplete_details = SimpleNamespace(reason=incomplete_reason)
|
||||
llm._extract_from_responses_output(response)
|
||||
assert llm._last_stop_reason == reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('mode', ['sync', 'async'])
|
||||
@pytest.mark.parametrize('shape', ['sdk', 'dict'])
|
||||
async def test_public_refusal_delta_without_completed_event(mode, shape, monkeypatch):
|
||||
import litellm
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
monkeypatch.setattr('praisonaiagents.llm.llm.check_model_request', lambda *args: None)
|
||||
event = {'type': 'response.refusal.delta', 'delta': 'Cannot comply.'}
|
||||
if shape == 'sdk':
|
||||
event = SimpleNamespace(**event)
|
||||
|
||||
def stream(**kwargs):
|
||||
return iter([event])
|
||||
|
||||
async def astream(**kwargs):
|
||||
async def iterate():
|
||||
yield event
|
||||
return iterate()
|
||||
|
||||
monkeypatch.setattr(litellm, 'responses', stream)
|
||||
monkeypatch.setattr(litellm, 'aresponses', astream)
|
||||
kwargs = dict(stream=True, verbose=False)
|
||||
answer = llm.get_response('question', **kwargs) if mode == 'sync' else await llm.get_response_async('question', **kwargs)
|
||||
assert answer == ''
|
||||
assert llm._last_stop_reason == 'refused'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('mode', ['sync', 'async'])
|
||||
async def test_following_normal_turn_clears_refusal(mode):
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
responses = iter([output_response('sdk', 'refusal'), output_response('sdk', 'normal')])
|
||||
|
||||
def respond(**kwargs):
|
||||
return next(responses)
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return respond(**kwargs)
|
||||
|
||||
llm._call_responses_api = respond
|
||||
llm._call_responses_api_async = arespond
|
||||
for expected_text, expected_reason in [('', 'refused'), ('answer', 'completed')]:
|
||||
kwargs = dict(stream=False, verbose=False)
|
||||
answer = llm.get_response('question', **kwargs) if mode == 'sync' else await llm.get_response_async('question', **kwargs)
|
||||
assert answer == expected_text
|
||||
assert llm._last_stop_reason == expected_reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('mode', ['sync', 'async'])
|
||||
@pytest.mark.parametrize('stream', [False, True])
|
||||
@pytest.mark.parametrize('shape', ['sdk', 'dict'])
|
||||
@pytest.mark.parametrize('final_profile', ['normal', 'refusal'])
|
||||
@pytest.mark.parametrize('incomplete_reason', [None, 'max_output_tokens', 'content_filter', 'unknown'])
|
||||
async def test_tool_continuation_uses_final_response_refusal(mode, stream, shape, final_profile, incomplete_reason, monkeypatch):
|
||||
import litellm
|
||||
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
tool = {'type': 'function_call', 'id': 'fc_test', 'call_id': 'call_test',
|
||||
'name': 'lookup', 'arguments': '{}'}
|
||||
tool_item = SimpleNamespace(**tool) if shape == 'sdk' else tool
|
||||
first = output_response(shape, 'refusal')
|
||||
first.output.append(tool_item)
|
||||
if incomplete_reason:
|
||||
first.status = 'incomplete'
|
||||
first.incomplete_details = SimpleNamespace(reason=incomplete_reason)
|
||||
if shape == 'dict':
|
||||
first = {'output': first.output, 'status': first.status,
|
||||
'incomplete_details': {'reason': incomplete_reason}}
|
||||
final = output_response(shape, final_profile)
|
||||
responses = iter([first, final])
|
||||
requests = []
|
||||
dispatched = []
|
||||
monkeypatch.setattr('praisonaiagents.llm.llm.check_model_request', lambda *args: None)
|
||||
|
||||
def respond(**kwargs):
|
||||
requests.append(kwargs)
|
||||
return next(responses)
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return respond(**kwargs)
|
||||
|
||||
def events(**kwargs):
|
||||
response = respond(**kwargs)
|
||||
result = []
|
||||
if response is first or final_profile == 'refusal':
|
||||
result.append({'type': 'response.refusal.delta', 'delta': 'Cannot comply.'})
|
||||
else:
|
||||
result.append({'type': 'response.output_text.delta', 'delta': 'answer'})
|
||||
if response is first:
|
||||
result.append({'type': 'response.output_item.done', 'output_index': 1, 'item': tool_item})
|
||||
event_type = 'response.incomplete' if response is first and incomplete_reason else 'response.completed'
|
||||
result.append({'type': event_type, 'response': response})
|
||||
return [SimpleNamespace(**event) for event in result] if shape == 'sdk' else result
|
||||
|
||||
def stream_response(**kwargs):
|
||||
return iter(events(**kwargs))
|
||||
|
||||
async def astream_response(**kwargs):
|
||||
batch = events(**kwargs)
|
||||
|
||||
async def iterate():
|
||||
for event in batch:
|
||||
yield event
|
||||
|
||||
return iterate()
|
||||
|
||||
def execute(name, arguments, *args, **kwargs):
|
||||
dispatched.append((name, arguments))
|
||||
return 'lookup result'
|
||||
|
||||
llm._call_responses_api = respond
|
||||
llm._call_responses_api_async = arespond
|
||||
monkeypatch.setattr(litellm, 'responses', stream_response)
|
||||
monkeypatch.setattr(litellm, 'aresponses', astream_response)
|
||||
kwargs = dict(stream=stream, verbose=False, execute_tool_fn=execute)
|
||||
answer = llm.get_response('question', **kwargs) if mode == 'sync' else await llm.get_response_async('question', **kwargs)
|
||||
assert answer == ('answer' if final_profile == 'normal' else '')
|
||||
assert dispatched == [('lookup', {})]
|
||||
assert len(requests) == 2
|
||||
assert any(item.get('type') == 'function_call_output' for item in requests[1]['input'])
|
||||
expected_reason = {'max_output_tokens': 'length_truncated', 'content_filter': 'content_filtered'}.get(incomplete_reason)
|
||||
assert llm._last_stop_reason == (expected_reason or ('completed' if final_profile == 'normal' else 'refused'))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("profile, expected_reason", [
|
||||
("refusal", "refused"),
|
||||
("max_output_tokens", "length_truncated"),
|
||||
("content_filter", "content_filtered"),
|
||||
])
|
||||
@pytest.mark.parametrize("interleaving", ["after_return", "overlapping_provider"])
|
||||
async def test_shared_llm_keeps_each_tasks_responses_outcome(profile, expected_reason, interleaving):
|
||||
import asyncio
|
||||
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
terminal_returned = asyncio.Event()
|
||||
normal_dispatched = asyncio.Event()
|
||||
normal_returned = asyncio.Event()
|
||||
terminal = output_response('sdk', 'refusal')
|
||||
if profile != "refusal":
|
||||
terminal.status = 'incomplete'
|
||||
terminal.incomplete_details = SimpleNamespace(reason=profile)
|
||||
normal = output_response('sdk', 'normal')
|
||||
|
||||
async def respond(**kwargs):
|
||||
if asyncio.current_task().get_name() == "terminal-request":
|
||||
if interleaving == "overlapping_provider":
|
||||
await normal_dispatched.wait()
|
||||
return terminal
|
||||
normal_dispatched.set()
|
||||
if interleaving == "overlapping_provider":
|
||||
await terminal_returned.wait()
|
||||
return normal
|
||||
|
||||
llm._call_responses_api_async = respond
|
||||
|
||||
async def terminal_call():
|
||||
answer = await llm.get_response_async('terminal', stream=False, verbose=False)
|
||||
terminal_returned.set()
|
||||
await normal_returned.wait()
|
||||
return answer, llm._last_stop_reason
|
||||
|
||||
async def normal_call():
|
||||
if interleaving == "after_return":
|
||||
await terminal_returned.wait()
|
||||
answer = await llm.get_response_async('normal', stream=False, verbose=False)
|
||||
normal_returned.set()
|
||||
return answer, llm._last_stop_reason
|
||||
|
||||
terminal_task = asyncio.create_task(terminal_call(), name="terminal-request")
|
||||
normal_task = asyncio.create_task(normal_call(), name="normal-request")
|
||||
terminal_result, normal_result = await asyncio.wait_for(
|
||||
asyncio.gather(terminal_task, normal_task), timeout=10)
|
||||
assert terminal_result == ('', expected_reason)
|
||||
assert normal_result == ('answer', 'completed')
|
||||
|
||||
|
||||
def test_copied_llm_has_an_independent_terminal_outcome():
|
||||
import copy
|
||||
|
||||
original = LLM(model='gpt-4o-mini')
|
||||
original._last_stop_reason = 'refused'
|
||||
clone = copy.deepcopy(original)
|
||||
assert clone._last_stop_reason == 'completed'
|
||||
clone._last_stop_reason = 'content_filtered'
|
||||
assert original._last_stop_reason == 'refused'
|
||||
assert clone._last_stop_reason == 'content_filtered'
|
||||
|
||||
|
||||
def test_response_outcome_remains_visible_after_asyncio_run():
|
||||
import asyncio
|
||||
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
|
||||
async def respond(**kwargs):
|
||||
return output_response('sdk', 'refusal')
|
||||
|
||||
llm._call_responses_api_async = respond
|
||||
assert asyncio.run(llm.get_response_async('question', stream=False, verbose=False)) == ''
|
||||
assert llm._last_stop_reason == 'refused'
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile, expected_reason", [
|
||||
("refusal", "refused"),
|
||||
("max_output_tokens", "length_truncated"),
|
||||
("content_filter", "content_filtered"),
|
||||
])
|
||||
def test_shared_llm_keeps_each_threads_responses_outcome(profile, expected_reason):
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
terminal_returned = threading.Event()
|
||||
normal_returned = threading.Event()
|
||||
terminal = output_response('sdk', 'refusal')
|
||||
if profile != "refusal":
|
||||
terminal.status = 'incomplete'
|
||||
terminal.incomplete_details = SimpleNamespace(reason=profile)
|
||||
|
||||
def respond(**kwargs):
|
||||
return terminal if threading.current_thread().name.startswith("terminal") else output_response('sdk', 'normal')
|
||||
|
||||
llm._call_responses_api = respond
|
||||
|
||||
def terminal_call():
|
||||
answer = llm.get_response('terminal', stream=False, verbose=False)
|
||||
terminal_returned.set()
|
||||
assert normal_returned.wait(10)
|
||||
return answer, llm._last_stop_reason
|
||||
|
||||
def normal_call():
|
||||
assert terminal_returned.wait(10)
|
||||
answer = llm.get_response('normal', stream=False, verbose=False)
|
||||
normal_returned.set()
|
||||
return answer, llm._last_stop_reason
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1, thread_name_prefix="terminal") as terminals, ThreadPoolExecutor(max_workers=1, thread_name_prefix="normal") as normals:
|
||||
terminal_task = terminals.submit(terminal_call)
|
||||
normal_task = normals.submit(normal_call)
|
||||
assert terminal_task.result(timeout=15) == ('', expected_reason)
|
||||
assert normal_task.result(timeout=15) == ('answer', 'completed')
|
||||
# Observers outside either request retain the latest completed result.
|
||||
assert llm._last_stop_reason == 'completed'
|
||||
|
||||
|
||||
@pytest.mark.parametrize("first_profile, first_reason", [
|
||||
("refusal", "refused"),
|
||||
("max_output_tokens", "length_truncated"),
|
||||
("content_filter", "content_filtered"),
|
||||
])
|
||||
@pytest.mark.parametrize("second_profile, second_reason", [
|
||||
("normal", "completed"),
|
||||
("refusal", "refused"),
|
||||
])
|
||||
def test_later_async_turn_replaces_a_sync_contexts_last_outcome(first_profile, first_reason, second_profile, second_reason):
|
||||
import asyncio
|
||||
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
first = output_response('sdk', 'refusal')
|
||||
if first_profile != "refusal":
|
||||
first.status = 'incomplete'
|
||||
first.incomplete_details = SimpleNamespace(reason=first_profile)
|
||||
|
||||
def respond(**kwargs):
|
||||
return first
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return output_response('sdk', second_profile)
|
||||
|
||||
llm._call_responses_api = respond
|
||||
llm._call_responses_api_async = arespond
|
||||
assert llm.get_response('first', stream=False, verbose=False) == ''
|
||||
assert llm._last_stop_reason == first_reason
|
||||
answer = asyncio.run(llm.get_response_async('second', stream=False, verbose=False))
|
||||
assert answer == ('answer' if second_profile == "normal" else '')
|
||||
assert llm._last_stop_reason == second_reason
|
||||
|
||||
|
||||
@pytest.mark.parametrize('outer_mode', ['sync', 'async'])
|
||||
@pytest.mark.parametrize('inner_mode', ['sync', 'async'])
|
||||
@pytest.mark.parametrize('outer_profile, expected_reason', [
|
||||
('normal', 'completed'),
|
||||
('max_output_tokens', 'length_truncated'),
|
||||
('content_filter', 'content_filtered'),
|
||||
])
|
||||
@pytest.mark.parametrize('inner_profile', ['normal', 'refusal'])
|
||||
@pytest.mark.parametrize('stream', [False, True])
|
||||
def test_nested_tool_request_preserves_the_outer_responses_outcome(
|
||||
outer_mode, inner_mode, outer_profile, expected_reason, inner_profile,
|
||||
stream, monkeypatch,
|
||||
):
|
||||
import asyncio
|
||||
import litellm
|
||||
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
first = output_response('sdk', 'normal')
|
||||
first.output = []
|
||||
first.output.append(SimpleNamespace(
|
||||
type='function_call', id='fc_nested', call_id='call_nested',
|
||||
name='lookup', arguments='{}',
|
||||
))
|
||||
if outer_profile != 'normal':
|
||||
first.status = 'incomplete'
|
||||
first.incomplete_details = SimpleNamespace(reason=outer_profile)
|
||||
outer_responses = iter([first, output_response('sdk', 'normal')])
|
||||
in_tool = False
|
||||
inner_answers = []
|
||||
|
||||
def respond(**kwargs):
|
||||
if in_tool:
|
||||
return output_response('sdk', inner_profile)
|
||||
return next(outer_responses)
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return respond(**kwargs)
|
||||
|
||||
def stream_response(**kwargs):
|
||||
response = respond(**kwargs)
|
||||
events = []
|
||||
if response is first:
|
||||
events.append({'type': 'response.output_item.done',
|
||||
'output_index': 0, 'item': first.output[0]})
|
||||
else:
|
||||
events.append({'type': 'response.output_text.delta', 'delta': 'answer'})
|
||||
events.append({'type': 'response.incomplete' if response is first
|
||||
and outer_profile != 'normal' else 'response.completed',
|
||||
'response': response})
|
||||
return iter(events)
|
||||
|
||||
async def astream_response(**kwargs):
|
||||
async def iterate():
|
||||
for event in stream_response(**kwargs):
|
||||
yield event
|
||||
return iterate()
|
||||
|
||||
def execute(name, arguments, *args, **kwargs):
|
||||
nonlocal in_tool
|
||||
assert name == 'lookup'
|
||||
in_tool = True
|
||||
try:
|
||||
if inner_mode == 'async':
|
||||
answer = asyncio.run(llm.get_response_async(
|
||||
'inner', stream=False, verbose=False))
|
||||
else:
|
||||
answer = llm.get_response('inner', stream=False, verbose=False)
|
||||
inner_answers.append(answer)
|
||||
finally:
|
||||
in_tool = False
|
||||
return 'lookup result'
|
||||
|
||||
async def aexecute(name, arguments, **kwargs):
|
||||
nonlocal in_tool
|
||||
assert name == 'lookup'
|
||||
in_tool = True
|
||||
try:
|
||||
if inner_mode == 'async':
|
||||
answer = await llm.get_response_async(
|
||||
'inner', stream=False, verbose=False)
|
||||
else:
|
||||
answer = llm.get_response('inner', stream=False, verbose=False)
|
||||
inner_answers.append(answer)
|
||||
finally:
|
||||
in_tool = False
|
||||
return 'lookup result'
|
||||
|
||||
llm._call_responses_api = respond
|
||||
llm._call_responses_api_async = arespond
|
||||
monkeypatch.setattr('praisonaiagents.llm.llm.check_model_request', lambda *args: None)
|
||||
monkeypatch.setattr(litellm, 'responses', stream_response)
|
||||
monkeypatch.setattr(litellm, 'aresponses', astream_response)
|
||||
if outer_mode == 'async':
|
||||
answer = asyncio.run(llm.get_response_async(
|
||||
'outer', stream=stream, verbose=False, execute_tool_fn=aexecute))
|
||||
else:
|
||||
answer = llm.get_response(
|
||||
'outer', stream=stream, verbose=False, execute_tool_fn=execute)
|
||||
assert answer == 'answer'
|
||||
assert inner_answers == ['answer' if inner_profile == 'normal' else '']
|
||||
assert llm._last_stop_reason == expected_reason
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Responses streams release their transport on every consumer exit."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.llm.llm import LLM, LLMResponseError
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('mode', ['sync', 'async'])
|
||||
@pytest.mark.parametrize('shape', ['direct', 'litellm'])
|
||||
@pytest.mark.parametrize('exit_kind', ['completed', 'failed', 'callback', 'cancelled'])
|
||||
async def test_responses_stream_releases_transport(mode, shape, exit_kind, monkeypatch):
|
||||
import httpx
|
||||
import litellm
|
||||
if mode == 'sync' and exit_kind == 'cancelled':
|
||||
pytest.skip('Task cancellation is an asynchronous exit')
|
||||
closed = []
|
||||
|
||||
class SyncTransport(httpx.SyncByteStream):
|
||||
def __iter__(self):
|
||||
return iter([])
|
||||
|
||||
def close(self):
|
||||
closed.append('closed')
|
||||
|
||||
class AsyncTransport(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
if False:
|
||||
yield b''
|
||||
|
||||
async def aclose(self):
|
||||
closed.append('closed')
|
||||
|
||||
response = httpx.Response(200, stream=SyncTransport() if mode == 'sync' else AsyncTransport())
|
||||
event = {'type': 'response.output_text.delta', 'delta': 'answer'}
|
||||
if exit_kind == 'failed':
|
||||
event = {'type': 'response.failed', 'response': {'error': {'code': 'server_error', 'message': 'Provider failed.'}}}
|
||||
|
||||
class Stream:
|
||||
def __init__(self):
|
||||
self.events = iter([event])
|
||||
if shape == 'litellm':
|
||||
self.response = response
|
||||
elif mode == 'sync':
|
||||
self.close = response.close
|
||||
else:
|
||||
self.aclose = response.aclose
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
return next(self.events)
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if exit_kind == 'cancelled':
|
||||
raise asyncio.CancelledError()
|
||||
try:
|
||||
return next(self.events)
|
||||
except StopIteration:
|
||||
raise StopAsyncIteration
|
||||
|
||||
stream = Stream()
|
||||
monkeypatch.setattr('praisonaiagents.llm.llm.check_model_request', lambda *args: None)
|
||||
monkeypatch.setattr(litellm, 'responses', lambda **kwargs: stream)
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return stream
|
||||
|
||||
monkeypatch.setattr(litellm, 'aresponses', arespond)
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
|
||||
def callback(event):
|
||||
if exit_kind == 'callback':
|
||||
raise ValueError('Callback failed.')
|
||||
|
||||
async def invoke():
|
||||
kwargs = dict(stream_callback=callback, emit_events=True)
|
||||
if mode == 'sync':
|
||||
return llm._stream_responses_api({'model': llm.model}, **kwargs)
|
||||
return await llm._stream_responses_api_async({'model': llm.model}, **kwargs)
|
||||
|
||||
if exit_kind == 'failed':
|
||||
with pytest.raises(LLMResponseError, match='Provider failed.'):
|
||||
await invoke()
|
||||
elif exit_kind == 'callback':
|
||||
with pytest.raises(ValueError, match='Callback failed.'):
|
||||
await invoke()
|
||||
elif exit_kind == 'cancelled':
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await invoke()
|
||||
else:
|
||||
assert (await invoke())[0] == 'answer'
|
||||
assert response.is_closed
|
||||
assert closed == ['closed']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('mode', ['sync', 'async'])
|
||||
async def test_cleanup_error_does_not_mask_provider_failure(mode, monkeypatch):
|
||||
import litellm
|
||||
event = {'type': 'error', 'message': 'Original failure.'}
|
||||
|
||||
def fail_close():
|
||||
raise ValueError('Cleanup failed.')
|
||||
|
||||
async def fail_aclose():
|
||||
fail_close()
|
||||
|
||||
class AsyncStream:
|
||||
aclose = staticmethod(fail_aclose)
|
||||
|
||||
def __aiter__(self):
|
||||
async def events():
|
||||
yield event
|
||||
return events()
|
||||
|
||||
monkeypatch.setattr('praisonaiagents.llm.llm.check_model_request', lambda *args: None)
|
||||
class SyncStream:
|
||||
close = staticmethod(fail_close)
|
||||
|
||||
def __iter__(self):
|
||||
return iter([event])
|
||||
|
||||
monkeypatch.setattr(litellm, 'responses', lambda **kwargs: SyncStream())
|
||||
|
||||
async def arespond(**kwargs):
|
||||
return AsyncStream()
|
||||
|
||||
monkeypatch.setattr(litellm, 'aresponses', arespond)
|
||||
llm = LLM(model='gpt-4o-mini')
|
||||
with pytest.raises(LLMResponseError, match='Original failure.'):
|
||||
if mode == 'sync':
|
||||
llm._stream_responses_api({'model': llm.model})
|
||||
else:
|
||||
await llm._stream_responses_api_async({'model': llm.model})
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Invalid provider retry hints must not produce invalid delays or holds."""
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.llm.error_classifier import classify_llm_error, extract_retry_after
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["nan", "inf", "-inf", "-1"])
|
||||
def test_invalid_header_is_ignored(value):
|
||||
error = Exception("rate limited")
|
||||
error.headers = {"retry-after": value}
|
||||
assert extract_retry_after(error) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf"), -1.0])
|
||||
def test_invalid_attribute_is_ignored(value):
|
||||
error = Exception("rate limited")
|
||||
error.retry_after = value
|
||||
assert extract_retry_after(error) is None
|
||||
|
||||
|
||||
def test_invalid_header_falls_through_to_valid_attribute():
|
||||
error = Exception("rate limited")
|
||||
error.headers = {"retry-after": "nan"}
|
||||
error.retry_after = 7
|
||||
assert extract_retry_after(error) == 7
|
||||
|
||||
|
||||
def test_invalid_attribute_falls_through_to_valid_message():
|
||||
error = Exception("rate limited; retry after 12 seconds")
|
||||
error.retry_after = float("nan")
|
||||
assert extract_retry_after(error) == 12
|
||||
|
||||
|
||||
def test_nan_hint_does_not_poison_classified_backoff():
|
||||
error = Exception("rate limited")
|
||||
error.headers = {"retry-after": "nan"}
|
||||
result = classify_llm_error(error, provider="openai", model="gpt-4o-mini")
|
||||
assert math.isfinite(result.backoff_seconds)
|
||||
assert result.backoff_seconds >= 0
|
||||
|
||||
|
||||
def test_nan_hint_does_not_poison_scheduler_hold():
|
||||
from praisonaiagents.scheduler.due import quota_hold_from_failure
|
||||
|
||||
error = Exception("rate limited")
|
||||
error.headers = {"retry-after": "nan"}
|
||||
assert quota_hold_from_failure(error, now=1000) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["0", "0.5", "9999"])
|
||||
def test_valid_header_keeps_zero_fraction_and_cap(value):
|
||||
error = Exception("rate limited")
|
||||
error.headers = {"retry-after": value}
|
||||
assert extract_retry_after(error) == min(float(value), 300)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"message",
|
||||
[
|
||||
"rate limited; retry after -1 seconds",
|
||||
"rate limited; retry-after: -3600 seconds",
|
||||
"please wait -5 seconds",
|
||||
],
|
||||
)
|
||||
def test_negative_message_hint_is_ignored(message):
|
||||
assert extract_retry_after(Exception(message)) is None
|
||||
|
||||
|
||||
def test_negative_message_hint_does_not_poison_scheduler_hold():
|
||||
from praisonaiagents.scheduler.due import quota_hold_from_failure
|
||||
|
||||
error = Exception("rate limited; retry after -3600 seconds")
|
||||
assert quota_hold_from_failure(error, now=1000) is None
|
||||
|
||||
|
||||
def test_valid_fractional_message_hint_is_preserved():
|
||||
assert extract_retry_after(Exception("retry after 2.5 seconds")) == 2.5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("token, expected", [
|
||||
("-.5", None), (".5", 0.5), ("+.5", 0.5),
|
||||
("-1e3", None), ("1e309", None), ("1e-1", 0.1),
|
||||
("1.2.3", None), ("+-.5", None), ("0", 0),
|
||||
])
|
||||
@pytest.mark.parametrize("prefix", ["wait ", "retry after ", "retry: ", ""])
|
||||
def test_message_numeric_token_is_not_partially_matched(token, expected, prefix):
|
||||
assert extract_retry_after(Exception(f"rate limited; {prefix}{token} seconds")) == expected
|
||||
|
||||
|
||||
def test_negative_leading_dot_hint_does_not_create_scheduler_hold():
|
||||
from praisonaiagents.scheduler.due import quota_hold_from_failure
|
||||
|
||||
error = Exception("rate limited; wait -.5 seconds")
|
||||
error.headers = {"retry-after": "-.5"}
|
||||
assert quota_hold_from_failure(error, now=1000) is None
|
||||
|
||||
|
||||
def test_message_delay_without_space_before_unit_is_preserved():
|
||||
assert extract_retry_after(Exception("wait 5seconds")) == 5
|
||||
@@ -0,0 +1,124 @@
|
||||
"""The retry loop must consume complete numeric hints like the shared parser."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
|
||||
|
||||
@pytest.mark.parametrize("template", [
|
||||
'rate limit: retry after {} seconds',
|
||||
'rate limit: try again in {}',
|
||||
'rate limit: Retry-After: {}',
|
||||
'rate limit: "retryDelay": "{}s"',
|
||||
'rate limit: retryDelay: {}s',
|
||||
])
|
||||
@pytest.mark.parametrize("token,expected", [
|
||||
(".5", .5), ("1e2", 100), ("0", 0), ("999", 300),
|
||||
("-.5", 60), ("-5", 60), ("1e309", 60), ("1.2.3", 60),
|
||||
])
|
||||
def test_retry_decision_uses_complete_numeric_hint(template, token, expected):
|
||||
llm = LLM(model="fake")
|
||||
assert llm._parse_retry_delay(template.format(token)) == expected
|
||||
decision = llm.resolve_failover_decision(
|
||||
Exception(template.format(token)),
|
||||
{"attempt": 1, "max_retries": 3, "side_effecting": False},
|
||||
)
|
||||
assert decision.reason == "rate_limit"
|
||||
# The existing decision applies its exponential floor to a zero hint.
|
||||
assert decision.backoff_ms == int((expected or 1) * 1000)
|
||||
|
||||
|
||||
def test_retry_parser_keeps_configured_cap_and_default():
|
||||
from types import SimpleNamespace
|
||||
|
||||
llm = LLM(model="fake")
|
||||
llm._rate_limiter = SimpleNamespace(max_retry_delay=12)
|
||||
llm._retry_delay = 7
|
||||
assert llm._parse_retry_delay('try again in 1e2') == 12
|
||||
assert llm._parse_retry_delay('no retry hint') == 7
|
||||
|
||||
|
||||
def test_sync_retry_waits_for_fractional_hint():
|
||||
from types import SimpleNamespace
|
||||
|
||||
waits = []
|
||||
llm = LLM(model="fake")
|
||||
llm._max_retries = 1
|
||||
llm._rate_limiter = SimpleNamespace(
|
||||
acquire=lambda: None, wait_for_retry=waits.append,
|
||||
)
|
||||
calls = 0
|
||||
|
||||
def completion(**kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls == 1:
|
||||
raise Exception('rate limit: retry after .5 seconds')
|
||||
return 'done'
|
||||
|
||||
assert llm._call_with_retry(completion) == 'done'
|
||||
assert waits == [.5]
|
||||
assert calls == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_retry_waits_for_exponent_hint():
|
||||
from types import SimpleNamespace
|
||||
|
||||
waits = []
|
||||
|
||||
async def acquire():
|
||||
pass
|
||||
|
||||
async def wait(delay):
|
||||
waits.append(delay)
|
||||
|
||||
llm = LLM(model="fake")
|
||||
llm._max_retries = 1
|
||||
llm._rate_limiter = SimpleNamespace(acquire_async=acquire, wait_for_retry_async=wait)
|
||||
calls = 0
|
||||
|
||||
async def completion(**kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls == 1:
|
||||
raise Exception('rate limit: try again in 1e2')
|
||||
return 'done'
|
||||
|
||||
assert await llm._call_with_retry_async(completion) == 'done'
|
||||
assert waits == [100]
|
||||
assert calls == 2
|
||||
|
||||
|
||||
def test_negative_retry_cap_never_reaches_sync_sleep():
|
||||
from praisonaiagents.llm.rate_limiter import RateLimiter
|
||||
|
||||
limiter = RateLimiter(max_retry_delay=-1)
|
||||
waits = []
|
||||
limiter._sleep = waits.append
|
||||
llm = LLM(model="fake")
|
||||
llm._rate_limiter = limiter
|
||||
delay = llm._parse_retry_delay("rate limit: retry after 5 seconds")
|
||||
limiter.wait_for_retry(delay)
|
||||
assert delay == 0
|
||||
assert waits == [0]
|
||||
decision = llm.resolve_failover_decision(
|
||||
Exception("rate limit: retry after 5 seconds"),
|
||||
{"attempt": 1, "max_retries": 3, "side_effecting": False},
|
||||
)
|
||||
assert decision.backoff_ms >= 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_negative_retry_cap_never_reaches_async_sleep():
|
||||
from praisonaiagents.llm.rate_limiter import RateLimiter
|
||||
|
||||
limiter = RateLimiter(max_retry_delay=-1)
|
||||
waits = []
|
||||
|
||||
async def sleep(delay):
|
||||
waits.append(delay)
|
||||
|
||||
limiter._async_sleep = sleep
|
||||
await limiter.wait_for_retry_async(5)
|
||||
assert waits == [0]
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Retry hints must distinguish punctuation from malformed field contents."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.llm.error_classifier import extract_retry_after
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
|
||||
|
||||
@pytest.mark.parametrize('token,expected', [('0.5.', .5), ('1e-1.', .1), ('.5.', .5)])
|
||||
@pytest.mark.parametrize('prefix', ['try again in ', 'retry after ', 'wait '])
|
||||
def test_sentence_period_keeps_fractional_hint(token, expected, prefix):
|
||||
error = Exception('rate limit: ' + prefix + token)
|
||||
assert extract_retry_after(error) == expected
|
||||
decision = LLM(model='fake').resolve_failover_decision(
|
||||
error, {'attempt': 1, 'max_retries': 3, 'side_effecting': False},
|
||||
)
|
||||
assert decision.backoff_ms == int(expected * 1000)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('key', ['"retryDelay"', 'retryDelay'])
|
||||
@pytest.mark.parametrize('value', ['"5s more"', '"5sfoo"', '"5s'])
|
||||
def test_incomplete_quoted_field_does_not_hide_valid_hint(key, value):
|
||||
message = f'rate limit: {key}: {value}, retry after 20 seconds'
|
||||
assert extract_retry_after(Exception(message)) == 20
|
||||
assert LLM(model='fake')._parse_retry_delay(message) == 20
|
||||
|
||||
|
||||
@pytest.mark.parametrize('message', [
|
||||
'retryDelay: 5s', 'retryDelay: "5s"', '"retryDelay": "5s"',
|
||||
])
|
||||
def test_complete_duration_field_is_preserved(message):
|
||||
assert extract_retry_after(Exception(message)) == 5
|
||||
|
||||
|
||||
@pytest.mark.parametrize('token', ['1.2.3', '1e-1.2', '0.5..'])
|
||||
def test_embedded_period_remains_invalid(token):
|
||||
assert extract_retry_after(Exception('try again in ' + token)) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize('token', ['1e2.', '0.5.', '1.', '-5', '1e309'])
|
||||
@pytest.mark.parametrize('prefix', ['retry after ', 'wait ', 'try again in '])
|
||||
def test_invalid_first_duration_does_not_hide_later_hint(token, prefix):
|
||||
message = f'rate limit: {prefix}{token} seconds, {prefix}20 seconds'
|
||||
assert extract_retry_after(Exception(message)) == 20
|
||||
assert LLM(model='fake')._parse_retry_delay(message) == 20
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Caller edits must not rewrite previously stored in-memory records."""
|
||||
|
||||
import pytest
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import praisonaiagents.memory.adapters.in_memory_adapter as adapter_module
|
||||
|
||||
from praisonaiagents.memory.adapters.in_memory_adapter import InMemoryAdapter
|
||||
|
||||
|
||||
@pytest.mark.parametrize('tier', ['short', 'long'])
|
||||
def test_store_snapshots_nested_metadata(tier):
|
||||
adapter = InMemoryAdapter()
|
||||
metadata = {'labels': ['original'], 'details': {'count': 1}}
|
||||
memory_id = getattr(adapter, f'store_{tier}_term')('original text', metadata)
|
||||
metadata['labels'].append('caller edit')
|
||||
metadata['details']['count'] = 2
|
||||
record = getattr(adapter, f'search_{tier}_term')('original')[0]
|
||||
assert record['id'] == memory_id
|
||||
assert record['metadata'] == {'labels': ['original'], 'details': {'count': 1}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize('tier', ['short', 'long'])
|
||||
@pytest.mark.parametrize('read', ['search', 'all'])
|
||||
def test_read_results_do_not_alias_stored_records(tier, read):
|
||||
adapter = InMemoryAdapter()
|
||||
memory_id = getattr(adapter, f'store_{tier}_term')(
|
||||
'original text', {'labels': ['original'], 'details': {'count': 1}},
|
||||
)
|
||||
search = getattr(adapter, f'search_{tier}_term')
|
||||
records = search('original') if read == 'search' else adapter.get_all_memories()
|
||||
records[0]['id'] = 'caller id'
|
||||
records[0]['text'] = 'caller text'
|
||||
records[0]['metadata']['labels'].append('caller edit')
|
||||
records[0]['metadata']['details']['count'] = 2
|
||||
records.clear()
|
||||
assert search('original') == [{
|
||||
'id': memory_id, 'text': 'original text', 'type': tier,
|
||||
'metadata': {'labels': ['original'], 'details': {'count': 1}},
|
||||
}]
|
||||
assert adapter.delete_memory(memory_id, tier=tier)
|
||||
assert adapter.get_all_memories() == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize('tier', ['short', 'long'])
|
||||
@pytest.mark.parametrize('metadata', [None, {}])
|
||||
def test_metadata_defaults_and_search_limits_remain_unchanged(tier, metadata):
|
||||
adapter = InMemoryAdapter()
|
||||
write = getattr(adapter, f'store_{tier}_term')
|
||||
expected_ids = [write('matching text', metadata) for _ in range(6)]
|
||||
search = getattr(adapter, f'search_{tier}_term')
|
||||
assert [record['id'] for record in search('matching')] == expected_ids[:5]
|
||||
assert [record['id'] for record in search('matching', limit=2)] == expected_ids[:2]
|
||||
assert all(record['metadata'] == metadata for record in adapter.get_all_memories())
|
||||
|
||||
|
||||
def test_enumeration_copy_allows_concurrent_eviction_and_preserves_snapshot(monkeypatch):
|
||||
"""A slow nested copy must not keep writers behind the collection lock."""
|
||||
adapter = InMemoryAdapter(max_size=1)
|
||||
original_id = adapter.store_short_term('original', {'nested': ['original']})
|
||||
copying = threading.Event()
|
||||
resume = threading.Event()
|
||||
real_copy = adapter_module.deepcopy
|
||||
|
||||
def paused_copy(value):
|
||||
if isinstance(value, list):
|
||||
copying.set()
|
||||
assert resume.wait(5)
|
||||
return real_copy(value)
|
||||
|
||||
monkeypatch.setattr(adapter_module, 'deepcopy', paused_copy)
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
read = executor.submit(adapter.get_all_memories)
|
||||
try:
|
||||
assert copying.wait(2)
|
||||
new_id = executor.submit(adapter.store_long_term, 'new', {'nested': ['new']}).result(timeout=1)
|
||||
finally:
|
||||
resume.set()
|
||||
snapshot = read.result(timeout=5)
|
||||
assert snapshot == [{'id': original_id, 'text': 'original', 'type': 'short', 'metadata': {'nested': ['original']}}]
|
||||
snapshot[0]['metadata']['nested'].append('caller edit')
|
||||
assert adapter.get_all_memories() == [{'id': new_id, 'text': 'new', 'type': 'long', 'metadata': {'nested': ['new']}}]
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Enumerating SQLite memories must include rows beyond search-sized caps."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.memory.adapters.sqlite_adapter import SqliteMemoryAdapter
|
||||
from praisonaiagents.memory.memory import Memory
|
||||
|
||||
|
||||
@pytest.mark.parametrize('facade', [False, True])
|
||||
@pytest.mark.parametrize('short_count,long_count', [(0, 0), (1000, 1), (1, 1000), (1001, 1001)])
|
||||
def test_get_all_memories_includes_every_stored_row(tmp_path, monkeypatch, facade, short_count, long_count):
|
||||
monkeypatch.setenv('PRAISONAI_HOME', str(tmp_path / 'home'))
|
||||
monkeypatch.chdir(tmp_path)
|
||||
store = Memory(config={'provider': 'sqlite'}) if facade else SqliteMemoryAdapter(
|
||||
short_db=str(tmp_path / 'short.db'), long_db=str(tmp_path / 'long.db'),
|
||||
)
|
||||
adapter = store.memory_adapter if facade else store
|
||||
expected = set()
|
||||
try:
|
||||
for tier, count in [('short', short_count), ('long', long_count)]:
|
||||
write = getattr(store, f'store_{tier}_term')
|
||||
for index in range(count):
|
||||
text = f'{tier} entry {index}'
|
||||
memory_id = write(text, metadata={'index': index})
|
||||
expected.add((f'{tier}_term', memory_id, text, index))
|
||||
records = store.get_all_memories()
|
||||
assert len(records) == short_count + long_count
|
||||
assert {
|
||||
(record['memory_type'], record['id'], record['text'], record['metadata']['index'])
|
||||
for record in records
|
||||
} == expected
|
||||
if facade:
|
||||
assert all(record['type'] == record['memory_type'] for record in records)
|
||||
# Enumeration must not widen the public search defaults or explicit limit.
|
||||
for tier, count in [('short', short_count), ('long', long_count)]:
|
||||
search = getattr(adapter, f'search_{tier}_term')
|
||||
assert len(search('')) == min(5, count)
|
||||
assert len(search('', limit=2)) == min(2, count)
|
||||
finally:
|
||||
adapter.close_connections()
|
||||
if facade:
|
||||
store.close_connections()
|
||||
|
||||
|
||||
@pytest.mark.parametrize('tier', ['short', 'long'])
|
||||
@pytest.mark.parametrize('value', [float('nan'), float('inf'), -float('inf')])
|
||||
def test_filtered_search_accepts_writer_nonfinite_metadata(tmp_path, tier, value):
|
||||
adapter = SqliteMemoryAdapter(short_db=str(tmp_path / 'short.db'), long_db=str(tmp_path / 'long.db'))
|
||||
try:
|
||||
write = getattr(adapter, f'store_{tier}_term')
|
||||
search = getattr(adapter, f'search_{tier}_term')
|
||||
write('other record', metadata={'user_id': 'other', 'value': value})
|
||||
wanted = write('wanted record', metadata={'user_id': 'wanted', 'value': value})
|
||||
assert [item['id'] for item in search('', user_id='wanted', limit=1)] == [wanted]
|
||||
assert search('', user_id='absent') == []
|
||||
finally:
|
||||
adapter.close_connections()
|
||||
|
||||
|
||||
@pytest.mark.parametrize('tier', ['short', 'long'])
|
||||
def test_filtered_search_ignores_corrupt_peer_metadata_and_keeps_string_matching(tmp_path, tier):
|
||||
adapter = SqliteMemoryAdapter(short_db=str(tmp_path / 'short.db'), long_db=str(tmp_path / 'long.db'))
|
||||
try:
|
||||
write = getattr(adapter, f'store_{tier}_term')
|
||||
wanted = write('wanted', metadata={'user_id': True})
|
||||
conn = getattr(adapter, f'_get_{"stm" if tier == "short" else "ltm"}_conn')()
|
||||
conn.execute(f'INSERT INTO {tier}_term_memory (content, metadata) VALUES (?, ?)', ('broken peer', '{invalid'))
|
||||
conn.commit()
|
||||
search = getattr(adapter, f'search_{tier}_term')
|
||||
assert [item['id'] for item in search('', user_id='True', limit=1)] == [wanted]
|
||||
finally:
|
||||
adapter.close_connections()
|
||||
|
||||
|
||||
@pytest.mark.parametrize('tier', ['short', 'long'])
|
||||
def test_null_user_is_distinct_from_literal_none_user(tmp_path, tier):
|
||||
adapter = SqliteMemoryAdapter(short_db=str(tmp_path / 'short.db'), long_db=str(tmp_path / 'long.db'))
|
||||
try:
|
||||
write = getattr(adapter, f'store_{tier}_term')
|
||||
write('unowned', metadata={'user_id': None})
|
||||
wanted = write('owned', metadata={'user_id': 'None'})
|
||||
search = getattr(adapter, f'search_{tier}_term')
|
||||
assert [item['id'] for item in search('', user_id='None')] == [wanted]
|
||||
finally:
|
||||
adapter.close_connections()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tier", ["short_term", "long_term"])
|
||||
@pytest.mark.parametrize("user_id", [0, False, ""])
|
||||
def test_falsey_supplied_user_id_still_filters_before_limit(tmp_path, tier, user_id):
|
||||
adapter = SqliteMemoryAdapter(short_db=str(tmp_path / 'short.db'), long_db=str(tmp_path / 'long.db'))
|
||||
try:
|
||||
write = getattr(adapter, f"store_{tier}")
|
||||
search = getattr(adapter, f"search_{tier}")
|
||||
write("other", metadata={"user_id": "other"})
|
||||
wanted = write("wanted", metadata={"user_id": user_id})
|
||||
assert [item["id"] for item in search("", user_id=user_id, limit=1)] == [wanted]
|
||||
finally:
|
||||
adapter.close_connections()
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Continuation prompts must distinguish running work from proven completion."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session import build_handoff_prompt
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"checkpoint,expected,absent",
|
||||
[
|
||||
({"current_step": "deploy"}, "PROGRESS: current step: deploy", ["last completed", "steps done"]),
|
||||
({"completed_steps": 0, "total_steps": 5, "current_step": "deploy"},
|
||||
"PROGRESS: 0 of 5 workflow steps done; current step: deploy", ["last completed"]),
|
||||
({"completed_steps": 2, "total_steps": 5, "last_step": "build", "current_step": "deploy"},
|
||||
"PROGRESS: 2 of 5 workflow steps done; last completed: build; current step: deploy", []),
|
||||
({"last_step": "build"}, "PROGRESS: last completed: build", ["steps done"]),
|
||||
({"completed_steps": 0, "total_steps": 0, "step_count": 9},
|
||||
"PROGRESS: 0 of 0 workflow steps done", ["of 9"]),
|
||||
({"completed_steps": 2, "step_count": 5, "last_step": "build"},
|
||||
"PROGRESS: 2 of 5 workflow steps done; last completed: build", []),
|
||||
],
|
||||
)
|
||||
def test_handoff_reports_only_the_supplied_progress(checkpoint, expected, absent):
|
||||
prompt = build_handoff_prompt(workflow_checkpoint=checkpoint)
|
||||
|
||||
assert expected in prompt
|
||||
for text in absent:
|
||||
assert text not in prompt
|
||||
assert prompt.endswith("before acting.")
|
||||
@@ -477,6 +477,41 @@ class TestHierarchicalSessionStore:
|
||||
assert len(session1.messages) == len(session2.messages)
|
||||
assert session1.session_id == session2.session_id
|
||||
|
||||
def test_extended_cache_is_bounded_and_reloads_after_eviction(self):
|
||||
"""The extended-session cache (and its mtimes) must stay bounded.
|
||||
|
||||
A long-lived process using HierarchicalSessionStore previously kept
|
||||
every loaded session in ``_extended_cache`` forever. It now shares the
|
||||
base LRU bound, so touching more sessions than ``cache_maxsize`` evicts
|
||||
the oldest, drops its companion mtime entry, and the next access
|
||||
reloads from disk (the source of truth).
|
||||
"""
|
||||
store = HierarchicalSessionStore(session_dir=self.temp_dir, cache_maxsize=2)
|
||||
|
||||
ids = []
|
||||
for i in range(5):
|
||||
sid = store.create_session(title=f"S{i}")
|
||||
store.add_message(sid, "user", f"msg-{i}")
|
||||
store.get_extended_session(sid) # populate extended cache
|
||||
ids.append(sid)
|
||||
|
||||
# Cache never exceeds the configured bound, and mtimes track the cache.
|
||||
assert len(store._extended_cache) <= 2
|
||||
assert len(store._cache_mtimes) <= 2
|
||||
assert set(store._cache_mtimes).issubset(set(store._extended_cache))
|
||||
|
||||
# An evicted session still reads correctly (reloads from disk).
|
||||
reloaded = store.get_extended_session(ids[0])
|
||||
assert [m.content for m in reloaded.messages] == ["msg-0"]
|
||||
|
||||
def test_extended_cache_unbounded_when_maxsize_zero(self):
|
||||
"""cache_maxsize=0 keeps the legacy unbounded extended cache."""
|
||||
store = HierarchicalSessionStore(session_dir=self.temp_dir, cache_maxsize=0)
|
||||
for i in range(5):
|
||||
sid = store.create_session(title=f"S{i}")
|
||||
store.get_extended_session(sid)
|
||||
assert len(store._extended_cache) == 5
|
||||
|
||||
|
||||
class TestGlobalHierarchicalStore:
|
||||
"""Tests for global hierarchical store."""
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
"""Hierarchy deserialization retains the complete base session schema."""
|
||||
|
||||
from dataclasses import fields
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.hierarchy import ExtendedSessionData, HierarchicalSessionStore
|
||||
from praisonaiagents.session.store import DefaultSessionStore, SessionData
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_session(tmp_path):
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path), active_window=3)
|
||||
assert store.add_message("session", "user", "archived opening")
|
||||
for index in range(8):
|
||||
assert store.add_message("session", "user", f"later turn {index}")
|
||||
assert store.set_gateway_info("session", gateway_session_id="gateway", agent_id="agent")
|
||||
assert store.set_runtime_state("session", "native", "turn", {"status": "saved"})
|
||||
assert store.append_compaction_checkpoint("session", "resume summary", tokens_before=20, tokens_after=5)
|
||||
session = store.get_session("session")
|
||||
assert session.archived_messages
|
||||
return store, session
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["read", "set_title", "add_message"])
|
||||
@pytest.mark.parametrize("field_name", ["gateway_session_id", "agent_id", "runtime_state", "last_compaction", "archived_messages"])
|
||||
def test_hierarchy_load_and_mutation_keep_base_fields(base_session, operation, field_name):
|
||||
plain, before = base_session
|
||||
# Preserve the existing archive without triggering a new compaction whose
|
||||
# legitimate anchor adjustment would obscure the deserialization contract.
|
||||
hierarchy = HierarchicalSessionStore(session_dir=plain.session_dir, active_window=100)
|
||||
if operation == "set_title":
|
||||
assert hierarchy.set_title("session", "Hierarchy title")
|
||||
elif operation == "add_message":
|
||||
assert hierarchy.add_message("session", "assistant", "new turn")
|
||||
after = hierarchy.get_extended_session("session", force_reload=True)
|
||||
assert getattr(after, field_name) == getattr(before, field_name)
|
||||
plain.invalidate_cache()
|
||||
assert getattr(plain.get_session("session"), field_name) == getattr(before, field_name)
|
||||
|
||||
|
||||
def test_conversion_from_base_session_keeps_every_base_field(base_session):
|
||||
_, before = base_session
|
||||
after = ExtendedSessionData.from_session_data(before)
|
||||
for descriptor in fields(SessionData):
|
||||
assert getattr(after, descriptor.name) == getattr(before, descriptor.name), descriptor.name
|
||||
|
||||
|
||||
def test_hierarchy_parser_keeps_legacy_metadata_and_precedence():
|
||||
payload = {"session_id": "legacy", "model": "recorded-model", "reasoning_effort": "high",
|
||||
"metadata": {"model": "preferred-model"}}
|
||||
base = SessionData.from_dict(payload)
|
||||
extended = ExtendedSessionData.from_dict(payload)
|
||||
assert extended.metadata == base.metadata == {"model": "preferred-model", "reasoning_effort": "high"}
|
||||
assert payload["metadata"] == {"model": "preferred-model"}
|
||||
|
||||
|
||||
def test_extended_fields_still_round_trip():
|
||||
payload = {"session_id": "extended", "parent_id": "parent", "forked_from_message_id": "2",
|
||||
"children_ids": ["child"], "is_shared": True, "title": "Title",
|
||||
"snapshots": [{"id": "snapshot", "session_id": "extended", "message_index": 2}]}
|
||||
after = ExtendedSessionData.from_dict(payload).to_dict()
|
||||
for name in ("parent_id", "forked_from_message_id", "children_ids", "is_shared", "title"):
|
||||
assert after[name] == payload[name]
|
||||
assert after["snapshots"][0]["id"] == "snapshot"
|
||||
@@ -0,0 +1,156 @@
|
||||
"""Hierarchical caches must describe the bytes actually loaded, not later stats."""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.hierarchy import HierarchicalSessionStore
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
from praisonaiagents.session.store import FileLock
|
||||
|
||||
|
||||
def warm_reader(tmp_path):
|
||||
reader = HierarchicalSessionStore(session_dir=str(tmp_path))
|
||||
assert reader.add_message("parent", "user", "first")
|
||||
assert [message.content for message in reader.get_extended_session("parent").messages] == ["first"]
|
||||
return reader
|
||||
|
||||
|
||||
@pytest.mark.parametrize("offset", [0, -1_000_000_000])
|
||||
def test_cross_instance_write_with_unchanged_or_older_mtime(tmp_path, offset):
|
||||
reader = warm_reader(tmp_path)
|
||||
path = tmp_path / "parent.json"
|
||||
stamp = path.stat()
|
||||
writer = HierarchicalSessionStore(session_dir=str(tmp_path))
|
||||
assert writer.add_message("parent", "user", "second")
|
||||
os.utime(path, ns=(stamp.st_atime_ns, stamp.st_mtime_ns + offset))
|
||||
|
||||
assert [message.content for message in reader.get_extended_session("parent").messages] == ["first", "second"]
|
||||
|
||||
|
||||
def test_same_size_in_place_change_with_unchanged_mtime(tmp_path):
|
||||
reader = warm_reader(tmp_path)
|
||||
path = tmp_path / "parent.json"
|
||||
stamp = path.stat()
|
||||
path.write_bytes(path.read_bytes().replace(b'"first"', b'"other"'))
|
||||
os.utime(path, ns=(stamp.st_atime_ns, stamp.st_mtime_ns))
|
||||
assert path.stat().st_size == stamp.st_size
|
||||
|
||||
assert reader.get_extended_session("parent").messages[0].content == "other"
|
||||
|
||||
|
||||
def test_unchanged_bytes_reuse_parsed_session(tmp_path, monkeypatch):
|
||||
reader = warm_reader(tmp_path)
|
||||
cached = reader.get_extended_session("parent")
|
||||
|
||||
def unexpected_parse(*args):
|
||||
pytest.fail("Unchanged bytes should not be parsed again")
|
||||
|
||||
monkeypatch.setattr(reader, "_load_session_from_disk", unexpected_parse)
|
||||
assert reader.get_extended_session("parent") is cached
|
||||
assert reader._is_cache_valid("parent")
|
||||
|
||||
|
||||
def test_invalidation_after_validation_reloads_instead_of_raising(tmp_path, monkeypatch):
|
||||
reader = warm_reader(tmp_path)
|
||||
original = reader._is_cache_valid
|
||||
|
||||
def invalidate_after_check(session_id):
|
||||
valid = original(session_id)
|
||||
reader.invalidate_cache(session_id)
|
||||
return valid
|
||||
|
||||
monkeypatch.setattr(reader, "_is_cache_valid", invalidate_after_check)
|
||||
assert reader.get_extended_session("parent").messages[0].content == "first"
|
||||
|
||||
|
||||
def test_write_after_base_read_does_not_certify_old_payload(tmp_path, monkeypatch):
|
||||
reader = warm_reader(tmp_path)
|
||||
writer = HierarchicalSessionStore(session_dir=str(tmp_path))
|
||||
original = DefaultSessionStore._read_session_fresh
|
||||
fired = False
|
||||
|
||||
def interleaved_read(self, session_id):
|
||||
nonlocal fired
|
||||
session = original(self, session_id)
|
||||
if self is reader and not fired:
|
||||
fired = True
|
||||
assert writer.add_message(session_id, "user", "second")
|
||||
return session
|
||||
|
||||
monkeypatch.setattr(DefaultSessionStore, "_read_session_fresh", interleaved_read)
|
||||
assert [message.content for message in reader.get_extended_session("parent", force_reload=True).messages] == ["first"]
|
||||
assert fired
|
||||
assert [message.content for message in reader.get_extended_session("parent").messages] == ["first", "second"]
|
||||
|
||||
|
||||
def test_transient_read_fallback_retries_after_recovery(tmp_path, monkeypatch):
|
||||
reader = warm_reader(tmp_path)
|
||||
writer = HierarchicalSessionStore(session_dir=str(tmp_path))
|
||||
assert writer.add_message("parent", "user", "second")
|
||||
original = reader._load_session_from_disk
|
||||
|
||||
def unavailable(*args):
|
||||
raise OSError("temporary read failure")
|
||||
|
||||
monkeypatch.setattr(reader, "_load_session_from_disk", unavailable)
|
||||
assert [message.content for message in reader.get_extended_session("parent", force_reload=True).messages] == ["first"]
|
||||
monkeypatch.setattr(reader, "_load_session_from_disk", original)
|
||||
|
||||
assert [message.content for message in reader.get_extended_session("parent").messages] == ["first", "second"]
|
||||
|
||||
|
||||
def test_corrupt_bytes_with_unchanged_mtime_are_quarantined(tmp_path):
|
||||
reader = warm_reader(tmp_path)
|
||||
path = tmp_path / "parent.json"
|
||||
stamp = path.stat()
|
||||
corrupt = b'{"messages": ['
|
||||
path.write_bytes(corrupt)
|
||||
os.utime(path, ns=(stamp.st_atime_ns, stamp.st_mtime_ns))
|
||||
|
||||
assert reader.get_extended_session("parent").messages == []
|
||||
quarantined = list(tmp_path.glob("parent.json.corrupt-*"))
|
||||
assert len(quarantined) == 1
|
||||
assert quarantined[0].read_bytes() == corrupt
|
||||
|
||||
|
||||
def test_cached_read_under_file_lock_contention_recovers(tmp_path):
|
||||
reader = warm_reader(tmp_path)
|
||||
cached = reader.get_extended_session("parent")
|
||||
reader.lock_timeout = 0.01
|
||||
with FileLock(str(tmp_path / "parent.json"), 1):
|
||||
assert reader.get_extended_session("parent") is cached
|
||||
with pytest.raises(OSError, match="Failed to acquire file lock"):
|
||||
reader.get_extended_session("parent", force_reload=True)
|
||||
with pytest.raises(OSError, match="Failed to acquire file lock"):
|
||||
reader.add_message("parent", "user", "must not be written")
|
||||
writer = HierarchicalSessionStore(session_dir=str(tmp_path))
|
||||
assert writer.add_message("parent", "user", "second")
|
||||
assert len(reader.get_extended_session("parent").messages) == 2
|
||||
|
||||
|
||||
def test_uncached_read_does_not_fabricate_success_under_contention(tmp_path):
|
||||
writer = warm_reader(tmp_path)
|
||||
reader = HierarchicalSessionStore(session_dir=str(tmp_path), lock_timeout=0.01)
|
||||
with FileLock(writer._get_session_path("parent"), 1):
|
||||
with pytest.raises(OSError, match="Failed to acquire file lock"):
|
||||
reader.get_extended_session("parent")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("write", ["save", "modify"])
|
||||
def test_local_write_reuses_written_parsed_session(tmp_path, monkeypatch, write):
|
||||
reader = warm_reader(tmp_path)
|
||||
if write == "save":
|
||||
cached = reader.get_extended_session("parent")
|
||||
cached.title = "écriture\nsecond line"
|
||||
assert reader._save_extended_session(cached)
|
||||
else:
|
||||
assert reader.add_message("parent", "user", "écriture\nsecond line")
|
||||
cached = reader._extended_cache["parent"]
|
||||
|
||||
def unexpected_parse(*args):
|
||||
pytest.fail("A successful local write should reuse its parsed object")
|
||||
|
||||
monkeypatch.setattr(reader, "_load_session_from_disk", unexpected_parse)
|
||||
assert reader.get_extended_session("parent") is cached
|
||||
assert reader._is_cache_valid("parent")
|
||||
@@ -0,0 +1,47 @@
|
||||
"""Public fork boundaries must identify an actual message before writing."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.hierarchy import HierarchicalSessionStore
|
||||
|
||||
|
||||
def make_parent(tmp_path, count):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), max_messages=100)
|
||||
store.create_session(session_id="parent", title="Parent")
|
||||
for index in range(count):
|
||||
assert store.add_message("parent", "user", f"message {index}")
|
||||
return store
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"count,index", [(3, -1), (3, -2), (3, -5), (3, 3), (3, 99), (0, 0), (0, -1)]
|
||||
)
|
||||
def test_invalid_fork_index_does_not_write(tmp_path, count, index):
|
||||
store = make_parent(tmp_path, count)
|
||||
parent_path = tmp_path / "parent.json"
|
||||
before = parent_path.read_bytes()
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid message index"):
|
||||
store.fork_session("parent", from_message_index=index)
|
||||
|
||||
assert parent_path.read_bytes() == before
|
||||
assert sorted(path.name for path in tmp_path.glob("*.json")) == ["parent.json"]
|
||||
assert store.get_children("parent") == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("count,index", [(3, None), (3, 0), (3, 2), (0, None)])
|
||||
def test_valid_fork_index_is_inclusive_and_durable(tmp_path, count, index):
|
||||
store = make_parent(tmp_path, count)
|
||||
fork_id = store.fork_session("parent", from_message_index=index)
|
||||
|
||||
fork = json.loads((tmp_path / f"{fork_id}.json").read_text(encoding="utf-8"))
|
||||
parent = json.loads((tmp_path / "parent.json").read_text(encoding="utf-8"))
|
||||
expected_count = count if index is None else index + 1
|
||||
assert [message["content"] for message in fork["messages"]] == [
|
||||
f"message {number}" for number in range(expected_count)
|
||||
]
|
||||
assert fork["forked_from_message_id"] == (None if index is None else str(index))
|
||||
assert fork["parent_id"] == "parent"
|
||||
assert parent["children_ids"] == [fork_id]
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Hierarchy writes must honor the inherited retention policy before slicing."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.hierarchy import HierarchicalSessionStore
|
||||
from praisonaiagents.session.store import DefaultSessionStore, SessionData, SessionMessage
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["add", "fork", "import"])
|
||||
@pytest.mark.parametrize(
|
||||
"retention,window,active,archived",
|
||||
[
|
||||
("keep_all", 2, ["0", "1", "2", "3", "4"], []),
|
||||
("compact", 2, ["3", "4"], ["0", "1", "2"]),
|
||||
("truncate", 4, ["1", "2", "3", "4"], []),
|
||||
(None, None, ["3", "4"], []),
|
||||
],
|
||||
)
|
||||
def test_retention_survives_public_hierarchy_writes(
|
||||
tmp_path, operation, retention, window, active, archived
|
||||
):
|
||||
store = HierarchicalSessionStore(
|
||||
session_dir=str(tmp_path), max_messages=2,
|
||||
retention=retention, active_window=window,
|
||||
)
|
||||
messages = [SessionMessage("user", str(index), timestamp=index) for index in range(5)]
|
||||
if operation == "add":
|
||||
result_id = store.create_session(session_id="parent")
|
||||
for message in messages:
|
||||
assert store.add_message(result_id, message.role, message.content)
|
||||
elif operation == "fork":
|
||||
source = DefaultSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
for message in messages:
|
||||
assert source.add_message("parent", message.role, message.content)
|
||||
result_id = store.fork_session("parent")
|
||||
else:
|
||||
result_id = store.import_session(
|
||||
SessionData(session_id="source", messages=messages).to_dict(),
|
||||
new_session_id="imported",
|
||||
)
|
||||
active = ["0", "1", "2", "3", "4"]
|
||||
archived = []
|
||||
|
||||
result = json.loads((tmp_path / f"{result_id}.json").read_text(encoding="utf-8"))
|
||||
raw_active = [message for message in result["messages"] if not message["metadata"].get("compaction")]
|
||||
assert [message["content"] for message in raw_active] == active
|
||||
assert [message["content"] for message in result["archived_messages"]] == archived
|
||||
if retention == "compact" and operation != "import":
|
||||
summary = result["messages"][0]
|
||||
assert summary["role"] == "system"
|
||||
assert summary["metadata"]["compaction"] is True
|
||||
assert summary["metadata"]["compacted_count"] == 3
|
||||
reopened = HierarchicalSessionStore(session_dir=str(tmp_path))
|
||||
fresh = reopened.get_extended_session(result_id, force_reload=True)
|
||||
assert [message.content for message in fresh.archived_messages] == archived
|
||||
|
||||
|
||||
@pytest.mark.parametrize("retention", ["keep_all", "compact", "truncate"])
|
||||
def test_import_restores_export_larger_than_destination_window(tmp_path, retention):
|
||||
store = HierarchicalSessionStore(
|
||||
session_dir=str(tmp_path), max_messages=100, active_window=2,
|
||||
retention=retention,
|
||||
)
|
||||
data = SessionData(
|
||||
session_id="source",
|
||||
messages=[SessionMessage("user", str(index), timestamp=index) for index in range(5)],
|
||||
archived_messages=[SessionMessage("user", "archived", timestamp=-1)],
|
||||
).to_dict()
|
||||
|
||||
imported_id = store.import_session(data)
|
||||
|
||||
saved = json.loads((tmp_path / f"{imported_id}.json").read_text(encoding="utf-8"))
|
||||
assert saved["messages"] == data["messages"]
|
||||
assert saved["archived_messages"] == data["archived_messages"]
|
||||
|
||||
|
||||
def test_truncate_does_not_persist_orphaned_tool_result(tmp_path):
|
||||
store = HierarchicalSessionStore(
|
||||
session_dir=str(tmp_path), max_messages=2, retention="truncate"
|
||||
)
|
||||
calls = [{"id": "call1", "type": "function", "function": {"name": "read", "arguments": "{}"}}]
|
||||
assert store.add_message("parent", "assistant", "", tool_calls=calls)
|
||||
assert store.add_message("parent", "tool", "result", tool_call_id="call1")
|
||||
assert store.add_message("parent", "user", "next")
|
||||
|
||||
result = json.loads((tmp_path / "parent.json").read_text(encoding="utf-8"))
|
||||
assert [message["role"] for message in result["messages"]] == ["user"]
|
||||
assert result["messages"][0]["content"] == "next"
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Snapshots must restore their transcript after the live window moves."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.hierarchy import HierarchicalSessionStore, SessionSnapshot
|
||||
from praisonaiagents.session.store import CompactionCheckpoint, SessionData, SessionMessage
|
||||
|
||||
|
||||
@pytest.mark.parametrize("retention", ["compact", "truncate", "keep_all"])
|
||||
def test_snapshot_survives_retention_and_reopen(tmp_path, retention):
|
||||
store = HierarchicalSessionStore(
|
||||
session_dir=str(tmp_path), max_messages=100, active_window=2,
|
||||
retention=retention,
|
||||
)
|
||||
for content in ["one", "two"]:
|
||||
assert store.add_message("parent", "user", content)
|
||||
before = json.loads((tmp_path / "parent.json").read_text(encoding="utf-8"))
|
||||
snapshot_id = store.create_snapshot("parent")
|
||||
for content in ["three", "four"]:
|
||||
assert store.add_message("parent", "user", content)
|
||||
reopened = HierarchicalSessionStore(
|
||||
session_dir=str(tmp_path), max_messages=100, active_window=2,
|
||||
retention=retention,
|
||||
)
|
||||
|
||||
assert reopened.revert_to_snapshot("parent", snapshot_id)
|
||||
|
||||
restored = json.loads((tmp_path / "parent.json").read_text(encoding="utf-8"))
|
||||
assert restored["messages"] == before["messages"]
|
||||
assert restored["archived_messages"] == before["archived_messages"]
|
||||
assert restored.get("last_compaction") == before.get("last_compaction")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("retention", ["compact", "truncate", "keep_all"])
|
||||
def test_snapshot_restores_large_import_without_reapplying_window(tmp_path, retention):
|
||||
store = HierarchicalSessionStore(
|
||||
session_dir=str(tmp_path), max_messages=100, active_window=2, retention=retention
|
||||
)
|
||||
exported = SessionData(
|
||||
session_id="source",
|
||||
messages=[SessionMessage("user", str(index), timestamp=index) for index in range(5)],
|
||||
archived_messages=[SessionMessage("user", "archived", timestamp=-1)],
|
||||
).to_dict()
|
||||
session_id = store.import_session(exported)
|
||||
snapshot_id = store.create_snapshot(session_id)
|
||||
live = json.loads((tmp_path / f"{session_id}.json").read_text(encoding="utf-8"))
|
||||
assert live["messages"] == exported["messages"]
|
||||
assert live["archived_messages"] == exported["archived_messages"]
|
||||
assert store.add_message(session_id, "user", "later")
|
||||
|
||||
assert store.revert_to_snapshot(session_id, snapshot_id)
|
||||
|
||||
saved = json.loads((tmp_path / f"{session_id}.json").read_text(encoding="utf-8"))
|
||||
assert saved["messages"] == exported["messages"]
|
||||
assert saved["archived_messages"] == exported["archived_messages"]
|
||||
|
||||
|
||||
def test_empty_snapshot_removes_later_archive_and_messages(tmp_path):
|
||||
store = HierarchicalSessionStore(
|
||||
session_dir=str(tmp_path), max_messages=100, active_window=2, retention="compact"
|
||||
)
|
||||
store.create_session(session_id="parent")
|
||||
snapshot_id = store.create_snapshot("parent")
|
||||
for content in ["one", "two", "three", "four"]:
|
||||
assert store.add_message("parent", "user", content)
|
||||
assert store.export_session("parent")["archived_messages"]
|
||||
|
||||
assert store.revert_to_snapshot("parent", snapshot_id)
|
||||
|
||||
restored = store.export_session("parent")
|
||||
assert restored["messages"] == []
|
||||
assert restored["archived_messages"] == []
|
||||
|
||||
|
||||
def test_legacy_snapshot_without_window_change_remains_compatible(tmp_path):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
assert store.add_message("parent", "user", "one")
|
||||
path = tmp_path / "parent.json"
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
legacy = SessionSnapshot(id="legacy", session_id="parent", message_index=0).to_dict()
|
||||
assert "transcript" not in legacy
|
||||
data["snapshots"] = [legacy]
|
||||
path.write_text(json.dumps(data), encoding="utf-8")
|
||||
assert store.add_message("parent", "user", "two")
|
||||
|
||||
assert store.revert_to_snapshot("parent", "legacy")
|
||||
assert [message["content"] for message in store.export_session("parent")["messages"]] == ["one"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("retention", ["compact", "truncate"])
|
||||
def test_legacy_snapshot_is_rejected_after_window_change(tmp_path, retention):
|
||||
store = HierarchicalSessionStore(
|
||||
session_dir=str(tmp_path), max_messages=100, active_window=2, retention=retention
|
||||
)
|
||||
for content in ["one", "two"]:
|
||||
assert store.add_message("parent", "user", content)
|
||||
path = tmp_path / "parent.json"
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
data["snapshots"] = [SessionSnapshot(id="legacy", session_id="parent", message_index=1).to_dict()]
|
||||
path.write_text(json.dumps(data), encoding="utf-8")
|
||||
for content in ["three", "four"]:
|
||||
assert store.add_message("parent", "user", content)
|
||||
before = path.read_bytes()
|
||||
|
||||
assert store.revert_to_snapshot("parent", "legacy") is False
|
||||
assert path.read_bytes() == before
|
||||
|
||||
|
||||
def test_persisted_snapshot_restores_structured_turns_archive_and_checkpoint(tmp_path):
|
||||
calls = [{"id": "call1", "type": "function", "function": {"name": "read", "arguments": "{}"}}]
|
||||
captured = SessionData(
|
||||
session_id="parent",
|
||||
messages=[SessionMessage("assistant", "", timestamp=1, tool_calls=calls),
|
||||
SessionMessage("tool", "result", timestamp=2, tool_call_id="call1")],
|
||||
archived_messages=[SessionMessage("user", "early", timestamp=0)],
|
||||
last_compaction=CompactionCheckpoint(summary="saved summary", message_index=1),
|
||||
).to_dict()
|
||||
transcript = {key: captured[key] for key in ["messages", "archived_messages", "last_compaction"]}
|
||||
snapshot = SessionSnapshot(id="saved", session_id="parent", transcript=transcript)
|
||||
live = SessionData(session_id="parent", messages=[SessionMessage("user", "later")]).to_dict()
|
||||
live["snapshots"] = [snapshot.to_dict()]
|
||||
path = tmp_path / "parent.json"
|
||||
path.write_text(json.dumps(live), encoding="utf-8")
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), active_window=1, retention="truncate")
|
||||
|
||||
assert store.revert_to_snapshot("parent", "saved")
|
||||
|
||||
restored = json.loads(path.read_text(encoding="utf-8"))
|
||||
for key, value in transcript.items():
|
||||
assert restored[key] == value
|
||||
|
||||
|
||||
def test_failed_restore_preserves_durable_file_and_snapshot(tmp_path, monkeypatch):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
assert store.add_message("parent", "user", "one")
|
||||
snapshot_id = store.create_snapshot("parent")
|
||||
assert store.add_message("parent", "user", "two")
|
||||
path = tmp_path / "parent.json"
|
||||
before = path.read_bytes()
|
||||
original = store._atomic_write_json
|
||||
monkeypatch.setattr(store, "_atomic_write_json", lambda *args: False)
|
||||
|
||||
assert store.revert_to_snapshot("parent", snapshot_id) is False
|
||||
assert path.read_bytes() == before
|
||||
monkeypatch.setattr(store, "_atomic_write_json", original)
|
||||
assert store.revert_to_snapshot("parent", snapshot_id)
|
||||
assert [message["content"] for message in store.export_session("parent")["messages"]] == ["one"]
|
||||
|
||||
|
||||
def test_legacy_snapshot_stays_invalid_after_manual_revert_and_append(tmp_path):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
for content in ["one", "two"]:
|
||||
assert store.add_message("parent", "user", content)
|
||||
path = tmp_path / "parent.json"
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
data["snapshots"] = [SessionSnapshot(id="legacy", session_id="parent", message_index=1).to_dict()]
|
||||
path.write_text(json.dumps(data), encoding="utf-8")
|
||||
assert store.revert_to_message("parent", 0)
|
||||
assert store.add_message("parent", "user", "replacement")
|
||||
before = path.read_bytes()
|
||||
|
||||
assert store.revert_to_snapshot("parent", "legacy") is False
|
||||
assert path.read_bytes() == before
|
||||
|
||||
|
||||
def test_new_snapshot_captures_existing_durable_compaction_checkpoint(tmp_path):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
assert store.add_message("parent", "user", "one")
|
||||
path = tmp_path / "parent.json"
|
||||
assert store.append_compaction_checkpoint("parent", "checkpoint")
|
||||
checkpoint = json.loads(path.read_text(encoding="utf-8"))["last_compaction"]
|
||||
|
||||
snapshot_id = store.create_snapshot("parent")
|
||||
assert store.add_message("parent", "user", "later")
|
||||
assert store.revert_to_snapshot("parent", snapshot_id)
|
||||
|
||||
assert json.loads(path.read_text(encoding="utf-8"))["last_compaction"] == checkpoint
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Hierarchy creation reports failed durable writes instead of returning IDs."""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.hierarchy import HierarchicalSessionStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path):
|
||||
result = HierarchicalSessionStore(session_dir=str(tmp_path))
|
||||
result.create_session("parent", title="Original")
|
||||
assert result.add_message("parent", "user", "original turn")
|
||||
return result
|
||||
|
||||
|
||||
def invoke(store, operation):
|
||||
if operation == "create":
|
||||
return store.create_session("new", title="New")
|
||||
if operation == "fork":
|
||||
return store.fork_session("parent")
|
||||
if operation == "snapshot":
|
||||
return store.create_snapshot("parent", label="Snapshot")
|
||||
return store.import_session({"session_id": "source", "messages": []}, new_session_id="new")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["create", "fork", "snapshot", "import"])
|
||||
def test_failed_replace_raises_instead_of_returning_id(store, monkeypatch, operation):
|
||||
def fail(*args):
|
||||
raise OSError("injected replace failure")
|
||||
|
||||
monkeypatch.setattr(os, "replace", fail)
|
||||
with pytest.raises(OSError):
|
||||
invoke(store, operation)
|
||||
store.invalidate_cache()
|
||||
parent = store.get_extended_session("parent", force_reload=True)
|
||||
assert parent.children_ids == []
|
||||
assert parent.snapshots == []
|
||||
assert [message.content for message in parent.messages] == ["original turn"]
|
||||
assert not store.session_exists("new")
|
||||
|
||||
|
||||
def test_failed_child_save_does_not_register_missing_child(store, monkeypatch):
|
||||
original = os.replace
|
||||
|
||||
def fail_child(source, destination):
|
||||
if os.path.basename(destination) == "child.json":
|
||||
raise OSError("injected child replace failure")
|
||||
return original(source, destination)
|
||||
|
||||
monkeypatch.setattr(os, "replace", fail_child)
|
||||
with pytest.raises(OSError):
|
||||
store.create_session("child", parent_id="parent")
|
||||
assert store.get_children("parent") == []
|
||||
assert not store.session_exists("child")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["create", "fork"])
|
||||
def test_failed_parent_registration_reports_retained_child(store, monkeypatch, operation):
|
||||
original = os.replace
|
||||
|
||||
def fail_parent(source, destination):
|
||||
if os.path.basename(destination) == "parent.json":
|
||||
raise OSError("injected parent replace failure")
|
||||
return original(source, destination)
|
||||
|
||||
monkeypatch.setattr(os, "replace", fail_parent)
|
||||
with pytest.raises(OSError, match="parent"):
|
||||
if operation == "create":
|
||||
store.create_session("child", parent_id="parent")
|
||||
else:
|
||||
store.fork_session("parent")
|
||||
store.invalidate_cache()
|
||||
assert store.get_children("parent") == []
|
||||
children = [row["session_id"] for row in store.list_sessions() if row["session_id"] != "parent"]
|
||||
assert len(children) == 1
|
||||
assert store.get_parent(children[0]) == "parent"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["create", "fork", "snapshot", "import"])
|
||||
def test_successful_operation_returns_persisted_id(store, operation):
|
||||
identifier = invoke(store, operation)
|
||||
assert isinstance(identifier, str) and identifier
|
||||
if operation == "snapshot":
|
||||
assert identifier in [snapshot.id for snapshot in store.get_snapshots("parent")]
|
||||
else:
|
||||
assert store.session_exists(identifier)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("session_id", ["new", "parent"])
|
||||
def test_self_parent_is_rejected_before_any_session_write(store, session_id):
|
||||
parent_path = Path(store.session_dir) / "parent.json"
|
||||
before = parent_path.read_bytes()
|
||||
with pytest.raises(ValueError, match="parent"):
|
||||
store.create_session(session_id, parent_id=session_id)
|
||||
assert parent_path.read_bytes() == before
|
||||
assert not store.session_exists("new")
|
||||
assert store.get_session_tree("parent")["children"] == []
|
||||
@@ -0,0 +1,65 @@
|
||||
"""Backfill callbacks may query the same indexed store without restarting a pass."""
|
||||
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.sqlite_store import SqliteSessionStore
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query", ["route", "search"])
|
||||
def test_corruption_callback_can_query_same_store(tmp_path, monkeypatch, query):
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
assert store.add_message("healthy", "user", "narwhal record")
|
||||
assert store.set_gateway_info("healthy", gateway_session_id="gateway")
|
||||
assert store.get_by_gateway_session("gateway") is not None
|
||||
(tmp_path / "broken.json").write_text("{invalid", encoding="utf-8")
|
||||
store._backfilled = False
|
||||
results, failures = [], []
|
||||
|
||||
def on_corruption(*args):
|
||||
if query == "route":
|
||||
results.append(store.get_by_gateway_session("gateway").session_id)
|
||||
else:
|
||||
results.append([hit.session_id for hit in store.search("narwhal")])
|
||||
|
||||
monkeypatch.setattr(store, "_fire_corruption_hook", on_corruption)
|
||||
|
||||
def backfill():
|
||||
try:
|
||||
store._ensure_backfilled()
|
||||
except BaseException as exc:
|
||||
failures.append(exc)
|
||||
|
||||
worker = threading.Thread(target=backfill, daemon=True)
|
||||
worker.start()
|
||||
worker.join(2)
|
||||
try:
|
||||
assert not worker.is_alive(), "callback query deadlocked during backfill"
|
||||
assert failures == []
|
||||
expected = ["healthy"] if query == "route" else [["healthy"]]
|
||||
assert results == expected
|
||||
assert store._backfilled
|
||||
finally:
|
||||
if not worker.is_alive() and store._conn is not None:
|
||||
store._conn.close()
|
||||
|
||||
|
||||
def test_raised_backfill_releases_running_guard_for_retry(tmp_path, monkeypatch):
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
original = store._reindex_all
|
||||
|
||||
def fail():
|
||||
raise RuntimeError("failed pass")
|
||||
|
||||
monkeypatch.setattr(store, "_reindex_all", fail)
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="failed pass"):
|
||||
store._ensure_backfilled()
|
||||
assert not store._backfilled
|
||||
monkeypatch.setattr(store, "_reindex_all", original)
|
||||
store._ensure_backfilled()
|
||||
assert store._backfilled
|
||||
finally:
|
||||
if store._conn is not None:
|
||||
store._conn.close()
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Indexed recall releases candidate payloads while preserving lineage ranking."""
|
||||
|
||||
import json
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.sqlite_store import SqliteSessionStore
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
def test_search_does_not_retain_every_continuation(tmp_path, monkeypatch, backend):
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
try:
|
||||
for index in range(30):
|
||||
sid = f"continuation{index:02}"
|
||||
assert store.add_message(sid, "user", "needle " + "filler " * 100)
|
||||
assert store.update_session_metadata(sid, lineage_id="conversation")
|
||||
assert store.add_message("other", "user", "needle " + "filler " * 100)
|
||||
assert store.add_message("continuation29", "assistant", "needle best continuation")
|
||||
store._ensure_backfilled()
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
|
||||
class Payload(dict):
|
||||
pass
|
||||
|
||||
original = json.load
|
||||
references = []
|
||||
peak = 0
|
||||
|
||||
def tracked_load(handle, *args, **kwargs):
|
||||
nonlocal peak
|
||||
value = original(handle, *args, **kwargs)
|
||||
if isinstance(value, dict) and str(handle.name).endswith(".json"):
|
||||
value = Payload(value)
|
||||
references.append(weakref.ref(value))
|
||||
peak = max(peak, sum(ref() is not None for ref in references))
|
||||
return value
|
||||
|
||||
monkeypatch.setattr(json, "load", tracked_load)
|
||||
hits = store.search("needle", limit=5)
|
||||
assert len(references) == 31
|
||||
assert {hit.session_id for hit in hits} == {"continuation29", "other"}
|
||||
assert hits[0].session_id == "continuation29"
|
||||
assert peak <= 2
|
||||
finally:
|
||||
store._conn.close()
|
||||
|
||||
|
||||
def test_lineage_winner_keeps_original_tie_order(tmp_path):
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
try:
|
||||
for sid, lineage, count in [("00", "a", 1), ("01", "b", 2), ("02", "a", 2)]:
|
||||
for _ in range(count):
|
||||
assert store.add_message(sid, "user", "needle")
|
||||
assert store.update_session_metadata(sid, lineage_id=lineage)
|
||||
path = tmp_path / f"{sid}.json"
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
data["updated_at"] = "2026-01-01T00:00:00+00:00"
|
||||
path.write_text(json.dumps(data), encoding="utf-8")
|
||||
store._ensure_backfilled()
|
||||
store._fts_available = False
|
||||
assert [hit.session_id for hit in store.search("needle")] == ["01", "02"]
|
||||
finally:
|
||||
store._conn.close()
|
||||
@@ -0,0 +1,82 @@
|
||||
"""An explicitly empty agent transcript must override older saved history."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.api import AGENT_HISTORY_KEY, Session
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
|
||||
|
||||
OLD = [{"role": "user", "content": "old conversation"}]
|
||||
PEER = [{"role": "user", "content": "other agent"}]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def setup(tmp_path, monkeypatch):
|
||||
import praisonaiagents.session.store as store_module
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setenv("PRAISONAI_HOME", str(tmp_path / "home"))
|
||||
directory = str(tmp_path / "sessions")
|
||||
store = DefaultSessionStore(session_dir=directory)
|
||||
monkeypatch.setattr(store_module, "_default_store", store)
|
||||
return Session(session_id="parent"), store, directory, store_module
|
||||
|
||||
|
||||
def test_close_persists_cleared_history_without_erasing_peer(setup, monkeypatch):
|
||||
session, store, directory, store_module = setup
|
||||
store.update_session_metadata("parent", **{AGENT_HISTORY_KEY: {
|
||||
"helper:Assistant": OLD, "peer:Assistant": PEER,
|
||||
}})
|
||||
session._agents["helper:Assistant"] = {
|
||||
"agent": SimpleNamespace(chat_history=[]), "chat_history": OLD,
|
||||
}
|
||||
session.close()
|
||||
fresh = DefaultSessionStore(session_dir=directory)
|
||||
monkeypatch.setattr(store_module, "_default_store", fresh)
|
||||
assert fresh.get_session("parent").metadata[AGENT_HISTORY_KEY] == {
|
||||
"helper:Assistant": [], "peer:Assistant": PEER,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("legacy", [False, True], ids=["memory", "tagged-legacy"])
|
||||
def test_empty_current_history_prevents_fallback_resurrection(setup, legacy):
|
||||
session, store, _, _ = setup
|
||||
store.update_session_metadata("parent", **{AGENT_HISTORY_KEY: {
|
||||
"helper:Assistant": [],
|
||||
}})
|
||||
if legacy:
|
||||
store.set_chat_history("parent_helper:Assistant", OLD)
|
||||
store.update_session_metadata(
|
||||
"parent_helper:Assistant", parent_session_id="parent",
|
||||
agent_key="helper:Assistant",
|
||||
)
|
||||
session._memory = SimpleNamespace(search_short_term=lambda **kwargs: [{
|
||||
"metadata": {"type": "agent_chat_history", "session_id": "parent",
|
||||
"agent_key": "helper:Assistant", "chat_history": OLD},
|
||||
}])
|
||||
assert session._restore_agent_chat_history("helper:Assistant") == []
|
||||
session._restore_agent_chat_histories()
|
||||
assert session._agents["helper:Assistant"]["chat_history"] == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("replacement", [[], [{"role": "user", "content": "new"}]])
|
||||
def test_verified_replacement_succeeds_on_first_write(setup, monkeypatch, replacement):
|
||||
session, store, _, _ = setup
|
||||
store.update_session_metadata("parent", **{AGENT_HISTORY_KEY: {
|
||||
"helper:Assistant": OLD, "peer:Assistant": PEER,
|
||||
}})
|
||||
original = store.merge_session_metadata_map
|
||||
writes = []
|
||||
|
||||
def merge(*args, **kwargs):
|
||||
writes.append((args, kwargs))
|
||||
return original(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(store, "merge_session_metadata_map", merge)
|
||||
assert session._store_agent_history(store, "helper:Assistant", replacement)
|
||||
assert len(writes) == 1
|
||||
assert store.get_session("parent").metadata[AGENT_HISTORY_KEY] == {
|
||||
"helper:Assistant": replacement, "peer:Assistant": PEER,
|
||||
}
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Default imports must not overwrite a session created after the first check."""
|
||||
"""Imports must not overwrite a session created after the first check."""
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from threading import Event
|
||||
@@ -6,13 +6,16 @@ from threading import Event
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
from praisonaiagents.session.sqlite_store import SqliteSessionStore
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", [DefaultSessionStore, SqliteSessionStore])
|
||||
@pytest.mark.parametrize("overwrite", [False, True])
|
||||
def test_import_rechecks_new_peer_session_under_write_lock(tmp_path, monkeypatch, overwrite):
|
||||
@pytest.mark.parametrize("peer_kind", [DefaultSessionStore, SqliteSessionStore])
|
||||
def test_import_rechecks_new_peer_session_under_write_lock(tmp_path, monkeypatch, kind, overwrite, peer_kind):
|
||||
directory = str(tmp_path / "sessions")
|
||||
importer = DefaultSessionStore(session_dir=directory)
|
||||
peer = DefaultSessionStore(session_dir=directory)
|
||||
stores = [kind(session_dir=directory), peer_kind(session_dir=directory)]
|
||||
importer, peer = stores
|
||||
ready, completed = Event(), Event()
|
||||
original = importer._save_imported_session
|
||||
|
||||
@@ -35,18 +38,25 @@ def test_import_rechecks_new_peer_session_under_write_lock(tmp_path, monkeypatch
|
||||
finally:
|
||||
completed.set()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
future = pool.submit(create_peer)
|
||||
report = importer.import_sessions(payload, overwrite=overwrite)
|
||||
future.result(timeout=15)
|
||||
fresh = DefaultSessionStore(session_dir=directory)
|
||||
if overwrite:
|
||||
assert report.imported == 1
|
||||
assert not report.skipped
|
||||
assert fresh.get_chat_history("shared") == [{"role": "user", "content": "imported"}]
|
||||
else:
|
||||
assert report.imported == 0
|
||||
assert report.skipped == [{
|
||||
"session_id": "shared", "reason": "already exists (use overwrite)",
|
||||
}]
|
||||
assert fresh.get_chat_history("shared") == [{"role": "user", "content": "peer-created"}]
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
future = pool.submit(create_peer)
|
||||
report = importer.import_sessions(payload, overwrite=overwrite)
|
||||
future.result(timeout=15)
|
||||
fresh = kind(session_dir=directory)
|
||||
stores.append(fresh)
|
||||
if overwrite:
|
||||
assert report.imported == 1
|
||||
assert not report.skipped
|
||||
assert fresh.get_chat_history("shared") == [{"role": "user", "content": "imported"}]
|
||||
else:
|
||||
assert report.imported == 0
|
||||
assert report.skipped == [{
|
||||
"session_id": "shared", "reason": "already exists (use overwrite)",
|
||||
}]
|
||||
assert fresh.get_chat_history("shared") == [{"role": "user", "content": "peer-created"}]
|
||||
finally:
|
||||
for store in stores:
|
||||
conn = getattr(store, "_conn", None)
|
||||
if conn is not None:
|
||||
conn.close()
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Independent WAL readers must never observe half an index replacement."""
|
||||
|
||||
import sqlite3
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session import SqliteSessionStore
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", [False, True])
|
||||
def test_wal_reader_retains_session_during_or_after_failed_index_refresh(tmp_path, monkeypatch, failure):
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path))
|
||||
deleted = threading.Event()
|
||||
resume = threading.Event()
|
||||
worker = None
|
||||
reader = None
|
||||
outcomes = []
|
||||
try:
|
||||
assert store.add_message("session", "user", "needle old")
|
||||
conn = store._conn
|
||||
session = store._read_session_fresh("session")
|
||||
session.messages[0].content = "needle new"
|
||||
# Refresh indexes the committed file generation, not a caller's
|
||||
# possibly stale object. Persist the new transcript without refreshing
|
||||
# its index so the WAL reader still starts with the old index row.
|
||||
assert store._atomic_write_json(
|
||||
store._get_session_path("session"), session.to_dict()
|
||||
)
|
||||
reader = sqlite3.connect(store.db_path, isolation_level=None)
|
||||
|
||||
class PausedConnection:
|
||||
@property
|
||||
def in_transaction(self):
|
||||
return conn.in_transaction
|
||||
|
||||
def execute(self, sql, parameters=()):
|
||||
if failure and sql.startswith("INSERT INTO session_fts"):
|
||||
raise sqlite3.OperationalError("injected index insert failure")
|
||||
result = conn.execute(sql, parameters)
|
||||
if sql.startswith("DELETE FROM session_fts"):
|
||||
deleted.set()
|
||||
if not resume.wait(5):
|
||||
raise TimeoutError("reader did not resume index writer")
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(store, "_connect", lambda: PausedConnection())
|
||||
worker = threading.Thread(target=lambda: outcomes.append(store._index_session(session)))
|
||||
worker.start()
|
||||
assert deleted.wait(5)
|
||||
# This is a second real connection, outside the store's Python lock.
|
||||
during = reader.execute("SELECT content FROM session_fts WHERE session_id = ?", ("session",)).fetchall()
|
||||
resume.set()
|
||||
worker.join(5)
|
||||
assert not worker.is_alive()
|
||||
after = reader.execute("SELECT content FROM session_fts WHERE session_id = ?", ("session",)).fetchall()
|
||||
assert during == [("needle old",)]
|
||||
assert after == [("needle old" if failure else "needle new",)]
|
||||
assert outcomes == [not failure]
|
||||
assert not conn.in_transaction
|
||||
finally:
|
||||
resume.set()
|
||||
if worker is not None:
|
||||
worker.join(5)
|
||||
if reader is not None:
|
||||
reader.close()
|
||||
if store._conn is not None:
|
||||
store._conn.close()
|
||||
@@ -0,0 +1,453 @@
|
||||
"""Unreadable transcript bytes must not interrupt recall of valid sessions."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session import SqliteSessionStore
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
def test_wal_writer_can_finish_while_matching_ids_are_spooled(tmp_path, monkeypatch, backend):
|
||||
import tempfile
|
||||
import threading
|
||||
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path))
|
||||
threads = []
|
||||
completed = []
|
||||
failures = []
|
||||
try:
|
||||
assert store.add_message("first", "user", "needle")
|
||||
assert store.add_message("other", "user", "needle")
|
||||
store._ensure_backfilled()
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
original = tempfile.TemporaryFile
|
||||
|
||||
class Spool:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.handle = original(*args, **kwargs)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
self.handle.close()
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.handle)
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self.handle, name)
|
||||
|
||||
def write(self, content):
|
||||
if not threads:
|
||||
done = threading.Event()
|
||||
|
||||
def write_index():
|
||||
try:
|
||||
store._deindex_session("other")
|
||||
except Exception as exc:
|
||||
failures.append(exc)
|
||||
finally:
|
||||
done.set()
|
||||
|
||||
thread = threading.Thread(target=write_index)
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
completed.append(done.wait(2))
|
||||
return self.handle.write(content)
|
||||
|
||||
monkeypatch.setattr(tempfile, "TemporaryFile", Spool)
|
||||
hits = store.search("needle")
|
||||
for thread in threads:
|
||||
thread.join(3)
|
||||
assert completed == [True]
|
||||
assert not failures
|
||||
assert {hit.session_id for hit in hits} == {"first", "other"}
|
||||
finally:
|
||||
for thread in threads:
|
||||
thread.join(3)
|
||||
store._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["json", "fts", "like"])
|
||||
def test_search_skips_invalid_utf8_candidate(tmp_path, backend):
|
||||
store = (
|
||||
DefaultSessionStore(session_dir=str(tmp_path))
|
||||
if backend == "json"
|
||||
else SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
)
|
||||
try:
|
||||
assert store.add_message("bad", "user", "migration needle")
|
||||
assert store.add_message("valid", "user", "migration needle")
|
||||
if backend == "fts" and not store._fts_available:
|
||||
pytest.skip("SQLite build lacks FTS5")
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
assert {hit.session_id for hit in store.search("needle")} == {"bad", "valid"}
|
||||
bad = tmp_path / "bad.json"
|
||||
payload = b"\xff\xfe"
|
||||
bad.write_bytes(payload)
|
||||
hits = store.search("needle")
|
||||
assert [hit.session_id for hit in hits] == ["valid"]
|
||||
assert hits[0].messages[0]["content"] == "migration needle"
|
||||
assert bad.read_bytes() == payload
|
||||
finally:
|
||||
conn = getattr(store, "_conn", None)
|
||||
if conn is not None:
|
||||
conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
def test_unreadable_candidates_do_not_hide_later_valid_hit(tmp_path, backend):
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
try:
|
||||
for index in range(25):
|
||||
assert store.add_message(f"bad{index:02}", "user", "needle")
|
||||
assert store.add_message("valid", "user", "needle " + "filler " * 100)
|
||||
if backend == "fts" and not store._fts_available:
|
||||
pytest.skip("SQLite build lacks FTS5")
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
store._ensure_backfilled()
|
||||
candidates = store._candidate_ids("needle", 5)
|
||||
assert len(candidates) == 25 and "valid" not in candidates
|
||||
for index in range(25):
|
||||
(tmp_path / f"bad{index:02}.json").write_bytes(b"\xff\xfe")
|
||||
assert [hit.session_id for hit in store.search("needle")] == ["valid"]
|
||||
finally:
|
||||
if store._conn is not None:
|
||||
store._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["json", "fts"])
|
||||
def test_search_reports_corruption_without_moving_file(tmp_path, monkeypatch, caplog, backend):
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path)) if backend == "json" else SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
events = []
|
||||
monkeypatch.setattr(store, "_fire_corruption_hook", lambda *args: events.append(args))
|
||||
try:
|
||||
assert store.add_message("bad", "user", "needle")
|
||||
assert store.add_message("valid", "user", "needle")
|
||||
store.search("needle")
|
||||
bad = tmp_path / "bad.json"
|
||||
bad.write_bytes(b"\xff\xfe")
|
||||
assert [hit.session_id for hit in store.search("needle")] == ["valid"]
|
||||
assert events and events[0][0] == "bad" and events[0][2] is None
|
||||
assert "Skipping unreadable session file" in caplog.text
|
||||
assert bad.read_bytes() == b"\xff\xfe"
|
||||
finally:
|
||||
conn = getattr(store, "_conn", None)
|
||||
if conn is not None:
|
||||
conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
def test_broad_query_stops_loading_after_readable_allowance(tmp_path, monkeypatch, backend):
|
||||
import builtins
|
||||
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
try:
|
||||
for index in range(60):
|
||||
assert store.add_message(f"s{index:02}", "user", "needle")
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
store._ensure_backfilled()
|
||||
opened = []
|
||||
original = builtins.open
|
||||
|
||||
def tracked(path, *args, **kwargs):
|
||||
if str(path).endswith(".json"):
|
||||
opened.append(str(path))
|
||||
return original(path, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "open", tracked)
|
||||
assert len(store.search("needle", limit=5)) == 5
|
||||
assert len(opened) == 25
|
||||
finally:
|
||||
store._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
def test_transcript_reads_do_not_pin_wal_checkpoint(tmp_path, monkeypatch, backend):
|
||||
import builtins
|
||||
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path))
|
||||
peer = SqliteSessionStore(session_dir=str(tmp_path))
|
||||
checkpoints = []
|
||||
try:
|
||||
assert store.add_message("first", "user", "needle")
|
||||
assert store.add_message("second", "user", "needle")
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
store._ensure_backfilled()
|
||||
peer_conn = peer._connect()
|
||||
peer_conn.execute("PRAGMA busy_timeout=0")
|
||||
original = builtins.open
|
||||
|
||||
def opened(path, *args, **kwargs):
|
||||
if str(path).endswith("first.json") and not checkpoints:
|
||||
assert peer.add_message("late", "user", "needle")
|
||||
checkpoints.append(peer_conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone())
|
||||
return original(path, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "open", opened)
|
||||
assert {hit.session_id for hit in store.search("needle")} == {"first", "second"}
|
||||
assert checkpoints and checkpoints[0][0] == 0
|
||||
finally:
|
||||
for current in (store, peer):
|
||||
if current._conn is not None:
|
||||
current._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("journal", ["delete", "memory"])
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
def test_search_does_not_copy_whole_index(tmp_path, monkeypatch, journal, backend):
|
||||
import sqlite3
|
||||
|
||||
backups = []
|
||||
|
||||
class Connection(sqlite3.Connection):
|
||||
def backup(self, *args, **kwargs):
|
||||
backups.append(True)
|
||||
return super().backup(*args, **kwargs)
|
||||
|
||||
original = sqlite3.connect
|
||||
|
||||
def connect(*args, **kwargs):
|
||||
kwargs["factory"] = Connection
|
||||
return original(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(sqlite3, "connect", connect)
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:" if journal == "memory" else None)
|
||||
try:
|
||||
assert store.add_message("valid", "user", "needle")
|
||||
assert store.add_message("unmatched", "user", "filler " * 10000)
|
||||
store._ensure_backfilled()
|
||||
if journal == "delete":
|
||||
store._conn.execute("PRAGMA journal_mode=DELETE")
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
assert [hit.session_id for hit in store.search("needle")] == ["valid"]
|
||||
assert backups == []
|
||||
finally:
|
||||
store._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
def test_index_deletion_during_read_does_not_shift_candidates(tmp_path, monkeypatch, backend):
|
||||
import builtins
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path))
|
||||
peer = SqliteSessionStore(session_dir=str(tmp_path))
|
||||
try:
|
||||
for index in range(25):
|
||||
assert store.add_message(f"bad{index:02}", "user", "needle")
|
||||
assert store.add_message("valid", "user", "needle " + "filler " * 100)
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
store._ensure_backfilled()
|
||||
peer_conn = peer._connect()
|
||||
for index in range(25):
|
||||
(tmp_path / f"bad{index:02}.json").write_bytes(b"\xff\xfe")
|
||||
changed = []
|
||||
original = builtins.open
|
||||
|
||||
def opened(path, *args, **kwargs):
|
||||
if str(path).endswith("bad00.json") and not changed:
|
||||
peer_conn.execute("DELETE FROM session_fts WHERE session_id = ?", ("bad00",))
|
||||
changed.append(True)
|
||||
return original(path, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "open", opened)
|
||||
assert [hit.session_id for hit in store.search("needle")] == ["valid"]
|
||||
assert changed
|
||||
finally:
|
||||
for current in (store, peer):
|
||||
if current._conn is not None:
|
||||
current._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["list", "agent", "gateway", "gateway_agent"])
|
||||
def test_directory_reads_skip_invalid_utf8(tmp_path, monkeypatch, operation):
|
||||
import praisonaiagents.session.store as module
|
||||
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path))
|
||||
assert store.add_message("valid", "user", "durable history")
|
||||
assert store.set_agent_info("valid", agent_name="support")
|
||||
assert store.set_gateway_info("valid", gateway_session_id="gateway", agent_id="agent")
|
||||
assert store.update_session_metadata("valid", lineage_id="thread")
|
||||
bad = tmp_path / "bad.json"
|
||||
payload = b"\xff\xfe"
|
||||
bad.write_bytes(payload)
|
||||
original = module.os.listdir
|
||||
monkeypatch.setattr(module.os, "listdir", lambda path: [bad.name, "valid.json"] if str(path) == str(tmp_path) else original(path))
|
||||
|
||||
if operation == "list":
|
||||
assert [row["session_id"] for row in store.list_sessions()] == ["valid"]
|
||||
elif operation == "agent":
|
||||
assert store.list_sessions_by_agent("support") == ["valid"]
|
||||
elif operation == "gateway":
|
||||
assert store.get_by_gateway_session("gateway").session_id == "valid"
|
||||
elif operation == "gateway_agent":
|
||||
assert store.list_sessions_by_gateway_agent("agent") == ["valid"]
|
||||
assert bad.read_bytes() == payload
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["export_all", "lineage"])
|
||||
@pytest.mark.parametrize("payload", [b"\xff\xfe", b"{invalid"])
|
||||
def test_export_refuses_to_report_incomplete_backup(tmp_path, operation, payload):
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path))
|
||||
assert store.add_message("valid", "user", "durable history")
|
||||
assert store.update_session_metadata("valid", lineage_id="thread")
|
||||
bad = tmp_path / "bad.json"
|
||||
bad.write_bytes(payload)
|
||||
export = store.export_all if operation == "export_all" else lambda: store.export_session("valid")
|
||||
with pytest.raises(OSError, match="Incomplete session export"):
|
||||
export()
|
||||
assert bad.read_bytes() == payload
|
||||
assert store.get_chat_history("valid")[0]["content"] == "durable history"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
def test_continuations_do_not_hide_another_lineage(tmp_path, backend):
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
try:
|
||||
for index in range(25):
|
||||
sid = f"continuation{index:02}"
|
||||
assert store.add_message(sid, "user", "needle")
|
||||
assert store.update_session_metadata(sid, lineage_id="conversation")
|
||||
assert store.add_message("other", "user", "needle " + "filler " * 100)
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
store._ensure_backfilled()
|
||||
assert "other" not in store._candidate_ids("needle", 5)
|
||||
hits = store.search("needle", limit=5)
|
||||
assert len(hits) == 2
|
||||
assert "other" in {hit.session_id for hit in hits}
|
||||
finally:
|
||||
store._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
@pytest.mark.parametrize("journal", ["wal", "delete", "memory"])
|
||||
def test_search_snapshots_matching_ids_in_bounded_batches(tmp_path, monkeypatch, backend, journal):
|
||||
import sqlite3
|
||||
|
||||
reads = []
|
||||
all_reads = []
|
||||
batches = []
|
||||
|
||||
class Cursor(sqlite3.Cursor):
|
||||
tracked = False
|
||||
|
||||
def execute(self, sql, parameters=()):
|
||||
self.tracked = sql.startswith("SELECT session_id FROM session_fts")
|
||||
return super().execute(sql, parameters)
|
||||
|
||||
def fetchall(self):
|
||||
rows = super().fetchall()
|
||||
if self.tracked:
|
||||
all_reads.append(len(rows))
|
||||
return rows
|
||||
|
||||
def fetchone(self):
|
||||
row = super().fetchone()
|
||||
if self.tracked and row is not None:
|
||||
reads.append(row[0])
|
||||
return row
|
||||
|
||||
def fetchmany(self, size=None):
|
||||
rows = super().fetchmany(size) if size is not None else super().fetchmany()
|
||||
if self.tracked:
|
||||
batches.append(len(rows))
|
||||
return rows
|
||||
|
||||
class Connection(sqlite3.Connection):
|
||||
def execute(self, sql, parameters=()):
|
||||
return self.cursor(factory=Cursor).execute(sql, parameters)
|
||||
|
||||
original = sqlite3.connect
|
||||
|
||||
def connect(*args, **kwargs):
|
||||
kwargs["factory"] = Connection
|
||||
return original(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(sqlite3, "connect", connect)
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:" if journal == "memory" else None)
|
||||
try:
|
||||
for index in range(60):
|
||||
assert store.add_message(f"s{index:02}", "user", "needle")
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
store._ensure_backfilled()
|
||||
if journal == "delete":
|
||||
assert store._conn.execute("PRAGMA journal_mode=DELETE").fetchone()[0] == "delete"
|
||||
assert len(store.search("needle", limit=5)) == 5
|
||||
assert all_reads == []
|
||||
assert reads == []
|
||||
assert sum(batches) == 60
|
||||
assert max(batches) <= 128
|
||||
finally:
|
||||
store._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["fts", "like"])
|
||||
@pytest.mark.parametrize("pause", ["file", "hook"])
|
||||
@pytest.mark.parametrize("journal", ["wal", "delete", "memory"])
|
||||
def test_index_writer_can_finish_during_transcript_work(tmp_path, monkeypatch, backend, pause, journal):
|
||||
import builtins
|
||||
import threading
|
||||
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:" if journal == "memory" else None)
|
||||
threads = []
|
||||
completed_during_pause = []
|
||||
failures = []
|
||||
try:
|
||||
assert store.add_message("first", "user", "needle")
|
||||
assert store.add_message("other", "user", "needle " + "filler " * 100)
|
||||
if backend == "like":
|
||||
store._fts_available = False
|
||||
store._ensure_backfilled()
|
||||
if journal == "delete":
|
||||
assert store._conn.execute("PRAGMA journal_mode=DELETE").fetchone()[0] == "delete"
|
||||
|
||||
def pause_for_writer():
|
||||
if threads:
|
||||
return
|
||||
done = threading.Event()
|
||||
|
||||
def write():
|
||||
try:
|
||||
store._deindex_session("other")
|
||||
except Exception as exc:
|
||||
failures.append(exc)
|
||||
finally:
|
||||
done.set()
|
||||
|
||||
thread = threading.Thread(target=write)
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
completed_during_pause.append(done.wait(2))
|
||||
|
||||
if pause == "file":
|
||||
original = builtins.open
|
||||
|
||||
def opened(path, *args, **kwargs):
|
||||
if str(path).endswith("first.json"):
|
||||
pause_for_writer()
|
||||
return original(path, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "open", opened)
|
||||
else:
|
||||
(tmp_path / "first.json").write_bytes(b"\xff\xfe")
|
||||
monkeypatch.setattr(store, "_fire_corruption_hook", lambda *args: pause_for_writer())
|
||||
|
||||
hits = store.search("needle")
|
||||
for thread in threads:
|
||||
thread.join(3)
|
||||
assert completed_during_pause == [True]
|
||||
assert not failures
|
||||
assert "other" in {hit.session_id for hit in hits}
|
||||
finally:
|
||||
for thread in threads:
|
||||
thread.join(3)
|
||||
store._conn.close()
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Salvaged JSON content and one corrupt file must not break session loading."""
|
||||
|
||||
from pathlib import Path
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.store import DefaultSessionStore, SessionMessage
|
||||
|
||||
|
||||
@pytest.mark.parametrize("incoming,expected", [
|
||||
({"b": [2], "a": [1]}, 1),
|
||||
({"a": [[1]], "b": [2]}, 2),
|
||||
])
|
||||
def test_spill_key_preserves_object_order_and_array_structure(tmp_path, monkeypatch, incoming, expected):
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path / "sessions"))
|
||||
monkeypatch.setattr(store, "_spill_dir", lambda: str(tmp_path / "spill"))
|
||||
assert store.add_message("s", "user", {"a": [1], "b": [2]})
|
||||
timestamp = store.get_session("s").messages[0].timestamp
|
||||
spill = store._spill("s", [SessionMessage("user", incoming, timestamp=timestamp)])
|
||||
|
||||
assert len(store.get_chat_history("s")) == expected
|
||||
assert not Path(spill).exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("existing", [False, True])
|
||||
def test_deep_json_content_does_not_abort_recovery(tmp_path, monkeypatch, existing):
|
||||
directory = tmp_path / "sessions"
|
||||
store = DefaultSessionStore(session_dir=str(directory), retention="keep_all")
|
||||
monkeypatch.setattr(store, "_spill_dir", lambda: str(tmp_path / "spill"))
|
||||
# Valid JSON whose recursive key conversion exceeded Python's stack limit.
|
||||
content = json.loads("[" * 500 + "0" + "]" * 500)
|
||||
message = SessionMessage(role="user", content=content, timestamp=123)
|
||||
if existing:
|
||||
assert store.add_message("s", "user", content)
|
||||
message = store.get_session("s").messages[0]
|
||||
spill = store._spill("s", [message])
|
||||
neighbor = store._spill("s", [SessionMessage("user", "neighbor", timestamp=124)])
|
||||
|
||||
history = store.get_chat_history("s")
|
||||
|
||||
assert len(history) == 2
|
||||
assert json.dumps(history[0]["content"]) == json.dumps(content)
|
||||
assert history[1]["content"] == "neighbor"
|
||||
assert not Path(spill).exists()
|
||||
assert not Path(neighbor).exists()
|
||||
reopened = DefaultSessionStore(session_dir=str(directory), retention="keep_all")
|
||||
assert json.dumps(reopened.get_chat_history("s")) == json.dumps(history)
|
||||
|
||||
|
||||
def test_json_parse_recursion_failure_preserves_spill_and_recovers_neighbor(tmp_path, monkeypatch):
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path / "sessions"))
|
||||
directory = tmp_path / "spill"
|
||||
monkeypatch.setattr(store, "_spill_dir", lambda: str(directory))
|
||||
neighbor = store._spill("s", [SessionMessage("user", "neighbor", timestamp=124)])
|
||||
bad = directory / "s.000.json"
|
||||
raw = json.dumps({"session_id": "s", "messages": []})
|
||||
bad.write_text(raw, encoding="utf-8")
|
||||
load = json.load
|
||||
|
||||
def decoder(file, *args, **kwargs):
|
||||
if Path(file.name) == bad:
|
||||
raise RecursionError("decoder nesting limit")
|
||||
return load(file, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(json, "load", decoder)
|
||||
|
||||
assert store.get_chat_history("s") == [{"role": "user", "content": "neighbor"}]
|
||||
assert bad.read_text(encoding="utf-8") == raw
|
||||
assert not Path(neighbor).exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("existing", [False, True])
|
||||
def test_structured_content_recovers_once(tmp_path, monkeypatch, existing):
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path / "sessions"))
|
||||
monkeypatch.setattr(store, "_spill_dir", lambda: str(tmp_path / "spill"))
|
||||
content = [{"type": "text", "text": "hello"}, {
|
||||
"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="},
|
||||
}]
|
||||
if existing:
|
||||
assert store.add_message("s", "user", content)
|
||||
msg = store.get_session("s").messages[0]
|
||||
else:
|
||||
msg = SessionMessage(role="user", content=content, timestamp=123)
|
||||
spill = store._spill("s", [msg])
|
||||
assert spill is not None
|
||||
history = store.get_chat_history("s")
|
||||
assert history == [{"role": "user", "content": content}]
|
||||
assert not Path(spill).exists()
|
||||
assert store.get_chat_history("s") == history
|
||||
|
||||
|
||||
def test_invalid_utf8_spill_does_not_block_valid_neighbor(tmp_path, monkeypatch):
|
||||
store = DefaultSessionStore(session_dir=str(tmp_path / "sessions"))
|
||||
directory = tmp_path / "spill"
|
||||
monkeypatch.setattr(store, "_spill_dir", lambda: str(directory))
|
||||
spill = store._spill("s", [SessionMessage(role="user", content="recovered", timestamp=123)])
|
||||
bad = directory / "s.000.json"
|
||||
bad.write_bytes(b"\xff\xfeinvalid")
|
||||
assert store.get_chat_history("s") == [{"role": "user", "content": "recovered"}]
|
||||
assert not Path(spill).exists()
|
||||
assert bad.read_bytes() == b"\xff\xfeinvalid"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ["object", "array"])
|
||||
@pytest.mark.parametrize("previous,incoming", [(True, 1), (1, 1.0), (False, 0), (True, True)])
|
||||
def test_spill_identity_preserves_json_scalar_types(tmp_path, monkeypatch, shape, previous, incoming):
|
||||
directory = tmp_path / "sessions"
|
||||
store = DefaultSessionStore(session_dir=str(directory))
|
||||
monkeypatch.setattr(store, "_spill_dir", lambda: str(tmp_path / "spill"))
|
||||
|
||||
def content(value):
|
||||
return {"nested": {"value": value}} if shape == "object" else [{"value": value}]
|
||||
|
||||
assert store.add_message("s", "user", content(previous))
|
||||
timestamp = store.get_session("s").messages[0].timestamp
|
||||
spill = store._spill("s", [SessionMessage(role="user", content=content(incoming), timestamp=timestamp)])
|
||||
assert spill is not None
|
||||
history = store.get_chat_history("s")
|
||||
expected = 1 if type(previous) is type(incoming) else 2
|
||||
assert len(history) == expected
|
||||
|
||||
reopened = DefaultSessionStore(session_dir=str(directory))
|
||||
persisted = reopened.get_chat_history("s")
|
||||
assert len(persisted) == expected
|
||||
values = [m["content"]["nested"]["value"] if shape == "object" else m["content"][0]["value"] for m in persisted]
|
||||
assert [type(value) for value in values] == ([type(previous)] if expected == 1 else [type(previous), type(incoming)])
|
||||
assert not Path(spill).exists()
|
||||
@@ -1702,6 +1702,54 @@ class TestCompactedHistoryRecall:
|
||||
hits = store.search("XZ99-SECRETVALUE")
|
||||
assert {h.session_id for h in hits} == {"good"}
|
||||
|
||||
def test_invalid_utf8_transcript_does_not_abort_search(self):
|
||||
"""A transcript with invalid UTF-8 bytes must be skipped during the
|
||||
direct JSON scan, not raise ``UnicodeDecodeError`` and abort search
|
||||
across the remaining valid sessions (Issue #5556)."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
store = DefaultSessionStore(session_dir=tmpdir, active_window=4)
|
||||
store.add_user_message("good", "the deploy token is XZ99-SECRETVALUE")
|
||||
store.add_user_message("bad", "the deploy token is XZ99-SECRETVALUE")
|
||||
# Corrupt one transcript with bytes that are not valid UTF-8.
|
||||
with open(os.path.join(tmpdir, "bad.json"), "wb") as f:
|
||||
f.write(b"\xff\xfe")
|
||||
|
||||
hits = store.search("XZ99-SECRETVALUE")
|
||||
assert {h.session_id for h in hits} == {"good"}
|
||||
|
||||
def test_sqlite_invalid_utf8_candidate_does_not_abort_search(self):
|
||||
"""An indexed candidate whose transcript is invalid UTF-8 must be
|
||||
skipped during candidate loading, not raise ``UnicodeDecodeError``
|
||||
(Issue #5556). Deterministically exercises both the FTS5 and the LIKE
|
||||
fallback candidate branches: ``force_like`` patches ``_init_schema`` to
|
||||
report FTS5 unavailable so the LIKE path is covered even on a SQLite
|
||||
build that ships FTS5."""
|
||||
from praisonaiagents.session.sqlite_store import SqliteSessionStore
|
||||
|
||||
for force_like in (False, True):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
store = SqliteSessionStore(session_dir=tmpdir, active_window=4)
|
||||
if force_like:
|
||||
original_init = store._init_schema
|
||||
|
||||
def _like_only(conn, _orig=original_init):
|
||||
_orig(conn)
|
||||
return False
|
||||
|
||||
store._init_schema = _like_only
|
||||
store.add_user_message("good", "the deploy token is XZ99-SECRETVALUE")
|
||||
store.add_user_message("bad", "the deploy token is XZ99-SECRETVALUE")
|
||||
# Confirm we are really exercising the intended candidate branch.
|
||||
assert store._fts_available is (not force_like)
|
||||
# Warm the index so both ids are indexed candidates.
|
||||
store.search("XZ99-SECRETVALUE")
|
||||
# Corrupt one candidate's transcript on disk.
|
||||
with open(store._get_session_path("bad"), "wb") as f:
|
||||
f.write(b"\xff\xfe")
|
||||
|
||||
hits = store.search("XZ99-SECRETVALUE")
|
||||
assert {h.session_id for h in hits} == {"good"}
|
||||
|
||||
def test_sqlite_upgrade_rebuilds_stale_archived_index(self):
|
||||
"""An index built by a pre-#5031 release (archived turns absent) must be
|
||||
rebuilt once on upgrade so archived-only queries work immediately,
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
"""Durable snapshot pooling must preserve portable and independent transcripts."""
|
||||
|
||||
import copy
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.hierarchy import (
|
||||
ExtendedSessionData, HierarchicalSessionStore, SessionSnapshot,
|
||||
)
|
||||
from praisonaiagents.session.store import SessionData, SessionMessage
|
||||
|
||||
|
||||
def test_repeated_snapshots_share_archive_on_disk_and_export_in_full(tmp_path):
|
||||
source = SessionData(
|
||||
session_id="source",
|
||||
messages=[SessionMessage("user", "live", timestamp=1)],
|
||||
archived_messages=[SessionMessage("user", "archive-marker-" + str(i) + "x" * 2048,
|
||||
timestamp=i) for i in range(100)],
|
||||
).to_dict()
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
sid = store.import_session(source, new_session_id="s")
|
||||
path = tmp_path / "s.json"
|
||||
initial = path.stat().st_size
|
||||
ids = [store.create_snapshot(sid) for _ in range(10)]
|
||||
assert store.add_message(sid, "user", "later")
|
||||
|
||||
assert path.read_text(encoding="utf-8").count("archive-marker-0x") == 2
|
||||
assert path.stat().st_size < initial * 3
|
||||
reopened = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
exported = reopened.export_session(sid)
|
||||
assert "snapshot_storage" not in exported
|
||||
assert all(s["transcript"]["archived_messages"] == source["archived_messages"]
|
||||
for s in exported["snapshots"])
|
||||
assert all("transcript_refs" not in s for s in exported["snapshots"])
|
||||
imported = reopened.import_session(json.loads(json.dumps(exported)), new_session_id="copy")
|
||||
assert reopened.revert_to_snapshot(imported, ids[0])
|
||||
assert reopened.export_session(imported)["messages"] == source["messages"]
|
||||
|
||||
|
||||
def test_ordinary_write_does_not_expand_portable_snapshot_dicts(tmp_path, monkeypatch):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
assert store.add_message("s", "user", "first")
|
||||
snapshot_id = store.create_snapshot("s")
|
||||
|
||||
def no_portable_copy(self):
|
||||
raise AssertionError("ordinary write expanded a portable snapshot")
|
||||
|
||||
monkeypatch.setattr(SessionSnapshot, "to_dict", no_portable_copy)
|
||||
assert store.add_message("s", "user", "second")
|
||||
assert store.revert_to_snapshot("s", snapshot_id)
|
||||
|
||||
|
||||
def test_positional_reader_rejects_pooled_snapshot_without_losing_real_flag(tmp_path):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
assert store.add_message("s", "user", "first")
|
||||
snapshot_id = store.create_snapshot("s")
|
||||
record = json.loads((tmp_path / "s.json").read_text(encoding="utf-8"))
|
||||
legacy_view = SessionSnapshot.from_dict(record["snapshots"][0])
|
||||
assert legacy_view.transcript is None
|
||||
assert legacy_view.invalidated is True
|
||||
assert store.revert_to_snapshot("s", snapshot_id)
|
||||
record["snapshots"][0]["transcript_invalidated"] = True
|
||||
assert ExtendedSessionData.from_dict(record).snapshots[0].invalidated is True
|
||||
|
||||
|
||||
def test_legacy_transcripts_upgrade_and_pooled_values_remain_independent(tmp_path):
|
||||
transcript = {"messages": [{"role": "user", "content": [{"text": "old"}], "timestamp": 1}],
|
||||
"archived_messages": [], "last_compaction": None}
|
||||
record = ExtendedSessionData(session_id="s", snapshots=[
|
||||
SessionSnapshot(id="a", transcript=copy.deepcopy(transcript)),
|
||||
SessionSnapshot(id="b", transcript=copy.deepcopy(transcript)),
|
||||
SessionSnapshot(id="empty", transcript={}),
|
||||
]).to_dict()
|
||||
(tmp_path / "s.json").write_text(json.dumps(record), encoding="utf-8")
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
assert store.add_message("s", "user", "new")
|
||||
reopened = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
snapshots = reopened.get_snapshots("s")
|
||||
snapshots[0].transcript["messages"][0]["content"][0]["text"] = "changed"
|
||||
assert snapshots[1].transcript == transcript
|
||||
assert snapshots[2].transcript == {}
|
||||
assert reopened.export_session("s")["snapshots"][0]["transcript"] == transcript
|
||||
|
||||
|
||||
@pytest.mark.parametrize("damage", ["missing", "version", "descriptor"])
|
||||
def test_invalid_pool_preserves_durable_record_on_write(tmp_path, damage):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
assert store.add_message("s", "user", "first")
|
||||
store.create_snapshot("s")
|
||||
path = tmp_path / "s.json"
|
||||
record = json.loads(path.read_text(encoding="utf-8"))
|
||||
if damage == "missing":
|
||||
record["snapshot_storage"]["records"] = {}
|
||||
elif damage == "version":
|
||||
record["snapshot_storage"]["version"] = 2
|
||||
else:
|
||||
record["snapshots"][0]["transcript_refs"]["messages"] = {"bad": "ref"}
|
||||
path.write_text(json.dumps(record), encoding="utf-8")
|
||||
before = path.read_bytes()
|
||||
assert store.add_message("s", "user", "must not replace the record") is False
|
||||
assert path.read_bytes() == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("missing", ["messages", "archived_messages", "last_compaction"])
|
||||
def test_incomplete_pool_cannot_replace_live_history(tmp_path, missing):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
assert store.add_message("s", "user", "first")
|
||||
snapshot_id = store.create_snapshot("s")
|
||||
assert store.add_message("s", "user", "later")
|
||||
path = tmp_path / "s.json"
|
||||
record = json.loads(path.read_text(encoding="utf-8"))
|
||||
del record["snapshots"][0]["transcript_refs"][missing]
|
||||
path.write_text(json.dumps(record), encoding="utf-8")
|
||||
before = path.read_bytes()
|
||||
with pytest.raises(ValueError, match="incomplete snapshot"):
|
||||
ExtendedSessionData.from_dict(record)
|
||||
assert store.revert_to_snapshot("s", snapshot_id) is False
|
||||
assert path.read_bytes() == before
|
||||
assert store.add_message("s", "user", "must not overwrite") is False
|
||||
assert path.read_bytes() == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("damage", ["message_descriptor", "archive_descriptor", "message_record", "checkpoint_record"])
|
||||
def test_wrong_reference_types_cannot_replace_live_history(tmp_path, damage):
|
||||
store = HierarchicalSessionStore(session_dir=str(tmp_path), retention="keep_all")
|
||||
assert store.add_message("s", "user", "first")
|
||||
snapshot_id = store.create_snapshot("s")
|
||||
assert store.add_message("s", "user", "later")
|
||||
path = tmp_path / "s.json"
|
||||
record = json.loads(path.read_text(encoding="utf-8"))
|
||||
records = record["snapshot_storage"]["records"]
|
||||
refs = record["snapshots"][0]["transcript_refs"]
|
||||
if damage.endswith("descriptor"):
|
||||
records["damaged"] = []
|
||||
field = "messages" if damage == "message_descriptor" else "archived_messages"
|
||||
refs[field] = {"value": "damaged"}
|
||||
elif damage == "message_record":
|
||||
records[refs["messages"]["items"][0]] = "not a message"
|
||||
else:
|
||||
records[refs["last_compaction"]["value"]] = []
|
||||
path.write_text(json.dumps(record), encoding="utf-8")
|
||||
before = path.read_bytes()
|
||||
with pytest.raises(ValueError, match="snapshot"):
|
||||
ExtendedSessionData.from_dict(record)
|
||||
assert store.revert_to_snapshot("s", snapshot_id) is False
|
||||
assert store.add_message("s", "user", "must not overwrite") is False
|
||||
assert path.read_bytes() == before
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Clearing active turns preserves the session's routing and archived recall."""
|
||||
|
||||
import pytest
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from praisonaiagents.session.sqlite_store import SqliteSessionStore
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def make_store(tmp_path):
|
||||
stores = []
|
||||
|
||||
def create(kind, **kwargs):
|
||||
store = kind(session_dir=str(tmp_path / str(len(stores))), **kwargs)
|
||||
stores.append(store)
|
||||
return store
|
||||
|
||||
yield create
|
||||
for store in stores:
|
||||
conn = getattr(store, "_conn", None)
|
||||
if conn is not None:
|
||||
conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", [DefaultSessionStore, SqliteSessionStore])
|
||||
@pytest.mark.parametrize("warm", ["cold", "search", "route"])
|
||||
def test_clear_preserves_gateway_and_agent_routes(make_store, kind, warm):
|
||||
store = make_store(kind)
|
||||
assert store.add_message("session", "user", "pelican notes")
|
||||
assert store.set_gateway_info("session", gateway_session_id="gateway", agent_id="agent")
|
||||
if warm == "search":
|
||||
assert store.search("pelican")
|
||||
elif warm == "route":
|
||||
assert store.get_by_gateway_session("gateway") is not None
|
||||
|
||||
assert store.clear_session("session")
|
||||
session = store.get_session("session")
|
||||
assert session.messages == []
|
||||
assert session.gateway_session_id == "gateway"
|
||||
assert session.agent_id == "agent"
|
||||
found = store.get_by_gateway_session("gateway")
|
||||
assert found is not None and found.session_id == "session"
|
||||
assert store.list_sessions_by_gateway_agent("agent") == ["session"]
|
||||
assert store.search("pelican") == []
|
||||
|
||||
assert store.add_message("session", "user", "narwhal notes")
|
||||
assert [hit.session_id for hit in store.search("narwhal")] == ["session"]
|
||||
assert store.delete_session("session")
|
||||
assert store.get_by_gateway_session("gateway") is None
|
||||
assert store.list_sessions_by_gateway_agent("agent") == []
|
||||
assert store.search("narwhal") == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", [DefaultSessionStore, SqliteSessionStore])
|
||||
def test_clear_keeps_archived_turns_searchable(make_store, kind):
|
||||
store = make_store(kind, active_window=3)
|
||||
assert store.add_message("session", "user", "archived pelican marker")
|
||||
for index in range(8):
|
||||
assert store.add_message("session", "user", f"later turn {index}")
|
||||
before = store.get_session("session")
|
||||
assert any("pelican" in message.content for message in before.archived_messages)
|
||||
assert store.search("pelican")
|
||||
|
||||
assert store.clear_session("session")
|
||||
after = store.get_session("session")
|
||||
assert after.messages == []
|
||||
assert after.archived_messages == before.archived_messages
|
||||
hits = store.search("pelican")
|
||||
assert [hit.session_id for hit in hits] == ["session"]
|
||||
assert any(message["archived"] for message in hits[0].messages)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("separate_store", [False, True])
|
||||
def test_delete_before_clear_index_refresh_does_not_restore_route(make_store, monkeypatch, separate_store):
|
||||
store = make_store(SqliteSessionStore)
|
||||
assert store.add_message("session", "user", "pelican notes")
|
||||
assert store.set_gateway_info("session", gateway_session_id="gateway", agent_id="agent")
|
||||
assert store.get_by_gateway_session("gateway") is not None
|
||||
deleting = store
|
||||
if separate_store:
|
||||
deleting = SqliteSessionStore(session_dir=store.session_dir, db_path=store.db_path)
|
||||
assert deleting.get_by_gateway_session("gateway") is not None
|
||||
|
||||
read_complete = threading.Event()
|
||||
resume = threading.Event()
|
||||
original = store._index_session
|
||||
|
||||
def delayed_index(session):
|
||||
read_complete.set()
|
||||
assert resume.wait(5), "delete did not release the delayed refresh"
|
||||
original(session)
|
||||
|
||||
monkeypatch.setattr(store, "_index_session", delayed_index)
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
clearing = executor.submit(store.clear_session, "session")
|
||||
try:
|
||||
assert read_complete.wait(5), "clear did not reach index refresh"
|
||||
assert deleting.delete_session("session")
|
||||
finally:
|
||||
resume.set()
|
||||
assert clearing.result(timeout=5)
|
||||
assert not store.session_exists("session")
|
||||
assert store.list_sessions_by_gateway_agent("agent") == []
|
||||
assert store.get_by_gateway_session("gateway") is None
|
||||
finally:
|
||||
if separate_store and deleting._conn is not None:
|
||||
deleting._conn.close()
|
||||
@@ -2,10 +2,12 @@
|
||||
|
||||
import pytest
|
||||
import threading
|
||||
import os
|
||||
import sqlite3
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from praisonaiagents.session.sqlite_store import SqliteSessionStore
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
from praisonaiagents.session.store import DefaultSessionStore, FileLock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -129,6 +131,13 @@ def test_write_between_import_save_and_index_keeps_newer_content(make_store, mon
|
||||
finally:
|
||||
resume.set()
|
||||
assert importing.result(timeout=5).imported == 1
|
||||
assert destination._cache["session"].messages[0].content == "narwhal newer turn"
|
||||
with monkeypatch.context() as fallback:
|
||||
def unavailable(*args):
|
||||
raise OSError("temporary read failure")
|
||||
|
||||
fallback.setattr(destination, "_load_session_from_disk", unavailable)
|
||||
assert destination._read_session_fresh("session").messages[0].content == "narwhal newer turn"
|
||||
destination.invalidate_cache()
|
||||
assert destination.get_session("session").messages[0].content == "narwhal newer turn"
|
||||
assert [hit.session_id for hit in destination.search("narwhal")] == ["session"]
|
||||
@@ -181,7 +190,593 @@ def test_post_import_read_failure_does_not_restore_old_routes(make_store, monkey
|
||||
assert saved.gateway_session_id is None
|
||||
assert saved.agent_id is None
|
||||
assert saved.messages[0].content == "narwhal replacement"
|
||||
assert destination.search("narwhal") == []
|
||||
assert [hit.session_id for hit in destination.search("narwhal")] == ["session"]
|
||||
assert destination.add_message("session", "user", "later update")
|
||||
assert [hit.session_id for hit in destination.search("narwhal")] == ["session"]
|
||||
assert destination.get_by_gateway_session("old-gateway") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("existing_index", [False, True])
|
||||
@pytest.mark.parametrize("stat_failure", [False, True])
|
||||
def test_failed_import_refresh_preserves_new_peer_index(make_store, monkeypatch, existing_index, stat_failure):
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "imported pelican")
|
||||
destination = make_store(SqliteSessionStore)
|
||||
if existing_index:
|
||||
assert destination.add_message("session", "user", "old otter")
|
||||
destination.search("pelican")
|
||||
peer = SqliteSessionStore(session_dir=destination.session_dir, db_path=destination.db_path)
|
||||
original = DefaultSessionStore._save_imported_session
|
||||
original_stat = os.stat
|
||||
|
||||
def unavailable_stat(path, *args, **kwargs):
|
||||
if os.fspath(path) == destination._get_session_path("session"):
|
||||
raise OSError("path stat temporarily unavailable")
|
||||
return original_stat(path, *args, **kwargs)
|
||||
|
||||
def save_then_peer_write(store, session, **kwargs):
|
||||
saved = original(store, session, **kwargs)
|
||||
assert peer.set_chat_history("session", [{"role": "user", "content": "newer narwhal"}])
|
||||
assert peer.set_gateway_info("session", gateway_session_id="new-route", agent_id="new-agent")
|
||||
if stat_failure:
|
||||
monkeypatch.setattr(os, "stat", unavailable_stat)
|
||||
return saved
|
||||
|
||||
def unavailable(*args):
|
||||
raise OSError("post-import refresh unavailable")
|
||||
|
||||
monkeypatch.setattr(DefaultSessionStore, "_save_imported_session", save_then_peer_write)
|
||||
monkeypatch.setattr(destination, "_load_session_from_disk", unavailable)
|
||||
try:
|
||||
assert destination.import_sessions(source.export_all(), overwrite=existing_index).imported == 1
|
||||
conn = destination._connect()
|
||||
assert conn.execute("SELECT content FROM session_fts WHERE session_id = ?", ("session",)).fetchone() == ("newer narwhal",)
|
||||
assert conn.execute("SELECT gateway_session_id, agent_id FROM session_route WHERE session_id = ?", ("session",)).fetchone() == ("new-route", "new-agent")
|
||||
finally:
|
||||
if peer._conn is not None:
|
||||
peer._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("peer_write", [False, True])
|
||||
def test_post_import_sql_failure_invalidates_all_old_index_views(make_store, monkeypatch, peer_write):
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "narwhal replacement")
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "pelican original")
|
||||
assert destination.set_gateway_info("session", gateway_session_id="old-gateway", agent_id="old-agent")
|
||||
assert destination.search("pelican")
|
||||
conn = destination._connect()
|
||||
conn.execute("CREATE TRIGGER fail_import_metadata BEFORE INSERT ON session_meta "
|
||||
"BEGIN SELECT RAISE(FAIL, 'injected index failure'); END")
|
||||
if peer_write:
|
||||
peer = DefaultSessionStore(session_dir=destination.session_dir)
|
||||
original = DefaultSessionStore._save_imported_session
|
||||
|
||||
def save_then_peer_write(store, session, **kwargs):
|
||||
saved = original(store, session, **kwargs)
|
||||
assert peer.set_chat_history("session", [{"role": "user", "content": "narwhal peer replacement"}])
|
||||
return saved
|
||||
|
||||
monkeypatch.setattr(DefaultSessionStore, "_save_imported_session", save_then_peer_write)
|
||||
report = destination.import_sessions(source.export_all(), overwrite=True)
|
||||
assert report.imported == 1
|
||||
assert report.skipped == []
|
||||
for table in ("session_fts", "session_meta", "session_route"):
|
||||
assert conn.execute(f"SELECT session_id FROM {table} WHERE session_id = ?", ("session",)).fetchall() == []
|
||||
expected = "narwhal peer replacement" if peer_write else "narwhal replacement"
|
||||
assert destination.get_session("session").messages[0].content == expected
|
||||
assert destination.get_by_gateway_session("old-gateway") is None
|
||||
assert destination.list_sessions_by_gateway_agent("old-agent") == []
|
||||
conn.execute("DROP TRIGGER fail_import_metadata")
|
||||
assert destination.add_message("session", "user", "recovered update")
|
||||
assert [hit.session_id for hit in destination.search("narwhal")] == ["session"]
|
||||
|
||||
|
||||
def test_failed_path_stat_does_not_preserve_stale_import_routes(make_store, monkeypatch):
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "new imported content")
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old indexed content")
|
||||
assert destination.set_gateway_info("session", gateway_session_id="old-route", agent_id="old-agent")
|
||||
assert destination.search("old")
|
||||
filepath = destination._get_session_path("session")
|
||||
original_stat = os.stat
|
||||
|
||||
def unavailable_stat(path, *args, **kwargs):
|
||||
if os.fspath(path) == filepath:
|
||||
raise OSError("path stat temporarily unavailable")
|
||||
return original_stat(path, *args, **kwargs)
|
||||
|
||||
def unavailable_read(*args):
|
||||
raise OSError("post-import read temporarily unavailable")
|
||||
|
||||
with monkeypatch.context() as failures:
|
||||
failures.setattr(os, "stat", unavailable_stat)
|
||||
failures.setattr(destination, "_load_session_from_disk", unavailable_read)
|
||||
report = destination.import_sessions(source.export_all(), overwrite=True)
|
||||
assert report.imported == 1
|
||||
assert report.skipped == []
|
||||
for table in ("session_fts", "session_meta", "session_route"):
|
||||
assert destination._conn.execute(f"SELECT session_id FROM {table} WHERE session_id = ?", ("session",)).fetchall() == []
|
||||
assert destination.get_session("session").messages[0].content == "new imported content"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("descriptor_identity", [False, True])
|
||||
def test_failed_refresh_keeps_peer_index_of_same_file_generation(make_store, monkeypatch, descriptor_identity):
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "imported pelican")
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old otter")
|
||||
peer = SqliteSessionStore(session_dir=destination.session_dir, db_path=destination.db_path)
|
||||
original_save = DefaultSessionStore._save_imported_session
|
||||
original_index = destination._index_session
|
||||
original_identity = destination._session_file_identity
|
||||
|
||||
def save_then_peer_write(store, session, **kwargs):
|
||||
saved = original_save(store, session, **kwargs)
|
||||
assert peer.set_chat_history("session", [{"role": "user", "content": "newer narwhal"}])
|
||||
assert peer.set_gateway_info("session", gateway_session_id="new-route", agent_id="new-agent")
|
||||
return saved
|
||||
|
||||
def failed_index_then_peer_indexes_same_generation(fresh):
|
||||
conn = destination._connect()
|
||||
conn.execute("CREATE TRIGGER fail_once BEFORE INSERT ON session_meta "
|
||||
"BEGIN SELECT RAISE(FAIL, 'injected index failure'); END")
|
||||
failed = original_index(fresh)
|
||||
conn.execute("DROP TRIGGER fail_once")
|
||||
assert failed is False
|
||||
assert peer._index_session(fresh)
|
||||
return failed
|
||||
|
||||
def identity_using_descriptor(filepath):
|
||||
def unavailable(*args, **kwargs):
|
||||
raise OSError("path stat temporarily unavailable")
|
||||
with monkeypatch.context() as path_failure:
|
||||
path_failure.setattr(os, "stat", unavailable)
|
||||
return original_identity(filepath)
|
||||
|
||||
monkeypatch.setattr(DefaultSessionStore, "_save_imported_session", save_then_peer_write)
|
||||
monkeypatch.setattr(destination, "_index_session", failed_index_then_peer_indexes_same_generation)
|
||||
if descriptor_identity:
|
||||
monkeypatch.setattr(destination, "_session_file_identity", identity_using_descriptor)
|
||||
try:
|
||||
assert destination.import_sessions(source.export_all(), overwrite=True).imported == 1
|
||||
conn = destination._connect()
|
||||
assert conn.execute("SELECT content FROM session_fts WHERE session_id = ?", ("session",)).fetchall() == [("newer narwhal",)]
|
||||
assert conn.execute("SELECT gateway_session_id, agent_id FROM session_route WHERE session_id = ?", ("session",)).fetchall() == [("new-route", "new-agent")]
|
||||
finally:
|
||||
if peer._conn is not None:
|
||||
peer._conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("peer_kind", [DefaultSessionStore, SqliteSessionStore])
|
||||
def test_failed_refresh_compares_index_to_replacement_peer_file(make_store, monkeypatch, peer_kind):
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "imported pelican")
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old otter")
|
||||
assert destination.set_gateway_info("session", gateway_session_id="old-route", agent_id="old-agent")
|
||||
kwargs = {"db_path": destination.db_path} if peer_kind is SqliteSessionStore else {}
|
||||
peer = peer_kind(session_dir=destination.session_dir, **kwargs)
|
||||
invalidate = destination._invalidate_import_index
|
||||
|
||||
def unavailable(*args):
|
||||
raise OSError("post-import read unavailable")
|
||||
|
||||
def replace_then_invalidate(session):
|
||||
# The failed refresh has released FileLock. A JSON-only peer cannot
|
||||
# refresh SQLite, whereas an indexed peer can publish a correct view.
|
||||
assert peer.set_chat_history("session", [{"role": "user", "content": "newer narwhal"}])
|
||||
assert peer.set_gateway_info("session", gateway_session_id="new-route", agent_id="new-agent")
|
||||
invalidate(session)
|
||||
|
||||
monkeypatch.setattr(destination, "_load_session_from_disk", unavailable)
|
||||
monkeypatch.setattr(destination, "_invalidate_import_index", replace_then_invalidate)
|
||||
try:
|
||||
assert destination.import_sessions(source.export_all(), overwrite=True).imported == 1
|
||||
conn = destination._connect()
|
||||
if peer_kind is DefaultSessionStore:
|
||||
for table in ("session_fts", "session_meta", "session_route"):
|
||||
assert conn.execute(f"SELECT session_id FROM {table} WHERE session_id = ?", ("session",)).fetchall() == []
|
||||
else:
|
||||
assert conn.execute("SELECT content FROM session_fts WHERE session_id = ?", ("session",)).fetchall() == [("newer narwhal",)]
|
||||
assert conn.execute("SELECT gateway_session_id, agent_id FROM session_route WHERE session_id = ?", ("session",)).fetchall() == [("new-route", "new-agent")]
|
||||
assert DefaultSessionStore(session_dir=destination.session_dir).get_chat_history("session") == [{"role": "user", "content": "newer narwhal"}]
|
||||
finally:
|
||||
conn = getattr(peer, "_conn", None)
|
||||
if conn is not None:
|
||||
conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body_failure", [False, True])
|
||||
def test_failed_index_commit_releases_transaction_and_preserves_old_views(make_store, body_failure):
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old otter")
|
||||
assert destination.set_gateway_info("session", gateway_session_id="old-route", agent_id="old-agent")
|
||||
conn = destination._connect()
|
||||
conn.execute("PRAGMA journal_mode=DELETE")
|
||||
conn.execute("PRAGMA busy_timeout=10")
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "new narwhal")
|
||||
assert source.set_gateway_info("session", gateway_session_id="new-route", agent_id="new-agent")
|
||||
replacement = source.get_session("session")
|
||||
if body_failure:
|
||||
conn.execute("CREATE TRIGGER fail_metadata BEFORE INSERT ON session_meta "
|
||||
"BEGIN SELECT RAISE(FAIL, 'injected body failure'); END")
|
||||
reader = sqlite3.connect(destination.db_path, isolation_level=None)
|
||||
try:
|
||||
reader.execute("BEGIN")
|
||||
assert reader.execute("SELECT content FROM session_fts WHERE session_id = 'session'").fetchone() == ("old otter",)
|
||||
assert destination._index_session(replacement) is False
|
||||
assert not conn.in_transaction, "failed refresh leaked its outer write transaction"
|
||||
finally:
|
||||
reader.close()
|
||||
if body_failure:
|
||||
conn.execute("DROP TRIGGER fail_metadata")
|
||||
assert conn.execute("SELECT content FROM session_fts WHERE session_id = 'session'").fetchone() == ("old otter",)
|
||||
assert conn.execute("SELECT gateway_session_id, agent_id FROM session_route WHERE session_id = 'session'").fetchone() == ("old-route", "old-agent")
|
||||
# A later refresh must commit independently and become visible to peers.
|
||||
assert destination._index_session(replacement)
|
||||
peer = sqlite3.connect(destination.db_path, isolation_level=None)
|
||||
try:
|
||||
assert peer.execute("SELECT content FROM session_fts WHERE session_id = 'session'").fetchone() == ("new narwhal",)
|
||||
assert peer.execute("SELECT gateway_session_id, agent_id FROM session_route WHERE session_id = 'session'").fetchone() == ("new-route", "new-agent")
|
||||
peer.execute("BEGIN IMMEDIATE")
|
||||
peer.execute("ROLLBACK")
|
||||
finally:
|
||||
peer.close()
|
||||
|
||||
|
||||
def test_failed_nested_index_refresh_preserves_callers_transaction(make_store):
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old otter")
|
||||
conn = destination._connect()
|
||||
conn.execute("CREATE TRIGGER fail_metadata BEFORE INSERT ON session_meta "
|
||||
"WHEN NEW.session_id = 'session' "
|
||||
"BEGIN SELECT RAISE(FAIL, 'injected body failure'); END")
|
||||
conn.execute("BEGIN")
|
||||
conn.execute("INSERT INTO session_meta VALUES ('unrelated', 'caller-pending')")
|
||||
assert destination._index_session(destination.get_session("session")) is False
|
||||
assert conn.in_transaction
|
||||
assert conn.execute("SELECT updated_at FROM session_meta WHERE session_id = 'unrelated'").fetchone() == ("caller-pending",)
|
||||
assert conn.execute("SELECT content FROM session_fts WHERE session_id = 'session'").fetchone() == ("old otter",)
|
||||
conn.execute("COMMIT")
|
||||
|
||||
|
||||
def test_import_and_initial_backfill_do_not_invert_file_database_locks(make_store, monkeypatch):
|
||||
from praisonaiagents.session import sqlite_store as indexed_module
|
||||
from praisonaiagents.session import store as store_module
|
||||
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "imported narwhal")
|
||||
destination = make_store(SqliteSessionStore, lock_timeout=0.3)
|
||||
destination._connect()
|
||||
import_holds_file = threading.Event()
|
||||
backfill_attempts_file = threading.Event()
|
||||
release_import = threading.Event()
|
||||
failed_locks = []
|
||||
original_load = destination._load_session_from_disk
|
||||
|
||||
class ObservedFileLock(FileLock):
|
||||
def acquire(self):
|
||||
if threading.current_thread().name.startswith("backfill"):
|
||||
backfill_attempts_file.set()
|
||||
acquired = super().acquire()
|
||||
if not acquired:
|
||||
failed_locks.append(self.filepath)
|
||||
return acquired
|
||||
|
||||
def load(session_id, filepath):
|
||||
if threading.current_thread().name.startswith("import"):
|
||||
import_holds_file.set()
|
||||
assert release_import.wait(5)
|
||||
return original_load(session_id, filepath)
|
||||
|
||||
monkeypatch.setattr(indexed_module, "FileLock", ObservedFileLock)
|
||||
monkeypatch.setattr(store_module, "FileLock", ObservedFileLock)
|
||||
monkeypatch.setattr(destination, "_load_session_from_disk", load)
|
||||
with ThreadPoolExecutor(max_workers=1, thread_name_prefix="import") as imports, ThreadPoolExecutor(max_workers=1, thread_name_prefix="backfill") as reads:
|
||||
restoring = imports.submit(destination.import_sessions, source.export_all())
|
||||
try:
|
||||
assert import_holds_file.wait(5)
|
||||
searching = reads.submit(destination.search, "narwhal")
|
||||
assert backfill_attempts_file.wait(5)
|
||||
finally:
|
||||
release_import.set()
|
||||
assert restoring.result(timeout=5).imported == 1
|
||||
assert [hit.session_id for hit in searching.result(timeout=5)] == ["session"]
|
||||
assert failed_locks == [], "backfill retained the DB lock while waiting for import's file lock"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["file_lock", "sql_refresh"])
|
||||
def test_initial_backfill_retries_after_a_failed_transcript(make_store, failure):
|
||||
destination = make_store(SqliteSessionStore, lock_timeout=0.05)
|
||||
plain = DefaultSessionStore(session_dir=destination.session_dir)
|
||||
assert plain.add_message("legacy", "user", "legacy narwhal")
|
||||
assert plain.set_gateway_info("legacy", gateway_session_id="legacy-route", agent_id="legacy-agent")
|
||||
conn = destination._connect()
|
||||
if failure == "file_lock":
|
||||
with FileLock(plain._get_session_path("legacy")):
|
||||
assert destination.search("narwhal") == []
|
||||
else:
|
||||
conn.execute("CREATE TRIGGER fail_metadata BEFORE INSERT ON session_meta "
|
||||
"BEGIN SELECT RAISE(FAIL, 'backfill failed'); END")
|
||||
assert destination.search("narwhal") == []
|
||||
conn.execute("DROP TRIGGER fail_metadata")
|
||||
assert [hit.session_id for hit in destination.search("narwhal")] == ["legacy"]
|
||||
assert destination.get_by_gateway_session("legacy-route").session_id == "legacy"
|
||||
assert destination.list_sessions_by_gateway_agent("legacy-agent") == ["legacy"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["file_lock", "sql_refresh"])
|
||||
def test_upgrade_backfill_does_not_rebuild_successful_sessions_on_retry(make_store, failure):
|
||||
destination = make_store(SqliteSessionStore, lock_timeout=0.02)
|
||||
plain = DefaultSessionStore(session_dir=destination.session_dir)
|
||||
for sid in ("healthy", "unavailable"):
|
||||
assert plain.add_message(sid, "user", sid + " narwhal")
|
||||
assert plain.set_gateway_info(sid, gateway_session_id=sid + "-route")
|
||||
conn = destination._connect()
|
||||
destination._set_index_content_version(conn, destination.INDEX_CONTENT_VERSION - 1)
|
||||
conn.execute("CREATE TABLE backfill_writes (session_id TEXT)")
|
||||
conn.execute("CREATE TRIGGER observe_backfill AFTER INSERT ON session_meta "
|
||||
"BEGIN INSERT INTO backfill_writes VALUES (NEW.session_id); END")
|
||||
lock = FileLock(plain._get_session_path("unavailable"))
|
||||
if failure == "file_lock":
|
||||
lock.acquire()
|
||||
else:
|
||||
conn.execute("CREATE TRIGGER fail_backfill BEFORE INSERT ON session_meta "
|
||||
"WHEN NEW.session_id = 'unavailable' "
|
||||
"BEGIN SELECT RAISE(FAIL, 'unavailable transcript'); END")
|
||||
try:
|
||||
for _ in range(3):
|
||||
assert [hit.session_id for hit in destination.search("healthy")] == ["healthy"]
|
||||
assert destination.get_by_gateway_session("healthy-route").session_id == "healthy"
|
||||
assert conn.execute("SELECT COUNT(*) FROM backfill_writes WHERE session_id = 'healthy'").fetchone() == (1,)
|
||||
assert destination._index_content_version(conn) == destination.INDEX_CONTENT_VERSION - 1
|
||||
finally:
|
||||
if failure == "file_lock":
|
||||
lock.release()
|
||||
else:
|
||||
conn.execute("DROP TRIGGER fail_backfill")
|
||||
assert destination.get_by_gateway_session("unavailable-route").session_id == "unavailable"
|
||||
assert conn.execute("SELECT COUNT(*) FROM backfill_writes WHERE session_id = 'healthy'").fetchone() == (1,)
|
||||
assert conn.execute("SELECT COUNT(*) FROM backfill_writes WHERE session_id = 'unavailable'").fetchone() == (1,)
|
||||
assert destination._index_content_version(conn) == destination.INDEX_CONTENT_VERSION
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["file_lock", "sql_refresh"])
|
||||
@pytest.mark.parametrize("peer_kind", [DefaultSessionStore, SqliteSessionStore])
|
||||
def test_upgrade_retry_reindexes_a_peer_changed_transcript(make_store, failure, peer_kind):
|
||||
destination = make_store(SqliteSessionStore, lock_timeout=0.02)
|
||||
plain = DefaultSessionStore(session_dir=destination.session_dir)
|
||||
assert plain.add_message("healthy", "user", "old otter")
|
||||
assert plain.set_gateway_info("healthy", gateway_session_id="old-route", agent_id="old-agent")
|
||||
assert plain.add_message("unavailable", "user", "blocked pelican")
|
||||
conn = destination._connect()
|
||||
destination._set_index_content_version(conn, destination.INDEX_CONTENT_VERSION - 1)
|
||||
lock = FileLock(plain._get_session_path("unavailable"))
|
||||
if failure == "file_lock":
|
||||
lock.acquire()
|
||||
else:
|
||||
conn.execute("CREATE TRIGGER fail_backfill BEFORE INSERT ON session_meta "
|
||||
"WHEN NEW.session_id = 'unavailable' "
|
||||
"BEGIN SELECT RAISE(FAIL, 'unavailable transcript'); END")
|
||||
try:
|
||||
assert [hit.session_id for hit in destination.search("otter")] == ["healthy"]
|
||||
assert destination._index_content_version(conn) == destination.INDEX_CONTENT_VERSION - 1
|
||||
kwargs = {"db_path": destination.db_path} if peer_kind is SqliteSessionStore else {}
|
||||
peer = peer_kind(session_dir=destination.session_dir, **kwargs)
|
||||
try:
|
||||
assert peer.set_chat_history("healthy", [{"role": "user", "content": "newer narwhal"}])
|
||||
assert peer.set_gateway_info("healthy", gateway_session_id="new-route", agent_id="new-agent")
|
||||
finally:
|
||||
peer_conn = getattr(peer, "_conn", None)
|
||||
if peer_conn is not None:
|
||||
peer_conn.close()
|
||||
finally:
|
||||
if failure == "file_lock":
|
||||
lock.release()
|
||||
else:
|
||||
conn.execute("DROP TRIGGER fail_backfill")
|
||||
assert [hit.session_id for hit in destination.search("narwhal")] == ["healthy"]
|
||||
assert destination.search("otter") == []
|
||||
assert destination.get_by_gateway_session("old-route") is None
|
||||
assert destination.get_by_gateway_session("new-route").session_id == "healthy"
|
||||
assert destination.list_sessions_by_gateway_agent("old-agent") == []
|
||||
assert destination.list_sessions_by_gateway_agent("new-agent") == ["healthy"]
|
||||
assert destination._index_content_version(conn) == destination.INDEX_CONTENT_VERSION
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reader", ["gateway", "agent"])
|
||||
@pytest.mark.parametrize("unreadable", [False, True])
|
||||
def test_unknown_import_generation_does_not_reuse_obsolete_routes(make_store, monkeypatch, reader, unreadable):
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "replacement")
|
||||
assert source.set_gateway_info("session", gateway_session_id="new-gateway", agent_id="new-agent")
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old")
|
||||
assert destination.set_gateway_info("session", gateway_session_id="old-gateway", agent_id="old-agent")
|
||||
assert destination.get_by_gateway_session("old-gateway") is not None
|
||||
monkeypatch.setattr(destination, "_session_file_identity", lambda *_: None)
|
||||
monkeypatch.setattr(destination, "_index_session", lambda *_: False)
|
||||
assert destination.import_sessions(source.export_all(), overwrite=True).imported == 1
|
||||
assert destination._connect().execute("SELECT gateway_session_id FROM session_route WHERE session_id = ?", ("session",)).fetchone() == ("old-gateway",)
|
||||
if unreadable:
|
||||
import builtins
|
||||
original_open = builtins.open
|
||||
path = destination._get_session_path("session")
|
||||
|
||||
def fail_transcript_open(filename, *args, **kwargs):
|
||||
if os.fspath(filename) == path and (not args or "r" in args[0]):
|
||||
raise PermissionError("temporarily unreadable transcript")
|
||||
return original_open(filename, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "open", fail_transcript_open)
|
||||
if reader == "gateway":
|
||||
assert destination.get_by_gateway_session("old-gateway") is None
|
||||
else:
|
||||
assert destination.list_sessions_by_gateway_agent("old-agent") == []
|
||||
|
||||
|
||||
def test_gateway_agent_limit_counts_verified_routes(make_store):
|
||||
destination = make_store(SqliteSessionStore)
|
||||
for sid in ["stale", "valid"]:
|
||||
assert destination.add_message(sid, "user", sid)
|
||||
assert destination.set_gateway_info(sid, agent_id="agent")
|
||||
assert len(destination.list_sessions_by_gateway_agent("agent")) == 2
|
||||
peer = DefaultSessionStore(session_dir=destination.session_dir)
|
||||
assert peer.set_gateway_info("stale", agent_id="other-agent")
|
||||
assert destination.list_sessions_by_gateway_agent("agent", limit=1) == ["valid"]
|
||||
|
||||
|
||||
def test_gateway_route_read_failure_does_not_use_warm_cache(make_store, monkeypatch):
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old")
|
||||
assert destination.set_gateway_info("session", gateway_session_id="old-gateway", agent_id="old-agent")
|
||||
assert destination.get_by_gateway_session("old-gateway") is not None
|
||||
peer = DefaultSessionStore(session_dir=destination.session_dir)
|
||||
assert peer.set_gateway_info("session", gateway_session_id="new-gateway", agent_id="new-agent")
|
||||
|
||||
def fail_read(*_):
|
||||
raise PermissionError("transient route read failure")
|
||||
|
||||
with monkeypatch.context() as blocked:
|
||||
blocked.setattr(destination, "_load_session_from_disk", fail_read)
|
||||
assert destination.get_by_gateway_session("old-gateway") is None
|
||||
assert destination.list_sessions_by_gateway_agent("old-agent") == []
|
||||
assert destination.get_by_gateway_session("old-gateway") is None
|
||||
assert destination.list_sessions_by_gateway_agent("old-agent") == []
|
||||
assert peer.get_by_gateway_session("new-gateway").session_id == "session"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("restart", [False, True])
|
||||
@pytest.mark.parametrize("failure", ["read", "index", "unknown"])
|
||||
def test_failed_import_recovers_new_recall_and_routes(make_store, monkeypatch, restart, failure):
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "new narwhal")
|
||||
assert source.set_gateway_info("session", gateway_session_id="new-route", agent_id="new-agent")
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old otter")
|
||||
assert destination.set_gateway_info("session", gateway_session_id="old-route", agent_id="old-agent")
|
||||
assert destination.get_by_gateway_session("old-route") is not None
|
||||
with monkeypatch.context() as blocked:
|
||||
if failure == "read":
|
||||
def fail_read(*_):
|
||||
raise PermissionError("temporarily unreadable")
|
||||
blocked.setattr(destination, "_load_session_from_disk", fail_read)
|
||||
else:
|
||||
blocked.setattr(destination, "_index_session", lambda *_: False)
|
||||
if failure == "unknown":
|
||||
blocked.setattr(destination, "_session_file_identity", lambda *_: None)
|
||||
assert destination.import_sessions(source.export_all(), overwrite=True, reset_live_fields=False).imported == 1
|
||||
assert destination.get_by_gateway_session("old-route") is None
|
||||
if restart:
|
||||
reopened = SqliteSessionStore(session_dir=destination.session_dir, db_path=destination.db_path)
|
||||
try:
|
||||
assert [hit.session_id for hit in reopened.search("narwhal")] == ["session"]
|
||||
assert reopened.get_by_gateway_session("new-route").session_id == "session"
|
||||
assert reopened.list_sessions_by_gateway_agent("new-agent") == ["session"]
|
||||
assert reopened.get_by_gateway_session("old-route") is None
|
||||
assert reopened.search("otter") == []
|
||||
finally:
|
||||
reopened._conn.close()
|
||||
return
|
||||
assert [hit.session_id for hit in destination.search("narwhal")] == ["session"]
|
||||
assert destination.search("otter") == []
|
||||
assert destination.get_by_gateway_session("new-route").session_id == "session"
|
||||
assert destination.list_sessions_by_gateway_agent("new-agent") == ["session"]
|
||||
assert destination.get_by_gateway_session("old-route") is None
|
||||
assert destination._connect().execute("SELECT content FROM session_fts WHERE session_id = ?", ("session",)).fetchall() == [("new narwhal",)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("peer_kind", [DefaultSessionStore, SqliteSessionStore])
|
||||
def test_import_retry_indexes_current_peer_and_does_not_rewrite_healthy_rows(make_store, peer_kind):
|
||||
destination = make_store(SqliteSessionStore)
|
||||
for sid in ("session", "healthy"):
|
||||
assert destination.add_message(sid, "user", "old otter")
|
||||
assert destination.set_gateway_info(sid, gateway_session_id=sid + "-route")
|
||||
assert destination.search("otter")
|
||||
conn = destination._connect()
|
||||
conn.execute("CREATE TABLE observed_writes (session_id TEXT)")
|
||||
conn.execute("CREATE TRIGGER observe_retry AFTER INSERT ON session_meta "
|
||||
"BEGIN INSERT INTO observed_writes VALUES (NEW.session_id); END")
|
||||
conn.execute("CREATE TRIGGER fail_retry BEFORE INSERT ON session_meta "
|
||||
"WHEN NEW.session_id = 'session' BEGIN SELECT RAISE(FAIL, 'retry unavailable'); END")
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "imported narwhal")
|
||||
assert destination.import_sessions(source.export_all(), overwrite=True).imported == 1
|
||||
for _ in range(2):
|
||||
assert destination.search("narwhal") == []
|
||||
assert destination.get_by_gateway_session("healthy-route").session_id == "healthy"
|
||||
assert conn.execute("SELECT COUNT(*) FROM observed_writes").fetchone() == (0,)
|
||||
conn.execute("DROP TRIGGER fail_retry")
|
||||
kwargs = {"db_path": destination.db_path} if peer_kind is SqliteSessionStore else {}
|
||||
peer = peer_kind(session_dir=destination.session_dir, **kwargs)
|
||||
try:
|
||||
assert peer.set_chat_history("session", [{"role": "user", "content": "newer penguin"}])
|
||||
assert peer.set_gateway_info("session", gateway_session_id="peer-route", agent_id="peer-agent")
|
||||
finally:
|
||||
if getattr(peer, "_conn", None) is not None:
|
||||
peer._conn.close()
|
||||
assert destination.get_by_gateway_session("peer-route").session_id == "session"
|
||||
assert [hit.session_id for hit in destination.search("penguin")] == ["session"]
|
||||
assert destination.search("narwhal") == []
|
||||
assert destination.get_by_gateway_session("session-route") is None
|
||||
assert conn.execute("SELECT key FROM session_index_meta WHERE key = 'import_pending:session'").fetchone() is None
|
||||
assert conn.execute("SELECT COUNT(*) FROM observed_writes WHERE session_id = 'healthy'").fetchone() == (0,)
|
||||
|
||||
|
||||
def test_failed_import_cleanup_preserves_caller_transaction(make_store):
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "new narwhal")
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old otter")
|
||||
assert destination.set_gateway_info("session", gateway_session_id="old-route")
|
||||
conn = destination._connect()
|
||||
conn.execute("CREATE TRIGGER fail_import BEFORE INSERT ON session_meta "
|
||||
"WHEN NEW.session_id = 'session' BEGIN SELECT RAISE(FAIL, 'refresh unavailable'); END")
|
||||
conn.execute("BEGIN")
|
||||
conn.execute("INSERT INTO session_meta VALUES ('unrelated', 'caller-pending')")
|
||||
assert destination.import_sessions(source.export_all(), overwrite=True).imported == 1
|
||||
assert conn.in_transaction
|
||||
assert conn.execute("SELECT updated_at FROM session_meta WHERE session_id = 'unrelated'").fetchone() == ("caller-pending",)
|
||||
conn.execute("COMMIT")
|
||||
for table in ("session_fts", "session_meta", "session_route"):
|
||||
assert conn.execute(f"SELECT session_id FROM {table} WHERE session_id = 'session'").fetchall() == []
|
||||
conn.execute("DROP TRIGGER fail_import")
|
||||
assert [hit.session_id for hit in destination.search("narwhal")] == ["session"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("missing", [False, True])
|
||||
@pytest.mark.parametrize("restart", [False, True])
|
||||
def test_import_retry_keeps_marker_when_transcript_is_unavailable(make_store, monkeypatch, missing, restart):
|
||||
source = make_store(DefaultSessionStore)
|
||||
assert source.add_message("session", "user", "new narwhal")
|
||||
assert source.set_gateway_info("session", gateway_session_id="new-route", agent_id="new-agent")
|
||||
destination = make_store(SqliteSessionStore)
|
||||
assert destination.add_message("session", "user", "old otter")
|
||||
assert destination.set_gateway_info("session", gateway_session_id="old-route")
|
||||
assert destination.search("otter")
|
||||
with monkeypatch.context() as blocked:
|
||||
blocked.setattr(destination, "_index_session", lambda *_: False)
|
||||
assert destination.import_sessions(source.export_all(), overwrite=True, reset_live_fields=False).imported == 1
|
||||
from pathlib import Path
|
||||
|
||||
path = Path(destination._get_session_path("session"))
|
||||
payload = path.read_bytes()
|
||||
if missing:
|
||||
path.unlink()
|
||||
else:
|
||||
path.write_bytes(b"{invalid json")
|
||||
if restart:
|
||||
destination._conn.close()
|
||||
destination._conn = None
|
||||
destination._db_ready = False
|
||||
destination._backfilled = False
|
||||
destination.invalidate_cache()
|
||||
assert destination.search("narwhal") == []
|
||||
assert destination._connect().execute("SELECT key FROM session_index_meta WHERE key = 'import_pending:session'").fetchone() is not None
|
||||
path.write_bytes(payload)
|
||||
assert [hit.session_id for hit in destination.search("narwhal")] == ["session"]
|
||||
assert destination.get_by_gateway_session("new-route").session_id == "session"
|
||||
assert destination.list_sessions_by_gateway_agent("new-agent") == ["session"]
|
||||
@@ -42,17 +42,18 @@ class TestFileBackend:
|
||||
|
||||
backend = FileBackend(storage_dir=str(tmp_path))
|
||||
|
||||
for raw in (b"\xff", b"\xe4\xb8", b"\xed\xa0\x80"):
|
||||
backend.save("bad", {"message": "original"})
|
||||
file_path = Path(tmp_path) / "bad.json"
|
||||
for idx, raw in enumerate((b"\xff", b"\xe4\xb8", b"\xed\xa0\x80")):
|
||||
key = f"bad_{idx}"
|
||||
backend.save(key, {"message": "original"})
|
||||
file_path = Path(tmp_path) / f"{key}.json"
|
||||
file_path.write_bytes(raw)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="praisonaiagents.storage.backends"):
|
||||
caplog.clear()
|
||||
assert backend.load("bad") is None
|
||||
assert backend.load(key) is None
|
||||
# A warning is emitted identifying the failed key
|
||||
assert any(
|
||||
rec.levelno == logging.WARNING and "Failed to load bad" in rec.getMessage()
|
||||
rec.levelno == logging.WARNING and f"Failed to load {key}" in rec.getMessage()
|
||||
for rec in caplog.records
|
||||
)
|
||||
# Original bytes are preserved (load is non-destructive)
|
||||
|
||||
@@ -287,6 +287,71 @@ class TestBaseJSONStore:
|
||||
assert store._data == {"items": [], "version": 1}
|
||||
|
||||
|
||||
class TestInvalidUtf8Recovery:
|
||||
"""Regression tests for invalid-UTF-8 recovery (Issue #5577)."""
|
||||
|
||||
def test_load_recovers_from_invalid_utf8_unlocked(self):
|
||||
"""Invalid UTF-8 bytes recover to defaults without a rewrite."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
store_path = Path(tmpdir) / "test.json"
|
||||
store_path.write_bytes(b"\xff\xfe\xfa")
|
||||
|
||||
store = BaseJSONStore(store_path, use_file_lock=False)
|
||||
assert store.load() == {}
|
||||
# Original bytes preserved on disk (no corruption rewrite on load).
|
||||
assert store_path.read_bytes() == b"\xff\xfe\xfa"
|
||||
|
||||
def test_load_recovers_from_invalid_utf8_locked(self):
|
||||
"""Invalid UTF-8 recovers with file locking enabled too."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
store_path = Path(tmpdir) / "test.json"
|
||||
store_path.write_bytes(b"\xff")
|
||||
|
||||
store = BaseJSONStore(store_path, use_file_lock=True)
|
||||
assert store.load() == {}
|
||||
assert store_path.read_bytes() == b"\xff"
|
||||
|
||||
def test_file_backend_recovers_from_invalid_utf8(self):
|
||||
"""FileBackend.load() returns None (not raise) on invalid UTF-8."""
|
||||
from praisonaiagents.storage.backends import FileBackend
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
backend = FileBackend(storage_dir=Path(tmpdir))
|
||||
key = "session"
|
||||
(Path(tmpdir) / f"{key}.json").write_bytes(b"\xff")
|
||||
|
||||
assert backend.load(key) is None
|
||||
|
||||
def test_store_with_file_backend_recovers_from_invalid_utf8(self):
|
||||
"""BaseJSONStore over a FileBackend recovers to defaults on bad bytes."""
|
||||
from praisonaiagents.storage.backends import FileBackend
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
store_path = Path(tmpdir) / "session.json"
|
||||
backend = FileBackend(storage_dir=Path(tmpdir))
|
||||
(Path(tmpdir) / "session.json").write_bytes(b"\xff")
|
||||
|
||||
store = BaseJSONStore(store_path, backend=backend)
|
||||
assert store.load() == {}
|
||||
|
||||
def test_async_load_recovers_from_invalid_utf8(self):
|
||||
"""AsyncBaseJSONStore.load_async() recovers from invalid UTF-8."""
|
||||
import asyncio
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("aiofiles")
|
||||
from praisonaiagents.storage.base import AsyncBaseJSONStore
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
store_path = Path(tmpdir) / "test.json"
|
||||
store_path.write_bytes(b"\xff")
|
||||
|
||||
store = AsyncBaseJSONStore(store_path)
|
||||
result = asyncio.run(store.load_async())
|
||||
assert result == {}
|
||||
assert store_path.read_bytes() == b"\xff"
|
||||
|
||||
|
||||
class TestListJsonSessions:
|
||||
"""Tests for list_json_sessions utility."""
|
||||
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
"""Base stores keep malformed-input fallback consistent for undecodable bytes."""
|
||||
|
||||
import asyncio
|
||||
import builtins
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.storage.base import AsyncBaseJSONStore, BaseJSONStore
|
||||
|
||||
|
||||
class DefaultStore(BaseJSONStore):
|
||||
def _default_data(self):
|
||||
return {"items": []}
|
||||
|
||||
|
||||
class AsyncDefaultStore(AsyncBaseJSONStore):
|
||||
def _default_data(self):
|
||||
return {"items": []}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ensure_aiofiles(monkeypatch):
|
||||
"""Guarantee the async recovery branch runs even when aiofiles is absent.
|
||||
|
||||
aiofiles is not a declared dependency, so CI may lack it and silently skip
|
||||
every async case. Rather than skip (hiding the regression), inject a minimal
|
||||
real-decode stub that mirrors aiofiles.open's text-read behaviour so the
|
||||
UnicodeDecodeError/JSONDecodeError recovery is always exercised.
|
||||
"""
|
||||
try:
|
||||
importlib.import_module("aiofiles")
|
||||
return
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
fake = types.ModuleType("aiofiles")
|
||||
|
||||
class _AsyncFile:
|
||||
def __init__(self, path, mode, encoding):
|
||||
self._path = path
|
||||
self._mode = mode
|
||||
self._encoding = encoding
|
||||
self._fh = None
|
||||
|
||||
async def __aenter__(self):
|
||||
self._fh = open(self._path, self._mode, encoding=self._encoding)
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
if self._fh is not None:
|
||||
self._fh.close()
|
||||
return False
|
||||
|
||||
async def read(self):
|
||||
return self._fh.read()
|
||||
|
||||
async def write(self, data):
|
||||
return self._fh.write(data)
|
||||
|
||||
def _open(path, mode="r", encoding=None):
|
||||
return _AsyncFile(path, mode, encoding)
|
||||
|
||||
fake.open = _open
|
||||
monkeypatch.setitem(sys.modules, "aiofiles", fake)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flavor", ["locked", "unlocked", "async"])
|
||||
@pytest.mark.parametrize("payload", [b"\xff", b'{"items":["\xe2\x82', b'{"items":["\xed\xa0\x80"]}', b"{invalid"])
|
||||
def test_unreadable_bytes_use_defaults_without_rewriting_file(tmp_path, payload, flavor, ensure_aiofiles):
|
||||
path = tmp_path / "store.json"
|
||||
path.write_bytes(payload)
|
||||
if flavor == "async":
|
||||
store = AsyncDefaultStore(path)
|
||||
|
||||
async def read_twice():
|
||||
assert await store.load_async() == {"items": []}
|
||||
assert await store.load_async() == {"items": []}
|
||||
|
||||
asyncio.run(read_twice())
|
||||
else:
|
||||
store = DefaultStore(path, use_file_lock=flavor == "locked")
|
||||
assert store.load() == {"items": []}
|
||||
assert path.read_bytes() == payload
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flavor", ["locked", "unlocked", "async"])
|
||||
def test_valid_unicode_and_missing_file_keep_existing_behavior(tmp_path, flavor, ensure_aiofiles):
|
||||
path = tmp_path / "store.json"
|
||||
expected = {"items": ["你好 café 😀"]}
|
||||
if flavor == "async":
|
||||
store = AsyncDefaultStore(path)
|
||||
|
||||
async def round_trip():
|
||||
assert await store.load_async() == {"items": []}
|
||||
await store.save_async(expected)
|
||||
assert await store.load_async() == expected
|
||||
|
||||
asyncio.run(round_trip())
|
||||
else:
|
||||
store = DefaultStore(path, use_file_lock=flavor == "locked")
|
||||
assert store.load() == {"items": []}
|
||||
store.save(expected)
|
||||
assert store.load() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_file_lock", [True, False])
|
||||
def test_io_error_fallback_cannot_overwrite_source_until_successful_reload(tmp_path, monkeypatch, use_file_lock):
|
||||
path = tmp_path / "store.json"
|
||||
payload = b'{"items":["original"]}'
|
||||
path.write_bytes(payload)
|
||||
real_open = builtins.open
|
||||
|
||||
def fail_source_read(file, mode="r", *args, **kwargs):
|
||||
if file == path and mode == "r":
|
||||
raise PermissionError("read temporarily denied")
|
||||
return real_open(file, mode, *args, **kwargs)
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
patch.setattr(builtins, "open", fail_source_read)
|
||||
store = DefaultStore(path, use_file_lock=use_file_lock)
|
||||
with pytest.raises(OSError, match="unreadable"):
|
||||
store.save({"items": ["new"]})
|
||||
assert path.read_bytes() == payload
|
||||
assert store.load() == {"items": ["original"]}
|
||||
store.save({"items": ["new"]})
|
||||
assert store.load() == {"items": ["new"]}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flavor", ["locked", "unlocked", "async"])
|
||||
@pytest.mark.parametrize("payload", [b"\xff", b"{invalid"])
|
||||
def test_save_after_default_fallback_preserves_unreadable_source(tmp_path, flavor, payload, ensure_aiofiles):
|
||||
path = tmp_path / "store.json"
|
||||
path.write_bytes(payload)
|
||||
if flavor == "async":
|
||||
store = AsyncDefaultStore(path)
|
||||
|
||||
async def attempt_save():
|
||||
assert await store.load_async() == {"items": []}
|
||||
with pytest.raises(OSError, match="unreadable"):
|
||||
await store.save_async({"items": ["new"]})
|
||||
|
||||
asyncio.run(attempt_save())
|
||||
else:
|
||||
store = DefaultStore(path, use_file_lock=flavor == "locked")
|
||||
with pytest.raises(OSError, match="unreadable"):
|
||||
store.save({"items": ["new"]})
|
||||
assert path.read_bytes() == payload
|
||||
assert not list(tmp_path.glob("*.tmp"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flavor", ["locked", "unlocked", "async"])
|
||||
@pytest.mark.parametrize("recovery", ["reload", "delete"])
|
||||
def test_unreadable_source_can_be_repaired_or_explicitly_deleted(tmp_path, flavor, recovery, ensure_aiofiles):
|
||||
path = tmp_path / "store.json"
|
||||
path.write_bytes(b"\xff")
|
||||
expected = {"items": ["updated"]}
|
||||
if flavor == "async":
|
||||
store = AsyncDefaultStore(path)
|
||||
|
||||
async def recover():
|
||||
await store.load_async()
|
||||
if recovery == "reload":
|
||||
path.write_text('{"items":["repaired"]}', encoding="utf-8")
|
||||
assert await store.load_async() == {"items": ["repaired"]}
|
||||
else:
|
||||
assert await store.delete_async()
|
||||
await store.save_async(expected)
|
||||
assert await store.load_async() == expected
|
||||
|
||||
asyncio.run(recover())
|
||||
else:
|
||||
store = DefaultStore(path, use_file_lock=flavor == "locked")
|
||||
if recovery == "reload":
|
||||
path.write_text('{"items":["repaired"]}', encoding="utf-8")
|
||||
assert store.load() == {"items": ["repaired"]}
|
||||
else:
|
||||
assert store.delete()
|
||||
store.save(expected)
|
||||
assert store.load() == expected
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Unreadable UTF-8 follows FileBackend's malformed-JSON load boundary."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.storage.backends import FileBackend
|
||||
from praisonaiagents.storage.base import BaseJSONStore
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [b"\xff", b'{"message":"\xe2\x82', b'{"message":"\xed\xa0\x80"}', b"{invalid"])
|
||||
def test_load_reports_unreadable_file_without_modifying_it(tmp_path, caplog, payload):
|
||||
backend = FileBackend(storage_dir=str(tmp_path))
|
||||
backend.save("bad", {"message": "original"})
|
||||
backend.save("valid", {"message": "neighbor"})
|
||||
path = tmp_path / "bad.json"
|
||||
path.write_bytes(payload)
|
||||
|
||||
assert backend.load("bad") is None
|
||||
assert "Failed to load bad" in caplog.text
|
||||
assert path.read_bytes() == payload
|
||||
assert backend.exists("bad")
|
||||
assert backend.load("valid") == {"message": "neighbor"}
|
||||
|
||||
|
||||
def test_load_preserves_valid_unicode_and_missing_file_behavior(tmp_path, caplog):
|
||||
backend = FileBackend(storage_dir=str(tmp_path))
|
||||
data = {"message": "你好 café 😀"}
|
||||
backend.save("valid", data)
|
||||
assert backend.load("valid") == data
|
||||
assert backend.load("missing") is None
|
||||
assert not caplog.records
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [b"\xff", b"{invalid"])
|
||||
def test_default_store_save_cannot_overwrite_unreadable_backend_record(tmp_path, payload):
|
||||
backend = FileBackend(storage_dir=str(tmp_path))
|
||||
path = tmp_path / "session.json"
|
||||
path.write_bytes(payload)
|
||||
store = BaseJSONStore("session.json", backend=backend)
|
||||
assert store.load() == {}
|
||||
with pytest.raises(OSError, match="unreadable record"):
|
||||
store.save({"iterations": [{"result": "new"}]})
|
||||
assert path.read_bytes() == payload
|
||||
assert not list(tmp_path.glob("*.tmp"))
|
||||
backend.save("neighbor", {"message": "valid"})
|
||||
assert backend.load("neighbor") == {"message": "valid"}
|
||||
|
||||
|
||||
def test_write_protection_uses_file_identity_and_clears_after_successful_read(tmp_path):
|
||||
backend = FileBackend(storage_dir=str(tmp_path))
|
||||
path = tmp_path / "alias_key.json"
|
||||
path.write_bytes(b"\xff")
|
||||
assert backend.load("alias/key") is None
|
||||
with pytest.raises(OSError, match="unreadable record"):
|
||||
backend.save("alias_key", {"message": "new"})
|
||||
path.write_text('{"message":"repaired"}', encoding="utf-8")
|
||||
assert backend.load("alias_key") == {"message": "repaired"}
|
||||
backend.save("alias/key", {"message": "updated"})
|
||||
assert backend.load("alias_key") == {"message": "updated"}
|
||||
|
||||
|
||||
def test_deleted_unreadable_record_can_be_created_again(tmp_path):
|
||||
backend = FileBackend(storage_dir=str(tmp_path))
|
||||
path = tmp_path / "bad.json"
|
||||
path.write_bytes(b"\xff")
|
||||
assert backend.load("bad") is None
|
||||
assert backend.delete("bad")
|
||||
backend.save("bad", {"message": "new"})
|
||||
assert backend.load("bad") == {"message": "new"}
|
||||
@@ -0,0 +1,62 @@
|
||||
"""Explicit reset recovers corrupt stores without weakening ordinary writes."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.memory.learn.stores import PersonaStore
|
||||
from praisonaiagents.storage import base
|
||||
from praisonaiagents.storage.base import BaseJSONStore
|
||||
|
||||
|
||||
@pytest.fixture(params=["base", "learn", "training"])
|
||||
def corrupt_store(request, tmp_path, monkeypatch):
|
||||
path = tmp_path / "session.json"
|
||||
path.write_bytes(b"\xff")
|
||||
if request.param == "base":
|
||||
store = BaseJSONStore(path)
|
||||
underlying = store
|
||||
expected = {}
|
||||
elif request.param == "learn":
|
||||
store = PersonaStore(store_path=str(path))
|
||||
underlying = store._store
|
||||
expected = {}
|
||||
else:
|
||||
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[4] / "praisonai-train"))
|
||||
from praisonai_train.train.agents.storage import TrainingStorage
|
||||
|
||||
store = TrainingStorage("session", storage_dir=tmp_path)
|
||||
underlying = store._store
|
||||
expected = {"session_id": "session", "scenarios": [], "iterations": [], "report": None}
|
||||
return store, underlying, path, expected
|
||||
|
||||
|
||||
def test_explicit_clear_recovers_corrupt_store(corrupt_store):
|
||||
store, underlying, path, expected = corrupt_store
|
||||
store.clear()
|
||||
saved = json.loads(path.read_text(encoding="utf-8"))
|
||||
assert all(saved[key] == value for key, value in expected.items())
|
||||
if not expected:
|
||||
assert saved == {}
|
||||
underlying.save({"recovered": True})
|
||||
assert json.loads(path.read_text(encoding="utf-8")) == {"recovered": True}
|
||||
|
||||
|
||||
def test_failed_explicit_clear_keeps_corrupt_source_guard(corrupt_store, monkeypatch):
|
||||
store, underlying, path, _ = corrupt_store
|
||||
replace = base.os.replace
|
||||
|
||||
def fail_replace(*args):
|
||||
raise OSError("replacement failed")
|
||||
|
||||
monkeypatch.setattr(base.os, "replace", fail_replace)
|
||||
with pytest.raises(OSError, match="replacement failed"):
|
||||
store.clear()
|
||||
assert path.read_bytes() == b"\xff"
|
||||
with pytest.raises(OSError, match="unreadable"):
|
||||
underlying.save({"recovered": True})
|
||||
assert path.read_bytes() == b"\xff"
|
||||
monkeypatch.setattr(base.os, "replace", replace)
|
||||
store.clear()
|
||||
assert not underlying._read_failed
|
||||
@@ -0,0 +1,49 @@
|
||||
"""clear reports the rows its DELETE removes, including committed peer changes."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.storage.backends import SQLiteBackend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ["insert", "delete", "none", "empty"])
|
||||
def test_clear_counts_its_delete_with_a_second_connection(tmp_path, monkeypatch, change):
|
||||
path = str(tmp_path / "store.db")
|
||||
backend = SQLiteBackend(db_path=path)
|
||||
peer = SQLiteBackend(db_path=path)
|
||||
deleted_keys = []
|
||||
try:
|
||||
if change != "empty":
|
||||
backend.save("first", {"value": 1})
|
||||
backend.save("second", {"value": 2})
|
||||
conn = backend._get_conn()
|
||||
|
||||
class Cursor:
|
||||
def __init__(self):
|
||||
self.cursor = conn.cursor()
|
||||
|
||||
def execute(self, sql, *args):
|
||||
if sql.startswith("DELETE FROM"):
|
||||
if change == "insert":
|
||||
peer.save("late", {"value": 3})
|
||||
elif change == "delete":
|
||||
assert peer.delete("first")
|
||||
deleted_keys.extend(peer.list_keys())
|
||||
return self.cursor.execute(sql, *args)
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self.cursor, name)
|
||||
|
||||
class Connection:
|
||||
def cursor(self):
|
||||
return Cursor()
|
||||
|
||||
def commit(self):
|
||||
conn.commit()
|
||||
|
||||
monkeypatch.setattr(backend, "_get_conn", lambda: Connection())
|
||||
assert backend.clear() == len(deleted_keys)
|
||||
assert len(deleted_keys) == {"insert": 3, "delete": 1, "none": 2, "empty": 0}[change]
|
||||
assert peer.list_keys() == []
|
||||
finally:
|
||||
backend.close()
|
||||
peer.close()
|
||||
@@ -0,0 +1,18 @@
|
||||
"""SQLite key enumeration uses literal, case-sensitive prefixes."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.storage.backends import SQLiteBackend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["task_", "task%", "Task", "", "ordinary", "键", "slash\\", "nul\x00"])
|
||||
def test_list_keys_matches_literal_prefix(tmp_path, prefix):
|
||||
keys = ["task_one", "taskXone", "task%one", "TaskOne", "ordinary", "键一", "键二", "slash\\one", "slashXone", "nul\x00one", "nulXone"]
|
||||
store = SQLiteBackend(db_path=str(tmp_path / "storage.db"))
|
||||
try:
|
||||
for key in keys:
|
||||
store.save(key, {"key": key})
|
||||
assert store.list_keys(prefix) == sorted(key for key in keys if key.startswith(prefix))
|
||||
assert store.list_keys("missing") == []
|
||||
finally:
|
||||
store.close()
|
||||
@@ -7,7 +7,7 @@ and ContextCompactor is called when enabled.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class TestExecutionConfigCompactionFields:
|
||||
@@ -80,15 +80,34 @@ class TestCompactionMaxTokensResolution:
|
||||
# Must be far above the old flat 8000 for a large-context model.
|
||||
assert resolved > 8000
|
||||
|
||||
@staticmethod
|
||||
def _content_for_tokens(compactor, target_tokens):
|
||||
"""Build text whose size is measured by the compactor's own tokeniser.
|
||||
|
||||
A fixed chars-per-token ratio is unsafe: a degenerate ``"A" * n`` string
|
||||
collapses to far fewer tokens under tiktoken (available when its vocab is
|
||||
cached, e.g. on CI) than the 4-chars/token heuristic assumes, so sizing by
|
||||
a hard-coded ratio makes the assertion environment-dependent. Measuring
|
||||
with ``compactor.estimate_tokens`` keeps the test deterministic under both
|
||||
the accurate and heuristic token-counting paths.
|
||||
"""
|
||||
word = "lorem ipsum dolor sit amet "
|
||||
text = word * max(1, target_tokens // max(1, compactor.estimate_tokens(word)))
|
||||
while compactor.estimate_tokens(text) < target_tokens:
|
||||
text += word * 256
|
||||
return text
|
||||
|
||||
def test_unset_does_not_trigger_within_window(self):
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig
|
||||
from praisonaiagents.compaction import ContextCompactor
|
||||
agent = self._make_agent("gpt-4o") # 128k window
|
||||
cfg = ExecutionConfig(context_compaction=True)
|
||||
resolved = agent._resolve_compaction_max_tokens(cfg)
|
||||
# ~30k tokens of conversation should fit comfortably in a 128k model.
|
||||
msgs = [{"role": "user", "content": "A" * 4 * 30000}]
|
||||
compactor = ContextCompactor(max_tokens=resolved)
|
||||
# A conversation well under the resolved budget must not trigger.
|
||||
content = self._content_for_tokens(compactor, resolved // 2)
|
||||
assert compactor.estimate_tokens(content) < resolved
|
||||
msgs = [{"role": "user", "content": content}]
|
||||
assert compactor.needs_compaction(msgs) is False
|
||||
|
||||
def test_unset_triggers_when_budget_approached(self):
|
||||
@@ -97,11 +116,11 @@ class TestCompactionMaxTokensResolution:
|
||||
agent = self._make_agent("gpt-4o")
|
||||
cfg = ExecutionConfig(context_compaction=True)
|
||||
resolved = agent._resolve_compaction_max_tokens(cfg)
|
||||
# Well beyond the working budget must trigger compaction. Repeated
|
||||
# characters compress heavily under tiktoken (~8 chars/token), so use
|
||||
# 8× not the heuristic 4× when sizing synthetic overflow content.
|
||||
msgs = [{"role": "user", "content": "A" * 8 * (resolved + 50000)}]
|
||||
compactor = ContextCompactor(max_tokens=resolved)
|
||||
# Well beyond the working budget must trigger compaction.
|
||||
content = self._content_for_tokens(compactor, resolved + 50000)
|
||||
assert compactor.estimate_tokens(content) > resolved
|
||||
msgs = [{"role": "user", "content": content}]
|
||||
assert compactor.needs_compaction(msgs) is True
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
"""Regression tests for the shared lazy run-lock helper (``_run_lock``).
|
||||
|
||||
AgentTeam and Workflow/AgentFlow both delegate their ``_execution_lock`` to
|
||||
``ensure_run_lock``; these tests pin the invariants that delegation relies on:
|
||||
stable per-instance identity, support for both the plain-attribute and
|
||||
preset-``None`` dataclass-field cases, and a single shared lock object under a
|
||||
concurrent first-access race (double-checked locking).
|
||||
"""
|
||||
import threading
|
||||
|
||||
from praisonaiagents._run_lock import ensure_run_lock
|
||||
|
||||
|
||||
class _PlainAttr:
|
||||
"""Mimics AgentTeam: ``_run_lock`` is a plain (missing) attribute."""
|
||||
|
||||
|
||||
class _PresetNoneField:
|
||||
"""Mimics the Workflow dataclass field defaulting to ``None``."""
|
||||
|
||||
def __init__(self):
|
||||
self._run_lock = None
|
||||
|
||||
|
||||
def test_mints_and_returns_stable_lock_plain_attribute():
|
||||
obj = _PlainAttr()
|
||||
lock = ensure_run_lock(obj)
|
||||
assert lock is not None
|
||||
assert ensure_run_lock(obj) is lock
|
||||
assert obj._run_lock is lock
|
||||
|
||||
|
||||
def test_mints_and_returns_stable_lock_preset_none_field():
|
||||
obj = _PresetNoneField()
|
||||
lock = ensure_run_lock(obj)
|
||||
assert lock is not None
|
||||
assert ensure_run_lock(obj) is lock
|
||||
assert obj._run_lock is lock
|
||||
|
||||
|
||||
def test_distinct_instances_get_distinct_locks():
|
||||
assert ensure_run_lock(_PlainAttr()) is not ensure_run_lock(_PlainAttr())
|
||||
|
||||
|
||||
def test_concurrent_first_access_observes_single_lock():
|
||||
"""Threads racing into the first access must all observe the same lock."""
|
||||
obj = _PlainAttr()
|
||||
start = threading.Event()
|
||||
observed = []
|
||||
observed_guard = threading.Lock()
|
||||
|
||||
def _race():
|
||||
start.wait(5)
|
||||
lock = ensure_run_lock(obj)
|
||||
with observed_guard:
|
||||
observed.append(lock)
|
||||
|
||||
threads = [threading.Thread(target=_race) for _ in range(16)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
start.set()
|
||||
for t in threads:
|
||||
t.join(5)
|
||||
|
||||
assert len(observed) == 16
|
||||
assert all(lock is observed[0] for lock in observed)
|
||||
assert obj._run_lock is observed[0]
|
||||
@@ -13,7 +13,6 @@ Everything that needs a live daemon is marked, so it is skipped loudly rather
|
||||
than passing quietly when Docker is not running.
|
||||
"""
|
||||
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
import pytest
|
||||
@@ -308,11 +307,12 @@ def test_a_hosted_backend_reclaims_its_instance_when_collected():
|
||||
instance kept billing until the provider's own idle timer noticed, which
|
||||
docker and flyio do not have."""
|
||||
import gc
|
||||
import time
|
||||
import threading
|
||||
|
||||
from praisonai.integrations.compute_managed_agent import ComputeManagedAgent
|
||||
|
||||
released = []
|
||||
release_done = threading.Event()
|
||||
|
||||
class Fake:
|
||||
provider_name = "fake"
|
||||
@@ -328,6 +328,7 @@ def test_a_hosted_backend_reclaims_its_instance_when_collected():
|
||||
|
||||
async def shutdown(self, instance_id):
|
||||
released.append(instance_id)
|
||||
release_done.set()
|
||||
|
||||
import asyncio
|
||||
|
||||
@@ -339,12 +340,7 @@ def test_a_hosted_backend_reclaims_its_instance_when_collected():
|
||||
|
||||
del backend
|
||||
gc.collect()
|
||||
# During live GC the finalizer reclaims fire-and-forget on the shared
|
||||
# AsyncBridge's background loop, so the shutdown lands a beat after
|
||||
# gc.collect() returns. Wait (bounded) for it rather than racing it.
|
||||
deadline = time.monotonic() + 5.0
|
||||
while not released and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert release_done.wait(5), "the background instance release did not complete"
|
||||
assert released == ["inst-1"], "the instance outlived the backend that owned it"
|
||||
|
||||
|
||||
@@ -368,25 +364,57 @@ def test_release_reclaims_even_on_an_event_loop_thread():
|
||||
ever runs, silently leaking the instance -- so _release must bridge the loop
|
||||
the way SharedCompute does, not call asyncio.run() directly."""
|
||||
import asyncio
|
||||
import time
|
||||
import threading
|
||||
|
||||
from praisonai.integrations.compute_managed_agent import _release
|
||||
|
||||
released = []
|
||||
release_done = threading.Event()
|
||||
|
||||
class Fake:
|
||||
async def shutdown(self, instance_id):
|
||||
released.append(instance_id)
|
||||
release_done.set()
|
||||
|
||||
async def collect_on_the_loop():
|
||||
_release(Fake(), "inst-loop", "docker")
|
||||
assert await asyncio.to_thread(release_done.wait, 5), "background release did not complete"
|
||||
|
||||
asyncio.run(collect_on_the_loop())
|
||||
# _release routes through the shared AsyncBridge's background loop
|
||||
# fire-and-forget, so the shutdown lands a beat after asyncio.run()
|
||||
# returns. Wait (bounded) for it rather than racing it — the assertion
|
||||
# still fails if _release ever leaks (never schedules the shutdown).
|
||||
deadline = time.monotonic() + 5.0
|
||||
while not released and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert released == ["inst-loop"], "an instance leaked when GC ran on a live loop"
|
||||
|
||||
|
||||
def test_release_returns_before_a_slow_provider_finishes():
|
||||
"""A live event loop must remain runnable while the cloud teardown waits."""
|
||||
import asyncio
|
||||
import threading
|
||||
|
||||
from praisonai.integrations.compute_managed_agent import _release
|
||||
|
||||
started = threading.Event()
|
||||
allow_completion = threading.Event()
|
||||
completed = threading.Event()
|
||||
released = []
|
||||
|
||||
class Fake:
|
||||
async def shutdown(self, instance_id):
|
||||
started.set()
|
||||
try:
|
||||
assert await asyncio.to_thread(allow_completion.wait, 5)
|
||||
released.append(instance_id)
|
||||
finally:
|
||||
completed.set()
|
||||
|
||||
async def collect_on_the_loop():
|
||||
try:
|
||||
_release(Fake(), "slow-instance", "docker")
|
||||
assert await asyncio.to_thread(started.wait, 5)
|
||||
assert not completed.is_set(), "release blocked the caller until teardown ended"
|
||||
# This caller's loop can still advance before permitting teardown.
|
||||
await asyncio.sleep(0)
|
||||
finally:
|
||||
allow_completion.set()
|
||||
assert await asyncio.to_thread(completed.wait, 5)
|
||||
|
||||
asyncio.run(collect_on_the_loop())
|
||||
assert released == ["slow-instance"]
|
||||
@@ -88,6 +88,31 @@ def test_slack_is_clamped_non_negative():
|
||||
assert quota_hold_from_failure(_RateLimited(100), now, slack_seconds=-5.0) == 100.0
|
||||
|
||||
|
||||
def test_nan_retry_after_does_not_produce_nan_hold():
|
||||
# A 429 carrying a malformed ``nan`` Retry-After must not park the job with
|
||||
# a NaN hold instant (issue #5424) — it carries no usable window → no park.
|
||||
import math
|
||||
|
||||
hold = quota_hold_from_failure(_RateLimited("nan"), 1000.0, slack_seconds=60.0)
|
||||
assert hold is None or math.isfinite(hold)
|
||||
assert hold is None
|
||||
|
||||
|
||||
def test_negative_retry_after_does_not_hold():
|
||||
# A negative Retry-After is not a usable future window → no park.
|
||||
assert quota_hold_from_failure(_RateLimited(-30), 1000.0) is None
|
||||
|
||||
|
||||
def test_echoed_negative_retry_after_in_message_does_not_hold():
|
||||
# A provider that echoes a negative Retry-After in the error message (e.g.
|
||||
# "retry after -3600 seconds") must not be read as a positive window that
|
||||
# parks the job on an invalid reset hint (issue #5424 / greptile #1).
|
||||
hold = quota_hold_from_failure(
|
||||
"429 Too Many Requests, retry after -3600 seconds", 1000.0, slack_seconds=60.0
|
||||
)
|
||||
assert hold is None
|
||||
|
||||
|
||||
# ── is_due hold guard ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -11,7 +11,8 @@ from praisonaiagents.workflows.step_cache import (
|
||||
make_step_key,
|
||||
resolve_step_cache,
|
||||
)
|
||||
from praisonaiagents.workflows.workflows import AgentFlow, Parallel
|
||||
from praisonaiagents.task.task import Task
|
||||
from praisonaiagents.workflows.workflows import AgentFlow, Parallel, StepResult
|
||||
|
||||
|
||||
def _counting_step(name, calls):
|
||||
@@ -55,6 +56,169 @@ class TestCachingWorks:
|
||||
|
||||
|
||||
class TestCorrectnessOnAHit:
|
||||
@pytest.mark.parametrize("boundary", ["miss", "hit"])
|
||||
@pytest.mark.parametrize("field", ["output", "variables", "steps"])
|
||||
def test_reference_cache_does_not_share_live_payloads(self, boundary, field):
|
||||
class ReferenceCache:
|
||||
def __init__(self):
|
||||
self.entries = {}
|
||||
|
||||
def get(self, key):
|
||||
return self.entries.get(key)
|
||||
|
||||
def set(self, key, value):
|
||||
self.entries[key] = value
|
||||
|
||||
calls = []
|
||||
|
||||
def produce(ctx):
|
||||
calls.append("produce")
|
||||
return StepResult(output={"items": [1]}, variables={"records": [{"value": 1}]})
|
||||
|
||||
flow = AgentFlow(steps=[produce], cache=ReferenceCache())
|
||||
result = flow.run("same", verbose=False)
|
||||
if boundary == "hit":
|
||||
result = flow.run("same", verbose=False)
|
||||
if field == "output":
|
||||
result["output"]["items"].append(999)
|
||||
elif field == "variables":
|
||||
result["variables"]["records"][0]["value"] = 999
|
||||
else:
|
||||
result["steps"][0]["status"] = "corrupted"
|
||||
later = flow.run("same", verbose=False)
|
||||
assert later["output"] == {"items": [1]}
|
||||
assert later["variables"]["records"] == [{"value": 1}]
|
||||
assert later["steps"][0]["status"] == "completed"
|
||||
assert calls == ["produce"]
|
||||
|
||||
def test_uncopyable_handler_variables_are_not_cached_by_reference(self):
|
||||
import threading
|
||||
|
||||
calls = []
|
||||
|
||||
def produce(ctx):
|
||||
calls.append("produce")
|
||||
return StepResult(
|
||||
output="ready", variables={"lock": threading.Lock(), "items": [1]}
|
||||
)
|
||||
|
||||
flow = AgentFlow(steps=[produce], cache=True)
|
||||
first = flow.run("same", verbose=False)
|
||||
first["variables"]["items"].append(999)
|
||||
second = flow.run("same", verbose=False)
|
||||
assert second["variables"]["items"] == [1]
|
||||
assert calls == ["produce", "produce"]
|
||||
|
||||
@pytest.mark.parametrize("boundary", ["miss", "hit"])
|
||||
def test_nested_reference_cache_snapshots_variables(self, boundary):
|
||||
class ReferenceCache:
|
||||
def __init__(self):
|
||||
self.entries = {}
|
||||
|
||||
def get(self, key):
|
||||
return self.entries.get(key)
|
||||
|
||||
def set(self, key, value):
|
||||
self.entries[key] = value
|
||||
|
||||
calls = []
|
||||
|
||||
def produce(ctx):
|
||||
calls.append("produce")
|
||||
return StepResult(output="ready", variables={"records": [{"value": 1}]})
|
||||
|
||||
flow = AgentFlow(steps=[Parallel(steps=[produce])], cache=ReferenceCache())
|
||||
result = flow.run("same", verbose=False)
|
||||
if boundary == "hit":
|
||||
result = flow.run("same", verbose=False)
|
||||
result["variables"]["records"][0]["value"] = 999
|
||||
assert flow.run("same", verbose=False)["variables"]["records"] == [{"value": 1}]
|
||||
assert calls == ["produce"]
|
||||
|
||||
def test_custom_cache_rejecting_variables_does_not_abort_workflow(self):
|
||||
import json
|
||||
import threading
|
||||
|
||||
class JsonCache:
|
||||
def get(self, key):
|
||||
return None
|
||||
|
||||
def set(self, key, value):
|
||||
json.dumps(value)
|
||||
|
||||
def produce(ctx):
|
||||
return StepResult(output="ready", variables={"lock": threading.Lock()})
|
||||
|
||||
flow = AgentFlow(steps=[produce], cache=JsonCache())
|
||||
result = flow.run("same", verbose=False)
|
||||
assert result["output"] == "ready"
|
||||
assert "lock" in result["variables"]
|
||||
assert result["steps"][0]["status"] == "completed"
|
||||
|
||||
def test_legacy_cache_entry_restores_output_variable_and_status(self):
|
||||
cache = InMemoryStepCache()
|
||||
cache.set(make_step_key("upstream", None, "same", {"input": "same"}), {
|
||||
"output": "cached", "variables": {"upstream_output": "cached"},
|
||||
})
|
||||
calls = []
|
||||
flow = AgentFlow(steps=[_counting_step("upstream", calls)], cache=cache)
|
||||
result = flow.run("same", verbose=False)
|
||||
assert calls == []
|
||||
assert result["output"] == "cached"
|
||||
assert result["variables"]["upstream_output"] == "cached"
|
||||
assert result["steps"] == [{
|
||||
"step": "upstream", "output": "cached", "status": "completed", "retries": 0,
|
||||
}]
|
||||
assert flow.step_statuses["upstream"] == "completed"
|
||||
|
||||
def test_a_cache_hit_preserves_early_stop(self):
|
||||
calls = []
|
||||
|
||||
def finish(ctx):
|
||||
calls.append("finish")
|
||||
return StepResult(
|
||||
output="done", stop_workflow=True, variables={"finished": True}
|
||||
)
|
||||
|
||||
def downstream(ctx):
|
||||
calls.append("downstream")
|
||||
return "unexpected"
|
||||
|
||||
flow = AgentFlow(steps=[finish, downstream], cache=True)
|
||||
first = flow.run("same", verbose=False)
|
||||
second = flow.run("same", verbose=False)
|
||||
|
||||
assert calls == ["finish"]
|
||||
assert second == first
|
||||
assert "finish_output" not in second["variables"]
|
||||
|
||||
def test_a_cache_hit_restores_handler_variables_for_downstream(self):
|
||||
calls = []
|
||||
|
||||
def produce(ctx):
|
||||
calls.append("produce")
|
||||
return StepResult(output="ready", variables={"answer": 42})
|
||||
|
||||
def consume(ctx):
|
||||
calls.append("consume")
|
||||
return str(ctx.variables.get("answer", "missing"))
|
||||
|
||||
flow = AgentFlow(steps=[produce, consume], cache=True)
|
||||
first = flow.run("same", verbose=False)
|
||||
second = flow.run("same", verbose=False)
|
||||
|
||||
assert first["output"] == second["output"] == "42"
|
||||
assert second["variables"]["answer"] == 42
|
||||
assert calls == ["produce", "consume"]
|
||||
|
||||
def test_a_cache_hit_reports_completed_step_status(self):
|
||||
flow = AgentFlow(steps=[_counting_step("upstream", [])], cache=True)
|
||||
first = flow.run("same", verbose=False)
|
||||
second = flow.run("same", verbose=False)
|
||||
|
||||
assert second["steps"] == first["steps"]
|
||||
assert flow.step_statuses["upstream"] == "completed"
|
||||
|
||||
def test_a_cache_hit_still_exposes_the_step_output_variable(self):
|
||||
"""A fresh run starts with empty working variables. A hit that restored
|
||||
only `output` would leave `<step>_output` missing from the run's
|
||||
@@ -82,6 +246,71 @@ class TestCorrectnessOnAHit:
|
||||
flow.run("same", verbose=False) # must retry, not serve cached failure
|
||||
assert attempts["n"] == 2
|
||||
|
||||
def test_a_cached_stop_still_stops_the_workflow(self):
|
||||
"""A handler that stopped the cold run must also stop the cached run.
|
||||
Losing the stop flag on a hit ran downstream steps that the cold run
|
||||
never reached."""
|
||||
calls = []
|
||||
|
||||
def finish(ctx):
|
||||
calls.append("finish")
|
||||
return StepResult(output="done", stop_workflow=True)
|
||||
|
||||
finish.__name__ = "finish"
|
||||
|
||||
def downstream(ctx):
|
||||
calls.append("downstream")
|
||||
return "unexpected"
|
||||
|
||||
downstream.__name__ = "downstream"
|
||||
|
||||
flow = AgentFlow(steps=[finish, downstream], cache=True)
|
||||
flow.run("same", verbose=False)
|
||||
flow.run("same", verbose=False)
|
||||
assert calls == ["finish"]
|
||||
|
||||
def test_a_cached_hit_preserves_handler_variables(self):
|
||||
"""A producer returning variables followed by a consumer reading them
|
||||
must give the same output on the cached run as the cold one."""
|
||||
def producer(ctx):
|
||||
return StepResult(output="p", variables={"answer": 42})
|
||||
|
||||
producer.__name__ = "producer"
|
||||
|
||||
def consumer(ctx):
|
||||
return f"answer={ctx.variables.get('answer')}"
|
||||
|
||||
consumer.__name__ = "consumer"
|
||||
|
||||
flow = AgentFlow(steps=[producer, consumer], cache=True)
|
||||
first = flow.run("same", verbose=False)
|
||||
second = flow.run("same", verbose=False)
|
||||
assert first["output"] == "answer=42"
|
||||
assert second["output"] == first["output"]
|
||||
|
||||
def test_a_cached_hit_preserves_completion_metadata(self):
|
||||
"""Step status and retries must survive a cache hit, not disappear from
|
||||
the result records."""
|
||||
flow = AgentFlow(steps=[_counting_step("s", [])], cache=True)
|
||||
flow.run("same", verbose=False)
|
||||
second = flow.run("same", verbose=False)
|
||||
record = second["steps"][0]
|
||||
assert record["status"] == "completed"
|
||||
assert record["retries"] == 0
|
||||
|
||||
def test_an_early_stop_hit_does_not_insert_the_output_variable(self):
|
||||
"""The cold path stops before writing `<step>_output`; a cached early
|
||||
stop must not insert it either."""
|
||||
def finish(ctx):
|
||||
return StepResult(output="done", stop_workflow=True)
|
||||
|
||||
finish.__name__ = "finish"
|
||||
|
||||
flow = AgentFlow(steps=[finish], cache=True)
|
||||
flow.run("same", verbose=False)
|
||||
second = flow.run("same", verbose=False)
|
||||
assert "finish_output" not in second["variables"]
|
||||
|
||||
def test_a_mutated_hit_does_not_corrupt_the_cache(self):
|
||||
"""Values handed back on a hit are snapshots; mutating one must not
|
||||
change what the next hit returns."""
|
||||
@@ -94,6 +323,66 @@ class TestCorrectnessOnAHit:
|
||||
assert second["variables"]["a"] == 1
|
||||
assert second["output"] == "v"
|
||||
|
||||
def test_a_cached_hit_preserves_variables_from_rejected_retries(self):
|
||||
"""A handler whose earlier (guardrail-rejected) attempt wrote extra
|
||||
variables leaves those in the live run; the accepted attempt's variables
|
||||
add to them. A hit must replay the union, not only the last attempt, or
|
||||
a downstream step sees a different variable set than on the cold run."""
|
||||
attempts = {"n": 0}
|
||||
|
||||
def producer(ctx):
|
||||
attempts["n"] += 1
|
||||
if attempts["n"] == 1:
|
||||
return StepResult(
|
||||
output="bad",
|
||||
variables={"shared": 1, "only_first": 2},
|
||||
)
|
||||
return StepResult(output="good", variables={"shared": 3})
|
||||
|
||||
def reject_first(result):
|
||||
return (result.output != "bad", "retry")
|
||||
|
||||
producer_task = Task(
|
||||
name="producer", handler=producer, guardrails=reject_first
|
||||
)
|
||||
|
||||
def consumer(ctx):
|
||||
return (
|
||||
f"shared={ctx.variables.get('shared')},"
|
||||
f"only_first={ctx.variables.get('only_first')}"
|
||||
)
|
||||
|
||||
consumer.__name__ = "consumer"
|
||||
|
||||
flow = AgentFlow(steps=[producer_task, consumer], cache=True)
|
||||
first = flow.run("same", verbose=False)
|
||||
assert first["output"] == "shared=3,only_first=2"
|
||||
attempts["n"] = 99 # a hit must not re-run the producer
|
||||
second = flow.run("same", verbose=False)
|
||||
assert second["output"] == first["output"]
|
||||
|
||||
def test_a_cached_mutable_variable_is_isolated_between_runs(self):
|
||||
"""A handler returns a mutable object; a downstream step on the cold run
|
||||
mutates it. That mutation must not leak back into the cache, so a later
|
||||
hit replays the pristine value the producer originally returned."""
|
||||
def producer(ctx):
|
||||
return StepResult(output="p", variables={"items": [1, 2]})
|
||||
|
||||
producer.__name__ = "producer"
|
||||
|
||||
def consumer(ctx):
|
||||
ctx.variables["items"].append(999) # mutate the live copy
|
||||
return f"len={len(ctx.variables['items'])}"
|
||||
|
||||
consumer.__name__ = "consumer"
|
||||
|
||||
flow = AgentFlow(steps=[producer, consumer], cache=True)
|
||||
flow.run("same", verbose=False) # cold run mutates items -> [1,2,999]
|
||||
second = flow.run("same", verbose=False) # served from cache
|
||||
# The producer's cached snapshot must be unaffected by the cold run's
|
||||
# in-place mutation; without isolation it would replay [1, 2, 999].
|
||||
assert second["variables"]["items"] == [1, 2]
|
||||
|
||||
|
||||
class TestKeys:
|
||||
def test_the_same_inputs_key_the_same(self):
|
||||
@@ -125,6 +414,14 @@ class TestKeys:
|
||||
|
||||
|
||||
class TestTheCacheItself:
|
||||
def test_uncopyable_replacement_removes_the_old_entry(self):
|
||||
import threading
|
||||
|
||||
cache = InMemoryStepCache()
|
||||
cache.set("k", {"output": "old"})
|
||||
cache.set("k", {"output": "new", "lock": threading.Lock()})
|
||||
assert cache.get("k") is None
|
||||
|
||||
def test_it_is_bounded(self):
|
||||
"""Unbounded would be a memory leak that only shows up in production."""
|
||||
cache = InMemoryStepCache(max_entries=2)
|
||||
|
||||
@@ -52,7 +52,12 @@ def test_no_isolation_shares_root(tmp_path):
|
||||
adapter.remove("a")
|
||||
|
||||
|
||||
def test_git_worktree_degrades_when_not_a_repo(tmp_path):
|
||||
def test_git_worktree_degrades_when_not_a_repo(tmp_path, monkeypatch):
|
||||
# A temporary directory can have a repository ancestor (including /).
|
||||
# Keep this negative fixture outside that ancestor's discovery scope.
|
||||
monkeypatch.setenv("GIT_CEILING_DIRECTORIES", str(tmp_path.parent))
|
||||
monkeypatch.delenv("GIT_DIR", raising=False)
|
||||
monkeypatch.delenv("GIT_WORK_TREE", raising=False)
|
||||
adapter = GitWorktreeAdapter(root=tmp_path)
|
||||
assert adapter.available is False
|
||||
assert adapter.create("run") == str(tmp_path.resolve())
|
||||
|
||||
Generated
+10
-6
@@ -1,5 +1,5 @@
|
||||
version = 1
|
||||
revision = 3
|
||||
revision = 5
|
||||
requires-python = ">=3.10"
|
||||
resolution-markers = [
|
||||
"python_full_version >= '3.14'",
|
||||
@@ -2263,7 +2263,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "mcp"
|
||||
version = "1.25.0"
|
||||
version = "1.30.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "anyio" },
|
||||
@@ -2281,9 +2281,9 @@ dependencies = [
|
||||
{ name = "typing-inspection" },
|
||||
{ name = "uvicorn", marker = "sys_platform != 'emscripten'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d5/2d/649d80a0ecf6a1f82632ca44bec21c0461a9d9fc8934d38cb5b319f2db5e/mcp-1.25.0.tar.gz", hash = "sha256:56310361ebf0364e2d438e5b45f7668cbb124e158bb358333cd06e49e83a6802", size = 605387, upload-time = "2025-12-19T10:19:56.985Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/ba/93/0142dc84a666daf8ad51a34268f34c12fd6fda4f3810c4be2504eecc8212/mcp-1.30.0.tar.gz", hash = "sha256:445414625fce5c295faa505bb11bacece661ab6f4028d57c935db57820b7a3e4", size = 680511, upload-time = "2026-09-07T14:34:15.845Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/e2/fc/6dc7659c2ae5ddf280477011f4213a74f806862856b796ef08f028e664bf/mcp-1.25.0-py3-none-any.whl", hash = "sha256:b37c38144a666add0862614cc79ec276e97d72aa8ca26d622818d4e278b9721a", size = 233076, upload-time = "2025-12-19T10:19:55.416Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f5/f4/e58bc33317c92a0203664daaf00bf6f41166cc0149e5d6870a03f7cd004a/mcp-1.30.0-py3-none-any.whl", hash = "sha256:666edb5009503e1047c9d60346a756f94b261f05cc2625f23d41c728ffc484d0", size = 234581, upload-time = "2026-09-07T14:34:14.266Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3123,7 +3123,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.9"
|
||||
version = "1.7.10"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -3219,6 +3219,9 @@ os = [
|
||||
{ name = "pyjwt" },
|
||||
{ name = "uvicorn" },
|
||||
]
|
||||
parallel-search = [
|
||||
{ name = "mcp" },
|
||||
]
|
||||
sandbox = [
|
||||
{ name = "praisonai-sandbox" },
|
||||
]
|
||||
@@ -3250,6 +3253,7 @@ requires-dist = [
|
||||
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14" },
|
||||
{ name = "markitdown", extras = ["all"], marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
|
||||
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.20.0" },
|
||||
{ name = "mcp", marker = "extra == 'parallel-search'", specifier = ">=1.30.0,<2" },
|
||||
{ name = "mem0ai", marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
|
||||
{ name = "mem0ai", extras = ["graph"], marker = "extra == 'graph'", specifier = ">=0.1.0" },
|
||||
{ name = "motor", marker = "extra == 'mongodb'", specifier = ">=3.4.0" },
|
||||
@@ -3288,7 +3292,7 @@ requires-dist = [
|
||||
{ name = "uvicorn", marker = "extra == 'os'", specifier = ">=0.34.0" },
|
||||
{ name = "websockets", marker = "extra == 'mcp'", specifier = ">=12.0" },
|
||||
]
|
||||
provides-extras = ["sandbox", "mcp", "memory", "knowledge", "graph", "llm", "api", "os", "telemetry", "mongodb", "dakera", "auth", "autonomy", "search", "crawl", "a2ui", "all"]
|
||||
provides-extras = ["sandbox", "mcp", "parallel-search", "memory", "knowledge", "graph", "llm", "api", "os", "telemetry", "mongodb", "dakera", "auth", "autonomy", "search", "crawl", "a2ui", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "primp"
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
# yaml-language-server: $schema=https://raw.githubusercontent.com/MervinPraison/PraisonAI/main/src/praisonai-bot/praisonai_bot/bots/gateway.schema.json
|
||||
gateway:
|
||||
host: 127.0.0.1
|
||||
port: 8765
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.1.8"
|
||||
__version__ = "0.1.9"
|
||||
@@ -31,6 +31,7 @@ finally to ``None``.
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import threading
|
||||
from typing import Any, Dict, Iterator, List, Optional, Tuple
|
||||
|
||||
_NON_DIGITS = re.compile(r"[^0-9]")
|
||||
@@ -56,12 +57,18 @@ def normalize_number(number: Optional[str]) -> Optional[str]:
|
||||
|
||||
|
||||
class AgentRegistry:
|
||||
"""A thread-unsafe, in-memory phone-number → agent routing table.
|
||||
"""A thread-safe, in-memory phone-number → agent routing table.
|
||||
|
||||
Lightweight by design: it holds references to already-constructed agents
|
||||
and a normalised-number index. It performs no I/O and adds no dependencies,
|
||||
so a gateway can build one per process and consult it on every inbound
|
||||
message.
|
||||
and a normalised-number index. It performs no I/O and adds no dependencies
|
||||
beyond the stdlib, so a gateway can build one per process and consult it on
|
||||
every inbound message.
|
||||
|
||||
Access is guarded by a per-instance :class:`threading.Lock` so concurrent
|
||||
inbound turns (the gateway dispatches turns concurrently) cannot race on
|
||||
the shared mapping: a simultaneous ``assign``/``unassign`` can never be
|
||||
observed as a torn read by ``resolve``, and the ``__len__``/``__iter__``/
|
||||
``numbers`` snapshots are taken atomically.
|
||||
|
||||
Args:
|
||||
default_agent: Optional fallback returned by :meth:`resolve` when an
|
||||
@@ -72,6 +79,7 @@ class AgentRegistry:
|
||||
def __init__(self, default_agent: Optional[Any] = None) -> None:
|
||||
self._by_number: Dict[str, Any] = {}
|
||||
self._default_agent = default_agent
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def assign(self, number: str, agent: Any) -> str:
|
||||
"""Assign ``number`` to ``agent`` and return the normalised key.
|
||||
@@ -84,7 +92,8 @@ class AgentRegistry:
|
||||
key = normalize_number(number)
|
||||
if key is None:
|
||||
raise ValueError(f"invalid phone number: {number!r}")
|
||||
self._by_number[key] = agent
|
||||
with self._lock:
|
||||
self._by_number[key] = agent
|
||||
return key
|
||||
|
||||
def unassign(self, number: str) -> bool:
|
||||
@@ -95,7 +104,8 @@ class AgentRegistry:
|
||||
key = normalize_number(number)
|
||||
if key is None:
|
||||
return False
|
||||
return self._by_number.pop(key, None) is not None
|
||||
with self._lock:
|
||||
return self._by_number.pop(key, None) is not None
|
||||
|
||||
def resolve(self, number: Optional[str]) -> Optional[Any]:
|
||||
"""Return the agent assigned to ``number``.
|
||||
@@ -105,26 +115,33 @@ class AgentRegistry:
|
||||
"""
|
||||
key = normalize_number(number)
|
||||
if key is not None:
|
||||
agent = self._by_number.get(key)
|
||||
with self._lock:
|
||||
agent = self._by_number.get(key)
|
||||
if agent is not None:
|
||||
return agent
|
||||
return self._default_agent
|
||||
|
||||
def numbers(self) -> List[str]:
|
||||
"""Return all assigned (normalised) numbers."""
|
||||
return list(self._by_number.keys())
|
||||
with self._lock:
|
||||
return list(self._by_number.keys())
|
||||
|
||||
def __contains__(self, number: object) -> bool:
|
||||
if not isinstance(number, str):
|
||||
return False
|
||||
key = normalize_number(number)
|
||||
return key is not None and key in self._by_number
|
||||
if key is None:
|
||||
return False
|
||||
with self._lock:
|
||||
return key in self._by_number
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._by_number)
|
||||
with self._lock:
|
||||
return len(self._by_number)
|
||||
|
||||
def __iter__(self) -> Iterator[Tuple[str, Any]]:
|
||||
return iter(self._by_number.items())
|
||||
with self._lock:
|
||||
return iter(dict(self._by_number).items())
|
||||
|
||||
|
||||
__all__ = ["AgentRegistry", "normalize_number"]
|
||||
@@ -23,6 +23,41 @@ from pydantic import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
#: Stable, published URL for the ``gateway.yaml`` / ``bot.yaml`` JSON Schema,
|
||||
#: mirroring the ``agents.yaml`` convention (``AGENTS_SCHEMA_URL``). Editors that
|
||||
#: speak the YAML language server use this via a leading
|
||||
#: ``# yaml-language-server: $schema=<url>`` header to provide autocomplete,
|
||||
#: inline validation, and hover docs while authoring the gateway/bot config.
|
||||
GATEWAY_SCHEMA_URL = (
|
||||
"https://raw.githubusercontent.com/MervinPraison/PraisonAI/main/"
|
||||
"src/praisonai-bot/praisonai_bot/bots/gateway.schema.json"
|
||||
)
|
||||
|
||||
#: Leading YAML comment prepended to scaffolded ``gateway.yaml`` / ``bot.yaml``
|
||||
#: files so editors wire up validation out of the box. A leading comment is
|
||||
#: ignored by ``yaml.safe_load``, so execution is unaffected.
|
||||
GATEWAY_SCHEMA_HEADER = f"# yaml-language-server: $schema={GATEWAY_SCHEMA_URL}\n"
|
||||
|
||||
|
||||
def resolve_model_by_provider() -> str:
|
||||
"""Resolve a default model for whichever provider credential is present.
|
||||
|
||||
Reuses the existing detection helper
|
||||
(``praisonai_code.llm.env.default_model_for_available_provider``) through
|
||||
the lazy cross-tier bridge so the gateway path leads with the provider the
|
||||
operator already has (Anthropic/Gemini/Ollama/…) instead of assuming
|
||||
OpenAI. Falls back to ``gpt-4o-mini`` only when no provider is detected or
|
||||
the helper is unavailable.
|
||||
"""
|
||||
try:
|
||||
from praisonai_bot._code_bridge import import_code_module
|
||||
|
||||
env = import_code_module("praisonai_code.llm.env")
|
||||
return env.default_model_for_available_provider()
|
||||
except Exception: # pragma: no cover - detection is best-effort
|
||||
return "gpt-4o-mini"
|
||||
|
||||
|
||||
def _register_redaction(value: str) -> None:
|
||||
"""Register a resolved secret value for log redaction (best-effort).
|
||||
|
||||
@@ -41,7 +76,12 @@ class AgentConfigSchema(BaseModel):
|
||||
"""Schema for agent configuration in bot.yaml."""
|
||||
name: str = "assistant"
|
||||
instructions: str = ""
|
||||
model: Optional[str] = "gpt-4o-mini"
|
||||
# Provider-agnostic by default (Issue #5609). ``None`` / ``"auto"`` resolve
|
||||
# at load time to the model for whichever provider credential is actually
|
||||
# present (via ``resolved_model()``), so an operator with only an
|
||||
# Anthropic/Gemini/Ollama key is not silently handed an OpenAI model. An
|
||||
# explicit ``model:`` always wins.
|
||||
model: Optional[str] = None
|
||||
llm: Optional[str] = None # Alias for model
|
||||
memory: bool = False
|
||||
tools: List[str] = Field(default_factory=list)
|
||||
@@ -58,6 +98,19 @@ class AgentConfigSchema(BaseModel):
|
||||
self.model = self.llm
|
||||
return self
|
||||
|
||||
def resolved_model(self) -> str:
|
||||
"""Return the effective model, detecting the provider when unset.
|
||||
|
||||
An explicit, non-``"auto"`` ``model`` (or ``llm`` alias) always wins.
|
||||
When the model is unset or ``"auto"``, resolve it to the detected
|
||||
provider's default (Issue #5609), with ``gpt-4o-mini`` as the
|
||||
last-resort fallback handled inside ``resolve_model_by_provider``.
|
||||
"""
|
||||
model = self.model
|
||||
if model and model.strip().lower() != "auto":
|
||||
return model
|
||||
return resolve_model_by_provider()
|
||||
|
||||
|
||||
class SessionResetConfigSchema(BaseModel):
|
||||
"""Schema for session reset policy configuration."""
|
||||
@@ -1506,6 +1559,184 @@ class GatewayConfigSchema(BaseModel):
|
||||
BotYamlSchema = GatewayConfigSchema
|
||||
|
||||
|
||||
def _allow_integer_owner_id(node: Any) -> None:
|
||||
"""Widen an ``owner_user_id`` property to also accept an integer (#5414).
|
||||
|
||||
``GatewayConfigSchema``/``ChannelConfigSchema`` type ``owner_user_id`` as
|
||||
``Optional[str]`` but a ``mode="before"`` validator (``_coerce_owner_user_id``)
|
||||
accepts an unquoted numeric YAML id (``owner_user_id: 987654321``) by casting
|
||||
it to ``str``. Pydantic emits the schema from the *field type* only, so the
|
||||
published schema would reject a config the runtime explicitly accepts. Mirror
|
||||
the runtime coercion by adding ``integer`` to the allowed types so editor/CI
|
||||
validation agrees with the gateway.
|
||||
"""
|
||||
if not isinstance(node, dict):
|
||||
return
|
||||
props = node.get("properties")
|
||||
if not isinstance(props, dict):
|
||||
return
|
||||
owner = props.get("owner_user_id")
|
||||
if not isinstance(owner, dict):
|
||||
return
|
||||
variants = owner.get("anyOf")
|
||||
if isinstance(variants, list):
|
||||
has_int = any(
|
||||
isinstance(v, dict) and v.get("type") == "integer" for v in variants
|
||||
)
|
||||
if not has_int:
|
||||
variants.append({"type": "integer"})
|
||||
elif owner.get("type") == "string":
|
||||
owner.pop("type", None)
|
||||
owner["anyOf"] = [{"type": "string"}, {"type": "integer"}]
|
||||
|
||||
|
||||
def _allow_env_ref_list(node: Any, field_names: tuple) -> None:
|
||||
"""Allow an array-typed field to also be a single ``${ENV}`` string (#5414).
|
||||
|
||||
``${VAR}`` references are expanded *before* validation at load time, so an
|
||||
operator may author ``allowed_users: ${TELEGRAM_ALLOWED_USERS}`` and have it
|
||||
expand to a list. Editors validate the *unexpanded* YAML, so the published
|
||||
schema must accept the string form for these fields too; otherwise it
|
||||
false-flags a config the runtime happily accepts.
|
||||
"""
|
||||
if not isinstance(node, dict):
|
||||
return
|
||||
props = node.get("properties")
|
||||
if not isinstance(props, dict):
|
||||
return
|
||||
for name in field_names:
|
||||
field = props.get(name)
|
||||
if not isinstance(field, dict):
|
||||
continue
|
||||
if field.get("type") == "array":
|
||||
array_variant = {k: v for k, v in field.items() if k not in ("default", "title")}
|
||||
new_field: Dict[str, Any] = {
|
||||
"anyOf": [array_variant, {"type": "string"}]
|
||||
}
|
||||
if "default" in field:
|
||||
new_field["default"] = field["default"]
|
||||
if "title" in field:
|
||||
new_field["title"] = field["title"]
|
||||
props[name] = new_field
|
||||
elif isinstance(field.get("anyOf"), list):
|
||||
variants = field["anyOf"]
|
||||
if not any(
|
||||
isinstance(v, dict) and v.get("type") == "string" for v in variants
|
||||
):
|
||||
variants.append({"type": "string"})
|
||||
|
||||
|
||||
def gateway_config_json_schema() -> Dict[str, Any]:
|
||||
"""Canonical JSON Schema for ``gateway.yaml`` / ``bot.yaml`` (draft-07).
|
||||
|
||||
Derived directly from :class:`GatewayConfigSchema` (Pydantic's
|
||||
``model_json_schema``) and decorated with the standard ``$schema``/``$id``/
|
||||
``title`` metadata so the artefact is self-describing and mirrors the
|
||||
``agents.yaml`` convention (``generate_agents_schema``). This is the single
|
||||
source of truth surfaced to editors (via the ``# yaml-language-server``
|
||||
header) and to CI linting, so a mistyped key is flagged at author time
|
||||
instead of only at gateway startup.
|
||||
|
||||
The raw ``model_json_schema()`` is post-processed so the *published* schema
|
||||
faithfully reflects what the runtime actually validates (#5414):
|
||||
|
||||
- ``gateway`` / ``hooks`` are stored as opaque ``Dict``/``List[Dict]`` on the
|
||||
model for downstream dict-style access, but validated field-by-field at
|
||||
load time via :class:`GatewayServerSchema` / :class:`HookSchema`
|
||||
(``extra="forbid"`` for the server block). Inline those typed schemas so a
|
||||
misspelled server knob (``drain_timout``) or hook key is caught at author
|
||||
time, not only at startup.
|
||||
- ``owner_user_id`` (top-level and per-channel) also accepts an integer,
|
||||
mirroring the runtime ``str(...)`` coercion of unquoted numeric ids.
|
||||
- Unknown top-level keys are rejected so a typo (``auto_enable_from_en``)
|
||||
that the runtime silently ignores is surfaced while authoring. Channel
|
||||
configs keep ``additionalProperties`` open for plugin-declared keys.
|
||||
"""
|
||||
base = GatewayConfigSchema.model_json_schema()
|
||||
defs = base.setdefault("$defs", {})
|
||||
|
||||
# Inline the typed gateway-server + hook schemas so the published schema
|
||||
# carries the same field-by-field constraints the runtime enforces, instead
|
||||
# of the opaque ``{"type": "object"}`` Pydantic emits for ``Dict[str, Any]``.
|
||||
server_schema = GatewayServerSchema.model_json_schema(
|
||||
ref_template="#/$defs/{model}"
|
||||
)
|
||||
for name, definition in server_schema.pop("$defs", {}).items():
|
||||
defs.setdefault(name, definition)
|
||||
defs["GatewayServerSchema"] = server_schema
|
||||
|
||||
hook_schema = HookSchema.model_json_schema(ref_template="#/$defs/{model}")
|
||||
for name, definition in hook_schema.pop("$defs", {}).items():
|
||||
defs.setdefault(name, definition)
|
||||
defs["HookSchema"] = hook_schema
|
||||
|
||||
props = base.get("properties", {})
|
||||
if "gateway" in props:
|
||||
props["gateway"] = {
|
||||
"anyOf": [{"$ref": "#/$defs/GatewayServerSchema"}, {"type": "null"}],
|
||||
"default": None,
|
||||
"title": props["gateway"].get("title", "Gateway"),
|
||||
}
|
||||
if "hooks" in props:
|
||||
props["hooks"] = {
|
||||
"anyOf": [
|
||||
{"type": "array", "items": {"$ref": "#/$defs/HookSchema"}},
|
||||
{"type": "null"},
|
||||
],
|
||||
"default": None,
|
||||
"title": props["hooks"].get("title", "Hooks"),
|
||||
}
|
||||
|
||||
# Accept unquoted numeric owner ids wherever they appear (top level +
|
||||
# ChannelConfigSchema), matching the runtime coercion.
|
||||
_allow_integer_owner_id(base)
|
||||
channel_def = defs.get("ChannelConfigSchema")
|
||||
if isinstance(channel_def, dict):
|
||||
_allow_integer_owner_id(channel_def)
|
||||
# ``${VAR}`` env references are substituted BEFORE validation at load time
|
||||
# (``validate_gateway_config``), so a list-typed field may legitimately be
|
||||
# authored as a single ``${ENV}`` string that expands to a list. The
|
||||
# scaffolder writes ``allowed_users: ${TELEGRAM_ALLOWED_USERS}``, so an
|
||||
# editor validating the *authored* (unexpanded) YAML must accept the
|
||||
# string form too — otherwise the schema false-flags a config the runtime
|
||||
# accepts. Widen the array-typed env-ref channel fields accordingly.
|
||||
_allow_env_ref_list(channel_def, ("allowed_users", "allowlist", "blocklist"))
|
||||
|
||||
# Reject unknown top-level keys so an author-time typo is caught (the runtime
|
||||
# silently ignores it, keeping the mistyped opt-out ineffective). Channels
|
||||
# keep ``additionalProperties`` open for plugin-declared config keys.
|
||||
#
|
||||
# ``config_version`` is not a ``GatewayConfigSchema`` field — it is stamped at
|
||||
# the top level by the doctor/start config-version machinery (and shipped in
|
||||
# the sample ``gateway.yaml``). Declare it so the stricter schema accepts a
|
||||
# stamped/migrated config instead of false-flagging it.
|
||||
base.setdefault("properties", {}).setdefault(
|
||||
"config_version",
|
||||
{
|
||||
"type": "integer",
|
||||
"title": "Config Version",
|
||||
"description": (
|
||||
"Stamped by `praisonai gateway doctor --fix` / `start` to track "
|
||||
"config migrations. Managed automatically."
|
||||
),
|
||||
},
|
||||
)
|
||||
base.setdefault("additionalProperties", False)
|
||||
|
||||
schema: Dict[str, Any] = {
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"$id": GATEWAY_SCHEMA_URL,
|
||||
**base,
|
||||
"title": "PraisonAI Gateway / Bot Configuration",
|
||||
"description": (
|
||||
"Schema for gateway.yaml / bot.yaml consumed by the PraisonAI "
|
||||
"gateway and single-bot runtime (gateway server, agents, channels, "
|
||||
"routing, schedules, hooks, streaming and delivery)."
|
||||
),
|
||||
}
|
||||
return schema
|
||||
|
||||
|
||||
def validate_gateway_config(raw: Dict[str, Any], apply_env_substitution: bool = True) -> GatewayConfigSchema:
|
||||
"""Validate a raw YAML dict against the unified gateway config schema.
|
||||
|
||||
|
||||
@@ -105,6 +105,78 @@ def resolve_activity_category(tool_name: Optional[str]) -> str:
|
||||
return "default"
|
||||
|
||||
|
||||
# Separator between the reply body and the runtime footer line. A blank line
|
||||
# then an em-dash keeps the footer visually distinct without markup that could
|
||||
# break on a given platform.
|
||||
_FOOTER_SEP = "\n\n— "
|
||||
|
||||
|
||||
def build_footer_line(
|
||||
model: Optional[str] = None,
|
||||
served_model: Optional[str] = None,
|
||||
context_pct: Optional[float] = None,
|
||||
latency_s: Optional[float] = None,
|
||||
) -> str:
|
||||
"""Build a privacy-safe, per-reply runtime footer line.
|
||||
|
||||
Renders only the fields that are available, e.g. ``gpt-4o · 38% ctx · 4.2s``.
|
||||
A missing field is omitted rather than shown as a placeholder, so a partial
|
||||
footer is always well-formed. Only the model name, context-window fill
|
||||
percentage and turn latency are surfaced — never tool args, URLs, paths or
|
||||
message content. Returns ``""`` when no field is available.
|
||||
|
||||
Args:
|
||||
model: Configured model id (e.g. ``"gpt-4o"``).
|
||||
served_model: Deployment that actually served the turn, shown as
|
||||
``model (served: X)`` when it differs from ``model`` (failover /
|
||||
routing visibility).
|
||||
context_pct: Session context-window fill, 0–100.
|
||||
latency_s: Turn latency in seconds.
|
||||
|
||||
Returns:
|
||||
A single footer line (without the leading separator) or ``""``.
|
||||
"""
|
||||
parts: list = []
|
||||
|
||||
if model:
|
||||
model_str = str(model).strip()
|
||||
if model_str:
|
||||
served = str(served_model).strip() if served_model else ""
|
||||
if served and served != model_str:
|
||||
parts.append(f"{model_str} (served: {served})")
|
||||
else:
|
||||
parts.append(model_str)
|
||||
|
||||
if context_pct is not None:
|
||||
try:
|
||||
pct = float(context_pct)
|
||||
if pct == pct and pct not in (float("inf"), float("-inf")):
|
||||
parts.append(f"{int(round(pct))}% ctx")
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
|
||||
if latency_s is not None:
|
||||
try:
|
||||
secs = float(latency_s)
|
||||
if secs == secs and secs >= 0 and secs not in (float("inf"), float("-inf")):
|
||||
parts.append(f"{secs:.1f}s")
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
|
||||
return " · ".join(parts)
|
||||
|
||||
|
||||
def append_footer(reply: str, footer_line: str) -> str:
|
||||
"""Append a footer line to a final reply, separated by a blank line.
|
||||
|
||||
Returns ``reply`` unchanged when ``footer_line`` is empty. Only ever used on
|
||||
the final reply — never on partial/streamed frames.
|
||||
"""
|
||||
if not footer_line:
|
||||
return reply
|
||||
return f"{reply}{_FOOTER_SEP}{footer_line}"
|
||||
|
||||
|
||||
def strip_reasoning_tags(text: str) -> str:
|
||||
"""Remove ``<think>``/``<reasoning>`` spans from streamed content.
|
||||
|
||||
@@ -201,6 +273,25 @@ class StreamingConfig:
|
||||
)
|
||||
|
||||
|
||||
def build_streaming_config(config: Any) -> Optional["StreamingConfig"]:
|
||||
"""Derive a :class:`StreamingConfig` from a ``BotConfig``.
|
||||
|
||||
Shared by every channel adapter (Telegram, Discord, Slack, …) so a
|
||||
configured ``streaming`` flag is wired to the same ``DraftStreamer`` engine
|
||||
identically everywhere, instead of each adapter re-deriving (or silently
|
||||
ignoring) the knob. Returns ``None`` when streaming is off so the caller's
|
||||
non-streaming path is unchanged.
|
||||
"""
|
||||
if not getattr(config, "streaming", False):
|
||||
return None
|
||||
interval_ms = getattr(config, "stream_edit_interval_ms", 700) or 700
|
||||
return StreamingConfig(
|
||||
mode=StreamingMode.DRAFT,
|
||||
min_interval=interval_ms / 1000.0,
|
||||
min_delta=50,
|
||||
)
|
||||
|
||||
|
||||
class BotAdapter(Protocol):
|
||||
"""Protocol for bot adapters that support message editing."""
|
||||
|
||||
|
||||
@@ -45,6 +45,12 @@ from ._commands import (
|
||||
)
|
||||
from ._failure import failure_reply_text
|
||||
from ._session import BotSessionManager
|
||||
from ._streaming import (
|
||||
StreamingConfig,
|
||||
StreamingMode,
|
||||
DraftStreamer,
|
||||
build_streaming_config,
|
||||
)
|
||||
from ._debounce import InboundDebouncer
|
||||
from ._ack import AckReactor
|
||||
from ._unknown_user import UnknownUserHandler, BotContext
|
||||
@@ -102,7 +108,12 @@ class DiscordBot(OutboundResilienceMixin, ChatCommandMixin, MessageHookMixin):
|
||||
|
||||
# Initialize allow_silence from config
|
||||
self._allow_silence = getattr(self.config, 'allow_silence', False)
|
||||
|
||||
|
||||
# Initialize streaming config based on BotConfig (shared with Telegram /
|
||||
# Slack so ``streaming.mode: draft|auto`` is honoured identically across
|
||||
# channels that support in-place edits — Issue #5415).
|
||||
self._streaming_config = build_streaming_config(self.config)
|
||||
|
||||
self._is_running = False
|
||||
self._bot_user: Optional[BotUser] = None
|
||||
self._client = None
|
||||
@@ -139,6 +150,15 @@ class DiscordBot(OutboundResilienceMixin, ChatCommandMixin, MessageHookMixin):
|
||||
self._register_interactive_handlers()
|
||||
self._bot_context: Optional[BotContext] = None
|
||||
|
||||
def configure_streaming(self, config: StreamingConfig) -> None:
|
||||
"""Configure streaming reply mode.
|
||||
|
||||
Args:
|
||||
config: Streaming configuration. Set mode=StreamingMode.OFF to disable.
|
||||
"""
|
||||
self._streaming_config = config
|
||||
logger.debug("DiscordBot: streaming configured, mode=%s", config.mode)
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
return self._is_running
|
||||
@@ -474,20 +494,95 @@ class DiscordBot(OutboundResilienceMixin, ChatCommandMixin, MessageHookMixin):
|
||||
|
||||
async def _send_agent_response():
|
||||
text_to_send = await self._debouncer.debounce(user_id, bot_message.text)
|
||||
channel_id = str(message.channel.id)
|
||||
|
||||
# Streaming path: progressively edit a placeholder in
|
||||
# place using the shared DraftStreamer, mirroring Telegram
|
||||
# (Issue #5415). Only engaged when streaming is configured.
|
||||
streaming_enabled = (
|
||||
self._streaming_config is not None
|
||||
and self._streaming_config.mode != StreamingMode.OFF
|
||||
)
|
||||
if streaming_enabled:
|
||||
streamer = DraftStreamer(
|
||||
adapter=self,
|
||||
channel_id=channel_id,
|
||||
config=self._streaming_config,
|
||||
platform="discord",
|
||||
)
|
||||
placeholder_message_id = await streamer.start()
|
||||
try:
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, text_to_send,
|
||||
chat_id=channel_id,
|
||||
user_name=str(getattr(message.author, "name", "")),
|
||||
message_id=str(message.id),
|
||||
account=getattr(self.config, "account", "default"),
|
||||
stream_callback=streamer.on_event,
|
||||
)
|
||||
send_result = self.fire_message_sending(
|
||||
channel_id, str(response),
|
||||
)
|
||||
if send_result["cancel"]:
|
||||
try:
|
||||
await self.delete_message(
|
||||
channel_id, placeholder_message_id
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
final_content = send_result["content"]
|
||||
# A streamed answer longer than Discord's hard
|
||||
# cap cannot be delivered by editing the single
|
||||
# placeholder (both the edit and the send
|
||||
# fallback 400 on >2000 chars), which would drop
|
||||
# the whole reply. Fall back to the chunked
|
||||
# reply path so long answers are split and keep
|
||||
# the reply reference (Greptile #3 / #5).
|
||||
_cap = min(self.config.max_message_length, 2000)
|
||||
if len(final_content) > _cap:
|
||||
try:
|
||||
await self.delete_message(
|
||||
channel_id, placeholder_message_id
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
await self._send_long_message(
|
||||
message.channel, final_content,
|
||||
reference=message,
|
||||
)
|
||||
else:
|
||||
await streamer.finalize(final_content)
|
||||
self.fire_message_sent(channel_id, final_content)
|
||||
except Exception:
|
||||
try:
|
||||
await self.delete_message(
|
||||
channel_id, placeholder_message_id
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
if ack_ctx:
|
||||
await self._ack.done(
|
||||
ack_ctx, react_fn=_discord_react,
|
||||
unreact_fn=_discord_unreact,
|
||||
)
|
||||
return
|
||||
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, text_to_send,
|
||||
chat_id=str(message.channel.id),
|
||||
chat_id=channel_id,
|
||||
user_name=str(getattr(message.author, "name", "")),
|
||||
message_id=str(message.id),
|
||||
account=getattr(self.config, "account", "default"),
|
||||
)
|
||||
send_result = self.fire_message_sending(
|
||||
str(message.channel.id), str(response),
|
||||
channel_id, str(response),
|
||||
)
|
||||
if send_result["cancel"]:
|
||||
return
|
||||
await self._send_long_message(message.channel, send_result["content"], reference=message)
|
||||
self.fire_message_sent(str(message.channel.id), send_result["content"])
|
||||
self.fire_message_sent(channel_id, send_result["content"])
|
||||
# Done reaction - show completion
|
||||
if ack_ctx:
|
||||
await self._ack.done(ack_ctx, react_fn=_discord_react, unreact_fn=_discord_unreact)
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -71,6 +71,12 @@ from ._commands import (
|
||||
build_command_access_policy,
|
||||
)
|
||||
from ._session import BotSessionManager
|
||||
from ._streaming import (
|
||||
StreamingConfig,
|
||||
StreamingMode,
|
||||
DraftStreamer,
|
||||
build_streaming_config,
|
||||
)
|
||||
from ._debounce import InboundDebouncer
|
||||
from ._ack import AckReactor
|
||||
from ._unknown_user import UnknownUserHandler, BotContext
|
||||
@@ -144,7 +150,12 @@ class SlackBot(OutboundResilienceMixin, ChatCommandMixin, MessageHookMixin):
|
||||
|
||||
# Initialize allow_silence from config
|
||||
self._allow_silence = getattr(self.config, 'allow_silence', False)
|
||||
|
||||
|
||||
# Initialize streaming config based on BotConfig (shared with Telegram /
|
||||
# Discord so ``streaming.mode: draft|auto`` is honoured identically
|
||||
# across channels that support in-place edits — Issue #5415).
|
||||
self._streaming_config = build_streaming_config(self.config)
|
||||
|
||||
self._is_running = False
|
||||
self._bot_user: Optional[BotUser] = None
|
||||
self._app = None
|
||||
@@ -186,6 +197,111 @@ class SlackBot(OutboundResilienceMixin, ChatCommandMixin, MessageHookMixin):
|
||||
# Audio capabilities
|
||||
self._stt_enabled: bool = False
|
||||
|
||||
def configure_streaming(self, config: StreamingConfig) -> None:
|
||||
"""Configure streaming reply mode.
|
||||
|
||||
Args:
|
||||
config: Streaming configuration. Set mode=StreamingMode.OFF to disable.
|
||||
"""
|
||||
self._streaming_config = config
|
||||
logger.debug("SlackBot: streaming configured, mode=%s", config.mode)
|
||||
|
||||
def _reply_needs_thread(self, thread_id: str) -> bool:
|
||||
"""Whether a reply must be posted into a thread rather than the root.
|
||||
|
||||
The shared ``DraftStreamer`` posts/edits its placeholder at the channel
|
||||
root (it only knows ``send_message(channel, text)``), so when a reply
|
||||
belongs in a thread — the inbound message was itself in a thread, or
|
||||
``reply_in_thread`` is configured — streaming in place would strand the
|
||||
answer in the main channel (Greptile #2 / #5417). In that case the
|
||||
caller falls back to the existing threaded send path.
|
||||
"""
|
||||
return bool(thread_id) or bool(getattr(self.config, "reply_in_thread", False))
|
||||
|
||||
async def _maybe_stream_reply(
|
||||
self,
|
||||
channel_id: str,
|
||||
user_id: str,
|
||||
text: str,
|
||||
*,
|
||||
thread_id: str = "",
|
||||
message_id: str = "",
|
||||
) -> Optional[str]:
|
||||
"""Progressively stream a reply via DraftStreamer when configured.
|
||||
|
||||
Returns the final (hook-processed) response text when the streaming
|
||||
path handled delivery, or ``None`` when streaming is disabled — or when
|
||||
the reply must be threaded (the channel-agnostic streamer can only post
|
||||
at the channel root) — so the caller falls back to its normal
|
||||
single-message, thread-aware, media-aware send (Issue #5415).
|
||||
"""
|
||||
streaming_enabled = (
|
||||
self._streaming_config is not None
|
||||
and self._streaming_config.mode != StreamingMode.OFF
|
||||
)
|
||||
if not streaming_enabled:
|
||||
return None
|
||||
|
||||
# Defer to the standard send path when the reply must land in a thread:
|
||||
# the shared streamer posts/edits at the channel root and would
|
||||
# otherwise strand a threaded answer in the main channel (Greptile #2).
|
||||
if self._reply_needs_thread(thread_id):
|
||||
return None
|
||||
|
||||
streamer = DraftStreamer(
|
||||
adapter=self,
|
||||
channel_id=channel_id,
|
||||
config=self._streaming_config,
|
||||
platform="slack",
|
||||
)
|
||||
placeholder_message_id = await streamer.start()
|
||||
try:
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, text,
|
||||
chat_id=channel_id,
|
||||
thread_id=thread_id,
|
||||
message_id=message_id,
|
||||
account=getattr(self.config, "account", "default"),
|
||||
stream_callback=streamer.on_event,
|
||||
)
|
||||
send_result = self.fire_message_sending(channel_id, str(response))
|
||||
if send_result["cancel"]:
|
||||
try:
|
||||
await self.delete_message(channel_id, placeholder_message_id)
|
||||
except Exception:
|
||||
pass
|
||||
return ""
|
||||
|
||||
# Extract any MEDIA: directives before finalizing so an audio
|
||||
# reply is uploaded via Slack's file path instead of being edited
|
||||
# into the placeholder as literal "MEDIA:/path" text (Greptile #4).
|
||||
parsed = split_media_from_output(send_result["content"])
|
||||
text_content = parsed["text"]
|
||||
media_urls = parsed.get("media_urls", [])
|
||||
|
||||
await streamer.finalize(text_content if text_content else send_result["content"])
|
||||
|
||||
for media_path in media_urls:
|
||||
if not os.path.exists(media_path):
|
||||
logger.warning(f"Media file not found: {media_path}")
|
||||
continue
|
||||
if is_audio_file(media_path) and self._client:
|
||||
try:
|
||||
await self._client.files_upload_v2(
|
||||
channel=channel_id,
|
||||
file=media_path,
|
||||
title=os.path.basename(media_path),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to upload audio to Slack: {e}")
|
||||
return send_result["content"]
|
||||
except Exception:
|
||||
try:
|
||||
await self.delete_message(channel_id, placeholder_message_id)
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
def enable_stt(self, enabled: bool = True) -> None:
|
||||
"""Enable STT for audio file transcription."""
|
||||
self._stt_enabled = enabled
|
||||
@@ -622,6 +738,26 @@ class SlackBot(OutboundResilienceMixin, ChatCommandMixin, MessageHookMixin):
|
||||
user_id = event.get("user", "unknown")
|
||||
logger.info(f"Message received: {text[:100]}...")
|
||||
text = await self._debouncer.debounce(user_id, text)
|
||||
|
||||
# Streaming path: progressively edit a placeholder in place
|
||||
# using the shared DraftStreamer, mirroring Telegram/Discord
|
||||
# (Issue #5415). Returns None when streaming is disabled.
|
||||
streamed = await self._maybe_stream_reply(
|
||||
str(channel_id) if channel_id else "",
|
||||
user_id, text,
|
||||
thread_id=event.get("thread_ts", "") or "",
|
||||
message_id=event.get("ts", ""),
|
||||
)
|
||||
if streamed is not None:
|
||||
if streamed:
|
||||
self.fire_message_sent(channel_id, streamed)
|
||||
if ack_ctx and self._client:
|
||||
await self._ack.done(
|
||||
ack_ctx, react_fn=_slack_react,
|
||||
unreact_fn=_slack_unreact,
|
||||
)
|
||||
return
|
||||
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, text,
|
||||
chat_id=str(channel_id) if channel_id else "",
|
||||
@@ -707,6 +843,21 @@ class SlackBot(OutboundResilienceMixin, ChatCommandMixin, MessageHookMixin):
|
||||
try:
|
||||
user_id = event.get("user", "unknown")
|
||||
logger.info(f"@mention received: {text[:100]}...")
|
||||
|
||||
# Streaming path: progressively edit a placeholder in place
|
||||
# using the shared DraftStreamer (Issue #5415). Returns None
|
||||
# when streaming is disabled so the plain path runs.
|
||||
streamed = await self._maybe_stream_reply(
|
||||
str(event.get("channel", "")),
|
||||
user_id, text,
|
||||
thread_id=event.get("thread_ts", "") or "",
|
||||
message_id=event.get("ts", ""),
|
||||
)
|
||||
if streamed is not None:
|
||||
if streamed:
|
||||
self.fire_message_sent(event.get("channel", ""), streamed)
|
||||
return
|
||||
|
||||
response = await self._session.chat(
|
||||
self._agent, user_id, text,
|
||||
chat_id=str(event.get("channel", "")),
|
||||
|
||||
Loaded 100 of 176 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user