feat(js): ai.countTokens support

[ifielker]

Checklist (if applicable):

• PR title is following https://www.conventionalcommits.org/en/v1.0.0/
• Tested (manually, unit tested, etc.)
• Docs updated (updated docs or a docs bug required)

[gemini-code-assist]

Code Review

This pull request introduces a new countTokens API to the Genkit AI framework, allowing users to calculate token usage for generative models. Key changes include the addition of a countTokens function in generate.ts, integration into the GenkitAI class with multiple overloads, and updates to the model definition logic to support optional token counting actions. The PR also refactors core middleware to support non-generation outputs and updates the GenerationUsageSchema to deprecate the custom field in favor of raw. Review feedback suggests improving the robustness of the token counter lookup for dynamic models, ensuring background models also register their token counters, and refining the configuration handling when preparing token count requests. A simplification for handling PromiseLike options in the GenkitAI class was also recommended.

[gemini-code-assist]

For dynamic models (created via the model() helper) that are not explicitly registered in the registry, lookupAction will fail. Since the ModelAction object already contains the __tokenCounterAction property when countTokens is provided, you should check this property first before attempting a registry lookup.

Suggested change

	const counterActionName = `/model-token-counter/${resolvedModel.modelAction.__action.name}`;
	const counterAction = await registry.lookupAction<
	typeof GenerateRequestSchema,
	typeof GenerationUsageSchema,
	TokenCounterAction<CustomOptions>
	>(counterActionName);
	const counterAction =
	resolvedModel.modelAction.__tokenCounterAction ??
	(await registry.lookupAction<
	typeof GenerateRequestSchema,
	typeof GenerationUsageSchema,
	TokenCounterAction<CustomOptions>
	>(`/model-token-counter/${resolvedModel.modelAction.__action.name}`));

[gemini-code-assist]

This logic only registers token counters for standard models. Background models also support token counting (as DefineBackgroundModelOptions extends DefineModelOptions), but their counters are not being registered here because they are handled in a separate block (line 752). Ensure that background models also have their __tokenCounterAction registered in the registry.

[ifielker]

There are currently no background models with token counting capability. We will add it when those models are available.

[gemini-code-assist]

The configuration logic here should ensure that the resolved model's configuration replaces the original request's configuration rather than merging with it, to ensure the model operates with its intended settings. Additionally, include the model version and use stripUndefinedOptions to ensure that undefined values don't cause issues.

Suggested change

	request.config = {
	...(resolvedModel.config || {}),
	...(request.config || {}),
	};
	request.config = {
	version: resolvedModel.version,
	...stripUndefinedOptions(resolvedModel.config),
	};

References

1. By default, a fallback model's configuration should replace the original request's configuration, not merge with it. This ensures that the fallback operates with its intended, specific configuration.

[gemini-code-assist]

The instanceof Promise check is unreliable for PromiseLike objects (thenables) and adds unnecessary complexity. Since the underlying countTokens function already handles PromiseLike arguments, you can simplify this implementation to match the pattern used in generateStream.

  async countTokens<
    O extends z.ZodTypeAny = z.ZodTypeAny,
    CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
  >(
    options:
      | string
      | Part[]
      | GenerateOptions<O, CustomOptions>
      | PromiseLike<GenerateOptions<O, CustomOptions>>
  ): Promise<GenerationUsage> {
    if (typeof options === 'string' || Array.isArray(options)) {
      return countTokens(this.registry, { prompt: options });
    }
    return countTokens(this.registry, options);
  }