You've already forked agentic-coding-workflow
bugfix
This commit is contained in:
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 owner 或 admin。
|
||||
// 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(),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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 }
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
@@ -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,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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"))
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 }
|
||||
Reference in New Issue
Block a user