fix: expose terminal task_updated events as typed TaskUpdatedMessage

[maxim092001]

Fixes #1019

Summary

Background tasks can leave consumers with stale active-task state. A background Bash task sometimes finishes by emitting only a generic system/task_updated message whose patch.status is terminal (completed/failed/stopped), with no typed TaskNotificationMessage:

TaskStartedMessage id=bs2r8eew4
SystemMessage subtype=task_updated id=bs2r8eew4 status=completed   # <- only this
AssistantMessage text=BG_WAIT_FINAL
ResultMessage active=['bs2r8eew4']                                  # consumer still thinks it's running
TIMEOUT active=['bs2r8eew4']

Consumers that track active task IDs from TaskStartedMessage and clear them only on TaskNotificationMessage then believe a finished task is still active. If they use that active set to bound receive_messages() (which is open-ended for the persistent client), they keep draining the stream until an outer timeout.

The parser previously mapped system/task_updated to a generic SystemMessage, so terminal task state — semantically lifecycle data — was never represented by a typed lifecycle message. This breaks the contract: a typed start event (TaskStartedMessage) was not guaranteed a typed terminal event.

Fix

Expose system/task_updated as a typed TaskUpdatedMessage(SystemMessage), mirroring the TypeScript SDK's SDKTaskUpdatedMessage:

field	source
task_id	data.task_id
patch	data.patch (full dict preserved)
status	patch.status
session_id	data.session_id
uuid	data.uuid

Plus:

• TaskUpdatedStatus — Literal["running","completed","failed","stopped","cancelled"].
• TERMINAL_TASK_STATUSES — shared frozenset of finished statuses, so consumers can clear active task IDs on a terminal status from either TaskNotificationMessage or TaskUpdatedMessage.
• Lifecycle contract documented on both TaskUpdatedMessage and TaskNotificationMessage.

After this change, a bounded drain works without hanging:

if isinstance(message, (TaskNotificationMessage, TaskUpdatedMessage)) and message.status in TERMINAL_TASK_STATUSES:
    active.discard(message.task_id)
elif isinstance(message, ResultMessage) and not active:
    break

Defensive parsing

task_updated parsing uses .get() throughout, guards a non-dict patch, and derives status from patch.status — it can never raise on a lifecycle event (the observed CLI payload omits uuid/session_id). A patch carrying only end_time/result/error is left non-terminal (status=None) with the full patch preserved for callers that need more.

Backward compatibility

• TaskUpdatedMessage subclasses SystemMessage, so existing isinstance(msg, SystemMessage) and case SystemMessage() checks still match; subtype and data remain populated with the raw payload.
• Follows the existing convention where SystemMessage subclasses are covered structurally by the Message union rather than listed individually.

Tests

Added parser tests for: terminal completed; the minimal observed shape (no uuid/session_id); running/non-terminal; missing patch; patch present without status (preserved verbatim); non-dict/None patch (parametrized — never raises); all terminal statuses completed/failed/stopped/cancelled (parametrized); and SystemMessage backward-compat.

ruff check, ruff format, mypy src/, and the full suite (775 passed, 5 skipped) all green.

🤖 Generated with Claude Code

[maxim092001]

Hey @qing-ant, could you please check it out? Thanks!

[maxim092001]

Or maybe @ashwin-ant you can take a look?

[ashwin-ant]

Thanks for digging into this, the root cause analysis is right. Terminal state for a background shell can show up as only a task_updated patch when no follow-up turn drains the notification, so typing this message is the correct fix.

One blocking issue though: the status enum doesn't match what the CLI actually sends. The set of values that can appear in patch.status is

pending | running | completed | failed | killed | paused

This PR declares TaskUpdatedStatus = Literal["running", "completed", "failed", "stopped", "cancelled"], which has a few problems:

1. stopped and cancelled never appear in a task_updated patch. stopped only exists on task_notification, where the CLI maps killed to stopped before emitting. cancelled doesn't exist at all.
2. killed is missing, and it's terminal. A task stopped via TaskStop emits task_updated with status: "killed". With the proposed TERMINAL_TASK_STATUSES that event isn't recognized as terminal, so the stale-active-task hang this PR fixes comes right back for stopped tasks. In some kill paths the notification is suppressed entirely, so task_updated is the only terminal signal. Easy to reproduce: start a background bash task, stop it with TaskStop, and watch the patch.
3. pending and paused are missing, so a paused task's patch produces a value that violates the declared Literal.

Suggested fix:

• TaskUpdatedStatus = Literal["pending", "running", "completed", "failed", "killed", "paused"]
• Terminal set for task_updated is {completed, failed, killed}. If you want one constant shared with TaskNotificationMessage, it needs both vocabularies: {completed, failed, stopped, killed}. Either way, drop cancelled.

Two smaller things:

• I'd soften the lifecycle contract docstring. "Every TaskStartedMessage gets a typed terminal event" is a stronger guarantee than the CLI makes today.
• Current CLI builds stamp uuid and session_id on every drained event, so the defensive handling is fine to keep but probably reflects an older build.

The shape, tests, and backward-compat approach all look good. Happy to approve once the enum is fixed.