Skip to content

Allow tuning wp-ai-client connect timeout #1796

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chubes4 merged 1 commit into from
May 5, 2026
Merged

Allow tuning wp-ai-client connect timeout #1796

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions inc/Abilities/SettingsAbilities.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,10 @@ private function registerUpdateSettings(): void {
'description' => 'Per-mode provider/model overrides keyed by mode id',
),
'max_turns' => array( 'type' => 'integer' ),
'wp_ai_client_connect_timeout' => array(
'type' => 'number',
'description' => 'Connection timeout in seconds for wp-ai-client provider requests.',
),
'disabled_tools' => array( 'type' => 'object' ),
'ai_provider_keys' => array( 'type' => 'object' ),
'queue_tuning' => array(
Expand Down Expand Up @@ -353,6 +357,7 @@ public function executeGetSettings( array $input ): array {
'default_model' => $settings['default_model'] ?? '',
'mode_models' => $settings['mode_models'] ?? array(),
'max_turns' => $settings['max_turns'] ?? $defaults['max_turns'],
'wp_ai_client_connect_timeout' => $settings['wp_ai_client_connect_timeout'] ?? $defaults['wp_ai_client_connect_timeout'],
'disabled_tools' => $settings['disabled_tools'] ?? array(),
'ai_provider_keys' => $masked_keys,
'queue_tuning' => wp_parse_args( $settings['queue_tuning'] ?? array(), $defaults['queue_tuning'] ),
Expand Down Expand Up @@ -457,6 +462,11 @@ public function executeUpdateSettings( array $input ): array {
$handled_keys[] = 'max_turns';
}

if ( isset( $input['wp_ai_client_connect_timeout'] ) && is_numeric( $input['wp_ai_client_connect_timeout'] ) ) {
$all_settings['wp_ai_client_connect_timeout'] = max( 0.0, min( 300.0, (float) $input['wp_ai_client_connect_timeout'] ) );
$handled_keys[] = 'wp_ai_client_connect_timeout';
}

if ( isset( $input['disabled_tools'] ) ) {
$all_settings['disabled_tools'] = array();
foreach ( $input['disabled_tools'] as $tool_id => $disabled ) {
Expand Down
8 changes: 5 additions & 3 deletions inc/Core/PluginSettings.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
class PluginSettings {

public const DEFAULT_MAX_TURNS = 25;
public const DEFAULT_WP_AI_CLIENT_CONNECT_TIMEOUT = 30.0;

private static ?array $cache = null;
private static array $agent_model_cache = array();
Expand Down Expand Up @@ -54,12 +55,13 @@ public static function getDefaultQueueTuning(): array {
/**
* Get centralized plugin defaults used by backend and admin UI.
*
* @return array{max_turns:int,queue_tuning:array{concurrent_batches:int,batch_size:int,time_limit:int,chunk_size:int,chunk_delay:int}}
* @return array{max_turns:int,wp_ai_client_connect_timeout:float,queue_tuning:array{concurrent_batches:int,batch_size:int,time_limit:int,chunk_size:int,chunk_delay:int}}
*/
public static function getDefaults(): array {
return array(
'max_turns' => self::DEFAULT_MAX_TURNS,
'queue_tuning' => self::getDefaultQueueTuning(),
'max_turns' => self::DEFAULT_MAX_TURNS,
'wp_ai_client_connect_timeout' => self::DEFAULT_WP_AI_CLIENT_CONNECT_TIMEOUT,
'queue_tuning' => self::getDefaultQueueTuning(),
);
}

Expand Down
55 changes: 48 additions & 7 deletions inc/Engine/AI/RequestBuilder.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

namespace DataMachine\Engine\AI;

use DataMachine\Core\PluginSettings;
use DataMachine\Engine\AI\Directives\DirectivePolicyResolver;

defined( 'ABSPATH' ) || exit;
Expand Down Expand Up @@ -130,7 +131,9 @@ public static function build(
}

$result = null;
$request_options = null;
$request_timeout = self::wpAiClientRequestTimeout( $mode, $provider, $model, $payload );
$connect_timeout = self::wpAiClientConnectTimeout( $mode, $provider, $model, $payload, $request_timeout );
$timeout_filter = static function ( $default_timeout ) use ( $request_timeout ) {
return max( (float) $default_timeout, $request_timeout );
};
Expand Down Expand Up @@ -169,15 +172,13 @@ public static function build(
/** @var callable $model_resolver wp-ai-client exposes this through __call() in some versions. */
$model_resolver = array( $registry, 'getProviderModel' );
$model_instance = call_user_func( $model_resolver, $provider_id, $model, null );
if (
is_object( $model_instance )
&& method_exists( $model_instance, 'setRequestOptions' )
&& class_exists( '\WordPress\AiClient\Providers\Http\DTO\RequestOptions' )
) {
if ( class_exists( '\WordPress\AiClient\Providers\Http\DTO\RequestOptions' ) ) {
$request_options = new \WordPress\AiClient\Providers\Http\DTO\RequestOptions();
$request_options->setTimeout( $request_timeout );
$request_options->setConnectTimeout( min( 30.0, $request_timeout ) );
$model_instance->setRequestOptions( $request_options );
$request_options->setConnectTimeout( $connect_timeout );
if ( is_object( $model_instance ) && method_exists( $model_instance, 'setRequestOptions' ) ) {
$model_instance->setRequestOptions( $request_options );
}
}

// wp-ai-client refuses to construct a MessagePart from an empty
Expand All @@ -187,6 +188,9 @@ public static function build(
$builder = '' !== $prompt_context['prompt']
? \wp_ai_client_prompt( $prompt_context['prompt'] )
: \wp_ai_client_prompt();
if ( null !== $request_options && is_callable( array( $builder, 'using_request_options' ) ) ) {
$builder = $builder->using_request_options( $request_options );
}
$builder = $builder->using_provider( $provider_id )
->using_model( $model_instance );

Expand Down Expand Up @@ -400,6 +404,43 @@ private static function wpAiClientRequestTimeout( string $mode, string $provider
return max( 0.0, (float) $timeout );
}

/**
* Resolve the connection timeout Data Machine applies to wp-ai-client calls.
*
* @param string $mode Execution mode.
* @param string $provider Provider identifier.
* @param string $model Model identifier.
* @param array $payload Step payload.
* @param float $request_timeout Resolved full request timeout in seconds.
* @return float Timeout in seconds.
*/
private static function wpAiClientConnectTimeout( string $mode, string $provider, string $model, array $payload, float $request_timeout ): float {
$setting_default = PluginSettings::get(
'wp_ai_client_connect_timeout',
PluginSettings::DEFAULT_WP_AI_CLIENT_CONNECT_TIMEOUT
);
if ( ! is_numeric( $setting_default ) ) {
$setting_default = PluginSettings::DEFAULT_WP_AI_CLIENT_CONNECT_TIMEOUT;
}

$default_timeout = min( max( 0.0, (float) $setting_default ), $request_timeout );
$timeout = apply_filters(
'datamachine_wp_ai_client_connect_timeout',
$default_timeout,
$mode,
$provider,
$model,
$payload,
$request_timeout
);

if ( ! is_numeric( $timeout ) ) {
return $default_timeout;
}

return max( 0.0, (float) $timeout );
}

/**
* Assemble a provider request without dispatching it.
*
Expand Down
33 changes: 31 additions & 2 deletions tests/wp-ai-client-request-timeout-smoke.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ class TimeoutPromptBuilderDouble {

private string $provider = '';
private mixed $model = null;
private mixed $request_options = null;
private array $history = array();
private float $request_timeout = 30.0;

Expand All @@ -118,6 +119,11 @@ public function using_model_config( $model_config ): self {
return $this;
}

public function using_request_options( $request_options ): self {
$this->request_options = $request_options;
return $this;
}

public function using_system_instruction( string $system_instruction ): self {
$this->history[] = array( 'role' => 'system', 'content' => $system_instruction );
return $this;
Expand All @@ -139,6 +145,7 @@ public function generate_text_result() {
self::$captured_request = array(
'provider' => $this->provider,
'model' => $this->model,
'request_options' => $this->request_options,
'prompt' => $this->prompt,
'timeout' => $this->request_timeout,
'history' => $this->history,
Expand Down Expand Up @@ -194,7 +201,8 @@ function timeout_smoke_filter_count( string $tag ): int {
return $count;
}

$timeout_context = null;
$timeout_context = null;
$connect_timeout_context = null;
$GLOBALS['datamachine_test_wp_ai_client_model_with_request_options'] = true;

add_filter(
Expand All @@ -207,6 +215,16 @@ function ( float $timeout, string $mode, string $provider, string $model, array
5
);

add_filter(
'datamachine_wp_ai_client_connect_timeout',
function ( float $timeout, string $mode, string $provider, string $model, array $payload, float $request_timeout ) use ( &$connect_timeout_context ): float {
$connect_timeout_context = compact( 'timeout', 'mode', 'provider', 'model', 'payload', 'request_timeout' );
return 120.0;
},
10,
6
);

$result = RequestBuilder::build(
array(
array(
Expand Down Expand Up @@ -240,7 +258,18 @@ function ( float $timeout, string $mode, string $provider, string $model, array
$captured_request_options = is_object( $captured_model ) && method_exists( $captured_model, 'getRequestOptions' ) ? $captured_model->getRequestOptions() : null;
assert_timeout_smoke( $captured_request_options instanceof \WordPress\AiClient\Providers\Http\DTO\RequestOptions, 'Data Machine applies wp-ai-client RequestOptions to API-based models' );
assert_timeout_smoke( 240.0 === $captured_request_options?->getTimeout(), 'Data Machine sets RequestOptions timeout from scoped request timeout' );
assert_timeout_smoke( 30.0 === $captured_request_options?->getConnectTimeout(), 'Data Machine caps RequestOptions connect timeout at 30 seconds' );
assert_timeout_smoke( 120.0 === $captured_request_options?->getConnectTimeout(), 'Data Machine sets RequestOptions connect timeout from scoped connect timeout' );

$captured_builder_request_options = TimeoutPromptBuilderDouble::$captured_request['request_options'] ?? null;
assert_timeout_smoke( $captured_builder_request_options instanceof \WordPress\AiClient\Providers\Http\DTO\RequestOptions, 'Data Machine applies wp-ai-client RequestOptions to PromptBuilder' );
assert_timeout_smoke( 240.0 === $captured_builder_request_options?->getTimeout(), 'Data Machine sets PromptBuilder RequestOptions timeout from scoped request timeout' );
assert_timeout_smoke( 120.0 === $captured_builder_request_options?->getConnectTimeout(), 'Data Machine sets PromptBuilder RequestOptions connect timeout from scoped connect timeout' );

assert_timeout_smoke( 30.0 === ( $connect_timeout_context['timeout'] ?? null ), 'Data Machine connect timeout filter receives product default' );
assert_timeout_smoke( 240.0 === ( $connect_timeout_context['request_timeout'] ?? null ), 'Data Machine connect timeout filter receives resolved request timeout' );
assert_timeout_smoke( 'pipeline' === ( $connect_timeout_context['mode'] ?? null ), 'Data Machine connect timeout filter receives execution mode' );
assert_timeout_smoke( 'openai' === ( $connect_timeout_context['provider'] ?? null ), 'Data Machine connect timeout filter receives provider' );
assert_timeout_smoke( 'gpt-smoke' === ( $connect_timeout_context['model'] ?? null ), 'Data Machine connect timeout filter receives model' );

$captured_history = TimeoutPromptBuilderDouble::$captured_request['history'] ?? array();
$captured_prompt = TimeoutPromptBuilderDouble::$captured_request['prompt'] ?? null;
Expand Down
Loading