From 734a2d5969a508bc45a58608e64a18fb99306008 Mon Sep 17 00:00:00 2001 From: cgycorey <4724788+cgycorey@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:00:32 +0100 Subject: [PATCH] fix(ui): preserve durable org identity for runtime approvals --- opc/engine.py | 57 +++++- opc/layer2_organization/company_mode.py | 42 ++++- opc/plugins/office_ui/execution_identity.py | 15 +- opc/plugins/office_ui/services/session.py | 25 +-- opc/plugins/office_ui/ws_handler.py | 129 +++++++++++-- tests/test_company_runtime_identity.py | 57 ++++++ tests/test_lock_free_checkpoint_answer.py | 35 +++- tests/test_office_execution_identity.py | 54 ++++++ tests/test_office_session_org_fallback.py | 36 ++-- tests/test_runtime_config_enforcement.py | 194 ++++++++++++++++++++ tests/test_session_integration.py | 142 +++++++++++++- 11 files changed, 719 insertions(+), 67 deletions(-) diff --git a/opc/engine.py b/opc/engine.py index 12d0c99..097faa0 100644 --- a/opc/engine.py +++ b/opc/engine.py @@ -3563,6 +3563,33 @@ class OPCEngine: return False return bool(dict(getattr(task, "metadata", {}) or {}).get("shared_role_session", False)) + @staticmethod + def _runtime_org_id_for_identity( + decision: RouterDecision | None, + metadata: dict[str, Any] | None, + org_config: Any | None, + ) -> str | None: + """Return the durable custom-org ID for a company runtime task.""" + task_metadata = dict(metadata or {}) + profile = str( + getattr(decision, "company_profile", "") + or task_metadata.get("company_profile", "") + or getattr(org_config, "company_profile", "") + or "" + ).strip().lower() + if profile != "custom": + return None + for candidate in ( + getattr(decision, "org_id", None), + task_metadata.get("org_id"), + task_metadata.get("organization_id"), + getattr(org_config, "organization_id", None), + ): + normalized = str(candidate or "").strip() + if normalized: + return normalized + return None + @staticmethod def _shared_company_role_session_id( parent_session_id: str, @@ -3601,6 +3628,11 @@ class OPCEngine: root_session: bool = False, ) -> Task: assert self.store and self.memory + runtime_org_id = self._runtime_org_id_for_identity( + decision, + getattr(work_item, "metadata", None), + getattr(getattr(self, "config", None), "org", None), + ) role_id = str(work_item.role_id or "").strip() seat_id = str((work_item.metadata or {}).get("seat_id", "") or "").strip() team_id = str((work_item.metadata or {}).get("team_id", "") or work_item.cell_id or "").strip() @@ -3653,6 +3685,10 @@ class OPCEngine: set_linked_work_item_id(existing, work_item.work_item_id) existing.session_id = session_id existing.metadata = dict(existing.metadata or {}) + if runtime_org_id: + existing.org_id = runtime_org_id + existing.metadata["org_id"] = runtime_org_id + existing.metadata["organization_id"] = runtime_org_id existing.metadata["shared_role_session"] = True existing.metadata["shared_role_id"] = role_id existing.metadata["company_runtime_root_session_id"] = parent_session_id @@ -3744,6 +3780,24 @@ class OPCEngine: owner_execution_copy = build_work_item_owner_execution_copy(work_item) owner_execution_copy.setdefault("delegation_role_session_id", role_session_id) owner_execution_copy["work_kind"] = work_item_turn_type + runtime_company_profile = str( + getattr(decision, "company_profile", "") + or (work_item.metadata or {}).get("company_profile", "") + or getattr(getattr(self.config, "org", None), "company_profile", "") + or "" + ).strip().lower() + runtime_identity_metadata = ( + { + "org_id": runtime_org_id or "", + "organization_id": runtime_org_id or "", + } + if runtime_company_profile == "custom" + else { + "organization_id": str( + getattr(getattr(self.config, "org", None), "organization_id", "") or "" + ).strip(), + } + ) task = Task( title=str(work_item.title or work_item_projection_ref or "Runtime Work Item").strip(), description=( @@ -3757,6 +3811,7 @@ class OPCEngine: session_id=session_id, parent_session_id=parent_session_id, assigned_external_agent=assigned_external_agent, + org_id=runtime_org_id, metadata=mark_work_item_projection(mark_work_item_runtime({ "mode": "company", "execution_mode": decision.mode.value, @@ -3765,7 +3820,6 @@ class OPCEngine: "original_message": original_message, "router_preferred_agent": decision.preferred_agent, "company_profile": decision.company_profile or getattr(self.config.org, "company_profile", "corporate"), - "organization_id": getattr(self.config.org, "organization_id", ""), "organization_name": getattr(self.config.org, "organization_name", ""), "organization_config_file": getattr(self.config.org, "organization_config_file", ""), "delegation_playbook": dict(delegation_playbook), @@ -3780,6 +3834,7 @@ class OPCEngine: ), "runtime_topology": copy.deepcopy(runtime_topology), **owner_execution_copy, + **runtime_identity_metadata, "work_item_projection_ref": work_item_projection_ref, "seat_manager_role_id": str(topology_seat.get("manager_role_id", "") or "").strip(), "manager_role_id": str(topology_seat.get("manager_role_id", "") or "").strip(), diff --git a/opc/layer2_organization/company_mode.py b/opc/layer2_organization/company_mode.py index eafea48..2fce38a 100644 --- a/opc/layer2_organization/company_mode.py +++ b/opc/layer2_organization/company_mode.py @@ -4247,7 +4247,21 @@ class CompanyWorkItemExecutor: existing_task_ids = {str(task.id or "").strip() for task in existing_tasks if str(task.id or "").strip()} existing_work_item_ids = set(task_by_linked_work_item_id(existing_tasks)) root_task = sorted(existing_tasks, key=lambda item: (item.created_at, item.id))[0] - runtime_topology = dict((root_task.metadata or {}).get("runtime_topology", {}) or {}) + root_metadata = dict(root_task.metadata or {}) + custom_runtime = str(root_metadata.get("company_profile", "") or "").strip().lower() == "custom" + runtime_org_id = str( + getattr(root_task, "org_id", "") + or root_metadata.get("org_id") + or root_metadata.get("organization_id") + or "" + ).strip() or None + if not custom_runtime: + runtime_org_id = None + if runtime_org_id: + for existing_task in existing_tasks: + if self._sync_runtime_org_identity(existing_task, runtime_org_id): + await self.store.save_task(existing_task) + runtime_topology = dict(root_metadata.get("runtime_topology", {}) or {}) root_parent_session_id = str( root_task.parent_session_id or root_task.session_id @@ -4283,6 +4297,8 @@ class CompanyWorkItemExecutor: persisted = await get_runtime_task(work_item_id) if persisted is not None: set_linked_work_item_id(persisted, work_item_id) + if self._sync_runtime_org_identity(persisted, runtime_org_id): + await self.store.save_task(persisted) self._raise_for_runtime_projection_issues(persisted, work_item, work_item_by_id) if persisted.id not in existing_task_ids: existing_tasks.append(persisted) @@ -4413,6 +4429,9 @@ class CompanyWorkItemExecutor: task_metadata.update(copy_work_item_execution_metadata(work_item)) task_metadata.update(owner_execution_copy) task_metadata[WORK_ITEM_TURN_TYPE_KEY] = turn_type + if custom_runtime: + task_metadata["org_id"] = runtime_org_id or "" + task_metadata["organization_id"] = runtime_org_id or "" temp_task = Task( id=str(uuid.uuid4()), title=str(getattr(work_item, "title", "") or projection_id or "Runtime Work Item").strip(), @@ -4423,6 +4442,7 @@ class CompanyWorkItemExecutor: session_id=session_id, parent_session_id=root_parent_session_id, assigned_external_agent=assigned_external_agent, + org_id=runtime_org_id, metadata=task_metadata, ) dependency_projection_ids: list[str] = [] @@ -4454,6 +4474,7 @@ class CompanyWorkItemExecutor: parent_session_id=temp_task.parent_session_id, assigned_external_agent=temp_task.assigned_external_agent, dependencies=dependency_projection_ids, + org_id=runtime_org_id, metadata=task_metadata, ) set_linked_work_item_id(task, work_item_id) @@ -4471,6 +4492,8 @@ class CompanyWorkItemExecutor: "failed to link new runtime Task " f"{task.id} for WorkItem {work_item_id}" ) + if self._sync_runtime_org_identity(task, runtime_org_id): + await self.store.save_task(task) set_linked_work_item_id(task, work_item_id) self._raise_for_runtime_projection_issues(task, work_item, work_item_by_id) if self.memory is not None and task.session_id: @@ -4510,6 +4533,23 @@ class CompanyWorkItemExecutor: await self.save_task(task) return existing_tasks + @staticmethod + def _sync_runtime_org_identity(task: Task, org_id: str | None) -> bool: + normalized_org_id = str(org_id or "").strip() + if not normalized_org_id: + return False + metadata = dict(task.metadata or {}) + changed = str(getattr(task, "org_id", "") or "").strip() != normalized_org_id + changed = changed or metadata.get("org_id") != normalized_org_id + changed = changed or metadata.get("organization_id") != normalized_org_id + if not changed: + return False + task.org_id = normalized_org_id + metadata["org_id"] = normalized_org_id + metadata["organization_id"] = normalized_org_id + task.metadata = metadata + return True + @staticmethod def _runtime_work_kind_to_work_item_turn_type(work_kind: str) -> str: return canonical_work_item_turn_type_for_kind(work_kind) diff --git a/opc/plugins/office_ui/execution_identity.py b/opc/plugins/office_ui/execution_identity.py index aaab91a..7201ab7 100644 --- a/opc/plugins/office_ui/execution_identity.py +++ b/opc/plugins/office_ui/execution_identity.py @@ -17,6 +17,7 @@ from dataclasses import dataclass from typing import Any from opc.core.config import validate_organization_id +from opc.layer2_organization.company_runtime_identity import is_company_runtime_task PREFERRED_AGENTS: frozenset[str] = frozenset({ "native", @@ -154,11 +155,12 @@ def execution_identity_from_task( company_profile = metadata.get("company_profile") metadata_profile = str(company_profile or "").strip().lower() metadata_org_id = ( - metadata.get("org_id") + getattr(task, "org_id", None) + or metadata.get("org_id") or metadata.get("organization_id") - or getattr(task, "org_id", None) or "" ) + mode_hint = str(metadata.get("mode", "") or "").strip().lower() if raw_exec_mode: exec_mode = raw_exec_mode @@ -169,6 +171,15 @@ def execution_identity_from_task( elif execution_mode == "company_mode" or metadata_profile: exec_mode = "company" explicit = True + elif ( + mode_hint in {"company", "org", "custom"} + or is_company_runtime_task(task) + ): + # Older company/runtime rows may only retain a mode marker or the + # runtime marker itself. Treat those rows as explicit company + # identity so they cannot fall through to task-mode/global defaults. + exec_mode = "org" if metadata_org_id else "company" + explicit = True elif metadata_org_id: exec_mode = "org" explicit = True diff --git a/opc/plugins/office_ui/services/session.py b/opc/plugins/office_ui/services/session.py index c4229ad..778b52e 100644 --- a/opc/plugins/office_ui/services/session.py +++ b/opc/plugins/office_ui/services/session.py @@ -618,25 +618,6 @@ class SessionService: default_preferred_agent=self.context.mode_state.task_preferred_agent, explicit_exec_mode=True, ) - if identity.is_custom_org and not identity.org_id: - # Role-task rows may lack org_id; mirror create()'s active-org fallback - fallback_org_id = "" - if self.context.get_active_saved_org_name is not None: - try: - fallback_org_id = await self.context.get_active_saved_org_name() - except Exception: - logger.opt(exception=True).debug( - "persist_session_config: failed to resolve active saved org" - ) - if fallback_org_id: - identity = canonicalize_execution_identity( - exec_mode=exec_mode, - company_profile=company_profile, - preferred_agent=preferred_agent, - org_id=fallback_org_id, - default_preferred_agent=self.context.mode_state.task_preferred_agent, - explicit_exec_mode=True, - ) if identity.is_custom_org and not identity.org_id: raise ServiceError("org_id_required", "org_id_required", { "task_id": str(getattr(task, "id", "") or ""), @@ -1182,6 +1163,12 @@ class SessionService: else: checkpoint = None org_id = self.resolve_task_org_id(config_task) if engine_mode == "org" else "" + if engine_mode == "org" and not org_id: + raise ServiceError( + "org_id_required", + "org_id_required", + {"project_id": project_id, "task_id": resolved_task_id}, + ) message_metadata: dict[str, Any] = {"ui_force_resume": True} if checkpoint is not None: message_metadata.update({ diff --git a/opc/plugins/office_ui/ws_handler.py b/opc/plugins/office_ui/ws_handler.py index 72e3dc1..cac83d2 100644 --- a/opc/plugins/office_ui/ws_handler.py +++ b/opc/plugins/office_ui/ws_handler.py @@ -4978,16 +4978,38 @@ class WSHandler: if task is None: return None exec_mode, _ = self._resolve_task_session_config(task) - if not self._is_company_session_exec_mode(exec_mode): + runtime_bound = is_company_runtime_task(task) + if not runtime_bound and not self._is_company_session_exec_mode(exec_mode): return task try: target = await self._resolve_company_runtime_target(task_id, engine=engine) - except Exception: + except ServiceError: + if runtime_bound: + raise logger.opt(exception=True).debug( "failed to resolve durable session config task" ) return task - return (target or {}).get("config_task") or task + except Exception as exc: + if runtime_bound: + raise ServiceError( + "company_runtime_identity_mismatch", + "Company runtime identity could not be resolved", + {"task_id": str(task_id or "").strip()}, + ) from exc + logger.opt(exception=True).debug( + "failed to resolve durable session config task" + ) + return task + if target is not None: + return target.get("config_task") or task + if runtime_bound: + raise ServiceError( + "company_runtime_identity_mismatch", + "Company runtime identity could not be resolved", + {"task_id": str(task_id or "").strip()}, + ) + return task @staticmethod def _is_company_session_exec_mode(exec_mode: Any) -> bool: @@ -5059,14 +5081,24 @@ class WSHandler: # Look up session_id from task session_id: str | None = None task = None + config_task = None preferred_agent = self._task_preferred_agent session_org_id = self._normalize_session_org_id(org_id) if task_id and getattr(engine, "store", None): task = await engine.store.get_task(task_id) if task: session_id = task.session_id - identity = self._resolve_task_identity( + + company_runtime_target: dict[str, Any] | None = None + try: + if task is not None: + config_task = await self._resolve_session_runtime_config_task( + task_id, task, + engine=engine, + ) + identity = self._resolve_task_identity( + config_task, default_exec_mode=mode, default_company_profile=profile, default_preferred_agent=preferred_agent, @@ -5076,9 +5108,12 @@ class WSHandler: profile = identity.company_profile session_org_id = identity.org_id preferred_agent = identity.preferred_agent - - company_runtime_target: dict[str, Any] | None = None - try: + if identity.is_custom_org and not identity.org_id: + raise ServiceError( + "org_id_required", + "org_id_required", + {"project_id": pid, "task_id": task_id}, + ) content = f"{title}\n{description}".strip() engine_mode, company_profile = self._resolve_engine_mode(mode, profile) engine_preferred_agent = preferred_agent if engine_mode == "project" else None @@ -8124,17 +8159,6 @@ class WSHandler: pid = self._normalize_project_id(run_project_id or getattr(run_engine, "project_id", None)) async with lock: try: - try: - await self._set_company_runtime_control( - target, - state="resuming", - checkpoint_id=str( - getattr(checkpoint, "checkpoint_id", "") or "" - ).strip(), - ) - except Exception: - logger.opt(exception=True).debug("failed to broadcast company suspend reply routing state") - config_task = target.get("config_task") session_exec_mode = self._normalize_session_exec_mode(self._exec_mode) session_company_profile = self._normalize_session_company_profile(self._company_profile) @@ -8146,6 +8170,27 @@ class WSHandler: session_exec_mode, session_company_profile, ) + if engine_mode == "org" and not session_org_id: + raise ServiceError( + "org_id_required", + "org_id_required", + { + "project_id": pid, + "task_id": str(target.get("config_source_task_id", "") or "").strip(), + }, + ) + + try: + await self._set_company_runtime_control( + target, + state="resuming", + checkpoint_id=str( + getattr(checkpoint, "checkpoint_id", "") or "" + ).strip(), + ) + except Exception: + logger.opt(exception=True).debug("failed to broadcast company suspend reply routing state") + engine_message_metadata = dict(message_metadata or {}) engine_message_metadata.update({ "response_to_checkpoint_id": str(getattr(checkpoint, "checkpoint_id", "") or ""), @@ -8612,13 +8657,37 @@ class WSHandler: item for item in pending or [] if str(getattr(item, "checkpoint_id", "") or "").strip() == checkpoint_id - and str(getattr(item, "checkpoint_type", "") or "").strip() - in self._LOCK_FREE_CHECKPOINT_ANSWER_TYPES + and str(getattr(item, "checkpoint_type", "") or "").strip() == checkpoint_type ), None, ) if checkpoint is None: return False + checkpoint_payload = dict(getattr(checkpoint, "payload", {}) or {}) + checkpoint_task_id = str( + getattr(checkpoint, "task_id", "") + or checkpoint_payload.get("waiting_task_id") + or checkpoint_payload.get("task_id") + or "" + ).strip() + checkpoint_session_id = str( + getattr(checkpoint, "session_id", "") + or checkpoint_payload.get("session_id") + or "" + ).strip() + payload_task_ids = { + str(item).strip() + for item in list(checkpoint_payload.get("task_ids", []) or []) + if str(item).strip() + } + if ( + not checkpoint_task_id + or checkpoint_task_id != str(task_id or "").strip() + or not checkpoint_session_id + or checkpoint_session_id != str(session_id or "").strip() + or (payload_task_ids and str(task_id or "").strip() not in payload_task_ids) + ): + return False logger.info( f"Lock-free checkpoint answer: task lock for {task_id} is held by a " f"live turn; delivering reply to pending checkpoint {checkpoint_id} " @@ -8732,6 +8801,16 @@ class WSHandler: ) session_org_id = self._resolve_task_org_id(config_task) session_preferred_agent = self._resolve_task_preferred_agent(config_task) + if ( + session_exec_mode in {"org", "custom"} + and not session_org_id + and is_company_runtime_task(task) + ): + raise ServiceError( + "org_id_required", + "org_id_required", + {"project_id": pid, "task_id": task_id}, + ) if await self._try_lock_free_parked_checkpoint_answer( task_id=task_id, @@ -8769,6 +8848,16 @@ class WSHandler: ) session_org_id = self._resolve_task_org_id(config_task) session_preferred_agent = self._resolve_task_preferred_agent(config_task) + if ( + session_exec_mode in {"org", "custom"} + and not session_org_id + and is_company_runtime_task(task) + ): + raise ServiceError( + "org_id_required", + "org_id_required", + {"project_id": pid, "task_id": task_id}, + ) if task.status == TaskStatus.DONE and self._is_company_session_exec_mode(session_exec_mode): task.status = TaskStatus.IDLE task.metadata = dict(getattr(task, "metadata", {}) or {}) diff --git a/tests/test_company_runtime_identity.py b/tests/test_company_runtime_identity.py index abf82fe..dd8aa2b 100644 --- a/tests/test_company_runtime_identity.py +++ b/tests/test_company_runtime_identity.py @@ -260,6 +260,63 @@ def test_work_item_chat_resume_uses_canonical_ui_anchor_as_engine_origin() -> No asyncio.run(scenario()) +def test_suspend_reply_missing_custom_org_id_fails_closed_before_engine_call() -> None: + async def scenario() -> None: + _tasks, checkpoint = _runtime_records() + handler = WSHandler.__new__(WSHandler) + handler._exec_mode = "task" + handler._company_profile = "corporate" + handler._shutting_down = False + handler.chat_store = None + handler._company_suspend_reply_locks = { + "runtime-session": asyncio.Lock(), + } + handler._task_bg_context = {} + handler._active_runtime_children = {} + handler._session_to_task = {} + handler._set_company_runtime_control = AsyncMock() + handler._normalize_session_exec_mode = MagicMock(return_value="task") + handler._normalize_session_company_profile = MagicMock(return_value="corporate") + handler._resolve_task_session_config = MagicMock(return_value=("org", "custom")) + handler._resolve_task_org_id = MagicMock(return_value="") + handler._extract_checkpoint_metadata = AsyncMock(return_value=None) + handler._sync_task_transcript_messages = AsyncMock() + handler.on_kanban_changed = AsyncMock() + handler._flush_progress = AsyncMock() + run_engine = SimpleNamespace( + project_id="project-a", + process_message=AsyncMock(), + ) + config_task = SimpleNamespace( + metadata={"exec_mode": "org", "company_profile": "custom"}, + org_id=None, + ) + target = { + "ui_anchor_task_id": "ui-anchor", + "config_task": config_task, + } + + await handler._process_company_suspend_reply( + ui_task_id="final-decider", + runtime_session_id="runtime-session", + content="continue", + attachment_refs=None, + message_metadata=None, + user_message_id=None, + user_message_created_at=None, + run_engine=run_engine, + run_project_id="project-a", + target=target, + checkpoint=checkpoint, + lock=handler._company_suspend_reply_locks["runtime-session"], + ) + + run_engine.process_message.assert_not_awaited() + handler._set_company_runtime_control.assert_not_awaited() + + asyncio.run(scenario()) + + def test_delivery_feedback_rejects_missing_canonical_identity_without_first_task_fallback() -> None: async def scenario() -> None: tasks, checkpoint = _runtime_records() diff --git a/tests/test_lock_free_checkpoint_answer.py b/tests/test_lock_free_checkpoint_answer.py index 09f0f35..82663b9 100644 --- a/tests/test_lock_free_checkpoint_answer.py +++ b/tests/test_lock_free_checkpoint_answer.py @@ -49,11 +49,20 @@ class _EngineStub: return self.reply -def _pending_checkpoint(checkpoint_id: str, checkpoint_type: str = "task_user_input") -> Any: +def _pending_checkpoint( + checkpoint_id: str, + checkpoint_type: str = "task_user_input", + *, + task_id: str = "chat-task", + session_id: str = "session-1", +) -> Any: return SimpleNamespace( checkpoint_id=checkpoint_id, checkpoint_type=checkpoint_type, status="pending", + task_id=task_id, + session_id=session_id, + payload={"task_ids": [task_id], "waiting_task_id": task_id, "session_id": session_id}, ) @@ -189,6 +198,30 @@ class LockFreeCheckpointAnswerTests(unittest.IsolatedAsyncioTestCase): holder.release_event.set() # type: ignore[attr-defined] await holder + async def test_lock_free_requires_exact_checkpoint_type_and_owner(self) -> None: + engine = _EngineStub( + _StoreStub([ + _pending_checkpoint( + "ckpt-park", + "company_work_item_gate", + task_id="other-task", + session_id="other-session", + ) + ]) + ) + handler = _make_handler(engine) + holder = await self._hold_lock(handler, "chat-task") + try: + handled = await handler._try_lock_free_parked_checkpoint_answer( + engine=engine, + **_answer_kwargs(), + ) + self.assertFalse(handled) + self.assertEqual(engine.calls, []) + finally: + holder.release_event.set() # type: ignore[attr-defined] + await holder + async def test_engine_failure_surfaces_error_without_queueing(self) -> None: engine = _EngineStub( _StoreStub([_pending_checkpoint("ckpt-park")]), diff --git a/tests/test_office_execution_identity.py b/tests/test_office_execution_identity.py index 5b1e4fa..e98bd89 100644 --- a/tests/test_office_execution_identity.py +++ b/tests/test_office_execution_identity.py @@ -112,6 +112,21 @@ def test_task_org_id_field_is_org_identity_fallback() -> None: assert identity.org_id == "quantum_harbor" +def test_durable_task_org_id_wins_over_stale_metadata_org_id() -> None: + task = SimpleNamespace( + metadata={ + "exec_mode": "org", + "company_profile": "custom", + "organization_id": "active-org", + }, + org_id="selected-org", + ) + + identity = execution_identity_from_task(task) + + assert identity.org_id == "selected-org" + + def test_default_org_id_applies_only_when_task_has_no_persisted_identity() -> None: task = SimpleNamespace(metadata={}, org_id=None) @@ -146,3 +161,42 @@ def test_explicit_company_identity_ignores_default_org_id() -> None: assert identity.exec_mode == "company" assert identity.company_profile == "corporate" assert identity.org_id == "" + + +def test_company_mode_marker_without_profile_is_company_identity() -> None: + task = SimpleNamespace( + metadata={"mode": "company"}, + org_id=None, + ) + + identity = execution_identity_from_task(task) + + assert identity.exec_mode == "company" + assert identity.company_profile == "corporate" + assert identity.org_id == "" + + +def test_work_item_runtime_marker_without_profile_is_company_identity() -> None: + task = SimpleNamespace( + metadata={"work_item_runtime": True}, + org_id=None, + ) + + identity = execution_identity_from_task(task) + + assert identity.exec_mode == "company" + assert identity.company_profile == "corporate" + assert identity.org_id == "" + + +def test_company_runtime_marker_with_custom_profile_is_org_identity() -> None: + task = SimpleNamespace( + metadata={"mode": "company", "work_item_runtime": True, "company_profile": "custom"}, + org_id=None, + ) + + identity = execution_identity_from_task(task) + + assert identity.exec_mode == "org" + assert identity.company_profile == "custom" + assert identity.org_id == "" diff --git a/tests/test_office_session_org_fallback.py b/tests/test_office_session_org_fallback.py index 78dfa01..1f4fbd1 100644 --- a/tests/test_office_session_org_fallback.py +++ b/tests/test_office_session_org_fallback.py @@ -3,6 +3,7 @@ from __future__ import annotations import unittest from types import SimpleNamespace from typing import Any +from unittest.mock import AsyncMock from opc.plugins.office_ui.services.context import OfficeServiceContext from opc.plugins.office_ui.services.models import ServiceError @@ -40,7 +41,7 @@ def _context(*, hook: Any | None = None) -> OfficeServiceContext: return context -class TestPersistSessionConfigOrgFallback(unittest.IsolatedAsyncioTestCase): +class TestPersistSessionConfigOrgIdentity(unittest.IsolatedAsyncioTestCase): async def _persist( self, context: OfficeServiceContext, @@ -58,26 +59,20 @@ class TestPersistSessionConfigOrgFallback(unittest.IsolatedAsyncioTestCase): org_id=org_id, ) - async def test_falls_back_to_active_saved_org_when_task_lacks_org_id(self) -> None: - async def active_org() -> str: - return "vc-investment-firm" - + async def test_rejects_missing_org_id_without_active_org_fallback(self) -> None: + active_org = AsyncMock(return_value="vc-investment-firm") context = _context(hook=active_org) task = _task() - await self._persist(context, task) + with self.assertRaises(ServiceError) as ctx: + await self._persist(context, task) - assert task.metadata["org_id"] == "vc-investment-firm" - assert task.metadata["organization_id"] == "vc-investment-firm" - assert task.org_id == "vc-investment-firm" - assert task.metadata["exec_mode"] == "org" - assert task.metadata["company_profile"] == "custom" - assert context.engine.store.saved == [task] + assert ctx.exception.code == "org_id_required" + active_org.assert_not_awaited() + assert context.engine.store.saved == [] async def test_explicit_org_id_still_used_when_present(self) -> None: - async def active_org() -> str: - return "other-org" - + active_org = AsyncMock(return_value="other-org") context = _context(hook=active_org) task = _task() @@ -85,17 +80,17 @@ class TestPersistSessionConfigOrgFallback(unittest.IsolatedAsyncioTestCase): assert task.metadata["org_id"] == "vc-investment-firm" assert task.org_id == "vc-investment-firm" + active_org.assert_not_awaited() async def test_raises_when_no_active_org_available(self) -> None: - async def empty_org() -> str: - return "" - + empty_org = AsyncMock(return_value="") context = _context(hook=empty_org) task = _task() with self.assertRaises(ServiceError) as ctx: await self._persist(context, task) assert ctx.exception.code == "org_id_required" + empty_org.assert_not_awaited() async def test_raises_when_hook_unset(self) -> None: context = _context() @@ -106,15 +101,14 @@ class TestPersistSessionConfigOrgFallback(unittest.IsolatedAsyncioTestCase): assert ctx.exception.code == "org_id_required" async def test_raises_when_hook_fails(self) -> None: - async def broken() -> str: - raise RuntimeError("org index unreadable") - + broken = AsyncMock(side_effect=RuntimeError("org index unreadable")) context = _context(hook=broken) task = _task() with self.assertRaises(ServiceError) as ctx: await self._persist(context, task) assert ctx.exception.code == "org_id_required" + broken.assert_not_awaited() async def test_company_mode_clears_org_fields_without_fallback(self) -> None: fallback_called = False diff --git a/tests/test_runtime_config_enforcement.py b/tests/test_runtime_config_enforcement.py index db0c969..7c4b530 100644 --- a/tests/test_runtime_config_enforcement.py +++ b/tests/test_runtime_config_enforcement.py @@ -323,6 +323,74 @@ class RuntimeConfigEnforcementTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(task.metadata["preferred_external_agent"], "opencode") self.assertEqual(task.metadata["work_item_execution_strategy"], WorkItemExecutionStrategy.EXTERNAL.value) + async def test_company_root_task_persists_decision_org_over_active_config(self) -> None: + engine = OPCEngine(config=OPCConfig(), project_id="proj1") + engine.config.org.company_profile = "custom" + engine.config.org.organization_id = "active-org" + engine.store = SimpleNamespace( + get_runtime_task_for_work_item=AsyncMock(return_value=None), + save_delegation_work_item=AsyncMock(), + save_task=AsyncMock(), + link_work_item_runtime_task=AsyncMock(return_value=True), + ) + engine.memory = SimpleNamespace(ensure_session=AsyncMock()) + engine.org_engine = SimpleNamespace( + current_org_version=MagicMock(return_value=1), + current_runtime_topology_version=MagicMock(return_value=1), + ) + engine._requests_explicit_project_knowledge = MagicMock(return_value=False) + work_item = DelegationWorkItem( + work_item_id="wi-selected-org", + run_id="run-selected-org", + cell_id="team::engineering", + team_instance_id="team-instance-1", + role_id="engineer", + seat_id="seat-engineer", + title="Engineering execution", + summary="Implement the requested change.", + kind="execute", + projection_id="engineering-execute", + metadata={"seat_id": "seat-engineer", "team_id": "team::engineering"}, + ) + + task = await engine._ensure_runtime_work_item_task( + work_item=work_item, + parent_session_id="sess-company", + original_message="Build the thing.", + decision=RouterDecision( + mode=ExecutionMode.COMPANY_MODE, + domains=[], + company_profile="custom", + org_id="selected-org", + ), + runtime_topology={ + "final_decider_role_id": "lead", + "seats": [ + { + "seat_id": "seat-engineer", + "team_id": "team::engineering", + "role_id": "engineer", + "employee_assignment": {"employee_id": "eng-1", "name": "Engineer"}, + "metadata": {"role_name": "Engineer"}, + } + ], + }, + delegation_playbook={}, + secretary_context="", + target_output_dir=None, + origin_channel="cli", + origin_chat_id="", + origin_thread_id="", + origin_task_id=None, + attachment_refs=[], + attachment_context="", + force_native_execution=False, + ) + + self.assertEqual(task.org_id, "selected-org") + self.assertEqual(task.metadata["org_id"], "selected-org") + self.assertEqual(task.metadata["organization_id"], "selected-org") + async def test_company_materialized_work_item_uses_selected_agent_over_template_preference(self) -> None: saved_tasks: list[Task] = [] @@ -389,6 +457,132 @@ class RuntimeConfigEnforcementTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(created.metadata["work_item_execution_strategy"], WorkItemExecutionStrategy.EXTERNAL.value) self.assertEqual(saved_tasks[0].assigned_external_agent, "cursor") + async def test_company_materialization_persists_durable_org_on_new_role_task(self) -> None: + executor = CompanyWorkItemExecutor( + org_engine=SimpleNamespace(), + communication=SimpleNamespace(), + approval_engine=SimpleNamespace(), + memory=SimpleNamespace(ensure_session=AsyncMock()), + execute_task=AsyncMock(), + save_task=AsyncMock(), + ) + executor.store = SimpleNamespace( + get_runtime_task_for_work_item=AsyncMock(return_value=None), + save_delegation_work_item=AsyncMock(), + save_task=AsyncMock(), + link_work_item_runtime_task=AsyncMock(return_value=True), + ) + root_task = Task( + id="root-custom-org", + title="Root", + project_id="proj1", + session_id="sess-company", + org_id="selected-org", + metadata={ + "execution_mode": "company_mode", + "runtime_model": "multi_team_org", + "company_profile": "custom", + "organization_id": "active-org", + "runtime_topology": { + "seats": [ + { + "seat_id": "seat-engineer", + "team_id": "team::engineering", + "role_id": "engineer", + "metadata": {"role_name": "Engineer"}, + } + ] + }, + }, + ) + work_item = DelegationWorkItem( + work_item_id="wi-new-custom-org", + run_id="run-custom-org", + cell_id="team::engineering", + team_instance_id="team-instance-1", + role_id="engineer", + seat_id="seat-engineer", + title="Engineering execution", + summary="Implement the requested change.", + kind="execute", + projection_id="engineering-execute", + metadata={"seat_id": "seat-engineer", "team_id": "team::engineering"}, + ) + + tasks = await executor._materialize_work_item_tasks([root_task], [work_item]) + created = next(task for task in tasks if task.id != root_task.id) + + self.assertEqual(created.org_id, "selected-org") + self.assertEqual(created.metadata["org_id"], "selected-org") + self.assertEqual(created.metadata["organization_id"], "selected-org") + + async def test_company_materialization_repairs_existing_role_task_org_identity(self) -> None: + existing = Task( + id="existing-custom-role", + title="Engineering execution", + project_id="proj1", + session_id="sess-company:wi-existing-custom-org", + assigned_to="engineer", + org_id="active-org", + metadata={ + "execution_mode": "company_mode", + "runtime_model": "multi_team_org", + "work_item_runtime": True, + "work_item_projection_id": "engineering-execute", + "work_item_turn_type": "execute", + "company_profile": "custom", + "organization_id": "active-org", + "delegation_seat_id": "seat-engineer", + }, + ) + executor = CompanyWorkItemExecutor( + org_engine=SimpleNamespace(), + communication=SimpleNamespace(), + approval_engine=SimpleNamespace(), + memory=SimpleNamespace(ensure_session=AsyncMock()), + execute_task=AsyncMock(), + save_task=AsyncMock(), + ) + executor.store = SimpleNamespace( + get_runtime_task_for_work_item=AsyncMock(return_value=existing), + save_delegation_work_item=AsyncMock(), + save_task=AsyncMock(), + ) + root_task = Task( + id="root-custom-org-existing", + title="Root", + project_id="proj1", + session_id="sess-company", + org_id="selected-org", + metadata={ + "execution_mode": "company_mode", + "runtime_model": "multi_team_org", + "company_profile": "custom", + "organization_id": "active-org", + }, + ) + work_item = DelegationWorkItem( + work_item_id="wi-existing-custom-org", + run_id="run-custom-org", + cell_id="team::engineering", + team_instance_id="team-instance-1", + role_id="engineer", + seat_id="seat-engineer", + title="Engineering execution", + summary="Implement the requested change.", + kind="execute", + projection_id="engineering-execute", + metadata={"seat_id": "seat-engineer", "team_id": "team::engineering"}, + ) + + tasks = await executor._materialize_work_item_tasks([root_task], [work_item]) + repaired = next(task for task in tasks if task.id == existing.id) + + self.assertEqual(repaired.org_id, "selected-org") + self.assertEqual(repaired.metadata["org_id"], "selected-org") + self.assertEqual(repaired.metadata["organization_id"], "selected-org") + executor.store.save_task.assert_any_await(existing) + async def test_task_mode_external_followup_reuses_primary_session_external_agent_session(self) -> None: engine = OPCEngine(config=OPCConfig(), project_id="proj1") engine.store = SimpleNamespace( diff --git a/tests/test_session_integration.py b/tests/test_session_integration.py index 26b04e1..4d2e3b0 100644 --- a/tests/test_session_integration.py +++ b/tests/test_session_integration.py @@ -2007,8 +2007,10 @@ class TestWSHandlerSessionSend(unittest.IsolatedAsyncioTestCase): "exec_mode": "org", "mode": "company", "company_profile": "custom", - "org_id": "vc-investment-firm", + "org_id": "wrong-active-org", + "organization_id": "wrong-active-org", } + anchor.org_id = "vc-investment-firm" await self.store.save_task(anchor) role_task = Task( @@ -2017,10 +2019,13 @@ class TestWSHandlerSessionSend(unittest.IsolatedAsyncioTestCase): project_id="test-project", session_id=f"{runtime_session_id}:role:sector-analyst", parent_session_id=runtime_session_id, + org_id="vc-investment-firm", metadata={ "exec_mode": "org", "mode": "company", "company_profile": "custom", + "org_id": "wrong-active-org", + "organization_id": "wrong-active-org", "shared_role_session": True, "shared_role_id": "sector_analyst", "company_runtime_root_session_id": runtime_session_id, @@ -2043,6 +2048,73 @@ class TestWSHandlerSessionSend(unittest.IsolatedAsyncioTestCase): self.assertEqual(call_kwargs["company_profile"], "custom") self.assertEqual(call_kwargs["org_id"], "vc-investment-firm") + async def test_process_session_message_fails_closed_when_runtime_identity_is_missing(self) -> None: + task = await self.store.get_task(self.task_id) + assert task is not None + task.metadata = { + "mode": "company", + "company_profile": "custom", + } + await self.store.save_task(task) + self.handler._resolve_company_runtime_target = AsyncMock(return_value=None) + self.handler.services_context.get_active_saved_org_name = AsyncMock( + return_value="wrong-active-org" + ) + + with self.assertRaises(ServiceError) as context: + await self.handler._process_session_message( + self.task_id, + "approve", + session_id=self.session_id, + ) + + self.engine.process_message.assert_not_called() + self.handler.services_context.get_active_saved_org_name.assert_not_awaited() + self.assertEqual(context.exception.code, "company_runtime_identity_mismatch") + + async def test_process_session_message_rejects_runtime_org_without_durable_org_id(self) -> None: + runtime_session_id = "runtime-org-missing-id-session" + anchor = await self.store.get_task(self.task_id) + assert anchor is not None + anchor.session_id = runtime_session_id + anchor.metadata = { + "exec_mode": "org", + "mode": "company", + "company_profile": "custom", + } + await self.store.save_task(anchor) + + role_task = Task( + id="role-task-missing-org-id", + title="Sector Analyst", + project_id="test-project", + session_id=f"{runtime_session_id}:role:sector-analyst", + parent_session_id=runtime_session_id, + metadata={ + "exec_mode": "org", + "mode": "company", + "company_profile": "custom", + "shared_role_session": True, + "shared_role_id": "sector_analyst", + "company_runtime_root_session_id": runtime_session_id, + }, + ) + await self.store.save_task(role_task) + self.handler.services_context.get_active_saved_org_name = AsyncMock( + return_value="wrong-active-org" + ) + + with self.assertRaises(ServiceError) as context: + await self.handler._process_session_message( + role_task.id, + "approve", + session_id=role_task.session_id, + ) + + self.engine.process_message.assert_not_called() + self.handler.services_context.get_active_saved_org_name.assert_not_awaited() + self.assertEqual(context.exception.code, "org_id_required") + async def test_lock_free_process_session_message_uses_durable_org_for_role_task(self) -> None: runtime_session_id = "runtime-org-lock-free-session" anchor = await self.store.get_task(self.task_id) @@ -2075,10 +2147,15 @@ class TestWSHandlerSessionSend(unittest.IsolatedAsyncioTestCase): checkpoint = ExecutionCheckpoint( checkpoint_id="org-lock-free-gate", project_id="test-project", - session_id=runtime_session_id, + session_id=role_task.session_id, checkpoint_type="company_work_item_gate", status="pending", task_id=role_task.id, + payload={ + "waiting_task_id": role_task.id, + "task_ids": [role_task.id], + "session_id": role_task.session_id, + }, ) await self.store.save_execution_checkpoint(checkpoint) self.engine._load_execution_checkpoint_by_id = AsyncMock(return_value=checkpoint) @@ -5405,6 +5482,41 @@ class TestOfficeServiceExecutionIdentity(unittest.IsolatedAsyncioTestCase): assert persisted is not None self.assertEqual(persisted.status, TaskStatus.RUNNING) + async def test_continue_rejects_runtime_custom_org_without_durable_org_id(self) -> None: + task = Task( + id="service-continue-org-missing-id", + title="Broken custom runtime", + session_id="service-continue-org-runtime", + project_id="test-project", + metadata={"exec_mode": "org", "company_profile": "custom"}, + ) + checkpoint = ExecutionCheckpoint( + checkpoint_id="service-continue-org-checkpoint", + project_id="test-project", + session_id=task.session_id, + checkpoint_type="company_runtime_suspended", + status="pending", + task_id=task.id, + ) + await self.store.save_task(task) + self.session_service._resolve_company_runtime_target = AsyncMock(return_value={ + "runtime_session_id": task.session_id, + "ui_anchor_task_id": task.id, + "config_task": task, + "checkpoint": checkpoint, + "affected_task_ids": [task.id], + }) + + with self.assertRaises(ServiceError) as raised: + await self.session_service.continue_run( + project_id="test-project", + task_id=task.id, + content="continue", + ) + + self.assertEqual(raised.exception.code, "org_id_required") + self.engine.process_message.assert_not_called() + async def test_session_send_from_work_item_uses_runtime_checkpoint_identity(self) -> None: anchor = Task( id="service-send-anchor", @@ -5835,6 +5947,32 @@ class TestWSHandlerRunTask(unittest.IsolatedAsyncioTestCase): self.assertEqual(call_kwargs["company_profile"], "custom") self.assertEqual(call_kwargs["org_id"], "quantum_harbor") + async def test_run_task_fails_closed_for_custom_identity_without_org_id(self) -> None: + task_id = str(uuid.uuid4()) + session_id = str(uuid.uuid4()) + task = Task( + id=task_id, + title="Missing Org Runtime", + session_id=session_id, + project_id="test-project", + metadata={"exec_mode": "org", "company_profile": "custom"}, + ) + await self.store.save_task(task) + self.handler.services_context.get_active_saved_org_name = AsyncMock( + return_value="wrong-active-org" + ) + + await self.handler._run_task( + "Missing Org Runtime", + "Description", + "org", + "custom", + task_id, + ) + + self.engine.process_message.assert_not_called() + self.handler.services_context.get_active_saved_org_name.assert_not_awaited() + async def test_run_task_explicit_company_clears_stale_custom_fields(self) -> None: task_id = str(uuid.uuid4()) session_id = str(uuid.uuid4())