datarobot_genai.dragent.agent_card_registry_backends

agent_card_registry_backends

Pluggable cache backends for the central agent card registry.

RegistryIds

Bases: NamedTuple

The registry identifiers a single agent card entry carries.

A card is keyed by exactly one of deployment_id / workload_id (the platform enforces this) and may additionally publish an external_id. Named fields rather than a positional tuple because three optional strings in a row are indistinguishable at a call site.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

class RegistryIds(NamedTuple):
    """The registry identifiers a single agent card entry carries.

    A card is keyed by exactly one of ``deployment_id`` / ``workload_id`` (the
    platform enforces this) and may additionally publish an ``external_id``.
    Named fields rather than a positional tuple because three optional strings
    in a row are indistinguishable at a call site.
    """

    deployment_id: str | None = None
    external_id: str | None = None
    workload_id: str | None = None

AgentCardCacheRecord

Bases: BaseModel

Serialized agent card cache entry shared across backends.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

class AgentCardCacheRecord(BaseModel):
    """Serialized agent card cache entry shared across backends."""

    version: int = 1
    fetched_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
    card: AgentCard
    source: str = "registry"
    deployment_id: str | None = None
    external_id: str | None = None
    workload_id: str | None = None

    def age_seconds(self) -> float:
        """Return the entry age in seconds (wall clock, safe across processes)."""
        return max(0.0, (datetime.now(UTC) - self.fetched_at).total_seconds())

    def is_fresh(self, cache_ttl: int) -> bool:
        """Return *True* if this entry is within the soft TTL.

        *cache_ttl* is the caller's soft-TTL bound (not the hard stale-if-error limit).
        """
        if cache_ttl == 0:
            return False
        return self.age_seconds() < cache_ttl

    def is_within_staleness(self, max_staleness_seconds: int) -> bool:
        """Return *True* if this entry may be served under stale-if-error."""
        if max_staleness_seconds == 0:
            return False
        return self.age_seconds() <= max_staleness_seconds

age_seconds

age_seconds() -> float

Return the entry age in seconds (wall clock, safe across processes).

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

def age_seconds(self) -> float:
    """Return the entry age in seconds (wall clock, safe across processes)."""
    return max(0.0, (datetime.now(UTC) - self.fetched_at).total_seconds())

is_fresh

is_fresh(cache_ttl: int) -> bool

Return True if this entry is within the soft TTL.

cache_ttl is the caller's soft-TTL bound (not the hard stale-if-error limit).

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

def is_fresh(self, cache_ttl: int) -> bool:
    """Return *True* if this entry is within the soft TTL.

    *cache_ttl* is the caller's soft-TTL bound (not the hard stale-if-error limit).
    """
    if cache_ttl == 0:
        return False
    return self.age_seconds() < cache_ttl

is_within_staleness

is_within_staleness(max_staleness_seconds: int) -> bool

Return True if this entry may be served under stale-if-error.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

def is_within_staleness(self, max_staleness_seconds: int) -> bool:
    """Return *True* if this entry may be served under stale-if-error."""
    if max_staleness_seconds == 0:
        return False
    return self.age_seconds() <= max_staleness_seconds

AgentCardCacheBackend

Bases: Protocol

Async cache backend for agent card registry entries.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

class AgentCardCacheBackend(Protocol):
    """Async cache backend for agent card registry entries."""

    async def get_fresh(
        self,
        lookup_key: str,
        *,
        cache_ttl: int,
        key_type: LookupKeyType | None = None,
    ) -> AgentCardCacheRecord | None:
        """Return a cached record when within the soft TTL."""

    async def get_stale(
        self,
        lookup_key: str,
        *,
        max_staleness_seconds: int,
        key_type: LookupKeyType | None = None,
    ) -> AgentCardCacheRecord | None:
        """Return a cached record within the hard staleness bound."""

    async def store(
        self,
        cards: dict[str, AgentCard],
        *,
        key_types: dict[str, LookupKeyType],
        registry_ids: dict[str, RegistryIds] | None = None,
    ) -> None:
        """Persist one or more cards keyed by lookup ID."""

    async def evict(
        self,
        lookup_key: str,
        *,
        key_type: LookupKeyType | None = None,
    ) -> None:
        """Remove a cached entry so stale-if-error cannot resurrect it."""

get_fresh async

get_fresh(lookup_key: str, *, cache_ttl: int, key_type: LookupKeyType | None = None) -> AgentCardCacheRecord | None

Return a cached record when within the soft TTL.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

async def get_fresh(
    self,
    lookup_key: str,
    *,
    cache_ttl: int,
    key_type: LookupKeyType | None = None,
) -> AgentCardCacheRecord | None:
    """Return a cached record when within the soft TTL."""

get_stale async

get_stale(lookup_key: str, *, max_staleness_seconds: int, key_type: LookupKeyType | None = None) -> AgentCardCacheRecord | None

Return a cached record within the hard staleness bound.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

async def get_stale(
    self,
    lookup_key: str,
    *,
    max_staleness_seconds: int,
    key_type: LookupKeyType | None = None,
) -> AgentCardCacheRecord | None:
    """Return a cached record within the hard staleness bound."""

store async

store(cards: dict[str, AgentCard], *, key_types: dict[str, LookupKeyType], registry_ids: dict[str, RegistryIds] | None = None) -> None

Persist one or more cards keyed by lookup ID.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

async def store(
    self,
    cards: dict[str, AgentCard],
    *,
    key_types: dict[str, LookupKeyType],
    registry_ids: dict[str, RegistryIds] | None = None,
) -> None:
    """Persist one or more cards keyed by lookup ID."""

evict async

evict(lookup_key: str, *, key_type: LookupKeyType | None = None) -> None

Remove a cached entry so stale-if-error cannot resurrect it.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

async def evict(
    self,
    lookup_key: str,
    *,
    key_type: LookupKeyType | None = None,
) -> None:
    """Remove a cached entry so stale-if-error cannot resurrect it."""

MemoryAgentCardCacheBackend

In-process dict cache (L1).

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

class MemoryAgentCardCacheBackend:
    """In-process dict cache (L1)."""

    def __init__(self) -> None:
        self._entries: dict[str, AgentCardCacheRecord] = {}

    async def get_fresh(
        self,
        lookup_key: str,
        *,
        cache_ttl: int,
        key_type: LookupKeyType | None = None,
    ) -> AgentCardCacheRecord | None:
        record = self._entries.get(lookup_key)
        if record is None or not record.is_fresh(cache_ttl):
            return None
        return record

    async def get_stale(
        self,
        lookup_key: str,
        *,
        max_staleness_seconds: int,
        key_type: LookupKeyType | None = None,
    ) -> AgentCardCacheRecord | None:
        record = self._entries.get(lookup_key)
        if record is None or not record.is_within_staleness(max_staleness_seconds):
            return None
        return record

    async def store(
        self,
        cards: dict[str, AgentCard],
        *,
        key_types: dict[str, LookupKeyType],
        registry_ids: dict[str, RegistryIds] | None = None,
    ) -> None:
        id_sets = registry_ids or {}
        for lookup_key, card in cards.items():
            key_type = key_types.get(lookup_key, "deployment")
            ids = id_sets.get(lookup_key, RegistryIds())
            self._entries[lookup_key] = build_cache_record(
                card,
                lookup_key=lookup_key,
                key_type=key_type,
                deployment_id=ids.deployment_id,
                external_id=ids.external_id,
                workload_id=ids.workload_id,
            )

    async def store_record(self, lookup_key: str, record: AgentCardCacheRecord) -> None:
        """Insert *record* as-is, preserving ``fetched_at`` for L1 read-through."""
        self._entries[lookup_key] = record.model_copy()

    def has_entry(self, lookup_key: str) -> bool:
        """Return *True* when *lookup_key* is present in L1 (fresh or stale)."""
        return lookup_key in self._entries

    async def evict(
        self,
        lookup_key: str,
        *,
        key_type: LookupKeyType | None = None,
    ) -> None:
        keys_to_evict = {lookup_key}
        record = self._entries.get(lookup_key)
        if record is not None:
            if record.deployment_id:
                keys_to_evict.add(record.deployment_id)
            if record.external_id:
                keys_to_evict.add(record.external_id)
            if record.workload_id:
                keys_to_evict.add(record.workload_id)
        for key in keys_to_evict:
            self._entries.pop(key, None)

    def age_entry_for_test(self, lookup_key: str, seconds: float) -> None:
        """Shift *lookup_key* fetch time backward (tests only)."""
        record = self._entries.get(lookup_key)
        if record is None:
            return
        record.fetched_at -= timedelta(seconds=seconds)

store_record async

store_record(lookup_key: str, record: AgentCardCacheRecord) -> None

Insert record as-is, preserving fetched_at for L1 read-through.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

async def store_record(self, lookup_key: str, record: AgentCardCacheRecord) -> None:
    """Insert *record* as-is, preserving ``fetched_at`` for L1 read-through."""
    self._entries[lookup_key] = record.model_copy()

has_entry

has_entry(lookup_key: str) -> bool

Return True when lookup_key is present in L1 (fresh or stale).

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

def has_entry(self, lookup_key: str) -> bool:
    """Return *True* when *lookup_key* is present in L1 (fresh or stale)."""
    return lookup_key in self._entries

age_entry_for_test

age_entry_for_test(lookup_key: str, seconds: float) -> None

Shift lookup_key fetch time backward (tests only).

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

def age_entry_for_test(self, lookup_key: str, seconds: float) -> None:
    """Shift *lookup_key* fetch time backward (tests only)."""
    record = self._entries.get(lookup_key)
    if record is None:
        return
    record.fetched_at -= timedelta(seconds=seconds)

MemorySpaceAgentCardCacheBackend

Shared DataRobot MemorySpace L2 cache using JSON payloads.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

class MemorySpaceAgentCardCacheBackend:
    """Shared DataRobot MemorySpace L2 cache using JSON payloads."""

    def __init__(
        self,
        *,
        kv_cache: MemorySpaceKVCache,
    ) -> None:
        self._kv = kv_cache

    def _storage_keys_for_record(self, record: AgentCardCacheRecord) -> list[str]:
        keys: list[str] = []
        if record.deployment_id:
            keys.append(f"deployment:{record.deployment_id}")
        if record.external_id:
            keys.append(f"external:{record.external_id}")
        if record.workload_id:
            keys.append(f"workload:{record.workload_id}")
        return keys

    def _storage_keys_for_lookup(
        self,
        lookup_key: str,
        *,
        key_type: LookupKeyType | None = None,
    ) -> list[str]:
        if key_type == "deployment":
            return [f"deployment:{lookup_key}"]
        if key_type == "external":
            return [f"external:{lookup_key}"]
        if key_type == "workload":
            return [f"workload:{lookup_key}"]
        return [
            f"deployment:{lookup_key}",
            f"external:{lookup_key}",
            f"workload:{lookup_key}",
        ]

    async def _get_record(self, storage_key: str) -> AgentCardCacheRecord | None:
        payload = await self._kv.get_value(storage_key)
        if payload is None:
            return None
        try:
            return AgentCardCacheRecord.model_validate_json(payload)
        except Exception:
            logger.warning("Ignoring invalid MemorySpace agent card payload for %s", storage_key)
            return None

    async def get_fresh(
        self,
        lookup_key: str,
        *,
        cache_ttl: int,
        key_type: LookupKeyType | None = None,
    ) -> AgentCardCacheRecord | None:
        for storage_key in self._storage_keys_for_lookup(lookup_key, key_type=key_type):
            record = await self._get_record(storage_key)
            if record is not None and record.is_fresh(cache_ttl):
                return record
        return None

    async def get_stale(
        self,
        lookup_key: str,
        *,
        max_staleness_seconds: int,
        key_type: LookupKeyType | None = None,
    ) -> AgentCardCacheRecord | None:
        for storage_key in self._storage_keys_for_lookup(lookup_key, key_type=key_type):
            record = await self._get_record(storage_key)
            if record is not None and record.is_within_staleness(max_staleness_seconds):
                return record
        return None

    async def store(
        self,
        cards: dict[str, AgentCard],
        *,
        key_types: dict[str, LookupKeyType],
        registry_ids: dict[str, RegistryIds] | None = None,
    ) -> None:
        id_sets = registry_ids or {}
        for lookup_key, card in cards.items():
            key_type = key_types.get(lookup_key, "deployment")
            ids = id_sets.get(lookup_key, RegistryIds())
            record = build_cache_record(
                card,
                lookup_key=lookup_key,
                key_type=key_type,
                deployment_id=ids.deployment_id,
                external_id=ids.external_id,
                workload_id=ids.workload_id,
            )
            payload = record.model_dump_json()
            for storage_key in self._storage_keys_for_record(record):
                await self._kv.set_value(storage_key, payload)

    async def evict(
        self,
        lookup_key: str,
        *,
        key_type: LookupKeyType | None = None,
    ) -> None:
        keys_to_evict = set(self._storage_keys_for_lookup(lookup_key, key_type=key_type))
        record: AgentCardCacheRecord | None = None
        for storage_key in keys_to_evict:
            if record is None:
                record = await self._get_record(storage_key)
        if record is not None:
            keys_to_evict.update(self._storage_keys_for_record(record))
        for storage_key in keys_to_evict:
            await self._kv.delete_value(storage_key)

LayeredAgentCardCacheBackend

L1 memory read-through / write-behind over an L2 backend.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

class LayeredAgentCardCacheBackend:
    """L1 memory read-through / write-behind over an L2 backend."""

    def __init__(
        self,
        l1: MemoryAgentCardCacheBackend,
        l2: AgentCardCacheBackend,
        *,
        l2_read_timeout: float = _DEFAULT_L2_READ_TIMEOUT_SECONDS,
    ) -> None:
        self._l1 = l1
        self._l2 = l2
        self._l2_read_timeout = l2_read_timeout
        self._l2_tasks: set[asyncio.Task[Any]] = set()

    def _schedule_l2(self, coro: Coroutine[Any, Any, Any]) -> None:
        """Run an L2 operation in the background without blocking callers."""
        task = asyncio.create_task(coro)
        self._l2_tasks.add(task)
        task.add_done_callback(self._l2_tasks.discard)
        task.add_done_callback(self._log_l2_task_failure)

    @staticmethod
    def _log_l2_task_failure(task: asyncio.Task[Any]) -> None:
        if task.cancelled():
            return
        try:
            exc = task.exception()
        except asyncio.CancelledError:
            return
        if exc is not None:
            logger.warning("Agent card registry L2 background task failed", exc_info=exc)

    async def flush_l2_tasks(self) -> None:
        """Await in-flight L2 writes/evictions (tests and graceful shutdown)."""
        if not self._l2_tasks:
            return
        tasks = list(self._l2_tasks)
        await asyncio.gather(*tasks, return_exceptions=True)

    async def _read_l2_with_timeout(
        self,
        coro: Coroutine[Any, Any, AgentCardCacheRecord | None],
        *,
        lookup_key: str,
        op: str,
    ) -> AgentCardCacheRecord | None:
        """Run an L2 read with a bounded wait; return ``None`` on timeout."""
        if self._l2_read_timeout <= 0:
            return await coro
        try:
            return await asyncio.wait_for(coro, timeout=self._l2_read_timeout)
        except TimeoutError:
            logger.warning(
                "Agent card registry cache: MemorySpace L2 %s read timed out after %.2fs for %s",
                op,
                self._l2_read_timeout,
                lookup_key,
            )
            return None

    async def get_fresh(
        self,
        lookup_key: str,
        *,
        cache_ttl: int,
        key_type: LookupKeyType | None = None,
    ) -> AgentCardCacheRecord | None:
        if record := await self._l1.get_fresh(lookup_key, cache_ttl=cache_ttl, key_type=key_type):
            return record
        # L1 holds a stale entry with the same fetched_at L2 would return — skip L2.
        if self._l1.has_entry(lookup_key):
            logger.debug(
                "Agent card registry cache: L1 stale entry for %s; skipping MemorySpace L2 "
                "fresh read",
                lookup_key,
            )
            return None
        logger.info(
            "Agent card registry cache: L1 miss for %s (key_type=%s); reading MemorySpace L2",
            lookup_key,
            key_type or "any",
        )
        if record := await self._read_l2_with_timeout(
            self._l2.get_fresh(
                lookup_key,
                cache_ttl=cache_ttl,
                key_type=key_type,
            ),
            lookup_key=lookup_key,
            op="fresh",
        ):
            logger.info(
                "Agent card registry cache: MemorySpace L2 hit for %s (fresh, age=%.0fs); "
                "promoted to L1",
                lookup_key,
                record.age_seconds(),
            )
            await self._promote_to_l1(lookup_key, record)
            return record
        logger.info(
            "Agent card registry cache: MemorySpace L2 miss for %s (no fresh entry)",
            lookup_key,
        )
        return None

    async def get_stale(
        self,
        lookup_key: str,
        *,
        max_staleness_seconds: int,
        key_type: LookupKeyType | None = None,
    ) -> AgentCardCacheRecord | None:
        if record := await self._l1.get_stale(
            lookup_key,
            max_staleness_seconds=max_staleness_seconds,
            key_type=key_type,
        ):
            logger.info(
                "Agent card registry cache: L1 hit for %s (stale, age=%.0fs)",
                lookup_key,
                record.age_seconds(),
            )
            return record
        # Stale-if-error runs only after a registry fetch failed — wait for L2.
        logger.info(
            "Agent card registry cache: L1 miss for %s; reading MemorySpace L2 (stale-if-error)",
            lookup_key,
        )
        if record := await self._l2.get_stale(
            lookup_key,
            max_staleness_seconds=max_staleness_seconds,
            key_type=key_type,
        ):
            logger.info(
                "Agent card registry cache: MemorySpace L2 hit for %s (stale, age=%.0fs); "
                "promoted to L1",
                lookup_key,
                record.age_seconds(),
            )
            await self._promote_to_l1(lookup_key, record)
            return record
        logger.info(
            "Agent card registry cache: MemorySpace L2 miss for %s (no stale entry within bound)",
            lookup_key,
        )
        return None

    async def _promote_to_l1(self, lookup_key: str, record: AgentCardCacheRecord) -> None:
        await self._l1.store_record(lookup_key, record)

    async def store(
        self,
        cards: dict[str, AgentCard],
        *,
        key_types: dict[str, LookupKeyType],
        registry_ids: dict[str, RegistryIds] | None = None,
    ) -> None:
        await self._l1.store(cards, key_types=key_types, registry_ids=registry_ids)
        logger.info(
            "Agent card registry cache: write-behind to MemorySpace L2 for %d key(s)",
            len(cards),
        )
        self._schedule_l2(self._l2.store(cards, key_types=key_types, registry_ids=registry_ids))

    async def evict(
        self,
        lookup_key: str,
        *,
        key_type: LookupKeyType | None = None,
    ) -> None:
        await self._l1.evict(lookup_key, key_type=key_type)
        # Evict L2 before returning — write-behind store must not let a later
        # get_fresh read a deregistered card from L2 after L1 was cleared.
        await self._l2.evict(lookup_key, key_type=key_type)

    @property
    def memory(self) -> MemoryAgentCardCacheBackend:
        """Expose the L1 backend (tests)."""
        return self._l1

memory property

memory: MemoryAgentCardCacheBackend

Expose the L1 backend (tests).

flush_l2_tasks async

flush_l2_tasks() -> None

Await in-flight L2 writes/evictions (tests and graceful shutdown).

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

async def flush_l2_tasks(self) -> None:
    """Await in-flight L2 writes/evictions (tests and graceful shutdown)."""
    if not self._l2_tasks:
        return
    tasks = list(self._l2_tasks)
    await asyncio.gather(*tasks, return_exceptions=True)

build_cache_record

build_cache_record(card: AgentCard, *, lookup_key: str, key_type: LookupKeyType, deployment_id: str | None = None, external_id: str | None = None, workload_id: str | None = None) -> AgentCardCacheRecord

Build a cache record for lookup_key with optional registry ID metadata.

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

def build_cache_record(
    card: AgentCard,
    *,
    lookup_key: str,
    key_type: LookupKeyType,
    deployment_id: str | None = None,
    external_id: str | None = None,
    workload_id: str | None = None,
) -> AgentCardCacheRecord:
    """Build a cache record for *lookup_key* with optional registry ID metadata."""
    resolved_dep = (
        deployment_id
        if deployment_id is not None
        else (lookup_key if key_type == "deployment" else None)
    )
    resolved_ext = (
        external_id if external_id is not None else (lookup_key if key_type == "external" else None)
    )
    resolved_wl = (
        workload_id if workload_id is not None else (lookup_key if key_type == "workload" else None)
    )
    return AgentCardCacheRecord(
        card=card,
        deployment_id=resolved_dep,
        external_id=resolved_ext,
        workload_id=resolved_wl,
    )

create_agent_card_cache_backend

create_agent_card_cache_backend(cache_ttl: int) -> AgentCardCacheBackend

Instantiate the agent card cache backend (L1, plus MemorySpace L2 when available).

Source code in datarobot_genai/dragent/agent_card_registry_backends.py

def create_agent_card_cache_backend(cache_ttl: int) -> AgentCardCacheBackend:
    """Instantiate the agent card cache backend (L1, plus MemorySpace L2 when available)."""
    l1 = MemoryAgentCardCacheBackend()
    if cache_ttl == 0:
        logger.debug("Agent card registry cache: L1 only (cache_ttl=0)")
        return l1

    memory_space_id = try_resolve_memory_space_id()
    if memory_space_id is None:
        if is_enclave_l2_workload():
            logger.warning(
                "Agent card registry cache: L1 only (enclave workload but MemorySpace id "
                "is not yet provisioned; L2 write-behind disabled)"
            )
        else:
            logger.debug("Agent card registry cache: L1 only (not on an enclave workload)")
        return l1

    kv_cache = MemorySpaceKVCache(memory_space_id=memory_space_id)
    logger.info(
        "Agent card registry cache: L1 + MemorySpace L2 enabled (space_id=%s)",
        memory_space_id,
    )
    return LayeredAgentCardCacheBackend(
        l1,
        MemorySpaceAgentCardCacheBackend(kv_cache=kv_cache),
    )