diff --git a/tests/Unit/Abilities/ImageGenerationPromptRefinementTest.php b/tests/Unit/Abilities/ImageGenerationPromptRefinementTest.php
index b5ca6c8a..eb5fa392 100644
--- a/tests/Unit/Abilities/ImageGenerationPromptRefinementTest.php
+++ b/tests/Unit/Abilities/ImageGenerationPromptRefinementTest.php
@@ -151,12 +151,17 @@ public function test_refine_prompt_includes_post_context_when_provided(): void {
 		ImageGenerationAbilities::refine_prompt( 'Crane meaning', 'This article explores the spiritual symbolism of cranes in various cultures.' );
 
 		$this->assertNotNull( $captured_request );
-		// Directives prepend messages, so find the last user message (our refinement request).
-		$user_message = '';
-		foreach ( array_reverse( $captured_request['messages'] ) as $msg ) {
-			if ( ( $msg['role'] ?? '' ) === 'user' ) {
-				$user_message = $msg['content'] ?? '';
-				break;
+		// `RequestBuilder::wpAiClientPromptContext()` lifts the latest user
+		// message into wp-ai-client's `prompt` slot and routes earlier turns
+		// through `with_history()`. Search both surfaces so the assertion
+		// stays true regardless of how many user messages were sent.
+		$user_message = (string) ( $captured_request['prompt'] ?? '' );
+		if ( '' === $user_message || ! str_contains( $user_message, 'Article context:' ) ) {
+			foreach ( array_reverse( $captured_request['messages'] ) as $msg ) {
+				if ( ( $msg['role'] ?? '' ) === 'user' ) {
+					$user_message = $msg['content'] ?? '';
+					break;
+				}
 			}
 		}
 		$this->assertStringContainsString( 'Article context:', $user_message );