Add collaborative_tool decorator for AI tool awareness

- Add collaborative_tool decorator in utils.py that enables real-time collaborative awareness
- Decorator automatically sets user presence in global and notebook-specific awareness systems
- Supports automatic file_path detection from function parameters
- Graceful error handling ensures tool execution continues if awareness fails
- Add comprehensive test suite with 15 tests covering all functionality
- Update README with example usage for creating collaborative tools
- Fix mypy type errors in notebook.py and utils.py
- Add pytest-asyncio dependency for async testing

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

Author: Zsailer

README.md

@@ -34,6 +34,40 @@ These tools are ideal for agents that assist users with code editing, version co
 
 ______________________________________________________________________
 
+## 🔧 Creating Collaborative Tools
+
+For developers building AI tools that need collaborative awareness, `jupyter_ai_tools` provides a `collaborative_tool` decorator that automatically enables real-time collaboration features:
+
+```python
+from jupyter_ai_tools.utils import collaborative_tool
+
+# Define user information
+user_info = {
+    "name": "Alice",
+    "color": "var(--jp-collaborator-color1)",
+    "display_name": "Alice Smith"
+}
+
+# Apply collaborative awareness to your tool
+@collaborative_tool(user=user_info)
+async def my_notebook_tool(file_path: str, content: str):
+    """Your tool implementation here"""
+    # Tool automatically sets user awareness for:
+    # - Global awareness system (all users can see Alice is active)
+    # - Notebook-specific awareness (for .ipynb files)
+    return f"Processed {file_path}"
+
+# For tools without user context, simply omit the user parameter
+@collaborative_tool()
+async def my_tool_no_user(file_path: str):
+    """Tool without collaborative awareness"""
+    return f"Processed {file_path}"
+```
+
+This decorator enables other users in the same Jupyter environment to see when your AI tool is actively working on shared notebooks, improving the collaborative experience.
+
+______________________________________________________________________
+
 ## Requirements
 
 - Jupyter Server

jupyter_ai_tools/toolkits/notebook.py

@@ -9,7 +9,7 @@
 from pycrdt import Awareness, Doc, Text, Assoc
 from jupyter_ydoc import YNotebook
 
-from ..utils import cell_to_md, get_file_id, get_jupyter_ydoc, notebook_json_to_md, get_global_awareness
+from ..utils import cell_to_md, get_file_id, get_jupyter_ydoc, notebook_json_to_md, get_global_awareness, collaborative_tool
 import re
 
 
@@ -161,7 +161,7 @@ async def read_cell_json(file_path: str, cell_id: str) -> Tuple[Dict[str, Any],
         raise
 
 
-async def get_cell_id_from_index(file_path: str, cell_index: int) -> Optional[int]:
+async def get_cell_id_from_index(file_path: str, cell_index: int) -> str:
     """Finds the cell_id of the cell at a specific cell index.
 
     This function reads a Jupyter notebook file and returns the UUID of the cell
@@ -250,7 +250,7 @@ async def add_cell(
                 ydoc.ycells.append(ycell)
             else:
                 ydoc.ycells.insert(insert_index, ycell)
-            await write_to_cell_collaboratively(ydoc, ycell, content)
+            await write_to_cell_collaboratively(ydoc, ycell, content or "")
         else:
             with open(file_path, "r", encoding="utf-8") as f:
                 notebook = nbformat.read(f, as_version=nbformat.NO_CONVERT)
@@ -316,7 +316,7 @@ async def insert_cell(
                 ydoc.ycells.append(ycell)
             else:
                 ydoc.ycells.insert(insert_index, ycell)
-            await write_to_cell_collaboratively(ydoc, ycell, content)
+            await write_to_cell_collaboratively(ydoc, ycell, content or "")
         else:
             with open(file_path, "r", encoding="utf-8") as f:
                 notebook = nbformat.read(f, as_version=nbformat.NO_CONVERT)
@@ -385,7 +385,7 @@ async def delete_cell(file_path: str, cell_id: str):
         raise
 
 
-def get_cursor_details(cell_source, start_index: int, stop_index: int = None) -> dict:
+def get_cursor_details(cell_source: Text, start_index: int, stop_index: Optional[int] = None) -> Dict[str, Any]:
     """
     Creates cursor details for collaborative notebook cursor positioning.
     
@@ -410,7 +410,7 @@ def get_cursor_details(cell_source, start_index: int, stop_index: int = None) ->
     head_sticky_index_data = head_sticky_index.to_json()
 
     # Initialize cursor details with default values
-    cursor_details = {"primary": True, "empty": True}
+    cursor_details: Dict[str, Any] = {"primary": True, "empty": True}
 
     # Set the head position (where cursor starts)
     cursor_details["head"] = {
@@ -438,7 +438,7 @@ def get_cursor_details(cell_source, start_index: int, stop_index: int = None) ->
     return cursor_details
 
 
-def set_cursor_in_ynotebook(ynotebook, cell_source, start_index: int, stop_index: int = None) -> None:
+def set_cursor_in_ynotebook(ynotebook: YNotebook, cell_source: Text, start_index: int, stop_index: Optional[int] = None) -> None:
     """
     Sets the cursor position in a collaborative notebook environment.
     
@@ -468,7 +468,8 @@ def set_cursor_in_ynotebook(ynotebook, cell_source, start_index: int, stop_index
         details = get_cursor_details(cell_source, start_index, stop_index=stop_index)
 
         # Update the awareness system with the cursor position
-        ynotebook.awareness.set_local_state_field("cursors", [details])
+        if ynotebook.awareness:
+            ynotebook.awareness.set_local_state_field("cursors", [details])
     except Exception:
         # Silently ignore cursor setting errors to avoid breaking main operations
         # This is intentional - cursor positioning is a visual enhancement, not critical
@@ -704,7 +705,7 @@ async def _handle_replace_operation(ynotebook, cell_source, cursor_position: int
     return cursor_position
 
 
-def _safe_set_cursor(ynotebook, cell_source, cursor_position: int, stop_cursor: int = None) -> None:
+def _safe_set_cursor(ynotebook: YNotebook, cell_source: Text, cursor_position: int, stop_cursor: Optional[int] = None) -> None:
     """
     Safely set cursor position with error handling.
     

jupyter_ai_tools/utils.py

@@ -1,4 +1,9 @@
 from jupyter_server.serverapp import ServerApp
+from pycrdt import Awareness
+import inspect
+import functools
+import typing
+from typing import Optional
 
 
 async def get_serverapp():
@@ -27,6 +32,19 @@ async def get_jupyter_ydoc(file_id: str):
         return notebook
 
 
+async def get_global_awareness() -> Optional[Awareness]:
+    serverapp = await get_serverapp()
+    yroom_manager = serverapp.web_app.settings["yroom_manager"]
+    
+    room_id = "JupyterLab:globalAwareness" 
+    if yroom_manager.has_room(room_id):
+        yroom = yroom_manager.get_room(room_id)
+        return yroom.get_awareness()
+    
+    # Return None if room doesn't exist
+    return None
+
+
 async def get_file_id(file_path: str) -> str:
     """Returns the file_id for the document
 
@@ -45,6 +63,94 @@ async def get_file_id(file_path: str) -> str:
     return file_id
 
 
+def collaborative_tool(user: typing.Optional[typing.Dict[str, typing.Any]] = None):
+    """
+    Decorator factory to enable collaborative awareness for toolkit functions.
+    
+    This decorator automatically sets up user awareness in the global
+    and notebook-specific awareness systems when functions are called.
+    It enables real-time collaborative features by making the user's
+    presence visible to other users in the same Jupyter environment.
+    
+    Args:
+        user: Optional user dictionary with user information. If None, no awareness is set.
+              Should contain keys like 'name', 'color', 'display_name', etc.
+    
+    Returns:
+        Decorator function that wraps the target function with collaborative awareness.
+        
+    Example:
+        >>> user_info = {
+        ...     "name": "Alice",
+        ...     "color": "var(--jp-collaborator-color1)",
+        ...     "display_name": "Alice Smith"
+        ... }
+        >>> 
+        >>> @collaborative_tool(user=user_info)
+        ... async def my_notebook_tool(file_path: str, content: str):
+        ...     # Your tool implementation here
+        ...     return f"Processed {file_path}"
+    """
+    def decorator(tool_func):
+        @functools.wraps(tool_func)
+        async def wrapper(*args, **kwargs):
+            # Skip awareness if no user provided
+            if user is None:
+                return await tool_func(*args, **kwargs)
+            
+            # Extract file_path from tool function arguments for notebook-specific awareness
+            file_path = None
+            try:
+                # Try to find file_path in kwargs first
+                if 'file_path' in kwargs:
+                    file_path = kwargs['file_path']
+                else:
+                    # Try to find file_path in positional args by inspecting the function signature
+                    sig = inspect.signature(tool_func)
+                    param_names = list(sig.parameters.keys())
+                    
+                    # Look for file_path parameter
+                    if 'file_path' in param_names and len(args) > param_names.index('file_path'):
+                        file_path = args[param_names.index('file_path')]
+                
+                # Set notebook-specific collaborative awareness if we have a file_path
+                if file_path and file_path.endswith('.ipynb'):
+                    try:
+                        file_id = await get_file_id(file_path)
+                        ydoc = await get_jupyter_ydoc(file_id)
+                        
+                        if ydoc:
+                            # Set the local user field in the notebook's awareness
+                            ydoc.awareness.set_local_state_field("user", user)
+                            
+                    except Exception:
+                        # Log but don't block tool execution
+                        pass
+                        
+            except Exception:
+                # Catch any errors in file_path detection
+                pass
+            
+            # Set global awareness
+            try:
+                g_awareness = await get_global_awareness()
+                if g_awareness:
+                    g_awareness.set_local_state({
+                        "user": user,
+                        "current": file_path or "",
+                        "documents": [file_path] if file_path else []
+                    })
+            except Exception:
+                # Log but don't block tool execution
+                pass
+            
+            # Execute the original tool function
+            return await tool_func(*args, **kwargs)
+        
+        return wrapper
+    return decorator
+
+
 def notebook_json_to_md(notebook_json: dict, include_outputs: bool = True) -> str:
     """Converts a notebook json dict to markdown string using a custom format.
 

pyproject.toml

@@ -30,7 +30,8 @@ dependencies = [
 [project.optional-dependencies]
 test = [
   "pytest>=7.0",
-  "pytest-jupyter[server]>=0.6"
+  "pytest-jupyter[server]>=0.6",
+  "pytest-asyncio"
 ]
 lint = [
   "black>=22.6.0",