Compare commits

..
Author SHA1 Message Date
praisonai-triage-agent[bot] 33aab37e80 Merge remote-tracking branch 'origin/main' into test/compute-release-completion 2026-10-07 10:13:14 +00:00
MervinPraison e283f35215 docs: auto-update api.md [skip ci] 2026-10-07 09:49:35 +00:00
MervinPraison 1b4ff4538b docs: auto-update documentation parity trackers [skip ci] 2026-10-07 09:46:43 +00:00
cb8cbaffbe feat: Ollama System One (Jev) decision models + praisonai-ts (#5401)
* feat: Ollama System One (Jev) decision models and triage routing

Add /v1/systemone client, ticket triage, CLI, praisonai-ts, optional typesafe-sdk extra, unit tests. Live-validated with tev1:0.8b.

* fix: harden System One decisions (security, routing, parity trackers)

- system_one: scope TYPESAFE_API_KEY to TYPESAFE_BASE_URL hosts only; validate
  base URL (reject non-http(s) and plain-HTTP to non-loopback); fix timeout
  sentinel so explicit 120s is honoured (None is the unset sentinel).
- routing: reject empty/non-choice route questions; trigger fallback when
  confidence is missing; preserve the prompt over extra_state in triage state.
- CLI: validate --route-question is a choice question; print JSON via builtin so
  global --json mode no longer swallows decision output.
- pyproject: include decisions extra in [all].
- TS: normalize bare OLLAMA_HOST with http scheme; use getChoice + own-key route
  check; await async function targets; export routing API from package root.
- Regenerate parity tracker artifacts (ts + rust) for new decision symbols.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: Mervin Praison <[email protected]>

* chore: regenerate TS parity tracker after rebase on main

---------

Co-authored-by: Dhivya-Bharathy <[email protected]>
Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-07 10:44:37 +01:00
praisonai-triage-agent[bot] acf83fc72d Merge remote-tracking branch 'origin/main' into test/compute-release-completion 2026-10-07 09:43:05 +00:00
352823ebd8 fix(session): persist explicitly cleared agent histories (#5543)
* fix(session): persist explicitly cleared agent histories

* fix(session): forward overwrite kwarg through SqliteSessionStore import

- SqliteSessionStore._save_imported_session now accepts and forwards the
  overwrite keyword to the base DefaultSessionStore, fixing a pre-existing
  TypeError regression surfaced after overwrite= was added to the base.
- Update delayed_save monkeypatch in test_sqlite_import_index to forward
  kwargs so the concurrent-write test matches the new signature.
- Retarget test_verified_replacement_succeeds_on_first_write at the
  merge_session_metadata_map locked-merge path (#5545) it now takes,
  preserving the single-write + peer-preservation guarantee.

Co-authored-by: Mervin Praison <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-07 07:33:13 +01:00
Wu Shuwen 75fb69c6ce fix(session): validate hierarchy fork message indices (#5530) 2026-10-07 07:33:05 +01:00
praisonai-triage-agent[bot]andMervin Praison d7538082e3 fix: atomic locked writes for MCP auth & todos, isolate sync hooks (#5658)
* fix: atomic locked writes for MCP auth & todos, isolate sync hooks (fixes #5657)

- utils/atomic_io: add update_json, a locked (in-process mutex + FileLock)
  read-modify-write over atomic_write_json; corrupt files are moved aside,
  never silently treated as empty.
- mcp/mcp_auth_storage: route mutations through update_json so concurrent
  OAuth flows no longer lose entries and a crash mid-write keeps the previous
  file; tokens are never world-readable (temp file is 0600); corrupt store is
  moved to .corrupt and logged.
- tools/todo_tools: serialise todo_add/todo_update through update_json and
  derive ids from max(id)+1, so parallel adds all persist with unique ids and
  a torn write no longer resets the plan.
- hooks/runner: run sync hooks in a dedicated pool (not the loop default
  executor), preserve contextvars, and fail closed on a later call while a
  previous invocation is still stuck, capping the leak at one thread per hook.

Co-authored-by: Mervin Praison <[email protected]>

* fix: restore _save_todos for MCP adapter; close quarantine & hook races

- todo_tools: re-add TodoTools._save_todos (atomic_write_json) — the MCP CLI
  adapter does its own read-modify-write via _load_todos/_save_todos; removing
  it broke test_cli_surface_defects (AttributeError).
- atomic_io: quarantine corrupt files under a unique <path>.<pid>.<ns>.corrupt
  sidecar and raise if the move fails (never overwrite recoverable bytes);
  add shared quarantine_corrupt_file helper.
- mcp_auth_storage._read: quarantine under the writer FileLock with a re-read,
  so a store another process just repaired is never moved aside.
- hooks/runner: reserve the single in-flight slot under an asyncio lock BEFORE
  submitting and clear it via a done-callback, so concurrent invocations can't
  both launch a worker (cap bypass) and a cancelled/timed-out hook still holds
  its slot.

Co-authored-by: Mervin Praison <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-07 07:32:56 +01:00
bde2501f6f fix: drain compute teardown before process exit (#5440)
* Drain compute teardown before the process bridge exits

* fix: synchronize GC-to-exit teardown handoff for compute release

Close the unsynchronised window where the GC finalizer submits the
teardown future before recording it in pending_release, letting a
concurrent exit callback return early and the bridge cancel the
in-flight shutdown. Guard both reclaim and reclaim_at_exit with a lock.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: Mervin Praison <[email protected]>

* Coordinate release ownership and future publication at exit

* fix: retain process-owned bridge for compute cleanup

* test: collect coverage from compute exit subprocesses

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-05 21:41:24 +01:00
praisonai-triage-agent[bot]andMervin Praison f02e58fee7 fix: harden ACP session store, atomic JSON writes, and scheduler PID safety (#5660)
* fix: harden ACP session store, atomic JSON writes, and scheduler PID safety (fixes #5659)

Three validated correctness gaps in src/praisonai/praisonai/:

1. ACP SessionStore: validate client-supplied session_id against a
   server-minted format before interpolating into a filesystem path
   (blocks arbitrary *.json read/delete via traversal) and write session
   files + the .last_session pointer atomically.

2. Non-atomic JSON writes: route JSONConversationStore and
   SchedulerStateManager saves through atomic temp-file + os.replace so a
   crash/Ctrl-C mid-write can't zero conversation history or scheduler
   state; sanitize session_id / name before filename interpolation
   (reuses the existing _atomic_write_json helper for scheduler state).

3. Scheduler lifecycle races: record process start_time and compare it in
   is_process_alive / stop_daemon so a recycled PID is never signalled
   (prevents killing an innocent process), and make generate_unique_name
   race-free via O_CREAT|O_EXCL.

Changes are backward-compatible (start_time optional) and reuse existing
in-tree helpers/patterns; no new user-facing params.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: MervinPraison <[email protected]>

* fix: close scheduler lifecycle + conversation-index gaps from review

- scheduler state: validate name in load_state/delete_state so a crafted
  name can't traverse out of the state dir to read/unlink an arbitrary
  *.json (greptile #2)
- schedule start: validate name BEFORE launching the daemon so an invalid
  name can't orphan a running daemon with no stoppable state file
  (greptile #1); use is_process_alive(pid, start_time) instead of a
  PID-only probe so a reused PID no longer falsely blocks a fresh start
  (greptile #5)
- stop_daemon/astop_daemon: re-verify PID ownership before SIGKILL so a
  PID reused during the SIGTERM wait is never force-killed (greptile #4)
- conversation json_store: reserve the "_sessions_index" id so a session
  can't overwrite the store's index file (greptile #6)

Co-authored-by: Mervin Praison <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: MervinPraison <[email protected]>
2026-10-05 21:35:49 +01:00
Wu Shuwen cfa1d131cc fix: retain indexed session routes when clearing turns (#5522)
Merged after rebase onto current main with minimal sqlite_store fix (stop de-indexing on clear_session).
2026-10-05 06:10:56 +01:00
Wu Shuwen 2135d03120 fix: isolate provider tool additions to each request (#5451) 2026-10-05 05:01:23 +01:00
praisonai-triage-agent[bot]andMervinPraison 1424ca016d fix: add resolved-config snapshot/export for gateway (fixes #5646) (#5648)
A running gateway's posture was split between the declared gateway.yaml and a
hidden per-host:port start-flags side file, so the YAML never described how the
gateway was actually running and no command rendered the resolved, merged,
redacted configuration. Add resolve_effective_config()/export_effective_config()
(wrapper) that merge declared YAML with the persisted CLI overrides into one
promotable document, and a `praisonai gateway config --resolved/--export` CLI.
Secrets are redacted by default via the existing diagnostics _SECRET_KEYS set
while non-secret values are preserved so the output stays a usable, commitable
config.

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: MervinPraison <[email protected]>
2026-10-05 05:01:05 +01:00
Wu Shuwen 3e7fd41807 fix(storage): match SQLite key prefixes literally (#5560) 2026-10-05 04:29:00 +01:00
Wu Shuwen ca332ceebe Return SQLiteBackend clear count from DELETE rowcount (#5573) 2026-10-05 04:28:23 +01:00
Wu Shuwen 295ebf33ad fix: persist portable SQLite transcript operations in the database (#5518)
* fix: use SQLite transcript rows for portable session operations

* fix: preserve transcript import no-clobber and cache contracts

* fix(session): filter lineage exports before loading transcript rows
2026-10-05 04:27:47 +01:00
Wu Shuwen 1f2c9c58e6 fix(session): honor hierarchy retention and preserve snapshot transcripts (#5534)
* fix(session): honor hierarchy retention policies

* fix(session): restore hierarchy imports without retention loss

* fix(session): restore independent snapshot transcripts after retention

* fix: preserve imported live history when taking snapshots

* fix: pool durable snapshot transcript records

* fix: reject pooled snapshots in positional readers

* fix: reject incomplete pooled snapshot transcripts

* fix: validate pooled snapshot reference types
2026-10-05 04:27:11 +01:00
Wu Shuwen e8657b709b fix(session): recover structured spill messages and skip invalid UTF-8 (#5551)
* fix(session): recover structured spill messages and skip invalid UTF-8

* fix: distinguish JSON scalar types during spill recovery

* fix: avoid recursive keys during spill recovery
2026-10-05 04:26:35 +01:00
MervinPraison c739ae2f00 docs: auto-update feature parity trackers [skip ci] 2026-10-05 03:26:04 +00:00
c5c1f90932 fix(session): recover failed import indexes and preserve peer updates (#5628)
* fix: forward overwrite through indexed session imports

* fix: retain refreshed portable import cache generations

* fix: preserve newer peer indexes after failed import refresh

* fix(session): invalidate import indexes after SQL refresh failures

* fix(run): guard profiled YAML path against legacy re-dispatch recursion

The profiled YAML runner (_run_from_file_profiled) delegated to the legacy
PraisonAI.run() without setting the PRAISONAI_IN_MODERN_RUN sentinel, so the
legacy dispatcher would re-enter the modern Typer run app and recurse. Set and
restore the sentinel around the profiled run just like _run_from_file, and add
a focused regression.

Co-authored-by: Mervin Praison <[email protected]>

* fix(session): retain atomic index refresh in portable import repair

* fix(session): recover import file identity after path stat failures

* fix(session): preserve a peer refreshed index of the same generation

* fix(session): invalidate stale indexes after plain peer replacements

* fix: roll back failed outermost index refresh commits

* fix: avoid import backfill lock inversion and retry incomplete passes

* fix: retain successful backfill upgrades across retries

* fix: check transcript generation before skipping backfill

* fix(session): allow callback queries during import backfill

* refactor(session): isolate profiled YAML routing in separate PR

* test(session): proxy in_transaction in atomic-refresh fake connection

Main's new WAL atomic-refresh test monkeypatches a fake connection, but
this PR's _index_session reads conn.in_transaction before the savepoint
(for the owns_transaction rollback contract). Proxy that attribute to the
real connection so the merged test exercises both intents.

Co-authored-by: Mervin Praison <[email protected]>

* Verify imported session routes before using indexed bindings

* Retry failed import indexes from current durable transcripts

* Preserve import recovery through nested cleanup and unavailable transcripts

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-05 04:25:20 +01:00
dajiaohuang 54ffe0d4d6 test: wait for durable training completion before restart 2026-10-04 22:17:59 +08:00
dajiaohuang 06d78d542c Isolate the non-repository fixture from ancestor Git metadata 2026-10-04 21:32:42 +08:00
dajiaohuang 65cd31feec test: scan hook event references once per source file 2026-10-04 21:31:25 +08:00
MervinPraison e49b254509 docs: auto-update documentation parity trackers [skip ci] 2026-10-03 18:05:16 +00:00
praisonai-triage-agent[bot] 55aff2bc1c Merge pull request #5656 from MervinPraison/claude/issue-5653-20261003-1713
refactor: extract shared lazy run-lock helper
2026-10-03 18:04:38 +00:00
praisonai-triage-agent[bot] c7be8d9047 Merge pull request #5652 from MervinPraison/claude/issue-5650-20261003-1616
fix: forward native tool-timeout to core (fixes #5650)
2026-10-03 17:56:00 +00:00
praisonai-triage-agent[bot]andMervin Praison 3bdfbb8d79 test: add regression tests for shared run-lock helper
Covers stable per-instance identity, both plain-attribute (AgentTeam) and
preset-None dataclass-field (Workflow) cases, and single-lock observation
under a concurrent first-access race (double-checked locking). Addresses
Greptile P2 finding that the shared concurrency path lacked a regression test.

Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 17:22:34 +00:00
praisonai-triage-agent[bot]andMervinPraison 22ee4fe369 refactor: extract shared lazy run-lock helper (fixes #5653)
Consolidate the duplicate double-checked-locking idiom (_RUN_LOCK_INIT_GUARD
+ _execution_lock) from agents/agents.py and workflows/workflows.py into a
single praisonaiagents/_run_lock.py helper. Behaviour-preserving: getattr-based
access works for both the dataclass-field (Workflow) and plain-attribute
(AgentTeam) cases. No public API or behaviour changes.

Co-authored-by: MervinPraison <[email protected]>
2026-10-03 17:17:49 +00:00
praisonai-triage-agent[bot]andMervin Praison 1ff45385b5 fix: preserve subsecond native tool-timeout when forwarding to core (#5650)
The native adapter rounded the resolved per-agent budget via int(round())
before building ToolConfig, collapsing subsecond budgets (e.g. 0.4 -> 0).
Core treats a non-positive timeout as "no timeout" and the native callable
is left unwrapped, so enforcement was silently lost. Core consumes
ToolConfig.timeout directly as seconds (float) on both sync and async paths,
so forward the budget unrounded. Adds a subsecond regression test.

Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 16:31:27 +00:00
praisonai-triage-agent[bot]andMervinPraison 6b6c3f5113 fix: forward native tool-timeout to core instead of wrapper-side pool (fixes #5650)
On the praisonai-native execution path, a plain callable already flows
through core's tool executor (agent/tool_execution.py), which enforces
the per-tool timeout with the same abandon-on-expiry semantics. The
wrapper no longer stands up a second ThreadPoolExecutor for those tools:
it leaves native callables unwrapped and forwards the resolved per-agent
tool_timeout to the core Agent via tool_config=ToolConfig(timeout=...).

Framework tool objects (CrewAI/AutoGen/LangChain BaseTool) never reach
core's executor, so they are still wrapped in place to preserve their
args_schema. Non-native frameworks keep the existing wrapper-side guard.
No new API surface; backward-compatible (flag defaults off).

Co-authored-by: MervinPraison <[email protected]>
2026-10-03 16:23:56 +00:00
MervinPraison 566e94edf2 docs: auto-update documentation parity trackers [skip ci] 2026-10-03 11:09:21 +00:00
praisonai-triage-agent[bot] 6a7d65202b Merge pull request #5426 from MervinPraison/claude/issue-5424-20261001-0941
fix: ignore invalid Retry-After numeric hints that poison backoff and scheduler holds
2026-10-03 11:08:48 +00:00
praisonai-triage-agent[bot] 6fdad4d964 Merge pull request #5411 from MervinPraison/claude/issue-5409-20261001-0902
fix: preserve StepResult stop flag and handler variables on workflow cache hits
2026-10-03 11:08:45 +00:00
praisonai-triage-agent[bot] 7886e17f26 Merge pull request #5558 from MervinPraison/claude/issue-5556-20261002-0141
fix: skip invalid UTF-8 transcripts during session search
2026-10-03 10:38:01 +00:00
praisonai-triage-agent[bot] 27c6045e50 Merge pull request #5579 from MervinPraison/claude/issue-5577-20261002-0304
fix: recover from invalid UTF-8 in JSON stores
2026-10-03 10:37:51 +00:00
praisonai-triage-agent[bot] d38b3cbcf6 Merge pull request #5647 from MervinPraison/claude/issue-5645-20261003-0914
fix: make AgentRegistry inbound routing thread-safe
2026-10-03 10:37:47 +00:00
praisonai-triage-agent[bot]andMervin Praison 921608ebe0 test: add shared-key contention regression test for AgentRegistry lock
Address greptile P2: the existing concurrent test used a distinct number per
worker (no shared-key contention) and accepted the fallback, so it could pass
even with the lock removed. Add a reader/writer test where writers churn an
overlapping key set while many readers snapshot concurrently — this surfaces
the exact race (RuntimeError: dictionary changed size during iteration / torn
resolve) that the per-instance lock + __iter__ snapshot prevent.

Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 09:21:38 +00:00
praisonai-triage-agent[bot]andMervinPraison 58ee5bc269 fix: make AgentRegistry inbound routing thread-safe (fixes #5645)
Guard the phone-number→agent routing table with a per-instance
threading.Lock so concurrent inbound turns cannot race on the shared
dict, matching the lock pattern every sibling gateway store already uses.
__iter__ now snapshots over a copy to avoid torn reads under mutation.

Co-authored-by: MervinPraison <[email protected]>
2026-10-03 09:16:10 +00:00
MervinPraison 8a5f427208 docs: auto-update documentation parity trackers [skip ci] 2026-10-03 08:56:37 +00:00
praisonai-triage-agent[bot]andMervinPraison 945d3ea3b2 fix: async-safe sync hooks/compaction + bounded session cache (#5641)
* fix: async-safe sync hooks/compaction + bounded session cache (fixes #5640)

- hooks/runner.execute_sync: route through the canonical
  run_coroutine_from_any_context bridge instead of raising inside a
  running event loop, so blocking policy/guardrail hooks still run (and
  can deny) from FastAPI/aiohttp/Jupyter/bot callbacks.
- compaction/compactor.compact: route LLM_SUMMARIZE through the same
  bridge so the user-configured summariser runs in both sync and
  async-host contexts; naive fallback is kept only for real failures and
  is now logged.
- session/store.DefaultSessionStore: bound the in-memory cache with an
  LRU (default 256, cache_maxsize=0 disables) so a long-lived
  multi-tenant process no longer retains every session in RAM.

Co-authored-by: MervinPraison <[email protected]>

* fix: bound hierarchical extended cache + add bridge timeout backstops

- hooks/runner.execute_sync: replace timeout=None with a bounded bridge
  timeout (default_timeout per hook + headroom) so a stuck hook can never
  pin the caller's thread indefinitely (Greptile P1).
- compaction/compactor.compact: drop timeout=None on the LLM summariser
  bridge call so a hung summariser falls back to the naive summary via the
  existing except handler instead of blocking forever (Greptile P1).
- session/hierarchy.HierarchicalSessionStore: bound _extended_cache with the
  same LRU policy/maxsize as the base _cache and prune _cache_mtimes on
  eviction via a new on_evict hook, so the subclass no longer retains every
  loaded session in RAM (Greptile P2).
- tests: assert extended cache stays bounded, reloads after eviction, and
  stays unbounded when cache_maxsize=0.

Co-authored-by: MervinPraison <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: MervinPraison <[email protected]>
2026-10-03 09:55:54 +01:00
praisonai-triage-agent[bot]andMervin Praison 2131cae365 fix: accumulate retried handler variables and isolate mutable cache values
Addresses two correctness gaps in the linear workflow step cache (PR #5411):

- P1 (retried variables): accumulate handler variables across guardrail
  retries instead of replacing, so a cache hit replays the same union of
  variables the cold run left in scope.
- P2 (mutable corruption): deep-copy cached variable values on set and on
  replay (falling back to the live reference when deepcopy fails), so a
  downstream step mutating a replayed object cannot corrupt what later hits see.

Adds 2 regression tests (rejected-retry variable union, mutable isolation).

Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 08:13:49 +00:00
praisonai-triage-agent[bot]anddajiaohuang e86069943b fix: preserve StepResult stop flag and handler variables on workflow cache hits (fixes #5409)
The linear executor cached only output and the generated output variable,
and on a hit appended only {step, output} then continued unconditionally.
This dropped stop_workflow, handler-supplied variables, status, and retries.

Now the cache snapshots the full step delta (stop signal, handler variables,
status, retries, and -- unless the step stopped first -- the output variable)
and the hit path replays it exactly, so a cached run is indistinguishable
from re-executing.

Co-authored-by: dajiaohuang <[email protected]>
2026-10-03 08:11:46 +00:00
praisonai-triage-agent[bot]andMervin Praison f59b6e9404 fix: preserve valid message Retry-After hints while rejecting negatives (#5424)
The sign-aware message patterns from the prior fix introduced two regressions
flagged by Greptile (P1):
- Valid delays went unmatched: the "(-?\d+)...second" pattern required the
  number adjacent to "second", dropping windows like "retry in 30 more seconds".
- Later valid hints were skipped: re.search saw only the first match, so an
  echoed negative ("retry after -5 seconds; retry after 30 seconds") shadowed a
  valid later 30s hint and returned None.

Fix: allow intervening word tokens before "second", capture fractional seconds
inside the group, and scan every match per pattern (re.finditer) so an earlier
invalid hint no longer masks a later usable one. Negatives/nan/inf remain
rejected via _usable_retry_delay. Added regression tests for all three cases.

Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 08:10:03 +00:00
praisonai-triage-agent[bot]andMervin Praison 859edd76de fix: reject echoed negative Retry-After in message fallback (#5424)
The message-fallback parser captured only digits, so a provider that echoed a
negative Retry-After in its error text (e.g. "retry after -3600 seconds") was
read as a positive delay and could park a scheduler hold on an invalid reset
window. Capture the optional leading sign and route the result through
_usable_retry_delay so negatives are rejected consistently with the structured
header/attribute sites, falling through to None.

Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 08:08:33 +00:00
praisonai-triage-agent[bot]anddajiaohuang fce6751700 fix: ignore invalid Retry-After numeric hints to protect backoff/scheduler (fixes #5424)
extract_retry_after() accepted non-finite (nan/inf) and negative
structured numeric hints, which poisoned classify_llm_error() backoff
and quota_hold_from_failure() scheduler holds. Validate structured
header and attribute hints are finite and non-negative; invalid values
are ignored and fall through the existing fallback order (attribute ->
message -> None). Preserves cap, HTTP-date handling and zero/fractional
delays.

Co-authored-by: dajiaohuang <[email protected]>
2026-10-03 08:08:22 +00:00
praisonai-triage-agent[bot]andMervin Praison c4471bd58d test: deterministically cover LIKE candidate branch for invalid-UTF-8 skip
Patch _init_schema to report FTS5 unavailable in one loop iteration so the
LIKE fallback branch in _candidate_ids is exercised even on SQLite builds that
ship FTS5, and assert _fts_available to prove the intended branch runs
(addresses greptile finding on #5558 / Issue #5556).

Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 08:05:57 +00:00
praisonai-triage-agent[bot]anddajiaohuang eeffecf339 fix: skip invalid UTF-8 transcripts during session search (fixes #5556)
DefaultSessionStore.search and SqliteSessionStore.search skipped malformed
JSON and file I/O errors but UnicodeDecodeError (a ValueError subclass, not
OSError) escaped the two transcript-reading boundaries, aborting search when
any candidate contained invalid UTF-8. Add UnicodeDecodeError to both except
tuples so unreadable transcripts are skipped, consistent with existing
malformed-JSON handling. Covers the direct JSON scan and indexed candidate
loading (FTS5 and LIKE branches).

Co-authored-by: dajiaohuang <[email protected]>
2026-10-03 08:05:57 +00:00
praisonai-triage-agent[bot]andMervin Praison 1bed1178af fix: recover from invalid UTF-8 in FileBackend + add regression tests
- FileBackend.load() now catches UnicodeDecodeError (returns None), matching
  the direct-file path and closing the backend-route gap noted in review.
- Add TestInvalidUtf8Recovery covering sync (locked/unlocked), FileBackend,
  BaseJSONStore-over-FileBackend, and async load_async() paths; all assert
  recovery to defaults while preserving the original bytes on disk.

Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 08:05:47 +00:00
MervinPraison 595addffdb docs: auto-update feature parity trackers [skip ci] 2026-10-03 07:48:08 +00:00
MervinPraison 0b65ff7ed5 docs: auto-update documentation parity trackers [skip ci] 2026-10-03 07:47:25 +00:00
Wu Shuwen bf15ea50f2 fix: surface hierarchy persistence failures before returning IDs (#5528)
* Report hierarchy persistence failures before returning identifiers

* Reject self-parent session creation before persistence
2026-10-03 08:44:09 +01:00
40eae2d0e3 fix(storage): preserve unreadable file-backed JSON after fallback (#5578)
* fix: recover default JSON stores from invalid UTF-8

* test(storage): exercise async UTF-8 recovery without requiring aiofiles

Replace importorskip with a fixture that injects a minimal real-decode
aiofiles stub when the optional package is absent, so the async
UnicodeDecodeError/JSONDecodeError recovery branch is always covered in
CI instead of silently skipped.

Co-authored-by: Mervin Praison <[email protected]>

* test: install real aiofiles for JSON store unit coverage

* fix(storage): refuse default saves over unreadable JSON files

* fix: preserve explicit JSON store clearing after read failures

* Handle invalid UTF-8 in FileBackend loads

* fix(storage): protect unreadable records from default saves

* test(storage): align main's invalid-UTF-8 test with unreadable-record guard

Resolve semantic rebase conflict: main added test_file_backend_load_invalid_utf8_returns_none
which re-saved the same key after a failed load; the PR's guard now refuses that
overwrite. Use distinct keys per iteration to preserve the test's full intent
(return None, warn, preserve bytes, neighbor loads) under the new guard.

Co-authored-by: Mervin Praison <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 08:43:37 +01:00
Wu Shuwen 99b1242319 fix: preserve workflow step state on cache hits (#5410)
* fix: preserve workflow step state on cache hits

* fix: keep uncacheable workflow state out of shared entries

* fix(workflows): snapshot custom cache payloads at each boundary
2026-10-03 08:43:08 +01:00
0a483ed08f fix: reject invalid Retry-After delay hints (#5425)
* fix: reject invalid retry-after delay hints

* fix: reject negative Retry-After hints in message fallback

Message-fallback patterns captured only unsigned digits, so a negative
hint like "-3600 seconds" was read as a positive delay despite the new
header/attribute validation. Capture the full signed (and fractional)
token so negatives are rejected by the finite/non-negative check and
fall through, preventing an invalid scheduler hold.

Co-authored-by: Mervin Praison <[email protected]>

* Parse complete numeric tokens in retry delay messages

* fix: reuse retry hint parsing in LLM retry loops

* fix: validate retry duration field and sentence boundaries

* fix: skip invalid retry hints before later valid durations

* fix(llm): normalize negative retry delay caps

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 08:43:05 +01:00
Wu Shuwen 3dc6a994bd fix: preserve Responses refusals and terminal outcomes (#5472)
* fix: preserve Responses refusal signals and terminal outcomes

* fix: clear response refusal before tool continuation

* test: isolate Responses refusal endpoints

* fix: retain incomplete outcomes alongside Responses refusals

* fix: isolate concurrent LLM terminal outcomes

* fix: observe async outcomes after a synchronous turn

* fix: preserve outer outcomes across nested tool requests

* fix: initialize outcome scopes for constructor-free callers
2026-10-03 08:43:01 +01:00
38cf60970e fix(session): preserve recall and reject incomplete transcript exports (#5557)
* fix(session): skip invalid UTF-8 transcript search candidates

* fix(session): isolate invalid UTF-8 in directory reads and exports

* fix(session): log skipped unreadable transcripts during search

Surface invalid-UTF-8/corrupt transcript skips via a debug log in both
the JSON scan and indexed candidate loader so a corrupt session is
observable instead of an indistinguishable silent miss. Search stays
read-only (no quarantine/rewrite), preserving corrupt bytes.

Co-authored-by: Mervin Praison <[email protected]>

* fix(session): page recall candidates and report corrupt scan entries

* fix(session): bound readable recall candidates with a stable cursor

* Preserve search lineages and release index lock before transcript work

* Stream search IDs from isolated SQLite readers

* Spool matching search IDs before transcript work

* Enumerate WAL search IDs on a separate reader connection

* fix: atomically refresh session indexes and reject incomplete exports

* fix(session): report atomic index refresh outcomes

* fix(session): score indexed continuations without retaining payloads

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 08:41:50 +01:00
b549d3cac6 fix(session): validate hierarchy cache against loaded JSON content (#5536)
* fix(session): validate hierarchy caches against loaded JSON bytes

* fix: retain hierarchy cache across contention and local writes

* fix(session): forward overwrite kwarg in SqliteSessionStore import

SqliteSessionStore._save_imported_session did not accept the overwrite
keyword that import_sessions forwards to the base implementation, causing
TypeError on every SQLite import. Forward the kwarg and update the stale
test monkeypatch wrapper to match the base signature.

Co-authored-by: Mervin Praison <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-03 08:41:46 +01:00
Wu Shuwen 2216795fd0 fix: preserve base session fields in hierarchy deserialization (#5526)
* Preserve the full base session schema in hierarchy conversions

* Verify archives survive hierarchy reload and mutation
2026-10-03 08:41:43 +01:00
dajiaohuang 1927a1f714 Merge main to include fork conflict-check handling 2026-10-03 06:16:10 +08:00
701bfba56d fix(ci): tolerate 403 clearing conflict-pending on fork PRs (#5638)
* fix(ci): tolerate 403 when clearing conflict-pending label on fork PRs

Fork pull requests can deny label removal; swallow 403/404 so detect-and-trigger does not fail the required check.

Co-authored-by: Cursor <[email protected]>

* fix(ci): only tolerate 403/404 clearing conflict label on fork PRs

Narrow the swallow to genuine forks (headRepo != baseRepo) so upstream
permission failures surface instead of silently leaving the merge-blocking
claude-conflict-pending label in place.

Co-authored-by: Mervin Praison <[email protected]>

---------

Co-authored-by: Cursor <[email protected]>
Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
2026-10-02 18:38:40 +01:00
b981fcf54a fix: default gateway agent model to detected provider, not OpenAI (#5611)
* fix: default gateway agent model to detected provider, not OpenAI (fixes #5609)

The gateway agent config schema defaulted `model` to the hardcoded
`gpt-4o-mini`, and onboarding wrote that literal into every generated
config, so an operator with only a non-OpenAI credential was handed a
config that could not work.

- AgentConfigSchema.model now defaults to None; new resolved_model()
  (and "auto") resolve to the detected provider's model via the existing
  default_model_for_available_provider() helper (reached through the lazy
  _code_bridge), with gpt-4o-mini only as last-resort fallback.
- Gateway server, bot CLI and the startup preflight route an unset/"auto"
  model through the same detection instead of assuming OpenAI.
- Onboarding YAML leads with the detected model, and the failure hint
  names the detected provider's API key env var instead of OPENAI_API_KEY.

An explicit model: (or --model) always wins; detection only fills the gap.

Co-authored-by: Mervin Praison <[email protected]>

* chore: regenerate gateway.schema.json after rebase onto #5416

Co-authored-by: Cursor <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: Mervin Praison <[email protected]>
Co-authored-by: MervinPraison <[email protected]>
Co-authored-by: Cursor <[email protected]>
2026-10-02 18:29:40 +01:00
praisonai-triage-agent[bot] 3983e45d08 Merge pull request #5417 from MervinPraison/claude/issue-5415-20261001-0914
fix: wire progressive streaming replies for Discord & Slack
2026-10-02 17:29:26 +00:00
praisonai-triage-agent[bot] 88e858adb9 Release v4.7.12 2026-10-02 16:34:56 +00:00
MervinPraison f9ed8ff862 docs: auto-update documentation parity trackers [skip ci] 2026-10-02 16:33:53 +00:00
Wu Shuwen dadac293e3 fix: surface Responses failures before accepting partial output (#5477)
* fix: surface Responses failures before accepting partial output

* fix: close Responses streams on consumer exit

* fix: surface native failed Responses without fallback replay

* test: defer optional HTTPX import until stream test execution

* fix: preserve dictionary Responses output in failure controls

* fix: surface failed native Responses during streaming

* Preserve token usage from dictionary Responses
2026-10-02 17:33:11 +01:00
praisonai-triage-agent[bot] 1ead4131e8 Merge pull request #5626 from MervinPraison/claude/issue-5617-20261002-1135
fix: add CI perf-gate flags to agent benchmark (#5617)
2026-10-02 16:24:59 +00:00
praisonai-triage-agent[bot] 2432197239 Merge pull request #5636 from MervinPraison/claude/issue-5632-20261002-1333
fix: capture re-entrancy sentinel before guarded import in YAML runner
2026-10-02 16:17:09 +00:00
praisonai-triage-agent[bot] 82574988f8 Merge pull request #5630 from MervinPraison/claude/issue-5624-20261002-1138
perf(agents): skip plugin-manager scan in Agent init when no plugins registered
2026-10-02 16:16:55 +00:00
MervinPraison 698aef54b8 docs: auto-update documentation parity trackers [skip ci] 2026-10-02 15:59:46 +00:00
praisonai-triage-agent[bot]andMervin Praison 96e76b535a fix: publish gateway.yaml/bot.yaml JSON Schema for editor + CI validation (#5416)
* fix: publish gateway.yaml/bot.yaml JSON Schema for editor + CI validation (fixes #5414)

Bring the gateway/bot config surface to parity with agents.yaml:
- Add gateway_config_json_schema() deriving a draft-07 schema from the
  canonical Pydantic GatewayConfigSchema (single source of truth).
- Add GATEWAY_SCHEMA_URL / GATEWAY_SCHEMA_HEADER mirroring AGENTS_SCHEMA_*.
- Add `praisonai gateway schema [-o FILE]` CLI command.
- Scaffold the yaml-language-server header into onboarded configs and the
  shipped gateway.yaml sample.
- Publish gateway.schema.json artefact and include it as package-data.

Co-authored-by: MervinPraison <[email protected]>

* fix: align published gateway schema with runtime validation (#5414)

Post-process the Pydantic-derived gateway/bot JSON Schema so the published
artefact faithfully reflects what the gateway actually validates, closing the
author-time/runtime gaps flagged in review:

- Inline GatewayServerSchema/HookSchema (extra="forbid") into gateway/hooks so
  a misspelled server knob (drain_timout) or hook key is caught while authoring,
  not only at startup.
- owner_user_id (top-level + channel) accepts an integer, mirroring the runtime
  str() coercion of unquoted numeric Telegram/Discord ids.
- Reject unknown top-level keys so a typo (auto_enable_from_en) the runtime
  silently ignores is surfaced; config_version (stamped) is allowed and channels
  keep additionalProperties open for plugin-declared keys.
- Allow array env-ref channel fields (allowed_users/allowlist/blocklist) to be a
  single ${VAR} string so scaffolded configs validate pre-substitution.

Regenerate committed gateway.schema.json. Shipped sample + onboard scaffolds now
validate clean; all three findings fixed.

Co-authored-by: Mervin Praison <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: MervinPraison <[email protected]>
2026-10-02 16:56:14 +01:00
praisonai-triage-agent[bot]andMervin Praison fdbbcbf52d fix: opt-in per-reply runtime footer for gateway bots (#5437)
* fix: add opt-in per-reply runtime footer for gateway bots (fixes #5428)

Surface a privacy-safe footer (model · latency) on the final reply only,
reusing the already-collected model name and measuring turn latency at the
bot turn site. Off by default via the existing BotConfig.footer knob.

Co-authored-by: MervinPraison <[email protected]>

* fix: exclude runtime footer from voice/TTS and cover slash-command replies

Address reviewer findings on the opt-in gateway-bot footer:

- Voice replies no longer narrate the footer. The footer is now appended to
  the delivered TEXT only (inside _send_response_with_media via a footer_line
  param), never to the pre-footer text handed to TTS/media synthesis. This also
  makes the streaming and non-streaming voice paths behave identically.
- File-based / unknown slash-command agent replies now include the footer,
  routed through the same delivery call with a per-turn latency start.
- Add handler-level delivery tests (opt-in gating, model/latency capture, and
  footer-in-text-not-voice) alongside the existing formatting unit tests.

Scope intentionally kept minimal per AGENTS.md: gateway channel-config wiring
and non-Telegram channel adapters remain out of scope (the PR is a deliberate
Telegram-only opt-in; the shared BotConfig.footer flag is a no-op on channels
that have not wired it, matching other per-channel BotConfig knobs).

Co-authored-by: Mervin Praison <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: MervinPraison <[email protected]>
2026-10-02 16:52:55 +01:00
praisonai-triage-agent[bot]andMervin Praison 2c91e0cc4f fix: add --json to gateway status and report real version (#5368)
* fix: add --json to gateway status and report real version (fixes #5363)

- Add --json to `gateway status`, emitting a single machine-readable
  snapshot (version, protocol_version, uptime, sessions, clients,
  per-channel health, degraded owners, delivery back-pressure, config
  revision + drift) built from /health + /info + daemon status.
- Replace the hardcoded "1.0.0" version on GET /info and MCP serverInfo
  with the installed praisonai-bot version via a single helper.
- Add protocol_version (GATEWAY_PROTOCOL_VERSION) to /info.

Co-authored-by: MervinPraison <[email protected]>

* fix: load persisted gateway token for status --json /info + guard version regressions

- Resolve GATEWAY_AUTH_TOKEN from ~/.praisonai/.env when not exported so an
  authenticated gateway's /info answers and status --json reports version
  (greptile P1). Surface version_unavailable="info_unreachable" when /health
  is reachable but /info is not, so a missing version is explicit not silent.
- Strengthen MCP initialize test to assert serverInfo.version == installed
  version and != "1.0.0" (greptile P2), guarding the endpoint from regressing
  back to a hardcoded value.
- Add token-resolution + version-unavailable snapshot tests.

Co-authored-by: Mervin Praison <[email protected]>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: MervinPraison <[email protected]>
2026-10-02 16:49:42 +01:00
MervinPraison 27206ed5bd docs: auto-update documentation parity trackers [skip ci] 2026-10-02 15:40:59 +00:00
Wu Shuwen 1fca8e2706 fix(memory): enumerate SQLite records and preserve MCP user filters (#5479)
* fix: enumerate all SQLite memory records without a search cap

* fix: filter SQLite memory users before display limits

* fix: use compatible metadata parsing for SQLite user filters

* fix: keep null memory owners distinct from literal user IDs

* fix(memory): filter explicitly supplied falsey user identifiers

* fix(mcp): preserve explicit memory user filters
2026-10-02 16:23:10 +01:00
MervinPraison ccd5849008 docs: auto-update documentation parity trackers [skip ci] 2026-10-02 14:02:15 +00:00
praisonai-triage-agent[bot]andMervin Praison 6cf57513c9 ci: re-trigger unit suite after flaky native segfault (no code change)
The failed run crashed with exit code 139 (Fatal Python error:
Segmentation fault) inside an unrelated tool-retry test
(tests/unit/tools/test_async_tool_retry_parity.py), with background
threads from posthog/asyncio/heartbeat and C-extensions on the stack.
The test passes in isolation and the crash point does not touch this
PR's plugin-manager change (has_registered_plugins + merge gates), which
is a pure, allocation-free dict read. Empty commit to re-run the full
suite on identical code.

Co-authored-by: Mervin Praison <[email protected]>
2026-10-02 13:53:58 +00:00
Wu Shuwen 8641da4633 fix: retain incomplete Responses outcomes and streamed usage (#5475)
* fix: retain incomplete Responses outcomes and streamed usage

* fix: preserve Responses usage mappings and verify real metrics
2026-10-02 14:53:11 +01:00
praisonai-triage-agent[bot] f7b2ff1be0 Merge pull request #5438 from MervinPraison/claude/issue-5433-20261001-0955
fix: implement live LangSmith trace delivery (PAI-TS-002)
2026-10-02 13:40:37 +00:00
praisonai-triage-agent[bot]anddajiaohuang ebca47a8ff fix: capture re-entrancy sentinel before guarded import in YAML runner (fixes #5632)
The _run_from_file YAML runner assigned _prev_in_run inside the try block,
right after the `from praisonai_code.cli.main import PraisonAI` import. When
that import raised, the broad except re-raised typer.Exit(1) and the finally
block then read the never-assigned _prev_in_run, raising UnboundLocalError and
masking the original diagnostic. Capture the sentinel's prior value before the
guarded block so the finally restore is always safe.

Co-authored-by: dajiaohuang <[email protected]>
2026-10-02 13:35:53 +00:00
Mervin PraisonandCursor 4d604e7716 ci: extend optimised smoke job timeout to 16 minutes (#5634)
Smoke runs were being cancelled at 12 minutes under load; 16 minutes matches headroom used elsewhere without widening test scope.

Co-authored-by: Cursor <[email protected]>
2026-10-02 14:22:07 +01:00
praisonai-triage-agent[bot]anddajiaohuang ce458b2ba5 fix: forward overwrite flag to indexed imported-session save (fixes #5627) (#5631)
Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: dajiaohuang <[email protected]>
2026-10-02 14:19:42 +01:00
praisonai-triage-agent[bot]andMervin Praison 3e048ab457 fix(session): forward overwrite kwarg in SqliteSessionStore import override
The base DefaultSessionStore._save_imported_session now takes an
overwrite keyword (import no-clobber work) and import_sessions calls it
with overwrite=overwrite, but the SqliteSessionStore override and a test
monkeypatch still used the old positional-only signature, raising
"unexpected keyword argument 'overwrite'" and breaking portable imports.

- sqlite_store.py: match base signature and forward overwrite to super()
- test_sqlite_import_index.py: update delayed_save mock to accept/forward
  overwrite so the concurrent-write test exercises the real code path

Co-authored-by: Mervin Praison <[email protected]>
2026-10-02 12:46:59 +00:00
praisonai-triage-agent[bot] 0a7195dac0 Merge remote-tracking branch 'origin/main' into claude/issue-5617-20261002-1135 2026-10-02 12:44:44 +00:00
Wu Shuwen 8dd23bf1c1 fix: isolate in-memory metadata and read results from caller edits (#5483)
* fix: isolate in-memory metadata and read results from caller edits

* fix: copy memory enumeration outside the write lock
2026-10-02 13:44:06 +01:00
Wu Shuwen 7fb721f3f3 fix(session): distinguish current and completed handoff progress (#5539) 2026-10-02 13:43:55 +01:00
praisonai-triage-agent[bot]andMervin Praison 92add1ef85 test(plugins): cover uncreated-singleton fast path in has_registered_plugins
Adds an isolated regression test that patches the module-global
_default_manager to None and asserts has_registered_plugins() returns
False without constructing the singleton. Closes the greptile P2 gap;
protects the allocation-free hot-path optimisation from regression.

Co-authored-by: Mervin Praison <[email protected]>
2026-10-02 12:41:30 +00:00
praisonai-triage-agent[bot]andMervin Praison 151d74096c fix: keep --json output parseable with --save and honor thresholds iterations/warmup
- Route save_results/update_readme status messages through an injectable
  logger; under --json they go to stderr so stdout stays pure JSON when
  --json and --save are combined (greptile P1).
- Make thresholds.json iterations/warmup authoritative for both the
  measurement and the emitted JSON metadata so configuration is no longer
  silently ignored (greptile P2).

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: Mervin Praison <[email protected]>
2026-10-02 12:13:57 +00:00
praisonai-triage-agent[bot]andDhivya-Bharathy 42c410d326 perf(agents): skip plugin-manager scan in Agent init when no plugins registered (fixes #5624)
Add a cheap, allocation-free has_registered_plugins() fast-path so the two
unconditional plugin merge calls in Agent.__init__ (_merge_plugin_tools and
_merge_plugin_subsystems) early-return when no plugins are registered — the
common case. Removes manager construction, lock acquisition and list scans
from the hot init path.

Tool-schema building was already lazy (first run()/chat, cached), so this
targets the remaining per-init overhead profiling surfaced. Benchmark (500
iters, warmed): min ~121us -> ~100us, mean ~135us -> ~111us; function calls
170k -> 134k. No behavioral change: plugins still merge when registered.

Co-authored-by: Dhivya-Bharathy <[email protected]>
2026-10-02 11:46:35 +00:00
praisonai-triage-agent[bot]andDhivya-Bharathy e9a5538a9c fix: add CI perf-gate flags to agent benchmark (fixes #5617)
Add --skip-crewai, --json, and --fail-over-threshold flags to
benchmarks/simple_benchmark.py plus a thresholds.json budget file so CI
can baseline and gate agent instantiation performance regressions.

- --skip-crewai avoids the optional-framework import hang
- --json emits machine-readable results for baseline artifacts
- --fail-over-threshold exits non-zero when PraisonAI init exceeds budget

Co-authored-by: Dhivya-Bharathy <[email protected]>
2026-10-02 11:37:21 +00:00
dajiaohuang c5f008f960 Merge main and preserve event-driven compute release tests 2026-10-02 19:05:25 +08:00
dajiaohuang 9081c77040 Merge main and retain measured compaction and release completion tests 2026-10-02 12:32:24 +08:00
praisonai-triage-agent[bot]andMervin Praison 07ac1d0547 fix(observability): harden LangSmith adapter per review (#5438)
- Surface failed run deliveries in flush() instead of resolving silently
- Reject plaintext http endpoints to non-loopback hosts (API key leak)
- Export span events into the run payload (extra.events)
- Evict completed runs/ordering state to bound memory in long processes
- Recognise legacy LANGCHAIN_API_KEY in CLI readiness via altEnvKeys
- Tests save/restore credential env; add coverage for all of the above

Co-authored-by: Mervin Praison <[email protected]>
2026-10-01 10:11:39 +00:00
praisonai-triage-agent[bot]andMervin Praison cf0043c585 fix(test): make compaction budget tests tokeniser-independent
test_unset_triggers_when_budget_approached assumed a flat 4 chars/token
ratio via "A" * 4 * n. Under tiktoken (available when its vocab is cached,
e.g. on CI) a degenerate all-A string collapses to ~8 chars/token, falling
under the resolved budget so needs_compaction returned False and the
assertion failed. Size payloads with the compactor's own estimate_tokens
so both the trigger and no-trigger assertions hold under the accurate and
heuristic token-counting paths. Production code unchanged.

Co-authored-by: Mervin Praison <[email protected]>
2026-10-01 10:08:15 +00:00
praisonai-triage-agent[bot]andDhivya-Bharathy c3f2676958 fix: implement live LangSmith trace delivery (fixes #5433)
Upgrade the LangSmith observability adapter from an honest stub to a real
deliverer using the existing axios dependency (no new deps). isEnabled is
derived from a live, credential-validated HTTP client, spans map to
LangSmith runs with tool/LLM child runs under the agent root, flush is
batched async, and secrets are read from env only with http(s)-only
endpoint validation. All other external vendors remain honest stubs.

Co-authored-by: Dhivya-Bharathy <[email protected]>
2026-10-01 10:00:14 +00:00
dajiaohuang c4def1b970 Wait for background compute release in lifecycle tests 2026-10-01 17:53:18 +08:00
praisonai-triage-agent[bot]andMervin Praison c2b84b84c0 fix: harden Discord/Slack streaming delivery guardrails (#5417)
Address Greptile review findings on the progressive-streaming wiring:
- Slack: defer to the thread-aware send path when a reply must be threaded
  (the root-only DraftStreamer would otherwise strand it in the channel),
  and extract MEDIA: directives so streamed audio is uploaded via Slack's
  file path instead of being edited in as literal text.
- Discord: when a streamed answer exceeds the 2000-char cap, delete the
  placeholder and fall back to the chunked _send_long_message path so long
  replies are split and keep their reply reference.
- Tests: add regression coverage for thread deferral, Slack media upload,
  and the Discord long-reply chunked fallback.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: Mervin Praison <[email protected]>
2026-10-01 09:34:36 +00:00
praisonai-triage-agent[bot]andMervinPraison 9bec7d3227 fix: wire progressive streaming replies for Discord & Slack (fixes #5415)
The channel-agnostic DraftStreamer was only instantiated by Telegram, so a
configured streaming.mode: draft|auto was silently inert on Discord and Slack
despite both supporting in-place edits. Reuse the existing engine in their
reply paths via a shared build_streaming_config factory.

Co-authored-by: MervinPraison <[email protected]>
2026-10-01 09:20:28 +00:00
176 changed files with 14760 additions and 1041 deletions

No files matched your search

+1 -1
View File
@@ -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:
+19 -7
View File
@@ -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' };
+1 -1
View File
@@ -42,7 +42,7 @@ jobs:
# ============================================
smoke:
runs-on: ubuntu-latest
timeout-minutes: 12
timeout-minutes: 16
steps:
- uses: actions/checkout@v4
with:
+2
View File
@@ -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
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=4.7.11" \
"praisonai>=4.7.12" \
"praisonai[chat]" \
"embedchain[github,youtube]"
+1 -1
View File
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=4.7.11" \
"praisonai>=4.7.12" \
"praisonai[ui]" \
"praisonai[chat]" \
"praisonai[realtime]" \
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=4.7.11" \
"praisonai>=4.7.12" \
"praisonai[ui]" \
"praisonai[crewai]"
+13 -13
View File
@@ -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
+364 -185
View File
@@ -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(
+7 -1
View File
@@ -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())
+10 -6
View File
@@ -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
View File
@@ -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
View File
@@ -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
+152 -1
View File
@@ -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