Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions dare_framework/agent/react_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -444,6 +444,8 @@ async def _execute_with_smart_context(
else None
)
messages = self._context.order_messages_for_llm(messages, sys_prompt_message)
if injected_reflection_prompt is not None:
messages.append(injected_reflection_prompt)

# Inject critical_block from plan_provider (maintained by plan tools)
# Disabled: skip injection to observe plan agent behavior without it
Expand Down
225 changes: 225 additions & 0 deletions dare_framework/checkpoint/defaults.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,225 @@
"""Supported default checkpoint implementations and contributors.

This module keeps legacy checkpoint symbols importable after checkpoint internals
were consolidated into ``dare_framework.checkpoint.kernel``.
"""

from __future__ import annotations

from copy import deepcopy
from dataclasses import asdict
from typing import Any
from uuid import uuid4

from dare_framework.context.types import Message

STM = "stm"
SESSION_STATE = "session_state"
SESSION_CONTEXT = "session_context"
WORKSPACE_FILES = "workspace_files"


class MemoryCheckpointStore:
"""In-memory checkpoint payload store kept for facade compatibility."""

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

def put(self, checkpoint_id: str, payload: dict[str, Any]) -> None:
self._store[checkpoint_id] = _clone_payload(payload)

def get(self, checkpoint_id: str) -> dict[str, Any] | None:
if checkpoint_id not in self._store:
return None
return _clone_payload(self._store[checkpoint_id])

def delete(self, checkpoint_id: str) -> bool:
if checkpoint_id in self._store:
del self._store[checkpoint_id]
return True
return False


def _scope_keys(scope: Any, method_name: str) -> list[str]:
method = getattr(scope, method_name, None)
if callable(method):
values = method()
if isinstance(values, (list, tuple, set)):
return [str(v) for v in values]
if isinstance(scope, (list, tuple, set)):
return [str(v) for v in scope]
return []


def _clone_payload(value: Any) -> Any:
"""Clone nested checkpoint payloads so save/restore stays side-effect free."""
return deepcopy(value)


class DefaultCheckpointSaveRestore:
"""Legacy save/restore coordinator over contributor payload components."""

def __init__(self, store: MemoryCheckpointStore, contributors: list[Any]) -> None:
self._store = store
self._contributors = {
str(c.component_key): c
for c in contributors
if getattr(c, "component_key", None) is not None
}

def save(self, scope: Any, ctx: Any) -> str:
payload: dict[str, Any] = {}
for key in _scope_keys(scope, "keys_for_save"):
contributor = self._contributors.get(key)
if contributor is None:
continue
payload[key] = contributor.serialize(ctx)
checkpoint_id = uuid4().hex[:16]
self._store.put(checkpoint_id, payload)
return checkpoint_id

def restore(self, checkpoint_id: str, scope: Any, ctx: Any) -> None:
payload = self._store.get(checkpoint_id)
if payload is None:
raise LookupError(f"Checkpoint not found: {checkpoint_id!r}")
for key in _scope_keys(scope, "keys_for_restore"):
if key not in payload:
continue
contributor = self._contributors.get(key)
if contributor is None:
continue
contributor.deserialize_and_apply(payload[key], ctx)


class StmContributor:
"""Serialize/restore short-term memory messages."""

component_key = STM

def serialize(self, ctx: Any) -> list[dict[str, Any]]:
context = getattr(ctx, "context", None)
if context is None:
return []
messages = context.stm_get()
return [
{
"role": m.role,
"kind": m.kind,
"text": m.text,
"attachments": [
{
"kind": attachment.kind,
"uri": attachment.uri,
"mime_type": attachment.mime_type,
"filename": attachment.filename,
"metadata": _clone_payload(getattr(attachment, "metadata", {}) or {}),
}
for attachment in m.attachments
],
"data": _clone_payload(m.data),
"name": m.name,
"metadata": _clone_payload(getattr(m, "metadata", {}) or {}),
"mark": getattr(m, "mark", None),
"id": getattr(m, "id", None),
}
for m in messages
]

def deserialize_and_apply(self, payload: list[Any], ctx: Any) -> None:
context = getattr(ctx, "context", None)
if context is None:
return
context.stm_clear()
for item in payload or []:
if not isinstance(item, dict):
continue
context.stm_add(
Message(
role=item.get("role", "user"),
kind=item.get("kind", "chat"),
text=item.get("text") or item.get("content", ""),
attachments=_clone_payload(item.get("attachments")) or [],
data=_clone_payload(item.get("data")),
name=item.get("name"),
metadata=dict(item.get("metadata") or {}),
Comment on lines +137 to +144

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve message kind/data in STM checkpoint restore

The new StmContributor restore path reconstructs each message with only role, text, name, and metadata, so any checkpoint containing non-chat STM entries (for example tool_call/tool_result messages with structured data or attachments) is silently downgraded to default chat messages and loses payload fields on restore. This corrupts checkpoint round-trips for realistic agent sessions that rely on typed message state after resume.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Handled in 22cb1da.

StmContributor now serializes/restores kind, attachments, data, mark, and id in addition to the legacy text/name/metadata fields, so typed STM entries survive checkpoint round-trips instead of being downgraded to default chat messages. Added a regression in tests/unit/test_checkpoint_defaults.py that round-trips attachments plus a TOOL_CALL message with structured payload.

mark=item.get("mark", "temporary"),
id=item.get("id"),
)
)


class SessionStateContributor:
"""Serialize/restore minimal session-state fields."""

component_key = SESSION_STATE

def serialize(self, ctx: Any) -> dict[str, Any] | None:
state = getattr(ctx, "session_state", None)
if state is None:
return None
try:
return asdict(state)
except Exception:
return {
"current_milestone_idx": getattr(state, "current_milestone_idx", None),
"task_id": getattr(state, "task_id", None),
"run_id": getattr(state, "run_id", None),
}

def deserialize_and_apply(self, payload: Any, ctx: Any) -> None:
state = getattr(ctx, "session_state", None)
if state is None or not isinstance(payload, dict):
return
if "current_milestone_idx" in payload and hasattr(state, "current_milestone_idx"):
state.current_milestone_idx = payload["current_milestone_idx"]
Comment on lines +173 to +174

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Restore task/run identifiers from session-state payload

SessionStateContributor.serialize() persists task_id and run_id, but deserialize_and_apply() only reapplies current_milestone_idx, so checkpoint restore loses the saved session identity fields. In any resume flow that expects full session-state round-tripping (for example, restoring into a fresh context object), the restored state will keep whatever task_id/run_id happened to be on the target object, which can desynchronize traceability and subsequent orchestration logic from the checkpoint that was actually loaded.

Useful? React with 👍 / 👎.



class SessionContextContributor:
"""Serialize session context for audit trails (restore intentionally no-op)."""

component_key = SESSION_CONTEXT

def serialize(self, ctx: Any) -> dict[str, Any] | None:
session_context = getattr(ctx, "session_context", None)
if session_context is None:
return None
try:
serialized = asdict(session_context)
except Exception:
return {
"session_id": getattr(session_context, "session_id", None),
"task_id": getattr(session_context, "task_id", None),
}
config_value = serialized.get("config")
if config_value is not None and not isinstance(config_value, dict):
try:
serialized["config"] = asdict(config_value)
except Exception:
serialized["config"] = None
return serialized

def deserialize_and_apply(self, payload: Any, ctx: Any) -> None:
_ = (payload, ctx)
# SessionContext is construction-time state; legacy restore path keeps this as no-op.


class WorkspaceGitContributor:
"""Compatibility no-op contributor after workspace-git checkpoint removal."""

component_key = WORKSPACE_FILES

def serialize(self, ctx: Any) -> dict[str, Any]:
_ = ctx
return {}

def deserialize_and_apply(self, payload: Any, ctx: Any) -> None:
_ = (payload, ctx)

__all__ = [
"MemoryCheckpointStore",
"DefaultCheckpointSaveRestore",
"StmContributor",
"WorkspaceGitContributor",
"SessionStateContributor",
"SessionContextContributor",
]
49 changes: 43 additions & 6 deletions dare_framework/compression/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from __future__ import annotations

from dataclasses import asdict, is_dataclass
from typing import TYPE_CHECKING, Any, List, Tuple

from dare_framework.context.types import Message as CtxMessage, MessageKind, MessageMark
Expand All @@ -21,6 +22,36 @@
_UNCHANGED = object()


def _freeze_value(value: Any) -> Any:
"""Build a hashable structural key for nested message payloads."""
if isinstance(value, dict):
return tuple(sorted((str(key), _freeze_value(item)) for key, item in value.items()))
if isinstance(value, list):
return tuple(_freeze_value(item) for item in value)
if isinstance(value, tuple):
return tuple(_freeze_value(item) for item in value)
if isinstance(value, (set, frozenset)):
frozen_items = [_freeze_value(item) for item in value]
return tuple(sorted(frozen_items, key=repr))
if is_dataclass(value) and not isinstance(value, type):
return ("dataclass", type(value).__qualname__, _freeze_value(asdict(value)))
if hasattr(value, "kind") and hasattr(value, "uri"):
return (
getattr(value.kind, "value", value.kind),
value.uri,
value.mime_type,
value.filename,
_freeze_value(getattr(value, "metadata", {})),
)
try:
hash(value)
except TypeError:
if hasattr(value, "__dict__"):
return (type(value).__qualname__, _freeze_value(vars(value)))
return (type(value).__qualname__, repr(value))
return value

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Normalize unhashable payload values before dedup

In the dedup_then_truncate path, _dedup_messages stores keys containing _freeze_value(msg.data) in a set, but _freeze_value falls through to return value for unsupported types. If a message payload contains an unhashable value (for example a set or an unhashable dataclass instance), compression raises TypeError during set membership and aborts the run. This is reachable in production when auto-compression is enabled with the default dedup_then_truncate strategy and tool outputs include non-JSON-native Python values.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Handled in f32e736.

_freeze_value() now normalizes sets/frozensets, dataclass instances, and other unhashable objects before they participate in the dedup identity key, so dedup_then_truncate no longer raises TypeError on non-JSON-native payload values. Added a regression in tests/unit/test_context_compression.py that deduplicates messages containing set payloads.



def _copy_message(
message: Message,
*,
Expand All @@ -41,18 +72,24 @@ def _copy_message(


def _dedup_messages(messages: List[Message]) -> Tuple[List[Message], int]:
"""Lightweight de-duplication on (role, text)."""
seen: set[int] = set()
"""De-duplicate only when the full public message payload matches."""
seen: set[Any] = set()
result: List[Message] = []
removed = 0

for msg in messages:
key = (msg.role, msg.text)
digest = hash(key)
if digest in seen:
key = (
msg.role,
msg.kind,
msg.text,
msg.name,
_freeze_value(msg.attachments),
_freeze_value(msg.data),
)
if key in seen:
removed += 1
continue
seen.add(digest)
seen.add(key)
result.append(msg)

return result, removed
Expand Down
12 changes: 8 additions & 4 deletions dare_framework/context/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,13 +176,17 @@ def budget_remaining(self, resource: str) -> float:
"""Get remaining budget for a resource."""
b = self._budget
if resource == "tokens":
return (b.max_tokens - b.used_tokens) if b.max_tokens else float("inf")
return (b.max_tokens - b.used_tokens) if b.max_tokens is not None else float("inf")
elif resource == "cost":
return (b.max_cost - b.used_cost) if b.max_cost else float("inf")
return (b.max_cost - b.used_cost) if b.max_cost is not None else float("inf")
elif resource == "tool_calls":
return (b.max_tool_calls - b.used_tool_calls) if b.max_tool_calls else float("inf")
return (b.max_tool_calls - b.used_tool_calls) if b.max_tool_calls is not None else float("inf")
elif resource == "time_seconds":
return (b.max_time_seconds - b.used_time_seconds) if b.max_time_seconds else float("inf")
return (
(b.max_time_seconds - b.used_time_seconds)
if b.max_time_seconds is not None
else float("inf")
)
return float("inf")

# ========== Tool Methods ==========
Expand Down
6 changes: 3 additions & 3 deletions dare_framework/embedding/__init__.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
"""embedding domain facade."""

from dare_framework.embedding.interfaces import IEmbeddingAdapter
from dare_framework.embedding.types import EmbeddingOptions, EmbeddingResult
from dare_framework.embedding._internal.openai_embedding import OpenAIEmbeddingAdapter
from dare_framework.embedding.interfaces import IEmbeddingAdapter
from dare_framework.embedding.types import EmbeddingOptions, EmbeddingResult
from dare_framework.embedding.defaults import OpenAIEmbeddingAdapter

__all__ = [
"IEmbeddingAdapter",
Expand Down
5 changes: 5 additions & 0 deletions dare_framework/embedding/defaults.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
"""Supported default embedding implementations."""

from dare_framework.embedding._internal.openai_embedding import OpenAIEmbeddingAdapter

__all__ = ["OpenAIEmbeddingAdapter"]
2 changes: 1 addition & 1 deletion dare_framework/event/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""event domain facade."""

from dare_framework.event._internal.sqlite_event_log import DefaultEventLog, SQLiteEventLog
from dare_framework.event.defaults import DefaultEventLog, SQLiteEventLog
from dare_framework.event.kernel import IEventLog
from dare_framework.event.types import Event, RuntimeSnapshot

Expand Down
5 changes: 5 additions & 0 deletions dare_framework/event/defaults.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
"""Supported default event-log implementations."""

from dare_framework.event._internal.sqlite_event_log import DefaultEventLog, SQLiteEventLog

__all__ = ["SQLiteEventLog", "DefaultEventLog"]
2 changes: 1 addition & 1 deletion dare_framework/hook/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from dare_framework.hook.interfaces import IHookManager
from dare_framework.hook.kernel import IExtensionPoint, IHook, HookFn
from dare_framework.hook._internal.hook_extension_point import HookExtensionPoint
from dare_framework.hook.defaults import HookExtensionPoint
from dare_framework.hook.types import HookDecision, HookEnvelope, HookPhase, HookResult

__all__ = [
Expand Down
5 changes: 5 additions & 0 deletions dare_framework/hook/defaults.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
"""Supported default extension-point implementation."""

from dare_framework.hook._internal.hook_extension_point import HookExtensionPoint

__all__ = ["HookExtensionPoint"]
8 changes: 7 additions & 1 deletion dare_framework/memory/in_memory_stm.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,13 @@ def compress(self, max_messages: int | None = None, **kwargs) -> int:
Returns:
Number of messages removed.
"""
if max_messages is None or len(self._messages) <= max_messages:
if max_messages is None:
return 0
if max_messages <= 0:
removed_count = len(self._messages)
self._messages = []
return removed_count
if len(self._messages) <= max_messages:
return 0

removed_count = len(self._messages) - max_messages
Expand Down
Loading