fix(odin): Fixed cancelling an in-flight action

[vikramlc-cognite]

Summary

Fixes cancelling an in-flight start_task/custom action, plus three follow-up correctness fixes surfaced during review: status/message accuracy when oversized metadata coincides with cancellation, and a token-registration leak on early failure.

Type of change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that changes existing behavior)
• Refactor (no functional change)
• Documentation update
• Chore / tooling / CI

What changed

Core fix — cancel an in-flight action instead of re-dispatching it:

• base.py: added _running_action_tokens (keyed by Action.external_id, mirroring _running_task_tokens), populated by _handle_start_task_action/_handle_custom_action.
• base.py: _dispatch_single_action now checks action.status up front — a cancel_pending re-delivery (Odin's signal that a user cancelled an already-dispatched action) cancels the tracked token by external_id and returns, instead of re-running the handler.
• base.py: _handle_custom_action reports ActionStatus.canceled (not succeeded) if the action's token was cancelled before the target returned.
• actions.py: ActionContext gains a required cancellation_token: CancellationToken, so custom actions can cooperatively check ctx.cancellation_token.is_cancelled/.wait(...).

Follow-up fixes (from review):

• base.py: behavioral change — the oversized-metadata branch in _handle_custom_action used to always report ActionStatus.failed, regardless of whether the action itself succeeded. It now reports the action's real outcome (succeeded/canceled) instead — metadata truncation is reflected only in the message and dropped fields, not the status. Any consumer (dashboards, alerting, automation) relying on "oversized metadata → failed" as a signal will see a different status for these actions after this PR. Message wording also no longer claims "completed successfully" when the action was actually cancelled.
• base.py: token registration (_running_action_tokens, and _running_task_tokens for start_task) now happens inside the same try/finally that cleans it up, in both _handle_start_task_action and _handle_custom_action — previously a failure between registration and the actual work (e.g. constructing the "running" ActionUpdate) would leave the task/action stuck "running" forever with no error ever reported.
• base.py: extracted the repeated "release token if still owned" lock/get/pop idiom (5 call sites) into a shared _release_if_owned helper.

Tests: updated all ActionContext(...) call sites (test_actions.py, test_log_upload_action.py) to pass a token; added 8 new tests to test_action_dispatch.py covering cancel-pending re-delivery (custom and start_task actions), cooperative cancellation reporting, oversized-metadata status/wording under cancellation, and registration cleanup on early failure.

Why it changed

Odin propagates a cancel of a running action by flipping it to cancel_pending and re-sending the same external_id on the next checkin. The SDK never read Action.status, so this was treated as a fresh dispatch — re-running custom action side effects a second time, or spuriously failing an already-running start_task action. The review follow-ups close two adjacent gaps in the same area: status/message accuracy was still wrong in one corner case (oversized metadata + cancellation), and the new token-tracking dicts had the same "register before try" fragility as existing code nearby.

What to focus on during review

• Behavioral change: oversized-metadata custom actions now report succeeded/canceled instead of always failed — see "Follow-up fixes" above. Flag if any downstream consumer depends on the old always-failed behavior for this case.
• ActionContext.cancellation_token is a new required constructor argument — any external code constructing ActionContext directly needs to pass one.
• Cancellation remains cooperative — a custom action that never checks the token still runs to completion; this PR only prevents double-execution and misreported status, not forced interruption.
• _release_if_owned is a pure mechanical refactor (verified via test diff against pre-refactor code) — no behavior change intended.

Test evidence

• pytest tests/test_unstable/test_action_dispatch.py tests/test_unstable/test_actions.py tests/test_unstable/test_log_upload_action.py -q → 96 passed.
• pytest tests/test_unstable/ -q → 248 passed (48 pre-existing, unrelated errors from missing local env vars; 0 new failures).
• mypy / ruff check / ruff format → clean.

Risks and unknowns

• The oversized-metadata status change (see above) is the one behavior change with real external-visibility risk in this PR — worth confirming no dashboard/alert depends on the old failed-always behavior before merge.
• Custom actions that don't check the cancellation token get no benefit beyond "no longer double-executed" — inherent cooperative-cancellation limitation, not addressed here.
• stop_task action behavior is unchanged (completes near-instantly; cancel-pending race window is negligible there).

Rollout and rollback

No flags/migrations. ActionContext's new required parameter is the only signature change; plain code revert if needed.

Checklist

• Self-reviewed the diff
• Tests added or updated
• Docs updated (or N/A) — N/A, internal SDK behavior; documented via docstrings
• No secrets, credentials, or PII committed
• Breaking changes called out above — ActionContext now requires cancellation_token; oversized-metadata status behavior change called out under "What to focus on during review"

[vikramlc-cognite]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for cancelling long-running custom actions and start-task actions via Odin's cancel_pending status. It integrates CancellationToken into ActionContext and tracks active actions using a new _running_action_tokens registry. When a cancel_pending action is dispatched, the corresponding token is cancelled instead of re-running the action. Comprehensive tests have been added to verify the cancellation behavior. I have no feedback to provide on these changes.

[gemini-code-assist]

Code Review

This pull request introduces cancellation support for custom actions and start-task actions by integrating CancellationToken into ActionContext and tracking in-flight actions using their external_id. When a cancel_pending action is received, the corresponding token is cancelled to stop the action early. The review feedback highlights a potential race condition where a cancel_pending action might be processed before the pending action's thread has registered its token. To prevent this, the reviewer suggested introducing a set to track cancelled action IDs so they can be cancelled immediately upon registration.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.09%. Comparing base (19f4209) to head (0a9b5c0).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #572      +/-   ##
==========================================
+ Coverage   83.97%   84.09%   +0.12%     
==========================================
  Files          46       46              
  Lines        4686     4709      +23     
==========================================
+ Hits         3935     3960      +25     
+ Misses        751      749       -2     

Files with missing lines	Coverage Δ
cognite/extractorutils/unstable/core/actions.py	100.00% <100.00%> (ø)
cognite/extractorutils/unstable/core/base.py	86.90% <100.00%> (+1.40%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[vikramlc-cognite]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces cancellation support for custom actions and start-task actions by integrating cancellation tokens into ActionContext and tracking active actions in the base extractor. If a cancel_pending action is received, the corresponding token is cancelled, and the action's final status is updated to ActionStatus.canceled. The review feedback highlights two potential issues where exceptions raised during cooperative aborts (for both tasks and custom actions) are caught and incorrectly reported as ActionStatus.failed instead of ActionStatus.canceled. The reviewer recommends checking the cancellation token's status in the exception handlers to ensure accurate status reporting.

[Jacob-Eliat-Eliat]

This might be very theoretical since there's nothing likely to throw in between, but ideally the assignments to _running_task_tokens and _running_action_tokens should happen inside the same try that cleans them up in finally.

Right now the only thing preventing a stuck task/action entry is that nothing between registration and the try currently throws which is not very future proof.

It's in my opinion a bit debatable: one one hand best practice is to have assignment to the map inside the try-finally because if finally is not run then there is a mem-leak, on the other hand, it's bad practice to wrap too much non-throwing code in a try block. Looking forwards to mentor review on this.

[vikramlc-cognite]

Agreed, I traced through it and this is worth fixing, for a slightly sharper reason than the leaked-entry framing: if something in that gap does throw, it doesn't just leak a dict entry, it leaves the task permanently stuck running until the whole extractor restarts.

[Jacob-Eliat-Eliat]

🐴 next step: will need to do a mentor review, flagging anything that I could have missed before it can be merged.

[nithinb]

🦄 LGTM