This commit is contained in:
2026-06-22 08:55:57 +08:00
parent dbb87823e8
commit 6ade6e8fa9
325 changed files with 41131 additions and 855 deletions
+9 -6
View File
@@ -1,9 +1,11 @@
// agent_home.go 实现 agent CLI 受管 home 目录的物化与重定向 env 注入。
//
// 每个 AgentKind 一个受管 home:<agent_homes_dir>/<agent_kind_id>/。spawn 前把
// DB 中的配置文件解密落盘(全量覆写受管文件,不清空目录 —— agent 运行期写入的
// 缓存 / 凭据得以保留),再按 client_type 注入配置目录重定向变量。受管 home 位于
// DataDir 下,容器场景随持久卷存活。
// 受管 home 按 user + agent-kind 维度隔离:<agent_homes_dir>/<user_id>/<agent_kind_id>/。
// 早期实现仅按 agent-kind 维度(多个用户共享同一 home 与缓存凭据),沙箱化要求
// per-USER home,故 home 路径以 user_id 为一级目录、agent_kind_id 为二级目录。
// spawn 前把 DB 中的配置文件解密落盘(全量覆写受管文件,不清空目录 —— agent 运行期
// 写入的缓存 / 凭据得以保留),再按 client_type 注入配置目录重定向变量。受管 home 位于
// DataDir 下,容器场景随持久卷存活。旧的 per-kind 目录无害遗留,下次 spawn 用新路径。
package acp
import (
@@ -16,9 +18,10 @@ import (
)
// materializeAgentHome 创建受管 home 并物化该 AgentKind 的全部配置文件。
// home 按 user + agent-kind 隔离:<AgentHomesDir>/<user_id>/<kind_id>。
// 返回 home 绝对路径。配置文件为空时仍创建目录(agent 需要可写 home)。
func (s *Supervisor) materializeAgentHome(ctx context.Context, kind *AgentKind) (string, error) {
home, err := filepath.Abs(filepath.Join(s.cfg.AgentHomesDir, kind.ID.String()))
func (s *Supervisor) materializeAgentHome(ctx context.Context, sess *Session, kind *AgentKind) (string, error) {
home, err := filepath.Abs(filepath.Join(s.cfg.AgentHomesDir, sess.UserID.String(), kind.ID.String()))
if err != nil {
return "", fmt.Errorf("resolve agent home: %w", err)
}
+39 -4
View File
@@ -41,16 +41,50 @@ func TestMaterializeAgentHome(t *testing.T) {
cfg: SupervisorConfig{AgentHomesDir: root},
}
home, err := s.materializeAgentHome(context.Background(), kind)
sess := &Session{ID: uuid.New(), UserID: uuid.New()}
home, err := s.materializeAgentHome(context.Background(), sess, kind)
require.NoError(t, err)
assert.Equal(t, filepath.Join(root, kind.ID.String()), home)
// per-user home:<root>/<user_id>/<kind_id>
assert.Equal(t, filepath.Join(root, sess.UserID.String(), kind.ID.String()), home)
got, err := os.ReadFile(filepath.Join(home, ".codex", "config.toml"))
require.NoError(t, err)
assert.Equal(t, content, got)
// 二次物化幂等(覆写)
_, err = s.materializeAgentHome(context.Background(), kind)
_, err = s.materializeAgentHome(context.Background(), sess, kind)
require.NoError(t, err)
}
func TestMaterializeAgentHome_PerUserIsolation(t *testing.T) {
t.Parallel()
enc, err := crypto.NewEncryptor([]byte("0123456789abcdef0123456789abcdef"))
require.NoError(t, err)
kind := &AgentKind{ID: uuid.New(), ClientType: ClientCodex}
encrypted, err := enc.Encrypt([]byte("x = 1"))
require.NoError(t, err)
root := t.TempDir()
s := &Supervisor{
repo: cfgRepoStub{files: []*ConfigFile{{AgentKindID: kind.ID, RelPath: ".codex/config.toml", EncryptedContent: encrypted}}},
crypto: enc,
cfg: SupervisorConfig{AgentHomesDir: root},
}
// 同一 agent-kind、两个不同用户 → 两个互不相同的 home 目录。
sessA := &Session{ID: uuid.New(), UserID: uuid.New()}
sessB := &Session{ID: uuid.New(), UserID: uuid.New()}
homeA, err := s.materializeAgentHome(context.Background(), sessA, kind)
require.NoError(t, err)
homeB, err := s.materializeAgentHome(context.Background(), sessB, kind)
require.NoError(t, err)
assert.NotEqual(t, homeA, homeB)
// 两个 home 各自物化了配置文件。
_, err = os.Stat(filepath.Join(homeA, ".codex", "config.toml"))
require.NoError(t, err)
_, err = os.Stat(filepath.Join(homeB, ".codex", "config.toml"))
require.NoError(t, err)
}
@@ -67,7 +101,8 @@ func TestMaterializeAgentHome_RejectsEscape(t *testing.T) {
crypto: enc,
cfg: SupervisorConfig{AgentHomesDir: t.TempDir()},
}
_, err = s.materializeAgentHome(context.Background(), kind)
sess := &Session{ID: uuid.New(), UserID: uuid.New()}
_, err = s.materializeAgentHome(context.Background(), sess, kind)
require.Error(t, err)
assert.Contains(t, err.Error(), "escapes agent home")
}
+80 -3
View File
@@ -12,9 +12,16 @@ import (
)
type agentKindService struct {
repo Repository
enc *crypto.Encryptor
audit audit.Recorder
repo Repository
enc *crypto.Encryptor
audit audit.Recorder
modelCheck ModelValidator // 校验 model_id 引用一个启用的 llm_model;可为 nil(跳过校验)
}
// ModelValidator 校验某 model_id 引用一个存在且启用的 llm_model。装配期由 app.go
// 注入一个委托给 chat 的 adapter,避免 acp 构造期直接依赖 chat。
type ModelValidator interface {
IsEnabledModel(ctx context.Context, modelID uuid.UUID) (bool, error)
}
// NewAgentKindService 构造 AgentKindService。
@@ -22,6 +29,9 @@ func NewAgentKindService(repo Repository, enc *crypto.Encryptor, rec audit.Recor
return &agentKindService{repo: repo, enc: enc, audit: rec}
}
// SetModelValidator 注入 model 校验器(装配期回填)。
func (s *agentKindService) SetModelValidator(v ModelValidator) { s.modelCheck = v }
func (s *agentKindService) Create(ctx context.Context, c Caller, in CreateAgentKindInput) (*AgentKind, error) {
if !c.IsAdmin {
return nil, errs.New(errs.CodeForbidden, "admin only")
@@ -52,6 +62,12 @@ func (s *agentKindService) Create(ctx context.Context, c Caller, in CreateAgentK
if allowlist == nil {
allowlist = []string{}
}
if err := s.validateBudget(in.MaxCostUSD, in.MaxTokens, in.MaxWallClockSeconds); err != nil {
return nil, err
}
if err := s.validateModel(ctx, in.ModelID); err != nil {
return nil, err
}
k := &AgentKind{
ID: uuid.New(), Name: in.Name, DisplayName: in.DisplayName, Description: in.Description,
BinaryPath: in.BinaryPath, Args: args, EncryptedEnv: encrypted, Enabled: in.Enabled,
@@ -59,6 +75,10 @@ func (s *agentKindService) Create(ctx context.Context, c Caller, in CreateAgentK
ClientType: clientType,
EncryptedMCPServers: encryptedMCP,
CreatedBy: c.UserID,
ModelID: in.ModelID,
MaxCostUSD: in.MaxCostUSD,
MaxTokens: in.MaxTokens,
MaxWallClockSeconds: in.MaxWallClockSeconds,
}
out, err := s.repo.CreateAgentKind(ctx, k)
if err != nil {
@@ -146,6 +166,34 @@ func (s *agentKindService) Update(ctx context.Context, c Caller, id uuid.UUID, i
cur.Enabled = *in.Enabled
changed["enabled"] = *in.Enabled
}
if in.SetModelID {
if err := s.validateModel(ctx, in.ModelID); err != nil {
return nil, err
}
cur.ModelID = in.ModelID
changed["model_id"] = "<changed>"
}
if in.SetMaxCostUSD {
if err := s.validateBudget(in.MaxCostUSD, nil, nil); err != nil {
return nil, err
}
cur.MaxCostUSD = in.MaxCostUSD
changed["max_cost_usd"] = "<changed>"
}
if in.SetMaxTokens {
if err := s.validateBudget(nil, in.MaxTokens, nil); err != nil {
return nil, err
}
cur.MaxTokens = in.MaxTokens
changed["max_tokens"] = "<changed>"
}
if in.SetMaxWallClock {
if err := s.validateBudget(nil, nil, in.MaxWallClockSeconds); err != nil {
return nil, err
}
cur.MaxWallClockSeconds = in.MaxWallClockSeconds
changed["max_wall_clock_seconds"] = "<changed>"
}
if len(changed) == 0 {
return cur, nil
}
@@ -174,6 +222,35 @@ func (s *agentKindService) Delete(ctx context.Context, c Caller, id uuid.UUID) e
return nil
}
// validateBudget 校验预算上限:非空值必须为正。
func (s *agentKindService) validateBudget(maxCost *float64, maxTokens *int64, maxWall *int32) error {
if maxCost != nil && *maxCost <= 0 {
return errs.New(errs.CodeInvalidInput, "max_cost_usd must be positive")
}
if maxTokens != nil && *maxTokens <= 0 {
return errs.New(errs.CodeInvalidInput, "max_tokens must be positive")
}
if maxWall != nil && *maxWall <= 0 {
return errs.New(errs.CodeInvalidInput, "max_wall_clock_seconds must be positive")
}
return nil
}
// validateModel 校验 model_id 引用一个启用的 llm_model(modelCheck 为 nil 时跳过)。
func (s *agentKindService) validateModel(ctx context.Context, modelID *uuid.UUID) error {
if modelID == nil || s.modelCheck == nil {
return nil
}
ok, err := s.modelCheck.IsEnabledModel(ctx, *modelID)
if err != nil {
return err
}
if !ok {
return errs.New(errs.CodeInvalidInput, "model_id does not reference an enabled model")
}
return nil
}
func (s *agentKindService) recordAudit(ctx context.Context, c Caller, action, targetID string, meta map[string]any) {
if s.audit == nil {
return
+69
View File
@@ -116,6 +116,15 @@ func (f *fakeAgentKindRepo) InsertSession(context.Context, *acp.Session) (*acp.S
func (f *fakeAgentKindRepo) GetSessionByID(context.Context, uuid.UUID) (*acp.Session, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) GetSessionByStepID(context.Context, uuid.UUID) (*acp.Session, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) GetCrashedSessionForResume(context.Context, uuid.UUID) (*acp.Session, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) ResetSessionForResume(context.Context, uuid.UUID) error {
panic("n/a")
}
func (f *fakeAgentKindRepo) ListSessionsByUser(context.Context, uuid.UUID, *uuid.UUID) ([]*acp.Session, error) {
panic("n/a")
}
@@ -125,6 +134,9 @@ func (f *fakeAgentKindRepo) ListAllSessions(context.Context, *uuid.UUID) ([]*acp
func (f *fakeAgentKindRepo) UpdateSessionRunning(context.Context, uuid.UUID, string, int32) error {
panic("n/a")
}
func (f *fakeAgentKindRepo) UpdateSessionSandboxMode(context.Context, uuid.UUID, string) error {
panic("n/a")
}
func (f *fakeAgentKindRepo) UpdateSessionFinished(context.Context, uuid.UUID, acp.SessionStatus, *int32, *string) error {
panic("n/a")
}
@@ -155,6 +167,63 @@ func (f *fakeAgentKindRepo) ListEventsSince(context.Context, uuid.UUID, int64, i
func (f *fakeAgentKindRepo) PurgeEventsBefore(context.Context, time.Time) (int64, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) InsertSessionUsage(context.Context, *acp.SessionUsageRecord) (bool, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) AddSessionUsageTotals(context.Context, uuid.UUID, int64, int64, int64, float64) (int64, int64, int64, float64, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) SumProjectCostUSD(context.Context, uuid.UUID, time.Time) (float64, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) MarkSessionTerminatedReason(context.Context, uuid.UUID, string) error {
panic("n/a")
}
func (f *fakeAgentKindRepo) ListSessionsForReaper(context.Context, time.Time, time.Time) ([]acp.ReaperSession, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) ListSessionUsage(context.Context, uuid.UUID, int32) ([]*acp.SessionUsageRecord, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) SummarizeAcpUsageByUser(context.Context, time.Time, time.Time) ([]acp.AcpUsageUserRow, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) SummarizeAcpUsageByModel(context.Context, time.Time, time.Time) ([]acp.AcpUsageModelRow, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) SummarizeAcpUsageByDay(context.Context, time.Time, time.Time) ([]acp.AcpUsageDayRow, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) SessionRollup(context.Context, acp.DashboardFilter) (acp.SessionRollup, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) SessionsByDay(context.Context, acp.DashboardFilter) ([]acp.SessionDayRow, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) SessionTimeline(context.Context, uuid.UUID) ([]acp.SessionTimelineBucket, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) InsertTurn(context.Context, *acp.Turn) (*acp.Turn, error) { panic("n/a") }
func (f *fakeAgentKindRepo) MarkTurnCompleted(context.Context, uuid.UUID, int, string, int) error {
panic("n/a")
}
func (f *fakeAgentKindRepo) MarkTurnAborted(context.Context, uuid.UUID, int) error { panic("n/a") }
func (f *fakeAgentKindRepo) IncrementTurnUpdateCount(context.Context, uuid.UUID, int) error {
panic("n/a")
}
func (f *fakeAgentKindRepo) GetLatestTurnBySession(context.Context, uuid.UUID) (*acp.Turn, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) ListTurnsBySession(context.Context, uuid.UUID) ([]*acp.Turn, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) UpdateSessionLastStopReason(context.Context, uuid.UUID, string) error {
panic("n/a")
}
func (f *fakeAgentKindRepo) PurgeTurnsBefore(context.Context, time.Time) (int64, error) {
panic("n/a")
}
func (f *fakeAgentKindRepo) InTx(context.Context, func(context.Context, pgx.Tx) error) error {
panic("n/a")
+137
View File
@@ -0,0 +1,137 @@
package acp_test
import (
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/go-chi/chi/v5"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/acp"
)
// mountDashboardHandler builds a Handler whose SessionService is backed by the
// in-memory fakeAcpRepo so run-history endpoints can be exercised without a DB.
func mountDashboardHandler(t *testing.T, callerUID uuid.UUID, isAdmin bool, repo *fakeAcpRepo) chi.Router {
t.Helper()
enc := testEncryptor(t)
ss := acp.NewSessionService(repo, nil, nil, nil, nil, nil, nil, nil, nil,
acp.SessionServiceConfig{}, nil)
r := chi.NewRouter()
h := acp.NewHandler(
nil, // akSvc — not exercised
ss,
nil, // sup — not exercised
repo,
nil, // permSvc — not exercised
fakeResolver{uid: callerUID},
fakeAdminLookup{admin: map[uuid.UUID]bool{callerUID: isAdmin}},
enc,
acp.WSConfig{},
)
h.Mount(r)
return r
}
func getJSONArray(t *testing.T, r chi.Router, path, token string) (*httptest.ResponseRecorder, []map[string]any) {
t.Helper()
req := httptest.NewRequest("GET", path, nil)
if token != "" {
req.Header.Set("Authorization", "Bearer "+token)
}
rr := httptest.NewRecorder()
r.ServeHTTP(rr, req)
var out []map[string]any
if rr.Body.Len() > 0 {
_ = json.Unmarshal(rr.Body.Bytes(), &out)
}
return rr, out
}
func TestDashboardHandler_Unauthenticated_401(t *testing.T) {
t.Parallel()
repo := &fakeAcpRepo{}
r := mountDashboardHandler(t, uuid.New(), false, repo)
rr, _ := doJSON(t, r, "GET", "/api/v1/acp/dashboard", "", nil)
assert.Equal(t, http.StatusUnauthorized, rr.Code)
}
func TestDashboardHandler_OwnerScoped(t *testing.T) {
t.Parallel()
uid := uuid.New()
repo := &fakeAcpRepo{
rollupResult: acp.SessionRollup{Total: 2, Succeeded: 2, TotalCostUSD: 0.4},
byDayResult: []acp.SessionDayRow{{Day: time.Now().UTC(), Total: 2, Succeeded: 2, TotalCostUSD: 0.4}},
}
r := mountDashboardHandler(t, uid, false, repo)
rr, body := doJSON(t, r, "GET", "/api/v1/acp/dashboard", "valid", nil)
require.Equal(t, http.StatusOK, rr.Code, "body=%s", rr.Body.String())
rollup := body["rollup"].(map[string]any)
assert.Equal(t, float64(2), rollup["total"])
assert.Equal(t, float64(1), rollup["success_rate"])
// 非 admin → repo filter scoped to caller.
require.NotNil(t, repo.rollupFilter.UserID)
assert.Equal(t, uid, *repo.rollupFilter.UserID)
}
func TestDashboardHandler_Admin_NotScoped_WithProjectFilter(t *testing.T) {
t.Parallel()
uid := uuid.New()
pid := uuid.New()
repo := &fakeAcpRepo{rollupResult: acp.SessionRollup{Total: 5, Succeeded: 4}}
r := mountDashboardHandler(t, uid, true, repo)
rr, _ := doJSON(t, r, "GET", "/api/v1/acp/dashboard?project_id="+pid.String(), "valid", nil)
require.Equal(t, http.StatusOK, rr.Code)
assert.Nil(t, repo.rollupFilter.UserID) // admin → 全量
require.NotNil(t, repo.rollupFilter.ProjectID)
assert.Equal(t, pid, *repo.rollupFilter.ProjectID)
}
func TestDashboardHandler_BadProjectID_400(t *testing.T) {
t.Parallel()
repo := &fakeAcpRepo{}
r := mountDashboardHandler(t, uuid.New(), true, repo)
rr, _ := doJSON(t, r, "GET", "/api/v1/acp/dashboard?project_id=not-a-uuid", "valid", nil)
assert.Equal(t, http.StatusBadRequest, rr.Code)
}
func TestTimelineHandler_NonOwner_NotFound(t *testing.T) {
t.Parallel()
sid := uuid.New()
repo := &fakeAcpRepo{}
repo.sessions = []*acp.Session{{ID: sid, UserID: uuid.New()}}
r := mountDashboardHandler(t, uuid.New(), false, repo)
rr, _ := doJSON(t, r, "GET", "/api/v1/acp/sessions/"+sid.String()+"/timeline", "valid", nil)
assert.Equal(t, http.StatusNotFound, rr.Code)
}
func TestTimelineHandler_Owner_OK(t *testing.T) {
t.Parallel()
uid := uuid.New()
sid := uuid.New()
repo := &fakeAcpRepo{
timelineResult: []acp.SessionTimelineBucket{
{Direction: "out", RPCKind: "request", Method: "session/prompt", EventCount: 3},
},
}
repo.sessions = []*acp.Session{{ID: sid, UserID: uid}}
r := mountDashboardHandler(t, uid, false, repo)
rr, arr := getJSONArray(t, r, "/api/v1/acp/sessions/"+sid.String()+"/timeline", "valid")
require.Equal(t, http.StatusOK, rr.Code, "body=%s", rr.Body.String())
require.Len(t, arr, 1)
assert.Equal(t, "session/prompt", arr[0]["method"])
assert.Equal(t, "request", arr[0]["kind"])
assert.Equal(t, sid, repo.timelineSID)
}
+242
View File
@@ -118,6 +118,12 @@ type AgentKind struct {
CreatedBy uuid.UUID
CreatedAt time.Time
UpdatedAt time.Time
// 成本核算:关联的 llm_models 行(价格来源)+ 每 kind 默认预算上限。
// 全部 NULLABLE:未配置 model 时成本按 0/agent 自报;未配置上限时无限制。
ModelID *uuid.UUID
MaxCostUSD *float64
MaxTokens *int64
MaxWallClockSeconds *int32
}
// ConfigFile 是 AgentKind 维度的 agent CLI 配置文件(如 .claude/settings.json、
@@ -153,6 +159,191 @@ type Session struct {
LastError *string
StartedAt time.Time
EndedAt *time.Time
// LastStopReason 是 agent 最近一次回合完成时上报的 stopReason 原始字符串
// (nil = 尚无完成的回合)。归一化枚举见 NormalizeStopReason,但落库为原值。
LastStopReason *string
// OrchestratorStepID 非 nil 时表示该 session 由编排器某个 step 驱动;onExit /
// 启动 reaper 据此把崩溃 session 映射回 step 并触发再驱动。manual session 为 nil。
OrchestratorStepID *uuid.UUID
// 成本核算运行时累加器(落库快照)。
PromptTokens int64
CompletionTokens int64
ThinkingTokens int64
TotalCostUSD float64
LastActivityAt time.Time
// 创建时快照的有效预算上限(session 覆盖 > agent-kind 默认 > 全局默认)。
BudgetMaxCostUSD *float64
BudgetMaxTokens *int64
BudgetMaxWallClockSeconds *int32
// TerminatedReason 非 nil 时记录会话被强制终止的原因(budget_*/reaper_*/manual)。
TerminatedReason *string
}
// BudgetCaps 是一次会话的有效预算上限快照。nil 字段 = 该维度无限制。
type BudgetCaps struct {
MaxCostUSD *float64
MaxTokens *int64
MaxWallClockSeconds *int32
}
// BreachReason 标识预算被突破的维度,落 acp_sessions.terminated_reason。
type BreachReason string
const (
BreachCost BreachReason = "budget_cost"
BreachTokens BreachReason = "budget_tokens"
BreachWallClock BreachReason = "budget_wall_clock"
)
// SessionUsageRecord 是 acp_session_usage 表一行(per-turn 账目,镜像 llm_usage)。
type SessionUsageRecord struct {
ID int64
SessionID uuid.UUID
UserID uuid.UUID
ProjectID uuid.UUID
AgentKindID uuid.UUID
ModelID *uuid.UUID
PromptTokens int64
CompletionTokens int64
ThinkingTokens int64
CostUSD float64
SourceEventID *int64
CreatedAt time.Time
}
// AcpUsageUserRow 按用户聚合的 ACP 用量。
type AcpUsageUserRow struct {
UserID uuid.UUID
PromptTokens, CompletionTokens, ThinkingTokens int64
CostUSD float64
}
// AcpUsageModelRow 按模型聚合的 ACP 用量。
type AcpUsageModelRow struct {
ModelID uuid.UUID
PromptTokens, CompletionTokens, ThinkingTokens int64
CostUSD float64
}
// AcpUsageDayRow 按天聚合的 ACP 用量。
type AcpUsageDayRow struct {
Day time.Time
PromptTokens, CompletionTokens, ThinkingTokens int64
CostUSD float64
}
// ===== run-history dashboard 读模型 =====
// DashboardFilter 限定 run-history 汇总的范围。UserID 非 nil 时仅统计该用户的
// 会话(owner scope);nil 表示全量(admin scope)。ProjectID/RequirementID
// 为可选过滤;From/To 为半开时间窗 [From, To)。
type DashboardFilter struct {
UserID *uuid.UUID
ProjectID *uuid.UUID
RequirementID *uuid.UUID
From time.Time
To time.Time
}
// SessionRollup 是 run-history 汇总:会话计数 + 成功率 + 总成本 + 平均时长。
type SessionRollup struct {
Total int64
Succeeded int64
Crashed int64
Active int64
TotalCostUSD float64
TotalTokens int64
AvgDurationSeconds float64
}
// SuccessRate 返回 succeeded/total(total=0 时为 0),便于 handler/MCP 直接使用。
func (r SessionRollup) SuccessRate() float64 {
if r.Total == 0 {
return 0
}
return float64(r.Succeeded) / float64(r.Total)
}
// SessionDayRow 是 run-history 按天的时间序列一行。
type SessionDayRow struct {
Day time.Time
Total int64
Succeeded int64
Crashed int64
TotalCostUSD float64
}
// SessionTimelineBucket 是单会话事件按 (direction, rpc_kind, method) 分桶的一行。
type SessionTimelineBucket struct {
Direction string
RPCKind string
Method string
EventCount int64
FirstAt time.Time
LastAt time.Time
}
// ReaperSession 是 reaper 需要的最小 session 投影。
type ReaperSession struct {
ID uuid.UUID
UserID uuid.UUID
StartedAt time.Time
LastActivityAt time.Time
MaxWallClockSeconds *int32
}
// StopReason 是回合完成的归一化原因枚举。agent 可能上报枚举外的 vendor-specific
// 值;这些值归一化为 StopOther,但原始字符串仍会原样落库(stop_reason /
// last_stop_reason),避免信息丢失。
type StopReason string
const (
StopEndTurn StopReason = "end_turn"
StopMaxTokens StopReason = "max_tokens"
StopRefusal StopReason = "refusal"
StopCancelled StopReason = "cancelled"
StopOther StopReason = "other"
)
// NormalizeStopReason 把 agent 上报的原始 stopReason 映射到归一化枚举。
// 未知值(含空串)映射为 StopOther。
func NormalizeStopReason(raw string) StopReason {
switch StopReason(raw) {
case StopEndTurn:
return StopEndTurn
case StopMaxTokens:
return StopMaxTokens
case StopRefusal:
return StopRefusal
case StopCancelled:
return StopCancelled
default:
return StopOther
}
}
// TurnStatus 是 acp_turns.status 枚举。
type TurnStatus string
const (
TurnInProgress TurnStatus = "in_progress"
TurnCompleted TurnStatus = "completed"
TurnAborted TurnStatus = "aborted"
)
// Turn 是一次 agent 回合的审计记录。一个 session 内 turn_index 单调递增(0-based)。
// PromptRequestID 是发起该回合的 session/prompt 的 JSON-RPC id(用于响应相关联)。
// StopReason 在完成前为 nil;存储 agent 上报的原始字符串。
type Turn struct {
ID int64
SessionID uuid.UUID
TurnIndex int
PromptRequestID *string
Status TurnStatus
StopReason *string
UpdateCount int
StartedAt time.Time
CompletedAt *time.Time
}
// Event 是一条入库的 JSON-RPC 消息。Payload 是 JSON-RPC envelope 完整 JSON
@@ -204,6 +395,39 @@ type SessionService interface {
Get(ctx context.Context, c Caller, id uuid.UUID) (*Session, error)
List(ctx context.Context, c Caller, f SessionListFilter) ([]*Session, error)
Terminate(ctx context.Context, c Caller, id uuid.UUID) error
// Dashboard 返回 run-history 汇总 + 按天序列。非 admin 调用强制按 caller
// scope(UserID 覆盖为 caller),admin 才能查看全量。
Dashboard(ctx context.Context, c Caller, in DashboardInput) (*DashboardResult, error)
// Timeline 返回单会话的事件时间线(先经 Get 做 owner/admin 访问校验)。
Timeline(ctx context.Context, c Caller, sessionID uuid.UUID) ([]SessionTimelineBucket, error)
}
// DashboardInput 是 run-history 汇总的入参(过滤维度由 service 在 scope 内归一化)。
type DashboardInput struct {
ProjectID *uuid.UUID
RequirementID *uuid.UUID
From time.Time
To time.Time
}
// DashboardResult 是 run-history 汇总响应:总览 + 按天序列。
type DashboardResult struct {
Rollup SessionRollup
ByDay []SessionDayRow
From time.Time
To time.Time
}
// TurnRunner 是编排器驱动一个阻塞回合所需的窄接口(不扩大 SessionService 的
// HTTP/MCP 暴露面)。sessionService 同时实现 SessionService 与 TurnRunner。
//
// - SendPromptAndWait 发送一条 session/prompt 并阻塞等待回合完成,返回原始
// stopReason(来自 relay.Call 的 server-initiated 请求响应)。
// - ResumeSession 在崩溃 session 的同一 cwd/worktree 上重新拉起 agent。
type TurnRunner interface {
SendPromptAndWait(ctx context.Context, sessionID uuid.UUID, prompt string) (stopReason string, err error)
ResumeSession(ctx context.Context, c Caller, sessionID uuid.UUID) (*Session, error)
}
// SessionListFilter 控制 sessions 列表的过滤维度。
@@ -222,6 +446,9 @@ type CreateSessionInput struct {
IssueNumber *int
RequirementNumber *int
InitialPrompt *string
// OrchestratorStepID 非 nil 时把 session 与编排器 step 反向关联(仅由 orchestrator
// 内部调用设置;HTTP/MCP 入口不暴露)。
OrchestratorStepID *uuid.UUID
}
// CreateAgentKindInput 携带创建必需字段。Args / Env 可为 nil。
@@ -237,6 +464,11 @@ type CreateAgentKindInput struct {
ToolAllowlist []string
ClientType ClientType
MCPServers []McpServerSpec
// 成本核算:model 价格来源 + 默认预算上限(均可选)。
ModelID *uuid.UUID
MaxCostUSD *float64
MaxTokens *int64
MaxWallClockSeconds *int32
}
// UpdateAgentKindInput 是 patch 输入。Name 不可改(创建后稳定,引用约束)。
@@ -251,6 +483,16 @@ type UpdateAgentKindInput struct {
ToolAllowlist []string // nil = 不改;非 nil(含空)= 替换
ClientType *ClientType // nil = 不改
MCPServers []McpServerSpec // nil = 不改;非 nil(含空)= 替换
// 成本核算字段:三态指针,nil = 不改。SetModelID 控制 ModelID 是否参与更新
// (ModelID 本身可被显式置空,故用独立 flag 区分"不改"与"清空")。
SetModelID bool
ModelID *uuid.UUID
SetMaxCostUSD bool
MaxCostUSD *float64
SetMaxTokens bool
MaxTokens *int64
SetMaxWallClock bool
MaxWallClockSeconds *int32
}
// PermissionStatus 是 acp_permission_requests.status 枚举。
+171 -3
View File
@@ -21,9 +21,14 @@ type AgentKindAdminDTO struct {
ClientType string `json:"client_type"`
// MCPServers 明文返回(admin only,编辑所需);条目可能含 header/env 密钥。
MCPServers []McpServerSpec `json:"mcp_servers"`
CreatedBy uuid.UUID `json:"created_by"`
CreatedAt string `json:"created_at"`
UpdatedAt string `json:"updated_at"`
// 成本核算:model 价格来源 + 默认预算上限(均可空)。
ModelID *string `json:"model_id"`
MaxCostUSD *float64 `json:"max_cost_usd"`
MaxTokens *int64 `json:"max_tokens"`
MaxWallClockSeconds *int32 `json:"max_wall_clock_seconds"`
CreatedBy uuid.UUID `json:"created_by"`
CreatedAt string `json:"created_at"`
UpdatedAt string `json:"updated_at"`
}
// AgentKindPublicDTO 是普通用户视图,不暴露 binary_path / args / env_keys。
@@ -46,6 +51,11 @@ type CreateAgentKindReq struct {
ToolAllowlist []string `json:"tool_allowlist"`
ClientType string `json:"client_type"`
MCPServers []McpServerSpec `json:"mcp_servers"`
// 成本核算:model_id 为字符串 UUID(空/缺省 = 不设置);预算上限均可选。
ModelID *string `json:"model_id"`
MaxCostUSD *float64 `json:"max_cost_usd"`
MaxTokens *int64 `json:"max_tokens"`
MaxWallClockSeconds *int32 `json:"max_wall_clock_seconds"`
}
type UpdateAgentKindReq struct {
@@ -58,6 +68,12 @@ type UpdateAgentKindReq struct {
ToolAllowlist []string `json:"tool_allowlist,omitempty"`
ClientType *string `json:"client_type,omitempty"`
MCPServers []McpServerSpec `json:"mcp_servers,omitempty"` // nil = 不改;非 nil(含空)= 替换
// 成本核算三态:字段存在(非 nil)= 设置;缺省(nil)= 不改。
// ModelID 空串 = 清空;预算上限值 <=0 = 清空。
ModelID *string `json:"model_id,omitempty"`
MaxCostUSD *float64 `json:"max_cost_usd,omitempty"`
MaxTokens *int64 `json:"max_tokens,omitempty"`
MaxWallClockSeconds *int32 `json:"max_wall_clock_seconds,omitempty"`
}
// ConfigFileDTO 是 AgentKind 配置文件的对外视图(admin only,content 明文)。
@@ -133,6 +149,15 @@ type SessionDTO struct {
LastError *string `json:"last_error"`
StartedAt string `json:"started_at"`
EndedAt *string `json:"ended_at"`
// 成本核算:运行时累计 + 有效预算上限 + 终止原因。
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
TotalCostUSD float64 `json:"total_cost_usd"`
BudgetMaxCostUSD *float64 `json:"budget_max_cost_usd"`
BudgetMaxTokens *int64 `json:"budget_max_tokens"`
BudgetMaxWallClockSeconds *int32 `json:"budget_max_wall_clock_seconds"`
TerminatedReason *string `json:"terminated_reason"`
}
func sessionToDTO(s *Session) SessionDTO {
@@ -167,5 +192,148 @@ func sessionToDTO(s *Session) SessionDTO {
LastError: s.LastError,
StartedAt: s.StartedAt.UTC().Format("2006-01-02T15:04:05Z"),
EndedAt: ended,
PromptTokens: s.PromptTokens,
CompletionTokens: s.CompletionTokens,
ThinkingTokens: s.ThinkingTokens,
TotalCostUSD: s.TotalCostUSD,
BudgetMaxCostUSD: s.BudgetMaxCostUSD,
BudgetMaxTokens: s.BudgetMaxTokens,
BudgetMaxWallClockSeconds: s.BudgetMaxWallClockSeconds,
TerminatedReason: s.TerminatedReason,
}
}
// SessionUsageDTO 是单条 per-turn 账目的对外视图。
type SessionUsageDTO struct {
ID int64 `json:"id"`
ModelID *string `json:"model_id"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
CostUSD float64 `json:"cost_usd"`
CreatedAt string `json:"created_at"`
}
// SessionUsageResponse 是 GET /sessions/{id}/usage 的响应:会话累计 + 最近账目。
type SessionUsageResponse struct {
SessionID string `json:"session_id"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
TotalCostUSD float64 `json:"total_cost_usd"`
BudgetMaxCostUSD *float64 `json:"budget_max_cost_usd"`
BudgetMaxTokens *int64 `json:"budget_max_tokens"`
Recent []SessionUsageDTO `json:"recent"`
}
// AcpUsageItem 是 admin ACP 用量聚合的一行(镜像 chat AdminUsageItem)。
type AcpUsageItem struct {
Key string `json:"key"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
CostUSD float64 `json:"cost_usd"`
}
// AcpUsageResponse 包裹分组用量项。
type AcpUsageResponse struct {
Items []AcpUsageItem `json:"items"`
}
// ===== run-history dashboard DTO =====
// DashboardResponse 是 GET /api/v1/acp/dashboard 的响应:总览 + 按天序列。
type DashboardResponse struct {
From string `json:"from"`
To string `json:"to"`
Rollup SessionRollupDTO `json:"rollup"`
ByDay []SessionDayDTO `json:"by_day"`
}
// SessionRollupDTO 是会话总览(计数 + 成功率 + 成本 + 平均时长)。
type SessionRollupDTO struct {
Total int64 `json:"total"`
Succeeded int64 `json:"succeeded"`
Crashed int64 `json:"crashed"`
Active int64 `json:"active"`
SuccessRate float64 `json:"success_rate"`
TotalCostUSD float64 `json:"total_cost_usd"`
TotalTokens int64 `json:"total_tokens"`
AvgDurationSeconds float64 `json:"avg_duration_seconds"`
}
// SessionDayDTO 是按天的时间序列一行。
type SessionDayDTO struct {
Day string `json:"day"`
Total int64 `json:"total"`
Succeeded int64 `json:"succeeded"`
Crashed int64 `json:"crashed"`
TotalCostUSD float64 `json:"total_cost_usd"`
}
// SessionTimelineDTO 是单会话事件时间线一行(按 method/kind 分桶)。
type SessionTimelineDTO struct {
Direction string `json:"direction"`
Kind string `json:"kind"`
Method string `json:"method"`
EventCount int64 `json:"event_count"`
FirstAt string `json:"first_at"`
LastAt string `json:"last_at"`
}
func dashboardToResponse(r *DashboardResult) DashboardResponse {
resp := DashboardResponse{
From: r.From.UTC().Format(time.RFC3339),
To: r.To.UTC().Format(time.RFC3339),
Rollup: SessionRollupDTO{
Total: r.Rollup.Total,
Succeeded: r.Rollup.Succeeded,
Crashed: r.Rollup.Crashed,
Active: r.Rollup.Active,
SuccessRate: r.Rollup.SuccessRate(),
TotalCostUSD: r.Rollup.TotalCostUSD,
TotalTokens: r.Rollup.TotalTokens,
AvgDurationSeconds: r.Rollup.AvgDurationSeconds,
},
ByDay: make([]SessionDayDTO, 0, len(r.ByDay)),
}
for _, d := range r.ByDay {
resp.ByDay = append(resp.ByDay, SessionDayDTO{
Day: d.Day.UTC().Format(time.RFC3339),
Total: d.Total,
Succeeded: d.Succeeded,
Crashed: d.Crashed,
TotalCostUSD: d.TotalCostUSD,
})
}
return resp
}
func sessionTimelineToDTO(b SessionTimelineBucket) SessionTimelineDTO {
return SessionTimelineDTO{
Direction: b.Direction,
Kind: b.RPCKind,
Method: b.Method,
EventCount: b.EventCount,
FirstAt: b.FirstAt.UTC().Format(time.RFC3339),
LastAt: b.LastAt.UTC().Format(time.RFC3339),
}
}
func sessionUsageToDTO(r *SessionUsageRecord) SessionUsageDTO {
var modelID *string
if r.ModelID != nil {
v := r.ModelID.String()
modelID = &v
}
return SessionUsageDTO{
ID: r.ID,
ModelID: modelID,
PromptTokens: r.PromptTokens,
CompletionTokens: r.CompletionTokens,
ThinkingTokens: r.ThinkingTokens,
CostUSD: r.CostUSD,
CreatedAt: r.CreatedAt.UTC().Format(time.RFC3339),
}
}
+329 -25
View File
@@ -6,6 +6,7 @@ package acp
import (
"context"
"encoding/json"
"log/slog"
"net/http"
"sort"
"strconv"
@@ -83,12 +84,27 @@ func (h *Handler) Mount(r chi.Router) {
r.Get("/{id}/events", h.getEvents)
r.Get("/{id}/ws", h.sessionWS)
r.Get("/{id}/permissions", h.listSessionPermissions)
r.Get("/{id}/usage", h.getSessionUsage)
// run-history 单会话事件时间线(owner 或 admin)。
r.Get("/{id}/timeline", h.getSessionTimeline)
})
// run-history dashboard 汇总(owner scope;admin 可全量)。
r.Group(func(r chi.Router) {
r.Use(middleware.Auth(h.resolver, middleware.AuthOptions{}))
r.Get("/api/v1/acp/dashboard", h.getDashboard)
})
r.Route("/api/v1/acp/permissions", func(r chi.Router) {
r.Use(middleware.Auth(h.resolver, middleware.AuthOptions{}))
// 跨会话审批 inbox:调用者所有会话的待审权限请求(HITL)。
r.Get("/", h.listPendingPermissions)
r.Post("/{reqID}/approve", h.approvePermission)
r.Post("/{reqID}/deny", h.denyPermission)
})
// Admin ACP 用量聚合(镜像 chat adminUsage);service 层 requireAdmin。
r.Group(func(r chi.Router) {
r.Use(middleware.Auth(h.resolver, middleware.AuthOptions{}))
r.Get("/api/v1/acp/usage", h.adminAcpUsage)
})
}
// listSessionPermissions 返回某会话的待审权限请求(owner 或 admin)。
@@ -115,6 +131,25 @@ func (h *Handler) listSessionPermissions(w http.ResponseWriter, r *http.Request)
httpx.WriteJSON(w, http.StatusOK, out)
}
// listPendingPermissions 返回调用者跨所有会话的待审权限请求(HITL inbox)。
func (h *Handler) listPendingPermissions(w http.ResponseWriter, r *http.Request) {
c, err := h.caller(r)
if err != nil {
writeErr(w, r, err)
return
}
reqs, err := h.permSvc.ListPendingForUser(r.Context(), c)
if err != nil {
writeErr(w, r, err)
return
}
out := make([]PermissionRequestDTO, 0, len(reqs))
for _, p := range reqs {
out = append(out, permissionRequestToDTO(p))
}
httpx.WriteJSON(w, http.StatusOK, out)
}
type approvePermissionReq struct {
OptionID string `json:"option_id"`
}
@@ -205,17 +240,26 @@ func (h *Handler) createAgentKind(w http.ResponseWriter, r *http.Request) {
if req.Enabled != nil {
enabled = *req.Enabled
}
modelID, perr := parseOptionalUUID(req.ModelID)
if perr != nil {
writeErr(w, r, errs.New(errs.CodeInvalidInput, "bad model_id"))
return
}
in := CreateAgentKindInput{
Name: strings.TrimSpace(req.Name),
DisplayName: req.DisplayName,
Description: req.Description,
BinaryPath: strings.TrimSpace(req.BinaryPath),
Args: req.Args,
Env: req.Env,
Enabled: enabled,
ToolAllowlist: req.ToolAllowlist,
ClientType: ClientType(req.ClientType),
MCPServers: req.MCPServers,
Name: strings.TrimSpace(req.Name),
DisplayName: req.DisplayName,
Description: req.Description,
BinaryPath: strings.TrimSpace(req.BinaryPath),
Args: req.Args,
Env: req.Env,
Enabled: enabled,
ToolAllowlist: req.ToolAllowlist,
ClientType: ClientType(req.ClientType),
MCPServers: req.MCPServers,
ModelID: modelID,
MaxCostUSD: req.MaxCostUSD,
MaxTokens: req.MaxTokens,
MaxWallClockSeconds: req.MaxWallClockSeconds,
}
out, err := h.akSvc.Create(r.Context(), c, in)
if err != nil {
@@ -278,6 +322,35 @@ func (h *Handler) updateAgentKind(w http.ResponseWriter, r *http.Request) {
ct := ClientType(*req.ClientType)
in.ClientType = &ct
}
if req.ModelID != nil {
in.SetModelID = true
if *req.ModelID == "" {
in.ModelID = nil // 显式清空
} else {
mid, merr := uuid.Parse(*req.ModelID)
if merr != nil {
writeErr(w, r, errs.New(errs.CodeInvalidInput, "bad model_id"))
return
}
in.ModelID = &mid
}
}
if req.MaxCostUSD != nil {
in.SetMaxCostUSD = true
in.MaxCostUSD = positiveOrNil(req.MaxCostUSD)
}
if req.MaxTokens != nil {
in.SetMaxTokens = true
if *req.MaxTokens > 0 {
in.MaxTokens = req.MaxTokens
}
}
if req.MaxWallClockSeconds != nil {
in.SetMaxWallClock = true
if *req.MaxWallClockSeconds > 0 {
in.MaxWallClockSeconds = req.MaxWallClockSeconds
}
}
out, err := h.akSvc.Update(r.Context(), c, id, in)
if err != nil {
writeErr(w, r, err)
@@ -457,6 +530,123 @@ func (h *Handler) getSession(w http.ResponseWriter, r *http.Request) {
httpx.WriteJSON(w, http.StatusOK, sessionToDTO(s))
}
// getSessionUsage 返回会话累计成本/token + 最近账目(owner 或 admin)。
func (h *Handler) getSessionUsage(w http.ResponseWriter, r *http.Request) {
c, err := h.caller(r)
if err != nil {
writeErr(w, r, err)
return
}
id, perr := uuid.Parse(chi.URLParam(r, "id"))
if perr != nil {
writeErr(w, r, errs.New(errs.CodeInvalidInput, "bad id"))
return
}
// 复用 sessSvc.Get 做 owner/admin 访问校验(NotFound/Forbidden 由其返回)。
s, err := h.sessSvc.Get(r.Context(), c, id)
if err != nil {
writeErr(w, r, err)
return
}
const recentLimit = 50
ledger, err := h.repo.ListSessionUsage(r.Context(), id, recentLimit)
if err != nil {
writeErr(w, r, err)
return
}
recent := make([]SessionUsageDTO, 0, len(ledger))
for _, rec := range ledger {
recent = append(recent, sessionUsageToDTO(rec))
}
httpx.WriteJSON(w, http.StatusOK, SessionUsageResponse{
SessionID: s.ID.String(),
PromptTokens: s.PromptTokens,
CompletionTokens: s.CompletionTokens,
ThinkingTokens: s.ThinkingTokens,
TotalCostUSD: s.TotalCostUSD,
BudgetMaxCostUSD: s.BudgetMaxCostUSD,
BudgetMaxTokens: s.BudgetMaxTokens,
Recent: recent,
})
}
// adminAcpUsage 返回 ACP 用量聚合(admin only),group=user|model|day,镜像 chat。
func (h *Handler) adminAcpUsage(w http.ResponseWriter, r *http.Request) {
c, err := h.caller(r)
if err != nil {
writeErr(w, r, err)
return
}
if !c.IsAdmin {
writeErr(w, r, errs.New(errs.CodeForbidden, "admin only"))
return
}
q := r.URL.Query()
to := time.Now().UTC()
from := to.AddDate(0, 0, -30)
if v := q.Get("from"); v != "" {
if t, e := time.Parse(time.RFC3339, v); e == nil {
from = t
}
}
if v := q.Get("to"); v != "" {
if t, e := time.Parse(time.RFC3339, v); e == nil {
to = t
}
}
groupBy := q.Get("group")
if groupBy == "" {
groupBy = q.Get("group_by")
}
if groupBy == "" {
groupBy = "day"
}
items := make([]AcpUsageItem, 0)
switch groupBy {
case "user":
rows, rerr := h.repo.SummarizeAcpUsageByUser(r.Context(), from, to)
if rerr != nil {
writeErr(w, r, rerr)
return
}
for _, row := range rows {
items = append(items, AcpUsageItem{
Key: row.UserID.String(), PromptTokens: row.PromptTokens,
CompletionTokens: row.CompletionTokens, ThinkingTokens: row.ThinkingTokens, CostUSD: row.CostUSD,
})
}
case "model":
rows, rerr := h.repo.SummarizeAcpUsageByModel(r.Context(), from, to)
if rerr != nil {
writeErr(w, r, rerr)
return
}
for _, row := range rows {
items = append(items, AcpUsageItem{
Key: row.ModelID.String(), PromptTokens: row.PromptTokens,
CompletionTokens: row.CompletionTokens, ThinkingTokens: row.ThinkingTokens, CostUSD: row.CostUSD,
})
}
case "day":
rows, rerr := h.repo.SummarizeAcpUsageByDay(r.Context(), from, to)
if rerr != nil {
writeErr(w, r, rerr)
return
}
for _, row := range rows {
items = append(items, AcpUsageItem{
Key: row.Day.Format(time.RFC3339), PromptTokens: row.PromptTokens,
CompletionTokens: row.CompletionTokens, ThinkingTokens: row.ThinkingTokens, CostUSD: row.CostUSD,
})
}
default:
writeErr(w, r, errs.New(errs.CodeInvalidInput, "invalid group"))
return
}
httpx.WriteJSON(w, http.StatusOK, AcpUsageResponse{Items: items})
}
func (h *Handler) terminateSession(w http.ResponseWriter, r *http.Request) {
c, err := h.caller(r)
if err != nil {
@@ -519,6 +709,80 @@ func (h *Handler) getEvents(w http.ResponseWriter, r *http.Request) {
httpx.WriteJSON(w, http.StatusOK, out)
}
// getDashboard 返回 run-history 汇总(owner scope;admin 全量)。可选 query:
// project_id / requirement_id / from / to(RFC3339)。
func (h *Handler) getDashboard(w http.ResponseWriter, r *http.Request) {
c, err := h.caller(r)
if err != nil {
writeErr(w, r, err)
return
}
q := r.URL.Query()
var in DashboardInput
if v := q.Get("project_id"); v != "" {
id, perr := uuid.Parse(v)
if perr != nil {
writeErr(w, r, errs.New(errs.CodeInvalidInput, "bad project_id"))
return
}
in.ProjectID = &id
}
if v := q.Get("requirement_id"); v != "" {
id, perr := uuid.Parse(v)
if perr != nil {
writeErr(w, r, errs.New(errs.CodeInvalidInput, "bad requirement_id"))
return
}
in.RequirementID = &id
}
if v := q.Get("from"); v != "" {
t, terr := time.Parse(time.RFC3339, v)
if terr != nil {
writeErr(w, r, errs.New(errs.CodeInvalidInput, "bad from"))
return
}
in.From = t
}
if v := q.Get("to"); v != "" {
t, terr := time.Parse(time.RFC3339, v)
if terr != nil {
writeErr(w, r, errs.New(errs.CodeInvalidInput, "bad to"))
return
}
in.To = t
}
res, err := h.sessSvc.Dashboard(r.Context(), c, in)
if err != nil {
writeErr(w, r, err)
return
}
httpx.WriteJSON(w, http.StatusOK, dashboardToResponse(res))
}
// getSessionTimeline 返回单会话事件时间线(owner 或 admin)。
func (h *Handler) getSessionTimeline(w http.ResponseWriter, r *http.Request) {
c, err := h.caller(r)
if err != nil {
writeErr(w, r, err)
return
}
id, perr := uuid.Parse(chi.URLParam(r, "id"))
if perr != nil {
writeErr(w, r, errs.New(errs.CodeInvalidInput, "bad id"))
return
}
buckets, err := h.sessSvc.Timeline(r.Context(), c, id)
if err != nil {
writeErr(w, r, err)
return
}
out := make([]SessionTimelineDTO, 0, len(buckets))
for _, b := range buckets {
out = append(out, sessionTimelineToDTO(b))
}
httpx.WriteJSON(w, http.StatusOK, out)
}
// ===== WebSocket =====
func (h *Handler) sessionWS(w http.ResponseWriter, r *http.Request) {
@@ -675,6 +939,17 @@ func (h *Handler) sessionWS(w http.ResponseWriter, r *http.Request) {
})
continue
}
// 回合相关联:用户从前端发起 session/prompt 时也登记回合,使其完成时
// 走与 auto-prompt 相同的检测路径。client 自生成的字符串 id 作为相关联键。
if m.Method == "session/prompt" {
if tr := relay.Tracker(); tr != nil {
if pid, ok := m.IDString(); ok && pid != "" {
if _, err := tr.StartTurn(ctx, pid); err != nil {
slog.Warn("acp.ws.start_turn", "session_id", sess.ID, "err", err.Error())
}
}
}
}
if err := relay.SendClient(m); err != nil {
_ = wsjson.Write(ctx, conn, map[string]any{
"kind": "client_error",
@@ -702,22 +977,51 @@ func (h *Handler) agentKindToAdminDTO(k *AgentKind) AgentKindAdminDTO {
if servers, err := decryptMCPServers(h.enc, k.EncryptedMCPServers); err == nil {
mcpServers = servers
}
return AgentKindAdminDTO{
ID: k.ID,
Name: k.Name,
DisplayName: k.DisplayName,
Description: k.Description,
BinaryPath: k.BinaryPath,
Args: append([]string(nil), k.Args...),
EnvKeys: envKeys,
Enabled: k.Enabled,
ToolAllowlist: append([]string(nil), k.ToolAllowlist...),
ClientType: string(k.ClientType),
MCPServers: mcpServers,
CreatedBy: k.CreatedBy,
CreatedAt: k.CreatedAt.UTC().Format("2006-01-02T15:04:05Z"),
UpdatedAt: k.UpdatedAt.UTC().Format("2006-01-02T15:04:05Z"),
var modelID *string
if k.ModelID != nil {
v := k.ModelID.String()
modelID = &v
}
return AgentKindAdminDTO{
ID: k.ID,
Name: k.Name,
DisplayName: k.DisplayName,
Description: k.Description,
BinaryPath: k.BinaryPath,
Args: append([]string(nil), k.Args...),
EnvKeys: envKeys,
Enabled: k.Enabled,
ToolAllowlist: append([]string(nil), k.ToolAllowlist...),
ClientType: string(k.ClientType),
MCPServers: mcpServers,
ModelID: modelID,
MaxCostUSD: k.MaxCostUSD,
MaxTokens: k.MaxTokens,
MaxWallClockSeconds: k.MaxWallClockSeconds,
CreatedBy: k.CreatedBy,
CreatedAt: k.CreatedAt.UTC().Format("2006-01-02T15:04:05Z"),
UpdatedAt: k.UpdatedAt.UTC().Format("2006-01-02T15:04:05Z"),
}
}
// parseOptionalUUID 解析一个可选的字符串 UUID。nil/空串 → (nil, nil);非法 → 错误。
func parseOptionalUUID(s *string) (*uuid.UUID, error) {
if s == nil || *s == "" {
return nil, nil
}
id, err := uuid.Parse(*s)
if err != nil {
return nil, err
}
return &id, nil
}
// positiveOrNil 返回正数指针,否则 nil(用于"值 <=0 视为清空"语义)。
func positiveOrNil(f *float64) *float64 {
if f == nil || *f <= 0 {
return nil
}
return f
}
func agentKindToPublicDTO(k *AgentKind) AgentKindPublicDTO {
+139 -1
View File
@@ -33,6 +33,7 @@ import (
"github.com/yan1h/agent-coding-workflow/internal/acp"
"github.com/yan1h/agent-coding-workflow/internal/audit"
"github.com/yan1h/agent-coding-workflow/internal/infra/crypto"
"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
"github.com/yan1h/agent-coding-workflow/internal/project"
"github.com/yan1h/agent-coding-workflow/internal/workspace"
@@ -152,7 +153,7 @@ func TestHandler_Session_CreateAndWS_E2E(t *testing.T) {
paStub := &stubProjectAccess{pj: &project.Project{ID: pid, Slug: "p-test"}}
svc := acp.NewSessionService(
repo, wsStub, nil, nil, paStub, sup,
repo, wsStub, nil, nil, paStub, nil, nil, sup,
audit.NewPostgresRecorder(pool),
acp.SessionServiceConfig{MaxActiveGlobal: 10, MaxActivePerUser: 5},
nil, // mcpTokens — not exercised by handler e2e
@@ -276,3 +277,140 @@ func TestHandler_Session_CreateAndWS_E2E(t *testing.T) {
}
}
}
// TestSession_TurnCompletion_E2E drives a full session with an InitialPrompt
// against the fake-acp-agent (which emits session/update chunks then a
// session/prompt response with stopReason=end_turn). It asserts:
// - an acp_turns row is created and completed with stop_reason=end_turn
// - update_count == FAKE_ACP_PROMPT_UPDATES
// - acp_sessions.last_stop_reason is set to end_turn
// - a test TurnBus subscriber receives turn_completed + session_idle
func TestSession_TurnCompletion_E2E(t *testing.T) {
if testing.Short() {
t.Skip("e2e: spawns subprocess + testcontainer pg")
}
t.Parallel()
ctx := context.Background()
repo, pool := setupRepo(t)
uid := mustInsertUser(t, ctx, pool, "turn-e2e@local", false)
pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)
cwd := t.TempDir()
_, err := pool.Exec(ctx, `UPDATE workspaces SET main_path = $1 WHERE id = $2`, cwd, wsid)
require.NoError(t, err)
bin := fakeAgentBinary(t)
enc := testEncryptor(t)
const wantUpdates = 4
kind, err := repo.CreateAgentKind(ctx, &acp.AgentKind{
ID: uuid.New(), Name: "turn-e2e-fake", DisplayName: "Fake",
BinaryPath: bin, Args: []string{}, Enabled: true, CreatedBy: uid,
// FAKE_ACP_PROMPT_UPDATES via encrypted env: emit N session/update before response.
EncryptedEnv: mustEncryptEnv(t, enc, map[string]string{"FAKE_ACP_PROMPT_UPDATES": "4"}),
})
require.NoError(t, err)
wsRow, err := workspace.NewPostgresRepository(pool).GetWorkspaceByID(ctx, wsid)
require.NoError(t, err)
disp := notify.NewDispatcher(notifyRepoStub{}, slogDiscard())
sup := acp.NewSupervisor(
repo, audit.NewPostgresRecorder(pool), disp,
nil, nil, nil, enc,
acp.SupervisorConfig{
SpawnTimeout: 5 * time.Second,
KillGrace: 1 * time.Second,
StderrBufferLines: 50,
StderrTailBytes: 1000,
EventMaxPayload: 65536,
EventTruncateField: 4096,
ShutdownGrace: 5 * time.Second,
},
slogDiscard(),
)
defer sup.ShutdownAll(context.Background())
// Test TurnBus subscriber records published events.
bus := acp.NewTurnBus(slogDiscard())
evCh := make(chan acp.TurnEvent, 8)
bus.Subscribe("test", func(_ context.Context, ev acp.TurnEvent) { evCh <- ev })
sup.SetTurnBus(bus)
wsStub := &stubWsService{wsRow: wsRow}
paStub := &stubProjectAccess{pj: &project.Project{ID: pid, Slug: "p-turn"}}
svc := acp.NewSessionService(
repo, wsStub, nil, nil, paStub, nil, nil, sup,
audit.NewPostgresRecorder(pool),
acp.SessionServiceConfig{MaxActiveGlobal: 10, MaxActivePerUser: 5},
nil,
)
permSvc := acp.NewPermissionService(repo, nil, 0, nil)
sup.SetPermissionService(permSvc)
permSvc.SetRelayLocator(sup)
prompt := "do the thing"
branch := "main"
sess, err := svc.Create(ctx, acp.Caller{UserID: uid}, acp.CreateSessionInput{
WorkspaceID: wsid, AgentKindID: kind.ID, Branch: &branch, InitialPrompt: &prompt,
})
require.NoError(t, err)
// Poll for the turn to complete.
var completedTurn *acp.Turn
deadline := time.After(15 * time.Second)
for completedTurn == nil {
ts, _ := repo.ListTurnsBySession(ctx, sess.ID)
for _, tn := range ts {
if tn.Status == acp.TurnCompleted {
completedTurn = tn
}
}
if completedTurn != nil {
break
}
select {
case <-deadline:
t.Fatalf("turn never completed; turns=%+v", ts)
case <-time.After(150 * time.Millisecond):
}
}
require.NotNil(t, completedTurn.StopReason)
assert.Equal(t, "end_turn", *completedTurn.StopReason)
assert.Equal(t, wantUpdates, completedTurn.UpdateCount)
got, _ := repo.GetSessionByID(ctx, sess.ID)
require.NotNil(t, got.LastStopReason)
assert.Equal(t, "end_turn", *got.LastStopReason)
// Collect bus events: expect a turn_completed and a session_idle.
var sawCompleted, sawIdle bool
evDeadline := time.After(3 * time.Second)
for !(sawCompleted && sawIdle) {
select {
case ev := <-evCh:
switch ev.Kind {
case acp.TurnCompletedEvent:
sawCompleted = true
assert.Equal(t, acp.StopEndTurn, ev.StopReason)
case acp.SessionIdleEvent:
sawIdle = true
}
case <-evDeadline:
t.Fatalf("did not receive both turn_completed and session_idle (completed=%v idle=%v)", sawCompleted, sawIdle)
}
}
}
// mustEncryptEnv marshals env to JSON and encrypts it (mirrors acp.encryptEnv,
// which is unexported), so the supervisor can decrypt it back into the spawn env.
func mustEncryptEnv(t *testing.T, enc *crypto.Encryptor, env map[string]string) []byte {
t.Helper()
b, err := json.Marshal(env)
require.NoError(t, err)
out, err := enc.Encrypt(b)
require.NoError(t, err)
return out
}
+8 -1
View File
@@ -43,9 +43,16 @@ type FsError struct {
type FsHandler struct {
Read *FsReadHandler
Write *FsWriteHandler
List *FsListHandler
Grep *GrepHandler
}
// NewFsHandler constructs a composite FsHandler with default sub-handlers.
func NewFsHandler() *FsHandler {
return &FsHandler{Read: NewFsReadHandler(), Write: NewFsWriteHandler()}
return &FsHandler{
Read: NewFsReadHandler(),
Write: NewFsWriteHandler(),
List: NewFsListHandler(),
Grep: NewGrepHandler(),
}
}
+88
View File
@@ -0,0 +1,88 @@
package handlers
import (
"context"
"encoding/json"
"errors"
"io/fs"
"os"
"sort"
)
// DefaultListMaxEntries caps the number of directory entries returned to bound
// payload size for large directories.
const DefaultListMaxEntries = 1000
// FsListHandler handles fs/list: a bounded directory listing inside the worktree.
type FsListHandler struct {
MaxEntries int
}
// NewFsListHandler constructs an FsListHandler with default caps.
func NewFsListHandler() *FsListHandler { return &FsListHandler{MaxEntries: DefaultListMaxEntries} }
func (h *FsListHandler) maxEntries() int {
if h.MaxEntries > 0 {
return h.MaxEntries
}
return DefaultListMaxEntries
}
type fsListParams struct {
SessionID string `json:"sessionId"`
Path string `json:"path"`
}
type fsDirEntry struct {
Name string `json:"name"`
IsDir bool `json:"is_dir"`
Size int64 `json:"size"`
}
type fsListResult struct {
Entries []fsDirEntry `json:"entries"`
Truncated bool `json:"truncated"`
}
// Handle lists the directory at sess.CwdPath/path. Returns FsError for path
// escape (-32001) or read failure. Results are sorted (dirs first, then name).
func (h *FsListHandler) Handle(_ context.Context, sess SessionContext, paramsJSON json.RawMessage) FsResult {
var p fsListParams
if err := json.Unmarshal(paramsJSON, &p); err != nil {
return FsResult{Err: &FsError{Code: -32602, Message: "invalid params: " + err.Error()}}
}
abs, err := resolveSafePath(sess.CwdPath, p.Path)
if err != nil {
return FsResult{Err: &FsError{Code: -32001, Message: err.Error()}}
}
dirents, err := os.ReadDir(abs)
if err != nil {
if errors.Is(err, fs.ErrNotExist) {
return FsResult{Err: &FsError{Code: -32603, Message: "directory not found"}}
}
return FsResult{Err: &FsError{Code: -32603, Message: "list: " + err.Error()}}
}
max := h.maxEntries()
out := fsListResult{Entries: make([]fsDirEntry, 0, len(dirents))}
for _, de := range dirents {
if len(out.Entries) >= max {
out.Truncated = true
break
}
e := fsDirEntry{Name: de.Name(), IsDir: de.IsDir()}
if info, ierr := de.Info(); ierr == nil && !de.IsDir() {
e.Size = info.Size()
}
out.Entries = append(out.Entries, e)
}
sort.SliceStable(out.Entries, func(i, j int) bool {
if out.Entries[i].IsDir != out.Entries[j].IsDir {
return out.Entries[i].IsDir // dirs first
}
return out.Entries[i].Name < out.Entries[j].Name
})
res, _ := json.Marshal(out)
return FsResult{OK: res}
}
@@ -0,0 +1,82 @@
package handlers_test
import (
"context"
"encoding/json"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
)
func TestFsList_HappyPath(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
require.NoError(t, os.WriteFile(filepath.Join(tmp, "a.txt"), []byte("hi"), 0o644))
require.NoError(t, os.MkdirAll(filepath.Join(tmp, "sub"), 0o755))
h := handlers.NewFsListHandler()
sess := handlers.SessionContext{CwdPath: tmp}
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "path": "."}))
require.Nil(t, res.Err)
var out struct {
Entries []struct {
Name string `json:"name"`
IsDir bool `json:"is_dir"`
Size int64 `json:"size"`
} `json:"entries"`
Truncated bool `json:"truncated"`
}
require.NoError(t, json.Unmarshal(res.OK, &out))
require.Len(t, out.Entries, 2)
// Dirs sort first.
assert.Equal(t, "sub", out.Entries[0].Name)
assert.True(t, out.Entries[0].IsDir)
assert.Equal(t, "a.txt", out.Entries[1].Name)
assert.Equal(t, int64(2), out.Entries[1].Size)
}
func TestFsList_PathOutsideWorktree(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
h := handlers.NewFsListHandler()
sess := handlers.SessionContext{CwdPath: tmp}
outside := filepath.Join(filepath.VolumeName(tmp)+string(filepath.Separator), "etc")
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "path": outside}))
require.NotNil(t, res.Err)
assert.Equal(t, -32001, res.Err.Code)
}
func TestFsList_NotFound(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
h := handlers.NewFsListHandler()
sess := handlers.SessionContext{CwdPath: tmp}
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "path": "missing"}))
require.NotNil(t, res.Err)
assert.Equal(t, -32603, res.Err.Code)
}
func TestFsList_MaxEntriesCap(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
for i := 0; i < 5; i++ {
require.NoError(t, os.WriteFile(filepath.Join(tmp, string(rune('a'+i))+".txt"), []byte("x"), 0o644))
}
h := &handlers.FsListHandler{MaxEntries: 2}
sess := handlers.SessionContext{CwdPath: tmp}
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "path": "."}))
require.Nil(t, res.Err)
var out struct {
Entries []any `json:"entries"`
Truncated bool `json:"truncated"`
}
require.NoError(t, json.Unmarshal(res.OK, &out))
require.Len(t, out.Entries, 2)
require.True(t, out.Truncated)
}
+209
View File
@@ -0,0 +1,209 @@
package handlers
import (
"bufio"
"context"
"encoding/json"
"io/fs"
"os"
"path/filepath"
"regexp"
"strings"
)
// Grep caps. These bound the cost of a single grep request to avoid a DoS via a
// crafted pattern / huge tree. All are overridable on the handler.
const (
DefaultGrepMaxMatches = 200
DefaultGrepMaxFiles = 5000
DefaultGrepMaxFileSize = 1 << 20 // 1 MiB: skip larger files
DefaultGrepMaxLineLen = 1000 // truncate long matched lines
)
// GrepHandler handles fs/grep: a bounded regexp search under the worktree using a
// pure-Go walk (no shell, no git grep interpolation) confined to the pathsafe
// root. Binary files and oversize files are skipped.
type GrepHandler struct {
MaxMatches int
MaxFiles int
MaxFileSize int64
}
// NewGrepHandler constructs a GrepHandler with default caps.
func NewGrepHandler() *GrepHandler {
return &GrepHandler{
MaxMatches: DefaultGrepMaxMatches,
MaxFiles: DefaultGrepMaxFiles,
MaxFileSize: DefaultGrepMaxFileSize,
}
}
func (h *GrepHandler) maxMatches() int {
if h.MaxMatches > 0 {
return h.MaxMatches
}
return DefaultGrepMaxMatches
}
func (h *GrepHandler) maxFiles() int {
if h.MaxFiles > 0 {
return h.MaxFiles
}
return DefaultGrepMaxFiles
}
func (h *GrepHandler) maxFileSize() int64 {
if h.MaxFileSize > 0 {
return h.MaxFileSize
}
return DefaultGrepMaxFileSize
}
type fsGrepParams struct {
SessionID string `json:"sessionId"`
Pattern string `json:"pattern"`
Path string `json:"path,omitempty"` // sub-path to scope the search; defaults to cwd root
CaseSensitive bool `json:"case_sensitive,omitempty"`
}
type grepMatch struct {
File string `json:"file"` // worktree-relative
Line int `json:"line"` // 1-based
Text string `json:"text"`
}
type fsGrepResult struct {
Matches []grepMatch `json:"matches"`
Truncated bool `json:"truncated"`
}
// Handle runs a regexp search. Returns FsError for path escape (-32001), invalid
// pattern (-32602), or read failure. Match/file caps are honored.
func (h *GrepHandler) Handle(ctx context.Context, sess SessionContext, paramsJSON json.RawMessage) FsResult {
var p fsGrepParams
if err := json.Unmarshal(paramsJSON, &p); err != nil {
return FsResult{Err: &FsError{Code: -32602, Message: "invalid params: " + err.Error()}}
}
if strings.TrimSpace(p.Pattern) == "" {
return FsResult{Err: &FsError{Code: -32602, Message: "pattern required"}}
}
expr := p.Pattern
if !p.CaseSensitive {
expr = "(?i)" + expr
}
re, err := regexp.Compile(expr)
if err != nil {
return FsResult{Err: &FsError{Code: -32602, Message: "invalid pattern: " + err.Error()}}
}
searchPath := p.Path
if searchPath == "" {
searchPath = "."
}
root, err := resolveSafePath(sess.CwdPath, searchPath)
if err != nil {
return FsResult{Err: &FsError{Code: -32001, Message: err.Error()}}
}
cwd := filepath.Clean(sess.CwdPath)
out := fsGrepResult{Matches: make([]grepMatch, 0, 16)}
maxMatches := h.maxMatches()
maxFiles := h.maxFiles()
maxSize := h.maxFileSize()
filesScanned := 0
walkErr := filepath.WalkDir(root, func(path string, d fs.DirEntry, werr error) error {
if werr != nil {
return nil // skip unreadable entries
}
select {
case <-ctx.Done():
return filepath.SkipAll
default:
}
if d.IsDir() {
if isSkippedDir(d.Name()) {
return filepath.SkipDir
}
return nil
}
if filesScanned >= maxFiles {
out.Truncated = true
return filepath.SkipAll
}
info, ierr := d.Info()
if ierr != nil || info.Size() > maxSize {
return nil
}
filesScanned++
rel, rerr := filepath.Rel(cwd, path)
if rerr != nil {
rel = path
}
rel = filepath.ToSlash(rel)
if grepFile(path, rel, re, &out, maxMatches) {
out.Truncated = true
return filepath.SkipAll
}
return nil
})
if walkErr != nil && walkErr != filepath.SkipAll {
return FsResult{Err: &FsError{Code: -32603, Message: "grep walk: " + walkErr.Error()}}
}
res, _ := json.Marshal(out)
return FsResult{OK: res}
}
// grepFile scans one file line-by-line. Returns true when the global match cap
// was reached (caller should stop the walk).
func grepFile(abs, rel string, re *regexp.Regexp, out *fsGrepResult, maxMatches int) bool {
f, err := os.Open(abs)
if err != nil {
return false
}
defer f.Close()
sc := bufio.NewScanner(f)
sc.Buffer(make([]byte, 0, 64*1024), 1<<20)
lineNo := 0
for sc.Scan() {
lineNo++
line := sc.Bytes()
// Skip binary content (NUL in line).
if hasNUL(line) {
return false
}
if re.Match(line) {
text := string(line)
if len(text) > DefaultGrepMaxLineLen {
text = text[:DefaultGrepMaxLineLen] + "...[truncated]"
}
out.Matches = append(out.Matches, grepMatch{File: rel, Line: lineNo, Text: text})
if len(out.Matches) >= maxMatches {
return true
}
}
}
return false
}
func hasNUL(b []byte) bool {
for _, c := range b {
if c == 0 {
return true
}
}
return false
}
// skippedDirs are never descended into during grep.
var skippedDirs = map[string]struct{}{
".git": {}, "node_modules": {}, "vendor": {}, "dist": {}, "build": {},
".venv": {}, "__pycache__": {}, "third_party": {},
}
func isSkippedDir(name string) bool {
_, ok := skippedDirs[name]
return ok
}
+125
View File
@@ -0,0 +1,125 @@
package handlers_test
import (
"context"
"encoding/json"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
)
func TestGrep_HappyPath(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
require.NoError(t, os.WriteFile(filepath.Join(tmp, "a.go"), []byte("package a\n\nfunc Foo() {}\n"), 0o644))
require.NoError(t, os.WriteFile(filepath.Join(tmp, "b.go"), []byte("package b\n\nfunc Bar() {}\n"), 0o644))
h := handlers.NewGrepHandler()
sess := handlers.SessionContext{CwdPath: tmp}
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "pattern": "func Foo"}))
require.Nil(t, res.Err)
var out struct {
Matches []struct {
File string `json:"file"`
Line int `json:"line"`
Text string `json:"text"`
} `json:"matches"`
Truncated bool `json:"truncated"`
}
require.NoError(t, json.Unmarshal(res.OK, &out))
require.Len(t, out.Matches, 1)
assert.Equal(t, "a.go", out.Matches[0].File)
assert.Equal(t, 3, out.Matches[0].Line)
assert.Contains(t, out.Matches[0].Text, "func Foo")
}
func TestGrep_CaseInsensitiveByDefault(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
require.NoError(t, os.WriteFile(filepath.Join(tmp, "a.go"), []byte("HELLO world\n"), 0o644))
h := handlers.NewGrepHandler()
sess := handlers.SessionContext{CwdPath: tmp}
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "pattern": "hello"}))
require.Nil(t, res.Err)
var out struct {
Matches []any `json:"matches"`
}
require.NoError(t, json.Unmarshal(res.OK, &out))
require.Len(t, out.Matches, 1)
}
func TestGrep_PathOutsideWorktree(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
h := handlers.NewGrepHandler()
sess := handlers.SessionContext{CwdPath: tmp}
outside := filepath.Join(filepath.VolumeName(tmp)+string(filepath.Separator), "etc")
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "pattern": "x", "path": outside}))
require.NotNil(t, res.Err)
assert.Equal(t, -32001, res.Err.Code)
}
func TestGrep_InvalidPattern(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
h := handlers.NewGrepHandler()
sess := handlers.SessionContext{CwdPath: tmp}
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "pattern": "[invalid("}))
require.NotNil(t, res.Err)
assert.Equal(t, -32602, res.Err.Code)
}
func TestGrep_EmptyPattern(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
h := handlers.NewGrepHandler()
sess := handlers.SessionContext{CwdPath: tmp}
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "pattern": ""}))
require.NotNil(t, res.Err)
assert.Equal(t, -32602, res.Err.Code)
}
func TestGrep_MaxMatchesCap(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
// 10 matching lines; cap at 3.
content := "m\nm\nm\nm\nm\nm\nm\nm\nm\nm\n"
require.NoError(t, os.WriteFile(filepath.Join(tmp, "a.txt"), []byte(content), 0o644))
h := &handlers.GrepHandler{MaxMatches: 3, MaxFiles: 100, MaxFileSize: 1 << 20}
sess := handlers.SessionContext{CwdPath: tmp}
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "pattern": "m"}))
require.Nil(t, res.Err)
var out struct {
Matches []any `json:"matches"`
Truncated bool `json:"truncated"`
}
require.NoError(t, json.Unmarshal(res.OK, &out))
require.Len(t, out.Matches, 3)
require.True(t, out.Truncated)
}
func TestGrep_SkipsGitDir(t *testing.T) {
t.Parallel()
tmp := t.TempDir()
require.NoError(t, os.MkdirAll(filepath.Join(tmp, ".git"), 0o755))
require.NoError(t, os.WriteFile(filepath.Join(tmp, ".git", "config"), []byte("secret pattern\n"), 0o644))
require.NoError(t, os.WriteFile(filepath.Join(tmp, "a.txt"), []byte("normal pattern\n"), 0o644))
h := handlers.NewGrepHandler()
sess := handlers.SessionContext{CwdPath: tmp}
res := h.Handle(context.Background(), sess, mustJSON(t, map[string]any{"sessionId": "x", "pattern": "pattern"}))
require.Nil(t, res.Err)
var out struct {
Matches []struct {
File string `json:"file"`
} `json:"matches"`
}
require.NoError(t, json.Unmarshal(res.OK, &out))
require.Len(t, out.Matches, 1)
assert.Equal(t, "a.txt", out.Matches[0].File)
}
+104 -15
View File
@@ -12,15 +12,17 @@ import (
)
// MCPBridge 适配 AgentKind/Session/Permission 三个服务到 mcp.ACPBridge。
// repo 仅用于只读的回合列表(ListSessionTurns)——scope 校验先经 sessions.Get 完成。
type MCPBridge struct {
kinds AgentKindService
sessions SessionService
perms *PermissionService
repo Repository
}
// NewMCPBridge 构造桥接适配器,由 app.go 装配 mcp.ServerDeps 时调用。
func NewMCPBridge(kinds AgentKindService, sessions SessionService, perms *PermissionService) *MCPBridge {
return &MCPBridge{kinds: kinds, sessions: sessions, perms: perms}
func NewMCPBridge(kinds AgentKindService, sessions SessionService, perms *PermissionService, repo Repository) *MCPBridge {
return &MCPBridge{kinds: kinds, sessions: sessions, perms: perms, repo: repo}
}
var _ mcp.ACPBridge = (*MCPBridge)(nil)
@@ -109,6 +111,86 @@ func (b *MCPBridge) ListPendingPermissions(ctx context.Context, userID uuid.UUID
return out, nil
}
// ListPendingApprovals 返回调用者跨所有会话的待审权限请求(HITL inbox,read-only)。
func (b *MCPBridge) ListPendingApprovals(ctx context.Context, userID uuid.UUID, isAdmin bool) ([]mcp.ACPPermission, error) {
ps, err := b.perms.ListPendingForUser(ctx, Caller{UserID: userID, IsAdmin: isAdmin})
if err != nil {
return nil, err
}
out := make([]mcp.ACPPermission, 0, len(ps))
for _, p := range ps {
out = append(out, mcp.ACPPermission{
ID: p.ID,
SessionID: p.SessionID,
AgentRequestID: p.AgentRequestID,
ToolName: p.ToolName,
Status: string(p.Status),
CreatedAt: p.CreatedAt,
})
}
return out, nil
}
// ListSessionTurns 返回会话的回合列表(只读)。先经 sessions.Get 做用户 + 项目
// scope 校验(与 GetSession 同语义),再用 repo 读 acp_turns。
func (b *MCPBridge) ListSessionTurns(ctx context.Context, userID uuid.UUID, isAdmin bool, sessionID uuid.UUID) ([]mcp.ACPTurn, error) {
if _, err := b.sessions.Get(ctx, Caller{UserID: userID, IsAdmin: isAdmin}, sessionID); err != nil {
return nil, err
}
ts, err := b.repo.ListTurnsBySession(ctx, sessionID)
if err != nil {
return nil, err
}
out := make([]mcp.ACPTurn, 0, len(ts))
for _, t := range ts {
out = append(out, mcp.ACPTurn{
TurnIndex: t.TurnIndex,
Status: string(t.Status),
StopReason: t.StopReason,
UpdateCount: t.UpdateCount,
StartedAt: t.StartedAt,
CompletedAt: t.CompletedAt,
})
}
return out, nil
}
// GetRunDashboard 返回 run-history 汇总(owner scope;admin 全量,read-only)。
func (b *MCPBridge) GetRunDashboard(ctx context.Context, userID uuid.UUID, isAdmin bool, in mcp.ACPRunDashboardInput) (*mcp.ACPRunDashboard, error) {
res, err := b.sessions.Dashboard(ctx, Caller{UserID: userID, IsAdmin: isAdmin}, DashboardInput{
ProjectID: in.ProjectID,
RequirementID: in.RequirementID,
From: in.From,
To: in.To,
})
if err != nil {
return nil, err
}
out := &mcp.ACPRunDashboard{
From: res.From,
To: res.To,
Total: res.Rollup.Total,
Succeeded: res.Rollup.Succeeded,
Crashed: res.Rollup.Crashed,
Active: res.Rollup.Active,
SuccessRate: res.Rollup.SuccessRate(),
TotalCostUSD: res.Rollup.TotalCostUSD,
TotalTokens: res.Rollup.TotalTokens,
AvgDurationSeconds: res.Rollup.AvgDurationSeconds,
ByDay: make([]mcp.ACPRunDayRow, 0, len(res.ByDay)),
}
for _, d := range res.ByDay {
out.ByDay = append(out.ByDay, mcp.ACPRunDayRow{
Day: d.Day,
Total: d.Total,
Succeeded: d.Succeeded,
Crashed: d.Crashed,
TotalCostUSD: d.TotalCostUSD,
})
}
return out, nil
}
func agentKindToMCP(k *AgentKind) mcp.ACPAgentKind {
return mcp.ACPAgentKind{
ID: k.ID,
@@ -122,18 +204,25 @@ func agentKindToMCP(k *AgentKind) mcp.ACPAgentKind {
func sessionToMCP(s *Session) mcp.ACPSession {
return mcp.ACPSession{
ID: s.ID,
WorkspaceID: s.WorkspaceID,
ProjectID: s.ProjectID,
AgentKindID: s.AgentKindID,
UserID: s.UserID,
IssueID: s.IssueID,
RequirementID: s.RequirementID,
Branch: s.Branch,
IsMainWorktree: s.IsMainWorktree,
Status: string(s.Status),
LastError: s.LastError,
StartedAt: s.StartedAt,
EndedAt: s.EndedAt,
ID: s.ID,
WorkspaceID: s.WorkspaceID,
ProjectID: s.ProjectID,
AgentKindID: s.AgentKindID,
UserID: s.UserID,
IssueID: s.IssueID,
RequirementID: s.RequirementID,
Branch: s.Branch,
IsMainWorktree: s.IsMainWorktree,
Status: string(s.Status),
LastError: s.LastError,
LastStopReason: s.LastStopReason,
StartedAt: s.StartedAt,
EndedAt: s.EndedAt,
PromptTokens: s.PromptTokens,
CompletionTokens: s.CompletionTokens,
ThinkingTokens: s.ThinkingTokens,
TotalCostUSD: s.TotalCostUSD,
BudgetMaxCostUSD: s.BudgetMaxCostUSD,
BudgetMaxTokens: s.BudgetMaxTokens,
}
}
+153
View File
@@ -0,0 +1,153 @@
package acp_test
import (
"context"
"sync"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/acp"
"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
)
// spyNotifier captures Dispatch calls for the HITL pending path.
type spyNotifier struct {
mu sync.Mutex
msgs []notify.Message
}
func (s *spyNotifier) Dispatch(_ context.Context, m notify.Message) error {
s.mu.Lock()
defer s.mu.Unlock()
s.msgs = append(s.msgs, m)
return nil
}
func (s *spyNotifier) count() int {
s.mu.Lock()
defer s.mu.Unlock()
return len(s.msgs)
}
func (s *spyNotifier) first() notify.Message {
s.mu.Lock()
defer s.mu.Unlock()
return s.msgs[0]
}
// spyMetrics captures permission decision outcomes.
type spyMetrics struct {
mu sync.Mutex
outcomes []string
}
func (s *spyMetrics) RecordPermissionDecision(outcome string) {
s.mu.Lock()
defer s.mu.Unlock()
s.outcomes = append(s.outcomes, outcome)
}
func (s *spyMetrics) snapshot() []string {
s.mu.Lock()
defer s.mu.Unlock()
out := make([]string, len(s.outcomes))
copy(out, s.outcomes)
return out
}
func contains(ss []string, want string) bool {
for _, s := range ss {
if s == want {
return true
}
}
return false
}
func TestPermission_Pending_DispatchesNotifyBeforeBlocking(t *testing.T) {
repo := newPermFakeRepo()
uid := uuid.New()
sid := uuid.New()
repo.sessions[sid] = &acp.Session{ID: sid, UserID: uid}
spy := &spyNotifier{}
met := &spyMetrics{}
svc := acp.NewPermissionService(repo, nil, time.Minute, nil)
svc.SetNotifier(spy)
svc.SetMetrics(met)
sess := handlers.SessionContext{SessionID: sid, UserID: uid} // no allowlist -> pending
resCh := make(chan string, 1)
go func() {
resCh <- svc.Decide(context.Background(), sess, "req-1",
[]byte(`{"name":"danger.tool"}`), optionsAllowReject())
}()
// Notification must fire while Decide is still blocked (before resolve).
require.Eventually(t, func() bool { return spy.count() == 1 }, time.Second, 5*time.Millisecond)
m := spy.first()
require.Equal(t, "acp.permission_pending", m.Topic)
require.Equal(t, uid, m.UserID)
require.Equal(t, notify.SeverityWarning, m.Severity)
require.Equal(t, "/approvals", m.Link)
require.Equal(t, sid.String(), m.Metadata["session_id"])
require.Equal(t, "danger.tool", m.Metadata["tool_name"])
// Resolve so Decide unblocks; exactly one dispatch (not on resolve).
var reqID uuid.UUID
require.Eventually(t, func() bool {
reqID = repo.onlyReqID()
return reqID != uuid.Nil
}, time.Second, 5*time.Millisecond)
require.NoError(t, svc.Resolve(context.Background(), acp.Caller{UserID: uid}, reqID, true, "allow_once"))
<-resCh
require.Equal(t, 1, spy.count(), "dispatch must fire once (pending), not on resolve")
// Metrics: approved recorded.
require.True(t, contains(met.snapshot(), "approved"))
}
func TestPermission_Metrics_AutoAndExpired(t *testing.T) {
repo := newPermFakeRepo()
uid := uuid.New()
sid := uuid.New()
repo.sessions[sid] = &acp.Session{ID: sid, UserID: uid}
met := &spyMetrics{}
svc := acp.NewPermissionService(repo, nil, 30*time.Millisecond, nil)
svc.SetMetrics(met)
// auto path
svc.Decide(context.Background(), handlers.SessionContext{
SessionID: sid, UserID: uid, ToolAllowlist: []string{"*"},
}, "req-auto", []byte(`{"name":"safe.tool"}`), optionsAllowReject())
require.True(t, contains(met.snapshot(), "auto"))
// expired path (no allowlist, times out)
chosen := svc.Decide(context.Background(), handlers.SessionContext{
SessionID: sid, UserID: uid,
}, "req-exp", []byte(`{"name":"danger.tool"}`), optionsAllowReject())
require.Equal(t, "", chosen)
require.True(t, contains(met.snapshot(), "expired"))
}
func TestPermission_ListPendingForUser(t *testing.T) {
repo := newPermFakeRepo()
uid := uuid.New()
other := uuid.New()
sidA := uuid.New()
sidB := uuid.New()
repo.sessions[sidA] = &acp.Session{ID: sidA, UserID: uid}
repo.sessions[sidB] = &acp.Session{ID: sidB, UserID: other}
// Two pending for uid (across sessions), one for other.
repo.reqs[uuid.New()] = &acp.PermissionRequest{ID: uuid.New(), SessionID: sidA, Status: acp.PermissionPending}
repo.reqs[uuid.New()] = &acp.PermissionRequest{ID: uuid.New(), SessionID: sidA, Status: acp.PermissionPending}
repo.reqs[uuid.New()] = &acp.PermissionRequest{ID: uuid.New(), SessionID: sidB, Status: acp.PermissionPending}
svc := acp.NewPermissionService(repo, nil, time.Minute, nil)
got, err := svc.ListPendingForUser(context.Background(), acp.Caller{UserID: uid})
require.NoError(t, err)
require.Len(t, got, 2)
}
+114 -3
View File
@@ -24,17 +24,42 @@ import (
"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
"github.com/yan1h/agent-coding-workflow/internal/audit"
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
)
// DefaultPermissionTimeout 是未配置时的人工决策超时;到期默认拒绝。
const DefaultPermissionTimeout = 5 * time.Minute
// permissionTopic 是 HITL 待审权限通知的 topic(inbox + 跨通道路由用)。
const permissionTopic = "acp.permission_pending"
// PermissionNotifier 是 PermissionService 投递"待人工审批"通知所需的窄接口。
// notify.Dispatcher 满足它;用窄接口便于测试断言 Dispatch 调用。
type PermissionNotifier interface {
Dispatch(ctx context.Context, msg notify.Message) error
}
// PermissionMetrics 是 PermissionService 记录决策计数所需的窄接口。
// metrics.Metrics 满足它;nil 时不记录(部分测试 / metrics 关闭)。
type PermissionMetrics interface {
RecordPermissionDecision(outcome string)
}
// RelayLocator 让 PermissionService 按 session id 取到 *Relay 以推送 WS 事件。
// Supervisor 满足该接口(GetRelay)。用窄接口避免 permSvc 依赖整个 Supervisor。
type RelayLocator interface {
GetRelay(sid uuid.UUID) *Relay
}
// ProjectMaintainer 让 PermissionService 判定某用户是否对会话所属 project 有写权限
// (project maintainer:owner / global-admin / project admin|member 角色)。允许任意
// project maintainer 处理待审权限请求,而不仅是会话 owner(autonomy roadmap §11)。
// app 层用 projectAccessAdapter.ResolveByID 实现;nil 时退化为仅 owner/admin。
type ProjectMaintainer interface {
// CanWriteProject 报告 (callerID, isAdmin) 是否对 projectID 有写权限。
CanWriteProject(ctx context.Context, callerID uuid.UUID, isAdmin bool, projectID uuid.UUID) (bool, error)
}
// decision 是投递给阻塞中 Decide 的人工/系统决策结果。
type decision struct {
optionID string
@@ -58,8 +83,27 @@ type PermissionService struct {
locMu sync.RWMutex
loc RelayLocator
// HITL:待审权限通知 + 决策计数。装配期经 SetNotifier/SetMetrics 注入,
// 二者均可为 nil(部分测试 / metrics 关闭)。
notify PermissionNotifier
metrics PermissionMetrics
// maintainer 允许任意 project maintainer(非仅会话 owner)处理待审权限请求。
// 装配期经 SetProjectMaintainer 注入;nil 时退化为仅 owner/global-admin。
maintainer ProjectMaintainer
}
// SetNotifier 注入待审权限通知投递器(装配期,best-effort)。
func (s *PermissionService) SetNotifier(n PermissionNotifier) { s.notify = n }
// SetMetrics 注入权限决策计数器(装配期)。
func (s *PermissionService) SetMetrics(m PermissionMetrics) { s.metrics = m }
// SetProjectMaintainer 注入 project maintainer 判定器(装配期)。注入后,会话
// 所属 project 的任意 maintainer 均可处理该会话的待审权限请求。
func (s *PermissionService) SetProjectMaintainer(m ProjectMaintainer) { s.maintainer = m }
// NewPermissionService 构造 PermissionService。timeout<=0 时用 DefaultPermissionTimeout。
func NewPermissionService(repo Repository, rec audit.Recorder, timeout time.Duration, log *slog.Logger) *PermissionService {
if log == nil {
@@ -109,6 +153,7 @@ func (s *PermissionService) Decide(ctx context.Context, sess handlers.SessionCon
s.log.Error("acp.permission.insert_auto", "session_id", sess.SessionID, "err", err.Error())
}
s.recordDecision(ctx, sess.UserID, sess.SessionID, toolName, "auto", nil)
s.recordMetric("auto")
return chosen
}
@@ -140,6 +185,10 @@ func (s *PermissionService) Decide(ctx context.Context, sess handlers.SessionCon
Payload: mustJSON(permissionRequestPayload(row)),
})
// HITL:在阻塞前 best-effort 推送 inbox 通知,使关闭页面的用户也能被提醒去
// /approvals 审批(fire-and-forget,错误不影响 deny-on-timeout 安全默认)。
s.notifyPending(ctx, sess, row, toolName)
select {
case d := <-w.ch:
return d.optionID
@@ -147,6 +196,7 @@ func (s *PermissionService) Decide(ctx context.Context, sess handlers.SessionCon
// 超时 → CAS 标记 expired(系统决策,decided_by=NULL)。
if _, ok, _ := s.repo.DecidePermissionRequest(context.WithoutCancel(ctx), reqID, PermissionExpired, nil, nil); ok {
s.fanout(sess.SessionID, resolvedEnvelope(reqID, PermissionExpired, ""))
s.recordMetric("expired")
}
return ""
case <-ctx.Done():
@@ -155,7 +205,8 @@ func (s *PermissionService) Decide(ctx context.Context, sess handlers.SessionCon
}
}
// Resolve 落地一次人工决策(approve/deny)。鉴权:session owneradmin
// Resolve 落地一次人工决策(approve/deny)。鉴权:session owner、global-admin,或
// 会话所属 project 的任意 maintainer(autonomy roadmap §11)。
func (s *PermissionService) Resolve(ctx context.Context, c Caller, reqID uuid.UUID, approve bool, optionID string) error {
req, err := s.repo.GetPermissionRequestByID(ctx, reqID)
if err != nil {
@@ -165,7 +216,7 @@ func (s *PermissionService) Resolve(ctx context.Context, c Caller, reqID uuid.UU
if err != nil {
return err
}
if !c.IsAdmin && sessRow.UserID != c.UserID {
if !s.canManageSession(ctx, c, sessRow) {
return errs.New(errs.CodeAcpSessionNotOwned, "无权处理该会话的权限请求")
}
@@ -200,21 +251,50 @@ func (s *PermissionService) Resolve(ctx context.Context, c Caller, reqID uuid.UU
s.wake(reqID, decision{optionID: chosen})
s.fanout(req.SessionID, resolvedEnvelope(reqID, status, chosen))
s.recordDecision(ctx, c.UserID, req.SessionID, req.ToolName, string(status), &c.UserID)
if approve {
s.recordMetric("approved")
} else {
s.recordMetric("denied")
}
return nil
}
// ListPendingForUser 返回该用户跨所有会话的待审权限请求(HITL 跨会话审批 inbox)。
// 包装 repo.ListPendingPermissionRequestsByUser(已 JOIN acp_sessions.user_id)。
// admin 仅看自己拥有会话的待审项(与现有 inbox 语义一致;全局视图走 dashboard)。
func (s *PermissionService) ListPendingForUser(ctx context.Context, c Caller) ([]*PermissionRequest, error) {
return s.repo.ListPendingPermissionRequestsByUser(ctx, c.UserID)
}
// ListPendingBySession 返回某会话的待审请求。鉴权同 Resolve。
func (s *PermissionService) ListPendingBySession(ctx context.Context, c Caller, sessionID uuid.UUID) ([]*PermissionRequest, error) {
sessRow, err := s.repo.GetSessionByID(ctx, sessionID)
if err != nil {
return nil, err
}
if !c.IsAdmin && sessRow.UserID != c.UserID {
if !s.canManageSession(ctx, c, sessRow) {
return nil, errs.New(errs.CodeAcpSessionNotOwned, "无权查看该会话")
}
return s.repo.ListPendingPermissionRequestsBySession(ctx, sessionID)
}
// canManageSession 报告 caller 是否可处理/查看该会话的权限请求:session owner、
// global-admin,或会话所属 project 的任意 maintainer(注入 ProjectMaintainer 时)。
// maintainer 查询失败按"无权"处理(fail-closed),但 owner/admin 短路不受影响。
func (s *PermissionService) canManageSession(ctx context.Context, c Caller, sess *Session) bool {
if c.IsAdmin || sess.UserID == c.UserID {
return true
}
if s.maintainer == nil {
return false
}
ok, err := s.maintainer.CanWriteProject(ctx, c.UserID, c.IsAdmin, sess.ProjectID)
if err != nil {
return false
}
return ok
}
// CancelSession 在会话终止时把其全部 pending 请求标 expired,并唤醒阻塞的 Decide
// goroutine 使其立即返回(agent 收默认拒绝)。supervisor.onExit 调用。
func (s *PermissionService) CancelSession(ctx context.Context, sessionID uuid.UUID) {
@@ -291,6 +371,37 @@ func (s *PermissionService) recordDecision(ctx context.Context, actor, sessionID
})
}
// recordMetric 记录一次权限决策计数(auto/approved/denied/expired)。metrics 为 nil 时 no-op。
func (s *PermissionService) recordMetric(outcome string) {
if s.metrics == nil {
return
}
s.metrics.RecordPermissionDecision(outcome)
}
// notifyPending best-effort 向会话所有者推送"待人工审批"通知。错误仅记日志,
// 绝不影响 Decide 的阻塞 / deny-on-timeout 安全默认。
func (s *PermissionService) notifyPending(ctx context.Context, sess handlers.SessionContext, row *PermissionRequest, toolName string) {
if s.notify == nil {
return
}
if err := s.notify.Dispatch(context.WithoutCancel(ctx), notify.Message{
UserID: sess.UserID,
Topic: permissionTopic,
Severity: notify.SeverityWarning,
Title: "Approval needed",
Body: "An agent is requesting permission to run a tool and needs your approval.",
Link: "/approvals",
Metadata: map[string]any{
"request_id": row.ID.String(),
"session_id": sess.SessionID.String(),
"tool_name": toolName,
},
}); err != nil {
s.log.Warn("acp.permission.notify_pending_failed", "session_id", sess.SessionID, "err", err.Error())
}
}
func permissionRequestPayload(p *PermissionRequest) map[string]any {
return map[string]any{
"request_id": p.ID.String(),
+90
View File
@@ -61,6 +61,13 @@ func (r *permFakeRepo) GetSessionByID(_ context.Context, id uuid.UUID) (*acp.Ses
out := *s
return &out, nil
}
func (r *permFakeRepo) GetSessionByStepID(context.Context, uuid.UUID) (*acp.Session, error) {
return nil, errs.New(errs.CodeAcpSessionNotFound, "not found")
}
func (r *permFakeRepo) GetCrashedSessionForResume(context.Context, uuid.UUID) (*acp.Session, error) {
return nil, errs.New(errs.CodeAcpSessionNotFound, "not found")
}
func (r *permFakeRepo) ResetSessionForResume(context.Context, uuid.UUID) error { return nil }
func (r *permFakeRepo) DecidePermissionRequest(_ context.Context, id uuid.UUID, status acp.PermissionStatus, chosen *string, decidedBy *uuid.UUID) (*acp.PermissionRequest, bool, error) {
r.mu.Lock()
@@ -89,6 +96,24 @@ func (r *permFakeRepo) ExpirePendingPermissionRequestsBySession(_ context.Contex
return n, nil
}
func (r *permFakeRepo) ListPendingPermissionRequestsByUser(_ context.Context, userID uuid.UUID) ([]*acp.PermissionRequest, error) {
r.mu.Lock()
defer r.mu.Unlock()
var out []*acp.PermissionRequest
for _, p := range r.reqs {
if p.Status != acp.PermissionPending {
continue
}
sess, ok := r.sessions[p.SessionID]
if !ok || sess.UserID != userID {
continue
}
cp := *p
out = append(out, &cp)
}
return out, nil
}
func (r *permFakeRepo) statusOf(id uuid.UUID) acp.PermissionStatus {
r.mu.Lock()
defer r.mu.Unlock()
@@ -310,3 +335,68 @@ func TestPermission_CancelSession_UnblocksDecide(t *testing.T) {
}
require.Equal(t, acp.PermissionExpired, repo.statusOf(repo.onlyReqID()))
}
// fakeMaintainer 满足 acp.ProjectMaintainer:仅当 projectID 命中 allow 集合时放行。
type fakeMaintainer struct {
allow map[uuid.UUID]bool
}
func (f fakeMaintainer) CanWriteProject(_ context.Context, callerID uuid.UUID, isAdmin bool, projectID uuid.UUID) (bool, error) {
if isAdmin {
return true, nil
}
return f.allow[projectID], nil
}
// 注入 ProjectMaintainer 后:会话所属 project 的任意 maintainer(非会话 owner)
// 也可处理待审权限请求(autonomy roadmap §11)。
func TestPermission_Resolve_ProjectMaintainer_Allowed(t *testing.T) {
repo := newPermFakeRepo()
owner := uuid.New()
maintainer := uuid.New()
pid := uuid.New()
sid := uuid.New()
repo.sessions[sid] = &acp.Session{ID: sid, UserID: owner, ProjectID: pid}
svc := acp.NewPermissionService(repo, nil, time.Minute, nil)
svc.SetProjectMaintainer(fakeMaintainer{allow: map[uuid.UUID]bool{pid: true}})
sess := handlers.SessionContext{SessionID: sid, UserID: owner}
go svc.Decide(context.Background(), sess, "req-1",
[]byte(`{"name":"danger.tool"}`), optionsAllowReject())
var reqID uuid.UUID
require.Eventually(t, func() bool {
reqID = repo.onlyReqID()
return reqID != uuid.Nil && repo.statusOf(reqID) == acp.PermissionPending
}, time.Second, 5*time.Millisecond)
// 非 owner 但是 project maintainer → 可处理。
require.NoError(t, svc.Resolve(context.Background(), acp.Caller{UserID: maintainer}, reqID, true, "allow_once"))
require.Equal(t, acp.PermissionApproved, repo.statusOf(reqID))
}
// 非 maintainer 的外部用户仍被拒(fail-closed);项目不在 allow 集合时也拒。
func TestPermission_Resolve_NonMaintainer_Forbidden(t *testing.T) {
repo := newPermFakeRepo()
owner := uuid.New()
pid := uuid.New()
otherPid := uuid.New()
sid := uuid.New()
repo.sessions[sid] = &acp.Session{ID: sid, UserID: owner, ProjectID: pid}
svc := acp.NewPermissionService(repo, nil, time.Minute, nil)
// maintainer 仅对 otherPid 有权,而会话属于 pid。
svc.SetProjectMaintainer(fakeMaintainer{allow: map[uuid.UUID]bool{otherPid: true}})
sess := handlers.SessionContext{SessionID: sid, UserID: owner}
go svc.Decide(context.Background(), sess, "req-1",
[]byte(`{"name":"danger.tool"}`), optionsAllowReject())
var reqID uuid.UUID
require.Eventually(t, func() bool {
reqID = repo.onlyReqID()
return reqID != uuid.Nil && repo.statusOf(reqID) == acp.PermissionPending
}, time.Second, 5*time.Millisecond)
err := svc.Resolve(context.Background(), acp.Caller{UserID: uuid.New()}, reqID, true, "allow_once")
ae, ok := errs.As(err)
require.True(t, ok)
require.Equal(t, errs.CodeAcpSessionNotOwned, ae.Code)
}
@@ -1,5 +1,5 @@
-- name: ListAgentKindConfigFiles :many
SELECT id, agent_kind_id, rel_path, encrypted_content, updated_by, created_at, updated_at
SELECT id, agent_kind_id, rel_path, encrypted_content, updated_by, created_at, updated_at, key_version
FROM acp_agent_kind_config_files
WHERE agent_kind_id = $1
ORDER BY rel_path ASC;
@@ -12,7 +12,7 @@ ON CONFLICT (agent_kind_id, rel_path) DO UPDATE
SET encrypted_content = EXCLUDED.encrypted_content,
updated_by = EXCLUDED.updated_by,
updated_at = now()
RETURNING id, agent_kind_id, rel_path, encrypted_content, updated_by, created_at, updated_at;
RETURNING id, agent_kind_id, rel_path, encrypted_content, updated_by, created_at, updated_at, key_version;
-- name: DeleteAgentKindConfigFile :execrows
DELETE FROM acp_agent_kind_config_files
+29 -18
View File
@@ -1,50 +1,61 @@
-- name: CreateAgentKind :one
INSERT INTO acp_agent_kinds (
id, name, display_name, description, binary_path, args, encrypted_env, enabled, created_by, tool_allowlist, client_type, encrypted_mcp_servers
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
id, name, display_name, description, binary_path, args, encrypted_env, enabled, created_by, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
RETURNING id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers;
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version;
-- name: GetAgentKindByID :one
SELECT id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
FROM acp_agent_kinds
WHERE id = $1;
-- name: GetAgentKindByName :one
SELECT id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
FROM acp_agent_kinds
WHERE name = $1;
-- name: ListAgentKinds :many
SELECT id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
FROM acp_agent_kinds
ORDER BY created_at DESC;
-- name: ListEnabledAgentKinds :many
SELECT id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
FROM acp_agent_kinds
WHERE enabled = TRUE
ORDER BY name ASC;
-- name: UpdateAgentKind :one
UPDATE acp_agent_kinds
SET display_name = $2,
description = $3,
binary_path = $4,
args = $5,
encrypted_env = $6,
enabled = $7,
tool_allowlist = $8,
client_type = $9,
encrypted_mcp_servers = $10,
updated_at = now()
SET display_name = $2,
description = $3,
binary_path = $4,
args = $5,
encrypted_env = $6,
enabled = $7,
tool_allowlist = $8,
client_type = $9,
encrypted_mcp_servers = $10,
model_id = $11,
max_cost_usd = $12,
max_tokens = $13,
max_wall_clock_seconds = $14,
updated_at = now()
WHERE id = $1
RETURNING id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers;
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version;
-- name: DeleteAgentKind :exec
DELETE FROM acp_agent_kinds WHERE id = $1;
+59
View File
@@ -0,0 +1,59 @@
-- dashboard.sql: run-history dashboard 读模型(只读聚合)。
-- - SessionTimeline: 单会话的 RPC 事件按 (direction, rpc_kind, method) 分桶,
-- 给出条数 + 首/末时间戳,用于会话时间线概览。
-- - SessionRollup: 在 owner/admin scope + 可选 project/requirement 过滤下,
-- 统计会话总数、成功率(exited 视为成功)、总成本、平均时长。
-- - SessionsByDay: 按天的 总数 / 成功 / 失败 / 成本 时间序列。
-- scope 约定(与 ListSessionsByUser/ListAllSessions 一致):
-- $1 user_id 为 NULL 时不按用户过滤(admin 全量),非 NULL 时仅该用户。
-- name: SessionTimeline :many
-- 单会话事件时间线:按 method/kind 分桶聚合(条数 + 首末时间)。
SELECT direction,
rpc_kind,
COALESCE(method, '') AS method,
COUNT(*)::bigint AS event_count,
MIN(created_at)::timestamptz AS first_at,
MAX(created_at)::timestamptz AS last_at
FROM acp_events
WHERE session_id = $1
GROUP BY direction, rpc_kind, COALESCE(method, '')
ORDER BY first_at ASC;
-- name: SessionRollup :one
-- owner/admin scope + 可选 project/requirement + 时间窗内的会话汇总。
-- success = status='exited';avg_duration_seconds 仅统计已结束(ended_at 非空)会话。
SELECT
COUNT(*)::bigint AS total,
COUNT(*) FILTER (WHERE status = 'exited')::bigint AS succeeded,
COUNT(*) FILTER (WHERE status = 'crashed')::bigint AS crashed,
COUNT(*) FILTER (WHERE status IN ('starting','running'))::bigint AS active,
COALESCE(SUM(cost_usd), 0)::numeric AS total_cost_usd,
COALESCE(SUM(tokens_total), 0)::bigint AS total_tokens,
COALESCE(
AVG(EXTRACT(EPOCH FROM (ended_at - started_at)))
FILTER (WHERE ended_at IS NOT NULL),
0)::float8 AS avg_duration_seconds
FROM acp_sessions
WHERE ($1::uuid IS NULL OR user_id = $1)
AND ($2::uuid IS NULL OR project_id = $2)
AND ($3::uuid IS NULL OR requirement_id = $3)
AND started_at >= $4
AND started_at < $5;
-- name: SessionsByDay :many
-- 按天的会话时间序列(总数 / 成功 / 失败 / 成本),同 SessionRollup 的 scope。
SELECT
date_trunc('day', started_at)::timestamptz AS day,
COUNT(*)::bigint AS total,
COUNT(*) FILTER (WHERE status = 'exited')::bigint AS succeeded,
COUNT(*) FILTER (WHERE status = 'crashed')::bigint AS crashed,
COALESCE(SUM(cost_usd), 0)::numeric AS total_cost_usd
FROM acp_sessions
WHERE ($1::uuid IS NULL OR user_id = $1)
AND ($2::uuid IS NULL OR project_id = $2)
AND ($3::uuid IS NULL OR requirement_id = $3)
AND started_at >= $4
AND started_at < $5
GROUP BY day
ORDER BY day ASC;
+9
View File
@@ -0,0 +1,9 @@
-- name: InsertSessionUsage :one
-- 写一条 per-turn 用量账目。source_event_id 非空时按唯一索引去重(同一事件
-- 重复 Observe 不会重复入账);返回受影响行数判定是否实际插入。
INSERT INTO acp_session_usage (
session_id, user_id, project_id, agent_kind_id, model_id,
prompt_tokens, completion_tokens, thinking_tokens, cost_usd, source_event_id
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
ON CONFLICT (source_event_id) WHERE source_event_id IS NOT NULL DO NOTHING
RETURNING id;
+170 -6
View File
@@ -1,26 +1,75 @@
-- name: InsertSession :one
INSERT INTO acp_sessions (
id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, branch, cwd_path, is_main_worktree, status
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
issue_id, requirement_id, branch, cwd_path, is_main_worktree, status,
orchestrator_step_id,
budget_max_cost_usd, budget_max_tokens, budget_max_wall_clock_seconds
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
RETURNING id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at;
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total;
-- name: GetSessionByID :one
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE id = $1;
-- name: GetSessionByStepID :one
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE orchestrator_step_id = $1
ORDER BY started_at DESC
LIMIT 1;
-- name: GetCrashedSessionForResume :one
-- 返回某 step 最近一次处于可恢复终态(crashed/exited)的 session。
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE orchestrator_step_id = $1
AND status IN ('crashed','exited')
ORDER BY started_at DESC
LIMIT 1;
-- name: ListSessionsByUser :many
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE user_id = $1
AND ($2::uuid IS NULL OR requirement_id = $2)
@@ -30,7 +79,12 @@ ORDER BY started_at DESC;
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE ($1::uuid IS NULL OR requirement_id = $1)
ORDER BY started_at DESC;
@@ -40,6 +94,17 @@ UPDATE acp_sessions
SET status = 'running', agent_session_id = $2, pid = $3
WHERE id = $1;
-- name: UpdateSessionLastStopReason :exec
UPDATE acp_sessions
SET last_stop_reason = $2
WHERE id = $1;
-- name: UpdateSessionSandboxMode :exec
-- 记录本 session 运行所用的沙箱模式(审计/取证)。
UPDATE acp_sessions
SET sandbox_mode = $2
WHERE id = $1;
-- name: UpdateSessionFinished :exec
UPDATE acp_sessions
SET status = $2, exit_code = $3, last_error = $4, ended_at = now()
@@ -50,6 +115,20 @@ UPDATE acp_sessions
SET status = 'crashed', last_error = $2, ended_at = now()
WHERE id = $1 AND status IN ('starting', 'running');
-- name: ResetSessionForResume :exec
-- 把崩溃/退出的 session 复用为新一轮:回到 starting、清空上次运行时字段。
-- 编排器 resume 在重新 Spawn 前调用,使同一行 session 在同 cwd/worktree 上复用。
UPDATE acp_sessions
SET status = 'starting',
agent_session_id = NULL,
pid = NULL,
exit_code = NULL,
last_error = NULL,
ended_at = NULL,
terminated_reason = NULL,
last_activity_at = now()
WHERE id = $1;
-- name: CountActiveSessions :one
SELECT COUNT(*) FROM acp_sessions WHERE status IN ('starting', 'running');
@@ -72,3 +151,88 @@ WHERE id = $1 AND active_main_session_id = $2;
-- name: AcquireMainWorktree :execrows
UPDATE workspaces SET active_main_session_id = $1
WHERE id = $2 AND active_main_session_id IS NULL;
-- name: AddSessionUsageTotals :one
-- 单 round-trip 累加 session 运行时 token/cost 总量并刷新 last_activity_at,
-- 返回累加后的权威总量供预算判定。
-- 同时把 dashboard rollup 列(cost_usd / tokens_total)与运行时累加器对齐,
-- 使 run-history dashboard 在会话进行中及退出后都能读到 agent 上报的成本。
UPDATE acp_sessions
SET prompt_tokens = prompt_tokens + $2,
completion_tokens = completion_tokens + $3,
thinking_tokens = thinking_tokens + $4,
total_cost_usd = total_cost_usd + $5,
cost_usd = cost_usd + $5,
tokens_total = tokens_total + $2 + $3 + $4,
last_activity_at = now()
WHERE id = $1
RETURNING prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd;
-- name: SumProjectCostUSD :one
-- 某 project 自 since 起的 ACP 累计花费(per-project 软预算)。
SELECT COALESCE(SUM(cost_usd), 0)::numeric AS cost_usd
FROM acp_session_usage
WHERE project_id = $1 AND created_at >= $2;
-- name: MarkSessionTerminatedReason :exec
-- 写入终止原因;已有非空 terminated_reason 时不覆盖(reaper / budget 互斥保护)。
UPDATE acp_sessions
SET terminated_reason = $2
WHERE id = $1 AND terminated_reason IS NULL;
-- name: ListSessionsForReaper :many
-- 返回应被回收的活跃 session:
-- 超过自身 wall-clock 上限(budget_max_wall_clock_seconds 非空且 started_at 过早),或
-- 超过全局 wall-clock 上限($1::timestamptz = now - 全局上限),或
-- 空闲超时(last_activity_at < $2::timestamptz)。
SELECT id, user_id, started_at, last_activity_at, budget_max_wall_clock_seconds
FROM acp_sessions
WHERE status IN ('starting','running')
AND (
(budget_max_wall_clock_seconds IS NOT NULL
AND started_at < now() - make_interval(secs => budget_max_wall_clock_seconds))
OR ($1::timestamptz IS NOT NULL AND started_at < $1::timestamptz)
OR last_activity_at < $2::timestamptz
);
-- name: SummarizeAcpUsageByUser :many
SELECT user_id,
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(thinking_tokens)::bigint AS thinking_tokens,
SUM(cost_usd)::numeric AS cost_usd
FROM acp_session_usage
WHERE created_at >= $1 AND created_at < $2
GROUP BY user_id
ORDER BY cost_usd DESC;
-- name: SummarizeAcpUsageByModel :many
SELECT model_id,
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(thinking_tokens)::bigint AS thinking_tokens,
SUM(cost_usd)::numeric AS cost_usd
FROM acp_session_usage
WHERE created_at >= $1 AND created_at < $2
GROUP BY model_id
ORDER BY cost_usd DESC;
-- name: SummarizeAcpUsageByDay :many
SELECT date_trunc('day', created_at)::timestamptz AS day,
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(thinking_tokens)::bigint AS thinking_tokens,
SUM(cost_usd)::numeric AS cost_usd
FROM acp_session_usage
WHERE created_at >= $1 AND created_at < $2
GROUP BY day
ORDER BY day;
-- name: ListSessionUsageBySession :many
SELECT id, session_id, user_id, project_id, agent_kind_id, model_id,
prompt_tokens, completion_tokens, thinking_tokens, cost_usd,
source_event_id, created_at
FROM acp_session_usage
WHERE session_id = $1
ORDER BY created_at DESC
LIMIT $2;
+39
View File
@@ -0,0 +1,39 @@
-- name: InsertTurn :one
INSERT INTO acp_turns (
session_id, turn_index, prompt_request_id, status
) VALUES ($1, $2, $3, $4)
RETURNING id, session_id, turn_index, prompt_request_id, status,
stop_reason, update_count, started_at, completed_at;
-- name: MarkTurnCompleted :exec
UPDATE acp_turns
SET status = 'completed', stop_reason = $3, update_count = $4, completed_at = now()
WHERE session_id = $1 AND turn_index = $2 AND status = 'in_progress';
-- name: MarkTurnAborted :exec
UPDATE acp_turns
SET status = 'aborted', completed_at = now()
WHERE session_id = $1 AND turn_index = $2 AND status = 'in_progress';
-- name: IncrementTurnUpdateCount :exec
UPDATE acp_turns
SET update_count = update_count + 1
WHERE session_id = $1 AND turn_index = $2 AND status = 'in_progress';
-- name: GetLatestTurnBySession :one
SELECT id, session_id, turn_index, prompt_request_id, status,
stop_reason, update_count, started_at, completed_at
FROM acp_turns
WHERE session_id = $1
ORDER BY turn_index DESC
LIMIT 1;
-- name: ListTurnsBySession :many
SELECT id, session_id, turn_index, prompt_request_id, status,
stop_reason, update_count, started_at, completed_at
FROM acp_turns
WHERE session_id = $1
ORDER BY turn_index ASC;
-- name: PurgeTurnsBefore :execrows
DELETE FROM acp_turns WHERE started_at < $1;
+46 -2
View File
@@ -71,6 +71,14 @@ type Relay struct {
log *slog.Logger
cfg RelayConfig
// 回合完成检测:被动解析 session/update + session/prompt 响应的 stopReason。
// 由 supervisor.Spawn 通过 SetTracker 注入;可为 nil(部分测试)。
tracker *TurnTracker
// 成本核算:被动解析 agent→client 消息中的 token/cost 用量并累加 + 预算判定。
// 由 supervisor.Spawn 通过 SetUsageSink 注入;可为 nil(部分测试 / 未配置)。
usageSink usageSink
// 终止控制
closed atomic.Bool
done chan struct{}
@@ -106,6 +114,17 @@ func NewRelay(sessionID uuid.UUID, dec *Decoder, enc *Encoder,
}
}
// SetTracker 注入回合追踪器。由 supervisor.Spawn 在 NewRelay 之后调用。
func (r *Relay) SetTracker(t *TurnTracker) { r.tracker = t }
// SetUsageSink 注入成本核算 sink。由 supervisor.Spawn 在 NewRelay 之后调用;
// 可为 nil(部分测试 / 未配置 model 价格时由 supervisor 自行决定是否注入)。
func (r *Relay) SetUsageSink(s usageSink) { r.usageSink = s }
// Tracker 返回回合追踪器(可能为 nil)。session_service / handler 在发送
// session/prompt 前用它登记回合 id。
func (r *Relay) Tracker() *TurnTracker { return r.tracker }
// Subscribe 注册一个 WS tab 订阅。
func (r *Relay) Subscribe(userID uuid.UUID) *Subscriber {
sub := &Subscriber{
@@ -228,6 +247,12 @@ func (r *Relay) reader(ctx context.Context, sess handlers.SessionContext) {
envelope := r.persistAndEnvelope(ctx, msg, "out")
if envelope != nil {
r.fanout(envelope)
// 成本核算:在事件落库后用持久化的 event id 喂给 usageSink。usageSink
// 内部仅在真实可计费回合时累加,并在预算突破时以 detached goroutine 触发
// kill(绝不阻塞 reader)。
if r.usageSink != nil {
r.usageSink.Observe(ctx, msg, envelope.ID)
}
}
switch {
@@ -238,10 +263,23 @@ func (r *Relay) reader(ctx context.Context, sess handlers.SessionContext) {
} else {
r.log.Warn("acp.relay.orphan_response", "session_id", r.sessionID, "id", id)
}
} else if sid, ok := msg.IDString(); ok && r.tracker.IsTrackedPrompt(sid) {
// 字符串-id 的 session/prompt 响应:回合完成检测。
// 仅当 id 匹配当前开放回合时拦截;其它字符串 id 仍走既有 fanout(上方已 fanout)。
if msg.Error != nil {
// JSON-RPC error 响应:本回合失败,标 aborted 并发 session_idle(reason=error),
// 而非误记为 completed(stopReason 空)。agent 仍存活,可接受下一条 prompt。
r.tracker.Abort(ctx, "error")
} else {
r.tracker.OnPromptResponse(ctx, sid, ParseStopReason(msg.Result))
}
}
// String IDs go to clients via fanout (handled above); no server action.
// 其它 String IDs go to clients via fanout (handled above); no server action.
case msg.IsNotification():
// Already persisted + fanout'd. No reply.
// Already persisted + fanout'd. No reply. 回合检测:累加 update 计数。
if msg.Method == "session/update" {
r.tracker.OnUpdate(ctx)
}
case msg.IsRequest():
go r.handleAgentRequest(ctx, sess, msg)
}
@@ -272,6 +310,12 @@ func (r *Relay) handleAgentRequest(ctx context.Context, sess handlers.SessionCon
case "fs/write_text_file":
out := r.fsHandler.Write.Handle(ctx, sess, req.Params)
resp = r.fsResultToResponse(req.ID, out)
case "fs/list":
out := r.fsHandler.List.Handle(ctx, sess, req.Params)
resp = r.fsResultToResponse(req.ID, out)
case "fs/grep":
out := r.fsHandler.Grep.Handle(ctx, sess, req.Params)
resp = r.fsResultToResponse(req.ID, out)
case "session/request_permission":
out := r.permHandler.Handle(ctx, sess, string(req.ID), req.Params)
resp = r.fsResultToResponse(req.ID, out)
+107
View File
@@ -0,0 +1,107 @@
// relay_discovery_test.go: relay dispatch tests for the fs/list + fs/grep agent
// discovery handlers (mirrors TestRelay_AgentFsRequest_Echoed).
package acp_test
import (
"encoding/json"
"errors"
"io"
"os"
"path/filepath"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/acp"
)
func relayRoundTrip(t *testing.T, rig *relayTestRig, id int64, method string, params map[string]any) *acp.Message {
t.Helper()
agentDec := acp.NewDecoder(rig.agentReader, 0)
agentEnc := acp.NewEncoder(rig.agentWriter)
req, err := acp.NewRequestInt(id, method, params)
require.NoError(t, err)
require.NoError(t, agentEnc.Encode(req))
respCh := make(chan *acp.Message, 1)
errCh := make(chan error, 1)
go func() {
m, derr := agentDec.DecodeMessage()
if derr != nil {
errCh <- derr
return
}
respCh <- m
}()
select {
case resp := <-respCh:
return resp
case err := <-errCh:
if errors.Is(err, io.EOF) {
t.Fatal("agent saw EOF before relay responded")
}
t.Fatalf("agent decoder error: %v", err)
case <-time.After(3 * time.Second):
t.Fatalf("relay did not respond to %s within 3s", method)
}
return nil
}
func TestRelay_AgentFsList_Dispatched(t *testing.T) {
t.Parallel()
rig := newRelayTestRig(t)
cancel, doneCh := rig.runRelay(t)
t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })
require.NoError(t, os.WriteFile(filepath.Join(rig.sess.CwdPath, "f.txt"), []byte("hi"), 0o644))
resp := relayRoundTrip(t, rig, 201, "fs/list", map[string]any{"path": ".", "sessionId": "a"})
require.True(t, resp.IsResponse())
require.Nil(t, resp.Error)
var out struct {
Entries []struct {
Name string `json:"name"`
} `json:"entries"`
}
require.NoError(t, json.Unmarshal(resp.Result, &out))
require.Len(t, out.Entries, 1)
assert.Equal(t, "f.txt", out.Entries[0].Name)
}
func TestRelay_AgentFsGrep_Dispatched(t *testing.T) {
t.Parallel()
rig := newRelayTestRig(t)
cancel, doneCh := rig.runRelay(t)
t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })
require.NoError(t, os.WriteFile(filepath.Join(rig.sess.CwdPath, "code.go"), []byte("package x\nfunc Needle() {}\n"), 0o644))
resp := relayRoundTrip(t, rig, 202, "fs/grep", map[string]any{"pattern": "Needle", "sessionId": "a"})
require.True(t, resp.IsResponse())
require.Nil(t, resp.Error)
var out struct {
Matches []struct {
File string `json:"file"`
Line int `json:"line"`
} `json:"matches"`
}
require.NoError(t, json.Unmarshal(resp.Result, &out))
require.Len(t, out.Matches, 1)
assert.Equal(t, "code.go", out.Matches[0].File)
assert.Equal(t, 2, out.Matches[0].Line)
}
func TestRelay_AgentFsGrep_PathEscapeRejected(t *testing.T) {
t.Parallel()
rig := newRelayTestRig(t)
cancel, doneCh := rig.runRelay(t)
t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })
outside := filepath.Join(filepath.VolumeName(rig.sess.CwdPath)+string(filepath.Separator), "etc")
resp := relayRoundTrip(t, rig, 203, "fs/grep", map[string]any{"pattern": "x", "path": outside, "sessionId": "a"})
require.True(t, resp.IsResponse())
require.NotNil(t, resp.Error)
assert.Equal(t, -32001, resp.Error.Code)
}
+267
View File
@@ -31,6 +31,18 @@ import (
type fakeEventRepo struct {
mu sync.Mutex
events []*acp.Event
// 回合状态机:relay turn 测试用。
turns []*acp.Turn
lastStopReason *string
// 成本核算:relay usage 测试用。
usageLedger []*acp.SessionUsageRecord
totPrompt int64
totComplete int64
totThink int64
totCost float64
dedupSources map[int64]struct{}
}
func (f *fakeEventRepo) InsertEvent(_ context.Context, e *acp.Event) (*acp.Event, error) {
@@ -81,6 +93,15 @@ func (f *fakeEventRepo) InsertSession(context.Context, *acp.Session) (*acp.Sessi
func (f *fakeEventRepo) GetSessionByID(context.Context, uuid.UUID) (*acp.Session, error) {
panic("n/a")
}
func (f *fakeEventRepo) GetSessionByStepID(context.Context, uuid.UUID) (*acp.Session, error) {
panic("n/a")
}
func (f *fakeEventRepo) GetCrashedSessionForResume(context.Context, uuid.UUID) (*acp.Session, error) {
panic("n/a")
}
func (f *fakeEventRepo) ResetSessionForResume(context.Context, uuid.UUID) error {
panic("n/a")
}
func (f *fakeEventRepo) ListSessionsByUser(context.Context, uuid.UUID, *uuid.UUID) ([]*acp.Session, error) {
panic("n/a")
}
@@ -90,6 +111,9 @@ func (f *fakeEventRepo) ListAllSessions(context.Context, *uuid.UUID) ([]*acp.Ses
func (f *fakeEventRepo) UpdateSessionRunning(context.Context, uuid.UUID, string, int32) error {
panic("n/a")
}
func (f *fakeEventRepo) UpdateSessionSandboxMode(context.Context, uuid.UUID, string) error {
panic("n/a")
}
func (f *fakeEventRepo) UpdateSessionFinished(context.Context, uuid.UUID, acp.SessionStatus, *int32, *string) error {
panic("n/a")
}
@@ -115,6 +139,67 @@ func (f *fakeEventRepo) ListEventsSince(context.Context, uuid.UUID, int64, int32
func (f *fakeEventRepo) PurgeEventsBefore(context.Context, time.Time) (int64, error) {
panic("n/a")
}
func (f *fakeEventRepo) InsertSessionUsage(_ context.Context, rec *acp.SessionUsageRecord) (bool, error) {
f.mu.Lock()
defer f.mu.Unlock()
if rec.SourceEventID != nil {
if f.dedupSources == nil {
f.dedupSources = map[int64]struct{}{}
}
if _, seen := f.dedupSources[*rec.SourceEventID]; seen {
return false, nil
}
f.dedupSources[*rec.SourceEventID] = struct{}{}
}
cp := *rec
f.usageLedger = append(f.usageLedger, &cp)
return true, nil
}
func (f *fakeEventRepo) AddSessionUsageTotals(_ context.Context, _ uuid.UUID, dp, dc, dt int64, dCost float64) (int64, int64, int64, float64, error) {
f.mu.Lock()
defer f.mu.Unlock()
f.totPrompt += dp
f.totComplete += dc
f.totThink += dt
f.totCost += dCost
return f.totPrompt, f.totComplete, f.totThink, f.totCost, nil
}
func (f *fakeEventRepo) SumProjectCostUSD(context.Context, uuid.UUID, time.Time) (float64, error) {
f.mu.Lock()
defer f.mu.Unlock()
return f.totCost, nil
}
func (f *fakeEventRepo) MarkSessionTerminatedReason(context.Context, uuid.UUID, string) error {
return nil
}
func (f *fakeEventRepo) ListSessionsForReaper(context.Context, time.Time, time.Time) ([]acp.ReaperSession, error) {
panic("n/a")
}
func (f *fakeEventRepo) ListSessionUsage(context.Context, uuid.UUID, int32) ([]*acp.SessionUsageRecord, error) {
f.mu.Lock()
defer f.mu.Unlock()
out := make([]*acp.SessionUsageRecord, len(f.usageLedger))
copy(out, f.usageLedger)
return out, nil
}
func (f *fakeEventRepo) SummarizeAcpUsageByUser(context.Context, time.Time, time.Time) ([]acp.AcpUsageUserRow, error) {
panic("n/a")
}
func (f *fakeEventRepo) SummarizeAcpUsageByModel(context.Context, time.Time, time.Time) ([]acp.AcpUsageModelRow, error) {
panic("n/a")
}
func (f *fakeEventRepo) SummarizeAcpUsageByDay(context.Context, time.Time, time.Time) ([]acp.AcpUsageDayRow, error) {
panic("n/a")
}
func (f *fakeEventRepo) SessionRollup(context.Context, acp.DashboardFilter) (acp.SessionRollup, error) {
panic("n/a")
}
func (f *fakeEventRepo) SessionsByDay(context.Context, acp.DashboardFilter) ([]acp.SessionDayRow, error) {
panic("n/a")
}
func (f *fakeEventRepo) SessionTimeline(context.Context, uuid.UUID) ([]acp.SessionTimelineBucket, error) {
panic("n/a")
}
func (f *fakeEventRepo) ListConfigFiles(context.Context, uuid.UUID) ([]*acp.ConfigFile, error) {
panic("n/a")
}
@@ -133,6 +218,83 @@ func (f *fakeEventRepo) WithTx(pgx.Tx) acp.Repository {
return f
}
// ----- Turn 方法(relay turn 测试用,功能性实现)-----
func (f *fakeEventRepo) InsertTurn(_ context.Context, t *acp.Turn) (*acp.Turn, error) {
f.mu.Lock()
defer f.mu.Unlock()
cp := *t
cp.ID = int64(len(f.turns) + 1)
cp.Status = acp.TurnInProgress
f.turns = append(f.turns, &cp)
out := cp
return &out, nil
}
func (f *fakeEventRepo) MarkTurnCompleted(_ context.Context, _ uuid.UUID, turnIndex int, stop string, updateCount int) error {
f.mu.Lock()
defer f.mu.Unlock()
for _, t := range f.turns {
if t.TurnIndex == turnIndex {
t.Status = acp.TurnCompleted
s := stop
t.StopReason = &s
t.UpdateCount = updateCount
}
}
return nil
}
func (f *fakeEventRepo) MarkTurnAborted(_ context.Context, _ uuid.UUID, turnIndex int) error {
f.mu.Lock()
defer f.mu.Unlock()
for _, t := range f.turns {
if t.TurnIndex == turnIndex {
t.Status = acp.TurnAborted
}
}
return nil
}
func (f *fakeEventRepo) IncrementTurnUpdateCount(context.Context, uuid.UUID, int) error { return nil }
func (f *fakeEventRepo) GetLatestTurnBySession(context.Context, uuid.UUID) (*acp.Turn, error) {
f.mu.Lock()
defer f.mu.Unlock()
if len(f.turns) == 0 {
return nil, nil
}
out := *f.turns[len(f.turns)-1]
return &out, nil
}
func (f *fakeEventRepo) ListTurnsBySession(context.Context, uuid.UUID) ([]*acp.Turn, error) {
f.mu.Lock()
defer f.mu.Unlock()
out := make([]*acp.Turn, len(f.turns))
copy(out, f.turns)
return out, nil
}
func (f *fakeEventRepo) UpdateSessionLastStopReason(_ context.Context, _ uuid.UUID, reason string) error {
f.mu.Lock()
defer f.mu.Unlock()
r := reason
f.lastStopReason = &r
return nil
}
func (f *fakeEventRepo) PurgeTurnsBefore(context.Context, time.Time) (int64, error) { return 0, nil }
func (f *fakeEventRepo) turnSnapshot() []*acp.Turn {
f.mu.Lock()
defer f.mu.Unlock()
out := make([]*acp.Turn, len(f.turns))
for i, t := range f.turns {
cp := *t
out[i] = &cp
}
return out
}
func (f *fakeEventRepo) lastStop() *string {
f.mu.Lock()
defer f.mu.Unlock()
return f.lastStopReason
}
// relayTestRig wires up a Relay with two io.Pipe pairs that simulate the agent
// subprocess. Returns:
// - r: the relay under test
@@ -470,3 +632,108 @@ func (f *fakeEventRepo) ExpirePendingPermissionRequestsBySession(context.Context
func (f *fakeEventRepo) ExpireStalePermissionRequests(context.Context, time.Time) (int64, error) {
return 0, nil
}
// ---------- Turn detection ----------
// TestRelay_TurnDetection feeds a session/update notification then a string-id
// session/prompt response with stopReason, asserting the tracker hooks fire:
// the turn is created, completed with the stop reason, and update_count tracked.
func TestRelay_TurnDetection(t *testing.T) {
t.Parallel()
rig := newRelayTestRig(t)
bus := &fakeBus{}
tr := acp.NewTurnTracker(rig.repo, bus, rig.sess, nil)
rig.r.SetTracker(tr)
cancel, doneCh := rig.runRelay(t)
t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })
// Register the turn (as session_service would before sending the prompt).
_, err := tr.StartTurn(context.Background(), "client-init")
require.NoError(t, err)
agentEnc := acp.NewEncoder(rig.agentWriter)
// Agent emits two session/update notifications then the prompt response.
notif, err := acp.NewNotification("session/update", map[string]any{
"sessionId": "agent-sess-1",
"update": map[string]any{"sessionUpdate": "agent_message_chunk"},
})
require.NoError(t, err)
require.NoError(t, agentEnc.Encode(notif))
require.NoError(t, agentEnc.Encode(notif))
idJSON, _ := json.Marshal("client-init")
resp := &acp.Message{
JSONRPC: "2.0", ID: idJSON,
Result: json.RawMessage(`{"stopReason":"max_tokens"}`),
}
require.NoError(t, agentEnc.Encode(resp))
// The turn must complete with stop_reason=max_tokens and update_count=2.
require.Eventually(t, func() bool {
ts := rig.repo.turnSnapshot()
if len(ts) != 1 {
return false
}
return ts[0].Status == acp.TurnCompleted &&
ts[0].StopReason != nil && *ts[0].StopReason == "max_tokens" &&
ts[0].UpdateCount == 2
}, 3*time.Second, 20*time.Millisecond, "turn not completed as expected")
// last_stop_reason persisted on the session.
require.Eventually(t, func() bool {
ls := rig.repo.lastStop()
return ls != nil && *ls == "max_tokens"
}, 2*time.Second, 20*time.Millisecond)
// Bus received turn_completed then session_idle.
require.Eventually(t, func() bool {
return len(bus.snapshot()) >= 2
}, 2*time.Second, 20*time.Millisecond)
evs := bus.snapshot()
assert.Equal(t, acp.TurnCompletedEvent, evs[0].Kind)
assert.Equal(t, acp.SessionIdleEvent, evs[1].Kind)
assert.Equal(t, acp.StopMaxTokens, evs[0].StopReason)
}
// TestRelay_IntResponseStillRoutesToInflight is a regression guard: with a
// tracker installed, an int64-id response must still route to inflight (the
// existing Call mechanism), not the turn path.
func TestRelay_IntResponseStillRoutesToInflight(t *testing.T) {
t.Parallel()
rig := newRelayTestRig(t)
rig.r.SetTracker(acp.NewTurnTracker(rig.repo, &fakeBus{}, rig.sess, nil))
cancel, doneCh := rig.runRelay(t)
t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })
agentDec := acp.NewDecoder(rig.agentReader, 0)
agentEnc := acp.NewEncoder(rig.agentWriter)
agentDone := make(chan error, 1)
go func() {
req, err := agentDec.DecodeMessage()
if err != nil {
agentDone <- err
return
}
respMsg, err := acp.NewResponseOK(req.ID, map[string]any{"protocolVersion": 1})
if err != nil {
agentDone <- err
return
}
agentDone <- agentEnc.Encode(respMsg)
}()
ctx, cancelCall := context.WithTimeout(context.Background(), 3*time.Second)
defer cancelCall()
resp, err := rig.r.Call(ctx, "initialize", map[string]any{"protocolVersion": 1})
require.NoError(t, err)
require.True(t, resp.IsResponse())
select {
case err := <-agentDone:
require.NoError(t, err)
case <-time.After(2 * time.Second):
t.Fatal("agent goroutine did not finish")
}
}
+230 -28
View File
@@ -36,14 +36,22 @@ type Repository interface {
// Session
InsertSession(ctx context.Context, s *Session) (*Session, error)
GetSessionByID(ctx context.Context, id uuid.UUID) (*Session, error)
// GetSessionByStepID 返回某编排器 step 最近的 session(任意状态);无则 NotFound。
GetSessionByStepID(ctx context.Context, stepID uuid.UUID) (*Session, error)
// GetCrashedSessionForResume 返回某编排器 step 最近一次 crashed/exited 的 session(用于 resume)。
GetCrashedSessionForResume(ctx context.Context, stepID uuid.UUID) (*Session, error)
// requirementID 非 nil 时仅返回关联该需求的 sessions。
ListSessionsByUser(ctx context.Context, userID uuid.UUID, requirementID *uuid.UUID) ([]*Session, error)
ListAllSessions(ctx context.Context, requirementID *uuid.UUID) ([]*Session, error)
UpdateSessionRunning(ctx context.Context, id uuid.UUID, agentSessionID string, pid int32) error
// UpdateSessionSandboxMode 记录本 session 运行所用的沙箱模式(审计/取证)。
UpdateSessionSandboxMode(ctx context.Context, id uuid.UUID, mode string) error
UpdateSessionFinished(ctx context.Context, id uuid.UUID, status SessionStatus, exitCode *int32, lastErr *string) error
// MarkSessionFailedIfActive 把仍处于 starting/running 的 session 标记为 crashed
// 并写入 lastErr;已是终态时不更新(守卫式,避免覆盖 onExit 写入的终态)。
MarkSessionFailedIfActive(ctx context.Context, id uuid.UUID, lastErr string) (bool, error)
// ResetSessionForResume 把崩溃/退出 session 复用为新一轮(回 starting、清运行时字段)。
ResetSessionForResume(ctx context.Context, id uuid.UUID) error
CountActiveSessions(ctx context.Context) (int64, error)
CountActiveSessionsByUser(ctx context.Context, userID uuid.UUID) (int64, error)
ResetStuckSessionsOnRestart(ctx context.Context) ([]*Session, error)
@@ -57,6 +65,32 @@ type Repository interface {
ListEventsSince(ctx context.Context, sessionID uuid.UUID, sinceID int64, limit int32) ([]*Event, error)
PurgeEventsBefore(ctx context.Context, before time.Time) (int64, error)
// 成本核算 / 预算 / reaper
InsertSessionUsage(ctx context.Context, r *SessionUsageRecord) (inserted bool, err error)
AddSessionUsageTotals(ctx context.Context, sid uuid.UUID, dp, dc, dt int64, dCost float64) (totPrompt, totCompletion, totThinking int64, totCost float64, err error)
SumProjectCostUSD(ctx context.Context, projectID uuid.UUID, since time.Time) (float64, error)
MarkSessionTerminatedReason(ctx context.Context, sid uuid.UUID, reason string) error
ListSessionsForReaper(ctx context.Context, wallClockStartedBefore, idleSince time.Time) ([]ReaperSession, error)
ListSessionUsage(ctx context.Context, sessionID uuid.UUID, limit int32) ([]*SessionUsageRecord, error)
SummarizeAcpUsageByUser(ctx context.Context, from, to time.Time) ([]AcpUsageUserRow, error)
SummarizeAcpUsageByModel(ctx context.Context, from, to time.Time) ([]AcpUsageModelRow, error)
SummarizeAcpUsageByDay(ctx context.Context, from, to time.Time) ([]AcpUsageDayRow, error)
// run-history dashboard(只读聚合)
SessionRollup(ctx context.Context, f DashboardFilter) (SessionRollup, error)
SessionsByDay(ctx context.Context, f DashboardFilter) ([]SessionDayRow, error)
SessionTimeline(ctx context.Context, sessionID uuid.UUID) ([]SessionTimelineBucket, error)
// Turn(回合完成检测)
InsertTurn(ctx context.Context, t *Turn) (*Turn, error)
MarkTurnCompleted(ctx context.Context, sessionID uuid.UUID, turnIndex int, stop string, updateCount int) error
MarkTurnAborted(ctx context.Context, sessionID uuid.UUID, turnIndex int) error
IncrementTurnUpdateCount(ctx context.Context, sessionID uuid.UUID, turnIndex int) error
GetLatestTurnBySession(ctx context.Context, sessionID uuid.UUID) (*Turn, error)
ListTurnsBySession(ctx context.Context, sessionID uuid.UUID) ([]*Turn, error)
UpdateSessionLastStopReason(ctx context.Context, sessionID uuid.UUID, reason string) error
PurgeTurnsBefore(ctx context.Context, before time.Time) (int64, error)
// PermissionRequest
InsertPermissionRequest(ctx context.Context, p *PermissionRequest) (*PermissionRequest, error)
GetPermissionRequestByID(ctx context.Context, id uuid.UUID) (*PermissionRequest, error)
@@ -159,6 +193,10 @@ func (r *pgRepo) CreateAgentKind(ctx context.Context, k *AgentKind) (*AgentKind,
ToolAllowlist: normalizeStrSlice(k.ToolAllowlist),
ClientType: string(k.ClientType),
EncryptedMcpServers: k.EncryptedMCPServers,
ModelID: toPgUUIDPtr(k.ModelID),
MaxCostUsd: float64PtrToNumeric(k.MaxCostUSD),
MaxTokens: k.MaxTokens,
MaxWallClockSeconds: k.MaxWallClockSeconds,
})
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "create agent kind")
@@ -218,6 +256,10 @@ func (r *pgRepo) UpdateAgentKind(ctx context.Context, k *AgentKind) (*AgentKind,
ToolAllowlist: normalizeStrSlice(k.ToolAllowlist),
ClientType: string(k.ClientType),
EncryptedMcpServers: k.EncryptedMCPServers,
ModelID: toPgUUIDPtr(k.ModelID),
MaxCostUsd: float64PtrToNumeric(k.MaxCostUSD),
MaxTokens: k.MaxTokens,
MaxWallClockSeconds: k.MaxWallClockSeconds,
})
if err != nil {
return nil, pgxNoRowsToErrNotFound(err, errs.CodeAcpAgentKindNotFound, "agent kind not found")
@@ -300,6 +342,10 @@ func rowToAgentKind(row acpsqlc.AcpAgentKind) *AgentKind {
CreatedBy: fromPgUUID(row.CreatedBy),
CreatedAt: row.CreatedAt.Time,
UpdatedAt: row.UpdatedAt.Time,
ModelID: fromPgUUIDPtr(row.ModelID),
MaxCostUSD: numericToFloat64Ptr(row.MaxCostUsd),
MaxTokens: row.MaxTokens,
MaxWallClockSeconds: row.MaxWallClockSeconds,
}
}
@@ -328,17 +374,21 @@ func normalizeStrSlice(s []string) []string {
func (r *pgRepo) InsertSession(ctx context.Context, s *Session) (*Session, error) {
row, err := r.q.InsertSession(ctx, acpsqlc.InsertSessionParams{
ID: toPgUUID(s.ID),
WorkspaceID: toPgUUID(s.WorkspaceID),
ProjectID: toPgUUID(s.ProjectID),
AgentKindID: toPgUUID(s.AgentKindID),
UserID: toPgUUID(s.UserID),
IssueID: toPgUUIDPtr(s.IssueID),
RequirementID: toPgUUIDPtr(s.RequirementID),
Branch: s.Branch,
CwdPath: s.CwdPath,
IsMainWorktree: s.IsMainWorktree,
Status: string(s.Status),
ID: toPgUUID(s.ID),
WorkspaceID: toPgUUID(s.WorkspaceID),
ProjectID: toPgUUID(s.ProjectID),
AgentKindID: toPgUUID(s.AgentKindID),
UserID: toPgUUID(s.UserID),
IssueID: toPgUUIDPtr(s.IssueID),
RequirementID: toPgUUIDPtr(s.RequirementID),
Branch: s.Branch,
CwdPath: s.CwdPath,
IsMainWorktree: s.IsMainWorktree,
Status: string(s.Status),
OrchestratorStepID: toPgUUIDPtr(s.OrchestratorStepID),
BudgetMaxCostUsd: float64PtrToNumeric(s.BudgetMaxCostUSD),
BudgetMaxTokens: s.BudgetMaxTokens,
BudgetMaxWallClockSeconds: s.BudgetMaxWallClockSeconds,
})
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "insert session")
@@ -354,6 +404,22 @@ func (r *pgRepo) GetSessionByID(ctx context.Context, id uuid.UUID) (*Session, er
return rowToSession(row), nil
}
func (r *pgRepo) GetSessionByStepID(ctx context.Context, stepID uuid.UUID) (*Session, error) {
row, err := r.q.GetSessionByStepID(ctx, toPgUUID(stepID))
if err != nil {
return nil, pgxNoRowsToErrNotFound(err, errs.CodeAcpSessionNotFound, "session not found for step")
}
return rowToSession(row), nil
}
func (r *pgRepo) GetCrashedSessionForResume(ctx context.Context, stepID uuid.UUID) (*Session, error) {
row, err := r.q.GetCrashedSessionForResume(ctx, toPgUUID(stepID))
if err != nil {
return nil, pgxNoRowsToErrNotFound(err, errs.CodeAcpSessionNotFound, "no crashed session to resume")
}
return rowToSession(row), nil
}
func (r *pgRepo) ListSessionsByUser(ctx context.Context, userID uuid.UUID, requirementID *uuid.UUID) ([]*Session, error) {
rows, err := r.q.ListSessionsByUser(ctx, acpsqlc.ListSessionsByUserParams{
UserID: toPgUUID(userID),
@@ -392,6 +458,17 @@ func (r *pgRepo) UpdateSessionRunning(ctx context.Context, id uuid.UUID, agentSe
return nil
}
func (r *pgRepo) UpdateSessionSandboxMode(ctx context.Context, id uuid.UUID, mode string) error {
m := mode
if err := r.q.UpdateSessionSandboxMode(ctx, acpsqlc.UpdateSessionSandboxModeParams{
ID: toPgUUID(id),
SandboxMode: &m,
}); err != nil {
return errs.Wrap(err, errs.CodeInternal, "update session sandbox mode")
}
return nil
}
func (r *pgRepo) UpdateSessionFinished(ctx context.Context, id uuid.UUID, status SessionStatus, exitCode *int32, lastErr *string) error {
if err := r.q.UpdateSessionFinished(ctx, acpsqlc.UpdateSessionFinishedParams{
ID: toPgUUID(id),
@@ -415,6 +492,13 @@ func (r *pgRepo) MarkSessionFailedIfActive(ctx context.Context, id uuid.UUID, la
return n > 0, nil
}
func (r *pgRepo) ResetSessionForResume(ctx context.Context, id uuid.UUID) error {
if err := r.q.ResetSessionForResume(ctx, toPgUUID(id)); err != nil {
return errs.Wrap(err, errs.CodeInternal, "reset session for resume")
}
return nil
}
func (r *pgRepo) CountActiveSessions(ctx context.Context) (int64, error) {
n, err := r.q.CountActiveSessions(ctx)
if err != nil {
@@ -480,23 +564,34 @@ func rowToSession(row acpsqlc.AcpSession) *Session {
endedAt = &t
}
return &Session{
ID: fromPgUUID(row.ID),
WorkspaceID: fromPgUUID(row.WorkspaceID),
ProjectID: fromPgUUID(row.ProjectID),
AgentKindID: fromPgUUID(row.AgentKindID),
UserID: fromPgUUID(row.UserID),
IssueID: fromPgUUIDPtr(row.IssueID),
RequirementID: fromPgUUIDPtr(row.RequirementID),
AgentSessionID: row.AgentSessionID,
Branch: row.Branch,
CwdPath: row.CwdPath,
IsMainWorktree: row.IsMainWorktree,
Status: SessionStatus(row.Status),
PID: row.Pid,
ExitCode: row.ExitCode,
LastError: row.LastError,
StartedAt: row.StartedAt.Time,
EndedAt: endedAt,
ID: fromPgUUID(row.ID),
WorkspaceID: fromPgUUID(row.WorkspaceID),
ProjectID: fromPgUUID(row.ProjectID),
AgentKindID: fromPgUUID(row.AgentKindID),
UserID: fromPgUUID(row.UserID),
IssueID: fromPgUUIDPtr(row.IssueID),
RequirementID: fromPgUUIDPtr(row.RequirementID),
AgentSessionID: row.AgentSessionID,
Branch: row.Branch,
CwdPath: row.CwdPath,
IsMainWorktree: row.IsMainWorktree,
Status: SessionStatus(row.Status),
PID: row.Pid,
ExitCode: row.ExitCode,
LastError: row.LastError,
StartedAt: row.StartedAt.Time,
EndedAt: endedAt,
LastStopReason: row.LastStopReason,
OrchestratorStepID: fromPgUUIDPtr(row.OrchestratorStepID),
PromptTokens: row.PromptTokens,
CompletionTokens: row.CompletionTokens,
ThinkingTokens: row.ThinkingTokens,
TotalCostUSD: numericToFloat64(row.TotalCostUsd),
LastActivityAt: row.LastActivityAt.Time,
BudgetMaxCostUSD: numericToFloat64Ptr(row.BudgetMaxCostUsd),
BudgetMaxTokens: row.BudgetMaxTokens,
BudgetMaxWallClockSeconds: row.BudgetMaxWallClockSeconds,
TerminatedReason: row.TerminatedReason,
}
}
@@ -556,6 +651,113 @@ func rowToEvent(row acpsqlc.AcpEvent) *Event {
}
}
// ===== Turn(回合完成检测)=====
func (r *pgRepo) InsertTurn(ctx context.Context, t *Turn) (*Turn, error) {
row, err := r.q.InsertTurn(ctx, acpsqlc.InsertTurnParams{
SessionID: toPgUUID(t.SessionID),
TurnIndex: int32(t.TurnIndex),
PromptRequestID: t.PromptRequestID,
Status: string(t.Status),
})
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "insert acp turn")
}
return rowToTurn(row), nil
}
func (r *pgRepo) MarkTurnCompleted(ctx context.Context, sessionID uuid.UUID, turnIndex int, stop string, updateCount int) error {
if err := r.q.MarkTurnCompleted(ctx, acpsqlc.MarkTurnCompletedParams{
SessionID: toPgUUID(sessionID),
TurnIndex: int32(turnIndex),
StopReason: &stop,
UpdateCount: int32(updateCount),
}); err != nil {
return errs.Wrap(err, errs.CodeInternal, "mark turn completed")
}
return nil
}
func (r *pgRepo) MarkTurnAborted(ctx context.Context, sessionID uuid.UUID, turnIndex int) error {
if err := r.q.MarkTurnAborted(ctx, acpsqlc.MarkTurnAbortedParams{
SessionID: toPgUUID(sessionID),
TurnIndex: int32(turnIndex),
}); err != nil {
return errs.Wrap(err, errs.CodeInternal, "mark turn aborted")
}
return nil
}
func (r *pgRepo) IncrementTurnUpdateCount(ctx context.Context, sessionID uuid.UUID, turnIndex int) error {
if err := r.q.IncrementTurnUpdateCount(ctx, acpsqlc.IncrementTurnUpdateCountParams{
SessionID: toPgUUID(sessionID),
TurnIndex: int32(turnIndex),
}); err != nil {
return errs.Wrap(err, errs.CodeInternal, "increment turn update count")
}
return nil
}
func (r *pgRepo) GetLatestTurnBySession(ctx context.Context, sessionID uuid.UUID) (*Turn, error) {
row, err := r.q.GetLatestTurnBySession(ctx, toPgUUID(sessionID))
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return nil, nil
}
return nil, errs.Wrap(err, errs.CodeInternal, "get latest turn by session")
}
return rowToTurn(row), nil
}
func (r *pgRepo) ListTurnsBySession(ctx context.Context, sessionID uuid.UUID) ([]*Turn, error) {
rows, err := r.q.ListTurnsBySession(ctx, toPgUUID(sessionID))
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "list turns by session")
}
out := make([]*Turn, 0, len(rows))
for _, row := range rows {
out = append(out, rowToTurn(row))
}
return out, nil
}
func (r *pgRepo) UpdateSessionLastStopReason(ctx context.Context, sessionID uuid.UUID, reason string) error {
if err := r.q.UpdateSessionLastStopReason(ctx, acpsqlc.UpdateSessionLastStopReasonParams{
ID: toPgUUID(sessionID),
LastStopReason: &reason,
}); err != nil {
return errs.Wrap(err, errs.CodeInternal, "update session last stop reason")
}
return nil
}
func (r *pgRepo) PurgeTurnsBefore(ctx context.Context, before time.Time) (int64, error) {
n, err := r.q.PurgeTurnsBefore(ctx, pgtype.Timestamptz{Time: before, Valid: true})
if err != nil {
return 0, errs.Wrap(err, errs.CodeInternal, "purge acp turns")
}
return n, nil
}
func rowToTurn(row acpsqlc.AcpTurn) *Turn {
var completedAt *time.Time
if row.CompletedAt.Valid {
t := row.CompletedAt.Time
completedAt = &t
}
return &Turn{
ID: row.ID,
SessionID: fromPgUUID(row.SessionID),
TurnIndex: int(row.TurnIndex),
PromptRequestID: row.PromptRequestID,
Status: TurnStatus(row.Status),
StopReason: row.StopReason,
UpdateCount: int(row.UpdateCount),
StartedAt: row.StartedAt.Time,
CompletedAt: completedAt,
}
}
// ===== PermissionRequest =====
func (r *pgRepo) InsertPermissionRequest(ctx context.Context, p *PermissionRequest) (*PermissionRequest, error) {
+218
View File
@@ -234,6 +234,86 @@ func TestPgRepo_Event_OrderAndPurge(t *testing.T) {
assert.Equal(t, int64(3), n)
}
func TestPgRepo_Turn_Lifecycle(t *testing.T) {
t.Parallel()
ctx := context.Background()
repo, pool := setupRepo(t)
uid := mustInsertUser(t, ctx, pool, "u6@local", false)
pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)
k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{
ID: uuid.New(), Name: "kind6", DisplayName: "x",
BinaryPath: "x", Enabled: true, CreatedBy: uid,
})
sess, _ := repo.InsertSession(ctx, &acp.Session{
ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
AgentKindID: k.ID, UserID: uid,
Branch: "main", CwdPath: "/tmp", IsMainWorktree: true, Status: acp.SessionRunning,
})
// 新 session 无回合
latest, err := repo.GetLatestTurnBySession(ctx, sess.ID)
require.NoError(t, err)
assert.Nil(t, latest)
// 插入回合 0
pidStr := "client-init"
turn, err := repo.InsertTurn(ctx, &acp.Turn{
SessionID: sess.ID, TurnIndex: 0, PromptRequestID: &pidStr, Status: acp.TurnInProgress,
})
require.NoError(t, err)
assert.Equal(t, 0, turn.TurnIndex)
assert.Equal(t, acp.TurnInProgress, turn.Status)
// IncrementTurnUpdateCount ×2(开放回合)
require.NoError(t, repo.IncrementTurnUpdateCount(ctx, sess.ID, 0))
require.NoError(t, repo.IncrementTurnUpdateCount(ctx, sess.ID, 0))
// 完成回合:MarkTurnCompleted 用内存计数覆盖 update_count
require.NoError(t, repo.MarkTurnCompleted(ctx, sess.ID, 0, "end_turn", 5))
require.NoError(t, repo.UpdateSessionLastStopReason(ctx, sess.ID, "end_turn"))
ts, err := repo.ListTurnsBySession(ctx, sess.ID)
require.NoError(t, err)
require.Len(t, ts, 1)
assert.Equal(t, acp.TurnCompleted, ts[0].Status)
require.NotNil(t, ts[0].StopReason)
assert.Equal(t, "end_turn", *ts[0].StopReason)
assert.Equal(t, 5, ts[0].UpdateCount)
require.NotNil(t, ts[0].CompletedAt)
// last_stop_reason round-trips on session SELECT
got, _ := repo.GetSessionByID(ctx, sess.ID)
require.NotNil(t, got.LastStopReason)
assert.Equal(t, "end_turn", *got.LastStopReason)
// GetLatestTurnBySession 返回 index 最大的回合
latest, err = repo.GetLatestTurnBySession(ctx, sess.ID)
require.NoError(t, err)
require.NotNil(t, latest)
assert.Equal(t, 0, latest.TurnIndex)
// UNIQUE(session_id, turn_index) 约束
_, err = repo.InsertTurn(ctx, &acp.Turn{SessionID: sess.ID, TurnIndex: 0, Status: acp.TurnInProgress})
require.Error(t, err, "duplicate (session_id, turn_index) must violate UNIQUE")
// 插入并中止回合 1
_, err = repo.InsertTurn(ctx, &acp.Turn{SessionID: sess.ID, TurnIndex: 1, Status: acp.TurnInProgress})
require.NoError(t, err)
require.NoError(t, repo.MarkTurnAborted(ctx, sess.ID, 1))
latest, _ = repo.GetLatestTurnBySession(ctx, sess.ID)
assert.Equal(t, 1, latest.TurnIndex)
assert.Equal(t, acp.TurnAborted, latest.Status)
// Purge before now+1m → 全删
purged, err := repo.PurgeTurnsBefore(ctx, time.Now().Add(time.Minute))
require.NoError(t, err)
assert.Equal(t, int64(2), purged)
// 幂等
purged2, _ := repo.PurgeTurnsBefore(ctx, time.Now().Add(time.Minute))
assert.Equal(t, int64(0), purged2)
}
// ===== test helpers =====
func mustInsertUser(t *testing.T, ctx context.Context, pool *pgxpool.Pool, email string, admin bool) uuid.UUID {
@@ -259,3 +339,141 @@ func mustInsertProjectWorkspace(t *testing.T, ctx context.Context, pool *pgxpool
require.NoError(t, err)
return pid, wsid
}
// TestPgRepo_UsagePersist_BumpsDashboardColumns 验证 AddSessionUsageTotals 同时
// 把 dashboard rollup 列(cost_usd / tokens_total)与运行时累加器对齐。
func TestPgRepo_UsagePersist_BumpsDashboardColumns(t *testing.T) {
t.Parallel()
ctx := context.Background()
repo, pool := setupRepo(t)
uid := mustInsertUser(t, ctx, pool, "dash-usage@local", false)
pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)
k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{
ID: uuid.New(), Name: "kind-dash-usage", DisplayName: "x",
BinaryPath: "x", Enabled: true, CreatedBy: uid,
})
sess, _ := repo.InsertSession(ctx, &acp.Session{
ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
AgentKindID: k.ID, UserID: uid,
Branch: "main", CwdPath: "/tmp", IsMainWorktree: true, Status: acp.SessionRunning,
})
// 两次回合累加。
_, _, _, _, err := repo.AddSessionUsageTotals(ctx, sess.ID, 100, 50, 10, 0.30)
require.NoError(t, err)
_, _, _, totCost, err := repo.AddSessionUsageTotals(ctx, sess.ID, 200, 40, 0, 0.20)
require.NoError(t, err)
assert.InDelta(t, 0.50, totCost, 1e-9)
got, err := repo.GetSessionByID(ctx, sess.ID)
require.NoError(t, err)
assert.InDelta(t, 0.50, got.TotalCostUSD, 1e-9)
// 运行时累加器 totals。
assert.Equal(t, int64(300), got.PromptTokens)
assert.Equal(t, int64(90), got.CompletionTokens)
assert.Equal(t, int64(10), got.ThinkingTokens)
// dashboard 汇总应读到与运行时一致的 cost / tokens。
rollup, err := repo.SessionRollup(ctx, acp.DashboardFilter{
UserID: &uid, From: time.Now().Add(-time.Hour), To: time.Now().Add(time.Hour),
})
require.NoError(t, err)
assert.Equal(t, int64(1), rollup.Total)
assert.Equal(t, int64(1), rollup.Active)
assert.InDelta(t, 0.50, rollup.TotalCostUSD, 1e-9)
assert.Equal(t, int64(400), rollup.TotalTokens) // 100+50+10 + 200+40+0
}
// TestPgRepo_Dashboard_RollupTimelineByDay 覆盖三个 dashboard 查询:
// SessionRollup(按状态计数 + 成本 + scope)、SessionsByDay、SessionTimeline。
func TestPgRepo_Dashboard_RollupTimelineByDay(t *testing.T) {
t.Parallel()
ctx := context.Background()
repo, pool := setupRepo(t)
owner := mustInsertUser(t, ctx, pool, "dash-owner@local", false)
other := mustInsertUser(t, ctx, pool, "dash-other@local", false)
pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, owner)
k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{
ID: uuid.New(), Name: "kind-dash", DisplayName: "x",
BinaryPath: "x", Enabled: true, CreatedBy: owner,
})
mk := func(u uuid.UUID, status acp.SessionStatus) *acp.Session {
s, err := repo.InsertSession(ctx, &acp.Session{
ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
AgentKindID: k.ID, UserID: u,
Branch: "b", CwdPath: "/tmp", IsMainWorktree: false, Status: status,
})
require.NoError(t, err)
return s
}
exited := mk(owner, acp.SessionRunning)
crashed := mk(owner, acp.SessionRunning)
mk(other, acp.SessionRunning) // 另一用户:owner-scope 下不计入
// owner 的两个会话各累加成本并结束。
_, _, _, _, err := repo.AddSessionUsageTotals(ctx, exited.ID, 100, 100, 0, 1.00)
require.NoError(t, err)
_, _, _, _, err = repo.AddSessionUsageTotals(ctx, crashed.ID, 50, 50, 0, 0.50)
require.NoError(t, err)
ec := int32(0)
require.NoError(t, repo.UpdateSessionFinished(ctx, exited.ID, acp.SessionExited, &ec, nil))
require.NoError(t, repo.UpdateSessionFinished(ctx, crashed.ID, acp.SessionCrashed, &ec, nil))
from := time.Now().Add(-time.Hour)
to := time.Now().Add(time.Hour)
// owner scope:只统计 owner 的 2 个会话。
ownerScope := acp.DashboardFilter{UserID: &owner, From: from, To: to}
rollup, err := repo.SessionRollup(ctx, ownerScope)
require.NoError(t, err)
assert.Equal(t, int64(2), rollup.Total)
assert.Equal(t, int64(1), rollup.Succeeded) // exited
assert.Equal(t, int64(1), rollup.Crashed)
assert.Equal(t, int64(0), rollup.Active)
assert.InDelta(t, 1.50, rollup.TotalCostUSD, 1e-9)
assert.True(t, rollup.AvgDurationSeconds >= 0)
// admin scope(UserID nil):跨用户 3 个会话。
adminRollup, err := repo.SessionRollup(ctx, acp.DashboardFilter{From: from, To: to})
require.NoError(t, err)
assert.Equal(t, int64(3), adminRollup.Total)
// SessionsByDay:owner scope 当天 1 行,total=2。
byDay, err := repo.SessionsByDay(ctx, ownerScope)
require.NoError(t, err)
require.Len(t, byDay, 1)
assert.Equal(t, int64(2), byDay[0].Total)
assert.Equal(t, int64(1), byDay[0].Succeeded)
assert.InDelta(t, 1.50, byDay[0].TotalCostUSD, 1e-9)
// SessionTimeline:插入两类事件,按桶聚合。
method := "session/update"
for i := 0; i < 3; i++ {
_, err := repo.InsertEvent(ctx, &acp.Event{
SessionID: exited.ID, Direction: acp.DirectionOut, RPCKind: acp.RPCNotification,
Method: &method, Payload: []byte(`{}`), PayloadSize: 2,
})
require.NoError(t, err)
}
reqMethod := "session/prompt"
_, err = repo.InsertEvent(ctx, &acp.Event{
SessionID: exited.ID, Direction: acp.DirectionOut, RPCKind: acp.RPCRequest,
Method: &reqMethod, Payload: []byte(`{}`), PayloadSize: 2,
})
require.NoError(t, err)
timeline, err := repo.SessionTimeline(ctx, exited.ID)
require.NoError(t, err)
require.Len(t, timeline, 2) // (out,notification,session/update) + (out,request,session/prompt)
total := int64(0)
for _, b := range timeline {
total += b.EventCount
assert.False(t, b.FirstAt.IsZero())
assert.False(t, b.LastAt.IsZero())
}
assert.Equal(t, int64(4), total)
}
+115
View File
@@ -0,0 +1,115 @@
// Package sandbox provides a pluggable per-session execution sandbox that sits
// between acp.session_service.Create and acp.Supervisor.Spawn. It confines an
// agent child process to a low-privilege UID, bind-mounts only that session's
// worktree + the per-user agent home, applies OS resource limits, and routes
// egress through an allowlist proxy — all WITHOUT touching proc.Group's
// Setpgid-based tree-kill (Apply only ADDS to cmd.SysProcAttr, never replaces).
//
// Three modes, config-gated via acp.sandbox.mode:
// - none : no-op (default; Windows dev + CI + any non-Linux box)
// - uid : drop to a low-priv UID + setrlimit (+ optional namespace/bind
// mounts) — Linux only
// - container : run the agent inside a rootless runc/bwrap or docker container
// on an egress-restricted network — Linux only
//
// The Linux-specific mechanics live in build-tagged files; this file holds the
// cross-platform contract and the factory so app.go and supervisor.go compile
// on every platform.
package sandbox
import "os/exec"
// Mode selects the sandbox strategy.
type Mode string
const (
// ModeNone disables sandboxing (no-op Apply). Default everywhere except
// production Linux.
ModeNone Mode = "none"
// ModeUID drops to a low-priv UID + setrlimit (+ optional bind mounts).
ModeUID Mode = "uid"
// ModeContainer runs the child inside a rootless container.
ModeContainer Mode = "container"
)
// Limits captures the OS resource limits applied to the child (and its tree).
// A zero field means "do not set this limit" (inherit the parent's).
type Limits struct {
AddressSpaceBytes uint64 // RLIMIT_AS — caps total virtual memory (anti big-alloc)
NProc uint64 // RLIMIT_NPROC — caps process/thread count (anti fork-bomb)
CPUSeconds uint64 // RLIMIT_CPU — caps CPU seconds
PIDs uint64 // cgroup pids.max (container mode); mirrors NProc for uid
DiskBytes uint64 // RLIMIT_FSIZE — caps max file size the child can write
}
// Spec is the per-spawn sandbox request, derived from the session + agent kind.
type Spec struct {
Mode Mode
UID int // target low-priv uid (uid mode); base+offset computed by caller
GID int // target low-priv gid
HomeDir string // per-user agent home (writable, bind-mounted in)
WorktreeDir string // this session's worktree (writable, bind-mounted in)
DataRoot string // /data root to mask everything else under
Rlimits Limits
ProxyURL string // HTTP(S)_PROXY value injected into the child env (egress allowlist)
NoProxy string // NO_PROXY value (e.g. localhost,127.0.0.1)
}
// Sandbox applies a Spec to an exec.Cmd before it is started.
type Sandbox interface {
// Apply mutates cmd.SysProcAttr (Credential, namespace/chroot flags) and
// cmd.Env (HTTP(S)_PROXY/NO_PROXY) BEFORE proc.Group.Prepare and cmd.Start.
// It returns a cleanup closure (unmount binds, remove ephemeral dirs, stop
// container) that the supervisor calls from monitorProcess AFTER
// group.Close(). cleanup is always non-nil (a no-op when there is nothing to
// clean) and must be safe to call exactly once. Apply MUST NOT set Setpgid —
// that is proc.Group's job; on Linux it ADDS to the existing SysProcAttr.
Apply(cmd *exec.Cmd, spec Spec) (cleanup func(), err error)
}
// New returns the Sandbox implementation for the given mode. On non-Linux
// platforms (or mode=none) it always returns the no-op sandbox. The actual
// uid/container constructors are provided by build-tagged files; newPlatform
// returns nil when the mode is unsupported on this OS, in which case we fall
// back to no-op so dev/CI never break.
func New(mode Mode) Sandbox {
switch mode {
case ModeUID, ModeContainer:
if sb := newPlatform(mode); sb != nil {
return sb
}
// Unsupported on this platform → safe no-op fallback.
return noopSandbox{}
default:
return noopSandbox{}
}
}
// noopSandbox makes no changes and returns a no-op cleanup. It is the default
// and the fallback for unsupported modes/platforms.
type noopSandbox struct{}
func (noopSandbox) Apply(_ *exec.Cmd, _ Spec) (func(), error) {
return func() {}, nil
}
// injectProxyEnv appends egress-proxy env vars to cmd.Env when a proxy URL is
// configured. Shared by uid + container modes (both honor HTTP(S)_PROXY). Kept
// here so the cross-platform contract is testable without build tags.
func injectProxyEnv(cmd *exec.Cmd, spec Spec) {
if spec.ProxyURL == "" {
return
}
noProxy := spec.NoProxy
if noProxy == "" {
noProxy = "localhost,127.0.0.1,::1"
}
cmd.Env = append(cmd.Env,
"HTTP_PROXY="+spec.ProxyURL,
"HTTPS_PROXY="+spec.ProxyURL,
"http_proxy="+spec.ProxyURL,
"https_proxy="+spec.ProxyURL,
"NO_PROXY="+noProxy,
"no_proxy="+noProxy,
)
}
@@ -0,0 +1,118 @@
//go:build linux
package sandbox
import (
"fmt"
"os/exec"
"strconv"
"time"
"github.com/google/uuid"
)
// containerSandbox runs the agent binary inside a rootless container with an
// egress-restricted network and the two writable paths bind-mounted in. This is
// the hard isolation boundary (UID + filesystem mask + network policy + cgroup
// limits) — unlike uid mode, a non-cooperating binary cannot bypass the egress
// allowlist because the container network only routes through the proxy/network
// policy.
//
// It rewrites cmd to `docker run --rm --name <n> --user uid:gid --read-only
// --network <egress-net> --pids-limit --memory --cpus --mount <home> --mount
// <worktree> -w <worktree> <image> <binary> <args>` where the image is the same
// runtime image (so the agent CLI is present). The cleanup closure runs
// `docker stop <n>` so proc.Group's tree-kill (which signals the `docker run`
// client's pgid) is reinforced by stopping the actual container — otherwise
// killing the client could orphan the container.
//
// Tree-kill invariant: Apply does NOT set Setpgid (proc.Group owns it). docker
// run becomes the group leader; killing -pgid signals it, and cleanup() issues
// docker stop for the container it spawned. The mapping client->container is via
// the deterministic --name we assign.
//
// The container runtime + egress network are provisioned by the deployment
// (Dockerfile installs the runtime; docker-compose defines the egress-net + the
// forward-proxy sidecar). Image/network/runtime are read from the spec's
// DataRoot convention or sensible defaults; here we keep them parameterized via
// package-level defaults so app.go can override without a code change later.
type containerSandbox struct{}
// container runtime defaults; overridable by deployment via env in a later pass.
const (
defaultRuntime = "docker"
defaultImage = "agent-coding-workflow-sandbox:latest"
defaultEgressNet = "egress-net"
containerStopWait = 5 * time.Second
)
func (s *containerSandbox) Apply(cmd *exec.Cmd, spec Spec) (func(), error) {
if spec.WorktreeDir == "" {
return func() {}, fmt.Errorf("sandbox container: worktree dir required")
}
runtime, err := exec.LookPath(defaultRuntime)
if err != nil {
// Runtime missing: fall back to no-op cleanup but surface the error so
// the supervisor can decide (it currently logs and proceeds unsandboxed
// only if configured to; otherwise spawn fails).
return func() {}, fmt.Errorf("sandbox container: %s not found: %w", defaultRuntime, err)
}
name := "acw-sbx-" + uuid.NewString()
orig := append([]string{cmd.Path}, cmd.Args[1:]...)
runArgs := []string{
"run", "--rm", "--name", name,
"--read-only",
"--network", defaultEgressNet,
"-w", spec.WorktreeDir,
"--mount", "type=bind,source=" + spec.WorktreeDir + ",target=" + spec.WorktreeDir,
}
if spec.HomeDir != "" {
runArgs = append(runArgs,
"--mount", "type=bind,source="+spec.HomeDir+",target="+spec.HomeDir,
"-e", "HOME="+spec.HomeDir)
}
if spec.UID > 0 {
gid := spec.GID
if gid <= 0 {
gid = spec.UID
}
runArgs = append(runArgs, "--user", strconv.Itoa(spec.UID)+":"+strconv.Itoa(gid))
}
if spec.Rlimits.AddressSpaceBytes > 0 {
runArgs = append(runArgs, "--memory", strconv.FormatUint(spec.Rlimits.AddressSpaceBytes, 10))
}
if spec.Rlimits.PIDs > 0 {
runArgs = append(runArgs, "--pids-limit", strconv.FormatUint(spec.Rlimits.PIDs, 10))
} else if spec.Rlimits.NProc > 0 {
runArgs = append(runArgs, "--pids-limit", strconv.FormatUint(spec.Rlimits.NProc, 10))
}
if spec.Rlimits.CPUSeconds > 0 {
// docker has no direct CPU-seconds cap; approximate with --cpus=1 and
// rely on the wall-clock/budget reaper for hard CPU-time bounds.
runArgs = append(runArgs, "--cpus", "1")
}
// Egress proxy env (the proxy itself enforces the allowlist on egress-net).
if spec.ProxyURL != "" {
noProxy := spec.NoProxy
if noProxy == "" {
noProxy = "localhost,127.0.0.1,::1"
}
runArgs = append(runArgs,
"-e", "HTTP_PROXY="+spec.ProxyURL,
"-e", "HTTPS_PROXY="+spec.ProxyURL,
"-e", "NO_PROXY="+noProxy)
}
runArgs = append(runArgs, defaultImage)
runArgs = append(runArgs, orig...)
cmd.Path = runtime
cmd.Args = append([]string{runtime}, runArgs...)
cleanup := func() {
stop := exec.Command(runtime, "stop", "-t", strconv.Itoa(int(containerStopWait.Seconds())), name)
_ = stop.Run() // best-effort; --rm removes it after stop
}
return cleanup, nil
}
+15
View File
@@ -0,0 +1,15 @@
//go:build linux
package sandbox
// newPlatform dispatches the Linux sandbox modes to their concrete impls.
func newPlatform(mode Mode) Sandbox {
switch mode {
case ModeUID:
return &uidSandbox{}
case ModeContainer:
return &containerSandbox{}
default:
return nil
}
}
+9
View File
@@ -0,0 +1,9 @@
//go:build !linux
package sandbox
// newPlatform returns nil on non-Linux platforms: UID drop, setrlimit and
// bind-mount namespaces are Linux-only, so Windows dev and macOS always fall
// back to the no-op sandbox (see New). This is the documented contract that
// mode=none is the only supported mode off Linux.
func newPlatform(_ Mode) Sandbox { return nil }
+59
View File
@@ -0,0 +1,59 @@
package sandbox
import (
"os/exec"
"runtime"
"testing"
"github.com/stretchr/testify/require"
)
func TestNew_NoneAlwaysNoop(t *testing.T) {
sb := New(ModeNone)
require.IsType(t, noopSandbox{}, sb)
}
func TestNew_UnsupportedPlatformFallsBackToNoop(t *testing.T) {
if runtime.GOOS == "linux" {
t.Skip("uid/container are supported on linux; this asserts the non-linux fallback")
}
// On non-linux, uid/container must fall back to no-op so dev/CI never break.
require.IsType(t, noopSandbox{}, New(ModeUID))
require.IsType(t, noopSandbox{}, New(ModeContainer))
}
func TestNoopSandbox_MutatesNothing(t *testing.T) {
cmd := exec.Command("echo", "hi")
before := cmd.SysProcAttr
cleanup, err := noopSandbox{}.Apply(cmd, Spec{Mode: ModeNone, ProxyURL: "http://proxy:8080"})
require.NoError(t, err)
require.NotNil(t, cleanup)
// no-op must NOT inject proxy env nor touch SysProcAttr.
require.Equal(t, before, cmd.SysProcAttr)
require.Empty(t, cmd.Env)
// cleanup is safe to call.
cleanup()
}
func TestInjectProxyEnv(t *testing.T) {
cmd := exec.Command("echo", "hi")
cmd.Env = []string{"PATH=/usr/bin"}
injectProxyEnv(cmd, Spec{ProxyURL: "http://proxy:3128"})
require.Contains(t, cmd.Env, "HTTP_PROXY=http://proxy:3128")
require.Contains(t, cmd.Env, "HTTPS_PROXY=http://proxy:3128")
require.Contains(t, cmd.Env, "NO_PROXY=localhost,127.0.0.1,::1")
// original env preserved.
require.Contains(t, cmd.Env, "PATH=/usr/bin")
}
func TestInjectProxyEnv_NoProxyURLIsNoop(t *testing.T) {
cmd := exec.Command("echo", "hi")
injectProxyEnv(cmd, Spec{})
require.Empty(t, cmd.Env)
}
func TestInjectProxyEnv_CustomNoProxy(t *testing.T) {
cmd := exec.Command("echo", "hi")
injectProxyEnv(cmd, Spec{ProxyURL: "http://p", NoProxy: "internal.local"})
require.Contains(t, cmd.Env, "NO_PROXY=internal.local")
}
+92
View File
@@ -0,0 +1,92 @@
//go:build linux
package sandbox
import (
"fmt"
"os/exec"
"strconv"
"syscall"
)
// uidSandbox confines the agent child to a low-privilege UID/GID and applies OS
// resource limits, while leaving proc.Group's Setpgid (tree-kill) intact.
//
// Tree-kill invariant: we ONLY add Credential to the existing SysProcAttr. We do
// NOT touch Setpgid (proc.Group sets it after Apply via group.Prepare). The
// negative-pgid SIGTERM/SIGKILL path keeps working because the new process group
// leader is the same child — it simply runs as the low-priv uid now.
//
// Resource limits are applied by wrapping the real binary in `prlimit(1)`:
// `prlimit --as=N --nproc=N --cpu=N --fsize=N -- <binary> <args...>`. prlimit
// sets the limits on the child it execs, so they apply to the agent and (since
// limits are inherited across fork) its whole subtree — exactly the tree the
// pgid covers. Wrapping is additive to Credential/Setpgid and does not change
// the process-group topology (prlimit execs the target in place, becoming the
// group leader). If prlimit is unavailable, UID drop still applies and we log no
// rlimits (best-effort; container mode is the hard boundary).
//
// Bind-mount masking (CLONE_NEWNS + remount /data with only HomeDir+WorktreeDir
// visible) is intentionally NOT auto-enabled here: it requires the server to
// hold CAP_SYS_ADMIN / a user namespace, which the rootless production container
// (USER app) does not have by default. The robust isolation path is
// ModeContainer. uidSandbox therefore provides UID drop + rlimits + egress proxy
// as the minimum viable confinement, and documents that filesystem masking is a
// container-mode guarantee.
type uidSandbox struct{}
func (s *uidSandbox) Apply(cmd *exec.Cmd, spec Spec) (func(), error) {
if spec.UID <= 0 {
return func() {}, fmt.Errorf("sandbox uid: invalid target uid %d", spec.UID)
}
// UID/GID drop — ADDITIVE to whatever proc.Group will set (Setpgid).
if cmd.SysProcAttr == nil {
cmd.SysProcAttr = &syscall.SysProcAttr{}
}
gid := spec.GID
if gid <= 0 {
gid = spec.UID
}
cmd.SysProcAttr.Credential = &syscall.Credential{
Uid: uint32(spec.UID),
Gid: uint32(gid),
}
// Resource limits via prlimit wrapper (if available).
if prlimit, err := exec.LookPath("prlimit"); err == nil {
args := prlimitArgs(spec.Rlimits)
if len(args) > 0 {
// Re-point the command through prlimit, preserving the original
// binary + args after the `--` separator.
orig := append([]string{cmd.Path}, cmd.Args[1:]...)
cmd.Path = prlimit
cmd.Args = append(append([]string{prlimit}, args...), append([]string{"--"}, orig...)...)
}
}
// Egress allowlist: inject HTTP(S)_PROXY/NO_PROXY into the stripped child env.
injectProxyEnv(cmd, spec)
// No ephemeral resources to tear down in uid mode (no bind mounts here);
// return a no-op cleanup.
return func() {}, nil
}
// prlimitArgs builds the prlimit flags for the non-zero limits in l.
func prlimitArgs(l Limits) []string {
var args []string
if l.AddressSpaceBytes > 0 {
args = append(args, "--as="+strconv.FormatUint(l.AddressSpaceBytes, 10))
}
if l.NProc > 0 {
args = append(args, "--nproc="+strconv.FormatUint(l.NProc, 10))
}
if l.CPUSeconds > 0 {
args = append(args, "--cpu="+strconv.FormatUint(l.CPUSeconds, 10))
}
if l.DiskBytes > 0 {
args = append(args, "--fsize="+strconv.FormatUint(l.DiskBytes, 10))
}
return args
}
+148
View File
@@ -0,0 +1,148 @@
package acp
import (
"context"
"sync"
"testing"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/audit"
"github.com/yan1h/agent-coding-workflow/internal/brief"
"github.com/yan1h/agent-coding-workflow/internal/project"
)
// recordingAudit 捕获 Record 调用,供断言 brief_composed 是否写入。
type recordingAudit struct {
mu sync.Mutex
entries []audit.Entry
}
func (r *recordingAudit) Record(_ context.Context, e audit.Entry) error {
r.mu.Lock()
defer r.mu.Unlock()
r.entries = append(r.entries, e)
return nil
}
func (r *recordingAudit) actions() []string {
r.mu.Lock()
defer r.mu.Unlock()
out := make([]string, 0, len(r.entries))
for _, e := range r.entries {
out = append(out, e.Action)
}
return out
}
// fakeComposer 是 brief.Composer 的测试替身:返回固定文本并记录是否被调用。
type fakeComposer struct {
res brief.BriefResult
err error
called bool
gotIn brief.BriefInput
}
func (f *fakeComposer) ComposeTaskBrief(_ context.Context, in brief.BriefInput) (brief.BriefResult, error) {
f.called = true
f.gotIn = in
return f.res, f.err
}
// fakeArtifactAccess 是 ArtifactAccess 的测试替身。
type fakeArtifactAccess struct {
arts []*project.Artifact
}
func (f fakeArtifactAccess) ListLatestArtifacts(_ context.Context, _ uuid.UUID) ([]*project.Artifact, error) {
return f.arts, nil
}
func testSession() *Session {
return &Session{
ID: uuid.New(),
ProjectID: uuid.New(),
Branch: "req/demo-7",
CwdPath: "/tmp/wt",
}
}
func TestComposeInitialPrompt_RequirementContextComposed(t *testing.T) {
t.Parallel()
rec := &recordingAudit{}
fc := &fakeComposer{res: brief.BriefResult{
Text: "你是一名资深产品规划助手。\n\n## 当前需求\n需求 #7:导出报表",
Tokens: 20,
Sections: []string{"role", "requirement_core"},
}}
s := &sessionService{
audit: rec,
composer: fc,
aa: fakeArtifactAccess{arts: []*project.Artifact{{Phase: project.PhasePlanning, Version: 2}}},
cfg: SessionServiceConfig{BriefTokenBudget: 6000},
}
req := &project.Requirement{ID: uuid.New(), Number: 7, Title: "导出报表", Phase: project.PhasePlanning}
c := Caller{UserID: uuid.New()}
out := s.composeInitialPrompt(context.Background(), c, testSession(), nil, req, nil, "请开始")
require.True(t, fc.called)
assert.Contains(t, out, "需求 #7:导出报表")
assert.Contains(t, out, "资深产品规划助手")
// 组装入参带上了 requirement + artifacts + cwd + 用户文本 + budget。
assert.Equal(t, req, fc.gotIn.Requirement)
assert.Equal(t, "请开始", fc.gotIn.UserPrompt)
assert.Equal(t, "/tmp/wt", fc.gotIn.RepoCwd)
assert.Equal(t, 6000, fc.gotIn.TokenBudget)
assert.Equal(t, brief.KindACPSession, fc.gotIn.Kind)
assert.Len(t, fc.gotIn.Artifacts, 1)
// 写了 brief_composed audit。
assert.Contains(t, rec.actions(), "acp.session.brief_composed")
}
func TestComposeInitialPrompt_IssueContextComposed(t *testing.T) {
t.Parallel()
rec := &recordingAudit{}
fc := &fakeComposer{res: brief.BriefResult{Text: "## 当前 Issue\nIssue #3:bug", Tokens: 8}}
s := &sessionService{audit: rec, composer: fc}
iss := &project.Issue{ID: uuid.New(), Number: 3, Title: "bug"}
out := s.composeInitialPrompt(context.Background(), Caller{UserID: uuid.New()}, testSession(), nil, nil, iss, "fix it")
require.True(t, fc.called)
assert.Contains(t, out, "Issue #3:bug")
assert.Equal(t, iss, fc.gotIn.Issue)
// issue-only 时不调用 ArtifactAccess(aa 为 nil 也不 panic)。
}
func TestComposeInitialPrompt_NoPMContextFallsBackToRaw(t *testing.T) {
t.Parallel()
fc := &fakeComposer{res: brief.BriefResult{Text: "SHOULD NOT BE USED"}}
s := &sessionService{audit: &recordingAudit{}, composer: fc}
out := s.composeInitialPrompt(context.Background(), Caller{UserID: uuid.New()}, testSession(), nil, nil, nil, "raw text only")
assert.Equal(t, "raw text only", out)
assert.False(t, fc.called, "composer must not be called without PM context")
}
func TestComposeInitialPrompt_NilComposerFallsBackToRaw(t *testing.T) {
t.Parallel()
s := &sessionService{audit: &recordingAudit{}, composer: nil}
req := &project.Requirement{ID: uuid.New(), Number: 7, Title: "x", Phase: project.PhasePlanning}
out := s.composeInitialPrompt(context.Background(), Caller{UserID: uuid.New()}, testSession(), nil, req, nil, "raw")
assert.Equal(t, "raw", out)
}
func TestComposeInitialPrompt_EmptyComposedTextFallsBackToRaw(t *testing.T) {
t.Parallel()
rec := &recordingAudit{}
fc := &fakeComposer{res: brief.BriefResult{Text: ""}}
s := &sessionService{audit: rec, composer: fc}
req := &project.Requirement{ID: uuid.New(), Number: 1, Title: "t", Phase: project.PhasePlanning}
out := s.composeInitialPrompt(context.Background(), Caller{UserID: uuid.New()}, testSession(), nil, req, nil, "fallback raw")
assert.Equal(t, "fallback raw", out)
// 空简报不写 brief_composed audit。
assert.NotContains(t, rec.actions(), "acp.session.brief_composed")
}
+314 -8
View File
@@ -6,6 +6,7 @@ import (
"context"
"encoding/json"
"fmt"
"log/slog"
"strings"
"time"
@@ -13,6 +14,7 @@ import (
"github.com/jackc/pgx/v5"
"github.com/yan1h/agent-coding-workflow/internal/audit"
"github.com/yan1h/agent-coding-workflow/internal/brief"
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
"github.com/yan1h/agent-coding-workflow/internal/mcp"
"github.com/yan1h/agent-coding-workflow/internal/project"
@@ -25,6 +27,7 @@ type SessionServiceConfig struct {
MaxActivePerUser int
SystemTokenTTL time.Duration
MCPPublicURL string
BriefTokenBudget int // 初始 prompt 简报的 token 硬上限(<=0 用 brief 默认值)
}
// ProjectAccess 是 acp 模块对 project 模块的窄接口;adapter 在 app.go 实现。
@@ -34,26 +37,37 @@ type ProjectAccess interface {
GetRequirementByNumber(ctx context.Context, projectID uuid.UUID, number int) (*project.Requirement, error)
}
// ArtifactAccess 是 acp 模块拉取某 requirement 各阶段最新产物的窄接口;adapter 在
// app.go 实现(基于 project.Repository.ListArtifactsByRequirement,按阶段取最大版本)。
type ArtifactAccess interface {
ListLatestArtifacts(ctx context.Context, requirementID uuid.UUID) ([]*project.Artifact, error)
}
type sessionService struct {
repo Repository
wsSvc workspace.WorkspaceService
wtSvc workspace.WorktreeService
wsRepo workspace.Repository
pa ProjectAccess
aa ArtifactAccess // 可为 nil(无产物注入)
composer brief.Composer // 可为 nil(退回 raw passthrough)
sup *Supervisor
audit audit.Recorder
cfg SessionServiceConfig
mcpTokens mcp.TokenService // spec §6.1 — system token issuance
}
// NewSessionService 构造 SessionService。
// NewSessionService 构造 SessionService。aa / composer 可为 nil:缺任一时
// initial prompt 退回原始透传(不组装简报),保持旧行为。
func NewSessionService(repo Repository, wsSvc workspace.WorkspaceService,
wtSvc workspace.WorktreeService, wsRepo workspace.Repository,
pa ProjectAccess, sup *Supervisor, rec audit.Recorder,
pa ProjectAccess, aa ArtifactAccess, composer brief.Composer,
sup *Supervisor, rec audit.Recorder,
cfg SessionServiceConfig, mcpTokens mcp.TokenService) SessionService {
return &sessionService{
repo: repo, wsSvc: wsSvc, wtSvc: wtSvc, wsRepo: wsRepo,
pa: pa, sup: sup, audit: rec, cfg: cfg, mcpTokens: mcpTokens,
pa: pa, aa: aa, composer: composer, sup: sup, audit: rec,
cfg: cfg, mcpTokens: mcpTokens,
}
}
@@ -199,28 +213,41 @@ func (s *sessionService) Create(ctx context.Context, c Caller, in CreateSessionI
worktreeID = &wid
}
// 6. PM 关联(best-effort:拿不到 ID 不阻塞 session 创建)
var issueID, reqID *uuid.UUID
// 6. PM 关联(best-effort:拿不到 ID 不阻塞 session 创建)
// 同时保留解析出的 issue/requirement 对象,供初始 prompt 组装简报复用。
var (
issueID, reqID *uuid.UUID
pmIssue *project.Issue
pmReq *project.Requirement
)
if in.IssueNumber != nil {
if iss, _ := s.pa.GetIssueByNumber(ctx, ws.ProjectID, *in.IssueNumber); iss != nil {
id := iss.ID
issueID = &id
pmIssue = iss
}
}
if in.RequirementNumber != nil {
if req, _ := s.pa.GetRequirementByNumber(ctx, ws.ProjectID, *in.RequirementNumber); req != nil {
id := req.ID
reqID = &id
pmReq = req
}
}
// 7. INSERT acp_sessions + system MCP token(spec §6.1 atomic guarantee)
// 预算快照:当前无 session 级覆盖入口,故直接继承 agent-kind 默认上限。快照到
// session 行使 supervisor / reaper 读取自洽,不必每次回查 agent-kind。
sess := &Session{
ID: sessionID, WorkspaceID: ws.ID, ProjectID: ws.ProjectID,
AgentKindID: kind.ID, UserID: c.UserID,
IssueID: issueID, RequirementID: reqID,
Branch: branch, CwdPath: cwdPath, IsMainWorktree: isMain,
Status: SessionStarting,
Status: SessionStarting,
OrchestratorStepID: in.OrchestratorStepID,
BudgetMaxCostUSD: kind.MaxCostUSD,
BudgetMaxTokens: kind.MaxTokens,
BudgetMaxWallClockSeconds: kind.MaxWallClockSeconds,
}
var (
out *Session
@@ -279,14 +306,91 @@ func (s *sessionService) Create(ctx context.Context, c Caller, in CreateSessionI
}
}()
// 10. initial prompt:握手成功后异步转发
// 10. initial prompt:握手成功后异步转发。先在请求路径同步组装服务端简报
// (需求/Issue + 阶段产物 + 仓库定向 + 操作者自由文本),失败/无 PM 上下文时
// 退回原始透传,保持旧行为。
if in.InitialPrompt != nil && *in.InitialPrompt != "" {
go s.sendInitialPrompt(out.ID, *in.InitialPrompt)
text := s.composeInitialPrompt(ctx, c, out, pj, pmReq, pmIssue, *in.InitialPrompt)
go s.sendInitialPrompt(out.ID, text)
}
return out, nil
}
// composeInitialPrompt 组装 ACP session 的初始 prompt 简报。无 composer 或无 PM 上下文
// (既无 requirement 也无 issue)时返回原始 userPrompt(raw passthrough)。组装失败也
// 静默退回原始文本,绝不阻塞 session。成功组装后写一条 audit。
func (s *sessionService) composeInitialPrompt(ctx context.Context, c Caller, sess *Session,
pj *project.Project, req *project.Requirement, iss *project.Issue, userPrompt string) string {
if s.composer == nil || (req == nil && iss == nil) {
return userPrompt
}
var artifacts []*project.Artifact
if req != nil && s.aa != nil {
if arts, err := s.aa.ListLatestArtifacts(ctx, req.ID); err == nil {
artifacts = arts
}
}
res, err := s.composer.ComposeTaskBrief(ctx, brief.BriefInput{
Project: pj,
Requirement: req,
Issue: iss,
Artifacts: artifacts,
RepoCwd: sess.CwdPath,
Branch: sess.Branch,
UserPrompt: userPrompt,
TokenBudget: s.cfg.BriefTokenBudget,
Kind: brief.KindACPSession,
})
if err != nil || res.Text == "" {
return userPrompt
}
s.recordAudit(ctx, c, "acp.session.brief_composed", sess.ID.String(), map[string]any{
"requirement_id": idOrNil(req),
"issue_id": issueIDOrNil(iss),
"artifact_versions": artifactVersions(artifacts),
"brief_tokens": res.Tokens,
"truncated": res.Truncated,
"sections": res.Sections,
})
return res.Text
}
func idOrNil(req *project.Requirement) *uuid.UUID {
if req == nil {
return nil
}
id := req.ID
return &id
}
func issueIDOrNil(iss *project.Issue) *uuid.UUID {
if iss == nil {
return nil
}
id := iss.ID
return &id
}
// artifactVersions 把产物列表归约为 {phase: version} 映射(每阶段取出现的最大版本),
// 仅用于 audit 元数据。
func artifactVersions(arts []*project.Artifact) map[string]int {
out := make(map[string]int, len(arts))
for _, a := range arts {
if a == nil {
continue
}
if v, ok := out[string(a.Phase)]; !ok || a.Version > v {
out[string(a.Phase)] = a.Version
}
}
return out
}
// findOrCreateWorktree 在指定 workspace 内查找 branch 同名 worktree,
// 不存在则创建。caller 已通过 wsSvc.Get 完成鉴权。
func (s *sessionService) findOrCreateWorktree(ctx context.Context, c workspace.Caller, wsID uuid.UUID, branch string) (*workspace.Worktree, error) {
@@ -342,11 +446,167 @@ func (s *sessionService) sendInitialPrompt(sid uuid.UUID, text string) {
msg := &Message{
JSONRPC: "2.0", ID: idJSON, Method: "session/prompt", Params: params,
}
// 回合相关联:在发送前登记回合 id("client-init"),使响应的 stopReason 可被
// 相关联到这个回合。tracker 可能为 nil(部分测试)——StartTurn 已做 nil 守卫。
if tr := relay.Tracker(); tr != nil {
if _, err := tr.StartTurn(context.Background(), "client-init"); err != nil {
slog.Warn("acp.session.start_turn", "session_id", sid, "err", err.Error())
}
}
_ = relay.SendClient(msg)
return
}
}
// 编译期断言:sessionService 同时实现 SessionService 与 TurnRunner(编排器依赖后者)。
var _ TurnRunner = (*sessionService)(nil)
// promptResult 是 session/prompt 响应 Result 中我们关心的字段。
type promptResult struct {
StopReason string `json:"stopReason"`
}
// SendPromptAndWait 发送一条 session/prompt 并阻塞等待回合完成,返回 agent 上报的原始
// stopReason。这是编排器驱动一个阻塞回合的核心原语:先轮询等待 supervisor handshake
// 完成(relay + agent_session_id 就绪,复用 sendInitialPrompt 的就绪检测),随后走
// server-initiated relay.Call('session/prompt') 取响应里的 stopReason。
//
// ctx 控制整体超时(编排器以 cfg.TurnTimeout 约束,须 < jobReaper LeaseTimeout)。
func (s *sessionService) SendPromptAndWait(ctx context.Context, sessionID uuid.UUID, prompt string) (string, error) {
relay, agentSessionID, err := s.waitForRelayReady(ctx, sessionID)
if err != nil {
return "", err
}
params := map[string]any{
"sessionId": agentSessionID,
"prompt": []map[string]any{{"type": "text", "text": prompt}},
}
resp, err := relay.Call(ctx, "session/prompt", params)
if err != nil {
return "", errs.Wrap(err, errs.CodeInternal, "session/prompt call")
}
var pr promptResult
if len(resp.Result) > 0 {
if uerr := json.Unmarshal(resp.Result, &pr); uerr != nil {
return "", errs.Wrap(uerr, errs.CodeInternal, "parse session/prompt result")
}
}
if pr.StopReason != "" {
// 落库最近一次 stopReason(best-effort),供观测/网关复用。
_ = s.repo.UpdateSessionLastStopReason(context.WithoutCancel(ctx), sessionID, pr.StopReason)
}
return pr.StopReason, nil
}
// waitForRelayReady 轮询直到 relay 与 agent_session_id 同时就绪或 ctx 截止。
// 200ms 一次,与 sendInitialPrompt 一致。
func (s *sessionService) waitForRelayReady(ctx context.Context, sessionID uuid.UUID) (*Relay, string, error) {
tick := time.NewTicker(200 * time.Millisecond)
defer tick.Stop()
for {
relay := s.sup.GetRelay(sessionID)
if relay != nil {
sess, _ := s.repo.GetSessionByID(ctx, sessionID)
if sess != nil && sess.AgentSessionID != nil {
return relay, *sess.AgentSessionID, nil
}
}
select {
case <-ctx.Done():
return nil, "", ctx.Err()
case <-tick.C:
}
}
}
// ResumeSession 在崩溃/退出 session 的同一 cwd/worktree 上重新拉起 agent。复用原
// session 的 workspace/branch/cwd/main 标记与编排器 step 关联:重新占用 worktree
// (main → AcquireMainWorktree CAS;子 worktree → 按 holder Acquire),签发新 system
// token,再 sup.Spawn。Spawn 失败时回滚 worktree(releaseOnFailure)。
//
// 返回的是被复用并重置为 starting 状态的同一行 session(同 ID),供编排器随后
// SendPromptAndWait。仅供编排器内部调用(caller 为 run.owner)。
func (s *sessionService) ResumeSession(ctx context.Context, c Caller, sessionID uuid.UUID) (*Session, error) {
sess, err := s.repo.GetSessionByID(ctx, sessionID)
if err != nil {
return nil, err
}
if sess.Status.IsActive() {
// 已经活跃:无需 resume,直接返回。
return sess, nil
}
kind, err := s.repo.GetAgentKindByID(ctx, sess.AgentKindID)
if err != nil {
return nil, err
}
sessCaller := workspace.Caller{UserID: c.UserID, IsAdmin: c.IsAdmin, SessionID: &sess.ID}
if sess.IsMainWorktree {
ok, aerr := s.repo.AcquireMainWorktree(ctx, sess.ID, sess.WorkspaceID)
if aerr != nil {
return nil, aerr
}
if !ok {
return nil, errs.New(errs.CodeAcpSessionBranchConflict, "main worktree in use; cannot resume")
}
} else {
wt, ferr := s.findOrCreateWorktree(ctx, sessCaller, sess.WorkspaceID, sess.Branch)
if ferr != nil {
return nil, ferr
}
if _, aerr := s.wtSvc.Acquire(ctx, sessCaller, wt.ID); aerr != nil {
return nil, aerr
}
}
// 重新签发 system token + 标记 session 回到 starting(清空上次终态字段)。
var (
tokenRes mcp.IssueResult
)
txErr := s.repo.InTx(ctx, func(txCtx context.Context, tx pgx.Tx) error {
if rerr := s.repo.WithTx(tx).ResetSessionForResume(txCtx, sess.ID); rerr != nil {
return rerr
}
res, ierr := s.mcpTokens.IssueSystemTokenWithTx(txCtx, tx, mcp.IssueSystemTokenInput{
UserID: c.UserID,
ACPSessionID: sess.ID,
ProjectIDs: []uuid.UUID{sess.ProjectID},
ExpiresIn: s.cfg.SystemTokenTTL,
})
if ierr != nil {
return ierr
}
tokenRes = res
return nil
})
if txErr != nil {
s.releaseOnFailure(ctx, sess, sessCaller, nil)
return nil, txErr
}
sess.Status = SessionStarting
extraEnv := map[string]string{
EnvMCPToken: tokenRes.Plaintext,
EnvMCPURL: s.cfg.MCPPublicURL,
}
if _, serr := s.sup.Spawn(ctx, sess, kind, extraEnv); serr != nil {
if _, merr := s.repo.MarkSessionFailedIfActive(ctx, sess.ID, serr.Error()); merr != nil {
s.sup.log.Error("acp.resume.mark_session", "session_id", sess.ID, "err", merr.Error())
}
_ = s.mcpTokens.RevokeBySession(ctx, sess.ID)
s.releaseOnFailure(ctx, sess, sessCaller, nil)
return nil, serr
}
s.recordAudit(ctx, c, "acp.session.resume", sess.ID.String(), map[string]any{
"branch": sess.Branch,
"cwd_path": sess.CwdPath,
})
return sess, nil
}
func (s *sessionService) Get(ctx context.Context, c Caller, id uuid.UUID) (*Session, error) {
sess, err := s.repo.GetSessionByID(ctx, id)
if err != nil {
@@ -380,6 +640,52 @@ func (s *sessionService) Terminate(ctx context.Context, c Caller, id uuid.UUID)
return nil
}
// defaultDashboardWindow 是 run-history 汇总在调用方未指定时间窗时的默认回溯。
const defaultDashboardWindow = 30 * 24 * time.Hour
// Dashboard 返回 run-history 汇总 + 按天序列。非 admin 强制按 caller scope。
func (s *sessionService) Dashboard(ctx context.Context, c Caller, in DashboardInput) (*DashboardResult, error) {
to := in.To
if to.IsZero() {
to = time.Now().UTC()
}
from := in.From
if from.IsZero() {
from = to.Add(-defaultDashboardWindow)
}
if !from.Before(to) {
return nil, errs.New(errs.CodeInvalidInput, "from must be before to")
}
f := DashboardFilter{
ProjectID: in.ProjectID,
RequirementID: in.RequirementID,
From: from,
To: to,
}
// 非 admin 只能看自己的会话:UserID 固定为 caller。admin 留空(全量)。
if !c.IsAdmin {
uid := c.UserID
f.UserID = &uid
}
rollup, err := s.repo.SessionRollup(ctx, f)
if err != nil {
return nil, err
}
byDay, err := s.repo.SessionsByDay(ctx, f)
if err != nil {
return nil, err
}
return &DashboardResult{Rollup: rollup, ByDay: byDay, From: from, To: to}, nil
}
// Timeline 返回单会话事件时间线(先经 Get 做 owner/admin 访问校验)。
func (s *sessionService) Timeline(ctx context.Context, c Caller, sessionID uuid.UUID) ([]SessionTimelineBucket, error) {
if _, err := s.Get(ctx, c, sessionID); err != nil {
return nil, err
}
return s.repo.SessionTimeline(ctx, sessionID)
}
func (s *sessionService) recordAudit(ctx context.Context, c Caller, action, targetID string, meta map[string]any) {
if s.audit == nil {
return
+304 -5
View File
@@ -123,6 +123,14 @@ type fakeAcpRepo struct {
// when fn returns nil (simulating real transaction semantics).
txStaging []*acp.Session
inTx bool
// run-history dashboard: 记录入参 filter 供断言;返回预置结果。
rollupFilter acp.DashboardFilter
byDayFilter acp.DashboardFilter
rollupResult acp.SessionRollup
byDayResult []acp.SessionDayRow
timelineResult []acp.SessionTimelineBucket
timelineSID uuid.UUID
}
func (r *fakeAcpRepo) InTx(_ context.Context, fn func(context.Context, pgx.Tx) error) error {
@@ -175,6 +183,43 @@ func (r *fakeAcpRepo) GetSessionByID(_ context.Context, id uuid.UUID) (*acp.Sess
}
return nil, fmt.Errorf("not found")
}
func (r *fakeAcpRepo) GetSessionByStepID(_ context.Context, stepID uuid.UUID) (*acp.Session, error) {
r.mu.Lock()
defer r.mu.Unlock()
for _, s := range r.sessions {
if s.OrchestratorStepID != nil && *s.OrchestratorStepID == stepID {
return s, nil
}
}
return nil, fmt.Errorf("not found")
}
func (r *fakeAcpRepo) GetCrashedSessionForResume(_ context.Context, stepID uuid.UUID) (*acp.Session, error) {
r.mu.Lock()
defer r.mu.Unlock()
for _, s := range r.sessions {
if s.OrchestratorStepID != nil && *s.OrchestratorStepID == stepID &&
(s.Status == acp.SessionCrashed || s.Status == acp.SessionExited) {
return s, nil
}
}
return nil, fmt.Errorf("not found")
}
func (r *fakeAcpRepo) ResetSessionForResume(_ context.Context, id uuid.UUID) error {
r.mu.Lock()
defer r.mu.Unlock()
for _, s := range r.sessions {
if s.ID == id {
s.Status = acp.SessionStarting
s.AgentSessionID = nil
s.PID = nil
s.ExitCode = nil
s.LastError = nil
s.EndedAt = nil
return nil
}
}
return fmt.Errorf("not found")
}
func (r *fakeAcpRepo) GetAgentKindByID(_ context.Context, id uuid.UUID) (*acp.AgentKind, error) {
for _, k := range r.kinds {
if k.ID == id {
@@ -252,6 +297,9 @@ func (r *fakeAcpRepo) ListAllSessions(_ context.Context, reqID *uuid.UUID) ([]*a
func (r *fakeAcpRepo) UpdateSessionRunning(_ context.Context, _ uuid.UUID, _ string, _ int32) error {
return nil
}
func (r *fakeAcpRepo) UpdateSessionSandboxMode(_ context.Context, _ uuid.UUID, _ string) error {
return nil
}
func (r *fakeAcpRepo) UpdateSessionFinished(_ context.Context, _ uuid.UUID, _ acp.SessionStatus, _ *int32, _ *string) error {
return nil
}
@@ -276,6 +324,69 @@ func (r *fakeAcpRepo) ListEventsSince(_ context.Context, _ uuid.UUID, _ int64, _
return nil, nil
}
func (r *fakeAcpRepo) PurgeEventsBefore(_ context.Context, _ time.Time) (int64, error) { return 0, nil }
func (r *fakeAcpRepo) InsertSessionUsage(_ context.Context, _ *acp.SessionUsageRecord) (bool, error) {
return true, nil
}
func (r *fakeAcpRepo) AddSessionUsageTotals(_ context.Context, _ uuid.UUID, dp, dc, dt int64, dCost float64) (int64, int64, int64, float64, error) {
return dp, dc, dt, dCost, nil
}
func (r *fakeAcpRepo) SumProjectCostUSD(_ context.Context, _ uuid.UUID, _ time.Time) (float64, error) {
return 0, nil
}
func (r *fakeAcpRepo) MarkSessionTerminatedReason(_ context.Context, _ uuid.UUID, _ string) error {
return nil
}
func (r *fakeAcpRepo) ListSessionsForReaper(_ context.Context, _, _ time.Time) ([]acp.ReaperSession, error) {
return nil, nil
}
func (r *fakeAcpRepo) ListSessionUsage(_ context.Context, _ uuid.UUID, _ int32) ([]*acp.SessionUsageRecord, error) {
return nil, nil
}
func (r *fakeAcpRepo) SummarizeAcpUsageByUser(_ context.Context, _, _ time.Time) ([]acp.AcpUsageUserRow, error) {
return nil, nil
}
func (r *fakeAcpRepo) SummarizeAcpUsageByModel(_ context.Context, _, _ time.Time) ([]acp.AcpUsageModelRow, error) {
return nil, nil
}
func (r *fakeAcpRepo) SummarizeAcpUsageByDay(_ context.Context, _, _ time.Time) ([]acp.AcpUsageDayRow, error) {
return nil, nil
}
func (r *fakeAcpRepo) SessionRollup(_ context.Context, f acp.DashboardFilter) (acp.SessionRollup, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.rollupFilter = f
return r.rollupResult, nil
}
func (r *fakeAcpRepo) SessionsByDay(_ context.Context, f acp.DashboardFilter) ([]acp.SessionDayRow, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.byDayFilter = f
return r.byDayResult, nil
}
func (r *fakeAcpRepo) SessionTimeline(_ context.Context, sid uuid.UUID) ([]acp.SessionTimelineBucket, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.timelineSID = sid
return r.timelineResult, nil
}
func (r *fakeAcpRepo) InsertTurn(_ context.Context, t *acp.Turn) (*acp.Turn, error) { return t, nil }
func (r *fakeAcpRepo) MarkTurnCompleted(_ context.Context, _ uuid.UUID, _ int, _ string, _ int) error {
return nil
}
func (r *fakeAcpRepo) MarkTurnAborted(_ context.Context, _ uuid.UUID, _ int) error { return nil }
func (r *fakeAcpRepo) IncrementTurnUpdateCount(_ context.Context, _ uuid.UUID, _ int) error {
return nil
}
func (r *fakeAcpRepo) GetLatestTurnBySession(_ context.Context, _ uuid.UUID) (*acp.Turn, error) {
return nil, nil
}
func (r *fakeAcpRepo) ListTurnsBySession(_ context.Context, _ uuid.UUID) ([]*acp.Turn, error) {
return nil, nil
}
func (r *fakeAcpRepo) UpdateSessionLastStopReason(_ context.Context, _ uuid.UUID, _ string) error {
return nil
}
func (r *fakeAcpRepo) PurgeTurnsBefore(_ context.Context, _ time.Time) (int64, error) { return 0, nil }
// fakeMCPTokenSvc satisfies mcp.TokenService for unit tests.
type fakeMCPTokenSvc struct {
@@ -382,7 +493,7 @@ func TestSessionService_Create_IssuesSystemToken(t *testing.T) {
)
svc := acp.NewSessionService(
repo, wsSvc, nil, nil, pa, sup, nil,
repo, wsSvc, nil, nil, pa, nil, nil, sup, nil,
acp.SessionServiceConfig{
MaxActiveGlobal: 10,
MaxActivePerUser: 5,
@@ -441,7 +552,7 @@ func TestSessionService_Create_TokenIssueFails_RollsBack(t *testing.T) {
)
svc := acp.NewSessionService(
repo, wsSvc, nil, nil, pa, sup, nil,
repo, wsSvc, nil, nil, pa, nil, nil, sup, nil,
acp.SessionServiceConfig{
MaxActiveGlobal: 10,
MaxActivePerUser: 5,
@@ -495,7 +606,7 @@ func TestSessionService_Create_SpawnFails_MarksSessionCrashed(t *testing.T) {
)
svc := acp.NewSessionService(
repo, wsSvc, nil, nil, pa, sup, nil,
repo, wsSvc, nil, nil, pa, nil, nil, sup, nil,
acp.SessionServiceConfig{
MaxActiveGlobal: 10,
MaxActivePerUser: 5,
@@ -535,7 +646,7 @@ func TestSessionService_Create_MutuallyExclusiveInputs(t *testing.T) {
akID := uuid.New()
svc := acp.NewSessionService(
&fakeAcpRepo{}, nil, nil, nil, nil, nil, nil,
&fakeAcpRepo{}, nil, nil, nil, nil, nil, nil, nil, nil,
acp.SessionServiceConfig{}, nil,
)
@@ -598,7 +709,7 @@ func TestSessionService_List_RequirementFilter(t *testing.T) {
}
svc := acp.NewSessionService(
repo, nil, nil, nil, nil, nil, nil,
repo, nil, nil, nil, nil, nil, nil, nil, nil,
acp.SessionServiceConfig{}, nil,
)
@@ -616,3 +727,191 @@ func TestSessionService_List_RequirementFilter(t *testing.T) {
require.NoError(t, err)
require.Len(t, list, 2)
}
// TestSessionService_ResumeSession_ReacquiresMainWorktree 验证:在崩溃的 main-worktree
// session 上 ResumeSession 会重新占用 main worktree(CAS)、重新签发 system token、把
// session 复位为 starting,然后尝试 spawn。binary 为空致 spawn 失败,但前置的占用/签发/
// 复位都应已发生。
func TestSessionService_ResumeSession_ReacquiresMainWorktree(t *testing.T) {
t.Parallel()
uid := uuid.New()
wsID := uuid.New()
pid := uuid.New()
akID := uuid.New()
sessID := uuid.New()
exitCode := int32(1)
crashed := &acp.Session{
ID: sessID, WorkspaceID: wsID, ProjectID: pid, AgentKindID: akID, UserID: uid,
Branch: "main", CwdPath: "/tmp/main", IsMainWorktree: true,
Status: acp.SessionCrashed, ExitCode: &exitCode,
}
repo := &fakeAcpRepo{
kinds: []*acp.AgentKind{{ID: akID, Name: "claude", Enabled: true}},
sessions: []*acp.Session{crashed},
}
mcpTokens := &fakeMCPTokenSvc{}
sup := acp.NewSupervisor(
repo, nil, nil, nil, nil, mcpTokens, nil,
acp.SupervisorConfig{SpawnTimeout: 2 * time.Second, KillGrace: time.Second},
nil,
)
svc := acp.NewSessionService(
repo, &fakeWsSvc{ws: &workspace.Workspace{ID: wsID, ProjectID: pid, DefaultBranch: "main", MainPath: "/tmp/main"}},
nil, nil, fakeProjectAccess{pj: &project.Project{ID: pid, Slug: "p"}}, nil, nil, sup, nil,
acp.SessionServiceConfig{MaxActiveGlobal: 10, MaxActivePerUser: 5, SystemTokenTTL: time.Hour, MCPPublicURL: "http://x/mcp"},
mcpTokens,
)
tr, ok := svc.(acp.TurnRunner)
require.True(t, ok, "sessionService must implement acp.TurnRunner")
// spawn 会失败(binary 空),ResumeSession 返回错误是预期的。
_, _ = tr.ResumeSession(context.Background(), acp.Caller{UserID: uid}, sessID)
// 但 token 应已重新签发(占用 + InTx 复位在 spawn 之前)。
mcpTokens.mu.Lock()
issuedCount := len(mcpTokens.issued)
mcpTokens.mu.Unlock()
require.GreaterOrEqual(t, issuedCount, 1, "resume should re-issue a system token")
assert.Equal(t, sessID, mcpTokens.issued[0].ACPSessionID)
}
// TestSessionService_ResumeSession_ActiveSessionNoOp 验证:对已活跃 session 调用
// ResumeSession 直接返回该 session,不重新占用/签发。
func TestSessionService_ResumeSession_ActiveSessionNoOp(t *testing.T) {
t.Parallel()
uid := uuid.New()
wsID := uuid.New()
pid := uuid.New()
akID := uuid.New()
sessID := uuid.New()
active := &acp.Session{
ID: sessID, WorkspaceID: wsID, ProjectID: pid, AgentKindID: akID, UserID: uid,
Branch: "main", CwdPath: "/tmp/main", IsMainWorktree: true, Status: acp.SessionRunning,
}
repo := &fakeAcpRepo{
kinds: []*acp.AgentKind{{ID: akID, Name: "claude", Enabled: true}},
sessions: []*acp.Session{active},
}
mcpTokens := &fakeMCPTokenSvc{}
sup := acp.NewSupervisor(repo, nil, nil, nil, nil, mcpTokens, nil,
acp.SupervisorConfig{SpawnTimeout: time.Second, KillGrace: time.Second}, nil)
svc := acp.NewSessionService(repo, &fakeWsSvc{ws: &workspace.Workspace{ID: wsID, ProjectID: pid, DefaultBranch: "main", MainPath: "/tmp/main"}},
nil, nil, fakeProjectAccess{pj: &project.Project{ID: pid, Slug: "p"}}, nil, nil, sup, nil,
acp.SessionServiceConfig{MaxActiveGlobal: 10, MaxActivePerUser: 5, SystemTokenTTL: time.Hour, MCPPublicURL: "http://x/mcp"},
mcpTokens)
tr := svc.(acp.TurnRunner)
out, err := tr.ResumeSession(context.Background(), acp.Caller{UserID: uid}, sessID)
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, sessID, out.ID)
mcpTokens.mu.Lock()
defer mcpTokens.mu.Unlock()
assert.Len(t, mcpTokens.issued, 0, "active session resume must not re-issue tokens")
}
// ===== run-history dashboard =====
func newDashboardSvc(repo *fakeAcpRepo) acp.SessionService {
return acp.NewSessionService(repo, nil, nil, nil, nil, nil, nil, nil, nil,
acp.SessionServiceConfig{}, nil)
}
func TestDashboard_NonAdmin_ScopedToCaller(t *testing.T) {
t.Parallel()
uid := uuid.New()
repo := &fakeAcpRepo{
rollupResult: acp.SessionRollup{Total: 4, Succeeded: 3, Crashed: 1, TotalCostUSD: 1.5, TotalTokens: 900},
byDayResult: []acp.SessionDayRow{{Day: time.Now().UTC(), Total: 4, Succeeded: 3, Crashed: 1, TotalCostUSD: 1.5}},
}
svc := newDashboardSvc(repo)
res, err := svc.Dashboard(context.Background(), acp.Caller{UserID: uid}, acp.DashboardInput{})
require.NoError(t, err)
require.NotNil(t, res)
// 非 admin 强制按 caller scope。
require.NotNil(t, repo.rollupFilter.UserID)
assert.Equal(t, uid, *repo.rollupFilter.UserID)
require.NotNil(t, repo.byDayFilter.UserID)
assert.Equal(t, uid, *repo.byDayFilter.UserID)
// 默认时间窗为半开且 from 早于 to。
assert.True(t, repo.rollupFilter.From.Before(repo.rollupFilter.To))
// rollup 透传 + 成功率计算。
assert.Equal(t, int64(4), res.Rollup.Total)
assert.InDelta(t, 0.75, res.Rollup.SuccessRate(), 1e-9)
assert.InDelta(t, 1.5, res.Rollup.TotalCostUSD, 1e-9)
assert.Len(t, res.ByDay, 1)
}
func TestDashboard_Admin_NoUserScope_WithFilters(t *testing.T) {
t.Parallel()
pid := uuid.New()
reqID := uuid.New()
from := time.Now().Add(-48 * time.Hour).UTC()
to := time.Now().UTC()
repo := &fakeAcpRepo{rollupResult: acp.SessionRollup{Total: 10, Succeeded: 10}}
svc := newDashboardSvc(repo)
res, err := svc.Dashboard(context.Background(), acp.Caller{UserID: uuid.New(), IsAdmin: true},
acp.DashboardInput{ProjectID: &pid, RequirementID: &reqID, From: from, To: to})
require.NoError(t, err)
// admin → 无 user scope;project/requirement/时间窗透传。
assert.Nil(t, repo.rollupFilter.UserID)
require.NotNil(t, repo.rollupFilter.ProjectID)
assert.Equal(t, pid, *repo.rollupFilter.ProjectID)
require.NotNil(t, repo.rollupFilter.RequirementID)
assert.Equal(t, reqID, *repo.rollupFilter.RequirementID)
assert.Equal(t, from, repo.rollupFilter.From)
assert.Equal(t, to, repo.rollupFilter.To)
assert.InDelta(t, 1.0, res.Rollup.SuccessRate(), 1e-9)
}
func TestDashboard_InvalidWindow(t *testing.T) {
t.Parallel()
to := time.Now().UTC()
from := to.Add(time.Hour) // from after to
svc := newDashboardSvc(&fakeAcpRepo{})
_, err := svc.Dashboard(context.Background(), acp.Caller{UserID: uuid.New()},
acp.DashboardInput{From: from, To: to})
require.Error(t, err)
}
func TestTimeline_OwnerAccess(t *testing.T) {
t.Parallel()
uid := uuid.New()
sid := uuid.New()
buckets := []acp.SessionTimelineBucket{
{Direction: "out", RPCKind: "request", Method: "session/prompt", EventCount: 2},
{Direction: "in", RPCKind: "notification", Method: "session/update", EventCount: 5},
}
repo := &fakeAcpRepo{timelineResult: buckets}
repo.sessions = []*acp.Session{{ID: sid, UserID: uid}}
svc := newDashboardSvc(repo)
got, err := svc.Timeline(context.Background(), acp.Caller{UserID: uid}, sid)
require.NoError(t, err)
assert.Equal(t, sid, repo.timelineSID)
require.Len(t, got, 2)
assert.Equal(t, "session/prompt", got[0].Method)
}
func TestTimeline_NonOwner_NotFound(t *testing.T) {
t.Parallel()
sid := uuid.New()
repo := &fakeAcpRepo{}
repo.sessions = []*acp.Session{{ID: sid, UserID: uuid.New()}}
svc := newDashboardSvc(repo)
_, err := svc.Timeline(context.Background(), acp.Caller{UserID: uuid.New()}, sid)
require.Error(t, err) // Get 拒绝 → 不泄漏会话
}
@@ -30,7 +30,7 @@ func (q *Queries) DeleteAgentKindConfigFile(ctx context.Context, arg DeleteAgent
}
const listAgentKindConfigFiles = `-- name: ListAgentKindConfigFiles :many
SELECT id, agent_kind_id, rel_path, encrypted_content, updated_by, created_at, updated_at
SELECT id, agent_kind_id, rel_path, encrypted_content, updated_by, created_at, updated_at, key_version
FROM acp_agent_kind_config_files
WHERE agent_kind_id = $1
ORDER BY rel_path ASC
@@ -53,6 +53,7 @@ func (q *Queries) ListAgentKindConfigFiles(ctx context.Context, agentKindID pgty
&i.UpdatedBy,
&i.CreatedAt,
&i.UpdatedAt,
&i.KeyVersion,
); err != nil {
return nil, err
}
@@ -72,7 +73,7 @@ ON CONFLICT (agent_kind_id, rel_path) DO UPDATE
SET encrypted_content = EXCLUDED.encrypted_content,
updated_by = EXCLUDED.updated_by,
updated_at = now()
RETURNING id, agent_kind_id, rel_path, encrypted_content, updated_by, created_at, updated_at
RETURNING id, agent_kind_id, rel_path, encrypted_content, updated_by, created_at, updated_at, key_version
`
type UpsertAgentKindConfigFileParams struct {
@@ -100,6 +101,7 @@ func (q *Queries) UpsertAgentKindConfigFile(ctx context.Context, arg UpsertAgent
&i.UpdatedBy,
&i.CreatedAt,
&i.UpdatedAt,
&i.KeyVersion,
)
return i, err
}
+97 -40
View File
@@ -24,25 +24,31 @@ func (q *Queries) CountAgentKindUsage(ctx context.Context, agentKindID pgtype.UU
const createAgentKind = `-- name: CreateAgentKind :one
INSERT INTO acp_agent_kinds (
id, name, display_name, description, binary_path, args, encrypted_env, enabled, created_by, tool_allowlist, client_type, encrypted_mcp_servers
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
id, name, display_name, description, binary_path, args, encrypted_env, enabled, created_by, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
RETURNING id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
`
type CreateAgentKindParams struct {
ID pgtype.UUID `json:"id"`
Name string `json:"name"`
DisplayName string `json:"display_name"`
Description string `json:"description"`
BinaryPath string `json:"binary_path"`
Args []string `json:"args"`
EncryptedEnv []byte `json:"encrypted_env"`
Enabled bool `json:"enabled"`
CreatedBy pgtype.UUID `json:"created_by"`
ToolAllowlist []string `json:"tool_allowlist"`
ClientType string `json:"client_type"`
EncryptedMcpServers []byte `json:"encrypted_mcp_servers"`
ID pgtype.UUID `json:"id"`
Name string `json:"name"`
DisplayName string `json:"display_name"`
Description string `json:"description"`
BinaryPath string `json:"binary_path"`
Args []string `json:"args"`
EncryptedEnv []byte `json:"encrypted_env"`
Enabled bool `json:"enabled"`
CreatedBy pgtype.UUID `json:"created_by"`
ToolAllowlist []string `json:"tool_allowlist"`
ClientType string `json:"client_type"`
EncryptedMcpServers []byte `json:"encrypted_mcp_servers"`
ModelID pgtype.UUID `json:"model_id"`
MaxCostUsd pgtype.Numeric `json:"max_cost_usd"`
MaxTokens *int64 `json:"max_tokens"`
MaxWallClockSeconds *int32 `json:"max_wall_clock_seconds"`
}
func (q *Queries) CreateAgentKind(ctx context.Context, arg CreateAgentKindParams) (AcpAgentKind, error) {
@@ -59,6 +65,10 @@ func (q *Queries) CreateAgentKind(ctx context.Context, arg CreateAgentKindParams
arg.ToolAllowlist,
arg.ClientType,
arg.EncryptedMcpServers,
arg.ModelID,
arg.MaxCostUsd,
arg.MaxTokens,
arg.MaxWallClockSeconds,
)
var i AcpAgentKind
err := row.Scan(
@@ -76,6 +86,11 @@ func (q *Queries) CreateAgentKind(ctx context.Context, arg CreateAgentKindParams
&i.ToolAllowlist,
&i.ClientType,
&i.EncryptedMcpServers,
&i.ModelID,
&i.MaxCostUsd,
&i.MaxTokens,
&i.MaxWallClockSeconds,
&i.KeyVersion,
)
return i, err
}
@@ -91,7 +106,8 @@ func (q *Queries) DeleteAgentKind(ctx context.Context, id pgtype.UUID) error {
const getAgentKindByID = `-- name: GetAgentKindByID :one
SELECT id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
FROM acp_agent_kinds
WHERE id = $1
`
@@ -114,13 +130,19 @@ func (q *Queries) GetAgentKindByID(ctx context.Context, id pgtype.UUID) (AcpAgen
&i.ToolAllowlist,
&i.ClientType,
&i.EncryptedMcpServers,
&i.ModelID,
&i.MaxCostUsd,
&i.MaxTokens,
&i.MaxWallClockSeconds,
&i.KeyVersion,
)
return i, err
}
const getAgentKindByName = `-- name: GetAgentKindByName :one
SELECT id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
FROM acp_agent_kinds
WHERE name = $1
`
@@ -143,13 +165,19 @@ func (q *Queries) GetAgentKindByName(ctx context.Context, name string) (AcpAgent
&i.ToolAllowlist,
&i.ClientType,
&i.EncryptedMcpServers,
&i.ModelID,
&i.MaxCostUsd,
&i.MaxTokens,
&i.MaxWallClockSeconds,
&i.KeyVersion,
)
return i, err
}
const listAgentKinds = `-- name: ListAgentKinds :many
SELECT id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
FROM acp_agent_kinds
ORDER BY created_at DESC
`
@@ -178,6 +206,11 @@ func (q *Queries) ListAgentKinds(ctx context.Context) ([]AcpAgentKind, error) {
&i.ToolAllowlist,
&i.ClientType,
&i.EncryptedMcpServers,
&i.ModelID,
&i.MaxCostUsd,
&i.MaxTokens,
&i.MaxWallClockSeconds,
&i.KeyVersion,
); err != nil {
return nil, err
}
@@ -191,7 +224,8 @@ func (q *Queries) ListAgentKinds(ctx context.Context) ([]AcpAgentKind, error) {
const listEnabledAgentKinds = `-- name: ListEnabledAgentKinds :many
SELECT id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
FROM acp_agent_kinds
WHERE enabled = TRUE
ORDER BY name ASC
@@ -221,6 +255,11 @@ func (q *Queries) ListEnabledAgentKinds(ctx context.Context) ([]AcpAgentKind, er
&i.ToolAllowlist,
&i.ClientType,
&i.EncryptedMcpServers,
&i.ModelID,
&i.MaxCostUsd,
&i.MaxTokens,
&i.MaxWallClockSeconds,
&i.KeyVersion,
); err != nil {
return nil, err
}
@@ -234,32 +273,41 @@ func (q *Queries) ListEnabledAgentKinds(ctx context.Context) ([]AcpAgentKind, er
const updateAgentKind = `-- name: UpdateAgentKind :one
UPDATE acp_agent_kinds
SET display_name = $2,
description = $3,
binary_path = $4,
args = $5,
encrypted_env = $6,
enabled = $7,
tool_allowlist = $8,
client_type = $9,
encrypted_mcp_servers = $10,
updated_at = now()
SET display_name = $2,
description = $3,
binary_path = $4,
args = $5,
encrypted_env = $6,
enabled = $7,
tool_allowlist = $8,
client_type = $9,
encrypted_mcp_servers = $10,
model_id = $11,
max_cost_usd = $12,
max_tokens = $13,
max_wall_clock_seconds = $14,
updated_at = now()
WHERE id = $1
RETURNING id, name, display_name, description, binary_path, args, encrypted_env,
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers
enabled, created_by, created_at, updated_at, tool_allowlist, client_type, encrypted_mcp_servers,
model_id, max_cost_usd, max_tokens, max_wall_clock_seconds, key_version
`
type UpdateAgentKindParams struct {
ID pgtype.UUID `json:"id"`
DisplayName string `json:"display_name"`
Description string `json:"description"`
BinaryPath string `json:"binary_path"`
Args []string `json:"args"`
EncryptedEnv []byte `json:"encrypted_env"`
Enabled bool `json:"enabled"`
ToolAllowlist []string `json:"tool_allowlist"`
ClientType string `json:"client_type"`
EncryptedMcpServers []byte `json:"encrypted_mcp_servers"`
ID pgtype.UUID `json:"id"`
DisplayName string `json:"display_name"`
Description string `json:"description"`
BinaryPath string `json:"binary_path"`
Args []string `json:"args"`
EncryptedEnv []byte `json:"encrypted_env"`
Enabled bool `json:"enabled"`
ToolAllowlist []string `json:"tool_allowlist"`
ClientType string `json:"client_type"`
EncryptedMcpServers []byte `json:"encrypted_mcp_servers"`
ModelID pgtype.UUID `json:"model_id"`
MaxCostUsd pgtype.Numeric `json:"max_cost_usd"`
MaxTokens *int64 `json:"max_tokens"`
MaxWallClockSeconds *int32 `json:"max_wall_clock_seconds"`
}
func (q *Queries) UpdateAgentKind(ctx context.Context, arg UpdateAgentKindParams) (AcpAgentKind, error) {
@@ -274,6 +322,10 @@ func (q *Queries) UpdateAgentKind(ctx context.Context, arg UpdateAgentKindParams
arg.ToolAllowlist,
arg.ClientType,
arg.EncryptedMcpServers,
arg.ModelID,
arg.MaxCostUsd,
arg.MaxTokens,
arg.MaxWallClockSeconds,
)
var i AcpAgentKind
err := row.Scan(
@@ -291,6 +343,11 @@ func (q *Queries) UpdateAgentKind(ctx context.Context, arg UpdateAgentKindParams
&i.ToolAllowlist,
&i.ClientType,
&i.EncryptedMcpServers,
&i.ModelID,
&i.MaxCostUsd,
&i.MaxTokens,
&i.MaxWallClockSeconds,
&i.KeyVersion,
)
return i, err
}
+201
View File
@@ -0,0 +1,201 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.31.1
// source: dashboard.sql
package acpsqlc
import (
"context"
"github.com/jackc/pgx/v5/pgtype"
)
const sessionRollup = `-- name: SessionRollup :one
SELECT
COUNT(*)::bigint AS total,
COUNT(*) FILTER (WHERE status = 'exited')::bigint AS succeeded,
COUNT(*) FILTER (WHERE status = 'crashed')::bigint AS crashed,
COUNT(*) FILTER (WHERE status IN ('starting','running'))::bigint AS active,
COALESCE(SUM(cost_usd), 0)::numeric AS total_cost_usd,
COALESCE(SUM(tokens_total), 0)::bigint AS total_tokens,
COALESCE(
AVG(EXTRACT(EPOCH FROM (ended_at - started_at)))
FILTER (WHERE ended_at IS NOT NULL),
0)::float8 AS avg_duration_seconds
FROM acp_sessions
WHERE ($1::uuid IS NULL OR user_id = $1)
AND ($2::uuid IS NULL OR project_id = $2)
AND ($3::uuid IS NULL OR requirement_id = $3)
AND started_at >= $4
AND started_at < $5
`
type SessionRollupParams struct {
Column1 pgtype.UUID `json:"column_1"`
Column2 pgtype.UUID `json:"column_2"`
Column3 pgtype.UUID `json:"column_3"`
StartedAt pgtype.Timestamptz `json:"started_at"`
StartedAt_2 pgtype.Timestamptz `json:"started_at_2"`
}
type SessionRollupRow struct {
Total int64 `json:"total"`
Succeeded int64 `json:"succeeded"`
Crashed int64 `json:"crashed"`
Active int64 `json:"active"`
TotalCostUsd pgtype.Numeric `json:"total_cost_usd"`
TotalTokens int64 `json:"total_tokens"`
AvgDurationSeconds float64 `json:"avg_duration_seconds"`
}
// owner/admin scope + 可选 project/requirement + 时间窗内的会话汇总。
// success = status='exited';avg_duration_seconds 仅统计已结束(ended_at 非空)会话。
func (q *Queries) SessionRollup(ctx context.Context, arg SessionRollupParams) (SessionRollupRow, error) {
row := q.db.QueryRow(ctx, sessionRollup,
arg.Column1,
arg.Column2,
arg.Column3,
arg.StartedAt,
arg.StartedAt_2,
)
var i SessionRollupRow
err := row.Scan(
&i.Total,
&i.Succeeded,
&i.Crashed,
&i.Active,
&i.TotalCostUsd,
&i.TotalTokens,
&i.AvgDurationSeconds,
)
return i, err
}
const sessionTimeline = `-- name: SessionTimeline :many
SELECT direction,
rpc_kind,
COALESCE(method, '') AS method,
COUNT(*)::bigint AS event_count,
MIN(created_at)::timestamptz AS first_at,
MAX(created_at)::timestamptz AS last_at
FROM acp_events
WHERE session_id = $1
GROUP BY direction, rpc_kind, COALESCE(method, '')
ORDER BY first_at ASC
`
type SessionTimelineRow struct {
Direction string `json:"direction"`
RpcKind string `json:"rpc_kind"`
Method string `json:"method"`
EventCount int64 `json:"event_count"`
FirstAt pgtype.Timestamptz `json:"first_at"`
LastAt pgtype.Timestamptz `json:"last_at"`
}
// dashboard.sql: run-history dashboard 读模型(只读聚合)。
// - SessionTimeline: 单会话的 RPC 事件按 (direction, rpc_kind, method) 分桶,
// 给出条数 + 首/末时间戳,用于会话时间线概览。
// - SessionRollup: 在 owner/admin scope + 可选 project/requirement 过滤下,
// 统计会话总数、成功率(exited 视为成功)、总成本、平均时长。
// - SessionsByDay: 按天的 总数 / 成功 / 失败 / 成本 时间序列。
//
// scope 约定(与 ListSessionsByUser/ListAllSessions 一致):
//
// $1 user_id 为 NULL 时不按用户过滤(admin 全量),非 NULL 时仅该用户。
//
// 单会话事件时间线:按 method/kind 分桶聚合(条数 + 首末时间)。
func (q *Queries) SessionTimeline(ctx context.Context, sessionID pgtype.UUID) ([]SessionTimelineRow, error) {
rows, err := q.db.Query(ctx, sessionTimeline, sessionID)
if err != nil {
return nil, err
}
defer rows.Close()
var items []SessionTimelineRow
for rows.Next() {
var i SessionTimelineRow
if err := rows.Scan(
&i.Direction,
&i.RpcKind,
&i.Method,
&i.EventCount,
&i.FirstAt,
&i.LastAt,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const sessionsByDay = `-- name: SessionsByDay :many
SELECT
date_trunc('day', started_at)::timestamptz AS day,
COUNT(*)::bigint AS total,
COUNT(*) FILTER (WHERE status = 'exited')::bigint AS succeeded,
COUNT(*) FILTER (WHERE status = 'crashed')::bigint AS crashed,
COALESCE(SUM(cost_usd), 0)::numeric AS total_cost_usd
FROM acp_sessions
WHERE ($1::uuid IS NULL OR user_id = $1)
AND ($2::uuid IS NULL OR project_id = $2)
AND ($3::uuid IS NULL OR requirement_id = $3)
AND started_at >= $4
AND started_at < $5
GROUP BY day
ORDER BY day ASC
`
type SessionsByDayParams struct {
Column1 pgtype.UUID `json:"column_1"`
Column2 pgtype.UUID `json:"column_2"`
Column3 pgtype.UUID `json:"column_3"`
StartedAt pgtype.Timestamptz `json:"started_at"`
StartedAt_2 pgtype.Timestamptz `json:"started_at_2"`
}
type SessionsByDayRow struct {
Day pgtype.Timestamptz `json:"day"`
Total int64 `json:"total"`
Succeeded int64 `json:"succeeded"`
Crashed int64 `json:"crashed"`
TotalCostUsd pgtype.Numeric `json:"total_cost_usd"`
}
// 按天的会话时间序列(总数 / 成功 / 失败 / 成本),同 SessionRollup 的 scope。
func (q *Queries) SessionsByDay(ctx context.Context, arg SessionsByDayParams) ([]SessionsByDayRow, error) {
rows, err := q.db.Query(ctx, sessionsByDay,
arg.Column1,
arg.Column2,
arg.Column3,
arg.StartedAt,
arg.StartedAt_2,
)
if err != nil {
return nil, err
}
defer rows.Close()
var items []SessionsByDayRow
for rows.Next() {
var i SessionsByDayRow
if err := rows.Scan(
&i.Day,
&i.Total,
&i.Succeeded,
&i.Crashed,
&i.TotalCostUsd,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
+240 -17
View File
@@ -8,6 +8,7 @@ import (
"net/netip"
"github.com/jackc/pgx/v5/pgtype"
"github.com/pgvector/pgvector-go"
)
type AcpAgentKind struct {
@@ -25,6 +26,11 @@ type AcpAgentKind struct {
ToolAllowlist []string `json:"tool_allowlist"`
ClientType string `json:"client_type"`
EncryptedMcpServers []byte `json:"encrypted_mcp_servers"`
ModelID pgtype.UUID `json:"model_id"`
MaxCostUsd pgtype.Numeric `json:"max_cost_usd"`
MaxTokens *int64 `json:"max_tokens"`
MaxWallClockSeconds *int32 `json:"max_wall_clock_seconds"`
KeyVersion int16 `json:"key_version"`
}
type AcpAgentKindConfigFile struct {
@@ -35,6 +41,7 @@ type AcpAgentKindConfigFile struct {
UpdatedBy pgtype.UUID `json:"updated_by"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
KeyVersion int16 `json:"key_version"`
}
type AcpEvent struct {
@@ -64,23 +71,64 @@ type AcpPermissionRequest struct {
}
type AcpSession struct {
ID pgtype.UUID `json:"id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
ProjectID pgtype.UUID `json:"project_id"`
AgentKindID pgtype.UUID `json:"agent_kind_id"`
UserID pgtype.UUID `json:"user_id"`
IssueID pgtype.UUID `json:"issue_id"`
RequirementID pgtype.UUID `json:"requirement_id"`
AgentSessionID *string `json:"agent_session_id"`
Branch string `json:"branch"`
CwdPath string `json:"cwd_path"`
IsMainWorktree bool `json:"is_main_worktree"`
Status string `json:"status"`
Pid *int32 `json:"pid"`
ExitCode *int32 `json:"exit_code"`
LastError *string `json:"last_error"`
StartedAt pgtype.Timestamptz `json:"started_at"`
EndedAt pgtype.Timestamptz `json:"ended_at"`
ID pgtype.UUID `json:"id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
ProjectID pgtype.UUID `json:"project_id"`
AgentKindID pgtype.UUID `json:"agent_kind_id"`
UserID pgtype.UUID `json:"user_id"`
IssueID pgtype.UUID `json:"issue_id"`
RequirementID pgtype.UUID `json:"requirement_id"`
AgentSessionID *string `json:"agent_session_id"`
Branch string `json:"branch"`
CwdPath string `json:"cwd_path"`
IsMainWorktree bool `json:"is_main_worktree"`
Status string `json:"status"`
Pid *int32 `json:"pid"`
ExitCode *int32 `json:"exit_code"`
LastError *string `json:"last_error"`
StartedAt pgtype.Timestamptz `json:"started_at"`
EndedAt pgtype.Timestamptz `json:"ended_at"`
LastStopReason *string `json:"last_stop_reason"`
OrchestratorStepID pgtype.UUID `json:"orchestrator_step_id"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
TotalCostUsd pgtype.Numeric `json:"total_cost_usd"`
LastActivityAt pgtype.Timestamptz `json:"last_activity_at"`
BudgetMaxCostUsd pgtype.Numeric `json:"budget_max_cost_usd"`
BudgetMaxTokens *int64 `json:"budget_max_tokens"`
BudgetMaxWallClockSeconds *int32 `json:"budget_max_wall_clock_seconds"`
TerminatedReason *string `json:"terminated_reason"`
SandboxMode *string `json:"sandbox_mode"`
CostUsd pgtype.Numeric `json:"cost_usd"`
TokensTotal int64 `json:"tokens_total"`
}
type AcpSessionUsage struct {
ID int64 `json:"id"`
SessionID pgtype.UUID `json:"session_id"`
UserID pgtype.UUID `json:"user_id"`
ProjectID pgtype.UUID `json:"project_id"`
AgentKindID pgtype.UUID `json:"agent_kind_id"`
ModelID pgtype.UUID `json:"model_id"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
CostUsd pgtype.Numeric `json:"cost_usd"`
SourceEventID *int64 `json:"source_event_id"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type AcpTurn struct {
ID int64 `json:"id"`
SessionID pgtype.UUID `json:"session_id"`
TurnIndex int32 `json:"turn_index"`
PromptRequestID *string `json:"prompt_request_id"`
Status string `json:"status"`
StopReason *string `json:"stop_reason"`
UpdateCount int32 `json:"update_count"`
StartedAt pgtype.Timestamptz `json:"started_at"`
CompletedAt pgtype.Timestamptz `json:"completed_at"`
}
type AuditLog struct {
@@ -94,6 +142,81 @@ type AuditLog struct {
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type ChangeRequest struct {
ID pgtype.UUID `json:"id"`
ProjectID pgtype.UUID `json:"project_id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
RequirementID pgtype.UUID `json:"requirement_id"`
IssueID pgtype.UUID `json:"issue_id"`
Number int32 `json:"number"`
Title string `json:"title"`
Description string `json:"description"`
SourceBranch string `json:"source_branch"`
TargetBranch string `json:"target_branch"`
State string `json:"state"`
ReviewVerdict string `json:"review_verdict"`
CiState string `json:"ci_state"`
Provider string `json:"provider"`
ExternalID *int64 `json:"external_id"`
ExternalUrl string `json:"external_url"`
MergeCommitSha string `json:"merge_commit_sha"`
ReviewedBy pgtype.UUID `json:"reviewed_by"`
ReviewedAt pgtype.Timestamptz `json:"reviewed_at"`
CreatedBy pgtype.UUID `json:"created_by"`
MergedAt pgtype.Timestamptz `json:"merged_at"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
}
type CodeChunk struct {
ID pgtype.UUID `json:"id"`
RunID pgtype.UUID `json:"run_id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
CommitSha string `json:"commit_sha"`
FilePath string `json:"file_path"`
Lang *string `json:"lang"`
StartLine int32 `json:"start_line"`
EndLine int32 `json:"end_line"`
Content string `json:"content"`
ContentSha []byte `json:"content_sha"`
TokenCount *int32 `json:"token_count"`
Embedding *pgvector.Vector `json:"embedding"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type CodeIndexRun struct {
ID pgtype.UUID `json:"id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
WorktreeID pgtype.UUID `json:"worktree_id"`
Branch string `json:"branch"`
CommitSha string `json:"commit_sha"`
Status string `json:"status"`
EmbeddingEndpointID pgtype.UUID `json:"embedding_endpoint_id"`
EmbeddingModel string `json:"embedding_model"`
Dims int32 `json:"dims"`
FilesIndexed int32 `json:"files_indexed"`
ChunksIndexed int32 `json:"chunks_indexed"`
Error *string `json:"error"`
StartedAt pgtype.Timestamptz `json:"started_at"`
FinishedAt pgtype.Timestamptz `json:"finished_at"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type CommitSummary struct {
ID pgtype.UUID `json:"id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
WorktreeID pgtype.UUID `json:"worktree_id"`
Branch string `json:"branch"`
CommitSha string `json:"commit_sha"`
Kind string `json:"kind"`
Title *string `json:"title"`
BodyMd string `json:"body_md"`
Model *string `json:"model"`
Status string `json:"status"`
Error *string `json:"error"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type Conversation struct {
ID pgtype.UUID `json:"id"`
UserID pgtype.UUID `json:"user_id"`
@@ -110,6 +233,14 @@ type Conversation struct {
RequirementID pgtype.UUID `json:"requirement_id"`
}
type CryptoKey struct {
Version int32 `json:"version"`
Provider string `json:"provider"`
KeyRef *string `json:"key_ref"`
Status string `json:"status"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type GitCredential struct {
ID pgtype.UUID `json:"id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
@@ -119,6 +250,7 @@ type GitCredential struct {
Fingerprint *string `json:"fingerprint"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
KeyVersion int16 `json:"key_version"`
}
type Issue struct {
@@ -135,6 +267,17 @@ type Issue struct {
CreatedAt pgtype.Timestamptz `json:"created_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
ParentID pgtype.UUID `json:"parent_id"`
Priority int32 `json:"priority"`
}
type IssueDependency struct {
ID pgtype.UUID `json:"id"`
ProjectID pgtype.UUID `json:"project_id"`
BlockedID pgtype.UUID `json:"blocked_id"`
BlockerID pgtype.UUID `json:"blocker_id"`
CreatedBy pgtype.UUID `json:"created_by"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type Job struct {
@@ -162,6 +305,7 @@ type LlmEndpoint struct {
CreatedBy pgtype.UUID `json:"created_by"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
KeyVersion int16 `json:"key_version"`
}
type LlmModel struct {
@@ -252,6 +396,50 @@ type Notification struct {
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type NotificationPref struct {
UserID pgtype.UUID `json:"user_id"`
EmailEnabled bool `json:"email_enabled"`
ImEnabled bool `json:"im_enabled"`
ImWebhookUrlEnc []byte `json:"im_webhook_url_enc"`
MinSeverity string `json:"min_severity"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
}
type OrchestratorRun struct {
ID pgtype.UUID `json:"id"`
ProjectID pgtype.UUID `json:"project_id"`
RequirementID pgtype.UUID `json:"requirement_id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
OwnerID pgtype.UUID `json:"owner_id"`
Status string `json:"status"`
CurrentPhase string `json:"current_phase"`
Config []byte `json:"config"`
LastError *string `json:"last_error"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
EndedAt pgtype.Timestamptz `json:"ended_at"`
}
type OrchestratorStep struct {
ID pgtype.UUID `json:"id"`
RunID pgtype.UUID `json:"run_id"`
Phase string `json:"phase"`
Attempt int32 `json:"attempt"`
ParentStepID pgtype.UUID `json:"parent_step_id"`
Depth int32 `json:"depth"`
Status string `json:"status"`
AgentKindID pgtype.UUID `json:"agent_kind_id"`
AcpSessionID pgtype.UUID `json:"acp_session_id"`
Prompt string `json:"prompt"`
StopReason *string `json:"stop_reason"`
GateResult []byte `json:"gate_result"`
LastError *string `json:"last_error"`
JobID pgtype.UUID `json:"job_id"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
EndedAt pgtype.Timestamptz `json:"ended_at"`
}
type Project struct {
ID pgtype.UUID `json:"id"`
Slug string `json:"slug"`
@@ -264,6 +452,30 @@ type Project struct {
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
}
type ProjectMember struct {
ProjectID pgtype.UUID `json:"project_id"`
UserID pgtype.UUID `json:"user_id"`
Role string `json:"role"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
AddedBy pgtype.UUID `json:"added_by"`
}
type ProjectMemory struct {
ID pgtype.UUID `json:"id"`
ProjectID pgtype.UUID `json:"project_id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
Kind string `json:"kind"`
Title string `json:"title"`
Body string `json:"body"`
Tags []string `json:"tags"`
Source string `json:"source"`
Embedding *pgvector.Vector `json:"embedding"`
EmbeddingModel *string `json:"embedding_model"`
CreatedBy pgtype.UUID `json:"created_by"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
}
type PromptTemplate struct {
ID pgtype.UUID `json:"id"`
Name string `json:"name"`
@@ -299,6 +511,16 @@ type RequirementArtifact struct {
SourceMessageID *int64 `json:"source_message_id"`
CreatedBy pgtype.UUID `json:"created_by"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
Verdict string `json:"verdict"`
}
type RequirementPhasePointer struct {
RequirementID pgtype.UUID `json:"requirement_id"`
Phase string `json:"phase"`
ApprovedArtifactID pgtype.UUID `json:"approved_artifact_id"`
ApprovedBy pgtype.UUID `json:"approved_by"`
ApprovedAt pgtype.Timestamptz `json:"approved_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
}
type User struct {
@@ -354,6 +576,7 @@ type WorkspaceRunProfile struct {
CreatedBy pgtype.UUID `json:"created_by"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
KeyVersion int16 `json:"key_version"`
}
type WorkspaceWorktree struct {
+54
View File
@@ -0,0 +1,54 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.31.1
// source: session_usage.sql
package acpsqlc
import (
"context"
"github.com/jackc/pgx/v5/pgtype"
)
const insertSessionUsage = `-- name: InsertSessionUsage :one
INSERT INTO acp_session_usage (
session_id, user_id, project_id, agent_kind_id, model_id,
prompt_tokens, completion_tokens, thinking_tokens, cost_usd, source_event_id
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
ON CONFLICT (source_event_id) WHERE source_event_id IS NOT NULL DO NOTHING
RETURNING id
`
type InsertSessionUsageParams struct {
SessionID pgtype.UUID `json:"session_id"`
UserID pgtype.UUID `json:"user_id"`
ProjectID pgtype.UUID `json:"project_id"`
AgentKindID pgtype.UUID `json:"agent_kind_id"`
ModelID pgtype.UUID `json:"model_id"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
CostUsd pgtype.Numeric `json:"cost_usd"`
SourceEventID *int64 `json:"source_event_id"`
}
// 写一条 per-turn 用量账目。source_event_id 非空时按唯一索引去重(同一事件
// 重复 Observe 不会重复入账);返回受影响行数判定是否实际插入。
func (q *Queries) InsertSessionUsage(ctx context.Context, arg InsertSessionUsageParams) (int64, error) {
row := q.db.QueryRow(ctx, insertSessionUsage,
arg.SessionID,
arg.UserID,
arg.ProjectID,
arg.AgentKindID,
arg.ModelID,
arg.PromptTokens,
arg.CompletionTokens,
arg.ThinkingTokens,
arg.CostUsd,
arg.SourceEventID,
)
var id int64
err := row.Scan(&id)
return id, err
}
+611 -17
View File
@@ -29,6 +29,56 @@ func (q *Queries) AcquireMainWorktree(ctx context.Context, arg AcquireMainWorktr
return result.RowsAffected(), nil
}
const addSessionUsageTotals = `-- name: AddSessionUsageTotals :one
UPDATE acp_sessions
SET prompt_tokens = prompt_tokens + $2,
completion_tokens = completion_tokens + $3,
thinking_tokens = thinking_tokens + $4,
total_cost_usd = total_cost_usd + $5,
cost_usd = cost_usd + $5,
tokens_total = tokens_total + $2 + $3 + $4,
last_activity_at = now()
WHERE id = $1
RETURNING prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd
`
type AddSessionUsageTotalsParams struct {
ID pgtype.UUID `json:"id"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
TotalCostUsd pgtype.Numeric `json:"total_cost_usd"`
}
type AddSessionUsageTotalsRow struct {
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
TotalCostUsd pgtype.Numeric `json:"total_cost_usd"`
}
// 单 round-trip 累加 session 运行时 token/cost 总量并刷新 last_activity_at,
// 返回累加后的权威总量供预算判定。
// 同时把 dashboard rollup 列(cost_usd / tokens_total)与运行时累加器对齐,
// 使 run-history dashboard 在会话进行中及退出后都能读到 agent 上报的成本。
func (q *Queries) AddSessionUsageTotals(ctx context.Context, arg AddSessionUsageTotalsParams) (AddSessionUsageTotalsRow, error) {
row := q.db.QueryRow(ctx, addSessionUsageTotals,
arg.ID,
arg.PromptTokens,
arg.CompletionTokens,
arg.ThinkingTokens,
arg.TotalCostUsd,
)
var i AddSessionUsageTotalsRow
err := row.Scan(
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.TotalCostUsd,
)
return i, err
}
const countActiveSessions = `-- name: CountActiveSessions :one
SELECT COUNT(*) FROM acp_sessions WHERE status IN ('starting', 'running')
`
@@ -52,11 +102,73 @@ func (q *Queries) CountActiveSessionsByUser(ctx context.Context, userID pgtype.U
return count, err
}
const getCrashedSessionForResume = `-- name: GetCrashedSessionForResume :one
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE orchestrator_step_id = $1
AND status IN ('crashed','exited')
ORDER BY started_at DESC
LIMIT 1
`
// 返回某 step 最近一次处于可恢复终态(crashed/exited)的 session。
func (q *Queries) GetCrashedSessionForResume(ctx context.Context, orchestratorStepID pgtype.UUID) (AcpSession, error) {
row := q.db.QueryRow(ctx, getCrashedSessionForResume, orchestratorStepID)
var i AcpSession
err := row.Scan(
&i.ID,
&i.WorkspaceID,
&i.ProjectID,
&i.AgentKindID,
&i.UserID,
&i.IssueID,
&i.RequirementID,
&i.AgentSessionID,
&i.Branch,
&i.CwdPath,
&i.IsMainWorktree,
&i.Status,
&i.Pid,
&i.ExitCode,
&i.LastError,
&i.StartedAt,
&i.EndedAt,
&i.LastStopReason,
&i.OrchestratorStepID,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.TotalCostUsd,
&i.LastActivityAt,
&i.BudgetMaxCostUsd,
&i.BudgetMaxTokens,
&i.BudgetMaxWallClockSeconds,
&i.TerminatedReason,
&i.SandboxMode,
&i.CostUsd,
&i.TokensTotal,
)
return i, err
}
const getSessionByID = `-- name: GetSessionByID :one
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE id = $1
`
@@ -82,6 +194,75 @@ func (q *Queries) GetSessionByID(ctx context.Context, id pgtype.UUID) (AcpSessio
&i.LastError,
&i.StartedAt,
&i.EndedAt,
&i.LastStopReason,
&i.OrchestratorStepID,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.TotalCostUsd,
&i.LastActivityAt,
&i.BudgetMaxCostUsd,
&i.BudgetMaxTokens,
&i.BudgetMaxWallClockSeconds,
&i.TerminatedReason,
&i.SandboxMode,
&i.CostUsd,
&i.TokensTotal,
)
return i, err
}
const getSessionByStepID = `-- name: GetSessionByStepID :one
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE orchestrator_step_id = $1
ORDER BY started_at DESC
LIMIT 1
`
func (q *Queries) GetSessionByStepID(ctx context.Context, orchestratorStepID pgtype.UUID) (AcpSession, error) {
row := q.db.QueryRow(ctx, getSessionByStepID, orchestratorStepID)
var i AcpSession
err := row.Scan(
&i.ID,
&i.WorkspaceID,
&i.ProjectID,
&i.AgentKindID,
&i.UserID,
&i.IssueID,
&i.RequirementID,
&i.AgentSessionID,
&i.Branch,
&i.CwdPath,
&i.IsMainWorktree,
&i.Status,
&i.Pid,
&i.ExitCode,
&i.LastError,
&i.StartedAt,
&i.EndedAt,
&i.LastStopReason,
&i.OrchestratorStepID,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.TotalCostUsd,
&i.LastActivityAt,
&i.BudgetMaxCostUsd,
&i.BudgetMaxTokens,
&i.BudgetMaxWallClockSeconds,
&i.TerminatedReason,
&i.SandboxMode,
&i.CostUsd,
&i.TokensTotal,
)
return i, err
}
@@ -89,26 +270,37 @@ func (q *Queries) GetSessionByID(ctx context.Context, id pgtype.UUID) (AcpSessio
const insertSession = `-- name: InsertSession :one
INSERT INTO acp_sessions (
id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, branch, cwd_path, is_main_worktree, status
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
issue_id, requirement_id, branch, cwd_path, is_main_worktree, status,
orchestrator_step_id,
budget_max_cost_usd, budget_max_tokens, budget_max_wall_clock_seconds
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
RETURNING id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
`
type InsertSessionParams struct {
ID pgtype.UUID `json:"id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
ProjectID pgtype.UUID `json:"project_id"`
AgentKindID pgtype.UUID `json:"agent_kind_id"`
UserID pgtype.UUID `json:"user_id"`
IssueID pgtype.UUID `json:"issue_id"`
RequirementID pgtype.UUID `json:"requirement_id"`
Branch string `json:"branch"`
CwdPath string `json:"cwd_path"`
IsMainWorktree bool `json:"is_main_worktree"`
Status string `json:"status"`
ID pgtype.UUID `json:"id"`
WorkspaceID pgtype.UUID `json:"workspace_id"`
ProjectID pgtype.UUID `json:"project_id"`
AgentKindID pgtype.UUID `json:"agent_kind_id"`
UserID pgtype.UUID `json:"user_id"`
IssueID pgtype.UUID `json:"issue_id"`
RequirementID pgtype.UUID `json:"requirement_id"`
Branch string `json:"branch"`
CwdPath string `json:"cwd_path"`
IsMainWorktree bool `json:"is_main_worktree"`
Status string `json:"status"`
OrchestratorStepID pgtype.UUID `json:"orchestrator_step_id"`
BudgetMaxCostUsd pgtype.Numeric `json:"budget_max_cost_usd"`
BudgetMaxTokens *int64 `json:"budget_max_tokens"`
BudgetMaxWallClockSeconds *int32 `json:"budget_max_wall_clock_seconds"`
}
func (q *Queries) InsertSession(ctx context.Context, arg InsertSessionParams) (AcpSession, error) {
@@ -124,6 +316,10 @@ func (q *Queries) InsertSession(ctx context.Context, arg InsertSessionParams) (A
arg.CwdPath,
arg.IsMainWorktree,
arg.Status,
arg.OrchestratorStepID,
arg.BudgetMaxCostUsd,
arg.BudgetMaxTokens,
arg.BudgetMaxWallClockSeconds,
)
var i AcpSession
err := row.Scan(
@@ -144,6 +340,20 @@ func (q *Queries) InsertSession(ctx context.Context, arg InsertSessionParams) (A
&i.LastError,
&i.StartedAt,
&i.EndedAt,
&i.LastStopReason,
&i.OrchestratorStepID,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.TotalCostUsd,
&i.LastActivityAt,
&i.BudgetMaxCostUsd,
&i.BudgetMaxTokens,
&i.BudgetMaxWallClockSeconds,
&i.TerminatedReason,
&i.SandboxMode,
&i.CostUsd,
&i.TokensTotal,
)
return i, err
}
@@ -152,7 +362,12 @@ const listAllSessions = `-- name: ListAllSessions :many
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE ($1::uuid IS NULL OR requirement_id = $1)
ORDER BY started_at DESC
@@ -185,6 +400,68 @@ func (q *Queries) ListAllSessions(ctx context.Context, dollar_1 pgtype.UUID) ([]
&i.LastError,
&i.StartedAt,
&i.EndedAt,
&i.LastStopReason,
&i.OrchestratorStepID,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.TotalCostUsd,
&i.LastActivityAt,
&i.BudgetMaxCostUsd,
&i.BudgetMaxTokens,
&i.BudgetMaxWallClockSeconds,
&i.TerminatedReason,
&i.SandboxMode,
&i.CostUsd,
&i.TokensTotal,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const listSessionUsageBySession = `-- name: ListSessionUsageBySession :many
SELECT id, session_id, user_id, project_id, agent_kind_id, model_id,
prompt_tokens, completion_tokens, thinking_tokens, cost_usd,
source_event_id, created_at
FROM acp_session_usage
WHERE session_id = $1
ORDER BY created_at DESC
LIMIT $2
`
type ListSessionUsageBySessionParams struct {
SessionID pgtype.UUID `json:"session_id"`
Limit int32 `json:"limit"`
}
func (q *Queries) ListSessionUsageBySession(ctx context.Context, arg ListSessionUsageBySessionParams) ([]AcpSessionUsage, error) {
rows, err := q.db.Query(ctx, listSessionUsageBySession, arg.SessionID, arg.Limit)
if err != nil {
return nil, err
}
defer rows.Close()
var items []AcpSessionUsage
for rows.Next() {
var i AcpSessionUsage
if err := rows.Scan(
&i.ID,
&i.SessionID,
&i.UserID,
&i.ProjectID,
&i.AgentKindID,
&i.ModelID,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.CostUsd,
&i.SourceEventID,
&i.CreatedAt,
); err != nil {
return nil, err
}
@@ -200,7 +477,12 @@ const listSessionsByUser = `-- name: ListSessionsByUser :many
SELECT id, workspace_id, project_id, agent_kind_id, user_id,
issue_id, requirement_id, agent_session_id,
branch, cwd_path, is_main_worktree, status,
pid, exit_code, last_error, started_at, ended_at
pid, exit_code, last_error, started_at, ended_at, last_stop_reason,
orchestrator_step_id,
prompt_tokens, completion_tokens, thinking_tokens, total_cost_usd,
last_activity_at, budget_max_cost_usd, budget_max_tokens,
budget_max_wall_clock_seconds, terminated_reason, sandbox_mode,
cost_usd, tokens_total
FROM acp_sessions
WHERE user_id = $1
AND ($2::uuid IS NULL OR requirement_id = $2)
@@ -239,6 +521,76 @@ func (q *Queries) ListSessionsByUser(ctx context.Context, arg ListSessionsByUser
&i.LastError,
&i.StartedAt,
&i.EndedAt,
&i.LastStopReason,
&i.OrchestratorStepID,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.TotalCostUsd,
&i.LastActivityAt,
&i.BudgetMaxCostUsd,
&i.BudgetMaxTokens,
&i.BudgetMaxWallClockSeconds,
&i.TerminatedReason,
&i.SandboxMode,
&i.CostUsd,
&i.TokensTotal,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const listSessionsForReaper = `-- name: ListSessionsForReaper :many
SELECT id, user_id, started_at, last_activity_at, budget_max_wall_clock_seconds
FROM acp_sessions
WHERE status IN ('starting','running')
AND (
(budget_max_wall_clock_seconds IS NOT NULL
AND started_at < now() - make_interval(secs => budget_max_wall_clock_seconds))
OR ($1::timestamptz IS NOT NULL AND started_at < $1::timestamptz)
OR last_activity_at < $2::timestamptz
)
`
type ListSessionsForReaperParams struct {
Column1 pgtype.Timestamptz `json:"column_1"`
Column2 pgtype.Timestamptz `json:"column_2"`
}
type ListSessionsForReaperRow struct {
ID pgtype.UUID `json:"id"`
UserID pgtype.UUID `json:"user_id"`
StartedAt pgtype.Timestamptz `json:"started_at"`
LastActivityAt pgtype.Timestamptz `json:"last_activity_at"`
BudgetMaxWallClockSeconds *int32 `json:"budget_max_wall_clock_seconds"`
}
// 返回应被回收的活跃 session:
//
// 超过自身 wall-clock 上限(budget_max_wall_clock_seconds 非空且 started_at 过早),或
// 超过全局 wall-clock 上限($1::timestamptz = now - 全局上限),或
// 空闲超时(last_activity_at < $2::timestamptz)。
func (q *Queries) ListSessionsForReaper(ctx context.Context, arg ListSessionsForReaperParams) ([]ListSessionsForReaperRow, error) {
rows, err := q.db.Query(ctx, listSessionsForReaper, arg.Column1, arg.Column2)
if err != nil {
return nil, err
}
defer rows.Close()
var items []ListSessionsForReaperRow
for rows.Next() {
var i ListSessionsForReaperRow
if err := rows.Scan(
&i.ID,
&i.UserID,
&i.StartedAt,
&i.LastActivityAt,
&i.BudgetMaxWallClockSeconds,
); err != nil {
return nil, err
}
@@ -269,6 +621,23 @@ func (q *Queries) MarkSessionFailedIfActive(ctx context.Context, arg MarkSession
return result.RowsAffected(), nil
}
const markSessionTerminatedReason = `-- name: MarkSessionTerminatedReason :exec
UPDATE acp_sessions
SET terminated_reason = $2
WHERE id = $1 AND terminated_reason IS NULL
`
type MarkSessionTerminatedReasonParams struct {
ID pgtype.UUID `json:"id"`
TerminatedReason *string `json:"terminated_reason"`
}
// 写入终止原因;已有非空 terminated_reason 时不覆盖(reaper / budget 互斥保护)。
func (q *Queries) MarkSessionTerminatedReason(ctx context.Context, arg MarkSessionTerminatedReasonParams) error {
_, err := q.db.Exec(ctx, markSessionTerminatedReason, arg.ID, arg.TerminatedReason)
return err
}
const releaseMainWorktree = `-- name: ReleaseMainWorktree :execrows
UPDATE workspaces SET active_main_session_id = NULL
WHERE id = $1 AND active_main_session_id = $2
@@ -287,6 +656,26 @@ func (q *Queries) ReleaseMainWorktree(ctx context.Context, arg ReleaseMainWorktr
return result.RowsAffected(), nil
}
const resetSessionForResume = `-- name: ResetSessionForResume :exec
UPDATE acp_sessions
SET status = 'starting',
agent_session_id = NULL,
pid = NULL,
exit_code = NULL,
last_error = NULL,
ended_at = NULL,
terminated_reason = NULL,
last_activity_at = now()
WHERE id = $1
`
// 把崩溃/退出的 session 复用为新一轮:回到 starting、清空上次运行时字段。
// 编排器 resume 在重新 Spawn 前调用,使同一行 session 在同 cwd/worktree 上复用。
func (q *Queries) ResetSessionForResume(ctx context.Context, id pgtype.UUID) error {
_, err := q.db.Exec(ctx, resetSessionForResume, id)
return err
}
const resetStuckSessionsOnRestart = `-- name: ResetStuckSessionsOnRestart :many
UPDATE acp_sessions
SET status = 'crashed',
@@ -332,6 +721,178 @@ func (q *Queries) ResetStuckSessionsOnRestart(ctx context.Context) ([]ResetStuck
return items, nil
}
const sumProjectCostUSD = `-- name: SumProjectCostUSD :one
SELECT COALESCE(SUM(cost_usd), 0)::numeric AS cost_usd
FROM acp_session_usage
WHERE project_id = $1 AND created_at >= $2
`
type SumProjectCostUSDParams struct {
ProjectID pgtype.UUID `json:"project_id"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
// 某 project 自 since 起的 ACP 累计花费(per-project 软预算)。
func (q *Queries) SumProjectCostUSD(ctx context.Context, arg SumProjectCostUSDParams) (pgtype.Numeric, error) {
row := q.db.QueryRow(ctx, sumProjectCostUSD, arg.ProjectID, arg.CreatedAt)
var cost_usd pgtype.Numeric
err := row.Scan(&cost_usd)
return cost_usd, err
}
const summarizeAcpUsageByDay = `-- name: SummarizeAcpUsageByDay :many
SELECT date_trunc('day', created_at)::timestamptz AS day,
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(thinking_tokens)::bigint AS thinking_tokens,
SUM(cost_usd)::numeric AS cost_usd
FROM acp_session_usage
WHERE created_at >= $1 AND created_at < $2
GROUP BY day
ORDER BY day
`
type SummarizeAcpUsageByDayParams struct {
CreatedAt pgtype.Timestamptz `json:"created_at"`
CreatedAt_2 pgtype.Timestamptz `json:"created_at_2"`
}
type SummarizeAcpUsageByDayRow struct {
Day pgtype.Timestamptz `json:"day"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
CostUsd pgtype.Numeric `json:"cost_usd"`
}
func (q *Queries) SummarizeAcpUsageByDay(ctx context.Context, arg SummarizeAcpUsageByDayParams) ([]SummarizeAcpUsageByDayRow, error) {
rows, err := q.db.Query(ctx, summarizeAcpUsageByDay, arg.CreatedAt, arg.CreatedAt_2)
if err != nil {
return nil, err
}
defer rows.Close()
var items []SummarizeAcpUsageByDayRow
for rows.Next() {
var i SummarizeAcpUsageByDayRow
if err := rows.Scan(
&i.Day,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.CostUsd,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const summarizeAcpUsageByModel = `-- name: SummarizeAcpUsageByModel :many
SELECT model_id,
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(thinking_tokens)::bigint AS thinking_tokens,
SUM(cost_usd)::numeric AS cost_usd
FROM acp_session_usage
WHERE created_at >= $1 AND created_at < $2
GROUP BY model_id
ORDER BY cost_usd DESC
`
type SummarizeAcpUsageByModelParams struct {
CreatedAt pgtype.Timestamptz `json:"created_at"`
CreatedAt_2 pgtype.Timestamptz `json:"created_at_2"`
}
type SummarizeAcpUsageByModelRow struct {
ModelID pgtype.UUID `json:"model_id"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
CostUsd pgtype.Numeric `json:"cost_usd"`
}
func (q *Queries) SummarizeAcpUsageByModel(ctx context.Context, arg SummarizeAcpUsageByModelParams) ([]SummarizeAcpUsageByModelRow, error) {
rows, err := q.db.Query(ctx, summarizeAcpUsageByModel, arg.CreatedAt, arg.CreatedAt_2)
if err != nil {
return nil, err
}
defer rows.Close()
var items []SummarizeAcpUsageByModelRow
for rows.Next() {
var i SummarizeAcpUsageByModelRow
if err := rows.Scan(
&i.ModelID,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.CostUsd,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const summarizeAcpUsageByUser = `-- name: SummarizeAcpUsageByUser :many
SELECT user_id,
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(thinking_tokens)::bigint AS thinking_tokens,
SUM(cost_usd)::numeric AS cost_usd
FROM acp_session_usage
WHERE created_at >= $1 AND created_at < $2
GROUP BY user_id
ORDER BY cost_usd DESC
`
type SummarizeAcpUsageByUserParams struct {
CreatedAt pgtype.Timestamptz `json:"created_at"`
CreatedAt_2 pgtype.Timestamptz `json:"created_at_2"`
}
type SummarizeAcpUsageByUserRow struct {
UserID pgtype.UUID `json:"user_id"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
ThinkingTokens int64 `json:"thinking_tokens"`
CostUsd pgtype.Numeric `json:"cost_usd"`
}
func (q *Queries) SummarizeAcpUsageByUser(ctx context.Context, arg SummarizeAcpUsageByUserParams) ([]SummarizeAcpUsageByUserRow, error) {
rows, err := q.db.Query(ctx, summarizeAcpUsageByUser, arg.CreatedAt, arg.CreatedAt_2)
if err != nil {
return nil, err
}
defer rows.Close()
var items []SummarizeAcpUsageByUserRow
for rows.Next() {
var i SummarizeAcpUsageByUserRow
if err := rows.Scan(
&i.UserID,
&i.PromptTokens,
&i.CompletionTokens,
&i.ThinkingTokens,
&i.CostUsd,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const updateSessionFinished = `-- name: UpdateSessionFinished :exec
UPDATE acp_sessions
SET status = $2, exit_code = $3, last_error = $4, ended_at = now()
@@ -355,6 +916,22 @@ func (q *Queries) UpdateSessionFinished(ctx context.Context, arg UpdateSessionFi
return err
}
const updateSessionLastStopReason = `-- name: UpdateSessionLastStopReason :exec
UPDATE acp_sessions
SET last_stop_reason = $2
WHERE id = $1
`
type UpdateSessionLastStopReasonParams struct {
ID pgtype.UUID `json:"id"`
LastStopReason *string `json:"last_stop_reason"`
}
func (q *Queries) UpdateSessionLastStopReason(ctx context.Context, arg UpdateSessionLastStopReasonParams) error {
_, err := q.db.Exec(ctx, updateSessionLastStopReason, arg.ID, arg.LastStopReason)
return err
}
const updateSessionRunning = `-- name: UpdateSessionRunning :exec
UPDATE acp_sessions
SET status = 'running', agent_session_id = $2, pid = $3
@@ -371,3 +948,20 @@ func (q *Queries) UpdateSessionRunning(ctx context.Context, arg UpdateSessionRun
_, err := q.db.Exec(ctx, updateSessionRunning, arg.ID, arg.AgentSessionID, arg.Pid)
return err
}
const updateSessionSandboxMode = `-- name: UpdateSessionSandboxMode :exec
UPDATE acp_sessions
SET sandbox_mode = $2
WHERE id = $1
`
type UpdateSessionSandboxModeParams struct {
ID pgtype.UUID `json:"id"`
SandboxMode *string `json:"sandbox_mode"`
}
// 记录本 session 运行所用的沙箱模式(审计/取证)。
func (q *Queries) UpdateSessionSandboxMode(ctx context.Context, arg UpdateSessionSandboxModeParams) error {
_, err := q.db.Exec(ctx, updateSessionSandboxMode, arg.ID, arg.SandboxMode)
return err
}
+180
View File
@@ -0,0 +1,180 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.31.1
// source: turns.sql
package acpsqlc
import (
"context"
"github.com/jackc/pgx/v5/pgtype"
)
const getLatestTurnBySession = `-- name: GetLatestTurnBySession :one
SELECT id, session_id, turn_index, prompt_request_id, status,
stop_reason, update_count, started_at, completed_at
FROM acp_turns
WHERE session_id = $1
ORDER BY turn_index DESC
LIMIT 1
`
func (q *Queries) GetLatestTurnBySession(ctx context.Context, sessionID pgtype.UUID) (AcpTurn, error) {
row := q.db.QueryRow(ctx, getLatestTurnBySession, sessionID)
var i AcpTurn
err := row.Scan(
&i.ID,
&i.SessionID,
&i.TurnIndex,
&i.PromptRequestID,
&i.Status,
&i.StopReason,
&i.UpdateCount,
&i.StartedAt,
&i.CompletedAt,
)
return i, err
}
const incrementTurnUpdateCount = `-- name: IncrementTurnUpdateCount :exec
UPDATE acp_turns
SET update_count = update_count + 1
WHERE session_id = $1 AND turn_index = $2 AND status = 'in_progress'
`
type IncrementTurnUpdateCountParams struct {
SessionID pgtype.UUID `json:"session_id"`
TurnIndex int32 `json:"turn_index"`
}
func (q *Queries) IncrementTurnUpdateCount(ctx context.Context, arg IncrementTurnUpdateCountParams) error {
_, err := q.db.Exec(ctx, incrementTurnUpdateCount, arg.SessionID, arg.TurnIndex)
return err
}
const insertTurn = `-- name: InsertTurn :one
INSERT INTO acp_turns (
session_id, turn_index, prompt_request_id, status
) VALUES ($1, $2, $3, $4)
RETURNING id, session_id, turn_index, prompt_request_id, status,
stop_reason, update_count, started_at, completed_at
`
type InsertTurnParams struct {
SessionID pgtype.UUID `json:"session_id"`
TurnIndex int32 `json:"turn_index"`
PromptRequestID *string `json:"prompt_request_id"`
Status string `json:"status"`
}
func (q *Queries) InsertTurn(ctx context.Context, arg InsertTurnParams) (AcpTurn, error) {
row := q.db.QueryRow(ctx, insertTurn,
arg.SessionID,
arg.TurnIndex,
arg.PromptRequestID,
arg.Status,
)
var i AcpTurn
err := row.Scan(
&i.ID,
&i.SessionID,
&i.TurnIndex,
&i.PromptRequestID,
&i.Status,
&i.StopReason,
&i.UpdateCount,
&i.StartedAt,
&i.CompletedAt,
)
return i, err
}
const listTurnsBySession = `-- name: ListTurnsBySession :many
SELECT id, session_id, turn_index, prompt_request_id, status,
stop_reason, update_count, started_at, completed_at
FROM acp_turns
WHERE session_id = $1
ORDER BY turn_index ASC
`
func (q *Queries) ListTurnsBySession(ctx context.Context, sessionID pgtype.UUID) ([]AcpTurn, error) {
rows, err := q.db.Query(ctx, listTurnsBySession, sessionID)
if err != nil {
return nil, err
}
defer rows.Close()
var items []AcpTurn
for rows.Next() {
var i AcpTurn
if err := rows.Scan(
&i.ID,
&i.SessionID,
&i.TurnIndex,
&i.PromptRequestID,
&i.Status,
&i.StopReason,
&i.UpdateCount,
&i.StartedAt,
&i.CompletedAt,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const markTurnAborted = `-- name: MarkTurnAborted :exec
UPDATE acp_turns
SET status = 'aborted', completed_at = now()
WHERE session_id = $1 AND turn_index = $2 AND status = 'in_progress'
`
type MarkTurnAbortedParams struct {
SessionID pgtype.UUID `json:"session_id"`
TurnIndex int32 `json:"turn_index"`
}
func (q *Queries) MarkTurnAborted(ctx context.Context, arg MarkTurnAbortedParams) error {
_, err := q.db.Exec(ctx, markTurnAborted, arg.SessionID, arg.TurnIndex)
return err
}
const markTurnCompleted = `-- name: MarkTurnCompleted :exec
UPDATE acp_turns
SET status = 'completed', stop_reason = $3, update_count = $4, completed_at = now()
WHERE session_id = $1 AND turn_index = $2 AND status = 'in_progress'
`
type MarkTurnCompletedParams struct {
SessionID pgtype.UUID `json:"session_id"`
TurnIndex int32 `json:"turn_index"`
StopReason *string `json:"stop_reason"`
UpdateCount int32 `json:"update_count"`
}
func (q *Queries) MarkTurnCompleted(ctx context.Context, arg MarkTurnCompletedParams) error {
_, err := q.db.Exec(ctx, markTurnCompleted,
arg.SessionID,
arg.TurnIndex,
arg.StopReason,
arg.UpdateCount,
)
return err
}
const purgeTurnsBefore = `-- name: PurgeTurnsBefore :execrows
DELETE FROM acp_turns WHERE started_at < $1
`
func (q *Queries) PurgeTurnsBefore(ctx context.Context, startedAt pgtype.Timestamptz) (int64, error) {
result, err := q.db.Exec(ctx, purgeTurnsBefore, startedAt)
if err != nil {
return 0, err
}
return result.RowsAffected(), nil
}
+283 -33
View File
@@ -17,6 +17,7 @@ import (
"log/slog"
"os"
"os/exec"
"path/filepath"
"sync"
"sync/atomic"
"time"
@@ -24,6 +25,7 @@ import (
"github.com/google/uuid"
"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
"github.com/yan1h/agent-coding-workflow/internal/acp/sandbox"
"github.com/yan1h/agent-coding-workflow/internal/audit"
"github.com/yan1h/agent-coding-workflow/internal/infra/crypto"
"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
@@ -44,6 +46,25 @@ type SupervisorConfig struct {
// AgentHomesDir 是 agent CLI 受管 home 的根目录(<DataDir>/agent-homes)。
// 空串时跳过配置物化与重定向 env 注入(部分测试)。
AgentHomesDir string
// DefaultProjectBudgetUSD 是 per-project ACP 累计花费的全局软上限(0 = 不限)。
// 当 session 与 agent-kind 均未设置成本上限时,用作 fallback 的预算判定基线。
DefaultProjectBudgetUSD float64
// ===== 沙箱(per-session 隔离)=====
// Sandbox 是 pluggable 沙箱实现(none|uid|container),在 cmd/env 构建后、
// proc.Group.Prepare 之前 Apply。为 nil 时等同 mode=none(无隔离)。
Sandbox sandbox.Sandbox
// SandboxMode 记录当前 mode,用于落库 acp_sessions.sandbox_mode(审计/取证)。
SandboxMode sandbox.Mode
// SandboxBaseUID 是 per-session 低权 UID 的基址:实际 uid = BaseUID + offset,
// offset 由 session 派生(保证不同活跃 session 不复用同一 uid)。0 = 不分配。
SandboxBaseUID int
// DataRoot 是要被屏蔽的 /data 根(除本 session worktree + per-user home 外)。
DataRoot string
// EgressProxyURL 是注入子进程 env 的 HTTP(S)_PROXY(egress 白名单代理)。
EgressProxyURL string
// SandboxRlimits 是套用到子进程(及其进程树)的 OS 资源上限。
SandboxRlimits sandbox.Limits
}
// Supervisor 是 ACP 子进程总管。
@@ -51,22 +72,87 @@ type Supervisor struct {
mu sync.RWMutex
procs map[uuid.UUID]*Process
repo Repository
audit audit.Recorder
notify *notify.Dispatcher
wtSvc workspace.WorktreeService
wsRepo workspace.Repository
mcpTokens mcp.TokenService // spec §6.3 — onExit revokes system token
crypto *crypto.Encryptor
permSvc *PermissionService // 工具调用权限审批;可为 nil(部分测试)
cfg SupervisorConfig
log *slog.Logger
repo Repository
audit audit.Recorder
notify *notify.Dispatcher
wtSvc workspace.WorktreeService
wsRepo workspace.Repository
mcpTokens mcp.TokenService // spec §6.3 — onExit revokes system token
crypto *crypto.Encryptor
permSvc *PermissionService // 工具调用权限审批;可为 nil(部分测试)
turnBus TurnBus // 回合完成事件总线;可为 nil(部分测试)
orchRedrive OrchestratorRedrive // 崩溃的编排器 session 再驱动回调;可为 nil
prices ModelPriceLookup // model 价格查询;可为 nil(无 model 价格时成本=0/agent 自报)
agentsMD AgentsMDRenderer // AGENTS.md 渲染回调(project memory);可为 nil
metrics SessionExitRecorder // 会话退出计数;可为 nil(metrics 关闭)
cfg SupervisorConfig
log *slog.Logger
}
// SessionExitRecorder 是 Supervisor 上报会话退出指标所需的窄接口。
// metrics.Metrics 满足它;nil 时不记录。
type SessionExitRecorder interface {
RecordSessionExit(status string)
}
// SetMetrics 注入会话退出计数器(装配期)。
func (s *Supervisor) SetMetrics(m SessionExitRecorder) { s.metrics = m }
// AgentsMDRenderer renders the AGENTS.md content for a (project, workspace) from
// project memory. Injected via SetAgentsMDRenderer to avoid an acp →
// projectmemory import cycle. May be nil (no seeding).
type AgentsMDRenderer func(ctx context.Context, projectID uuid.UUID, wsID *uuid.UUID) (string, error)
// SetAgentsMDRenderer 注入 AGENTS.md 渲染回调(装配期)。
func (s *Supervisor) SetAgentsMDRenderer(fn AgentsMDRenderer) { s.agentsMD = fn }
// agentsMDFileName is the generated file written into the worktree root.
const agentsMDFileName = "AGENTS.md"
// seedAgentsMD renders project memory into <cwd>/AGENTS.md before spawn. Skipped
// when no renderer is configured, CwdPath is empty, or ProjectID is nil. Failures
// are logged, never fatal — the agent still spawns without seeded knowledge.
func (s *Supervisor) seedAgentsMD(ctx context.Context, sess *Session) {
if s.agentsMD == nil || sess.CwdPath == "" || sess.ProjectID == uuid.Nil {
return
}
var wsID *uuid.UUID
if sess.WorkspaceID != uuid.Nil {
id := sess.WorkspaceID
wsID = &id
}
content, err := s.agentsMD(ctx, sess.ProjectID, wsID)
if err != nil {
s.log.Warn("acp.agents_md.render_failed", "session_id", sess.ID, "err", err.Error())
return
}
dst := filepath.Join(sess.CwdPath, agentsMDFileName)
if err := os.WriteFile(dst, []byte(content), 0o644); err != nil {
s.log.Warn("acp.agents_md.write_failed", "session_id", sess.ID, "path", dst, "err", err.Error())
}
}
// OrchestratorRedrive 在一个由编排器 step 驱动的 session 崩溃退出时被调用,由编排器
// 据此把对应 step 重新入队(带 backoff),让回合在同一 worktree 上恢复。注入为 func
// 以避免 acp → orchestrator 的 import 循环(与 notify dispatcher / permission service
// 的注入方式一致)。实现必须是非阻塞、不持 supervisor 锁的纯 DB 写(jobs Enqueue)。
type OrchestratorRedrive func(ctx context.Context, stepID uuid.UUID, reason string)
// SetOrchestratorRedrive 回填编排器再驱动回调(装配期)。
func (s *Supervisor) SetOrchestratorRedrive(fn OrchestratorRedrive) { s.orchRedrive = fn }
// SetPermissionService 注入权限审批服务。与 Supervisor 互相依赖,故装配时回填
// (permSvc 需要 Supervisor 作为 RelayLocator,Supervisor 需要 permSvc 注入 handler)。
func (s *Supervisor) SetPermissionService(p *PermissionService) { s.permSvc = p }
// SetTurnBus 注入回合事件总线。Spawn 时每个 relay 构造一个 TurnTracker 复用它。
// 可为 nil——此时回合仍落库(若 repo 可用),但不发布内部事件。
func (s *Supervisor) SetTurnBus(b TurnBus) { s.turnBus = b }
// SetModelPriceLookup 注入 model 价格查询(成本核算用)。装配期回填,避免 acp →
// chat 的构造期依赖。可为 nil——此时无 model 价格的会话成本为 0 或取 agent 自报。
func (s *Supervisor) SetModelPriceLookup(p ModelPriceLookup) { s.prices = p }
// NewSupervisor 构造空 Supervisor。Start 不需要 — supervisor 是被动的(spawn 时
// 起 goroutine,无主循环)。Stop 在 ShutdownAll 中实现。
func NewSupervisor(repo Repository, rec audit.Recorder, disp *notify.Dispatcher,
@@ -114,6 +200,10 @@ type Process struct {
// 返回后立即取消),故 relay 的 reader/writer/persist/permission Decide 在整个
// 会话生命周期内都使用未取消的 ctx。onExit 调用 relayCancel 完成 relay teardown。
relayCancel context.CancelFunc
// sandboxCleanup 拆除沙箱临时资源(卸载 bind mount、停止容器、删临时目录)。
// 由 monitorProcess 在 group.Close() 之后调用,仅调用一次。可为 nil(mode=none)。
sandboxCleanup func()
}
// stderrRing 是固定行数 ring buffer,monitor 退出时取 tail 写 last_error。
@@ -163,14 +253,20 @@ func (s *Supervisor) GetRelay(sid uuid.UUID) *Relay {
// handshake 失败时调用 Kill 回滚并返回错误。调用方(SessionService)负责前置的
// worktree 占用与 acp_sessions 行 INSERT;本方法仅负责 OS 层资源与协议层握手。
func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind, extraEnv map[string]string) (*Process, error) {
// AGENTS.md seeding:在 spawn 前把 project memory 渲染成 worktree 根的
// AGENTS.md,让 agent 启动即有结构化项目知识。best-effort,失败仅告警不阻断。
s.seedAgentsMD(ctx, sess)
// 受管 home:物化 DB 中的配置文件并注入重定向 env(优先级最低,可被
// AgentKind env / extraEnv 覆盖)。AgentHomesDir 未配置时跳过(部分测试)。
envMap := map[string]string{}
var agentHome string
if s.cfg.AgentHomesDir != "" {
home, err := s.materializeAgentHome(ctx, kind)
home, err := s.materializeAgentHome(ctx, sess, kind)
if err != nil {
return nil, fmt.Errorf("materialize agent home: %w", err)
}
agentHome = home
envMap = agentHomeEnv(home, kind.ClientType)
}
@@ -197,6 +293,28 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
cmd.Dir = sess.CwdPath
cmd.Env = osEnv
// 沙箱:在 cmd/env 构建完成后、proc.Group.Prepare(Setpgid)之前 Apply。
// Apply 只 ADD 到 cmd.SysProcAttr(Credential / namespace),绝不替换,
// 故不破坏后续 group.Prepare 的 Setpgid 与负 pgid 整树终止。返回的 cleanup
// 在 monitorProcess 的 group.Close() 之后调用一次。mode=none 时为 no-op。
sandboxCleanup := func() {}
if s.cfg.Sandbox != nil {
spec := s.buildSandboxSpec(sess, agentHome)
cleanup, serr := s.cfg.Sandbox.Apply(cmd, spec)
if serr != nil {
return nil, fmt.Errorf("sandbox apply: %w", serr)
}
if cleanup != nil {
sandboxCleanup = cleanup
}
// 记录本 session 运行所用的 sandbox mode(审计/取证)。best-effort。
if s.cfg.SandboxMode != "" {
if err := s.repo.UpdateSessionSandboxMode(ctx, sess.ID, string(s.cfg.SandboxMode)); err != nil {
s.log.Warn("acp.sandbox.record_mode", "session_id", sess.ID, "err", err.Error())
}
}
}
// 进程树管理:Prepare 须在 Start 前(Unix Setpgid),Adopt 须在 Start 后
// (Windows Job Object)。终止时整树清理,避免 agent 派生的子孙进程残留。
group := procgrp.NewGroup()
@@ -243,6 +361,50 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
s.log,
)
// 回合追踪器:被动解析 session/update + session/prompt 响应。复用同一 acpRepo
// 与共享 TurnBus,无需额外连接池。onExit 通过 relay.Tracker().Abort 闭合悬挂回合。
sessCtx := handlers.SessionContext{
SessionID: sess.ID,
WorkspaceID: sess.WorkspaceID,
UserID: sess.UserID,
CwdPath: sess.CwdPath,
AgentSessionID: sess.AgentSessionID,
ToolAllowlist: kind.ToolAllowlist,
}
relay.SetTracker(NewTurnTracker(s.repo, s.turnBus, sessCtx, s.log))
// 成本核算:构造 per-session 累加器 + usage sink,seeded with 会话快照的有效
// 预算上限与 agent-kind 的 model 价格来源。预算突破时以 detached goroutine 标记
// terminated_reason 后 Kill(绝不阻塞 relay.reader,且 Kill 幂等)。
caps := BudgetCaps{
MaxCostUSD: sess.BudgetMaxCostUSD,
MaxTokens: sess.BudgetMaxTokens,
MaxWallClockSeconds: sess.BudgetMaxWallClockSeconds,
}
acc := newUsageAccumulator(AccumulatorParams{
Repo: s.repo,
Prices: s.prices,
Log: s.log,
SessionID: sess.ID,
UserID: sess.UserID,
ProjectID: sess.ProjectID,
AgentKindID: sess.AgentKindID,
ModelID: kind.ModelID,
StartedAt: time.Now(),
Caps: caps,
ProjectBudgetUSD: s.cfg.DefaultProjectBudgetUSD,
})
sid := sess.ID
killForBudget := func(reason BreachReason) {
bgCtx := context.Background()
if err := s.repo.MarkSessionTerminatedReason(bgCtx, sid, string(reason)); err != nil {
s.log.Error("acp.budget.mark_terminated", "session_id", sid, "err", err.Error())
}
s.log.Warn("acp.budget.breach_kill", "session_id", sid, "reason", string(reason))
s.Kill(bgCtx, sid, s.cfg.KillGrace)
}
relay.SetUsageSink(newAccumulatorSink(acc, killForBudget, s.log))
// relay 的会话级 context:独立于调用方的 spawn ctx(生产调用方在 Spawn 返回后
// 立即 cancel spawn ctx)。relay 的 reader/writer/persist/permission Decide 必须
// 在整个会话生命周期内使用未取消的 ctx,否则 prompt 被丢、事件停止持久化、权限
@@ -250,20 +412,21 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
relayCtx, relayCancel := context.WithCancel(context.Background())
proc := &Process{
SessionID: sess.ID,
WorkspaceID: sess.WorkspaceID,
UserID: sess.UserID,
IsMain: sess.IsMainWorktree,
Cmd: cmd,
group: group,
Stdin: stdin,
Stdout: stdout,
Stderr: stderr,
Relay: relay,
StderrBuf: newStderrRing(s.cfg.StderrBufferLines),
StartedAt: time.Now(),
done: make(chan struct{}),
relayCancel: relayCancel,
SessionID: sess.ID,
WorkspaceID: sess.WorkspaceID,
UserID: sess.UserID,
IsMain: sess.IsMainWorktree,
Cmd: cmd,
group: group,
Stdin: stdin,
Stdout: stdout,
Stderr: stderr,
Relay: relay,
StderrBuf: newStderrRing(s.cfg.StderrBufferLines),
StartedAt: time.Now(),
done: make(chan struct{}),
relayCancel: relayCancel,
sandboxCleanup: sandboxCleanup,
}
s.mu.Lock()
@@ -272,14 +435,7 @@ func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind,
go s.drainStderr(proc)
go s.monitorProcess(proc)
go relay.Run(relayCtx, handlers.SessionContext{
SessionID: sess.ID,
WorkspaceID: sess.WorkspaceID,
UserID: sess.UserID,
CwdPath: sess.CwdPath,
AgentSessionID: sess.AgentSessionID,
ToolAllowlist: kind.ToolAllowlist,
})
go relay.Run(relayCtx, sessCtx)
// 握手仍走调用方的 spawn ctx,保留 SpawnTimeout 截止语义。
if err := s.handshake(ctx, sess, kind, relay, proc, extraEnv); err != nil {
@@ -354,6 +510,30 @@ func (s *Supervisor) BuildSpawnEnv(kind *AgentKind, extraEnv map[string]string)
return out, nil
}
// buildSandboxSpec 从 session + 受管 home 派生本次 spawn 的沙箱 Spec。
// per-session UID = BaseUID + offset,offset 由 session ID 派生(低 16 位),
// 保证不同 session 大概率落在不同 uid(碰撞仅影响隔离粒度,不影响正确性)。
// HomeDir = 受管 home(per-user),WorktreeDir = session cwd,二者为唯一可写路径。
func (s *Supervisor) buildSandboxSpec(sess *Session, agentHome string) sandbox.Spec {
uid := 0
if s.cfg.SandboxBaseUID > 0 {
// 取 session ID 的低 16 位作为 offset,限制 uid 漂移范围。
b := sess.ID
offset := int(b[14])<<8 | int(b[15])
uid = s.cfg.SandboxBaseUID + offset
}
return sandbox.Spec{
Mode: s.cfg.SandboxMode,
UID: uid,
GID: uid,
HomeDir: agentHome,
WorktreeDir: sess.CwdPath,
DataRoot: s.cfg.DataRoot,
Rlimits: s.cfg.SandboxRlimits,
ProxyURL: s.cfg.EgressProxyURL,
}
}
// Kill 终止指定 session 的子进程。在 Unix 平台上先发 SIGTERM 等待 grace 时间,
// 仍未退出再 Process.Kill;Windows 没有可靠 SIGTERM 实现,直接 Process.Kill。
// 始终阻塞直到 monitorProcess 关闭 done channel。Kill 是幂等的——目标 session
@@ -440,6 +620,12 @@ func (s *Supervisor) monitorProcess(proc *Process) {
// 进程已 Wait 返回,此处 Close 不影响 onExit-must-not-relock 不变量。
proc.group.Close()
// 沙箱清理:在 group.Close() 之后调用一次(卸载 bind mount、停容器、删临时目录)。
// 不持 supervisor.mu,遵守 onExit-must-not-relock 不变量。mode=none 时为 no-op。
if proc.sandboxCleanup != nil {
proc.sandboxCleanup()
}
status := SessionExited
if !proc.KilledByUs.Load() {
status = SessionCrashed
@@ -451,6 +637,9 @@ func (s *Supervisor) monitorProcess(proc *Process) {
// 不能再 acquire s.mu(monitorProcess 已 delete map 后调用本函数,但 Kill
// 也持锁等 done,会形成 monitor → Kill → onExit 的链式 deadlock)。
func (s *Supervisor) onExit(ctx context.Context, proc *Process, status SessionStatus, exitCode int32, waitErr error) {
if s.metrics != nil {
s.metrics.RecordSessionExit(string(status))
}
stderrTail := proc.StderrBuf.Tail(s.cfg.StderrTailBytes)
var lastErr *string
if status == SessionCrashed && stderrTail != "" {
@@ -460,6 +649,14 @@ func (s *Supervisor) onExit(ctx context.Context, proc *Process, status SessionSt
s.log.Error("acp.on_exit.update_session", "session_id", proc.SessionID, "err", err.Error())
}
// 成本/预算终止:若 terminated_reason 已被预算门或 reaper 写入,发 audit + notify。
// 读 DB(UpdateSessionFinished 之后)拿最新 terminated_reason;不覆盖既有值。
if finished, gerr := s.repo.GetSessionByID(ctx, proc.SessionID); gerr == nil && finished != nil &&
finished.TerminatedReason != nil && *finished.TerminatedReason != "" {
s.recordBudgetTerminated(ctx, proc.UserID, proc.SessionID, *finished.TerminatedReason,
finished.TotalCostUSD, finished.PromptTokens+finished.CompletionTokens+finished.ThinkingTokens)
}
// Release worktree。主 worktree 走 acp 表自己的 CAS;子 worktree 按 holder
// 释放(holder = "session:<sid>",见 workspace.Caller.Holder),与启动 reaper
// 一致,不需要知道 worktree ID。
@@ -530,6 +727,17 @@ func (s *Supervisor) onExit(ctx context.Context, proc *Process, status SessionSt
}
}
// 回合检测:闭合任何开放的 in_progress 回合并发布 session_idle,确保崩溃/被杀时
// 编排层也能收到空闲信号。Abort 只触碰 repo+bus(不持 supervisor.mu),且必须在
// Relay.Close 之前——Close 后 tracker 仍可用,但语义上回合应在 teardown 前闭合。
if tracker := proc.Relay.Tracker(); tracker != nil {
reason := string(StopCancelled)
if status == SessionCrashed {
reason = "crashed"
}
tracker.Abort(ctx, reason)
}
// 取消 relay 的会话级 context:停止 reader/writer,并让仍阻塞在 Decide / 已派生
// 的 handleAgentRequest goroutine 收到 ctx.Done 后退出。在 Relay.Close 之前调,
// 二者共同完成 teardown(relayCancel 停 ctx 派生的 goroutine,Close 关 r.done +
@@ -540,4 +748,46 @@ func (s *Supervisor) onExit(ctx context.Context, proc *Process, status SessionSt
// 最后关 relay:广播 session_terminated 给 WS subscribers。
proc.Relay.Close(string(status), &exitCode)
// 编排器再驱动:若该 session 由编排器某 step 驱动且本次为崩溃退出,把对应 step
// 重新入队(带 backoff),让回合在同一 worktree 上恢复。只做 DB 读 + jobs Enqueue,
// 不持 supervisor 锁(保持 onExit-must-not-relock 不变量)。
if status == SessionCrashed && s.orchRedrive != nil {
if sess, gerr := s.repo.GetSessionByID(ctx, proc.SessionID); gerr == nil &&
sess != nil && sess.OrchestratorStepID != nil {
s.orchRedrive(ctx, *sess.OrchestratorStepID, "session_crashed")
}
}
}
// recordBudgetTerminated 在会话因预算/reaper 被强制终止时发 audit + owner notify。
func (s *Supervisor) recordBudgetTerminated(ctx context.Context, userID, sessionID uuid.UUID, reason string, totalCost float64, totalTokens int64) {
if s.audit != nil {
uid := userID
_ = s.audit.Record(ctx, audit.Entry{
UserID: &uid,
Action: "acp.session.budget_exceeded",
TargetType: "acp_session",
TargetID: sessionID.String(),
Metadata: map[string]any{
"reason": reason,
"total_cost_usd": totalCost,
"total_tokens": totalTokens,
},
})
}
if s.notify != nil {
_ = s.notify.Dispatch(ctx, notify.Message{
UserID: userID,
Topic: "acp.session_budget_exceeded",
Severity: notify.SeverityWarning,
Title: "ACP session terminated (budget/reaper)",
Body: fmt.Sprintf("Session %s was terminated: %s.", sessionID, reason),
Metadata: map[string]any{
"session_id": sessionID.String(),
"reason": reason,
"total_cost_usd": totalCost,
},
})
}
}
+242
View File
@@ -0,0 +1,242 @@
// turn.go 实现 TurnTracker:单 relay 维度的回合状态机,外加 stopReason 解析辅助。
//
// 回合(turn)定义:一次 session/prompt 请求到其响应之间的窗口。relay reader 在解析到
// - session/update 通知 → OnUpdate(累加 update_count,内存计数,完成时一次性写库)
// - 该 prompt 的字符串-id 响应 → OnPromptResponse(标完成 + 写 last_stop_reason + 发事件)
//
// MVP 假设:同一 session 同时只有一个开放回合。若在已有开放回合时收到新的 StartTurn,
// 自动中止前一个(auto-abort),避免悬挂。所有状态由 mu 保护;Abort 只触碰 repo+bus,
// 不持有 supervisor 锁,故可安全地从 onExit 调用(不违反 onExit-must-not-relock 不变量)。
package acp
import (
"context"
"encoding/json"
"log/slog"
"sync"
"time"
"github.com/google/uuid"
"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
)
// TurnRepo 是 TurnTracker 需要的最小 repo 子集,由 acp.Repository 满足。
type TurnRepo interface {
InsertTurn(ctx context.Context, t *Turn) (*Turn, error)
MarkTurnCompleted(ctx context.Context, sessionID uuid.UUID, turnIndex int, stop string, updateCount int) error
MarkTurnAborted(ctx context.Context, sessionID uuid.UUID, turnIndex int) error
GetLatestTurnBySession(ctx context.Context, sessionID uuid.UUID) (*Turn, error)
UpdateSessionLastStopReason(ctx context.Context, sessionID uuid.UUID, reason string) error
}
// TurnTracker 是单 relay 的回合状态机。
type TurnTracker struct {
repo TurnRepo
bus TurnBus
sess handlers.SessionContext
log *slog.Logger
mu sync.Mutex
open bool // 是否有开放回合
turnIndex int // 当前/最近回合的 index
promptID string // 当前开放回合的 prompt 请求 id(用于响应相关联)
updateCount int // 当前回合的 session/update 计数(内存累计,完成时落库)
nextIndex int // 下一个回合的 index(构造时由 GetLatestTurnBySession 播种)
seeded bool // nextIndex 是否已播种
}
// NewTurnTracker 构造一个 TurnTracker。bus / repo 为 nil 时方法均为安全 no-op
// (部分测试可不注入)。
func NewTurnTracker(repo TurnRepo, bus TurnBus, sess handlers.SessionContext, log *slog.Logger) *TurnTracker {
if log == nil {
log = slog.Default()
}
return &TurnTracker{repo: repo, bus: bus, sess: sess, log: log}
}
// seedLocked 懒播种 nextIndex:从 DB 取最近回合 index+1。失败时从 0 开始(best-effort)。
// 调用方必须已持有 t.mu。
func (t *TurnTracker) seedLocked(ctx context.Context) {
if t.seeded {
return
}
t.seeded = true
if t.repo == nil {
return
}
latest, err := t.repo.GetLatestTurnBySession(ctx, t.sess.SessionID)
if err != nil {
t.log.Warn("acp.turn.seed_index", "session_id", t.sess.SessionID, "err", err.Error())
return
}
if latest != nil {
t.nextIndex = latest.TurnIndex + 1
}
}
// StartTurn 在发送 session/prompt 之前调用,登记一个 in_progress 回合并把 promptRequestID
// 与之关联。若已有开放回合,先 auto-abort 之(MVP 单回合假设)。返回新回合的 index。
func (t *TurnTracker) StartTurn(ctx context.Context, promptRequestID string) (int, error) {
if t == nil {
return 0, nil
}
t.mu.Lock()
defer t.mu.Unlock()
t.seedLocked(ctx)
// 已有开放回合 → auto-abort,避免悬挂。
if t.open {
t.abortLocked(ctx, string(StopCancelled))
}
idx := t.nextIndex
if t.repo != nil {
var pid *string
if promptRequestID != "" {
p := promptRequestID
pid = &p
}
if _, err := t.repo.InsertTurn(ctx, &Turn{
SessionID: t.sess.SessionID,
TurnIndex: idx,
PromptRequestID: pid,
Status: TurnInProgress,
}); err != nil {
return 0, err
}
}
t.open = true
t.turnIndex = idx
t.promptID = promptRequestID
t.updateCount = 0
t.nextIndex = idx + 1
return idx, nil
}
// OnUpdate 在收到一条 session/update 通知时调用,累加当前开放回合的 update_count
// (内存计数,避免每 chunk 一次 UPDATE;完成时一次性落库)。无开放回合时 no-op。
func (t *TurnTracker) OnUpdate(_ context.Context) {
if t == nil {
return
}
t.mu.Lock()
defer t.mu.Unlock()
if t.open {
t.updateCount++
}
}
// IsTrackedPrompt 报告 id 是否为当前开放回合关联的 prompt 请求 id。relay reader 用它
// 判断一条字符串-id 响应是否应触发 OnPromptResponse(而非沿用既有 fanout 路径)。
func (t *TurnTracker) IsTrackedPrompt(id string) bool {
if t == nil || id == "" {
return false
}
t.mu.Lock()
defer t.mu.Unlock()
return t.open && t.promptID == id
}
// OnPromptResponse 在收到匹配当前回合的 session/prompt 响应时调用:标完成、写
// session.last_stop_reason、发布 turn_completed 与 session_idle 事件。id 不匹配时 no-op。
func (t *TurnTracker) OnPromptResponse(ctx context.Context, promptRequestID string, rawStopReason string) {
if t == nil {
return
}
t.mu.Lock()
if !t.open || t.promptID != promptRequestID {
t.mu.Unlock()
return
}
idx := t.turnIndex
updateCount := t.updateCount
t.open = false
t.promptID = ""
t.mu.Unlock()
norm := NormalizeStopReason(rawStopReason)
completedAt := time.Now()
if t.repo != nil {
if err := t.repo.MarkTurnCompleted(ctx, t.sess.SessionID, idx, rawStopReason, updateCount); err != nil {
t.log.Error("acp.turn.mark_completed", "session_id", t.sess.SessionID, "turn_index", idx, "err", err.Error())
}
if err := t.repo.UpdateSessionLastStopReason(ctx, t.sess.SessionID, rawStopReason); err != nil {
t.log.Error("acp.turn.update_last_stop_reason", "session_id", t.sess.SessionID, "err", err.Error())
}
}
if t.bus != nil {
base := TurnEvent{
SessionID: t.sess.SessionID,
UserID: t.sess.UserID,
WorkspaceID: t.sess.WorkspaceID,
TurnIndex: idx,
StopReason: norm,
RawStopReason: rawStopReason,
CompletedAt: completedAt,
}
completed := base
completed.Kind = TurnCompletedEvent
t.bus.Publish(ctx, completed)
idle := base
idle.Kind = SessionIdleEvent
t.bus.Publish(ctx, idle)
}
}
// Abort 中止当前开放回合(如果有)并发布 session_idle 事件。从 supervisor.onExit
// 在子进程退出时调用,确保崩溃/被杀时也能闭合悬挂回合。无开放回合时 no-op。
func (t *TurnTracker) Abort(ctx context.Context, reason string) {
if t == nil {
return
}
t.mu.Lock()
defer t.mu.Unlock()
if !t.open {
return
}
t.abortLocked(ctx, reason)
}
// abortLocked 闭合当前开放回合并发布 session_idle。调用方必须已持有 t.mu。
func (t *TurnTracker) abortLocked(ctx context.Context, reason string) {
idx := t.turnIndex
t.open = false
t.promptID = ""
if t.repo != nil {
if err := t.repo.MarkTurnAborted(ctx, t.sess.SessionID, idx); err != nil {
t.log.Error("acp.turn.mark_aborted", "session_id", t.sess.SessionID, "turn_index", idx, "err", err.Error())
}
}
if t.bus != nil {
norm := NormalizeStopReason(reason)
t.bus.Publish(ctx, TurnEvent{
SessionID: t.sess.SessionID,
UserID: t.sess.UserID,
WorkspaceID: t.sess.WorkspaceID,
TurnIndex: idx,
Kind: SessionIdleEvent,
StopReason: norm,
RawStopReason: reason,
CompletedAt: time.Now(),
})
}
}
// ParseStopReason 从 session/prompt 响应的 Result JSON 中提取 stopReason 字段。
// 字段缺失或解析失败时返回空串(调用方归一化为 StopOther)。
func ParseStopReason(result json.RawMessage) string {
if len(result) == 0 {
return ""
}
var r struct {
StopReason string `json:"stopReason"`
}
if err := json.Unmarshal(result, &r); err != nil {
return ""
}
return r.StopReason
}
+248
View File
@@ -0,0 +1,248 @@
package acp_test
import (
"context"
"encoding/json"
"sync"
"testing"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/acp"
"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
)
// ----- NormalizeStopReason -----
func TestNormalizeStopReason(t *testing.T) {
t.Parallel()
cases := []struct {
raw string
want acp.StopReason
}{
{"end_turn", acp.StopEndTurn},
{"max_tokens", acp.StopMaxTokens},
{"refusal", acp.StopRefusal},
{"cancelled", acp.StopCancelled},
{"vendor_specific_thing", acp.StopOther},
{"", acp.StopOther},
}
for _, c := range cases {
assert.Equal(t, c.want, acp.NormalizeStopReason(c.raw), "raw=%q", c.raw)
}
}
func TestParseStopReason(t *testing.T) {
t.Parallel()
assert.Equal(t, "end_turn", acp.ParseStopReason(json.RawMessage(`{"stopReason":"end_turn"}`)))
assert.Equal(t, "max_tokens", acp.ParseStopReason(json.RawMessage(`{"stopReason":"max_tokens","other":1}`)))
assert.Equal(t, "", acp.ParseStopReason(json.RawMessage(`{}`)))
assert.Equal(t, "", acp.ParseStopReason(nil))
assert.Equal(t, "", acp.ParseStopReason(json.RawMessage(`not json`)))
}
// ----- fakeTurnRepo: 记录 turn 调用 -----
type fakeTurnRepo struct {
mu sync.Mutex
inserted []*acp.Turn
completed []completedCall
aborted []int
lastStopUpdates []string
}
type completedCall struct {
turnIndex int
stop string
updateCount int
}
func (f *fakeTurnRepo) InsertTurn(_ context.Context, t *acp.Turn) (*acp.Turn, error) {
f.mu.Lock()
defer f.mu.Unlock()
cp := *t
cp.ID = int64(len(f.inserted) + 1)
cp.Status = acp.TurnInProgress
f.inserted = append(f.inserted, &cp)
out := cp
return &out, nil
}
func (f *fakeTurnRepo) MarkTurnCompleted(_ context.Context, _ uuid.UUID, turnIndex int, stop string, updateCount int) error {
f.mu.Lock()
defer f.mu.Unlock()
f.completed = append(f.completed, completedCall{turnIndex, stop, updateCount})
return nil
}
func (f *fakeTurnRepo) MarkTurnAborted(_ context.Context, _ uuid.UUID, turnIndex int) error {
f.mu.Lock()
defer f.mu.Unlock()
f.aborted = append(f.aborted, turnIndex)
return nil
}
func (f *fakeTurnRepo) GetLatestTurnBySession(context.Context, uuid.UUID) (*acp.Turn, error) {
return nil, nil
}
func (f *fakeTurnRepo) UpdateSessionLastStopReason(_ context.Context, _ uuid.UUID, reason string) error {
f.mu.Lock()
defer f.mu.Unlock()
f.lastStopUpdates = append(f.lastStopUpdates, reason)
return nil
}
// ----- fakeBus: 记录发布的事件,保留顺序 -----
type fakeBus struct {
mu sync.Mutex
events []acp.TurnEvent
}
func (b *fakeBus) Subscribe(string, func(context.Context, acp.TurnEvent)) {}
func (b *fakeBus) Publish(_ context.Context, ev acp.TurnEvent) {
b.mu.Lock()
defer b.mu.Unlock()
b.events = append(b.events, ev)
}
func (b *fakeBus) snapshot() []acp.TurnEvent {
b.mu.Lock()
defer b.mu.Unlock()
out := make([]acp.TurnEvent, len(b.events))
copy(out, b.events)
return out
}
func newTestSessCtx() handlers.SessionContext {
return handlers.SessionContext{
SessionID: uuid.New(),
WorkspaceID: uuid.New(),
UserID: uuid.New(),
}
}
// ----- TurnTracker lifecycle -----
func TestTurnTracker_Lifecycle(t *testing.T) {
t.Parallel()
repo := &fakeTurnRepo{}
bus := &fakeBus{}
sess := newTestSessCtx()
tr := acp.NewTurnTracker(repo, bus, sess, nil)
ctx := context.Background()
idx, err := tr.StartTurn(ctx, "client-init")
require.NoError(t, err)
assert.Equal(t, 0, idx)
require.Len(t, repo.inserted, 1)
assert.Equal(t, 0, repo.inserted[0].TurnIndex)
// 3 个 update
const n = 3
for i := 0; i < n; i++ {
tr.OnUpdate(ctx)
}
// id 不匹配 → 不应完成
tr.OnPromptResponse(ctx, "wrong-id", "end_turn")
assert.Empty(t, repo.completed, "mismatched id must not complete")
// 匹配 → 完成
tr.OnPromptResponse(ctx, "client-init", "end_turn")
require.Len(t, repo.completed, 1)
assert.Equal(t, 0, repo.completed[0].turnIndex)
assert.Equal(t, "end_turn", repo.completed[0].stop)
assert.Equal(t, n, repo.completed[0].updateCount, "update_count must equal observed updates")
require.Len(t, repo.lastStopUpdates, 1)
assert.Equal(t, "end_turn", repo.lastStopUpdates[0])
// 事件顺序:turn_completed 然后 session_idle
evs := bus.snapshot()
require.Len(t, evs, 2)
assert.Equal(t, acp.TurnCompletedEvent, evs[0].Kind)
assert.Equal(t, acp.SessionIdleEvent, evs[1].Kind)
assert.Equal(t, acp.StopEndTurn, evs[0].StopReason)
assert.Equal(t, "end_turn", evs[0].RawStopReason)
assert.Equal(t, sess.SessionID, evs[0].SessionID)
assert.Equal(t, sess.UserID, evs[0].UserID)
// 重复响应 → no-op(回合已闭合)
tr.OnPromptResponse(ctx, "client-init", "end_turn")
assert.Len(t, repo.completed, 1, "duplicate response must be ignored")
}
func TestTurnTracker_UnknownStopReason(t *testing.T) {
t.Parallel()
repo := &fakeTurnRepo{}
bus := &fakeBus{}
tr := acp.NewTurnTracker(repo, bus, newTestSessCtx(), nil)
ctx := context.Background()
_, err := tr.StartTurn(ctx, "p1")
require.NoError(t, err)
tr.OnPromptResponse(ctx, "p1", "weird_vendor_reason")
require.Len(t, repo.completed, 1)
// 原始字符串原样落库
assert.Equal(t, "weird_vendor_reason", repo.completed[0].stop)
assert.Equal(t, "weird_vendor_reason", repo.lastStopUpdates[0])
// 归一化为 other
evs := bus.snapshot()
require.Len(t, evs, 2)
assert.Equal(t, acp.StopOther, evs[0].StopReason)
assert.Equal(t, "weird_vendor_reason", evs[0].RawStopReason)
}
func TestTurnTracker_Abort(t *testing.T) {
t.Parallel()
repo := &fakeTurnRepo{}
bus := &fakeBus{}
tr := acp.NewTurnTracker(repo, bus, newTestSessCtx(), nil)
ctx := context.Background()
_, err := tr.StartTurn(ctx, "p1")
require.NoError(t, err)
tr.Abort(ctx, "crashed")
require.Len(t, repo.aborted, 1)
assert.Equal(t, 0, repo.aborted[0])
evs := bus.snapshot()
require.Len(t, evs, 1)
assert.Equal(t, acp.SessionIdleEvent, evs[0].Kind)
// 无开放回合再 Abort → no-op
tr.Abort(ctx, "crashed")
assert.Len(t, repo.aborted, 1, "abort with no open turn must be no-op")
}
func TestTurnTracker_AutoAbortPriorTurn(t *testing.T) {
t.Parallel()
repo := &fakeTurnRepo{}
bus := &fakeBus{}
tr := acp.NewTurnTracker(repo, bus, newTestSessCtx(), nil)
ctx := context.Background()
idx0, err := tr.StartTurn(ctx, "p1")
require.NoError(t, err)
assert.Equal(t, 0, idx0)
// 第二次 StartTurn 在前一个开放时 → auto-abort 前一个
idx1, err := tr.StartTurn(ctx, "p2")
require.NoError(t, err)
assert.Equal(t, 1, idx1)
require.Len(t, repo.aborted, 1)
assert.Equal(t, 0, repo.aborted[0], "prior open turn must be auto-aborted")
require.Len(t, repo.inserted, 2)
}
func TestTurnTracker_NilSafe(t *testing.T) {
t.Parallel()
var tr *acp.TurnTracker
ctx := context.Background()
// 所有方法在 nil receiver 上安全 no-op
_, err := tr.StartTurn(ctx, "x")
assert.NoError(t, err)
tr.OnUpdate(ctx)
tr.OnPromptResponse(ctx, "x", "end_turn")
tr.Abort(ctx, "x")
assert.False(t, tr.IsTrackedPrompt("x"))
}
+131
View File
@@ -0,0 +1,131 @@
// turnbus.go 实现 TurnBus:进程内的极简 observer / pub-sub,作为"回合完成 /
// 会话空闲"的内部事件原语。与 notify.Dispatcher(面向用户的外发通知)解耦——
// TurnBus 是给未来的编排层(auto-continue / gate / hand-off)订阅用的内部信号,
// 不落库、不外发。
//
// 关键不变量:Publish 必须永不阻塞 relay reader 热路径。订阅者以 fire-and-forget
// 方式在独立 goroutine 中调用,并用 recover() 隔离 panic——一个慢/有 bug 的
// 订阅者不能拖垮协议 I/O。语义对齐 notify.Dispatcher 的非阻塞 fan-out。
package acp
import (
"context"
"fmt"
"log/slog"
"sync"
"time"
"github.com/google/uuid"
)
// TurnEventKind 标识 TurnEvent 的类别。
type TurnEventKind string
const (
// TurnCompletedEvent 表示一次回合正常完成(收到 session/prompt 响应的 stopReason)。
TurnCompletedEvent TurnEventKind = "turn_completed"
// SessionIdleEvent 表示会话当前空闲(回合完成或被中止后),编排层可据此决策。
SessionIdleEvent TurnEventKind = "session_idle"
)
// TurnEvent 是 TurnBus 上发布的事件载荷。RawStopReason 是 agent 上报的原始字符串
// (可能为空,例如 abort 路径),StopReason 是其归一化枚举。
type TurnEvent struct {
SessionID uuid.UUID
UserID uuid.UUID
WorkspaceID uuid.UUID
TurnIndex int
Kind TurnEventKind
StopReason StopReason
RawStopReason string
CompletedAt time.Time
}
// TurnBus 是回合事件的进程内 pub-sub 抽象。
type TurnBus interface {
// Subscribe 注册一个具名订阅者。name 仅用于日志/排障。重复 name 会追加,
// 不去重(与 notify 一致,调用方自行保证唯一)。
Subscribe(name string, fn func(context.Context, TurnEvent))
// Publish 把事件 fire-and-forget 派发给所有订阅者,永不阻塞调用方。
Publish(ctx context.Context, ev TurnEvent)
}
type deliver struct {
ctx context.Context
ev TurnEvent
}
type subscriber struct {
name string
fn func(context.Context, TurnEvent)
ch chan deliver
}
// subBuffer 是每个订阅者的投递缓冲容量。
const subBuffer = 256
// defaultTurnBus 是 TurnBus 的默认实现:每个订阅者一个常驻 worker goroutine,从带缓冲
// channel 顺序消费事件——因此对同一订阅者保证 per-subscriber FIFO(turn_completed 必先于
// 同一次 Publish 序列里随后的 session_idle 到达)。panic 用 recover() 隔离。Publish 以非阻塞
// 方式投递(缓冲满则丢弃并告警),永不阻塞 relay reader 热路径。
type defaultTurnBus struct {
mu sync.RWMutex
subs []*subscriber
log *slog.Logger
}
// NewTurnBus 构造默认 TurnBus。
func NewTurnBus(log *slog.Logger) TurnBus {
if log == nil {
log = slog.Default()
}
return &defaultTurnBus{log: log}
}
func (b *defaultTurnBus) Subscribe(name string, fn func(context.Context, TurnEvent)) {
if fn == nil {
return
}
s := &subscriber{name: name, fn: fn, ch: make(chan deliver, subBuffer)}
b.mu.Lock()
b.subs = append(b.subs, s)
b.mu.Unlock()
// 每个订阅者一个常驻 worker,顺序消费保证 FIFO。worker 与 app 同生命周期。
go b.worker(s)
}
func (b *defaultTurnBus) worker(s *subscriber) {
for d := range s.ch {
b.dispatch(s, d)
}
}
func (b *defaultTurnBus) dispatch(s *subscriber, d deliver) {
defer func() {
if rec := recover(); rec != nil {
b.log.Error("acp.turnbus.subscriber_panic",
"subscriber", s.name,
"session_id", d.ev.SessionID,
"panic", fmt.Sprintf("%v", rec))
}
}()
s.fn(d.ctx, d.ev)
}
func (b *defaultTurnBus) Publish(ctx context.Context, ev TurnEvent) {
b.mu.RLock()
subs := b.subs
b.mu.RUnlock()
// 解绑调用方 ctx 的取消(relay teardown 会 cancel ctx,但订阅者应能完成它的短暂处理)。
d := deliver{ctx: context.WithoutCancel(ctx), ev: ev}
for _, s := range subs {
select {
case s.ch <- d:
default:
// 缓冲满:丢弃以保证 Publish 永不阻塞 relay reader(正常负载不应发生)。
b.log.Warn("acp.turnbus.drop_full",
"subscriber", s.name, "session_id", ev.SessionID, "kind", string(ev.Kind))
}
}
}
+75
View File
@@ -0,0 +1,75 @@
package acp_test
import (
"context"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/acp"
)
func TestTurnBus_FanoutToMultipleSubscribers(t *testing.T) {
t.Parallel()
bus := acp.NewTurnBus(nil)
var wg sync.WaitGroup
wg.Add(2)
var a, b atomic.Int32
bus.Subscribe("a", func(_ context.Context, _ acp.TurnEvent) { a.Add(1); wg.Done() })
bus.Subscribe("b", func(_ context.Context, _ acp.TurnEvent) { b.Add(1); wg.Done() })
bus.Publish(context.Background(), acp.TurnEvent{SessionID: uuid.New(), Kind: acp.TurnCompletedEvent})
done := make(chan struct{})
go func() { wg.Wait(); close(done) }()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("subscribers not invoked in time")
}
assert.Equal(t, int32(1), a.Load())
assert.Equal(t, int32(1), b.Load())
}
func TestTurnBus_PanicIsolatedAndNonBlocking(t *testing.T) {
t.Parallel()
bus := acp.NewTurnBus(nil)
var good atomic.Int32
var wg sync.WaitGroup
wg.Add(1)
bus.Subscribe("panicky", func(_ context.Context, _ acp.TurnEvent) { panic("boom") })
bus.Subscribe("good", func(_ context.Context, _ acp.TurnEvent) { good.Add(1); wg.Done() })
// Publish 必须立即返回(非阻塞),且 panicky 订阅者的 panic 被 recover 隔离,
// 不影响 good 订阅者执行。
bus.Publish(context.Background(), acp.TurnEvent{Kind: acp.SessionIdleEvent})
done := make(chan struct{})
go func() { wg.Wait(); close(done) }()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("good subscriber not invoked despite panicky peer")
}
assert.Equal(t, int32(1), good.Load())
}
func TestTurnBus_PublishDoesNotBlockOnSlowSubscriber(t *testing.T) {
t.Parallel()
bus := acp.NewTurnBus(nil)
release := make(chan struct{})
bus.Subscribe("slow", func(_ context.Context, _ acp.TurnEvent) { <-release })
start := time.Now()
bus.Publish(context.Background(), acp.TurnEvent{Kind: acp.SessionIdleEvent})
// Publish 应远快于订阅者的阻塞时长。
require.Less(t, time.Since(start), 500*time.Millisecond, "Publish must not block on slow subscriber")
close(release)
}
+396
View File
@@ -0,0 +1,396 @@
// usage.go implements ACP session cost accounting:
// - ExtractUsage: a tolerant parser pulling per-turn billable token usage out
// of agent messages (session/prompt RESPONSE result.usage, snake/camel) and
// opportunistic cost out of usage_update notifications.
// - usageAccumulator: applies one parsed turn — prices it (ModelPriceLookup or
// agent self-reported cost fallback), persists ledger + running totals via
// the repo, and reports a budget breach when a cap is now exceeded.
//
// Billing source of truth (spec §7 risk #1): claude-agent-acp returns billable
// usage in the session/prompt response ({stopReason, usage:{input_tokens,...}}).
// The ACP usage_update notification carries a context-window GAUGE (used/size),
// NOT a per-turn delta — summing it would massively overcount, so it is used
// only for the optional cost{amount} fallback, never for token deltas.
package acp
import (
"context"
"encoding/json"
"log/slog"
"sync"
"time"
"github.com/google/uuid"
)
// ParsedUsage is one extracted turn. HasUsage=false means no billable usage was
// present (ordinary chunk / malformed payload) and the caller must skip it.
type ParsedUsage struct {
PromptTokens int
CompletionTokens int
ThinkingTokens int
// CostUSD is the agent-self-reported cost (usage_update.cost.amount), used as
// a pricing fallback when no model price is configured. nil = not reported.
CostUSD *float64
HasUsage bool
}
// ModelPrice is the per-million-token price triple for a model.
type ModelPrice struct {
PromptPerM float64
CompletionPerM float64
ThinkingPerM float64
Found bool
}
// ModelPriceLookup resolves a model_id to its prices. Implemented in app.go by
// an adapter over the chat endpoint/model service.
type ModelPriceLookup interface {
PriceForModel(ctx context.Context, modelID uuid.UUID) (ModelPrice, error)
}
// UsageAccumulator applies parsed turns to one session: persists the ledger +
// running totals and reports budget breaches.
type UsageAccumulator interface {
// Observe applies one parsed turn. sourceEventID (>0) de-dups via the ledger
// unique index. Returns a non-empty breach reason when a cap is now exceeded.
Observe(ctx context.Context, p ParsedUsage, sourceEventID int64) (breach BreachReason, breached bool)
}
// usageSink is fed raw messages from relay.reader after persistence; it extracts
// usage and forwards to the UsageAccumulator. Separated from the accumulator so
// the relay does not depend on pricing/repo details.
type usageSink interface {
Observe(ctx context.Context, msg *Message, eventID int64)
}
// accumulatorSink adapts a UsageAccumulator to the relay's usageSink: it parses
// each agent->client message and, on a real billable turn, forwards to Observe.
// A breach invokes onBreach exactly once on a detached goroutine (so the relay
// reader never blocks on sup.Kill).
type accumulatorSink struct {
acc UsageAccumulator
onBreach func(reason BreachReason)
fired sync.Once
log *slog.Logger
}
// newAccumulatorSink builds the relay usage sink.
func newAccumulatorSink(acc UsageAccumulator, onBreach func(reason BreachReason), log *slog.Logger) *accumulatorSink {
if log == nil {
log = slog.Default()
}
return &accumulatorSink{acc: acc, onBreach: onBreach, log: log}
}
func (s *accumulatorSink) Observe(ctx context.Context, msg *Message, eventID int64) {
if s == nil || s.acc == nil {
return
}
p := ExtractUsage(msg)
if !p.HasUsage {
return
}
breach, ok := s.acc.Observe(ctx, p, eventID)
if !ok || s.onBreach == nil {
return
}
// Fire the kill exactly once, detached so the reader keeps draining.
s.fired.Do(func() {
reason := breach
go s.onBreach(reason)
})
}
// ===== ExtractUsage =====
// promptResponseUsage matches both snake_case (claude-agent-acp stable) and
// camelCase (unstable v2 Usage) usage objects on a session/prompt response.
type promptResponseResult struct {
StopReason string `json:"stopReason"`
Usage *struct {
// snake_case (stable)
InputTokens *int `json:"input_tokens"`
OutputTokens *int `json:"output_tokens"`
// camelCase (unstable v2)
InputTokensCamel *int `json:"inputTokens"`
OutputTokensCamel *int `json:"outputTokens"`
ThoughtTokens *int `json:"thoughtTokens"`
CachedReadTokens *int `json:"cachedReadTokens"`
CachedWriteTokens *int `json:"cachedWriteTokens"`
// thinking_tokens snake fallback
ThinkingTokens *int `json:"thinking_tokens"`
} `json:"usage"`
}
// usageUpdateParams matches a session/update notification carrying usage_update.
type usageUpdateParams struct {
Update *struct {
SessionUpdate string `json:"sessionUpdate"`
Cost *struct {
Amount float64 `json:"amount"`
Currency string `json:"currency"`
} `json:"cost"`
} `json:"update"`
}
// ExtractUsage is the tolerant entry point. It inspects an agent->client message
// for billable usage. Never panics on garbage — returns HasUsage=false instead.
func ExtractUsage(msg *Message) ParsedUsage {
if msg == nil {
return ParsedUsage{}
}
// 1) session/prompt RESPONSE result.usage — the reliable billable source.
if len(msg.Result) > 0 {
if p, ok := parsePromptResponseUsage(msg.Result); ok {
return p
}
}
// 2) usage_update notification — opportunistic cost only (used/size is a
// context-window gauge, NOT a billable token delta; do not sum it).
if msg.Method == "session/update" && len(msg.Params) > 0 {
if p, ok := parseUsageUpdateCost(msg.Params); ok {
return p
}
}
return ParsedUsage{}
}
func parsePromptResponseUsage(result json.RawMessage) (ParsedUsage, bool) {
var r promptResponseResult
if err := json.Unmarshal(result, &r); err != nil || r.Usage == nil {
return ParsedUsage{}, false
}
u := r.Usage
prompt := firstNonNil(u.InputTokens, u.InputTokensCamel)
completion := firstNonNil(u.OutputTokens, u.OutputTokensCamel)
thinking := firstNonNil(u.ThoughtTokens, u.ThinkingTokens)
// cachedReadTokens/cachedWriteTokens are reported but not separately priced;
// fold cached writes into prompt-equivalent only if no input token present.
if prompt == 0 && completion == 0 && thinking == 0 {
// No billable token signal in the usage object.
return ParsedUsage{}, false
}
return ParsedUsage{
PromptTokens: prompt,
CompletionTokens: completion,
ThinkingTokens: thinking,
HasUsage: true,
}, true
}
func parseUsageUpdateCost(params json.RawMessage) (ParsedUsage, bool) {
var p usageUpdateParams
if err := json.Unmarshal(params, &p); err != nil || p.Update == nil {
return ParsedUsage{}, false
}
if p.Update.SessionUpdate != "usage_update" || p.Update.Cost == nil {
return ParsedUsage{}, false
}
amt := p.Update.Cost.Amount
if amt <= 0 {
return ParsedUsage{}, false
}
return ParsedUsage{CostUSD: &amt, HasUsage: true}, true
}
func firstNonNil(a, b *int) int {
if a != nil {
return *a
}
if b != nil {
return *b
}
return 0
}
// ===== usageAccumulator (per-session) =====
// usageRepo is the narrow repo surface the accumulator needs.
type usageRepo interface {
InsertSessionUsage(ctx context.Context, r *SessionUsageRecord) (bool, error)
AddSessionUsageTotals(ctx context.Context, sid uuid.UUID, dp, dc, dt int64, dCost float64) (int64, int64, int64, float64, error)
SumProjectCostUSD(ctx context.Context, projectID uuid.UUID, since time.Time) (float64, error)
}
// pgUsageAccumulator is the concrete per-session accumulator. Thread-safe: the
// relay reader is single-goroutine today, but Observe is guarded so future
// concurrent callers (or tests) stay correct.
type pgUsageAccumulator struct {
repo usageRepo
prices ModelPriceLookup
log *slog.Logger
now func() time.Time
mu sync.Mutex
dedup map[int64]struct{}
sessionID uuid.UUID
userID uuid.UUID
projectID uuid.UUID
agentKindID uuid.UUID
modelID *uuid.UUID
startedAt time.Time
caps BudgetCaps
// projectBudgetUSD (>0) enforces a cumulative per-project ACP spend soft cap
// over [projectSince, now]; 0 disables it.
projectBudgetUSD float64
projectSince time.Time
}
// AccumulatorParams seeds a per-session accumulator.
type AccumulatorParams struct {
Repo usageRepo
Prices ModelPriceLookup
Log *slog.Logger
SessionID uuid.UUID
UserID uuid.UUID
ProjectID uuid.UUID
AgentKindID uuid.UUID
ModelID *uuid.UUID
StartedAt time.Time
Caps BudgetCaps
// ProjectBudgetUSD, when >0, enforces a per-project soft cap on cumulative
// ACP spend since ProjectSince.
ProjectBudgetUSD float64
ProjectSince time.Time
Now func() time.Time
}
func newUsageAccumulator(p AccumulatorParams) *pgUsageAccumulator {
log := p.Log
if log == nil {
log = slog.Default()
}
now := p.Now
if now == nil {
now = time.Now
}
since := p.ProjectSince
if since.IsZero() {
since = now().Add(-30 * 24 * time.Hour)
}
return &pgUsageAccumulator{
repo: p.Repo,
prices: p.Prices,
log: log,
now: now,
dedup: map[int64]struct{}{},
sessionID: p.SessionID,
userID: p.UserID,
projectID: p.ProjectID,
agentKindID: p.AgentKindID,
modelID: p.ModelID,
startedAt: p.StartedAt,
caps: p.Caps,
projectBudgetUSD: p.ProjectBudgetUSD,
projectSince: since,
}
}
func (a *pgUsageAccumulator) Observe(ctx context.Context, p ParsedUsage, sourceEventID int64) (BreachReason, bool) {
if !p.HasUsage {
return "", false
}
a.mu.Lock()
defer a.mu.Unlock()
// In-memory de-dup guard (DB unique index is the durable de-dup).
if sourceEventID > 0 {
if _, seen := a.dedup[sourceEventID]; seen {
return "", false
}
}
cost := a.price(ctx, p)
// Persist ledger row first (idempotent per source_event_id), then accumulate
// session totals only when the ledger row was actually inserted (no re-add on
// a duplicate event).
var srcID *int64
if sourceEventID > 0 {
v := sourceEventID
srcID = &v
}
inserted, err := a.repo.InsertSessionUsage(ctx, &SessionUsageRecord{
SessionID: a.sessionID,
UserID: a.userID,
ProjectID: a.projectID,
AgentKindID: a.agentKindID,
ModelID: a.modelID,
PromptTokens: int64(p.PromptTokens),
CompletionTokens: int64(p.CompletionTokens),
ThinkingTokens: int64(p.ThinkingTokens),
CostUSD: cost,
SourceEventID: srcID,
})
if err != nil {
a.log.Error("acp.usage.insert_ledger", "session_id", a.sessionID, "err", err.Error())
return "", false
}
if !inserted {
// Duplicate event id — already accounted.
if sourceEventID > 0 {
a.dedup[sourceEventID] = struct{}{}
}
return "", false
}
if sourceEventID > 0 {
a.dedup[sourceEventID] = struct{}{}
}
totPrompt, totCompletion, totThinking, totCost, err := a.repo.AddSessionUsageTotals(
ctx, a.sessionID,
int64(p.PromptTokens), int64(p.CompletionTokens), int64(p.ThinkingTokens), cost)
if err != nil {
a.log.Error("acp.usage.add_totals", "session_id", a.sessionID, "err", err.Error())
return "", false
}
return a.checkBreach(ctx, totPrompt, totCompletion, totThinking, totCost)
}
// price computes the turn cost: prefer the configured model price; otherwise
// fall back to the agent-self-reported cost; otherwise 0 (unpriced).
func (a *pgUsageAccumulator) price(ctx context.Context, p ParsedUsage) float64 {
if a.modelID != nil && a.prices != nil {
if mp, err := a.prices.PriceForModel(ctx, *a.modelID); err == nil && mp.Found {
return (float64(p.PromptTokens)*mp.PromptPerM +
float64(p.CompletionTokens)*mp.CompletionPerM +
float64(p.ThinkingTokens)*mp.ThinkingPerM) / 1_000_000
}
}
if p.CostUSD != nil {
return *p.CostUSD
}
return 0
}
func (a *pgUsageAccumulator) checkBreach(ctx context.Context, totPrompt, totCompletion, totThinking int64, totCost float64) (BreachReason, bool) {
// Wall-clock: enforced here too (the reaper is the periodic backstop).
if a.caps.MaxWallClockSeconds != nil {
elapsed := a.now().Sub(a.startedAt)
if elapsed >= time.Duration(*a.caps.MaxWallClockSeconds)*time.Second {
return BreachWallClock, true
}
}
if a.caps.MaxTokens != nil {
if totPrompt+totCompletion+totThinking >= *a.caps.MaxTokens {
return BreachTokens, true
}
}
if a.caps.MaxCostUSD != nil && totCost >= *a.caps.MaxCostUSD {
return BreachCost, true
}
// Per-project soft cap (advisory kill, slight overshoot tolerated under
// concurrency — spec §7 risk). Only queried when a project budget is set.
if a.projectBudgetUSD > 0 {
if projCost, err := a.repo.SumProjectCostUSD(ctx, a.projectID, a.projectSince); err == nil {
if projCost >= a.projectBudgetUSD {
return BreachCost, true
}
} else {
a.log.Warn("acp.usage.project_cost_check", "project_id", a.projectID, "err", err.Error())
}
}
return "", false
}
+317
View File
@@ -0,0 +1,317 @@
package acp
import (
"context"
"errors"
"fmt"
"math/big"
"time"
"github.com/google/uuid"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgtype"
acpsqlc "github.com/yan1h/agent-coding-workflow/internal/acp/sqlc"
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
)
func isNoRows(err error) bool { return errors.Is(err, pgx.ErrNoRows) }
// ===== NUMERIC <-> float64 转换 helper(复制 chat/repository.go 的实现)=====
func numericToFloat64(n pgtype.Numeric) float64 {
if !n.Valid {
return 0
}
f8, err := n.Float64Value()
if err != nil {
return 0
}
return f8.Float64
}
func numericToFloat64Ptr(n pgtype.Numeric) *float64 {
if !n.Valid {
return nil
}
v := numericToFloat64(n)
return &v
}
func float64ToNumeric(f float64) pgtype.Numeric {
var n pgtype.Numeric
if err := n.ScanScientific(fmt.Sprintf("%.10e", f)); err != nil {
bi := new(big.Int)
bi.SetInt64(int64(f))
return pgtype.Numeric{Int: bi, Exp: 0, Valid: true}
}
return n
}
func float64PtrToNumeric(f *float64) pgtype.Numeric {
if f == nil {
return pgtype.Numeric{}
}
return float64ToNumeric(*f)
}
// ===== 成本核算 / 预算 / reaper 仓储方法 =====
func (r *pgRepo) InsertSessionUsage(ctx context.Context, rec *SessionUsageRecord) (bool, error) {
var srcID *int64
if rec.SourceEventID != nil {
v := *rec.SourceEventID
srcID = &v
}
id, err := r.q.InsertSessionUsage(ctx, acpsqlc.InsertSessionUsageParams{
SessionID: toPgUUID(rec.SessionID),
UserID: toPgUUID(rec.UserID),
ProjectID: toPgUUID(rec.ProjectID),
AgentKindID: toPgUUID(rec.AgentKindID),
ModelID: toPgUUIDPtr(rec.ModelID),
PromptTokens: rec.PromptTokens,
CompletionTokens: rec.CompletionTokens,
ThinkingTokens: rec.ThinkingTokens,
CostUsd: float64ToNumeric(rec.CostUSD),
SourceEventID: srcID,
})
if err != nil {
// ON CONFLICT DO NOTHING + no row -> not inserted (idempotent de-dup).
if isNoRows(err) {
return false, nil
}
return false, errs.Wrap(err, errs.CodeInternal, "insert session usage")
}
rec.ID = id
return true, nil
}
func (r *pgRepo) AddSessionUsageTotals(ctx context.Context, sid uuid.UUID, dp, dc, dt int64, dCost float64) (int64, int64, int64, float64, error) {
row, err := r.q.AddSessionUsageTotals(ctx, acpsqlc.AddSessionUsageTotalsParams{
ID: toPgUUID(sid),
PromptTokens: dp,
CompletionTokens: dc,
ThinkingTokens: dt,
TotalCostUsd: float64ToNumeric(dCost),
})
if err != nil {
return 0, 0, 0, 0, errs.Wrap(err, errs.CodeInternal, "add session usage totals")
}
return row.PromptTokens, row.CompletionTokens, row.ThinkingTokens, numericToFloat64(row.TotalCostUsd), nil
}
func (r *pgRepo) SumProjectCostUSD(ctx context.Context, projectID uuid.UUID, since time.Time) (float64, error) {
n, err := r.q.SumProjectCostUSD(ctx, acpsqlc.SumProjectCostUSDParams{
ProjectID: toPgUUID(projectID),
CreatedAt: pgtype.Timestamptz{Time: since, Valid: true},
})
if err != nil {
return 0, errs.Wrap(err, errs.CodeInternal, "sum project cost")
}
return numericToFloat64(n), nil
}
func (r *pgRepo) MarkSessionTerminatedReason(ctx context.Context, sid uuid.UUID, reason string) error {
if err := r.q.MarkSessionTerminatedReason(ctx, acpsqlc.MarkSessionTerminatedReasonParams{
ID: toPgUUID(sid),
TerminatedReason: &reason,
}); err != nil {
return errs.Wrap(err, errs.CodeInternal, "mark session terminated reason")
}
return nil
}
func (r *pgRepo) ListSessionsForReaper(ctx context.Context, wallClockStartedBefore, idleSince time.Time) ([]ReaperSession, error) {
var wallTS pgtype.Timestamptz
if !wallClockStartedBefore.IsZero() {
wallTS = pgtype.Timestamptz{Time: wallClockStartedBefore, Valid: true}
}
rows, err := r.q.ListSessionsForReaper(ctx, acpsqlc.ListSessionsForReaperParams{
Column1: wallTS,
Column2: pgtype.Timestamptz{Time: idleSince, Valid: true},
})
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "list sessions for reaper")
}
out := make([]ReaperSession, 0, len(rows))
for _, row := range rows {
out = append(out, ReaperSession{
ID: fromPgUUID(row.ID),
UserID: fromPgUUID(row.UserID),
StartedAt: row.StartedAt.Time,
LastActivityAt: row.LastActivityAt.Time,
MaxWallClockSeconds: row.BudgetMaxWallClockSeconds,
})
}
return out, nil
}
func (r *pgRepo) ListSessionUsage(ctx context.Context, sessionID uuid.UUID, limit int32) ([]*SessionUsageRecord, error) {
if limit <= 0 {
limit = 50
}
rows, err := r.q.ListSessionUsageBySession(ctx, acpsqlc.ListSessionUsageBySessionParams{
SessionID: toPgUUID(sessionID),
Limit: limit,
})
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "list session usage")
}
out := make([]*SessionUsageRecord, 0, len(rows))
for _, row := range rows {
out = append(out, rowToSessionUsage(row))
}
return out, nil
}
func rowToSessionUsage(row acpsqlc.AcpSessionUsage) *SessionUsageRecord {
var srcID *int64
if row.SourceEventID != nil {
v := *row.SourceEventID
srcID = &v
}
return &SessionUsageRecord{
ID: row.ID,
SessionID: fromPgUUID(row.SessionID),
UserID: fromPgUUID(row.UserID),
ProjectID: fromPgUUID(row.ProjectID),
AgentKindID: fromPgUUID(row.AgentKindID),
ModelID: fromPgUUIDPtr(row.ModelID),
PromptTokens: row.PromptTokens,
CompletionTokens: row.CompletionTokens,
ThinkingTokens: row.ThinkingTokens,
CostUSD: numericToFloat64(row.CostUsd),
SourceEventID: srcID,
CreatedAt: row.CreatedAt.Time,
}
}
func (r *pgRepo) SummarizeAcpUsageByUser(ctx context.Context, from, to time.Time) ([]AcpUsageUserRow, error) {
rows, err := r.q.SummarizeAcpUsageByUser(ctx, acpsqlc.SummarizeAcpUsageByUserParams{
CreatedAt: pgtype.Timestamptz{Time: from, Valid: true},
CreatedAt_2: pgtype.Timestamptz{Time: to, Valid: true},
})
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "summarize acp usage by user")
}
out := make([]AcpUsageUserRow, 0, len(rows))
for _, row := range rows {
out = append(out, AcpUsageUserRow{
UserID: fromPgUUID(row.UserID),
PromptTokens: row.PromptTokens,
CompletionTokens: row.CompletionTokens,
ThinkingTokens: row.ThinkingTokens,
CostUSD: numericToFloat64(row.CostUsd),
})
}
return out, nil
}
func (r *pgRepo) SummarizeAcpUsageByModel(ctx context.Context, from, to time.Time) ([]AcpUsageModelRow, error) {
rows, err := r.q.SummarizeAcpUsageByModel(ctx, acpsqlc.SummarizeAcpUsageByModelParams{
CreatedAt: pgtype.Timestamptz{Time: from, Valid: true},
CreatedAt_2: pgtype.Timestamptz{Time: to, Valid: true},
})
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "summarize acp usage by model")
}
out := make([]AcpUsageModelRow, 0, len(rows))
for _, row := range rows {
out = append(out, AcpUsageModelRow{
ModelID: fromPgUUID(row.ModelID),
PromptTokens: row.PromptTokens,
CompletionTokens: row.CompletionTokens,
ThinkingTokens: row.ThinkingTokens,
CostUSD: numericToFloat64(row.CostUsd),
})
}
return out, nil
}
func (r *pgRepo) SummarizeAcpUsageByDay(ctx context.Context, from, to time.Time) ([]AcpUsageDayRow, error) {
rows, err := r.q.SummarizeAcpUsageByDay(ctx, acpsqlc.SummarizeAcpUsageByDayParams{
CreatedAt: pgtype.Timestamptz{Time: from, Valid: true},
CreatedAt_2: pgtype.Timestamptz{Time: to, Valid: true},
})
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "summarize acp usage by day")
}
out := make([]AcpUsageDayRow, 0, len(rows))
for _, row := range rows {
out = append(out, AcpUsageDayRow{
Day: row.Day.Time,
PromptTokens: row.PromptTokens,
CompletionTokens: row.CompletionTokens,
ThinkingTokens: row.ThinkingTokens,
CostUSD: numericToFloat64(row.CostUsd),
})
}
return out, nil
}
// ===== run-history dashboard 聚合 =====
func (r *pgRepo) SessionRollup(ctx context.Context, f DashboardFilter) (SessionRollup, error) {
row, err := r.q.SessionRollup(ctx, acpsqlc.SessionRollupParams{
Column1: toPgUUIDPtr(f.UserID),
Column2: toPgUUIDPtr(f.ProjectID),
Column3: toPgUUIDPtr(f.RequirementID),
StartedAt: pgtype.Timestamptz{Time: f.From, Valid: true},
StartedAt_2: pgtype.Timestamptz{Time: f.To, Valid: true},
})
if err != nil {
return SessionRollup{}, errs.Wrap(err, errs.CodeInternal, "session rollup")
}
return SessionRollup{
Total: row.Total,
Succeeded: row.Succeeded,
Crashed: row.Crashed,
Active: row.Active,
TotalCostUSD: numericToFloat64(row.TotalCostUsd),
TotalTokens: row.TotalTokens,
AvgDurationSeconds: row.AvgDurationSeconds,
}, nil
}
func (r *pgRepo) SessionsByDay(ctx context.Context, f DashboardFilter) ([]SessionDayRow, error) {
rows, err := r.q.SessionsByDay(ctx, acpsqlc.SessionsByDayParams{
Column1: toPgUUIDPtr(f.UserID),
Column2: toPgUUIDPtr(f.ProjectID),
Column3: toPgUUIDPtr(f.RequirementID),
StartedAt: pgtype.Timestamptz{Time: f.From, Valid: true},
StartedAt_2: pgtype.Timestamptz{Time: f.To, Valid: true},
})
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "sessions by day")
}
out := make([]SessionDayRow, 0, len(rows))
for _, row := range rows {
out = append(out, SessionDayRow{
Day: row.Day.Time,
Total: row.Total,
Succeeded: row.Succeeded,
Crashed: row.Crashed,
TotalCostUSD: numericToFloat64(row.TotalCostUsd),
})
}
return out, nil
}
func (r *pgRepo) SessionTimeline(ctx context.Context, sessionID uuid.UUID) ([]SessionTimelineBucket, error) {
rows, err := r.q.SessionTimeline(ctx, toPgUUID(sessionID))
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "session timeline")
}
out := make([]SessionTimelineBucket, 0, len(rows))
for _, row := range rows {
out = append(out, SessionTimelineBucket{
Direction: row.Direction,
RPCKind: row.RpcKind,
Method: row.Method,
EventCount: row.EventCount,
FirstAt: row.FirstAt.Time,
LastAt: row.LastAt.Time,
})
}
return out, nil
}
+260
View File
@@ -0,0 +1,260 @@
package acp
import (
"context"
"encoding/json"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func msgWithResult(t *testing.T, result string) *Message {
t.Helper()
return &Message{JSONRPC: "2.0", ID: json.RawMessage(`"client-prompt-1"`), Result: json.RawMessage(result)}
}
func msgNotification(t *testing.T, method, params string) *Message {
t.Helper()
return &Message{JSONRPC: "2.0", Method: method, Params: json.RawMessage(params)}
}
func TestExtractUsage(t *testing.T) {
tests := []struct {
name string
msg *Message
wantHas bool
wantPrompt int
wantComp int
wantThink int
wantCost *float64
}{
{
name: "snake_case prompt response",
msg: msgWithResult(t, `{"stopReason":"end_turn","usage":{"input_tokens":120,"output_tokens":48}}`),
wantHas: true,
wantPrompt: 120,
wantComp: 48,
},
{
name: "camelCase v2 usage with thoughtTokens + cached",
msg: msgWithResult(t, `{"stopReason":"end_turn","usage":{"inputTokens":200,"outputTokens":90,"thoughtTokens":15,"cachedReadTokens":50}}`),
wantHas: true,
wantPrompt: 200,
wantComp: 90,
wantThink: 15,
},
{
name: "usage_update notification with cost only",
msg: msgNotification(t, "session/update", `{"update":{"sessionUpdate":"usage_update","used":1000,"size":200000,"cost":{"amount":0.0123,"currency":"USD"}}}`),
wantHas: true,
wantCost: f64ptr(0.0123),
},
{
name: "usage_update without cost -> no usage (gauge ignored)",
msg: msgNotification(t, "session/update", `{"update":{"sessionUpdate":"usage_update","used":1000,"size":200000}}`),
wantHas: false,
},
{
name: "ordinary agent_message_chunk -> no usage",
msg: msgNotification(t, "session/update", `{"update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"hi"}}}`),
wantHas: false,
},
{
name: "malformed result -> no panic, no usage",
msg: msgWithResult(t, `{"usage":not-json`),
wantHas: false,
},
{
name: "result without usage object -> no usage",
msg: msgWithResult(t, `{"stopReason":"end_turn"}`),
wantHas: false,
},
{
name: "nil message -> no usage",
msg: nil,
wantHas: false,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
got := ExtractUsage(tc.msg)
assert.Equal(t, tc.wantHas, got.HasUsage)
if !tc.wantHas {
return
}
assert.Equal(t, tc.wantPrompt, got.PromptTokens)
assert.Equal(t, tc.wantComp, got.CompletionTokens)
assert.Equal(t, tc.wantThink, got.ThinkingTokens)
if tc.wantCost == nil {
assert.Nil(t, got.CostUSD)
} else {
require.NotNil(t, got.CostUSD)
assert.InDelta(t, *tc.wantCost, *got.CostUSD, 1e-9)
}
})
}
}
func f64ptr(f float64) *float64 { return &f }
func i64ptr(i int64) *int64 { return &i }
// ===== usageAccumulator =====
type fakeUsageRepo struct {
ledger []*SessionUsageRecord
dedup map[int64]struct{}
totPrompt int64
totComp int64
totThink int64
totCost float64
projectCost float64
}
func (f *fakeUsageRepo) InsertSessionUsage(_ context.Context, r *SessionUsageRecord) (bool, error) {
if r.SourceEventID != nil {
if f.dedup == nil {
f.dedup = map[int64]struct{}{}
}
if _, ok := f.dedup[*r.SourceEventID]; ok {
return false, nil
}
f.dedup[*r.SourceEventID] = struct{}{}
}
cp := *r
f.ledger = append(f.ledger, &cp)
return true, nil
}
func (f *fakeUsageRepo) AddSessionUsageTotals(_ context.Context, _ uuid.UUID, dp, dc, dt int64, dCost float64) (int64, int64, int64, float64, error) {
f.totPrompt += dp
f.totComp += dc
f.totThink += dt
f.totCost += dCost
return f.totPrompt, f.totComp, f.totThink, f.totCost, nil
}
func (f *fakeUsageRepo) SumProjectCostUSD(_ context.Context, _ uuid.UUID, _ time.Time) (float64, error) {
return f.projectCost, nil
}
type fakePriceLookup struct {
price ModelPrice
}
func (f fakePriceLookup) PriceForModel(_ context.Context, _ uuid.UUID) (ModelPrice, error) {
return f.price, nil
}
func TestAccumulator_PricesViaModelLookup(t *testing.T) {
repo := &fakeUsageRepo{}
mid := uuid.New()
acc := newUsageAccumulator(AccumulatorParams{
Repo: repo,
Prices: fakePriceLookup{price: ModelPrice{PromptPerM: 3, CompletionPerM: 15, ThinkingPerM: 0, Found: true}},
ModelID: &mid,
Caps: BudgetCaps{},
})
_, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 1_000_000, CompletionTokens: 1_000_000, HasUsage: true}, 1)
assert.False(t, breached)
require.Len(t, repo.ledger, 1)
// 3 USD (prompt) + 15 USD (completion) = 18.
assert.InDelta(t, 18.0, repo.ledger[0].CostUSD, 1e-6)
assert.InDelta(t, 18.0, repo.totCost, 1e-6)
}
func TestAccumulator_FallsBackToAgentCost(t *testing.T) {
repo := &fakeUsageRepo{}
acc := newUsageAccumulator(AccumulatorParams{Repo: repo}) // no model price
cost := 0.5
_, _ = acc.Observe(context.Background(), ParsedUsage{CostUSD: &cost, HasUsage: true}, 1)
require.Len(t, repo.ledger, 1)
assert.InDelta(t, 0.5, repo.ledger[0].CostUSD, 1e-9)
}
func TestAccumulator_SkipsWhenNoUsage(t *testing.T) {
repo := &fakeUsageRepo{}
acc := newUsageAccumulator(AccumulatorParams{Repo: repo})
_, breached := acc.Observe(context.Background(), ParsedUsage{HasUsage: false}, 1)
assert.False(t, breached)
assert.Empty(t, repo.ledger)
}
func TestAccumulator_DedupBySourceEventID(t *testing.T) {
repo := &fakeUsageRepo{}
mid := uuid.New()
acc := newUsageAccumulator(AccumulatorParams{
Repo: repo,
Prices: fakePriceLookup{price: ModelPrice{CompletionPerM: 10, Found: true}},
ModelID: &mid,
})
p := ParsedUsage{CompletionTokens: 1000, HasUsage: true}
_, _ = acc.Observe(context.Background(), p, 7)
_, _ = acc.Observe(context.Background(), p, 7) // same event id -> no-op
assert.Len(t, repo.ledger, 1)
assert.InDelta(t, 0.01, repo.totCost, 1e-9)
}
func TestAccumulator_BreachCost(t *testing.T) {
repo := &fakeUsageRepo{}
mid := uuid.New()
acc := newUsageAccumulator(AccumulatorParams{
Repo: repo,
Prices: fakePriceLookup{price: ModelPrice{CompletionPerM: 1000, Found: true}},
ModelID: &mid,
Caps: BudgetCaps{MaxCostUSD: f64ptr(0.5)},
})
// 1000 completion tokens * 1000/M = 1.0 USD >= 0.5 cap.
reason, breached := acc.Observe(context.Background(), ParsedUsage{CompletionTokens: 1000, HasUsage: true}, 1)
assert.True(t, breached)
assert.Equal(t, BreachCost, reason)
}
func TestAccumulator_BreachTokens(t *testing.T) {
repo := &fakeUsageRepo{}
acc := newUsageAccumulator(AccumulatorParams{
Repo: repo,
Caps: BudgetCaps{MaxTokens: i64ptr(1000)},
})
reason, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 600, CompletionTokens: 600, HasUsage: true}, 1)
assert.True(t, breached)
assert.Equal(t, BreachTokens, reason)
}
func TestAccumulator_BreachWallClock(t *testing.T) {
repo := &fakeUsageRepo{}
start := time.Now().Add(-2 * time.Hour)
acc := newUsageAccumulator(AccumulatorParams{
Repo: repo,
StartedAt: start,
Caps: BudgetCaps{MaxWallClockSeconds: i32ptr(60)}, // 1 minute cap, 2h elapsed
})
reason, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 1, HasUsage: true}, 1)
assert.True(t, breached)
assert.Equal(t, BreachWallClock, reason)
}
func TestAccumulator_BreachProjectBudget(t *testing.T) {
repo := &fakeUsageRepo{projectCost: 100} // already over the project cap
acc := newUsageAccumulator(AccumulatorParams{
Repo: repo,
ProjectBudgetUSD: 50,
})
reason, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 1, HasUsage: true}, 1)
assert.True(t, breached)
assert.Equal(t, BreachCost, reason)
}
func TestAccumulator_NoBreachUnderCaps(t *testing.T) {
repo := &fakeUsageRepo{}
acc := newUsageAccumulator(AccumulatorParams{
Repo: repo,
Caps: BudgetCaps{MaxTokens: i64ptr(1_000_000), MaxCostUSD: f64ptr(100)},
})
_, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 10, CompletionTokens: 10, HasUsage: true}, 1)
assert.False(t, breached)
}
func i32ptr(i int32) *int32 { return &i }