feat(go): experimental agent middleware for sub-agent delegation and artifacts

.gitignore

@@ -31,10 +31,6 @@ js/testapps/firebase-functions-sample1/public/config.js
 .genkit
 js/**/.genkit
 samples/**/.genkit
-go/**/.genkit
-# Compiled Go sample binaries (e.g. `go build ./samples/<name>` writes go/<name>)
-/go/custom-agent
-/go/x-agent-interrupts
 ui-debug.log
 firebase-debug.log
 firestore-debug.log

go/.gitignore

@@ -0,0 +1,4 @@
+# Compiled sample binaries produced by `go build ./samples/<name>` from this
+# directory (e.g. ./samples/basic-agents). They are build artifacts, not source.
+/basic-agents
+/basic-agents-server

go/README.md

@@ -298,6 +298,59 @@ Detach requires a store that implements `SnapshotSubscriber` (both bundled local
 
 [See full example](samples/basic-agents)
 
+### Delegate to Sub-Agents
+
+The experimental `Agents` middleware (in `plugins/middleware/exp`) lets one agent delegate to others. It injects one `delegate_to_<name>` tool per sub-agent and a `<sub-agents>` listing into the orchestrator's system prompt, then runs the chosen sub-agent and returns its result when the model calls the tool. Each sub-agent's `aix.WithDescription` (captured by `agent.Ref()`) tells the orchestrator when to reach for it:
+
+```go
+import (
+    "github.com/firebase/genkit/go/ai"
+    aix "github.com/firebase/genkit/go/ai/exp"
+    "github.com/firebase/genkit/go/ai/exp/localstore"
+    genkitx "github.com/firebase/genkit/go/genkit/exp"
+    middlewarex "github.com/firebase/genkit/go/plugins/middleware/exp"
+)
+
+researcher := genkitx.DefineAgent(g, "researcher",
+    aix.InlinePrompt{
+        ai.WithModelName("googleai/gemini-flash-latest"),
+        ai.WithSystem("You are a thorough research assistant. Summarize well-sourced findings."),
+    },
+    aix.WithDescription[any]("Researches a topic and summarizes well-sourced findings."),
+)
+
+engineer := genkitx.DefineAgent(g, "engineer",
+    aix.InlinePrompt{
+        ai.WithModelName("googleai/gemini-flash-latest"),
+        ai.WithSystem("You are an expert programmer. Write clean, well-commented code."),
+    },
+    aix.WithDescription[any]("Writes and explains code."),
+)
+
+// The orchestrator delegates instead of answering directly: the model calls
+// delegate_to_researcher / delegate_to_engineer and the middleware runs them.
+orchestrator := genkitx.DefineAgent(g, "orchestrator",
+    aix.InlinePrompt{
+        ai.WithModelName("googleai/gemini-flash-latest"),
+        ai.WithSystem("You are a project coordinator. Delegate to the right sub-agent, " +
+            "then synthesize a final answer."),
+        ai.WithUse(&middlewarex.Agents{
+            Agents:         []aix.AgentRef{researcher.Ref(), engineer.Ref()},
+            MaxDelegations: 5, // cap delegation tool calls per turn (0 = unlimited)
+            HistoryLength:  4, // recent messages forwarded to client-managed sub-agents
+        }),
+    },
+    aix.WithSessionStore(localstore.NewInMemorySessionStore[any]()),
+)
+
+out, _ := orchestrator.RunText(ctx, "Research goroutine scheduling, then sketch a worker pool.")
+fmt.Println(out.Message.Text())
+```
+
+Sub-agents are named by `aix.AgentRef`, either captured from an agent value with `agent.Ref()` or written by hand (`aix.AgentRef{Name: "researcher"}`). The middleware composes with the `Artifacts` middleware: give the sub-agents `&middlewarex.Artifacts{}` so they can save output, set `ArtifactStrategy: middlewarex.ArtifactStrategySession` to merge those artifacts into the orchestrator's session instead of inlining them in the tool result, and add `&middlewarex.Artifacts{Readonly: true}` on the orchestrator so it can review them before answering.
+
+[See full example](samples/basic-agents)
+
 ### Serve Agents over HTTP
 
 An `Agent` is an `api.BidiAction`, so it serves over HTTP one turn per request. The `genkit/exp` package lays out a default route surface for every registered agent, including the snapshot companion endpoints for store-backed agents:

go/ai/exp/agent.go

@@ -39,6 +39,7 @@ import (
 	"github.com/firebase/genkit/go/core/logger"
 	"github.com/firebase/genkit/go/core/tracing"
 	"github.com/firebase/genkit/go/internal/base"
+	"github.com/firebase/genkit/go/internal/genkitbridge"
 	"github.com/google/uuid"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
@@ -727,6 +728,24 @@ func DefineAgent[State any](
 	return DefineCustomAgent(r, name, agentLoop[State](r, p, nil), opts...)
 }
 
+// genkitContextSeed returns a context decorator that seeds the host Genkit
+// instance into each agent invocation, so the agent's prompt, tools, and
+// middleware can retrieve it via genkit.FromContext and resolve or run other
+// actions. The instance is reconstructed from r by the genkit package through
+// the internal bridge, so the registry-level constructors below wire seeding up
+// themselves and there is no public option for it.
+//
+// It returns nil when the genkit package is not linked into the build, leaving
+// an agent defined directly on a bare registry untouched.
+func genkitContextSeed(r api.Registry) func(context.Context) context.Context {
+	if genkitbridge.SeedContextForRegistry == nil {
+		return nil
+	}
+	return func(ctx context.Context) context.Context {
+		return genkitbridge.SeedContextForRegistry(ctx, r)
+	}
+}
+
 // DefinePromptAgent defines a prompt-backed agent and registers it, sourcing
 // its prompt from the registry by name. Each turn renders the prompt, appends
 // conversation history, calls the model with streaming, and updates session
@@ -754,6 +773,9 @@ func DefinePromptAgent[State any](
 	name string,
 	opts ...PromptAgentOption[State],
 ) *Agent[State] {
+	if seed := genkitContextSeed(r); seed != nil {
+		opts = append(opts, &agentOptions[State]{contextFunc: seed})
+	}
 	cfg := &promptAgentOptions[State]{}
 	for _, opt := range opts {
 		if err := opt.applyPromptAgent(cfg); err != nil {
@@ -818,7 +840,7 @@ func newCustomAgent[State any](
 	cfg *agentOptions[State],
 ) *Agent[State] {
 	// Typed under ActionTypeAgent so agents surface as their own action
-	// kind rather than as flows (genkit.ListAgents vs ListFlows). Built on
+	// kind rather than as flows (genkit/exp.ListAgents vs genkit.ListFlows). Built on
 	// NewBidiAction so the agent capability metadata is set at construction
 	// time; actions must be immutable once registered. WithFlowContext
 	// below preserves the flow-context wrapping that makes core.Run work
@@ -841,6 +863,12 @@ func newCustomAgent[State any](
 			outCh chan<- *AgentStreamChunk,
 		) (*AgentOutput[State], error) {
 			ctx = core.WithFlowContext(ctx, name)
+			// Apply any context decorators (e.g. the genkit package seeding its
+			// instance) before the runtime derives the per-turn work context, so
+			// the decorated values reach each turn's prompt, tools, and middleware.
+			if cfg.contextFunc != nil {
+				ctx = cfg.contextFunc(ctx)
+			}
 			rt, err := newAgentRuntime(ctx, name, cfg, in, inCh, outCh)
 			if err != nil {
 				// Init failures (a rejected init payload, a failed
@@ -880,6 +908,9 @@ func DefineCustomAgent[State any](
 	fn AgentFunc[State],
 	opts ...AgentOption[State],
 ) *Agent[State] {
+	if seed := genkitContextSeed(r); seed != nil {
+		opts = append(opts, &agentOptions[State]{contextFunc: seed})
+	}
 	a := NewCustomAgent(name, fn, opts...)
 	a.Register(r)
 	return a

go/ai/exp/option.go

@@ -111,6 +111,11 @@ type agentOptions[State any] struct {
 	transform       StateTransform[State]
 	streamTransform StreamTransform
 	description     string
+	// contextFunc decorates each invocation's context once before the turn
+	// loop runs. It has no public option: the registry-level constructors set
+	// it internally to seed the genkit instance (see genkitContextSeed in
+	// agent.go), so callers reach the instance via genkit.FromContext.
+	contextFunc func(context.Context) context.Context
 }
 
 func (o *agentOptions[State]) applyAgent(opts *agentOptions[State]) error {
@@ -138,6 +143,12 @@ func (o *agentOptions[State]) applyAgent(opts *agentOptions[State]) error {
 		}
 		opts.description = o.description
 	}
+	if o.contextFunc != nil {
+		// Seeded internally by the registry-level constructors
+		// (genkitContextSeed), at most once per agent, so this is a plain set
+		// rather than a compose of multiple decorators.
+		opts.contextFunc = o.contextFunc
+	}
 	return nil
 }
 

go/ai/exp/ref.go

@@ -0,0 +1,49 @@
+// 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.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package exp
+
+// AgentRef refers to an agent by name, optionally carrying a description. It is
+// the agent analog of [ai.ModelRef] / [ai.ToolRef]: a small, JSON-serializable
+// value that names an agent for resolution against a registry. Like those, it
+// resolves by name (the path the Dev UI, HTTP serving, and ListAgents all use),
+// so the referenced agent must be registered wherever the ref is consumed.
+//
+// Build one by name with a struct literal, or derive it from an agent value
+// with [Agent.Ref], which fills in the name and description so callers need not
+// restate either:
+//
+//	aix.AgentRef{Name: "researcher"}
+//	coderAgent.Ref()
+type AgentRef struct {
+	// Name identifies the agent, resolved as /agent/<Name>. Required.
+	Name string `json:"name"`
+	// Description is a human-readable description used by consumers that list
+	// agents (e.g. the agents middleware's system prompt). [Agent.Ref] fills it
+	// from the agent's descriptor. Optional.
+	Description string `json:"description,omitempty"`
+}
+
+// Ref returns an [AgentRef] for this agent, capturing its name and description
+// so callers can reference it without restating either, and without a name
+// string that can drift from the agent. Resolution remains by name, so the
+// agent must be registered (as [DefineAgent] does) wherever the ref is used.
+func (a *Agent[State]) Ref() AgentRef {
+	return AgentRef{
+		Name:        a.Name(),
+		Description: a.Desc().Description,
+	}
+}

go/ai/exp/session.go

@@ -510,3 +510,26 @@ func SessionFromContext[State any](ctx context.Context) *Session[State] {
 	session, _ := sessionCtxKey.FromContext(ctx).(*Session[State])
 	return session
 }
+
+// ArtifactStore is the State-agnostic view of a session's artifact collection.
+// Every [Session] satisfies it regardless of its State type, since artifact
+// operations do not touch custom state. Middleware and tools that work with
+// artifacts without knowing the agent's State type use it via
+// [ArtifactStoreFromContext], where [SessionFromContext] cannot help because it
+// requires the concrete State.
+type ArtifactStore interface {
+	// Artifacts returns a snapshot of the session's current artifacts.
+	Artifacts() []*Artifact
+	// AddArtifacts adds artifacts, replacing any existing artifact of the same
+	// name.
+	AddArtifacts(artifacts ...*Artifact)
+}
+
+// ArtifactStoreFromContext returns the active session's artifacts as a
+// State-agnostic [ArtifactStore], or nil if there is no active session in ctx.
+// Unlike [SessionFromContext] it does not require knowing the session's State
+// type, so it is the accessor for middleware and tools.
+func ArtifactStoreFromContext(ctx context.Context) ArtifactStore {
+	store, _ := sessionCtxKey.FromContext(ctx).(ArtifactStore)
+	return store
+}

go/ai/exp/tools.go

@@ -20,9 +20,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"reflect"
 
 	"github.com/firebase/genkit/go/ai"
 	"github.com/firebase/genkit/go/ai/exp/tool"
+	"github.com/firebase/genkit/go/core"
 	"github.com/firebase/genkit/go/core/api"
 	"github.com/firebase/genkit/go/internal/base"
 )
@@ -45,15 +47,42 @@ type Tool[In, Out any] struct {
 	inner *ai.ToolDef[In, *ai.MultipartToolResponse] // DEPRECATED(breaking): remove wrapper; Tool owns the action directly.
 }
 
-// DEPRECATED(breaking): All methods below exist only to implement ai.Tool by
-// delegating to the wrapped ai.ToolDef. With breaking changes, Tool would own
-// the action directly and implement these natively without delegation.
+// DEPRECATED(breaking): The methods below exist to implement ai.Tool on top of
+// the wrapped ai.ToolDef. Most are pure delegation; Definition additionally
+// restores the real output schema (see its comment). With breaking changes, Tool
+// would own the action directly and implement these natively, inferring the
+// output schema from Out without the override.
 
 // Name returns the name of the tool.
 func (t *Tool[In, Out]) Name() string { return t.inner.Name() }
 
 // Definition returns the [ai.ToolDefinition] for this tool.
-func (t *Tool[In, Out]) Definition() *ai.ToolDefinition { return t.inner.Definition() }
+//
+// The inner tool is built on [ai.NewMultipartTool], whose function returns
+// *[ai.MultipartToolResponse], so the inner definition would advertise that
+// envelope as the output schema. We override OutputSchema with the schema
+// inferred from the Out type parameter, making the definition equivalent to what
+// [ai.NewTool] exposes (the real output type) rather than leaking the multipart
+// envelope to the model and Dev UI. Genkit infers schemas with DoNotReference,
+// so the result is fully inlined and needs no registry resolution.
+func (t *Tool[In, Out]) Definition() *ai.ToolDefinition {
+	def := t.inner.Definition()
+	if schema := inferOutputSchema[Out](); schema != nil {
+		def.OutputSchema = schema
+	}
+	return def
+}
+
+// inferOutputSchema returns the inlined JSON schema for the Out type parameter,
+// or nil when Out carries no schema (e.g. any), mirroring how [ai.NewTool]
+// derives its output schema from the output type.
+func inferOutputSchema[Out any]() map[string]any {
+	var zero Out
+	if reflect.TypeOf(zero) == nil {
+		return nil
+	}
+	return core.InferSchemaMap(zero)
+}
 
 // RunRaw runs the tool with raw input.
 func (t *Tool[In, Out]) RunRaw(ctx context.Context, input any) (any, error) {

go/ai/exp/tools_test.go

@@ -18,6 +18,7 @@ package exp
 
 import (
 	"context"
+	"reflect"
 	"strings"
 	"sync"
 	"testing"
@@ -118,6 +119,57 @@ func TestTool_AttachParts(t *testing.T) {
 	}
 }
 
+type reportItem struct {
+	Name string `json:"name"`
+}
+
+type reportOut struct {
+	Title string       `json:"title"`
+	Items []reportItem `json:"items"`
+}
+
+// TestTool_OutputSchemaMatchesClassic guards against the multipart envelope
+// leaking into the tool definition. aix.NewTool wraps ai.NewMultipartTool, whose
+// function returns *MultipartToolResponse; without Definition restoring the real
+// output schema, the model and Dev UI would see the envelope ({content, output,
+// metadata}) instead of the actual Out type. The exp tool must advertise the
+// same output schema ai.NewTool would, including for the interruptible variant
+// (which embeds Tool) and for a nested struct that exercises schema inlining.
+func TestTool_OutputSchemaMatchesClassic(t *testing.T) {
+	classic := ai.NewTool("classic", "d",
+		func(tc *ai.ToolContext, _ weatherIn) (reportOut, error) { return reportOut{}, nil })
+	want := classic.Definition().OutputSchema
+	if want == nil {
+		t.Fatal("ai.NewTool unexpectedly produced a nil output schema")
+	}
+
+	simple := NewTool("exp-simple", "d",
+		func(ctx context.Context, _ weatherIn) (reportOut, error) { return reportOut{}, nil })
+	interruptible := NewInterruptibleTool("exp-interruptible", "d",
+		func(ctx context.Context, _ weatherIn, _ *confirmation) (reportOut, error) { return reportOut{}, nil })
+
+	for _, tc := range []struct {
+		name string
+		got  any
+	}{
+		{"NewTool", simple.Definition().OutputSchema},
+		{"NewInterruptibleTool", interruptible.Definition().OutputSchema},
+	} {
+		if !reflect.DeepEqual(tc.got, want) {
+			t.Errorf("%s output schema = %#v\nwant %#v (matching ai.NewTool)", tc.name, tc.got, want)
+		}
+		// Explicit guard on intent: the real Out fields are present and the
+		// MultipartToolResponse envelope fields are not.
+		props, _ := tc.got.(map[string]any)["properties"].(map[string]any)
+		if _, ok := props["title"]; !ok {
+			t.Errorf("%s output schema missing the real %q field: %#v", tc.name, "title", tc.got)
+		}
+		if _, ok := props["content"]; ok {
+			t.Errorf("%s output schema leaked the multipart envelope (has %q): %#v", tc.name, "content", tc.got)
+		}
+	}
+}
+
 // TestTool_SendPartialNoOpWithoutStreaming confirms SendPartial is a safe no-op
 // when no streaming callback is wired (here, a direct RunRaw).
 func TestTool_SendPartialNoOpWithoutStreaming(t *testing.T) {