From f6e03e33a65454ca71cd3dd4c21ea3a9c96e5faf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mice=20P=C3=A1pai?= Date: Wed, 22 Jul 2026 17:34:10 +0200 Subject: [PATCH 1/5] feat(agentex): materialized Task.current_state for reactive state observability Expose an agent's StateMachine current state to REST/SSE consumers via a nullable, opaque `Task.current_state` label. It is written through the existing UPDATE_TASK / PUT /tasks/{id} path, which already publishes a `task_updated` SSE event carrying the full task, so updates are reactive (push) while GET /tasks/{id} remains the authoritative point-read for load/reconnect reconciliation. The field is framework-agnostic and default-null: agents opt in by writing it; existing tasks and stream consumers are unaffected (additive, backward compatible). - ORM column + Alembic migration (nullable add, down_revision a1b2c3d4e5f6) - Task response schema, UpdateTaskRequest, TaskEntity + converter - update_mutable_fields_on_task applies it in the same update_task write as task_metadata (single task_updated publish) - both PUT routes thread it through - regenerated openapi.yaml - tests: route PUT/null/no-clobber, task_updated carries current_state, service-layer emit+persist, extra="ignore" forward-compat pin --- ...200_add_task_current_state_b2c3d4e5f6a7.py | 28 +++++++ agentex/openapi.yaml | 17 ++++ agentex/src/adapters/orm.py | 4 + agentex/src/api/routes/tasks.py | 2 + agentex/src/api/schemas/tasks.py | 11 +++ agentex/src/domain/entities/tasks.py | 5 ++ .../src/domain/use_cases/tasks_use_case.py | 16 +++- .../integration/api/tasks/test_tasks_api.py | 79 +++++++++++++++++++ agentex/tests/integration/test_task_stream.py | 58 ++++++++++++++ .../tests/unit/services/test_task_service.py | 26 ++++++ 10 files changed, 242 insertions(+), 4 deletions(-) create mode 100644 agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py diff --git a/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py b/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py new file mode 100644 index 00000000..b332ff44 --- /dev/null +++ b/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py @@ -0,0 +1,28 @@ +"""add task current_state + +Revision ID: b2c3d4e5f6a7 +Revises: a1b2c3d4e5f6 +Create Date: 2026-07-22 12:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'b2c3d4e5f6a7' +down_revision: Union[str, None] = 'a1b2c3d4e5f6' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +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)) + + +def downgrade() -> None: + op.drop_column('tasks', 'current_state') diff --git a/agentex/openapi.yaml b/agentex/openapi.yaml index 904f1615..1a3339c7 100644 --- a/agentex/openapi.yaml +++ b/agentex/openapi.yaml @@ -6588,6 +6588,12 @@ components: type: object - type: 'null' title: Task metadata + current_state: + anyOf: + - type: string + - type: 'null' + title: Opaque label mirroring the agent's StateMachine current state; null + when the agent does not emit one. Orthogonal to 'status'. type: object required: - id @@ -6811,6 +6817,12 @@ components: type: object - type: 'null' title: Task metadata + current_state: + anyOf: + - type: string + - type: 'null' + title: Opaque label mirroring the agent's StateMachine current state; null + when the agent does not emit one. Orthogonal to 'status'. agents: anyOf: - items: @@ -7396,6 +7408,11 @@ components: - type: 'null' title: Optional shallow-merge patch applied to the task's params column. Top-level keys overwrite; pass full nested objects to change subfields. + current_state: + anyOf: + - type: string + - type: 'null' + title: If provided, replaces the task's current_state label. type: object title: UpdateTaskRequest ValidationError: diff --git a/agentex/src/adapters/orm.py b/agentex/src/adapters/orm.py index e5f7b139..2671a1fa 100644 --- a/agentex/src/adapters/orm.py +++ b/agentex/src/adapters/orm.py @@ -75,6 +75,10 @@ class TaskORM(BaseORM): cleaned_at = Column(DateTime(timezone=True), nullable=True) params = Column(JSONB, nullable=True) 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) # Many-to-Many relationship with agents agents = relationship("AgentORM", secondary="task_agents", back_populates="tasks") diff --git a/agentex/src/api/routes/tasks.py b/agentex/src/api/routes/tasks.py index 86abb22a..4b491d71 100644 --- a/agentex/src/api/routes/tasks.py +++ b/agentex/src/api/routes/tasks.py @@ -196,6 +196,7 @@ async def update_task( id=task_id, task_metadata=request.task_metadata, merge_params=request.merge_params, + current_state=request.current_state, ) return Task.model_validate(updated_task_entity) @@ -217,6 +218,7 @@ async def update_task_by_name( name=task_name, task_metadata=request.task_metadata, merge_params=request.merge_params, + current_state=request.current_state, ) return Task.model_validate(updated_task_entity) diff --git a/agentex/src/api/schemas/tasks.py b/agentex/src/api/schemas/tasks.py index a79d9fef..6db6786a 100644 --- a/agentex/src/api/schemas/tasks.py +++ b/agentex/src/api/schemas/tasks.py @@ -63,6 +63,13 @@ class Task(BaseModel): None, title="Task metadata", ) + current_state: str | None = Field( + None, + title=( + "Opaque label mirroring the agent's StateMachine current state; " + "null when the agent does not emit one. Orthogonal to 'status'." + ), + ) class TaskResponse(Task): @@ -87,6 +94,10 @@ class UpdateTaskRequest(BaseModel): "subfields." ), ) + current_state: str | None = Field( + None, + title="If provided, replaces the task's current_state label.", + ) class TaskStatusReasonRequest(BaseModel): diff --git a/agentex/src/domain/entities/tasks.py b/agentex/src/domain/entities/tasks.py index 76e93cf9..be5bdc86 100644 --- a/agentex/src/domain/entities/tasks.py +++ b/agentex/src/domain/entities/tasks.py @@ -66,6 +66,10 @@ class TaskEntity(BaseModel): None, title="Task metadata", ) + current_state: str | None = Field( + None, + title="Opaque label mirroring the agent's StateMachine current state", + ) # allow extra fields for agents relationships model_config = ConfigDict(extra="allow") @@ -84,4 +88,5 @@ def convert_task_to_entity(task: Task) -> TaskEntity: cleaned_at=task.cleaned_at, params=task.params, task_metadata=task.task_metadata, + current_state=task.current_state, ) diff --git a/agentex/src/domain/use_cases/tasks_use_case.py b/agentex/src/domain/use_cases/tasks_use_case.py index 7c114744..6f2257dc 100644 --- a/agentex/src/domain/use_cases/tasks_use_case.py +++ b/agentex/src/domain/use_cases/tasks_use_case.py @@ -95,6 +95,7 @@ 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, ) -> TaskEntity: """Update mutable fields on a task entity. This is used by our API since not all fields should be mutable.""" @@ -109,8 +110,8 @@ async def update_mutable_fields_on_task( else: raise ItemDoesNotExist(f"Task {name} not found") - # No-op if neither field was supplied. - if task_metadata is None and merge_params is None: + # No-op if no mutable field was supplied. + if task_metadata is None and merge_params is None and current_state is None: return task_entity # `merge_params` is a separate atomic JSONB shallow-merge so concurrent @@ -126,8 +127,15 @@ async def update_mutable_fields_on_task( if merged is not None: task_entity = merged - if task_metadata is not None: - task_entity.task_metadata = task_metadata + # 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) return task_entity diff --git a/agentex/tests/integration/api/tasks/test_tasks_api.py b/agentex/tests/integration/api/tasks/test_tasks_api.py index bdf857bd..eca1775b 100644 --- a/agentex/tests/integration/api/tasks/test_tasks_api.py +++ b/agentex/tests/integration/api/tasks/test_tasks_api.py @@ -683,6 +683,85 @@ async def test_update_task_endpoint_success( assert response_data["task_metadata"]["configuration"]["version"] == "2.0.0" assert response_data["task_metadata"]["metrics"]["complexity_score"] == 75 + async def test_update_task_current_state( + self, isolated_client, isolated_repositories + ): + """PUT /tasks/{id} writes current_state; null/omitted leaves it unset.""" + agent_repo = isolated_repositories["agent_repository"] + agent = AgentEntity( + id=orm_id(), + name="current-state-agent", + description="Agent for current_state update testing", + acp_url="http://test-acp:8000", + acp_type=ACPType.SYNC, + ) + await agent_repo.create(agent) + + task_repo = isolated_repositories["task_repository"] + task = TaskEntity( + id=orm_id(), + name="task-for-current-state", + status=TaskStatus.RUNNING, + status_reason="Test task for current_state", + ) + created_task = await task_repo.create(agent_id=agent.id, task=task) + + # Fresh task: current_state present in response and null by default. + response = await isolated_client.get(f"/tasks/{created_task.id}") + assert response.status_code == 200 + assert response.json()["current_state"] is None + + # Setting current_state persists and echoes back. + response = await isolated_client.put( + f"/tasks/{created_task.id}", json={"current_state": "awaiting_input"} + ) + assert response.status_code == 200 + assert response.json()["current_state"] == "awaiting_input" + + # Point-read reflects the committed value (source of truth). + response = await isolated_client.get(f"/tasks/{created_task.id}") + assert response.status_code == 200 + assert response.json()["current_state"] == "awaiting_input" + + # Updating only task_metadata leaves current_state untouched (no clobber). + response = await isolated_client.put( + f"/tasks/{created_task.id}", json={"task_metadata": {"k": "v"}} + ) + assert response.status_code == 200 + assert response.json()["current_state"] == "awaiting_input" + + async def test_update_task_request_ignores_unknown_fields( + self, isolated_client, isolated_repositories + ): + """Guards the extra="ignore" assumption a newer SDK relies on: a payload + carrying fields the server doesn't model must 200 (drop them), not 422. + Breaks loudly if UpdateTaskRequest ever switches to extra="forbid".""" + agent_repo = isolated_repositories["agent_repository"] + agent = AgentEntity( + id=orm_id(), + name="unknown-fields-agent", + description="Agent for unknown-field compat testing", + acp_url="http://test-acp:8000", + acp_type=ACPType.SYNC, + ) + await agent_repo.create(agent) + + task_repo = isolated_repositories["task_repository"] + task = TaskEntity( + id=orm_id(), + name="task-for-unknown-fields", + status=TaskStatus.RUNNING, + status_reason="Test task for unknown-field compat", + ) + created_task = await task_repo.create(agent_id=agent.id, task=task) + + response = await isolated_client.put( + f"/tasks/{created_task.id}", + json={"current_state": "working", "field_from_a_newer_sdk": "ignored"}, + ) + assert response.status_code == 200 + assert response.json()["current_state"] == "working" + async def test_update_task_endpoint_validation( self, isolated_client, isolated_repositories ): diff --git a/agentex/tests/integration/test_task_stream.py b/agentex/tests/integration/test_task_stream.py index 7da23edd..a1edf5b0 100644 --- a/agentex/tests/integration/test_task_stream.py +++ b/agentex/tests/integration/test_task_stream.py @@ -230,6 +230,64 @@ async def collect_stream_events(): print("✅ Task metadata update successfully triggered stream event") + async def test_current_state_update_triggers_stream_event( + self, test_agent_and_task, tasks_use_case, streams_use_case + ): + """current_state rides the existing task_updated event, so a client + subscribed to the task stream is pushed the new state reactively.""" + _agent, task = test_agent_and_task + + stream_events = [] + + async def collect_stream_events(): + try: + async for event_data in streams_use_case.stream_task_events( + task_id=task.id + ): + if event_data.startswith("data: "): + import json + + event_json = event_data[6:].strip() + if event_json: + try: + event = json.loads(event_json) + stream_events.append(event) + if event.get("type") == "task_updated": + break + except json.JSONDecodeError: + pass + except asyncio.CancelledError: + pass + + stream_task = asyncio.create_task(collect_stream_events()) + await asyncio.sleep(0.1) + + updated_task = await tasks_use_case.update_mutable_fields_on_task( + id=task.id, current_state="awaiting_input" + ) + + await asyncio.sleep(0.5) + + stream_task.cancel() + try: + await stream_task + except asyncio.CancelledError: + pass + + task_updated_events = [ + e for e in stream_events if e.get("type") == "task_updated" + ] + assert len(task_updated_events) >= 1, ( + f"Expected task_updated event, got events: {[e.get('type') for e in stream_events]}" + ) + event_task = task_updated_events[0]["task"] + assert event_task["id"] == task.id + assert event_task["current_state"] == "awaiting_input" + + assert updated_task.current_state == "awaiting_input" + + print("✅ current_state update successfully triggered stream event") + async def test_get_task_returns_updated_metadata_after_stream_update( self, test_agent_and_task, tasks_use_case ): diff --git a/agentex/tests/unit/services/test_task_service.py b/agentex/tests/unit/services/test_task_service.py index 91caa8d6..74f880fc 100644 --- a/agentex/tests/unit/services/test_task_service.py +++ b/agentex/tests/unit/services/test_task_service.py @@ -938,6 +938,32 @@ async def test_update_task_with_task_metadata_changes( assert event_data["type"] == "task_updated" assert event_data["task"]["task_metadata"] == updated_metadata + async def test_update_task_current_state_publishes_stream_event( + self, task_service, agent_repository, sample_agent, redis_stream_repository + ): + """update_task persists current_state and carries it on the published + task_updated event, so subscribed clients see the new state.""" + await create_or_get_agent(agent_repository, sample_agent) + created_task = await task_service.create_task( + agent=sample_agent, task_name="task-for-current-state" + ) + + created_task.current_state = "working" + redis_stream_repository.send_data = AsyncMock() + + result = await task_service.update_task(created_task) + + assert result.current_state == "working" + retrieved_task = await task_service.get_task(id=created_task.id) + assert retrieved_task.current_state == "working" + + redis_stream_repository.send_data.assert_called_once() + call_args = redis_stream_repository.send_data.call_args + assert call_args[0][0] == f"task:{created_task.id}" + event_data = call_args[0][1] + assert event_data["type"] == "task_updated" + assert event_data["task"]["current_state"] == "working" + async def test_get_task_preserves_task_metadata( self, task_service, agent_repository, sample_agent ): From 8bcc19f81e5e97e1a8e8ec1c4f896f2ddb6f4296 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mice=20P=C3=A1pai?= Date: Wed, 22 Jul 2026 19:37:00 +0200 Subject: [PATCH 2/5] =?UTF-8?q?fix(agentex):=20address=20audit=20=E2=80=94?= =?UTF-8?q?=20atomic=20current=5Fstate=20write,=20clear-to-null,=20bounds?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- CLAUDE.md | 4 +- ...200_add_task_current_state_b2c3d4e5f6a7.py | 8 +- agentex/openapi.yaml | 6 +- agentex/src/adapters/orm.py | 6 +- agentex/src/api/routes/tasks.py | 18 ++- agentex/src/api/schemas/tasks.py | 12 +- agentex/src/domain/entities/tasks.py | 5 +- .../domain/repositories/task_repository.py | 38 ++++- agentex/src/domain/services/task_service.py | 34 +++++ .../src/domain/use_cases/tasks_use_case.py | 59 +++++--- .../integration/api/tasks/test_tasks_api.py | 136 +++++++++++++++++- agentex/tests/integration/test_task_stream.py | 14 +- .../tests/unit/services/test_task_service.py | 56 ++++++++ .../unit/use_cases/test_tasks_use_case.py | 96 +++++++++++++ 14 files changed, 453 insertions(+), 39 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 99a62276..f8530835 100644 --- a/CLAUDE.md +++ b/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 diff --git a/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py b/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py index b332ff44..38986c48 100644 --- a/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py +++ b/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") diff --git a/agentex/openapi.yaml b/agentex/openapi.yaml index 1a3339c7..f448d8e0 100644 --- a/agentex/openapi.yaml +++ b/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: diff --git a/agentex/src/adapters/orm.py b/agentex/src/adapters/orm.py index 2671a1fa..45c09fa4 100644 --- a/agentex/src/adapters/orm.py +++ b/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") diff --git a/agentex/src/api/routes/tasks.py b/agentex/src/api/routes/tasks.py index 4b491d71..f1f297e7 100644 --- a/agentex/src/api/routes/tasks.py +++ b/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) diff --git a/agentex/src/api/schemas/tasks.py b/agentex/src/api/schemas/tasks.py index 6db6786a..260f2c86 100644 --- a/agentex/src/api/schemas/tasks.py +++ b/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." + ), ) diff --git a/agentex/src/domain/entities/tasks.py b/agentex/src/domain/entities/tasks.py index be5bdc86..82cc3776 100644 --- a/agentex/src/domain/entities/tasks.py +++ b/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 diff --git a/agentex/src/domain/repositories/task_repository.py b/agentex/src/domain/repositories/task_repository.py index 45662015..795770d6 100644 --- a/agentex/src/domain/repositories/task_repository.py +++ b/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 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, diff --git a/agentex/src/domain/services/task_service.py b/agentex/src/domain/services/task_service.py index d5d1291f..90cb76a5 100644 --- a/agentex/src/domain/services/task_service.py +++ b/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. diff --git a/agentex/src/domain/use_cases/tasks_use_case.py b/agentex/src/domain/use_cases/tasks_use_case.py index 6f2257dc..b1a87cdc 100644 --- a/agentex/src/domain/use_cases/tasks_use_case.py +++ b/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,15 +127,17 @@ 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 @@ -127,16 +145,21 @@ async def update_mutable_fields_on_task( 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 diff --git a/agentex/tests/integration/api/tasks/test_tasks_api.py b/agentex/tests/integration/api/tasks/test_tasks_api.py index eca1775b..492d87d7 100644 --- a/agentex/tests/integration/api/tasks/test_tasks_api.py +++ b/agentex/tests/integration/api/tasks/test_tasks_api.py @@ -686,7 +686,8 @@ async def test_update_task_endpoint_success( async def test_update_task_current_state( self, isolated_client, isolated_repositories ): - """PUT /tasks/{id} writes current_state; null/omitted leaves it unset.""" + """PUT /tasks/{id} writes current_state; explicit null clears it while an + omitted field leaves it untouched; point-read is source of truth.""" agent_repo = isolated_repositories["agent_repository"] agent = AgentEntity( id=orm_id(), @@ -723,13 +724,144 @@ async def test_update_task_current_state( assert response.status_code == 200 assert response.json()["current_state"] == "awaiting_input" - # Updating only task_metadata leaves current_state untouched (no clobber). + # Updating only task_metadata (current_state omitted) does not clobber it. response = await isolated_client.put( f"/tasks/{created_task.id}", json={"task_metadata": {"k": "v"}} ) assert response.status_code == 200 assert response.json()["current_state"] == "awaiting_input" + # Explicit null clears the label (distinct from omitting the field). + response = await isolated_client.put( + f"/tasks/{created_task.id}", json={"current_state": None} + ) + assert response.status_code == 200 + assert response.json()["current_state"] is None + response = await isolated_client.get(f"/tasks/{created_task.id}") + assert response.json()["current_state"] is None + + async def test_update_task_current_state_and_metadata_together( + self, isolated_client, isolated_repositories + ): + """current_state + task_metadata in one PUT both persist (single write), + without clobbering status.""" + agent_repo = isolated_repositories["agent_repository"] + agent = AgentEntity( + id=orm_id(), + name="current-state-combined-agent", + description="Agent for combined update testing", + acp_url="http://test-acp:8000", + acp_type=ACPType.SYNC, + ) + await agent_repo.create(agent) + + task_repo = isolated_repositories["task_repository"] + task = TaskEntity( + id=orm_id(), + name="task-for-combined-update", + status=TaskStatus.RUNNING, + status_reason="Test task for combined update", + ) + created_task = await task_repo.create(agent_id=agent.id, task=task) + + response = await isolated_client.put( + f"/tasks/{created_task.id}", + json={"current_state": "step_2", "task_metadata": {"stage": "two"}}, + ) + assert response.status_code == 200 + body = response.json() + assert body["current_state"] == "step_2" + assert body["task_metadata"] == {"stage": "two"} + assert body["status"] == "RUNNING" + + async def test_update_task_current_state_by_name( + self, isolated_client, isolated_repositories + ): + """PUT /tasks/name/{name} forwards current_state too.""" + agent_repo = isolated_repositories["agent_repository"] + agent = AgentEntity( + id=orm_id(), + name="current-state-by-name-agent", + description="Agent for by-name current_state testing", + acp_url="http://test-acp:8000", + acp_type=ACPType.SYNC, + ) + await agent_repo.create(agent) + + task_repo = isolated_repositories["task_repository"] + task = TaskEntity( + id=orm_id(), + name="task-for-current-state-by-name", + status=TaskStatus.RUNNING, + status_reason="Test task for by-name current_state", + ) + await task_repo.create(agent_id=agent.id, task=task) + + response = await isolated_client.put( + "/tasks/name/task-for-current-state-by-name", + json={"current_state": "working"}, + ) + assert response.status_code == 200 + assert response.json()["current_state"] == "working" + + async def test_update_task_current_state_empty_string( + self, isolated_client, isolated_repositories + ): + """Empty string is a valid label distinct from null (guards against a + future falsy check treating "" as "unset").""" + agent_repo = isolated_repositories["agent_repository"] + agent = AgentEntity( + id=orm_id(), + name="current-state-empty-agent", + description="Agent for empty-string current_state testing", + acp_url="http://test-acp:8000", + acp_type=ACPType.SYNC, + ) + await agent_repo.create(agent) + + task_repo = isolated_repositories["task_repository"] + task = TaskEntity( + id=orm_id(), + name="task-for-current-state-empty", + status=TaskStatus.RUNNING, + status_reason="Test task for empty current_state", + ) + created_task = await task_repo.create(agent_id=agent.id, task=task) + + response = await isolated_client.put( + f"/tasks/{created_task.id}", json={"current_state": ""} + ) + assert response.status_code == 200 + assert response.json()["current_state"] == "" + + async def test_update_task_current_state_too_long_rejected( + self, isolated_client, isolated_repositories + ): + """current_state exceeding the max length is rejected with 422.""" + agent_repo = isolated_repositories["agent_repository"] + agent = AgentEntity( + id=orm_id(), + name="current-state-toolong-agent", + description="Agent for max-length current_state testing", + acp_url="http://test-acp:8000", + acp_type=ACPType.SYNC, + ) + await agent_repo.create(agent) + + task_repo = isolated_repositories["task_repository"] + task = TaskEntity( + id=orm_id(), + name="task-for-current-state-toolong", + status=TaskStatus.RUNNING, + status_reason="Test task for over-long current_state", + ) + created_task = await task_repo.create(agent_id=agent.id, task=task) + + response = await isolated_client.put( + f"/tasks/{created_task.id}", json={"current_state": "x" * 256} + ) + assert response.status_code == 422 + async def test_update_task_request_ignores_unknown_fields( self, isolated_client, isolated_repositories ): diff --git a/agentex/tests/integration/test_task_stream.py b/agentex/tests/integration/test_task_stream.py index a1edf5b0..d888109c 100644 --- a/agentex/tests/integration/test_task_stream.py +++ b/agentex/tests/integration/test_task_stream.py @@ -260,19 +260,21 @@ async def collect_stream_events(): pass stream_task = asyncio.create_task(collect_stream_events()) + # Let the subscription establish before the update; the stream is + # tail-only, so an event published before we subscribe would be missed. await asyncio.sleep(0.1) updated_task = await tasks_use_case.update_mutable_fields_on_task( id=task.id, current_state="awaiting_input" ) - await asyncio.sleep(0.5) - - stream_task.cancel() + # Wait until the collector sees task_updated (it breaks on it) rather + # than guessing a fixed delay; the timeout is only a failure ceiling. try: - await stream_task - except asyncio.CancelledError: - pass + async with asyncio.timeout(5): + await stream_task + except (TimeoutError, asyncio.CancelledError): + stream_task.cancel() task_updated_events = [ e for e in stream_events if e.get("type") == "task_updated" diff --git a/agentex/tests/unit/services/test_task_service.py b/agentex/tests/unit/services/test_task_service.py index 74f880fc..df6ef9fd 100644 --- a/agentex/tests/unit/services/test_task_service.py +++ b/agentex/tests/unit/services/test_task_service.py @@ -964,6 +964,62 @@ async def test_update_task_current_state_publishes_stream_event( assert event_data["type"] == "task_updated" assert event_data["task"]["current_state"] == "working" + async def test_update_mutable_fields_persists_and_publishes( + self, task_service, agent_repository, sample_agent, redis_stream_repository + ): + """update_mutable_fields writes only the given columns and publishes a + task_updated event carrying them.""" + await create_or_get_agent(agent_repository, sample_agent) + created_task = await task_service.create_task( + agent=sample_agent, task_name="task-for-mutable-fields" + ) + + redis_stream_repository.send_data = AsyncMock() + result = await task_service.update_mutable_fields( + created_task.id, + {"current_state": "working", "task_metadata": {"a": 1}}, + ) + + assert result.current_state == "working" + assert result.task_metadata == {"a": 1} + retrieved = await task_service.get_task(id=created_task.id) + assert retrieved.current_state == "working" + assert retrieved.task_metadata == {"a": 1} + + redis_stream_repository.send_data.assert_called_once() + event_data = redis_stream_repository.send_data.call_args[0][1] + assert event_data["type"] == "task_updated" + assert event_data["task"]["current_state"] == "working" + + async def test_update_mutable_fields_does_not_clobber_status( + self, task_service, agent_repository, sample_agent, redis_stream_repository + ): + """Regression (blocker): a current_state write must not revert a status + changed by another writer. update_mutable_fields is column-scoped, so a + task that went terminal is not resurrected to its earlier status.""" + await create_or_get_agent(agent_repository, sample_agent) + created_task = await task_service.create_task( + agent=sample_agent, task_name="task-for-noclobber" + ) + + # Another writer moves the task to a terminal status. + await task_service.transition_task_status( + task_id=created_task.id, + expected_status=TaskStatus.RUNNING, + new_status=TaskStatus.COMPLETED, + status_reason="done", + ) + + redis_stream_repository.send_data = AsyncMock() + result = await task_service.update_mutable_fields( + created_task.id, {"current_state": "late"} + ) + + assert result.current_state == "late" + assert result.status == TaskStatus.COMPLETED + retrieved = await task_service.get_task(id=created_task.id) + assert retrieved.status == TaskStatus.COMPLETED + async def test_get_task_preserves_task_metadata( self, task_service, agent_repository, sample_agent ): diff --git a/agentex/tests/unit/use_cases/test_tasks_use_case.py b/agentex/tests/unit/use_cases/test_tasks_use_case.py index 798cb7cd..a2d6e737 100644 --- a/agentex/tests/unit/use_cases/test_tasks_use_case.py +++ b/agentex/tests/unit/use_cases/test_tasks_use_case.py @@ -3,6 +3,7 @@ methods (complete_task, fail_task, etc.) and metadata updates. """ +from unittest.mock import AsyncMock from uuid import uuid4 import pytest @@ -592,6 +593,101 @@ async def test_update_metadata_and_merge_params_both_persist( assert updated.task_metadata == {"stage": "tuned"} assert updated.params == {"model": "gpt-4", "temperature": 0.7} + async def test_update_current_state( + self, tasks_use_case, task_service, agent_repository, sample_agent + ): + """current_state persists and leaves status/task_metadata untouched.""" + await create_or_get_agent(agent_repository, sample_agent) + task = await task_service.create_task( + agent=sample_agent, + task_name="current-state-test", + task_metadata={"keep": "me"}, + ) + + updated = await tasks_use_case.update_mutable_fields_on_task( + id=task.id, current_state="working" + ) + + assert updated.current_state == "working" + assert updated.status == TaskStatus.RUNNING + assert updated.task_metadata == {"keep": "me"} + + async def test_update_current_state_noop_when_omitted( + self, tasks_use_case, task_service, agent_repository, sample_agent + ): + """Omitting current_state (the UNSET default) leaves it untouched.""" + await create_or_get_agent(agent_repository, sample_agent) + task = await task_service.create_task( + agent=sample_agent, task_name="current-state-omitted-test" + ) + await tasks_use_case.update_mutable_fields_on_task( + id=task.id, current_state="set-once" + ) + + # A later metadata-only update must not clear current_state. + updated = await tasks_use_case.update_mutable_fields_on_task( + id=task.id, task_metadata={"a": 1} + ) + + assert updated.current_state == "set-once" + + async def test_update_current_state_clears_on_explicit_null( + self, tasks_use_case, task_service, agent_repository, sample_agent + ): + """Passing current_state=None explicitly clears the label (vs omitting).""" + await create_or_get_agent(agent_repository, sample_agent) + task = await task_service.create_task( + agent=sample_agent, task_name="current-state-clear-test" + ) + await tasks_use_case.update_mutable_fields_on_task( + id=task.id, current_state="working" + ) + + cleared = await tasks_use_case.update_mutable_fields_on_task( + id=task.id, current_state=None + ) + + assert cleared.current_state is None + + async def test_update_current_state_and_metadata_single_atomic_write( + self, tasks_use_case, task_service, agent_repository, sample_agent + ): + """Supplying current_state + task_metadata together persists both via a + single atomic write (one publish), and does not clobber status.""" + await create_or_get_agent(agent_repository, sample_agent) + task = await task_service.create_task( + agent=sample_agent, task_name="current-state-combined-test" + ) + + spy = AsyncMock(wraps=task_service.update_mutable_fields) + task_service.update_mutable_fields = spy + + updated = await tasks_use_case.update_mutable_fields_on_task( + id=task.id, + task_metadata={"stage": "two"}, + current_state="working", + ) + + spy.assert_awaited_once() + assert updated.current_state == "working" + assert updated.task_metadata == {"stage": "two"} + assert updated.status == TaskStatus.RUNNING + + async def test_update_current_state_on_deleted_task_raises( + self, tasks_use_case, task_service, agent_repository, sample_agent + ): + """Updating current_state on a deleted task raises not found.""" + await create_or_get_agent(agent_repository, sample_agent) + task = await task_service.create_task( + agent=sample_agent, task_name="current-state-deleted-test" + ) + await tasks_use_case.delete_task(id=task.id) + + with pytest.raises(ItemDoesNotExist): + await tasks_use_case.update_mutable_fields_on_task( + id=task.id, current_state="working" + ) + async def test_update_metadata_on_deleted_task_raises( self, tasks_use_case, task_service, agent_repository, sample_agent ): From 8e0fda6d1c097ddfe4eae6158d3f63ca0ca9613a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mice=20P=C3=A1pai?= Date: Thu, 23 Jul 2026 13:34:33 +0200 Subject: [PATCH 3/5] test/fix(agentex): address round-2 audit on current_state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add the real use-case-level clobber regression test: it forces a stale read (entity fetched RUNNING, task transitioned to COMPLETED before the write) and asserts status stays COMPLETED. This fails on the old whole-row-merge path and passes on the column-scoped update — the prior service-level test could not catch the regression (it exercised the new primitive directly, which is clobber-proof by construction); its docstring is corrected accordingly. - update_mutable_fields_on_task now raises ItemDoesNotExist when the atomic update reports the row vanished, instead of silently returning stale data (consistent with the not-found contract and sibling transition methods; defensive — no live hard-delete path reaches it). - Drop max_length from the response Task.current_state field: input is already bounded by UpdateTaskRequest + the String(255) column, and enforcing it on the read path would turn a future column-widening into a 500 on every read. - Make the repository update_mutable_fields empty-fields branch's contract honest in its docstring (it's an unreachable defensive no-op). - Assert the intermediate set in the clear-to-null test so the clear is a real transition, not a null→null no-op. - Re-await the cancelled collector in the stream test to avoid a dangling-task warning at teardown. - Remove the now-unused `import sqlalchemy as sa` from the migration. - Regenerated openapi.yaml (maxLength now only on the request schema). --- ...200_add_task_current_state_b2c3d4e5f6a7.py | 1 - agentex/openapi.yaml | 2 - agentex/src/api/schemas/tasks.py | 4 +- .../domain/repositories/task_repository.py | 8 +++- .../src/domain/use_cases/tasks_use_case.py | 10 +++- agentex/tests/integration/test_task_stream.py | 3 ++ .../tests/unit/services/test_task_service.py | 9 ++-- .../unit/use_cases/test_tasks_use_case.py | 47 ++++++++++++++++++- 8 files changed, 71 insertions(+), 13 deletions(-) diff --git a/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py b/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py index 38986c48..2bb06d94 100644 --- a/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py +++ b/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py @@ -8,7 +8,6 @@ from typing import Sequence, Union from alembic import op -import sqlalchemy as sa # revision identifiers, used by Alembic. diff --git a/agentex/openapi.yaml b/agentex/openapi.yaml index f448d8e0..056cc3dc 100644 --- a/agentex/openapi.yaml +++ b/agentex/openapi.yaml @@ -6591,7 +6591,6 @@ 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'. @@ -6821,7 +6820,6 @@ 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'. diff --git a/agentex/src/api/schemas/tasks.py b/agentex/src/api/schemas/tasks.py index 260f2c86..934b7296 100644 --- a/agentex/src/api/schemas/tasks.py +++ b/agentex/src/api/schemas/tasks.py @@ -68,9 +68,11 @@ class Task(BaseModel): None, title="Task metadata", ) + # No max_length on this response field on purpose: input is already bounded + # by UpdateTaskRequest + the String(255) column, and enforcing the bound on + # the read path would turn a future column-widening into a 500 on every read. 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'." diff --git a/agentex/src/domain/repositories/task_repository.py b/agentex/src/domain/repositories/task_repository.py index 795770d6..3cbba0a0 100644 --- a/agentex/src/domain/repositories/task_repository.py +++ b/agentex/src/domain/repositories/task_repository.py @@ -264,11 +264,15 @@ async def update_mutable_fields( changed ``status``/``params``), this issues a single ``UPDATE ... SET 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. + or param merge is never reverted. With a non-empty ``fields`` this 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). + + An empty ``fields`` is a defensive no-op that returns the current task + (raising ``ItemDoesNotExist`` if absent); the sole caller never reaches it, + as it gates the call behind a non-empty field set. """ if not fields: return await self.get(id=task_id) diff --git a/agentex/src/domain/use_cases/tasks_use_case.py b/agentex/src/domain/use_cases/tasks_use_case.py index b1a87cdc..32a11280 100644 --- a/agentex/src/domain/use_cases/tasks_use_case.py +++ b/agentex/src/domain/use_cases/tasks_use_case.py @@ -158,8 +158,14 @@ async def update_mutable_fields_on_task( updated = await self.task_service.update_mutable_fields( task_entity.id, fields ) - if updated is not None: - task_entity = updated + if updated is None: + # The row vanished between the initial read and the write. No live + # hard-delete path reaches here today (delete is a soft status + # update, already guarded above), so this is defensive — but raise + # rather than return stale data, matching the not-found contract + # above and the sibling terminal-transition methods. + raise ItemDoesNotExist(f"Task {id or name} not found") + task_entity = updated return task_entity diff --git a/agentex/tests/integration/test_task_stream.py b/agentex/tests/integration/test_task_stream.py index d888109c..9a22adb1 100644 --- a/agentex/tests/integration/test_task_stream.py +++ b/agentex/tests/integration/test_task_stream.py @@ -275,6 +275,9 @@ async def collect_stream_events(): await stream_task except (TimeoutError, asyncio.CancelledError): stream_task.cancel() + # Re-await so the task's own cancellation cleanup runs and no + # "Task exception was never retrieved" warning leaks at teardown. + await asyncio.gather(stream_task, return_exceptions=True) task_updated_events = [ e for e in stream_events if e.get("type") == "task_updated" diff --git a/agentex/tests/unit/services/test_task_service.py b/agentex/tests/unit/services/test_task_service.py index df6ef9fd..9813af66 100644 --- a/agentex/tests/unit/services/test_task_service.py +++ b/agentex/tests/unit/services/test_task_service.py @@ -991,12 +991,13 @@ async def test_update_mutable_fields_persists_and_publishes( assert event_data["type"] == "task_updated" assert event_data["task"]["current_state"] == "working" - async def test_update_mutable_fields_does_not_clobber_status( + async def test_update_mutable_fields_leaves_status_untouched( self, task_service, agent_repository, sample_agent, redis_stream_repository ): - """Regression (blocker): a current_state write must not revert a status - changed by another writer. update_mutable_fields is column-scoped, so a - task that went terminal is not resurrected to its earlier status.""" + """The primitive is column-scoped: writing current_state does not touch + status. (The use-case-level clobber regression — a stale read racing a + status transition — is guarded in test_tasks_use_case.py; this only pins + that the repository UPDATE sets no columns beyond those supplied.)""" await create_or_get_agent(agent_repository, sample_agent) created_task = await task_service.create_task( agent=sample_agent, task_name="task-for-noclobber" diff --git a/agentex/tests/unit/use_cases/test_tasks_use_case.py b/agentex/tests/unit/use_cases/test_tasks_use_case.py index a2d6e737..a62bec28 100644 --- a/agentex/tests/unit/use_cases/test_tasks_use_case.py +++ b/agentex/tests/unit/use_cases/test_tasks_use_case.py @@ -639,9 +639,12 @@ async def test_update_current_state_clears_on_explicit_null( task = await task_service.create_task( agent=sample_agent, task_name="current-state-clear-test" ) - await tasks_use_case.update_mutable_fields_on_task( + was_set = await tasks_use_case.update_mutable_fields_on_task( id=task.id, current_state="working" ) + # Confirm it was actually set, so the clear below is a real transition + # (not a null→null no-op that would pass trivially). + assert was_set.current_state == "working" cleared = await tasks_use_case.update_mutable_fields_on_task( id=task.id, current_state=None @@ -673,6 +676,48 @@ async def test_update_current_state_and_metadata_single_atomic_write( assert updated.task_metadata == {"stage": "two"} assert updated.status == TaskStatus.RUNNING + async def test_update_current_state_does_not_clobber_concurrent_status( + self, tasks_use_case, task_service, task_repository, agent_repository, sample_agent + ): + """Regression (round-one blocker): a current_state write must not revert a + status changed concurrently. Reproduces the exact staleness the old + whole-row-merge path needed — the use case reads the entity while RUNNING, + the task goes COMPLETED before the write lands, and the write must keep + COMPLETED. Fails on the old update_task/merge path (reverts to RUNNING), + passes on the column-scoped update. + """ + await create_or_get_agent(agent_repository, sample_agent) + task = await task_service.create_task( + agent=sample_agent, task_name="current-state-clobber-test" + ) + + # Snapshot the entity as the use case would have read it, BEFORE the + # concurrent transition — this is the stale read the bug wrote back. + stale_entity = await task_service.get_task(id=task.id) + assert stale_entity.status == TaskStatus.RUNNING + + # Another writer moves the task to a terminal status after that read. + await task_service.transition_task_status( + task_id=task.id, + expected_status=TaskStatus.RUNNING, + new_status=TaskStatus.COMPLETED, + status_reason="done", + ) + + # Force the use case to operate on the stale (pre-transition) read. + task_service.get_task = AsyncMock(return_value=stale_entity) + + updated = await tasks_use_case.update_mutable_fields_on_task( + id=task.id, current_state="late" + ) + + # The write set current_state without reverting the terminal status. + assert updated.current_state == "late" + assert updated.status == TaskStatus.COMPLETED + persisted = await task_repository.get(id=task.id) + assert persisted.status == TaskStatus.COMPLETED + assert persisted.current_state == "late" + async def test_update_current_state_on_deleted_task_raises( self, tasks_use_case, task_service, agent_repository, sample_agent ): From 0f544773447d18a74fc8f115ef1bf767b5e34b67 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mice=20P=C3=A1pai?= Date: Thu, 23 Jul 2026 15:33:09 +0200 Subject: [PATCH 4/5] style(agentex): collapse verbose comments/docstrings to one line Self-review cleanup: tighten the multi-line comments and docstrings added for current_state down to single lines (drop AI-slop narration; keep only the non-obvious why), per the repo's comment-discipline convention. --- ...200_add_task_current_state_b2c3d4e5f6a7.py | 5 +--- agentex/src/adapters/orm.py | 6 +---- agentex/src/api/routes/tasks.py | 6 ++--- agentex/src/api/schemas/tasks.py | 8 ++---- .../domain/repositories/task_repository.py | 21 ++++----------- agentex/src/domain/services/task_service.py | 10 ++----- .../src/domain/use_cases/tasks_use_case.py | 27 +++++-------------- .../integration/api/tasks/test_tasks_api.py | 13 +++------ agentex/tests/integration/test_task_stream.py | 12 +++------ .../tests/unit/services/test_task_service.py | 13 ++++----- .../unit/use_cases/test_tasks_use_case.py | 17 ++++-------- 11 files changed, 37 insertions(+), 101 deletions(-) diff --git a/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py b/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py index 2bb06d94..468e51d5 100644 --- a/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py +++ b/agentex/database/migrations/alembic/versions/2026_07_22_1200_add_task_current_state_b2c3d4e5f6a7.py @@ -18,10 +18,7 @@ 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. - # 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)). + # Nullable additive column; idempotent, metadata-only, non-blocking. op.execute( "ALTER TABLE tasks ADD COLUMN IF NOT EXISTS current_state VARCHAR(255)" ) diff --git a/agentex/src/adapters/orm.py b/agentex/src/adapters/orm.py index 45c09fa4..6e12a586 100644 --- a/agentex/src/adapters/orm.py +++ b/agentex/src/adapters/orm.py @@ -75,11 +75,7 @@ class TaskORM(BaseORM): cleaned_at = Column(DateTime(timezone=True), nullable=True) params = Column(JSONB, nullable=True) 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). 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. + # Opaque agent-state label, orthogonal to `status`; capped since it rides every task_updated SSE payload. current_state = Column(String(255), nullable=True) # Many-to-Many relationship with agents agents = relationship("AgentORM", secondary="task_agents", back_populates="tasks") diff --git a/agentex/src/api/routes/tasks.py b/agentex/src/api/routes/tasks.py index f1f297e7..b87503d3 100644 --- a/agentex/src/api/routes/tasks.py +++ b/agentex/src/api/routes/tasks.py @@ -196,8 +196,7 @@ async def update_task( id=task_id, task_metadata=request.task_metadata, merge_params=request.merge_params, - # Pass through only when the client actually sent the key, so an explicit - # null clears the label while an omitted field leaves it untouched. + # Forward only when the client sent the key: explicit null clears, omitted leaves unchanged. current_state=( request.current_state if "current_state" in request.model_fields_set @@ -224,8 +223,7 @@ async def update_task_by_name( name=task_name, task_metadata=request.task_metadata, merge_params=request.merge_params, - # Pass through only when the client actually sent the key, so an explicit - # null clears the label while an omitted field leaves it untouched. + # Forward only when the client sent the key: explicit null clears, omitted leaves unchanged. current_state=( request.current_state if "current_state" in request.model_fields_set diff --git a/agentex/src/api/schemas/tasks.py b/agentex/src/api/schemas/tasks.py index 934b7296..5b995f40 100644 --- a/agentex/src/api/schemas/tasks.py +++ b/agentex/src/api/schemas/tasks.py @@ -7,9 +7,7 @@ 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. +# Bound on the current_state label; mirrors the String(255) column width. CURRENT_STATE_MAX_LENGTH = 255 @@ -68,9 +66,7 @@ class Task(BaseModel): None, title="Task metadata", ) - # No max_length on this response field on purpose: input is already bounded - # by UpdateTaskRequest + the String(255) column, and enforcing the bound on - # the read path would turn a future column-widening into a 500 on every read. + # No bound on the read path: enforcing it would 500 if the column is ever widened (writes are already bounded). current_state: str | None = Field( None, title=( diff --git a/agentex/src/domain/repositories/task_repository.py b/agentex/src/domain/repositories/task_repository.py index 3cbba0a0..1f5f797f 100644 --- a/agentex/src/domain/repositories/task_repository.py +++ b/agentex/src/domain/repositories/task_repository.py @@ -257,22 +257,11 @@ async def merge_params(self, task_id: str, patch: dict) -> TaskEntity | None: 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 WHERE id = :id RETURNING *``. - Touching only the supplied columns means a concurrent status transition - or param merge is never reverted. With a non-empty ``fields`` this 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). - - An empty ``fields`` is a defensive no-op that returns the current task - (raising ``ItemDoesNotExist`` if absent); the sole caller never reaches it, - as it gates the call behind a non-empty field set. + """Atomically ``UPDATE ... SET WHERE id`` — column-scoped, + so (unlike ``update``'s whole-row merge) it can't clobber a concurrently changed + ``status``/``params``. Values apply verbatim (``current_state=None`` clears). Returns + the updated entity, or ``None`` if the task doesn't exist. Empty ``fields`` is an + unreachable defensive no-op (the sole caller gates on a non-empty set). """ if not fields: return await self.get(id=task_id) diff --git a/agentex/src/domain/services/task_service.py b/agentex/src/domain/services/task_service.py index 90cb76a5..347266ad 100644 --- a/agentex/src/domain/services/task_service.py +++ b/agentex/src/domain/services/task_service.py @@ -245,14 +245,8 @@ async def update_task(self, task: TaskEntity) -> TaskEntity: 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. + """Column-scoped atomic update of the given columns, then publish task_updated. + Returns the updated entity, or ``None`` if the task no longer exists. """ updated_task = await self.task_repository.update_mutable_fields( task_id, fields diff --git a/agentex/src/domain/use_cases/tasks_use_case.py b/agentex/src/domain/use_cases/tasks_use_case.py index 32a11280..a54ff321 100644 --- a/agentex/src/domain/use_cases/tasks_use_case.py +++ b/agentex/src/domain/use_cases/tasks_use_case.py @@ -12,9 +12,7 @@ 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.""" + """Sentinel: PATCH field omitted (untouched) vs. explicitly null (cleared).""" UNSET: Any = _Unset() @@ -106,11 +104,8 @@ async def update_mutable_fields_on_task( merge_params: dict[str, Any] | 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. - - ``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"). + """Update mutable fields on a task. ``current_state`` uses the UNSET sentinel + (explicit null clears, omitted leaves it); ``task_metadata`` None means "not supplied". """ if not id and not name: @@ -134,10 +129,7 @@ async def update_mutable_fields_on_task( ): return task_entity - # `merge_params` is a separate atomic JSONB shallow-merge so concurrent - # 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. + # Atomic JSONB shallow-merge; run first so its refreshed entity is the fallback return. if merge_params: merged = await self.task_service.merge_task_params( task_entity.id, merge_params @@ -145,10 +137,7 @@ async def update_mutable_fields_on_task( if merged is not None: task_entity = merged - # 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. + # Single column-scoped write → one task_updated publish, no whole-row clobber. fields: dict[str, Any] = {} if task_metadata is not None: fields["task_metadata"] = task_metadata @@ -159,11 +148,7 @@ async def update_mutable_fields_on_task( task_entity.id, fields ) if updated is None: - # The row vanished between the initial read and the write. No live - # hard-delete path reaches here today (delete is a soft status - # update, already guarded above), so this is defensive — but raise - # rather than return stale data, matching the not-found contract - # above and the sibling terminal-transition methods. + # Row vanished mid-flight (defensive; no live hard-delete path). Raise, don't return stale. raise ItemDoesNotExist(f"Task {id or name} not found") task_entity = updated diff --git a/agentex/tests/integration/api/tasks/test_tasks_api.py b/agentex/tests/integration/api/tasks/test_tasks_api.py index 492d87d7..f292434d 100644 --- a/agentex/tests/integration/api/tasks/test_tasks_api.py +++ b/agentex/tests/integration/api/tasks/test_tasks_api.py @@ -686,8 +686,7 @@ async def test_update_task_endpoint_success( async def test_update_task_current_state( self, isolated_client, isolated_repositories ): - """PUT /tasks/{id} writes current_state; explicit null clears it while an - omitted field leaves it untouched; point-read is source of truth.""" + """PUT current_state: explicit null clears, omitted leaves it, point-read reconciles.""" agent_repo = isolated_repositories["agent_repository"] agent = AgentEntity( id=orm_id(), @@ -743,8 +742,7 @@ async def test_update_task_current_state( async def test_update_task_current_state_and_metadata_together( self, isolated_client, isolated_repositories ): - """current_state + task_metadata in one PUT both persist (single write), - without clobbering status.""" + """current_state + task_metadata in one PUT both persist without clobbering status.""" agent_repo = isolated_repositories["agent_repository"] agent = AgentEntity( id=orm_id(), @@ -807,8 +805,7 @@ async def test_update_task_current_state_by_name( async def test_update_task_current_state_empty_string( self, isolated_client, isolated_repositories ): - """Empty string is a valid label distinct from null (guards against a - future falsy check treating "" as "unset").""" + """Empty string is a valid label distinct from null (guards a falsy-check regression).""" agent_repo = isolated_repositories["agent_repository"] agent = AgentEntity( id=orm_id(), @@ -865,9 +862,7 @@ async def test_update_task_current_state_too_long_rejected( async def test_update_task_request_ignores_unknown_fields( self, isolated_client, isolated_repositories ): - """Guards the extra="ignore" assumption a newer SDK relies on: a payload - carrying fields the server doesn't model must 200 (drop them), not 422. - Breaks loudly if UpdateTaskRequest ever switches to extra="forbid".""" + """Unknown fields are ignored (200, not 422) — guards the extra="ignore" SDK-compat assumption.""" agent_repo = isolated_repositories["agent_repository"] agent = AgentEntity( id=orm_id(), diff --git a/agentex/tests/integration/test_task_stream.py b/agentex/tests/integration/test_task_stream.py index 9a22adb1..ae0786a6 100644 --- a/agentex/tests/integration/test_task_stream.py +++ b/agentex/tests/integration/test_task_stream.py @@ -233,8 +233,7 @@ async def collect_stream_events(): async def test_current_state_update_triggers_stream_event( self, test_agent_and_task, tasks_use_case, streams_use_case ): - """current_state rides the existing task_updated event, so a client - subscribed to the task stream is pushed the new state reactively.""" + """current_state rides the existing task_updated event (reactive push to subscribers).""" _agent, task = test_agent_and_task stream_events = [] @@ -260,23 +259,20 @@ async def collect_stream_events(): pass stream_task = asyncio.create_task(collect_stream_events()) - # Let the subscription establish before the update; the stream is - # tail-only, so an event published before we subscribe would be missed. + # Let the tail-only subscription establish before the update, or the event is missed. await asyncio.sleep(0.1) updated_task = await tasks_use_case.update_mutable_fields_on_task( id=task.id, current_state="awaiting_input" ) - # Wait until the collector sees task_updated (it breaks on it) rather - # than guessing a fixed delay; the timeout is only a failure ceiling. + # Wait for the collector to see task_updated (it breaks on it); timeout is only a ceiling. try: async with asyncio.timeout(5): await stream_task except (TimeoutError, asyncio.CancelledError): stream_task.cancel() - # Re-await so the task's own cancellation cleanup runs and no - # "Task exception was never retrieved" warning leaks at teardown. + # Re-await so cancellation cleanup runs and no dangling-task warning leaks at teardown. await asyncio.gather(stream_task, return_exceptions=True) task_updated_events = [ diff --git a/agentex/tests/unit/services/test_task_service.py b/agentex/tests/unit/services/test_task_service.py index 9813af66..68f4865e 100644 --- a/agentex/tests/unit/services/test_task_service.py +++ b/agentex/tests/unit/services/test_task_service.py @@ -941,8 +941,7 @@ async def test_update_task_with_task_metadata_changes( async def test_update_task_current_state_publishes_stream_event( self, task_service, agent_repository, sample_agent, redis_stream_repository ): - """update_task persists current_state and carries it on the published - task_updated event, so subscribed clients see the new state.""" + """update_task persists current_state and carries it on the task_updated event.""" await create_or_get_agent(agent_repository, sample_agent) created_task = await task_service.create_task( agent=sample_agent, task_name="task-for-current-state" @@ -967,8 +966,7 @@ async def test_update_task_current_state_publishes_stream_event( async def test_update_mutable_fields_persists_and_publishes( self, task_service, agent_repository, sample_agent, redis_stream_repository ): - """update_mutable_fields writes only the given columns and publishes a - task_updated event carrying them.""" + """update_mutable_fields persists the given columns and publishes task_updated.""" await create_or_get_agent(agent_repository, sample_agent) created_task = await task_service.create_task( agent=sample_agent, task_name="task-for-mutable-fields" @@ -994,10 +992,9 @@ async def test_update_mutable_fields_persists_and_publishes( async def test_update_mutable_fields_leaves_status_untouched( self, task_service, agent_repository, sample_agent, redis_stream_repository ): - """The primitive is column-scoped: writing current_state does not touch - status. (The use-case-level clobber regression — a stale read racing a - status transition — is guarded in test_tasks_use_case.py; this only pins - that the repository UPDATE sets no columns beyond those supplied.)""" + """The primitive is column-scoped: writing current_state leaves status untouched + (the use-case stale-read clobber regression is guarded in test_tasks_use_case.py). + """ await create_or_get_agent(agent_repository, sample_agent) created_task = await task_service.create_task( agent=sample_agent, task_name="task-for-noclobber" diff --git a/agentex/tests/unit/use_cases/test_tasks_use_case.py b/agentex/tests/unit/use_cases/test_tasks_use_case.py index a62bec28..cb207212 100644 --- a/agentex/tests/unit/use_cases/test_tasks_use_case.py +++ b/agentex/tests/unit/use_cases/test_tasks_use_case.py @@ -642,8 +642,7 @@ async def test_update_current_state_clears_on_explicit_null( was_set = await tasks_use_case.update_mutable_fields_on_task( id=task.id, current_state="working" ) - # Confirm it was actually set, so the clear below is a real transition - # (not a null→null no-op that would pass trivially). + # Confirm it was set, so the clear below is a real transition (not a trivial null→null). assert was_set.current_state == "working" cleared = await tasks_use_case.update_mutable_fields_on_task( @@ -655,8 +654,7 @@ async def test_update_current_state_clears_on_explicit_null( async def test_update_current_state_and_metadata_single_atomic_write( self, tasks_use_case, task_service, agent_repository, sample_agent ): - """Supplying current_state + task_metadata together persists both via a - single atomic write (one publish), and does not clobber status.""" + """current_state + task_metadata together persist via one atomic write (one publish).""" await create_or_get_agent(agent_repository, sample_agent) task = await task_service.create_task( agent=sample_agent, task_name="current-state-combined-test" @@ -679,20 +677,15 @@ async def test_update_current_state_and_metadata_single_atomic_write( async def test_update_current_state_does_not_clobber_concurrent_status( self, tasks_use_case, task_service, task_repository, agent_repository, sample_agent ): - """Regression (round-one blocker): a current_state write must not revert a - status changed concurrently. Reproduces the exact staleness the old - whole-row-merge path needed — the use case reads the entity while RUNNING, - the task goes COMPLETED before the write lands, and the write must keep - COMPLETED. Fails on the old update_task/merge path (reverts to RUNNING), - passes on the column-scoped update. + """Regression: a stale RUNNING read racing a COMPLETED transition must not revert + status on the current_state write. Fails on the old whole-row merge, passes column-scoped. """ await create_or_get_agent(agent_repository, sample_agent) task = await task_service.create_task( agent=sample_agent, task_name="current-state-clobber-test" ) - # Snapshot the entity as the use case would have read it, BEFORE the - # concurrent transition — this is the stale read the bug wrote back. + # Snapshot the entity BEFORE the transition — the stale read the old bug wrote back. stale_entity = await task_service.get_task(id=task.id) assert stale_entity.status == TaskStatus.RUNNING From 7a5907071405814894926b64e0ffdc50e3337d0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mice=20P=C3=A1pai?= Date: Fri, 24 Jul 2026 22:33:14 +0200 Subject: [PATCH 5/5] refactor(agentex): address review on current_state write path - Guardrail update_mutable_fields against writing non-mutable columns: an allowlist (task_metadata, current_state) rejects anything else, so a future caller can't route status/params through it and bypass their atomic CAS. - Dedupe merge_params and update_mutable_fields into a shared _update_returning helper (one UPDATE ... RETURNING body); drop the unreachable empty-fields branch that contradicted the None-on-missing contract. - Single-source the current_state bound as CURRENT_STATE_MAX_LENGTH (new src/utils/task_constants.py) so the request-validation limit and the String(...) column width can't drift; both import it. - Correct the sentinel docstring wording (partial update, not PATCH) and ruff-format the two files the formatter had left unwrapped. --- agentex/src/adapters/orm.py | 3 +- agentex/src/api/schemas/tasks.py | 4 +- .../domain/repositories/task_repository.py | 57 +++++++++---------- agentex/src/domain/services/task_service.py | 10 ++-- .../src/domain/use_cases/tasks_use_case.py | 2 +- agentex/src/utils/task_constants.py | 3 + .../unit/use_cases/test_tasks_use_case.py | 7 ++- 7 files changed, 46 insertions(+), 40 deletions(-) create mode 100644 agentex/src/utils/task_constants.py diff --git a/agentex/src/adapters/orm.py b/agentex/src/adapters/orm.py index 6e12a586..0851afad 100644 --- a/agentex/src/adapters/orm.py +++ b/agentex/src/adapters/orm.py @@ -24,6 +24,7 @@ from src.domain.entities.deployments import DeploymentStatus from src.domain.entities.tasks import TaskStatus from src.utils.ids import orm_id +from src.utils.task_constants import CURRENT_STATE_MAX_LENGTH BaseORM = declarative_base() @@ -76,7 +77,7 @@ class TaskORM(BaseORM): params = Column(JSONB, nullable=True) task_metadata = Column(JSONB, nullable=True) # Opaque agent-state label, orthogonal to `status`; capped since it rides every task_updated SSE payload. - current_state = Column(String(255), nullable=True) + current_state = Column(String(CURRENT_STATE_MAX_LENGTH), nullable=True) # Many-to-Many relationship with agents agents = relationship("AgentORM", secondary="task_agents", back_populates="tasks") diff --git a/agentex/src/api/schemas/tasks.py b/agentex/src/api/schemas/tasks.py index 5b995f40..5b53f6e1 100644 --- a/agentex/src/api/schemas/tasks.py +++ b/agentex/src/api/schemas/tasks.py @@ -6,9 +6,7 @@ from src.api.schemas.agents import Agent from src.utils.model_utils import BaseModel - -# Bound on the current_state label; mirrors the String(255) column width. -CURRENT_STATE_MAX_LENGTH = 255 +from src.utils.task_constants import CURRENT_STATE_MAX_LENGTH class TaskRelationships(str, Enum): diff --git a/agentex/src/domain/repositories/task_repository.py b/agentex/src/domain/repositories/task_repository.py index 1f5f797f..439a6802 100644 --- a/agentex/src/domain/repositories/task_repository.py +++ b/agentex/src/domain/repositories/task_repository.py @@ -21,6 +21,10 @@ logger = make_logger(__name__) +# Columns update_mutable_fields is allowed to set — a code guardrail (not a comment) +# so a future caller can't route status/params through it and bypass their atomic CAS. +_MUTABLE_TASK_COLUMNS = frozenset({"task_metadata", "current_state"}) + class TaskRepository(PostgresCRUDRepository[TaskORM, TaskEntity, TaskRelationships]): """Repository for Task entity with relationship loading support""" @@ -231,41 +235,34 @@ async def merge_params(self, task_id: str, patch: dict) -> TaskEntity | None: general-purpose updater. """ - async with ( - self.start_async_db_session(True) as session, - async_sql_exception_handler(), - ): - # ``COALESCE(params, '{}'::jsonb)`` so a NULL existing value - # doesn't poison the concat to NULL. Both operands cast to - # JSONB explicitly so Postgres picks the JSONB ``||`` operator - # (not the text concat overload). - existing = func.coalesce(TaskORM.params, cast({}, JSONB)) - merged = existing.op("||", return_type=JSONB)(cast(patch, JSONB)) - stmt = ( - update(TaskORM) - .where(TaskORM.id == task_id) - .values(params=merged) - .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) + # ``COALESCE(params, '{}'::jsonb)`` so a NULL existing value doesn't poison the + # concat; explicit JSONB casts so Postgres picks the jsonb ``||`` (not text concat). + existing = func.coalesce(TaskORM.params, cast({}, JSONB)) + merged = existing.op("||", return_type=JSONB)(cast(patch, JSONB)) + return await self._update_returning(task_id, {"params": merged}) async def update_mutable_fields( self, task_id: str, fields: dict[str, Any] ) -> TaskEntity | None: - """Atomically ``UPDATE ... SET WHERE id`` — column-scoped, - so (unlike ``update``'s whole-row merge) it can't clobber a concurrently changed - ``status``/``params``. Values apply verbatim (``current_state=None`` clears). Returns - the updated entity, or ``None`` if the task doesn't exist. Empty ``fields`` is an - unreachable defensive no-op (the sole caller gates on a non-empty set). + """Atomically set the given caller-mutable columns on one task row via + ``UPDATE ... RETURNING`` — column-scoped, so (unlike ``update``'s whole-row merge) it + can't clobber a concurrently changed ``status``/``params``. Values apply verbatim + (``current_state=None`` clears). Returns the updated entity, or ``None`` if absent. """ - if not fields: - return await self.get(id=task_id) + unknown = fields.keys() - _MUTABLE_TASK_COLUMNS + if unknown: + raise ValueError( + f"update_mutable_fields may only set {sorted(_MUTABLE_TASK_COLUMNS)}; " + f"got disallowed columns {sorted(unknown)}" + ) + return await self._update_returning(task_id, fields) + async def _update_returning( + self, task_id: str, values: dict[str, Any] + ) -> TaskEntity | None: + """``UPDATE tasks SET WHERE id`` → the updated entity, or ``None`` if no + such task exists. Shared by merge_params and update_mutable_fields. + """ async with ( self.start_async_db_session(True) as session, async_sql_exception_handler(), @@ -273,7 +270,7 @@ async def update_mutable_fields( stmt = ( update(TaskORM) .where(TaskORM.id == task_id) - .values(**fields) + .values(**values) .returning(TaskORM) ) result = await session.execute(stmt) diff --git a/agentex/src/domain/services/task_service.py b/agentex/src/domain/services/task_service.py index 347266ad..59e2e6cf 100644 --- a/agentex/src/domain/services/task_service.py +++ b/agentex/src/domain/services/task_service.py @@ -248,9 +248,7 @@ async def update_mutable_fields( """Column-scoped atomic update of the given columns, then publish task_updated. Returns the updated entity, or ``None`` if the task no longer exists. """ - updated_task = await self.task_repository.update_mutable_fields( - task_id, fields - ) + updated_task = await self.task_repository.update_mutable_fields(task_id, fields) if updated_task is None: return None @@ -402,7 +400,11 @@ async def cancel_task( new_status=TaskStatus.CANCELED, status_reason="Task canceled by user", ) - return updated if updated is not None else await self.task_repository.get(id=task.id) + return ( + updated + if updated is not None + else await self.task_repository.get(id=task.id) + ) async def interrupt_task( self, agent: AgentEntity, task: TaskEntity, acp_url: str diff --git a/agentex/src/domain/use_cases/tasks_use_case.py b/agentex/src/domain/use_cases/tasks_use_case.py index a54ff321..3812b348 100644 --- a/agentex/src/domain/use_cases/tasks_use_case.py +++ b/agentex/src/domain/use_cases/tasks_use_case.py @@ -12,7 +12,7 @@ class _Unset: - """Sentinel: PATCH field omitted (untouched) vs. explicitly null (cleared).""" + """Sentinel for partial updates: field omitted (untouched) vs. explicitly null (cleared).""" UNSET: Any = _Unset() diff --git a/agentex/src/utils/task_constants.py b/agentex/src/utils/task_constants.py new file mode 100644 index 00000000..d1f0b46f --- /dev/null +++ b/agentex/src/utils/task_constants.py @@ -0,0 +1,3 @@ +# Single source for the current_state bound: the request-validation limit +# (UpdateTaskRequest) and the storage limit (TaskORM.current_state column width). +CURRENT_STATE_MAX_LENGTH = 255 diff --git a/agentex/tests/unit/use_cases/test_tasks_use_case.py b/agentex/tests/unit/use_cases/test_tasks_use_case.py index cb207212..bc523076 100644 --- a/agentex/tests/unit/use_cases/test_tasks_use_case.py +++ b/agentex/tests/unit/use_cases/test_tasks_use_case.py @@ -675,7 +675,12 @@ async def test_update_current_state_and_metadata_single_atomic_write( assert updated.status == TaskStatus.RUNNING async def test_update_current_state_does_not_clobber_concurrent_status( - self, tasks_use_case, task_service, task_repository, agent_repository, sample_agent + self, + tasks_use_case, + task_service, + task_repository, + agent_repository, + sample_agent, ): """Regression: a stale RUNNING read racing a COMPLETED transition must not revert status on the current_state write. Fails on the old whole-row merge, passes column-scoped.