From 17a758090ee7488f149c5a670daf13547de5a6f8 Mon Sep 17 00:00:00 2001 From: Herik Webb Date: Sun, 5 Jul 2026 08:27:23 -0400 Subject: [PATCH] fix(claude): derive context token budget from the Claude model in Claude Code mode In Claude Code mode ai_service_manager.chat_model reflects the user's most recent non-Claude provider selection and is None on a Claude-only setup. The attachment budget fell through to the legacy 100-token floor (80 tokens after the 0.8 factor), which silently skipped cell-output context and dropped every attachment after the first via the budget break in on_message. Resolve the budget from the configured Claude model instead, via model_info_from_id (200K-window fallback for unknown or default model ids). --- notebook_intelligence/extension.py | 24 +++- tests/test_websocket_handler_integration.py | 132 ++++++++++++++++++++ 2 files changed, 155 insertions(+), 1 deletion(-) diff --git a/notebook_intelligence/extension.py b/notebook_intelligence/extension.py index 60c793b..b411c60 100644 --- a/notebook_intelligence/extension.py +++ b/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 diff --git a/tests/test_websocket_handler_integration.py b/tests/test_websocket_handler_integration.py index f3ba3c0..9685a5d 100644 --- a/tests/test_websocket_handler_integration.py +++ b/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