Agent/cyber_agent/application/context.py

"""Deterministic TaskContextProvider resolution into existing ContextRefs."""

from __future__ import annotations

from hashlib import sha256
from typing import Any

from cyber_agent.application.models import canonical_json
from cyber_agent.application.ports import (
    ContextMaterial,
    ContextRequest,
    ContextResourceDescriptor,
)
from cyber_agent.core.context_policy import (
    MAX_CONTEXT_REFS,
    MAX_CONTEXT_SNAPSHOT_CHARS,
    MAX_CONTEXT_SNAPSHOTS_CHARS,
    ContextPolicyError,
    ContextSnapshot,
    create_context_snapshot,
    ensure_context_access,
    replace_context_access,
)
from cyber_agent.core.task_protocol import RootTaskAnchor, TaskBrief


def _descriptor_identity(
    descriptor: ContextResourceDescriptor,
) -> tuple[str, str]:
    return descriptor.source_id, descriptor.source_version


def _authorize(
    request: ContextRequest,
    descriptor: ContextResourceDescriptor,
) -> None:
    expected = (
        request.application_ref,
        request.root_trace_id,
        request.trace_id,
        request.uid,
        request.role_id,
    )
    actual = (
        descriptor.application_ref,
        descriptor.root_trace_id,
        descriptor.trace_id,
        descriptor.uid,
        descriptor.role_id,
    )
    if actual != expected:
        raise ContextPolicyError(
            "Application context descriptor belongs to another run, task, or role"
        )


async def load_application_context(
    binding: Any,
    context: dict[str, Any],
    request: ContextRequest,
    *,
    root_task_anchor: RootTaskAnchor,
    task_brief: TaskBrief | dict[str, Any] | None,
    granted_at_sequence: int,
) -> list[ContextSnapshot]:
    """Resolve, verify, and atomically add a stable bounded resource set."""
    provider = binding.services.context_provider
    if provider is None:
        return []
    raw_descriptors = await provider.list_context(request)
    descriptors = [ContextResourceDescriptor.model_validate(item) for item in raw_descriptors]
    unique: dict[tuple[str, str], ContextResourceDescriptor] = {}
    for descriptor in descriptors:
        _authorize(request, descriptor)
        key = _descriptor_identity(descriptor)
        prior = unique.get(key)
        if prior is not None and prior.content_hash != descriptor.content_hash:
            raise ContextPolicyError(
                f"Application context identity has conflicting hashes: {key}"
            )
        unique[key] = descriptor
    ordered = sorted(
        unique.values(),
        key=lambda item: (-item.priority, item.source_id, item.source_version),
    )
    access = ensure_context_access(context)
    existing = [
        ContextSnapshot.model_validate(item)
        for item in access["snapshots"].values()
    ]
    available_slots = max(0, MAX_CONTEXT_REFS - len(existing))
    current_chars = sum(len(canonical_json(item.content)) for item in existing)
    selected: list[ContextSnapshot] = []
    for descriptor in ordered:
        if len(selected) >= available_slots:
            break
        if descriptor.estimated_chars > MAX_CONTEXT_SNAPSHOT_CHARS:
            continue
        material = ContextMaterial.model_validate(
            await provider.resolve_context(request, descriptor)
        )
        if material.descriptor != descriptor:
            raise ContextPolicyError(
                f"Application context resolver changed descriptor: {descriptor.source_id}"
            )
        encoded_content = canonical_json(material.content)
        actual_hash = sha256(encoded_content.encode("utf-8")).hexdigest()
        if actual_hash != descriptor.content_hash:
            raise ContextPolicyError(
                f"Application context hash mismatch: {descriptor.source_id}"
            )
        wrapped = {
            "application_resource": {
                "application_ref": request.application_ref.model_dump(mode="json"),
                "source_id": descriptor.source_id,
                "source_version": descriptor.source_version,
                "kind": descriptor.kind,
                "inheritance": descriptor.inheritance,
                "authorization": descriptor.authorization,
            },
            "content": material.content,
        }
        chars = len(canonical_json(wrapped))
        if chars > MAX_CONTEXT_SNAPSHOT_CHARS:
            continue
        if current_chars + chars > MAX_CONTEXT_SNAPSHOTS_CHARS:
            continue
        source_trace_id = (
            f"application:{request.application_ref.application_id}:"
            f"{descriptor.source_id}@{descriptor.source_version}"
        )
        selected.append(create_context_snapshot(
            kind="application_resource",
            summary=descriptor.summary,
            source_trace_id=source_trace_id,
            root_trace_id=request.root_trace_id,
            uid=request.uid,
            content=wrapped,
            granted_at_sequence=granted_at_sequence,
        ))
        current_chars += chars
    replace_context_access(
        context,
        [*existing, *selected],
        root_task_anchor=root_task_anchor,
        task_brief=task_brief,
    )
    return selected