fix(claude): derive context token budget from the Claude model in Claude Code mode

notebook_intelligence/extension.py

@@ -60,6 +60,7 @@
     claude_bypass_disabled_by_managed_settings,
     claude_managed_default_permission_mode,
     fetch_claude_models,
+    model_info_from_id,
     resolve_permission_mode,
 )
 from notebook_intelligence.claude_mcp_manager import ClaudeMCPManager
@@ -163,6 +164,27 @@ def _resolve_supports_vision(ai_service_manager) -> bool:
     return chat_model.supports_vision if chat_model is not None else False
 
 
+def _resolve_context_token_limit(ai_service_manager) -> int:
+    """Token budget source for attachment/output context.
+
+    In Claude Code mode the active model is Claude, not
+    ``ai_service_manager.chat_model`` — that property still reflects the
+    user's most recent non-Claude provider selection and is ``None`` on a
+    Claude-only setup. Falling through to the legacy 100-token floor in
+    that state shrank the context budget to 80 tokens, which silently
+    dropped cell-output attachments and truncated every attachment after
+    the first. Resolve the budget from the configured Claude model
+    instead (``model_info_from_id`` falls back to a 200K window for
+    unknown or default model ids).
+    """
+    if ai_service_manager.is_claude_code_mode:
+        model_id = ai_service_manager.nbi_config.claude_settings.get('chat_model', '')
+        model_id = model_id.strip() if isinstance(model_id, str) else ''
+        return model_info_from_id(model_id)["context_window"]
+    chat_model = ai_service_manager.chat_model
+    return 100 if chat_model is None else chat_model.context_window
+
+
 def _resolve_policy_with_env(env_var_name: str, traitlet_value: str) -> str:
     """Resolve a feature policy: env var wins if valid, else traitlet.
 
@@ -2325,7 +2347,7 @@ def on_message(self, message):
                     current_directory_file_msg += f" and current file is: '{filename}'"
                 chat_history.append({"role": "user", "content": current_directory_file_msg})
 
-            token_limit = 100 if ai_service_manager.chat_model is None else ai_service_manager.chat_model.context_window
+            token_limit = _resolve_context_token_limit(ai_service_manager)
             remaining_token_budget = int(0.8 * token_limit)
 
             # Resolve once; reused for sandbox containment and for

tests/test_websocket_handler_integration.py

@@ -948,3 +948,135 @@ def test_missing_mode_defaults(
         handler = self._handler(bypass_allowed=True)
         req = self._send(handler, mock_ai_manager, mock_nb_intel, None)
         assert req.permission_mode == 'default'
+
+
+class TestContextTokenLimit:
+    """The attachment/output-context budget must come from the Claude model
+    in Claude Code mode, not from the (possibly ``None``) legacy chat_model.
+
+    Regression guard: on a Claude-only setup ``chat_model`` is ``None``, and
+    the old ``100 if chat_model is None else ...`` expression shrank the
+    budget to 80 tokens — cell-output context was skipped and every
+    attachment after the first was dropped by the budget break.
+    """
+
+    def _mock_manager(self, claude_mode: bool, chat_model):
+        manager = Mock()
+        manager.is_claude_code_mode = claude_mode
+        manager.chat_model = chat_model
+        manager.nbi_config.claude_settings = {'chat_model': ''}
+        return manager
+
+    def test_non_claude_mode_without_model_keeps_floor(self):
+        from notebook_intelligence.extension import _resolve_context_token_limit
+        assert _resolve_context_token_limit(self._mock_manager(False, None)) == 100
+
+    def test_non_claude_mode_uses_model_context_window(self):
+        from notebook_intelligence.extension import _resolve_context_token_limit
+        chat_model = Mock()
+        chat_model.context_window = 4096
+        assert _resolve_context_token_limit(self._mock_manager(False, chat_model)) == 4096
+
+    def test_claude_mode_ignores_missing_legacy_model(self):
+        from notebook_intelligence.extension import _resolve_context_token_limit
+        with patch(
+            'notebook_intelligence.extension.model_info_from_id',
+            return_value={'id': '', 'name': '', 'context_window': 123456},
+        ) as mock_info:
+            assert _resolve_context_token_limit(self._mock_manager(True, None)) == 123456
+        mock_info.assert_called_once_with('')
+
+    def test_claude_mode_passes_configured_model_id(self):
+        from notebook_intelligence.extension import _resolve_context_token_limit
+        manager = self._mock_manager(True, None)
+        manager.nbi_config.claude_settings = {'chat_model': ' claude-sonnet-4-5 '}
+        with patch(
+            'notebook_intelligence.extension.model_info_from_id',
+            return_value={'id': 'claude-sonnet-4-5', 'name': 'Sonnet', 'context_window': 200000},
+        ) as mock_info:
+            assert _resolve_context_token_limit(manager) == 200000
+        mock_info.assert_called_once_with('claude-sonnet-4-5')
+
+    def test_claude_mode_default_window_without_model_cache(self):
+        # No patching: the real model_info_from_id falls back to a 200K
+        # window for an unknown/default id, so the budget is never the
+        # legacy 100-token floor in Claude Code mode.
+        from notebook_intelligence.extension import _resolve_context_token_limit
+        assert _resolve_context_token_limit(self._mock_manager(True, None)) >= 200000
+
+    @patch('notebook_intelligence.extension.ai_service_manager')
+    @patch('notebook_intelligence.extension.NotebookIntelligence')
+    @patch('notebook_intelligence.extension.threading.Thread')
+    def test_claude_mode_output_context_survives_without_legacy_model(
+        self, mock_thread, mock_nb_intel, mock_ai_manager
+    ):
+        """A cell-output attachment larger than the legacy 80-token budget
+        must still reach chat history in Claude Code mode."""
+        mock_nb_intel.root_dir = "/workspace"
+        mock_ai_manager.handle_chat_request = Mock()
+        mock_ai_manager.is_claude_code_mode = True
+        mock_ai_manager.chat_model = None  # Claude-only setup
+        mock_ai_manager.nbi_config.claude_settings = {'chat_model': ''}
+
+        mock_factory = Mock(spec=RuleContextFactory)
+        mock_factory.create.return_value = Mock(spec=RuleContext)
+
+        with patch('notebook_intelligence.extension.ThreadSafeWebSocketConnector'):
+            handler = WebsocketCopilotHandler(
+                self._create_mock_application(),
+                self._create_mock_request(),
+                context_factory=mock_factory
+            )
+
+        message = {
+            'id': 'test-message-id',
+            'type': 'chat-request',
+            'data': {
+                'chatId': 'test-chat-id',
+                'prompt': 'Explain this output',
+                'language': 'python',
+                'filename': 'notebook.ipynb',
+                'chatMode': 'agent',
+                'toolSelections': {},
+                'additionalContext': [
+                    {
+                        'filePath': 'notebook.ipynb',
+                        'outputContext': {
+                            'cellSource': 'df.describe()',
+                            'mimeBundles': [
+                                {
+                                    'mimeType': 'text/plain',
+                                    'data': 'count 100\nmean 4.2',
+                                    # Over the legacy 80-token budget: the
+                                    # old code path skipped this bundle.
+                                    'sizeTokens': 500,
+                                }
+                            ],
+                            'isError': False,
+                            'truncated': False,
+                        },
+                    }
+                ]
+            }
+        }
+
+        handler.on_message(json.dumps(message))
+
+        mock_ai_manager.handle_chat_request.assert_called_once()
+        chat_request = mock_ai_manager.handle_chat_request.call_args[0][0]
+        history_text = json.dumps(chat_request.chat_history)
+        assert 'df.describe()' in history_text
+        assert 'mean 4.2' in history_text
+
+    def _create_mock_application(self):
+        app = Mock(spec=Application)
+        app.settings = {"jinja2_env": None, "headers": {}}
+        app.ui_methods = {}
+        app.ui_modules = {}
+        app.transforms = []
+        return app
+
+    def _create_mock_request(self):
+        request = Mock(spec=HTTPServerRequest)
+        request.connection = Mock()
+        return request