fix: unify company runtime recovery lifecycle
This commit is contained in:
@@ -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")
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user