"""Bounded reconstruction of agent history from the action stream.
The store turns ``AgentAction`` events back into the runs, actions and attempts
they describe. It is development history, not an audit log: it is bounded, it
is in memory, and it is allowed to forget. What it is not allowed to do is
pretend.
Three properties follow from that, and most of this module exists to keep them:
* **Nothing is invented.** An attempt with no terminal event stays unfinished;
an action whose parent was never seen keeps ``parent_missing``; an action that
names no run is not filed under one. Absence is reported as absence.
* **Duplicates are free.** EDRI delivers an event once, but a store that
advanced its revision on a replay would hand a subscriber a change that never
happened. Re-applying one ``Event._id`` is a no-op.
* **Forgetting is visible.** Evicting attempts marks the surviving action
``truncated``; evicting a run invalidates it, so a client cannot read the
next generation's low revision as an old update.
Run identity is ``root_action_id``, never the upstream ``run_id``: one
``run_id`` may be handed to two concurrent invocations, and an action's parent
may be missing exactly when it is needed.
"""
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from enum import Enum
from typing import Any, Callable, Final, Mapping
from uuid import UUID, uuid4
from edri.integrations.pydantic_ai.observability import (
AgentAction,
AgentActionOutcome,
AgentActionPhase,
)
__all__ = [
"ActionState",
"AttemptState",
"ChangeKind",
"InspectorStore",
"RunState",
"StoreChange",
"StoreLimits",
"StoreStats",
]
[docs]
class ChangeKind(Enum):
"""What a consumer must do with a change record.
``UPDATED`` carries a revision a subscriber can compare against what it
holds. ``INVALIDATED`` says the run's retained state is gone -- evicted,
expired or cleared -- so comparing revisions is meaningless and the only
correct response is to stop trusting what is held.
"""
UPDATED = "updated"
INVALIDATED = "invalidated"
[docs]
@dataclass(frozen=True, slots=True)
class StoreChange:
"""One thing that happened to the store, for the manager to publish.
``generation_id`` travels with every record because ``revision`` means
nothing without it: a run re-materialized after eviction restarts at 1, and
a subscriber holding revision 40 of the previous generation must not read
that as stale.
"""
kind: ChangeKind
root_action_id: UUID
generation_id: UUID
revision: int
action_id: UUID | None = None
#: The run's upstream identity, carried on the record rather than looked
#: up: an invalidation is published after the run it names is gone.
run_id: str | None = None
agent: str | None = None
def _utc_now() -> datetime:
return datetime.now(timezone.utc)
#: Charged per retained value so that a run made of many tiny records still
#: costs something. The budget is an approximation of retained payload size,
#: not a measurement of Python object overhead -- what it has to bound is the
#: history an agent can make the inspector hold, and that is dominated by
#: captured prompts, results and errors.
_VALUE_BYTES: Final = 16
_ATTEMPT_BYTES: Final = 128
_ACTION_BYTES: Final = 192
_RUN_BYTES: Final = 256
_EVENT_ID_BYTES: Final = 32
_MAX_ESTIMATE_DEPTH: Final = 12
#: Outcomes an operator means by "show me what went wrong". Cancellation and
#: deferral end an attempt without it having failed, so neither is here.
_FAILURE_OUTCOMES: Final = frozenset(
{AgentActionOutcome.ERROR, AgentActionOutcome.TIMEOUT}
)
#: Sorting fallback for a run with no timestamps at all.
_EPOCH: Final = datetime(1970, 1, 1, tzinfo=timezone.utc)
def _estimate(value: Any, depth: int = 0) -> int:
"""Approximate retained size of one normalized payload.
Payloads arrive already bounded by the producer's normalization limits, so
this walk is over small built-in structures; the depth guard is for a
consumer replaying a stream this integration did not produce.
"""
if value is None or depth > _MAX_ESTIMATE_DEPTH:
return 0
if isinstance(value, str):
return len(value) + _VALUE_BYTES
if isinstance(value, (bool, int, float)):
return _VALUE_BYTES
if isinstance(value, Mapping):
return sum(
len(str(key)) + _VALUE_BYTES + _estimate(item, depth + 1)
for key, item in value.items()
)
if isinstance(value, (list, tuple, set, frozenset)):
return _VALUE_BYTES + sum(_estimate(item, depth + 1) for item in value)
return len(str(value)) + _VALUE_BYTES
[docs]
@dataclass(frozen=True, slots=True)
class StoreLimits:
"""What the store is allowed to keep.
Every limit applies to the *whole* retained state, auxiliary indexes
included: a bookkeeping dict that grows while ``max_actions`` is still
satisfied is the same leak, one level down.
"""
max_runs: int = 256
max_actions: int = 8_192
max_attempts_per_action: int = 32
#: A budget too small to hold even one run's root action degrades to
#: publish-and-drop: every action is applied, announced and then evicted,
#: which is bounded and correct but costs two update messages per action
#: instead of one. Size this above the largest root payload an agent
#: produces rather than to the smallest number that still runs.
max_bytes: int = 16 * 1_024 * 1_024
#: Age of a run's last update after which it is dropped. ``None`` keeps
#: runs until a count or byte limit displaces them.
ttl: timedelta | None = None
def __post_init__(self) -> None:
for name in ("max_runs", "max_actions", "max_attempts_per_action", "max_bytes"):
value = getattr(self, name)
if value < 1:
raise ValueError(f"{name} must be positive, got {value}")
if self.ttl is not None and self.ttl <= timedelta(0):
raise ValueError(f"ttl must be positive, got {self.ttl}")
[docs]
@dataclass(frozen=True, slots=True)
class StoreStats:
"""What the store currently holds, and what it has already lost.
``evicted_runs`` is retention's own count. It is not a guess about actions
the producer never published: the store cannot see those and does not
pretend to.
"""
runs: int
actions: int
attempts: int
bytes: int
incomplete_runs: int
#: Deduplication entries held on behalf of retained runs. Named because it
#: is the index that leaks invisibly: it is bounded by nothing a run or
#: action count would show.
dedupe_entries: int
evicted_runs: int
trimmed_runs: int
limits: StoreLimits
[docs]
@dataclass(slots=True)
class AttemptState:
"""One execution attempt of one logical action.
START and the terminal event are separate routed events and either may be
absent, so both halves are optional and their absence is a reported flag
rather than a reason to drop the attempt.
"""
attempt: int
started_at: datetime | None = None
ended_at: datetime | None = None
phase: AgentActionPhase | None = None
outcome: AgentActionOutcome | None = None
duration: float | None = None
usage: dict[str, Any] | None = None
request: Any = None
response_data: Any = None
error: dict[str, Any] | None = None
#: Approximate retained size. Zero until the store prices it, so that the
#: first measurement adds the whole record to the store's total.
bytes: int = 0
[docs]
def measure(self) -> int:
return _ATTEMPT_BYTES + sum(
_estimate(value)
for value in (self.request, self.response_data, self.error, self.usage)
)
@property
def start_missing(self) -> bool:
return self.started_at is None
@property
def terminal_missing(self) -> bool:
return self.outcome is None
[docs]
@dataclass(slots=True)
class ActionState:
"""One logical operation, with every attempt of it the store still holds."""
action_id: UUID
root_action_id: UUID
parent_action_id: UUID | None = None
namespace: str = ""
operation: str = ""
name: str | None = None
agent: str | None = None
model: str | None = None
provider: str | None = None
metadata: dict[str, Any] = field(default_factory=dict)
attempts: dict[int, AttemptState] = field(default_factory=dict)
#: Older attempts were evicted; the retry history shown is incomplete.
truncated: bool = False
#: Approximate retained size; see :attr:`AttemptState.bytes`.
bytes: int = 0
[docs]
def measure(self) -> int:
return (
_ACTION_BYTES
+ _estimate(self.name)
+ _estimate(self.metadata)
+ sum(attempt.bytes for attempt in self.attempts.values())
)
@property
def terminal_missing(self) -> bool:
return any(attempt.terminal_missing for attempt in self.attempts.values())
@property
def start_missing(self) -> bool:
return any(attempt.start_missing for attempt in self.attempts.values())
[docs]
@dataclass(slots=True)
class RunState:
"""One agent invocation: its identity, its actions and its revision."""
root_action_id: UUID
#: Minted per materialization. A run re-created after eviction is a new
#: generation, so its revisions are never comparable with the old ones.
generation_id: UUID
revision: int = 0
run_id: str | None = None
conversation_id: str | None = None
agent: str | None = None
created_at: datetime | None = None
updated_at: datetime | None = None
actions: dict[UUID, ActionState] = field(default_factory=dict)
truncated: bool = False
#: Every ``Event._id`` this run consumed, so eviction can drop the store's
#: deduplication history for it instead of leaking one entry per event.
event_ids: set[UUID] = field(default_factory=set)
#: Approximate retained size; see :attr:`AttemptState.bytes`.
bytes: int = 0
[docs]
def measure(self) -> int:
return (
_RUN_BYTES
+ len(self.event_ids) * _EVENT_ID_BYTES
+ sum(action.bytes for action in self.actions.values())
)
@property
def root(self) -> ActionState | None:
return self.actions.get(self.root_action_id)
@property
def completed(self) -> bool:
root = self.root
return root is not None and not root.terminal_missing and bool(root.attempts)
@property
def outcome(self) -> AgentActionOutcome | None:
root = self.root
if root is None or not root.attempts:
return None
return root.attempts[max(root.attempts)].outcome
@property
def failed(self) -> bool:
"""Something in this run ended badly.
A timed-out or errored action counts even when the run itself finished:
a retry that eventually succeeded still describes a run an operator
looking for failures wants to find. Cancellation and deferral are not
failures -- they are how a run ends when something asked it to.
"""
return any(
attempt.outcome in _FAILURE_OUTCOMES
for action in self.actions.values()
for attempt in action.attempts.values()
)
@property
def partial(self) -> bool:
"""Something this run should contain is not here.
Either the run itself was never fully observed, or one of its actions
was not -- a missing START, a missing terminal, a parent that never
arrived, or history retention removed.
"""
if self.root is None or self.truncated:
return True
return any(
action.truncated
or action.start_missing
or action.terminal_missing
or (
action.parent_action_id is not None
and action.parent_action_id not in self.actions
)
for action in self.actions.values()
)
[docs]
class InspectorStore:
"""Bounded in-memory history of observed agent runs.
Mutation is synchronous on the caller's thread: the manager owns the store
and ingests on its own event loop, so there is no background thread and no
lock. ``clock`` is injected because retention is a function of time and
tests must control it without sleeping.
"""
def __init__(
self,
*,
limits: StoreLimits | None = None,
clock: Callable[[], datetime] = _utc_now,
) -> None:
self._limits = limits if limits is not None else StoreLimits()
self._clock = clock
self._bytes = 0
self._actions = 0
self._evicted_runs = 0
self._trimmed_runs = 0
#: Invalidations produced outside an ``ingest`` call -- a TTL expiry
#: noticed while answering a query. A subscriber must still hear about
#: them, so they wait here for the manager to drain rather than being
#: dropped on the floor of whichever read happened to notice.
self._pending: list[StoreChange] = []
self._runs: dict[UUID, RunState] = {}
#: ``Event._id`` -> the run that consumed it. Owned by that run, so
#: evicting the run drops its dedupe history with it.
self._seen: dict[UUID, UUID] = {}
[docs]
def ingest(self, action: AgentAction) -> tuple[StoreChange, ...]:
"""Apply one observed action and report what changed.
Returns one ``UPDATED`` record when the run changed, nothing when the
event was already applied, and additionally an ``INVALIDATED`` record
for anything retention removed while making room.
"""
root_action_id = action.root_action_id
if root_action_id is None:
# Unattributed: the producer did not name a run, and the store does
# not guess one from timing or upstream ``run_id``.
return ()
if action._id in self._seen:
return ()
run = self._runs.get(root_action_id)
if run is None:
run = RunState(
root_action_id=root_action_id,
generation_id=uuid4(),
created_at=self._clock(),
)
self._runs[root_action_id] = run
self._seen[action._id] = root_action_id
run.event_ids.add(action._id)
self._apply(run, action)
run.updated_at = self._clock()
run.revision += 1
# Retention runs after the update, so the limits bound what is retained
# rather than what was offered: an action is applied and may then
# displace older history, including some of its own run's.
self._enforce(run)
if root_action_id in self._runs:
# Retention may have dropped the very run this action was applied
# to; there is nothing left to price.
self._measure(run)
# The update first: it describes the event just applied. Anything
# retention removed to make room for it follows.
return (
StoreChange(
kind=ChangeKind.UPDATED,
root_action_id=root_action_id,
generation_id=run.generation_id,
revision=run.revision,
action_id=action.action_id,
run_id=run.run_id,
agent=run.agent,
),
*self.drain_changes(),
)
[docs]
def drain_changes(self) -> tuple[StoreChange, ...]:
"""Take the invalidations the store has accumulated.
A query can expire a run, and a client holding that run has to be told.
The manager drains after serving a query for exactly that reason.
"""
drained = tuple(self._pending)
self._pending.clear()
return drained
[docs]
def stats(self) -> StoreStats:
"""What is retained now, and what retention has already removed."""
self._expire()
return StoreStats(
runs=len(self._runs),
actions=self._actions,
attempts=sum(
len(action.attempts)
for run in self._runs.values()
for action in run.actions.values()
),
bytes=self._bytes,
incomplete_runs=sum(1 for run in self._runs.values() if run.partial),
dedupe_entries=len(self._seen),
evicted_runs=self._evicted_runs,
trimmed_runs=self._trimmed_runs,
limits=self._limits,
)
# -- retention ---------------------------------------------------------
def _measure(self, run: RunState) -> None:
"""Re-price one run and carry the difference into the store total."""
for action in run.actions.values():
for attempt in action.attempts.values():
attempt.bytes = attempt.measure()
action.bytes = action.measure()
measured = run.measure()
self._bytes += measured - run.bytes
run.bytes = measured
def _over_budget(self) -> bool:
return (
len(self._runs) > self._limits.max_runs
or self._actions > self._limits.max_actions
or self._bytes > self._limits.max_bytes
)
def _enforce(self, current: RunState) -> tuple[StoreChange, ...]:
"""Bring the store back inside its limits, newest history first.
Whole runs go before anything inside the run just written to: an
operator watching a live run should not lose its history to a run that
finished ten minutes ago. Only when the current run is all that is left
does the store trim inside it, and only when trimming has run out of
anything to drop does the run itself go.
"""
self._measure(current)
self._expire()
changes: list[StoreChange] = []
while self._over_budget():
victim = self._victim(current)
if victim is not None:
changes.append(self._invalidate(victim)) # queued for the manager
self._evicted_runs += 1
continue
if self._trim(current):
self._measure(current)
continue
# Nothing is left to drop but the run just written to, and it
# still does not fit: trimming stops at the root action, and a
# root whose own payload exceeds the budget never could. Keeping
# it would mean the store living above its limit for as long as
# this run does -- and, worse, its deduplication history growing
# by an entry per event, which trimming actions never reclaims.
# The run goes instead. The next action naming the same root
# starts a new generation, so a subscriber refetches rather than
# comparing revisions across the gap.
changes.append(self._invalidate(current))
self._evicted_runs += 1
break
return tuple(changes)
def _victim(self, current: RunState) -> RunState | None:
"""The run to drop next: oldest completed, else oldest still running."""
candidates = [
run for run in self._runs.values() if run is not current
]
if not candidates:
return None
return min(
candidates,
key=lambda run: (not run.completed, run.updated_at or run.created_at),
)
def _trim(self, run: RunState) -> bool:
"""Drop one action from this run, oldest finished first.
The root action is never dropped -- it is what identifies the run -- and
an operation still in flight outlives a finished one, the same order
used between runs. Returns ``False`` when nothing may be removed, so
the caller stops instead of looping on a run it cannot shrink.
"""
droppable = [
action_id
for action_id in run.actions
if action_id != run.root_action_id
]
if not droppable:
return False
finished = [
action_id
for action_id in droppable
if not run.actions[action_id].terminal_missing
]
# Insertion order is arrival order, which is the only ordering the
# store can trust: source timestamps come from several producers.
run.actions.pop((finished or droppable)[0])
self._actions -= 1
if not run.truncated:
self._trimmed_runs += 1
run.truncated = True
return True
def _expire(self) -> tuple[StoreChange, ...]:
"""Drop runs whose last update is older than the TTL.
Enforced on ingest and on query rather than by a timer: the store has
no thread of its own, and an expired run must not be readable merely
because nothing has happened since it expired.
"""
ttl = self._limits.ttl
if ttl is None:
return ()
deadline = self._clock() - ttl
expired = [
run
for run in list(self._runs.values())
if (run.updated_at or run.created_at or deadline) < deadline
]
changes = []
for run in expired:
changes.append(self._invalidate(run))
self._evicted_runs += 1
return tuple(changes)
# -- queries -----------------------------------------------------------
[docs]
def runs(
self,
*,
agent: str | None = None,
run_id: str | None = None,
errors_only: bool = False,
since: datetime | None = None,
until: datetime | None = None,
offset: int = 0,
limit: int | None = None,
) -> tuple[tuple[RunState, ...], int]:
"""Matching runs, newest first, and how many matched in total.
The count is of matches, not of the page: a client paging through needs
to know what it has not seen yet. Every returned run is a detached copy.
"""
self._expire()
matched = [
run
for run in self._runs.values()
if self._matches(run, agent, run_id, errors_only, since, until)
]
matched.sort(
key=lambda run: (run.updated_at or run.created_at or _EPOCH),
reverse=True,
)
page = matched[offset:] if limit is None else matched[offset : offset + limit]
return tuple(deepcopy(run) for run in page), len(matched)
[docs]
def action(self, action_id: UUID) -> tuple[RunState, ActionState] | None:
"""One action and the run it belongs to, both detached."""
self._expire()
for run in self._runs.values():
found = run.actions.get(action_id)
if found is not None:
return deepcopy(run), deepcopy(found)
return None
@staticmethod
def _matches(
run: RunState,
agent: str | None,
run_id: str | None,
errors_only: bool,
since: datetime | None,
until: datetime | None,
) -> bool:
if agent is not None and run.agent != agent:
return False
if run_id is not None and run.run_id != run_id:
return False
if errors_only and not run.failed:
return False
stamp = run.updated_at or run.created_at
if since is not None and (stamp is None or stamp < since):
return False
if until is not None and (stamp is None or stamp > until):
return False
return True
[docs]
def clear(self) -> tuple[StoreChange, ...]:
"""Forget everything, invalidating every run a client may be holding."""
changes = tuple(self._invalidate(run) for run in list(self._runs.values()))
# The caller is holding them; queueing them again would publish each
# invalidation twice.
self._pending.clear()
self._runs.clear()
self._seen.clear()
# Belt and braces: ownership-aware eviction already returned every
# byte, and a residue here would mean the accounting had drifted.
self._bytes = 0
self._actions = 0
return changes
def _invalidate(self, run: RunState) -> StoreChange:
"""Drop one run and everything the store holds on its behalf.
The record is both returned and queued: a caller that wants it
immediately gets it, and one that does not still cannot lose it.
"""
self._runs.pop(run.root_action_id, None)
for event_id in run.event_ids:
self._seen.pop(event_id, None)
self._actions -= len(run.actions)
self._bytes -= run.bytes
change = StoreChange(
kind=ChangeKind.INVALIDATED,
root_action_id=run.root_action_id,
generation_id=run.generation_id,
revision=run.revision,
run_id=run.run_id,
agent=run.agent,
)
self._pending.append(change)
return change
[docs]
def run(self, root_action_id: UUID) -> RunState | None:
"""A detached snapshot of one run, or ``None`` if it is not retained."""
self._expire()
run = self._runs.get(root_action_id)
return deepcopy(run) if run is not None else None
def _apply(self, run: RunState, action: AgentAction) -> None:
if action.run_id is not None:
run.run_id = action.run_id
if action.conversation_id is not None:
run.conversation_id = action.conversation_id
if action.agent is not None:
run.agent = action.agent
state = run.actions.get(action.action_id)
if state is None:
state = ActionState(
action_id=action.action_id,
root_action_id=run.root_action_id,
parent_action_id=action.parent_action_id,
)
run.actions[action.action_id] = state
self._actions += 1
state.namespace = action.namespace
state.operation = action.operation
if action.parent_action_id is not None:
state.parent_action_id = action.parent_action_id
for name in ("name", "agent", "model", "provider"):
value = getattr(action, name)
if value is not None:
setattr(state, name, value)
if action.metadata:
# Detached on the way in as well as on the way out: the action is
# not frozen, and in-process routing hands the store the producer's
# own object. History must not be editable from the outside.
state.metadata.update(deepcopy(action.metadata))
attempt = state.attempts.get(action.attempt)
if attempt is None:
attempt = AttemptState(attempt=action.attempt)
state.attempts[action.attempt] = attempt
if action.phase is AgentActionPhase.START:
attempt.started_at = action.timestamp
if action.request is not None:
attempt.request = deepcopy(action.request)
else:
attempt.ended_at = action.timestamp
attempt.phase = action.phase
attempt.outcome = action.outcome
attempt.duration = action.duration
if action.response_data is not None:
attempt.response_data = deepcopy(action.response_data)
if action.error is not None:
attempt.error = deepcopy(action.error)
if action.usage is not None:
attempt.usage = deepcopy(action.usage)
self._cap_attempts(state)
def _cap_attempts(self, state: ActionState) -> None:
"""Bound one action's retry history.
Finished attempts go first: an attempt still running is the one an
operator is watching. Whatever is dropped, the action says so -- a
retry list that silently lost its beginning reads as a shorter retry
story than actually happened.
"""
limit = self._limits.max_attempts_per_action
while len(state.attempts) > limit:
completed = [
number
for number, attempt in state.attempts.items()
if not attempt.terminal_missing
]
state.attempts.pop(min(completed) if completed else min(state.attempts))
state.truncated = True