feat(authlib): Add ContentSource capability for guardrail plugins

authbridge/authlib/contracts/content.go

@@ -0,0 +1,89 @@
+// Package contracts defines capability interfaces that protocol
+// extensions implement to participate in framework services like content
+// inspection, session bucketing, and summarization. The package is
+// deliberately dependency-free — parser plugins and consumer plugins
+// (guardrails) both import it without importing each other, so no
+// implicit coupling forms between producers and consumers of a
+// capability.
+//
+// Capability interfaces are checked by consumers via type assertion. A
+// protocol extension that doesn't implement one simply isn't offered
+// that capability; guardrails skip it.
+package contracts
+
+// ContentSource is implemented by protocol extensions whose payload
+// contains user-visible text that guardrails might inspect — PII
+// scrubbing, jailbreak detection, content classification, credential
+// leakage scanning, and similar content-oriented checks.
+//
+// Usage on the consumer side:
+//
+//	for _, src := range pctx.ContentSources() {
+//	    for _, f := range src.Fragments() {
+//	        if f.Role == contracts.RoleUser {
+//	            scan(f.Text)
+//	        }
+//	    }
+//	}
+//
+// The guardrail does not import any parser package — it knows only
+// this contract. When a new protocol ships with a Fragments
+// implementation, existing guardrails pick it up automatically.
+//
+// Protocols that carry no inspectable text (control-plane RPCs like
+// MCP initialize, binary protocols, identity-only auth messages) simply
+// do not implement this interface.
+type ContentSource interface {
+	// Fragments returns every inspectable text fragment on the current
+	// request or response phase, in document order. Returns nil or an
+	// empty slice when there's nothing to inspect at this phase.
+	// Fragments with empty Text are filtered by the producer.
+	Fragments() []Fragment
+}
+
+// Fragment is one inspectable piece of text, tagged with the role that
+// produced it.
+//
+// The struct may grow additional fields in future versions (e.g., a
+// Path locator for violation citations, or a Kind hint for
+// format-aware scanners). Use **named-field initialization** when
+// constructing literals — `Fragment{Role: ..., Text: ...}` rather
+// than `Fragment{"user", "..."}` — so existing call sites and tests
+// remain unaffected when fields are added. Tests that compare
+// `[]Fragment` literals with reflect.DeepEqual should likewise rely
+// on named fields so new zero-valued fields compare equal without
+// test churn.
+type Fragment struct {
+	// Role identifies who produced the text. Use the standard values
+	// below when the semantic fit is clear; protocols that don't fit
+	// any standard value may emit their own role strings.
+	Role string
+
+	// Text is the raw text content. Never empty — producers filter
+	// empty fragments before returning.
+	Text string
+}
+
+// Standard Role values. Parsers reference these instead of string
+// literals so spelling is compiler-checked. Guardrails may compare
+// against either the constant or a string literal — the over-the-wire
+// representation is the same.
+//
+//   - RoleUser: end-user input, human-authored.
+//   - RoleAssistant: model or agent output.
+//   - RoleSystem: system prompt or instructions to the model.
+//   - RoleTool: the name of a tool being invoked.
+//   - RoleToolArgs: an argument value passed to a tool invocation.
+//   - RoleToolResult: content returned from a tool invocation.
+//
+// The vocabulary is open. A protocol that invents a role outside this
+// list is valid; guardrails that don't recognize a role treat it per
+// their own policy (typically: scan anyway, or skip).
+const (
+	RoleUser       = "user"
+	RoleAssistant  = "assistant"
+	RoleSystem     = "system"
+	RoleTool       = "tool"
+	RoleToolArgs   = "tool_args"
+	RoleToolResult = "tool_result"
+)

authbridge/authlib/pipeline/content.go

@@ -0,0 +1,253 @@
+package pipeline
+
+import (
+	"encoding/json"
+	"fmt"
+	"log/slog"
+
+	"github.com/kagenti/kagenti-extensions/authbridge/authlib/contracts"
+)
+
+// This file wires the named protocol extensions (A2AExtension,
+// MCPExtension, InferenceExtension) to contracts.ContentSource. The
+// methods live alongside their receiver types rather than with the
+// parser plugins because Go only allows defining methods on a type
+// from that type's own package.
+//
+// Compile-time assertions make the implementation visible to
+// grep/LSP and catch interface drift early.
+var (
+	_ contracts.ContentSource = (*A2AExtension)(nil)
+	_ contracts.ContentSource = (*MCPExtension)(nil)
+	_ contracts.ContentSource = (*InferenceExtension)(nil)
+)
+
+// Fragments implements contracts.ContentSource for A2A messages.
+//
+// Request-phase: iterates message Parts, emitting text and data parts
+// tagged with the message role (normalized: A2A's native "agent" role
+// is rewritten to "assistant" so guardrails match a single vocabulary
+// across A2A and Inference). File parts carry URIs or base64 blobs,
+// not prose; they're skipped.
+//
+// Response-phase: the final artifact is assistant-authored text.
+func (e *A2AExtension) Fragments() []contracts.Fragment {
+	if e == nil {
+		return nil
+	}
+	var out []contracts.Fragment
+
+	role := normalizeA2ARole(e.Role)
+	for _, p := range e.Parts {
+		switch p.Kind {
+		case "text", "data":
+			if p.Content != "" {
+				out = append(out, contracts.Fragment{Role: role, Text: p.Content})
+			}
+		case "file":
+			// File parts carry URIs or base64 blobs; not inspectable as
+			// prose. A dedicated file-scanning guardrail can type-assert
+			// to *A2AExtension and access the raw Parts directly.
+		}
+	}
+
+	if e.Artifact != "" {
+		out = append(out, contracts.Fragment{Role: contracts.RoleAssistant, Text: e.Artifact})
+	}
+	return out
+}
+
+// normalizeA2ARole rewrites A2A's native role vocabulary (user/agent)
+// to the standard cross-protocol vocabulary used by Inference and by
+// the Role constants in authlib/contracts (user/assistant). Uniform
+// role names across every protocol that implements ContentSource is
+// the design goal: a jailbreak detector, PII scrubber, or content
+// classifier compares `f.Role == contracts.RoleUser` once and works
+// on A2A, MCP, and Inference without per-protocol branching.
+//
+// Fragment.Role IS the role information consumers need — no
+// type-assertion required. The only situation where reading the raw
+// A2A-native string ("agent") would be appropriate is A2A-specialized
+// tooling (e.g., an A2A-protocol inspector that wants to display the
+// wire-level value verbatim); such callers hold a concrete
+// *A2AExtension and read .Role directly. Framework-generic consumers
+// should not do that.
+//
+// Conceptual home: this helper (and Fragments on A2AExtension) more
+// naturally belongs in the a2a-parser package — it's A2A-specific
+// logic that the framework itself has no stake in. It lives here
+// because Go requires methods on a type to be declared in the type's
+// own package, and A2AExtension lives in pipeline/ as a named slot
+// on Extensions. A planned protocols-registry refactor will move the
+// extension types into their parser packages (removing the named
+// slots on Extensions); this helper moves with A2AExtension at that
+// point. The equivalent move applies to MCPExtension.Fragments and
+// InferenceExtension.Fragments below.
+func normalizeA2ARole(r string) string {
+	switch r {
+	case "agent":
+		return contracts.RoleAssistant
+	case "user":
+		return contracts.RoleUser
+	default:
+		// Unknown / unset roles pass through so guardrails at least
+		// see something to filter on. Empty string is tolerated too.
+		return r
+	}
+}
+
+// Fragments implements contracts.ContentSource for MCP messages.
+//
+// Request-phase: only tools/call is modeled — it's the one MCP method
+// carrying user-intent content. Control-plane calls (initialize, ping,
+// tools/list, resources/list, etc.) return nil. The tool name is
+// emitted as role=tool; each argument value is emitted as
+// role=tool_args, JSON-stringified if non-string. On JSON-RPC errors,
+// the error message is emitted as role=tool_result so guardrails see
+// content that could leak through the error channel (credentials,
+// stack traces, PII) the same as they see normal tool output.
+//
+// Response-phase: MCP tool results are conventionally shaped as
+// {"content": [{"type":"text","text":"..."}, {"type":"image",...}, ...]}.
+// Text items are emitted with role=tool_result; non-text items are
+// skipped as not inspectable.
+//
+// Type-assertion misses on Params["arguments"] / Result["content"] /
+// content-items get a DEBUG log. These shapes are what the MCP parser
+// produced from a JSON-validated body, so misses typically indicate
+// a malformed or protocol-non-conforming payload worth surfacing to
+// operators debugging an odd client — rather than a silent skip.
+func (e *MCPExtension) Fragments() []contracts.Fragment {
+	if e == nil {
+		return nil
+	}
+	var out []contracts.Fragment
+
+	if e.Method == "tools/call" && e.Params != nil {
+		if name, _ := e.Params["name"].(string); name != "" {
+			out = append(out, contracts.Fragment{Role: contracts.RoleTool, Text: name})
+		}
+		if raw, present := e.Params["arguments"]; present {
+			args, ok := raw.(map[string]any)
+			if !ok {
+				slog.Debug("pipeline/content: MCP tools/call arguments not a map; skipping",
+					"type", fmt.Sprintf("%T", raw))
+			} else {
+				for _, v := range args {
+					text := stringifyAny(v)
+					if text != "" {
+						out = append(out, contracts.Fragment{Role: contracts.RoleToolArgs, Text: text})
+					}
+				}
+			}
+		}
+	}
+
+	// Tool result content, when present. Errors that arrive as the
+	// JSON-RPC-level Err field (not content[]) are handled below.
+	if e.Result != nil {
+		if raw, present := e.Result["content"]; present {
+			items, ok := raw.([]any)
+			if !ok {
+				slog.Debug("pipeline/content: MCP result.content not an array; skipping",
+					"type", fmt.Sprintf("%T", raw))
+			} else {
+				for i, it := range items {
+					m, ok := it.(map[string]any)
+					if !ok {
+						slog.Debug("pipeline/content: MCP result.content item not an object; skipping",
+							"index", i, "type", fmt.Sprintf("%T", it))
+						continue
+					}
+					if m["type"] != "text" {
+						continue
+					}
+					if t, _ := m["text"].(string); t != "" {
+						out = append(out, contracts.Fragment{Role: contracts.RoleToolResult, Text: t})
+					}
+				}
+			}
+		}
+	}
+
+	// JSON-RPC-level errors carry a Message that may contain
+	// inspectable text (leaked credentials, stack traces, PII from
+	// failed DB lookups, etc.). Emit as tool_result so a PII scrubber
+	// or credential detector covers the error channel uniformly.
+	if e.Err != nil && e.Err.Message != "" {
+		out = append(out, contracts.Fragment{Role: contracts.RoleToolResult, Text: e.Err.Message})
+	}
+
+	return out
+}
+
+// Fragments implements contracts.ContentSource for Inference messages.
+//
+// Request-phase: walks the Messages slice. OpenAI's role vocabulary
+// maps to our standard values directly, except that OpenAI's "tool"
+// role marks a tool RESULT in the conversation history — remapped to
+// "tool_result" so it lines up with MCP's tool result semantics.
+//
+// Response-phase: the model's completion (assistant) plus any tool
+// calls the model emitted (tool name + arguments as separate fragments).
+func (e *InferenceExtension) Fragments() []contracts.Fragment {
+	if e == nil {
+		return nil
+	}
+	// Use a nil slice so an empty result returns nil, consistent with
+	// A2AExtension.Fragments and MCPExtension.Fragments — append
+	// tolerates nil and the cap hint isn't measurable on this path.
+	var out []contracts.Fragment
+
+	for _, m := range e.Messages {
+		if m.Content == "" {
+			continue
+		}
+		role := m.Role
+		if role == "tool" {
+			role = contracts.RoleToolResult
+		}
+		out = append(out, contracts.Fragment{Role: role, Text: m.Content})
+	}
+
+	if e.Completion != "" {
+		out = append(out, contracts.Fragment{Role: contracts.RoleAssistant, Text: e.Completion})
+	}
+	for _, tc := range e.ToolCalls {
+		if tc.Name != "" {
+			out = append(out, contracts.Fragment{Role: contracts.RoleTool, Text: tc.Name})
+		}
+		if tc.Arguments != "" {
+			out = append(out, contracts.Fragment{Role: contracts.RoleToolArgs, Text: tc.Arguments})
+		}
+	}
+
+	return out
+}
+
+// stringifyAny renders an arbitrary argument value as a string suitable
+// for text scanning. Strings pass through unchanged; anything else goes
+// through JSON so nested maps / slices become flat inspectable text.
+//
+// Precondition: v should be JSON-origin data (values that came out of
+// json.Unmarshal into map[string]any / []any / primitives). Those
+// round-trip through json.Marshal without error in practice. Values
+// with unmarshalable types (channels, funcs, cyclic refs) will hit the
+// error path — the function returns "" and logs at DEBUG so the skip
+// is observable in verbose runs rather than silent. Callers filter
+// empty strings regardless.
+func stringifyAny(v any) string {
+	if v == nil {
+		return ""
+	}
+	if s, ok := v.(string); ok {
+		return s
+	}
+	b, err := json.Marshal(v)
+	if err != nil {
+		slog.Debug("pipeline/content: stringifyAny marshal error, returning empty",
+			"error", err)
+		return ""
+	}
+	return string(b)
+}