agentengine-sdk-python 0.8.5__py3-none-any.whl → 0.8.6__py3-none-any.whl

ksadk/toolsets/workspace_editing.py

@@ -0,0 +1,291 @@
+"""Workspace edit validation, read-state reconciliation and bounded diff evidence.
+
+Filesystem authorization and tool entry points remain in ``workspace``.
+"""
+
+from __future__ import annotations
+
+import difflib
+import hashlib
+from dataclasses import replace
+from pathlib import Path
+from typing import Any
+
+from ksadk.tools.result_budget import budget_tool_output, default_tool_result_budget
+from ksadk.toolsets.workspace_mutation import publish_workspace_bytes
+from ksadk.toolsets.workspace_state import WorkspaceReadState, record_read_state
+
+
+def _dict_result(value: Any, *, tool_name: str) -> dict[str, Any]:
+    if not isinstance(value, dict):
+        raise TypeError(f"{tool_name} returned {type(value).__name__}, expected dict")
+    return dict(value)
+
+
+def _read_workspace_text(target: Path) -> tuple[str | None, dict[str, Any] | None]:
+    if not target.is_file():
+        return None, {"ok": False, "error_message": "workspace file not found"}
+    if target.stat().st_size > 2 * 1024 * 1024:
+        return None, {"ok": False, "error_message": "file is larger than 2MB"}
+    try:
+        # Editing must retain the original separators and UTF-8 BOM. The
+        # model's line-oriented read view is not the file's canonical bytes.
+        with target.open("rb") as source:
+            payload = source.read(2 * 1024 * 1024 + 1)
+        if len(payload) > 2 * 1024 * 1024:
+            return None, {"ok": False, "error_message": "file is larger than 2MB"}
+        return payload.decode("utf-8"), None
+    except UnicodeDecodeError:
+        return None, {"ok": False, "error_message": "file is not utf-8 text"}
+
+
+def _validate_read_state(
+    target: Path,
+    relative: str,
+    *,
+    state: WorkspaceReadState | None,
+    content: str | None = None,
+) -> dict[str, Any] | None:
+    if (
+        state is None
+        or state.canonical_path != str(target.resolve())
+        or not state.content_sha256
+    ):
+        return {
+            "ok": False,
+            "path": relative,
+            "error_type": "file_not_read",
+            "error_message": "file must be read before editing",
+            "suggested_action": f"Call read_workspace_file(path={relative!r}) before editing.",
+        }
+    stat = target.stat() if target.exists() else None
+    if (
+        stat is None
+        or stat.st_mtime_ns != state.mtime_ns
+        or stat.st_size != state.size_bytes
+        or (content is not None
+            and hashlib.sha256(content.encode("utf-8")).hexdigest() != state.content_sha256)
+    ):
+        return {
+            "ok": False,
+            "path": relative,
+            "error_type": "file_modified_since_read",
+            "error_message": "workspace file changed since last read",
+            "suggested_action": (
+                f"Re-read the file with read_workspace_file(path={relative!r}) and retry."
+            ),
+        }
+    return None
+
+
+def _commit_workspace_edit(
+    target: Path, relative: str, updated: str, *, state: WorkspaceReadState
+) -> tuple[Any | None, dict[str, Any] | None]:
+    """Recheck the pinned source immediately before writing, never a new read state.
+
+    This detects intervening changes during matching/diff generation. It is
+    not a filesystem CAS against unrelated writers after the final check.
+    """
+    current, error = _read_workspace_text(target)
+    if error is not None:
+        return None, {**error, "path": relative}
+    if current is None:
+        return None, {"ok": False, "path": relative, "error_message": "workspace file read failed"}
+    error = _validate_read_state(target, relative, state=state, content=current)
+    if error is not None:
+        return None, error
+    publish_workspace_bytes(target, updated.encode("utf-8"))
+    stat = target.stat()
+    _refresh_read_state(relative, stat, updated, target=target)
+    return stat, None
+
+
+def _apply_single_edit(
+    text: str,
+    old_text: Any,
+    new_text: Any,
+    *,
+    expected_replacements: int = 1,
+    replace_all: bool = False,
+    relative: str,
+) -> dict[str, Any]:
+    request_error = _validate_single_edit_request(old_text, relative=relative)
+    if request_error is not None:
+        return request_error
+    # Match LF/CRLF snippets against the file's convention without converting
+    # any unrelated bytes. In particular, never normalize the whole file.
+    ending = "\r\n" if "\r\n" in text else "\n"
+    old = str(old_text).replace("\r\n", "\n").replace("\n", ending)
+    replacement = str(new_text or "").replace("\r\n", "\n").replace("\n", ending)
+    expected = max(1, int(expected_replacements or 1))
+    actual_old, used_quote_normalization = _find_actual_string(text, old)
+    matches = _match_previews(text, actual_old or old) if actual_old else []
+    occurrences = text.count(actual_old) if actual_old else 0
+    if occurrences == 0:
+        return {
+            "ok": False,
+            "path": relative,
+            "error_type": "snippet_not_found",
+            "error_message": "old_text was not found in the workspace file",
+            "occurrences": 0,
+            "nearby_candidates": _nearby_candidates(text, old),
+            "suggested_action": (
+                "Use one nearby candidate as old_text, or call read_workspace_file "
+                "with line numbers and retry."
+            ),
+        }
+    if actual_old is None:
+        return {
+            "ok": False,
+            "path": relative,
+            "error_type": "snippet_not_found",
+            "error_message": "old_text was not found in the workspace file",
+        }
+    if not replace_all and occurrences != expected:
+        return {
+            "ok": False,
+            "path": relative,
+            "error_type": "ambiguous_edit",
+            "error_message": f"old_text matched {occurrences} occurrences, expected {expected}",
+            "occurrences": occurrences,
+            "expected_replacements": expected,
+            "matches": matches,
+            "suggested_action": (
+                "Provide a larger unique old_text snippet, set expected_replacements, "
+                "or use replace_all=true."
+            ),
+        }
+    replacements = occurrences if replace_all else expected
+    return {
+        "ok": True,
+        "updated": text.replace(actual_old, replacement, replacements),
+        "replacements": replacements,
+        "used_quote_normalization": used_quote_normalization,
+        "matches": matches,
+    }
+
+
+def _validate_single_edit_request(old_text: Any, *, relative: str) -> dict[str, Any] | None:
+    if not str(old_text or ""):
+        return {
+            "ok": False,
+            "path": relative,
+            "error_type": "missing_old_text",
+            "error_message": "old_text is required",
+        }
+    return None
+
+
+def _match_previews(text: str, needle: str, limit: int = 20) -> list[dict[str, Any]]:
+    if not needle:
+        return []
+    matches: list[dict[str, Any]] = []
+    for line_no, line in enumerate(text.splitlines(), start=1):
+        if needle in line:
+            matches.append({"line": line_no, "text": line.strip(), "snippet": line})
+            if len(matches) >= limit:
+                break
+    return matches
+
+
+def _nearby_candidates(text: str, old: str, limit: int = 5) -> list[dict[str, Any]]:
+    old_norm = _normalize_quotes(old).strip()
+    candidates: list[dict[str, Any]] = []
+    for line_no, line in enumerate(text.splitlines(), start=1):
+        ratio = difflib.SequenceMatcher(None, old_norm, _normalize_quotes(line).strip()).ratio()
+        if old_norm and (
+            old_norm[: max(1, min(4, len(old_norm)))] in _normalize_quotes(line) or ratio > 0.45
+        ):
+            candidates.append(
+                {"line": line_no, "text": line.strip(), "snippet": line, "score": round(ratio, 3)}
+            )
+    candidates.sort(key=lambda item: (-float(item["score"]), int(item["line"])))
+    return candidates[:limit]
+
+
+def _build_diff(relative: str, before: str, after: str) -> str:
+    return "".join(
+        difflib.unified_diff(
+            before.splitlines(keepends=True),
+            after.splitlines(keepends=True),
+            fromfile=f"{relative} before",
+            tofile=f"{relative} after",
+        )
+    )
+
+
+def _budget_diff(tool_name: str, diff: str, relative: str) -> dict[str, Any]:
+    # The model supplied the edit, so echoing a whole-file diff back into its
+    # next context wastes the same budget it needs to verify and repair. Keep
+    # the full diff in a distinct audit file and return a bounded preview.
+    default = default_tool_result_budget()
+    budget = replace(
+        default,
+        max_chars=min(default.max_chars, 4096),
+        preview_chars=min(default.preview_chars, 2048),
+        persist_threshold_chars=min(default.persist_threshold_chars, 4096),
+    )
+    diff_digest = hashlib.sha256(diff.encode("utf-8")).hexdigest()
+    path_digest = hashlib.sha256(relative.encode("utf-8")).hexdigest()[:12]
+    result = _dict_result(
+        budget_tool_output(
+            tool_name=tool_name,
+            field_name="diff",
+            value=diff,
+            metadata={"tool_use_id": f"{tool_name}_{path_digest}_{diff_digest[:16]}"},
+            budget=budget,
+        ),
+        tool_name=tool_name,
+    )
+    result["diff_sha256"] = diff_digest
+    result["diff_line_count"] = diff.count("\n")
+    if result["truncated"]:
+        result["suggested_action"] = (
+            "Edit succeeded. Read the affected file and run relevant checks; "
+            "the full diff is retained at the persisted path for audit."
+        )
+    return result
+
+
+def _refresh_read_state(relative: str, stat: Any, text: str, *, target: Path) -> None:
+    lines = text.splitlines()
+    record_read_state(
+        WorkspaceReadState(
+            path=relative,
+            mtime_ns=stat.st_mtime_ns,
+            size_bytes=stat.st_size,
+            start_line=1,
+            end_line=len(lines),
+            total_lines=len(lines),
+            partial=False,
+            canonical_path=str(target.resolve()),
+            content_sha256=hashlib.sha256(text.encode("utf-8")).hexdigest(),
+        )
+    )
+
+
+def _find_actual_string(text: str, old: str) -> tuple[str | None, bool]:
+    if old in text:
+        return old, False
+    normalized_old = _normalize_quotes(old)
+    smart_quote_map: dict[str, str | int | None] = {'"': "“", "'": "‘"}
+    for candidate in {normalized_old, old.translate(str.maketrans(smart_quote_map))}:
+        if candidate and candidate in text:
+            return candidate, True
+    normalized_text = _normalize_quotes(text)
+    index = normalized_text.find(normalized_old)
+    if index < 0:
+        return None, False
+    return text[index : index + len(old)], True
+
+
+def _normalize_quotes(value: str) -> str:
+    quote_map: dict[str, str | int | None] = {
+        "“": '"',
+        "”": '"',
+        "„": '"',
+        "‘": "'",
+        "’": "'",
+        "‚": "'",
+    }
+    return value.translate(str.maketrans(quote_map))

ksadk/toolsets/workspace_model_observation.py

@@ -0,0 +1,104 @@
+"""Source-reader model view; the full SDK result remains the audit record.
+
+Only host-selected builtin readers use this projection. Source text, visible
+range, omissions, continuation and overflow locators stay intact. Internal
+accounting and duplicate display counters need not compete with the source
+for the model's window. Unknown fields are preserved, never silently dropped.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+
+WORKSPACE_SOURCE_MODEL_VERSION = "workspace-source-model/v1"
+_AUDIT_FIELDS = frozenset(
+    {"size", "size_bytes", "mtime_ns", "source_output_budget", "original_chars", "preview_chars"}
+)
+
+
+def source_model_reply(result: object) -> str | None:
+    """Return a detached, structured view of a recognized successful page.
+
+    No file access, source extraction, interpretation or acceptance occurs.
+    Malformed/failed results keep the existing raw projection instead.
+    """
+    if (
+        not isinstance(result, Mapping)
+        or result.get("ok") is not True
+        or not isinstance(result.get("path"), str)
+        or not result["path"]
+        or not isinstance(result.get("content"), str)
+    ):
+        return None
+    start, end, total, count = (
+        result.get(key) for key in ("start_line", "end_line", "total_lines", "line_count")
+    )
+    carrier = result.get("requested_source")
+    if result.get("artifact_only") is True:
+        if (
+            not isinstance(carrier, Mapping) or carrier.get("path") != result["path"]
+            or not isinstance(carrier.get("content"), str) or not carrier["content"]
+            or not isinstance(carrier.get("read_range"), dict)
+            or type(start) is not int or type(total) is not int or not 1 <= start <= total
+            or type(count) is not int or count != 0
+            or end is not None or result.get("read_range") is not None
+            or result["content"] != "" or result.get("truncated") is not True
+            or result.get("partial") is not True
+        ):
+            return None
+        requested = carrier["read_range"]
+        if (
+            set(requested) != {"start", "end"}
+            or any(type(v) is not int for v in requested.values())
+            or requested["start"] != start or not start <= requested["end"] <= total
+        ):
+            return None
+        return json.dumps({"status": "ok", "content": {
+            "ok": True, "path": result["path"], "source_state": "unread",
+            "source_read_range": requested,
+        }}, ensure_ascii=False, separators=(",", ":"))
+    if any(type(value) is not int for value in (start, total, count)):
+        return None
+    if start < 1 or total < 0 or count < 0:
+        return None
+    read_range = result.get("read_range")
+    if count:
+        if (
+            type(end) is not int
+            or end != start + count - 1
+            or end > total
+            or not isinstance(read_range, dict)
+            or any(type(read_range.get(key)) is not int for key in ("start", "end"))
+            or read_range != {"start": start, "end": end}
+        ):
+            return None
+    elif total:
+        if (
+            start > total
+            or end is not None
+            or read_range is not None
+            or result.get("partial_line") != start
+            or result.get("truncated") is not True
+            or not result["content"]
+        ):
+            return None
+    elif end != 0 or result["content"] or start != 1 or read_range != {"start": 1, "end": 0}:
+        return None
+    if any(type(result.get(key)) is not bool for key in ("partial", "truncated")):
+        return None
+    hidden = _AUDIT_FIELDS
+    if "requested_source" in result:
+        # The opt-in carrier crosses the host boundary, but never the model
+        # window. SDK filesystem locators cannot grant Native artifact access.
+        hidden = hidden | {"requested_source", "persisted"}
+    view = {k: v for k, v in result.items() if k not in hidden}
+    if "requested_source" in result:
+        view["suggested_action"] = (
+            "Inspect the complete requested source through the host artifact reference."
+        )
+    return json.dumps(
+        {"status": "ok", "content": view},
+        ensure_ascii=False,
+        separators=(",", ":"),
+    )

ksadk/toolsets/workspace_mutation.py

@@ -0,0 +1,75 @@
+"""Local file mutation primitives; neither a Run ledger nor a distributed lock.
+
+SDK callers in one process serialize by canonical path. Builtin workers retain
+their host's SQLite serialization. Unrelated external writers do not obey these
+locks: atomic publication prevents partial files, not arbitrary filesystem CAS.
+"""
+
+from __future__ import annotations
+
+import os
+import secrets
+import stat
+from contextlib import ExitStack, contextmanager
+from pathlib import Path
+from threading import Lock, RLock
+from weakref import WeakValueDictionary
+
+_LOCKS: WeakValueDictionary[str, RLock] = WeakValueDictionary()
+_LOCKS_GUARD = Lock()
+
+
+@contextmanager
+def locked_workspace_paths(paths):
+    """Acquire a deterministic set of local path locks, including nested calls."""
+    keys = sorted({os.path.normcase(str(Path(path).resolve())) for path in paths})
+    # Hold strong references before acquiring any lock. Otherwise an unused
+    # entry could disappear while another caller prepares the same set.
+    with _LOCKS_GUARD:
+        locks = []
+        for key in keys:
+            lock = _LOCKS.get(key)
+            if lock is None:
+                lock = RLock()
+                _LOCKS[key] = lock
+            locks.append(lock)
+    with ExitStack() as stack:
+        for lock in locks:
+            stack.enter_context(lock)
+        yield
+
+
+def publish_workspace_bytes(target: Path, payload: bytes, *, create_only: bool = False) -> None:
+    """Publish a complete sibling file; never fall back to truncating the target.
+
+    The caller must hold its local mutation lock and validate the observed
+    revision. Create-only publication cannot replace a racing new file. A
+    failed replace leaves the original content and read-state unchanged.
+    """
+    target.parent.mkdir(parents=True, exist_ok=True)
+    mode = None
+    try:
+        info = target.lstat()
+    except FileNotFoundError:
+        pass
+    else:
+        if not stat.S_ISREG(info.st_mode):
+            raise ValueError("workspace mutation requires a regular file")
+        mode = stat.S_IMODE(info.st_mode)
+    temporary = target.parent / f".{target.name}.ksadk-{secrets.token_hex(12)}"
+    # Preserve ordinary new-file permissions (0666 filtered by the process
+    # umask) without reading/changing the process-wide umask in other threads.
+    descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o666)
+    try:
+        with os.fdopen(descriptor, "wb") as stream:
+            stream.write(payload)
+            if mode is not None:
+                os.chmod(temporary, mode)
+            stream.flush()
+            os.fsync(stream.fileno())
+        if create_only:
+            os.link(temporary, target)
+        else:
+            os.replace(temporary, target)
+    finally:
+        temporary.unlink(missing_ok=True)

ksadk/toolsets/workspace_read_budget.py

@@ -0,0 +1,216 @@
+"""Optional host budget for a source page, independent of Agent frameworks."""
+
+from __future__ import annotations
+
+import json
+import math
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Callable
+
+from ksadk.conversations.model_context import estimate_text_tokens
+from ksadk.tools.artifact_observation import (
+    ARCHIVE_MODEL_POLICY,
+    ARCHIVE_SEARCH_MODEL_POLICY,
+    ARCHIVE_URI_MAX_CHARS,
+    CANONICAL_ARCHIVE_MODEL_POLICY,
+    OBSERVED_ARCHIVE_MODEL_POLICY,
+    SOURCE_ARCHIVE_MODEL_POLICY,
+    archived_observation,
+)
+from ksadk.toolsets.workspace_model_observation import (
+    WORKSPACE_SOURCE_MODEL_VERSION,
+    source_model_reply,
+)
+
+SOURCE_COUNTER_ID = "studio-cjk-ascii-estimate-125/v1"
+
+
+def estimated_source_units(text: str) -> int:
+    return max(1, math.ceil(estimate_text_tokens(text) * 1.25))
+
+
+@dataclass(frozen=True)
+class SourceReadBudget:
+    max_units: int
+    count_units: Callable[[str], int]
+    counter_id: str
+    max_result_units: int | None = None
+    result_projection: str | None = None
+    retain_requested_source: bool = False
+    artifact_only_fallback: bool = False
+    archive_model_policy: str = ARCHIVE_MODEL_POLICY
+
+    def __post_init__(self):
+        if type(self.max_units) is not int or self.max_units < 1:
+            raise ValueError("source_read_budget_unavailable")
+        if not callable(self.count_units) or not self.counter_id.strip():
+            raise ValueError("source_read_counter_unavailable")
+        if self.max_result_units is not None and (
+            type(self.max_result_units) is not int or self.max_result_units < 1
+        ):
+            raise ValueError("source_read_budget_unavailable")
+        if self.result_projection not in (None, WORKSPACE_SOURCE_MODEL_VERSION):
+            raise ValueError("source_read_projection_unavailable")
+        if type(self.retain_requested_source) is not bool or (
+            self.retain_requested_source
+            and self.result_projection != WORKSPACE_SOURCE_MODEL_VERSION
+        ):
+            raise ValueError("source_read_archive_policy_invalid")
+        if type(self.artifact_only_fallback) is not bool or (
+            self.artifact_only_fallback and not self.retain_requested_source
+        ):
+            raise ValueError("source_read_archive_policy_invalid")
+        if self.archive_model_policy not in {
+            ARCHIVE_MODEL_POLICY, ARCHIVE_SEARCH_MODEL_POLICY, CANONICAL_ARCHIVE_MODEL_POLICY,
+            OBSERVED_ARCHIVE_MODEL_POLICY, SOURCE_ARCHIVE_MODEL_POLICY,
+        } or (
+            self.archive_model_policy in {
+                ARCHIVE_SEARCH_MODEL_POLICY, CANONICAL_ARCHIVE_MODEL_POLICY,
+                OBSERVED_ARCHIVE_MODEL_POLICY, SOURCE_ARCHIVE_MODEL_POLICY,
+            }
+            and not (self.retain_requested_source and self.artifact_only_fallback)
+        ):
+            raise ValueError("source_read_archive_policy_invalid")
+
+    def character_limit(self, text: str, requested_limit: int) -> int:
+        def fits(length):
+            fragment = text[:length]
+            cost = self.count_units(fragment)
+            if type(cost) is not int or cost < 0 or (fragment and cost == 0):
+                raise ValueError("source_read_counter_invalid")
+            return cost <= self.max_units
+
+        high = min(len(text), requested_limit)
+        if fits(high):
+            return max(1, high)
+        low = 0
+        while low < high:
+            middle = (low + high + 1) // 2
+            if fits(middle):
+                low = middle
+            else:
+                high = middle - 1
+        if not low:
+            raise ValueError("source_read_budget_unavailable")
+        return low
+
+    def result_fits(self, result: dict) -> bool:
+        if self.max_result_units is None:
+            return True
+        # A JSON reply is quoted inside the model tool-message text. Charge
+        # that escaping as well as the page metadata, not just source text.
+        text = (
+            source_model_reply(result)
+            if self.result_projection == WORKSPACE_SOURCE_MODEL_VERSION
+            else None
+        )
+        if text is None:
+            text = json.dumps({"status": "ok", "content": result}, ensure_ascii=False)
+        if self.retain_requested_source:
+            # The URI is issued only after persistence. Reserve the maximum
+            # declared locator here, including its complete model framing.
+            text = archived_observation(
+                text,
+                "artifact://" + "a" * (ARCHIVE_URI_MAX_CHARS - 11),
+                policy=self.archive_model_policy,
+            )
+        cost = self.count_units(json.dumps(text, ensure_ascii=False))
+        if type(cost) is not int or cost <= 0:
+            raise ValueError("source_read_counter_invalid")
+        return cost <= self.max_result_units
+
+    def transport(self) -> dict:
+        """Only a versioned host counter can cross an isolated SDK boundary."""
+        if self.counter_id != SOURCE_COUNTER_ID:
+            raise ValueError("source_read_counter_transport_unavailable")
+        result = {
+            "version": "source-read-budget/v2"
+            if self.result_projection is None
+            else "source-read-budget/v3",
+            "counter_id": self.counter_id,
+            "max_units": self.max_units,
+            "max_result_units": self.max_result_units,
+        }
+        if self.result_projection is not None:
+            result["result_projection"] = self.result_projection
+        if self.retain_requested_source:
+            result.update(version="source-read-budget/v4", retain_requested_source=True)
+        if self.artifact_only_fallback:
+            result.update(version="source-read-budget/v5", artifact_only_fallback=True)
+        if self.archive_model_policy == ARCHIVE_SEARCH_MODEL_POLICY:
+            result.update(
+                version="source-read-budget/v6", archive_model_policy=self.archive_model_policy
+            )
+        if self.archive_model_policy == CANONICAL_ARCHIVE_MODEL_POLICY:
+            result.update(
+                version="source-read-budget/v7", archive_model_policy=self.archive_model_policy
+            )
+        versions = {
+            OBSERVED_ARCHIVE_MODEL_POLICY: "source-read-budget/v8",
+            SOURCE_ARCHIVE_MODEL_POLICY: "source-read-budget/v9",
+        }
+        if self.archive_model_policy in versions:
+            result.update(version=versions[self.archive_model_policy],
+                          archive_model_policy=self.archive_model_policy)
+        return result
+
+    @classmethod
+    def from_transport(cls, value: object) -> SourceReadBudget | None:
+        if value is None:
+            return None
+        if not isinstance(value, dict):
+            raise ValueError("source_read_budget_transport_invalid")
+        version = value.get("version")
+        known = {f"source-read-budget/v{i}": i for i in range(2, 10)}
+        rank = known.get(version) if isinstance(version, str) else None
+        if rank is None:
+            raise ValueError("source_read_budget_transport_invalid")
+        expected = {"version", "counter_id", "max_units", "max_result_units"}
+        if rank >= 3:
+            expected.add("result_projection")
+        if rank >= 4:
+            expected.add("retain_requested_source")
+        if rank >= 5:
+            expected.add("artifact_only_fallback")
+        if rank >= 6:
+            expected.add("archive_model_policy")
+        policies = {
+            6: ARCHIVE_SEARCH_MODEL_POLICY, 7: CANONICAL_ARCHIVE_MODEL_POLICY,
+            8: OBSERVED_ARCHIVE_MODEL_POLICY, 9: SOURCE_ARCHIVE_MODEL_POLICY,
+        }
+        if (
+            set(value) != expected or value["counter_id"] != SOURCE_COUNTER_ID
+            or (rank >= 3 and value["result_projection"] != WORKSPACE_SOURCE_MODEL_VERSION)
+            or (rank >= 4 and value["retain_requested_source"] is not True)
+            or (rank >= 5 and value["artifact_only_fallback"] is not True)
+            or (rank >= 6 and value["archive_model_policy"] != policies[rank])
+        ):
+            raise ValueError("source_read_budget_transport_invalid")
+        return cls(
+            value["max_units"],
+            estimated_source_units,
+            SOURCE_COUNTER_ID,
+            value["max_result_units"],
+            value.get("result_projection"),
+            value.get("retain_requested_source", False),
+            value.get("artifact_only_fallback", False),
+            value.get("archive_model_policy", ARCHIVE_MODEL_POLICY),
+        )
+
+
+_READ_BUDGET: ContextVar[SourceReadBudget | None] = ContextVar("source_read_budget", default=None)
+
+
+@contextmanager
+def source_read_budget_scope(budget: SourceReadBudget | None):
+    token = _READ_BUDGET.set(budget)
+    try:
+        yield
+    finally:
+        _READ_BUDGET.reset(token)
+
+
+def current_source_read_budget() -> SourceReadBudget | None:
+    return _READ_BUDGET.get()