Honor ctx cancellation and deadlines on every outbound call

Call, remote invokes, proxy hops and CLI commands all called Client.Do,
which takes no ctx, so a cancelled or expired caller kept waiting until
the transport's own read timeout. zap-proto/http v0.3.11 adds
Transport.DoContext; a client that implements the new ContextClient
interface is now driven through it, and Call returns an error wrapping
context.Canceled or context.DeadlineExceeded.

Co-authored-by: Hanzo Dev <[email protected]>

call.go

@@ -2,6 +2,7 @@ package zip
 
 import (
 	"context"
+	"fmt"
 	"strings"
 
 	"github.com/zap-proto/zip/internal/sock"
@@ -267,10 +268,10 @@ func (c *Conn) Close() error {
 // socket's peer credential (see [Peer]), not by anything this client could
 // set.
 //
-// ctx is honored to the extent the transport allows: an already-cancelled ctx
-// fails before the wire, and the transport's own read timeout bounds the call
-// (zap: 30s). It is not cancellable mid-flight, because abandoning a call
-// would mean abandoning the pooled buffers it is writing into.
+// ctx bounds the call: its deadline caps the exchange and cancelling it stops
+// the call mid-flight, over any transport whose client is a [ContextClient]
+// (zap is). Either way the error wraps ctx's cause, so errors.Is sees
+// context.Canceled or context.DeadlineExceeded.
 func Call[In, Out any](ctx context.Context, c *Conn, op string, in *In) (*Out, error) {
 	if c == nil || c.client == nil {
 		return nil, ErrInternal("zip: Call on a nil Conn")
@@ -279,7 +280,7 @@ func Call[In, Out any](ctx context.Context, c *Conn, op string, in *In) (*Out, e
 		return nil, ErrBadRequest("zip: invalid op name " + op)
 	}
 	if err := ctx.Err(); err != nil {
-		return nil, Errorf(499, "zip: call %s: %v", op, err)
+		return nil, fmt.Errorf("zip: call %s: %w", op, context.Cause(ctx))
 	}
 
 	var body []byte
@@ -303,7 +304,10 @@ func Call[In, Out any](ctx context.Context, c *Conn, op string, in *In) (*Out, e
 	req.SetBody(body)
 	forwardIdentity(ctx, req)
 
-	if err := c.client.Do(req, resp); err != nil {
+	if err := do(ctx, c.client, req, resp); err != nil {
+		if ended(err) {
+			return nil, fmt.Errorf("zip: call %s at %s: %w", op, c.addr, err)
+		}
 		// The far end, not this hop, is what failed — same reading Proxy gives a
 		// broken upstream.
 		return nil, Errorf(502, "zip: call %s at %s: %v", op, c.addr, err)

call_test.go

@@ -422,3 +422,56 @@ func TestCall_CancelledContextNeverReachesTheWire(t *testing.T) {
 		t.Fatal("a cancelled context must not reach the wire")
 	}
 }
+
+// slowApp serves one op that holds the call for seconds, so only the caller's
+// ctx can end it early.
+func slowApp(t *testing.T) *zip.App {
+	t.Helper()
+	app := zip.New(zip.Config{AppName: "slow", DisableStartupMessage: true})
+	app.Post("/v1/slow/wait", func(_ context.Context, _ *retireIn) (*boolOut, error) {
+		time.Sleep(5 * time.Second)
+		return &boolOut{}, nil
+	}, zip.WithOperationID("slow_wait"))
+	return app
+}
+
+// TestCall_CancelStopsCallInFlight — cancelling ctx while the callee is still
+// working ends the call at once with context.Canceled.
+func TestCall_CancelStopsCallInFlight(t *testing.T) {
+	c, err := zip.Dial(serveUDS(t, slowApp(t)))
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer func() { _ = c.Close() }()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	time.AfterFunc(50*time.Millisecond, cancel)
+	start := time.Now()
+	_, err = zip.Call[retireIn, boolOut](ctx, c, "slow_wait", &retireIn{Flag: "x"})
+	if took := time.Since(start); took > time.Second {
+		t.Fatalf("cancel took %v to stop the call", took)
+	}
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("err = %v, want context.Canceled", err)
+	}
+}
+
+// TestCall_DeadlineBoundsCall — a ctx deadline caps the call.
+func TestCall_DeadlineBoundsCall(t *testing.T) {
+	c, err := zip.Dial(serveUDS(t, slowApp(t)))
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer func() { _ = c.Close() }()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+	start := time.Now()
+	_, err = zip.Call[retireIn, boolOut](ctx, c, "slow_wait", &retireIn{Flag: "x"})
+	if took := time.Since(start); took > time.Second {
+		t.Fatalf("deadline took %v to stop the call", took)
+	}
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("err = %v, want context.DeadlineExceeded", err)
+	}
+}

cli.go

@@ -484,9 +484,7 @@ func (r Remote) Invoke(ctx context.Context, c Command, path map[string]string, b
 	return json.RawMessage(out), nil
 }
 
-// do performs one request over the transport Base names. ctx is honoured up to
-// the point the request is handed to the transport, which owns its own
-// deadlines from there.
+// do performs one request over the transport Base names, bounded by ctx.
 func (r Remote) do(ctx context.Context, method, path string, body []byte) ([]byte, error) {
 	if err := ctx.Err(); err != nil {
 		return nil, err
@@ -535,7 +533,7 @@ func (r Remote) do(ctx context.Context, method, path string, body []byte) ([]byt
 	for k, v := range r.Header {
 		req.Header.Set(k, v)
 	}
-	if err := t.Dial(host).Do(req, resp); err != nil {
+	if err := do(ctx, t.Dial(host), req, resp); err != nil {
 		return nil, fmt.Errorf("%s %s: %w", method, path, err)
 	}
 	out := append([]byte(nil), resp.Body()...)

go.mod

@@ -24,7 +24,7 @@ require (
 	github.com/valyala/fasthttp v1.70.0
 	github.com/zap-proto/fiber/v3 v3.2.1
 	github.com/zap-proto/go v1.8.3
-	github.com/zap-proto/http v0.3.5
+	github.com/zap-proto/http v0.3.11
 	github.com/zap-proto/mcp v1.0.5
 	golang.org/x/sys v0.45.0
 	golang.org/x/tools v0.45.0

go.sum

@@ -87,8 +87,8 @@ github.com/zap-proto/fiber/v3 v3.2.1 h1:k45oKyTwySPtGt8sPz2Ao8OUHc7pDEhai8Np2Ym6
 github.com/zap-proto/fiber/v3 v3.2.1/go.mod h1:eDm2z+ufJrkuE4MeX0Mea4oc/p7/HpXiZjVB+BXCKOA=
 github.com/zap-proto/go v1.8.3 h1:GWhb8y7ylgVcaVOR8XN/It56Xovm5Mkv+nQqiblnnoA=
 github.com/zap-proto/go v1.8.3/go.mod h1:fbNeWcssRRbTkQEDCrmUmzqJA1gKcAlgi1z4QrtEFZg=
-github.com/zap-proto/http v0.3.5 h1:ezSuUbF05jxFGpfmScBnzIVTeH1ZGhcDSm7d5yL8Byc=
-github.com/zap-proto/http v0.3.5/go.mod h1:UYfGhDDCetgxs65XSev8Lpf65COg5vKQK+cWwZGh4zQ=
+github.com/zap-proto/http v0.3.11 h1:LdaTSQtxcHRooQMwz0bZGjw+7IuONEwCG3QtLH/fHr0=
+github.com/zap-proto/http v0.3.11/go.mod h1:Gw4Kn/PpF1P+CHJyBQQ3aoFDI3P4p0wsjURFwjPnHzE=
 github.com/zap-proto/mcp v1.0.5 h1:1Hi0KNi4YlWW2zkQUMDK9n3jP3wZaRJ1rseMW9p9F0c=
 github.com/zap-proto/mcp v1.0.5/go.mod h1:cYOsq+33NyW4X8nUVtIs/qtezVovYazw1s0V7mrQBYA=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=

host_forward_test.go

@@ -3,6 +3,7 @@ package zip
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"strings"
 	"testing"
 
@@ -51,7 +52,7 @@ func TestForwardKeepsTheCallersHost(t *testing.T) {
 	req.Header.SetHost("lux.id")
 	req.Header.SetRequestURI("/.well-known/openid-configuration")
 
-	if err := forward(req, resp, rec, "/var/run/iam.sock", "", "mount /v1/iam"); err != nil {
+	if err := forward(context.Background(), req, resp, rec, "/var/run/iam.sock", "", "mount /v1/iam"); err != nil {
 		t.Fatalf("forward: %v", err)
 	}
 	if rec.sawHost != "lux.id" {
@@ -70,7 +71,7 @@ func TestForwardFallsBackToTheDialAddress(t *testing.T) {
 
 	req.Header.SetRequestURI("/mcp")
 
-	if err := forward(req, resp, rec, "plugin.sock", "", "mcp iam"); err != nil {
+	if err := forward(context.Background(), req, resp, rec, "plugin.sock", "", "mcp iam"); err != nil {
 		t.Fatalf("forward: %v", err)
 	}
 	if rec.sawHost != "plugin.sock" {
@@ -88,7 +89,7 @@ func TestForwardStillRewritesThePath(t *testing.T) {
 	req.Header.SetHost("hanzo.ai")
 	req.Header.SetRequestURI("/v1/iam/mcp")
 
-	if err := forward(req, resp, rec, "iam.sock", "/mcp", "mcp iam"); err != nil {
+	if err := forward(context.Background(), req, resp, rec, "iam.sock", "/mcp", "mcp iam"); err != nil {
 		t.Fatalf("forward: %v", err)
 	}
 	if got := string(req.URI().Path()); got != "/mcp" {
@@ -104,7 +105,7 @@ func TestForwardWithNoClientIs503(t *testing.T) {
 	req, resp := fasthttp.AcquireRequest(), fasthttp.AcquireResponse()
 	defer fasthttp.ReleaseRequest(req)
 	defer fasthttp.ReleaseResponse(resp)
-	err := forward(req, resp, nil, "iam.sock", "", "mount /v1/iam")
+	err := forward(context.Background(), req, resp, nil, "iam.sock", "", "mount /v1/iam")
 	if err == nil || !strings.Contains(err.Error(), "no instance running") {
 		t.Fatalf("err = %v, want a 503 naming the mount", err)
 	}

mcp.go

@@ -817,7 +817,7 @@ func (a *App) relay(ctx context.Context, f *zapmcp.Frame, p *plugin) *zapmcp.Fra
 	req.Header.SetContentType(mimeJSON)
 	req.SetBody(body)
 	forwardIdentity(ctx, req)
-	if err := forward(req, resp, client, host, p.spec.mcpPath(), "mcp "+p.name); err != nil {
+	if err := forward(ctx, req, resp, client, host, p.spec.mcpPath(), "mcp "+p.name); err != nil {
 		return a.answer(f, nil, err)
 	}
 	var ans zapmcp.Frame

remote.go

@@ -104,7 +104,7 @@ func remoteApp(parent *App, prefix, addr string, d Declaration) (*App, error) {
 	r := newApp(cfg)
 
 	proxy := func(c *Ctx) error {
-		return forward(c.fc.Request(), c.fc.Response(), client, hostport, "", "mount "+prefix)
+		return forward(c.Context(), c.fc.Request(), c.fc.Response(), client, hostport, "", "mount "+prefix)
 	}
 	if len(d.Routes) == 0 {
 		// Nothing declared: the prefix and everything under it, which is what a
@@ -233,7 +233,10 @@ func remoteInvoke(client Client, host, method, pattern string) func(context.Cont
 		}
 		forwardIdentity(ctx, req)
 
-		if err := client.Do(req, resp); err != nil {
+		if err := do(ctx, client, req, resp); err != nil {
+			if ended(err) {
+				return nil, fmt.Errorf("zip: call %s %s at %s: %w", method, url, host, err)
+			}
 			return nil, Errorf(502, "zip: call %s %s at %s: %v", method, url, host, err)
 		}
 		switch code := resp.StatusCode(); {

transport.go

@@ -3,6 +3,7 @@ package zip
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -51,6 +52,34 @@ type Client interface {
 	Do(req *fasthttp.Request, resp *fasthttp.Response) error
 }
 
+// ContextClient is a Client that can be cancelled mid-flight and bounded by a
+// deadline. *http.Transport implements it; a Client that does not is still
+// usable, it just cannot be stopped once the request is on the wire.
+type ContextClient interface {
+	DoContext(ctx context.Context, req *fasthttp.Request, resp *fasthttp.Response) error
+}
+
+var _ ContextClient = (*http.Transport)(nil)
+
+// do sends req over client under ctx. A client without DoContext gets the
+// ctx checked once, before the wire.
+func do(ctx context.Context, client Client, req *fasthttp.Request, resp *fasthttp.Response) error {
+	if cc, ok := client.(ContextClient); ok {
+		return cc.DoContext(ctx, req, resp)
+	}
+	if err := ctx.Err(); err != nil {
+		return context.Cause(ctx)
+	}
+	return client.Do(req, resp)
+}
+
+// ended reports whether err is a ctx ending rather than a far-end failure. The
+// conn's deadline can fire a moment before ctx's own timer, so ctx.Err() alone
+// can still be nil here.
+func ended(err error) bool {
+	return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)
+}
+
 // Transport is one address scheme in both directions: Serve terminates bytes
 // arriving here, Dial originates bytes going there. They are the same concept
 // — a wire — so they live in one value rather than two registries that drift.
@@ -333,7 +362,7 @@ func (a *App) mountVia(prefix string, to func() (Client, string)) {
 		if upgrading(c.fc.Request()) {
 			return relay(c.fc.RequestCtx(), host, "mount "+prefix)
 		}
-		return forward(c.fc.Request(), c.fc.Response(), client, host, "", "mount "+prefix)
+		return forward(c.Context(), c.fc.Request(), c.fc.Response(), client, host, "", "mount "+prefix)
 	}
 	prefix = strings.TrimSuffix(normPath(prefix), "/")
 	site := here(1)
@@ -359,7 +388,7 @@ func (a *App) mountVia(prefix string, to func() (Client, string)) {
 //
 // what names the caller in the error, because "no instance running" is worth
 // nothing without knowing what did not run.
-func forward(req *fasthttp.Request, resp *fasthttp.Response, client Client, host, path, what string) error {
+func forward(ctx context.Context, req *fasthttp.Request, resp *fasthttp.Response, client Client, host, path, what string) error {
 	if client == nil {
 		return Errorf(503, "%s: no instance running", what)
 	}
@@ -374,7 +403,7 @@ func forward(req *fasthttp.Request, resp *fasthttp.Response, client Client, host
 	if path != "" {
 		req.URI().SetPath(path)
 	}
-	if err := client.Do(req, resp); err != nil {
+	if err := do(ctx, client, req, resp); err != nil {
 		// The upstream, not this hop, is what failed.
 		return Errorf(502, "%s: %v", what, err)
 	}