You've already forked agentic-coding-workflow
bugfix
This commit is contained in:
@@ -68,7 +68,10 @@ func sliceLines(s string, line, limit *int) string {
|
||||
}
|
||||
end := len(lines)
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ package handlers_test
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
@@ -81,6 +82,37 @@ func TestFsRead_LineLimit(t *testing.T) {
|
||||
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 {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(v)
|
||||
|
||||
@@ -314,8 +314,13 @@ func resolvedEnvelope(reqID uuid.UUID, status PermissionStatus, chosen string) *
|
||||
}
|
||||
}
|
||||
|
||||
// extractToolName 尝试从 toolCall JSON 中提取工具标识。多个 agent 实现的字段名不同,
|
||||
// 依次尝试 name/toolName/kind/title;都取不到则返回空(保守走人工审批)。
|
||||
// extractToolName 尝试从 toolCall JSON 中提取真实工具标识。
|
||||
//
|
||||
// 安全约束(自动放行依据):只接受真实工具名字段(name / toolName)。绝不回退到
|
||||
// ACP toolCall 的 'kind'(粗粒度枚举 read/edit/execute…)或 'title'(自由文本)——
|
||||
// 二者都由不可信 agent 自行填写,agent 可把一个 execute/shell 调用标成 kind='read'
|
||||
// 或冒用白名单里的 title 来绕过人工审批闸门。取不到真实名时返回空串,使 allowlistHit
|
||||
// 一律不放行(保守走人工审批),同时保留空名默认拒绝的语义。
|
||||
func extractToolName(toolCall json.RawMessage) string {
|
||||
if len(toolCall) == 0 {
|
||||
return ""
|
||||
@@ -324,7 +329,7 @@ func extractToolName(toolCall json.RawMessage) string {
|
||||
if err := json.Unmarshal(toolCall, &m); err != nil {
|
||||
return ""
|
||||
}
|
||||
for _, k := range []string{"name", "toolName", "kind", "title"} {
|
||||
for _, k := range []string{"name", "toolName"} {
|
||||
if v, ok := m[k].(string); ok && v != "" {
|
||||
return v
|
||||
}
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
package acp
|
||||
|
||||
import "testing"
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// 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 伪装的调用不应命中白名单(含通配符)")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -249,7 +249,21 @@ func (r *Relay) reader(ctx context.Context, sess handlers.SessionContext) {
|
||||
}
|
||||
|
||||
// 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) {
|
||||
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
|
||||
switch req.Method {
|
||||
case "fs/read_text_file":
|
||||
|
||||
@@ -108,6 +108,12 @@ type Process struct {
|
||||
StartedAt time.Time
|
||||
KilledByUs atomic.Bool // true → exited,false → crashed
|
||||
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。
|
||||
@@ -237,6 +243,12 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
|
||||
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{
|
||||
SessionID: sess.ID,
|
||||
WorkspaceID: sess.WorkspaceID,
|
||||
@@ -251,6 +263,7 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
|
||||
StderrBuf: newStderrRing(s.cfg.StderrBufferLines),
|
||||
StartedAt: time.Now(),
|
||||
done: make(chan struct{}),
|
||||
relayCancel: relayCancel,
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
@@ -259,7 +272,7 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
|
||||
|
||||
go s.drainStderr(proc)
|
||||
go s.monitorProcess(proc)
|
||||
go relay.Run(ctx, handlers.SessionContext{
|
||||
go relay.Run(relayCtx, handlers.SessionContext{
|
||||
SessionID: sess.ID,
|
||||
WorkspaceID: sess.WorkspaceID,
|
||||
UserID: sess.UserID,
|
||||
@@ -268,6 +281,7 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
|
||||
ToolAllowlist: kind.ToolAllowlist,
|
||||
})
|
||||
|
||||
// 握手仍走调用方的 spawn ctx,保留 SpawnTimeout 截止语义。
|
||||
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.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。
|
||||
proc.Relay.Close(string(status), &exitCode)
|
||||
}
|
||||
|
||||
+41
-17
@@ -70,26 +70,50 @@ type subscription struct {
|
||||
// the stream closes, or ctx is cancelled.
|
||||
// It returns all events from nextEventID to the current tail.
|
||||
func (sub *subscription) Next(ctx context.Context, nextEventID int64) ([]SSEEvent, bool, error) {
|
||||
sub.st.mu.Lock()
|
||||
defer sub.st.mu.Unlock()
|
||||
st := sub.st
|
||||
|
||||
for int64(len(sub.st.events)) <= nextEventID && !sub.st.closed {
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
sub.st.cond.Broadcast()
|
||||
case <-done:
|
||||
}
|
||||
}()
|
||||
sub.st.cond.Wait()
|
||||
close(done)
|
||||
if ctx.Err() != nil {
|
||||
return nil, false, ctx.Err()
|
||||
// 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{})
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
st.mu.Lock()
|
||||
st.cond.Broadcast()
|
||||
st.mu.Unlock()
|
||||
case <-done:
|
||||
}
|
||||
}()
|
||||
|
||||
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)
|
||||
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()
|
||||
}
|
||||
out := append([]SSEEvent(nil), sub.st.events[nextEventID:]...)
|
||||
return out, sub.st.closed, nil
|
||||
|
||||
if int64(len(st.events)) <= nextEventID && !st.closed && ctx.Err() != nil {
|
||||
return nil, false, ctx.Err()
|
||||
}
|
||||
|
||||
out := append([]SSEEvent(nil), st.events[nextEventID:]...)
|
||||
return out, st.closed, nil
|
||||
}
|
||||
|
||||
// streamerHub manages all in-flight streamers.
|
||||
|
||||
@@ -176,7 +176,7 @@ func (panicRepo) SoftDeleteConversation(_ context.Context, _ uuid.UUID) error {
|
||||
func (panicRepo) InsertMessage(_ context.Context, _ uuid.UUID, _ MessageRole, _ string, _ MessageStatus) (*Message, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (panicRepo) GetMessage(_ context.Context, _ int64) (*Message, error) { panic("not implemented") }
|
||||
func (panicRepo) GetMessage(_ context.Context, _ int64) (*Message, error) { panic("not implemented") }
|
||||
func (panicRepo) ListMessagesByConversation(_ context.Context, _ uuid.UUID, _ *int64, _ int) ([]Message, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
@@ -667,6 +667,77 @@ func TestStreamer_TitleGenerationTriggered(t *testing.T) {
|
||||
}, 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 =====
|
||||
|
||||
// spyConvSvc is a minimal conversationService-shaped struct that calls a hook
|
||||
|
||||
@@ -41,8 +41,12 @@ func (g *unixGroup) TerminateGroup(done <-chan struct{}, grace time.Duration) {
|
||||
_ = syscall.Kill(-g.pgid, syscall.SIGTERM)
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
// 组长已被回收(Cmd.Wait 返回),但其子孙(如 trap SIGTERM 的
|
||||
// vite/esbuild/node)可能仍存活于同一进程组。无条件向整组补发一次
|
||||
// SIGKILL 以确保不留下任何孤儿;若组已全灭,内核返回 ESRCH,无害。
|
||||
_ = syscall.Kill(-g.pgid, syscall.SIGKILL)
|
||||
case <-time.After(grace):
|
||||
// 组长自身超过 grace 仍未退出:直接对整组 SIGKILL。
|
||||
_ = syscall.Kill(-g.pgid, syscall.SIGKILL)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ package proc
|
||||
import (
|
||||
"errors"
|
||||
"os/exec"
|
||||
"sync"
|
||||
"time"
|
||||
"unsafe"
|
||||
|
||||
@@ -15,7 +16,8 @@ import (
|
||||
// KILL_ON_JOB_CLOSE 的 Job,此后该进程派生的所有子孙自动归入同一 Job,
|
||||
// 关闭 Job 句柄即整树终止。优于 Process.Kill(只杀单 PID,漏 node 等子进程)。
|
||||
type windowsGroup struct {
|
||||
job windows.Handle
|
||||
job windows.Handle
|
||||
closeOnce sync.Once
|
||||
}
|
||||
|
||||
func newGroup() Group { return &windowsGroup{} }
|
||||
@@ -65,9 +67,14 @@ func (g *windowsGroup) TerminateGroup(_ <-chan struct{}, _ time.Duration) {
|
||||
g.Close()
|
||||
}
|
||||
|
||||
// Close 幂等且并发安全:TerminateGroup 与 monitor goroutine 可能并发调用,
|
||||
// 用 sync.Once 保证 CloseHandle 至多执行一次,避免重复关闭句柄(可能误关被复用
|
||||
// 的句柄)及对 g.job 的数据竞争。不改变 KILL_ON_JOB_CLOSE 语义:句柄关闭即整树终止。
|
||||
func (g *windowsGroup) Close() {
|
||||
if g.job != 0 {
|
||||
_ = windows.CloseHandle(g.job)
|
||||
g.job = 0
|
||||
}
|
||||
g.closeOnce.Do(func() {
|
||||
if g.job != 0 {
|
||||
_ = windows.CloseHandle(g.job)
|
||||
g.job = 0
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -83,6 +83,14 @@ func IsPermanent(err error) bool {
|
||||
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.
|
||||
// Zero Interval means "use the runner's documented default".
|
||||
type RunnerConfig struct {
|
||||
|
||||
+21
-6
@@ -54,15 +54,17 @@ func NewModule(cfg Config, deps ModuleDeps) *Module {
|
||||
sch.Add(name, entry.interval, entry.enabled, fn)
|
||||
}
|
||||
|
||||
handlerTimeout := handlerTimeoutFor(cfg.JobReaper)
|
||||
workers := make([]*Worker, cfg.Workers)
|
||||
for i := 0; i < cfg.Workers; i++ {
|
||||
workers[i] = NewWorker(deps.Repo, deps.Handlers, WorkerOptions{
|
||||
ID: workerID(i),
|
||||
PollInterval: 200 * time.Millisecond,
|
||||
Backoff: deps.WorkerBackoff,
|
||||
Clock: clk,
|
||||
Log: log,
|
||||
OnDeadLetter: deps.OnDeadLetter,
|
||||
ID: workerID(i),
|
||||
PollInterval: 200 * time.Millisecond,
|
||||
Backoff: deps.WorkerBackoff,
|
||||
Clock: clk,
|
||||
Log: log,
|
||||
HandlerTimeout: handlerTimeout,
|
||||
OnDeadLetter: deps.OnDeadLetter,
|
||||
})
|
||||
}
|
||||
return &Module{cfg: cfg, scheduler: sch, workers: workers, log: log}
|
||||
@@ -112,6 +114,19 @@ func (m *Module) Stop(ctx context.Context) {
|
||||
|
||||
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 {
|
||||
interval time.Duration
|
||||
enabled bool
|
||||
|
||||
@@ -20,28 +20,31 @@ UPDATE jobs
|
||||
)
|
||||
RETURNING *;
|
||||
|
||||
-- name: CompleteJob :exec
|
||||
-- name: CompleteJob :execrows
|
||||
-- 围栏谓词:只有持有当前租约的 worker 才能写回,防止被 reaper 重新租出后旧 worker 覆盖新一轮状态。
|
||||
UPDATE jobs
|
||||
SET status='completed', completed_at=now(),
|
||||
leased_at=NULL, leased_by=NULL, last_error=NULL,
|
||||
updated_at=now()
|
||||
WHERE id=$1;
|
||||
WHERE id=$1 AND leased_by=$2 AND status='running';
|
||||
|
||||
-- name: FailJobWithRetry :exec
|
||||
-- name: FailJobWithRetry :execrows
|
||||
-- 围栏谓词同上。
|
||||
UPDATE jobs
|
||||
SET status='pending',
|
||||
scheduled_at=$2,
|
||||
leased_at=NULL, leased_by=NULL,
|
||||
last_error=$3,
|
||||
updated_at=now()
|
||||
WHERE id=$1;
|
||||
WHERE id=$1 AND leased_by=$4 AND status='running';
|
||||
|
||||
-- name: DeadLetterJob :exec
|
||||
-- name: DeadLetterJob :execrows
|
||||
-- 围栏谓词同上。
|
||||
UPDATE jobs
|
||||
SET status='dead', completed_at=now(),
|
||||
leased_at=NULL, leased_by=NULL, last_error=$2,
|
||||
updated_at=now()
|
||||
WHERE id=$1;
|
||||
WHERE id=$1 AND leased_by=$3 AND status='running';
|
||||
|
||||
-- name: GetJob :one
|
||||
SELECT * FROM jobs WHERE id=$1;
|
||||
|
||||
+38
-13
@@ -22,9 +22,13 @@ import (
|
||||
type Repository interface {
|
||||
Enqueue(ctx context.Context, typ JobType, payload json.RawMessage, scheduledAt time.Time, maxAttempts int32) (*Job, error)
|
||||
Lease(ctx context.Context, leasedBy string) (*Job, error)
|
||||
Complete(ctx context.Context, id uuid.UUID) error
|
||||
FailWithRetry(ctx context.Context, id uuid.UUID, scheduledAt time.Time, lastErr string) error
|
||||
DeadLetter(ctx context.Context, id uuid.UUID, lastErr string) error
|
||||
// Complete, FailWithRetry and DeadLetter take the leasing worker id and apply
|
||||
// a fencing predicate (id + leased_by + status='running'); if the lease has
|
||||
// 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)
|
||||
ReapStuckRunningJobs(ctx context.Context, leasedBefore time.Time) ([]uuid.UUID, 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
|
||||
}
|
||||
|
||||
// Complete marks a job as completed and clears its lease/error fields.
|
||||
func (r *PgRepository) Complete(ctx context.Context, id uuid.UUID) error {
|
||||
if err := r.q.CompleteJob(ctx, toPgUUID(id)); err != nil {
|
||||
// Complete marks a job as completed and clears its lease/error fields. The
|
||||
// update is fenced on the leasing worker id and status='running'; if no row
|
||||
// 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")
|
||||
}
|
||||
if n == 0 {
|
||||
return ErrLeaseLost
|
||||
}
|
||||
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
|
||||
// worker can compare against MaxAttempts to decide whether to retry or
|
||||
// dead-letter on the next failure.
|
||||
func (r *PgRepository) FailWithRetry(ctx context.Context, id uuid.UUID, scheduledAt time.Time, lastErr string) error {
|
||||
if err := r.q.FailJobWithRetry(ctx, jobssqlc.FailJobWithRetryParams{
|
||||
func (r *PgRepository) FailWithRetry(ctx context.Context, id uuid.UUID, leasedBy string, scheduledAt time.Time, lastErr string) error {
|
||||
n, err := r.q.FailJobWithRetry(ctx, jobssqlc.FailJobWithRetryParams{
|
||||
ID: toPgUUID(id),
|
||||
ScheduledAt: pgtype.Timestamptz{Time: scheduledAt, Valid: true},
|
||||
LastError: &lastErr,
|
||||
}); err != nil {
|
||||
LeasedBy: &leasedBy,
|
||||
})
|
||||
if err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "fail job with retry")
|
||||
}
|
||||
if n == 0 {
|
||||
return ErrLeaseLost
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeadLetter terminally marks a job as dead and stores the last error.
|
||||
func (r *PgRepository) DeadLetter(ctx context.Context, id uuid.UUID, lastErr string) error {
|
||||
if err := r.q.DeadLetterJob(ctx, jobssqlc.DeadLetterJobParams{
|
||||
// DeadLetter terminally marks a job as dead and stores the last error. Fenced
|
||||
// on the leasing worker id and status='running'; returns ErrLeaseLost on a
|
||||
// 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),
|
||||
LastError: &lastErr,
|
||||
}); err != nil {
|
||||
LeasedBy: &leasedBy,
|
||||
})
|
||||
if err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "dead-letter job")
|
||||
}
|
||||
if n == 0 {
|
||||
return ErrLeaseLost
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -125,7 +125,7 @@ func TestComplete_ClearsLease(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
leased, err := repo.Lease(ctx, "h")
|
||||
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)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, jobs.StatusCompleted, got.Status)
|
||||
@@ -133,6 +133,49 @@ func TestComplete_ClearsLease(t *testing.T) {
|
||||
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) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
@@ -142,7 +185,7 @@ func TestFailWithRetry_ReschedulesAndKeepsAttempts(t *testing.T) {
|
||||
leased, err := repo.Lease(ctx, "h")
|
||||
require.NoError(t, err)
|
||||
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)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, jobs.StatusPending, got.Status)
|
||||
@@ -160,7 +203,7 @@ func TestDeadLetter_FinalState(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
leased, err := repo.Lease(ctx, "h")
|
||||
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)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, jobs.StatusDead, got.Status)
|
||||
@@ -204,7 +247,10 @@ func TestPurge_DeletesOldCompletedAndDead(t *testing.T) {
|
||||
deadOld := mkOld(jobs.StatusDead)
|
||||
freshComp, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5)
|
||||
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))
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -11,35 +11,49 @@ import (
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
)
|
||||
|
||||
const completeJob = `-- name: CompleteJob :exec
|
||||
const completeJob = `-- name: CompleteJob :execrows
|
||||
UPDATE jobs
|
||||
SET status='completed', completed_at=now(),
|
||||
leased_at=NULL, leased_by=NULL, last_error=NULL,
|
||||
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 {
|
||||
_, err := q.db.Exec(ctx, completeJob, id)
|
||||
return err
|
||||
type CompleteJobParams struct {
|
||||
ID pgtype.UUID `json:"id"`
|
||||
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
|
||||
SET status='dead', completed_at=now(),
|
||||
leased_at=NULL, leased_by=NULL, last_error=$2,
|
||||
updated_at=now()
|
||||
WHERE id=$1
|
||||
WHERE id=$1 AND leased_by=$3 AND status='running'
|
||||
`
|
||||
|
||||
type DeadLetterJobParams struct {
|
||||
ID pgtype.UUID `json:"id"`
|
||||
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)
|
||||
return err
|
||||
// 围栏谓词同上。
|
||||
func (q *Queries) DeadLetterJob(ctx context.Context, arg DeadLetterJobParams) (int64, error) {
|
||||
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
|
||||
@@ -83,25 +97,35 @@ func (q *Queries) EnqueueJob(ctx context.Context, arg EnqueueJobParams) (Job, er
|
||||
return i, err
|
||||
}
|
||||
|
||||
const failJobWithRetry = `-- name: FailJobWithRetry :exec
|
||||
const failJobWithRetry = `-- name: FailJobWithRetry :execrows
|
||||
UPDATE jobs
|
||||
SET status='pending',
|
||||
scheduled_at=$2,
|
||||
leased_at=NULL, leased_by=NULL,
|
||||
last_error=$3,
|
||||
updated_at=now()
|
||||
WHERE id=$1
|
||||
WHERE id=$1 AND leased_by=$4 AND status='running'
|
||||
`
|
||||
|
||||
type FailJobWithRetryParams struct {
|
||||
ID pgtype.UUID `json:"id"`
|
||||
ScheduledAt pgtype.Timestamptz `json:"scheduled_at"`
|
||||
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)
|
||||
return err
|
||||
// 围栏谓词同上。
|
||||
func (q *Queries) FailJobWithRetry(ctx context.Context, arg FailJobWithRetryParams) (int64, error) {
|
||||
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
|
||||
|
||||
+41
-5
@@ -20,6 +20,12 @@ type WorkerOptions struct {
|
||||
Backoff []time.Duration
|
||||
Clock clock.Clock
|
||||
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
|
||||
// (either because IsPermanent(err) or attempts exhausted). The hook runs in
|
||||
// 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]
|
||||
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", "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)
|
||||
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.Info("job.complete",
|
||||
@@ -113,7 +123,9 @@ func (w *Worker) execute(ctx context.Context, job Job) {
|
||||
// Permanent or attempts exhausted -> dead-letter
|
||||
if IsPermanent(err) || job.Attempts >= job.MaxAttempts {
|
||||
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",
|
||||
@@ -126,7 +138,9 @@ func (w *Worker) execute(ctx context.Context, job Job) {
|
||||
delay := backoffFor(w.opts.Backoff, int(job.Attempts))
|
||||
next := w.opts.Clock.Now().Add(delay)
|
||||
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.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))
|
||||
}
|
||||
|
||||
// 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) {
|
||||
if w.opts.HandlerTimeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, w.opts.HandlerTimeout)
|
||||
defer cancel()
|
||||
}
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
err = fmt.Errorf("handler panic: %v", r)
|
||||
|
||||
+17
-8
@@ -144,14 +144,14 @@ func postToken(deps AdminHandlerDeps) http.HandlerFunc {
|
||||
}
|
||||
|
||||
type tokenRow struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Issuer string `json:"issuer"`
|
||||
Scope Scope `json:"scope"`
|
||||
ExpiresAt *string `json:"expires_at,omitempty"`
|
||||
RevokedAt *string `json:"revoked_at,omitempty"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Issuer string `json:"issuer"`
|
||||
Scope Scope `json:"scope"`
|
||||
ExpiresAt *string `json:"expires_at,omitempty"`
|
||||
RevokedAt *string `json:"revoked_at,omitempty"`
|
||||
LastUsedAt *string `json:"last_used_at,omitempty"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
}
|
||||
|
||||
func listTokens(deps AdminHandlerDeps) http.HandlerFunc {
|
||||
@@ -200,7 +200,7 @@ func listTokens(deps AdminHandlerDeps) http.HandlerFunc {
|
||||
func deleteToken(deps AdminHandlerDeps) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
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 {
|
||||
httpx.WriteError(w, rid, err)
|
||||
return
|
||||
@@ -211,6 +211,15 @@ func deleteToken(deps AdminHandlerDeps) http.HandlerFunc {
|
||||
httpx.WriteError(w, rid, errs.Wrap(err, errs.CodeInvalidInput, "id parse"))
|
||||
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 {
|
||||
httpx.WriteError(w, rid, err)
|
||||
return
|
||||
|
||||
@@ -8,11 +8,15 @@ RETURNING id, token_hash, user_id, name, issuer, scope,
|
||||
created_by, created_at;
|
||||
|
||||
-- name: GetMcpTokenByHash :one
|
||||
SELECT id, token_hash, user_id, name, issuer, scope,
|
||||
acp_session_id, expires_at, last_used_at, revoked_at,
|
||||
created_by, created_at
|
||||
FROM mcp_tokens
|
||||
WHERE token_hash = $1;
|
||||
-- 鉴权用:仅返回属于 enabled 用户的 token,禁用用户的 MCP token 立即失效
|
||||
-- (与 web session 解析 GetSessionByTokenHash 的 enabled=true 约束一致)。
|
||||
SELECT t.id, t.token_hash, t.user_id, t.name, t.issuer, t.scope,
|
||||
t.acp_session_id, t.expires_at, t.last_used_at, t.revoked_at,
|
||||
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
|
||||
SELECT id, token_hash, user_id, name, issuer, scope,
|
||||
|
||||
@@ -74,13 +74,17 @@ func (q *Queries) DeleteMcpTokenByID(ctx context.Context, id pgtype.UUID) error
|
||||
}
|
||||
|
||||
const getMcpTokenByHash = `-- name: GetMcpTokenByHash :one
|
||||
SELECT id, token_hash, user_id, name, issuer, scope,
|
||||
acp_session_id, expires_at, last_used_at, revoked_at,
|
||||
created_by, created_at
|
||||
FROM mcp_tokens
|
||||
WHERE token_hash = $1
|
||||
SELECT t.id, t.token_hash, t.user_id, t.name, t.issuer, t.scope,
|
||||
t.acp_session_id, t.expires_at, t.last_used_at, t.revoked_at,
|
||||
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
|
||||
`
|
||||
|
||||
// 鉴权用:仅返回属于 enabled 用户的 token,禁用用户的 MCP token 立即失效
|
||||
// (与 web session 解析 GetSessionByTokenHash 的 enabled=true 约束一致)。
|
||||
func (q *Queries) GetMcpTokenByHash(ctx context.Context, tokenHash []byte) (McpToken, error) {
|
||||
row := q.db.QueryRow(ctx, getMcpTokenByHash, tokenHash)
|
||||
var i McpToken
|
||||
|
||||
@@ -208,6 +208,7 @@ type fakeGit struct {
|
||||
cloneErr error
|
||||
fetchErr error
|
||||
cloneSeen []string
|
||||
removed []string
|
||||
}
|
||||
|
||||
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) MergeFFOnly(_ 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) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
wt, err := s.repo.GetWorktreeByID(ctx, wtID)
|
||||
ws, err := s.authorizeWorktreeWrite(ctx, c, wtID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ws, err := s.repo.GetWorkspaceByID(ctx, wt.WorkspaceID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, canWrite, err := s.pa.ResolveByID(ctx, c.UserID, c.IsAdmin, ws.ProjectID); err != nil {
|
||||
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")
|
||||
}
|
||||
// 先 set pruning,再调 git rm,然后 DELETE row。中途失败 row 留 pruning。
|
||||
// check-and-transition 必须原子:在 tx 内对行加 FOR UPDATE,确认 idle 后才转 pruning。
|
||||
// 否则并发 Acquire(FOR UPDATE)可能刚把它置为 active,而这里用非锁读看到旧的 idle,
|
||||
// 进而 git rm 掉 agent 正在使用的目录。仅 idle→pruning 在锁下完成。
|
||||
var wtPath string
|
||||
err = s.repo.InTx(ctx, func(tx Tx) error {
|
||||
_, e := tx.SetWorktreePruning(ctx, wtID)
|
||||
return e
|
||||
wt, e := tx.GetWorktreeByIDForUpdate(ctx, wtID)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
if wt.Status != WorktreeStatusIdle {
|
||||
return errs.New(errs.CodeWorktreeAlreadyActive, "release before delete")
|
||||
}
|
||||
pruned, e := tx.SetWorktreePruning(ctx, wtID)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
wtPath = pruned.Path
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
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")
|
||||
}
|
||||
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) {
|
||||
if _, err := s.authorizeWorktreeWrite(ctx, c, wtID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
holder := c.Holder()
|
||||
var out *Worktree
|
||||
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) {
|
||||
if _, err := s.authorizeWorktreeWrite(ctx, c, wtID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
holder := c.Holder()
|
||||
var out *Worktree
|
||||
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)
|
||||
}
|
||||
|
||||
// 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) {
|
||||
_ = s.rec.Record(ctx, audit.Entry{
|
||||
UserID: uid, Action: action, TargetType: ttype, TargetID: tid, Metadata: meta,
|
||||
|
||||
@@ -6,8 +6,17 @@ import (
|
||||
|
||||
"github.com/google/uuid"
|
||||
"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) {
|
||||
t.Helper()
|
||||
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)
|
||||
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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user