@openlucaskaka/kagent 0.1.7

package/src/kagent/runtime/approval.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import json
+from typing import Any, Dict
+
+from kagent.runtime.types import parse_agent_plan
+
+
+def build_resumable_plan(
+    plan: Dict[str, Any],
+    pending_approval: Dict[str, Any],
+) -> Dict[str, Any] | None:
+    """Return the pending action and every action that follows it.
+
+    Actions before the approval boundary already completed in the persisted run.
+    Their dependencies are therefore removed while dependencies between remaining
+    actions stay intact.
+    """
+
+    try:
+        parsed_plan = parse_agent_plan(
+            json.dumps(plan, ensure_ascii=False, sort_keys=True)
+        )
+    except (TypeError, ValueError):
+        return None
+
+    pending_action_id = str(pending_approval.get("id", "")).strip()
+    raw_actions = plan.get("actions")
+    if not pending_action_id or not isinstance(raw_actions, list):
+        return None
+
+    matching_indexes = [
+        index
+        for index, action in enumerate(raw_actions)
+        if isinstance(action, dict) and action.get("id") == pending_action_id
+    ]
+    if len(matching_indexes) != 1:
+        return None
+
+    pending_index = matching_indexes[0]
+    if raw_actions[pending_index] != pending_approval:
+        return None
+
+    completed_action_ids = {
+        action.id for action in parsed_plan.actions[:pending_index]
+    }
+    remaining_actions = []
+    for action in parsed_plan.actions[pending_index:]:
+        payload = action.to_dict()
+        remaining_dependencies = [
+            dependency
+            for dependency in action.depends_on
+            if dependency not in completed_action_ids
+        ]
+        if remaining_dependencies:
+            payload["depends_on"] = remaining_dependencies
+        else:
+            payload.pop("depends_on", None)
+        remaining_actions.append(payload)
+
+    resumable_plan: Dict[str, Any] = {"actions": remaining_actions}
+    if parsed_plan.final_answer:
+        resumable_plan["final_answer"] = parsed_plan.final_answer
+    return resumable_plan

package/src/kagent/runtime/cancellation.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from threading import Event, Lock
+from typing import Any, Callable, Dict, Mapping, Optional
+
+ExternalCancellationProbe = Callable[[], Optional[Mapping[str, Any]]]
+
+
+class RuntimeCancellationToken:
+    """Thread-safe cooperative cancellation shared across runtime boundaries."""
+
+    def __init__(
+        self,
+        *,
+        external_cancellation_probe: ExternalCancellationProbe | None = None,
+    ) -> None:
+        self._event = Event()
+        self._lock = Lock()
+        self._reason = ""
+        self._cancelled_at = ""
+        self._external_cancellation_probe = external_cancellation_probe
+
+    def cancel(self, reason: str = "", *, cancelled_at: str = "") -> bool:
+        normalized_reason = reason.strip()
+        normalized_cancelled_at = cancelled_at.strip()
+        with self._lock:
+            if self._event.is_set():
+                return False
+            self._reason = normalized_reason
+            self._cancelled_at = (
+                normalized_cancelled_at or datetime.now(timezone.utc).isoformat()
+            )
+            self._event.set()
+            return True
+
+    def is_cancelled(self) -> bool:
+        self._sync_external_cancellation()
+        return self._event.is_set()
+
+    def snapshot(self) -> Dict[str, str]:
+        self._sync_external_cancellation()
+        with self._lock:
+            return {
+                "cancelled": str(self._event.is_set()).lower(),
+                "reason": self._reason,
+                "cancelled_at": self._cancelled_at,
+            }
+
+    def _sync_external_cancellation(self) -> None:
+        if self._event.is_set() or self._external_cancellation_probe is None:
+            return
+        try:
+            cancellation = self._external_cancellation_probe()
+        except Exception:
+            return
+        if not cancellation:
+            return
+        reason = cancellation.get("reason", cancellation.get("cancel_reason", ""))
+        cancelled_at = cancellation.get("cancelled_at", "")
+        self.cancel(
+            str(reason) if reason is not None else "",
+            cancelled_at=str(cancelled_at) if cancelled_at is not None else "",
+        )

package/src/kagent/runtime/checkpoint_state.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import time
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, List, TypedDict
+
+from kagent.runtime.cancellation import RuntimeCancellationToken
+from kagent.runtime.policy import RuntimePolicy
+from kagent.runtime.redaction import redact_runtime_text
+from kagent.runtime.steering import RuntimeSteeringBuffer
+from kagent.runtime.tools import RuntimeToolSpec
+
+_CHECKPOINT_SECRET_KEY_PARTS = (
+    "api_key",
+    "apikey",
+    "authorization",
+    "bearer",
+    "password",
+    "secret",
+    "token",
+)
+
+RuntimeEventSink = Callable[[Dict[str, Any]], None]
+
+
+class RuntimeGraphState(TypedDict, total=False):
+    goal: str
+    run_id: str
+    max_iterations: int
+    approved_action_ids: List[str]
+    metadata: Dict[str, str]
+    tags: List[str]
+    started_at: str
+    initial_events: List[Dict[str, Any]]
+    initial_hook_failure_count: int
+    initial_planner: Dict[str, Any]
+    initial_progress_events: List[Dict[str, Any]]
+    initial_progress_event_sink_failure_count: int
+    initial_action_prepared: Dict[str, Any]
+    initial_action_phase: Dict[str, Any]
+    initial_action_outcome: Dict[str, Any]
+    result: Dict[str, Any]
+    graph_phases: List[Dict[str, str]]
+
+
+class RuntimeGraphContext(TypedDict, total=False):
+    goal: str
+    provider: Any
+    policy: RuntimePolicy
+    tools: Dict[str, RuntimeToolSpec]
+    cancellation_token: RuntimeCancellationToken
+    steering_buffer: RuntimeSteeringBuffer
+    event_sink: RuntimeEventSink
+    hooks: List[Any]
+    runtime_workspace_dir: str
+    redis_url: str
+    milvus_url: str
+    embedding_base_url: str
+    embedding_api_key: str
+    embedding_model: str
+    embedding_timeout_seconds: float
+    embedding_max_retries: int
+    embedding_retry_backoff_seconds: float
+    external_backend_timeout_seconds: float
+    stream_answers: bool
+    planner_plan_cache: Dict[str, Dict[str, Any]]
+    prepared_action_cache: Dict[str, Dict[str, Any]]
+    executing_action_cache: Dict[str, Dict[str, Any]]
+
+
+def checkpoint_safe_value(value: Any) -> Any:
+    if value is None or isinstance(value, (bool, int, float)):
+        return value
+    if isinstance(value, str):
+        return redact_runtime_text(value)
+    if isinstance(value, dict):
+        return {
+            redact_runtime_text(str(key)): checkpoint_safe_value(item)
+            for key, item in value.items()
+        }
+    if isinstance(value, (list, tuple)):
+        return [checkpoint_safe_value(item) for item in value]
+    return f"[unsupported {type(value).__name__}]"
+
+
+def checkpoint_plan_projection(value: Any, *, key: str = "") -> tuple[Any, bool]:
+    if key and any(part in key.lower() for part in _CHECKPOINT_SECRET_KEY_PARTS):
+        return "[REDACTED]", True
+    if value is None or isinstance(value, (bool, int, float)):
+        return value, False
+    if isinstance(value, str):
+        redacted = redact_runtime_text(value)
+        return redacted, redacted != value
+    if isinstance(value, dict):
+        projection: Dict[str, Any] = {}
+        changed = False
+        for raw_key, item in value.items():
+            projected_key = str(raw_key)
+            projected_item, item_changed = checkpoint_plan_projection(
+                item,
+                key=projected_key,
+            )
+            projection[projected_key] = projected_item
+            changed = changed or item_changed
+        return projection, changed
+    if isinstance(value, (list, tuple)):
+        projection = []
+        changed = False
+        for item in value:
+            projected_item, item_changed = checkpoint_plan_projection(item)
+            projection.append(projected_item)
+            changed = changed or item_changed
+        return projection, changed
+    return f"[unsupported {type(value).__name__}]", True
+
+
+def append_graph_phase(
+    phases: List[Dict[str, str]] | None,
+    node: str,
+    started_at: str,
+    started_timer: float,
+) -> List[Dict[str, str]]:
+    return [
+        *(phases or []),
+        {
+            "node": node,
+            "status": "ok",
+            **timing_fields(started_at, started_timer),
+        },
+    ]
+
+
+def utc_timestamp() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def timing_fields(started_at: str, started_timer: float) -> Dict[str, str]:
+    return {
+        "started_at": started_at,
+        "completed_at": utc_timestamp(),
+        "duration_seconds": duration_since(started_timer),
+    }
+
+
+def duration_since(started_at: float) -> str:
+    return f"{time.perf_counter() - started_at:.4f}"

package/src/kagent/runtime/checkpoint_storage.py

@@ -0,0 +1,270 @@
+from __future__ import annotations
+
+import hashlib
+import hmac
+import json
+import os
+import re
+import secrets
+import tempfile
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+DIRECTORY_MODE = 0o700
+FILE_MODE = 0o600
+MANIFEST_VERSION = 1
+_SIGNING_KEY_NAME = ".checkpoint-signing-key"
+_CHECKPOINT_ID_PATTERN = re.compile(
+    r"^\d{8}T\d{6}\.\d{6}Z-[0-9a-f]{12}$"
+)
+_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$")
+
+
+def new_checkpoint_id() -> str:
+    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ")
+    return f"{timestamp}-{secrets.token_hex(6)}"
+
+
+def validate_checkpoint_id(checkpoint_id: str) -> None:
+    if not isinstance(checkpoint_id, str) or not _CHECKPOINT_ID_PATTERN.fullmatch(
+        checkpoint_id
+    ):
+        raise ValueError("invalid checkpoint_id")
+
+
+def resolve_checkpoint_target(workspace_root: Path, relative_path: Any) -> Path:
+    if not isinstance(relative_path, str):
+        raise ValueError("checkpoint manifest is invalid")
+    candidate = Path(relative_path)
+    if candidate.is_absolute() or any(
+        part in {"", ".", ".."} for part in candidate.parts
+    ):
+        raise ValueError("checkpoint path must stay inside the workspace")
+    target = workspace_root.joinpath(*candidate.parts)
+    try:
+        target.relative_to(workspace_root)
+    except ValueError as exc:
+        raise ValueError("checkpoint path must stay inside the workspace") from exc
+    return target
+
+
+def ensure_store_layout(
+    state_root: Path,
+    workspace_directory: Path,
+    checkpoint_directory: Path,
+) -> None:
+    reject_existing_symlink_chain(state_root)
+    directories = (
+        state_root,
+        workspace_directory,
+        workspace_directory / "checkpoints",
+        checkpoint_directory,
+    )
+    for directory in directories:
+        if directory.exists() and directory.is_symlink():
+            raise ValueError("checkpoint path must not contain symlinks")
+        directory.mkdir(parents=True, exist_ok=True)
+        directory.chmod(DIRECTORY_MODE)
+
+
+def reject_existing_symlink_chain(path: Path) -> None:
+    absolute_path = path.absolute()
+    current = Path(absolute_path.anchor)
+    for part in absolute_path.parts[1:]:
+        current = current / part
+        if current.exists() and current.is_symlink():
+            raise ValueError("checkpoint path must not contain symlinks")
+
+
+def reject_symlink(path: Path) -> None:
+    if path.is_symlink():
+        raise ValueError("checkpoint path must not contain symlinks")
+
+
+def remove_checkpoint_directory(checkpoint_directory: Path) -> None:
+    reject_symlink(checkpoint_directory)
+    for path in checkpoint_directory.iterdir():
+        reject_symlink(path)
+        if not path.is_file():
+            raise ValueError("checkpoint directory contains an unexpected entry")
+        path.unlink()
+    checkpoint_directory.rmdir()
+
+
+def atomic_write_text(target: Path, content: str) -> None:
+    descriptor, temporary_name = tempfile.mkstemp(
+        prefix=f".{target.name}.",
+        suffix=".tmp",
+        dir=target.parent,
+    )
+    temporary_path = Path(temporary_name)
+    try:
+        os.chmod(temporary_path, FILE_MODE)
+        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
+            descriptor = -1
+            handle.write(content)
+            handle.flush()
+            os.fsync(handle.fileno())
+        temporary_path.replace(target)
+        _fsync_directory(target.parent)
+    except Exception:
+        if descriptor != -1:
+            os.close(descriptor)
+        temporary_path.unlink(missing_ok=True)
+        raise
+
+
+def write_signed_manifest(
+    target: Path,
+    manifest: dict[str, Any],
+    state_root: Path,
+) -> None:
+    signing_key = load_or_create_signing_key(state_root)
+    atomic_write_text(
+        target,
+        json.dumps(
+            signed_manifest_payload(manifest, signing_key),
+            ensure_ascii=False,
+            separators=(",", ":"),
+        ),
+    )
+
+
+def read_signed_manifest(
+    manifest_path: Path,
+    state_root: Path,
+    checkpoint_id: str,
+) -> dict[str, Any]:
+    try:
+        envelope = json.loads(manifest_path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError) as exc:
+        raise ValueError("checkpoint manifest is invalid") from exc
+    if (
+        not isinstance(envelope, dict)
+        or set(envelope) != {"manifest", "hmac_sha256"}
+        or not isinstance(envelope["manifest"], dict)
+        or not isinstance(envelope["hmac_sha256"], str)
+    ):
+        raise ValueError("checkpoint manifest is invalid")
+    manifest = envelope["manifest"]
+    expected_hmac = _manifest_hmac(manifest, read_signing_key(state_root))
+    if not hmac.compare_digest(envelope["hmac_sha256"], expected_hmac):
+        raise ValueError("checkpoint manifest integrity check failed")
+    _validate_manifest(manifest, checkpoint_id)
+    return manifest
+
+
+def signed_manifest_payload(
+    manifest: dict[str, Any],
+    signing_key: bytes,
+) -> dict[str, Any]:
+    return {
+        "manifest": manifest,
+        "hmac_sha256": _manifest_hmac(manifest, signing_key),
+    }
+
+
+def load_or_create_signing_key(state_root: Path) -> bytes:
+    key_path = state_root / _SIGNING_KEY_NAME
+    reject_symlink(key_path)
+    if key_path.exists():
+        return read_signing_key(state_root)
+    signing_key = secrets.token_bytes(32)
+    try:
+        descriptor = os.open(
+            key_path,
+            os.O_WRONLY | os.O_CREAT | os.O_EXCL,
+            FILE_MODE,
+        )
+    except FileExistsError:
+        return read_signing_key(state_root)
+    try:
+        os.write(descriptor, signing_key)
+        os.fsync(descriptor)
+    finally:
+        os.close(descriptor)
+    _fsync_directory(state_root)
+    return signing_key
+
+
+def read_signing_key(state_root: Path) -> bytes:
+    key_path = state_root / _SIGNING_KEY_NAME
+    reject_existing_symlink_chain(state_root)
+    reject_symlink(key_path)
+    if not key_path.is_file():
+        raise ValueError("checkpoint signing key is missing")
+    signing_key = key_path.read_bytes()
+    if len(signing_key) != 32:
+        raise ValueError("checkpoint signing key is invalid")
+    return signing_key
+
+
+def _manifest_hmac(manifest: dict[str, Any], signing_key: bytes) -> str:
+    canonical = json.dumps(
+        manifest,
+        ensure_ascii=False,
+        separators=(",", ":"),
+        sort_keys=True,
+    ).encode("utf-8")
+    return hmac.new(signing_key, canonical, hashlib.sha256).hexdigest()
+
+
+def _validate_manifest(manifest: Any, checkpoint_id: str) -> None:
+    if (
+        not isinstance(manifest, dict)
+        or manifest.get("version") != MANIFEST_VERSION
+        or manifest.get("checkpoint_id") != checkpoint_id
+        or manifest.get("status") not in {"prepared", "committed"}
+        or not isinstance(manifest.get("created_at"), str)
+        or not isinstance(manifest.get("entries"), list)
+        or not manifest["entries"]
+    ):
+        raise ValueError("checkpoint manifest is invalid")
+    paths = []
+    for entry in manifest["entries"]:
+        if not isinstance(entry, dict) or set(entry) != {"path", "before", "after"}:
+            raise ValueError("checkpoint manifest is invalid")
+        if not isinstance(entry["path"], str):
+            raise ValueError("checkpoint manifest is invalid")
+        _validate_state(entry["before"])
+        _validate_state(entry["after"])
+        paths.append(entry["path"])
+    if len(paths) != len(set(paths)):
+        raise ValueError("checkpoint manifest contains duplicate paths")
+
+
+def _validate_state(state: Any) -> None:
+    if not isinstance(state, dict) or set(state) != {
+        "exists",
+        "sha256",
+        "mode",
+        "blob",
+    }:
+        raise ValueError("checkpoint manifest is invalid")
+    exists = state["exists"]
+    mode = state["mode"]
+    blob = state["blob"]
+    sha256 = state["sha256"]
+    if not isinstance(exists, bool):
+        raise ValueError("checkpoint manifest is invalid")
+    if isinstance(mode, bool) or not isinstance(mode, int) or not 0 <= mode <= 0o7777:
+        raise ValueError("checkpoint manifest is invalid")
+    if not isinstance(blob, str) or not isinstance(sha256, str):
+        raise ValueError("checkpoint manifest is invalid")
+    if exists:
+        if Path(blob).name != blob or not blob or not _SHA256_PATTERN.fullmatch(sha256):
+            raise ValueError("checkpoint manifest is invalid")
+    elif blob or sha256:
+        raise ValueError("checkpoint manifest is invalid")
+
+
+def _fsync_directory(directory: Path) -> None:
+    flags = os.O_RDONLY
+    if hasattr(os, "O_DIRECTORY"):
+        flags |= os.O_DIRECTORY
+    descriptor = os.open(directory, flags)
+    try:
+        os.fsync(descriptor)
+    finally:
+        os.close(descriptor)

package/src/kagent/runtime/context.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict
+
+
+@dataclass
+class RuntimeContextStats:
+    truncated_string_count: int = 0
+    truncated_chars: int = 0
+    omitted_artifact_content_count: int = 0
+
+
+class RuntimeContextManager:
+    def __init__(self, *, max_string_chars: int = 500) -> None:
+        if max_string_chars < 1:
+            raise ValueError("max_string_chars must be positive")
+        self.max_string_chars = max_string_chars
+        self._stats = RuntimeContextStats()
+
+    def compact_observation_output(self, output: Any) -> Any:
+        if isinstance(output, dict) and str(output.get("artifact_id", "")).strip():
+            self._stats.omitted_artifact_content_count += 1
+            return self._artifact_prompt_metadata(output)
+        return self.compact_value(output)
+
+    def compact_value(self, value: Any) -> Any:
+        if isinstance(value, str):
+            if len(value) <= self.max_string_chars:
+                return value
+            self._stats.truncated_string_count += 1
+            truncated_chars = len(value) - self.max_string_chars
+            self._stats.truncated_chars += truncated_chars
+            return {
+                "text_prefix": value[: self.max_string_chars],
+                "original_chars": len(value),
+                "truncated_chars": truncated_chars,
+            }
+        if isinstance(value, dict):
+            return {key: self.compact_value(item) for key, item in value.items()}
+        if isinstance(value, list):
+            return [self.compact_value(item) for item in value]
+        return value
+
+    def report(self) -> Dict[str, Any]:
+        return {
+            "strategy": "runtime_context_manager",
+            "artifact_content_omitted": True,
+            "max_string_chars": str(self.max_string_chars),
+            "long_string_shape": "text_prefix/original_chars/truncated_chars",
+            "truncated_string_count": str(self._stats.truncated_string_count),
+            "truncated_chars": str(self._stats.truncated_chars),
+            "omitted_artifact_content_count": str(
+                self._stats.omitted_artifact_content_count
+            ),
+        }
+
+    def _artifact_prompt_metadata(self, output: Dict[str, Any]) -> Dict[str, Any]:
+        metadata = {
+            key: output[key]
+            for key in ["artifact_id", "title", "kind", "format", "tags", "bytes"]
+            if key in output
+        }
+        metadata["content_omitted"] = True
+        return metadata