39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
"""Runtime audit context shared by user-facing PMFlow AI and LLM logs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from typing import Iterator
|
|
|
|
|
|
_agent_run_id: ContextVar[str | None] = ContextVar("pmflow_audit_agent_run_id", default=None)
|
|
_agent_step_key: ContextVar[str | None] = ContextVar("pmflow_audit_agent_step_key", default=None)
|
|
|
|
|
|
@contextmanager
|
|
def bind_agent_audit_context(*, agent_run_id: str | None = None, agent_step_key: str | None = None) -> Iterator[None]:
|
|
run_token = _agent_run_id.set(str(agent_run_id or "").strip() or None) if agent_run_id is not None else None
|
|
step_token = _agent_step_key.set(str(agent_step_key or "").strip() or None) if agent_step_key is not None else None
|
|
try:
|
|
yield
|
|
finally:
|
|
if step_token is not None:
|
|
_agent_step_key.reset(step_token)
|
|
if run_token is not None:
|
|
_agent_run_id.reset(run_token)
|
|
|
|
|
|
def current_agent_audit_context() -> dict[str, str]:
|
|
payload: dict[str, str] = {}
|
|
run_id = str(_agent_run_id.get() or "").strip()
|
|
step_key = str(_agent_step_key.get() or "").strip()
|
|
if run_id:
|
|
payload["agent_run_id"] = run_id[:80]
|
|
if step_key:
|
|
payload["agent_step_key"] = step_key[:120]
|
|
return payload
|
|
|
|
|
|
__all__ = ["bind_agent_audit_context", "current_agent_audit_context"]
|