fix(agentex): address audit — atomic current_state write, clear-to-null, bounds

Resolve the review findings on the current_state observability change:

- Blocker: route current_state (and task_metadata) through a new column-scoped
  atomic UPDATE (TaskRepository/AgentTaskService.update_mutable_fields) instead
  of update_task's whole-row session.merge of a stale entity. Touching only the
  supplied columns means a concurrent status transition or params merge can no
  longer be silently reverted (terminal task revert / deleted task resurrection
  / lost param edit). update_task keeps its full-row semantics for the
  status-writing callers (delete/fail/forward) that need them.
- Support clearing current_state: an UNSET sentinel distinguishes an explicit
  null (clears the label) from an omitted field (left untouched), driven off the
  request's model_fields_set.
- Bound the label: String(255) column + max_length on the request/response
  schema (new field, no back-compat risk) so it can't amplify unboundedly onto
  every task_updated SSE payload. Single source of truth: CURRENT_STATE_MAX_LENGTH.
- Make the migration idempotent (ADD/DROP COLUMN IF [NOT] EXISTS), width 255.
- Align the domain-entity field title with the API schema title.
- Fix stale migration-linter path/flag in CLAUDE.md
  (scripts/ci_tools/migration_lint.py --base).
- Tests: single-atomic-write (spy), status no-clobber regression, clear-to-null,
  by-name route, empty-string, over-length 422, combined update; robust bounded
  wait in the new stream test. Regenerated openapi.yaml.

Author: qwhex

CLAUDE.md

@@ -358,7 +358,7 @@ Use the escape hatch for "this needs a maintenance window with traffic shifted a
 
 ##### Anti-pattern → linter rule reference
 
-The migration linter at `agentex/scripts/lint_migrations.py` enforces these rules at PR time via `.github/workflows/migration-lint.yml`. It only checks files changed vs the PR base, so existing migrations are not retro-flagged. The mapping below is what the linter catches:
+The migration linter at `agentex/scripts/ci_tools/migration_lint.py` enforces these rules at PR time via `.github/workflows/migration-lint.yml`. It only checks files changed vs the PR base, so existing migrations are not retro-flagged. The mapping below is what the linter catches:
 
 | Anti-pattern | Linter rule |
 |---|---|
@@ -371,7 +371,7 @@ The migration linter at `agentex/scripts/lint_migrations.py` enforces these rule
 Run the linter locally before pushing:
 
 ```bash
-agentex/scripts/lint_migrations.py --base-ref origin/main
+agentex/scripts/ci_tools/migration_lint.py --base origin/main
 ```
 
 ##### Other rules

agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py

@@ -21,8 +21,12 @@
 def upgrade() -> None:
     # Opaque label mirroring an agent's StateMachine current state. Nullable and
     # additive; agents opt in by emitting it. Metadata-only add, non-blocking.
-    op.add_column('tasks', sa.Column('current_state', sa.String(), nullable=True))
+    # Idempotent (IF NOT EXISTS) so re-running on an environment that already has
+    # the column is a no-op. Width matches TaskORM.current_state (String(255)).
+    op.execute(
+        "ALTER TABLE tasks ADD COLUMN IF NOT EXISTS current_state VARCHAR(255)"
+    )
 
 
 def downgrade() -> None:
-    op.drop_column('tasks', 'current_state')
+    op.execute("ALTER TABLE tasks DROP COLUMN IF EXISTS current_state")

agentex/openapi.yaml

@@ -6591,6 +6591,7 @@ components:
         current_state:
           anyOf:
           - type: string
+            maxLength: 255
           - type: 'null'
           title: Opaque label mirroring the agent's StateMachine current state; null
             when the agent does not emit one. Orthogonal to 'status'.
@@ -6820,6 +6821,7 @@ components:
         current_state:
           anyOf:
           - type: string
+            maxLength: 255
           - type: 'null'
           title: Opaque label mirroring the agent's StateMachine current state; null
             when the agent does not emit one. Orthogonal to 'status'.
@@ -7411,8 +7413,10 @@ components:
         current_state:
           anyOf:
           - type: string
+            maxLength: 255
           - type: 'null'
-          title: If provided, replaces the task's current_state label.
+          title: If provided, replaces the task's current_state label; pass null to
+            clear it, omit to leave it unchanged.
       type: object
       title: UpdateTaskRequest
     ValidationError:

agentex/src/adapters/orm.py

@@ -77,8 +77,10 @@ class TaskORM(BaseORM):
     task_metadata = Column(JSONB, nullable=True)
     # Opaque, framework-agnostic label mirroring an agent's StateMachine current
     # state. Written best-effort by the agent on each transition; reconciled on
-    # read. Orthogonal to `status` (Temporal workflow lifecycle).
-    current_state = Column(String, nullable=True)
+    # read. Orthogonal to `status` (Temporal workflow lifecycle). Bounded: a
+    # state label is short, and the value is echoed onto every task_updated SSE
+    # payload, so we cap it to avoid unbounded stream amplification.
+    current_state = Column(String(255), nullable=True)
     # Many-to-Many relationship with agents
     agents = relationship("AgentORM", secondary="task_agents", back_populates="tasks")
 

agentex/src/api/routes/tasks.py

@@ -22,7 +22,7 @@
 from src.domain.entities.tasks import TaskStatus as DomainTaskStatus
 from src.domain.services.authorization_service import DAuthorizationService
 from src.domain.use_cases.streams_use_case import DStreamsUseCase
-from src.domain.use_cases.tasks_use_case import DTaskUseCase
+from src.domain.use_cases.tasks_use_case import UNSET, DTaskUseCase
 from src.utils.authorization_shortcuts import (
     DAuthorizedId,
     DAuthorizedName,
@@ -196,7 +196,13 @@ async def update_task(
         id=task_id,
         task_metadata=request.task_metadata,
         merge_params=request.merge_params,
-        current_state=request.current_state,
+        # Pass through only when the client actually sent the key, so an explicit
+        # null clears the label while an omitted field leaves it untouched.
+        current_state=(
+            request.current_state
+            if "current_state" in request.model_fields_set
+            else UNSET
+        ),
     )
     return Task.model_validate(updated_task_entity)
 
@@ -218,7 +224,13 @@ async def update_task_by_name(
         name=task_name,
         task_metadata=request.task_metadata,
         merge_params=request.merge_params,
-        current_state=request.current_state,
+        # Pass through only when the client actually sent the key, so an explicit
+        # null clears the label while an omitted field leaves it untouched.
+        current_state=(
+            request.current_state
+            if "current_state" in request.model_fields_set
+            else UNSET
+        ),
     )
     return Task.model_validate(updated_task_entity)
 

agentex/src/api/schemas/tasks.py

@@ -7,6 +7,11 @@
 from src.api.schemas.agents import Agent
 from src.utils.model_utils import BaseModel
 
+# Upper bound on the opaque `current_state` label. Kept in sync with the
+# `TaskORM.current_state` column width (String(255)); a state label is short and
+# the value rides every task_updated SSE payload, so it is capped.
+CURRENT_STATE_MAX_LENGTH = 255
+
 
 class TaskRelationships(str, Enum):
     """Task relationships that can be loaded"""
@@ -65,6 +70,7 @@ class Task(BaseModel):
     )
     current_state: str | None = Field(
         None,
+        max_length=CURRENT_STATE_MAX_LENGTH,
         title=(
             "Opaque label mirroring the agent's StateMachine current state; "
             "null when the agent does not emit one. Orthogonal to 'status'."
@@ -96,7 +102,11 @@ class UpdateTaskRequest(BaseModel):
     )
     current_state: str | None = Field(
         None,
-        title="If provided, replaces the task's current_state label.",
+        max_length=CURRENT_STATE_MAX_LENGTH,
+        title=(
+            "If provided, replaces the task's current_state label; "
+            "pass null to clear it, omit to leave it unchanged."
+        ),
     )
 
 

agentex/src/domain/entities/tasks.py

@@ -68,7 +68,10 @@ class TaskEntity(BaseModel):
     )
     current_state: str | None = Field(
         None,
-        title="Opaque label mirroring the agent's StateMachine current state",
+        title=(
+            "Opaque label mirroring the agent's StateMachine current state; "
+            "null when the agent does not emit one. Orthogonal to 'status'."
+        ),
     )
 
     # allow extra fields for agents relationships

agentex/src/domain/repositories/task_repository.py

@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from datetime import UTC, datetime, timedelta
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
 from fastapi import Depends
 from sqlalchemy import cast, distinct, func, select, update
@@ -254,6 +254,42 @@ async def merge_params(self, task_id: str, patch: dict) -> TaskEntity | None:
                 return None
             return TaskEntity.model_validate(row)
 
+    async def update_mutable_fields(
+        self, task_id: str, fields: dict[str, Any]
+    ) -> TaskEntity | None:
+        """Atomically set only the given scalar columns on a single task row.
+
+        Unlike ``update`` (a whole-row ``session.merge`` of a possibly-stale
+        entity, which rewrites every column and can clobber a concurrently
+        changed ``status``/``params``), this issues a single
+        ``UPDATE ... SET <only these columns> WHERE id = :id RETURNING *``.
+        Touching only the supplied columns means a concurrent status transition
+        or param merge is never reverted. Returns the updated entity, or
+        ``None`` if no task with ``task_id`` exists.
+
+        ``fields`` values are applied verbatim, so passing ``current_state=None``
+        clears the column (callers distinguish "clear" from "omit" upstream).
+        """
+        if not fields:
+            return await self.get(id=task_id)
+
+        async with (
+            self.start_async_db_session(True) as session,
+            async_sql_exception_handler(),
+        ):
+            stmt = (
+                update(TaskORM)
+                .where(TaskORM.id == task_id)
+                .values(**fields)
+                .returning(TaskORM)
+            )
+            result = await session.execute(stmt)
+            row = result.scalar_one_or_none()
+            await session.commit()
+            if row is None:
+                return None
+            return TaskEntity.model_validate(row)
+
     async def transition_status(
         self,
         task_id: str,

agentex/src/domain/services/task_service.py

@@ -242,6 +242,40 @@ async def update_task(self, task: TaskEntity) -> TaskEntity:
 
         return updated_task
 
+    async def update_mutable_fields(
+        self, task_id: str, fields: dict[str, Any]
+    ) -> TaskEntity | None:
+        """Atomically update only the given scalar columns on a task, then
+        publish a task_updated event.
+
+        Column-scoped (see ``TaskRepository.update_mutable_fields``) so it cannot
+        clobber a concurrently changed ``status``/``params`` — unlike
+        ``update_task``'s whole-row merge, which the status-writing callers
+        (delete/fail/forward) still rely on. Returns the updated entity, or
+        ``None`` if the task no longer exists.
+        """
+        updated_task = await self.task_repository.update_mutable_fields(
+            task_id, fields
+        )
+        if updated_task is None:
+            return None
+
+        try:
+            topic = get_task_event_stream_topic(task_id=task_id)
+            await self.stream_repository.send_data(
+                topic,
+                TaskStreamTaskUpdatedEventEntity(
+                    type="task_updated", task=updated_task
+                ).model_dump(mode="json"),
+            )
+            logger.info(f"task_updated event published to topic: {topic}")
+        except Exception as e:
+            logger.error(
+                f"Error sending task_updated event to stream: {e}", exc_info=True
+            )
+
+        return updated_task
+
     async def merge_task_params(self, task_id: str, patch: dict) -> TaskEntity | None:
         """Atomically shallow-merge ``patch`` into ``tasks.params``. Returns
         the updated entity, or ``None`` if no task with ``task_id`` exists.

agentex/src/domain/use_cases/tasks_use_case.py

@@ -11,6 +11,15 @@
 logger = make_logger(__name__)
 
 
+class _Unset:
+    """Sentinel distinguishing "field omitted" from "field explicitly set to null"
+    in a PATCH-style update, so an explicit null can clear a column while an
+    omitted field is left untouched."""
+
+
+UNSET: Any = _Unset()
+
+
 class TasksUseCase:
     """
     Use case for managing tasks. Handles CRUD operations and delegates task operations to ACP servers.
@@ -95,13 +104,20 @@ async def update_mutable_fields_on_task(
         name: str | None = None,
         task_metadata: dict[str, Any] | None = None,
         merge_params: dict[str, Any] | None = None,
-        current_state: str | None = None,
+        current_state: str | None | _Unset = UNSET,
     ) -> TaskEntity:
-        """Update mutable fields on a task entity. This is used by our API since not all fields should be mutable."""
+        """Update mutable fields on a task entity. This is used by our API since not all fields should be mutable.
+
+        ``current_state`` uses the UNSET sentinel so an explicit null clears the
+        column while an omitted field leaves it untouched. ``task_metadata`` keeps
+        its legacy semantics (``None`` means "not supplied").
+        """
 
         if not id and not name:
             raise ClientError("Either id or name must be provided")
 
+        current_state_supplied = not isinstance(current_state, _Unset)
+
         # todo: make this a transaction?
         task_entity = await self.task_service.get_task(id=id, name=name)
         if task_entity.status == TaskStatus.DELETED:
@@ -111,32 +127,39 @@ async def update_mutable_fields_on_task(
                 raise ItemDoesNotExist(f"Task {name} not found")
 
         # No-op if no mutable field was supplied.
-        if task_metadata is None and merge_params is None and current_state is None:
+        if (
+            task_metadata is None
+            and merge_params is None
+            and not current_state_supplied
+        ):
             return task_entity
 
         # `merge_params` is a separate atomic JSONB shallow-merge so concurrent
-        # callers don't overwrite each other's fields (vs reading→mutating→writing
-        # the whole params dict on task_entity). Run it first so the refreshed
-        # entity it returns becomes the base we apply `task_metadata` on top of;
-        # otherwise the `task_entity = merged` reassignment would discard an
-        # in-memory metadata change made before the merge.
+        # callers don't overwrite each other's fields. Run it first so the
+        # refreshed entity it returns becomes the value we return if no scalar
+        # field updates follow.
         if merge_params:
             merged = await self.task_service.merge_task_params(
                 task_entity.id, merge_params
             )
             if merged is not None:
                 task_entity = merged
 
-        # Apply the whole-row field updates (task_metadata, current_state) in a
-        # single update_task write so they emit one task_updated event rather than
-        # two. current_state rides the same publish that already powers reactive
-        # task_updated delivery to stream consumers.
-        if task_metadata is not None or current_state is not None:
-            if task_metadata is not None:
-                task_entity.task_metadata = task_metadata
-            if current_state is not None:
-                task_entity.current_state = current_state
-            task_entity = await self.task_service.update_task(task=task_entity)
+        # Apply task_metadata/current_state via a single column-scoped atomic
+        # update — one task_updated publish, and (unlike a whole-row merge) no
+        # risk of clobbering a concurrently changed status/params. current_state
+        # rides that publish, powering reactive delivery to stream consumers.
+        fields: dict[str, Any] = {}
+        if task_metadata is not None:
+            fields["task_metadata"] = task_metadata
+        if current_state_supplied:
+            fields["current_state"] = current_state
+        if fields:
+            updated = await self.task_service.update_mutable_fields(
+                task_entity.id, fields
+            )
+            if updated is not None:
+                task_entity = updated
 
         return task_entity