feat: 添加资源缓存开关、自动刷新AT fix: 修复邀请码复制、修复移动端无法滑动、优化图片、视频格式输出

This commit is contained in:
TheSmallHanCat
2025-11-09 18:26:37 +08:00
parent 7269e3fa79
commit b6cedb0ece
11 changed files with 382 additions and 113 deletions

View File

@@ -729,12 +729,49 @@ async def get_cache_config(token: str = Depends(verify_admin_token)):
return {
"success": True,
"config": {
"enabled": config.cache_enabled,
"timeout": config.cache_timeout,
"base_url": config.cache_base_url, # 返回实际配置的值,可能为空字符串
"effective_base_url": config.cache_base_url or f"http://{config.server_host}:{config.server_port}" # 实际生效的值
}
}
@router.post("/api/cache/enabled")
async def update_cache_enabled(
request: dict,
token: str = Depends(verify_admin_token)
):
"""Update cache enabled status"""
try:
enabled = request.get("enabled", True)
# Update config file
config_path = Path("config/setting.toml")
with open(config_path, "r", encoding="utf-8") as f:
config_data = toml.load(f)
if "cache" not in config_data:
config_data["cache"] = {}
config_data["cache"]["enabled"] = enabled
with open(config_path, "w", encoding="utf-8") as f:
toml.dump(config_data, f)
# Update in-memory config
config.set_cache_enabled(enabled)
# Reload config to ensure consistency
config.reload_config()
return {
"success": True,
"message": f"Cache {'enabled' if enabled else 'disabled'} successfully",
"enabled": enabled
}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to update cache enabled status: {str(e)}")
# Generation timeout config endpoints
@router.get("/api/generation/timeout")
async def get_generation_timeout(token: str = Depends(verify_admin_token)):
@@ -862,3 +899,53 @@ async def update_video_length_config(
raise
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to update video length config: {str(e)}")
# AT auto refresh config endpoints
@router.get("/api/token-refresh/config")
async def get_at_auto_refresh_config(token: str = Depends(verify_admin_token)):
"""Get AT auto refresh configuration"""
# Reload config from file to get latest values
config.reload_config()
return {
"success": True,
"config": {
"at_auto_refresh_enabled": config.at_auto_refresh_enabled
}
}
@router.post("/api/token-refresh/enabled")
async def update_at_auto_refresh_enabled(
request: dict,
token: str = Depends(verify_admin_token)
):
"""Update AT auto refresh enabled status"""
try:
enabled = request.get("enabled", False)
# Update config file
config_path = Path("config/setting.toml")
with open(config_path, "r", encoding="utf-8") as f:
config_data = toml.load(f)
if "token_refresh" not in config_data:
config_data["token_refresh"] = {}
config_data["token_refresh"]["at_auto_refresh_enabled"] = enabled
with open(config_path, "w", encoding="utf-8") as f:
toml.dump(config_data, f)
# Update in-memory config
config.set_at_auto_refresh_enabled(enabled)
# Reload config to ensure consistency
config.reload_config()
return {
"success": True,
"message": f"AT auto refresh {'enabled' if enabled else 'disabled'} successfully",
"enabled": enabled
}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to update AT auto refresh enabled status: {str(e)}")

View File

@@ -120,6 +120,17 @@ class Config:
self._config["cache"] = {}
self._config["cache"]["base_url"] = base_url
@property
def cache_enabled(self) -> bool:
"""Get cache enabled status"""
return self._config.get("cache", {}).get("enabled", True)
def set_cache_enabled(self, enabled: bool):
"""Set cache enabled status"""
if "cache" not in self._config:
self._config["cache"] = {}
self._config["cache"]["enabled"] = enabled
@property
def image_timeout(self) -> int:
"""Get image generation timeout in seconds"""
@@ -168,5 +179,16 @@ class Config:
"""Get custom parse server access token"""
return self._config.get("watermark_free", {}).get("custom_parse_token", "")
@property
def at_auto_refresh_enabled(self) -> bool:
"""Get AT auto refresh enabled status"""
return self._config.get("token_refresh", {}).get("at_auto_refresh_enabled", False)
def set_at_auto_refresh_enabled(self, enabled: bool):
"""Set AT auto refresh enabled/disabled"""
if "token_refresh" not in self._config:
self._config["token_refresh"] = {}
self._config["token_refresh"]["at_auto_refresh_enabled"] = enabled
# Global config instance
config = Config()

View File

@@ -229,6 +229,8 @@ class GenerationHandler:
start_time = time.time()
last_heartbeat_time = start_time # Track last heartbeat for image generation
heartbeat_interval = 10 # Send heartbeat every 10 seconds for image generation
last_status_output_time = start_time # Track last status output time for video generation
video_status_interval = 30 # Output status every 30 seconds for video generation
debug_logger.log_info(f"Starting task polling: task_id={task_id}, is_video={is_video}, timeout={timeout}s, max_attempts={max_attempts}")
@@ -275,16 +277,18 @@ class GenerationHandler:
else:
progress_pct = int(progress_pct * 100)
# Only yield progress update if it changed
if progress_pct != last_progress:
last_progress = progress_pct
status = task.get("status", "processing")
debug_logger.log_info(f"Task {task_id} progress: {progress_pct}% (status: {status})")
# Update last_progress for tracking
last_progress = progress_pct
status = task.get("status", "processing")
if stream:
yield self._format_stream_chunk(
reasoning_content=f"**Video Generation Progress**: {progress_pct}% ({status})\n"
)
# Output status every 30 seconds (not just when progress changes)
current_time = time.time()
if stream and (current_time - last_status_output_time >= video_status_interval):
last_status_output_time = current_time
debug_logger.log_info(f"Task {task_id} progress: {progress_pct}% (status: {status})")
yield self._format_stream_chunk(
reasoning_content=f"**Video Generation Progress**: {progress_pct}% ({status})\n"
)
break
# If task not found in pending tasks, it's completed - fetch from drafts
@@ -356,43 +360,51 @@ class GenerationHandler:
if stream:
yield self._format_stream_chunk(
reasoning_content=f"Video published successfully. Post ID: {post_id}\nNow caching watermark-free video...\n"
reasoning_content=f"Video published successfully. Post ID: {post_id}\nNow {'caching' if config.cache_enabled else 'preparing'} watermark-free video...\n"
)
# Cache watermark-free video
try:
cached_filename = await self.file_cache.download_and_cache(watermark_free_url, "video")
local_url = f"{self._get_base_url()}/tmp/{cached_filename}"
if stream:
yield self._format_stream_chunk(
reasoning_content="Watermark-free video cached successfully. Preparing final response...\n"
)
# Delete the published post after caching
# Cache watermark-free video (if cache enabled)
if config.cache_enabled:
try:
debug_logger.log_info(f"Deleting published post: {post_id}")
await self.sora_client.delete_post(post_id, token)
debug_logger.log_info(f"Published post deleted successfully: {post_id}")
cached_filename = await self.file_cache.download_and_cache(watermark_free_url, "video")
local_url = f"{self._get_base_url()}/tmp/{cached_filename}"
if stream:
yield self._format_stream_chunk(
reasoning_content="Published post deleted successfully.\n"
reasoning_content="Watermark-free video cached successfully. Preparing final response...\n"
)
except Exception as delete_error:
debug_logger.log_error(
error_message=f"Failed to delete published post {post_id}: {str(delete_error)}",
status_code=500,
response_text=str(delete_error)
)
# Delete the published post after caching
try:
debug_logger.log_info(f"Deleting published post: {post_id}")
await self.sora_client.delete_post(post_id, token)
debug_logger.log_info(f"Published post deleted successfully: {post_id}")
if stream:
yield self._format_stream_chunk(
reasoning_content="Published post deleted successfully.\n"
)
except Exception as delete_error:
debug_logger.log_error(
error_message=f"Failed to delete published post {post_id}: {str(delete_error)}",
status_code=500,
response_text=str(delete_error)
)
if stream:
yield self._format_stream_chunk(
reasoning_content=f"Warning: Failed to delete published post - {str(delete_error)}\n"
)
except Exception as cache_error:
# Fallback to watermark-free URL if caching fails
local_url = watermark_free_url
if stream:
yield self._format_stream_chunk(
reasoning_content=f"Warning: Failed to delete published post - {str(delete_error)}\n"
reasoning_content=f"Warning: Failed to cache file - {str(cache_error)}\nUsing original watermark-free URL instead...\n"
)
except Exception as cache_error:
# Fallback to watermark-free URL if caching fails
else:
# Cache disabled: use watermark-free URL directly
local_url = watermark_free_url
if stream:
yield self._format_stream_chunk(
reasoning_content=f"Warning: Failed to cache file - {str(cache_error)}\nUsing original watermark-free URL instead...\n"
reasoning_content="Cache is disabled. Using watermark-free URL directly...\n"
)
except Exception as publish_error:
@@ -410,34 +422,45 @@ class GenerationHandler:
url = item.get("downloadable_url") or item.get("url")
if not url:
raise Exception("Video URL not found")
try:
cached_filename = await self.file_cache.download_and_cache(url, "video")
local_url = f"{self._get_base_url()}/tmp/{cached_filename}"
except Exception as cache_error:
if config.cache_enabled:
try:
cached_filename = await self.file_cache.download_and_cache(url, "video")
local_url = f"{self._get_base_url()}/tmp/{cached_filename}"
except Exception as cache_error:
local_url = url
else:
local_url = url
else:
# Normal mode: use downloadable_url instead of url
url = item.get("downloadable_url") or item.get("url")
if url:
# Cache video file
if stream:
yield self._format_stream_chunk(
reasoning_content="**Video Generation Completed**\n\nVideo generation successful. Now caching the video file...\n"
)
try:
cached_filename = await self.file_cache.download_and_cache(url, "video")
local_url = f"{self._get_base_url()}/tmp/{cached_filename}"
# Cache video file (if cache enabled)
if config.cache_enabled:
if stream:
yield self._format_stream_chunk(
reasoning_content="Video file cached successfully. Preparing final response...\n"
reasoning_content="**Video Generation Completed**\n\nVideo generation successful. Now caching the video file...\n"
)
except Exception as cache_error:
# Fallback to original URL if caching fails
try:
cached_filename = await self.file_cache.download_and_cache(url, "video")
local_url = f"{self._get_base_url()}/tmp/{cached_filename}"
if stream:
yield self._format_stream_chunk(
reasoning_content="Video file cached successfully. Preparing final response...\n"
)
except Exception as cache_error:
# Fallback to original URL if caching fails
local_url = url
if stream:
yield self._format_stream_chunk(
reasoning_content=f"Warning: Failed to cache file - {str(cache_error)}\nUsing original URL instead...\n"
)
else:
# Cache disabled: use original URL directly
local_url = url
if stream:
yield self._format_stream_chunk(
reasoning_content=f"Warning: Failed to cache file - {str(cache_error)}\nUsing original URL instead...\n"
reasoning_content="**Video Generation Completed**\n\nCache is disabled. Using original URL directly...\n"
)
# Task completed
@@ -482,27 +505,37 @@ class GenerationHandler:
base_url = self._get_base_url()
local_urls = []
for idx, url in enumerate(urls):
try:
cached_filename = await self.file_cache.download_and_cache(url, "image")
local_url = f"{base_url}/tmp/{cached_filename}"
local_urls.append(local_url)
if stream and len(urls) > 1:
yield self._format_stream_chunk(
reasoning_content=f"Cached image {idx + 1}/{len(urls)}...\n"
)
except Exception as cache_error:
# Fallback to original URL if caching fails
local_urls.append(url)
if stream:
yield self._format_stream_chunk(
reasoning_content=f"Warning: Failed to cache image {idx + 1} - {str(cache_error)}\nUsing original URL instead...\n"
)
if stream and all(u.startswith(base_url) for u in local_urls):
yield self._format_stream_chunk(
reasoning_content="All images cached successfully. Preparing final response...\n"
)
# Check if cache is enabled
if config.cache_enabled:
for idx, url in enumerate(urls):
try:
cached_filename = await self.file_cache.download_and_cache(url, "image")
local_url = f"{base_url}/tmp/{cached_filename}"
local_urls.append(local_url)
if stream and len(urls) > 1:
yield self._format_stream_chunk(
reasoning_content=f"Cached image {idx + 1}/{len(urls)}...\n"
)
except Exception as cache_error:
# Fallback to original URL if caching fails
local_urls.append(url)
if stream:
yield self._format_stream_chunk(
reasoning_content=f"Warning: Failed to cache image {idx + 1} - {str(cache_error)}\nUsing original URL instead...\n"
)
if stream and all(u.startswith(base_url) for u in local_urls):
yield self._format_stream_chunk(
reasoning_content="All images cached successfully. Preparing final response...\n"
)
else:
# Cache disabled: use original URLs directly
local_urls = urls
if stream:
yield self._format_stream_chunk(
reasoning_content="Cache is disabled. Using original URLs directly...\n"
)
await self.db.update_task(
task_id, "completed", 100.0,
@@ -510,10 +543,10 @@ class GenerationHandler:
)
if stream:
# Final response with content
content_html = "".join([f"<img src='{url}' />" for url in local_urls])
# Final response with content (Markdown format)
content_markdown = "\n".join([f"![Generated Image]({url})" for url in local_urls])
yield self._format_stream_chunk(
content=content_html,
content=content_markdown,
finish_reason="STOP"
)
yield "data: [DONE]\n\n"
@@ -544,7 +577,7 @@ class GenerationHandler:
last_heartbeat_time = current_time
elapsed = int(current_time - start_time)
yield self._format_stream_chunk(
reasoning_content=f"**Generating**\n\nImage generation in progress... ({elapsed}s elapsed)\n"
reasoning_content=f"Image generation in progress... ({elapsed}s elapsed)\n"
)
# If task not found in response, send heartbeat for image generation
@@ -554,7 +587,7 @@ class GenerationHandler:
last_heartbeat_time = current_time
elapsed = int(current_time - start_time)
yield self._format_stream_chunk(
reasoning_content=f"**Generating**\n\nImage generation in progress... ({elapsed}s elapsed)\n"
reasoning_content=f"Image generation in progress... ({elapsed}s elapsed)\n"
)
# Progress update for stream mode (fallback if no status from API)
@@ -638,7 +671,7 @@ class GenerationHandler:
if media_type == "video":
content = f"```html\n<video src='{url}' controls></video>\n```"
else:
content = f"<img src='{url}' />"
content = f"![Generated Image]({url})"
response = {
"id": f"chatcmpl-{datetime.now().timestamp()}",

View File

@@ -25,6 +25,18 @@ class LoadBalancer:
Returns:
Selected token or None if no available tokens
"""
# Try to auto-refresh tokens expiring within 24 hours if enabled
if config.at_auto_refresh_enabled:
all_tokens = await self.token_manager.get_all_tokens()
for token in all_tokens:
if token.is_active and token.expiry_time:
from datetime import datetime
time_until_expiry = token.expiry_time - datetime.now()
hours_until_expiry = time_until_expiry.total_seconds() / 3600
# Refresh if expiry is within 24 hours
if hours_until_expiry <= 24:
await self.token_manager.auto_refresh_expiring_token(token.id)
active_tokens = await self.token_manager.get_active_tokens()
if not active_tokens:

View File

@@ -871,3 +871,77 @@ class TokenManager:
print(f"Failed to refresh Sora2 remaining count: {e}")
except Exception as e:
print(f"Error in refresh_sora2_remaining_if_cooldown_expired: {e}")
async def auto_refresh_expiring_token(self, token_id: int) -> bool:
"""
Auto refresh token when expiry time is within 24 hours using ST or RT
Returns:
True if refresh successful, False otherwise
"""
try:
token_data = await self.db.get_token(token_id)
if not token_data:
return False
# Check if token is expiring within 24 hours
if not token_data.expiry_time:
return False # No expiry time set
time_until_expiry = token_data.expiry_time - datetime.now()
hours_until_expiry = time_until_expiry.total_seconds() / 3600
# Only refresh if expiry is within 24 hours (1440 minutes)
if hours_until_expiry > 24:
return False # Token not expiring soon
if hours_until_expiry < 0:
# Token already expired, still try to refresh
print(f"🔄 Token {token_id} 已过期,尝试自动刷新...")
else:
print(f"🔄 Token {token_id} 将在 {hours_until_expiry:.1f} 小时后过期,尝试自动刷新...")
# Priority: ST > RT
new_at = None
new_st = None
new_rt = None
if token_data.st:
# Try to refresh using ST
try:
print(f"📝 使用 ST 刷新 Token {token_id}...")
result = await self.st_to_at(token_data.st)
new_at = result.get("access_token")
# ST refresh doesn't return new ST, so keep the old one
new_st = token_data.st
print(f"✅ 使用 ST 刷新成功")
except Exception as e:
print(f"❌ 使用 ST 刷新失败: {e}")
new_at = None
if not new_at and token_data.rt:
# Try to refresh using RT
try:
print(f"📝 使用 RT 刷新 Token {token_id}...")
result = await self.rt_to_at(token_data.rt)
new_at = result.get("access_token")
new_rt = result.get("refresh_token", token_data.rt) # RT might be updated
print(f"✅ 使用 RT 刷新成功")
except Exception as e:
print(f"❌ 使用 RT 刷新失败: {e}")
new_at = None
if new_at:
# Update token with new AT
await self.update_token(token_id, token=new_at, st=new_st, rt=new_rt)
print(f"✅ Token {token_id} 已自动刷新")
return True
else:
# No ST or RT, disable token
print(f"⚠️ Token {token_id} 无法刷新(无 ST 或 RT已禁用")
await self.disable_token(token_id)
return False
except Exception as e:
print(f"❌ 自动刷新 Token {token_id} 失败: {e}")
return False