okstra 0.174.0 → 0.176.0

package/runtime/python/okstra_ctl/assignment_resolver.py

@@ -0,0 +1,790 @@
+"""Deterministic role-instance model assignment and host binding."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Literal, Mapping, Sequence
+
+from .application.resolve_assignment import (
+    resolve_assignment_runner,
+    resolve_lead_provider_id,
+    resolve_lead_provider_from_specs,
+)
+from .domain.host import HostDescriptor, HostSessionContext
+from .domain.provider import (
+    ExecutionNormalization,
+    HostModelBinding,
+    LeadProviderAssignment,
+)
+from .domain.role import normalize_role
+from .domain.worker_exec import WorkerWriteCapability
+from .model_defaults import (
+    ModelDefaultError,
+    ModelDefaultScopes,
+    default_candidates,
+    default_scope,
+)
+from .model_pool import ModelPool, ResolvedModel
+from .ports.host_model import (
+    HostModelBindingError,
+    HostModelBindingPort,
+    HostModelBindingRequest,
+)
+from .role_requirements import RoleProfile, RoleRequirement
+
+
+TerminalBackend = Literal[
+    "cli-wrapper",
+    "cmux-pane",
+    "mixed",
+]
+_TERMINAL_BACKENDS = frozenset({
+    "cli-wrapper",
+    "cmux-pane",
+    "mixed",
+})
+
+
+@dataclass(frozen=True)
+class AssignmentEnvironment:
+    """Immutable host and dispatch facts consumed by assignment resolution."""
+
+    host_descriptor: HostDescriptor
+    host_model_port: HostModelBindingPort
+    native_worker_write_capability: WorkerWriteCapability
+    terminal_backend: TerminalBackend
+    execution_resolutions: Mapping[
+        tuple[str, str],
+        ExecutionNormalization,
+    ]
+
+    def __post_init__(self) -> None:
+        if self.terminal_backend not in _TERMINAL_BACKENDS:
+            allowed = ", ".join(sorted(_TERMINAL_BACKENDS))
+            raise ValueError(
+                f"unknown terminal backend {self.terminal_backend!r}. "
+                f"Allowed values: {allowed}"
+            )
+        resolutions: dict[tuple[str, str], ExecutionNormalization] = {}
+        for key, resolution in self.execution_resolutions.items():
+            if (
+                not isinstance(key, tuple)
+                or len(key) != 2
+                or not all(isinstance(value, str) and value for value in key)
+                or not isinstance(resolution, ExecutionNormalization)
+            ):
+                raise ValueError(
+                    "execution resolutions require (model_ref, role) keys and "
+                    "ExecutionNormalization values"
+                )
+            resolutions[key] = resolution
+        object.__setattr__(
+            self,
+            "execution_resolutions",
+            MappingProxyType(resolutions),
+        )
+
+
+@dataclass(frozen=True)
+class AssignmentContext:
+    pool: ModelPool
+    environment: AssignmentEnvironment
+
+
+@dataclass(frozen=True)
+class RoleCountSelection:
+    role: str
+    count: int
+
+
+@dataclass(frozen=True)
+class RoleModelSelection:
+    role: str
+    model_refs: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class RoleInstance:
+    role: str
+    duty_id: str
+    ordinal: int
+
+
+@dataclass(frozen=True)
+class CandidateAssessment:
+    role: str
+    ordinal: int
+    model_ref: str
+    scope: str
+    available: bool
+    reason: str = ""
+
+
+@dataclass(frozen=True)
+class UnsatisfiedConstraint:
+    role: str
+    constraint: str
+    required: int
+    available: int
+
+
+@dataclass(frozen=True)
+class ResolvedAssignment:
+    instance: RoleInstance
+    model_ref: str | None
+    provider_id: str
+    model_id: str
+    display_name: str
+    host_runtime: str
+    entry_mode: Literal["current-session", "new-session"]
+    binding: HostModelBinding | None
+
+    @property
+    def role(self) -> str:
+        return self.instance.role
+
+    @property
+    def duty_id(self) -> str:
+        return self.instance.duty_id
+
+    @property
+    def ordinal(self) -> int:
+        return self.instance.ordinal
+
+    @property
+    def display(self) -> str:
+        return self.display_name
+
+    @property
+    def execution(self) -> str:
+        return self.binding.resolved_execution_value if self.binding else ""
+
+    @property
+    def execution_label(self) -> str:
+        return f"{self.role}-{self.provider_id}-{self.model_id}-{self.ordinal}"
+
+
+@dataclass(frozen=True)
+class AssignmentPlan:
+    assignments: tuple[ResolvedAssignment, ...]
+    candidates: tuple[CandidateAssessment, ...] = ()
+
+    def by_role(self, role: str) -> tuple[ResolvedAssignment, ...]:
+        return tuple(row for row in self.assignments if row.role == role)
+
+
+class AssignmentResolutionError(ValueError):
+    """Raised with all evaluated candidates and failed constraints."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        candidates: Sequence[CandidateAssessment] = (),
+        unsatisfied_constraints: Sequence[UnsatisfiedConstraint] = (),
+    ) -> None:
+        self.candidates = tuple(candidates)
+        self.unsatisfied_constraints = tuple(unsatisfied_constraints)
+        super().__init__(message)
+
+
+def _validate_environment(
+    host: HostSessionContext,
+    environment: AssignmentEnvironment,
+) -> None:
+    if environment.host_descriptor.id != host.host_id:
+        raise AssignmentResolutionError(
+            "assignment environment host does not match the session context: "
+            f"{environment.host_descriptor.id!r} != {host.host_id!r}"
+        )
+
+
+def resolve_assignments(
+    *,
+    profile: RoleProfile,
+    role_counts: Mapping[str, int],
+    role_models: Mapping[str, tuple[str, ...]],
+    scopes: ModelDefaultScopes,
+    pool: ModelPool,
+    host: HostSessionContext,
+    environment: AssignmentEnvironment,
+) -> AssignmentPlan:
+    """Resolve the first complete assignment in deterministic DFS order."""
+    _validate_environment(host, environment)
+    requirements = _static_requirements(profile)
+    instances = _role_instances(requirements, role_counts)
+    _validate_model_selections(instances, role_models)
+    current_leader = _current_leader_assignment(
+        instances,
+        role_models,
+        pool,
+        host,
+        environment,
+    )
+    options, assessments = _assignment_options(
+        instances,
+        role_models,
+        scopes,
+        pool,
+        host,
+        profile,
+        requirements,
+        current_leader,
+        environment,
+    )
+    resolved = _first_complete_assignment(options, requirements)
+    if resolved is None:
+        constraints = _unmet_constraints(options, requirements)
+        raise AssignmentResolutionError(
+            "no complete role assignment satisfies the profile constraints",
+            candidates=assessments,
+            unsatisfied_constraints=constraints,
+        )
+    return AssignmentPlan(tuple(resolved), tuple(assessments))
+
+
+def resolve_model_assignment(
+    *,
+    instance: RoleInstance,
+    model_ref: str,
+    pool: ModelPool,
+    host: HostSessionContext,
+    environment: AssignmentEnvironment,
+    entry_mode: Literal["current-session", "new-session"] = "new-session",
+) -> ResolvedAssignment:
+    """Resolve one canonical model through the selected host execution path."""
+    _validate_environment(host, environment)
+    model = pool.resolve(model_ref)
+    availability = pool.availability(
+        model_ref,
+        instance.role,
+        host.host_id,
+        entry_mode,
+    )
+    if not availability.available:
+        raise AssignmentResolutionError(
+            f"model {model_ref!r} is unavailable: {availability.reason}"
+        )
+    binding = _first_supported_binding(
+        instance,
+        model,
+        pool,
+        environment,
+    )
+    return _resolved_assignment(instance, model, host, entry_mode, binding)
+
+
+def resolve_dispatch_assignment(
+    *,
+    context: AssignmentContext,
+    host_runtime: str,
+    role: str,
+    duty_id: str,
+    provider: str,
+    model: str = "",
+) -> ResolvedAssignment:
+    """Resolve one legacy provider/model input through the shared core."""
+    canonical_role = normalize_role(role)
+    pool = context.pool
+    selected = resolve_model_selection(pool, provider, canonical_role, model)
+    host = HostSessionContext(
+        host_id=host_runtime,
+        entry_mode="spawn-process",
+        available_functions=frozenset(),
+        interaction_surface="application-dispatch",
+    )
+    return resolve_model_assignment(
+        instance=RoleInstance(canonical_role, duty_id, 1),
+        model_ref=str(selected.model_ref),
+        pool=pool,
+        host=host,
+        environment=context.environment,
+    )
+
+
+def resolve_context_lead_provider(
+    *,
+    context: AssignmentContext,
+    requested_provider: str,
+) -> LeadProviderAssignment:
+    """Resolve a lead provider from the same immutable assignment facts."""
+    host = context.environment.host_descriptor
+    provider_id = resolve_lead_provider_id(
+        host=host,
+        requested_provider=requested_provider,
+    )
+    provider = context.pool.resolve_provider(provider_id)
+    return resolve_lead_provider_from_specs(
+        host=host,
+        requested_provider=requested_provider,
+        provider=provider,
+    )
+
+
+def resolve_model_selection(
+    pool: ModelPool,
+    provider: str,
+    role: str,
+    raw_model: str,
+) -> ResolvedModel:
+    """Resolve one explicit or provider-default catalog model without binding."""
+    if raw_model:
+        return pool.resolve_alias(provider, raw_model)
+    selected = pool.default_candidate(provider, role)
+    if selected is None:
+        raise AssignmentResolutionError(
+            f"provider {provider!r} has no default model for role {role!r}"
+        )
+    return selected
+
+
+def _static_requirements(profile: RoleProfile) -> tuple[RoleRequirement, ...]:
+    leader = RoleRequirement("leader", 1, 1, 1, "lead")
+    return (leader, *(row for row in profile.roles if not row.dynamic))
+
+
+def _role_instances(
+    requirements: Sequence[RoleRequirement],
+    role_counts: Mapping[str, int],
+) -> tuple[RoleInstance, ...]:
+    allowed = {row.role for row in requirements}
+    unknown = tuple(sorted(set(role_counts) - allowed))
+    if unknown:
+        raise AssignmentResolutionError(
+            f"role counts contain roles outside the profile: {', '.join(unknown)}"
+        )
+    instances: list[RoleInstance] = []
+    for requirement in requirements:
+        fixed = requirement.min_count == requirement.max_count
+        if requirement.role in role_counts and fixed:
+            raise AssignmentResolutionError(
+                f"role {requirement.role!r} has a fixed quantity and cannot "
+                "receive an explicit count"
+            )
+        count = role_counts.get(requirement.role, requirement.recommended_count)
+        if count < requirement.min_count or count > requirement.max_count:
+            raise AssignmentResolutionError(
+                f"role {requirement.role!r} count must be in "
+                f"{requirement.min_count}..{requirement.max_count}: {count}"
+            )
+        instances.extend(
+            RoleInstance(requirement.role, requirement.duty, ordinal)
+            for ordinal in range(1, count + 1)
+        )
+    return tuple(instances)
+
+
+def _validate_model_selections(
+    instances: Sequence[RoleInstance],
+    role_models: Mapping[str, tuple[str, ...]],
+) -> None:
+    counts = {
+        role: sum(instance.role == role for instance in instances)
+        for role in {instance.role for instance in instances}
+    }
+    unknown = tuple(sorted(set(role_models) - counts.keys()))
+    if unknown:
+        raise AssignmentResolutionError(
+            f"role models contain roles outside the profile: {', '.join(unknown)}"
+        )
+    for role, models in role_models.items():
+        if len(models) > counts[role]:
+            raise AssignmentResolutionError(
+                f"role {role!r} has {len(models)} models for {counts[role]} instances"
+            )
+
+
+def _current_leader_assignment(
+    instances: Sequence[RoleInstance],
+    role_models: Mapping[str, tuple[str, ...]],
+    pool: ModelPool,
+    host: HostSessionContext,
+    environment: AssignmentEnvironment,
+) -> ResolvedAssignment | None:
+    if host.entry_mode != "current-session":
+        return None
+    leader = next(row for row in instances if row.role == "leader")
+    attestation = host.current_model
+    model = _validated_current_model(pool, host)
+    explicit = role_models.get("leader", ())
+    if explicit and (model is None or explicit[0] != str(model.model_ref)):
+        return None
+    if model is None:
+        provider_id = (
+            attestation.provider_id
+            or environment.host_descriptor.native_provider_id
+        )
+        return ResolvedAssignment(
+            leader,
+            None,
+            provider_id,
+            "unknown",
+            "unknown",
+            host.host_id,
+            "current-session",
+            None,
+        )
+    binding = HostModelBinding(
+        runner="current-session",
+        catalog_execution_value=model.execution_value,
+        resolved_execution_value=model.execution_value,
+        host_model_value=attestation.observed_model,
+        binding_fidelity=attestation.level,
+        worker_write_capability=None,
+    )
+    return _resolved_assignment(leader, model, host, "current-session", binding)
+
+
+def _validated_current_model(
+    pool: ModelPool,
+    host: HostSessionContext,
+) -> ResolvedModel | None:
+    attestation = host.current_model
+    if attestation.level == "unknown":
+        return None
+    model_ref = attestation.normalized_model_ref or ""
+    try:
+        model = pool.resolve(model_ref)
+    except ValueError as exc:
+        action = "update the provider adapter catalog or start a registered leader model"
+        raise AssignmentResolutionError(
+            "unregistered current-session model attestation: "
+            f"provider={attestation.provider_id!r}, "
+            f"model={attestation.observed_model!r}, source={attestation.source!r}; "
+            f"{action}"
+        ) from exc
+    if model.provider_id != attestation.provider_id:
+        raise AssignmentResolutionError(
+            "current-session model attestation provider does not match its model ref"
+        )
+    if attestation.level == "exact" and model.version_kind != "pinned":
+        raise AssignmentResolutionError(
+            "exact current-session attestation must reference a pinned model"
+        )
+    if attestation.level == "channel" and model.version_kind != "channel":
+        raise AssignmentResolutionError(
+            "channel current-session attestation must reference a channel model"
+        )
+    availability = pool.availability(
+        str(model.model_ref),
+        "leader",
+        host.host_id,
+        "current-session",
+    )
+    if not availability.available:
+        raise AssignmentResolutionError(
+            "current-session model is unavailable for leader: "
+            f"{availability.reason}"
+        )
+    return model
+
+
+def _assignment_options(
+    instances: Sequence[RoleInstance],
+    role_models: Mapping[str, tuple[str, ...]],
+    scopes: ModelDefaultScopes,
+    pool: ModelPool,
+    host: HostSessionContext,
+    profile: RoleProfile,
+    requirements: Sequence[RoleRequirement],
+    current_leader: ResolvedAssignment | None,
+    environment: AssignmentEnvironment,
+) -> tuple[list[tuple[ResolvedAssignment, ...]], list[CandidateAssessment]]:
+    options: list[tuple[ResolvedAssignment, ...]] = []
+    assessments: list[CandidateAssessment] = []
+    for instance in instances:
+        if current_leader is not None and instance.role == "leader":
+            options.append((current_leader,))
+            continue
+        rows, evaluated = _instance_options(
+            instance,
+            role_models,
+            scopes,
+            pool,
+            host,
+            profile,
+            environment,
+        )
+        options.append(rows)
+        assessments.extend(evaluated)
+    if any(not rows for rows in options):
+        raise AssignmentResolutionError(
+            "one or more role instances have no executable model candidates",
+            candidates=assessments,
+            unsatisfied_constraints=_unmet_constraints(options, requirements),
+        )
+    return options, assessments
+
+
+def _instance_options(
+    instance: RoleInstance,
+    role_models: Mapping[str, tuple[str, ...]],
+    scopes: ModelDefaultScopes,
+    pool: ModelPool,
+    host: HostSessionContext,
+    profile: RoleProfile,
+    environment: AssignmentEnvironment,
+) -> tuple[tuple[ResolvedAssignment, ...], list[CandidateAssessment]]:
+    explicit = role_models.get(instance.role, ())
+    if instance.ordinal <= len(explicit):
+        candidates, scope = (explicit[instance.ordinal - 1],), "explicit"
+    else:
+        candidates, scope = _rotated_defaults(instance, scopes, pool)
+    resolved: list[ResolvedAssignment] = []
+    assessments: list[CandidateAssessment] = []
+    for model_ref in candidates:
+        try:
+            assignment = resolve_model_assignment(
+                instance=instance,
+                model_ref=model_ref,
+                pool=pool,
+                host=host,
+                environment=environment,
+            )
+            resolved.append(assignment)
+            assessments.append(_assessment(instance, model_ref, scope, True, ""))
+        except (AssignmentResolutionError, ValueError, RuntimeError) as exc:
+            assessments.append(
+                _assessment(instance, model_ref, scope, False, str(exc))
+            )
+    if scope != "bundled" and any(not row.available for row in assessments):
+        return (), assessments
+    return tuple(resolved), assessments
+
+
+def _rotated_defaults(
+    instance: RoleInstance,
+    scopes: ModelDefaultScopes,
+    pool: ModelPool,
+) -> tuple[tuple[str, ...], str]:
+    try:
+        candidates = default_candidates(instance.role, scopes)
+    except ModelDefaultError as exc:
+        raise AssignmentResolutionError(str(exc)) from exc
+    scope = default_scope(instance.role, scopes)
+    if not candidates:
+        candidates = tuple(
+            str(model.model_ref) for model in pool.default_candidates(instance.role)
+        )
+        scope = "bundled"
+    if not candidates:
+        raise AssignmentResolutionError(
+            f"role {instance.role!r} has no default model candidates"
+        )
+    offset = (instance.ordinal - 1) % len(candidates)
+    return (*candidates[offset:], *candidates[:offset]), scope or "bundled"
+
+
+def _assessment(
+    instance: RoleInstance,
+    model_ref: str,
+    scope: str,
+    available: bool,
+    reason: str,
+) -> CandidateAssessment:
+    return CandidateAssessment(
+        instance.role,
+        instance.ordinal,
+        model_ref,
+        scope,
+        available,
+        reason,
+    )
+
+
+def _first_supported_binding(
+    instance: RoleInstance,
+    model: ResolvedModel,
+    pool: ModelPool,
+    environment: AssignmentEnvironment,
+) -> HostModelBinding:
+    provider = pool.resolve_provider(model.provider_id)
+    errors: list[str] = []
+    for runner in _runner_candidates(instance, environment, provider):
+        try:
+            resolve_assignment_runner(
+                host=environment.host_descriptor,
+                provider=provider,
+                role=instance.role,
+                requested_runner=runner,
+            )
+            binding = _binding_for_runner(
+                instance,
+                model,
+                provider,
+                environment,
+                runner,
+            )
+            _validate_binding_fidelity(model, binding)
+            return binding
+        except (ValueError, RuntimeError, AttributeError) as exc:
+            errors.append(f"{runner}: {exc}")
+    raise AssignmentResolutionError(
+        f"model {model.model_ref!r} has no executable binding: {'; '.join(errors)}"
+    )
+
+
+def _runner_candidates(
+    instance: RoleInstance,
+    environment: AssignmentEnvironment,
+    provider,
+) -> tuple[str, ...]:
+    if (
+        instance.role != "leader"
+        and environment.terminal_backend == "cmux-pane"
+    ):
+        return ("cli-wrapper",) if provider.wrapper else ()
+    runners: list[str] = []
+    if provider.provider == environment.host_descriptor.native_provider_id:
+        runners.append("native-session")
+    if provider.wrapper:
+        runners.append("cli-wrapper")
+    return tuple(runners or ("cli-wrapper",))
+
+
+def _binding_for_runner(
+    instance: RoleInstance,
+    model: ResolvedModel,
+    provider,
+    environment: AssignmentEnvironment,
+    runner: str,
+) -> HostModelBinding:
+    execution = model.execution_value
+    if runner == "cli-wrapper":
+        resolution = environment.execution_resolutions.get(
+            (str(model.model_ref), instance.role)
+        )
+        if resolution is None and provider.execution_normalization_required:
+            raise AssignmentResolutionError(
+                f"missing execution snapshot for model {model.model_ref!r} "
+                f"and role {instance.role!r}"
+            )
+        if resolution is not None:
+            if resolution.error:
+                raise AssignmentResolutionError(resolution.error)
+            execution = resolution.resolved_execution_value
+    capability = _worker_write_capability(
+        instance,
+        provider,
+        environment,
+        runner,
+    )
+    request = HostModelBindingRequest(
+        host_runtime=environment.host_descriptor.id,
+        provider=model.provider_id,
+        model_ref=str(model.model_ref),
+        version_kind=model.version_kind,
+        model_execution_value=model.execution_value,
+        resolved_execution_value=execution,
+        runner=runner,
+        worker_write_capability=capability,
+    )
+    try:
+        return environment.host_model_port.resolve(request)
+    except HostModelBindingError as exc:
+        raise AssignmentResolutionError(str(exc)) from exc
+
+
+def _worker_write_capability(instance, provider, environment, runner):
+    if instance.role == "leader":
+        return None
+    if runner == "native-session":
+        return environment.native_worker_write_capability
+    if provider.cli_worker_write_capability is None:
+        raise AssignmentResolutionError(
+            f"provider {provider.provider!r} has no worker execution strategy"
+        )
+    return provider.cli_worker_write_capability
+
+
+def _validate_binding_fidelity(
+    model: ResolvedModel,
+    binding: HostModelBinding,
+) -> None:
+    allowed = {"exact"} if model.version_kind == "pinned" else {"exact", "channel"}
+    if binding.binding_fidelity not in allowed:
+        raise AssignmentResolutionError(
+            f"model {model.model_ref!r} cannot bind with fidelity "
+            f"{binding.binding_fidelity!r}"
+        )
+
+
+def _resolved_assignment(instance, model, host, entry_mode, binding):
+    return ResolvedAssignment(
+        instance=instance,
+        model_ref=str(model.model_ref),
+        provider_id=model.provider_id,
+        model_id=model.model_ref.model_id,
+        display_name=model.display_name,
+        host_runtime=host.host_id,
+        entry_mode=entry_mode,
+        binding=binding,
+    )
+
+
+def _first_complete_assignment(
+    options: Sequence[Sequence[ResolvedAssignment]],
+    requirements: Sequence[RoleRequirement],
+) -> list[ResolvedAssignment] | None:
+    chosen: list[ResolvedAssignment] = []
+
+    def visit(index: int) -> list[ResolvedAssignment] | None:
+        if index == len(options):
+            return list(chosen) if not _constraint_failures(chosen, requirements) else None
+        for assignment in options[index]:
+            chosen.append(assignment)
+            if (complete := visit(index + 1)) is not None:
+                return complete
+            chosen.pop()
+        return None
+
+    return visit(0)
+
+
+def _constraint_failures(
+    assignments: Sequence[ResolvedAssignment],
+    requirements: Sequence[RoleRequirement],
+) -> tuple[UnsatisfiedConstraint, ...]:
+    failures: list[UnsatisfiedConstraint] = []
+    for requirement in requirements:
+        role_rows = [row for row in assignments if row.role == requirement.role]
+        if not role_rows:
+            continue
+        refs = {row.model_ref for row in role_rows}
+        if len(refs) < len(role_rows):
+            failures.append(UnsatisfiedConstraint(
+                requirement.role,
+                "uniqueModelRefs",
+                len(role_rows),
+                len(refs),
+            ))
+    return tuple(failures)
+
+
+def _unmet_constraints(
+    options: Sequence[Sequence[ResolvedAssignment]],
+    requirements: Sequence[RoleRequirement],
+) -> tuple[UnsatisfiedConstraint, ...]:
+    failures: list[UnsatisfiedConstraint] = []
+    for requirement in requirements:
+        role_groups = [
+            candidates
+            for candidates in options
+            if candidates and candidates[0].role == requirement.role
+        ]
+        if not role_groups:
+            continue
+        worker_count = len(role_groups)
+        unique_refs = {
+            row.model_ref
+            for candidates in role_groups
+            for row in candidates
+        }
+        if len(unique_refs) < worker_count:
+            failures.append(UnsatisfiedConstraint(
+                requirement.role,
+                "uniqueModelRefs",
+                worker_count,
+                len(unique_refs),
+            ))
+    return tuple(failures)