This commit is contained in:
2026-06-20 19:16:44 +08:00
parent db3d030169
commit dbb87823e8
26 changed files with 676 additions and 115 deletions
+4 -1
View File
@@ -68,7 +68,10 @@ func sliceLines(s string, line, limit *int) string {
} }
end := len(lines) end := len(lines)
if limit != nil && *limit >= 0 { if limit != nil && *limit >= 0 {
if start+*limit < end { // overflow-safe:用剩余行数为上限做钳制,避免 start+*limit 在 *limit 很大
// (如 math.MaxInt)时整数溢出为负,绕过下界判断后导致 lines[start:end] panic。
remaining := len(lines) - start
if *limit < remaining {
end = start + *limit end = start + *limit
} }
} }
@@ -3,6 +3,7 @@ package handlers_test
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"math"
"os" "os"
"path/filepath" "path/filepath"
"testing" "testing"
@@ -81,6 +82,37 @@ func TestFsRead_LineLimit(t *testing.T) {
assert.Equal(t, "L1\nL2", out.Content) assert.Equal(t, "L1\nL2", out.Content)
} }
// TestFsRead_HugeLimit_NoOverflow 锁定 bug #2:limit 为 math.MaxInt 时
// start+limit 会整数溢出为负,旧逻辑下 lines[start:end] 会 panic。修复后应安全
// 钳制到文件末尾,返回从 line 起的剩余内容。
func TestFsRead_HugeLimit_NoOverflow(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
require.NoError(t, os.WriteFile(filepath.Join(tmp, "f.txt"),
[]byte("L0\nL1\nL2\nL3\n"), 0o644))
h := handlers.NewFsReadHandler()
sess := handlers.SessionContext{CwdPath: tmp}
line := 1
limit := math.MaxInt
params := mustJSON(t, map[string]any{
"sessionId": "x", "path": filepath.Join(tmp, "f.txt"),
"line": line, "limit": limit,
})
var res handlers.FsResult
require.NotPanics(t, func() {
res = h.Handle(context.Background(), sess, params)
})
require.Nil(t, res.Err)
var out struct {
Content string `json:"content"`
}
require.NoError(t, json.Unmarshal(res.OK, &out))
// 从 line=1 到末尾(split 出 5 个元素:L0 L1 L2 L3 "")。
assert.Equal(t, "L1\nL2\nL3\n", out.Content)
}
func mustJSON(t *testing.T, v any) json.RawMessage { func mustJSON(t *testing.T, v any) json.RawMessage {
t.Helper() t.Helper()
b, err := json.Marshal(v) b, err := json.Marshal(v)
+8 -3
View File
@@ -314,8 +314,13 @@ func resolvedEnvelope(reqID uuid.UUID, status PermissionStatus, chosen string) *
} }
} }
// extractToolName 尝试从 toolCall JSON 中提取工具标识。多个 agent 实现的字段名不同, // extractToolName 尝试从 toolCall JSON 中提取真实工具标识。
// 依次尝试 name/toolName/kind/title;都取不到则返回空(保守走人工审批)。 //
// 安全约束(自动放行依据):只接受真实工具名字段(name / toolName)。绝不回退到
// ACP toolCall 的 'kind'(粗粒度枚举 read/edit/execute…)或 'title'(自由文本)——
// 二者都由不可信 agent 自行填写,agent 可把一个 execute/shell 调用标成 kind='read'
// 或冒用白名单里的 title 来绕过人工审批闸门。取不到真实名时返回空串,使 allowlistHit
// 一律不放行(保守走人工审批),同时保留空名默认拒绝的语义。
func extractToolName(toolCall json.RawMessage) string { func extractToolName(toolCall json.RawMessage) string {
if len(toolCall) == 0 { if len(toolCall) == 0 {
return "" return ""
@@ -324,7 +329,7 @@ func extractToolName(toolCall json.RawMessage) string {
if err := json.Unmarshal(toolCall, &m); err != nil { if err := json.Unmarshal(toolCall, &m); err != nil {
return "" return ""
} }
for _, k := range []string{"name", "toolName", "kind", "title"} { for _, k := range []string{"name", "toolName"} {
if v, ok := m[k].(string); ok && v != "" { if v, ok := m[k].(string); ok && v != "" {
return v return v
} }
@@ -1,6 +1,9 @@
package acp package acp
import "testing" import (
"encoding/json"
"testing"
)
// TestAllowlistHit 覆盖白名单匹配的边界:空 toolName 一律不放行(含 "*"), // TestAllowlistHit 覆盖白名单匹配的边界:空 toolName 一律不放行(含 "*"),
// 白名单中的空串条目不误匹配。 // 白名单中的空串条目不误匹配。
@@ -26,3 +29,44 @@ func TestAllowlistHit(t *testing.T) {
}) })
} }
} }
// TestExtractToolName_IgnoresAgentAssertedKindTitle 锁定安全约束:extractToolName
// 只信任真实工具名(name / toolName),绝不回退到 agent 自填的 kind / title。否则
// agent 可把 execute/shell 调用伪装成 kind='read' 或冒用白名单 title 绕过人工审批。
func TestExtractToolName_IgnoresAgentAssertedKindTitle(t *testing.T) {
cases := []struct {
name string
toolCall string
want string
}{
{"真实 name", `{"name":"shell.exec"}`, "shell.exec"},
{"真实 toolName", `{"toolName":"shell.exec"}`, "shell.exec"},
{"仅 kind 不取", `{"kind":"read"}`, ""},
{"仅 title 不取", `{"title":"Read file"}`, ""},
{"kind+title 都不取", `{"kind":"execute","title":"safe.tool"}`, ""},
{"name 优先于 kind/title", `{"name":"real.tool","kind":"read","title":"x"}`, "real.tool"},
{"空 toolCall", ``, ""},
{"无可识别字段", `{"unknownField":1}`, ""},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
got := extractToolName(json.RawMessage(c.toolCall))
if got != c.want {
t.Fatalf("extractToolName(%q) = %q, want %q", c.toolCall, got, c.want)
}
})
}
}
// TestAllowlist_DoesNotHonorAgentKind 端到端验证 bug #3:即使 allowlist 含 'read'
// 这类 kind 值,agent 把 execute 调用标成 kind='read' 也不会自动放行(extractToolName
// 返回空 → allowlistHit 为 false)。
func TestAllowlist_DoesNotHonorAgentKind(t *testing.T) {
toolName := extractToolName(json.RawMessage(`{"kind":"read","title":"safe"}`))
if toolName != "" {
t.Fatalf("kind/title 不应被当作工具名, got %q", toolName)
}
if allowlistHit([]string{"read", "safe", "*"}, toolName) {
t.Fatal("kind/title 伪装的调用不应命中白名单(含通配符)")
}
}
+14
View File
@@ -249,7 +249,21 @@ func (r *Relay) reader(ctx context.Context, sess handlers.SessionContext) {
} }
// handleAgentRequest dispatches an agent → server request to the appropriate handler. // handleAgentRequest dispatches an agent → server request to the appropriate handler.
// 在独立 goroutine 中运行,故 defer recover 兜底任何 handler panic:记录日志并回送
// JSON-RPC internal error,避免 agent 构造的请求 panic 掉整个进程(DoS)。
func (r *Relay) handleAgentRequest(ctx context.Context, sess handlers.SessionContext, req *Message) { func (r *Relay) handleAgentRequest(ctx context.Context, sess handlers.SessionContext, req *Message) {
defer func() {
if rec := recover(); rec != nil {
r.log.Error("acp.relay.handle_request_panic",
"session_id", r.sessionID, "method", req.Method, "panic", fmt.Sprintf("%v", rec))
resp := NewResponseErr(req.ID, JSONRPCInternalError, "internal error")
select {
case r.serverInitiated <- resp:
case <-ctx.Done():
}
}
}()
var resp *Message var resp *Message
switch req.Method { switch req.Method {
case "fs/read_text_file": case "fs/read_text_file":
+23 -1
View File
@@ -108,6 +108,12 @@ type Process struct {
StartedAt time.Time StartedAt time.Time
KilledByUs atomic.Bool // true → exited,false → crashed KilledByUs atomic.Bool // true → exited,false → crashed
done chan struct{} // monitor 退出时 close done chan struct{} // monitor 退出时 close
// relayCancel 取消 relay 的会话级 context。该 context 由 Spawn 用
// context.Background() 派生,独立于调用方的 spawn ctx(spawn ctx 在 Spawn
// 返回后立即取消),故 relay 的 reader/writer/persist/permission Decide 在整个
// 会话生命周期内都使用未取消的 ctx。onExit 调用 relayCancel 完成 relay teardown。
relayCancel context.CancelFunc
} }
// stderrRing 是固定行数 ring buffer,monitor 退出时取 tail 写 last_error。 // stderrRing 是固定行数 ring buffer,monitor 退出时取 tail 写 last_error。
@@ -237,6 +243,12 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
s.log, s.log,
) )
// relay 的会话级 context:独立于调用方的 spawn ctx(生产调用方在 Spawn 返回后
// 立即 cancel spawn ctx)。relay 的 reader/writer/persist/permission Decide 必须
// 在整个会话生命周期内使用未取消的 ctx,否则 prompt 被丢、事件停止持久化、权限
// 审批被强制拒绝。teardown 由 onExit 调 relayCancel + Relay.Close(关 r.done)完成。
relayCtx, relayCancel := context.WithCancel(context.Background())
proc := &Process{ proc := &Process{
SessionID: sess.ID, SessionID: sess.ID,
WorkspaceID: sess.WorkspaceID, WorkspaceID: sess.WorkspaceID,
@@ -251,6 +263,7 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
StderrBuf: newStderrRing(s.cfg.StderrBufferLines), StderrBuf: newStderrRing(s.cfg.StderrBufferLines),
StartedAt: time.Now(), StartedAt: time.Now(),
done: make(chan struct{}), done: make(chan struct{}),
relayCancel: relayCancel,
} }
s.mu.Lock() s.mu.Lock()
@@ -259,7 +272,7 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
go s.drainStderr(proc) go s.drainStderr(proc)
go s.monitorProcess(proc) go s.monitorProcess(proc)
go relay.Run(ctx, handlers.SessionContext{ go relay.Run(relayCtx, handlers.SessionContext{
SessionID: sess.ID, SessionID: sess.ID,
WorkspaceID: sess.WorkspaceID, WorkspaceID: sess.WorkspaceID,
UserID: sess.UserID, UserID: sess.UserID,
@@ -268,6 +281,7 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
ToolAllowlist: kind.ToolAllowlist, ToolAllowlist: kind.ToolAllowlist,
}) })
// 握手仍走调用方的 spawn ctx,保留 SpawnTimeout 截止语义。
if err := s.handshake(ctx, sess, kind, relay, proc, extraEnv); err != nil { if err := s.handshake(ctx, sess, kind, relay, proc, extraEnv); err != nil {
s.log.Error("acp.handshake_failed", "session_id", sess.ID, "err", err.Error()) s.log.Error("acp.handshake_failed", "session_id", sess.ID, "err", err.Error())
s.Kill(ctx, sess.ID, s.cfg.KillGrace) s.Kill(ctx, sess.ID, s.cfg.KillGrace)
@@ -516,6 +530,14 @@ func (s *Supervisor) onExit(ctx context.Context, proc *Process, status SessionSt
} }
} }
// 取消 relay 的会话级 context:停止 reader/writer,并让仍阻塞在 Decide / 已派生
// 的 handleAgentRequest goroutine 收到 ctx.Done 后退出。在 Relay.Close 之前调,
// 二者共同完成 teardown(relayCancel 停 ctx 派生的 goroutine,Close 关 r.done +
// 广播 session_terminated)。relayCancel 在早期 spawn 失败路径也可能为 nil。
if proc.relayCancel != nil {
proc.relayCancel()
}
// 最后关 relay:广播 session_terminated 给 WS subscribers。 // 最后关 relay:广播 session_terminated 给 WS subscribers。
proc.Relay.Close(string(status), &exitCode) proc.Relay.Close(string(status), &exitCode)
} }
+33 -9
View File
@@ -70,26 +70,50 @@ type subscription struct {
// the stream closes, or ctx is cancelled. // the stream closes, or ctx is cancelled.
// It returns all events from nextEventID to the current tail. // It returns all events from nextEventID to the current tail.
func (sub *subscription) Next(ctx context.Context, nextEventID int64) ([]SSEEvent, bool, error) { func (sub *subscription) Next(ctx context.Context, nextEventID int64) ([]SSEEvent, bool, error) {
sub.st.mu.Lock() st := sub.st
defer sub.st.mu.Unlock()
for int64(len(sub.st.events)) <= nextEventID && !sub.st.closed { // Launch a single watcher that observes ctx cancellation and wakes any
// waiter. It broadcasts under the lock so the signal cannot be missed by a
// waiter that has not yet parked on cond.Wait(); the predicate loop below
// re-checks ctx.Err() after every wake, so a broadcast delivered before a
// waiter parks is still observed on the next predicate check.
//
// done is closed (under the lock) when Next returns, which makes the watcher
// exit even if ctx never cancels — so neither goroutine can leak.
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
select { select {
case <-ctx.Done(): case <-ctx.Done():
sub.st.cond.Broadcast() st.mu.Lock()
st.cond.Broadcast()
st.mu.Unlock()
case <-done: case <-done:
} }
}() }()
sub.st.cond.Wait()
st.mu.Lock()
defer func() {
// Signal the watcher to exit, then release the lock. Closing done under
// the lock pairs with the watcher's Broadcast-under-lock so there is no
// window in which the watcher broadcasts to an unlocked, already-returned
// Next.
close(done) close(done)
if ctx.Err() != nil { st.mu.Unlock()
}()
// Wait until there is data at/after nextEventID, the stream closes, or ctx
// is cancelled. Each predicate component is re-checked after every wake so a
// missed/early broadcast cannot wedge the wait.
for int64(len(st.events)) <= nextEventID && !st.closed && ctx.Err() == nil {
st.cond.Wait()
}
if int64(len(st.events)) <= nextEventID && !st.closed && ctx.Err() != nil {
return nil, false, ctx.Err() return nil, false, ctx.Err()
} }
}
out := append([]SSEEvent(nil), sub.st.events[nextEventID:]...) out := append([]SSEEvent(nil), st.events[nextEventID:]...)
return out, sub.st.closed, nil return out, st.closed, nil
} }
// streamerHub manages all in-flight streamers. // streamerHub manages all in-flight streamers.
+71
View File
@@ -667,6 +667,77 @@ func TestStreamer_TitleGenerationTriggered(t *testing.T) {
}, 1*time.Second, 10*time.Millisecond, "GenerateTitle should have been triggered") }, 1*time.Second, 10*time.Millisecond, "GenerateTitle should have been triggered")
} }
// TestStreamer_Next_EarlyCancel_NoLeak exercises the early-cancel race in
// subscription.Next: when ctx is already cancelled (or cancels in the window
// before the waiter parks on cond.Wait) and there are no new events, Next must
// return ctx.Err() promptly instead of blocking forever, and must not leak the
// Next goroutine or its ctx watcher.
func TestStreamer_Next_EarlyCancel_NoLeak(t *testing.T) {
t.Parallel()
st := newStreamer()
sub := &subscription{st: st}
// Case 1: ctx already cancelled before Next is called. With no events and an
// open stream, the predicate is unsatisfied except for ctx, so Next must
// observe cancellation immediately rather than parking forever.
ctx, cancel := context.WithCancel(context.Background())
cancel()
done := make(chan struct{})
go func() {
defer close(done)
evs, closed, err := sub.Next(ctx, 0)
require.ErrorIs(t, err, context.Canceled)
require.Nil(t, evs)
require.False(t, closed)
}()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("Next did not return on an already-cancelled ctx; goroutine leaked")
}
// Case 2: cancel concurrently while a waiter is (about to be) parked. Repeat
// to widen the window between launching the watcher and parking on Wait.
for i := 0; i < 200; i++ {
ctx, cancel := context.WithCancel(context.Background())
ret := make(chan error, 1)
go func() {
_, _, err := sub.Next(ctx, 0)
ret <- err
}()
cancel()
select {
case err := <-ret:
require.ErrorIs(t, err, context.Canceled)
case <-time.After(2 * time.Second):
t.Fatalf("iteration %d: Next wedged on concurrent cancel", i)
}
}
}
// TestStreamer_Next_DataDeliveredDespiteCancel verifies that when events are
// already buffered, Next returns them even if ctx is cancelled (data wins over
// a late cancel, preserving the original behavior).
func TestStreamer_Next_DataDeliveredDespiteCancel(t *testing.T) {
t.Parallel()
st := newStreamer()
sub := &subscription{st: st}
st.append(SSEEvent{Event: "text_delta", Data: map[string]string{"text": "hi"}})
ctx, cancel := context.WithCancel(context.Background())
cancel()
evs, closed, err := sub.Next(ctx, 0)
require.NoError(t, err)
require.False(t, closed)
require.Len(t, evs, 1)
require.Equal(t, "text_delta", evs[0].Event)
}
// ===== helpers for title generation test ===== // ===== helpers for title generation test =====
// spyConvSvc is a minimal conversationService-shaped struct that calls a hook // spyConvSvc is a minimal conversationService-shaped struct that calls a hook
+5 -1
View File
@@ -41,8 +41,12 @@ func (g *unixGroup) TerminateGroup(done <-chan struct{}, grace time.Duration) {
_ = syscall.Kill(-g.pgid, syscall.SIGTERM) _ = syscall.Kill(-g.pgid, syscall.SIGTERM)
select { select {
case <-done: case <-done:
return // 组长已被回收(Cmd.Wait 返回),但其子孙(如 trap SIGTERM 的
// vite/esbuild/node)可能仍存活于同一进程组。无条件向整组补发一次
// SIGKILL 以确保不留下任何孤儿;若组已全灭,内核返回 ESRCH,无害。
_ = syscall.Kill(-g.pgid, syscall.SIGKILL)
case <-time.After(grace): case <-time.After(grace):
// 组长自身超过 grace 仍未退出:直接对整组 SIGKILL。
_ = syscall.Kill(-g.pgid, syscall.SIGKILL) _ = syscall.Kill(-g.pgid, syscall.SIGKILL)
} }
} }
+7
View File
@@ -5,6 +5,7 @@ package proc
import ( import (
"errors" "errors"
"os/exec" "os/exec"
"sync"
"time" "time"
"unsafe" "unsafe"
@@ -16,6 +17,7 @@ import (
// 关闭 Job 句柄即整树终止。优于 Process.Kill(只杀单 PID,漏 node 等子进程)。 // 关闭 Job 句柄即整树终止。优于 Process.Kill(只杀单 PID,漏 node 等子进程)。
type windowsGroup struct { type windowsGroup struct {
job windows.Handle job windows.Handle
closeOnce sync.Once
} }
func newGroup() Group { return &windowsGroup{} } func newGroup() Group { return &windowsGroup{} }
@@ -65,9 +67,14 @@ func (g *windowsGroup) TerminateGroup(_ <-chan struct{}, _ time.Duration) {
g.Close() g.Close()
} }
// Close 幂等且并发安全:TerminateGroup 与 monitor goroutine 可能并发调用,
// 用 sync.Once 保证 CloseHandle 至多执行一次,避免重复关闭句柄(可能误关被复用
// 的句柄)及对 g.job 的数据竞争。不改变 KILL_ON_JOB_CLOSE 语义:句柄关闭即整树终止。
func (g *windowsGroup) Close() { func (g *windowsGroup) Close() {
g.closeOnce.Do(func() {
if g.job != 0 { if g.job != 0 {
_ = windows.CloseHandle(g.job) _ = windows.CloseHandle(g.job)
g.job = 0 g.job = 0
} }
})
} }
+8
View File
@@ -83,6 +83,14 @@ func IsPermanent(err error) bool {
return errors.As(err, &pe) return errors.As(err, &pe)
} }
// ErrLeaseLost is returned by the terminal/retry write-back methods
// (Complete, FailWithRetry, DeadLetter) when the fencing predicate matches no
// row: the job is no longer leased by this worker (e.g. the reaper re-leased it
// to another worker after the lease timeout). The caller must treat this as a
// no-op — the write was correctly dropped to avoid clobbering the new run's
// state — and must NOT surface it as a failure that triggers further retries.
var ErrLeaseLost = errors.New("job lease lost; write-back skipped")
// RunnerConfig is the per-runner configuration loaded from config.yaml. // RunnerConfig is the per-runner configuration loaded from config.yaml.
// Zero Interval means "use the runner's documented default". // Zero Interval means "use the runner's documented default".
type RunnerConfig struct { type RunnerConfig struct {
+15
View File
@@ -54,6 +54,7 @@ func NewModule(cfg Config, deps ModuleDeps) *Module {
sch.Add(name, entry.interval, entry.enabled, fn) sch.Add(name, entry.interval, entry.enabled, fn)
} }
handlerTimeout := handlerTimeoutFor(cfg.JobReaper)
workers := make([]*Worker, cfg.Workers) workers := make([]*Worker, cfg.Workers)
for i := 0; i < cfg.Workers; i++ { for i := 0; i < cfg.Workers; i++ {
workers[i] = NewWorker(deps.Repo, deps.Handlers, WorkerOptions{ workers[i] = NewWorker(deps.Repo, deps.Handlers, WorkerOptions{
@@ -62,6 +63,7 @@ func NewModule(cfg Config, deps ModuleDeps) *Module {
Backoff: deps.WorkerBackoff, Backoff: deps.WorkerBackoff,
Clock: clk, Clock: clk,
Log: log, Log: log,
HandlerTimeout: handlerTimeout,
OnDeadLetter: deps.OnDeadLetter, OnDeadLetter: deps.OnDeadLetter,
}) })
} }
@@ -112,6 +114,19 @@ func (m *Module) Stop(ctx context.Context) {
func workerID(i int) string { return "worker#" + strconv.Itoa(i) } func workerID(i int) string { return "worker#" + strconv.Itoa(i) }
// handlerTimeoutFor derives the per-job handler timeout from the reaper's lease
// timeout, leaving a margin below it so a healthy long handler finishes (or is
// canceled) before the reaper re-leases the job and double-runs it. Returns 0
// (no timeout) when reaping is disabled or no positive lease timeout is set.
func handlerTimeoutFor(rc JobReaperConfig) time.Duration {
if !rc.Enabled || rc.LeaseTimeout <= 0 {
return 0
}
// 90% of the lease window leaves headroom for the write-back to land before
// the reaper's cutoff.
return rc.LeaseTimeout - rc.LeaseTimeout/10
}
type schedEntry struct { type schedEntry struct {
interval time.Duration interval time.Duration
enabled bool enabled bool
+9 -6
View File
@@ -20,28 +20,31 @@ UPDATE jobs
) )
RETURNING *; RETURNING *;
-- name: CompleteJob :exec -- name: CompleteJob :execrows
-- 围栏谓词:只有持有当前租约的 worker 才能写回,防止被 reaper 重新租出后旧 worker 覆盖新一轮状态。
UPDATE jobs UPDATE jobs
SET status='completed', completed_at=now(), SET status='completed', completed_at=now(),
leased_at=NULL, leased_by=NULL, last_error=NULL, leased_at=NULL, leased_by=NULL, last_error=NULL,
updated_at=now() updated_at=now()
WHERE id=$1; WHERE id=$1 AND leased_by=$2 AND status='running';
-- name: FailJobWithRetry :exec -- name: FailJobWithRetry :execrows
-- 围栏谓词同上。
UPDATE jobs UPDATE jobs
SET status='pending', SET status='pending',
scheduled_at=$2, scheduled_at=$2,
leased_at=NULL, leased_by=NULL, leased_at=NULL, leased_by=NULL,
last_error=$3, last_error=$3,
updated_at=now() updated_at=now()
WHERE id=$1; WHERE id=$1 AND leased_by=$4 AND status='running';
-- name: DeadLetterJob :exec -- name: DeadLetterJob :execrows
-- 围栏谓词同上。
UPDATE jobs UPDATE jobs
SET status='dead', completed_at=now(), SET status='dead', completed_at=now(),
leased_at=NULL, leased_by=NULL, last_error=$2, leased_at=NULL, leased_by=NULL, last_error=$2,
updated_at=now() updated_at=now()
WHERE id=$1; WHERE id=$1 AND leased_by=$3 AND status='running';
-- name: GetJob :one -- name: GetJob :one
SELECT * FROM jobs WHERE id=$1; SELECT * FROM jobs WHERE id=$1;
+38 -13
View File
@@ -22,9 +22,13 @@ import (
type Repository interface { type Repository interface {
Enqueue(ctx context.Context, typ JobType, payload json.RawMessage, scheduledAt time.Time, maxAttempts int32) (*Job, error) Enqueue(ctx context.Context, typ JobType, payload json.RawMessage, scheduledAt time.Time, maxAttempts int32) (*Job, error)
Lease(ctx context.Context, leasedBy string) (*Job, error) Lease(ctx context.Context, leasedBy string) (*Job, error)
Complete(ctx context.Context, id uuid.UUID) error // Complete, FailWithRetry and DeadLetter take the leasing worker id and apply
FailWithRetry(ctx context.Context, id uuid.UUID, scheduledAt time.Time, lastErr string) error // a fencing predicate (id + leased_by + status='running'); if the lease has
DeadLetter(ctx context.Context, id uuid.UUID, lastErr string) error // been lost (e.g. the reaper re-leased the job) they return ErrLeaseLost so
// the caller drops the write instead of clobbering the new run's state.
Complete(ctx context.Context, id uuid.UUID, leasedBy string) error
FailWithRetry(ctx context.Context, id uuid.UUID, leasedBy string, scheduledAt time.Time, lastErr string) error
DeadLetter(ctx context.Context, id uuid.UUID, leasedBy string, lastErr string) error
Get(ctx context.Context, id uuid.UUID) (*Job, error) Get(ctx context.Context, id uuid.UUID) (*Job, error)
ReapStuckRunningJobs(ctx context.Context, leasedBefore time.Time) ([]uuid.UUID, error) ReapStuckRunningJobs(ctx context.Context, leasedBefore time.Time) ([]uuid.UUID, error)
PurgeCompleted(ctx context.Context, completedBefore time.Time) (int64, error) PurgeCompleted(ctx context.Context, completedBefore time.Time) (int64, error)
@@ -90,11 +94,20 @@ func (r *PgRepository) Lease(ctx context.Context, leasedBy string) (*Job, error)
return rowToJob(row), nil return rowToJob(row), nil
} }
// Complete marks a job as completed and clears its lease/error fields. // Complete marks a job as completed and clears its lease/error fields. The
func (r *PgRepository) Complete(ctx context.Context, id uuid.UUID) error { // update is fenced on the leasing worker id and status='running'; if no row
if err := r.q.CompleteJob(ctx, toPgUUID(id)); err != nil { // matches (lease lost to the reaper) it returns ErrLeaseLost as a no-op signal.
func (r *PgRepository) Complete(ctx context.Context, id uuid.UUID, leasedBy string) error {
n, err := r.q.CompleteJob(ctx, jobssqlc.CompleteJobParams{
ID: toPgUUID(id),
LeasedBy: &leasedBy,
})
if err != nil {
return errs.Wrap(err, errs.CodeInternal, "complete job") return errs.Wrap(err, errs.CodeInternal, "complete job")
} }
if n == 0 {
return ErrLeaseLost
}
return nil return nil
} }
@@ -102,25 +115,37 @@ func (r *PgRepository) Complete(ctx context.Context, id uuid.UUID) error {
// the last error. Attempts is preserved (it was incremented on lease) so the // the last error. Attempts is preserved (it was incremented on lease) so the
// worker can compare against MaxAttempts to decide whether to retry or // worker can compare against MaxAttempts to decide whether to retry or
// dead-letter on the next failure. // dead-letter on the next failure.
func (r *PgRepository) FailWithRetry(ctx context.Context, id uuid.UUID, scheduledAt time.Time, lastErr string) error { func (r *PgRepository) FailWithRetry(ctx context.Context, id uuid.UUID, leasedBy string, scheduledAt time.Time, lastErr string) error {
if err := r.q.FailJobWithRetry(ctx, jobssqlc.FailJobWithRetryParams{ n, err := r.q.FailJobWithRetry(ctx, jobssqlc.FailJobWithRetryParams{
ID: toPgUUID(id), ID: toPgUUID(id),
ScheduledAt: pgtype.Timestamptz{Time: scheduledAt, Valid: true}, ScheduledAt: pgtype.Timestamptz{Time: scheduledAt, Valid: true},
LastError: &lastErr, LastError: &lastErr,
}); err != nil { LeasedBy: &leasedBy,
})
if err != nil {
return errs.Wrap(err, errs.CodeInternal, "fail job with retry") return errs.Wrap(err, errs.CodeInternal, "fail job with retry")
} }
if n == 0 {
return ErrLeaseLost
}
return nil return nil
} }
// DeadLetter terminally marks a job as dead and stores the last error. // DeadLetter terminally marks a job as dead and stores the last error. Fenced
func (r *PgRepository) DeadLetter(ctx context.Context, id uuid.UUID, lastErr string) error { // on the leasing worker id and status='running'; returns ErrLeaseLost on a
if err := r.q.DeadLetterJob(ctx, jobssqlc.DeadLetterJobParams{ // 0-row match so a stale worker does not overwrite a re-leased run.
func (r *PgRepository) DeadLetter(ctx context.Context, id uuid.UUID, leasedBy string, lastErr string) error {
n, err := r.q.DeadLetterJob(ctx, jobssqlc.DeadLetterJobParams{
ID: toPgUUID(id), ID: toPgUUID(id),
LastError: &lastErr, LastError: &lastErr,
}); err != nil { LeasedBy: &leasedBy,
})
if err != nil {
return errs.Wrap(err, errs.CodeInternal, "dead-letter job") return errs.Wrap(err, errs.CodeInternal, "dead-letter job")
} }
if n == 0 {
return ErrLeaseLost
}
return nil return nil
} }
+50 -4
View File
@@ -125,7 +125,7 @@ func TestComplete_ClearsLease(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
leased, err := repo.Lease(ctx, "h") leased, err := repo.Lease(ctx, "h")
require.NoError(t, err) require.NoError(t, err)
require.NoError(t, repo.Complete(ctx, leased.ID)) require.NoError(t, repo.Complete(ctx, leased.ID, "h"))
got, err := repo.Get(ctx, job.ID) got, err := repo.Get(ctx, job.ID)
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, jobs.StatusCompleted, got.Status) assert.Equal(t, jobs.StatusCompleted, got.Status)
@@ -133,6 +133,49 @@ func TestComplete_ClearsLease(t *testing.T) {
assert.Nil(t, got.LeasedBy) assert.Nil(t, got.LeasedBy)
} }
// TestComplete_LeaseLostIsNoOp verifies the fencing predicate: a write-back
// from a worker that no longer holds the lease (e.g. the reaper re-leased the
// job) matches 0 rows and returns ErrLeaseLost without clobbering the new run.
func TestComplete_LeaseLostIsNoOp(t *testing.T) {
t.Parallel()
ctx := context.Background()
repo, pool := setupRepo(t)
_, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5)
require.NoError(t, err)
// Worker A leases the job, then its handler stalls past the lease timeout.
leased, err := repo.Lease(ctx, "workerA")
require.NoError(t, err)
// Reaper resets it to pending; worker B re-leases it (now leased_by=workerB).
_, err = pool.Exec(ctx, `UPDATE jobs SET leased_at = now() - INTERVAL '10 minutes' WHERE id=$1`, leased.ID)
require.NoError(t, err)
ids, err := repo.ReapStuckRunningJobs(ctx, time.Now().Add(-5*time.Minute))
require.NoError(t, err)
require.Len(t, ids, 1)
reLeased, err := repo.Lease(ctx, "workerB")
require.NoError(t, err)
require.Equal(t, leased.ID, reLeased.ID)
// Stale worker A's terminal write-backs must be no-ops (ErrLeaseLost) and
// must not change the row that workerB now owns.
assert.ErrorIs(t, repo.Complete(ctx, leased.ID, "workerA"), jobs.ErrLeaseLost)
assert.ErrorIs(t, repo.FailWithRetry(ctx, leased.ID, "workerA", time.Now(), "stale"), jobs.ErrLeaseLost)
assert.ErrorIs(t, repo.DeadLetter(ctx, leased.ID, "workerA", "stale"), jobs.ErrLeaseLost)
got, err := repo.Get(ctx, leased.ID)
require.NoError(t, err)
assert.Equal(t, jobs.StatusRunning, got.Status, "row must still be owned by workerB's run")
require.NotNil(t, got.LeasedBy)
assert.Equal(t, "workerB", *got.LeasedBy)
assert.Nil(t, got.LastError, "stale write must not have set an error")
// Worker B (the true owner) completes successfully.
require.NoError(t, repo.Complete(ctx, leased.ID, "workerB"))
got, err = repo.Get(ctx, leased.ID)
require.NoError(t, err)
assert.Equal(t, jobs.StatusCompleted, got.Status)
}
func TestFailWithRetry_ReschedulesAndKeepsAttempts(t *testing.T) { func TestFailWithRetry_ReschedulesAndKeepsAttempts(t *testing.T) {
t.Parallel() t.Parallel()
ctx := context.Background() ctx := context.Background()
@@ -142,7 +185,7 @@ func TestFailWithRetry_ReschedulesAndKeepsAttempts(t *testing.T) {
leased, err := repo.Lease(ctx, "h") leased, err := repo.Lease(ctx, "h")
require.NoError(t, err) require.NoError(t, err)
next := time.Now().Add(30 * time.Second) next := time.Now().Add(30 * time.Second)
require.NoError(t, repo.FailWithRetry(ctx, leased.ID, next, "boom")) require.NoError(t, repo.FailWithRetry(ctx, leased.ID, "h", next, "boom"))
got, err := repo.Get(ctx, leased.ID) got, err := repo.Get(ctx, leased.ID)
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, jobs.StatusPending, got.Status) assert.Equal(t, jobs.StatusPending, got.Status)
@@ -160,7 +203,7 @@ func TestDeadLetter_FinalState(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
leased, err := repo.Lease(ctx, "h") leased, err := repo.Lease(ctx, "h")
require.NoError(t, err) require.NoError(t, err)
require.NoError(t, repo.DeadLetter(ctx, leased.ID, "max attempts")) require.NoError(t, repo.DeadLetter(ctx, leased.ID, "h", "max attempts"))
got, err := repo.Get(ctx, leased.ID) got, err := repo.Get(ctx, leased.ID)
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, jobs.StatusDead, got.Status) assert.Equal(t, jobs.StatusDead, got.Status)
@@ -204,7 +247,10 @@ func TestPurge_DeletesOldCompletedAndDead(t *testing.T) {
deadOld := mkOld(jobs.StatusDead) deadOld := mkOld(jobs.StatusDead)
freshComp, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5) freshComp, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5)
require.NoError(t, err) require.NoError(t, err)
require.NoError(t, repo.Complete(ctx, freshComp.ID)) leasedFresh, err := repo.Lease(ctx, "h")
require.NoError(t, err)
require.Equal(t, freshComp.ID, leasedFresh.ID)
require.NoError(t, repo.Complete(ctx, freshComp.ID, "h"))
rows, err := repo.PurgeCompleted(ctx, time.Now().Add(-7*24*time.Hour)) rows, err := repo.PurgeCompleted(ctx, time.Now().Add(-7*24*time.Hour))
require.NoError(t, err) require.NoError(t, err)
+39 -15
View File
@@ -11,35 +11,49 @@ import (
"github.com/jackc/pgx/v5/pgtype" "github.com/jackc/pgx/v5/pgtype"
) )
const completeJob = `-- name: CompleteJob :exec const completeJob = `-- name: CompleteJob :execrows
UPDATE jobs UPDATE jobs
SET status='completed', completed_at=now(), SET status='completed', completed_at=now(),
leased_at=NULL, leased_by=NULL, last_error=NULL, leased_at=NULL, leased_by=NULL, last_error=NULL,
updated_at=now() updated_at=now()
WHERE id=$1 WHERE id=$1 AND leased_by=$2 AND status='running'
` `
func (q *Queries) CompleteJob(ctx context.Context, id pgtype.UUID) error { type CompleteJobParams struct {
_, err := q.db.Exec(ctx, completeJob, id) ID pgtype.UUID `json:"id"`
return err LeasedBy *string `json:"leased_by"`
} }
const deadLetterJob = `-- name: DeadLetterJob :exec // 围栏谓词:只有持有当前租约的 worker 才能写回,防止被 reaper 重新租出后旧 worker 覆盖新一轮状态。
func (q *Queries) CompleteJob(ctx context.Context, arg CompleteJobParams) (int64, error) {
result, err := q.db.Exec(ctx, completeJob, arg.ID, arg.LeasedBy)
if err != nil {
return 0, err
}
return result.RowsAffected(), nil
}
const deadLetterJob = `-- name: DeadLetterJob :execrows
UPDATE jobs UPDATE jobs
SET status='dead', completed_at=now(), SET status='dead', completed_at=now(),
leased_at=NULL, leased_by=NULL, last_error=$2, leased_at=NULL, leased_by=NULL, last_error=$2,
updated_at=now() updated_at=now()
WHERE id=$1 WHERE id=$1 AND leased_by=$3 AND status='running'
` `
type DeadLetterJobParams struct { type DeadLetterJobParams struct {
ID pgtype.UUID `json:"id"` ID pgtype.UUID `json:"id"`
LastError *string `json:"last_error"` LastError *string `json:"last_error"`
LeasedBy *string `json:"leased_by"`
} }
func (q *Queries) DeadLetterJob(ctx context.Context, arg DeadLetterJobParams) error { // 围栏谓词同上。
_, err := q.db.Exec(ctx, deadLetterJob, arg.ID, arg.LastError) func (q *Queries) DeadLetterJob(ctx context.Context, arg DeadLetterJobParams) (int64, error) {
return err result, err := q.db.Exec(ctx, deadLetterJob, arg.ID, arg.LastError, arg.LeasedBy)
if err != nil {
return 0, err
}
return result.RowsAffected(), nil
} }
const enqueueJob = `-- name: EnqueueJob :one const enqueueJob = `-- name: EnqueueJob :one
@@ -83,25 +97,35 @@ func (q *Queries) EnqueueJob(ctx context.Context, arg EnqueueJobParams) (Job, er
return i, err return i, err
} }
const failJobWithRetry = `-- name: FailJobWithRetry :exec const failJobWithRetry = `-- name: FailJobWithRetry :execrows
UPDATE jobs UPDATE jobs
SET status='pending', SET status='pending',
scheduled_at=$2, scheduled_at=$2,
leased_at=NULL, leased_by=NULL, leased_at=NULL, leased_by=NULL,
last_error=$3, last_error=$3,
updated_at=now() updated_at=now()
WHERE id=$1 WHERE id=$1 AND leased_by=$4 AND status='running'
` `
type FailJobWithRetryParams struct { type FailJobWithRetryParams struct {
ID pgtype.UUID `json:"id"` ID pgtype.UUID `json:"id"`
ScheduledAt pgtype.Timestamptz `json:"scheduled_at"` ScheduledAt pgtype.Timestamptz `json:"scheduled_at"`
LastError *string `json:"last_error"` LastError *string `json:"last_error"`
LeasedBy *string `json:"leased_by"`
} }
func (q *Queries) FailJobWithRetry(ctx context.Context, arg FailJobWithRetryParams) error { // 围栏谓词同上。
_, err := q.db.Exec(ctx, failJobWithRetry, arg.ID, arg.ScheduledAt, arg.LastError) func (q *Queries) FailJobWithRetry(ctx context.Context, arg FailJobWithRetryParams) (int64, error) {
return err result, err := q.db.Exec(ctx, failJobWithRetry,
arg.ID,
arg.ScheduledAt,
arg.LastError,
arg.LeasedBy,
)
if err != nil {
return 0, err
}
return result.RowsAffected(), nil
} }
const getJob = `-- name: GetJob :one const getJob = `-- name: GetJob :one
+41 -5
View File
@@ -20,6 +20,12 @@ type WorkerOptions struct {
Backoff []time.Duration Backoff []time.Duration
Clock clock.Clock Clock clock.Clock
Log *slog.Logger Log *slog.Logger
// HandlerTimeout bounds a single handler invocation. It must be set slightly
// below the JobReaper's LeaseTimeout so a healthy long-running handler
// finishes (or is canceled) before the reaper re-leases the job to another
// worker and double-runs it. When <= 0 no per-job timeout is applied (the
// module always sets it from JobReaper.LeaseTimeout when reaping is enabled).
HandlerTimeout time.Duration
// OnDeadLetter, if non-nil, is invoked once after a job is moved to dead-letter // OnDeadLetter, if non-nil, is invoked once after a job is moved to dead-letter
// (either because IsPermanent(err) or attempts exhausted). The hook runs in // (either because IsPermanent(err) or attempts exhausted). The hook runs in
// the worker goroutine; panics are recovered and logged but do not stop the // the worker goroutine; panics are recovered and logged but do not stop the
@@ -91,7 +97,9 @@ func (w *Worker) execute(ctx context.Context, job Job) {
h, ok := w.handlers[job.Type] h, ok := w.handlers[job.Type]
if !ok { if !ok {
if derr := w.repo.DeadLetter(ctx, job.ID, fmt.Sprintf("unknown job type: %s", job.Type)); derr != nil { if derr := w.repo.DeadLetter(ctx, job.ID, w.opts.ID, fmt.Sprintf("unknown job type: %s", job.Type)); w.leaseLost(derr, job, "dead_letter") {
return
} else if derr != nil {
w.opts.Log.Error("job.dead_letter_write_failed", "job_id", job.ID, "err", derr.Error()) w.opts.Log.Error("job.dead_letter_write_failed", "job_id", job.ID, "err", derr.Error())
} }
w.opts.Log.Error("job.dead_letter", "job_id", job.ID, "job_type", job.Type, "reason", "unknown_type") w.opts.Log.Error("job.dead_letter", "job_id", job.ID, "job_type", job.Type, "reason", "unknown_type")
@@ -101,7 +109,9 @@ func (w *Worker) execute(ctx context.Context, job Job) {
err := w.runHandler(ctx, h, job) err := w.runHandler(ctx, h, job)
if err == nil { if err == nil {
if cerr := w.repo.Complete(ctx, job.ID); cerr != nil { if cerr := w.repo.Complete(ctx, job.ID, w.opts.ID); w.leaseLost(cerr, job, "complete") {
return
} else if cerr != nil {
w.opts.Log.Error("job.complete_write_failed", "job_id", job.ID, "err", cerr.Error()) w.opts.Log.Error("job.complete_write_failed", "job_id", job.ID, "err", cerr.Error())
} }
w.opts.Log.Info("job.complete", w.opts.Log.Info("job.complete",
@@ -113,7 +123,9 @@ func (w *Worker) execute(ctx context.Context, job Job) {
// Permanent or attempts exhausted -> dead-letter // Permanent or attempts exhausted -> dead-letter
if IsPermanent(err) || job.Attempts >= job.MaxAttempts { if IsPermanent(err) || job.Attempts >= job.MaxAttempts {
msg := truncate(err.Error(), 2000) msg := truncate(err.Error(), 2000)
if derr := w.repo.DeadLetter(ctx, job.ID, msg); derr != nil { if derr := w.repo.DeadLetter(ctx, job.ID, w.opts.ID, msg); w.leaseLost(derr, job, "dead_letter") {
return
} else if derr != nil {
w.opts.Log.Error("job.dead_letter_write_failed", "job_id", job.ID, "err", derr.Error()) w.opts.Log.Error("job.dead_letter_write_failed", "job_id", job.ID, "err", derr.Error())
} }
w.opts.Log.Error("job.dead_letter", w.opts.Log.Error("job.dead_letter",
@@ -126,7 +138,9 @@ func (w *Worker) execute(ctx context.Context, job Job) {
delay := backoffFor(w.opts.Backoff, int(job.Attempts)) delay := backoffFor(w.opts.Backoff, int(job.Attempts))
next := w.opts.Clock.Now().Add(delay) next := w.opts.Clock.Now().Add(delay)
msg := truncate(err.Error(), 2000) msg := truncate(err.Error(), 2000)
if ferr := w.repo.FailWithRetry(ctx, job.ID, next, msg); ferr != nil { if ferr := w.repo.FailWithRetry(ctx, job.ID, w.opts.ID, next, msg); w.leaseLost(ferr, job, "fail_retry") {
return
} else if ferr != nil {
w.opts.Log.Error("job.fail_retry_write_failed", "job_id", job.ID, "err", ferr.Error()) w.opts.Log.Error("job.fail_retry_write_failed", "job_id", job.ID, "err", ferr.Error())
} }
w.opts.Log.Warn("job.fail_retry", w.opts.Log.Warn("job.fail_retry",
@@ -134,8 +148,30 @@ func (w *Worker) execute(ctx context.Context, job Job) {
"err", msg, "next_at", next.Format(time.RFC3339)) "err", msg, "next_at", next.Format(time.RFC3339))
} }
// runHandler wraps Handle with panic recovery, converting panics into retryable errors. // leaseLost reports whether writeErr is ErrLeaseLost. The fencing predicate on
// the terminal/retry updates matched no row, meaning this worker's lease was
// stolen (e.g. its handler outran LeaseTimeout and the reaper re-leased the
// job). The write is correctly dropped: we only log and drop, and must NOT
// treat it as a failure that schedules further work — the new run owns the job.
func (w *Worker) leaseLost(writeErr error, job Job, phase string) bool {
if errors.Is(writeErr, ErrLeaseLost) {
w.opts.Log.Warn("job.lease_lost",
"job_id", job.ID, "job_type", job.Type, "attempt", job.Attempts, "phase", phase)
return true
}
return false
}
// runHandler wraps Handle with panic recovery (converting panics into retryable
// errors) and bounds the handler to HandlerTimeout. Deriving a per-job timeout
// below the reaper's LeaseTimeout ensures a healthy long-running handler is
// canceled before the reaper re-leases the job, so it is never double-run.
func (w *Worker) runHandler(ctx context.Context, h Handler, job Job) (err error) { func (w *Worker) runHandler(ctx context.Context, h Handler, job Job) (err error) {
if w.opts.HandlerTimeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, w.opts.HandlerTimeout)
defer cancel()
}
defer func() { defer func() {
if r := recover(); r != nil { if r := recover(); r != nil {
err = fmt.Errorf("handler panic: %v", r) err = fmt.Errorf("handler panic: %v", r)
+10 -1
View File
@@ -200,7 +200,7 @@ func listTokens(deps AdminHandlerDeps) http.HandlerFunc {
func deleteToken(deps AdminHandlerDeps) http.HandlerFunc { func deleteToken(deps AdminHandlerDeps) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) { return func(w http.ResponseWriter, r *http.Request) {
rid := middleware.RequestIDFromContext(r.Context()) rid := middleware.RequestIDFromContext(r.Context())
caller, _, err := resolveSessionUser(r.Context(), deps.Users, r) caller, isAdmin, err := resolveSessionUser(r.Context(), deps.Users, r)
if err != nil { if err != nil {
httpx.WriteError(w, rid, err) httpx.WriteError(w, rid, err)
return return
@@ -211,6 +211,15 @@ func deleteToken(deps AdminHandlerDeps) http.HandlerFunc {
httpx.WriteError(w, rid, errs.Wrap(err, errs.CodeInvalidInput, "id parse")) httpx.WriteError(w, rid, errs.Wrap(err, errs.CodeInvalidInput, "id parse"))
return return
} }
// 授权:仅 admin 或 token 持有者可撤销。非 admin 时按 id 反查 token
// 校验归属;不归属(或不存在)一律返 NotFound,避免暴露 token 是否存在。
if !isAdmin {
tok, err := deps.Tokens.GetByID(r.Context(), id)
if err != nil || tok.UserID != caller {
httpx.WriteError(w, rid, errs.New(errs.CodeMcpTokenNotFound, "mcp token not found"))
return
}
}
if err := deps.Tokens.Revoke(r.Context(), id, caller); err != nil { if err := deps.Tokens.Revoke(r.Context(), id, caller); err != nil {
httpx.WriteError(w, rid, err) httpx.WriteError(w, rid, err)
return return
+9 -5
View File
@@ -8,11 +8,15 @@ RETURNING id, token_hash, user_id, name, issuer, scope,
created_by, created_at; created_by, created_at;
-- name: GetMcpTokenByHash :one -- name: GetMcpTokenByHash :one
SELECT id, token_hash, user_id, name, issuer, scope, -- 鉴权用:仅返回属于 enabled 用户的 token,禁用用户的 MCP token 立即失效
acp_session_id, expires_at, last_used_at, revoked_at, -- (与 web session 解析 GetSessionByTokenHash 的 enabled=true 约束一致)。
created_by, created_at SELECT t.id, t.token_hash, t.user_id, t.name, t.issuer, t.scope,
FROM mcp_tokens t.acp_session_id, t.expires_at, t.last_used_at, t.revoked_at,
WHERE token_hash = $1; t.created_by, t.created_at
FROM mcp_tokens t
JOIN users u ON u.id = t.user_id
WHERE t.token_hash = $1
AND u.enabled = true;
-- name: GetMcpTokenByID :one -- name: GetMcpTokenByID :one
SELECT id, token_hash, user_id, name, issuer, scope, SELECT id, token_hash, user_id, name, issuer, scope,
+9 -5
View File
@@ -74,13 +74,17 @@ func (q *Queries) DeleteMcpTokenByID(ctx context.Context, id pgtype.UUID) error
} }
const getMcpTokenByHash = `-- name: GetMcpTokenByHash :one const getMcpTokenByHash = `-- name: GetMcpTokenByHash :one
SELECT id, token_hash, user_id, name, issuer, scope, SELECT t.id, t.token_hash, t.user_id, t.name, t.issuer, t.scope,
acp_session_id, expires_at, last_used_at, revoked_at, t.acp_session_id, t.expires_at, t.last_used_at, t.revoked_at,
created_by, created_at t.created_by, t.created_at
FROM mcp_tokens FROM mcp_tokens t
WHERE token_hash = $1 JOIN users u ON u.id = t.user_id
WHERE t.token_hash = $1
AND u.enabled = true
` `
// 鉴权用:仅返回属于 enabled 用户的 token,禁用用户的 MCP token 立即失效
// (与 web session 解析 GetSessionByTokenHash 的 enabled=true 约束一致)。
func (q *Queries) GetMcpTokenByHash(ctx context.Context, tokenHash []byte) (McpToken, error) { func (q *Queries) GetMcpTokenByHash(ctx context.Context, tokenHash []byte) (McpToken, error) {
row := q.db.QueryRow(ctx, getMcpTokenByHash, tokenHash) row := q.db.QueryRow(ctx, getMcpTokenByHash, tokenHash)
var i McpToken var i McpToken
+5 -1
View File
@@ -208,6 +208,7 @@ type fakeGit struct {
cloneErr error cloneErr error
fetchErr error fetchErr error
cloneSeen []string cloneSeen []string
removed []string
} }
func (f *fakeGit) Clone(_ context.Context, dst, _, _ string, _ *git.Credential) error { func (f *fakeGit) Clone(_ context.Context, dst, _, _ string, _ *git.Credential) error {
@@ -222,7 +223,10 @@ func (f *fakeGit) Commit(_ context.Context, _, _ string, _ bool) (string, error)
func (f *fakeGit) Status(_ context.Context, _ string) ([]git.FileStatus, error) { return nil, nil } func (f *fakeGit) Status(_ context.Context, _ string) ([]git.FileStatus, error) { return nil, nil }
func (f *fakeGit) MergeFFOnly(_ context.Context, _ string) error { return nil } func (f *fakeGit) MergeFFOnly(_ context.Context, _ string) error { return nil }
func (f *fakeGit) WorktreeAdd(_ context.Context, _, _, _, _ string) error { return nil } func (f *fakeGit) WorktreeAdd(_ context.Context, _, _, _, _ string) error { return nil }
func (f *fakeGit) WorktreeRemove(_ context.Context, _, _ string) error { return nil } func (f *fakeGit) WorktreeRemove(_ context.Context, _, path string) error {
f.removed = append(f.removed, path)
return nil
}
func (f *fakeGit) WorktreeList(_ context.Context, _ string) ([]git.Worktree, error) { func (f *fakeGit) WorktreeList(_ context.Context, _ string) ([]git.Worktree, error) {
return nil, nil return nil, nil
} }
+44 -14
View File
@@ -88,31 +88,34 @@ func (s *worktreeService) List(ctx context.Context, c Caller, wsID uuid.UUID) ([
} }
func (s *worktreeService) Delete(ctx context.Context, c Caller, wtID uuid.UUID) error { func (s *worktreeService) Delete(ctx context.Context, c Caller, wtID uuid.UUID) error {
wt, err := s.repo.GetWorktreeByID(ctx, wtID) ws, err := s.authorizeWorktreeWrite(ctx, c, wtID)
if err != nil { if err != nil {
return err return err
} }
ws, err := s.repo.GetWorkspaceByID(ctx, wt.WorkspaceID) // check-and-transition 必须原子:在 tx 内对行加 FOR UPDATE,确认 idle 后才转 pruning。
if err != nil { // 否则并发 Acquire(FOR UPDATE)可能刚把它置为 active,而这里用非锁读看到旧的 idle,
return err // 进而 git rm 掉 agent 正在使用的目录。仅 idle→pruning 在锁下完成。
var wtPath string
err = s.repo.InTx(ctx, func(tx Tx) error {
wt, e := tx.GetWorktreeByIDForUpdate(ctx, wtID)
if e != nil {
return e
} }
if _, canWrite, err := s.pa.ResolveByID(ctx, c.UserID, c.IsAdmin, ws.ProjectID); err != nil { if wt.Status != WorktreeStatusIdle {
return err
} else if !canWrite {
return errs.New(errs.CodeForbidden, "no write access")
}
if wt.Status == WorktreeStatusActive {
return errs.New(errs.CodeWorktreeAlreadyActive, "release before delete") return errs.New(errs.CodeWorktreeAlreadyActive, "release before delete")
} }
// 先 set pruning,再调 git rm,然后 DELETE row。中途失败 row 留 pruning。 pruned, e := tx.SetWorktreePruning(ctx, wtID)
err = s.repo.InTx(ctx, func(tx Tx) error { if e != nil {
_, e := tx.SetWorktreePruning(ctx, wtID)
return e return e
}
wtPath = pruned.Path
return nil
}) })
if err != nil { if err != nil {
return err return err
} }
if rmErr := s.gitr.WorktreeRemove(ctx, ws.MainPath, wt.Path); rmErr != nil { // 仅在守卫后的转移成功后才动磁盘。中途失败 row 留 pruning。
if rmErr := s.gitr.WorktreeRemove(ctx, ws.MainPath, wtPath); rmErr != nil {
return errs.Wrap(rmErr, errs.CodeGitCmdFailed, "git worktree remove") return errs.Wrap(rmErr, errs.CodeGitCmdFailed, "git worktree remove")
} }
if err := s.repo.DeleteWorktree(ctx, wtID); err != nil { if err := s.repo.DeleteWorktree(ctx, wtID); err != nil {
@@ -123,6 +126,9 @@ func (s *worktreeService) Delete(ctx context.Context, c Caller, wtID uuid.UUID)
} }
func (s *worktreeService) Acquire(ctx context.Context, c Caller, wtID uuid.UUID) (*Worktree, error) { func (s *worktreeService) Acquire(ctx context.Context, c Caller, wtID uuid.UUID) (*Worktree, error) {
if _, err := s.authorizeWorktreeWrite(ctx, c, wtID); err != nil {
return nil, err
}
holder := c.Holder() holder := c.Holder()
var out *Worktree var out *Worktree
err := s.repo.InTx(ctx, func(tx Tx) error { err := s.repo.InTx(ctx, func(tx Tx) error {
@@ -157,6 +163,9 @@ func (s *worktreeService) Acquire(ctx context.Context, c Caller, wtID uuid.UUID)
} }
func (s *worktreeService) Release(ctx context.Context, c Caller, wtID uuid.UUID) (*Worktree, error) { func (s *worktreeService) Release(ctx context.Context, c Caller, wtID uuid.UUID) (*Worktree, error) {
if _, err := s.authorizeWorktreeWrite(ctx, c, wtID); err != nil {
return nil, err
}
holder := c.Holder() holder := c.Holder()
var out *Worktree var out *Worktree
err := s.repo.InTx(ctx, func(tx Tx) error { err := s.repo.InTx(ctx, func(tx Tx) error {
@@ -190,6 +199,27 @@ func (s *worktreeService) ReleaseByHolder(ctx context.Context, holder string) ([
return s.repo.ReleaseWorktreesByHolder(ctx, holder) return s.repo.ReleaseWorktreesByHolder(ctx, holder)
} }
// authorizeWorktreeWrite 把 worktree → workspace → project 解析出来并要求 write 权限,
// 与 Delete/Create 的鉴权方式一致。返回 worktree 所属的 workspace 供调用方使用
// (如取 MainPath)。Acquire/Release 此前缺失该校验,导致任意已登录用户可凭 UUID
// 跨租户锁定/释放任意 worktree。
func (s *worktreeService) authorizeWorktreeWrite(ctx context.Context, c Caller, wtID uuid.UUID) (*Workspace, error) {
wt, err := s.repo.GetWorktreeByID(ctx, wtID)
if err != nil {
return nil, err
}
ws, err := s.repo.GetWorkspaceByID(ctx, wt.WorkspaceID)
if err != nil {
return nil, err
}
if _, canWrite, err := s.pa.ResolveByID(ctx, c.UserID, c.IsAdmin, ws.ProjectID); err != nil {
return nil, err
} else if !canWrite {
return nil, errs.New(errs.CodeForbidden, "no write access")
}
return ws, nil
}
func (s *worktreeService) audit(ctx context.Context, uid *uuid.UUID, action, ttype, tid string, meta map[string]any) { func (s *worktreeService) audit(ctx context.Context, uid *uuid.UUID, action, ttype, tid string, meta map[string]any) {
_ = s.rec.Record(ctx, audit.Entry{ _ = s.rec.Record(ctx, audit.Entry{
UserID: uid, Action: action, TargetType: ttype, TargetID: tid, Metadata: meta, UserID: uid, Action: action, TargetType: ttype, TargetID: tid, Metadata: meta,
@@ -6,8 +6,17 @@ import (
"github.com/google/uuid" "github.com/google/uuid"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
) )
func requireCode(t *testing.T, err error, want errs.Code) {
t.Helper()
ae, ok := errs.As(err)
require.True(t, ok, "expected *errs.AppError, got %T", err)
require.Equal(t, want, ae.Code)
}
func newWtSvc(t *testing.T) (*worktreeService, *fakeRepo, *fakePA, *fakeGit, uuid.UUID) { func newWtSvc(t *testing.T) (*worktreeService, *fakeRepo, *fakePA, *fakeGit, uuid.UUID) {
t.Helper() t.Helper()
pa := &fakePA{pid: uuid.New(), canRead: true, canWrite: true} pa := &fakePA{pid: uuid.New(), canRead: true, canWrite: true}
@@ -81,3 +90,55 @@ func TestWorktreeService_Delete_BlockedWhenActive(t *testing.T) {
err = svc.Delete(context.Background(), caller, wt.ID) err = svc.Delete(context.Background(), caller, wt.ID)
require.Error(t, err) require.Error(t, err)
} }
// Bug #5: Delete 的 check-and-transition 必须在锁下原子完成。即使 worktree 处于
// active(被持有),Delete 也必须在事务内(GetWorktreeByIDForUpdate)发现并拒绝,
// 返回 CodeWorktreeAlreadyActive,且不触碰磁盘 / 不删行。
func TestWorktreeService_Delete_AtomicGuard_RefusesActive(t *testing.T) {
t.Parallel()
svc, repo, _, gitr, wsID := newWtSvc(t)
caller := Caller{UserID: uuid.New()}
wt, err := svc.Create(context.Background(), caller, wsID, CreateWorktreeInput{Branch: "feat-g"})
require.NoError(t, err)
_, err = svc.Acquire(context.Background(), caller, wt.ID)
require.NoError(t, err)
err = svc.Delete(context.Background(), caller, wt.ID)
require.Error(t, err)
requireCode(t, err, errs.CodeWorktreeAlreadyActive)
// 磁盘未被触碰,行仍存在且保持 active。
require.Empty(t, gitr.removed)
got, err := repo.GetWorktreeByID(context.Background(), wt.ID)
require.NoError(t, err)
require.Equal(t, WorktreeStatusActive, got.Status)
}
// Bug #6: Acquire 必须做 project 鉴权。无 write 权限的 caller 不得锁定 worktree。
func TestWorktreeService_Acquire_DeniedWithoutWriteAccess(t *testing.T) {
t.Parallel()
svc, _, pa, _, wsID := newWtSvc(t)
caller := Caller{UserID: uuid.New()}
wt, err := svc.Create(context.Background(), caller, wsID, CreateWorktreeInput{Branch: "feat-a"})
require.NoError(t, err)
pa.canWrite = false
_, err = svc.Acquire(context.Background(), caller, wt.ID)
require.Error(t, err)
requireCode(t, err, errs.CodeForbidden)
}
// Bug #6: Release 同样必须做 project 鉴权(在 holder-match 之前)。
func TestWorktreeService_Release_DeniedWithoutWriteAccess(t *testing.T) {
t.Parallel()
svc, _, pa, _, wsID := newWtSvc(t)
caller := Caller{UserID: uuid.New()}
wt, err := svc.Create(context.Background(), caller, wsID, CreateWorktreeInput{Branch: "feat-rr"})
require.NoError(t, err)
_, err = svc.Acquire(context.Background(), caller, wt.ID)
require.NoError(t, err)
pa.canWrite = false
_, err = svc.Release(context.Background(), caller, wt.ID)
require.Error(t, err)
requireCode(t, err, errs.CodeForbidden)
}
@@ -0,0 +1,24 @@
-- Revert to the original (buggy) bare ON DELETE SET NULL composite FKs so that
-- this migration is a faithful inverse of 0015 up. Constraint names and column
-- references match the definitions in 0003_workspace.up.sql and 0006_acp.up.sql.
-- ============ acp_sessions -> requirements ============
ALTER TABLE acp_sessions DROP CONSTRAINT acp_sessions_req_project_fk;
ALTER TABLE acp_sessions ADD CONSTRAINT acp_sessions_req_project_fk
FOREIGN KEY (project_id, requirement_id)
REFERENCES requirements (project_id, id) MATCH SIMPLE
ON DELETE SET NULL;
-- ============ acp_sessions -> issues ============
ALTER TABLE acp_sessions DROP CONSTRAINT acp_sessions_issue_project_fk;
ALTER TABLE acp_sessions ADD CONSTRAINT acp_sessions_issue_project_fk
FOREIGN KEY (project_id, issue_id)
REFERENCES issues (project_id, id) MATCH SIMPLE
ON DELETE SET NULL;
-- ============ issues -> workspaces ============
ALTER TABLE issues DROP CONSTRAINT issues_project_workspace_fk;
ALTER TABLE issues ADD CONSTRAINT issues_project_workspace_fk
FOREIGN KEY (project_id, workspace_id)
REFERENCES workspaces (project_id, id) MATCH SIMPLE
ON DELETE SET NULL;
@@ -0,0 +1,34 @@
-- Fix data-integrity bug: composite FKs with bare ON DELETE SET NULL null ALL
-- referencing columns on delete, including the NOT NULL project_id column,
-- which raises 23502 and makes the referenced-row delete impossible.
-- PostgreSQL 15+ supports a column-subset form: ON DELETE SET NULL (col, ...)
-- so only the nullable column is nulled.
-- ============ issues -> workspaces ============
-- Original (0003): FOREIGN KEY (project_id, workspace_id)
-- REFERENCES workspaces (project_id, id) MATCH SIMPLE ON DELETE SET NULL
-- Deleting a workspace still referenced by an issue tried to null issues.project_id
-- (NOT NULL) -> blocked workspace deletion entirely.
ALTER TABLE issues DROP CONSTRAINT issues_project_workspace_fk;
ALTER TABLE issues ADD CONSTRAINT issues_project_workspace_fk
FOREIGN KEY (project_id, workspace_id)
REFERENCES workspaces (project_id, id) MATCH SIMPLE
ON DELETE SET NULL (workspace_id);
-- ============ acp_sessions -> issues ============
-- Original (0006): FOREIGN KEY (project_id, issue_id)
-- REFERENCES issues (project_id, id) MATCH SIMPLE ON DELETE SET NULL
ALTER TABLE acp_sessions DROP CONSTRAINT acp_sessions_issue_project_fk;
ALTER TABLE acp_sessions ADD CONSTRAINT acp_sessions_issue_project_fk
FOREIGN KEY (project_id, issue_id)
REFERENCES issues (project_id, id) MATCH SIMPLE
ON DELETE SET NULL (issue_id);
-- ============ acp_sessions -> requirements ============
-- Original (0006): FOREIGN KEY (project_id, requirement_id)
-- REFERENCES requirements (project_id, id) MATCH SIMPLE ON DELETE SET NULL
ALTER TABLE acp_sessions DROP CONSTRAINT acp_sessions_req_project_fk;
ALTER TABLE acp_sessions ADD CONSTRAINT acp_sessions_req_project_fk
FOREIGN KEY (project_id, requirement_id)
REFERENCES requirements (project_id, id) MATCH SIMPLE
ON DELETE SET NULL (requirement_id);
@@ -114,6 +114,14 @@ export function useConversationSession() {
const id = streamingMessageId.value const id = streamingMessageId.value
if (id == null) return if (id == null) return
await chatApi.cancelMessage(id) await chatApi.cancelMessage(id)
// 取消会 abort SSE,后端不再下发 done/error/message_meta 终止事件,
// 故须在此手动收尾(对齐 done/error/onError 路径):将仍处于 pending 的
// 助手消息置为终态 'cancelled',并清空 streamingMessageId 解除 isStreaming。
const msg = messages.value.find((m) => m.id === id)
if (msg && msg.status === 'pending') {
msg.status = 'cancelled'
}
streamingMessageId.value = null
detachStream() detachStream()
} }