agentpause 0.2.0__tar.gz → 0.3.0__tar.gz

{agentpause-0.2.0/src/agentpause.egg-info → agentpause-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentpause
-Version: 0.2.0
+Version: 0.3.0
 Summary: Predictive scheduling for autonomous LLM agents: suspend gracefully before rate limits, resume without losing work.
 Author: Andrea Lelio Ciampolillo
 License: MIT
@@ -40,11 +40,49 @@ plugin for self-hosted runtimes (llama.cpp, vLLM).
 | End-to-end A/B (thinking models, 4B/8B) | recovery **54×–93× faster**; total task time −19% |
 
 Reproduce them yourself with a free key: `python scripts/benchmark_groq.py`.
-
-> Status: **early alpha (0.1)**. Core components, the high-level
-> `PredictiveScheduler` API, the LiteLLM adapter (any provider), the LangGraph
-> adapter, runnable examples, and a full test suite are in place; the optional
-> KV-cache plugin is next.
+Live run (2026-07-08, Groq free tier, 12 steps, ~1k-token context,
+refill-aware chunked waiting, leveled windows):
+
+|  | reactive baseline | agentpause |
+|---|---|---|
+| 429 errors suffered | 7 | **0** |
+| steps redone | 7 | **0** |
+| tokens re-sent (waste) | 12,840 | **0** |
+| telemetry overhead (pings) | 0 | 259 |
+| wall-clock | 82 s | 86 s |
+
+Zero errors and zero waste at wall-clock parity: the sensor costs ~2% of
+what crashes waste (259 vs 12,840 tokens) — and waste is money on paid tiers.
+
+**Context slimming vs. answer quality** (`python scripts/quality_ab.py`,
+live 2026-07-10, Groq `llama-3.1-8b-instant`): six facts planted early in a
+verbose conversation, same final quiz under three histories.
+
+| condition | prompt chars | facts recalled |
+|---|---|---|
+| A — full history | 8,984 | **6/6** |
+| B — `compact()` (blind truncation) | 4,370 | **0/6** — and the model *invented* plausible replacements |
+| C — `summarize_with()` (one summary call) | 3,608 | **6/6** |
+
+A 2.2× larger run (`--big`, 19,145 chars — near the free tier's whole TPM
+window, the physical ceiling for a single call) repeated the pattern:
+A 6/6, B 0/6, C 6/6 at a third of the prompt.
+
+The instructive failure is B's, and it is *erratic*: in one run it answered
+confidently with fabricated values (fake codename, fake budget, fake city);
+in the larger run it honestly declined. You cannot know in advance which
+failure you get — and the hallucinating one is the dangerous one. Blind
+truncation is an emergency exit (§8.6 overflow, no LLM available), not a
+strategy; semantic summarization recalled everything at a fraction of the
+prompt, earning its one extra call.
+A 20-step live stress test (2026-07-10) closed the loop: at the §8.6 wall the
+scheduler suspended, compacted the checkpoint offline, resumed slim, and
+completed 20/20 steps with zero 429s.
+
+> Status: **0.3.0**. Core scheduler, direct + LiteLLM + Anthropic adapters,
+> LangGraph integration, multi-provider routing, multi-agent shared budgets,
+> feature-based cost/latency estimation, runnable examples, and a full test
+> suite are in place; the optional KV-cache plugin is next.
 
 ## Quick example
 
@@ -167,6 +205,54 @@ Resume the paused thread with `graph.invoke(Command(resume=True), config)`.
 On resume the guard re-reads telemetry fresh — never from the checkpoint.
 Install with `pip install -e ".[langgraph]"`; see `examples/langgraph_quickstart.py`.
 
+## Scaling up: many providers, many agents, smarter estimates (v0.3)
+
+The same predictive idea — read the budget first, act before the error —
+extends in three directions, and they compose:
+
+```python
+from agentpause import BudgetRouter, MultiAgentCoordinator, FeatureEstimator
+from agentpause.adapters.openai_compat import OpenAICompatAdapter
+from agentpause.adapters.anthropic import AnthropicAdapter
+
+# 1. Route each call to the provider with the most headroom (predictive
+#    fallback: switch BEFORE the 429, not after). Providers in cooldown
+#    after a real 429 are skipped until their window resets.
+router = BudgetRouter(
+    ("groq",   OpenAICompatAdapter.for_model("groq/llama-3.1-8b-instant")),
+    ("claude", AnthropicAdapter("claude-haiku-4-5")),
+)
+
+# 2. Share ONE rate-limit window across a fleet: every granted call
+#    reserves its predicted cost (estimate + k·σ) from the shared pool,
+#    so agents can't overcommit the window together. Contention is
+#    arbitrated by priority, then longest-waiting.
+coord = MultiAgentCoordinator(telemetry=router.budget)
+coord.register("researcher", priority=1)
+coord.register("summarizer")
+
+est = FeatureEstimator()                     # 3. learn cost from features,
+est.set_context(tool="web_search")           #    not just context size
+
+d = coord.request("researcher", estimated=est.estimate(1200),
+                  sigma=est.sigma(fallback_estimate=1200))
+if d.action == "continue":
+    reply, used = router.backend(messages)   # router picks the provider
+    coord.complete("researcher", actual_tokens=used)
+    est.record(1200, used)
+```
+
+`FeatureEstimator` is a drop-in for the default estimator
+(`PredictiveScheduler(estimator=FeatureEstimator())`): a dependency-free
+ridge regression over features you declare (tool, model, temperature, …)
+that also learns per-step **latency**, feeding the optional time budget
+(`PredictiveScheduler(time_budget_s=...)`) — if the predicted step can't
+finish before the deadline, the answer is checkpoint, never wait (time,
+unlike tokens, does not refill).
+
+Runnable demo without any key: `python examples/fleet_quickstart.py`
+(routing switch + predictive WAIT under a shared window, in 30 lines of loop).
+
 ## Why
 
 Current agent frameworks persist state *reactively*: they checkpoint after a step
@@ -178,16 +264,23 @@ from the exact step.
 This library is the engineering counterpart of the research preprint *"A
 Resource-Aware Predictive Scheduler for Autonomous LLM Agents"*.
 
-## Components (v0.1)
+## Components
 
 | Module | Role |
 |--------|------|
 | `PredictiveScheduler` / `Session` | the high-level API: `session()`, `should_suspend()`, `call()`, `checkpoint()` |
 | `Estimator` | predicts next-step token cost with a moving-average error correction (ε) and tracks σ |
+| `FeatureEstimator` | drop-in replacement that learns cost *and latency* from declared features (tool, model, …) via dependency-free ridge regression |
 | `should_checkpoint` / `RiskModel` | the suspension rule (`remaining < estimated + k·σ`) and a diagnostic risk score |
-| `Budget` / `decide` | three-dimensional telemetry (TPM, RPM, reset time) and the three-valued rule: continue / wait / checkpoint |
-| `StateStore` / `Checkpoint` | atomic logical checkpointing with idempotency keys — works on any provider |
+| `Budget` / `decide` | multi-dimensional telemetry (TPM, RPM, reset time, deadline) and the three-valued rule: continue / wait / checkpoint |
+| `StateStore` / `Checkpoint` | atomic logical checkpointing with idempotency keys — works on any provider; `compact()` / `summarize_with()` shrink a suspended checkpoint offline |
+| `BudgetRouter` | predictive multi-provider routing: reads every provider's budget first, routes to the most headroom, cools down 429'd providers |
+| `MultiAgentCoordinator` | one shared rate-limit window across many agents: granted calls reserve their predicted cost; `arbitrate()` resolves contention by priority + fairness |
+| `ToolQuota` | client-side sliding window for rate-limited tools that expose no headers |
+| `CircuitBreaker` / `FallbackBackend` | reactive safety nets: fail fast on a broken provider, try the next one in order |
 | `adapters.litellm.LiteLLMAdapter` | backend + telemetry for any LiteLLM-supported provider (headers → budget, stale reading → 1-token ping) |
+| `adapters.openai_compat.OpenAICompatAdapter` | direct HTTP adapter for OpenAI-compatible APIs (Groq, OpenAI, …) — no litellm dependency |
+| `adapters.anthropic.AnthropicAdapter` | direct adapter for the Anthropic Messages API, with `cache_control` prompt caching and measured `cache_read/write_tokens` |
 | `adapters.langgraph.AgentPauseGuard` | predictive gate for LangGraph nodes: `check()` interrupts the graph before the fatal call, `record()` trains the estimator |
 
 ## Install (from source, during development)
@@ -206,6 +299,11 @@ pytest
 - [x] Runnable quickstart example (no keys)
 - [x] LiteLLM adapter (works with any provider)
 - [x] LangGraph adapter (interrupt + checkpointer)
+- [x] Direct adapters (OpenAI-compatible, Anthropic with prompt caching)
+- [x] Predictive multi-provider routing (`BudgetRouter`)
+- [x] Shared budget across agents (`MultiAgentCoordinator`)
+- [x] Feature-based cost & latency estimator (`FeatureEstimator`)
+- [ ] CrewAI / AutoGen / LlamaIndex adapters
 - [ ] Optional KV-cache plugin for llama.cpp / vLLM (true warm start)
 
 ## License

{agentpause-0.2.0 → agentpause-0.3.0}/README.md

@@ -17,11 +17,49 @@ plugin for self-hosted runtimes (llama.cpp, vLLM).
 | End-to-end A/B (thinking models, 4B/8B) | recovery **54×–93× faster**; total task time −19% |
 
 Reproduce them yourself with a free key: `python scripts/benchmark_groq.py`.
-
-> Status: **early alpha (0.1)**. Core components, the high-level
-> `PredictiveScheduler` API, the LiteLLM adapter (any provider), the LangGraph
-> adapter, runnable examples, and a full test suite are in place; the optional
-> KV-cache plugin is next.
+Live run (2026-07-08, Groq free tier, 12 steps, ~1k-token context,
+refill-aware chunked waiting, leveled windows):
+
+|  | reactive baseline | agentpause |
+|---|---|---|
+| 429 errors suffered | 7 | **0** |
+| steps redone | 7 | **0** |
+| tokens re-sent (waste) | 12,840 | **0** |
+| telemetry overhead (pings) | 0 | 259 |
+| wall-clock | 82 s | 86 s |
+
+Zero errors and zero waste at wall-clock parity: the sensor costs ~2% of
+what crashes waste (259 vs 12,840 tokens) — and waste is money on paid tiers.
+
+**Context slimming vs. answer quality** (`python scripts/quality_ab.py`,
+live 2026-07-10, Groq `llama-3.1-8b-instant`): six facts planted early in a
+verbose conversation, same final quiz under three histories.
+
+| condition | prompt chars | facts recalled |
+|---|---|---|
+| A — full history | 8,984 | **6/6** |
+| B — `compact()` (blind truncation) | 4,370 | **0/6** — and the model *invented* plausible replacements |
+| C — `summarize_with()` (one summary call) | 3,608 | **6/6** |
+
+A 2.2× larger run (`--big`, 19,145 chars — near the free tier's whole TPM
+window, the physical ceiling for a single call) repeated the pattern:
+A 6/6, B 0/6, C 6/6 at a third of the prompt.
+
+The instructive failure is B's, and it is *erratic*: in one run it answered
+confidently with fabricated values (fake codename, fake budget, fake city);
+in the larger run it honestly declined. You cannot know in advance which
+failure you get — and the hallucinating one is the dangerous one. Blind
+truncation is an emergency exit (§8.6 overflow, no LLM available), not a
+strategy; semantic summarization recalled everything at a fraction of the
+prompt, earning its one extra call.
+A 20-step live stress test (2026-07-10) closed the loop: at the §8.6 wall the
+scheduler suspended, compacted the checkpoint offline, resumed slim, and
+completed 20/20 steps with zero 429s.
+
+> Status: **0.3.0**. Core scheduler, direct + LiteLLM + Anthropic adapters,
+> LangGraph integration, multi-provider routing, multi-agent shared budgets,
+> feature-based cost/latency estimation, runnable examples, and a full test
+> suite are in place; the optional KV-cache plugin is next.
 
 ## Quick example
 
@@ -144,6 +182,54 @@ Resume the paused thread with `graph.invoke(Command(resume=True), config)`.
 On resume the guard re-reads telemetry fresh — never from the checkpoint.
 Install with `pip install -e ".[langgraph]"`; see `examples/langgraph_quickstart.py`.
 
+## Scaling up: many providers, many agents, smarter estimates (v0.3)
+
+The same predictive idea — read the budget first, act before the error —
+extends in three directions, and they compose:
+
+```python
+from agentpause import BudgetRouter, MultiAgentCoordinator, FeatureEstimator
+from agentpause.adapters.openai_compat import OpenAICompatAdapter
+from agentpause.adapters.anthropic import AnthropicAdapter
+
+# 1. Route each call to the provider with the most headroom (predictive
+#    fallback: switch BEFORE the 429, not after). Providers in cooldown
+#    after a real 429 are skipped until their window resets.
+router = BudgetRouter(
+    ("groq",   OpenAICompatAdapter.for_model("groq/llama-3.1-8b-instant")),
+    ("claude", AnthropicAdapter("claude-haiku-4-5")),
+)
+
+# 2. Share ONE rate-limit window across a fleet: every granted call
+#    reserves its predicted cost (estimate + k·σ) from the shared pool,
+#    so agents can't overcommit the window together. Contention is
+#    arbitrated by priority, then longest-waiting.
+coord = MultiAgentCoordinator(telemetry=router.budget)
+coord.register("researcher", priority=1)
+coord.register("summarizer")
+
+est = FeatureEstimator()                     # 3. learn cost from features,
+est.set_context(tool="web_search")           #    not just context size
+
+d = coord.request("researcher", estimated=est.estimate(1200),
+                  sigma=est.sigma(fallback_estimate=1200))
+if d.action == "continue":
+    reply, used = router.backend(messages)   # router picks the provider
+    coord.complete("researcher", actual_tokens=used)
+    est.record(1200, used)
+```
+
+`FeatureEstimator` is a drop-in for the default estimator
+(`PredictiveScheduler(estimator=FeatureEstimator())`): a dependency-free
+ridge regression over features you declare (tool, model, temperature, …)
+that also learns per-step **latency**, feeding the optional time budget
+(`PredictiveScheduler(time_budget_s=...)`) — if the predicted step can't
+finish before the deadline, the answer is checkpoint, never wait (time,
+unlike tokens, does not refill).
+
+Runnable demo without any key: `python examples/fleet_quickstart.py`
+(routing switch + predictive WAIT under a shared window, in 30 lines of loop).
+
 ## Why
 
 Current agent frameworks persist state *reactively*: they checkpoint after a step
@@ -155,16 +241,23 @@ from the exact step.
 This library is the engineering counterpart of the research preprint *"A
 Resource-Aware Predictive Scheduler for Autonomous LLM Agents"*.
 
-## Components (v0.1)
+## Components
 
 | Module | Role |
 |--------|------|
 | `PredictiveScheduler` / `Session` | the high-level API: `session()`, `should_suspend()`, `call()`, `checkpoint()` |
 | `Estimator` | predicts next-step token cost with a moving-average error correction (ε) and tracks σ |
+| `FeatureEstimator` | drop-in replacement that learns cost *and latency* from declared features (tool, model, …) via dependency-free ridge regression |
 | `should_checkpoint` / `RiskModel` | the suspension rule (`remaining < estimated + k·σ`) and a diagnostic risk score |
-| `Budget` / `decide` | three-dimensional telemetry (TPM, RPM, reset time) and the three-valued rule: continue / wait / checkpoint |
-| `StateStore` / `Checkpoint` | atomic logical checkpointing with idempotency keys — works on any provider |
+| `Budget` / `decide` | multi-dimensional telemetry (TPM, RPM, reset time, deadline) and the three-valued rule: continue / wait / checkpoint |
+| `StateStore` / `Checkpoint` | atomic logical checkpointing with idempotency keys — works on any provider; `compact()` / `summarize_with()` shrink a suspended checkpoint offline |
+| `BudgetRouter` | predictive multi-provider routing: reads every provider's budget first, routes to the most headroom, cools down 429'd providers |
+| `MultiAgentCoordinator` | one shared rate-limit window across many agents: granted calls reserve their predicted cost; `arbitrate()` resolves contention by priority + fairness |
+| `ToolQuota` | client-side sliding window for rate-limited tools that expose no headers |
+| `CircuitBreaker` / `FallbackBackend` | reactive safety nets: fail fast on a broken provider, try the next one in order |
 | `adapters.litellm.LiteLLMAdapter` | backend + telemetry for any LiteLLM-supported provider (headers → budget, stale reading → 1-token ping) |
+| `adapters.openai_compat.OpenAICompatAdapter` | direct HTTP adapter for OpenAI-compatible APIs (Groq, OpenAI, …) — no litellm dependency |
+| `adapters.anthropic.AnthropicAdapter` | direct adapter for the Anthropic Messages API, with `cache_control` prompt caching and measured `cache_read/write_tokens` |
 | `adapters.langgraph.AgentPauseGuard` | predictive gate for LangGraph nodes: `check()` interrupts the graph before the fatal call, `record()` trains the estimator |
 
 ## Install (from source, during development)
@@ -183,6 +276,11 @@ pytest
 - [x] Runnable quickstart example (no keys)
 - [x] LiteLLM adapter (works with any provider)
 - [x] LangGraph adapter (interrupt + checkpointer)
+- [x] Direct adapters (OpenAI-compatible, Anthropic with prompt caching)
+- [x] Predictive multi-provider routing (`BudgetRouter`)
+- [x] Shared budget across agents (`MultiAgentCoordinator`)
+- [x] Feature-based cost & latency estimator (`FeatureEstimator`)
+- [ ] CrewAI / AutoGen / LlamaIndex adapters
 - [ ] Optional KV-cache plugin for llama.cpp / vLLM (true warm start)
 
 ## License

{agentpause-0.2.0 → agentpause-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "agentpause"
-version = "0.2.0"
+version = "0.3.0"
 description = "Predictive scheduling for autonomous LLM agents: suspend gracefully before rate limits, resume without losing work."
 readme = "README.md"
 requires-python = ">=3.9"

{agentpause-0.2.0 → agentpause-0.3.0}/src/agentpause/__init__.py

@@ -14,18 +14,24 @@ from .errors import (
 )
 from .breaker import CircuitBreaker, CircuitOpenError
 from .estimator import Estimator
+from .regression import FeatureEstimator
 from .fallback import FallbackBackend
+from .refill import RegimeDetector
 from .retry import RetryPolicy
 from .risk import Budget, Decision, RiskModel, decide, should_checkpoint
+from .router import BudgetRouter
+from .coordinator import MultiAgentCoordinator
 from .state import Checkpoint, StateStore
+from .tools import ToolQuota
 from .scheduler import PredictiveScheduler, Session
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 __all__ = [
     "PredictiveScheduler",
     "Session",
     "Estimator",
+    "FeatureEstimator",
     "RiskModel",
     "Budget",
     "Decision",
@@ -35,8 +41,12 @@ __all__ = [
     "StateStore",
     "RetryPolicy",
     "FallbackBackend",
+    "BudgetRouter",
+    "MultiAgentCoordinator",
     "CircuitBreaker",
     "CircuitOpenError",
+    "RegimeDetector",
+    "ToolQuota",
     "AgentPauseError",
     "RateLimitHit",
     "TelemetryError",

agentpause-0.3.0/src/agentpause/adapters/anthropic.py

@@ -0,0 +1,118 @@
+"""Direct adapter for Anthropic's Messages API (F9.6), cache-aware (F9.5).
+
+Anthropic is not OpenAI-compatible — different endpoint (``/v1/messages``),
+different auth (``x-api-key``), ``system`` as a top-level field, usage split
+into input/output tokens — so it gets its own thin adapter, built on the
+same telemetry machinery as :class:`OpenAICompatAdapter` (whose header
+parser already understands ``anthropic-ratelimit-*`` names and RFC 3339
+resets).
+
+Prompt caching — the cloud analog of the KV warm start
+------------------------------------------------------
+True KV-cache export is impossible on cloud APIs, but Anthropic caches
+byte-identical prompt prefixes server-side (cache reads cost ~10% of the
+input price). agentpause's resume already re-sends the prefix verbatim by
+design, so with ``cache_prompt=True`` (default) the adapter plants
+``cache_control`` breakpoints on the stable prefix — the system message and
+the last message before the fresh question — and MEASURES what came back
+from cache (``cache_read_tokens``). Two honest caveats: the cache TTL is
+minutes (helps waits and short suspensions, not tomorrow's resume), and
+cache reads discount price/latency, not the rate-limit token count.
+
+    from agentpause.adapters.anthropic import AnthropicAdapter
+
+    adapter = AnthropicAdapter("claude-haiku-4-5")
+    sched = PredictiveScheduler(backend=adapter.backend, telemetry=adapter.budget)
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+
+from .openai_compat import OpenAICompatAdapter
+
+__all__ = ["AnthropicAdapter"]
+
+ANTHROPIC_VERSION = "2023-06-01"
+
+
+class AnthropicAdapter(OpenAICompatAdapter):
+    """Backend + telemetry for Anthropic models, with prompt-cache breakpoints.
+
+    Args:
+        model: e.g. ``"claude-haiku-4-5"``.
+        api_key: literal key, or None to read ``ANTHROPIC_API_KEY``.
+        cache_prompt: plant ``cache_control`` breakpoints on the stable
+            prefix (system + last message before the fresh question).
+        default_max_tokens: the Messages API requires ``max_tokens``.
+        Everything else as in :class:`OpenAICompatAdapter`.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        api_key: Optional[str] = None,
+        cache_prompt: bool = True,
+        default_max_tokens: int = 1024,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            model,
+            base_url="https://api.anthropic.com/v1",
+            api_key=api_key,
+            api_key_env="ANTHROPIC_API_KEY",
+            **kwargs,
+        )
+        self._auth = {
+            "x-api-key": self._api_key,
+            "anthropic-version": ANTHROPIC_VERSION,
+            "content-type": "application/json",
+        }
+        self.cache_prompt = cache_prompt
+        self.default_max_tokens = default_max_tokens
+        # measured cache effect (F9.5): never assumed, always counted
+        self.cache_read_tokens = 0
+        self.cache_write_tokens = 0
+
+    # -- protocol translation ---------------------------------------------------
+
+    def _call(self, messages: List[Dict[str, str]], **extra: Any) -> Dict[str, Any]:
+        system: Optional[str] = None
+        msgs = messages
+        if msgs and msgs[0].get("role") == "system":
+            system = msgs[0].get("content") or ""
+            msgs = msgs[1:]
+        payload: Dict[str, Any] = {
+            "model": self.model,
+            "max_tokens": extra.pop("max_tokens", self.default_max_tokens),
+            "messages": [dict(m) for m in msgs],
+            **extra,
+        }
+        if system is not None:
+            block: Dict[str, Any] = {"type": "text", "text": system}
+            if self.cache_prompt:
+                block["cache_control"] = {"type": "ephemeral"}
+            payload["system"] = [block]
+        if self.cache_prompt and len(payload["messages"]) >= 2:
+            # the stable prefix ends just before the fresh question: closing
+            # the breakpoint there lets a resumed (byte-identical) history
+            # hit the server-side cache
+            closer = payload["messages"][-2]
+            content = closer.get("content")
+            if isinstance(content, str):
+                closer["content"] = [{"type": "text", "text": content,
+                                      "cache_control": {"type": "ephemeral"}}]
+        body = self._post_and_check(f"{self.base_url}/messages", payload)
+        usage = body.get("usage") or {}
+        self.cache_read_tokens += int(usage.get("cache_read_input_tokens") or 0)
+        self.cache_write_tokens += int(usage.get("cache_creation_input_tokens") or 0)
+        return body
+
+    def backend(self, messages: List[Dict[str, str]]) -> Tuple[str, int]:
+        body = self._call(messages, **self.request_kwargs)
+        self._prev_ping = self._sample_point()   # real call starts a sampling chain
+        text = "".join(b.get("text", "") for b in (body.get("content") or [])
+                       if b.get("type") == "text")
+        usage = body.get("usage") or {}
+        used = int((usage.get("input_tokens") or 0) + (usage.get("output_tokens") or 0))
+        return text, used

{agentpause-0.2.0 → agentpause-0.3.0}/src/agentpause/adapters/langgraph.py

@@ -97,6 +97,8 @@ class AgentPauseGuard:
         wait_threshold_s: float = 15.0,
         sleep_fn: Callable[[float], None] = time.sleep,
         async_sleep_fn: Optional[Callable[[float], Any]] = None,
+        chunk_s: float = 10.0,
+        while_waiting: Optional[Callable[[float], None]] = None,
         max_resumes: int = 100,
     ) -> None:
         self.telemetry = telemetry
@@ -107,6 +109,8 @@ class AgentPauseGuard:
         self.wait_threshold_s = wait_threshold_s
         self.sleep_fn = sleep_fn
         self.async_sleep_fn = async_sleep_fn
+        self.chunk_s = chunk_s
+        self.while_waiting = while_waiting
         self.max_resumes = max_resumes
 
     # -- internals -----------------------------------------------------------
@@ -144,8 +148,18 @@ class AgentPauseGuard:
             if d.action == "continue":
                 return
             if d.action == "wait":
-                # reset is imminent: sleeping beats a suspend/resume cycle
-                self.sleep_fn((budget.reset_seconds or 1.0) + 0.5)
+                # reset is imminent: sleeping beats a suspend/resume cycle.
+                # Sleep in chunks and re-read: if the bucket refills sooner,
+                # resume sooner (and the samples feed regime detection).
+                # Useful waits (F9.2): if the app registered while_waiting,
+                # hand it the time — indexing, memory compression, prompt
+                # prep are all work that needs no LLM.
+                wait_s = min(d.wait_seconds or budget.reset_seconds or 1.0,
+                             self.chunk_s) + 0.5
+                if self.while_waiting is not None:
+                    self.while_waiting(wait_s)
+                else:
+                    self.sleep_fn(wait_s)
                 continue
             # interrupt() raises on the first pass (the node aborts here);
             # on a resume pass it RETURNS, and the loop re-reads telemetry.
@@ -177,7 +191,8 @@ class AgentPauseGuard:
             if d.action == "continue":
                 return
             if d.action == "wait":
-                delay = (budget.reset_seconds or 1.0) + 0.5
+                wait_s = d.wait_seconds or budget.reset_seconds or 1.0
+                delay = min(wait_s, self.chunk_s) + 0.5
                 if self.async_sleep_fn is not None:
                     await self.async_sleep_fn(delay)
                 else:

{agentpause-0.2.0 → agentpause-0.3.0}/src/agentpause/adapters/litellm.py

@@ -71,11 +71,23 @@ _INSTALL_HINT = ("The LiteLLM adapter needs the 'litellm' package. "
                  "Install it with:  pip install agentpause[litellm]")
 
 
+def _enable_headers() -> None:
+    """Ask litellm to surface provider response headers.
+
+    Without this flag litellm does NOT populate
+    ``_hidden_params["additional_headers"]`` — telemetry would be blind.
+    (Found in the field on the very first real-provider run.)
+    """
+    import litellm
+    litellm.return_response_headers = True
+
+
 def _default_completion() -> Callable[..., Any]:
     try:
         from litellm import completion
     except ImportError as exc:  # pragma: no cover - depends on environment
         raise ImportError(_INSTALL_HINT) from exc
+    _enable_headers()
     return completion
 
 
@@ -84,6 +96,7 @@ def _default_acompletion() -> Callable[..., Any]:
         from litellm import acompletion
     except ImportError as exc:  # pragma: no cover - depends on environment
         raise ImportError(_INSTALL_HINT) from exc
+    _enable_headers()
     return acompletion
 
 
@@ -207,11 +220,23 @@ class LiteLLMAdapter:
         return Budget(
             remaining_tokens=await self.atelemetry(),
             remaining_requests=self._remaining_requests,
-            reset_seconds=self._reset_seconds,
+            reset_seconds=self._aged_reset(),
             remaining_input_tokens=self._remaining_input,
             remaining_output_tokens=self._remaining_output,
         )
 
+    def _aged_reset(self) -> Optional[float]:
+        """Reset countdown advanced by the age of the cached reading.
+
+        The header value was true at capture time; served from cache up to
+        ``max_age_s`` later, it overstates the wait. Subtract the elapsed
+        time (floor 0) so callers never wait on already-elapsed seconds.
+        """
+        if self._reset_seconds is None or self._read_at is None:
+            return self._reset_seconds
+        age = max(0.0, self._clock() - self._read_at)
+        return max(0.0, self._reset_seconds - age)
+
     def count_tokens(self, text: str) -> int:
         """Per-model token count via litellm, with a safe heuristic fallback.
 
@@ -239,13 +264,17 @@ class LiteLLMAdapter:
         return Budget(
             remaining_tokens=self.telemetry(),
             remaining_requests=self._remaining_requests,
-            reset_seconds=self._reset_seconds,
+            reset_seconds=self._aged_reset(),
             remaining_input_tokens=self._remaining_input,
             remaining_output_tokens=self._remaining_output,
         )
 
     # -- internals -------------------------------------------------------------
 
+    def invalidate(self) -> None:
+        """Mark the cached telemetry stale; the next read will ping."""
+        self._read_at = None
+
     def ping(self) -> None:
         """Issue a deliberately tiny call whose only purpose is fresh headers."""
         response = self._call(