fix: unify company runtime recovery lifecycle
This commit is contained in:
@@ -0,0 +1,391 @@
|
||||
"""Canonical company-runtime identity derived from durable records.
|
||||
|
||||
Company-mode Tasks are execution envelopes, not the identity of a run. A
|
||||
runtime is owned by its root session and an active suspend checkpoint. This
|
||||
module deliberately has no UI dependencies so every surface can resolve the
|
||||
same scope without relying on process-local task maps.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Iterable
|
||||
|
||||
from opc.layer2_organization.work_item_links import linked_work_item_id_for_task
|
||||
from opc.layer2_organization.work_item_runtime import is_work_item_runtime_metadata
|
||||
|
||||
|
||||
COMPANY_RUNTIME_CHECKPOINT_TYPES: frozenset[str] = frozenset({
|
||||
"company_runtime_suspended",
|
||||
"company_runtime_interrupted",
|
||||
})
|
||||
ACTIVE_COMPANY_RUNTIME_CHECKPOINT_STATUSES: frozenset[str] = frozenset({
|
||||
"pending",
|
||||
"resuming",
|
||||
})
|
||||
|
||||
|
||||
def _text(value: Any) -> str:
|
||||
return str(value or "").strip()
|
||||
|
||||
|
||||
def _metadata(task: Any) -> dict[str, Any]:
|
||||
return dict(getattr(task, "metadata", {}) or {})
|
||||
|
||||
|
||||
def _task_id(task: Any) -> str:
|
||||
return _text(getattr(task, "id", ""))
|
||||
|
||||
|
||||
def _task_session_id(task: Any) -> str:
|
||||
return _text(getattr(task, "session_id", ""))
|
||||
|
||||
|
||||
def _task_parent_session_id(task: Any) -> str:
|
||||
metadata = _metadata(task)
|
||||
return _text(
|
||||
getattr(task, "parent_session_id", "")
|
||||
or metadata.get("company_runtime_root_session_id")
|
||||
or metadata.get("parent_session_id")
|
||||
)
|
||||
|
||||
|
||||
def _has_company_runtime_marker(task: Any) -> bool:
|
||||
metadata = _metadata(task)
|
||||
exec_mode = _text(metadata.get("exec_mode")).lower()
|
||||
mode = _text(metadata.get("mode")).lower()
|
||||
execution_mode = _text(metadata.get("execution_mode")).lower()
|
||||
if exec_mode in {"company", "org", "custom"} or mode in {"company", "org", "custom"}:
|
||||
return True
|
||||
if execution_mode in {"company", "company_mode", "multi_team_org"}:
|
||||
return True
|
||||
if is_work_item_runtime_metadata(metadata):
|
||||
return True
|
||||
if linked_work_item_id_for_task(task):
|
||||
return True
|
||||
if any(
|
||||
metadata.get(key) not in (None, "", [], {})
|
||||
for key in (
|
||||
"company_work_item_plan",
|
||||
"company_runtime_root_session_id",
|
||||
"delegation_run_id",
|
||||
"work_item_projection_id",
|
||||
"work_item_projection_ref",
|
||||
"work_item_role_id",
|
||||
"shared_role_session",
|
||||
)
|
||||
):
|
||||
return True
|
||||
# Old company records may predate exec_mode. An explicit task-mode marker
|
||||
# wins over the legacy profile hint.
|
||||
explicitly_task_mode = (
|
||||
exec_mode in {"task", "project", "single"}
|
||||
or mode == "task"
|
||||
or execution_mode in {"task", "task_mode", "project"}
|
||||
or _text(metadata.get("task_mode_contract")) == "single_full_capability_main_agent"
|
||||
)
|
||||
return not explicitly_task_mode and bool(_text(metadata.get("company_profile")))
|
||||
|
||||
|
||||
def is_company_runtime_task(task: Any) -> bool:
|
||||
"""Return whether durable Task metadata identifies company-owned work."""
|
||||
|
||||
return _has_company_runtime_marker(task)
|
||||
|
||||
|
||||
def is_pure_company_ui_anchor(task: Any, runtime_session_id: str) -> bool:
|
||||
"""Return whether *task* is the user-facing container for a runtime.
|
||||
|
||||
A shared final-decider Task can have the same ``session_id`` as the UI
|
||||
anchor. Work-item, role, or parent links therefore disqualify a Task even
|
||||
when its session id is an exact match.
|
||||
"""
|
||||
|
||||
session_id = _text(runtime_session_id)
|
||||
if not session_id or _task_session_id(task) != session_id:
|
||||
return False
|
||||
if _text(getattr(task, "parent_session_id", "")) or _text(getattr(task, "parent_id", "")):
|
||||
return False
|
||||
if linked_work_item_id_for_task(task):
|
||||
return False
|
||||
metadata = _metadata(task)
|
||||
return not any(
|
||||
metadata.get(key) not in (None, "", [], {}, False)
|
||||
for key in (
|
||||
"work_item_runtime",
|
||||
"work_item_projection_id",
|
||||
"work_item_projection_ref",
|
||||
"work_item_id",
|
||||
"work_item_role_id",
|
||||
"delegation_role_session_id",
|
||||
"shared_role_session",
|
||||
"shared_role_id",
|
||||
"company_runtime_root_session_id",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _created_sort_key(value: Any) -> tuple[float, str]:
|
||||
created_at = getattr(value, "created_at", None)
|
||||
if isinstance(created_at, datetime):
|
||||
timestamp = created_at.timestamp()
|
||||
elif hasattr(created_at, "timestamp"):
|
||||
try:
|
||||
timestamp = float(created_at.timestamp())
|
||||
except Exception:
|
||||
timestamp = 0.0
|
||||
else:
|
||||
timestamp = 0.0
|
||||
return timestamp, _task_id(value)
|
||||
|
||||
|
||||
def _checkpoint_sort_key(checkpoint: Any) -> tuple[float, float, str]:
|
||||
def _timestamp(value: Any) -> float:
|
||||
if isinstance(value, datetime):
|
||||
return value.timestamp()
|
||||
if hasattr(value, "timestamp"):
|
||||
try:
|
||||
return float(value.timestamp())
|
||||
except Exception:
|
||||
return 0.0
|
||||
return 0.0
|
||||
|
||||
return (
|
||||
_timestamp(getattr(checkpoint, "updated_at", None)),
|
||||
_timestamp(getattr(checkpoint, "created_at", None)),
|
||||
_text(getattr(checkpoint, "checkpoint_id", "")),
|
||||
)
|
||||
|
||||
|
||||
def _checkpoint_runtime_session_id(checkpoint: Any) -> str:
|
||||
payload = dict(getattr(checkpoint, "payload", {}) or {})
|
||||
return _text(
|
||||
getattr(checkpoint, "session_id", "")
|
||||
or payload.get("parent_session_id")
|
||||
or payload.get("session_id")
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CompanyRuntimeIdentity:
|
||||
"""Resolved identity for one company runtime scope."""
|
||||
|
||||
project_id: str
|
||||
runtime_session_id: str
|
||||
runtime_task_ids: tuple[str, ...]
|
||||
ui_anchor_task_id: str = ""
|
||||
config_source_task_id: str = ""
|
||||
pending_checkpoint_id: str = ""
|
||||
pending_checkpoint_type: str = ""
|
||||
pending_checkpoint_status: str = ""
|
||||
resumable: bool = False
|
||||
checkpoint: Any | None = field(default=None, repr=False, compare=False)
|
||||
|
||||
|
||||
class CompanyRuntimeIdentityIndex:
|
||||
"""Session-first index over preloaded Tasks and checkpoints."""
|
||||
|
||||
def __init__(self, tasks: Iterable[Any], checkpoints: Iterable[Any] = ()) -> None:
|
||||
self.tasks = tuple(tasks or ())
|
||||
self.checkpoints = tuple(checkpoints or ())
|
||||
self.tasks_by_id = {
|
||||
_task_id(task): task
|
||||
for task in self.tasks
|
||||
if _task_id(task)
|
||||
}
|
||||
self.checkpoints_by_id = {
|
||||
_text(getattr(checkpoint, "checkpoint_id", "")): checkpoint
|
||||
for checkpoint in self.checkpoints
|
||||
if _text(getattr(checkpoint, "checkpoint_id", ""))
|
||||
}
|
||||
self._identities_by_session = self._build_identities()
|
||||
self._runtime_session_by_task_id: dict[str, str] = {}
|
||||
runtime_sessions_by_task_session_id: dict[str, set[str]] = {}
|
||||
for runtime_session_id, identity in self._identities_by_session.items():
|
||||
for task_id in identity.runtime_task_ids:
|
||||
self._runtime_session_by_task_id[task_id] = runtime_session_id
|
||||
task_session_id = _task_session_id(self.tasks_by_id.get(task_id))
|
||||
if task_session_id:
|
||||
runtime_sessions_by_task_session_id.setdefault(
|
||||
task_session_id,
|
||||
set(),
|
||||
).add(runtime_session_id)
|
||||
self._runtime_session_by_task_session_id = {
|
||||
task_session_id: next(iter(runtime_session_ids))
|
||||
for task_session_id, runtime_session_ids in runtime_sessions_by_task_session_id.items()
|
||||
if len(runtime_session_ids) == 1
|
||||
}
|
||||
|
||||
@property
|
||||
def identities(self) -> tuple[CompanyRuntimeIdentity, ...]:
|
||||
return tuple(self._identities_by_session.values())
|
||||
|
||||
def task(self, task_id: str) -> Any | None:
|
||||
return self.tasks_by_id.get(_text(task_id))
|
||||
|
||||
def resolve(
|
||||
self,
|
||||
*,
|
||||
task_id: str = "",
|
||||
task_session_id: str = "",
|
||||
runtime_session_id: str = "",
|
||||
checkpoint_id: str = "",
|
||||
) -> CompanyRuntimeIdentity | None:
|
||||
requested_task_id = _text(task_id)
|
||||
requested_task_session_id = _text(task_session_id)
|
||||
requested_session_id = _text(runtime_session_id)
|
||||
requested_checkpoint_id = _text(checkpoint_id)
|
||||
|
||||
task_scope_id = self._runtime_session_by_task_id.get(requested_task_id, "")
|
||||
session_scope_id = self._runtime_session_by_task_session_id.get(
|
||||
requested_task_session_id,
|
||||
"",
|
||||
)
|
||||
checkpoint = self.checkpoints_by_id.get(requested_checkpoint_id) if requested_checkpoint_id else None
|
||||
checkpoint_session_id = _checkpoint_runtime_session_id(checkpoint) if checkpoint is not None else ""
|
||||
|
||||
candidates = {
|
||||
value
|
||||
for value in (
|
||||
requested_session_id,
|
||||
task_scope_id,
|
||||
session_scope_id,
|
||||
checkpoint_session_id,
|
||||
)
|
||||
if value
|
||||
}
|
||||
if requested_task_session_id and not session_scope_id:
|
||||
return None
|
||||
if len(candidates) != 1:
|
||||
return None
|
||||
resolved_session_id = next(iter(candidates))
|
||||
identity = self._identities_by_session.get(resolved_session_id)
|
||||
if identity is None:
|
||||
return None
|
||||
if requested_task_id and requested_task_id not in identity.runtime_task_ids:
|
||||
return None
|
||||
if requested_checkpoint_id and requested_checkpoint_id != identity.pending_checkpoint_id:
|
||||
return None
|
||||
return identity
|
||||
|
||||
def _build_identities(self) -> dict[str, CompanyRuntimeIdentity]:
|
||||
active_checkpoints_by_session: dict[str, list[Any]] = {}
|
||||
for checkpoint in self.checkpoints:
|
||||
checkpoint_type = _text(getattr(checkpoint, "checkpoint_type", ""))
|
||||
checkpoint_status = _text(getattr(checkpoint, "status", "")).lower()
|
||||
if (
|
||||
checkpoint_type not in COMPANY_RUNTIME_CHECKPOINT_TYPES
|
||||
or checkpoint_status not in ACTIVE_COMPANY_RUNTIME_CHECKPOINT_STATUSES
|
||||
):
|
||||
continue
|
||||
runtime_session_id = _checkpoint_runtime_session_id(checkpoint)
|
||||
if runtime_session_id:
|
||||
active_checkpoints_by_session.setdefault(runtime_session_id, []).append(checkpoint)
|
||||
|
||||
known_sessions = set(active_checkpoints_by_session)
|
||||
for task in self.tasks:
|
||||
if not _has_company_runtime_marker(task):
|
||||
continue
|
||||
runtime_session_id = _task_parent_session_id(task) or _task_session_id(task)
|
||||
if runtime_session_id:
|
||||
known_sessions.add(runtime_session_id)
|
||||
|
||||
tasks_by_session: dict[str, list[Any]] = {session_id: [] for session_id in known_sessions}
|
||||
for task in self.tasks:
|
||||
task_id = _task_id(task)
|
||||
if not task_id:
|
||||
continue
|
||||
parent_session_id = _task_parent_session_id(task)
|
||||
own_session_id = _task_session_id(task)
|
||||
runtime_session_id = parent_session_id or own_session_id
|
||||
if runtime_session_id not in known_sessions:
|
||||
continue
|
||||
if not (
|
||||
_has_company_runtime_marker(task)
|
||||
or runtime_session_id in active_checkpoints_by_session
|
||||
or is_pure_company_ui_anchor(task, runtime_session_id)
|
||||
):
|
||||
continue
|
||||
tasks_by_session.setdefault(runtime_session_id, []).append(task)
|
||||
|
||||
identities: dict[str, CompanyRuntimeIdentity] = {}
|
||||
for runtime_session_id in sorted(known_sessions):
|
||||
group = sorted(tasks_by_session.get(runtime_session_id, []), key=_created_sort_key)
|
||||
anchor = next(
|
||||
(task for task in group if is_pure_company_ui_anchor(task, runtime_session_id)),
|
||||
None,
|
||||
)
|
||||
def _has_runtime_config(task: Any) -> bool:
|
||||
metadata = _metadata(task)
|
||||
return any(
|
||||
metadata.get(key) not in (None, "", [], {})
|
||||
for key in (
|
||||
"exec_mode",
|
||||
"mode",
|
||||
"company_profile",
|
||||
"org_id",
|
||||
"organization_id",
|
||||
"preferred_agent",
|
||||
"selected_execution_agent",
|
||||
)
|
||||
)
|
||||
|
||||
config_source = (
|
||||
anchor if anchor is not None and _has_runtime_config(anchor) else None
|
||||
) or next(
|
||||
(
|
||||
task for task in group
|
||||
if _has_runtime_config(task)
|
||||
),
|
||||
anchor or (group[0] if group else None),
|
||||
)
|
||||
checkpoint_candidates = active_checkpoints_by_session.get(runtime_session_id, [])
|
||||
checkpoint = max(checkpoint_candidates, key=_checkpoint_sort_key) if checkpoint_candidates else None
|
||||
checkpoint_status = _text(getattr(checkpoint, "status", "")).lower() if checkpoint is not None else ""
|
||||
project_id = _text(
|
||||
getattr(checkpoint, "project_id", "") if checkpoint is not None else ""
|
||||
) or _text(getattr(config_source, "project_id", "") if config_source is not None else "") or "default"
|
||||
identities[runtime_session_id] = CompanyRuntimeIdentity(
|
||||
project_id=project_id,
|
||||
runtime_session_id=runtime_session_id,
|
||||
runtime_task_ids=tuple(_task_id(task) for task in group if _task_id(task)),
|
||||
ui_anchor_task_id=_task_id(anchor) if anchor is not None else "",
|
||||
config_source_task_id=_task_id(config_source) if config_source is not None else "",
|
||||
pending_checkpoint_id=_text(getattr(checkpoint, "checkpoint_id", "")) if checkpoint is not None else "",
|
||||
pending_checkpoint_type=_text(getattr(checkpoint, "checkpoint_type", "")) if checkpoint is not None else "",
|
||||
pending_checkpoint_status=checkpoint_status,
|
||||
resumable=checkpoint_status == "pending",
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
return identities
|
||||
|
||||
|
||||
def build_company_runtime_identity_index(
|
||||
tasks: Iterable[Any],
|
||||
checkpoints: Iterable[Any] = (),
|
||||
) -> CompanyRuntimeIdentityIndex:
|
||||
return CompanyRuntimeIdentityIndex(tasks, checkpoints)
|
||||
|
||||
|
||||
async def load_company_runtime_identity_index(
|
||||
store: Any,
|
||||
project_id: str,
|
||||
) -> CompanyRuntimeIdentityIndex:
|
||||
"""Load durable records once and build the canonical runtime index."""
|
||||
|
||||
tasks = await store.get_tasks(project_id=project_id)
|
||||
checkpoint_getter = getattr(store, "get_execution_checkpoints", None)
|
||||
if callable(checkpoint_getter):
|
||||
checkpoints = await checkpoint_getter(
|
||||
project_id=project_id,
|
||||
checkpoint_types=sorted(COMPANY_RUNTIME_CHECKPOINT_TYPES),
|
||||
statuses=sorted(ACTIVE_COMPANY_RUNTIME_CHECKPOINT_STATUSES),
|
||||
)
|
||||
else:
|
||||
checkpoint_getter = getattr(store, "get_pending_checkpoints", None)
|
||||
checkpoints = await checkpoint_getter(
|
||||
project_id=project_id,
|
||||
checkpoint_types=sorted(COMPANY_RUNTIME_CHECKPOINT_TYPES),
|
||||
) if callable(checkpoint_getter) else []
|
||||
return build_company_runtime_identity_index(tasks, checkpoints)
|
||||
Reference in New Issue
Block a user