feat(agent-runner): support scoped token counting

src/langbot/pkg/agent/runner/context_builder.py

@@ -184,7 +184,7 @@ class AgentRunContextBuilder:
     def _is_llm_model_resource(model_resource: ModelResource) -> bool:
         operations = model_resource.get('operations')
         if isinstance(operations, list) and operations:
-            return bool({'invoke', 'stream'} & {str(operation) for operation in operations})
+            return bool({'invoke', 'stream', 'count_tokens'} & {str(operation) for operation in operations})
         return model_resource.get('model_type') != 'rerank'
 
     async def _build_model_context_window_tokens(self, resources: AgentResources) -> int | None:

src/langbot/pkg/agent/runner/resource_builder.py

@@ -101,9 +101,9 @@ class AgentResourceBuilder:
         seen_model_ids: set[str] = set()
 
         model_perms = set(manifest_perms.models)
-        include_llm = bool({'invoke', 'stream'} & model_perms)
+        include_llm = bool({'invoke', 'stream', 'count_tokens'} & model_perms)
         include_rerank = 'rerank' in model_perms
-        llm_operations = [operation for operation in ('invoke', 'stream') if operation in model_perms]
+        llm_operations = [operation for operation in ('invoke', 'stream', 'count_tokens') if operation in model_perms]
         if not include_llm and not include_rerank:
             return models
 

src/langbot/pkg/agent/runner/session_registry.py

@@ -13,7 +13,7 @@ from .context_builder import AgentResources
 MAX_STEERING_QUEUE_ITEMS = 100
 
 DEFAULT_RESOURCE_OPERATIONS: dict[str, set[str]] = {
-    'model': {'invoke', 'stream', 'rerank'},
+    'model': {'invoke', 'stream', 'rerank', 'count_tokens'},
     'tool': {'detail', 'call'},
     'knowledge_base': {'list', 'retrieve'},
     'skill': {'activate'},

src/langbot/pkg/plugin/handler.py

@@ -565,6 +565,55 @@ class RuntimeConnectionHandler(handler.Handler):
                 },
             )
 
+        @self.action(PluginToRuntimeAction.COUNT_TOKENS)
+        async def count_tokens(data: dict[str, Any]) -> handler.ActionResponse:
+            """Count model input tokens.
+
+            For AgentRunner calls: requires run_id and validates model_uuid against session.resources.models.
+            For regular plugin calls: no run_id, unrestricted access (backward compatibility).
+            """
+            llm_model_uuid = data['llm_model_uuid']
+            messages = data['messages']
+            funcs = data.get('funcs', [])
+            extra_args = data.get('extra_args', {})
+            run_id = data.get('run_id')
+            caller_plugin_identity = data.get('caller_plugin_identity')
+
+            if run_id:
+                _session, error = await _validate_run_authorization(
+                    run_id, 'model', llm_model_uuid, self.ap, caller_plugin_identity, operation='count_tokens'
+                )
+                if error:
+                    return error
+
+            llm_model = await self.ap.model_mgr.get_model_by_uuid(llm_model_uuid)
+            if llm_model is None:
+                return handler.ActionResponse.error(
+                    message=f'LLM model with llm_model_uuid {llm_model_uuid} not found',
+                )
+
+            messages_obj = [provider_message.Message.model_validate(message) for message in messages]
+
+            async def _placeholder_func(**kwargs):
+                pass
+
+            funcs_obj = [resource_tool.LLMTool.model_validate({**func, 'func': _placeholder_func}) for func in funcs]
+            count_tokens_method = getattr(llm_model.provider.requester, 'count_tokens', None)
+            if not callable(count_tokens_method):
+                return handler.ActionResponse.error(message='LLM provider does not support token counting')
+
+            try:
+                tokens = await count_tokens_method(
+                    model=llm_model,
+                    messages=messages_obj,
+                    funcs=funcs_obj,
+                    extra_args=extra_args,
+                )
+            except Exception as exc:
+                return handler.ActionResponse.error(message=f'Token counting failed: {exc}')
+
+            return handler.ActionResponse.success(data={'tokens': tokens})
+
         @self.action(PluginToRuntimeAction.INVOKE_LLM)
         async def invoke_llm(data: dict[str, Any]) -> handler.ActionResponse:
             """Invoke llm

src/langbot/pkg/provider/modelmgr/requester.py

@@ -411,6 +411,20 @@ class ProviderAPIRequester(metaclass=abc.ABCMeta):
         """
         pass
 
+    async def count_tokens(
+        self,
+        model: RuntimeLLMModel,
+        messages: typing.List[provider_message.Message],
+        funcs: typing.List[resource_tool.LLMTool] = None,
+        extra_args: dict[str, typing.Any] = {},
+    ) -> int:
+        """Count model input tokens before invoking the model.
+
+        Requesters should use the same provider/model conversion path as
+        ``invoke_llm`` so the preflight count matches the actual request shape.
+        """
+        raise NotImplementedError('This requester does not support token counting')
+
     async def invoke_llm_stream(
         self,
         query: pipeline_query.Query,

src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py

@@ -521,6 +521,33 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
 
         return args
 
+    async def count_tokens(
+        self,
+        model: requester.RuntimeLLMModel,
+        messages: typing.List[provider_message.Message],
+        funcs: typing.List[resource_tool.LLMTool] = None,
+        extra_args: dict[str, typing.Any] = {},
+    ) -> int:
+        """Count input tokens with LiteLLM's model-aware tokenizer."""
+        args = await self._build_completion_args(model, messages, funcs, extra_args, stream=False)
+        count_args: dict[str, typing.Any] = {
+            'model': args['model'],
+            'messages': args['messages'],
+        }
+        if 'tools' in args:
+            count_args['tools'] = args['tools']
+        if 'tool_choice' in args:
+            count_args['tool_choice'] = args['tool_choice']
+
+        try:
+            tokens = litellm.token_counter(**count_args)
+        except Exception as e:
+            self._handle_litellm_error(e)
+
+        if isinstance(tokens, bool) or not isinstance(tokens, int) or tokens < 0:
+            raise errors.RequesterError(f'token counter returned invalid value: {tokens!r}')
+        return tokens
+
     async def invoke_llm(
         self,
         query: pipeline_query.Query,