worker.js

/**
 * AI Image Proxy - Cloudflare Worker v4.0
 * 
 * 将 AI 生成的 base64 图片转换为 URL，解决客户端堆栈溢出问题
 * 支持流式响应和历史图片上下文保留
 * 
 * @author Your Name
 * @license MIT
 */

export default {
    async fetch(request, env, ctx) {
        if (request.method === 'OPTIONS') {
            return new Response(null, {
                headers: {
                    'Access-Control-Allow-Origin': '*',
                    'Access-Control-Allow-Methods': 'POST, GET, OPTIONS',
                    'Access-Control-Allow-Headers': 'Content-Type, Authorization',
                },
            });
        }

        if (request.method !== 'POST') {
            return new Response('AI Image Proxy v4.0', { status: 200 });
        }

        function arrayBufferToBase64(buffer) {
            const bytes = new Uint8Array(buffer);
            let binary = '';
            const chunkSize = 8192;
            for (let i = 0; i < bytes.length; i += chunkSize) {
                const chunk = bytes.slice(i, i + chunkSize);
                binary += String.fromCharCode.apply(null, chunk);
            }
            return btoa(binary);
        }

        try {
            const requestBody = await request.json();

            const cleanBody = (obj) => {
                if (Array.isArray(obj)) return obj.map(cleanBody);
                if (obj && typeof obj === 'object') {
                    const cleaned = {};
                    for (const [key, value] of Object.entries(obj)) {
                        if (value !== "[undefined]" && value !== undefined && value !== null) {
                            cleaned[key] = cleanBody(value);
                        }
                    }
                    return cleaned;
                }
                return obj;
            };

            const cleanedBody = cleanBody(requestBody);
            cleanedBody.stream = true;

            // ✅ 修正：将 imageSize 移到 imageConfig 内部
            cleanedBody.providerOptions = {
                google: {
                    imageConfig: {
                        aspectRatio: "16:9",
                        imageSize: "2K"
                    }
                }
            };

            const R2_DOMAIN = env.R2_PUBLIC_URL.replace('https://', '').replace('http://', '');

            // 🔥 收集历史中所有生成的图片
            let historyImages = [];

            if (cleanedBody.messages) {
                // 第一遍：从 assistant 消息中提取图片
                for (const message of cleanedBody.messages) {
                    if (message.role === 'assistant' && typeof message.content === 'string') {
                        const imgRegex = /!\[.*?\]\((https?:\/\/[^)]+)\)/g;
                        const matches = [...message.content.matchAll(imgRegex)];

                        for (const match of matches) {
                            const imgUrl = match[1];
                            try {
                                const imgResponse = await fetch(imgUrl);
                                if (imgResponse.ok) {
                                    const imgBuffer = await imgResponse.arrayBuffer();
                                    const base64 = arrayBufferToBase64(imgBuffer);
                                    const contentType = imgResponse.headers.get('content-type') || 'image/png';
                                    historyImages.push({
                                        type: 'image_url',
                                        image_url: { url: `data:${contentType};base64,${base64}` }
                                    });
                                }
                            } catch (e) {
                                console.error('Failed to fetch image:', e);
                            }

                            // 从 assistant 消息中移除图片 markdown，保留文本
                            message.content = message.content.replace(match[0], '[之前生成的图片]');
                        }
                    }
                }

                // � 第二遍：如果有历史图片，添加到最后一条 user 消息中
                if (historyImages.length > 0) {
                    // 找到最后一条 user 消息
                    for (let i = cleanedBody.messages.length - 1; i >= 0; i--) {
                        const message = cleanedBody.messages[i];

                        if (message.role === 'user') {
                            // 获取原始文本内容
                            let textContent = '';
                            let existingImages = [];

                            if (typeof message.content === 'string') {
                                textContent = message.content;
                            } else if (Array.isArray(message.content)) {
                                for (const part of message.content) {
                                    if (part.type === 'text') {
                                        textContent += part.text;
                                    } else if (part.type === 'image_url') {
                                        existingImages.push(part);
                                    }
                                }
                            }

                            // 🔥 构建新的 content：历史图片 + 用户上传的图片 + 文本
                            const newContent = [
                                ...historyImages,           // 历史生成的图片
                                ...existingImages,          // 用户新上传的图片
                                { type: 'text', text: `请基于上面的图片进行修改：${textContent}` }
                            ];

                            cleanedBody.messages[i].content = newContent;
                            break;
                        }
                    }
                }
            }

            const aiResponse = await fetch('https://ai-gateway.vercel.sh/v1/chat/completions', {
                method: 'POST',
                headers: {
                    'Authorization': `Bearer ${env.AI_GATEWAY_KEY}`,
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify(cleanedBody),
            });

            if (!aiResponse.ok) {
                const errorText = await aiResponse.text();
                return new Response(JSON.stringify({
                    error: 'AI Gateway error',
                    status: aiResponse.status,
                    details: errorText.substring(0, 1000)
                }), {
                    status: 500,
                    headers: { 'Content-Type': 'application/json', 'Access-Control-Allow-Origin': '*' },
                });
            }

            const { readable, writable } = new TransformStream();
            const writer = writable.getWriter();
            const reader = aiResponse.body.getReader();
            const decoder = new TextDecoder();
            const encoder = new TextEncoder();

            ctx.waitUntil((async () => {
                let buffer = '';

                try {
                    while (true) {
                        const { done, value } = await reader.read();
                        if (done) {
                            await writer.close();
                            break;
                        }

                        buffer += decoder.decode(value, { stream: true });
                        const lines = buffer.split('\n');
                        buffer = lines.pop() || '';

                        for (const line of lines) {
                            if (!line.trim() || !line.startsWith('data: ')) {
                                await writer.write(encoder.encode(line + '\n'));
                                continue;
                            }

                            const jsonStr = line.slice(6);
                            if (jsonStr === '[DONE]') {
                                await writer.write(encoder.encode(line + '\n'));
                                continue;
                            }

                            try {
                                const data = JSON.parse(jsonStr);
                                const delta = data.choices?.[0]?.delta;

                                if (delta?.images && Array.isArray(delta.images)) {
                                    let imageMarkdown = '';

                                    for (const img of delta.images) {
                                        const imageUrl = img.image_url?.url || '';

                                        if (imageUrl.startsWith('data:image/')) {
                                            const matches = imageUrl.match(/^data:image\/(\w+);base64,(.+)$/);
                                            if (matches) {
                                                const imageType = matches[1];
                                                const base64Data = matches[2];

                                                const binaryString = atob(base64Data);
                                                const bytes = new Uint8Array(binaryString.length);
                                                for (let j = 0; j < binaryString.length; j++) {
                                                    bytes[j] = binaryString.charCodeAt(j);
                                                }

                                                const fileName = `output-${Date.now()}-${Math.random().toString(36).slice(2)}.${imageType}`;
                                                await env.R2_BUCKET.put(fileName, bytes, {
                                                    httpMetadata: { contentType: `image/${imageType}` },
                                                });

                                                const publicUrl = `${env.R2_PUBLIC_URL}/${fileName}`;
                                                imageMarkdown += `\n\n![image](${publicUrl})\n\n`;
                                            }
                                        }
                                    }

                                    delete delta.images;
                                    delta.content = (delta.content || '') + imageMarkdown;
                                }

                                await writer.write(encoder.encode(`data: ${JSON.stringify(data)}\n`));

                            } catch (e) {
                                await writer.write(encoder.encode(line + '\n'));
                            }
                        }
                    }
                } catch (e) {
                    await writer.close();
                }
            })());

            return new Response(readable, {
                headers: {
                    'Content-Type': 'text/event-stream',
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                    'Access-Control-Allow-Origin': '*',
                },
            });

        } catch (error) {
            return new Response(JSON.stringify({ error: error.message }), {
                status: 500,
                headers: { 'Content-Type': 'application/json', 'Access-Control-Allow-Origin': '*' },
            });
        }
    },
};