diff --git a/authbridge/authlib/pipeline/extensions.go b/authbridge/authlib/pipeline/extensions.go index 627cc2e90..b79ffaf9b 100644 --- a/authbridge/authlib/pipeline/extensions.go +++ b/authbridge/authlib/pipeline/extensions.go @@ -40,17 +40,20 @@ type MCPPromptMetadata struct { Args map[string]string } -// A2AExtension carries parsed A2A protocol metadata. +// A2AExtension carries parsed A2A protocol metadata from inbound requests. type A2AExtension struct { - TaskID string - Method string // "tasks/send", "tasks/get", etc. - Parts []A2APart + Method string // JSON-RPC method: "message/send", "message/stream" + RPCID any // JSON-RPC id for request-response correlation + SessionID string // conversation session (from params.sessionId) + MessageID string // unique message ID (from params.message.messageId) + Role string // "user" or "assistant" + Parts []A2APart } // A2APart represents a message part in an A2A request. type A2APart struct { - Type string // "text", "file", "data" - Content string + Kind string // "text", "file", "data" + Content string // text content, file URI, or serialized data } // SecurityExtension carries guardrail output. diff --git a/authbridge/authlib/plugins/a2aparser.go b/authbridge/authlib/plugins/a2aparser.go new file mode 100644 index 000000000..74e1ba4a8 --- /dev/null +++ b/authbridge/authlib/plugins/a2aparser.go @@ -0,0 +1,107 @@ +package plugins + +import ( + "context" + "encoding/json" + "log/slog" + + "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" +) + +// A2AParser parses A2A JSON-RPC 2.0 request bodies and populates +// pctx.Extensions.A2A with the parsed method, session ID, message parts, +// and role for downstream policy plugins (e.g., guardrails). +type A2AParser struct{} + +func NewA2AParser() *A2AParser { return &A2AParser{} } + +func (p *A2AParser) Name() string { return "a2a-parser" } + +func (p *A2AParser) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{ + Writes: []string{"a2a"}, + BodyAccess: true, + } +} + +func (p *A2AParser) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { + if len(pctx.Body) == 0 { + slog.Debug("a2a-parser: no body, skipping") + return pipeline.Action{Type: pipeline.Continue} + } + + var rpc jsonRPCRequest + if err := json.Unmarshal(pctx.Body, &rpc); err != nil { + slog.Debug("a2a-parser: invalid JSON-RPC", "error", err, "bodyLen", len(pctx.Body)) + return pipeline.Action{Type: pipeline.Continue} + } + + ext := &pipeline.A2AExtension{ + Method: rpc.Method, + RPCID: rpc.ID, + } + + // Extract message fields generically — any method with params.message + // gets full extraction (forward-compatible with future A2A methods). + ext.SessionID = rpc.stringParam("sessionId") + if msg := rpc.mapParam("message"); msg != nil { + if role, ok := msg["role"].(string); ok { + ext.Role = role + } + if messageID, ok := msg["messageId"].(string); ok { + ext.MessageID = messageID + } + if rawParts, ok := msg["parts"].([]any); ok { + ext.Parts = parseA2AParts(rawParts) + } + } + + pctx.Extensions.A2A = ext + + slog.Info("a2a-parser", "method", rpc.Method) + slog.Debug("a2a-parser: extracted", + "method", rpc.Method, + "sessionId", ext.SessionID, + "role", ext.Role, + "messageId", ext.MessageID, + "parts", len(ext.Parts), + ) + return pipeline.Action{Type: pipeline.Continue} +} + +func (p *A2AParser) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} + +func parseA2AParts(rawParts []any) []pipeline.A2APart { + parts := make([]pipeline.A2APart, 0, len(rawParts)) + for _, raw := range rawParts { + partMap, ok := raw.(map[string]any) + if !ok { + continue + } + kind, ok := partMap["kind"].(string) + if !ok || kind == "" { + continue + } + var content string + switch kind { + case "text": + content, _ = partMap["text"].(string) + case "file": + // TODO: update when A2A spec stabilizes — canonical Part uses mediaType + content field presence, not "kind". + content, _ = partMap["data"].(string) + if content == "" { + content, _ = partMap["uri"].(string) + } + case "data": + if dataVal, ok := partMap["data"]; ok && dataVal != nil { + if b, err := json.Marshal(dataVal); err == nil { + content = string(b) + } + } + } + parts = append(parts, pipeline.A2APart{Kind: kind, Content: content}) + } + return parts +} diff --git a/authbridge/authlib/plugins/a2aparser_test.go b/authbridge/authlib/plugins/a2aparser_test.go new file mode 100644 index 000000000..c15d71bb4 --- /dev/null +++ b/authbridge/authlib/plugins/a2aparser_test.go @@ -0,0 +1,373 @@ +package plugins + +import ( + "context" + "testing" + + "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" +) + +func TestA2AParser_Capabilities(t *testing.T) { + p := NewA2AParser() + + if p.Name() != "a2a-parser" { + t.Errorf("Name() = %q, want %q", p.Name(), "a2a-parser") + } + + caps := p.Capabilities() + if !caps.BodyAccess { + t.Error("BodyAccess should be true") + } + if len(caps.Writes) != 1 || caps.Writes[0] != "a2a" { + t.Errorf("Writes = %v, want [a2a]", caps.Writes) + } +} + +func TestA2AParser_MessageSend(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/send","id":"req-1","params":{"message":{"role":"user","parts":[{"kind":"text","text":"Hello agent"}],"messageId":"msg-001"},"sessionId":"sess-abc"}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + if pctx.Extensions.A2A == nil { + t.Fatal("Extensions.A2A is nil") + } + ext := pctx.Extensions.A2A + if ext.Method != "message/send" { + t.Errorf("Method = %q, want %q", ext.Method, "message/send") + } + if ext.RPCID != "req-1" { + t.Errorf("RPCID = %v, want %q", ext.RPCID, "req-1") + } + if ext.SessionID != "sess-abc" { + t.Errorf("SessionID = %q, want %q", ext.SessionID, "sess-abc") + } + if ext.MessageID != "msg-001" { + t.Errorf("MessageID = %q, want %q", ext.MessageID, "msg-001") + } + if ext.Role != "user" { + t.Errorf("Role = %q, want %q", ext.Role, "user") + } + if len(ext.Parts) != 1 { + t.Fatalf("Parts len = %d, want 1", len(ext.Parts)) + } + if ext.Parts[0].Kind != "text" { + t.Errorf("Parts[0].Kind = %q, want %q", ext.Parts[0].Kind, "text") + } + if ext.Parts[0].Content != "Hello agent" { + t.Errorf("Parts[0].Content = %q, want %q", ext.Parts[0].Content, "Hello agent") + } +} + +func TestA2AParser_MessageStream(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/stream","id":42,"params":{"message":{"role":"user","parts":[{"kind":"text","text":"What is the weather?"}],"messageId":"msg-002"},"sessionId":"sess-xyz"}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if ext.Method != "message/stream" { + t.Errorf("Method = %q, want %q", ext.Method, "message/stream") + } + if ext.RPCID != float64(42) { + t.Errorf("RPCID = %v, want 42", ext.RPCID) + } + if ext.SessionID != "sess-xyz" { + t.Errorf("SessionID = %q, want %q", ext.SessionID, "sess-xyz") + } +} + +func TestA2AParser_MultipleParts(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/send","id":"req-3","params":{"message":{"role":"user","parts":[{"kind":"text","text":"First"},{"kind":"text","text":"Second"}],"messageId":"msg-003"}}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if len(ext.Parts) != 2 { + t.Fatalf("Parts len = %d, want 2", len(ext.Parts)) + } + if ext.Parts[0].Content != "First" { + t.Errorf("Parts[0].Content = %q, want %q", ext.Parts[0].Content, "First") + } + if ext.Parts[1].Content != "Second" { + t.Errorf("Parts[1].Content = %q, want %q", ext.Parts[1].Content, "Second") + } +} + +func TestA2AParser_FilePart(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/send","id":"req-4","params":{"message":{"role":"user","parts":[{"kind":"file","data":"base64-encoded-content"}],"messageId":"msg-004"}}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if len(ext.Parts) != 1 { + t.Fatalf("Parts len = %d, want 1", len(ext.Parts)) + } + if ext.Parts[0].Kind != "file" { + t.Errorf("Parts[0].Kind = %q, want %q", ext.Parts[0].Kind, "file") + } + if ext.Parts[0].Content != "base64-encoded-content" { + t.Errorf("Parts[0].Content = %q, want %q", ext.Parts[0].Content, "base64-encoded-content") + } +} + +func TestA2AParser_AnyMethod(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"tasks/get","id":"req-5","params":{"taskId":"task-123"}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if ext.Method != "tasks/get" { + t.Errorf("Method = %q, want %q", ext.Method, "tasks/get") + } + if ext.RPCID != "req-5" { + t.Errorf("RPCID = %v, want %q", ext.RPCID, "req-5") + } + if len(ext.Parts) != 0 { + t.Errorf("Parts should be empty when no params.message, got %d", len(ext.Parts)) + } +} + +func TestA2AParser_FutureMethodWithMessage(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/resume","id":"req-7","params":{"message":{"role":"user","parts":[{"kind":"text","text":"Continue"}],"messageId":"msg-007"},"sessionId":"sess-future"}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if ext.Method != "message/resume" { + t.Errorf("Method = %q, want %q", ext.Method, "message/resume") + } + if ext.SessionID != "sess-future" { + t.Errorf("SessionID = %q, want %q", ext.SessionID, "sess-future") + } + if ext.Role != "user" { + t.Errorf("Role = %q, want %q", ext.Role, "user") + } + if len(ext.Parts) != 1 { + t.Fatalf("Parts len = %d, want 1", len(ext.Parts)) + } + if ext.Parts[0].Content != "Continue" { + t.Errorf("Parts[0].Content = %q, want %q", ext.Parts[0].Content, "Continue") + } +} + +func TestA2AParser_DataPart(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/send","id":"req-8","params":{"message":{"role":"user","parts":[{"kind":"data","data":{"key":"value","num":42}}],"messageId":"msg-008"}}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if len(ext.Parts) != 1 { + t.Fatalf("Parts len = %d, want 1", len(ext.Parts)) + } + if ext.Parts[0].Kind != "data" { + t.Errorf("Parts[0].Kind = %q, want %q", ext.Parts[0].Kind, "data") + } + if ext.Parts[0].Content == "" { + t.Error("Parts[0].Content should not be empty for data part") + } +} + +func TestA2AParser_FilePartURI(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/send","id":"req-9","params":{"message":{"role":"user","parts":[{"kind":"file","uri":"https://example.com/doc.pdf"}],"messageId":"msg-009"}}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if len(ext.Parts) != 1 { + t.Fatalf("Parts len = %d, want 1", len(ext.Parts)) + } + if ext.Parts[0].Kind != "file" { + t.Errorf("Parts[0].Kind = %q, want %q", ext.Parts[0].Kind, "file") + } + if ext.Parts[0].Content != "https://example.com/doc.pdf" { + t.Errorf("Parts[0].Content = %q, want %q", ext.Parts[0].Content, "https://example.com/doc.pdf") + } +} + +func TestA2AParser_NilBody(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{Body: nil} + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + if pctx.Extensions.A2A != nil { + t.Error("Extensions.A2A should be nil when body is nil") + } +} + +func TestA2AParser_EmptyBody(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{Body: []byte{}} + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + if pctx.Extensions.A2A != nil { + t.Error("Extensions.A2A should be nil when body is empty") + } +} + +func TestA2AParser_InvalidJSON(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{Body: []byte("not json")} + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + if pctx.Extensions.A2A != nil { + t.Error("Extensions.A2A should be nil for invalid JSON") + } +} + +func TestA2AParser_MissingMessage(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/send","id":"req-6","params":{"sessionId":"sess-1"}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if ext.Method != "message/send" { + t.Errorf("Method = %q, want %q", ext.Method, "message/send") + } + if ext.SessionID != "sess-1" { + t.Errorf("SessionID = %q, want %q", ext.SessionID, "sess-1") + } + if ext.Role != "" { + t.Errorf("Role = %q, want empty", ext.Role) + } + if len(ext.Parts) != 0 { + t.Errorf("Parts len = %d, want 0", len(ext.Parts)) + } +} + +func TestA2AParser_MalformedParts(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/send","id":"req-10","params":{"message":{"role":"user","parts":[{"kind":0,"text":"bad"},{"text":"missing-kind"},{"kind":"text","text":"valid"}],"messageId":"msg-010"}}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if len(ext.Parts) != 1 { + t.Fatalf("Parts len = %d, want 1 (only the valid part)", len(ext.Parts)) + } + if ext.Parts[0].Kind != "text" { + t.Errorf("Parts[0].Kind = %q, want %q", ext.Parts[0].Kind, "text") + } + if ext.Parts[0].Content != "valid" { + t.Errorf("Parts[0].Content = %q, want %q", ext.Parts[0].Content, "valid") + } +} + +func TestA2AParser_MalformedContentValues(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Body: []byte(`{"jsonrpc":"2.0","method":"message/send","id":"req-11","params":{"message":{"role":"user","parts":[{"kind":"text","text":false},{"kind":"file","data":0,"uri":{}},{"kind":"data","data":null}],"messageId":"msg-011"}}}`), + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + ext := pctx.Extensions.A2A + if ext == nil { + t.Fatal("Extensions.A2A is nil") + } + if len(ext.Parts) != 3 { + t.Fatalf("Parts len = %d, want 3", len(ext.Parts)) + } + if ext.Parts[0].Content != "" { + t.Errorf("text part with bool value: Content = %q, want empty", ext.Parts[0].Content) + } + if ext.Parts[1].Content != "" { + t.Errorf("file part with numeric data and object uri: Content = %q, want empty", ext.Parts[1].Content) + } + if ext.Parts[2].Content != "" { + t.Errorf("data part with null data: Content = %q, want empty", ext.Parts[2].Content) + } +} + +func TestA2AParser_OnResponse(t *testing.T) { + p := NewA2AParser() + action := p.OnResponse(context.Background(), &pipeline.Context{}) + if action.Type != pipeline.Continue { + t.Errorf("OnResponse should return Continue, got %v", action.Type) + } +} diff --git a/authbridge/authlib/plugins/registry.go b/authbridge/authlib/plugins/registry.go index 39aed031b..5033f8626 100644 --- a/authbridge/authlib/plugins/registry.go +++ b/authbridge/authlib/plugins/registry.go @@ -14,6 +14,7 @@ var registry = map[string]PluginFactory{ "jwt-validation": func(a *auth.Auth) pipeline.Plugin { return NewJWTValidation(a) }, "token-exchange": func(a *auth.Auth) pipeline.Plugin { return NewTokenExchange(a) }, "mcp-parser": func(_ *auth.Auth) pipeline.Plugin { return NewMCPParser() }, + "a2a-parser": func(_ *auth.Auth) pipeline.Plugin { return NewA2AParser() }, } // Build constructs a pipeline from an ordered list of plugin names.