Move registry to sub module.

Makefile

@@ -4,7 +4,7 @@ test:
 	go test -v -cover ./...
 
 test-registry:
-	REGISTRY_TEST_ENABLE=true go test -v -cover ./registry/...
+	REGISTRY_TEST_ENABLE=true go test -v -cover ./clients/registry/...
 
 check-moc:
 	find ic -type f -name '*.mo' -print0 | xargs -0 $(shell dfx cache show)/moc --check
@@ -16,7 +16,7 @@ test-cover:
 gen:
 	cd candid && go generate
 	cd pocketic && go generate
-	cd registry && go generate
+	cd clients/registry && go generate
 
 gen-ic:
 	go run ic/testdata/gen.go

agent.go

@@ -1,11 +1,13 @@
 package agent
 
 import (
+	"context"
 	"crypto/rand"
 	"encoding/binary"
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"github.com/aviate-labs/agent-go/candid/idl"
 	"net/url"
 	"reflect"
 	"time"
@@ -89,9 +91,66 @@ func uint64FromBytes(raw []byte) uint64 {
 	}
 }
 
+type APIRequest[In, Out any] struct {
+	a                   *Agent
+	unmarshal           func([]byte, Out) error
+	typ                 RequestType
+	methodName          string
+	effectiveCanisterID principal.Principal
+	requestID           RequestID
+	data                []byte
+}
+
+func CreateAPIRequest[In, Out any](
+	a *Agent,
+	marshal func(In) ([]byte, error),
+	unmarshal func([]byte, Out) error,
+	typ RequestType,
+	canisterID principal.Principal,
+	methodName string,
+	in In,
+) (*APIRequest[In, Out], error) {
+	rawArgs, err := marshal(in)
+	if err != nil {
+		return nil, err
+	}
+	nonce, err := newNonce()
+	if err != nil {
+		return nil, err
+	}
+	requestID, data, err := a.sign(Request{
+		Type:          typ,
+		Sender:        a.Sender(),
+		CanisterID:    canisterID,
+		MethodName:    methodName,
+		Arguments:     rawArgs,
+		IngressExpiry: a.expiryDate(),
+		Nonce:         nonce,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &APIRequest[In, Out]{
+		a:                   a,
+		unmarshal:           unmarshal,
+		typ:                 typ,
+		methodName:          methodName,
+		effectiveCanisterID: canisterID,
+		requestID:           *requestID,
+		data:                data,
+	}, nil
+}
+
+// WithEffectiveCanisterID sets the effective canister ID for the Call.
+func (c *APIRequest[In, Out]) WithEffectiveCanisterID(canisterID principal.Principal) *APIRequest[In, Out] {
+	c.effectiveCanisterID = canisterID
+	return c
+}
+
 // Agent is a client for the Internet Computer.
 type Agent struct {
 	client           Client
+	ctx              context.Context
 	identity         identity.Identity
 	ingressExpiry    time.Duration
 	rootKey          []byte
@@ -103,7 +162,7 @@ type Agent struct {
 // New returns a new Agent based on the given configuration.
 func New(cfg Config) (*Agent, error) {
 	if cfg.IngressExpiry == 0 {
-		cfg.IngressExpiry = time.Minute
+		cfg.IngressExpiry = 5 * time.Minute
 	}
 	// By default, use the anonymous identity.
 	var id identity.Identity = new(identity.AnonymousIdentity)
@@ -139,6 +198,7 @@ func New(cfg Config) (*Agent, error) {
 	}
 	return &Agent{
 		client:           client,
+		ctx:              context.Background(),
 		identity:         id,
 		ingressExpiry:    cfg.IngressExpiry,
 		rootKey:          rootKey,
@@ -154,6 +214,19 @@ func (a Agent) Client() *Client {
 	return &a.client
 }
 
+// CreateCandidAPIRequest creates a new api request to the given canister and method.
+func (a *Agent) CreateCandidAPIRequest(typ RequestType, canisterID principal.Principal, methodName string, args ...any) (*CandidAPIRequest, error) {
+	return CreateAPIRequest[[]any, []any](
+		a,
+		idl.Marshal,
+		idl.Unmarshal,
+		typ,
+		effectiveCanisterID(canisterID, args),
+		methodName,
+		args,
+	)
+}
+
 // GetCanisterControllers returns the list of principals that can control the given canister.
 func (a Agent) GetCanisterControllers(canisterID principal.Principal) ([]principal.Principal, error) {
 	resp, err := a.GetCanisterInfo(canisterID, "controllers")
@@ -252,7 +325,9 @@ func (a Agent) Sender() principal.Principal {
 }
 
 func (a Agent) call(ecID principal.Principal, data []byte) ([]byte, error) {
-	return a.client.Call(ecID, data)
+	ctx, cancel := context.WithTimeout(a.ctx, a.ingressExpiry)
+	defer cancel()
+	return a.client.Call(ctx, ecID, data)
 }
 
 func (a Agent) expiryDate() uint64 {
@@ -299,7 +374,9 @@ func (a Agent) poll(ecID principal.Principal, requestID RequestID) ([]byte, erro
 }
 
 func (a Agent) readState(ecID principal.Principal, data []byte) (map[string][]byte, error) {
-	resp, err := a.client.ReadState(ecID, data)
+	ctx, cancel := context.WithTimeout(a.ctx, a.ingressExpiry)
+	defer cancel()
+	resp, err := a.client.ReadState(ctx, ecID, data)
 	if err != nil {
 		return nil, err
 	}
@@ -336,7 +413,9 @@ func (a Agent) readStateCertificate(ecID principal.Principal, paths [][]hashtree
 }
 
 func (a Agent) readSubnetState(subnetID principal.Principal, data []byte) (map[string][]byte, error) {
-	resp, err := a.client.ReadSubnetState(subnetID, data)
+	ctx, cancel := context.WithTimeout(a.ctx, a.ingressExpiry)
+	defer cancel()
+	resp, err := a.client.ReadSubnetState(ctx, subnetID, data)
 	if err != nil {
 		return nil, err
 	}
@@ -385,12 +464,14 @@ func (a Agent) sign(request Request) (*RequestID, []byte, error) {
 	return &requestID, data, nil
 }
 
+type CandidAPIRequest = APIRequest[[]any, []any]
+
 // Config is the configuration for an Agent.
 type Config struct {
 	// Identity is the identity used by the Agent.
 	Identity identity.Identity
 	// IngressExpiry is the duration for which an ingress message is valid.
-	// The default is set to 1 minute.
+	// The default is set to 5 minutes.
 	IngressExpiry time.Duration
 	// ClientConfig is the configuration for the underlying Client.
 	ClientConfig *ClientConfig

call.go

@@ -1,18 +1,41 @@
 package agent
 
 import (
-	"github.com/aviate-labs/agent-go/candid/idl"
 	"github.com/aviate-labs/agent-go/principal"
 	"google.golang.org/protobuf/proto"
 )
 
+// Call calls a method on a canister, it does not wait for the result.
+func (c APIRequest[_, _]) Call() error {
+	c.a.logger.Printf("[AGENT] CALL %s %s (%x)", c.effectiveCanisterID, c.methodName, c.requestID)
+	_, err := c.a.call(c.effectiveCanisterID, c.data)
+	return err
+}
+
+// CallAndWait calls a method on a canister and waits for the result.
+func (c APIRequest[_, Out]) CallAndWait(out Out) error {
+	if err := c.Call(); err != nil {
+		return err
+	}
+	return c.Wait(out)
+}
+
+// Wait waits for the result of the Call and unmarshals it into the given values.
+func (c APIRequest[_, Out]) Wait(out Out) error {
+	raw, err := c.a.poll(c.effectiveCanisterID, c.requestID)
+	if err != nil {
+		return err
+	}
+	return c.unmarshal(raw, out)
+}
+
 // Call calls a method on a canister and unmarshals the result into the given values.
-func (a Agent) Call(canisterID principal.Principal, methodName string, args []any, values []any) error {
-	call, err := a.CreateCall(canisterID, methodName, args...)
+func (a Agent) Call(canisterID principal.Principal, methodName string, in []any, out []any) error {
+	call, err := a.CreateCandidAPIRequest(RequestTypeCall, canisterID, methodName, in...)
 	if err != nil {
 		return err
 	}
-	return call.CallAndWait(values...)
+	return call.CallAndWait(out)
 }
 
 // CallProto calls a method on a canister and unmarshals the result into the given proto message.
@@ -41,77 +64,3 @@ func (a Agent) CallProto(canisterID principal.Principal, methodName string, in,
 	}
 	return proto.Unmarshal(raw, out)
 }
-
-// CreateCall creates a new Call to the given canister and method.
-func (a *Agent) CreateCall(canisterID principal.Principal, methodName string, args ...any) (*Call, error) {
-	rawArgs, err := idl.Marshal(args)
-	if err != nil {
-		return nil, err
-	}
-	if len(args) == 0 {
-		// Default to the empty Candid argument list.
-		rawArgs = []byte{'D', 'I', 'D', 'L', 0, 0}
-	}
-	nonce, err := newNonce()
-	if err != nil {
-		return nil, err
-	}
-	requestID, data, err := a.sign(Request{
-		Type:          RequestTypeCall,
-		Sender:        a.Sender(),
-		CanisterID:    canisterID,
-		MethodName:    methodName,
-		Arguments:     rawArgs,
-		IngressExpiry: a.expiryDate(),
-		Nonce:         nonce,
-	})
-	if err != nil {
-		return nil, err
-	}
-	return &Call{
-		a:                   a,
-		methodName:          methodName,
-		effectiveCanisterID: effectiveCanisterID(canisterID, args),
-		requestID:           *requestID,
-		data:                data,
-	}, nil
-}
-
-// Call is an intermediate representation of a Call to a canister.
-type Call struct {
-	a                   *Agent
-	methodName          string
-	effectiveCanisterID principal.Principal
-	requestID           RequestID
-	data                []byte
-}
-
-// Call calls a method on a canister, it does not wait for the result.
-func (c Call) Call() error {
-	c.a.logger.Printf("[AGENT] CALL %s %s (%x)", c.effectiveCanisterID, c.methodName, c.requestID)
-	_, err := c.a.call(c.effectiveCanisterID, c.data)
-	return err
-}
-
-// CallAndWait calls a method on a canister and waits for the result.
-func (c Call) CallAndWait(values ...any) error {
-	if err := c.Call(); err != nil {
-		return err
-	}
-	return c.Wait(values...)
-}
-
-// Wait waits for the result of the Call and unmarshals it into the given values.
-func (c Call) Wait(values ...any) error {
-	raw, err := c.a.poll(c.effectiveCanisterID, c.requestID)
-	if err != nil {
-		return err
-	}
-	return idl.Unmarshal(raw, values)
-}
-
-// WithEffectiveCanisterID sets the effective canister ID for the Call.
-func (c *Call) WithEffectiveCanisterID(canisterID principal.Principal) *Call {
-	c.effectiveCanisterID = canisterID
-	return c
-}

client.go

@@ -2,6 +2,7 @@ package agent
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -40,10 +41,14 @@ func NewClientWithLogger(cfg ClientConfig, logger Logger) Client {
 	}
 }
 
-func (c Client) Call(canisterID principal.Principal, data []byte) ([]byte, error) {
+func (c Client) Call(ctx context.Context, canisterID principal.Principal, data []byte) ([]byte, error) {
 	u := c.url(fmt.Sprintf("/api/v2/canister/%s/call", canisterID.Encode()))
 	c.logger.Printf("[CLIENT] CALL %s", u)
-	resp, err := c.client.Post(u, "application/cbor", bytes.NewBuffer(data))
+	req, err := c.newRequest(ctx, "POST", u, bytes.NewBuffer(data))
+	if err != nil {
+		return nil, err
+	}
+	resp, err := c.client.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -63,16 +68,16 @@ func (c Client) Call(canisterID principal.Principal, data []byte) ([]byte, error
 	}
 }
 
-func (c Client) Query(canisterID principal.Principal, data []byte) ([]byte, error) {
-	return c.post("query", canisterID, data)
+func (c Client) Query(ctx context.Context, canisterID principal.Principal, data []byte) ([]byte, error) {
+	return c.post(ctx, "query", canisterID, data)
 }
 
-func (c Client) ReadState(canisterID principal.Principal, data []byte) ([]byte, error) {
-	return c.post("read_state", canisterID, data)
+func (c Client) ReadState(ctx context.Context, canisterID principal.Principal, data []byte) ([]byte, error) {
+	return c.post(ctx, "read_state", canisterID, data)
 }
 
-func (c Client) ReadSubnetState(subnetID principal.Principal, data []byte) ([]byte, error) {
-	return c.postSubnet("read_state", subnetID, data)
+func (c Client) ReadSubnetState(ctx context.Context, subnetID principal.Principal, data []byte) ([]byte, error) {
+	return c.postSubnet(ctx, "read_state", subnetID, data)
 }
 
 // Status returns the status of the IC.
@@ -94,10 +99,23 @@ func (c Client) get(path string) ([]byte, error) {
 	return io.ReadAll(resp.Body)
 }
 
-func (c Client) post(path string, canisterID principal.Principal, data []byte) ([]byte, error) {
+func (c Client) newRequest(ctx context.Context, method, url string, body io.Reader) (*http.Request, error) {
+	req, err := http.NewRequestWithContext(ctx, method, url, body)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Content-Type", "application/cbor")
+	return req, nil
+}
+
+func (c Client) post(ctx context.Context, path string, canisterID principal.Principal, data []byte) ([]byte, error) {
 	u := c.url(fmt.Sprintf("/api/v2/canister/%s/%s", canisterID.Encode(), path))
 	c.logger.Printf("[CLIENT] POST %s", u)
-	resp, err := c.client.Post(u, "application/cbor", bytes.NewBuffer(data))
+	req, err := c.newRequest(ctx, "POST", u, bytes.NewBuffer(data))
+	if err != nil {
+		return nil, err
+	}
+	resp, err := c.client.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -110,10 +128,14 @@ func (c Client) post(path string, canisterID principal.Principal, data []byte) (
 	}
 }
 
-func (c Client) postSubnet(path string, subnetID principal.Principal, data []byte) ([]byte, error) {
+func (c Client) postSubnet(ctx context.Context, path string, subnetID principal.Principal, data []byte) ([]byte, error) {
 	u := c.url(fmt.Sprintf("/api/v2/subnet/%s/%s", subnetID.Encode(), path))
 	c.logger.Printf("[CLIENT] POST %s", u)
-	resp, err := c.client.Post(u, "application/cbor", bytes.NewBuffer(data))
+	req, err := c.newRequest(ctx, "POST", u, bytes.NewBuffer(data))
+	if err != nil {
+		return nil, err
+	}
+	resp, err := c.client.Do(req)
 	if err != nil {
 		return nil, err
 	}