diff --git a/go/a2a/README.md b/go/a2a/README.md index ff9f8d62f..8942f1884 100644 --- a/go/a2a/README.md +++ b/go/a2a/README.md @@ -1,20 +1,21 @@ # 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. | | | |---|---| | 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. | @@ -31,15 +32,25 @@ 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 // 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}) ``` @@ -72,7 +83,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 +92,67 @@ 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. +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 := traceaia2a.WrapHandler(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. + +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 := 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 +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 +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. + +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 -| 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 +166,17 @@ 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. +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 @@ -116,7 +194,7 @@ is when a2a-go closes the stream. | `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 @@ -129,13 +207,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 @@ -145,13 +225,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 @@ -159,9 +244,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/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/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 new file mode 100644 index 000000000..8be13b68c --- /dev/null +++ b/go/a2a/client_snapshot_test.go @@ -0,0 +1,102 @@ +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" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" +) + +// 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) { + 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)) + } + }) + } + } +} 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 5793b6d1b..5f244d007 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 @@ -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. // @@ -27,5 +30,10 @@ // 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. 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/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 be8b4b8a7..945e10e1e 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" @@ -136,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) @@ -174,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 { @@ -208,11 +208,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 +229,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 := 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} go func() { _ = srv.Serve(ln) }() @@ -258,39 +258,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/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 25f23b1d3..0a5079890 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) { @@ -67,6 +67,10 @@ 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 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 @@ -92,6 +96,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{} @@ -155,43 +164,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), @@ -211,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. @@ -234,24 +212,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 @@ -265,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 @@ -275,6 +261,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 { @@ -331,7 +360,10 @@ func (s *callState) applyTask(t *a2a.Task) { for _, a := range t.Artifacts { s.addArtifact(a) } - if endsStream(t.Status.State) && t.Status.Message != nil { + // 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) } } @@ -373,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")) @@ -387,12 +419,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 = s.endsStream(v.Status.State) } case *a2a.TaskArtifactUpdateEvent: ev = append(ev, attribute.String(A2AEventType, "artifact-update")) @@ -410,10 +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 the two interrupted ones. -func endsStream(state a2a.TaskState) bool { - return state.Terminal() || state == a2a.TaskStateInputRequired || state == a2a.TaskStateAuthRequired +// 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() { @@ -445,7 +478,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, "") } @@ -466,7 +499,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() } @@ -484,3 +517,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 new file mode 100644 index 000000000..4495e4951 --- /dev/null +++ b/go/a2a/server.go @@ -0,0 +1,154 @@ +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. +// 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. +// +// 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. 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 := traceaia2a.WrapHandler(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...)} +} + +// 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{} + +// 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 + } + info, ok := methods[callCtx.Method()] + if !ok { + return ctx, nil, nil + } + 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 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() + st.stopWatch = context.AfterFunc(ctx, func() { st.endWithContextError(spanCtx) }) + st.mu.Unlock() + attachCompletion(ctx, serverCompletionKey, st) + 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 + } + // 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.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_nested_test.go b/go/a2a/server_nested_test.go new file mode 100644 index 000000000..e3fa3302b --- /dev/null +++ b/go/a2a/server_nested_test.go @@ -0,0 +1,160 @@ +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", "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) + 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" || 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 + 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", "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) + } + 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" || nested == "policy-error" { + 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" || 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 + 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) + } + }) + } + } + } + } +} diff --git a/go/a2a/server_otlp_test.go b/go/a2a/server_otlp_test.go new file mode 100644 index 000000000..027176efd --- /dev/null +++ b/go/a2a/server_otlp_test.go @@ -0,0 +1,182 @@ +package traceaia2a_test + +import ( + "compress/gzip" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "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/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, 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") + } + 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") + } + 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) + } + 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 + 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) != 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 { + 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 + } + 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") + } + } + 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") + } +} diff --git a/go/a2a/server_test.go b/go/a2a/server_test.go new file mode 100644 index 000000000..04c163143 --- /dev/null +++ b/go/a2a/server_test.go @@ -0,0 +1,127 @@ +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, serverOpts ...a2asrv.RequestHandlerOption) *traceaia2a.Client { + t.Helper() + serverOpts = append([]a2asrv.RequestHandlerOption{a2asrv.WithCallInterceptors(traceaia2a.NewServerInterceptor( + traceaia2a.WithTracerProvider(tp), traceaia2a.WithTraceConfig(tc)))}, 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)))} + 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 traceaia2a.WrapClient(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())) + } + }) + } + } +} diff --git a/go/a2a/stream.go b/go/a2a/stream.go new file mode 100644 index 000000000..bba3db75d --- /dev/null +++ b/go/a2a/stream.go @@ -0,0 +1,146 @@ +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 +} + +// 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) { + 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