src/ai_company/api/controllers/tasks.py:136

• The controller claims “full CRUD via TaskEngine”, but there is no cancel endpoint and no API support for expected_version/optimistic concurrency (UpdateTaskRequest/TransitionTaskRequest don’t expose it and the controller doesn’t accept it). This diverges from the PR description; either add the missing endpoints/parameters or adjust the description/scope to match the actual API surface.

class TaskController(Controller):
    """Full CRUD for tasks via ``TaskEngine``."""

    path = "/tasks"
    tags = ("tasks",)
    guards = [require_read_access]  # noqa: RUF012

src/ai_company/engine/task_engine_models.py (2)

332-374: ⚠️ Potential issue | 🟠 Major

Add a stable task_id to TaskStateChanged.

task is explicitly None on delete, and _publish_snapshot() currently has no other identifier to serialize. Consumers cannot tell which task was removed, so delete snapshots are not actionable.

Suggested model change

 class TaskStateChanged(BaseModel):
@@
+    task_id: NotBlankStr = Field(description="Affected task identifier")
     task: Task | None = Field(
         default=None,
         description="Task snapshot after mutation",
     )

Populate it from result.task.id when present, or from mutation.task_id for deletes.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/ai_company/engine/task_engine_models.py` around lines 332 - 374,
TaskStateChanged needs a stable task identifier so delete events remain
actionable: add a new field task_id: NotBlankStr | None (or appropriate type) to
the TaskStateChanged model and ensure _publish_snapshot() populates it from
result.task.id when result.task exists, otherwise use mutation.task_id for
delete events; update any constructors/creation sites that build
TaskStateChanged to set task_id accordingly while leaving task (snapshot) None
on deletes.

---

148-170: ⚠️ Potential issue | 🟠 Major

Freeze updates and overrides after validation.

The deep copy only severs aliasing with caller input. Both fields remain directly mutable on the frozen model, so code can still modify a queued mutation after validation and before the engine consumes it.

Suggested fix

+from collections.abc import Mapping
+from types import MappingProxyType
@@
-    updates: dict[str, object] = Field(description="Field-value pairs to apply")
+    updates: Mapping[str, object] = Field(description="Field-value pairs to apply")
@@
-        object.__setattr__(self, "updates", copy.deepcopy(self.updates))
+        object.__setattr__(
+            self,
+            "updates",
+            MappingProxyType(copy.deepcopy(dict(self.updates))),
+        )
@@
-    overrides: dict[str, object] = Field(
+    overrides: Mapping[str, object] = Field(
         default_factory=dict,
         description="Additional field overrides",
     )
@@
-        object.__setattr__(self, "overrides", copy.deepcopy(self.overrides))
+        object.__setattr__(
+            self,
+            "overrides",
+            MappingProxyType(copy.deepcopy(dict(self.overrides))),
+        )

As per coding guidelines "Use immutability by creating new objects, never mutating existing ones. For non-Pydantic internal collections, use copy.deepcopy() at construction + MappingProxyType wrapping."

Also applies to: 202-227

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/ai_company/engine/task_engine_models.py` around lines 148 - 170, The
updates (and similarly overrides) dicts are deep-copied but left mutable on the
model; after validation freeze them by wrapping the deepcopy in
types.MappingProxyType and assign via object.__setattr__ so the model exposes an
immutable mapping; ensure MappingProxyType is imported and apply the same
pattern for both the "updates" field handling in __init__ and the analogous code
for "overrides" (referencing the _reject_immutable_fields validator, the
"updates" attribute, and the overrides handling block around lines ~202-227).

src/ai_company/engine/task_engine.py (3)

622-631: 🛠️ Refactor suggestion | 🟠 Major

Log successful mutation application at INFO.

This is the single authoritative writer, but _process_one() only emits a DEBUG receipt and optional snapshot logs. The committed mutation itself should produce an INFO event with task/version/status context.

Suggested fix

+# also import TASK_ENGINE_MUTATION_APPLIED from
+# ai_company.observability.events.task_engine
@@
             if not envelope.future.done():
                 envelope.future.set_result(result)
+            if result.success:
+                logger.info(
+                    TASK_ENGINE_MUTATION_APPLIED,
+                    mutation_type=mutation.mutation_type,
+                    request_id=mutation.request_id,
+                    task_id=(
+                        result.task.id
+                        if result.task is not None
+                        else getattr(mutation, "task_id", None)
+                    ),
+                    previous_status=(
+                        result.previous_status.value
+                        if result.previous_status is not None
+                        else None
+                    ),
+                    new_status=(
+                        result.task.status.value
+                        if result.task is not None
+                        else None
+                    ),
+                    version=result.version,
+                )
             if result.success and self._config.publish_snapshots:
                 await self._publish_snapshot(mutation, result)

Based on learnings "All state transitions must log at INFO level" and "Use DEBUG logging for object creation, internal flow, entry/exit of key functions."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/ai_company/engine/task_engine.py` around lines 622 - 631, Add an
INFO-level log when a mutation is successfully applied inside _process_one after
awaiting _dispatch_mutation (the block handling result and envelope.future);
specifically, when result is returned and before calling _publish_snapshot, emit
an info log that includes the mutation id/context (from mutation), task id or
name (if available on mutation), the resulting version (from self._versions or
result), and the result.status/success to provide authoritative state-transition
telemetry; keep the existing DEBUG receipt and snapshot logs intact and only add
this single INFO log entry adjacent to the existing envelope.future.set_result
and _publish_snapshot calls.

---

166-183: ⚠️ Potential issue | 🔴 Critical

Don't let wait_for() clear _in_flight before the timeout handler runs.

asyncio.wait_for(self._processing_task, ...) cancels the worker on timeout before the except TimeoutError block executes. By the time saved_in_flight is read, _process_one() has already hit its finally block and cleared it to None, so the blocked caller's future is never resolved.

Suggested fix

             try:
                 await asyncio.wait_for(
-                    self._processing_task,
+                    asyncio.shield(self._processing_task),
                     timeout=effective_timeout,
                 )
                 logger.info(TASK_ENGINE_DRAIN_COMPLETE)
             except TimeoutError:
                 logger.warning(
                     TASK_ENGINE_DRAIN_TIMEOUT,
                     remaining=self._queue.qsize(),
                 )
-                # Capture in-flight ref before cancel — the finally block
-                # in _process_one clears self._in_flight on CancelledError.
+                # Shield keeps the task alive long enough for us to capture
+                # the in-flight envelope before explicit cancellation.
                 saved_in_flight = self._in_flight
                 self._processing_task.cancel()
                 with contextlib.suppress(asyncio.CancelledError):
                     await self._processing_task
                 self._fail_remaining_futures(saved_in_flight)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/ai_company/engine/task_engine.py` around lines 166 - 183, The worker task
is being cancelled by asyncio.wait_for on timeout which lets _process_one reach
its finally block and clear self._in_flight before the except handler can
capture it; change the wait_for call to await
asyncio.wait_for(asyncio.shield(self._processing_task),
timeout=effective_timeout) so the underlying _processing_task is protected from
automatic cancellation, then in the TimeoutError handler capture saved_in_flight
= self._in_flight, cancel self._processing_task explicitly, await it inside
contextlib.suppress(asyncio.CancelledError), and call
_fail_remaining_futures(saved_in_flight) as you already do.

---

101-122: ⚠️ Potential issue | 🔴 Critical

Guard shutdown and enqueue with the same lifecycle lock.

A submitter can pass the running check, then stop() can drain and tear down the worker, and only afterward the envelope gets queued. That future then has no consumer and can hang forever.

Suggested fix

     def __init__(
         self,
         *,
         persistence: PersistenceBackend,
         message_bus: MessageBus | None = None,
         config: TaskEngineConfig | None = None,
     ) -> None:
+        self._lifecycle_lock = asyncio.Lock()
         self._persistence = persistence
         self._message_bus = message_bus
@@
-        if not self._running:
-            return
-        self._running = False
+        async with self._lifecycle_lock:
+            if not self._running or self._processing_task is None:
+                return
+            self._running = False
+            processing_task = self._processing_task
@@
-                await asyncio.wait_for(
-                    self._processing_task,
-                    timeout=effective_timeout,
-                )
+                await asyncio.wait_for(processing_task, timeout=effective_timeout)
@@
-        if not self._running:
-            logger.warning(
-                TASK_ENGINE_NOT_RUNNING,
-                mutation_type=mutation.mutation_type,
-                request_id=mutation.request_id,
-            )
-            msg = "TaskEngine is not running"
-            raise TaskEngineNotRunningError(msg)
-
-        envelope = _MutationEnvelope(mutation=mutation)
-        try:
-            self._queue.put_nowait(envelope)
-        except asyncio.QueueFull:
-            logger.warning(
-                TASK_ENGINE_QUEUE_FULL,
-                mutation_type=mutation.mutation_type,
-                request_id=mutation.request_id,
-                queue_size=self._queue.qsize(),
-            )
-            msg = "TaskEngine queue is full"
-            raise TaskEngineQueueFullError(msg) from None
+        async with self._lifecycle_lock:
+            if not self._running or self._processing_task is None:
+                logger.warning(
+                    TASK_ENGINE_NOT_RUNNING,
+                    mutation_type=mutation.mutation_type,
+                    request_id=mutation.request_id,
+                )
+                msg = "TaskEngine is not running"
+                raise TaskEngineNotRunningError(msg)
+
+            envelope = _MutationEnvelope(mutation=mutation)
+            try:
+                self._queue.put_nowait(envelope)
+            except asyncio.QueueFull:
+                logger.warning(
+                    TASK_ENGINE_QUEUE_FULL,
+                    mutation_type=mutation.mutation_type,
+                    request_id=mutation.request_id,
+                    queue_size=self._queue.qsize(),
+                )
+                msg = "TaskEngine queue is full"
+                raise TaskEngineQueueFullError(msg) from None

Also applies to: 146-184, 237-272

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/ai_company/engine/task_engine.py` around lines 101 - 122, The
submit/enqueue race occurs because the running check and queue.put are not
atomic with stop(); add a lifecycle asyncio.Lock (e.g., self._lifecycle_lock =
asyncio.Lock()) in __init__ and use it to guard the check-and-enqueue path in
the submit/enqueue functions and the shutdown path in stop() so that stop() can
drain/teardown only after acquiring the same lock; ensure you hold the lock
while checking self._running and performing self._queue.put_nowait/_queue.put
and while stop() flips self._running, cancels/awaits self._processing_task, and
drains the queue to avoid enqueuing envelopes with no consumer (also apply same
locking to other enqueue/submit sites referenced in the ranges 146-184 and
237-272).

tests/unit/engine/test_task_engine_integration.py (1)

99-115: 🧹 Nitpick | 🔵 Trivial

Prefer TaskGroup for this concurrent fan-out.

If one create fails before the final join, the remaining tasks keep running until gather(). TaskGroup gives the same coverage with structured cancellation and cleaner failure handling.

Suggested refactor

-        create_tasks = [
-            asyncio.create_task(
-                eng.create_task(make_create_data(), requested_by="alice")
-            )
-            for _ in range(5)
-        ]
-
-        # Yield to the event loop so the tasks can enqueue their mutations before stop()
-        await asyncio.sleep(0)
-
-        # Stop while tasks may still be in flight — drain timeout is generous
-        await eng.stop(timeout=5.0)
-
-        # All futures resolved (drained during stop or completed before stop)
-        results = await asyncio.gather(*create_tasks)
+        async with asyncio.TaskGroup() as tg:
+            create_tasks = [
+                tg.create_task(
+                    eng.create_task(make_create_data(), requested_by="alice")
+                )
+                for _ in range(5)
+            ]
+
+            # Yield to the event loop so the tasks can enqueue their mutations before stop()
+            await asyncio.sleep(0)
+
+            # Stop while tasks may still be in flight — drain timeout is generous
+            await eng.stop(timeout=5.0)
+
+        # All futures resolved (drained during stop or completed before stop)
+        results = [task.result() for task in create_tasks]

As per coding guidelines "Prefer asyncio.TaskGroup for fan-out/fan-in parallel operations in new code (multiple tool invocations, parallel agent calls). Prefer structured concurrency over bare create_task."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit/engine/test_task_engine_integration.py` around lines 99 - 115,
Replace the ad-hoc create_task/gather pattern with structured concurrency using
asyncio.TaskGroup: instead of building create_tasks with asyncio.create_task and
later awaiting asyncio.gather, open an asyncio.TaskGroup, spawn tasks calling
eng.create_task(make_create_data(), requested_by="alice") inside it, and rely on
the TaskGroup to propagate failures and cancel siblings; keep the yield to the
event loop (await asyncio.sleep(0)) if needed before calling await
eng.stop(timeout=5.0), and remove the final asyncio.gather by letting the
TaskGroup scope ensure all tasks are resolved before asserting the number of
completed tasks results from the group.

src/ai_company/api/controllers/tasks.py (1)

185-186: 🧹 Nitpick | 🔵 Trivial

Use the repository's Python 3.14 multi-except form consistently.

These handlers still use tuple-style except (...) as exc: blocks instead of the except A, B as exc: form required by the project rule.

As per coding guidelines "Use except A, B: syntax (no parentheses) for multiple exception types — PEP 758 on Python 3.14."

Also applies to: 211-212, 261-267, 305-312, 354-362, 395-402

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/ai_company/api/controllers/tasks.py` around lines 185 - 186, Replace
tuple-style multi-except blocks with the Python 3.14 multi-except form used
across the repo: change occurrences like the except (TaskInternalError,
TaskEngineNotRunningError) as exc: to the form except TaskInternalError,
TaskEngineNotRunningError as exc: so the same exception object is captured;
apply the same change to the other listed handlers in this file (the blocks
around lines referenced: the handlers at 211-212, 261-267, 305-312, 354-362,
395-402) and ensure each uses the unparenthesized `except A, B as exc:` syntax
and still re-raises via raise _map_task_engine_errors(exc) from exc.

tests/unit/api/controllers/test_task_helpers.py (1)

72-88: ⚠️ Potential issue | 🟡 Minor

Assert the exact sanitized 503 strings.

These cases only prove the raw engine text is hidden. They'll still pass if the mapping regresses to some other 503 message, which weakens the API contract the controller is deliberately sanitizing.

Suggested assertions

     def test_not_running_maps_to_service_unavailable(self) -> None:
         exc = TaskEngineNotRunningError("not running")
         result = _map_task_engine_errors(exc)
         assert isinstance(result, ServiceUnavailableError)
-        assert "not running" not in str(result)
+        assert str(result) == "Service temporarily unavailable"

     def test_queue_full_maps_to_service_unavailable(self) -> None:
         exc = TaskEngineQueueFullError("queue full")
         result = _map_task_engine_errors(exc)
         assert isinstance(result, ServiceUnavailableError)
-        assert "queue full" not in str(result)
+        assert str(result) == "Service temporarily unavailable"

     def test_internal_error_maps_to_service_unavailable(self) -> None:
         exc = TaskInternalError("internal fault")
         result = _map_task_engine_errors(exc)
         assert isinstance(result, ServiceUnavailableError)
-        assert "internal fault" not in str(result)
+        assert str(result) == "Internal server error"
+        assert "internal fault" not in str(result)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit/api/controllers/test_task_helpers.py` around lines 72 - 88, The
tests only check that the original engine text is removed; change each test to
assert the exact sanitized 503 text produced by the mapping function so the API
contract is enforced. Locate the sanitized message constant (the string used
inside _map_task_engine_errors, e.g. SANITIZED_SERVICE_UNAVAILABLE_MESSAGE or
similar) and replace the current negative assertions with equality checks like
assert str(result) == SANITIZED_SERVICE_UNAVAILABLE_MESSAGE for
TaskEngineNotRunningError, TaskEngineQueueFullError, and TaskInternalError to
ensure the mapping returns the exact expected ServiceUnavailableError message.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes