From 6dd12425d6b4b987a057237f10e1b2f2c7b1cd29 Mon Sep 17 00:00:00 2001 From: wunianxiao Date: Tue, 2 Jun 2026 11:03:31 +0800 Subject: [PATCH] =?UTF-8?q?fix:=205=20=E4=B8=AA=E5=AE=89=E5=85=A8/?= =?UTF-8?q?=E6=AD=A3=E7=A1=AE=E6=80=A7=E9=97=AE=E9=A2=98=EF=BC=88=E7=AC=AC?= =?UTF-8?q?=E4=BA=8C=E8=BD=AE=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. (High) user_id NOT NULL 导致 add 任务失败 → ALTER TABLE 加 DEFAULT '', DROP NOT NULL 2. (High) 调度器任务卡住阻塞后续 → timeout(120s) + kill_on_drop 3. (Medium) 多实例重复执行 → SELECT ... FOR UPDATE SKIP LOCKED 4. (Medium) SQL 注入（字符串拼接） → psql -v 参数化 (:name, :cmd, :interval) 5. (Medium) 非工具对话 usage 未记录 → consume() 改为 &mut self，消费后再读 usage --- ...0601000008_fix_scheduled_tasks_user_id.sql | 3 +++ .../manage_scheduled_task/scripts/manage.sh | 15 ++++++----- src/llm/conversation.rs | 3 ++- src/main.rs | 6 +++-- src/scheduler.rs | 27 +++++++++++-------- 5 files changed, 33 insertions(+), 21 deletions(-) create mode 100644 migrations/20260601000008_fix_scheduled_tasks_user_id.sql diff --git a/migrations/20260601000008_fix_scheduled_tasks_user_id.sql b/migrations/20260601000008_fix_scheduled_tasks_user_id.sql new file mode 100644 index 0000000..526a07f --- /dev/null +++ b/migrations/20260601000008_fix_scheduled_tasks_user_id.sql @@ -0,0 +1,3 @@ +-- scheduled_tasks.user_id 默认值修复 +ALTER TABLE scheduled_tasks ALTER COLUMN user_id SET DEFAULT ''; +ALTER TABLE scheduled_tasks ALTER COLUMN user_id DROP NOT NULL; diff --git a/resources/skills/manage_scheduled_task/scripts/manage.sh b/resources/skills/manage_scheduled_task/scripts/manage.sh index 337c07b..d0e86af 100755 --- a/resources/skills/manage_scheduled_task/scripts/manage.sh +++ b/resources/skills/manage_scheduled_task/scripts/manage.sh @@ -1,5 +1,5 @@ #!/bin/bash -# 管理定时任务 → JSON 输出 +# 管理定时任务 → JSON 输出（参数化查询） read -r input ACTION=$(echo "$input" | python3 -c "import sys,json; print(json.load(sys.stdin).get('action',''))" 2>/dev/null) NAME=$(echo "$input" | python3 -c "import sys,json; print(json.load(sys.stdin).get('name',''))" 2>/dev/null) @@ -14,22 +14,23 @@ DB_URL="${DATABASE_URL}" case "$ACTION" in add) [ -z "$CMD" ] && echo '{"error":"add 需要 command 参数"}' && exit 0 - psql "$DB_URL" -c " + psql "$DB_URL" -v name="$NAME" -v desc="$DESC" -v cmd="$CMD" -v interval="$INTERVAL" << 'SQL' INSERT INTO scheduled_tasks (id, name, description, command, interval_seconds, enabled, next_run_at) - VALUES (gen_random_uuid(), '$NAME', '$DESC', '$CMD', $INTERVAL, true, NOW() + interval '$INTERVAL seconds') + VALUES (gen_random_uuid(), :'name', :'desc', :'cmd', :interval, true, NOW() + interval ':interval seconds') ON CONFLICT (name) DO UPDATE SET command=EXCLUDED.command, interval_seconds=EXCLUDED.interval_seconds, enabled=true, description=EXCLUDED.description - " 2>/dev/null && echo '{"ok":true,"action":"add","name":"'"$NAME"'"}' || echo '{"error":"添加失败"}' +SQL + echo '{"ok":true,"action":"add","name":"'"$NAME"'"}' ;; delete) - psql "$DB_URL" -c "DELETE FROM scheduled_tasks WHERE name = '$NAME'" 2>/dev/null + psql "$DB_URL" -v name="$NAME" -c "DELETE FROM scheduled_tasks WHERE name = :'name'" 2>/dev/null echo '{"ok":true,"action":"delete","name":"'"$NAME"'"}' ;; enable) - psql "$DB_URL" -c "UPDATE scheduled_tasks SET enabled = true WHERE name = '$NAME'" 2>/dev/null + psql "$DB_URL" -v name="$NAME" -c "UPDATE scheduled_tasks SET enabled = true WHERE name = :'name'" 2>/dev/null echo '{"ok":true,"action":"enable","name":"'"$NAME"'"}' ;; disable) - psql "$DB_URL" -c "UPDATE scheduled_tasks SET enabled = false WHERE name = '$NAME'" 2>/dev/null + psql "$DB_URL" -v name="$NAME" -c "UPDATE scheduled_tasks SET enabled = false WHERE name = :'name'" 2>/dev/null echo '{"ok":true,"action":"disable","name":"'"$NAME"'"}' ;; *) diff --git a/src/llm/conversation.rs b/src/llm/conversation.rs index d35185c..69e1531 100644 --- a/src/llm/conversation.rs +++ b/src/llm/conversation.rs @@ -222,7 +222,8 @@ impl ChatHandle { chunk } - pub async fn consume(mut self) -> Result { + /// 消费所有流式块，返回完整文本 + pub async fn consume(&mut self) -> Result { while let Some(chunk) = self.next_chunk().await { if let StreamChunk::Error(e) = chunk { return Err(e); } } diff --git a/src/main.rs b/src/main.rs index c1f274d..7aa11c1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -598,7 +598,9 @@ async fn generate_reply( user_text: String, db: &Option>, ) -> Result { - let handle = conv.chat(user_text).await?; + let mut handle = conv.chat(user_text).await?; + // 先消费流（usage 在 Done chunk 中才可用） + let reply = handle.consume().await?; if let Some(u) = handle.get_usage() { info!( "📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}", @@ -610,7 +612,7 @@ async fn generate_reply( ); store_usage(db, "", conv.model(), conv.provider_name(), u).await; } - handle.consume().await + Ok(reply) } async fn generate_reply_with_tools( diff --git a/src/scheduler.rs b/src/scheduler.rs index 599acf5..cbeff9f 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -76,7 +76,8 @@ async fn fetch_due_tasks(pool: &PgPool) -> Result, String> { FROM scheduled_tasks WHERE enabled = true AND next_run_at <= NOW() ORDER BY next_run_at ASC - LIMIT 10 + LIMIT 5 + FOR UPDATE SKIP LOCKED "#, ) .fetch_all(pool) @@ -110,17 +111,20 @@ async fn execute_task(task: &DueTask) -> String { task.cwd.clone() }; - let result = tokio::process::Command::new(shell) - .arg("-c") - .arg(&task.command) - .current_dir(&cwd) - .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::piped()) - .output() - .await; + let result = tokio::time::timeout( + std::time::Duration::from_secs(120), + tokio::process::Command::new(shell) + .arg("-c") + .arg(&task.command) + .current_dir(&cwd) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .output() + ).await; match result { - Ok(output) => { + Ok(Ok(output)) => { let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); @@ -138,7 +142,8 @@ async fn execute_task(task: &DueTask) -> String { ) } } - Err(e) => format!("执行失败: {}", e), + Ok(Err(e)) => format!("进程错误: {}", e), + Err(_timeout) => "执行超时 (120s)".to_string(), } }