Allow agent to control light/dark color theme, and display screenshots in the chat

apps/cli/ai/runtimes/pi/index.ts

@@ -37,7 +37,7 @@ import {
 } from '@studio/common/lib/site-runtime';
 import { getAiPayloadsPath, getConfigDirectory } from '@studio/common/lib/well-known-paths';
 import { buildSystemPrompt } from 'cli/ai/system-prompt';
-import { resolveStudioToolDefinitions } from 'cli/ai/tools';
+import { resolveStudioToolDefinitions, withChatArtifactEmission } from 'cli/ai/tools';
 import { createAskUserQuestionTool } from 'cli/ai/tools/ask-user-question';
 import { createSiteTool } from 'cli/ai/tools/create-site';
 import { pullSiteTool } from 'cli/ai/tools/pull-site';
@@ -527,11 +527,12 @@ function buildAgentTools(
 	];
 
 	if ( isRemoteSite ) {
+		const remoteStudioTools = [ takeScreenshotTool, createSiteTool, pullSiteTool ].map( ( tool ) =>
+			withChatArtifactEmission( tool, chatArtifactsEnabled )
+		);
 		return [
 			createWpcomRequestTool( config.wpcomAccessToken!, config.activeSite!.wpcomSiteId! ),
-			takeScreenshotTool,
-			createSiteTool,
-			pullSiteTool,
+			...remoteStudioTools,
 			...remoteScratchTools,
 			...askUserTool,
 			...skillTool,

apps/cli/ai/screenshot-storage.ts

@@ -0,0 +1,38 @@
+import { mkdir, mkdtemp } from 'fs/promises';
+import os from 'os';
+import path from 'path';
+
+/**
+ * Ambient provider for the directory screenshot captures are saved to,
+ * mirroring `setChatArtifactCallback`. The `studio code` command points it at
+ * a per-session sidecar directory (`<session>.screenshots/` next to the
+ * session JSONL) so screenshots referenced by persisted chat artifacts
+ * survive OS temp cleanup and are removed together with the session.
+ * Standalone CLI/MCP runs never set a provider and fall back to a throwaway
+ * temp directory.
+ */
+type ScreenshotDirectoryProvider = () => Promise< string | null > | string | null;
+
+let screenshotDirectoryProvider: ScreenshotDirectoryProvider | null = null;
+
+export function setScreenshotDirectoryProvider(
+	provider: ScreenshotDirectoryProvider | null
+): void {
+	screenshotDirectoryProvider = provider;
+}
+
+export async function resolveScreenshotDirectory(): Promise< string > {
+	try {
+		const directory = await screenshotDirectoryProvider?.();
+		if ( directory ) {
+			await mkdir( directory, { recursive: true } );
+			return directory;
+		}
+	} catch ( error ) {
+		// Fall through to the temp directory so the capture still succeeds, but
+		// say so: files there are purged by the OS, so persisted artifacts will
+		// eventually show as unavailable and this line is the only clue why.
+		console.warn( '[screenshots] falling back to a temporary directory:', error );
+	}
+	return mkdtemp( path.join( os.tmpdir(), 'studio-screenshot-' ) );
+}

apps/cli/ai/system-prompt.ts

@@ -39,6 +39,7 @@ ${ REMOTE_DESIGN_GUIDELINES }${ remoteSessionAddendum }
 
 	return `${ buildLocalIntro( {
 		chatArtifactsEnabled: options?.chatArtifactsEnabled ?? false,
+		remoteSession: options?.remoteSession ?? false,
 		runtime: options?.runtime,
 	} ) }
 
@@ -97,9 +98,20 @@ function getPostContentGuidance( runtime?: SiteRuntime ): string {
 
 function buildLocalIntro( options: {
 	chatArtifactsEnabled: boolean;
+	remoteSession: boolean;
 	runtime?: SiteRuntime;
 } ): string {
 	const postContentGuidance = getPostContentGuidance( options.runtime );
+	// Remote-bridge sessions also run without chat artifacts, but their user is
+	// on the other end of a messaging bridge: local file paths are unreachable
+	// and REMOTE_SESSION_GUIDANCE (share_screenshot) already covers delivery.
+	const terminalScreenshotSection = options.remoteSession
+		? ''
+		: `
+
+## Screenshots
+
+This session runs in a terminal, which may not be able to display images. Screenshots you capture are for your own visual verification; the user may only see a link to the saved image file in the transcript. Do not respond as though the user is looking at the capture (e.g. "Here's your site!") — instead, state what you verified and describe notable findings, and point to the saved screenshot file when it helps.`;
 	const automaticArtifactSection = options.chatArtifactsEnabled
 		? `
 
@@ -114,7 +126,7 @@ ${ getStudioPresentationRulesPrompt() }
 
 Available desks widget types:
 ${ getStudioWidgetPromptManifest() }`
-		: '';
+		: terminalScreenshotSection;
 	const studioPresentToolBullet = options.chatArtifactsEnabled
 		? `
 - studio_present: Show one or more Studio desks widgets as inline visual artifacts.`

apps/cli/ai/tests/screenshot-helpers.test.ts

@@ -1,22 +1,77 @@
-import { readFile, rm } from 'node:fs/promises';
+import { mkdtemp, readFile, readdir, rm } from 'node:fs/promises';
 import os from 'node:os';
 import path from 'node:path';
-import { describe, expect, it } from 'vitest';
-import { saveScreenshotToTempFile } from '../tools/screenshot-helpers';
+import { afterEach, describe, expect, it, vi } from 'vitest';
+import { setScreenshotDirectoryProvider } from '../screenshot-storage';
+import { saveScreenshotFile } from '../tools/screenshot-helpers';
 
 describe( 'screenshot helpers', () => {
-	it( 'saves screenshots to a temporary local file', async () => {
+	afterEach( () => {
+		setScreenshotDirectoryProvider( null );
+	} );
+
+	it( 'falls back to a temporary directory when no provider is set', async () => {
 		const buffer = Buffer.from( 'not-really-a-png' );
-		const result = await saveScreenshotToTempFile( buffer, { viewportType: 'desktop' } );
+		const result = await saveScreenshotFile( buffer, { viewportType: 'desktop' } );
 
 		try {
 			expect( result.path.startsWith( os.tmpdir() ) ).toBe( true );
 			expect( result.fileUrl.startsWith( 'file://' ) ).toBe( true );
-			expect( result.name ).toBe( 'screenshot-desktop.png' );
+			expect( result.name ).toMatch( /^screenshot-desktop-[0-9a-f]{8}\.png$/ );
 			expect( result.mimeType ).toBe( 'image/png' );
 			await expect( readFile( result.path ) ).resolves.toEqual( buffer );
 		} finally {
 			await rm( path.dirname( result.path ), { recursive: true, force: true } );
 		}
 	} );
+
+	it( 'saves captures into the provided session directory with unique names', async () => {
+		const sessionRoot = await mkdtemp( path.join( os.tmpdir(), 'studio-session-' ) );
+		const screenshotsDirectory = path.join( sessionRoot, 'session.screenshots' );
+		setScreenshotDirectoryProvider( () => screenshotsDirectory );
+
+		try {
+			const first = await saveScreenshotFile( Buffer.from( 'first' ), {
+				viewportType: 'desktop',
+				format: 'jpeg',
+				colorScheme: 'dark',
+			} );
+			const second = await saveScreenshotFile( Buffer.from( 'second' ), {
+				viewportType: 'desktop',
+				format: 'jpeg',
+				colorScheme: 'dark',
+			} );
+
+			expect( path.dirname( first.path ) ).toBe( screenshotsDirectory );
+			expect( first.name ).toMatch( /^screenshot-desktop-dark-[0-9a-f]{8}\.jpg$/ );
+			expect( first.name ).not.toBe( second.name );
+			await expect( readFile( first.path, 'utf8' ) ).resolves.toBe( 'first' );
+			await expect( readFile( second.path, 'utf8' ) ).resolves.toBe( 'second' );
+			await expect( readdir( screenshotsDirectory ) ).resolves.toHaveLength( 2 );
+		} finally {
+			await rm( sessionRoot, { recursive: true, force: true } );
+		}
+	} );
+
+	it( 'falls back to a temporary directory when the provider throws, and says so', async () => {
+		const warnSpy = vi.spyOn( console, 'warn' ).mockImplementation( () => {} );
+		setScreenshotDirectoryProvider( () => {
+			throw new Error( 'no session' );
+		} );
+
+		const result = await saveScreenshotFile( Buffer.from( 'fallback' ), {
+			viewportType: 'mobile',
+		} );
+
+		try {
+			expect( result.path.startsWith( os.tmpdir() ) ).toBe( true );
+			expect( warnSpy ).toHaveBeenCalledWith(
+				expect.stringContaining( '[screenshots] falling back to a temporary directory' ),
+				expect.any( Error )
+			);
+		} finally {
+			warnSpy.mockRestore();
+			await rm( path.dirname( result.path ), { recursive: true, force: true } );
+		}
+	} );
 } );

apps/cli/ai/tests/system-prompt.test.ts

@@ -41,9 +41,8 @@ describe( 'buildSystemPrompt', () => {
 		expect( prompt ).toContain( 'For generated SVGs, write a complete .svg file' );
 		expect( prompt ).toContain( 'Do not present generated SVG code as a drawing widget' );
 		expect( prompt ).not.toContain( '- drawing:' );
-		expect( prompt ).toContain( '- screenshot-local-media:' );
-		expect( prompt ).toContain( 'present the actual captured PNG' );
-		expect( prompt ).toContain( 'Do not substitute a site-preview widget for a screenshot' );
+		expect( prompt ).toContain( '- screenshot-auto-artifact:' );
+		expect( prompt ).toContain( 'Never call studio_present for a screenshot' );
 		expect( prompt ).toContain( 'site-preview is for live previews, not captured screenshots' );
 		expect( prompt ).toContain( '- theme:' );
 		expect( prompt ).toContain( '- theme-template:' );
@@ -147,7 +146,30 @@ describe( 'buildSystemPrompt', () => {
 		expect( prompt ).not.toContain( '## Visual artifacts' );
 		expect( prompt ).not.toContain( '- site-code-scratchpad:' );
 		expect( prompt ).not.toContain( '- saved-local-media:' );
-		expect( prompt ).not.toContain( '- screenshot-local-media:' );
+		expect( prompt ).not.toContain( '- screenshot-auto-artifact:' );
 		expect( prompt ).not.toContain( 'studio_present' );
 	} );
+
+	it( 'warns that terminal users may not see screenshots when chat artifacts are disabled', () => {
+		const prompt = buildSystemPrompt( { chatArtifactsEnabled: false } );
+
+		expect( prompt ).toContain( '## Screenshots' );
+		expect( prompt ).toContain( 'Do not respond as though the user is looking at the capture' );
+	} );
+
+	it( 'omits the terminal screenshot caveat when chat artifacts are enabled', () => {
+		const prompt = buildSystemPrompt( { chatArtifactsEnabled: true } );
+
+		expect( prompt ).not.toContain( 'Do not respond as though the user is looking at the capture' );
+	} );
+
+	it( 'omits the terminal screenshot caveat for remote-bridge sessions', () => {
+		// The Telegram user cannot open local file paths; delivery is covered
+		// by the remote-session share_screenshot guidance instead.
+		const prompt = buildSystemPrompt( { chatArtifactsEnabled: false, remoteSession: true } );
+
+		expect( prompt ).not.toContain( '## Screenshots' );
+		expect( prompt ).not.toContain( 'Do not respond as though the user is looking at the capture' );
+		expect( prompt ).toContain( '## Telegram remote session' );
+	} );
 } );