fix: harden workflow smoke issues

This commit is contained in:
boos
2026-06-06 03:23:39 +08:00
parent e72e323313
commit e4630c4cb6
9 changed files with 32 additions and 51 deletions
+2
View File
@@ -368,6 +368,8 @@ async fn status_handler(
"keyName": key_name,
"userId": auth.user_id,
"role": auth.role,
"apiBase": upstream_base,
"upstreamBase": upstream_base,
"upstreamLabel": upstream_label,
"comfyui": state.comfyui.status_response().await,
})))
+1
View File
@@ -1570,6 +1570,7 @@ mod tests {
site_id INTEGER NOT NULL,
access_token TEXT NOT NULL,
api_token TEXT,
unit_cost REAL,
status TEXT NOT NULL DEFAULT 'active'
)",
)
+21 -46
View File
@@ -5183,55 +5183,30 @@ async fn workflow_context(state: &AppState, caller: &WorkflowCaller) -> AppResul
}
fn inject_bootstrap(html: String, config: &crate::config::ComfyUIConfig) -> String {
let cloud_api_key = workflow_cloud_api_key(config);
if config.uses_native_backend() && !config.uses_cloud_backend() {
return inject_external_bootstrap(html, cloud_api_key);
return inject_external_bootstrap(html);
}
inject_enhanced_bootstrap(html, cloud_api_key)
inject_enhanced_bootstrap(html)
}
fn workflow_api_key_bootstrap(api_key: Option<&str>) -> String {
let Some(api_key) = api_key else {
return String::new();
};
let api_key = serde_json::to_string(api_key).unwrap_or_else(|_| "null".to_string());
format!(
r#"<script>
(function () {{
const apiKey = {api_key};
if (!apiKey) return;
const installComfyApiKey = () => {{
try {{
localStorage.setItem('comfy_api_key', apiKey);
sessionStorage.setItem('comfy_api_key', apiKey);
for (const api of [window.app?.api, window.comfyAPI?.api, window.api]) {{
if (api) {{
api.apiKey = api.apiKey || apiKey;
if (!api.__boosapiCloudApiKeyQueuePatch && typeof api.queuePrompt === 'function') {{
const originalQueuePrompt = api.queuePrompt.bind(api);
api.queuePrompt = function() {{
if (!this.apiKey) this.apiKey = apiKey;
return originalQueuePrompt(...arguments);
}};
api.__boosapiCloudApiKeyQueuePatch = true;
}}
}}
}}
document.cookie = 'comfy_api_key=' + encodeURIComponent(apiKey) + '; Path=/workflow-ui; SameSite=Lax';
}} catch (err) {{}}
}};
installComfyApiKey();
window.addEventListener('DOMContentLoaded', installComfyApiKey);
window.setTimeout(installComfyApiKey, 300);
window.setTimeout(installComfyApiKey, 1200);
window.setInterval(installComfyApiKey, 5000);
}})();
fn workflow_secret_cleanup_bootstrap() -> &'static str {
r#"<script>
(function () {
const clearLeakedComfyApiKey = () => {
try {
localStorage.removeItem('comfy_api_key');
sessionStorage.removeItem('comfy_api_key');
document.cookie = 'comfy_api_key=; Path=/workflow-ui; Max-Age=0; SameSite=Lax';
} catch (err) {}
};
clearLeakedComfyApiKey();
window.addEventListener('DOMContentLoaded', clearLeakedComfyApiKey);
})();
</script>"#
)
}
fn inject_external_bootstrap(html: String, cloud_api_key: Option<&str>) -> String {
let api_key_bootstrap = workflow_api_key_bootstrap(cloud_api_key);
fn inject_external_bootstrap(html: String) -> String {
let secret_cleanup_bootstrap = workflow_secret_cleanup_bootstrap();
let shell = r#"<script>
(function () {
const renameComfyUiBrand = () => {
@@ -5310,12 +5285,12 @@ fn inject_external_bootstrap(html: String, cloud_api_key: Option<&str>) -> Strin
window.setTimeout(installWorkflowValidationButton, 1800);
})();
</script>"#;
let bootstrap = format!("{api_key_bootstrap}{shell}");
let bootstrap = format!("{secret_cleanup_bootstrap}{shell}");
html.replace("</head>", &format!("{bootstrap}\n</head>"))
}
fn inject_enhanced_bootstrap(html: String, cloud_api_key: Option<&str>) -> String {
let api_key_bootstrap = workflow_api_key_bootstrap(cloud_api_key);
fn inject_enhanced_bootstrap(html: String) -> String {
let secret_cleanup_bootstrap = workflow_secret_cleanup_bootstrap();
let shell = r#"<style id="boosapi-workflow-shell-style">
</style>
<script>
@@ -5713,7 +5688,7 @@ fn inject_enhanced_bootstrap(html: String, cloud_api_key: Option<&str>) -> Strin
window.setTimeout(patchWorkflowShell, 1200);
})();
</script>"#;
let bootstrap = format!("{api_key_bootstrap}{shell}");
let bootstrap = format!("{secret_cleanup_bootstrap}{shell}");
html.replace("</head>", &format!("{bootstrap}\n</head>"))
}