Fix agent errors

[3coins]

Summary

This PR refactors the Jupyternaut agent implementation to improve reliability, performance, and user experience.

Key Improvements:

• Better disambiguation between chat responses and notebook operations
• Incremental, iterative code generation instead of overwhelming users with extensive examples
• Automatic code execution and validation after cell creation/editing
• Proper handling of active notebook/cell context with new helper methods
• Enhanced error messages and output processing
• Context-aware username injection for multi-user scenarios

[dlqqq]

Rebased onto latest main to include #21; reviewing this now.

[dlqqq]

I asked it to create a notebook, and it does successfully open the notebook. However, it emits an error stating it cannot find the current cell.

Logged output:

[E 2025-12-01 10:31:37.400 PersonaManagerExtension] Persona 'Jupyternaut' encountered an exception printed below when attempting to stream output.
[E 2025-12-01 10:31:37.400 PersonaManagerExtension] No active cell found. Make sure the notebook is focused.
    Traceback (most recent call last):
      File "/Users/dlq/workplace/jupyter-ai-devrepo/jupyter-ai-persona-manager/jupyter_ai_persona_manager/base_persona.py", line 243, in stream_message
        async for chunk in reply_stream:
        ...<25 lines>...
            )
      File "/Users/dlq/workplace/jupyter-ai-devrepo/jupyter-ai-jupyternaut/jupyter_ai_jupyternaut/jupyternaut/jupyternaut.py", line 99, in create_aiter
        async for token, metadata in agent.astream(
        ...<8 lines>...
                    yield token.text
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/pregel/main.py", line 2956, in astream
        async for _ in runner.atick(
        ...<13 lines>...
                yield o
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/pregel/_runner.py", line 410, in atick
        _panic_or_proceed(
        ~~~~~~~~~~~~~~~~~^
            futures.done.union(f for f, t in futures.items() if t is not None),
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
            timeout_exc_cls=asyncio.TimeoutError,
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
            panic=reraise,
            ^^^^^^^^^^^^^^
        )
        ^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/pregel/_runner.py", line 520, in _panic_or_proceed
        raise exc
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/pregel/_retry.py", line 137, in arun_with_retry
        return await task.proc.ainvoke(task.input, config)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/_internal/_runnable.py", line 705, in ainvoke
        input = await asyncio.create_task(
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
            step.ainvoke(input, config, **kwargs), context=context
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        )
        ^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/_internal/_runnable.py", line 473, in ainvoke
        ret = await self.afunc(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/prebuilt/tool_node.py", line 760, in _afunc
        outputs = await asyncio.gather(*coros)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/prebuilt/tool_node.py", line 1099, in _arun_one
        return await self._execute_tool_async(tool_request, input_type, config)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/prebuilt/tool_node.py", line 1048, in _execute_tool_async
        content = _handle_tool_error(e, flag=self._handle_tool_errors)
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/prebuilt/tool_node.py", line 404, in _handle_tool_error
        content = flag(e)  # type: ignore [assignment, call-arg]
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/prebuilt/tool_node.py", line 361, in _default_handle_tool_errors
        raise e
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langgraph/prebuilt/tool_node.py", line 1001, in _execute_tool_async
        response = await tool.ainvoke(call_args, config)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langchain_core/tools/structured.py", line 66, in ainvoke
        return await super().ainvoke(input, config, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langchain_core/tools/base.py", line 615, in ainvoke
        return await self.arun(tool_input, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langchain_core/tools/base.py", line 1061, in arun
        raise error_to_raise
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langchain_core/tools/base.py", line 1027, in arun
        response = await coro_with_context(coro, context)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/.venv/lib/python3.13/site-packages/langchain_core/tools/structured.py", line 120, in _arun
        return await self.coroutine(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/Users/dlq/workplace/jupyter-ai-devrepo/jupyter-ai-jupyternaut/jupyter_ai_jupyternaut/jupyternaut/toolkits/jupyterlab.py", line 71, in run_cell
        await select_cell(cell_id, username)
      File "/Users/dlq/workplace/jupyter-ai-devrepo/jupyter-ai-jupyternaut/jupyter_ai_jupyternaut/jupyternaut/toolkits/notebook.py", line 1013, in select_cell
        raise RuntimeError(
            "No active cell found. Make sure the notebook is focused."
        )
    RuntimeError: No active cell found. Make sure the notebook is focused.
    During task with name 'tools' and id '20b1eaba-0de7-6cf6-f981-70d40b50a056'

[3coins]

@dlqqq
Can you try the version pre merging PR #21.

[dlqqq]

Can you try the version pre merging PR #21.

@3coins The commits are the same as what had existed before, minus the first 2 commits that were merged in #21; the only merge conflicts I addressed were minor whitespace issues, no code changes.

I was debugging this and found that after restarting JupyterLab, the notebook now contains an input cell that was added by Jupyternaut in the previous session. So Jupyternaut did add the cell correctly, I just needed to rebuild jupyter_server_documents:

However, the agent is getting stuck on the same error when trying to run the cell that was created (RuntimeError: No active cell found. Make sure the notebook is focused.).

[dlqqq]

Interestingly, if I re-open the notebook and ask Jupyternaut in a new chat to run all the cells, Jupyternaut can do that. Seems like the user's awareness state does not always contain the activeCellId if Jupyternaut opens the notebook for them. Here's what is shown when I log ydoc.awareness.states right before that exception is raised:

username: None
awareness_states: {3608698533: {'file_id': '52f88ec2-eed2-4df7-ab2c-591f90e51834', 'kernel': {'execution_state': 'idle'}, 'cursors': [{'primary': True, 'empty': True, 'head': {'type': {'client': 3608698533, 'clock': 74}, 'tname': None, 'item': {'client': 3608698533, 'clock': 74}, 'assoc': 0}, 'anchor': {'type': {'client': 3608698533, 'clock': 74}, 'tname': None, 'item': {'client': 3608698533, 'clock': 74}, 'assoc': 0}}]}, 2342954947: {'user': {'username': '16f6a1a0595549fa86865e5a04328087', 'name': 'Anonymous Io', 'display_name': 'Anonymous Io', 'initials': 'AI', 'avatar_url': None, 'color': 'var(--jp-collaborator-color4)'}}}

[dlqqq]

@3coins I see the problem. jupyterlab_notebook_awareness is not listed as a dependency, so I did not have it installed locally. I think you installed it manually so it works in your environment but not mine.

This also explains the weird behavior I recorded above. There are two separate tools: run_cell() and run_all_cells(). When Jupyternaut runs all cells, it just runs "notebook:run-all-cells", which works fine. However, when Jupyternaut runs a single cell, it calls select_cell(), which requires jupyterlab_notebook_awareness to get the active cell ID.

[dlqqq]

@3coins Thank you for working so hard on this over Thanksgiving week! The new agent is much better than before, and the real-time navigation capabilities are really impressive. I pushed a commit to add jupyterlab_notebook_awareness as a dependency so this works out of the box for users.

I noticed a couple of issues, but I think we can iterate on these as follow-up items & merge this for now.
Follow-up issues:

• #25
• #26

[dlqqq]

This may be causing #26. Updating ydoc._cells[cell_index]["source"] may not cause a YDoc update to be broadcasted immediately.