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