feat(runtime-core): add optional signal to AgentBlockContext for cancelling agent runs by tkislan · Pull Request #445 · deepnote/deepnote

tkislan · 2026-08-04T19:43:16Z

Adds an optional signal?: AbortSignal to AgentBlockContext so a host with a Stop control — the VS Code extension's cell Stop button — can halt an agent run in progress. Additive and backwards compatible: the field is optional, no in-repo caller passes it, and no exported type is narrowed.

Why a host can't do this today

Host attempt	What actually happens
`onAgentEvent` throws	`fullStream` is a `tee()` branch, so cancelling the consumer never reaches the fetch. The request completes and is billed.
Tool callback throws	Becomes a `tool-error` stream part fed back to the model, which keeps looping.
Return a sentinel string	Suppresses side effects, but the model still runs to `stepCountIs(maxTurns)`.

In the extension that means: after Stop, agent-generated code can still be submitted to the kernel, new cells can still be inserted, and spend continues for up to ~10 turns.

What the change does

Five edits in packages/runtime-core/src/agent-handler.ts (+18 −3):

signal?: AbortSignal on AgentBlockContext
Entry guard, before MCP subprocess creation, so a pre-aborted run spawns nothing
A guard inside both host tool execute wrappers
abortSignal: context.signal forwarded to agent.stream()
A guard after await streamResult.text

Item 3 is load-bearing and not obvious. Measured A/B with three tool calls in one step, aborting from inside tool 1:

no guard:   ["enter-1","SIDE-EFFECT-1","enter-2","SIDE-EFFECT-2","enter-3","SIDE-EFFECT-3"]
with guard: ["enter-1","enter-2","enter-3"]

The SDK dispatches every remaining tool call of the step after the signal aborts. Forwarding abortSignal alone does not stop it — without the in-execute guard a cancelled run keeps adding and executing cells, one per queued call.

Item 5 matters because a late abort leaves .text resolving with the previous step's text, which would otherwise be returned as { finalOutput: 'stale text' } — a success.

Cancellation always rejects with signal.reason rather than resolving. Resolving would be fail-open: the engine marks the block success: true and continues to the next block, which is exactly what Stop must not do.

🤖 Generated with Claude Code

https://claude.ai/code/session_01RwFZ8m75wzt2MgYTJope9j

Summary by CodeRabbit

New Features
- Added cancellation support for agent block execution.
- Requests can now stop before or during tool and model operations.
Bug Fixes
- Prevented unnecessary client creation and continued tool callbacks after cancellation.
- Improved cleanup of resources when execution is interrupted.
Tests
- Added coverage for cancellation before execution, during tools, and during final model responses.

…elling agent runs A host with a Stop control had no way to halt an agent run in progress. Throwing from onAgentEvent leaves the model request billed, throwing from a tool callback is fed back to the model as context, and returning a sentinel still burns turns up to maxTurns. In the VS Code extension that means Stop still submits generated code to the kernel and still inserts cells. Passing abortSignal to agent.stream() alone is not enough: the SDK hands over every remaining tool call of a step after the signal aborts, so each callback needs its own guard or a cancelled run keeps mutating the notebook. A late abort also leaves streamResult.text resolving with the previous step's text, which would be returned as a successful result. Cancellation always rejects with signal.reason rather than resolving, so the execution engine records a failed block and breaks the run instead of counting the block executed and continuing. Additive and backwards compatible: the field is optional, no in-repo caller passes it, and no exported type is narrowed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RwFZ8m75wzt2MgYTJope9j

codecov · 2026-08-04T19:45:47Z

Codecov Report

❌ Patch coverage is 94.73684% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 88.11%. Comparing base (f9f89d5) to head (a3a3096).

Files with missing lines	Patch %	Lines
packages/runtime-core/src/agent-handler.ts	94.73%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #445      +/-   ##
==========================================
+ Coverage   87.78%   88.11%   +0.33%     
==========================================
  Files         187      187              
  Lines        9927     9936       +9     
  Branches     2768     2847      +79     
==========================================
+ Hits         8714     8755      +41     
+ Misses       1212     1180      -32     
  Partials        1        1

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Added the MCPClient type to the import statement for better type safety. Refactored the execution logic to initialize mcpClients as an empty array and wrapped the Promise.all call in a try-catch block for improved error handling. This change ensures that the agent block can handle tool execution more robustly, maintaining the existing functionality while enhancing code clarity and maintainability.

The abort tests drove the run from outside and then spun the event loop waiting for it to settle, which needed a gated model stream, a deferred helper, a promise-state observer and a fixed-round drain loop. Aborting from inside a callback the SDK already invokes — a host tool callback, or doStream itself — orders the abort deterministically against the code under test, so plain `await expect(...).rejects` works and all four helpers go away. Mutation-checked: removing any one of the four abort guards (entry, agent.stream abortSignal, per-tool, post-stream) still fails exactly one test, same as before. Drops the MCP limitation test: no mutation killed it because it asserted ai core stays broken, and it would go red on an ai upgrade that fixed it. That caveat now lives on the AgentBlockContext.signal docstring, where a caller will see it. Also replaces the manual OPENAI_BASE_URL save/restore with vi.stubEnv, matching the convention already used elsewhere in the repo. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RwFZ8m75wzt2MgYTJope9j

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RwFZ8m75wzt2MgYTJope9j

tkislan · 2026-08-06T07:58:26Z

@coderabbitai full review

coderabbitai · 2026-08-06T07:58:33Z

✅ Action performed

Full review finished.

coderabbitai · 2026-08-06T07:58:47Z

📝 Walkthrough

Walkthrough

AgentBlockContext now accepts an optional AbortSignal. executeAgentBlock checks cancellation before resource creation, during notebook tool execution, during agent streaming, and before returning output. MCP clients are tracked inside resource cleanup. New tests cover successful execution and multiple cancellation points.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant executeAgentBlock
  participant MCPClient
  participant OpenAI
  participant NotebookTool
  Caller->>executeAgentBlock: Provide AgentBlockContext.signal
  executeAgentBlock->>executeAgentBlock: Check abort state
  executeAgentBlock->>MCPClient: Create and track clients
  executeAgentBlock->>OpenAI: Start agent stream with signal
  OpenAI-->>executeAgentBlock: Request notebook tool
  executeAgentBlock->>NotebookTool: Invoke when not aborted
  NotebookTool-->>executeAgentBlock: Return tool result
  Caller-->>executeAgentBlock: Abort signal
  executeAgentBlock->>executeAgentBlock: Reject aborted execution
  executeAgentBlock->>MCPClient: Clean up tracked clients

Possibly related PRs

deepnote/deepnote#342: Both changes modify AgentBlockContext and executeAgentBlock in agent-handler.ts.

Suggested reviewers: jamesbhobbs

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Updates Docs	⚠️ Warning	PR commits `23edb83`–22aed5f change only runtime-core source and tests; no Agent documentation update exists, and deepnote-internal is not available locally.	Update deepnote/deepnote Agent documentation with signal cancellation semantics, then update the deepnote-internal landing-page roadmap and verify it separately.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: adding an optional abort signal to AgentBlockContext for cancelling agent runs.

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@packages/runtime-core/src/agent-handler.ts`:
- Around line 175-211: In the MCP initialization flow, add
context.signal?.throwIfAborted() immediately after the awaited createMCPClient
Promise.all and again after the awaited mcpClients.map(client => client.tools())
discovery, before continuing. Add a regression test that keeps createMCPClient
pending, aborts the signal, resolves creation, and verifies client.tools() is
never invoked.
- Around line 175-186: Update the MCP client initialization around
createMCPClient to collect individual outcomes instead of allowing Promise.all
to discard fulfilled clients when another rejects. Assign all fulfilled clients
to mcpClients, then rethrow the startup error so existing cleanup closes partial
clients; add a test covering one successful and one rejected client and
verifying the successful client is retained and cleaned up.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: f44a2559-ff67-473a-ae2d-e63bb2a31042

📥 Commits

Reviewing files that changed from the base of the PR and between 0e58ac8 and 22aed5f.

📒 Files selected for processing (2)

packages/runtime-core/src/agent-handler.ts
packages/runtime-core/src/execute-agent-block.test.ts

m1so

nice 👍

tkislan and others added 3 commits August 5, 2026 07:15

refactor(runtime-core): drop the signal field comment

22aed5f

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RwFZ8m75wzt2MgYTJope9j

coderabbitai Bot requested changes Aug 6, 2026

View reviewed changes

Comment thread packages/runtime-core/src/agent-handler.ts

Comment thread packages/runtime-core/src/agent-handler.ts

This was referenced Aug 6, 2026

Close partially initialized MCP clients when startup fails #449

Open

Avoid MCP tool discovery after agent cancellation #450

Open

tkislan marked this pull request as ready for review August 6, 2026 13:34

tkislan requested a review from a team as a code owner August 6, 2026 13:34

Merge branch 'main' into feat/agent-block-abort-signal

a3a3096

coderabbitai Bot approved these changes Aug 6, 2026

View reviewed changes

m1so approved these changes Aug 6, 2026

View reviewed changes

tkislan merged commit b48db84 into main Aug 6, 2026
21 checks passed

tkislan deleted the feat/agent-block-abort-signal branch August 6, 2026 16:19

tkislan mentioned this pull request Aug 6, 2026

chore(release): bump package versions #444

Draft

coderabbitai Bot mentioned this pull request Aug 10, 2026

fix(runtime-core): close partially initialized MCP clients when startup fails #452

Open

Navigation Menu

Uh oh!

feat(runtime-core): add optional signal to AgentBlockContext for cancelling agent runs - #445

feat(runtime-core): add optional signal to AgentBlockContext for cancelling agent runs#445
tkislan merged 5 commits into
mainfrom
feat/agent-block-abort-signal

tkislan commented Aug 4, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

codecov Bot commented Aug 4, 2026 •

edited

Loading

Uh oh!

tkislan commented Aug 6, 2026

Uh oh!

coderabbitai Bot commented Aug 6, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Aug 6, 2026 •

edited

Loading

Walkthrough

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

m1so left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

tkislan commented Aug 4, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why a host can't do this today

What the change does

Summary by CodeRabbit

Uh oh!

codecov Bot commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

tkislan commented Aug 6, 2026

Uh oh!

coderabbitai Bot commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

m1so left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tkislan commented Aug 4, 2026 •

edited by coderabbitai Bot

Loading

codecov Bot commented Aug 4, 2026 •

edited

Loading

coderabbitai Bot commented Aug 6, 2026 •

edited

Loading

coderabbitai Bot commented Aug 6, 2026 •

edited

Loading