Add GatewayClient.cancel_session() to cancel all active tasks in a session

[zhchxiao123]

Problem Statement

When a caller wants to abort an in-flight agent interaction, GatewayClient.cancel_task() requires the caller to know the specific message_id of the task (or one of its ancestors) they want to cancel, then walks the parent_message_id tree from that node outward via BFS to find the sub-tree to cancel. This works well when the caller knows exactly which task/turn to target, but there's no way to say "stop everything currently running in this session" without first looking up a message_id to seed the walk. Callers that just want to abandon a session outright (e.g. the end user closed the chat, or a new turn should supersede an old one) have to reconstruct or guess a message_id before they can cancel anything.

Solution

Add a GatewayClient.cancel_session(session_id, reason="") method that cancels every non-terminal execution registered under a given session_id, without requiring a starting message_id. Since WorkerRegistry.get_all_session_executions(session_id) already returns every execution record for a session in one flat call (it's the same data cancel_task's BFS reconstructs a tree from), cancel_session skips the tree-walking step entirely: it partitions the flat list into terminal vs. non-terminal executions and applies the same per-node cancel/flag actions cancel_task already applies to its to_cancel and terminal_ancestors buckets — just without needing a target node to seed the walk.

User Stories

1. As a backend developer integrating by-framework, I want to cancel every active task in a session with one call, so that I don't have to look up or track a message_id just to abandon a session.
2. As a backend developer, I want cancel_session to cancel every node in a multi-agent call chain (A→B→C→...) currently active under that session, so that abandoning a conversation turn doesn't leave orphaned sub-agent tasks running.
3. As a backend developer, I want executions that are already COMPLETED/FAILED/CANCELLED to be left alone (not re-marked as cancelling) but still flagged so that a still-running descendant can't wake them back up via callback, so that stopping a session doesn't resurrect finished work.
4. As a backend developer, I want tasks that are queued but not yet claimed by any worker (no worker_id assigned yet) to still be cancelled, so that in-flight-but-unstarted work doesn't slip through and start after I asked for the session to stop.
5. As a backend developer, I want a clear response telling me how many executions were actively cancelled vs. how many were already finished, so that I can log/report the outcome without re-querying the registry.
6. As a backend developer, I want cancel_session to behave predictably when the session has no executions at all (never existed, or fully expired), so that I can distinguish "nothing to cancel" from a real failure.
7. As a backend developer, I want cancel_session to raise the same error cancel_task raises when the client was constructed without a WorkerRegistry, so that the failure mode is consistent across the client's cancel APIs.
8. As a backend developer using ByaiGatewayClient, I want cancel_session to work without any extra setup, so that I get the same session-cancel capability regardless of which client subclass I use.
9. As a backend developer, I want to pass a reason and cancel_mode through to every CancelTaskCommand dispatched by cancel_session, so that downstream plugins/traces/logs see a consistent cancellation reason across the whole session, just like they would with an individual cancel_task call.
10. As a backend developer, I want to know that cancel_session is intentionally broader than cancel_task — it cancels all active executions in the session, not just one task's sub-tree — so that I don't reach for it when I actually mean to cancel one conversation turn while others continue in the background.
11. As a maintainer, I want cancel_session to reuse the existing CancelTaskCommand protocol and worker-side handle_cancel_task handling unchanged, so that this feature requires no worker-side or wire-protocol changes.

Implementation Decisions

• New method: GatewayClient.cancel_session(self, session_id: str, reason: str = "", requested_by: str = "client", cancel_mode: str = CancelMode.GRACEFUL) -> CancelSessionResponse, defined alongside cancel_task in by_framework/client/client.py.
• Registry guard: same precondition as cancel_task — if self.registry is None, raise the same error cancel_task raises today (ValueError("GatewayClient requires a WorkerRegistry to cancel tasks")) so the two cancel APIs fail the same way under the same misconfiguration.
• No tree walk: fetch all_executions = await self.registry.get_all_session_executions(session_id) once. Do not reconstruct a parent_message_id tree and do not BFS — the full session's executions are already flat in this result, which is everything cancel_task's BFS was assembling piece by piece.
• Partitioning: split all_executions into terminal (status in {"COMPLETED", "FAILED", "CANCELLED"}) and non_terminal using the same terminal-state set cancel_task uses.
• Non-terminal executions: for each, call registry.mark_execution_cancelling(execution_id, session_id, reason); if the execution record has a non-empty worker_id, build a CancelTaskCommand (same fields cancel_task populates: target_message_id, target_execution_id, target_worker_id, reason, requested_by, cancel_mode) and xadd it to RedisKeys.worker_ctrl_stream(worker_id). If worker_id is empty (task queued but not yet claimed), skip sending the command — rely on the existing runner-side cancel_requested pre-check that already short-circuits unclaimed tasks before they start processing.
• Terminal executions: for each, call registry.mark_cancel_requested(execution_id, session_id, reason) only — do not change their status. This is the same protective flagging cancel_task applies to "terminal ancestors" of its BFS target, applied here to every terminal execution in the session (there's no single target node to compute ancestry from, so it's unconditional rather than ancestry-scoped).
• New response type: add CancelSessionResponse (frozen dataclass) to by_framework/core/protocol/responses.py:

  @dataclass(frozen=True)
  class CancelSessionResponse:
      success: bool
      session_id: str
      status: str
      timestamp: int
      cancelled_count: int = 0
      already_finished_count: int = 0
      error: str = ""

  Reuse ExecutionStatus.NOT_FOUND, ExecutionStatus.ALREADY_FINISHED, and ExecutionStatus.CANCEL_REQUESTED for status — no new status constants needed.
• Result semantics:
  • all_executions empty → success=False, status=NOT_FOUND, cancelled_count=0, already_finished_count=0.
  • all_executions non-empty but non_terminal empty (everything already terminal) → success=False, status=ALREADY_FINISHED, cancelled_count=0, already_finished_count=len(terminal).
  • Otherwise → success=True, status=CANCEL_REQUESTED, cancelled_count=len(non_terminal), already_finished_count=len(terminal).
• ByaiGatewayClient: no override needed. It only overrides send_message today; cancel_task is inherited as-is and cancel_session will be too.
• No worker-side changes: handle_cancel_task and the CancelTaskCommand wire format are untouched — cancel_session is purely a client-side/registry-read change that dispatches the existing command.

Testing Decisions

• Tests only observe external behavior: given a mocked registry.get_all_session_executions return value, assert on (a) the CancelSessionResponse returned, (b) which registry.mark_execution_cancelling / registry.mark_cancel_requested calls were made and with what arguments, and (c) which redis.xadd calls were made and to which worker_id streams — not on any internal partitioning helper.
• Location: tests/client/test_client.py, next to the existing cancel_task tests. Same seam, same fixtures (AsyncMock() for registry and redis_client, constructing a bare GatewayClient(redis_client=..., registry=...)).
• Prior art to follow directly: test_client_cascading_cancel_a_b_c (mixed terminal/non-terminal partitioning and per-worker xadd assertions), test_client_cancel_task_marks_registry_when_worker_not_assigned (empty worker_id → registry marked but no xadd), test_client_cancel_task_returns_already_finished (all-terminal → success=False).
• Cases to cover:
  1. Mixed session (some terminal, some non-terminal, e.g. an A→B→C chain plus one unrelated completed root) → correct split, correct xadd targets, correct counts.
  2. All executions non-terminal → cancelled_count == len(all_executions), one xadd per execution with a worker_id.
  3. All executions terminal → success=False, status=ALREADY_FINISHED, zero xadd calls, mark_cancel_requested called once per execution.
  4. Empty session (get_all_session_executions returns []) → success=False, status=NOT_FOUND.
  5. Non-terminal execution with empty worker_id (queued, unclaimed) → mark_execution_cancelling called, no xadd for that node.
  6. registry is None → raises the same error cancel_task raises.
  7. ByaiGatewayClient instance calling cancel_session behaves identically to GatewayClient (confirms no override is required).

Out of Scope

• Porting cancel_session to the TypeScript (by-framework-ts) or Java (by-framework-java) SDKs — those are separate repos with their own issue trackers and would need their own issues per the workspace's cross-SDK convention.
• Fixing the pre-existing gap where cancel_mode (GRACEFUL vs FORCE) is threaded through CancelTaskCommand but never actually changes worker-side cancellation behavior (handle_cancel_task always does an immediate task.cancel() regardless of mode). cancel_session inherits this exact same pre-existing limitation from cancel_task; not this PRD's job to fix it.
• Locking or atomicity guarantees against tasks dispatched concurrently with cancel_session (e.g. a node calling call_agent to spawn a new child in the window between the registry snapshot and the cancel commands being sent). cancel_task's BFS has the identical limitation today; this PRD does not change that.
• Any worker-side or wire-protocol changes — CancelTaskCommand and handle_cancel_task are reused unmodified.
• A REST/HTTP-facing endpoint for session cancellation — this PRD only covers the GatewayClient Python API surface.

Further Notes

• cancel_session is deliberately coarser-grained than cancel_task: it will cancel every active execution in a session, including unrelated concurrent top-level tasks that share the same session_id but aren't part of the same conversation turn (a session can have more than one root execution over its lifetime). Callers that only want to abandon one turn while leaving other in-flight work alone should keep using cancel_task with that turn's root message_id.
• This was scoped out of a conversation about how cancel_task's BFS cascade already works for multi-hop call_agent chains (A→B→C→D→E); the observation that motivated this PRD is that get_all_session_executions already returns the flattened data the BFS reconstructs, so a session-wide cancel is strictly simpler to implement than the existing per-task cascade, not harder.