diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..33d12f3 --- /dev/null +++ b/.gitignore @@ -0,0 +1,12 @@ +# 运行数据（usage.db 与 blobs/，含请求/响应记录，可能含敏感信息） +/data/ + +# 编译产物 +*.exe +/build/ + +# 运行日志 +*.log + +# 本地说明文件（含上游密钥，不入库） +说明.md diff --git a/build.ps1 b/build.ps1 new file mode 100644 index 0000000..cb7bda8 --- /dev/null +++ b/build.ps1 @@ -0,0 +1,12 @@ +# 编译脚本：源码在 src/go/，产物输出到 build/relay.exe +$ErrorActionPreference = 'Stop' +$root = $PSScriptRoot +$out = Join-Path $root 'build\relay.exe' + +Push-Location (Join-Path $root 'src\go') +try { + go build -o $out . + Write-Host "Built: $out" +} finally { + Pop-Location +} diff --git a/src/go/entry.go b/src/go/entry.go new file mode 100644 index 0000000..15d4c30 --- /dev/null +++ b/src/go/entry.go @@ -0,0 +1,47 @@ +package main + +import ( + "context" + "time" +) + +// Entry 一次完整的请求-响应留存记录。 +type Entry struct { + ID int64 + Timestamp time.Time + ClientIP string + + Method string + Path string + + ReqHeader map[string]string + ReqBody []byte + + Status int + RespHeader map[string]string + RespBody []byte + DurationMs int64 + + // 解析出的用量信息 + APIFormat string // openai | claude | ollama | unknown + Model string + InputTokens int + OutputTokens int + TotalTokens int + IsStream bool + + // 原文落盘后的相对路径 + BlobPath string +} + +// entryCtxKey 用于在 request context 里传递 Entry，避免污染全局。 +type entryCtxKey struct{} + +func withEntry(ctx context.Context, e *Entry) context.Context { + return context.WithValue(ctx, entryCtxKey{}, e) +} + +func entryFrom(ctx context.Context) *Entry { + e, _ := ctx.Value(entryCtxKey{}).(*Entry) + return e +} diff --git a/src/go/go.mod b/src/go/go.mod new file mode 100644 index 0000000..0c5f08f --- /dev/null +++ b/src/go/go.mod @@ -0,0 +1,17 @@ +module goaiapi + +go 1.26.3 + +require modernc.org/sqlite v1.51.0 + +require ( + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + golang.org/x/sys v0.42.0 // indirect + modernc.org/libc v1.72.3 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect +) diff --git a/src/go/go.sum b/src/go/go.sum new file mode 100644 index 0000000..0705615 --- /dev/null +++ b/src/go/go.sum @@ -0,0 +1,51 @@ +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8= +golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= +golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k= +golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0= +modernc.org/cc/v4 v4.28.2 h1:3tQ0lf2ADtoby2EtSP+J7IE2SHwEJdP8ioR59wx7XpY= +modernc.org/cc/v4 v4.28.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= +modernc.org/ccgo/v4 v4.34.0 h1:yRLPFZieg532OT4rp4JFNIVcquwalMX26G95WQDqwCQ= +modernc.org/ccgo/v4 v4.34.0/go.mod h1:AS5WYMyBakQ+fhsHhtP8mWB82KTGPkNNJDGfGQCe0/A= +modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= +modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo= +modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.72.3 h1:ZnDF4tXn4NBXFutMMQC4vtbTFSXhhKzR73fv0beZEAU= +modernc.org/libc v1.72.3/go.mod h1:dn0dZNnnn1clLyvRxLxYExxiKRZIRENOfqQ8XEeg4Qs= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg= +modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.51.0 h1:aH/MMSoayAIhozZ7uJbVTT9QO/VhzBf0J9tymmmuC/U= +modernc.org/sqlite v1.51.0/go.mod h1:tcNzv5p84E0skkmJn038y+hWJbLQXQqEnQfeh5r2JLM= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/src/go/main.go b/src/go/main.go new file mode 100644 index 0000000..996010f --- /dev/null +++ b/src/go/main.go @@ -0,0 +1,186 @@ +package main + +import ( + "bytes" + "flag" + "io" + "log" + "net/http" + "net/http/httputil" + "net/url" + "strings" + "sync/atomic" + "time" +) + +// 命令行参数：本地监听地址 + 上游地址 + 数据目录 +var ( + listenAddr = flag.String("l", "127.0.0.1:8080", "本地监听地址 host:port") + upstream = flag.String("u", "https://api.ai.pulsareon.com", "上游地址 scheme://host[:port]") + dataDir = flag.String("d", "./data", "留存数据目录") + webAddr = flag.String("a", "127.0.0.1:8081", "Web 管理界面地址 host:port（空字符串关闭）") +) + +var connID atomic.Int64 + +func main() { + flag.Parse() + + target, err := url.Parse(*upstream) + if err != nil { + log.Fatalf("上游地址解析失败 %s: %v", *upstream, err) + } + + rec, err := NewRecorder(*dataDir) + if err != nil { + log.Fatalf("初始化留存模块失败: %v", err) + } + defer rec.Close() + + // Web 管理界面（独立端口，只读）。空地址则关闭。 + if *webAddr != "" { + go startWebUI(*webAddr, rec) + } + + proxy := &httputil.ReverseProxy{ + // Director 改写出站请求：指向上游，保持原路径 + Director: func(req *http.Request) { + req.URL.Scheme = target.Scheme + req.URL.Host = target.Host + req.Host = target.Host // 让上游看到正确的 Host 头 + }, + // ModifyResponse 在响应返回客户端前介入，做全量留存 + ModifyResponse: recordResponse(rec), + ErrorHandler: func(w http.ResponseWriter, r *http.Request, e error) { + log.Printf("[!] 转发失败 %s %s: %v", r.Method, r.URL.Path, e) + http.Error(w, "proxy error: "+e.Error(), http.StatusBadGateway) + }, + } + + // 外层 handler：抓请求原文 + 计时，再交给 proxy + handler := func(w http.ResponseWriter, r *http.Request) { + // seq 仅是本进程内的日志序号，方便把同一次请求的进/出两条日志对上； + // 它不是持久化主键——真正的主键由 SQLite AUTOINCREMENT 在 Save 时分配。 + seq := connID.Add(1) + start := time.Now() + + // 读取并缓存请求体（读完要还回去，否则 proxy 转发不到） + var reqBody []byte + if r.Body != nil { + reqBody, _ = io.ReadAll(r.Body) + r.Body = io.NopCloser(bytes.NewReader(reqBody)) + } + + // 把本次调用的上下文挂到 request 上，供 ModifyResponse 取用。 + // ID 留空，由 Recorder.Save 写库后回填权威主键。 + entry := &Entry{ + Timestamp: start, + ClientIP: clientIP(r), + Method: r.Method, + Path: r.URL.Path, + ReqHeader: flattenHeader(r.Header), + ReqBody: reqBody, + } + r = r.WithContext(withEntry(r.Context(), entry)) + + log.Printf("[+] #%d %s %s from %s", seq, r.Method, r.URL.Path, entry.ClientIP) + proxy.ServeHTTP(w, r) + } + + srv := &http.Server{ + Addr: *listenAddr, + Handler: http.HandlerFunc(handler), + } + log.Printf("反向代理启动: %s -> %s (数据目录 %s)", *listenAddr, *upstream, *dataDir) + if err := srv.ListenAndServe(); err != nil { + log.Fatalf("服务退出: %v", err) + } +} + +// recordResponse 返回一个 ModifyResponse 回调，做全量留存 + usage 解析。 +// +// 非流式：直接 io.ReadAll 读完响应体、解析 JSON usage、落库。 +// 流式（SSE）：不能 io.ReadAll（会缓冲整条流，客户端要等所有 chunk 到齐才收到第一个字）。 +// 改用 streamRecorder 包住 resp.Body——数据被客户端读取时同步流过， +// 一边原样转发、一边拼回全量原文并解析 usage，读到 EOF 时回调落库。 +func recordResponse(rec *Recorder) func(*http.Response) error { + return func(resp *http.Response) error { + entry := entryFrom(resp.Request.Context()) + if entry == nil { + return nil // 没有上下文，跳过留存但不影响转发 + } + + entry.Status = resp.StatusCode + entry.RespHeader = flattenHeader(resp.Header) + entry.IsStream = isStream(resp) + + if entry.IsStream { + // 流式：包装 body，落库延迟到流结束（onDone）时进行。 + format := guessFormatByPath(entry.Path) + src := resp.Body + resp.Body = newStreamRecorder(src, format, func(raw []byte, in, out int, model string) { + entry.RespBody = raw + entry.DurationMs = time.Since(entry.Timestamp).Milliseconds() + entry.InputTokens = in + entry.OutputTokens = out + entry.TotalTokens = in + out + if model != "" { + entry.Model = model + } + // 流式响应体里没有顶层 usage 时，仍按路径标注格式 + entry.APIFormat = format + saveEntry(rec, entry) + }) + return nil + } + + // 非流式：读完响应体并还回去（客户端还要收） + respBody, err := io.ReadAll(resp.Body) + resp.Body.Close() + if err != nil { + log.Printf("[!] #%d 读取响应体失败: %v", entry.ID, err) + } + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + + entry.RespBody = respBody + entry.DurationMs = time.Since(entry.Timestamp).Milliseconds() + parseUsage(entry, respBody) + saveEntry(rec, entry) + return nil + } +} + +// saveEntry 落库并打印一条结果日志，统一非流式/流式两条路径的收尾。 +func saveEntry(rec *Recorder, entry *Entry) { + if err := rec.Save(entry); err != nil { + log.Printf("[!] 留存失败: %v", err) + return + } + log.Printf("[-] #%d done status=%d %dms model=%s in=%d out=%d stream=%v", + entry.ID, entry.Status, entry.DurationMs, entry.Model, + entry.InputTokens, entry.OutputTokens, entry.IsStream) +} + +// isStream 判断响应是否为 SSE 流式。 +func isStream(resp *http.Response) bool { + ct := resp.Header.Get("Content-Type") + return strings.Contains(strings.ToLower(ct), "text/event-stream") +} + +// clientIP 取客户端 IP（去掉端口）。 +func clientIP(r *http.Request) string { + ip := r.RemoteAddr + if i := strings.LastIndex(ip, ":"); i >= 0 { + ip = ip[:i] + } + return ip +} + +// flattenHeader 把 http.Header 拍平成单层 map，便于 JSON 留存。 +func flattenHeader(h http.Header) map[string]string { + m := make(map[string]string, len(h)) + for k, v := range h { + m[k] = strings.Join(v, ", ") + } + return m +} diff --git a/src/go/proxy_test.go b/src/go/proxy_test.go new file mode 100644 index 0000000..fa5f136 --- /dev/null +++ b/src/go/proxy_test.go @@ -0,0 +1,188 @@ +package main + +import ( + "bytes" + "context" + "database/sql" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "net/http/httputil" + "net/url" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" +) + +// newTestProxy 用被测的 Director/ModifyResponse 组装一个反向代理，指向 mock 上游。 +// 这一段刻意复刻 main() 里的装配逻辑，保证测的是真实代码路径。 +func newTestProxy(t *testing.T, upstreamURL string, rec *Recorder) http.Handler { + t.Helper() + target, err := url.Parse(upstreamURL) + if err != nil { + t.Fatalf("parse upstream: %v", err) + } + var id atomic.Int64 + proxy := &httputil.ReverseProxy{ + Director: func(req *http.Request) { + req.URL.Scheme = target.Scheme + req.URL.Host = target.Host + req.Host = target.Host + }, + ModifyResponse: recordResponse(rec), + } + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + start := time.Now() + var reqBody []byte + if r.Body != nil { + reqBody, _ = io.ReadAll(r.Body) + r.Body = io.NopCloser(bytes.NewReader(reqBody)) + } + entry := &Entry{ + ID: id.Add(1), + Timestamp: start, + ClientIP: clientIP(r), + Method: r.Method, + Path: r.URL.Path, + ReqHeader: flattenHeader(r.Header), + ReqBody: reqBody, + } + r = r.WithContext(withEntry(r.Context(), entry)) + proxy.ServeHTTP(w, r) + }) +} + +// TestProxyForwardAndRecord 验证非流式全链路： +// client -> proxy -> mock上游 -> proxy(留存) -> client， +// 并检查：响应原样回传、原文落盘完整、SQLite 索引含正确 usage。 +func TestProxyForwardAndRecord(t *testing.T) { + // 1) mock 上游：模拟 OpenAI 非流式响应，带 usage + respBody := `{"id":"chatcmpl-1","model":"gpt-4o-mini","choices":[{"message":{"role":"assistant","content":"你好"}}],"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18}}` + var gotAuth, gotReqBody string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + b, _ := io.ReadAll(r.Body) + gotReqBody = string(b) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + io.WriteString(w, respBody) + })) + defer upstream.Close() + + // 2) recorder 落到临时目录 + dir := t.TempDir() + rec, err := NewRecorder(dir) + if err != nil { + t.Fatalf("NewRecorder: %v", err) + } + defer rec.Close() + + // 3) 起代理 + proxy := httptest.NewServer(newTestProxy(t, upstream.URL, rec)) + defer proxy.Close() + + // 4) 客户端发请求（自带 key，验证透传不丢） + reqPayload := `{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}` + req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost, + proxy.URL+"/v1/chat/completions", strings.NewReader(reqPayload)) + req.Header.Set("Authorization", "Bearer sk-test-123") + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("client do: %v", err) + } + gotResp, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + // --- 断言 A：响应原样回传给客户端 --- + if string(gotResp) != respBody { + t.Fatalf("响应未原样回传:\n want %q\n got %q", respBody, gotResp) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + + // --- 断言 B：key 透传到上游、请求体完整 --- + if gotAuth != "Bearer sk-test-123" { + t.Fatalf("key 未透传到上游: got %q", gotAuth) + } + if gotReqBody != reqPayload { + t.Fatalf("请求体未完整转发:\n want %q\n got %q", reqPayload, gotReqBody) + } + + // ModifyResponse 是同步执行的（在响应回客户端前），此时记录已落盘。 + // 5) 校验 SQLite 索引 + db, err := sql.Open("sqlite", filepath.Join(dir, "usage.db")) + if err != nil { + t.Fatalf("open db: %v", err) + } + defer db.Close() + + var ( + apiFormat, model, endpoint, blobPath string + in, out, total, status int + isStream int + ) + row := db.QueryRow(`SELECT api_format, model, endpoint, input_tokens, output_tokens, + total_tokens, is_stream, status, blob_path FROM traffic LIMIT 1`) + if err := row.Scan(&apiFormat, &model, &endpoint, &in, &out, &total, &isStream, &status, &blobPath); err != nil { + t.Fatalf("scan index row: %v", err) + } + if apiFormat != "openai" { + t.Errorf("api_format = %q, want openai", apiFormat) + } + if model != "gpt-4o-mini" { + t.Errorf("model = %q, want gpt-4o-mini", model) + } + if endpoint != "/v1/chat/completions" { + t.Errorf("endpoint = %q", endpoint) + } + if in != 11 || out != 7 || total != 18 { + t.Errorf("usage = in:%d out:%d total:%d, want 11/7/18", in, out, total) + } + if isStream != 0 { + t.Errorf("is_stream = %d, want 0", isStream) + } + if status != 200 { + t.Errorf("status = %d, want 200", status) + } + + // --- 断言 C：原文落盘完整，请求/响应原文都在 --- + raw, err := os.ReadFile(filepath.Join(dir, blobPath)) + if err != nil { + t.Fatalf("read blob: %v", err) + } + var b blob + if err := json.Unmarshal(raw, &b); err != nil { + t.Fatalf("unmarshal blob: %v", err) + } + if b.ReqBody != reqPayload { + t.Errorf("blob 请求原文不完整:\n want %q\n got %q", reqPayload, b.ReqBody) + } + if b.RespBody != respBody { + t.Errorf("blob 响应原文不完整:\n want %q\n got %q", respBody, b.RespBody) + } + // 安全留存的关键：客户端的 key 出现在留存的请求头里（全量留存要求） + if got := b.ReqHeader["Authorization"]; got != "Bearer sk-test-123" { + t.Errorf("留存的请求头缺少 Authorization: got %q", got) + } + t.Logf("全链路通过: model=%s in=%d out=%d blob=%s", model, in, out, blobPath) +} + +// TestParseUsageClaude 单独验证 Claude 字段命名也能解析。 +func TestParseUsageClaude(t *testing.T) { + body := []byte(`{"model":"claude-3-5-sonnet","usage":{"input_tokens":25,"output_tokens":40}}`) + var e Entry + parseUsage(&e, body) + if e.APIFormat != "claude" { + t.Fatalf("api_format = %q, want claude", e.APIFormat) + } + if e.InputTokens != 25 || e.OutputTokens != 40 || e.TotalTokens != 65 { + t.Fatalf("claude usage = %d/%d/%d, want 25/40/65", e.InputTokens, e.OutputTokens, e.TotalTokens) + } +} diff --git a/src/go/query.go b/src/go/query.go new file mode 100644 index 0000000..2d25622 --- /dev/null +++ b/src/go/query.go @@ -0,0 +1,162 @@ +package main + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" +) + +// Record 是 traffic 索引表的一行（不含原文，原文在 blob 文件里，按需取）。 +type Record struct { + ID int64 `json:"id"` + TS int64 `json:"ts"` // unix 毫秒 + ClientIP string `json:"client_ip"` + APIFormat string `json:"api_format"` + Endpoint string `json:"endpoint"` + Model string `json:"model"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + TotalTokens int `json:"total_tokens"` + IsStream bool `json:"is_stream"` + DurationMs int64 `json:"duration_ms"` + Status int `json:"status"` + BlobPath string `json:"blob_path"` +} + +// ListOptions 列表查询过滤条件。空值表示不过滤。 +type ListOptions struct { + Limit int + Offset int + Model string // 精确匹配 + Format string // 精确匹配 + Status int // >0 时精确匹配 +} + +// List 按条件查询索引记录，按 id 倒序（最新在前）；返回本页记录 + 满足条件的总数。 +func (r *Recorder) List(opts ListOptions) ([]Record, int, error) { + var where []string + var args []any + if opts.Model != "" { + where = append(where, "model = ?") + args = append(args, opts.Model) + } + if opts.Format != "" { + where = append(where, "api_format = ?") + args = append(args, opts.Format) + } + if opts.Status > 0 { + where = append(where, "status = ?") + args = append(args, opts.Status) + } + clause := "" + if len(where) > 0 { + clause = "WHERE " + strings.Join(where, " AND ") + } + + // 总数（用于前端分页显示） + var total int + if err := r.db.QueryRow("SELECT COUNT(*) FROM traffic "+clause, args...).Scan(&total); err != nil { + return nil, 0, fmt.Errorf("统计总数: %w", err) + } + + limit := opts.Limit + if limit <= 0 || limit > 500 { + limit = 100 + } + q := fmt.Sprintf(`SELECT id, ts, client_ip, api_format, endpoint, model, + input_tokens, output_tokens, total_tokens, is_stream, duration_ms, status, blob_path + FROM traffic %s ORDER BY id DESC LIMIT ? OFFSET ?`, clause) + args = append(args, limit, opts.Offset) + + rows, err := r.db.Query(q, args...) + if err != nil { + return nil, 0, fmt.Errorf("查询列表: %w", err) + } + defer rows.Close() + + var out []Record + for rows.Next() { + var rec Record + var isStream int + if err := rows.Scan(&rec.ID, &rec.TS, &rec.ClientIP, &rec.APIFormat, &rec.Endpoint, + &rec.Model, &rec.InputTokens, &rec.OutputTokens, &rec.TotalTokens, + &isStream, &rec.DurationMs, &rec.Status, &rec.BlobPath); err != nil { + return nil, 0, fmt.Errorf("扫描行: %w", err) + } + rec.IsStream = isStream != 0 + out = append(out, rec) + } + return out, total, rows.Err() +} + +// ReadBlob 按记录 id 找到 blob_path 并读回完整原文。 +func (r *Recorder) ReadBlob(id int64) (*blob, error) { + var rel string + if err := r.db.QueryRow("SELECT blob_path FROM traffic WHERE id = ?", id).Scan(&rel); err != nil { + return nil, fmt.Errorf("查记录 %d: %w", id, err) + } + abs := filepath.Join(r.dataDir, rel) + data, err := os.ReadFile(abs) + if err != nil { + return nil, fmt.Errorf("读原文 %s: %w", rel, err) + } + var b blob + if err := json.Unmarshal(data, &b); err != nil { + return nil, fmt.Errorf("解析原文: %w", err) + } + return &b, nil +} + +// ModelStat 单个模型的聚合统计。 +type ModelStat struct { + Model string `json:"model"` + Count int `json:"count"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + TotalTokens int64 `json:"total_tokens"` +} + +// Stats 整体聚合：总请求数、成功/失败数、token 总量、按模型明细。 +type Stats struct { + TotalRequests int `json:"total_requests"` + SuccessCount int `json:"success_count"` + ErrorCount int `json:"error_count"` + TotalTokens int64 `json:"total_tokens"` + ByModel []ModelStat `json:"by_model"` +} + +// Stats 计算整体用量聚合。 +func (r *Recorder) Stats() (*Stats, error) { + s := &Stats{ByModel: []ModelStat{}} + row := r.db.QueryRow(`SELECT + COUNT(*), + COALESCE(SUM(CASE WHEN status >= 200 AND status < 400 THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN status >= 400 OR status = 0 THEN 1 ELSE 0 END), 0), + COALESCE(SUM(total_tokens), 0) + FROM traffic`) + if err := row.Scan(&s.TotalRequests, &s.SuccessCount, &s.ErrorCount, &s.TotalTokens); err != nil { + return nil, fmt.Errorf("汇总: %w", err) + } + + rows, err := r.db.Query(`SELECT + COALESCE(NULLIF(model, ''), '(未知)') AS m, + COUNT(*), + COALESCE(SUM(input_tokens), 0), + COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(total_tokens), 0) + FROM traffic GROUP BY m ORDER BY SUM(total_tokens) DESC, COUNT(*) DESC`) + if err != nil { + return nil, fmt.Errorf("按模型聚合: %w", err) + } + defer rows.Close() + for rows.Next() { + var ms ModelStat + if err := rows.Scan(&ms.Model, &ms.Count, &ms.InputTokens, &ms.OutputTokens, &ms.TotalTokens); err != nil { + return nil, fmt.Errorf("扫描模型统计: %w", err) + } + s.ByModel = append(s.ByModel, ms) + } + return s, rows.Err() +} diff --git a/src/go/recorder.go b/src/go/recorder.go new file mode 100644 index 0000000..b2f7df5 --- /dev/null +++ b/src/go/recorder.go @@ -0,0 +1,169 @@ +package main + +import ( + "database/sql" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sync" + "time" + + _ "modernc.org/sqlite" +) + +// Recorder 负责全量留存：原文落盘 + SQLite 索引。 +type Recorder struct { + dataDir string + db *sql.DB + mu sync.Mutex // 串行化写库，sqlite 单写者 +} + +// blob 落盘的原文结构：请求 + 响应完整内容。 +type blob struct { + ID int64 `json:"id"` + Timestamp string `json:"timestamp"` + ClientIP string `json:"client_ip"` + Method string `json:"method"` + Path string `json:"path"` + ReqHeader map[string]string `json:"req_header"` + ReqBody string `json:"req_body"` + Status int `json:"status"` + RespHeader map[string]string `json:"resp_header"` + RespBody string `json:"resp_body"` + DurationMs int64 `json:"duration_ms"` +} + +// NewRecorder 创建数据目录、打开 sqlite、建表。 +func NewRecorder(dir string) (*Recorder, error) { + if err := os.MkdirAll(dir, 0o700); err != nil { + return nil, fmt.Errorf("创建数据目录: %w", err) + } + dbPath := filepath.Join(dir, "usage.db") + db, err := sql.Open("sqlite", dbPath) + if err != nil { + return nil, fmt.Errorf("打开 sqlite: %w", err) + } + // WAL 提升并发读写表现，busy_timeout 避免写锁瞬时冲突直接报错 + for _, pragma := range []string{ + "PRAGMA journal_mode=WAL;", + "PRAGMA busy_timeout=5000;", + } { + if _, err := db.Exec(pragma); err != nil { + db.Close() + return nil, fmt.Errorf("设置 pragma %q: %w", pragma, err) + } + } + + schema := ` +CREATE TABLE IF NOT EXISTS traffic ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + ts INTEGER NOT NULL, + client_ip TEXT, + api_format TEXT, + endpoint TEXT, + model TEXT, + input_tokens INTEGER, + output_tokens INTEGER, + total_tokens INTEGER, + is_stream INTEGER, + duration_ms INTEGER, + status INTEGER, + blob_path TEXT +); +CREATE INDEX IF NOT EXISTS idx_traffic_ts ON traffic(ts); +CREATE INDEX IF NOT EXISTS idx_traffic_model ON traffic(model);` + if _, err := db.Exec(schema); err != nil { + db.Close() + return nil, fmt.Errorf("建表: %w", err) + } + + return &Recorder{dataDir: dir, db: db}, nil +} + +// Save 把一条记录写入 SQLite 索引并将原文落盘。 +// +// 顺序很关键：先 INSERT（不带 id，由 SQLite AUTOINCREMENT 分配）， +// 拿回 LastInsertId() 作为权威主键，再用它命名 blob 文件并回填 e.ID。 +// 这样 id 由 DB 全局唯一递增，跨进程重启不再冲突；blob 文件名也随之唯一。 +func (r *Recorder) Save(e *Entry) error { + // 1) 先写 SQLite 索引，blob_path 暂留空，拿到自增 id 后再回填。 + r.mu.Lock() + res, err := r.db.Exec( + `INSERT INTO traffic + (ts, client_ip, api_format, endpoint, model, + input_tokens, output_tokens, total_tokens, is_stream, + duration_ms, status, blob_path) + VALUES (?,?,?,?,?,?,?,?,?,?,?,?)`, + e.Timestamp.UnixMilli(), e.ClientIP, e.APIFormat, e.Path, e.Model, + e.InputTokens, e.OutputTokens, e.TotalTokens, boolToInt(e.IsStream), + e.DurationMs, e.Status, "", + ) + if err != nil { + r.mu.Unlock() + return fmt.Errorf("写索引: %w", err) + } + id, err := res.LastInsertId() + if err != nil { + r.mu.Unlock() + return fmt.Errorf("取自增 id: %w", err) + } + e.ID = id // 回填权威主键（供调用方日志/后续引用） + + // 2) 原文落盘：按 日期/ID.json 组织，用 DB 分配的 id 命名。 + day := e.Timestamp.Format("2006-01-02") + blobDir := filepath.Join(r.dataDir, "blobs", day) + if err := os.MkdirAll(blobDir, 0o700); err != nil { + r.mu.Unlock() + return fmt.Errorf("创建 blob 目录: %w", err) + } + rel := filepath.Join("blobs", day, fmt.Sprintf("%d.json", id)) + abs := filepath.Join(r.dataDir, rel) + + b := blob{ + ID: id, + Timestamp: e.Timestamp.Format(time.RFC3339Nano), + ClientIP: e.ClientIP, + Method: e.Method, + Path: e.Path, + ReqHeader: e.ReqHeader, + ReqBody: string(e.ReqBody), + Status: e.Status, + RespHeader: e.RespHeader, + RespBody: string(e.RespBody), + DurationMs: e.DurationMs, + } + data, err := json.MarshalIndent(b, "", " ") + if err != nil { + r.mu.Unlock() + return fmt.Errorf("序列化原文: %w", err) + } + if err := os.WriteFile(abs, data, 0o600); err != nil { + r.mu.Unlock() + return fmt.Errorf("写原文: %w", err) + } + e.BlobPath = rel + + // 3) 回填 blob_path 到刚插入的行。 + _, err = r.db.Exec("UPDATE traffic SET blob_path = ? WHERE id = ?", rel, id) + r.mu.Unlock() + if err != nil { + return fmt.Errorf("回填 blob_path: %w", err) + } + return nil +} + +// Close 关闭数据库。 +func (r *Recorder) Close() error { + if r.db != nil { + return r.db.Close() + } + return nil +} + +func boolToInt(b bool) int { + if b { + return 1 + } + return 0 +} diff --git a/src/go/sse.go b/src/go/sse.go new file mode 100644 index 0000000..e221177 --- /dev/null +++ b/src/go/sse.go @@ -0,0 +1,185 @@ +package main + +import ( + "bytes" + "encoding/json" + "io" + "strings" +) + +// streamRecorder 包在 resp.Body 外层：数据被客户端读取时同步流过， +// 一边原样转发（客户端实时收到每个 chunk），一边把全量原文拼回 + 解析 SSE usage。 +// 读到 EOF（流结束）时回调 onDone，由它落库。 +// +// 这是流式代理的标准做法：用一个 io.ReadCloser 包装上游 body， +// 不做 io.ReadAll，避免缓冲整条流导致客户端要等所有 chunk 到齐才收到第一个字。 +type streamRecorder struct { + src io.ReadCloser + raw bytes.Buffer // 全量原文（拼回所有 chunk） + parser *sseParser + onDone func(raw []byte, in, out int, model string) + done bool +} + +func newStreamRecorder(src io.ReadCloser, format string, onDone func(raw []byte, in, out int, model string)) *streamRecorder { + return &streamRecorder{ + src: src, + parser: newSSEParser(format), + onDone: onDone, + } +} + +func (s *streamRecorder) Read(p []byte) (int, error) { + n, err := s.src.Read(p) + if n > 0 { + // 原样留一份，并喂给 SSE 解析器累积 usage + s.raw.Write(p[:n]) + s.parser.feed(p[:n]) + } + if err == io.EOF && !s.done { + s.done = true + s.parser.flush() + in, out, model := s.parser.result() + s.onDone(s.raw.Bytes(), in, out, model) + } + return n, err +} + +func (s *streamRecorder) Close() error { + // 客户端提前断开等情况：若还没触发 onDone，这里兜底落库一次。 + if !s.done { + s.done = true + s.parser.flush() + in, out, model := s.parser.result() + s.onDone(s.raw.Bytes(), in, out, model) + } + return s.src.Close() +} + +// sseParser 按行累积 SSE 数据，跨 chunk 解析 OpenAI / Claude 两家流式 usage。 +// +// OpenAI 流式：usage 出现在末尾带 stream_options.include_usage 的 chunk 里， +// 形如 data: {...,"usage":{"prompt_tokens":..,"completion_tokens":..}}； +// 没开 include_usage 则 usage 落空（如实记 0）。 +// Claude 流式：message_start 事件给 input_tokens（usage.input_tokens）， +// message_delta 事件累加 output_tokens（usage.output_tokens），需跨事件攒。 +type sseParser struct { + format string + buf bytes.Buffer // 未成行的残留字节 + + in int + out int + model string +} + +func newSSEParser(format string) *sseParser { + return &sseParser{format: format} +} + +// feed 接收任意切片的字节，按行切出完整的 SSE 行交给 parseLine。 +func (p *sseParser) feed(b []byte) { + p.buf.Write(b) + data := p.buf.Bytes() + for { + i := bytes.IndexByte(data, '\n') + if i < 0 { + break + } + line := data[:i] + p.parseLine(line) + data = data[i+1:] + } + // 保留未成行的尾部 + rest := make([]byte, len(data)) + copy(rest, data) + p.buf.Reset() + p.buf.Write(rest) +} + +// flush 处理流结束时缓冲里残留的最后一行（可能没有结尾换行）。 +func (p *sseParser) flush() { + if p.buf.Len() > 0 { + p.parseLine(p.buf.Bytes()) + p.buf.Reset() + } +} + +// parseLine 解析单行 SSE。只关心 "data:" 行里的 JSON。 +func (p *sseParser) parseLine(line []byte) { + s := strings.TrimRight(string(line), "\r") + s = strings.TrimSpace(s) + if !strings.HasPrefix(s, "data:") { + return + } + payload := strings.TrimSpace(strings.TrimPrefix(s, "data:")) + if payload == "" || payload == "[DONE]" { + return + } + + var m map[string]json.RawMessage + if err := json.Unmarshal([]byte(payload), &m); err != nil { + return // 非 JSON 行，跳过 + } + + // model 字段（两家都有，取到就记） + if raw, ok := m["model"]; ok { + var mod string + if json.Unmarshal(raw, &mod) == nil && mod != "" { + p.model = mod + } + } + + // usage 可能直接在顶层（OpenAI 末尾 chunk）， + // 也可能在 message.usage 里（Claude message_start）。 + if raw, ok := m["usage"]; ok { + p.applyUsage(raw) + } + if raw, ok := m["message"]; ok { + var msg map[string]json.RawMessage + if json.Unmarshal(raw, &msg) == nil { + if u, ok := msg["usage"]; ok { + p.applyUsage(u) + } + if mraw, ok := msg["model"]; ok { + var mod string + if json.Unmarshal(mraw, &mod) == nil && mod != "" { + p.model = mod + } + } + } + } +} + +// applyUsage 从一段 usage JSON 里抽 token 数，兼容两家命名。 +// Claude 的 output_tokens 在 message_delta 里是累计值（取最大即可）， +// input_tokens 在 message_start 给定后不再变；OpenAI 末尾一次性给全量。 +func (p *sseParser) applyUsage(raw json.RawMessage) { + var u struct { + // OpenAI + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + // Claude + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + } + if json.Unmarshal(raw, &u) != nil { + return + } + // 取较大值，避免后到的事件覆盖成 0（Claude 各事件字段不全）。 + if u.PromptTokens > p.in { + p.in = u.PromptTokens + } + if u.InputTokens > p.in { + p.in = u.InputTokens + } + if u.CompletionTokens > p.out { + p.out = u.CompletionTokens + } + if u.OutputTokens > p.out { + p.out = u.OutputTokens + } +} + +func (p *sseParser) result() (in, out int, model string) { + return p.in, p.out, p.model +} diff --git a/src/go/sse_test.go b/src/go/sse_test.go new file mode 100644 index 0000000..0302eb6 --- /dev/null +++ b/src/go/sse_test.go @@ -0,0 +1,103 @@ +package main + +import ( + "io" + "strings" + "testing" +) + +// feedInChunks 把一段完整 SSE 文本按指定大小切片喂给 streamRecorder， +// 模拟真实流式逐 chunk 到达（且切片边界可能落在一行中间）， +// 验证跨 chunk 的行重组与 usage 累积都正确。 +func drainInChunks(t *testing.T, full string, format string, chunk int) (raw []byte, in, out int, model string) { + t.Helper() + src := io.NopCloser(strings.NewReader(full)) + var gotIn, gotOut int + var gotModel string + var gotRaw []byte + sr := newStreamRecorder(src, format, func(b []byte, i, o int, m string) { + gotRaw = append([]byte(nil), b...) + gotIn, gotOut, gotModel = i, o, m + }) + // 像客户端那样按小 buffer 反复读，直到 EOF + buf := make([]byte, chunk) + for { + _, err := sr.Read(buf) + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("read: %v", err) + } + } + return gotRaw, gotIn, gotOut, gotModel +} + +// TestSSEOpenAI 验证 OpenAI 流式：usage 在末尾带 include_usage 的 chunk 里。 +func TestSSEOpenAI(t *testing.T) { + full := strings.Join([]string{ + `data: {"model":"gpt-5.5","choices":[{"delta":{"content":"O"}}]}`, + `data: {"model":"gpt-5.5","choices":[{"delta":{"content":"K"}}]}`, + `data: {"model":"gpt-5.5","choices":[],"usage":{"prompt_tokens":42,"completion_tokens":3,"total_tokens":45}}`, + `data: [DONE]`, + "", + }, "\n") + + // 用一个会落在行中间的小 chunk，逼真地考验跨 chunk 重组 + raw, in, out, model := drainInChunks(t, full, "openai", 7) + if in != 42 || out != 3 { + t.Errorf("openai usage = in:%d out:%d, want 42/3", in, out) + } + if model != "gpt-5.5" { + t.Errorf("model = %q, want gpt-5.5", model) + } + if string(raw) != full { + t.Errorf("原文未完整拼回:\n want %q\n got %q", full, raw) + } +} + +// TestSSEClaude 验证 Claude 流式：input_tokens 在 message_start， +// output_tokens 在 message_delta 累计（取最大值）。 +func TestSSEClaude(t *testing.T) { + full := strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"model":"claude-3-5-sonnet","usage":{"input_tokens":58,"output_tokens":1}}}`, + `event: content_block_delta`, + `data: {"type":"content_block_delta","delta":{"text":"hi"}}`, + `event: message_delta`, + `data: {"type":"message_delta","usage":{"output_tokens":27}}`, + `event: message_stop`, + `data: {"type":"message_stop"}`, + "", + }, "\n") + + raw, in, out, model := drainInChunks(t, full, "claude", 13) + if in != 58 { + t.Errorf("claude input_tokens = %d, want 58", in) + } + if out != 27 { + t.Errorf("claude output_tokens = %d, want 27 (message_delta 累计)", out) + } + if model != "claude-3-5-sonnet" { + t.Errorf("model = %q, want claude-3-5-sonnet", model) + } + if string(raw) != full { + t.Errorf("原文未完整拼回") + } +} + +// TestSSENoUsage 验证未开 include_usage 时 usage 落空但不报错、原文仍完整。 +func TestSSENoUsage(t *testing.T) { + full := strings.Join([]string{ + `data: {"model":"gpt-5.5","choices":[{"delta":{"content":"hi"}}]}`, + `data: [DONE]`, + "", + }, "\n") + raw, in, out, _ := drainInChunks(t, full, "openai", 1) + if in != 0 || out != 0 { + t.Errorf("无 usage 时应为 0/0, got %d/%d", in, out) + } + if string(raw) != full { + t.Errorf("原文未完整拼回") + } +} diff --git a/src/go/usage.go b/src/go/usage.go new file mode 100644 index 0000000..2fc3e01 --- /dev/null +++ b/src/go/usage.go @@ -0,0 +1,84 @@ +package main + +import ( + "encoding/json" + "strings" +) + +// parseUsage 从非流式响应体里解析用量信息，结果写回 entry。 +// 兼容 OpenAI、Claude 两种字段命名；识别不出就标 unknown，不报错。 +func parseUsage(entry *Entry, body []byte) { + if len(body) == 0 { + entry.APIFormat = "unknown" + return + } + + var m map[string]json.RawMessage + if err := json.Unmarshal(body, &m); err != nil { + entry.APIFormat = "unknown" + return + } + + // model 字段（OpenAI 和 Claude 都叫 model） + if raw, ok := m["model"]; ok { + var s string + if json.Unmarshal(raw, &s) == nil { + entry.Model = s + } + } + + usageRaw, ok := m["usage"] + if !ok { + entry.APIFormat = "unknown" + return + } + + // usage 字段名两家不同，都尝试一遍 + var u struct { + // OpenAI + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + // Claude + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + } + if err := json.Unmarshal(usageRaw, &u); err != nil { + entry.APIFormat = "unknown" + return + } + + switch { + case u.PromptTokens > 0 || u.CompletionTokens > 0 || u.TotalTokens > 0: + // OpenAI 风格 + entry.APIFormat = "openai" + entry.InputTokens = u.PromptTokens + entry.OutputTokens = u.CompletionTokens + entry.TotalTokens = u.TotalTokens + if entry.TotalTokens == 0 { + entry.TotalTokens = u.PromptTokens + u.CompletionTokens + } + case u.InputTokens > 0 || u.OutputTokens > 0: + // Claude 风格 + entry.APIFormat = "claude" + entry.InputTokens = u.InputTokens + entry.OutputTokens = u.OutputTokens + entry.TotalTokens = u.InputTokens + u.OutputTokens + default: + entry.APIFormat = "unknown" + } +} + +// guessFormatByPath 在拿不到 usage 时，按路径粗略猜测 API 格式（仅作辅助标注）。 +func guessFormatByPath(path string) string { + switch { + case strings.Contains(path, "/v1/messages"): + return "claude" + case strings.Contains(path, "/v1/chat/completions"), strings.Contains(path, "/v1/completions"): + return "openai" + case strings.Contains(path, "/api/chat"), strings.Contains(path, "/api/generate"): + return "ollama" + default: + return "unknown" + } +} diff --git a/src/go/web.go b/src/go/web.go new file mode 100644 index 0000000..d32e0ad --- /dev/null +++ b/src/go/web.go @@ -0,0 +1,103 @@ +package main + +import ( + "encoding/json" + "log" + "net/http" + "strconv" +) + +// startWebUI 在独立端口起一个只读管理界面，对外提供 dashboard + JSON API。 +// 单独端口避免和反代端口冲突（反代会把所有路径转发上游）。 +func startWebUI(addr string, rec *Recorder) { + mux := webMux(rec) + log.Printf("Web 管理界面启动: http://%s", addr) + if err := http.ListenAndServe(addr, mux); err != nil { + log.Printf("[!] Web 管理界面退出: %v", err) + } +} + +// webMux 组装管理界面的路由，单独抽出便于测试。 +func webMux(rec *Recorder) http.Handler { + mux := http.NewServeMux() + + // 首页：dashboard HTML + mux.HandleFunc("GET /", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/" { + http.NotFound(w, r) + return + } + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.Write([]byte(dashboardHTML)) + }) + + // 聚合统计 + mux.HandleFunc("GET /api/stats", func(w http.ResponseWriter, r *http.Request) { + s, err := rec.Stats() + if err != nil { + writeErr(w, err) + return + } + writeJSON(w, s) + }) + + // 记录列表（支持 model/format/status/limit/offset 过滤） + mux.HandleFunc("GET /api/records", func(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + opts := ListOptions{ + Model: q.Get("model"), + Format: q.Get("format"), + Limit: atoiDefault(q.Get("limit"), 100), + Offset: atoiDefault(q.Get("offset"), 0), + Status: atoiDefault(q.Get("status"), 0), + } + recs, total, err := rec.List(opts) + if err != nil { + writeErr(w, err) + return + } + if recs == nil { + recs = []Record{} + } + writeJSON(w, map[string]any{"total": total, "records": recs}) + }) + + // 单条记录原文详情（含完整请求/响应） + mux.HandleFunc("GET /api/records/{id}", func(w http.ResponseWriter, r *http.Request) { + id, err := strconv.ParseInt(r.PathValue("id"), 10, 64) + if err != nil { + http.Error(w, "无效 id", http.StatusBadRequest) + return + } + b, err := rec.ReadBlob(id) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + writeJSON(w, b) + }) + + return mux +} + +func writeJSON(w http.ResponseWriter, v any) { + w.Header().Set("Content-Type", "application/json; charset=utf-8") + if err := json.NewEncoder(w).Encode(v); err != nil { + log.Printf("[!] 写 JSON 响应失败: %v", err) + } +} + +func writeErr(w http.ResponseWriter, err error) { + http.Error(w, err.Error(), http.StatusInternalServerError) +} + +func atoiDefault(s string, def int) int { + if s == "" { + return def + } + n, err := strconv.Atoi(s) + if err != nil { + return def + } + return n +} diff --git a/src/go/web_test.go b/src/go/web_test.go new file mode 100644 index 0000000..84c7ad3 --- /dev/null +++ b/src/go/web_test.go @@ -0,0 +1,188 @@ +package main + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// seedRecords 往 recorder 里塞几条记录，供 web 接口测试使用。 +func seedRecords(t *testing.T, rec *Recorder) { + t.Helper() + now := time.Now() + seeds := []*Entry{ + {ID: 1, Timestamp: now, ClientIP: "127.0.0.1", Method: "POST", Path: "/v1/chat/completions", + ReqHeader: map[string]string{"Authorization": "Bearer sk-aaa"}, ReqBody: []byte(`{"model":"gpt-5.5"}`), + Status: 200, RespHeader: map[string]string{"Content-Type": "application/json"}, RespBody: []byte(`{"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}`), + APIFormat: "openai", Model: "gpt-5.5", InputTokens: 10, OutputTokens: 5, TotalTokens: 15}, + {ID: 2, Timestamp: now, ClientIP: "127.0.0.1", Method: "POST", Path: "/v1/messages", + ReqHeader: map[string]string{}, ReqBody: []byte(`{"model":"claude-3-5-sonnet"}`), + Status: 200, RespHeader: map[string]string{}, RespBody: []byte(`{}`), + APIFormat: "claude", Model: "claude-3-5-sonnet", InputTokens: 100, OutputTokens: 200, TotalTokens: 300}, + {ID: 3, Timestamp: now, ClientIP: "10.0.0.2", Method: "POST", Path: "/v1/chat/completions", + ReqHeader: map[string]string{}, ReqBody: []byte(`{}`), + Status: 503, RespHeader: map[string]string{}, RespBody: []byte(``), + APIFormat: "unknown", Model: "", InputTokens: 0, OutputTokens: 0, TotalTokens: 0}, + } + for _, e := range seeds { + if err := rec.Save(e); err != nil { + t.Fatalf("seed save #%d: %v", e.ID, err) + } + } +} + +// TestWebStats 验证 /api/stats 聚合正确。 +func TestWebStats(t *testing.T) { + rec, err := NewRecorder(t.TempDir()) + if err != nil { + t.Fatalf("NewRecorder: %v", err) + } + defer rec.Close() + seedRecords(t, rec) + + srv := httptest.NewServer(webMux(rec)) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/api/stats") + if err != nil { + t.Fatalf("get stats: %v", err) + } + defer resp.Body.Close() + var s Stats + if err := json.NewDecoder(resp.Body).Decode(&s); err != nil { + t.Fatalf("decode stats: %v", err) + } + if s.TotalRequests != 3 { + t.Errorf("total_requests = %d, want 3", s.TotalRequests) + } + if s.SuccessCount != 2 { + t.Errorf("success_count = %d, want 2", s.SuccessCount) + } + if s.ErrorCount != 1 { + t.Errorf("error_count = %d, want 1", s.ErrorCount) + } + if s.TotalTokens != 315 { + t.Errorf("total_tokens = %d, want 315", s.TotalTokens) + } + if len(s.ByModel) != 3 { + t.Errorf("by_model 种类 = %d, want 3", len(s.ByModel)) + } +} + +// TestWebRecordsFilter 验证 /api/records 的过滤与分页。 +func TestWebRecordsFilter(t *testing.T) { + rec, err := NewRecorder(t.TempDir()) + if err != nil { + t.Fatalf("NewRecorder: %v", err) + } + defer rec.Close() + seedRecords(t, rec) + + srv := httptest.NewServer(webMux(rec)) + defer srv.Close() + + // 无过滤：3 条 + if total := recordsTotal(t, srv.URL+"/api/records"); total != 3 { + t.Errorf("无过滤 total = %d, want 3", total) + } + // 按格式过滤：openai 只有 1 条 + if total := recordsTotal(t, srv.URL+"/api/records?format=openai"); total != 1 { + t.Errorf("format=openai total = %d, want 1", total) + } + // 按状态过滤：503 只有 1 条 + if total := recordsTotal(t, srv.URL+"/api/records?status=503"); total != 1 { + t.Errorf("status=503 total = %d, want 1", total) + } + // 按模型过滤 + if total := recordsTotal(t, srv.URL+"/api/records?model=gpt-5.5"); total != 1 { + t.Errorf("model=gpt-5.5 total = %d, want 1", total) + } +} + +func recordsTotal(t *testing.T, url string) int { + t.Helper() + resp, err := http.Get(url) + if err != nil { + t.Fatalf("get %s: %v", url, err) + } + defer resp.Body.Close() + var out struct { + Total int `json:"total"` + Records []Record `json:"records"` + } + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + t.Fatalf("decode records: %v", err) + } + return out.Total +} + +// TestWebRecordDetail 验证 /api/records/{id} 返回完整原文（含留存的请求头）。 +func TestWebRecordDetail(t *testing.T) { + rec, err := NewRecorder(t.TempDir()) + if err != nil { + t.Fatalf("NewRecorder: %v", err) + } + defer rec.Close() + seedRecords(t, rec) + + srv := httptest.NewServer(webMux(rec)) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/api/records/1") + if err != nil { + t.Fatalf("get detail: %v", err) + } + defer resp.Body.Close() + var b blob + if err := json.NewDecoder(resp.Body).Decode(&b); err != nil { + t.Fatalf("decode blob: %v", err) + } + if b.ID != 1 { + t.Errorf("id = %d, want 1", b.ID) + } + if !strings.Contains(b.ReqBody, "gpt-5.5") { + t.Errorf("req_body 缺内容: %q", b.ReqBody) + } + // 全量留存：请求头里的凭证原样可查 + if b.ReqHeader["Authorization"] != "Bearer sk-aaa" { + t.Errorf("留存请求头缺 Authorization: %q", b.ReqHeader["Authorization"]) + } + + // 不存在的 id 应 404 + resp2, _ := http.Get(srv.URL + "/api/records/999") + if resp2 != nil { + resp2.Body.Close() + if resp2.StatusCode != http.StatusNotFound { + t.Errorf("不存在 id 状态 = %d, want 404", resp2.StatusCode) + } + } +} + +// TestWebDashboardHTML 验证首页返回 HTML。 +func TestWebDashboardHTML(t *testing.T) { + rec, err := NewRecorder(t.TempDir()) + if err != nil { + t.Fatalf("NewRecorder: %v", err) + } + defer rec.Close() + + srv := httptest.NewServer(webMux(rec)) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/") + if err != nil { + t.Fatalf("get /: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "text/html") { + t.Errorf("Content-Type = %q, want text/html", ct) + } + if !strings.Contains(string(body), "AI 通讯留存") { + t.Errorf("首页缺标题") + } +} diff --git a/src/go/webui.go b/src/go/webui.go new file mode 100644 index 0000000..e7da9c1 --- /dev/null +++ b/src/go/webui.go @@ -0,0 +1,231 @@ +package main + +// dashboardHTML 是只读管理界面的单页前端（原生 JS，无外部依赖）。 +// 通过 /api/stats、/api/records、/api/records/{id} 拉数据。 +const dashboardHTML = ` + + + + +AI 通讯留存 - 管理界面 + + + +
+

AI 通讯留存

+ 全量审计 · 只读管理界面 + + +
+
+
+ +
+ + + + + + +
+ + + + + + + + +
#时间客户端格式模型路径状态输入输出合计流式耗时
+ +
+ + + +
+
+ +
+
+ +
+
+ + + +` diff --git a/说明.md b/说明.md deleted file mode 100644 index e69de29..0000000