Go codegen: address review feedback from #22057

src/Aspire.Hosting.CodeGeneration.Go/AtsGoCodeGenerator.ApiProjection.cs

@@ -1012,7 +1012,14 @@ private List<GoApiMember> ProjectSynthesizedMembers(ApiExportScope scope, string
             members.AddRange(ProjectMembers(interfaceName, helpers.Where(helper => helper.Kind == GoDeclarationKind.Method), GoApiErrorModel.Returned));
         }
 
-        members.AddRange(ProjectMembers(interfaceName, [scope.ErrMethod], GoApiErrorModel.None));
+        var errMembers = ProjectMembers(interfaceName, [scope.ErrMethod], GoApiErrorModel.None);
+        if (string.Equals(typeId, AtsConstants.BuilderTypeId, StringComparison.Ordinal))
+        {
+            // The builder's generated Err overrides the promoted handle Err and reports every client error.
+            errMembers = errMembers.Select(err => err with { Summary = BuilderErrSummary }).ToList();
+        }
+
+        members.AddRange(errMembers);
         return members;
     }
 

src/Aspire.Hosting.CodeGeneration.Go/AtsGoCodeGenerator.cs

@@ -2078,6 +2078,17 @@ private void EmitHandleParamErrorChecks(string indent, AtsCapabilityInfo capabil
             var paramName = GetLocalIdentifier(p.Name);
             WriteLine($"{indent}if {paramName} != nil {{ if err := {paramName}.Err(); err != nil {{ {errorAction} }} }}");
         }
+
+        // Handle-valued unions are plain `any` locals, so the wrapper's Err is read through
+        // an interface assertion on handleReference so unrelated values with an Err method
+        // (for example a context.Context) still reach union validation. Without the guard
+        // serializeValue turns an errored wrapper into nil.
+        foreach (var p in GetUnionParameters(capability))
+        {
+            if (p.Type!.UnionTypes is not { } members || !members.Any(m => m.Category == AtsTypeCategory.Handle)) { continue; }
+            var paramName = GetLocalIdentifier(p.Name);
+            WriteLine($"{indent}if e, ok := {paramName}.(handleReference); ok {{ if err := e.Err(); err != nil {{ {errorAction} }} }}");
+        }
     }
 
     private void EmitArgsConstruction(
@@ -2186,18 +2197,18 @@ private void EmitCallbackRegistration(string indent, AtsParameterInfo p, string
 
     // Emits a strongly-typed DTO callback property into the DTO's ToMap output. Method parameters
     // pre-register their callback (putting a string id into reqArgs), but DTO properties instead
-    // embed the shim func directly in the map; the client's marshalTransportValue walks the
-    // serialized args, finds the func(...any) any value, and registers it. This keeps the public
+    // embed a clientCallback in the map; the client's marshalTransportValue walks the
+    // serialized args, binds itself to that shim, and registers it. This keeps the public
     // DTO field strongly typed (e.g. func(ctx InputsDialogValidationContext)) while reusing the
     // exact arg-decoding/return/writeback behavior of method-parameter callbacks.
     private void EmitDtoCallbackToMap(AtsDtoPropertyInfo property, string propertyName)
     {
         var p = DtoPropertyToParameterInfo(property);
         WriteLine($"\tif d.{propertyName} != nil {{");
         WriteLine($"\t\tcb := d.{propertyName}");
-        WriteLine($"\t\tm[{GoStringLiteral(JsonNamingPolicy.CamelCase.ConvertName(property.Name))}] = func(args ...any) any {{");
-        EmitGoCallbackShimBody("\t\t\t", p, clientExpression: null);
-        WriteLine($"\t\t}}");
+        WriteLine($"\t\tm[{GoStringLiteral(JsonNamingPolicy.CamelCase.ConvertName(property.Name))}] = clientCallback(func(c *client, args ...any) any {{");
+        EmitGoCallbackShimBody("\t\t\t", p, clientExpression: "c");
+        WriteLine($"\t\t}})");
         WriteLine($"\t}}");
     }
 

src/Aspire.Hosting.CodeGeneration.Go/Resources/transport.go

@@ -95,6 +95,10 @@ func (e *CapabilityError) Error() string      { return e.err.Error() }
 //	}
 //
 // Non-SDK errors are returned via err.Error() unchanged.
+//
+// Callers commonly add context with fmt.Errorf("...: %w", err); the joined
+// failures are still found underneath, and the added context is printed on its
+// own line before them.
 func FormatError(err error) string {
 	if err == nil {
 		return ""
@@ -103,9 +107,15 @@ func FormatError(err error) string {
 	if !errors.As(err, &capErr) {
 		return err.Error()
 	}
-	if joined, ok := err.(interface{ Unwrap() []error }); ok {
+	if aggregate, joined := findJoinedErrors(err); joined != nil {
 		var b strings.Builder
-		for _, child := range joined.Unwrap() {
+		// "building app: " + aggregate.Error() -> "building app"
+		if prefix, ok := strings.CutSuffix(err.Error(), aggregate.Error()); ok {
+			if prefix = strings.TrimSuffix(prefix, ": "); prefix != "" {
+				b.WriteString(prefix)
+			}
+		}
+		for _, child := range joined {
 			formatted := FormatError(child)
 			if formatted == "" {
 				continue
@@ -146,6 +156,23 @@ func FormatError(err error) string {
 	return b.String()
 }
 
+// findJoinedErrors walks single-error wrappers (fmt.Errorf with one %w) down to
+// the first error that wraps several, such as the result of errors.Join. It
+// returns that aggregate error and the errors it wraps.
+func findJoinedErrors(err error) (error, []error) {
+	for err != nil {
+		switch u := err.(type) {
+		case interface{ Unwrap() []error }:
+			return err, u.Unwrap()
+		case interface{ Unwrap() error }:
+			err = u.Unwrap()
+		default:
+			return nil, nil
+		}
+	}
+	return nil, nil
+}
+
 // CancellationToken wraps a context.Context to provide cooperative cancellation.
 // Use NewCancellationToken to create one; call Cancel() to cancel it.
 // Pass token.Context() to any standard Go library that accepts a context.
@@ -154,6 +181,13 @@ type CancellationToken struct {
 	cancel     context.CancelFunc
 	cancelOnce sync.Once
 	notify     func()
+	// remoteID is the host-issued ID of a token received in a callback. Forwarding
+	// the token to another capability reuses it so the host links the nested
+	// operation to the original cancellation source.
+	remoteID string
+	// remoteClient is the connection that issued remoteID; cancellation IDs are
+	// only meaningful on that connection.
+	remoteClient *client
 }
 
 // NewCancellationToken creates a new cancellation token backed by a cancellable
@@ -171,6 +205,8 @@ func newCancellationTokenFrom(parent context.Context) *CancellationToken {
 
 func newRemoteCancellationToken(c *client, id string) *CancellationToken {
 	token := newCancellationTokenFrom(context.Background())
+	token.remoteID = id
+	token.remoteClient = c
 	token.notify = func() {
 		go c.cancelToken(id)
 	}
@@ -750,6 +786,12 @@ func (c *client) registerCancellation(token *CancellationToken) string {
 	if token == nil {
 		return ""
 	}
+	// The host already owns this ID on this connection, so there is nothing to register
+	// or watch locally: Cancel on the remote token notifies the host itself. A token
+	// forwarded through a different client takes the normal local registration below.
+	if token.remoteID != "" && token.remoteClient == c {
+		return token.remoteID
+	}
 	id := fmt.Sprintf("ct_%d_%d", c.cancellationCounter.Add(1), time.Now().UnixMilli())
 	c.cancellationsMu.Lock()
 	c.cancellations[id] = &cancellationRegistration{token: token}
@@ -875,6 +917,11 @@ func (c *client) invokeCapability(ctx context.Context, capabilityID string, args
 	return wrapIfHandle(result, c)
 }
 
+// clientCallback is a callback shim that needs the registering client, for
+// example to build remote cancellation tokens. DTO ToMap methods have no client
+// in scope, so they emit this type and marshalTransportValue binds the client.
+type clientCallback func(c *client, args ...any) any
+
 func (c *client) marshalTransportValue(value any) any {
 	if callback, ok := value.(func(...any) any); ok {
 		return c.registerCallback(callback)
@@ -884,6 +931,8 @@ func (c *client) marshalTransportValue(value any) any {
 	switch v := serialized.(type) {
 	case func(...any) any:
 		return c.registerCallback(v)
+	case clientCallback:
+		return c.registerCallback(func(args ...any) any { return v(c, args...) })
 	case map[string]any:
 		result := make(map[string]any, len(v))
 		for key, nestedValue := range v {

tests/Aspire.Hosting.CodeGeneration.Go.Tests/AtsGoCodeGeneratorTests.cs

@@ -1359,6 +1359,7 @@ package aspire
                 import (
                     "bufio"
                     "bytes"
+                    "context"
                     "encoding/json"
                     "errors"
                     "fmt"
@@ -1709,6 +1710,80 @@ func TestPropagatedChildErrorIsNotRecordedAgain(t *testing.T) {
                     }
                 }
 
+                func TestFormatErrorIncludesJoinedErrorsBehindWrappers(t *testing.T) {
+                    first := &CapabilityError{err: atsError{Code: "FIRST", Message: "first failed", Capability: "addFirst"}}
+                    second := &CapabilityError{err: atsError{Code: "SECOND", Message: "second failed", Capability: "addSecond"}}
+
+                    want := FormatError(errors.Join(first, second))
+                    wrapped := fmt.Errorf("building app: %w", fmt.Errorf("resources: %w", errors.Join(first, second)))
+                    if got, wantWrapped := FormatError(wrapped), "building app: resources\n"+want; got != wantWrapped {
+                        t.Errorf("FormatError() = %q, want %q", got, wantWrapped)
+                    }
+                }
+
+                func TestForwardedCallbackTokenFromOtherClientIsRegisteredLocally(t *testing.T) {
+                    other := newClient("")
+                    remote := newRemoteCancellationToken(newClient(""), "ct_host")
+                    if id := other.registerCancellation(remote); id == "ct_host" || id == "" {
+                        t.Fatalf("registerCancellation() = %q, want a locally issued ID", id)
+                    }
+                }
+
+                func TestForwardedCallbackTokenReusesHostID(t *testing.T) {
+                    c := newClient("")
+                    remote := newRemoteCancellationToken(c, "ct_host")
+                    if id := c.registerCancellation(remote); id != "ct_host" {
+                        t.Fatalf("registerCancellation() = %q, want the host-issued ID", id)
+                    }
+                    c.cancellationsMu.Lock()
+                    _, registered := c.cancellations["ct_host"]
+                    c.cancellationsMu.Unlock()
+                    if registered {
+                        t.Fatal("a host-owned token ID was registered as a local cancellation")
+                    }
+                }
+
+                func TestClientCallbackIsBoundToRegisteringClient(t *testing.T) {
+                    c := newClient("")
+                    var got *client
+                    marshaled, _ := c.marshalTransportValue(map[string]any{
+                        "cb": clientCallback(func(cl *client, args ...any) any { got = cl; return nil }),
+                    }).(map[string]any)
+                    id, _ := marshaled["cb"].(string)
+                    c.callbacksMu.Lock()
+                    callback := c.callbacks[id]
+                    c.callbacksMu.Unlock()
+                    if callback == nil {
+                        t.Fatalf("callback %q was not registered", id)
+                    }
+                    callback()
+                    if got != c {
+                        t.Fatalf("callback received client %p, want %p", got, c)
+                    }
+                }
+
+                func TestErroredHandleInUnionIsNotSent(t *testing.T) {
+                    ctx, cancel := context.WithCancel(context.Background())
+                    cancel()
+                    unrelated := newTestRedisResourceFromHandle(&handle{HandleID: "u", TypeID: testRedisTypeID}, newClient(""))
+                    unrelated.WithUnionDependency(ctx)
+                    if err := unrelated.Err(); err == nil || !strings.Contains(err.Error(), "must be one of") {
+                        t.Errorf("a canceled context got %v, want the union guard error", err)
+                    }
+                    c, received := connectChildErrorHost(t, "addTestRedis")
+                    builder := newBuilderOn(c)
+
+                    failed := builder.AddTestRedis("failed")
+                    target := newTestRedisResourceFromHandle(&handle{HandleID: "target", TypeID: testRedisTypeID}, c)
+                    target.WithUnionDependency(failed)
+                    if !errors.Is(target.Err(), failed.Err()) {
+                        t.Errorf("target.Err() = %v, want the dependency's error %v", target.Err(), failed.Err())
+                    }
+                    if got, want := received(), []string{"addTestRedis"}; !slices.Equal(got, want) {
+                        t.Errorf("the AppHost received %v, want %v (withUnionDependency must not be sent)", got, want)
+                    }
+                }
+
                 func TestFailedWithCallOnChildReachesBuilder(t *testing.T) {
                     c, _ := connectChildErrorHost(t, "withConnectionStringDirect")
                     builder := newBuilderOn(c)