507 lines
20 KiB
Python
507 lines
20 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
"""ContextFactory centralizes context construction and one-shot checkpoint hydration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
import pickle
|
|
from typing import Any, Generic, TypeVar
|
|
|
|
from band_aider.context.alpha_context import AlphaContext
|
|
from band_aider.context.beta_context import BetaContext
|
|
from band_aider.context.bootstrap_context import BootstrapContext
|
|
from band_aider.context.contex_bag import ContextBag
|
|
from band_aider.context.global_context import GlobalContext, SovereignState, StoryCursor, TaskCursor
|
|
from band_aider.context.resume_state import resume_settings
|
|
from band_aider.context.stabilization_context import StabilizationContext
|
|
from band_aider.context.topological_liveliness import TopologicalLiveliness
|
|
from band_aider.context.working_backlog import BackLog
|
|
from band_aider.judgement.code_dependency_graph_maintainer import (
|
|
CodeDependencyGraphMaintainer,
|
|
)
|
|
from band_aider.judgement.dependency_discovery_strategy import (
|
|
DependencyDiscoveryStrategy,
|
|
)
|
|
from band_aider.judgement.file_manager import FileManager
|
|
from band_aider.strategies import LanguageStrategySet
|
|
from band_aider.judgement.repo_exec import RepoExec
|
|
|
|
T = TypeVar("T", bound=ContextBag)
|
|
|
|
|
|
def _rebind_file_manager_dependency_runtime(file_manager: FileManager) -> None:
|
|
"""Restore current FileManager dependency collaborators and graph state."""
|
|
|
|
language_strategy_set = LanguageStrategySet.get(
|
|
file_manager,
|
|
maintainer=file_manager.dependency_graph_maintainer,
|
|
)
|
|
if file_manager.dependency_graph_maintainer is None:
|
|
file_manager.dependency_graph_maintainer = language_strategy_set[
|
|
CodeDependencyGraphMaintainer
|
|
]
|
|
if file_manager.dependency_strategy is None:
|
|
file_manager.dependency_strategy = language_strategy_set[
|
|
DependencyDiscoveryStrategy
|
|
]
|
|
elif hasattr(file_manager.dependency_strategy, "maintainer"):
|
|
file_manager.dependency_strategy.maintainer = (
|
|
file_manager.dependency_graph_maintainer
|
|
)
|
|
|
|
maintainer = file_manager.dependency_graph_maintainer
|
|
graph = maintainer.graph
|
|
if (
|
|
graph.source_by_location
|
|
or graph.test_by_location
|
|
or not (file_manager.src_classes or file_manager.test_classes)
|
|
):
|
|
return
|
|
|
|
# Hydration repairs an absent runtime graph only. It does not attempt stale
|
|
# detection for source changes made between stop and resume.
|
|
maintainer.build(
|
|
file_manager.src_classes,
|
|
file_manager.test_classes,
|
|
file_manager.repo_root,
|
|
file_manager.src_root,
|
|
file_manager.test_root,
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LoadedCheckpoint:
|
|
"""LoadedCheckpoint carries one checkpoint payload plus its schema version."""
|
|
|
|
version: int
|
|
payload: dict[str, Any]
|
|
|
|
|
|
def _declared_attributes(context_type: type[object]) -> set[str]:
|
|
"""Return the public and annotated attributes supported by a context type."""
|
|
|
|
attributes = set(dir(context_type))
|
|
for base in context_type.__mro__:
|
|
attributes.update(getattr(base, "__annotations__", {}))
|
|
return attributes
|
|
|
|
|
|
def _checkpoint_shape_is_current(
|
|
context_type: type[ContextBag],
|
|
payload: dict[str, Any],
|
|
) -> bool:
|
|
"""Accept only checkpoints whose object shape belongs to the current code.
|
|
|
|
Checkpoints are temporary run state, not a compatibility database. If a
|
|
payload contains retired state or omits current required state, hydration
|
|
is intentionally rejected and the caller may start without it.
|
|
This policy check applies to the context and its nested cursor objects;
|
|
it does not grow a field-by-field migration layer.
|
|
"""
|
|
|
|
# Reject retired context fields before any nested state is considered.
|
|
if not set(payload).issubset(_declared_attributes(context_type)):
|
|
return False
|
|
# Reject checkpoints missing the context-owned BFS and deferred-topology lifecycle state.
|
|
if context_type is StabilizationContext:
|
|
# Require every durable stabilization owner introduced by the current shape.
|
|
if not {
|
|
"topology",
|
|
"prior_traversed_frontiers",
|
|
"frontier_open_hashes",
|
|
}.issubset(payload):
|
|
return False
|
|
topology = payload["topology"]
|
|
# Reject values that cannot enforce TopologicalLiveliness invariants.
|
|
if not isinstance(topology, TopologicalLiveliness):
|
|
return False
|
|
# Require the durable admitted-frontier state needed to resume the active COT.
|
|
if not {
|
|
"frontiers",
|
|
"bfs_depth",
|
|
"visited",
|
|
"whitelist",
|
|
"deferred_topology_locations",
|
|
"unchanged_perimeter",
|
|
}.issubset(vars(topology)):
|
|
return False
|
|
frontier_open_hashes = payload["frontier_open_hashes"]
|
|
# Depth zero identifies context.t0; positive depths identify later admitted frontiers.
|
|
initial_bfs_depth = 0
|
|
# Accept only depth-keyed string hash baselines captured before Alpha.
|
|
if not isinstance(frontier_open_hashes, dict) or not all(
|
|
map(
|
|
lambda item: (
|
|
isinstance(item[0], int)
|
|
and item[0] >= initial_bfs_depth
|
|
and isinstance(item[1], dict)
|
|
and all(map(
|
|
lambda hash_item: (
|
|
isinstance(hash_item[0], str)
|
|
and isinstance(hash_item[1], str)
|
|
),
|
|
item[1].items(),
|
|
))
|
|
),
|
|
frontier_open_hashes.items(),
|
|
)
|
|
):
|
|
return False
|
|
# The active topological position cannot outpace recorded frontiers.
|
|
if (
|
|
topology.bfs_depth != len(topology.frontiers)
|
|
):
|
|
return False
|
|
# Every recorded frontier retains its pre-Alpha hash baseline.
|
|
if (
|
|
topology.frontiers
|
|
and topology.bfs_depth not in frontier_open_hashes
|
|
):
|
|
return False
|
|
# Contexts without a story cursor need no further nested-shape validation.
|
|
if context_type is not GlobalContext:
|
|
return True
|
|
|
|
# A current GlobalContext always owns a resumable story cursor.
|
|
if "story_cursor" not in payload:
|
|
return False
|
|
cursor = payload["story_cursor"]
|
|
# Reject payloads that cannot enforce StoryCursor lifecycle invariants.
|
|
if not isinstance(cursor, StoryCursor):
|
|
return False
|
|
# Reject retired cursor fields instead of silently retaining stale state.
|
|
if not set(vars(cursor)).issubset(_declared_attributes(StoryCursor)):
|
|
return False
|
|
# A story cursor must carry the current task cursor implementation.
|
|
if not isinstance(cursor.task_cursor, TaskCursor):
|
|
return False
|
|
task_cursor_attributes = vars(cursor.task_cursor)
|
|
# Reject retired task-cursor fields before accepting its traversal lineage.
|
|
if not set(task_cursor_attributes).issubset(_declared_attributes(TaskCursor)):
|
|
return False
|
|
# Reject a cursor missing lineage required by the current COT scheduler.
|
|
if "prior_traversed_frontiers" not in task_cursor_attributes:
|
|
return False
|
|
return True
|
|
|
|
|
|
class Olaf(Generic[T], ABC):
|
|
"""Olaf thaws persisted fields and delegates runtime reassembly per context."""
|
|
|
|
def thaw(self, context: T, checkpoint: LoadedCheckpoint, kwargs: dict[str, Any]) -> T:
|
|
"""Hydrate one context from persisted payload and caller-supplied runtime wiring."""
|
|
|
|
for name, value in checkpoint.payload.items():
|
|
setattr(context, name, value)
|
|
self.rebind_runtime(context, checkpoint, kwargs)
|
|
self.post_thaw(context)
|
|
return context
|
|
|
|
@abstractmethod
|
|
def rebind_runtime(
|
|
self,
|
|
context: T,
|
|
checkpoint: LoadedCheckpoint,
|
|
kwargs: dict[str, Any],
|
|
) -> None:
|
|
"""Rebind runtime-owned fields after persisted payload has been restored."""
|
|
|
|
def post_thaw(self, context: T) -> None:
|
|
"""Run any post-thaw normalization."""
|
|
|
|
|
|
class DefaultReassembly(Olaf[T]):
|
|
"""Default reassembly applies persisted payload without runtime rebinding."""
|
|
|
|
def rebind_runtime(
|
|
self,
|
|
context: T,
|
|
checkpoint: LoadedCheckpoint,
|
|
kwargs: dict[str, Any],
|
|
) -> None:
|
|
"""Leave runtime state unchanged for contexts without special ownership rules."""
|
|
|
|
|
|
class GlobalContextReassembly(Olaf[GlobalContext]):
|
|
"""GlobalContextReassembly restores top-level runtime services from CLI wiring."""
|
|
|
|
def rebind_runtime(
|
|
self,
|
|
context: GlobalContext,
|
|
checkpoint: LoadedCheckpoint,
|
|
kwargs: dict[str, Any],
|
|
) -> None:
|
|
"""Rebind GlobalContext-owned runtime services and the owned FileManager collaborators."""
|
|
|
|
context.parser = kwargs["parser"]
|
|
context.parser_factory = kwargs.get("parser_factory")
|
|
context.client = kwargs["client"]
|
|
context.prompt_formatter_builder = kwargs["prompt_formatter_builder"]
|
|
context.intent_check_factory = kwargs["intent_check_factory"]
|
|
context.repo_root = kwargs["repo_root"]
|
|
context.src_root = kwargs["src_root"]
|
|
context.test_root = kwargs["test_root"]
|
|
if context._file_manager is not None:
|
|
repo_exec = kwargs.get("repo_exec") or context.parser.repo_exec
|
|
context._file_manager.repo_root = context.repo_root
|
|
context._file_manager.src_root = context.src_root
|
|
context._file_manager.test_root = context.test_root
|
|
rehydrate = getattr(context._file_manager, "rehydrate", None)
|
|
if rehydrate is not None:
|
|
rehydrate(
|
|
repo_exec,
|
|
context.parser,
|
|
context.client,
|
|
context.prompt_formatter_builder,
|
|
)
|
|
_rebind_file_manager_dependency_runtime(context._file_manager)
|
|
|
|
def post_thaw(self, context: GlobalContext) -> None:
|
|
"""Restore enum values for a checkpoint with the current object shape."""
|
|
|
|
context.story_cursor.sovereign_phase = SovereignState.coerce(
|
|
context.story_cursor.sovereign_phase,
|
|
SovereignState.SUPREME_ACTIVE,
|
|
)
|
|
context.supreme_completion_state = SovereignState.coerce(
|
|
context.supreme_completion_state,
|
|
SovereignState.SUPREME_ACTIVE,
|
|
)
|
|
|
|
|
|
class BootstrapContextReassembly(Olaf[BootstrapContext]):
|
|
"""BootstrapContextReassembly restores runtime FileManager ownership without resyncing."""
|
|
|
|
def rebind_runtime(
|
|
self,
|
|
context: BootstrapContext,
|
|
checkpoint: LoadedCheckpoint,
|
|
kwargs: dict[str, Any],
|
|
) -> None:
|
|
"""Rebind the current global FileManager for a thawed bootstrap context."""
|
|
|
|
context.upstream_file_manager = kwargs["upstream_file_manager"]
|
|
_rebind_file_manager_dependency_runtime(context.upstream_file_manager)
|
|
|
|
def post_thaw(self, context: BootstrapContext) -> None:
|
|
"""Retain Bootstrap work history, defaulting older checkpoints to an empty record list."""
|
|
|
|
context.editor_records = list(getattr(context, "editor_records", ()))
|
|
|
|
|
|
class StabilizationContextReassembly(Olaf[StabilizationContext]):
|
|
"""StabilizationContextReassembly restores scoped FileManager wiring."""
|
|
|
|
def rebind_runtime(
|
|
self,
|
|
context: StabilizationContext,
|
|
checkpoint: LoadedCheckpoint,
|
|
kwargs: dict[str, Any],
|
|
) -> None:
|
|
"""Rebind the current scoped FileManager checkpoint."""
|
|
|
|
context.upstream_file_manager = kwargs["upstream_file_manager"]
|
|
_rebind_file_manager_dependency_runtime(context.upstream_file_manager)
|
|
if context.file_manager is not None:
|
|
context.file_manager.upstream_file_manager = context.upstream_file_manager
|
|
context.file_manager.rehydrate(
|
|
context.upstream_file_manager.repo_exec,
|
|
context.upstream_file_manager.parser,
|
|
context.upstream_file_manager.client,
|
|
context.upstream_file_manager.prompt_formatter_builder,
|
|
)
|
|
context.file_manager.dependency_graph_maintainer = (
|
|
context.upstream_file_manager.dependency_graph_maintainer
|
|
)
|
|
context.file_manager.dependency_strategy = (
|
|
context.upstream_file_manager.dependency_strategy
|
|
)
|
|
context.dependency_graph_maintainer = (
|
|
context.file_manager.dependency_graph_maintainer
|
|
)
|
|
else:
|
|
context.dependency_graph_maintainer = (
|
|
context.upstream_file_manager.dependency_graph_maintainer
|
|
)
|
|
|
|
def post_thaw(self, context: StabilizationContext) -> None:
|
|
"""Reconcile current-shape backlog state after runtime rebinding."""
|
|
|
|
if context.backlog is None:
|
|
context.backlog = BackLog(completion=context.completion)
|
|
elif context.backlog.completion is None:
|
|
context.backlog = BackLog(
|
|
completion=context.completion,
|
|
remediation=list(context.backlog.remediation),
|
|
inspection=list(context.backlog.inspection),
|
|
testing=list(context.backlog.testing),
|
|
)
|
|
else:
|
|
context.completion = context.backlog.completion
|
|
if context.backlog.judgement_service is None:
|
|
context.backlog.judgement_service = BackLog(
|
|
completion=context.completion
|
|
).judgement_service
|
|
if context.final_exam_started and not context.final_exam_finished:
|
|
context.backlog.clear_work_queues()
|
|
context.final_exam_started = False
|
|
context.final_exam_finished = False
|
|
context.fully_explored = False
|
|
|
|
|
|
class AlphaContextReassembly(Olaf[AlphaContext]):
|
|
"""AlphaContextReassembly restores runtime backlog wrapper and shared scoped FileManager."""
|
|
|
|
def rebind_runtime(
|
|
self,
|
|
context: AlphaContext,
|
|
checkpoint: LoadedCheckpoint,
|
|
kwargs: dict[str, Any],
|
|
) -> None:
|
|
"""Rebind Alpha runtime collaborators from the current stabilization scope."""
|
|
|
|
context.file_manager = kwargs["file_manager"]
|
|
context.backlog = kwargs["backlog"]
|
|
if "runtime_queue_depth_callback" in kwargs:
|
|
context.runtime_queue_depth_callback = kwargs["runtime_queue_depth_callback"]
|
|
if "alpha_postflight_callback" in kwargs:
|
|
context.alpha_postflight_callback = kwargs["alpha_postflight_callback"]
|
|
|
|
class BetaContextReassembly(Olaf[BetaContext]):
|
|
"""BetaContextReassembly restores runtime backlog wrapper, strategies, and shared FileManager."""
|
|
|
|
def rebind_runtime(
|
|
self,
|
|
context: BetaContext,
|
|
checkpoint: LoadedCheckpoint,
|
|
kwargs: dict[str, Any],
|
|
) -> None:
|
|
"""Rebind Beta runtime collaborators from the current stabilization scope."""
|
|
|
|
context.file_manager = kwargs["file_manager"]
|
|
context.backlog = kwargs["backlog"]
|
|
if "runtime_queue_depth_callback" in kwargs:
|
|
context.runtime_queue_depth_callback = kwargs["runtime_queue_depth_callback"]
|
|
|
|
|
|
class ContextFactory:
|
|
"""ContextFactory creates contexts and thaws at most one checkpoint per context type."""
|
|
|
|
_reassemblies: dict[type[ContextBag], Olaf[Any]] = {
|
|
AlphaContext: AlphaContextReassembly(),
|
|
BetaContext: BetaContextReassembly(),
|
|
BootstrapContext: BootstrapContextReassembly(),
|
|
GlobalContext: GlobalContextReassembly(),
|
|
StabilizationContext: StabilizationContextReassembly(),
|
|
}
|
|
|
|
def _latest_checkpoint(
|
|
self,
|
|
context_type: type[ContextBag],
|
|
) -> Path | None:
|
|
"""Return the newest checkpoint for a context type when one exists."""
|
|
|
|
settings = resume_settings()
|
|
repo_exec = settings.repo_exec
|
|
if repo_exec is None:
|
|
raise RuntimeError("Checkpoint discovery requires the authoritative RepoExec.")
|
|
checkpoint_dir = settings.checkpoint_dir
|
|
checkpoints = sorted(
|
|
path for path in repo_exec.list_directory(checkpoint_dir)
|
|
if path.name.startswith(f"{context_type.__name__}-") and path.name.endswith(".pkl")
|
|
) if repo_exec.is_dir(checkpoint_dir) else []
|
|
return checkpoints[-1] if checkpoints else None
|
|
|
|
def _load_checkpoint(self, checkpoint: Path) -> LoadedCheckpoint:
|
|
"""Load a checkpoint payload and schema version from disk."""
|
|
|
|
repo_exec = resume_settings().repo_exec
|
|
if repo_exec is None:
|
|
raise RuntimeError("Checkpoint loading requires the authoritative RepoExec.")
|
|
data = pickle.loads(repo_exec.read_bytes(checkpoint))
|
|
payload = data.get("payload", {})
|
|
version = data.get("checkpoint_version", 1)
|
|
return LoadedCheckpoint(
|
|
version=version if isinstance(version, int) else 1,
|
|
payload=payload if isinstance(payload, dict) else {},
|
|
)
|
|
|
|
def _load_payload(self, checkpoint: Path) -> dict[str, Any]:
|
|
"""Return only the payload for older callers that inspect checkpoints directly."""
|
|
|
|
return self._load_checkpoint(checkpoint).payload
|
|
|
|
def latest_payload_for(
|
|
self,
|
|
context_type: type[ContextBag],
|
|
) -> dict[str, Any] | None:
|
|
"""Return the latest checkpoint payload for one context type, when it exists."""
|
|
|
|
checkpoint = self._latest_checkpoint(context_type)
|
|
if checkpoint is None:
|
|
return None
|
|
return self._load_payload(checkpoint)
|
|
|
|
def _reassembly_for(self, context_type: type[T]) -> Olaf[T]:
|
|
"""Return the reassembly rule set for one context type."""
|
|
|
|
reassembly = self._reassemblies.get(context_type)
|
|
if reassembly is None:
|
|
return DefaultReassembly()
|
|
return reassembly
|
|
|
|
def _thaw(
|
|
self,
|
|
context: T,
|
|
checkpoint: LoadedCheckpoint,
|
|
kwargs: dict[str, Any],
|
|
) -> T | None:
|
|
"""Apply persisted attributes, rebind runtime ownership, and validate resume coordinates."""
|
|
|
|
if not _checkpoint_shape_is_current(type(context), checkpoint.payload):
|
|
return None
|
|
thawed = self._reassembly_for(type(context)).thaw(context, checkpoint, kwargs)
|
|
root_change = kwargs.get("root_change")
|
|
cot_index = kwargs.get("cot_index")
|
|
total_cot = kwargs.get("total_cot")
|
|
if root_change is not None:
|
|
if thawed.root_change.location != root_change.location:
|
|
raise ValueError(
|
|
f"Checkpoint root {thawed.root_change.location} does not match requested root {root_change.location}."
|
|
)
|
|
if cot_index is not None and thawed.cot_index != cot_index:
|
|
raise ValueError(
|
|
f"Checkpoint COT {thawed.cot_index} does not match requested COT {cot_index}."
|
|
)
|
|
if total_cot is not None:
|
|
if thawed.total_cot != total_cot:
|
|
raise ValueError(
|
|
f"Checkpoint total COT {thawed.total_cot} does not match requested total COT {total_cot}."
|
|
)
|
|
return thawed
|
|
|
|
def create(self, context_type: type[T], **kwargs: Any) -> T:
|
|
"""Create a fresh context from constructor arguments."""
|
|
|
|
return context_type(**kwargs)
|
|
|
|
def thaw(
|
|
self,
|
|
context_type: type[T],
|
|
**kwargs: Any,
|
|
) -> T | None:
|
|
"""Thaw current-shape state, returning ``None`` for discarded checkpoints."""
|
|
|
|
checkpoint = self._latest_checkpoint(context_type)
|
|
if checkpoint is None:
|
|
raise FileNotFoundError(f"No checkpoint found for {context_type.__name__}.")
|
|
|
|
context = context_type.__new__(context_type)
|
|
return self._thaw(context, self._load_checkpoint(checkpoint), kwargs)
|
|
|
|
context_factory = ContextFactory()
|