From 22136a55fc7e95325caa4d4ccd65ce6dd3235a87 Mon Sep 17 00:00:00 2001 From: Bohan Jiang <52446949+Bohan-J@users.noreply.github.com> Date: Wed, 29 Apr 2026 15:00:00 +0800 Subject: [PATCH] fix(server/heartbeat): split auth_ms into decode/runtime_lookup/workspace_check + auth_path (#1822) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Prod slow-log on the deployed v0.2.17 fix shows total_ms=4012, auth_ms=4010, update_ms=1, all skill stages = 0 — meaning the bottleneck on /api/daemon/heartbeat is now the auth section, not the Redis claim path. To pinpoint which sub-stage dominates, decompose auth_ms into: - decode_ms — JSON body decode - runtime_lookup_ms — Queries.GetAgentRuntime (PG PK select) - workspace_check_ms — requireDaemonWorkspaceAccess (string compare for daemon-token, requireWorkspaceMember for PAT/JWT) Also add auth_path ("daemon_token" | "pat" | "jwt") set by DaemonAuth middleware so slow-logs disambiguate which token kind was used. PAT/JWT takes an extra DB round-trip via requireWorkspaceMember and is a candidate cause of long auth tails on daemons that haven't migrated to mdt_ tokens. The handler keeps the same external behavior; the change inlines and instruments requireDaemonRuntimeAccess in DaemonHeartbeat only — other callers of the helper are untouched. logHeartbeatEndpointSlow gains the new fields. Existing heartbeat tests pass; the slow-probe test output now shows the new auth_path / decode_ms / runtime_lookup_ms / workspace_check_ms fields populated. --- server/internal/handler/daemon.go | 48 +++++++++++++++++++---- server/internal/middleware/daemon_auth.go | 24 +++++++++++- 2 files changed, 62 insertions(+), 10 deletions(-) diff --git a/server/internal/handler/daemon.go b/server/internal/handler/daemon.go index 1054e06a76..cca9ec3600 100644 --- a/server/internal/handler/daemon.go +++ b/server/internal/handler/daemon.go @@ -517,18 +517,23 @@ const heartbeatHasPendingTimeout = 1 * time.Second func (h *Handler) DaemonHeartbeat(w http.ResponseWriter, r *http.Request) { start := time.Now() + authPath := middleware.DaemonAuthPathFromContext(r.Context()) var ( outcome = "unauth" runtimeID string + decodeMs, runtimeLookupMs, workspaceCheckMs int64 authMs, updateMs, probeSkillsMs, popSkillsMs, probeImportMs, popImportMs int64 probeSkillsTimedOut, probeImportTimedOut bool ) defer func() { - logHeartbeatEndpointSlow(runtimeID, outcome, start, authMs, updateMs, probeSkillsMs, popSkillsMs, probeImportMs, popImportMs, probeSkillsTimedOut, probeImportTimedOut) + logHeartbeatEndpointSlow(runtimeID, outcome, authPath, start, decodeMs, runtimeLookupMs, workspaceCheckMs, authMs, updateMs, probeSkillsMs, popSkillsMs, probeImportMs, popImportMs, probeSkillsTimedOut, probeImportTimedOut) }() + decodeStart := time.Now() var req DaemonHeartbeatRequest - if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + decodeErr := json.NewDecoder(r.Body).Decode(&req) + decodeMs = time.Since(decodeStart).Milliseconds() + if decodeErr != nil { outcome = "bad_body" writeError(w, http.StatusBadRequest, "invalid request body") return @@ -541,17 +546,38 @@ func (h *Handler) DaemonHeartbeat(w http.ResponseWriter, r *http.Request) { } runtimeID = req.RuntimeID - // Verify the caller owns this runtime's workspace. - rt, ok := h.requireDaemonRuntimeAccess(w, r, req.RuntimeID) + // Inlined and instrumented version of requireDaemonRuntimeAccess so we + // can attribute the runtime-lookup and workspace-check sub-stages + // independently in slow-logs. Together with the auth_path label set by + // DaemonAuth middleware, this lets us tell whether prod heartbeat tail + // latency is in pgx pool acquisition (runtime_lookup_ms), in the PAT + // fallback workspace-membership query (workspace_check_ms), or upstream. + runtimeUUID, ok := parseUUIDOrBadRequest(w, req.RuntimeID, "runtime_id") if !ok { + outcome = "bad_runtime_id" + return + } + lookupStart := time.Now() + rt, lookupErr := h.Queries.GetAgentRuntime(r.Context(), runtimeUUID) + runtimeLookupMs = time.Since(lookupStart).Milliseconds() + if lookupErr != nil { + outcome = "runtime_not_found" + writeError(w, http.StatusNotFound, "runtime not found") + return + } + wsCheckStart := time.Now() + wsOK := h.requireDaemonWorkspaceAccess(w, r, uuidToString(rt.WorkspaceID)) + workspaceCheckMs = time.Since(wsCheckStart).Milliseconds() + if !wsOK { + outcome = "workspace_denied" return } authMs = time.Since(start).Milliseconds() updateStart := time.Now() - _, err := h.Queries.UpdateAgentRuntimeHeartbeat(r.Context(), rt.ID) + _, updateErr := h.Queries.UpdateAgentRuntimeHeartbeat(r.Context(), rt.ID) updateMs = time.Since(updateStart).Milliseconds() - if err != nil { + if updateErr != nil { outcome = "error_update" writeError(w, http.StatusInternalServerError, "heartbeat failed") return @@ -637,8 +663,10 @@ func (h *Handler) DaemonHeartbeat(w http.ResponseWriter, r *http.Request) { // logHeartbeatEndpointSlow emits one structured log when /api/daemon/heartbeat // exceeds 500ms, splitting auth / update / probe / pop phases for both queues // so the prod tail can be attributed without flooding logs at normal rates. -// Mirrors logClaimEndpointSlow for consistency. -func logHeartbeatEndpointSlow(runtimeID, outcome string, start time.Time, authMs, updateMs, probeSkillsMs, popSkillsMs, probeImportMs, popImportMs int64, probeSkillsTimedOut, probeImportTimedOut bool) { +// auth_ms is further decomposed into decode_ms, runtime_lookup_ms, and +// workspace_check_ms; auth_path labels which token kind authenticated the +// request ("daemon_token", "pat", or "jwt"). Mirrors logClaimEndpointSlow. +func logHeartbeatEndpointSlow(runtimeID, outcome, authPath string, start time.Time, decodeMs, runtimeLookupMs, workspaceCheckMs, authMs, updateMs, probeSkillsMs, popSkillsMs, probeImportMs, popImportMs int64, probeSkillsTimedOut, probeImportTimedOut bool) { totalMs := time.Since(start).Milliseconds() if totalMs < 500 && !probeSkillsTimedOut && !probeImportTimedOut { return @@ -646,8 +674,12 @@ func logHeartbeatEndpointSlow(runtimeID, outcome string, start time.Time, authMs slog.Info("heartbeat_endpoint slow", "runtime_id", runtimeID, "outcome", outcome, + "auth_path", authPath, "total_ms", totalMs, "auth_ms", authMs, + "decode_ms", decodeMs, + "runtime_lookup_ms", runtimeLookupMs, + "workspace_check_ms", workspaceCheckMs, "update_ms", updateMs, "probe_skills_ms", probeSkillsMs, "pop_skills_ms", popSkillsMs, diff --git a/server/internal/middleware/daemon_auth.go b/server/internal/middleware/daemon_auth.go index a534a20f8d..b53c4acaca 100644 --- a/server/internal/middleware/daemon_auth.go +++ b/server/internal/middleware/daemon_auth.go @@ -17,6 +17,14 @@ type daemonContextKey int const ( ctxKeyDaemonWorkspaceID daemonContextKey = iota ctxKeyDaemonID + ctxKeyDaemonAuthPath +) + +// Daemon auth path labels exposed via context for slow-log attribution. +const ( + DaemonAuthPathDaemonToken = "daemon_token" + DaemonAuthPathPAT = "pat" + DaemonAuthPathJWT = "jwt" ) // DaemonWorkspaceIDFromContext returns the workspace ID set by DaemonAuth middleware. @@ -31,11 +39,20 @@ func DaemonIDFromContext(ctx context.Context) string { return id } +// DaemonAuthPathFromContext returns which token kind authenticated this +// request — "daemon_token", "pat", or "jwt" — for telemetry. Empty when the +// request did not pass through DaemonAuth. +func DaemonAuthPathFromContext(ctx context.Context) string { + p, _ := ctx.Value(ctxKeyDaemonAuthPath).(string) + return p +} + // WithDaemonContext returns a new context with the daemon workspace ID and daemon ID set. // This is used by tests to simulate daemon token authentication. func WithDaemonContext(ctx context.Context, workspaceID, daemonID string) context.Context { ctx = context.WithValue(ctx, ctxKeyDaemonWorkspaceID, workspaceID) ctx = context.WithValue(ctx, ctxKeyDaemonID, daemonID) + ctx = context.WithValue(ctx, ctxKeyDaemonAuthPath, DaemonAuthPathDaemonToken) return ctx } @@ -71,6 +88,7 @@ func DaemonAuth(queries *db.Queries) func(http.Handler) http.Handler { ctx := context.WithValue(r.Context(), ctxKeyDaemonWorkspaceID, uuidToString(dt.WorkspaceID)) ctx = context.WithValue(ctx, ctxKeyDaemonID, dt.DaemonID) + ctx = context.WithValue(ctx, ctxKeyDaemonAuthPath, DaemonAuthPathDaemonToken) next.ServeHTTP(w, r.WithContext(ctx)) return } @@ -86,7 +104,8 @@ func DaemonAuth(queries *db.Queries) func(http.Handler) http.Handler { } r.Header.Set("X-User-ID", uuidToString(pat.UserID)) go queries.UpdatePersonalAccessTokenLastUsed(context.Background(), pat.ID) - next.ServeHTTP(w, r) + ctx := context.WithValue(r.Context(), ctxKeyDaemonAuthPath, DaemonAuthPathPAT) + next.ServeHTTP(w, r.WithContext(ctx)) return } @@ -114,7 +133,8 @@ func DaemonAuth(queries *db.Queries) func(http.Handler) http.Handler { return } r.Header.Set("X-User-ID", sub) - next.ServeHTTP(w, r) + ctx := context.WithValue(r.Context(), ctxKeyDaemonAuthPath, DaemonAuthPathJWT) + next.ServeHTTP(w, r.WithContext(ctx)) }) } }