[Enhancement]: Implement worker invocation cancellation across Rust and Python SDKs

[bbednarski9]

Affected area

• Rust core runtime
• Plugins
• Python worker SDK
• Documentation and examples

Problem or opportunity

The grpc-v1 worker protocol reserves an explicit cancellation operation and assigns every callback invocation an invocation_id:

• PluginWorker.CancelInvocation
• InvokeRequest.invocation_id and CancelInvocationRequest

Cancellation is not implemented end to end today:

• The Rust worker SDK returns accepted = false.
• The Python worker SDK introduced in PR feat(python): add nemo-relay-plugin worker SDK package #310 likewise returns accepted = false.
• The Rust host creates invocation IDs but does not currently send CancelInvocation. Unary worker calls instead pass through a fixed 30-second RPC timeout.

Transport-level cancellation may drop an async handler future in some cases, but the worker contract does not currently track callback tasks by invocation_id or guarantee cancellation behavior. In particular, synchronous Python callbacks cannot be preempted by cancelling an asyncio task. Long-running callbacks may therefore continue work or side effects after the host caller has timed out or stopped consuming a stream, while host-owned continuations and invocation scope mappings are being removed.

This limitation was identified and accepted as a follow-up in PR #310's review thread; see also the author response.

Proposed enhancement

Define and implement invocation cancellation across the Relay worker host, Rust worker SDK, and Python worker SDK:

1. Track active unary and streaming invocations by invocation_id in each worker SDK.
2. Have the host send CancelInvocation when a managed call is cancelled, a worker invocation times out, or a worker-backed stream is abandoned when explicit cancellation is required.
3. Cancel the corresponding Rust future/task or Python asyncio task cooperatively and return a meaningful WorkerAck.
4. Make repeated cancellation and cancellation racing with normal completion safe and deterministic.
5. Coordinate continuation and scope-stack cleanup so cancellation cannot leave stale host state or allow continuation use after the invocation ends.
6. Represent cancellation distinctly from generic worker failure in host-facing errors and stream termination.
7. Decide whether cancellation is mandatory for grpc-v1 or negotiated as a worker capability so older SDKs remain compatible.

Synchronous callbacks that cannot be safely preempted should have explicit documented behavior. The contract should not claim that a cancellation acknowledgment guarantees interruption of arbitrary blocking user code.

Runtime contract and binding impact

• Rust core host: propagate caller cancellation and timeout into the worker protocol; own final continuation and scope cleanup.
• Rust worker SDK: maintain active invocation handles and cancel asynchronous handler work by invocation ID.
• Python worker SDK: maintain active asyncio.Task instances and cancel asynchronous handlers by invocation ID; document the limitation for synchronous blocking callbacks.
• Other worker SDKs: future worker implementations should follow the same grpc-v1 cancellation semantics or advertise that the capability is unsupported.
• Native/in-process plugins: not directly affected because they do not cross the worker RPC boundary, although managed-call cancellation semantics should remain consistent where possible.

Alternatives considered

• Rely only on the 30-second host timeout: this bounds how long the host waits for unary RPCs but does not define worker-side cancellation, synchronous callback behavior, or cleanup semantics.
• Rely only on gRPC transport cancellation: transport cancellation is useful but does not replace an explicit cross-language contract keyed by invocation_id, especially for blocking callbacks and coordinated host state cleanup.
• Document cancellation as unsupported: acceptable as a short-term beta limitation, but incomplete for long-running execution intercepts and worker processes.
• Terminate the entire worker process: provides a coarse stop mechanism but disrupts unrelated concurrent invocations and loses plugin state.

Acceptance criteria

• The host invokes CancelInvocation for the defined cancellation and timeout paths.
• Rust and Python worker SDKs track active unary and streaming invocations by invocation_id.
• Active asynchronous callbacks can be cancelled, and workers return an accurate acknowledgment.
• Cancellation of unknown, completed, or already-cancelled invocation IDs has documented, idempotent behavior.
• Continuations, invocation scope mappings, and stream resources are released exactly once across completion/cancellation races.
• Rust host and SDK tests cover timeout, caller cancellation, streaming cancellation, repeated cancellation, and completion races.
• Python SDK tests cover asynchronous callback cancellation and the documented behavior of synchronous callbacks.
• At least one end-to-end worker integration test verifies host-to-worker cancellation.
• grpc-v1 compatibility or capability-negotiation behavior is defined for workers that do not support cancellation.
• Worker authoring documentation explains cancellation guarantees and limitations.

References

• Relates to feat(python): add nemo-relay-plugin worker SDK package #310
• PR #310 cancellation discussion
• Rust worker SDK cancellation stub