A batch message with task_id and subtask_id attaches to the app-owned task 🟩

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Khalim Conn-Kowlessar 2026-07-07 11:48:53 +00:00
parent ca6ca2a220
commit f01131a064
2 changed files with 39 additions and 25 deletions

View file

@ -5,7 +5,7 @@ from typing import Any
from uuid import UUID
import pytest
from sqlalchemy import Engine
from sqlalchemy import Engine, text
from sqlmodel import Session
from domain.tasks.tasks import Source
@ -148,14 +148,9 @@ def test_task_handler_attaches_to_a_supplied_task_and_subtask(
{"task_id": str(task.id), "subtask_id": str(subtask.id)}
]
assert harness.subtasks.get(subtask.id).status.value == "complete"
with Session(db_engine) as session:
from sqlmodel import select
from backend.app.db.models.tasks import SubTask as SubTaskRow
from backend.app.db.models.tasks import Task as TaskRow
assert len(session.exec(select(TaskRow)).all()) == 1
assert len(session.exec(select(SubTaskRow)).all()) == 1
with db_engine.connect() as conn:
assert conn.execute(text("SELECT count(*) FROM tasks")).scalar_one() == 1
assert conn.execute(text("SELECT count(*) FROM sub_task")).scalar_one() == 1
def test_task_handler_leaves_cloudwatch_url_unset_outside_lambda(

View file

@ -9,6 +9,7 @@ import logging
from contextlib import AbstractContextManager
from functools import wraps
from typing import Any, Callable, Optional, cast
from uuid import UUID
from utilities.aws_lambda.cloud_logs import cloudwatch_url
from utilities.aws_lambda.default_orchestrator import default_orchestrator
@ -50,39 +51,47 @@ def task_handler(
for record in _records(event):
body = _parse_body(record)
raw_source_id = body.get(source.value)
source_id = (
str(raw_source_id) if raw_source_id is not None else None
)
task, subtask = orchestrator.create_task_with_subtask(
task_source=task_source,
inputs=body,
source=source,
source_id=source_id,
)
# Attach mode (ADR-0055): a Modelling Run batch carries the
# app-owned task_id and its distributor-pre-created
# subtask_id — run under those; create nothing.
supplied = _supplied_ids(body)
if supplied is not None:
task_id, subtask_id = supplied
else:
raw_source_id = body.get(source.value)
source_id = (
str(raw_source_id) if raw_source_id is not None else None
)
task, subtask = orchestrator.create_task_with_subtask(
task_source=task_source,
inputs=body,
source=source,
source_id=source_id,
)
task_id, subtask_id = task.id, subtask.id
task_ids.append(
{"task_id": str(task.id), "subtask_id": str(subtask.id)}
{"task_id": str(task_id), "subtask_id": str(subtask_id)}
)
try:
if pass_task_orchestrator:
orchestrator.run_subtask(
subtask.id,
work=lambda: func(body, context, orchestrator, task.id),
subtask_id,
work=lambda: func(body, context, orchestrator, task_id),
cloud_logs_url=cloud_logs_url,
)
else:
orchestrator.run_subtask(
subtask.id,
subtask_id,
work=lambda: func(body, context),
cloud_logs_url=cloud_logs_url,
)
except Exception:
logger.exception(
"subtask failed (task_source=%s source_id=%s)",
"subtask failed (task_source=%s subtask_id=%s)",
task_source,
source_id,
subtask_id,
)
if "Records" in event:
message_id = record.get("messageId", "")
@ -99,6 +108,16 @@ def task_handler(
return decorator
def _supplied_ids(body: dict[str, Any]) -> Optional[tuple[UUID, UUID]]:
"""The (task_id, subtask_id) an attach-mode message carries, or None for
the classic create-your-own-task message."""
raw_task_id = body.get("task_id")
raw_subtask_id = body.get("subtask_id")
if raw_task_id is None or raw_subtask_id is None:
return None
return UUID(str(raw_task_id)), UUID(str(raw_subtask_id))
def _parse_body(record: dict[str, Any]) -> dict[str, Any]:
raw = record.get("body", record)
if isinstance(raw, str):