feat(go/a2ui): add A2UI plugin for streaming interactive UI - #5850
feat(go/a2ui): add A2UI plugin for streaming interactive UI#5850pavelgj wants to merge 2 commits into
Conversation
Add experimental A2UI ("Agent to UI") plugin for Genkit Go, enabling
agents to stream rich, interactive UI surfaces via the JSON-based A2UI
protocol.
The plugin provides middleware that injects catalog capabilities into
the system prompt, intercepts model output to extract and validate a2ui
fenced blocks, and rewrites them into a2ui data parts using the mime
type `application/a2ui+json`. This maps 1:1 onto the A2A binding, making
envelopes byte-compatible with the JS and Dart plugins.
Includes a bundled basic catalog, custom catalog support via the Genkit
registry, configurable validation modes, and helpers for extracting
envelopes from message parts.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the experimental a2ui ("Agent to UI") plugin and middleware for Genkit Go, enabling agents to stream rich, interactive UI surfaces. It updates the core Part struct to support arbitrary structured data and adds the necessary parser, catalog, and loader components. The review feedback provides valuable recommendations to ensure thread safety in surfaceIDReplay using a mutex, optimize streaming performance by reducing character lag and caching component lookups, prevent potential nil-pointer panics, and improve the clarity of the README documentation.
| type surfaceIDReplay struct { | ||
| base func() string | ||
| generated []string | ||
| cursor int | ||
| } | ||
|
|
||
| func replayableSurfaceIDs(base func() string) *surfaceIDReplay { | ||
| return &surfaceIDReplay{base: base} | ||
| } | ||
|
|
||
| func (r *surfaceIDReplay) next() string { | ||
| id := r.base() | ||
| r.generated = append(r.generated, id) | ||
| return id | ||
| } | ||
|
|
||
| func (r *surfaceIDReplay) replayNext() string { | ||
| if r.cursor < len(r.generated) { | ||
| id := r.generated[r.cursor] | ||
| r.cursor++ | ||
| return id | ||
| } | ||
| return r.next() | ||
| } | ||
|
|
||
| func (r *surfaceIDReplay) reset() { r.cursor = 0 } |
There was a problem hiding this comment.
surfaceIDReplay is not thread-safe. If the streaming callback is invoked concurrently or across multiple goroutines, concurrent access to r.generated and r.cursor can cause a data race. Adding a sync.Mutex ensures thread safety.
type surfaceIDReplay struct {
mu sync.Mutex
base func() string
generated []string
cursor int
}
func replayableSurfaceIDs(base func() string) *surfaceIDReplay {
return &surfaceIDReplay{base: base}
}
func (r *surfaceIDReplay) next() string {
r.mu.Lock()
defer r.mu.Unlock()
return r.nextLocked()
}
func (r *surfaceIDReplay) nextLocked() string {
id := r.base()
r.generated = append(r.generated, id)
return id
}
func (r *surfaceIDReplay) replayNext() string {
r.mu.Lock()
defer r.mu.Unlock()
if r.cursor 8< len(r.generated) {
id := r.generated[r.cursor]
r.cursor++
return id
}
return r.nextLocked()
}
func (r *surfaceIDReplay) reset() {
r.mu.Lock()
defer r.mu.Unlock()
r.cursor = 0
}There was a problem hiding this comment.
Done in 6ab2519 - added a sync.Mutex to surfaceIDReplay (with a nextLocked helper). In practice the streamed parse and the final-message parse never overlap, but the lock is cheap insurance. Verified with go test -race.
| import ( | ||
| "context" | ||
| "encoding/json" | ||
| "fmt" | ||
| "strings" | ||
|
|
||
| "github.com/firebase/genkit/go/ai" | ||
| "github.com/firebase/genkit/go/core/api" | ||
| "github.com/firebase/genkit/go/genkit" | ||
| "github.com/google/uuid" | ||
| ) |
There was a problem hiding this comment.
Add the sync package to the imports to support thread-safe operations in surfaceIDReplay.
| import ( | |
| "context" | |
| "encoding/json" | |
| "fmt" | |
| "strings" | |
| "github.com/firebase/genkit/go/ai" | |
| "github.com/firebase/genkit/go/core/api" | |
| "github.com/firebase/genkit/go/genkit" | |
| "github.com/google/uuid" | |
| ) | |
| import ( | |
| "context" | |
| "encoding/json" | |
| "fmt" | |
| "strings" | |
| "sync" | |
| "github.com/firebase/genkit/go/ai" | |
| "github.com/firebase/genkit/go/core/api" | |
| "github.com/firebase/genkit/go/genkit" | |
| "github.com/google/uuid" | |
| ) |
| type streamParser struct { | ||
| opts parserOptions | ||
| buffer string | ||
| inBlock bool | ||
| // currentSurfaceID is the stable surface id for the current block | ||
| // (placeholders map to this). | ||
| currentSurfaceID string | ||
| hasSurfaceID bool | ||
| } | ||
|
|
||
| func newStreamParser(opts parserOptions) *streamParser { | ||
| if opts.validate == "" { | ||
| opts.validate = ValidateWarn | ||
| } | ||
| if opts.version == "" { | ||
| opts.version = DefaultVersion | ||
| } | ||
| return &streamParser{opts: opts} | ||
| } |
There was a problem hiding this comment.
To avoid rebuilding the component names map on every updateComponents envelope validation, we can cache the map in streamParser during initialization.
type streamParser struct {
opts parserOptions
buffer string
inBlock bool
// currentSurfaceID is the stable surface id for the current block
// (placeholders map to this).
currentSurfaceID string
hasSurfaceID bool
knownComponents map[string]bool
}
func newStreamParser(opts parserOptions) *streamParser {
if opts.validate == "" {
opts.validate = ValidateWarn
}
if opts.version == "" {
opts.version = DefaultVersion
}
var known map[string]bool
if opts.catalog != nil {
known = opts.catalog.componentNames()
}
return &streamParser{opts: opts, knownComponents: known}
}There was a problem hiding this comment.
Done in 6ab2519 - the component-name set is now computed once in newStreamParser and cached on the parser as knownComponents.
| keep := maxPartialFence | ||
| if keep > len(p.buffer) { | ||
| keep = len(p.buffer) | ||
| } | ||
| safeLen := len(p.buffer) - keep | ||
| if safeLen > 0 { | ||
| proseBuf += p.buffer[:safeLen] | ||
| p.buffer = p.buffer[safeLen:] | ||
| } |
There was a problem hiding this comment.
The current implementation always holds back maxPartialFence (8) characters during streaming, which introduces a constant 8-character lag in the streamed output even when the tail of the buffer does not resemble a fence. We can optimize this by only holding back characters if the buffer ends with a potential prefix of the fence (i.e., contains a backtick within the last maxPartialFence characters).
keep := 0
if idx := strings.LastIndexByte(p.buffer, 96); idx >= 0 {
dist := len(p.buffer) - idx
if dist <= maxPartialFence {
keep = dist
}
}
safeLen := len(p.buffer) - keep
if safeLen > 0 {
proseBuf += p.buffer[:safeLen]
p.buffer = p.buffer[safeLen:]
}There was a problem hiding this comment.
Not taking this one - the suggested code is subtly incorrect. It keeps from the last backtick (LastIndexByte), so for a fence split mid-opener like "hello a2ui" it would flush "hello ``" and retain only "`a2ui"; once reassembled with the next chunk the opener has a single backtick and no longer matches a2ui, so the block is missed. The constant 8-char holdback is intentional (it mirrors the JS parser) and the lag is at most a few chars until the next chunk arrives, so I'd rather keep the simple, correct version.
| func (p *streamParser) validateComponents(components any) string { | ||
| catalog := p.opts.catalog | ||
| if catalog == nil { | ||
| return "" | ||
| } | ||
| arr, ok := components.([]any) | ||
| if !ok { | ||
| return "updateComponents.components must be an array." | ||
| } | ||
| known := catalog.componentNames() |
There was a problem hiding this comment.
Use the cached knownComponents map instead of calling catalog.componentNames() repeatedly.
| func (p *streamParser) validateComponents(components any) string { | |
| catalog := p.opts.catalog | |
| if catalog == nil { | |
| return "" | |
| } | |
| arr, ok := components.([]any) | |
| if !ok { | |
| return "updateComponents.components must be an array." | |
| } | |
| known := catalog.componentNames() | |
| func (p *streamParser) validateComponents(components any) string { | |
| catalog := p.opts.catalog | |
| if catalog == nil { | |
| return "" | |
| } | |
| arr, ok := components.([]any) | |
| if !ok { | |
| return "updateComponents.components must be an array." | |
| } | |
| known := p.knownComponents |
There was a problem hiding this comment.
Done - validateComponents now uses the cached knownComponents map (6ab2519).
| if v := genkit.LookupValue(g, catalogRegistryKey(id)); v != nil { | ||
| if c, ok := v.(*Catalog); ok { | ||
| return c, nil | ||
| } | ||
| return nil, fmt.Errorf("a2ui: registry value %q is not an *a2ui.Catalog", catalogRegistryKey(id)) |
There was a problem hiding this comment.
If a catalog is registered as a value instead of a pointer, v.(*Catalog) will fail. We can make this more robust by also checking for Catalog value and taking its address.
if v := genkit.LookupValue(g, catalogRegistryKey(id)); v != nil {
if c, ok := v.(*Catalog); ok {
return c, nil
}
if c, ok := v.(Catalog); ok {
return &c, nil
}
return nil, fmt.Errorf("a2ui: registry value %q is not an *a2ui.Catalog", catalogRegistryKey(id))There was a problem hiding this comment.
Leaving as-is. The only registration path (LoadCatalog) always stores a *Catalog, so a Catalog-by-value case would be guarding against a misuse this package never creates. The current type assertion already returns a clear error for anything unexpected.
| func RenderCatalogInstructions(catalog *Catalog) string { | ||
| var componentDocs strings.Builder | ||
| for i, c := range catalog.Components { |
There was a problem hiding this comment.
Since RenderCatalogInstructions is an exported function, it can be called with a nil catalog, which would cause a panic. Adding a nil check prevents this.
| func RenderCatalogInstructions(catalog *Catalog) string { | |
| var componentDocs strings.Builder | |
| for i, c := range catalog.Components { | |
| func RenderCatalogInstructions(catalog *Catalog) string { | |
| if catalog == nil { | |
| return "" | |
| } | |
| var componentDocs strings.Builder | |
| for i, c := range catalog.Components { |
There was a problem hiding this comment.
Done - added a nil-catalog guard returning "" in 6ab2519.
| | `Instructions` | `InstructionsSystem` | Where to inject catalog capabilities. `InstructionsNone` injects nothing. | | ||
| | `Validate` | `ValidateWarn` | Validate emitted envelopes. `ValidateWarn` logs and drops bad blocks; `ValidateStrict` returns an error; `ValidateOff` skips checking. | |
There was a problem hiding this comment.
It would be clearer to specify the actual string values (e.g., "system" and "warn") as the defaults in the table, rather than the Go constant names, to make it easier for users configuring the plugin via JSON/YAML or other non-Go environments.
| | `Instructions` | `InstructionsSystem` | Where to inject catalog capabilities. `InstructionsNone` injects nothing. | | |
| | `Validate` | `ValidateWarn` | Validate emitted envelopes. `ValidateWarn` logs and drops bad blocks; `ValidateStrict` returns an error; `ValidateOff` skips checking. | | |
| | Instructions | "system" | Where to inject catalog capabilities. InstructionsNone injects nothing. | | |
| | Validate | "warn" | Validate emitted envelopes. ValidateWarn logs and drops bad blocks; ValidateStrict returns an error; ValidateOff skips checking. | |
There was a problem hiding this comment.
Done - the options table now shows the string values ("system", "warn", etc.) instead of the Go constant names (6ab2519).
- Guard surfaceIDReplay with a mutex for thread safety. - Cache the catalog component-name set in the stream parser instead of rebuilding it per updateComponents envelope. - Add a nil-catalog guard to RenderCatalogInstructions. - Show string default values in the README options table.
1 participant
Add experimental A2UI ("Agent to UI") plugin for Genkit Go, enabling agents to stream rich, interactive UI surfaces via the JSON-based A2UI protocol.
The plugin provides middleware that injects catalog capabilities into the system prompt, intercepts model output to extract and validate a2ui fenced blocks, and rewrites them into a2ui data parts using the mime type
application/a2ui+json. This maps 1:1 onto the A2A binding, making envelopes byte-compatible with the JS and Dart plugins.Includes a bundled basic catalog, custom catalog support via the Genkit registry, configurable validation modes, and helpers for extracting envelopes from message parts.