refactor: split 4 oversized handler files into focused sub-files

workspace-server/internal/handlers/a2a_proxy_helpers.go

@@ -0,0 +1,372 @@
+package handlers
+
+// a2a_proxy_helpers.go — A2A proxy error handling, activity logging,
+// caller auth validation, token usage tracking, and SSRF safety checks.
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
+	"github.com/Molecule-AI/molecule-monorepo/platform/internal/wsauth"
+	"github.com/gin-gonic/gin"
+)
+// proxyDispatchBuildError is a sentinel wrapper for failures inside
+// http.NewRequestWithContext. handleA2ADispatchError unwraps it to emit the
+// "failed to create proxy request" 500 instead of the standard 502/503 paths.
+type proxyDispatchBuildError struct{ err error }
+
+func (e *proxyDispatchBuildError) Error() string { return e.err.Error() }
+
+// handleA2ADispatchError translates a forward-call failure into a proxyA2AError,
+// runs the reactive container-health check, and (when `logActivity` is true)
+// schedules a detached LogActivity goroutine for the failed attempt.
+func (h *WorkspaceHandler) handleA2ADispatchError(ctx context.Context, workspaceID, callerID string, body []byte, a2aMethod string, err error, durationMs int, logActivity bool) (int, []byte, *proxyA2AError) {
+	// Build-time failure (couldn't even create the http.Request) — return
+	// a 500 without the reactive-health / busy-retry paths.
+	if buildErr, ok := err.(*proxyDispatchBuildError); ok {
+		_ = buildErr
+		return 0, nil, &proxyA2AError{
+			Status:   http.StatusInternalServerError,
+			Response: gin.H{"error": "failed to create proxy request"},
+		}
+	}
+
+	log.Printf("ProxyA2A forward error: %v", err)
+
+	containerDead := h.maybeMarkContainerDead(ctx, workspaceID)
+
+	if logActivity {
+		h.logA2AFailure(ctx, workspaceID, callerID, body, a2aMethod, err, durationMs)
+	}
+	if containerDead {
+		return 0, nil, &proxyA2AError{
+			Status:   http.StatusServiceUnavailable,
+			Response: gin.H{"error": "workspace agent unreachable — container restart triggered", "restarting": true},
+		}
+	}
+	// Container is alive but upstream Do() failed with a timeout/EOF-
+	// shaped error — the agent is most likely mid-synthesis on a
+	// previous request (single-threaded main loop). Surface as 503
+	// Busy with a Retry-After hint so callers can distinguish this
+	// from a real unreachable-agent (502) and retry with backoff.
+	// Issue #110.
+	if isUpstreamBusyError(err) {
+		return 0, nil, &proxyA2AError{
+			Status:  http.StatusServiceUnavailable,
+			Headers: map[string]string{"Retry-After": strconv.Itoa(busyRetryAfterSeconds)},
+			Response: gin.H{
+				"error":       "workspace agent busy — retry after a short backoff",
+				"busy":        true,
+				"retry_after": busyRetryAfterSeconds,
+			},
+		}
+	}
+	return 0, nil, &proxyA2AError{
+		Status:   http.StatusBadGateway,
+		Response: gin.H{"error": "failed to reach workspace agent"},
+	}
+}
+
+// maybeMarkContainerDead runs the reactive health check after a forward error.
+// If the workspace's Docker container is no longer running (and the workspace
+// isn't external), it marks the workspace offline, clears Redis state,
+// broadcasts WORKSPACE_OFFLINE, and triggers an async restart. Returns true
+// when the container was found dead.
+func (h *WorkspaceHandler) maybeMarkContainerDead(ctx context.Context, workspaceID string) bool {
+	var wsRuntime string
+	db.DB.QueryRowContext(ctx, `SELECT COALESCE(runtime, 'langgraph') FROM workspaces WHERE id = $1`, workspaceID).Scan(&wsRuntime)
+	if h.provisioner == nil || wsRuntime == "external" {
+		return false
+	}
+	running, inspectErr := h.provisioner.IsRunning(ctx, workspaceID)
+	if inspectErr != nil {
+		// Transient Docker-daemon error (timeout, socket EOF, etc.). Post-
+		// #386, IsRunning returns (true, err) in this case — caller stays
+		// on the alive path and does not trigger a restart cascade. Log
+		// so the defect is visible without being destructive.
+		log.Printf("ProxyA2A: IsRunning for %s returned transient error (assuming alive): %v", workspaceID, inspectErr)
+	}
+	if running {
+		return false
+	}
+	log.Printf("ProxyA2A: container for %s is dead — marking offline and triggering restart", workspaceID)
+	if _, err := db.DB.ExecContext(ctx, `UPDATE workspaces SET status = 'offline', updated_at = now() WHERE id = $1 AND status NOT IN ('removed', 'provisioning')`, workspaceID); err != nil {
+		log.Printf("ProxyA2A: failed to mark workspace %s offline: %v", workspaceID, err)
+	}
+	db.ClearWorkspaceKeys(ctx, workspaceID)
+	h.broadcaster.RecordAndBroadcast(ctx, "WORKSPACE_OFFLINE", workspaceID, map[string]interface{}{})
+	go h.RestartByID(workspaceID)
+	return true
+}
+
+// logA2AFailure records a failed A2A attempt to activity_logs in a detached
+// goroutine (the request context may already be done by the time it runs).
+func (h *WorkspaceHandler) logA2AFailure(ctx context.Context, workspaceID, callerID string, body []byte, a2aMethod string, err error, durationMs int) {
+	errMsg := err.Error()
+	var errWsName string
+	db.DB.QueryRowContext(ctx, `SELECT name FROM workspaces WHERE id = $1`, workspaceID).Scan(&errWsName)
+	if errWsName == "" {
+		errWsName = workspaceID
+	}
+	summary := "A2A request to " + errWsName + " failed: " + errMsg
+	go func(parent context.Context) {
+		logCtx, cancel := context.WithTimeout(context.WithoutCancel(parent), 30*time.Second)
+		defer cancel()
+		LogActivity(logCtx, h.broadcaster, ActivityParams{
+			WorkspaceID:  workspaceID,
+			ActivityType: "a2a_receive",
+			SourceID:     nilIfEmpty(callerID),
+			TargetID:     &workspaceID,
+			Method:       &a2aMethod,
+			Summary:      &summary,
+			RequestBody:  json.RawMessage(body),
+			DurationMs:   &durationMs,
+			Status:       "error",
+			ErrorDetail:  &errMsg,
+		})
+	}(ctx)
+}
+
+// logA2ASuccess records a successful A2A round-trip and (for canvas-initiated
+// 2xx/3xx responses) broadcasts an A2A_RESPONSE event so the frontend can
+// receive the reply without polling.
+func (h *WorkspaceHandler) logA2ASuccess(ctx context.Context, workspaceID, callerID string, body, respBody []byte, a2aMethod string, statusCode, durationMs int) {
+	logStatus := "ok"
+	if statusCode >= 400 {
+		logStatus = "error"
+	}
+	var wsNameForLog string
+	db.DB.QueryRowContext(ctx, `SELECT name FROM workspaces WHERE id = $1`, workspaceID).Scan(&wsNameForLog)
+	if wsNameForLog == "" {
+		wsNameForLog = workspaceID
+	}
+
+	// #817: track outbound activity on the CALLER so orchestrators can detect
+	// silent workspaces. Only update when callerID is a real workspace (not
+	// canvas, not a system caller) and the target returned 2xx/3xx.
+	if callerID != "" && !isSystemCaller(callerID) && statusCode < 400 {
+		go func() {
+			bgCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+			if _, err := db.DB.ExecContext(bgCtx,
+				`UPDATE workspaces SET last_outbound_at = NOW() WHERE id = $1`, callerID); err != nil {
+				log.Printf("last_outbound_at update failed for %s: %v", callerID, err)
+			}
+		}()
+	}
+	summary := a2aMethod + " → " + wsNameForLog
+	go func(parent context.Context) {
+		logCtx, cancel := context.WithTimeout(context.WithoutCancel(parent), 30*time.Second)
+		defer cancel()
+		LogActivity(logCtx, h.broadcaster, ActivityParams{
+			WorkspaceID:  workspaceID,
+			ActivityType: "a2a_receive",
+			SourceID:     nilIfEmpty(callerID),
+			TargetID:     &workspaceID,
+			Method:       &a2aMethod,
+			Summary:      &summary,
+			RequestBody:  json.RawMessage(body),
+			ResponseBody: json.RawMessage(respBody),
+			DurationMs:   &durationMs,
+			Status:       logStatus,
+		})
+	}(ctx)
+
+	if callerID == "" && statusCode < 400 {
+		h.broadcaster.BroadcastOnly(workspaceID, "A2A_RESPONSE", map[string]interface{}{
+			"response_body": json.RawMessage(respBody),
+			"method":        a2aMethod,
+			"duration_ms":   durationMs,
+		})
+	}
+}
+
+func nilIfEmpty(s string) *string {
+	if s == "" {
+		return nil
+	}
+	return &s
+}
+
+// validateCallerToken enforces the Phase 30.5 auth-token contract on the
+// caller of an A2A proxy request. Same lazy-bootstrap shape as
+// registry.requireWorkspaceToken: if the caller workspace has any live
+// token on file, the Authorization header is mandatory and must match;
+// if the caller has zero live tokens, they're grandfathered through
+// (their next /registry/register will mint their first token, after
+// which this branch never fires again for them).
+//
+// On auth failure this writes the 401 via c and returns an error so the
+// handler aborts without running the proxy.
+func validateCallerToken(ctx context.Context, c *gin.Context, callerID string) error {
+	hasLive, err := wsauth.HasAnyLiveToken(ctx, db.DB, callerID)
+	if err != nil {
+		// Fail-open here matches the heartbeat path — A2A caller auth is
+		// defense-in-depth on top of access-control hierarchy, not the
+		// sole gate on the secret material. A DB hiccup shouldn't take
+		// the whole A2A path down.
+		log.Printf("wsauth: caller HasAnyLiveToken(%s) failed: %v — allowing A2A", callerID, err)
+		return nil
+	}
+	if !hasLive {
+		return nil // legacy / pre-upgrade caller
+	}
+	tok := wsauth.BearerTokenFromHeader(c.GetHeader("Authorization"))
+	if tok == "" {
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "missing caller auth token"})
+		return errInvalidCallerToken
+	}
+	if err := wsauth.ValidateToken(ctx, db.DB, callerID, tok); err != nil {
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid caller auth token"})
+		return err
+	}
+	return nil
+}
+
+// errInvalidCallerToken is a sentinel for validateCallerToken's "missing
+// token" branch so the handler-level guard can detect it without string
+// matching (the wsauth errors are typed for the invalid case).
+var errInvalidCallerToken = errors.New("missing caller auth token")
+
+// extractAndUpsertTokenUsage parses LLM usage from a raw A2A response body
+// and persists it via upsertTokenUsage. Safe to call in a goroutine — logs
+// errors but never panics. ctx must already be detached from the request.
+func extractAndUpsertTokenUsage(ctx context.Context, workspaceID string, respBody []byte) {
+	in, out := parseUsageFromA2AResponse(respBody)
+	if in > 0 || out > 0 {
+		upsertTokenUsage(ctx, workspaceID, in, out)
+	}
+}
+
+// parseUsageFromA2AResponse extracts input_tokens / output_tokens from an A2A
+// JSON-RPC response. Inspects two locations in order of preference:
+//  1. result.usage — the JSON-RPC 2.0 result envelope from workspace agents.
+//  2. usage — top-level, for non-JSON-RPC or direct Anthropic-shaped payloads.
+//
+// Returns (0, 0) when no recognisable usage data is found.
+func parseUsageFromA2AResponse(body []byte) (inputTokens, outputTokens int64) {
+	if len(body) == 0 {
+		return 0, 0
+	}
+	var top map[string]json.RawMessage
+	if err := json.Unmarshal(body, &top); err != nil {
+		return 0, 0
+	}
+
+	// 1. result.usage (JSON-RPC 2.0 wrapper produced by workspace agents).
+	if rawResult, ok := top["result"]; ok {
+		var result map[string]json.RawMessage
+		if err := json.Unmarshal(rawResult, &result); err == nil {
+			if in, out, ok := readUsageMap(result); ok {
+				return in, out
+			}
+		}
+	}
+
+	// 2. Fallback: top-level usage (direct Anthropic or non-JSON-RPC response).
+	if in, out, ok := readUsageMap(top); ok {
+		return in, out
+	}
+	return 0, 0
+}
+
+// isSafeURL validates that a URL resolves to a publicly-routable address,
+// preventing A2A requests from being redirected to internal/cloud-metadata
+// infrastructure (SSRF, CWE-918). Workspace URLs come from DB/Redis caches
+// so we validate before making any outbound HTTP call.
+func isSafeURL(rawURL string) error {
+	u, err := url.Parse(rawURL)
+	if err != nil {
+		return fmt.Errorf("invalid URL: %w", err)
+	}
+	// Reject non-HTTP(S) schemes.
+	if u.Scheme != "http" && u.Scheme != "https" {
+		return fmt.Errorf("forbidden scheme: %s (only http/https allowed)", u.Scheme)
+	}
+	host := u.Hostname()
+	if host == "" {
+		return fmt.Errorf("empty hostname")
+	}
+	// Block direct IP addresses.
+	if ip := net.ParseIP(host); ip != nil {
+		if ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() {
+			return fmt.Errorf("forbidden loopback/unspecified IP: %s", ip)
+		}
+		if isPrivateOrMetadataIP(ip) {
+			return fmt.Errorf("forbidden private/metadata IP: %s", ip)
+		}
+		return nil
+	}
+	// For hostnames, resolve and validate each returned IP.
+	addrs, err := net.LookupHost(host)
+	if err != nil {
+		// DNS resolution failure — block it. Could be an internal hostname.
+		return fmt.Errorf("DNS resolution blocked for hostname: %s (%v)", host, err)
+	}
+	if len(addrs) == 0 {
+		return fmt.Errorf("DNS returned no addresses for: %s", host)
+	}
+	for _, addr := range addrs {
+		ip := net.ParseIP(addr)
+		if ip != nil && (ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() || isPrivateOrMetadataIP(ip)) {
+			return fmt.Errorf("hostname %s resolves to forbidden IP: %s", host, ip)
+		}
+	}
+	return nil
+}
+
+// isPrivateOrMetadataIP returns true for RFC-1918 private, carrier-grade NAT,
+// link-local, and cloud metadata ranges.
+func isPrivateOrMetadataIP(ip net.IP) bool {
+	var privateRanges = []net.IPNet{
+		{IP: net.ParseIP("10.0.0.0"), Mask: net.CIDRMask(8, 32)},
+		{IP: net.ParseIP("172.16.0.0"), Mask: net.CIDRMask(12, 32)},
+		{IP: net.ParseIP("192.168.0.0"), Mask: net.CIDRMask(16, 32)},
+		{IP: net.ParseIP("169.254.0.0"), Mask: net.CIDRMask(16, 32)},
+		{IP: net.ParseIP("100.64.0.0"), Mask: net.CIDRMask(10, 32)},
+		{IP: net.ParseIP("192.0.2.0"), Mask: net.CIDRMask(24, 32)},
+		{IP: net.ParseIP("198.51.100.0"), Mask: net.CIDRMask(24, 32)},
+		{IP: net.ParseIP("203.0.113.0"), Mask: net.CIDRMask(24, 32)},
+	}
+	ip = ip.To4()
+	if ip == nil {
+		return false
+	}
+	for _, r := range privateRanges {
+		if r.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
+
+// readUsageMap extracts input_tokens / output_tokens from the "usage" key of m.
+// Returns (0, 0, false) when the key is absent or contains no non-zero values.
+func readUsageMap(m map[string]json.RawMessage) (inputTokens, outputTokens int64, ok bool) {
+	rawUsage, has := m["usage"]
+	if !has {
+		return 0, 0, false
+	}
+	var usage struct {
+		InputTokens  int64 `json:"input_tokens"`
+		OutputTokens int64 `json:"output_tokens"`
+	}
+	if err := json.Unmarshal(rawUsage, &usage); err != nil {
+		return 0, 0, false
+	}
+	if usage.InputTokens == 0 && usage.OutputTokens == 0 {
+		return 0, 0, false
+	}
+	return usage.InputTokens, usage.OutputTokens, true
+}