Fix session/image UX and align post cost totals

assets/js/block-image-generate.js

@@ -60,16 +60,13 @@
 			};
 			
 			const agentImageId = getAgentImageId();
-			if (!agentImageId) {
-				return wp.element.createElement(BlockEdit, props);
-			}
 
 			const openImageModal = () => {
 				// Dispatch custom event to open image generation modal
 				window.dispatchEvent(
 					new CustomEvent('wpaw:open-image-modal', {
 						detail: {
-							agentImageId: agentImageId,
+							agentImageId: agentImageId || null,
 							blockId: clientId,
 						},
 					})

assets/js/image-modal.js

@@ -8,7 +8,7 @@
 
 (function() {
 	const { Modal, Button, Spinner, TextControl, TextareaControl } = wp.components;
-	const { useState, useEffect, render } = wp.element;
+	const { useState, useEffect, render, createRoot } = wp.element;
 	
 	window.wpAgenticWriter = window.wpAgenticWriter || {};
 	
@@ -16,18 +16,75 @@
 	 * Image Review Modal
 	 * Shows after article generation with image recommendations
 	 */
-	window.wpAgenticWriter.ImageReviewModal = function({ postId, initialImageId, onClose, onComplete }) {
+	window.wpAgenticWriter.ImageReviewModal = function({ postId, initialImageId, initialBlockId, onClose, onComplete }) {
 		const [step, setStep] = useState('loading');
 		const [images, setImages] = useState([]);
+		const [generatedImages, setGeneratedImages] = useState([]);
 		const [selectedImages, setSelectedImages] = useState([]);
 		const [variantCounts, setVariantCounts] = useState({});
 		const [isGenerating, setIsGenerating] = useState(false);
+		const [committingVariantId, setCommittingVariantId] = useState(null);
+		const [extraVariantCounts, setExtraVariantCounts] = useState({});
+		const [generatingMoreFor, setGeneratingMoreFor] = useState(null);
 		const [error, setError] = useState(null);
 		
 		useEffect(() => {
 			loadImageRecommendations();
 		}, []);
 
+		const scopedImages = (() => {
+			if (initialImageId) {
+				const direct = images.filter((img) => String(img.agent_image_id || '').trim() === String(initialImageId).trim());
+				if (direct.length > 0) {
+					return direct;
+				}
+			}
+
+			if (!initialBlockId) {
+				return images;
+			}
+
+			const allBlocks = wp.data.select('core/block-editor').getBlocks() || [];
+			const flatImageBlocks = [];
+			const walk = (blocks) => {
+				blocks.forEach((block) => {
+					if (block?.name === 'core/image') {
+						flatImageBlocks.push(block);
+					}
+					if (Array.isArray(block?.innerBlocks) && block.innerBlocks.length > 0) {
+						walk(block.innerBlocks);
+					}
+				});
+			};
+			walk(allBlocks);
+
+			const slotIndex = flatImageBlocks.findIndex((block) => block?.clientId === initialBlockId);
+			if (slotIndex < 0) {
+				return images;
+			}
+
+			const placement = `slot_${slotIndex + 1}`;
+			const placementScoped = images.filter((img) => String(img.placement || '').toLowerCase() === placement);
+			return placementScoped.length > 0 ? placementScoped : images;
+		})();
+
+		useEffect(() => {
+			if ((!initialImageId && !initialBlockId) || scopedImages.length === 0) {
+				return;
+			}
+			const scopedIds = scopedImages.map((img) => img.agent_image_id).filter(Boolean);
+			if (scopedIds.length === 0) {
+				return;
+			}
+			const normalize = (items) => [...items].sort().join('|');
+			setSelectedImages((prev) => {
+				if (normalize(prev) === normalize(scopedIds)) {
+					return prev;
+				}
+				return scopedIds;
+			});
+		}, [initialImageId, initialBlockId, scopedImages.length]);
+		
 		const loadImageRecommendations = async () => {
 			try {
 				const response = await fetch(
@@ -40,7 +97,7 @@
 				);
 				
 				if (!response.ok) {
-					throw new Error('Failed to load image recommendations');
+					throw new Error(await extractApiErrorMessage(response, 'Failed to load image recommendations'));
 				}
 				
 				const data = await response.json();
@@ -53,6 +110,13 @@
 				});
 				setVariantCounts(initialCounts);
 
+				if (initialImageId) {
+					const preferred = imgs.filter((img) => String(img.agent_image_id || '').trim() === String(initialImageId).trim());
+					if (preferred.length > 0) {
+						setSelectedImages(preferred.map((img) => img.agent_image_id));
+					}
+				}
+				
 				setStep('review');
 			} catch (err) {
 				setError(err.message);
@@ -104,6 +168,14 @@
 			return total.toFixed(3);
 		};
 
+		const calculateTotalVariants = () => {
+			let total = 0;
+			selectedImages.forEach((imageId) => {
+				total += variantCounts[imageId] || 2;
+			});
+			return total;
+		};
+		
 		const handleGenerateSelected = async () => {
 			if (selectedImages.length === 0) {
 				alert('Please select at least one image to generate');
@@ -114,8 +186,12 @@
 			setStep('generating');
 			
 			try {
+				const generated = [];
 				for (const imageId of selectedImages) {
 					const image = images.find(img => img.agent_image_id === imageId);
+					if (!image) {
+						throw new Error(`Image recommendation was not found for ${imageId}`);
+					}
 					
 					const response = await fetch(
 						`${wpAgenticWriter.apiUrl}/generate-image`,
@@ -130,24 +206,31 @@
 								agent_image_id: imageId,
 								prompt: image.prompt_edited || image.prompt_initial,
 								alt: image.alt_text_edited || image.alt_text_initial,
-								variant_count: variantCounts[imageId] || 1,
+								variant_count: variantCounts[imageId] || 2,
 							}),
 						}
 					);
 					
 					if (!response.ok) {
-						throw new Error(`Failed to generate image: ${imageId}`);
+						const detailedError = await extractApiErrorMessage(
+							response,
+							`Failed to generate image: ${imageId}`
+						);
+						throw new Error(detailedError);
 					}
 					
 					const result = await response.json();
+					const imageWithVariants = { ...image, variants: result.variants || [] };
+					generated.push(imageWithVariants);
 					
 					setImages(prev => prev.map(img => 
 						img.agent_image_id === imageId 
-							? { ...img, variants: result.variants }
+							? imageWithVariants
 							: img
 					));
 				}
 				
+				setGeneratedImages(generated);
 				setStep('selecting');
 			} catch (err) {
 				setError(err.message);
@@ -158,9 +241,11 @@
 		};
 		
 		const handleSelectVariant = async (imageId, variantId) => {
-			const image = images.find(img => img.agent_image_id === imageId);
+			const image = generatedImages.find(img => img.agent_image_id === imageId)
+				|| images.find(img => img.agent_image_id === imageId);
 			
 			try {
+				setCommittingVariantId(variantId);
 				const response = await fetch(
 					`${wpAgenticWriter.apiUrl}/commit-image`,
 					{
@@ -179,41 +264,175 @@
 				);
 				
 				if (!response.ok) {
-					throw new Error('Failed to commit image');
+					throw new Error(await extractApiErrorMessage(response, 'Failed to commit image'));
 				}
 				
 				const result = await response.json();
 				
-				updateGutenbergBlock(imageId, result);
+				const didUpdateBlock = updateGutenbergBlock(imageId, result);
 				
 				setImages(prev => prev.map(img => 
 					img.agent_image_id === imageId 
 						? { ...img, status: 'committed', attachment_id: result.attachment_id }
 						: img
 				));
+				setGeneratedImages(prev => prev.map(img =>
+					img.agent_image_id === imageId
+						? { ...img, status: 'committed', attachment_id: result.attachment_id }
+						: img
+				));
+				if (!didUpdateBlock) {
+					throw new Error('Image committed to Media Library, but the matching image block could not be found.');
+				}
+				onComplete();
 			} catch (err) {
 				alert('Failed to commit image: ' + err.message);
+				setCommittingVariantId(null);
 			}
 		};
 
+		const handleGenerateMore = async (imageId) => {
+			const image = generatedImages.find(img => img.agent_image_id === imageId)
+				|| images.find(img => img.agent_image_id === imageId);
+			if (!image) {
+				alert('Image recommendation not found.');
+				return;
+			}
+
+			const moreCount = extraVariantCounts[imageId] || 1;
+			setGeneratingMoreFor(imageId);
+			try {
+				const response = await fetch(
+					`${wpAgenticWriter.apiUrl}/generate-image`,
+					{
+						method: 'POST',
+						headers: {
+							'Content-Type': 'application/json',
+							'X-WP-Nonce': wpAgenticWriter.nonce,
+						},
+						body: JSON.stringify({
+							post_id: postId,
+							agent_image_id: imageId,
+							prompt: image.prompt_edited || image.prompt_initial,
+							alt: image.alt_text_edited || image.alt_text_initial,
+							variant_count: moreCount,
+						}),
+					}
+				);
+
+				if (!response.ok) {
+					throw new Error(await extractApiErrorMessage(response, `Failed to generate more variants for ${imageId}`));
+				}
+
+				const result = await response.json();
+				const newVariants = Array.isArray(result?.variants) ? result.variants : [];
+				if (newVariants.length === 0) {
+					throw new Error('No additional variants returned.');
+				}
+
+				setGeneratedImages(prev => prev.map(img =>
+					img.agent_image_id === imageId
+						? { ...img, variants: [...(img.variants || []), ...newVariants] }
+						: img
+				));
+				setImages(prev => prev.map(img =>
+					img.agent_image_id === imageId
+						? { ...img, variants: [...(img.variants || []), ...newVariants] }
+						: img
+				));
+			} catch (err) {
+				alert('Failed to generate more variants: ' + err.message);
+			} finally {
+				setGeneratingMoreFor(null);
+			}
+		};
+
+		const extractApiErrorMessage = async (response, fallback) => {
+			try {
+				const data = await response.clone().json();
+				const message = data?.message || data?.error || '';
+				const details = data?.details || data?.reason || data?.provider_error || '';
+				const trace = data?.data?.trace || data?.trace || null;
+				const traceText = formatTraceForDisplay(trace);
+				const combined = [message, details, traceText].filter(Boolean).join('\n\n').trim();
+				if (combined) {
+					return combined;
+				}
+			} catch (jsonError) {
+				// Continue to text fallback.
+			}
+
+			try {
+				const text = (await response.text()).trim();
+				if (text) {
+					return text;
+				}
+			} catch (textError) {
+				// Ignore read error.
+			}
+
+			return fallback;
+		};
+
+		const formatTraceForDisplay = (trace) => {
+			if (!trace || typeof trace !== 'object') {
+				return '';
+			}
+
+			const openRouterMessage = trace.openrouter_response?.error?.message
+				|| trace.openrouter_response?.message
+				|| '';
+			const lines = [
+				'Trace:',
+				`selected_model: ${trace.selected_model || ''}`,
+				`settings_image_model: ${trace.settings_image_model || ''}`,
+				`image_task_provider: ${trace.image_task_provider || ''}`,
+				`custom_model_match: ${trace.custom_model_match ? 'yes' : 'no'}`,
+				`custom_model_ids: ${Array.isArray(trace.custom_model_ids) ? trace.custom_model_ids.join(', ') : ''}`,
+				`model_cache_loaded: ${trace.model_cache_loaded ? 'yes' : 'no'}`,
+				`model_cache_has_model: ${trace.model_cache_has_model ? 'yes' : 'no'}`,
+				`refreshed_has_model: ${trace.refreshed_has_model === null ? 'n/a' : (trace.refreshed_has_model ? 'yes' : 'no')}`,
+				`endpoint: ${trace.endpoint || ''}`,
+				`request_model: ${trace.request_model || ''}`,
+				`request_size: ${trace.request_size || ''}`,
+				`request_quality: ${trace.request_quality || ''}`,
+				`request_n: ${trace.request_n || ''}`,
+				`openrouter_http: ${trace.openrouter_http || ''}`,
+			];
+
+			if (openRouterMessage) {
+				lines.push(`openrouter_message: ${openRouterMessage}`);
+			}
+
+			return lines.filter((line) => !line.endsWith(': ')).join('\n');
+		};
+		
 		const updateGutenbergBlock = (agentImageId, attachmentData) => {
-			const blocks = wp.data.select('core/block-editor').getBlocks();
+			const editor = wp.data.select('core/block-editor');
+			const dispatcher = wp.data.dispatch('core/block-editor');
+			const blocks = editor.getBlocks();
+			const updateBlock = (clientId) => {
+				if (!clientId) {
+					return false;
+				}
+				dispatcher.updateBlockAttributes(
+					clientId,
+					{
+						id: attachmentData.attachment_id,
+						url: attachmentData.attachment_url,
+						alt: attachmentData.alt,
+						caption: '',
+						'data-agent-image-id': agentImageId,
+					}
+				);
+				return true;
+			};
 			
 			const findAndUpdateBlock = (blocks) => {
 				for (const block of blocks) {
 					if (block.name === 'core/image' && 
 						block.attributes['data-agent-image-id'] === agentImageId) {
-						
+						return updateBlock(block.clientId);
-						wp.data.dispatch('core/block-editor').updateBlockAttributes(
-							block.clientId,
-							{
-								id: attachmentData.attachment_id,
-								url: attachmentData.attachment_url,
-								alt: attachmentData.alt,
-								'data-agent-image-id': undefined,
-							}
-						);
-						return true;
 					}
 					
 					if (block.innerBlocks && block.innerBlocks.length > 0) {
@@ -225,7 +444,21 @@
 				return false;
 			};
 			
-			findAndUpdateBlock(blocks);
+			if (findAndUpdateBlock(blocks)) {
+				return true;
+			}
+
+			const initialBlock = initialBlockId ? editor.getBlock(initialBlockId) : null;
+			if (initialBlock?.name === 'core/image') {
+				return updateBlock(initialBlockId);
+			}
+
+			const selectedBlock = editor.getSelectedBlock();
+			if (selectedBlock?.name === 'core/image') {
+				return updateBlock(selectedBlock.clientId);
+			}
+
+			return false;
 		};
 		
 		if (step === 'loading') {
@@ -241,7 +474,7 @@
 		
 		if (step === 'review') {
 			return wp.element.createElement(Modal, {
-				title: `Image Recommendations (${images.length})`,
+				title: `Image Recommendations (${scopedImages.length})`,
 				onRequestClose: onClose,
 				style: { maxWidth: '800px' },
 			},
@@ -251,7 +484,7 @@
 						style: { marginBottom: '20px' }
 					}, error),
 					
-					images.length === 0 && wp.element.createElement('div', {
+					scopedImages.length === 0 && wp.element.createElement('div', {
 						style: {
 							padding: '40px 20px',
 							textAlign: 'center',
@@ -270,7 +503,7 @@
 						}, 'Continue Without Images')
 					),
 					
-					images.map(image => 
+					scopedImages.map(image => 
 						wp.element.createElement('div', {
 							key: image.agent_image_id,
 							className: 'wpaw-image-card',
@@ -311,6 +544,7 @@
 										padding: '5px',
 										borderRadius: '3px',
 										border: '1px solid #ddd',
+										minWidth: '150px'
 									}
 								},
 									wp.element.createElement('option', { value: '1' }, '1 variant'),
@@ -356,7 +590,7 @@
 							variant: 'primary',
 							onClick: handleGenerateSelected,
 							disabled: selectedImages.length === 0 || isGenerating,
-						}, `Generate ${selectedImages.length} Image(s) (~$${calculateTotalCost()})`)
+						}, `Generate ${calculateTotalVariants()} Image(s) (~$${calculateTotalCost()})`)
 					)
 				)
 			);
@@ -380,20 +614,56 @@
 		}
 		
 		if (step === 'selecting') {
+			const selectableImages = generatedImages.filter(img => img.variants && img.variants.length > 0);
 			return wp.element.createElement(Modal, {
 				title: 'Select Image Variants',
 				onRequestClose: onClose,
 				style: { maxWidth: '900px' },
 			},
 				wp.element.createElement('div', { className: 'wpaw-variant-selection' },
-					images
+					selectableImages.length === 0 && wp.element.createElement('div', {
-						.filter(img => img.variants && img.variants.length > 0)
+						style: { padding: '24px', color: '#666' }
+					}, 'No generated variants were returned. Please try generating again.'),
+					selectableImages
 						.map(image => 
 							wp.element.createElement('div', {
 								key: image.agent_image_id,
 								style: { marginBottom: '30px' },
 							},
 								wp.element.createElement('h3', null, image.section_title),
+								wp.element.createElement('div', {
+									style: {
+										display: 'flex',
+										alignItems: 'center',
+										gap: '8px',
+										marginBottom: '10px',
+									}
+								},
+									wp.element.createElement('label', { style: { fontSize: '12px', color: '#555' } }, 'Generate More'),
+									wp.element.createElement('select', {
+										value: extraVariantCounts[image.agent_image_id] || 1,
+										onChange: (e) => setExtraVariantCounts(prev => ({
+											...prev,
+											[image.agent_image_id]: parseInt(e.target.value, 10)
+										})),
+										disabled: generatingMoreFor !== null || committingVariantId !== null,
+										style: {
+											padding: '4px 6px',
+											borderRadius: '4px',
+											border: '1px solid #ddd',
+											minWidth: '60px'
+										}
+									},
+										wp.element.createElement('option', { value: '1' }, '+1'),
+										wp.element.createElement('option', { value: '2' }, '+2'),
+										wp.element.createElement('option', { value: '3' }, '+3')
+									),
+									wp.element.createElement(Button, {
+										variant: 'secondary',
+										onClick: () => handleGenerateMore(image.agent_image_id),
+										disabled: generatingMoreFor !== null || committingVariantId !== null,
+									}, generatingMoreFor === image.agent_image_id ? 'Generating...' : 'Generate More')
+								),
 								
 								wp.element.createElement('div', {
 									style: {
@@ -419,14 +689,15 @@
 											
 											wp.element.createElement('div', { style: { padding: '10px' } },
 												wp.element.createElement('p', { style: { fontSize: '12px', margin: '0 0 10px' } },
-													`Cost: $${variant.cost.toFixed(3)} • ${variant.generation_time}s`
+													`Cost: $${variant.cost.toFixed(3)} • ${Math.round(Number(variant.generation_time) || 0)}s`
 												),
 												
 												wp.element.createElement(Button, {
 													variant: 'primary',
 													onClick: () => handleSelectVariant(image.agent_image_id, variant.id),
+													disabled: committingVariantId !== null,
 													style: { width: '100%' },
-												}, 'Select')
+												}, committingVariantId === variant.id ? 'Selecting...' : 'Select')
 											)
 										)
 									)
@@ -450,6 +721,28 @@
 	// Initialize modal container and event listeners
 	let modalContainer = null;
 	let currentModalInstance = null;
+	let modalRoot = null;
+
+	const mountModal = (element) => {
+		if (typeof createRoot === 'function') {
+			if (!modalRoot) {
+				modalRoot = createRoot(modalContainer);
+			}
+			modalRoot.render(element);
+			return;
+		}
+		render(element, modalContainer);
+	};
+
+	const unmountModal = () => {
+		if (modalRoot && typeof modalRoot.unmount === 'function') {
+			modalRoot.unmount();
+			modalRoot = null;
+		} else if (modalContainer) {
+			render(null, modalContainer);
+		}
+		currentModalInstance = null;
+	};
 	
 	/**
 	 * Open image modal for review after article generation
@@ -463,24 +756,16 @@
 			document.body.appendChild(modalContainer);
 		}
 		
-		currentModalInstance = render(
+		currentModalInstance = wp.element.createElement(window.wpAgenticWriter.ImageReviewModal, {
-			wp.element.createElement(window.wpAgenticWriter.ImageReviewModal, {
 				postId: postId,
 				onClose: () => {
-					if (modalContainer) {
+					unmountModal();
-						render(null, modalContainer);
-						currentModalInstance = null;
-					}
 				},
 				onComplete: () => {
-					if (modalContainer) {
+					unmountModal();
-						render(null, modalContainer);
-						currentModalInstance = null;
-					}
 				},
-			}),
+			});
-			modalContainer
+		mountModal(currentModalInstance);
-		);
 	});
 	
 	/**
@@ -496,24 +781,17 @@
 			document.body.appendChild(modalContainer);
 		}
 		
-		currentModalInstance = render(
+		currentModalInstance = wp.element.createElement(window.wpAgenticWriter.ImageReviewModal, {
-			wp.element.createElement(window.wpAgenticWriter.ImageReviewModal, {
 				postId: postId,
 				initialImageId: agentImageId,
+				initialBlockId: blockId,
 				onClose: () => {
-					if (modalContainer) {
+					unmountModal();
-						render(null, modalContainer);
-						currentModalInstance = null;
-					}
 				},
 				onComplete: () => {
-					if (modalContainer) {
+					unmountModal();
-						render(null, modalContainer);
-						currentModalInstance = null;
-					}
 				},
-			}),
+			});
-			modalContainer
+		mountModal(currentModalInstance);
-		);
 	});
 })();

includes/class-admin-columns.php

@@ -77,16 +77,10 @@ class WP_Agentic_Writer_Admin_Columns {
 			return;
 		}
 
-		global $wpdb;
+		$total_cost = 0.0;
-		$table_name = $wpdb->prefix . 'wpaw_cost_tracking';
+		if ( class_exists( 'WP_Agentic_Writer_Cost_Tracker' ) ) {
-
+			$total_cost = WP_Agentic_Writer_Cost_Tracker::get_instance()->get_session_total( $post_id );
-		// Get total cost for this post
+		}
-		$total_cost = $wpdb->get_var(
-			$wpdb->prepare(
-				"SELECT SUM(cost) FROM {$table_name} WHERE post_id = %d",
-				$post_id
-			)
-		);
 
 		if ( $total_cost && $total_cost > 0 ) {
 			// Color code based on cost