@openlucaskaka/kagent 0.1.7

package/src/kagent/runtime/patch_checkpoints.py

@@ -0,0 +1,385 @@
+from __future__ import annotations
+
+import hashlib
+import os
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Mapping
+
+from kagent.runtime.checkpoint_storage import (
+    FILE_MODE,
+    MANIFEST_VERSION,
+    atomic_write_text,
+    ensure_store_layout,
+    new_checkpoint_id,
+    read_signed_manifest,
+    reject_existing_symlink_chain,
+    reject_symlink,
+    remove_checkpoint_directory,
+    resolve_checkpoint_target,
+    validate_checkpoint_id,
+    write_signed_manifest,
+)
+from kagent.runtime.file_transaction import (
+    TextFileState,
+    capture_text_states,
+    commit_text_changes,
+    validate_workspace_targets,
+    workspace_transaction,
+)
+from kagent.utils.paths import kagent_state_dir, migrate_legacy_kagent_state
+
+
+@dataclass(frozen=True)
+class PatchRevert:
+    checkpoint_id: str
+    expected_current: dict[Path, str | None]
+    staged_contents: dict[Path, str | None]
+    target_modes: dict[Path, int]
+
+
+class PatchCheckpointStore:
+    def __init__(self, state_root: str | Path) -> None:
+        self.state_root = Path(state_root).expanduser()
+
+    @classmethod
+    def from_environment(
+        cls,
+        env: Mapping[str, str] | None = None,
+    ) -> "PatchCheckpointStore":
+        source = os.environ if env is None else env
+        configured = source.get("KAGENT_PATCH_STATE_DIR", "").strip()
+        if configured:
+            return cls(configured)
+        migrate_legacy_kagent_state(source)
+        return cls(kagent_state_dir(source) / "patches")
+
+    def prepare(
+        self,
+        workspace_root: Path,
+        before: dict[Path, TextFileState],
+        after: dict[Path, str | None],
+    ) -> str:
+        if set(before) != set(after):
+            raise ValueError("checkpoint before and after paths must match")
+        workspace_root = workspace_root.resolve()
+        checkpoint_id = new_checkpoint_id()
+        checkpoint_directory = self._checkpoint_directory(
+            workspace_root,
+            checkpoint_id,
+        )
+        ensure_store_layout(
+            self.state_root,
+            self._workspace_directory(workspace_root),
+            checkpoint_directory,
+        )
+        try:
+            entries = []
+            for index, target in enumerate(after):
+                relative_path = target.relative_to(workspace_root).as_posix()
+                before_state = before[target]
+                after_content = after[target]
+                entries.append(
+                    {
+                        "path": relative_path,
+                        "before": self._write_state(
+                            checkpoint_directory,
+                            f"before-{index:04d}.txt",
+                            before_state.content,
+                            before_state.mode,
+                        ),
+                        "after": self._write_state(
+                            checkpoint_directory,
+                            f"after-{index:04d}.txt",
+                            after_content,
+                            before_state.mode
+                            if before_state.content is not None
+                            else FILE_MODE,
+                        ),
+                    }
+                )
+            manifest = {
+                "version": MANIFEST_VERSION,
+                "checkpoint_id": checkpoint_id,
+                "created_at": datetime.now(timezone.utc).isoformat(),
+                "status": "prepared",
+                "entries": entries,
+            }
+            self._write_manifest(checkpoint_directory, manifest)
+        except Exception:
+            remove_checkpoint_directory(checkpoint_directory)
+            raise
+        return checkpoint_id
+
+    def mark_committed(self, workspace_root: Path, checkpoint_id: str) -> None:
+        checkpoint_directory, manifest = self._load_manifest(
+            workspace_root,
+            checkpoint_id,
+        )
+        if manifest.get("status") != "prepared":
+            raise ValueError("checkpoint is not prepared")
+        manifest["status"] = "committed"
+        self._write_manifest(checkpoint_directory, manifest)
+
+    def discard(self, workspace_root: Path, checkpoint_id: str) -> None:
+        checkpoint_directory = self._checkpoint_directory(
+            workspace_root.resolve(),
+            checkpoint_id,
+        )
+        if not checkpoint_directory.exists():
+            return
+        remove_checkpoint_directory(checkpoint_directory)
+
+    def history(self, workspace_root: Path, *, limit: int = 20) -> dict[str, Any]:
+        if limit < 1:
+            raise ValueError("limit must be positive")
+        reject_existing_symlink_chain(self.state_root)
+        checkpoints_directory = (
+            self._workspace_directory(workspace_root.resolve()) / "checkpoints"
+        )
+        if not checkpoints_directory.exists():
+            return {"checkpoints": [], "checkpoint_count": 0}
+        reject_symlink(checkpoints_directory)
+        checkpoints = []
+        for directory in sorted(checkpoints_directory.iterdir(), reverse=True):
+            if len(checkpoints) >= limit:
+                break
+            if directory.is_symlink() or not directory.is_dir():
+                continue
+            try:
+                _, manifest = self._load_manifest(workspace_root, directory.name)
+            except ValueError:
+                continue
+            if manifest.get("status") != "committed":
+                continue
+            entries = manifest["entries"]
+            checkpoints.append(
+                {
+                    "checkpoint_id": manifest["checkpoint_id"],
+                    "created_at": manifest["created_at"],
+                    "file_count": len(entries),
+                    "paths": [entry["path"] for entry in entries],
+                }
+            )
+        return {"checkpoints": checkpoints, "checkpoint_count": len(checkpoints)}
+
+    def load_revert(self, workspace_root: Path, checkpoint_id: str) -> PatchRevert:
+        workspace_root = workspace_root.resolve()
+        checkpoint_directory, manifest = self._load_manifest(
+            workspace_root,
+            checkpoint_id,
+        )
+        if manifest.get("status") != "committed":
+            raise ValueError("checkpoint is not committed")
+        return self._revert_from_manifest(
+            workspace_root,
+            checkpoint_directory,
+            manifest,
+        )
+
+    def recover_prepared(self, workspace_root: Path) -> int:
+        workspace_root = workspace_root.resolve()
+        checkpoints_directory = self._workspace_directory(workspace_root) / "checkpoints"
+        if not checkpoints_directory.exists():
+            return 0
+        recovered = 0
+        with workspace_transaction(workspace_root):
+            for directory in sorted(checkpoints_directory.iterdir()):
+                if directory.is_symlink() or not directory.is_dir():
+                    continue
+                _, manifest = self._load_manifest(workspace_root, directory.name)
+                if manifest.get("status") != "prepared":
+                    continue
+                revert = self._revert_from_manifest(
+                    workspace_root,
+                    directory,
+                    manifest,
+                )
+                validate_workspace_targets(workspace_root, revert.staged_contents)
+                current_states = capture_text_states(revert.staged_contents)
+                for target, current_state in current_states.items():
+                    allowed_contents = {
+                        revert.expected_current[target],
+                        revert.staged_contents[target],
+                    }
+                    if current_state.content not in allowed_contents:
+                        relative_path = target.relative_to(workspace_root).as_posix()
+                        raise ValueError(
+                            f"checkpoint recovery SHA-256 conflict: {relative_path}"
+                        )
+                if any(
+                    current_states[target].content != revert.staged_contents[target]
+                    for target in revert.staged_contents
+                ):
+                    commit_text_changes(
+                        workspace_root,
+                        revert.staged_contents,
+                        original_states=current_states,
+                        target_modes=revert.target_modes,
+                    )
+                self.discard(workspace_root, directory.name)
+                recovered += 1
+        return recovered
+
+    def _revert_from_manifest(
+        self,
+        workspace_root: Path,
+        checkpoint_directory: Path,
+        manifest: dict[str, Any],
+    ) -> PatchRevert:
+        expected_current = {}
+        staged_contents = {}
+        target_modes = {}
+        for entry in manifest["entries"]:
+            target = resolve_checkpoint_target(workspace_root, entry["path"])
+            expected_current[target] = self._read_state(
+                checkpoint_directory,
+                entry["after"],
+            )
+            staged_contents[target] = self._read_state(
+                checkpoint_directory,
+                entry["before"],
+            )
+            target_modes[target] = int(entry["before"]["mode"])
+        return PatchRevert(
+            checkpoint_id=manifest["checkpoint_id"],
+            expected_current=expected_current,
+            staged_contents=staged_contents,
+            target_modes=target_modes,
+        )
+
+    def _workspace_directory(self, workspace_root: Path) -> Path:
+        identity = hashlib.sha256(str(workspace_root).encode("utf-8")).hexdigest()
+        return self.state_root / identity
+
+    def _checkpoint_directory(
+        self,
+        workspace_root: Path,
+        checkpoint_id: str,
+    ) -> Path:
+        validate_checkpoint_id(checkpoint_id)
+        return self._workspace_directory(workspace_root) / "checkpoints" / checkpoint_id
+
+    def _load_manifest(
+        self,
+        workspace_root: Path,
+        checkpoint_id: str,
+    ) -> tuple[Path, dict[str, Any]]:
+        checkpoint_directory = self._checkpoint_directory(
+            workspace_root.resolve(),
+            checkpoint_id,
+        )
+        reject_existing_symlink_chain(self.state_root)
+        manifest_path = checkpoint_directory / "manifest.json"
+        reject_symlink(checkpoint_directory)
+        reject_symlink(manifest_path)
+        if not manifest_path.is_file():
+            raise ValueError("checkpoint does not exist")
+        manifest = read_signed_manifest(
+            manifest_path,
+            self.state_root,
+            checkpoint_id,
+        )
+        return checkpoint_directory, manifest
+
+    def _write_manifest(
+        self,
+        checkpoint_directory: Path,
+        manifest: dict[str, Any],
+    ) -> None:
+        write_signed_manifest(
+            checkpoint_directory / "manifest.json",
+            manifest,
+            self.state_root,
+        )
+
+    def _write_state(
+        self,
+        checkpoint_directory: Path,
+        blob_name: str,
+        content: str | None,
+        mode: int,
+    ) -> dict[str, Any]:
+        if content is None:
+            return {"exists": False, "sha256": "", "mode": mode, "blob": ""}
+        encoded = content.encode("utf-8")
+        atomic_write_text(checkpoint_directory / blob_name, content)
+        return {
+            "exists": True,
+            "sha256": hashlib.sha256(encoded).hexdigest(),
+            "mode": mode,
+            "blob": blob_name,
+        }
+
+    def _read_state(
+        self,
+        checkpoint_directory: Path,
+        state: Any,
+    ) -> str | None:
+        if not isinstance(state, dict) or not isinstance(state.get("exists"), bool):
+            raise ValueError("checkpoint manifest is invalid")
+        if state["exists"] is False:
+            return None
+        blob_name = state.get("blob")
+        expected_sha256 = state.get("sha256")
+        if not isinstance(blob_name, str) or Path(blob_name).name != blob_name:
+            raise ValueError("checkpoint manifest is invalid")
+        blob_path = checkpoint_directory / blob_name
+        reject_symlink(blob_path)
+        if not blob_path.is_file():
+            raise ValueError("checkpoint content is missing")
+        content = blob_path.read_text(encoding="utf-8")
+        actual_sha256 = hashlib.sha256(content.encode("utf-8")).hexdigest()
+        if actual_sha256 != expected_sha256:
+            raise ValueError("checkpoint content SHA-256 mismatch")
+        return content
+
+
+def commit_checkpointed_text_changes(
+    store: PatchCheckpointStore,
+    workspace_root: Path,
+    staged_contents: dict[Path, str | None],
+    *,
+    target_modes: dict[Path, int] | None = None,
+) -> str:
+    original_states = capture_text_states(staged_contents)
+    checkpoint_id = store.prepare(workspace_root, original_states, staged_contents)
+    files_committed = False
+    try:
+        commit_text_changes(
+            workspace_root,
+            staged_contents,
+            original_states=original_states,
+            target_modes=target_modes,
+        )
+        files_committed = True
+        store.mark_committed(workspace_root, checkpoint_id)
+    except Exception as commit_error:
+        rollback_error = None
+        if files_committed:
+            try:
+                commit_text_changes(
+                    workspace_root,
+                    {
+                        target: state.content
+                        for target, state in original_states.items()
+                    },
+                    target_modes={
+                        target: state.mode for target, state in original_states.items()
+                    },
+                )
+            except Exception as exc:
+                rollback_error = exc
+        try:
+            store.discard(workspace_root, checkpoint_id)
+        except Exception as discard_error:
+            if rollback_error is None:
+                rollback_error = discard_error
+        if rollback_error is not None:
+            raise RuntimeError(
+                f"checkpoint commit failed: {commit_error}; "
+                f"rollback failed: {rollback_error}"
+            ) from commit_error
+        raise
+    return checkpoint_id

package/src/kagent/runtime/policy.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Set
+
+
+@dataclass(frozen=True)
+class PolicyDecision:
+    status: str
+    reason: str = ""
+
+
+class RuntimePolicy:
+    def __init__(self, *, allowed_tools: Optional[Set[str]] = None) -> None:
+        self.allowed_tools = set(
+            allowed_tools
+            or {
+                "apply_patch",
+                "artifact",
+                "delegate_task",
+                "decision_matrix",
+                "list_files",
+                "memory_get",
+                "memory_put",
+                "memory_recall",
+                "memory_remember",
+                "memory_search",
+                "memory_upsert",
+                "note",
+                "patch_history",
+                "read_file",
+                "rubric_score",
+                "skill_get",
+                "skill_list",
+                "task_list",
+                "task_transition",
+                "transform_text",
+                "workspace_diff",
+                "workspace_history",
+                "workspace_list",
+                "workspace_read",
+                "workspace_search",
+                "workspace_write",
+            }
+        )
+
+    def authorize(self, tool: str, _input_payload: Dict[str, Any]) -> PolicyDecision:
+        if tool not in self.allowed_tools:
+            return PolicyDecision(status="denied", reason="tool_not_allowed")
+        return PolicyDecision(status="allowed")

package/src/kagent/runtime/presentation.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from kagent.runtime.redaction import redact_runtime_text
+
+MAX_PRESENTATION_CONTENT_CHARS = 4000
+
+
+def project_runtime_presentation(
+    tool: str,
+    status: str,
+    output: Dict[str, Any],
+) -> Dict[str, Any]:
+    if status != "ok" or not isinstance(output, dict):
+        return {}
+
+    projectors = {
+        "artifact": _project_artifact,
+        "apply_patch": _project_apply_patch,
+        "revert_patch": _project_revert_patch,
+        "workspace_diff": _project_workspace_diff,
+        "workspace_restore": _project_workspace_restore,
+        "open_url": _project_open_url,
+        "open_app": _project_open_app,
+        "http_request": _project_http_request,
+        "shell_command": _project_shell_command,
+    }
+    projector = projectors.get(tool)
+    if projector is None:
+        return {}
+    return projector(output)
+
+
+def _project_artifact(output: Dict[str, Any]) -> Dict[str, Any]:
+    title = _string(output.get("title"))
+    content = _string(output.get("content"))
+    if not title or not content:
+        return {}
+    kind = _display_label(output.get("kind"))
+    artifact_format = _display_label(output.get("format"))
+    byte_count = output.get("bytes")
+    bytes_label = f"{byte_count} bytes" if isinstance(byte_count, int) else ""
+    detail = " · ".join(
+        part for part in (kind, artifact_format, bytes_label) if part
+    )
+    visible, truncated = _bounded_content(content)
+    return _presentation_with_content(
+        f"Created {title}",
+        detail,
+        visible,
+        truncated,
+    )
+
+
+def _project_apply_patch(output: Dict[str, Any]) -> Dict[str, Any]:
+    changed_files = output.get("changed_files")
+    if not isinstance(changed_files, list) or not changed_files:
+        return {}
+    paths = []
+    for item in changed_files:
+        if not isinstance(item, dict):
+            continue
+        path = _string(item.get("path"))
+        if path:
+            paths.append(path)
+    if not paths:
+        return {}
+    count = len(paths)
+    visible_paths = paths[:3]
+    path_detail = ", ".join(visible_paths)
+    if count > len(visible_paths):
+        path_detail += f", +{count - len(visible_paths)} more"
+    label = "file" if count == 1 else "files"
+    return _presentation("Updated files", f"{count} {label}: {path_detail}")
+
+
+def _project_revert_patch(output: Dict[str, Any]) -> Dict[str, Any]:
+    projected = _project_apply_patch(output)
+    if not projected:
+        return {}
+    projected["title"] = "Restored files"
+    return projected
+
+
+def _project_workspace_diff(output: Dict[str, Any]) -> Dict[str, Any]:
+    path = _string(output.get("path"))
+    diff = _string(output.get("diff"))
+    if not path or not diff:
+        return {}
+    kind = _string(output.get("kind"))
+    detail = f"{kind}/{path}" if kind else path
+    visible, truncated = _bounded_content(
+        diff,
+        already_truncated=output.get("truncated") is True,
+    )
+    return _presentation_with_content(
+        "Workspace changes",
+        detail,
+        visible,
+        truncated,
+    )
+
+
+def _project_workspace_restore(output: Dict[str, Any]) -> Dict[str, Any]:
+    path = _string(output.get("path"))
+    revision_id = _string(output.get("restored_revision_id"))
+    if not path or not revision_id:
+        return {}
+    kind = _string(output.get("kind"))
+    detail = f"{kind}/{path}" if kind else path
+    return _presentation("Restored workspace asset", detail)
+
+
+def _project_open_url(output: Dict[str, Any]) -> Dict[str, Any]:
+    url = _string(output.get("url"))
+    if output.get("opened") is not True or not url:
+        return {}
+    return _presentation("Opened URL", url)
+
+
+def _project_open_app(output: Dict[str, Any]) -> Dict[str, Any]:
+    application = _string(output.get("application"))
+    if output.get("opened") is not True or not application:
+        return {}
+    return _presentation("Opened application", application)
+
+
+def _project_http_request(output: Dict[str, Any]) -> Dict[str, Any]:
+    url = _string(output.get("url"))
+    status_code = output.get("status_code")
+    if not url or not isinstance(status_code, int):
+        return {}
+    content_type = _string(output.get("content_type"))
+    detail = " · ".join(
+        part for part in (str(status_code), content_type, url) if part
+    )
+    return _presentation("Fetched URL", detail)
+
+
+def _project_shell_command(output: Dict[str, Any]) -> Dict[str, Any]:
+    exit_code = output.get("exit_code")
+    if not isinstance(exit_code, int):
+        return {}
+    duration = output.get("duration_seconds")
+    detail_parts = [f"Exit {exit_code}"]
+    if isinstance(duration, (int, float)):
+        detail_parts.append(f"{duration}s")
+    stdout = _string(output.get("stdout"))
+    stderr = _string(output.get("stderr"))
+    content = "\n".join(part for part in (stdout, stderr) if part)
+    visible, truncated = _bounded_content(
+        content,
+        already_truncated=output.get("truncated") is True,
+    )
+    return _presentation_with_content(
+        "Command completed",
+        " · ".join(detail_parts),
+        visible,
+        truncated,
+    )
+
+
+def _presentation(title: str, detail: str) -> Dict[str, Any]:
+    return {
+        "title": redact_runtime_text(title),
+        "detail": redact_runtime_text(detail),
+    }
+
+
+def _presentation_with_content(
+    title: str,
+    detail: str,
+    content: str,
+    truncated: bool,
+) -> Dict[str, Any]:
+    return {
+        **_presentation(title, detail),
+        "content": content,
+        "truncated": truncated,
+    }
+
+
+def _bounded_content(
+    content: str,
+    *,
+    already_truncated: bool = False,
+) -> tuple[str, bool]:
+    redacted = redact_runtime_text(content)
+    was_bounded = len(redacted) > MAX_PRESENTATION_CONTENT_CHARS
+    return (
+        redacted[:MAX_PRESENTATION_CONTENT_CHARS],
+        already_truncated or was_bounded,
+    )
+
+
+def _display_label(value: Any) -> str:
+    text = _string(value)
+    if not text:
+        return ""
+    return " ".join(part.capitalize() for part in text.replace("_", " ").split())
+
+
+def _string(value: Any) -> str:
+    return value if isinstance(value, str) else ""