From ead9ee3a7a0c1c480631c0e8e52bb7a70090a8e6 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek <8347993+nik13@users.noreply.github.com> Date: Sat, 31 Oct 2026 18:21:28 -0700 Subject: [PATCH 1/6] feat(a2a): trace Go server calls through official SDK hooks Co-Authored-By: Claude Code --- .github/workflows/go-ci.yml | 7 + go/a2a/README.md | 48 ++++- go/a2a/doc.go | 7 +- go/a2a/example/main.go | 30 +-- go/a2a/example/main_test.go | 11 +- go/a2a/interceptor.go | 116 ++++++----- go/a2a/server.go | 119 ++++++++++++ go/a2a/server_lifecycle_test.go | 328 ++++++++++++++++++++++++++++++++ go/a2a/server_otlp_test.go | 128 +++++++++++++ go/a2a/server_test.go | 126 ++++++++++++ 10 files changed, 831 insertions(+), 89 deletions(-) create mode 100644 go/a2a/server.go create mode 100644 go/a2a/server_lifecycle_test.go create mode 100644 go/a2a/server_otlp_test.go create mode 100644 go/a2a/server_test.go diff --git a/.github/workflows/go-ci.yml b/.github/workflows/go-ci.yml index 498cf313f..5932d2675 100644 --- a/.github/workflows/go-ci.yml +++ b/.github/workflows/go-ci.yml @@ -64,3 +64,10 @@ jobs: run: go build -o "$RUNNER_TEMP/traceai-example" ./example - name: Render package documentation run: go doc . + - name: Test A2A current SDK compatibility + if: matrix.module == 'a2a' + run: | + # The preceding checks use the v2.5.0 floor from go.mod. + go get github.com/a2aproject/a2a-go/v2@v2.7.0 + go mod tidy + go test -race -count=1 ./... diff --git a/go/a2a/README.md b/go/a2a/README.md index ff9f8d62f..c5284855c 100644 --- a/go/a2a/README.md +++ b/go/a2a/README.md @@ -1,15 +1,16 @@ # traceAI A2A Go -OpenTelemetry tracing for A2A (Agent2Agent) protocol clients built with the official Go SDK, +OpenTelemetry tracing for A2A (Agent2Agent) protocol clients and servers built with the official Go SDK, [`github.com/a2aproject/a2a-go/v2`](https://github.com/a2aproject/a2a-go), sending spans to [Future AGI](https://futureagi.com). -It is an `a2aclient.CallInterceptor`, the SDK's own extension point: you pass it with +The client uses `a2aclient.CallInterceptor`, the SDK's own extension point: you pass it with `a2aclient.WithCallInterceptors` when you create the client. Each call to a remote agent becomes one client span with `gen_ai.span.kind=AGENT` and the A2A attribute names that the Python `traceai-a2a` package uses (`gen_ai.a2a.task.id`, `gen_ai.a2a.task.state`, `gen_ai.a2a.agent.url`, ...). The interceptor also sends the W3C `traceparent` (as an HTTP header on JSON-RPC and REST, as call -metadata on gRPC), so a remote agent that reads it continues the same trace. +metadata on gRPC). Add `NewServerInterceptor` to the receiving agent to continue the trace +with one server span per call, without copying context extraction into your application. | | | |---|---| @@ -72,7 +73,7 @@ remote agent; that request gets `traceparent` (and `tracestate` when set) only. [`example/main.go`](example/main.go) is a complete program: an agent span, A2A calls (streamed with `-stream`), a failing task (`-fail`), and a demo remote agent built with a2a-go's server -package whose own spans join the trace. When it calls `api.openai.com`, the demo marks its own LLM +package whose maintained server interceptor joins the trace. When it calls `api.openai.com`, the demo marks its own LLM span with `gen_ai.provider.name=openai`; other OpenAI-compatible hosts remain unlabelled rather than guessing a vendor. The A2A interceptor itself still never creates an LLM span or invents a provider. @@ -81,9 +82,39 @@ export FI_API_KEY=... FI_SECRET_KEY=... FI_PROJECT_NAME=my-service go run ./example -stream -fail ``` +## Trace the receiving agent + +Use the same tracer provider and content settings with the official server hook: + +```go +handler := a2asrv.NewHandler(executor, + a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor( + traceaia2a.WithTracerProvider(tp)))) +httpHandler := a2asrv.NewJSONRPCHandler(handler) +``` + +Here `executor` is your `a2asrv.AgentExecutor`. Serve `httpHandler` using `net/http`. +For REST use `a2asrv.NewRESTHandler(handler)`; for gRPC use +`a2agrpc.NewHandler(handler).RegisterWith(grpcServer)` from `a2agrpc/v1`. +Attach the interceptor once, first in the list. It traces protocol calls, not agent-card +requests, transport parsing/authentication failures that occur before the hook, or separate +executor lifetimes. A background task may outlive a nonblocking request's span. + +The executor receives a context containing the server span. Start any application/model +spans from that context; do not re-extract `traceparent` or add a second A2A server span. +An existing local transport span is kept as the parent. `WithPropagator(nil)` disables incoming +extraction; the default extracts W3C trace context only, not baggage. Session/user context helpers +are process-local: the A2A context ID becomes the server's `session.id` unless set on its context. +Headers, authenticated user details and request metadata are never copied to spans. + +The server uses the same privacy settings and stream/error handling as the client. HTTP and gRPC +cancel the request context on disconnect. When calling a handler directly, pass a cancellable +context and cancel it if you stop consuming a stream early: the SDK does not call `After` on +iterator close. + ## What is traced -| Client method | Span | `gen_ai.operation.name` | +| Client and server method | Span | `gen_ai.operation.name` | |---|---|---| | `SendMessage` | `a2a.SendMessage` | `a2a.send` | | `SendStreamingMessage` | `a2a.SendStreamingMessage` | `a2a.send` | @@ -97,8 +128,8 @@ a span. A stream is one span. Every streamed event is a span event named `a2a.event` (with `gen_ai.a2a.event.type`, and the task state or artifact type when the event has one), and the span ends at the agent's reply message or when the task reaches a terminal state (`completed`, -`failed`, `canceled`, `rejected`) or an interrupted one (`input-required`, `auth-required`), which -is when a2a-go closes the stream. +`failed`, `canceled`, `rejected`) or `input-required`, which is when a2a-go closes the stream. `auth-required` does not close +an active stream: later events and its final result remain on the same span. ## Attributes @@ -159,9 +190,6 @@ recorded, values are cut at 32 000 bytes, and the module adds no resource attrib ## Not supported -- Server-side tracing (`a2asrv`). The server interceptor is a different interface and is not part - of this module; a remote agent joins the trace only if it reads `traceparent` itself (the - example's demo agent shows how). - Push notifications and the a2a CLI. ## Test diff --git a/go/a2a/doc.go b/go/a2a/doc.go index 5793b6d1b..b448fb756 100644 --- a/go/a2a/doc.go +++ b/go/a2a/doc.go @@ -1,4 +1,4 @@ -// Package traceaia2a traces A2A (Agent2Agent) protocol client calls made with +// Package traceaia2a traces A2A (Agent2Agent) protocol client and server calls made with // the official Go SDK, github.com/a2aproject/a2a-go/v2, for Future AGI. // // [Interceptor] is an a2aclient.CallInterceptor, the SDK's own extension @@ -27,5 +27,8 @@ // pointed at Future AGI (see the example directory) and pass it with // [WithTracerProvider], or install it as the global provider. // -// Server-side handlers (a2asrv) are not instrumented by this package. +// [ServerInterceptor] attaches to a2asrv.WithCallInterceptors and traces the +// same methods as SERVER spans. It extracts W3C context from incoming service +// parameters and passes the span context to the executor. Existing local +// transport spans remain parents. No application tracing glue is required. package traceaia2a diff --git a/go/a2a/example/main.go b/go/a2a/example/main.go index be8b4b8a7..a0f612d08 100644 --- a/go/a2a/example/main.go +++ b/go/a2a/example/main.go @@ -44,7 +44,6 @@ import ( "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp" - "go.opentelemetry.io/otel/propagation" "go.opentelemetry.io/otel/sdk/resource" sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/trace" @@ -208,11 +207,9 @@ func printEvent(ev a2a.Event) { // startDemoAgent serves a small A2A agent on 127.0.0.1 with a2a-go's own // server package and returns the URL its agent card is served from. // -// The agent stands in for a remote agent. traceaia2a does not instrument -// a2asrv: the agent's spans below come from its own code, which reads the -// traceparent the interceptor sent so they join the caller's trace. A remote -// agent that does the same (or uses any OpenTelemetry HTTP middleware) shows up -// in the same trace; one that does not still leaves the client span. +// The maintained server interceptor extracts the incoming context and traces +// the protocol call. The executor receives that context directly, so any model +// spans it starts are children of the A2A server span; no tracing glue is needed. func startDemoAgent(tp *sdktrace.TracerProvider) (string, func(), error) { ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { @@ -231,7 +228,9 @@ func startDemoAgent(tp *sdktrace.TracerProvider) (string, func(), error) { } agent := &demoAgent{tracer: tp.Tracer("demo-remote-agent"), tc: traceaia2a.TraceConfigFromEnv()} mux := http.NewServeMux() - mux.Handle("/a2a", a2asrv.NewJSONRPCHandler(a2asrv.NewHandler(agent))) + handler := a2asrv.NewHandler(agent, a2asrv.WithCallInterceptors( + traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)))) + mux.Handle("/a2a", a2asrv.NewJSONRPCHandler(handler)) mux.Handle(a2asrv.WellKnownAgentCardPath, a2asrv.NewStaticAgentCardHandler(card)) srv := &http.Server{Handler: mux, ReadHeaderTimeout: 10 * time.Second} go func() { _ = srv.Serve(ln) }() @@ -258,39 +257,22 @@ func (d *demoAgent) Execute(ctx context.Context, ec *a2asrv.ExecutorContext) ite question += p.Text() } } - // Continue the caller's trace from the traceparent it sent. - carrier := propagation.MapCarrier{} - if v, ok := ec.ServiceParams.Get("traceparent"); ok && len(v) > 0 { - carrier.Set("traceparent", v[0]) - } - ctx = propagation.TraceContext{}.Extract(ctx, carrier) - ctx, span := d.tracer.Start(ctx, "trip-helper", trace.WithSpanKind(trace.SpanKindServer), - trace.WithAttributes( - attribute.String(traceaia2a.GenAISpanKind, traceaia2a.SpanKindAgent), - attribute.String("gen_ai.agent.name", "trip-helper"), - attribute.String(traceaia2a.A2ATaskID, string(ec.TaskID)), - )) - defer span.End() - d.tc.SetAttributes(span, attribute.String(traceaia2a.InputValue, question)) if !yield(a2a.NewSubmittedTask(ec, ec.Message), nil) || !yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateWorking, nil), nil) { return } if fail, _ := ec.Metadata["demo.fail"].(bool); fail { - span.SetStatus(codes.Error, "demo failure") yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateFailed, a2a.NewMessageForTask(a2a.MessageRoleAgent, ec, a2a.NewTextPart("could not plan: "+question))), nil) return } answer, err := d.answer(ctx, question) if err != nil { - span.SetStatus(codes.Error, "model call failed") yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateFailed, a2a.NewMessageForTask(a2a.MessageRoleAgent, ec, a2a.NewTextPart("model call failed"))), nil) return } - d.tc.SetAttributes(span, attribute.String(traceaia2a.OutputValue, answer)) _ = yield(a2a.NewArtifactEvent(ec, a2a.NewTextPart(answer)), nil) && yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateCompleted, a2a.NewMessageForTask(a2a.MessageRoleAgent, ec, a2a.NewTextPart("done"))), nil) diff --git a/go/a2a/example/main_test.go b/go/a2a/example/main_test.go index 6e515e049..89ea9b40b 100644 --- a/go/a2a/example/main_test.go +++ b/go/a2a/example/main_test.go @@ -106,8 +106,15 @@ func TestExampleExportsOneStitchedTraceToFutureAGI(t *testing.T) { t.Errorf("want one trace, got %d", len(traces)) } root := byName["trip-planner"] - calls := byName["a2a.SendStreamingMessage"] - remote := byName["trip-helper"] + var calls, remote []*tracepb.Span + for _, span := range byName["a2a.SendStreamingMessage"] { + if span.Kind == tracepb.Span_SPAN_KIND_CLIENT { + calls = append(calls, span) + } + if span.Kind == tracepb.Span_SPAN_KIND_SERVER { + remote = append(remote, span) + } + } if len(root) != 1 || len(calls) != 2 || len(remote) != 2 { t.Fatalf("spans: %d root, %d a2a, %d remote", len(root), len(calls), len(remote)) } diff --git a/go/a2a/interceptor.go b/go/a2a/interceptor.go index 25f23b1d3..e6fd54cfa 100644 --- a/go/a2a/interceptor.go +++ b/go/a2a/interceptor.go @@ -21,7 +21,7 @@ const ScopeName = "github.com/future-agi/traceAI/go/a2a" // version. const Version = "0.1.0" -// Option configures [NewInterceptor]. +// Option configures [NewInterceptor] and [NewServerInterceptor]. type Option func(*settings) type settings struct { @@ -45,9 +45,9 @@ func WithTraceConfig(c TraceConfig) Option { return func(s *settings) { s.tc = &c } } -// WithPropagator sets how trace context is written onto the outbound A2A -// request. The default is W3C Trace Context only (traceparent and -// tracestate), so no baggage leaves the process unless you ask for it. +// WithPropagator sets how trace context is injected by the client and extracted +// by the server. The default is W3C Trace Context only (traceparent and +// tracestate); baggage is neither injected nor extracted unless you ask for it. // Passing nil turns propagation off. func WithPropagator(p propagation.TextMapPropagator) Option { return func(s *settings) { @@ -155,43 +155,8 @@ func (i *Interceptor) Before(ctx context.Context, req *a2aclient.Request) (conte attrs = append(attrs, attribute.String(A2AAgentCardVer, c.Version)) } } - if id, ok := SessionIDFromContext(ctx); ok { - attrs = append(attrs, attribute.String(SessionID, id)) - st.sessionSet = true - } - if id, ok := UserIDFromContext(ctx); ok { - attrs = append(attrs, attribute.String(UserID, id)) - } - - var content []attribute.KeyValue - switch p := req.Payload.(type) { - case *a2a.SendMessageRequest: - if p != nil && p.Message != nil { - m := p.Message - if m.Role != a2a.MessageRoleUnspecified { - attrs = append(attrs, attribute.String(A2AMessageRole, string(m.Role))) - } - attrs = append(attrs, attribute.Int(A2AMessagePartsCnt, len(m.Parts))) - attrs = append(attrs, st.taskInfo(string(m.TaskID), m.ContextID)...) - if text := partsText(m.Parts); text != "" { - content = append(content, - attribute.String(InputValue, truncate(text)), - attribute.String(InputMimeType, MimeTypeText)) - } - } - case *a2a.GetTaskRequest: - if p != nil { - attrs = append(attrs, st.taskInfo(string(p.ID), "")...) - } - case *a2a.CancelTaskRequest: - if p != nil { - attrs = append(attrs, st.taskInfo(string(p.ID), "")...) - } - case *a2a.SubscribeToTaskRequest: - if p != nil { - attrs = append(attrs, st.taskInfo(string(p.ID), "")...) - } - } + requestAttrs, content := st.requestAttributes(ctx, req.Payload) + attrs = append(attrs, requestAttrs...) ctx, span := i.tracer.Start(ctx, "a2a."+req.Method, trace.WithSpanKind(trace.SpanKindClient), @@ -234,24 +199,29 @@ func (i *Interceptor) After(ctx context.Context, resp *a2aclient.Response) error if st == nil || resp.Method != st.method { return nil } + st.recordResponse(resp.Payload, resp.Err) + return nil +} + +// recordResponse shares stream, error and privacy handling between both sides. +func (st *callState) recordResponse(payload any, err error) { st.mu.Lock() defer st.mu.Unlock() if st.ended { - return nil + return } - if resp.Err != nil { - st.fail(resp.Err) - return nil + if err != nil { + st.fail(err) + return } if !st.streaming { - st.applyResult(resp.Payload) + st.applyResult(payload) st.finish() - return nil + return } - if st.applyEvent(resp.Payload) { + if st.applyEvent(payload) { st.finish() } - return nil } // callState is one traced call, shared by Before, every After and the @@ -275,6 +245,49 @@ type callState struct { final string // reply message or final status message text } +// requestAttributes records only protocol facts and applies the same content +// handling to inbound and outbound messages. Headers and metadata are excluded. +func (s *callState) requestAttributes(ctx context.Context, payload any) (attrs, content []attribute.KeyValue) { + if id, ok := SessionIDFromContext(ctx); ok { + attrs = append(attrs, attribute.String(SessionID, id)) + s.sessionSet = true + } + if id, ok := UserIDFromContext(ctx); ok { + attrs = append(attrs, attribute.String(UserID, id)) + } + + switch p := payload.(type) { + case *a2a.SendMessageRequest: + if p != nil && p.Message != nil { + m := p.Message + if m.Role != a2a.MessageRoleUnspecified { + attrs = append(attrs, attribute.String(A2AMessageRole, string(m.Role))) + } + attrs = append(attrs, attribute.Int(A2AMessagePartsCnt, len(m.Parts))) + attrs = append(attrs, s.taskInfo(string(m.TaskID), m.ContextID)...) + if text := partsText(m.Parts); text != "" { + content = append(content, + attribute.String(InputValue, truncate(text)), + attribute.String(InputMimeType, MimeTypeText)) + } + } + case *a2a.GetTaskRequest: + if p != nil { + attrs = append(attrs, s.taskInfo(string(p.ID), "")...) + } + case *a2a.CancelTaskRequest: + if p != nil { + attrs = append(attrs, s.taskInfo(string(p.ID), "")...) + } + case *a2a.SubscribeToTaskRequest: + if p != nil { + attrs = append(attrs, s.taskInfo(string(p.ID), "")...) + } + } + + return attrs, content +} + // taskInfo returns the task and context attributes not yet recorded. The A2A // context id becomes session.id unless the caller set a session id. func (s *callState) taskInfo(taskID, contextID string) []attribute.KeyValue { @@ -411,9 +424,10 @@ func (s *callState) applyEvent(payload any) bool { } // endsStream reports whether a2a-go closes a stream at this state: the -// terminal states and the two interrupted ones. +// terminal states and input-required. Auth-required can be followed by more +// events in the same stream in the supported SDK versions. func endsStream(state a2a.TaskState) bool { - return state.Terminal() || state == a2a.TaskStateInputRequired || state == a2a.TaskStateAuthRequired + return state.Terminal() || state == a2a.TaskStateInputRequired } func (s *callState) output() { diff --git a/go/a2a/server.go b/go/a2a/server.go new file mode 100644 index 000000000..d455660c0 --- /dev/null +++ b/go/a2a/server.go @@ -0,0 +1,119 @@ +package traceaia2a + +import ( + "context" + "strings" + + "github.com/a2aproject/a2a-go/v2/a2asrv" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" +) + +// ServerInterceptor traces incoming A2A calls through the official server hook. +// Attach it once, before other interceptors, with a2asrv.WithCallInterceptors. +// It covers the same five methods as [Interceptor], with one SERVER span per +// call or stream, and never modifies a request, response or error. +// +// The default propagator extracts W3C trace context from HTTP headers or gRPC +// metadata. If transport middleware has already started a local span, that span +// remains the parent so the transport and protocol spans form a single tree. +// Headers, credentials, baggage and authenticated user details are not recorded. +// Use [WithTraceConfig] or FI_HIDE_INPUTS / FI_HIDE_OUTPUTS to hide message content. +// +// a2asrv calls After for each streamed event, not on iterator close. Streams +// therefore end at a terminal/interrupted event or when the request context is +// canceled. Direct in-process callers must cancel their context if they stop +// reading early; HTTP and gRPC transports cancel it when the request ends. +// Detached background task execution can outlive the request span. +// This interceptor does not wrap transport handlers or create executor spans. +// +// handler := a2asrv.NewHandler(executor, +// a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor( +// traceaia2a.WithTracerProvider(tp)))) +// +// Use that handler with a2asrv.NewJSONRPCHandler, a2asrv.NewRESTHandler or +// a2agrpc/v1.NewHandler. +type ServerInterceptor struct { + config *Interceptor +} + +var _ a2asrv.CallInterceptor = (*ServerInterceptor)(nil) + +// NewServerInterceptor returns a server interceptor with the same options and +// privacy defaults as [NewInterceptor]. It does not register global telemetry. +func NewServerInterceptor(opts ...Option) *ServerInterceptor { + return &ServerInterceptor{config: NewInterceptor(opts...)} +} + +// Separate state keys keep outbound calls from an executor from overwriting +// the inbound call's lifecycle state. +type serverStateKey struct{} + +// Before extracts incoming context and starts a protocol SERVER span. +func (i *ServerInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallContext, req *a2asrv.Request) (context.Context, any, error) { + if callCtx == nil || req == nil { + return ctx, nil, nil + } + info, ok := methods[callCtx.Method()] + if !ok { + return ctx, nil, nil + } + if ctx == nil { + ctx = context.Background() + } + config := i.config + parent := trace.SpanContextFromContext(ctx) + if config.prop != nil && (!parent.IsValid() || parent.IsRemote()) { + ctx = config.prop.Extract(ctx, serviceParamsCarrier{callCtx.ServiceParams()}) + } + st := &callState{method: callCtx.Method(), tc: config.tc, streaming: info.streaming} + attrs, content := st.requestAttributes(ctx, req.Payload) + attrs = append(attrs, + attribute.String(GenAISpanKind, SpanKindAgent), + attribute.String(GenAIOperationName, info.op), + attribute.Bool(A2AStreaming, info.streaming)) + ctx, st.span = config.tracer.Start(ctx, "a2a."+st.method, + trace.WithSpanKind(trace.SpanKindServer), trace.WithAttributes(attrs...)) + config.tc.SetAttributes(st.span, content...) + + // Watch every inbound call: the SDK may reject a request after Before + // without running After (e.g. a required extension was not requested). + // Synchronize with an already-canceled request before publishing the stop fn. + spanCtx := ctx + st.mu.Lock() + st.stopWatch = context.AfterFunc(ctx, func() { st.endWithContextError(spanCtx) }) + st.mu.Unlock() + return context.WithValue(ctx, serverStateKey{}, st), nil, nil +} + +// After observes a response without changing protocol behavior. +func (i *ServerInterceptor) After(ctx context.Context, callCtx *a2asrv.CallContext, resp *a2asrv.Response) error { + if ctx == nil || callCtx == nil || resp == nil { + return nil + } + st, _ := ctx.Value(serverStateKey{}).(*callState) + if st != nil && st.method == callCtx.Method() { + st.recordResponse(resp.Payload, resp.Err) + } + return nil +} + +// This read-only carrier uses the SDK's case-insensitive lookup and never +// copies the full header map into span attributes. Joining repeated values +// makes duplicate traceparent headers invalid instead of trusting one of them. +type serviceParamsCarrier struct{ params *a2asrv.ServiceParams } + +func (c serviceParamsCarrier) Get(key string) string { + values, _ := c.params.Get(key) + return strings.Join(values, ",") +} + +func (serviceParamsCarrier) Set(string, string) {} + +func (c serviceParamsCarrier) Keys() []string { + var keys []string + for key := range c.params.List() { + keys = append(keys, key) + } + return keys +} diff --git a/go/a2a/server_lifecycle_test.go b/go/a2a/server_lifecycle_test.go new file mode 100644 index 000000000..2454dd64a --- /dev/null +++ b/go/a2a/server_lifecycle_test.go @@ -0,0 +1,328 @@ +package traceaia2a_test + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/baggage" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/propagation" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "go.opentelemetry.io/otel/trace" +) + +func serverRecorder(t *testing.T) (*sdktrace.TracerProvider, *tracetest.SpanRecorder) { + t.Helper() + rec := tracetest.NewSpanRecorder() + tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(rec)) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + return tp, rec +} + +func serverSpans(rec *tracetest.SpanRecorder) []sdktrace.ReadOnlySpan { + var result []sdktrace.ReadOnlySpan + for _, s := range rec.Ended() { + if s.SpanKind() == trace.SpanKindServer { + result = append(result, s) + } + } + return result +} + +func TestServerTaskLifecycle(t *testing.T) { + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + t.Run(transport, func(t *testing.T) { + tp, rec := serverRecorder(t) + client := serverClient(t, transport, blocks(), tp, traceaia2a.TraceConfig{}) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + req := userMessage("background") + req.Config = &a2a.SendMessageConfig{ReturnImmediately: true} + res, err := client.SendMessage(ctx, req) + if err != nil { + t.Fatal(err) + } + task := res.(*a2a.Task) + if _, err := client.GetTask(ctx, &a2a.GetTaskRequest{ID: task.ID}); err != nil { + t.Fatal(err) + } + n := 0 + for _, err := range client.SubscribeToTask(ctx, &a2a.SubscribeToTaskRequest{ID: task.ID}) { + if err != nil { + t.Fatal(err) + } + n++ + if n == 1 { + if _, err := client.CancelTask(ctx, &a2a.CancelTaskRequest{ID: task.ID}); err != nil { + t.Fatal(err) + } + } + } + spans := serverSpans(rec) + if len(spans) != 4 || len(rec.Ended()) != 8 { + t.Fatalf("want 4 client/server pairs, got %d server, %d total", len(spans), len(rec.Ended())) + } + want := map[string]string{"a2a.SendMessage": "a2a.send", "a2a.GetTask": "a2a.get_task", "a2a.CancelTask": "a2a.cancel_task", "a2a.SubscribeToTask": "a2a.subscribe_to_task"} + for _, s := range spans { + a := attrs(s) + if want[s.Name()] == "" || str(a, traceaia2a.GenAIOperationName) != want[s.Name()] || str(a, traceaia2a.A2ATaskID) != string(task.ID) || s.Status().Code != codes.Ok { + t.Fatalf("unexpected span: %s %v %v", s.Name(), a, s.Status()) + } + delete(want, s.Name()) + } + if len(want) != 0 { + t.Fatalf("missing calls: %v", want) + } + }) + } +} + +func TestServerTerminalStatesAndPrivacy(t *testing.T) { + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + for _, state := range []a2a.TaskState{a2a.TaskStateCompleted, a2a.TaskStateFailed, a2a.TaskStateRejected, a2a.TaskStateCanceled, a2a.TaskStateInputRequired} { + t.Run(transport+"/"+string(state), func(t *testing.T) { + tp, rec := serverRecorder(t) + client := serverClient(t, transport, endsIn(state, "private output "+marker), tp, traceaia2a.TraceConfig{HideInputs: true, HideOutputs: true}) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + for _, err := range client.SendStreamingMessage(ctx, userMessage("private input "+marker)) { + if err != nil { + t.Fatal(err) + } + } + spans := serverSpans(rec) + if len(spans) != 1 { + t.Fatalf("want one server span, got %d", len(spans)) + } + s := spans[0] + a := attrs(s) + if strings.Contains(everything(s), marker) || str(a, traceaia2a.InputValue) != "__REDACTED__" || str(a, traceaia2a.OutputValue) != "__REDACTED__" { + t.Fatalf("content privacy not applied: %v", a) + } + if str(a, traceaia2a.A2ATaskState) != string(state) { + t.Fatal("missing task state") + } + wantError := state == a2a.TaskStateFailed || state == a2a.TaskStateRejected + if (s.Status().Code == codes.Error) != wantError { + t.Fatalf("unexpected status %v for %s", s.Status(), state) + } + }) + } + } +} + +func TestServerProtocolErrors(t *testing.T) { + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + t.Run(transport, func(t *testing.T) { + tp, rec := serverRecorder(t) + client := serverClient(t, transport, completes(), tp, traceaia2a.TraceConfig{HideOutputs: true}) + _, err := client.GetTask(context.Background(), &a2a.GetTaskRequest{ID: "missing"}) + if !errors.Is(err, a2a.ErrTaskNotFound) { + t.Fatalf("protocol error was changed: %v", err) + } + spans := serverSpans(rec) + if len(spans) != 1 { + t.Fatalf("want one server span, got %d", len(spans)) + } + s := spans[0] + if s.Status().Code != codes.Error || str(attrs(s), traceaia2a.ErrorType) != "TASK_NOT_FOUND" || has(attrs(s), traceaia2a.ErrorMessage) { + t.Fatalf("unexpected error span: %v %v", s.Status(), attrs(s)) + } + }) + } +} + +func TestServerClientDisconnectEndsStream(t *testing.T) { + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + t.Run(transport, func(t *testing.T) { + tp, rec := serverRecorder(t) + client := serverClient(t, transport, blocks(), tp, traceaia2a.TraceConfig{}) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + for _, err := range client.SendStreamingMessage(ctx, userMessage("cancel me")) { + if err != nil { + t.Fatal(err) + } + break + } + cancel() + waitEnded(t, rec, 2) + spans := serverSpans(rec) + if len(spans) != 1 || spans[0].Status().Code != codes.Error || str(attrs(spans[0]), traceaia2a.ErrorType) != "canceled" { + t.Fatalf("missing canceled server span: %v", spans) + } + }) + } +} + +// Exercise extraction through the real handler hook without relying on an HTTP +// transport to normalize malformed metadata. Also covers in-process callers. +func TestServerPropagationPolicies(t *testing.T) { + remote := "00-11111111111111111111111111111111-2222222222222222-01" + for _, name := range []string{"w3c", "absent", "invalid", "duplicate", "disabled", "local-parent", "baggage-opt-in"} { + t.Run(name, func(t *testing.T) { + tp, rec := serverRecorder(t) + ctx := context.Background() + params := map[string][]string{"TraceParent": {remote}, "TraceState": {"vendor=value"}, "baggage": {"private=" + marker}, "Authorization": {"Bearer " + marker}} + opts := []traceaia2a.Option{traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(traceaia2a.TraceConfig{})} + var local trace.Span + switch name { + case "absent": + params = nil + case "invalid": + params["TraceParent"] = []string{"bad"} + case "duplicate": + params["TraceParent"] = []string{remote, remote} + case "disabled": + opts = append(opts, traceaia2a.WithPropagator(nil)) + case "local-parent": + ctx, local = tp.Tracer("transport").Start(ctx, "transport") + defer local.End() + case "baggage-opt-in": + opts = append(opts, traceaia2a.WithPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{}))) + } + exec := repliesWithMessage() + inner := exec.execute + exec.execute = func(ctx context.Context, ec *a2asrv.ExecutorContext, yield func(a2a.Event, error) bool) { + if got := baggage.FromContext(ctx).Member("private").Value(); (got == marker) != (name == "baggage-opt-in") { + t.Errorf("unexpected baggage extraction for %s", name) + } + inner(ctx, ec, yield) + } + h := a2asrv.NewHandler(exec, a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor(opts...))) + ctx, _ = a2asrv.NewCallContext(ctx, a2asrv.NewServiceParams(params)) + if _, err := h.SendMessage(ctx, userMessage("hello")); err != nil { + t.Fatal(err) + } + s := serverSpans(rec)[0] + switch name { + case "w3c", "baggage-opt-in": + if s.Parent().SpanID().String() != "2222222222222222" || !s.Parent().IsRemote() || s.SpanContext().TraceState().String() != "vendor=value" { + t.Fatal("W3C context not extracted") + } + case "local-parent": + if s.Parent().SpanID() != local.SpanContext().SpanID() || s.Parent().IsRemote() { + t.Fatal("local transport parent was replaced") + } + default: + if s.Parent().IsValid() { + t.Fatal("unexpected parent") + } + } + if strings.Contains(everything(s), marker) { + t.Fatal("headers or baggage reached span attributes") + } + }) + } +} + +// Rejection by another official interceptor must still finish our span, and +// return precisely the original result/error without creating executor spans. +type rejectingServer struct { + a2asrv.PassthroughCallInterceptor + err error +} + +func (i rejectingServer) Before(ctx context.Context, _ *a2asrv.CallContext, _ *a2asrv.Request) (context.Context, any, error) { + return ctx, nil, i.err +} + +func TestServerInterceptorRejectionAndUntracedMethods(t *testing.T) { + tp, rec := serverRecorder(t) + failure := a2a.NewError(a2a.ErrUnauthorized, "private "+marker) + h := a2asrv.NewHandler(completes(), a2asrv.WithCallInterceptors( + traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(traceaia2a.TraceConfig{HideInputs: true})), rejectingServer{err: failure})) + for _, streaming := range []bool{false, true} { + if streaming { + for _, err := range h.SendStreamingMessage(context.Background(), userMessage(marker)) { + if err != failure { + t.Fatalf("error changed: %v", err) + } + } + } else if _, err := h.SendMessage(context.Background(), userMessage(marker)); err != failure { + t.Fatalf("error changed: %v", err) + } + } + if _, err := h.GetExtendedAgentCard(context.Background(), &a2a.GetExtendedAgentCardRequest{}); err != failure { + t.Fatalf("error changed: %v", err) + } + spans := rec.Ended() + if len(spans) != 2 { + t.Fatalf("want exactly two protocol spans, got %d", len(spans)) + } + for _, s := range spans { + if s.Status().Code != codes.Error || strings.Contains(everything(s), marker) { + t.Fatal("incorrect rejection span") + } + } +} + +func TestServerNestedOutboundCallDoesNotCorruptState(t *testing.T) { + tp, rec := serverRecorder(t) + remote := serverClient(t, "jsonrpc", repliesWithMessage(), tp, traceaia2a.TraceConfig{}) + exec := repliesWithMessage() + inner := exec.execute + exec.execute = func(ctx context.Context, ec *a2asrv.ExecutorContext, yield func(a2a.Event, error) bool) { + if _, err := remote.SendMessage(ctx, userMessage("nested")); err != nil { + yield(nil, err) + return + } + inner(ctx, ec, yield) + } + h := a2asrv.NewHandler(exec, a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)))) + if _, err := h.SendMessage(context.Background(), userMessage("outer")); err != nil { + t.Fatal(err) + } + spans := rec.Ended() + if len(spans) != 3 { + t.Fatalf("want outer server -> client -> remote server, got %d spans", len(spans)) + } + byID := map[trace.SpanID]sdktrace.ReadOnlySpan{} + for _, s := range spans { + byID[s.SpanContext().SpanID()] = s + } + for _, s := range spans { + if s.Parent().IsValid() && byID[s.Parent().SpanID()] == nil { + t.Fatal("broken nested call tree") + } + if s.Status().Code != codes.Ok { + t.Fatal("nested call did not complete") + } + } +} + +func TestServerAuthRequiredDoesNotEndActiveStream(t *testing.T) { + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + t.Run(transport, func(t *testing.T) { + tp, rec := serverRecorder(t) + exec := &agent{execute: func(ctx context.Context, ec *a2asrv.ExecutorContext, yield func(a2a.Event, error) bool) { + _ = yield(a2a.NewSubmittedTask(ec, ec.Message), nil) && + yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateAuthRequired, nil), nil) && + yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateCompleted, a2a.NewMessageForTask(a2a.MessageRoleAgent, ec, a2a.NewTextPart("authenticated"))), nil) + }} + client := serverClient(t, transport, exec, tp, traceaia2a.TraceConfig{}) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + for _, err := range client.SendStreamingMessage(ctx, userMessage("go")) { + if err != nil { + t.Fatal(err) + } + } + if len(rec.Ended()) != 2 { + t.Fatalf("want one client/server pair, got %d", len(rec.Ended())) + } + for _, s := range rec.Ended() { + if str(attrs(s), traceaia2a.A2ATaskState) != string(a2a.TaskStateCompleted) || str(attrs(s), traceaia2a.OutputValue) != "authenticated" || len(s.Events()) != 3 { + t.Fatalf("span ended at auth-required: %s %v", s.Name(), attrs(s)) + } + } + }) + } +} diff --git a/go/a2a/server_otlp_test.go b/go/a2a/server_otlp_test.go new file mode 100644 index 000000000..68126427a --- /dev/null +++ b/go/a2a/server_otlp_test.go @@ -0,0 +1,128 @@ +package traceaia2a_test + +import ( + "compress/gzip" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + collectortrace "go.opentelemetry.io/proto/otlp/collector/trace/v1" + tracepb "go.opentelemetry.io/proto/otlp/trace/v1" + "google.golang.org/protobuf/proto" +) + +// Use the stock exporter, not a hand-built approximation of OTLP attributes. +// JSON is not a supported output format of otlptracehttp; it emits protobuf. +func TestServerOTLPWirePrivacy(t *testing.T) { + for _, compression := range []otlptracehttp.Compression{otlptracehttp.NoCompression, otlptracehttp.GzipCompression} { + name := "plain" + if compression == otlptracehttp.GzipCompression { + name = "gzip" + } + t.Run(name, func(t *testing.T) { + requests := make(chan *collectortrace.ExportTraceServiceRequest, 8) + sink := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/traces" || r.Header.Get("Content-Type") != "application/x-protobuf" { + t.Error("wrong OTLP endpoint or content type") + } + var reader io.Reader = r.Body + if compression == otlptracehttp.GzipCompression { + if r.Header.Get("Content-Encoding") != "gzip" { + t.Error("missing gzip header") + } + gz, err := gzip.NewReader(r.Body) + if err != nil { + t.Error(err) + w.WriteHeader(400) + return + } + defer gz.Close() + reader = gz + } else if r.Header.Get("Content-Encoding") != "" { + t.Error("unexpected compression") + } + body, err := io.ReadAll(reader) + if err != nil { + t.Error(err) + w.WriteHeader(400) + return + } + if strings.Contains(string(body), marker) { + t.Error("private content reached OTLP wire") + } + msg := &collectortrace.ExportTraceServiceRequest{} + if err := proto.Unmarshal(body, msg); err != nil { + t.Error(err) + w.WriteHeader(400) + return + } + requests <- msg + w.Header().Set("Content-Type", "application/x-protobuf") + })) + defer sink.Close() + exp, err := otlptracehttp.New(context.Background(), otlptracehttp.WithEndpointURL(sink.URL+"/v1/traces"), otlptracehttp.WithCompression(compression)) + if err != nil { + t.Fatal(err) + } + tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exp)) + defer func() { _ = tp.Shutdown(context.Background()) }() + client := serverClient(t, "jsonrpc", completes(), tp, traceaia2a.TraceConfig{HideInputs: true, HideOutputs: true}) + for _, err := range client.SendStreamingMessage(context.Background(), userMessage(marker)) { + if err != nil { + t.Fatal(err) + } + } + if _, err := client.GetTask(context.Background(), &a2a.GetTaskRequest{ID: "missing"}); err == nil { + t.Fatal("expected protocol error") + } + if len(requests) != 4 { + t.Fatalf("want two client/server pairs, got %d exports", len(requests)) + } + byID := map[string]*tracepb.Span{} + serverCount := 0 + for len(requests) > 0 { + for _, rs := range (<-requests).ResourceSpans { + for _, ss := range rs.ScopeSpans { + if ss.Scope.Name != traceaia2a.ScopeName { + t.Fatal("wrong instrumentation scope") + } + for _, s := range ss.Spans { + byID[string(s.SpanId)] = s + if s.Kind == tracepb.Span_SPAN_KIND_SERVER { + serverCount++ + } + if strings.Contains(s.String(), marker) { + t.Fatal("private content in decoded OTLP") + } + } + } + } + } + if len(byID) != 4 || serverCount != 2 { + t.Fatalf("unexpected wire span counts: %d total, %d server", len(byID), serverCount) + } + for _, s := range byID { + if s.Kind != tracepb.Span_SPAN_KIND_SERVER { + continue + } + p := byID[string(s.ParentSpanId)] + if p == nil || p.Kind != tracepb.Span_SPAN_KIND_CLIENT || string(p.TraceId) != string(s.TraceId) { + t.Fatal("OTLP parent chain broken") + } + if s.Name == "a2a.SendStreamingMessage" && (len(s.Events) != 4 || s.Status.Code != tracepb.Status_STATUS_CODE_OK) { + t.Fatal("wire stream shape changed") + } + if s.Name == "a2a.GetTask" && s.Status.Code != tracepb.Status_STATUS_CODE_ERROR { + t.Fatal("wire error status missing") + } + } + }) + } +} diff --git a/go/a2a/server_test.go b/go/a2a/server_test.go new file mode 100644 index 000000000..459b31b92 --- /dev/null +++ b/go/a2a/server_test.go @@ -0,0 +1,126 @@ +package traceaia2a_test + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + a2agrpc "github.com/a2aproject/a2a-go/v2/a2agrpc/v1" + "github.com/a2aproject/a2a-go/v2/a2asrv" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/codes" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "go.opentelemetry.io/otel/trace" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" +) + +func serverClient(t *testing.T, transport string, executor *agent, tp *sdktrace.TracerProvider, tc traceaia2a.TraceConfig) *a2aclient.Client { + t.Helper() + handler := a2asrv.NewHandler(executor, a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor( + traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(tc)))) + var c *a2a.AgentCard + opts := []a2aclient.FactoryOption{a2aclient.WithCallInterceptors(traceaia2a.NewInterceptor( + traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(tc)))} + if transport == "grpc" { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + gs := grpc.NewServer() + a2agrpc.NewHandler(handler).RegisterWith(gs) + go func() { _ = gs.Serve(ln) }() + t.Cleanup(gs.Stop) + c = &a2a.AgentCard{SupportedInterfaces: []*a2a.AgentInterface{a2a.NewAgentInterface(ln.Addr().String(), a2a.TransportProtocolGRPC)}, Capabilities: a2a.AgentCapabilities{Streaming: true}} + opts = append(opts, a2agrpc.WithGRPCTransport(grpc.WithTransportCredentials(insecure.NewCredentials()))) + } else { + var h http.Handler = a2asrv.NewJSONRPCHandler(handler) + if transport == "rest" { + h = a2asrv.NewRESTHandler(handler) + } + srv := httptest.NewServer(h) + t.Cleanup(srv.Close) + c = card(srv.URL, transport == "rest", true) + } + client, err := a2aclient.NewFromCard(context.Background(), c, opts...) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = client.Destroy() }) + return client +} + +func TestServerTraceContinuesClient(t *testing.T) { + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + for _, streaming := range []bool{false, true} { + method := "SendMessage" + if streaming { + method = "SendStreamingMessage" + } + t.Run(transport+"/"+method, func(t *testing.T) { + rec := tracetest.NewSpanRecorder() + tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(rec)) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + exec := completes() + inner := exec.execute + exec.execute = func(ctx context.Context, ec *a2asrv.ExecutorContext, yield func(a2a.Event, error) bool) { + _, span := tp.Tracer("test-executor").Start(ctx, "executor") + span.End() + inner(ctx, ec, yield) + } + client := serverClient(t, transport, exec, tp, traceaia2a.TraceConfig{}) + ctx, root := tp.Tracer("test").Start(context.Background(), "root") + if streaming { + for _, err := range client.SendStreamingMessage(ctx, userMessage("hello "+marker)) { + if err != nil { + t.Fatal(err) + } + } + } else if _, err := client.SendMessage(ctx, userMessage("hello "+marker)); err != nil { + t.Fatal(err) + } + root.End() + spans := rec.Ended() + if len(spans) != 4 { + t.Fatalf("want root + client + server + executor (4 spans), got %d", len(spans)) + } + var clientSpan, serverSpan, child sdktrace.ReadOnlySpan + for _, s := range spans { + switch s.SpanKind() { + case trace.SpanKindClient: + clientSpan = s + case trace.SpanKindServer: + serverSpan = s + } + if s.Name() == "executor" { + child = s + } + } + if clientSpan == nil || serverSpan == nil || child == nil { + t.Fatal("missing client, server or executor span") + } + if clientSpan.Parent().SpanID() != root.SpanContext().SpanID() || serverSpan.Parent().SpanID() != clientSpan.SpanContext().SpanID() || child.Parent().SpanID() != serverSpan.SpanContext().SpanID() { + t.Fatal("broken root -> client -> server -> executor parent chain") + } + if !serverSpan.Parent().IsRemote() || serverSpan.SpanContext().TraceID() != root.SpanContext().TraceID() { + t.Fatal("server did not extract remote trace context") + } + if serverSpan.Name() != "a2a."+method || serverSpan.Status().Code != codes.Ok { + t.Fatalf("server name/status: %s %v", serverSpan.Name(), serverSpan.Status()) + } + a := attrs(serverSpan) + if str(a, traceaia2a.InputValue) != "hello "+marker || str(a, traceaia2a.GenAISpanKind) != "AGENT" || str(a, traceaia2a.A2ATaskState) != string(a2a.TaskStateCompleted) { + t.Fatalf("server attributes: %v", a) + } + if streaming && len(serverSpan.Events()) != 4 { + t.Fatalf("want 4 stream events, got %d", len(serverSpan.Events())) + } + }) + } + } +} From 63af829485f236a76790f20df3ccee3a33977d0d Mon Sep 17 00:00:00 2001 From: Nikhil Pareek <8347993+nik13@users.noreply.github.com> Date: Sat, 31 Oct 2026 19:10:01 -0700 Subject: [PATCH 2/6] fix(a2a): preserve unary snapshots and require reliable server hooks Keep unary status-message capture independent of stream termination. Require a2a-go v2.7.0 so required-extension rejections run completion hooks, with regression and OTLP privacy coverage. Co-Authored-By: Claude Code --- .github/workflows/go-ci.yml | 7 --- go/a2a/README.md | 14 ++++-- go/a2a/auth_snapshot_test.go | 80 ++++++++++++++++++++++++++++++++++ go/a2a/extension_error_test.go | 71 ++++++++++++++++++++++++++++++ go/a2a/go.mod | 2 +- go/a2a/go.sum | 4 +- go/a2a/interceptor.go | 4 +- go/a2a/server.go | 6 ++- go/a2a/server_otlp_test.go | 25 +++++++++-- go/a2a/server_test.go | 7 +-- 10 files changed, 197 insertions(+), 23 deletions(-) create mode 100644 go/a2a/auth_snapshot_test.go create mode 100644 go/a2a/extension_error_test.go diff --git a/.github/workflows/go-ci.yml b/.github/workflows/go-ci.yml index 5932d2675..498cf313f 100644 --- a/.github/workflows/go-ci.yml +++ b/.github/workflows/go-ci.yml @@ -64,10 +64,3 @@ jobs: run: go build -o "$RUNNER_TEMP/traceai-example" ./example - name: Render package documentation run: go doc . - - name: Test A2A current SDK compatibility - if: matrix.module == 'a2a' - run: | - # The preceding checks use the v2.5.0 floor from go.mod. - go get github.com/a2aproject/a2a-go/v2@v2.7.0 - go mod tidy - go test -race -count=1 ./... diff --git a/go/a2a/README.md b/go/a2a/README.md index c5284855c..f83a16cd5 100644 --- a/go/a2a/README.md +++ b/go/a2a/README.md @@ -15,7 +15,7 @@ with one server span per call, without copying context extraction into your appl | | | |---|---| | Module | `github.com/future-agi/traceAI/go/a2a` (package `traceaia2a`) | -| A2A SDK | `github.com/a2aproject/a2a-go/v2` v2.5.0 or later (the `go.mod` minimum) | +| A2A SDK | `github.com/a2aproject/a2a-go/v2` v2.7.0 or later (the `go.mod` minimum) | | Go | 1.26.0 or later, the common minimum across traceAI Go modules | | Transports | JSON-RPC, HTTP+JSON (REST) and gRPC (`a2agrpc/v1`), each tested against a2a-go's own server | | Dependencies | the OpenTelemetry API (`go.opentelemetry.io/otel`, `.../trace`) and a2a-go. No exporter is built in. | @@ -84,7 +84,12 @@ go run ./example -stream -fail ## Trace the receiving agent -Use the same tracer provider and content settings with the official server hook: +Use the same tracer provider and content settings with the official server hook. +Keep a2a-go at v2.7.0 or later: earlier versions skip the completion hook on required-extension +rejections, leaving unary server spans open or incorrectly marked as canceled. +The module's dependency minimum selects the fixed SDK automatically. + +Attach the interceptor when creating your handler: ```go handler := a2asrv.NewHandler(executor, @@ -129,7 +134,8 @@ A stream is one span. Every streamed event is a span event named `a2a.event` (wi `gen_ai.a2a.event.type`, and the task state or artifact type when the event has one), and the span ends at the agent's reply message or when the task reaches a terminal state (`completed`, `failed`, `canceled`, `rejected`) or `input-required`, which is when a2a-go closes the stream. `auth-required` does not close -an active stream: later events and its final result remain on the same span. +an active stream: later events and its final result remain on the same span. Unary task snapshots +include their current status message in `output.value`, including an authorization prompt. ## Attributes @@ -147,7 +153,7 @@ an active stream: later events and its final result remain on the same span. | `user.id` | the id from `ContextWithUserID` | | `gen_ai.a2a.stream.event_count` (Go only) | streamed calls | | `input.value` | the text and data parts you send (content) | -| `output.value` | artifact text, then the reply message or the final status message (content) | +| `output.value` | artifact text, then the reply or status message (current snapshot for unary calls, final status for streams; content) | | `error.type`, `error.message` | on errors, see below | Python `traceai-a2a` writes `gen_ai.span.kind=A2A_CLIENT` and lower-case v0.3 state names diff --git a/go/a2a/auth_snapshot_test.go b/go/a2a/auth_snapshot_test.go new file mode 100644 index 000000000..5df579df3 --- /dev/null +++ b/go/a2a/auth_snapshot_test.go @@ -0,0 +1,80 @@ +package traceaia2a_test + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/codes" +) + +// Unary calls return the authorization prompt while the executor is still alive. +// Capturing that snapshot must not depend on whether its state ends a stream. +func TestAuthRequiredUnarySnapshotOutput(t *testing.T) { + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + for _, hide := range []bool{false, true} { + name := transport + "/visible" + if hide { + name = transport + "/hidden" + } + t.Run(name, func(t *testing.T) { + release := make(chan struct{}) + defer close(release) + const challenge = "Please authorize the synthetic calendar " + marker + executor := &agent{execute: func(ctx context.Context, ec *a2asrv.ExecutorContext, yield func(a2a.Event, error) bool) { + if !yield(a2a.NewSubmittedTask(ec, ec.Message), nil) || + !yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateAuthRequired, + a2a.NewMessageForTask(a2a.MessageRoleAgent, ec, a2a.NewTextPart(challenge))), nil) { + return + } + <-release + yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateCompleted, nil), nil) + }} + tp, rec := serverRecorder(t) + client := serverClient(t, transport, executor, tp, traceaia2a.TraceConfig{HideOutputs: hide}) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + result, err := client.SendMessage(ctx, userMessage("start protected task")) + if err != nil { + t.Fatal(err) + } + task, ok := result.(*a2a.Task) + if !ok { + t.Fatalf("expected task, got %T", result) + } + fetched, err := client.GetTask(ctx, &a2a.GetTaskRequest{ID: task.ID}) + if err != nil { + t.Fatal(err) + } + for _, snapshot := range []*a2a.Task{task, fetched} { + if snapshot.Status.State != a2a.TaskStateAuthRequired || snapshot.Status.Message == nil || + len(snapshot.Status.Message.Parts) != 1 || snapshot.Status.Message.Parts[0].Text() != challenge { + t.Fatalf("framework did not return the auth challenge: %+v", snapshot) + } + } + if len(rec.Ended()) != 4 || len(serverSpans(rec)) != 2 { + t.Fatalf("want two completed client/server pairs, got %d spans", len(rec.Ended())) + } + want := challenge + if hide { + want = "__REDACTED__" + } + for _, span := range rec.Ended() { + if got := str(attrs(span), traceaia2a.OutputValue); got != want { + t.Errorf("%s %s output.value=%q, want %q", span.SpanKind(), span.Name(), got, want) + } + if span.Status().Code != codes.Ok || str(attrs(span), traceaia2a.A2ATaskState) != string(a2a.TaskStateAuthRequired) { + t.Errorf("auth-required snapshot must not be an error: %v", span.Status()) + } + if hide && strings.Contains(everything(span), marker) { + t.Error("authorization prompt bypassed output privacy") + } + } + }) + } + } +} diff --git a/go/a2a/extension_error_test.go b/go/a2a/extension_error_test.go new file mode 100644 index 000000000..e4e938b06 --- /dev/null +++ b/go/a2a/extension_error_test.go @@ -0,0 +1,71 @@ +package traceaia2a_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/codes" +) + +func TestServerRequiredExtensionProtocolError(t *testing.T) { + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + t.Run(transport, func(t *testing.T) { + tp, rec := serverRecorder(t) + client := serverClient(t, transport, completes(), tp, traceaia2a.TraceConfig{}, + a2asrv.WithCapabilityChecks(&a2a.AgentCapabilities{Extensions: []a2a.AgentExtension{{URI: "https://example.invalid/required-extension", Required: true}}})) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + _, err := client.GetTask(ctx, &a2a.GetTaskRequest{ID: "synthetic-task"}) + if !errors.Is(err, a2a.ErrExtensionSupportRequired) { + t.Fatalf("unexpected wire response: %v", err) + } + waitEnded(t, rec, 2) + if len(rec.Ended()) != 2 || len(serverSpans(rec)) != 1 { + t.Fatalf("want one client/server pair, got %d spans", len(rec.Ended())) + } + for _, span := range rec.Ended() { + if str(attrs(span), traceaia2a.ErrorType) != "EXTENSION_SUPPORT_REQUIRED" || span.Status().Code != codes.Error { + t.Errorf("%s span must record the protocol rejection, not context cancellation: %v %v", span.SpanKind(), attrs(span), span.Status()) + } + } + }) + } +} + +// With v2.5/v2.6 the SDK skipped After on required-extension rejections, leaving +// these completed unary calls open forever when the context was not canceled. +func TestServerRequiredExtensionEndsUnarySpan(t *testing.T) { + for _, method := range []string{"SendMessage", "GetTask", "CancelTask"} { + t.Run(method, func(t *testing.T) { + tp, rec := serverRecorder(t) + h := a2asrv.NewHandler(completes(), + a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp))), + a2asrv.WithCapabilityChecks(&a2a.AgentCapabilities{Extensions: []a2a.AgentExtension{{URI: "https://example.invalid/required-extension", Required: true}}})) + ctx := context.Background() + var err error + switch method { + case "SendMessage": + _, err = h.SendMessage(ctx, userMessage("start task")) + case "GetTask": + _, err = h.GetTask(ctx, &a2a.GetTaskRequest{ID: "synthetic-task"}) + case "CancelTask": + _, err = h.CancelTask(ctx, &a2a.CancelTaskRequest{ID: "synthetic-task"}) + } + if !errors.Is(err, a2a.ErrExtensionSupportRequired) { + t.Fatalf("unexpected handler result: %v", err) + } + if len(rec.Started()) != 1 || len(rec.Ended()) != 1 { + t.Fatalf("completed unary call left an open span: started=%d ended=%d", len(rec.Started()), len(rec.Ended())) + } + span := rec.Ended()[0] + if span.Name() != "a2a."+method || str(attrs(span), traceaia2a.ErrorType) != "EXTENSION_SUPPORT_REQUIRED" || span.Status().Code != codes.Error { + t.Fatalf("incorrect rejection span: %s %v %v", span.Name(), attrs(span), span.Status()) + } + }) + } +} diff --git a/go/a2a/go.mod b/go/a2a/go.mod index 50bfd6565..607bdcdd6 100644 --- a/go/a2a/go.mod +++ b/go/a2a/go.mod @@ -3,7 +3,7 @@ module github.com/future-agi/traceAI/go/a2a go 1.26.0 require ( - github.com/a2aproject/a2a-go/v2 v2.5.0 + github.com/a2aproject/a2a-go/v2 v2.7.0 go.opentelemetry.io/otel v1.47.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.47.0 go.opentelemetry.io/otel/sdk v1.47.0 diff --git a/go/a2a/go.sum b/go/a2a/go.sum index 76471a02d..eb2fc49b7 100644 --- a/go/a2a/go.sum +++ b/go/a2a/go.sum @@ -1,5 +1,5 @@ -github.com/a2aproject/a2a-go/v2 v2.5.0 h1:ZdcFoxv+nZTUV0i2ue5hES76YCANFPG9vjqd7vK8yWM= -github.com/a2aproject/a2a-go/v2 v2.5.0/go.mod h1:NcRp/ZHxgMzDj12/BteIC2gOjljuEBKaGRfEdJ2lNSI= +github.com/a2aproject/a2a-go/v2 v2.7.0 h1:rwIpzO5O9Mxu/Rm4S2SZfHYzNcvF3FK9aaaTxP35kAo= +github.com/a2aproject/a2a-go/v2 v2.7.0/go.mod h1:rvFIb0sBJl2qmJRYU2fQxI6hhIQksmGGMEhT/uRwWww= github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= diff --git a/go/a2a/interceptor.go b/go/a2a/interceptor.go index e6fd54cfa..8ceecadd9 100644 --- a/go/a2a/interceptor.go +++ b/go/a2a/interceptor.go @@ -344,7 +344,9 @@ func (s *callState) applyTask(t *a2a.Task) { for _, a := range t.Artifacts { s.addArtifact(a) } - if endsStream(t.Status.State) && t.Status.Message != nil { + // A unary snapshot returns its current status message, including an auth + // prompt. Only streaming calls wait for a stream-ending state. + if t.Status.Message != nil && (!s.streaming || endsStream(t.Status.State)) { s.final = partsText(t.Status.Message.Parts) } } diff --git a/go/a2a/server.go b/go/a2a/server.go index d455660c0..d19110edf 100644 --- a/go/a2a/server.go +++ b/go/a2a/server.go @@ -14,6 +14,9 @@ import ( // It covers the same five methods as [Interceptor], with one SERVER span per // call or stream, and never modifies a request, response or error. // +// Requires a2a-go v2.7.0 or later (the module minimum): older SDKs skip After +// when rejecting unary calls for missing required extensions. +// // The default propagator extracts W3C trace context from HTTP headers or gRPC // metadata. If transport middleware has already started a local span, that span // remains the parent so the transport and protocol spans form a single tree. @@ -76,8 +79,7 @@ func (i *ServerInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallCont trace.WithSpanKind(trace.SpanKindServer), trace.WithAttributes(attrs...)) config.tc.SetAttributes(st.span, content...) - // Watch every inbound call: the SDK may reject a request after Before - // without running After (e.g. a required extension was not requested). + // Watch for cancellation even if later middleware short-circuits After. // Synchronize with an already-canceled request before publishing the stop fn. spanCtx := ctx st.mu.Lock() diff --git a/go/a2a/server_otlp_test.go b/go/a2a/server_otlp_test.go index 68126427a..00f874aca 100644 --- a/go/a2a/server_otlp_test.go +++ b/go/a2a/server_otlp_test.go @@ -82,8 +82,16 @@ func TestServerOTLPWirePrivacy(t *testing.T) { if _, err := client.GetTask(context.Background(), &a2a.GetTaskRequest{ID: "missing"}); err == nil { t.Fatal("expected protocol error") } - if len(requests) != 4 { - t.Fatalf("want two client/server pairs, got %d exports", len(requests)) + authClient := serverClient(t, "jsonrpc", endsIn(a2a.TaskStateAuthRequired, marker), tp, traceaia2a.TraceConfig{HideInputs: true, HideOutputs: true}) + result, err := authClient.SendMessage(context.Background(), userMessage("authorize")) + if err != nil { + t.Fatal(err) + } + if task, ok := result.(*a2a.Task); !ok || task.Status.State != a2a.TaskStateAuthRequired { + t.Fatalf("expected auth-required snapshot, got %v", result) + } + if len(requests) != 6 { + t.Fatalf("want three client/server pairs, got %d exports", len(requests)) } byID := map[string]*tracepb.Span{} serverCount := 0 @@ -105,10 +113,21 @@ func TestServerOTLPWirePrivacy(t *testing.T) { } } } - if len(byID) != 4 || serverCount != 2 { + if len(byID) != 6 || serverCount != 3 { t.Fatalf("unexpected wire span counts: %d total, %d server", len(byID), serverCount) } for _, s := range byID { + if s.Name == "a2a.SendMessage" { + redacted := false + for _, attr := range s.Attributes { + if attr.Key == traceaia2a.OutputValue && attr.Value.GetStringValue() == "__REDACTED__" { + redacted = true + } + } + if !redacted || s.Status.Code != tracepb.Status_STATUS_CODE_OK { + t.Fatal("wire auth-required snapshot must retain redacted output and OK status") + } + } if s.Kind != tracepb.Span_SPAN_KIND_SERVER { continue } diff --git a/go/a2a/server_test.go b/go/a2a/server_test.go index 459b31b92..77b7f3678 100644 --- a/go/a2a/server_test.go +++ b/go/a2a/server_test.go @@ -20,10 +20,11 @@ import ( "google.golang.org/grpc/credentials/insecure" ) -func serverClient(t *testing.T, transport string, executor *agent, tp *sdktrace.TracerProvider, tc traceaia2a.TraceConfig) *a2aclient.Client { +func serverClient(t *testing.T, transport string, executor *agent, tp *sdktrace.TracerProvider, tc traceaia2a.TraceConfig, serverOpts ...a2asrv.RequestHandlerOption) *a2aclient.Client { t.Helper() - handler := a2asrv.NewHandler(executor, a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor( - traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(tc)))) + serverOpts = append([]a2asrv.RequestHandlerOption{a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor( + traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(tc)))}, serverOpts...) + handler := a2asrv.NewHandler(executor, serverOpts...) var c *a2a.AgentCard opts := []a2aclient.FactoryOption{a2aclient.WithCallInterceptors(traceaia2a.NewInterceptor( traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(tc)))} From 84798a2eba63d5eaed6870c3655af5c7b9097d74 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek <8347993+nik13@users.noreply.github.com> Date: Sat, 31 Oct 2026 19:48:54 -0700 Subject: [PATCH 3/6] fix(a2a): isolate nested server call completion state Finalize each server call from its context-local state instead of the SDK shared mutable method. Shadow inherited state for untraced nested calls so their responses cannot finish the outer span. Co-Authored-By: Claude Code --- go/a2a/server.go | 13 +++- go/a2a/server_nested_test.go | 142 +++++++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+), 4 deletions(-) create mode 100644 go/a2a/server_nested_test.go diff --git a/go/a2a/server.go b/go/a2a/server.go index d19110edf..c00467989 100644 --- a/go/a2a/server.go +++ b/go/a2a/server.go @@ -54,6 +54,12 @@ type serverStateKey struct{} // Before extracts incoming context and starts a protocol SERVER span. func (i *ServerInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallContext, req *a2asrv.Request) (context.Context, any, error) { + if ctx == nil { + ctx = context.Background() + } + // Even untraced nested calls need a state boundary: their After must not + // finalize the outer call's span inherited through the executor context. + ctx = context.WithValue(ctx, serverStateKey{}, (*callState)(nil)) if callCtx == nil || req == nil { return ctx, nil, nil } @@ -61,9 +67,6 @@ func (i *ServerInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallCont if !ok { return ctx, nil, nil } - if ctx == nil { - ctx = context.Background() - } config := i.config parent := trace.SpanContextFromContext(ctx) if config.prop != nil && (!parent.IsValid() || parent.IsRemote()) { @@ -93,8 +96,10 @@ func (i *ServerInterceptor) After(ctx context.Context, callCtx *a2asrv.CallConte if ctx == nil || callCtx == nil || resp == nil { return nil } + // The SDK reuses CallContext for nested in-process calls, mutating Method. + // Only the context returned by Before identifies this call's state. st, _ := ctx.Value(serverStateKey{}).(*callState) - if st != nil && st.method == callCtx.Method() { + if st != nil { st.recordResponse(resp.Payload, resp.Err) } return nil diff --git a/go/a2a/server_nested_test.go b/go/a2a/server_nested_test.go new file mode 100644 index 000000000..62451fa68 --- /dev/null +++ b/go/a2a/server_nested_test.go @@ -0,0 +1,142 @@ +package traceaia2a_test + +import ( + "context" + "errors" + "iter" + "strings" + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/codes" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/trace" +) + +// The SDK reuses and mutates CallContext for nested in-process calls. Completion +// must follow each call's context-local state, including calls we do not trace. +func TestServerNestedInProcessCalls(t *testing.T) { + for _, transport := range []string{"direct", "jsonrpc", "rest", "grpc"} { + for _, method := range []string{"SendMessage", "SendStreamingMessage"} { + for _, nested := range []string{"traced", "uninstrumented", "untraced-error", "untraced-success", "same-method"} { + for _, privacy := range []string{"visible", "redacted"} { + t.Run(transport+"/"+method+"/"+nested+"/"+privacy, func(t *testing.T) { + tp, rec := serverRecorder(t) + tc := traceaia2a.TraceConfig{HideInputs: privacy == "redacted", HideOutputs: privacy == "redacted"} + interceptor := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(tc)) + var opts []a2asrv.RequestHandlerOption + if nested != "uninstrumented" { + opts = append(opts, a2asrv.WithCallInterceptors(interceptor)) + } + if nested == "untraced-success" { + opts = append(opts, a2asrv.WithExtendedAgentCard(&a2a.AgentCard{Name: "nested " + marker})) + } + inner := a2asrv.NewHandler(repliesWithMessage(), opts...) + exec := repliesWithMessage() + reply := exec.execute + exec.execute = func(ctx context.Context, ec *a2asrv.ExecutorContext, yield func(a2a.Event, error) bool) { + var err, wantErr error + switch nested { + case "untraced-error": + _, err = inner.GetExtendedAgentCard(ctx, &a2a.GetExtendedAgentCardRequest{}) + wantErr = a2a.ErrExtendedCardNotConfigured + case "untraced-success": + _, err = inner.GetExtendedAgentCard(ctx, &a2a.GetExtendedAgentCardRequest{}) + case "same-method": + _, err = inner.SendMessage(ctx, userMessage("nested "+marker)) + default: + _, err = inner.GetTask(ctx, &a2a.GetTaskRequest{ID: "synthetic-missing"}) + wantErr = a2a.ErrTaskNotFound + } + if !errors.Is(err, wantErr) { + t.Errorf("nested call error = %v, want %v", err, wantErr) + } + if !trace.SpanFromContext(ctx).IsRecording() { + t.Error("nested call ended the outer server span before its response") + } + reply(ctx, ec, yield) + } + + var send func(context.Context, *a2a.SendMessageRequest) (a2a.SendMessageResult, error) + var stream func(context.Context, *a2a.SendMessageRequest) iter.Seq2[a2a.Event, error] + wantSpans := 1 + if transport == "direct" { + handler := a2asrv.NewHandler(exec, a2asrv.WithCallInterceptors(interceptor)) + send, stream = handler.SendMessage, handler.SendStreamingMessage + } else { + client := serverClient(t, transport, exec, tp, tc) + send, stream = client.SendMessage, client.SendStreamingMessage + wantSpans++ + } + if nested == "traced" || nested == "same-method" { + wantSpans++ + } + // No cancellation: a completed direct call must end its own span. + ctx := context.Background() + req := userMessage("outer " + marker) + if method == "SendStreamingMessage" { + events := 0 + for event, err := range stream(ctx, req) { + if err != nil { + t.Fatal(err) + } + if _, ok := event.(*a2a.Message); !ok { + t.Fatalf("unexpected event %T", event) + } + events++ + } + if events != 1 { + t.Fatalf("want one message event, got %d", events) + } + } else { + result, err := send(ctx, req) + if err != nil { + t.Fatal(err) + } + if _, ok := result.(*a2a.Message); !ok { + t.Fatalf("unexpected result %T", result) + } + } + if transport != "direct" { + waitEnded(t, rec, wantSpans) + } + if len(rec.Started()) != wantSpans || len(rec.Ended()) != wantSpans { + t.Fatalf("want %d finished spans: started=%d ended=%d", wantSpans, len(rec.Started()), len(rec.Ended())) + } + var outer, child sdktrace.ReadOnlySpan + for _, span := range serverSpans(rec) { + // Nested spans finish before their parent. + child, outer = outer, span + } + if outer == nil || outer.Name() != "a2a."+method || outer.Status().Code != codes.Ok { + t.Fatalf("successful outer call has incorrect span: %v", outer) + } + if child != nil { + if child.Parent().SpanID() != outer.SpanContext().SpanID() || child.SpanContext().TraceID() != outer.SpanContext().TraceID() || child.Parent().IsRemote() { + t.Error("nested server span lost its local parent") + } + if nested == "traced" && (child.Name() != "a2a.GetTask" || child.Status().Code != codes.Error || str(attrs(child), traceaia2a.ErrorType) != "TASK_NOT_FOUND") { + t.Error("nested error was not recorded on its own span") + } + } + input, output := "outer "+marker, "hello outer "+marker + if privacy == "redacted" { + input, output = "__REDACTED__", "__REDACTED__" + for _, span := range rec.Ended() { + if strings.Contains(everything(span), marker) { + t.Error("nested tracing leaked redacted content") + } + } + } + a := attrs(outer) + if str(a, traceaia2a.InputValue) != input || str(a, traceaia2a.OutputValue) != output || has(a, traceaia2a.ErrorType) { + t.Errorf("outer response was lost or replaced by nested response: %v", a) + } + }) + } + } + } + } +} From b1f56320f75006f59524ed4c38b87648788ad682 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek <8347993+nik13@users.noreply.github.com> Date: Sat, 31 Oct 2026 20:40:02 -0700 Subject: [PATCH 4/6] fix(a2a): observe server response middleware rejections Wrap response-rejecting middleware so tracing sees its actual error before the SDK short-circuits remaining callbacks. Preserve policy ordering, response mutation and error identity, with real SDK, nested-call and OTLP regression coverage. Resolve exporter dependencies in the setup instructions. Co-Authored-By: Claude Code --- go/a2a/README.md | 20 +++++ go/a2a/doc.go | 4 +- go/a2a/server.go | 27 ++++++ go/a2a/server_nested_test.go | 30 +++++-- go/a2a/server_otlp_test.go | 51 +++++++++-- go/a2a/server_policy_test.go | 164 +++++++++++++++++++++++++++++++++++ 6 files changed, 281 insertions(+), 15 deletions(-) create mode 100644 go/a2a/server_policy_test.go diff --git a/go/a2a/README.md b/go/a2a/README.md index f83a16cd5..947f74f72 100644 --- a/go/a2a/README.md +++ b/go/a2a/README.md @@ -32,8 +32,12 @@ differences. ```sh go get github.com/future-agi/traceAI/go/a2a@latest +go get go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp@v1.47.0 go.opentelemetry.io/otel/sdk@v1.47.0 ``` +The second command installs the exporter and SDK used by the setup below and resolves their +`go.sum` entries. The integration itself does not create an exporter. + ## Use ```go @@ -105,6 +109,22 @@ Attach the interceptor once, first in the list. It traces protocol calls, not ag requests, transport parsing/authentication failures that occur before the hook, or separate executor lifetimes. A background task may outlive a nonblocking request's span. +When you add middleware that can reject a response by returning an error from `After`, +wrap each such interceptor with `tracing.Wrap`. For example, with your existing `outputPolicy`: + +```go +tracing := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)) +handler := a2asrv.NewHandler(executor, + a2asrv.WithCallInterceptors(tracing, tracing.Wrap(outputPolicy))) +``` + +The SDK stops the remaining `After` callbacks on rejection. The wrapper records the error +and ends the server span before returning the same error; it does not change the response, +run skipped policy callbacks or add another span. Without the wrapper, a rejected call can +leave its server span open until cancellation, which cannot recover the policy error. +Middleware that instead sets `resp.Err`, clears `resp.Payload` and returns `nil` does not +short-circuit the callbacks and needs no wrapper. + The executor receives a context containing the server span. Start any application/model spans from that context; do not re-extract `traceparent` or add a second A2A server span. An existing local transport span is kept as the parent. `WithPropagator(nil)` disables incoming diff --git a/go/a2a/doc.go b/go/a2a/doc.go index b448fb756..672016dba 100644 --- a/go/a2a/doc.go +++ b/go/a2a/doc.go @@ -30,5 +30,7 @@ // [ServerInterceptor] attaches to a2asrv.WithCallInterceptors and traces the // same methods as SERVER spans. It extracts W3C context from incoming service // parameters and passes the span context to the executor. Existing local -// transport spans remain parents. No application tracing glue is required. +// transport spans remain parents. Use [ServerInterceptor.Wrap] for additional +// middleware that can reject a response by returning an error from After, so +// the rejection ends the server span even when the SDK skips subsequent hooks. package traceaia2a diff --git a/go/a2a/server.go b/go/a2a/server.go index c00467989..100d13057 100644 --- a/go/a2a/server.go +++ b/go/a2a/server.go @@ -13,6 +13,9 @@ import ( // Attach it once, before other interceptors, with a2asrv.WithCallInterceptors. // It covers the same five methods as [Interceptor], with one SERVER span per // call or stream, and never modifies a request, response or error. +// Use [ServerInterceptor.Wrap] for other interceptors that can return an After +// error: the SDK otherwise skips the tracing After, leaving only cancellation +// to finish the span without the rejection error. // // Requires a2a-go v2.7.0 or later (the module minimum): older SDKs skip After // when rejecting unary calls for missing required extensions. @@ -48,6 +51,30 @@ func NewServerInterceptor(opts ...Option) *ServerInterceptor { return &ServerInterceptor{config: NewInterceptor(opts...)} } +// Wrap observes errors returned by another interceptor's After, which otherwise +// stop the SDK from calling the tracing After. Register i first, then i.Wrap(policy) +// in the same a2asrv.WithCallInterceptors list. Wrap each interceptor that can +// reject a response this way. The wrapped interceptor's Before, response changes, +// returned error and SDK callback ordering are preserved; no extra span is created. +func (i *ServerInterceptor) Wrap(interceptor a2asrv.CallInterceptor) a2asrv.CallInterceptor { + return &serverResponseObserver{CallInterceptor: interceptor, tracing: i} +} + +type serverResponseObserver struct { + a2asrv.CallInterceptor + tracing *ServerInterceptor +} + +func (o *serverResponseObserver) After(ctx context.Context, callCtx *a2asrv.CallContext, resp *a2asrv.Response) error { + err := o.CallInterceptor.After(ctx, callCtx, resp) + if err != nil { + // The SDK discards the payload on rejection. Observe that outcome without + // changing the response or allowing earlier policy callbacks to run. + _ = o.tracing.After(ctx, callCtx, &a2asrv.Response{Err: err}) + } + return err +} + // Separate state keys keep outbound calls from an executor from overwriting // the inbound call's lifecycle state. type serverStateKey struct{} diff --git a/go/a2a/server_nested_test.go b/go/a2a/server_nested_test.go index 62451fa68..e3fa3302b 100644 --- a/go/a2a/server_nested_test.go +++ b/go/a2a/server_nested_test.go @@ -20,7 +20,7 @@ import ( func TestServerNestedInProcessCalls(t *testing.T) { for _, transport := range []string{"direct", "jsonrpc", "rest", "grpc"} { for _, method := range []string{"SendMessage", "SendStreamingMessage"} { - for _, nested := range []string{"traced", "uninstrumented", "untraced-error", "untraced-success", "same-method"} { + for _, nested := range []string{"traced", "uninstrumented", "untraced-error", "untraced-success", "same-method", "policy-error", "untraced-policy-error"} { for _, privacy := range []string{"visible", "redacted"} { t.Run(transport+"/"+method+"/"+nested+"/"+privacy, func(t *testing.T) { tp, rec := serverRecorder(t) @@ -30,9 +30,15 @@ func TestServerNestedInProcessCalls(t *testing.T) { if nested != "uninstrumented" { opts = append(opts, a2asrv.WithCallInterceptors(interceptor)) } - if nested == "untraced-success" { + if nested == "untraced-success" || nested == "untraced-policy-error" { opts = append(opts, a2asrv.WithExtendedAgentCard(&a2a.AgentCard{Name: "nested " + marker})) } + if nested == "policy-error" || nested == "untraced-policy-error" { + policy := responsePolicy{after: func(context.Context, *a2asrv.CallContext, *a2asrv.Response) error { + return a2a.NewError(a2a.ErrUnsupportedOperation, "nested policy "+marker) + }} + opts = append(opts, a2asrv.WithCallInterceptors(interceptor.Wrap(policy))) + } inner := a2asrv.NewHandler(repliesWithMessage(), opts...) exec := repliesWithMessage() reply := exec.execute @@ -42,13 +48,19 @@ func TestServerNestedInProcessCalls(t *testing.T) { case "untraced-error": _, err = inner.GetExtendedAgentCard(ctx, &a2a.GetExtendedAgentCardRequest{}) wantErr = a2a.ErrExtendedCardNotConfigured - case "untraced-success": + case "untraced-success", "untraced-policy-error": _, err = inner.GetExtendedAgentCard(ctx, &a2a.GetExtendedAgentCardRequest{}) + if nested == "untraced-policy-error" { + wantErr = a2a.ErrUnsupportedOperation + } case "same-method": _, err = inner.SendMessage(ctx, userMessage("nested "+marker)) default: _, err = inner.GetTask(ctx, &a2a.GetTaskRequest{ID: "synthetic-missing"}) wantErr = a2a.ErrTaskNotFound + if nested == "policy-error" { + wantErr = a2a.ErrUnsupportedOperation + } } if !errors.Is(err, wantErr) { t.Errorf("nested call error = %v, want %v", err, wantErr) @@ -70,7 +82,7 @@ func TestServerNestedInProcessCalls(t *testing.T) { send, stream = client.SendMessage, client.SendStreamingMessage wantSpans++ } - if nested == "traced" || nested == "same-method" { + if nested == "traced" || nested == "same-method" || nested == "policy-error" { wantSpans++ } // No cancellation: a completed direct call must end its own span. @@ -117,8 +129,14 @@ func TestServerNestedInProcessCalls(t *testing.T) { if child.Parent().SpanID() != outer.SpanContext().SpanID() || child.SpanContext().TraceID() != outer.SpanContext().TraceID() || child.Parent().IsRemote() { t.Error("nested server span lost its local parent") } - if nested == "traced" && (child.Name() != "a2a.GetTask" || child.Status().Code != codes.Error || str(attrs(child), traceaia2a.ErrorType) != "TASK_NOT_FOUND") { - t.Error("nested error was not recorded on its own span") + if nested == "traced" || nested == "policy-error" { + wantType := "TASK_NOT_FOUND" + if nested == "policy-error" { + wantType = "UNSUPPORTED_OPERATION" + } + if child.Name() != "a2a.GetTask" || child.Status().Code != codes.Error || str(attrs(child), traceaia2a.ErrorType) != wantType { + t.Error("nested error was not recorded on its own span") + } } } input, output := "outer "+marker, "hello outer "+marker diff --git a/go/a2a/server_otlp_test.go b/go/a2a/server_otlp_test.go index 00f874aca..027176efd 100644 --- a/go/a2a/server_otlp_test.go +++ b/go/a2a/server_otlp_test.go @@ -3,6 +3,7 @@ package traceaia2a_test import ( "compress/gzip" "context" + "errors" "io" "net/http" "net/http/httptest" @@ -10,6 +11,7 @@ import ( "testing" "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" traceaia2a "github.com/future-agi/traceAI/go/a2a" "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp" sdktrace "go.opentelemetry.io/otel/sdk/trace" @@ -27,7 +29,7 @@ func TestServerOTLPWirePrivacy(t *testing.T) { name = "gzip" } t.Run(name, func(t *testing.T) { - requests := make(chan *collectortrace.ExportTraceServiceRequest, 8) + requests := make(chan *collectortrace.ExportTraceServiceRequest, 10) sink := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/traces" || r.Header.Get("Content-Type") != "application/x-protobuf" { t.Error("wrong OTLP endpoint or content type") @@ -90,8 +92,22 @@ func TestServerOTLPWirePrivacy(t *testing.T) { if task, ok := result.(*a2a.Task); !ok || task.Status.State != a2a.TaskStateAuthRequired { t.Fatalf("expected auth-required snapshot, got %v", result) } - if len(requests) != 6 { - t.Fatalf("want three client/server pairs, got %d exports", len(requests)) + tracing := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)) + policy := responsePolicy{after: func(context.Context, *a2asrv.CallContext, *a2asrv.Response) error { + return a2a.NewError(a2a.ErrUnsupportedOperation, "policy rejection "+marker) + }} + policyClient := serverClient(t, "jsonrpc", repliesWithMessage(), tp, traceaia2a.TraceConfig{HideInputs: true, HideOutputs: true}, + a2asrv.WithCallInterceptors(tracing.Wrap(policy))) + if _, err := policyClient.SendMessage(context.Background(), userMessage(marker)); !errors.Is(err, a2a.ErrUnsupportedOperation) { + t.Fatalf("want policy rejection, got %v", err) + } + for _, err := range policyClient.SendStreamingMessage(context.Background(), userMessage(marker)) { + if !errors.Is(err, a2a.ErrUnsupportedOperation) { + t.Fatalf("want streamed policy rejection, got %v", err) + } + } + if len(requests) != 10 { + t.Fatalf("want five client/server pairs, got %d exports", len(requests)) } byID := map[string]*tracepb.Span{} serverCount := 0 @@ -113,10 +129,30 @@ func TestServerOTLPWirePrivacy(t *testing.T) { } } } - if len(byID) != 6 || serverCount != 3 { + if len(byID) != 10 || serverCount != 5 { t.Fatalf("unexpected wire span counts: %d total, %d server", len(byID), serverCount) } + policyErrors := 0 for _, s := range byID { + if s.Kind == tracepb.Span_SPAN_KIND_SERVER { + p := byID[string(s.ParentSpanId)] + if p == nil || p.Kind != tracepb.Span_SPAN_KIND_CLIENT || string(p.TraceId) != string(s.TraceId) { + t.Fatal("OTLP parent chain broken") + } + } + rejected := false + for _, attr := range s.Attributes { + if attr.Key == traceaia2a.ErrorType && attr.Value.GetStringValue() == "UNSUPPORTED_OPERATION" { + rejected = true + } + } + if rejected { + policyErrors++ + if s.Status.Code != tracepb.Status_STATUS_CODE_ERROR || len(s.Events) != 0 { + t.Fatal("policy rejection lost or rejected output recorded") + } + continue + } if s.Name == "a2a.SendMessage" { redacted := false for _, attr := range s.Attributes { @@ -131,10 +167,6 @@ func TestServerOTLPWirePrivacy(t *testing.T) { if s.Kind != tracepb.Span_SPAN_KIND_SERVER { continue } - p := byID[string(s.ParentSpanId)] - if p == nil || p.Kind != tracepb.Span_SPAN_KIND_CLIENT || string(p.TraceId) != string(s.TraceId) { - t.Fatal("OTLP parent chain broken") - } if s.Name == "a2a.SendStreamingMessage" && (len(s.Events) != 4 || s.Status.Code != tracepb.Status_STATUS_CODE_OK) { t.Fatal("wire stream shape changed") } @@ -142,6 +174,9 @@ func TestServerOTLPWirePrivacy(t *testing.T) { t.Fatal("wire error status missing") } } + if policyErrors != 4 { + t.Fatalf("want four client/server policy errors, got %d", policyErrors) + } }) } } diff --git a/go/a2a/server_policy_test.go b/go/a2a/server_policy_test.go new file mode 100644 index 000000000..0e70236b0 --- /dev/null +++ b/go/a2a/server_policy_test.go @@ -0,0 +1,164 @@ +package traceaia2a_test + +import ( + "context" + "errors" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/codes" +) + +type responsePolicy struct { + a2asrv.PassthroughCallInterceptor + after func(context.Context, *a2asrv.CallContext, *a2asrv.Response) error +} + +func (p responsePolicy) After(ctx context.Context, callCtx *a2asrv.CallContext, resp *a2asrv.Response) error { + return p.after(ctx, callCtx, resp) +} + +// After can reject a response by returning an error. The SDK then skips earlier +// After callbacks, so the wrapper must finish tracing without changing that order. +func TestServerResponsePolicy(t *testing.T) { + for _, transport := range []string{"direct", "jsonrpc", "rest", "grpc"} { + for _, config := range []struct { + name string + tc traceaia2a.TraceConfig + }{ + {"visible", traceaia2a.TraceConfig{}}, + {"hide-input", traceaia2a.TraceConfig{HideInputs: true}}, + {"hide-output", traceaia2a.TraceConfig{HideOutputs: true}}, + } { + for _, test := range []struct { + name string + streaming bool + rejectAt int32 + }{ + {"unary-rejection", false, 1}, + {"stream-rejection", true, 1}, + {"stream-late-rejection", true, 3}, + {"unary-allowed", false, 0}, + {"stream-allowed", true, 0}, + } { + t.Run(transport+"/"+config.name+"/"+test.name, func(t *testing.T) { + tp, rec := serverRecorder(t) + tracing := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(config.tc)) + rejection := a2a.NewError(a2a.ErrUnsupportedOperation, "output policy "+marker) + var earlier, policy, later atomic.Int32 + count := func(n *atomic.Int32) a2asrv.CallInterceptor { + return responsePolicy{after: func(context.Context, *a2asrv.CallContext, *a2asrv.Response) error { + n.Add(1) + return nil + }} + } + outputPolicy := responsePolicy{after: func(_ context.Context, _ *a2asrv.CallContext, resp *a2asrv.Response) error { + if resp.Payload == nil || resp.Err != nil { + t.Error("policy did not receive the original successful response") + } + if policy.Add(1) == test.rejectAt { + return rejection + } + return nil + }} + policies := a2asrv.WithCallInterceptors(tracing.Wrap(count(&earlier)), tracing.Wrap(outputPolicy), tracing.Wrap(count(&later))) + exec := repliesWithMessage() + if test.rejectAt > 1 { + exec = completes() + } + handler := a2asrv.NewHandler(exec, a2asrv.WithCallInterceptors(tracing), policies) + send, stream := handler.SendMessage, handler.SendStreamingMessage + expected := 1 + if transport != "direct" { + client := serverClient(t, transport, exec, tp, config.tc, policies) + send, stream = client.SendMessage, client.SendStreamingMessage + expected = 2 + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + var got error + if test.streaming { + for _, err := range stream(ctx, userMessage("request "+marker)) { + got = err + } + } else { + _, got = send(ctx, userMessage("request "+marker)) + } + if test.rejectAt == 0 { + if got != nil { + t.Fatal(got) + } + } else if !errors.Is(got, a2a.ErrUnsupportedOperation) || (transport == "direct" && got != rejection) { + t.Fatalf("policy rejection changed: %v", got) + } + // Assert before canceling: cancellation must not mask a leaked span. + if len(rec.Started()) != expected || len(rec.Ended()) != expected { + t.Fatalf("completed call: started=%d ended=%d, want %d", len(rec.Started()), len(rec.Ended()), expected) + } + wantCalls, wantEarlier := int32(1), int32(1) + if test.rejectAt > 0 { + wantCalls, wantEarlier = test.rejectAt, test.rejectAt-1 + } + if policy.Load() != wantCalls || later.Load() != wantCalls || earlier.Load() != wantEarlier { + t.Fatalf("SDK callback order changed: earlier=%d policy=%d later=%d", earlier.Load(), policy.Load(), later.Load()) + } + spans := serverSpans(rec) + if len(spans) != 1 { + t.Fatalf("want one server span, got %d", len(spans)) + } + s := spans[0] + a := attrs(s) + if test.rejectAt > 0 { + if s.Status().Code != codes.Error || str(a, traceaia2a.ErrorType) != "UNSUPPORTED_OPERATION" { + t.Fatalf("rejection lost: status=%v attrs=%v", s.Status(), a) + } + if !config.tc.HideInputs && !config.tc.HideOutputs && !strings.Contains(str(a, traceaia2a.ErrorMessage), marker) { + t.Fatal("visible protocol error text was lost") + } + if str(a, traceaia2a.OutputValue) != "" || (test.streaming && len(s.Events()) != int(test.rejectAt)-1) { + t.Fatal("rejected output was captured as a successful response") + } + } else if s.Status().Code != codes.Ok { + t.Fatalf("allowed response has status %v", s.Status()) + } + if (config.tc.HideInputs || config.tc.HideOutputs) && str(a, traceaia2a.ErrorMessage) != "" { + t.Fatal("hidden protocol error text leaked") + } + }) + } + } + } +} + +type earlyResponsePolicy struct { + a2asrv.PassthroughCallInterceptor + result *a2a.Message +} + +type policyContextKey struct{} + +func (p earlyResponsePolicy) Before(ctx context.Context, _ *a2asrv.CallContext, req *a2asrv.Request) (context.Context, any, error) { + req.Payload = p.result + return context.WithValue(ctx, policyContextKey{}, p.result), p.result, a2a.ErrUnsupportedOperation +} + +func TestServerWrapPreservesBefore(t *testing.T) { + tp, rec := serverRecorder(t) + tracing := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)) + result := a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("synthetic result")) + wrapped := tracing.Wrap(earlyResponsePolicy{result: result}) + req := &a2asrv.Request{} + ctx, got, err := wrapped.Before(context.Background(), nil, req) + if got != result || err != a2a.ErrUnsupportedOperation || req.Payload != result || ctx.Value(policyContextKey{}) != result { + t.Fatal("wrapped Before changed context, early response, error or request mutation") + } + resp := &a2asrv.Response{Payload: result} + if err := wrapped.After(ctx, nil, resp); err != nil || resp.Payload != result || resp.Err != nil || len(rec.Started()) != 0 { + t.Fatal("wrapped pass-through changed response or started an extra span") + } +} From 942a90783c9c89a5668691e468fc6beada1751aa Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 07:17:49 -0800 Subject: [PATCH 5/6] fix(a2a): honor error privacy and complete snapshot streams at EOF Derive error-detail hiding from content flags with an explicit override. Add maintained client and handler completion wrappers because SDK interceptors do not receive iterator EOF; keep active auth-required streams open and distinguish early abandonment. Cover privacy on actual OTLP exports and completion across all three transports. Co-Authored-By: Claude Code --- go/a2a/README.md | 62 +++++++--- go/a2a/client_snapshot_test.go | 46 +++++++ go/a2a/completion_regression_test.go | 95 ++++++++++++++ go/a2a/doc.go | 5 +- go/a2a/error_privacy_test.go | 90 ++++++++++++++ go/a2a/error_privacy_wire_test.go | 80 ++++++++++++ go/a2a/example/main.go | 9 +- go/a2a/interceptor.go | 36 ++++-- go/a2a/interceptor_test.go | 6 +- go/a2a/internal_test.go | 2 +- go/a2a/server.go | 13 +- go/a2a/server_test.go | 6 +- go/a2a/stream.go | 142 +++++++++++++++++++++ go/a2a/stream_completion_test.go | 179 +++++++++++++++++++++++++++ go/a2a/traceconfig.go | 40 ++++-- 15 files changed, 755 insertions(+), 56 deletions(-) create mode 100644 go/a2a/client_snapshot_test.go create mode 100644 go/a2a/completion_regression_test.go create mode 100644 go/a2a/error_privacy_test.go create mode 100644 go/a2a/error_privacy_wire_test.go create mode 100644 go/a2a/stream.go create mode 100644 go/a2a/stream_completion_test.go diff --git a/go/a2a/README.md b/go/a2a/README.md index 947f74f72..bdec5c828 100644 --- a/go/a2a/README.md +++ b/go/a2a/README.md @@ -42,9 +42,15 @@ The second command installs the exporter and SDK used by the setup below and res ```go // tp is an *sdktrace.TracerProvider that exports to Future AGI (see below). -client, err := a2aclient.NewFromCard(ctx, card, +sdkClient, err := a2aclient.NewFromCard(ctx, card, a2aclient.WithCallInterceptors(traceaia2a.NewInterceptor(traceaia2a.WithTracerProvider(tp)))) +if err != nil { + return err +} +client := traceaia2a.WrapClient(sdkClient) +defer client.Destroy() + ctx = traceaia2a.ContextWithSessionID(ctx, "session-123") // optional resp, err := client.SendMessage(ctx, &a2a.SendMessageRequest{Message: msg}) ``` @@ -96,9 +102,9 @@ The module's dependency minimum selects the fixed SDK automatically. Attach the interceptor when creating your handler: ```go -handler := a2asrv.NewHandler(executor, +handler := traceaia2a.WrapHandler(a2asrv.NewHandler(executor, a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor( - traceaia2a.WithTracerProvider(tp)))) + traceaia2a.WithTracerProvider(tp))))) httpHandler := a2asrv.NewJSONRPCHandler(handler) ``` @@ -114,8 +120,8 @@ wrap each such interceptor with `tracing.Wrap`. For example, with your existing ```go tracing := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)) -handler := a2asrv.NewHandler(executor, - a2asrv.WithCallInterceptors(tracing, tracing.Wrap(outputPolicy))) +handler := traceaia2a.WrapHandler(a2asrv.NewHandler(executor, + a2asrv.WithCallInterceptors(tracing, tracing.Wrap(outputPolicy)))) ``` The SDK stops the remaining `After` callbacks on rejection. The wrapper records the error @@ -132,10 +138,17 @@ extraction; the default extracts W3C trace context only, not baggage. Session/us are process-local: the A2A context ID becomes the server's `session.id` unless set on its context. Headers, authenticated user details and request metadata are never copied to spans. -The server uses the same privacy settings and stream/error handling as the client. HTTP and gRPC -cancel the request context on disconnect. When calling a handler directly, pass a cancellable -context and cancel it if you stop consuming a stream early: the SDK does not call `After` on -iterator close. +Use `WrapClient` and `WrapHandler` for streaming calls. They observe normal iterator exhaustion +and early abandonment because the SDK does not call `After` on iterator close. A persisted +nonterminal task snapshot can reach clean EOF without a final event; its span completes as OK, +not canceled. Breaking out early completes an unfinished span as ERROR/canceled and cancels only +the wrapper's child context. The caller's context is not canceled. Continue passing a deadline +or cancellable context for blocked reads and remote disconnects. + +The wrappers complete the interceptor's existing span, without adding spans or changing events. +Wrap the outermost server handler (after any `a2asrv.InterceptedHandler` layers). The bare +interceptor APIs remain available, but without the completion wrappers they can only finish +streams on a terminal event or context cancellation, not on clean EOF. ## What is traced @@ -154,8 +167,10 @@ A stream is one span. Every streamed event is a span event named `a2a.event` (wi `gen_ai.a2a.event.type`, and the task state or artifact type when the event has one), and the span ends at the agent's reply message or when the task reaches a terminal state (`completed`, `failed`, `canceled`, `rejected`) or `input-required`, which is when a2a-go closes the stream. `auth-required` does not close -an active stream: later events and its final result remain on the same span. Unary task snapshots -include their current status message in `output.value`, including an authorization prompt. +an active stream: later events and its final result remain on the same span. With the completion +wrappers, clean EOF also ends a stream, including a persisted `auth-required` snapshot with no +active executor. The latest snapshot/status message is included in `output.value`, including an +authorization prompt at EOF; a later status replaces it. ## Attributes @@ -186,13 +201,15 @@ here, and no token or cost attribute is invented. - An error a2a-go returns ends the span with status ERROR and `error.type`: the A2A reason for protocol errors (`TASK_NOT_FOUND`, `UNSUPPORTED_OPERATION`, ...), `canceled` or `timeout` for the context, and `transport` for everything else (connection, HTTP status, decoding). The remote - agent's own error message is added as `error.message` unless inputs or outputs are hidden. - Transport error text (URLs, host names) is never recorded. + agent's own error message is added as `error.message` unless the explicit or derived + error-details policy hides it. Hidden errors retain ERROR status and `error.type`, with an + empty status description. Transport error text (URLs, host names) is never recorded. - A task the agent reports as `failed` or `rejected` is ERROR with `error.type` set to the state. `canceled`, `input-required` and `auth-required` are recorded in `gen_ai.a2a.task.state` and are not errors. -- A stream ends without a callback when you stop reading early. The span then ends, as ERROR - `canceled`, when the context you passed is cancelled, so pass one you cancel (`defer cancel()`). +- With `WrapClient` / `WrapHandler`, normal EOF completes an unfinished span successfully; + stopping iteration early ends it as ERROR `canceled`. Without them, early abandonment requires + caller context cancellation to finish the span, and clean EOF cannot be distinguished. - The interceptor never changes a request, a response or an error, and never fails a call. ## Privacy @@ -202,13 +219,18 @@ switches, with the same names and defaults as Python, turn it off: | Variable | Effect | |---|---| -| `FI_HIDE_INPUTS` | `input.value` becomes `__REDACTED__`; `error.message` is dropped | -| `FI_HIDE_OUTPUTS` | `output.value` becomes `__REDACTED__`; `error.message` is dropped | -| `FI_HIDE_INPUT_MESSAGES`, `FI_HIDE_OUTPUT_MESSAGES`, `FI_HIDE_INPUT_TEXT`, `FI_HIDE_OUTPUT_TEXT`, `FI_HIDE_INPUT_IMAGES`, `FI_HIDE_LLM_INVOCATION_PARAMETERS`, `FI_BASE64_IMAGE_MAX_LENGTH` | read for parity; they act on attributes this module does not write (and on spans you start yourself through `TraceConfig.SetAttributes`) | +| `FI_HIDE_INPUTS` | `input.value` becomes `__REDACTED__`; derives error-detail hiding | +| `FI_HIDE_OUTPUTS` | `output.value` becomes `__REDACTED__`; derives error-detail hiding | +| `FI_HIDE_INPUT_MESSAGES`, `FI_HIDE_OUTPUT_MESSAGES`, `FI_HIDE_INPUT_TEXT`, `FI_HIDE_OUTPUT_TEXT` | hide matching message attributes on your own spans through `TraceConfig.SetAttributes`; also derive error-detail hiding | +| `FI_HIDE_ERROR_DETAILS` | `true` drops `error.message`, `exception.message`, `exception.stacktrace` and the ERROR status description; `false` explicitly allows error details, even when content is hidden | +| `FI_HIDE_INPUT_IMAGES`, `FI_HIDE_LLM_INVOCATION_PARAMETERS`, `FI_BASE64_IMAGE_MAX_LENGTH` | act on matching attributes you set through `TraceConfig.SetAttributes`; do not derive error-detail hiding | -A value counts only when it is `true` (any case). `WithTraceConfig(traceaia2a.TraceConfig{...})` +Content-hiding switches enable only for `true` (any case). `WithTraceConfig(traceaia2a.TraceConfig{...})` sets them in code instead. Task ids, states, the agent URL, the session id and the error class are -kept in every mode. +kept in every mode. When `FI_HIDE_ERROR_DETAILS` is unset, empty or invalid, error details +are hidden if any of the six input/output, message or text switches is true. In code, +`TraceConfig.HideErrorDetails` is a `*bool`: nil derives that policy; a pointer to true/false +overrides it. `WithTraceConfig` replaces environment settings for that interceptor. In every mode: file parts (raw bytes or file URLs) are never copied, credentials and query strings in the agent URL are stripped, request headers are never copied, transport error text is never diff --git a/go/a2a/client_snapshot_test.go b/go/a2a/client_snapshot_test.go new file mode 100644 index 000000000..a6b88793a --- /dev/null +++ b/go/a2a/client_snapshot_test.go @@ -0,0 +1,46 @@ +package traceaia2a_test + +import ( + "context" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/codes" + "net/http/httptest" + "testing" + "time" +) + +// Client-only coverage against an uninstrumented official server. +func TestClientAuthSnapshotEOF(t *testing.T) { + store := taskstore.NewInMemory(nil) + task := &a2a.Task{ID: "persisted-auth-task", ContextID: "persisted-context", Status: a2a.TaskStatus{State: a2a.TaskStateAuthRequired}} + if _, err := store.Create(context.Background(), task); err != nil { + t.Fatal(err) + } + handler := a2asrv.NewHandler(repliesWithMessage(), a2asrv.WithTaskStore(store)) + server := httptest.NewServer(a2asrv.NewJSONRPCHandler(handler)) + defer server.Close() + client, rec, _ := newClient(t, card(server.URL, false, true)) + defer client.Destroy() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + count := 0 + for _, err := range traceaia2a.WrapClient(client).SubscribeToTask(ctx, &a2a.SubscribeToTaskRequest{ID: task.ID}) { + if err != nil { + t.Fatal(err) + } + count++ + } + if count != 1 || ctx.Err() != nil { + t.Fatalf("not a clean EOF: events=%d ctx=%v", count, ctx.Err()) + } + t.Logf("clean subscription EOF: started=%d ended=%d", len(rec.Started()), len(rec.Ended())) + if len(rec.Ended()) != 1 { + t.Fatalf("successful subscription left client span unfinished") + } + if rec.Ended()[0].Status().Code != codes.Ok { + t.Fatalf("incorrect completion: %v", rec.Ended()[0].Status()) + } +} diff --git a/go/a2a/completion_regression_test.go b/go/a2a/completion_regression_test.go new file mode 100644 index 000000000..5cc948659 --- /dev/null +++ b/go/a2a/completion_regression_test.go @@ -0,0 +1,95 @@ +package traceaia2a_test + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/codes" +) + +// Exercise the newly introduced server spans with the shared FI error-privacy contract. +func TestServerErrorPrivacyFromEnv(t *testing.T) { + for _, flag := range []string{"FI_HIDE_ERROR_DETAILS", "FI_HIDE_INPUT_TEXT", "FI_HIDE_OUTPUT_TEXT", "FI_HIDE_INPUT_MESSAGES", "FI_HIDE_OUTPUT_MESSAGES"} { + t.Run(flag, func(t *testing.T) { + for _, name := range []string{"FI_HIDE_INPUTS", "FI_HIDE_OUTPUTS", "FI_HIDE_ERROR_DETAILS", "FI_HIDE_INPUT_TEXT", "FI_HIDE_OUTPUT_TEXT", "FI_HIDE_INPUT_MESSAGES", "FI_HIDE_OUTPUT_MESSAGES"} { + t.Setenv(name, "false") + } + if flag != "FI_HIDE_ERROR_DETAILS" { + t.Setenv("FI_HIDE_ERROR_DETAILS", "") + } + t.Setenv(flag, "true") + tp, rec := serverRecorder(t) + tracing := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)) + rejection := a2a.NewError(a2a.ErrUnauthorized, "synthetic-private-error-92831") + h := a2asrv.NewHandler(repliesWithMessage(), a2asrv.WithCallInterceptors(tracing, tracing.Wrap(responsePolicy{after: func(context.Context, *a2asrv.CallContext, *a2asrv.Response) error { return rejection }}))) + _, err := h.SendMessage(context.Background(), userMessage("safe input")) + if err != rejection { + t.Fatalf("error identity changed: %v", err) + } + spans := serverSpans(rec) + if len(spans) != 1 || spans[0].Status().Code != codes.Error { + t.Fatal("expected one ERROR server span") + } + if strings.Contains(everything(spans[0]), "synthetic-private-error-92831") { + t.Errorf("%s=true leaked error.message=%q", flag, str(attrs(spans[0]), traceaia2a.ErrorMessage)) + } + }) + } +} + +// A persisted auth-required task has no active execution after a restart. The +// official handler returns its snapshot and closes the subscription successfully. +func TestAuthSnapshotSubscriptionEOF(t *testing.T) { + for _, transport := range []string{"direct", "jsonrpc", "rest", "grpc"} { + t.Run(transport, func(t *testing.T) { + tp, rec := serverRecorder(t) + store := taskstore.NewInMemory(nil) + task := &a2a.Task{ID: "persisted-auth-task", ContextID: "persisted-context", Status: a2a.TaskStatus{State: a2a.TaskStateAuthRequired}} + if _, err := store.Create(context.Background(), task); err != nil { + t.Fatal(err) + } + tracing := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)) + h := a2asrv.NewHandler(repliesWithMessage(), a2asrv.WithTaskStore(store), a2asrv.WithCallInterceptors(tracing)) + subscribe := traceaia2a.WrapHandler(h).SubscribeToTask + want := 1 + if transport != "direct" { + client := serverClient(t, transport, repliesWithMessage(), tp, traceaia2a.TraceConfig{}, a2asrv.WithTaskStore(store)) + subscribe, want = client.SubscribeToTask, 2 + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + n := 0 + for ev, err := range subscribe(ctx, &a2a.SubscribeToTaskRequest{ID: task.ID}) { + if err != nil { + t.Fatalf("unexpected protocol failure: %v", err) + } + got, ok := ev.(*a2a.Task) + if !ok || got.Status.State != a2a.TaskStateAuthRequired { + t.Fatalf("unexpected snapshot: %v", ev) + } + n++ + } + if n != 1 || ctx.Err() != nil { + t.Fatalf("not a clean EOF: n=%d ctx=%v", n, ctx.Err()) + } + t.Logf("clean EOF before caller cancellation: started=%d ended=%d", len(rec.Started()), len(rec.Ended())) + if len(rec.Ended()) != want { + t.Errorf("completed subscription left spans open: started=%d ended=%d want=%d", len(rec.Started()), len(rec.Ended()), want) + } + cancel() + waitEnded(t, rec, want) + for _, s := range rec.Ended() { + t.Logf("after caller cancellation: %s status=%v error.type=%s", s.SpanKind(), s.Status(), str(attrs(s), traceaia2a.ErrorType)) + if s.Status().Code != codes.Ok { + t.Errorf("successful subscription was falsely reported as an error: %v", s.Status()) + } + } + }) + } +} diff --git a/go/a2a/doc.go b/go/a2a/doc.go index 672016dba..5f244d007 100644 --- a/go/a2a/doc.go +++ b/go/a2a/doc.go @@ -9,7 +9,10 @@ // gen_ai.a2a.task.state, gen_ai.a2a.agent.url, ...), each written only when // a2a-go reports that fact. A stream is one span: every event is a span event // on it, and the span ends at the reply message or at a terminal or -// interrupted task state. The W3C traceparent is written into the request's +// interrupted task state. Add [WrapClient] and [WrapHandler] to also observe +// clean EOF and early iterator abandonment, which the SDK's hooks do not report. +// These wrappers complete the existing span and never create a duplicate. +// The W3C traceparent is written into the request's // service parameters (HTTP headers on the JSON-RPC and REST transports, call // metadata on gRPC), so a remote agent that reads it continues the same trace. // diff --git a/go/a2a/error_privacy_test.go b/go/a2a/error_privacy_test.go new file mode 100644 index 000000000..2dd816815 --- /dev/null +++ b/go/a2a/error_privacy_test.go @@ -0,0 +1,90 @@ +package traceaia2a_test + +import ( + "context" + "strings" + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" +) + +func TestErrorDetailsPolicy(t *testing.T) { + contentFlags := []string{"FI_HIDE_INPUTS", "FI_HIDE_OUTPUTS", "FI_HIDE_INPUT_MESSAGES", "FI_HIDE_OUTPUT_MESSAGES", "FI_HIDE_INPUT_TEXT", "FI_HIDE_OUTPUT_TEXT"} + for _, flag := range append([]string{"none", "FI_HIDE_INPUT_IMAGES", "FI_HIDE_LLM_INVOCATION_PARAMETERS"}, contentFlags...) { + for _, override := range []string{"", "true", " TrUe ", "false", " FaLsE ", "invalid", "1"} { + t.Run(flag+"/"+override, func(t *testing.T) { + for _, name := range append(append([]string{}, contentFlags...), "FI_HIDE_INPUT_IMAGES", "FI_HIDE_LLM_INVOCATION_PARAMETERS") { + t.Setenv(name, "false") + } + if flag != "none" { + t.Setenv(flag, "true") + } + t.Setenv("FI_HIDE_ERROR_DETAILS", override) + wantHidden := false + for _, name := range contentFlags { + wantHidden = wantHidden || flag == name + } + switch strings.ToLower(strings.TrimSpace(override)) { + case "true": + wantHidden = true + case "false": + wantHidden = false + } + config := traceaia2a.TraceConfigFromEnv() + for _, key := range []string{traceaia2a.ErrorMessage, "exception.message", "exception.stacktrace"} { + if _, keep := config.Mask(attribute.String(key, marker)); keep == wantHidden { + t.Errorf("%s keep=%v, hidden=%v", key, keep, wantHidden) + } + } + for _, key := range []string{traceaia2a.ErrorType, "exception.type"} { + if _, keep := config.Mask(attribute.String(key, "UNAUTHORIZED")); !keep { + t.Errorf("error class %s was hidden", key) + } + } + // Exercise the same derived/default policy on actual client errors. + tp, rec := serverRecorder(t) + i := traceaia2a.NewInterceptor(traceaia2a.WithTracerProvider(tp)) + ctx, _, err := i.Before(context.Background(), &a2aclient.Request{Method: "GetTask"}) + if err != nil { + t.Fatal(err) + } + rejection := a2a.NewError(a2a.ErrUnauthorized, marker) + if err := i.After(ctx, &a2aclient.Response{Method: "GetTask", Err: rejection}); err != nil { + t.Fatal(err) + } + span := only(t, rec) + if span.Status().Code != codes.Error || str(attrs(span), traceaia2a.ErrorType) != "UNAUTHORIZED" { + t.Fatalf("error classification lost: %v", span.Status()) + } + if strings.Contains(everything(span), marker) == wantHidden { + t.Errorf("error text policy mismatch: hidden=%v", wantHidden) + } + if wantHidden && span.Status().Description != "" { + t.Error("hidden error retained a status description") + } + }) + } + } +} + +func TestErrorDetailsProgrammaticOverride(t *testing.T) { + for _, hide := range []bool{true, false} { + t.Run(map[bool]string{true: "hide", false: "allow"}[hide], func(t *testing.T) { + // Code wins over environment, and the constructor snapshots the pointer. + t.Setenv("FI_HIDE_ERROR_DETAILS", map[bool]string{true: "false", false: "true"}[hide]) + tp, rec := serverRecorder(t) + override := hide + i := traceaia2a.NewInterceptor(traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(traceaia2a.TraceConfig{HideInputs: true, HideErrorDetails: &override})) + override = !hide + ctx, _, _ := i.Before(context.Background(), &a2aclient.Request{Method: "GetTask"}) + _ = i.After(ctx, &a2aclient.Response{Method: "GetTask", Err: a2a.NewError(a2a.ErrUnauthorized, marker)}) + if strings.Contains(everything(only(t, rec)), marker) == hide { + t.Fatal("programmatic override did not win or changed after construction") + } + }) + } +} diff --git a/go/a2a/error_privacy_wire_test.go b/go/a2a/error_privacy_wire_test.go new file mode 100644 index 000000000..32a038305 --- /dev/null +++ b/go/a2a/error_privacy_wire_test.go @@ -0,0 +1,80 @@ +package traceaia2a_test + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + collectortrace "go.opentelemetry.io/proto/otlp/collector/trace/v1" + tracepb "go.opentelemetry.io/proto/otlp/trace/v1" + "google.golang.org/protobuf/proto" +) + +func TestErrorDetailsOTLPWire(t *testing.T) { + for _, name := range []string{"FI_HIDE_INPUTS", "FI_HIDE_OUTPUTS", "FI_HIDE_INPUT_TEXT", "FI_HIDE_OUTPUT_TEXT", "FI_HIDE_INPUT_MESSAGES", "FI_HIDE_OUTPUT_MESSAGES"} { + t.Setenv(name, "false") + } + t.Setenv("FI_HIDE_ERROR_DETAILS", "true") + bodies := make(chan []byte, 1) + sink := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + if err != nil { + t.Error(err) + } + bodies <- body + w.Header().Set("Content-Type", "application/x-protobuf") + })) + defer sink.Close() + exporter, err := otlptracehttp.New(context.Background(), otlptracehttp.WithEndpointURL(sink.URL+"/tracer/v1/traces")) + if err != nil { + t.Fatal(err) + } + tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) + defer tp.Shutdown(context.Background()) + tracing := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)) + rejection := a2a.NewError(a2a.ErrUnauthorized, "synthetic-private-error-92831") + h := a2asrv.NewHandler(repliesWithMessage(), a2asrv.WithCallInterceptors(tracing, tracing.Wrap(responsePolicy{after: func(context.Context, *a2asrv.CallContext, *a2asrv.Response) error { return rejection }}))) + if _, got := h.SendMessage(context.Background(), userMessage("safe input")); got != rejection { + t.Fatalf("error identity changed: %v", got) + } + if len(bodies) != 1 { + t.Fatalf("want exactly one server-span export, got %d", len(bodies)) + } + body := <-bodies + if bytes.Contains(body, []byte("synthetic-private-error-92831")) { + t.Fatal("FI_HIDE_ERROR_DETAILS=true leaked protocol error text in actual OTLP protobuf") + } + request := &collectortrace.ExportTraceServiceRequest{} + if err := proto.Unmarshal(body, request); err != nil { + t.Fatal(err) + } + count := 0 + for _, resource := range request.ResourceSpans { + for _, scope := range resource.ScopeSpans { + for _, span := range scope.Spans { + count++ + if span.Kind != tracepb.Span_SPAN_KIND_SERVER || span.Status.Code != tracepb.Status_STATUS_CODE_ERROR || span.Status.Message != "" { + t.Fatalf("wrong exported error status: %v", span.Status) + } + values := map[string]string{} + for _, attr := range span.Attributes { + values[attr.Key] = attr.Value.GetStringValue() + } + if values[traceaia2a.ErrorType] != "UNAUTHORIZED" || values[traceaia2a.InputValue] != "safe input" { + t.Fatal("error privacy dropped the class or unrelated input capture") + } + } + } + } + if count != 1 { + t.Fatalf("want one decoded server span, got %d", count) + } +} diff --git a/go/a2a/example/main.go b/go/a2a/example/main.go index a0f612d08..945e10e1e 100644 --- a/go/a2a/example/main.go +++ b/go/a2a/example/main.go @@ -135,11 +135,12 @@ func run() error { if err != nil { return fmt.Errorf("resolve agent card: %w", err) } - client, err := a2aclient.NewFromCard(ctx, card, + sdkClient, err := a2aclient.NewFromCard(ctx, card, a2aclient.WithCallInterceptors(traceaia2a.NewInterceptor(traceaia2a.WithTracerProvider(tp)))) if err != nil { return err } + client := traceaia2a.WrapClient(sdkClient) defer func() { _ = client.Destroy() }() sessionCtx := traceaia2a.ContextWithSessionID(ctx, *session) @@ -173,7 +174,7 @@ func run() error { } // send makes one A2A call, streamed or not, and prints what came back. -func send(ctx context.Context, client *a2aclient.Client, req *a2a.SendMessageRequest) error { +func send(ctx context.Context, client *traceaia2a.Client, req *a2a.SendMessageRequest) error { if *stream { for ev, err := range client.SendStreamingMessage(ctx, req) { if err != nil { @@ -228,8 +229,8 @@ func startDemoAgent(tp *sdktrace.TracerProvider) (string, func(), error) { } agent := &demoAgent{tracer: tp.Tracer("demo-remote-agent"), tc: traceaia2a.TraceConfigFromEnv()} mux := http.NewServeMux() - handler := a2asrv.NewHandler(agent, a2asrv.WithCallInterceptors( - traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)))) + handler := traceaia2a.WrapHandler(a2asrv.NewHandler(agent, a2asrv.WithCallInterceptors( + traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp))))) mux.Handle("/a2a", a2asrv.NewJSONRPCHandler(handler)) mux.Handle(a2asrv.WellKnownAgentCardPath, a2asrv.NewStaticAgentCardHandler(card)) srv := &http.Server{Handler: mux, ReadHeaderTimeout: 10 * time.Second} diff --git a/go/a2a/interceptor.go b/go/a2a/interceptor.go index 8ceecadd9..47a3f9bc3 100644 --- a/go/a2a/interceptor.go +++ b/go/a2a/interceptor.go @@ -67,6 +67,8 @@ func WithPropagator(p propagation.TextMapPropagator) Option { // never changes a request payload, a response or an error. It adds only the // traceparent (and tracestate) service parameters, which the JSON-RPC and REST // transports send as HTTP headers and the gRPC transport as call metadata. +// Use [WrapClient] around the created client for clean-EOF and early-abandonment +// completion; the SDK's interceptor hooks alone do not observe iterator close. type Interceptor struct { tracer trace.Tracer tc TraceConfig @@ -92,6 +94,11 @@ func NewInterceptor(opts ...Option) *Interceptor { if s.tc != nil { tc = *s.tc } + // Freeze the optional override just like the other value-type switches. + if tc.HideErrorDetails != nil { + hide := *tc.HideErrorDetails + tc.HideErrorDetails = &hide + } prop := s.prop if prop == nil && !s.noProp { prop = propagation.TraceContext{} @@ -186,6 +193,7 @@ func (i *Interceptor) Before(ctx context.Context, req *a2aclient.Request) (conte st.stopWatch = context.AfterFunc(ctx, func() { st.endWithContextError(spanCtx) }) st.mu.Unlock() } + attachCompletion(ctx, clientCompletionKey, st) return context.WithValue(ctx, stateKey{}, st), nil, nil } @@ -344,9 +352,10 @@ func (s *callState) applyTask(t *a2a.Task) { for _, a := range t.Artifacts { s.addArtifact(a) } - // A unary snapshot returns its current status message, including an auth - // prompt. Only streaming calls wait for a stream-ending state. - if t.Status.Message != nil && (!s.streaming || endsStream(t.Status.State)) { + // Keep the latest snapshot's status text, including an auth prompt at EOF. + // A later status replaces it so an active stream never retains a stale prompt. + s.final = "" + if t.Status.Message != nil { s.final = partsText(t.Status.Message.Parts) } } @@ -402,12 +411,11 @@ func (s *callState) applyEvent(payload any) bool { if v.Status.State != a2a.TaskStateUnspecified { ev = append(ev, attribute.String(A2ATaskState, string(v.Status.State))) } - if endsStream(v.Status.State) { - done = true - if v.Status.Message != nil { - s.final = partsText(v.Status.Message.Parts) - } + s.final = "" + if v.Status.Message != nil { + s.final = partsText(v.Status.Message.Parts) } + done = endsStream(v.Status.State) } case *a2a.TaskArtifactUpdateEvent: ev = append(ev, attribute.String(A2AEventType, "artifact-update")) @@ -461,7 +469,7 @@ func (s *callState) finish() { s.output() if s.state == a2a.TaskStateFailed || s.state == a2a.TaskStateRejected { s.span.SetAttributes(attribute.String(ErrorType, string(s.state))) - s.span.SetStatus(codes.Error, string(s.state)) + s.setErrorStatus(string(s.state)) } else { s.span.SetStatus(codes.Ok, "") } @@ -482,7 +490,7 @@ func (s *callState) fail(err error) { if agentMsg != "" && s.tc.errorMessageAllowed() { s.tc.SetAttributes(s.span, attribute.String(ErrorMessage, truncate(agentMsg))) } - s.span.SetStatus(codes.Error, errType) + s.setErrorStatus(errType) s.span.End() } @@ -500,3 +508,11 @@ func (s *callState) endWithContextError(ctx context.Context) { } s.fail(err) } + +// Error class attributes remain useful even when descriptions are hidden. +func (s *callState) setErrorStatus(description string) { + if !s.tc.errorMessageAllowed() { + description = "" + } + s.span.SetStatus(codes.Error, description) +} diff --git a/go/a2a/interceptor_test.go b/go/a2a/interceptor_test.go index 7365a7594..c314c0ca9 100644 --- a/go/a2a/interceptor_test.go +++ b/go/a2a/interceptor_test.go @@ -680,7 +680,11 @@ func TestProtocolErrorKeepsTheAgentMessageOnlyWhenContentIsOn(t *testing.T) { } s := only(t, rec) a := attrs(s) - if s.Status().Code != codes.Error || str(a, traceaia2a.ErrorType) != "TASK_NOT_FOUND" || s.Status().Description != "TASK_NOT_FOUND" { + wantDescription := "TASK_NOT_FOUND" + if hide { + wantDescription = "" + } + if s.Status().Code != codes.Error || str(a, traceaia2a.ErrorType) != "TASK_NOT_FOUND" || s.Status().Description != wantDescription { t.Errorf("status %+v error.type %q", s.Status(), str(a, traceaia2a.ErrorType)) } if has(a, traceaia2a.ErrorMessage) == hide { diff --git a/go/a2a/internal_test.go b/go/a2a/internal_test.go index 7bff838d0..f64644496 100644 --- a/go/a2a/internal_test.go +++ b/go/a2a/internal_test.go @@ -98,7 +98,7 @@ func TestClassifyError(t *testing.T) { func clearHideEnv(t *testing.T) { t.Helper() for _, k := range []string{EnvHideLLMInvocationParameters, EnvHideInputs, EnvHideOutputs, EnvHideInputMessages, - EnvHideOutputMessages, EnvHideInputImages, EnvHideInputText, EnvHideOutputText, EnvBase64ImageMaxLength} { + EnvHideOutputMessages, EnvHideInputImages, EnvHideInputText, EnvHideOutputText, EnvBase64ImageMaxLength, EnvHideErrorDetails} { t.Setenv(k, "") } } diff --git a/go/a2a/server.go b/go/a2a/server.go index 100d13057..4495e4951 100644 --- a/go/a2a/server.go +++ b/go/a2a/server.go @@ -26,16 +26,16 @@ import ( // Headers, credentials, baggage and authenticated user details are not recorded. // Use [WithTraceConfig] or FI_HIDE_INPUTS / FI_HIDE_OUTPUTS to hide message content. // -// a2asrv calls After for each streamed event, not on iterator close. Streams -// therefore end at a terminal/interrupted event or when the request context is -// canceled. Direct in-process callers must cancel their context if they stop -// reading early; HTTP and gRPC transports cancel it when the request ends. +// a2asrv calls After for each streamed event, not on iterator close. Use +// [WrapHandler] around the instrumented handler to complete spans at clean EOF +// (including persisted nonterminal task snapshots) and on early abandonment. +// Without it, streams end only at a terminal/interrupted event or cancellation. // Detached background task execution can outlive the request span. // This interceptor does not wrap transport handlers or create executor spans. // -// handler := a2asrv.NewHandler(executor, +// handler := traceaia2a.WrapHandler(a2asrv.NewHandler(executor, // a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor( -// traceaia2a.WithTracerProvider(tp)))) +// traceaia2a.WithTracerProvider(tp))))) // // Use that handler with a2asrv.NewJSONRPCHandler, a2asrv.NewRESTHandler or // a2agrpc/v1.NewHandler. @@ -115,6 +115,7 @@ func (i *ServerInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallCont st.mu.Lock() st.stopWatch = context.AfterFunc(ctx, func() { st.endWithContextError(spanCtx) }) st.mu.Unlock() + attachCompletion(ctx, serverCompletionKey, st) return context.WithValue(ctx, serverStateKey{}, st), nil, nil } diff --git a/go/a2a/server_test.go b/go/a2a/server_test.go index 77b7f3678..04c163143 100644 --- a/go/a2a/server_test.go +++ b/go/a2a/server_test.go @@ -20,11 +20,11 @@ import ( "google.golang.org/grpc/credentials/insecure" ) -func serverClient(t *testing.T, transport string, executor *agent, tp *sdktrace.TracerProvider, tc traceaia2a.TraceConfig, serverOpts ...a2asrv.RequestHandlerOption) *a2aclient.Client { +func serverClient(t *testing.T, transport string, executor *agent, tp *sdktrace.TracerProvider, tc traceaia2a.TraceConfig, serverOpts ...a2asrv.RequestHandlerOption) *traceaia2a.Client { t.Helper() serverOpts = append([]a2asrv.RequestHandlerOption{a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor( traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(tc)))}, serverOpts...) - handler := a2asrv.NewHandler(executor, serverOpts...) + handler := traceaia2a.WrapHandler(a2asrv.NewHandler(executor, serverOpts...)) var c *a2a.AgentCard opts := []a2aclient.FactoryOption{a2aclient.WithCallInterceptors(traceaia2a.NewInterceptor( traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(tc)))} @@ -53,7 +53,7 @@ func serverClient(t *testing.T, transport string, executor *agent, tp *sdktrace. t.Fatal(err) } t.Cleanup(func() { _ = client.Destroy() }) - return client + return traceaia2a.WrapClient(client) } func TestServerTraceContinuesClient(t *testing.T) { diff --git a/go/a2a/stream.go b/go/a2a/stream.go new file mode 100644 index 000000000..335b6307e --- /dev/null +++ b/go/a2a/stream.go @@ -0,0 +1,142 @@ +package traceaia2a + +import ( + "context" + "iter" + "sync" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + "github.com/a2aproject/a2a-go/v2/a2asrv" +) + +// Client preserves the SDK client's API and observes stream completion. Create +// the SDK client with [NewInterceptor], then call [WrapClient]. The wrapper does +// not create spans: it completes the interceptor's existing span at clean EOF, +// including snapshots with no terminal event, or on early iterator abandonment. +// All non-streaming methods are delegated unchanged to the embedded client. +type Client struct{ *a2aclient.Client } + +// WrapClient adds iterator-completion handling to an instrumented SDK client. +// Use the returned client for streaming calls; calling its embedded Client +// directly bypasses completion handling. Install [NewInterceptor] only once. +func WrapClient(client *a2aclient.Client) *Client { return &Client{Client: client} } + +// SendStreamingMessage observes completion without changing events or errors. +func (c *Client) SendStreamingMessage(ctx context.Context, req *a2a.SendMessageRequest) iter.Seq2[a2a.Event, error] { + return observeStream(ctx, clientCompletionKey, "SendStreamingMessage", func(ctx context.Context) iter.Seq2[a2a.Event, error] { + return c.Client.SendStreamingMessage(ctx, req) + }) +} + +// SubscribeToTask observes completion, including a persisted nonterminal snapshot. +func (c *Client) SubscribeToTask(ctx context.Context, req *a2a.SubscribeToTaskRequest) iter.Seq2[a2a.Event, error] { + return observeStream(ctx, clientCompletionKey, "SubscribeToTask", func(ctx context.Context) iter.Seq2[a2a.Event, error] { + return c.Client.SubscribeToTask(ctx, req) + }) +} + +// WrapHandler adds iterator-completion handling to an SDK handler configured +// with [NewServerInterceptor]. Wrap the outermost RequestHandler, after any +// InterceptedHandler layers, and pass the result to the HTTP/gRPC transport or +// direct caller. It completes existing spans, never creates an additional span, +// and delegates non-streaming methods unchanged. Without this wrapper the SDK +// provides no After callback on clean EOF or early iterator abandonment. +func WrapHandler(handler a2asrv.RequestHandler) a2asrv.RequestHandler { + return &completionHandler{RequestHandler: handler} +} + +type completionHandler struct{ a2asrv.RequestHandler } + +func (h *completionHandler) SendStreamingMessage(ctx context.Context, req *a2a.SendMessageRequest) iter.Seq2[a2a.Event, error] { + return observeStream(ctx, serverCompletionKey, "SendStreamingMessage", func(ctx context.Context) iter.Seq2[a2a.Event, error] { + return h.RequestHandler.SendStreamingMessage(ctx, req) + }) +} + +func (h *completionHandler) SubscribeToTask(ctx context.Context, req *a2a.SubscribeToTaskRequest) iter.Seq2[a2a.Event, error] { + return observeStream(ctx, serverCompletionKey, "SubscribeToTask", func(ctx context.Context) iter.Seq2[a2a.Event, error] { + return h.RequestHandler.SubscribeToTask(ctx, req) + }) +} + +type completionKey int + +const ( + clientCompletionKey completionKey = iota + serverCompletionKey +) + +// Each iteration owns a slot. Only its first matching Before can claim it, so +// nested calls inheriting this context cannot overwrite the outer call's state. +// Separate keys isolate inbound and outbound calls on the same context. +type streamCompletion struct { + mu sync.Mutex + method string + state *callState +} + +func attachCompletion(ctx context.Context, key completionKey, st *callState) { + if completion, _ := ctx.Value(key).(*streamCompletion); completion != nil { + completion.mu.Lock() + defer completion.mu.Unlock() + if completion.state == nil && completion.method == st.method { + completion.state = st + } + } +} + +func (c *streamCompletion) end(err error) { + c.mu.Lock() + st := c.state + c.mu.Unlock() + if st == nil { + return // No tracing interceptor was installed, or Before was skipped. + } + st.mu.Lock() + defer st.mu.Unlock() + if st.ended { + return + } + if err != nil { + st.fail(err) + } else { + st.finish() + } +} + +func observeStream(ctx context.Context, key completionKey, method string, stream func(context.Context) iter.Seq2[a2a.Event, error]) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + parent := ctx + if parent == nil { + parent = context.Background() + } + // Cancel only our child context on early abandonment, never the caller's. + // Complete the span before canceling on clean EOF, otherwise its context + // watcher would incorrectly turn a successful snapshot into an error. + callCtx, cancel := context.WithCancel(parent) + defer cancel() + completion := &streamCompletion{method: method} + callCtx = context.WithValue(callCtx, key, completion) + var streamErr error + completed := false + defer func() { + if streamErr == nil { + streamErr = callCtx.Err() + } + if streamErr == nil && !completed { + streamErr = context.Canceled + } + completion.end(streamErr) + }() + for event, err := range stream(callCtx) { + if err != nil { + streamErr = err + } + if !yield(event, err) { + return + } + } + completed = true + } +} diff --git a/go/a2a/stream_completion_test.go b/go/a2a/stream_completion_test.go new file mode 100644 index 000000000..06ee1d432 --- /dev/null +++ b/go/a2a/stream_completion_test.go @@ -0,0 +1,179 @@ +package traceaia2a_test + +import ( + "context" + "iter" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + traceaia2a "github.com/future-agi/traceAI/go/a2a" + "go.opentelemetry.io/otel/codes" +) + +type streamHandler struct { + a2asrv.RequestHandler + stream func(context.Context) iter.Seq2[a2a.Event, error] +} + +func (h streamHandler) SubscribeToTask(ctx context.Context, _ *a2a.SubscribeToTaskRequest) iter.Seq2[a2a.Event, error] { + return h.stream(ctx) +} + +func TestStreamCompletionLifecycle(t *testing.T) { + for _, outcome := range []string{"empty", "snapshot", "break", "error", "panic", "cancel", "timeout"} { + t.Run(outcome, func(t *testing.T) { + tp, rec := serverRecorder(t) + tracing := traceaia2a.NewServerInterceptor(traceaia2a.WithTracerProvider(tp)) + var innerCtx context.Context + rejection := a2a.NewError(a2a.ErrUnauthorized, marker) + calls := 0 + base := streamHandler{stream: func(ctx context.Context) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + calls++ + innerCtx = ctx + if outcome == "empty" { + return + } + if outcome == "panic" { + panic("synthetic panic") + } + if outcome == "error" { + yield(nil, rejection) + return + } + if outcome == "cancel" || outcome == "timeout" { + <-ctx.Done() + return // Completion must check context even without an error event. + } + yield(&a2a.Task{ID: "snapshot", Status: a2a.TaskStatus{State: a2a.TaskStateAuthRequired, + Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("authorize"))}}, nil) + } + }} + h := traceaia2a.WrapHandler(&a2asrv.InterceptedHandler{Handler: base, Interceptors: []a2asrv.CallInterceptor{tracing}}) + // Double-wrapping must not create or end duplicate spans. + h = traceaia2a.WrapHandler(h) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if outcome == "cancel" { + cancel() + } + if outcome == "timeout" { + var stop context.CancelFunc + ctx, stop = context.WithDeadline(ctx, time.Now().Add(-time.Second)) + defer stop() + } + seq := h.SubscribeToTask(ctx, &a2a.SubscribeToTaskRequest{ID: "snapshot"}) + if calls != 0 || len(rec.Started()) != 0 { + t.Fatal("constructing an iterator executed the call") + } + for attempt := 0; attempt < 2; attempt++ { + panicked := false + func() { + defer func() { + if p := recover(); p != nil { + if p != "synthetic panic" { + t.Fatalf("panic changed: %v", p) + } + panicked = true + } + }() + for event, err := range seq { + if outcome == "error" { + if err != rejection { + t.Fatal("error identity changed") + } + } else if err != nil || event == nil { + t.Fatalf("event or error changed: %v %v", event, err) + } + if outcome == "break" { + break + } + } + }() + if panicked != (outcome == "panic") { + t.Fatal("panic was swallowed or invented") + } + if calls != attempt+1 || len(rec.Ended()) != attempt+1 || len(rec.Started()) != attempt+1 { + t.Fatalf("iterator state not per-call: calls=%d started=%d ended=%d", calls, len(rec.Started()), len(rec.Ended())) + } + span := rec.Ended()[attempt] + wantError := outcome != "empty" && outcome != "snapshot" + if (span.Status().Code == codes.Error) != wantError { + t.Fatalf("incorrect completion: %v", span.Status()) + } + if outcome == "snapshot" && str(attrs(span), traceaia2a.OutputValue) != "authorize" { + t.Fatal("EOF lost the auth snapshot's output") + } + if outcome == "timeout" && str(attrs(span), traceaia2a.ErrorType) != "timeout" { + t.Fatal("deadline was misclassified") + } + if innerCtx.Err() == nil { + t.Fatal("wrapper did not release its child context") + } + if outcome != "cancel" && outcome != "timeout" && ctx.Err() != nil { + t.Fatal("wrapper canceled its caller's context") + } + } + }) + } +} + +func TestAuthStreamWaitsForCompletion(t *testing.T) { + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + for _, abandon := range []bool{false, true} { + t.Run(transport+map[bool]string{false: "/complete", true: "/abandon"}[abandon], func(t *testing.T) { + tp, rec := serverRecorder(t) + release := make(chan struct{}) + defer close(release) + exec := &agent{execute: func(ctx context.Context, ec *a2asrv.ExecutorContext, yield func(a2a.Event, error) bool) { + if !yield(a2a.NewSubmittedTask(ec, ec.Message), nil) || + !yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateAuthRequired, a2a.NewMessageForTask(a2a.MessageRoleAgent, ec, a2a.NewTextPart("stale auth prompt"))), nil) { + return + } + select { + case <-release: + case <-ctx.Done(): + return + } + yield(a2a.NewStatusUpdateEvent(ec, a2a.TaskStateCompleted, nil), nil) + }} + client := serverClient(t, transport, exec, tp, traceaia2a.TraceConfig{}) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + sawAuth := false + for event, err := range client.SendStreamingMessage(ctx, userMessage("protected task")) { + if err != nil { + t.Fatal(err) + } + if status, ok := event.(*a2a.TaskStatusUpdateEvent); ok && status.Status.State == a2a.TaskStateAuthRequired { + sawAuth = true + if len(rec.Ended()) != 0 { + t.Fatal("auth-required prematurely completed an active stream") + } + if abandon { + break + } + release <- struct{}{} + } + } + if !sawAuth || ctx.Err() != nil { + t.Fatal("missing auth event or caller context was canceled") + } + waitEnded(t, rec, 2) + if len(rec.Started()) != 2 { + t.Fatal("duplicate spans") + } + for _, span := range rec.Ended() { + if (span.Status().Code == codes.Error) != abandon { + t.Fatalf("wrong completion status: %v", span.Status()) + } + if !abandon && str(attrs(span), traceaia2a.OutputValue) != "" { + t.Fatal("completed stream retained a stale authorization prompt") + } + } + }) + } + } +} diff --git a/go/a2a/traceconfig.go b/go/a2a/traceconfig.go index 1e3603a7d..5b9efefbf 100644 --- a/go/a2a/traceconfig.go +++ b/go/a2a/traceconfig.go @@ -11,9 +11,11 @@ import ( // Environment variables read by [TraceConfigFromEnv]. They are the traceAI // TraceConfig switches, with the same names and defaults as the Python and -// TypeScript packages: everything is captured unless a switch is "true". +// TypeScript packages: content is captured unless its switch is "true". +// Error details use an explicit true/false override or derive from content flags. const ( EnvHideLLMInvocationParameters = "FI_HIDE_LLM_INVOCATION_PARAMETERS" + EnvHideErrorDetails = "FI_HIDE_ERROR_DETAILS" EnvHideInputs = "FI_HIDE_INPUTS" EnvHideOutputs = "FI_HIDE_OUTPUTS" EnvHideInputMessages = "FI_HIDE_INPUT_MESSAGES" @@ -35,11 +37,16 @@ const ( // TraceConfig decides which content attributes leave the process. The zero // value captures everything, like DEFAULT_HIDE_* = False in Python traceAI. type TraceConfig struct { + // HideErrorDetails drops error.message, exception.message and + // exception.stacktrace. Nil derives the policy from the six input/output, + // message and text switches below. An explicit false allows error details + // even when content is hidden. Error classes and ERROR status are kept. + HideErrorDetails *bool // HideLLMInvocationParameters drops gen_ai.request.parameters. HideLLMInvocationParameters bool // HideInputs redacts input.value (the message sent to the remote agent) // and drops input messages, system instructions and tool-call arguments. - // HideInputs or HideOutputs also drops the remote agent's error message. + // Content-hiding flags also hide error details unless explicitly overridden. HideInputs bool // HideOutputs redacts output.value (the agent's reply, artifacts and // final status message) and drops output messages and tool-call results. @@ -61,7 +68,8 @@ type TraceConfig struct { // TraceConfigFromEnv reads the FI_HIDE_* variables. A value is true only when // it is "true" (any case) and false otherwise, matching Python traceAI, which -// falls back to the default for anything it cannot parse. +// falls back to the default for anything it cannot parse. HideErrorDetails +// is optional: explicit true/false overrides derivation; other values derive it. func TraceConfigFromEnv() TraceConfig { c := TraceConfig{ HideLLMInvocationParameters: envBool(EnvHideLLMInvocationParameters), @@ -73,6 +81,14 @@ func TraceConfigFromEnv() TraceConfig { HideInputText: envBool(EnvHideInputText), HideOutputText: envBool(EnvHideOutputText), } + switch strings.ToLower(strings.TrimSpace(os.Getenv(EnvHideErrorDetails))) { + case "true": + v := true + c.HideErrorDetails = &v + case "false": + v := false + c.HideErrorDetails = &v + } if v, err := strconv.Atoi(strings.TrimSpace(os.Getenv(EnvBase64ImageMaxLength))); err == nil && v > 0 { c.Base64ImageMaxLength = v } @@ -85,7 +101,13 @@ func envBool(name string) bool { // errorMessageAllowed reports whether the remote agent's error message may be // written, mirroring the error.message rule in TraceConfig.Mask. -func (c TraceConfig) errorMessageAllowed() bool { return !c.HideInputs && !c.HideOutputs } +func (c TraceConfig) errorMessageAllowed() bool { + if c.HideErrorDetails != nil { + return !*c.HideErrorDetails + } + return !(c.HideInputs || c.HideOutputs || c.HideInputMessages || + c.HideOutputMessages || c.HideInputText || c.HideOutputText) +} func (c TraceConfig) imageMax() int { if c.Base64ImageMaxLength > 0 { @@ -97,11 +119,9 @@ func (c TraceConfig) imageMax() int { // Mask applies the config to one attribute. It returns the value to write and // false when the attribute must be dropped. The rules and their order follow // fi_instrumentation.TraceConfig.mask in Python traceAI. This package adds -// three rules for keys Python does not mask, the same as the other traceAI Go -// modules: system instructions and tool-call arguments go with the inputs, -// tool-call results with the outputs, and error.message is dropped when either -// inputs or outputs are hidden, because a remote agent's error can quote the -// request. +// A2A content rules keep system instructions and tool-call arguments with inputs, +// tool-call results with the outputs, and error details follow the explicit or +// derived HideErrorDetails policy because an error can quote the request. func (c TraceConfig) Mask(kv attribute.KeyValue) (attribute.KeyValue, bool) { key := string(kv.Key) in := strings.Contains(key, GenAIInputMessages) @@ -126,7 +146,7 @@ func (c TraceConfig) Mask(kv attribute.KeyValue) (attribute.KeyValue, bool) { return kv, false case c.HideOutputs && key == GenAIToolCallResult: return kv, false - case (c.HideInputs || c.HideOutputs) && key == ErrorMessage: + case !c.errorMessageAllowed() && (key == ErrorMessage || key == "exception.message" || key == "exception.stacktrace"): return kv, false case c.HideInputText && in && strings.Contains(key, MessageContent) && !strings.Contains(key, MessageContents): return redacted, true From 2b24ba9ff01095f5332d243adba8eb768a314088 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 08:10:13 -0800 Subject: [PATCH 6/6] fix(a2a): preserve interceptor-only auth snapshot completion Keep the legacy auth-required completion boundary only for client calls without their own stream-completion hook. Wrapped clients and server streams continue recording active authorization streams. Cover both client APIs across transports and nested completion-slot ownership. Co-Authored-By: Claude Code --- go/a2a/README.md | 12 +++- go/a2a/client_completion_test.go | 82 +++++++++++++++++++++ go/a2a/client_snapshot_test.go | 118 +++++++++++++++++++++++-------- go/a2a/interceptor.go | 27 ++++--- go/a2a/stream.go | 6 +- 5 files changed, 201 insertions(+), 44 deletions(-) create mode 100644 go/a2a/client_completion_test.go diff --git a/go/a2a/README.md b/go/a2a/README.md index bdec5c828..8942f1884 100644 --- a/go/a2a/README.md +++ b/go/a2a/README.md @@ -166,12 +166,18 @@ a span. A stream is one span. Every streamed event is a span event named `a2a.event` (with `gen_ai.a2a.event.type`, and the task state or artifact type when the event has one), and the span ends at the agent's reply message or when the task reaches a terminal state (`completed`, -`failed`, `canceled`, `rejected`) or `input-required`, which is when a2a-go closes the stream. `auth-required` does not close -an active stream: later events and its final result remain on the same span. With the completion -wrappers, clean EOF also ends a stream, including a persisted `auth-required` snapshot with no +`failed`, `canceled`, `rejected`) or `input-required`, which is when a2a-go closes the stream. +With `WrapClient` and `WrapHandler`, `auth-required` keeps an active stream's span open for later +events and its final result. Clean EOF also ends the span, including a persisted snapshot with no active executor. The latest snapshot/status message is included in `output.value`, including an authorization prompt at EOF; a later status replaces it. +Existing clients using only `WithCallInterceptors(NewInterceptor(...))` retain their previous +behavior: an `auth-required` task or status event ends the client span successfully. This keeps +persisted authorization snapshots complete without requiring a client migration, but misses later +events if the same stream continues. Use `WrapClient` to capture those events and to observe clean +EOF and early iterator abandonment. + ## Attributes | Attribute | When | diff --git a/go/a2a/client_completion_test.go b/go/a2a/client_completion_test.go new file mode 100644 index 000000000..32970dbbd --- /dev/null +++ b/go/a2a/client_completion_test.go @@ -0,0 +1,82 @@ +package traceaia2a + +import ( + "context" + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + "go.opentelemetry.io/otel/codes" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" +) + +// A wrapper context is not enough: only its owning call may wait for completion. +// In particular, nested unwrapped calls must retain the legacy auth boundary. +func TestClientAuthRequiredCompletionBoundary(t *testing.T) { + for _, method := range []string{"SubscribeToTask", "SendStreamingMessage"} { + for _, eventType := range []string{"task", "status"} { + for _, scope := range []string{"none", "owned", "claimed", "other-method", "server"} { + t.Run(method+"/"+eventType+"/"+scope, func(t *testing.T) { + rec := tracetest.NewSpanRecorder() + tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(rec)) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + tracing := NewInterceptor(WithTracerProvider(tp), WithTraceConfig(TraceConfig{})) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + slot := &streamCompletion{method: method} + key := clientCompletionKey + switch scope { + case "claimed": + slot.state = &callState{method: method} + case "other-method": + slot.method = "GetTask" + case "server": + key = serverCompletionKey + } + previousOwner := slot.state + if scope != "none" { + ctx = context.WithValue(ctx, key, slot) + } + ctx, _, err := tracing.Before(ctx, &a2aclient.Request{Method: method}) + if err != nil { + t.Fatal(err) + } + status := a2a.TaskStatus{State: a2a.TaskStateAuthRequired} + var event any = &a2a.Task{ID: "auth-task", Status: status} + if eventType == "status" { + event = &a2a.TaskStatusUpdateEvent{TaskID: "auth-task", Status: status} + } + if err := tracing.After(ctx, &a2aclient.Response{Method: method, Payload: event}); err != nil { + t.Fatal(err) + } + wantEvents := 1 + if scope == "owned" { + if len(rec.Ended()) != 0 { + t.Fatal("wrapped active authorization stream ended prematurely") + } + if slot.state != ctx.Value(stateKey{}).(*callState) { + t.Fatal("wrapper did not claim the call") + } + // Retain later events rather than ending at auth-required. + _ = tracing.After(ctx, &a2aclient.Response{Method: method, Payload: &a2a.TaskStatusUpdateEvent{ + TaskID: "auth-task", Status: a2a.TaskStatus{State: a2a.TaskStateWorking}, + }}) + slot.end(nil) + wantEvents = 2 + } else if slot.state != previousOwner { + t.Fatal("unwrapped call changed an unrelated completion slot") + } + if len(rec.Started()) != 1 || len(rec.Ended()) != 1 { + t.Fatalf("want one completed span: started=%d ended=%d", len(rec.Started()), len(rec.Ended())) + } + cancel() + span := rec.Ended()[0] + if span.Status().Code != codes.Ok || len(span.Events()) != wantEvents { + t.Fatalf("wrong completion: status=%v events=%d", span.Status(), len(span.Events())) + } + }) + } + } + } +} diff --git a/go/a2a/client_snapshot_test.go b/go/a2a/client_snapshot_test.go index a6b88793a..8be13b68c 100644 --- a/go/a2a/client_snapshot_test.go +++ b/go/a2a/client_snapshot_test.go @@ -2,45 +2,101 @@ package traceaia2a_test import ( "context" + "net" + "net/http" + "net/http/httptest" + "testing" + "time" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + a2agrpc "github.com/a2aproject/a2a-go/v2/a2agrpc/v1" "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" traceaia2a "github.com/future-agi/traceAI/go/a2a" "go.opentelemetry.io/otel/codes" - "net/http/httptest" - "testing" - "time" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" ) -// Client-only coverage against an uninstrumented official server. +// Client-only coverage against an uninstrumented official server. Existing +// interceptor-only clients must retain successful auth snapshot completion +// without adopting WrapClient; wrapped clients also observe actual clean EOF. func TestClientAuthSnapshotEOF(t *testing.T) { - store := taskstore.NewInMemory(nil) - task := &a2a.Task{ID: "persisted-auth-task", ContextID: "persisted-context", Status: a2a.TaskStatus{State: a2a.TaskStateAuthRequired}} - if _, err := store.Create(context.Background(), task); err != nil { - t.Fatal(err) - } - handler := a2asrv.NewHandler(repliesWithMessage(), a2asrv.WithTaskStore(store)) - server := httptest.NewServer(a2asrv.NewJSONRPCHandler(handler)) - defer server.Close() - client, rec, _ := newClient(t, card(server.URL, false, true)) - defer client.Destroy() - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - count := 0 - for _, err := range traceaia2a.WrapClient(client).SubscribeToTask(ctx, &a2a.SubscribeToTaskRequest{ID: task.ID}) { - if err != nil { - t.Fatal(err) + for _, transport := range []string{"jsonrpc", "rest", "grpc"} { + for _, mode := range []string{"interceptor-only", "wrapped"} { + t.Run(transport+"/"+mode, func(t *testing.T) { + tp, rec := serverRecorder(t) + store := taskstore.NewInMemory(nil) + task := &a2a.Task{ID: "persisted-auth-task", ContextID: "persisted-context", Status: a2a.TaskStatus{ + State: a2a.TaskStateAuthRequired, + Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("authorize this task")), + }} + if _, err := store.Create(context.Background(), task); err != nil { + t.Fatal(err) + } + handler := a2asrv.NewHandler(repliesWithMessage(), a2asrv.WithTaskStore(store)) + opts := []a2aclient.FactoryOption{a2aclient.WithCallInterceptors(traceaia2a.NewInterceptor( + traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(traceaia2a.TraceConfig{})))} + var c *a2a.AgentCard + if transport == "grpc" { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + gs := grpc.NewServer() + a2agrpc.NewHandler(handler).RegisterWith(gs) + go func() { _ = gs.Serve(ln) }() + t.Cleanup(gs.Stop) + c = &a2a.AgentCard{SupportedInterfaces: []*a2a.AgentInterface{a2a.NewAgentInterface(ln.Addr().String(), a2a.TransportProtocolGRPC)}, Capabilities: a2a.AgentCapabilities{Streaming: true}} + opts = append(opts, a2agrpc.WithGRPCTransport(grpc.WithTransportCredentials(insecure.NewCredentials()))) + } else { + var h http.Handler = a2asrv.NewJSONRPCHandler(handler) + if transport == "rest" { + h = a2asrv.NewRESTHandler(handler) + } + server := httptest.NewServer(h) + t.Cleanup(server.Close) + c = card(server.URL, transport == "rest", true) + } + client, err := a2aclient.NewFromCard(context.Background(), c, opts...) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = client.Destroy() }) + subscribe := client.SubscribeToTask + if mode == "wrapped" { + subscribe = traceaia2a.WrapClient(client).SubscribeToTask + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + count := 0 + for event, err := range subscribe(ctx, &a2a.SubscribeToTaskRequest{ID: task.ID}) { + if err != nil { + t.Fatal(err) + } + if got, ok := event.(*a2a.Task); !ok || got.ID != task.ID || got.Status.State != a2a.TaskStateAuthRequired { + t.Fatalf("unexpected snapshot: %v", event) + } + count++ + } + if count != 1 || ctx.Err() != nil || len(rec.Started()) != 1 { + t.Fatalf("not a clean single-span EOF: events=%d ctx=%v started=%d", count, ctx.Err(), len(rec.Started())) + } + if len(rec.Ended()) != 1 { + t.Error("successful subscription left client span unfinished at EOF") + } + cancel() + waitEnded(t, rec, 1) + span := rec.Ended()[0] + if span.Status().Code != codes.Ok { + t.Errorf("successful call falsely recorded as %v after caller cancellation", span.Status()) + } + if str(attrs(span), traceaia2a.A2ATaskState) != string(a2a.TaskStateAuthRequired) || + str(attrs(span), traceaia2a.OutputValue) != "authorize this task" || len(span.Events()) != 1 { + t.Fatalf("lost snapshot state, output or event: %v", attrs(span)) + } + }) } - count++ - } - if count != 1 || ctx.Err() != nil { - t.Fatalf("not a clean EOF: events=%d ctx=%v", count, ctx.Err()) - } - t.Logf("clean subscription EOF: started=%d ended=%d", len(rec.Started()), len(rec.Ended())) - if len(rec.Ended()) != 1 { - t.Fatalf("successful subscription left client span unfinished") - } - if rec.Ended()[0].Status().Code != codes.Ok { - t.Fatalf("incorrect completion: %v", rec.Ended()[0].Status()) } } diff --git a/go/a2a/interceptor.go b/go/a2a/interceptor.go index 47a3f9bc3..0a5079890 100644 --- a/go/a2a/interceptor.go +++ b/go/a2a/interceptor.go @@ -68,7 +68,9 @@ func WithPropagator(p propagation.TextMapPropagator) Option { // traceparent (and tracestate) service parameters, which the JSON-RPC and REST // transports send as HTTP headers and the gRPC transport as call metadata. // Use [WrapClient] around the created client for clean-EOF and early-abandonment -// completion; the SDK's interceptor hooks alone do not observe iterator close. +// completion and to keep recording after auth-required. Without it, auth-required +// ends the client span as in earlier releases; the SDK's interceptor hooks alone +// do not observe iterator close. type Interceptor struct { tracer trace.Tracer tc TraceConfig @@ -183,6 +185,10 @@ func (i *Interceptor) Before(ctx context.Context, req *a2aclient.Request) (conte } } + // Keep the interceptor-only client's historical auth-required completion. + // Only a call claimed by WrapClient can wait for actual iterator completion; + // an inherited wrapper slot may belong to a different (outer) call. + st.authRequiredEndsStream = !attachCompletion(ctx, clientCompletionKey, st) if st.streaming { // A stream ends without a callback when the caller stops reading // early; the span then ends when the caller's context is done. @@ -193,7 +199,6 @@ func (i *Interceptor) Before(ctx context.Context, req *a2aclient.Request) (conte st.stopWatch = context.AfterFunc(ctx, func() { st.endWithContextError(spanCtx) }) st.mu.Unlock() } - attachCompletion(ctx, clientCompletionKey, st) return context.WithValue(ctx, stateKey{}, st), nil, nil } @@ -243,6 +248,9 @@ type callState struct { ended bool stopWatch func() bool + // Legacy client-only fallback; servers and wrapped clients leave this false. + authRequiredEndsStream bool + sessionSet bool taskID string contextID string @@ -397,7 +405,7 @@ func (s *callState) applyEvent(payload any) bool { if v.Status.State != a2a.TaskStateUnspecified { ev = append(ev, attribute.String(A2ATaskState, string(v.Status.State))) } - done = endsStream(v.Status.State) + done = s.endsStream(v.Status.State) } case *a2a.Message: ev = append(ev, attribute.String(A2AEventType, "message")) @@ -415,7 +423,7 @@ func (s *callState) applyEvent(payload any) bool { if v.Status.Message != nil { s.final = partsText(v.Status.Message.Parts) } - done = endsStream(v.Status.State) + done = s.endsStream(v.Status.State) } case *a2a.TaskArtifactUpdateEvent: ev = append(ev, attribute.String(A2AEventType, "artifact-update")) @@ -433,11 +441,12 @@ func (s *callState) applyEvent(payload any) bool { return done } -// endsStream reports whether a2a-go closes a stream at this state: the -// terminal states and input-required. Auth-required can be followed by more -// events in the same stream in the supported SDK versions. -func endsStream(state a2a.TaskState) bool { - return state.Terminal() || state == a2a.TaskStateInputRequired +// endsStream reports a terminal/interrupted state, plus the historical +// auth-required boundary for interceptor-only clients. Wrapped clients and +// servers can receive more events after auth-required in the supported SDK. +func (s *callState) endsStream(state a2a.TaskState) bool { + return state.Terminal() || state == a2a.TaskStateInputRequired || + (s.authRequiredEndsStream && state == a2a.TaskStateAuthRequired) } func (s *callState) output() { diff --git a/go/a2a/stream.go b/go/a2a/stream.go index 335b6307e..bba3db75d 100644 --- a/go/a2a/stream.go +++ b/go/a2a/stream.go @@ -76,14 +76,18 @@ type streamCompletion struct { state *callState } -func attachCompletion(ctx context.Context, key completionKey, st *callState) { +// attachCompletion reports whether this call claimed its wrapper's completion +// slot. An inherited slot already owned by another call is not a completion hook. +func attachCompletion(ctx context.Context, key completionKey, st *callState) bool { if completion, _ := ctx.Value(key).(*streamCompletion); completion != nil { completion.mu.Lock() defer completion.mu.Unlock() if completion.state == nil && completion.method == st.method { completion.state = st + return true } } + return false } func (c *streamCompletion) end(err error) {