fix: unify company runtime recovery lifecycle

This commit is contained in:
LZH-YS1998
2026-07-14 14:35:43 +08:00
parent 5e02364eb4
commit b8202bbe9e
56 changed files with 8753 additions and 3542 deletions
+195 -142
View File
@@ -17,6 +17,10 @@ from typing import Any, Awaitable, Callable
from loguru import logger
from opc.core.active_task_runs import (
ActiveTaskRunAdmissionClosed,
ActiveTaskRunRegistry,
)
from opc.core.config import DEFAULT_EXTERNAL_AGENT_STARTUP_TIMEOUT_SECONDS, DEFAULT_ORGANIZATION_ID
from opc.core.models import (
AdaptiveRoleProfile,
@@ -456,6 +460,26 @@ class WorkItemOutputBundle:
summary: str = ""
@dataclass(frozen=True)
class CompanyExecutorDriverOwnership:
"""One registry attempt covering a complete company scheduler run."""
registry: ActiveTaskRunRegistry
project_id: str
task_id: str
attempt_token: str
def bind(self):
return self.registry.bind_driver_attempt(self.attempt_token)
def release(self) -> bool:
return self.registry.unregister(
self.project_id,
self.task_id,
self.attempt_token,
)
def serialize_company_runtime_spec(spec: CompanyRuntimeSpec | None) -> dict[str, Any]:
if spec is None:
return {}
@@ -1354,6 +1378,7 @@ class CompanyWorkItemExecutor:
store: Any | None = None,
llm: Any | None = None,
role_prompt_runner: Callable[[Task, str, dict[str, Any], str, bool], Awaitable[str | None]] | None = None,
active_task_run_registry: ActiveTaskRunRegistry | None = None,
) -> None:
self.org_engine = org_engine
self.communication = communication
@@ -1373,6 +1398,7 @@ class CompanyWorkItemExecutor:
self.on_kanban_changed = on_kanban_changed
self.work_item_timeout = work_item_timeout
self.role_prompt_runner = role_prompt_runner
self.active_task_run_registry = active_task_run_registry
self._default_run_state = CompanyExecutorRunState()
self._run_state_var: ContextVar[CompanyExecutorRunState | None] = ContextVar(
f"company-executor-run-state:{id(self)}",
@@ -3368,12 +3394,6 @@ class CompanyWorkItemExecutor:
"cell_id": work_item.cell_id,
"parent_work_item_id": work_item.parent_work_item_id,
}
if work_item.phase == Phase.PAUSED:
task.metadata.setdefault("interrupted_recovery", {
"reason": "work_item_interrupted",
"detected_at": datetime.now().isoformat(),
})
if task.metadata != before_metadata:
changed = True
return changed
@@ -4359,14 +4379,76 @@ class CompanyWorkItemExecutor:
return "dispatch_required"
return "worker_execute"
async def execute(self, plan: CompanyWorkItemRuntimePlan, tasks: list[Task]) -> str:
plan = _coerce_company_work_item_runtime_plan(plan) or CompanyWorkItemRuntimePlan()
plan.metadata = {
**dict(plan.metadata or {}),
"execution_model": "multi_team_org",
"runtime_model": "multi_team_org",
@staticmethod
def _driver_ownership_task(
tasks: list[Task],
*,
preferred_task_ids: set[str] | None = None,
) -> Task | None:
preferred = {
str(task_id or "").strip()
for task_id in set(preferred_task_ids or set())
if str(task_id or "").strip()
}
return await self._execute_multi_team_org(plan, tasks)
candidates = [
task
for task in tasks
if str(getattr(task, "id", "") or "").strip()
and (not preferred or task.id in preferred)
]
for task in candidates:
if linked_work_item_id_for_task(task):
return task
for task in candidates:
if is_work_item_runtime_metadata(dict(task.metadata or {})):
return task
return candidates[0] if candidates else None
def acquire_driver_ownership(
self,
tasks: list[Task],
*,
preferred_task_ids: set[str] | None = None,
) -> CompanyExecutorDriverOwnership | None:
registry = self.active_task_run_registry
task = self._driver_ownership_task(
tasks,
preferred_task_ids=preferred_task_ids,
)
if registry is None or task is None:
return None
project_id = str(task.project_id or "default").strip() or "default"
try:
attempt_token = registry.register(project_id, task.id)
except ActiveTaskRunAdmissionClosed as exc:
raise asyncio.CancelledError(str(exc)) from exc
return CompanyExecutorDriverOwnership(
registry=registry,
project_id=project_id,
task_id=task.id,
attempt_token=attempt_token,
)
async def execute(
self,
plan: CompanyWorkItemRuntimePlan,
tasks: list[Task],
) -> str:
ownership = self.acquire_driver_ownership(tasks)
try:
plan = _coerce_company_work_item_runtime_plan(plan) or CompanyWorkItemRuntimePlan()
plan.metadata = {
**dict(plan.metadata or {}),
"execution_model": "multi_team_org",
"runtime_model": "multi_team_org",
}
if ownership is None:
return await self._execute_multi_team_org(plan, tasks)
with ownership.bind():
return await self._execute_multi_team_org(plan, tasks)
finally:
if ownership is not None:
ownership.release()
async def _execute_multi_team_org(
self,
@@ -4501,11 +4583,11 @@ class CompanyWorkItemExecutor:
# Claim whatever is immediately claimable and spawn each
# work item as an independent asyncio.Task so the loop no
# longer blocks on the slowest sibling.
claims = await self.runtime.claim_runnable_tasks(tasks, work_items=work_items)
for member_session, claimed_task in claims:
work_item_coro = self._run_claimed_work_item(member_session, claimed_task, {})
work_item_task = asyncio.create_task(work_item_coro)
active_work_item_tasks[work_item_task] = (member_session, claimed_task)
claims = await self._claim_and_create_work_item_tasks(
tasks,
work_items,
active_work_item_tasks,
)
# Termination: only when nothing is in-flight AND nothing
# else is runnable. If work items are still running, even an
# "empty runnable" snapshot may become non-empty within
@@ -4608,10 +4690,6 @@ class CompanyWorkItemExecutor:
self._schedule_kanban_notification()
except asyncio.CancelledError:
claimed_pairs = list(active_work_item_tasks.values())
claimed_tasks = [
claimed_task
for _member_session, claimed_task in claimed_pairs
]
for work_item_task in list(active_work_item_tasks.keys()):
if not work_item_task.done():
work_item_task.cancel()
@@ -4660,63 +4738,6 @@ class CompanyWorkItemExecutor:
)
except Exception:
logger.opt(exception=True).debug("company runtime cancellation: failed session idle reset")
for claimed_task in claimed_tasks:
if claimed_task.status in {TaskStatus.DONE, TaskStatus.FAILED, TaskStatus.CANCELLED}:
continue
claimed_task.metadata = dict(claimed_task.metadata or {})
claimed_task.metadata["company_runtime_suspended_at"] = datetime.now().isoformat()
claimed_task.metadata.setdefault("last_stop_reason", "runtime_cancelled")
claimed_task.metadata["company_runtime_stop_state"] = "suspended"
claimed_task.metadata["company_runtime_stop_marked_at"] = (
claimed_task.metadata.get("company_runtime_stop_marked_at") or datetime.now().isoformat()
)
claimed_task.metadata.setdefault(
"suspended_task_status",
claimed_task.status.value if isinstance(claimed_task.status, TaskStatus) else str(claimed_task.status or ""),
)
work_item_id = linked_work_item_id_for_task(claimed_task)
try:
if work_item_id and self._store_is_ready(self.store) and hasattr(self.store, "get_delegation_work_item"):
work_item = await self.store.get_delegation_work_item(work_item_id)
else:
work_item = None
if work_item is not None and getattr(work_item, "phase", None) not in {Phase.APPROVED, Phase.FAILED, Phase.CANCELLED}:
phase = getattr(work_item, "phase", Phase.RUNNING)
phase_value = phase.value if isinstance(phase, Phase) else str(phase or "")
original_claim = {
"claimed_by_role_runtime_session_id": str(getattr(work_item, "claimed_by_role_runtime_session_id", "") or ""),
"claimed_by_seat_id": str(getattr(work_item, "claimed_by_seat_id", "") or ""),
"claimed_by_role_session_id": str((getattr(work_item, "metadata", {}) or {}).get("claimed_by_role_session_id", "") or ""),
"claimed_task_id": str((getattr(work_item, "metadata", {}) or {}).get("claimed_task_id", "") or claimed_task.id),
}
await self.store.update_delegation_work_item(
work_item_id,
metadata_updates={
"dispatch_hold": "company_runtime_suspended",
"suspended_at": datetime.now().isoformat(),
"suspend_reason": claimed_task.metadata.get("last_stop_reason", "runtime_cancelled"),
"suspended_phase": phase_value,
"suspended_task_status": claimed_task.metadata.get("suspended_task_status", ""),
"suspended_claim": original_claim,
"claimed_by_role_session_id": "",
"claimed_task_id": "",
},
claimed_by_role_runtime_session_id="",
claimed_by_seat_id="",
)
claimed_task.metadata["dispatch_hold"] = "company_runtime_suspended"
claimed_task.metadata["suspended_phase"] = phase_value
except Exception:
logger.opt(exception=True).debug(
"company runtime cancellation: failed suspend hold release",
)
if self.save_task and self._store_is_ready(self.store):
try:
await self.save_task(claimed_task)
except Exception:
logger.opt(exception=True).debug(
"company runtime cancellation: failed suspended task save",
)
raise
finally:
# Drain any work items still running (shouldn't happen given the
@@ -4739,6 +4760,97 @@ class CompanyWorkItemExecutor:
active_work_item_tasks.clear()
return self._summarize_multi_team_org_results(tasks)
@staticmethod
def _runtime_scope_for_tasks(tasks: list[Task]) -> tuple[str, str]:
project_id = "default"
runtime_session_id = ""
for task in tasks:
project_id = str(task.project_id or project_id).strip() or "default"
metadata = dict(task.metadata or {})
runtime_session_id = str(
getattr(task, "parent_session_id", "")
or metadata.get("company_runtime_root_session_id")
or metadata.get("parent_session_id")
or ""
).strip()
if runtime_session_id:
return project_id, runtime_session_id
for task in tasks:
runtime_session_id = str(getattr(task, "session_id", "") or "").strip()
if runtime_session_id:
return project_id, runtime_session_id
return project_id, ""
async def _claim_and_create_work_item_tasks(
self,
tasks: list[Task],
work_items: list[DelegationWorkItem],
active_work_item_tasks: dict[
asyncio.Task[TaskResult | None],
tuple[CompanyMemberSession, Task],
],
) -> list[tuple[CompanyMemberSession, Task]]:
"""Keep durable claim and coroutine ownership in one scope boundary."""
async def claim_and_create() -> list[tuple[CompanyMemberSession, Task]]:
claims = await self.runtime.claim_runnable_tasks(
tasks,
work_items=work_items,
)
for member_session, claimed_task in claims:
work_item_task = self._create_claimed_work_item_task(
member_session,
claimed_task,
{},
)
active_work_item_tasks[work_item_task] = (
member_session,
claimed_task,
)
return claims
registry = self.active_task_run_registry
project_id, runtime_session_id = self._runtime_scope_for_tasks(tasks)
if registry is None or not runtime_session_id:
return await claim_and_create()
async with registry.scope_lock(project_id, runtime_session_id):
return await claim_and_create()
def _create_claimed_work_item_task(
self,
member_session: CompanyMemberSession,
task: Task,
task_by_projection_id: dict[str, Task],
) -> asyncio.Task[TaskResult | None]:
"""Register ownership before scheduling the full claimed-item coroutine."""
registry = self.active_task_run_registry
project_id = str(task.project_id or "default").strip() or "default"
attempt_token = ""
if registry is not None:
try:
attempt_token = registry.register(project_id, task.id)
except ActiveTaskRunAdmissionClosed as exc:
raise asyncio.CancelledError(str(exc)) from exc
async def run_owned() -> TaskResult | None:
try:
return await self._run_claimed_work_item(
member_session,
task,
task_by_projection_id,
)
finally:
if registry is not None and attempt_token:
registry.unregister(project_id, task.id, attempt_token)
try:
return asyncio.create_task(run_owned())
except BaseException:
if registry is not None and attempt_token:
registry.unregister(project_id, task.id, attempt_token)
raise
async def _run_claimed_work_item(
self,
member_session: CompanyMemberSession,
@@ -5298,69 +5410,10 @@ class CompanyWorkItemExecutor:
timeout=self.work_item_timeout,
)
except asyncio.CancelledError:
task.metadata = dict(task.metadata or {})
task.metadata["company_runtime_suspended_at"] = datetime.now().isoformat()
task.metadata.setdefault("last_stop_reason", "runtime_cancelled")
task.metadata["company_runtime_stop_state"] = "suspended"
task.metadata["company_runtime_stop_marked_at"] = (
task.metadata.get("company_runtime_stop_marked_at") or datetime.now().isoformat()
)
task.metadata.setdefault(
"suspended_task_status",
task.status.value if isinstance(task.status, TaskStatus) else str(task.status or ""),
)
work_item_id = linked_work_item_id_for_task(task)
store_ready = self._store_is_ready(self.store)
if work_item_id and self.store and store_ready:
task.metadata.pop("progress_log", None)
await append_work_item_progress(
self.store,
work_item_id,
"Work item suspended by runtime cancellation.",
)
else:
progress = list(task.metadata.get("progress_log", []) or [])
progress.append("Work item suspended by runtime cancellation.")
task.metadata["progress_log"] = progress[-20:]
try:
work_item = (
await self.store.get_delegation_work_item(work_item_id)
if work_item_id and store_ready and hasattr(self.store, "get_delegation_work_item")
else None
)
if work_item is not None and getattr(work_item, "phase", None) not in {Phase.APPROVED, Phase.FAILED, Phase.CANCELLED}:
phase = getattr(work_item, "phase", Phase.RUNNING)
phase_value = phase.value if isinstance(phase, Phase) else str(phase or "")
original_claim = {
"claimed_by_role_runtime_session_id": str(getattr(work_item, "claimed_by_role_runtime_session_id", "") or ""),
"claimed_by_seat_id": str(getattr(work_item, "claimed_by_seat_id", "") or ""),
"claimed_by_role_session_id": str((getattr(work_item, "metadata", {}) or {}).get("claimed_by_role_session_id", "") or ""),
"claimed_task_id": str((getattr(work_item, "metadata", {}) or {}).get("claimed_task_id", "") or task.id),
}
await self.store.update_delegation_work_item(
work_item_id,
metadata_updates={
"dispatch_hold": "company_runtime_suspended",
"suspended_at": datetime.now().isoformat(),
"suspend_reason": task.metadata.get("last_stop_reason", "runtime_cancelled"),
"suspended_phase": phase_value,
"suspended_task_status": task.metadata.get("suspended_task_status", ""),
"suspended_claim": original_claim,
"claimed_by_role_session_id": "",
"claimed_task_id": "",
},
claimed_by_role_runtime_session_id="",
claimed_by_seat_id="",
)
task.metadata["dispatch_hold"] = "company_runtime_suspended"
task.metadata["suspended_phase"] = phase_value
task.status = task_status_for_phase(phase) if isinstance(phase, Phase) else task.status
except Exception:
logger.opt(exception=True).debug(
"company runtime cancellation: failed to apply suspend hold",
)
if self.save_task and self._store_is_ready(self.store):
await self.save_task(task)
# Suspension is a checkpoint transition owned by OPCEngine.
# This task object may be stale by the time cancellation is
# observed, so persisting it here can erase the canonical
# checkpoint type, stop intent, or WorkItem hold.
raise
except asyncio.TimeoutError:
logger.error(f"Company work item {projection_id} timed out after {self.work_item_timeout}s")
+73 -35
View File
@@ -1373,7 +1373,6 @@ class CompanyRuntime:
_skip("no task materialized for work_item this tick",
session=session_label, work_item_id=work_item_id)
continue
self._claimed_work_item_ids.add(work_item_id)
else:
task_id = queued_item_id
self._queued_task_ids.discard(task_id)
@@ -1392,6 +1391,30 @@ class CompanyRuntime:
status=getattr(task, "status", None))
continue
self._claimed_task_ids.add(task_id)
if work_item is not None:
claimed = await self._claim_role_session_work_item(
session,
work_item,
task,
)
if not claimed:
fresh_work_item = None
get_work_item = getattr(
self.store,
"get_delegation_work_item",
None,
)
if callable(get_work_item):
fresh_work_item = await get_work_item(work_item_id)
if fresh_work_item is not None:
work_item_map[work_item_id] = fresh_work_item
_skip(
"atomic WorkItem claim lost to a phase/hold/owner update",
session=session_label,
work_item_id=work_item_id,
)
continue
self._claimed_work_item_ids.add(work_item_id)
if can_soft_wake and (
bool((task.metadata or {}).get("review_task", False))
or bool((task.metadata or {}).get("review_execution_work_item", False))
@@ -1412,7 +1435,6 @@ class CompanyRuntime:
self.prepare_task_for_session(session, task)
await self._sync_current_turn_mode_to_work_item(task, session.current_turn_mode)
if work_item is not None:
await self._claim_role_session_work_item(session, work_item, task)
# #7: mirror the claim onto the in-memory work_item so
# subsequent iterations in the same claim pass see it
# and is_dispatchable returns False (race safety after
@@ -1908,21 +1930,65 @@ class CompanyRuntime:
self.role_sessions[role_session_id] = role_session
return role_session
async def _claim_role_session_work_item(self, session: CompanyMemberSession, work_item: Any, task: Task) -> None:
async def _claim_role_session_work_item(
self,
session: CompanyMemberSession,
work_item: Any,
task: Task,
) -> bool:
"""Atomically claim ``work_item`` for the role-instance behind
``session``.
In the role-instance model the claim identity is
``role_runtime_session_id``. Seat / manager-seat columns are
still written for org-chart lookups but they are NOT part of
the claim key — only the role session is.
the claim key — only the role session is. Pure in-memory runtimes have
no durable race to arbitrate and keep the same local claim semantics.
"""
work_item_id = str(getattr(work_item, "work_item_id", "") or "").strip()
if not work_item_id:
return
return False
role_session = self._ensure_role_session(task)
if role_session is None:
return
return False
work_item_revision = 0
try:
work_item_revision = int((getattr(work_item, "metadata", {}) or {}).get("manager_mutation_revision") or 0)
except (TypeError, ValueError):
work_item_revision = 0
store_ready = self.store is not None and bool(
getattr(self.store, "is_ready", False)
)
claim = (
getattr(self.store, "claim_delegation_work_item_if_dispatchable", None)
if store_ready
else None
)
if store_ready:
# A durable runtime must win the store CAS before it mutates any
# in-memory scheduling state. This is the Stop/shutdown race
# boundary: a missing CAS API is a failed claim, not permission to
# fall back to the former best-effort update path.
if not callable(claim):
return False
persisted = await claim(
work_item_id,
expected_phase=getattr(work_item, "phase", Phase.READY),
role_runtime_session_id=role_session.role_session_id,
seat_id=str(getattr(session, "seat_id", "") or "").strip(),
task_id=task.id,
work_item_revision=work_item_revision,
)
if persisted is None:
return False
work_item.phase = persisted.phase
work_item.role_runtime_session_id = persisted.role_runtime_session_id
work_item.claimed_by_role_runtime_session_id = (
persisted.claimed_by_role_runtime_session_id
)
work_item.claimed_by_seat_id = persisted.claimed_by_seat_id
work_item.metadata = dict(persisted.metadata or {})
ready_background_ids = [
item_id
for item_id in list(role_session.background_work_item_ids or [])
@@ -1937,39 +2003,11 @@ class CompanyRuntime:
session.background_work_item_ids = list(role_session.background_work_item_ids)
task.metadata = dict(task.metadata)
task.metadata["delegation_role_session_id"] = role_session.role_session_id
work_item_revision = 0
try:
work_item_revision = int((getattr(work_item, "metadata", {}) or {}).get("manager_mutation_revision") or 0)
except (TypeError, ValueError):
work_item_revision = 0
task.metadata["started_work_item_revision"] = work_item_revision
task.metadata["claimed_work_item_revision"] = work_item_revision
if self.store and bool(getattr(self.store, "is_ready", False)) and hasattr(self.store, "update_delegation_work_item"):
# Do not regress a work item that is already in a review
# phase (AWAITING_MANAGER_REVIEW / AWAITING_HUMAN) back to
# RUNNING: the DB phase validator rejects that transition
# and the error bubbles up to the session loop. This
# happens when the reactivation sweeper wakes a task whose
# work item has already been promoted to review — treat
# the claim as "refresh the task/role-session bindings
# only" and leave the phase alone.
current_phase = getattr(work_item, "phase", None)
phase_to_write: Phase | None = Phase.RUNNING
if current_phase in IN_REVIEW_PHASES:
phase_to_write = None
await self.store.update_delegation_work_item(
work_item_id,
phase=phase_to_write,
role_runtime_session_id=role_session.role_session_id,
claimed_by_role_runtime_session_id=role_session.role_session_id,
metadata_updates={
"claimed_by_role_session_id": role_session.role_session_id,
"claimed_task_id": task.id,
"claimed_work_item_revision": work_item_revision,
},
)
if self.store and bool(getattr(self.store, "is_ready", False)) and hasattr(self.store, "save_delegation_role_session"):
await self.store.save_delegation_role_session(role_session)
return True
def ensure_role_instance_session(
self, task: Task
@@ -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)
@@ -75,6 +75,8 @@ class CustomRuntimeRunner:
store=shared_store,
owns_store=shared_store is None,
run_startup_reconcile=shared_store is None,
active_task_run_registry=getattr(self.parent, "_active_task_run_registry", None),
owns_active_task_run_registry=False,
on_progress=self.parent.on_progress,
on_runtime_event=self.parent.on_runtime_event,
on_escalation=self.parent.on_escalation,
+4 -2
View File
@@ -144,8 +144,10 @@ class EngineSeatExecutor:
member_session: CompanyMemberSession | None = None,
) -> None:
_ = member_session
if hasattr(self.host, "_active_task_runs"):
self.host._active_task_runs.discard(task.id)
# The task coroutine owns its registry attempt token and removes it in
# ``finally``. Interrupt requests must not make a still-running
# coroutine appear inactive.
return None
async def shutdown(
self,