feat(agent-block): Add support for Agent block#358
Draft
tkislan wants to merge 9 commits intotk/deepnote-runtime-corefrom
Draft
feat(agent-block): Add support for Agent block#358tkislan wants to merge 9 commits intotk/deepnote-runtime-corefrom
tkislan wants to merge 9 commits intotk/deepnote-runtime-corefrom
Conversation
…ecution functions and improve ephemeral cell handling
…handling of ephemeral cells in serialization and decoration
Contributor
📝 WalkthroughWalkthroughAdds end-to-end agent cell support for Deepnote notebooks: agent cells are detected and executed immediately via a new agent execution handler that serializes notebook context, streams agent events into cell outputs, and manages ephemeral cells (insert, execute, cleanup). Kernel execution is preserved for non-agent cells and only starts when kernel-targeted cells remain. Introduces agent block conversion, status-bar and decoration providers for agent/ephemeral cells, secret store helpers and commands for an OpenAI API key, excludes ephemeral cells from serialization, extends the renderer for ephemeral markdown, and adds comprehensive unit tests. Sequence Diagram(s)sequenceDiagram
participant Controller as NotebookController
participant Filter as ExecutionFilter
participant AgentExec as AgentCellExecutor
participant Serializer as ContextSerializer
participant EphemeralMgr as EphemeralManager
participant Kernel as KernelExecutor
Controller->>Filter: executeCells(request)
Filter->>Filter: split cells -> agentCells + kernelCells
alt agentCells exist
Filter->>AgentExec: executeAgentCell(agentCell)
AgentExec->>Serializer: serializeNotebookContext(notebook)
AgentExec->>AgentExec: clear outputs, start execution
AgentExec->>EphemeralMgr: removeEphemeralCellsForAgent(agentBlockId)
AgentExec->>AgentExec: executeAgentBlock (stream events)
AgentExec->>AgentExec: stream text/tool deltas -> update outputs
AgentExec->>EphemeralMgr: insertEphemeralCell(code/markdown)
EphemeralMgr->>Kernel: executeEphemeralCell(ephemeralCell)
Kernel-->>EphemeralMgr: execution result
EphemeralMgr-->>AgentExec: propagate outputs
AgentExec->>AgentExec: end execution (success/error)
end
alt kernelCells exist
Filter->>Kernel: start kernel and execute kernelCells
Kernel-->>Controller: kernel execution complete
else no kernelCells
Filter-->>Controller: return (no kernel started)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 📝 Coding Plan
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/notebooks/deepnote/agentCellExecutionHandler.ts`:
- Around line 110-113: The code directly reads OPENAI_API_KEY from process.env
in agentCellExecutionHandler.ts (openAiToken = process.env.OPENAI_API_KEY) which
is unsafe for production; replace this direct env access with a secure secret
retrieval call (e.g., a new getOpenAiApiKey() that fetches from your secret
manager/credentials vault or from an injected secure config) and update callers
to inject the key instead of relying on process.env; ensure the secret is never
logged or included in error messages and keep the existing null-check/throw
behavior but reference the secure getter (getOpenAiApiKey) or injected parameter
in place of process.env.OPENAI_API_KEY.
- Around line 274-278: The success check in the return object of
agentCellExecutionHandler is too permissive—replace the current expression
`cell.executionSummary?.success !== false` with an explicit true check like
`cell.executionSummary?.success === true` (so only an explicit success is
reported; undefined/in-progress will not be treated as success); update the
return here (where `success`, `outputs:
cell.outputs.map(translateCellDisplayOutput)`, and `executionCount:
cell.executionSummary?.executionOrder ?? null` are constructed) to use that
strict equality.
In `@src/notebooks/deepnote/agentCellStatusBarProvider.ts`:
- Around line 69-71: The dispose() method currently uses an expression-bodied
arrow in this.disposables.forEach((d) => d.dispose()) which triggers the Biome
callback-return lint; change the callback to a block body or replace the forEach
with a for...of loop so the disposables are disposed without returning a
value—e.g., update dispose() to iterate over the disposables array and call
dispose() inside a statement block (reference: dispose method and disposables
property).
- Around line 142-149: getMaxIterations currently only enforces a lower bound;
add an upper-bound check so the returned value is an integer between
MIN_ITERATIONS and MAX_ITERATIONS (e.g., require value <= MAX_ITERATIONS). In
setMaxIterations replace permissive parseInt usage with strict integer
validation (use a full-match regex like /^\d+$/) and then parse with Number() so
inputs like "5.5" or "10abc" are rejected; after parsing ensure the numeric
value is an integer and within MIN_ITERATIONS..MAX_ITERATIONS before accepting
or falling back to DEFAULT_MAX_ITERATIONS. Update both occurrences in
setMaxIterations that currently call parseInt to use this strict validation and
range check, and reference the getMaxIterations and setMaxIterations functions
when making the change.
In `@src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts`:
- Around line 1-5: Reorder the imports so third-party modules are grouped
together and local imports come after: move the dedent import to be alongside
the other external imports (DeepnoteBlock, chai's assert, and vscode's
NotebookCellData/NotebookCellKind) and place the local AgentBlockConverter
import ('./agentBlockConverter') after that group; ensure the symbols
DeepnoteBlock, assert, NotebookCellData, NotebookCellKind, and dedent remain
imported and only the order changes to comply with the "third-party then local"
guideline.
🪄 Autofix (Beta)
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: ASSERTIVE
Plan: Pro
Run ID: 45a324a0-84a7-463d-903d-d15c32e2b30d
📒 Files selected for processing (16)
src/notebooks/controllers/vscodeNotebookController.tssrc/notebooks/deepnote/agentCellExecutionHandler.tssrc/notebooks/deepnote/agentCellExecutionHandler.unit.test.tssrc/notebooks/deepnote/agentCellStatusBarProvider.tssrc/notebooks/deepnote/agentCellStatusBarProvider.unit.test.tssrc/notebooks/deepnote/converters/agentBlockConverter.tssrc/notebooks/deepnote/converters/agentBlockConverter.unit.test.tssrc/notebooks/deepnote/deepnoteDataConverter.tssrc/notebooks/deepnote/deepnoteKernelAutoSelector.node.tssrc/notebooks/deepnote/deepnoteTestHelpers.tssrc/notebooks/deepnote/ephemeralCellDecorationProvider.tssrc/notebooks/deepnote/ephemeralCellStatusBarProvider.tssrc/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.tssrc/notebooks/serviceRegistry.node.tssrc/notebooks/serviceRegistry.web.tssrc/renderers/client/markdown.ts
…k/deepnote-agent-block
- Introduced `getOpenAiApiKey` function to retrieve the OpenAI API key from configuration, improving error handling when the key is not set. - Updated `executeAgentCell` and `executeEphemeralCell` functions to utilize the new API key retrieval method and handle cancellation tokens. - Enhanced `AgentCellStatusBarProvider` to validate max iterations using Zod schema, ensuring robust input handling and defaulting to safe values. - Added unit tests for new functionality and edge cases in both execution handling and status bar provider.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/notebooks/deepnote/agentCellExecutionHandler.ts`:
- Around line 142-147: The onLog callback in agentCellExecutionHandler.ts
contains commented-out accumulation code and a TODO; either remove the dead code
or implement it: add an accumulated string variable in the enclosing scope, make
onLog async (or forward logs to an async helper), append incoming message to
accumulated, then call
execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output)
to update the cell output; if you choose to drop it, delete the commented lines
and the TODO and keep only logger.info('Agent log', message). Reference: onLog
callback, accumulated variable, execution.replaceOutputItems,
NotebookCellOutputItem.text, and output.
- Around line 41-64: serializeNotebookContext instantiates a new
DeepnoteDataConverter on every call which is wasteful if called frequently;
modify serializeNotebookContext to use a shared or injected converter instance
instead of creating one per invocation (e.g., accept a DeepnoteDataConverter
parameter or read from a module-scoped singleton), and update callers to pass or
rely on the shared converter so convertCellToBlock usage inside
serializeNotebookContext reuses the same converter.
In `@src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts`:
- Around line 310-324: The test creates a CancellationTokenSource named
tokenSource and cancels it but never disposes it; update the test for 'returns
success false immediately when token is pre-cancelled' to ensure
tokenSource.dispose() is called after use (e.g., in a finally block or via
afterEach cleanup) so the CancellationTokenSource is properly disposed; locate
the tokenSource variable in this test and add the dispose call around
executeEphemeralCell(tokenSource.token) to clean up resources.
In `@src/notebooks/deepnote/agentCellStatusBarProvider.ts`:
- Line 27: MaxIterationsSchema currently only enforces a minimum via
MIN_ITERATIONS so values >100 slip through; update MaxIterationsSchema to also
enforce an upper bound (e.g., .max(100)) or reference a new constant like
MAX_ITERATIONS = 100 if you prefer a named limit, ensuring you use
z.coerce.number().int().min(MIN_ITERATIONS).max(MAX_ITERATIONS) (or .max(100))
to validate both ends; modify the schema definition where MaxIterationsSchema is
declared and add the MAX_ITERATIONS constant if not already present.
In `@src/notebooks/deepnote/ephemeralCellDecorationProvider.ts`:
- Around line 69-71: The dispose method on EphemeralCellDecorationProvider
currently iterates disposables with this.disposables.forEach((d) =>
d.dispose());—replace the forEach with a for...of loop to align with the pattern
used in AgentCellStatusBarProvider and to ensure proper synchronous disposal and
error handling: iterate over this.disposables using for (const d of
this.disposables) and call d.dispose() inside the loop (referencing the dispose
method and the disposables array to locate the change).
🪄 Autofix (Beta)
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: ASSERTIVE
Plan: Pro
Run ID: 5e9fa43e-8179-4408-8722-29b13fbca570
📒 Files selected for processing (10)
build/esbuild/build.tssrc/notebooks/deepnote/agentCellExecutionHandler.tssrc/notebooks/deepnote/agentCellExecutionHandler.unit.test.tssrc/notebooks/deepnote/agentCellStatusBarProvider.tssrc/notebooks/deepnote/agentCellStatusBarProvider.unit.test.tssrc/notebooks/deepnote/dataConversionUtils.tssrc/notebooks/deepnote/deepnoteSerializer.tssrc/notebooks/deepnote/deepnoteSerializer.unit.test.tssrc/notebooks/deepnote/ephemeralCellDecorationProvider.tssrc/notebooks/deepnote/ephemeralCellStatusBarProvider.ts
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@package.json`:
- Around line 1641-1646: The package.json setting "deepnote.agent.openAiApiKey"
stores the API key in plain settings; remove that configuration entry and
instead read/write the key via VS Code SecretStorage (use
context.secrets.get/set) like the existing apiAccess.ts usage; update the code
that previously read configuration for deepnote.agent.openAiApiKey to check
context.secrets.get("openAiApiKey") and, if missing, prompt the user with an
input dialog (and offer a command to set/clear the secret), and reuse the helper
functions or patterns from apiAccess.ts to centralize secret handling and
prompting.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
- Added commands to set and clear the OpenAI API key, enhancing user interaction. - Introduced a new `deepnoteSecretStore` module for managing secrets, including functions to get, set, and clear the OpenAI API key. - Updated `agentCellExecutionHandler` to utilize the new secret management functions, improving error handling when the API key is not set. - Enhanced unit tests to cover the new secret management functionality and ensure robust error handling.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (1)
src/notebooks/deepnote/agentCellStatusBarProvider.ts (1)
207-224: 🛠️ Refactor suggestion | 🟠 MajorReuse
MaxIterationsSchemafor consistent validation.
parseIntis lenient:"5.5"becomes5,"10abc"becomes10. The existing Zod schema handles this properly and is already used ingetMaxIterations.,
♻️ Suggested fix
validateInput: (value) => { - const num = parseInt(value, 10); - if (isNaN(num) || !Number.isInteger(num)) { - return l10n.t('Please enter a whole number'); - } - if (num < MIN_ITERATIONS || num > MAX_ITERATIONS) { + const result = MaxIterationsSchema.safeParse(value); + if (!result.success) { return l10n.t('Value must be between {0} and {1}', MIN_ITERATIONS, MAX_ITERATIONS); } return undefined; }- const newValue = parseInt(input, 10); + const newValue = MaxIterationsSchema.parse(input); if (newValue === currentValue) {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/notebooks/deepnote/agentCellStatusBarProvider.ts` around lines 207 - 224, The validateInput logic should reuse the existing MaxIterationsSchema instead of using parseInt; replace the parseInt/isNaN checks in validateInput with MaxIterationsSchema.safeParse(input) (or parse and catch) and return l10n.t(...) on failure, ensuring the schema enforces integer-only and range constraints consistent with MIN_ITERATIONS and MAX_ITERATIONS; after the prompt returns, set newValue from the validated schema result (the parsed numeric value) rather than calling parseInt again; refer to validateInput, MaxIterationsSchema, getMaxIterations, MIN_ITERATIONS, MAX_ITERATIONS, and newValue when making these changes.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/notebooks/deepnote/agentCellExecutionHandler.ts`:
- Around line 240-245: The code currently assumes workspace.applyEdit(edit)
succeeded and returns insertIndex blindly; change it to check the boolean result
of await workspace.applyEdit(edit) and verify the notebook now contains the
inserted cell (e.g. notebook.cellCount > insertIndex or try
notebook.cellAt(insertIndex) exists). If applyEdit returns false or the
verification fails, throw an Error (or return a sentinel/failure value as per
project convention) instead of returning insertIndex so callers won't operate on
an invalid index; use the same local symbols edit, insertIndex, notebook,
WorkspaceEdit, NotebookEdit.insertCells and workspace.applyEdit to locate and
implement the checks.
- Around line 136-138: The handler onAgentEvent currently logs the full
serialized AgentStreamEvent (logger.info('Agent event', JSON.stringify(event)))
which can leak user/tool content and bloat logs; change this to log only minimal
metadata such as event.type, any safe IDs or timestamps, and the transition
detected using lastAgentEventType (e.g., logger.info('Agent event', { type:
event.type, prevType: lastAgentEventType, timestamp: ... })) and remove
JSON.stringify(event) so no full payload is written to logs.
- Around line 264-283: The code rejects completionDeferred when
token.isCancellationRequested but still proceeds to run
commands.executeCommand('notebook.cell.execute'), allowing work after
cancellation; update the handler (around token, completionDeferred,
CancellationError and before commands.executeCommand) to short-circuit: if token
&& token.isCancellationRequested (or if completionDeferred has already been
rejected/settled) then clear the timeout, dispose any disposables, and
return/throw so commands.executeCommand is not invoked; ensure the same
early-exit path is taken when token.onCancellationRequested fires so cancelled
executions never call notebook.cell.execute.
In `@src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts`:
- Around line 366-384: The test for executeEphemeralCell should also assert that
no execution request was sent when the token is pre-cancelled: after calling
executeEphemeralCell with the pre-cancelled CancellationTokenSource, add an
assertion that notebook.cell.execute was never invoked (i.e., verify/expect the
mocked notebook cell execution method did not get called), and keep the existing
assertion on the returned result; refer to executeEphemeralCell,
mockedVSCodeNamespaces.commands.executeCommand and the notebook.cell.execute
mock when adding this check.
In `@src/notebooks/deepnote/ephemeralCellDecorationProvider.ts`:
- Around line 117-123: The current loop in ephemeralCellDecorationProvider
builds a Range per line (lineRanges) and calls
editor.setDecorations(this.ephemeralDecorationType, lineRanges), which is
wasteful; replace it by creating a single full-cell Range spanning from the
start of the first line to the end of the last line (use
editor.document.lineAt(0).range.start and
editor.document.lineAt(editor.document.lineCount - 1).range.end) and pass an
array with that single Range to
editor.setDecorations(this.ephemeralDecorationType, [fullRange]) so you avoid
allocating per-line Range objects while preserving the same decoration coverage.
---
Duplicate comments:
In `@src/notebooks/deepnote/agentCellStatusBarProvider.ts`:
- Around line 207-224: The validateInput logic should reuse the existing
MaxIterationsSchema instead of using parseInt; replace the parseInt/isNaN checks
in validateInput with MaxIterationsSchema.safeParse(input) (or parse and catch)
and return l10n.t(...) on failure, ensuring the schema enforces integer-only and
range constraints consistent with MIN_ITERATIONS and MAX_ITERATIONS; after the
prompt returns, set newValue from the validated schema result (the parsed
numeric value) rather than calling parseInt again; refer to validateInput,
MaxIterationsSchema, getMaxIterations, MIN_ITERATIONS, MAX_ITERATIONS, and
newValue when making these changes.
🪄 Autofix (Beta)
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: ASSERTIVE
Plan: Pro
Run ID: e1298466-ae5e-4a9e-aaf3-1c4f03b06f10
📒 Files selected for processing (8)
package.jsonpackage.nls.jsonsrc/notebooks/deepnote/agentCellExecutionHandler.tssrc/notebooks/deepnote/agentCellExecutionHandler.unit.test.tssrc/notebooks/deepnote/agentCellStatusBarProvider.tssrc/notebooks/deepnote/deepnoteSecretStore.tssrc/notebooks/deepnote/deepnoteSecretStore.unit.test.tssrc/notebooks/deepnote/ephemeralCellDecorationProvider.ts
Comment on lines
+136
to
+138
| onAgentEvent: async (event: AgentStreamEvent) => { | ||
| logger.info('Agent event', JSON.stringify(event)); | ||
| if (lastAgentEventType != null && lastAgentEventType !== event.type) { |
Contributor
There was a problem hiding this comment.
Avoid logging full streamed agent events.
Line 137 logs full serialized payloads; this can expose user content/tool outputs in logs and increase log volume.
💡 Proposed fix
- logger.info('Agent event', JSON.stringify(event));
+ logger.info('Agent event', {
+ type: event.type,
+ toolName: 'toolName' in event ? event.toolName : undefined,
+ textLength: 'text' in event ? event.text.length : undefined,
+ outputLength: 'output' in event ? event.output?.length : undefined
+ });🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/notebooks/deepnote/agentCellExecutionHandler.ts` around lines 136 - 138,
The handler onAgentEvent currently logs the full serialized AgentStreamEvent
(logger.info('Agent event', JSON.stringify(event))) which can leak user/tool
content and bloat logs; change this to log only minimal metadata such as
event.type, any safe IDs or timestamps, and the transition detected using
lastAgentEventType (e.g., logger.info('Agent event', { type: event.type,
prevType: lastAgentEventType, timestamp: ... })) and remove
JSON.stringify(event) so no full payload is written to logs.
Comment on lines
+240
to
+245
| const edit = new WorkspaceEdit(); | ||
| edit.set(notebook.uri, [NotebookEdit.insertCells(insertIndex, [cellData])]); | ||
| await workspace.applyEdit(edit); | ||
|
|
||
| return insertIndex; | ||
| } |
Contributor
There was a problem hiding this comment.
Insertion path assumes edit success and returns a possibly invalid index.
If the edit is not applied, caller logic can read/execute the wrong cell at the returned index.
💡 Proposed fix
const edit = new WorkspaceEdit();
edit.set(notebook.uri, [NotebookEdit.insertCells(insertIndex, [cellData])]);
- await workspace.applyEdit(edit);
+ const applied = await workspace.applyEdit(edit);
+ if (!applied) {
+ throw new Error('Failed to insert ephemeral cell');
+ }
return insertIndex;
}As per coding guidelines "Verify state after async setup operations - methods can return early without throwing, so check expected state was created".
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| const edit = new WorkspaceEdit(); | |
| edit.set(notebook.uri, [NotebookEdit.insertCells(insertIndex, [cellData])]); | |
| await workspace.applyEdit(edit); | |
| return insertIndex; | |
| } | |
| const edit = new WorkspaceEdit(); | |
| edit.set(notebook.uri, [NotebookEdit.insertCells(insertIndex, [cellData])]); | |
| const applied = await workspace.applyEdit(edit); | |
| if (!applied) { | |
| throw new Error('Failed to insert ephemeral cell'); | |
| } | |
| return insertIndex; | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/notebooks/deepnote/agentCellExecutionHandler.ts` around lines 240 - 245,
The code currently assumes workspace.applyEdit(edit) succeeded and returns
insertIndex blindly; change it to check the boolean result of await
workspace.applyEdit(edit) and verify the notebook now contains the inserted cell
(e.g. notebook.cellCount > insertIndex or try notebook.cellAt(insertIndex)
exists). If applyEdit returns false or the verification fails, throw an Error
(or return a sentinel/failure value as per project convention) instead of
returning insertIndex so callers won't operate on an invalid index; use the same
local symbols edit, insertIndex, notebook, WorkspaceEdit,
NotebookEdit.insertCells and workspace.applyEdit to locate and implement the
checks.
Comment on lines
+264
to
+283
| if (token) { | ||
| if (token.isCancellationRequested) { | ||
| completionDeferred.reject(new CancellationError()); | ||
| } else { | ||
| disposables.push(token.onCancellationRequested(() => completionDeferred.reject(new CancellationError()))); | ||
| } | ||
| } | ||
|
|
||
| const timeout = setTimeout(() => { | ||
| completionDeferred.reject(new Error('Ephemeral cell execution timed out')); | ||
| }, EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS); | ||
|
|
||
| try { | ||
| const cellIndex = cell.index; | ||
|
|
||
| await commands.executeCommand('notebook.cell.execute', { | ||
| ranges: [{ start: cellIndex, end: cellIndex + 1 }], | ||
| document: cell.notebook.uri | ||
| }); | ||
|
|
Contributor
There was a problem hiding this comment.
Pre-cancelled executions still run the cell command.
Line 265 rejects the deferred, but Line 279 still executes notebook.cell.execute. That can run code after cancellation.
💡 Proposed fix
export async function executeEphemeralCell(
cell: NotebookCell,
token?: CancellationToken
): Promise<{ success: boolean; outputs: unknown[]; executionCount: number | null }> {
+ if (token?.isCancellationRequested) {
+ return {
+ success: false,
+ outputs: [],
+ executionCount: null
+ };
+ }
+
const completionDeferred = createDeferred<void>();
const disposables: IDisposable[] = [];
@@
- if (token) {
- if (token.isCancellationRequested) {
- completionDeferred.reject(new CancellationError());
- } else {
- disposables.push(token.onCancellationRequested(() => completionDeferred.reject(new CancellationError())));
- }
+ if (token) {
+ disposables.push(token.onCancellationRequested(() => completionDeferred.reject(new CancellationError())));
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/notebooks/deepnote/agentCellExecutionHandler.ts` around lines 264 - 283,
The code rejects completionDeferred when token.isCancellationRequested but still
proceeds to run commands.executeCommand('notebook.cell.execute'), allowing work
after cancellation; update the handler (around token, completionDeferred,
CancellationError and before commands.executeCommand) to short-circuit: if token
&& token.isCancellationRequested (or if completionDeferred has already been
rejected/settled) then clear the timeout, dispose any disposables, and
return/throw so commands.executeCommand is not invoked; ensure the same
early-exit path is taken when token.onCancellationRequested fires so cancelled
executions never call notebook.cell.execute.
Comment on lines
+366
to
+384
| test('returns success false immediately when token is pre-cancelled', async () => { | ||
| const cell = createMockCell({ index: 0 }); | ||
| const tokenSource = new CancellationTokenSource(); | ||
| tokenSource.cancel(); | ||
|
|
||
| when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); | ||
|
|
||
| try { | ||
| const result = await executeEphemeralCell(cell, tokenSource.token); | ||
|
|
||
| expect(result).to.deep.equal({ | ||
| success: false, | ||
| outputs: [], | ||
| executionCount: null | ||
| }); | ||
| } finally { | ||
| tokenSource.dispose(); | ||
| } | ||
| }); |
Contributor
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Strengthen cancellation test by asserting no execution command is sent.
The current assertion checks return value only. It should also verify notebook.cell.execute is not invoked for a pre-cancelled token.
💡 Proposed test addition
import { anything, capture, instance, mock, reset, when } from 'ts-mockito';
+import { verify } from 'ts-mockito';
@@
test('returns success false immediately when token is pre-cancelled', async () => {
@@
try {
const result = await executeEphemeralCell(cell, tokenSource.token);
expect(result).to.deep.equal({
success: false,
outputs: [],
executionCount: null
});
+ verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never();
} finally {
tokenSource.dispose();
}
});🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts` around lines
366 - 384, The test for executeEphemeralCell should also assert that no
execution request was sent when the token is pre-cancelled: after calling
executeEphemeralCell with the pre-cancelled CancellationTokenSource, add an
assertion that notebook.cell.execute was never invoked (i.e., verify/expect the
mocked notebook cell execution method did not get called), and keep the existing
assertion on the returned result; refer to executeEphemeralCell,
mockedVSCodeNamespaces.commands.executeCommand and the notebook.cell.execute
mock when adding this check.
Comment on lines
+117
to
+123
| const lineRanges: Range[] = []; | ||
| for (let i = 0; i < editor.document.lineCount; i++) { | ||
| const line = editor.document.lineAt(i); | ||
| lineRanges.push(line.range); | ||
| } | ||
|
|
||
| editor.setDecorations(this.ephemeralDecorationType, lineRanges); |
Contributor
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Use a single full-cell range instead of per-line ranges.
This loop allocates one Range per line on every refresh. For long cells, that is avoidable overhead.
💡 Proposed refactor
- const lineRanges: Range[] = [];
- for (let i = 0; i < editor.document.lineCount; i++) {
- const line = editor.document.lineAt(i);
- lineRanges.push(line.range);
- }
-
- editor.setDecorations(this.ephemeralDecorationType, lineRanges);
+ if (editor.document.lineCount === 0) {
+ editor.setDecorations(this.ephemeralDecorationType, []);
+ continue;
+ }
+
+ const first = editor.document.lineAt(0).range.start;
+ const last = editor.document.lineAt(editor.document.lineCount - 1).range.end;
+ editor.setDecorations(this.ephemeralDecorationType, [new Range(first, last)]);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/notebooks/deepnote/ephemeralCellDecorationProvider.ts` around lines 117 -
123, The current loop in ephemeralCellDecorationProvider builds a Range per line
(lineRanges) and calls editor.setDecorations(this.ephemeralDecorationType,
lineRanges), which is wasteful; replace it by creating a single full-cell Range
spanning from the start of the first line to the end of the last line (use
editor.document.lineAt(0).range.start and
editor.document.lineAt(editor.document.lineCount - 1).range.end) and pass an
array with that single Range to
editor.setDecorations(this.ephemeralDecorationType, [fullRange]) so you avoid
allocating per-line Range objects while preserving the same decoration coverage.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Behavior Changes
Tests