aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/context/metadata.py

@@ -0,0 +1,169 @@
+"""Safe compaction metadata contracts for observability.
+
+Provides the CompactionMetadata dataclass and helpers to build
+safe A2A/SSE event metadata for compaction lifecycle events.
+
+Authors:
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+_COMPACTION_ACTIVITY_ID = "context_compaction"
+_AGENT_THINKING_STEP_KIND = "agent_thinking_step"
+
+_STAGE_MESSAGES = {
+    "started": "Compacting context...",
+    "completed": "Context compacted",
+    "skipped": "Compaction skipped",
+    "failed": "Context compaction failed",
+}
+
+
+@dataclass(frozen=True)
+class CompactionMetadata:
+    """Safe observable metadata for a single compaction cycle.
+
+    All string fields are safe for logging, traces, and stream events.
+    Raw payloads, secrets, signed URLs, and base64 blobs are excluded
+    by construction.
+
+    Attributes:
+        summary_id: Opaque identifier for the compaction summary.
+        tokens_before: Estimated prompt tokens before compaction.
+        tokens_after: Estimated prompt tokens after compaction.
+        offloaded_ref_count: Number of payload references offloaded.
+        preservation_topic: Optional topic preserved during compaction.
+    """
+
+    summary_id: str | None
+    tokens_before: int | None
+    tokens_after: int | None
+    offloaded_ref_count: int
+    preservation_topic: str | None
+
+
+def build_compaction_event_metadata(
+    *,
+    stage: str,
+    trigger_reason: str | None = None,
+    summary_status: str | None = None,
+    tokens_before: int | None = None,
+    tokens_after: int | None = None,
+    offloaded_ref_count: int = 0,
+    summary_id: str | None = None,
+    preservation_topic: str | None = None,
+) -> dict[str, Any]:
+    """Build safe A2A/SSE metadata dict for a compaction lifecycle event.
+
+    Args:
+        stage: Compaction lifecycle stage (started, completed, skipped, failed).
+        trigger_reason: Why compaction was triggered.
+        summary_status: Status of the summary generation.
+        tokens_before: Estimated tokens before compaction.
+        tokens_after: Estimated tokens after compaction.
+        offloaded_ref_count: Number of offloaded payload references.
+        summary_id: Opaque summary identifier.
+        preservation_topic: Optional topic preserved during compaction.
+
+    Returns:
+        Safe metadata dict for A2A/SSE status_update events.
+        Guarantees no raw payloads, secrets, signed URLs, or base64 blobs.
+    """
+    metadata: dict[str, Any] = {
+        "kind": _AGENT_THINKING_STEP_KIND,
+        "step_type": "context_compaction",
+        "stage": stage,
+    }
+
+    if summary_status:
+        metadata["summary_status"] = summary_status
+    if tokens_before is not None:
+        metadata["tokens_before"] = tokens_before
+    if tokens_after is not None:
+        metadata["tokens_after"] = tokens_after
+    if tokens_before is not None and tokens_after is not None:
+        metadata["estimated_token_savings"] = max(0, tokens_before - tokens_after)
+    if offloaded_ref_count:
+        metadata["offloaded_ref_count"] = offloaded_ref_count
+    if summary_id:
+        metadata["summary_id"] = summary_id
+    if trigger_reason:
+        metadata["trigger_reason"] = trigger_reason
+    if preservation_topic:
+        metadata["preservation_topic"] = preservation_topic
+
+    metadata["thinking_and_activity_info"] = _build_compaction_activity_info(metadata)
+    return metadata
+
+
+def _build_compaction_activity_info(metadata: dict[str, Any]) -> dict[str, str]:
+    """Build the existing glchat sidebar activity payload shape."""
+    payload: dict[str, Any] = {
+        "type": "context_compaction",
+        "message": safe_compaction_content(str(metadata.get("stage", ""))),
+        "stage": metadata.get("stage", ""),
+    }
+    for key in (
+        "summary_status",
+        "trigger_reason",
+        "tokens_before",
+        "tokens_after",
+        "estimated_token_savings",
+        "offloaded_ref_count",
+        "summary_id",
+        "preservation_topic",
+    ):
+        if key in metadata:
+            payload[key] = metadata[key]
+
+    return {
+        "data_type": "activity",
+        "id": _COMPACTION_ACTIVITY_ID,
+        "data_value": json.dumps(payload, separators=(",", ":")),
+    }
+
+
+def safe_compaction_content(stage: str) -> str:
+    """Return a safe, short content string for a compaction stream event.
+
+    Args:
+        stage: Compaction lifecycle stage.
+
+    Returns:
+        Short non-sensitive status text safe for stream event content.
+    """
+    return _STAGE_MESSAGES.get(stage, "Context compaction")
+
+
+def build_compaction_metadata(
+    state: Any,
+    *,
+    tokens_before: int | None = None,
+    tokens_after: int | None = None,
+    preservation_topic: str | None = None,
+    summary_id_for_state: str | None = None,
+) -> CompactionMetadata:
+    """Build a CompactionMetadata from a CompactionState and extra params.
+
+    Args:
+        state: CompactionState from the context preparation result.
+        tokens_before: Estimated prompt tokens before compaction (overrides state).
+        tokens_after: Estimated prompt tokens after compaction (overrides state).
+        preservation_topic: Optional topic preserved during compaction.
+        summary_id_for_state: Optional summary ID from result.summary.summary_id.
+
+    Returns:
+        CompactionMetadata with safe fields only.
+    """
+    return CompactionMetadata(
+        summary_id=summary_id_for_state or getattr(state, "summary_id", None),
+        tokens_before=tokens_before,
+        tokens_after=tokens_after,
+        offloaded_ref_count=getattr(state, "offloaded_ref_count", 0),
+        preservation_topic=preservation_topic,
+    )

aip_agents/context/offload.py

@@ -0,0 +1,449 @@
+"""Helpers for lossless payload offloading and safe context pruning.
+
+Authors:
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+import copy
+import hashlib
+import json
+import re
+import uuid
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:  # pragma: no cover - only for type checking
+    from aip_agents.middleware.backends.protocol import BackendProtocol
+
+
+def _estimate_tokens(text: str, chars_per_token: int = 4) -> int:
+    # deferred to #809: Use shared aip_agents.context.utils.estimate_tokens
+    if not text:
+        return 0
+    return max(1, (len(text) + chars_per_token - 1) // chars_per_token)
+
+
+SourceKind = Literal["tool_input", "tool_result", "subagent_output", "attachment", "artifact"]
+
+_SIGNED_URL_PATTERN = re.compile(r"(?i)(x-amz-signature|sig|signature)=")
+_SIGNED_QUERY_KV_PATTERN = re.compile(r"(?i)\b(x-amz-signature|sig|signature)\s*=\s*[^&\s;]+")
+_SECRET_KV_PATTERN = re.compile(r"(?i)(api[_-]?key|token|password|secret)\s*[:=]\s*[^\s]+")
+_LONG_BASE64_PATTERN = re.compile(r"(?=[^+/=]*[+/=])[A-Za-z0-9+/=]{120,}")
+
+
+@dataclass(frozen=True)
+class OffloadedPayloadRef:
+    """Reference metadata for an offloaded payload.
+
+    Attributes:
+        ref_id: Stable identifier returned by the offload backend.
+        source_kind: Payload origin used for audit and downstream handling.
+        backend_uri: Backend-specific location for retrieving the full payload.
+        display_name: Safe human-readable name for the payload, if available.
+        mime_type: Payload MIME type, if known.
+        size_bytes: Original payload size in bytes, if known.
+        token_estimate: Estimated token cost of the original payload, if known.
+        checksum: Integrity checksum for the original payload, if available.
+        preview: Redacted inline preview used in the transformed message.
+        redaction_applied: Whether the preview was changed by secret redaction.
+    """
+
+    ref_id: str
+    source_kind: SourceKind
+    backend_uri: str
+    display_name: str | None
+    mime_type: str | None
+    size_bytes: int | None
+    token_estimate: int | None
+    checksum: str | None
+    preview: str
+    redaction_applied: bool
+
+
+def offload_oversized_payloads(
+    messages: Sequence[dict[str, Any]],
+    *,
+    backend: BackendProtocol | None = None,
+    max_inline_chars: int = 4_000,
+    preview_chars: int = 400,
+) -> tuple[list[dict[str, Any]], list[OffloadedPayloadRef]]:
+    """Offload oversized or sensitive payload content from messages.
+
+    Message order and tool-call IDs are preserved.
+
+    Usage::
+
+        from aip_agents.context.offload import offload_oversized_payloads
+
+        messages = [
+            {"role": "tool", "name": "search", "content": "x" * 5000},
+        ]
+        transformed, refs = offload_oversized_payloads(
+            messages, backend=backend, max_inline_chars=4000
+        )
+        # transformed[0]["content"] contains a preview + ref link
+        # refs[0] has the OffloadedPayloadRef metadata
+    """
+    transformed: list[dict[str, Any]] = []
+    refs: list[OffloadedPayloadRef] = []
+
+    for message in messages:
+        msg = copy.deepcopy(message)
+        content = _extract_text_content(msg)
+
+        content_ref = _offload_content_field(
+            msg,
+            content,
+            backend=backend,
+            max_inline_chars=max_inline_chars,
+            preview_chars=preview_chars,
+        )
+        if content_ref is not None:
+            refs.append(content_ref)
+
+        refs.extend(
+            _offload_tool_call_arguments(
+                msg,
+                backend=backend,
+                max_inline_chars=max_inline_chars,
+                preview_chars=preview_chars,
+            )
+        )
+        transformed.append(msg)
+
+    return transformed, refs
+
+
+def _offload_content_field(
+    message: dict[str, Any],
+    content: str,
+    *,
+    backend: BackendProtocol | None,
+    max_inline_chars: int,
+    preview_chars: int,
+) -> OffloadedPayloadRef | None:
+    if not _should_offload_content(content, max_inline_chars=max_inline_chars):
+        return None
+    if backend is None:
+        return None
+
+    ref = offload_payload(
+        content,
+        source_kind=_infer_source_kind(message),
+        display_name=_infer_display_name(message),
+        backend=backend,
+        preview_chars=preview_chars,
+    )
+    if ref is None:
+        return None
+    message["content"] = render_offloaded_reference_message(ref)
+    message["offloaded_ref_id"] = ref.ref_id
+    return ref
+
+
+def _offload_tool_call_arguments(
+    message: dict[str, Any],
+    *,
+    backend: BackendProtocol | None,
+    max_inline_chars: int,
+    preview_chars: int,
+) -> list[OffloadedPayloadRef]:
+    tool_calls = message.get("tool_calls")
+    if not isinstance(tool_calls, list):
+        return []
+
+    refs: list[OffloadedPayloadRef] = []
+    for call in tool_calls:
+        if not isinstance(call, dict):
+            continue
+        arguments = _tool_call_arguments(call)
+        if arguments is None or not _should_offload_content(arguments, max_inline_chars=max_inline_chars):
+            continue
+        if backend is None:
+            continue
+        ref = offload_payload(
+            arguments,
+            source_kind="tool_input",
+            display_name=_tool_call_name(call) or _infer_display_name(message),
+            backend=backend,
+            preview_chars=preview_chars,
+        )
+        if ref is None:
+            continue
+        _replace_tool_call_arguments(
+            call, json.dumps({"context_payload_ref": ref.ref_id, "source": "tool_input", "preview": ref.preview})
+        )
+        refs.append(ref)
+
+    if refs:
+        message["offloaded_tool_call_ref_ids"] = [ref.ref_id for ref in refs]
+    return refs
+
+
+def _tool_call_arguments(call: dict[str, Any]) -> str | None:
+    function = call.get("function")
+    if isinstance(function, dict) and isinstance(function.get("arguments"), str):
+        return function["arguments"]
+    args = call.get("args")
+    return args if isinstance(args, str) else None
+
+
+def _replace_tool_call_arguments(call: dict[str, Any], replacement: str) -> None:
+    function = call.get("function")
+    if isinstance(function, dict) and isinstance(function.get("arguments"), str):
+        function["arguments"] = replacement
+        return
+    if isinstance(call.get("args"), str):
+        call["args"] = replacement
+
+
+def _tool_call_name(call: dict[str, Any]) -> str | None:
+    if isinstance(call.get("name"), str):
+        return call["name"]
+    function = call.get("function")
+    if isinstance(function, dict) and isinstance(function.get("name"), str):
+        return function["name"]
+    return None
+
+
+def offload_payload(
+    payload_text: str,
+    *,
+    source_kind: SourceKind,
+    display_name: str | None,
+    backend: BackendProtocol | None = None,
+    preview_chars: int = 400,
+    mime_type: str | None = "text/plain",
+) -> OffloadedPayloadRef | None:
+    """Create an offloaded payload reference and optional backend write.
+
+    Returns None when no backend is provided or the backend write fails.
+    """
+    ref_id = f"ctx_ref_{uuid.uuid4().hex[:12]}"
+    checksum = hashlib.sha256(payload_text.encode("utf-8")).hexdigest()
+    size_bytes = len(payload_text.encode("utf-8"))
+    preview, redaction_applied = _safe_preview(payload_text, max_chars=preview_chars)
+
+    if backend is None:
+        return None
+
+    backend_path = f"/context-offloads/{ref_id}.txt"
+    try:
+        write_result = backend.write(backend_path, payload_text)
+    except Exception:
+        return None
+
+    if not getattr(write_result, "success", False):
+        return None
+
+    token_estimate = _estimate_tokens(payload_text)
+    return OffloadedPayloadRef(
+        ref_id=ref_id,
+        source_kind=source_kind,
+        backend_uri=backend_path,
+        display_name=display_name,
+        mime_type=mime_type,
+        size_bytes=size_bytes,
+        token_estimate=token_estimate,
+        checksum=checksum,
+        preview=preview,
+        redaction_applied=redaction_applied,
+    )
+
+
+def render_offloaded_reference_message(ref: OffloadedPayloadRef) -> str:
+    """Render model-visible replacement content for an offloaded payload."""
+    return (
+        "[Context payload offloaded]\n"
+        f"source: {ref.source_kind}\n"
+        f"ref: {ref.ref_id}\n"
+        f"name: {ref.display_name or 'unknown'}\n"
+        f"size: {ref.size_bytes or 0} bytes\n"
+        "preview:\n"
+        f"{ref.preview}\n"
+        "Use the runtime/backend ref to retrieve details if needed."
+    )
+
+
+def prune_duplicate_tool_outputs(
+    messages: Sequence[dict[str, Any]],
+    *,
+    protected_tail_start_index: int,
+) -> list[dict[str, Any]]:
+    """Prune duplicate older tool outputs while preserving provider-valid pairs.
+
+    Usage::
+
+        from aip_agents.context.offload import prune_duplicate_tool_outputs
+
+        # Only tool results at or after protected_tail_start_index are kept;
+        # older duplicates are replaced with back-reference placeholders.
+        pruned = prune_duplicate_tool_outputs(messages, protected_tail_start_index=10)
+    """
+    transformed = [copy.deepcopy(message) for message in messages]
+    latest_by_fingerprint: dict[str, int] = {}
+
+    for index in range(len(transformed) - 1, -1, -1):
+        _replace_duplicate_tool_output(
+            transformed,
+            latest_by_fingerprint,
+            index,
+            protected_tail_start_index=protected_tail_start_index,
+        )
+
+    return transformed
+
+
+def _replace_duplicate_tool_output(
+    messages: list[dict[str, Any]],
+    latest_by_fingerprint: dict[str, int],
+    index: int,
+    *,
+    protected_tail_start_index: int,
+) -> None:
+    message = messages[index]
+    if not _is_tool_result_message(message):
+        return
+
+    fingerprint = _tool_output_fingerprint(message)
+    if index >= protected_tail_start_index:
+        latest_by_fingerprint.setdefault(fingerprint, index)
+        return
+
+    newest_index = latest_by_fingerprint.get(fingerprint)
+    if newest_index is None:
+        latest_by_fingerprint[fingerprint] = index
+        return
+
+    _replace_with_tool_output_back_reference(message, newest_index, messages[newest_index].get("offloaded_ref_id"))
+
+
+def _replace_with_tool_output_back_reference(message: dict[str, Any], newest_index: int, newest_ref: Any) -> None:
+    content = message.get("content")
+    if not isinstance(content, str):
+        return
+    message["content"] = (
+        "[Tool output back-reference]\n"
+        f"original_index: {newest_index}\n"
+        f"ref: {newest_ref or 'none'}\n"
+        "Use the newest protected-tail tool output for full details."
+    )
+
+
+def repair_tool_result_pairing(messages: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Repair missing tool-result IDs when a single prior tool call is unambiguous.
+
+    Usage::
+
+        from aip_agents.context.offload import repair_tool_result_pairing
+
+        # After pruning or offloading, tool-call IDs may be missing from
+        # tool-result messages. This restores them when one unambiguous call
+        # precedes each result.
+        repaired = repair_tool_result_pairing(messages)
+    """
+    repaired = [copy.deepcopy(message) for message in messages]
+    recent_tool_call_ids: list[str] = []
+
+    for message in repaired:
+        recent_tool_call_ids = _collect_tool_call_ids(message, recent_tool_call_ids)
+        recent_tool_call_ids = _repair_tool_result_id(message, recent_tool_call_ids)
+
+    return repaired
+
+
+def _collect_tool_call_ids(message: dict[str, Any], current: list[str]) -> list[str]:
+    tool_calls = message.get("tool_calls")
+    if not isinstance(tool_calls, list):
+        return current
+    fresh: list[str] = []
+    for call in tool_calls:
+        if isinstance(call, dict):
+            call_id = call.get("id")
+            if isinstance(call_id, str) and call_id:
+                fresh.append(call_id)
+    return fresh
+
+
+def _repair_tool_result_id(message: dict[str, Any], recent_tool_call_ids: list[str]) -> list[str]:
+    if not _is_tool_result_message(message):
+        return recent_tool_call_ids
+    tool_call_id = message.get("tool_call_id")
+    if isinstance(tool_call_id, str) and tool_call_id:
+        return []
+    if len(recent_tool_call_ids) == 1:
+        message["tool_call_id"] = recent_tool_call_ids[0]
+        return []
+    return recent_tool_call_ids
+
+
+def _extract_text_content(message: dict[str, Any]) -> str:
+    content = message.get("content")
+    if isinstance(content, str):
+        return content
+    return str(content or "")
+
+
+def _should_offload_content(content: str, *, max_inline_chars: int) -> bool:
+    if len(content) > max_inline_chars:
+        return True
+    if _SECRET_KV_PATTERN.search(content):
+        return True
+    if _SIGNED_URL_PATTERN.search(content):
+        return True
+    if _LONG_BASE64_PATTERN.search(content):
+        return True
+    return False
+
+
+def _infer_source_kind(message: dict[str, Any]) -> SourceKind:
+    explicit = message.get("source_kind")
+    if explicit in {"tool_input", "tool_result", "subagent_output", "attachment", "artifact"}:
+        return explicit
+
+    role = str(message.get("role", "")).lower()
+    if role in {"tool", "tool_result"}:
+        return "tool_result"
+    if role in {"user", "human"}:
+        return "tool_input"
+    return "artifact"
+
+
+def _infer_display_name(message: dict[str, Any]) -> str | None:
+    for key in ("name", "tool_name", "title"):
+        value = message.get(key)
+        if isinstance(value, str) and value.strip():
+            return value.strip()
+    return None
+
+
+def _safe_preview(
+    content: str, *, max_chars: int
+) -> tuple[str, bool]:  # NOSONAR - intentionally redacts sensitive data
+    # deferred to #809: Use shared aip_agents.context.utils.safe_preview
+    redacted = _SIGNED_QUERY_KV_PATTERN.sub(lambda m: f"{m.group(1)}=[REDACTED]", content)
+    redacted = _SECRET_KV_PATTERN.sub(lambda m: f"{m.group(1)}=[REDACTED]", redacted)
+    redacted = _LONG_BASE64_PATTERN.sub("[BASE64_REDACTED]", redacted)
+    redacted = re.sub(r"(?i)(https?://[^\s?]+)\?[^\s]+", r"\1?[REDACTED_QUERY]", redacted)
+    redaction_applied = redacted != content
+
+    clipped = redacted[:max_chars]
+    if len(redacted) > max_chars:
+        clipped += "..."
+
+    return clipped, redaction_applied
+
+
+def _is_tool_result_message(message: dict[str, Any]) -> bool:
+    role = str(message.get("role", "")).lower()
+    return role in {"tool", "tool_result"} or "tool_call_id" in message
+
+
+def _tool_output_fingerprint(message: dict[str, Any]) -> str:
+    tool_name = str(message.get("name") or message.get("tool_name") or "")
+    content = _extract_text_content(message)
+    normalized = f"{tool_name}\n{content}".encode()
+    return hashlib.sha256(normalized).hexdigest()