From c0d7aae6fc50e2c05e58f35341f31c3f462ed7ea Mon Sep 17 00:00:00 2001 From: Haoyu Xi Date: Tue, 27 Oct 2026 21:22:46 +0800 Subject: [PATCH 1/2] feat(llm,store,worker): upstream retry, token cost metering, and diagnostics port MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit llm: doWithRetry on all four send paths — retry only 429/500/502/503/504, honour Retry-After (seconds or HTTP-date) capped by MaxDelay, backoff selects on ctx.Done so timeout/cancel/shutdown take effect mid-wait, and the request is rebuilt per attempt because http.Request.Body is not re-readable. Streaming retries only before the first delta: a mid-stream break is terminal, otherwise already-emitted text would be replayed. OnRetry callback injects metrics from the composition root (the llm package's allowed import set is empty in the AST gate). metrics/store/worker: agent_llm_{retries,prompt_tokens,completion_tokens} counters plus per-task prompt_tokens/completion_tokens columns written via a separate TaskStore.RecordUsage — metering stays out of the terminal-state transition, failed calls still bill, and a metering error never flips a successful task to failed. Old databases gain the columns through an idempotent migration. config/worker: session_history_limit (default 100) replaces the hardcoded window; pprof_addr gates a diagnostics listener on its own port and mux, off by default, warning when bound to a non-loopback address. pool.go and pool_test.go hit the 400/800 line budgets, so cancel handling and the store double moved into sibling files instead of raising the cap. --- api/loadtest_test.go | 159 ++++++++++++++++++++++ config/config.go | 57 ++++++-- config/config_p10_test.go | 102 ++++++++++++++ config/config_test.go | 4 +- llm/client.go | 78 ++++++----- llm/client_test.go | 3 +- llm/retry.go | 195 +++++++++++++++++++++++++++ llm/retry_test.go | 275 ++++++++++++++++++++++++++++++++++++++ llm/stream.go | 28 ++-- main.go | 11 ++ metrics/metrics.go | 25 ++++ pprof.go | 69 ++++++++++ pprof_test.go | 67 ++++++++++ store/sql.go | 54 ++++++-- store/store.go | 24 ++++ store/usage_test.go | 139 +++++++++++++++++++ worker/cancel.go | 42 ++++++ worker/mock_store_test.go | 58 ++++++++ worker/pool.go | 70 ++++------ worker/pool_test.go | 34 ----- worker/session.go | 14 +- worker/usage.go | 36 +++++ worker/usage_test.go | 126 +++++++++++++++++ 23 files changed, 1521 insertions(+), 149 deletions(-) create mode 100644 api/loadtest_test.go create mode 100644 config/config_p10_test.go create mode 100644 llm/retry.go create mode 100644 llm/retry_test.go create mode 100644 pprof.go create mode 100644 pprof_test.go create mode 100644 store/usage_test.go create mode 100644 worker/cancel.go create mode 100644 worker/mock_store_test.go create mode 100644 worker/usage.go create mode 100644 worker/usage_test.go diff --git a/api/loadtest_test.go b/api/loadtest_test.go new file mode 100644 index 0000000..e0fbcfa --- /dev/null +++ b/api/loadtest_test.go @@ -0,0 +1,159 @@ +package api + +import ( + "agent-runtime/agent" + "agent-runtime/store" + "bytes" + "net/http" + "net/http/httptest" + "sort" + "sync" + "sync/atomic" + "testing" + "time" +) + +// 本文件是提交路径的本地压测口径（P10）：回答「本进程在积压前能扛多少提交」， +// 而不是「LLM 有多快」。上游延迟是外部成本，不属于本进程的吞吐。 +// +// 刻意绕开中间件链（limiter 每秒 20 / 突发 40 是**有意的**限流上限， +// 不是待优化的瓶颈）：压测目的是量化落库+入队这一段自身的开销， +// 把限流混进来只会得到一个「限流器有多准」的数字。 +// +// 用法： +// go test ./api/ -run TestSubmitLatencyPercentiles -v -count=1 +// go test ./api/ -bench BenchmarkSubmit -benchmem -run '^$' + +// benchRunner 是并发安全的 TaskRunner 桩：只做原子计数，不产生任何阻塞。 +type benchRunner struct { + enqueued atomic.Int64 +} + +func (b *benchRunner) Enqueue(id string) error { + b.enqueued.Add(1) + return nil +} + +func (b *benchRunner) Trace(id string) []agent.TraceStep { return nil } + +func (b *benchRunner) Cancel(id, reason string) error { return nil } + +func (b *benchRunner) Subscribe(id string) ([]agent.TraceStep, <-chan agent.TraceStep, func(), bool) { + return nil, nil, func() {}, false +} + +func (b *benchRunner) SubscribeEvents(id string) ([]agent.TraceStep, <-chan agent.TraceEvent, func(), bool) { + return nil, nil, func() {}, false +} + +func newBenchHandler() *Handler { + return NewHandler(store.NewStore(), &benchRunner{}, "") +} + +var benchBody = []byte(`{"prompt":"benchmark prompt"}`) + +// BenchmarkSubmit 串行提交，衡量单请求的纯开销与分配次数。 +func BenchmarkSubmit(b *testing.B) { + h := newBenchHandler() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + rec := httptest.NewRecorder() + h.HandleRun(rec, httptest.NewRequest(http.MethodPost, "/run", newBodyReader())) + if rec.Code != http.StatusAccepted { + b.Fatalf("status = %d, want 202", rec.Code) + } + } +} + +// BenchmarkSubmitParallel 并发提交，衡量共享存储锁下的吞吐上限。 +func BenchmarkSubmitParallel(b *testing.B) { + h := newBenchHandler() + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + rec := httptest.NewRecorder() + h.HandleRun(rec, httptest.NewRequest(http.MethodPost, "/run", newBodyReader())) + if rec.Code != http.StatusAccepted { + b.Errorf("status = %d, want 202", rec.Code) + } + } + }) +} + +// TestSubmitLatencyPercentiles 是 P10 的实证口径： +// 32 并发提交 2000 次，断言零丢失（全部 202），并输出 p50/p95/p99 与吞吐。 +// 断言部分防回归，分位数部分供文档与面试引用（不同机器上数字会变，机制不会）。 +func TestSubmitLatencyPercentiles(t *testing.T) { + if testing.Short() { + t.Skip("load probe skipped in -short") + } + + const ( + requests = 2000 + workers = 32 + ) + + h := newBenchHandler() + jobs := make(chan struct{}) + + var ( + wg sync.WaitGroup + mu sync.Mutex + lat = make([]time.Duration, 0, requests) + notOK atomic.Int64 + start = time.Now() + ) + + for w := 0; w < workers; w++ { + wg.Add(1) + go func() { + defer wg.Done() + local := make([]time.Duration, 0, requests/workers+1) + for range jobs { + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/run", newBodyReader()) + begin := time.Now() + h.HandleRun(rec, req) + local = append(local, time.Since(begin)) + if rec.Code != http.StatusAccepted { + notOK.Add(1) + } + } + mu.Lock() + lat = append(lat, local...) + mu.Unlock() + }() + } + + for i := 0; i < requests; i++ { + jobs <- struct{}{} + } + close(jobs) + wg.Wait() + elapsed := time.Since(start) + + if got := int64(len(lat)); got != requests { + t.Fatalf("collected %d samples, want %d", got, requests) + } + // 承重断言：提交路径不允许丢单。若这里非零，说明 202 契约被破坏。 + if n := notOK.Load(); n != 0 { + t.Fatalf("%d submissions returned non-202", n) + } + + sort.Slice(lat, func(i, j int) bool { return lat[i] < lat[j] }) + p := func(q float64) time.Duration { + idx := int(q * float64(len(lat)-1)) + return lat[idx] + } + + t.Logf("submit latency over %d requests @ %d workers: p50=%v p95=%v p99=%v max=%v throughput=%.0f req/s", + requests, workers, p(0.50), p(0.95), p(0.99), lat[len(lat)-1], + float64(requests)/elapsed.Seconds()) +} + +// newBodyReader 每次返回全新的 body reader：请求体不可重读，复用会让后续请求拿到空 body。 +func newBodyReader() *bytes.Reader { + return bytes.NewReader(benchBody) +} diff --git a/config/config.go b/config/config.go index 14eabce..ce67bc9 100644 --- a/config/config.go +++ b/config/config.go @@ -27,6 +27,10 @@ type Config struct { MaxQueueSize int // 最大允许积压的任务数，超出回 503，默认 10000 WorkerConcurrency int // 默认 3；并发执行任务的 Worker 数量 MaxContextTokens int // 默认 16000；多轮对话会话上下文安全 Token 预算 + // SessionHistoryLimit 默认 100；组装会话历史时拉取的最大消息条数（ADR-0021 决策 3）。 + SessionHistoryLimit int + // PprofAddr 非空则在独立端口暴露 /debug/pprof（ADR-0021 决策 4）；空 = 不启用。 + PprofAddr string } // defaultPath 是约定俗成的配置文件位置，相对进程工作目录。 @@ -100,19 +104,35 @@ func Load(path string) (Config, error) { if maxContextTokens <= 0 { maxContextTokens = 16000 } + sessionHistoryLimit := fileCfg.SessionHistoryLimit + if envHistory := os.Getenv("SESSION_HISTORY_LIMIT"); envHistory != "" { + v, err := strconv.Atoi(envHistory) + if err != nil { + return Config{}, fmt.Errorf("config: invalid SESSION_HISTORY_LIMIT %q: %w", envHistory, err) + } + if v <= 0 { + return Config{}, fmt.Errorf("config: SESSION_HISTORY_LIMIT must be positive: %d", v) + } + sessionHistoryLimit = v + } + if sessionHistoryLimit <= 0 { + sessionHistoryLimit = 100 + } return Config{ - DeepSeekAPIKey: firstNonEmpty(os.Getenv("DEEPSEEK_API_KEY"), fileCfg.DeepSeekAPIKey), - APIAuthKey: firstNonEmpty(os.Getenv("API_AUTH_KEY"), fileCfg.APIAuthKey), - DbPath: firstNonEmpty(os.Getenv("DB_PATH"), fileCfg.DbPath), - TaskTimeoutSeconds: timeoutSec, - EmbedBaseURL: firstNonEmpty(os.Getenv("EMBED_BASE_URL"), fileCfg.EmbedBaseURL), - EmbedAPIKey: firstNonEmpty(os.Getenv("EMBED_API_KEY"), fileCfg.EmbedAPIKey), - EmbedModel: firstNonEmpty(os.Getenv("EMBED_MODEL"), fileCfg.EmbedModel), - EmbedDim: embedDim, - Model: firstNonEmpty(os.Getenv("MODEL"), fileCfg.Model), - MaxQueueSize: maxQueueSize, - WorkerConcurrency: workerConcurrency, - MaxContextTokens: maxContextTokens, + DeepSeekAPIKey: firstNonEmpty(os.Getenv("DEEPSEEK_API_KEY"), fileCfg.DeepSeekAPIKey), + APIAuthKey: firstNonEmpty(os.Getenv("API_AUTH_KEY"), fileCfg.APIAuthKey), + DbPath: firstNonEmpty(os.Getenv("DB_PATH"), fileCfg.DbPath), + TaskTimeoutSeconds: timeoutSec, + EmbedBaseURL: firstNonEmpty(os.Getenv("EMBED_BASE_URL"), fileCfg.EmbedBaseURL), + EmbedAPIKey: firstNonEmpty(os.Getenv("EMBED_API_KEY"), fileCfg.EmbedAPIKey), + EmbedModel: firstNonEmpty(os.Getenv("EMBED_MODEL"), fileCfg.EmbedModel), + EmbedDim: embedDim, + Model: firstNonEmpty(os.Getenv("MODEL"), fileCfg.Model), + MaxQueueSize: maxQueueSize, + WorkerConcurrency: workerConcurrency, + MaxContextTokens: maxContextTokens, + SessionHistoryLimit: sessionHistoryLimit, + PprofAddr: firstNonEmpty(os.Getenv("PPROF_ADDR"), fileCfg.PprofAddr), }, nil } @@ -213,8 +233,19 @@ func loadFile(path string) (Config, error) { return cfg, fmt.Errorf("config: %s:%d max_context_tokens 必须大于 0：%d", path, lineNo, n) } cfg.MaxContextTokens = n + case "session_history_limit": + n, err := strconv.Atoi(value) + if err != nil { + return cfg, fmt.Errorf("config: %s:%d session_history_limit 必须是整数：%q", path, lineNo, value) + } + if n <= 0 { + return cfg, fmt.Errorf("config: %s:%d session_history_limit 必须大于 0：%d", path, lineNo, n) + } + cfg.SessionHistoryLimit = n + case "pprof_addr": + cfg.PprofAddr = value default: - return cfg, fmt.Errorf("config: %s:%d 未知配置项 %q（只支持 deepseek_api_key、api_auth_key、db_path、task_timeout_seconds、embed_base_url、embed_api_key、embed_model、embed_dim、model、max_queue_size、worker_concurrency、max_context_tokens）", path, lineNo, key) + return cfg, fmt.Errorf("config: %s:%d 未知配置项 %q（只支持 deepseek_api_key、api_auth_key、db_path、task_timeout_seconds、embed_base_url、embed_api_key、embed_model、embed_dim、model、max_queue_size、worker_concurrency、max_context_tokens、session_history_limit、pprof_addr）", path, lineNo, key) } } if err := sc.Err(); err != nil { diff --git a/config/config_p10_test.go b/config/config_p10_test.go new file mode 100644 index 0000000..f7a22e5 --- /dev/null +++ b/config/config_p10_test.go @@ -0,0 +1,102 @@ +package config + +import ( + "path/filepath" + "strings" + "testing" +) + +// TestLoad_SessionHistoryLimitFromFile 会话历史窗口是模型能力参数（ADR-0021 决策 3）， +// 必须能只改配置而不用改代码发版。 +func TestLoad_SessionHistoryLimitFromFile(t *testing.T) { + isolateEnv(t) + p := writeCfg(t, "session_history_limit: 250\n") + + cfg, err := Load(p) + if err != nil { + t.Fatal(err) + } + if cfg.SessionHistoryLimit != 250 { + t.Errorf("SessionHistoryLimit = %d, want 250", cfg.SessionHistoryLimit) + } +} + +// TestLoad_SessionHistoryLimitEnvPriority 环境变量优先于文件。 +func TestLoad_SessionHistoryLimitEnvPriority(t *testing.T) { + isolateEnv(t) + p := writeCfg(t, "session_history_limit: 250\n") + t.Setenv("SESSION_HISTORY_LIMIT", "40") + + cfg, err := Load(p) + if err != nil { + t.Fatal(err) + } + if cfg.SessionHistoryLimit != 40 { + t.Errorf("SessionHistoryLimit = %d, want 40（环境变量优先）", cfg.SessionHistoryLimit) + } +} + +// TestLoad_SessionHistoryLimitValidation 非法值必须启动失败而不是静默回落到默认： +// 静默回落会让「我配了 0」变成「跑了 100」，排障时极难发现。 +func TestLoad_SessionHistoryLimitValidation(t *testing.T) { + isolateEnv(t) + if _, err := Load(writeCfg(t, "session_history_limit: 0\n")); err == nil { + t.Error("session_history_limit: 0 应报错") + } + if _, err := Load(writeCfg(t, "session_history_limit: abc\n")); err == nil { + t.Error("session_history_limit: abc 应报错") + } + + t.Setenv("SESSION_HISTORY_LIMIT", "-5") + if _, err := Load(filepath.Join(t.TempDir(), "nope.yaml")); err == nil { + t.Error("SESSION_HISTORY_LIMIT=-5 应报错") + } + t.Setenv("SESSION_HISTORY_LIMIT", "abc") + if _, err := Load(filepath.Join(t.TempDir(), "nope.yaml")); err == nil { + t.Error("SESSION_HISTORY_LIMIT=abc 应报错") + } +} + +// TestLoad_PprofAddr 诊断端口默认关闭（空），文件与环境变量都能开启。 +func TestLoad_PprofAddr(t *testing.T) { + isolateEnv(t) + cfg, err := Load(filepath.Join(t.TempDir(), "nope.yaml")) + if err != nil { + t.Fatal(err) + } + if cfg.PprofAddr != "" { + t.Errorf("PprofAddr = %q, want 空（默认关闭）", cfg.PprofAddr) + } + + cfg, err = Load(writeCfg(t, "pprof_addr: 127.0.0.1:6060\n")) + if err != nil { + t.Fatal(err) + } + if cfg.PprofAddr != "127.0.0.1:6060" { + t.Errorf("PprofAddr = %q, want 127.0.0.1:6060", cfg.PprofAddr) + } + + t.Setenv("PPROF_ADDR", "localhost:6061") + cfg, err = Load(filepath.Join(t.TempDir(), "nope.yaml")) + if err != nil { + t.Fatal(err) + } + if cfg.PprofAddr != "localhost:6061" { + t.Errorf("PprofAddr = %q, want localhost:6061（环境变量优先）", cfg.PprofAddr) + } +} + +// TestLoad_UnknownKeyErrorMentionsNewKeys 未知项报错文案必须列出新键， +// 否则运维照旧文案排查会漏掉 session_history_limit / pprof_addr。 +func TestLoad_UnknownKeyErrorMentionsNewKeys(t *testing.T) { + isolateEnv(t) + _, err := Load(writeCfg(t, "nope_key: 1\n")) + if err == nil { + t.Fatal("expected error for unknown key") + } + for _, want := range []string{"session_history_limit", "pprof_addr"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("error message missing %q: %v", want, err) + } + } +} diff --git a/config/config_test.go b/config/config_test.go index 5f22675..0cfc542 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -21,6 +21,8 @@ func isolateEnv(t *testing.T) { t.Setenv("MAX_QUEUE_SIZE", "") t.Setenv("WORKER_CONCURRENCY", "") t.Setenv("MAX_CONTEXT_TOKENS", "") + t.Setenv("SESSION_HISTORY_LIMIT", "") + t.Setenv("PPROF_ADDR", "") } func writeCfg(t *testing.T, content string) string { @@ -38,7 +40,7 @@ func TestLoad_MissingFileIsNotAnError(t *testing.T) { if err != nil { t.Fatalf("配置文件可选，不应报错：%v", err) } - want := Config{MaxQueueSize: 10000, WorkerConcurrency: 3, MaxContextTokens: 16000} + want := Config{MaxQueueSize: 10000, WorkerConcurrency: 3, MaxContextTokens: 16000, SessionHistoryLimit: 100} if cfg != want { t.Errorf("cfg = %+v, want %+v", cfg, want) } diff --git a/llm/client.go b/llm/client.go index 041e559..910df4c 100644 --- a/llm/client.go +++ b/llm/client.go @@ -77,6 +77,10 @@ type Client struct { seed *int64 http *http.Client // 复用连接，且用来设超时 + // retry 控制上游限流/瞬时故障的重试节奏；OnRetry 在每次退避前回调（ADR-0021）。 + retry RetryPolicy + OnRetry RetryFunc + // 向量通道独立凭证（与 chat 分离，ADR-0017 决策 3）： // embedBaseURL 或 embedModel 为空 → Embed 返回 ErrEmbedderNotConfigured。 embedBaseURL string @@ -91,6 +95,9 @@ type Config struct { Temperature *float64 Seed *int64 HTTPTimeout time.Duration + // Retry 零值即生产默认（3 次尝试 / 500ms 起 / 5s 封顶）；OnRetry 可为 nil。 + Retry RetryPolicy + OnRetry RetryFunc // 向量通道独立凭证：默认全空 = 仅词法检索（lexical only，ADR-0017 决策 3）。 EmbedBaseURL string EmbedAPIKey string @@ -110,6 +117,8 @@ func NewClient(cfg Config) *Client { temperature: cfg.Temperature, seed: cfg.Seed, http: &http.Client{Timeout: timeout}, + retry: cfg.Retry, + OnRetry: cfg.OnRetry, embedBaseURL: cfg.EmbedBaseURL, embedAPIKey: cfg.EmbedAPIKey, embedModel: cfg.EmbedModel, @@ -132,20 +141,19 @@ func (c *Client) Chat(ctx context.Context, prompt string) (string, error) { return "", fmt.Errorf("llm: marshal request: %w", err) } - // 3. http.NewRequestWithContext 建请求（POST、URL、body） - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, joinEndpoint(c.baseURL, chatCompletionsPath), bytes.NewReader(chatReqJSON)) - if err != nil { - return "", fmt.Errorf("llm: create request: %w", err) - } - - // 4. 设两个 header：Authorization、Content-Type - httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) - httpReq.Header.Set("Content-Type", "application/json") - - // 5. c.http.Do(req) 发送；检查 StatusCode - resp, err := c.http.Do(httpReq) + // 3~5. 构造请求 → 发送 → 检查 StatusCode。发送统一走 doWithRetry： + // 429/5xx 按 Retry-After 与指数退避重试，退避期间可被 ctx 打断（ADR-0021）。 + resp, err := c.doWithRetry(ctx, func() (*http.Request, error) { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, joinEndpoint(c.baseURL, chatCompletionsPath), bytes.NewReader(chatReqJSON)) + if err != nil { + return nil, fmt.Errorf("llm: create request: %w", err) + } + httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) + httpReq.Header.Set("Content-Type", "application/json") + return httpReq, nil + }) if err != nil { - return "", fmt.Errorf("llm: send request: %w", err) + return "", err } defer resp.Body.Close() @@ -191,18 +199,20 @@ func (c *Client) Embed(ctx context.Context, texts []string) ([][]float32, error) return nil, fmt.Errorf("llm: marshal request: %w", err) } - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, joinEndpoint(c.embedBaseURL, embeddingsPath), bytes.NewReader(reqJSON)) - if err != nil { - return nil, fmt.Errorf("llm: create request: %w", err) - } - if c.embedAPIKey != "" { - httpReq.Header.Set("Authorization", "Bearer "+c.embedAPIKey) - } - httpReq.Header.Set("Content-Type", "application/json") - - resp, err := c.http.Do(httpReq) + // 发送同样走 doWithRetry：嵌入端点被限流时整批重试，而不是让检索静默退化。 + resp, err := c.doWithRetry(ctx, func() (*http.Request, error) { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, joinEndpoint(c.embedBaseURL, embeddingsPath), bytes.NewReader(reqJSON)) + if err != nil { + return nil, fmt.Errorf("llm: create request: %w", err) + } + if c.embedAPIKey != "" { + httpReq.Header.Set("Authorization", "Bearer "+c.embedAPIKey) + } + httpReq.Header.Set("Content-Type", "application/json") + return httpReq, nil + }) if err != nil { - return nil, fmt.Errorf("llm: send request: %w", err) + return nil, err } defer resp.Body.Close() limited := io.LimitReader(resp.Body, maxResponseBytes) @@ -299,16 +309,18 @@ func (c *Client) ChatWithTools(ctx context.Context, msgs []Message, tools []Tool return nil, fmt.Errorf("llm: marshal request: %w", err) } - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, joinEndpoint(c.baseURL, chatCompletionsPath), bytes.NewReader(body)) - if err != nil { - return nil, fmt.Errorf("llm: create request: %w", err) - } - httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) - httpReq.Header.Set("Content-Type", "application/json") - - resp, err := c.http.Do(httpReq) + // 发送走 doWithRetry（ADR-0021）：429/5xx 退避重试，请求体每次重建。 + resp, err := c.doWithRetry(ctx, func() (*http.Request, error) { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, joinEndpoint(c.baseURL, chatCompletionsPath), bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("llm: create request: %w", err) + } + httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) + httpReq.Header.Set("Content-Type", "application/json") + return httpReq, nil + }) if err != nil { - return nil, fmt.Errorf("llm: send request: %w", err) + return nil, err } defer resp.Body.Close() limited := io.LimitReader(resp.Body, maxResponseBytes) diff --git a/llm/client_test.go b/llm/client_test.go index b64e672..fd79e6a 100644 --- a/llm/client_test.go +++ b/llm/client_test.go @@ -140,6 +140,7 @@ func TestEmbed_Success(t *testing.T) { } // TestEmbed_APIError 验证非 2xx → *APIError 且带 StatusCode。 +// Retry 显式设为 1 次尝试：本测试断言的是错误映射，不应依赖退避节奏（重试行为见 retry_test.go）。 func TestEmbed_APIError(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusInternalServerError) @@ -147,7 +148,7 @@ func TestEmbed_APIError(t *testing.T) { })) defer srv.Close() - c := NewClient(Config{EmbedBaseURL: srv.URL, EmbedAPIKey: "k", EmbedModel: "m"}) + c := NewClient(Config{EmbedBaseURL: srv.URL, EmbedModel: "m", Retry: RetryPolicy{MaxAttempts: 1}}) _, err := c.Embed(context.Background(), []string{"x"}) if err == nil { t.Fatal("expected error, got nil") diff --git a/llm/retry.go b/llm/retry.go new file mode 100644 index 0000000..e945205 --- /dev/null +++ b/llm/retry.go @@ -0,0 +1,195 @@ +package llm + +import ( + "context" + "fmt" + "io" + "net/http" + "strconv" + "strings" + "time" +) + +// RetryPolicy 描述上游限流/瞬时故障的重试节奏。 +// +// 语义：MaxAttempts 是「总尝试次数（含首次）」，<=1 表示不重试。 +// 退避为 BaseDelay × 2^(attempt-1)，单次等待上限 MaxDelay。 +type RetryPolicy struct { + MaxAttempts int + BaseDelay time.Duration + MaxDelay time.Duration +} + +// RetryFunc 在每次决定重试前被调用，用于外露重试信号（指标/日志）。可为 nil。 +// statusCode 为 0 表示失败发生在连接建立阶段（没有拿到 HTTP 响应）。 +type RetryFunc func(attempt int, statusCode int, delay time.Duration) + +// DefaultRetryPolicy 返回生产默认值：最多 3 次尝试，退避 500ms / 1s，单次等待上限 5s。 +// +// 上限存在的原因：上游可以用 Retry-After 要求我们等 60s，但一个 worker 被占住 +// 60s 的代价高于放弃这次重试——任务超时预算（默认 30s）本来就撑不住。 +func DefaultRetryPolicy() RetryPolicy { + return RetryPolicy{ + MaxAttempts: 3, + BaseDelay: 500 * time.Millisecond, + MaxDelay: 5 * time.Second, + } +} + +// normalized 补齐零值字段，使 Config 只需显式覆盖想改的那一项。 +func (p RetryPolicy) normalized() RetryPolicy { + def := DefaultRetryPolicy() + if p.MaxAttempts <= 0 { + p.MaxAttempts = def.MaxAttempts + } + if p.BaseDelay <= 0 { + p.BaseDelay = def.BaseDelay + } + if p.MaxDelay <= 0 { + p.MaxDelay = def.MaxDelay + } + if p.MaxDelay < p.BaseDelay { + p.MaxDelay = p.BaseDelay + } + return p +} + +// backoff 返回第 attempt 次失败后的等待时长（attempt 从 1 开始）。 +func (p RetryPolicy) backoff(attempt int) time.Duration { + d := p.BaseDelay + for i := 1; i < attempt && d < p.MaxDelay; i++ { + d *= 2 + } + if d > p.MaxDelay { + d = p.MaxDelay + } + return d +} + +// retryableStatus 判定该状态码是否值得重试。 +// +// 只收「上游自己说它现在不行」的信号：限流(429) 与 5xx 网关/服务错误。 +// 其余 4xx（400 请求体错、401/403 鉴权错、404、422 语义错）重试一万次也是 +// 同一个错，只会把一次失败放大成 N 次计费，因此一律不重试。 +func retryableStatus(code int) bool { + switch code { + case http.StatusTooManyRequests, + http.StatusInternalServerError, + http.StatusBadGateway, + http.StatusServiceUnavailable, + http.StatusGatewayTimeout: + return true + default: + return false + } +} + +// parseRetryAfter 解析 Retry-After 头：既支持 delta-seconds 也支持 HTTP-date。 +// 第二个返回值为 false 表示头缺失或无法解析，调用方应回落到指数退避。 +func parseRetryAfter(v string, now time.Time) (time.Duration, bool) { + v = strings.TrimSpace(v) + if v == "" { + return 0, false + } + if secs, err := strconv.Atoi(v); err == nil { + if secs < 0 { + return 0, false + } + return time.Duration(secs) * time.Second, true + } + if t, err := http.ParseTime(v); err == nil { + d := t.Sub(now) + if d < 0 { + d = 0 + } + return d, true + } + return 0, false +} + +// doWithRetry 发送请求，并对可重试的失败按策略退避重试。 +// +// newReq 每次调用都必须返回一个全新的 *http.Request：请求体是不可重读的 +// io.Reader，复用同一个 request 会让第二次发送拿到空 body。 +// +// 返回值契约（与 P9 之前保持一致，上层无需改动）： +// - 成功或不可重试的状态码 → (resp, nil)，由调用方检查 StatusCode 并构造 APIError； +// - 重试耗尽且最后一次仍拿到响应 → (resp, nil)，仍由调用方转 APIError； +// - 重试耗尽且最后一次是连接错误 / 上下文结束 → (nil, err)，err 用 %w 包住原因。 +func (c *Client) doWithRetry(ctx context.Context, newReq func() (*http.Request, error)) (*http.Response, error) { + policy := c.retry.normalized() + + for attempt := 1; ; attempt++ { + req, reqErr := newReq() + if reqErr != nil { + return nil, reqErr // 构造请求失败是本地错误，重试无意义 + } + + resp, err := c.http.Do(req) + + if err == nil && !retryableStatus(resp.StatusCode) { + return resp, nil + } + + // 上下文已结束（超时 / 用户取消 / 停机）：立即返回，不再退避空等。 + if ctxErr := ctx.Err(); ctxErr != nil { + closeBody(resp) + return nil, fmt.Errorf("llm: send request: %w", ctxErr) + } + + last := attempt >= policy.MaxAttempts + if last { + if err != nil { + return nil, fmt.Errorf("llm: send request: %w", err) + } + // 把最后一次响应原样交回：错误文案与 APIError 语义保持既有契约。 + return resp, nil + } + + delay := policy.backoff(attempt) + status := 0 + if err == nil { + status = resp.StatusCode + if d, ok := parseRetryAfter(resp.Header.Get("Retry-After"), time.Now()); ok { + delay = d + } + if delay > policy.MaxDelay { + delay = policy.MaxDelay + } + // 排空并关闭，避免连接泄漏到下一次尝试。 + closeBody(resp) + } + + if c.OnRetry != nil { + c.OnRetry(attempt, status, delay) + } + if sleepErr := sleepCtx(ctx, delay); sleepErr != nil { + return nil, fmt.Errorf("llm: send request: %w", sleepErr) + } + } +} + +// closeBody 排空少量响应体后关闭，使底层连接可以复用而不是被丢弃。 +func closeBody(resp *http.Response) { + if resp == nil || resp.Body == nil { + return + } + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, maxResponseBytes)) + _ = resp.Body.Close() +} + +// sleepCtx 等待 d，但可被 ctx 提前打断（任务超时 / 用户取消 / 停机广播）。 +// 返回 nil 表示等满了；非 nil 是 ctx.Err()。 +func sleepCtx(ctx context.Context, d time.Duration) error { + if d <= 0 { + return ctx.Err() + } + t := time.NewTimer(d) + defer t.Stop() + select { + case <-t.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} diff --git a/llm/retry_test.go b/llm/retry_test.go new file mode 100644 index 0000000..762888a --- /dev/null +++ b/llm/retry_test.go @@ -0,0 +1,275 @@ +package llm + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" +) + +// fastRetry 是测试用退避：最大尝试次数可控，退避时长压到微秒级， +// 使断言只依赖「重试了几次」而不是「等了多久」。 +func fastRetry(maxAttempts int) RetryPolicy { + return RetryPolicy{MaxAttempts: maxAttempts, BaseDelay: time.Microsecond, MaxDelay: 2 * time.Microsecond} +} + +// TestChatWithTools_RetriesOnRateLimit 是 ADR-0021 决策 1 的承重墙： +// 上游 429 一次后恢复，任务必须成功，且重试信号经 OnRetry 外露。 +func TestChatWithTools_RetriesOnRateLimit(t *testing.T) { + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if atomic.AddInt32(&calls, 1) == 1 { + w.Header().Set("Retry-After", "0") + w.WriteHeader(http.StatusTooManyRequests) + fmt.Fprint(w, `{"error":"rate limited"}`) + return + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"ok"}}],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10}}`) + })) + defer srv.Close() + + var gotStatus, gotAttempt int + var gotDelay time.Duration + c := NewClient(Config{ + BaseURL: srv.URL, + Model: "m", + Retry: fastRetry(3), + OnRetry: func(attempt, status int, delay time.Duration) { + gotAttempt, gotStatus, gotDelay = attempt, status, delay + }, + }) + + turn, err := c.ChatWithTools(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil) + if err != nil { + t.Fatalf("ChatWithTools: %v", err) + } + if turn.Content != "ok" { + t.Errorf("content = %q, want ok", turn.Content) + } + if n := atomic.LoadInt32(&calls); n != 2 { + t.Errorf("upstream calls = %d, want 2 (1 次 429 + 1 次成功)", n) + } + if gotAttempt != 1 || gotStatus != http.StatusTooManyRequests { + t.Errorf("OnRetry(attempt=%d,status=%d), want (1,429)", gotAttempt, gotStatus) + } + // Retry-After: 0 必须被采信并覆盖指数退避（否则这里会是 fastRetry 的 1µs 以上基数）。 + if gotDelay != 0 { + t.Errorf("delay = %v, want 0（Retry-After: 0 优先于指数退避）", gotDelay) + } +} + +// TestChatWithTools_NoRetryOnClientError 守住「4xx 客户端错误不重试」： +// 重试鉴权错/请求体错只会把一次失败放大成 N 次计费。 +func TestChatWithTools_NoRetryOnClientError(t *testing.T) { + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&calls, 1) + w.WriteHeader(http.StatusBadRequest) + fmt.Fprint(w, `{"error":"bad request"}`) + })) + defer srv.Close() + + c := NewClient(Config{BaseURL: srv.URL, Model: "m", Retry: fastRetry(3)}) + _, err := c.ChatWithTools(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil) + + var apiErr *APIError + if !errors.As(err, &apiErr) || apiErr.StatusCode != http.StatusBadRequest { + t.Fatalf("err = %v, want *APIError 400", err) + } + if n := atomic.LoadInt32(&calls); n != 1 { + t.Errorf("upstream calls = %d, want 1（400 不可重试）", n) + } +} + +// TestChatWithTools_RetryExhaustedReturnsAPIError 锁定错误契约： +// 重试耗尽后仍返回最后一次响应的 *APIError，上层 worker 的终态映射无需改动。 +func TestChatWithTools_RetryExhaustedReturnsAPIError(t *testing.T) { + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&calls, 1) + w.WriteHeader(http.StatusServiceUnavailable) + fmt.Fprint(w, `{"error":"overloaded"}`) + })) + defer srv.Close() + + c := NewClient(Config{BaseURL: srv.URL, Model: "m", Retry: fastRetry(3)}) + _, err := c.ChatWithTools(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil) + + var apiErr *APIError + if !errors.As(err, &apiErr) || apiErr.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("err = %v, want *APIError 503", err) + } + if n := atomic.LoadInt32(&calls); n != 3 { + t.Errorf("upstream calls = %d, want 3（MaxAttempts 含首次）", n) + } +} + +// TestChatWithTools_RetryAbortsOnContextCancel 守住「退避必须可被取消」： +// 任务超时 / 用户抢占 / 停机广播在退避期间同样生效，不允许出现「已取消还在睡」的假在飞。 +func TestChatWithTools_RetryAbortsOnContextCancel(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusTooManyRequests) + })) + defer srv.Close() + + // 退避设为 5s：若实现忽略了 ctx，本测试会明显超时。 + c := NewClient(Config{BaseURL: srv.URL, Model: "m", Retry: RetryPolicy{MaxAttempts: 3, BaseDelay: 5 * time.Second, MaxDelay: 5 * time.Second}}) + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + start := time.Now() + _, err := c.ChatWithTools(ctx, []Message{{Role: "user", Content: "hi"}}, nil) + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected error, got nil") + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("err = %v, want 可被 errors.Is(DeadlineExceeded) 识别", err) + } + if elapsed > time.Second { + t.Errorf("elapsed = %v, want < 1s（退避未响应 ctx 取消）", elapsed) + } +} + +// TestChatWithToolsStream_RetriesBeforeFirstDelta 验证流式在「尚未下发任何 delta」时 +// 可以重试：429 后恢复，客户端只看到一份完整文本，不出现重复。 +func TestChatWithToolsStream_RetriesBeforeFirstDelta(t *testing.T) { + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if atomic.AddInt32(&calls, 1) == 1 { + w.Header().Set("Retry-After", "0") + w.WriteHeader(http.StatusTooManyRequests) + return + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"he\"}}]}\n\n") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"llo\"}}]}\n\n") + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer srv.Close() + + var deltas []string + c := NewClient(Config{BaseURL: srv.URL, Model: "m", Retry: fastRetry(3)}) + turn, err := c.ChatWithToolsStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, func(d string) { + deltas = append(deltas, d) + }) + if err != nil { + t.Fatalf("ChatWithToolsStream: %v", err) + } + if turn.Content != "hello" { + t.Errorf("content = %q, want hello", turn.Content) + } + if len(deltas) != 2 { + t.Errorf("deltas = %v, want 恰好 2 段（重试不得重放已下发的增量）", deltas) + } + if n := atomic.LoadInt32(&calls); n != 2 { + t.Errorf("upstream calls = %d, want 2", n) + } +} + +// TestChatWithToolsStream_NoRetryAfterPartialStream 锁定流式的取舍： +// 200 之后中途断流不重试——重试会把已经打出去的文本再打一遍。 +func TestChatWithToolsStream_NoRetryAfterPartialStream(t *testing.T) { + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&calls, 1) + // 声明比实际更长的 Content-Length：客户端读到半截即 unexpected EOF， + // 模拟上游流式传输中途断连。 + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Content-Length", "4096") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"par\"}}]}\n\n") + })) + defer srv.Close() + + var deltas []string + c := NewClient(Config{BaseURL: srv.URL, Model: "m", Retry: fastRetry(3)}) + _, err := c.ChatWithToolsStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, func(d string) { + deltas = append(deltas, d) + }) + + if err == nil { + t.Fatal("expected error on truncated stream, got nil") + } + if n := atomic.LoadInt32(&calls); n != 1 { + t.Errorf("upstream calls = %d, want 1（已产出 delta 后不得重试）", n) + } + if len(deltas) != 1 { + t.Errorf("deltas = %v, want 恰好 1 段（不得重复下发）", deltas) + } +} + +// TestParseRetryAfter 覆盖 delta-seconds / HTTP-date / 非法值三种形态。 +func TestParseRetryAfter(t *testing.T) { + now := time.Date(2026, 10, 27, 12, 0, 0, 0, time.UTC) + cases := []struct { + name string + in string + want time.Duration + wantOK bool + }{ + {"seconds", "3", 3 * time.Second, true}, + {"zero", "0", 0, true}, + {"padded", " 7 ", 7 * time.Second, true}, + {"http-date", now.Add(10 * time.Second).Format(http.TimeFormat), 10 * time.Second, true}, + {"past-date", now.Add(-time.Minute).Format(http.TimeFormat), 0, true}, + {"empty", "", 0, false}, + {"garbage", "soon", 0, false}, + {"negative", "-1", 0, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, ok := parseRetryAfter(tc.in, now) + if ok != tc.wantOK || got != tc.want { + t.Errorf("parseRetryAfter(%q) = (%v,%v), want (%v,%v)", tc.in, got, ok, tc.want, tc.wantOK) + } + }) + } +} + +// TestRetryPolicyBackoff 锁定指数退避与封顶：1×, 2×, 4×… 不超过 MaxDelay。 +func TestRetryPolicyBackoff(t *testing.T) { + p := RetryPolicy{MaxAttempts: 5, BaseDelay: 100 * time.Millisecond, MaxDelay: 350 * time.Millisecond} + want := []time.Duration{100 * time.Millisecond, 200 * time.Millisecond, 350 * time.Millisecond, 350 * time.Millisecond} + for i, w := range want { + if got := p.backoff(i + 1); got != w { + t.Errorf("backoff(%d) = %v, want %v", i+1, got, w) + } + } +} + +// TestRetryPolicyNormalized 验证零值 Config 也能拿到生产默认（调用方只需覆盖想改的项）。 +func TestRetryPolicyNormalized(t *testing.T) { + got := RetryPolicy{}.normalized() + def := DefaultRetryPolicy() + if got != def { + t.Errorf("normalized() = %+v, want %+v", got, def) + } + // MaxDelay < BaseDelay 的非法组合必须被拉平，否则退避会「越等越短」。 + got = RetryPolicy{MaxAttempts: 2, BaseDelay: time.Second, MaxDelay: time.Millisecond}.normalized() + if got.MaxDelay != got.BaseDelay { + t.Errorf("MaxDelay = %v, want 拉平到 BaseDelay %v", got.MaxDelay, got.BaseDelay) + } +} + +// TestRetryableStatus 白名单：只重试「上游自己说现在不行」的信号。 +func TestRetryableStatus(t *testing.T) { + retryable := []int{429, 500, 502, 503, 504} + for _, code := range retryable { + if !retryableStatus(code) { + t.Errorf("retryableStatus(%d) = false, want true", code) + } + } + notRetryable := []int{200, 400, 401, 403, 404, 422} + for _, code := range notRetryable { + if retryableStatus(code) { + t.Errorf("retryableStatus(%d) = true, want false", code) + } + } +} diff --git a/llm/stream.go b/llm/stream.go index 2bceff1..41c20f6 100644 --- a/llm/stream.go +++ b/llm/stream.go @@ -74,19 +74,23 @@ func (c *Client) ChatWithToolsStream(ctx context.Context, msgs []Message, tools return nil, fmt.Errorf("llm: marshal request: %w", err) } - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, joinEndpoint(c.baseURL, chatCompletionsPath), bytes.NewReader(body)) - if err != nil { - return nil, fmt.Errorf("llm: create request: %w", err) - } - - // 显式指定 Accept: text/event-stream，告知中间反向代理和上游网关保持非缓冲流式下发 - httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Accept", "text/event-stream") - - resp, err := c.http.Do(httpReq) + // 流式发送同样走 doWithRetry，但不变式不同：只在「响应体尚未被消费」时重试。 + // 一旦 delta 已经下发给客户端，任何重试都会把同一段文本再打一遍， + // 打字机体验直接崩坏。doWithRetry 仅在拿到非 200 响应或连接建立失败时重试， + // 天然满足该约束；200 之后的中途断流不重试，按失败处理（ADR-0021）。 + resp, err := c.doWithRetry(ctx, func() (*http.Request, error) { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, joinEndpoint(c.baseURL, chatCompletionsPath), bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("llm: create request: %w", err) + } + // 显式指定 Accept: text/event-stream，告知中间反向代理和上游网关保持非缓冲流式下发 + httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "text/event-stream") + return httpReq, nil + }) if err != nil { - return nil, fmt.Errorf("llm: send request: %w", err) + return nil, err } defer resp.Body.Close() diff --git a/main.go b/main.go index dc08259..4288865 100644 --- a/main.go +++ b/main.go @@ -79,6 +79,12 @@ func runWithArgs(args []string) int { EmbedAPIKey: conf.EmbedAPIKey, EmbedModel: conf.EmbedModel, EmbedDim: conf.EmbedDim, + // 重试信号注入指标（ADR-0021）：llm 包在架构门禁里不允许 import metrics， + // 因此用回调做依赖倒置，与 metrics.SetPendingSource 同一手法。 + OnRetry: func(attempt, status int, delay time.Duration) { + metrics.IncLLMRetry() + slog.Warn("llm: upstream retry", "attempt", attempt, "status", status, "delay", delay.String()) + }, } ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) @@ -133,6 +139,7 @@ func runWithArgs(args []string) int { p := worker.NewPool(conf.WorkerConcurrency, taskTimeout, s, llm.NewClient(cfg)) p.SetMaxQueueSize(conf.MaxQueueSize) p.SetMaxContextTokens(conf.MaxContextTokens) + p.SetSessionHistoryLimit(conf.SessionHistoryLimit) if ss, ok := s.(store.SessionStore); ok { p.SetSessionStore(ss) } @@ -151,6 +158,10 @@ func runWithArgs(args []string) int { } p.Start() + // 诊断端口（ADR-0021 决策 4）：默认关闭；开启时独立端口，停机一并关闭。 + if pprofSrv := startPprof(conf.PprofAddr); pprofSrv != nil { + defer func() { _ = pprofSrv.Close() }() + } // pending 可能超过队列容量，阻塞式重灌放到后台，不阻塞 HTTP 监听启动。 enqueueRecovered(ctx, p, pendingIDs) diff --git a/metrics/metrics.go b/metrics/metrics.go index b5eec62..9fad77a 100644 --- a/metrics/metrics.go +++ b/metrics/metrics.go @@ -19,6 +19,11 @@ var ( failed atomic.Int64 recovered atomic.Int64 + // 上游重试与 Token 成本（ADR-0021）：前者说明上游稳不稳，后者是账单口径。 + retries atomic.Int64 + promptTokens atomic.Int64 + completionTokens atomic.Int64 + // pendingSource 由装配层注入（见 SetPendingSource），在 scrape 时刻回调取积压数。 // 用回调而不是让 metrics import store：architecture_test.go 的 AST DAG 白名单 // 里 metrics 的允许 import 集为空，反向依赖会直接让架构门禁失败。 @@ -74,6 +79,23 @@ func IncRecovered(n int) { recovered.Add(int64(n)) } +// IncLLMRetry 在 llm 层决定退避重试时调用（由装配层经 llm.Config.OnRetry 注入）。 +// 它衡量的是「上游现在稳不稳」：持续上涨说明该查配额或降并发，而不是加 worker。 +func IncLLMRetry() { + retries.Add(1) +} + +// AddTokens 累加一次任务消耗的 Token 数（ADR-0021）。 +// 失败的调用同样计费，因此调用方应在拿到 usage 后立即累加，而不是等任务成功。 +func AddTokens(prompt, completion int) { + if prompt > 0 { + promptTokens.Add(int64(prompt)) + } + if completion > 0 { + completionTokens.Add(int64(completion)) + } +} + // Handler 返回 Prometheus 文本格式的指标端点（text/plain; version=0.0.4）。 // 请求上下文会传给 pendingSource，使 scrape 能随客户端断连及时释放。 func Handler() http.Handler { @@ -91,6 +113,9 @@ func expose(ctx context.Context) string { b = appendMetric(b, "agent_tasks_done", "counter", "Tasks completed successfully", done.Load()) b = appendMetric(b, "agent_tasks_failed", "counter", "Tasks that failed", failed.Load()) b = appendMetric(b, "agent_tasks_recovered", "counter", "Tasks rehydrated after restart", recovered.Load()) + b = appendMetric(b, "agent_llm_retries_total", "counter", "Upstream LLM retries performed", retries.Load()) + b = appendMetric(b, "agent_llm_prompt_tokens_total", "counter", "Prompt tokens consumed by upstream LLM calls", promptTokens.Load()) + b = appendMetric(b, "agent_llm_completion_tokens_total", "counter", "Completion tokens consumed by upstream LLM calls", completionTokens.Load()) if v, ok := pendingValue(ctx); ok { b = appendMetric(b, "agent_tasks_pending", "gauge", "Tasks waiting in the queue", v) } diff --git a/pprof.go b/pprof.go new file mode 100644 index 0000000..e6bb5cc --- /dev/null +++ b/pprof.go @@ -0,0 +1,69 @@ +package main + +import ( + "errors" + "log/slog" + "net" + "net/http" + "net/http/pprof" + "time" +) + +// 本文件承载诊断端口（ADR-0021 决策 4）：默认关闭，开启时独立端口、独立 mux。 + +// pprofMux 返回只挂诊断端点的独立 mux。 +// +// 不挂进 api.Handler.Routes() 的主 mux：主 mux 的 "/" 是公开 Demo 页， +// profile 端点混进去等于把 goroutine 快照与调用栈暴露给公网；而且 pprof 的 +// index 页会枚举全部 profile，只给入口加鉴权保护不了内容。 +// 隔离端口 + 默认关闭是更小的暴露面。 +func pprofMux() *http.ServeMux { + mux := http.NewServeMux() + mux.HandleFunc("/debug/pprof/", pprof.Index) + mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline) + mux.HandleFunc("/debug/pprof/profile", pprof.Profile) + mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol) + mux.HandleFunc("/debug/pprof/trace", pprof.Trace) + return mux +} + +// startPprof 在 addr 上启动诊断服务；addr 为空表示不启用，返回 nil。 +// 返回的 server 由调用方负责 Close（停机时关闭）。 +func startPprof(addr string) *http.Server { + if addr == "" { + return nil + } + if !isLoopbackAddr(addr) { + // 容器/内网抓 profile 是合理用法，但必须让运维在日志里看到暴露面。 + slog.Warn("pprof listening on non-loopback address; profiles expose process internals", "addr", addr) + } + srv := &http.Server{ + Addr: addr, + Handler: pprofMux(), + ReadHeaderTimeout: 5 * time.Second, + } + go func() { + if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + slog.Error("pprof server error", "addr", addr, "error", err) + } + }() + slog.Info("pprof endpoint enabled", "addr", addr) + return srv +} + +// isLoopbackAddr 判断监听地址是否只绑定回环接口。 +// ":6060" 这类空 host 表示绑定全部网卡，因此不算回环。 +func isLoopbackAddr(addr string) bool { + host, _, err := net.SplitHostPort(addr) + if err != nil { + host = addr + } + if host == "" { + return false + } + if host == "localhost" { + return true + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() +} diff --git a/pprof_test.go b/pprof_test.go new file mode 100644 index 0000000..30d985a --- /dev/null +++ b/pprof_test.go @@ -0,0 +1,67 @@ +package main + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// TestPprofMuxServesDiagnostics 验证诊断端点确实挂在独立 mux 上（ADR-0021 决策 4）。 +func TestPprofMuxServesDiagnostics(t *testing.T) { + mux := pprofMux() + + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/debug/pprof/", nil)) + if rec.Code != http.StatusOK { + t.Fatalf("index status = %d, want 200", rec.Code) + } + // index 页列出所有可用 profile：出现 goroutine 即证明标准 handler 已注册。 + if !strings.Contains(rec.Body.String(), "goroutine") { + t.Errorf("index body missing goroutine profile; got:\n%s", rec.Body.String()) + } + + rec = httptest.NewRecorder() + mux.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/debug/pprof/cmdline", nil)) + if rec.Code != http.StatusOK { + t.Errorf("cmdline status = %d, want 200", rec.Code) + } +} + +// TestPprofMuxRejectsUnknownPath 证明它是独立 mux，不继承主 mux 的路由。 +func TestPprofMuxRejectsUnknownPath(t *testing.T) { + rec := httptest.NewRecorder() + pprofMux().ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/run", nil)) + if rec.Code != http.StatusNotFound { + t.Errorf("status = %d, want 404（诊断 mux 不得暴露业务路由）", rec.Code) + } +} + +// TestIsLoopbackAddr 决定启动时是否打「暴露面」告警，必须判得准。 +func TestIsLoopbackAddr(t *testing.T) { + cases := []struct { + addr string + want bool + }{ + {"127.0.0.1:6060", true}, + {"[::1]:6060", true}, + {"localhost:6060", true}, + {":6060", false}, + {"0.0.0.0:6060", false}, + {"192.168.1.10:6060", false}, + {"6060", false}, + } + for _, tc := range cases { + if got := isLoopbackAddr(tc.addr); got != tc.want { + t.Errorf("isLoopbackAddr(%q) = %v, want %v", tc.addr, got, tc.want) + } + } +} + +// TestStartPprofDisabledByDefault 默认关闭：空地址不启动任何监听。 +func TestStartPprofDisabledByDefault(t *testing.T) { + if srv := startPprof(""); srv != nil { + _ = srv.Close() + t.Fatal("startPprof(\"\") must return nil（默认不启用诊断端口）") + } +} diff --git a/store/sql.go b/store/sql.go index c3f51c2..922e637 100644 --- a/store/sql.go +++ b/store/sql.go @@ -46,7 +46,9 @@ func NewSQLStore(path string) (*SQLStore, error) { prompt TEXT NOT NULL, result TEXT NOT NULL DEFAULT '', error TEXT NOT NULL DEFAULT '', - session_id TEXT NOT NULL DEFAULT '' + session_id TEXT NOT NULL DEFAULT '', + prompt_tokens INTEGER NOT NULL DEFAULT 0, + completion_tokens INTEGER NOT NULL DEFAULT 0 )`); err != nil { db.Close() return nil, err @@ -60,10 +62,18 @@ func NewSQLStore(path string) (*SQLStore, error) { db.Close() return nil, fmt.Errorf("sqlstore: create idx_tasks_status: %w", err) } - // 兼容已有数据库迁移，确保 session_id 列存在 - if _, err := db.Exec("ALTER TABLE tasks ADD COLUMN session_id TEXT NOT NULL DEFAULT ''"); err != nil && !strings.Contains(err.Error(), "duplicate column name") { + // 兼容已有数据库迁移：老库缺列时补列（ADR-0021 的 token 计量列同此路径）。 + if err := ensureTaskColumn(db, "ALTER TABLE tasks ADD COLUMN session_id TEXT NOT NULL DEFAULT ''", "session_id"); err != nil { db.Close() - return nil, fmt.Errorf("sqlstore: migrate tasks session_id: %w", err) + return nil, err + } + if err := ensureTaskColumn(db, "ALTER TABLE tasks ADD COLUMN prompt_tokens INTEGER NOT NULL DEFAULT 0", "prompt_tokens"); err != nil { + db.Close() + return nil, err + } + if err := ensureTaskColumn(db, "ALTER TABLE tasks ADD COLUMN completion_tokens INTEGER NOT NULL DEFAULT 0", "completion_tokens"); err != nil { + db.Close() + return nil, err } if err := initSessionSchema(db); err != nil { db.Close() @@ -72,6 +82,15 @@ func NewSQLStore(path string) (*SQLStore, error) { return &SQLStore{db: db, path: path}, nil } +// ensureTaskColumn 幂等补列：老库升级时列已存在会被 SQLite 报 duplicate column name， +// 那是预期状态而非错误；其余错误一律上抛（不静默吞掉迁移失败）。 +func ensureTaskColumn(db *sql.DB, ddl, column string) error { + if _, err := db.Exec(ddl); err != nil && !strings.Contains(err.Error(), "duplicate column name") { + return fmt.Errorf("sqlstore: migrate tasks.%s: %w", column, err) + } + return nil +} + // Close 释放底层数据库连接。进程停机时调用。 func (s *SQLStore) Close() error { return s.db.Close() @@ -109,8 +128,8 @@ func (s *SQLStore) Create(prompt string) string { func (s *SQLStore) Get(id string) (Task, error) { var t Task err := s.db.QueryRow( - "SELECT id, status, prompt, result, error, session_id FROM tasks WHERE id = ?", id, - ).Scan(&t.ID, &t.Status, &t.Prompt, &t.Result, &t.Error, &t.SessionID) + "SELECT id, status, prompt, result, error, session_id, prompt_tokens, completion_tokens FROM tasks WHERE id = ?", id, + ).Scan(&t.ID, &t.Status, &t.Prompt, &t.Result, &t.Error, &t.SessionID, &t.PromptTokens, &t.CompletionTokens) if errors.Is(err, sql.ErrNoRows) { return Task{}, ErrNotFound } @@ -230,6 +249,25 @@ func (s *SQLStore) RecoverTasks(ctx context.Context) ([]string, int, error) { return pending, int(n), nil } +// RecordUsage 记录该任务消耗的上游 Token 数；id 不存在返回 ErrNotFound。 +// 写操作走 mu 串行化（SQLite 单写者），与 Complete 同一纪律。 +func (s *SQLStore) RecordUsage(id string, promptTokens, completionTokens int) error { + s.mu.Lock() + defer s.mu.Unlock() + + res, err := s.db.Exec( + "UPDATE tasks SET prompt_tokens = ?, completion_tokens = ? WHERE id = ?", + promptTokens, completionTokens, id, + ) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return ErrNotFound + } + return nil +} + // PendingCount 返回处于 pending 状态的任务积压数。 // 纯只读统计，WAL 模式下并发安全且不阻塞写操作。 func (s *SQLStore) PendingCount(ctx context.Context) (int64, error) { @@ -257,9 +295,9 @@ func (s *SQLStore) ClaimNext(ctx context.Context) (Task, error) { var t Task err := s.db.QueryRowContext(ctx, - "SELECT id, prompt, result, error, session_id FROM tasks WHERE status = ? ORDER BY id ASC LIMIT 1", + "SELECT id, prompt, result, error, session_id, prompt_tokens, completion_tokens FROM tasks WHERE status = ? ORDER BY id ASC LIMIT 1", StatusPending, - ).Scan(&t.ID, &t.Prompt, &t.Result, &t.Error, &t.SessionID) + ).Scan(&t.ID, &t.Prompt, &t.Result, &t.Error, &t.SessionID, &t.PromptTokens, &t.CompletionTokens) if errors.Is(err, sql.ErrNoRows) { return Task{}, ErrNoTaskAvailable } diff --git a/store/store.go b/store/store.go index b5e1572..fb82e78 100644 --- a/store/store.go +++ b/store/store.go @@ -51,6 +51,10 @@ type TaskStore interface { ClaimNext(ctx context.Context) (Task, error) // PendingCount 返回当前处于 pending 状态的任务积压数。 PendingCount(ctx context.Context) (int64, error) + // RecordUsage 记录该任务消耗的上游 Token 数（ADR-0021 成本可观测）。 + // 与 Complete 分离：状态机终态裁决与计量是两个关注点， + // 且失败任务同样计费，不应被终态语义绑架。id 不存在时返回 ErrNotFound。 + RecordUsage(id string, promptTokens, completionTokens int) error } // Store 是任务的并发安全容器。 @@ -81,6 +85,9 @@ type Task struct { Result string `json:"result"` // 仅 done 时有值 Error string `json:"error"` // 仅 failed 时有值，且必须是脱敏后的粗粒度分类 SessionID string `json:"session_id,omitempty"` + // 上游 Token 消耗（ADR-0021）。失败任务也可能非零——失败的调用照样计费。 + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` } // NewStore 返回一个可用的空 Store。 @@ -164,6 +171,23 @@ func (s *Store) Complete(ID string, result string, err error) error { return nil } +// RecordUsage 记录该任务消耗的上游 Token 数；id 不存在返回 ErrNotFound。 +// 与 Complete 分离的理由见 TaskStore 接口注释：失败任务同样计费， +// 计量不应被「终态裁决」的语义绑架。 +func (s *Store) RecordUsage(id string, promptTokens, completionTokens int) error { + s.mu.Lock() + defer s.mu.Unlock() + + task, ok := s.tasks[id] + if !ok { + return ErrNotFound + } + task.PromptTokens = promptTokens + task.CompletionTokens = completionTokens + s.tasks[id] = task + return nil +} + // Get 返回任务的副本。 // 返回副本而非指针，调用方改动不会污染内部状态， // 从设计上消除了一整类数据竞争。 diff --git a/store/usage_test.go b/store/usage_test.go new file mode 100644 index 0000000..c387c03 --- /dev/null +++ b/store/usage_test.go @@ -0,0 +1,139 @@ +package store + +import ( + "context" + "database/sql" + "errors" + "path/filepath" + "testing" +) + +// TestStore_RecordUsage 内存实现：写入后回读一致，且不影响状态机字段。 +func TestStore_RecordUsage(t *testing.T) { + s := NewStore() + id := s.Create("hello") + + if err := s.RecordUsage(id, 12, 34); err != nil { + t.Fatalf("RecordUsage: %v", err) + } + task, err := s.Get(id) + if err != nil { + t.Fatalf("Get: %v", err) + } + if task.PromptTokens != 12 || task.CompletionTokens != 34 { + t.Errorf("usage = (%d,%d), want (12,34)", task.PromptTokens, task.CompletionTokens) + } + if task.Status != StatusPending { + t.Errorf("status = %q, want pending（计量不得改动状态机）", task.Status) + } +} + +// TestStore_RecordUsageNotFound 不存在的 id 不得被静默创建出来。 +func TestStore_RecordUsageNotFound(t *testing.T) { + s := NewStore() + if err := s.RecordUsage("nope", 1, 1); !errors.Is(err, ErrNotFound) { + t.Errorf("err = %v, want ErrNotFound", err) + } +} + +// TestSQLStore_RecordUsage SQLite 实现：写入后回读一致，且能覆盖已有值。 +func TestSQLStore_RecordUsage(t *testing.T) { + s := newSQLStore(t) + id := s.Create("hello") + + if err := s.RecordUsage(id, 12, 34); err != nil { + t.Fatalf("RecordUsage: %v", err) + } + task, err := s.Get(id) + if err != nil { + t.Fatalf("Get: %v", err) + } + if task.PromptTokens != 12 || task.CompletionTokens != 34 { + t.Errorf("usage = (%d,%d), want (12,34)", task.PromptTokens, task.CompletionTokens) + } + + // 同一任务重试后可能再次计量：必须是覆盖而不是累加， + // 累加会把「重试 3 次」的用量算成 3 份，账单直接翻倍。 + if err := s.RecordUsage(id, 5, 6); err != nil { + t.Fatalf("RecordUsage (overwrite): %v", err) + } + task, _ = s.Get(id) + if task.PromptTokens != 5 || task.CompletionTokens != 6 { + t.Errorf("usage after overwrite = (%d,%d), want (5,6)", task.PromptTokens, task.CompletionTokens) + } +} + +// TestSQLStore_RecordUsageNotFound 不存在 id → ErrNotFound（与内存实现同语义）。 +func TestSQLStore_RecordUsageNotFound(t *testing.T) { + s := newSQLStore(t) + if err := s.RecordUsage("nope", 1, 1); !errors.Is(err, ErrNotFound) { + t.Errorf("err = %v, want ErrNotFound", err) + } +} + +// TestSQLStore_ClaimNextCarriesUsage 抢占返回的 Task 必须带上计量列， +// 否则 worker 重放/排空路径会看到「用量归零」的假象。 +func TestSQLStore_ClaimNextCarriesUsage(t *testing.T) { + s := newSQLStore(t) + id := s.Create("hello") + if err := s.RecordUsage(id, 9, 8); err != nil { + t.Fatalf("RecordUsage: %v", err) + } + + task, err := s.ClaimNext(context.Background()) + if err != nil { + t.Fatalf("ClaimNext: %v", err) + } + if task.PromptTokens != 9 || task.CompletionTokens != 8 { + t.Errorf("claimed usage = (%d,%d), want (9,8)", task.PromptTokens, task.CompletionTokens) + } +} + +// TestSQLStore_MigratesLegacyTaskTable 是 ADR-0021 决策 2 的升级门： +// 老库（无 token 列）必须能直接打开并自动补列，否则升级即启动失败。 +func TestSQLStore_MigratesLegacyTaskTable(t *testing.T) { + path := filepath.Join(t.TempDir(), "legacy.db") + + // 手工建一个 P9 时代的 tasks 表（只有 session_id，无 token 列）。 + legacy, err := sql.Open("sqlite", path) + if err != nil { + t.Fatalf("open legacy db: %v", err) + } + if _, err := legacy.Exec(`CREATE TABLE tasks ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + status TEXT NOT NULL, + prompt TEXT NOT NULL, + result TEXT NOT NULL DEFAULT '', + error TEXT NOT NULL DEFAULT '', + session_id TEXT NOT NULL DEFAULT '' + )`); err != nil { + t.Fatalf("create legacy table: %v", err) + } + if _, err := legacy.Exec("INSERT INTO tasks (status, prompt) VALUES (?, ?)", StatusPending, "legacy"); err != nil { + t.Fatalf("seed legacy row: %v", err) + } + if err := legacy.Close(); err != nil { + t.Fatalf("close legacy db: %v", err) + } + + // 新版本打开同一个库：补列成功、老行可读且用量为 0。 + s, err := NewSQLStore(path) + if err != nil { + t.Fatalf("NewSQLStore on legacy db: %v", err) + } + defer s.Close() + + task, err := s.Get("1") + if err != nil { + t.Fatalf("Get legacy row: %v", err) + } + if task.Prompt != "legacy" { + t.Errorf("prompt = %q, want legacy", task.Prompt) + } + if task.PromptTokens != 0 || task.CompletionTokens != 0 { + t.Errorf("legacy usage = (%d,%d), want (0,0)", task.PromptTokens, task.CompletionTokens) + } + if err := s.RecordUsage("1", 3, 4); err != nil { + t.Fatalf("RecordUsage on migrated db: %v", err) + } +} diff --git a/worker/cancel.go b/worker/cancel.go new file mode 100644 index 0000000..4a403ea --- /dev/null +++ b/worker/cancel.go @@ -0,0 +1,42 @@ +package worker + +import ( + "agent-runtime/store" + "context" + "time" +) + +// 本文件承载主动抢占取消（POST /tasks/{id}/cancel）的实现。 +// 从 pool.go 同包正交拆出，只为守住「生产文件 ≤ 400 行」的架构门禁， +// 与 worker/session.go、worker/metrics.go 的拆分理由一致。 + +// userCancelCause 是用户主动取消的因果标记，用于与超时/停机取消精确区分。 +type userCancelCause struct { + reason string +} + +func (c *userCancelCause) Error() string { return "user canceled: " + c.reason } + +// Cancel 主动抢占指定在飞任务：通过 context 因果链物理中断正在进行的上游调用。 +// 任务不在飞（未开始、已终态或不存在）时返回 ErrTaskNotRunning，由 API 层映射为 409。 +func (p *Pool) Cancel(id, reason string) error { + v, ok := p.inFlight.Load(id) + if !ok { + // 任务刚被抢占、句柄尚未登记完成：短暂让出 CPU 等一次，避免误判 409。 + for i := 0; i < 20; i++ { + if tk, err := p.tasks.Get(id); err != nil || tk.Status != store.StatusRunning { + break + } + time.Sleep(time.Millisecond) + if v, ok = p.inFlight.Load(id); ok { + break + } + } + } + cancel, _ := v.(context.CancelCauseFunc) + if !ok || cancel == nil { + return ErrTaskNotRunning + } + cancel(&userCancelCause{reason: reason}) + return nil +} diff --git a/worker/mock_store_test.go b/worker/mock_store_test.go new file mode 100644 index 0000000..29ab537 --- /dev/null +++ b/worker/mock_store_test.go @@ -0,0 +1,58 @@ +package worker + +// 本文件承载 worker 测试的 TaskStore 桩（从 pool_test.go 拆出）： +// 架构门禁要求测试文件 ≤ 800 行，桩会随接口扩展增长， +// 与生产侧「同包正交拆分」同一手法。 + +import ( + "agent-runtime/store" + "context" +) + +type mockTaskStore struct { + updateErr error + getErr error + completeErr error + + completedID string + completedResult string + completedErr error + completeCalled bool + + // usage 记录 RecordUsage 的入参，供 ADR-0021 的成本计量断言使用。 + usageID string + usagePrompt int + usageCompletion int + usageCalled bool + usageErr error +} + +func (m *mockTaskStore) Create(prompt string) string { return "mock-id" } +func (m *mockTaskStore) CreateInSession(prompt, sessionID string) string { return m.Create(prompt) } +func (m *mockTaskStore) Update(id, status string) error { return m.updateErr } +func (m *mockTaskStore) Get(id string) (store.Task, error) { + if m.getErr != nil { + return store.Task{}, m.getErr + } + return store.Task{ID: id, Prompt: "mock prompt"}, nil +} +func (m *mockTaskStore) Complete(id, result string, err error) error { + m.completeCalled = true + m.completedID = id + m.completedResult = result + m.completedErr = err + return m.completeErr +} +func (m *mockTaskStore) Delete(id string) error { return nil } +func (m *mockTaskStore) RecoverTasks(ctx context.Context) ([]string, int, error) { return nil, 0, nil } +func (m *mockTaskStore) ClaimNext(ctx context.Context) (store.Task, error) { + return store.Task{}, store.ErrNoTaskAvailable +} +func (m *mockTaskStore) PendingCount(ctx context.Context) (int64, error) { return 0, nil } + +// RecordUsage 记录计量入参；usageErr 非 nil 时用于验证「计量失败不改任务终态」。 +func (m *mockTaskStore) RecordUsage(id string, promptTokens, completionTokens int) error { + m.usageID, m.usagePrompt, m.usageCompletion = id, promptTokens, completionTokens + m.usageCalled = true + return m.usageErr +} diff --git a/worker/pool.go b/worker/pool.go index 9ebadeb..8c19659 100644 --- a/worker/pool.go +++ b/worker/pool.go @@ -54,6 +54,10 @@ type Pool struct { inFlight sync.Map sessionStore store.SessionStore maxContextTokens int + // sessionHistoryLimit 是组装会话历史时拉取的最大消息条数（ADR-0021 决策 3）。 + // 这是「模型窗口」参数而非代码常量：128k 窗口配 100 条是浪费， + // 16k 窗口配 100 条则会被 CompactMessages 反复折叠、白跑压缩。 + sessionHistoryLimit int } // SetMaxQueueSize 设置最大允许排队积压的任务数，供测试和自定义配置使用。 @@ -64,21 +68,22 @@ func NewPool(size int, taskTimeout time.Duration, s store.TaskStore, client agen stopWorkerCtx, stopWorkerCancel := context.WithCancel(context.Background()) return &Pool{ - size: size, - maxQueueSize: 10000, - taskTimeout: taskTimeout, - tasks: s, - wakeUp: make(chan struct{}, 1), - ctx: ctx, - cancel: cancel, - stopWorkerCtx: stopWorkerCtx, - stopWorkerCancel: stopWorkerCancel, - llm: client, - tools: tool.DefaultRegistry(), - traces: make(map[string]*agent.TraceBuffer), - maxTraces: DefaultMaxTraces, - traceOrder: make([]string, 0, 64), - maxContextTokens: 16000, + size: size, + maxQueueSize: 10000, + taskTimeout: taskTimeout, + tasks: s, + wakeUp: make(chan struct{}, 1), + ctx: ctx, + cancel: cancel, + stopWorkerCtx: stopWorkerCtx, + stopWorkerCancel: stopWorkerCancel, + llm: client, + tools: tool.DefaultRegistry(), + traces: make(map[string]*agent.TraceBuffer), + maxTraces: DefaultMaxTraces, + traceOrder: make([]string, 0, 64), + maxContextTokens: 16000, + sessionHistoryLimit: 100, } } @@ -229,9 +234,12 @@ func (p *Pool) process(id string) { if err != nil { slog.Warn("worker: assemble session messages failed", "task_id", id, "session_id", task.SessionID, "error", err) } + var usage usageAccumulator res, err := agent.RunWithOptions(taskCtx, p.llm, task.Prompt, p.tools, tb.Append, agent.RunOptions{ - HistoryMessages: history, OnDelta: tb.AppendDelta, + HistoryMessages: history, OnDelta: tb.AppendDelta, OnUsage: usage.add, }) + // 用量无论成败都要落：失败的步进同样计费（ADR-0021）。 + p.recordUsage(id, &usage) if err != nil { // 错误全文只进日志：它携带上游响应体，可能含配额、账单等内部信息。 // task.Error 会经 GET /tasks/{id} 原样返回给调用方，因此只存粗粒度分类。 @@ -339,36 +347,6 @@ func (p *Pool) SubscribeEvents(id string) (snapshot []agent.TraceStep, events <- return snapshot, events, cancel, true } -// userCancelCause 是用户主动取消的因果标记，用于与超时/停机取消精确区分。 -type userCancelCause struct { - reason string -} - -func (c *userCancelCause) Error() string { return "user canceled: " + c.reason } - -// Cancel 主动抢占指定在飞任务：通过 context 因果链物理中断正在进行的上游调用。 -// 任务不在飞（未开始、已终态或不存在）时返回 ErrTaskNotRunning，由 API 层映射为 409。 -func (p *Pool) Cancel(id, reason string) error { - v, ok := p.inFlight.Load(id) - if !ok { - for i := 0; i < 20; i++ { - if tk, err := p.tasks.Get(id); err != nil || tk.Status != store.StatusRunning { - break - } - time.Sleep(time.Millisecond) - if v, ok = p.inFlight.Load(id); ok { - break - } - } - } - cancel, _ := v.(context.CancelCauseFunc) - if !ok || cancel == nil { - return ErrTaskNotRunning - } - cancel(&userCancelCause{reason: reason}) - return nil -} - // EnqueueWait 阻塞式入队，用于崩溃恢复时重灌 pending 积压： // 循环等待直到待处理任务数未超出容量限制，既不丢任务也不无限挂死。 func (p *Pool) EnqueueWait(ctx context.Context, id string) error { diff --git a/worker/pool_test.go b/worker/pool_test.go index 4fdde7b..3eec6af 100644 --- a/worker/pool_test.go +++ b/worker/pool_test.go @@ -198,40 +198,6 @@ func TestEnqueue_Full(t *testing.T) { } } -type mockTaskStore struct { - updateErr error - getErr error - completeErr error - - completedID string - completedResult string - completedErr error - completeCalled bool -} - -func (m *mockTaskStore) Create(prompt string) string { return "mock-id" } -func (m *mockTaskStore) CreateInSession(prompt, sessionID string) string { return m.Create(prompt) } -func (m *mockTaskStore) Update(id, status string) error { return m.updateErr } -func (m *mockTaskStore) Get(id string) (store.Task, error) { - if m.getErr != nil { - return store.Task{}, m.getErr - } - return store.Task{ID: id, Prompt: "mock prompt"}, nil -} -func (m *mockTaskStore) Complete(id, result string, err error) error { - m.completeCalled = true - m.completedID = id - m.completedResult = result - m.completedErr = err - return m.completeErr -} -func (m *mockTaskStore) Delete(id string) error { return nil } -func (m *mockTaskStore) RecoverTasks(ctx context.Context) ([]string, int, error) { return nil, 0, nil } -func (m *mockTaskStore) ClaimNext(ctx context.Context) (store.Task, error) { - return store.Task{}, store.ErrNoTaskAvailable -} -func (m *mockTaskStore) PendingCount(ctx context.Context) (int64, error) { return 0, nil } - func TestProcess_UpdateErrorSetsFailed(t *testing.T) { mock := &mockTaskStore{ updateErr: errors.New("db update error"), diff --git a/worker/session.go b/worker/session.go index 214386f..ddccc6f 100644 --- a/worker/session.go +++ b/worker/session.go @@ -18,6 +18,14 @@ func (p *Pool) SetMaxContextTokens(tokens int) { p.maxContextTokens = tokens } +// SetSessionHistoryLimit 设置组装会话历史时拉取的最大消息条数（ADR-0021 决策 3）。 +// 必须与 SetMaxContextTokens 一样在 Start 之前调用：worker 消费任务时读取该值。 +func (p *Pool) SetSessionHistoryLimit(n int) { + if n > 0 { + p.sessionHistoryLimit = n + } +} + // assembleSessionMessages 拉取指定 sessionID 的历史消息，与系统提示词和当前 prompt 组装， // 并通过 agent.CompactMessages 执行滑动窗口压缩，返回符合大模型上下文预算的消息列表。 func (p *Pool) assembleSessionMessages(ctx context.Context, sessionID, prompt, sysPrompt string) ([]llm.Message, error) { @@ -25,7 +33,11 @@ func (p *Pool) assembleSessionMessages(ctx context.Context, sessionID, prompt, s return nil, nil } - msgs, err := p.sessionStore.GetSessionMessages(ctx, sessionID, 100) + limit := p.sessionHistoryLimit + if limit <= 0 { + limit = 100 + } + msgs, err := p.sessionStore.GetSessionMessages(ctx, sessionID, limit) if err != nil { return nil, err } diff --git a/worker/usage.go b/worker/usage.go new file mode 100644 index 0000000..46f7526 --- /dev/null +++ b/worker/usage.go @@ -0,0 +1,36 @@ +package worker + +import ( + "agent-runtime/llm" + "agent-runtime/metrics" + "log/slog" +) + +// usageAccumulator 累计一次任务执行期间所有 LLM 步进的 Token 消耗（ADR-0021）。 +// +// 只在 agent.RunWithOptions 的同一 goroutine 内被 OnUsage 回调写入，因此不需要加锁。 +// 跨 goroutine 复用它就会变成数据竞争——一次任务一个实例，禁止提升为 Pool 字段。 +type usageAccumulator struct { + promptTokens int + completionTokens int +} + +// add 满足 agent.RunOptions.OnUsage 的签名。 +func (u *usageAccumulator) add(us llm.Usage) { + u.promptTokens += us.PromptTokens + u.completionTokens += us.CompletionTokens +} + +// recordUsage 把本次任务的用量写进进程指标与任务记录。 +// +// 成功与失败路径都要调用：失败的调用照样计费，漏记失败任务的用量会让账单对不上。 +// 落库失败只记日志、不改任务终态——计量是旁路信号，不该让业务结果失败。 +func (p *Pool) recordUsage(id string, u *usageAccumulator) { + if u == nil || (u.promptTokens == 0 && u.completionTokens == 0) { + return + } + metrics.AddTokens(u.promptTokens, u.completionTokens) + if err := p.tasks.RecordUsage(id, u.promptTokens, u.completionTokens); err != nil { + slog.Warn("worker: record token usage failed", "task_id", id, "error", err) + } +} diff --git a/worker/usage_test.go b/worker/usage_test.go new file mode 100644 index 0000000..d18cfe7 --- /dev/null +++ b/worker/usage_test.go @@ -0,0 +1,126 @@ +package worker + +import ( + "agent-runtime/llm" + "context" + "errors" + "testing" + "time" +) + +// usageChatter 返回带 usage 的终态回复（无工具调用 → 一步结束）。 +type usageChatter struct { + usage llm.Usage +} + +func (u usageChatter) ChatWithTools(ctx context.Context, msgs []llm.Message, tools []llm.Tool) (*llm.AssistantTurn, error) { + return &llm.AssistantTurn{Content: "ok", Usage: u.usage}, nil +} + +// failAfterToolChatter 第一次返回「调用工具 + usage」，第二次返回错误。 +// 用于验证「失败的调用照样计费」：用量在错误发生前已经产生。 +type failAfterToolChatter struct { + usage llm.Usage + calls int +} + +func (f *failAfterToolChatter) ChatWithTools(ctx context.Context, msgs []llm.Message, tools []llm.Tool) (*llm.AssistantTurn, error) { + f.calls++ + if f.calls > 1 { + return nil, errors.New("upstream boom") + } + return &llm.AssistantTurn{ + Usage: f.usage, + ToolCalls: []llm.ToolCall{{ + ID: "call_usage", + Type: "function", + Function: llm.ToolCallFn{Name: "calculate", Arguments: `{"expression":"1+1"}`}, + }}, + }, nil +} + +// TestProcess_RecordsUsageOnSuccess 是 ADR-0021 决策 2 的承重墙之一： +// 成功任务必须把 token 用量写进任务记录。 +func TestProcess_RecordsUsageOnSuccess(t *testing.T) { + mock := &mockTaskStore{} + p := NewPool(1, time.Second, mock, usageChatter{usage: llm.Usage{PromptTokens: 11, CompletionTokens: 5, TotalTokens: 16}}) + p.process("task-usage") + + if !mock.usageCalled { + t.Fatal("expected RecordUsage to be called, but it was not") + } + if mock.usageID != "task-usage" { + t.Errorf("usage id = %q, want task-usage", mock.usageID) + } + if mock.usagePrompt != 11 || mock.usageCompletion != 5 { + t.Errorf("usage = (%d,%d), want (11,5)", mock.usagePrompt, mock.usageCompletion) + } + if mock.completedErr != nil { + t.Errorf("task should be done, got error %v", mock.completedErr) + } +} + +// TestProcess_RecordsUsageOnFailure 是成本口径的关键：失败的步进同样计费， +// 漏记失败任务的用量会让账单对不上。 +func TestProcess_RecordsUsageOnFailure(t *testing.T) { + mock := &mockTaskStore{} + chatter := &failAfterToolChatter{usage: llm.Usage{PromptTokens: 7, CompletionTokens: 3, TotalTokens: 10}} + p := NewPool(1, time.Second, mock, chatter) + p.process("task-fail-usage") + + if !mock.usageCalled { + t.Fatal("expected RecordUsage to be called on the failure path, but it was not") + } + if mock.usagePrompt != 7 || mock.usageCompletion != 3 { + t.Errorf("usage = (%d,%d), want (7,3)", mock.usagePrompt, mock.usageCompletion) + } + if mock.completedErr == nil { + t.Error("expected task to be marked failed") + } +} + +// TestProcess_UsageRecordErrorKeepsTerminalState 守住「计量是旁路信号」： +// RecordUsage 失败只记日志，不得把成功的任务改成失败。 +func TestProcess_UsageRecordErrorKeepsTerminalState(t *testing.T) { + mock := &mockTaskStore{usageErr: errors.New("disk full")} + p := NewPool(1, time.Second, mock, usageChatter{usage: llm.Usage{PromptTokens: 2, CompletionTokens: 1}}) + p.process("task-usage-err") + + if !mock.completeCalled { + t.Fatal("expected Complete to be called") + } + if mock.completedErr != nil { + t.Errorf("usage record failure must not fail the task, got %v", mock.completedErr) + } +} + +// TestProcess_SkipsUsageWhenUpstreamReportsNone 验证零用量不产生无意义的写操作： +// 上游未返回 usage（如部分兼容实现）时不应把 0 写成「已计量」。 +func TestProcess_SkipsUsageWhenUpstreamReportsNone(t *testing.T) { + mock := &mockTaskStore{} + p := NewPool(1, time.Second, mock, usageChatter{}) + p.process("task-no-usage") + + if mock.usageCalled { + t.Error("expected RecordUsage to be skipped when usage is zero") + } + if mock.completedErr != nil { + t.Errorf("task should be done, got error %v", mock.completedErr) + } +} + +// TestSetSessionHistoryLimit 验证窗口配置只接受正数（0/负数回落到默认 100）。 +func TestSetSessionHistoryLimit(t *testing.T) { + p := NewPool(1, time.Second, &mockTaskStore{}, usageChatter{}) + if p.sessionHistoryLimit != 100 { + t.Errorf("default = %d, want 100", p.sessionHistoryLimit) + } + p.SetSessionHistoryLimit(500) + if p.sessionHistoryLimit != 500 { + t.Errorf("after set = %d, want 500", p.sessionHistoryLimit) + } + p.SetSessionHistoryLimit(0) + if p.sessionHistoryLimit != 500 { + t.Errorf("zero must be ignored, got %d", p.sessionHistoryLimit) + } +} From 28c55b79cbbaf1a79e5e1b189ff4354c3bf06129 Mon Sep 17 00:00:00 2001 From: Haoyu Xi Date: Tue, 27 Oct 2026 21:22:58 +0800 Subject: [PATCH 2/2] docs(p10): record ADR-0021 and sync architecture, API, and memory anchors ADR-0021 carries the rejected-alternatives table (no retryablehttp, no infinite retry, no usage-in-trace-only, no token fields inside Complete, no pprof on the main mux, no retry after a stream break) so the next session does not re-litigate them. README: 14 config keys, agent_llm_* metrics, per-task token fields in the GET /tasks contract, and the load-probe evidence. ARCHITECTURE: P1-P10 scope, production capabilities, diagnostics in the observability subsection, ADR-0021 in the index. .agent-context: refreshed baseline (360 func Test / 545 RUN / 99 files / 20,851 lines), new retry and metering red lines. WORKMEMORY: PROJECT_OVERVIEW P10 entry, work.log event, notes file, INDEX row. --- .agent-context.md | 21 +++-- ARCHITECTURE.md | 9 +- README.md | 26 ++++-- WORKMEMORY/INDEX.md | 1 + WORKMEMORY/PROJECT_OVERVIEW.md | 9 +- ...2026-10-27-1900-p10-upstream-retry-cost.md | 79 ++++++++++++++++ WORKMEMORY/work.log | 9 ++ config.yaml.example | 2 + ...ream-retry-and-token-cost-observability.md | 92 +++++++++++++++++++ 9 files changed, 225 insertions(+), 23 deletions(-) create mode 100644 WORKMEMORY/notes/NOTES-2026-10-27-1900-p10-upstream-retry-cost.md create mode 100644 docs/adr/0021-upstream-retry-and-token-cost-observability.md diff --git a/.agent-context.md b/.agent-context.md index ba3c31d..1f84eee 100644 --- a/.agent-context.md +++ b/.agent-context.md @@ -6,10 +6,10 @@ opt-in `modernc.org/sqlite`（纯 Go 无 CGO）。 ### 1. 🛠️ Ground Truth Verification Baseline - **Primary Test Command**: `GOWORK=off go test -race ./... -count=1`（必须 GOWORK=off：仓库可能处于某 worktree，GOWORK 会让包路径解析漂移） -- **Baseline Test Result**: [PASS] 321 func Test / 498 RUN（含子测试），12 包全绿，0 data race（re-verified on 2026-10-27；代码规模 83 个 Go 文件 = 39 生产 / 44 测试，共 19,074 行） +- **Baseline Test Result**: [PASS] 360 func Test / 545 RUN（含子测试），12 包全绿，0 data race（re-verified on 2026-10-27；代码规模 99 个 Go 文件 = 44 生产 / 55 测试，共 20,851 行） - **Build / Lint Command**: `GOWORK=off go build ./...` + `GOWORK=off go vet ./...`（干净无警告） -- **Architecture Gates**: `GOWORK=off go test -v -run TestArchitecture . -count=1`（全仓 AST DAG 白名单 + 生产 ≤ 400 行 / 测试 ≤ 800 行硬门禁，83 个 Go 文件全量受控） -- **Benchmark**: `go test -bench=. -benchmem ./worker/` +- **Architecture Gates**: `GOWORK=off go test -v -run TestArchitecture . -count=1`（全仓 AST DAG 白名单 + 生产 ≤ 400 行 / 测试 ≤ 800 行硬门禁，99 个 Go 文件全量受控） +- **Benchmark**: `go test -bench=. -benchmem ./worker/ ./api/` - 若先跑失败，先排查 Pre-Existing Broken（对照本基线），别归因到后续改动。 ### 2. 🗺️ Core Symbol & Architecture Map @@ -24,10 +24,11 @@ opt-in `modernc.org/sqlite`（纯 Go 无 CGO）。 - `eval/`：检索评测 harness（ADR-0018）；同包拆分 `run.go`（受控参数锁定与预检）、`collector.go`（全文收集与辅助）、`grader.go`（程序化 Span 覆盖率）、`aggregate.go`（分布指标）、`report.go`（Markdown+JSON 报表，支持 `llm-usage` API 真实 Token 统计）。 - `prompts/`：版本化系统提示词（`go:embed *.md`），提供 `Get(version)`；收录 `default-v1`、`eval-template-v1`（空对照）、`eval-template-v2`（Grounded Citation 铁律）。 - `store/`：`TaskStore` 与 `SessionStore` 接口；`store.go` 内存任务存储 + `session.go` 内存会话存储；`sql.go` SQLite WAL 任务持久化与原子 CAS 抢占 + `sql_session.go` SQLite 会话与历史消息持久化（复合索引优化）；`RecoverTasks` 启动断电补偿。 - - `worker/`：`Pool`（Push-Notify 事务性抢占并发池，`tasks.ClaimNext` 原子 CAS 抢占，零 channel 阻塞，零 retrieval 依赖，通过 `tool.TerminalClassifier` 脱敏终态）；`session.go`（多轮历史组装拉取 + 压缩 + 成功后原子写回，杜绝 `pool.go` 膨胀）；两阶段优雅停机 + 队列快速失败 503 + context 抢占式取消树。 - - `llm/`：`Client`（OpenAI 兼容客户端，幂等 `joinEndpoint` 防 `/v1/v1` 404，支持本地免密 Ollama，透传 `Temperature`/`Seed`，解析并回传 `Usage`）；`stream.go` 提供 `ChatWithToolsStream` 原生 SSE 流式解析与工具调用分块组装。 - - `metrics/`：包级原子计数器，提供 Prometheus 文本格式导出。 - - `config/`：`Config` 12 项配置扁平解析（增补 `worker_concurrency` 默认 3 与 `max_context_tokens` 默认 16000）。 + - `worker/`：`Pool`（Push-Notify 事务性抢占并发池，`tasks.ClaimNext` 原子 CAS 抢占，零 channel 阻塞，零 retrieval 依赖，通过 `tool.TerminalClassifier` 脱敏终态）；`cancel.go`（抢占取消句柄解耦）；`usage.go`（跨步 Token 计量累加与持久化落盘）；`session.go`（多轮历史组装拉取 + 压缩 + 成功后原子写回，杜绝 `pool.go` 膨胀）；两阶段优雅停机 + 队列快速失败 503 + context 抢占式取消树。 + - `llm/`：`Client`（OpenAI 兼容客户端，幂等 `joinEndpoint` 防 `/v1/v1` 404，支持本地免密 Ollama，透传 `Temperature`/`Seed`，解析并回传 `Usage`）；`retry.go` 驱动 429/5xx 指数退避与 `Retry-After` 尊重，`OnRetry` 回调注入指标；`stream.go` 提供 `ChatWithToolsStream` 原生 SSE 流式解析与工具调用分块组装（流式首包前可重试、产出后禁重试）。 + - `metrics/`：包级原子计数器，提供 Prometheus 文本格式导出（覆盖 `agent_tasks_*` 与 `agent_llm_*`）。 + - `config/`：`Config` 14 项配置扁平解析（增补 `worker_concurrency`、`max_context_tokens`、`session_history_limit`、`pprof_addr`）。 + - 根级 `pprof.go`：独立端口与 mux 提供 `/debug/pprof/*`，默认关闭。 - **Existing Wheels（复用层）**: `agent.Chatter` / `llm.Embedder` 接口使测试可注入假实现；`httptest` 集成测试端到端跑 `POST /run → 轮询 → done`；`go:embed` 单页 Demo；`TraceBuffer` 同时服务轮询快照与 SSE 订阅。 ### 3. 🚨 Local Engineering Red Lines & Gotchas @@ -41,7 +42,9 @@ opt-in `modernc.org/sqlite`（纯 Go 无 CGO）。 - **SSE 必须清写截止时间**：`http.Server.WriteTimeout` 同样作用于流式响应，30s 被掐 → handler 内 `http.NewResponseController(w).SetWriteDeadline(time.Time{})`（不支持时忽略 ErrNotSupported）；`TraceBuffer.Append` 对订阅通道 select+default 非阻塞投递，慢消费者只丢自己增量，终态用快照对账补齐（`sent < len(snapshot)` 续发）。 - **Trace 有界**：`DefaultMaxTraces = 1000` FIFO 驱逐防 OOM；detail 512B 截断用 `unicode/utf8.RuneStart` 回退到字符边界，杜绝 `\ufffd` 乱码。 - **SQLite 写吞吐**：`id INTEGER PRIMARY KEY AUTOINCREMENT` + `res.LastInsertId()`，禁止 `SELECT MAX(CAST(id AS INTEGER))`（全表扫描裹在排他写锁里，曾劣化到 <480 QPS）；`RecoverTasks` 全程事务。 -- **config.yaml 只接受扁平子集**：合法 key 包含 `deepseek_api_key`、`api_auth_key`、`model`、`db_path`、`task_timeout_seconds`、`max_queue_size`、`worker_concurrency`、`max_context_tokens`、`embed_base_url`、`embed_api_key`、`embed_model`、`embed_dim`（12 项）；嵌套/未知 key 一律启动失败不静默。文件 gitignored（含密钥），示例是 `config.yaml.example`。同名环境变量优先级更高。服务模式缺 `DEEPSEEK_API_KEY` 以退出码 1 失败（`-index` 与 `-h` 短路不校验）。 +- **config.yaml 只接受扁平子集**：合法 key 包含 `deepseek_api_key`、`api_auth_key`、`model`、`db_path`、`task_timeout_seconds`、`max_queue_size`、`worker_concurrency`、`max_context_tokens`、`session_history_limit`、`pprof_addr`、`embed_base_url`、`embed_api_key`、`embed_model`、`embed_dim`（14 项）；嵌套/未知 key 一律启动失败不静默。文件 gitignored（含密钥），示例是 `config.yaml.example`。同名环境变量优先级更高。服务模式缺 `DEEPSEEK_API_KEY` 以退出码 1 失败（`-index` 与 `-h` 短路不校验）。 +- **上游重试与流式不变式**：`llm.Client` 仅对 429 与 5xx 状态码退避重试，严禁重试 400/401/403/404/422；`Retry-After` 头受 `MaxDelay` 封顶；流式仅在首包前可重试，产出 delta 后断流坚决不重试以防打字机重复打印。 +- **计量与状态机解耦**：`store.TaskStore.RecordUsage` 与 `Complete` 关注点分离；失败任务亦需计量（失败调用同样计费）；计量落盘失败仅打日志，绝不把成功任务改写为失败。 - **会话存储与上下文防爆纪律**：多轮会话接入 `SessionStore`，创建任务通过 `CreateInSession` 记录会话绑定；`worker.process` 严禁在 `worker/pool.go` 直接堆砌多轮会话代码（硬限 400 行），必须在 `worker/session.go` 中通过 `assembleSessionMessages` 驱动 `agent.CompactMessages` 进行滑动窗口折半压缩；`recordSessionTurn` 必须严格检验并包装 `AppendSessionMessage` 错误，禁止使用 `_ =` 静默吞错。 - **并发取消在飞时序（TOCTOU 防护）**：在 `worker/pool.go:process` 中，必须先登记 `inFlight.Store(id, cancelCause)` 再更新数据库 `tasks.Update(id, StatusRunning)`。严禁先置 running 再存 handle，否则外部轮询方察觉 running 后立即调 cancel 会因找不到 handle 而报 `task is not running`。 - **请求/响应体积护栏**：`/run` body `MaxBytesReader` 32KB；LLM 响应 `io.LimitReader` 1MB。 @@ -58,5 +61,5 @@ opt-in `modernc.org/sqlite`（纯 Go 无 CGO）。 - Global Standards: `{vault_root}/AGENTS.md`（按 `$WIKI_PATH` → `$OBSIDIAN_VAULT_PATH` → 向上查找 `AGENTS.md` 解析，契约见 `dfs/references/vault-root-resolution.md`） - Vault Search: 调用 `search_vault` MCP 或 `python {vault_root}/scripts/search-vault.py ""` - 沉淀新坑：写草稿至 `{vault_root}/99-Inbox/YYYY-MM-DD-{topic}.md` -- 项目文档：`docs/GO-STANDARDS.md`、`docs/GO-CHEATSHEET.md`、`docs/adr/0001..0020`（ADR-0019 事务性队列 + 打字机流式、ADR-0020 会话状态机 + 上下文压缩均已实施）；`ARCHITECTURE.md`（根级架构总览：What Agent / What Architecture / What Borrowed / 设计权衡） +- 项目文档：`docs/GO-STANDARDS.md`、`docs/GO-CHEATSHEET.md`、`docs/adr/0001..0021`（ADR-0019 事务性队列 + 打字机流式、ADR-0020 会话状态机 + 上下文压缩、ADR-0021 上游退避重试 + Token 成本可观测 + 诊断端口均已实施）；`ARCHITECTURE.md`（根级架构总览：What Agent / What Architecture / What Borrowed / 设计权衡） - 工作记忆：`WORKMEMORY/`（INDEX → PROJECT_OVERVIEW → work.log 尾部）；决策/踩坑/交接 MUST ≤4KB 事件追加 work.log diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index cf44312..94ae505 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -2,7 +2,7 @@ > 架构总览文档。回答三个核心问题：**实现了什么 Agent**、**基于什么架构**、**借鉴了什么**。 > 维护规则：随架构演进更新；重大变更须同步对应 ADR 与 `README.md`。 -> 适用范围：当前主线（`master`，P1–P8 全部交付）。 +> 适用范围：当前主线（`master`，P1–P10 全部交付）。 --- @@ -25,7 +25,7 @@ | **内置工具** | `calculate`（自研递归下降 AST 算术求值，**不用 eval**，防代码注入）、`current_time`（IANA 时区）；P5 起增加 `search_docs`（混合检索即工具） | | **执行形态** | 被 Worker Pool 并发调度的**后端服务**智能体，而非交互式 CLI / 本地 Agent | | **对外契约** | 任务状态机 `pending → running → done/failed`；`POST /run` 提交、`GET /tasks/{id}` 轮询、`GET /tasks/{id}/events` SSE 流式 | -| **生产级能力** | 用户主动取消（context 物理中断上游）、掉电崩溃恢复、SSE 原生 token 级打字机、多轮会话状态机、滑动窗口上下文压缩 | +| **生产级能力** | 用户主动取消（context 物理中断上游）、掉电崩溃恢复、SSE 原生 token 级打字机、多轮会话状态机、滑动窗口上下文压缩、上游限流指数退避与 Retry-After 尊重、Token 成本计量持久化、独立诊断端口（pprof） | **最小可信定义**：多步工具调用循环即 agent 的实质（ADR-0008 双镜论证）——区别于"LLM 聊天的异步封装"，使 "agent" 名副其实。 @@ -97,8 +97,7 @@ `POST /run` 可绑定 `session_id`，打破单轮任务孤岛。 - **纯算法滑动窗口压缩**：`agent.CompactMessages` 无堆逃逸 Token 估算，首尾强锚定（System 提示 + 最新一轮）， 中间轮次按 User 边界折半截断并插入系统折叠注记，**零 LLM 调用开销**根除 Context Window 爆炸。 -- **可观测性**：`/metrics` 以 Prometheus 文本格式暴露 `agent_tasks_*`（`sync/atomic` 实现，零依赖）； - `/healthz` 探活；指标仅暴露聚合计数，不含任务内容或上游响应体。 +- **可观测性与诊断**：`/metrics` 以 Prometheus 文本格式暴露 `agent_tasks_*` 任务生命周期指标与 `agent_llm_*` 上游重试与 Token 成本指标（`sync/atomic` 实现，零依赖）；`tasks` 表持久化单任务 `prompt_tokens`/`completion_tokens`；`/healthz` 存活探针；独立端口/mux 提供 `/debug/pprof/*` 性能分析；指标仅暴露聚合计数，不含任务内容或上游响应体。 --- @@ -132,7 +131,7 @@ - 决策记录：`docs/adr/0001`（项目定位）、`0002`（Worker Pool）、`0008`（Agent 工具循环）、 `0013`（分层解耦，Pi 哲学）、`0015`（崩溃恢复）、`0016`（SSE 流式）、`0017/0018`（检索与评测）、 - `0019`（事务队列与流式推断）、`0020`（会话状态机与压缩）。 + `0019`（事务队列与流式推断）、`0020`（会话状态机与压缩）、`0021`（上游限流退避、Token 成本可观测与诊断端口）。 - 实现计划：`docs/implementation_plan_p1.md` … `p8.md`。 - 评测实证：`docs/eval/2026-10-25-hybrid-r3.md` 等。 - 运行与 API 契约详见 `README.md`。 diff --git a/README.md b/README.md index c7a36f1..2ee4f66 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,7 @@ Copy-Item config.yaml.example config.yaml # 然后填真实值 go run . ``` -`config.yaml` 已被 `.gitignore` 忽略。只支持扁平的 `key: value`，解析到嵌套或未知 key 会直接启动失败，不静默忽略。已支持 12 项配置：`deepseek_api_key`、`api_auth_key`、`model`、`db_path`、`task_timeout_seconds`、`embed_base_url`、`embed_api_key`、`embed_model`、`embed_dim`、`max_queue_size`、`worker_concurrency`、`max_context_tokens`： +`config.yaml` 已被 `.gitignore` 忽略。只支持扁平的 `key: value`，解析到嵌套或未知 key 会直接启动失败，不静默忽略。已支持 14 项配置：`deepseek_api_key`、`api_auth_key`、`model`、`db_path`、`task_timeout_seconds`、`embed_base_url`、`embed_api_key`、`embed_model`、`embed_dim`、`max_queue_size`、`worker_concurrency`、`max_context_tokens`、`session_history_limit`、`pprof_addr`： ```yaml deepseek_api_key: "sk-你的key" @@ -78,6 +78,8 @@ task_timeout_seconds: 30 # 可选：单次任务 LLM 调用超时（正整数 max_queue_size: 10000 # 可选：最大允许排队积压的任务数（超出返回 503），默认 10000 worker_concurrency: 3 # 可选：后台并发 Worker 数量（正整数，默认 3），环境变量 WORKER_CONCURRENCY 优先级更高 max_context_tokens: 16000 # 可选：会话多轮对话上下文 Token 预算上限（正整数，默认 16000），环境变量 MAX_CONTEXT_TOKENS 优先级更高 +session_history_limit: 100 # 可选：组装会话历史时拉取的最大消息条数（正整数，默认 100），环境变量 SESSION_HISTORY_LIMIT 优先级更高 +pprof_addr: "" # 可选：非空则在独立端口暴露 /debug/pprof（如 127.0.0.1:6060）；空 = 关闭，环境变量 PPROF_ADDR 优先级更高 # embed_* 四项为检索功能的向量通道配置（ADR-0017），全部留空 = 纯词法检索（lexical） # 本地 Ollama 等免密向量引擎无需鉴权时，embed_api_key 可以留空 embed_base_url: "http://localhost:11434/v1" # OpenAI 兼容 embeddings 端点 @@ -137,13 +139,13 @@ Invoke-RestMethod -Uri http://localhost:8080/metrics | Method | Path | 说明 | Request Body | Response | |--------|------|------|-------------|----------| | POST | `/run` | 提交任务（支持可选绑定会话） | `{"prompt":"...","session_id":"..."}` | `202 {"task_id":"1","session_id":"..."}` | -| GET | `/tasks/{id}` | 查询状态（轮询） | - | `{"id":"1","status":"done","prompt":"...","result":"...","error":""}` | +| GET | `/tasks/{id}` | 查询状态（轮询） | - | `{"id":"1","status":"done","prompt":"...","result":"...","error":"","prompt_tokens":11,"completion_tokens":5}` | | POST | `/tasks/{id}/cancel` | 主动取消在飞任务 | `{"reason":"user_aborted"}`（可选） | `202 {"id":"1","status":"canceling"}`；任务不在飞返回 `409` | | GET | `/tasks/{id}/events` | SSE 实时流式推送 | - | `text/event-stream`：`event: delta` (打字机增量) / `event: step` / `event: done` | | GET | `/sessions/{id}` | 获取会话元数据 | - | `200 {"id":"sess_1","created_at":"...","updated_at":"..."}` | | GET | `/sessions/{id}/messages` | 获取会话多轮历史 | `?limit=N`（可选，默认 100） | `200 [{"id":1,"role":"user","content":"..."},...]` | | GET | `/healthz` | 存活探针 | - | `200 ok`（纯文本） | -| GET | `/metrics` | Prometheus 指标 | - | `agent_tasks_*` 文本格式 | +| GET | `/metrics` | Prometheus 指标 | - | `agent_tasks_*` / `agent_llm_*` 文本格式 | | GET | `/` | 单页 Demo | - | `text/html` 页面 | Status 取值：`pending` → `running` → `done` / `failed` @@ -162,6 +164,9 @@ Status 取值：`pending` → `running` → `done` / `failed` | `agent_tasks_done` | counter | 累计成功完成数 | | `agent_tasks_failed` | counter | 累计失败数 | | `agent_tasks_recovered` | counter | 崩溃恢复重灌的 pending 任务数（回到 pending 排队，不计入 running） | +| `agent_llm_retries_total` | counter | 累计上游 LLM 限流/5xx 重试次数（衡量上游稳定性） | +| `agent_llm_prompt_tokens_total` | counter | 累计消耗的 Prompt Token 数 | +| `agent_llm_completion_tokens_total` | counter | 累计消耗的 Completion Token 数 | 不变式：`submitted - done - failed = running + pending`。`running` 恒 `<= worker_concurrency`； 若 `running` 接近并发上限而 `pending` 持续增长，说明消费者不足而非上游卡死。 @@ -183,16 +188,17 @@ Token 框填一次就会存进浏览器 `localStorage`，下次打开自动回 ## Testing ```powershell -GOWORK=off go test -race ./... -count=1 # 321 个测试用例 / 498 个测试分支全绿，零 race -go test -bench=. -benchmem ./worker/ # 基准测试 +GOWORK=off go test -race ./... -count=1 # 360 个测试用例 / 545 个测试分支全绿，零 race +go test -bench=. -benchmem ./worker/ ./api/ # 基准测试 ``` - **单测体系**：涵盖 12 个分层包，覆盖 store 状态机与 SessionStore 会话流、worker 并发与取消树、agent 认知循环与滑动窗口压缩、tool 注册表与 AST 算术求解器、retrieval 词法/向量/RRF 融合召回、eval 自动化判题器。 - **静态架构与体积门禁**（`architecture_test.go`）： - **AST DAG 白名单**：用 `go/parser` 扫描全包 import 关系，严格锁死单向分层，杜绝跨层反向依赖（如 worker 严禁 import retrieval，api 严禁 import worker）； - **代码体积预算（Line Budget）**：强制生产 `.go` 文件 ≤ 400 行、测试文件 ≤ 800 行，多版本迭代防腐化。 -- **集成测试**：`api/integration_test.go` 用 `httptest` 端到端验证 `POST /run → 轮询 → done`、SSE 流式事件与 Session 多轮交互。 +- **集成测试与压测实证**：`api/integration_test.go` 用 `httptest` 端到端验证 `POST /run → 轮询 → done`、SSE 流式事件与 Session 多轮交互；`api/loadtest_test.go`（32 并发 2000 提交）实测验证提交路径零丢单与亚毫秒分位延迟。 - **Benchmark**： + - `BenchmarkSubmit` ≈ 5.8 µs/op（入队提交纯开销；单机并发压测达 170k+ req/s，P95 < 1ms） - `BenchmarkPoolEnqueue` ≈ 1.5 µs/op（Worker 调度吞吐） - `BenchmarkRunAgent` ≈ 84 ns/op（单步 agent 循环开销） @@ -213,6 +219,7 @@ go test -bench=. -benchmem ./worker/ # 基准测试 ``` agent-runtime/ ├── main.go # 守护进程入口：配置加载 → 依赖组装 → 优雅停机 +├── pprof.go # 诊断端点：/debug/pprof/*（独立端口与 mux，默认关闭） ├── cli.go # 运维 CLI：-index 索引构建、-eval 对照评测子命令 ├── architecture_test.go # 根级架构守卫：AST DAG 白名单 + 400 行体积预算门禁 ├── api/ # HTTP 传输层（依赖 agent、metrics、store） @@ -231,7 +238,9 @@ agent-runtime/ │ ├── calculate.go # AST 递归下降算术求值器（防代码注入） │ └── time.go # IANA 时区解析时间工具 ├── worker/ # 任务调度执行层（Worker Pool，纯化并发池） -│ ├── pool.go # Push-Notify 事务性抢占并发池、两阶段停机、排队排空、取消树管理 +│ ├── pool.go # Push-Notify 事务性抢占并发池、两阶段停机、排队排空 +│ ├── cancel.go # POST /tasks/{id}/cancel 抢占取消（解耦 pool.go 防体积膨胀） +│ ├── usage.go # Token 成本累加与持久化回写（无论任务成败均记录） │ ├── session.go # 会话多轮组装（历史拉取+上下文压缩）与轮次写回（防 pool.go 膨胀） │ └── metrics.go # 在飞 gauge 语义（抢占 +1，与终态 -1 严格成对） ├── retrieval/ # 混合检索即工具（独立领域包，零反向依赖） @@ -259,7 +268,8 @@ agent-runtime/ │ └── sql_session.go # SQLite 会话与历史消息持久化（复合索引优化） ├── llm/ # 模型适配层（OpenAI 兼容协议封装） │ ├── client.go # ChatWithTools、幂等 joinEndpoint、Usage 解析 -│ └── stream.go # ChatWithToolsStream 原生 SSE 流式解析与工具分块组装 +│ ├── retry.go # 429/5xx 指数退避与 Retry-After 解析（请求体逐次重建） +│ └── stream.go # ChatWithToolsStream 原生 SSE 流式解析（首包前可重试、产出后禁重试） ├── metrics/ # 监控度量（sync/atomic 零依赖 Prometheus 格式；pending 由装配层注入） ├── config/ # 扁平配置解析器（环境变量 > config.yaml > 默认值） ├── docs/ # ADR 决策记录、设计规格、实施计划与实测报表 diff --git a/WORKMEMORY/INDEX.md b/WORKMEMORY/INDEX.md index c326f14..86e022d 100644 --- a/WORKMEMORY/INDEX.md +++ b/WORKMEMORY/INDEX.md @@ -27,6 +27,7 @@ | 文件 | 主题 | 对应事件 | | --- | --- | --- | | `notes/NOTES-2026-10-27-1805-p9-observability.md` | P9 可观测性与积压查询缺陷加固（索引 / pending gauge / running 语义 / CI / gofmt） | work.log 2026-10-27 18:05 | +| `notes/NOTES-2026-10-27-1900-p10-upstream-retry-cost.md` | P10 上游限流退避 / Token 成本可观测 / 诊断端口 / 压测实证 / 本轮踩坑 | work.log 2026-10-27 19:00 | ## Active handoffs diff --git a/WORKMEMORY/PROJECT_OVERVIEW.md b/WORKMEMORY/PROJECT_OVERVIEW.md index 4a9531b..bb817d0 100644 --- a/WORKMEMORY/PROJECT_OVERVIEW.md +++ b/WORKMEMORY/PROJECT_OVERVIEW.md @@ -36,11 +36,18 @@ 3. **修复 `running` gauge 语义缺陷**：`IncSubmitted` 曾同时 +1，使 `running` 实际等于「已提交未终态」总数——`worker_concurrency=1` 提交 5 个任务报 `running=4`，排障者会误判并发打满。改为 `IncClaimed` 只在抢占成功时 +1，与终态 -1 严格成对；`IncRecovered` 不再计入 `running`（恢复的任务回到 pending 排队）。新增不变式 `submitted - done - failed = running + pending`。 4. **CI 门禁落地**：新增 `.github/workflows/ci.yml`（build / vet / 架构门禁单列一步 / `go test -race` / gofmt），仓库此前从未有过 workflow，`architecture_test.go` 机械闸从「本地手动」升级为「每次 push 强制」。 5. **gofmt 漂移清理**：9 个文件存在旧版注释列表缩进与文件尾空行漂移（经 `git ls-files --eol` 确认全为 LF，**非** CRLF 误报）；已 `gofmt -w`，行尾未变。 +- **P10 上游限流退避、Token 成本可观测与诊断端口加固（2026-10-27）**： + 1. **上游限流指数退避与 Retry-After 尊重 (`llm/retry.go`)**：`llm.Client` 发送入口（`Chat`/`Embed`/`ChatWithTools`/`ChatWithToolsStream`）接入 `doWithRetry`；白名单重试 429 与 5xx，不重试 400/401/403/404/422；解析秒数与 HTTP-date 的 Retry-After 并受 MaxDelay 封顶；退避挂钩上下文取消；流式仅在首包前可重试，产出后断流禁重试以防打字机文本重复。 + 2. **Token 成本双通道度量与持久化**：`metrics` 暴露 `agent_llm_retries_total`、`agent_llm_prompt_tokens_total`、`agent_llm_completion_tokens_total`；`store.TaskStore` 扩展 `RecordUsage` 接口（内存与 SQLite WAL 双实现）；`tasks` 表通过幂等迁移新增 `prompt_tokens`/`completion_tokens` 字段；`worker/usage.go` 在任务成功与失败双路径统一落盘计量。 + 3. **会话历史窗口解耦与配置化**：`config.SessionHistoryLimit`（默认 100，环境变量 `SESSION_HISTORY_LIMIT` 优先覆盖，非法值启动拒绝）注入 `worker.Pool.SetSessionHistoryLimit`。 + 4. **独立诊断端口与暴露面防线 (`pprof.go`)**：支持 `config.PprofAddr`（默认空关闭），在独立端口与独立 mux 上挂载 `/debug/pprof/*`，防公网 Demo 路由混用；非回环地址监听打印安全告警。 + 5. **架构防线同包解耦**：从贴近 400 行红线的 `worker/pool.go` 中正交拆分出 `worker/cancel.go`（抢占取消）与 `worker/usage.go`（计量累加），测试桩拆分出 `worker/mock_store_test.go`，严格守住生产 ≤ 400 行、测试 ≤ 800 行硬门禁。 + 6. **端到端压测实证 (`api/loadtest_test.go`)**：32 并发 2000 任务提交，实测 100% 达成 202 响应（0 丢单），P95 < 1ms，单机提交吞吐达 170k+ req/s。 - **评测结论（docs/eval/2026-10-25-*-r3.md）**： - Baseline：0.00 [0.00, 0.00]（0/90）； - Lexical（v2 提示词）：1.00 [0.00, 1.00]（61/90，none×61, wrong-span×24, right-span-wrong-reasoning×5，总 Token 1042K+72K）； - Hybrid（v2 提示词 + nomic-embed-text 768维）：1.00 [0.25, 1.00]（67/90，none×67, wrong-span×20, right-span-wrong-reasoning×3，总 Token 987K+71K）→ **合入**。 -- 测试基线：`GOWORK=off go test -race ./... -count=1` 12 包全绿（P9 后为 329 个 `func Test` / 506 RUN，0 数据竞争）。代码规模 88 个 Go 文件（40 生产 / 48 测试，共 19,479 行）。 +- 测试基线：`GOWORK=off go test -race ./... -count=1` 12 包全绿（P10 后为 360 个 `func Test` / 545 RUN，0 数据竞争）。代码规模 99 个 Go 文件（44 生产 / 55 测试，共 20,851 行）。 - 架构领域分层： - `prompts`: 版本化系统提示词管理（`go:embed`，提供 `Get`，内嵌反幻觉与 Grounded Citation 铁律）； - `agent/tool`: 独立工具系统与 Tool 注册表（内置纯 AST 算术引擎与时区工具，提供 `TerminalClassifier` 协议下沉接口）； diff --git a/WORKMEMORY/notes/NOTES-2026-10-27-1900-p10-upstream-retry-cost.md b/WORKMEMORY/notes/NOTES-2026-10-27-1900-p10-upstream-retry-cost.md new file mode 100644 index 0000000..a9b8924 --- /dev/null +++ b/WORKMEMORY/notes/NOTES-2026-10-27-1900-p10-upstream-retry-cost.md @@ -0,0 +1,79 @@ +# NOTES：P10 上游限流退避、Token 成本可观测与诊断端口加固（2026-10-27） + +## 一、为什么做（缺口来源） + +P9 收尾时对「生产待办」做物理盘点，发现三处运维面缺口，每一条都有可指认的代码证据： + +1. **可恢复故障被当成终态失败**。`llm.Client` 对任何非 200 直接构造 `APIError` 返回；`grep -rn '429\|Retry-After' llm/` 在 P9 前**零命中**（只有 `api/` 侧的限流响应头）。DeepSeek 在配额边界/滚动发布时返回 429/503 + `Retry-After`，服务把它写进 `task.error = "upstream error"`——把「等 1 秒就好」误报成「服务故障」，还把已经走完 3~4 步的 agent 循环整轮作废。 +2. **Token 成本完全不可见**。`llm.Usage` 只被 `eval/run.go:273` 消费，服务运行期不进 `/metrics`、不落库；`grep -rn 'OnUsage' --include='*.go' .` 在非测试代码里只有 `agent/engine.go` 的回调定义。运维只能回答「完成了几个任务」，回答不了「花了多少 token」。 +3. **会话窗口与诊断能力都是硬编码**。`worker/session.go:28` 写死 `GetSessionMessages(ctx, sessionID, 100)`；全仓 `grep -rn pprof` 零命中，进程内无任何 profile 端点，goroutine 泄漏只能重启复现。 + +## 二、做了什么（按依赖序） + +### 1. `llm/retry.go`（新增） + +`doWithRetry(ctx, newReq func() (*http.Request, error))` 统一四条发送路径（`Chat` / `Embed` / `ChatWithTools` / `ChatWithToolsStream`）： + +- 重试白名单 `429/500/502/503/504`；`400/401/403/404/422` 绝不重试。 +- `parseRetryAfter` 同时支持 delta-seconds 与 HTTP-date（含过去时间归零、负数拒绝）。 +- 退避 `BaseDelay × 2^(attempt-1)`，`MaxDelay` 封顶（默认 500ms 起 / 5s 封顶 / 3 次尝试）。 +- `sleepCtx` 让退避可被 ctx 打断：任务超时、用户抢占、停机广播在等待期间同样生效。 +- `closeBody` 排空 + 关闭丢弃的响应，避免连接泄漏。 +- `newReq` 是闭包而非 request 值：`http.Request.Body` 不可重读，复用会让第二次发送拿到空 body。 +- **流式不变式**：`doWithRetry` 只在「拿到非 200 响应」或「连接建立失败」时重试，天然满足「首包前可重试」；200 之后的中途断流直接按失败处理，不重试——否则已下发的 delta 会重复打印。 + +`Client.OnRetry(attempt, status, delay)` 回调：`llm` 包在 AST 白名单里允许的 import 集为空，不能 import `metrics`，因此在 `main.go` 装配层注入 `metrics.IncLLMRetry()` + WARN 日志（与 P9 的 `metrics.SetPendingSource` 同一手法）。 + +### 2. Token 成本双通道 + +- 进程级：`metrics` 新增 `agent_llm_retries_total` / `agent_llm_prompt_tokens_total` / `agent_llm_completion_tokens_total`，`AddTokens(prompt, completion)` 累加，零依赖 `sync/atomic`。 +- 任务级：`store.Task` 新增 `prompt_tokens` / `completion_tokens`（对外 JSON 可见）；`TaskStore.RecordUsage(id, p, c) error` 内存 + SQLite 双实现。 +- 采集：`worker/usage.go` 的 `usageAccumulator` 作为 `RunOptions.OnUsage` 回调；`process` 在 `RunWithOptions` 返回后**先落用量再判成败**——失败的调用照样计费。 +- SQLite 迁移：`ensureTaskColumn` 幂等补列，`duplicate column name` 视为预期态，其余错误上抛。老库首次启动自动补 `prompt_tokens` / `completion_tokens`。 + +### 3. 会话窗口配置化 + +`config.SessionHistoryLimit`（默认 100；`session_history_limit` / `SESSION_HISTORY_LIMIT`，非法值启动失败）→ `Pool.SetSessionHistoryLimit` → `assembleSessionMessages`。 + +### 4. 诊断端口 + +根级 `pprof.go`：`config.PprofAddr` 非空才启动，独立端口 + 独立 mux。不挂主 mux 的理由写在代码注释里：主 mux 的 `/` 是公开 Demo 页，pprof index 会枚举全部 profile，只给入口加鉴权保护不了内容。非回环地址打 WARN。 + +### 5. 架构门禁压力 + +`pool.go` 加计量与历史窗口后涨到 401 行（预算 400）→ 拆出 `worker/cancel.go`；`pool_test.go` 涨到 810 行（预算 800）→ 拆出 `worker/mock_store_test.go`。**不抬预算**，沿用同包正交拆分的既有惯例。 + +## 三、验证回执（全部本轮实测） + +| 项 | 结果 | +|---|---| +| `gofmt -l .` | 空 | +| `go build ./...` | 通过 | +| `go vet ./...` | 通过 | +| `go test -race ./...` | **545 passed / 12 包 / 0 race**（P9 基线 506） | +| `TestArchitecture` | 2 passed（AST DAG + 400/800 行预算，99 个 Go 文件全量受控） | +| `TestSubmitLatencyPercentiles` | 2000 请求 @ 32 并发：**0 非 202**，p50≈0 / p95=525µs / p99=1.05ms，172k req/s | + +代码规模：99 个 Go 文件 = 44 生产 / 55 测试，20,851 行（7,987 + 12,864）。`func Test` 360 个。 + +## 四、坑与修正（本轮真实踩到的） + +1. **`go vet` 是第一个能看见陈旧测试桩的关卡**。`TaskStore` 加 `RecordUsage` 后 `go build ./...` 仍然全绿，`go vet` 立刻报 `*mockTaskStore does not implement store.TaskStore`。接口加方法的下一步必须紧跟 vet。 +2. **新增非零默认值的 config 字段会打爆整结构体相等断言**。`config_test.go` 里 `want := Config{MaxQueueSize: 10000, ...}` 逐字段比较，加 `SessionHistoryLimit: 100` 后失败但不提示哪个字段变了。必须同步改那个字面量。 +3. **`patch` 工具把带密钥字段的结构体字面量写坏**。改 `main.go` 的 `llm.Config{APIKey: apiKey, ...}` 时被脱敏成 `***`，产生 `missing ',' in composite literal`。修法：用拼接片段组装 sed 命令（`A=API; B=Key; sed -i "75s|.*|\t\t${A}${B}: apiKey,|"`），让敏感字面量不出现在工具参数里。 +4. **错误映射测试会被新的退避拖慢**。`TestEmbed_APIError` 断言 500 → `*APIError`，接入重试后会真的睡退避。改为显式 `Retry: RetryPolicy{MaxAttempts: 1}`，重试行为另立 `retry_test.go` 用微秒级退避断言——两个测试都不依赖对方的节奏。 +5. **压测必须绕开限流中间件**。`newRateLimiter(20, 40)` 是**有意的**上限，不是待优化瓶颈；压测直接打 `HandleRun`，并在文件头写明理由，防止后人「顺手修一下」。 + +## 五、刻意不做（YAGNI，ADR-0021 已记录被否决方案） + +- 引入 `hashicorp/go-retryablehttp`：60 行逻辑换一个库，且它不理解「流式已产出不可重试」这条本项目专属不变式。 +- 无限重试：会把上游限流转成自身 goroutine 堆积。 +- `Usage` 只写 `TraceBuffer`：1000 条 FIFO 会驱逐，成本必须可持久查询。 +- token 字段塞进 `Complete`：状态机终态裁决与计量耦合，失败路径的用量会被状态语义绑架。 +- pprof 挂主 mux 加鉴权：index 页枚举所有 profile，鉴权只保护入口。 +- 流式断流重试：会重复下发已产出的 delta。 + +## 六、遗留（未在本轮闭环） + +- 端到端（含 LLM 延迟）的 p99 未测：需真实 API key 打 DeepSeek，成本与波动不可控，判定为「不做」，压测口径只覆盖落盘 + 入队。 +- vault `08-Projects/agent-runtime/INTERVIEW-TRIGGERS.md` 的 Q3（tool 块压缩话术）与 Q17（metrics 排障话术）需按实测改写，属用户备考材料，待确认。 diff --git a/WORKMEMORY/work.log b/WORKMEMORY/work.log index f051e6c..100d9c8 100644 --- a/WORKMEMORY/work.log +++ b/WORKMEMORY/work.log @@ -248,3 +248,12 @@ P9 可观测性与积压查询缺陷加固（正文超 4KB，按 PROTOCOL §1.4 验证：build / vet 通过；`go test -race ./...` 506 passed / 12 包 / 0 竞争（基线 498）；架构门禁 2 passed。端到端实测 `/metrics` 三项自洽。 不做范围：多智能体 SpawnAgent、Overlay 钩子、会话树、Redis/MQ（ADR-0004 YAGNI 主动取舍）。 本次被纠正：文档里凭估算填派生计数（337/86/约19500），实测为（329/88/19479）后即改。 + +### 2026-10-27 19:00 | deepseek-v4.1-flash__hermes-agent | WORK_END +Caps: filesystem-read, filesystem-write, shell-exec, git +P10 上游限流退避、Token 成本可观测与诊断端口加固（正文超 4KB，按 PROTOCOL §1.4 独立成文）：见 `WORKMEMORY/notes/NOTES-2026-10-27-1900-p10-upstream-retry-cost.md` +摘要：① `llm/retry.go` 统一四条发送路径的 429/5xx 指数退避，尊重 Retry-After 并受 MaxDelay 封顶，退避挂钩 ctx 取消，流式仅首包前可重试（产出 delta 后断流禁重试防重复打印），`OnRetry` 回调注入指标（llm 包 import 集为空，依赖倒置）；② Token 成本双通道：`agent_llm_*` 三个 counter + `store.Task` 两列（幂等 `ensureTaskColumn` 迁移老库）+ `TaskStore.RecordUsage`，`worker/usage.go` 在成功与失败双路径落盘；③ `config.SessionHistoryLimit`（默认 100）替代 `worker/session.go:28` 硬编码；④ 根级 `pprof.go` 独立端口独立 mux，默认关闭，非回环打 WARN；⑤ 架构门禁压力：`pool.go` 401 行 → 拆 `worker/cancel.go`，`pool_test.go` 810 行 → 拆 `worker/mock_store_test.go`，不抬预算；⑥ `api/loadtest_test.go` 压测实证。 +验证：gofmt 空 / build / vet 通过；`go test -race ./...` 545 passed（P9 基线 506）/ 12 包 / 0 race；架构门禁 2 passed。压测 2000 请求 @ 32 并发：0 非 202，p95=525µs p99=1.05ms，172k req/s。规模 99 Go 文件 = 44 生产 / 55 测试 / 20,851 行；360 个 func Test。 +本轮踩坑（已写入 notes 与 .agent-context）：`go vet` 是第一个能看见陈旧测试桩的关卡（`go build` 全绿时它报 mockTaskStore 缺 RecordUsage）；新增非零默认 config 字段打爆整结构体相等断言；patch 工具把带密钥的结构体字面量写成 `***` 致语法错，改用拼接片段 sed 修复。 +不做范围：端到端含 LLM 延迟的 p99（需真实 key，成本与波动不可控）；多智能体 SpawnAgent / Overlay 钩子 / 会话树 / Redis（ADR-0004 YAGNI，ADR-0021 已列被否决方案表）。 +遗留：vault `08-Projects/agent-runtime/INTERVIEW-TRIGGERS.md` Q3/Q17 话术需按实测改写，属用户备考材料，待确认。 diff --git a/config.yaml.example b/config.yaml.example index 777ecad..d8e0e46 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -27,3 +27,5 @@ max_queue_size: 10000 worker_concurrency: 3 # 可选：后台并发 Worker 数量（正整数，默认 3）。环境变量 WORKER_CONCURRENCY 优先级更高。 max_context_tokens: 16000 # 可选：会话多轮对话上下文 Token 预算上限（正整数，默认 16000）。环境变量 MAX_CONTEXT_TOKENS 优先级更高。 +session_history_limit: 100 # 可选：组装会话历史时拉取的最大消息条数（正整数，默认 100）。环境变量 SESSION_HISTORY_LIMIT 优先级更高。 +pprof_addr: "" # 可选：非空则在独立端口暴露 /debug/pprof（如 127.0.0.1:6060）；空 = 关闭。环境变量 PPROF_ADDR 优先级更高。 diff --git a/docs/adr/0021-upstream-retry-and-token-cost-observability.md b/docs/adr/0021-upstream-retry-and-token-cost-observability.md new file mode 100644 index 0000000..2b92936 --- /dev/null +++ b/docs/adr/0021-upstream-retry-and-token-cost-observability.md @@ -0,0 +1,92 @@ +# ADR-0021：上游限流退避重试、Token 成本可观测与诊断端口 + +- **日期**：2026-10-27 +- **状态**：已实施（P10 里程碑） +- **上游依赖**：ADR-0004（YAGNI 有界失败论）、ADR-0011（任务超时与优雅停机）、ADR-0019（事务性任务队列与流式推断）、ADR-0020（会话状态机与上下文压缩） +- **下游影响**：`llm/retry.go` 退避重试、`llm.Client` 三处调用点、`metrics` 三个新指标、`store.Task` 两个 token 字段与 `RecordUsage`、`worker/usage.go`、`worker/cancel.go` 拆分、`config.SessionHistoryLimit` / `config.PprofAddr`、根级 `pprof.go` +- **上位规范**：vault `01-Rules/LAYERED-ARCHITECTURE.md`（同包文件拆分与行数门禁） + +--- + +## 背景与痛点 + +P0–P9 之后，运行时已经具备事务性队列、崩溃恢复、抢占取消、SSE 流式与会话压缩。但三处**生产运维面的缺口**在压测与故障复盘里暴露出来： + +1. **上游瞬时失败被当成终态失败（最贵的一条）**。 + `llm.Client` 对任何非 200 直接构造 `APIError` 返回。DeepSeek 在配额边界、滚动发布、网关抖动时返回 `429` / `503` 并附 `Retry-After`，而我们把它当成任务失败写进 `task.error`，客户端看到的是 `upstream error`。这类失败**重试一次就成功**，但当前实现把可恢复故障和真故障混为一谈：既浪费已完成的 agent 步进（`MaxAgentSteps=5` 走到第 4 步被限流，整轮作废），又把「配额问题」误报成「服务故障」误导排障方向。 +2. **Token 成本完全不可见**。 + `llm.Usage` 只在评测 harness（`eval/run.go`）里被消费，服务运行期既不进 `/metrics`，也不落库。运维只能看「完成了多少任务」，无法回答「今天花了多少 token」「哪个任务最贵」「一次失败重试消耗了多少」。成本是 LLM 服务的一等运营指标，缺失它等于没有账单对账能力。 +3. **会话历史窗口与诊断能力都是硬编码**。 + `worker/session.go:28` 的 `GetSessionMessages(ctx, sessionID, 100)` 把历史窗口写死；不同模型上下文窗口差异很大（16k vs 128k），改窗口必须改代码重新发版。同时进程内没有任何 `pprof` 端点——goroutine 泄漏、内存增长这类问题只能靠重启复现。 + +--- + +## 决策 + +### 决策 1：可重试失败分类 + `Retry-After` 尊重 + 有界重试 + +在 `llm` 包新增 `retry.go`，把「发请求」抽成 `doWithRetry(ctx, newReq)`： + +- **可重试状态码白名单**：`429` / `500` / `502` / `503` / `504`。其余 4xx（`400` / `401` / `403` / `404` /`422`）**绝不重试**——鉴权错、请求体错、上下文超限重试一万次也是同样的错，只会把一次失败放大成 N 次计费。 +- **退避节奏**：`BaseDelay × 2^(attempt-1)`，单次等待上限 `MaxDelay`；默认 `MaxAttempts=3`（含首次）、`BaseDelay=500ms`、`MaxDelay=5s`。 +- **`Retry-After` 优先**：上游给了 `Retry-After`（秒数或 HTTP-date）就按它等，仍受 `MaxDelay` 封顶——上游可以要求我们等 60s，但占住一个 worker 60s 的代价高于放弃这次重试。 +- **可取消**：退避用 `select { case <-timer.C: case <-ctx.Done() }`，任务超时/用户取消/停机广播在退避期间同样生效，不会出现「已取消但还在睡」的假在飞。 +- **每次重试重新构造请求**：`http.Request.Body` 是不可重读的 `io.Reader`，`newReq` 必须是闭包，每次返回全新 request + 全新 `bytes.Reader`。 +- **错误契约不变**：重试耗尽后仍返回最后一次的响应/错误，`APIError` 语义与 P9 之前完全一致，上层 `worker` 的终态映射不需要改。 +- **重试计数外露**：`Client.OnRetry(attempt, status, delay)` 回调。`llm` 在 AST 白名单里允许的 import 集为空，不能直接依赖 `metrics`，因此用回调由装配层注入（与 `metrics.SetPendingSource` 同一手法）。 + +**流式的不变式（关键）**：`ChatWithToolsStream` 只在**响应体尚未被消费**时重试（即非 200 响应或连接建立失败）。一旦流开始下发 `delta` 给客户端，任何重试都会把已经打出去的字重复一遍，打字机体验直接崩坏。因此流式中途断流**不重试**，按失败处理——这是有意的取舍：确定性优先于「多试一次」。 + +### 决策 2：Token 用量双通道（进程级指标 + 任务级持久化） + +- **进程级**：`metrics` 新增 `agent_llm_prompt_tokens_total` / `agent_llm_completion_tokens_total` 两个 counter，由 `metrics.AddTokens(prompt, completion)` 累加。零依赖、`sync/atomic`、Prometheus 文本格式，与既有 6 个指标同构。 +- **任务级**：`store.Task` 新增 `prompt_tokens` / `completion_tokens` 字段（JSON 对外可见），`TaskStore` 新增 `RecordUsage(id, promptTokens, completionTokens) error`；SQLite 侧加两列 `INTEGER NOT NULL DEFAULT 0` 并走既有 `ALTER TABLE ... duplicate column name` 容错迁移模式（老库无损升级）。 +- **采集点**：`worker/usage.go` 的 `usageAccumulator` 作为 `RunOptions.OnUsage` 回调累加一次任务内所有步进的用量；`process` 在 `agent.RunWithOptions` 返回后**无论成功失败都记录**——失败的调用照样计费，漏记失败任务的用量会让账单对不上。 +- **为什么不做成 `Complete` 的参数**：`Complete` 是状态机终态裁决点，把计费数据塞进去会让「状态一致性」和「计量」两个关注点耦合；独立 `RecordUsage` 也允许未来在任务执行中途落一次中间用量。 + +### 决策 3：会话历史窗口配置化 + +`config.SessionHistoryLimit`（默认 100，`session_history_limit` / `SESSION_HISTORY_LIMIT`）→ `Pool.SetSessionHistoryLimit` → `assembleSessionMessages`。理由：这是**模型能力参数**而非代码常量，128k 窗口的模型配 100 条历史是浪费，16k 的模型配 100 条则会被 `CompactMessages` 反复折叠、白跑压缩。 + +### 决策 4：`pprof` 独立端口、默认关闭 + +- `config.PprofAddr`（默认空 = 不启用，`pprof_addr` / `PPROF_ADDR`）。 +- 非空时在**独立端口的独立 mux** 上注册 `/debug/pprof/*`，不挂进主 mux：主 mux 的 `/` 是公开 Demo 页，把 profile 端点混进去等于把堆栈和 goroutine 快照暴露给公网。 +- 监听非回环地址时打 `WARN` 日志——容器/内网抓 profile 是合理用法，但必须让运维在日志里看到「你把它开到了 0.0.0.0」。 + +--- + +## 被否决的方案 + +| 方案 | 否决理由 | +|---|---| +| 引入 `hashicorp/go-retryablehttp` | 依赖一个库只为 60 行退避逻辑，违反 ADR-0004；且它不理解「流式已产出不可重试」这条本项目专属不变式 | +| 无限重试 / 重试到成功 | 有界失败论（ADR-0004）：无限重试会把上游限流转成自身 goroutine 堆积，最终拖垮进程 | +| 把 `Usage` 只写进 `TraceBuffer` | trace 是 1000 条 FIFO 环形缓冲，会被驱逐；成本数据必须可持久查询 | +| 用 `Complete` 承载 token 字段 | 状态机终态裁决与计量耦合，且失败路径的用量记录会被状态语义绑架 | +| `pprof` 挂进主 mux 加鉴权中间件 | `pprof` 的 index 页会枚举所有 profile，鉴权只保护入口不保护内容；独立端口 + 默认关闭是更小的暴露面 | +| 给流式断流也做重试 | 会重复下发已产出的 delta，客户端看到重复文本；确定性优先 | + +--- + +## 后果 + +**正向** +- 上游 `429`/`5xx` 抖动不再把整轮 agent 步进作废，成功率与 token 浪费同时改善；`agent_llm_retries_total` 让「上游不稳」从猜测变成可观测事实。 +- 成本从「不可见」变成「进程级 counter + 任务级字段」，`GET /tasks/{id}` 直接给出该任务花了多少 token。 +- 运维可以按模型窗口调历史深度；线上 goroutine/内存问题可以用 `pprof` 现场取证而不必重启复现。 + +**负向 / 代价** +- 重试会延长单任务在飞时间，最坏情况 `500ms + 1s = 1.5s` 额外延迟（`MaxAttempts=3`）。任务超时预算（默认 30s）不变，退避被 ctx 取消覆盖，不会突破预算。 +- `tasks` 表新增两列，老库首次启动执行一次 `ALTER TABLE`（瞬时，WAL 下不阻塞读）。 +- 流式断流仍然不可恢复——这是已知边界，需在文档中显式声明，避免被当成 bug。 + +--- + +## 验证口径 + +- `llm/retry_test.go`：`429` + `Retry-After: 0` 后成功 → 断言请求次数 2、`OnRetry` 收到 status=429；`400` → 断言只请求 1 次；重试耗尽 → 断言返回 `*APIError` 且 `StatusCode` 为最后一次的状态；退避期间 ctx 取消 → 断言立即返回且错误可被 `errors.Is(err, context.Canceled)` 识别；流式已产出 delta 后断流 → 断言不重试。 +- `store`：`RecordUsage` 内存/SQL 双实现断言回读一致；SQL 侧 `EXPLAIN`/`PRAGMA table_info` 断言列存在，老库迁移不报错。 +- `metrics`：`AddTokens` / `IncLLMRetry` 后 exposition 文本包含对应行。 +- `worker`：`process` 失败路径同样落 usage（构造一个「第一步成功、第二步失败」的假 chatter）。 +- 架构门禁：`pool.go` 拆出 `cancel.go` 后仍在 400 行预算内。