diff --git a/internal/acp/handlers/server_fs_read.go b/internal/acp/handlers/server_fs_read.go index 82b818b..91b4d63 100644 --- a/internal/acp/handlers/server_fs_read.go +++ b/internal/acp/handlers/server_fs_read.go @@ -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 } } diff --git a/internal/acp/handlers/server_fs_read_test.go b/internal/acp/handlers/server_fs_read_test.go index 4498fd2..5707049 100644 --- a/internal/acp/handlers/server_fs_read_test.go +++ b/internal/acp/handlers/server_fs_read_test.go @@ -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) diff --git a/internal/acp/permission_service.go b/internal/acp/permission_service.go index 11ce8b3..2aea883 100644 --- a/internal/acp/permission_service.go +++ b/internal/acp/permission_service.go @@ -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 } diff --git a/internal/acp/permission_service_internal_test.go b/internal/acp/permission_service_internal_test.go index d854790..2e53bdc 100644 --- a/internal/acp/permission_service_internal_test.go +++ b/internal/acp/permission_service_internal_test.go @@ -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 伪装的调用不应命中白名单(含通配符)") + } +} diff --git a/internal/acp/relay.go b/internal/acp/relay.go index 78f7a46..bce8fcb 100644 --- a/internal/acp/relay.go +++ b/internal/acp/relay.go @@ -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": diff --git a/internal/acp/supervisor.go b/internal/acp/supervisor.go index 6c864ee..65bf907 100644 --- a/internal/acp/supervisor.go +++ b/internal/acp/supervisor.go @@ -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) } diff --git a/internal/chat/streamer.go b/internal/chat/streamer.go index efd839c..5674d44 100644 --- a/internal/chat/streamer.go +++ b/internal/chat/streamer.go @@ -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. diff --git a/internal/chat/streamer_test.go b/internal/chat/streamer_test.go index b1186e9..daeecb0 100644 --- a/internal/chat/streamer_test.go +++ b/internal/chat/streamer_test.go @@ -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 diff --git a/internal/infra/proc/proc_unix.go b/internal/infra/proc/proc_unix.go index 00ac383..c4c6a13 100644 --- a/internal/infra/proc/proc_unix.go +++ b/internal/infra/proc/proc_unix.go @@ -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) } } diff --git a/internal/infra/proc/proc_windows.go b/internal/infra/proc/proc_windows.go index 2dcc2d3..f63a570 100644 --- a/internal/infra/proc/proc_windows.go +++ b/internal/infra/proc/proc_windows.go @@ -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 + } + }) } diff --git a/internal/jobs/domain.go b/internal/jobs/domain.go index 436aa35..bb5e3c0 100644 --- a/internal/jobs/domain.go +++ b/internal/jobs/domain.go @@ -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 { diff --git a/internal/jobs/module.go b/internal/jobs/module.go index 1a3d962..4810da0 100644 --- a/internal/jobs/module.go +++ b/internal/jobs/module.go @@ -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 diff --git a/internal/jobs/queries/jobs.sql b/internal/jobs/queries/jobs.sql index b697b8e..7725c1a 100644 --- a/internal/jobs/queries/jobs.sql +++ b/internal/jobs/queries/jobs.sql @@ -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; diff --git a/internal/jobs/repository.go b/internal/jobs/repository.go index 02b8cb4..7644beb 100644 --- a/internal/jobs/repository.go +++ b/internal/jobs/repository.go @@ -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 } diff --git a/internal/jobs/repository_test.go b/internal/jobs/repository_test.go index 8e1db30..8611cab 100644 --- a/internal/jobs/repository_test.go +++ b/internal/jobs/repository_test.go @@ -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) diff --git a/internal/jobs/sqlc/jobs.sql.go b/internal/jobs/sqlc/jobs.sql.go index 0e7c18e..d63fc87 100644 --- a/internal/jobs/sqlc/jobs.sql.go +++ b/internal/jobs/sqlc/jobs.sql.go @@ -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 diff --git a/internal/jobs/worker.go b/internal/jobs/worker.go index 6116e8b..466c88a 100644 --- a/internal/jobs/worker.go +++ b/internal/jobs/worker.go @@ -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) diff --git a/internal/mcp/handler.go b/internal/mcp/handler.go index defd2ca..958cad1 100644 --- a/internal/mcp/handler.go +++ b/internal/mcp/handler.go @@ -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 diff --git a/internal/mcp/queries/mcp_tokens.sql b/internal/mcp/queries/mcp_tokens.sql index defacfe..b80d361 100644 --- a/internal/mcp/queries/mcp_tokens.sql +++ b/internal/mcp/queries/mcp_tokens.sql @@ -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, diff --git a/internal/mcp/sqlc/mcp_tokens.sql.go b/internal/mcp/sqlc/mcp_tokens.sql.go index f545e32..526857c 100644 --- a/internal/mcp/sqlc/mcp_tokens.sql.go +++ b/internal/mcp/sqlc/mcp_tokens.sql.go @@ -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 diff --git a/internal/workspace/workspace_service_test.go b/internal/workspace/workspace_service_test.go index 8989957..209a4ae 100644 --- a/internal/workspace/workspace_service_test.go +++ b/internal/workspace/workspace_service_test.go @@ -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 } diff --git a/internal/workspace/worktree_service.go b/internal/workspace/worktree_service.go index 97c5372..4b22486 100644 --- a/internal/workspace/worktree_service.go +++ b/internal/workspace/worktree_service.go @@ -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, diff --git a/internal/workspace/worktree_service_test.go b/internal/workspace/worktree_service_test.go index c4aca9e..6794a36 100644 --- a/internal/workspace/worktree_service_test.go +++ b/internal/workspace/worktree_service_test.go @@ -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) +} diff --git a/migrations/0015_fix_issue_workspace_fk.down.sql b/migrations/0015_fix_issue_workspace_fk.down.sql new file mode 100644 index 0000000..ad08134 --- /dev/null +++ b/migrations/0015_fix_issue_workspace_fk.down.sql @@ -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; diff --git a/migrations/0015_fix_issue_workspace_fk.up.sql b/migrations/0015_fix_issue_workspace_fk.up.sql new file mode 100644 index 0000000..ebbcf8d --- /dev/null +++ b/migrations/0015_fix_issue_workspace_fk.up.sql @@ -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); diff --git a/web/src/composables/useConversationSession.ts b/web/src/composables/useConversationSession.ts index 0e090f8..ce7ae67 100644 --- a/web/src/composables/useConversationSession.ts +++ b/web/src/composables/useConversationSession.ts @@ -114,6 +114,14 @@ export function useConversationSession() { const id = streamingMessageId.value if (id == null) return 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() }