diff --git a/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php b/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php
index 596656d36f516..9f3d8909afbfe 100644
--- a/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php
+++ b/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php
@@ -75,7 +75,19 @@ public function is_ability_call( FunctionCall $call ): bool {
return false;
}
- return str_starts_with( $name, self::ABILITY_PREFIX );
+ return self::is_ability_function_name( $name );
+ }
+
+ /**
+ * Checks if a function name refers to an ability.
+ *
+ * @since 7.2.0
+ *
+ * @param string $function_name The function name to check.
+ * @return bool True if the function name refers to an ability, false otherwise.
+ */
+ public static function is_ability_function_name( string $function_name ): bool {
+ return str_starts_with( $function_name, self::ABILITY_PREFIX );
}
/**
diff --git a/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php b/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php
index a64957fe73157..a683b54755cbc 100644
--- a/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php
+++ b/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php
@@ -11,6 +11,8 @@
use WordPress\AiClient\Builders\PromptBuilder;
use WordPress\AiClient\Common\Exception\InvalidArgumentException;
use WordPress\AiClient\Common\Exception\TokenLimitReachedException;
+use WordPress\AiClient\Events\AfterGenerateResultEvent;
+use WordPress\AiClient\Events\BeforeGenerateResultEvent;
use WordPress\AiClient\Files\DTO\File;
use WordPress\AiClient\Files\Enums\FileTypeEnum;
use WordPress\AiClient\Files\Enums\MediaOrientationEnum;
@@ -24,8 +26,11 @@
use WordPress\AiClient\Providers\Models\Contracts\ModelInterface;
use WordPress\AiClient\Providers\Models\DTO\ModelConfig;
use WordPress\AiClient\Providers\Models\Enums\CapabilityEnum;
+use WordPress\AiClient\Providers\Models\TextGeneration\Contracts\TextGenerationModelInterface;
use WordPress\AiClient\Providers\ProviderRegistry;
use WordPress\AiClient\Results\DTO\GenerativeAiResult;
+use WordPress\AiClient\Results\DTO\TokenUsage;
+use WordPress\AiClient\Tools\DTO\FunctionCall;
use WordPress\AiClient\Tools\DTO\FunctionDeclaration;
use WordPress\AiClient\Tools\DTO\FunctionResponse;
use WordPress\AiClient\Tools\DTO\WebSearch;
@@ -122,6 +127,14 @@ class WP_AI_Client_Prompt_Builder {
*/
private ?WP_Error $error = null;
+ /**
+ * Options for automatic ability resolution, or null when disabled.
+ *
+ * @since 7.2.0
+ * @var array{max_iterations: int}|null
+ */
+ private ?array $ability_resolution_options = null;
+
/**
* List of methods that generate a result from the prompt.
*
@@ -277,12 +290,69 @@ public function using_abilities( ...$abilities ): self {
return $this;
}
+ /**
+ * Enables automatic resolution of ability function calls.
+ *
+ * When enabled, the text generation methods run a resolution loop instead of
+ * a single request. Each round executes the ability calls requested by the
+ * model, appends the results to the conversation, and requests a follow-up
+ * response. The loop ends when the model produces a response without ability
+ * calls, when it requests a function that is not a registered ability, or
+ * when the maximum number of rounds is reached.
+ *
+ * Only abilities that were exposed to the model as function declarations,
+ * typically with {@see self::using_abilities()}, can be executed.
+ * Resolution follows the first response candidate and supports the
+ * generate_text_result() and generate_text() methods. Details about the loop
+ * are exposed under the `ability_resolution` key of the additional data of
+ * the final result.
+ *
+ * @since 7.2.0
+ *
+ * @param array $options {
+ * Optional. Options controlling the resolution loop.
+ *
+ * @type int $max_iterations Maximum number of resolution rounds. Each round executes
+ * the ability calls from one model response and requests a
+ * follow-up response. Default 5.
+ * }
+ * @return self The current instance for method chaining.
+ */
+ public function using_ability_resolution( array $options = array() ): self {
+ $options = wp_parse_args(
+ $options,
+ array(
+ 'max_iterations' => 5,
+ )
+ );
+
+ if ( ! is_int( $options['max_iterations'] ) || $options['max_iterations'] < 1 ) {
+ _doing_it_wrong(
+ __METHOD__,
+ sprintf(
+ /* translators: %s: max_iterations */
+ __( 'The %s option must be a positive integer.' ),
+ 'max_iterations'
+ ),
+ '7.2.0'
+ );
+ $options['max_iterations'] = 5;
+ }
+
+ $this->ability_resolution_options = array(
+ 'max_iterations' => $options['max_iterations'],
+ );
+
+ return $this;
+ }
+
/**
* Magic method to proxy snake_case method calls to their PHP AI Client camelCase counterparts.
*
* This allows WordPress developers to use snake_case naming conventions. It catches
* any exceptions thrown, stores them, and returns a WP_Error when a terminate method
- * is called.
+ * is called. When automatic ability resolution is enabled, the supported text
+ * generation methods run the resolution loop instead of a single request.
*
* @since 7.0.0
*
@@ -291,6 +361,37 @@ public function using_abilities( ...$abilities ): self {
* @return mixed The result of the method call.
*/
public function __call( string $name, array $arguments ) {
+ if ( null !== $this->ability_resolution_options && self::is_generating_method( $name ) ) {
+ if ( 'generate_text_result' === $name || 'generate_text' === $name ) {
+ return $this->generate_with_ability_resolution( $name );
+ }
+
+ _doing_it_wrong(
+ __METHOD__,
+ sprintf(
+ /* translators: 1: generate_text_result, 2: generate_text, 3: the method that was called. */
+ __( 'Automatic ability resolution supports only the %1$s and %2$s methods. The %3$s method runs without it.' ),
+ 'generate_text_result()',
+ 'generate_text()',
+ '' . esc_html( $name ) . '()'
+ ),
+ '7.2.0'
+ );
+ }
+
+ return $this->call_builder( $name, $arguments );
+ }
+
+ /**
+ * Proxies a method call to the wrapped prompt builder with WordPress-specific guards.
+ *
+ * @since 7.2.0
+ *
+ * @param string $name The method name in snake_case.
+ * @param array $arguments The method arguments.
+ * @return mixed The result of the method call.
+ */
+ private function call_builder( string $name, array $arguments ) {
/*
* If an error occurred in a previous method call, either return the error for terminate methods,
* or return the same instance for other methods to maintain the fluent interface.
@@ -307,39 +408,16 @@ public function __call( string $name, array $arguments ) {
// Check if the prompt should be prevented for is_supported* and generate_*/convert_text_to_speech* methods.
if ( self::is_support_check_method( $name ) || self::is_generating_method( $name ) ) {
- // If AI is not supported, then there's no need to apply the filter as the prompt will be prevented anyway.
- $is_ai_disabled = ! wp_supports_ai();
- $prevent = $is_ai_disabled;
- if ( ! $prevent ) {
- /**
- * Filters whether to prevent the prompt from being executed.
- *
- * @since 7.0.0
- *
- * @param bool $prevent Whether to prevent the prompt. Default false.
- * @param WP_AI_Client_Prompt_Builder $builder A clone of the prompt builder instance (read-only).
- */
- $prevent = (bool) apply_filters( 'wp_ai_client_prevent_prompt', false, clone $this );
- }
+ $prevented = $this->get_prompt_prevented_error();
- if ( $prevent ) {
+ if ( null !== $prevented ) {
// For is_supported* methods, return false.
if ( self::is_support_check_method( $name ) ) {
return false;
}
- $error_message = $is_ai_disabled
- ? __( 'AI features are not supported in this environment.' )
- : __( 'Prompt execution was prevented by a filter.' );
-
- // For generate_* and convert_text_to_speech* methods, create a WP_Error.
- $this->error = new WP_Error(
- 'prompt_prevented',
- $error_message,
- array(
- 'status' => 503,
- )
- );
+ // For generate_* and convert_text_to_speech* methods, store the WP_Error.
+ $this->error = $prevented;
if ( self::is_generating_method( $name ) ) {
return $this->error;
@@ -368,6 +446,328 @@ public function __call( string $name, array $arguments ) {
}
}
+ /**
+ * Checks whether the prompt is prevented from being executed.
+ *
+ * @since 7.2.0
+ *
+ * @return WP_Error|null A WP_Error when the prompt is prevented, null otherwise.
+ */
+ private function get_prompt_prevented_error(): ?WP_Error {
+ // If AI is not supported, then there's no need to apply the filter as the prompt will be prevented anyway.
+ $is_ai_disabled = ! wp_supports_ai();
+ $prevent = $is_ai_disabled;
+ if ( ! $prevent ) {
+ /**
+ * Filters whether to prevent the prompt from being executed.
+ *
+ * @since 7.0.0
+ *
+ * @param bool $prevent Whether to prevent the prompt. Default false.
+ * @param WP_AI_Client_Prompt_Builder $builder A clone of the prompt builder instance (read-only).
+ */
+ $prevent = (bool) apply_filters( 'wp_ai_client_prevent_prompt', false, clone $this );
+ }
+
+ if ( ! $prevent ) {
+ return null;
+ }
+
+ $error_message = $is_ai_disabled
+ ? __( 'AI features are not supported in this environment.' )
+ : __( 'Prompt execution was prevented by a filter.' );
+
+ return new WP_Error(
+ 'prompt_prevented',
+ $error_message,
+ array(
+ 'status' => 503,
+ )
+ );
+ }
+
+ /**
+ * Generates a text result while automatically resolving ability function calls.
+ *
+ * Runs the resolution loop: each round executes the ability calls requested
+ * by the model, appends the results to the conversation, and requests a
+ * follow-up response. See {@see self::using_ability_resolution()} for the
+ * termination conditions.
+ *
+ * @since 7.2.0
+ *
+ * @param string $method Either 'generate_text_result' or 'generate_text'.
+ * @return GenerativeAiResult|string|WP_Error The final result, the final text, or a WP_Error on failure.
+ */
+ private function generate_with_ability_resolution( string $method ) {
+ $options = $this->ability_resolution_options;
+
+ /*
+ * The PHP AI Client prompt builder does not expose its message list, nor
+ * a way to append messages to it. The first request therefore captures
+ * the sent messages and the resolved model from the lifecycle event that
+ * the builder dispatches. Later rounds call the captured model directly
+ * with an extended copy of that transcript. A message append method in
+ * the PHP AI Client would simplify this.
+ */
+ $captured = null;
+ $capture = static function ( $event ) use ( &$captured ) {
+ if ( null === $captured && $event instanceof BeforeGenerateResultEvent ) {
+ $captured = $event;
+ }
+ };
+
+ add_action( 'wp_ai_client_before_generate_result', $capture );
+ $result = $this->call_builder( 'generate_text_result', array() );
+ remove_action( 'wp_ai_client_before_generate_result', $capture );
+
+ if ( is_wp_error( $result ) ) {
+ return $result;
+ }
+
+ if ( null === $captured || ! $captured->getModel() instanceof TextGenerationModelInterface ) {
+ // Without the captured context the conversation cannot be continued.
+ return $this->to_generation_return_value( $result, $method );
+ }
+
+ $model = $captured->getModel();
+ $capability = $captured->getCapability();
+ $transcript = $captured->getMessages();
+ $dispatcher = AiClient::getEventDispatcher();
+
+ /*
+ * The allow-list for execution is derived from the function declarations
+ * that were sent to the model. A response may name any function, so the
+ * resolver enforces that only explicitly exposed abilities can run.
+ */
+ $ability_names = array();
+ $declarations = $model->getConfig()->getFunctionDeclarations() ?? array();
+ foreach ( $declarations as $declaration ) {
+ $function_name = $declaration->getName();
+ if ( WP_AI_Client_Ability_Function_Resolver::is_ability_function_name( $function_name ) ) {
+ $ability_names[] = WP_AI_Client_Ability_Function_Resolver::function_name_to_ability_name( $function_name );
+ }
+ }
+
+ if ( empty( $ability_names ) ) {
+ _doing_it_wrong(
+ __METHOD__,
+ sprintf(
+ /* translators: 1: using_ability_resolution, 2: using_abilities */
+ __( '%1$s requires abilities registered with %2$s.' ),
+ 'using_ability_resolution()',
+ 'using_abilities()'
+ ),
+ '7.2.0'
+ );
+ return $this->to_generation_return_value( $result, $method );
+ }
+
+ $resolver = new WP_AI_Client_Ability_Function_Resolver( ...$ability_names );
+
+ $rounds = 0;
+ $usage = $result->getTokenUsage();
+ $resolved_calls = array();
+
+ while ( true ) {
+ $message = $result->toMessage();
+ $calls = $this->get_function_calls( $message );
+
+ if ( empty( $calls ) ) {
+ $stop_reason = 'completed';
+ break;
+ }
+
+ $ability_calls = array_filter( $calls, array( $resolver, 'is_ability_call' ) );
+
+ if ( count( $ability_calls ) < count( $calls ) ) {
+ // The response requests functions that are not registered abilities.
+ // Hand the round back to the caller to resolve them.
+ $stop_reason = 'unresolved_function_calls';
+ break;
+ }
+
+ if ( $rounds >= $options['max_iterations'] ) {
+ $stop_reason = 'max_iterations';
+ break;
+ }
+
+ $prevented = $this->get_prompt_prevented_error();
+ if ( null !== $prevented ) {
+ $this->error = $prevented;
+ return $this->error;
+ }
+
+ $responses = $resolver->execute_abilities( $message );
+
+ foreach ( $ability_calls as $call ) {
+ $resolved_calls[] = array(
+ 'id' => $call->getId(),
+ 'ability' => WP_AI_Client_Ability_Function_Resolver::function_name_to_ability_name( (string) $call->getName() ),
+ );
+ }
+
+ $transcript[] = $message;
+ $transcript[] = $responses;
+ ++$rounds;
+
+ if ( null !== $dispatcher ) {
+ $dispatcher->dispatch( new BeforeGenerateResultEvent( $transcript, $model, $capability ) );
+ }
+
+ try {
+ $result = $model->generateTextResult( $transcript );
+ } catch ( Exception $e ) {
+ $this->error = $this->exception_to_wp_error( $e );
+ return $this->error;
+ }
+
+ if ( null !== $dispatcher ) {
+ $dispatcher->dispatch( new AfterGenerateResultEvent( $transcript, $model, $capability, $result ) );
+ }
+
+ $usage = $this->aggregate_token_usage( $usage, $result->getTokenUsage() );
+ }
+
+ return $this->finish_ability_resolution( $result, $method, $stop_reason, $rounds, $usage, $resolved_calls, $transcript );
+ }
+
+ /**
+ * Converts a result into the return value of the called generation method.
+ *
+ * Used when the resolution loop exits early with a plain result, so that
+ * generate_text() still returns a string or a WP_Error.
+ *
+ * @since 7.2.0
+ *
+ * @param GenerativeAiResult $result The result to convert.
+ * @param string $method Either 'generate_text_result' or 'generate_text'.
+ * @return GenerativeAiResult|string|WP_Error The result, its text, or a WP_Error on failure.
+ */
+ private function to_generation_return_value( GenerativeAiResult $result, string $method ) {
+ if ( 'generate_text' !== $method ) {
+ return $result;
+ }
+
+ try {
+ return $result->toText();
+ } catch ( Exception $e ) {
+ $this->error = $this->exception_to_wp_error( $e );
+ return $this->error;
+ }
+ }
+
+ /**
+ * Retrieves the function calls contained in a message.
+ *
+ * @since 7.2.0
+ *
+ * @param Message $message The message to inspect.
+ * @return FunctionCall[] The function calls in the message.
+ */
+ private function get_function_calls( Message $message ): array {
+ $calls = array();
+
+ foreach ( $message->getParts() as $part ) {
+ if ( $part->getType()->isFunctionCall() ) {
+ $call = $part->getFunctionCall();
+ if ( $call instanceof FunctionCall ) {
+ $calls[] = $call;
+ }
+ }
+ }
+
+ return $calls;
+ }
+
+ /**
+ * Adds up two token usage objects.
+ *
+ * @since 7.2.0
+ *
+ * @param TokenUsage $total The running total.
+ * @param TokenUsage $addition The usage to add.
+ * @return TokenUsage The combined token usage.
+ */
+ private function aggregate_token_usage( TokenUsage $total, TokenUsage $addition ): TokenUsage {
+ $thought_tokens = null;
+ if ( null !== $total->getThoughtTokens() || null !== $addition->getThoughtTokens() ) {
+ $thought_tokens = (int) $total->getThoughtTokens() + (int) $addition->getThoughtTokens();
+ }
+
+ return new TokenUsage(
+ $total->getPromptTokens() + $addition->getPromptTokens(),
+ $total->getCompletionTokens() + $addition->getCompletionTokens(),
+ $total->getTotalTokens() + $addition->getTotalTokens(),
+ $thought_tokens
+ );
+ }
+
+ /**
+ * Builds the final value of an ability resolution loop.
+ *
+ * Rebuilds the result with the aggregated token usage and details about the
+ * loop under the `ability_resolution` key of the additional data.
+ *
+ * @since 7.2.0
+ *
+ * @param GenerativeAiResult $result The result of the last round.
+ * @param string $method Either 'generate_text_result' or 'generate_text'.
+ * @param string $stop_reason Why the loop ended. One of 'completed',
+ * 'unresolved_function_calls', or 'max_iterations'.
+ * @param int $rounds Number of resolution rounds that ran.
+ * @param TokenUsage $usage Aggregated token usage across all rounds.
+ * @param array $resolved_calls The ability calls that were resolved.
+ * @param Message[] $transcript The conversation before the final response.
+ * @return GenerativeAiResult|string|WP_Error The final result or text, or a WP_Error on failure.
+ */
+ private function finish_ability_resolution( GenerativeAiResult $result, string $method, string $stop_reason, int $rounds, TokenUsage $usage, array $resolved_calls, array $transcript ) {
+ $messages = $transcript;
+ $messages[] = $result->toMessage();
+
+ $additional_data = $result->getAdditionalData();
+ $additional_data['ability_resolution'] = array(
+ 'rounds' => $rounds,
+ 'stop_reason' => $stop_reason,
+ 'resolved_calls' => $resolved_calls,
+ 'messages' => array_map(
+ static function ( Message $message ) {
+ return $message->toArray();
+ },
+ $messages
+ ),
+ );
+
+ $final = new GenerativeAiResult(
+ $result->getId(),
+ $result->getCandidates(),
+ $usage,
+ $result->getProviderMetadata(),
+ $result->getModelMetadata(),
+ $additional_data
+ );
+
+ if ( 'generate_text_result' === $method ) {
+ return $final;
+ }
+
+ // generate_text() returns the plain text of the final answer.
+ if ( 'completed' !== $stop_reason ) {
+ $this->error = new WP_Error(
+ 'ability_resolution_incomplete',
+ __( 'The model did not produce a final answer within the ability resolution limits.' ),
+ array(
+ 'status' => 500,
+ 'stop_reason' => $stop_reason,
+ 'rounds' => $rounds,
+ )
+ );
+ return $this->error;
+ }
+
+ return $this->to_generation_return_value( $final, $method );
+ }
+
/**
* Converts an exception into a WP_Error with a structured error code and message.
*
diff --git a/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php b/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php
index 7513df3ff0fd5..3ec9b26d3f55b 100644
--- a/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php
+++ b/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php
@@ -223,6 +223,95 @@ public function streamGenerateTextResult( array $prompt ): Generator { // phpcs:
};
}
+ /**
+ * Creates a mock text generation model that returns scripted consecutive results.
+ *
+ * Each generateTextResult() call returns the next result from the list and
+ * records the message list it received in the referenced capture array.
+ * When the results run out, the last result is returned again.
+ *
+ * @param array $results Scripted results or exceptions, in order.
+ * @param array $captured_prompts Receives the message list of each call.
+ * @param ModelMetadata|null $metadata Optional metadata.
+ * @return ModelInterface&TextGenerationModelInterface The mock model.
+ * @throws InvalidArgumentException If no results are provided.
+ */
+ protected function create_scripted_text_generation_model(
+ array $results,
+ array &$captured_prompts,
+ ?ModelMetadata $metadata = null
+ ): ModelInterface {
+ if ( empty( $results ) ) {
+ throw new InvalidArgumentException( 'At least one scripted result is required.' );
+ }
+
+ $metadata = $metadata ?? $this->create_test_text_model_metadata();
+
+ $provider_metadata = new ProviderMetadata(
+ 'mock',
+ 'Mock Provider',
+ ProviderTypeEnum::cloud()
+ );
+
+ return new class( $metadata, $provider_metadata, $results, $captured_prompts ) implements ModelInterface, TextGenerationModelInterface {
+
+ private ModelMetadata $metadata;
+ private ProviderMetadata $provider_metadata;
+ private array $results;
+ private array $captured_prompts;
+ private ModelConfig $config;
+
+ public function __construct(
+ ModelMetadata $metadata,
+ ProviderMetadata $provider_metadata,
+ array $results,
+ array &$captured_prompts
+ ) {
+ $this->metadata = $metadata;
+ $this->provider_metadata = $provider_metadata;
+ $this->results = $results;
+ $this->captured_prompts = &$captured_prompts;
+ $this->config = new ModelConfig();
+ }
+
+ public function metadata(): ModelMetadata {
+ return $this->metadata;
+ }
+
+ public function providerMetadata(): ProviderMetadata {
+ return $this->provider_metadata;
+ }
+
+ public function setConfig( ModelConfig $config ): void {
+ $this->config = $config;
+ }
+
+ public function getConfig(): ModelConfig {
+ return $this->config;
+ }
+
+ public function generateTextResult( array $prompt ): GenerativeAiResult {
+ $this->captured_prompts[] = $prompt;
+
+ if ( count( $this->results ) > 1 ) {
+ $result = array_shift( $this->results );
+ } else {
+ $result = $this->results[0];
+ }
+
+ if ( $result instanceof Exception ) {
+ throw $result;
+ }
+
+ return $result;
+ }
+
+ public function streamGenerateTextResult( array $prompt ): Generator {
+ yield $this->generateTextResult( $prompt );
+ }
+ };
+ }
+
/**
* Creates a mock image generation model using anonymous class.
*
diff --git a/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php b/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php
new file mode 100644
index 0000000000000..42b11ad6ecfeb
--- /dev/null
+++ b/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php
@@ -0,0 +1,784 @@
+registry = $this->createMock( ProviderRegistry::class );
+ }
+
+ /**
+ * Creates a result whose message consists of the given function calls.
+ *
+ * @param array $calls List of id, function name, and arguments triples.
+ * @return GenerativeAiResult The result containing the function calls.
+ */
+ private function create_function_call_result( array $calls ): GenerativeAiResult {
+ $parts = array();
+ foreach ( $calls as $call ) {
+ $parts[] = new MessagePart( new FunctionCall( $call[0], $call[1], $call[2] ) );
+ }
+
+ $candidate = new Candidate(
+ new ModelMessage( $parts ),
+ FinishReasonEnum::toolCalls()
+ );
+
+ return new GenerativeAiResult(
+ 'function-call-result',
+ array( $candidate ),
+ new TokenUsage( 5, 7, 12 ),
+ new ProviderMetadata( 'mock', 'Mock Provider', ProviderTypeEnum::cloud() ),
+ $this->create_test_text_model_metadata()
+ );
+ }
+
+ /**
+ * Creates a prompt builder backed by a scripted model with resolution enabled.
+ *
+ * @param array $results Scripted results or exceptions, in order.
+ * @param array $captured_prompts Receives each model call's message list.
+ * @param string ...$abilities Ability names to register on the builder.
+ * @return WP_AI_Client_Prompt_Builder The prompt builder.
+ */
+ private function create_resolution_builder( array $results, array &$captured_prompts, string ...$abilities ): WP_AI_Client_Prompt_Builder {
+ $model = $this->create_scripted_text_generation_model( $results, $captured_prompts );
+
+ $builder = new WP_AI_Client_Prompt_Builder( $this->registry, 'Test prompt' );
+ $builder->using_model( $model );
+
+ if ( ! empty( $abilities ) ) {
+ $builder->using_abilities( ...$abilities );
+ }
+
+ return $builder;
+ }
+
+ /**
+ * Returns the function name for a test ability.
+ *
+ * @param string $ability_name The ability name.
+ * @return string The function name exposed to the model.
+ */
+ private function function_name( string $ability_name ): string {
+ return WP_AI_Client_Ability_Function_Resolver::ability_name_to_function_name( $ability_name );
+ }
+
+ /**
+ * Test that using_ability_resolution() is chainable.
+ *
+ * @ticket 64865
+ */
+ public function test_using_ability_resolution_is_chainable() {
+ $builder = new WP_AI_Client_Prompt_Builder( $this->registry, 'Test prompt' );
+
+ $this->assertSame( $builder, $builder->using_ability_resolution() );
+ }
+
+ /**
+ * Test that the scripted model requires at least one result.
+ *
+ * @ticket 64865
+ */
+ public function test_scripted_model_requires_at_least_one_result() {
+ $captured = array();
+
+ $this->expectException( InvalidArgumentException::class );
+ $this->expectExceptionMessage( 'At least one scripted result is required.' );
+
+ $this->create_scripted_text_generation_model( array(), $captured );
+ }
+
+ /**
+ * Test that an invalid max_iterations option is rejected.
+ *
+ * @ticket 64865
+ * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::using_ability_resolution
+ */
+ public function test_using_ability_resolution_rejects_invalid_max_iterations() {
+ $builder = new WP_AI_Client_Prompt_Builder( $this->registry, 'Test prompt' );
+
+ $this->assertSame( $builder, $builder->using_ability_resolution( array( 'max_iterations' => 0 ) ) );
+ }
+
+ /**
+ * Test that an invalid option falls back to the default.
+ *
+ * @ticket 64865
+ * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::using_ability_resolution
+ */
+ public function test_invalid_max_iterations_falls_back_to_default() {
+ $captured = array();
+ $call_result = $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ );
+
+ // The scripted model keeps returning the function call result.
+ $builder = $this->create_resolution_builder( array( $call_result ), $captured, 'wpaiclienttests/simple' );
+ $result = $builder
+ ->using_ability_resolution( array( 'max_iterations' => 0 ) )
+ ->generate_text_result();
+
+ $this->assertSame( 5, $result->getAdditionalData()['ability_resolution']['rounds'] );
+ }
+
+ /**
+ * Test that a response without function calls passes through with loop metadata.
+ *
+ * @ticket 64865
+ */
+ public function test_result_without_function_calls_passes_through_with_metadata() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array( $this->create_test_result( 'Plain answer' ) ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertInstanceOf( GenerativeAiResult::class, $result );
+ $this->assertSame( 'Plain answer', $result->toText() );
+ $this->assertCount( 1, $captured );
+
+ $resolution = $result->getAdditionalData()['ability_resolution'];
+ $this->assertSame( 'completed', $resolution['stop_reason'] );
+ $this->assertSame( 0, $resolution['rounds'] );
+ $this->assertSame( array(), $resolution['resolved_calls'] );
+ $this->assertCount( 2, $resolution['messages'], 'The transcript should contain the prompt and the final response.' );
+ }
+
+ /**
+ * Test that an ability call is executed and the final answer is returned.
+ *
+ * @ticket 64865
+ */
+ public function test_resolves_ability_call_and_returns_final_answer() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ ),
+ $this->create_test_result( 'Final answer' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertInstanceOf( GenerativeAiResult::class, $result );
+ $this->assertSame( 'Final answer', $result->toText() );
+ $this->assertCount( 2, $captured );
+
+ $resolution = $result->getAdditionalData()['ability_resolution'];
+ $this->assertSame( 'completed', $resolution['stop_reason'] );
+ $this->assertSame( 1, $resolution['rounds'] );
+ $this->assertSame(
+ array(
+ array(
+ 'id' => 'call-1',
+ 'ability' => 'wpaiclienttests/simple',
+ ),
+ ),
+ $resolution['resolved_calls']
+ );
+ $this->assertCount( 4, $resolution['messages'], 'The transcript should contain the prompt, the call, the response, and the final answer.' );
+ }
+
+ /**
+ * Test that the follow-up request contains the expected conversation.
+ *
+ * @ticket 64865
+ */
+ public function test_second_request_contains_expected_transcript() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ ),
+ $this->create_test_result( 'Final answer' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $builder->using_ability_resolution()->generate_text_result();
+
+ $messages = $captured[1];
+ $this->assertCount( 3, $messages );
+ $this->assertTrue( $messages[0]->getRole()->isUser(), 'The first message should be the user prompt.' );
+ $this->assertTrue( $messages[1]->getRole()->isModel(), 'The second message should be the model response.' );
+ $this->assertTrue( $messages[2]->getRole()->isUser(), 'The third message should carry the function responses.' );
+
+ $parts = $messages[2]->getParts();
+ $this->assertCount( 1, $parts );
+
+ $response = $parts[0]->getFunctionResponse();
+ $this->assertInstanceOf( FunctionResponse::class, $response );
+ $this->assertSame( 'call-1', $response->getId() );
+ $this->assertSame( $this->function_name( 'wpaiclienttests/simple' ), $response->getName() );
+ $this->assertSame( array( 'success' => true ), $response->getResponse() );
+ }
+
+ /**
+ * Test that ability arguments from the model reach the ability.
+ *
+ * @ticket 64865
+ */
+ public function test_resolves_ability_call_with_arguments() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/with-params' ), array( 'title' => 'Hello' ) ) )
+ ),
+ $this->create_test_result( 'Done' ),
+ ),
+ $captured,
+ 'wpaiclienttests/with-params'
+ );
+
+ $builder->using_ability_resolution()->generate_text_result();
+
+ $response = $captured[1][2]->getParts()[0]->getFunctionResponse();
+ $this->assertSame(
+ array(
+ 'success' => true,
+ 'title' => 'Hello',
+ ),
+ $response->getResponse()
+ );
+ }
+
+ /**
+ * Test that all calls from one response are answered in a single message.
+ *
+ * @ticket 64865
+ */
+ public function test_answers_all_calls_from_one_response() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array(
+ array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ),
+ array( 'call-2', $this->function_name( 'wpaiclienttests/with-params' ), array( 'title' => 'Hello' ) ),
+ )
+ ),
+ $this->create_test_result( 'Done' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple',
+ 'wpaiclienttests/with-params'
+ );
+
+ $builder->using_ability_resolution()->generate_text_result();
+
+ $parts = $captured[1][2]->getParts();
+ $this->assertCount( 2, $parts );
+ $this->assertSame( 'call-1', $parts[0]->getFunctionResponse()->getId() );
+ $this->assertSame( 'call-2', $parts[1]->getFunctionResponse()->getId() );
+ }
+
+ /**
+ * Test that the loop stops after the configured number of rounds.
+ *
+ * @ticket 64865
+ */
+ public function test_stops_after_max_iterations() {
+ $captured = array();
+ $call_result = $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ );
+
+ // The scripted model keeps returning the function call result.
+ $builder = $this->create_resolution_builder( array( $call_result ), $captured, 'wpaiclienttests/simple' );
+ $result = $builder
+ ->using_ability_resolution( array( 'max_iterations' => 2 ) )
+ ->generate_text_result();
+
+ $this->assertInstanceOf( GenerativeAiResult::class, $result );
+ $this->assertCount( 3, $captured, 'The model should be called once initially and once per allowed round.' );
+
+ $resolution = $result->getAdditionalData()['ability_resolution'];
+ $this->assertSame( 'max_iterations', $resolution['stop_reason'] );
+ $this->assertSame( 2, $resolution['rounds'] );
+ }
+
+ /**
+ * Test that generate_text() returns an error when the loop is incomplete.
+ *
+ * @ticket 64865
+ */
+ public function test_generate_text_returns_error_when_max_iterations_reached() {
+ $captured = array();
+ $call_result = $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ );
+
+ $builder = $this->create_resolution_builder( array( $call_result ), $captured, 'wpaiclienttests/simple' );
+ $result = $builder
+ ->using_ability_resolution( array( 'max_iterations' => 1 ) )
+ ->generate_text();
+
+ $this->assertWPError( $result );
+ $this->assertSame( 'ability_resolution_incomplete', $result->get_error_code() );
+ $this->assertSame( 'max_iterations', $result->get_error_data()['stop_reason'] );
+ }
+
+ /**
+ * Test that generate_text() returns the final answer through the loop.
+ *
+ * @ticket 64865
+ */
+ public function test_generate_text_returns_final_answer() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ ),
+ $this->create_test_result( 'Final answer' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text();
+
+ $this->assertSame( 'Final answer', $result );
+ }
+
+ /**
+ * Test that the loop stops without executing anything when unknown functions are requested.
+ *
+ * @ticket 64865
+ */
+ public function test_stops_when_response_contains_unknown_function_calls() {
+ $invoked_abilities = array();
+ add_action(
+ 'wp_ability_invoked',
+ static function ( $ability_name ) use ( &$invoked_abilities ) {
+ $invoked_abilities[] = $ability_name;
+ }
+ );
+
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array(
+ array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ),
+ array( 'call-2', 'custom_function', array() ),
+ )
+ ),
+ $this->create_test_result( 'Never returned' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertInstanceOf( GenerativeAiResult::class, $result );
+ $this->assertCount( 1, $captured, 'The loop should not request a follow-up response.' );
+ $this->assertSame( array(), $invoked_abilities, 'No ability should be executed when unknown functions are requested.' );
+
+ $resolution = $result->getAdditionalData()['ability_resolution'];
+ $this->assertSame( 'unresolved_function_calls', $resolution['stop_reason'] );
+ $this->assertSame( 0, $resolution['rounds'] );
+ }
+
+ /**
+ * Test that an ability error is sent back to the model and the loop continues.
+ *
+ * @ticket 64865
+ */
+ public function test_error_from_ability_is_sent_back_to_model() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/returns-error' ), array() ) )
+ ),
+ $this->create_test_result( 'Recovered' ),
+ ),
+ $captured,
+ 'wpaiclienttests/returns-error'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertSame( 'Recovered', $result->toText() );
+
+ $response = $captured[1][2]->getParts()[0]->getFunctionResponse()->getResponse();
+ $this->assertSame( 'test_error', $response['code'] );
+ }
+
+ /**
+ * Test that a call to an ability outside the allowed list is answered with an error.
+ *
+ * @ticket 64865
+ */
+ public function test_not_allowed_ability_error_is_sent_back_to_model() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/with-params' ), array( 'title' => 'Hello' ) ) )
+ ),
+ $this->create_test_result( 'Done' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertSame( 'Done', $result->toText() );
+
+ $response = $captured[1][2]->getParts()[0]->getFunctionResponse()->getResponse();
+ $this->assertSame( 'ability_not_allowed', $response['code'] );
+ }
+
+ /**
+ * Test that the prevent filter also stops the loop between rounds.
+ *
+ * @ticket 64865
+ */
+ public function test_prevent_filter_stops_the_loop_between_rounds() {
+ $evaluations = 0;
+ $invoked_abilities = array();
+ add_filter(
+ 'wp_ai_client_prevent_prompt',
+ static function ( $prevent ) use ( &$evaluations ) { // phpcs:ignore Generic.CodeAnalysis.UnusedFunctionParameter
+ ++$evaluations;
+ return $evaluations > 1;
+ }
+ );
+ add_action(
+ 'wp_ability_invoked',
+ static function ( $ability_name ) use ( &$invoked_abilities ) {
+ $invoked_abilities[] = $ability_name;
+ }
+ );
+
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ ),
+ $this->create_test_result( 'Never returned' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertWPError( $result );
+ $this->assertSame( 'prompt_prevented', $result->get_error_code() );
+ $this->assertCount( 1, $captured, 'The follow-up request should be prevented.' );
+ $this->assertSame( array(), $invoked_abilities, 'No ability should be executed after prompt execution is prevented.' );
+ }
+
+ /**
+ * Test that token usage is aggregated across all rounds.
+ *
+ * @ticket 64865
+ */
+ public function test_token_usage_is_aggregated_across_rounds() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ // Uses 5 prompt, 7 completion, and 12 total tokens.
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ ),
+ // Uses 10 prompt, 20 completion, and 30 total tokens.
+ $this->create_test_result( 'Final answer' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+ $usage = $result->getTokenUsage();
+
+ $this->assertSame( 15, $usage->getPromptTokens() );
+ $this->assertSame( 27, $usage->getCompletionTokens() );
+ $this->assertSame( 42, $usage->getTotalTokens() );
+ }
+
+ /**
+ * Test that lifecycle events fire for every round.
+ *
+ * @ticket 64865
+ */
+ public function test_lifecycle_events_fire_for_each_round() {
+ $before_events = array();
+ $after_events = array();
+ add_action(
+ 'wp_ai_client_before_generate_result',
+ static function ( $event ) use ( &$before_events ) {
+ $before_events[] = $event;
+ }
+ );
+ add_action(
+ 'wp_ai_client_after_generate_result',
+ static function ( $event ) use ( &$after_events ) {
+ $after_events[] = $event;
+ }
+ );
+
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ ),
+ $this->create_test_result( 'Final answer' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertSame( 'Final answer', $result->toText() );
+ $this->assertCount( 2, $before_events, 'The before event should fire for the initial request and each round.' );
+ $this->assertCount( 2, $after_events, 'The after event should fire after the initial request and each successful round.' );
+
+ $this->assertInstanceOf( BeforeGenerateResultEvent::class, $before_events[0] );
+ $this->assertInstanceOf( BeforeGenerateResultEvent::class, $before_events[1] );
+ $this->assertCount( 1, $before_events[0]->getMessages() );
+ $this->assertCount( 3, $before_events[1]->getMessages() );
+ $this->assertEquals( $captured[0], $before_events[0]->getMessages() );
+ $this->assertEquals( $captured[1], $before_events[1]->getMessages() );
+
+ $this->assertInstanceOf( AfterGenerateResultEvent::class, $after_events[0] );
+ $this->assertInstanceOf( AfterGenerateResultEvent::class, $after_events[1] );
+ $this->assertCount( 1, $after_events[0]->getMessages() );
+ $this->assertCount( 3, $after_events[1]->getMessages() );
+ $this->assertCount( 1, $after_events[0]->getResult()->toMessage()->getParts() );
+ $this->assertSame( 'call-1', $after_events[0]->getResult()->toMessage()->getParts()[0]->getFunctionCall()->getId() );
+ $this->assertSame( 'Final answer', $after_events[1]->getResult()->toText() );
+ }
+
+ /**
+ * Test that a failed follow-up request does not dispatch an after event.
+ *
+ * @ticket 64865
+ */
+ public function test_failed_follow_up_request_does_not_fire_after_event() {
+ $before_events = array();
+ $after_events = array();
+ add_action(
+ 'wp_ai_client_before_generate_result',
+ static function ( $event ) use ( &$before_events ) {
+ $before_events[] = $event;
+ }
+ );
+ add_action(
+ 'wp_ai_client_after_generate_result',
+ static function ( $event ) use ( &$after_events ) {
+ $after_events[] = $event;
+ }
+ );
+
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ ),
+ new RuntimeException( 'Follow-up failed.' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertWPError( $result );
+ $this->assertSame( 'prompt_builder_error', $result->get_error_code() );
+ $this->assertSame( 'Follow-up failed.', $result->get_error_message() );
+ $this->assertCount( 2, $captured );
+ $this->assertCount( 2, $before_events, 'The before event should fire before the failed follow-up request.' );
+ $this->assertCount( 1, $after_events, 'The after event should only fire for the successful initial request.' );
+ $this->assertInstanceOf( AfterGenerateResultEvent::class, $after_events[0] );
+ }
+
+ /**
+ * Test that ability declarations added directly are resolvable too.
+ *
+ * The allow-list is derived from the function declarations exposed to the
+ * model, so declarations built without using_abilities() participate.
+ *
+ * @ticket 64865
+ */
+ public function test_directly_declared_abilities_are_resolvable() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ ),
+ $this->create_test_result( 'Final answer' ),
+ ),
+ $captured
+ );
+
+ $builder->using_function_declarations(
+ new FunctionDeclaration( $this->function_name( 'wpaiclienttests/simple' ), 'A simple test ability.' )
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertSame( 'Final answer', $result->toText() );
+ $this->assertSame( 1, $result->getAdditionalData()['ability_resolution']['rounds'] );
+ }
+
+ /**
+ * Test that replacing the declarations also replaces the allow-list.
+ *
+ * @ticket 64865
+ */
+ public function test_replaced_declarations_limit_the_allow_list() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array(
+ $this->create_function_call_result(
+ array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) )
+ ),
+ $this->create_test_result( 'Done' ),
+ ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ // Replaces the declarations from using_abilities() above.
+ $builder->using_function_declarations(
+ new FunctionDeclaration( $this->function_name( 'wpaiclienttests/with-params' ), 'Another test ability.' )
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertSame( 'Done', $result->toText() );
+
+ // The call to the no longer declared ability must not execute.
+ $response = $captured[1][2]->getParts()[0]->getFunctionResponse()->getResponse();
+ $this->assertSame( 'ability_not_allowed', $response['code'] );
+ }
+
+ /**
+ * Test that resolution without registered abilities falls back to plain generation.
+ *
+ * @ticket 64865
+ * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::generate_with_ability_resolution
+ */
+ public function test_resolution_without_abilities_falls_back_to_plain_generation() {
+ $captured = array();
+ $text_result = $this->create_test_result( 'Plain answer' );
+
+ $builder = $this->create_resolution_builder( array( $text_result ), $captured );
+ $result = $builder->using_ability_resolution()->generate_text_result();
+
+ $this->assertSame( $text_result, $result, 'The unmodified result should be returned.' );
+ $this->assertCount( 1, $captured );
+ }
+
+ /**
+ * Test that the generate_text() fallback still returns a string.
+ *
+ * @ticket 64865
+ * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::generate_with_ability_resolution
+ */
+ public function test_generate_text_without_abilities_falls_back_to_plain_text() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array( $this->create_test_result( 'Plain answer' ) ),
+ $captured
+ );
+
+ $result = $builder->using_ability_resolution()->generate_text();
+
+ $this->assertSame( 'Plain answer', $result );
+ }
+
+ /**
+ * Test that unsupported generation methods warn and run without resolution.
+ *
+ * @ticket 64865
+ * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::__call
+ */
+ public function test_resolution_warns_for_unsupported_generation_methods() {
+ $captured = array();
+ $builder = $this->create_resolution_builder(
+ array( $this->create_test_result( 'Plain answer' ) ),
+ $captured,
+ 'wpaiclienttests/simple'
+ );
+
+ $result = $builder->using_ability_resolution()->generate_image_result();
+
+ $this->assertWPError( $result );
+ }
+}