Greasy Fork

// ==UserScript==
// @name Google AI Studio 模型注入器 (优化版)
// @namespace http://tampermonkey.net/
// @version 1.8.0
// @description 向 Google AI Studio 注入自定义模型，性能优化版。支持手动添加ID（无需输入 models/）。拦截 XHR/Fetch 请求。
// @author Generated by AI / HCPTangHY / Mozi / wisdgod / UserModified / Optimized by AI
// @match https://aistudio.google.com/*
// @icon https://www.google.com/s2/favicons?sz=64&domain=aistudio.google.com
// @grant none
// @run-at document-start
// @license MIT
// @downloadURL https://update.greasyfork.org/scripts/539399/Google%20AI%20Studio%20%E6%A8%A1%E5%9E%8B%E6%B3%A8%E5%85%A5%E5%99%A8.user.js
// @updateURL https://update.greasyfork.org/scripts/539399/Google%20AI%20Studio%20%E6%A8%A1%E5%9E%8B%E6%B3%A8%E5%85%A5%E5%99%A8.meta.js
// ==/UserScript==

(function() {
'use strict';

// ==================== 配置区域 ====================
const SCRIPT_VERSION = "v1.8.0";
const CONFIG = Object.freeze({
VERSION: SCRIPT_VERSION,
LOG_PREFIX: `[AI Studio 注入器 ${SCRIPT_VERSION}]`,
ANTI_HIJACK_PREFIX: ")]}'\n",
STORAGE_KEY: 'AI_STUDIO_INJECTOR_CUSTOM_MODELS',
API_TARGET_PATH: '/ListModels',
API_TARGET_HOST: 'alkalimakersuite',
// 用于从响应数组中定位数据的字段索引
MODEL_FIELDS: {
NAME: 0,
DISPLAY_NAME: 3,
DESCRIPTION: 4,
METHODS: 7
}
});

// 日志工具
const Log = {
info: (...args) => console.log(CONFIG.LOG_PREFIX, ...args),
warn: (...args) => console.warn(CONFIG.LOG_PREFIX, ...args),
error: (...args) => console.error(CONFIG.LOG_PREFIX, ...args)
};

// 特殊“动作”模型定义
const ACTIONS = Object.freeze({
ADD: {
name: 'models/---script-action-add-custom---',
displayName: `➕ 添加自定义模型 (点击此处)`,
description: '点击此项以手动输入新的模型 ID 并保存。'
},
CLEAR: {
name: 'models/---script-action-clear-custom---',
displayName: `🗑️ 清除手动添加的模型 (点击此处)`,
description: '点击此项以清除所有您手动添加的模型。'
}
});

// 预设模型配置列表 (保持 name 唯一)
const PREDEFINED_MODELS = Object.freeze([
{
name: 'models/blacktooth-ab-test',
displayName: `🏴‍☠️ Blacktooth`,
description: `由脚本预设注入的模型`
},
{
name: 'models/jfdksal98a',
displayName: `🪐 jfdksal98a`,
description: `由脚本预设注入的模型`
},
{
name: 'models/gemini-2.5-pro-exp-03-25',
displayName: `✨ Gemini 2.5 Pro 03-25`,
description: `由脚本预设注入的模型`
},
{
name: 'models/goldmane-ab-test',
displayName: `🦁 Goldmane`,
description: `由脚本预设注入的模型`
},
{
name: 'models/claybrook-ab-test',
displayName: `💧 Claybrook`,
description: `由脚本预设注入的模型`
},
{
name: 'models/frostwind-ab-test',
displayName: `❄️ Frostwind`,
description: `由脚本预设注入的模型`
},
{
name: 'models/calmriver-ab-test',
displayName: `🌊 Calmriver`,
description: `由脚本预设注入的模型`
}
]);

// ==================== 模型管理 & UI 逻辑 ====================

// 从 localStorage 加载并清洗自定义模型
function getStoredCustomModels() {
const storedData = localStorage.getItem(CONFIG.STORAGE_KEY);
if (!storedData) return [];
try {
const parsedModels = JSON.parse(storedData);
if (!Array.isArray(parsedModels)) return [];
// 过滤无效数据，确保包含必需字段
return parsedModels.filter(m => m && typeof m === 'object' && typeof m.name === 'string' && m.name.startsWith('models/'));
} catch (e) {
Log.error('解析存储的自定义模型失败:', e);
localStorage.removeItem(CONFIG.STORAGE_KEY); // 清除损坏的数据
return [];
}
}

function saveCustomModels(models) {
try {
localStorage.setItem(CONFIG.STORAGE_KEY, JSON.stringify(models));
} catch (e) {
Log.error('保存自定义模型失败:', e);
}
}

// 给模型名称和描述添加脚本版本后缀
function formatModelData(model) {
return {
...model,
// 清除旧的版本号，添加新的
displayName: model.displayName.replace(/ \(脚本 v[\d.]+\)$/, '') + ` (脚本 ${CONFIG.VERSION})`,
description: model.description.replace(/脚本 v[\d.]+/, '') + ` (脚本 ${CONFIG.VERSION})`
};
}

// 提示用户添加新模型
function promptForNewModel(allKnownModelNames) {
const rawInputId = prompt("【添加自定义模型】\n请输入新的模型ID (例如: my-custom-model):\n(会自动添加 'models/' 前缀)");
if (!rawInputId || !rawInputId.trim()) return;

const trimmedId = rawInputId.trim();
const fullModelId = trimmedId.startsWith('models/') ? trimmedId : 'models/' + trimmedId;
const modelName = fullModelId.split('/').pop();

// 性能优化：使用 Set 进行存在性检查
if (allKnownModelNames.has(fullModelId)) {
alert(`错误：模型ID ${fullModelId} 已存在于预设或自定义列表中。`);
return;
}

const displayNameInput = prompt("请输入该模型的显示名称 (例如: 🤖 我的模型):", `🤖 ${modelName}`);
if (!displayNameInput || !displayNameInput.trim()) return;

const newModel = {
name: fullModelId,
displayName: displayNameInput.trim(),
description: `由用户手动添加并通过脚本注入的模型`
};

const currentCustomModels = getStoredCustomModels();
currentCustomModels.push(newModel); // 存储时不带版本号，加载时再加
saveCustomModels(currentCustomModels);

alert(`模型 ${fullModelId} 添加成功！\n\n页面将自动刷新以应用更改。`);
window.location.reload();
}

// 清除自定义模型
function clearAllCustomModels() {
if (confirm("⚠️ 确定要清除所有您手动添加的自定义模型吗？\n\n(脚本预设的模型不会被删除)")) {
saveCustomModels([]);
alert("所有手动添加的自定义模型已清除。页面将自动刷新。");
window.location.reload();
}
}

// ==================== 初始化与全局变量 ====================

const CUSTOM_MODELS = getStoredCustomModels();
// 预先格式化好所有需要注入的模型（带版本号）
const ALL_MODELS_TO_INJECT = [
...PREDEFINED_MODELS.map(formatModelData),
...CUSTOM_MODELS.map(formatModelData)
];

// 存储所有已知模型名称的 Set，用于快速查重 (O(1))
const KNOWN_MODEL_NAMES = new Set(ALL_MODELS_TO_INJECT.map(m => m.name));

Log.info(`预设模型: ${PREDEFINED_MODELS.length}, 自定义模型: ${CUSTOM_MODELS.length}. 脚本已激活。`);

// ==================== 页面交互拦截 ====================

// 拦截模型列表的点击事件，用于触发自定义动作
function setupActionInterceptor() {
// 使用捕获阶段，以便在 Angular 等框架处理之前拦截事件
document.body.addEventListener('click', (event) => {
const optionElement = event.target.closest('[role="option"], mat-option');
if (!optionElement || !optionElement.textContent) return;

const text = optionElement.textContent;
let action = null;

if (text.includes(ACTIONS.ADD.displayName)) {
Log.info("触发 '添加模型' 动作。");
action = () => promptForNewModel(KNOWN_MODEL_NAMES);
} else if (text.includes(ACTIONS.CLEAR.displayName)) {
Log.info("触发 '清除模型' 动作。");
action = clearAllCustomModels;
}

if (action) {
event.preventDefault();
event.stopPropagation();
// 延迟执行以避免 UI 状态冲突（例如下拉菜单关闭）
setTimeout(action, 100);
// 尝试关闭打开的下拉菜单
if (document.activeElement) document.activeElement.blur();
}
}, true);
Log.info('UI 点击拦截器已设置。');
}

// 在 DOMContentLoaded 时设置 UI 拦截器
if (document.readyState === 'loading') {
document.addEventListener('DOMContentLoaded', setupActionInterceptor);
} else {
// 如果 document-start 运行时 DOM 已经准备好，则立即执行
setupActionInterceptor();
}

// ==================== 数据处理工具 ====================

// 判断是否为目标请求 URL
function isTargetURL(url) {
return url && url.includes(CONFIG.API_TARGET_HOST) && url.includes(CONFIG.API_TARGET_PATH);
}

// 递归查找响应体中的模型列表数组
// 优化点：由于 Google API 返回的数据结构通常是固定的，深度递归可能是多余的，
// 但保留它可以增加脚本对微小结构变化的鲁棒性。
function findModelListArray(obj) {
if (!obj || typeof obj !== 'object') return null;

// 检查是否为模型列表
if (Array.isArray(obj) && obj.length > 0 && obj.every(isModelEntry)) {
return obj;
}

// 递归遍历对象属性
for (const key in obj) {
if (Object.hasOwnProperty.call(obj, key)) {
const result = findModelListArray(obj[key]);
if (result) return result;
}
}
return null;
}

// 判断一个数组项是否符合模型数据的结构
function isModelEntry(item) {
return Array.isArray(item) &&
typeof item[CONFIG.MODEL_FIELDS.NAME] === 'string' &&
item[CONFIG.MODEL_FIELDS.NAME].startsWith('models/');
}

// 找到一个用于克隆的模板模型
function findTemplateModel(modelsArray) {
// 优先选择带有 "generateContent" 等方法的模型作为模板
const template = modelsArray.find(m =>
isModelEntry(m) &&
Array.isArray(m[CONFIG.MODEL_FIELDS.METHODS]) &&
m[CONFIG.MODEL_FIELDS.METHODS].includes('generateContent')
);
// 否则，选择第一个找到的模型
return template || modelsArray.find(isModelEntry);
}

// 根据模板创建新的模型数据项
function createModelEntry(templateModel, modelInfo, templateName) {
// 使用 structuredClone 进行深拷贝
const newEntry = structuredClone(templateModel);
newEntry[CONFIG.MODEL_FIELDS.NAME] = modelInfo.name;
newEntry[CONFIG.MODEL_FIELDS.DISPLAY_NAME] = modelInfo.displayName;
newEntry[CONFIG.MODEL_FIELDS.DESCRIPTION] = `${modelInfo.description} (基于 ${templateName} 结构)`;

// 如果是“动作”模型，清空方法列表，让它不可用
if (modelInfo.name.includes('---script-action-')) {
newEntry[CONFIG.MODEL_FIELDS.METHODS] = [];
} else if (!Array.isArray(newEntry[CONFIG.MODEL_FIELDS.METHODS])) {
// 如果模板的方法列表缺失，给予默认值
newEntry[CONFIG.MODEL_FIELDS.METHODS] = ["generateContent", "countTokens", "createCachedContent", "batchGenerateContent"];
}
return newEntry;
}

// ==================== 核心数据处理 ====================

function processJsonData(jsonData) {
const modelsArray = findModelListArray(jsonData);
if (!modelsArray) return { data: jsonData, modified: false };

const templateModel = findTemplateModel(modelsArray);
if (!templateModel) {
Log.warn('无法找到用于克隆的模板模型结构。');
return { data: jsonData, modified: false };
}
const templateName = templateModel[CONFIG.MODEL_FIELDS.NAME] || 'unknown';

// --- 性能优化点：使用 Map 存储现有模型以便 O(1) 查找 ---
// Key: model.name, Value: model entry array
const existingModelsMap = new Map(modelsArray.map(m => [m[CONFIG.MODEL_FIELDS.NAME], m]));
const modelsToAdd = [];
let modificationCount = 0;

// 1. 处理预设和自定义模型
ALL_MODELS_TO_INJECT.forEach(modelToInject => {
const existingEntry = existingModelsMap.get(modelToInject.name);

if (existingEntry) {
// 模型已存在于原始列表中
const currentDisplayName = existingEntry[CONFIG.MODEL_FIELDS.DISPLAY_NAME];
// 如果显示名称与当前脚本版本要求的不一致（例如脚本升级了），则更新它
if (currentDisplayName !== modelToInject.displayName) {
existingEntry[CONFIG.MODEL_FIELDS.DISPLAY_NAME] = modelToInject.displayName;
existingEntry[CONFIG.MODEL_FIELDS.DESCRIPTION] = `${modelToInject.description} (基于 ${templateName} 结构)`;
modificationCount++;
}
} else {
// 模型不存在，准备注入
modelsToAdd.push(createModelEntry(templateModel, modelToInject, templateName));
modificationCount++;
}
});

// 2. 添加“动作”模型
const actionsToAdd = [];
actionsToAdd.push(createModelEntry(templateModel, ACTIONS.ADD, templateName));
if (CUSTOM_MODELS.length > 0) {
actionsToAdd.push(createModelEntry(templateModel, ACTIONS.CLEAR, templateName));
}
modificationCount += actionsToAdd.length;

// 3. 将新模型和动作添加到列表开头
if (modificationCount > 0) {
modelsArray.unshift(...modelsToAdd, ...actionsToAdd);
Log.info(`模型列表已修改，共注入/更新了 ${modificationCount} 项。`);
return { data: jsonData, modified: true };
}

return { data: jsonData, modified: false };
}

// 处理 API 响应体
function modifyResponseBody(originalText, url) {
if (!originalText || typeof originalText !== 'string') return originalText;

let textBody = originalText;
const hasPrefix = textBody.startsWith(CONFIG.ANTI_HIJACK_PREFIX);
if (hasPrefix) {
textBody = textBody.substring(CONFIG.ANTI_HIJACK_PREFIX.length);
}

if (!textBody.trim()) return originalText;

try {
const jsonData = JSON.parse(textBody);
const { data, modified } = processJsonData(jsonData, url);

if (modified) {
let newBody = JSON.stringify(data);
return hasPrefix ? CONFIG.ANTI_HIJACK_PREFIX + newBody : newBody;
}
} catch (error) {
Log.error('解析或处理响应体失败:', url, error);
}
return originalText; // 如果未修改或出错，返回原始文本
}

// ==================== 请求拦截 ====================

// 拦截 Fetch API
const originalFetch = window.fetch;
window.fetch = async function(resource, init) {
const url = (resource instanceof Request) ? resource.url : String(resource);
const response = await originalFetch.apply(this, arguments);

if (isTargetURL(url) && response.ok) {
try {
// 克隆响应以读取原始数据，因为 response.body 只能读取一次
const originalText = await response.clone().text();
const newBody = modifyResponseBody(originalText, url);

if (newBody !== originalText) {
// 创建并返回修改后的响应
return new Response(newBody, {
status: response.status,
statusText: response.statusText,
headers: response.headers
});
}
} catch (e) {
Log.error('[Fetch] 拦截处理错误:', e);
}
}
// 返回原始响应
return response;
};

// 拦截 XMLHttpRequest (XHR)
const xhrProto = XMLHttpRequest.prototype;
const originalXHRopen = xhrProto.open;
// 获取原始 responseText 和 response 的属性描述符
const originalResponseTextDesc = Object.getOwnPropertyDescriptor(xhrProto, 'responseText');
const originalResponseDesc = Object.getOwnPropertyDescriptor(xhrProto, 'response');

// 覆盖 open 方法，记录 URL
xhrProto.open = function(method, url) {
this._injectorIsTarget = isTargetURL(url);
this._injectorUrl = url;
return originalXHRopen.apply(this, arguments);
};

// 优化点：为 XHR 增加缓存，避免重复计算 modifyResponseBody
function getModifiedXHRResponse(xhr, originalValue, type) {
if (!xhr._injectorIsTarget || xhr.readyState !== 4 || xhr.status !== 200) {
return originalValue;
}

// 使用缓存
if (!xhr._injectorResponseCache) {
const originalText = (typeof originalValue === 'string') ? originalValue : originalResponseTextDesc.get.call(xhr);
xhr._injectorResponseCache = modifyResponseBody(originalText, xhr._injectorUrl);
}

const modifiedText = xhr._injectorResponseCache;

if (type === 'json') {
try {
const cleanJson = modifiedText.replace(CONFIG.ANTI_HIJACK_PREFIX, '');
return JSON.parse(cleanJson || "null");
} catch(e) {
Log.error('[XHR] 无法将修改后的文本解析为 JSON，返回原始内容:', e);
return originalValue;
}
}
return modifiedText;
}

// 覆盖 responseText 属性
if (originalResponseTextDesc && originalResponseTextDesc.get) {
Object.defineProperty(xhrProto, 'responseText', {
get() {
const originalText = originalResponseTextDesc.get.call(this);
// 只有当 responseType 为 "" 或 "text" 时才尝试修改
if (this.responseType === '' || this.responseType === 'text') {
return getModifiedXHRResponse(this, originalText, 'text');
}
return originalText;
},
configurable: true
});
}

// 覆盖 response 属性 (支持 responseType="json")
if (originalResponseDesc && originalResponseDesc.get) {
Object.defineProperty(xhrProto, 'response', {
get() {
const originalResponse = originalResponseDesc.get.call(this);
if (this.responseType === 'json') {
return getModifiedXHRResponse(this, originalResponse, 'json');
}
if (this.responseType === '' || this.responseType === 'text') {
return getModifiedXHRResponse(this, originalResponse, 'text');
}
return originalResponse;
},
configurable: true
});
}

Log.info('Fetch 和 XHR 拦截已启用。');
})();