diff --git a/extras/scion-a2a-bridge/README.md b/extras/scion-a2a-bridge/README.md index 659c92c16..715e14e3e 100644 --- a/extras/scion-a2a-bridge/README.md +++ b/extras/scion-a2a-bridge/README.md @@ -75,6 +75,34 @@ docker run -p 8443:8443 -p 9090:9090 \ - `tasks/pushNotification/get` — list webhooks for a task - `tasks/pushNotification/delete` — remove a webhook +## Transports + +The bridge speaks all three A2A transports through the same a2a-go SDK request +handler, so behaviour (task lifecycle, streaming, push notifications) is +identical across them. + +| Transport | Enabled by | Routing | Streaming | +|-----------|-----------|---------|-----------| +| JSON-RPC 2.0 over HTTP | always (`bridge.listen_address`) | per-request, from the URL path | SSE | +| gRPC | `bridge.grpc_listen_address` | fixed: first exposed agent of the first project | server streaming | +| HTTP+JSON (REST) | `bridge.rest_listen_address` | fixed: first exposed agent of the first project | SSE | + +Authentication is the same on every transport: the configured `auth.scheme` +(`apiKey`, `bearer`, `hubUAT`, `hubJWT`, or `none`) is enforced for JSON-RPC, +gRPC, and REST alike. gRPC clients pass credentials as metadata (`x-api-key` +or `authorization: Bearer …`); REST clients pass the equivalent HTTP headers. +Per-user schemes (`hubUAT`, `hubJWT`) resolve the caller identity and use it +for Hub writes on all transports. + +Auth can only be disabled per transport by explicitly setting +`bridge.grpc_insecure: true` / `bridge.rest_insecure: true`; the bridge refuses +to start with a transport enabled, `auth.scheme: none`, and no matching +insecure flag. + +Both extra transports enforce the same 1 MB request body limit as JSON-RPC and +apply a rolling write deadline to SSE streams so long-lived streams are not cut +by the server write timeout. + ## TLS The A2A server listens on plain HTTP. **TLS must be terminated at a reverse proxy** (e.g. Caddy, nginx, or a cloud load balancer) in front of the bridge. The bridge logs a `WARN` at startup as a reminder. Do not expose the bridge port directly to the internet without a TLS-terminating proxy. @@ -84,6 +112,8 @@ The A2A server listens on plain HTTP. **TLS must be terminated at a reverse prox | Port | Purpose | |------|---------| | 8443 | A2A HTTP server (JSON-RPC, agent cards, health/metrics) | +| 8444 | A2A gRPC transport (optional, `bridge.grpc_listen_address`) | +| 8445 | A2A HTTP+JSON/REST transport (optional, `bridge.rest_listen_address`) | | 9090 | Broker plugin RPC (Hub connects here to push agent messages) | ## Setup and onboarding (agent instructions) @@ -131,7 +161,11 @@ Edit `scion-a2a-bridge.yaml`. The required fields are: | `hub.user` | Admin identity the bridge uses for Hub API calls | `a2a-bridge@example.com` | | `hub.signing_key` | Path to a file containing the Hub's base64-encoded HS256 signing key. Mutually exclusive with `hub.signing_key_secret`. | `/path/to/signing-key.b64` | | `hub.signing_key_secret` | GCP Secret Manager resource name for the signing key. Mutually exclusive with `hub.signing_key`. | `projects/my-project/secrets/hub-signing-key` | -| `bridge.listen_address` | Address for the A2A HTTP server | `:8443` | +| `bridge.listen_address` | Address for the A2A HTTP server (JSON-RPC) | `:8443` | +| `bridge.grpc_listen_address` | Optional. Address for the A2A gRPC transport. Empty disables it. | `:8444` | +| `bridge.rest_listen_address` | Optional. Address for the A2A HTTP+JSON (REST) transport. Empty disables it. | `:8445` | +| `bridge.grpc_insecure` | Optional. Disable auth on the gRPC transport (explicit opt-in; required if `auth.scheme` is `none`). | `false` | +| `bridge.rest_insecure` | Optional. Disable auth on the REST transport (explicit opt-in; required if `auth.scheme` is `none`). | `false` | | `bridge.external_url` | Public URL where A2A clients reach the bridge | `https://a2a.example.com` | | `auth.api_key` | Static API key clients pass in the `X-API-Key` header. Supports env var expansion. | `${A2A_API_KEY}` | | `projects[].slug` | Grove slug to expose. Add one entry per project. | `my-project` | @@ -271,7 +305,7 @@ The container runs as non-root user `bridge` (UID 1000). The state database dire ## Known Limitations -- **No gRPC or REST transport.** The bridge only supports JSON-RPC 2.0 over HTTP. gRPC and HTTP+JSON/REST transports are not implemented. +- **gRPC and REST transports are single-agent.** JSON-RPC routes per request via `/projects/{project}/agents/{agent}/jsonrpc`. The gRPC and HTTP+JSON/REST transports have no per-request routing in the A2A spec, so every request is routed to the first exposed agent of the first configured project. Run one bridge instance per exposed agent if you need more. - **Blocking-mode `input-required` flows.** In blocking mode, state-change messages are skipped for waiters so the actual content reply is delivered. A blocking `message/send` against an agent that transitions to `input-required` without sending content will time out (default 120s). Use non-blocking mode with push notifications or SSE for `input-required` flows. ## Security considerations diff --git a/extras/scion-a2a-bridge/cmd/scion-a2a-bridge/main.go b/extras/scion-a2a-bridge/cmd/scion-a2a-bridge/main.go index 23e63246e..5ed73b319 100644 --- a/extras/scion-a2a-bridge/cmd/scion-a2a-bridge/main.go +++ b/extras/scion-a2a-bridge/cmd/scion-a2a-bridge/main.go @@ -22,6 +22,7 @@ import ( "flag" "fmt" "log/slog" + "net" "net/http" "os" "os/signal" @@ -32,9 +33,12 @@ import ( secretmanager "cloud.google.com/go/secretmanager/apiv1" smpb "cloud.google.com/go/secretmanager/apiv1/secretmanagerpb" "github.com/a2aproject/a2a-go/v2/a2a" + a2agrpc "github.com/a2aproject/a2a-go/v2/a2agrpc/v0" "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" "github.com/prometheus/client_golang/prometheus" + "google.golang.org/grpc" + "google.golang.org/grpc/keepalive" "gopkg.in/yaml.v3" "github.com/GoogleCloudPlatform/scion/extras/scion-a2a-bridge/internal/bridge" @@ -176,16 +180,22 @@ func main() { srv := bridge.NewServer(b, cfg, metrics, log.With("component", "a2a-server"), sdkJSONRPCHandler) srv.WarnOnOpenAuth() + // Bound non-streaming responses: blocking sends can take up to + // Timeouts.SendMessage, so allow that plus margin. SSE responses opt out of + // this deadline via SSEWriteDeadlineMiddleware, which installs a rolling + // per-write deadline instead. + writeTimeout := cfg.Timeouts.SendMessage + 30*time.Second + httpServer := &http.Server{ Addr: listenAddr, Handler: srv.Handler(), ReadTimeout: 30 * time.Second, - WriteTimeout: 30 * time.Second, + WriteTimeout: writeTimeout, IdleTimeout: 120 * time.Second, MaxHeaderBytes: 1 << 20, } - errCh := make(chan error, 1) + errCh := make(chan error, 3) // capacity matches the 3 server goroutines (HTTP, gRPC, REST) go func() { log.Warn("A2A server starting WITHOUT TLS — ensure TLS is terminated at a reverse proxy (e.g. Caddy, nginx, cloud LB)", "address", listenAddr) log.Info("A2A protocol server starting", "address", listenAddr) @@ -194,8 +204,125 @@ func main() { } }() + // Start gRPC server if configured. + // NOTE: gRPC and REST transports require a single-project, single-agent + // configuration because they lack per-request project/agent routing. The + // executor injects the configured default route into every request context. + // Auth uses the same schemes as the JSON-RPC transport (including the + // per-user hubUAT/hubJWT schemes) and can only be disabled by explicitly + // setting bridge.grpc_insecure / bridge.rest_insecure. + var grpcServer *grpc.Server + if cfg.Bridge.GRPCListenAddress != "" { + if len(cfg.Projects) == 0 || len(cfg.Projects[0].ExposedAgents) == 0 { + log.Error("gRPC transport requires at least one project with exposed agents in config") + os.Exit(1) + } + defaultRoute := bridge.RouteInfo{ + ProjectSlug: cfg.Projects[0].Slug, + AgentSlug: cfg.Projects[0].ExposedAgents[0], + } + log.Warn("gRPC transport uses fixed routing — all requests go to the first configured agent", + "project", defaultRoute.ProjectSlug, "agent", defaultRoute.AgentSlug) + + if !cfg.Bridge.GRPCInsecure { + log.Warn("gRPC transport: auth enabled — clients must provide credentials via gRPC metadata", + "scheme", cfg.Auth.Scheme) + } else { + log.Warn("⚠ gRPC transport: auth DISABLED (grpc_insecure: true) — any client can send requests without credentials", + "address", cfg.Bridge.GRPCListenAddress) + } + + grpcServer = grpc.NewServer( + grpc.MaxRecvMsgSize(1<<20), // 1 MB, matching REST/JSON-RPC body limit + grpc.MaxConcurrentStreams(100), + grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{ + MinTime: 30 * time.Second, + PermitWithoutStream: true, + }), + grpc.ChainUnaryInterceptor( + srv.AuthUnaryInterceptor(), + bridge.RouteInfoUnaryInterceptor(defaultRoute), + ), + grpc.ChainStreamInterceptor( + srv.AuthStreamInterceptor(), + bridge.RouteInfoStreamInterceptor(defaultRoute), + ), + ) + grpcHandler := a2agrpc.NewHandler(sdkRequestHandler) + grpcHandler.RegisterWith(grpcServer) + + grpcListener, err := net.Listen("tcp", cfg.Bridge.GRPCListenAddress) + if err != nil { + log.Error("failed to listen for gRPC", "address", cfg.Bridge.GRPCListenAddress, "error", err) + os.Exit(1) + } + + go func() { + log.Info("gRPC transport starting", "address", cfg.Bridge.GRPCListenAddress) + if err := grpcServer.Serve(grpcListener); err != nil { + errCh <- fmt.Errorf("gRPC server: %w", err) + } + }() + } + + // Start REST server if configured. + var restServer *http.Server + if cfg.Bridge.RESTListenAddress != "" { + if len(cfg.Projects) == 0 || len(cfg.Projects[0].ExposedAgents) == 0 { + log.Error("REST transport requires at least one project with exposed agents in config") + os.Exit(1) + } + defaultRoute := bridge.RouteInfo{ + ProjectSlug: cfg.Projects[0].Slug, + AgentSlug: cfg.Projects[0].ExposedAgents[0], + } + log.Warn("REST transport uses fixed routing — all requests go to the first configured agent", + "project", defaultRoute.ProjectSlug, "agent", defaultRoute.AgentSlug) + if !cfg.Bridge.RESTInsecure { + log.Warn("REST transport: auth enabled — clients must provide credentials via HTTP headers", + "scheme", cfg.Auth.Scheme) + } else { + log.Warn("⚠ REST transport: auth DISABLED (rest_insecure: true) — any client can send requests without credentials", + "address", cfg.Bridge.RESTListenAddress) + } + + restHandler := srv.AuthHTTPMiddleware(bridge.MaxBytesReaderMiddleware(1<<20, + bridge.RouteInfoMiddleware(defaultRoute, + bridge.SSEWriteDeadlineMiddleware( + a2asrv.NewRESTHandler( + sdkRequestHandler, + a2asrv.WithTransportKeepAlive(cfg.Timeouts.SSEKeepalive), + ), + ), + ), + )) + + restServer = &http.Server{ + Addr: cfg.Bridge.RESTListenAddress, + Handler: restHandler, + ReadTimeout: 30 * time.Second, + WriteTimeout: writeTimeout, // SSE responses override this with a rolling deadline (SSEWriteDeadlineMiddleware). + IdleTimeout: 120 * time.Second, + MaxHeaderBytes: 1 << 20, + } + + go func() { + log.Info("REST transport starting", "address", cfg.Bridge.RESTListenAddress) + if err := restServer.ListenAndServe(); err != nil && err != http.ErrServerClosed { + errCh <- fmt.Errorf("REST server: %w", err) + } + }() + } + + transports := []string{"JSON-RPC"} + if cfg.Bridge.GRPCListenAddress != "" { + transports = append(transports, "gRPC") + } + if cfg.Bridge.RESTListenAddress != "" { + transports = append(transports, "REST") + } log.Info("scion-a2a-bridge ready", - "transport", "JSON-RPC", + "transports", transports, "sdk", "a2a-go/v2", ) @@ -217,6 +344,28 @@ func main() { log.Error("failed to stop A2A server", "error", err) } + if grpcServer != nil { + grpcStopped := make(chan struct{}) + go func() { + grpcServer.GracefulStop() + close(grpcStopped) + }() + select { + case <-grpcStopped: + log.Info("gRPC server stopped gracefully") + case <-shutdownCtx.Done(): + log.Warn("gRPC graceful shutdown timed out, forcing stop") + grpcServer.Stop() + } + } + + if restServer != nil { + if err := restServer.Shutdown(shutdownCtx); err != nil { + log.Error("failed to stop REST server", "error", err) + } + log.Info("REST server stopped") + } + // Drain background goroutines before closing the store. b.Shutdown() diff --git a/extras/scion-a2a-bridge/go.mod b/extras/scion-a2a-bridge/go.mod index b5cab5e47..dc744e86a 100644 --- a/extras/scion-a2a-bridge/go.mod +++ b/extras/scion-a2a-bridge/go.mod @@ -12,6 +12,7 @@ require ( github.com/mattn/go-sqlite3 v1.14.28 github.com/prometheus/client_golang v1.23.2 github.com/prometheus/client_model v0.6.2 + google.golang.org/grpc v1.81.1 gopkg.in/yaml.v3 v3.0.1 ) @@ -20,6 +21,7 @@ require ( cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect cloud.google.com/go/compute/metadata v0.9.0 // indirect cloud.google.com/go/iam v1.11.0 // indirect + github.com/a2aproject/a2a-go v0.3.15 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/fatih/color v1.16.0 // indirect @@ -46,6 +48,7 @@ require ( go.opentelemetry.io/otel/trace v1.43.0 // indirect go.yaml.in/yaml/v2 v2.4.3 // indirect golang.org/x/crypto v0.53.0 // indirect + golang.org/x/mod v0.36.0 // indirect golang.org/x/net v0.56.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect golang.org/x/sync v0.21.0 // indirect @@ -57,7 +60,6 @@ require ( google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260615183401-62b3387ff324 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260610212136-7ab31c22f7ad // indirect - google.golang.org/grpc v1.81.1 // indirect google.golang.org/protobuf v1.36.11 // indirect ) diff --git a/extras/scion-a2a-bridge/go.sum b/extras/scion-a2a-bridge/go.sum index 3e21bfd12..fe4db6fe9 100644 --- a/extras/scion-a2a-bridge/go.sum +++ b/extras/scion-a2a-bridge/go.sum @@ -10,6 +10,8 @@ cloud.google.com/go/iam v1.11.0 h1:KieQ9Pb+LLPak1O3Rv3GgCxhnmkYf7Xyh0P5HfF1jFM= cloud.google.com/go/iam v1.11.0/go.mod h1:KP+nKGugNJW4LcLx1uEZcq1ok5sQHFaQehQNl4QDgV4= cloud.google.com/go/secretmanager v1.16.0 h1:19QT7ZsLJ8FSP1k+4esQvuCD7npMJml6hYzilxVyT+k= cloud.google.com/go/secretmanager v1.16.0/go.mod h1://C/e4I8D26SDTz1f3TQcddhcmiC3rMEl0S1Cakvs3Q= +github.com/a2aproject/a2a-go v0.3.15 h1:h5YpCiPq3jxQ5rIns7oDjPag3ivP8u817AzdA4F+NiI= +github.com/a2aproject/a2a-go v0.3.15/go.mod h1:I7Cm+a1oL+UT6zMoP+roaRE5vdfUa1iQGVN8aSOuZ0I= github.com/a2aproject/a2a-go/v2 v2.3.1 h1:QWMdOX2UsJ8BJmjs952eo1FRyGsOVl0gFCKeM76AgGE= github.com/a2aproject/a2a-go/v2 v2.3.1/go.mod h1:mkZr8y2bUgAVQsjs/5fHK7xrRlAHDybMEyxWh2tKRC8= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= diff --git a/extras/scion-a2a-bridge/internal/bridge/bridge.go b/extras/scion-a2a-bridge/internal/bridge/bridge.go index 35839dd31..28e429a71 100644 --- a/extras/scion-a2a-bridge/internal/bridge/bridge.go +++ b/extras/scion-a2a-bridge/internal/bridge/bridge.go @@ -345,11 +345,14 @@ func (b *Bridge) SendMessage(ctx context.Context, projectSlug, agentSlug, contex aKey := agentKey(agentCtx.ProjectID, agentCtx.AgentSlug) b.registerActiveTask(taskID, aKey) responseCh := make(chan *messages.StructuredMessage, 1) - b.addWaiter(taskID, &waiter{ + if !b.addWaiter(taskID, &waiter{ ch: responseCh, agentSlug: agentCtx.AgentSlug, projectID: agentCtx.ProjectID, - }) + }) { + b.unregisterActiveTask(taskID, aKey) + return nil, fmt.Errorf("concurrent request for task %s", taskID) + } defer b.removeWaiter(taskID) // Keep task registered in activeTasks — the agent's eventual state-change // to completed/failed will close it via dispatchToActiveTask. @@ -468,7 +471,10 @@ func (b *Bridge) sendFollowUp(ctx context.Context, projectSlug, agentSlug, taskI aKey := agentKey(task.ProjectID, task.AgentSlug) b.registerActiveTask(taskID, aKey) responseCh := make(chan *messages.StructuredMessage, 1) - b.addWaiter(taskID, &waiter{ch: responseCh, agentSlug: task.AgentSlug, projectID: task.ProjectID}) + if !b.addWaiter(taskID, &waiter{ch: responseCh, agentSlug: task.AgentSlug, projectID: task.ProjectID}) { + b.unregisterActiveTask(taskID, aKey) + return nil, fmt.Errorf("concurrent follow-up for task %s: another blocking request is already waiting", taskID) + } defer b.removeWaiter(taskID) defer b.unregisterActiveTask(taskID, aKey) @@ -730,9 +736,34 @@ func (b *Bridge) dispatchBrokerMessage(topic string, msg *messages.StructuredMes // If the message carries a task correlation ID, dispatch only to that task // after verifying the message's agent matches the task's owner. if taskID := msg.Metadata["a2aTaskId"]; taskID != "" { + // Try waiter first — SDK-created tasks (via AgentExecutor) may not + // be stored in the local SQLite store, but they register a waiter + // for blocking response correlation. Check the waiter before the + // store to avoid dropping responses for SDK-managed tasks. + if b.dispatchToWaiter(taskID, msg) { + return + } + task, err := b.store.GetTask(taskID) if err != nil || task == nil { - b.log.Debug("ignoring message for unknown task", "task_id", taskID) + // Also check if the task is registered as active (SDK executor + // registers in activeTasks even without a store entry). + b.tasksMu.RLock() + entry, isActive := b.activeTasks[taskID] + b.tasksMu.RUnlock() + if !isActive { + b.log.Debug("ignoring message for unknown task", "task_id", taskID) + return + } + // Verify the message sender's agent slug matches the active + // task's registered agent key (format "projectID:agentSlug"). + if parts := strings.SplitN(entry.aKey, ":", 2); len(parts) == 2 && parts[1] != agentSlug { + b.log.Warn("dropping cross-agent message for SDK-managed task", + "task_agent", parts[1], "msg_agent", agentSlug, "task_id", taskID) + return + } + // Active but not in store — SDK-managed task, dispatch via active path. + b.dispatchToActiveTask(ctx, taskID, agentSlug, msg) return } if task.AgentSlug != agentSlug { @@ -741,9 +772,6 @@ func (b *Bridge) dispatchBrokerMessage(topic string, msg *messages.StructuredMes return } - if b.dispatchToWaiter(taskID, msg) { - return - } b.tasksMu.RLock() _, isActive := b.activeTasks[taskID] b.tasksMu.RUnlock() @@ -786,7 +814,8 @@ func (b *Bridge) dispatchBrokerMessage(topic string, msg *messages.StructuredMes // dispatchToWaiter sends a message to a blocking waiter for the given taskID. // Returns true if a waiter exists and handled the message (callers should skip // further dispatch). State-change messages are skipped so the actual reply -// lands in the buffer. +// lands in the buffer. Verifies the message sender's agent slug matches the +// waiter's expected agent to prevent cross-agent message injection. func (b *Bridge) dispatchToWaiter(taskID string, msg *messages.StructuredMessage) bool { b.mu.RLock() w, ok := b.waiters[taskID] @@ -794,6 +823,20 @@ func (b *Bridge) dispatchToWaiter(taskID string, msg *messages.StructuredMessage if !ok { return false } + // Verify agent ownership: the waiter's expected agent must match the + // message sender's agent slug. This prevents a response from Agent B + // being delivered to a task that was started for Agent A. + if w.agentSlug != "" { + senderAgent := extractAgentIDFromSender(msg.Sender) + if senderAgent != "" && senderAgent != w.agentSlug { + b.log.Warn("dropping cross-agent message for waiter", + "task_id", taskID, + "expected_agent", w.agentSlug, + "sender_agent", senderAgent, + ) + return true // consumed but rejected — don't fall through to other dispatch paths + } + } if msg.Type == messages.TypeStateChange { // Terminal state-changes must still be persisted to the DB even though // we skip the waiter — otherwise the task's stored state is never updated. @@ -1015,7 +1058,7 @@ func (b *Bridge) GenerateAgentCard(ctx context.Context, projectSlug, agentSlug s "version": "1.0.0", "capabilities": map[string]bool{ "streaming": true, - "pushNotifications": true, + "pushNotifications": false, }, "defaultInputModes": []string{"text/plain", "application/json"}, "defaultOutputModes": []string{"text/plain", "application/json"}, @@ -1212,10 +1255,17 @@ func (b *Bridge) unregisterActiveTask(taskID, aKey string) { } } -func (b *Bridge) addWaiter(taskID string, w *waiter) { +// addWaiter registers a blocking waiter for a task. Returns false if a waiter +// is already registered (concurrent follow-ups to the same task), in which case +// the caller should reject the request to prevent response mis-delivery. +func (b *Bridge) addWaiter(taskID string, w *waiter) bool { b.mu.Lock() defer b.mu.Unlock() + if _, exists := b.waiters[taskID]; exists { + return false + } b.waiters[taskID] = w + return true } func (b *Bridge) removeWaiter(taskID string) { diff --git a/extras/scion-a2a-bridge/internal/bridge/config.go b/extras/scion-a2a-bridge/internal/bridge/config.go index 306decb36..942466e54 100644 --- a/extras/scion-a2a-bridge/internal/bridge/config.go +++ b/extras/scion-a2a-bridge/internal/bridge/config.go @@ -34,10 +34,14 @@ type Config struct { // BridgeConfig holds the A2A protocol server settings. type BridgeConfig struct { - ListenAddress string `yaml:"listen_address"` - ExternalURL string `yaml:"external_url"` - MaxSubscribers int `yaml:"max_subscribers"` - Provider ProviderConfig `yaml:"provider"` + ListenAddress string `yaml:"listen_address"` + GRPCListenAddress string `yaml:"grpc_listen_address"` + RESTListenAddress string `yaml:"rest_listen_address"` + GRPCInsecure bool `yaml:"grpc_insecure"` + RESTInsecure bool `yaml:"rest_insecure"` + ExternalURL string `yaml:"external_url"` + MaxSubscribers int `yaml:"max_subscribers"` + Provider ProviderConfig `yaml:"provider"` } // ProviderConfig describes the bridge operator. diff --git a/extras/scion-a2a-bridge/internal/bridge/executor.go b/extras/scion-a2a-bridge/internal/bridge/executor.go index 940adb6a5..1061480f5 100644 --- a/extras/scion-a2a-bridge/internal/bridge/executor.go +++ b/extras/scion-a2a-bridge/internal/bridge/executor.go @@ -19,10 +19,16 @@ import ( "fmt" "iter" "log/slog" + "net/http" + "strings" "time" "github.com/a2aproject/a2a-go/v2/a2a" "github.com/a2aproject/a2a-go/v2/a2asrv" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" "github.com/GoogleCloudPlatform/scion/pkg/hubclient" "github.com/GoogleCloudPlatform/scion/pkg/messages" @@ -71,6 +77,10 @@ func NewScionExecutor(bridge *Bridge, log *slog.Logger) *ScionExecutor { // to a Scion agent and yields events as the agent responds. func (e *ScionExecutor) Execute(ctx context.Context, execCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { return func(yield func(a2a.Event, error) bool) { + if execCtx == nil { + yield(nil, fmt.Errorf("executor context is nil: %w", a2a.ErrInternalError)) + return + } route, ok := RouteInfoFrom(ctx) if !ok { yield(nil, fmt.Errorf("missing route info in context: %w", a2a.ErrInternalError)) @@ -85,6 +95,9 @@ func (e *ScionExecutor) Execute(ctx context.Context, execCtx *a2asrv.ExecutorCon } // Resolve the Scion agent context (agent ID, project ID). + // TODO(multi-turn): Pass execCtx.ContextID here to reuse existing + // Scion contexts for multi-turn conversations. Currently always creates + // a new context, breaking agents that use input-required → completed flows. agentCtx, err := e.bridge.resolveContext(ctx, route.ProjectSlug, route.AgentSlug, "") if err != nil { yield(nil, fmt.Errorf("resolve agent: %w", err)) @@ -135,13 +148,17 @@ func (e *ScionExecutor) Execute(ctx context.Context, execCtx *a2asrv.ExecutorCon e.bridge.registerActiveTask(string(taskID), aKey) defer e.bridge.unregisterActiveTask(string(taskID), aKey) - // Set up response channel. + // Set up response channel. Reject if a waiter already exists (concurrent + // request to the same task). responseCh := make(chan *messages.StructuredMessage, 1) - e.bridge.addWaiter(string(taskID), &waiter{ + if !e.bridge.addWaiter(string(taskID), &waiter{ ch: responseCh, agentSlug: agentCtx.AgentSlug, projectID: agentCtx.ProjectID, - }) + }) { + yield(nil, fmt.Errorf("concurrent request for task %s: %w", taskID, a2a.ErrInternalError)) + return + } defer e.bridge.removeWaiter(string(taskID)) // Send to Hub using the per-user or admin client. @@ -219,7 +236,7 @@ func (e *ScionExecutor) Cancel(ctx context.Context, execCtx *a2asrv.ExecutorCont return func(yield func(a2a.Event, error) bool) { taskID := execCtx.TaskID - // Look up the stored task to find the agent. + // Look up the stored task to find the agent and send an interrupt. if execCtx.StoredTask != nil && e.bridge.hubClient != nil { route, ok := RouteInfoFrom(ctx) if !ok { @@ -261,3 +278,175 @@ func (e *ScionExecutor) Cancel(ctx context.Context, execCtx *a2asrv.ExecutorCont yield(a2a.NewStatusUpdateEvent(execCtx, a2a.TaskStateCanceled, nil), nil) } } + +// SSEWriteDeadlineMiddleware wraps an http.Handler to clear the write deadline +// for SSE (text/event-stream) responses, allowing long-lived streaming +// connections while keeping WriteTimeout enabled for non-streaming endpoints. +func SSEWriteDeadlineMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + next.ServeHTTP(&sseDeadlineWriter{ResponseWriter: w}, r) + }) +} + +// sseWriteDeadline is the rolling per-write deadline for SSE connections. +// Each write resets the deadline, so an active stream stays alive while +// a stalled one is reaped after this duration. +const sseWriteDeadline = 60 * time.Second + +// sseDeadlineWriter intercepts WriteHeader and Write to apply a rolling +// per-write deadline for SSE streams (Content-Type: text/event-stream). +// Non-SSE responses are passed through unchanged. +type sseDeadlineWriter struct { + http.ResponseWriter + isSSE bool + checked bool +} + +// detectSSE checks the Content-Type header once and caches the result. +func (s *sseDeadlineWriter) detectSSE() { + if !s.checked { + ct := s.ResponseWriter.Header().Get("Content-Type") + s.isSSE = strings.HasPrefix(ct, "text/event-stream") + s.checked = true + } +} + +// extendDeadline sets a rolling write deadline for SSE connections. +func (s *sseDeadlineWriter) extendDeadline() { + s.detectSSE() + if s.isSSE { + rc := http.NewResponseController(s.ResponseWriter) + _ = rc.SetWriteDeadline(time.Now().Add(sseWriteDeadline)) + } +} + +func (s *sseDeadlineWriter) WriteHeader(code int) { + s.extendDeadline() + s.ResponseWriter.WriteHeader(code) +} + +func (s *sseDeadlineWriter) Write(b []byte) (int, error) { + s.extendDeadline() + return s.ResponseWriter.Write(b) +} + +func (s *sseDeadlineWriter) Flush() { + if f, ok := s.ResponseWriter.(http.Flusher); ok { + f.Flush() + } +} + +func (s *sseDeadlineWriter) Unwrap() http.ResponseWriter { + return s.ResponseWriter +} + +// MaxBytesReaderMiddleware wraps an http.Handler to limit the request body size, +// preventing memory exhaustion from oversized payloads. This mirrors the +// MaxBytesReader applied in the JSON-RPC path (server.go handleJSONRPC). +func MaxBytesReaderMiddleware(maxBytes int64, next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + r.Body = http.MaxBytesReader(w, r.Body, maxBytes) + next.ServeHTTP(w, r) + }) +} + +// RouteInfoMiddleware wraps an http.Handler to inject a fixed RouteInfo into the +// request context. Used for transports (REST) that don't have per-request +// project/agent routing. +func RouteInfoMiddleware(route RouteInfo, next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx := WithRouteInfo(r.Context(), route) + next.ServeHTTP(w, r.WithContext(ctx)) + }) +} + +// grpcHeaderLookup adapts gRPC incoming metadata to a headerLookup. gRPC +// metadata keys are always lowercase, so canonical HTTP names are lowered here. +func grpcHeaderLookup(ctx context.Context) (headerLookup, bool) { + md, ok := metadata.FromIncomingContext(ctx) + if !ok { + return nil, false + } + return func(name string) string { + if vals := md.Get(strings.ToLower(name)); len(vals) > 0 { + return vals[0] + } + return "" + }, true +} + +// authenticateGRPC runs the shared authenticator against gRPC metadata and maps +// failures onto gRPC status codes. +func (s *Server) authenticateGRPC(ctx context.Context) (context.Context, error) { + if s.config.Auth.Scheme == "none" || s.config.Bridge.GRPCInsecure { + return ctx, nil + } + lookup, ok := grpcHeaderLookup(ctx) + if !ok { + return nil, status.Error(codes.Unauthenticated, "missing metadata") + } + authCtx, authErr := s.authenticate(ctx, lookup) + if authErr != nil { + if authErr.internal { + return nil, status.Error(codes.Internal, authErr.msg) + } + return nil, status.Error(codes.Unauthenticated, authErr.msg) + } + return authCtx, nil +} + +// AuthUnaryInterceptor returns a gRPC unary interceptor that validates caller +// credentials using the same schemes as the HTTP transports (including the +// per-user hubUAT/hubJWT schemes, which inject a CallerIdentity into the +// context). Auth is skipped when auth.scheme is "none" or bridge.grpc_insecure +// is set. +func (s *Server) AuthUnaryInterceptor() grpc.UnaryServerInterceptor { + return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { + authCtx, err := s.authenticateGRPC(ctx) + if err != nil { + return nil, err + } + return handler(authCtx, req) + } +} + +// AuthStreamInterceptor returns a gRPC stream interceptor with the same +// semantics as AuthUnaryInterceptor. +func (s *Server) AuthStreamInterceptor() grpc.StreamServerInterceptor { + return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + authCtx, err := s.authenticateGRPC(ss.Context()) + if err != nil { + return err + } + return handler(srv, &ctxServerStream{ServerStream: ss, ctx: authCtx}) + } +} + +// RouteInfoUnaryInterceptor returns a gRPC unary server interceptor that injects +// a fixed RouteInfo into the request context. +func RouteInfoUnaryInterceptor(route RouteInfo) grpc.UnaryServerInterceptor { + return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { + return handler(WithRouteInfo(ctx, route), req) + } +} + +// RouteInfoStreamInterceptor returns a gRPC stream server interceptor that +// injects a fixed RouteInfo into the stream context. +func RouteInfoStreamInterceptor(route RouteInfo) grpc.StreamServerInterceptor { + return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + wrapped := &ctxServerStream{ServerStream: ss, ctx: WithRouteInfo(ss.Context(), route)} + return handler(srv, wrapped) + } +} + +// ctxServerStream wraps a grpc.ServerStream to override its Context, letting +// interceptors inject values (auth identity, route info) that downstream +// handlers observe. +type ctxServerStream struct { + grpc.ServerStream + ctx context.Context +} + +func (s *ctxServerStream) Context() context.Context { + return s.ctx +} diff --git a/extras/scion-a2a-bridge/internal/bridge/followup_test.go b/extras/scion-a2a-bridge/internal/bridge/followup_test.go index 320e9f496..6d25245c9 100644 --- a/extras/scion-a2a-bridge/internal/bridge/followup_test.go +++ b/extras/scion-a2a-bridge/internal/bridge/followup_test.go @@ -16,7 +16,6 @@ package bridge import ( "context" - "encoding/json" "errors" "fmt" "io" @@ -28,29 +27,15 @@ import ( "testing" "time" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" + "github.com/GoogleCloudPlatform/scion/extras/scion-a2a-bridge/internal/state" "github.com/GoogleCloudPlatform/scion/pkg/hubclient" "github.com/GoogleCloudPlatform/scion/pkg/messages" ) -// ErrCodeInvalidParams is the JSON-RPC error code for invalid params. -// Previously defined in handler.go; kept here for test compatibility. -const ErrCodeInvalidParams = -32602 - -// SendMessageParams mirrors the A2A message/send params for test assertions. -type SendMessageParams struct { - TaskID string `json:"taskId,omitempty"` - Message Message `json:"message"` - Configuration *SendMessageConfig `json:"configuration,omitempty"` -} - -// SendMessageConfig mirrors the A2A configuration block. -type SendMessageConfig struct { - Blocking *bool `json:"blocking,omitempty"` -} - -func boolPtr(b bool) *bool { return &b } - // --- Mock hubclient --- // mockAgentService implements hubclient.AgentService for testing. @@ -170,6 +155,14 @@ func (m *mockHubClient) Health(ctx context.Context) (*hubclient.HealthResponse, return &hubclient.HealthResponse{}, nil } +// DiscoverSkillsDirectory and HubPreStartHooks were added to hubclient.Client +// upstream (#914, #892) without updating this mock, which broke the bridge test +// build on main. +func (m *mockHubClient) HubPreStartHooks() hubclient.HubPreStartHookService { return nil } +func (m *mockHubClient) DiscoverSkillsDirectory(ctx context.Context, req hubclient.DiscoverSkillsDirectoryRequest) (*hubclient.DiscoverSkillsDirectoryResponse, error) { + return nil, fmt.Errorf("not implemented") +} + // --- Test helpers --- // newFollowUpTestBridge creates a Bridge wired to a mock hub client and real SQLite store. @@ -816,19 +809,59 @@ func TestSendFollowUp_ResolvesAgentIDViaLookup(t *testing.T) { } } -// --- Server-layer tests for handleSendMessage with TaskID --- +// --- Server-layer tests for message/send via SDK handler --- -func TestHandleSendMessage_PassesTaskIDToSendMessage(t *testing.T) { +// newFollowUpTestServerWithHub creates a test server wired to a custom mockHubClient, +// including the SDK executor and JSON-RPC handler (matching newTestServer pattern). +func newFollowUpTestServerWithHub(t *testing.T, hub hubclient.Client, cfg *Config) (*Server, *httptest.Server, *state.Store) { + t.Helper() dir := t.TempDir() store, err := state.New(filepath.Join(dir, "test.db")) if err != nil { t.Fatalf("state.New: %v", err) } - defer store.Close() + t.Cleanup(func() { store.Close() }) + log := slog.New(slog.NewTextHandler(io.Discard, nil)) + b := New(store, hub, nil, cfg, nil, log) + t.Cleanup(func() { b.Shutdown() }) + + executor := NewScionExecutor(b, log) + routeAuth := RouteKeyAuthenticator() + innerStore := taskstore.NewInMemory(&taskstore.InMemoryStoreConfig{ + Authenticator: routeAuth, + }) + scopedStore := NewScopedTaskStore(innerStore) + sdkRequestHandler := a2asrv.NewHandler( + executor, + a2asrv.WithLogger(log), + a2asrv.WithCapabilityChecks(&a2a.AgentCapabilities{ + Streaming: true, + PushNotifications: false, + }), + a2asrv.WithAgentInactivityTimeout(cfg.Timeouts.SendMessage), + a2asrv.WithTaskStore(scopedStore), + ) + b.SetSDKRequestHandler(sdkRequestHandler) + sdkJSONRPCHandler := a2asrv.NewJSONRPCHandler(sdkRequestHandler) + + srv := NewServer(b, cfg, nil, log, sdkJSONRPCHandler) + ts := httptest.NewServer(srv.Handler()) + t.Cleanup(ts.Close) + return srv, ts, store +} + +func TestHandleSendMessage_NewMessage_RoutesToExecutor(t *testing.T) { var mu sync.Mutex var capturedMeta map[string]string agents := &mockAgentService{ + listFn: func(ctx context.Context, opts *hubclient.ListAgentsOptions) (*hubclient.ListAgentsResponse, error) { + return &hubclient.ListAgentsResponse{ + Agents: []hubclient.Agent{ + {ID: "agent-id-1", Slug: "agent-a", ProjectID: "proj-1"}, + }, + }, nil + }, sendFn: func(ctx context.Context, agentID string, msg *messages.StructuredMessage, interrupt, notify, wake bool) (*hubclient.MessageResponse, error) { mu.Lock() defer mu.Unlock() @@ -842,31 +875,22 @@ func TestHandleSendMessage_PassesTaskIDToSendMessage(t *testing.T) { Hub: HubConfig{User: "test-user"}, Auth: AuthConfig{Scheme: "apiKey", APIKey: "test-key"}, Projects: []ProjectConfig{{Slug: "proj-1", ExposedAgents: []string{"agent-a"}}}, - Timeouts: TimeoutConfig{SendMessage: 2 * time.Second}, + Timeouts: TimeoutConfig{SendMessage: 5 * time.Second}, } - log := slog.New(slog.NewTextHandler(io.Discard, nil)) hub := &mockHubClient{agents: agents} - bridge := New(store, hub, nil, cfg, nil, log) - defer bridge.Shutdown() - srv := NewServer(bridge, cfg, nil, log, testHandler()) - ts := httptest.NewServer(srv.Handler()) - defer ts.Close() - - seedTask(t, store, "existing-task", "ctx-1", "proj-1", "agent-a", "aid", TaskStateWorking) - - params := SendMessageParams{ - TaskID: "existing-task", - Message: Message{ - Role: RoleUser, - Parts: []Part{{Text: "follow up"}}, - }, - Configuration: &SendMessageConfig{ - Blocking: boolPtr(false), + _, ts, _ := newFollowUpTestServerWithHub(t, hub, cfg) + + // SDK SendMessage: the message contains no taskId → new task. + params := map[string]interface{}{ + "message": map[string]interface{}{ + "messageId": "test-msg-1", + "role": "user", + "parts": []map[string]interface{}{{"text": "follow up"}}, }, } rpcResp := doRPC(t, ts, "/projects/proj-1/agents/agent-a/jsonrpc", - "message/send", params, "test-key") + "SendMessage", params, "test-key") if rpcResp.Error != nil { t.Fatalf("unexpected error: code=%d msg=%s", rpcResp.Error.Code, rpcResp.Error.Message) @@ -891,114 +915,43 @@ func TestHandleSendMessage_PassesTaskIDToSendMessage(t *testing.T) { mu.Lock() defer mu.Unlock() - if capturedMeta["a2aTaskId"] != "existing-task" { - t.Errorf("metadata a2aTaskId = %q, want %q", capturedMeta["a2aTaskId"], "existing-task") - } -} - -func TestHandleSendMessage_ErrTaskTerminal_ReturnsCorrectError(t *testing.T) { - dir := t.TempDir() - store, err := state.New(filepath.Join(dir, "test.db")) - if err != nil { - t.Fatalf("state.New: %v", err) - } - defer store.Close() - - agents := &mockAgentService{} - cfg := &Config{ - Bridge: BridgeConfig{ExternalURL: "https://test.example.com"}, - Hub: HubConfig{User: "test-user"}, - Auth: AuthConfig{Scheme: "apiKey", APIKey: "test-key"}, - Projects: []ProjectConfig{{Slug: "proj-1", ExposedAgents: []string{"agent-a"}}}, - } - log := slog.New(slog.NewTextHandler(io.Discard, nil)) - hub := &mockHubClient{agents: agents} - bridge := New(store, hub, nil, cfg, nil, log) - defer bridge.Shutdown() - srv := NewServer(bridge, cfg, nil, log, testHandler()) - ts := httptest.NewServer(srv.Handler()) - defer ts.Close() - - seedTask(t, store, "done-task", "ctx-1", "proj-1", "agent-a", "aid", TaskStateCompleted) - - params := SendMessageParams{ - TaskID: "done-task", - Message: Message{ - Role: RoleUser, - Parts: []Part{{Text: "try to follow up"}}, - }, - } - - rpcResp := doRPC(t, ts, "/projects/proj-1/agents/agent-a/jsonrpc", - "message/send", params, "test-key") - - if rpcResp.Error == nil { - t.Fatal("expected error for terminal task") - } - if rpcResp.Error.Code != ErrCodeInvalidParams { - t.Errorf("error code = %d, want %d", rpcResp.Error.Code, ErrCodeInvalidParams) - } - if rpcResp.Error.Message != "task is in a terminal state" { - t.Errorf("error message = %q, want %q", rpcResp.Error.Message, "task is in a terminal state") + if capturedMeta["a2aTaskId"] == "" { + t.Error("expected non-empty a2aTaskId in metadata") } } -func TestHandleSendMessage_UnknownTaskID_ReturnsAgentNotFound(t *testing.T) { - dir := t.TempDir() - store, err := state.New(filepath.Join(dir, "test.db")) - if err != nil { - t.Fatalf("state.New: %v", err) - } - defer store.Close() - - agents := &mockAgentService{} +func TestHandleSendMessage_NoHubClient_ReturnsError(t *testing.T) { cfg := &Config{ Bridge: BridgeConfig{ExternalURL: "https://test.example.com"}, Hub: HubConfig{User: "test-user"}, Auth: AuthConfig{Scheme: "apiKey", APIKey: "test-key"}, Projects: []ProjectConfig{{Slug: "proj-1", ExposedAgents: []string{"agent-a"}}}, } - log := slog.New(slog.NewTextHandler(io.Discard, nil)) - hub := &mockHubClient{agents: agents} - bridge := New(store, hub, nil, cfg, nil, log) - defer bridge.Shutdown() - srv := NewServer(bridge, cfg, nil, log, testHandler()) - ts := httptest.NewServer(srv.Handler()) - defer ts.Close() + // nil hub client → executor returns error + _, ts, _ := newFollowUpTestServerWithHub(t, nil, cfg) - params := SendMessageParams{ - TaskID: "no-such-task", - Message: Message{ - Role: RoleUser, - Parts: []Part{{Text: "follow up"}}, + params := map[string]interface{}{ + "message": map[string]interface{}{ + "messageId": "test-msg-2", + "role": "user", + "parts": []map[string]interface{}{{"text": "try it"}}, }, } rpcResp := doRPC(t, ts, "/projects/proj-1/agents/agent-a/jsonrpc", - "message/send", params, "test-key") + "SendMessage", params, "test-key") if rpcResp.Error == nil { - t.Fatal("expected error for unknown task ID") + t.Fatal("expected error when hub client is nil") } - if rpcResp.Error.Code != ErrCodeInvalidParams { - t.Errorf("error code = %d, want %d", rpcResp.Error.Code, ErrCodeInvalidParams) - } - if rpcResp.Error.Message != "agent not found" { - t.Errorf("error message = %q, want %q", rpcResp.Error.Message, "agent not found") + // The SDK wraps internal errors with a negative error code. + if rpcResp.Error.Code >= 0 { + t.Errorf("expected negative error code, got %d", rpcResp.Error.Code) } } func TestHandleSendMessage_NoTaskID_RoutesToNewTask(t *testing.T) { - // When TaskID is empty, SendMessage should try to create a new task (and fail - // because there's no real hub client to resolve the context). This verifies - // the router correctly falls through to the new-task path. - dir := t.TempDir() - store, err := state.New(filepath.Join(dir, "test.db")) - if err != nil { - t.Fatalf("state.New: %v", err) - } - defer store.Close() - + // When the message contains no taskId, the SDK executor creates a new task. agents := &mockAgentService{ listFn: func(ctx context.Context, opts *hubclient.ListAgentsOptions) (*hubclient.ListAgentsResponse, error) { return &hubclient.ListAgentsResponse{ @@ -1018,58 +971,28 @@ func TestHandleSendMessage_NoTaskID_RoutesToNewTask(t *testing.T) { Projects: []ProjectConfig{{Slug: "proj-1", ExposedAgents: []string{"agent-a"}}}, Timeouts: TimeoutConfig{SendMessage: 2 * time.Second}, } - log := slog.New(slog.NewTextHandler(io.Discard, nil)) hub := &mockHubClient{agents: agents} - bridge := New(store, hub, nil, cfg, nil, log) - defer bridge.Shutdown() - srv := NewServer(bridge, cfg, nil, log, testHandler()) - ts := httptest.NewServer(srv.Handler()) - defer ts.Close() - - params := SendMessageParams{ - Message: Message{ - Role: RoleUser, - Parts: []Part{{Text: "new message"}}, - }, - Configuration: &SendMessageConfig{ - Blocking: boolPtr(false), + _, ts, _ := newFollowUpTestServerWithHub(t, hub, cfg) + + // SDK SendMessage with no taskId → new task path. + params := map[string]interface{}{ + "message": map[string]interface{}{ + "messageId": "test-msg-3", + "role": "user", + "parts": []map[string]interface{}{{"text": "new message"}}, }, } rpcResp := doRPC(t, ts, "/projects/proj-1/agents/agent-a/jsonrpc", - "message/send", params, "test-key") + "SendMessage", params, "test-key") - // Should succeed — the new task path creates a context and task. + // Should succeed — the executor creates a context and task via the SDK. if rpcResp.Error != nil { t.Fatalf("unexpected error: code=%d msg=%s", rpcResp.Error.Code, rpcResp.Error.Message) } - resultBytes, err2 := json.Marshal(rpcResp.Result) - if err2 != nil { - t.Fatalf("marshal result: %v", err2) - } - var result TaskResult - if err2 = json.Unmarshal(resultBytes, &result); err2 != nil { - t.Fatalf("unmarshal result: %v", err2) - } - - if result.ID == "" { - t.Error("expected non-empty task ID for new task") - } - if result.Status.State != TaskStateSubmitted { - t.Errorf("status.state = %q, want %q", result.Status.State, TaskStateSubmitted) - } -} - -func TestSendFollowUp_SendMessageParams_TaskIDField(t *testing.T) { - // Verify the TaskID field is correctly parsed from JSON. - raw := `{"taskId":"my-task-123","message":{"role":"user","parts":[{"text":"hi"}]}}` - var params SendMessageParams - if err := json.Unmarshal([]byte(raw), ¶ms); err != nil { - t.Fatalf("unmarshal: %v", err) - } - if params.TaskID != "my-task-123" { - t.Errorf("TaskID = %q, want %q", params.TaskID, "my-task-123") + if rpcResp.Result == nil { + t.Fatal("expected non-nil result") } } diff --git a/extras/scion-a2a-bridge/internal/bridge/scoped_store.go b/extras/scion-a2a-bridge/internal/bridge/scoped_store.go index ad078f149..540c07655 100644 --- a/extras/scion-a2a-bridge/internal/bridge/scoped_store.go +++ b/extras/scion-a2a-bridge/internal/bridge/scoped_store.go @@ -74,6 +74,8 @@ func (s *ScopedTaskStore) Create(ctx context.Context, task *a2a.Task) (taskstore } // Update verifies ownership before delegating to the inner store. +// When a task reaches a terminal state, its ownership entry is removed to +// prevent unbounded growth of the ownership map. func (s *ScopedTaskStore) Update(ctx context.Context, update *taskstore.UpdateRequest) (taskstore.TaskVersion, error) { owner, ok := ownerKey(ctx) if !ok { @@ -88,7 +90,19 @@ func (s *ScopedTaskStore) Update(ctx context.Context, update *taskstore.UpdateRe return taskstore.TaskVersionMissing, a2a.ErrTaskNotFound } - return s.inner.Update(ctx, update) + version, err := s.inner.Update(ctx, update) + if err != nil { + return version, err + } + + // Clean up ownership entry when the task reaches a terminal state. + if isSDKTerminalState(update.Task.Status.State) { + s.mu.Lock() + delete(s.ownership, update.Task.ID) + s.mu.Unlock() + } + + return version, nil } // Get retrieves a task and verifies that the caller owns it. @@ -117,6 +131,17 @@ func (s *ScopedTaskStore) List(ctx context.Context, req *a2a.ListTasksRequest) ( return s.inner.List(ctx, req) } +// isSDKTerminalState returns true if the SDK task state is terminal (completed, +// failed, canceled, rejected). Used to clean up ownership entries. +func isSDKTerminalState(state a2a.TaskState) bool { + switch state { + case a2a.TaskStateCompleted, a2a.TaskStateFailed, a2a.TaskStateCanceled, a2a.TaskStateRejected: + return true + default: + return false + } +} + // RouteKeyAuthenticator returns a taskstore.Authenticator that derives the // "user" identity from the RouteInfo in the request context. This ensures // the in-memory task store's built-in user-filtering on List matches tasks diff --git a/extras/scion-a2a-bridge/internal/bridge/server.go b/extras/scion-a2a-bridge/internal/bridge/server.go index 785a08d9a..3b7db9452 100644 --- a/extras/scion-a2a-bridge/internal/bridge/server.go +++ b/extras/scion-a2a-bridge/internal/bridge/server.go @@ -15,6 +15,7 @@ package bridge import ( + "context" "crypto/sha256" "crypto/subtle" "encoding/json" @@ -117,6 +118,13 @@ func ValidateConfig(cfg *Config) error { return fmt.Errorf("bridge.provider.url is invalid: %w", err) } } + // Require explicit opt-in for unauthenticated gRPC/REST transports. + if cfg.Bridge.GRPCListenAddress != "" && !cfg.Bridge.GRPCInsecure && cfg.Auth.Scheme == "none" { + return fmt.Errorf("gRPC transport is configured but auth.scheme is \"none\"; set bridge.grpc_insecure: true to acknowledge unauthenticated gRPC access, or configure auth") + } + if cfg.Bridge.RESTListenAddress != "" && !cfg.Bridge.RESTInsecure && cfg.Auth.Scheme == "none" { + return fmt.Errorf("REST transport is configured but auth.scheme is \"none\"; set bridge.rest_insecure: true to acknowledge unauthenticated REST access, or configure auth") + } return nil } @@ -157,10 +165,15 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /readyz", s.handleReadyz) mux.Handle("GET /metrics", MetricsHandler()) - // Wrap with middleware chain: metrics -> rate limit -> auth. + // Wrap with middleware chain: SSE deadlines -> metrics -> rate limit -> auth. + // SSEWriteDeadlineMiddleware is outermost so it wraps the raw + // http.ResponseWriter: it replaces the server's fixed WriteTimeout with a + // rolling per-write deadline for text/event-stream responses, which keeps + // long-lived streams alive without disabling write deadlines globally. handler := s.authMiddleware(mux) handler = RateLimitMiddleware(handler, s.config.RateLimit) handler = InstrumentHandler(handler, s.metrics) + handler = SSEWriteDeadlineMiddleware(handler) return handler } @@ -288,8 +301,8 @@ func (s *Server) handleJSONRPC(w http.ResponseWriter, r *http.Request) { return } - // Enforce request body size limit to prevent memory exhaustion. - r.Body = http.MaxBytesReader(w, r.Body, 1<<20) // 1 MB + // Limit request body to 1 MB to prevent memory exhaustion from oversized payloads. + r.Body = http.MaxBytesReader(w, r.Body, 1<<20) // Inject routing info into context for the executor. ctx := WithRouteInfo(r.Context(), RouteInfo{ @@ -302,6 +315,20 @@ func (s *Server) handleJSONRPC(w http.ResponseWriter, r *http.Request) { s.sdkHandler.ServeHTTP(w, r) } +// normalizeJSONRPCID ensures only valid JSON-RPC 2.0 ID types (string, number, +// null) are echoed back. Arrays, objects, and booleans are replaced with null +// per JSON-RPC 2.0 §4.1. +func normalizeJSONRPCID(id interface{}) interface{} { + switch id.(type) { + case nil, string, float64, int, int64: + return id + case json.Number: + return id + default: + return nil + } +} + // writeJSONRPCError writes a minimal JSON-RPC error response. func writeJSONRPCError(w http.ResponseWriter, id interface{}, code int, message string) { type jsonrpcError struct { @@ -315,11 +342,102 @@ func writeJSONRPCError(w http.ResponseWriter, id interface{}, code int, message } resp := jsonrpcResponse{ JSONRPC: "2.0", - ID: id, + ID: normalizeJSONRPCID(id), Error: &jsonrpcError{Code: code, Message: message}, } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(resp) + if err := json.NewEncoder(w).Encode(resp); err != nil { + slog.Default().Error("failed to encode JSON-RPC error response", "error", err) + } +} + +// verifyCredential checks that the provided credential matches the configured API key. +func verifyCredential(provided, expected string) bool { + expectedHash := sha256.Sum256([]byte(expected)) + providedHash := sha256.Sum256([]byte(provided)) + return subtle.ConstantTimeCompare(expectedHash[:], providedHash[:]) == 1 +} + +// authError describes an authentication failure in a transport-neutral way so +// that HTTP, REST, and gRPC callers can map it onto their own error model. +type authError struct { + msg string + // internal marks a server-side misconfiguration (500) rather than a + // caller error (401). + internal bool +} + +func (e *authError) Error() string { return e.msg } + +// headerLookup returns the value of a request header/metadata entry. Keys are +// the canonical HTTP names ("Authorization", "X-API-Key"); implementations are +// responsible for any transport-specific casing (gRPC metadata is lowercase). +type headerLookup func(name string) string + +// authenticate validates caller credentials against the configured auth scheme. +// It is transport-neutral: the JSON-RPC HTTP middleware, the REST middleware, +// and the gRPC interceptors all go through it so every transport supports the +// same schemes. For the per-user schemes (hubUAT/hubJWT) the returned context +// carries the resolved CallerIdentity; legacy schemes return ctx unchanged. +func (s *Server) authenticate(ctx context.Context, header headerLookup) (context.Context, *authError) { + switch s.config.Auth.Scheme { + case "none": + return ctx, nil + + case "hubUAT": + token := bearerOrAPIKeyFrom(header) + if !strings.HasPrefix(token, "scion_pat_") { + return nil, &authError{msg: "unauthorized: expected scion_pat_* token"} + } + if s.uatValidator == nil { + s.log.Error("hubUAT scheme configured but UAT validator not initialized") + return nil, &authError{msg: "internal server error", internal: true} + } + caller, err := s.uatValidator.Validate(ctx, token) + if err != nil { + s.log.Debug("UAT validation failed", "error", err) + return nil, &authError{msg: "unauthorized"} + } + return withCallerIdentity(ctx, caller), nil + + case "hubJWT": + token := bearerFrom(header) + if token == "" { + return nil, &authError{msg: "unauthorized: missing bearer token"} + } + if s.jwtValidator == nil { + s.log.Error("hubJWT scheme configured but JWT validator not initialized") + return nil, &authError{msg: "internal server error", internal: true} + } + caller, err := s.jwtValidator.Validate(token) + if err != nil { + s.log.Debug("JWT validation failed", "error", err) + return nil, &authError{msg: "unauthorized"} + } + return withCallerIdentity(ctx, caller), nil + + default: + // Legacy schemes: "apiKey", "bearer", or "" (accept either header). + // No CallerIdentity is injected. + var apiKey string + switch s.config.Auth.Scheme { + case "apiKey": + apiKey = header("X-API-Key") + case "bearer": + apiKey = bearerFrom(header) + default: + // When auth.scheme is unset (empty), accept credentials from either + // X-API-Key or Authorization: Bearer headers for convenience. + apiKey = header("X-API-Key") + if apiKey == "" { + apiKey = bearerFrom(header) + } + } + if !verifyCredential(apiKey, s.config.Auth.APIKey) { + return nil, &authError{msg: "unauthorized"} + } + return ctx, nil + } } // authMiddleware validates authentication on non-public endpoints. @@ -338,79 +456,43 @@ func (s *Server) authMiddleware(next http.Handler) http.Handler { return } - switch s.config.Auth.Scheme { - case "none": - next.ServeHTTP(w, r) - return - - case "hubUAT": - token := extractBearerOrAPIKey(r) - if !strings.HasPrefix(token, "scion_pat_") { - http.Error(w, "unauthorized: expected scion_pat_* token", http.StatusUnauthorized) - return - } - caller, err := s.uatValidator.Validate(r.Context(), token) - if err != nil { - s.log.Debug("UAT validation failed", "error", err) - http.Error(w, "unauthorized", http.StatusUnauthorized) - return - } - ctx := withCallerIdentity(r.Context(), caller) - next.ServeHTTP(w, r.WithContext(ctx)) + ctx, authErr := s.authenticate(r.Context(), r.Header.Get) + if authErr != nil { + writeAuthError(w, authErr) return + } + next.ServeHTTP(w, r.WithContext(ctx)) + }) +} - case "hubJWT": - token := extractBearerToken(r) - if token == "" { - http.Error(w, "unauthorized: missing bearer token", http.StatusUnauthorized) - return - } - if s.jwtValidator == nil { - s.log.Error("hubJWT scheme configured but JWT validator not initialized") - http.Error(w, "internal server error", http.StatusInternalServerError) - return - } - caller, err := s.jwtValidator.Validate(token) - if err != nil { - s.log.Debug("JWT validation failed", "error", err) - http.Error(w, "unauthorized", http.StatusUnauthorized) - return - } - ctx := withCallerIdentity(r.Context(), caller) - next.ServeHTTP(w, r.WithContext(ctx)) +// AuthHTTPMiddleware wraps an http.Handler with the configured authentication, +// for transports served outside the main JSON-RPC mux (REST). It supports the +// same schemes as authMiddleware, including per-user hubUAT/hubJWT. Auth is +// skipped only when the operator explicitly opted out via bridge.rest_insecure +// or auth.scheme: "none". +func (s *Server) AuthHTTPMiddleware(next http.Handler) http.Handler { + if s.config.Auth.Scheme == "none" || s.config.Bridge.RESTInsecure { + return next + } + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx, authErr := s.authenticate(r.Context(), r.Header.Get) + if authErr != nil { + writeAuthError(w, authErr) return - - default: - // Legacy schemes: "apiKey", "bearer", or "" (accept either header). - // No CallerIdentity is injected. - var apiKey string - switch s.config.Auth.Scheme { - case "apiKey": - apiKey = r.Header.Get("X-API-Key") - case "bearer": - apiKey = extractBearerToken(r) - default: - // When auth.scheme is unset (empty), accept credentials from either - // X-API-Key or Authorization: Bearer headers for convenience. - apiKey = r.Header.Get("X-API-Key") - if apiKey == "" { - apiKey = extractBearerToken(r) - } - } - - // Compare SHA-256 hashes to avoid leaking key length via timing. - expectedHash := sha256.Sum256([]byte(s.config.Auth.APIKey)) - providedHash := sha256.Sum256([]byte(apiKey)) - if subtle.ConstantTimeCompare(expectedHash[:], providedHash[:]) != 1 { - http.Error(w, "unauthorized", http.StatusUnauthorized) - return - } - - next.ServeHTTP(w, r) } + next.ServeHTTP(w, r.WithContext(ctx)) }) } +// writeAuthError maps an authError onto an HTTP response. +func writeAuthError(w http.ResponseWriter, err *authError) { + if err.internal { + http.Error(w, err.msg, http.StatusInternalServerError) + return + } + http.Error(w, err.msg, http.StatusUnauthorized) +} + // extractBearerToken extracts the token from an Authorization: Bearer header. func extractBearerToken(r *http.Request) string { auth := r.Header.Get("Authorization") @@ -427,3 +509,22 @@ func extractBearerOrAPIKey(r *http.Request) string { } return r.Header.Get("X-API-Key") } + +// bearerFrom extracts the token from an Authorization: Bearer header using a +// transport-neutral lookup. +func bearerFrom(header headerLookup) string { + auth := header("Authorization") + if strings.HasPrefix(auth, "Bearer ") { + return strings.TrimPrefix(auth, "Bearer ") + } + return "" +} + +// bearerOrAPIKeyFrom extracts a token from Authorization: Bearer or X-API-Key +// using a transport-neutral lookup. +func bearerOrAPIKeyFrom(header headerLookup) string { + if token := bearerFrom(header); token != "" { + return token + } + return header("X-API-Key") +} diff --git a/extras/scion-a2a-bridge/internal/bridge/server_test.go b/extras/scion-a2a-bridge/internal/bridge/server_test.go index 52e9bfbbc..10c185821 100644 --- a/extras/scion-a2a-bridge/internal/bridge/server_test.go +++ b/extras/scion-a2a-bridge/internal/bridge/server_test.go @@ -31,6 +31,7 @@ import ( "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" "github.com/GoogleCloudPlatform/scion/extras/scion-a2a-bridge/internal/state" + "github.com/GoogleCloudPlatform/scion/pkg/messages" ) // jsonRPCRequest is a test helper for constructing JSON-RPC requests. @@ -217,8 +218,8 @@ func TestWellKnownAgentCard(t *testing.T) { if caps["streaming"] != true { t.Errorf("capabilities.streaming = %v, want true", caps["streaming"]) } - if caps["pushNotifications"] != true { - t.Errorf("capabilities.pushNotifications = %v, want true", caps["pushNotifications"]) + if caps["pushNotifications"] != false { + t.Errorf("capabilities.pushNotifications = %v, want false", caps["pushNotifications"]) } } @@ -254,8 +255,8 @@ func TestPerAgentCard(t *testing.T) { if caps["streaming"] != true { t.Errorf("capabilities.streaming = %v, want true", caps["streaming"]) } - if caps["pushNotifications"] != true { - t.Errorf("capabilities.pushNotifications = %v, want true", caps["pushNotifications"]) + if caps["pushNotifications"] != false { + t.Errorf("capabilities.pushNotifications = %v, want false", caps["pushNotifications"]) } } @@ -301,7 +302,7 @@ func TestAuthMiddleware(t *testing.T) { } // JSON-RPC without auth should be rejected. - rpcReq, _ := json.Marshal(jsonRPCRequest{JSONRPC: "2.0", ID: 1, Method: "tasks/get", Params: json.RawMessage(`{"id":"x"}`)}) + rpcReq, _ := json.Marshal(jsonRPCRequest{JSONRPC: "2.0", ID: 1, Method: "GetTask", Params: json.RawMessage(`{"id":"x"}`)}) httpReq, _ := http.NewRequest(http.MethodPost, ts.URL+"/projects/test-grove/agents/test-agent/jsonrpc", bytes.NewReader(rpcReq)) httpReq.Header.Set("Content-Type", "application/json") @@ -334,7 +335,7 @@ func TestGetTaskNotFound(t *testing.T) { // The SDK handler will return TaskNotFound via its own error handling. rpcResp := doRPC(t, ts, "/projects/test-grove/agents/test-agent/jsonrpc", - "tasks/get", map[string]interface{}{"id": "nonexistent-task"}, "test-api-key") + "GetTask", map[string]interface{}{"id": "nonexistent-task"}, "test-api-key") if rpcResp.Error == nil { t.Fatal("expected error for nonexistent task") @@ -349,7 +350,7 @@ func TestUnknownMethod(t *testing.T) { _, ts, _ := newTestServer(t) rpcResp := doRPC(t, ts, "/projects/test-grove/agents/test-agent/jsonrpc", - "unknown/method", map[string]string{}, "test-api-key") + "unknown.method", map[string]string{}, "test-api-key") if rpcResp.Error == nil { t.Fatal("expected error for unknown method") @@ -364,7 +365,7 @@ func TestCancelTaskNotFound(t *testing.T) { _, ts, _ := newTestServer(t) rpcResp := doRPC(t, ts, "/projects/test-grove/agents/test-agent/jsonrpc", - "tasks/cancel", map[string]string{"id": "nonexistent-task"}, "test-api-key") + "CancelTask", map[string]string{"id": "nonexistent-task"}, "test-api-key") if rpcResp.Error == nil { t.Fatal("expected error for cancel of nonexistent task") @@ -378,7 +379,7 @@ func TestInvalidJSONRPC(t *testing.T) { rpcReq, _ := json.Marshal(map[string]interface{}{ "jsonrpc": "1.0", "id": 1, - "method": "tasks/get", + "method": "GetTask", "params": map[string]string{"id": "x"}, }) httpReq, _ := http.NewRequest(http.MethodPost, ts.URL+"/projects/test-grove/agents/test-agent/jsonrpc", bytes.NewReader(rpcReq)) @@ -429,9 +430,9 @@ func TestJSONRPCDeniesNonExposedAgent(t *testing.T) { _, ts, _ := newTestServer(t) methods := []string{ - "message/send", - "tasks/get", - "tasks/cancel", + "SendMessage", + "GetTask", + "CancelTask", } for _, method := range methods { @@ -473,7 +474,7 @@ func TestLegacyGrovePath(t *testing.T) { } // Test legacy JSON-RPC path (requires auth) - rpcReq, _ := json.Marshal(jsonRPCRequest{JSONRPC: "2.0", ID: 1, Method: "tasks/get", Params: json.RawMessage(`{"id":"x"}`)}) + rpcReq, _ := json.Marshal(jsonRPCRequest{JSONRPC: "2.0", ID: 1, Method: "GetTask", Params: json.RawMessage(`{"id":"x"}`)}) httpReq, _ := http.NewRequest(http.MethodPost, ts.URL+"/groves/test-grove/agents/test-agent/jsonrpc", bytes.NewReader(rpcReq)) httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("X-API-Key", "test-api-key") @@ -555,3 +556,199 @@ func TestRouteInfoContextMissing(t *testing.T) { t.Fatal("expected no route info in empty context") } } + +func TestRouteInfoMiddleware(t *testing.T) { + route := RouteInfo{ProjectSlug: "test-proj", AgentSlug: "test-agent"} + + var capturedRoute RouteInfo + var capturedOK bool + + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedRoute, capturedOK = RouteInfoFrom(r.Context()) + w.WriteHeader(http.StatusOK) + }) + + handler := RouteInfoMiddleware(route, inner) + ts := httptest.NewServer(handler) + defer ts.Close() + + resp, err := http.Get(ts.URL + "/test") + if err != nil { + t.Fatalf("GET: %v", err) + } + resp.Body.Close() + + if !capturedOK { + t.Fatal("expected route info in context") + } + if capturedRoute.ProjectSlug != "test-proj" || capturedRoute.AgentSlug != "test-agent" { + t.Errorf("RouteInfo = %+v, want {test-proj, test-agent}", capturedRoute) + } +} + +func TestNormalizeJSONRPCID(t *testing.T) { + tests := []struct { + name string + id interface{} + want interface{} + }{ + {"nil", nil, nil}, + {"string", "abc", "abc"}, + {"float64", float64(42), float64(42)}, + {"int", int(7), int(7)}, + {"json.Number", json.Number("99"), json.Number("99")}, + {"array rejected", []int{1, 2}, nil}, + {"object rejected", map[string]int{"a": 1}, nil}, + {"bool rejected", true, nil}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := normalizeJSONRPCID(tt.id) + if got != tt.want { + t.Errorf("normalizeJSONRPCID(%v) = %v (%T), want %v (%T)", tt.id, got, got, tt.want, tt.want) + } + }) + } +} + +func TestMaxBytesReaderOnJSONRPC(t *testing.T) { + _, ts, _ := newTestServer(t) + + // Send a body larger than 1MB. + bigBody := make([]byte, 2<<20) // 2MB + for i := range bigBody { + bigBody[i] = 'a' + } + + httpReq, _ := http.NewRequest(http.MethodPost, ts.URL+"/projects/test-grove/agents/test-agent/jsonrpc", bytes.NewReader(bigBody)) + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("X-API-Key", "test-api-key") + + resp, err := http.DefaultClient.Do(httpReq) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + // MaxBytesReader causes the SDK handler's body read to fail. + // The response should either be an error status (413) or contain + // a JSON-RPC error (parse error). We accept any non-success outcome. + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode == http.StatusOK { + // If 200, verify the JSON-RPC response contains an error. + var rpcResp jsonRPCResponse + if json.Unmarshal(body, &rpcResp) == nil && rpcResp.Error == nil { + t.Error("expected error for oversized request body") + } + } +} + +func TestValidateConfigGRPCInsecureRequired(t *testing.T) { + cfg := &Config{ + Bridge: BridgeConfig{ + ExternalURL: "https://test.example.com", + GRPCListenAddress: ":50051", + // GRPCInsecure not set + }, + Hub: HubConfig{Endpoint: "https://hub.example.com", User: "test"}, + Auth: AuthConfig{Scheme: "none"}, + } + err := ValidateConfig(cfg) + if err == nil { + t.Fatal("expected error for gRPC without grpc_insecure when auth is none") + } + if !bytes.Contains([]byte(err.Error()), []byte("grpc_insecure")) { + t.Errorf("error should mention grpc_insecure: %v", err) + } + + // With GRPCInsecure set, validation should pass (for this check). + cfg.Bridge.GRPCInsecure = true + err = ValidateConfig(cfg) + if err != nil && bytes.Contains([]byte(err.Error()), []byte("grpc_insecure")) { + t.Errorf("should not error with grpc_insecure set: %v", err) + } +} + +func TestDispatchToWaiterAgentSlugVerification(t *testing.T) { + dir := t.TempDir() + store, err := state.New(filepath.Join(dir, "waiter-test.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + + cfg := &Config{ + Bridge: BridgeConfig{ExternalURL: "https://a2a.test.example.com"}, + Timeouts: TimeoutConfig{SendMessage: 10 * time.Second}, + } + log := slog.New(slog.NewTextHandler(io.Discard, nil)) + b := New(store, nil, nil, cfg, nil, log) + defer b.Shutdown() + + // Register a waiter for agent-a. + ch := make(chan *messages.StructuredMessage, 1) + b.addWaiter("task-1", &waiter{ + ch: ch, + agentSlug: "agent-a", + projectID: "proj-1", + }) + + // Message from the correct agent should be dispatched. + correctMsg := &messages.StructuredMessage{ + Sender: "agent:agent-a", + Msg: "hello from agent-a", + } + if !b.dispatchToWaiter("task-1", correctMsg) { + t.Error("dispatchToWaiter should return true for matching agent") + } + select { + case got := <-ch: + if got.Msg != "hello from agent-a" { + t.Errorf("expected message from agent-a, got %q", got.Msg) + } + default: + t.Error("expected message in channel for matching agent") + } + + // Message from a different agent should be rejected. + wrongMsg := &messages.StructuredMessage{ + Sender: "agent:agent-b", + Msg: "hello from agent-b", + } + if !b.dispatchToWaiter("task-1", wrongMsg) { + t.Error("dispatchToWaiter should return true (consumed) even for wrong agent") + } + select { + case got := <-ch: + t.Errorf("should not receive message from wrong agent, got %q", got.Msg) + default: + // correct — message was rejected + } + + b.removeWaiter("task-1") +} + +func TestValidateConfigRESTInsecureRequired(t *testing.T) { + cfg := &Config{ + Bridge: BridgeConfig{ + ExternalURL: "https://test.example.com", + RESTListenAddress: ":8080", + // RESTInsecure not set + }, + Hub: HubConfig{Endpoint: "https://hub.example.com", User: "test"}, + Auth: AuthConfig{Scheme: "none"}, + } + err := ValidateConfig(cfg) + if err == nil { + t.Fatal("expected error for REST without rest_insecure when auth is none") + } + if !bytes.Contains([]byte(err.Error()), []byte("rest_insecure")) { + t.Errorf("error should mention rest_insecure: %v", err) + } + + cfg.Bridge.RESTInsecure = true + err = ValidateConfig(cfg) + if err != nil && bytes.Contains([]byte(err.Error()), []byte("rest_insecure")) { + t.Errorf("should not error with rest_insecure set: %v", err) + } +} diff --git a/extras/scion-a2a-bridge/internal/bridge/sse_deadline_test.go b/extras/scion-a2a-bridge/internal/bridge/sse_deadline_test.go new file mode 100644 index 000000000..e76f5d41b --- /dev/null +++ b/extras/scion-a2a-bridge/internal/bridge/sse_deadline_test.go @@ -0,0 +1,135 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bridge + +import ( + "bufio" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// sseStreamHandler writes n SSE events spaced by gap. +func sseStreamHandler(n int, gap time.Duration) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + for i := 0; i < n; i++ { + time.Sleep(gap) + if _, err := fmt.Fprintf(w, "data: event-%d\n\n", i); err != nil { + return + } + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + } + }) +} + +// readSSEEvents counts "data:" lines until the stream ends. +func readSSEEvents(t *testing.T, url string) int { + t.Helper() + resp, err := http.Get(url) + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + count := 0 + sc := bufio.NewScanner(resp.Body) + for sc.Scan() { + if strings.HasPrefix(sc.Text(), "data: ") { + count++ + } + } + return count +} + +// TestSSEWriteDeadlineMiddleware_KeepsStreamAlive verifies that an SSE stream +// outlives the server's fixed WriteTimeout when the middleware is installed, +// and (as a control) that it is cut short without it. +func TestSSEWriteDeadlineMiddleware_KeepsStreamAlive(t *testing.T) { + const ( + events = 6 + gap = 100 * time.Millisecond // total ~600ms + writeTimeout = 250 * time.Millisecond + ) + + t.Run("without middleware the stream is cut short", func(t *testing.T) { + srv := httptest.NewUnstartedServer(sseStreamHandler(events, gap)) + srv.Config.WriteTimeout = writeTimeout + srv.Start() + defer srv.Close() + + if got := readSSEEvents(t, srv.URL); got >= events { + t.Fatalf("got %d events, expected the WriteTimeout to truncate the stream", got) + } + }) + + t.Run("with middleware the stream completes", func(t *testing.T) { + // The middleware installs a rolling per-write deadline; use a short one + // so the test stays fast while still exceeding the server WriteTimeout. + srv := httptest.NewUnstartedServer(SSEWriteDeadlineMiddleware(sseStreamHandler(events, gap))) + srv.Config.WriteTimeout = writeTimeout + srv.Start() + defer srv.Close() + + if got := readSSEEvents(t, srv.URL); got != events { + t.Fatalf("got %d events, want %d", got, events) + } + }) +} + +// TestSSEWriteDeadlineMiddleware_NonSSEUnaffected verifies that ordinary JSON +// responses still honour the server's WriteTimeout (the middleware must not +// disable write deadlines globally). +func TestSSEWriteDeadlineMiddleware_NonSSEUnaffected(t *testing.T) { + slow := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + for i := 0; i < 6; i++ { + time.Sleep(100 * time.Millisecond) + if _, err := fmt.Fprintf(w, `{"chunk":%d}`+"\n", i); err != nil { + return + } + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + } + }) + + srv := httptest.NewUnstartedServer(SSEWriteDeadlineMiddleware(slow)) + srv.Config.WriteTimeout = 250 * time.Millisecond + srv.Start() + defer srv.Close() + + resp, err := http.Get(srv.URL) + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + count := 0 + sc := bufio.NewScanner(resp.Body) + for sc.Scan() { + if strings.HasPrefix(sc.Text(), `{"chunk"`) { + count++ + } + } + if count >= 6 { + t.Fatalf("non-SSE response returned %d chunks; WriteTimeout should have truncated it", count) + } +} diff --git a/extras/scion-a2a-bridge/internal/bridge/translate.go b/extras/scion-a2a-bridge/internal/bridge/translate.go index 83cdd921e..294053b44 100644 --- a/extras/scion-a2a-bridge/internal/bridge/translate.go +++ b/extras/scion-a2a-bridge/internal/bridge/translate.go @@ -16,6 +16,7 @@ package bridge import ( "encoding/json" + "log/slog" "strings" "time" @@ -126,12 +127,14 @@ func TranslateA2AToScion(parts []Part) *messages.StructuredMessage { attachments = append(attachments, part.URL) case part.Data != nil: jsonBytes, err := json.Marshal(part.Data) - if err == nil { - if textContent.Len() > 0 { - textContent.WriteString("\n") - } - textContent.WriteString(string(jsonBytes)) + if err != nil { + slog.Default().Warn("failed to marshal A2A Data part, skipping", "error", err) + continue } + if textContent.Len() > 0 { + textContent.WriteString("\n") + } + textContent.WriteString(string(jsonBytes)) } } @@ -155,6 +158,9 @@ func TranslateA2AToScion(parts []Part) *messages.StructuredMessage { // TranslateScionToA2A converts a Scion StructuredMessage into an A2A Message and optional Artifacts. func TranslateScionToA2A(msg *messages.StructuredMessage) (Message, []Artifact) { + if msg == nil { + return Message{MessageID: uuid.New().String(), Role: RoleAgent}, nil + } parts := []Part{{Text: msg.Msg, MediaType: "text/plain"}} for _, att := range msg.Attachments { @@ -198,12 +204,14 @@ func TranslateA2APartsToScion(parts a2a.ContentParts) *messages.StructuredMessag attachments = append(attachments, string(v)) case a2a.Data: jsonBytes, err := json.Marshal(v.Value) - if err == nil { - if textContent.Len() > 0 { - textContent.WriteString("\n") - } - textContent.WriteString(string(jsonBytes)) + if err != nil { + slog.Default().Warn("failed to marshal A2A Data part, skipping", "error", err) + continue + } + if textContent.Len() > 0 { + textContent.WriteString("\n") } + textContent.WriteString(string(jsonBytes)) } } diff --git a/extras/scion-a2a-bridge/internal/bridge/translate_test.go b/extras/scion-a2a-bridge/internal/bridge/translate_test.go index 2094841f1..5926f6b62 100644 --- a/extras/scion-a2a-bridge/internal/bridge/translate_test.go +++ b/extras/scion-a2a-bridge/internal/bridge/translate_test.go @@ -17,6 +17,8 @@ package bridge import ( "testing" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/GoogleCloudPlatform/scion/pkg/messages" ) @@ -166,3 +168,154 @@ func TestTranslateScionToA2AStateChange(t *testing.T) { t.Errorf("Artifacts = %d, want 0 for state-change messages", len(artifacts)) } } + +// --- SDK translation function tests --- + +func TestTranslateA2APartsToScionText(t *testing.T) { + parts := a2a.ContentParts{ + {Content: a2a.Text("Hello"), MediaType: "text/plain"}, + {Content: a2a.Text("World"), MediaType: "text/plain"}, + } + + msg := TranslateA2APartsToScion(parts) + + if msg.Msg != "Hello\nWorld" { + t.Errorf("Msg = %q, want %q", msg.Msg, "Hello\nWorld") + } + if msg.Type != messages.TypeInstruction { + t.Errorf("Type = %q, want %q", msg.Type, messages.TypeInstruction) + } + if msg.Version != 1 { + t.Errorf("Version = %d, want 1", msg.Version) + } + if msg.Timestamp == "" { + t.Error("expected non-empty Timestamp") + } +} + +func TestTranslateA2APartsToScionURL(t *testing.T) { + parts := a2a.ContentParts{ + {Content: a2a.Text("See this file:"), MediaType: "text/plain"}, + {Content: a2a.URL("https://example.com/file.pdf")}, + } + + msg := TranslateA2APartsToScion(parts) + + if msg.Msg != "See this file:" { + t.Errorf("Msg = %q, want %q", msg.Msg, "See this file:") + } + if len(msg.Attachments) != 1 || msg.Attachments[0] != "https://example.com/file.pdf" { + t.Errorf("Attachments = %v, want [https://example.com/file.pdf]", msg.Attachments) + } +} + +func TestTranslateA2APartsToScionData(t *testing.T) { + parts := a2a.ContentParts{ + {Content: a2a.Data{Value: map[string]interface{}{"key": "value"}}}, + } + + msg := TranslateA2APartsToScion(parts) + + if msg.Msg != `{"key":"value"}` { + t.Errorf("Msg = %q, want JSON data", msg.Msg) + } +} + +func TestTranslateA2APartsToScionEmpty(t *testing.T) { + msg := TranslateA2APartsToScion(nil) + + if msg.Msg != "[empty A2A request]" { + t.Errorf("Msg = %q, want %q", msg.Msg, "[empty A2A request]") + } +} + +func TestTranslateA2APartsToScionAttachmentOnly(t *testing.T) { + parts := a2a.ContentParts{ + {Content: a2a.URL("https://example.com/data.csv")}, + } + + msg := TranslateA2APartsToScion(parts) + + if msg.Msg != "[A2A request with attachments only]" { + t.Errorf("Msg = %q, want attachment-only placeholder", msg.Msg) + } + if len(msg.Attachments) != 1 { + t.Errorf("Attachments = %d, want 1", len(msg.Attachments)) + } +} + +func TestTranslateScionToA2AParts(t *testing.T) { + scionMsg := &messages.StructuredMessage{ + Version: 1, + Msg: "Agent response text", + Type: messages.TypeAssistantReply, + Attachments: []string{"https://example.com/output.pdf"}, + } + + message, artifacts := TranslateScionToA2AParts(scionMsg) + + if message == nil { + t.Fatal("expected non-nil message") + } + if message.Role != a2a.MessageRoleAgent { + t.Errorf("Role = %v, want %v", message.Role, a2a.MessageRoleAgent) + } + if len(message.Parts) != 2 { + t.Fatalf("Parts = %d, want 2", len(message.Parts)) + } + if text, ok := message.Parts[0].Content.(a2a.Text); !ok || string(text) != "Agent response text" { + t.Errorf("Parts[0] = %v, want Text('Agent response text')", message.Parts[0].Content) + } + if url, ok := message.Parts[1].Content.(a2a.URL); !ok || string(url) != "https://example.com/output.pdf" { + t.Errorf("Parts[1] = %v, want URL attachment", message.Parts[1].Content) + } + + if len(artifacts) != 1 { + t.Fatalf("Artifacts = %d, want 1 for assistant reply", len(artifacts)) + } + if artifacts[0].ID == "" { + t.Error("expected non-empty artifact ID") + } +} + +func TestTranslateScionToA2APartsStateChange(t *testing.T) { + scionMsg := &messages.StructuredMessage{ + Version: 1, + Msg: "State changed", + Type: messages.TypeStateChange, + } + + _, artifacts := TranslateScionToA2AParts(scionMsg) + + if len(artifacts) != 0 { + t.Errorf("Artifacts = %d, want 0 for state-change", len(artifacts)) + } +} + +func TestMapActivityToSDKTaskState(t *testing.T) { + tests := []struct { + activity string + want a2a.TaskState + }{ + {"WORKING", a2a.TaskStateWorking}, + {"THINKING", a2a.TaskStateWorking}, + {"EXECUTING", a2a.TaskStateWorking}, + {"WAITING_FOR_INPUT", a2a.TaskStateInputRequired}, + {"COMPLETED", a2a.TaskStateCompleted}, + {"ERROR", a2a.TaskStateFailed}, + {"STALLED", a2a.TaskStateFailed}, + {"LIMITS_EXCEEDED", a2a.TaskStateFailed}, + {"OFFLINE", a2a.TaskStateFailed}, + {"UNKNOWN_ACTIVITY", a2a.TaskStateWorking}, + {"working", a2a.TaskStateWorking}, + } + + for _, tt := range tests { + t.Run(tt.activity, func(t *testing.T) { + got := MapActivityToSDKTaskState(tt.activity) + if got != tt.want { + t.Errorf("MapActivityToSDKTaskState(%q) = %q, want %q", tt.activity, got, tt.want) + } + }) + } +} diff --git a/extras/scion-a2a-bridge/internal/bridge/transport_auth_test.go b/extras/scion-a2a-bridge/internal/bridge/transport_auth_test.go new file mode 100644 index 000000000..3fb1cf62f --- /dev/null +++ b/extras/scion-a2a-bridge/internal/bridge/transport_auth_test.go @@ -0,0 +1,355 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bridge + +import ( + "context" + "encoding/json" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + + "github.com/GoogleCloudPlatform/scion/extras/scion-a2a-bridge/internal/state" +) + +// newMockUATHub returns a Hub stub that accepts "Bearer scion_pat_valid". +func newMockUATHub(t *testing.T) *httptest.Server { + t.Helper() + hub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/v1/auth/me" { + http.Error(w, "not found", http.StatusNotFound) + return + } + if r.Header.Get("Authorization") == "Bearer scion_pat_valid" { + _ = json.NewEncoder(w).Encode(userResponse{ + ID: "uat-user-1", + Email: "alice@example.com", + Role: "user", + }) + return + } + http.Error(w, "unauthorized", http.StatusUnauthorized) + })) + t.Cleanup(hub.Close) + return hub +} + +// newTransportAuthServer builds a Server with the given auth/bridge config. +func newTransportAuthServer(t *testing.T, hubURL string, auth AuthConfig, br BridgeConfig) *Server { + t.Helper() + dir := t.TempDir() + store, err := state.New(filepath.Join(dir, "test.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { store.Close() }) + + br.ExternalURL = "https://test" + cfg := &Config{ + Bridge: br, + Hub: HubConfig{Endpoint: hubURL, User: "admin@test"}, + Auth: auth, + } + log := slog.New(slog.NewTextHandler(io.Discard, nil)) + b := New(store, nil, nil, cfg, nil, log) + return NewServer(b, cfg, nil, log, testHandler()) +} + +// TestRESTAuthMiddleware_AllSchemes verifies the REST transport enforces the +// same auth schemes as JSON-RPC, including per-user hubUAT identity injection. +func TestRESTAuthMiddleware_AllSchemes(t *testing.T) { + hub := newMockUATHub(t) + + tests := []struct { + name string + auth AuthConfig + bridge BridgeConfig + header string + headerVal string + wantStatus int + wantCaller string + }{ + { + name: "apiKey/valid", + auth: AuthConfig{Scheme: "apiKey", APIKey: "my-secret"}, + header: "X-API-Key", + headerVal: "my-secret", + wantStatus: http.StatusOK, + }, + { + name: "apiKey/invalid", + auth: AuthConfig{Scheme: "apiKey", APIKey: "my-secret"}, + header: "X-API-Key", + headerVal: "nope", + wantStatus: http.StatusUnauthorized, + }, + { + name: "apiKey/missing", + auth: AuthConfig{Scheme: "apiKey", APIKey: "my-secret"}, + wantStatus: http.StatusUnauthorized, + }, + { + name: "bearer/valid", + auth: AuthConfig{Scheme: "bearer", APIKey: "my-secret"}, + header: "Authorization", + headerVal: "Bearer my-secret", + wantStatus: http.StatusOK, + }, + { + name: "hubUAT/valid", + auth: AuthConfig{Scheme: "hubUAT"}, + header: "Authorization", + headerVal: "Bearer scion_pat_valid", + wantStatus: http.StatusOK, + wantCaller: "uat-user-1", + }, + { + name: "hubUAT/invalid", + auth: AuthConfig{Scheme: "hubUAT"}, + header: "Authorization", + headerVal: "Bearer scion_pat_bad", + wantStatus: http.StatusUnauthorized, + }, + { + name: "hubUAT/x-api-key-header", + auth: AuthConfig{Scheme: "hubUAT"}, + header: "X-API-Key", + headerVal: "scion_pat_valid", + wantStatus: http.StatusOK, + wantCaller: "uat-user-1", + }, + { + name: "none/bypass", + auth: AuthConfig{Scheme: "none"}, + wantStatus: http.StatusOK, + }, + { + name: "rest_insecure/bypass", + auth: AuthConfig{Scheme: "apiKey", APIKey: "my-secret"}, + bridge: BridgeConfig{RESTInsecure: true}, + wantStatus: http.StatusOK, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + srv := newTransportAuthServer(t, hub.URL, tt.auth, tt.bridge) + h := srv.AuthHTTPMiddleware(testHandler()) + + req := httptest.NewRequest(http.MethodPost, "/v1/message:send", nil) + if tt.header != "" { + req.Header.Set(tt.header, tt.headerVal) + } + w := httptest.NewRecorder() + h.ServeHTTP(w, req) + + if w.Code != tt.wantStatus { + t.Fatalf("status = %d, want %d (body: %s)", w.Code, tt.wantStatus, w.Body.String()) + } + if tt.wantStatus != http.StatusOK { + return + } + var body map[string]string + _ = json.Unmarshal(w.Body.Bytes(), &body) + if got := body["user_id"]; got != tt.wantCaller { + t.Errorf("caller user_id = %q, want %q", got, tt.wantCaller) + } + }) + } +} + +// TestRESTAuthMiddleware_HubJWT verifies hubJWT works on the REST transport and +// injects the caller identity. +func TestRESTAuthMiddleware_HubJWT(t *testing.T) { + signingKey := testSigningKey(t) + srv := newTransportAuthServer(t, "http://hub", AuthConfig{Scheme: "hubJWT"}, BridgeConfig{}) + srv.SetJWTValidator(NewJWTValidator(signingKey)) + + h := srv.AuthHTTPMiddleware(testHandler()) + token := mintTestJWT(t, signingKey, validClaims()) + + req := httptest.NewRequest(http.MethodPost, "/v1/message:send", nil) + req.Header.Set("Authorization", "Bearer "+token) + w := httptest.NewRecorder() + h.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200 (body: %s)", w.Code, w.Body.String()) + } + var body map[string]string + _ = json.Unmarshal(w.Body.Bytes(), &body) + if body["user_id"] != "user-1" || body["token_type"] != "jwt" { + t.Errorf("caller = %v, want user-1/jwt", body) + } + + // Missing token must be rejected. + w2 := httptest.NewRecorder() + h.ServeHTTP(w2, httptest.NewRequest(http.MethodPost, "/v1/message:send", nil)) + if w2.Code != http.StatusUnauthorized { + t.Errorf("missing token: status = %d, want 401", w2.Code) + } +} + +// TestRESTAuthMiddleware_MissingValidator ensures a misconfigured server fails +// closed with 500 rather than allowing the request through. +func TestRESTAuthMiddleware_MissingValidator(t *testing.T) { + srv := newTransportAuthServer(t, "http://hub", AuthConfig{Scheme: "hubJWT"}, BridgeConfig{}) + // SetJWTValidator intentionally not called. + h := srv.AuthHTTPMiddleware(testHandler()) + + req := httptest.NewRequest(http.MethodPost, "/v1/message:send", nil) + req.Header.Set("Authorization", "Bearer something") + w := httptest.NewRecorder() + h.ServeHTTP(w, req) + + if w.Code != http.StatusInternalServerError { + t.Errorf("status = %d, want 500", w.Code) + } +} + +// TestGRPCAuthInterceptors_AllSchemes verifies the gRPC transport enforces the +// same auth schemes and injects CallerIdentity for per-user schemes. +func TestGRPCAuthInterceptors_AllSchemes(t *testing.T) { + hub := newMockUATHub(t) + + tests := []struct { + name string + auth AuthConfig + bridge BridgeConfig + md metadata.MD + noMetadata bool + wantCode codes.Code + wantCaller string + }{ + { + name: "apiKey/valid", + auth: AuthConfig{Scheme: "apiKey", APIKey: "my-secret"}, + md: metadata.Pairs("x-api-key", "my-secret"), + wantCode: codes.OK, + }, + { + name: "apiKey/invalid", + auth: AuthConfig{Scheme: "apiKey", APIKey: "my-secret"}, + md: metadata.Pairs("x-api-key", "nope"), + wantCode: codes.Unauthenticated, + }, + { + name: "bearer/valid", + auth: AuthConfig{Scheme: "bearer", APIKey: "my-secret"}, + md: metadata.Pairs("authorization", "Bearer my-secret"), + wantCode: codes.OK, + }, + { + name: "bearer/invalid", + auth: AuthConfig{Scheme: "bearer", APIKey: "my-secret"}, + md: metadata.Pairs("authorization", "Bearer wrong"), + wantCode: codes.Unauthenticated, + }, + { + name: "no-metadata", + auth: AuthConfig{Scheme: "apiKey", APIKey: "my-secret"}, + noMetadata: true, + wantCode: codes.Unauthenticated, + }, + { + name: "hubUAT/valid", + auth: AuthConfig{Scheme: "hubUAT"}, + md: metadata.Pairs("authorization", "Bearer scion_pat_valid"), + wantCode: codes.OK, + wantCaller: "uat-user-1", + }, + { + name: "hubUAT/invalid", + auth: AuthConfig{Scheme: "hubUAT"}, + md: metadata.Pairs("authorization", "Bearer scion_pat_bad"), + wantCode: codes.Unauthenticated, + }, + { + name: "none/bypass", + auth: AuthConfig{Scheme: "none"}, + wantCode: codes.OK, + }, + { + name: "grpc_insecure/bypass", + auth: AuthConfig{Scheme: "apiKey", APIKey: "my-secret"}, + bridge: BridgeConfig{GRPCInsecure: true}, + wantCode: codes.OK, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + srv := newTransportAuthServer(t, hub.URL, tt.auth, tt.bridge) + + ctx := context.Background() + if !tt.noMetadata { + md := tt.md + if md == nil { + md = metadata.MD{} + } + ctx = metadata.NewIncomingContext(ctx, md) + } + + var gotCaller string + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + if c := callerIdentityFromContext(ctx); c != nil { + gotCaller = c.UserID + } + return "ok", nil + } + + // Unary. + _, err := srv.AuthUnaryInterceptor()(ctx, nil, &grpc.UnaryServerInfo{}, handler) + if got := status.Code(err); got != tt.wantCode { + t.Fatalf("unary code = %v, want %v (err: %v)", got, tt.wantCode, err) + } + if tt.wantCode == codes.OK && gotCaller != tt.wantCaller { + t.Errorf("unary caller = %q, want %q", gotCaller, tt.wantCaller) + } + + // Stream: the interceptor must also propagate the auth context. + gotCaller = "" + streamHandler := func(srv interface{}, ss grpc.ServerStream) error { + if c := callerIdentityFromContext(ss.Context()); c != nil { + gotCaller = c.UserID + } + return nil + } + err = srv.AuthStreamInterceptor()(nil, &fakeServerStream{ctx: ctx}, &grpc.StreamServerInfo{}, streamHandler) + if got := status.Code(err); got != tt.wantCode { + t.Fatalf("stream code = %v, want %v (err: %v)", got, tt.wantCode, err) + } + if tt.wantCode == codes.OK && gotCaller != tt.wantCaller { + t.Errorf("stream caller = %q, want %q", gotCaller, tt.wantCaller) + } + }) + } +} + +// fakeServerStream is a minimal grpc.ServerStream carrying a context. +type fakeServerStream struct { + grpc.ServerStream + ctx context.Context +} + +func (f *fakeServerStream) Context() context.Context { return f.ctx } diff --git a/extras/scion-a2a-bridge/internal/bridge/transport_e2e_test.go b/extras/scion-a2a-bridge/internal/bridge/transport_e2e_test.go new file mode 100644 index 000000000..96584a716 --- /dev/null +++ b/extras/scion-a2a-bridge/internal/bridge/transport_e2e_test.go @@ -0,0 +1,490 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bridge + +import ( + "context" + "io" + "log/slog" + "net" + "net/http" + "net/http/httptest" + "net/url" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + a2agrpc "github.com/a2aproject/a2a-go/v2/a2agrpc/v0" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2asrv/taskstore" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/keepalive" + + "github.com/GoogleCloudPlatform/scion/extras/scion-a2a-bridge/internal/state" + "github.com/GoogleCloudPlatform/scion/pkg/hubclient" + "github.com/GoogleCloudPlatform/scion/pkg/messages" + "github.com/GoogleCloudPlatform/scion/pkg/projectcompat" +) + +const ( + e2eProject = "proj-1" + e2eAgentSlug = "agent-a" + e2eAgentID = "agent-id-1" + e2eAPIKey = "e2e-secret" + e2eHubUser = "admin@test" +) + +// e2eStack is a fully wired bridge: real state store, real ScionExecutor, real +// a2asrv handler, with the Scion Hub replaced by an in-process fake that echoes +// each message back through the broker path (Bridge.HandleBrokerMessage), which +// is exactly how a real Scion agent's reply reaches the bridge. +type e2eStack struct { + bridge *Bridge + server *Server + sdkHandler a2asrv.RequestHandler + route RouteInfo +} + +func newE2EStack(t *testing.T) *e2eStack { + t.Helper() + + dir := t.TempDir() + store, err := state.New(filepath.Join(dir, "e2e.db")) + if err != nil { + t.Fatalf("state.New: %v", err) + } + t.Cleanup(func() { store.Close() }) + + cfg := &Config{ + Bridge: BridgeConfig{ExternalURL: "https://e2e.test"}, + Hub: HubConfig{Endpoint: "http://hub.invalid", User: e2eHubUser}, + Auth: AuthConfig{Scheme: "apiKey", APIKey: e2eAPIKey}, + Projects: []ProjectConfig{{Slug: e2eProject, ExposedAgents: []string{e2eAgentSlug}}}, + Timeouts: TimeoutConfig{SendMessage: 10 * time.Second}, + } + + stack := &e2eStack{route: RouteInfo{ProjectSlug: e2eProject, AgentSlug: e2eAgentSlug}} + + // Fake Scion agent: on Hub send, reply asynchronously over the broker path. + agents := &mockAgentService{ + listFn: func(ctx context.Context, opts *hubclient.ListAgentsOptions) (*hubclient.ListAgentsResponse, error) { + return &hubclient.ListAgentsResponse{ + Agents: []hubclient.Agent{{ID: e2eAgentID, Slug: e2eAgentSlug, ProjectID: e2eProject}}, + }, nil + }, + sendFn: func(ctx context.Context, agentID string, msg *messages.StructuredMessage, interrupt, notify, wake bool) (*hubclient.MessageResponse, error) { + reply := &messages.StructuredMessage{ + Type: msg.Type, + Sender: "agent:" + e2eAgentSlug, + Recipient: msg.Sender, + Msg: "echo: " + msg.Msg, + Metadata: map[string]string{"a2aTaskId": msg.Metadata["a2aTaskId"]}, + } + go func() { + // Small delay so the reply arrives after the waiter is armed. + time.Sleep(20 * time.Millisecond) + topic := projectcompat.UserTopic(e2eProject, e2eHubUser) + if err := stack.bridge.HandleBrokerMessage(context.Background(), topic, reply); err != nil { + t.Logf("fake agent: HandleBrokerMessage: %v", err) + } + }() + return &hubclient.MessageResponse{}, nil + }, + } + + log := slog.New(slog.NewTextHandler(io.Discard, nil)) + b := New(store, &mockHubClient{agents: agents}, nil, cfg, nil, log) + t.Cleanup(func() { b.Shutdown() }) + stack.bridge = b + + executor := NewScionExecutor(b, log) + scopedStore := NewScopedTaskStore(taskstore.NewInMemory(&taskstore.InMemoryStoreConfig{ + Authenticator: RouteKeyAuthenticator(), + })) + sdkRequestHandler := a2asrv.NewHandler( + executor, + a2asrv.WithLogger(log), + a2asrv.WithCapabilityChecks(&a2a.AgentCapabilities{Streaming: true, PushNotifications: false}), + a2asrv.WithAgentInactivityTimeout(cfg.Timeouts.SendMessage), + a2asrv.WithTaskStore(scopedStore), + ) + b.SetSDKRequestHandler(sdkRequestHandler) + stack.sdkHandler = sdkRequestHandler + + stack.server = NewServer(b, cfg, nil, log, a2asrv.NewJSONRPCHandler(sdkRequestHandler)) + return stack +} + +// startGRPC starts the gRPC transport exactly as cmd/scion-a2a-bridge does. +func (s *e2eStack) startGRPC(t *testing.T) string { + t.Helper() + grpcServer := grpc.NewServer( + grpc.MaxRecvMsgSize(1<<20), + grpc.MaxConcurrentStreams(100), + grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{ + MinTime: 30 * time.Second, + PermitWithoutStream: true, + }), + grpc.ChainUnaryInterceptor( + s.server.AuthUnaryInterceptor(), + RouteInfoUnaryInterceptor(s.route), + ), + grpc.ChainStreamInterceptor( + s.server.AuthStreamInterceptor(), + RouteInfoStreamInterceptor(s.route), + ), + ) + a2agrpc.NewHandler(s.sdkHandler).RegisterWith(grpcServer) + + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + go func() { _ = grpcServer.Serve(lis) }() + t.Cleanup(grpcServer.Stop) + return lis.Addr().String() +} + +// startREST starts the REST transport with the same middleware chain as +// cmd/scion-a2a-bridge. +func (s *e2eStack) startREST(t *testing.T) string { + t.Helper() + handler := s.server.AuthHTTPMiddleware(MaxBytesReaderMiddleware(1<<20, + RouteInfoMiddleware(s.route, + SSEWriteDeadlineMiddleware( + a2asrv.NewRESTHandler(s.sdkHandler, a2asrv.WithTransportKeepAlive(15*time.Second)), + ), + ), + )) + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + return ts.URL +} + +// apiKeyRoundTripper injects the bridge API key into every REST request. +type apiKeyRoundTripper struct { + key string + base http.RoundTripper +} + +func (rt *apiKeyRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { + r = r.Clone(r.Context()) + r.Header.Set("X-API-Key", rt.key) + base := rt.base + if base == nil { + base = http.DefaultTransport + } + return base.RoundTrip(r) +} + +// grpcAPIKeyCreds injects the bridge API key into gRPC metadata. +type grpcAPIKeyCreds struct{ key string } + +func (c grpcAPIKeyCreds) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) { + return map[string]string{"x-api-key": c.key}, nil +} +func (c grpcAPIKeyCreds) RequireTransportSecurity() bool { return false } + +// textOf extracts the concatenated text of a task's final status message. +func textOf(t *testing.T, task *a2a.Task) string { + t.Helper() + if task.Status.Message == nil { + return "" + } + var sb strings.Builder + for _, p := range task.Status.Message.Parts { + sb.WriteString(p.Text()) + } + return sb.String() +} + +// TestE2E_GRPCTransport_RoundTrip drives the gRPC transport with the real +// a2a-go gRPC client over a real TCP socket and asserts a complete +// request/response round-trip through the executor and back. +func TestE2E_GRPCTransport_RoundTrip(t *testing.T) { + stack := newE2EStack(t) + addr := stack.startGRPC(t) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + client, err := a2aclient.NewFromEndpoints(ctx, + []*a2a.AgentInterface{{ + URL: addr, + ProtocolBinding: a2a.TransportProtocolGRPC, + ProtocolVersion: a2a.ProtocolVersion("0.3"), + }}, + a2aclient.WithDefaultsDisabled(), + a2agrpc.WithGRPCTransport( + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithPerRPCCredentials(grpcAPIKeyCreds{key: e2eAPIKey}), + ), + ) + if err != nil { + t.Fatalf("a2aclient (gRPC): %v", err) + } + defer client.Destroy() + + result, err := client.SendMessage(ctx, &a2a.SendMessageRequest{ + Message: a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("hello over grpc")), + }) + if err != nil { + t.Fatalf("SendMessage over gRPC: %v", err) + } + + task, ok := result.(*a2a.Task) + if !ok { + t.Fatalf("result type = %T, want *a2a.Task", result) + } + if task.Status.State != a2a.TaskStateCompleted { + t.Fatalf("task state = %q, want completed (message: %q)", task.Status.State, textOf(t, task)) + } + if got := textOf(t, task); got != "echo: hello over grpc" { + t.Errorf("agent reply = %q, want %q", got, "echo: hello over grpc") + } + t.Logf("gRPC round-trip OK: task=%s state=%s reply=%q", task.ID, task.Status.State, textOf(t, task)) +} + +// TestE2E_GRPCTransport_Streaming exercises the streaming (server-side stream) +// path over gRPC, which also proves the stream interceptors propagate both the +// auth context and the route info. +func TestE2E_GRPCTransport_Streaming(t *testing.T) { + stack := newE2EStack(t) + addr := stack.startGRPC(t) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + client, err := a2aclient.NewFromEndpoints(ctx, + []*a2a.AgentInterface{{ + URL: addr, + ProtocolBinding: a2a.TransportProtocolGRPC, + ProtocolVersion: a2a.ProtocolVersion("0.3"), + }}, + a2aclient.WithDefaultsDisabled(), + a2agrpc.WithGRPCTransport( + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithPerRPCCredentials(grpcAPIKeyCreds{key: e2eAPIKey}), + ), + ) + if err != nil { + t.Fatalf("a2aclient (gRPC): %v", err) + } + defer client.Destroy() + + var states []a2a.TaskState + var final string + for event, err := range client.SendStreamingMessage(ctx, &a2a.SendMessageRequest{ + Message: a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("stream over grpc")), + }) { + if err != nil { + t.Fatalf("streaming event error: %v", err) + } + switch e := event.(type) { + case *a2a.Task: + states = append(states, e.Status.State) + case *a2a.TaskStatusUpdateEvent: + states = append(states, e.Status.State) + if e.Status.Message != nil { + var sb strings.Builder + for _, p := range e.Status.Message.Parts { + sb.WriteString(p.Text()) + } + if sb.Len() > 0 { + final = sb.String() + } + } + } + } + + if len(states) == 0 || states[len(states)-1] != a2a.TaskStateCompleted { + t.Fatalf("streaming states = %v, want to end in completed", states) + } + if final != "echo: stream over grpc" { + t.Errorf("streamed reply = %q, want %q", final, "echo: stream over grpc") + } + t.Logf("gRPC streaming OK: states=%v reply=%q", states, final) +} + +// TestE2E_GRPCTransport_RejectsBadCredentials proves auth is actually enforced +// on the gRPC transport end to end. +func TestE2E_GRPCTransport_RejectsBadCredentials(t *testing.T) { + stack := newE2EStack(t) + addr := stack.startGRPC(t) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + client, err := a2aclient.NewFromEndpoints(ctx, + []*a2a.AgentInterface{{ + URL: addr, + ProtocolBinding: a2a.TransportProtocolGRPC, + ProtocolVersion: a2a.ProtocolVersion("0.3"), + }}, + a2aclient.WithDefaultsDisabled(), + a2agrpc.WithGRPCTransport( + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithPerRPCCredentials(grpcAPIKeyCreds{key: "wrong-key"}), + ), + ) + if err != nil { + t.Fatalf("a2aclient (gRPC): %v", err) + } + defer client.Destroy() + + if _, err := client.SendMessage(ctx, &a2a.SendMessageRequest{ + Message: a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("nope")), + }); err == nil { + t.Fatal("expected SendMessage with bad credentials to fail") + } else { + t.Logf("gRPC auth rejection OK: %v", err) + } +} + +// TestE2E_RESTTransport_RoundTrip drives the REST transport with the real +// a2a-go REST client over a real HTTP socket. +func TestE2E_RESTTransport_RoundTrip(t *testing.T) { + stack := newE2EStack(t) + baseURL := stack.startREST(t) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + httpClient := &http.Client{ + Timeout: 30 * time.Second, + Transport: &apiKeyRoundTripper{key: e2eAPIKey}, + } + + client, err := a2aclient.NewFromEndpoints(ctx, + []*a2a.AgentInterface{{ + URL: baseURL, + ProtocolBinding: a2a.TransportProtocolHTTPJSON, + ProtocolVersion: a2a.ProtocolVersion("1.0"), + }}, + a2aclient.WithDefaultsDisabled(), + a2aclient.WithRESTTransport(httpClient), + ) + if err != nil { + t.Fatalf("a2aclient (REST): %v", err) + } + defer client.Destroy() + + result, err := client.SendMessage(ctx, &a2a.SendMessageRequest{ + Message: a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("hello over rest")), + }) + if err != nil { + t.Fatalf("SendMessage over REST: %v", err) + } + + task, ok := result.(*a2a.Task) + if !ok { + t.Fatalf("result type = %T, want *a2a.Task", result) + } + if task.Status.State != a2a.TaskStateCompleted { + t.Fatalf("task state = %q, want completed (message: %q)", task.Status.State, textOf(t, task)) + } + if got := textOf(t, task); got != "echo: hello over rest" { + t.Errorf("agent reply = %q, want %q", got, "echo: hello over rest") + } + t.Logf("REST round-trip OK: task=%s state=%s reply=%q", task.ID, task.Status.State, textOf(t, task)) +} + +// TestE2E_RESTTransport_Streaming exercises the REST SSE streaming path, which +// also covers SSEWriteDeadlineMiddleware in the live chain. +func TestE2E_RESTTransport_Streaming(t *testing.T) { + stack := newE2EStack(t) + baseURL := stack.startREST(t) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + httpClient := &http.Client{Transport: &apiKeyRoundTripper{key: e2eAPIKey}} + client, err := a2aclient.NewFromEndpoints(ctx, + []*a2a.AgentInterface{{ + URL: baseURL, + ProtocolBinding: a2a.TransportProtocolHTTPJSON, + ProtocolVersion: a2a.ProtocolVersion("1.0"), + }}, + a2aclient.WithDefaultsDisabled(), + a2aclient.WithRESTTransport(httpClient), + ) + if err != nil { + t.Fatalf("a2aclient (REST): %v", err) + } + defer client.Destroy() + + var states []a2a.TaskState + var final string + for event, err := range client.SendStreamingMessage(ctx, &a2a.SendMessageRequest{ + Message: a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("stream over rest")), + }) { + if err != nil { + t.Fatalf("streaming event error: %v", err) + } + switch e := event.(type) { + case *a2a.Task: + states = append(states, e.Status.State) + case *a2a.TaskStatusUpdateEvent: + states = append(states, e.Status.State) + if e.Status.Message != nil { + var sb strings.Builder + for _, p := range e.Status.Message.Parts { + sb.WriteString(p.Text()) + } + if sb.Len() > 0 { + final = sb.String() + } + } + } + } + + if len(states) == 0 || states[len(states)-1] != a2a.TaskStateCompleted { + t.Fatalf("streaming states = %v, want to end in completed", states) + } + if final != "echo: stream over rest" { + t.Errorf("streamed reply = %q, want %q", final, "echo: stream over rest") + } + t.Logf("REST streaming OK: states=%v reply=%q", states, final) +} + +// TestE2E_RESTTransport_RejectsBadCredentials proves auth is enforced on REST. +func TestE2E_RESTTransport_RejectsBadCredentials(t *testing.T) { + stack := newE2EStack(t) + baseURL := stack.startREST(t) + + u, err := url.Parse(baseURL) + if err != nil { + t.Fatal(err) + } + transport := a2aclient.NewRESTTransport(u, &http.Client{ + Transport: &apiKeyRoundTripper{key: "wrong-key"}, + }) + defer transport.Destroy() + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + if _, err := transport.SendMessage(ctx, a2aclient.ServiceParams{}, &a2a.SendMessageRequest{ + Message: a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("nope")), + }); err == nil { + t.Fatal("expected REST SendMessage with bad credentials to fail") + } else { + t.Logf("REST auth rejection OK: %v", err) + } +} diff --git a/extras/scion-a2a-bridge/scion-a2a-bridge.yaml.sample b/extras/scion-a2a-bridge/scion-a2a-bridge.yaml.sample index 7cdca475b..62cf9d318 100644 --- a/extras/scion-a2a-bridge/scion-a2a-bridge.yaml.sample +++ b/extras/scion-a2a-bridge/scion-a2a-bridge.yaml.sample @@ -8,6 +8,20 @@ bridge: # Address for the A2A HTTP server (JSON-RPC + agent card endpoints). listen_address: ":8443" + # Optional: enable the A2A gRPC transport. Leave unset to disable. + # gRPC has no per-request routing, so every request is routed to the first + # exposed agent of the first project below. + # grpc_listen_address: ":8444" + + # Optional: enable the A2A HTTP+JSON (REST) transport. Leave unset to + # disable. Same fixed-routing caveat as gRPC. + # rest_listen_address: ":8445" + + # Explicit opt-in required to run gRPC/REST without authentication. + # Only meaningful together with auth.scheme: "none". + # grpc_insecure: false + # rest_insecure: false + # Public-facing URL where A2A clients reach this bridge. Used to # construct agent card URLs and self-referencing links. external_url: "https://a2a.example.com"